gitea_admin/sientia-dataops-laborious-legacy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

QTZPOC-20

Browse Source 
QTZPOC-20: Improve structured logging and reporting for all model import pipeline failures by introducing a dedicated activity and centralizing error classification. This ensures consistent observability, with all workflow failures now reported before cleanup. Also, adjusted import step timeouts and refined manual run script cleanup logic.
This commit is contained in:
Vitor Pimentel dos Santos
2026-08-27 14:44:57 -03:00
parent e95e65b72b
commit 853c83e662
18 changed files with 857 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

150
scripts/import_model_manual_run.py
View File

@@ -41,12 +41,25 @@ WebCrypto; `laborious/utils/import_password.py` is the decrypt counterpart.
# Reads only. No secret is printed here or anywhere else in this script.
from __future__ import annotations
import base64
import hashlib
import io
import json
import os
import secrets
import sys
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from urllib.parse import urlparse
import mlflow
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from dotenv import load_dotenv
from minio import Minio
from pymongo import MongoClient
from sqlalchemy import create_engine, text
from sqlalchemy.engine import URL
REPO_ROOT = Path(__file__).resolve().parent.parent
if str(REPO_ROOT) not in sys.path:
@@ -54,6 +67,11 @@ if str(REPO_ROOT) not in sys.path:
load_dotenv(REPO_ROOT / '.env')
# Below the `sys.path` insert on purpose: `laborious` is only importable once the repository root is
# on the path, which is what those three lines above just arranged. Hoisting this to the top of the
# cell is a `ModuleNotFoundError` when the script runs from anywhere but the repository root.
from laborious.utils.bundle.format import DIGEST_OFFSET, HEADER_SIZE, parse_header # noqa: E402
def env(name: str, default: str = '') -> str:
"""The environment variable, stripped, or `default`."""
@@ -68,6 +86,21 @@ def env_required(name: str) -> str:
return value
def env_secret_required(name: str) -> str:
"""The environment variable *verbatim*, or a `RuntimeError` naming it.
Unlike `env_required` this does not strip, and the difference is not cosmetic. A key and a
password are byte strings whose surrounding whitespace is part of the secret. The worker reads
the same variables with a plain `os.getenv`, so stripping here would seal the envelope under
bytes the worker never derives, and the import would fail at a gate with nothing in the message
to say why.
"""
value = os.getenv(name, '')
if not value:
raise RuntimeError(f'{name} is not set — see .env.example')
return value
def _line(label: str, value: object) -> str:
"""One aligned `label value` line, for output that stays readable in a terminal."""
return f' {label:<24} {value}'
@@ -114,7 +147,7 @@ MLFLOW_TRACKING_URI = f'{MLFLOW_HOST}:{MLFLOW_PORT}'
TEMPORAL_HOST = env_required('TEMPORAL_HOST')
TEMPORAL_NAMESPACE = env_required('TEMPORAL_NAMESPACE')
RUNTIME = env_required('RUNTIME')
IMPORT_TASK_QUEUE = f'import_model-{RUNTIME}-queue'
IMPORT_TASK_QUEUE = 'import_model-legacy-queue'
# This script's own inputs.
BUNDLE_PATH = Path(
@@ -145,7 +178,7 @@ print(f'to continue, set the cell 2 confirmation to: {STATUS_DB_HOST}')
# Restate the Postgres host cell 1 printed. Either set IMPORT_SCRIPT_CONFIRM_TARGET in .env, or
# replace the right-hand side below with the value cell 1 told you. It is a read-then-restate on
# purpose: a boolean flag can be flipped without ever looking at what it is pointed at.
CONFIRM_TARGET = env('IMPORT_SCRIPT_CONFIRM_TARGET')
CONFIRM_TARGET = 'localhost'
if CONFIRM_TARGET != STATUS_DB_HOST:
raise RuntimeError(
@@ -162,9 +195,7 @@ print(f'target confirmed: {STATUS_DB_HOST} — the write cells may run')
# `expected_digest` is the SHA-256 of the exact bytes that get uploaded, computed the way the
# frontend computes it: over the file, not over anything derived from it. The header fields printed
# here are plaintext at the front of the file — reading them decrypts nothing.
import hashlib
from laborious.utils.bundle.format import DIGEST_OFFSET, HEADER_SIZE, parse_header
bundle_bytes = BUNDLE_PATH.read_bytes()
expected_digest = hashlib.sha256(bundle_bytes).hexdigest()
@@ -202,16 +233,22 @@ print(_line('chunk stream from', f'offset {DIGEST_OFFSET}'))
# produce byte-identical payloads. The decrypt counterpart the worker runs is
# `laborious/utils/import_password.py` (`decrypt_password_envelope`).
#
# `IMPORT_PASSWORD_KEY` must be base64 of exactly 32 bytes — `openssl rand -base64 32`, or the
# value the deployment's Kubernetes Secret carries. Anything else fails here, which is the point:
# the worker resolves the same variable the same way, so a value this cell cannot decode is a value
# the worker could not have decoded either.
#
# The bundle password is read verbatim — `env_secret_required`, not `env_required` — because a
# trailing space is part of the password, and stripping it here would seal an envelope whose
# plaintext the bundle reader rejects at gate 3.
#
# Neither the key nor the password nor the plaintext is printed — only the envelope's length.
import base64
import secrets
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
envelope_key = base64.b64decode(env_required('IMPORT_PASSWORD_KEY'), validate=True)
nonce = secrets.token_bytes(12)
sealed = AESGCM(envelope_key).encrypt(
nonce, env_required('IMPORT_SCRIPT_BUNDLE_PASSWORD').encode('utf-8'), None
nonce, env_secret_required('IMPORT_SCRIPT_BUNDLE_PASSWORD').encode('utf-8'), None
)
password_envelope = base64.b64encode(nonce + sealed).decode('ascii')
@@ -224,10 +261,7 @@ print(_line('envelope', f'{len(password_envelope)} base64 characters (not printe
# The re-read is the point: `expected_digest` is a promise about what is in the bucket, and the
# worker's gate 1 holds the upload to it. If the digest of the stored object differs, the frontend
# contract is already broken and there is no reason to insert a row.
import io
from urllib.parse import urlparse
from minio import Minio
_endpoint = urlparse(MINIO_ENDPOINT_URL)
minio_client = Minio(
@@ -273,11 +307,7 @@ print(_line('re-read digest', f'{stored_digest} (matches expected_digest)'))
# so the worker overwrites it with the bundle's own experiment name as soon as the bundle is open
# (`record_import_names`). `file_name` is the object key, because the worker compares the two and
# refuses the import at the `received` step when they disagree.
import json
from datetime import UTC, datetime
from sqlalchemy import create_engine, text
from sqlalchemy.engine import URL
status_engine = create_engine(
URL.create(
@@ -358,7 +388,7 @@ print(json.dumps(workflow_input, indent=2))
# The coroutine runs on a loop of its own in a worker thread rather than through `asyncio.run`, so
# this cell behaves the same in the interactive window (which already has a running loop) as it
# does when the module is executed top to bottom (which does not).
START_THE_WORKFLOW = False
START_THE_WORKFLOW = True
if START_THE_WORKFLOW:
import asyncio
@@ -411,7 +441,6 @@ else:
print()
print('--- mlflow ---')
import mlflow
os.environ['MLFLOW_TRACKING_USERNAME'] = env('MLFLOW_USERNAME')
os.environ['MLFLOW_TRACKING_PASSWORD'] = os.getenv('MLFLOW_PASSWORD', '')
@@ -420,22 +449,44 @@ mlflow_client = mlflow.tracking.MlflowClient()
_experiment_name = row['experiment_name'] if row is not None else ''
_experiment = mlflow_client.get_experiment_by_name(_experiment_name) if _experiment_name else None
# The experiment is get-or-create, so it can perfectly well be the model's own experiment with
# months of training runs in it. `_run_ids` is everything in it — printed, and used to decide
# whether the experiment is this import's alone — while `_import_run_ids` is only the run this
# import created, matched on the run name the record carries. The undo acts on the second list.
_row_run_name = (row['run_name'] if row is not None else '') or ''
_run_ids: list[str] = []
_import_run_ids: list[str] = []
if _experiment is None:
print(f' no experiment named {_experiment_name!r}')
else:
print(_line('experiment', f'{_experiment.name} ({_experiment.experiment_id})'))
for _run in mlflow_client.search_runs([_experiment.experiment_id]):
print(_line('run', f'{_run.info.run_name} ({_run.info.run_id}) {_run.info.status}'))
_mine = bool(_row_run_name) and _run.info.run_name == _row_run_name
print(
_line(
'run',
f'{_run.info.run_name} ({_run.info.run_id}) {_run.info.status}'
f'{" <- this import" if _mine else ""}',
)
)
_run_ids.append(_run.info.run_id)
if _mine:
_import_run_ids.append(_run.info.run_id)
for _artifact in mlflow_client.list_artifacts(_run.info.run_id):
print(_line(' artifact', _artifact.path))
# The run name is `import-{model_name}-v{version}` (`derive_run_name`), so the registered model
# name is what sits between the prefix and the version suffix.
_run_name = (row['run_name'] if row is not None else '') or ''
_registered = _run_name.removeprefix('import-').rsplit('-v', 1)[0] if _run_name else ''
if _registered:
# The registered model is looked up by `run_id`, not by a name reconstructed from `run_name`. The
# run name is `import-{model_name}-v{model_version}` truncated to 50 characters (`derive_run_name`),
# and a 32-character `model_version` leaves ten characters for the name — so `Quartz-demo` comes
# back as `Quartz-dem` and the lookup finds nothing. The version's `run_id` is exact, and it is what
# ties the registered model to the run this import created whatever the name was truncated to.
_registered_names: set[str] = set()
_registered_versions: list[tuple[str, str]] = []
for _run_id in _import_run_ids:
try:
for _version in mlflow_client.search_model_versions(f"name='{_registered}'"):
for _version in mlflow_client.search_model_versions(f"run_id='{_run_id}'"):
_registered_names.add(_version.name)
_registered_versions.append((_version.name, str(_version.version)))
print(
_line(
'version',
@@ -447,12 +498,13 @@ if _registered:
except Exception as error:
# A read-back cell reports what it could not read; it does not fail the inspection.
print(f' registry read failed: {type(error).__name__}: {error}')
else:
print(' no run_name on the row yet, so no registered model to look up')
if not _import_run_ids:
print(' no run of this import in the experiment yet, so no registered model to look up')
elif not _registered_names:
print(" no registered version points at this import's run")
print()
print('--- mongodb model document ---')
from pymongo import MongoClient
_credentials = f'{MONGODB_USERNAME}:{MONGODB_PASSWORD}@' if MONGODB_USERNAME else ''
mongo_client: MongoClient = MongoClient(f'mongodb://{_credentials}{MONGODB_URL}')
@@ -466,6 +518,13 @@ for _document in _collection.find({}, {'_id': False}).sort('id', 1):
# The object and the row always: they are what this script created, and leaving them behind leaves a
# PENDING import nobody will ever run. What MLflow and MongoDB hold is left alone unless
# DELETE_PROVISIONING is set, because after a successful run that is the thing you wanted to look at.
#
# When it *is* set, the undo removes only what this import added. An import of a model the platform
# already had lands as a new *version* of that registered model and writes no document at all, so
# deleting the registered model by name would take versions trained months ago with it. The version
# this import created is deleted by number; the registered model and the document go only when
# nothing else is left pointing at that name — which is exactly the case where this import was the
# first one.
DELETE_PROVISIONING = False
try:
@@ -483,18 +542,43 @@ with status_engine.begin() as connection:
print(_line('deleted rows', _deleted))
if DELETE_PROVISIONING:
if _experiment is not None:
mlflow_client.delete_experiment(_experiment.experiment_id)
print(_line('deleted experiment', _experiment.experiment_id))
if _registered:
for _name, _number in _registered_versions:
try:
mlflow_client.delete_model_version(_name, _number)
print(_line('deleted version', f'{_name} v{_number}'))
except Exception as error:
# Same rule as above: report and carry on to the next one.
print(f' could not delete the model version: {type(error).__name__}: {error}')
for _registered in sorted(_registered_names):
try:
_remaining = list(mlflow_client.search_model_versions(f"name='{_registered}'"))
except Exception as error:
print(f' could not count the remaining versions: {type(error).__name__}: {error}')
continue
if _remaining:
# Versions this import did not create: the model was on the platform before it ran, so
# the registered model and its document are somebody else's and stay exactly as they are.
print(_line('kept model', f'{_registered} ({len(_remaining)} other versions)'))
continue
try:
mlflow_client.delete_registered_model(_registered)
print(_line('deleted model', _registered))
except Exception as error:
# Same rule as above: report and carry on to the document.
print(f' could not delete the registered model: {type(error).__name__}: {error}')
_removed = _collection.delete_many({'name': _registered}).deleted_count
print(_line('deleted documents', _removed))
print(_line('deleted documents', f'{_removed} named {_registered}'))
# Last, and only when the experiment holds nothing but this import's own run: the experiment is
# get-or-create, so an import into a model the platform already had reuses the experiment its
# training runs live in, and deleting it would take them all.
if _experiment is not None:
_others = len(_run_ids) - len(_import_run_ids)
if _others:
print(_line('kept experiment', f'{_experiment.name} ({_others} other runs)'))
else:
mlflow_client.delete_experiment(_experiment.experiment_id)
print(_line('deleted experiment', _experiment.experiment_id))
else:
print(' provisioning kept — set DELETE_PROVISIONING = True to remove it too')