gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions 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

35
laborious/activities/gates.py
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:
"""

41
laborious/activities/mlflow.py
View File

@@ -243,30 +243,29 @@ class MLFlow(BaseActivity):
data = data.dropna()
data.columns.name = None
try:
retrain_output, experiment = self.model_monitoring_repository.retrain_model(
data=data,
model_name=model_name,
model_config=model_config
)
retrain_output = self.model_monitoring_repository.retrain_model(
data=data,
model_name=model_name,
model_config=model_config
)
return {
'status': retrain_output,
'timestamp': timestamp,
'experiment': experiment
}
except Exception as e:
trace = traceback.format_exc()
if not retrain_output['success']:
trace = retrain_output['traceback']
self.send_notification(
metadata=metadata,
notification_id='RETRAIN_MODEL_ERROR',
message=f'Error retraining model {model_name}: {e}',
message=f'Error retraining model {model_name}: {retrain_output['message']}',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.error(trace, metadata=metadata)
raise e
return {
**retrain_output,
'timestamp': timestamp
}
@activity.defn(name="update_production_model")
async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
@@ -307,11 +306,7 @@ class MLFlow(BaseActivity):
"""
metadata = input_data['metadata']
model_name = input_data['model_name']
model_id = input_data['model_id']
experiment = input_data['experiment']
timestamp = input_data['timestamp']
status = input_data['status']
self.info(
f'Updating production model {model_name} from experiment {experiment}...', metadata)
@@ -321,15 +316,9 @@ class MLFlow(BaseActivity):
model_name=model_name
)
report = DataFrame([response])
report['model_id'] = model_id
report['model_name'] = model_name
report['timestamp'] = timestamp
report['status'] = status
self.info(
f'Production model {model_name} updated successfully', metadata)
return report.to_dict()
return response
except Exception as e:
trace = traceback.format_exc()

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:
"""

22
laborious/workflows/minimal_retrain.py
View File

@@ -91,13 +91,29 @@ class MinimalRetrain():
start_to_close_timeout=timedelta(seconds=60)
)
if experiment_response['success']:
update_report = await workflow.execute_activity_method(
Activities.update_production_model,
{
**metadata,
'model_name': model_name,
**experiment_response
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
else:
update_report = {}
report = await workflow.execute_activity_method(
Activities.update_production_model,
Activities.format_retrain_report,
{
**metadata,
'model_name': model_name,
'experiment_response': experiment_response,
'model_id': input_data['model_id'],
**experiment_response
'model_name': model_name,
'update_report': update_report
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)