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/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(