From 49b6e504aeb4266692ae3870ba281e7a7867178d Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 13 Oct 2025 10:36:11 -0300 Subject: [PATCH] SIENTIAPDE-1231 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. --- laborious/activities/gates.py | 18 + laborious/activities/mlflow.py | 1 - .../utils/repository/model_repository.py | 493 ++++++++++-------- laborious/worker/worker.py | 47 -- laborious/workflows/minimal_retrain.py | 6 +- 5 files changed, 291 insertions(+), 274 deletions(-) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 832cfbb..7e6f785 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -17,6 +17,8 @@ with workflow.unsafe.imports_passed_through(): ) from pandas import DataFrame from laborious import metrics + from os import path + from shutil import rmtree # Input filter function mappings input_filter_functions = { @@ -614,3 +616,19 @@ class Gates(BaseActivity): self.info( f"Metrics written for model {metadata['model_name']}", metadata) + + @activity.defn(name="clean_tmp_files") + async def clean_tmp_files(self, input_data: dict[str, Any]): + """ + Clean temporary files in the tmp directory. + """ + model_name = input_data['model_name'] + metadata = input_data['metadata'] + self.info(f"Cleaning tmp files for model {model_name}...", metadata) + + if path.exists(f"tmp/retrain_data/{model_name}"): + rmtree(f"tmp/retrain_data/{model_name}") + if path.exists(f"tmp/artifacts/{model_name}"): + rmtree(f"tmp/artifacts/{model_name}") + + self.info("Tmp files cleaned", metadata) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index a09c297..822cd69 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -304,7 +304,6 @@ class MLFlow(BaseActivity): data['timestamp'] = to_datetime( data['timestamp'], format=DATETIME_FORMAT) - # data = data.dropna() data.columns.name = None retrain_output = self.model_monitoring_repository.retrain_model( diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index b194a69..880602c 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -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) diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index bb74fb3..4658524 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -28,8 +28,6 @@ Environment Variables: from temporalio import workflow, client from temporalio.worker import Worker, PollerBehaviorAutoscaling from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig -from temporalio.converter import PayloadCodec, DataConverter -from temporalio.api.common.v1 import Payload with workflow.unsafe.imports_passed_through(): import os @@ -58,48 +56,6 @@ with workflow.unsafe.imports_passed_through(): POD_ID = os.getenv('POD_ID') SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091")) -LZMA_MIN_MB = float(os.getenv('LZMA_MIN_MB', "1.5")) - - -class LzmaPayloadCodec(PayloadCodec): - async def encode(self, payloads): - out = [] - for p in payloads: - if p.data: - old_len = len(p.data) / 1000000 - # Only compress payloads larger than 1.5 MB - if old_len > LZMA_MIN_MB: - compressed_data = lzma.compress(p.data) - new_len = len(compressed_data) / 1000000 - ratio = new_len / old_len if old_len else 0 - print( - f"[codec] encode lzma: {old_len} MB -> {new_len} MB ({ratio:.2f}x)") - meta = dict(p.metadata or {}) - meta[b"codec"] = b"lzma" - out.append(Payload(metadata=meta, data=compressed_data)) - else: - out.append(p) - - else: - out.append(p) - return out - - async def decode(self, payloads): - out = [] - for p in payloads: - if p.data and (p.metadata or {}).get(b"codec") == b"lzma": - # comp_len = len(p.data) - decomp = lzma.decompress(p.data) - # decomp_len = len(decomp) - # ratio = (decomp_len / comp_len) if comp_len else 0 - # print( - # f"[codec] decode lzma: {comp_len} B -> {decomp_len} B ({ratio:.2f}x)") - meta = dict(p.metadata or {}) - meta.pop(b"codec", None) - out.append(Payload(metadata=meta, data=decomp)) - else: - out.append(p) - return out async def main(): @@ -174,12 +130,9 @@ async def main(): logger.custom_info(f'Starting Temporal Client at {host}...', metadata) - codec_dc = DataConverter(payload_codec=LzmaPayloadCodec()) - temporal_client = await client.Client.connect( target_host=host, namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'), - data_converter=codec_dc, runtime=new_runtime ) diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py index ebdb0b2..71f72ee 100644 --- a/laborious/workflows/minimal_retrain.py +++ b/laborious/workflows/minimal_retrain.py @@ -95,7 +95,7 @@ class MinimalRetrain(): 'model_config': model_config }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=600) + start_to_close_timeout=timedelta(hours=1) ) if experiment_response['success']: @@ -108,7 +108,7 @@ class MinimalRetrain(): **experiment_response }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=600) + start_to_close_timeout=timedelta(seconds=60) ) else: update_report = {} @@ -123,7 +123,7 @@ class MinimalRetrain(): 'update_report': update_report }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=600) + start_to_close_timeout=timedelta(seconds=60) ) await workflow.execute_activity_method(