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:
67
e2e/fixtures/experiment_run.sql
Normal file
67
e2e/fixtures/experiment_run.sql
Normal 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);
|
||||||
@@ -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)
|
||||||
|
|||||||
1044
laborious/activities/model_import.py
Normal file
1044
laborious/activities/model_import.py
Normal file
File diff suppressed because it is too large
Load Diff
348
laborious/activities/model_import_errors.py
Normal file
348
laborious/activities/model_import_errors.py
Normal 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
|
||||||
@@ -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']
|
||||||
|
|||||||
27
laborious/utils/bundle/__init__.py
Normal file
27
laborious/utils/bundle/__init__.py
Normal 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',
|
||||||
|
]
|
||||||
216
laborious/utils/bundle/format.py
Normal file
216
laborious/utils/bundle/format.py
Normal 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()
|
||||||
473
laborious/utils/bundle/reader.py
Normal file
473
laborious/utils/bundle/reader.py
Normal 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)
|
||||||
55
laborious/utils/bundle/steps.py
Normal file
55
laborious/utils/bundle/steps.py
Normal 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
|
||||||
@@ -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',
|
||||||
|
}
|
||||||
|
|||||||
98
laborious/utils/import_password.py
Normal file
98
laborious/utils/import_password.py
Normal 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
|
||||||
@@ -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())
|
||||||
|
|||||||
416
laborious/workflows/import_model.py
Normal file
416
laborious/workflows/import_model.py
Normal 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}
|
||||||
@@ -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
|
||||||
|
|||||||
@@ -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
86
tests/fixtures/bundle/README.md
vendored
Normal 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.
|
||||||
24
tests/fixtures/bundle/golden_model_v3.metadata.json
vendored
Normal file
24
tests/fixtures/bundle/golden_model_v3.metadata.json
vendored
Normal 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"
|
||||||
|
}
|
||||||
|
}
|
||||||
BIN
tests/fixtures/bundle/golden_model_v3.sientia
vendored
Normal file
BIN
tests/fixtures/bundle/golden_model_v3.sientia
vendored
Normal file
Binary file not shown.
185
tests/fixtures/bundle/regenerate.py
vendored
Normal file
185
tests/fixtures/bundle/regenerate.py
vendored
Normal 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())
|
||||||
0
tests/helpers/__init__.py
Normal file
0
tests/helpers/__init__.py
Normal file
272
tests/helpers/bundle_factory.py
Normal file
272
tests/helpers/bundle_factory.py
Normal 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')
|
||||||
0
tests/laborious/utils/bundle/__init__.py
Normal file
0
tests/laborious/utils/bundle/__init__.py
Normal file
34
tests/laborious/utils/bundle/conftest.py
Normal file
34
tests/laborious/utils/bundle/conftest.py
Normal 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
|
||||||
227
tests/laborious/utils/bundle/test_format.py
Normal file
227
tests/laborious/utils/bundle/test_format.py
Normal 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'}
|
||||||
460
tests/laborious/utils/bundle/test_reader.py
Normal file
460
tests/laborious/utils/bundle/test_reader.py
Normal 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
|
||||||
@@ -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',
|
||||||
|
}
|
||||||
|
|||||||
99
tests/laborious/utils/test_import_password.py
Normal file
99
tests/laborious/utils/test_import_password.py
Normal 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
|
||||||
Reference in New Issue
Block a user