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,