QTZPOC-20

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:
2026-09-08 17:06:25 -03:00
parent e4fc12e529
commit 192834bde3
34 changed files with 1485 additions and 531 deletions

View File

@@ -0,0 +1,40 @@
# Two independent MinIO expiry mechanisms
This platform expires objects in MinIO through two mechanisms that were built independently, for
different features, and remain independent today. This note exists so a reader who finds one does
not assume it is the whole story.
## `ModelImport.ensure_import_bundle_retention`
`laborious/activities/model_import.py`. Applies a **bucket lifecycle rule** (a MinIO/S3-native
`Expiration` rule, `ENABLED`, keyed by a fixed `rule_id`) to the import bundle prefix
(`import_config['prefix']`), expiring objects after `import_config['retention_days']` days. The
rule is declarative: once set, MinIO itself removes the object when it ages out — the workflow does
not scan for or delete anything. The activity's job is only to ensure the rule exists and matches
the configured `days`/`prefix`, which is what makes "lifecycle active" a criterion the code can
verify rather than an assumption about a Helm value.
## `Storage.cleanup_minio_objects_expired`
`laborious/activities/storage.py`. An older, unrelated activity that **scans and deletes**: given a
list of key prefixes and a retention window (`retention_hours`), it lists objects under each prefix,
compares their age against `now() - retention_hours`, and deletes the ones that have aged out. There
is no bucket-side rule — the deletion happens in this code, on whatever schedule calls the activity.
## Why they are not unified today
They expire different prefix families, on different units (days vs. hours), through different
mechanisms (declarative rule vs. active scan), and were written for different features at different
times. Neither mechanism depends on or interferes with the other; a bucket can carry an
`ensure_import_bundle_retention` lifecycle rule for one prefix and still have
`cleanup_minio_objects_expired` scanning a different prefix on its own cadence.
## What unifying them would require
Either retiring the scan-and-delete path in favor of a lifecycle rule for every prefix it currently
handles (only possible if every one of its callers can tolerate a fixed number of retention *days*,
since S3-style lifecycle rules do not support hour-granularity windows), or teaching the
lifecycle-rule path to manage a per-prefix table of rules with arbitrary retention windows instead of
the one fixed rule it manages today. Either direction touches callers and configuration this change
does not, and neither is a refactor — this is tracked as a separate, larger change, not a task of
this one.

View File

