SIENTIAPDE-1231

Update model retraining and logging enhancements

- Changed the GITHUB_BRANCH value in values.yaml to 'main' for consistency.
- Refactored MLFlow class to improve timestamp handling and error messaging during model retraining.
- Enhanced MLFlowRepository methods to include metadata logging and improved model version retrieval.
- Updated minimal_retrain workflow to support extended timeout for activities and include model configuration in input data.

laborious/activities/mlflow.py

@@ -240,13 +240,19 @@ class MLFlow(BaseActivity):
         data.sort_index(inplace=True)
         data.reset_index(inplace=True)
 
-        data = data.dropna()
+        data['timestamp'] = to_datetime(
+            data['timestamp'], format=DATETIME_FORMAT_WITH_TZ).dt.strftime(DATETIME_FORMAT)
+        data['timestamp'] = to_datetime(
+            data['timestamp'], format=DATETIME_FORMAT)
+
+        # data = data.dropna()
         data.columns.name = None
 
         retrain_output = self.model_monitoring_repository.retrain_model(
             data=data,
             model_name=model_name,
-            model_config=model_config
+            model_config=model_config,
+            metadata=metadata
         )
 
         if not retrain_output['success']:
@@ -255,7 +261,7 @@ class MLFlow(BaseActivity):
             self.send_notification(
                 metadata=metadata,
                 notification_id='RETRAIN_MODEL_ERROR',
-                message=f'Error retraining model {model_name}: {retrain_output['message']}',
+                message=f"Error retraining model {model_name}: {retrain_output['message']}",
                 block='retrain_model',
                 level=NotificationLevel.ERROR,
                 attachment_content=trace

laborious/utils/repository/model_repository.py

@@ -23,6 +23,8 @@ from sientia_do.observability.logger import Logger
 import lzma
 import gzip
 import pickle
+from numpy import ndarray
+from typing import Any
 
 from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
@@ -86,18 +88,22 @@ class MLFlowRepository():
                 f"Model '{model_name}' not found in the Model Registry."
             )
 
-        # Get the latest version in the specified stage
+        # Get all versions of the model and filter by stage
         model_versions = self.client.search_model_versions(
-            filter_string=f"name='{model_name}' and stage='{stage}'"
+            filter_string=f"name='{model_name}'"
         )
 
-        if not model_versions:
+        # Filter versions by the desired stage using current_stage attribute
+        stage_versions = [
+            mv for mv in model_versions if mv.current_stage == stage]
+
+        if not stage_versions:
             raise mlflow.exceptions.MlflowException(
                 f"Model '{model_name}' in stage '{stage}' not found in the Model Registry."
             )
 
         # Sort by version number to get the latest
-        latest_version = max(model_versions, key=lambda v: int(v.version))
+        latest_version = max(stage_versions, key=lambda v: int(v.version))
         run_id = latest_version.source.split("/")
         return run_id[2]
 
@@ -236,7 +242,7 @@ class MLFlowRepository():
         )
 
     def load_predict_model(self, model_name: str, flavor: str = 'pyfunc',
-                           artifact_path: str | None = None):
+                           artifact_path: str | None = None) -> Any:
         """
         Downloads a predictive model from the MLflow Model Registry.
 
@@ -257,12 +263,14 @@ class MLFlowRepository():
 
         if artifact_path:
             self.logger.info(
-                f"Prediction model {model_name} is not compressed, loading from {artifact_path}")
+                f"Prediction model {model_name} is compressed, loading from {artifact_path}")
 
             model = self.load_model_with_compression(
                 artifact_path, "prediction")
 
         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':
@@ -276,7 +284,7 @@ class MLFlowRepository():
         return model
 
     def load_transform_model(self, model_name: str, flavor: str,
-                             artifact_path: str | None = None):
+                             artifact_path: str | None = None) -> Any:
         """
         Downloads the latest production version of a specified transformation model.
 
@@ -307,11 +315,13 @@ class MLFlowRepository():
         if artifact_path:
             # Download model artifacts
             self.logger.info(
-                f"Data model {model_name} is not compressed, loading from {artifact_path}")
+                f"Data model {model_name} is compressed, loading from {artifact_path}")
 
             model = self.load_model_with_compression(
                 artifact_path, "transformer")
         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':
@@ -323,7 +333,7 @@ class MLFlowRepository():
                     "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.")
         return model
 
