From 9b71ad7556f08a706b91511b2ce7ca03f8652c48 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 2 Oct 2025 16:36:18 -0300 Subject: [PATCH] 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 | 12 +- .../utils/repository/model_repository.py | 323 ++++++++++-------- laborious/worker/worker.py | 3 +- laborious/workflows/minimal_retrain.py | 14 +- values.yaml | 2 +- 5 files changed, 204 insertions(+), 150 deletions(-) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 7ad795b..3e4b492 100644 --- a/laborious/activities/mlflow.py +++ b/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 diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 1816c80..89922a8 100644 --- a/laborious/utils/repository/model_repository.py +++ b/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,16 +515,7 @@ 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'] + 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']) - self.logger.custom_debug( - f"Data received from model prediction: {data.to_csv()}", metadata) + + if isinstance(data, pd.DataFrame): + self.logger.custom_debug( + 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 diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 567a257..81c716b 100644 --- a/laborious/worker/worker.py +++ b/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, diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py index 98c1e16..607f3a5 100644 --- a/laborious/workflows/minimal_retrain.py +++ b/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) ) diff --git a/values.yaml b/values.yaml index c441f65..ef86d89 100644 --- a/values.yaml +++ b/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"