From 192834bde3a54a174348949e4f8f57089244c27e Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 8 Sep 2026 17:06:25 -0300 Subject: [PATCH] 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. --- docs/minio-object-expiry.md | 40 ++ e2e/import_harness.py | 21 +- e2e/test_import_model_digest.py | 2 +- e2e/test_import_model_gates.py | 2 +- e2e/test_import_model_happy_path.py | 1 + ...test_import_model_provisioning_failures.py | 12 +- e2e/test_import_model_second_import.py | 4 +- e2e/test_import_model_status_record.py | 28 +- laborious/activities/activities.py | 22 +- laborious/activities/mlflow.py | 37 +- laborious/activities/model_import.py | 394 ++++++++++++------ laborious/activities/storage.py | 15 +- laborious/utils/bundle/reader.py | 18 +- laborious/utils/connectors_config.py | 8 + laborious/utils/messages.py | 1 + laborious/utils/models/import_errors.py | 106 +++-- .../utils/repository/model_repository.py | 19 +- laborious/utils/repository/opc_repository.py | 9 +- laborious/utils/time.py | 6 + laborious/worker/worker.py | 8 +- laborious/workflows/import_model.py | 166 ++++---- tests/laborious/activities/test_activities.py | 70 +++- tests/laborious/activities/test_mlflow.py | 35 +- .../activities/test_model_import_bundle.py | 183 ++++++-- .../test_model_import_provisioning.py | 314 +++++++++++--- .../activities/test_model_import_status.py | 113 +++++ tests/laborious/test_import_absence.py | 28 +- tests/laborious/utils/bundle/test_reader.py | 34 ++ .../utils/models/test_import_errors.py | 112 ++++- .../utils/repository/test_model_repository.py | 22 + .../laborious/utils/test_connectors_config.py | 3 + tests/laborious/utils/test_time.py | 19 + .../worker/test_runtime_task_queues.py | 2 +- .../laborious/workflows/test_import_model.py | 162 +++++-- 34 files changed, 1485 insertions(+), 531 deletions(-) create mode 100644 docs/minio-object-expiry.md create mode 100644 laborious/utils/messages.py create mode 100644 laborious/utils/time.py create mode 100644 tests/laborious/utils/test_time.py diff --git a/docs/minio-object-expiry.md b/docs/minio-object-expiry.md new file mode 100644 index 0000000..ebac303 --- /dev/null +++ b/docs/minio-object-expiry.md @@ -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. diff --git a/e2e/import_harness.py b/e2e/import_harness.py index 8d4b99d..48c962f 100644 --- a/e2e/import_harness.py +++ b/e2e/import_harness.py @@ -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 diff --git a/e2e/test_import_model_digest.py b/e2e/test_import_model_digest.py index d3a6ee0..62c5c51 100644 --- a/e2e/test_import_model_digest.py +++ b/e2e/test_import_model_digest.py @@ -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 diff --git a/e2e/test_import_model_gates.py b/e2e/test_import_model_gates.py index fd3e29d..0739e22 100644 --- a/e2e/test_import_model_gates.py +++ b/e2e/test_import_model_gates.py @@ -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}' diff --git a/e2e/test_import_model_happy_path.py b/e2e/test_import_model_happy_path.py index 30bf253..293c586 100644 --- a/e2e/test_import_model_happy_path.py +++ b/e2e/test_import_model_happy_path.py @@ -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) diff --git a/e2e/test_import_model_provisioning_failures.py b/e2e/test_import_model_provisioning_failures.py index 44f1263..5853d97 100644 --- a/e2e/test_import_model_provisioning_failures.py +++ b/e2e/test_import_model_provisioning_failures.py @@ -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})) diff --git a/e2e/test_import_model_second_import.py b/e2e/test_import_model_second_import.py index 78ffddf..7580315 100644 --- a/e2e/test_import_model_second_import.py +++ b/e2e/test_import_model_second_import.py @@ -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. diff --git a/e2e/test_import_model_status_record.py b/e2e/test_import_model_status_record.py index b46238e..9e93ab8 100644 --- a/e2e/test_import_model_status_record.py +++ b/e2e/test_import_model_status_record.py @@ -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'] diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index a06eb16..1e09238 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -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, diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index d075508..f6901e1 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -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'] diff --git a/laborious/activities/model_import.py b/laborious/activities/model_import.py index 08015ae..1a6f547 100644 --- a/laborious/activities/model_import.py +++ b/laborious/activities/model_import.py @@ -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 # 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 WHERE id = :import_run_id [AND ]`. + + `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 )` 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( diff --git a/laborious/activities/storage.py b/laborious/activities/storage.py index 7c67466..7bafbe2 100644 --- a/laborious/activities/storage.py +++ b/laborious/activities/storage.py @@ -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': {}, diff --git a/laborious/utils/bundle/reader.py b/laborious/utils/bundle/reader.py index 61fcd90..aea9646 100644 --- a/laborious/utils/bundle/reader.py +++ b/laborious/utils/bundle/reader.py @@ -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(): diff --git a/laborious/utils/connectors_config.py b/laborious/utils/connectors_config.py index 4b31496..01e29e1 100644 --- a/laborious/utils/connectors_config.py +++ b/laborious/utils/connectors_config.py @@ -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')), } diff --git a/laborious/utils/messages.py b/laborious/utils/messages.py new file mode 100644 index 0000000..1b84ea3 --- /dev/null +++ b/laborious/utils/messages.py @@ -0,0 +1 @@ +MISSING_MINIO_REPOSITORY = 'Minio repository not initialized' diff --git a/laborious/utils/models/import_errors.py b/laborious/utils/models/import_errors.py index 8312bd4..8d3c88c 100644 --- a/laborious/utils/models/import_errors.py +++ b/laborious/utils/models/import_errors.py @@ -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: diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 15b4a10..1c32baf 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -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 diff --git a/laborious/utils/repository/opc_repository.py b/laborious/utils/repository/opc_repository.py index 2ecc27a..802c318 100644 --- a/laborious/utils/repository/opc_repository.py +++ b/laborious/utils/repository/opc_repository.py @@ -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__ diff --git a/laborious/utils/time.py b/laborious/utils/time.py new file mode 100644 index 0000000..eebe7c5 --- /dev/null +++ b/laborious/utils/time.py @@ -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 diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index f4dbf4d..09cc2c8 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -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, diff --git a/laborious/workflows/import_model.py b/laborious/workflows/import_model.py index b8aadd2..3a81df7 100644 --- a/laborious/workflows/import_model.py +++ b/laborious/workflows/import_model.py @@ -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], diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index dd07c1f..44dc194 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -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, diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index da17c94..e87c618 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -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, diff --git a/tests/laborious/activities/test_model_import_bundle.py b/tests/laborious/activities/test_model_import_bundle.py index edf4018..ba9fd4d 100644 --- a/tests/laborious/activities/test_model_import_bundle.py +++ b/tests/laborious/activities/test_model_import_bundle.py @@ -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', + ] diff --git a/tests/laborious/activities/test_model_import_provisioning.py b/tests/laborious/activities/test_model_import_provisioning.py index fe5d11f..e025898 100644 --- a/tests/laborious/activities/test_model_import_provisioning.py +++ b/tests/laborious/activities/test_model_import_provisioning.py @@ -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() diff --git a/tests/laborious/activities/test_model_import_status.py b/tests/laborious/activities/test_model_import_status.py index 2f24080..0233406 100644 --- a/tests/laborious/activities/test_model_import_status.py +++ b/tests/laborious/activities/test_model_import_status.py @@ -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 diff --git a/tests/laborious/test_import_absence.py b/tests/laborious/test_import_absence.py index 686d9c5..b0eab28 100644 --- a/tests/laborious/test_import_absence.py +++ b/tests/laborious/test_import_absence.py @@ -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) diff --git a/tests/laborious/utils/bundle/test_reader.py b/tests/laborious/utils/bundle/test_reader.py index ebff1bf..e912ad5 100644 --- a/tests/laborious/utils/bundle/test_reader.py +++ b/tests/laborious/utils/bundle/test_reader.py @@ -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 diff --git a/tests/laborious/utils/models/test_import_errors.py b/tests/laborious/utils/models/test_import_errors.py index 53472fa..173547f 100644 --- a/tests/laborious/utils/models/test_import_errors.py +++ b/tests/laborious/utils/models/test_import_errors.py @@ -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()) diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index d576b5f..a6ea506 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -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( diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index 08639c1..015b99b 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -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 diff --git a/tests/laborious/utils/test_time.py b/tests/laborious/utils/test_time.py new file mode 100644 index 0000000..1feb957 --- /dev/null +++ b/tests/laborious/utils/test_time.py @@ -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) diff --git a/tests/laborious/worker/test_runtime_task_queues.py b/tests/laborious/worker/test_runtime_task_queues.py index 3dff897..6dcea16 100644 --- a/tests/laborious/worker/test_runtime_task_queues.py +++ b/tests/laborious/worker/test_runtime_task_queues.py @@ -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) diff --git a/tests/laborious/workflows/test_import_model.py b/tests/laborious/workflows/test_import_model.py index 8d05c5c..1e68237 100644 --- a/tests/laborious/workflows/test_import_model.py +++ b/tests/laborious/workflows/test_import_model.py @@ -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,