-    def load_model_with_compression(self, artifact_path: str, type: str):
+    def load_model_with_compression(self, artifact_path: str, type: str) -> Any:
         """
         Load model from pickle file trying different compression methods.
 
@@ -378,7 +388,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,
-                       compressed: bool = False) -> dict:
+                       download_artifacts: bool = False) -> Any:
         """
         Download model based on type (predict or transform).
 
@@ -386,17 +396,20 @@ 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')
-            compressed (bool): Whether model is compressed
+            download_artifacts (bool): Whether to download artifacts
 
         Returns:
-            dict: Model configuration with model and artifact paths
+            Any: Model object
         """
 
+        self.logger.info(
+            f"Downloading {model_type} model {model_name} with flavor {flavor} and download_artifacts {download_artifacts}")
+
         if model_type not in ["predict", "transform"]:
             raise ValueError(
                 "Invalid model_type. Use 'predict' or 'transform'.")
 
-        if compressed:
+        if download_artifacts:
             target = "prediction_model" if model_type == "predict" else "data_model"
 
             artifact_path = self.dowload_artifacts(
@@ -411,10 +424,7 @@ class MLFlowRepository():
             model = self.load_transform_model(
                 model_name, flavor, artifact_path)
 
-        return {
-            "model": model,
-            "artifact_path": artifact_path
-        }
+        return model
 
     """
     Functions related to data format
@@ -469,22 +479,6 @@ class MLFlowRepository():
     Functions related to cache management of models
     """
 
-    def check_cache_config(self, cache: dict, new_config: dict) -> bool:
-        """
-        Check if cache configuration matches new configuration.
-
-        Args:
-            cache (dict): Cached model configuration
-            new_config (dict): New configuration to compare
-
-        Returns:
-            bool: True if configurations match, False otherwise
-        """
-        old_config = cache['config']
-        if old_config != new_config:
-            return False
-        return True
-
     def check_cache_retention(self, cache: dict, retention: int) -> bool:
         """
         Check if cache is still valid based on retention time.
@@ -503,7 +497,6 @@ class MLFlowRepository():
         return True
 
     def handle_valid_model(self, model_name: str, model_type: str,
-                           compressed: bool, retention_target: str,
                            cache: dict) -> dict:
         """
         Handle valid cached model by returning appropriate model configuration.
@@ -522,15 +515,6 @@ class MLFlowRepository():
             self.logger.debug(
                 f"Model {model_name} is still valid, using cached version")
 
-        # If model is compressed and retention target is artifact, load the model from pkl
-        if compressed and retention_target == "artifact":
-            model = self.load_model_with_compression(
-                cache['target']['artifact_path'], model_type)
-            return {
-                'model': model,
-                'artifact_path': cache['target']['artifact_path']
-            }
-        else:
         return cache['target']
 
     def handle_outdated_model(self, model_name: str, model_key: str) -> None:
@@ -549,13 +533,9 @@ class MLFlowRepository():
                 f"Model {model_name} is outdated, downloading a new one")
 
         del self.model_cache[model_key]['target']['model']
-        if path.exists(self.model_cache[model_key]['target']['artifact_path']):
-            remove(self.model_cache[model_key]
-                   ['target']['artifact_path'])
         del self.model_cache[model_key]
 
-    def get_model(self, model_name: str, retention: int, model_type: str, flavor: str,
-                  compressed: bool = False, retention_target: str = "model"):
+    def get_model(self, model_name: str, retention: int, model_type: str, flavor: str):
         """
         Get model with caching support based on retention policy.
 
@@ -564,7 +544,6 @@ 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')
-            compressed (bool): Whether model is compressed
             retention_target (str): What to cache ('model' or 'artifact')
 
         Returns:
