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 model import workflow, introduce atomic model ID generation, and clarify MinIO expiry.

Fuses bundle download, decryption, MLflow experiment creation, and artifact upload into a single activity to prevent cross-worker issues. Implements an atomic counter for MongoDB model IDs, replacing the max(id) + 1 scan to avoid collisions during concurrent imports. Adds comprehensive documentation for MinIO object expiry mechanisms and refines lifecycle rule application. Improves error handling with clearer responsibility classification and centralizes MLflow repository injection.
This commit is contained in:
Vitor Pimentel dos Santos
2026-09-08 17:06:25 -03:00
parent e4fc12e529
commit 192834bde3
34 changed files with 1485 additions and 531 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -9,6 +9,7 @@ Everything is asserted through the recorded activity calls, because that is exac
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
@@ -27,6 +28,7 @@ from laborious.utils.models.import_errors import (
ImportGateRejectedError,
ImportInputError,
ImportLogRowNotClaimableError,
ImportStatusRejectedError,
)
from laborious.workflows.import_model import (
REJECTED_INPUT_FIELDS,
@@ -60,20 +62,16 @@ STATUS_ACTIVITIES = (
CLAIM_AND_TERMINAL = (Activities.claim_import_status, Activities.record_import_terminal_status)
HINTS = (Activities.record_import_status, Activities.record_import_names)
# What a correct import does, in order.
# What a correct import does, in order. Download, decryption/extraction, experiment creation and
# artifact upload run inside one fused activity (bug 8), so only one hint precedes it, and there is
# no longer a separate activity call for each of those four phases.
EXPECTED_ORDER = [
Activities.claim_import_status,
Activities.record_import_status,
Activities.ensure_import_bundle_retention,
Activities.download_import_bundle,
Activities.record_import_status,
Activities.open_import_bundle,
Activities.prepare_and_upload_import_bundle,
Activities.record_import_names,
Activities.record_import_status,
Activities.create_import_experiment,
Activities.record_import_status,
Activities.upload_import_artifacts,
Activities.record_import_status,
Activities.register_import_model_version,
Activities.record_import_status,
Activities.write_import_model_document,
@@ -81,9 +79,9 @@ EXPECTED_ORDER = [
Activities.record_import_terminal_status,
]
# The two provisioning activities left as their own Temporal activities after the fusion (bug 8
# folded experiment creation and artifact upload into `prepare_and_upload_import_bundle`).
PROVISIONING = (
Activities.create_import_experiment,
Activities.upload_import_artifacts,
Activities.register_import_model_version,
Activities.write_import_model_document,
)
@@ -95,19 +93,16 @@ HAPPY_RESULTS = {
'file_name': None,
'request_digest': None,
},
Activities.download_import_bundle: {
'path': '/work/bundle.sientia',
'size': 4096,
'digest': DIGEST,
},
Activities.open_import_bundle: {
Activities.prepare_and_upload_import_bundle: {
'extracted_dir': '/work/extract',
'metadata': BUNDLE_METADATA,
'parameters': BUNDLE_PARAMETERS,
'digest': DIGEST,
'experiment_id': '17',
'experiment_name': 'factory_experiment',
'run_id': RUN_ID,
'run_name': 'import-factory_model-v2',
},
Activities.create_import_experiment: {'experiment_id': '17'},
Activities.upload_import_artifacts: {'run_id': RUN_ID, 'run_name': 'import-factory_model-v2'},
Activities.register_import_model_version: {'model_name': 'factory_model', 'version': '4'},
Activities.write_import_model_document: {'written': True, 'id': '8'},
Activities.cleanup_import_files: {'cleaned': True, 'existed': True},
@@ -210,21 +205,60 @@ async def test_the_claim_is_the_first_thing_that_happens(
@mark.asyncio
async def test_the_names_are_recorded_before_mlflow_is_touched(
async def test_the_names_are_recorded_after_the_bundle_is_prepared_and_before_registration(
import_workflow: ImportModel, workflow_env
):
"""So a person watching the import list sees which model is being imported, while it runs."""
"""The names are only known once the fused activity returns (bug 8); they still land before
the two provisioning steps that follow it — the registered version and the model document."""
await import_workflow.run(valid_input())
order = workflow_env.recorder.order
assert order.index(Activities.record_import_names) < order.index(
Activities.create_import_experiment
assert (
order.index(Activities.prepare_and_upload_import_bundle)
< order.index(Activities.record_import_names)
< order.index(Activities.register_import_model_version)
)
names = workflow_env.recorder.payload_of(Activities.record_import_names)
assert names['experiment_name'] == 'factory_experiment'
assert names['run_name'] == 'import-factory_model-v2'
@mark.asyncio
async def test_retention_and_the_bundle_activity_are_both_scheduled_before_either_finishes(
import_workflow: ImportModel,
):
"""Achado #9: `_ensure_retention` and the fused prepare-and-upload activity run under
`asyncio.gather`, so neither has to complete — or even be scheduled — before the other. Call
order between them isn't asserted (gather makes no such guarantee); only that the second
activity gets recorded while the first is still pending confirms neither blocks the other."""
retention_started = asyncio.Event()
release_retention = asyncio.Event()
class BlockingRecorder(Recorder):
async def __call__(self, method, payload, **kwargs):
self.calls.append(SimpleNamespace(method=method, payload=payload, kwargs=kwargs))
if method is Activities.ensure_import_bundle_retention:
retention_started.set()
await release_retention.wait()
if method in self.failures:
raise self.failures[method]
return self.results.get(method, {})
recorder = BlockingRecorder()
patcher, _ = run_with(recorder)
task = asyncio.ensure_future(import_workflow.run(valid_input()))
try:
await asyncio.wait_for(retention_started.wait(), timeout=1)
# Retention is still awaiting `release_retention`; the bundle activity must already have
# been scheduled and recorded rather than waiting for retention to finish first.
await asyncio.sleep(0)
assert Activities.prepare_and_upload_import_bundle in recorder.order
finally:
release_retention.set()
await task
patcher.stop()
@mark.asyncio
async def test_the_model_is_named_from_the_bundles_model_name(
import_workflow: ImportModel, workflow_env
@@ -247,11 +281,13 @@ async def test_the_model_is_named_from_the_bundles_model_name(
async def test_the_run_is_created_with_exactly_the_recorded_name(
import_workflow: ImportModel, workflow_env
):
"""The fused activity computes the run name and creates the run with it; the workflow then
records that same value, so the two cannot diverge."""
await import_workflow.run(valid_input())
recorded = workflow_env.recorder.payload_of(Activities.record_import_names)['run_name']
used = workflow_env.recorder.payload_of(Activities.upload_import_artifacts)['run_name']
assert used == recorded
created = workflow_env.recorder.results[Activities.prepare_and_upload_import_bundle]['run_name']
assert recorded == created
@mark.asyncio
@@ -296,7 +332,7 @@ async def test_the_download_is_handed_the_digest_the_request_carried(
):
await import_workflow.run(valid_input())
payload = workflow_env.recorder.payload_of(Activities.download_import_bundle)
payload = workflow_env.recorder.payload_of(Activities.prepare_and_upload_import_bundle)
assert payload['expected_digest'] == DIGEST
assert payload['object_key'] == OBJECT_KEY
@@ -308,10 +344,9 @@ async def test_the_download_is_handed_the_digest_the_request_carried(
@mark.parametrize(
('failing', 'step'),
[
(Activities.download_import_bundle, ImportStep.DOWNLOAD),
(Activities.open_import_bundle, ImportStep.DECRYPTION),
(Activities.create_import_experiment, ImportStep.EXPERIMENT_CREATION),
(Activities.upload_import_artifacts, ImportStep.ARTIFACT_UPLOAD),
# An unclassified failure from the fused activity (download/decryption/experiment/upload)
# falls back to its first phase, since that is the step the workflow was attempting.
(Activities.prepare_and_upload_import_bundle, ImportStep.DOWNLOAD),
(Activities.register_import_model_version, ImportStep.REGISTRATION),
(Activities.write_import_model_document, ImportStep.MODEL_DOCUMENT),
],
@@ -348,7 +383,7 @@ async def test_a_failing_step_is_recorded_as_that_step_and_stops_the_import(
async def test_nothing_is_provisioned_before_the_gates_pass(import_workflow: ImportModel):
recorder = Recorder(
failures={
Activities.open_import_bundle: ImportGateRejectedError(
Activities.prepare_and_upload_import_bundle: ImportGateRejectedError(
'rejected', step=ImportStep.ARCHIVE_INSPECTION, gate=4
)
}
@@ -373,7 +408,7 @@ async def test_a_gate_rejection_is_recorded_with_the_gate_that_rejected_it(
step the workflow was attempting."""
recorder = Recorder(
failures={
Activities.open_import_bundle: ImportGateRejectedError(
Activities.prepare_and_upload_import_bundle: ImportGateRejectedError(
'rejected', step=ImportStep.CONTENT_POLICY, gate=7
)
}
@@ -395,7 +430,7 @@ async def test_a_gate_rejection_is_recorded_with_the_gate_that_rejected_it(
async def test_an_ambiguous_decryption_failure_keeps_its_step(import_workflow: ImportModel):
recorder = Recorder(
failures={
Activities.open_import_bundle: ImportGateRejectedError(
Activities.prepare_and_upload_import_bundle: ImportGateRejectedError(
'the bundle could not be opened', step=ImportStep.DECRYPTION, gate=3
)
}
@@ -433,7 +468,7 @@ async def test_an_unknown_field_is_rejected_rather_than_ignored(
# After the claim, so the refusal lands on the record instead of only in the logs.
assert recorder.order[0] is Activities.claim_import_status
assert Activities.download_import_bundle not in recorder.order
assert Activities.prepare_and_upload_import_bundle not in recorder.order
terminal = recorder.payload_of(Activities.record_import_terminal_status)
assert terminal['step'] == ImportStep.RECEIVED.value
assert terminal['succeeded'] is False
@@ -457,7 +492,7 @@ async def test_a_missing_or_malformed_digest_fails_before_anything_is_downloaded
finally:
patcher.stop()
assert Activities.download_import_bundle not in recorder.order
assert Activities.prepare_and_upload_import_bundle not in recorder.order
assert (
recorder.payload_of(Activities.record_import_terminal_status)['step']
== ImportStep.RECEIVED.value
@@ -479,7 +514,7 @@ async def test_a_request_missing_a_required_field_fails_at_received(
finally:
patcher.stop()
assert Activities.download_import_bundle not in recorder.order
assert Activities.prepare_and_upload_import_bundle not in recorder.order
@mark.asyncio
@@ -503,7 +538,7 @@ async def test_a_record_naming_another_object_is_a_received_failure(
finally:
patcher.stop()
assert Activities.download_import_bundle not in recorder.order
assert Activities.prepare_and_upload_import_bundle not in recorder.order
@mark.asyncio
@@ -527,7 +562,7 @@ async def test_a_record_naming_another_digest_is_a_received_failure(
finally:
patcher.stop()
assert Activities.download_import_bundle not in recorder.order
assert Activities.prepare_and_upload_import_bundle not in recorder.order
@mark.asyncio
@@ -595,7 +630,7 @@ async def test_a_lost_names_write_does_not_fail_a_working_import(import_workflow
patcher.stop()
assert result['version'] == '4'
assert Activities.upload_import_artifacts in recorder.order
assert Activities.prepare_and_upload_import_bundle in recorder.order
@mark.asyncio
@@ -712,7 +747,7 @@ async def test_the_failure_that_ends_an_import_is_reported_through_the_channel(
nothing at all — only the SDK's own traceback for the attempt, with no id and no step."""
recorder = Recorder(
failures={
Activities.create_import_experiment: activity_failure(
Activities.prepare_and_upload_import_bundle: activity_failure(
'API request to http://localhost:5080/api/2.0/mlflow/experiments/get-by-name '
'failed with exception Connection refused',
'MlflowException',
@@ -728,7 +763,10 @@ async def test_the_failure_that_ends_an_import_is_reported_through_the_channel(
patcher.stop()
reported = recorder.payload_of(Activities.report_import_pipeline_failure)
assert reported['step'] == ImportStep.EXPERIMENT_CREATION.value
# An unclassified failure from the fused activity falls back to its first phase: the workflow
# has no more precise location for it, since `current_step` only advances past `DOWNLOAD` once
# the fused activity returns.
assert reported['step'] == ImportStep.DOWNLOAD.value
assert reported['error_type'] == 'MlflowException'
assert reported['detail'].splitlines()[0].startswith('MlflowException: API request to')
assert reported['import_run_id'] == 77
@@ -741,7 +779,9 @@ async def test_the_report_comes_before_the_cleanup_and_the_terminal_write(
):
"""So the log reads in the order things happened, and the row's step matches the line's."""
recorder = Recorder(
failures={Activities.upload_import_artifacts: activity_failure('bucket gone', 'S3Error')}
failures={
Activities.prepare_and_upload_import_bundle: activity_failure('bucket gone', 'S3Error')
}
)
patcher, _ = run_with(recorder)
payload = valid_input()
@@ -759,7 +799,7 @@ async def test_the_report_comes_before_the_cleanup_and_the_terminal_write(
)
reported = recorder.payload_of(Activities.report_import_pipeline_failure)
terminal = recorder.payload_of(Activities.record_import_terminal_status)
assert reported['step'] == terminal['step'] == ImportStep.ARTIFACT_UPLOAD.value
assert reported['step'] == terminal['step'] == ImportStep.DOWNLOAD.value
@mark.asyncio
@@ -768,7 +808,7 @@ async def test_the_gate_that_rejected_the_bundle_travels_into_the_report(
):
recorder = Recorder(
failures={
Activities.open_import_bundle: ImportGateRejectedError(
Activities.prepare_and_upload_import_bundle: ImportGateRejectedError(
'gate 3 rejected the bundle', step=ImportStep.DECRYPTION, gate=3
)
}
@@ -807,6 +847,33 @@ async def test_a_claim_that_matches_nothing_is_not_reported_twice(import_workflo
assert Activities.report_import_pipeline_failure not in recorder.order
@mark.asyncio
async def test_a_claim_write_that_is_refused_is_not_reported_twice_by_the_workflow(
import_workflow: ImportModel,
):
"""Bug 7: the claim's own write failure is now reported *inside* `claim_import_status` itself,
before it re-raises (`code=STATUS_NOT_RECORDED`). The claim still runs outside the workflow's
try/finally — there is no row this workflow owns yet — so the workflow's own pipeline-failure
channel must not report it a second time."""
recorder = Recorder(
failures={
Activities.claim_import_status: ImportStatusRejectedError(
'the import log database refused the status statement'
)
}
)
patcher, _ = run_with(recorder)
payload = valid_input()
try:
with raises(ImportStatusRejectedError):
await import_workflow.run(payload)
finally:
patcher.stop()
assert recorder.order == [Activities.claim_import_status]
assert Activities.report_import_pipeline_failure not in recorder.order
@mark.asyncio
async def test_a_report_that_cannot_be_made_does_not_mask_the_failure(
import_workflow: ImportModel,
@@ -842,7 +909,9 @@ async def test_the_report_gives_up_inside_the_short_bound(import_workflow: Impor
"""A failure already happened; spending 151 seconds telling the log about it would be worse."""
recorder = Recorder(
failures={
Activities.create_import_experiment: activity_failure('refused', 'MlflowException')
Activities.prepare_and_upload_import_bundle: activity_failure(
'refused', 'MlflowException'
)
}
)
patcher, _ = run_with(recorder)
@@ -904,7 +973,10 @@ async def test_no_provisioning_call_references_a_retrying_policy(
await import_workflow.run(valid_input())
for call in workflow_env.recorder.calls:
if call.method in PROVISIONING or call.method is Activities.open_import_bundle:
if (
call.method in PROVISIONING
or call.method is Activities.prepare_and_upload_import_bundle
):
assert call.kwargs['retry_policy'] not in (
IMPORT_STATUS_RETRY_POLICY,
IMPORT_STATUS_HINT_RETRY_POLICY,