diff --git a/e2e/fixtures/experiment_run.sql b/e2e/fixtures/experiment_run.sql new file mode 100644 index 0000000..d43cfd1 --- /dev/null +++ b/e2e/fixtures/experiment_run.sql @@ -0,0 +1,67 @@ +-- public.experiment_run as migrations V1 … V12 leave it, for the e2e Postgres container. +-- +-- Source of truth: sientia-dataops-database-migrations, +-- postgres/sientia-core-mlops-bff/migrations/ — this file mirrors the *final state* of: +-- +-- V1 create_table_experiment_run (13 columns, one index) +-- V2 update_check_constraint_on_status (named the status CHECK) +-- V3 update_check_constraints (named every other CHECK) +-- V4 update_status_column_size_and_add_orchestrator_validation_error (status -> VARCHAR(50)) +-- V5 migrate_mage_statuses_to_orchestrator (data only) +-- V6 rename_mage_response_data_column (-> orchestrator_response_data) +-- V7 migrate_mlflow_statuses_to_tracking (data only) +-- V8 remove_orchestrator_request_sent_status (status vocabulary) +-- V9 remove_file_upload_statuses (status vocabulary) +-- V10 remove_unused_statuses (status vocabulary) +-- V11 add_run_type_to_experiment_run (run_type + reindexed) +-- V12 add_import_statuses_to_experiment_run (IMPORT_SUCCESS / IMPORT_ERROR + pairing) +-- +-- The same drift trade already accepted for the bundle fixtures: this is a copy, and the tests +-- that read it are what make the copy auditable. If a migration lands there, update this file and +-- the version list above. Nothing in this repository owns the DDL — no import ever executes it +-- outside the test suite. + +CREATE TABLE IF NOT EXISTS public.experiment_run ( + id BIGSERIAL CONSTRAINT pk_experiment_run PRIMARY KEY, + experiment_name VARCHAR(50) NOT NULL, + run_name VARCHAR(50), + username VARCHAR(320) NOT NULL, + status VARCHAR(50) NOT NULL, + error_message VARCHAR(1024), + created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + updated_at TIMESTAMP WITHOUT TIME ZONE NOT NULL, + bucket_name VARCHAR(63), + file_name VARCHAR(1024), + request_data JSONB NOT NULL, + orchestrator_response_data JSONB, + run_type VARCHAR(20) NOT NULL DEFAULT 'TRAIN', + + CONSTRAINT chk_experiment_run_experiment_name_min_length CHECK (LENGTH(experiment_name) >= 3), + CONSTRAINT chk_experiment_run_run_name_min_length CHECK (LENGTH(run_name) >= 3), + CONSTRAINT chk_experiment_run_username_min_length CHECK (LENGTH(username) >= 3), + CONSTRAINT chk_experiment_run_username_pattern + CHECK (username ~* '^[a-zA-Z0-9_!#$%&''*+/=?`{|}~^.-]+@[a-zA-Z0-9.-]+$'), + CONSTRAINT chk_experiment_run_error_message_min_length CHECK (LENGTH(error_message) >= 3), + CONSTRAINT chk_experiment_run_bucket_name_min_length CHECK (LENGTH(bucket_name) >= 3), + CONSTRAINT chk_experiment_run_bucket_name_pattern CHECK (bucket_name ~* '^[a-z0-9-]+$'), + CONSTRAINT chk_experiment_run_file_name_min_length CHECK (LENGTH(file_name) >= 3), + CONSTRAINT chk_experiment_run_file_name_pattern + CHECK (file_name ~* '^[a-zA-Z0-9 /._-]+$'), + + -- V12: the two shared stage values plus the training and import terminals. + CONSTRAINT chk_experiment_run_status CHECK (status IN ('PENDING', 'ORCHESTRATOR_REQUEST_ERROR', + 'ORCHESTRATOR_WAITING_PROC', + 'ORCHESTRATOR_VALIDATION_ERROR', + 'TRAINING_SUCCESS', 'TRAINING_ERROR', + 'IMPORT_SUCCESS', 'IMPORT_ERROR')), + CONSTRAINT chk_experiment_run_run_type CHECK (run_type IN ('TRAIN', 'IMPORT')), + + -- V12: each pair of terminal values belongs to its own run_type. The shared values stay out. + CONSTRAINT chk_experiment_run_terminal_status_run_type CHECK ( + (run_type = 'IMPORT' OR status NOT IN ('IMPORT_SUCCESS', 'IMPORT_ERROR')) + AND (run_type = 'TRAIN' OR status NOT IN ('TRAINING_SUCCESS', 'TRAINING_ERROR')) + ) +); + +CREATE INDEX IF NOT EXISTS idx_experiment_run_user_exp_status_date_type + ON public.experiment_run (username, experiment_name, status, created_at, run_type); diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index df78a08..f4a9a2c 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -7,16 +7,22 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.observability.logger import Logger from sientia_do.repository.minio_repository import MinioRepository + from sientia_do.repository.mongodb_repository import MongoDBRepository from laborious.activities.api import API from laborious.activities.gates import Gates from laborious.activities.mlflow import MLFlow + from laborious.activities.model_import import ModelImport from laborious.activities.model_metrics import ModelMetrics from laborious.activities.opc import OPC from laborious.activities.storage import Storage + from laborious.utils.connectors_config import ( + build_import_config, + build_import_status_config, + ) -class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): +class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport): """ Main activities orchestrator for the Laborious system. @@ -50,6 +56,9 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): pi_web_api_config: dict[str, Any], logger: Logger, notification_handler: NotificationHandler, + mongo_config: dict[str, Any] | None = None, + import_config: dict[str, Any] | None = None, + import_status_config: dict[str, Any] | None = None, ): """ Initialize the Activities orchestrator with all required configurations. @@ -68,6 +77,13 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): Required keys: base_url, auth_type, auth_token logger: Logger instance for observability and debugging notification_handler: Notification handler for alerts and monitoring + mongo_config: MongoDB configuration for the model listing the import writes. When + absent, the import activities that need it refuse rather than guess — the + notification handler keeps its own client either way. + import_config: `build_import_config()`; read from the environment when absent + import_status_config: `build_import_status_config()`; read from the environment when + absent. This is the BFF database holding the import log, a second Postgres + connection that is never the one Storage holds. Raises: Exception: If any parent class initialization fails @@ -147,6 +163,32 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): metrics_controller=metrics_controller, ) + mongo_repository = ( + MongoDBRepository( + connection_string=mongo_config['connection_string'], + database_name=mongo_config['database_name'], + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + if mongo_config + else None + ) + + ModelImport.__init__( + self, + 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. + mlflow_repository=self.model_monitoring_repository, + mongo_repository=mongo_repository, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + async def shutdown(self): """ Gracefully shutdown all activities and clean up resources. @@ -167,3 +209,4 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): await OPC.aclose(self) ModelMetrics.close(self) API.close(self) + ModelImport.close(self) diff --git a/laborious/activities/model_import.py b/laborious/activities/model_import.py new file mode 100644 index 0000000..125ee97 --- /dev/null +++ b/laborious/activities/model_import.py @@ -0,0 +1,1044 @@ +"""Model import activities: the import log row, the bundle, MLflow, the model document, cleanup. + +The import log row is **not created here**. It is `public.experiment_run` in the Postgres database +`sientia-core-mlops-bff`; the frontend inserts it through the BFF with `run_type='IMPORT'` and +`status='PENDING'`, and passes its `id` in the workflow input. Every write in this module is an +`UPDATE ... WHERE id = :import_run_id`, on an engine of its own against that database — never the +`sientia` database Laborious itself uses, and never in a transaction shared with a provisioning +write. + +Three classes of status write, and only the first one is strict: + +- **A, the claim** — `WHERE id AND run_type='IMPORT' AND (PENDING OR my own workflow id)`. This is + the ownership decision, and a claim matching nothing fails the import explicitly. +- **B, progress** — `WHERE id`, nothing else. A hint. Its failure never fails the import. +- **C, terminal** — `WHERE id`, nothing else, in the workflow's `finally`, after cleanup. The + authority: it rewrites status, code, reason and step from one classification of the outcome. +""" + +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + import json + import shutil + from datetime import timedelta + from enum import StrEnum + from pathlib import Path + from typing import Any, Final + + from minio.commonconfig import ENABLED, Filter + from minio.lifecycleconfig import Expiration, LifecycleConfig, Rule + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.observability.logger import Logger + from sientia_do.observability.metrics_controller import MetricsController + from sientia_do.observability.sientia_monitoring import SientiaMonitoring + from sientia_do.repository.minio_repository import MinioRepository + from sientia_do.repository.mongodb_repository import MongoDBRepository + from sientia_do.temporal.constants import now + from sqlalchemy import create_engine, text + from sqlalchemy.engine import URL + from sqlalchemy.exc import DataError, IntegrityError, ProgrammingError + from temporalio.common import RetryPolicy + + from laborious import metrics + from laborious.activities.model_import_errors import ( + ImportErrorCode, + ImportGateRejectedError, + ImportLogRowNotClaimableError, + ImportStatusRejectedError, + classify_import_failure, + report_import_failure, + ) + from laborious.utils.bundle import format as fmt + from laborious.utils.bundle.reader import BundleGateError, BundleLimits, BundleReader + from laborious.utils.bundle.steps import ImportStep + from laborious.utils.import_password import ( + EnvelopeError, + decrypt_password_envelope, + load_envelope_key, + ) + from laborious.utils.repository.model_repository import MLFlowRepository + +# 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 +# statements below write the name out literally rather than interpolating this constant — an +# interpolation is a shape a reader (and bandit) has to reason about, and there is nothing to +# parameterise. The constant is what the tests assert every statement against. +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' + +# Where a digest may sit inside the BFF-owned `request_data`. Checked in this order; absence is not +# a failure, only a disagreement is. +REQUEST_DIGEST_KEYS: Final[tuple[str, ...]] = ( + 'expected_digest', + 'digest', + 'sha256', + 'file_digest', +) + + +class ImportState(StrEnum): + """The workflow's internal vocabulary, mapped onto the shared status column by D17.""" + + RUNNING = 'RUNNING' + COMPLETED = 'COMPLETED' + ERROR = 'ERROR' + + +class ExperimentRunStatus(StrEnum): + """The values this workflow writes into `experiment_run.status`. + + Two are **reused** from the training flow because they describe a stage, not a flow — the row's + flow is told by `run_type`, not by the status. Two are **new in V12**, because + `TRAINING_SUCCESS` / `TRAINING_ERROR` would lie about where the row came from. Nobody should + "fix" the reuse into an `IMPORT_RUNNING` twin: it would carry no information `run_type` does not + already carry, and it would widen the pairing CHECK and the frontend's switch for nothing. + """ + + ORCHESTRATOR_WAITING_PROC = 'ORCHESTRATOR_WAITING_PROC' # reused + ORCHESTRATOR_VALIDATION_ERROR = 'ORCHESTRATOR_VALIDATION_ERROR' # reused + IMPORT_SUCCESS = 'IMPORT_SUCCESS' # new in V12 + 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 +# other repositories share. +IMPORT_STATUS_RETRY_POLICY: Final[RetryPolicy] = RetryPolicy( + maximum_attempts=10, + initial_interval=timedelta(seconds=1), + backoff_coefficient=2.0, + maximum_interval=timedelta(seconds=30), + non_retryable_error_types=['ImportLogRowNotClaimableError', 'ImportStatusRejectedError'], +) + +# Progress writes are hints nobody reads — the terminal write rewrites them — so a 151-second pause +# before each of a dozen stage updates would stall a working import for twenty minutes because of a +# database the import itself does not need. +IMPORT_STATUS_HINT_RETRY_POLICY: Final[RetryPolicy] = RetryPolicy( + maximum_attempts=3, + initial_interval=timedelta(seconds=1), + backoff_coefficient=2.0, + maximum_interval=timedelta(seconds=2), + 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, + run_type, + status, + bucket_name, + file_name, + request_data, + orchestrator_response_data -> 'import' ->> 'workflow_id' AS claimed_by + FROM public.experiment_run + WHERE id = :import_run_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 + """ +) + +_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 + """ +) + +_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 + """ +) + + +def map_state_to_status( + state: ImportState, step: ImportStep | str | None = None +) -> ExperimentRunStatus: + """Map the workflow's internal state onto the shared status column (design D17). + + `RUNNING` is the orchestrator holding the row. `COMPLETED` is `IMPORT_SUCCESS`. An `ERROR` is + split in two by *who can act on it*: the user, for everything up to the content policy, or only + support, for everything after it — the same split the table already draws between + `ORCHESTRATOR_VALIDATION_ERROR` and `TRAINING_ERROR`. + """ + if state is ImportState.RUNNING: + 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: + return ExperimentRunStatus.ORCHESTRATOR_VALIDATION_ERROR + return ExperimentRunStatus.IMPORT_ERROR + + +def derive_run_name(model_name: str, model_version: str) -> str: + """`import-{model_name}-v{model_version}`, truncated inside the model-name segment. + + Deterministic, because the recorded name and the created run's name must not diverge: the run is + created with exactly this value, passed in as an argument rather than recomputed. The `import-` + prefix guarantees the column's length-3 floor; the truncation is what keeps it under 50. + """ + suffix = f'-v{model_version}' + room = RUN_NAME_MAX_LENGTH - len(RUN_NAME_PREFIX) - len(suffix) + if room < 1: + # A version string long enough to leave no room for a name: keep the prefix and as much of + # the version as fits, so the value is still a legible, deterministic 50 characters. + return f'{RUN_NAME_PREFIX}{suffix.lstrip("-")}'[:RUN_NAME_MAX_LENGTH] + return f'{RUN_NAME_PREFIX}{model_name[:room]}{suffix}' + + +def build_import_detail( + *, + workflow_id: str, + step: ImportStep | str | None = None, + code: ImportErrorCode | str | None = None, + gate: int | None = None, + source: dict[str, Any] | None = None, + cleanup_failed: bool = False, +) -> str: + """Build `orchestrator_response_data` as JSON text. + + The whole `import` object is rewritten on every write, which is what keeps the workflow + identifier the claim wrote from ever being dropped. There is no column for any of this (design + D21): the code the frontend switches on and the identifier support quotes live at fixed keys in + the JSON column the workflow already owns. + """ + detail: dict[str, Any] = {'workflow_id': workflow_id, 'gate': gate} + if step is not None: + detail['step'] = str(step) + if code is not None: + detail['code'] = str(code) + if source: + detail['source'] = source + if cleanup_failed: + detail['cleanup_failed'] = True + return json.dumps({'import': detail}) + + +class ModelImport(SientiaMonitoring): + """Activities turning a `.sientia` object into a registered model, or into a durable ERROR.""" + + def __init__( + self, + status_config: dict[str, Any], + import_config: dict[str, Any], + minio_repository: MinioRepository | None = None, + mlflow_repository: MLFlowRepository | None = None, + mongo_repository: MongoDBRepository | None = None, + logger: Logger | None = None, + notification_handler: NotificationHandler | None = None, + metrics_controller: MetricsController | None = None, + ): + """ + Args: + status_config: `build_import_status_config()` — the BFF database holding the import log. + This is a **second** Postgres connection, independent of the one Laborious already + holds, which is what makes the transactional independence of design D7 literal. + import_config: `build_import_config()` — bucket, prefix, ceilings, work directory, + envelope key and the models collection name. + minio_repository: the shared MinIO repository, for the uploaded bundle. + mlflow_repository: the shared MLflow repository — the same instance the rest of the + worker uses, because setting the tracking URI is a process-wide side effect. + mongo_repository: MongoDB, where the model listing lives. + logger: logger instance. + notification_handler: the notification path the rest of the repository uses. + metrics_controller: metrics controller. + """ + SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + self.import_config = import_config + self.status_config = status_config + self.import_minio_repository = minio_repository + self.import_mlflow_repository = mlflow_repository + self.mongo_repository = mongo_repository + self.status_engine = create_engine( + URL.create( + 'postgresql', + username=status_config['user'], + password=status_config['password'], + host=status_config['host'], + port=status_config['port'], + database=status_config['dbname'], + ), + pool_size=2, + max_overflow=3, + pool_pre_ping=True, + ) + + def close(self) -> None: + """Dispose the status engine. The import holds no other connection.""" + engine = getattr(self, 'status_engine', None) + if engine is not None: + engine.dispose() + + # ------------------------------------------------------------------ helpers + + @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 + + @staticmethod + def _digest_from_request_data(request_data: Any) -> str | None: + """Find the uploader's digest inside the BFF-owned request payload, if it recorded one.""" + if not isinstance(request_data, dict): + return None + for key in REQUEST_DIGEST_KEYS: + value = request_data.get(key) + if isinstance(value, str) and value: + return value.lower() + return None + + def _execute_status_write(self, statement: Any, parameters: dict[str, Any]) -> int: + """Run one status statement in its own transaction and return the rows it affected. + + Raises: + ImportStatusRejectedError: the database accepted the connection and refused the + statement — a CHECK violation, an undefined column, a value too long. Non-retryable, + so a missing `V12` fails on the first attempt instead of after the full envelope. + Exception: anything else (connection refused, reset, timeout, a pod being replaced) is + left to propagate, because repeating it is exactly what fixes it. + """ + try: + with self.status_engine.begin() as connection: + return connection.execute(statement, parameters).rowcount + except (IntegrityError, ProgrammingError, DataError) as error: + raise ImportStatusRejectedError( + 'the import log database refused the status statement' + ) from error + + def _read_row(self, import_run_id: Any) -> dict[str, Any] | None: + """Read the record. A read, not a write: nothing of somebody else's is overwritten.""" + with self.status_engine.connect() as connection: + row = connection.execute(_READ_ROW_SQL, {'import_run_id': import_run_id}).mappings() + found = row.first() + return dict(found) if found is not None else None + + @staticmethod + def _describe_unclaimable(row: dict[str, Any] | None, workflow_id: str) -> str: + """Name the cause of a zero-row claim: missing row, training row, or someone else's.""" + if row is None: + return 'no such import log row' + if row.get('run_type') != 'IMPORT': + return f'row is a {row.get("run_type")} run, not an import' + claimed_by = row.get('claimed_by') + if claimed_by and claimed_by != workflow_id: + return f'row is already claimed by workflow {claimed_by}' + return f'row is in status {row.get("status")} and cannot be claimed' + + # -------------------------------------------------------------- class A: claim + + @activity.defn(name='claim_import_status') + async def claim_import_status(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Claim the import log row the frontend inserted. The workflow's first write. + + One conditional `UPDATE`, whose second disjunct makes a re-execution of the same workflow + idempotent instead of a conflict — a lost response on attempt 1 must not become a false + "not mine" on attempt 2. Zero rows affected covers the id typo, the duplicated start and the + wrong `run_type` in one condition, with no read-then-write race, and it is what stops a + duplicated workflow start from provisioning twice. + + Args (input_data): + metadata (dict): workflow metadata for logging and metrics. + import_run_id (int): the record's primary key, from the workflow input. + workflow_id (str): this workflow's identifier, recorded in the JSON detail. + source (dict, optional): the bundle's declared origin, when already known. + + Returns: + dict: `claimed`, plus the row values the input guards are checked against — + `bucket_name`, `file_name` and `request_digest`. + + Raises: + ImportLogRowNotClaimableError: the row is not this workflow's. Non-retryable, reported + through the channel, and **written to no row** — the loudness has to go somewhere + that is not somebody else's audit record. + """ + metadata = input_data.get('metadata', {}) + 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, + ) + await report_import_failure( + self, + None, + error, + import_run_id=import_run_id, + workflow_id=workflow_id, + metadata=metadata, + code=ImportErrorCode.LOG_ROW_NOT_CLAIMABLE, + marker='not_claimable', + ) + raise error + + await self.emit_metric( + metric_object=metrics.MODEL_IMPORT_STARTED_COUNT, + tags=self.get_core_labels(metadata, operation_type=ImportStep.RECEIVED.value), + ) + self.info(f'import log row {import_run_id} claimed by {workflow_id}', metadata) + row = row or {} + return { + 'claimed': True, + 'bucket_name': row.get('bucket_name'), + 'file_name': row.get('file_name'), + 'request_digest': self._digest_from_request_data(row.get('request_data')), + } + + # ----------------------------------------------------------- class B: progress + + @activity.defn(name='record_import_status') + async def record_import_status(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Record the step the workflow is about to attempt. A hint, keyed on the id alone. + + Its failure is **not** an import failure: the workflow reports it through the channel and + carries on, because losing a hint costs a record pointing one step behind and nothing else — + the terminal write rewrites status, code, reason and step from one classification. + + Args (input_data): + metadata (dict): workflow metadata. + import_run_id (int): the record's primary key. + workflow_id (str): recorded in the detail, so it is never dropped. + step (str): the step about to be attempted. + source (dict, optional): the bundle's declared origin. + """ + affected = self._execute_status_write( + _PROGRESS_SQL, + { + 'status': ExperimentRunStatus.ORCHESTRATOR_WAITING_PROC.value, + 'updated_at': self._timestamp(), + 'detail': build_import_detail( + workflow_id=input_data['workflow_id'], + step=input_data['step'], + source=input_data.get('source'), + ), + 'import_run_id': input_data['import_run_id'], + }, + ) + return {'recorded': affected > 0, 'rows': affected} + + @activity.defn(name='record_import_names') + async def record_import_names(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Write the experiment and run names, before anything exists in MLflow. + + Both columns exist so a person watching the import list sees *which* model is being + imported. The earliest moment both values are known is when the bundle is open, so that is + when they are written — before the experiment exists, before the run exists. + + Args (input_data): + metadata (dict): workflow metadata. + import_run_id (int): the record's primary key. + workflow_id (str): recorded in the detail. + experiment_name (str): the bundle's own experiment name, verbatim. + run_name (str): the derived run name, which the run is then created with. + source (dict, optional): the bundle's declared origin. + """ + affected = self._execute_status_write( + _NAMES_SQL, + { + 'experiment_name': input_data['experiment_name'], + 'run_name': input_data['run_name'], + 'updated_at': self._timestamp(), + 'detail': build_import_detail( + workflow_id=input_data['workflow_id'], + step=input_data.get('step', ImportStep.EXPERIMENT_CREATION), + source=input_data.get('source'), + ), + 'import_run_id': input_data['import_run_id'], + }, + ) + return {'recorded': affected > 0, 'rows': affected} + + # ----------------------------------------------------------- class C: terminal + + @activity.defn(name='record_import_terminal_status') + async def record_import_terminal_status(self, input_data: dict[str, Any]) -> dict[str, Any]: + """The single authoritative write, run from the workflow's `finally` after cleanup. + + Every failure is classified **once**, here, from the step that failed. On success the reason + column is written as `NULL` — never `''`, which its length-3 CHECK would refuse — and the + detail carries no code. A cleanup failure is noted in the detail without changing a + successful status: the user's model was imported, and reporting ERROR would be false. + + Args (input_data): + metadata (dict): workflow metadata. + import_run_id (int): the record's primary key. + workflow_id (str): recorded in the detail. + succeeded (bool): whether every step succeeded. + step (str, optional): the failing step, or the last step reached on success. + gate (int, optional): the gate that rejected the bundle, when one did. + source (dict, optional): the bundle's declared origin. + cleanup_failed (bool, optional): whether the temporary-file cleanup failed. + + Returns: + dict: the verdict that was written — status, code and reason — so a caller that has to + report a failure to record it still holds the verdict it could not write. + """ + metadata = input_data.get('metadata', {}) + succeeded = bool(input_data.get('succeeded')) + step = input_data.get('step') + cleanup_failed = bool(input_data.get('cleanup_failed')) + + if succeeded: + status = map_state_to_status(ImportState.COMPLETED) + 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) + code, reason = failure.code, failure.reason + + verdict = { + 'status': status.value, + 'code': code.value if code else None, + 'reason': reason, + 'step': str(step) if step else None, + } + + self._execute_status_write( + _TERMINAL_SQL, + { + 'status': status.value, + 'error_message': reason, + 'updated_at': self._timestamp(), + 'detail': build_import_detail( + workflow_id=input_data['workflow_id'], + step=step, + code=code, + gate=input_data.get('gate'), + source=input_data.get('source'), + cleanup_failed=cleanup_failed, + ), + 'import_run_id': input_data['import_run_id'], + }, + ) + + if succeeded: + await self.emit_metric( + metric_object=metrics.MODEL_IMPORT_COMPLETED_COUNT, + tags=self.get_core_labels(metadata, operation_type=status.value), + ) + self.info(f'import {input_data["import_run_id"]} recorded as {status.value}', metadata) + return verdict + + @activity.defn(name='report_import_status_write_failure') + async def report_import_status_write_failure( + self, input_data: dict[str, Any] + ) -> dict[str, Any]: + """Report a status write that could not be made — the one failure with no row to land on. + + Called by the workflow when a progress or terminal write exhausted its attempts or was + refused. It carries the verdict it could not write, so the outcome survives in text even + when it cannot survive in the row, and it never turns a completed import into a failed one. + + Args (input_data): + metadata (dict): workflow metadata. + import_run_id (int): the record's primary key. + workflow_id (str): the identifier that correlates this line with the record. + detail (str): the underlying failure, as text — it crossed an activity boundary, so the + exception object itself is gone by now. + verdict (str, optional): the verdict the terminal write could not record. + """ + await report_import_failure( + self, + None, + None, + import_run_id=input_data['import_run_id'], + workflow_id=input_data['workflow_id'], + metadata=input_data.get('metadata', {}), + code=ImportErrorCode.STATUS_NOT_RECORDED, + marker='status_write', + verdict=input_data.get('verdict'), + detail_text=input_data.get('detail'), + ) + return {'reported': True} + + # ------------------------------------------------------------- bundle helpers + + def _bundle_reader(self) -> BundleReader: + """A reader carrying the configured ceilings.""" + return BundleReader( + BundleLimits( + max_object_bytes=self.import_config['max_object_bytes'], + max_entries=self.import_config['max_entries'], + max_uncompressed_bytes=self.import_config['max_uncompressed_bytes'], + max_compression_ratio=self.import_config['max_compression_ratio'], + object_prefix=self.import_config['prefix'], + ) + ) + + def _work_dir(self, workflow_id: str) -> Path: + """This import's own scratch directory. One per workflow, never shared.""" + return Path(self.import_config['work_dir']) / workflow_id + + def _effective_bucket(self, requested: str | None) -> str: + return requested or self.import_config['bucket'] + + # --------------------------------------------------------- download and open + + @activity.defn(name='download_import_bundle') + async def download_import_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. + The SHA-256 of what was downloaded is then compared against `expected_digest` — the value the + uploader computed over the bytes it uploaded — with **no path that skips the comparison**, + because a missing value was already rejected at `received`. A mismatch fails here, before the + header is parsed: it means the object behind that key is not the file the user chose. + + 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. + + Returns: + dict: `path`, `size` and `digest` of the downloaded file. + + Raises: + ImportGateRejectedError: gate 1 rejected the object, or the digests disagree. + """ + if self.import_minio_repository is None: + raise ValueError(_MISSING_MINIO) + + metadata = input_data.get('metadata', {}) + bucket = self._effective_bucket(input_data.get('bucket')) + object_key = input_data['object_key'] + expected_digest = input_data['expected_digest'] + reader = self._bundle_reader() + + client = self.import_minio_repository.minio_client + try: + stat = client.stat_object(bucket, object_key) + except Exception as error: + raise ImportGateRejectedError( + 'the uploaded object could not be read', step=ImportStep.DOWNLOAD, gate=1 + ) from error + + try: + reader.check_object_limits(size=stat.size or 0, key=object_key) + except BundleGateError as error: + raise ImportGateRejectedError(str(error), step=error.step, gate=error.gate) from error + + work_dir = self._work_dir(input_data['workflow_id']) + work_dir.mkdir(parents=True, exist_ok=True) + target = work_dir / 'bundle.sientia' + try: + client.fget_object(bucket, object_key, str(target)) + except Exception as error: + raise ImportGateRejectedError( + 'the uploaded object could not be downloaded', step=ImportStep.DOWNLOAD, gate=1 + ) from error + + digest = fmt.file_digest(target) + if digest != str(expected_digest).lower(): + target.unlink(missing_ok=True) + raise ImportGateRejectedError( + 'the downloaded object does not match the digest the uploader computed', + step=ImportStep.DOWNLOAD, + gate=1, + ) + + 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]: + """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 + no plaintext password is ever an activity input, an activity result or a workflow variable — + and therefore never lands in the Temporal event history. The recovered password is a local + value and nothing else. + + Args (input_data): + metadata (dict): workflow metadata. + workflow_id (str): names this import's scratch directory. + path (str): the downloaded bundle. + password_envelope (str): base64 `nonce || ciphertext || tag`, AES-256-GCM. + + Returns: + dict: `extracted_dir`, `metadata`, `parameters`, `digest` — nothing secret. + + Raises: + ImportGateRejectedError: with the step and the gate that rejected the bundle. + """ + metadata = input_data.get('metadata', {}) + work_dir = self._work_dir(input_data['workflow_id']) + extract_root = work_dir / 'extract' + + try: + key = load_envelope_key(self.import_config.get('password_key')) + password = decrypt_password_envelope(input_data['password_envelope'], key) + except EnvelopeError as error: + # A distinct failure from a wrong *bundle* password, and never confused with it: the + # envelope is the platform's own transport, and its message says so. Same step, because + # the step is where the import stopped, not who is at fault. + raise ImportGateRejectedError( + 'the import password envelope could not be opened', + step=ImportStep.DECRYPTION, + gate=3, + cause='envelope', + ) from error + + reader = self._bundle_reader() + try: + opened = reader.open(input_data['path'], password, extract_root, work_dir) + except BundleGateError as error: + raise ImportGateRejectedError(str(error), step=error.step, gate=error.gate) from error + finally: + del password + + await self.emit_metric( + metric_object=metrics.MODEL_IMPORT_KDF_DURATION, + tags=self.get_core_labels(metadata, operation_type=ImportStep.DECRYPTION.value), + method='observe', + value=opened.kdf_seconds, + ) + self.info(f'import bundle opened into {opened.extracted_dir}', metadata) + return { + 'extracted_dir': str(opened.extracted_dir), + 'metadata': opened.metadata, + 'parameters': opened.parameters, + 'digest': opened.digest, + } + + # ------------------------------------------------------------ MLflow provisioning + + @activity.defn(name='create_import_experiment') + async def create_import_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): + metadata (dict): workflow metadata. + experiment_name (str): the bundle's `metadata.experiment_name`. + + Returns: + dict: `experiment_id` and `experiment_name`. + """ + if self.import_mlflow_repository is None: + raise ValueError(_MISSING_MLFLOW) + + experiment_name = input_data['experiment_name'] + experiment = self.import_mlflow_repository.get_experiment( + experiment_name, create_if_not_exists=True + ) + self.info( + f'import experiment ready: {experiment_name} ({experiment.experiment_id})', + input_data.get('metadata', {}), + ) + 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]: + """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. + + 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. + 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.*`. + + Returns: + dict: `run_id` and `run_name`. + """ + if self.import_mlflow_repository is None: + raise ValueError(_MISSING_MLFLOW) + + client = self.import_mlflow_repository.client + run = client.create_run( + experiment_id=input_data['experiment_id'], run_name=input_data['run_name'] + ) + run_id = run.info.run_id + + 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) + + 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='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. + + The source is `runs://prediction_model`, because `get_model_run_id` extracts the run + id as `source.split('/')[2]` — a `models:/m-*` source, what the new stack emits, breaks it. + The stage is left untouched: a `COMPLETED` import is a registered model, not a servable one, + and promoting it to `Production` is a manual operator step outside this workflow. + + Args (input_data): + metadata (dict): workflow metadata. + model_name (str): the bundle's `metadata.model_name`, verbatim. + run_id (str): the run this import created. + + Returns: + dict: `model_name`, `version` and `source`. + """ + if self.import_mlflow_repository is None: + raise ValueError(_MISSING_MLFLOW) + + client = self.import_mlflow_repository.client + model_name = input_data['model_name'] + run_id = input_data['run_id'] + + try: + client.get_registered_model(model_name) + except Exception: + # 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) + + source = f'runs:/{run_id}/prediction_model' + version = client.create_model_version(name=model_name, source=source, run_id=run_id) + self.info( + f'import registered {model_name} version {version.version}', + input_data.get('metadata', {}), + ) + return {'model_name': model_name, 'version': str(version.version), 'source': source} + + # --------------------------------------------------- the MongoDB model document + + @activity.defn(name='write_import_model_document') + async def write_import_model_document(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Write the document that *is* the model listing, inert, last. + + 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. + + Args (input_data): + metadata (dict): workflow metadata. + model_name (str): the bundle's `metadata.model_name`, verbatim. + target (str): the bundle's `parameters.target_variable`. + + Returns: + dict: `written`, `id`, `skipped` and how many unusable ids were ignored. + """ + if self.mongo_repository is None: + raise ValueError(_MISSING_MONGO) + + metadata = input_data.get('metadata', {}) + collection = self.import_config['models_collection'] + model_name = input_data['model_name'] + + existing = await self.mongo_repository.find_one(collection, {'name': model_name}, metadata) + if existing is not None: + self.info( + f'model document for {model_name} already exists; leaving it untouched', metadata + ) + return {'written': False, 'skipped': True, 'id': existing.get('id')} + + next_id, ignored = await self._next_model_document_id(collection, metadata) + document = { + 'id': next_id, + 'name': model_name, + 'active': False, + 'model_config': { + 'transform_flavor': 'sklearn', + 'predict_flavor': 'joblib', + 'retention_minutes': 60, + 'target': input_data['target'], + }, + } + await self.mongo_repository.insert(collection, document, metadata) + self.info(f'model document written for {model_name} with id {next_id}', metadata) + return {'written': True, 'skipped': False, 'id': next_id, 'ignored_ids': ignored} + + 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 + + # ----------------------------------------------------------- cleanup and retention + + @activity.defn(name='cleanup_import_files') + async def cleanup_import_files(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Remove this import's scratch directory. Never the uploaded object. + + Runs on every path, success or failure, and **never raises**: a cleanup failure is reported + through the channel and noted in the record's detail, but the user's model was imported and + reporting an error would be false. + + Args (input_data): + metadata (dict): workflow metadata. + workflow_id (str): names the directory to remove. + import_run_id (int): for correlation when the removal fails. + + Returns: + dict: `cleaned` and, when it failed, the reason as text. + """ + metadata = input_data.get('metadata', {}) + work_dir = self._work_dir(input_data['workflow_id']) + try: + shutil.rmtree(work_dir, ignore_errors=False) + except FileNotFoundError: + # Nothing was downloaded, or a gate removed its own partial tree. Not a failure. + return {'cleaned': True, 'existed': False} + except OSError as error: + await report_import_failure( + self, + ImportStep.CLEANUP, + error, + import_run_id=input_data.get('import_run_id'), + workflow_id=input_data['workflow_id'], + metadata=metadata, + code=ImportErrorCode.CLEANUP_INCOMPLETE, + marker=ImportStep.CLEANUP.value, + ) + return {'cleaned': False, 'existed': True, 'reason': type(error).__name__} + return {'cleaned': True, 'existed': True} + + @activity.defn(name='ensure_import_bundle_retention') + async def ensure_import_bundle_retention(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Ensure the 7-day expiry rule on the import prefix, idempotently. + + Applied by the importer rather than assumed from a Helm value, so the acceptance criterion + ("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. + + Args (input_data): + metadata (dict): workflow metadata. + bucket (str, optional): the effective bucket; the configured one when absent. + + Returns: + dict: whether the rule was already there and what it expires. + """ + if self.import_minio_repository is None: + raise ValueError(_MISSING_MINIO) + + metadata = input_data.get('metadata', {}) + bucket = self._effective_bucket(input_data.get('bucket')) + prefix = self.import_config['prefix'] + days = self.import_config['retention_days'] + rule_id = 'sientia-import-bundle-expiry' + client = self.import_minio_repository.minio_client + + try: + existing = client.get_bucket_lifecycle(bucket) + except Exception: + # A bucket with no lifecycle configuration at all answers with an error on some + # deployments and with `None` on others; both mean "there is no rule yet". + existing = None + + rules = [rule for rule in (existing.rules if existing else []) if rule.rule_id != rule_id] + rules.append( + Rule( + ENABLED, + rule_filter=Filter(prefix=prefix), + rule_id=rule_id, + expiration=Expiration(days=days), + ) + ) + client.set_bucket_lifecycle(bucket, LifecycleConfig(rules)) + self.info(f'import retention rule ensured on {bucket} ({days} days)', metadata) + return {'bucket': bucket, 'prefix': prefix, 'days': days, 'rules': len(rules)} diff --git a/laborious/activities/model_import_errors.py b/laborious/activities/model_import_errors.py new file mode 100644 index 0000000..657f47f --- /dev/null +++ b/laborious/activities/model_import_errors.py @@ -0,0 +1,348 @@ +"""The import's failure vocabulary: a closed catalog, typed errors and one reporting channel. + +Three things live here and nowhere else: + +- **the catalog** — one stable `error_code` and one plain-language `error_reason` per pipeline step. + The reason is written for a person who does not know this system exists: what happened, what state + the model is in, what to do next. It never carries a path, a table name, an object key, a Python + exception class or a traceback, and it is never built from `str(exception)`. +- **the typed errors** — the two failures that must not be retried, because repeating them cannot + change the answer: a record this workflow does not own, and a statement the database refuses. +- **the channel** — `report_import_failure`, the single funnel every failure passes through. It is + the only place in the import path that logs an error or sends a notification, which a test + asserts, so no failure is reported twice or half-reported. +""" + +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + import traceback + from collections.abc import Mapping + from dataclasses import dataclass + from enum import StrEnum + from types import MappingProxyType + from typing import Any, Final, Protocol + + from temporalio.exceptions import ApplicationError + + from laborious import metrics + from laborious.utils.bundle.steps import ImportStep + + +class ImportErrorCode(StrEnum): + """The machine-readable half of a failure: the frontend's i18n key and support's triage key. + + Written under `import.code` in the record's JSON detail — not in a column of its own, so the + import adds nothing to a table it shares with the training flow. + """ + + REQUEST_INVALID = 'IMPORT_REQUEST_INVALID' + FILE_UNREADABLE = 'IMPORT_FILE_UNREADABLE' + BUNDLE_CANNOT_BE_OPENED = 'IMPORT_BUNDLE_CANNOT_BE_OPENED' + BUNDLE_REJECTED = 'IMPORT_BUNDLE_REJECTED' + BUNDLE_UNPACK_FAILED = 'IMPORT_BUNDLE_UNPACK_FAILED' + BUNDLE_INCOMPLETE = 'IMPORT_BUNDLE_INCOMPLETE' + BUNDLE_UNEXPECTED_CONTENT = 'IMPORT_BUNDLE_UNEXPECTED_CONTENT' + STORAGE_PREPARATION_FAILED = 'IMPORT_STORAGE_PREPARATION_FAILED' + MODEL_FILES_NOT_STORED = 'IMPORT_MODEL_FILES_NOT_STORED' + MODEL_NOT_REGISTERED = 'IMPORT_MODEL_NOT_REGISTERED' + MODEL_SETTINGS_NOT_SAVED = 'IMPORT_MODEL_SETTINGS_NOT_SAVED' + UNEXPECTED_ERROR = 'IMPORT_UNEXPECTED_ERROR' + # Codes that never reach a record's `error_message`, and are therefore absent from the catalog: + # one names the failure to own a row, one the failure to write one, one the temporary files left + # on the worker after an import that otherwise succeeded. All three live on the channel — log, + # metric, notification, Temporal error — and the third is additionally flagged in the record's + # detail as `cleanup_failed`, without a code and without changing the status. + LOG_ROW_NOT_CLAIMABLE = 'IMPORT_LOG_ROW_NOT_CLAIMABLE' + STATUS_NOT_RECORDED = 'IMPORT_STATUS_NOT_RECORDED' + CLEANUP_INCOMPLETE = 'IMPORT_CLEANUP_INCOMPLETE' + + +@dataclass(frozen=True) +class ImportFailure: + """One classification of one failure: the step, its code and the sentence it stores.""" + + step: ImportStep | None + code: ImportErrorCode + reason: str + + +_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( + { + 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.', + ), + ImportStep.DOWNLOAD: ( + ImportErrorCode.FILE_UNREADABLE, + 'The uploaded file could not be read, or it is not the file that was sent. ' + 'Upload it again.', + ), + 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.', + ), + 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}', + ), + ImportStep.EXTRACTION: ( + ImportErrorCode.BUNDLE_UNPACK_FAILED, + f'The file could not be unpacked safely and nothing from it was kept. {_RE_EXPORT}', + ), + 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}', + ), + ImportStep.CONTENT_POLICY: ( + ImportErrorCode.BUNDLE_UNEXPECTED_CONTENT, + 'The file contains something a model export should not contain, so it was ' + f'rejected. {_RE_EXPORT}', + ), + 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}', + ), + 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}', + ), + 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}', + ), + 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}', + ), + } +) + +UNCLASSIFIED_FAILURE: Final[tuple[ImportErrorCode, str]] = ( + ImportErrorCode.UNEXPECTED_ERROR, + f'The import stopped for an unexpected reason and no model was created. {_CONTACT_SUPPORT}', +) + +# `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: + """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. + + 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) + + +class ImportLogRowNotClaimableError(ApplicationError): + """The record named by `import_run_id` is not this workflow's to write. + + Raised when the claim affects zero rows: no such row, a `TRAIN` row, or a row already running + under another workflow. Non-retryable on purpose — a second attempt would read the same rows and + reach the same conclusion — and loud everywhere except the row itself, which belongs to someone + else. + """ + + def __init__(self, message: str, *, import_run_id: Any = None, observed: str = '') -> None: + super().__init__( + message, + type='ImportLogRowNotClaimableError', + non_retryable=True, + ) + self.import_run_id = import_run_id + self.observed = observed + + +class ImportStatusRejectedError(ApplicationError): + """The database accepted the connection and refused the statement. + + A `CHECK` violation, an undefined column, a value too long for its column. `V12` not being + applied yet lands here, at the terminal write, and it must fail on the first attempt instead of + spending the whole retry envelope to say so. + """ + + def __init__(self, message: str) -> None: + super().__init__(message, type='ImportStatusRejectedError', non_retryable=True) + + +class ImportInputError(ApplicationError): + """The request is incomplete or contradictory. Fails at `received`, on the claimed row.""" + + def __init__(self, message: str) -> None: + super().__init__( + message, + {'step': ImportStep.RECEIVED.value, 'gate': None}, + type='ImportInputError', + non_retryable=True, + ) + + +class ImportGateRejectedError(ApplicationError): + """A gate rejected the bundle, and the error says which one. + + Gates 2 to 7 all run inside one activity, so the step the workflow was *attempting* is not + precise enough to record. An exception loses its Python type crossing an activity boundary but + keeps its `details`, so the step and the gate travel there and the terminal write can name them. + """ + + def __init__( + self, + message: str, + *, + step: ImportStep, + gate: int | None = None, + cause: str | None = None, + ) -> None: + super().__init__( + message, + {'step': step.value, 'gate': gate, 'cause': cause}, + type='ImportGateRejectedError', + non_retryable=True, + ) + self.step = step + self.gate = gate + + +def failure_location(exception: BaseException | None) -> tuple[str | None, int | None]: + """Read the step and gate an activity's failure carries, if it carried any. + + Walks the `__cause__` chain because Temporal wraps an activity failure in an `ActivityError`. + 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: + 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 + + +class FailureSinks(Protocol): + """What the channel needs from its caller: the repository's existing observability surface.""" + + def error(self, message: str, metadata: dict | None = None) -> None: ... + + def get_core_labels(self, metadata: dict | None = None, operation_type: str = '-') -> dict: ... + + async def emit_metric( + self, + metric_object: Any, + tags: dict[str, Any], + method: str = 'inc', + value: float | None = None, + ) -> None: ... + + async def send_notification_async( + self, + metadata: dict, + notification_id: str, + message: str, + block: str, + level: str = 'INFO', + attachment_content: str | None = None, + ) -> None: ... + + +async def report_import_failure( + sinks: FailureSinks, + step: ImportStep | None, + exception: BaseException | None, + *, + import_run_id: Any, + workflow_id: str, + metadata: dict[str, Any] | None = None, + code: ImportErrorCode | None = None, + marker: str | None = None, + verdict: str | None = None, + detail_text: str | None = None, +) -> ImportFailure: + """Report one failure through every sink, exactly once. + + Sinks, in order: a structured log carrying the ids, the step, the code **and** the exception + with its traceback (the only sink that gets the exception text); the import error metric, with + the code in `operation_type`; a notification through the existing handler. The exception itself + reaches Temporal by being raised, and the record is the fifth, optional sink — the terminal + write, when the row can be written at all. + + Args: + sinks: the activity providing logger, metrics and notifications. + step: the pipeline step, or `None` for a failure that is not one (see `marker`). + exception: the failure. + import_run_id: the record's id, for correlation. + workflow_id: the identifier a user quotes to support. + metadata: workflow metadata for labels and the notification. + code: an explicit code, for the two failures that are not pipeline steps. + marker: `not_claimable` or `status_write` — what a step name would have been. + verdict: for a failed terminal write, the verdict it could not record, so the outcome + survives in text even when it cannot survive in the row. + detail_text: the failure as text, for the callers whose exception object is already gone — + anything that crossed an activity boundary. Used in place of a traceback. + + Returns: + ImportFailure: the classification that was reported. + """ + failure = classify_import_failure(step, exception) + if code is not None: + failure = ImportFailure(step=step, code=code, reason=failure.reason) + + where = marker or (step.value if step else 'unknown') + if exception is not None: + detail = ''.join(traceback.format_exception(exception)).strip() + else: + detail = detail_text or 'no exception object available' + verdict_line = f' verdict_not_recorded={verdict}' if verdict else '' + sinks.error( + f'model import failed: step={where} code={failure.code.value} ' + f'import_run_id={import_run_id} workflow_id={workflow_id}{verdict_line}\n{detail}', + metadata or {}, + ) + + labels = sinks.get_core_labels(metadata or {}, operation_type=failure.code.value) + await sinks.emit_metric(metric_object=metrics.MODEL_IMPORT_ERROR_COUNT, tags=labels) + await sinks.send_notification_async( + metadata=metadata or {}, + notification_id=failure.code.value, + message=failure.reason, + block=where, + level='ERROR', + ) + return failure diff --git a/laborious/metrics.py b/laborious/metrics.py index 74a9189..3f4de2c 100644 --- a/laborious/metrics.py +++ b/laborious/metrics.py @@ -183,6 +183,45 @@ MODEL_ANALYZE_ERROR_COUNT = Counter( ) +# ================== Model import metrics ================== +# +# The failing step and the error code travel in `operation_type`, the one label the shared +# `CORE_LABELS` leaves free. Nothing else about a failure is a label: an error reason is a sentence +# and an exception is a log line, neither of which belongs in a metric's cardinality. + +MODEL_IMPORT_STARTED_COUNT = Counter( + 'laborious_model_import_started_count', + 'Number of model imports that claimed their import log row', + CORE_LABELS, +) + +MODEL_IMPORT_COMPLETED_COUNT = Counter( + 'laborious_model_import_completed_count', + 'Number of model imports that reached a registered version and a model document', + CORE_LABELS, +) + +MODEL_IMPORT_ERROR_COUNT = Counter( + 'laborious_model_import_error_count', + 'Number of failed model imports, with the failing step or error code in operation_type', + CORE_LABELS, +) + +MODEL_IMPORT_KDF_DURATION = Histogram( + 'laborious_model_import_kdf_duration_seconds', + 'Time spent deriving the bundle key with Argon2id', + CORE_LABELS, + buckets=[0.5, 1.0, 2.0, 3.0, 5.0, 8.0, 13.0, 21.0], +) + +MODEL_IMPORT_STEP_DURATION = Histogram( + 'laborious_model_import_step_duration_seconds', + 'Duration of each model import step, named in operation_type', + CORE_LABELS, + buckets=[0.1, 0.5, 1.0, 5.0, 10.0, 30.0, 60.0, 300.0], +) + + # ================== PI Web API metrics ================== PI_WEB_API_LABELS = [*CORE_LABELS, 'tag_name'] diff --git a/laborious/utils/bundle/__init__.py b/laborious/utils/bundle/__init__.py new file mode 100644 index 0000000..77e7945 --- /dev/null +++ b/laborious/utils/bundle/__init__.py @@ -0,0 +1,27 @@ +"""Reading `.sientia` bundles: the wire format, the seven gates and the step vocabulary. + +There is no writer here. The producer is the Streamlit app in `sientia-projects-templates` +(`app/src/operations/model_export/`), and a second production writer for a format with one producer +would be a fork waiting to happen. What stands in for it in the tests is +`tests/helpers/bundle_factory.py` plus the golden fixtures under `tests/fixtures/bundle/`. +""" + +from laborious.utils.bundle.format import BundleFormatError, BundleHeader +from laborious.utils.bundle.reader import ( + BundleGateError, + BundleLimits, + BundleReader, + OpenedBundle, +) +from laborious.utils.bundle.steps import GATE_TO_STEP, ImportStep + +__all__ = [ + 'GATE_TO_STEP', + 'BundleFormatError', + 'BundleGateError', + 'BundleHeader', + 'BundleLimits', + 'BundleReader', + 'ImportStep', + 'OpenedBundle', +] diff --git a/laborious/utils/bundle/format.py b/laborious/utils/bundle/format.py new file mode 100644 index 0000000..f07d2f8 --- /dev/null +++ b/laborious/utils/bundle/format.py @@ -0,0 +1,216 @@ +"""The `.sientia` wire format, as the producer writes it. + +Every constant here is copied from the producer that emits the bundles this reader opens — the +Streamlit app in `sientia-projects-templates`, modules +`app/src/operations/model_export/crypto.py` and `app/src/operations/model_export/bundle.py`. The +producer file each value came from is named next to it. The golden fixtures under +`tests/fixtures/bundle/` are the anti-drift device: if the producer's layout moves, the tests here +fail against the committed bytes instead of failing in production. + +This module deliberately imports nothing from Laborious, Temporal, MLflow or libsodium: it is pure +`struct` and `hashlib` over bytes, so it stays liftable into a shared package if the format ever +gets one (the cancelled QTZPOC-15). `reader.py` is where the crypto lives. +""" + +from __future__ import annotations + +import hashlib +import struct +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +# --- Producer: app/src/operations/model_export/crypto.py --- +MAGIC: Final[bytes] = b'SIENTIA1' +FORMAT_VERSION: Final[int] = 1 +AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM: Final[int] = 1 +KDF_ID_ARGON2ID: Final[int] = 1 +SALT_SIZE: Final[int] = 16 +DIGEST_SIZE: Final[int] = 32 +RESERVED_SIGNATURE_SIZE: Final[int] = 64 +# `KDF_MEMLIMIT_BYTES` / `KDF_OPSLIMIT` on the producer: what it writes is also the ceiling the +# reader refuses to exceed, so header bytes cannot buy an arbitrary allocation on the worker. +MAX_KDF_MEMLIMIT: Final[int] = 256 * 1024 * 1024 # 268435456 +MAX_KDF_OPSLIMIT: Final[int] = 3 +# `KDF_PARALLELISM_RESERVED`: libsodium's `crypto_pwhash` has no lane parameter, so the field is +# reserved and must be exactly 1. +PARALLELISM: Final[int] = 1 +CHUNK_SIZE: Final[int] = 64 * 1024 # 65536, fixed by the format +KEY_SIZE: Final[int] = 32 # crypto_secretstream_xchacha20poly1305_KEYBYTES + +# magic(8s) version(B) aead_id(B) kdf_id(B) salt(16s) memlimit(Q) opslimit(Q) +# parallelism(B) digest(32s) reserved_signature(64s) — producer's `_HEADER_STRUCT`. +HEADER_STRUCT: Final[struct.Struct] = struct.Struct('>8sBBB16sQQB32s64s') +HEADER_SIZE: Final[int] = HEADER_STRUCT.size # 140 + +# --- libsodium constants the framing depends on --- +# Not producer constants, but constants of the primitive the producer uses. Kept as literals so +# this module stays import-free; `tests/laborious/utils/bundle/test_format.py` asserts they equal +# the values `nacl.bindings` reports, which is what makes the literals safe. +STREAM_HEADER_SIZE: Final[int] = 24 # crypto_secretstream_xchacha20poly1305_HEADERBYTES +ABYTES: Final[int] = 17 # crypto_secretstream_xchacha20poly1305_ABYTES + +# The pre-upload check prefix of `08-security-and-encryption.md` § 5.3: enough bytes to hold the +# header, the stream header and one full chunk. Derived from constants, never read from the file. +PREFIX_SIZE: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE + CHUNK_SIZE + ABYTES # 65717 + +# The digest covers the chunk stream only. The producer computes `sha256(ciphertext)` over the +# joined chunks and *then* writes `header + stream_header + ciphertext`, so on disk the covered +# range starts after both headers. Verified against a bundle produced by the producer's own code: +# `sha256(file[164:])` matches the header field and `sha256(file[140:])` does not. +DIGEST_OFFSET: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE # 164 + +MIN_FILE_SIZE: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE + +SUPPORTED_FORMAT_VERSIONS: Final[frozenset[int]] = frozenset({FORMAT_VERSION}) +SUPPORTED_AEAD_IDS: Final[frozenset[int]] = frozenset({AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM}) +SUPPORTED_KDF_IDS: Final[frozenset[int]] = frozenset({KDF_ID_ARGON2ID}) + +_HASH_BLOCK_SIZE: Final[int] = 1024 * 1024 + + +class BundleFormatError(Exception): + """The bytes are not a `.sientia` bundle this reader implements. + + Raised for every header-level rejection: wrong magic, unknown version, unknown primitive ids, + cost parameters above the ceiling, a reserved field carrying a value, a file too short to hold + a header, or a chunk stream whose digest does not match the header's. `reader.py` turns this + into a gate-2 rejection; nothing here knows about gates or import steps. + """ + + +@dataclass(frozen=True) +class BundleHeader: + """The 140-byte plaintext header, parsed and validated.""" + + format_version: int + aead_id: int + kdf_id: int + salt: bytes + kdf_memlimit_bytes: int + kdf_opslimit: int + kdf_parallelism: int + ciphertext_digest: bytes + reserved_signature: bytes + + +def parse_header(raw: bytes) -> BundleHeader: + """Parse and validate the plaintext header. + + Everything the reader trusts about the file's shape is decided here, before any offset past + the header is used and before a key is derived. + + Args: + raw: at least `HEADER_SIZE` bytes read from the start of the object. + + Returns: + BundleHeader: the validated header. + + Raises: + BundleFormatError: on any invalid or unsupported field. + """ + if len(raw) < HEADER_SIZE: + raise BundleFormatError( + f'file is shorter than the {HEADER_SIZE}-byte header ({len(raw)} bytes available)' + ) + + ( + magic, + format_version, + aead_id, + kdf_id, + salt, + kdf_memlimit_bytes, + kdf_opslimit, + kdf_parallelism, + ciphertext_digest, + reserved_signature, + ) = HEADER_STRUCT.unpack(raw[:HEADER_SIZE]) + + if magic != MAGIC: + raise BundleFormatError('file does not start with the expected bundle marker') + if format_version not in SUPPORTED_FORMAT_VERSIONS: + raise BundleFormatError(f'unsupported bundle format version: {format_version}') + if aead_id not in SUPPORTED_AEAD_IDS: + raise BundleFormatError(f'unsupported encryption identifier: {aead_id}') + if kdf_id not in SUPPORTED_KDF_IDS: + raise BundleFormatError(f'unsupported key derivation identifier: {kdf_id}') + if kdf_memlimit_bytes > MAX_KDF_MEMLIMIT: + raise BundleFormatError( + f'declared key derivation memory {kdf_memlimit_bytes} exceeds the ' + f'{MAX_KDF_MEMLIMIT} ceiling' + ) + if kdf_opslimit > MAX_KDF_OPSLIMIT: + raise BundleFormatError( + f'declared key derivation passes {kdf_opslimit} exceeds the {MAX_KDF_OPSLIMIT} ceiling' + ) + if kdf_memlimit_bytes <= 0 or kdf_opslimit <= 0: + raise BundleFormatError('key derivation parameters must be positive') + if kdf_parallelism != PARALLELISM: + raise BundleFormatError( + f'reserved parallelism field must be {PARALLELISM}, found {kdf_parallelism}' + ) + if reserved_signature != bytes(RESERVED_SIGNATURE_SIZE): + raise BundleFormatError('reserved signature field must be empty in this format version') + + return BundleHeader( + format_version=format_version, + aead_id=aead_id, + kdf_id=kdf_id, + salt=salt, + kdf_memlimit_bytes=kdf_memlimit_bytes, + kdf_opslimit=kdf_opslimit, + kdf_parallelism=kdf_parallelism, + ciphertext_digest=ciphertext_digest, + reserved_signature=reserved_signature, + ) + + +def read_header(path: Path | str) -> BundleHeader: + """Read and validate the header of the file at `path`. + + Raises: + BundleFormatError: when the file cannot hold a header plus a stream header, or when the + header itself is invalid. + """ + file_path = Path(path) + size = file_path.stat().st_size + if size < MIN_FILE_SIZE: + raise BundleFormatError( + f'file is {size} bytes, smaller than the {MIN_FILE_SIZE} bytes a bundle needs for its ' + 'header and encrypted stream header' + ) + with file_path.open('rb') as handle: + return parse_header(handle.read(HEADER_SIZE)) + + +def compute_ciphertext_digest(path: Path | str) -> bytes: + """SHA-256 of the chunk stream — the bytes from `DIGEST_OFFSET` to the end of the file.""" + digest = hashlib.sha256() + with Path(path).open('rb') as handle: + handle.seek(DIGEST_OFFSET) + while block := handle.read(_HASH_BLOCK_SIZE): + digest.update(block) + return digest.digest() + + +def verify_ciphertext_digest(path: Path | str, header: BundleHeader) -> None: + """Compare the file's chunk-stream digest against the header's. + + Raises: + BundleFormatError: when they differ — a truncated, overwritten or corrupted transfer. This + runs before any key derivation, so a damaged file never costs an Argon2id pass. + """ + if compute_ciphertext_digest(path) != header.ciphertext_digest: + raise BundleFormatError( + 'the encrypted content does not match the digest recorded in the bundle header' + ) + + +def file_digest(path: Path | str) -> str: + """Hex SHA-256 of the whole file, for the `expected_digest` comparison of gate 1.""" + digest = hashlib.sha256() + with Path(path).open('rb') as handle: + while block := handle.read(_HASH_BLOCK_SIZE): + digest.update(block) + return digest.hexdigest() diff --git a/laborious/utils/bundle/reader.py b/laborious/utils/bundle/reader.py new file mode 100644 index 0000000..61fcd90 --- /dev/null +++ b/laborious/utils/bundle/reader.py @@ -0,0 +1,473 @@ +"""Reading a `.sientia` bundle: the seven gates, in order, in one function. + +`BundleReader.open()` is the only entry point. Gate ordering is a property of that function rather +than of the caller, so it cannot be reordered by accident and can be unit-tested without Temporal: + +1. object-level limits (`check_object_limits`, called by the download activity before the body is + fetched — the only gate that does not need the file) +2. header validation, including the chunk-stream digest, before any key derivation +3. AEAD decryption — the integrity check; a wrong password and a tampered file fail identically +4. archive inspection over the central directory, writing nothing to disk +5. extraction into a fresh, isolated directory +6. structure and schema of the extracted tree +7. content policy over the file set + +Nothing here logs, notifies or emits a metric: rejections are raised as `BundleGateError` and the +activity layer owns the one failure channel. +""" + +from __future__ import annotations + +import json +import os +import shutil +import stat +import time +import unicodedata +import zipfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import nacl.bindings as sodium +import nacl.exceptions + +from laborious.utils.bundle import format as fmt +from laborious.utils.bundle.steps import GATE_TO_STEP, ImportStep + +# Gate 6 — the layout `bundle.build_zip` writes and `model_repository.py` reads back. +METADATA_FILE_NAME = 'metadata.json' +ARTIFACTS_DIR_NAME = 'artifacts' +REQUIRED_TOP_LEVEL_KEYS = frozenset({'parameters', 'metadata'}) +REQUIRED_METADATA_FIELDS = ( + 'model_name', + 'experiment_name', + 'run_id', + 'model_version', + 'model_project', + 'export_timestamp', +) +REQUIRED_PARAMETERS = ('target_variable',) +REQUIRED_MODEL_DIRS = ('prediction_model', 'data_model') +REQUIRED_MODEL_FILES = ('MLmodel', 'model.pkl') + +# The import record's `experiment_name` column is `VARCHAR(50)` with a length-3 floor, and the name +# is written onto the record before provisioning starts — so the column's limits are enforced here, +# not three activities later as a constraint violation. +EXPERIMENT_NAME_MIN_LENGTH = 3 +EXPERIMENT_NAME_MAX_LENGTH = 50 + +# Gate 7 — an MLflow 2.x artifact tree for this platform: the model directories' known files, the +# model card, the CSVs the training template logs and the bundle's own metadata. +ALLOWED_FILE_NAMES = frozenset( + { + 'MLmodel', + 'metadata.json', + 'model.pkl', + 'conda.yaml', + 'python_env.yaml', + 'requirements.txt', + 'model_card.json', + 'model_card.svg', + } +) +ALLOWED_FILE_SUFFIXES = frozenset({'.csv', '.json', '.yaml', '.yml', '.txt', '.pkl', '.svg', '.md'}) + +_DECRYPT_FAILURE_MESSAGE = ( + 'could not open this bundle: wrong password, or the file is damaged or was altered' +) + + +class BundleGateError(Exception): + """A gate rejected the bundle. + + Carries the gate number and the import step that gate maps onto, so the workflow never has to + guess which step to record. + """ + + def __init__(self, message: str, *, gate: int, step: ImportStep | None = None) -> None: + super().__init__(message) + self.gate = gate + self.step = step if step is not None else GATE_TO_STEP[gate] + + +@dataclass(frozen=True) +class BundleLimits: + """Ceilings enforced by gates 1 and 4. Values come from `build_import_config()`.""" + + max_object_bytes: int = 1024 * 1024 * 1024 + max_entries: int = 5000 + max_uncompressed_bytes: int = 4 * 1024 * 1024 * 1024 + max_compression_ratio: float = 200.0 + object_prefix: str = 'imported_models/' + object_suffix: str = '.sientia' + + +@dataclass(frozen=True) +class OpenedBundle: + """What a bundle yields once every gate has passed. Nothing secret is carried.""" + + extracted_dir: Path + metadata: dict[str, Any] = field(default_factory=dict) + parameters: dict[str, Any] = field(default_factory=dict) + digest: str = '' + header: fmt.BundleHeader | None = None + # How long Argon2id took, so the worker's most expensive step is measurable without the + # activity layer having to reach inside the reader. + kdf_seconds: float = 0.0 + + +class BundleReader: + """Opens `.sientia` bundles produced by the exporter in `sientia-projects-templates`.""" + + def __init__(self, limits: BundleLimits | None = None) -> None: + self.limits = limits or BundleLimits() + + # ------------------------------------------------------------------ gate 1 + + def check_object_limits(self, *, size: int, key: str) -> None: + """Gate 1, object level: reject before the object's body is fetched. + + Args: + size: the object's size as reported by `stat_object`. + key: the object key. + + Raises: + BundleGateError: oversized object, key outside the import prefix, or wrong suffix. + """ + if size > self.limits.max_object_bytes: + raise BundleGateError( + f'uploaded object is {size} bytes, above the ' + f'{self.limits.max_object_bytes} byte limit', + gate=1, + ) + if size <= 0: + raise BundleGateError('uploaded object is empty', gate=1) + if not key.startswith(self.limits.object_prefix): + raise BundleGateError('uploaded object is not under the import prefix', gate=1) + if not key.endswith(self.limits.object_suffix): + raise BundleGateError( + f'uploaded object does not end in {self.limits.object_suffix}', gate=1 + ) + + # --------------------------------------------------------------------- open + + def open( + self, + encrypted_path: Path | str, + password: str, + extract_root: Path | str, + work_dir: Path | str | None = None, + ) -> OpenedBundle: + """Run gates 2 to 7 in order and return the opened bundle. + + Args: + encrypted_path: the downloaded `.sientia` file. + password: the plaintext bundle password. Never stored, never logged, never returned. + extract_root: the directory to extract into. It MUST NOT exist yet (gate 5). + work_dir: where the decrypted zip is written. Defaults to `extract_root`'s parent. + + Returns: + OpenedBundle: extracted tree, `metadata`, `parameters`, digest and header. + + Raises: + BundleGateError: with the gate that rejected the bundle and its import step. + """ + encrypted = Path(encrypted_path) + root = Path(extract_root) + staging = Path(work_dir) if work_dir is not None else root.parent + + 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) + metadata, parameters = self._gate_6_structure(root) + self._gate_7_content_policy(root) + except BaseException: + # A rejection after decryption leaves nothing extracted behind (gate 5's contract) and + # never leaves a decrypted zip on the worker for the next import to trip over. + shutil.rmtree(root, ignore_errors=True) + zip_path.unlink(missing_ok=True) + raise + zip_path.unlink(missing_ok=True) + + return OpenedBundle( + extracted_dir=root, + metadata=metadata, + parameters=parameters, + digest=header.ciphertext_digest.hex(), + header=header, + kdf_seconds=kdf_seconds, + ) + + # ------------------------------------------------------------------ gate 2 + + def _gate_2_header(self, encrypted: Path) -> fmt.BundleHeader: + """Parse the header and check the chunk-stream digest, before deriving anything.""" + try: + header = fmt.read_header(encrypted) + fmt.verify_ciphertext_digest(encrypted, header) + except fmt.BundleFormatError as error: + raise BundleGateError(str(error), gate=2) from error + except OSError as error: + raise BundleGateError('downloaded bundle could not be read', gate=2) from error + return header + + # ------------------------------------------------------------------ gate 3 + + def _gate_3_decrypt( + self, encrypted: Path, password: str, header: fmt.BundleHeader, staging: Path + ) -> tuple[Path, float]: + """Derive the key with the header's parameters and decrypt the stream to a file. + + The AEAD *is* the integrity check: a wrong password and a tampered byte fail here with the + same message, which never claims to distinguish them. + """ + staging.mkdir(parents=True, exist_ok=True) + zip_path = staging / 'bundle.zip' + try: + started = time.monotonic() + key = self._derive_key(password, header) + kdf_seconds = time.monotonic() - started + self._pull_stream(encrypted, key, zip_path) + except BundleGateError: + zip_path.unlink(missing_ok=True) + raise + except (nacl.exceptions.CryptoError, ValueError, RuntimeError) as error: + zip_path.unlink(missing_ok=True) + raise BundleGateError(_DECRYPT_FAILURE_MESSAGE, gate=3) from error + except OSError as error: + zip_path.unlink(missing_ok=True) + raise BundleGateError('decrypted bundle could not be written', gate=3) from error + return zip_path, kdf_seconds + + @staticmethod + def _derive_key(password: str, header: fmt.BundleHeader) -> bytes: + """Argon2id over the NFC-normalised password, with the header's own cost parameters. + + The password is normalised and **not** trimmed, matching the producer's `_derive_key`: the + same password typed on another OS must derive the same key, while a trailing space is part + of the secret. + """ + normalized = unicodedata.normalize('NFC', password) + return sodium.crypto_pwhash_alg( + fmt.KEY_SIZE, + normalized.encode('utf-8'), + header.salt, + header.kdf_opslimit, + header.kdf_memlimit_bytes, + sodium.crypto_pwhash_ALG_ARGON2ID13, + ) + + @staticmethod + def _pull_stream(encrypted: Path, key: bytes, zip_path: Path) -> None: + """Decrypt the chunk stream into `zip_path`, requiring the FINAL tag.""" + chunk_size = fmt.CHUNK_SIZE + fmt.ABYTES + state = sodium.crypto_secretstream_xchacha20poly1305_state() + with encrypted.open('rb') as source, zip_path.open('wb') as target: + source.seek(fmt.HEADER_SIZE) + stream_header = source.read(fmt.STREAM_HEADER_SIZE) + sodium.crypto_secretstream_xchacha20poly1305_init_pull(state, stream_header, key) + last_tag: int | None = None + while chunk := source.read(chunk_size): + plaintext, last_tag = sodium.crypto_secretstream_xchacha20poly1305_pull( + state, chunk + ) + target.write(plaintext) + if last_tag != sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL: + # A stream that stops without its FINAL tag was cut short. The digest already catches + # truncation of a stored file; this catches a stream that was framed to look complete. + raise BundleGateError(_DECRYPT_FAILURE_MESSAGE, gate=3) + + # ------------------------------------------------------------------ gate 4 + + def _gate_4_inspect(self, zip_path: Path) -> None: + """Inspect the central directory only. Nothing is written to disk by this gate.""" + try: + with zipfile.ZipFile(zip_path) as archive: + entries = archive.infolist() + except zipfile.BadZipFile as error: + raise BundleGateError('bundle content is not a readable archive', gate=4) from error + + if len(entries) > self.limits.max_entries: + raise BundleGateError( + f'archive declares {len(entries)} entries, above the ' + f'{self.limits.max_entries} limit', + gate=4, + ) + if not entries: + raise BundleGateError('archive is empty', gate=4) + + uncompressed = sum(entry.file_size for entry in entries) + compressed = sum(entry.compress_size for entry in entries) + if uncompressed > self.limits.max_uncompressed_bytes: + raise BundleGateError( + f'archive declares {uncompressed} uncompressed bytes, above the ' + f'{self.limits.max_uncompressed_bytes} limit', + gate=4, + ) + ratio = uncompressed / max(compressed, 1) + if ratio > self.limits.max_compression_ratio: + raise BundleGateError( + f'archive compression ratio {ratio:.1f} is above the ' + f'{self.limits.max_compression_ratio} limit', + gate=4, + ) + + for entry in entries: + self._check_entry_name(entry.filename) + self._check_entry_mode(entry) + + @staticmethod + def _check_entry_name(name: str) -> None: + """Reject absolute names and traversal, from the declared name alone.""" + normalised = name.replace('\\', '/') + if normalised.startswith('/') or (len(normalised) > 1 and normalised[1] == ':'): + raise BundleGateError('archive declares an absolute entry path', gate=4) + if any(part == '..' for part in normalised.split('/')): + raise BundleGateError('archive declares an entry escaping its own tree', gate=4) + + @staticmethod + def _check_entry_mode(entry: zipfile.ZipInfo) -> None: + """Reject symlinks and anything that is neither a regular file nor a directory.""" + mode = entry.external_attr >> 16 + if stat.S_IFMT(mode) == 0: + # No file-type bits stored. This is the normal case for two kinds of entry: a zip + # written by a tool that records no Unix mode at all (`mode == 0`), and + # `ZipFile.writestr`, which stores permissions only — the producer's `metadata.json` + # arrives as `0o600 << 16`. Nothing to check: the member is extracted as a plain file + # with the worker's own permissions either way. + return + if stat.S_ISLNK(mode): + raise BundleGateError('archive declares a symbolic link', gate=4) + if not (stat.S_ISREG(mode) or stat.S_ISDIR(mode)): + raise BundleGateError('archive declares an entry that is not a regular file', gate=4) + + # ------------------------------------------------------------------ gate 5 + + def _gate_5_extract(self, zip_path: Path, root: Path) -> 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. + """ + if root.exists(): + raise BundleGateError('extraction directory already exists', gate=5) + root.mkdir(parents=True) + resolved_root = root.resolve() + try: + with zipfile.ZipFile(zip_path) as archive: + for entry in archive.infolist(): + self._extract_member(archive, entry, root, resolved_root) + except BundleGateError: + shutil.rmtree(root, ignore_errors=True) + raise + except (OSError, zipfile.BadZipFile) as error: + shutil.rmtree(root, ignore_errors=True) + raise BundleGateError('bundle content could not be unpacked', gate=5) from error + + @staticmethod + def _extract_member( + archive: zipfile.ZipFile, entry: zipfile.ZipInfo, root: Path, resolved_root: Path + ) -> None: + """Write one member, re-checking containment on the resolved path.""" + name = entry.filename.replace('\\', '/') + target = root / name + resolved = Path(os.path.realpath(target)) + if resolved != resolved_root and resolved_root not in resolved.parents: + raise BundleGateError('archive member resolves outside the extraction root', gate=5) + if entry.is_dir(): + resolved.mkdir(parents=True, exist_ok=True) + return + if resolved.exists(): + raise BundleGateError('archive member would overwrite an existing file', gate=5) + resolved.parent.mkdir(parents=True, exist_ok=True) + # Written as a plain file: the archive's stored permissions are never restored, so no + # member can arrive executable. + with archive.open(entry) as source, resolved.open('wb') as sink: + shutil.copyfileobj(source, sink) + + # ------------------------------------------------------------------ gate 6 + + def _gate_6_structure(self, root: Path) -> tuple[dict[str, Any], dict[str, Any]]: + """Validate the extracted layout and the metadata schema.""" + metadata_path = root / METADATA_FILE_NAME + artifacts_dir = root / ARTIFACTS_DIR_NAME + if not metadata_path.is_file(): + raise BundleGateError('bundle has no metadata document at its root', gate=6) + if not artifacts_dir.is_dir(): + raise BundleGateError('bundle has no artifacts directory', gate=6) + + document = self._load_metadata_document(metadata_path) + metadata = document['metadata'] + parameters = document['parameters'] + + self._check_metadata_fields(metadata) + self._check_parameters(parameters) + self._check_model_directories(artifacts_dir) + return metadata, parameters + + @staticmethod + def _load_metadata_document(metadata_path: Path) -> dict[str, Any]: + """Read `metadata.json` and check its two top-level keys.""" + try: + document = json.loads(metadata_path.read_text(encoding='utf-8')) + except (json.JSONDecodeError, UnicodeDecodeError, OSError) as error: + raise BundleGateError('bundle metadata is not readable JSON', gate=6) from error + if not isinstance(document, dict) or set(document) != REQUIRED_TOP_LEVEL_KEYS: + raise BundleGateError( + 'bundle metadata does not carry exactly the parameters and metadata blocks', gate=6 + ) + if not isinstance(document['metadata'], dict) or not isinstance( + document['parameters'], dict + ): + raise BundleGateError('bundle metadata blocks are not objects', gate=6) + return document + + @staticmethod + def _check_metadata_fields(metadata: dict[str, Any]) -> None: + """Every origin field the import records must be present and usable.""" + for field_name in REQUIRED_METADATA_FIELDS: + value = metadata.get(field_name) + if not isinstance(value, str) or not value.strip(): + raise BundleGateError(f'bundle metadata is missing the {field_name} field', gate=6) + experiment_name = metadata['experiment_name'] + if not (EXPERIMENT_NAME_MIN_LENGTH <= len(experiment_name) <= EXPERIMENT_NAME_MAX_LENGTH): + raise BundleGateError( + "bundle experiment name does not fit the platform's limits", gate=6 + ) + + @staticmethod + def _check_parameters(parameters: dict[str, Any]) -> None: + """`target_variable` is what `model_config.target` is built from; nothing else supplies it.""" + for name in REQUIRED_PARAMETERS: + value = parameters.get(name) + if value is None or not str(value).strip(): + raise BundleGateError(f'bundle parameters are missing {name}', gate=6) + + @staticmethod + def _check_model_directories(artifacts_dir: Path) -> None: + """`prediction_model` and `data_model` must each be a loadable MLflow model directory.""" + for directory in REQUIRED_MODEL_DIRS: + model_dir = artifacts_dir / directory + if not model_dir.is_dir(): + raise BundleGateError(f'bundle has no {directory} artifacts', gate=6) + for required in REQUIRED_MODEL_FILES: + if not (model_dir / required).is_file(): + raise BundleGateError(f'bundle {directory} artifacts are incomplete', gate=6) + + # ------------------------------------------------------------------ gate 7 + + 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('*')): + if path.is_dir(): + continue + if not path.is_file(): + raise BundleGateError('bundle contains an entry that is not a file', gate=7) + if path.name in ALLOWED_FILE_NAMES: + continue + if path.suffix.lower() in ALLOWED_FILE_SUFFIXES: + continue + raise BundleGateError('bundle contains a file a model export should not carry', gate=7) diff --git a/laborious/utils/bundle/steps.py b/laborious/utils/bundle/steps.py new file mode 100644 index 0000000..35c2506 --- /dev/null +++ b/laborious/utils/bundle/steps.py @@ -0,0 +1,55 @@ +"""The import pipeline's step vocabulary, and the gate-to-step mapping. + +`ImportStep` is the join between three things: the machine-readable step written into the import +record's JSON detail, the key of the error catalog (one code and one sentence per step) and the +value the frontend switches on. It lives in the bundle package rather than in the activities module +because the reader raises with it — `laborious/activities/model_import.py` imports it from here so +there is exactly one definition. + +It holds **pipeline** steps only. The two failures that are not steps of the pipeline — failing to +claim the record and failing to write it — are reported by the failure channel under their own +markers (`not_claimable`, `status_write`), deliberately outside this enum so the +catalog-completeness test keeps meaning "every pipeline step has a sentence". +""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Final + + +class ImportStep(StrEnum): + """Where an import stopped. Values are what gets written to the record.""" + + RECEIVED = 'received' + DOWNLOAD = 'download' + DECRYPTION = 'decryption' + ARCHIVE_INSPECTION = 'archive_inspection' + EXTRACTION = 'extraction' + STRUCTURE_VALIDATION = 'structure_validation' + CONTENT_POLICY = 'content_policy' + EXPERIMENT_CREATION = 'experiment_creation' + ARTIFACT_UPLOAD = 'artifact_upload' + REGISTRATION = 'registration' + MODEL_DOCUMENT = 'model_document' + CLEANUP = 'cleanup' + + +# The seven gates map onto the first seven steps, with gates 1 and 2 sharing one: both say "these +# are not the bytes that were sent" — an object over the size ceiling, a key with the wrong shape, a +# digest that disagrees with what the uploader computed, a wrong magic, an unsupported version, a +# chunk stream that does not match the header's digest. The catalog's `download` sentence ("The +# uploaded file could not be read, or it is not the file that was sent. Upload it again.") is true +# for all of them, and none of them may borrow the `decryption` sentence, which implicates the +# password. +GATE_TO_STEP: Final[dict[int, ImportStep]] = { + 1: ImportStep.DOWNLOAD, + 2: ImportStep.DOWNLOAD, + 3: ImportStep.DECRYPTION, + 4: ImportStep.ARCHIVE_INSPECTION, + 5: ImportStep.EXTRACTION, + 6: ImportStep.STRUCTURE_VALIDATION, + 7: ImportStep.CONTENT_POLICY, +} + +GATE_COUNT: Final[int] = 7 diff --git a/laborious/utils/connectors_config.py b/laborious/utils/connectors_config.py index ed3477a..aa2a019 100644 --- a/laborious/utils/connectors_config.py +++ b/laborious/utils/connectors_config.py @@ -1,5 +1,6 @@ import json -from os import getenv +import tempfile +from os import getenv, path from typing import Any @@ -90,3 +91,72 @@ def build_minio_config() -> dict[str, Any]: 'retention_hours': int(getenv('MINIO_RETENTION_HOURS', '24')), 'secure': getenv('MINIO_SECURE', 'false') == 'true', } + + +def build_import_config() -> dict[str, Any]: + """ + Build the `.sientia` model-import configuration from environment variables. + + The password key is read as the raw base64 string and **not** validated here: the worker must + start whether or not `IMPORT_MODEL_ENABLED` is set, and a missing or malformed key is an import + failure at the decryption step (`laborious.utils.import_password.load_envelope_key`), not a + boot failure. + + Environment Variables: + IMPORT_PASSWORD_KEY: base64 of the 32-byte AES-256-GCM key that wraps bundle passwords + IMPORT_BUNDLE_BUCKET: bucket holding the uploaded bundles (default: sientia) + IMPORT_BUNDLE_PREFIX: key prefix the uploads must sit under (default: imported_models/) + IMPORT_WORK_DIR: worker scratch directory (default: /sientia-import) + IMPORT_MAX_OBJECT_BYTES: gate 1 object size ceiling (default: 1 GiB) + IMPORT_MAX_ARCHIVE_ENTRIES: gate 4 entry count ceiling (default: 5000) + 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_BUNDLE_RETENTION_DAYS: lifecycle expiry for the uploaded bundle (default: 7) + + Returns: + dict: import configuration dictionary + """ + prefix = getenv('IMPORT_BUNDLE_PREFIX', 'imported_models/') + return { + 'password_key': getenv('IMPORT_PASSWORD_KEY'), + 'bucket': getenv('IMPORT_BUNDLE_BUCKET', 'sientia'), + 'prefix': prefix if prefix.endswith('/') else f'{prefix}/', + 'work_dir': getenv('IMPORT_WORK_DIR', path.join(tempfile.gettempdir(), 'sientia-import')), + 'max_object_bytes': int(getenv('IMPORT_MAX_OBJECT_BYTES', str(1024 * 1024 * 1024))), + 'max_entries': int(getenv('IMPORT_MAX_ARCHIVE_ENTRIES', '5000')), + 'max_uncompressed_bytes': int( + getenv('IMPORT_MAX_UNCOMPRESSED_BYTES', str(4 * 1024 * 1024 * 1024)) + ), + 'max_compression_ratio': float(getenv('IMPORT_MAX_COMPRESSION_RATIO', '200')), + 'models_collection': getenv('IMPORT_MODELS_COLLECTION', 'models'), + 'retention_days': int(getenv('IMPORT_BUNDLE_RETENTION_DAYS', '7')), + } + + +def build_import_status_config() -> dict[str, Any]: + """ + Build the connection to the import log database — a **second** Postgres connection. + + The import log is `public.experiment_run` in the database owned by the Spring Boot BFF + (`sientia-core-mlops-bff`), not the `sientia` database Laborious itself uses. The table and + schema are constants of the code, never inputs; only the connection is configuration. + + Environment Variables: + IMPORT_STATUS_DB_HOST: host (default: localhost) + IMPORT_STATUS_DB_PORT: port (default: 5432) + IMPORT_STATUS_DB_NAME: database (default: sientia-core-mlops-bff) + IMPORT_STATUS_DB_USER: user (default: mlops_bff_user) + IMPORT_STATUS_DB_PASSWORD: password (default: empty) + + Returns: + dict: import status database configuration dictionary + """ + return { + 'host': getenv('IMPORT_STATUS_DB_HOST', 'localhost'), + 'port': int(getenv('IMPORT_STATUS_DB_PORT', '5432')), + 'dbname': getenv('IMPORT_STATUS_DB_NAME', 'sientia-core-mlops-bff'), + 'user': getenv('IMPORT_STATUS_DB_USER', 'mlops_bff_user'), + 'password': getenv('IMPORT_STATUS_DB_PASSWORD', ''), + 'schema': 'public', + } diff --git a/laborious/utils/import_password.py b/laborious/utils/import_password.py new file mode 100644 index 0000000..5b61e8b --- /dev/null +++ b/laborious/utils/import_password.py @@ -0,0 +1,98 @@ +"""The bundle password envelope: AES-256-GCM, opened inside the activity that uses the password. + +The workflow only ever carries `password_envelope`, so no plaintext password reaches the Temporal +event history. This module is imported by the bundle-opening activity and by nothing else — in +particular not by `laborious/workflows/import_model.py`, which a test asserts. + +Envelope layout, chosen because WebCrypto gives the Angular frontend AES-GCM natively: + + base64( nonce(12 bytes) || ciphertext || GCM tag(16 bytes) ) + +The key is a 32-byte value held in a Kubernetes Secret and exposed as `IMPORT_PASSWORD_KEY`, +base64-encoded. Neither the key, the envelope nor the recovered password is ever logged, notified, +used as a metric label or written to the import record. +""" + +from __future__ import annotations + +import base64 +import binascii +from typing import Final + +from cryptography.exceptions import InvalidTag +from cryptography.hazmat.primitives.ciphers.aead import AESGCM + +NONCE_SIZE: Final[int] = 12 +TAG_SIZE: Final[int] = 16 +KEY_SIZE: Final[int] = 32 +MIN_ENVELOPE_SIZE: Final[int] = NONCE_SIZE + TAG_SIZE + + +class EnvelopeError(Exception): + """The envelope or its key is unusable. + + Deliberately one type for every cause — missing key, wrong key, malformed base64, altered + bytes — because the failure is reported at the decryption step either way. What it must never + be confused with is a wrong *bundle* password: that one comes from the reader's gate 3, and the + two carry different sentences. + """ + + +def load_envelope_key(raw_key: str | None) -> bytes: + """Decode the configured envelope key. + + Args: + raw_key: the value of `IMPORT_PASSWORD_KEY` — base64 of exactly 32 bytes. + + Returns: + bytes: the 32-byte key. + + Raises: + EnvelopeError: when the value is absent, not base64, or not 32 bytes long. There is no + fallback interpretation: an unset key never means "the envelope is plaintext". + """ + if not raw_key: + raise EnvelopeError('the import password key is not configured') + try: + key = base64.b64decode(raw_key, validate=True) + except (binascii.Error, ValueError) as error: + raise EnvelopeError('the import password key is not valid base64') from error + if len(key) != KEY_SIZE: + raise EnvelopeError(f'the import password key must be {KEY_SIZE} bytes') + return key + + +def decrypt_password_envelope(envelope: str, key: bytes) -> str: + """Recover the bundle password from its envelope. + + Args: + envelope: base64 of `nonce || ciphertext || tag`. + key: the 32-byte key from `load_envelope_key`. + + Returns: + str: the plaintext password, exactly as it was typed — not trimmed, not case-folded. NFC + normalisation happens at key derivation, where the producer does it. + + Raises: + EnvelopeError: malformed envelope, wrong key or altered bytes (AES-GCM authentication). + """ + if not envelope: + raise EnvelopeError('the import request carries no password envelope') + try: + payload = base64.b64decode(envelope, validate=True) + except (binascii.Error, ValueError) as error: + raise EnvelopeError('the password envelope is not valid base64') from error + if len(payload) < MIN_ENVELOPE_SIZE: + raise EnvelopeError('the password envelope is too short to hold a nonce and a tag') + + nonce, sealed = payload[:NONCE_SIZE], payload[NONCE_SIZE:] + try: + plaintext = AESGCM(key).decrypt(nonce, sealed, None) + except InvalidTag as error: + raise EnvelopeError('the password envelope could not be authenticated') from error + except ValueError as error: + raise EnvelopeError('the password envelope could not be opened') from error + try: + return plaintext.decode('utf-8') + except UnicodeDecodeError as error: + raise EnvelopeError('the recovered password is not valid UTF-8') from error diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 7a420ba..9f76141 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -10,6 +10,8 @@ The worker supports multiple runtime-scoped task queues (via ``sientia_do.tempor - minimal_retrain-{runtime}-queue: Model retraining workflows - drift-{runtime}-queue: Drift detection workflows - simple_metrics-{runtime}-queue: Simple metrics workflows +- import_model-{runtime}-queue: `.sientia` model import — registered **only** when + ``IMPORT_MODEL_ENABLED`` is truthy, so the containment boundary stays a deployment decision ``RUNTIME`` must be set; it is passed to every ``prepare_worker`` call. Schedulers must use the same queue names (breaking change vs legacy ``drift-queue`` / ``simple_metrics-queue``). @@ -30,6 +32,9 @@ Environment Variables: - HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090) - HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091) - PROJECT_NAME: Project name for notifications (default: laborious) +- IMPORT_MODEL_ENABLED: register the fifth (import) worker; unset by default +- IMPORT_PASSWORD_KEY / IMPORT_BUNDLE_* / IMPORT_STATUS_DB_*: import configuration, see + ``laborious.utils.connectors_config`` """ from temporalio import client, workflow @@ -53,11 +58,14 @@ with workflow.unsafe.imports_passed_through(): from laborious import metrics from laborious.activities.activities import Activities from laborious.utils.connectors_config import ( + build_import_config, + build_import_status_config, build_minio_config, build_mlflow_config, build_opc_config, ) from laborious.workflows.drift import Drift + from laborious.workflows.import_model import ImportModel from laborious.workflows.minimal_retrain import MinimalRetrain from laborious.workflows.predictions_batch import PredictionsBatch from laborious.workflows.simple_metrics import SimpleMetrics @@ -137,6 +145,9 @@ async def main(): pi_web_api_config=build_api_config(), logger=logger, notification_handler=notification_handler, + mongo_config=mongo_config, + import_config=build_import_config(), + import_status_config=build_import_status_config(), ) logger.custom_info('Initializing OPC...', metadata_runtime) @@ -235,6 +246,42 @@ async def main(): ), ] + # The fifth worker is opt-in: containment (QTZPOC-19) wants the importer in its own runtime with + # its own credentials, and a flag lets that be a deployment decision instead of forcing every + # Laborious runtime to poll the import queue today. Unset, the four workers above are unaffected + # and no import queue is polled at all. + if os.getenv('IMPORT_MODEL_ENABLED', '').strip().lower() in ('1', 'true', 'yes', 'on'): + logger.custom_info('Import model worker enabled', metadata_runtime) + workers.append( + prepare_worker( + temporal_client=temporal_client, + main_workflow=ImportModel, + other_workflows=[], + activities=[ + # Import log row (the BFF database, update-only) + activities.claim_import_status, + activities.record_import_status, + activities.record_import_names, + activities.record_import_terminal_status, + activities.report_import_status_write_failure, + # Bundle + activities.download_import_bundle, + activities.open_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, + # Cleanup and retention + activities.cleanup_import_files, + activities.ensure_import_bundle_retention, + ], + logger=logger, + runtime=runtime, + ) + ) + handlers = [] for w in workers: handlers.append(w.run()) diff --git a/laborious/workflows/import_model.py b/laborious/workflows/import_model.py new file mode 100644 index 0000000..39cbbf6 --- /dev/null +++ b/laborious/workflows/import_model.py @@ -0,0 +1,416 @@ +"""The `import_model` workflow: a `.sientia` object becomes a registered model, or a durable ERROR. + +Shape of this workflow, and the reasons it has that shape: + +- **The claim comes first.** The record already exists — the frontend inserted it through the BFF — + so the workflow's first act is to take ownership of it. Everything after that failure is recorded + *on* the record; a failed claim is recorded nowhere, because the workflow owns no row. +- **Fail-fast, and nothing else.** No rollback, no compensation, no resume, no reconciliation, no + `continue_as_new`. A failed import stays failed; trying again is a new import with a new upload. +- **Only the status writes retry**, and only because `UPDATE ... SET WHERE id` is + idempotent. Every provisioning call keeps `maximum_attempts=1`, because a repeated registry call or + Mongo insert can write twice. +- **One `finally`, two actions**: cleanup, then the single terminal write that classifies the + outcome once. No `except` block writes a verdict of its own. +- **The workflow holds ids, never contents.** The decrypted bundle never crosses an activity + boundary, and the password reaches this module only as an envelope it cannot open. +""" + +from datetime import timedelta + +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities + from laborious.activities.model_import import ( + IMPORT_STATUS_HINT_RETRY_POLICY, + IMPORT_STATUS_RETRY_POLICY, + derive_run_name, + ) + from laborious.activities.model_import_errors import ImportInputError, failure_location + from laborious.utils.bundle.steps import ImportStep + +# `start_to_close_timeout` is per attempt. The terminal write additionally gets a +# `schedule_to_close_timeout` larger than its retry envelope (10 attempts, ≈151 s of backoff, ≈251 s +# if every attempt also burns its timeout) — set it smaller and Temporal cuts the retries short, +# which is the one configuration mistake that would silently undo design D20. +STATUS_WRITE_TIMEOUT = timedelta(seconds=10) +STATUS_WRITE_ENVELOPE = timedelta(seconds=300) +DOWNLOAD_TIMEOUT = timedelta(seconds=600) +OPEN_TIMEOUT = timedelta(seconds=900) +EXPERIMENT_TIMEOUT = timedelta(seconds=60) +ARTIFACT_UPLOAD_TIMEOUT = timedelta(seconds=900) +REGISTRATION_TIMEOUT = timedelta(seconds=120) +DOCUMENT_TIMEOUT = timedelta(seconds=60) +CLEANUP_TIMEOUT = timedelta(seconds=120) +RETENTION_TIMEOUT = timedelta(seconds=60) + +# Fields no import may carry. The project is a frontend concern: the model listing is the MongoDB +# `models` document, no migrated database holds a project or model table, and nothing in this +# repository reads a project-to-model link (design D22). An input naming one is rejected rather than +# accepted and ignored, so no caller believes a project link was recorded. The same goes for a +# relational target: there is none, and for a plaintext password, which does not belong in an event +# history whatever else is true. +REJECTED_INPUT_FIELDS = ( + 'project_id', + 'project_name', + 'provisioning_target', + 'schema', + 'table_name', + 'status_table', + 'password', +) + +DIGEST_LENGTH = 64 +_HEX_DIGITS = frozenset('0123456789abcdef') + + +@workflow.defn(name='import_model') +class ImportModel: + """Turn an encrypted `.sientia` object in MinIO into a registered model, or into an ERROR.""" + + @workflow.run + async def run(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Execute the import. + + Args: + input_data: the import request. + Required keys: + - import_run_id (int): primary key of the import log row the frontend inserted + - object_key (str): the uploaded object's key + - expected_digest (str): the SHA-256 the uploader computed for those bytes + - password_envelope (str): the encrypted bundle password, never plaintext + Optional keys: + - bucket (str): defaults to the configured import bucket + + Returns: + dict: `import_run_id`, `run_id`, `model_name`, `version` and the model document's id. + + Raises: + ImportLogRowNotClaimableError: the record is not this workflow's — the run ends FAILED + with that error type and nothing is written to any row. + Exception: whatever failed. The record carries the failing step, its code and its + sentence; this exception carries the technical truth. + """ + workflow_id = workflow.info().workflow_id + import_run_id = self._require_import_run_id(input_data.get('import_run_id')) + metadata: dict[str, Any] = { + 'model_name': '-', + 'model_id': str(import_run_id), + 'workflow_name': 'import_model', + 'schedule_name': '-', + } + base = { + 'metadata': metadata, + 'import_run_id': import_run_id, + 'workflow_id': workflow_id, + } + + # Class A. Outside the try/finally on purpose: a claim that matches nothing leaves this + # workflow owning no row, so there is nothing to clean up and nothing to record. + claim = await workflow.execute_activity_method( + Activities.claim_import_status, + base, + retry_policy=IMPORT_STATUS_RETRY_POLICY, + start_to_close_timeout=STATUS_WRITE_TIMEOUT, + schedule_to_close_timeout=STATUS_WRITE_ENVELOPE, + ) + + current_step = ImportStep.RECEIVED + failed_step: str | None = None + gate: int | None = None + succeeded = False + source: dict[str, Any] = {} + result: dict[str, Any] = {'import_run_id': import_run_id} + + try: + self._validate_input(input_data, claim) + + 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, + ) + + 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, + ) + source = opened['metadata'] + parameters = opened['parameters'] + model_name = source['model_name'] + metadata['model_name'] = model_name + run_name = derive_run_name(model_name, source['model_version']) + + # 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, + { + **base, + 'experiment_id': experiment['experiment_id'], + 'run_name': run_name, + 'extracted_dir': opened['extracted_dir'], + 'parameters': parameters, + 'source': source, + }, + retry_policy=retry_policy, + start_to_close_timeout=ARTIFACT_UPLOAD_TIMEOUT, + ) + + current_step = ImportStep.REGISTRATION + await self._hint(base, current_step, source) + version = await workflow.execute_activity_method( + Activities.register_import_model_version, + {**base, 'model_name': model_name, 'run_id': run['run_id']}, + retry_policy=retry_policy, + start_to_close_timeout=REGISTRATION_TIMEOUT, + ) + + current_step = ImportStep.MODEL_DOCUMENT + await self._hint(base, current_step, source) + document = await workflow.execute_activity_method( + Activities.write_import_model_document, + { + **base, + 'model_name': model_name, + 'target': parameters['target_variable'], + }, + retry_policy=retry_policy, + start_to_close_timeout=DOCUMENT_TIMEOUT, + ) + + succeeded = True + result = { + 'import_run_id': import_run_id, + 'run_id': run['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 + gate = located_gate + raise + finally: + cleanup = await self._cleanup(base) + await self._record_terminal( + base, + succeeded=succeeded, + step=failed_step or current_step.value, + gate=gate, + source=source, + cleanup_failed=not cleanup.get('cleaned', True), + ) + + return result + + # ------------------------------------------------------------------ validation + + @staticmethod + def _require_import_run_id(value: Any) -> int: + """The one check that runs before the claim, because the claim needs its result. + + A missing or unusable id means there is no record to claim and none to write: the run ends + FAILED with this error type, and the failure is visible in the workflow's terminal state, its + error type, the error metric and a notification — everywhere except a row, which is the + point. + """ + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ImportInputError('import_run_id must be a positive integer') + return value + + @staticmethod + def _validate_input(input_data: dict[str, Any], claim: dict[str, Any]) -> None: + """Validate the rest of the request, and check it against the record it just claimed. + + This runs **after** the claim so its failure lands on the record rather than only in the + logs: by now the row is this workflow's, so a bad request is reported to the person who made + it. Everything here fails at step `received`. + """ + unknown = [name for name in REJECTED_INPUT_FIELDS if name in input_data] + if unknown: + raise ImportInputError( + f'the import request carries fields this workflow does not accept: ' + f'{", ".join(sorted(unknown))}' + ) + + for required in ('object_key', 'expected_digest', 'password_envelope'): + if not input_data.get(required): + raise ImportInputError(f'the import request is missing {required}') + + digest = str(input_data['expected_digest']).lower() + if len(digest) != DIGEST_LENGTH or not set(digest) <= _HEX_DIGITS: + raise ImportInputError('expected_digest must be a 64-character hexadecimal SHA-256') + + object_key = input_data['object_key'] + file_name = claim.get('file_name') + if file_name and file_name != object_key: + raise ImportInputError('the record names a different object than the request') + + request_digest = claim.get('request_digest') + if request_digest and request_digest != digest: + raise ImportInputError('the record names a different digest than the request') + + # --------------------------------------------------------------- status writes + + async def _hint( + self, + base: dict[str, Any], + step: ImportStep, + source: dict[str, Any] | None = None, + ) -> None: + """Record the step about to be attempted. 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}, + 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], + *, + succeeded: bool, + step: str | None, + gate: int | None, + source: dict[str, Any], + cleanup_failed: bool, + ) -> None: + """The single authoritative write, and the last thing this workflow does. + + Its own failure is reported through the channel with the verdict it could not write, and it + never turns a completed import into a failed one — nor masks the cause of a real failure. + """ + try: + await workflow.execute_activity_method( + Activities.record_import_terminal_status, + { + **base, + 'succeeded': succeeded, + 'step': step, + 'gate': gate, + 'source': source or None, + 'cleanup_failed': cleanup_failed, + }, + retry_policy=IMPORT_STATUS_RETRY_POLICY, + start_to_close_timeout=STATUS_WRITE_TIMEOUT, + schedule_to_close_timeout=STATUS_WRITE_ENVELOPE, + ) + except BaseException as error: # NOSONAR - the outcome stands even when it cannot be written + await self._report_status_write_failure( + base, + error, + verdict=f'succeeded={succeeded} step={step} gate={gate}', + ) + + async def _report_status_write_failure( + self, + base: dict[str, Any], + error: BaseException, + verdict: str | None = None, + ) -> None: + """Route a status write that could not be made through the one failure channel.""" + try: + await workflow.execute_activity_method( + Activities.report_import_status_write_failure, + {**base, 'detail': str(error), 'verdict': verdict}, + retry_policy=retry_policy, + start_to_close_timeout=STATUS_WRITE_TIMEOUT, + ) + except BaseException: # NOSONAR - reporting the report is where this stops + pass + + # ----------------------------------------------------------- side-effect steps + + @staticmethod + async def _ensure_retention(base: dict[str, Any], input_data: dict[str, Any]) -> None: + """Ensure the bundle bucket's expiry rule, without making it a reason to fail an import. + + The rule is applied by the importer rather than assumed from a Helm value, so the acceptance + criterion is verifiable from the code that depends on it. But a bucket policy the importer + cannot set is an operational problem, not a bad import: failing a valid model import because + a lifecycle API refused would be the wrong trade, so this is reported and dropped like a + lost status hint. + """ + try: + await workflow.execute_activity_method( + Activities.ensure_import_bundle_retention, + {**base, 'bucket': input_data.get('bucket')}, + retry_policy=retry_policy, + start_to_close_timeout=RETENTION_TIMEOUT, + ) + except BaseException: # NOSONAR - retention is not a precondition of a correct import + pass + + @staticmethod + async def _cleanup(base: dict[str, Any]) -> dict[str, Any]: + """Remove this import's temporary files. Runs on every path, before the terminal write.""" + try: + return await workflow.execute_activity_method( + Activities.cleanup_import_files, + base, + retry_policy=retry_policy, + start_to_close_timeout=CLEANUP_TIMEOUT, + ) + except BaseException: # NOSONAR - cleanup never replaces the failure that brought us here + return {'cleaned': False} diff --git a/requirements-local.txt b/requirements-local.txt index 852328a..8d060e3 100644 --- a/requirements-local.txt +++ b/requirements-local.txt @@ -4,11 +4,17 @@ sqlalchemy asyncua==1.0.6 redis git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.12.1 -git+ssh://git@github.com/Aignosi/sientia-model-library.git@0.10.0 +/home/grezewave/Documents/projects/sientia/sientia-mlops-library prometheus-client -botocore -boto3 -s3fs +# fsspec/s3fs/gcsfs pinned as a set: nannyml (via the mlops library) pulls both s3fs and +# gcsfs, and each pins fsspec exactly, so a free `s3fs` makes the resolver fail against +# evidently's fsspec>=2024.2.0. Versions match the library-distribution-controller freeze. +fsspec==2026.7.0 +s3fs==2026.7.0 +gcsfs==2026.8.0 +# aiobotocore 3.9.0 (pulled by s3fs) requires botocore<1.43.57,>=1.43.3. +botocore==1.43.56 +boto3==1.43.56 pyarrow kaleido hyperopt diff --git a/requirements.txt b/requirements.txt index 4b2a2a8..5b70c81 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,11 +4,17 @@ sqlalchemy asyncua==1.0.6 redis sientia_do>=1.12.1 -sientia>0.40.0 +sientia>=0.41.0 prometheus-client -botocore -boto3 -s3fs +# fsspec/s3fs/gcsfs pinned as a set: nannyml (via sientia) pulls both s3fs and gcsfs, and +# each pins fsspec exactly, so a free `s3fs` makes the resolver fail against evidently's +# fsspec>=2024.2.0. Versions match the library-distribution-controller freeze. +fsspec==2026.7.0 +s3fs==2026.7.0 +gcsfs==2026.8.0 +# aiobotocore 3.9.0 (pulled by s3fs) requires botocore<1.43.57,>=1.43.3. +botocore==1.43.56 +boto3==1.43.56 pyarrow kaleido hyperopt @@ -18,3 +24,10 @@ scipy<1.14.0 scikit-learn==1.4.2 pycaret==3.3.2 joblib +# .sientia bundle import (Argon2id + XChaCha20-Poly1305 secretstream): same pin as the +# producer in sientia-projects-templates, so both sides open the same bytes. +pynacl==1.5.0 +# AES-256-GCM password envelope; already present transitively, declared explicitly. +cryptography +# MongoDB models collection (the model listing); transitive via sientia_do, declared explicitly. +pymongo diff --git a/tests/fixtures/bundle/README.md b/tests/fixtures/bundle/README.md new file mode 100644 index 0000000..7a5cad8 --- /dev/null +++ b/tests/fixtures/bundle/README.md @@ -0,0 +1,86 @@ +# Golden `.sientia` fixtures + +Binary conformance vectors for `laborious/utils/bundle/`. They are **bytes the producer's own code +emitted**, not bytes this repository knows how to write: there is no production writer here (change +`import-model-workflow`, design D1). If the producer's layout moves, the reader's tests fail against +these files instead of failing in production. + +| File | What it is | +|---|---| +| `golden_model_v3.sientia` | `format_version` 1 bundle, 110767 bytes, plaintext zip 110569 bytes — **two chunks**, so `MESSAGE`/`FINAL` framing is exercised by the golden fixture itself | +| `golden_model_v3.metadata.json` | the `metadata.json` that went into that bundle, in plaintext, so a test can assert the recovered value without decrypting first | +| `regenerate.py` | the script that produced them, importing the producer's `bundle.py`/`crypto.py` by path | + +## Password + +``` +Import1Model!Fixture +``` + +Recorded deliberately: the fixture protects nothing. It satisfies the producer's export-side +strength rule (≥ 12 characters, upper, lower, digit, symbol), which is checked by +`crypto.validate_password_strength` at generation time. + +## Provenance + +- Producer repository: `sientia-projects-templates` +- Producer commit: `8d3c1fd117c697e8b352c34f11c732809a4a129b` + (`8d3c1fd QTZPOC-7: adapt to lint rules picked up from merging main`) +- Producer modules used: `app/src/operations/model_export/bundle.py`, + `app/src/operations/model_export/crypto.py` — loaded by path, so every constant, the `struct` + layout, the zip layout, the KDF parameters and the chunk framing come from the producer and none of + them is re-implemented here +- Producer pin honoured: `pynacl==1.5.0` (this repository pins the same version) + +One deviation from task 0.3, recorded rather than hidden: the bytes were generated with the +producer's **code** under *this* repository's venv, not inside the producer's own venv, because that +venv has neither `pynacl` nor `mlflow` installed. The dependency that matters for the bytes is +`pynacl`, pinned identically on both sides, and `mlflow` is not reached — `regenerate.py` builds the +artifact tree itself instead of calling `client.download_artifacts`. + +## Contents of the bundle + +``` +metadata.json (written first, no indent, by bundle.build_zip) +artifacts/prediction_model/MLmodel +artifacts/prediction_model/model.pkl (90 KiB of incompressible filler) +artifacts/prediction_model/conda.yaml +artifacts/prediction_model/python_env.yaml +artifacts/prediction_model/requirements.txt +artifacts/data_model/MLmodel +artifacts/data_model/model.pkl (8 KiB) +artifacts/data_model/conda.yaml +artifacts/data_model/python_env.yaml +artifacts/data_model/requirements.txt +artifacts/model_card/model_card.json +artifacts/train_data.csv +artifacts/test_data.csv +artifacts/raw_data.csv +``` + +No directory entries are stored — the producer only adds paths for which `is_file()` is true. + +`parameters` carries `target_variable: "SE"`, which gate 6 requires and which becomes +`model_config.target` in the MongoDB model document. + +## Digest boundary — the reason this fixture exists + +The header's `ciphertext_digest` is `sha256(file[164:])`: the producer computes +`hashlib.sha256(ciphertext)` over the joined chunks and only then writes +`header + stream_header + ciphertext`, so the 24-byte `secretstream` header is **not** covered. +Verified against this fixture and against a freshly produced bundle. Earlier prose (this change's own +first draft, and `aig-docs/03-sientia-bundle-format.md`) said "everything after the plaintext +header", i.e. from offset 140 — that is wrong, and `tests/laborious/utils/bundle/test_format.py` +asserts both directions so it cannot come back. + +## Regenerating + +```bash +source ./venv/bin/activate +PRODUCER_REPO=../sientia-projects-templates python tests/fixtures/bundle/regenerate.py +``` + +Two fields change on every run and are expected to: the 16-byte salt (`os.urandom`) and +`metadata.metadata.export_timestamp` (`datetime.now(UTC)`). Everything else is deterministic — +the filler bytes come from a seeded PRNG. After regenerating, update the producer commit above and +the byte counts in the table. diff --git a/tests/fixtures/bundle/golden_model_v3.metadata.json b/tests/fixtures/bundle/golden_model_v3.metadata.json new file mode 100644 index 0000000..cf30161 --- /dev/null +++ b/tests/fixtures/bundle/golden_model_v3.metadata.json @@ -0,0 +1,24 @@ +{ + "parameters": { + "model_name": "fixture_model", + "target_variable": "SE", + "input_variables": "['feature_a', 'feature_b']", + "models_params": "{'alpha': 0.1}", + "nan_treatment": "drop", + "scaler_name": "StandardScaler", + "train_size": "0.8", + "test_size": "0.2", + "include_ar": "False", + "retrain": "False", + "start_date": "2026-01-01 00:00:00", + "end_date": "2026-08-01 00:00:00" + }, + "metadata": { + "model_name": "fixture_model", + "experiment_name": "fixture_experiment", + "run_id": "8f14e45fceea167a5a36dedd4bea2543", + "model_version": "3", + "model_project": "fixture_project", + "export_timestamp": "2026-08-25T12:08:59.533441+00:00" + } +} diff --git a/tests/fixtures/bundle/golden_model_v3.sientia b/tests/fixtures/bundle/golden_model_v3.sientia new file mode 100644 index 0000000..2766539 Binary files /dev/null and b/tests/fixtures/bundle/golden_model_v3.sientia differ diff --git a/tests/fixtures/bundle/regenerate.py b/tests/fixtures/bundle/regenerate.py new file mode 100644 index 0000000..078a562 --- /dev/null +++ b/tests/fixtures/bundle/regenerate.py @@ -0,0 +1,185 @@ +"""Regenerate the golden `.sientia` fixtures using the producer's own source. + +This script is the auditable half of the anti-drift contract described in the change +`import-model-workflow` (design D1): the committed fixtures are bytes produced by +`sientia-projects-templates/app/src/operations/model_export/{bundle,crypto}.py`, loaded here by +path so no constant is re-implemented. Regenerating is a deliberate act — see the README next to +this file for the recorded producer commit and the fixture password. + +Usage (from the repository root, inside the venv): + + PRODUCER_REPO=../sientia-projects-templates ./venv/bin/python tests/fixtures/bundle/regenerate.py + +The producer's own venv is not used: it does not have `pynacl` installed, while this repository +pins the producer's exact version (`pynacl==1.5.0`). The bytes are therefore produced by the +producer's code on the producer's dependency pin, which is what the fixture is for. +""" + +from __future__ import annotations + +import importlib.util +import json +import os +import random +import shutil +import subprocess +import sys +import tempfile +import types +from pathlib import Path + +FIXTURE_DIR = Path(__file__).resolve().parent +DEFAULT_PRODUCER_REPO = FIXTURE_DIR.parents[2].parent / 'sientia-projects-templates' + +# Recorded in the README alongside the fixtures; the password satisfies the producer's own +# export-side strength rule (>= 12 chars, upper, lower, digit, symbol). +FIXTURE_PASSWORD = 'Import1Model!Fixture' + +# Deterministic filler so a regeneration differs only in the salt and the export timestamp. +FILLER_SEED = 20260825 + + +def _load_producer_module(repo: Path, name: str) -> types.ModuleType: + """Import one producer module by path, without importing the Streamlit app.""" + path = repo / 'app' / 'src' / 'operations' / 'model_export' / f'{name}.py' + if not path.is_file(): + raise SystemExit(f'producer module not found: {path}') + spec = importlib.util.spec_from_file_location(f'producer_{name}', path) + if spec is None or spec.loader is None: + raise SystemExit(f'cannot load producer module: {path}') + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _producer_commit(repo: Path) -> str: + result = subprocess.run( # noqa: S603 - fixed argv, no shell + ['git', '-C', str(repo), 'rev-parse', 'HEAD'], # noqa: S607 - git from PATH is intended + capture_output=True, + text=True, + check=False, + ) + return result.stdout.strip() or 'unknown' + + +def _write_model_dir(root: Path, flavor: str, pickle_bytes: bytes) -> None: + """Write the files an MLflow 2.x model directory carries for this platform.""" + root.mkdir(parents=True, exist_ok=True) + (root / 'MLmodel').write_text( + f'artifact_path: {flavor}\n' + 'flavors:\n' + ' python_function:\n' + ' loader_module: mlflow.sklearn\n' + ' model_path: model.pkl\n' + ' sklearn:\n' + ' pickled_model: model.pkl\n' + ' serialization_format: cloudpickle\n' + ' sklearn_version: 1.4.2\n' + 'mlflow_version: 2.10.1\n' + 'model_uuid: 0f9c1c1f9c1f4c0f9c1c1f9c1f4c0f9c\n' + 'utc_time_created: "2026-08-25 12:00:00.000000"\n', + encoding='utf-8', + ) + (root / 'model.pkl').write_bytes(pickle_bytes) + (root / 'conda.yaml').write_text( + 'channels:\n- conda-forge\ndependencies:\n- python=3.11\n- pip\n- pip:\n' + ' - mlflow==2.10.1\n - scikit-learn==1.4.2\nname: mlflow-env\n', + encoding='utf-8', + ) + (root / 'python_env.yaml').write_text( + 'python: 3.11.15\nbuild_dependencies:\n- pip==24.0\ndependencies:\n- -r requirements.txt\n', + encoding='utf-8', + ) + (root / 'requirements.txt').write_text( + 'mlflow==2.10.1\nscikit-learn==1.4.2\ncloudpickle==3.0.0\n', encoding='utf-8' + ) + + +def _write_csv(path: Path, rows: int, rng: random.Random) -> None: + lines = ['datetime,feature_a,feature_b,SE'] + for i in range(rows): + lines.append( + f'2026-08-{(i % 27) + 1:02d}T00:00:00,' + f'{rng.uniform(0, 100):.6f},{rng.uniform(0, 100):.6f},{rng.uniform(0, 100):.6f}' + ) + path.write_text('\n'.join(lines) + '\n', encoding='utf-8') + + +def _build_artifact_tree(root: Path) -> None: + rng = random.Random(FILLER_SEED) # noqa: S311 - test filler, not key material + # Incompressible filler keeps the zip above the 64 KiB chunk size, so the golden fixture + # exercises multi-chunk `MESSAGE`/`FINAL` framing with real producer bytes. + big_pickle = bytes(rng.getrandbits(8) for _ in range(90 * 1024)) + small_pickle = bytes(rng.getrandbits(8) for _ in range(8 * 1024)) + _write_model_dir(root / 'prediction_model', 'prediction_model', big_pickle) + _write_model_dir(root / 'data_model', 'data_model', small_pickle) + (root / 'model_card').mkdir(parents=True, exist_ok=True) + (root / 'model_card' / 'model_card.json').write_text( + json.dumps( + { + 'model_name': 'fixture_model', + 'target': 'SE', + 'created_by': 'fixture', + 'creation_method': 'time_series_template', + } + ), + encoding='utf-8', + ) + _write_csv(root / 'train_data.csv', 120, rng) + _write_csv(root / 'test_data.csv', 40, rng) + _write_csv(root / 'raw_data.csv', 200, rng) + + +RUN_PARAMETERS = { + 'model_name': 'fixture_model', + 'target_variable': 'SE', + 'input_variables': "['feature_a', 'feature_b']", + 'models_params': "{'alpha': 0.1}", + 'nan_treatment': 'drop', + 'scaler_name': 'StandardScaler', + 'train_size': '0.8', + 'test_size': '0.2', + 'include_ar': 'False', + 'retrain': 'False', + 'start_date': '2026-01-01 00:00:00', + 'end_date': '2026-08-01 00:00:00', +} + +METADATA_FIELDS = { + 'model_name': 'fixture_model', + 'experiment_name': 'fixture_experiment', + 'run_id': '8f14e45fceea167a5a36dedd4bea2543', + 'model_version': '3', + 'model_project': 'fixture_project', +} + + +def main() -> int: + repo = Path(os.getenv('PRODUCER_REPO', str(DEFAULT_PRODUCER_REPO))).resolve() + bundle = _load_producer_module(repo, 'bundle') + crypto = _load_producer_module(repo, 'crypto') + + crypto.validate_password_strength(FIXTURE_PASSWORD) + + staging = Path(tempfile.mkdtemp(prefix='sientia_fixture_')) + try: + _build_artifact_tree(staging) + metadata = bundle.build_metadata_json(run_parameters=RUN_PARAMETERS, **METADATA_FIELDS) + zip_bytes = bundle.build_zip(staging, metadata) + sientia_bytes = crypto.encrypt_archive(zip_bytes, FIXTURE_PASSWORD) + finally: + shutil.rmtree(staging, ignore_errors=True) + + target = FIXTURE_DIR / 'golden_model_v3.sientia' + target.write_bytes(sientia_bytes) + (FIXTURE_DIR / 'golden_model_v3.metadata.json').write_text( + json.dumps(metadata, indent=2) + '\n', encoding='utf-8' + ) + print(f'wrote {target} ({len(sientia_bytes)} bytes, zip {len(zip_bytes)} bytes)') + print(f'producer commit: {_producer_commit(repo)}') + print(f'password: {FIXTURE_PASSWORD}') + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/helpers/bundle_factory.py b/tests/helpers/bundle_factory.py new file mode 100644 index 0000000..db54be5 --- /dev/null +++ b/tests/helpers/bundle_factory.py @@ -0,0 +1,272 @@ +"""Test-only `.sientia` writer, mirrored field for field from the producer. + +There is no production writer in this repository (change `import-model-workflow`, design D1): the +producer is the Streamlit app in `sientia-projects-templates`. This factory exists for one reason — +building the **hostile** variants the gate tests need, which the producer cannot emit. Its output is +pinned against the golden fixture by `test_format.py::test_factory_header_matches_fixture_layout`, +so a drift in either direction shows up as a failing test. + +Mirrored from: +- `app/src/operations/model_export/crypto.py` — header layout, KDF, `secretstream` framing +- `app/src/operations/model_export/bundle.py` — zip layout and the two-top-level-key JSON + +Deliberate deviation: the default KDF cost is a fraction of the producer's (8 MiB / 1 pass instead +of 256 MiB / 3), because the reader honours whatever the header declares and a test suite that +derives thirty 256 MiB Argon2id keys is a test suite nobody runs. The golden fixture is what covers +the production cost parameters. +""" + +from __future__ import annotations + +import base64 +import hashlib +import io +import json +import os +import struct +import unicodedata +import zipfile +from dataclasses import dataclass, field +from typing import Any + +import nacl.bindings as sodium +from cryptography.hazmat.primitives.ciphers.aead import AESGCM + +# --- mirrored from the producer's crypto.py --- +MAGIC = b'SIENTIA1' +FORMAT_VERSION = 1 +AEAD_ID = 1 +KDF_ID = 1 +SALT_SIZE = 16 +RESERVED_SIGNATURE_SIZE = 64 +CHUNK_SIZE = 64 * 1024 +HEADER_STRUCT = struct.Struct('>8sBBB16sQQB32s64s') +HEADER_SIZE = HEADER_STRUCT.size +PRODUCER_KDF_MEMLIMIT_BYTES = 256 * 1024 * 1024 +PRODUCER_KDF_OPSLIMIT = 3 + +# Cheap-but-valid cost parameters for the suite; see the module docstring. +TEST_KDF_MEMLIMIT_BYTES = 8 * 1024 * 1024 +TEST_KDF_OPSLIMIT = 1 + +DEFAULT_PASSWORD = 'Factory1Model!Test' +DEFAULT_ENVELOPE_KEY = bytes(range(32)) + +MODEL_DIRS = ('prediction_model', 'data_model') +MODEL_FILES = ('MLmodel', 'model.pkl', 'conda.yaml', 'python_env.yaml', 'requirements.txt') +CSV_FILES = ('train_data.csv', 'test_data.csv', 'raw_data.csv') + +DEFAULT_METADATA = { + 'model_name': 'factory_model', + 'experiment_name': 'factory_experiment', + 'run_id': 'aa11bb22cc33dd44ee55ff6677889900', + 'model_version': '2', + 'model_project': 'factory_project', + 'export_timestamp': '2026-08-25T12:00:00+00:00', +} + +DEFAULT_PARAMETERS = { + 'model_name': 'factory_model', + 'target_variable': 'SE', + 'input_variables': "['feature_a', 'feature_b']", + 'retrain': 'False', +} + + +def default_tree() -> dict[str, bytes]: + """The artifact tree a well-formed bundle carries, as archive-name -> bytes.""" + files: dict[str, bytes] = {} + for directory in MODEL_DIRS: + for name in MODEL_FILES: + files[f'artifacts/{directory}/{name}'] = f'{directory}/{name}\n'.encode() + files['artifacts/model_card/model_card.json'] = json.dumps({'target': 'SE'}).encode() + for name in CSV_FILES: + files[f'artifacts/{name}'] = b'datetime,feature_a,SE\n2026-08-25T00:00:00,1.0,2.0\n' + return files + + +def build_metadata_document( + *, + metadata: dict[str, Any] | None = None, + parameters: dict[str, Any] | None = None, + drop_metadata_field: str | None = None, + drop_top_level_key: str | None = None, + drop_parameter: str | None = None, +) -> dict[str, Any]: + """Build `metadata.json`, with hooks for the malformed variants gate 6 must reject.""" + meta = {**DEFAULT_METADATA, **(metadata or {})} + params = {**DEFAULT_PARAMETERS, **(parameters or {})} + if drop_metadata_field: + meta.pop(drop_metadata_field, None) + if drop_parameter: + params.pop(drop_parameter, None) + document: dict[str, Any] = {'parameters': params, 'metadata': meta} + if drop_top_level_key: + document.pop(drop_top_level_key, None) + return document + + +@dataclass +class ZipSpec: + """How to lay out the zip inside the bundle.""" + + files: dict[str, bytes] = field(default_factory=default_tree) + document: dict[str, Any] | None = None + include_metadata: bool = True + symlinks: dict[str, str] = field(default_factory=dict) + absolute_entry: str | None = None + traversal_entry: str | None = None + bomb_bytes: int = 0 + extra_files: dict[str, bytes] = field(default_factory=dict) + + +def build_zip(spec: ZipSpec | None = None) -> bytes: + """Zip `artifacts/` plus `metadata.json`, mirroring the producer's `build_zip`. + + `metadata.json` is written first with `json.dumps` and no indent, entries use `/` separators, + and no directory entries are stored — all three matter to the reader. + """ + spec = spec or ZipSpec() + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, 'w', zipfile.ZIP_DEFLATED) as archive: + if spec.include_metadata: + document = spec.document if spec.document is not None else build_metadata_document() + archive.writestr('metadata.json', json.dumps(document)) + for name, content in {**spec.files, **spec.extra_files}.items(): + archive.writestr(_regular_file_info(name), content) + for name, target in spec.symlinks.items(): + info = zipfile.ZipInfo(name) + info.create_system = 3 + info.external_attr = (0o120777 << 16) | 0o200000 + archive.writestr(info, target) + if spec.absolute_entry: + archive.writestr(_regular_file_info(spec.absolute_entry), b'absolute\n') + if spec.traversal_entry: + archive.writestr(_regular_file_info(spec.traversal_entry), b'traversal\n') + if spec.bomb_bytes: + archive.writestr(_regular_file_info('artifacts/bomb.csv'), b'\0' * spec.bomb_bytes) + return buffer.getvalue() + + +def _regular_file_info(name: str) -> zipfile.ZipInfo: + """A zip entry carrying a Unix regular-file mode, as `ZipFile.write` would store it.""" + info = zipfile.ZipInfo(name) + info.create_system = 3 + info.compress_type = zipfile.ZIP_DEFLATED + info.external_attr = 0o100644 << 16 + return info + + +def derive_key(password: str, salt: bytes, opslimit: int, memlimit: int) -> bytes: + """Argon2id over the NFC-normalised password — the producer's `_derive_key`.""" + normalized = unicodedata.normalize('NFC', password) + return sodium.crypto_pwhash_alg( + sodium.crypto_secretstream_xchacha20poly1305_KEYBYTES, + normalized.encode('utf-8'), + salt, + opslimit, + memlimit, + sodium.crypto_pwhash_ALG_ARGON2ID13, + ) + + +def encrypt_archive( # noqa: PLR0913 - every knob here exists to build one hostile variant + plaintext: bytes, + password: str = DEFAULT_PASSWORD, + *, + format_version: int = FORMAT_VERSION, + aead_id: int = AEAD_ID, + kdf_id: int = KDF_ID, + memlimit: int = TEST_KDF_MEMLIMIT_BYTES, + opslimit: int = TEST_KDF_OPSLIMIT, + declared_memlimit: int | None = None, + declared_opslimit: int | None = None, + parallelism: int = 1, + reserved_signature: bytes | None = None, + magic: bytes = MAGIC, + digest_override: bytes | None = None, + salt: bytes | None = None, + drop_final_tag: bool = False, +) -> bytes: + """Produce a full `.sientia` byte stream, mirroring the producer's `encrypt_archive`. + + The `declared_*` arguments write cost parameters into the header that differ from the ones + actually used, which is how a "512 MiB header" is built without allocating 512 MiB. + `drop_final_tag` pushes every chunk as `MESSAGE`, so the stream is framed as if more were + coming — the truncation the digest cannot see. + """ + salt = salt if salt is not None else os.urandom(SALT_SIZE) + key = derive_key(password, salt, opslimit, memlimit) + + state = sodium.crypto_secretstream_xchacha20poly1305_state() + stream_header = sodium.crypto_secretstream_xchacha20poly1305_init_push(state, key) + + chunks = [] + offset = 0 + total = len(plaintext) + while offset < total: + chunk = plaintext[offset : offset + CHUNK_SIZE] + offset += len(chunk) + is_last = offset >= total and not drop_final_tag + tag = ( + sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL + if is_last + else sodium.crypto_secretstream_xchacha20poly1305_TAG_MESSAGE + ) + chunks.append(sodium.crypto_secretstream_xchacha20poly1305_push(state, chunk, None, tag)) + ciphertext = b''.join(chunks) + + digest = digest_override if digest_override is not None else hashlib.sha256(ciphertext).digest() + header = HEADER_STRUCT.pack( + magic, + format_version, + aead_id, + kdf_id, + salt, + declared_memlimit if declared_memlimit is not None else memlimit, + declared_opslimit if declared_opslimit is not None else opslimit, + parallelism, + digest, + reserved_signature if reserved_signature is not None else bytes(RESERVED_SIGNATURE_SIZE), + ) + return header + stream_header + ciphertext + + +def make_bundle( + zip_spec: ZipSpec | None = None, + password: str = DEFAULT_PASSWORD, + **header_kwargs: Any, +) -> bytes: + """The common case: a well-formed tree, encrypted with the test cost parameters.""" + return encrypt_archive(build_zip(zip_spec), password, **header_kwargs) + + +def tamper_ciphertext(bundle: bytes, offset_from_stream: int = 8) -> bytes: + """Flip one byte inside the chunk stream **and** repair the header digest. + + Repairing the digest is what makes this a gate-3 test rather than a gate-2 one: the file is + self-consistent and only the AEAD can tell it was altered. + """ + position = HEADER_SIZE + 24 + offset_from_stream + altered = bytearray(bundle) + altered[position] ^= 0xFF + repaired = hashlib.sha256(bytes(altered[HEADER_SIZE + 24 :])).digest() + altered[44:76] = repaired + return bytes(altered) + + +def truncate_bundle(bundle: bytes, drop_bytes: int = 32) -> bytes: + """Cut the tail off a bundle, leaving the header's digest describing bytes that are gone.""" + return bundle[:-drop_bytes] + + +def build_envelope(password: str = DEFAULT_PASSWORD, key: bytes = DEFAULT_ENVELOPE_KEY) -> str: + """`base64(nonce || ciphertext || tag)` under AES-256-GCM, as the frontend will send it.""" + nonce = os.urandom(12) + sealed = AESGCM(key).encrypt(nonce, password.encode('utf-8'), None) + return base64.b64encode(nonce + sealed).decode('ascii') + + +def envelope_key_b64(key: bytes = DEFAULT_ENVELOPE_KEY) -> str: + """The key as `IMPORT_PASSWORD_KEY` carries it.""" + return base64.b64encode(key).decode('ascii') diff --git a/tests/laborious/utils/bundle/__init__.py b/tests/laborious/utils/bundle/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/laborious/utils/bundle/conftest.py b/tests/laborious/utils/bundle/conftest.py new file mode 100644 index 0000000..ffc41f6 --- /dev/null +++ b/tests/laborious/utils/bundle/conftest.py @@ -0,0 +1,34 @@ +"""Shared fixtures for the `.sientia` bundle tests.""" + +from pathlib import Path + +import pytest + +FIXTURE_DIR = Path(__file__).resolve().parents[3] / 'fixtures' / 'bundle' +GOLDEN_BUNDLE = FIXTURE_DIR / 'golden_model_v3.sientia' +GOLDEN_METADATA = FIXTURE_DIR / 'golden_model_v3.metadata.json' +# Recorded in tests/fixtures/bundle/README.md; the fixture protects nothing. +GOLDEN_PASSWORD = 'Import1Model!Fixture' + + +@pytest.fixture +def golden_bundle() -> Path: + """The `.sientia` file the real producer emitted.""" + assert GOLDEN_BUNDLE.is_file(), ( + f'golden fixture missing: {GOLDEN_BUNDLE} — regenerate it with ' + 'tests/fixtures/bundle/regenerate.py' + ) + return GOLDEN_BUNDLE + + +@pytest.fixture +def golden_password() -> str: + return GOLDEN_PASSWORD + + +@pytest.fixture +def work_dir(tmp_path: Path) -> Path: + """A per-test stand-in for `IMPORT_WORK_DIR/`.""" + root = tmp_path / 'workflow-id' + root.mkdir() + return root diff --git a/tests/laborious/utils/bundle/test_format.py b/tests/laborious/utils/bundle/test_format.py new file mode 100644 index 0000000..e8c0f6b --- /dev/null +++ b/tests/laborious/utils/bundle/test_format.py @@ -0,0 +1,227 @@ +"""The wire format: header parsing, ceilings, digest coverage and the fixture contract.""" + +import ast +import hashlib +from pathlib import Path + +import nacl.bindings as sodium +import pytest + +from laborious.utils.bundle import format as fmt +from tests.helpers import bundle_factory as factory + + +def _write(tmp_path: Path, payload: bytes, name: str = 'bundle.sientia') -> Path: + target = tmp_path / name + target.write_bytes(payload) + return target + + +# --------------------------------------------------------------- constants + + +def test_header_size_is_140(): + assert fmt.HEADER_SIZE == 140 + assert fmt.HEADER_STRUCT.format == '>8sBBB16sQQB32s64s' + + +def test_libsodium_literals_match_the_binding(): + """The literals in `format.py` keep the module import-free; these assertions keep them true.""" + assert fmt.STREAM_HEADER_SIZE == sodium.crypto_secretstream_xchacha20poly1305_HEADERBYTES + assert fmt.ABYTES == sodium.crypto_secretstream_xchacha20poly1305_ABYTES + assert fmt.KEY_SIZE == sodium.crypto_secretstream_xchacha20poly1305_KEYBYTES + + +def test_producer_constants_match(): + assert fmt.MAGIC == b'SIENTIA1' + assert fmt.CHUNK_SIZE == 65536 + assert fmt.MAX_KDF_MEMLIMIT == 268435456 + assert fmt.MAX_KDF_OPSLIMIT == 3 + assert fmt.SALT_SIZE == 16 + assert fmt.DIGEST_SIZE == 32 + assert fmt.RESERVED_SIGNATURE_SIZE == 64 + + +def test_prefix_size_is_65717(): + """The pre-upload prefix is computed from constants, never read from a file.""" + assert fmt.PREFIX_SIZE == 140 + 24 + (65536 + 17) + assert fmt.PREFIX_SIZE == 65717 + + +def test_digest_offset_skips_both_headers(): + assert fmt.DIGEST_OFFSET == 164 + + +# ------------------------------------------------------------------- parse + + +def test_parses_a_producer_emitted_header(golden_bundle): + header = fmt.read_header(golden_bundle) + + assert header.format_version == 1 + assert header.aead_id == fmt.AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM + assert header.kdf_id == fmt.KDF_ID_ARGON2ID + assert len(header.salt) == 16 + assert header.kdf_memlimit_bytes == 268435456 + assert header.kdf_opslimit == 3 + assert header.kdf_parallelism == 1 + assert len(header.ciphertext_digest) == 32 + assert header.reserved_signature == bytes(64) + + +def test_wrong_magic_rejected(tmp_path): + path = _write(tmp_path, factory.make_bundle(magic=b'NOTSIENT')) + + with pytest.raises(fmt.BundleFormatError, match='expected bundle marker'): + fmt.read_header(path) + + +def test_unknown_format_version_rejected(tmp_path): + path = _write(tmp_path, factory.make_bundle(format_version=2)) + + with pytest.raises(fmt.BundleFormatError, match='format version'): + fmt.read_header(path) + + +def test_unknown_aead_id_rejected(tmp_path): + path = _write(tmp_path, factory.make_bundle(aead_id=7)) + + with pytest.raises(fmt.BundleFormatError, match='encryption identifier'): + fmt.read_header(path) + + +def test_unknown_kdf_id_rejected(tmp_path): + path = _write(tmp_path, factory.make_bundle(kdf_id=9)) + + with pytest.raises(fmt.BundleFormatError, match='key derivation identifier'): + fmt.read_header(path) + + +def test_memlimit_above_the_ceiling_rejected(tmp_path): + path = _write(tmp_path, factory.make_bundle(declared_memlimit=512 * 1024 * 1024)) + + with pytest.raises(fmt.BundleFormatError, match='memory'): + fmt.read_header(path) + + +def test_opslimit_above_the_ceiling_rejected(tmp_path): + path = _write(tmp_path, factory.make_bundle(declared_opslimit=4)) + + with pytest.raises(fmt.BundleFormatError, match='passes'): + fmt.read_header(path) + + +def test_zero_cost_parameters_rejected(tmp_path): + path = _write(tmp_path, factory.make_bundle(declared_opslimit=0)) + + with pytest.raises(fmt.BundleFormatError, match='positive'): + fmt.read_header(path) + + +def test_reserved_parallelism_must_be_one(tmp_path): + path = _write(tmp_path, factory.make_bundle(parallelism=2)) + + with pytest.raises(fmt.BundleFormatError, match='parallelism'): + fmt.read_header(path) + + +def test_non_zero_reserved_signature_rejected(tmp_path): + path = _write(tmp_path, factory.make_bundle(reserved_signature=b'\x01' + bytes(63))) + + with pytest.raises(fmt.BundleFormatError, match='reserved signature'): + fmt.read_header(path) + + +def test_file_shorter_than_header_plus_stream_header_rejected(tmp_path): + path = _write(tmp_path, b'SIENTIA1' + bytes(100)) + + with pytest.raises(fmt.BundleFormatError, match='smaller than'): + fmt.read_header(path) + + +def test_parse_header_needs_the_whole_header(): + with pytest.raises(fmt.BundleFormatError, match='shorter than'): + fmt.parse_header(b'SIENTIA1') + + +# ------------------------------------------------------------------ digest + + +def test_digest_covers_the_chunk_stream_not_the_stream_header(golden_bundle): + """The producer hashes the chunks before prepending either header — offset 164, not 140.""" + header = fmt.read_header(golden_bundle) + payload = golden_bundle.read_bytes() + + assert hashlib.sha256(payload[164:]).digest() == header.ciphertext_digest + assert hashlib.sha256(payload[140:]).digest() != header.ciphertext_digest + assert fmt.compute_ciphertext_digest(golden_bundle) == header.ciphertext_digest + + +def test_verify_ciphertext_digest_accepts_the_fixture(golden_bundle): + fmt.verify_ciphertext_digest(golden_bundle, fmt.read_header(golden_bundle)) + + +def test_digest_mismatch_rejected(tmp_path): + path = _write(tmp_path, factory.make_bundle(digest_override=bytes(32))) + header = fmt.read_header(path) + + with pytest.raises(fmt.BundleFormatError, match='does not match the digest'): + fmt.verify_ciphertext_digest(path, header) + + +def test_truncated_file_fails_the_digest(tmp_path): + path = _write(tmp_path, factory.truncate_bundle(factory.make_bundle())) + header = fmt.read_header(path) + + with pytest.raises(fmt.BundleFormatError, match='does not match the digest'): + fmt.verify_ciphertext_digest(path, header) + + +def test_file_digest_is_the_whole_object(tmp_path): + payload = factory.make_bundle() + path = _write(tmp_path, payload) + + assert fmt.file_digest(path) == hashlib.sha256(payload).hexdigest() + + +# ------------------------------------------------------- factory conformance + + +def test_factory_header_matches_fixture_layout(golden_bundle, tmp_path): + """The test-only writer must be structurally identical to the producer's output.""" + produced = fmt.read_header(_write(tmp_path, factory.make_bundle())) + golden = fmt.read_header(golden_bundle) + + assert (produced.format_version, produced.aead_id, produced.kdf_id) == ( + golden.format_version, + golden.aead_id, + golden.kdf_id, + ) + assert produced.kdf_parallelism == golden.kdf_parallelism + assert produced.reserved_signature == golden.reserved_signature + assert len(produced.salt) == len(golden.salt) + assert len(produced.ciphertext_digest) == len(golden.ciphertext_digest) + assert factory.HEADER_STRUCT.format == fmt.HEADER_STRUCT.format + assert factory.PRODUCER_KDF_MEMLIMIT_BYTES == golden.kdf_memlimit_bytes + assert factory.PRODUCER_KDF_OPSLIMIT == golden.kdf_opslimit + + +# --------------------------------------------------------------- liftability + +FORBIDDEN_IMPORT_ROOTS = frozenset( + {'laborious', 'temporalio', 'mlflow', 'nacl', 'cryptography', 'sqlalchemy', 'pymongo'} +) + + +def test_format_module_imports_nothing_platform_specific(): + """`format.py` stays liftable into a shared package: stdlib only, no platform imports.""" + tree = ast.parse(Path(fmt.__file__).read_text(encoding='utf-8')) + roots = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + roots.update(alias.name.split('.')[0] for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module: + roots.add(node.module.split('.')[0]) + + assert roots.isdisjoint(FORBIDDEN_IMPORT_ROOTS), roots & FORBIDDEN_IMPORT_ROOTS + assert roots == {'__future__', 'hashlib', 'struct', 'dataclasses', 'pathlib', 'typing'} diff --git a/tests/laborious/utils/bundle/test_reader.py b/tests/laborious/utils/bundle/test_reader.py new file mode 100644 index 0000000..9430173 --- /dev/null +++ b/tests/laborious/utils/bundle/test_reader.py @@ -0,0 +1,460 @@ +"""The seven gates: what each one rejects, in what order, and what it leaves on disk.""" + +import json +import random +from pathlib import Path + +import nacl.bindings as sodium +import pytest + +from laborious.utils.bundle import format as fmt +from laborious.utils.bundle.reader import BundleGateError, BundleLimits, BundleReader +from laborious.utils.bundle.steps import ImportStep +from tests.helpers import bundle_factory as factory + +OBJECT_KEY = 'imported_models/factory_model_v2.sientia' + + +@pytest.fixture +def reader() -> BundleReader: + return BundleReader() + + +def _bundle_file(work_dir: Path, payload: bytes) -> Path: + target = work_dir / 'download.sientia' + target.write_bytes(payload) + return target + + +def _open(reader: BundleReader, work_dir: Path, payload: bytes, password: str | None = None): + return reader.open( + _bundle_file(work_dir, payload), + password if password is not None else factory.DEFAULT_PASSWORD, + work_dir / 'extract', + work_dir, + ) + + +def _extracted_files(root: Path) -> list[str]: + if not root.exists(): + return [] + return sorted(str(p.relative_to(root)) for p in root.rglob('*') if p.is_file()) + + +# ------------------------------------------------------------------- gate 1 + + +def test_gate_1_accepts_a_well_shaped_object(reader): + reader.check_object_limits(size=1024, key=OBJECT_KEY) + + +def test_gate_1_rejects_an_oversized_object(reader): + limited = BundleReader(BundleLimits(max_object_bytes=1000)) + + with pytest.raises(BundleGateError) as error: + limited.check_object_limits(size=1001, key=OBJECT_KEY) + + assert error.value.gate == 1 + assert error.value.step is ImportStep.DOWNLOAD + + +def test_gate_1_rejects_an_empty_object(reader): + with pytest.raises(BundleGateError, match='empty'): + reader.check_object_limits(size=0, key=OBJECT_KEY) + + +def test_gate_1_rejects_a_key_outside_the_import_prefix(reader): + with pytest.raises(BundleGateError, match='import prefix'): + reader.check_object_limits(size=10, key='other_place/model.sientia') + + +def test_gate_1_rejects_a_key_with_the_wrong_suffix(reader): + with pytest.raises(BundleGateError, match='.sientia'): + reader.check_object_limits(size=10, key='imported_models/model.zip') + + +# ------------------------------------------------------------- golden bundle + + +def test_opens_the_producer_fixture(golden_bundle, golden_password, work_dir, reader): + """The producer's own bytes, at the producer's own KDF cost.""" + opened = reader.open(golden_bundle, golden_password, work_dir / 'extract', work_dir) + + assert opened.metadata['model_name'] == 'fixture_model' + assert opened.metadata['experiment_name'] == 'fixture_experiment' + assert set(opened.metadata) == { + 'model_name', + 'experiment_name', + 'run_id', + 'model_version', + 'model_project', + 'export_timestamp', + } + assert opened.parameters['target_variable'] == 'SE' + assert opened.header is not None + assert opened.digest == opened.header.ciphertext_digest.hex() + assert _extracted_files(opened.extracted_dir) == [ + 'artifacts/data_model/MLmodel', + 'artifacts/data_model/conda.yaml', + 'artifacts/data_model/model.pkl', + 'artifacts/data_model/python_env.yaml', + 'artifacts/data_model/requirements.txt', + 'artifacts/model_card/model_card.json', + 'artifacts/prediction_model/MLmodel', + 'artifacts/prediction_model/conda.yaml', + 'artifacts/prediction_model/model.pkl', + 'artifacts/prediction_model/python_env.yaml', + 'artifacts/prediction_model/requirements.txt', + 'artifacts/raw_data.csv', + 'artifacts/test_data.csv', + 'artifacts/train_data.csv', + 'metadata.json', + ] + assert list(work_dir.glob('*.zip')) == [] + + +def test_fixture_metadata_matches_its_committed_plaintext_copy( + golden_bundle, golden_password, work_dir, reader +): + from tests.laborious.utils.bundle.conftest import GOLDEN_METADATA + + opened = reader.open(golden_bundle, golden_password, work_dir / 'extract', work_dir) + committed = json.loads(GOLDEN_METADATA.read_text(encoding='utf-8')) + + assert opened.metadata == committed['metadata'] + assert opened.parameters == committed['parameters'] + + +# ------------------------------------------------------------------- gate 2 + + +def test_gate_2_rejects_a_wrong_magic_before_deriving(work_dir, reader, monkeypatch): + payload = factory.make_bundle(magic=b'NOTSIENT') + monkeypatch.setattr( + sodium, 'crypto_pwhash_alg', _never_called('key derivation must not run at gate 2') + ) + + with pytest.raises(BundleGateError) as error: + _open(reader, work_dir, payload) + + assert error.value.gate == 2 + assert error.value.step is ImportStep.DOWNLOAD + + +def test_gate_2_rejects_kdf_parameters_above_the_ceiling_without_allocating( + work_dir, reader, monkeypatch +): + payload = factory.make_bundle(declared_memlimit=512 * 1024 * 1024) + monkeypatch.setattr( + sodium, 'crypto_pwhash_alg', _never_called('the ceiling must be enforced before deriving') + ) + + with pytest.raises(BundleGateError) as error: + _open(reader, work_dir, payload) + + assert error.value.gate == 2 + + +def test_gate_2_rejects_a_digest_mismatch(work_dir, reader): + with pytest.raises(BundleGateError) as error: + _open(reader, work_dir, factory.make_bundle(digest_override=bytes(32))) + + assert error.value.gate == 2 + + +def test_gate_2_rejects_a_truncated_transfer(work_dir, reader): + with pytest.raises(BundleGateError) as error: + _open(reader, work_dir, factory.truncate_bundle(factory.make_bundle())) + + assert error.value.gate == 2 + + +def _never_called(message: str): + def _raise(*_args, **_kwargs): + raise AssertionError(message) + + return _raise + + +# ------------------------------------------------------------------- gate 3 + +AMBIGUOUS_REASON = 'wrong password, or the file is damaged or was altered' + + +def test_gate_3_wrong_password(work_dir, reader): + with pytest.raises(BundleGateError) as error: + _open(reader, work_dir, factory.make_bundle(), password='Wrong1Password!Here') + + assert error.value.gate == 3 + assert error.value.step is ImportStep.DECRYPTION + assert AMBIGUOUS_REASON in str(error.value) + + +def test_gate_3_altered_ciphertext_reads_the_same_as_a_wrong_password(work_dir, reader): + tampered = factory.tamper_ciphertext(factory.make_bundle()) + + with pytest.raises(BundleGateError) as altered: + _open(reader, work_dir, tampered) + with pytest.raises(BundleGateError) as wrong_password: + _open(reader, work_dir, factory.make_bundle(), password='Wrong1Password!Here') + + assert str(altered.value) == str(wrong_password.value) + assert altered.value.gate == wrong_password.value.gate == 3 + + +def test_gate_3_stream_without_a_final_tag_is_rejected(work_dir, reader): + with pytest.raises(BundleGateError) as error: + _open(reader, work_dir, factory.make_bundle(drop_final_tag=True)) + + assert error.value.gate == 3 + assert AMBIGUOUS_REASON in str(error.value) + + +def test_gate_3_leaves_no_decrypted_zip_behind(work_dir, reader): + with pytest.raises(BundleGateError): + _open(reader, work_dir, factory.make_bundle(), password='Wrong1Password!Here') + + assert list(work_dir.glob('*.zip')) == [] + + +def test_gate_3_normalises_the_password_to_nfc(work_dir, reader): + composed = 'Café1Model!Test' + decomposed = 'Café1Model!Test' + payload = factory.make_bundle(password=composed) + + opened = _open(reader, work_dir, payload, password=decomposed) + + assert opened.metadata['model_name'] == 'factory_model' + + +def test_gate_3_does_not_trim_the_password(work_dir, reader): + payload = factory.make_bundle(password='Trailing1Space!Test ') + + with pytest.raises(BundleGateError): + _open(reader, work_dir, payload, password='Trailing1Space!Test') + + +def test_multi_chunk_plaintext_round_trips(work_dir, reader): + """More than 64 KiB of plaintext, so `MESSAGE` and `FINAL` framing are both exercised.""" + spec = factory.ZipSpec() + # Incompressible on purpose: a compressible filler would deflate below one chunk and the test + # would silently stop covering multi-chunk framing. + spec.files['artifacts/prediction_model/model.pkl'] = random.Random( # noqa: S311 - filler + 20260825 + ).randbytes(200 * 1024) + payload = factory.make_bundle(spec) + assert len(payload) > fmt.CHUNK_SIZE + fmt.ABYTES + fmt.DIGEST_OFFSET + + opened = _open(reader, work_dir, payload) + + extracted = opened.extracted_dir / 'artifacts' / 'prediction_model' / 'model.pkl' + assert extracted.read_bytes() == spec.files['artifacts/prediction_model/model.pkl'] + + +# ------------------------------------------------------------------- gate 4 + + +@pytest.mark.parametrize( + ('spec_kwargs', 'match'), + [ + ({'bomb_bytes': 40 * 1024 * 1024}, 'compression ratio'), + ({'traversal_entry': 'artifacts/../../etc/passwd'}, 'escaping'), + ({'absolute_entry': '/etc/passwd'}, 'absolute'), + ({'symlinks': {'artifacts/link': '/etc/passwd'}}, 'symbolic link'), + ], +) +def test_gate_4_rejects_hostile_central_directories(work_dir, reader, spec_kwargs, match): + payload = factory.make_bundle(factory.ZipSpec(**spec_kwargs)) + + with pytest.raises(BundleGateError, match=match) as error: + _open(reader, work_dir, payload) + + assert error.value.gate == 4 + assert error.value.step is ImportStep.ARCHIVE_INSPECTION + assert _extracted_files(work_dir / 'extract') == [] + assert not (work_dir / 'extract').exists() + + +def test_gate_4_rejects_too_many_entries(work_dir): + limited = BundleReader(BundleLimits(max_entries=3)) + + with pytest.raises(BundleGateError, match='entries') as error: + _open(limited, work_dir, factory.make_bundle()) + + assert error.value.gate == 4 + + +def test_gate_4_rejects_an_oversized_uncompressed_total(work_dir): + limited = BundleReader(BundleLimits(max_uncompressed_bytes=64)) + + with pytest.raises(BundleGateError, match='uncompressed') as error: + _open(limited, work_dir, factory.make_bundle()) + + assert error.value.gate == 4 + + +def test_gate_4_rejects_content_that_is_not_an_archive(work_dir, reader): + payload = factory.encrypt_archive(b'not a zip at all') + + with pytest.raises(BundleGateError, match='readable archive') as error: + _open(reader, work_dir, payload) + + assert error.value.gate == 4 + + +# ------------------------------------------------------------------- gate 5 + + +def test_gate_5_refuses_a_pre_existing_extraction_root(work_dir, reader): + (work_dir / 'extract').mkdir() + + with pytest.raises(BundleGateError, match='already exists') as error: + _open(reader, work_dir, factory.make_bundle()) + + assert error.value.gate == 5 + assert error.value.step is ImportStep.EXTRACTION + + +def test_gate_5_extraction_writes_plain_files(work_dir, reader): + opened = _open(reader, work_dir, factory.make_bundle()) + + for path in opened.extracted_dir.rglob('*'): + if path.is_file(): + assert path.stat().st_mode & 0o111 == 0, path + + +# ------------------------------------------------------------------- gate 6 + + +@pytest.mark.parametrize( + ('spec', 'match'), + [ + (factory.ZipSpec(include_metadata=False), 'metadata document'), + ( + factory.ZipSpec( + document=factory.build_metadata_document(drop_top_level_key='parameters') + ), + 'parameters and metadata blocks', + ), + ( + factory.ZipSpec( + document=factory.build_metadata_document(drop_metadata_field='experiment_name') + ), + 'experiment_name', + ), + ( + factory.ZipSpec( + document=factory.build_metadata_document(drop_parameter='target_variable') + ), + 'target_variable', + ), + ( + factory.ZipSpec( + document=factory.build_metadata_document(metadata={'experiment_name': 'x' * 51}) + ), + 'experiment name', + ), + ( + factory.ZipSpec( + document=factory.build_metadata_document(metadata={'experiment_name': 'ab'}) + ), + 'experiment name', + ), + ], +) +def test_gate_6_rejects_a_malformed_metadata_document(work_dir, reader, spec, match): + with pytest.raises(BundleGateError, match=match) as error: + _open(reader, work_dir, factory.make_bundle(spec)) + + assert error.value.gate == 6 + assert error.value.step is ImportStep.STRUCTURE_VALIDATION + assert not (work_dir / 'extract').exists() + + +def test_gate_6_rejects_unreadable_metadata_json(work_dir, reader): + spec = factory.ZipSpec(include_metadata=False) + spec.extra_files['metadata.json'] = b'{not json' + + with pytest.raises(BundleGateError, match='readable JSON') as error: + _open(reader, work_dir, factory.make_bundle(spec)) + + assert error.value.gate == 6 + + +def test_gate_6_rejects_a_missing_artifacts_directory(work_dir, reader): + spec = factory.ZipSpec(files={}) + + with pytest.raises(BundleGateError, match='artifacts directory') as error: + _open(reader, work_dir, factory.make_bundle(spec)) + + assert error.value.gate == 6 + + +@pytest.mark.parametrize( + ('dropped', 'match'), + [ + ('artifacts/data_model/MLmodel', 'data_model artifacts are incomplete'), + ('artifacts/prediction_model/model.pkl', 'prediction_model artifacts are incomplete'), + ], +) +def test_gate_6_rejects_an_incomplete_model_directory(work_dir, reader, dropped, match): + spec = factory.ZipSpec() + spec.files.pop(dropped) + + with pytest.raises(BundleGateError, match=match) as error: + _open(reader, work_dir, factory.make_bundle(spec)) + + assert error.value.gate == 6 + + +def test_gate_6_rejects_a_bundle_without_data_model(work_dir, reader): + spec = factory.ZipSpec() + for name in list(spec.files): + if name.startswith('artifacts/data_model/'): + spec.files.pop(name) + + with pytest.raises(BundleGateError, match='no data_model artifacts') as error: + _open(reader, work_dir, factory.make_bundle(spec)) + + assert error.value.gate == 6 + + +# ------------------------------------------------------------------- gate 7 + + +@pytest.mark.parametrize('name', ['artifacts/run.sh', 'artifacts/tool.bin', 'artifacts/mod.so']) +def test_gate_7_rejects_content_a_model_export_should_not_carry(work_dir, reader, name): + spec = factory.ZipSpec() + spec.extra_files[name] = b'#!/bin/sh\n' + + with pytest.raises(BundleGateError, match='should not carry') as error: + _open(reader, work_dir, factory.make_bundle(spec)) + + assert error.value.gate == 7 + assert error.value.step is ImportStep.CONTENT_POLICY + assert not (work_dir / 'extract').exists() + + +# ------------------------------------------------------------------ ordering + + +def test_the_earlier_gate_wins(work_dir, reader): + """A bundle that would fail gate 4 and gate 6 must report gate 4, having extracted nothing.""" + spec = factory.ZipSpec( + traversal_entry='artifacts/../escape.csv', + document=factory.build_metadata_document(drop_parameter='target_variable'), + ) + + with pytest.raises(BundleGateError) as error: + _open(reader, work_dir, factory.make_bundle(spec)) + + assert error.value.gate == 4 + assert _extracted_files(work_dir / 'extract') == [] + + +def test_header_is_checked_before_the_password(work_dir, reader): + """A bundle with a bad header and a wrong password reports the header, not the password.""" + with pytest.raises(BundleGateError) as error: + _open(reader, work_dir, factory.make_bundle(magic=b'NOTSIENT'), password='Nope1!Nope1!') + + assert error.value.gate == 2 diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index b94bb8d..8dd82f5 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -1,11 +1,42 @@ -from os import environ +import tempfile +from os import environ, path + +import pytest from laborious.utils.connectors_config import ( + build_import_config, + build_import_status_config, build_minio_config, build_mlflow_config, build_opc_config, ) +IMPORT_ENV_VARS = ( + 'IMPORT_PASSWORD_KEY', + 'IMPORT_BUNDLE_BUCKET', + 'IMPORT_BUNDLE_PREFIX', + 'IMPORT_WORK_DIR', + 'IMPORT_MAX_OBJECT_BYTES', + 'IMPORT_MAX_ARCHIVE_ENTRIES', + 'IMPORT_MAX_UNCOMPRESSED_BYTES', + 'IMPORT_MAX_COMPRESSION_RATIO', + 'IMPORT_MODELS_COLLECTION', + 'IMPORT_BUNDLE_RETENTION_DAYS', + 'IMPORT_STATUS_DB_HOST', + 'IMPORT_STATUS_DB_PORT', + 'IMPORT_STATUS_DB_NAME', + 'IMPORT_STATUS_DB_USER', + 'IMPORT_STATUS_DB_PASSWORD', +) + + +@pytest.fixture +def clean_import_env(monkeypatch): + """Import configuration is read from a known-empty environment.""" + for name in IMPORT_ENV_VARS: + monkeypatch.delenv(name, raising=False) + return monkeypatch + def test_build_mlflow_config_with_env_vars(): # Arrange @@ -120,3 +151,83 @@ def test_build_minio_config_with_defaults(): 'retention_hours': 24, 'secure': False, } + + +def test_build_import_config_with_defaults(clean_import_env): + assert build_import_config() == { + 'password_key': None, + 'bucket': 'sientia', + 'prefix': 'imported_models/', + 'work_dir': path.join(tempfile.gettempdir(), 'sientia-import'), + 'max_object_bytes': 1024 * 1024 * 1024, + 'max_entries': 5000, + 'max_uncompressed_bytes': 4 * 1024 * 1024 * 1024, + 'max_compression_ratio': 200.0, + 'models_collection': 'models', + 'retention_days': 7, + } + + +def test_build_import_config_reads_the_environment(clean_import_env): + clean_import_env.setenv('IMPORT_PASSWORD_KEY', 'a2V5') + clean_import_env.setenv('IMPORT_BUNDLE_BUCKET', 'other-bucket') + clean_import_env.setenv('IMPORT_BUNDLE_PREFIX', 'incoming/models') + clean_import_env.setenv('IMPORT_WORK_DIR', '/var/lib/import') + clean_import_env.setenv('IMPORT_MAX_OBJECT_BYTES', '2048') + clean_import_env.setenv('IMPORT_MAX_ARCHIVE_ENTRIES', '12') + 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_BUNDLE_RETENTION_DAYS', '14') + + config = build_import_config() + + assert config['password_key'] == 'a2V5' + assert config['bucket'] == 'other-bucket' + # A prefix without its trailing slash would make `imported_models_backup/x` pass the gate-1 + # prefix check, so the separator is added rather than trusted. + assert config['prefix'] == 'incoming/models/' + assert config['work_dir'] == '/var/lib/import' + assert config['max_object_bytes'] == 2048 + assert config['max_entries'] == 12 + assert config['max_uncompressed_bytes'] == 4096 + assert config['max_compression_ratio'] == 12.5 + assert config['models_collection'] == 'models_v2' + assert config['retention_days'] == 14 + + +def test_build_import_config_does_not_validate_the_password_key(clean_import_env): + """A missing or malformed key must fail the import at decryption, never at worker start.""" + clean_import_env.setenv('IMPORT_PASSWORD_KEY', 'not base64 at all') + + assert build_import_config()['password_key'] == 'not base64 at all' + + +def test_build_import_status_config_with_defaults(clean_import_env): + assert build_import_status_config() == { + 'host': 'localhost', + 'port': 5432, + 'dbname': 'sientia-core-mlops-bff', + 'user': 'mlops_bff_user', + 'password': '', + 'schema': 'public', + } + + +def test_build_import_status_config_reads_the_environment(clean_import_env): + clean_import_env.setenv('IMPORT_STATUS_DB_HOST', 'bff-postgres') + clean_import_env.setenv('IMPORT_STATUS_DB_PORT', '6543') + clean_import_env.setenv('IMPORT_STATUS_DB_NAME', 'other-db') + clean_import_env.setenv('IMPORT_STATUS_DB_USER', 'other-user') + clean_import_env.setenv('IMPORT_STATUS_DB_PASSWORD', 'secret') + + config = build_import_status_config() + + assert config == { + 'host': 'bff-postgres', + 'port': 6543, + 'dbname': 'other-db', + 'user': 'other-user', + 'password': 'secret', + 'schema': 'public', + } diff --git a/tests/laborious/utils/test_import_password.py b/tests/laborious/utils/test_import_password.py new file mode 100644 index 0000000..5d01876 --- /dev/null +++ b/tests/laborious/utils/test_import_password.py @@ -0,0 +1,99 @@ +"""The password envelope: AES-256-GCM in, plaintext password out, nothing in between.""" + +import base64 + +import pytest + +from laborious.utils.import_password import ( + EnvelopeError, + decrypt_password_envelope, + load_envelope_key, +) +from tests.helpers import bundle_factory as factory + +KEY = factory.DEFAULT_ENVELOPE_KEY +OTHER_KEY = bytes(range(1, 33)) + + +def test_recovers_the_password_used_at_export(): + envelope = factory.build_envelope('Correct1Horse!Battery', KEY) + + assert decrypt_password_envelope(envelope, KEY) == 'Correct1Horse!Battery' + + +def test_wrong_key_fails_authentication(): + envelope = factory.build_envelope('Correct1Horse!Battery', KEY) + + with pytest.raises(EnvelopeError, match='authenticated'): + decrypt_password_envelope(envelope, OTHER_KEY) + + +def test_tampered_envelope_fails_authentication(): + payload = bytearray(base64.b64decode(factory.build_envelope('Correct1Horse!Battery', KEY))) + payload[-1] ^= 0xFF + + with pytest.raises(EnvelopeError, match='authenticated'): + decrypt_password_envelope(base64.b64encode(bytes(payload)).decode(), KEY) + + +def test_malformed_base64_rejected(): + with pytest.raises(EnvelopeError, match='base64'): + decrypt_password_envelope('this is not base64!!', KEY) + + +def test_envelope_too_short_rejected(): + with pytest.raises(EnvelopeError, match='too short'): + decrypt_password_envelope(base64.b64encode(bytes(20)).decode(), KEY) + + +def test_empty_envelope_rejected(): + with pytest.raises(EnvelopeError, match='no password envelope'): + decrypt_password_envelope('', KEY) + + +def test_password_is_not_trimmed(): + """A trailing space is part of the secret; the KDF must see it.""" + envelope = factory.build_envelope('Trailing1Space!Test ', KEY) + + assert decrypt_password_envelope(envelope, KEY) == 'Trailing1Space!Test ' + + +def test_password_keeps_its_unicode_form(): + """Normalisation happens at key derivation, not here — the envelope round-trips verbatim.""" + decomposed = 'Café1Model!Test' + envelope = factory.build_envelope(decomposed, KEY) + + assert decrypt_password_envelope(envelope, KEY) == decomposed + + +def test_loads_a_valid_key(): + assert load_envelope_key(factory.envelope_key_b64(KEY)) == KEY + + +@pytest.mark.parametrize( + ('raw_key', 'match'), + [ + (None, 'not configured'), + ('', 'not configured'), + ('not base64 at all', 'valid base64'), + (base64.b64encode(bytes(16)).decode(), '32 bytes'), + (base64.b64encode(bytes(64)).decode(), '32 bytes'), + ], +) +def test_rejects_an_unusable_key(raw_key, match): + """There is no fallback: an unusable key never means the envelope is plaintext.""" + with pytest.raises(EnvelopeError, match=match): + load_envelope_key(raw_key) + + +def test_no_secret_appears_in_the_error_text(): + password = 'Correct1Horse!Battery' + envelope = factory.build_envelope(password, KEY) + + with pytest.raises(EnvelopeError) as error: + decrypt_password_envelope(envelope, OTHER_KEY) + + message = str(error.value) + assert password not in message + assert envelope not in message + assert factory.envelope_key_b64(OTHER_KEY) not in message