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)