From d7bf36605d40d8cd9aae6d137461f2bb2f0a79cc Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 19 Aug 2026 11:01:52 -0300 Subject: [PATCH] QTZPOC-13 QTZPOC-13: Dynamically resolve MLflow pyfunc model code paths to support varied layouts during re-logging. --- .../utils/repository/model_repository.py | 39 ++++++++++++++++- .../utils/repository/test_model_repository.py | 43 ++++++++++++++++--- 2 files changed, 75 insertions(+), 7 deletions(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index d6dc2bc..ab007d4 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -1094,6 +1094,43 @@ class MLFlowRepository(SientiaMonitoring): } return retrain_data + def resolve_code_paths(self, artifact_path: str | None, metadata: dict) -> list[str]: + """Resolve the code dependencies to forward when re-logging a pyfunc model. + + MLflow flattens whatever was given as `code_paths` into `/code/`, + so the layout differs per model: some ship a `utils` package, others a + single self-contained module (`_model.py`). Re-logging must forward + the entries that are actually on disk — assuming one fixed layout makes + `_validate_and_copy_code_paths` raise `FileNotFoundError` for the other. + + Args: + artifact_path (str | None): Local directory of the downloaded model + artifact, e.g. `./tmp/artifacts//prediction_model`. + metadata (dict): Metadata for structured logging. + + Returns: + list[str]: Existing paths inside `/code`, empty when + the model carries no code directory. + """ + if not artifact_path: + self.warning('No artifact path available, logging model without code paths', metadata) + return [] + + code_dir = path.join(artifact_path, 'code') + + if not path.isdir(code_dir): + self.warning( + f'No code directory found at {code_dir}, logging model without code paths', + metadata, + ) + return [] + + return [ + path.join(code_dir, entry) + for entry in sorted(listdir(code_dir)) + if entry != '__pycache__' + ] + async def log_model(self, model_data: dict, flavor: str, model_type: str, metadata: dict): """Log a model into the active MLflow run. @@ -1114,7 +1151,7 @@ class MLFlowRepository(SientiaMonitoring): if flavor == 'sklearn': mlflow.sklearn.log_model(model, model_type) elif flavor == 'pyfunc': - code_path = [path.join(model_data['artifact_path'], 'code', 'utils')] + code_path = self.resolve_code_paths(model_data.get('artifact_path'), metadata) self.debug(f'Code path: {code_path}', metadata) diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 64c5e5f..457e29f 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -1049,21 +1049,52 @@ async def test_log_model_sklearn(mlflow, mlflow_repository): ) -@patch('laborious.utils.repository.model_repository.path') +def test_resolve_code_paths_package_layout(tmp_path, mlflow_repository): + code_dir = tmp_path / 'code' + (code_dir / 'utils').mkdir(parents=True) + (code_dir / '__pycache__').mkdir() + + assert mlflow_repository.resolve_code_paths(str(tmp_path), metadata['metadata']) == [ + str(code_dir / 'utils') + ] + + +def test_resolve_code_paths_single_module_layout(tmp_path, mlflow_repository): + code_dir = tmp_path / 'code' + code_dir.mkdir() + (code_dir / 'loop_valve_classifier_model.py').touch() + + assert mlflow_repository.resolve_code_paths(str(tmp_path), metadata['metadata']) == [ + str(code_dir / 'loop_valve_classifier_model.py') + ] + + +def test_resolve_code_paths_without_code_dir(tmp_path, mlflow_repository): + assert mlflow_repository.resolve_code_paths(str(tmp_path), metadata['metadata']) == [] + + +def test_resolve_code_paths_without_artifact_path(mlflow_repository): + assert mlflow_repository.resolve_code_paths(None, metadata['metadata']) == [] + + @pytest.mark.asyncio -async def test_log_model_pyfunc(path, mlflow, mlflow_repository): +async def test_log_model_pyfunc(tmp_path, mlflow, mlflow_repository): + code_dir = tmp_path / 'code' + code_dir.mkdir() + (code_dir / 'loop_valve_classifier_model.py').touch() + model_mock = MagicMock() - model_data = {'model': model_mock, 'artifact_path': 'artifact_path'} + model_data = {'model': model_mock, 'artifact_path': str(tmp_path)} await mlflow_repository.log_model( model_data, 'pyfunc', 'prediction_model', metadata['metadata'] ) mlflow.pyfunc.log_model.assert_not_called() - path.join.assert_called_once_with('artifact_path', 'code', 'utils') - model_mock.store_model.assert_called_once_with( - artifact_path='prediction_model', code_path=[path.join.return_value], to_disk=False + artifact_path='prediction_model', + code_path=[str(code_dir / 'loop_valve_classifier_model.py')], + to_disk=False, ) mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_WRITE_LAG, ANY)