diff --git a/e2e/test_drift.py b/e2e/test_drift.py index 830b60e..7941bf5 100644 --- a/e2e/test_drift.py +++ b/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 diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index e5e1af5..1d6bbc9 100644 --- a/laborious/activities/mlflow.py +++ b/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 ``/`` 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) diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py index b9a512a..ebe285c 100644 --- a/laborious/activities/opc.py +++ b/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, diff --git a/laborious/utils/repository/opc_repository.py b/laborious/utils/repository/opc_repository.py index 1c17020..a0856cd 100644 --- a/laborious/utils/repository/opc_repository.py +++ b/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. diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 3a619a4..fcc396b 100644 --- a/laborious/worker/worker.py +++ b/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) diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index 3b98c83..030ad16 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/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)