QTZPOC-13

QTZPOC-13: Dynamically resolve MLflow pyfunc model code paths to support varied layouts during re-logging.

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 `<artifact>/code/`,
+        so the layout differs per model: some ship a `utils` package, others a
+        single self-contained module (`<model>_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/<model>/prediction_model`.
+            metadata (dict): Metadata for structured logging.
+
+        Returns:
+            list[str]: Existing paths inside `<artifact_path>/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)
 

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)