gitea_admin/sientia-dataops-laborious-legacy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

QTZPOC-20

Browse Source 
QTZPOC-20: Remove PyTorch from supported model flavors as it was never functional due to missing runtime dependencies.
This commit is contained in:
Vitor Pimentel dos Santos
2026-08-27 14:58:43 -03:00
parent 853c83e662
commit dfb9deae1a
7 changed files with 70 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
README.md
View File

@@ -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

6
docs/model-flavors.md
View File

@@ -24,11 +24,13 @@ reads the artifact's `MLmodel` manifest to guess it.
|---|---|---|---|
| `sklearn` | yes | yes | `mlflow.sklearn.load_model(<uri>)` |
| `pyfunc` | yes | yes | `mlflow.pyfunc.load_model(<uri>)`; the wrapper path also unwraps `._model_impl.python_model` |
| `pytorch` | yes | yes | `mlflow.pytorch.load_model(<uri>)` |
| `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.

4
e2e/conftest.py
View File

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

2
e2e/scenarios.md
View File

@@ -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
---

2
e2e/test_predictions_batch_joblib_flavor.py
View File

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

28
laborious/utils/repository/model_repository.py
View File

@@ -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')

97
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -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