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.

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'

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