SIENTIAPDE-1646

Enhance MLFlow reference data handling and testing

- Updated the artifact handling in MLFlow to prioritize 'retrain_input.csv' over 'train_data.csv' when resolving reference data.
- Introduced new methods for resolving artifact names and locating downloaded CSV files.
- Modified the `get_reference_data` method to improve artifact resolution and error handling.
- Expanded unit tests to cover scenarios for missing artifacts and preference logic between retrain and train data CSVs.

e2e/test_drift.py

@@ -142,7 +142,7 @@ def _recent_minute_timestamps(count: int, offset_minutes: int = 6) -> list[str]:
 def _configure_reference_csv(mlflow_repository_stub, reference_rows: pd.DataFrame) -> None:
     """
     Wire ``mlflow_repository_stub`` so ``get_reference_data`` returns
-    ``reference_rows`` by writing them to ``dst_path/evaluation_data.csv``.
+    ``reference_rows`` by writing them to ``dst_path/retrain_input.csv``.
 
     Args:
      - mlflow_repository_stub: External MLflow repository fixture.
@@ -150,12 +150,16 @@ def _configure_reference_csv(mlflow_repository_stub, reference_rows: pd.DataFram
     """
 
     def _download(run_id: str, artifact_path: str, dst_path: str, metadata=None):
-        target = Path(dst_path) / 'evaluation_data.csv'
+        target = Path(dst_path) / artifact_path
+        target.parent.mkdir(parents=True, exist_ok=True)
         reference_rows.to_csv(target, index=False)
 
     mlflow_repository_stub._client.get_model_version_by_alias.return_value = MagicMock(
         run_id='fake-reference-run'
     )
+    file_info = MagicMock()
+    file_info.path = 'retrain_input.csv'
+    mlflow_repository_stub._client.list_artifacts.return_value = [file_info]
     mlflow_repository_stub.download_artifacts.side_effect = _download
 
 

laborious/activities/mlflow.py

@@ -53,6 +53,7 @@ class MLFlow(SientiaMonitoring):
 
     _MAX_DEBUG_DATAFRAME_ROWS = 100
     _DEFAULT_MODEL_ALIAS = 'production'
+    _REFERENCE_ARTIFACT_CANDIDATES = ('retrain_input.csv', 'train_data.csv')
 
     def __init__(
         self,
@@ -639,21 +640,59 @@ class MLFlow(SientiaMonitoring):
             self.error(trace, metadata=metadata)
             raise e
 
+    def _resolve_reference_artifact_name(self, run_id: str) -> str | None:
+        """
+        Pick the first available reference CSV artifact path from the MLflow run.
+
+        Candidates are checked in priority order: ``retrain_input.csv``, then ``train_data.csv``.
+        A path matches when it equals the candidate or ends with ``/<candidate>`` for nested layouts.
+
+        Args:
+         - run_id: MLflow run UUID linked to the production model version.
+
+        Return:
+            Artifact path string for ``download_artifacts``, or ``None`` if no candidate exists.
+        """
+        listed = self.mlflow_repository._client.list_artifacts(run_id)
+        paths = [file_info.path for file_info in listed]
+        for candidate in self._REFERENCE_ARTIFACT_CANDIDATES:
+            for path in paths:
+                if path == candidate or path.endswith(f'/{candidate}'):
+                    return path
+        return None
+
+    def _find_downloaded_csv(self, tmpdir: str, artifact_name: str) -> Path | None:
+        """
+        Locate a downloaded reference CSV in the temp directory.
+
+        Args:
+         - tmpdir: Directory where ``download_artifacts`` wrote files.
+         - artifact_name: Basename of the resolved artifact (e.g. ``retrain_input.csv``).
+
+        Return:
+            ``Path`` to the CSV file if found, else ``None``.
+        """
+        direct = Path(tmpdir) / artifact_name
+        if direct.exists():
+            return direct
+        matches = list(Path(tmpdir).rglob(artifact_name))
+        return matches[0] if matches else None
+
     @activity.defn(name='get_reference_data')
     def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
         """
-        Download ``evaluation_data.csv`` from the MLflow run linked to ``production`` and parse it.
+        Download reference training CSV from the MLflow run linked to the production alias.
 
-        Used by drift workflows to compare live data against the reference distribution logged with
-        the model. Artifacts are downloaded to a temp directory, discovered via ``rglob`` (nested
-        layout-safe), then timestamps are normalized to ``DATETIME_FORMAT`` string columns before
-        returning record-oriented dicts.
+        Resolves ``retrain_input.csv`` or ``train_data.csv`` via artifact listing before download.
+        ``retrain_input.csv`` is preferred when both exist (most recent retrain snapshot). Used by
+        drift workflows to compare live data against the reference distribution logged with the model.
+        Timestamps are normalized to ``DATETIME_FORMAT`` string columns before returning records.
 
         Args:
-         - input_data: ``metadata`` and ``model_name`` for registry lookup.
+         - input_data: ``metadata``, ``model_name``, and optional ``model_config`` with ``alias``.
 
         Return:
-            List of row dicts, or ``None`` if the artifact path is missing or any step fails.
+            List of row dicts with normalized timestamps, or ``None`` if resolution or load fails.
         """
         metadata = input_data['metadata']
         model_name = input_data['model_name']
@@ -666,20 +705,25 @@ class MLFlow(SientiaMonitoring):
             )
             run_id = mv.run_id
 
