vitor-aignosi
4c7cae502f
QTZPOC-20: Refactor ModelImport status engine for lazy initialization, add byte-counted size validation to bundle extraction, optimize MLflow run ID lookups, and ensure workflow cancellation propagates correctly.
1251 lines
58 KiB
Python
1251 lines
58 KiB
Python
"""Model import activities: the import log row, the bundle, MLflow, the model document, cleanup.
|
|
|
|
The import log row is **not created here**. It is `public.experiment_run` in the Postgres database
|
|
`sientia-core-mlops-bff`; the frontend inserts it through the BFF with `run_type='IMPORT'` and
|
|
`status='PENDING'`, and passes its `id` in the workflow input. Every write in this module is an
|
|
`UPDATE ... WHERE id = :import_run_id`, on an engine of its own against that database — never the
|
|
`sientia` database Laborious itself uses, and never in a transaction shared with a provisioning
|
|
write.
|
|
|
|
Three classes of status write, and only the first one is strict:
|
|
|
|
- **A, the claim** — `WHERE id AND run_type='IMPORT' AND (PENDING OR my own workflow id)`. This is
|
|
the ownership decision, and a claim matching nothing fails the import explicitly.
|
|
- **B, progress** — `WHERE id`, nothing else. A hint. Its failure never fails the import.
|
|
- **C, terminal** — `WHERE id`, nothing else, in the workflow's `finally`, after cleanup. The
|
|
authority: it rewrites status, code, reason and step from one classification of the outcome.
|
|
"""
|
|
|
|
from temporalio import activity, workflow
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
import json
|
|
import shutil
|
|
from datetime import timedelta
|
|
from enum import StrEnum
|
|
from pathlib import Path
|
|
from typing import Any, Final
|
|
|
|
from minio.commonconfig import ENABLED, Filter
|
|
from minio.lifecycleconfig import Expiration, LifecycleConfig, Rule
|
|
from mlflow.entities import Param
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST, ErrorCode
|
|
from pymongo import ReturnDocument
|
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
|
from sientia_do.observability.logger import Logger
|
|
from sientia_do.observability.metrics_controller import MetricsController
|
|
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
|
from sientia_do.repository.minio_repository import MinioRepository
|
|
from sientia_do.repository.mongodb_repository import MongoDBRepository
|
|
from sientia_do.temporal.constants import now
|
|
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.engine import URL, Engine
|
|
from sqlalchemy.exc import DataError, IntegrityError, ProgrammingError
|
|
from sqlalchemy.sql import Executable
|
|
from temporalio.common import RetryPolicy
|
|
from temporalio.exceptions import ApplicationError
|
|
|
|
from laborious import metrics
|
|
from laborious.utils.bundle import format as fmt
|
|
from laborious.utils.bundle.reader import BundleGateError, BundleLimits, BundleReader
|
|
from laborious.utils.bundle.steps import ImportStep
|
|
from laborious.utils.import_password import (
|
|
EnvelopeError,
|
|
decrypt_password_envelope,
|
|
load_envelope_key,
|
|
)
|
|
from laborious.utils.messages import MISSING_MINIO_REPOSITORY as _MISSING_MINIO
|
|
from laborious.utils.models.import_errors import (
|
|
ImportErrorCode,
|
|
ImportGateRejectedError,
|
|
ImportLogRowNotClaimableError,
|
|
ImportMisconfiguredError,
|
|
ImportResponsibility,
|
|
ImportStatusRejectedError,
|
|
classify_import_failure,
|
|
report_import_failure,
|
|
)
|
|
from laborious.utils.repository.model_repository import MLFlowRepository
|
|
from laborious.utils.time import to_naive
|
|
|
|
# The table is a constant of the code, never an input: there is exactly one import log on this
|
|
# platform, and letting a caller name it would be a SQL-injection surface bought for nothing. The
|
|
# statements below write the name out literally rather than interpolating this constant — an
|
|
# interpolation is a shape a reader (and bandit) has to reason about, and there is nothing to
|
|
# parameterise. The constant is what the tests assert every statement against.
|
|
IMPORT_LOG_TABLE: Final[str] = 'public.experiment_run'
|
|
|
|
RUN_NAME_MAX_LENGTH: Final[int] = 50
|
|
RUN_NAME_PREFIX: Final[str] = 'import-'
|
|
|
|
_MISSING_MLFLOW = 'MLflow repository not initialized'
|
|
_MISSING_MONGO = 'MongoDB repository not initialized'
|
|
|
|
# Where a digest may sit inside the BFF-owned `request_data`. Checked in this order; absence is not
|
|
# a failure, only a disagreement is.
|
|
REQUEST_DIGEST_KEYS: Final[tuple[str, ...]] = (
|
|
'expected_digest',
|
|
'digest',
|
|
'sha256',
|
|
'file_digest',
|
|
)
|
|
|
|
|
|
class ImportState(StrEnum):
|
|
"""The workflow's internal vocabulary, mapped onto the shared status column by D17."""
|
|
|
|
RUNNING = 'RUNNING'
|
|
COMPLETED = 'COMPLETED'
|
|
ERROR = 'ERROR'
|
|
|
|
|
|
class ExperimentRunStatus(StrEnum):
|
|
"""The values this workflow writes into `experiment_run.status`.
|
|
|
|
Two are **reused** from the training flow because they describe a stage, not a flow — the row's
|
|
flow is told by `run_type`, not by the status. Two are **new in V12**, because
|
|
`TRAINING_SUCCESS` / `TRAINING_ERROR` would lie about where the row came from. Nobody should
|
|
"fix" the reuse into an `IMPORT_RUNNING` twin: it would carry no information `run_type` does not
|
|
already carry, and it would widen the pairing CHECK and the frontend's switch for nothing.
|
|
"""
|
|
|
|
ORCHESTRATOR_WAITING_PROC = 'ORCHESTRATOR_WAITING_PROC' # reused
|
|
ORCHESTRATOR_VALIDATION_ERROR = 'ORCHESTRATOR_VALIDATION_ERROR' # reused
|
|
IMPORT_SUCCESS = 'IMPORT_SUCCESS' # new in V12
|
|
IMPORT_ERROR = 'IMPORT_ERROR' # new in V12
|
|
|
|
|
|
# Only the status writes retry, and only because they are idempotent: `UPDATE ... SET <fixed values>
|
|
# WHERE id` leaves the row the same however many times it runs. A provisioning call has no such
|
|
# shape, which is why it keeps `maximum_attempts=1`. Defined here rather than in `sientia_do`, which
|
|
# other repositories share.
|
|
IMPORT_STATUS_RETRY_POLICY: Final[RetryPolicy] = RetryPolicy(
|
|
maximum_attempts=10,
|
|
initial_interval=timedelta(seconds=1),
|
|
backoff_coefficient=2.0,
|
|
maximum_interval=timedelta(seconds=30),
|
|
non_retryable_error_types=['ImportLogRowNotClaimableError', 'ImportStatusRejectedError'],
|
|
)
|
|
|
|
# Progress writes are hints nobody reads — the terminal write rewrites them — so a 151-second pause
|
|
# before each of a dozen stage updates would stall a working import for twenty minutes because of a
|
|
# database the import itself does not need.
|
|
IMPORT_STATUS_HINT_RETRY_POLICY: Final[RetryPolicy] = RetryPolicy(
|
|
maximum_attempts=3,
|
|
initial_interval=timedelta(seconds=1),
|
|
backoff_coefficient=2.0,
|
|
maximum_interval=timedelta(seconds=2),
|
|
non_retryable_error_types=['ImportLogRowNotClaimableError', 'ImportStatusRejectedError'],
|
|
)
|
|
|
|
_READ_ROW_SQL = text(
|
|
"""
|
|
SELECT id,
|
|
run_type,
|
|
status,
|
|
bucket_name,
|
|
file_name,
|
|
request_data,
|
|
orchestrator_response_data -> 'import' ->> 'workflow_id' AS claimed_by
|
|
FROM public.experiment_run
|
|
WHERE id = :import_run_id
|
|
"""
|
|
)
|
|
|
|
|
|
def _build_update(
|
|
columns: dict[str, str], *, casts: dict[str, str] | None = None, extra_where: str = ''
|
|
) -> Executable:
|
|
"""`UPDATE public.experiment_run SET <columns> WHERE id = :import_run_id [AND <extra_where>]`.
|
|
|
|
`columns` maps each SET target to its bind parameter name, e.g. `{'status': 'status'}` renders
|
|
`status = :status`. `casts` names which of those columns need `CAST(:name AS <cast>)` instead of
|
|
a bare `:name`. Order is preserved (Python 3.7+ dict order), so the rendered SQL is deterministic
|
|
and diff-friendly.
|
|
"""
|
|
casts = casts or {}
|
|
assignments = ',\n '.join(
|
|
f'{column} = CAST(:{param} AS {casts[column]})'
|
|
if column in casts
|
|
else f'{column} = :{param}'
|
|
for column, param in columns.items()
|
|
)
|
|
where = 'WHERE id = :import_run_id'
|
|
if extra_where:
|
|
where = f'{where}\n AND {extra_where}'
|
|
return text(
|
|
f'UPDATE public.experiment_run\n SET {assignments}\n {where}' # nosec B608 - columns/extra_where come only from literal constants in this module, never user input
|
|
)
|
|
|
|
|
|
_CLAIM_SQL = _build_update(
|
|
{'status': 'status', 'updated_at': 'updated_at', 'orchestrator_response_data': 'detail'},
|
|
casts={'orchestrator_response_data': 'JSONB'},
|
|
extra_where=(
|
|
"run_type = 'IMPORT'\n"
|
|
" AND (status = 'PENDING'\n"
|
|
" OR (status = 'ORCHESTRATOR_WAITING_PROC'\n"
|
|
" AND orchestrator_response_data -> 'import' ->> 'workflow_id' "
|
|
'= :workflow_id))'
|
|
),
|
|
)
|
|
|
|
# Class B and C key on the id alone: ownership was decided once, by the claim. Re-checking it on
|
|
# every later write buys nothing and costs the thing that matters — a predicate miss turning a
|
|
# pipeline that already imported the model into a reported failure.
|
|
_PROGRESS_SQL = _build_update(
|
|
{'status': 'status', 'updated_at': 'updated_at', 'orchestrator_response_data': 'detail'},
|
|
casts={'orchestrator_response_data': 'JSONB'},
|
|
)
|
|
|
|
_NAMES_SQL = _build_update(
|
|
{
|
|
'experiment_name': 'experiment_name',
|
|
'run_name': 'run_name',
|
|
'updated_at': 'updated_at',
|
|
'orchestrator_response_data': 'detail',
|
|
},
|
|
casts={'orchestrator_response_data': 'JSONB'},
|
|
)
|
|
|
|
_TERMINAL_SQL = _build_update(
|
|
{
|
|
'status': 'status',
|
|
'error_message': 'error_message',
|
|
'updated_at': 'updated_at',
|
|
'orchestrator_response_data': 'detail',
|
|
},
|
|
casts={'orchestrator_response_data': 'JSONB'},
|
|
)
|
|
|
|
|
|
def map_state_to_status(
|
|
state: ImportState, step: ImportStep | str | None = None
|
|
) -> ExperimentRunStatus:
|
|
"""Map the workflow's internal state onto the shared status column (design D17).
|
|
|
|
`RUNNING` is the orchestrator holding the row. `COMPLETED` is `IMPORT_SUCCESS`. An `ERROR` is
|
|
split in two by *who can act on it*: the user, for everything up to the content policy, or only
|
|
support, for everything after it — the same split the table already draws between
|
|
`ORCHESTRATOR_VALIDATION_ERROR` and `TRAINING_ERROR`.
|
|
"""
|
|
if state is ImportState.RUNNING:
|
|
return ExperimentRunStatus.ORCHESTRATOR_WAITING_PROC
|
|
if state is ImportState.COMPLETED:
|
|
return ExperimentRunStatus.IMPORT_SUCCESS
|
|
failure = classify_import_failure(ImportStep(step) if step is not None else None)
|
|
if failure.responsibility is ImportResponsibility.USER:
|
|
return ExperimentRunStatus.ORCHESTRATOR_VALIDATION_ERROR
|
|
return ExperimentRunStatus.IMPORT_ERROR
|
|
|
|
|
|
def derive_run_name(model_name: str, model_version: str) -> str:
|
|
"""`import-{model_name}-v{model_version}`, truncated inside the model-name segment.
|
|
|
|
Deterministic, because the recorded name and the created run's name must not diverge: the run is
|
|
created with exactly this value, passed in as an argument rather than recomputed. The `import-`
|
|
prefix guarantees the column's length-3 floor; the truncation is what keeps it under 50.
|
|
"""
|
|
suffix = f'-v{model_version}'
|
|
room = RUN_NAME_MAX_LENGTH - len(RUN_NAME_PREFIX) - len(suffix)
|
|
if room < 1:
|
|
# A version string long enough to leave no room for a name: keep the prefix and as much of
|
|
# the version as fits, so the value is still a legible, deterministic 50 characters.
|
|
return f'{RUN_NAME_PREFIX}{suffix.lstrip("-")}'[:RUN_NAME_MAX_LENGTH]
|
|
return f'{RUN_NAME_PREFIX}{model_name[:room]}{suffix}'
|
|
|
|
|
|
def build_import_detail(
|
|
*,
|
|
workflow_id: str,
|
|
step: ImportStep | str | None = None,
|
|
code: ImportErrorCode | str | None = None,
|
|
gate: int | None = None,
|
|
source: dict[str, Any] | None = None,
|
|
cleanup_failed: bool = False,
|
|
) -> str:
|
|
"""Build `orchestrator_response_data` as JSON text.
|
|
|
|
The whole `import` object is rewritten on every write, which is what keeps the workflow
|
|
identifier the claim wrote from ever being dropped. There is no column for any of this (design
|
|
D21): the code the frontend switches on and the identifier support quotes live at fixed keys in
|
|
the JSON column the workflow already owns.
|
|
"""
|
|
detail: dict[str, Any] = {'workflow_id': workflow_id, 'gate': gate}
|
|
if step is not None:
|
|
detail['step'] = str(step)
|
|
if code is not None:
|
|
detail['code'] = str(code)
|
|
if source:
|
|
detail['source'] = source
|
|
if cleanup_failed:
|
|
detail['cleanup_failed'] = True
|
|
return json.dumps({'import': detail})
|
|
|
|
|
|
class ModelImport(SientiaMonitoring):
|
|
"""Activities turning a `.sientia` object into a registered model, or into a durable ERROR."""
|
|
|
|
def __init__(
|
|
self,
|
|
status_config: dict[str, Any],
|
|
import_config: dict[str, Any],
|
|
minio_repository: MinioRepository | None = None,
|
|
mlflow_repository: MLFlowRepository | None = None,
|
|
mongo_repository: MongoDBRepository | None = None,
|
|
logger: Logger | None = None,
|
|
notification_handler: NotificationHandler | None = None,
|
|
metrics_controller: MetricsController | None = None,
|
|
):
|
|
"""
|
|
Args:
|
|
status_config: `build_import_status_config()` — the BFF database holding the import log.
|
|
This is a **second** Postgres connection — another database, on the `POSTGRES_*`
|
|
server unless the IMPORT_STATUS_DB_* overrides move it — and its separateness from
|
|
the one Laborious already holds is what makes design D7's transactional
|
|
independence literal.
|
|
import_config: `build_import_config()` — bucket, prefix, ceilings, work directory,
|
|
envelope key and the models collection name.
|
|
minio_repository: the shared MinIO repository, for the uploaded bundle.
|
|
mlflow_repository: the shared MLflow repository — the same instance the rest of the
|
|
worker uses, because setting the tracking URI is a process-wide side effect.
|
|
mongo_repository: MongoDB, where the model listing lives.
|
|
logger: logger instance.
|
|
notification_handler: the notification path the rest of the repository uses.
|
|
metrics_controller: metrics controller.
|
|
"""
|
|
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
|
self.import_config = import_config
|
|
self.status_config = status_config
|
|
self.import_minio_repository = minio_repository
|
|
self.import_mlflow_repository = mlflow_repository
|
|
self.mongo_repository = mongo_repository
|
|
self._status_engine: Engine | None = None
|
|
|
|
@property
|
|
def status_engine(self) -> Engine:
|
|
"""The BFF status database's engine, created on first access and cached from then on.
|
|
|
|
A worker pod where the import feature is never exercised never opens a connection pool
|
|
against that Postgres: nothing in this module reaches `status_engine` unless an import
|
|
activity actually runs.
|
|
"""
|
|
if self._status_engine is None:
|
|
self._status_engine = create_engine(
|
|
URL.create(
|
|
'postgresql',
|
|
username=self.status_config['user'],
|
|
password=self.status_config['password'],
|
|
host=self.status_config['host'],
|
|
port=self.status_config['port'],
|
|
database=self.status_config['dbname'],
|
|
),
|
|
pool_size=2,
|
|
max_overflow=3,
|
|
pool_pre_ping=True,
|
|
)
|
|
return self._status_engine
|
|
|
|
def close(self) -> None:
|
|
"""Dispose the status engine, when one was ever created. The import holds no other
|
|
connection."""
|
|
if self._status_engine is not None:
|
|
self._status_engine.dispose()
|
|
|
|
# ------------------------------------------------------------------ helpers
|
|
|
|
@staticmethod
|
|
def _timestamp() -> Any:
|
|
"""`updated_at` is `TIMESTAMP WITHOUT TIME ZONE`, so the value handed over is naive."""
|
|
return to_naive(now())
|
|
|
|
@staticmethod
|
|
def _digest_from_request_data(request_data: Any) -> str | None:
|
|
"""Find the uploader's digest inside the BFF-owned request payload, if it recorded one."""
|
|
if not isinstance(request_data, dict):
|
|
return None
|
|
for key in REQUEST_DIGEST_KEYS:
|
|
value = request_data.get(key)
|
|
if isinstance(value, str) and value:
|
|
return value.lower()
|
|
return None
|
|
|
|
def _execute_status_write(self, statement: Any, parameters: dict[str, Any]) -> int:
|
|
"""Run one status statement in its own transaction and return the rows it affected.
|
|
|
|
Raises:
|
|
ImportStatusRejectedError: the database accepted the connection and refused the
|
|
statement — a CHECK violation, an undefined column, a value too long. Non-retryable,
|
|
so a missing `V12` fails on the first attempt instead of after the full envelope.
|
|
Exception: anything else (connection refused, reset, timeout, a pod being replaced) is
|
|
left to propagate, because repeating it is exactly what fixes it.
|
|
"""
|
|
try:
|
|
with self.status_engine.begin() as connection:
|
|
return connection.execute(statement, parameters).rowcount
|
|
except (IntegrityError, ProgrammingError, DataError) as error:
|
|
raise ImportStatusRejectedError(
|
|
'the import log database refused the status statement'
|
|
) from error
|
|
|
|
def _read_row(self, import_run_id: Any) -> dict[str, Any] | None:
|
|
"""Read the record. A read, not a write: nothing of somebody else's is overwritten."""
|
|
with self.status_engine.connect() as connection:
|
|
row = connection.execute(_READ_ROW_SQL, {'import_run_id': import_run_id}).mappings()
|
|
found = row.first()
|
|
return dict(found) if found is not None else None
|
|
|
|
@staticmethod
|
|
def _describe_unclaimable(row: dict[str, Any] | None, workflow_id: str) -> str:
|
|
"""Name the cause of a zero-row claim: missing row, training row, or someone else's."""
|
|
if row is None:
|
|
return 'no such import log row'
|
|
if row.get('run_type') != 'IMPORT':
|
|
return f'row is a {row.get("run_type")} run, not an import'
|
|
claimed_by = row.get('claimed_by')
|
|
if claimed_by and claimed_by != workflow_id:
|
|
return f'row is already claimed by workflow {claimed_by}'
|
|
return f'row is in status {row.get("status")} and cannot be claimed'
|
|
|
|
# -------------------------------------------------------------- class A: claim
|
|
|
|
@activity.defn(name='claim_import_status')
|
|
async def claim_import_status(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Claim the import log row the frontend inserted. The workflow's first write.
|
|
|
|
One conditional `UPDATE`, whose second disjunct makes a re-execution of the same workflow
|
|
idempotent instead of a conflict — a lost response on attempt 1 must not become a false
|
|
"not mine" on attempt 2. Zero rows affected covers the id typo, the duplicated start and the
|
|
wrong `run_type` in one condition, with no read-then-write race, and it is what stops a
|
|
duplicated workflow start from provisioning twice.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata for logging and metrics.
|
|
import_run_id (int): the record's primary key, from the workflow input.
|
|
workflow_id (str): this workflow's identifier, recorded in the JSON detail.
|
|
source (dict, optional): the bundle's declared origin, when already known.
|
|
|
|
Returns:
|
|
dict: `claimed`, plus the row values the input guards are checked against —
|
|
`bucket_name`, `file_name` and `request_digest`.
|
|
|
|
Raises:
|
|
ImportLogRowNotClaimableError: the row is not this workflow's. Non-retryable, reported
|
|
through the channel, and **written to no row** — the loudness has to go somewhere
|
|
that is not somebody else's audit record.
|
|
"""
|
|
metadata = input_data.get('metadata', {})
|
|
import_run_id = input_data['import_run_id']
|
|
workflow_id = input_data['workflow_id']
|
|
|
|
try:
|
|
affected = self._execute_status_write(
|
|
_CLAIM_SQL,
|
|
{
|
|
'status': ExperimentRunStatus.ORCHESTRATOR_WAITING_PROC.value,
|
|
'updated_at': self._timestamp(),
|
|
'detail': build_import_detail(
|
|
workflow_id=workflow_id,
|
|
step=ImportStep.RECEIVED,
|
|
source=input_data.get('source'),
|
|
),
|
|
'import_run_id': import_run_id,
|
|
'workflow_id': workflow_id,
|
|
},
|
|
)
|
|
except ImportStatusRejectedError as error:
|
|
await report_import_failure(
|
|
self,
|
|
None,
|
|
error,
|
|
import_run_id=import_run_id,
|
|
workflow_id=workflow_id,
|
|
metadata=metadata,
|
|
code=ImportErrorCode.STATUS_NOT_RECORDED,
|
|
marker='status_write',
|
|
)
|
|
raise
|
|
|
|
row = self._read_row(import_run_id)
|
|
if affected == 0:
|
|
observed = self._describe_unclaimable(row, workflow_id)
|
|
not_claimable = ImportLogRowNotClaimableError(
|
|
f'import log row {import_run_id} could not be claimed: {observed}',
|
|
import_run_id=import_run_id,
|
|
observed=observed,
|
|
)
|
|
await report_import_failure(
|
|
self,
|
|
None,
|
|
not_claimable,
|
|
import_run_id=import_run_id,
|
|
workflow_id=workflow_id,
|
|
metadata=metadata,
|
|
code=ImportErrorCode.LOG_ROW_NOT_CLAIMABLE,
|
|
marker='not_claimable',
|
|
)
|
|
raise not_claimable
|
|
|
|
await self.emit_metric(
|
|
metric_object=metrics.MODEL_IMPORT_STARTED_COUNT,
|
|
tags=self.get_core_labels(metadata, operation_type=ImportStep.RECEIVED.value),
|
|
)
|
|
self.info(f'import log row {import_run_id} claimed by {workflow_id}', metadata)
|
|
row = row or {}
|
|
return {
|
|
'claimed': True,
|
|
'bucket_name': row.get('bucket_name'),
|
|
'file_name': row.get('file_name'),
|
|
'request_digest': self._digest_from_request_data(row.get('request_data')),
|
|
}
|
|
|
|
# ----------------------------------------------------------- class B: progress
|
|
|
|
@activity.defn(name='record_import_status')
|
|
async def record_import_status(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Record the step the workflow is about to attempt. A hint, keyed on the id alone.
|
|
|
|
Its failure is **not** an import failure: the workflow reports it through the channel and
|
|
carries on, because losing a hint costs a record pointing one step behind and nothing else —
|
|
the terminal write rewrites status, code, reason and step from one classification.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
import_run_id (int): the record's primary key.
|
|
workflow_id (str): recorded in the detail, so it is never dropped.
|
|
step (str): the step about to be attempted.
|
|
source (dict, optional): the bundle's declared origin.
|
|
"""
|
|
affected = self._execute_status_write(
|
|
_PROGRESS_SQL,
|
|
{
|
|
'status': ExperimentRunStatus.ORCHESTRATOR_WAITING_PROC.value,
|
|
'updated_at': self._timestamp(),
|
|
'detail': build_import_detail(
|
|
workflow_id=input_data['workflow_id'],
|
|
step=input_data['step'],
|
|
source=input_data.get('source'),
|
|
),
|
|
'import_run_id': input_data['import_run_id'],
|
|
},
|
|
)
|
|
return {'recorded': affected > 0, 'rows': affected}
|
|
|
|
@activity.defn(name='record_import_names')
|
|
async def record_import_names(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Write the experiment and run names, before anything exists in MLflow.
|
|
|
|
Both columns exist so a person watching the import list sees *which* model is being
|
|
imported. The earliest moment both values are known is when the bundle is open, so that is
|
|
when they are written — before the experiment exists, before the run exists.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
import_run_id (int): the record's primary key.
|
|
workflow_id (str): recorded in the detail.
|
|
experiment_name (str): the bundle's own experiment name, verbatim.
|
|
run_name (str): the derived run name, which the run is then created with.
|
|
source (dict, optional): the bundle's declared origin.
|
|
"""
|
|
affected = self._execute_status_write(
|
|
_NAMES_SQL,
|
|
{
|
|
'experiment_name': input_data['experiment_name'],
|
|
'run_name': input_data['run_name'],
|
|
'updated_at': self._timestamp(),
|
|
'detail': build_import_detail(
|
|
workflow_id=input_data['workflow_id'],
|
|
step=input_data.get('step', ImportStep.EXPERIMENT_CREATION),
|
|
source=input_data.get('source'),
|
|
),
|
|
'import_run_id': input_data['import_run_id'],
|
|
},
|
|
)
|
|
return {'recorded': affected > 0, 'rows': affected}
|
|
|
|
# ----------------------------------------------------------- class C: terminal
|
|
|
|
@activity.defn(name='record_import_terminal_status')
|
|
async def record_import_terminal_status(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""The single authoritative write, run from the workflow's `finally` after cleanup.
|
|
|
|
Every failure is classified **once**, here, from the step that failed. On success the reason
|
|
column is written as `NULL` — never `''`, which its length-3 CHECK would refuse — and the
|
|
detail carries no code. A cleanup failure is noted in the detail without changing a
|
|
successful status: the user's model was imported, and reporting ERROR would be false.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
import_run_id (int): the record's primary key.
|
|
workflow_id (str): recorded in the detail.
|
|
succeeded (bool): whether every step succeeded.
|
|
step (str, optional): the failing step, or the last step reached on success.
|
|
gate (int, optional): the gate that rejected the bundle, when one did.
|
|
source (dict, optional): the bundle's declared origin.
|
|
cleanup_failed (bool, optional): whether the temporary-file cleanup failed.
|
|
|
|
Returns:
|
|
dict: the verdict that was written — status, code and reason — so a caller that has to
|
|
report a failure to record it still holds the verdict it could not write.
|
|
"""
|
|
metadata = input_data.get('metadata', {})
|
|
succeeded = bool(input_data.get('succeeded'))
|
|
step = input_data.get('step')
|
|
cleanup_failed = bool(input_data.get('cleanup_failed'))
|
|
|
|
if succeeded:
|
|
status = map_state_to_status(ImportState.COMPLETED)
|
|
code: ImportErrorCode | None = None
|
|
reason: str | None = None
|
|
else:
|
|
try:
|
|
step_enum = ImportStep(step) if step else None
|
|
except ValueError:
|
|
step_enum = None
|
|
failure = classify_import_failure(step_enum)
|
|
status = map_state_to_status(ImportState.ERROR, step_enum)
|
|
code, reason = failure.code, failure.reason
|
|
|
|
verdict = {
|
|
'status': status.value,
|
|
'code': code.value if code else None,
|
|
'reason': reason,
|
|
'step': str(step) if step else None,
|
|
}
|
|
|
|
self._execute_status_write(
|
|
_TERMINAL_SQL,
|
|
{
|
|
'status': status.value,
|
|
'error_message': reason,
|
|
'updated_at': self._timestamp(),
|
|
'detail': build_import_detail(
|
|
workflow_id=input_data['workflow_id'],
|
|
step=step,
|
|
code=code,
|
|
gate=input_data.get('gate'),
|
|
source=input_data.get('source'),
|
|
cleanup_failed=cleanup_failed,
|
|
),
|
|
'import_run_id': input_data['import_run_id'],
|
|
},
|
|
)
|
|
|
|
if succeeded:
|
|
await self.emit_metric(
|
|
metric_object=metrics.MODEL_IMPORT_COMPLETED_COUNT,
|
|
tags=self.get_core_labels(metadata, operation_type=status.value),
|
|
)
|
|
self.info(f'import {input_data["import_run_id"]} recorded as {status.value}', metadata)
|
|
return verdict
|
|
|
|
@activity.defn(name='report_import_status_write_failure')
|
|
async def report_import_status_write_failure(
|
|
self, input_data: dict[str, Any]
|
|
) -> dict[str, Any]:
|
|
"""Report a status write that could not be made — the one failure with no row to land on.
|
|
|
|
Called by the workflow when a progress or terminal write exhausted its attempts or was
|
|
refused. It carries the verdict it could not write, so the outcome survives in text even
|
|
when it cannot survive in the row, and it never turns a completed import into a failed one.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
import_run_id (int): the record's primary key.
|
|
workflow_id (str): the identifier that correlates this line with the record.
|
|
detail (str): the underlying failure, as text — it crossed an activity boundary, so the
|
|
exception object itself is gone by now.
|
|
verdict (str, optional): the verdict the terminal write could not record.
|
|
"""
|
|
await report_import_failure(
|
|
self,
|
|
None,
|
|
None,
|
|
import_run_id=input_data['import_run_id'],
|
|
workflow_id=input_data['workflow_id'],
|
|
metadata=input_data.get('metadata', {}),
|
|
code=ImportErrorCode.STATUS_NOT_RECORDED,
|
|
marker='status_write',
|
|
verdict=input_data.get('verdict'),
|
|
detail_text=input_data.get('detail'),
|
|
)
|
|
return {'reported': True}
|
|
|
|
@activity.defn(name='report_import_pipeline_failure')
|
|
async def report_import_pipeline_failure(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Report the failure that ended the import — the one every step shares.
|
|
|
|
Until this existed, a failure at `download` … `model_document` reached the observability
|
|
surface only as whatever the Temporal SDK printed for the failing attempt: a bare traceback,
|
|
with no `import_run_id`, no step, no code, and no pod or runtime label, so it could not be
|
|
joined to the record it failed. It also never incremented the import error metric and never
|
|
notified, which the two failures that *are* reported (a row that cannot be claimed, a verdict
|
|
that cannot be written) always did.
|
|
|
|
The exception object is gone by the time the workflow gets here — it crossed an activity
|
|
boundary — so the failure travels as text and its type as a name, and the classification comes
|
|
from the step, exactly as the record's does. Reporting is never a reason to fail: an
|
|
unrecognised step is reported as unclassified rather than raising over a spelling.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
import_run_id (int): the record's primary key.
|
|
workflow_id (str): the identifier a user quotes to support.
|
|
step (str, optional): the step that failed. Absent means unclassified.
|
|
gate (int, optional): the gate that rejected the bundle, when one did.
|
|
detail (str): the failure as text, root cause first.
|
|
error_type (str, optional): the exception type name that travelled with it.
|
|
"""
|
|
raw_step = input_data.get('step')
|
|
try:
|
|
step = ImportStep(raw_step) if raw_step else None
|
|
except ValueError:
|
|
step = None
|
|
await report_import_failure(
|
|
self,
|
|
step,
|
|
None,
|
|
import_run_id=input_data.get('import_run_id'),
|
|
workflow_id=input_data['workflow_id'],
|
|
metadata=input_data.get('metadata', {}),
|
|
marker=str(raw_step) if raw_step else None,
|
|
detail_text=input_data.get('detail'),
|
|
error_type=input_data.get('error_type'),
|
|
gate=input_data.get('gate'),
|
|
)
|
|
return {'reported': True}
|
|
|
|
# ------------------------------------------------------------- bundle helpers
|
|
|
|
def _bundle_reader(self) -> BundleReader:
|
|
"""A reader carrying the configured ceilings."""
|
|
return BundleReader(
|
|
BundleLimits(
|
|
max_object_bytes=self.import_config['max_object_bytes'],
|
|
max_entries=self.import_config['max_entries'],
|
|
max_uncompressed_bytes=self.import_config['max_uncompressed_bytes'],
|
|
max_compression_ratio=self.import_config['max_compression_ratio'],
|
|
object_prefix=self.import_config['prefix'],
|
|
)
|
|
)
|
|
|
|
def _work_dir(self, workflow_id: str) -> Path:
|
|
"""This import's own scratch directory. One per workflow, never shared."""
|
|
return Path(self.import_config['work_dir']) / workflow_id
|
|
|
|
def _effective_bucket(self, requested: str | None) -> str:
|
|
return requested or self.import_config['bucket']
|
|
|
|
# --------------------------------------------------------- download and open
|
|
|
|
async def _download_bundle(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Gate 1, then the download, then the mandatory digest comparison.
|
|
|
|
The object's size and key shape are checked from `stat_object`, before its body is fetched.
|
|
The SHA-256 of what was downloaded is then compared against `expected_digest` — the value the
|
|
uploader computed over the bytes it uploaded — with **no path that skips the comparison**,
|
|
because a missing value was already rejected at `received`. A mismatch fails here, before the
|
|
header is parsed: it means the object behind that key is not the file the user chose.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
workflow_id (str): names this import's scratch directory.
|
|
bucket (str, optional): the effective bucket; the configured one when absent.
|
|
object_key (str): the uploaded object's key.
|
|
expected_digest (str): the uploader's SHA-256, lowercase hex.
|
|
|
|
Returns:
|
|
dict: `path`, `size` and `digest` of the downloaded file.
|
|
|
|
Raises:
|
|
ImportGateRejectedError: gate 1 rejected the object, or the digests disagree.
|
|
"""
|
|
if self.import_minio_repository is None:
|
|
raise ImportMisconfiguredError(_MISSING_MINIO)
|
|
|
|
metadata = input_data.get('metadata', {})
|
|
bucket = self._effective_bucket(input_data.get('bucket'))
|
|
object_key = input_data['object_key']
|
|
expected_digest = input_data['expected_digest']
|
|
reader = self._bundle_reader()
|
|
|
|
client = self.import_minio_repository.minio_client
|
|
try:
|
|
stat = client.stat_object(bucket, object_key)
|
|
except Exception as error:
|
|
raise ImportGateRejectedError(
|
|
'the uploaded object could not be read', step=ImportStep.DOWNLOAD, gate=1
|
|
) from error
|
|
|
|
try:
|
|
reader.check_object_limits(size=stat.size or 0, key=object_key)
|
|
except BundleGateError as error:
|
|
raise ImportGateRejectedError(str(error), step=error.step, gate=error.gate) from error
|
|
|
|
work_dir = self._work_dir(input_data['workflow_id'])
|
|
work_dir.mkdir(parents=True, exist_ok=True)
|
|
target = work_dir / 'bundle.sientia'
|
|
try:
|
|
client.fget_object(bucket, object_key, str(target))
|
|
except Exception as error:
|
|
raise ImportGateRejectedError(
|
|
'the uploaded object could not be downloaded', step=ImportStep.DOWNLOAD, gate=1
|
|
) from error
|
|
|
|
digest = fmt.file_digest(target)
|
|
if digest != str(expected_digest).lower():
|
|
target.unlink(missing_ok=True)
|
|
raise ImportGateRejectedError(
|
|
'the downloaded object does not match the digest the uploader computed',
|
|
step=ImportStep.DOWNLOAD,
|
|
gate=1,
|
|
)
|
|
|
|
self.info(f'import bundle downloaded from {bucket} ({stat.size} bytes)', metadata)
|
|
return {'path': str(target), 'size': stat.size, 'digest': digest}
|
|
|
|
async def _open_bundle(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Open the envelope, then the bundle: gates 2 to 7, in the reader's own order.
|
|
|
|
The password envelope is decrypted **here**, inside the activity that uses the password, so
|
|
no plaintext password is ever an activity input, an activity result or a workflow variable —
|
|
and therefore never lands in the Temporal event history. The recovered password is a local
|
|
value and nothing else.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
workflow_id (str): names this import's scratch directory.
|
|
path (str): the downloaded bundle.
|
|
password_envelope (str): base64 `nonce || ciphertext || tag`, AES-256-GCM.
|
|
|
|
Returns:
|
|
dict: `extracted_dir`, `metadata`, `parameters`, `digest` — nothing secret.
|
|
|
|
Raises:
|
|
ImportGateRejectedError: with the step and the gate that rejected the bundle.
|
|
"""
|
|
metadata = input_data.get('metadata', {})
|
|
work_dir = self._work_dir(input_data['workflow_id'])
|
|
extract_root = work_dir / 'extract'
|
|
|
|
try:
|
|
key = load_envelope_key(self.import_config.get('password_key'))
|
|
password = decrypt_password_envelope(input_data['password_envelope'], key)
|
|
except EnvelopeError as error:
|
|
# A distinct failure from a wrong *bundle* password, and never confused with it: the
|
|
# envelope is the platform's own transport, and its message says so. Same step, because
|
|
# the step is where the import stopped, not who is at fault.
|
|
raise ImportGateRejectedError(
|
|
'the import password envelope could not be opened',
|
|
step=ImportStep.DECRYPTION,
|
|
gate=3,
|
|
cause='envelope',
|
|
) from error
|
|
|
|
reader = self._bundle_reader()
|
|
try:
|
|
opened = reader.open(input_data['path'], password, extract_root, work_dir)
|
|
except BundleGateError as error:
|
|
raise ImportGateRejectedError(str(error), step=error.step, gate=error.gate) from error
|
|
finally:
|
|
del password
|
|
|
|
await self.emit_metric(
|
|
metric_object=metrics.MODEL_IMPORT_KDF_DURATION,
|
|
tags=self.get_core_labels(metadata, operation_type=ImportStep.DECRYPTION.value),
|
|
method='observe',
|
|
value=opened.kdf_seconds,
|
|
)
|
|
self.info(f'import bundle opened into {opened.extracted_dir}', metadata)
|
|
return {
|
|
'extracted_dir': str(opened.extracted_dir),
|
|
'metadata': opened.metadata,
|
|
'parameters': opened.parameters,
|
|
'digest': opened.digest,
|
|
}
|
|
|
|
# ------------------------------------------------------------ MLflow provisioning
|
|
|
|
async def _create_experiment(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Get or create the experiment named by the bundle, keeping the origin's name.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
experiment_name (str): the bundle's `metadata.experiment_name`.
|
|
|
|
Returns:
|
|
dict: `experiment_id` and `experiment_name`.
|
|
"""
|
|
if self.import_mlflow_repository is None:
|
|
raise ImportMisconfiguredError(_MISSING_MLFLOW)
|
|
|
|
experiment_name = input_data['experiment_name']
|
|
experiment = self.import_mlflow_repository.get_experiment(
|
|
experiment_name, create_if_not_exists=True
|
|
)
|
|
self.info(
|
|
f'import experiment ready: {experiment_name} ({experiment.experiment_id})',
|
|
input_data.get('metadata', {}),
|
|
)
|
|
return {'experiment_id': experiment.experiment_id, 'experiment_name': experiment_name}
|
|
|
|
async def _upload_artifacts(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Create a new run with the recorded name and log the bundle's tree and parameters.
|
|
|
|
Always a new run — the importer never appends to an existing one. The run is created with
|
|
exactly the `run_name` the fused activity computed, and the workflow records that same value
|
|
onto the row right after the fused activity returns, so the two cannot diverge. The artifact
|
|
tree is logged verbatim, keeping `prediction_model` and `data_model` at the run's artifact
|
|
root, which is what the legacy loader resolves. The bundle's `parameters` are logged
|
|
unchanged, and the origin's own identity is logged alongside them under `source.*` — records
|
|
of where the model came from, not links.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
experiment_id (str): the experiment to create the run in.
|
|
run_name (str): the name to create the run with.
|
|
extracted_dir (str): the opened bundle's directory.
|
|
parameters (dict): the bundle's `parameters` block, verbatim.
|
|
source (dict): the bundle's `metadata` block, logged as `source.*`.
|
|
|
|
Returns:
|
|
dict: `run_id` and `run_name`.
|
|
"""
|
|
if self.import_mlflow_repository is None:
|
|
raise ImportMisconfiguredError(_MISSING_MLFLOW)
|
|
|
|
client = self.import_mlflow_repository.client
|
|
run = client.create_run(
|
|
experiment_id=input_data['experiment_id'], run_name=input_data['run_name']
|
|
)
|
|
run_id = run.info.run_id
|
|
|
|
artifacts_dir = Path(input_data['extracted_dir']) / 'artifacts'
|
|
client.log_artifacts(run_id, str(artifacts_dir))
|
|
|
|
params = [
|
|
Param(name, str(value)) for name, value in (input_data.get('parameters') or {}).items()
|
|
]
|
|
params += [
|
|
Param(f'source.{name}', str(value))
|
|
for name, value in (input_data.get('source') or {}).items()
|
|
]
|
|
for i in range(0, len(params), 100): # MLflow's log_batch limit is 100 params per call
|
|
client.log_batch(run_id, params=params[i : i + 100])
|
|
|
|
client.set_terminated(run_id, 'FINISHED')
|
|
self.info(f'import artifacts uploaded to run {run_id}', input_data.get('metadata', {}))
|
|
return {'run_id': run_id, 'run_name': input_data['run_name']}
|
|
|
|
@activity.defn(name='prepare_and_upload_import_bundle')
|
|
async def prepare_and_upload_import_bundle(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Download, decrypt, open, create the experiment, and upload artifacts — one disk, one try.
|
|
|
|
`_download_bundle`, `_open_bundle`, `_create_experiment` and `_upload_artifacts` used to be
|
|
four separately-schedulable Temporal activities sharing one worker-local scratch directory,
|
|
with nothing guaranteeing they landed on the same worker process — a real hazard once the
|
|
import worker runs more than one replica, since a later phase reading a file an earlier
|
|
phase wrote on a different pod fails with `FileNotFoundError`. Fusing them into one activity
|
|
means every disk write and every disk read this import makes happens inside one attempt, on
|
|
one process.
|
|
|
|
The run name is computed here, from the bundle's own `model_name`/`model_version`, because it
|
|
is needed to create the run — the workflow cannot compute it beforehand, since it does not
|
|
learn the bundle's metadata until this activity returns. The workflow records that same
|
|
`run_name` onto the row immediately after this activity returns (before the registered
|
|
version and the model document are written), rather than before the run is created, which is
|
|
the narrowed guarantee `model-import-status-record`'s names requirement now makes.
|
|
|
|
A retry of this activity redoes all four phases from the start. Re-uploading the same
|
|
artifacts to the same run is safe — MLflow overwrites identical content rather than
|
|
duplicating it — so the cost of a retry is bandwidth and time, not corruption.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
workflow_id (str): names this import's scratch directory.
|
|
bucket (str, optional): the effective bucket; the configured one when absent.
|
|
object_key (str): the uploaded object's key.
|
|
expected_digest (str): the uploader's SHA-256, lowercase hex.
|
|
password_envelope (str): base64 `nonce || ciphertext || tag`, AES-256-GCM.
|
|
|
|
Returns:
|
|
dict: `extracted_dir`, `metadata` (the bundle's own, i.e. `source`), `parameters`,
|
|
`digest`, `experiment_id`, `experiment_name`, `run_id`, `run_name`.
|
|
"""
|
|
downloaded = await self._download_bundle(input_data)
|
|
opened = await self._open_bundle({**input_data, **downloaded})
|
|
source = opened['metadata']
|
|
run_name = derive_run_name(source['model_name'], source['model_version'])
|
|
try:
|
|
experiment = await self._create_experiment(
|
|
{**input_data, 'experiment_name': source['experiment_name']}
|
|
)
|
|
except ImportMisconfiguredError:
|
|
raise
|
|
except Exception as error:
|
|
# Raw MLflow/network failures carry no location of their own (unlike the gates, which
|
|
# raise `ImportGateRejectedError`), so without this the workflow's fallback step —
|
|
# whatever it was attempting before this fused activity started — misclassifies the
|
|
# failure. Retryable: this is a broken dependency, not a rejected bundle.
|
|
raise ApplicationError(
|
|
str(error), {'step': ImportStep.EXPERIMENT_CREATION.value}
|
|
) from error
|
|
|
|
try:
|
|
run = await self._upload_artifacts(
|
|
{
|
|
**input_data,
|
|
**opened,
|
|
**experiment,
|
|
'run_name': run_name,
|
|
'source': source,
|
|
}
|
|
)
|
|
except ImportMisconfiguredError:
|
|
raise
|
|
except Exception as error:
|
|
raise ApplicationError(
|
|
str(error), {'step': ImportStep.ARTIFACT_UPLOAD.value}
|
|
) from error
|
|
return {**opened, **experiment, **run}
|
|
|
|
@activity.defn(name='register_import_model_version')
|
|
async def register_import_model_version(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Register a version whose source the legacy loader can parse. No stage change.
|
|
|
|
The source is the run's **resolved artifact URI** with `prediction_model` appended —
|
|
`mlflow-artifacts:/<experiment_id>/<run_id>/artifacts/prediction_model` on this deployment —
|
|
because `get_model_run_id` extracts the run id as `source.split('/')[2]`. That index is the
|
|
run id only in the resolved form: a literal `runs:/<run_id>/prediction_model` has just two
|
|
slashes, so index 2 is the string `prediction_model`, and a `models:/m-*` source (what the
|
|
new stack emits) has none of the parts at all. This is the same value the training flow
|
|
stores, because `mlflow.register_model` resolves the `runs:/` URI before writing it
|
|
(`model_repository.py`, `register_model`) — `create_model_version` does not.
|
|
|
|
The stage is left untouched: a `COMPLETED` import is a registered model, not a servable one,
|
|
and promoting it to `Production` is a manual operator step outside this workflow.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
model_name (str): the bundle's `metadata.model_name`, verbatim.
|
|
run_id (str): the run this import created.
|
|
|
|
Returns:
|
|
dict: `model_name`, `version` and `source`.
|
|
"""
|
|
if self.import_mlflow_repository is None:
|
|
raise ImportMisconfiguredError(_MISSING_MLFLOW)
|
|
|
|
client = self.import_mlflow_repository.client
|
|
model_name = input_data['model_name']
|
|
run_id = input_data['run_id']
|
|
|
|
try:
|
|
client.get_registered_model(model_name)
|
|
except MlflowException as error:
|
|
if error.error_code != ErrorCode.Name(RESOURCE_DOES_NOT_EXIST):
|
|
raise
|
|
# First import of this name: the registered model is created on demand. A second import
|
|
# lands as a new version of the same registered model, never a differently named one.
|
|
client.create_registered_model(model_name)
|
|
|
|
source = f'{str(client.get_run(run_id).info.artifact_uri).rstrip("/")}/prediction_model'
|
|
version = client.create_model_version(name=model_name, source=source, run_id=run_id)
|
|
self.info(
|
|
f'import registered {model_name} version {version.version}',
|
|
input_data.get('metadata', {}),
|
|
)
|
|
return {'model_name': model_name, 'version': str(version.version), 'source': source}
|
|
|
|
# --------------------------------------------------- the MongoDB model document
|
|
|
|
@activity.defn(name='write_import_model_document')
|
|
async def write_import_model_document(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Write the document that *is* the model listing, inert, last.
|
|
|
|
This is what the runtime reads to know how to load a model — `model_config.transform_flavor`,
|
|
`model_config.predict_flavor` and `model_config.target` — so without it an imported model is
|
|
invisible to every workflow. `id` continues the collection's own numbering, assigned by an
|
|
atomic counter (`_next_model_document_id`) rather than a scan of existing ids, so two
|
|
concurrent imports can never collide on the same id. A name already in the collection is left
|
|
exactly as it is: `active`, `target` and the flavors of a model already on the platform are
|
|
operator-owned.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
model_name (str): the bundle's `metadata.model_name`, verbatim.
|
|
target (str): the bundle's `parameters.target_variable`.
|
|
|
|
Returns:
|
|
dict: `written`, `id`, `skipped` and `ignored_ids` (always `0`, kept for shape
|
|
compatibility — no id is skipped under the counter).
|
|
"""
|
|
if self.mongo_repository is None:
|
|
raise ImportMisconfiguredError(_MISSING_MONGO)
|
|
|
|
metadata = input_data.get('metadata', {})
|
|
collection = self.import_config['models_collection']
|
|
model_name = input_data['model_name']
|
|
|
|
existing = await self.mongo_repository.find_one(collection, {'name': model_name}, metadata)
|
|
if existing is not None:
|
|
self.info(
|
|
f'model document for {model_name} already exists; leaving it untouched', metadata
|
|
)
|
|
return {'written': False, 'skipped': True, 'id': existing.get('id')}
|
|
|
|
next_id, ignored = await self._next_model_document_id(collection, metadata)
|
|
document = {
|
|
'id': next_id,
|
|
'name': model_name,
|
|
'active': False,
|
|
'model_config': {
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'joblib',
|
|
'retention_minutes': 60,
|
|
'target': input_data['target'],
|
|
'skip_predict_timestamp': True,
|
|
},
|
|
}
|
|
await self.mongo_repository.insert(collection, document, metadata)
|
|
self.info(f'model document written for {model_name} with id {next_id}', metadata)
|
|
return {'written': True, 'skipped': False, 'id': next_id, 'ignored_ids': ignored}
|
|
|
|
async def _next_model_document_id(
|
|
self, collection: str, metadata: dict[str, Any]
|
|
) -> tuple[str, int]:
|
|
"""Atomically increment a counter dedicated to `collection`, and return its new value.
|
|
|
|
Reaches MongoDB through the repository's own public `database` handle — the same access
|
|
pattern `find_one`/`distinct` already use internally — with pymongo's native
|
|
`find_one_and_update`, rather than the max-scan this replaces. That scan raced: two imports
|
|
writing at the same time could both read the same maximum and collide on the same id. An
|
|
atomic `$inc` cannot.
|
|
|
|
The counter's backing collection (`import_config['counters_collection']`) is an external
|
|
dependency: it must be created and seeded by a migration before this runs against a given
|
|
environment (see `openspec/changes/model-import-fix-bugs/design.md`, Decision 3). This code
|
|
does not create or seed it.
|
|
|
|
Assumes `self.mongo_repository` is not `None` — the only caller, `write_import_model_document`,
|
|
already raises before calling this when it is.
|
|
"""
|
|
counters = self.mongo_repository.database[ # type: ignore[union-attr]
|
|
self.import_config['counters_collection']
|
|
]
|
|
doc = counters.find_one_and_update(
|
|
{'_id': collection},
|
|
{'$inc': {'seq': 1}},
|
|
upsert=True,
|
|
return_document=ReturnDocument.AFTER,
|
|
)
|
|
return str(doc['seq']), 0 # type: ignore[index]
|
|
|
|
# ----------------------------------------------------------- cleanup and retention
|
|
|
|
@activity.defn(name='cleanup_import_files')
|
|
async def cleanup_import_files(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Remove this import's scratch directory. Never the uploaded object.
|
|
|
|
Runs on every path, success or failure, and **never raises**: a cleanup failure is reported
|
|
through the channel and noted in the record's detail, but the user's model was imported and
|
|
reporting an error would be false.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
workflow_id (str): names the directory to remove.
|
|
import_run_id (int): for correlation when the removal fails.
|
|
|
|
Returns:
|
|
dict: `cleaned` and, when it failed, the reason as text.
|
|
"""
|
|
metadata = input_data.get('metadata', {})
|
|
work_dir = self._work_dir(input_data['workflow_id'])
|
|
try:
|
|
shutil.rmtree(work_dir, ignore_errors=False)
|
|
except FileNotFoundError:
|
|
# Nothing was downloaded, or a gate removed its own partial tree. Not a failure.
|
|
return {'cleaned': True, 'existed': False}
|
|
except OSError as error:
|
|
await report_import_failure(
|
|
self,
|
|
ImportStep.CLEANUP,
|
|
error,
|
|
import_run_id=input_data.get('import_run_id'),
|
|
workflow_id=input_data['workflow_id'],
|
|
metadata=metadata,
|
|
code=ImportErrorCode.CLEANUP_INCOMPLETE,
|
|
marker=ImportStep.CLEANUP.value,
|
|
)
|
|
return {'cleaned': False, 'existed': True, 'reason': type(error).__name__}
|
|
return {'cleaned': True, 'existed': True}
|
|
|
|
@activity.defn(name='ensure_import_bundle_retention')
|
|
async def ensure_import_bundle_retention(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Ensure the 7-day expiry rule on the import prefix, idempotently.
|
|
|
|
Applied by the importer rather than assumed from a Helm value, so the acceptance criterion
|
|
("lifecycle active") is verifiable from the code that depends on it. The uploaded object is
|
|
never deleted by the import itself — it is left to this rule, as the record of an attempt.
|
|
|
|
See also: `Storage.cleanup_minio_objects_expired` (laborious/activities/storage.py) — an
|
|
older, unrelated scan-and-delete expiry mechanism for a different prefix family. The two are
|
|
independent by design for now; see docs/minio-object-expiry.md.
|
|
|
|
Args (input_data):
|
|
metadata (dict): workflow metadata.
|
|
bucket (str, optional): the effective bucket; the configured one when absent.
|
|
|
|
Returns:
|
|
dict: whether the rule was already there and what it expires.
|
|
"""
|
|
if self.import_minio_repository is None:
|
|
raise ImportMisconfiguredError(_MISSING_MINIO)
|
|
|
|
metadata = input_data.get('metadata', {})
|
|
bucket = self._effective_bucket(input_data.get('bucket'))
|
|
prefix = self.import_config['prefix']
|
|
days = self.import_config['retention_days']
|
|
rule_id = 'sientia-import-bundle-expiry'
|
|
client = self.import_minio_repository.minio_client
|
|
|
|
try:
|
|
existing = client.get_bucket_lifecycle(bucket)
|
|
except Exception:
|
|
# A bucket with no lifecycle configuration at all answers with an error on some
|
|
# deployments and with `None` on others; both mean "there is no rule yet".
|
|
existing = None
|
|
|
|
existing_rule = next(
|
|
(rule for rule in (existing.rules if existing else []) if rule.rule_id == rule_id), None
|
|
)
|
|
if (
|
|
existing_rule is not None
|
|
and existing_rule.status == ENABLED
|
|
and getattr(existing_rule.expiration, 'days', None) == days
|
|
and getattr(existing_rule.rule_filter, 'prefix', None) == prefix
|
|
):
|
|
# The read still runs every time (it is how drift — someone editing the rule outside
|
|
# this code — gets detected); only the write is skipped when nothing has drifted.
|
|
self.info(f'import retention rule already correct on {bucket} ({days} days)', metadata)
|
|
return {
|
|
'bucket': bucket,
|
|
'prefix': prefix,
|
|
'days': days,
|
|
'rules': len(existing.rules), # type: ignore[union-attr]
|
|
}
|
|
|
|
rules = [rule for rule in (existing.rules if existing else []) if rule.rule_id != rule_id]
|
|
rules.append(
|
|
Rule(
|
|
ENABLED,
|
|
rule_filter=Filter(prefix=prefix),
|
|
rule_id=rule_id,
|
|
expiration=Expiration(days=days),
|
|
)
|
|
)
|
|
client.set_bucket_lifecycle(bucket, LifecycleConfig(rules))
|
|
self.info(f'import retention rule ensured on {bucket} ({days} days)', metadata)
|
|
return {'bucket': bucket, 'prefix': prefix, 'days': days, 'rules': len(rules)}
|