SIENTIAPDE-1231

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

View File

@@ -207,13 +207,12 @@ class Gates(BaseActivity):
filter_output = []
self.debug(
f"Input data: \n {create_sample_dict(data, max_items=5, max_depth=2)}", metadata)
f"Input data: \n {create_sample_dict(data, max_items=5, max_depth=5)}", metadata)
self.debug(f"Filters: {filters}", metadata)
comments = []
for fil, config in filters.items():
if fil not in mlflow_response_filter_functions:
self.error(f"Filter {fil} not found", metadata)
continue
try:
if mlflow_response_filter_functions[fil](data, config):
@@ -293,7 +292,7 @@ class Gates(BaseActivity):
filter_output = []
self.debug(f"Input data:\n {data.head(5).to_string()}", metadata)
self.debug(f"Filters: \n {create_sample_dict(filters)}", metadata)
self.debug(f"Filters: \n {filters}", metadata)
for fil, config in filters.items():
if fil not in mlflow_content_filter_functions:
@@ -492,6 +491,36 @@ class Gates(BaseActivity):
self.info(f"Default prediction formatted: {data.size} rows", metadata)
return data.to_dict()
@activity.defn(name="format_retrain_report")
async def format_retrain_report(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Format retrain report data according to configured storage policies.
"""
metadata = input_data['metadata']
self.info("Formatting retrain report...", metadata)
experiment_response = input_data['experiment_response']
update_report = input_data['update_report']
model_id = input_data['model_id']
model_name = input_data['model_name']
report = DataFrame({
'model_id': [model_id],
'model_name': [model_name],
'timestamp': [experiment_response['timestamp']],
'status': [experiment_response['message']]
})
if experiment_response['success']:
# Retrain was successfull
report['version'] = update_report['version']
report['mlflow_run_id'] = update_report['mlflow_run_id']
report['mlflow_experiment_id'] = update_report['mlflow_experiment_id']
self.debug(f"Retrain report: {report.to_csv()}", metadata)
return report.to_dict()
@activity.defn(name="get_last_timestamp")
async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
"""