@@ -572,58 +551,42 @@ class MLFlowRepository():
         """
         # Retention is 0, download a new model
         if retention <= 0:
-            model_config = self.download_model(
-                model_name, model_type, flavor, compressed)
-            return model_config
+            return self.download_model(
+                model_name=model_name, model_type=model_type, flavor=flavor, download_artifacts=False)
 
         model_key = f'{model_name}_{model_type}'
-        config = {
-            'compressed': compressed,
-            'retention_target': retention_target
-        }
 
         if model_key in self.model_cache:
             cache = self.model_cache[model_key]
 
             # Check if config has changed or is outdated
-            if self.check_cache_config(cache, config) or self.check_cache_retention(cache, retention):
+            if self.check_cache_retention(cache, retention):
                 return self.handle_valid_model(
-                    model_name, model_type, compressed,
-                    retention_target, cache)
+                    model_name=model_name, model_type=model_type, cache=cache)
             else:
                 # Model is outdated, delete old model files
                 self.handle_outdated_model(model_name, model_key)
         else:
             if self.logger:
                 self.logger.debug(
-                    f"Model {model_name} is not in the cache, downloading a new one")
+                    f"Model {model_name} is not in {model_type} cache, downloading a new one")
 
         # Donwload new model
-        model_config = self.download_model(model_name, model_type, flavor,
-                                           compressed)
-        # If model is compressed and retention target is artifact,
-        # dont save the model in the cache
-        if compressed and retention_target == "artifact":
-            model_config_to_cache = {
-                'artifact_path': model_config['artifact_path'],
-                'model': None
-            }
-        else:
-            model_config_to_cache = {
-                **model_config
-            }
+        model = self.download_model(
+            model_name=model_name, model_type=model_type, flavor=flavor,
+            download_artifacts=False)
+
         cache = {
-            'target': model_config_to_cache,
-            'config': config,
+            'target': model,
             'timestamp': datetime.now()
         }
 
         self.model_cache[model_key] = cache
 
-        return model_config
+        return model
 
-    def get_cached_transform(self, model_name: str, data: pd.DataFrame, retention: int, flavor: str,
-                             compressed: bool = False, retention_target: str = "model", keyword: str = "predict") -> pd.DataFrame:
+    def get_cached_transform(self, model_name: str, data: pd.DataFrame,
+                             retention: int, flavor: str) -> pd.DataFrame:
         """
         Get transformed data using cached transform model.
 
@@ -632,35 +595,18 @@ class MLFlowRepository():
             data (pd.DataFrame): Data to transform
             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')
-            keyword (str): Method name to call on model (default: 'predict')
 
         Returns:
             pd.DataFrame: Transformed data
         """
-        model_config = self.get_model(model_name, retention,
-                                      "transform", flavor, compressed, retention_target)
-        model = model_config['model']
+        model = self.get_model(
+            model_name=model_name, retention=retention,
+            model_type="transform", flavor=flavor)
 
-        # Use getattr to dynamically call the method specified by keyword
-        method = getattr(model, keyword)
-        transformed_data = method(data)
-
-        # If model is compressed and retention target is artifact,
-        # delete the model after the prediction
-        if compressed and retention_target == "artifact":
-            del model
-
-        # If retention is 0, delete the artifacts after the prediction
-        if retention == 0 and model_config['artifact_path'] is not None:
-            if path.exists(model_config['artifact_path']):
-                remove(model_config['artifact_path'])
-
-        return transformed_data
+        return 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:
+                           compressed: bool = False, retention_target: str = "model") -> pd.DataFrame | ndarray:
         """
         Get predictions using cached prediction model.
 
@@ -675,9 +621,10 @@ class MLFlowRepository():
         Returns:
             pd.DataFrame: Model predictions
         """
-        model_config = self.get_model(model_name, retention, "predict",
-                                      flavor, compressed, retention_target)
-        model = model_config['model']
+        model = self.get_model(
+            model_name=model_name, retention=retention,
+            model_type="predict", flavor=flavor)
+
         return model.predict(data)
 
     """
@@ -686,7 +633,8 @@ class MLFlowRepository():
 
     def create_model_experiment(self, model_name: str, data: pd.DataFrame,
                                 transform_flavor: str = 'sklearn', predict_flavor: str = 'pyfunc',
-                                compressed: bool = False, fit_config: dict = {}, target_name: str = None) -> tuple:
+                                fit_config: dict = {}, target_name: str = None,
+                                metadata: dict = {}) -> tuple:
         """
         Create a new MLFlow experiment for model retraining.
 
@@ -702,7 +650,9 @@ class MLFlowRepository():
             data (pd.DataFrame): Training data for model retraining
             transform_flavor (str): Flavor for transformation model
             predict_flavor (str): Flavor for prediction model
-            compressed (bool): Whether models are compressed
+            fit_config (dict): Fit configuration
+            target_name (str): Target name
+            metadata (dict): Metadata for logging
 
         Returns:
             tuple: (prediction_model, data_model, experiment)
@@ -710,42 +660,94 @@ class MLFlowRepository():
                 - data_model: Fitted transformation model
                 - experiment: MLFlow experiment name
         """
