SIENTIAPDE-1646

Refactor MLFlow retraining logic to always use retrain method

- Removed the conditional logic for full retraining, ensuring the `retrain` method is always called.
- Updated the documentation in the `retrain_model` method to reflect the changes in retraining flow.
- Adjusted tests to verify that the `retrain` method is invoked correctly, while ensuring `train` is not called when the full retrain flag is set.
This commit is contained in:
vitor-aignosi
2026-05-07 09:03:33 -03:00
parent 424be007ef
commit aaf647efdf
2 changed files with 10 additions and 19 deletions

View File

@@ -26,7 +26,6 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.utils.formatters import create_sample_dict
from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
from sientia_model.model_repository.plugin_store import PluginStore
from sklearn.model_selection import train_test_split
from laborious.utils.dataframe_debug import build_dataframe_debug_message
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
@@ -418,18 +417,17 @@ class MLFlow(MinioManager):
@activity.defn(name='retrain_model')
async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Fit an updated wrapper from historical data, log a new run, and register a model version.
Fit an updated wrapper from historical data, then log and register in MLflow.
Flow: load long-format data from MinIO → dedupe/pivot like inference prep → require
``model_config['target']`` → read current ``production`` version for ``source_run_id`` tag →
``start_run`` with retrain tags → ``wrapper.retrain`` or ``wrapper.train`` when
``full_retrain`` is set (optional ``validation_fraction``) → log input CSV artifact →
``store_model`` and ``log_params``. Does not promote; the workflow calls
run ``wrapper.retrain`` outside run timing → ``start_run`` with retrain tags → log input
CSV artifact → ``store_model`` and ``log_params``. Does not promote; the workflow calls
``update_production_model`` after validation.
Args:
- input_data: Must include ``metadata``, ``model_name``, ``data`` (payload), and
``model_config`` with at least ``target``; optional ``full_retrain``, ``validation_fraction``.
``model_config`` with at least ``target``.
Return:
On success: ``success``, ``experiment`` (``run_id``, ``experiment_id``, ``experiment_name``),
@@ -522,6 +520,9 @@ class MLFlow(MinioManager):
metadata=metadata,
)
# Keep heavy model fitting outside MLflow run timing.
wrapper.retrain(data)
run_name = f'{model_name}-retrain-{datetime.now().strftime("%Y%m%d%H%M%S")}'
with self.mlflow_repository.start_run(
@@ -531,17 +532,6 @@ class MLFlow(MinioManager):
tags={'retrain': 'true', 'source_run_id': source_run_id},
metadata=metadata,
) as run_info:
if model_config.get('full_retrain'):
val_frac = float(model_config.get('validation_fraction', 0.2))
train_df, val_df = train_test_split(data, test_size=val_frac, random_state=42)
wrapper.train(
train_data=train_df,
val_data=val_df,
target=target,
)
else:
wrapper.retrain(data)
tmp_dir = tempfile.mkdtemp(prefix='laborious_retrain_')
try:
raw_csv = Path(tmp_dir) / 'retrain_input.csv'