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: Refactor ModelImport status engine for lazy initialization, add byte-counted size validation to bundle extraction, optimize MLflow run ID lookups, and ensure workflow cancellation propagates correctly.
This commit is contained in:
Vitor Pimentel dos Santos
2026-09-10 08:59:31 -03:00
parent d06ce0daae
commit 4c7cae502f
11 changed files with 452 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
tests/laborious/activities/test_model_import_bundle.py
View File

@@ -72,24 +72,25 @@ def minio_client() -> MagicMock:
def import_activity(tmp_path: Path, minio_client: MagicMock) -> ModelImport:
repository = MagicMock()
repository.minio_client = minio_client
with patch('laborious.activities.model_import.create_engine', return_value=FakeStatusEngine()):
activity = ModelImport(
status_config={
'host': 'h',
'port': 5432,
'dbname': 'sientia-core-mlops-bff',
'user': 'u',
'password': 'p',
'schema': 'public',
},
import_config=import_config(tmp_path),
minio_repository=repository,
mlflow_repository=MagicMock(),
mongo_repository=AsyncMock(),
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
)
activity = ModelImport(
status_config={
'host': 'h',
'port': 5432,
'dbname': 'sientia-core-mlops-bff',
'user': 'u',
'password': 'p',
'schema': 'public',
},
import_config=import_config(tmp_path),
minio_repository=repository,
mlflow_repository=MagicMock(),
mongo_repository=AsyncMock(),
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
)
# `status_engine` is created lazily on first access; pre-seed it with the recording double.
activity._status_engine = FakeStatusEngine()
activity.emit_metric = AsyncMock()
activity.send_notification_async = AsyncMock()
return activity

61
tests/laborious/activities/test_model_import_provisioning.py
View File

@@ -100,36 +100,37 @@ def import_activity(tmp_path: Path, mlflow_client: MagicMock, mongo: AsyncMock)
minio_repository = MagicMock()
mlflow_repository = MagicMock()
mlflow_repository.client = mlflow_client
with patch('laborious.activities.model_import.create_engine', return_value=FakeStatusEngine()):
activity = ModelImport(
status_config={
'host': 'h',
'port': 5432,
'dbname': 'sientia-core-mlops-bff',
'user': 'u',
'password': 'p',
'schema': 'public',
},
import_config={
'password_key': None,
'bucket': 'sientia',
'prefix': 'imported_models/',
'work_dir': str(tmp_path),
'max_object_bytes': 1024,
'max_entries': 10,
'max_uncompressed_bytes': 4096,
'max_compression_ratio': 200.0,
'models_collection': 'models',
'counters_collection': 'model_id_counters',
'retention_days': 7,
},
minio_repository=minio_repository,
mlflow_repository=mlflow_repository,
mongo_repository=mongo,
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
)
activity = ModelImport(
status_config={
'host': 'h',
'port': 5432,
'dbname': 'sientia-core-mlops-bff',
'user': 'u',
'password': 'p',
'schema': 'public',
},
import_config={
'password_key': None,
'bucket': 'sientia',
'prefix': 'imported_models/',
'work_dir': str(tmp_path),
'max_object_bytes': 1024,
'max_entries': 10,
'max_uncompressed_bytes': 4096,
'max_compression_ratio': 200.0,
'models_collection': 'models',
'counters_collection': 'model_id_counters',
'retention_days': 7,
},
minio_repository=minio_repository,
mlflow_repository=mlflow_repository,
mongo_repository=mongo,
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
)
# `status_engine` is created lazily on first access; pre-seed it with the recording double.
activity._status_engine = FakeStatusEngine()
activity.emit_metric = AsyncMock()
activity.send_notification_async = AsyncMock()
return activity

80
tests/laborious/activities/test_model_import_status.py
View File