+        self.logger.custom_info(
+            f"Starting model experiment creation for {model_name}", metadata)
+        self.logger.custom_debug(
+            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"
         )
+
+        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)
+
         data_model = self.download_model(
-            model_name, "transform", transform_flavor, compressed
-        )['model']
+            model_name=model_name, model_type="transform", flavor=transform_flavor,
+            download_artifacts=(transform_flavor == 'pyfunc')
+        )
+
+        self.logger.custom_info(
+            f"Loading prediction model for {model_name}", metadata)
+
         prediction_model = self.download_model(
-            model_name, "predict", predict_flavor, compressed
-        )['model']
+            model_name=model_name, model_type="predict", flavor=predict_flavor,
+            download_artifacts=(predict_flavor == 'pyfunc')
+        )
+
+        self.logger.custom_debug(
+            f"Fitting transformation model with training data (shape: {data.shape})", metadata)
+
         data_model = data_model.fit(data)
+
+        self.logger.custom_debug(
+            f"Applying transformation to training data", metadata)
+
         treated_data = data_model.predict(data)
 
+        self.logger.custom_debug(
+            f"Transformed data shape: {treated_data.shape}", metadata)
+
         if target_name is None:
             target_name = data_model.target_variable
+            self.logger.custom_debug(
+                f"Using target variable from data model: {target_name}", metadata)
+        else:
+            self.logger.custom_debug(
+                f"Using provided target variable: {target_name}", metadata)
 
         if fit_config.get('y_type', 'series').lower() == 'series':
-            y = treated_data[target_name]
+            y = data[target_name]
+            self.logger.custom_debug(
+                f"Extracting target as series, shape: {y.shape}", metadata)
         else:
-            y = treated_data[[target_name]]
+            y = data[[target_name]]
+            self.logger.custom_debug(
+                f"Extracting target as dataframe, shape: {y.shape}", metadata)
 
         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)
 
             prediction_model = prediction_model.fit(treated_data)
+            self.logger.custom_debug(
+                "Prediction model fitted with combined data", metadata)
         else:
             # Keep data separated
-            if fit_config.get('split_fit_first', 'x').lower() == 'x':
+            fit_order = fit_config.get('split_fit_first', 'x').lower()
+            self.logger.custom_debug(
+                f"Fitting prediction model with separated data, order: {fit_order}", metadata)
+            if fit_order == 'x':
                 prediction_model = prediction_model.fit(treated_data, y)
+                self.logger.custom_debug(
+                    "Prediction model fitted with X, y order", metadata)
             else:
                 prediction_model = prediction_model.fit(y, treated_data)
+                self.logger.custom_debug(
+                    "Prediction model fitted with y, X order", metadata)
 
         experiment = self.get_experiment_by_run_id(latest_production_id)
         mlflow.set_experiment(experiment)
 
+        self.logger.custom_info(
+            f"Model experiment creation completed successfully for {model_name}", metadata)
         return prediction_model, data_model, experiment
 
     def perform_model_retrain(self,
@@ -753,7 +755,8 @@ class MLFlowRepository():
                               data_model,
                               experiment: str,
                               model_name: str,
-                              data: pd.DataFrame):
+                              data: pd.DataFrame,
+                              metadata: dict = {}):
         """
         Execute the complete model retraining process in MLFlow.
 
@@ -770,12 +773,16 @@ 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
+            metadata (dict): Metadata for logging
 
         Returns:
             tuple: (status_message, experiment_name)
                 - status_message (str): Success confirmation message
                 - experiment_name (str): Name of the experiment
         """
+        self.logger.custom_info(
+            f"Starting model retraining process for {model_name} in experiment {experiment}", metadata)
+
         pred_model_atributes = vars(prediction_model)  # load class attributes
         data_model_atributes = vars(data_model)  # load class attributes
         experiment_description = f"Retrain model {model_name} with new data"
@@ -813,7 +820,7 @@ class MLFlowRepository():
 
             return experiment
 
-    def update_production_model_by_run_id(self, run_id: str, model_name: str) -> dict:
+    def update_production_model_by_run_id(self, run_id: str, model_name: str, metadata: dict = {}) -> dict:
         """
         Update production model with a specific MLFlow run.
 
@@ -824,6 +831,7 @@ class MLFlowRepository():
         Args:
             run_id (str): MLFlow run ID containing the model to promote
             model_name (str): Name of the MLFlow model
+            metadata (dict): Metadata for logging
 
         Returns:
             dict: Model update metadata containing:
@@ -837,6 +845,9 @@ class MLFlowRepository():
             3. Transitions the model to 'Production' stage
             4. Archives existing production versions
         """
