vitor-aignosi
4c7cae502f
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.
451 lines
20 KiB
Python
451 lines
20 KiB
Python
"""The `import_model` workflow: a `.sientia` object becomes a registered model, or a durable ERROR.
|
|
|
|
Shape of this workflow, and the reasons it has that shape:
|
|
|
|
- **The claim comes first.** The record already exists — the frontend inserted it through the BFF —
|
|
so the workflow's first act is to take ownership of it. Everything after that failure is recorded
|
|
*on* the record; a failed claim is recorded nowhere, because the workflow owns no row.
|
|
- **Fail-fast, and nothing else.** No rollback, no compensation, no resume, no reconciliation, no
|
|
`continue_as_new`. A failed import stays failed; trying again is a new import with a new upload.
|
|
- **Only the status writes retry**, and only because `UPDATE ... SET <fixed values> WHERE id` is
|
|
idempotent. Every provisioning call keeps `maximum_attempts=1`, because a repeated registry call or
|
|
Mongo insert can write twice.
|
|
- **One `finally`, two actions**: cleanup, then the single terminal write that classifies the
|
|
outcome once. No `except` block writes a verdict of its own.
|
|
- **The workflow holds ids, never contents.** The decrypted bundle never crosses an activity
|
|
boundary, and the password reaches this module only as an envelope it cannot open.
|
|
"""
|
|
|
|
import asyncio
|
|
from datetime import timedelta
|
|
|
|
from temporalio import workflow
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Any
|
|
|
|
from sientia_do.temporal.policies import retry_policy
|
|
from temporalio.exceptions import is_cancelled_exception
|
|
|
|
from laborious.activities.activities import Activities
|
|
from laborious.activities.model_import import (
|
|
IMPORT_STATUS_HINT_RETRY_POLICY,
|
|
IMPORT_STATUS_RETRY_POLICY,
|
|
)
|
|
from laborious.utils.bundle.steps import ImportStep
|
|
from laborious.utils.models.import_errors import (
|
|
ImportInputError,
|
|
describe_failure_chain,
|
|
failure_location,
|
|
)
|
|
|
|
# `start_to_close_timeout` is per attempt. The terminal write additionally gets a
|
|
# `schedule_to_close_timeout` larger than its retry envelope (10 attempts, ≈151 s of backoff, ≈451 s
|
|
# if every attempt also burns its timeout) — set it smaller and Temporal cuts the retries short,
|
|
# which is the one configuration mistake that would silently undo design D20.
|
|
STATUS_WRITE_TIMEOUT = timedelta(seconds=30)
|
|
STATUS_WRITE_ENVELOPE = timedelta(seconds=900)
|
|
# Download, decryption/extraction, experiment creation and artifact upload now run inside one
|
|
# activity attempt (bug 8): the budget is the sum of what the four separate activities used to get,
|
|
# since a retry of the fused activity redoes all four phases.
|
|
PREPARE_AND_UPLOAD_TIMEOUT = (
|
|
timedelta(seconds=1800)
|
|
+ timedelta(seconds=1800)
|
|
+ timedelta(seconds=300)
|
|
+ timedelta(seconds=1800)
|
|
)
|
|
REGISTRATION_TIMEOUT = timedelta(seconds=600)
|
|
DOCUMENT_TIMEOUT = timedelta(seconds=300)
|
|
CLEANUP_TIMEOUT = timedelta(seconds=600)
|
|
RETENTION_TIMEOUT = timedelta(seconds=300)
|
|
|
|
# Fields no import may carry. The project is a frontend concern: the model listing is the MongoDB
|
|
# `models` document, no migrated database holds a project or model table, and nothing in this
|
|
# repository reads a project-to-model link (design D22). An input naming one is rejected rather than
|
|
# accepted and ignored, so no caller believes a project link was recorded. The same goes for a
|
|
# relational target: there is none, and for a plaintext password, which does not belong in an event
|
|
# history whatever else is true.
|
|
REJECTED_INPUT_FIELDS = (
|
|
'project_id',
|
|
'project_name',
|
|
'provisioning_target',
|
|
'schema',
|
|
'table_name',
|
|
'status_table',
|
|
'password',
|
|
)
|
|
|
|
DIGEST_LENGTH = 64
|
|
_HEX_DIGITS = frozenset('0123456789abcdef')
|
|
|
|
|
|
@workflow.defn(name='import_model')
|
|
class ImportModel:
|
|
"""Turn an encrypted `.sientia` object in MinIO into a registered model, or into an ERROR."""
|
|
|
|
@workflow.run
|
|
async def run(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Execute the import.
|
|
|
|
Args:
|
|
input_data: the import request.
|
|
Required keys:
|
|
- import_run_id (int): primary key of the import log row the frontend inserted
|
|
- object_key (str): the uploaded object's key
|
|
- expected_digest (str): the SHA-256 the uploader computed for those bytes
|
|
- password_envelope (str): the encrypted bundle password, never plaintext
|
|
Optional keys:
|
|
- bucket (str): defaults to the configured import bucket
|
|
|
|
Returns:
|
|
dict: `import_run_id`, `run_id`, `model_name`, `version` and the model document's id.
|
|
|
|
Raises:
|
|
ImportLogRowNotClaimableError: the record is not this workflow's — the run ends FAILED
|
|
with that error type and nothing is written to any row.
|
|
Exception: whatever failed. The record carries the failing step, its code and its
|
|
sentence; this exception carries the technical truth.
|
|
"""
|
|
workflow_id = workflow.info().workflow_id
|
|
import_run_id = self._require_import_run_id(input_data.get('import_run_id'))
|
|
metadata: dict[str, Any] = {
|
|
'model_name': '-',
|
|
'model_id': str(import_run_id),
|
|
'workflow_name': 'import_model',
|
|
'schedule_name': '-',
|
|
}
|
|
base = {
|
|
'metadata': metadata,
|
|
'import_run_id': import_run_id,
|
|
'workflow_id': workflow_id,
|
|
}
|
|
|
|
# Class A. Outside the try/finally on purpose: a claim that matches nothing leaves this
|
|
# workflow owning no row, so there is nothing to clean up and nothing to record.
|
|
claim = await workflow.execute_activity_method(
|
|
Activities.claim_import_status,
|
|
base,
|
|
retry_policy=IMPORT_STATUS_RETRY_POLICY,
|
|
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
|
|
schedule_to_close_timeout=STATUS_WRITE_ENVELOPE,
|
|
)
|
|
|
|
step: str = ImportStep.RECEIVED.value
|
|
gate: int | None = None
|
|
succeeded = False
|
|
source: dict[str, Any] = {}
|
|
result: dict[str, Any] = {'import_run_id': import_run_id}
|
|
|
|
try:
|
|
self._validate_input(input_data, claim)
|
|
|
|
# Download, decryption/extraction, experiment creation and artifact upload run inside
|
|
# one activity attempt (bug 8), so no later phase can depend on a file a different,
|
|
# separately-scheduled activity wrote to a different worker's disk. One hint covers the
|
|
# whole fused step; an exception it raises with no location of its own falls back to
|
|
# this step, since that is what the workflow was attempting.
|
|
step = ImportStep.DOWNLOAD.value
|
|
await self._record_hint(
|
|
base, Activities.record_import_status, {'step': step, 'source': None}
|
|
)
|
|
# Achado #9: the retention check carries no data this step consumes and already
|
|
# swallows its own failure (`_ensure_retention`), so it runs alongside the fused
|
|
# download+decrypt+provision activity instead of blocking it. Achado #8 fused
|
|
# `download_import_bundle` into `prepare_and_upload_import_bundle` before this change
|
|
# landed, so this gathers against that fused activity, not the original isolated one.
|
|
_, prepared = await asyncio.gather(
|
|
self._ensure_retention(base, input_data),
|
|
workflow.execute_activity_method(
|
|
Activities.prepare_and_upload_import_bundle,
|
|
{
|
|
**base,
|
|
'bucket': input_data.get('bucket'),
|
|
'object_key': input_data['object_key'],
|
|
'expected_digest': input_data['expected_digest'],
|
|
'password_envelope': input_data['password_envelope'],
|
|
},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=PREPARE_AND_UPLOAD_TIMEOUT,
|
|
),
|
|
)
|
|
source = prepared['metadata']
|
|
parameters = prepared['parameters']
|
|
model_name = source['model_name']
|
|
metadata['model_name'] = model_name
|
|
run_name = prepared['run_name']
|
|
|
|
# The names are known only once the fused activity returns; they are still recorded
|
|
# before the two provisioning steps that follow (the registered version and the model
|
|
# document), which is the guarantee `model-import-status-record` now makes.
|
|
await self._record_hint(
|
|
base,
|
|
Activities.record_import_names,
|
|
{
|
|
'experiment_name': source['experiment_name'],
|
|
'run_name': run_name,
|
|
'source': source,
|
|
'step': ImportStep.EXPERIMENT_CREATION.value,
|
|
},
|
|
)
|
|
|
|
step = ImportStep.REGISTRATION.value
|
|
await self._record_hint(
|
|
base,
|
|
Activities.record_import_status,
|
|
{'step': step, 'source': source or None},
|
|
)
|
|
version = await workflow.execute_activity_method(
|
|
Activities.register_import_model_version,
|
|
{**base, 'model_name': model_name, 'run_id': prepared['run_id']},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=REGISTRATION_TIMEOUT,
|
|
)
|
|
|
|
step = ImportStep.MODEL_DOCUMENT.value
|
|
await self._record_hint(
|
|
base,
|
|
Activities.record_import_status,
|
|
{'step': step, 'source': source or None},
|
|
)
|
|
document = await workflow.execute_activity_method(
|
|
Activities.write_import_model_document,
|
|
{
|
|
**base,
|
|
'model_name': model_name,
|
|
'target': parameters['target_variable'],
|
|
},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=DOCUMENT_TIMEOUT,
|
|
)
|
|
|
|
succeeded = True
|
|
result = {
|
|
'import_run_id': import_run_id,
|
|
'run_id': prepared['run_id'],
|
|
'model_name': model_name,
|
|
'version': version['version'],
|
|
'document_id': document.get('id'),
|
|
}
|
|
except BaseException as error: # NOSONAR - classified once below, then re-raised untouched
|
|
located_step, located_gate = failure_location(error)
|
|
step = located_step or step
|
|
gate = located_gate
|
|
await self._report_failure(base, error, step=step, gate=gate)
|
|
raise
|
|
finally:
|
|
cleanup = await self._cleanup(base)
|
|
await self._record_terminal(
|
|
base,
|
|
succeeded=succeeded,
|
|
step=step,
|
|
gate=gate,
|
|
source=source,
|
|
cleanup_failed=not cleanup.get('cleaned', True),
|
|
)
|
|
|
|
return result
|
|
|
|
# ------------------------------------------------------------------ validation
|
|
|
|
@staticmethod
|
|
def _require_import_run_id(value: Any) -> int:
|
|
"""The one check that runs before the claim, because the claim needs its result.
|
|
|
|
A missing or unusable id means there is no record to claim and none to write: the run ends
|
|
FAILED with this error type, and the failure is visible in the workflow's terminal state, its
|
|
error type, the error metric and a notification — everywhere except a row, which is the
|
|
point.
|
|
"""
|
|
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
|
|
raise ImportInputError('import_run_id must be a positive integer')
|
|
return value
|
|
|
|
@staticmethod
|
|
def _validate_input(input_data: dict[str, Any], claim: dict[str, Any]) -> None:
|
|
"""Validate the rest of the request, and check it against the record it just claimed.
|
|
|
|
This runs **after** the claim so its failure lands on the record rather than only in the
|
|
logs: by now the row is this workflow's, so a bad request is reported to the person who made
|
|
it. Everything here fails at step `received`.
|
|
"""
|
|
unknown = [name for name in REJECTED_INPUT_FIELDS if name in input_data]
|
|
if unknown:
|
|
raise ImportInputError(
|
|
f'the import request carries fields this workflow does not accept: '
|
|
f'{", ".join(sorted(unknown))}'
|
|
)
|
|
|
|
for required in ('object_key', 'expected_digest', 'password_envelope'):
|
|
if not input_data.get(required):
|
|
raise ImportInputError(f'the import request is missing {required}')
|
|
|
|
digest = str(input_data['expected_digest']).lower()
|
|
if len(digest) != DIGEST_LENGTH or not set(digest) <= _HEX_DIGITS:
|
|
raise ImportInputError('expected_digest must be a 64-character hexadecimal SHA-256')
|
|
|
|
object_key = input_data['object_key']
|
|
file_name = claim.get('file_name')
|
|
if file_name and file_name != object_key:
|
|
raise ImportInputError('the record names a different object than the request')
|
|
|
|
request_digest = claim.get('request_digest')
|
|
if request_digest and request_digest != digest:
|
|
raise ImportInputError('the record names a different digest than the request')
|
|
|
|
# --------------------------------------------------------------- status writes
|
|
|
|
async def _record_hint(
|
|
self,
|
|
base: dict[str, Any],
|
|
activity: Callable[..., Awaitable[Any]],
|
|
extra: dict[str, Any],
|
|
) -> None:
|
|
"""Record a status hint. A hint: its failure never fails the import."""
|
|
try:
|
|
await workflow.execute_activity_method(
|
|
activity,
|
|
{**base, **extra},
|
|
retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY,
|
|
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
|
|
)
|
|
except BaseException as error: # NOSONAR - a lost hint must not fail a working import
|
|
if is_cancelled_exception(error):
|
|
raise
|
|
await self._report_status_write_failure(base, error)
|
|
|
|
async def _record_terminal(
|
|
self,
|
|
base: dict[str, Any],
|
|
*,
|
|
succeeded: bool,
|
|
step: str | None,
|
|
gate: int | None,
|
|
source: dict[str, Any],
|
|
cleanup_failed: bool,
|
|
) -> None:
|
|
"""The single authoritative write, and the last thing this workflow does.
|
|
|
|
Its own failure is reported through the channel with the verdict it could not write, and it
|
|
never turns a completed import into a failed one — nor masks the cause of a real failure.
|
|
"""
|
|
try:
|
|
await workflow.execute_activity_method(
|
|
Activities.record_import_terminal_status,
|
|
{
|
|
**base,
|
|
'succeeded': succeeded,
|
|
'step': step,
|
|
'gate': gate,
|
|
'source': source or None,
|
|
'cleanup_failed': cleanup_failed,
|
|
},
|
|
retry_policy=IMPORT_STATUS_RETRY_POLICY,
|
|
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
|
|
schedule_to_close_timeout=STATUS_WRITE_ENVELOPE,
|
|
)
|
|
except (
|
|
BaseException
|
|
) as error: # NOSONAR - the outcome stands even when it cannot be written
|
|
if is_cancelled_exception(error):
|
|
raise
|
|
await self._report_status_write_failure(
|
|
base,
|
|
error,
|
|
verdict=f'succeeded={succeeded} step={step} gate={gate}',
|
|
)
|
|
|
|
@staticmethod
|
|
async def _report_failure(
|
|
base: dict[str, Any],
|
|
error: BaseException,
|
|
*,
|
|
step: str | None,
|
|
gate: int | None,
|
|
) -> None:
|
|
"""Route the failure that ended this import through the one failure channel.
|
|
|
|
Runs from the `except`, before cleanup and before the terminal write, so the log reads in the
|
|
order things happened: what failed, then what was cleaned, then what was recorded. The step
|
|
and the gate are the ones already located from the failure, so the log line and the row can
|
|
never name different ones.
|
|
|
|
Its own failure is swallowed. The channel is the thing that logs, so a report that cannot be
|
|
made has nowhere to say so, and the failure it was reporting is re-raised either way — the
|
|
event history and the record both still carry it.
|
|
"""
|
|
error_type, detail = describe_failure_chain(error)
|
|
try:
|
|
await workflow.execute_activity_method(
|
|
Activities.report_import_pipeline_failure,
|
|
{**base, 'step': step, 'gate': gate, 'detail': detail, 'error_type': error_type},
|
|
retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY,
|
|
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
|
|
)
|
|
except BaseException as error: # NOSONAR - the report of a report is where this stops
|
|
if is_cancelled_exception(error):
|
|
raise
|
|
|
|
async def _report_status_write_failure(
|
|
self,
|
|
base: dict[str, Any],
|
|
error: BaseException,
|
|
verdict: str | None = None,
|
|
) -> None:
|
|
"""Route a status write that could not be made through the one failure channel."""
|
|
try:
|
|
await workflow.execute_activity_method(
|
|
Activities.report_import_status_write_failure,
|
|
{**base, 'detail': str(error), 'verdict': verdict},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
|
|
)
|
|
except BaseException as error: # NOSONAR - the report of a report is where this stops
|
|
if is_cancelled_exception(error):
|
|
raise
|
|
# Deliberately silent otherwise: the channel this was trying to reach is the thing that
|
|
# logs, so there is nowhere left to say it, and an import that worked must not fail
|
|
# over it.
|
|
|
|
# ----------------------------------------------------------- side-effect steps
|
|
|
|
@staticmethod
|
|
async def _ensure_retention(base: dict[str, Any], input_data: dict[str, Any]) -> None:
|
|
"""Ensure the bundle bucket's expiry rule, without making it a reason to fail an import.
|
|
|
|
The rule is applied by the importer rather than assumed from a Helm value, so the acceptance
|
|
criterion is verifiable from the code that depends on it. But a bucket policy the importer
|
|
cannot set is an operational problem, not a bad import: failing a valid model import because
|
|
a lifecycle API refused would be the wrong trade, so the failure is dropped here and left
|
|
where Temporal already records it — the activity's own failure in the event history.
|
|
"""
|
|
try:
|
|
await workflow.execute_activity_method(
|
|
Activities.ensure_import_bundle_retention,
|
|
{**base, 'bucket': input_data.get('bucket')},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=RETENTION_TIMEOUT,
|
|
)
|
|
except BaseException as error: # NOSONAR - retention is not a precondition of an import
|
|
if is_cancelled_exception(error):
|
|
raise
|
|
# Dropped, not reported otherwise: the activity's own failure is already in the event
|
|
# history, and a bucket policy the importer cannot set is an operational problem rather
|
|
# than a bad import. Failing a valid model import over a lifecycle API would be the
|
|
# wrong trade.
|
|
|
|
@staticmethod
|
|
async def _cleanup(base: dict[str, Any]) -> dict[str, Any]:
|
|
"""Remove this import's temporary files. Runs on every path, before the terminal write."""
|
|
try:
|
|
return await workflow.execute_activity_method(
|
|
Activities.cleanup_import_files,
|
|
base,
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=CLEANUP_TIMEOUT,
|
|
)
|
|
except BaseException as error: # NOSONAR - cleanup never replaces the failure that brought us here
|
|
if is_cancelled_exception(error):
|
|
raise
|
|
return {'cleaned': False}
|