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:
@@ -7,16 +7,22 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.repository.minio_repository import MinioRepository
|
||||
from sientia_do.repository.mongodb_repository import MongoDBRepository
|
||||
|
||||
from laborious.activities.api import API
|
||||
from laborious.activities.gates import Gates
|
||||
from laborious.activities.mlflow import MLFlow
|
||||
from laborious.activities.model_import import ModelImport
|
||||
from laborious.activities.model_metrics import ModelMetrics
|
||||
from laborious.activities.opc import OPC
|
||||
from laborious.activities.storage import Storage
|
||||
from laborious.utils.connectors_config import (
|
||||
build_import_config,
|
||||
build_import_status_config,
|
||||
)
|
||||
|
||||
|
||||
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
|
||||
"""
|
||||
Main activities orchestrator for the Laborious system.
|
||||
|
||||
@@ -50,6 +56,9 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
pi_web_api_config: dict[str, Any],
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
mongo_config: dict[str, Any] | None = None,
|
||||
import_config: dict[str, Any] | None = None,
|
||||
import_status_config: dict[str, Any] | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize the Activities orchestrator with all required configurations.
|
||||
@@ -68,6 +77,13 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
Required keys: base_url, auth_type, auth_token
|
||||
logger: Logger instance for observability and debugging
|
||||
notification_handler: Notification handler for alerts and monitoring
|
||||
mongo_config: MongoDB configuration for the model listing the import writes. When
|
||||
absent, the import activities that need it refuse rather than guess — the
|
||||
notification handler keeps its own client either way.
|
||||
import_config: `build_import_config()`; read from the environment when absent
|
||||
import_status_config: `build_import_status_config()`; read from the environment when
|
||||
absent. This is the BFF database holding the import log, a second Postgres
|
||||
connection that is never the one Storage holds.
|
||||
|
||||
Raises:
|
||||
Exception: If any parent class initialization fails
|
||||
@@ -147,6 +163,32 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
mongo_repository = (
|
||||
MongoDBRepository(
|
||||
connection_string=mongo_config['connection_string'],
|
||||
database_name=mongo_config['database_name'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
if mongo_config
|
||||
else None
|
||||
)
|
||||
|
||||
ModelImport.__init__(
|
||||
self,
|
||||
status_config=import_status_config or build_import_status_config(),
|
||||
import_config=import_config or build_import_config(),
|
||||
minio_repository=minio_repository,
|
||||
# The MLflow repository is shared rather than rebuilt: setting the tracking URI is a
|
||||
# process-wide side effect, so two instances would be two chances to disagree.
|
||||
mlflow_repository=self.model_monitoring_repository,
|
||||
mongo_repository=mongo_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
async def shutdown(self):
|
||||
"""
|
||||
Gracefully shutdown all activities and clean up resources.
|
||||
@@ -167,3 +209,4 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
await OPC.aclose(self)
|
||||
ModelMetrics.close(self)
|
||||
API.close(self)
|
||||
ModelImport.close(self)
|
||||
|
||||
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_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
|
||||
from os import getenv
|
||||
import tempfile
|
||||
from os import getenv, path
|
||||
from typing import Any
|
||||
|
||||
|
||||
@@ -90,3 +91,72 @@ def build_minio_config() -> dict[str, Any]:
|
||||
'retention_hours': int(getenv('MINIO_RETENTION_HOURS', '24')),
|
||||
'secure': getenv('MINIO_SECURE', 'false') == 'true',
|
||||
}
|
||||
|
||||
|
||||
def build_import_config() -> dict[str, Any]:
|
||||
"""
|
||||
Build the `.sientia` model-import configuration from environment variables.
|
||||
|
||||
The password key is read as the raw base64 string and **not** validated here: the worker must
|
||||
start whether or not `IMPORT_MODEL_ENABLED` is set, and a missing or malformed key is an import
|
||||
failure at the decryption step (`laborious.utils.import_password.load_envelope_key`), not a
|
||||
boot failure.
|
||||
|
||||
Environment Variables:
|
||||
IMPORT_PASSWORD_KEY: base64 of the 32-byte AES-256-GCM key that wraps bundle passwords
|
||||
IMPORT_BUNDLE_BUCKET: bucket holding the uploaded bundles (default: sientia)
|
||||
IMPORT_BUNDLE_PREFIX: key prefix the uploads must sit under (default: imported_models/)
|
||||
IMPORT_WORK_DIR: worker scratch directory (default: <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
|
||||
- drift-{runtime}-queue: Drift detection workflows
|
||||
- simple_metrics-{runtime}-queue: Simple metrics workflows
|
||||
- import_model-{runtime}-queue: `.sientia` model import — registered **only** when
|
||||
``IMPORT_MODEL_ENABLED`` is truthy, so the containment boundary stays a deployment decision
|
||||
|
||||
``RUNTIME`` must be set; it is passed to every ``prepare_worker`` call. Schedulers must use the
|
||||
same queue names (breaking change vs legacy ``drift-queue`` / ``simple_metrics-queue``).
|
||||
@@ -30,6 +32,9 @@ Environment Variables:
|
||||
- HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090)
|
||||
- HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091)
|
||||
- PROJECT_NAME: Project name for notifications (default: laborious)
|
||||
- IMPORT_MODEL_ENABLED: register the fifth (import) worker; unset by default
|
||||
- IMPORT_PASSWORD_KEY / IMPORT_BUNDLE_* / IMPORT_STATUS_DB_*: import configuration, see
|
||||
``laborious.utils.connectors_config``
|
||||
"""
|
||||
|
||||
from temporalio import client, workflow
|
||||
@@ -53,11 +58,14 @@ with workflow.unsafe.imports_passed_through():
|
||||
from laborious import metrics
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.utils.connectors_config import (
|
||||
build_import_config,
|
||||
build_import_status_config,
|
||||
build_minio_config,
|
||||
build_mlflow_config,
|
||||
build_opc_config,
|
||||
)
|
||||
from laborious.workflows.drift import Drift
|
||||
from laborious.workflows.import_model import ImportModel
|
||||
from laborious.workflows.minimal_retrain import MinimalRetrain
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
from laborious.workflows.simple_metrics import SimpleMetrics
|
||||
@@ -137,6 +145,9 @@ async def main():
|
||||
pi_web_api_config=build_api_config(),
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
mongo_config=mongo_config,
|
||||
import_config=build_import_config(),
|
||||
import_status_config=build_import_status_config(),
|
||||
)
|
||||
|
||||
logger.custom_info('Initializing OPC...', metadata_runtime)
|
||||
@@ -235,6 +246,42 @@ async def main():
|
||||
),
|
||||
]
|
||||
|
||||
# The fifth worker is opt-in: containment (QTZPOC-19) wants the importer in its own runtime with
|
||||
# its own credentials, and a flag lets that be a deployment decision instead of forcing every
|
||||
# Laborious runtime to poll the import queue today. Unset, the four workers above are unaffected
|
||||
# and no import queue is polled at all.
|
||||
if os.getenv('IMPORT_MODEL_ENABLED', '').strip().lower() in ('1', 'true', 'yes', 'on'):
|
||||
logger.custom_info('Import model worker enabled', metadata_runtime)
|
||||
workers.append(
|
||||
prepare_worker(
|
||||
temporal_client=temporal_client,
|
||||
main_workflow=ImportModel,
|
||||
other_workflows=[],
|
||||
activities=[
|
||||
# Import log row (the BFF database, update-only)
|
||||
activities.claim_import_status,
|
||||
activities.record_import_status,
|
||||
activities.record_import_names,
|
||||
activities.record_import_terminal_status,
|
||||
activities.report_import_status_write_failure,
|
||||
# Bundle
|
||||
activities.download_import_bundle,
|
||||
activities.open_import_bundle,
|
||||
# MLflow provisioning
|
||||
activities.create_import_experiment,
|
||||
activities.upload_import_artifacts,
|
||||
activities.register_import_model_version,
|
||||
# The model listing
|
||||
activities.write_import_model_document,
|
||||
# Cleanup and retention
|
||||
activities.cleanup_import_files,
|
||||
activities.ensure_import_bundle_retention,
|
||||
],
|
||||
logger=logger,
|
||||
runtime=runtime,
|
||||
)
|
||||
)
|
||||
|
||||
handlers = []
|
||||
for w in workers:
|
||||
handlers.append(w.run())
|
||||
|
||||
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}
|
||||
Reference in New Issue
Block a user