+        self.logger.custom_info(
+            f"Starting production model update for {model_name} with run ID: {run_id}", metadata)
+
         # Registrar o modelo
         # Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro.
         # Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso.
@@ -909,23 +920,24 @@ class MLFlowRepository():
                       and returned in the response structure rather than propagated.
         """
         self.logger.custom_debug(
-            f"Data received for model transformation: {data.to_csv()}", metadata)
+            f"Data received for model transformation: {data.head(5).to_csv()}", metadata)
+
+        data.to_csv(
+            f"tmp/data_{model_name}.csv", index=True)
 
         model_retention = model_config.get('retention_minutes', 0)
         flavor = model_config.get('transform_flavor', 'sklearn')
-        compressed = model_config.get('is_compressed', False)
-        retention_target = model_config.get('retention_target', 'model')
-        transform_keyword = model_config.get(
-            'transform_function_keyword', 'predict')
 
         try:
             transformed_data = self.get_cached_transform(
-                model_name, data, model_retention, flavor,
-                compressed, retention_target, transform_keyword
+                model_name, data, model_retention, flavor
             )
 
             self.logger.custom_debug(
-                f"Data received from model transformation: {transformed_data.to_csv()}", metadata)
+                f"Data received from model transformation: {transformed_data.head(5).to_csv()}", metadata)
+
+            transformed_data.to_csv(
+                f"tmp/transformed_data_{model_name}.csv", index=True)
 
             transformed_data = self.detect_and_parse_datetime_index(
                 transformed_data, metadata)
@@ -995,14 +1007,29 @@ class MLFlowRepository():
             start_time = datetime.now()
 
             self.logger.custom_debug(
-                f"Data received for model prediction: {data.to_csv()}", metadata)
+                f"Data received for model prediction: {data.head(5).to_csv()}", metadata)
+
+            data.to_csv(
+                f"tmp/treated_data_{model_name}.csv", index=True)
+
             data = self.get_cached_predict(
                 model_name, data, model_retention, flavor)
 
             end_time = datetime.now()
-            data = pd.DataFrame(data, columns=['prediction'])
+
+            if isinstance(data, pd.DataFrame):
                 self.logger.custom_debug(
-                f"Data received from model prediction: {data.to_csv()}", metadata)
+                    f"Data received from model prediction: {data.head(5).to_csv()}", metadata)
+
+                data.to_csv(
+                    f"tmp/predicted_data_{model_name}.csv", index=True)
+                data.columns = ['prediction']
+
+            else:
+                data = pd.DataFrame(data, columns=['prediction'])
+                data.to_csv(
+                    f"tmp/predicted_data_{model_name}.csv", index=True)
+
             data.index = input_index
             data['response_time'] = (end_time - start_time).total_seconds()
 
@@ -1053,6 +1080,7 @@ class MLFlowRepository():
             model_name (str): Name of the MLFlow model to retrain. Must exist
                             in the MLFlow Model Registry in Production stage.
             model_config (dict): Model configuration parameters
+            metadata (dict): Metadata for logging
 
         Returns:
             tuple: Retraining operation results containing:
@@ -1065,13 +1093,15 @@ class MLFlowRepository():
             Exception: Any other exception during the retraining process
         """
 
-        self.logger.debug(
+        self.logger.custom_info(
+            f"Starting model retraining workflow for {model_name}", metadata)
+        self.logger.custom_debug(
             f"Data received for model retraining: {data.to_csv()}", metadata)
 
+        target_name = model_config.get('target', None)
+
         transform_flavor = model_config.get('transform_flavor', 'sklearn')
         predict_flavor = model_config.get('predict_flavor', 'pyfunc')
-        compressed = model_config.get('is_compressed', False)
-        target_name = model_config.get('target', None)
 
         fit_config = {
             'split_fit_data': model_config.get('split_fit_data', False),
@@ -1079,13 +1109,24 @@ class MLFlowRepository():
             'y_type': model_config.get('y_type', 'series').lower()
         }
 
-        try:
+        self.logger.custom_debug(
+            f"Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, fit_config: {fit_config}, target_name: {target_name}", metadata)
 
+        try:
+            self.logger.custom_info(
+                "Creating model experiment environment", metadata)
             prediction_model, data_model, experiment = self.create_model_experiment(
-                model_name, data, transform_flavor, predict_flavor, compressed,
-                fit_config, target_name)
+                model_name=model_name, data=data, transform_flavor=transform_flavor,
+                predict_flavor=predict_flavor, fit_config=fit_config, target_name=target_name,
+                metadata=metadata)
+            self.logger.custom_info(
+                f"Model experiment created successfully: {experiment}", metadata)
+
+            self.logger.custom_info("Performing model retraining", metadata)
             experiment = self.perform_model_retrain(
-                prediction_model, data_model, experiment, model_name, data)
+                prediction_model, data_model, experiment, model_name, data, metadata)
+            self.logger.custom_info(
+                f"Model retraining completed successfully for experiment: {experiment}", metadata)
 
             return {
                 'success': True,
@@ -1093,14 +1134,16 @@ class MLFlowRepository():
                 'message': 'Model retrained successfully.'
             }
         except Exception as e:
