QTZPOC-20

QTZPOC-20: Remove PyTorch from supported model flavors as it was never functional due to missing runtime dependencies.
This commit is contained in:
2026-08-27 14:58:43 -03:00
parent 853c83e662
commit dfb9deae1a
7 changed files with 70 additions and 73 deletions

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