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

SIENTIAPDE-1231

Browse Source 
Update model retraining and reporting functionality

- Changed the GITHUB_BRANCH value in values.yaml to reflect the latest adjustments for retraining the courier.
- Enhanced the Gates class with a new method `format_retrain_report` to format retraining report data according to storage policies.
- Refactored the MLFlow class to improve error handling during model retraining and return structured output.
- Updated the model_repository to utilize the latest MLFlow API for retrieving model versions and improved logging.
- Modified the minimal_retrain workflow to conditionally update the production model based on retraining success.
This commit is contained in:
vitor-aignosi
2025-09-29 17:34:47 -03:00
parent 1aede51dc1
commit c6f004d20d
7 changed files with 197 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
laborious/utils/filters/conditional_filters.py
View File

@@ -20,6 +20,10 @@ def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool
False if none of the specified variables contain null values.
"""
if data.empty:
return False
return not data[
data['variable'].isin(config['variables']) & data['value'].isna()].empty

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

@@ -76,16 +76,30 @@ class MLFlowRepository():
Returns:
str: The run_id of the model.
"""
latest_versions = self.client.get_latest_versions(
name=model_name, stages=[stage]
# Use search_registered_models instead of deprecated get_latest_versions
registered_models = self.client.search_registered_models(
filter_string=f"name='{model_name}'"
)
if not latest_versions:
if not registered_models:
raise mlflow.exceptions.MlflowException(
f"Model '{model_name}' not found in the Model Registry."
)
# Get the latest version in the specified stage
model_versions = self.client.search_model_versions(
filter_string=f"name='{model_name}' and stage='{stage}'"
)
if not model_versions:
raise mlflow.exceptions.MlflowException(
f"Model '{model_name}' in stage '{stage}' not found in the Model Registry."
)
else:
run_id = latest_versions[0].source.split("/")
return run_id[2]
# Sort by version number to get the latest
latest_version = max(model_versions, key=lambda v: int(v.version))
run_id = latest_version.source.split("/")
return run_id[2]
def get_experiment_by_run_id(self, run_id: str) -> str:
"""
@@ -482,7 +496,7 @@ class MLFlowRepository():
Returns:
bool: True if cache is still valid, False if expired
"""
current_time = self.now()
current_time = datetime.now()
cache_time = cache['timestamp']
if current_time - cache_time >= timedelta(minutes=retention):
return False
@@ -601,7 +615,7 @@ class MLFlowRepository():
cache = {
'target': model_config_to_cache,
'config': config,
'timestamp': self.now()
'timestamp': datetime.now()
}
self.model_cache[model_key] = cache
@@ -672,7 +686,7 @@ class MLFlowRepository():
def create_model_experiment(self, model_name: str, data: pd.DataFrame,
transform_flavor: str = 'sklearn', predict_flavor: str = 'pyfunc',
compressed: bool = False) -> tuple:
compressed: bool = False, fit_config: dict = {}, target_name: str = None) -> tuple:
"""
Create a new MLFlow experiment for model retraining.
@@ -701,18 +715,34 @@ class MLFlowRepository():
)
data_model = self.download_model(
model_name, "transform", transform_flavor, compressed
)
)['model']
prediction_model = self.download_model(
model_name, "predict", predict_flavor, compressed
)
)['model']
data_model = data_model.fit(data)
treated_data = data_model.predict(data)
target_name = data_model.target_variable
y = data[target_name]
treated_data = pd.merge(
treated_data, y, left_index=True, right_index=True)
prediction_model = prediction_model.fit(treated_data)
if target_name is None:
target_name = data_model.target_variable
if fit_config.get('y_type', 'series').lower() == 'series':
y = treated_data[target_name]
else:
y = treated_data[[target_name]]
if not fit_config.get('split_fit_data', False):
# Merge treated data with target
treated_data = pd.merge(
treated_data, y, left_index=True, right_index=True)
prediction_model = prediction_model.fit(treated_data)
else:
# Keep data separated
if fit_config.get('split_fit_first', 'x').lower() == '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)
@@ -781,7 +811,7 @@ class MLFlowRepository():
if path.exists(file_path):
remove(file_path)
return "Model retrained successfully", experiment
return experiment
def update_production_model_by_run_id(self, run_id: str, model_name: str) -> dict:
"""
@@ -844,7 +874,7 @@ class MLFlowRepository():
Functions that provide the interface to model operations
"""
def transform(self, model_name: str, data: pd.DataFrame, model_retention: int,
def transform(self, model_name: str, data: pd.DataFrame,
model_config: dict, metadata: dict):
"""
Transform data using a cached transformation model.
@@ -863,7 +893,6 @@ class MLFlowRepository():
Parameters:
model_name (str): The name of the MLFlow model to use for transformation.
data (pd.DataFrame): The input data to be transformed by the model.
model_retention (int): Cache retention time in minutes (0 = no caching).
model_config (dict): Model configuration parameters
metadata (dict): Metadata for logging
@@ -915,7 +944,7 @@ class MLFlowRepository():
}
}
def predict(self, model_name: str, data: pd.DataFrame, model_retention: int,
def predict(self, model_name: str, data: pd.DataFrame,
model_config: dict, metadata: dict):
"""
Generate predictions using a cached prediction model.
@@ -992,7 +1021,7 @@ class MLFlowRepository():
}
def retrain_model(self, data: pd.DataFrame, model_name: str,
model_config: dict) -> tuple:
model_config: dict, metadata: dict) -> tuple:
"""
Orchestrate the complete model retraining workflow.
@@ -1035,15 +1064,41 @@ class MLFlowRepository():
ValueError: If experiment cannot be created or models cannot be loaded
Exception: Any other exception during the retraining process
"""
self.logger.debug(
f"Data received for model retraining: {data.to_csv()}", metadata)
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)
prediction_model, data_model, experiment = self.create_model_experiment(
model_name, data, transform_flavor, predict_flavor, compressed)
retrain_result = self.perform_model_retrain(
prediction_model, data_model, experiment, model_name, data)
return retrain_result
fit_config = {
'split_fit_data': model_config.get('split_fit_data', False),
'split_fit_first': model_config.get('split_fit_first', 'x').lower(),
'y_type': model_config.get('y_type', 'series').lower()
}
try:
prediction_model, data_model, experiment = self.create_model_experiment(
model_name, data, transform_flavor, predict_flavor, compressed,
fit_config, target_name)
experiment = self.perform_model_retrain(
prediction_model, data_model, experiment, model_name, data)
return {
'success': True,
'experiment': experiment,
'message': 'Model retrained successfully.'
}
except Exception as e:
return {
'success': False,
'experiment': None,
'message': f'Error retraining model {model_name}: {e}',
'traceback': traceback.format_exc()
}
def update_production_model(self, experiment: str, model_name: str) -> dict:
"""