QTZPOC-20

QTZPOC-20: Implement .sientia model bundle import workflow

Introduces a new Temporal workflow and activities to import .sientia model bundles. This enables:
- Secure decryption and validation of bundles through a seven-gate process.
- Provisioning imported models in MLflow and updating the MongoDB model listing.
- Comprehensive status tracking in the public.experiment_run table with new IMPORT_SUCCESS/IMPORT_ERROR states.
- Dedicated metrics, robust error handling, and configurable settings.
- Includes updated e2e database fixtures and extensive unit/integration tests for bundle format, reader, and overall workflow.
This commit is contained in:
2026-08-25 10:40:46 -03:00
parent 291d16efa5
commit 0e445de386
27 changed files with 4471 additions and 11 deletions

View File

@@ -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);

View File

@@ -7,16 +7,22 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger from sientia_do.observability.logger import Logger
from sientia_do.repository.minio_repository import MinioRepository 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.api import API
from laborious.activities.gates import Gates from laborious.activities.gates import Gates
from laborious.activities.mlflow import MLFlow from laborious.activities.mlflow import MLFlow
from laborious.activities.model_import import ModelImport
from laborious.activities.model_metrics import ModelMetrics from laborious.activities.model_metrics import ModelMetrics
from laborious.activities.opc import OPC from laborious.activities.opc import OPC
from laborious.activities.storage import Storage 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. 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], pi_web_api_config: dict[str, Any],
logger: Logger, logger: Logger,
notification_handler: NotificationHandler, 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. 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 Required keys: base_url, auth_type, auth_token
logger: Logger instance for observability and debugging logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring 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: Raises:
Exception: If any parent class initialization fails Exception: If any parent class initialization fails
@@ -147,6 +163,32 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
metrics_controller=metrics_controller, 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): async def shutdown(self):
""" """
Gracefully shutdown all activities and clean up resources. Gracefully shutdown all activities and clean up resources.
@@ -167,3 +209,4 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
await OPC.aclose(self) await OPC.aclose(self)
ModelMetrics.close(self) ModelMetrics.close(self)
API.close(self) API.close(self)
ModelImport.close(self)

File diff suppressed because it is too large Load Diff

View File

@@ -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

View File

@@ -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 metrics ==================
PI_WEB_API_LABELS = [*CORE_LABELS, 'tag_name'] PI_WEB_API_LABELS = [*CORE_LABELS, 'tag_name']

View File

@@ -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',
]

View File

@@ -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()

View File

@@ -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)

View File

@@ -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

View File

@@ -1,5 +1,6 @@
import json import json
from os import getenv import tempfile
from os import getenv, path
from typing import Any from typing import Any
@@ -90,3 +91,72 @@ def build_minio_config() -> dict[str, Any]:
'retention_hours': int(getenv('MINIO_RETENTION_HOURS', '24')), 'retention_hours': int(getenv('MINIO_RETENTION_HOURS', '24')),
'secure': getenv('MINIO_SECURE', 'false') == 'true', '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: <tempdir>/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',
}

View File

@@ -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

View File

