gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1231

Browse Source 
Enhance Gates and MLFlowRepository with new functionalities and improvements

- Added a new method `clean_tmp_files` in the Gates class to remove temporary files associated with model retraining.
- Updated MLFlowRepository methods to improve experiment handling, including dynamic parameter logging and model retrieval.
- Refactored model loading methods to streamline the process and enhance error handling.
- Improved logging for model operations and added support for model parameter retrieval.
- Adjusted minimal_retrain workflow to extend timeouts for activities and ensure proper model configuration handling.
This commit is contained in:
vitor-aignosi
2025-10-13 10:36:11 -03:00
parent 7512963e19
commit 49b6e504ae
5 changed files with 291 additions and 274 deletions
Download Patch File Download Diff File Expand all files Collapse all files

493
laborious/utils/repository/model_repository.py
View File

@@ -15,9 +15,11 @@ The repository provides comprehensive functionality for:
"""
from datetime import datetime, timedelta
import traceback
from mlflow.entities import Experiment, experiment
import pandas as pd
import mlflow
from os import makedirs, path, remove, environ
from shutil import rmtree
from sys import path as sys_path
from sientia_do.observability.logger import Logger
import lzma
@@ -29,6 +31,8 @@ from typing import Any
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
ARTIFACTS_PATH = "./tmp/artifacts"
TRANSFORMED_COMPRESSED_PATH = "artifacts/training_transformer.pkl"
PREDICTION_COMPRESSED_PATH = "artifacts/stacking_model.pkl"
class MLFlowRepository():
@@ -107,7 +111,7 @@ class MLFlowRepository():
run_id = latest_version.source.split("/")
return run_id[2]
def get_experiment_by_run_id(self, run_id: str) -> str:
def get_experiment_by_run_id(self, run_id: str) -> Experiment:
"""
Get experiment name by run ID.
@@ -124,9 +128,7 @@ class MLFlowRepository():
experiment_id = run.info.experiment_id
# Get the experiment details using the experiment ID
experiment = mlflow.get_experiment(experiment_id)
experiment_name = experiment.name
return experiment_name
return mlflow.get_experiment(experiment_id)
def get_next_run_name(self, model_name: str) -> str:
"""
@@ -147,7 +149,7 @@ class MLFlowRepository():
next_run_number = len(runs) + 1
return f"{model_name}-{next_run_number}"
def get_experiment(self, experiment_name: str) -> int:
def get_experiment(self, experiment_name: str, create_if_not_exists: bool = False) -> Experiment:
"""
Retrieve MLFlow experiment ID by experiment name.
@@ -167,9 +169,12 @@ class MLFlowRepository():
experiment = mlflow.get_experiment_by_name(experiment_name)
if experiment is None:
raise ValueError(f'Experiment {experiment_name} not found')
if create_if_not_exists:
experiment = mlflow.create_experiment(experiment_name)
else:
raise ValueError(f'Experiment {experiment_name} not found')
return int(experiment.experiment_id)
return experiment
def get_experiment_last_run(self, experiment_id: int) -> str:
"""
@@ -212,6 +217,11 @@ class MLFlowRepository():
return latest_run_id
def get_model_params(self, run_id: str):
"""Obtém os parâmetros de uma run"""
run_info = mlflow.get_run(run_id)
return run_info.data.params
"""
Functions related to download and load models
"""
@@ -232,8 +242,12 @@ class MLFlowRepository():
)
output_dir = f"{ARTIFACTS_PATH}/{model_name}"
if not path.exists(output_dir):
makedirs(output_dir)
full_path = path.join(output_dir, artifact_path)
if path.exists(full_path):
# Remove the directory and create a new one
rmtree(full_path)
makedirs(output_dir, exist_ok=True)
self.logger.info(
f"Downloading artifacts from {run_id} to {output_dir}")
@@ -244,8 +258,7 @@ class MLFlowRepository():
output_dir
)
def load_predict_model(self, model_name: str, flavor: str = 'sklearn',
artifact_path: str | None = None) -> Any:
def load_predict_model(self, model_name: str, flavor: str = 'sklearn') -> Any:
"""
Downloads a predictive model from the MLflow Model Registry.
@@ -261,33 +274,22 @@ class MLFlowRepository():
- The model is fetched from the "production" stage of the MLflow Model Registry.
- Warnings during the model loading process are suppressed.
"""
model_uri = f"models:/{model_name}/production"
if artifact_path:
self.logger.info(
f"Prediction model {model_name} is compressed, loading from {artifact_path}")
model = self.load_model_with_compression(
artifact_path, "prediction")
self.logger.info(
f"Loading prediction model {model_name} from {model_uri}")
if flavor == 'pyfunc':
model = mlflow.pyfunc.load_model(model_uri)
elif flavor == 'sklearn':
model = mlflow.sklearn.load_model(model_uri)
elif flavor == 'pytorch':
model = mlflow.pytorch.load_model(model_uri)
else:
self.logger.info(
f"Prediction model {model_name} is not compressed, loading from {model_uri}")
if flavor == 'pyfunc':
model = mlflow.pyfunc.load_model(model_uri)
elif flavor == 'sklearn':
model = mlflow.sklearn.load_model(model_uri)
elif flavor == 'pytorch':
model = mlflow.pytorch.load_model(model_uri)
else:
raise ValueError(
"Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.")
raise ValueError(
"Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.")
return model
def load_transform_model(self, model_name: str, flavor: str,
artifact_path: str | None = None) -> Any:
def load_transform_model(self, model_name: str, flavor: str) -> Any:
"""
Downloads the latest production version of a specified transformation model.
@@ -306,37 +308,27 @@ class MLFlowRepository():
Exception: If the model run ID or URI cannot be retrieved, or if the
model cannot be loaded.
"""
latest_production_id = self.get_model_run_id(
model_name=model_name, stage="Production"
)
model_uri = self.get_model_uri(
latest_production_id, prediction=False)
self.logger.debug(
f"Model {model_name} is at {model_uri} and latest production id is {latest_production_id}")
if artifact_path:
# Download model artifacts
self.logger.info(
f"Data model {model_name} is compressed, loading from {artifact_path}")
model = self.load_model_with_compression(
artifact_path, "transformer")
self.logger.info(
f"Loading data model {model_name} from {model_uri}")
if flavor == 'sklearn':
model = mlflow.sklearn.load_model(model_uri)
elif flavor == 'pyfunc':
model = mlflow.pyfunc.load_model(model_uri)
elif flavor == 'pytorch':
model = mlflow.pytorch.load_model(model_uri)
else:
self.logger.info(
f"Data model {model_name} is not compressed, loading from {model_uri}")
if flavor == 'sklearn':
model = mlflow.sklearn.load_model(model_uri)
elif flavor == 'pyfunc':
model = mlflow.pyfunc.load_model(model_uri)
elif flavor == 'pytorch':
model = mlflow.pytorch.load_model(model_uri)
else:
raise ValueError(
"Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.")
raise ValueError(
"Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.")
return model
def load_model_with_compression(self, artifact_path: str, type: str) -> Any:
def load_model_with_compression(self, artifact_path: str, model_type: str) -> Any:
"""
Load model from pickle file trying different compression methods.
@@ -353,10 +345,12 @@ class MLFlowRepository():
code_path = path.join(
artifact_path, "code")
pickle_file = "training_transformer.pkl" if type == "transformer" else "stacking_model.pkl"
pickle_file = TRANSFORMED_COMPRESSED_PATH if model_type == "transform" else PREDICTION_COMPRESSED_PATH
pickle_path = path.join(
artifact_path, "artifacts", pickle_file)
pickle_path = path.join(artifact_path, pickle_file)
self.logger.info(
f"Loading model with type {model_type} from {pickle_path}")
if code_path not in sys_path:
sys_path.insert(0, code_path)
@@ -371,7 +365,8 @@ class MLFlowRepository():
for format_name, open_func in loading_methods:
try:
self.logger.info(f"Trying to load with {format_name}...")
self.logger.inf
(f"Trying to load with {format_name}...")
with open_func(pickle_path) as f:
model = pickle.load(f)
self.logger.info(
@@ -386,7 +381,7 @@ class MLFlowRepository():
f"Could not load model from {pickle_path} - unknown or corrupted format")
def download_model(self, model_name: str, model_type: str, flavor: str,
download_artifacts: bool = False) -> tuple[Any, str]:
load_wrapper: bool = False) -> tuple[Any, str]:
"""
Download model based on type (predict or transform).
@@ -394,33 +389,44 @@ class MLFlowRepository():
model_name (str): Name of the model to download
model_type (str): Type of model ('predict' or 'transform')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
download_artifacts (bool): Whether to download artifacts
load_wrapper (bool): Whether to load wrapper
Returns:
tuple[Any, str]: Model object and artifact path if model is compressed
"""
self.logger.info(
f"Downloading {model_type} model {model_name} with flavor {flavor} and download_artifacts {download_artifacts}")
f"Downloading {model_type} model {model_name} with flavor {flavor} and load_wrapper {load_wrapper}")
if model_type not in ["predict", "transform"]:
raise ValueError(
"Invalid model_type. Use 'predict' or 'transform'.")
if download_artifacts:
artifact_path = None
if load_wrapper:
self.logger.info(
f"Loading wrapper for {model_type} model {model_name} with flavor {flavor}")
target = "prediction_model" if model_type == "predict" else "data_model"
artifact_path = self.dowload_artifacts(
model_name, target)
self.logger.info(
f"Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}")
raw_model = mlflow.pyfunc.load_model(artifact_path)
model = raw_model._model_impl.python_model
else:
artifact_path = None
if model_type == "predict":
model = self.load_predict_model(model_name, flavor, artifact_path)
if model_type == "predict":
model = self.load_predict_model(
model_name, flavor)
elif model_type == "transform":
model = self.load_transform_model(
model_name, flavor, artifact_path)
elif model_type == "transform":
model = self.load_transform_model(
model_name, flavor)
return model, artifact_path
@@ -536,7 +542,7 @@ class MLFlowRepository():
del self.model_cache[model_key]['target']['model']
del self.model_cache[model_key]
def get_model(self, model_name: str, retention: int, model_type: str, flavor: str):
def get_model(self, model_name: str, retention: int, model_type: str, flavor: str) -> Any:
"""
Get model with caching support based on retention policy.
@@ -545,15 +551,15 @@ class MLFlowRepository():
retention (int): Cache retention time in minutes (0 = no cache)
model_type (str): Type of model ('predict' or 'transform')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
retention_target (str): What to cache ('model' or 'artifact')
Returns:
dict: Model configuration with model and artifact paths
Any: Model object
"""
# Retention is 0, download a new model
if retention <= 0:
return self.download_model(
model_name=model_name, model_type=model_type, flavor=flavor, download_artifacts=False)
model, _artifact_path = self.download_model(
model_name=model_name, model_type=model_type, flavor=flavor, load_wrapper=False)
return model
model_key = f'{model_name}_{model_type}'
@@ -573,9 +579,9 @@ class MLFlowRepository():
f"Model {model_name} is not in {model_type} cache, downloading a new one")
# Donwload new model
model = self.download_model(
model, _artifact_path = self.download_model(
model_name=model_name, model_type=model_type, flavor=flavor,
download_artifacts=False)
load_wrapper=False)
cache = {
'target': model,
@@ -604,10 +610,15 @@ class MLFlowRepository():
model_name=model_name, retention=retention,
model_type="transform", flavor=flavor)
return model.predict(data)
prediction = model.predict(data)
def get_cached_predict(self, model_name: str, data: pd.DataFrame, retention: int, flavor: str,
compressed: bool = False, retention_target: str = "model") -> pd.DataFrame | ndarray:
if retention == 0:
del model
return prediction
def get_cached_predict(self, model_name: str, data: pd.DataFrame, retention: int,
flavor: str) -> ndarray:
"""
Get predictions using cached prediction model.
@@ -616,8 +627,6 @@ class MLFlowRepository():
data (pd.DataFrame): Data to make predictions on
retention (int): Cache retention time in minutes
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
compressed (bool): Whether model is compressed
retention_target (str): What to cache ('model' or 'artifact')
Returns:
pd.DataFrame: Model predictions
@@ -626,16 +635,21 @@ class MLFlowRepository():
model_name=model_name, retention=retention,
model_type="predict", flavor=flavor)
return model.predict(data)
prediction = model.predict(data)
if retention == 0:
del model
return prediction
"""
Functions related to model retraining
"""
def create_model_experiment(self, model_name: str, data: pd.DataFrame,
transform_config: dict = {}, predict_config: dict = {},
def create_model_experiment(self, model_name: str, data: pd.DataFrame, latest_production_id: str,
transform_flavor: str = 'sklearn', predict_flavor: str = 'sklearn',
fit_config: dict = {}, target_name: str = None,
is_compressed: bool = False, metadata: dict = {}) -> tuple:
metadata: dict = {}) -> tuple:
"""
Create a new MLFlow experiment for model retraining.
@@ -649,11 +663,10 @@ class MLFlowRepository():
Args:
model_name (str): Name of the MLFlow model to retrain
data (pd.DataFrame): Training data for model retraining
transform_config (dict): Configuration for transformation model
predict_config (dict): Configuration for prediction model
transform_flavor (str): Flavor for transformation model
predict_flavor (str): Flavor for prediction model
fit_config (dict): Fit configuration
target_name (str): Target name
is_compressed (bool): Whether model is compressed
metadata (dict): Metadata for logging
Returns:
@@ -665,49 +678,34 @@ class MLFlowRepository():
self.logger.custom_info(
f"Starting model experiment creation for {model_name}", metadata)
self.logger.custom_debug(
f"Model configuration - transform_config: {transform_config}, predict_config: {predict_config}, fit_config: {fit_config}, target_name: {target_name}", metadata)
f"Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, fit_config: {fit_config}, target_name: {target_name}", metadata)
latest_production_id = self.get_model_run_id(
model_name, stage="Production"
)
data.to_csv(
f"tmp/retrain_data_{model_name}.csv", index=True)
self.logger.custom_info(
f"Retrieved latest production run ID: {latest_production_id}", metadata)
self.logger.custom_info(
f"Loading transformation model for {model_name}", metadata)
transform_flavor = transform_config.get('flavor', 'sklearn')
transformed_is_compressed = transform_config.get(
'is_compressed', False)
load_transform_wrapper = transform_flavor == 'pyfunc'
data_model, data_artifact_path = self.download_model(
model_name=model_name, model_type="transform", flavor=transform_flavor,
download_artifacts=transformed_is_compressed
load_wrapper=load_transform_wrapper
)
self.logger.custom_info(
f"Loading prediction model for {model_name}", metadata)
predict_flavor = predict_config.get('flavor', 'pyfunc')
predicted_is_compressed = predict_config.get('is_compressed', False)
load_predict_wrapper = predict_flavor == 'pyfunc'
prediction_model, prediction_artifact_path = self.download_model(
model_name=model_name, model_type="predict", flavor=predict_flavor,
download_artifacts=predicted_is_compressed
load_wrapper=load_predict_wrapper
)
self.logger.custom_debug(
f"Fitting transformation model with training data (shape: {data.shape})", metadata)
data.to_csv(
f"tmp/retrain_data_{model_name}.csv", index=True)
data_model = data_model.fit(data)
self.logger.custom_debug(
f"Applying transformation to training data", metadata)
treated_data = data_model.predict(data)
treated_data = data_model.fit(data)
# Stores current index as timestamp, Courier model expects a timestamp column
# with specific format
@@ -743,86 +741,105 @@ class MLFlowRepository():
f"Target variable {target_name} not found in treated data, aligning data with treated data indexes", metadata)
# Aligns data with treated data indexes to get target variable
aligned_data = data.loc[treated_data.index]
retrain_dataset = pd.merge(
treated_data, aligned_data, left_index=True, right_index=True)
else:
# Uses target variable from treated data
self.logger.custom_debug(
f"Target variable {target_name} found in treated data, using it", metadata)
aligned_data = treated_data
retrain_dataset = treated_data
if fit_config.get('y_type', 'series').lower() == 'series':
y = aligned_data[target_name]
self.logger.custom_debug(
f"Extracting target as series, shape: {y.shape}", metadata)
else:
y = aligned_data[[target_name]]
self.logger.custom_debug(
f"Extracting target as dataframe, shape: {y.shape}", metadata)
retrain_dataset.to_csv(
f"tmp/retrain_retrain_dataset_{model_name}.csv", index=True)
if not fit_config.get('split_fit_data', False):
# Merge treated data with target
self.logger.custom_debug(
"Merging treated data with target for combined fit", metadata)
treated_data = pd.merge(
treated_data, y, left_index=True, right_index=True)
self.logger.custom_debug(
f"Combined data shape for prediction model fit: {treated_data.shape}", metadata)
treated_data.to_csv(
f"tmp/retrain_treated_data_combined_{model_name}.csv", index=True)
prediction_model = prediction_model.fit(treated_data)
self.logger.custom_debug(
"Prediction model fitted with combined data", metadata)
else:
# Keep data separated
fit_order = fit_config.get('split_fit_first', 'x').lower()
y.to_csv(
f"tmp/retrain_y_{model_name}.csv", index=True)
self.logger.custom_debug(
f"Fitting prediction model with separated data, first arg: {fit_order}", metadata)
if fit_order == 'x':
prediction_model = prediction_model.fit(treated_data, y)
else:
prediction_model = prediction_model.fit(y, treated_data)
experiment = self.get_experiment_by_run_id(latest_production_id)
mlflow.set_experiment(experiment)
prediction_model.fit(retrain_dataset)
self.logger.custom_info(
f"Model experiment creation completed successfully for {model_name}", metadata)
retrain_data = {
'prediction_model': prediction_model,
'prediction_artifact_path': prediction_artifact_path,
'data_model': data_model,
'data_artifact_path': data_artifact_path,
'experiment': experiment
'prediction_model': {
'model': prediction_model,
'artifact_path': prediction_artifact_path
},
'data_model': {
'model': data_model,
'artifact_path': data_artifact_path
}
}
return retrain_data
def log_model(self, model: Any, artifact_local_path: str, flavor: str, model_type: str):
def export_model_to_pkl(self, model_data: dict, model_type: str, metadata: dict = {}):
artifact_local_path = model_data['artifact_path']
model = model_data['model']
if artifact_local_path:
mlflow.log_artifact(artifact_local_path, artifact_path="")
pickle_file = TRANSFORMED_COMPRESSED_PATH if model_type == "data_model" else PREDICTION_COMPRESSED_PATH
pickle_path = path.join(artifact_local_path, pickle_file)
self.logger.custom_info(
f"Saving model to {pickle_path}", metadata)
with open(pickle_path, "wb") as f:
pickle.dump(model, f)
def export_model_to_lzma(self, model_data: dict, model_type: str, metadata: dict = {}):
artifact_local_path = model_data['artifact_path']
model = model_data['model']
if artifact_local_path:
pickle_file = TRANSFORMED_COMPRESSED_PATH if model_type == "data_model" else PREDICTION_COMPRESSED_PATH
pickle_path = path.join(artifact_local_path, pickle_file)
self.logger.custom_info(
f"Saving model to {pickle_path}", metadata)
if path.exists(pickle_path):
remove(pickle_path)
with lzma.open(pickle_path + ".xz", "wb") as f:
self.logger.custom_info(
"Compressing model", metadata)
pickle.dump(model, f)
def log_model(self, model_data: dict, flavor: str, model_type: str,
metadata: dict = {}):
model = model_data['model']
self.logger.custom_debug(
f"Logging {model_type} model to {model_type}", metadata)
if flavor == 'sklearn':
mlflow.sklearn.log_model(model, model_type)
elif flavor == 'pyfunc':
code_path = [path.join(
model_data['artifact_path'], 'code', "utils")]
self.logger.custom_debug(
f"Code path: {code_path}", metadata)
model.store_model(
artifact_path=model_type,
code_path=code_path,
to_disk=False
)
self.logger.custom_debug(
f"Model uploaded successfully", metadata)
elif flavor == 'pytorch':
mlflow.pytorch.log_model(model, model_type)
else:
if flavor == 'sklearn':
mlflow.sklearn.log_model(model, model_type)
elif flavor == 'pyfunc':
mlflow.pyfunc.log_model(model, model_type)
elif flavor == 'pytorch':
mlflow.pytorch.log_model(model, model_type)
else:
raise ValueError(
"Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.")
raise ValueError(
"Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.")
def perform_model_retrain(self,
model_name: str,
data: pd.DataFrame,
retrain_data: dict,
transform_config: dict = {},
predict_config: dict = {},
metadata: dict = {}):
transform_flavor: str = 'sklearn',
predict_flavor: str = 'sklearn',
metadata: dict = {},
latest_production_id: str = None) -> dict:
"""
Execute the complete model retraining process in MLFlow.
@@ -839,7 +856,8 @@ class MLFlowRepository():
experiment (str): MLFlow experiment name for the retraining
model_name (str): Name of the model being retrained
data (pd.DataFrame): Training data used for retraining
is_compressed (bool): Whether model is compressed
transform_flavor (str): Flavor for transformation model
predict_flavor (str): Flavor for prediction model
metadata (dict): Metadata for logging
Returns:
@@ -850,63 +868,90 @@ class MLFlowRepository():
prediction_model = retrain_data['prediction_model']
data_model = retrain_data['data_model']
experiment = retrain_data['experiment']
prediction_artifact_path = retrain_data['prediction_artifact_path']
data_artifact_path = retrain_data['data_artifact_path']
model_temp_path = path.join(
ARTIFACTS_PATH, model_name)
self.logger.custom_info(
f"Starting model retraining process for {model_name} in experiment {experiment}", metadata)
f"Starting model retraining process for {model_name}", metadata)
transform_flavor = transform_config.get('flavor', 'sklearn')
predict_flavor = predict_config.get('flavor', 'sklearn')
pred_model_atributes = vars(prediction_model) # load class attributes
data_model_atributes = vars(data_model) # load class attributes
original_params = self.get_model_params(latest_production_id)
retrain_params = {
**original_params,
"retrain": True,
'retrain_date': datetime.now().isoformat(),
'source_run_id': latest_production_id,
'retrain_samples': str(data.shape),
}
experiment_description = f"Retrain model {model_name} with new data"
current_run_name = self.get_next_run_name(experiment)
attributes = {}
experiment = self.get_experiment(model_name,
create_if_not_exists=True)
experiment_name = experiment.name
for name_atribute, val_atribute in pred_model_atributes.items():
if name_atribute != "model":
attributes[name_atribute] = val_atribute
for name_atribute, val_atribute in data_model_atributes.items():
if name_atribute != "model":
attributes[name_atribute] = val_atribute
current_run_name = self.get_next_run_name(experiment_name)
self.logger.custom_debug(
f"Attributes: {attributes}", metadata)
f"Attributes: {retrain_params}", metadata)
data_path = f"{model_temp_path}/retrain_data.csv"
data.to_csv(data_path, index=True)
self.logger.custom_info(
f"Starting model upload for {experiment_name} with run name {current_run_name}", metadata)
with mlflow.start_run(
run_name=current_run_name, description=experiment_description
experiment_id=experiment.experiment_id,
run_name=current_run_name,
description=experiment_description
) as _run:
run_id = _run.info.run_id
self.logger.custom_info(
f"Logging data model", metadata)
# dynamic parameters, including model itself
self.log_model(data_model, transform_flavor,
"data_model", metadata)
# dynamic parameters, including model itself
self.logger.custom_info(
f"Logging prediction model", metadata)
self.log_model(prediction_model, predict_flavor,
"prediction_model", metadata)
self.logger.custom_info(
f"Model logged successfully for {model_name}", metadata)
self.logger.custom_info(
f"Logging remaining parameters for {model_name}", metadata)
# update transfomation model
# fixed parameters
for name_atribute, val_atribute in attributes.items():
mlflow.log_param(name_atribute, val_atribute)
# dynamic parameters, including model itself
self.log_model(data_model, data_artifact_path,
transform_flavor, "data_model")
makedirs("tmp/retrain_data", exist_ok=True)
file_path = f"tmp/retrain_data/retrain_data_{model_name}.csv"
data.to_csv(file_path, index=True)
mlflow.log_params(retrain_params)
# log the data raw
mlflow.log_artifact(file_path)
mlflow.log_artifact(data_path)
# dynamic parameters, including model itself
self.log_model(prediction_model, prediction_artifact_path,
predict_flavor, "prediction_model")
mlflow.log_param("retrain", True)
self.logger.custom_info(
f"Deleting model from filesystem", metadata)
if path.exists(model_temp_path):
rmtree(model_temp_path)
# clear temp file
if path.exists(file_path):
remove(file_path)
self.logger.custom_info(
f"Deleting prediction model from memory", metadata)
del prediction_model['model']
del prediction_model
return experiment
self.logger.custom_info(
f"Deleting data model from memory", metadata)
del data_model['model']
del data_model
return {
'run_id': run_id,
'experiment_id': experiment.experiment_id,
'experiment_name': experiment.name
}
def update_production_model_by_run_id(self, run_id: str, model_name: str, metadata: dict = {}) -> dict:
"""
@@ -942,12 +987,8 @@ class MLFlowRepository():
mlflow.register_model(
f"runs:/{run_id}/prediction_model", model_name)
# Colocar a versão do modelo em produção
# Depois de registrar o modelo, precisamos pegar a versão mais recente do modelo e movê-lo para o estágio 'Production'
client = mlflow.tracking.MlflowClient()
# Obter a versão mais recente registrada do modelo
model_versions = client.get_registered_model(
model_versions = self.client.get_registered_model(
model_name).latest_versions
if not isinstance(model_versions, list):
@@ -956,7 +997,7 @@ class MLFlowRepository():
max_version = max(model_versions, key=lambda x: int(x.version)).version
# Mover a versão mais recente do modelo para o estágio de 'Production'
client.transition_model_version_stage(
self.client.transition_model_version_stage(
name=model_name,
version=max_version,
stage="Production",
@@ -1190,8 +1231,8 @@ class MLFlowRepository():
target_name = model_config.get('target', None)
transform_config = model_config.get('transform_config', {})
predict_config = model_config.get('predict_config', {})
transform_flavor = model_config.get('transform_flavor', 'sklearn')
predict_flavor = model_config.get('predict_flavor', 'sklearn')
fit_config = {
'split_fit_data': model_config.get('split_fit_data', False),
@@ -1200,19 +1241,25 @@ class MLFlowRepository():
}
self.logger.custom_debug(
f"Model configuration - transform_config: {transform_config}, predict_config: {predict_config}, fit_config: {fit_config}, target_name: {target_name}", metadata)
f"Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, fit_config: {fit_config}, target_name: {target_name}", metadata)
try:
latest_production_id = self.get_model_run_id(
model_name, stage="Production"
)
self.logger.custom_info(
"Creating model experiment environment", metadata)
retrain_data = self.create_model_experiment(
model_name=model_name, data=data, transform_config=transform_config, predict_config=predict_config, fit_config=fit_config, target_name=target_name, metadata=metadata)
model_name=model_name, data=data, transform_flavor=transform_flavor,
predict_flavor=predict_flavor, fit_config=fit_config, target_name=target_name,
metadata=metadata, latest_production_id=latest_production_id)
self.logger.custom_info(
f"Model experiment created successfully: {retrain_data}", metadata)
self.logger.custom_info("Saving model retrain", metadata)
experiment = self.perform_model_retrain(
model_name, data, retrain_data, transform_config, predict_config, metadata)
model_name=model_name, data=data, retrain_data=retrain_data,
transform_flavor=transform_flavor, predict_flavor=predict_flavor, metadata=metadata, latest_production_id=latest_production_id)
self.logger.custom_info(
f"Model retraining completed successfully for experiment: {experiment}", metadata)
@@ -1276,8 +1323,8 @@ class MLFlowRepository():
This operation is irreversible. The previous production model will
be automatically archived when the new version is promoted.
"""
experiment_id = self.get_experiment(experiment)
run_id = self.get_experiment_last_run(experiment_id)
run_id = experiment['run_id']
experiment_id = experiment['experiment_id']
metadata_result = self.update_production_model_by_run_id(
run_id, model_name, metadata)