@@ -164,17 +164,19 @@ def test_build_update_with_no_casts_or_extra_where_renders_a_bare_update():
def build_activity(engine: FakeStatusEngine) -> ModelImport:
"""A `ModelImport` whose only connection is the recording engine double."""
with patch('laborious.activities.model_import.create_engine', return_value=engine):
activity = ModelImport(
status_config=STATUS_CONFIG,
import_config=IMPORT_CONFIG,
minio_repository=MagicMock(),
mlflow_repository=MagicMock(),
mongo_repository=AsyncMock(),
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
)
activity = ModelImport(
status_config=STATUS_CONFIG,
import_config=IMPORT_CONFIG,
minio_repository=MagicMock(),
mlflow_repository=MagicMock(),
mongo_repository=AsyncMock(),
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
)
# `status_engine` is created lazily on first access; setting the private attribute directly
# pre-seeds it with the recording double instead of exercising the real `create_engine`.
activity._status_engine = engine
activity.emit_metric = AsyncMock()
activity.send_notification_async = AsyncMock()
return activity
@@ -983,8 +985,62 @@ async def test_the_engine_is_disposed_on_shutdown(import_activity: ModelImport,
@mark.asyncio
async def test_close_is_a_no_op_without_an_engine(import_activity: ModelImport, engine):
"""A partially-constructed instance (or a second close) must not blow up on a missing engine."""
del import_activity.status_engine
import_activity._status_engine = None
import_activity.close()
assert engine.disposed is False
# ------------------------------------------------------------- the lazy status engine
def test_no_engine_is_created_when_status_engine_is_never_accessed():
"""A worker pod that never runs an import activity never opens a pool against the BFF Postgres."""
with patch('laborious.activities.model_import.create_engine') as create_engine_mock:
activity = ModelImport(
status_config=STATUS_CONFIG,
import_config=IMPORT_CONFIG,
minio_repository=MagicMock(),
mlflow_repository=MagicMock(),
mongo_repository=AsyncMock(),
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
)
create_engine_mock.assert_not_called()
assert activity._status_engine is None
def test_first_access_creates_the_engine_and_later_accesses_reuse_it():
with patch(
'laborious.activities.model_import.create_engine', return_value=FakeStatusEngine()
) as create_engine_mock:
activity = ModelImport(
status_config=STATUS_CONFIG,
import_config=IMPORT_CONFIG,
minio_repository=MagicMock(),
mlflow_repository=MagicMock(),
mongo_repository=AsyncMock(),
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
)
first = activity.status_engine
second = activity.status_engine
create_engine_mock.assert_called_once()
assert first is second
@mark.asyncio
async def test_close_never_accessed_the_engine_does_not_create_one(import_activity: ModelImport):
"""`close()` must not force the lazy engine into existence just to dispose it."""
import_activity._status_engine = None
with patch('laborious.activities.model_import.create_engine') as create_engine_mock:
import_activity.close()
create_engine_mock.assert_not_called()

48
tests/laborious/utils/bundle/test_reader.py
View File

@@ -437,7 +437,7 @@ def test_gate_5_wraps_an_os_error_during_unpacking(work_dir, reader, monkeypatch
def _raise_os_error(*_args, **_kwargs):
raise OSError('disk full')
monkeypatch.setattr('laborious.utils.bundle.reader.shutil.copyfileobj', _raise_os_error)
monkeypatch.setattr(zipfile.ZipExtFile, 'read', _raise_os_error)
with pytest.raises(BundleGateError, match='could not be unpacked') as error:
_open(reader, work_dir, payload)
@@ -463,6 +463,52 @@ def test_extract_member_rechecks_containment_on_the_resolved_path(work_dir, read
assert error.value.gate == 5
def test_extract_member_rejects_a_stream_larger_than_its_declared_size(work_dir, reader, monkeypatch):
"""A decompressed stream longer than the entry's declared `file_size` is rejected mid-write,
instead of trusting the central-directory metadata gate 4 already checked in aggregate.
A single crafted DEFLATE chunk can make `ZipExtFile.read()` return far more bytes than
`file_size` declares before its own EOF/CRC bookkeeping catches up, so the stand-in here
forces exactly that response shape rather than fighting zipfile's own CRC validation."""
root = work_dir / 'root'
root.mkdir()
buffer = io.BytesIO()
with zipfile.ZipFile(buffer, 'w') as archive:
archive.writestr('payload.bin', b'x' * 10)
call_count = {'n': 0}
def _oversized_read(self, *_args, **_kwargs):
call_count['n'] += 1
return b'x' * 1000 if call_count['n'] == 1 else b''
monkeypatch.setattr(zipfile.ZipExtFile, 'read', _oversized_read)
with zipfile.ZipFile(io.BytesIO(buffer.getvalue())) as archive:
entry = archive.infolist()[0]
with pytest.raises(BundleGateError, match='more than its declared') as error:
reader._extract_member(archive, entry, root, root.resolve())
assert error.value.gate == 5
def test_extract_member_writes_a_well_formed_member_in_full(work_dir, reader):
"""A member whose decompressed size matches its declared `file_size` is extracted byte for
byte, exactly as before the byte-counting check was added."""
root = work_dir / 'root'
root.mkdir()
payload = b'y' * 100
buffer = io.BytesIO()
with zipfile.ZipFile(buffer, 'w') as archive:
archive.writestr('payload.bin', payload)
with zipfile.ZipFile(io.BytesIO(buffer.getvalue())) as archive:
entry = archive.infolist()[0]
reader._extract_member(archive, entry, root, root.resolve())
assert (root / 'payload.bin').read_bytes() == payload
# ------------------------------------------------------------------- gate 6

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

@@ -557,12 +557,15 @@ async def test_download_model_invalid_model_type(mlflow_repository):
'model_type', [('predict', 'prediction_model'), ('transform', 'data_model')]
)
async def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type, tmp_path):
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path))
result = await mlflow_repository.download_model('test_model', {}, model_type[0], 'pyfunc', True)
mlflow_repository.dowload_artifacts.assert_called_once_with('test_model', {}, model_type[1])
mlflow_repository.dowload_artifacts.assert_called_once_with(
'test_model', {}, model_type[1], run_id='run-123'
)
mlflow.pyfunc.load_model.assert_called_once_with(
mlflow_repository.dowload_artifacts.return_value
@@ -576,13 +579,16 @@ async def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type
@pytest.mark.asyncio
async def test_download_model_predict(mlflow_repository):
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
mlflow_repository.load_predict_model = AsyncMock(return_value=('model', 'local_dir'))
mlflow_repository.load_transform_model = AsyncMock()
result = await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', False)
mlflow_repository.load_predict_model.assert_called_once_with('test_model', {}, 'pyfunc')
mlflow_repository.load_predict_model.assert_called_once_with(
'test_model', {}, 'pyfunc', run_id='run-123'
)
mlflow_repository.load_transform_model.assert_not_called()
assert result == ('model', 'local_dir')
@@ -590,6 +596,7 @@ async def test_download_model_predict(mlflow_repository):
@pytest.mark.asyncio
async def test_download_model_transform(mlflow_repository):
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
mlflow_repository.load_predict_model = AsyncMock()
mlflow_repository.load_transform_model = AsyncMock(return_value=('model', 'local_dir'))
@@ -597,7 +604,9 @@ async def test_download_model_transform(mlflow_repository):
result = await mlflow_repository.download_model('test_model', {}, 'transform', 'pyfunc', False)
mlflow_repository.load_predict_model.assert_not_called()
mlflow_repository.load_transform_model.assert_called_once_with('test_model', {}, 'pyfunc')
mlflow_repository.load_transform_model.assert_called_once_with(
'test_model', {}, 'pyfunc', run_id='run-123'
)
assert result == ('model', 'local_dir')
@@ -616,10 +625,11 @@ async def test_download_model_prepares_the_code_dir_first(
calls.append('code')
return 'code-artifact-dir'
async def record_load(*_a):
async def record_load(*_a, **_k):
calls.append('load')
return 'model', None
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.download_artifact_code_dir = AsyncMock(side_effect=record_code_dir)
loader = AsyncMock(side_effect=record_load)
mlflow_repository.load_predict_model = loader
@@ -628,7 +638,7 @@ async def test_download_model_prepares_the_code_dir_first(
result = await mlflow_repository.download_model('test_model', {}, model_type, 'sklearn', False)
mlflow_repository.download_artifact_code_dir.assert_called_once_with(
'test_model', artifact_name, 'sklearn', {}
'test_model', artifact_name, 'sklearn', {}, run_id='run-123'
)
assert calls == ['code', 'load']
assert result == ('model', 'code-artifact-dir')
@@ -637,6 +647,7 @@ async def test_download_model_prepares_the_code_dir_first(
@pytest.mark.asyncio
async def test_download_model_wrapper_returns_the_code_artifact_dir(mlflow, mlflow_repository):
"""With a `code` directory downloaded, the wrapper branch returns it over its own download."""
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.download_artifact_code_dir = AsyncMock(return_value='code-artifact-dir')
mlflow_repository.dowload_artifacts = AsyncMock(return_value='wrapper-dir')
@@ -647,6 +658,42 @@ async def test_download_model_wrapper_returns_the_code_artifact_dir(mlflow, mlfl
assert artifact_dir == 'code-artifact-dir'
@pytest.mark.asyncio
async def test_download_model_resolves_run_id_once_non_wrapper(mlflow_repository):
"""The `run_id` `download_model` resolves once is the same one every downstream call uses."""
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
mlflow_repository.load_predict_model = AsyncMock(return_value=('model', None))
await mlflow_repository.download_model('test_model', {}, 'predict', 'sklearn', False)
mlflow_repository.get_model_run_id.assert_called_once_with(model_name='test_model', stage='Production')
mlflow_repository.download_artifact_code_dir.assert_called_once_with(
'test_model', 'prediction_model', 'sklearn', {}, run_id='run-123'
)
mlflow_repository.load_predict_model.assert_called_once_with(
'test_model', {}, 'sklearn', run_id='run-123'
)
@pytest.mark.asyncio
async def test_download_model_resolves_run_id_once_wrapper(mlflow, mlflow_repository, tmp_path):
"""Same single resolution, for the `load_wrapper=True` branch."""
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path))
await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)
mlflow_repository.get_model_run_id.assert_called_once_with(model_name='test_model', stage='Production')
mlflow_repository.download_artifact_code_dir.assert_called_once_with(
'test_model', 'prediction_model', 'pyfunc', {}, run_id='run-123'
)
mlflow_repository.dowload_artifacts.assert_called_once_with(
'test_model', {}, 'prediction_model', run_id='run-123'
)
def test_detect_and_parse_datetime_index_empty(mlflow_repository):
input_data = DataFrame()
response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
@@ -1965,6 +2012,7 @@ async def test_load_transform_model_flavor_failure_propagates(
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_wrapper_failure_propagates(joblib_mock, mlflow, mlflow_repository):
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
mlflow_repository.dowload_artifacts = AsyncMock()
loader_error = mlflow_lib.exceptions.MlflowException('wrapper boom')
@@ -1983,6 +2031,7 @@ async def test_download_model_wrapper_failure_propagates(joblib_mock, mlflow, ml
async def test_download_model_wrapper_failure_emits_error_metric(
joblib_mock, mlflow, mlflow_repository
):
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
mlflow_repository.dowload_artifacts = AsyncMock()
loader_error = mlflow_lib.exceptions.MlflowException('wrapper boom')
@@ -2503,13 +2552,14 @@ async def test_download_model_delegates_every_flavor_to_the_loaders(
mlflow_repository, model_type, loader
):
"""`download_model` owns no flavor branch: joblib delegates like the other two."""
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
delegate = AsyncMock(return_value=('model', 'local-dir'))
setattr(mlflow_repository, loader, delegate)
result = await mlflow_repository.download_model('test_model', {}, model_type, 'joblib')
delegate.assert_called_once_with('test_model', {}, 'joblib')
delegate.assert_called_once_with('test_model', {}, 'joblib', run_id='run-123')
assert result == ('model', 'local-dir')

91
tests/laborious/workflows/test_import_model.py
View File

@@ -12,7 +12,7 @@ workflow is: a sequence of activity invocations with policies attached.
import asyncio
from datetime import timedelta
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
from pytest import fixture, mark, raises
from sientia_do.temporal.policies import retry_policy
@@ -725,6 +725,95 @@ async def test_reporting_the_report_is_where_it_stops(import_workflow: ImportMod
assert result['run_id'] == RUN_ID
# --------------------------------------------------- cancellation is never swallowed
def cancelling_activities():
"""Patch the workflow module's `workflow` object so every activity call raises cancellation."""
patcher = patch('laborious.workflows.import_model.workflow')
workflow_mock = patcher.start()
workflow_mock.execute_activity_method = AsyncMock(side_effect=asyncio.CancelledError())
return patcher
BASE = {'metadata': {}, 'import_run_id': 77, 'workflow_id': WORKFLOW_ID}
@mark.asyncio
async def test_a_cancelled_hint_propagates_instead_of_being_reported(import_workflow: ImportModel):
patcher = cancelling_activities()
try:
with raises(asyncio.CancelledError):
await import_workflow._record_hint(
BASE, Activities.record_import_status, {'step': ImportStep.DOWNLOAD.value}
)
finally:
patcher.stop()
@mark.asyncio
async def test_a_cancelled_terminal_write_propagates_instead_of_being_reported(
import_workflow: ImportModel,
):
patcher = cancelling_activities()
try:
with raises(asyncio.CancelledError):
await import_workflow._record_terminal(
BASE,
succeeded=True,
step=ImportStep.MODEL_DOCUMENT.value,
gate=None,
source={},
cleanup_failed=False,
)
finally:
patcher.stop()
@mark.asyncio
async def test_a_cancelled_pipeline_failure_report_propagates(import_workflow: ImportModel):
patcher = cancelling_activities()
try:
with raises(asyncio.CancelledError):
await import_workflow._report_failure(
BASE, RuntimeError('boom'), step=ImportStep.REGISTRATION.value, gate=None
)
finally:
patcher.stop()
@mark.asyncio
async def test_a_cancelled_status_write_failure_report_propagates(import_workflow: ImportModel):
patcher = cancelling_activities()
try:
with raises(asyncio.CancelledError):
await import_workflow._report_status_write_failure(BASE, RuntimeError('db down'))
finally:
patcher.stop()
@mark.asyncio
async def test_a_cancelled_retention_check_propagates(import_workflow: ImportModel):
patcher = cancelling_activities()
try:
with raises(asyncio.CancelledError):
await import_workflow._ensure_retention(BASE, valid_input())
finally:
patcher.stop()
@mark.asyncio
async def test_a_cancelled_cleanup_propagates_instead_of_returning_not_cleaned(
import_workflow: ImportModel,
):
patcher = cancelling_activities()
try:
with raises(asyncio.CancelledError):
await import_workflow._cleanup(BASE)
finally:
patcher.stop()
# --------------------------------------------------------------- the failure report