@@ -10,6 +10,8 @@ The worker supports multiple runtime-scoped task queues (via ``sientia_do.tempor
- minimal_retrain-{runtime}-queue: Model retraining workflows - minimal_retrain-{runtime}-queue: Model retraining workflows
- drift-{runtime}-queue: Drift detection workflows - drift-{runtime}-queue: Drift detection workflows
- simple_metrics-{runtime}-queue: Simple metrics 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 ``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``). 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_METRICS_PORT: Prometheus metrics server port (default: 9090)
- HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091) - HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091)
- PROJECT_NAME: Project name for notifications (default: laborious) - 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 from temporalio import client, workflow
@@ -53,11 +58,14 @@ with workflow.unsafe.imports_passed_through():
from laborious import metrics from laborious import metrics
from laborious.activities.activities import Activities from laborious.activities.activities import Activities
from laborious.utils.connectors_config import ( from laborious.utils.connectors_config import (
build_import_config,
build_import_status_config,
build_minio_config, build_minio_config,
build_mlflow_config, build_mlflow_config,
build_opc_config, build_opc_config,
) )
from laborious.workflows.drift import Drift from laborious.workflows.drift import Drift
from laborious.workflows.import_model import ImportModel
from laborious.workflows.minimal_retrain import MinimalRetrain from laborious.workflows.minimal_retrain import MinimalRetrain
from laborious.workflows.predictions_batch import PredictionsBatch from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.workflows.simple_metrics import SimpleMetrics from laborious.workflows.simple_metrics import SimpleMetrics
@@ -137,6 +145,9 @@ async def main():
pi_web_api_config=build_api_config(), pi_web_api_config=build_api_config(),
logger=logger, logger=logger,
notification_handler=notification_handler, 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) 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 = [] handlers = []
for w in workers: for w in workers:
handlers.append(w.run()) handlers.append(w.run())

View File

@@ -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 <fixed values> WHERE id` is
idempotent. Every provisioning call keeps `maximum_attempts=1`, because a repeated registry call or
Mongo insert can write twice.
- **One `finally`, two actions**: cleanup, then the single terminal write that classifies the
outcome once. No `except` block writes a verdict of its own.
- **The workflow holds ids, never contents.** The decrypted bundle never crosses an activity
boundary, and the password reaches this module only as an envelope it cannot open.
"""
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}

View File

@@ -4,11 +4,17 @@ sqlalchemy
asyncua==1.0.6 asyncua==1.0.6
redis redis
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.12.1 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 prometheus-client
botocore # fsspec/s3fs/gcsfs pinned as a set: nannyml (via the mlops library) pulls both s3fs and
boto3 # gcsfs, and each pins fsspec exactly, so a free `s3fs` makes the resolver fail against
s3fs # 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 pyarrow
kaleido kaleido
hyperopt hyperopt

View File

@@ -4,11 +4,17 @@ sqlalchemy
asyncua==1.0.6 asyncua==1.0.6
redis redis
sientia_do>=1.12.1 sientia_do>=1.12.1
sientia>0.40.0 sientia>=0.41.0
prometheus-client prometheus-client
botocore # fsspec/s3fs/gcsfs pinned as a set: nannyml (via sientia) pulls both s3fs and gcsfs, and
boto3 # each pins fsspec exactly, so a free `s3fs` makes the resolver fail against evidently's
s3fs # 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 pyarrow
kaleido kaleido
hyperopt hyperopt
@@ -18,3 +24,10 @@ scipy<1.14.0
scikit-learn==1.4.2 scikit-learn==1.4.2
pycaret==3.3.2 pycaret==3.3.2
joblib 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

86
tests/fixtures/bundle/README.md vendored Normal file
View File

@@ -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.

View File

@@ -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"
}
}

Binary file not shown.

185
tests/fixtures/bundle/regenerate.py vendored Normal file
View File

@@ -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())

View File

View File

@@ -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')

View File

View File

@@ -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/<workflow_id>`."""
root = tmp_path / 'workflow-id'
root.mkdir()
return root

View File

@@ -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'}

View File

@@ -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

View File

@@ -1,11 +1,42 @@
from os import environ import tempfile
from os import environ, path
import pytest
from laborious.utils.connectors_config import ( from laborious.utils.connectors_config import (
build_import_config,
build_import_status_config,
build_minio_config, build_minio_config,
build_mlflow_config, build_mlflow_config,
build_opc_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(): def test_build_mlflow_config_with_env_vars():
# Arrange # Arrange
@@ -120,3 +151,83 @@ def test_build_minio_config_with_defaults():
'retention_hours': 24, 'retention_hours': 24,
'secure': False, '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',
}

View File

@@ -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