From dfb9deae1a0b11dcae0acdc4edfe04f5e74675d7 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 27 Aug 2026 14:58:43 -0300 Subject: [PATCH] QTZPOC-20 QTZPOC-20: Remove PyTorch from supported model flavors as it was never functional due to missing runtime dependencies. --- README.md | 4 +- docs/model-flavors.md | 6 +- e2e/conftest.py | 4 - e2e/scenarios.md | 2 +- e2e/test_predictions_batch_joblib_flavor.py | 2 +- .../utils/repository/model_repository.py | 28 +++--- .../utils/repository/test_model_repository.py | 97 ++++++++++--------- 7 files changed, 70 insertions(+), 73 deletions(-) diff --git a/README.md b/README.md index 48c75d6..d889abd 100644 --- a/README.md +++ b/README.md @@ -1109,8 +1109,8 @@ Legacy MinIO object layout (relative key): ## Model Flavors -`model_config.predict_flavor` / `model_config.transform_flavor` accept `sklearn`, `pyfunc`, -`pytorch` and `joblib`. The flavor is configuration, never discovered from the artifact, and there is +`model_config.predict_flavor` / `model_config.transform_flavor` accept `sklearn`, `pyfunc` and +`joblib`. The flavor is configuration, never discovered from the artifact, and there is no fallback between flavors: a failing loader propagates. `joblib` reads the artifact directly — download, first top-level `.pkl`, `joblib.load` — with the diff --git a/docs/model-flavors.md b/docs/model-flavors.md index a8db934..82c0f60 100644 --- a/docs/model-flavors.md +++ b/docs/model-flavors.md @@ -24,11 +24,13 @@ reads the artifact's `MLmodel` manifest to guess it. |---|---|---|---| | `sklearn` | yes | yes | `mlflow.sklearn.load_model()` | | `pyfunc` | yes | yes | `mlflow.pyfunc.load_model()`; the wrapper path also unwraps `._model_impl.python_model` | -| `pytorch` | yes | yes | `mlflow.pytorch.load_model()` | | `joblib` | yes | yes | downloads the artifact and `joblib.load`s its pickle — **no MLflow flavor loader involved** | Anything else raises `ValueError` with `INVALID_FLAVOR_MESSAGE`, **before** any artifact is resolved -or downloaded. The same message and the same four values govern reading and writing. +or downloaded. The same message and the same three values govern reading and writing. + +`pytorch` was accepted until this change and is now rejected like any other unknown value: `torch` is +not installed in the runtime, so `mlflow.pytorch.load_model` could never have loaded anything there. **There is no fallback between flavors.** A failing loader propagates: one attempt, one cause. A model configured with the wrong flavor fails visibly instead of being rescued silently. diff --git a/e2e/conftest.py b/e2e/conftest.py index 51af260..29076b3 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -546,10 +546,6 @@ 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) - # Mock pytorch.load_model - mock_mlflow.pytorch = MagicMock() - mock_mlflow.pytorch.load_model = MagicMock(return_value=mock_mlflow_models['predict_model']) - # Mock other mlflow methods that might be called mock_mlflow.set_tracking_uri = MagicMock() mock_mlflow.get_run = MagicMock( diff --git a/e2e/scenarios.md b/e2e/scenarios.md index 71abd5c..1654303 100644 --- a/e2e/scenarios.md +++ b/e2e/scenarios.md @@ -1124,7 +1124,7 @@ restores `sys.path` in teardown, because the flavor deliberately never removes w **Assertions**: - Exactly one prediction row, `prediction_status='Good'`, value matching the artifact's model - `download_artifacts` called once, with `'prediction_model'` -- No `mlflow.{sklearn,pyfunc,pytorch}.load_model` call names this model +- No `mlflow.{sklearn,pyfunc}.load_model` call names this model --- diff --git a/e2e/test_predictions_batch_joblib_flavor.py b/e2e/test_predictions_batch_joblib_flavor.py index 87b6f9d..0993e8f 100644 --- a/e2e/test_predictions_batch_joblib_flavor.py +++ b/e2e/test_predictions_batch_joblib_flavor.py @@ -26,7 +26,7 @@ from laborious.workflows.predictions_batch import PredictionsBatch base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}' -FLAVOR_LOADERS = ('sklearn', 'pyfunc', 'pytorch') +FLAVOR_LOADERS = ('sklearn', 'pyfunc') def build_input_data( diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 500622f..54e3391 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -45,7 +45,7 @@ ARTIFACTS_PATH = './tmp/artifacts' TRANSFORMED_COMPRESSED_PATH = 'artifacts/training_transformer.pkl' PREDICTION_COMPRESSED_PATH = 'artifacts/stacking_model.pkl' -INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch' or 'joblib'." +INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'joblib'." def force_memory_release(logger: Logger): @@ -550,7 +550,7 @@ class MLFlowRepository(SientiaMonitoring): Args: model_name (str): The name of the model to download from the registry. metadata (dict[str, Any]): Metadata used for structured logging. - flavor (str): Model flavor ('pyfunc', 'sklearn', 'pytorch', 'joblib') + flavor (str): Model flavor ('pyfunc', 'sklearn', 'joblib') artifact_path (str | None): Path to compressed artifacts if model is compressed Returns: @@ -564,7 +564,7 @@ class MLFlowRepository(SientiaMonitoring): `joblib.load`, which needs no `MLmodel` manifest. The artifact's `code` directory, when present, is put on `sys.path` first. """ - if flavor not in ('pyfunc', 'sklearn', 'pytorch', 'joblib'): + if flavor not in ('pyfunc', 'sklearn', 'joblib'): raise ValueError(INVALID_FLAVOR_MESSAGE) if flavor == 'joblib': @@ -584,10 +584,8 @@ class MLFlowRepository(SientiaMonitoring): try: if flavor == 'pyfunc': model = mlflow.pyfunc.load_model(model_uri) - elif flavor == 'sklearn': - model = mlflow.sklearn.load_model(model_uri) else: - model = mlflow.pytorch.load_model(model_uri) + model = mlflow.sklearn.load_model(model_uri) except Exception as e: await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels) raise e @@ -608,7 +606,7 @@ class MLFlowRepository(SientiaMonitoring): Args: model_name (str): The name of the model to download. metadata (dict[str, Any]): Metadata used for structured logging. - flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch', 'joblib') + flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib') artifact_path (str | None): Path to compressed artifacts if model is compressed Returns: @@ -626,7 +624,7 @@ class MLFlowRepository(SientiaMonitoring): directory, when present, is put on `sys.path` first. """ - if flavor not in ('sklearn', 'pyfunc', 'pytorch', 'joblib'): + if flavor not in ('sklearn', 'pyfunc', 'joblib'): raise ValueError(INVALID_FLAVOR_MESSAGE) latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production') @@ -650,10 +648,8 @@ class MLFlowRepository(SientiaMonitoring): try: if flavor == 'sklearn': model = mlflow.sklearn.load_model(model_uri) - elif flavor == 'pyfunc': - model = mlflow.pyfunc.load_model(model_uri) else: - model = mlflow.pytorch.load_model(model_uri) + model = mlflow.pyfunc.load_model(model_uri) except Exception as e: await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels) raise e @@ -677,7 +673,7 @@ class MLFlowRepository(SientiaMonitoring): model_name (str): Name of the model to download metadata (dict[str, Any]): Metadata used for structured logging. model_type (str): Type of model ('predict' or 'transform') - flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch', 'joblib') + flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib') load_wrapper (bool): Whether to load wrapper Returns: @@ -862,7 +858,7 @@ class MLFlowRepository(SientiaMonitoring): metadata (dict[str, Any]): Metadata used for structured logging. retention (int): Cache retention time in minutes (0 = no cache). model_type (str): Type of model ('predict' or 'transform') - flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch') + flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib') Returns: Any: Model object. @@ -950,7 +946,7 @@ class MLFlowRepository(SientiaMonitoring): model_name (str): Registered model name. data (pd.DataFrame): Input data. retention (int): Cache retention in minutes. - flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch'). + flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib'). metadata (dict[str, Any]): Metadata used for structured logging. Returns: pd.DataFrame | ndarray: Operation result. @@ -1228,7 +1224,7 @@ class MLFlowRepository(SientiaMonitoring): Args: model_data (dict): Model holder with keys 'model' and optional 'artifact_path'. - flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch', 'joblib'). + flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib'). model_type (str): Artifact name, e.g., 'prediction_model' or 'data_model'. metadata (dict): Metadata for structured logging. @@ -1257,8 +1253,6 @@ class MLFlowRepository(SientiaMonitoring): model.store_model(artifact_path=model_type, code_path=code_path, to_disk=False) self.debug('Model uploaded successfully', metadata) - elif flavor == 'pytorch': - mlflow.pytorch.log_model(model, model_type) elif flavor == 'joblib': with tempfile.TemporaryDirectory() as tmp_dir: pkl_path = path.join(tmp_dir, 'model.pkl') diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 2c449de..02e8cd5 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -455,17 +455,6 @@ async def test_load_predict_model_pyfunc(mlflow, mlflow_repository): ) -@pytest.mark.asyncio -async def test_load_predict_model_pytorch(mlflow, mlflow_repository): - result = await mlflow_repository.load_predict_model('test_model', {}, 'pytorch') - assert result == mlflow.pytorch.load_model.return_value - mlflow.pytorch.load_model.assert_called_once_with('models:/test_model/production') - mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY) - mlflow_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MODEL_READ_COUNT, tags=ANY - ) - - @pytest.mark.asyncio async def test_load_predict_model_error(mlflow_repository): with pytest.raises(ValueError) as e: @@ -523,23 +512,6 @@ async def test_load_transform_model_pyfunc(mlflow, mlflow_repository): ) -@pytest.mark.asyncio -async def test_load_transform_model_pytorch(mlflow, mlflow_repository): - mlflow_repository.get_model_run_id = MagicMock() - mlflow_repository.get_model_uri = MagicMock() - - result = await mlflow_repository.load_transform_model('test_model', {}, 'pytorch') - validate_common_load_transform_model_mocks(mlflow_repository, 'test_model') - - assert result == mlflow.pytorch.load_model.return_value - mlflow.pytorch.load_model.assert_called_once_with(mlflow_repository.get_model_uri.return_value) - - mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY) - mlflow_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MODEL_READ_COUNT, tags=ANY - ) - - @pytest.mark.asyncio async def test_load_transform_model_error(mlflow_repository): mlflow_repository.get_model_run_id = MagicMock() @@ -1108,20 +1080,6 @@ async def test_log_model_pyfunc(tmp_path, mlflow, mlflow_repository): ) -@pytest.mark.asyncio -async def test_log_model_pytorch(mlflow, mlflow_repository): - model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'} - await mlflow_repository.log_model( - model_data, 'pytorch', 'prediction_model', metadata['metadata'] - ) - mlflow.pytorch.log_model.assert_called_once_with(model_data['model'], 'prediction_model') - - mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_WRITE_LAG, ANY) - mlflow_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MODEL_WRITE_COUNT, tags=ANY - ) - - @pytest.mark.asyncio async def test_log_model_error(mlflow_repository): model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'} @@ -1777,7 +1735,7 @@ Tests for the joblib flavor on the model load and write paths """ JOBLIB_PATH = 'laborious.utils.repository.model_repository.joblib' -FLAVOR_LOADERS = ('sklearn', 'pyfunc', 'pytorch') +FLAVOR_LOADERS = ('sklearn', 'pyfunc') @pytest.fixture @@ -1889,7 +1847,7 @@ async def test_load_transform_model_joblib_resolves_run_once( @patch(JOBLIB_PATH) @pytest.mark.asyncio -@pytest.mark.parametrize('flavor', ['sklearn', 'pyfunc', 'pytorch']) +@pytest.mark.parametrize('flavor', ['sklearn', 'pyfunc']) async def test_load_predict_model_flavor_failure_propagates( joblib_mock, mlflow, mlflow_repository, flavor ): @@ -1910,7 +1868,7 @@ async def test_load_predict_model_flavor_failure_propagates( @patch(JOBLIB_PATH) @pytest.mark.asyncio -@pytest.mark.parametrize('flavor', ['sklearn', 'pyfunc', 'pytorch']) +@pytest.mark.parametrize('flavor', ['sklearn', 'pyfunc']) async def test_load_transform_model_flavor_failure_propagates( joblib_mock, mlflow, mlflow_repository, flavor ): @@ -2361,10 +2319,57 @@ async def test_log_model_invalid_flavor(joblib_mock, mlflow, mlflow_repository): mlflow_repository.observe_lag.assert_not_called() +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_load_predict_model_rejects_pytorch(joblib_mock, mlflow, mlflow_repository): + with pytest.raises(ValueError) as exc_info: + await mlflow_repository.load_predict_model('test_model', {}, 'pytorch') + + assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE + mlflow.pytorch.load_model.assert_not_called() + joblib_mock.load.assert_not_called() + mlflow_repository.client.download_artifacts.assert_not_called() + mlflow_repository.emit_metric.assert_not_called() + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_load_transform_model_rejects_pytorch(joblib_mock, mlflow, mlflow_repository): + with pytest.raises(ValueError) as exc_info: + await mlflow_repository.load_transform_model('test_model', {}, 'pytorch') + + assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE + mlflow.pytorch.load_model.assert_not_called() + joblib_mock.load.assert_not_called() + mlflow_repository.client.download_artifacts.assert_not_called() + mlflow_repository.emit_metric.assert_not_called() + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_log_model_rejects_pytorch(joblib_mock, mlflow, mlflow_repository): + model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'} + + with pytest.raises(ValueError) as exc_info: + await mlflow_repository.log_model( + model_data, 'pytorch', 'prediction_model', metadata['metadata'] + ) + + assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE + mlflow.pytorch.log_model.assert_not_called() + joblib_mock.dump.assert_not_called() + mlflow_repository.emit_metric.assert_called_once_with( + metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY + ) + mlflow_repository.observe_lag.assert_not_called() + + def test_invalid_flavor_message_lists_every_accepted_flavor(): - for flavor in ('sklearn', 'pyfunc', 'pytorch', 'joblib'): + for flavor in ('sklearn', 'pyfunc', 'joblib'): assert f"'{flavor}'" in INVALID_FLAVOR_MESSAGE + assert 'pytorch' not in INVALID_FLAVOR_MESSAGE + @patch(JOBLIB_PATH) @pytest.mark.asyncio