diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index da622d9..b995907 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -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' diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index c427ca5..be17a12 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -477,7 +477,7 @@ async def test_retrain_model_success_with_payload_data( @patch('laborious.activities.mlflow.tempfile.mkdtemp') @patch('laborious.activities.mlflow.rmtree') @patch('laborious.activities.mlflow.to_datetime') -async def test_retrain_model_success_full_retrain_branch( +async def test_retrain_model_always_uses_retrain_even_with_full_retrain_flag( mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow ): mock_mkdtemp.return_value = 'tmp' @@ -510,7 +510,8 @@ async def test_retrain_model_success_full_retrain_branch( } ) - wrapper.train.assert_called_once() + wrapper.retrain.assert_called_once() + wrapper.train.assert_not_called() assert response['success'] is True