QTZPOC-20

QTZPOC-20: Prioritize code/ directory download and sys.path injection before model loading, enabling MLflow flavors to stream models directly from URIs and simplifying joblib artifact handling.

e2e/conftest.py

@@ -546,14 +546,7 @@ def patch_mlflow(mock_mlflow_models, joblib_flavor_artifacts):
     mock_mlflow.pyfunc = MagicMock()
     mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model)
 
-    # The sklearn and pyfunc paths resolve the URI to a local directory before loading, so that the
-    # artifact's `code/` can go on `sys.path` whatever the `MLmodel` manifest says. Echoing the URI
-    # keeps the loader stubs above — which branch on the URI text — working unchanged, and points
-    # the code-directory lookup at a path that does not exist.
     mock_mlflow.artifacts = MagicMock()
-    mock_mlflow.artifacts.download_artifacts = MagicMock(
-        side_effect=lambda artifact_uri: str(artifact_uri)
-    )
 
     # Mock other mlflow methods that might be called
     mock_mlflow.set_tracking_uri = MagicMock()
@@ -572,9 +565,23 @@ def patch_mlflow(mock_mlflow_models, joblib_flavor_artifacts):
         ),
     )
 
-    # Only the joblib flavor calls download_artifacts directly on the client, one artifact
-    # per model type — a single return value would serve `data_model` from `prediction_model`.
+    # `download_model` asks whether the run's artifact carries a `code` directory before it loads
+    # anything, and downloads only that directory when it does. Answering "no" here keeps the
+    # download counts of these tests about the model itself; a test that exercises the `code` step
+    # overrides `list_artifacts` to answer "yes".
+    def mock_list_artifacts(_run_id, artifact_path=None):
+        if artifact_path:
+            return []
+        return [MagicMock(path=artifact_name) for artifact_name in joblib_flavor_artifacts]
+
+    mock_mlflow_client.list_artifacts = MagicMock(side_effect=mock_list_artifacts)
+
+    # The joblib flavor downloads a whole model artifact, one per model type — a single return
+    # value would serve `data_model` from `prediction_model`. The `code` step asks for the nested
+    # `<artifact>/code` instead, and gets that directory back, as MLflow would return it.
     def mock_download_artifacts(_run_id, artifact_name):
+        if artifact_name.endswith('/code'):
+            return joblib_flavor_artifacts[artifact_name.removesuffix('/code')]['code_dir']
         return joblib_flavor_artifacts[artifact_name]['path']
 
     mock_mlflow_client.download_artifacts = MagicMock(side_effect=mock_download_artifacts)

e2e/scenarios.md

@@ -1149,9 +1149,34 @@ restores `sys.path` in teardown, because the flavor deliberately never removes w
 #### Scenario 5.1.3: An MLflow Flavor Is Unaffected
 **Description**: A model declaring `sklearn` behaves exactly as before
 
+**Input**:
+- `client.list_artifacts` answering that the run's artifacts carry no `code` directory
+
 **Assertions**:
 - The prediction lands with the mocked model's value
-- `download_artifacts` NOT called — the joblib path is never entered as a safety net
+- `download_artifacts` NOT called — the joblib path is never entered as a safety net, and with no
+  `code` directory to fetch there is nothing else to download either
+
+---
+
+#### Scenario 5.1.4: Only the Artifact's `code` Directory Is Downloaded
+**Description**: A model on an MLflow flavor whose artifact **does** carry `code/`. The flavor
+loaders read the model URI themselves and fetch what they need, so the `code` directory — the one
+thing MLflow will not put on `sys.path` unless the `MLmodel` manifest declares it — is the whole of
+what this workflow downloads through the client
+
+**Input**:
+- `model_config` on the default `sklearn` flavors
+- `client.list_artifacts` answering that both artifacts carry `code`
+
+**Expected Behavior**:
+- The `code` directory of each artifact is downloaded before its model is loaded, and goes on
+  `sys.path`
+- Neither model's weights travel through `client.download_artifacts`
+
+**Assertions**:
+- The prediction lands with the mocked model's value
+- `download_artifacts` called exactly twice, with `'data_model/code'` and `'prediction_model/code'`
 
 ---
 

e2e/test_predictions_batch_joblib_flavor.py

@@ -7,6 +7,8 @@ A passing run therefore proves the flavor decompresses the pickle **and** resolv
 without any `mlflow.*.load_model` being called for that model.
 """
 
+from unittest.mock import MagicMock
+
 import pandas as pd
 import pytest
 from sqlalchemy import text
@@ -208,6 +210,41 @@ async def test_prediction_workflow_with_sklearn_flavor_is_unchanged(
     mlflow_client.download_artifacts.assert_not_called()
 
 
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_prediction_workflow_downloads_only_the_artifact_code_directory(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    joblib_flavor_artifacts,
+    mlflow_client,
+    mock_mlflow_models,
+):
+    """A flavor model whose artifact carries `code` downloads that directory — and only it.
+
+    The flavor loaders read the model URI themselves, so nothing else is fetched through the
+    client: the `code` directory is the whole of what this workflow downloads.
+    """
+    mlflow_client.list_artifacts.side_effect = lambda _run_id, artifact_path=None: (
+        [MagicMock(path=f'{artifact_path}/code')] if artifact_path else []
+    )
+
+    client = temporal_test_env.client
+    model_id = 275
+    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
+    input_data = build_input_data(model_id, 'test_model')
+
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-code-only')
+    )
+
+    assert_prediction(postgres_engine, model_id, prediction=0.5)
+
+    downloaded = [loaded.args[1] for loaded in mlflow_client.download_artifacts.call_args_list]
+    assert sorted(downloaded) == ['data_model/code', 'prediction_model/code']
+
+
 @pytest.mark.asyncio
 @pytest.mark.integration
 async def test_prediction_workflow_joblib_flavor_failure_reports_single_cause(

e2e/test_retrain_joblib_flavor.py

@@ -117,7 +117,9 @@ async def test_retrain_joblib_model_writes_pickle_and_code_and_reloads(
 
     mlflow_client.download_artifacts.side_effect = lambda _run_id, _artifact_name: str(written)
 
-    reloaded = await repository.load_predict_model(JOBLIB_FLAVOR_MODEL_NAME, METADATA, 'joblib')
+    reloaded, _local_dir = await repository.load_predict_model(
+        JOBLIB_FLAVOR_MODEL_NAME, METADATA, 'joblib'
+    )
 
     features = pd.DataFrame({'feature_1': [0.234], 'feature_2': [0.783]})
     expected = joblib_flavor_artifacts['prediction_model']['model'].predict(features)