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.

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)