gitea_admin/sientia-dataops-laborious-legacy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
72a55c09002b28caf400f464cda25de923209673
sientia-dataops-laborious-l…/e2e/test_import_model_retry_recovery.py
vitor-aignosi e95e65b72b QTZPOC-20 
QTZPOC-20: Implement explicit joblib model flavor and enhance model import workflow with extensive E2E tests.

This commit introduces a first-class joblib model flavor for direct loading and saving of Python pickle artifacts, including custom code within the artifact's code/ directory, bypassing MLflow's flavor loaders. The model import workflow's status database configuration is now streamlined to default to the POSTGRES_* server, with explicit IMPORT_STATUS_DB_* overrides.

Comprehensive end-to-end tests have been added for the model import workflow, covering success paths, gate rejections, digest enforcement, provisioning failures, status record integrity, cleanup, and retry/recovery scenarios using a real-clock environment. New E2
2026-08-26 15:45:33 -03:00

515 lines
23 KiB
Python
Raw Blame History

"""What happens when the *record* cannot be written — the cases that gate this feature's release.
Everything else in this suite asks "did the import do the right thing". This file asks the harder
question: the import did the right thing, and then could not say so. Design D20 answers it in two
halves, and both are release-critical.
**A database that comes back.** The terminal write retries, so a status database restarted during a
deploy still records the correct verdict. Without this, every rolling restart would leave imports
looking unfinished forever.
**A database that does not come back.** The model is fully provisioned, and the workflow must not
claim otherwise. It **completes**, and the verdict it could not write goes out through the failure
channel — log, metric, notification — under `IMPORT_STATUS_NOT_RECORDED`. An import that worked must
never be reported as an import that failed just because the audit row was unreachable: somebody
would re-import a model that is already there.
That last point was worth discovering the hard way. The first version of this file asserted the run
*failed*; it does not, and it should not. The spec says so in as many words — "the run is not
reported as failed" — and the implementation agrees with the spec.
And the case that is the evidence for a migration in another repository: **a table without `V12`**.
The claim succeeds (`V11` already gave the table `run_type`, and `PENDING` and the running status are
pre-`V12` values), the model is fully provisioned, and only the *terminal* value is new — so the last
write of the import is the first thing that fails. `ImportStatusRejectedError` is non-retryable, so
it lands on the first attempt with no wait, and the row is left running forever. That is what
"apply `V12` before turning `IMPORT_MODEL_ENABLED` on" buys, as a test rather than a warning in a
document.
Everything here runs on a **real clock** (design E9): under time skipping the backoff sleeps are
virtual, so a restarted container would never get wall-clock time to come back and the recovery test
would race by construction. The retry intervals are shortened by `import_status_retry_shortening`;
the production values stay asserted by unit test 12.1 of `import-model-workflow`, because what these
tests are about is the behaviour, not the arithmetic.
"""
from __future__ import annotations
import asyncio
import json
import time
from typing import Any
import mlflow
import pytest
from temporalio.client import WorkflowFailureError
from e2e.import_harness import failure_text, failure_types
pytestmark = [pytest.mark.import_e2e, pytest.mark.slow, pytest.mark.asyncio]
FACTORY_PASSWORD = 'Factory1Model!Test'
RUNNING_STATUS = 'ORCHESTRATOR_WAITING_PROC'
NOT_RECORDED_CODE = 'IMPORT_STATUS_NOT_RECORDED'
STATUS_WRITE_MARKER = 'step=status_write'
PROVISIONING_STEPS = ('experiment_creation', 'artifact_upload', 'registration', 'model_document')
def _as_dict(value: Any) -> dict[str, Any]:
return json.loads(value) if isinstance(value, str) else (value or {})
def _client(server: dict[str, Any]) -> Any:
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
def _channel_text(mock_logger: Any) -> str:
"""Everything the failure channel said, as one string.
The channel is a log line, a metric and a notification, and the log line is the only sink that
carries the whole classification — code, marker and the verdict it could not record. Reading the
logger is therefore reading the channel, and it is synchronous, which the notification path is
not.
"""
return '\n'.join(str(call) for call in mock_logger.mock_calls)
def _assert_fully_provisioned(
*, server: dict[str, Any], names: dict[str, str], collection: Any
) -> str:
"""Everything the import was supposed to create exists. Returns the run id."""
client = _client(server)
experiment = client.get_experiment_by_name(names['experiment_name'])
assert experiment is not None, 'no experiment: the import never got as far as provisioning'
runs = client.search_runs([experiment.experiment_id])
assert len(runs) == 1, f'expected one run, found {len(runs)}'
assert {a.path for a in client.list_artifacts(runs[0].info.run_id)} >= {
'prediction_model',
'data_model',
}
versions = list(client.search_model_versions(f"name='{names['model_name']}'"))
assert len(versions) == 1, f'expected one registered version, found {len(versions)}'
assert collection.find_one({'name': names['model_name']}) is not None, 'no model document'
return runs[0].info.run_id
def _assert_no_provisioning_step_is_blamed(text_of_failure: str) -> None:
"""The import's own steps all succeeded, so none of them may appear as the cause."""
for step in PROVISIONING_STEPS:
assert step not in text_of_failure, f'{step!r} is blamed, but that step succeeded'
async def _wait_for_claim(read_row: Any, import_run_id: int, timeout: float = 30.0) -> None:
"""Block until the workflow has taken ownership of the row.
The earliest deterministic point after which every status write belongs to this import. Used to
place an outage that must cover the whole rest of the run.
"""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
row = read_row(import_run_id)
if row is not None and row['status'] == RUNNING_STATUS:
return
await asyncio.sleep(0.1)
raise AssertionError('the row was never claimed, so the outage could not be placed')
async def _wait_for_provisioning(
*, server: dict[str, Any], names: dict[str, str], collection: Any, timeout: float = 90.0
) -> None:
"""Block until the model document exists — the observable end of provisioning.
Polling the two destinations rather than patching an activity: "the document exists" *is* the
end of provisioning, so there is nothing to hook and nothing to fake. This is what lets the
outage below be placed exactly on the terminal write.
"""
client = _client(server)
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
document = collection.find_one({'name': names['model_name']})
if document is not None and list(
client.search_model_versions(f"name='{names['model_name']}'")
):
return
await asyncio.sleep(0.2)
raise AssertionError('provisioning never finished, so the outage could not be placed')
class TestTheDatabaseComesBack:
"""Stopped and restarted inside the retry envelope: the verdict still lands."""
async def test_a_restart_inside_the_envelope_still_records_the_verdict(
self,
import_status_table: Any,
import_models_collection: Any,
import_names: dict[str, str],
import_mlflow_server: dict[str, Any],
import_bundle_factory: Any,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute_real_clock: Any,
import_envelope: Any,
import_read_status_row: Any,
import_bff_stopped: Any,
) -> None:
"""The outage lands on the terminal write, and the server is back before the retries run out.
The stop is placed after provisioning rather than at the start, so what is being tested is
the *terminal* write's retry envelope specifically. A stop at the start would also exercise
the claim's envelope, which is a different (and already covered) code path.
"""
object_key, digest = import_upload_bundle(import_bundle_factory())
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
handle = await import_execute_real_clock(
import_build_input(
import_run_id=import_run_id,
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
await _wait_for_provisioning(
server=import_mlflow_server,
names=import_names,
collection=import_models_collection,
)
# Down for two seconds — several retry attempts, well inside the shortened envelope.
async with import_bff_stopped():
await asyncio.sleep(2)
result = await handle.result()
assert result['model_name'] == import_names['model_name']
run_id = _assert_fully_provisioned(
server=import_mlflow_server,
names=import_names,
collection=import_models_collection,
)
assert result['run_id'] == run_id
row = import_read_status_row(import_run_id)
assert row is not None
assert row['status'] == 'IMPORT_SUCCESS', (
f'the verdict did not land after the restart: the row says {row["status"]!r}'
)
assert row['error_message'] is None
assert _as_dict(row['orchestrator_response_data'])['import']['workflow_id'] == handle.id
async def test_a_broken_progress_update_does_not_delay_the_verdict(
self,
import_status_table: Any,
import_models_collection: Any,
import_names: dict[str, str],
import_mlflow_server: dict[str, Any],
import_bundle_factory: Any,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute_real_clock: Any,
import_envelope: Any,
import_read_status_row: Any,
import_bff_stopped: Any,
) -> None:
"""A mid-flight outage that only the progress hints see. The import does not care.
Progress writes are hints nobody reads — the terminal write rewrites status, code, reason and
step from one classification — so they get their own short policy. If they shared the terminal
write's envelope, a database the import does not need would stall a working import for
minutes at every stage update.
The outage is placed early, so it lands on hints rather than on the terminal write, and the
assertion is that the import still finished with the correct verdict.
"""
object_key, digest = import_upload_bundle(import_bundle_factory())
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
handle = await import_execute_real_clock(
import_build_input(
import_run_id=import_run_id,
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
# Early and brief: the claim has landed, the gates are running, the terminal write is far off.
await asyncio.sleep(1.0)
async with import_bff_stopped():
await asyncio.sleep(1.5)
started = time.monotonic()
result = await handle.result()
elapsed = time.monotonic() - started
assert result['model_name'] == import_names['model_name']
row = import_read_status_row(import_run_id)
assert row is not None
assert row['status'] == 'IMPORT_SUCCESS'
assert elapsed < 60, (
f'the import took {elapsed:.1f}s after the outage ended; a lost hint should cost '
f'its own short policy and nothing more'
)
class TestTheDatabaseStaysDown:
"""Past the envelope. The model is real, the record is not, and the run says which is which."""
@pytest.fixture
def import_status_retry_shortening(self) -> dict[str, Any]:
"""Shorter still, because these tests wait for the envelope to *expire*.
Six half-second attempts inside an eight-second envelope: long enough that retries really
happen, short enough that giving up is something a test suite can wait for.
"""
return {
'initial_interval': 0.5,
'backoff_coefficient': 1.0,
'maximum_attempts': 6,
'hint_maximum_attempts': 2,
'envelope_seconds': 8,
'write_timeout_seconds': 3,
}
async def test_a_successful_import_is_not_reported_as_failed(
self,
import_status_table: Any,
import_models_collection: Any,
import_names: dict[str, str],
import_mlflow_server: dict[str, Any],
import_bundle_factory: Any,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute_real_clock: Any,
import_envelope: Any,
import_read_status_row: Any,
import_bff_stopped: Any,
mock_logger: Any,
) -> None:
"""The most important assertion in this file.
The model is registered, the document is written, the artifacts are stored — and the row
cannot be updated. The workflow **completes**: reporting failure here would tell somebody to
re-import a model that is already on the platform. The verdict goes out through the failure
channel instead, under the not-recorded code, carrying what it could not write.
The outage is placed as soon as the row has been **claimed**, not once provisioning has
finished. Waiting for provisioning looks more precise and is in fact a race: the terminal
write follows the model document by milliseconds, while `docker stop` takes seconds, so the
write lands before the database is actually gone and the test proves nothing. Stopping right
after the claim covers everything that follows — the progress hints (which are allowed to
fail) and the terminal write (which is the subject) — with no timing to get right.
"""
object_key, digest = import_upload_bundle(import_bundle_factory())
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
handle = await import_execute_real_clock(
import_build_input(
import_run_id=import_run_id,
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
await _wait_for_claim(import_read_status_row, import_run_id)
async with import_bff_stopped():
result = await handle.result()
# The run completed. This is the assertion, not an incidental detail.
assert result['model_name'] == import_names['model_name']
assert result['import_run_id'] == import_run_id
_assert_fully_provisioned(
server=import_mlflow_server,
names=import_names,
collection=import_models_collection,
)
channel = _channel_text(mock_logger)
assert NOT_RECORDED_CODE in channel, (
'the unrecordable verdict never reached the failure channel'
)
assert STATUS_WRITE_MARKER in channel, (
'the channel does not mark this as a status-write failure'
)
assert 'verdict_not_recorded=' in channel, (
'the channel reports the failure but not the verdict it could not write'
)
_assert_no_provisioning_step_is_blamed(
'\n'.join(
line for line in channel.splitlines() if 'model import failed' in line
).replace('step=model_document', '')
)
row = import_read_status_row(import_run_id)
assert row is not None
assert row['status'] == RUNNING_STATUS, (
f'the row should still be running, not {row["status"]!r}'
)
async def test_a_failed_import_keeps_its_original_cause(
self,
monkeypatch: pytest.MonkeyPatch,
import_status_table_pre_v12: Any,
import_models_collection: Any,
import_names: dict[str, str],
import_bundle_factory: Any,
import_activities: Any,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute_real_clock: Any,
import_envelope: Any,
import_read_status_row: Any,
mock_logger: Any,
) -> None:
"""An import that already failed, whose failure then cannot be written down either.
Two failures, and only one of them answers "why did this import not work". The status write
is the second, and it must not displace the first: somebody told "the status could not be
recorded" would go and look at the database, when what actually happened is that their model's
files could not be stored.
**Why a pre-`V12` table instead of stopping the database.** The outage has to cover the
terminal write and nothing earlier, and an import that fails at a gate reaches its terminal
write within a second of starting — there is no window to aim at, and the first version of
this test proved only that a stopped database *after* a completed write changes nothing.
A pre-`V12` table has no timing at all: `IMPORT_ERROR` is a value it has never heard of, so
the terminal write of a *failed* import is refused deterministically while every write before
it succeeds. The failure being tested is the same one either way — the record cannot take the
verdict — and this version cannot race.
"""
client = import_activities.import_mlflow_repository.client
def refuse(*_args: Any, **_kwargs: Any) -> None:
raise OSError('artifact store is unreachable')
monkeypatch.setattr(client, 'log_artifacts', refuse)
object_key, digest = import_upload_bundle(import_bundle_factory())
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
handle = await import_execute_real_clock(
import_build_input(
import_run_id=import_run_id,
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
with pytest.raises(WorkflowFailureError) as failure:
await handle.result()
# The reported cause is the step that actually failed.
text_of_failure = failure_text(failure.value)
assert 'artifact store is unreachable' in text_of_failure, (
f'the reported cause lost the original failure: {text_of_failure}'
)
assert 'ImportStatusRejectedError' not in failure_types(failure.value), (
f'the status write displaced the original cause: {failure_types(failure.value)}'
)
assert 'the import log database refused' not in text_of_failure, (
'the run is reported as a status-write failure rather than an artifact-upload failure'
)
# The status write still reported *itself*, through the channel, as the secondary failure.
channel = _channel_text(mock_logger)
assert NOT_RECORDED_CODE in channel, 'the unrecordable verdict was never reported'
assert STATUS_WRITE_MARKER in channel
# And the row is left running, because the verdict never landed.
row = import_read_status_row(import_run_id)
assert row is not None
assert row['status'] == RUNNING_STATUS
class TestMissingV12:
"""A table without `V12`. The evidence for why the migration gates enabling the worker."""
async def test_the_terminal_write_is_refused_on_its_first_attempt(
self,
import_status_table_pre_v12: Any,
import_models_collection: Any,
import_names: dict[str, str],
import_mlflow_server: dict[str, Any],
import_bundle_factory: Any,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute_real_clock: Any,
import_envelope: Any,
import_read_status_row: Any,
import_bff_engine: Any,
mock_logger: Any,
) -> None:
"""Provisioned, then unrecordable — and unrecordable *immediately*.
The shape matters as much as the outcome. `ImportStatusRejectedError` is non-retryable: a
`CHECK` violation is not going to start passing because it was tried again, and waiting the
full 151-second envelope would turn a misconfiguration into a two-and-a-half-minute stall on
every single import. The elapsed time is what proves it — the shortened policy still allows
24 attempts, so a retried failure could not finish this fast.
"""
from sqlalchemy import text as sql_text
object_key, digest = import_upload_bundle(import_bundle_factory())
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
started = time.monotonic()
handle = await import_execute_real_clock(
import_build_input(
import_run_id=import_run_id,
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
result = await handle.result()
elapsed = time.monotonic() - started
# The run completed: an import that provisioned a model is not a failed import, even when
# the last thing it wanted to do was impossible.
assert result['model_name'] == import_names['model_name']
# The claim and the names write both landed — the pre-V12 table was usable right up to the
# last statement, which is exactly what makes this failure mode easy to miss.
row = import_read_status_row(import_run_id)
assert row is not None
assert row['status'] == RUNNING_STATUS, (
f'the row is {row["status"]!r}; a pre-V12 table cannot hold a terminal import status, '
f'so the row is left looking as though the import never finished'
)
assert row['experiment_name'] == import_names['experiment_name']
assert row['run_name'] == f'import-{import_names["model_name"]}-v2'
# The model is fully provisioned. This is the whole problem: the work is done and invisible.
_assert_fully_provisioned(
server=import_mlflow_server,
names=import_names,
collection=import_models_collection,
)
# The verdict reached the failure channel under the not-recorded code.
channel = _channel_text(mock_logger)
assert NOT_RECORDED_CODE in channel
assert STATUS_WRITE_MARKER in channel
assert 'verdict_not_recorded=' in channel
# Non-retryable: refused on the first attempt, with no retry wait.
assert elapsed < 60, (
f'the import took {elapsed:.1f}s; a non-retryable rejection must not be retried'
)
# And the constraint that refused it is the one V12 adds — asserted directly, so the test
# says *why* the write was impossible rather than only that it was.
with import_bff_engine.begin() as connection:
statuses = connection.execute(
sql_text(
'SELECT pg_get_constraintdef(oid) FROM pg_constraint '
"WHERE conname = 'chk_experiment_run_status'"
)
).scalar_one()
assert 'IMPORT_SUCCESS' not in statuses, (
'the pre-V12 fixture already allows the import statuses, so this proved nothing'
)
Reference in New Issue View Git Blame Copy Permalink