diff --git a/scripts/import_model_manual_run.py b/scripts/import_model_manual_run.py new file mode 100644 index 0000000..b260313 --- /dev/null +++ b/scripts/import_model_manual_run.py @@ -0,0 +1,502 @@ +"""Run one `.sientia` model import by hand, standing in for the frontend that does not exist yet. + +**What this is.** The Angular screen plus the Spring Boot BFF will, between them, pick a `.sientia` +file, hash it, wrap the bundle password for the worker, upload the object and insert one `PENDING` +row into `public.experiment_run` before starting `import_model`. None of that exists yet. This +script does exactly those four things, against whatever environment `.env` points at, so the import +path can be exercised for real and so whoever writes the frontend has a working reference for the +wire shapes (the digest, the envelope, the object key, the row). + +**What it writes.** Two things, and it can undo both: one object under +`{IMPORT_BUNDLE_BUCKET}/{IMPORT_BUNDLE_PREFIX}` in MinIO, and one row in `public.experiment_run` in +the BFF Postgres database. Everything else it does is a read. If the workflow runs, *it* creates the +MLflow experiment, run and registered version and the MongoDB model document — cell 10 can remove +those too, behind its own flag. + +**The safety rule.** Every target comes from the environment, and this file contains no host, user, +password or key literal. A `.env` inherited from a cluster deployment can perfectly well name a +production database and a production bucket, so cell 5 and cell 6 — the only cells that write — are +gated on cell 2, where the operator has to restate the Postgres host that cell 1 printed. This is a +development tool. It is not for production, and the guard is a mitigation, not a guarantee. + +**Cells.** Run them one at a time in the VS Code interactive window, or run the module top to +bottom. Cell order is the contract: + + 1 resolve every target and print it 6 insert the PENDING row -> import_run_id + 2 confirmation guard (first write is after) 7 print the workflow input + 3 pick the bundle, compute expected_digest 8 optionally start the workflow and wait + 4 build the password envelope 9 read the outcome back + 5 upload the object, verify the stored bytes 10 undo + +**Boundaries.** This script imports nothing from `laborious.activities` or `laborious.workflows` — +it stands in for the frontend, so it must not borrow the consumer's code. The one exception is +`laborious.utils.bundle.format`, which is pure `struct`/`hashlib` over bytes and lets cell 3 print +the header of the file it is about to upload. The envelope is built here in four visible lines +rather than imported, because those four lines are what the frontend has to transliterate to +WebCrypto; `laborious/utils/import_password.py` is the decrypt counterpart. +""" + +# %% +# Cell 1 — resolve every target from the environment and print what we are pointed at. +# Reads only. No secret is printed here or anywhere else in this script. +from __future__ import annotations + +import os +import sys +from pathlib import Path +from typing import Any + +from dotenv import load_dotenv + +REPO_ROOT = Path(__file__).resolve().parent.parent +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +load_dotenv(REPO_ROOT / '.env') + + +def env(name: str, default: str = '') -> str: + """The environment variable, stripped, or `default`.""" + return os.getenv(name, default).strip() + + +def env_required(name: str) -> str: + """The environment variable, or a `RuntimeError` naming it.""" + value = env(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}' + + +# Every connection target below is *required* rather than defaulted. The worker's own config falls +# back to localhost, which is right for a service that must boot; it is wrong for a script that +# writes, because a silent fallback is a write to somewhere nobody read out loud. +# +# The import log database: the POSTGRES_* server from .env, with only the database name of its own, +# because the log lives in the Spring Boot BFF's database rather than in `sientia`. The schema and +# table are constants of the worker's code, never configuration. +# +# `build_import_status_config` also honours IMPORT_STATUS_DB_{HOST,PORT,USER,PASSWORD} as overrides, +# for a deployment whose log sits on another server. This script does not: it is a development tool +# pointed at a real environment, and the override exists for the e2e suite's two containers. +STATUS_DB_HOST = env_required('POSTGRES_HOST') +STATUS_DB_PORT = int(env_required('POSTGRES_PORT')) +STATUS_DB_NAME = env_required('IMPORT_STATUS_DB_NAME') +STATUS_DB_USER = env_required('POSTGRES_USER') +STATUS_DB_PASSWORD = os.getenv('POSTGRES_PASSWORD', '') +STATUS_TABLE = 'public.experiment_run' + +# Where the frontend uploads the bundle. The prefix keeps the worker's own default, which is a +# format decision rather than a target. +MINIO_ENDPOINT_URL = env_required('MINIO_ENDPOINT_URL') +MINIO_ACCESS_KEY = env_required('MINIO_ACCESS_KEY') +MINIO_SECRET_KEY = env_required('MINIO_SECRET_KEY') +BUNDLE_BUCKET = env_required('IMPORT_BUNDLE_BUCKET') +_raw_prefix = env('IMPORT_BUNDLE_PREFIX', 'imported_models/') +BUNDLE_PREFIX = _raw_prefix if _raw_prefix.endswith('/') else f'{_raw_prefix}/' + +# Read-back targets. This script never writes to these; the workflow does. +MONGODB_URL = env_required('MONGODB_URL') +MONGODB_USERNAME = env('MONGODB_USERNAME') +MONGODB_PASSWORD = os.getenv('MONGODB_PASSWORD', '') +MONGODB_DATABASE = env_required('MONGODB_DATABASE') +MODELS_COLLECTION = env('IMPORT_MODELS_COLLECTION', 'models') + +MLFLOW_HOST = env_required('MLFLOW_HOST') +MLFLOW_PORT = env_required('MLFLOW_PORT') +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' + +# This script's own inputs. +BUNDLE_PATH = Path( + env('IMPORT_SCRIPT_BUNDLE_PATH', 'tests/fixtures/bundle/golden_model_v3.sientia') +) +if not BUNDLE_PATH.is_absolute(): + BUNDLE_PATH = REPO_ROOT / BUNDLE_PATH +IMPORT_USERNAME = env_required('IMPORT_SCRIPT_USERNAME') + +print('resolved targets — nothing has been written yet') +print(_line('status postgres', f'{STATUS_DB_HOST}:{STATUS_DB_PORT}/{STATUS_DB_NAME}')) +print(_line('status user', STATUS_DB_USER)) +print(_line('status table', STATUS_TABLE)) +print(_line('minio endpoint', MINIO_ENDPOINT_URL)) +print(_line('minio bucket/prefix', f'{BUNDLE_BUCKET}/{BUNDLE_PREFIX}')) +print(_line('mongo', f'{MONGODB_URL}/{MONGODB_DATABASE}.{MODELS_COLLECTION}')) +print(_line('mlflow', MLFLOW_TRACKING_URI)) +print(_line('temporal', f'{TEMPORAL_HOST} ns={TEMPORAL_NAMESPACE}')) +print(_line('task queue', IMPORT_TASK_QUEUE)) +print(_line('bundle', BUNDLE_PATH)) +print(_line('username', IMPORT_USERNAME)) +print() +print(f'to continue, set the cell 2 confirmation to: {STATUS_DB_HOST}') + +# %% +# Cell 2 — the confirmation guard. Cells 5 and 6 write; this is the last cell before them. +# +# 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') + +if CONFIRM_TARGET != STATUS_DB_HOST: + raise RuntimeError( + 'refusing to write: the confirmation does not match the resolved status database host. ' + f'confirmation={CONFIRM_TARGET!r}, resolved POSTGRES_HOST={STATUS_DB_HOST!r}. ' + 'Read the targets cell 1 printed, then restate the host here.' + ) + +print(f'target confirmed: {STATUS_DB_HOST} — the write cells may run') + +# %% +# Cell 3 — pick the bundle, compute the digest the frontend must send, print its header. +# +# `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() +object_name = BUNDLE_PATH.name +object_key = f'{BUNDLE_PREFIX}{object_name}' + +header = parse_header(bundle_bytes[:HEADER_SIZE]) + +print(_line('file', BUNDLE_PATH)) +print(_line('size', f'{len(bundle_bytes)} bytes')) +print(_line('expected_digest', expected_digest)) +print(_line('object key', f'{BUNDLE_BUCKET}/{object_key}')) +print(_line('format version', header.format_version)) +print(_line('aead / kdf id', f'{header.aead_id} / {header.kdf_id}')) +print(_line('kdf memlimit', f'{header.kdf_memlimit_bytes} bytes')) +print(_line('kdf opslimit', header.kdf_opslimit)) +print(_line('kdf parallelism', header.kdf_parallelism)) +print(_line('salt', header.salt.hex())) +print(_line('ciphertext digest', header.ciphertext_digest.hex())) +print(_line('chunk stream from', f'offset {DIGEST_OFFSET}')) + +# %% +# Cell 4 — build the password envelope, exactly as the frontend must build it. +# +# Wire shape: base64( nonce(12) || ciphertext || GCM tag(16) ) under AES-256-GCM with the 32-byte +# key that `IMPORT_PASSWORD_KEY` carries base64-encoded. The four lines below are the whole of it, +# written out rather than imported so the Angular side can transliterate them to WebCrypto: +# +# const k = await crypto.subtle.importKey('raw', keyBytes, 'AES-GCM', false, ['encrypt']); +# const nonce = crypto.getRandomValues(new Uint8Array(12)); +# const sealed = await crypto.subtle.encrypt({name: 'AES-GCM', iv: nonce}, k, passwordBytes); +# const envelope = base64(concat(nonce, new Uint8Array(sealed))); +# +# WebCrypto appends the GCM tag to the ciphertext, and so does `cryptography`'s AESGCM, so the two +# produce byte-identical payloads. The decrypt counterpart the worker runs is +# `laborious/utils/import_password.py` (`decrypt_password_envelope`). +# +# 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 +) +password_envelope = base64.b64encode(nonce + sealed).decode('ascii') + +print(_line('envelope key', f'{len(envelope_key)} bytes (not printed)')) +print(_line('envelope', f'{len(password_envelope)} base64 characters (not printed)')) + +# %% +# Cell 5 — FIRST WRITE. Upload the object, then prove the stored bytes are the ones we hashed. +# +# 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( + _endpoint.netloc or MINIO_ENDPOINT_URL, + access_key=MINIO_ACCESS_KEY, + secret_key=MINIO_SECRET_KEY, + secure=_endpoint.scheme == 'https', +) + +minio_client.put_object( + BUNDLE_BUCKET, + object_key, + io.BytesIO(bundle_bytes), + length=len(bundle_bytes), + content_type='application/octet-stream', +) + +stat = minio_client.stat_object(BUNDLE_BUCKET, object_key) +_response = minio_client.get_object(BUNDLE_BUCKET, object_key) +try: + stored_bytes = _response.read() +finally: + _response.close() + _response.release_conn() + +stored_digest = hashlib.sha256(stored_bytes).hexdigest() +if stored_digest != expected_digest: + raise RuntimeError( + f'the stored object does not hash to the digest we computed: stored {stored_digest}, ' + f'expected {expected_digest} — do not insert a row for this upload' + ) + +print(_line('uploaded', f'{BUNDLE_BUCKET}/{object_key}')) +print(_line('reported size', f'{stat.size} bytes')) +print(_line('etag', stat.etag)) +print(_line('re-read digest', f'{stored_digest} (matches expected_digest)')) + +# %% +# Cell 6 — SECOND WRITE. Insert the import log row: the BFF's half of the contract. +# +# `URL.create` rather than a URL string, so a password containing `@ : /` survives. `experiment_name` +# is a placeholder: the column is NOT NULL and the frontend cannot see inside an encrypted bundle, +# 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( + 'postgresql', + username=STATUS_DB_USER, + password=STATUS_DB_PASSWORD, + host=STATUS_DB_HOST, + port=STATUS_DB_PORT, + database=STATUS_DB_NAME, + ), + pool_pre_ping=True, +) + +request_data = { + 'expected_digest': expected_digest, + 'bucket': BUNDLE_BUCKET, + 'object_key': object_key, + 'file_name': object_name, + 'source': 'scripts/import_model_manual_run.py', +} + +_now = datetime.now(UTC).replace(tzinfo=None) +_placeholder_experiment = f'import-{Path(object_name).stem}'[:50] + +INSERT_ROW_SQL = text( + 'INSERT INTO public.experiment_run ' + '(experiment_name, username, status, created_at, updated_at, ' + ' bucket_name, file_name, request_data, run_type) ' + 'VALUES (:experiment_name, :username, :status, :created_at, :updated_at, ' + ' :bucket_name, :file_name, CAST(:request_data AS JSONB), :run_type) ' + 'RETURNING id' +) + +with status_engine.begin() as connection: + import_run_id = connection.execute( + INSERT_ROW_SQL, + { + 'experiment_name': _placeholder_experiment, + 'username': IMPORT_USERNAME, + 'status': 'PENDING', + 'created_at': _now, + 'updated_at': _now, + 'bucket_name': BUNDLE_BUCKET, + 'file_name': object_key, + 'request_data': json.dumps(request_data), + 'run_type': 'IMPORT', + }, + ).scalar_one() + +print(_line('import_run_id', import_run_id)) +print(_line('status / run_type', 'PENDING / IMPORT')) +print(_line('experiment_name', f'{_placeholder_experiment} (the worker overwrites this)')) +print(_line('file_name', object_key)) + +# %% +# Cell 7 — the workflow input, ready to paste into the Temporal UI. +# +# These five keys are the whole contract. `bucket` is optional — the worker falls back to +# IMPORT_BUNDLE_BUCKET — and is sent explicitly so the printed input is self-describing. No model +# name is sent: the worker takes it from the bundle's own metadata, verbatim. +workflow_input = { + 'import_run_id': import_run_id, + 'bucket': BUNDLE_BUCKET, + 'object_key': object_key, + 'expected_digest': expected_digest, + 'password_envelope': password_envelope, +} + +print(f'workflow: import_model task queue: {IMPORT_TASK_QUEUE}') +print(json.dumps(workflow_input, indent=2)) + +# %% +# Cell 8 — optionally start the workflow and wait for it. +# +# Skipping this cell is normal: the operator may prefer to start the run from the Temporal UI with +# the JSON cell 7 printed. Cell 9 works either way, because it looks the outcome up by row id. +# +# 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 + +if START_THE_WORKFLOW: + import asyncio + from concurrent.futures import ThreadPoolExecutor + + from temporalio.client import Client + + async def start_and_wait() -> Any: + """Start `import_model` on the import queue and return its result.""" + client = await Client.connect(TEMPORAL_HOST, namespace=TEMPORAL_NAMESPACE) + handle = await client.start_workflow( + 'import_model', + workflow_input, + id=f'manual-import-{import_run_id}', + task_queue=IMPORT_TASK_QUEUE, + ) + print(_line('workflow id', handle.id)) + print(_line('run id', handle.result_run_id)) + return await handle.result() + + with ThreadPoolExecutor(max_workers=1) as _pool: + workflow_result = _pool.submit(asyncio.run, start_and_wait()).result() + print(json.dumps(workflow_result, indent=2, default=str)) +else: + print('not started — paste the input from cell 7 into the Temporal UI, or set') + print(f'START_THE_WORKFLOW = True and re-run this cell. Task queue: {IMPORT_TASK_QUEUE}') + +# %% +# Cell 9 — read the outcome back from all three places it lands. +# +# The row is authoritative for the verdict; MLflow and MongoDB are what the verdict is about. Each +# block is labelled, so the output can be pasted into a ticket as it is. +print('--- import log row ---') +READ_ROW_SQL = text( + 'SELECT id, experiment_name, run_name, status, error_message, run_type, ' + ' bucket_name, file_name, created_at, updated_at, ' + ' request_data, orchestrator_response_data ' + 'FROM public.experiment_run WHERE id = :id' +) + +with status_engine.begin() as connection: + _result = connection.execute(READ_ROW_SQL, {'id': import_run_id}) + row = _result.mappings().one_or_none() + +if row is None: + print(f' no row with id {import_run_id}') +else: + for _column, _value in row.items(): + print(_line(_column, _value)) + +print() +print('--- mlflow ---') +import mlflow + +os.environ['MLFLOW_TRACKING_USERNAME'] = env('MLFLOW_USERNAME') +os.environ['MLFLOW_TRACKING_PASSWORD'] = os.getenv('MLFLOW_PASSWORD', '') +mlflow.set_tracking_uri(MLFLOW_TRACKING_URI) +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 +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}')) + 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: + try: + for _version in mlflow_client.search_model_versions(f"name='{_registered}'"): + print( + _line( + 'version', + f'{_version.name} v{_version.version} stage={_version.current_stage} ' + f'run={_version.run_id}', + ) + ) + print(_line(' source', _version.source)) + 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') + +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}') +_collection = mongo_client[MONGODB_DATABASE][MODELS_COLLECTION] +for _document in _collection.find({}, {'_id': False}).sort('id', 1): + print(_line('document', _document)) + +# %% +# Cell 10 — undo. +# +# 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. +DELETE_PROVISIONING = False + +try: + minio_client.remove_object(BUNDLE_BUCKET, object_key) + print(_line('deleted object', f'{BUNDLE_BUCKET}/{object_key}')) +except Exception as error: + # Undo reports what it could not remove and carries on to the row. + print(f' could not delete the object: {type(error).__name__}: {error}') + +with status_engine.begin() as connection: + _deleted = connection.execute( + text('DELETE FROM public.experiment_run WHERE id = :id'), + {'id': import_run_id}, + ).rowcount +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: + 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)) +else: + print(' provisioning kept — set DELETE_PROVISIONING = True to remove it too') + +status_engine.dispose() +mongo_client.close() diff --git a/scripts/run_coverage.sh b/scripts/run_coverage.sh new file mode 100755 index 0000000..f9af4cb --- /dev/null +++ b/scripts/run_coverage.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +# Exit on any error +set -e + +echo "Activating virtual environment..." +source ./venv/bin/activate + +pytest --cov=laborious --cov-report=html + +xdg-open htmlcov/index.html \ No newline at end of file diff --git a/scripts/run_local.sh b/scripts/run_local.sh new file mode 100755 index 0000000..2bbd5c2 --- /dev/null +++ b/scripts/run_local.sh @@ -0,0 +1,18 @@ +#!/bin/bash + +# Exit on any error +set -e + +echo "Activating virtual environment..." +source ./venv/bin/activate + +echo "Loading environment variables from .env..." +if [ -f .env ]; then + export $(cat .env | grep -v '^#' | xargs) + echo "Environment variables loaded from .env" +else + echo "Warning: .env file not found. Continuing without environment variables." +fi + +echo "Starting ingestor application..." +python -m laborious.worker.worker