+            error_msg = f'Error retraining model {model_name}: {e}'
+            self.logger.custom_info(error_msg, metadata)
             return {
                 'success': False,
                 'experiment': None,
-                'message': f'Error retraining model {model_name}: {e}',
+                'message': error_msg,
                 'traceback': traceback.format_exc()
             }
 
-    def update_production_model(self, experiment: str, model_name: str) -> dict:
+    def update_production_model(self, experiment: str, model_name: str, metadata: dict = {}) -> dict:
         """
         Update production model using the latest retraining run.
 
@@ -1126,6 +1169,7 @@ class MLFlowRepository():
                            Must be a valid experiment that exists in MLFlow.
             model_name (str): Name of the MLFlow model to update. Must exist
                             in the MLFlow Model Registry.
+            metadata (dict): Metadata for logging
 
         Returns:
             dict: Complete model update metadata containing:
@@ -1146,8 +1190,9 @@ class MLFlowRepository():
         """
         experiment_id = self.get_experiment(experiment)
         run_id = self.get_experiment_last_run(experiment_id)
-        metadata = self.update_production_model_by_run_id(run_id, model_name)
+        metadata_result = self.update_production_model_by_run_id(
+            run_id, model_name, metadata)
 
-        metadata['mlflow_experiment_id'] = experiment_id
+        metadata_result['mlflow_experiment_id'] = experiment_id
 
-        return metadata
+        return metadata_result

laborious/worker/worker.py

@@ -142,7 +142,8 @@ async def main():
                 activities.load_custom_query,
                 activities.retrain_model,
                 activities.update_production_model,
-                activities.export_data_to_postgres
+                activities.format_retrain_report,
+                activities.export_data_to_postgres,
             ],
             max_concurrent_workflow_tasks=50,
             max_concurrent_activities=50,

laborious/workflows/minimal_retrain.py

@@ -68,6 +68,7 @@ class MinimalRetrain():
         }
 
         model_name = input_data['model_name']
+        model_config = input_data.get('model_config', {})
 
         data = await workflow.execute_local_activity_method(
             Activities.load_custom_query,
@@ -77,7 +78,7 @@ class MinimalRetrain():
                 'datetime_columns': input_data.get('datetime_columns', [])
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=600)
         )
 
         experiment_response = await workflow.execute_activity_method(
@@ -85,10 +86,11 @@ class MinimalRetrain():
             {
                 **metadata,
                 'data': data,
-                'model_name': model_name
+                'model_name': model_name,
+                'model_config': model_config
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=600)
         )
 
         if experiment_response['success']:
@@ -101,7 +103,7 @@ class MinimalRetrain():
                     **experiment_response
                 },
                 retry_policy=retry_policy,
-                start_to_close_timeout=timedelta(seconds=60)
+                start_to_close_timeout=timedelta(seconds=600)
             )
         else:
             update_report = {}
@@ -116,7 +118,7 @@ class MinimalRetrain():
                 'update_report': update_report
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=600)
         )
 
         await workflow.execute_activity_method(
@@ -128,5 +130,5 @@ class MinimalRetrain():
                 'table_name': input_data['table_name']
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=600)
         )

values.yaml

@@ -151,7 +151,7 @@ env:
   - name: GITHUB_REPO_URL
     value: "git@github.com:Aignosi/sientia-dataops-laborious_temporal.git"
   - name: GITHUB_BRANCH
-    value: SIENTIAPDE-1231-ajustar-o-retreino-do-courier-no-laborious
+    value: main
   - name: PYTHON_APP
     value: "laborious.worker.worker"