+            artifact_path = self._resolve_reference_artifact_name(run_id)
+            if artifact_path is None:
+                self.warning(f'Reference data not found for model {model_name}', metadata)
+                return None
+
+            artifact_name = Path(artifact_path).name
             tmpdir = tempfile.mkdtemp(prefix='laborious_eval_')
             try:
                 self.mlflow_repository.download_artifacts(
                     run_id=run_id,
-                    artifact_path='evaluation_data.csv',
+                    artifact_path=artifact_path,
                     dst_path=tmpdir,
                     metadata=metadata,
                 )
-                csv_candidates = list(Path(tmpdir).rglob('evaluation_data.csv'))
-                if not csv_candidates:
+                csv_path = self._find_downloaded_csv(tmpdir, artifact_name)
+                if csv_path is None:
                     self.warning(f'Reference data not found for model {model_name}', metadata)
                     return None
-
-                reference_data = pd.read_csv(csv_candidates[0])
+                reference_data = pd.read_csv(csv_path)
             finally:
                 rmtree(tmpdir, ignore_errors=True)
 

laborious/activities/opc.py

@@ -143,9 +143,7 @@ class OPC(SientiaMonitoring):
                     attachment_content=error_data.get('attachment_content', None),
                 )
             else:
-                self.info(
-                    f'OPC server {opc_id}:{server["server_name"]} connected successfully.'
-                )
+                self.info(f'OPC server {opc_id}:{server["server_name"]} connected successfully.')
 
     def write_data(
         self,

laborious/utils/repository/opc_repository.py

@@ -755,9 +755,7 @@ class OpcRepository(SientiaMonitoring):
                 pass
         node_obj.set_value(data, variant_type)
 
-    def _write_reconnect_in_progress(
-        self, metadata: dict[str, Any]
-    ) -> tuple[bool, dict[str, Any]]:
+    def _write_reconnect_in_progress(self, metadata: dict[str, Any]) -> tuple[bool, dict[str, Any]]:
         """
         Fail a write because a background reconnect thread is already running.
 

laborious/worker/worker.py

@@ -147,9 +147,13 @@ async def main():
         to_install_runtime = runtime
 
     try:
-        await plugin_store.install_runtime(runtime_name=to_install_runtime, metadata=metadata_runtime)
+        await plugin_store.install_runtime(
+            runtime_name=to_install_runtime, metadata=metadata_runtime
+        )
     except Exception as exc:
-        logger.custom_critical(f'Failed to install runtime {to_install_runtime}: {exc}', metadata_runtime)
+        logger.custom_critical(
+            f'Failed to install runtime {to_install_runtime}: {exc}', metadata_runtime
+        )
         metrics.APP_UP.labels(pod_id=POD_ID).set(0)
         sys.exit(1)
 

tests/laborious/activities/test_mlflow.py

@@ -652,6 +652,12 @@ def test_update_production_model_error(mlflow):
         raise AssertionError('Expected exception')
 
 
+def _artifact_file_info(path: str) -> MagicMock:
+    file_info = MagicMock()
+    file_info.path = path
+    return file_info
+
+
 @patch('laborious.activities.mlflow.to_datetime')
 def test_get_reference_data_success(mock_to_datetime, mlflow):
     input_data = {
@@ -661,6 +667,9 @@ def test_get_reference_data_success(mock_to_datetime, mlflow):
 
     mv = MagicMock(run_id='run1')
     mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('retrain_input.csv'),
+    ]
 
     mock_reference_data = MagicMock()
     mock_reference_data.__getitem__.return_value = MagicMock()
@@ -672,10 +681,19 @@ def test_get_reference_data_success(mock_to_datetime, mlflow):
     with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
         with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
             with patch('laborious.activities.mlflow.rmtree'):
-                with patch('laborious.activities.mlflow.Path') as mp:
-                    mp.return_value.rglob.return_value = [MagicMock()]
+                with patch.object(
+                    mlflow,
+                    '_find_downloaded_csv',
+                    return_value=MagicMock(),
+                ):
                     result = mlflow.get_reference_data(input_data)
 
+    mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
+        run_id='run1',
+        artifact_path='retrain_input.csv',
+        dst_path='/t',
+        metadata=metadata['metadata'],
+    )
     mock_reference_data.to_dict.assert_called_once_with(orient='records')
     assert result == mock_reference_data.to_dict.return_value
 
@@ -694,6 +712,97 @@ def test_get_reference_data_not_found(mlflow):
     assert result is None
 
 
+def test_get_reference_data_only_train_data_csv(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+    mv = MagicMock(run_id='run1')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('train_data.csv'),
+    ]
+
+    mock_reference_data = MagicMock()
+    mock_reference_data.__getitem__.return_value = MagicMock()
+    mock_reference_data.to_dict.return_value = [{'timestamp': '2023-05-26 11:12:27', 'value': 1.0}]
+
+    with patch('laborious.activities.mlflow.to_datetime') as mock_to_datetime:
+        mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
+        with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
+            with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
+                with patch('laborious.activities.mlflow.rmtree'):
+                    with patch.object(
+                        mlflow,
+                        '_find_downloaded_csv',
+                        return_value=MagicMock(),
+                    ):
+                        result = mlflow.get_reference_data(input_data)
+
+    mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
+        run_id='run1',
+        artifact_path='train_data.csv',
+        dst_path='/t',
+        metadata=metadata['metadata'],
+    )
+    assert result == mock_reference_data.to_dict.return_value
+
+
+def test_get_reference_data_prefers_retrain_input_when_both_listed(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+    mv = MagicMock(run_id='run1')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('train_data.csv'),
+        _artifact_file_info('retrain_input.csv'),
+    ]
+
+    mock_reference_data = MagicMock()
+    mock_reference_data.__getitem__.return_value = MagicMock()
+    mock_reference_data.to_dict.return_value = [{'timestamp': '2023-05-26 11:12:27', 'value': 1.0}]
+
+    with patch('laborious.activities.mlflow.to_datetime') as mock_to_datetime:
+        mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
+        with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
+            with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
+                with patch('laborious.activities.mlflow.rmtree'):
+                    with patch.object(
+                        mlflow,
+                        '_find_downloaded_csv',
+                        return_value=MagicMock(),
+                    ):
+                        result = mlflow.get_reference_data(input_data)
+
+    mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
+        run_id='run1',
+        artifact_path='retrain_input.csv',
+        dst_path='/t',
+        metadata=metadata['metadata'],
+    )
+    assert result == mock_reference_data.to_dict.return_value
+
+
+def test_get_reference_data_no_candidate_artifacts_returns_none(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+    mv = MagicMock(run_id='run1')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('other_artifact.csv'),
+    ]
+
+    result = mlflow.get_reference_data(input_data)
+
+    mlflow.mlflow_repository.download_artifacts.assert_not_called()
+    mlflow.warning.assert_called()
+    assert result is None
+
+
 def test_get_reference_data_missing_csv_file_returns_none(mlflow):
     input_data = {
         **metadata,
@@ -701,13 +810,17 @@ def test_get_reference_data_missing_csv_file_returns_none(mlflow):
     }
     mv = MagicMock(run_id='run1')
     mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('retrain_input.csv'),
+    ]
 
     with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='tmp'):
         with patch('laborious.activities.mlflow.rmtree'):
-            with patch('laborious.activities.mlflow.Path') as mp:
-                mp.return_value.rglob.return_value = []
+            with patch.object(mlflow, '_find_downloaded_csv', return_value=None):
                 result = mlflow.get_reference_data(input_data)
 
+    mlflow.mlflow_repository.download_artifacts.assert_called_once()
+    mlflow.warning.assert_called()
     assert result is None
 
 
@@ -719,6 +832,9 @@ def test_get_reference_data_exception(mlflow):
 
     mv = MagicMock(run_id='run1')
     mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('retrain_input.csv'),
+    ]
     mlflow.mlflow_repository.download_artifacts.side_effect = Exception('dl fail')
 
     result = mlflow.get_reference_data(input_data)