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

@@ -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)

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(

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)
 

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
     )
 

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(