@@ -102,6 +102,9 @@ IMPORT_BUCKET = 'sientia-import'
IMPORT_PREFIX = 'imported_models/'
MODELS_COLLECTION = 'models'
MONGO_DATABASE = 'sientia'
# Must track `IMPORT_MODEL_ID_COUNTERS_COLLECTION`'s default in connectors_config.py — the harness
# never sets that env var, so production's default is what the worker under test actually reads from.
MODEL_ID_COUNTERS_COLLECTION = 'model_id_counters'
IMPORT_TASK_QUEUE = 'import-model-e2e-queue'
# Deliberately gapped, and deliberately strings: the baseline migration stores `id` as a string and
@@ -119,10 +122,7 @@ IMPORT_ACTIVITY_NAMES = (
'record_import_terminal_status',
'report_import_status_write_failure',
'report_import_pipeline_failure',
'download_import_bundle',
'open_import_bundle',
'create_import_experiment',
'upload_import_artifacts',
'prepare_and_upload_import_bundle',
'register_import_model_version',
'write_import_model_document',
'cleanup_import_files',
@@ -331,6 +331,12 @@ def import_models_collection(import_mongo_client: MongoClient) -> Any:
The index is the assertion's teeth: without it "the id continues the collection" and "a name
already present is left alone" would both pass against a collection that quietly took a
duplicate.
New ids no longer come from scanning this collection — `_next_model_document_id` reads an
atomic counter in a separate `model_id_counters` collection instead (a different database, no
less), so it survives this drop untouched. Reset it here too, to the same seeded high-water
mark, or the next id handed out silently carries over from whatever an earlier test in the same
run already claimed.
"""
collection = import_mongo_client[MONGO_DATABASE][MODELS_COLLECTION]
collection.drop()
@@ -351,6 +357,13 @@ def import_models_collection(import_mongo_client: MongoClient) -> Any:
]
)
collection.create_index('id', unique=True)
counters = import_mongo_client[MONGO_DATABASE][MODEL_ID_COUNTERS_COLLECTION]
counters.replace_one(
{'_id': MODELS_COLLECTION},
{'_id': MODELS_COLLECTION, 'seq': max(int(seeded) for seeded in SEEDED_MODEL_IDS)},
upsert=True,
)
return collection

View File

@@ -39,7 +39,7 @@ def _as_dict(value: Any) -> dict[str, Any]:
def _assert_step(row: dict[str, Any], step: ImportStep, workflow_id: str) -> None:
code, reason = IMPORT_FAILURE_CATALOG[step]
code, reason, _responsibility = IMPORT_FAILURE_CATALOG[step]
detail = _as_dict(row['orchestrator_response_data'])['import']
assert row['status'] == VALIDATION_ERROR_STATUS, row['status']
assert detail['step'] == step.value, detail

View File

@@ -92,7 +92,7 @@ async def _assert_rejected(
f'terminal status is {row["status"]!r}, expected {VALIDATION_ERROR_STATUS!r}'
)
expected_code, expected_reason = IMPORT_FAILURE_CATALOG[step]
expected_code, expected_reason, _expected_responsibility = IMPORT_FAILURE_CATALOG[step]
detail = _as_dict(row['orchestrator_response_data'])['import']
assert detail['step'] == step.value, (
f'recorded step is {detail["step"]!r}, expected {step.value!r}'

View File

@@ -244,6 +244,7 @@ class TestFactoryBundle:
'predict_flavor': 'joblib',
'retention_minutes': 60,
'target': 'SE',
'skip_predict_timestamp': True,
}
row = import_read_status_row(import_run_id)

View File

@@ -69,7 +69,7 @@ def _assert_platform_failure(
f"status is {row['status']!r}; a provisioning failure is not the user's to fix"
)
code, reason = IMPORT_FAILURE_CATALOG[step]
code, reason, _responsibility = IMPORT_FAILURE_CATALOG[step]
detail = _as_dict(row['orchestrator_response_data'])['import']
assert detail['step'] == step.value, detail
assert detail['code'] == code.value, detail
@@ -152,11 +152,13 @@ class TestExperimentCreation:
object_key=object_key,
)
# The bundle *was* opened before MLflow was reached, so the names are on the row — which is
# what task 9.3 asks for, observed from a failure that stopped exactly there.
# Download, decryption, experiment creation and artifact upload are one fused activity
# attempt (`model-import-fix-bugs`, bug 8): the workflow only learns the bundle's names once
# that whole attempt returns, and an experiment-creation failure means it never does. The
# placeholder names inserted with the row are therefore still there — narrowed from the old
# "names before MLflow is touched" invariant, per `model-import-status-record`'s updated spec.
assert row is not None
assert row['experiment_name'] == import_names['experiment_name']
assert row['run_name'] == f'import-{import_names["model_name"]}-v2'
assert row['experiment_name'] == 'import-e2e-placeholder'
# Nothing in Mongo: the document is the last step and never ran.
documents = list(import_models_collection.find({}, {'_id': False}))

View File

@@ -34,6 +34,7 @@ CONTRACTED_MODEL_CONFIG_KEYS = {
'predict_flavor',
'retention_minutes',
'target',
'skip_predict_timestamp',
}
@@ -203,7 +204,7 @@ class TestModelDocumentContent:
assert document['active'] is False
assert set(document['model_config']) == CONTRACTED_MODEL_CONFIG_KEYS, (
f'model_config holds {sorted(document["model_config"])}, not exactly the four '
f'model_config holds {sorted(document["model_config"])}, not exactly the '
f'contracted keys'
)
assert document['model_config'] == {
@@ -211,6 +212,7 @@ class TestModelDocumentContent:
'predict_flavor': 'joblib',
'retention_minutes': 60,
'target': 'PRODUCT_QUALITY',
'skip_predict_timestamp': True,
}
# The document has exactly these fields and no others — nothing was smuggled in alongside.

View File

@@ -155,24 +155,30 @@ class TestOrdering:
import_envelope: Any,
import_read_status_row: Any,
) -> None:
"""Observed at the first call experiment creation makes, so nothing exists in MLflow yet.
"""Observed at the first call registration makes, so the experiment and run already exist.
Download, decryption, experiment creation and artifact upload are fused into one activity
attempt (`model-import-status-record`'s design, decision on the fused step): the names are
not known to the workflow, and so cannot be recorded, until that whole activity returns —
the guarantee is anchored to the two steps that follow it, registration and the model
document, not to the experiment or run, which by construction already exist by then.
This is the stronger form of the same claim task 7.1 makes from a failure: here the import
*succeeds*, and the names were still already written — so their presence is not an artefact
of the run having stopped.
*succeeds*, and the names were still already written before registration ran — so their
presence is not an artefact of the run having stopped.
"""
repository = import_activities.import_mlflow_repository
real_get_experiment = repository.get_experiment
client = import_activities.import_mlflow_repository.client
real_get_registered_model = client.get_registered_model
snapshots: list[dict[str, Any] | None] = []
observed_run_id: list[int] = [0]
def observing_get_experiment(*args: Any, **kwargs: Any) -> Any:
def observing_get_registered_model(*args: Any, **kwargs: Any) -> Any:
snapshots.append(import_read_status_row(observed_run_id[0]))
return real_get_experiment(*args, **kwargs)
return real_get_registered_model(*args, **kwargs)
object_key, digest = import_upload_bundle(import_bundle_factory())
observed_run_id[0] = import_insert_status_row(object_key=object_key, expected_digest=digest)
repository.get_experiment = observing_get_experiment
client.get_registered_model = observing_get_registered_model
try:
handle = await import_execute(
import_build_input(
@@ -184,9 +190,9 @@ class TestOrdering:
)
await handle.result()
finally:
repository.get_experiment = real_get_experiment
client.get_registered_model = real_get_registered_model
assert snapshots, 'experiment creation never ran'
assert snapshots, 'registration never ran'
mid_flight = snapshots[0]
assert mid_flight is not None
assert mid_flight['status'] == RUNNING_STATUS
@@ -436,7 +442,7 @@ class TestInputAgainstTheRow:
row = import_read_status_row(import_run_id)
assert row is not None
code, reason = IMPORT_FAILURE_CATALOG[ImportStep.RECEIVED]
code, reason, _responsibility = IMPORT_FAILURE_CATALOG[ImportStep.RECEIVED]
assert row['status'] == VALIDATION_ERROR_STATUS
assert row['error_message'] == reason
detail = _as_dict(row['orchestrator_response_data'])['import']

View File

@@ -20,6 +20,7 @@ with workflow.unsafe.imports_passed_through():
build_import_config,
build_import_status_config,
)
from laborious.utils.repository.model_repository import MLFlowRepository
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
@@ -101,6 +102,15 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
secure=minio_config['secure'],
)
mlflow_repository = MLFlowRepository(
f'{mlflow_config["host"]}:{mlflow_config["port"]}',
mlflow_config['username'],
mlflow_config['password'],
logger,
notification_handler,
metrics_controller,
)
# Initialize parent classes
Storage.__init__(
self,
@@ -120,10 +130,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
MLFlow.__init__(
self,
mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
mlflow_repository=mlflow_repository,
minio_repository=minio_repository,
logger=logger,
notification_handler=notification_handler,
@@ -180,12 +187,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
status_config=import_status_config or build_import_status_config(),
import_config=import_config or build_import_config(),
minio_repository=minio_repository,
# The MLflow repository is shared rather than rebuilt: setting the tracking URI is a
# process-wide side effect, so two instances would be two chances to disagree. Read
# through `getattr` because `MLFlow.__init__` is what creates it — if that ever stops
# running before this line, the import activities refuse with "MLflow repository not
# initialized" at the point of use rather than silently provisioning nothing.
mlflow_repository=getattr(self, 'model_monitoring_repository', None),
mlflow_repository=mlflow_repository,
mongo_repository=mongo_repository,
logger=logger,
notification_handler=notification_handler,

View File

@@ -20,6 +20,7 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.utils.formatters import create_sample_dict
from laborious.utils.dataframe_debug import build_dataframe_debug_message
from laborious.utils.messages import MISSING_MINIO_REPOSITORY
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
from laborious.utils.repository.minio_manager import MinioManager
from laborious.utils.repository.model_repository import MLFlowRepository
@@ -37,10 +38,6 @@ class MLFlow(MinioManager):
MLFlow operations, ensuring reliable model inference in production environments.
Attributes:
mlflow_host (str): MLFlow server hostname
mlflow_port (int): MLFlow server port
mlflow_username (str): MLFlow authentication username
mlflow_password (str): MLFlow authentication password
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
"""
@@ -48,45 +45,25 @@ class MLFlow(MinioManager):
def __init__(
self,
mlflow_host: str,
mlflow_port: int,
mlflow_username: str,
mlflow_password: str,
mlflow_repository: MLFlowRepository,
minio_repository: MinioRepository | None = None,
logger: Logger | None = None,
notification_handler: NotificationHandler | None = None,
metrics_controller: MetricsController | None = None,
):
"""
Initialize MLFlow activities with server configuration.
Initialize MLFlow activities with an already-built repository.
Args:
mlflow_host: MLFlow server hostname or IP address
mlflow_port: MLFlow server port number
mlflow_username: Username for MLFlow authentication
mlflow_password: Password for MLFlow authentication
mlflow_repository: Repository for MLFlow operations, built and owned by the caller
(shared with other activities rather than rebuilt here)
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
Raises:
Exception: If MLFlowRepository initialization fails
"""
MinioManager.__init__(
self, minio_repository, logger, notification_handler, metrics_controller
)
self.mlflow_host = mlflow_host
self.mlflow_port = mlflow_port
self.mlflow_username = mlflow_username
self.mlflow_password = mlflow_password
self.model_monitoring_repository = MLFlowRepository(
f'{mlflow_host}:{mlflow_port}',
mlflow_username,
mlflow_password,
logger,
notification_handler,
metrics_controller,
)
self.model_monitoring_repository = mlflow_repository
def close(self) -> None:
"""
@@ -338,7 +315,7 @@ class MLFlow(MinioManager):
"""
if self.minio_repository is None:
raise ValueError('Minio repository not initialized')
raise ValueError(MISSING_MINIO_REPOSITORY)
metadata = input_data['metadata']

View File

@@ -28,6 +28,10 @@ with workflow.unsafe.imports_passed_through():
from minio.commonconfig import ENABLED, Filter
from minio.lifecycleconfig import Expiration, LifecycleConfig, Rule
from mlflow.entities import Param
from mlflow.exceptions import MlflowException
from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST, ErrorCode
from pymongo import ReturnDocument
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
@@ -38,7 +42,9 @@ with workflow.unsafe.imports_passed_through():
from sqlalchemy import create_engine, text
from sqlalchemy.engine import URL
from sqlalchemy.exc import DataError, IntegrityError, ProgrammingError
from sqlalchemy.sql import Executable
from temporalio.common import RetryPolicy
from temporalio.exceptions import ApplicationError
from laborious import metrics
from laborious.utils.bundle import format as fmt
@@ -49,15 +55,19 @@ with workflow.unsafe.imports_passed_through():
decrypt_password_envelope,
load_envelope_key,
)
from laborious.utils.messages import MISSING_MINIO_REPOSITORY as _MISSING_MINIO
from laborious.utils.models.import_errors import (
ImportErrorCode,
ImportGateRejectedError,
ImportLogRowNotClaimableError,
ImportMisconfiguredError,
ImportResponsibility,
ImportStatusRejectedError,
classify_import_failure,
report_import_failure,
)
from laborious.utils.repository.model_repository import MLFlowRepository
from laborious.utils.time import to_naive
# The table is a constant of the code, never an input: there is exactly one import log on this
# platform, and letting a caller name it would be a SQL-injection surface bought for nothing. The
@@ -69,7 +79,6 @@ IMPORT_LOG_TABLE: Final[str] = 'public.experiment_run'
RUN_NAME_MAX_LENGTH: Final[int] = 50
RUN_NAME_PREFIX: Final[str] = 'import-'
_MISSING_MINIO = 'Minio repository not initialized'
_MISSING_MLFLOW = 'MLflow repository not initialized'
_MISSING_MONGO = 'MongoDB repository not initialized'
@@ -107,20 +116,6 @@ class ExperimentRunStatus(StrEnum):
IMPORT_ERROR = 'IMPORT_ERROR' # new in V12
# A failure at these steps is the user's to act on — re-export, re-type the password, re-submit —
# which is the one thing the import screen can branch on without reading the error code.
USER_ACTIONABLE_STEPS: Final[frozenset[ImportStep]] = frozenset(
{
ImportStep.RECEIVED,
ImportStep.DOWNLOAD,
ImportStep.DECRYPTION,
ImportStep.ARCHIVE_INSPECTION,
ImportStep.EXTRACTION,
ImportStep.STRUCTURE_VALIDATION,
ImportStep.CONTENT_POLICY,
}
)
# Only the status writes retry, and only because they are idempotent: `UPDATE ... SET <fixed values>
# WHERE id` leaves the row the same however many times it runs. A provisioning call has no such
# shape, which is why it keeps `maximum_attempts=1`. Defined here rather than in `sientia_do`, which
@@ -144,20 +139,6 @@ IMPORT_STATUS_HINT_RETRY_POLICY: Final[RetryPolicy] = RetryPolicy(
non_retryable_error_types=['ImportLogRowNotClaimableError', 'ImportStatusRejectedError'],
)
_CLAIM_SQL = text(
"""
UPDATE public.experiment_run
SET status = :status,
updated_at = :updated_at,
orchestrator_response_data = CAST(:detail AS JSONB)
WHERE id = :import_run_id
AND run_type = 'IMPORT'
AND (status = 'PENDING'
OR (status = 'ORCHESTRATOR_WAITING_PROC'
AND orchestrator_response_data -> 'import' ->> 'workflow_id' = :workflow_id))
"""
)
_READ_ROW_SQL = text(
"""
SELECT id,
@@ -172,39 +153,70 @@ _READ_ROW_SQL = text(
"""
)
def _build_update(
columns: dict[str, str], *, casts: dict[str, str] | None = None, extra_where: str = ''
) -> Executable:
"""`UPDATE public.experiment_run SET <columns> WHERE id = :import_run_id [AND <extra_where>]`.
`columns` maps each SET target to its bind parameter name, e.g. `{'status': 'status'}` renders
`status = :status`. `casts` names which of those columns need `CAST(:name AS <cast>)` instead of
a bare `:name`. Order is preserved (Python 3.7+ dict order), so the rendered SQL is deterministic
and diff-friendly.
"""
casts = casts or {}
assignments = ',\n '.join(
f'{column} = CAST(:{param} AS {casts[column]})'
if column in casts
else f'{column} = :{param}'
for column, param in columns.items()
)
where = 'WHERE id = :import_run_id'
if extra_where:
where = f'{where}\n AND {extra_where}'
return text(
f'UPDATE public.experiment_run\n SET {assignments}\n {where}' # nosec B608 - columns/extra_where come only from literal constants in this module, never user input
)
_CLAIM_SQL = _build_update(
{'status': 'status', 'updated_at': 'updated_at', 'orchestrator_response_data': 'detail'},
casts={'orchestrator_response_data': 'JSONB'},
extra_where=(
"run_type = 'IMPORT'\n"
" AND (status = 'PENDING'\n"
" OR (status = 'ORCHESTRATOR_WAITING_PROC'\n"
" AND orchestrator_response_data -> 'import' ->> 'workflow_id' "
'= :workflow_id))'
),
)
# Class B and C key on the id alone: ownership was decided once, by the claim. Re-checking it on
# every later write buys nothing and costs the thing that matters — a predicate miss turning a
# pipeline that already imported the model into a reported failure.
_PROGRESS_SQL = text(
"""
UPDATE public.experiment_run
SET status = :status,
updated_at = :updated_at,
orchestrator_response_data = CAST(:detail AS JSONB)
WHERE id = :import_run_id
"""
_PROGRESS_SQL = _build_update(
{'status': 'status', 'updated_at': 'updated_at', 'orchestrator_response_data': 'detail'},
casts={'orchestrator_response_data': 'JSONB'},
)
_NAMES_SQL = text(
"""
UPDATE public.experiment_run
SET experiment_name = :experiment_name,
run_name = :run_name,
updated_at = :updated_at,
orchestrator_response_data = CAST(:detail AS JSONB)
WHERE id = :import_run_id
"""
_NAMES_SQL = _build_update(
{
'experiment_name': 'experiment_name',
'run_name': 'run_name',
'updated_at': 'updated_at',
'orchestrator_response_data': 'detail',
},
casts={'orchestrator_response_data': 'JSONB'},
)
_TERMINAL_SQL = text(
"""
UPDATE public.experiment_run
SET status = :status,
error_message = :error_message,
updated_at = :updated_at,
orchestrator_response_data = CAST(:detail AS JSONB)
WHERE id = :import_run_id
"""
_TERMINAL_SQL = _build_update(
{
'status': 'status',
'error_message': 'error_message',
'updated_at': 'updated_at',
'orchestrator_response_data': 'detail',
},
casts={'orchestrator_response_data': 'JSONB'},
)
@@ -222,7 +234,8 @@ def map_state_to_status(
return ExperimentRunStatus.ORCHESTRATOR_WAITING_PROC
if state is ImportState.COMPLETED:
return ExperimentRunStatus.IMPORT_SUCCESS
if step is not None and ImportStep(step) in USER_ACTIONABLE_STEPS:
failure = classify_import_failure(ImportStep(step) if step is not None else None)
if failure.responsibility is ImportResponsibility.USER:
return ExperimentRunStatus.ORCHESTRATOR_VALIDATION_ERROR
return ExperimentRunStatus.IMPORT_ERROR
@@ -333,8 +346,7 @@ class ModelImport(SientiaMonitoring):
@staticmethod
def _timestamp() -> Any:
"""`updated_at` is `TIMESTAMP WITHOUT TIME ZONE`, so the value handed over is naive."""
moment = now()
return moment.replace(tzinfo=None) if moment.tzinfo else moment
return to_naive(now())
@staticmethod
def _digest_from_request_data(request_data: Any) -> str | None:
@@ -415,29 +427,22 @@ class ModelImport(SientiaMonitoring):
import_run_id = input_data['import_run_id']
workflow_id = input_data['workflow_id']
affected = self._execute_status_write(
_CLAIM_SQL,
{
'status': ExperimentRunStatus.ORCHESTRATOR_WAITING_PROC.value,
'updated_at': self._timestamp(),
'detail': build_import_detail(
workflow_id=workflow_id,
step=ImportStep.RECEIVED,
source=input_data.get('source'),
),
'import_run_id': import_run_id,
'workflow_id': workflow_id,
},
)
row = self._read_row(import_run_id)
if affected == 0:
observed = self._describe_unclaimable(row, workflow_id)
error = ImportLogRowNotClaimableError(
f'import log row {import_run_id} could not be claimed: {observed}',
import_run_id=import_run_id,
observed=observed,
try:
affected = self._execute_status_write(
_CLAIM_SQL,
{
'status': ExperimentRunStatus.ORCHESTRATOR_WAITING_PROC.value,
'updated_at': self._timestamp(),
'detail': build_import_detail(
workflow_id=workflow_id,
step=ImportStep.RECEIVED,
source=input_data.get('source'),
),
'import_run_id': import_run_id,
'workflow_id': workflow_id,
},
)
except ImportStatusRejectedError as error:
await report_import_failure(
self,
None,
@@ -445,10 +450,30 @@ class ModelImport(SientiaMonitoring):
import_run_id=import_run_id,
workflow_id=workflow_id,
metadata=metadata,
code=ImportErrorCode.STATUS_NOT_RECORDED,
marker='status_write',
)
raise
row = self._read_row(import_run_id)
if affected == 0:
observed = self._describe_unclaimable(row, workflow_id)
not_claimable = ImportLogRowNotClaimableError(
f'import log row {import_run_id} could not be claimed: {observed}',
import_run_id=import_run_id,
observed=observed,
)
await report_import_failure(
self,
None,
not_claimable,
import_run_id=import_run_id,
workflow_id=workflow_id,
metadata=metadata,
code=ImportErrorCode.LOG_ROW_NOT_CLAIMABLE,
marker='not_claimable',
)
raise error
raise not_claimable
await self.emit_metric(
metric_object=metrics.MODEL_IMPORT_STARTED_COUNT,
@@ -562,8 +587,12 @@ class ModelImport(SientiaMonitoring):
code: ImportErrorCode | None = None
reason: str | None = None
else:
failure = classify_import_failure(ImportStep(step) if step else None)
status = map_state_to_status(ImportState.ERROR, step)
try:
step_enum = ImportStep(step) if step else None
except ValueError:
step_enum = None
failure = classify_import_failure(step_enum)
status = map_state_to_status(ImportState.ERROR, step_enum)
code, reason = failure.code, failure.reason
verdict = {
@@ -698,8 +727,7 @@ class ModelImport(SientiaMonitoring):
# --------------------------------------------------------- download and open
@activity.defn(name='download_import_bundle')
async def download_import_bundle(self, input_data: dict[str, Any]) -> dict[str, Any]:
async def _download_bundle(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""Gate 1, then the download, then the mandatory digest comparison.
The object's size and key shape are checked from `stat_object`, before its body is fetched.
@@ -722,7 +750,7 @@ class ModelImport(SientiaMonitoring):
ImportGateRejectedError: gate 1 rejected the object, or the digests disagree.
"""
if self.import_minio_repository is None:
raise ValueError(_MISSING_MINIO)
raise ImportMisconfiguredError(_MISSING_MINIO)
metadata = input_data.get('metadata', {})
bucket = self._effective_bucket(input_data.get('bucket'))
@@ -765,8 +793,7 @@ class ModelImport(SientiaMonitoring):
self.info(f'import bundle downloaded from {bucket} ({stat.size} bytes)', metadata)
return {'path': str(target), 'size': stat.size, 'digest': digest}
@activity.defn(name='open_import_bundle')
async def open_import_bundle(self, input_data: dict[str, Any]) -> dict[str, Any]:
async def _open_bundle(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""Open the envelope, then the bundle: gates 2 to 7, in the reader's own order.
The password envelope is decrypted **here**, inside the activity that uses the password, so
@@ -828,8 +855,7 @@ class ModelImport(SientiaMonitoring):
# ------------------------------------------------------------ MLflow provisioning
@activity.defn(name='create_import_experiment')
async def create_import_experiment(self, input_data: dict[str, Any]) -> dict[str, Any]:
async def _create_experiment(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""Get or create the experiment named by the bundle, keeping the origin's name.
Args (input_data):
@@ -840,7 +866,7 @@ class ModelImport(SientiaMonitoring):
dict: `experiment_id` and `experiment_name`.
"""
if self.import_mlflow_repository is None:
raise ValueError(_MISSING_MLFLOW)
raise ImportMisconfiguredError(_MISSING_MLFLOW)
experiment_name = input_data['experiment_name']
experiment = self.import_mlflow_repository.get_experiment(
@@ -852,21 +878,21 @@ class ModelImport(SientiaMonitoring):
)
return {'experiment_id': experiment.experiment_id, 'experiment_name': experiment_name}
@activity.defn(name='upload_import_artifacts')
async def upload_import_artifacts(self, input_data: dict[str, Any]) -> dict[str, Any]:
async def _upload_artifacts(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""Create a new run with the recorded name and log the bundle's tree and parameters.
Always a new run — the importer never appends to an existing one. The run is created with
exactly the `run_name` already written on the record, passed in rather than recomputed, so
the two cannot diverge. The artifact tree is logged verbatim, keeping `prediction_model` and
`data_model` at the run's artifact root, which is what the legacy loader resolves. The
bundle's `parameters` are logged unchanged, and the origin's own identity is logged
alongside them under `source.*` — records of where the model came from, not links.
exactly the `run_name` the fused activity computed, and the workflow records that same value
onto the row right after the fused activity returns, so the two cannot diverge. The artifact
tree is logged verbatim, keeping `prediction_model` and `data_model` at the run's artifact
root, which is what the legacy loader resolves. The bundle's `parameters` are logged
unchanged, and the origin's own identity is logged alongside them under `source.*` — records
of where the model came from, not links.
Args (input_data):
metadata (dict): workflow metadata.
experiment_id (str): the experiment to create the run in.
run_name (str): the name already on the record.
run_name (str): the name to create the run with.
extracted_dir (str): the opened bundle's directory.
parameters (dict): the bundle's `parameters` block, verbatim.
source (dict): the bundle's `metadata` block, logged as `source.*`.
@@ -875,7 +901,7 @@ class ModelImport(SientiaMonitoring):
dict: `run_id` and `run_name`.
"""
if self.import_mlflow_repository is None:
raise ValueError(_MISSING_MLFLOW)
raise ImportMisconfiguredError(_MISSING_MLFLOW)
client = self.import_mlflow_repository.client
run = client.create_run(
@@ -886,15 +912,92 @@ class ModelImport(SientiaMonitoring):
artifacts_dir = Path(input_data['extracted_dir']) / 'artifacts'
client.log_artifacts(run_id, str(artifacts_dir))
for name, value in (input_data.get('parameters') or {}).items():
client.log_param(run_id, name, value)
for name, value in (input_data.get('source') or {}).items():
client.log_param(run_id, f'source.{name}', value)
params = [
Param(name, str(value)) for name, value in (input_data.get('parameters') or {}).items()
]
params += [
Param(f'source.{name}', str(value))
for name, value in (input_data.get('source') or {}).items()
]
for i in range(0, len(params), 100): # MLflow's log_batch limit is 100 params per call
client.log_batch(run_id, params=params[i : i + 100])
client.set_terminated(run_id, 'FINISHED')
self.info(f'import artifacts uploaded to run {run_id}', input_data.get('metadata', {}))
return {'run_id': run_id, 'run_name': input_data['run_name']}
@activity.defn(name='prepare_and_upload_import_bundle')
async def prepare_and_upload_import_bundle(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""Download, decrypt, open, create the experiment, and upload artifacts — one disk, one try.
`_download_bundle`, `_open_bundle`, `_create_experiment` and `_upload_artifacts` used to be
four separately-schedulable Temporal activities sharing one worker-local scratch directory,
with nothing guaranteeing they landed on the same worker process — a real hazard once the
import worker runs more than one replica, since a later phase reading a file an earlier
phase wrote on a different pod fails with `FileNotFoundError`. Fusing them into one activity
means every disk write and every disk read this import makes happens inside one attempt, on
one process.
The run name is computed here, from the bundle's own `model_name`/`model_version`, because it
is needed to create the run — the workflow cannot compute it beforehand, since it does not
learn the bundle's metadata until this activity returns. The workflow records that same
`run_name` onto the row immediately after this activity returns (before the registered
version and the model document are written), rather than before the run is created, which is
the narrowed guarantee `model-import-status-record`'s names requirement now makes.
A retry of this activity redoes all four phases from the start. Re-uploading the same
artifacts to the same run is safe — MLflow overwrites identical content rather than
duplicating it — so the cost of a retry is bandwidth and time, not corruption.
Args (input_data):
metadata (dict): workflow metadata.
workflow_id (str): names this import's scratch directory.
bucket (str, optional): the effective bucket; the configured one when absent.
object_key (str): the uploaded object's key.
expected_digest (str): the uploader's SHA-256, lowercase hex.
password_envelope (str): base64 `nonce || ciphertext || tag`, AES-256-GCM.
Returns:
dict: `extracted_dir`, `metadata` (the bundle's own, i.e. `source`), `parameters`,
`digest`, `experiment_id`, `experiment_name`, `run_id`, `run_name`.
"""
downloaded = await self._download_bundle(input_data)
opened = await self._open_bundle({**input_data, **downloaded})
source = opened['metadata']
run_name = derive_run_name(source['model_name'], source['model_version'])
try:
experiment = await self._create_experiment(
{**input_data, 'experiment_name': source['experiment_name']}
)
except ImportMisconfiguredError:
raise
except Exception as error:
# Raw MLflow/network failures carry no location of their own (unlike the gates, which
# raise `ImportGateRejectedError`), so without this the workflow's fallback step —
# whatever it was attempting before this fused activity started — misclassifies the
# failure. Retryable: this is a broken dependency, not a rejected bundle.
raise ApplicationError(
str(error), {'step': ImportStep.EXPERIMENT_CREATION.value}
) from error
try:
run = await self._upload_artifacts(
{
**input_data,
**opened,
**experiment,
'run_name': run_name,
'source': source,
}
)
except ImportMisconfiguredError:
raise
except Exception as error:
raise ApplicationError(
str(error), {'step': ImportStep.ARTIFACT_UPLOAD.value}
) from error
return {**opened, **experiment, **run}
@activity.defn(name='register_import_model_version')
async def register_import_model_version(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""Register a version whose source the legacy loader can parse. No stage change.
@@ -920,7 +1023,7 @@ class ModelImport(SientiaMonitoring):
dict: `model_name`, `version` and `source`.
"""
if self.import_mlflow_repository is None:
raise ValueError(_MISSING_MLFLOW)
raise ImportMisconfiguredError(_MISSING_MLFLOW)
client = self.import_mlflow_repository.client
model_name = input_data['model_name']
@@ -928,7 +1031,9 @@ class ModelImport(SientiaMonitoring):
try:
client.get_registered_model(model_name)
except Exception:
except MlflowException as error:
if error.error_code != ErrorCode.Name(RESOURCE_DOES_NOT_EXIST):
raise
# First import of this name: the registered model is created on demand. A second import
# lands as a new version of the same registered model, never a differently named one.
client.create_registered_model(model_name)
@@ -949,10 +1054,11 @@ class ModelImport(SientiaMonitoring):
This is what the runtime reads to know how to load a model — `model_config.transform_flavor`,
`model_config.predict_flavor` and `model_config.target` — so without it an imported model is
invisible to every workflow. `id` continues the collection's own numbering: the largest
existing id **as an integer** plus one, never a count and never a lexicographic sort, which
would answer `"9" > "30"`. A name already in the collection is left exactly as it is:
`active`, `target` and the flavors of a model already on the platform are operator-owned.
invisible to every workflow. `id` continues the collection's own numbering, assigned by an
atomic counter (`_next_model_document_id`) rather than a scan of existing ids, so two
concurrent imports can never collide on the same id. A name already in the collection is left
exactly as it is: `active`, `target` and the flavors of a model already on the platform are
operator-owned.
Args (input_data):
metadata (dict): workflow metadata.
@@ -960,10 +1066,11 @@ class ModelImport(SientiaMonitoring):
target (str): the bundle's `parameters.target_variable`.
Returns:
dict: `written`, `id`, `skipped` and how many unusable ids were ignored.
dict: `written`, `id`, `skipped` and `ignored_ids` (always `0`, kept for shape
compatibility — no id is skipped under the counter).
"""
if self.mongo_repository is None:
raise ValueError(_MISSING_MONGO)
raise ImportMisconfiguredError(_MISSING_MONGO)
metadata = input_data.get('metadata', {})
collection = self.import_config['models_collection']
@@ -986,6 +1093,7 @@ class ModelImport(SientiaMonitoring):
'predict_flavor': 'joblib',
'retention_minutes': 60,
'target': input_data['target'],
'skip_predict_timestamp': True,
},
}
await self.mongo_repository.insert(collection, document, metadata)
@@ -995,22 +1103,32 @@ class ModelImport(SientiaMonitoring):
async def _next_model_document_id(
self, collection: str, metadata: dict[str, Any]
) -> tuple[str, int]:
"""`max(int(id)) + 1` as a decimal string, tolerating ids that are not numbers."""
if self.mongo_repository is None:
raise ValueError(_MISSING_MONGO)
raw_ids = await self.mongo_repository.distinct(collection, 'id', {}, metadata)
numeric: list[int] = []
ignored = 0
for value in raw_ids:
try:
numeric.append(int(str(value)))
except (TypeError, ValueError):
# A document whose id is not a decimal number cannot take part in the numbering.
# Ignored rather than fatal, and counted so the log says how many there were.
ignored += 1
if ignored:
self.info(f'ignored {ignored} non-numeric model document ids', metadata)
return str(max(numeric, default=0) + 1), ignored
"""Atomically increment a counter dedicated to `collection`, and return its new value.
Reaches MongoDB through the repository's own public `database` handle — the same access
pattern `find_one`/`distinct` already use internally — with pymongo's native
`find_one_and_update`, rather than the max-scan this replaces. That scan raced: two imports
writing at the same time could both read the same maximum and collide on the same id. An
atomic `$inc` cannot.
The counter's backing collection (`import_config['counters_collection']`) is an external
dependency: it must be created and seeded by a migration before this runs against a given
environment (see `openspec/changes/model-import-fix-bugs/design.md`, Decision 3). This code
does not create or seed it.
Assumes `self.mongo_repository` is not `None` — the only caller, `write_import_model_document`,
already raises before calling this when it is.
"""
counters = self.mongo_repository.database[ # type: ignore[union-attr]
self.import_config['counters_collection']
]
doc = counters.find_one_and_update(
{'_id': collection},
{'$inc': {'seq': 1}},
upsert=True,
return_document=ReturnDocument.AFTER,
)
return str(doc['seq']), 0
# ----------------------------------------------------------- cleanup and retention
@@ -1059,6 +1177,10 @@ class ModelImport(SientiaMonitoring):
("lifecycle active") is verifiable from the code that depends on it. The uploaded object is
never deleted by the import itself — it is left to this rule, as the record of an attempt.
See also: `Storage.cleanup_minio_objects_expired` (laborious/activities/storage.py) — an
older, unrelated scan-and-delete expiry mechanism for a different prefix family. The two are
independent by design for now; see docs/minio-object-expiry.md.
Args (input_data):
metadata (dict): workflow metadata.
bucket (str, optional): the effective bucket; the configured one when absent.
@@ -1067,7 +1189,7 @@ class ModelImport(SientiaMonitoring):
dict: whether the rule was already there and what it expires.
"""
if self.import_minio_repository is None:
raise ValueError(_MISSING_MINIO)
raise ImportMisconfiguredError(_MISSING_MINIO)
metadata = input_data.get('metadata', {})
bucket = self._effective_bucket(input_data.get('bucket'))
@@ -1083,6 +1205,20 @@ class ModelImport(SientiaMonitoring):
# deployments and with `None` on others; both mean "there is no rule yet".
existing = None
existing_rule = next(
(rule for rule in (existing.rules if existing else []) if rule.rule_id == rule_id), None
)
if (
existing_rule is not None
and existing_rule.status == ENABLED
and getattr(existing_rule.expiration, 'days', None) == days
and getattr(existing_rule.rule_filter, 'prefix', None) == prefix
):
# The read still runs every time (it is how drift — someone editing the rule outside
# this code — gets detected); only the write is skipped when nothing has drifted.
self.info(f'import retention rule already correct on {bucket} ({days} days)', metadata)
return {'bucket': bucket, 'prefix': prefix, 'days': days, 'rules': len(existing.rules)}
rules = [rule for rule in (existing.rules if existing else []) if rule.rule_id != rule_id]
rules.append(
Rule(

View File

@@ -17,7 +17,9 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.temporal.activities.postgres import Postgres
from sientia_do.temporal.constants import now
from laborious.utils.messages import MISSING_MINIO_REPOSITORY
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
from laborious.utils.time import to_naive
_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'})
@@ -83,7 +85,7 @@ class Storage(Postgres, MinioManager):
dict[str, Any]: Flat ``MinioDataFramePayload`` dict or ``success: False`` on failure.
"""
if self.minio_repository is None:
raise ValueError('Minio repository not initialized')
raise ValueError(MISSING_MINIO_REPOSITORY)
metadata: dict = input_data.get('metadata', {})
model_name = input_data['model_name']
@@ -129,6 +131,11 @@ class Storage(Postgres, MinioManager):
"""
Delete objects under the given prefixes that are older than the retention window.
See also: `ModelImport.ensure_import_bundle_retention`
(laborious/activities/model_import.py) — a newer, unrelated lifecycle-policy expiry
mechanism for the import bundle prefix. The two are independent by design for now; see
docs/minio-object-expiry.md.
Args (input_data):
metadata (dict): Workflow metadata for logging and metrics.
prefixes (list[str]): Key prefixes to scan (one level or subtree per prefix).
@@ -137,15 +144,13 @@ class Storage(Postgres, MinioManager):
dict[str, Any]: ``success``, ``deleted_count``, and optional ``message``.
"""
if self.minio_repository is None:
raise ValueError('Minio repository not initialized')
raise ValueError(MISSING_MINIO_REPOSITORY)
metadata = input_data.get('metadata', {})
payload = MinioDataFramePayload.from_dict(input_data['data'])
prefix = payload.cleanup_prefix()
base = now()
cutoff = (base.replace(tzinfo=None) if base.tzinfo else base) - timedelta(
hours=self.retention_hours
)
cutoff = to_naive(base) - timedelta(hours=self.retention_hours)
report: dict[str, Any] = {
'failed': {},

View File

@@ -180,8 +180,8 @@ class BundleReader:
header = self._gate_2_header(encrypted)
zip_path, kdf_seconds = self._gate_3_decrypt(encrypted, password, header, staging)
try:
self._gate_4_inspect(zip_path)
self._gate_5_extract(zip_path, root)
entries = self._gate_4_inspect(zip_path)
self._gate_5_extract(zip_path, root, entries)
metadata, parameters = self._gate_6_structure(root)
self._gate_7_content_policy(root)
except BaseException:
@@ -282,7 +282,7 @@ class BundleReader:
# ------------------------------------------------------------------ gate 4
def _gate_4_inspect(self, zip_path: Path) -> None:
def _gate_4_inspect(self, zip_path: Path) -> list[zipfile.ZipInfo]:
"""Inspect the central directory only. Nothing is written to disk by this gate."""
try:
with zipfile.ZipFile(zip_path) as archive:
@@ -319,6 +319,8 @@ class BundleReader:
self._check_entry_name(entry.filename)
self._check_entry_mode(entry)
return entries
@staticmethod
def _check_entry_name(name: str) -> None:
"""Reject absolute names and traversal, from the declared name alone."""
@@ -346,11 +348,13 @@ class BundleReader:
# ------------------------------------------------------------------ gate 5
def _gate_5_extract(self, zip_path: Path, root: Path) -> None:
def _gate_5_extract(self, zip_path: Path, root: Path, entries: list[zipfile.ZipInfo]) -> None:
"""Extract into a fresh directory, containing every member by resolved path.
A pre-existing extraction root is a failure rather than something to clear: it means
another import, or a previous attempt, owns that path.
another import, or a previous attempt, owns that path. `entries` is the central-directory
listing gate 4 already read; the archive is still opened here to read member bytes, but its
`infolist()` is not read a second time.
"""
if root.exists():
raise BundleGateError('extraction directory already exists', gate=5)
@@ -358,7 +362,7 @@ class BundleReader:
resolved_root = root.resolve()
try:
with zipfile.ZipFile(zip_path) as archive:
for entry in archive.infolist():
for entry in entries:
self._extract_member(archive, entry, root, resolved_root)
except BundleGateError:
shutil.rmtree(root, ignore_errors=True)
@@ -461,7 +465,7 @@ class BundleReader:
def _gate_7_content_policy(self, root: Path) -> None:
"""Only the file kinds an MLflow artifact tree calls for may be present."""
for path in sorted(root.rglob('*')):
for path in root.rglob('*'):
if path.is_dir():
continue
if not path.is_file():

View File

@@ -112,6 +112,13 @@ def build_import_config() -> dict[str, Any]:
IMPORT_MAX_UNCOMPRESSED_BYTES: gate 4 uncompressed size ceiling (default: 4 GiB)
IMPORT_MAX_COMPRESSION_RATIO: gate 4 compression ratio ceiling (default: 200)
IMPORT_MODELS_COLLECTION: MongoDB collection holding the model listing (default: models)
IMPORT_MODEL_ID_COUNTERS_COLLECTION: MongoDB collection holding the atomic id counters that
assign `models` document ids (default: model_id_counters). **External dependency**: this
collection and its seed value (`seq = max(int(id) for id in models)`) are not created by
this repository — a migration must add them to
`sientia-dataops-database-migrations/mongodb/migrations/`, following the `models`
baseline's `ensureCollection`/`ensureIndexes` convention, before any environment's counter
is first incremented. See `openspec/changes/model-import-fix-bugs/design.md`, Decision 3.
IMPORT_BUNDLE_RETENTION_DAYS: lifecycle expiry for the uploaded bundle (default: 7)
Returns:
@@ -130,6 +137,7 @@ def build_import_config() -> dict[str, Any]:
),
'max_compression_ratio': float(getenv('IMPORT_MAX_COMPRESSION_RATIO', '200')),
'models_collection': getenv('IMPORT_MODELS_COLLECTION', 'models'),
'counters_collection': getenv('IMPORT_MODEL_ID_COUNTERS_COLLECTION', 'model_id_counters'),
'retention_days': int(getenv('IMPORT_BUNDLE_RETENTION_DAYS', '7')),
}

View File

@@ -0,0 +1 @@
MISSING_MINIO_REPOSITORY = 'Minio repository not initialized'

View File

@@ -17,7 +17,8 @@ from temporalio import workflow
with workflow.unsafe.imports_passed_through():
import traceback
from collections.abc import Mapping
from collections import deque
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from enum import StrEnum
from types import MappingProxyType
@@ -58,74 +59,95 @@ class ImportErrorCode(StrEnum):
CLEANUP_INCOMPLETE = 'IMPORT_CLEANUP_INCOMPLETE'
class ImportResponsibility(StrEnum):
"""Who can act on a step's failure: the user (re-export, re-type, re-submit) or only support."""
USER = 'user'
SUPPORT = 'support'
@dataclass(frozen=True)
class ImportFailure:
"""One classification of one failure: the step, its code and the sentence it stores."""
"""One classification of one failure: the step, its code, its sentence and whose to fix."""
step: ImportStep | None
code: ImportErrorCode
reason: str
responsibility: ImportResponsibility = ImportResponsibility.SUPPORT
_CONTACT_SUPPORT = "Contact support with this import's identifier."
_RE_EXPORT = 'Export the model again from the origin platform.'
IMPORT_FAILURE_CATALOG: Final[Mapping[ImportStep, tuple[ImportErrorCode, str]]] = MappingProxyType(
IMPORT_FAILURE_CATALOG: Final[
Mapping[ImportStep, tuple[ImportErrorCode, str, ImportResponsibility]]
] = MappingProxyType(
{
ImportStep.RECEIVED: (
ImportErrorCode.REQUEST_INVALID,
'The import could not be started because the request was incomplete or '
'contradictory. Nothing was created. Submit the import again.',
ImportResponsibility.USER,
),
ImportStep.DOWNLOAD: (
ImportErrorCode.FILE_UNREADABLE,
'The uploaded file could not be read, or it is not the file that was sent. '
'Upload it again.',
ImportResponsibility.USER,
),
ImportStep.DECRYPTION: (
ImportErrorCode.BUNDLE_CANNOT_BE_OPENED,
'This file could not be opened: the password is wrong, or the file is damaged or '
'was altered. Confirm the password with whoever exported the model, then try '
'again.',
ImportResponsibility.USER,
),
ImportStep.ARCHIVE_INSPECTION: (
ImportErrorCode.BUNDLE_REJECTED,
"The file's contents did not pass the platform's safety checks, so it was not "
f'opened. {_RE_EXPORT}',
ImportResponsibility.USER,
),
ImportStep.EXTRACTION: (
ImportErrorCode.BUNDLE_UNPACK_FAILED,
f'The file could not be unpacked safely and nothing from it was kept. {_RE_EXPORT}',
ImportResponsibility.USER,
),
ImportStep.STRUCTURE_VALIDATION: (
ImportErrorCode.BUNDLE_INCOMPLETE,
'This file is not a complete model export — part of what the platform needs is '
f'missing from it. {_RE_EXPORT}',
ImportResponsibility.USER,
),
ImportStep.CONTENT_POLICY: (
ImportErrorCode.BUNDLE_UNEXPECTED_CONTENT,
'The file contains something a model export should not contain, so it was '
f'rejected. {_RE_EXPORT}',
ImportResponsibility.USER,
),
ImportStep.EXPERIMENT_CREATION: (
ImportErrorCode.STORAGE_PREPARATION_FAILED,
'The platform could not prepare a place to keep this model. The import stopped '
f'and no model was created. {_CONTACT_SUPPORT}',
ImportResponsibility.SUPPORT,
),
ImportStep.ARTIFACT_UPLOAD: (
ImportErrorCode.MODEL_FILES_NOT_STORED,
"The model's files could not be stored on this platform. The import stopped and "
f'the model is not available. {_CONTACT_SUPPORT}',
ImportResponsibility.SUPPORT,
),
ImportStep.REGISTRATION: (
ImportErrorCode.MODEL_NOT_REGISTERED,
"The model's files were stored, but the model itself could not be registered, so "
f'it cannot be used. {_CONTACT_SUPPORT}',
ImportResponsibility.SUPPORT,
),
ImportStep.MODEL_DOCUMENT: (
ImportErrorCode.MODEL_SETTINGS_NOT_SAVED,
'The model was registered, but the settings that tell the platform how to run it '
f'could not be saved. {_CONTACT_SUPPORT}',
ImportResponsibility.SUPPORT,
),
}
)
@@ -135,36 +157,44 @@ UNCLASSIFIED_FAILURE: Final[tuple[ImportErrorCode, str]] = (
f'The import stopped for an unexpected reason and no model was created. {_CONTACT_SUPPORT}',
)
# The three explicit-code failures are not pipeline steps, so `classify_import_failure` falls
# through to `UNCLASSIFIED_FAILURE`'s reason for all of them. That sentence says "no model was
# created", which is false for a cleanup-only failure (the model was imported) and misleading for
# the other two (neither ever attempted the pipeline). Each gets its own reason here instead.
_OVERRIDE_REASONS: Final[Mapping[ImportErrorCode, str]] = MappingProxyType(
{
ImportErrorCode.CLEANUP_INCOMPLETE: (
'The model was imported successfully; only its temporary files could not be removed '
'from the worker. No action needed.'
),
ImportErrorCode.LOG_ROW_NOT_CLAIMABLE: 'The import request could not be started.',
ImportErrorCode.STATUS_NOT_RECORDED: 'The import outcome could not be recorded.',
}
)
# `cleanup` deletes temporary files on the worker. It is never the reported failure: when every
# other step succeeded and only cleanup failed, the import is a success and the leak is a log line.
# So it deliberately has no user-facing sentence, and the catalog-completeness test knows it.
STEPS_WITHOUT_A_USER_FACING_REASON: Final[frozenset[ImportStep]] = frozenset({ImportStep.CLEANUP})
def classify_import_failure(
step: ImportStep | None, exception: BaseException | None = None
) -> ImportFailure:
def classify_import_failure(step: ImportStep | None) -> ImportFailure:
"""Turn a failure into the one code and the one sentence that will be recorded.
The exception's text is **ignored entirely** — it is what carries paths, table names and type
names, and it belongs in the log. What decides the outcome is the step.
The step alone decides the outcome — an exception's text carries paths, table names and type
names, and belongs in the log, never in the record.
Args:
step: the step that was being attempted, or `None` when none had been entered.
exception: the failure, when there is one to hand. It is accepted so callers read as
"classify this failure" and never absent because the classification needs it: the
catalog is keyed by step, and an exception cannot cross an activity boundary with its
type intact anyway — the terminal write classifies from the step alone.
Returns:
ImportFailure: step, code and the reason to store.
"""
del exception # the catalog is keyed by step; the exception's text never reaches the record
if step is None or step not in IMPORT_FAILURE_CATALOG:
code, reason = UNCLASSIFIED_FAILURE
return ImportFailure(step=step, code=code, reason=reason)
code, reason = IMPORT_FAILURE_CATALOG[step]
return ImportFailure(step=step, code=code, reason=reason)
code, reason, responsibility = IMPORT_FAILURE_CATALOG[step]
return ImportFailure(step=step, code=code, reason=reason, responsibility=responsibility)
class ImportLogRowNotClaimableError(ApplicationError):
@@ -198,6 +228,13 @@ class ImportStatusRejectedError(ApplicationError):
super().__init__(message, type='ImportStatusRejectedError', non_retryable=True)
class ImportMisconfiguredError(ApplicationError):
"""A required repository was not injected into this worker. A deploy defect, never transient."""
def __init__(self, message: str) -> None:
super().__init__(message, type='ImportMisconfiguredError', non_retryable=True)
class ImportInputError(ApplicationError):
"""The request is incomplete or contradictory. Fails at `received`, on the claimed row."""
@@ -241,6 +278,18 @@ _CHAIN_LIMIT = 10
_UNKNOWN_ERROR_TYPE = 'unavailable'
def walk_cause_chain(
exception: BaseException, limit: int = _CHAIN_LIMIT
) -> Iterator[BaseException]:
"""Yield `exception`, then each `__cause__`, up to `limit` links. Bounded against a cycle."""
candidate: BaseException | None = exception
for _ in range(limit):
if candidate is None:
return
yield candidate
candidate = candidate.__cause__
def _root_cause(exception: BaseException) -> BaseException:
"""The deepest `__cause__` in the chain: the failure that actually happened.
@@ -249,12 +298,7 @@ def _root_cause(exception: BaseException) -> BaseException:
says nothing an operator can act on. The chain is walked to its end, and bounded, because a cycle
in `__cause__` would otherwise hang the reporting of a failure.
"""
candidate = exception
for _ in range(_CHAIN_LIMIT):
if candidate.__cause__ is None:
break
candidate = candidate.__cause__
return candidate
return deque(walk_cause_chain(exception), maxlen=1)[-1]
def _type_name(exception: BaseException) -> str:
@@ -316,10 +360,7 @@ def describe_failure_chain(exception: BaseException) -> tuple[str, str]:
`ApplicationError` whose `type` says so.
"""
links: list[str] = []
candidate: BaseException | None = exception
for _ in range(_CHAIN_LIMIT):
if candidate is None:
break
for candidate in walk_cause_chain(exception):
kind = _type_name(candidate)
message = _one_line(str(candidate))
# `ApplicationError.__str__` already prefixes its own type, and `MlflowException:
@@ -330,7 +371,6 @@ def describe_failure_chain(exception: BaseException) -> tuple[str, str]:
links.append(message)
else:
links.append(f'{kind}: {message}')
candidate = candidate.__cause__
links.reverse()
root = links[0].split(':', 1)[0] if links else _UNKNOWN_ERROR_TYPE
return root, '\n'.join(links)
@@ -379,16 +419,14 @@ def failure_location(exception: BaseException | None) -> tuple[str | None, int |
Returns `(None, None)` when the failure did not say where it happened, and the caller then falls
back to the step it was attempting.
"""
seen = 0
candidate: BaseException | None = exception
while candidate is not None and seen < 10:
if exception is None:
return None, None
for candidate in walk_cause_chain(exception):
details = getattr(candidate, 'details', None)
if details and isinstance(details[0], dict) and 'step' in details[0]:
location = details[0]
gate = location.get('gate')
return location.get('step'), int(gate) if isinstance(gate, int) else None
candidate = candidate.__cause__
seen += 1
return None, None
@@ -466,9 +504,11 @@ async def report_import_failure(
Returns:
ImportFailure: the classification that was reported.
"""
failure = classify_import_failure(step, exception)
failure = classify_import_failure(step)
if code is not None:
failure = ImportFailure(step=step, code=code, reason=failure.reason)
failure = ImportFailure(
step=step, code=code, reason=_OVERRIDE_REASONS.get(code, failure.reason)
)
where = marker or (step.value if step else 'unknown')
if exception is not None:

View File

@@ -807,16 +807,23 @@ class MLFlowRepository(SientiaMonitoring):
)
if load_wrapper:
core_labels = self.get_core_labels(metadata, operation_type='download_model_wrapper')
self.info(f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}')
artifact_path = await self.dowload_artifacts(model_name, metadata, target)
try:
artifact_path = await self.dowload_artifacts(model_name, metadata, target)
self.info(
f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}'
)
self.info(
f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}'
)
raw_model = mlflow.pyfunc.load_model(artifact_path)
model = raw_model._model_impl.python_model
raw_model = mlflow.pyfunc.load_model(artifact_path)
model = raw_model._model_impl.python_model
except Exception as e:
await self.emit_metric(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels
)
raise e
self.debug(
f'Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}', metadata

View File

@@ -17,6 +17,7 @@ from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from laborious import metrics
from laborious.utils.models.import_errors import walk_cause_chain
# Requested session and secure channel lifetime (ms) before server revision; 10 minutes.
OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS = 10 * 60 * 1000
@@ -92,11 +93,9 @@ def _opc_status_from_exception(exc: BaseException) -> str:
Return:
str: Status class name or generic Python exception name.
"""
current: BaseException | None = exc
while current is not None:
if isinstance(current, UaStatusCodeError):
return type(current).__name__
current = current.__cause__
for candidate in walk_cause_chain(exc):
if isinstance(candidate, UaStatusCodeError):
return type(candidate).__name__
return type(exc).__name__

6
laborious/utils/time.py Normal file
View File

@@ -0,0 +1,6 @@
from datetime import datetime
def to_naive(moment: datetime) -> datetime:
"""Strip tzinfo, if any. Postgres `TIMESTAMP WITHOUT TIME ZONE` columns need naive values."""
return moment.replace(tzinfo=None) if moment.tzinfo else moment

View File

@@ -278,12 +278,10 @@ async def main():
activities.record_import_terminal_status,
activities.report_import_status_write_failure,
activities.report_import_pipeline_failure,
# Bundle
activities.download_import_bundle,
activities.open_import_bundle,
# Bundle download/open, MLflow experiment creation and artifact upload — one
# fused activity, so they never split across worker processes.
activities.prepare_and_upload_import_bundle,
# MLflow provisioning
activities.create_import_experiment,
activities.upload_import_artifacts,
activities.register_import_model_version,
# The model listing
activities.write_import_model_document,

View File

@@ -16,11 +16,13 @@ Shape of this workflow, and the reasons it has that shape:
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
@@ -29,7 +31,6 @@ with workflow.unsafe.imports_passed_through():
from laborious.activities.model_import import (
IMPORT_STATUS_HINT_RETRY_POLICY,
IMPORT_STATUS_RETRY_POLICY,
derive_run_name,
)
from laborious.utils.bundle.steps import ImportStep
from laborious.utils.models.import_errors import (
@@ -44,10 +45,15 @@ with workflow.unsafe.imports_passed_through():
# 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_TIMEOUT = timedelta(seconds=1800)
OPEN_TIMEOUT = timedelta(seconds=1800)
EXPERIMENT_TIMEOUT = timedelta(seconds=300)
ARTIFACT_UPLOAD_TIMEOUT = timedelta(seconds=1800)
# 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)
@@ -124,8 +130,7 @@ class ImportModel:
schedule_to_close_timeout=STATUS_WRITE_ENVELOPE,
)
current_step = ImportStep.RECEIVED
failed_step: str | None = None
step: str = ImportStep.RECEIVED.value
gate: int | None = None
succeeded = False
source: dict[str, Any] = {}
@@ -134,79 +139,74 @@ class ImportModel:
try:
self._validate_input(input_data, claim)
current_step = ImportStep.DOWNLOAD
await self._hint(base, current_step)
await self._ensure_retention(base, input_data)
download = await workflow.execute_activity_method(
Activities.download_import_bundle,
{
**base,
'bucket': input_data.get('bucket'),
'object_key': input_data['object_key'],
'expected_digest': input_data['expected_digest'],
},
retry_policy=retry_policy,
start_to_close_timeout=DOWNLOAD_TIMEOUT,
# 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}
)
current_step = ImportStep.DECRYPTION
await self._hint(base, current_step)
opened = await workflow.execute_activity_method(
Activities.open_import_bundle,
{
**base,
'path': download['path'],
'password_envelope': input_data['password_envelope'],
},
retry_policy=retry_policy,
start_to_close_timeout=OPEN_TIMEOUT,
# 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 = opened['metadata']
parameters = opened['parameters']
source = prepared['metadata']
parameters = prepared['parameters']
model_name = source['model_name']
metadata['model_name'] = model_name
run_name = derive_run_name(model_name, source['model_version'])
run_name = prepared['run_name']
# The names go on the record before anything exists in MLflow, so a person watching the
# import list sees *which* model is being imported while it is still running.
await self._record_names(base, source, run_name)
current_step = ImportStep.EXPERIMENT_CREATION
await self._hint(base, current_step, source)
experiment = await workflow.execute_activity_method(
Activities.create_import_experiment,
{**base, 'experiment_name': source['experiment_name']},
retry_policy=retry_policy,
start_to_close_timeout=EXPERIMENT_TIMEOUT,
)
current_step = ImportStep.ARTIFACT_UPLOAD
await self._hint(base, current_step, source)
run = await workflow.execute_activity_method(
Activities.upload_import_artifacts,
# 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,
{
**base,
'experiment_id': experiment['experiment_id'],
'experiment_name': source['experiment_name'],
'run_name': run_name,
'extracted_dir': opened['extracted_dir'],
'parameters': parameters,
'source': source,
'step': ImportStep.EXPERIMENT_CREATION.value,
},
retry_policy=retry_policy,
start_to_close_timeout=ARTIFACT_UPLOAD_TIMEOUT,
)
current_step = ImportStep.REGISTRATION
await self._hint(base, current_step, source)
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': run['run_id']},
{**base, 'model_name': model_name, 'run_id': prepared['run_id']},
retry_policy=retry_policy,
start_to_close_timeout=REGISTRATION_TIMEOUT,
)
current_step = ImportStep.MODEL_DOCUMENT
await self._hint(base, current_step, source)
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,
{
@@ -221,23 +221,23 @@ class ImportModel:
succeeded = True
result = {
'import_run_id': import_run_id,
'run_id': run['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)
failed_step = located_step or current_step.value
step = located_step or step
gate = located_gate
await self._report_failure(base, error, step=failed_step, gate=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=failed_step or current_step.value,
step=step,
gate=gate,
source=source,
cleanup_failed=not cleanup.get('cleaned', True),
@@ -294,43 +294,23 @@ class ImportModel:
# --------------------------------------------------------------- status writes
async def _hint(
async def _record_hint(
self,
base: dict[str, Any],
step: ImportStep,
source: dict[str, Any] | None = None,
activity: Callable[..., Awaitable[Any]],
extra: dict[str, Any],
) -> None:
"""Record the step about to be attempted. A hint: its failure never fails the import."""
"""Record a status hint. A hint: its failure never fails the import."""
try:
await workflow.execute_activity_method(
Activities.record_import_status,
{**base, 'step': step.value, 'source': source or None},
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
await self._report_status_write_failure(base, error)
async def _record_names(
self, base: dict[str, Any], source: dict[str, Any], run_name: str
) -> None:
"""Write the experiment and run names. Also a hint: the run is created either way."""
try:
await workflow.execute_activity_method(
Activities.record_import_names,
{
**base,
'experiment_name': source['experiment_name'],
'run_name': run_name,
'source': source,
'step': ImportStep.EXPERIMENT_CREATION.value,
},
retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY,
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
)
except BaseException as error: # NOSONAR - the names are for a human reading the list
await self._report_status_write_failure(base, error)
async def _record_terminal(
self,
base: dict[str, Any],

View File

@@ -17,11 +17,13 @@ from laborious.activities.storage import Storage
@patch('laborious.activities.activities.Gates.__init__')
@patch('laborious.activities.activities.ModelMetrics.__init__')
@patch('laborious.activities.activities.API.__init__')
@patch('laborious.activities.activities.MLFlowRepository')
@patch('laborious.activities.activities.MinioRepository')
@patch('laborious.activities.activities.MetricsController')
def test___init__(
mock_metrics_controller,
mock_minio_repository,
mock_mlflow_repository,
mock_api_init,
mock_model_metrics_init,
mock_gates_init,
@@ -101,16 +103,22 @@ def test___init__(
mock_mlflow_init.assert_called_once_with(
ANY,
mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
mlflow_repository=mock_mlflow_repository.return_value,
minio_repository=mock_minio_repository.return_value,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
mock_mlflow_repository.assert_called_once_with(
f'{mlflow_config["host"]}:{mlflow_config["port"]}',
mlflow_config['username'],
mlflow_config['password'],
logger,
notification_handler,
mock_metrics_controller.return_value,
)
mock_opc_init.assert_called_once_with(
ANY,
opc_servers=opc_config,
@@ -198,11 +206,13 @@ IMPORT_STATUS_CONFIG = {
@patch('laborious.activities.activities.ModelMetrics.__init__')
@patch('laborious.activities.activities.API.__init__')
@patch('laborious.activities.activities.MongoDBRepository')
@patch('laborious.activities.activities.MLFlowRepository')
@patch('laborious.activities.activities.MinioRepository')
@patch('laborious.activities.activities.MetricsController')
def test_the_import_activities_share_the_worker_s_mlflow_repository( # noqa: PLR0913
def test_mlflow_and_import_activities_receive_the_same_mlflow_repository_instance( # noqa: PLR0913
_mock_metrics_controller,
mock_minio_repository,
mock_mlflow_repository,
mock_mongo_repository,
_mock_api_init,
_mock_model_metrics_init,
@@ -214,11 +224,6 @@ def test_the_import_activities_share_the_worker_s_mlflow_repository( # noqa: PL
):
"""Setting the MLflow tracking URI is a process-wide side effect, so two repository instances
would be two chances to disagree about where MLflow is."""
repository = MagicMock(name='model_monitoring_repository')
mock_mlflow_init.side_effect = lambda instance, **_: setattr(
instance, 'model_monitoring_repository', repository
)
Activities(
**BASE_CONFIGS,
logger=MagicMock(),
@@ -228,6 +233,9 @@ def test_the_import_activities_share_the_worker_s_mlflow_repository( # noqa: PL
import_status_config=IMPORT_STATUS_CONFIG,
)
repository = mock_mlflow_repository.return_value
assert mock_mlflow_init.call_args.kwargs['mlflow_repository'] is repository
passed = mock_model_import_init.call_args.kwargs
assert passed['mlflow_repository'] is repository
assert passed['minio_repository'] is mock_minio_repository.return_value
@@ -244,11 +252,13 @@ def test_the_import_activities_share_the_worker_s_mlflow_repository( # noqa: PL
@patch('laborious.activities.activities.ModelMetrics.__init__')
@patch('laborious.activities.activities.API.__init__')
@patch('laborious.activities.activities.MongoDBRepository')
@patch('laborious.activities.activities.MLFlowRepository')
@patch('laborious.activities.activities.MinioRepository')
@patch('laborious.activities.activities.MetricsController')
def test_without_a_mongo_config_the_import_gets_no_repository( # noqa: PLR0913
_mock_metrics_controller,
_mock_minio_repository,
_mock_mlflow_repository,
mock_mongo_repository,
_mock_api_init,
_mock_model_metrics_init,
@@ -271,6 +281,44 @@ def test_without_a_mongo_config_the_import_gets_no_repository( # noqa: PLR0913
mock_mongo_repository.assert_not_called()
@patch('laborious.activities.activities.ModelImport.__init__')
@patch('laborious.activities.activities.Storage.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
@patch('laborious.activities.activities.ModelMetrics.__init__')
@patch('laborious.activities.activities.API.__init__')
@patch('laborious.activities.activities.MongoDBRepository')
@patch('laborious.activities.activities.MLFlowRepository')
@patch('laborious.activities.activities.MinioRepository')
@patch('laborious.activities.activities.MetricsController')
def test_exactly_one_mlflow_repository_is_built_even_with_import_active( # noqa: PLR0913
_mock_metrics_controller,
_mock_minio_repository,
mock_mlflow_repository,
_mock_mongo_repository,
_mock_api_init,
_mock_model_metrics_init,
_mock_gates_init,
_mock_opc_init,
_mock_mlflow_init,
_mock_storage_init,
_mock_model_import_init,
):
"""No second `MLFlowRepository` is ever created on any path: `MLFlow` and `ModelImport` only
receive the one instance built here, they never build their own."""
Activities(
**BASE_CONFIGS,
logger=MagicMock(),
notification_handler=MagicMock(),
mongo_config={'connection_string': 'mongodb://mongo', 'database_name': 'sientia'},
import_config=IMPORT_CONFIG,
import_status_config=IMPORT_STATUS_CONFIG,
)
mock_mlflow_repository.assert_called_once()
@mark.asyncio
@patch('laborious.activities.activities.Storage')
@patch('laborious.activities.activities.MLFlow')
@@ -278,9 +326,11 @@ def test_without_a_mongo_config_the_import_gets_no_repository( # noqa: PLR0913
@patch('laborious.activities.activities.Gates')
@patch('laborious.activities.activities.ModelMetrics')
@patch('laborious.activities.activities.API')
@patch('laborious.activities.activities.MLFlowRepository')
@patch('laborious.activities.activities.MinioRepository')
async def test_shutdown(
_mock_minio_repository,
_mock_mlflow_repository,
mock_api_init,
mock_model_metrics_init,
mock_gates_init,

View File

@@ -16,12 +16,12 @@ def _passthrough_from_dict():
yield
@patch('laborious.activities.mlflow.MLFlowRepository')
@patch('laborious.activities.mlflow.MinioRepository')
def test___init__(mock_minio_repository, mock_mlflow_repository):
def test___init__(mock_minio_repository):
logger = MagicMock()
notification_handler = MagicMock()
metrics_controller = AsyncMock()
mock_mlflow_repository = MagicMock()
minio_repo = mock_minio_repository(
endpoint='localhost:9000',
access_key='minio',
@@ -32,24 +32,14 @@ def test___init__(mock_minio_repository, mock_mlflow_repository):
bucket='test',
)
mlflow = MLFlow(
mlflow_host='http://localhost',
mlflow_port=5000,
mlflow_username='admin',
mlflow_password='admin',
mlflow_repository=mock_mlflow_repository,
minio_repository=minio_repo,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
assert mlflow.mlflow_host == 'http://localhost'
assert mlflow.mlflow_port == 5000
assert mlflow.mlflow_username == 'admin'
assert mlflow.mlflow_password == 'admin'
mock_mlflow_repository.assert_called_once_with(
'http://localhost:5000', 'admin', 'admin', ANY, ANY, ANY
)
assert mlflow.model_monitoring_repository is mock_mlflow_repository
mock_minio_repository.assert_called_once_with(
endpoint='localhost:9000',
@@ -62,10 +52,18 @@ def test___init__(mock_minio_repository, mock_mlflow_repository):
)
def test___init___requires_mlflow_repository():
"""The failure moves to the signature, at construction time, instead of a `None` that only
surfaces silently at first use."""
minio_repo = MagicMock()
logger = MagicMock()
with raises(TypeError):
MLFlow(minio_repository=minio_repo, logger=logger)
@fixture
@patch('laborious.activities.mlflow.MLFlowRepository')
@patch('laborious.activities.mlflow.MinioRepository')
def mlflow(mock_minio_repository, mock_mlflow_repository):
def mlflow(mock_minio_repository):
logger = MagicMock()
notification_handler = MagicMock()
metrics_controller = AsyncMock()
@@ -79,10 +77,7 @@ def mlflow(mock_minio_repository, mock_mlflow_repository):
bucket='test',
)
mlflow = MLFlow(
mlflow_host='http://localhost:5000',
mlflow_port=5000,
mlflow_username='admin',
mlflow_password='admin',
mlflow_repository=MagicMock(),
minio_repository=minio_repo,
logger=logger,
notification_handler=notification_handler,

View File

@@ -21,11 +21,12 @@ from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch
from pytest import fixture, mark, raises
from temporalio.exceptions import ApplicationError
from laborious import metrics
from laborious.activities.model_import import ModelImport
from laborious.activities.model_import import ModelImport, derive_run_name
from laborious.utils.bundle.steps import ImportStep
from laborious.utils.models.import_errors import ImportGateRejectedError
from laborious.utils.models.import_errors import ImportGateRejectedError, ImportMisconfiguredError
from tests.helpers import bundle_factory as factory
from tests.helpers.status_engine import FakeStatusEngine
@@ -150,7 +151,7 @@ async def test_a_matching_digest_yields_the_downloaded_file(
serve_object(minio_client, payload)
digest = hashlib.sha256(payload).hexdigest()
result = await import_activity.download_import_bundle(download_input(digest))
result = await import_activity._download_bundle(download_input(digest))
assert Path(result['path']).read_bytes() == payload
assert result['digest'] == digest
@@ -165,9 +166,7 @@ async def test_the_configured_bucket_is_used_when_the_request_names_none(
payload = factory.make_bundle(password=BUNDLE_PASSWORD)
serve_object(minio_client, payload)
await import_activity.download_import_bundle(
download_input(hashlib.sha256(payload).hexdigest())
)
await import_activity._download_bundle(download_input(hashlib.sha256(payload).hexdigest()))
assert minio_client.stat_object.call_args.args[0] == 'sientia'
@@ -179,7 +178,7 @@ async def test_the_request_may_name_another_bucket(
payload = factory.make_bundle(password=BUNDLE_PASSWORD)
serve_object(minio_client, payload)
await import_activity.download_import_bundle(
await import_activity._download_bundle(
download_input(hashlib.sha256(payload).hexdigest(), bucket='other-bucket')
)
@@ -197,7 +196,7 @@ async def test_a_digest_from_another_file_fails_before_the_header_is_parsed(
request = download_input(hashlib.sha256(other).hexdigest())
with raises(ImportGateRejectedError) as failure:
await import_activity.download_import_bundle(request)
await import_activity._download_bundle(request)
assert failure.value.details[0]['step'] == ImportStep.DOWNLOAD.value
assert failure.value.details[0]['gate'] == 1
@@ -212,7 +211,7 @@ async def test_a_file_that_does_not_match_is_not_left_on_the_worker(
request = download_input('0' * 64)
with raises(ImportGateRejectedError):
await import_activity.download_import_bundle(request)
await import_activity._download_bundle(request)
assert not (tmp_path / WORKFLOW_ID / 'bundle.sientia').exists()
@@ -224,7 +223,7 @@ async def test_the_digest_comparison_is_case_insensitive_on_the_input(
payload = factory.make_bundle(password=BUNDLE_PASSWORD)
serve_object(minio_client, payload)
result = await import_activity.download_import_bundle(
result = await import_activity._download_bundle(
download_input(hashlib.sha256(payload).hexdigest().upper())
)
@@ -239,7 +238,7 @@ async def test_an_object_over_the_ceiling_is_refused_before_its_body_is_fetched(
request = download_input('a' * 64)
with raises(ImportGateRejectedError) as failure:
await import_activity.download_import_bundle(request)
await import_activity._download_bundle(request)
assert failure.value.details[0]['step'] == ImportStep.DOWNLOAD.value
minio_client.fget_object.assert_not_called()
@@ -256,7 +255,7 @@ async def test_a_key_of_the_wrong_shape_is_refused(
request = download_input('a' * 64, object_key=object_key)
with raises(ImportGateRejectedError):
await import_activity.download_import_bundle(request)
await import_activity._download_bundle(request)
minio_client.fget_object.assert_not_called()
@@ -269,7 +268,7 @@ async def test_an_unreadable_object_is_a_download_failure(
request = download_input('a' * 64)
with raises(ImportGateRejectedError) as failure:
await import_activity.download_import_bundle(request)
await import_activity._download_bundle(request)
assert failure.value.details[0]['step'] == ImportStep.DOWNLOAD.value
@@ -282,7 +281,7 @@ async def test_a_failed_transfer_is_a_download_failure(
request = download_input('a' * 64)
with raises(ImportGateRejectedError) as failure:
await import_activity.download_import_bundle(request)
await import_activity._download_bundle(request)
assert failure.value.details[0]['step'] == ImportStep.DOWNLOAD.value
@@ -307,8 +306,11 @@ async def test_a_missing_minio_repository_is_a_configuration_error(tmp_path: Pat
request = download_input('a' * 64)
with raises(ValueError, match='Minio repository not initialized'):
await activity.download_import_bundle(request)
with raises(ImportMisconfiguredError, match='Minio repository not initialized') as failure:
await activity._download_bundle(request)
assert failure.value.non_retryable is True
assert failure.value.type == 'ImportMisconfiguredError'
# ----------------------------------------------------------------------- opening the bundle
@@ -321,7 +323,7 @@ async def test_a_well_formed_bundle_opens_and_returns_nothing_secret(
path = tmp_path / 'bundle.sientia'
path.write_bytes(factory.make_bundle(password=BUNDLE_PASSWORD))
result = await import_activity.open_import_bundle(open_input(str(path)))
result = await import_activity._open_bundle(open_input(str(path)))
assert result['metadata']['model_name'] == 'factory_model'
assert result['parameters']['target_variable'] == 'SE'
@@ -336,7 +338,7 @@ async def test_opening_a_bundle_measures_the_kdf(import_activity: ModelImport, t
path = tmp_path / 'bundle.sientia'
path.write_bytes(factory.make_bundle(password=BUNDLE_PASSWORD))
await import_activity.open_import_bundle(open_input(str(path)))
await import_activity._open_bundle(open_input(str(path)))
observed = [
call.kwargs
@@ -358,7 +360,7 @@ async def test_a_wrong_envelope_key_is_not_confused_with_a_wrong_bundle_password
request = open_input(str(path), envelope)
with raises(ImportGateRejectedError) as failure:
await import_activity.open_import_bundle(request)
await import_activity._open_bundle(request)
details = failure.value.details[0]
assert details['step'] == ImportStep.DECRYPTION.value
@@ -377,7 +379,7 @@ async def test_a_wrong_bundle_password_is_a_decryption_failure_of_its_own(
request = open_input(str(path))
with raises(ImportGateRejectedError) as failure:
await import_activity.open_import_bundle(request)
await import_activity._open_bundle(request)
details = failure.value.details[0]
assert details['step'] == ImportStep.DECRYPTION.value
@@ -394,7 +396,7 @@ async def test_a_missing_password_key_stops_the_import_at_decryption(
request = open_input(str(path))
with raises(ImportGateRejectedError) as failure:
await import_activity.open_import_bundle(request)
await import_activity._open_bundle(request)
assert failure.value.details[0]['cause'] == 'envelope'
@@ -413,7 +415,7 @@ async def test_a_tampered_or_malformed_envelope_is_refused(
request = open_input(str(path), envelope)
with raises(ImportGateRejectedError) as failure:
await import_activity.open_import_bundle(request)
await import_activity._open_bundle(request)
assert failure.value.details[0]['cause'] == 'envelope'
@@ -486,7 +488,7 @@ async def test_a_gate_failure_reaches_the_activity_as_its_own_step(
request = open_input(str(path))
with raises(ImportGateRejectedError) as failure:
await import_activity.open_import_bundle(request)
await import_activity._open_bundle(request)
assert failure.value.details[0]['step'] == step.value
assert failure.value.details[0]['gate'] == gate
@@ -501,7 +503,7 @@ async def test_a_rejected_bundle_leaves_no_extracted_tree(
request = open_input(str(path))
with raises(ImportGateRejectedError):
await import_activity.open_import_bundle(request)
await import_activity._open_bundle(request)
extract_root = tmp_path / WORKFLOW_ID / 'extract'
assert not extract_root.exists() or list(extract_root.rglob('*')) == []
@@ -545,7 +547,7 @@ async def test_no_failure_path_says_the_password_out_loud(
raised: BaseException | None = None
result = None
try:
result = await import_activity.open_import_bundle(open_input(str(path), envelope))
result = await import_activity._open_bundle(open_input(str(path), envelope))
except BaseException as error: # noqa: BLE001 - the point is that *whatever* it is says nothing
raised = error
@@ -572,7 +574,136 @@ async def test_the_envelope_never_leaves_the_activity_that_opens_it(
path.write_bytes(factory.make_bundle(password=BUNDLE_PASSWORD))
envelope = factory.build_envelope(BUNDLE_PASSWORD, ENVELOPE_KEY)
result = await import_activity.open_import_bundle(open_input(str(path), envelope))
result = await import_activity._open_bundle(open_input(str(path), envelope))
assert envelope not in str(result)
assert 'password' not in str(result).lower()
# ------------------------------------------------------- the fused activity (bug 8)
@mark.asyncio
async def test_prepare_and_upload_import_bundle_runs_all_four_phases_end_to_end(
import_activity: ModelImport, minio_client: MagicMock
):
"""Download, open, experiment creation and artifact upload, one activity, one attempt."""
payload = factory.make_bundle(password=BUNDLE_PASSWORD)
serve_object(minio_client, payload)
digest = hashlib.sha256(payload).hexdigest()
mlflow_repository = import_activity.import_mlflow_repository
mlflow_repository.get_experiment.return_value = MagicMock(experiment_id='17')
client = mlflow_repository.client
client.create_run.return_value.info.run_id = 'imported-run-id'
request = {
'metadata': METADATA,
'import_run_id': 77,
'workflow_id': WORKFLOW_ID,
'object_key': OBJECT_KEY,
'expected_digest': digest,
'password_envelope': factory.build_envelope(BUNDLE_PASSWORD, ENVELOPE_KEY),
}
result = await import_activity.prepare_and_upload_import_bundle(request)
expected_run_name = derive_run_name('factory_model', result['metadata']['model_version'])
assert result['metadata']['model_name'] == 'factory_model'
assert result['experiment_id'] == '17'
assert result['run_id'] == 'imported-run-id'
assert result['run_name'] == expected_run_name
mlflow_repository.get_experiment.assert_called_once_with(
'factory_experiment', create_if_not_exists=True
)
client.create_run.assert_called_once_with(experiment_id='17', run_name=expected_run_name)
client.log_artifacts.assert_called_once()
client.set_terminated.assert_called_once_with('imported-run-id', 'FINISHED')
@mark.asyncio
async def test_a_gate_failure_inside_the_fused_activity_still_names_its_step_and_gate(
import_activity: ModelImport, minio_client: MagicMock
):
"""A gate rejection during the fused activity's decryption/extraction phase is unaffected."""
payload = tampered()
serve_object(minio_client, payload)
digest = hashlib.sha256(payload).hexdigest()
request = {
'metadata': METADATA,
'import_run_id': 77,
'workflow_id': WORKFLOW_ID,
'object_key': OBJECT_KEY,
'expected_digest': digest,
'password_envelope': factory.build_envelope(BUNDLE_PASSWORD, ENVELOPE_KEY),
}
with raises(ImportGateRejectedError) as failure:
await import_activity.prepare_and_upload_import_bundle(request)
assert failure.value.details[0]['step'] == ImportStep.DECRYPTION.value
assert failure.value.details[0]['gate'] == 3
@mark.asyncio
async def test_a_retry_of_the_fused_activity_redoes_all_four_phases(import_activity: ModelImport):
"""A retry re-runs download, open, experiment creation and upload from the start.
Re-sending the same artifacts to the same run is safe (MLflow overwrites rather than
duplicates), so the cost of a retry is bandwidth and time, not corruption — but it is a real
cost, since nothing here is skipped on a second attempt.
"""
calls: list[str] = []
async def fake_download(_input_data: dict) -> dict:
calls.append('download')
return {'path': 'p', 'size': 1, 'digest': 'd'}
async def fake_open(_input_data: dict) -> dict:
calls.append('open')
return {
'extracted_dir': 'e',
'metadata': {
'model_name': 'factory_model',
'model_version': '1',
'experiment_name': 'factory_experiment',
},
'parameters': {},
'digest': 'd',
}
async def fake_create_experiment(_input_data: dict) -> dict:
calls.append('experiment')
return {'experiment_id': '1', 'experiment_name': 'factory_experiment'}
attempts = {'count': 0}
async def fake_upload(_input_data: dict) -> dict:
calls.append('upload')
attempts['count'] += 1
if attempts['count'] == 1:
raise RuntimeError('artifact upload failed on the first attempt')
return {'run_id': 'r', 'run_name': 'import-factory_model-v1'}
with (
patch.object(import_activity, '_download_bundle', side_effect=fake_download),
patch.object(import_activity, '_open_bundle', side_effect=fake_open),
patch.object(import_activity, '_create_experiment', side_effect=fake_create_experiment),
patch.object(import_activity, '_upload_artifacts', side_effect=fake_upload),
):
with raises(ApplicationError, match='artifact upload failed on the first attempt'):
await import_activity.prepare_and_upload_import_bundle({'metadata': METADATA})
await import_activity.prepare_and_upload_import_bundle({'metadata': METADATA})
assert calls == [
'download',
'open',
'experiment',
'upload',
'download',
'open',
'experiment',
'upload',
]

View File

@@ -16,9 +16,13 @@ from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch
import mlflow as mlflow_lib
from minio.commonconfig import ENABLED
from mlflow.exceptions import MlflowException
from mlflow.protos.databricks_pb2 import INTERNAL_ERROR, RESOURCE_DOES_NOT_EXIST
from pytest import fixture, mark, raises
from laborious.activities.model_import import ModelImport
from laborious.utils.models.import_errors import ImportMisconfiguredError
from laborious.utils.repository.model_repository import MLFlowRepository
from tests.helpers.status_engine import FakeStatusEngine
@@ -74,9 +78,23 @@ def mongo() -> AsyncMock:
repository = AsyncMock()
repository.find_one.return_value = None
repository.distinct.return_value = []
# `_next_model_document_id` reaches Mongo through the repository's plain pymongo `database`
# handle, not through the repository's own async methods, with a *synchronous* call
# (`find_one_and_update` is never awaited). A sane default so tests that don't care about the
# exact id still get one; tests that do override this return value (or its `side_effect`) per
# collection access.
repository.database.__getitem__.return_value.find_one_and_update.return_value = {
'_id': 'models',
'seq': 1,
}
return repository
def counter(mongo: AsyncMock) -> MagicMock:
"""The mock pymongo `Collection` behind `mongo_repository.database[collection]`."""
return mongo.database.__getitem__.return_value.find_one_and_update
@fixture
def import_activity(tmp_path: Path, mlflow_client: MagicMock, mongo: AsyncMock) -> ModelImport:
minio_repository = MagicMock()
@@ -102,6 +120,7 @@ def import_activity(tmp_path: Path, mlflow_client: MagicMock, mongo: AsyncMock)
'max_uncompressed_bytes': 4096,
'max_compression_ratio': 200.0,
'models_collection': 'models',
'counters_collection': 'model_id_counters',
'retention_days': 7,
},
minio_repository=minio_repository,
@@ -133,9 +152,7 @@ async def test_the_experiment_keeps_the_origin_s_name(import_activity: ModelImpo
repository = import_activity.import_mlflow_repository
repository.get_experiment.return_value = MagicMock(experiment_id='17')
result = await import_activity.create_import_experiment(
base(experiment_name='factory_experiment')
)
result = await import_activity._create_experiment(base(experiment_name='factory_experiment'))
repository.get_experiment.assert_called_once_with(
'factory_experiment', create_if_not_exists=True
@@ -148,8 +165,8 @@ async def test_a_missing_mlflow_repository_is_a_configuration_error(import_activ
import_activity.import_mlflow_repository = None
payload = base(experiment_name='e')
with raises(ValueError, match='MLflow repository not initialized'):
await import_activity.create_import_experiment(payload)
with raises(ImportMisconfiguredError, match='MLflow repository not initialized'):
await import_activity._create_experiment(payload)
# ------------------------------------------------------------------- the run and its tree
@@ -173,8 +190,8 @@ async def test_upload_artifacts_requires_the_mlflow_repository(
import_activity.import_mlflow_repository = None
payload = upload_input(tmp_path / 'extract')
with raises(ValueError, match='MLflow repository not initialized'):
await import_activity.upload_import_artifacts(payload)
with raises(ImportMisconfiguredError, match='MLflow repository not initialized'):
await import_activity._upload_artifacts(payload)
@mark.asyncio
@@ -182,7 +199,7 @@ async def test_every_import_creates_a_new_run_with_the_recorded_name(
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
):
"""The name is passed in, never recomputed, so the record and the run cannot diverge."""
result = await import_activity.upload_import_artifacts(upload_input(tmp_path / 'extract'))
result = await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
mlflow_client.create_run.assert_called_once_with(
experiment_id='17', run_name='import-factory_model-v2'
@@ -196,21 +213,30 @@ async def test_the_artifact_tree_is_logged_from_the_bundle_s_artifacts_directory
):
extracted = tmp_path / 'extract'
await import_activity.upload_import_artifacts(upload_input(extracted))
await import_activity._upload_artifacts(upload_input(extracted))
mlflow_client.log_artifacts.assert_called_once_with(RUN_ID, str(extracted / 'artifacts'))
def _logged_params(mlflow_client: MagicMock) -> dict:
"""Flatten every `client.log_batch(run_id, params=[...])` call into one key/value map."""
return {
param.key: param.value
for call in mlflow_client.log_batch.call_args_list
for param in call.kwargs['params']
}
@mark.asyncio
async def test_the_bundle_s_parameters_are_logged_verbatim(
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
):
await import_activity.upload_import_artifacts(upload_input(tmp_path / 'extract'))
await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
logged = {
call.args[1]: call.args[2]
for call in mlflow_client.log_param.call_args_list
if not call.args[1].startswith('source.')
key: value
for key, value in _logged_params(mlflow_client).items()
if not key.startswith('source.')
}
assert logged == BUNDLE_PARAMETERS
assert logged['target_variable'] == 'SE'
@@ -220,12 +246,12 @@ async def test_the_bundle_s_parameters_are_logged_verbatim(
async def test_the_origin_is_logged_as_source_parameters(
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
):
await import_activity.upload_import_artifacts(upload_input(tmp_path / 'extract'))
await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
logged = {
call.args[1]: call.args[2]
for call in mlflow_client.log_param.call_args_list
if call.args[1].startswith('source.')
key: value
for key, value in _logged_params(mlflow_client).items()
if key.startswith('source.')
}
assert logged == {f'source.{key}': value for key, value in BUNDLE_METADATA.items()}
# A record of where the model came from, not a link: the origin run id is never dereferenced.
@@ -234,11 +260,45 @@ async def test_the_origin_is_logged_as_source_parameters(
mlflow_client.download_artifacts.assert_not_called()
@mark.asyncio
async def test_parameters_are_logged_in_batches_of_at_most_100(
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
):
"""More params than `log_batch`'s per-call limit must span multiple calls, each within the
limit, and every key must still be logged exactly once across all of them."""
many_parameters = {f'param_{i}': str(i) for i in range(150)}
payload = upload_input(tmp_path / 'extract')
payload['parameters'] = many_parameters
await import_activity._upload_artifacts(payload)
calls = mlflow_client.log_batch.call_args_list
assert len(calls) > 1
for call in calls:
assert len(call.kwargs['params']) <= 100
logged = _logged_params(mlflow_client)
expected = {
**many_parameters,
**{f'source.{key}': value for key, value in BUNDLE_METADATA.items()},
}
assert logged == expected
@mark.asyncio
async def test_upload_artifacts_never_uses_log_param(
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
):
await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
mlflow_client.log_param.assert_not_called()
@mark.asyncio
async def test_the_run_is_closed_when_its_tree_is_in(
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
):
await import_activity.upload_import_artifacts(upload_input(tmp_path / 'extract'))
await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
mlflow_client.set_terminated.assert_called_once_with(RUN_ID, 'FINISHED')
@@ -251,7 +311,7 @@ async def test_register_version_requires_the_mlflow_repository(import_activity:
import_activity.import_mlflow_repository = None
payload = base(model_name='factory_model', run_id=RUN_ID)
with raises(ValueError, match='MLflow repository not initialized'):
with raises(ImportMisconfiguredError, match='MLflow repository not initialized'):
await import_activity.register_import_model_version(payload)
@@ -283,11 +343,19 @@ async def test_a_literal_runs_uri_would_not_be_parseable_by_the_legacy_loader():
async def test_a_registered_model_is_created_only_when_it_is_absent(
import_activity: ModelImport, mlflow_client: MagicMock
):
mlflow_client.get_registered_model.side_effect = RuntimeError('not found')
"""The only exception `get_registered_model` may treat as "does not exist yet": an
`MlflowException` whose `error_code` is literally `RESOURCE_DOES_NOT_EXIST`."""
mlflow_client.get_registered_model.side_effect = MlflowException(
'not found', error_code=RESOURCE_DOES_NOT_EXIST
)
await import_activity.register_import_model_version(base(model_name='new_model', run_id=RUN_ID))
result = await import_activity.register_import_model_version(
base(model_name='new_model', run_id=RUN_ID)
)
mlflow_client.create_registered_model.assert_called_once_with('new_model')
mlflow_client.create_model_version.assert_called_once()
assert result['model_name'] == 'new_model'
@mark.asyncio
@@ -303,13 +371,46 @@ async def test_a_second_import_reuses_the_registered_model(
mlflow_client.create_registered_model.assert_not_called()
@mark.asyncio
async def test_a_different_mlflow_error_is_not_treated_as_absent(
import_activity: ModelImport, mlflow_client: MagicMock
):
"""Only `RESOURCE_DOES_NOT_EXIST` means "create it". Any other MLflow error code — a transient
failure, an internal error — must not be swallowed and must not trigger a spurious create."""
error = MlflowException('backend is unavailable', error_code=INTERNAL_ERROR)
mlflow_client.get_registered_model.side_effect = error
payload = base(model_name='factory_model', run_id=RUN_ID)
with raises(MlflowException, match='backend is unavailable') as excinfo:
await import_activity.register_import_model_version(payload)
assert excinfo.value is error
mlflow_client.create_registered_model.assert_not_called()
mlflow_client.create_model_version.assert_not_called()
@mark.asyncio
async def test_a_non_mlflow_exception_is_never_treated_as_absent(
import_activity: ModelImport, mlflow_client: MagicMock
):
"""A bare `ConnectionError` (or any non-`MlflowException`) is not this function's business to
interpret at all: it must propagate unmodified, exactly as any other unexpected error would."""
mlflow_client.get_registered_model.side_effect = ConnectionError('connection refused')
payload = base(model_name='factory_model', run_id=RUN_ID)
with raises(ConnectionError, match='connection refused'):
await import_activity.register_import_model_version(payload)
mlflow_client.create_registered_model.assert_not_called()
@mark.asyncio
async def test_the_import_never_promotes_anything(
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
):
"""A `COMPLETED` import is a registered model, not a servable one."""
await import_activity.create_import_experiment(base(experiment_name='factory_experiment'))
await import_activity.upload_import_artifacts(upload_input(tmp_path / 'extract'))
await import_activity._create_experiment(base(experiment_name='factory_experiment'))
await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
await import_activity.register_import_model_version(
base(model_name='factory_model', run_id=RUN_ID)
)
@@ -322,7 +423,7 @@ async def test_the_import_never_promotes_anything(
async def test_no_provisioning_call_passes_a_stage(
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
):
await import_activity.upload_import_artifacts(upload_input(tmp_path / 'extract'))
await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
result = await import_activity.register_import_model_version(
base(model_name='factory_model', run_id=RUN_ID)
)
@@ -391,10 +492,12 @@ def test_the_legacy_loader_cannot_read_a_logged_model_uri():
@mark.asyncio
async def test_the_document_id_continues_the_collection_across_gaps(
async def test_the_document_id_is_whatever_the_atomic_counter_returns(
import_activity: ModelImport, mongo: AsyncMock
):
mongo.distinct.return_value = [str(value) for value in [*range(1, 19), *range(20, 25), 30]]
"""`id` is no longer computed here at all: it is `str(seq)` off the counter document
`find_one_and_update` hands back, whatever gaps or history the collection has."""
counter(mongo).return_value = {'_id': 'models', 'seq': 31}
result = await import_activity.write_import_model_document(
base(model_name='factory_model', target='SE')
@@ -404,20 +507,26 @@ async def test_the_document_id_continues_the_collection_across_gaps(
@mark.asyncio
async def test_the_document_id_is_numeric_not_lexicographic(
async def test_the_document_id_is_the_counter_s_seq_verbatim_not_reparsed(
import_activity: ModelImport, mongo: AsyncMock
):
mongo.distinct.return_value = ['9', '30']
"""A value large enough that the old max-scan's `"9" > "30"` lexicographic bug would have
mangled it, demonstrating the new id is never routed through any string comparison at all."""
counter(mongo).return_value = {'_id': 'models', 'seq': 100}
result = await import_activity.write_import_model_document(
base(model_name='factory_model', target='SE')
)
assert result['id'] == '31'
assert result['id'] == '100'
@mark.asyncio
async def test_an_empty_collection_starts_at_one(import_activity: ModelImport, mongo: AsyncMock):
"""`$inc` on a counter document that does not exist yet starts from zero, so the first import
against a fresh collection gets id `1`."""
counter(mongo).return_value = {'_id': 'models', 'seq': 1}
result = await import_activity.write_import_model_document(
base(model_name='factory_model', target='SE')
)
@@ -429,7 +538,7 @@ async def test_an_empty_collection_starts_at_one(import_activity: ModelImport, m
async def test_the_document_is_exactly_the_contracted_shape(
import_activity: ModelImport, mongo: AsyncMock
):
mongo.distinct.return_value = ['7']
counter(mongo).return_value = {'_id': 'models', 'seq': 8}
await import_activity.write_import_model_document(base(model_name='factory_model', target='SE'))
@@ -445,10 +554,28 @@ async def test_the_document_is_exactly_the_contracted_shape(
'predict_flavor': 'joblib',
'retention_minutes': 60,
'target': 'SE',
'skip_predict_timestamp': True,
},
}
@mark.asyncio
@mark.parametrize(
'model_name, target',
[('factory_model', 'SE'), ('another_model', 'temperature')],
)
async def test_the_document_always_carries_skip_predict_timestamp_true(
import_activity: ModelImport, mongo: AsyncMock, model_name: str, target: str
):
"""The opt-out is decided by the importer at write time, never left for an operator to set."""
counter(mongo).return_value = {'_id': 'models', 'seq': 8}
await import_activity.write_import_model_document(base(model_name=model_name, target=target))
_, document, _ = mongo.insert.await_args.args
assert document['model_config']['skip_predict_timestamp'] is True
@mark.asyncio
async def test_the_document_carries_no_field_the_runtime_does_not_read(
import_activity: ModelImport, mongo: AsyncMock
@@ -474,23 +601,31 @@ async def test_a_name_already_in_the_collection_is_left_exactly_as_it_is(
assert result == {'written': False, 'skipped': True, 'id': '12'}
mongo.insert.assert_not_awaited()
mongo.distinct.assert_not_awaited()
# No counter increment is spent on a name that was already there: a duplicate import must
# never move the sequence forward for a document it isn't going to write.
counter(mongo).assert_not_called()
@mark.asyncio
async def test_unusable_ids_are_ignored_and_counted(import_activity: ModelImport, mongo: AsyncMock):
mongo.distinct.return_value = ['1', 'legacy-model', None, '3']
async def test_two_imports_in_a_row_never_reuse_an_id(
import_activity: ModelImport, mongo: AsyncMock
):
"""No dependency on `distinct` at all: two sequential calls simply relay whatever an atomic
(fake, here) counter hands back, strictly increasing, with no scan of existing documents."""
seq = {'value': 40}
result = await import_activity.write_import_model_document(
base(model_name='factory_model', target='SE')
)
def increment(*args, **kwargs): # noqa: ARG001 - mirrors pymongo's call signature
seq['value'] += 1
return {'_id': 'models', 'seq': seq['value']}
assert result['id'] == '4'
assert result['ignored_ids'] == 2
assert any(
'ignored 2 non-numeric' in str(call)
for call in import_activity.logger.custom_info.call_args_list
)
counter(mongo).side_effect = increment
first_id, _ = await import_activity._next_model_document_id('models', METADATA)
second_id, _ = await import_activity._next_model_document_id('models', METADATA)
assert (first_id, second_id) == ('41', '42')
assert int(second_id) > int(first_id)
mongo.distinct.assert_not_called()
@mark.asyncio
@@ -498,14 +633,14 @@ async def test_a_duplicate_id_fails_the_step_without_recomputing_it(
import_activity: ModelImport, mongo: AsyncMock
):
"""No retry and no second guess: the unique index says the collection moved under us."""
mongo.distinct.return_value = ['1']
counter(mongo).return_value = {'_id': 'models', 'seq': 2}
mongo.insert.side_effect = RuntimeError('E11000 duplicate key error: id_1')
payload = base(model_name='factory_model', target='SE')
with raises(RuntimeError, match='duplicate key'):
await import_activity.write_import_model_document(payload)
assert mongo.distinct.await_count == 1
assert counter(mongo).call_count == 1
assert mongo.insert.await_count == 1
@@ -513,11 +648,18 @@ async def test_a_duplicate_id_fails_the_step_without_recomputing_it(
async def test_the_import_only_reads_and_inserts_documents(
import_activity: ModelImport, mongo: AsyncMock
):
"""No update, no replace, no upsert: an import cannot change a model already on the platform."""
"""No update, no replace, no upsert on the model document itself: an import cannot change a
model already on the platform. The atomic counter's own `$inc` against a dedicated counters
collection is a different concern and is explicitly allowed."""
await import_activity.write_import_model_document(base(model_name='factory_model', target='SE'))
used = {name for name, _, _ in mongo.mock_calls}
assert used <= {'find_one', 'distinct', 'insert'}
assert used <= {
'find_one',
'insert',
'database.__getitem__',
'database.__getitem__().find_one_and_update',
}
@mark.asyncio
@@ -527,29 +669,27 @@ async def test_a_missing_mongo_repository_is_a_configuration_error(
import_activity.mongo_repository = None
payload = base(model_name='m', target='SE')
with raises(ValueError, match='MongoDB repository not initialized'):
with raises(ImportMisconfiguredError, match='MongoDB repository not initialized'):
await import_activity.write_import_model_document(payload)
@mark.asyncio
async def test_next_model_document_id_also_guards_a_missing_mongo_repository(
import_activity: ModelImport,
):
"""The helper repeats the guard so it stays safe if ever called from somewhere new."""
import_activity.mongo_repository = None
with raises(ValueError, match='MongoDB repository not initialized'):
await import_activity._next_model_document_id('models', {})
def test_no_import_module_writes_an_active_document():
"""Asserted from the source too, so `active: false` cannot be flipped by a later edit."""
"""Asserted from the source too, so `active: false` cannot be flipped by a later edit.
`upsert` is checked against `write_import_model_document`'s own body only, not the whole file:
`_next_model_document_id`'s atomic counter legitimately upserts a *counter* document in a
dedicated collection — a different concern from ever updating, replacing or upserting the model
document itself, which this test still forbids absolutely.
"""
source = Path('laborious/activities/model_import.py').read_text()
start = source.index('async def write_import_model_document')
end = source.index('async def _next_model_document_id')
document_write = source[start:end]
assert "'active': False" in source
assert "'active': True" not in source
for forbidden in ('update_one', 'replace_one', 'update_many', 'upsert'):
assert forbidden not in source
assert forbidden not in document_write
# ------------------------------------------------------------------ cleanup and retention
@@ -613,7 +753,7 @@ async def test_ensure_retention_requires_the_minio_repository(import_activity: M
import_activity.import_minio_repository = None
payload = base()
with raises(ValueError, match='Minio repository not initialized'):
with raises(ImportMisconfiguredError, match='Minio repository not initialized'):
await import_activity.ensure_import_bundle_retention(payload)
@@ -688,6 +828,58 @@ async def test_the_retention_rule_can_be_applied_to_a_requested_bucket(
assert client.set_bucket_lifecycle.call_args.args[0] == 'other-bucket'
@mark.asyncio
async def test_a_correct_existing_rule_is_not_rewritten(import_activity: ModelImport):
client = import_activity.import_minio_repository.minio_client
existing = MagicMock()
matching_rule = MagicMock(rule_id='sientia-import-bundle-expiry', status=ENABLED)
matching_rule.expiration.days = 7
matching_rule.rule_filter.prefix = 'imported_models/'
existing.rules = [matching_rule]
client.get_bucket_lifecycle.return_value = existing
result = await import_activity.ensure_import_bundle_retention(base())
client.set_bucket_lifecycle.assert_not_called()
assert result == {'bucket': 'sientia', 'prefix': 'imported_models/', 'days': 7, 'rules': 1}
@mark.asyncio
async def test_a_drifted_existing_rule_is_rewritten(import_activity: ModelImport):
client = import_activity.import_minio_repository.minio_client
existing = MagicMock()
drifted_rule = MagicMock(rule_id='sientia-import-bundle-expiry', status=ENABLED)
drifted_rule.expiration.days = 3
drifted_rule.rule_filter.prefix = 'imported_models/'
existing.rules = [drifted_rule]
client.get_bucket_lifecycle.return_value = existing
await import_activity.ensure_import_bundle_retention(base())
config = client.set_bucket_lifecycle.call_args.args[1]
corrected = next(
rule for rule in config.rules if rule.rule_id == 'sientia-import-bundle-expiry'
)
assert corrected.expiration.days == 7
@mark.asyncio
async def test_an_unrelated_rule_survives_a_skipped_rewrite(import_activity: ModelImport):
client = import_activity.import_minio_repository.minio_client
existing = MagicMock()
matching_rule = MagicMock(rule_id='sientia-import-bundle-expiry', status=ENABLED)
matching_rule.expiration.days = 7
matching_rule.rule_filter.prefix = 'imported_models/'
unrelated_rule = MagicMock(rule_id='predictions-offload-expiry')
existing.rules = [unrelated_rule, matching_rule]
client.get_bucket_lifecycle.return_value = existing
result = await import_activity.ensure_import_bundle_retention(base())
client.set_bucket_lifecycle.assert_not_called()
assert result['rules'] == 2
def test_the_model_document_shape_is_documented_in_one_place():
"""A guard against the shape drifting between the code and what the runtime reads."""
source = Path('laborious/activities/model_import.py').read_text()

View File

@@ -21,6 +21,10 @@ from sqlalchemy.exc import DataError, IntegrityError, OperationalError, Programm
from laborious import metrics
from laborious.activities.model_import import (
_CLAIM_SQL, # noqa: PLC2701 - the SQL-shape test below asserts against these directly
_NAMES_SQL,
_PROGRESS_SQL,
_TERMINAL_SQL,
IMPORT_LOG_TABLE,
IMPORT_STATUS_HINT_RETRY_POLICY,
IMPORT_STATUS_RETRY_POLICY,
@@ -28,14 +32,18 @@ from laborious.activities.model_import import (
ExperimentRunStatus,
ImportState,
ModelImport,
_build_update,
build_import_detail,
derive_run_name,
map_state_to_status,
)
from laborious.utils.bundle.steps import ImportStep
from laborious.utils.models.import_errors import (
IMPORT_FAILURE_CATALOG,
UNCLASSIFIED_FAILURE,
ImportErrorCode,
ImportLogRowNotClaimableError,
ImportResponsibility,
ImportStatusRejectedError,
)
from tests.helpers.status_engine import FakeStatusEngine
@@ -93,6 +101,67 @@ EXPECTED_ERROR_STATUS = {
}
# The four statements as they were written out literally before `_build_update` existed. The test
# below proves the helper reproduces each one byte-for-byte (modulo whitespace).
_ORIGINAL_CLAIM_SQL = """
UPDATE public.experiment_run
SET status = :status,
updated_at = :updated_at,
orchestrator_response_data = CAST(:detail AS JSONB)
WHERE id = :import_run_id
AND run_type = 'IMPORT'
AND (status = 'PENDING'
OR (status = 'ORCHESTRATOR_WAITING_PROC'
AND orchestrator_response_data -> 'import' ->> 'workflow_id' = :workflow_id))
"""
_ORIGINAL_PROGRESS_SQL = """
UPDATE public.experiment_run
SET status = :status,
updated_at = :updated_at,
orchestrator_response_data = CAST(:detail AS JSONB)
WHERE id = :import_run_id
"""
_ORIGINAL_NAMES_SQL = """
UPDATE public.experiment_run
SET experiment_name = :experiment_name,
run_name = :run_name,
updated_at = :updated_at,
orchestrator_response_data = CAST(:detail AS JSONB)
WHERE id = :import_run_id
"""
_ORIGINAL_TERMINAL_SQL = """
UPDATE public.experiment_run
SET status = :status,
error_message = :error_message,
updated_at = :updated_at,
orchestrator_response_data = CAST(:detail AS JSONB)
WHERE id = :import_run_id
"""
def _normalized(sql: object) -> str:
return ' '.join(str(sql).split())
def test_build_update_reproduces_each_original_statement_byte_for_byte():
assert _normalized(_CLAIM_SQL) == _normalized(_ORIGINAL_CLAIM_SQL)
assert _normalized(_PROGRESS_SQL) == _normalized(_ORIGINAL_PROGRESS_SQL)
assert _normalized(_NAMES_SQL) == _normalized(_ORIGINAL_NAMES_SQL)
assert _normalized(_TERMINAL_SQL) == _normalized(_ORIGINAL_TERMINAL_SQL)
def test_build_update_with_no_casts_or_extra_where_renders_a_bare_update():
statement = _build_update({'status': 'status'})
assert (
_normalized(statement)
== 'UPDATE public.experiment_run SET status = :status WHERE id = :import_run_id'
)
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):
@@ -172,6 +241,27 @@ def test_an_unclassified_error_is_a_platform_error():
assert map_state_to_status(ImportState.ERROR) == ExperimentRunStatus.IMPORT_ERROR
@mark.parametrize('step', list(IMPORT_FAILURE_CATALOG))
def test_responsibility_comes_from_the_catalog_alone(step: ImportStep):
"""The regression test for `specs/model-import-status-record/spec.md`'s single-source rule."""
assert map_state_to_status(ImportState.ERROR, step) == EXPECTED_ERROR_STATUS[step]
def test_adding_a_step_to_the_catalog_is_enough_to_resolve_its_status():
"""A fake catalog entry needs nothing but its own `responsibility` to resolve correctly."""
fake_step = ImportStep.RECEIVED
for responsibility, expected in (
(ImportResponsibility.USER, ExperimentRunStatus.ORCHESTRATOR_VALIDATION_ERROR),
(ImportResponsibility.SUPPORT, ExperimentRunStatus.IMPORT_ERROR),
):
fake_catalog = {
**IMPORT_FAILURE_CATALOG,
fake_step: (ImportErrorCode.UNEXPECTED_ERROR, 'fake reason', responsibility),
}
with patch('laborious.utils.models.import_errors.IMPORT_FAILURE_CATALOG', fake_catalog):
assert map_state_to_status(ImportState.ERROR, fake_step) == expected
# ------------------------------------------------------------------------- the run name
@@ -583,6 +673,29 @@ async def test_a_failure_is_recorded_with_its_step_code_and_reason(
assert detail['gate'] == 4
@mark.asyncio
async def test_a_failure_with_an_unrecognized_step_is_recorded_unclassified(
import_activity: ModelImport, engine: FakeStatusEngine
):
"""A step string that is not an `ImportStep` member must not crash the one authoritative write.
It runs in the workflow's `finally`, after cleanup, so raising here would leave the row without
a verdict. The unrecognized step falls back to the unclassified failure instead.
"""
verdict = await import_activity.record_import_terminal_status(
{**claim_input(), 'succeeded': False, 'step': 'no_longer_a_step'}
)
write = engine.writes[0]
assert write.parameters['status'] == ExperimentRunStatus.IMPORT_ERROR.value
assert write.parameters['error_message'] == UNCLASSIFIED_FAILURE[1]
detail = detail_of(write)
assert detail['code'] == ImportErrorCode.UNEXPECTED_ERROR.value
assert verdict['status'] == ExperimentRunStatus.IMPORT_ERROR.value
assert verdict['code'] == ImportErrorCode.UNEXPECTED_ERROR.value
assert verdict['reason'] == UNCLASSIFIED_FAILURE[1]
@mark.asyncio
async def test_a_cleanup_failure_does_not_turn_a_success_into_an_error(
import_activity: ModelImport, engine: FakeStatusEngine

View File

@@ -9,6 +9,7 @@ behaviour is the point of not existing.
"""
import ast
import inspect
from pathlib import Path
from pytest import mark
@@ -182,15 +183,24 @@ def test_the_plaintext_password_field_is_refused_by_name():
def _error_message_values(tree: ast.Module) -> list[ast.expr]:
"""Every expression written into the record's `error_message`."""
"""Every expression written into the record's `error_message`, at an actual status write.
Scoped to `_execute_status_write`'s own parameter dict — not every dict in the module — because
`_build_update`'s `columns` argument also carries an `'error_message'` key (naming the column and
its bind parameter, not a value being written), and that is not a write this test is about.
"""
values = []
for node in ast.walk(tree):
if isinstance(node, ast.Dict):
for key, value in zip(node.keys, node.values, strict=True):
if not (
isinstance(node, ast.Call) and ast.unparse(node.func).endswith('_execute_status_write')
):
continue
for arg in node.args:
if not isinstance(arg, ast.Dict):
continue
for key, value in zip(arg.keys, arg.values, strict=True):
if isinstance(key, ast.Constant) and key.value == 'error_message':
values.append(value)
if isinstance(node, ast.keyword) and node.arg == 'error_message':
values.append(node.value)
return values
@@ -211,10 +221,12 @@ def test_the_reason_written_is_the_classified_one():
def test_the_classification_ignores_the_exception_by_construction():
"""`classify_import_failure` deletes its exception argument; nothing can leak from it."""
source = ERRORS_MODULE.read_text()
"""`classify_import_failure` has no exception parameter at all; nothing can leak from it."""
from laborious.utils.models.import_errors import classify_import_failure
assert 'del exception' in source
parameters = inspect.signature(classify_import_failure).parameters
assert 'exception' not in parameters
assert list(parameters) == ['step']
@mark.parametrize('module_path', (WORKFLOW_MODULE, ACTIVITIES_MODULE), ids=lambda p: p.name)

View File

@@ -367,6 +367,25 @@ def test_gate_4_rejects_an_entry_that_is_neither_a_file_nor_a_directory(work_dir
assert error.value.gate == 4
def test_gate_4_and_gate_5_read_the_central_directory_once(
golden_bundle, golden_password, work_dir, reader, monkeypatch
):
"""Achado #11: gate 5 extracts using the entries gate 4 already read, instead of reopening the
archive to call `infolist()` a second time."""
original_infolist = zipfile.ZipFile.infolist
calls: list[None] = []
def _counting_infolist(self):
calls.append(None)
return original_infolist(self)
monkeypatch.setattr(zipfile.ZipFile, 'infolist', _counting_infolist)
reader.open(golden_bundle, golden_password, work_dir / 'extract', work_dir)
assert len(calls) == 1
# ------------------------------------------------------------------- gate 5
@@ -580,6 +599,21 @@ def test_gate_7_rejects_a_non_regular_file_left_in_the_tree(work_dir, reader):
assert error.value.gate == 7
def test_gate_7_rejects_disallowed_content_regardless_of_traversal_order(work_dir, reader):
"""Achado #12: gate 7 no longer sorts the tree before walking it, so a disallowed file that
would not be first or last alphabetically must still be caught."""
root = work_dir / 'unsorted_root'
root.mkdir()
for name in ('a_file.txt', 'b_file.txt', 'y_file.txt', 'z_file.txt'):
(root / name).write_bytes(b'ok')
(root / 'n_disallowed.sh').write_bytes(b'#!/bin/sh\n')
with pytest.raises(BundleGateError, match='should not carry') as error:
reader._gate_7_content_policy(root)
assert error.value.gate == 7
# ------------------------------------------------------------------ ordering

View File

@@ -34,6 +34,7 @@ from laborious.utils.models.import_errors import (
failure_location,
report_import_failure,
summarize_failure,
walk_cause_chain,
)
REPOSITORY_ROOT = Path(__file__).resolve().parents[4]
@@ -137,12 +138,12 @@ def test_every_gate_maps_to_a_step_that_has_a_sentence():
def test_each_step_has_its_own_code():
codes = [code for code, _ in IMPORT_FAILURE_CATALOG.values()]
codes = [code for code, _, _ in IMPORT_FAILURE_CATALOG.values()]
assert len(codes) == len(set(codes))
def test_the_codes_that_never_reach_a_record_are_absent_from_the_catalog():
catalogued = {code for code, _ in IMPORT_FAILURE_CATALOG.values()}
catalogued = {code for code, _, _ in IMPORT_FAILURE_CATALOG.values()}
for code in (
ImportErrorCode.LOG_ROW_NOT_CLAIMABLE,
ImportErrorCode.STATUS_NOT_RECORDED,
@@ -153,7 +154,7 @@ def test_the_codes_that_never_reach_a_record_are_absent_from_the_catalog():
@mark.parametrize('step', list(IMPORT_FAILURE_CATALOG))
def test_a_reason_never_carries_the_exception(step: ImportStep):
failure = classify_import_failure(step, RuntimeError(LEAKY_MESSAGE))
failure = classify_import_failure(step)
for leak in LEAKS_THAT_MUST_NOT_SURFACE:
assert leak not in failure.reason
@@ -162,12 +163,11 @@ def test_a_reason_never_carries_the_exception(step: ImportStep):
@mark.parametrize('step', list(IMPORT_FAILURE_CATALOG))
def test_the_classification_is_a_function_of_the_step_alone(step: ImportStep):
"""Which is what makes a reworded sentence keep its code: neither is read from the exception."""
first = classify_import_failure(step, RuntimeError(LEAKY_MESSAGE))
second = classify_import_failure(step, KeyError('target_variable'))
third = classify_import_failure(step, None)
"""Classification does not read the exception at all — it was already dropped as a parameter."""
first = classify_import_failure(step)
second = classify_import_failure(step)
assert first == second == third
assert first == second
assert first.code is IMPORT_FAILURE_CATALOG[step][0]
@@ -187,7 +187,7 @@ def test_a_reason_the_user_cannot_act_on_names_the_import_identifier(step: Impor
def test_a_failure_with_no_step_is_unclassified():
failure = classify_import_failure(None, RuntimeError('boom'))
failure = classify_import_failure(None)
assert failure.step is None
assert failure.code is ImportErrorCode.UNEXPECTED_ERROR
@@ -196,7 +196,7 @@ def test_a_failure_with_no_step_is_unclassified():
def test_every_reason_is_long_enough_for_the_column():
"""`error_message` carries a `LENGTH(...) >= 3` CHECK; every sentence here clears it."""
for _, reason in IMPORT_FAILURE_CATALOG.values():
for _, reason, _ in IMPORT_FAILURE_CATALOG.values():
assert len(reason) >= 3
@@ -343,8 +343,13 @@ async def test_a_report_without_an_exception_object_says_so():
@mark.asyncio
async def test_an_explicit_code_keeps_the_step_s_sentence():
"""The cleanup case: its own code, and a sentence a person can still read."""
async def test_an_explicit_code_gets_its_own_sentence_not_the_unclassified_one():
"""The cleanup case: its own code, and a sentence a person can still read.
A cleanup-only failure means the model *was* imported — only its temporary files were left
behind. `UNCLASSIFIED_FAILURE`'s sentence ("no model was created") would be false here, so the
override in `_OVERRIDE_REASONS` must be what is reported, not the catalog's fallback.
"""
sinks = FakeSinks()
failure = await report_import_failure(
@@ -358,10 +363,64 @@ async def test_an_explicit_code_keeps_the_step_s_sentence():
)
assert failure.code is ImportErrorCode.CLEANUP_INCOMPLETE
assert failure.reason == UNCLASSIFIED_FAILURE[1]
assert failure.reason != UNCLASSIFIED_FAILURE[1]
assert 'no model was created' not in failure.reason
assert 'unexpected reason' not in failure.reason
assert 'imported successfully' in failure.reason
assert sinks.notifications[0]['message'] == failure.reason
assert sinks.notifications[0]['level'] == 'ERROR'
@mark.asyncio
async def test_an_unclaimable_row_gets_its_own_reason_not_the_generic_one():
"""`LOG_ROW_NOT_CLAIMABLE` never attempted the pipeline, so it must not say "no model was
created" either — it gets its own sentence about the request never starting."""
sinks = FakeSinks()
failure = await report_import_failure(
sinks,
None,
ImportLogRowNotClaimableError('row 29 is in status IMPORT_ERROR', import_run_id=29),
import_run_id=29,
workflow_id='manual-import-29',
code=ImportErrorCode.LOG_ROW_NOT_CLAIMABLE,
marker='not_claimable',
)
assert failure.code is ImportErrorCode.LOG_ROW_NOT_CLAIMABLE
assert failure.reason != UNCLASSIFIED_FAILURE[1]
assert 'no model was created' not in failure.reason
assert 'unexpected reason' not in failure.reason
assert failure.reason == 'The import request could not be started.'
assert sinks.notifications[0]['message'] == failure.reason
@mark.asyncio
async def test_a_status_write_failure_gets_its_own_reason_not_the_generic_one():
"""`STATUS_NOT_RECORDED` is a terminal-write failure, not a stalled pipeline — its sentence
says the outcome could not be recorded, not that nothing was created."""
sinks = FakeSinks()
failure = await report_import_failure(
sinks,
None,
None,
import_run_id=77,
workflow_id='wf',
code=ImportErrorCode.STATUS_NOT_RECORDED,
marker='status_write',
verdict='succeeded=True step=model_document gate=None',
detail_text='OperationalError: connection refused',
)
assert failure.code is ImportErrorCode.STATUS_NOT_RECORDED
assert failure.reason != UNCLASSIFIED_FAILURE[1]
assert 'no model was created' not in failure.reason
assert 'unexpected reason' not in failure.reason
assert failure.reason == 'The import outcome could not be recorded.'
assert sinks.notifications[0]['message'] == failure.reason
# --------------------------------------------- the fields the two log lines carry
#
# The message text is for a person reading a log; these fields are for the query that finds it. Both
@@ -443,7 +502,7 @@ async def test_the_user_facing_line_is_greppable_by_the_same_fields():
)
notified = sinks.notifications[0]
assert notified['message'] == UNCLASSIFIED_FAILURE[1]
assert notified['message'] == 'The import request could not be started.'
assert notified['metadata']['import_code'] == ImportErrorCode.LOG_ROW_NOT_CLAIMABLE.value
assert notified['metadata']['import_step'] == 'not_claimable'
assert notified['metadata']['import_run_id'] == 29
@@ -498,6 +557,31 @@ async def test_no_field_can_shadow_the_formatter_s_own_keys():
assert not {'at', 'msg', 'timestamp'} & set(fields)
def test_walk_cause_chain_yields_just_the_exception_when_there_is_no_cause():
error = RuntimeError('alone')
assert list(walk_cause_chain(error)) == [error]
def test_walk_cause_chain_yields_the_whole_chain_when_shorter_than_the_limit():
root = RuntimeError('root')
middle = RuntimeError('middle')
middle.__cause__ = root
outer = RuntimeError('outer')
outer.__cause__ = middle
assert list(walk_cause_chain(outer, limit=10)) == [outer, middle, root]
def test_walk_cause_chain_stops_at_the_limit_for_a_chain_longer_than_it():
first = RuntimeError('first')
second = RuntimeError('second')
first.__cause__ = second
second.__cause__ = first
assert list(walk_cause_chain(first, limit=3)) == [first, second, first]
def test_the_chain_is_flattened_root_cause_first():
"""Because the first line is what becomes `error_summary`, and it has to be the cause."""
kind, detail = describe_failure_chain(_wrapped_activity_failure())

View File

@@ -1978,6 +1978,28 @@ async def test_download_model_wrapper_failure_propagates(joblib_mock, mlflow, ml
mlflow_repository.client.download_artifacts.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_wrapper_failure_emits_error_metric(
joblib_mock, mlflow, mlflow_repository
):
mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
mlflow_repository.dowload_artifacts = AsyncMock()
loader_error = mlflow_lib.exceptions.MlflowException('wrapper boom')
mlflow.pyfunc.load_model.side_effect = loader_error
with pytest.raises(mlflow_lib.exceptions.MlflowException) as exc_info:
await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)
assert exc_info.value is loader_error
joblib_mock.load.assert_not_called()
mlflow_repository.client.download_artifacts.assert_not_called()
error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT)
assert len(error_tags) == 1
assert error_tags[0]['operation_type'] == 'download_model_wrapper'
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_predict_model_joblib_success_metrics(

View File

@@ -173,6 +173,7 @@ def test_build_import_config_with_defaults(clean_import_env):
'max_uncompressed_bytes': 4 * 1024 * 1024 * 1024,
'max_compression_ratio': 200.0,
'models_collection': 'models',
'counters_collection': 'model_id_counters',
'retention_days': 7,
}
@@ -187,6 +188,7 @@ def test_build_import_config_reads_the_environment(clean_import_env):
clean_import_env.setenv('IMPORT_MAX_UNCOMPRESSED_BYTES', '4096')
clean_import_env.setenv('IMPORT_MAX_COMPRESSION_RATIO', '12.5')
clean_import_env.setenv('IMPORT_MODELS_COLLECTION', 'models_v2')
clean_import_env.setenv('IMPORT_MODEL_ID_COUNTERS_COLLECTION', 'model_id_counters_v2')
clean_import_env.setenv('IMPORT_BUNDLE_RETENTION_DAYS', '14')
config = build_import_config()
@@ -202,6 +204,7 @@ def test_build_import_config_reads_the_environment(clean_import_env):
assert config['max_uncompressed_bytes'] == 4096
assert config['max_compression_ratio'] == 12.5
assert config['models_collection'] == 'models_v2'
assert config['counters_collection'] == 'model_id_counters_v2'
assert config['retention_days'] == 14

View File

@@ -0,0 +1,19 @@
from datetime import UTC, datetime
from laborious.utils.time import to_naive
def test_to_naive_returns_naive_input_unchanged() -> None:
moment = datetime(2024, 1, 1, 12, 30, 45)
assert to_naive(moment) == moment
assert to_naive(moment).tzinfo is None
def test_to_naive_strips_tzinfo_from_aware_input() -> None:
moment = datetime(2024, 1, 1, 12, 30, 45, tzinfo=UTC)
result = to_naive(moment)
assert result.tzinfo is None
assert result == datetime(2024, 1, 1, 12, 30, 45)

View File

@@ -109,6 +109,6 @@ def test_the_import_worker_registers_only_the_import_activities():
if isinstance(inner, ast.keyword) and inner.arg == 'other_workflows':
assert ast.unparse(inner.value) == '[]'
assert len(activities) == 14
assert len(activities) == 11
assert all(name.startswith('activities.') for name in activities)
assert all('import' in name for name in activities)

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,