Compare commits
73 Commits
fcc8920a8b
...
0.9.1
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
47d10f6a0e | ||
| f027e98560 | |||
| 9e27688a70 | |||
| d91a0d069c | |||
| fdeafe9768 | |||
| 25f11fa943 | |||
| 8ef7ed800a | |||
|
|
ac67d5b447 | ||
|
|
ce91a0bbb6 | ||
|
|
959c78170c | ||
|
|
5690e792a9 | ||
|
|
76351fdf8a | ||
|
|
070956def9 | ||
| 85fe4663b1 | |||
| e2fbf841fc | |||
| b65cf2ba68 | |||
| 1c32e9d678 | |||
| bf48a18c7b | |||
| 6c0d178f99 | |||
| 4cccc45c34 | |||
| 512a49f78f | |||
| bda83e7943 | |||
| e600c7922e | |||
| 9e0da03717 | |||
| 7c45dbb407 | |||
| 4c7cae502f | |||
| d06ce0daae | |||
| 37544edafe | |||
| abfd474e4d | |||
| eb89917c5a | |||
| 192834bde3 | |||
| e4fc12e529 | |||
| 61db575cdc | |||
| 72a55c0900 | |||
| d257e70465 | |||
| 7b2e90a9f9 | |||
| bde614e2a6 | |||
| 6a412147ee | |||
| 220ee5509c | |||
| 09027bc93b | |||
| 086b65c45f | |||
| 30ff378cb7 | |||
| c1a44a9d7a | |||
| f4c110cd8a | |||
| dfb9deae1a | |||
| 853c83e662 | |||
| e95e65b72b | |||
| dac4f1718f | |||
| 4e3f26f1c1 | |||
| 83b025ff9f | |||
| 0e445de386 | |||
| 4576502d78 | |||
| 291d16efa5 | |||
| 03ac4b31e5 | |||
| d7bf36605d | |||
|
|
695e4c07a6 | ||
| c12f2996c2 | |||
| 6a8c41b328 | |||
| 773c980fc3 | |||
| 35efb67c89 | |||
| acd925cf2a | |||
| 310ceea0d8 | |||
|
|
2d73ec9ec2 | ||
| 2142143ab9 | |||
| c07457bfbe | |||
| 086b12492e | |||
| 4ea0754f0c | |||
| ddb1618209 | |||
| 90f8bdda61 | |||
| 2ccda3e440 | |||
| d856150e24 | |||
| 6569810756 | |||
| 7153f1da0d |
40
.env.example
40
.env.example
@@ -32,4 +32,42 @@ MINIO_ENDPOINT_URL="http://localhost:9000"
|
||||
MINIO_ACCESS_KEY="sientia"
|
||||
MINIO_SECRET_KEY="sientia"
|
||||
MINIO_REGION_NAME="sa-east-1"
|
||||
MINIO_DEFAULT_BUCKET="sientia"
|
||||
MINIO_DEFAULT_BUCKET="sientia"
|
||||
|
||||
# Suffix for every Temporal task queue name; required, non-empty.
|
||||
# The import worker listens on "import_model-${RUNTIME}-queue".
|
||||
RUNTIME="local"
|
||||
|
||||
# --- Model import (.sientia bundles) -----------------------------------------
|
||||
# The import worker is registered only when IMPORT_MODEL_ENABLED is truthy
|
||||
# (1/true/yes/on). Do not turn it on before the V12 migration is applied to
|
||||
# public.experiment_run in the BFF database.
|
||||
IMPORT_MODEL_ENABLED="false"
|
||||
|
||||
# base64 of the 32-byte AES-256-GCM key that wraps bundle passwords.
|
||||
# Generate a throwaway one with: openssl rand -base64 32
|
||||
IMPORT_PASSWORD_KEY="replace-with-base64-of-32-random-bytes"
|
||||
|
||||
# Where the frontend uploads the bundle. The object key must sit under the prefix.
|
||||
IMPORT_BUNDLE_BUCKET="sientia"
|
||||
IMPORT_BUNDLE_PREFIX="imported_models/"
|
||||
|
||||
# MongoDB collection holding the model listing the import writes its document into.
|
||||
IMPORT_MODELS_COLLECTION="models"
|
||||
|
||||
# The import log lives in the Spring Boot BFF database on the POSTGRES_* server
|
||||
# above: same host, port, user and password, another database name. Only that
|
||||
# name is configuration here; the schema (public) and the table (experiment_run)
|
||||
# are constants of the code. Set IMPORT_STATUS_DB_{HOST,PORT,USER,PASSWORD} only
|
||||
# if the log really is on a different server than POSTGRES_* — they override it.
|
||||
IMPORT_STATUS_DB_NAME="sientia-core-mlops-bff"
|
||||
|
||||
# --- scripts/import_model_manual_run.py only ---------------------------------
|
||||
# Inputs for the manual run that stands in for the frontend. The script writes
|
||||
# one object and one row to whatever the variables above name.
|
||||
IMPORT_SCRIPT_BUNDLE_PATH="tests/fixtures/bundle/golden_model_v3.sientia"
|
||||
IMPORT_SCRIPT_BUNDLE_PASSWORD="the-password-the-bundle-was-written-with"
|
||||
IMPORT_SCRIPT_USERNAME="manual.run@example.com"
|
||||
# Cell 2's guard: restate POSTGRES_HOST here (or in the cell) after reading what cell 1
|
||||
# printed. Left empty on purpose — the script refuses to write until it matches.
|
||||
IMPORT_SCRIPT_CONFIRM_TARGET=""
|
||||
35
.github/workflows/deploy.yml
vendored
Normal file
35
.github/workflows/deploy.yml
vendored
Normal file
@@ -0,0 +1,35 @@
|
||||
name: Deploy DataOps Worker
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [closed]
|
||||
branches: [main, 'release/**', 'feature/**']
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
id-token: write
|
||||
pull-requests: read
|
||||
|
||||
jobs:
|
||||
deploy:
|
||||
name: Deploy
|
||||
if: github.event_name == 'workflow_dispatch' || github.event.pull_request.merged == true
|
||||
uses: Aignosi/github_workflow_templates/.github/workflows/reusable-deploy-dataops.yml@main
|
||||
with:
|
||||
branch_name: ${{ github.event.pull_request.head.ref }}
|
||||
project_name: 'laborious'
|
||||
chart_name: 'laborious-worker'
|
||||
# Only the "legacy" runtime, for now (card 13 correction) — laborious-worker has two
|
||||
# Deployments (basic, legacy). basic's GITHUB_BRANCH lives in the chart base (.global.env)
|
||||
# and is deliberately left untouched here. legacy's own override exists ONLY at the profile
|
||||
# level — profiles/gcp.yaml's laborious-worker.runtimes.legacy.env, not the chart base at
|
||||
# all — hence profile/profile_env_list_path below instead of env_list_path.
|
||||
profile: 'gcp'
|
||||
profile_env_list_path: '[".runtimes.legacy.env"]'
|
||||
update_version_in: '["pyproject"]'
|
||||
create_tarball: ${{ github.event_name == 'pull_request' }}
|
||||
mirror: true
|
||||
bump_chart: true
|
||||
deploy: true
|
||||
secrets: inherit
|
||||
18
.github/workflows/quality-gate.yml
vendored
18
.github/workflows/quality-gate.yml
vendored
@@ -1,15 +1,29 @@
|
||||
name: Quality gate
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- 'release/**'
|
||||
- 'feature/**'
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- 'release/**'
|
||||
- 'feature/**'
|
||||
types: [ opened, synchronize, reopened ]
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
quality-gate:
|
||||
uses: Aignosi/github_workflow_templates/.github/workflows/dataops-module-quality-gate.yml@main
|
||||
permissions: write-all
|
||||
uses: Aignosi/github_workflow_templates/.github/workflows/python-quality-gate.yml@main
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
issues: write
|
||||
with:
|
||||
project_name: 'laborious'
|
||||
repositories: 'sientia-dataops-library, sientia-mlops-library'
|
||||
|
||||
25
.github/workflows/release.yml
vendored
25
.github/workflows/release.yml
vendored
@@ -1,25 +0,0 @@
|
||||
name: Create Release on Merge to Main
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [closed]
|
||||
branches:
|
||||
- main
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
version:
|
||||
description: 'Version to release'
|
||||
required: false
|
||||
type: string
|
||||
|
||||
jobs:
|
||||
release:
|
||||
if: |
|
||||
(github.event_name == 'pull_request' && github.event.pull_request.merged == true) ||
|
||||
github.event_name == 'workflow_dispatch'
|
||||
uses: Aignosi/github_workflow_templates/.github/workflows/dataops-module-release.yml@main
|
||||
permissions: write-all
|
||||
with:
|
||||
project_name: 'laborious'
|
||||
release_version: ${{ github.event.inputs.version || '' }}
|
||||
secrets: inherit
|
||||
16
.github/workflows/require-chart-ref-label.yml
vendored
Normal file
16
.github/workflows/require-chart-ref-label.yml
vendored
Normal file
@@ -0,0 +1,16 @@
|
||||
name: Require chart-ref label
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [opened, synchronize, reopened, labeled, unlabeled]
|
||||
branches:
|
||||
- main
|
||||
- 'release/**'
|
||||
- 'feature/**'
|
||||
|
||||
permissions:
|
||||
pull-requests: write
|
||||
|
||||
jobs:
|
||||
require-chart-ref-label:
|
||||
uses: Aignosi/github_workflow_templates/.github/workflows/require-chart-ref-label.yml@main
|
||||
3
.gitignore
vendored
3
.gitignore
vendored
@@ -55,4 +55,5 @@ mlruns/
|
||||
|
||||
relatorio*
|
||||
openspec/
|
||||
.cursor/
|
||||
.cursor/
|
||||
.claude/
|
||||
7
Makefile
7
Makefile
@@ -1,7 +0,0 @@
|
||||
VERSION = 1.0.8
|
||||
name = sientia-laborious
|
||||
# ENVIRONMENT = production
|
||||
|
||||
docker-hub:
|
||||
@docker build --no-cache -t aignosi.azurecr.io/$(name):$(VERSION) .
|
||||
@docker push aignosi.azurecr.io/$(name):$(VERSION)
|
||||
239
README.md
239
README.md
@@ -21,6 +21,7 @@ A comprehensive, Temporal-based ML orchestration system for industrial data proc
|
||||
- [Minimal Retrain Workflow](#4-minimal-retrain-workflow-minimal_retrainpy)
|
||||
- [Drift Workflow](#5-drift-workflow-driftpy)
|
||||
- [Simple Metrics Workflow](#6-simple-metrics-workflow-simple_metricspy)
|
||||
- [Import Model Workflow](#7-import-model-workflow-import_modelpy)
|
||||
- [Installation & Setup](#installation--setup)
|
||||
- [Prerequisites](#prerequisites)
|
||||
- [Environment Setup](#environment-setup)
|
||||
@@ -31,6 +32,7 @@ A comprehensive, Temporal-based ML orchestration system for industrial data proc
|
||||
- [Running Tests and Coverage](#running-tests-and-coverage)
|
||||
- [Manual Test Execution](#manual-test-execution)
|
||||
- [Manual Application Execution](#manual-application-execution)
|
||||
- [Running One Model Import by Hand](#running-one-model-import-by-hand)
|
||||
- [Code Quality & Validation](#code-quality--validation)
|
||||
- [Overview](#overview)
|
||||
- [Validation Tools](#validation-tools)
|
||||
@@ -43,6 +45,7 @@ A comprehensive, Temporal-based ML orchestration system for industrial data proc
|
||||
- [Testing](#testing)
|
||||
- [Test Structure](#test-structure)
|
||||
- [Test Execution](#test-execution)
|
||||
- [End-to-End Tests](#end-to-end-tests-e2e)
|
||||
- [Monitoring and Metrics](#monitoring-and-metrics)
|
||||
- [Application Health Metrics](#application-health-metrics)
|
||||
- [Prediction Operation Metrics](#prediction-operation-metrics)
|
||||
@@ -126,10 +129,26 @@ Laborious uses a Temporal-based architecture with strong separation of concerns
|
||||
### Key Components
|
||||
|
||||
#### **Worker (`laborious/worker/worker.py`)**
|
||||
- Temporal client setup, worker lifecycle, task queues
|
||||
- Temporal client setup, four workers via `sientia_do.temporal.worker.prepare_worker`, plus a
|
||||
fifth registered only when `IMPORT_MODEL_ENABLED` is truthy (`1`, `true`, `yes`, `on`)
|
||||
- Runtime-scoped task queues: `{workflow}-{RUNTIME}-queue` for all workflows
|
||||
- Metrics server initialization, notification handler setup
|
||||
- Graceful shutdown and autoscaling-friendly behavior
|
||||
|
||||
**Breaking (schedulers):** drift and simple_metrics queues are no longer `drift-queue` /
|
||||
`simple_metrics-queue`. Use `drift-{RUNTIME}-queue` and `simple_metrics-{RUNTIME}-queue`
|
||||
matching the worker pod `RUNTIME` env (same as `predictions_batch` / `minimal_retrain`).
|
||||
|
||||
Task queues, one per worker:
|
||||
|
||||
| Queue | Worker | Registered |
|
||||
|---|---|---|
|
||||
| `predictions_batch-{RUNTIME}-queue` | PredictionsBatch (+ its two sub-workflows) | always |
|
||||
| `minimal_retrain-{RUNTIME}-queue` | MinimalRetrain | always |
|
||||
| `drift-{RUNTIME}-queue` | Drift | always |
|
||||
| `simple_metrics-{RUNTIME}-queue` | SimpleMetrics | always |
|
||||
| `import_model-{RUNTIME}-queue` | ImportModel | only when `IMPORT_MODEL_ENABLED` is truthy |
|
||||
|
||||
#### **Workflows (`laborious/workflows/`)**
|
||||
- `predictions_batch.py`: Batch prediction entry point
|
||||
- `sub_workflows/prediction_process.py`: Core prediction pipeline
|
||||
@@ -137,6 +156,8 @@ Laborious uses a Temporal-based architecture with strong separation of concerns
|
||||
- `minimal_retrain.py`: Automated model retraining and production update
|
||||
- `drift.py`: Data drift detection and monitoring
|
||||
- `simple_metrics.py`: Regression metrics calculation (RMSE, MSE, MAE, R²)
|
||||
- `import_model.py`: `.sientia` bundle import — registered model + model document, behind
|
||||
`IMPORT_MODEL_ENABLED`
|
||||
|
||||
#### **Activities (`laborious/activities/`)**
|
||||
- `gates.py`: Data quality validation, filtering, and data formatting operations
|
||||
@@ -504,6 +525,55 @@ flowchart LR
|
||||
}
|
||||
```
|
||||
|
||||
### 7. Import Model Workflow (`import_model.py`)
|
||||
|
||||
The **ImportModel** workflow turns an encrypted `.sientia` bundle already uploaded to MinIO into a
|
||||
registered MLflow model plus one document in the MongoDB `models` collection — or into a durable
|
||||
`ERROR` row carrying a sentence a person can read. It is registered **only** when
|
||||
`IMPORT_MODEL_ENABLED` is set, so the importer can live in its own runtime with its own credentials.
|
||||
|
||||
Full contract for the frontend and the BFF: [`docs/model-import.md`](docs/model-import.md).
|
||||
The file format: [`docs/sientia-bundle-format.md`](docs/sientia-bundle-format.md).
|
||||
|
||||
#### Execution Flow
|
||||
1. **Claim**: takes ownership of the `public.experiment_run` row the BFF inserted (it never inserts one)
|
||||
2. **Download**: object limits, then the mandatory `expected_digest` comparison
|
||||
3. **Open**: seven ordered gates — header, decryption, archive inspection, extraction, structure, content policy
|
||||
4. **Provision**: experiment, a new run with the recorded name, the artifact tree, a registered version (stage untouched)
|
||||
5. **List**: the MongoDB model document, `active: false`
|
||||
6. **Finish**: cleanup, then one terminal write classifying the outcome once
|
||||
|
||||
#### Architecture Diagram
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[1. claim_import_status] --> B[2. download_import_bundle] --> C[3. open_import_bundle]
|
||||
C --> D[4. create_import_experiment] --> E[5. upload_import_artifacts]
|
||||
E --> F[6. register_import_model_version] --> G[7. write_import_model_document]
|
||||
G --> H[8. cleanup_import_files] --> I[9. record_import_terminal_status]
|
||||
|
||||
A -.-> BFF[(sientia-core-mlops-bff)]
|
||||
B -.-> MinIO[(MinIO)]
|
||||
F -.-> MLflow[(MLflow)]
|
||||
G -.-> Mongo[(MongoDB)]
|
||||
I -.-> BFF
|
||||
```
|
||||
|
||||
#### Input Parameters
|
||||
```json
|
||||
{
|
||||
"import_run_id": 41,
|
||||
"object_key": "imported_models/sales_forecast_v3.sientia",
|
||||
"expected_digest": "9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08",
|
||||
"password_envelope": "Base64OfNonceCiphertextTag==",
|
||||
"bucket": "sientia"
|
||||
}
|
||||
```
|
||||
|
||||
A successful import leaves the registered version **unstaged** and the model document **inactive**:
|
||||
promoting and activating are manual operator steps (QTZPOC-21). The `V12` migration of
|
||||
`sientia-core-mlops-bff` must be applied before enabling the worker — see
|
||||
[`docs/model-import.md`](docs/model-import.md) § 7.
|
||||
|
||||
## 📋 Prerequisites
|
||||
|
||||
- Python 3.11+
|
||||
@@ -584,10 +654,10 @@ Use the provided script to run the application locally:
|
||||
|
||||
```bash
|
||||
# Make script executable (first time only)
|
||||
chmod +x run_local.sh
|
||||
chmod +x scripts/run_local.sh
|
||||
|
||||
# Run the application
|
||||
./run_local.sh
|
||||
# Run the application (from the repository root)
|
||||
./scripts/run_local.sh
|
||||
```
|
||||
|
||||
The script will:
|
||||
@@ -601,10 +671,10 @@ Use the provided script to run tests with coverage:
|
||||
|
||||
```bash
|
||||
# Make script executable (first time only)
|
||||
chmod +x run_coverage.sh
|
||||
chmod +x scripts/run_coverage.sh
|
||||
|
||||
# Run tests with coverage
|
||||
./run_coverage.sh
|
||||
# Run tests with coverage (from the repository root)
|
||||
./scripts/run_coverage.sh
|
||||
```
|
||||
|
||||
The script will:
|
||||
@@ -649,6 +719,62 @@ fi
|
||||
python -m laborious.worker.worker
|
||||
```
|
||||
|
||||
### Running One Model Import by Hand
|
||||
|
||||
[`scripts/import_model_manual_run.py`](scripts/import_model_manual_run.py) stands in for the frontend
|
||||
that does not exist yet. The Angular screen plus the Spring Boot BFF will, between them, hash a
|
||||
`.sientia` file, wrap the bundle password, upload the object and insert one `PENDING` row before
|
||||
starting `import_model`; this script does exactly those four things against whatever `.env` points
|
||||
at, then reads the outcome back and can undo itself.
|
||||
|
||||
It is a **development tool**. It writes to whatever `.env` names — a real object in a real bucket and
|
||||
a real row in a real `public.experiment_run` — and a `.env` copied from a cluster deployment can
|
||||
perfectly well name production. It is not for production use.
|
||||
|
||||
**Prerequisites**
|
||||
|
||||
- A reachable environment: the BFF Postgres, MinIO, MongoDB, MLflow and Temporal that `.env` names.
|
||||
- The `V12` migration applied to `public.experiment_run` in the BFF database. Without it the import
|
||||
provisions the model and then cannot record the outcome — see
|
||||
[`docs/model-import.md`](docs/model-import.md) § 7.
|
||||
- A worker running with `IMPORT_MODEL_ENABLED` truthy and the same `RUNTIME`, if you want cell 8 (or
|
||||
the Temporal UI) to actually execute the run.
|
||||
- The import variables from [`.env.example`](.env.example) filled in, including `IMPORT_PASSWORD_KEY`
|
||||
and the script's own `IMPORT_SCRIPT_BUNDLE_PATH`, `IMPORT_SCRIPT_BUNDLE_PASSWORD` and
|
||||
`IMPORT_SCRIPT_USERNAME`.
|
||||
|
||||
**Running it.** The file is split into `# %%` cells: run them one at a time in the VS Code
|
||||
interactive window, or run the module top to bottom.
|
||||
|
||||
```bash
|
||||
source ./venv/bin/activate
|
||||
|
||||
# top to bottom (cells 5 and 6 write; cell 2 must pass first)
|
||||
python scripts/import_model_manual_run.py
|
||||
```
|
||||
|
||||
| Cell | Does | Writes? |
|
||||
|---|---|---|
|
||||
| 1 | Resolves and prints every target — Postgres, MinIO, Mongo, MLflow, Temporal, the task queue | no |
|
||||
| 2 | **Confirmation guard** | no |
|
||||
| 3 | Picks the bundle, computes `expected_digest`, prints the parsed header | no |
|
||||
| 4 | Builds the AES-256-GCM password envelope | no |
|
||||
| 5 | Uploads the object, re-reads it and checks the digest still matches | **yes** |
|
||||
| 6 | Inserts the `PENDING` / `run_type='IMPORT'` row, prints `import_run_id` | **yes** |
|
||||
| 7 | Prints the workflow input as JSON, ready for the Temporal UI | no |
|
||||
| 8 | Optionally starts `import_model` and waits (`START_THE_WORKFLOW`) | starts a run |
|
||||
| 9 | Reads the outcome back: the row, the MLflow experiment/run/version, the model document | no |
|
||||
| 10 | Undo: deletes the object and the row always; the provisioning behind `DELETE_PROVISIONING` | deletes |
|
||||
|
||||
**The guard.** Cells 5 and 6 are the only cells that write, and neither runs until cell 2 passes.
|
||||
Cell 2 compares a confirmation value against the Postgres host cell 1 resolved: set
|
||||
`IMPORT_SCRIPT_CONFIRM_TARGET` in `.env`, or paste the host into the cell. It is a read-then-restate
|
||||
rather than a boolean on purpose — a flag can be flipped without ever looking at what it points at.
|
||||
|
||||
No credential is written in the script and none is printed: it prints hosts, ports, database, bucket
|
||||
and object names, the row id and the workflow input, and never the bundle password, the envelope key
|
||||
or the envelope's plaintext.
|
||||
|
||||
## Code Quality & Validation
|
||||
|
||||
### Overview
|
||||
@@ -756,6 +882,57 @@ pytest tests/laborious/activities/test_gates.py
|
||||
pytest tests/laborious/workflows/test_predictions_batch.py
|
||||
```
|
||||
|
||||
### End-to-End Tests (`e2e/`)
|
||||
|
||||
The `e2e/` suite runs against real dependencies in Docker containers via
|
||||
[testcontainers](https://testcontainers.com/): PostgreSQL, MinIO, MongoDB, and — for the model-import
|
||||
suite — an `mlflow server` subprocess started from this repository's venv. It lives outside
|
||||
`testpaths`, so it is never collected by the unit gate.
|
||||
|
||||
**One-time setup, required before running anything under `e2e/`:**
|
||||
|
||||
```bash
|
||||
echo "ryuk.disabled=true" >> ~/.testcontainers.properties
|
||||
```
|
||||
|
||||
Without it, **every** e2e test fails at setup with a Docker mount error and never reaches an
|
||||
assertion:
|
||||
|
||||
```
|
||||
docker.errors.APIError: 500 Server Error ... /containers/<id>/start:
|
||||
Mounts denied: The path /socket_mnt/home/<user>/.docker/desktop/docker.sock
|
||||
is not shared from the host and is not known to Docker.
|
||||
```
|
||||
|
||||
The cause is not the tests. Before starting any test container, testcontainers starts a reaper
|
||||
container ("Ryuk") that bind-mounts the Docker socket so it can clean up if the test process dies
|
||||
without teardown. Under Docker Desktop on Linux that socket is `~/.docker/desktop/docker.sock`, a path
|
||||
Docker Desktop's File Sharing does not expose, so the reaper's own container cannot start. Note the
|
||||
symptom is `ERROR`, not `FAILED` — the failure is in fixture setup.
|
||||
|
||||
Disabling it costs nothing in a normal run: every fixture stops its own containers in teardown. Ryuk
|
||||
only matters when the process is killed outright (`kill -9`, a crash, a hard stop from the IDE). After
|
||||
one of those, clean up with:
|
||||
|
||||
```bash
|
||||
docker ps -a --filter "label=org.testcontainers=true" -q | xargs -r docker rm -f
|
||||
```
|
||||
|
||||
**Running it:**
|
||||
|
||||
```bash
|
||||
source ./venv/bin/activate
|
||||
|
||||
# the predictions and OPC suites
|
||||
pytest e2e -m "not import_e2e"
|
||||
|
||||
# the .sientia model-import suite (opt-in; see e2e/scenarios.md § 4)
|
||||
pytest e2e -m import_e2e
|
||||
```
|
||||
|
||||
Full scenario catalogue and the import suite's fidelity contract:
|
||||
[`e2e/scenarios.md`](e2e/scenarios.md).
|
||||
|
||||
## 📊 Monitoring and Metrics
|
||||
|
||||
The Laborious system exposes comprehensive Prometheus metrics for operational visibility and performance monitoring:
|
||||
@@ -802,6 +979,7 @@ See [OPC UA Communication](#opc-ua-communication) for semantics, concurrency, an
|
||||
|----------|-------------|---------|----------|
|
||||
| `TEMPORAL_HOST` | Temporal server address | `localhost:7233` | Yes |
|
||||
| `TEMPORAL_NAMESPACE` | Temporal namespace | `laborious` | No |
|
||||
| `RUNTIME` | Task queue suffix for all workflows (`{workflow}-{RUNTIME}-queue`) | _(none)_ | Yes |
|
||||
| `POSTGRES_HOST` | PostgreSQL hostname | `localhost` | Yes |
|
||||
| `POSTGRES_PORT` | PostgreSQL port | `5432` | Yes |
|
||||
| `POSTGRES_USER` | PostgreSQL username | `sientia` | Yes |
|
||||
@@ -813,6 +991,19 @@ See [OPC UA Communication](#opc-ua-communication) for semantics, concurrency, an
|
||||
| `MLFLOW_PORT` | MLFlow server port | `5080` | Yes |
|
||||
| `MLFLOW_USERNAME` | MLFlow username | `aignosi` | Yes |
|
||||
| `MLFLOW_PASSWORD` | MLFlow password | `aignosi` | Yes |
|
||||
| `IMPORT_MODEL_ENABLED` | Register the import worker on `import_model-{RUNTIME}-queue` | _(unset)_ | No |
|
||||
| `IMPORT_PASSWORD_KEY` | base64 of the 32-byte AES-256-GCM key wrapping bundle passwords | _(none)_ | If importing |
|
||||
| `IMPORT_BUNDLE_BUCKET` | Bucket holding uploaded `.sientia` objects | `sientia` | No |
|
||||
| `IMPORT_BUNDLE_PREFIX` | Key prefix the uploads must sit under | `imported_models/` | No |
|
||||
| `IMPORT_BUNDLE_RETENTION_DAYS` | Lifecycle expiry applied to that prefix | `7` | No |
|
||||
| `IMPORT_WORK_DIR` | Worker scratch directory for bundles | `<tempdir>/sientia-import` | No |
|
||||
| `IMPORT_MAX_OBJECT_BYTES` | Gate 1 object size ceiling | `1073741824` | No |
|
||||
| `IMPORT_MAX_ARCHIVE_ENTRIES` | Gate 4 entry count ceiling | `5000` | No |
|
||||
| `IMPORT_MAX_UNCOMPRESSED_BYTES` | Gate 4 uncompressed size ceiling | `4294967296` | No |
|
||||
| `IMPORT_MAX_COMPRESSION_RATIO` | Gate 4 compression ratio ceiling | `200` | No |
|
||||
| `IMPORT_MODELS_COLLECTION` | MongoDB collection holding the model listing | `models` | No |
|
||||
| `IMPORT_STATUS_DB_NAME` | Import log database name (the BFF's, on the `POSTGRES_*` server) | `sientia-core-mlops-bff` | If importing |
|
||||
| `IMPORT_STATUS_DB_{HOST,PORT,USER,PASSWORD}` | Override, only when the import log is not on the `POSTGRES_*` server | the `POSTGRES_*` value | No |
|
||||
| `OPC_CONFIG` | OPC server configuration (JSON) | `{}` | No |
|
||||
| `OPC_ID` | OPC server identifier | `1` | No |
|
||||
| `OPC_URL` | OPC server URL | `opc.tcp://localhost:4840` | No |
|
||||
@@ -913,6 +1104,33 @@ Legacy MinIO object layout (relative key):
|
||||
`training_datasets/{model_name}/{object_prefix}_{timestamp}.parquet` where `object_prefix` is sanitized
|
||||
(slashes replaced by underscores) to keep a stable model-level directory.
|
||||
|
||||
## Model Flavors
|
||||
|
||||
`model_config.predict_flavor` / `model_config.transform_flavor` accept `sklearn`, `pyfunc` and
|
||||
`joblib`. The flavor is configuration, never discovered from the artifact, and there is
|
||||
no fallback between flavors: a failing loader propagates.
|
||||
|
||||
`model_config` also carries ad-hoc behavioral flags, unrelated to flavor, e.g.
|
||||
`skip_predict_timestamp` (default `False`): when `True`, `request_predict` does not reinsert the
|
||||
`timestamp` column derived from the index into the transformed data before calling predict. Full
|
||||
list: [docs/model-flavors.md § Other `model_config` flags](docs/model-flavors.md#other-model_config-flags).
|
||||
|
||||
`joblib` reads the artifact directly — download, `<artifact>/model.pkl`, `joblib.load` — and writes on
|
||||
retrain (`model.pkl` + `code/`, no `MLmodel`, so only this runtime reads it back).
|
||||
|
||||
Every flavor puts the artifact's own `code/` on `sys.path` before loading, **when that directory
|
||||
exists**: MLflow's loaders only do it when the `MLmodel` manifest declares `code`. `download_model`
|
||||
owns that step for all of them — it asks the registry whether the artifact carries `code/` and, only
|
||||
then, downloads that one directory. The models themselves are never pre-downloaded: `sklearn` and
|
||||
`pyfunc` load straight from the model URI and fetch what they need.
|
||||
|
||||
`download_model` does not branch on the flavor: past that shared first step it delegates to
|
||||
`load_predict_model` / `load_transform_model`, and its only branch left is `load_wrapper`, which is
|
||||
orthogonal to the flavor. It returns `(model, artifact_dir)`, the directory the retrain path re-logs
|
||||
`code/` from.
|
||||
|
||||
Full reference, including how to switch an existing model over: **[docs/model-flavors.md](docs/model-flavors.md)**.
|
||||
|
||||
## OPC UA Communication
|
||||
|
||||
Full reference: **[docs/opc-communication.md](docs/opc-communication.md)** (connection lifecycle, Tier-1 `Bad*` reconnect, connection lock / session readiness, metrics, PostgreSQL confidence **12** vs **14**, tests).
|
||||
@@ -1097,8 +1315,7 @@ laborious/
|
||||
│ ├── prediction_process.py # Core prediction workflow
|
||||
│ └── format_and_export_prediction.py # Export workflow
|
||||
├── worker/ # Worker implementation
|
||||
│ ├── worker.py # Main worker orchestrator
|
||||
│ └── prepare_worker.py # Worker factory with autoscaling config
|
||||
│ └── worker.py # Main worker orchestrator (uses sientia_do prepare_worker)
|
||||
├── utils/ # Utility functions
|
||||
│ ├── connectors_config.py # Environment-driven config builders
|
||||
│ ├── models/ # Data models
|
||||
@@ -1179,7 +1396,9 @@ export LOG_LEVEL=DEBUG
|
||||
### Scaling Considerations
|
||||
|
||||
- **Horizontal Scaling**: Deploy multiple worker instances
|
||||
- **Task Queue Distribution**: Use multiple task queues for different workflow types
|
||||
- **Task Queue Distribution**: One worker pod per `RUNTIME`; queues are
|
||||
`predictions_batch-{RUNTIME}-queue`, `minimal_retrain-{RUNTIME}-queue`,
|
||||
`drift-{RUNTIME}-queue`, `simple_metrics-{RUNTIME}-queue`
|
||||
- **Database Performance**: Optimize indexes and connection pooling
|
||||
- **MLFlow Performance**: Configure appropriate model serving resources
|
||||
|
||||
|
||||
40
docs/minio-object-expiry.md
Normal file
40
docs/minio-object-expiry.md
Normal file
@@ -0,0 +1,40 @@
|
||||
# Two independent MinIO expiry mechanisms
|
||||
|
||||
This platform expires objects in MinIO through two mechanisms that were built independently, for
|
||||
different features, and remain independent today. This note exists so a reader who finds one does
|
||||
not assume it is the whole story.
|
||||
|
||||
## `ModelImport.ensure_import_bundle_retention`
|
||||
|
||||
`laborious/activities/model_import.py`. Applies a **bucket lifecycle rule** (a MinIO/S3-native
|
||||
`Expiration` rule, `ENABLED`, keyed by a fixed `rule_id`) to the import bundle prefix
|
||||
(`import_config['prefix']`), expiring objects after `import_config['retention_days']` days. The
|
||||
rule is declarative: once set, MinIO itself removes the object when it ages out — the workflow does
|
||||
not scan for or delete anything. The activity's job is only to ensure the rule exists and matches
|
||||
the configured `days`/`prefix`, which is what makes "lifecycle active" a criterion the code can
|
||||
verify rather than an assumption about a Helm value.
|
||||
|
||||
## `Storage.cleanup_minio_objects_expired`
|
||||
|
||||
`laborious/activities/storage.py`. An older, unrelated activity that **scans and deletes**: given a
|
||||
list of key prefixes and a retention window (`retention_hours`), it lists objects under each prefix,
|
||||
compares their age against `now() - retention_hours`, and deletes the ones that have aged out. There
|
||||
is no bucket-side rule — the deletion happens in this code, on whatever schedule calls the activity.
|
||||
|
||||
## Why they are not unified today
|
||||
|
||||
They expire different prefix families, on different units (days vs. hours), through different
|
||||
mechanisms (declarative rule vs. active scan), and were written for different features at different
|
||||
times. Neither mechanism depends on or interferes with the other; a bucket can carry an
|
||||
`ensure_import_bundle_retention` lifecycle rule for one prefix and still have
|
||||
`cleanup_minio_objects_expired` scanning a different prefix on its own cadence.
|
||||
|
||||
## What unifying them would require
|
||||
|
||||
Either retiring the scan-and-delete path in favor of a lifecycle rule for every prefix it currently
|
||||
handles (only possible if every one of its callers can tolerate a fixed number of retention *days*,
|
||||
since S3-style lifecycle rules do not support hour-granularity windows), or teaching the
|
||||
lifecycle-rule path to manage a per-prefix table of rules with arbitrary retention windows instead of
|
||||
the one fixed rule it manages today. Either direction touches callers and configuration this change
|
||||
does not, and neither is a refactor — this is tracked as a separate, larger change, not a task of
|
||||
this one.
|
||||
166
docs/model-flavors.md
Normal file
166
docs/model-flavors.md
Normal file
@@ -0,0 +1,166 @@
|
||||
# Model flavors
|
||||
|
||||
The *flavor* is how this runtime is told to deserialize a model. It is **configuration, never
|
||||
discovery**: the value comes from the model document in the `models` collection and the code never
|
||||
reads the artifact's `MLmodel` manifest to guess it.
|
||||
|
||||
```json
|
||||
{
|
||||
"model_config": {
|
||||
"transform_flavor": "sklearn",
|
||||
"predict_flavor": "joblib",
|
||||
"retention_minutes": 60,
|
||||
"target": "SE"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
`predict_flavor` drives `load_predict_model` (artifact `prediction_model`); `transform_flavor` drives
|
||||
`load_transform_model` (artifact `data_model`). Both default to `sklearn` when the key is absent.
|
||||
|
||||
## Other `model_config` flags
|
||||
|
||||
Ad-hoc behavioral overrides, read directly with `.get(key, default)` — not part of the model
|
||||
document written by import (`model-import.md` § 12), and not validated against a schema:
|
||||
|
||||
| Key | Default | Effect |
|
||||
|---|---|---|
|
||||
| `skip_transform` | `False` | `fit_models` (retrain) skips calling `transform` before fitting. |
|
||||
| `skip_predict_timestamp` | `False` | `request_predict` skips reinserting the `timestamp` column (from the index) into the transformed data before calling predict. |
|
||||
|
||||
## Accepted values
|
||||
|
||||
| Flavor | Read | Write (retrain) | How it loads |
|
||||
|---|---|---|---|
|
||||
| `sklearn` | yes | yes | `mlflow.sklearn.load_model(<uri>)` — the loader fetches what it needs |
|
||||
| `pyfunc` | yes | yes | `mlflow.pyfunc.load_model(<uri>)`; the wrapper path downloads the artifact instead and unwraps `._model_impl.python_model` |
|
||||
| `joblib` | yes | yes | downloads the same artifact and `joblib.load`s its `<artifact>/model.pkl` — **no MLflow flavor loader involved** |
|
||||
|
||||
Anything else raises `ValueError` with `INVALID_FLAVOR_MESSAGE`, **before** any artifact is resolved
|
||||
or downloaded. The same message and the same three values govern reading and writing.
|
||||
|
||||
`pytorch` was accepted until this change and is now rejected like any other unknown value: `torch` is
|
||||
not installed in the runtime, so `mlflow.pytorch.load_model` could never have loaded anything there.
|
||||
|
||||
**There is no fallback between flavors.** A failing loader propagates: one attempt, one cause. A
|
||||
model configured with the wrong flavor fails visibly instead of being rescued silently.
|
||||
|
||||
## Same artifact, three readers
|
||||
|
||||
The flavor changes the reader, never the path. All three resolve the same two artifact directories of
|
||||
the Production run:
|
||||
|
||||
```
|
||||
<run>/artifacts/
|
||||
├── data_model/ ← transform_flavor loads this one
|
||||
│ ├── MLmodel ← parsed by sklearn/pyfunc; ignored by joblib
|
||||
│ ├── model.pkl ← what joblib.load reads
|
||||
│ └── code/ ← put on sys.path by every flavor, when present
|
||||
├── prediction_model/ ← predict_flavor loads this one
|
||||
└── model_card/, *.csv
|
||||
```
|
||||
|
||||
The run's `artifacts/` root holds **no** `.pkl` — every model pickle lives one level down, inside its
|
||||
artifact directory. Predict resolves that directory through the registry URI
|
||||
(`models:/<name>/production`, which MLflow points at the same `prediction_model`); transform builds
|
||||
the run artifact URI directly with `get_model_uri`, since the run is already resolved. Observability
|
||||
is the same for the three: `MODEL_READ_*` on load, `MODEL_WRITE_*` on `log_model`, under the base
|
||||
`operation_type` — no per-flavor metric, operation type or label.
|
||||
|
||||
## `code/` on `sys.path`, for every flavor
|
||||
|
||||
An artifact may ship the modules its own classes live in, under `<artifact>/code/`. Without that
|
||||
directory on `sys.path`, unpickling raises `ModuleNotFoundError` — the artifact is intact, the class
|
||||
is simply not importable.
|
||||
|
||||
MLflow's loaders do add it, but **from the manifest**: `_add_code_from_conf_to_system_path` acts only
|
||||
when the `MLmodel` flavor config declares a `code` entry. An artifact carrying `code/` on disk
|
||||
without that entry — code logged with `log_artifacts`, a re-log that dropped the config, anything a
|
||||
`joblib` retrain wrote — loads with no module path at all. So this runtime reads the **directory**,
|
||||
which no manifest can misdescribe:
|
||||
|
||||
- `prepend_artifact_code_dir` inserts `<artifact>/code` at the front of `sys.path`, **only when that
|
||||
directory exists**, and only once (an entry already present is not duplicated).
|
||||
- `download_artifact_code_dir` is what feeds it on the download paths, and it is deliberately cheap:
|
||||
the remote listing (`list_artifacts(run_id, '<artifact>')`) decides first, so a run without `code`
|
||||
downloads **nothing**, and a run with it downloads only `<artifact>/code` — a handful of `.py`
|
||||
files, never the model's weights.
|
||||
- `download_model` runs that step first, for every flavor and both of its branches, *before* any
|
||||
model is loaded — the insert has to be in place before the loader unpickles.
|
||||
- The models themselves are not pre-downloaded: `mlflow.sklearn.load_model` and
|
||||
`mlflow.pyfunc.load_model` are handed the model URI and fetch what they need. `joblib` downloads
|
||||
its artifact because it must read the pickle, and adds the entry again inside `_load_model_pickle`
|
||||
through the same helper, so the flavor works when called directly too.
|
||||
- When the manifest *does* declare `code`, MLflow inserts the same directory a second time. Harmless.
|
||||
|
||||
The entry is never removed: a loaded model may import lazily, long after the load returned.
|
||||
|
||||
The directory `download_artifact_code_dir` downloaded into is also what `download_model` returns
|
||||
next to the model — `tuple[Any, str | None]`, the second element being the artifact directory the
|
||||
retrain path re-logs `code/` from, or `None` when the artifact has no `code/`. It falls back to
|
||||
whatever the loader downloaded on its own, which is a real artifact directory for `joblib` and for
|
||||
the wrapper branch. `load_predict_model` and `load_transform_model` stay the **single** place where
|
||||
the flavor is decided: `download_model` has no `joblib` branch, it delegates by `model_type`. Its
|
||||
only remaining branch is `load_wrapper`, which is derived from `flavor == 'pyfunc'` but is otherwise
|
||||
orthogonal to the flavor.
|
||||
|
||||
> `sys.path` is process-global: two models embedding a package with the same name (`utils/`) on the
|
||||
> same worker resolve to whichever entered first. Mitigation, if it bites: one worker per model.
|
||||
|
||||
## `sklearn`
|
||||
|
||||
Default. A standard MLflow sklearn artifact: `mlflow.sklearn.log_model` writes `MLmodel` +
|
||||
`model.pkl`, `mlflow.sklearn.load_model` reads it, and every outside consumer (MLflow UI, MLflow
|
||||
serving) recognizes it as a model.
|
||||
|
||||
## `pyfunc`
|
||||
|
||||
For models logged as `mlflow.pyfunc`, including the wrapper models this pipeline produces. Retrain
|
||||
writes through the wrapper's own `store_model(artifact_path=<artifact>, code_path=…)`. Two read
|
||||
paths, and the second is **orthogonal to the flavor**:
|
||||
|
||||
- Plain: `mlflow.pyfunc.load_model(<uri>)`, called as `predict(context, model_input)`.
|
||||
- Wrapper (`download_model(load_wrapper=True)`): downloads to `./tmp/artifacts/<model>/<artifact>`
|
||||
via `dowload_artifacts`, loads as pyfunc, then unwraps `._model_impl.python_model`. Inside such an
|
||||
artifact the real pickle is **nested** — `<artifact>/artifacts/stacking_model.pkl` for predict,
|
||||
`<artifact>/artifacts/training_transformer.pkl` for transform (`PREDICTION_COMPRESSED_PATH` /
|
||||
`TRANSFORMED_COMPRESSED_PATH`). That nesting is why the `joblib` listing does not recurse.
|
||||
|
||||
## `joblib`
|
||||
|
||||
For models whose `.pkl` is in the compressed form `joblib.dump` writes, which MLflow's `pickle.load`
|
||||
does not read — the artifact is intact, but no flavor loader opens it.
|
||||
|
||||
- **Read**: `client.download_artifacts(run_id, <artifact>)`, then
|
||||
`joblib.load(<artifact>/model.pkl)`. `model.pkl` is `MODEL_PICKLE_NAME`, the same constant
|
||||
`log_model` writes; another `.pkl` directly inside the artifact is a fallback, in name order.
|
||||
- **Write**: `joblib.dump(model, '<artifact>/model.pkl', compress=3)`, plus a re-log of the source
|
||||
artifact's `code/` to `<artifact>/code`.
|
||||
- **`code/` goes on `sys.path`** before the load, when present — see the section above; the joblib
|
||||
flavor is where that started, and it is now shared by all three.
|
||||
- **No `MLmodel` is written**, so only this runtime reads the artifact back —
|
||||
`mlflow.sklearn.load_model` and `models:/<name>/production` both fail on it. A model retrained in
|
||||
joblib stays joblib.
|
||||
- **The pickled object must expose `predict(data)`**, plus `fit(data)` if retrained. A pyfunc
|
||||
`PythonModel` wrapper expecting `predict(context, model_input)` raises `TypeError`.
|
||||
|
||||
## Switching a model to `joblib`
|
||||
|
||||
The runtime does not migrate documents; editing `models` is an operator step.
|
||||
|
||||
1. **Inventory before the deploy** — run predict/transform and note which models raise. There is no
|
||||
discovery afterwards: the failure names the loader, not the fix.
|
||||
2. **Edit the document** — `predict_flavor` and/or `transform_flavor` to `'joblib'`. Imported models
|
||||
arrive ready: `import_model` writes it itself, see [`model-import.md`](model-import.md) § 12.
|
||||
3. **Check the first retrain** — the new run must carry `prediction_model/model.pkl` and, when the
|
||||
source had one, `prediction_model/code/`.
|
||||
4. **Rollback is not just a revert** — artifacts a joblib retrain already wrote stop being loadable,
|
||||
so it also means promoting the previous version back and reverting the documents' flavors.
|
||||
|
||||
## Known gap: how the run id is resolved
|
||||
|
||||
`get_model_run_id` takes `source.split('/')[2]` of the registered version's `source`. **Not confirmed
|
||||
against the real tracking server**: no scheme reproduced locally puts the run id there (a file store
|
||||
yields `''`, S3 the bucket, `mlflow-artifacts:/` the experiment id), yet predict works in production.
|
||||
Every flavor that resolves a run goes through this line — confirm it with a **read-only** query
|
||||
before relying on it.
|
||||
417
docs/model-import.md
Normal file
417
docs/model-import.md
Normal file
@@ -0,0 +1,417 @@
|
||||
# Importing a model (`import_model`)
|
||||
|
||||
The `import_model` workflow turns an encrypted `.sientia` object in MinIO into a registered MLflow
|
||||
model plus one document in the MongoDB `models` collection — or into a row that says, in a sentence a
|
||||
person can read, why it did not.
|
||||
|
||||
This document is the contract for the two sides that talk to it: the **frontend**, which uploads the
|
||||
object and builds the request, and the **BFF**, which owns the import log row. Everything here is
|
||||
enforced by code in this repository; where a statement is about a column or a constraint, the
|
||||
authority is the migration named in §7.
|
||||
|
||||
Format of the file itself: [`sientia-bundle-format.md`](sientia-bundle-format.md).
|
||||
|
||||
---
|
||||
|
||||
## 1. What the caller does, in order
|
||||
|
||||
1. **Upload** the `.sientia` object to the import bucket, under the configured prefix (default
|
||||
`imported_models/`). The key must end in `.sientia`.
|
||||
2. **Compute `expected_digest`** — the SHA-256, lowercase hex, of the **exact bytes that were
|
||||
uploaded**. Not of a re-read, not of a re-serialisation: of what went over the wire.
|
||||
3. **Encrypt the bundle password** into a `password_envelope` (§3).
|
||||
4. **Insert the import log row** (§5) and take its `id`.
|
||||
5. **Start the workflow** on `import_model-{RUNTIME}-queue` with the input of §2.
|
||||
|
||||
The workflow claims the row, opens the bundle, provisions MLflow, writes the model document, cleans
|
||||
up and records one terminal verdict. It never inserts a row, never promotes a version and never
|
||||
activates a model.
|
||||
|
||||
## 2. The workflow input
|
||||
|
||||
```json
|
||||
{
|
||||
"import_run_id": 41,
|
||||
"object_key": "imported_models/sales_forecast_v3.sientia",
|
||||
"expected_digest": "9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08",
|
||||
"password_envelope": "Base64OfNonceCiphertextTag==",
|
||||
"bucket": "sientia"
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Required | Meaning |
|
||||
|---|---|---|
|
||||
| `import_run_id` | yes | primary key of the `public.experiment_run` row the BFF inserted. Must be a positive integer |
|
||||
| `object_key` | yes | the uploaded object's key. Must match the row's `file_name` when the row carries one |
|
||||
| `expected_digest` | **yes** | SHA-256 of the uploaded bytes, 64 lowercase hex characters. There is no path that skips the comparison: a missing value fails at `received`, a mismatching one at `download`, before the header is parsed |
|
||||
| `password_envelope` | yes | the bundle password, encrypted (§3). A plaintext `password` field is **refused** |
|
||||
| `bucket` | no | defaults to `IMPORT_BUNDLE_BUCKET` |
|
||||
|
||||
**Fields that are refused rather than ignored**: `project_id`, `project_name`,
|
||||
`provisioning_target`, `schema`, `table_name`, `status_table`, `password`. An input carrying any of
|
||||
them fails at `received`, so no caller can believe it recorded something the import does not record
|
||||
(see §10).
|
||||
|
||||
## 3. The password envelope
|
||||
|
||||
AES-256-GCM, and the ciphertext layout is what the WebCrypto API returns natively:
|
||||
|
||||
```
|
||||
envelope = base64( nonce(12 bytes) || ciphertext || tag(16 bytes) )
|
||||
```
|
||||
|
||||
`crypto.subtle.encrypt({name: 'AES-GCM', iv: nonce}, key, plaintext)` returns `ciphertext || tag`
|
||||
already concatenated, so the browser side is: generate 12 random bytes, encrypt, concatenate the
|
||||
nonce in front, base64 the result.
|
||||
|
||||
The key is a 32-byte AES-256 key shared with the worker through `IMPORT_PASSWORD_KEY` (base64). The
|
||||
worker decrypts the envelope **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 appears in
|
||||
the Temporal event history. A test asserts that for every failure path.
|
||||
|
||||
A failure to open the envelope is recorded as a `decryption` failure, distinct in the logs from a
|
||||
wrong *bundle* password but sharing its step, because the step says where the import stopped, not who
|
||||
is at fault.
|
||||
|
||||
## 4. Environment
|
||||
|
||||
| Variable | Default | What it is |
|
||||
|---|---|---|
|
||||
| `IMPORT_MODEL_ENABLED` | unset | registers the fifth (import) worker. Truthy values: `1`, `true`, `yes`, `on` |
|
||||
| `IMPORT_PASSWORD_KEY` | — | base64 of the 32-byte AES-256-GCM key that wraps bundle passwords |
|
||||
| `IMPORT_BUNDLE_BUCKET` | `sientia` | bucket holding the uploaded bundles |
|
||||
| `IMPORT_BUNDLE_PREFIX` | `imported_models/` | key prefix the uploads must sit under |
|
||||
| `IMPORT_BUNDLE_RETENTION_DAYS` | `7` | lifecycle expiry the importer applies to that prefix |
|
||||
| `IMPORT_WORK_DIR` | `<tempdir>/sientia-import` | worker scratch directory |
|
||||
| `IMPORT_MAX_OBJECT_BYTES` | 1 GiB | gate 1 |
|
||||
| `IMPORT_MAX_ARCHIVE_ENTRIES` | `5000` | gate 4 |
|
||||
| `IMPORT_MAX_UNCOMPRESSED_BYTES` | 4 GiB | gate 4 |
|
||||
| `IMPORT_MAX_COMPRESSION_RATIO` | `200` | gate 4 |
|
||||
| `IMPORT_MODELS_COLLECTION` | `models` | the MongoDB collection that *is* the model listing |
|
||||
| `IMPORT_STATUS_DB_NAME` | `sientia-core-mlops-bff` | the BFF database holding the import log |
|
||||
| `IMPORT_STATUS_DB_HOST` | `POSTGRES_HOST` | override: the log is on another server |
|
||||
| `IMPORT_STATUS_DB_PORT` | `POSTGRES_PORT` | override |
|
||||
| `IMPORT_STATUS_DB_USER` | `POSTGRES_USER` | override |
|
||||
| `IMPORT_STATUS_DB_PASSWORD` | `POSTGRES_PASSWORD` | override |
|
||||
|
||||
The import log is a **second connection but the same server**: host, port, user and password come
|
||||
from `POSTGRES_*`, and only the database name is import configuration, because the log lives in the
|
||||
BFF's database rather than in the one Laborious holds for predictions. The import writes to that
|
||||
database and no other, and to one table in it. The four `IMPORT_STATUS_DB_*` connection variables
|
||||
are overrides for a deployment whose log genuinely sits elsewhere — the e2e suite, which runs the
|
||||
two databases as two containers, is exactly that case.
|
||||
|
||||
A missing or malformed `IMPORT_PASSWORD_KEY` is an import failure at the decryption step, not a boot
|
||||
failure: the worker starts either way.
|
||||
|
||||
## 5. The import log row — the BFF's half
|
||||
|
||||
The row is `public.experiment_run` in `sientia-core-mlops-bff`. **The workflow never inserts it.**
|
||||
The inserting side writes:
|
||||
|
||||
| Column | Value |
|
||||
|---|---|
|
||||
| `run_type` | `'IMPORT'` |
|
||||
| `status` | `'PENDING'` |
|
||||
| `username` | whoever asked for the import |
|
||||
| `bucket_name` | the bucket the object was uploaded to |
|
||||
| `file_name` | the object key |
|
||||
| `request_data` | the request payload. If it carries a digest under `expected_digest`, `digest`, `sha256` or `file_digest`, the workflow checks it against the input's |
|
||||
|
||||
Then it passes the returned `id` as `import_run_id`.
|
||||
|
||||
The workflow's first act is to **claim** that row:
|
||||
|
||||
```sql
|
||||
UPDATE public.experiment_run
|
||||
SET status = 'ORCHESTRATOR_WAITING_PROC', updated_at = :now,
|
||||
orchestrator_response_data = :detail
|
||||
WHERE id = :import_run_id
|
||||
AND run_type = 'IMPORT'
|
||||
AND (status = 'PENDING'
|
||||
OR (status = 'ORCHESTRATOR_WAITING_PROC'
|
||||
AND orchestrator_response_data -> 'import' ->> 'workflow_id' = :workflow_id))
|
||||
```
|
||||
|
||||
The second disjunct makes a re-execution of the same workflow idempotent instead of a conflict. Every
|
||||
later write keys on `id` alone, because ownership was settled here.
|
||||
|
||||
**Columns the workflow never writes**: `username`, `request_data`, `bucket_name`, `file_name`,
|
||||
`run_type`, `created_at`. **Columns it does write**: `status`, `error_message`, `updated_at`,
|
||||
`orchestrator_response_data`, and — as soon as the bundle is open, before anything exists in MLflow —
|
||||
`experiment_name` (the bundle's own, verbatim) and `run_name` (`import-{model_name}-v{model_version}`,
|
||||
truncated inside the name segment to fit 50 characters). Those two are **overwritten** by the
|
||||
workflow; whatever the inserting side put there is replaced.
|
||||
|
||||
### When the row cannot be claimed
|
||||
|
||||
If the claim matches nothing — no such id, a `run_type='TRAIN'` row, or a row already running under
|
||||
another workflow — the import **writes nothing anywhere**, because the row belongs to someone else.
|
||||
Instead it:
|
||||
|
||||
- fails the Temporal run with error type `ImportLogRowNotClaimableError` and code
|
||||
`IMPORT_LOG_ROW_NOT_CLAIMABLE`, non-retryable;
|
||||
- increments the import error metric and fires a notification;
|
||||
- logs the id and the observed state (which of the three causes it was).
|
||||
|
||||
So a row left `PENDING` with no `workflow_id` in its detail JSON means the workflow never owned it.
|
||||
**Closing that row out belongs to whichever side inserted it** — by timeout, by a sweeper, by a person
|
||||
— however that side chooses. The importer deliberately does not guess.
|
||||
|
||||
## 6. Status vocabulary
|
||||
|
||||
| Status | Written when | New in `V12`? |
|
||||
|---|---|---|
|
||||
| `ORCHESTRATOR_WAITING_PROC` | claimed, and while running | no — **reused** from the training flow |
|
||||
| `ORCHESTRATOR_VALIDATION_ERROR` | failure at `received` … `content_policy` — the user can act on it | no — **reused** |
|
||||
| `IMPORT_SUCCESS` | every step succeeded | yes |
|
||||
| `IMPORT_ERROR` | failure at `experiment_creation` … `model_document`, or unclassified — only support can act on it | yes |
|
||||
|
||||
The two reused values describe a *stage*, not a flow: what says the row is an import is `run_type`,
|
||||
not the status. There is deliberately no `IMPORT_RUNNING`; it would carry no information `run_type`
|
||||
does not already carry, and it would widen both the pairing `CHECK` and the frontend's switch for
|
||||
nothing.
|
||||
|
||||
`error_message` is the user-facing sentence, taken from the catalog in §8. On success it is written as
|
||||
`NULL`, never `''` — the column carries a `LENGTH(error_message) >= 3` check.
|
||||
|
||||
## 7. The migration this depends on
|
||||
|
||||
`sientia-dataops-database-migrations`,
|
||||
`postgres/sientia-core-mlops-bff/migrations/V12__add_import_statuses_to_experiment_run.sql`:
|
||||
|
||||
- adds `IMPORT_SUCCESS` and `IMPORT_ERROR` to `chk_experiment_run_status`;
|
||||
- adds `chk_experiment_run_terminal_status_run_type`, pairing only the **terminal** values with their
|
||||
`run_type` (`IMPORT_*` on `'IMPORT'`, `TRAINING_*` on `'TRAIN'`) and leaving the shared
|
||||
`ORCHESTRATOR_*` values out of the rule;
|
||||
- **adds no columns at all.**
|
||||
|
||||
`V12` must be merged and applied before `IMPORT_MODEL_ENABLED` is set anywhere. Without it, an import
|
||||
provisions the model and then cannot record its verdict: the claim still succeeds (no new columns are
|
||||
needed for it), the old `CHECK` refuses `IMPORT_SUCCESS`, and the failure surfaces as
|
||||
`IMPORT_STATUS_NOT_RECORDED` on the first attempt — no retry wait — leaving the row in
|
||||
`ORCHESTRATOR_WAITING_PROC` with a fully provisioned model behind it.
|
||||
|
||||
## 8. `orchestrator_response_data`, and the error catalog
|
||||
|
||||
There is no column for the error code and none for the workflow identifier (decision D21). Both live
|
||||
at fixed keys in the JSON column the workflow already owns, and the whole `import` object is rewritten
|
||||
on every write, so the identifier the claim wrote is never dropped:
|
||||
|
||||
```json
|
||||
{
|
||||
"import": {
|
||||
"workflow_id": "import_model-41-01HW…",
|
||||
"code": "IMPORT_BUNDLE_INCOMPLETE",
|
||||
"step": "structure_validation",
|
||||
"gate": 6,
|
||||
"source": { "model_name": "sales_forecast", "experiment_name": "…", "…": "…" },
|
||||
"cleanup_failed": false
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- `workflow_id` — what a user quotes to support, and what the claim's idempotency disjunct reads.
|
||||
- `code` — the frontend's i18n key. **Absent on success.**
|
||||
- `step` — where the import stopped (the values of §9).
|
||||
- `gate` — which of the seven gates rejected it, when one did.
|
||||
- `source` — the bundle's own metadata block, once it is known.
|
||||
- `cleanup_failed` — temporary files were left on the worker. Present on a **successful** import too:
|
||||
it never changes the status, because the user's model was imported.
|
||||
|
||||
### Code → reason
|
||||
|
||||
`error_message` is exactly the sentence in this table; nothing is interpolated into it, and the
|
||||
exception's text — paths, table names, exception classes, tracebacks — reaches only the logs.
|
||||
|
||||
| Code | Step | Sentence written to `error_message` |
|
||||
|---|---|---|
|
||||
| `IMPORT_REQUEST_INVALID` | `received` | The import could not be started because the request was incomplete or contradictory. Nothing was created. Submit the import again. |
|
||||
| `IMPORT_FILE_UNREADABLE` | `download` | The uploaded file could not be read, or it is not the file that was sent. Upload it again. |
|
||||
| `IMPORT_BUNDLE_CANNOT_BE_OPENED` | `decryption` | 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. |
|
||||
| `IMPORT_BUNDLE_REJECTED` | `archive_inspection` | The file's contents did not pass the platform's safety checks, so it was not opened. Export the model again from the origin platform. |
|
||||
| `IMPORT_BUNDLE_UNPACK_FAILED` | `extraction` | The file could not be unpacked safely and nothing from it was kept. Export the model again from the origin platform. |
|
||||
| `IMPORT_BUNDLE_INCOMPLETE` | `structure_validation` | This file is not a complete model export — part of what the platform needs is missing from it. Export the model again from the origin platform. |
|
||||
| `IMPORT_BUNDLE_UNEXPECTED_CONTENT` | `content_policy` | The file contains something a model export should not contain, so it was rejected. Export the model again from the origin platform. |
|
||||
| `IMPORT_STORAGE_PREPARATION_FAILED` | `experiment_creation` | The platform could not prepare a place to keep this model. The import stopped and no model was created. Contact support with this import's identifier. |
|
||||
| `IMPORT_MODEL_FILES_NOT_STORED` | `artifact_upload` | The model's files could not be stored on this platform. The import stopped and the model is not available. Contact support with this import's identifier. |
|
||||
| `IMPORT_MODEL_NOT_REGISTERED` | `registration` | The model's files were stored, but the model itself could not be registered, so it cannot be used. Contact support with this import's identifier. |
|
||||
| `IMPORT_MODEL_SETTINGS_NOT_SAVED` | `model_document` | The model was registered, but the settings that tell the platform how to run it could not be saved. Contact support with this import's identifier. |
|
||||
| `IMPORT_UNEXPECTED_ERROR` | — | The import stopped for an unexpected reason and no model was created. Contact support with this import's identifier. |
|
||||
|
||||
Three codes never reach a row's `error_message`, because there is no row to write or no failure to
|
||||
report to a user. They exist in the logs, the metric and the notification only:
|
||||
|
||||
| Code | When |
|
||||
|---|---|
|
||||
| `IMPORT_LOG_ROW_NOT_CLAIMABLE` | the row is not this workflow's (§5) |
|
||||
| `IMPORT_STATUS_NOT_RECORDED` | the verdict itself could not be written; the log line carries the verdict it could not record |
|
||||
| `IMPORT_CLEANUP_INCOMPLETE` | temporary files were left behind; also flagged as `cleanup_failed` in the detail |
|
||||
|
||||
## 9. Steps
|
||||
|
||||
`received`, `download`, `decryption`, `archive_inspection`, `extraction`, `structure_validation`,
|
||||
`content_policy`, `experiment_creation`, `artifact_upload`, `registration`, `model_document`,
|
||||
`cleanup`. The first seven are the gates (see the format document); a failure in any of them is a
|
||||
`ORCHESTRATOR_VALIDATION_ERROR`. The next four are provisioning; a failure there is an
|
||||
`IMPORT_ERROR`. `cleanup` is never a reported failure.
|
||||
|
||||
## 10. The import knows nothing about projects
|
||||
|
||||
The model listing **is** the MongoDB `models` document. No migrated relational database in this
|
||||
platform holds a project table or a model table, nothing in this repository reads a project-to-model
|
||||
link, and the workflow input carries no project target — which is why naming one is refused rather
|
||||
than accepted and ignored. Whatever the import screen does with projects is the frontend's business.
|
||||
|
||||
The bundle's own `metadata.model_project` is logged as an MLflow parameter (`source.model_project`),
|
||||
as a record of where the model came from. It is not a link.
|
||||
|
||||
## 11. What MLflow ends up holding
|
||||
|
||||
| Thing | Value |
|
||||
|---|---|
|
||||
| Experiment | `metadata.experiment_name` from the bundle, created if absent, reused if present |
|
||||
| Run | always **new**, created with the `run_name` already written on the record |
|
||||
| Artifacts | the bundle's whole `artifacts/` tree, with `prediction_model` and `data_model` at the run's artifact root |
|
||||
| Parameters | the bundle's `parameters`, verbatim, plus the origin metadata as `source.*` (including `source.run_id`, which is never dereferenced) |
|
||||
| Registered version | created on demand; source is the run's **resolved artifact URI** with `/prediction_model` appended, e.g. `mlflow-artifacts:/17/<run_id>/artifacts/prediction_model` |
|
||||
| Stage | **untouched** |
|
||||
|
||||
The source shape matters and is easy to get wrong: the legacy loader extracts the run id as
|
||||
`source.split('/')[2]`, which is the run id in the resolved artifact URI but the *artifact path* in a
|
||||
literal `runs:/<run_id>/prediction_model`. Registering with the unresolved URI breaks loading
|
||||
silently — nothing raises, the loader just returns the wrong string. This is the same value
|
||||
`mlflow.register_model` stores for the training flow, because it resolves the `runs:/` URI before
|
||||
writing it; `MlflowClient.create_model_version` does not, so the import resolves it itself.
|
||||
|
||||
## 12. The model document
|
||||
|
||||
Written **last**, into the collection named by `IMPORT_MODELS_COLLECTION` (default `models`):
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "31",
|
||||
"name": "sales_forecast",
|
||||
"active": false,
|
||||
"model_config": {
|
||||
"transform_flavor": "sklearn",
|
||||
"predict_flavor": "joblib",
|
||||
"retention_minutes": 60,
|
||||
"target": "SE"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- `id` is `str(max(int(id) for id in collection) + 1)` — computed in Python, from the ids that are
|
||||
numbers. Never a `count`, never a lexicographic sort (which would answer `"9" > "30"`). Ids that
|
||||
are not decimal numbers are ignored for the maximum and counted in the log.
|
||||
- `name` is `metadata.model_name` from the bundle, verbatim.
|
||||
- `active` is `false`, always. No code path in the import writes or updates a document with
|
||||
`active: true`.
|
||||
- `target` comes from `parameters.target_variable`; nothing else supplies it.
|
||||
- The `models` collection carries a **unique sparse index on `id`**, created by the MongoDB baseline
|
||||
migration. A duplicate-key insert is reported as a `model_document` failure — no retry, no id
|
||||
recomputation.
|
||||
|
||||
**A name already in the collection is left exactly as it is.** The import writes nothing, modifies
|
||||
nothing, logs the skip and succeeds: `active`, `target` and the flavors of a model already on the
|
||||
platform are operator-owned.
|
||||
|
||||
## 13. After a successful import — the manual steps
|
||||
|
||||
A `COMPLETED` import is **not** a servable model. Stated plainly:
|
||||
|
||||
- the registered version is **unstaged**, and `load_predict_model` only sees a version once it is in
|
||||
stage `Production`;
|
||||
- the model document is **inactive** (`active: false`).
|
||||
|
||||
So an operator must promote the version to `Production` and set `active: true`. This workflow never
|
||||
promotes and never activates — that runbook is QTZPOC-21's.
|
||||
|
||||
One more thing to know before serving an imported model: `predict_flavor: 'joblib'` is what the
|
||||
artifact actually is, and it loads, predicts **and** retrains — `log_model` writes that flavor too,
|
||||
as `model.pkl` + `code/`. The artifact it writes carries no `MLmodel` manifest, so an imported model
|
||||
stays readable only by this runtime, before and after a retrain. Full contract:
|
||||
[`model-flavors.md`](model-flavors.md).
|
||||
|
||||
## 14. Retry, and what does not retry
|
||||
|
||||
Only the status writes retry, and only because `UPDATE ... SET <fixed values> WHERE id` leaves the row
|
||||
the same however many times it runs:
|
||||
|
||||
| Call | Attempts | Why |
|
||||
|---|---|---|
|
||||
| claim, terminal write | 10, 1 s → 30 s backoff (≈151 s envelope) | the verdict must land; the database being briefly unavailable is exactly what a retry fixes |
|
||||
| progress writes, names | 3, capped at 2 s | a hint nobody reads — the terminal write rewrites it — must not stall a working import |
|
||||
| everything else | 1 | a repeated registry call or Mongo insert can write twice |
|
||||
|
||||
`ImportLogRowNotClaimableError` and `ImportStatusRejectedError` are non-retryable: the first would
|
||||
read the same rows and reach the same conclusion, and the second means the statement itself was
|
||||
refused (a `CHECK` violation, an undefined column, a value too long), which a second attempt cannot
|
||||
change.
|
||||
|
||||
There is no rollback, no compensation, no resume and no reconciliation. A failed import stays failed;
|
||||
trying again is a new upload and a new row.
|
||||
|
||||
## 15. What a failure looks like in the logs
|
||||
|
||||
Every failure — a row that cannot be claimed, a gate that rejected the bundle, an MLflow that refused
|
||||
the connection, a verdict that could not be written — goes through the one channel in
|
||||
`laborious/utils/models/import_errors.py`, which produces exactly two `logfmt` lines: the technical
|
||||
one, whose message
|
||||
carries the traceback (or the failure chain, root cause first, for anything that crossed an activity
|
||||
boundary), and the user-facing sentence the notification logs. **Both carry the same fields**, so one
|
||||
query returns the pair:
|
||||
|
||||
| Field | What it holds |
|
||||
|---|---|
|
||||
| `import_run_id` | the record's primary key — the id a user quotes to support |
|
||||
| `import_workflow_id` | the Temporal workflow id |
|
||||
| `import_step` | the step that failed, or `not_claimable` / `status_write` |
|
||||
| `import_code` | the code of § 8, the same one written to the row |
|
||||
| `error_type` | the failing exception's own type — `MlflowException`, not the `ApplicationError` wrapping it |
|
||||
| `error_summary` | one bounded line of what it said |
|
||||
| `import_gate` | the gate that rejected the bundle, when one did |
|
||||
| `import_verdict_not_recorded` | the verdict a failed terminal write could not store |
|
||||
|
||||
The workflow metadata every other line already carries (`pod_id`, `runtime`, `model_id`, `model_name`,
|
||||
`workflow_name`, `schedule_name`) is kept alongside them.
|
||||
|
||||
A failure at any pipeline step is reported from the workflow's `except`, before cleanup and before the
|
||||
terminal write, so the log reads in the order things happened and the line's step can never disagree
|
||||
with the row's. That report never changes the outcome: its own failure is swallowed, and the failure it
|
||||
was reporting is re-raised untouched. The Temporal SDK still prints its own traceback for each failed
|
||||
attempt — that is the per-attempt record, and it is not structured; the channel's line is the one to
|
||||
query.
|
||||
|
||||
## 16. How this contract is verified
|
||||
|
||||
Two suites, and they prove different things.
|
||||
|
||||
**Unit** — `pytest tests/`, part of the standard gate (`validate --project-name=laborious`). Covers
|
||||
the bundle reader against committed golden bytes, every gate, the failure catalog, the status SQL and
|
||||
the workflow's ordering, all against fakes.
|
||||
|
||||
**End to end** — `pytest e2e -m import_e2e`, opt-in and outside `testpaths`. Runs the real workflow on
|
||||
a real worker against real MinIO, a real Postgres holding this document's `public.experiment_run`
|
||||
(with and without `V12`), a real MongoDB with the unique index on `models.id`, and a real
|
||||
`mlflow server`. Nothing on the import path is mocked. This is what proves the claims a fake cannot:
|
||||
that the `CHECK` constraints of § 6 accept every status written, that the registry source of § 11
|
||||
resolves through the existing loader, that the `V12` failure shape of § 7 is what it says, and that no
|
||||
plaintext password reaches the Temporal event history.
|
||||
|
||||
Requires Docker plus one one-time setup per machine:
|
||||
|
||||
```bash
|
||||
echo "ryuk.disabled=true" >> ~/.testcontainers.properties
|
||||
```
|
||||
|
||||
Without it every e2e test fails in fixture setup with a Docker mount error — see
|
||||
[`README.md` § End-to-End Tests](../README.md#end-to-end-tests-e2e) for why and what it costs.
|
||||
|
||||
Scenario catalogue: [`e2e/scenarios.md`](../e2e/scenarios.md) § 4.
|
||||
Requirement-to-test mapping: `openspec/changes/import-model-e2e-suite/traceability.md`.
|
||||
|
||||
Doing one import by hand, standing in for the frontend that does not exist yet:
|
||||
[`scripts/import_model_manual_run.py`](../scripts/import_model_manual_run.py) — see
|
||||
[`README.md` § Running One Model Import by Hand](../README.md#running-one-model-import-by-hand).
|
||||
100
docs/sientia-bundle-format.md
Normal file
100
docs/sientia-bundle-format.md
Normal file
@@ -0,0 +1,100 @@
|
||||
# The `.sientia` bundle format
|
||||
|
||||
A `.sientia` file is one exported MLflow model: its whole artifact tree, its run parameters and its
|
||||
origin metadata, zipped and then encrypted under a password the exporter chose.
|
||||
|
||||
**It is written by the Streamlit platform, `sientia-projects-templates`** (`app/src/operations/
|
||||
model_export/`). That application owns the export — its screen, its flow, its code — and documents
|
||||
it. This page is the short version: the parameters both sides have to agree on, and what this
|
||||
repository's reader accepts.
|
||||
|
||||
## How it is generated
|
||||
|
||||
The producer downloads the run's artifact tree whole, copies `run.data.params` **verbatim** (with
|
||||
`target_variable` among them, or the import fails at gate 6), writes `metadata.json` first into a
|
||||
`ZIP_DEFLATED` archive with the artifacts following as `artifacts/<relpath>` and no directory
|
||||
entries, encrypts that zip, and offers it as `{model_name}_v{model_version}.sientia`.
|
||||
|
||||
`metadata.json` has exactly two keys: `parameters` — the run's parameters verbatim — and `metadata`,
|
||||
six fields: `model_name`, `experiment_name`, `run_id`, `model_version`, `model_project`,
|
||||
`export_timestamp` (UTC ISO 8601).
|
||||
|
||||
## Encryption parameters
|
||||
|
||||
The reader derives the key with the values the header carries, but refuses anything above the
|
||||
producer's own cost — a header declaring more is a resource-exhaustion attempt, refused before any
|
||||
memory is allocated.
|
||||
|
||||
| | Value |
|
||||
|---|---|
|
||||
| KDF | Argon2id (`crypto_pwhash_ALG_ARGON2ID13`) |
|
||||
| memlimit | `268435456` (256 MiB) — a **ceiling** on the reader's side |
|
||||
| opslimit | `3` — likewise a ceiling |
|
||||
| parallelism | `1` (reserved, not configurable) |
|
||||
| salt | 16 bytes, fresh per bundle |
|
||||
| key | 32 bytes |
|
||||
| AEAD | XChaCha20-Poly1305 `secretstream`, 64 KiB plaintext chunks, last tagged `FINAL` |
|
||||
| digest | SHA-256 of the **pushed chunks only** — on disk, `sha256(file[164:])`, past the 140-byte header and the 24-byte stream header. Hashing from 140 rejects every valid bundle |
|
||||
| password | normalised `NFC`, encoded UTF-8 — no trimming, no case folding |
|
||||
|
||||
The export-side password rule (12+ characters, mixed case, digit, symbol) is enforced there only: by
|
||||
the time a bundle exists, the password is whatever it was, and the reader cannot re-check it.
|
||||
|
||||
## Shape of the file
|
||||
|
||||
`HEADER_STRUCT` is `>8sBBB16sQQB32s64s`. Big-endian, no alignment padding, so the offsets are exact:
|
||||
|
||||
```
|
||||
offset width field constraint the reader enforces
|
||||
------ ----- -------------------- ------------------------------------------------------
|
||||
0 8 magic == b'SIENTIA1'
|
||||
8 1 format_version == 1
|
||||
9 1 aead_id == 1 (XChaCha20-Poly1305 secretstream)
|
||||
10 1 kdf_id == 1 (Argon2id)
|
||||
11 16 salt passed to the KDF as-is
|
||||
27 8 kdf_memlimit_bytes <= 268435456
|
||||
35 8 kdf_opslimit <= 3
|
||||
43 1 kdf_parallelism == 1
|
||||
44 32 ciphertext_digest == sha256(file[164:])
|
||||
76 64 reserved_signature all zero
|
||||
------ -----
|
||||
140 24 secretstream header libsodium HEADERBYTES
|
||||
164 ... chunk stream 64 KiB plaintext chunks, each +17 bytes of AEAD tag
|
||||
```
|
||||
|
||||
`format_version` is checked for equality, not for a range: an unknown version is refused at gate 2
|
||||
rather than parsed optimistically. `reserved_signature` must be all zero today, so a signed bundle
|
||||
will be a new format version and not a silent change of meaning.
|
||||
|
||||
`laborious/utils/bundle/format.py` is the normative layout on this side, and the golden fixture is
|
||||
what keeps it and the producer from drifting apart.
|
||||
|
||||
## How it is read
|
||||
|
||||
`BundleReader.open()` runs seven ordered gates, each mapped to an import step and to one message a
|
||||
person can read:
|
||||
|
||||
| Gate | Step | Refuses |
|
||||
|---|---|---|
|
||||
| 1 | `download` | object too large, wrong prefix or suffix, digest ≠ the one the uploader sent |
|
||||
| 2 | `download` | header that is not this format or not this version |
|
||||
| 3 | `decryption` | wrong password, or altered bytes — the AEAD cannot tell them apart, and the message says so |
|
||||
| 4 | `archive_inspection` | zip bomb, absolute or `..` entries, symlinks — read from the central directory, before extracting |
|
||||
| 5 | `extraction` | anything that escapes the fresh work directory, or overwrites a file |
|
||||
| 6 | `structure_validation` | missing `metadata.json` fields, missing `target_variable`, missing `prediction_model`/`data_model` |
|
||||
| 7 | `content_policy` | a file that is not a known model file by name or suffix |
|
||||
|
||||
Gates 1 and 2 share the `download` step on purpose: neither may borrow the `decryption` sentence,
|
||||
which implicates the password. A gate-6 rejection naming `target_variable` means the origin run never
|
||||
logged it — the fix is a re-export, on the producer's side.
|
||||
|
||||
Error codes, sentences and the ceilings the gates read from configuration:
|
||||
[`model-import.md`](model-import.md).
|
||||
|
||||
## Fixtures
|
||||
|
||||
- `tests/fixtures/bundle/` — a real bundle emitted by the producer, with its password and producer
|
||||
commit recorded in that directory's `README.md`. It pays the producer's real Argon2id cost, so it
|
||||
is used where the real bytes matter and not in every test.
|
||||
- `tests/helpers/bundle_factory.py` — the test-only writer: same header at the same offsets, cheap
|
||||
KDF parameters, and a knob for every hostile variant the gates above refuse.
|
||||
343
e2e/conftest.py
343
e2e/conftest.py
@@ -2,9 +2,19 @@
|
||||
Pytest configuration and fixtures for E2E tests.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
# Load every BSON codec now, for its side effect only. `pymongo` imports `bson.objectid`,
|
||||
# `bson.timestamp`, `bson.decimal128` and a dozen more lazily, and the first time is while an
|
||||
# activity is running — after Temporal's workflow sandbox has snapshotted `sys.modules`. The sandbox
|
||||
# then warns once per submodule ("Module bson.timestamp was imported after initial workflow load"),
|
||||
# which was ~20 lines of noise per test that touches MongoDB, enough to bury a warning worth reading.
|
||||
# This has to happen **before** the workflow modules are imported below, which is why it is here and
|
||||
# not in `import_harness.py` — pytest loads `pytest_plugins` only after this module's body has run.
|
||||
import bson.json_util # noqa: F401
|
||||
|
||||
# E2E workflows under test do not run ModelAnalysis; stub before Activities import.
|
||||
_model_analysis_module = MagicMock()
|
||||
_model_analysis_module.ModelAnalysis = MagicMock
|
||||
@@ -12,28 +22,84 @@ sys.modules.setdefault('sientia', MagicMock())
|
||||
sys.modules.setdefault('sientia.ModelAnalysis', _model_analysis_module)
|
||||
from io import BytesIO
|
||||
|
||||
import joblib
|
||||
import pandas as pd
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sqlalchemy import create_engine, text
|
||||
from testcontainers.minio import MinioContainer
|
||||
from testcontainers.postgres import PostgresContainer
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
from testcontainers.community.minio import MinioContainer
|
||||
from testcontainers.community.postgres import PostgresContainer
|
||||
|
||||
from e2e.opc_test_server import OpcE2ETestServer
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
|
||||
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
|
||||
|
||||
# The `.sientia` import fixtures live in their own module rather than here: four containers' worth
|
||||
# of setup, an MLflow server subprocess and an ImportModel worker would double this file. Every
|
||||
# fixture it adds is `import_`-prefixed, so nothing below changes meaning.
|
||||
pytest_plugins = ['e2e.import_harness']
|
||||
|
||||
# Test constants
|
||||
TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
|
||||
TEST_DATABASE_NAME = 'test_db'
|
||||
|
||||
# Dedicated model served through the explicit `joblib` flavor, so no MLflow flavor loader runs.
|
||||
JOBLIB_FLAVOR_MODEL_NAME = 'joblib_flavor_model'
|
||||
|
||||
# Module carrying the classes of the joblib artifacts' models. It ships inside the artifact's
|
||||
# `code/` directory and is importable only while that directory is on `sys.path` — which is what
|
||||
# the joblib flavor arranges before deserializing.
|
||||
JOBLIB_FLAVOR_CODE_MODULE = 'joblib_flavor_probe'
|
||||
|
||||
JOBLIB_FLAVOR_WEIGHTS = {'feature_1': 2.0, 'feature_2': 3.0}
|
||||
|
||||
JOBLIB_FLAVOR_CODE_SOURCE = '''"""Model classes that exist only inside the artifact's `code/` directory."""
|
||||
|
||||
import pandas as pd
|
||||
|
||||
FEATURES = ('feature_1', 'feature_2')
|
||||
|
||||
|
||||
class JoblibFlavorPredictor:
|
||||
"""Prediction model: a weighted sum of the transform output columns."""
|
||||
|
||||
def __init__(self, weights):
|
||||
self.weights = weights
|
||||
|
||||
def fit(self, data):
|
||||
return self
|
||||
|
||||
def predict(self, data):
|
||||
frame = pd.DataFrame(data)
|
||||
return [
|
||||
sum(row[name] * self.weights[name] for name in FEATURES)
|
||||
for _, row in frame.iterrows()
|
||||
]
|
||||
|
||||
|
||||
class JoblibFlavorTransformer:
|
||||
"""Transform model: emits the fixed feature frame the prediction model expects."""
|
||||
|
||||
target_variable = 'target'
|
||||
|
||||
def fit(self, data):
|
||||
return self.predict(data)
|
||||
|
||||
def predict(self, data):
|
||||
frame = pd.DataFrame(data)
|
||||
result = pd.DataFrame(
|
||||
{'feature_1': [0.234] * len(frame), 'feature_2': [0.783] * len(frame)}
|
||||
)
|
||||
result.index = frame.index
|
||||
return result
|
||||
'''
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope='session')
|
||||
def minio_container():
|
||||
@@ -50,11 +116,15 @@ def minio_container():
|
||||
def postgres_container():
|
||||
"""
|
||||
Create a PostgreSQL container using testcontainers.
|
||||
|
||||
|
||||
This fixture creates a real PostgreSQL database in a Docker container
|
||||
that will be used for all tests in the session.
|
||||
|
||||
`log_statement=all` is the one thing the import suite needs from this fixture: proving an import
|
||||
issues no statement at all against the predictions database requires a log that would have
|
||||
shown one. It changes what the server writes to its own log and nothing a test observes.
|
||||
"""
|
||||
postgres = PostgresContainer('postgres:15')
|
||||
postgres = PostgresContainer('postgres:15').with_command('postgres -c log_statement=all')
|
||||
postgres.start()
|
||||
yield postgres
|
||||
postgres.stop()
|
||||
@@ -64,21 +134,21 @@ def postgres_container():
|
||||
def postgres_engine(postgres_container):
|
||||
"""
|
||||
Create SQLAlchemy engine for PostgreSQL test database.
|
||||
|
||||
|
||||
This fixture creates a connection to the PostgreSQL container
|
||||
created by the postgres_container fixture.
|
||||
"""
|
||||
engine = create_engine(postgres_container.get_connection_url())
|
||||
|
||||
|
||||
yield engine
|
||||
|
||||
|
||||
engine.dispose()
|
||||
|
||||
|
||||
def _create_schema_and_tables(engine):
|
||||
"""
|
||||
Helper function to create schema and tables in the given engine.
|
||||
|
||||
|
||||
Creates predictions_schema with:
|
||||
- laborious_data: Input data table for queries
|
||||
- predictions: Output predictions table
|
||||
@@ -87,8 +157,8 @@ def _create_schema_and_tables(engine):
|
||||
# Use begin() to ensure transaction is properly committed
|
||||
with engine.begin() as conn:
|
||||
# Create predictions_schema
|
||||
conn.execute(text("CREATE SCHEMA IF NOT EXISTS predictions_schema"))
|
||||
|
||||
conn.execute(text('CREATE SCHEMA IF NOT EXISTS predictions_schema'))
|
||||
|
||||
# Create laborious_data table (input data from sensors)
|
||||
create_laborious_data_sql = """
|
||||
CREATE TABLE IF NOT EXISTS predictions_schema.laborious_data (
|
||||
@@ -102,7 +172,7 @@ def _create_schema_and_tables(engine):
|
||||
);
|
||||
"""
|
||||
conn.execute(text(create_laborious_data_sql))
|
||||
|
||||
|
||||
# Create predictions table
|
||||
create_predictions_sql = """
|
||||
CREATE TABLE if not exists predictions_schema.predictions (
|
||||
@@ -119,7 +189,7 @@ def _create_schema_and_tables(engine):
|
||||
);
|
||||
"""
|
||||
conn.execute(text(create_predictions_sql))
|
||||
|
||||
|
||||
# Create transformed_data table
|
||||
create_transformed_sql = """
|
||||
CREATE TABLE IF NOT EXISTS predictions_schema.transformed_data (
|
||||
@@ -139,7 +209,7 @@ def _create_schema_and_tables(engine):
|
||||
def setup_postgres_schema_and_tables(postgres_engine):
|
||||
"""
|
||||
Automatically create necessary schema and tables before each test.
|
||||
|
||||
|
||||
This fixture runs automatically (autouse=True) and ensures
|
||||
that the predictions_schema and tables exist with the correct structure.
|
||||
"""
|
||||
@@ -150,35 +220,22 @@ def setup_postgres_schema_and_tables(postgres_engine):
|
||||
@pytest_asyncio.fixture
|
||||
def mock_logger():
|
||||
"""Mock logger for testing."""
|
||||
|
||||
def message(message):
|
||||
print(f"[LOG] {message}")
|
||||
def custom_message(message, _metadata={}):
|
||||
print(f"[LOG] {message}")
|
||||
print(f'[LOG] {message}')
|
||||
|
||||
def custom_message(message, _metadata=None):
|
||||
print(f'[LOG] {message}')
|
||||
|
||||
logger = MagicMock()
|
||||
logger.info = MagicMock(
|
||||
side_effect=message
|
||||
)
|
||||
logger.debug = MagicMock(
|
||||
side_effect=message
|
||||
)
|
||||
logger.error = MagicMock(
|
||||
side_effect=message
|
||||
)
|
||||
logger.warning = MagicMock(
|
||||
side_effect=message
|
||||
)
|
||||
logger.custom_info = MagicMock(
|
||||
side_effect=custom_message
|
||||
)
|
||||
logger.custom_debug = MagicMock(
|
||||
side_effect=custom_message
|
||||
)
|
||||
logger.custom_error = MagicMock(
|
||||
side_effect=custom_message
|
||||
)
|
||||
logger.custom_warning = MagicMock(
|
||||
side_effect=custom_message
|
||||
)
|
||||
logger.info = MagicMock(side_effect=message)
|
||||
logger.debug = MagicMock(side_effect=message)
|
||||
logger.error = MagicMock(side_effect=message)
|
||||
logger.warning = MagicMock(side_effect=message)
|
||||
logger.custom_info = MagicMock(side_effect=custom_message)
|
||||
logger.custom_debug = MagicMock(side_effect=custom_message)
|
||||
logger.custom_error = MagicMock(side_effect=custom_message)
|
||||
logger.custom_warning = MagicMock(side_effect=custom_message)
|
||||
return logger
|
||||
|
||||
|
||||
@@ -186,24 +243,24 @@ def mock_logger():
|
||||
def mock_mongo_client():
|
||||
"""
|
||||
Mock MongoDB client to avoid real connections.
|
||||
|
||||
|
||||
This fixture mocks the pymongo.MongoClient used by CoreNotificationHandler,
|
||||
allowing us to use a real NotificationHandler instance without connecting to MongoDB.
|
||||
"""
|
||||
mock_client = MagicMock()
|
||||
mock_db = MagicMock()
|
||||
mock_collection = MagicMock()
|
||||
|
||||
|
||||
# Configure the mock chain: client[database] -> db[collection] -> collection
|
||||
mock_client.__getitem__.return_value = mock_db
|
||||
mock_db.__getitem__.return_value = mock_collection
|
||||
|
||||
|
||||
# Mock server_info() to avoid connection attempts
|
||||
mock_client.server_info = MagicMock()
|
||||
|
||||
|
||||
# Mock insert_one for notifications
|
||||
mock_collection.insert_one = MagicMock()
|
||||
|
||||
|
||||
return mock_client
|
||||
|
||||
|
||||
@@ -225,7 +282,7 @@ def notification_inserts(mock_mongo_client):
|
||||
def notification_handler(mock_logger, mock_mongo_client):
|
||||
"""
|
||||
Create a real NotificationHandler instance with mocked MongoDB client.
|
||||
|
||||
|
||||
This fixture creates a real CoreNotificationHandler instance but mocks
|
||||
the underlying MongoDB connection to avoid real database connections.
|
||||
"""
|
||||
@@ -258,7 +315,7 @@ def mock_minio_repository():
|
||||
parquet_buffer = BytesIO()
|
||||
parquet_df.to_parquet(parquet_buffer, engine='pyarrow', index=True)
|
||||
parquet_bytes = parquet_buffer.getvalue()
|
||||
|
||||
|
||||
# sientia_do MinioRepository API
|
||||
mock_repo.bucket = 'test-bucket'
|
||||
mock_repo.upload_file = AsyncMock(
|
||||
@@ -273,7 +330,7 @@ def mock_minio_repository():
|
||||
mock_repo.list_objects = AsyncMock(return_value=[])
|
||||
mock_repo.delete_file = AsyncMock()
|
||||
mock_repo.close = MagicMock()
|
||||
|
||||
|
||||
return mock_repo
|
||||
|
||||
|
||||
@@ -295,13 +352,12 @@ def mock_pi_web_api_repository():
|
||||
mock_repo.close = MagicMock()
|
||||
return mock_repo
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
def mock_opc_repository():
|
||||
"""Mock OPC repository for OPC operations."""
|
||||
mock_repo = MagicMock()
|
||||
mock_repo.write_data = AsyncMock(
|
||||
return_value=(True, {'response_time': 0.1})
|
||||
)
|
||||
mock_repo.write_data = AsyncMock(return_value=(True, {'response_time': 0.1}))
|
||||
mock_repo.disconnect = AsyncMock()
|
||||
return mock_repo
|
||||
|
||||
@@ -318,10 +374,13 @@ async def opc_e2e_server():
|
||||
finally:
|
||||
await server.stop()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
def patch_create_engine(postgres_engine):
|
||||
"""Patch create_engine to return test postgres_engine."""
|
||||
with patch('sientia_do.temporal.activities.postgres.create_engine', return_value=postgres_engine):
|
||||
with patch(
|
||||
'sientia_do.temporal.activities.postgres.create_engine', return_value=postgres_engine
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@@ -329,44 +388,54 @@ def patch_create_engine(postgres_engine):
|
||||
def patch_minio_repository(mock_minio_repository):
|
||||
"""Patch MinioRepository to return mock."""
|
||||
# Patch where Activities resolves the symbol (import binds the original class).
|
||||
with patch('laborious.activities.activities.MinioRepository', return_value=mock_minio_repository):
|
||||
with patch(
|
||||
'laborious.activities.activities.MinioRepository', return_value=mock_minio_repository
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
def patch_pi_web_api_repository(mock_pi_web_api_repository):
|
||||
"""Patch MLflowRepository to return mock."""
|
||||
with patch('laborious.activities.api.PIWebAPIClient', return_value=mock_pi_web_api_repository):
|
||||
yield
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
def mock_mlflow_models():
|
||||
"""Create mock models for MLflow load_model methods."""
|
||||
# Mock transform model - returns DataFrame with same index as input
|
||||
mock_transform_model = MagicMock()
|
||||
|
||||
def mock_transform_predict(data):
|
||||
num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
|
||||
print(data.to_csv())
|
||||
print(data.index)
|
||||
result = pd.DataFrame({
|
||||
'feature_1': [0.234] * num_rows,
|
||||
'feature_2': [0.783] * num_rows,
|
||||
})
|
||||
result = pd.DataFrame(
|
||||
{
|
||||
'feature_1': [0.234] * num_rows,
|
||||
'feature_2': [0.783] * num_rows,
|
||||
}
|
||||
)
|
||||
result.index = data.index
|
||||
return result
|
||||
|
||||
mock_transform_model.predict = MagicMock(side_effect=mock_transform_predict)
|
||||
|
||||
|
||||
# Mock predict model - returns array/list of predictions
|
||||
mock_predict_model = MagicMock()
|
||||
|
||||
def mock_predict_predict(data):
|
||||
num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
|
||||
return [0.5] * num_rows
|
||||
|
||||
mock_predict_model.predict = MagicMock(side_effect=mock_predict_predict)
|
||||
|
||||
|
||||
# Mock PyFuncModel for compressed models
|
||||
mock_pyfunc_model = MagicMock()
|
||||
mock_pyfunc_model._model_impl = MagicMock()
|
||||
mock_pyfunc_model._model_impl.python_model = mock_transform_model
|
||||
|
||||
|
||||
return {
|
||||
'transform_model': mock_transform_model,
|
||||
'predict_model': mock_predict_model,
|
||||
@@ -374,19 +443,98 @@ def mock_mlflow_models():
|
||||
}
|
||||
|
||||
|
||||
def _write_artifact_code_dir(artifact_dir):
|
||||
"""Write the artifact's `code/` directory, carrying the module with the model classes."""
|
||||
code_dir = artifact_dir / 'code'
|
||||
code_dir.mkdir()
|
||||
(code_dir / f'{JOBLIB_FLAVOR_CODE_MODULE}.py').write_text(JOBLIB_FLAVOR_CODE_SOURCE)
|
||||
return code_dir
|
||||
|
||||
|
||||
def _build_joblib_flavor_artifact(artifact_dir, class_name, *args):
|
||||
"""
|
||||
Build one joblib artifact: `code/` with the class module, plus the compressed `model.pkl`.
|
||||
|
||||
Instantiating **and** dumping happen while `code/` is on `sys.path` — `pickle` refuses to
|
||||
serialize a class it cannot import. Both are undone on the way out, so from here the pickled
|
||||
class is unreachable until something puts the artifact's `code/` back on the path. That
|
||||
something is the flavor under test.
|
||||
"""
|
||||
code_dir = _write_artifact_code_dir(artifact_dir)
|
||||
|
||||
sys.path.insert(0, str(code_dir))
|
||||
try:
|
||||
module = importlib.import_module(JOBLIB_FLAVOR_CODE_MODULE)
|
||||
model = getattr(module, class_name)(*args)
|
||||
joblib.dump(model, artifact_dir / 'model.pkl', compress=3)
|
||||
finally:
|
||||
sys.path.remove(str(code_dir))
|
||||
sys.modules.pop(JOBLIB_FLAVOR_CODE_MODULE, None)
|
||||
importlib.invalidate_caches()
|
||||
|
||||
return {'path': str(artifact_dir), 'code_dir': str(code_dir), 'model': model}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def joblib_flavor_artifacts(tmp_path):
|
||||
"""
|
||||
Real artifact directories for the `joblib` flavor, one per model type.
|
||||
|
||||
Each holds a top-level `model.pkl` written by `joblib.dump(..., compress=3)` and a real
|
||||
`code/` directory with the module where the model's class lives. No `MLmodel` is written:
|
||||
the artifact this flavor reads — and the one a joblib retrain writes — has none.
|
||||
|
||||
Return:
|
||||
dict keyed by artifact name ('prediction_model', 'data_model'), each value carrying
|
||||
'path', 'code_dir' and 'model' (the object serialized, so tests can compute the
|
||||
expected prediction).
|
||||
"""
|
||||
original_sys_path = list(sys.path)
|
||||
|
||||
builds = {
|
||||
'prediction_model': ('JoblibFlavorPredictor', (JOBLIB_FLAVOR_WEIGHTS,)),
|
||||
'data_model': ('JoblibFlavorTransformer', ()),
|
||||
}
|
||||
|
||||
artifacts = {}
|
||||
|
||||
for artifact_name, (class_name, args) in builds.items():
|
||||
artifact_dir = tmp_path / artifact_name
|
||||
artifact_dir.mkdir()
|
||||
|
||||
artifacts[artifact_name] = _build_joblib_flavor_artifact(artifact_dir, class_name, *args)
|
||||
|
||||
yield artifacts
|
||||
|
||||
# The flavor prepends `code/` to `sys.path` and deliberately never removes it.
|
||||
sys.path[:] = original_sys_path
|
||||
sys.modules.pop(JOBLIB_FLAVOR_CODE_MODULE, None)
|
||||
importlib.invalidate_caches()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def joblib_flavor_artifact_without_pickle(tmp_path):
|
||||
"""Artifact directory with no top-level `.pkl`, so the joblib flavor fails on it."""
|
||||
artifact_dir = tmp_path / 'prediction_model_without_pickle'
|
||||
artifact_dir.mkdir()
|
||||
(artifact_dir / 'MLmodel').write_text('flavors:\n sklearn: {}\n')
|
||||
return str(artifact_dir)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
def patch_mlflow(mock_mlflow_models):
|
||||
def patch_mlflow(mock_mlflow_models, joblib_flavor_artifacts):
|
||||
"""Patch mlflow module in repository with load_model mocks."""
|
||||
mock_mlflow = MagicMock()
|
||||
|
||||
|
||||
# Mock sklearn.load_model
|
||||
def mock_sklearn_load_model(model_uri):
|
||||
if 'data_model' in model_uri or 'transform' in model_uri.lower():
|
||||
return mock_mlflow_models['transform_model']
|
||||
return mock_mlflow_models['predict_model']
|
||||
|
||||
mock_mlflow.sklearn = MagicMock()
|
||||
mock_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
|
||||
|
||||
|
||||
# Mock pyfunc.load_model
|
||||
def mock_pyfunc_load_model(model_uri):
|
||||
if 'artifacts' in model_uri or 'tmp' in model_uri:
|
||||
@@ -394,30 +542,61 @@ def patch_mlflow(mock_mlflow_models):
|
||||
if 'data_model' in model_uri or 'transform' in model_uri.lower():
|
||||
return mock_mlflow_models['transform_model']
|
||||
return mock_mlflow_models['predict_model']
|
||||
|
||||
mock_mlflow.pyfunc = MagicMock()
|
||||
mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model)
|
||||
|
||||
# Mock pytorch.load_model
|
||||
mock_mlflow.pytorch = MagicMock()
|
||||
mock_mlflow.pytorch.load_model = MagicMock(return_value=mock_mlflow_models['predict_model'])
|
||||
|
||||
|
||||
mock_mlflow.artifacts = MagicMock()
|
||||
|
||||
# Mock other mlflow methods that might be called
|
||||
mock_mlflow.set_tracking_uri = MagicMock()
|
||||
mock_mlflow.get_run = MagicMock(return_value=MagicMock(info=MagicMock(artifact_uri='mlflow-artifacts:/test_run_id')))
|
||||
mock_mlflow.get_run = MagicMock(
|
||||
return_value=MagicMock(info=MagicMock(artifact_uri='mlflow-artifacts:/test_run_id'))
|
||||
)
|
||||
mock_mlflow.tracking = MagicMock()
|
||||
mock_mlflow.tracking.MlflowClient = MagicMock(return_value=MagicMock(
|
||||
mock_mlflow_client = MagicMock(
|
||||
search_registered_models=MagicMock(return_value=[MagicMock(name='test_model')]),
|
||||
search_model_versions=MagicMock(return_value=[MagicMock(
|
||||
current_stage='Production',
|
||||
version='1',
|
||||
source='runs:/artifacts/test_run_id'
|
||||
)])
|
||||
))
|
||||
|
||||
search_model_versions=MagicMock(
|
||||
return_value=[
|
||||
MagicMock(
|
||||
current_stage='Production', version='1', source='runs:/artifacts/test_run_id'
|
||||
)
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
# `download_model` asks whether the run's artifact carries a `code` directory before it loads
|
||||
# anything, and downloads only that directory when it does. Answering "no" here keeps the
|
||||
# download counts of these tests about the model itself; a test that exercises the `code` step
|
||||
# overrides `list_artifacts` to answer "yes".
|
||||
def mock_list_artifacts(_run_id, artifact_path=None):
|
||||
if artifact_path:
|
||||
return []
|
||||
return [MagicMock(path=artifact_name) for artifact_name in joblib_flavor_artifacts]
|
||||
|
||||
mock_mlflow_client.list_artifacts = MagicMock(side_effect=mock_list_artifacts)
|
||||
|
||||
# The joblib flavor downloads a whole model artifact, one per model type — a single return
|
||||
# value would serve `data_model` from `prediction_model`. The `code` step asks for the nested
|
||||
# `<artifact>/code` instead, and gets that directory back, as MLflow would return it.
|
||||
def mock_download_artifacts(_run_id, artifact_name):
|
||||
if artifact_name.endswith('/code'):
|
||||
return joblib_flavor_artifacts[artifact_name.removesuffix('/code')]['code_dir']
|
||||
return joblib_flavor_artifacts[artifact_name]['path']
|
||||
|
||||
mock_mlflow_client.download_artifacts = MagicMock(side_effect=mock_download_artifacts)
|
||||
mock_mlflow.tracking.MlflowClient = MagicMock(return_value=mock_mlflow_client)
|
||||
|
||||
with patch('laborious.utils.repository.model_repository.mlflow', new=mock_mlflow):
|
||||
yield mock_mlflow
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mlflow_client(patch_mlflow):
|
||||
"""The MlflowClient mock the repository holds as `self.client`."""
|
||||
return patch_mlflow.tracking.MlflowClient.return_value
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope='function')
|
||||
async def test_activities(
|
||||
postgres_engine,
|
||||
@@ -430,7 +609,7 @@ async def test_activities(
|
||||
patch_minio_repository,
|
||||
patch_mlflow,
|
||||
patch_pi_web_api_repository,
|
||||
mock_opc_repository
|
||||
mock_opc_repository,
|
||||
):
|
||||
"""
|
||||
Create Activities instance with test dependencies.
|
||||
|
||||
65
e2e/fixtures/experiment_run.sql
Normal file
65
e2e/fixtures/experiment_run.sql
Normal file
@@ -0,0 +1,65 @@
|
||||
-- public.experiment_run as migrations V1 … V11 leave it, for the e2e Postgres container.
|
||||
--
|
||||
-- Source of truth: sientia-dataops-database-migrations,
|
||||
-- postgres/sientia-core-mlops-bff/migrations/ — this file mirrors the *final state* of:
|
||||
--
|
||||
-- V1 create_table_experiment_run (13 columns, one index)
|
||||
-- V2 update_check_constraint_on_status (named the status CHECK)
|
||||
-- V3 update_check_constraints (named every other CHECK)
|
||||
-- V4 update_status_column_size_and_add_orchestrator_validation_error (status -> VARCHAR(50))
|
||||
-- V5 migrate_mage_statuses_to_orchestrator (data only)
|
||||
-- V6 rename_mage_response_data_column (-> orchestrator_response_data)
|
||||
-- V7 migrate_mlflow_statuses_to_tracking (data only)
|
||||
-- V8 remove_orchestrator_request_sent_status (status vocabulary)
|
||||
-- V9 remove_file_upload_statuses (status vocabulary)
|
||||
-- V10 remove_unused_statuses (status vocabulary)
|
||||
-- V11 add_run_type_to_experiment_run (run_type + reindexed)
|
||||
--
|
||||
-- V12 (add_import_statuses_to_experiment_run) lives in `experiment_run_v12.sql`, applied on top of
|
||||
-- this file. It is separate because one e2e test — the evidence for why the migration gates
|
||||
-- enabling the import worker — has to run against a table that does *not* have it, and a second
|
||||
-- whole-schema copy with V12 removed would rot the moment either copy were edited.
|
||||
--
|
||||
-- The same drift trade already accepted for the bundle fixtures: this is a copy, and the tests
|
||||
-- that read it are what make the copy auditable. If a migration lands there, update this file and
|
||||
-- the version list above. Nothing in this repository owns the DDL — no import ever executes it
|
||||
-- outside the test suite.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.experiment_run (
|
||||
id BIGSERIAL CONSTRAINT pk_experiment_run PRIMARY KEY,
|
||||
experiment_name VARCHAR(50) NOT NULL,
|
||||
run_name VARCHAR(50),
|
||||
username VARCHAR(320) NOT NULL,
|
||||
status VARCHAR(50) NOT NULL,
|
||||
error_message VARCHAR(1024),
|
||||
created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL,
|
||||
updated_at TIMESTAMP WITHOUT TIME ZONE NOT NULL,
|
||||
bucket_name VARCHAR(63),
|
||||
file_name VARCHAR(1024),
|
||||
request_data JSONB NOT NULL,
|
||||
orchestrator_response_data JSONB,
|
||||
run_type VARCHAR(20) NOT NULL DEFAULT 'TRAIN',
|
||||
|
||||
CONSTRAINT chk_experiment_run_experiment_name_min_length CHECK (LENGTH(experiment_name) >= 3),
|
||||
CONSTRAINT chk_experiment_run_run_name_min_length CHECK (LENGTH(run_name) >= 3),
|
||||
CONSTRAINT chk_experiment_run_username_min_length CHECK (LENGTH(username) >= 3),
|
||||
CONSTRAINT chk_experiment_run_username_pattern
|
||||
CHECK (username ~* '^[a-zA-Z0-9_!#$%&''*+/=?`{|}~^.-]+@[a-zA-Z0-9.-]+$'),
|
||||
CONSTRAINT chk_experiment_run_error_message_min_length CHECK (LENGTH(error_message) >= 3),
|
||||
CONSTRAINT chk_experiment_run_bucket_name_min_length CHECK (LENGTH(bucket_name) >= 3),
|
||||
CONSTRAINT chk_experiment_run_bucket_name_pattern CHECK (bucket_name ~* '^[a-z0-9-]+$'),
|
||||
CONSTRAINT chk_experiment_run_file_name_min_length CHECK (LENGTH(file_name) >= 3),
|
||||
CONSTRAINT chk_experiment_run_file_name_pattern
|
||||
CHECK (file_name ~* '^[a-zA-Z0-9 /._-]+$'),
|
||||
|
||||
-- The V1…V11 status vocabulary: the two shared stage values plus the training terminals.
|
||||
-- V12 adds IMPORT_SUCCESS and IMPORT_ERROR to this list.
|
||||
CONSTRAINT chk_experiment_run_status CHECK (status IN ('PENDING', 'ORCHESTRATOR_REQUEST_ERROR',
|
||||
'ORCHESTRATOR_WAITING_PROC',
|
||||
'ORCHESTRATOR_VALIDATION_ERROR',
|
||||
'TRAINING_SUCCESS', 'TRAINING_ERROR')),
|
||||
CONSTRAINT chk_experiment_run_run_type CHECK (run_type IN ('TRAIN', 'IMPORT'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_experiment_run_user_exp_status_date_type
|
||||
ON public.experiment_run (username, experiment_name, status, created_at, run_type);
|
||||
36
e2e/fixtures/experiment_run_v12.sql
Normal file
36
e2e/fixtures/experiment_run_v12.sql
Normal file
@@ -0,0 +1,36 @@
|
||||
-- V12 add_import_statuses_to_experiment_run, as a delta applied on top of `experiment_run.sql`.
|
||||
--
|
||||
-- Source of truth: sientia-dataops-database-migrations,
|
||||
-- postgres/sientia-core-mlops-bff/migrations/V12__add_import_statuses_to_experiment_run.sql
|
||||
--
|
||||
-- Two things, and nothing else:
|
||||
--
|
||||
-- 1. the status vocabulary gains IMPORT_SUCCESS and IMPORT_ERROR;
|
||||
-- 2. chk_experiment_run_terminal_status_run_type pairs each terminal value with its own
|
||||
-- run_type, so a TRAIN row can never end IMPORT_SUCCESS and an IMPORT row can never end
|
||||
-- TRAINING_ERROR. The two shared stage values (PENDING, ORCHESTRATOR_WAITING_PROC) stay out
|
||||
-- of the pairing — both run types legitimately pass through them.
|
||||
--
|
||||
-- This file is applied separately because one e2e test runs an import against a table *without*
|
||||
-- it: that failure shape — provisioned, then unrecordable, refused on the first attempt — is the
|
||||
-- evidence for why this migration has to land before IMPORT_MODEL_ENABLED is turned on anywhere.
|
||||
|
||||
ALTER TABLE public.experiment_run
|
||||
DROP CONSTRAINT IF EXISTS chk_experiment_run_status;
|
||||
|
||||
ALTER TABLE public.experiment_run
|
||||
ADD CONSTRAINT chk_experiment_run_status CHECK (status IN ('PENDING',
|
||||
'ORCHESTRATOR_REQUEST_ERROR',
|
||||
'ORCHESTRATOR_WAITING_PROC',
|
||||
'ORCHESTRATOR_VALIDATION_ERROR',
|
||||
'TRAINING_SUCCESS', 'TRAINING_ERROR',
|
||||
'IMPORT_SUCCESS', 'IMPORT_ERROR'));
|
||||
|
||||
ALTER TABLE public.experiment_run
|
||||
DROP CONSTRAINT IF EXISTS chk_experiment_run_terminal_status_run_type;
|
||||
|
||||
ALTER TABLE public.experiment_run
|
||||
ADD CONSTRAINT chk_experiment_run_terminal_status_run_type CHECK (
|
||||
(run_type = 'IMPORT' OR status NOT IN ('IMPORT_SUCCESS', 'IMPORT_ERROR'))
|
||||
AND (run_type = 'TRAIN' OR status NOT IN ('TRAINING_SUCCESS', 'TRAINING_ERROR'))
|
||||
);
|
||||
@@ -11,7 +11,9 @@ from sqlalchemy import text
|
||||
from sqlalchemy.engine import Engine
|
||||
|
||||
|
||||
async def start_and_await_workflow(client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0):
|
||||
async def start_and_await_workflow(
|
||||
client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0
|
||||
):
|
||||
"""
|
||||
Start a workflow and wait for its result.
|
||||
|
||||
@@ -44,7 +46,9 @@ def insert_sample_data(postgres_engine: Engine, model_id: int, values: list[Any]
|
||||
values: Per-sensor values; use string 'NULL' for SQL NULL.
|
||||
"""
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
|
||||
)
|
||||
values_sql = []
|
||||
for i, value in enumerate(values):
|
||||
values_sql.append(f"""
|
||||
@@ -79,7 +83,6 @@ def assert_prediction(
|
||||
comments: Expected exact comments string (optional).
|
||||
comments_contains: Substring expected in comments when queued (optional).
|
||||
"""
|
||||
import pytest
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
result_query = conn.execute(
|
||||
@@ -90,7 +93,9 @@ def assert_prediction(
|
||||
)
|
||||
)
|
||||
prediction_rows = result_query.fetchall()
|
||||
assert len(prediction_rows) == 1, f'Expected one prediction record, got {len(prediction_rows)}'
|
||||
assert len(prediction_rows) == 1, (
|
||||
f'Expected one prediction record, got {len(prediction_rows)}'
|
||||
)
|
||||
row = prediction_rows[0]
|
||||
assert row[0] == model_id, f'Expected model_id={model_id}, got {row[0]}'
|
||||
assert row[1] == prediction or Decimal(str(row[1])) == Decimal(str(prediction)), (
|
||||
@@ -99,7 +104,9 @@ def assert_prediction(
|
||||
assert row[2] == prediction_confidence or Decimal(str(row[2])) == Decimal(
|
||||
str(prediction_confidence)
|
||||
), f'Expected prediction_confidence={prediction_confidence}, got {row[2]}'
|
||||
assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}"
|
||||
assert row[3] == prediction_status, (
|
||||
f"Expected prediction_status='{prediction_status}', got {row[3]}"
|
||||
)
|
||||
if comments is not None:
|
||||
assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
|
||||
if comments_contains is not None:
|
||||
@@ -135,7 +142,6 @@ def assert_continue(
|
||||
|
||||
def assert_stop(postgres_engine: Engine, model_id: int) -> None:
|
||||
"""Assert no prediction rows for model_id."""
|
||||
import pytest
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
result_query = conn.execute(
|
||||
@@ -156,7 +162,6 @@ def assert_repeat(postgres_engine: Engine, model_id: int, last_prediction: tuple
|
||||
model_id: Model id.
|
||||
last_prediction: Tuple (model_id, prediction, confidence, status) to match both rows.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
result_query = conn.execute(
|
||||
|
||||
1343
e2e/import_harness.py
Normal file
1343
e2e/import_harness.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -147,8 +147,11 @@ class OpcE2ETestServer:
|
||||
value = await self._confidence_node.read_value()
|
||||
return float(value)
|
||||
|
||||
async def _pre_write_callback(self, _event, _service) -> None:
|
||||
if self._session_bad_on_write:
|
||||
async def _pre_write_callback(self, event, _service) -> None:
|
||||
# Only client writes are faulted. asyncua's own clock task writes CurrentTime
|
||||
# through this same PreWrite dispatch, and raising there kills its time_task;
|
||||
# the stored exception then resurfaces as a teardown failure in server.stop().
|
||||
if self._session_bad_on_write and getattr(event, 'is_external', False):
|
||||
raise ServiceError(ua.StatusCodes.BadSessionIdInvalid)
|
||||
|
||||
|
||||
|
||||
667
e2e/scenarios.md
667
e2e/scenarios.md
@@ -4,12 +4,71 @@ This document describes all possible test scenarios for the `predictions_batch`
|
||||
|
||||
## Running automated E2E tests (`e2e/`)
|
||||
|
||||
### One-time setup: disable the testcontainers reaper
|
||||
|
||||
Do this once per machine, before running anything under `e2e/`:
|
||||
|
||||
```bash
|
||||
echo "ryuk.disabled=true" >> ~/.testcontainers.properties
|
||||
```
|
||||
|
||||
**Why.** Before starting any test container, testcontainers starts a reaper container ("Ryuk") whose
|
||||
job is to clean up containers if the test process dies without running its teardown. Ryuk bind-mounts
|
||||
the Docker socket. Under **Docker Desktop on Linux** the socket lives at
|
||||
`~/.docker/desktop/docker.sock`, which Docker Desktop's File Sharing does not expose, so the mount is
|
||||
denied and the reaper cannot start:
|
||||
|
||||
```
|
||||
docker.errors.APIError: 500 Server Error ... /containers/<id>/start:
|
||||
Mounts denied: The path /socket_mnt/home/<user>/.docker/desktop/docker.sock
|
||||
is not shared from the host and is not known to Docker.
|
||||
```
|
||||
|
||||
Every test then errors **at setup** — `ERROR`, not `FAILED`, and never inside an assertion. If you see
|
||||
that, this is why.
|
||||
|
||||
**What you give up.** Nothing in a normal run: every fixture in `conftest.py` and
|
||||
`import_harness.py` stops its own containers in teardown. Ryuk only covers the case where the process
|
||||
is killed without teardown (`kill -9`, a crash, a hard stop from the IDE). After one of those, clean
|
||||
up by hand:
|
||||
|
||||
```bash
|
||||
docker ps -a --filter "label=org.testcontainers=true" -q | xargs -r docker rm -f
|
||||
```
|
||||
|
||||
The setting is read by `testcontainers/core/config.py`, which accepts either `ryuk.disabled` in
|
||||
`~/.testcontainers.properties` or the `TESTCONTAINERS_RYUK_DISABLED` environment variable. The
|
||||
properties file is preferred here because it is per-machine and applies to every project and every
|
||||
shell, so nobody has to remember a prefix on the command line.
|
||||
|
||||
*Alternative, if you want to keep the reaper:* point Docker at the classic socket with
|
||||
`export DOCKER_HOST=unix:///var/run/docker.sock`. That only works if a native Docker Engine is
|
||||
running alongside Docker Desktop; if the Desktop daemon is the only one, disabling Ryuk is the way.
|
||||
|
||||
### Requirements and invocation
|
||||
|
||||
- **Runtime**: Docker (or a Docker-compatible daemon) must be available so [testcontainers](https://testcontainers.com/) can start **PostgreSQL** and **MinIO** containers.
|
||||
- **Dependencies**: install dev requirements (includes `testcontainers[postgres,minio]`).
|
||||
- **Dependencies**: install dev requirements (includes `testcontainers[postgres,minio,mongodb]`).
|
||||
- **Invocation**: run only integration-marked tests, for example: `pytest e2e/ -m integration`.
|
||||
- **MinIO tests**: `e2e/test_minio_offload.py` exercises real S3 uploads; other E2E modules continue to mock MinIO on the worker used by most scenarios.
|
||||
- **OPC tests (real server)**: `e2e/test_opc_real_server.py` uses an in-process **asyncua** server and real `OpcRepository` (`test_activities_real_opc`). Scenarios 3.1.2, 3.2.2, 3.2.4, and 3.2.5 are covered there. Other E2E modules keep the OPC mock.
|
||||
- Run only OPC real-server tests: `pytest e2e/test_opc_real_server.py -m "integration and opc"`.
|
||||
- **Model import tests**: the `import_model` suite is opt-in and has its own harness, containers and
|
||||
fidelity contract — see [§4.0](#40-running-the-import-e2e-tests). Run it with
|
||||
`pytest e2e -m import_e2e`.
|
||||
- **joblib flavor tests**: `e2e/test_predictions_batch_joblib_flavor.py` covers reading a model
|
||||
through the explicit `joblib` flavor (predict and transform), and
|
||||
`e2e/test_retrain_joblib_flavor.py` covers writing one back on a retrain — see
|
||||
[§5](#5-joblib-flavor-scenarios). Neither `joblib.dump` nor `joblib.load` is mocked in either
|
||||
module: the artifact fixture holds a real `joblib.dump(..., compress=3)` pickle whose class ships
|
||||
inside the artifact's `code/` directory, so only the flavor's own `sys.path` handling makes it
|
||||
deserialize. Run them with
|
||||
`pytest e2e/test_predictions_batch_joblib_flavor.py e2e/test_retrain_joblib_flavor.py -m integration`.
|
||||
- **Nothing under `e2e/` is covered by `validate`**: that script lints `laborious/` and `tests/`
|
||||
only, and its test step is `pytest tests/` (`testpaths = ["tests"]`). Lint these files by hand
|
||||
(`ruff format --check e2e/ && ruff check e2e/`), and remember that `addopts` carries
|
||||
`--strict-markers`: a test without `@pytest.mark.integration` is not collected, and a run that
|
||||
collects nothing still exits 0. Check the collected count, not the exit code.
|
||||
|
||||
## Workflow Overview
|
||||
|
||||
@@ -561,4 +620,608 @@ These paths do **not** rely on Temporal activity retries for export failures: th
|
||||
- Error notification sent
|
||||
- Workflow completes
|
||||
|
||||
---
|
||||
---
|
||||
## 4. Model Import Workflow Scenarios (`import_model`)
|
||||
|
||||
The `import_model` workflow turns an encrypted `.sientia` bundle already in MinIO into a registered
|
||||
MLflow model plus one MongoDB `models` document — or into a durable error on the row the BFF
|
||||
inserted. It shares no code path with the predictions workflows above and has its own harness.
|
||||
|
||||
### 4.0 Running the import E2E tests
|
||||
|
||||
- **Marker**: every test below carries `import_e2e`; the real-clock retry tests also carry `slow`.
|
||||
`testpaths` stays `["tests"]`, so none of this is collected by the unit gate.
|
||||
- **Invocation**: `pytest e2e -m import_e2e`.
|
||||
- **Runtime**: Docker for three containers (the shared MinIO, a dedicated Postgres for
|
||||
`sientia-core-mlops-bff`, MongoDB) **plus** an `mlflow server` subprocess started from this
|
||||
repository's venv. Dev requirements install `testcontainers[postgres,minio,mongodb]`.
|
||||
- **Fidelity**: on the import path nothing is mocked — real MinIO, real Postgres, real MongoDB with a
|
||||
unique index, a real MLflow tracking server with a working registry, real Argon2id and real
|
||||
`secretstream`. Only `CoreNotificationHandler`'s Mongo client stays mocked, because notifications
|
||||
are a side channel this suite counts rather than a dependency it tests.
|
||||
- **Fixtures**: `e2e/import_harness.py`, registered as a pytest plugin from `e2e/conftest.py`. Every
|
||||
fixture is `import_`-prefixed.
|
||||
- **Prerequisite**: the one-time reaper setup at the top of this document
|
||||
([Running automated E2E tests](#one-time-setup-disable-the-testcontainers-reaper)). Without it every
|
||||
test here errors at setup with a Docker mount error and never reaches an assertion.
|
||||
|
||||
### 4.1 Success Scenarios
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.1.1: Producer-Emitted Bundle Imports End to End
|
||||
**Description**: The committed golden `.sientia` file — bytes the producer's own code wrote — is
|
||||
imported at the producer's real Argon2id cost (256 MiB, 3 passes)
|
||||
|
||||
**Input**:
|
||||
- `tests/fixtures/bundle/golden_model_v3.sientia` uploaded under `imported_models/`
|
||||
- One `PENDING` / `run_type='IMPORT'` row carrying the object key and the uploader's digest
|
||||
- The workflow input: `import_run_id`, `bucket`, `object_key`, `expected_digest`, `password_envelope`
|
||||
|
||||
**Expected Behavior**:
|
||||
- The row is claimed, the seven gates pass, the experiment and a new run are created
|
||||
- The artifact tree is logged with `prediction_model` and `data_model` at the run's root
|
||||
- A registry version is created and left unstaged; the model document is written inactive
|
||||
- One terminal write sets `IMPORT_SUCCESS`
|
||||
|
||||
**Assertions**:
|
||||
- Exactly one run named `import-fixture_model-v3`, status `FINISHED`
|
||||
- `prediction_model` contains `MLmodel` and `model.pkl`
|
||||
- Bundle parameters logged, origin logged under `source.*`
|
||||
- Destination version is `1`, independent of the bundle's `model_version` of `3`
|
||||
- Model document `id` is the seeded maximum plus one; `active` is `false`
|
||||
- Row: `IMPORT_SUCCESS`, `error_message` null, both names filled
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.1.2: Factory-Built Bundle Imports End to End
|
||||
**Description**: The same assertions over a synthesised tree with names unique to the test
|
||||
|
||||
**Input**: A bundle from `tests/helpers/bundle_factory.py` at the cheap KDF cost
|
||||
|
||||
**Expected Behavior**: As 4.1.1
|
||||
|
||||
**Assertions**:
|
||||
- As 4.1.1, plus `model_config` holds exactly `transform_flavor`, `predict_flavor`,
|
||||
`retention_minutes` and `target`
|
||||
- `import.workflow_id` in the row's detail JSON matches the run
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.1.3: The Registry Source Resolves Through the Legacy Loader
|
||||
**Description**: The `source` the importer writes must be one the *existing* loader can parse
|
||||
|
||||
**Input**: A successful import's registered version
|
||||
|
||||
**Expected Behavior**: `MLFlowRepository.get_model_run_id` extracts the run id from the source
|
||||
|
||||
**Assertions**:
|
||||
- The source ends in `/prediction_model`
|
||||
- `get_model_run_id` returns the run the import created
|
||||
- Both artifact paths are present under that run
|
||||
- Called with the stage the version is actually in, because an import promotes nothing
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.1.4: Second Import of the Same Model
|
||||
**Description**: The same model imported twice, with different declared `model_version`s
|
||||
|
||||
**Input**: Two bundles declaring `model_version` `2` then `7`, same model and experiment names
|
||||
|
||||
**Expected Behavior**: Additive — a new run and a new registry version, nothing promoted, the
|
||||
existing model document untouched
|
||||
|
||||
**Assertions**:
|
||||
- Two versions (`1`, `2`) of one registered model; destination numbering ignores the bundle's
|
||||
- Two runs in the same experiment, each logging its own `source.model_version`
|
||||
- No version in `Production`
|
||||
- The model document is byte-identical to what the first import wrote
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.1.5: A Model Name Already in the Collection
|
||||
**Description**: An operator-owned document already carries the bundle's model name
|
||||
|
||||
**Input**: A document with `active: True`, `predict_flavor: pyfunc` and a hand-set target
|
||||
|
||||
**Expected Behavior**: The import notices the name and leaves the document alone; the import still
|
||||
succeeds
|
||||
|
||||
**Assertions**:
|
||||
- The collection is unchanged, exactly one document carries that name, no duplicate
|
||||
- Row: `IMPORT_SUCCESS` — a document that is already there is not a failure
|
||||
|
||||
---
|
||||
|
||||
### 4.2 Gate Rejection Scenarios
|
||||
|
||||
Every case asserts the same five things: terminal status `ORCHESTRATOR_VALIDATION_ERROR`, the
|
||||
recorded step, the recorded code, `error_message` **equal** to the catalog sentence for that step,
|
||||
and that no experiment, run, artifact, registry version, model document or extracted file remains.
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.2.1: Object Above the Size Ceiling (gate 1)
|
||||
**Description**: An object larger than `IMPORT_MAX_OBJECT_BYTES`, rejected from `stat_object` before
|
||||
its body is fetched
|
||||
|
||||
**Input**: A factory bundle against a ceiling narrowed to 1 KiB for this test class
|
||||
|
||||
**Expected Behavior**: Rejected at `download`
|
||||
|
||||
**Assertions**: The five shared assertions
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.2.2: Malformed Header (gate 2)
|
||||
**Description**: Five variants — wrong magic, unsupported format version, KDF memory above the
|
||||
ceiling, KDF passes above the ceiling, a header digest that does not describe the chunk stream
|
||||
|
||||
**Expected Behavior**: Each rejected at `download`, sharing gate 1's step
|
||||
|
||||
**Assertions**:
|
||||
- The five shared assertions
|
||||
- The recorded sentence does **not** mention the password: gates 1 and 2 must not implicate a
|
||||
password that is fine
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.2.3: Decryption Failures (gate 3)
|
||||
**Description**: Three variants — a wrong bundle password inside a valid envelope, a tampered chunk
|
||||
with the header digest repaired, and a stream whose every chunk is framed `MESSAGE`
|
||||
|
||||
**Expected Behavior**: Each rejected at `decryption`
|
||||
|
||||
**Assertions**:
|
||||
- The five shared assertions
|
||||
- The tampered case proves the file is self-consistent and only the AEAD can tell
|
||||
- The framing case proves the digest cannot see an incomplete stream
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.2.4: Hostile Archive (gate 4)
|
||||
**Description**: Four variants — an entry escaping the tree, an absolute entry path, a symbolic
|
||||
link, and a declared expansion over the ceiling
|
||||
|
||||
**Expected Behavior**: Each rejected at `archive_inspection`, from the central directory alone
|
||||
|
||||
**Assertions**:
|
||||
- The five shared assertions
|
||||
- **Nothing extracted**: the worker work directory is empty, not merely free of provisioning
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.2.5: A Member Extraction Cannot Place (gate 5)
|
||||
**Description**: An archive declaring the same member name twice — passes gate 4, which checks names
|
||||
and modes but not duplicates
|
||||
|
||||
**Expected Behavior**: Rejected at `extraction` on the second member
|
||||
|
||||
**Assertions**: The five shared assertions
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.2.6: Incomplete Structure (gate 6)
|
||||
**Description**: Five variants — no `metadata.json`, metadata missing `run_id`, parameters missing
|
||||
`target_variable`, no `prediction_model` directory, and a 60-character experiment name
|
||||
|
||||
**Expected Behavior**: Each rejected at `structure_validation`
|
||||
|
||||
**Assertions**:
|
||||
- The five shared assertions
|
||||
- For the over-long name: the row's `experiment_name` is still the inserted placeholder, so no write
|
||||
reached the `VARCHAR(50)` column, and that column is confirmed still 50 wide
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.2.7: Forbidden File Type (gate 7)
|
||||
**Description**: An `artifacts/payload.exe` member with a suffix outside the allow-list
|
||||
|
||||
**Expected Behavior**: Rejected at `content_policy`
|
||||
|
||||
**Assertions**: The five shared assertions
|
||||
|
||||
---
|
||||
|
||||
### 4.3 Digest Enforcement Scenarios
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.3.1: A Digest Computed Over Other Bytes
|
||||
**Description**: A perfectly valid bundle whose `expected_digest` belongs to a different file
|
||||
|
||||
**Expected Behavior**: Fails at `download`, before any header is parsed
|
||||
|
||||
**Assertions**:
|
||||
- Terminal status, step, code and catalog sentence
|
||||
- The failure names the digest, not the magic or the version
|
||||
- `run_name` is still null and `experiment_name` still the placeholder — the bundle was never opened
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.3.2: The Object Is Replaced After Upload
|
||||
**Description**: The same key, different bytes — also a valid bundle, so nothing but the digest can
|
||||
tell
|
||||
|
||||
**Expected Behavior**: Fails at `download`
|
||||
|
||||
**Assertions**: As 4.3.1; only a real object store can show this at all
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.3.3: No Digest at All
|
||||
**Description**: `expected_digest` absent from the input
|
||||
|
||||
**Expected Behavior**: Fails at `received`, before anything is downloaded
|
||||
|
||||
**Assertions**: The request-invalid code, and the work directory was never created
|
||||
|
||||
---
|
||||
|
||||
### 4.4 Provisioning Failure Scenarios
|
||||
|
||||
Each asserts the platform terminal status `IMPORT_ERROR` — not the validation status, because a
|
||||
provisioning failure is not the user's to fix — its own step and code, the workflow identifier in
|
||||
the detail JSON, and `error_message` **equal** to the catalog sentence.
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.4.1: The Tracking Server Is Stopped
|
||||
**Description**: The `mlflow server` process is genuinely terminated for the duration
|
||||
|
||||
**Expected Behavior**: Fails at `experiment_creation`
|
||||
|
||||
**Assertions**:
|
||||
- The shared assertions
|
||||
- The names are already on the row: the bundle was opened before MLflow was reached
|
||||
- Nothing in Mongo; nothing provisioned
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.4.2: Artifact Upload Fails
|
||||
**Description**: `log_artifacts` is patched — the one step where neither hostile input nor a broken
|
||||
dependency can isolate the failure
|
||||
|
||||
**Expected Behavior**: Fails at `artifact_upload`
|
||||
|
||||
**Assertions**:
|
||||
- The shared assertions
|
||||
- The experiment exists (the previous step succeeded); no version, no document
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.4.3: Registration Fails
|
||||
**Description**: `create_model_version` is patched. MLflow rejects only an empty registered-model
|
||||
name, and one process serves both tracking and the registry, so nothing else isolates this step
|
||||
|
||||
**Expected Behavior**: Fails at `registration`
|
||||
|
||||
**Assertions**:
|
||||
- The shared assertions
|
||||
- The run and both artifact paths survive — "the files were stored, the model was not registered"
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.4.4: The Model Document Is Refused
|
||||
**Description**: A real unique index on `model_config.target` refuses the activity's own insert.
|
||||
(A duplicate `id` cannot be constructed: the activity derives `max(int(id)) + 1` from the same
|
||||
collection, so the id it computes can never already exist)
|
||||
|
||||
**Expected Behavior**: Fails at `model_document`
|
||||
|
||||
**Assertions**:
|
||||
- The shared assertions
|
||||
- One registry version exists — registered, and unusable
|
||||
|
||||
---
|
||||
|
||||
### 4.5 Status Record Scenarios
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.5.1: The Claim Precedes Every Gate
|
||||
**Description**: The row is observed from inside `stat_object`, the first call gate 1 makes
|
||||
|
||||
**Assertions**: Running status and `import.workflow_id` already present at that moment
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.5.2: The Names Are Visible Mid-Flight
|
||||
**Description**: The row is observed from inside the first call experiment creation makes, during an
|
||||
import that then succeeds
|
||||
|
||||
**Assertions**: `experiment_name` and `run_name` already populated before anything exists in MLflow
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.5.3: The Inserting Side's Columns Are Never Rewritten
|
||||
**Description**: After a success and after a failure
|
||||
|
||||
**Assertions**: `username`, `request_data`, `bucket_name`, `file_name`, `run_type`, `created_at`
|
||||
unchanged; `status` and `updated_at` did change
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.5.4: Unclaimable Rows
|
||||
**Description**: Three variants — an unknown id, a `run_type='TRAIN'` row, and a row already held by
|
||||
another workflow id
|
||||
|
||||
**Expected Behavior**: The Temporal run ends failed with the not-claimable error type, and **nothing
|
||||
is written to any row**
|
||||
|
||||
**Assertions**: The error type from the history, the whole table byte-identical, nothing provisioned
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.5.5: The Row Names a Different Object
|
||||
**Description**: `file_name` disagrees with the input's `object_key`
|
||||
|
||||
**Expected Behavior**: Fails at `received` **with the error recorded on the row** — unlike 4.5.4,
|
||||
this row was successfully claimed
|
||||
|
||||
**Assertions**: The received code and sentence on the row; nothing provisioned
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.5.6: Constraint Conformance
|
||||
**Description**: The real constraints, exercised directly
|
||||
|
||||
**Assertions**:
|
||||
- Every status a normal import writes is accepted on an `IMPORT` row
|
||||
- All four terminal values are refused across run types, naming
|
||||
`chk_experiment_run_terminal_status_run_type`; the shared stage values are accepted on both
|
||||
- A training row still walks its own vocabulary
|
||||
- The table has exactly the 13 columns `V1`…`V12` leave (no `ALTER TABLE` by this feature)
|
||||
|
||||
---
|
||||
|
||||
### 4.6 Absence Scenarios
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.6.1: One Table, One Database
|
||||
**Description**: A full import bracketed by marker statements on both Postgres servers, read from
|
||||
`log_statement=all`
|
||||
|
||||
**Assertions**:
|
||||
- Every statement touching a user relation names `experiment_run`
|
||||
- No `project`, `models`, `import_run`, `model_project` or `model_registry` reference (with quoted
|
||||
literals stripped, so the bundle's own `model_project` *value* is not mistaken for a table)
|
||||
- The predictions database logs **no** statement at all, though the import holds working credentials
|
||||
for it
|
||||
- A rejected import — which writes more status rows than a successful one — behaves the same
|
||||
|
||||
---
|
||||
|
||||
### 4.7 Cleanup, Retention and History Scenarios
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.7.1: The Worker Filesystem Is Clean
|
||||
**Description**: After a success and after rejections at two different steps (decryption and
|
||||
structure validation, which unwind on different code paths)
|
||||
|
||||
**Assertions**: No extracted tree, no downloaded object, no temporary archive
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.7.2: The Retention Rule Is Present
|
||||
**Assertions**: `sientia-import-bundle-expiry` on the import bucket, enabled, filtered on
|
||||
`imported_models/`, expiring in 7 days — and the uploaded object still there, because the import
|
||||
never deletes it
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.7.3: The Record Outlives the Object
|
||||
**Description**: The uploaded object is deleted, as retention will eventually do
|
||||
|
||||
**Assertions**: The row still reads `IMPORT_SUCCESS` with its names
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.7.4: No Plaintext Secret in the Event History
|
||||
**Assertions**:
|
||||
- The bundle password, the base64 envelope key and the key's hex form appear nowhere in the history
|
||||
- `password_envelope` **is** in the history, so the negative assertions above cannot pass vacuously
|
||||
|
||||
---
|
||||
|
||||
### 4.8 Retry Envelope and Missing-`V12` Scenarios
|
||||
|
||||
These run on a **real clock** (`WorkflowEnvironment.start_local()`) with the status retry intervals
|
||||
shortened; the production values stay asserted by the unit suite. All are marked `slow`.
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.8.1: A Database That Comes Back
|
||||
**Description**: The status database is stopped for two seconds after provisioning and restarted
|
||||
inside the retry envelope
|
||||
|
||||
**Assertions**: The terminal write lands `IMPORT_SUCCESS` with the correct workflow id
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.8.2: A Broken Progress Update
|
||||
**Description**: A brief early outage that only the progress hints see
|
||||
|
||||
**Assertions**: The import completes with the correct verdict and does not spend the terminal
|
||||
write's envelope on a hint nobody reads
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.8.3: A Database That Stays Down After a Successful Import
|
||||
**Description**: The database is taken away as soon as the row is claimed and never comes back
|
||||
|
||||
**Expected Behavior**: The model is fully provisioned and the workflow **completes** — reporting
|
||||
failure would tell somebody to re-import a model that is already on the platform
|
||||
|
||||
**Assertions**:
|
||||
- The run returns its result; the model, version and document all exist
|
||||
- The failure channel carries `IMPORT_STATUS_NOT_RECORDED`, `step=status_write` and
|
||||
`verdict_not_recorded=`
|
||||
- The row is left in the running status
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.8.4: A Failed Import Whose Failure Cannot Be Written Either
|
||||
**Description**: The import fails at `artifact_upload` and the terminal write is then refused
|
||||
|
||||
**Assertions**:
|
||||
- The reported cause is the artifact-upload failure, not the status write
|
||||
- The status write reports itself only through the channel
|
||||
- The row is left running
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 4.8.5: A Table Without `V12`
|
||||
**Description**: The evidence for why the migration gates enabling the import worker
|
||||
|
||||
**Expected Behavior**: The claim succeeds (`V11` gave the table `run_type`; `PENDING` and the running
|
||||
status are pre-`V12` values), the model is fully provisioned, and only the terminal value is new — so
|
||||
the last write of the import is the first thing that fails, non-retryably
|
||||
|
||||
**Assertions**:
|
||||
- The run completes; the model, version and document all exist
|
||||
- The row is left in the running status, with both names filled
|
||||
- The channel carries the not-recorded code and the verdict
|
||||
- Elapsed time shows the refusal was **not** retried
|
||||
- `pg_constraint` confirms the fixture's `chk_experiment_run_status` really does not allow
|
||||
`IMPORT_SUCCESS`, so the test cannot pass vacuously
|
||||
|
||||
---
|
||||
|
||||
## 5. joblib Flavor Scenarios
|
||||
|
||||
`joblib` is an explicit model flavor, declared in `model_config.predict_flavor` /
|
||||
`model_config.transform_flavor` like any other. It reads the artifact directly — download the
|
||||
`.pkl`, put the artifact's `code/` on `sys.path`, `joblib.load` — and **never** calls an MLflow
|
||||
flavor loader. There is no fallback in either direction: a failing loader propagates.
|
||||
|
||||
Covered by `e2e/test_predictions_batch_joblib_flavor.py` (read) and
|
||||
`e2e/test_retrain_joblib_flavor.py` (write). Both use the `joblib_flavor_artifacts` fixture, which
|
||||
builds one artifact directory per model type holding a real compressed `model.pkl` plus a `code/`
|
||||
directory carrying the module where the model's class lives. The fixture takes that directory off
|
||||
`sys.path` again before yielding, so the pickle is **unloadable** until the flavor puts it back —
|
||||
that is what makes these tests prove module resolution rather than assume it. The fixture also
|
||||
restores `sys.path` in teardown, because the flavor deliberately never removes what it inserted.
|
||||
|
||||
### 5.1 Read Scenarios
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 5.1.1: Predict Through the joblib Flavor
|
||||
**Description**: `predict_flavor: 'joblib'` on a normal `predictions_batch` run
|
||||
|
||||
**Input**:
|
||||
- `model_config.predict_flavor = 'joblib'`, `transform_flavor = 'sklearn'`
|
||||
- `client.download_artifacts` serving the real `prediction_model` artifact
|
||||
|
||||
**Expected Behavior**:
|
||||
- The artifact is downloaded once and deserialized with `joblib.load`
|
||||
- No `models:/<name>/production` URI is built, so no flavor loader ever sees this model
|
||||
- The prediction the artifact's own object produces is what lands in Postgres
|
||||
|
||||
**Assertions**:
|
||||
- Exactly one prediction row, `prediction_status='Good'`, value matching the artifact's model
|
||||
- `download_artifacts` called once, with `'prediction_model'`
|
||||
- No `mlflow.{sklearn,pyfunc}.load_model` call names this model
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 5.1.2: Transform Through the joblib Flavor
|
||||
**Description**: `transform_flavor: 'joblib'`, which had no E2E coverage before
|
||||
|
||||
**Input**:
|
||||
- `model_config.transform_flavor = 'joblib'`, `predict_flavor = 'sklearn'`
|
||||
- `client.download_artifacts` serving the real `data_model` artifact
|
||||
|
||||
**Expected Behavior**:
|
||||
- `data_model` is downloaded and deserialized; the prediction model stays on its MLflow flavor
|
||||
- No artifact URI is resolved for the transform, so `mlflow.get_run` is never called
|
||||
|
||||
**Assertions**:
|
||||
- The prediction lands with the mocked prediction model's value
|
||||
- `download_artifacts` called once, with `'data_model'`
|
||||
- `mlflow.get_run` not called; no flavor loader was given a `/data_model` URI
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 5.1.3: An MLflow Flavor Is Unaffected
|
||||
**Description**: A model declaring `sklearn` behaves exactly as before
|
||||
|
||||
**Input**:
|
||||
- `client.list_artifacts` answering that the run's artifacts carry no `code` directory
|
||||
|
||||
**Assertions**:
|
||||
- The prediction lands with the mocked model's value
|
||||
- `download_artifacts` NOT called — the joblib path is never entered as a safety net, and with no
|
||||
`code` directory to fetch there is nothing else to download either
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 5.1.4: Only the Artifact's `code` Directory Is Downloaded
|
||||
**Description**: A model on an MLflow flavor whose artifact **does** carry `code/`. The flavor
|
||||
loaders read the model URI themselves and fetch what they need, so the `code` directory — the one
|
||||
thing MLflow will not put on `sys.path` unless the `MLmodel` manifest declares it — is the whole of
|
||||
what this workflow downloads through the client
|
||||
|
||||
**Input**:
|
||||
- `model_config` on the default `sklearn` flavors
|
||||
- `client.list_artifacts` answering that both artifacts carry `code`
|
||||
|
||||
**Expected Behavior**:
|
||||
- The `code` directory of each artifact is downloaded before its model is loaded, and goes on
|
||||
`sys.path`
|
||||
- Neither model's weights travel through `client.download_artifacts`
|
||||
|
||||
**Assertions**:
|
||||
- The prediction lands with the mocked model's value
|
||||
- `download_artifacts` called exactly twice, with `'data_model/code'` and `'prediction_model/code'`
|
||||
|
||||
---
|
||||
|
||||
### 5.2 Error Scenarios
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 5.2.1: joblib Flavor Failure Has a Single Cause
|
||||
**Description**: The downloaded artifact has no top-level `.pkl`
|
||||
|
||||
**Expected Behavior**: The workflow stops. The reported cause is the missing pickle and nothing
|
||||
else — with no fallback there is no second exception chained to a first
|
||||
|
||||
**Assertions**:
|
||||
- The run stops (`assert_stop`)
|
||||
- The log carries `No .pkl file found in prediction_model artifact`
|
||||
- The log does **not** carry `Original error`, the marker of a chained cause
|
||||
|
||||
---
|
||||
|
||||
### 5.3 Write Scenarios
|
||||
|
||||
---
|
||||
|
||||
#### Scenario 5.3.1: Retraining a joblib Model Round-Trips
|
||||
**Description**: `retrain_model` on a model whose `predict_flavor` and `transform_flavor` are both
|
||||
`joblib` — write and read in one test
|
||||
|
||||
**Input**:
|
||||
- Training frame indexed by `DATETIME_FORMAT_WITH_TZ` timestamps, carrying the target column
|
||||
- MLflow's artifact I/O replaced by a fake store that copies into an inspectable run directory
|
||||
|
||||
**Expected Behavior**:
|
||||
- Both models load through the joblib flavor; no flavor loader runs
|
||||
- `log_model` writes `prediction_model/model.pkl` (a `joblib.dump(..., compress=3)` pickle) and
|
||||
re-logs `prediction_model/code/`, without an `MLmodel` manifest
|
||||
- Loading the **written** artifact through the joblib flavor deserializes it and predicts, after
|
||||
the module the first load imported has been dropped from `sys.modules`
|
||||
|
||||
**Assertions**:
|
||||
- `retrain_model` returns `success: True`
|
||||
- `download_artifacts` was asked for both `data_model` and `prediction_model`
|
||||
- `sklearn`/`pyfunc` loaders have zero calls
|
||||
- The written run has `prediction_model/model.pkl`, `prediction_model/code/<module>.py` and
|
||||
`data_model/model.pkl`, and **no** `prediction_model/MLmodel`
|
||||
- The reloaded model's prediction equals the original's, and the written `code/` is on `sys.path`
|
||||
|
||||
---
|
||||
|
||||
@@ -26,7 +26,9 @@ async def test_format_and_export_prediction_default_path_e2e(
|
||||
client = temporal_test_env.client
|
||||
model_id = 401
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
||||
)
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
|
||||
196
e2e/test_import_harness_smoke.py
Normal file
196
e2e/test_import_harness_smoke.py
Normal file
@@ -0,0 +1,196 @@
|
||||
"""Does the harness itself come up? Diagnosable separately from the suite that runs on it.
|
||||
|
||||
Four containers, an `mlflow server` subprocess and a Temporal worker is a lot of moving parts, and
|
||||
when one of them fails to start the first *real* test fails somewhere deep inside an activity with
|
||||
a message about something else. These tests exist so that a broken harness says so in its own
|
||||
words: each one checks one piece is alive and does nothing with the import path.
|
||||
|
||||
The last test is the exception — it runs a real workflow, but with an input designed to be rejected
|
||||
at the very first step, so what it proves is only that the wiring reaches the first activity.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import mlflow
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import (
|
||||
BFF_DATABASE,
|
||||
IMPORT_ACTIVITY_NAMES,
|
||||
IMPORT_BUCKET,
|
||||
IMPORT_PREFIX,
|
||||
IMPORT_TASK_QUEUE,
|
||||
SEEDED_MODEL_IDS,
|
||||
failure_text,
|
||||
failure_types,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
|
||||
class TestContainers:
|
||||
"""Each dependency is real, reachable, and the one the harness says it is."""
|
||||
|
||||
async def test_bff_container_holds_the_status_table(self, import_status_table: Any) -> None:
|
||||
with import_status_table.begin() as connection:
|
||||
database = connection.execute(text('SELECT current_database()')).scalar_one()
|
||||
columns = {
|
||||
row[0]
|
||||
for row in connection.execute(
|
||||
text(
|
||||
'SELECT column_name FROM information_schema.columns '
|
||||
"WHERE table_schema = 'public' AND table_name = 'experiment_run'"
|
||||
)
|
||||
)
|
||||
}
|
||||
assert database == BFF_DATABASE
|
||||
assert {'id', 'status', 'run_type', 'request_data', 'orchestrator_response_data'} <= columns
|
||||
|
||||
async def test_bff_container_logs_every_statement(self, import_bff_container: Any) -> None:
|
||||
"""`log_statement=all` — without it the absence assertions have no vantage point."""
|
||||
result = import_bff_container.exec(
|
||||
['psql', '-U', 'bff', '-d', BFF_DATABASE, '-tAc', 'SHOW log_statement']
|
||||
)
|
||||
assert b'all' in result.output
|
||||
|
||||
async def test_the_v12_delta_is_what_separates_the_two_tables(
|
||||
self, import_status_table_pre_v12: Any
|
||||
) -> None:
|
||||
"""The pre-`V12` table has `run_type` and refuses the import terminal statuses."""
|
||||
with import_status_table_pre_v12.begin() as connection:
|
||||
constraints = {
|
||||
row[0]
|
||||
for row in connection.execute(
|
||||
text(
|
||||
'SELECT conname FROM pg_constraint '
|
||||
"WHERE conrelid = 'public.experiment_run'::regclass"
|
||||
)
|
||||
)
|
||||
}
|
||||
columns = {
|
||||
row[0]
|
||||
for row in connection.execute(
|
||||
text(
|
||||
'SELECT column_name FROM information_schema.columns '
|
||||
"WHERE table_schema = 'public' AND table_name = 'experiment_run'"
|
||||
)
|
||||
)
|
||||
}
|
||||
assert 'run_type' in columns
|
||||
assert 'chk_experiment_run_terminal_status_run_type' not in constraints
|
||||
|
||||
async def test_the_seeded_collection_has_gapped_ids_and_a_unique_index(
|
||||
self, import_models_collection: Any
|
||||
) -> None:
|
||||
ids = {
|
||||
str(document['id']) for document in import_models_collection.find({}, {'_id': False})
|
||||
}
|
||||
assert ids == set(SEEDED_MODEL_IDS)
|
||||
unique_keys = {
|
||||
tuple(index['key'].keys())
|
||||
for index in import_models_collection.list_indexes()
|
||||
if index.get('unique')
|
||||
}
|
||||
assert ('id',) in unique_keys
|
||||
|
||||
async def test_the_import_bucket_exists_and_its_prefix_is_empty(
|
||||
self, import_bucket: Any
|
||||
) -> None:
|
||||
assert import_bucket.bucket_exists(IMPORT_BUCKET)
|
||||
assert not list(
|
||||
import_bucket.list_objects(IMPORT_BUCKET, prefix=IMPORT_PREFIX, recursive=True)
|
||||
)
|
||||
|
||||
async def test_uploading_a_bundle_returns_a_key_and_its_digest(
|
||||
self, import_upload_bundle: Any, import_bucket: Any
|
||||
) -> None:
|
||||
object_key, digest = import_upload_bundle(b'not a real bundle, just bytes')
|
||||
assert object_key.startswith(IMPORT_PREFIX)
|
||||
assert len(digest) == 64
|
||||
assert import_bucket.stat_object(IMPORT_BUCKET, object_key).size > 0
|
||||
|
||||
async def test_the_mlflow_server_answers_and_serves_a_registry(
|
||||
self, import_mlflow_server: dict[str, Any]
|
||||
) -> None:
|
||||
client = mlflow.tracking.MlflowClient(tracking_uri=import_mlflow_server['tracking_uri'])
|
||||
experiment_id = client.create_experiment('harness_smoke_experiment')
|
||||
run = client.create_run(experiment_id, run_name='harness-smoke-run')
|
||||
client.create_registered_model('harness_smoke_model')
|
||||
version = client.create_model_version(
|
||||
'harness_smoke_model', f'runs:/{run.info.run_id}/prediction_model', run.info.run_id
|
||||
)
|
||||
assert version.version == '1'
|
||||
assert version.current_stage in (None, 'None')
|
||||
|
||||
|
||||
class TestActivitiesAndWorker:
|
||||
"""The object under test is built from the real dependencies, and the worker registers it."""
|
||||
|
||||
async def test_activities_point_at_the_real_dependencies(
|
||||
self, import_activities: Any, import_mlflow_server: dict[str, Any]
|
||||
) -> None:
|
||||
assert mlflow.get_tracking_uri() == import_mlflow_server['tracking_uri']
|
||||
assert import_activities.import_config['bucket'] == IMPORT_BUCKET
|
||||
assert import_activities.import_config['prefix'] == IMPORT_PREFIX
|
||||
assert import_activities.status_config['dbname'] == BFF_DATABASE
|
||||
assert import_activities.mongo_repository is not None
|
||||
# Not a mock: a MagicMock would satisfy any attribute access, so the type is the assertion.
|
||||
assert type(import_activities.import_minio_repository).__name__ == 'MinioRepository'
|
||||
|
||||
async def test_every_import_activity_is_registered(self, import_activities: Any) -> None:
|
||||
for name in IMPORT_ACTIVITY_NAMES:
|
||||
assert callable(getattr(import_activities, name)), f'{name} is missing'
|
||||
|
||||
async def test_the_worker_polls_the_import_queue(self, import_worker: Any) -> None:
|
||||
assert import_worker.task_queue == IMPORT_TASK_QUEUE
|
||||
|
||||
async def test_a_workflow_input_reaches_the_first_activity(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_execute: Any,
|
||||
import_build_input: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""The claim runs, and the row it could not claim proves the wiring is live.
|
||||
|
||||
An id no row carries is the cheapest input that gets past workflow validation and into
|
||||
`claim_import_status`, so a run that ends failed *there* means the client, the worker, the
|
||||
workflow, the activity and the BFF connection are all in place. Nothing else is exercised.
|
||||
"""
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=987654321,
|
||||
object_key=f'{IMPORT_PREFIX}nothing.sientia',
|
||||
expected_digest='0' * 64,
|
||||
password_envelope='not-a-real-envelope',
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
text_of_failure = failure_text(failure.value)
|
||||
assert 'claim' in text_of_failure.lower()
|
||||
assert '987654321' in text_of_failure
|
||||
assert 'ImportLogRowNotClaimableError' in failure_types(failure.value)
|
||||
assert import_read_status_row(987654321) is None
|
||||
|
||||
async def test_the_insert_helper_writes_a_claimable_row(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
object_key = f'{IMPORT_PREFIX}claimable.sientia'
|
||||
row_id = import_insert_status_row(object_key=object_key, expected_digest='a' * 64)
|
||||
row = import_read_status_row(row_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'PENDING'
|
||||
assert row['run_type'] == 'IMPORT'
|
||||
assert row['file_name'] == object_key
|
||||
request_data = row['request_data']
|
||||
payload = json.loads(request_data) if isinstance(request_data, str) else request_data
|
||||
assert payload['expected_digest'] == 'a' * 64
|
||||
184
e2e/test_import_model_absence.py
Normal file
184
e2e/test_import_model_absence.py
Normal file
@@ -0,0 +1,184 @@
|
||||
"""One table, one database — proven from the servers' own logs, not from the application's side.
|
||||
|
||||
Decision D22 says an import writes to `public.experiment_run` and to nothing else relational: no
|
||||
project table, no model table, no row in the predictions database. That claim is only worth making
|
||||
if it is checked from somewhere that sees *every* client, and there is exactly one such place.
|
||||
|
||||
Why not the obvious alternatives:
|
||||
|
||||
- a SQLAlchemy `before_cursor_execute` listener sees the status engine and is blind to the psycopg
|
||||
pool `Storage` holds — which is precisely the client that would betray the claim;
|
||||
- `pg_stat_statements` is aggregated, so it cannot be scoped to the window of one import.
|
||||
|
||||
So both containers run with `log_statement=all` and the window is bracketed by a marker statement on
|
||||
each server. The assertions are about *which relation names appear*, not about exact statement text,
|
||||
because the text is the driver's business and changes with it.
|
||||
|
||||
The predictions database is deliberately real and reachable here: `import_activities` is built with
|
||||
working credentials for it, so "no statement reached it" is a claim about the import's behaviour and
|
||||
not about a missing connection.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import names_no_relation, strip_literals
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
|
||||
# Relations that exist in the predictions database and in nobody's import path.
|
||||
FOREIGN_RELATIONS = (
|
||||
'predictions_schema',
|
||||
'laborious_data',
|
||||
'transformed_data',
|
||||
'predictions',
|
||||
)
|
||||
|
||||
# Tables a reader might expect an import to want, and which decision D22 says do not exist.
|
||||
RELATIONS_D22_DENIES = ('project', 'model_registry', 'models', 'import_run', 'model_project')
|
||||
|
||||
|
||||
class TestOneTableOneDatabase:
|
||||
"""A full import, from claim to terminal write, watched from both servers."""
|
||||
|
||||
async def test_every_statement_the_import_issues_names_experiment_run(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_log_window: Any,
|
||||
) -> None:
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
|
||||
with import_log_window() as captured:
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS', (
|
||||
'the window has to bracket a *working* import, or it proves nothing'
|
||||
)
|
||||
|
||||
bff_statements = captured['bff']
|
||||
assert bff_statements, 'the status database logged nothing at all — the window is wrong'
|
||||
|
||||
touching_a_relation = [
|
||||
statement for statement in bff_statements if not names_no_relation(statement)
|
||||
]
|
||||
assert touching_a_relation, 'no relational statement was logged, so nothing was checked'
|
||||
for statement in touching_a_relation:
|
||||
assert 'experiment_run' in statement, (
|
||||
f'a statement naming another relation reached the status database: {statement!r}'
|
||||
)
|
||||
|
||||
# And no relation a reader might expect the feature to want appears anywhere in the window.
|
||||
# Literals are stripped first: the import writes the bundle's own `model_project` value into
|
||||
# the JSON detail, and `factory_project` is data, not a reference to a `project` table.
|
||||
window_text = '\n'.join(strip_literals(s) for s in bff_statements).lower()
|
||||
for relation in RELATIONS_D22_DENIES:
|
||||
assert relation not in window_text, (
|
||||
f'{relation!r} appears in the import window: D22 says no such table is written'
|
||||
)
|
||||
for relation in FOREIGN_RELATIONS:
|
||||
assert relation not in window_text, f'{relation!r} appears in the import window'
|
||||
|
||||
async def test_the_predictions_database_sees_no_statement_from_an_import(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_log_window: Any,
|
||||
) -> None:
|
||||
"""Not "no statement against our tables" — no statement at all.
|
||||
|
||||
The import holds a working connection to this server (`Activities` is built with one), so
|
||||
this is about what the import chooses to do, not about what it is able to do.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
|
||||
with import_log_window() as captured:
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
await handle.result()
|
||||
|
||||
assert import_read_status_row(import_run_id)['status'] == 'IMPORT_SUCCESS'
|
||||
|
||||
predictions_statements = [
|
||||
statement for statement in captured['predictions'] if not names_no_relation(statement)
|
||||
]
|
||||
assert not predictions_statements, (
|
||||
f'the import issued statements against the predictions database: '
|
||||
f'{predictions_statements}'
|
||||
)
|
||||
|
||||
async def test_a_rejected_import_also_touches_nothing_else(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_log_window: Any,
|
||||
) -> None:
|
||||
"""The failure path writes more status rows than the success path, so it is worth its own case.
|
||||
|
||||
A rejected import runs the claim, the progress hints, cleanup and the terminal write — more
|
||||
statements than a success, and every one of them still has to name the same table.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory(include_metadata=False))
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
|
||||
with import_log_window() as captured:
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
for statement in captured['bff']:
|
||||
if names_no_relation(statement):
|
||||
continue
|
||||
assert 'experiment_run' in statement, statement
|
||||
assert not [s for s in captured['predictions'] if not names_no_relation(s)]
|
||||
247
e2e/test_import_model_cleanup.py
Normal file
247
e2e/test_import_model_cleanup.py
Normal file
@@ -0,0 +1,247 @@
|
||||
"""What the worker is left holding, what the bucket is left promising, and what the history reveals.
|
||||
|
||||
Three things an import must be true about after it finishes, whichever way it finished:
|
||||
|
||||
- **The worker's disk is empty of it.** No extracted tree, no downloaded object, no decrypted zip.
|
||||
The bundle is decrypted onto local disk, so a leak here is a plaintext model sitting in a scratch
|
||||
directory for the next import to trip over — and after a *failure* is exactly when that would go
|
||||
unnoticed.
|
||||
- **The bucket carries the retention rule.** The import never deletes the uploaded object: it is the
|
||||
record of an attempt, and the lifecycle rule is what eventually removes it. That rule is applied
|
||||
by the importer rather than assumed from a Helm value, so it is checkable from here.
|
||||
- **The status record outlives the object**, because it is the audit trail and the object is not.
|
||||
|
||||
And the one claim that can only be made against a real Temporal history: no plaintext password
|
||||
anywhere in it. The workflow carries an envelope it cannot open, and the activity that opens it
|
||||
never returns what it found. That is the design; this reads the history and checks it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import IMPORT_BUCKET, IMPORT_PREFIX
|
||||
from tests.helpers.bundle_factory import DEFAULT_ENVELOPE_KEY, envelope_key_b64
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
RETENTION_RULE_ID = 'sientia-import-bundle-expiry'
|
||||
|
||||
|
||||
async def _run(
|
||||
*,
|
||||
payload: bytes,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
password: str = FACTORY_PASSWORD,
|
||||
) -> tuple[Any, int, str]:
|
||||
object_key, digest = import_upload_bundle(payload)
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(password),
|
||||
)
|
||||
)
|
||||
return handle, import_run_id, object_key
|
||||
|
||||
|
||||
class TestWorkerFilesystem:
|
||||
"""Nothing of the import is left on the worker — on the success path or on either failure path."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('case', 'factory_kwargs', 'expect_success'),
|
||||
[
|
||||
('success', {}, True),
|
||||
('rejected at decryption', {'drop_final_tag': True}, False),
|
||||
('rejected at structure validation', {'include_metadata': False}, False),
|
||||
],
|
||||
)
|
||||
async def test_the_work_directory_is_empty_afterwards(
|
||||
self,
|
||||
case: str,
|
||||
factory_kwargs: dict[str, Any],
|
||||
expect_success: bool,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
tmp_path: Any,
|
||||
) -> None:
|
||||
"""Two different failure steps, because they clean up at different points.
|
||||
|
||||
A decryption failure leaves a partially written zip in the staging directory; a structure
|
||||
failure has already extracted a whole tree. The reader unwinds both, and the cleanup activity
|
||||
removes the scratch directory after either — but they are separate code paths, so asserting
|
||||
one would not cover the other.
|
||||
"""
|
||||
handle, _, _ = await _run(
|
||||
payload=import_bundle_factory(**factory_kwargs),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
if expect_success:
|
||||
await handle.result()
|
||||
else:
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
work_dir = tmp_path / 'import-work'
|
||||
leftovers = sorted(str(path) for path in work_dir.rglob('*')) if work_dir.exists() else []
|
||||
assert not leftovers, f'{case}: the worker still holds {leftovers}'
|
||||
|
||||
|
||||
class TestRetention:
|
||||
"""The rule the uploaded object's eventual removal depends on."""
|
||||
|
||||
async def test_the_import_bucket_carries_the_retention_rule(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_bucket: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
) -> None:
|
||||
handle, _, object_key = await _run(
|
||||
payload=import_bundle_factory(),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await handle.result()
|
||||
|
||||
lifecycle = import_bucket.get_bucket_lifecycle(IMPORT_BUCKET)
|
||||
assert lifecycle is not None, 'the bucket carries no lifecycle configuration at all'
|
||||
rules = {rule.rule_id: rule for rule in lifecycle.rules}
|
||||
assert RETENTION_RULE_ID in rules, (
|
||||
f'no {RETENTION_RULE_ID!r} rule; the bucket has {sorted(rules)}'
|
||||
)
|
||||
rule = rules[RETENTION_RULE_ID]
|
||||
assert rule.status == 'Enabled'
|
||||
assert rule.rule_filter.prefix == IMPORT_PREFIX
|
||||
assert rule.expiration.days == 7
|
||||
|
||||
# The import did not delete the uploaded object: it is the record of an attempt, and the
|
||||
# rule above is what removes it later.
|
||||
assert import_bucket.stat_object(IMPORT_BUCKET, object_key).size > 0
|
||||
|
||||
async def test_the_status_record_outlives_the_uploaded_object(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_bucket: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""Deleting the object is what the retention rule will eventually do; the verdict must survive.
|
||||
|
||||
The record is the audit trail. If reading a finished import's outcome depended on the bundle
|
||||
still being in the bucket, every import would become unreadable after seven days.
|
||||
"""
|
||||
handle, import_run_id, object_key = await _run(
|
||||
payload=import_bundle_factory(),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await handle.result()
|
||||
|
||||
import_bucket.remove_object(IMPORT_BUCKET, object_key)
|
||||
with pytest.raises(Exception): # noqa: B017 - MinIO's own not-found type is not the subject
|
||||
import_bucket.stat_object(IMPORT_BUCKET, object_key)
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS'
|
||||
assert row['file_name'] == object_key
|
||||
assert row['run_name'] is not None
|
||||
|
||||
|
||||
class TestEventHistorySecrecy:
|
||||
"""The claim that only a real event history can settle."""
|
||||
|
||||
async def test_no_plaintext_password_or_key_appears_in_the_history(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_history_events: Any,
|
||||
) -> None:
|
||||
"""Dump every event of a completed import and search the whole thing.
|
||||
|
||||
Three values must be absent, and they are absent for three different reasons:
|
||||
|
||||
- the **bundle password** never crosses an activity boundary — `open_import_bundle` decrypts
|
||||
the envelope inside itself and returns metadata, never the password;
|
||||
- the **envelope key** is configuration, read from the environment inside the activity, so it
|
||||
is never an argument to anything;
|
||||
- the **decrypted password** is the first of these two by another name, and is searched for
|
||||
separately because a future refactor that returned it "just for logging" would pass a test
|
||||
that only looked for the literal input.
|
||||
|
||||
The envelope itself *is* in the history, and must be: it is how the password reaches the
|
||||
activity. That it is there is asserted too, so this test cannot pass by finding nothing at
|
||||
all — which is the way an assertion like this usually rots.
|
||||
"""
|
||||
password = 'Secret1Password!InHistory'
|
||||
handle, _, _ = await _run(
|
||||
payload=import_bundle_factory(password=password),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
password=password,
|
||||
)
|
||||
await handle.result()
|
||||
|
||||
events, blob = await import_history_events(handle)
|
||||
assert events, 'the history is empty'
|
||||
|
||||
assert password not in blob, 'the bundle password appears in the event history'
|
||||
assert envelope_key_b64(DEFAULT_ENVELOPE_KEY) not in blob, (
|
||||
'the envelope key appears in the event history'
|
||||
)
|
||||
assert DEFAULT_ENVELOPE_KEY.hex() not in blob.lower(), (
|
||||
'the envelope key appears in the event history, hex-encoded'
|
||||
)
|
||||
|
||||
# The negative assertions above are only meaningful if the history is actually readable and
|
||||
# actually carries the envelope. Without this, an empty or opaque history would "pass".
|
||||
assert 'password_envelope' in blob, (
|
||||
'the history does not even carry the envelope, so the search above proved nothing'
|
||||
)
|
||||
214
e2e/test_import_model_digest.py
Normal file
214
e2e/test_import_model_digest.py
Normal file
@@ -0,0 +1,214 @@
|
||||
"""The digest is a promise about the bytes in the bucket, and it is kept against the bucket.
|
||||
|
||||
Three ways the promise can be broken, and where each is caught:
|
||||
|
||||
- the uploader computed the digest over *other* bytes — caught at `download`, after the object is
|
||||
fetched and before its header is parsed, so a wrong file never costs an Argon2id pass;
|
||||
- the object was replaced between the upload and the run — the same failure, which is the point:
|
||||
the check is against what is stored now, not against what was stored when the row was written;
|
||||
- no digest at all — caught at `received`, before anything is downloaded, because a request with no
|
||||
digest is not a request this workflow can honour.
|
||||
|
||||
The first two matter because MinIO is real here. A mocked object store would return whatever the
|
||||
test told it to, and "the digest is enforced against the bytes that were uploaded" would be a
|
||||
statement about the mock.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import IMPORT_BUCKET, failure_text
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from laborious.utils.models.import_errors import IMPORT_FAILURE_CATALOG
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
VALIDATION_ERROR_STATUS = 'ORCHESTRATOR_VALIDATION_ERROR'
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> dict[str, Any]:
|
||||
return json.loads(value) if isinstance(value, str) else (value or {})
|
||||
|
||||
|
||||
def _assert_step(row: dict[str, Any], step: ImportStep, workflow_id: str) -> None:
|
||||
code, reason, _responsibility = IMPORT_FAILURE_CATALOG[step]
|
||||
detail = _as_dict(row['orchestrator_response_data'])['import']
|
||||
assert row['status'] == VALIDATION_ERROR_STATUS, row['status']
|
||||
assert detail['step'] == step.value, detail
|
||||
assert detail['code'] == code.value, detail
|
||||
assert detail['workflow_id'] == workflow_id
|
||||
assert row['error_message'] == reason
|
||||
|
||||
|
||||
class TestWrongDigest:
|
||||
"""A digest computed over bytes that are not the ones in the bucket."""
|
||||
|
||||
async def test_a_digest_from_another_file_fails_at_download(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""The bundle is perfectly valid; only the digest is another file's.
|
||||
|
||||
The row carries the same wrong digest as the input, so the request and the record agree and
|
||||
the `received` guard passes — otherwise this would be a `received` failure about a
|
||||
contradictory request instead of the download-step digest check it is meant to exercise.
|
||||
"""
|
||||
payload = import_bundle_factory()
|
||||
other_digest = hashlib.sha256(
|
||||
import_bundle_factory(metadata={'run_id': 'other'})
|
||||
).hexdigest()
|
||||
assert other_digest != hashlib.sha256(payload).hexdigest()
|
||||
|
||||
object_key, _ = import_upload_bundle(payload)
|
||||
import_run_id = import_insert_status_row(
|
||||
object_key=object_key, expected_digest=other_digest
|
||||
)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=other_digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
_assert_step(row, ImportStep.DOWNLOAD, handle.id)
|
||||
|
||||
# No header was parsed: the failure names the digest, not the format. If the digest check
|
||||
# ran *after* the header the message here would be about magic bytes or a version.
|
||||
text_of_failure = failure_text(failure.value).lower()
|
||||
assert 'digest' in text_of_failure, text_of_failure
|
||||
assert 'marker' not in text_of_failure and 'version' not in text_of_failure
|
||||
|
||||
# The names are still unwritten, which is the other half of "no header was parsed": the
|
||||
# bundle was never opened, so nothing knew what to call it.
|
||||
assert row['run_name'] is None
|
||||
assert row['experiment_name'] == 'import-e2e-placeholder'
|
||||
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
|
||||
class TestOverwrittenObject:
|
||||
"""The bytes changed after the row was written. Only a real object store can show this."""
|
||||
|
||||
async def test_replacing_the_object_after_upload_fails_at_download(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_bucket: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
original = import_bundle_factory()
|
||||
object_key, digest = import_upload_bundle(original)
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
|
||||
# Same key, different bytes — also a valid bundle, so nothing but the digest can tell.
|
||||
replacement = import_bundle_factory(metadata={'run_id': 'ff00ff00ff00ff00ff00ff00ff00ff00'})
|
||||
assert hashlib.sha256(replacement).hexdigest() != digest
|
||||
import_bucket.put_object(
|
||||
IMPORT_BUCKET,
|
||||
object_key,
|
||||
io.BytesIO(replacement),
|
||||
length=len(replacement),
|
||||
content_type='application/octet-stream',
|
||||
)
|
||||
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
_assert_step(row, ImportStep.DOWNLOAD, handle.id)
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
|
||||
class TestMissingDigest:
|
||||
"""No digest is refused at the door, before a single byte is fetched."""
|
||||
|
||||
async def test_an_absent_digest_fails_at_received_with_nothing_downloaded(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
tmp_path: Any,
|
||||
) -> None:
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=None,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
_assert_step(row, ImportStep.RECEIVED, handle.id)
|
||||
assert 'expected_digest' in failure_text(failure.value)
|
||||
|
||||
# Nothing downloaded: the work directory was never even created.
|
||||
work_dir = tmp_path / 'import-work'
|
||||
assert not work_dir.exists() or not list(work_dir.rglob('*'))
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
679
e2e/test_import_model_gates.py
Normal file
679
e2e/test_import_model_gates.py
Normal file
@@ -0,0 +1,679 @@
|
||||
"""One rejection per gate, against a real object, a real archive and a real filesystem.
|
||||
|
||||
The gates are unit-tested already. What only a real run can show is that a rejection is *recorded*
|
||||
correctly and that it provisions nothing: the row ends `ORCHESTRATOR_VALIDATION_ERROR` naming that
|
||||
gate's step, its catalog code and its catalog sentence, and no experiment, run, artifact, registry
|
||||
version, model document or extracted file survives.
|
||||
|
||||
Every case here is hostile **input** — design E7's first preference. Nothing is patched, no
|
||||
dependency is broken; the bundles are built by the shared factory and the gates do the rest. The
|
||||
one exception to "input only" is the object-size ceiling, where the ceiling itself is narrowed
|
||||
rather than a gigabyte being uploaded; that is a configuration change, not a patched code path, and
|
||||
it is scoped to the one class that needs it.
|
||||
|
||||
`_assert_rejected` is where the shared claim of task 5.8 lives — terminal status, recorded step,
|
||||
recorded code, catalog sentence, and `assert_nothing_provisioned` — so no case can quietly assert
|
||||
four fifths of it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import (
|
||||
IMPORT_BUCKET,
|
||||
IMPORT_PREFIX,
|
||||
failure_text,
|
||||
)
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from laborious.utils.models.import_errors import IMPORT_FAILURE_CATALOG
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
VALIDATION_ERROR_STATUS = 'ORCHESTRATOR_VALIDATION_ERROR'
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> dict[str, Any]:
|
||||
return json.loads(value) if isinstance(value, str) else (value or {})
|
||||
|
||||
|
||||
async def _run_rejected_import(
|
||||
*,
|
||||
payload: bytes,
|
||||
password: str,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
envelope_password: str | None = None,
|
||||
) -> tuple[Any, int, str]:
|
||||
"""Upload a hostile bundle, insert the row, run the import and expect it to fail."""
|
||||
object_key, digest = import_upload_bundle(payload)
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(
|
||||
envelope_password if envelope_password is not None else password
|
||||
),
|
||||
)
|
||||
)
|
||||
return handle, import_run_id, object_key
|
||||
|
||||
|
||||
async def _assert_rejected(
|
||||
*,
|
||||
handle: Any,
|
||||
import_run_id: int,
|
||||
object_key: str,
|
||||
step: ImportStep,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
import_names: dict[str, str],
|
||||
experiment_name: str | None = None,
|
||||
model_name: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""The whole of task 5.8, in one call. Returns the row so a case can add to it."""
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None, 'the rejection lost the row it was recorded on'
|
||||
assert row['status'] == VALIDATION_ERROR_STATUS, (
|
||||
f'terminal status is {row["status"]!r}, expected {VALIDATION_ERROR_STATUS!r}'
|
||||
)
|
||||
|
||||
expected_code, expected_reason, _expected_responsibility = IMPORT_FAILURE_CATALOG[step]
|
||||
detail = _as_dict(row['orchestrator_response_data'])['import']
|
||||
assert detail['step'] == step.value, (
|
||||
f'recorded step is {detail["step"]!r}, expected {step.value!r}'
|
||||
)
|
||||
assert detail['code'] == expected_code.value, (
|
||||
f'recorded code is {detail["code"]!r}, expected {expected_code.value!r}'
|
||||
)
|
||||
assert detail['workflow_id'] == handle.id
|
||||
assert row['error_message'] == expected_reason, (
|
||||
'error_message is not the catalog sentence for this step'
|
||||
)
|
||||
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=experiment_name or import_names['experiment_name'],
|
||||
model_name=model_name or import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
return row
|
||||
|
||||
|
||||
class TestGate1ObjectLimits:
|
||||
"""Gate 1 rejects before the object's body is fetched at all."""
|
||||
|
||||
@pytest.fixture
|
||||
def import_limits_override(self) -> dict[str, str]:
|
||||
"""A 1 KiB ceiling, so a factory bundle is over it.
|
||||
|
||||
The production ceiling is 1 GiB. Narrowing it is the only way to cross gate 1 without
|
||||
uploading a gigabyte, and it changes a configured limit rather than a code path — the same
|
||||
comparison runs, against a different number.
|
||||
"""
|
||||
return {'IMPORT_MAX_OBJECT_BYTES': '1024'}
|
||||
|
||||
async def test_an_oversized_object_is_rejected_at_download(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
payload = import_bundle_factory()
|
||||
assert len(payload) > 1024, 'the fixture bundle must exceed the narrowed ceiling'
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=payload,
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.DOWNLOAD,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
|
||||
class TestGate2Header:
|
||||
"""Four header rejections, all landing on `download` — the step that says "not these bytes"."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('case', 'header_kwargs'),
|
||||
[
|
||||
('wrong magic', {'magic': b'NOTSIENT'}),
|
||||
('unsupported version', {'format_version': 9}),
|
||||
('kdf memory above the ceiling', {'declared_memlimit': 512 * 1024 * 1024}),
|
||||
('kdf passes above the ceiling', {'declared_opslimit': 9}),
|
||||
('digest not of the chunk stream', {'digest_override': bytes(32)}),
|
||||
],
|
||||
)
|
||||
async def test_a_malformed_header_is_rejected_at_download(
|
||||
self,
|
||||
case: str,
|
||||
header_kwargs: dict[str, Any],
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(**header_kwargs),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
row = await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.DOWNLOAD,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
# None of these may borrow the decryption sentence: the password is not implicated.
|
||||
assert 'password' not in (row['error_message'] or '').lower(), case
|
||||
|
||||
|
||||
class TestGate3Decryption:
|
||||
"""The password, the tag and the framing — the three ways a stream refuses to open."""
|
||||
|
||||
async def test_a_wrong_password_is_rejected_at_decryption(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""The envelope opens — it is the *bundle* password inside it that is wrong.
|
||||
|
||||
Worth separating: an envelope the worker cannot open is also a decryption-step failure, but
|
||||
for a different reason, and only this one proves the KDF actually ran and produced a key
|
||||
the stream rejected.
|
||||
"""
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(password=FACTORY_PASSWORD),
|
||||
password=FACTORY_PASSWORD,
|
||||
envelope_password='Wrong1Password!Entirely',
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.DECRYPTION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
async def test_a_tampered_chunk_is_rejected_at_decryption(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""One flipped byte inside the stream, with the header digest repaired to match.
|
||||
|
||||
Repairing the digest is what makes this gate 3 rather than gate 2: the file is
|
||||
self-consistent, so only the AEAD can tell it was altered.
|
||||
"""
|
||||
from tests.helpers.bundle_factory import tamper_ciphertext
|
||||
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=tamper_ciphertext(import_bundle_factory()),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.DECRYPTION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
async def test_a_stream_without_its_final_tag_is_rejected_at_decryption(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""Every chunk framed as MESSAGE, so the stream claims more is coming and never ends.
|
||||
|
||||
The digest cannot see this — the bytes on disk are exactly the ones it covers. Only the
|
||||
`secretstream` framing knows the file is incomplete.
|
||||
"""
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(drop_final_tag=True),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.DECRYPTION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
|
||||
class TestGate4ArchiveInspection:
|
||||
"""The central directory alone decides — nothing has been written to disk yet."""
|
||||
|
||||
@pytest.fixture
|
||||
def import_limits_override(self) -> dict[str, str]:
|
||||
"""A 1 MiB declared-expansion ceiling, for the oversized-expansion case in this class."""
|
||||
return {'IMPORT_MAX_UNCOMPRESSED_BYTES': str(1024 * 1024)}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('case', 'kwargs'),
|
||||
[
|
||||
('entry escaping the tree', {'traversal_entry': 'artifacts/../../escaped.csv'}),
|
||||
('absolute entry path', {'absolute_entry': '/opt/absolute.csv'}),
|
||||
('symbolic link', {'symlinks': {'artifacts/link.csv': 'artifacts/train_data.csv'}}),
|
||||
('declared expansion over the ceiling', {'bomb_bytes': 4 * 1024 * 1024}),
|
||||
],
|
||||
)
|
||||
async def test_a_hostile_archive_is_rejected_before_anything_is_extracted(
|
||||
self,
|
||||
case: str,
|
||||
kwargs: dict[str, Any],
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
tmp_path: Any,
|
||||
) -> None:
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(**kwargs),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.ARCHIVE_INSPECTION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
# "Nothing extracted" is stronger than "nothing provisioned": not one member reached disk.
|
||||
work_dir = tmp_path / 'import-work'
|
||||
assert not list(work_dir.rglob('*')) if work_dir.exists() else True, case
|
||||
|
||||
|
||||
class TestGate5Extraction:
|
||||
"""A member that gate 4 passed and extraction still cannot place.
|
||||
|
||||
Gate 5 is about *placement*, not size: containment on the resolved path, and never overwriting
|
||||
a file already written. A zip carrying the same name twice passes gate 4 — which checks names
|
||||
and modes, not duplicates — and is refused here on the second member.
|
||||
"""
|
||||
|
||||
# `zipfile` warns while *writing* the duplicate entry. That is the fixture doing exactly what
|
||||
# this test needs, so it is silenced here rather than left to look like a problem — and scoped to
|
||||
# this one test, so a duplicate name appearing anywhere else still gets noticed.
|
||||
@pytest.mark.filterwarnings('ignore:Duplicate name:UserWarning')
|
||||
async def test_a_duplicate_member_is_rejected_at_extraction(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
from tests.helpers.bundle_factory import default_tree
|
||||
|
||||
# `bomb_bytes` writes `artifacts/bomb.csv`, and `files` already carries that name, so the
|
||||
# archive declares it twice. Small enough that no size ceiling is involved.
|
||||
files = dict(default_tree())
|
||||
files['artifacts/bomb.csv'] = b'first\n'
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(files=files, bomb_bytes=64),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.EXTRACTION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
|
||||
class TestGate6Structure:
|
||||
"""The layout and the metadata schema, checked on the extracted tree."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('case', 'kwargs'),
|
||||
[
|
||||
('no metadata.json', {'include_metadata': False}),
|
||||
('metadata missing run_id', {'drop_metadata_field': 'run_id'}),
|
||||
('parameters missing target_variable', {'drop_parameter': 'target_variable'}),
|
||||
],
|
||||
)
|
||||
async def test_an_incomplete_bundle_is_rejected_at_structure_validation(
|
||||
self,
|
||||
case: str,
|
||||
kwargs: dict[str, Any],
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(**kwargs),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.STRUCTURE_VALIDATION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
async def test_a_bundle_without_prediction_model_is_rejected(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""Without `prediction_model` there is nothing to register a version against."""
|
||||
from tests.helpers.bundle_factory import default_tree
|
||||
|
||||
files = {
|
||||
name: content
|
||||
for name, content in default_tree().items()
|
||||
if not name.startswith('artifacts/prediction_model/')
|
||||
}
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(files=files),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.STRUCTURE_VALIDATION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
async def test_an_over_long_experiment_name_never_reaches_the_varchar_column(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
import_bff_engine: Any,
|
||||
) -> None:
|
||||
"""A 60-character experiment name is refused at gate 6, not by Postgres three steps later.
|
||||
|
||||
`experiment_name` is `VARCHAR(50)` and the import writes the bundle's own name onto the row
|
||||
*before* provisioning. If the gate did not enforce the column's limit, the names write would
|
||||
be a `DataError` — non-retryable, recorded as a status failure, on a row whose import was
|
||||
otherwise fine. Asserting the row's `experiment_name` is untouched is what proves the write
|
||||
never happened.
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
long_name = 'x' * 60
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(metadata={'experiment_name': long_name}),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
row = await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.STRUCTURE_VALIDATION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
experiment_name=long_name,
|
||||
)
|
||||
assert row['experiment_name'] == 'import-e2e-placeholder', (
|
||||
'the over-long name reached the row, so the gate did not hold the column'
|
||||
)
|
||||
with import_bff_engine.begin() as connection:
|
||||
width = connection.execute(
|
||||
text(
|
||||
'SELECT character_maximum_length FROM information_schema.columns '
|
||||
"WHERE table_name = 'experiment_run' AND column_name = 'experiment_name'"
|
||||
)
|
||||
).scalar_one()
|
||||
assert width == 50, 'the column this gate exists to protect is no longer VARCHAR(50)'
|
||||
|
||||
|
||||
class TestGate7ContentPolicy:
|
||||
"""A file kind an MLflow artifact tree has no reason to carry."""
|
||||
|
||||
async def test_a_forbidden_file_type_is_rejected_at_content_policy(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(extra_files={'artifacts/payload.exe': b'MZ\x00\x00'}),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.CONTENT_POLICY,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
|
||||
class TestTheRejectionSaysNothingTechnical:
|
||||
"""`error_message` is a sentence for a person: no path, no object key, no traceback.
|
||||
|
||||
One case is enough to make the point structurally — the catalog is keyed by step and
|
||||
`classify_import_failure` discards the exception's text entirely — and the per-gate cases above
|
||||
each assert the sentence is *exactly* the catalog's, which is the same guarantee stated
|
||||
positively.
|
||||
"""
|
||||
|
||||
async def test_no_technical_detail_reaches_the_recorded_reason(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(include_metadata=False),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
message = row['error_message'] or ''
|
||||
for forbidden in (
|
||||
object_key,
|
||||
IMPORT_BUCKET,
|
||||
IMPORT_PREFIX,
|
||||
'experiment_run',
|
||||
'Traceback',
|
||||
'BundleGateError',
|
||||
tempfile.gettempdir(),
|
||||
'gate',
|
||||
):
|
||||
assert forbidden not in message, f'{forbidden!r} leaked into error_message: {message!r}'
|
||||
|
||||
# The technical truth is not lost — it is on the Temporal error, where support can read it.
|
||||
assert 'metadata' in failure_text(failure.value).lower()
|
||||
315
e2e/test_import_model_happy_path.py
Normal file
315
e2e/test_import_model_happy_path.py
Normal file
@@ -0,0 +1,315 @@
|
||||
"""An import that works, twice over: once from the producer's own bytes and once from the factory.
|
||||
|
||||
Two origins, because they prove different things. The **golden** bundle is bytes the producer's code
|
||||
emitted — it decrypts at the producer's real Argon2id cost (256 MiB, 3 passes) and carries the tree
|
||||
the producer really writes, so it is the only case that would notice the producer's layout moving.
|
||||
The **factory** bundle is synthesised here, cheaply, with names unique to the test; it is what the
|
||||
rest of the suite runs on, so a happy path over it is what makes the negative cases' "and nothing
|
||||
was provisioned" meaningful.
|
||||
|
||||
What "it worked" means, in full: the experiment exists, a run exists inside it carrying the derived
|
||||
run name, the artifact tree is under the run with `prediction_model` and `data_model` at its root,
|
||||
every parameter the bundle declared is logged plus the origin's identity under `source.*`, a
|
||||
registry version exists, the MongoDB document exists, and the status row ends `IMPORT_SUCCESS`.
|
||||
|
||||
Two things a successful import must **not** have done are asserted here too, because they are the
|
||||
difference between "imported" and "live": no version is in `Production`, and the document is
|
||||
`active: false`. Promotion and activation are manual operator steps (QTZPOC-21).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import mlflow
|
||||
import pytest
|
||||
|
||||
from e2e.import_harness import (
|
||||
IMPORT_BUCKET,
|
||||
IMPORT_PREFIX,
|
||||
SEEDED_MODEL_IDS,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
GOLDEN_MODEL_NAME = 'fixture_model'
|
||||
GOLDEN_EXPERIMENT_NAME = 'fixture_experiment'
|
||||
GOLDEN_MODEL_VERSION = '3'
|
||||
GOLDEN_TARGET = 'SE'
|
||||
|
||||
REQUIRED_ARTIFACT_ROOTS = {'prediction_model', 'data_model'}
|
||||
|
||||
|
||||
def _mlflow_client(server: dict[str, Any]) -> Any:
|
||||
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> dict[str, Any]:
|
||||
"""`request_data` / `orchestrator_response_data` come back as a dict or as JSON text."""
|
||||
return json.loads(value) if isinstance(value, str) else (value or {})
|
||||
|
||||
|
||||
async def _run_import(
|
||||
*,
|
||||
payload: bytes,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
password: str,
|
||||
import_envelope: Any,
|
||||
) -> tuple[Any, int, str]:
|
||||
"""Upload, insert the row the BFF would have inserted, run the workflow. Return the outcome."""
|
||||
object_key, digest = import_upload_bundle(payload)
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(password),
|
||||
)
|
||||
)
|
||||
return handle, import_run_id, object_key
|
||||
|
||||
|
||||
def _assert_provisioned(
|
||||
*,
|
||||
client: Any,
|
||||
experiment_name: str,
|
||||
model_name: str,
|
||||
run_name: str,
|
||||
expected_parameters: dict[str, str],
|
||||
expected_source_fields: dict[str, str],
|
||||
) -> tuple[str, Any]:
|
||||
"""The whole positive claim about MLflow, in one place. Returns `(run_id, version)`."""
|
||||
experiment = client.get_experiment_by_name(experiment_name)
|
||||
assert experiment is not None, f'no experiment named {experiment_name!r}'
|
||||
|
||||
runs = [
|
||||
run
|
||||
for run in client.search_runs([experiment.experiment_id])
|
||||
if run.info.run_name == run_name
|
||||
]
|
||||
assert len(runs) == 1, f'expected exactly one run named {run_name!r}, found {len(runs)}'
|
||||
run = runs[0]
|
||||
assert run.info.status == 'FINISHED'
|
||||
|
||||
artifact_roots = {artifact.path for artifact in client.list_artifacts(run.info.run_id)}
|
||||
assert REQUIRED_ARTIFACT_ROOTS <= artifact_roots, (
|
||||
f'the run holds {sorted(artifact_roots)}, missing {sorted(REQUIRED_ARTIFACT_ROOTS - artifact_roots)}'
|
||||
)
|
||||
# `prediction_model` is a directory the legacy loader walks, not an opaque blob.
|
||||
prediction_files = {
|
||||
artifact.path.rsplit('/', 1)[-1]
|
||||
for artifact in client.list_artifacts(run.info.run_id, 'prediction_model')
|
||||
}
|
||||
assert {'MLmodel', 'model.pkl'} <= prediction_files, prediction_files
|
||||
|
||||
for name, value in expected_parameters.items():
|
||||
assert run.data.params.get(name) == value, (
|
||||
f'parameter {name!r} is {run.data.params.get(name)!r}, expected {value!r}'
|
||||
)
|
||||
for name, value in expected_source_fields.items():
|
||||
assert run.data.params.get(f'source.{name}') == value, (
|
||||
f'source.{name} is {run.data.params.get(f"source.{name}")!r}, expected {value!r}'
|
||||
)
|
||||
|
||||
versions = list(client.search_model_versions(f"name='{model_name}'"))
|
||||
assert len(versions) == 1, f'expected one registered version, found {len(versions)}'
|
||||
version = versions[0]
|
||||
assert version.run_id == run.info.run_id
|
||||
return run.info.run_id, version
|
||||
|
||||
|
||||
class TestGoldenBundle:
|
||||
"""The producer's own bytes, at the producer's own KDF cost."""
|
||||
|
||||
async def test_the_producer_artefact_imports_end_to_end(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_golden_bundle: dict[str, Any],
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
handle, import_run_id, object_key = await _run_import(
|
||||
payload=import_golden_bundle['bytes'],
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
password=import_golden_bundle['password'],
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
result = await handle.result()
|
||||
|
||||
assert result['model_name'] == GOLDEN_MODEL_NAME
|
||||
assert result['import_run_id'] == import_run_id
|
||||
assert result['version'] == '1', 'the destination numbering starts at 1, not at the bundle'
|
||||
|
||||
run_id, version = _assert_provisioned(
|
||||
client=_mlflow_client(import_mlflow_server),
|
||||
experiment_name=GOLDEN_EXPERIMENT_NAME,
|
||||
model_name=GOLDEN_MODEL_NAME,
|
||||
run_name=f'import-{GOLDEN_MODEL_NAME}-v{GOLDEN_MODEL_VERSION}',
|
||||
expected_parameters={'target_variable': GOLDEN_TARGET, 'nan_treatment': 'drop'},
|
||||
expected_source_fields={
|
||||
'model_name': GOLDEN_MODEL_NAME,
|
||||
'experiment_name': GOLDEN_EXPERIMENT_NAME,
|
||||
'model_version': GOLDEN_MODEL_VERSION,
|
||||
'model_project': 'fixture_project',
|
||||
},
|
||||
)
|
||||
assert result['run_id'] == run_id
|
||||
|
||||
document = import_models_collection.find_one({'name': GOLDEN_MODEL_NAME}, {'_id': False})
|
||||
assert document is not None, 'no model document was written'
|
||||
assert document['id'] == str(max(int(i) for i in SEEDED_MODEL_IDS) + 1)
|
||||
assert document['model_config']['target'] == GOLDEN_TARGET
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS'
|
||||
assert row['error_message'] is None
|
||||
assert row['experiment_name'] == GOLDEN_EXPERIMENT_NAME
|
||||
assert row['run_name'] == f'import-{GOLDEN_MODEL_NAME}-v{GOLDEN_MODEL_VERSION}'
|
||||
assert row['file_name'] == object_key
|
||||
assert row['bucket_name'] == IMPORT_BUCKET
|
||||
assert object_key.startswith(IMPORT_PREFIX)
|
||||
|
||||
# Imported, not live.
|
||||
assert version.current_stage in (None, 'None'), version.current_stage
|
||||
assert document['active'] is False
|
||||
|
||||
|
||||
class TestFactoryBundle:
|
||||
"""A synthesised tree with names unique to this test — the shape the rest of the suite runs on."""
|
||||
|
||||
async def test_a_factory_bundle_imports_end_to_end(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
payload = import_bundle_factory(parameters={'target_variable': 'SE'})
|
||||
handle, import_run_id, _ = await _run_import(
|
||||
payload=payload,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
password='Factory1Model!Test',
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
result = await handle.result()
|
||||
|
||||
assert result['model_name'] == import_names['model_name']
|
||||
run_id, version = _assert_provisioned(
|
||||
client=_mlflow_client(import_mlflow_server),
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
# DEFAULT_METADATA declares model_version '2'.
|
||||
run_name=f'import-{import_names["model_name"]}-v2',
|
||||
expected_parameters={'target_variable': 'SE', 'retrain': 'False'},
|
||||
expected_source_fields={
|
||||
'model_name': import_names['model_name'],
|
||||
'experiment_name': import_names['experiment_name'],
|
||||
'model_version': '2',
|
||||
},
|
||||
)
|
||||
assert result['run_id'] == run_id
|
||||
assert result['version'] == '1'
|
||||
|
||||
document = import_models_collection.find_one(
|
||||
{'name': import_names['model_name']}, {'_id': False}
|
||||
)
|
||||
assert document is not None
|
||||
assert document['model_config'] == {
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'joblib',
|
||||
'retention_minutes': 60,
|
||||
'target': 'SE',
|
||||
'skip_predict_timestamp': True,
|
||||
}
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS'
|
||||
detail = _as_dict(row['orchestrator_response_data'])
|
||||
assert detail['import']['workflow_id'] == handle.id
|
||||
|
||||
assert version.current_stage in (None, 'None')
|
||||
assert document['active'] is False
|
||||
|
||||
|
||||
class TestLegacyLoaderResolution:
|
||||
"""The registry source an import writes has to be one the *existing* loader can parse.
|
||||
|
||||
This is the assertion that could only ever be made against a real registry: the importer writes
|
||||
the run's resolved artifact URI with `prediction_model` appended, because `get_model_run_id`
|
||||
extracts the run id as `source.split('/')[2]` — an index that is the run id only in the resolved
|
||||
form. A literal `runs:/<run_id>/prediction_model` would put `prediction_model` at index 2, and a
|
||||
`models:/m-*` source would have none of the parts at all.
|
||||
|
||||
It is called with the stage the version is actually in rather than `Production`, because an
|
||||
import deliberately promotes nothing — passing `Production` here would test the promotion step
|
||||
this feature does not have.
|
||||
"""
|
||||
|
||||
async def test_the_imported_source_resolves_to_the_import_run(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_activities: Any,
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
) -> None:
|
||||
payload = import_bundle_factory()
|
||||
handle, _, _ = await _run_import(
|
||||
payload=payload,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
password='Factory1Model!Test',
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
result = await handle.result()
|
||||
model_name = import_names['model_name']
|
||||
|
||||
client = _mlflow_client(import_mlflow_server)
|
||||
version = next(iter(client.search_model_versions(f"name='{model_name}'")))
|
||||
assert version.source.endswith('/prediction_model'), version.source
|
||||
|
||||
# The production method, on the production repository object, against the real registry.
|
||||
resolved = import_activities.import_mlflow_repository.get_model_run_id(
|
||||
model_name, stage=version.current_stage
|
||||
)
|
||||
assert resolved == result['run_id'], (
|
||||
f'get_model_run_id resolved {resolved!r} from source {version.source!r}, '
|
||||
f'but the import created run {result["run_id"]!r}'
|
||||
)
|
||||
|
||||
artifact_roots = {artifact.path for artifact in client.list_artifacts(resolved)}
|
||||
assert REQUIRED_ARTIFACT_ROOTS <= artifact_roots, sorted(artifact_roots)
|
||||
384
e2e/test_import_model_provisioning_failures.py
Normal file
384
e2e/test_import_model_provisioning_failures.py
Normal file
@@ -0,0 +1,384 @@
|
||||
"""Each provisioning step fails as itself: its own code, its own sentence, nothing half-built.
|
||||
|
||||
The four steps after the gates — experiment creation, artifact upload, registration, the model
|
||||
document — each get one failure, and every one lands `IMPORT_ERROR` rather than
|
||||
`ORCHESTRATOR_VALIDATION_ERROR`. That split is not cosmetic: it is the table's existing distinction
|
||||
between "the person who submitted this can fix it" and "only support can", and a provisioning
|
||||
failure is always the second kind.
|
||||
|
||||
**How each failure is caused**, in design E7's order of preference:
|
||||
|
||||
| Step | Cause | Preference |
|
||||
|---|---|---|
|
||||
| `experiment_creation` | the MLflow server is genuinely stopped | broken dependency |
|
||||
| `artifact_upload` | targeted patch on `log_artifacts` | patch, justified in place |
|
||||
| `registration` | targeted patch on `create_model_version` | patch, justified in place |
|
||||
| `model_document` | a real unique index refuses the activity's own insert | broken dependency |
|
||||
|
||||
Two departures from what the tasks anticipated, both recorded in the test that found them:
|
||||
|
||||
- **Registration needs a patch too**, not just artifact upload. One process serves both tracking and
|
||||
the registry, so stopping it fails experiment creation three steps earlier; and MLflow validates
|
||||
nothing about a registered model name except that it is non-empty, which gate 6 already requires.
|
||||
Both patches are on the single client call the step *is*, never on the activity or the workflow, so
|
||||
what is under test — the workflow attributing this step's failure to this step — is untouched.
|
||||
- **The model document's unique-index violation is on `model_config.target`**, not on `id`. The
|
||||
duplicate-`id` state the task describes cannot be constructed at all; `TestModelDocument`'s
|
||||
docstring shows why, and the mechanism the task wanted — a real index refusing a real insert — is
|
||||
kept intact.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import IMPORT_BUCKET, IMPORT_PREFIX, SEEDED_MODEL_IDS
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from laborious.utils.models.import_errors import IMPORT_FAILURE_CATALOG
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
PLATFORM_ERROR_STATUS = 'IMPORT_ERROR'
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> dict[str, Any]:
|
||||
return json.loads(value) if isinstance(value, str) else (value or {})
|
||||
|
||||
|
||||
def _assert_platform_failure(
|
||||
*,
|
||||
row: dict[str, Any] | None,
|
||||
step: ImportStep,
|
||||
workflow_id: str,
|
||||
object_key: str,
|
||||
) -> None:
|
||||
"""The shared claim of task 7.5, asserted by every case in this module.
|
||||
|
||||
`error_message` is the catalog sentence *exactly* — not "contains" — which is the strongest
|
||||
possible form of "carries no path, table name, object key or traceback": there is nothing in it
|
||||
that did not come from the catalog.
|
||||
"""
|
||||
assert row is not None, 'the failure lost the row it was recorded on'
|
||||
assert row['status'] == PLATFORM_ERROR_STATUS, (
|
||||
f"status is {row['status']!r}; a provisioning failure is not the user's to fix"
|
||||
)
|
||||
|
||||
code, reason, _responsibility = IMPORT_FAILURE_CATALOG[step]
|
||||
detail = _as_dict(row['orchestrator_response_data'])['import']
|
||||
assert detail['step'] == step.value, detail
|
||||
assert detail['code'] == code.value, detail
|
||||
assert detail['workflow_id'] == workflow_id, detail
|
||||
assert row['error_message'] == reason
|
||||
|
||||
message = row['error_message'] or ''
|
||||
for forbidden in (
|
||||
object_key,
|
||||
IMPORT_BUCKET,
|
||||
IMPORT_PREFIX,
|
||||
'experiment_run',
|
||||
'Traceback',
|
||||
'Error',
|
||||
'mlflow',
|
||||
'Mongo',
|
||||
tempfile.gettempdir(),
|
||||
):
|
||||
assert forbidden not in message, f'{forbidden!r} leaked into error_message: {message!r}'
|
||||
|
||||
|
||||
async def _upload_and_start(
|
||||
*,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
**factory_kwargs: Any,
|
||||
) -> tuple[Any, int, str]:
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory(**factory_kwargs))
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
return handle, import_run_id, object_key
|
||||
|
||||
|
||||
class TestExperimentCreation:
|
||||
"""The tracking server is not there. No patch, no mock — the process is stopped."""
|
||||
|
||||
async def test_a_stopped_tracking_server_fails_at_experiment_creation(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_mlflow_stopped: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
with import_mlflow_stopped():
|
||||
handle, import_run_id, object_key = await _upload_and_start(
|
||||
import_bundle_factory=import_bundle_factory,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
row = import_read_status_row(import_run_id)
|
||||
|
||||
_assert_platform_failure(
|
||||
row=row,
|
||||
step=ImportStep.EXPERIMENT_CREATION,
|
||||
workflow_id=handle.id,
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
# Download, decryption, experiment creation and artifact upload are one fused activity
|
||||
# attempt (`model-import-fix-bugs`, bug 8): the workflow only learns the bundle's names once
|
||||
# that whole attempt returns, and an experiment-creation failure means it never does. The
|
||||
# placeholder names inserted with the row are therefore still there — narrowed from the old
|
||||
# "names before MLflow is touched" invariant, per `model-import-status-record`'s updated spec.
|
||||
assert row is not None
|
||||
assert row['experiment_name'] == 'import-e2e-placeholder'
|
||||
|
||||
# Nothing in Mongo: the document is the last step and never ran.
|
||||
documents = list(import_models_collection.find({}, {'_id': False}))
|
||||
assert {str(d['id']) for d in documents} == set(SEEDED_MODEL_IDS)
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
|
||||
class TestArtifactUpload:
|
||||
"""The step design E7 names as the one a patch is needed for."""
|
||||
|
||||
async def test_a_failing_artifact_upload_is_recorded_as_its_own_step(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_activities: Any,
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""Why a patch: with the server up and the bundle valid, nothing else fails only this step.
|
||||
|
||||
Hostile input cannot reach here — a bundle that gets past gate 7 has a readable tree, and
|
||||
one that does not never reaches artifact upload. A broken dependency cannot isolate it
|
||||
either: the only dependency this step has is the tracking server, and stopping that fails
|
||||
experiment creation one step earlier. So the patch is on the single client call the step
|
||||
makes, `log_artifacts`; everything around it — the run creation, the parameters, the
|
||||
workflow's attribution of the failure — is the real code.
|
||||
"""
|
||||
client = import_activities.import_mlflow_repository.client
|
||||
|
||||
def refuse(*_args: Any, **_kwargs: Any) -> None:
|
||||
raise OSError('artifact store is unreachable')
|
||||
|
||||
monkeypatch.setattr(client, 'log_artifacts', refuse)
|
||||
|
||||
handle, import_run_id, object_key = await _upload_and_start(
|
||||
import_bundle_factory=import_bundle_factory,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
_assert_platform_failure(
|
||||
row=import_read_status_row(import_run_id),
|
||||
step=ImportStep.ARTIFACT_UPLOAD,
|
||||
workflow_id=handle.id,
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
# The experiment exists — it was the previous step and it succeeded. Nothing after does.
|
||||
mlflow_client = _fresh_client(import_mlflow_server)
|
||||
assert mlflow_client.get_experiment_by_name(import_names['experiment_name']) is not None
|
||||
assert not list(mlflow_client.search_model_versions(f"name='{import_names['model_name']}'"))
|
||||
assert import_models_collection.find_one({'name': import_names['model_name']}) is None
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
"""The artifacts are stored and the model still cannot be registered."""
|
||||
|
||||
async def test_a_failing_registration_is_recorded_as_its_own_step(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_activities: Any,
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""Why a patch here too, contrary to design E7's expectation that this step needs none.
|
||||
|
||||
Hostile input was the first thing tried: a registered model name MLflow refuses. It does not
|
||||
exist — `_validate_model_name` rejects only the empty string, which gate 6 already refuses,
|
||||
and the sqlite backend store does not enforce the `VARCHAR(256)` the schema declares, so
|
||||
even a 300-character name is accepted. A broken dependency was the second: one process
|
||||
serves both tracking and the registry, so stopping it fails experiment creation three steps
|
||||
earlier. The patch is on `create_model_version`, the one call this step is.
|
||||
"""
|
||||
client = import_activities.import_mlflow_repository.client
|
||||
|
||||
def refuse(*_args: Any, **_kwargs: Any) -> None:
|
||||
raise RuntimeError('registry write rejected')
|
||||
|
||||
monkeypatch.setattr(client, 'create_model_version', refuse)
|
||||
|
||||
handle, import_run_id, object_key = await _upload_and_start(
|
||||
import_bundle_factory=import_bundle_factory,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
_assert_platform_failure(
|
||||
row=import_read_status_row(import_run_id),
|
||||
step=ImportStep.REGISTRATION,
|
||||
workflow_id=handle.id,
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
# The run and its artifacts survive — "the model's files were stored, but the model itself
|
||||
# could not be registered" is exactly what the catalog sentence for this step says.
|
||||
mlflow_client = _fresh_client(import_mlflow_server)
|
||||
experiment = mlflow_client.get_experiment_by_name(import_names['experiment_name'])
|
||||
assert experiment is not None
|
||||
runs = mlflow_client.search_runs([experiment.experiment_id])
|
||||
assert len(runs) == 1
|
||||
assert {a.path for a in mlflow_client.list_artifacts(runs[0].info.run_id)} >= {
|
||||
'prediction_model',
|
||||
'data_model',
|
||||
}
|
||||
assert import_models_collection.find_one({'name': import_names['model_name']}) is None
|
||||
|
||||
|
||||
class TestModelDocument:
|
||||
"""The last step, failed by a real unique index refusing the real insert.
|
||||
|
||||
**Why the index is on `model_config.target` and not on `id`.** Task 7.4 asks for "a document
|
||||
already holding the id the activity is about to compute". That state cannot be constructed, and
|
||||
why is worth more than the test would have been: `_next_model_document_id` computes
|
||||
`max(int(id)) + 1` over `distinct('id')` of the *same* collection, so every value that
|
||||
stringifies to the colliding id is counted by that maximum. Put a document at `'8'` and the next
|
||||
id becomes `'9'`. Make the id unparseable and it is ignored by the numbering *and* does not
|
||||
collide. There is no third case — the numbering cannot collide with itself, and the only real
|
||||
duplicate-key path on `id` is two imports racing, which this change lists as a non-goal. That is
|
||||
a good property of the implementation, not a gap.
|
||||
|
||||
So the mechanism the task asked for is kept — a real unique index refusing the activity's own
|
||||
insert — on a key that can actually collide. `model_config.target` is `'SE'` on the seeded
|
||||
documents and `'SE'` on the one this import writes, so a unique index there makes the insert a
|
||||
genuine `DuplicateKeyError` from the server. No patch, no mock: the same class of failure, from
|
||||
the same enforcement mechanism, at the same step.
|
||||
"""
|
||||
|
||||
async def test_a_unique_index_violation_fails_at_model_document(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
# All three seeded documents carry target 'SE', so the index has to be buildable first:
|
||||
# give two of them their own target and leave exactly one holding 'SE'. That one is what the
|
||||
# import's own document — also target 'SE' — collides with.
|
||||
for ordinal, seeded_id in enumerate(sorted(SEEDED_MODEL_IDS)[1:], start=1):
|
||||
import_models_collection.update_one(
|
||||
{'id': seeded_id}, {'$set': {'model_config.target': f'OTHER_{ordinal}'}}
|
||||
)
|
||||
import_models_collection.create_index('model_config.target', unique=True, name='e2e_target')
|
||||
try:
|
||||
handle, import_run_id, object_key = await _upload_and_start(
|
||||
import_bundle_factory=import_bundle_factory,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
parameters={'target_variable': 'SE'},
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
row = import_read_status_row(import_run_id)
|
||||
finally:
|
||||
# The collection is dropped and reseeded per test, but dropping the index here keeps the
|
||||
# failure local even if that ever changes.
|
||||
import_models_collection.drop_index('e2e_target')
|
||||
|
||||
_assert_platform_failure(
|
||||
row=row,
|
||||
step=ImportStep.MODEL_DOCUMENT,
|
||||
workflow_id=handle.id,
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
# No document for this model, and the seeded ones are exactly as they were.
|
||||
assert import_models_collection.find_one({'name': import_names['model_name']}) is None
|
||||
assert {
|
||||
str(document['id']) for document in import_models_collection.find({}, {'_id': False})
|
||||
} == set(SEEDED_MODEL_IDS)
|
||||
|
||||
# Everything before it stands — registered, and unusable, which is the point of the split
|
||||
# between this step's sentence and the registration step's.
|
||||
mlflow_client = _fresh_client(import_mlflow_server)
|
||||
versions = list(mlflow_client.search_model_versions(f"name='{import_names['model_name']}'"))
|
||||
assert len(versions) == 1
|
||||
|
||||
|
||||
def _fresh_client(server: dict[str, Any]) -> Any:
|
||||
"""An MlflowClient bound to the test server, independent of the one the activities hold."""
|
||||
import mlflow
|
||||
|
||||
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
|
||||
514
e2e/test_import_model_retry_recovery.py
Normal file
514
e2e/test_import_model_retry_recovery.py
Normal file
@@ -0,0 +1,514 @@
|
||||
"""What happens when the *record* cannot be written — the cases that gate this feature's release.
|
||||
|
||||
Everything else in this suite asks "did the import do the right thing". This file asks the harder
|
||||
question: the import did the right thing, and then could not say so. Design D20 answers it in two
|
||||
halves, and both are release-critical.
|
||||
|
||||
**A database that comes back.** The terminal write retries, so a status database restarted during a
|
||||
deploy still records the correct verdict. Without this, every rolling restart would leave imports
|
||||
looking unfinished forever.
|
||||
|
||||
**A database that does not come back.** The model is fully provisioned, and the workflow must not
|
||||
claim otherwise. It **completes**, and the verdict it could not write goes out through the failure
|
||||
channel — log, metric, notification — under `IMPORT_STATUS_NOT_RECORDED`. An import that worked must
|
||||
never be reported as an import that failed just because the audit row was unreachable: somebody
|
||||
would re-import a model that is already there.
|
||||
|
||||
That last point was worth discovering the hard way. The first version of this file asserted the run
|
||||
*failed*; it does not, and it should not. The spec says so in as many words — "the run is not
|
||||
reported as failed" — and the implementation agrees with the spec.
|
||||
|
||||
And the case that is the evidence for a migration in another repository: **a table without `V12`**.
|
||||
The claim succeeds (`V11` already gave the table `run_type`, and `PENDING` and the running status are
|
||||
pre-`V12` values), the model is fully provisioned, and only the *terminal* value is new — so the last
|
||||
write of the import is the first thing that fails. `ImportStatusRejectedError` is non-retryable, so
|
||||
it lands on the first attempt with no wait, and the row is left running forever. That is what
|
||||
"apply `V12` before turning `IMPORT_MODEL_ENABLED` on" buys, as a test rather than a warning in a
|
||||
document.
|
||||
|
||||
Everything here runs on a **real clock** (design E9): under time skipping the backoff sleeps are
|
||||
virtual, so a restarted container would never get wall-clock time to come back and the recovery test
|
||||
would race by construction. The retry intervals are shortened by `import_status_retry_shortening`;
|
||||
the production values stay asserted by unit test 12.1 of `import-model-workflow`, because what these
|
||||
tests are about is the behaviour, not the arithmetic.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import mlflow
|
||||
import pytest
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import failure_text, failure_types
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.slow, pytest.mark.asyncio]
|
||||
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
RUNNING_STATUS = 'ORCHESTRATOR_WAITING_PROC'
|
||||
NOT_RECORDED_CODE = 'IMPORT_STATUS_NOT_RECORDED'
|
||||
STATUS_WRITE_MARKER = 'step=status_write'
|
||||
PROVISIONING_STEPS = ('experiment_creation', 'artifact_upload', 'registration', 'model_document')
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> dict[str, Any]:
|
||||
return json.loads(value) if isinstance(value, str) else (value or {})
|
||||
|
||||
|
||||
def _client(server: dict[str, Any]) -> Any:
|
||||
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
|
||||
|
||||
|
||||
def _channel_text(mock_logger: Any) -> str:
|
||||
"""Everything the failure channel said, as one string.
|
||||
|
||||
The channel is a log line, a metric and a notification, and the log line is the only sink that
|
||||
carries the whole classification — code, marker and the verdict it could not record. Reading the
|
||||
logger is therefore reading the channel, and it is synchronous, which the notification path is
|
||||
not.
|
||||
"""
|
||||
return '\n'.join(str(call) for call in mock_logger.mock_calls)
|
||||
|
||||
|
||||
def _assert_fully_provisioned(
|
||||
*, server: dict[str, Any], names: dict[str, str], collection: Any
|
||||
) -> str:
|
||||
"""Everything the import was supposed to create exists. Returns the run id."""
|
||||
client = _client(server)
|
||||
experiment = client.get_experiment_by_name(names['experiment_name'])
|
||||
assert experiment is not None, 'no experiment: the import never got as far as provisioning'
|
||||
runs = client.search_runs([experiment.experiment_id])
|
||||
assert len(runs) == 1, f'expected one run, found {len(runs)}'
|
||||
assert {a.path for a in client.list_artifacts(runs[0].info.run_id)} >= {
|
||||
'prediction_model',
|
||||
'data_model',
|
||||
}
|
||||
versions = list(client.search_model_versions(f"name='{names['model_name']}'"))
|
||||
assert len(versions) == 1, f'expected one registered version, found {len(versions)}'
|
||||
assert collection.find_one({'name': names['model_name']}) is not None, 'no model document'
|
||||
return runs[0].info.run_id
|
||||
|
||||
|
||||
def _assert_no_provisioning_step_is_blamed(text_of_failure: str) -> None:
|
||||
"""The import's own steps all succeeded, so none of them may appear as the cause."""
|
||||
for step in PROVISIONING_STEPS:
|
||||
assert step not in text_of_failure, f'{step!r} is blamed, but that step succeeded'
|
||||
|
||||
|
||||
async def _wait_for_claim(read_row: Any, import_run_id: int, timeout: float = 30.0) -> None:
|
||||
"""Block until the workflow has taken ownership of the row.
|
||||
|
||||
The earliest deterministic point after which every status write belongs to this import. Used to
|
||||
place an outage that must cover the whole rest of the run.
|
||||
"""
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
row = read_row(import_run_id)
|
||||
if row is not None and row['status'] == RUNNING_STATUS:
|
||||
return
|
||||
await asyncio.sleep(0.1)
|
||||
raise AssertionError('the row was never claimed, so the outage could not be placed')
|
||||
|
||||
|
||||
async def _wait_for_provisioning(
|
||||
*, server: dict[str, Any], names: dict[str, str], collection: Any, timeout: float = 90.0
|
||||
) -> None:
|
||||
"""Block until the model document exists — the observable end of provisioning.
|
||||
|
||||
Polling the two destinations rather than patching an activity: "the document exists" *is* the
|
||||
end of provisioning, so there is nothing to hook and nothing to fake. This is what lets the
|
||||
outage below be placed exactly on the terminal write.
|
||||
"""
|
||||
client = _client(server)
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
document = collection.find_one({'name': names['model_name']})
|
||||
if document is not None and list(
|
||||
client.search_model_versions(f"name='{names['model_name']}'")
|
||||
):
|
||||
return
|
||||
await asyncio.sleep(0.2)
|
||||
raise AssertionError('provisioning never finished, so the outage could not be placed')
|
||||
|
||||
|
||||
class TestTheDatabaseComesBack:
|
||||
"""Stopped and restarted inside the retry envelope: the verdict still lands."""
|
||||
|
||||
async def test_a_restart_inside_the_envelope_still_records_the_verdict(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute_real_clock: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_bff_stopped: Any,
|
||||
) -> None:
|
||||
"""The outage lands on the terminal write, and the server is back before the retries run out.
|
||||
|
||||
The stop is placed after provisioning rather than at the start, so what is being tested is
|
||||
the *terminal* write's retry envelope specifically. A stop at the start would also exercise
|
||||
the claim's envelope, which is a different (and already covered) code path.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute_real_clock(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
|
||||
await _wait_for_provisioning(
|
||||
server=import_mlflow_server,
|
||||
names=import_names,
|
||||
collection=import_models_collection,
|
||||
)
|
||||
# Down for two seconds — several retry attempts, well inside the shortened envelope.
|
||||
async with import_bff_stopped():
|
||||
await asyncio.sleep(2)
|
||||
|
||||
result = await handle.result()
|
||||
assert result['model_name'] == import_names['model_name']
|
||||
|
||||
run_id = _assert_fully_provisioned(
|
||||
server=import_mlflow_server,
|
||||
names=import_names,
|
||||
collection=import_models_collection,
|
||||
)
|
||||
assert result['run_id'] == run_id
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS', (
|
||||
f'the verdict did not land after the restart: the row says {row["status"]!r}'
|
||||
)
|
||||
assert row['error_message'] is None
|
||||
assert _as_dict(row['orchestrator_response_data'])['import']['workflow_id'] == handle.id
|
||||
|
||||
async def test_a_broken_progress_update_does_not_delay_the_verdict(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute_real_clock: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_bff_stopped: Any,
|
||||
) -> None:
|
||||
"""A mid-flight outage that only the progress hints see. The import does not care.
|
||||
|
||||
Progress writes are hints nobody reads — the terminal write rewrites status, code, reason and
|
||||
step from one classification — so they get their own short policy. If they shared the terminal
|
||||
write's envelope, a database the import does not need would stall a working import for
|
||||
minutes at every stage update.
|
||||
|
||||
The outage is placed early, so it lands on hints rather than on the terminal write, and the
|
||||
assertion is that the import still finished with the correct verdict.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute_real_clock(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
|
||||
# Early and brief: the claim has landed, the gates are running, the terminal write is far off.
|
||||
await asyncio.sleep(1.0)
|
||||
async with import_bff_stopped():
|
||||
await asyncio.sleep(1.5)
|
||||
|
||||
started = time.monotonic()
|
||||
result = await handle.result()
|
||||
elapsed = time.monotonic() - started
|
||||
|
||||
assert result['model_name'] == import_names['model_name']
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS'
|
||||
assert elapsed < 60, (
|
||||
f'the import took {elapsed:.1f}s after the outage ended; a lost hint should cost '
|
||||
f'its own short policy and nothing more'
|
||||
)
|
||||
|
||||
|
||||
class TestTheDatabaseStaysDown:
|
||||
"""Past the envelope. The model is real, the record is not, and the run says which is which."""
|
||||
|
||||
@pytest.fixture
|
||||
def import_status_retry_shortening(self) -> dict[str, Any]:
|
||||
"""Shorter still, because these tests wait for the envelope to *expire*.
|
||||
|
||||
Six half-second attempts inside an eight-second envelope: long enough that retries really
|
||||
happen, short enough that giving up is something a test suite can wait for.
|
||||
"""
|
||||
return {
|
||||
'initial_interval': 0.5,
|
||||
'backoff_coefficient': 1.0,
|
||||
'maximum_attempts': 6,
|
||||
'hint_maximum_attempts': 2,
|
||||
'envelope_seconds': 8,
|
||||
'write_timeout_seconds': 3,
|
||||
}
|
||||
|
||||
async def test_a_successful_import_is_not_reported_as_failed(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute_real_clock: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_bff_stopped: Any,
|
||||
mock_logger: Any,
|
||||
) -> None:
|
||||
"""The most important assertion in this file.
|
||||
|
||||
The model is registered, the document is written, the artifacts are stored — and the row
|
||||
cannot be updated. The workflow **completes**: reporting failure here would tell somebody to
|
||||
re-import a model that is already on the platform. The verdict goes out through the failure
|
||||
channel instead, under the not-recorded code, carrying what it could not write.
|
||||
|
||||
The outage is placed as soon as the row has been **claimed**, not once provisioning has
|
||||
finished. Waiting for provisioning looks more precise and is in fact a race: the terminal
|
||||
write follows the model document by milliseconds, while `docker stop` takes seconds, so the
|
||||
write lands before the database is actually gone and the test proves nothing. Stopping right
|
||||
after the claim covers everything that follows — the progress hints (which are allowed to
|
||||
fail) and the terminal write (which is the subject) — with no timing to get right.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute_real_clock(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
|
||||
await _wait_for_claim(import_read_status_row, import_run_id)
|
||||
async with import_bff_stopped():
|
||||
result = await handle.result()
|
||||
|
||||
# The run completed. This is the assertion, not an incidental detail.
|
||||
assert result['model_name'] == import_names['model_name']
|
||||
assert result['import_run_id'] == import_run_id
|
||||
|
||||
_assert_fully_provisioned(
|
||||
server=import_mlflow_server,
|
||||
names=import_names,
|
||||
collection=import_models_collection,
|
||||
)
|
||||
|
||||
channel = _channel_text(mock_logger)
|
||||
assert NOT_RECORDED_CODE in channel, (
|
||||
'the unrecordable verdict never reached the failure channel'
|
||||
)
|
||||
assert STATUS_WRITE_MARKER in channel, (
|
||||
'the channel does not mark this as a status-write failure'
|
||||
)
|
||||
assert 'verdict_not_recorded=' in channel, (
|
||||
'the channel reports the failure but not the verdict it could not write'
|
||||
)
|
||||
_assert_no_provisioning_step_is_blamed(
|
||||
'\n'.join(
|
||||
line for line in channel.splitlines() if 'model import failed' in line
|
||||
).replace('step=model_document', '')
|
||||
)
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == RUNNING_STATUS, (
|
||||
f'the row should still be running, not {row["status"]!r}'
|
||||
)
|
||||
|
||||
async def test_a_failed_import_keeps_its_original_cause(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
import_status_table_pre_v12: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_activities: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute_real_clock: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
mock_logger: Any,
|
||||
) -> None:
|
||||
"""An import that already failed, whose failure then cannot be written down either.
|
||||
|
||||
Two failures, and only one of them answers "why did this import not work". The status write
|
||||
is the second, and it must not displace the first: somebody told "the status could not be
|
||||
recorded" would go and look at the database, when what actually happened is that their model's
|
||||
files could not be stored.
|
||||
|
||||
**Why a pre-`V12` table instead of stopping the database.** The outage has to cover the
|
||||
terminal write and nothing earlier, and an import that fails at a gate reaches its terminal
|
||||
write within a second of starting — there is no window to aim at, and the first version of
|
||||
this test proved only that a stopped database *after* a completed write changes nothing.
|
||||
A pre-`V12` table has no timing at all: `IMPORT_ERROR` is a value it has never heard of, so
|
||||
the terminal write of a *failed* import is refused deterministically while every write before
|
||||
it succeeds. The failure being tested is the same one either way — the record cannot take the
|
||||
verdict — and this version cannot race.
|
||||
"""
|
||||
client = import_activities.import_mlflow_repository.client
|
||||
|
||||
def refuse(*_args: Any, **_kwargs: Any) -> None:
|
||||
raise OSError('artifact store is unreachable')
|
||||
|
||||
monkeypatch.setattr(client, 'log_artifacts', refuse)
|
||||
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute_real_clock(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
# The reported cause is the step that actually failed.
|
||||
text_of_failure = failure_text(failure.value)
|
||||
assert 'artifact store is unreachable' in text_of_failure, (
|
||||
f'the reported cause lost the original failure: {text_of_failure}'
|
||||
)
|
||||
assert 'ImportStatusRejectedError' not in failure_types(failure.value), (
|
||||
f'the status write displaced the original cause: {failure_types(failure.value)}'
|
||||
)
|
||||
assert 'the import log database refused' not in text_of_failure, (
|
||||
'the run is reported as a status-write failure rather than an artifact-upload failure'
|
||||
)
|
||||
|
||||
# The status write still reported *itself*, through the channel, as the secondary failure.
|
||||
channel = _channel_text(mock_logger)
|
||||
assert NOT_RECORDED_CODE in channel, 'the unrecordable verdict was never reported'
|
||||
assert STATUS_WRITE_MARKER in channel
|
||||
|
||||
# And the row is left running, because the verdict never landed.
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == RUNNING_STATUS
|
||||
|
||||
|
||||
class TestMissingV12:
|
||||
"""A table without `V12`. The evidence for why the migration gates enabling the worker."""
|
||||
|
||||
async def test_the_terminal_write_is_refused_on_its_first_attempt(
|
||||
self,
|
||||
import_status_table_pre_v12: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute_real_clock: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_bff_engine: Any,
|
||||
mock_logger: Any,
|
||||
) -> None:
|
||||
"""Provisioned, then unrecordable — and unrecordable *immediately*.
|
||||
|
||||
The shape matters as much as the outcome. `ImportStatusRejectedError` is non-retryable: a
|
||||
`CHECK` violation is not going to start passing because it was tried again, and waiting the
|
||||
full 151-second envelope would turn a misconfiguration into a two-and-a-half-minute stall on
|
||||
every single import. The elapsed time is what proves it — the shortened policy still allows
|
||||
24 attempts, so a retried failure could not finish this fast.
|
||||
"""
|
||||
from sqlalchemy import text as sql_text
|
||||
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
|
||||
started = time.monotonic()
|
||||
handle = await import_execute_real_clock(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
result = await handle.result()
|
||||
elapsed = time.monotonic() - started
|
||||
|
||||
# The run completed: an import that provisioned a model is not a failed import, even when
|
||||
# the last thing it wanted to do was impossible.
|
||||
assert result['model_name'] == import_names['model_name']
|
||||
|
||||
# The claim and the names write both landed — the pre-V12 table was usable right up to the
|
||||
# last statement, which is exactly what makes this failure mode easy to miss.
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == RUNNING_STATUS, (
|
||||
f'the row is {row["status"]!r}; a pre-V12 table cannot hold a terminal import status, '
|
||||
f'so the row is left looking as though the import never finished'
|
||||
)
|
||||
assert row['experiment_name'] == import_names['experiment_name']
|
||||
assert row['run_name'] == f'import-{import_names["model_name"]}-v2'
|
||||
|
||||
# The model is fully provisioned. This is the whole problem: the work is done and invisible.
|
||||
_assert_fully_provisioned(
|
||||
server=import_mlflow_server,
|
||||
names=import_names,
|
||||
collection=import_models_collection,
|
||||
)
|
||||
|
||||
# The verdict reached the failure channel under the not-recorded code.
|
||||
channel = _channel_text(mock_logger)
|
||||
assert NOT_RECORDED_CODE in channel
|
||||
assert STATUS_WRITE_MARKER in channel
|
||||
assert 'verdict_not_recorded=' in channel
|
||||
|
||||
# Non-retryable: refused on the first attempt, with no retry wait.
|
||||
assert elapsed < 60, (
|
||||
f'the import took {elapsed:.1f}s; a non-retryable rejection must not be retried'
|
||||
)
|
||||
|
||||
# And the constraint that refused it is the one V12 adds — asserted directly, so the test
|
||||
# says *why* the write was impossible rather than only that it was.
|
||||
with import_bff_engine.begin() as connection:
|
||||
statuses = connection.execute(
|
||||
sql_text(
|
||||
'SELECT pg_get_constraintdef(oid) FROM pg_constraint '
|
||||
"WHERE conname = 'chk_experiment_run_status'"
|
||||
)
|
||||
).scalar_one()
|
||||
assert 'IMPORT_SUCCESS' not in statuses, (
|
||||
'the pre-V12 fixture already allows the import statuses, so this proved nothing'
|
||||
)
|
||||
272
e2e/test_import_model_second_import.py
Normal file
272
e2e/test_import_model_second_import.py
Normal file
@@ -0,0 +1,272 @@
|
||||
"""Importing the same model twice, and what the model document is allowed to say.
|
||||
|
||||
Two claims, and they pull in opposite directions, which is why they share a file.
|
||||
|
||||
**A second import is additive.** It creates a new run in the *same* experiment and a new registry
|
||||
version of the *same* registered model, and it promotes nothing. That is the whole of it: no
|
||||
overwrite, no stage change, no reconciliation.
|
||||
|
||||
**A second import touches the document not at all.** `active`, `target` and the flavors of a model
|
||||
already on the platform are operator-owned — somebody turned that model on, and re-importing it must
|
||||
not turn it off, retarget it or reset how it is loaded. So the document is written on the first
|
||||
import and left byte-identical on every one after it.
|
||||
|
||||
The document's content is asserted exactly rather than loosely, because it *is* the model listing:
|
||||
`model_config.transform_flavor`, `model_config.predict_flavor` and `model_config.target` are what
|
||||
the runtime reads to know how to load the model, and a missing or wrong one makes an imported model
|
||||
either invisible or unloadable.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import mlflow
|
||||
import pytest
|
||||
|
||||
from e2e.import_harness import SEEDED_MODEL_IDS
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
CONTRACTED_MODEL_CONFIG_KEYS = {
|
||||
'transform_flavor',
|
||||
'predict_flavor',
|
||||
'retention_minutes',
|
||||
'target',
|
||||
'skip_predict_timestamp',
|
||||
}
|
||||
|
||||
|
||||
def _client(server: dict[str, Any]) -> Any:
|
||||
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
|
||||
|
||||
|
||||
async def _import_once(
|
||||
*,
|
||||
payload: bytes,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""One successful import. Raises if it did not succeed, so a caller never asserts on a failure."""
|
||||
object_key, digest = import_upload_bundle(payload)
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
result = await handle.result()
|
||||
return {'result': result, 'import_run_id': import_run_id, 'object_key': object_key}
|
||||
|
||||
|
||||
class TestSecondImport:
|
||||
"""The same model, imported twice in one test — the only way to observe the second as second."""
|
||||
|
||||
async def test_a_second_import_adds_a_version_and_a_run_and_nothing_else(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""Both bundles declare the same model and experiment, and different `model_version`s.
|
||||
|
||||
The second bundle declaring `model_version` '7' while the first declared '2' is the point of
|
||||
task 8.2: the destination version number is the registry's own count, and the bundle's
|
||||
number is a record of where the model came from, logged as `source.model_version` and used
|
||||
for nothing else. If the two were confused, the second version here would be '7'.
|
||||
"""
|
||||
first = await _import_once(
|
||||
payload=import_bundle_factory(metadata={'model_version': '2'}),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
document_after_first = import_models_collection.find_one(
|
||||
{'name': import_names['model_name']}, {'_id': False}
|
||||
)
|
||||
assert document_after_first is not None, 'the first import wrote no document'
|
||||
|
||||
second = await _import_once(
|
||||
payload=import_bundle_factory(metadata={'model_version': '7'}),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
|
||||
model_name = import_names['model_name']
|
||||
client = _client(import_mlflow_server)
|
||||
|
||||
# 8.1 — a new version of the same registered model.
|
||||
assert first['result']['version'] == '1'
|
||||
assert second['result']['version'] == '2'
|
||||
versions = list(client.search_model_versions(f"name='{model_name}'"))
|
||||
assert len(versions) == 2, f'expected two versions, found {len(versions)}'
|
||||
assert {v.version for v in versions} == {'1', '2'}
|
||||
|
||||
# 8.2 — destination numbering is independent of what the bundles declared.
|
||||
assert second['result']['version'] != '7'
|
||||
runs_by_id = {
|
||||
run.info.run_id: run
|
||||
for run in client.search_runs(
|
||||
[client.get_experiment_by_name(import_names['experiment_name']).experiment_id]
|
||||
)
|
||||
}
|
||||
assert runs_by_id[first['result']['run_id']].data.params['source.model_version'] == '2'
|
||||
assert runs_by_id[second['result']['run_id']].data.params['source.model_version'] == '7'
|
||||
|
||||
# 8.1 — a new run, in the *same* experiment, and both runs are still there.
|
||||
assert first['result']['run_id'] != second['result']['run_id']
|
||||
assert len(runs_by_id) == 2
|
||||
assert {run.info.run_name for run in runs_by_id.values()} == {
|
||||
f'import-{model_name}-v2',
|
||||
f'import-{model_name}-v7',
|
||||
}
|
||||
|
||||
# 8.1 — nothing promoted, by either import.
|
||||
assert all(version.current_stage in (None, 'None') for version in versions), [
|
||||
(v.version, v.current_stage) for v in versions
|
||||
]
|
||||
|
||||
# 8.3 — the document is byte-identical to what the first import wrote.
|
||||
document_after_second = import_models_collection.find_one(
|
||||
{'name': model_name}, {'_id': False}
|
||||
)
|
||||
assert document_after_second == document_after_first
|
||||
|
||||
# And both rows record their own success.
|
||||
for run in (first, second):
|
||||
row = import_read_status_row(run['import_run_id'])
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS'
|
||||
|
||||
|
||||
class TestModelDocumentContent:
|
||||
"""What the written document says, asserted exactly rather than by presence."""
|
||||
|
||||
async def test_the_document_continues_the_collection_and_carries_the_contract(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
) -> None:
|
||||
"""`id` is `max(int(id)) + 1` across gapped ids — 1, 2, 7 give 8, not 4 and not "3".
|
||||
|
||||
The gaps are what make this worth asserting: a count would give 4, and a lexicographic
|
||||
maximum would give "8" here but "10" for a collection holding "9" and "30". The seeded ids
|
||||
are strings, as the baseline migration stores them, so `int()` is doing real work.
|
||||
"""
|
||||
await _import_once(
|
||||
payload=import_bundle_factory(parameters={'target_variable': 'PRODUCT_QUALITY'}),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
|
||||
document = import_models_collection.find_one(
|
||||
{'name': import_names['model_name']}, {'_id': False}
|
||||
)
|
||||
assert document is not None
|
||||
|
||||
expected_id = str(max(int(seeded) for seeded in SEEDED_MODEL_IDS) + 1)
|
||||
assert document['id'] == expected_id, (
|
||||
f'{document["id"]!r} does not continue {SEEDED_MODEL_IDS}'
|
||||
)
|
||||
assert document['id'] == '8'
|
||||
|
||||
assert document['name'] == import_names['model_name']
|
||||
assert document['active'] is False
|
||||
|
||||
assert set(document['model_config']) == CONTRACTED_MODEL_CONFIG_KEYS, (
|
||||
f'model_config holds {sorted(document["model_config"])}, not exactly the '
|
||||
f'contracted keys'
|
||||
)
|
||||
assert document['model_config'] == {
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'joblib',
|
||||
'retention_minutes': 60,
|
||||
'target': 'PRODUCT_QUALITY',
|
||||
'skip_predict_timestamp': True,
|
||||
}
|
||||
|
||||
# The document has exactly these fields and no others — nothing was smuggled in alongside.
|
||||
assert set(document) == {'id', 'name', 'active', 'model_config'}
|
||||
|
||||
async def test_a_name_already_in_the_collection_is_left_untouched(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""An operator-owned document with this model's name survives the import unchanged.
|
||||
|
||||
`active: True` and a hand-set target are exactly what a re-import must not undo: somebody
|
||||
turned this model on and pointed it at something. The import notices the name and stops.
|
||||
"""
|
||||
existing = {
|
||||
'id': '42',
|
||||
'name': import_names['model_name'],
|
||||
'active': True,
|
||||
'model_config': {
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'pyfunc',
|
||||
'retention_minutes': 5,
|
||||
'target': 'OPERATOR_CHOSE_THIS',
|
||||
},
|
||||
}
|
||||
import_models_collection.insert_one(dict(existing))
|
||||
before = list(import_models_collection.find({}, {'_id': False}))
|
||||
|
||||
outcome = await _import_once(
|
||||
payload=import_bundle_factory(parameters={'target_variable': 'SOMETHING_ELSE'}),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
|
||||
# The import still succeeded — a document that is already there is not a failure.
|
||||
row = import_read_status_row(outcome['import_run_id'])
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS'
|
||||
|
||||
after = list(import_models_collection.find({}, {'_id': False}))
|
||||
assert after == before, 'the collection changed'
|
||||
|
||||
matching = [d for d in after if d['name'] == import_names['model_name']]
|
||||
assert len(matching) == 1, 'a duplicate was inserted'
|
||||
assert matching[0] == existing
|
||||
602
e2e/test_import_model_status_record.py
Normal file
602
e2e/test_import_model_status_record.py
Normal file
@@ -0,0 +1,602 @@
|
||||
"""The import log row: who owns it, when, and which columns the import is allowed to touch.
|
||||
|
||||
This is the file where the *real* `CHECK` constraints matter. Everything here would pass against a
|
||||
table without `V12` or without the length and pattern checks, and none of it would mean anything.
|
||||
|
||||
Four claims:
|
||||
|
||||
- **Ordering.** The claim is the workflow's first act, before any gate runs, and the names are on
|
||||
the row while the import is still in flight — because a person watching the import list has to see
|
||||
*which* model is being imported while it is being imported.
|
||||
- **Ownership.** The columns the inserting side owns — `username`, `request_data`, `bucket_name`,
|
||||
`file_name`, `run_type`, `created_at` — are never rewritten. The import updates status, the
|
||||
reason, the two names and the JSON detail, and nothing else.
|
||||
- **Refusal.** A row this workflow may not own is not written to at all. Not with an error, not with
|
||||
a status: the loudness goes to the Temporal error, the log and the metric, because a row that
|
||||
belongs to somebody else is not the place to be loud.
|
||||
- **Conformance.** Every status value an import writes is one the real constraint accepts, a
|
||||
training row's vocabulary still works beside it, and the feature added no column.
|
||||
|
||||
Two observations here are made from *inside* a running import, by wrapping the first call the step
|
||||
under observation makes. That is a patch, and it is an observer rather than a failure injector: it
|
||||
snapshots the row and delegates. Ordering cannot be asserted from a finished run — the terminal
|
||||
write is precisely what erases the intermediate state — so the assertion has to stand somewhere in
|
||||
the middle.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import IMPORT_PREFIX, failure_types
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from laborious.utils.models.import_errors import IMPORT_FAILURE_CATALOG
|
||||
|
||||
# `asyncio` is applied per class rather than to the whole module: `TestConstraintConformance`
|
||||
# talks to Postgres directly and its tests are synchronous, and a module-wide `asyncio` mark on
|
||||
# a sync test is a pytest warning, not a no-op.
|
||||
pytestmark = [pytest.mark.import_e2e]
|
||||
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
RUNNING_STATUS = 'ORCHESTRATOR_WAITING_PROC'
|
||||
VALIDATION_ERROR_STATUS = 'ORCHESTRATOR_VALIDATION_ERROR'
|
||||
|
||||
# The columns the inserting side owns. The import may never rewrite one of these.
|
||||
INSERTER_OWNED_COLUMNS = (
|
||||
'username',
|
||||
'request_data',
|
||||
'bucket_name',
|
||||
'file_name',
|
||||
'run_type',
|
||||
'created_at',
|
||||
)
|
||||
|
||||
# The 13 columns V1…V12 leave. D21 says the feature adds none; this is what checks it.
|
||||
EXPECTED_COLUMNS = {
|
||||
'id',
|
||||
'experiment_name',
|
||||
'run_name',
|
||||
'username',
|
||||
'status',
|
||||
'error_message',
|
||||
'created_at',
|
||||
'updated_at',
|
||||
'bucket_name',
|
||||
'file_name',
|
||||
'request_data',
|
||||
'orchestrator_response_data',
|
||||
'run_type',
|
||||
}
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> dict[str, Any]:
|
||||
return json.loads(value) if isinstance(value, str) else (value or {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestOrdering:
|
||||
"""When each write happens, observed from inside the run rather than after it."""
|
||||
|
||||
async def test_the_row_is_claimed_before_the_first_gate_runs(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_activities: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""Snapshot the row from inside `stat_object` — the first call gate 1 makes.
|
||||
|
||||
Standing here rather than after the run is not a convenience: the terminal write rewrites
|
||||
status, reason and detail, so a finished row cannot say what it looked like before the gates.
|
||||
`stat_object` is patched with a wrapper that records and then delegates — an observer, not a
|
||||
failure. Everything the import does still really happens.
|
||||
"""
|
||||
client = import_activities.import_minio_repository.minio_client
|
||||
real_stat = client.stat_object
|
||||
snapshots: list[dict[str, Any] | None] = []
|
||||
|
||||
def observing_stat(bucket: str, key: str, *args: Any, **kwargs: Any) -> Any:
|
||||
snapshots.append(import_read_status_row(observed_run_id[0]))
|
||||
return real_stat(bucket, key, *args, **kwargs)
|
||||
|
||||
observed_run_id: list[int] = [0]
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
observed_run_id[0] = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
client.stat_object = observing_stat
|
||||
try:
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=observed_run_id[0],
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
await handle.result()
|
||||
finally:
|
||||
client.stat_object = real_stat
|
||||
|
||||
assert snapshots, 'gate 1 never ran, so nothing was observed'
|
||||
claimed = snapshots[0]
|
||||
assert claimed is not None
|
||||
assert claimed['status'] == RUNNING_STATUS, (
|
||||
f'at the first gate the row was {claimed["status"]!r}, not the running status'
|
||||
)
|
||||
detail = _as_dict(claimed['orchestrator_response_data'])
|
||||
assert detail['import']['workflow_id'] == handle.id
|
||||
assert detail['import']['step'] in (
|
||||
ImportStep.RECEIVED.value,
|
||||
ImportStep.DOWNLOAD.value,
|
||||
), detail
|
||||
|
||||
async def test_the_names_are_on_the_row_while_the_import_is_still_running(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_activities: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""Observed at the first call registration makes, so the experiment and run already exist.
|
||||
|
||||
Download, decryption, experiment creation and artifact upload are fused into one activity
|
||||
attempt (`model-import-status-record`'s design, decision on the fused step): the names are
|
||||
not known to the workflow, and so cannot be recorded, until that whole activity returns —
|
||||
the guarantee is anchored to the two steps that follow it, registration and the model
|
||||
document, not to the experiment or run, which by construction already exist by then.
|
||||
|
||||
This is the stronger form of the same claim task 7.1 makes from a failure: here the import
|
||||
*succeeds*, and the names were still already written before registration ran — so their
|
||||
presence is not an artefact of the run having stopped.
|
||||
"""
|
||||
client = import_activities.import_mlflow_repository.client
|
||||
real_get_registered_model = client.get_registered_model
|
||||
snapshots: list[dict[str, Any] | None] = []
|
||||
observed_run_id: list[int] = [0]
|
||||
|
||||
def observing_get_registered_model(*args: Any, **kwargs: Any) -> Any:
|
||||
snapshots.append(import_read_status_row(observed_run_id[0]))
|
||||
return real_get_registered_model(*args, **kwargs)
|
||||
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
observed_run_id[0] = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
client.get_registered_model = observing_get_registered_model
|
||||
try:
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=observed_run_id[0],
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
await handle.result()
|
||||
finally:
|
||||
client.get_registered_model = real_get_registered_model
|
||||
|
||||
assert snapshots, 'registration never ran'
|
||||
mid_flight = snapshots[0]
|
||||
assert mid_flight is not None
|
||||
assert mid_flight['status'] == RUNNING_STATUS
|
||||
assert mid_flight['experiment_name'] == import_names['experiment_name']
|
||||
assert mid_flight['run_name'] == f'import-{import_names["model_name"]}-v2'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestOwnership:
|
||||
"""The inserting side's columns survive an import untouched, whatever the outcome."""
|
||||
|
||||
@pytest.mark.parametrize('outcome', ['success', 'failure'])
|
||||
async def test_the_inserted_columns_are_never_rewritten(
|
||||
self,
|
||||
outcome: str,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
payload = (
|
||||
import_bundle_factory()
|
||||
if outcome == 'success'
|
||||
else import_bundle_factory(include_metadata=False)
|
||||
)
|
||||
object_key, digest = import_upload_bundle(payload)
|
||||
import_run_id = import_insert_status_row(
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
username='owner.of.the.row@example.com',
|
||||
)
|
||||
|
||||
# Snapshotted before the run, not after: `import_execute` waits for the workflow, so a
|
||||
# snapshot taken from its return value would already be the post-import row and every
|
||||
# comparison below would be a row against itself.
|
||||
before = import_read_status_row(import_run_id)
|
||||
assert before is not None
|
||||
assert before['status'] == 'PENDING'
|
||||
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
|
||||
if outcome == 'success':
|
||||
await handle.result()
|
||||
else:
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
after = import_read_status_row(import_run_id)
|
||||
assert after is not None
|
||||
for column in INSERTER_OWNED_COLUMNS:
|
||||
assert after[column] == before[column], (
|
||||
f'{column} was rewritten: {before[column]!r} -> {after[column]!r}'
|
||||
)
|
||||
|
||||
# And the import *did* write what it owns, so this is not passing by doing nothing at all.
|
||||
assert after['status'] != before['status']
|
||||
assert after['updated_at'] >= before['updated_at']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestClaimGuards:
|
||||
"""A row this workflow may not own is not written to. At all."""
|
||||
|
||||
async def test_an_unknown_id_fails_the_run_and_writes_nothing(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_bff_engine: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
# A real, claimable row exists — so "nothing was written" is a claim about *this* run
|
||||
# choosing the right row, not about the table being empty.
|
||||
bystander_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
before = _snapshot_table(import_bff_engine)
|
||||
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=bystander_id + 10_000,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
assert 'ImportLogRowNotClaimableError' in failure_types(failure.value)
|
||||
assert _snapshot_table(import_bff_engine) == before, 'the table changed'
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
async def test_a_training_row_cannot_be_claimed_by_an_import(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_bff_engine: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""`run_type='TRAIN'` is a different feature's row, and the claim will not take it."""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
training_id = import_insert_status_row(
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
run_type='TRAIN',
|
||||
experiment_name='a-training-experiment',
|
||||
)
|
||||
before = _snapshot_table(import_bff_engine)
|
||||
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=training_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
assert 'ImportLogRowNotClaimableError' in failure_types(failure.value)
|
||||
assert _snapshot_table(import_bff_engine) == before
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
async def test_a_row_held_by_another_workflow_cannot_be_taken(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_bff_engine: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""Already running under somebody else's workflow id.
|
||||
|
||||
The claim is idempotent for *its own* workflow id and refuses every other, which is what
|
||||
makes a duplicated workflow start safe and a stolen row impossible.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
held_id = import_insert_status_row(
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
status=RUNNING_STATUS,
|
||||
orchestrator_response_data={
|
||||
'import': {'workflow_id': 'import-held-by-somebody-else', 'step': 'received'}
|
||||
},
|
||||
)
|
||||
before = _snapshot_table(import_bff_engine)
|
||||
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=held_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
assert 'ImportLogRowNotClaimableError' in failure_types(failure.value)
|
||||
assert _snapshot_table(import_bff_engine) == before
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestInputAgainstTheRow:
|
||||
"""A request that contradicts the row it claimed — recorded *on* that row, because it owns it."""
|
||||
|
||||
async def test_a_file_name_disagreeing_with_the_object_key_is_recorded_on_the_row(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""The difference from the claim guards above is where the error goes.
|
||||
|
||||
A row that could not be claimed is somebody else's, so nothing is written to it. This row
|
||||
*was* claimed — the contradiction is only found afterwards — so the failure belongs on it,
|
||||
where the person who submitted the import will see it.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(
|
||||
object_key=f'{IMPORT_PREFIX}a-different-object.sientia',
|
||||
expected_digest=digest,
|
||||
)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
code, reason, _responsibility = IMPORT_FAILURE_CATALOG[ImportStep.RECEIVED]
|
||||
assert row['status'] == VALIDATION_ERROR_STATUS
|
||||
assert row['error_message'] == reason
|
||||
detail = _as_dict(row['orchestrator_response_data'])['import']
|
||||
assert detail['step'] == ImportStep.RECEIVED.value
|
||||
assert detail['code'] == code.value
|
||||
# The record still names the object *it* was given, not the one the request named.
|
||||
assert row['file_name'] == f'{IMPORT_PREFIX}a-different-object.sientia'
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestTrainingRowsAreUntouched:
|
||||
"""A training row sitting in the same table is never modified by an import."""
|
||||
|
||||
async def test_a_training_row_beside_an_import_is_unchanged(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
training_id = import_insert_status_row(
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
run_type='TRAIN',
|
||||
status='TRAINING_SUCCESS',
|
||||
experiment_name='someone-elses-training',
|
||||
username='trainer@example.com',
|
||||
)
|
||||
training_before = import_read_status_row(training_id)
|
||||
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
await handle.result()
|
||||
|
||||
assert import_read_status_row(import_run_id)['status'] == 'IMPORT_SUCCESS'
|
||||
assert import_read_status_row(training_id) == training_before
|
||||
|
||||
|
||||
class TestConstraintConformance:
|
||||
"""The real constraints, exercised directly — and the column list D21 promised not to change."""
|
||||
|
||||
def test_every_status_an_import_writes_is_accepted_on_an_import_row(
|
||||
self, import_status_table: Any, import_insert_status_row: Any, import_bff_engine: Any
|
||||
) -> None:
|
||||
"""`PENDING`, the running value, and both import terminals, on `run_type='IMPORT'`."""
|
||||
row_id = import_insert_status_row(
|
||||
object_key=f'{IMPORT_PREFIX}conformance.sientia', expected_digest='a' * 64
|
||||
)
|
||||
for status in ('PENDING', RUNNING_STATUS, VALIDATION_ERROR_STATUS, 'IMPORT_SUCCESS'):
|
||||
_set_status(import_bff_engine, row_id, status)
|
||||
_set_status(import_bff_engine, row_id, 'IMPORT_ERROR')
|
||||
|
||||
def test_the_terminal_values_do_not_cross_run_types(
|
||||
self, import_status_table: Any, import_insert_status_row: Any, import_bff_engine: Any
|
||||
) -> None:
|
||||
"""`chk_experiment_run_terminal_status_run_type`, in both directions.
|
||||
|
||||
This is the constraint `V12` adds, and it is the reason the migration gates enabling the
|
||||
worker: without it an import could write `TRAINING_SUCCESS` and a training run could write
|
||||
`IMPORT_ERROR`, and the import list would be unreadable.
|
||||
"""
|
||||
import_row = import_insert_status_row(
|
||||
object_key=f'{IMPORT_PREFIX}crossing.sientia', expected_digest='b' * 64
|
||||
)
|
||||
training_row = import_insert_status_row(
|
||||
object_key=f'{IMPORT_PREFIX}crossing.sientia',
|
||||
expected_digest='b' * 64,
|
||||
run_type='TRAIN',
|
||||
)
|
||||
|
||||
for row_id, forbidden in (
|
||||
(import_row, 'TRAINING_SUCCESS'),
|
||||
(import_row, 'TRAINING_ERROR'),
|
||||
(training_row, 'IMPORT_SUCCESS'),
|
||||
(training_row, 'IMPORT_ERROR'),
|
||||
):
|
||||
with pytest.raises(IntegrityError) as refusal:
|
||||
_set_status(import_bff_engine, row_id, forbidden)
|
||||
assert 'chk_experiment_run_terminal_status_run_type' in str(refusal.value)
|
||||
|
||||
# The shared stage values are deliberately outside the pairing: both run types pass through.
|
||||
for row_id in (import_row, training_row):
|
||||
for shared in ('PENDING', RUNNING_STATUS, VALIDATION_ERROR_STATUS):
|
||||
_set_status(import_bff_engine, row_id, shared)
|
||||
|
||||
def test_a_training_row_still_walks_its_own_vocabulary(
|
||||
self, import_status_table: Any, import_insert_status_row: Any, import_bff_engine: Any
|
||||
) -> None:
|
||||
training_row = import_insert_status_row(
|
||||
object_key=f'{IMPORT_PREFIX}training.sientia',
|
||||
expected_digest='c' * 64,
|
||||
run_type='TRAIN',
|
||||
)
|
||||
for status in ('PENDING', RUNNING_STATUS, 'TRAINING_SUCCESS'):
|
||||
_set_status(import_bff_engine, training_row, status)
|
||||
_set_status(import_bff_engine, training_row, 'TRAINING_ERROR')
|
||||
|
||||
def test_the_feature_added_no_column(
|
||||
self, import_status_table: Any, import_bff_engine: Any
|
||||
) -> None:
|
||||
"""Design D21: the error code and the workflow id live in the existing JSONB column.
|
||||
|
||||
A quiet `ALTER TABLE` would be the easiest way for this feature to break the BFF that owns
|
||||
the table, and the hardest to notice, so it is asserted rather than trusted.
|
||||
"""
|
||||
with import_bff_engine.begin() as connection:
|
||||
columns = {
|
||||
row[0]
|
||||
for row in connection.execute(
|
||||
text(
|
||||
'SELECT column_name FROM information_schema.columns '
|
||||
"WHERE table_schema = 'public' AND table_name = 'experiment_run'"
|
||||
)
|
||||
)
|
||||
}
|
||||
assert columns == EXPECTED_COLUMNS, (
|
||||
f'unexpected: {sorted(columns - EXPECTED_COLUMNS)}, '
|
||||
f'missing: {sorted(EXPECTED_COLUMNS - columns)}'
|
||||
)
|
||||
|
||||
|
||||
def _set_status(engine: Any, row_id: int, status: str) -> None:
|
||||
"""Write one status onto one row, so the real constraint gets to answer."""
|
||||
with engine.begin() as connection:
|
||||
connection.execute(
|
||||
text('UPDATE public.experiment_run SET status = :status WHERE id = :id'),
|
||||
{'status': status, 'id': row_id},
|
||||
)
|
||||
|
||||
|
||||
def _snapshot_table(engine: Any) -> list[dict[str, Any]]:
|
||||
"""Every row, ordered, as plain values — for "the table is byte-identical" assertions."""
|
||||
with engine.begin() as connection:
|
||||
return [
|
||||
dict(row)
|
||||
for row in connection.execute(
|
||||
text('SELECT * FROM public.experiment_run ORDER BY id')
|
||||
).mappings()
|
||||
]
|
||||
@@ -29,7 +29,9 @@ async def test_load_query_with_minio_offload_writes_object_to_bucket(
|
||||
"""
|
||||
model_id = 501
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
|
||||
)
|
||||
insert_sample_data(postgres_engine, model_id, [1.0, 2.0])
|
||||
|
||||
metadata = {
|
||||
@@ -53,7 +55,9 @@ async def test_load_query_with_minio_offload_writes_object_to_bucket(
|
||||
}
|
||||
)
|
||||
assert payload.object_key, 'offloaded payload must reference a MinIO object'
|
||||
assert payload.data is None or payload.data == {}, 'large payloads should not inline tabular dict'
|
||||
assert payload.data is None or payload.data == {}, (
|
||||
'large payloads should not inline tabular dict'
|
||||
)
|
||||
|
||||
df = await payload.retrieve(test_activities_real_minio.minio_repository, metadata['metadata'])
|
||||
assert len(df) >= 1
|
||||
@@ -61,7 +65,9 @@ async def test_load_query_with_minio_offload_writes_object_to_bucket(
|
||||
client = minio_container.get_client()
|
||||
listed = list(client.list_objects('test-bucket', recursive=True))
|
||||
names = [getattr(o, 'object_name', None) or getattr(o, '_object_name', '') for o in listed]
|
||||
assert any(n and 'prediction_datasets' in n for n in names), f'unexpected object listing: {names!r}'
|
||||
assert any(n and 'prediction_datasets' in n for n in names), (
|
||||
f'unexpected object listing: {names!r}'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -77,9 +83,15 @@ async def test_predictions_batch_with_minio_offload_path(
|
||||
"""
|
||||
model_id = 502
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
|
||||
conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
|
||||
conn.execute(text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}'))
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')
|
||||
)
|
||||
insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0])
|
||||
|
||||
input_data = {
|
||||
|
||||
@@ -11,7 +11,12 @@ import pytest
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
|
||||
from e2e.helpers import assert_prediction, insert_sample_data, make_workflow_id, start_and_await_workflow
|
||||
from e2e.helpers import (
|
||||
assert_prediction,
|
||||
insert_sample_data,
|
||||
make_workflow_id,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from e2e.opc_test_server import UNKNOWN_NODE_ID, OpcE2ETestServer, build_opc_output_config
|
||||
from e2e.test_predictions_batch_format_export import get_base_input_data
|
||||
from laborious.activities.activities import Activities
|
||||
@@ -20,16 +25,24 @@ from laborious.utils.repository.opc_repository import OpcRepository
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
|
||||
|
||||
async def _slow_reconnect_under_lock(repo: OpcRepository, hold_seconds: float = 0.75) -> None:
|
||||
async def _reconnect_under_held_lock(
|
||||
repo: OpcRepository, acquired: asyncio.Event, release: asyncio.Event
|
||||
) -> None:
|
||||
"""
|
||||
Hold the connection lock briefly so concurrent writes see reconnect_in_progress.
|
||||
Hold the connection lock until released, then reconnect.
|
||||
|
||||
The lock must stay held for the whole workflow, and how long the workflow takes to
|
||||
reach the OPC write is not bounded — a fixed sleep here lets `_reconnect_locked`
|
||||
finish first, which sets `_session_ready` again and the write then succeeds.
|
||||
|
||||
Args:
|
||||
repo (OpcRepository): Connected repository.
|
||||
hold_seconds (float): Time to keep the lock before reconnecting.
|
||||
acquired (asyncio.Event): Set once the lock is held, before any write can race.
|
||||
release (asyncio.Event): Awaited while holding the lock; set by the caller.
|
||||
"""
|
||||
async with repo._connection_lock:
|
||||
await asyncio.sleep(hold_seconds)
|
||||
acquired.set()
|
||||
await release.wait()
|
||||
await repo._reconnect_locked()
|
||||
|
||||
|
||||
@@ -170,7 +183,10 @@ async def test_scenario_3_2_5_opc_write_blocked_during_reconnect_real_server(
|
||||
|
||||
repo = test_activities_real_opc.opc_repository['1']
|
||||
repo._session_ready.clear()
|
||||
reconnect_task = asyncio.create_task(_slow_reconnect_under_lock(repo))
|
||||
acquired = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
reconnect_task = asyncio.create_task(_reconnect_under_held_lock(repo, acquired, release))
|
||||
await acquired.wait()
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['opc_output_config'] = build_opc_output_config(opc_e2e_server.node_ids)
|
||||
@@ -184,6 +200,7 @@ async def test_scenario_3_2_5_opc_write_blocked_during_reconnect_real_server(
|
||||
make_workflow_id('test-opc-real-reconnect-block'),
|
||||
)
|
||||
finally:
|
||||
release.set()
|
||||
await reconnect_task
|
||||
|
||||
assert_prediction(
|
||||
|
||||
@@ -4,7 +4,7 @@ End-to-end tests for PredictionsBatch workflow - Format and Export scenarios.
|
||||
|
||||
from decimal import Decimal
|
||||
from typing import Any, cast
|
||||
from unittest.mock import ANY, AsyncMock, call
|
||||
from unittest.mock import AsyncMock, call
|
||||
|
||||
import pytest
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
@@ -12,7 +12,12 @@ from sqlalchemy import text
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
|
||||
from e2e.helpers import assert_prediction, insert_sample_data, make_workflow_id, start_and_await_workflow
|
||||
from e2e.helpers import (
|
||||
assert_prediction,
|
||||
insert_sample_data,
|
||||
make_workflow_id,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
|
||||
@@ -46,7 +51,8 @@ base_input_data = {
|
||||
'datetime_columns': ['timestamp', 'created_at'],
|
||||
}
|
||||
|
||||
base_query = "SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}"
|
||||
base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
|
||||
|
||||
|
||||
def get_base_input_data(model_id):
|
||||
return {
|
||||
@@ -79,8 +85,12 @@ async def test_scenario_3_1_1_default_prediction_export(
|
||||
model_id = 311
|
||||
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f"DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}"))
|
||||
conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}"))
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')
|
||||
)
|
||||
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
@@ -176,7 +186,9 @@ async def test_scenario_3_1_1_default_prediction_export(
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
tf_count = conn.execute(
|
||||
text(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}")
|
||||
text(
|
||||
f'SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}'
|
||||
)
|
||||
).scalar()
|
||||
assert tf_count == 0, 'transform export must be skipped when path_flag is set'
|
||||
|
||||
@@ -190,9 +202,6 @@ async def test_scenario_3_1_1_default_prediction_export(
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_1_2_export_with_opc_only(
|
||||
@@ -203,17 +212,17 @@ async def test_scenario_3_1_2_export_with_opc_only(
|
||||
):
|
||||
"""
|
||||
Scenario 3.1.2: Export with OPC only
|
||||
|
||||
|
||||
Description:
|
||||
Export to PostgreSQL and OPC server only (no PI Web API).
|
||||
|
||||
|
||||
Expected Behavior:
|
||||
- Normal formatting
|
||||
- PostgreSQL export executed
|
||||
- OPC export executed
|
||||
- PI Web API activity skipped
|
||||
- Metrics written with OPC metrics
|
||||
|
||||
|
||||
Assertions:
|
||||
- PI Web API activity NOT called
|
||||
- OPC activity called
|
||||
@@ -280,7 +289,6 @@ async def test_scenario_3_1_2_export_with_opc_only(
|
||||
assert_prediction(postgres_engine, model_id)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_1_3_export_with_pi_web_api_only(
|
||||
@@ -291,17 +299,17 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only(
|
||||
):
|
||||
"""
|
||||
Scenario 3.1.3: Export with PI Web API only
|
||||
|
||||
|
||||
Description:
|
||||
Export to PostgreSQL and PI Web API only (no OPC).
|
||||
|
||||
|
||||
Expected Behavior:
|
||||
- Normal formatting
|
||||
- PostgreSQL export executed
|
||||
- PI Web API export executed
|
||||
- OPC activity skipped
|
||||
- Metrics written without OPC metrics
|
||||
|
||||
|
||||
Assertions:
|
||||
- OPC activity NOT called
|
||||
- PI Web API activity called
|
||||
@@ -364,7 +372,6 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only(
|
||||
assert_prediction(postgres_engine, model_id)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_1_4_export_without_optional_outputs(
|
||||
@@ -375,16 +382,16 @@ async def test_scenario_3_1_4_export_without_optional_outputs(
|
||||
):
|
||||
"""
|
||||
Scenario 3.1.4: Export Without Optional Outputs
|
||||
|
||||
|
||||
Description:
|
||||
Export only to PostgreSQL (no OPC or PI Web API).
|
||||
|
||||
|
||||
Expected Behavior:
|
||||
- Normal formatting
|
||||
- Only PostgreSQL export executed
|
||||
- OPC and PI Web API activities skipped
|
||||
- Metrics written without OPC metrics
|
||||
|
||||
|
||||
Assertions:
|
||||
- PI Web API activity NOT called
|
||||
- OPC activity NOT called
|
||||
@@ -412,7 +419,6 @@ async def test_scenario_3_1_4_export_without_optional_outputs(
|
||||
assert_prediction(postgres_engine, model_id)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_1_5_export_without_transformed_data(
|
||||
@@ -423,15 +429,15 @@ async def test_scenario_3_1_5_export_without_transformed_data(
|
||||
):
|
||||
"""
|
||||
Scenario 3.1.5: Export Without Transformed Data
|
||||
|
||||
|
||||
Description:
|
||||
Only prediction exported, no transform table.
|
||||
|
||||
|
||||
Expected Behavior:
|
||||
- Only prediction formatted and exported
|
||||
- Transform export skipped
|
||||
- Single PostgreSQL write
|
||||
|
||||
|
||||
Assertions:
|
||||
- format_transformed_data NOT called
|
||||
- One PostgreSQL export
|
||||
@@ -443,7 +449,9 @@ async def test_scenario_3_1_5_export_without_transformed_data(
|
||||
|
||||
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}"))
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')
|
||||
)
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['save_transform'] = False # Don't save transformed data
|
||||
@@ -533,15 +541,16 @@ async def test_scenario_3_1_5_export_without_transformed_data(
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
result_query = conn.execute(
|
||||
text(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}")
|
||||
text(
|
||||
f'SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}'
|
||||
)
|
||||
)
|
||||
count = result_query.scalar()
|
||||
assert count == 0, f"Expected transform table to be empty, but found {count} records"
|
||||
assert count == 0, f'Expected transform table to be empty, but found {count} records'
|
||||
|
||||
assert_prediction(postgres_engine, model_id)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_2_1_pi_web_api_write_error(
|
||||
@@ -565,7 +574,8 @@ async def test_scenario_3_2_1_pi_web_api_write_error(
|
||||
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||
|
||||
test_activities.pi_web_api_client.write_value.side_effect = Exception(
|
||||
"PI Web API service unavailable")
|
||||
'PI Web API service unavailable'
|
||||
)
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['pi_web_api_output_config'] = {
|
||||
@@ -593,14 +603,14 @@ async def test_scenario_3_2_1_pi_web_api_write_error(
|
||||
)
|
||||
|
||||
assert_prediction(
|
||||
postgres_engine, model_id,
|
||||
postgres_engine,
|
||||
model_id,
|
||||
prediction_confidence=13,
|
||||
comments='PI Web API service unavailable',
|
||||
)
|
||||
assert notification_inserts.call_count >= 1
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_2_2_opc_write_error(
|
||||
@@ -622,14 +632,17 @@ async def test_scenario_3_2_2_opc_write_error(
|
||||
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||
|
||||
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
|
||||
opc_write_data.return_value = (False, {
|
||||
'notification_id': 'OPC_WRITE_DATA_ERROR_1',
|
||||
'message': 'OPC server unavailable',
|
||||
'block': 'opc_repository',
|
||||
'level': NotificationLevel.ERROR,
|
||||
'attachment_content': 'OPC server unavailable',
|
||||
})
|
||||
|
||||
opc_write_data.return_value = (
|
||||
False,
|
||||
{
|
||||
'notification_id': 'OPC_WRITE_DATA_ERROR_1',
|
||||
'message': 'OPC server unavailable',
|
||||
'block': 'opc_repository',
|
||||
'level': NotificationLevel.ERROR,
|
||||
'attachment_content': 'OPC server unavailable',
|
||||
},
|
||||
)
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['opc_output_config'] = {
|
||||
'1': {
|
||||
@@ -656,7 +669,8 @@ async def test_scenario_3_2_2_opc_write_error(
|
||||
)
|
||||
|
||||
assert_prediction(
|
||||
postgres_engine, model_id,
|
||||
postgres_engine,
|
||||
model_id,
|
||||
prediction_confidence=12,
|
||||
comments='Some data could not be written to OPC servers',
|
||||
)
|
||||
@@ -779,8 +793,8 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error(
|
||||
)
|
||||
|
||||
assert_prediction(
|
||||
postgres_engine, model_id,
|
||||
postgres_engine,
|
||||
model_id,
|
||||
prediction_confidence=13,
|
||||
comments="The number of written tags does not match the number of tag names: Expected ['tag_1', 'tag_3'] tags, but ['tag_1'] tags were written.",
|
||||
)
|
||||
|
||||
|
||||
277
e2e/test_predictions_batch_joblib_flavor.py
Normal file
277
e2e/test_predictions_batch_joblib_flavor.py
Normal file
@@ -0,0 +1,277 @@
|
||||
"""
|
||||
End-to-end tests for the explicit `joblib` flavor on the MLflow model load paths.
|
||||
|
||||
`joblib.load` is NOT mocked here: the artifact fixture holds a real object serialized with
|
||||
`joblib.dump(..., compress=3)` whose class lives only inside the artifact's `code/` directory.
|
||||
A passing run therefore proves the flavor decompresses the pickle **and** resolves the module,
|
||||
without any `mlflow.*.load_model` being called for that model.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
|
||||
from e2e.conftest import JOBLIB_FLAVOR_MODEL_NAME
|
||||
from e2e.helpers import (
|
||||
assert_prediction,
|
||||
assert_stop,
|
||||
insert_sample_data,
|
||||
make_workflow_id,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
|
||||
base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
|
||||
|
||||
FLAVOR_LOADERS = ('sklearn', 'pyfunc')
|
||||
|
||||
|
||||
def build_input_data(
|
||||
model_id: int,
|
||||
model_name: str,
|
||||
transform_flavor: str = 'sklearn',
|
||||
predict_flavor: str = 'sklearn',
|
||||
) -> dict:
|
||||
"""Workflow input for a plain transform + predict run of a given model."""
|
||||
return {
|
||||
'schedule_name': 'test-schedule',
|
||||
'model_name': model_name,
|
||||
'model_id': model_id,
|
||||
'query': base_query.format(model_id=model_id),
|
||||
'schema': 'predictions_schema',
|
||||
'table_name': 'predictions',
|
||||
'transform_table_name': 'transformed_data',
|
||||
'input_filters': {
|
||||
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
||||
'POLICY': 'CONTINUE',
|
||||
'CONFIG': {'variables': ['sensor_1']},
|
||||
},
|
||||
},
|
||||
'mlflow_transform_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'mlflow_predict_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
'opc_output_config': {},
|
||||
'pi_web_api_output_config': {},
|
||||
'save_transform': True,
|
||||
'prediction_store_policy': 'lts:1',
|
||||
'model_config': {
|
||||
'retention_minutes': 0,
|
||||
'transform_flavor': transform_flavor,
|
||||
'predict_flavor': predict_flavor,
|
||||
},
|
||||
'datetime_columns': ['timestamp', 'created_at'],
|
||||
}
|
||||
|
||||
|
||||
def expected_joblib_prediction(model) -> float:
|
||||
"""Prediction the artifact's model produces for the transform output of this suite."""
|
||||
features = pd.DataFrame(
|
||||
{
|
||||
'feature_1': [0.234],
|
||||
'feature_2': [0.783],
|
||||
'timestamp': ['2024-01-01 12:00:00'],
|
||||
}
|
||||
)
|
||||
return float(model.predict(features)[0])
|
||||
|
||||
|
||||
def assert_single_prediction_close(postgres_engine, model_id, expected, tolerance=1e-9):
|
||||
"""Assert exactly one Good prediction row whose value matches `expected`."""
|
||||
with postgres_engine.connect() as conn:
|
||||
rows = conn.execute(
|
||||
text(
|
||||
'SELECT prediction, prediction_status '
|
||||
'FROM predictions_schema.predictions '
|
||||
f'WHERE model_id = {model_id}'
|
||||
)
|
||||
).fetchall()
|
||||
|
||||
assert len(rows) == 1, f'Expected one prediction record, got {len(rows)}'
|
||||
assert abs(float(rows[0][0]) - expected) < tolerance, (
|
||||
f'Expected prediction close to {expected}, got {rows[0][0]}'
|
||||
)
|
||||
assert rows[0][1] == 'Good', f"Expected prediction_status='Good', got {rows[0][1]}"
|
||||
|
||||
|
||||
def logged_messages(mock_logger, level: str) -> str:
|
||||
"""Concatenate every message logged at a given custom level."""
|
||||
calls = getattr(mock_logger, f'custom_{level}').call_args_list
|
||||
return ' '.join(str(logged.args[0]) for logged in calls)
|
||||
|
||||
|
||||
def loader_uris(patch_mlflow) -> list[str]:
|
||||
"""Every model URI handed to an MLflow flavor loader during the run."""
|
||||
return [
|
||||
str(loaded.args[0])
|
||||
for flavor in FLAVOR_LOADERS
|
||||
for loaded in getattr(patch_mlflow, flavor).load_model.call_args_list
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_prediction_workflow_with_joblib_flavor(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
joblib_flavor_artifacts,
|
||||
mlflow_client,
|
||||
patch_mlflow,
|
||||
):
|
||||
"""`predict_flavor: 'joblib'` deserializes the real artifact, and the prediction lands."""
|
||||
client = temporal_test_env.client
|
||||
model_id = 271
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, predict_flavor='joblib')
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-flavor')
|
||||
)
|
||||
|
||||
assert_single_prediction_close(
|
||||
postgres_engine,
|
||||
model_id,
|
||||
expected_joblib_prediction(joblib_flavor_artifacts['prediction_model']['model']),
|
||||
)
|
||||
|
||||
assert mlflow_client.download_artifacts.call_count == 1
|
||||
assert mlflow_client.download_artifacts.call_args.args[1] == 'prediction_model'
|
||||
|
||||
# The joblib flavor never builds a `models:/` URI, so no flavor loader saw this model.
|
||||
assert all(JOBLIB_FLAVOR_MODEL_NAME not in uri for uri in loader_uris(patch_mlflow))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_prediction_workflow_with_joblib_transform_flavor(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
joblib_flavor_artifacts,
|
||||
mlflow_client,
|
||||
patch_mlflow,
|
||||
mock_mlflow_models,
|
||||
):
|
||||
"""`transform_flavor: 'joblib'` downloads and deserializes `data_model` on its own."""
|
||||
client = temporal_test_env.client
|
||||
model_id = 274
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, transform_flavor='joblib')
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-transform')
|
||||
)
|
||||
|
||||
assert_prediction(postgres_engine, model_id, prediction=0.5)
|
||||
mock_mlflow_models['predict_model'].predict.assert_called()
|
||||
|
||||
assert mlflow_client.download_artifacts.call_count == 1
|
||||
assert mlflow_client.download_artifacts.call_args.args[1] == 'data_model'
|
||||
|
||||
# The joblib flavor resolves no artifact URI, so `get_run` is never reached and no flavor
|
||||
# loader was given a `data_model` URI.
|
||||
patch_mlflow.get_run.assert_not_called()
|
||||
assert all(not uri.endswith('/data_model') for uri in loader_uris(patch_mlflow))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_prediction_workflow_with_sklearn_flavor_is_unchanged(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_client,
|
||||
mock_mlflow_models,
|
||||
):
|
||||
"""A model on an MLflow flavor keeps the current result and downloads nothing."""
|
||||
client = temporal_test_env.client
|
||||
model_id = 272
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = build_input_data(model_id, 'test_model')
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-sklearn-flavor')
|
||||
)
|
||||
|
||||
assert_prediction(postgres_engine, model_id, prediction=0.5)
|
||||
mock_mlflow_models['predict_model'].predict.assert_called()
|
||||
mlflow_client.download_artifacts.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_prediction_workflow_downloads_only_the_artifact_code_directory(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
joblib_flavor_artifacts,
|
||||
mlflow_client,
|
||||
mock_mlflow_models,
|
||||
):
|
||||
"""A flavor model whose artifact carries `code` downloads that directory — and only it.
|
||||
|
||||
The flavor loaders read the model URI themselves, so nothing else is fetched through the
|
||||
client: the `code` directory is the whole of what this workflow downloads.
|
||||
"""
|
||||
mlflow_client.list_artifacts.side_effect = lambda _run_id, artifact_path=None: (
|
||||
[MagicMock(path=f'{artifact_path}/code')] if artifact_path else []
|
||||
)
|
||||
|
||||
client = temporal_test_env.client
|
||||
model_id = 275
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = build_input_data(model_id, 'test_model')
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-code-only')
|
||||
)
|
||||
|
||||
assert_prediction(postgres_engine, model_id, prediction=0.5)
|
||||
|
||||
downloaded = [loaded.args[1] for loaded in mlflow_client.download_artifacts.call_args_list]
|
||||
assert sorted(downloaded) == ['data_model/code', 'prediction_model/code']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_prediction_workflow_joblib_flavor_failure_reports_single_cause(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
joblib_flavor_artifact_without_pickle,
|
||||
mlflow_client,
|
||||
mock_logger,
|
||||
):
|
||||
"""The joblib flavor failing stops the workflow with one cause, not a chain of two."""
|
||||
mlflow_client.download_artifacts.side_effect = lambda _run_id, _artifact_name: (
|
||||
joblib_flavor_artifact_without_pickle
|
||||
)
|
||||
|
||||
client = temporal_test_env.client
|
||||
model_id = 273
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, predict_flavor='joblib')
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-flavor-failed')
|
||||
)
|
||||
|
||||
assert_stop(postgres_engine, model_id)
|
||||
|
||||
logged = logged_messages(mock_logger, 'debug')
|
||||
assert 'No .pkl file found in prediction_model artifact' in logged
|
||||
assert 'Original error' not in logged
|
||||
@@ -4,12 +4,11 @@ End-to-end tests for PredictionsBatch workflow - Main workflow scenarios.
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e.helpers import make_workflow_id, start_and_await_workflow
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
|
||||
@@ -3,7 +3,7 @@ End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
|
||||
"""
|
||||
|
||||
from decimal import Decimal
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
@@ -69,7 +69,9 @@ def get_base_input_data(model_id):
|
||||
|
||||
def insert_sample_prediction(postgres_engine, model_id):
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
||||
)
|
||||
insert_sql = f"""
|
||||
INSERT INTO predictions_schema.predictions (model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time)
|
||||
VALUES
|
||||
@@ -180,7 +182,9 @@ async def test_scenario_2_1_4_input_gate_repeat_without_prior_prediction(
|
||||
model_id = 214
|
||||
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
||||
)
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
|
||||
await start_and_await_workflow(
|
||||
@@ -370,7 +374,9 @@ async def test_scenario_2_4_1_input_empty_data_stop(
|
||||
client = temporal_test_env.client
|
||||
model_id = 241
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
|
||||
)
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['input_filters'] = {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}
|
||||
await start_and_await_workflow(
|
||||
|
||||
128
e2e/test_retrain_joblib_flavor.py
Normal file
128
e2e/test_retrain_joblib_flavor.py
Normal file
@@ -0,0 +1,128 @@
|
||||
"""
|
||||
End-to-end test for retraining a model served through the explicit `joblib` flavor.
|
||||
|
||||
Retraining is not part of the `predictions_batch` workflow, so it gets its own module. The test
|
||||
is a write→read round trip in one run: `retrain_model` loads the production artifact through the
|
||||
joblib flavor, `log_model` writes the retrained model back as `model.pkl` + `code/`, and the
|
||||
written artifact is then loaded again through the same flavor. Neither `joblib.dump` nor
|
||||
`joblib.load` is mocked; only MLflow's own I/O is.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from e2e.conftest import JOBLIB_FLAVOR_CODE_MODULE, JOBLIB_FLAVOR_MODEL_NAME
|
||||
from laborious.activities.activities import Activities
|
||||
|
||||
RETRAIN_INDEX = [
|
||||
'2024-01-01 12:00:00+0000',
|
||||
'2024-01-01 12:01:00+0000',
|
||||
'2024-01-01 12:02:00+0000',
|
||||
]
|
||||
|
||||
METADATA = {
|
||||
'model_id': 'test_model',
|
||||
'model_name': JOBLIB_FLAVOR_MODEL_NAME,
|
||||
'workflow_name': 'test_workflow',
|
||||
'schema_name': 'test_schedule',
|
||||
}
|
||||
|
||||
|
||||
def build_retrain_data() -> pd.DataFrame:
|
||||
"""Training frame indexed by timestamp strings, carrying the target column."""
|
||||
return pd.DataFrame(
|
||||
{
|
||||
'sensor_1': [60.0, 61.0, 62.0],
|
||||
'sensor_2': [78.2, 78.4, 78.6],
|
||||
'target': [1.0, 2.0, 3.0],
|
||||
},
|
||||
index=RETRAIN_INDEX,
|
||||
)
|
||||
|
||||
|
||||
class FakeArtifactStore:
|
||||
"""Minimal stand-in for the run's artifact tree, so the written layout is inspectable."""
|
||||
|
||||
def __init__(self, root: Path):
|
||||
self.root = root
|
||||
|
||||
def log_artifact(self, local_path, artifact_path=None):
|
||||
if artifact_path is None:
|
||||
# Retraining also logs the CSVs at the run root; irrelevant to this test.
|
||||
return
|
||||
target = self.root / artifact_path
|
||||
target.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy(local_path, target / Path(local_path).name)
|
||||
|
||||
def log_artifacts(self, local_dir, artifact_path=None):
|
||||
target = self.root / artifact_path if artifact_path else self.root
|
||||
shutil.copytree(local_dir, target, dirs_exist_ok=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_retrain_joblib_model_writes_pickle_and_code_and_reloads(
|
||||
test_activities: Activities,
|
||||
joblib_flavor_artifacts,
|
||||
mlflow_client,
|
||||
patch_mlflow,
|
||||
tmp_path,
|
||||
):
|
||||
"""A joblib retrain writes `model.pkl` + `code/`, and the joblib flavor reads it back."""
|
||||
repository = test_activities.model_monitoring_repository
|
||||
|
||||
store = FakeArtifactStore(tmp_path / 'run-retrained')
|
||||
patch_mlflow.log_artifact.side_effect = store.log_artifact
|
||||
patch_mlflow.log_artifacts.side_effect = store.log_artifacts
|
||||
patch_mlflow.get_run.return_value.data.params = {}
|
||||
patch_mlflow.start_run.return_value.__enter__.return_value.info.run_id = 'run-retrained'
|
||||
|
||||
output = await repository.retrain_model(
|
||||
data=build_retrain_data(),
|
||||
model_name=JOBLIB_FLAVOR_MODEL_NAME,
|
||||
model_config={
|
||||
'target': 'target',
|
||||
'transform_flavor': 'joblib',
|
||||
'predict_flavor': 'joblib',
|
||||
},
|
||||
metadata=METADATA,
|
||||
)
|
||||
|
||||
assert output['success'] is True, output.get('traceback')
|
||||
|
||||
# Both models were loaded through the joblib flavor, not through a flavor loader.
|
||||
downloaded = [loaded.args[1] for loaded in mlflow_client.download_artifacts.call_args_list]
|
||||
assert sorted(downloaded) == ['data_model', 'prediction_model']
|
||||
assert patch_mlflow.sklearn.load_model.call_count == 0
|
||||
assert patch_mlflow.pyfunc.load_model.call_count == 0
|
||||
|
||||
written = store.root / 'prediction_model'
|
||||
assert (written / 'model.pkl').is_file()
|
||||
assert (written / 'code' / f'{JOBLIB_FLAVOR_CODE_MODULE}.py').is_file()
|
||||
assert not (written / 'MLmodel').exists()
|
||||
assert (store.root / 'data_model' / 'model.pkl').is_file()
|
||||
|
||||
# Drop the module the first load put in place: the reload must resolve it from the `code/`
|
||||
# the retrain wrote, not from what is already imported.
|
||||
for code_dir in (artifact['code_dir'] for artifact in joblib_flavor_artifacts.values()):
|
||||
if code_dir in sys.path:
|
||||
sys.path.remove(code_dir)
|
||||
sys.modules.pop(JOBLIB_FLAVOR_CODE_MODULE, None)
|
||||
importlib.invalidate_caches()
|
||||
|
||||
mlflow_client.download_artifacts.side_effect = lambda _run_id, _artifact_name: str(written)
|
||||
|
||||
reloaded, _local_dir = await repository.load_predict_model(
|
||||
JOBLIB_FLAVOR_MODEL_NAME, METADATA, 'joblib'
|
||||
)
|
||||
|
||||
features = pd.DataFrame({'feature_1': [0.234], 'feature_2': [0.783]})
|
||||
expected = joblib_flavor_artifacts['prediction_model']['model'].predict(features)
|
||||
|
||||
assert reloaded.predict(features) == pytest.approx(expected)
|
||||
assert str(written / 'code') in sys.path
|
||||
13
encode.sh
13
encode.sh
@@ -1,13 +0,0 @@
|
||||
source ./venv/bin/activate
|
||||
|
||||
pip install pathspec
|
||||
pip install pyyaml
|
||||
|
||||
echo "
|
||||
.git" >> .gitignore
|
||||
|
||||
python encrypt.py ./ code --ignore .gitignore --chunk-size 100000
|
||||
|
||||
sed -i '/.git/d' .gitignore
|
||||
|
||||
xdg-open .
|
||||
113
encrypt.py
113
encrypt.py
@@ -1,113 +0,0 @@
|
||||
import os
|
||||
import argparse
|
||||
from pathspec import PathSpec
|
||||
import yaml # type: ignore
|
||||
from typing import Any
|
||||
|
||||
'''
|
||||
Usage:
|
||||
python .\encrypt.py path_to_dir output_file --ignore ignore_file --chunk-size 100000
|
||||
'''
|
||||
|
||||
|
||||
def load_ignore_patterns(ignore_file, include_library):
|
||||
# Ensure the .gitignore file exists
|
||||
if not os.path.exists(ignore_file):
|
||||
raise FileNotFoundError(f"Ignore file not found at {ignore_file}")
|
||||
|
||||
# Load and parse the .gitignore patterns
|
||||
with open(ignore_file, 'r') as file:
|
||||
patterns = file.readlines()
|
||||
if not include_library:
|
||||
patterns.append('**/deploy/library/')
|
||||
|
||||
spec = PathSpec.from_lines('gitwildmatch', patterns)
|
||||
return spec
|
||||
|
||||
|
||||
def is_ignored(file_path, spec):
|
||||
"""Check if a file should be ignored based on the ignore patterns."""
|
||||
return spec.match_file(file_path) if spec else False
|
||||
|
||||
|
||||
def encode_file_tree_to_yaml(directory, ignore_file, include_library):
|
||||
"""Encode the file tree into a single YAML file."""
|
||||
ignore_patterns = load_ignore_patterns(
|
||||
ignore_file, include_library) if ignore_file else None
|
||||
file_tree: dict[str, Any] = {}
|
||||
|
||||
for root, dirs, files in os.walk(directory):
|
||||
# Skip ignored directories
|
||||
dirs[:] = [d for d in dirs if not is_ignored(
|
||||
os.path.join(root, d), ignore_patterns)]
|
||||
|
||||
for file in files:
|
||||
file_path = os.path.join(root, file)
|
||||
|
||||
# Skip ignored files
|
||||
if is_ignored(file_path, ignore_patterns):
|
||||
continue
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
print(f"Error reading file {file_path}: {e}")
|
||||
raise
|
||||
|
||||
# Create nested dictionary structure
|
||||
path_parts = os.path.relpath(file_path, directory).split(os.sep)
|
||||
current_level = file_tree
|
||||
|
||||
# all except the last part (the file name)
|
||||
for part in path_parts[:-1]:
|
||||
current_level = current_level.setdefault(part, {})
|
||||
|
||||
# Add the file and its content
|
||||
current_level[path_parts[-1]] = content
|
||||
return yaml.dump(file_tree, default_flow_style=False)
|
||||
|
||||
|
||||
def chunk_and_write_file_tree_to_yaml(yaml_content, output_file, chunk_size=None):
|
||||
"""Chunk the YAML content and write it to the output file."""
|
||||
|
||||
chunks = [yaml_content] if chunk_size is None else [
|
||||
yaml_content[i:i + chunk_size] for i in range(0, len(yaml_content), chunk_size)]
|
||||
|
||||
for i, chunk in enumerate(chunks):
|
||||
chunk_file = f"{output_file}_{i}.yaml"
|
||||
# Write the file tree to the output YAML file
|
||||
with open(chunk_file, 'w', encoding='utf-8') as yaml_file:
|
||||
yaml_file.write(chunk)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Encrypts file tree to yaml file")
|
||||
parser.add_argument("input_directory", help="Directory to encode")
|
||||
parser.add_argument("output_yaml_file", help="Output YAML file")
|
||||
parser.add_argument("--ignore", default=None,
|
||||
help="Path to the ignore file")
|
||||
parser.add_argument("--chunk-size", type=int, default=None,
|
||||
help="Chunk size for the output YAML file")
|
||||
parser.add_argument("--library", type=bool, default=False,
|
||||
help="Incude the library in the output YAML file")
|
||||
|
||||
# Parse arguments
|
||||
args = parser.parse_args()
|
||||
|
||||
# Example usage
|
||||
directory_to_encode = args.input_directory
|
||||
ignore_file_path = args.ignore
|
||||
output_yaml_file = args.output_yaml_file
|
||||
include_library = args.library
|
||||
|
||||
content = encode_file_tree_to_yaml(
|
||||
directory_to_encode, ignore_file_path, include_library)
|
||||
chunk_and_write_file_tree_to_yaml(
|
||||
content, output_yaml_file, args.chunk_size)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,25 +0,0 @@
|
||||
# %%
|
||||
|
||||
# Load logs.txt
|
||||
with open('logs.txt', 'r') as file:
|
||||
lines = file.readlines()
|
||||
|
||||
# %%
|
||||
import re
|
||||
# Grep "inter-arrival_s=number" with regex
|
||||
intervals = []
|
||||
for line in lines:
|
||||
match = re.search(r'inter-arrival_s=([0-9.]+)', line)
|
||||
if match:
|
||||
intervals.append(float(match.group(1)))
|
||||
# %%
|
||||
|
||||
print(intervals)
|
||||
# %%
|
||||
import matplotlib.pyplot as plt
|
||||
plt.plot(intervals)
|
||||
plt.ylabel('Inter-arrival time (s)')
|
||||
plt.xlabel('Sample')
|
||||
plt.title('Inter-arrival time distribution')
|
||||
plt.show()
|
||||
# %%
|
||||
@@ -7,16 +7,23 @@ 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,
|
||||
)
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
|
||||
|
||||
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 +57,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 +78,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: another database on the
|
||||
`POSTGRES_*` server, and a second connection that is never the one Storage holds.
|
||||
|
||||
Raises:
|
||||
Exception: If any parent class initialization fails
|
||||
@@ -85,6 +102,15 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
secure=minio_config['secure'],
|
||||
)
|
||||
|
||||
mlflow_repository = MLFlowRepository(
|
||||
f'{mlflow_config["host"]}:{mlflow_config["port"]}',
|
||||
mlflow_config['username'],
|
||||
mlflow_config['password'],
|
||||
logger,
|
||||
notification_handler,
|
||||
metrics_controller,
|
||||
)
|
||||
|
||||
# Initialize parent classes
|
||||
Storage.__init__(
|
||||
self,
|
||||
@@ -104,10 +130,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
|
||||
MLFlow.__init__(
|
||||
self,
|
||||
mlflow_host=mlflow_config['host'],
|
||||
mlflow_port=mlflow_config['port'],
|
||||
mlflow_username=mlflow_config['username'],
|
||||
mlflow_password=mlflow_config['password'],
|
||||
mlflow_repository=mlflow_repository,
|
||||
minio_repository=minio_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
@@ -147,6 +170,30 @@ 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,
|
||||
mlflow_repository=mlflow_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.
|
||||
@@ -164,6 +211,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
Storage.close(self)
|
||||
MLFlow.close(self)
|
||||
Gates.close(self)
|
||||
await OPC.close(self)
|
||||
await OPC.aclose(self)
|
||||
ModelMetrics.close(self)
|
||||
API.close(self)
|
||||
ModelImport.close(self)
|
||||
|
||||
@@ -20,6 +20,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.utils.formatters import create_sample_dict
|
||||
|
||||
from laborious.utils.dataframe_debug import build_dataframe_debug_message
|
||||
from laborious.utils.messages import MISSING_MINIO_REPOSITORY
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
from laborious.utils.repository.minio_manager import MinioManager
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
@@ -37,10 +38,6 @@ class MLFlow(MinioManager):
|
||||
MLFlow operations, ensuring reliable model inference in production environments.
|
||||
|
||||
Attributes:
|
||||
mlflow_host (str): MLFlow server hostname
|
||||
mlflow_port (int): MLFlow server port
|
||||
mlflow_username (str): MLFlow authentication username
|
||||
mlflow_password (str): MLFlow authentication password
|
||||
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
|
||||
"""
|
||||
|
||||
@@ -48,45 +45,25 @@ class MLFlow(MinioManager):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mlflow_host: str,
|
||||
mlflow_port: int,
|
||||
mlflow_username: str,
|
||||
mlflow_password: str,
|
||||
mlflow_repository: MLFlowRepository,
|
||||
minio_repository: MinioRepository | None = None,
|
||||
logger: Logger | None = None,
|
||||
notification_handler: NotificationHandler | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize MLFlow activities with server configuration.
|
||||
Initialize MLFlow activities with an already-built repository.
|
||||
|
||||
Args:
|
||||
mlflow_host: MLFlow server hostname or IP address
|
||||
mlflow_port: MLFlow server port number
|
||||
mlflow_username: Username for MLFlow authentication
|
||||
mlflow_password: Password for MLFlow authentication
|
||||
mlflow_repository: Repository for MLFlow operations, built and owned by the caller
|
||||
(shared with other activities rather than rebuilt here)
|
||||
logger: Logger instance for observability and debugging
|
||||
notification_handler: Notification handler for alerts and monitoring
|
||||
|
||||
Raises:
|
||||
Exception: If MLFlowRepository initialization fails
|
||||
"""
|
||||
MinioManager.__init__(
|
||||
self, minio_repository, logger, notification_handler, metrics_controller
|
||||
)
|
||||
self.mlflow_host = mlflow_host
|
||||
self.mlflow_port = mlflow_port
|
||||
self.mlflow_username = mlflow_username
|
||||
self.mlflow_password = mlflow_password
|
||||
|
||||
self.model_monitoring_repository = MLFlowRepository(
|
||||
f'{mlflow_host}:{mlflow_port}',
|
||||
mlflow_username,
|
||||
mlflow_password,
|
||||
logger,
|
||||
notification_handler,
|
||||
metrics_controller,
|
||||
)
|
||||
self.model_monitoring_repository = mlflow_repository
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
@@ -237,6 +214,9 @@ class MLFlow(MinioManager):
|
||||
- data (dict): Transformed data for prediction
|
||||
- model_name (str): Name of the MLFlow model to use
|
||||
- model_retention (int): Model retention period in minutes
|
||||
Optional keys (in model_config):
|
||||
- skip_predict_timestamp (bool): If True, don't reinsert the
|
||||
timestamp column from the index before predicting. Default False.
|
||||
|
||||
Returns:
|
||||
dict: Prediction results from MLFlow model
|
||||
@@ -258,10 +238,11 @@ class MLFlow(MinioManager):
|
||||
# Convert numpy.nan to None for model compatibility
|
||||
data.replace(np.nan, None, inplace=True)
|
||||
|
||||
data['timestamp'] = data.index
|
||||
data['timestamp'] = to_datetime(
|
||||
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
|
||||
).dt.strftime(DATETIME_FORMAT)
|
||||
if not model_config.get('skip_predict_timestamp', False):
|
||||
data['timestamp'] = data.index
|
||||
data['timestamp'] = to_datetime(
|
||||
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
|
||||
).dt.strftime(DATETIME_FORMAT)
|
||||
|
||||
# Request prediction from MLFlow model
|
||||
response_data = await self.model_monitoring_repository.predict(
|
||||
@@ -334,7 +315,7 @@ class MLFlow(MinioManager):
|
||||
"""
|
||||
|
||||
if self.minio_repository is None:
|
||||
raise ValueError('Minio repository not initialized')
|
||||
raise ValueError(MISSING_MINIO_REPOSITORY)
|
||||
|
||||
metadata = input_data['metadata']
|
||||
|
||||
|
||||
1250
laborious/activities/model_import.py
Normal file
1250
laborious/activities/model_import.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -106,15 +106,23 @@ class OPC(SientiaMonitoring):
|
||||
4. Logs connection status for operational visibility
|
||||
|
||||
Raises:
|
||||
Exception: If OPC repository initialization fails or connection
|
||||
establishment encounters critical errors
|
||||
ValueError: A server whose `cert_path` is set but whose `server_uri`
|
||||
(or private key) is missing fails `OpcRepository.set_security()`
|
||||
with no try/except between here and `main()`. This is a
|
||||
configuration error, not a connectivity one, and is intentionally
|
||||
fatal to the whole worker process — fix the deployment config
|
||||
rather than catch it here.
|
||||
|
||||
Note:
|
||||
Connection failures are logged and reported but do not prevent
|
||||
the initialization of other OPC servers. Each server is handled
|
||||
independently to ensure maximum availability.
|
||||
Connection failures (server unreachable, timeout, etc.) are logged
|
||||
and reported but do not prevent the initialization of other OPC
|
||||
servers — each server's `connect()` call catches those on its own.
|
||||
This isolation does NOT extend to configuration errors raised before
|
||||
a connection is attempted (see `Raises` above): a single
|
||||
misconfigured server intentionally stops every other server, and the
|
||||
whole worker, from starting.
|
||||
"""
|
||||
self.logger.info('Initializing OPC servers...')
|
||||
self.info('Initializing OPC servers...')
|
||||
for opc_id, server in self.opc_servers.items():
|
||||
self.opc_repository[opc_id] = OpcRepository(
|
||||
opc_id=server['id'],
|
||||
@@ -145,9 +153,7 @@ class OPC(SientiaMonitoring):
|
||||
attachment_content=error_data.get('attachment_content', None),
|
||||
)
|
||||
else:
|
||||
self.logger.info(
|
||||
f'OPC server {opc_id}:{server["server_name"]} connected successfully.'
|
||||
)
|
||||
self.info(f'OPC server {opc_id}:{server["server_name"]} connected successfully.')
|
||||
|
||||
async def write_data(
|
||||
self,
|
||||
@@ -491,7 +497,7 @@ class OPC(SientiaMonitoring):
|
||||
|
||||
return data.to_dict()
|
||||
|
||||
async def close(self):
|
||||
async def aclose(self):
|
||||
"""
|
||||
Gracefully shutdown all OPC server connections and cleanup resources.
|
||||
|
||||
|
||||
@@ -17,7 +17,9 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.temporal.activities.postgres import Postgres
|
||||
from sientia_do.temporal.constants import now
|
||||
|
||||
from laborious.utils.messages import MISSING_MINIO_REPOSITORY
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
from laborious.utils.time import to_naive
|
||||
|
||||
_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'})
|
||||
|
||||
@@ -83,7 +85,7 @@ class Storage(Postgres, MinioManager):
|
||||
dict[str, Any]: Flat ``MinioDataFramePayload`` dict or ``success: False`` on failure.
|
||||
"""
|
||||
if self.minio_repository is None:
|
||||
raise ValueError('Minio repository not initialized')
|
||||
raise ValueError(MISSING_MINIO_REPOSITORY)
|
||||
|
||||
metadata: dict = input_data.get('metadata', {})
|
||||
model_name = input_data['model_name']
|
||||
@@ -129,6 +131,11 @@ class Storage(Postgres, MinioManager):
|
||||
"""
|
||||
Delete objects under the given prefixes that are older than the retention window.
|
||||
|
||||
See also: `ModelImport.ensure_import_bundle_retention`
|
||||
(laborious/activities/model_import.py) — a newer, unrelated lifecycle-policy expiry
|
||||
mechanism for the import bundle prefix. The two are independent by design for now; see
|
||||
docs/minio-object-expiry.md.
|
||||
|
||||
Args (input_data):
|
||||
metadata (dict): Workflow metadata for logging and metrics.
|
||||
prefixes (list[str]): Key prefixes to scan (one level or subtree per prefix).
|
||||
@@ -137,15 +144,13 @@ class Storage(Postgres, MinioManager):
|
||||
dict[str, Any]: ``success``, ``deleted_count``, and optional ``message``.
|
||||
"""
|
||||
if self.minio_repository is None:
|
||||
raise ValueError('Minio repository not initialized')
|
||||
raise ValueError(MISSING_MINIO_REPOSITORY)
|
||||
|
||||
metadata = input_data.get('metadata', {})
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
prefix = payload.cleanup_prefix()
|
||||
base = now()
|
||||
cutoff = (base.replace(tzinfo=None) if base.tzinfo else base) - timedelta(
|
||||
hours=self.retention_hours
|
||||
)
|
||||
cutoff = to_naive(base) - timedelta(hours=self.retention_hours)
|
||||
|
||||
report: dict[str, Any] = {
|
||||
'failed': {},
|
||||
@@ -202,7 +207,8 @@ class Storage(Postgres, MinioManager):
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close Storage resources (MinIO client and Postgres engine)."""
|
||||
Postgres.close(self)
|
||||
if hasattr(self, 'engine'):
|
||||
Postgres.close(self)
|
||||
MinioManager.close(self)
|
||||
|
||||
def __del__(self):
|
||||
|
||||
@@ -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()
|
||||
495
laborious/utils/bundle/reader.py
Normal file
495
laborious/utils/bundle/reader.py
Normal file
@@ -0,0 +1,495 @@
|
||||
"""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 5 — chunk size for the byte-counted extraction copy loop.
|
||||
_EXTRACT_CHUNK_BYTES = 1024 * 1024
|
||||
|
||||
# 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:
|
||||
entries = self._gate_4_inspect(zip_path)
|
||||
self._gate_5_extract(zip_path, root, entries)
|
||||
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) -> list[zipfile.ZipInfo]:
|
||||
"""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)
|
||||
|
||||
return entries
|
||||
|
||||
@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, entries: list[zipfile.ZipInfo]) -> 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. `entries` is the central-directory
|
||||
listing gate 4 already read; the archive is still opened here to read member bytes, but its
|
||||
`infolist()` is not read a second time.
|
||||
"""
|
||||
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 entries:
|
||||
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.
|
||||
#
|
||||
# Copied in chunks, counting bytes actually written, instead of `shutil.copyfileobj`:
|
||||
# `entry.file_size` is central-directory metadata gate 4 already checked in aggregate, but a
|
||||
# DEFLATE stream can be crafted to decompress to far more bytes than its own header
|
||||
# declares. Aborting as soon as the running count exceeds `file_size` catches that member
|
||||
# before the excess ever reaches disk, rather than trusting the declared size alone.
|
||||
written = 0
|
||||
with archive.open(entry) as source, resolved.open('wb') as sink:
|
||||
while chunk := source.read(_EXTRACT_CHUNK_BYTES):
|
||||
written += len(chunk)
|
||||
if written > entry.file_size:
|
||||
raise BundleGateError(
|
||||
f'archive member {entry.filename!r} decompresses to more than its '
|
||||
f'declared {entry.file_size} bytes',
|
||||
gate=5,
|
||||
)
|
||||
sink.write(chunk)
|
||||
|
||||
# ------------------------------------------------------------------ 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 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,89 @@ 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_MODEL_ID_COUNTERS_COLLECTION: MongoDB collection holding the atomic id counters that
|
||||
assign `models` document ids (default: model_id_counters). **External dependency**: this
|
||||
collection and its seed value (`seq = max(int(id) for id in models)`) are not created by
|
||||
this repository — a migration must add them to
|
||||
`sientia-dataops-database-migrations/mongodb/migrations/`, following the `models`
|
||||
baseline's `ensureCollection`/`ensureIndexes` convention, before any environment's counter
|
||||
is first incremented. See `openspec/changes/model-import-fix-bugs/design.md`, Decision 3.
|
||||
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'),
|
||||
'counters_collection': getenv('IMPORT_MODEL_ID_COUNTERS_COLLECTION', 'model_id_counters'),
|
||||
'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 — the POSTGRES_* server, another database.
|
||||
|
||||
The import log is `public.experiment_run` in the database owned by the Spring Boot BFF
|
||||
(`sientia-core-mlops-bff`), which sits on the same Postgres server as the `sientia` database
|
||||
Laborious itself uses. Only the database name differs, so only the database name has a variable
|
||||
of its own; host, port, user and password come from POSTGRES_*. The table and schema are
|
||||
constants of the code, never inputs.
|
||||
|
||||
The four IMPORT_STATUS_DB_* connection variables remain as overrides, for a deployment where the
|
||||
log genuinely lives on another server — the e2e suite is exactly that, two containers.
|
||||
|
||||
Environment Variables:
|
||||
IMPORT_STATUS_DB_NAME: database (default: sientia-core-mlops-bff)
|
||||
POSTGRES_HOST: server host, unless IMPORT_STATUS_DB_HOST overrides it (default: localhost)
|
||||
POSTGRES_PORT: server port, unless IMPORT_STATUS_DB_PORT overrides it (default: 5432)
|
||||
POSTGRES_USER: user, unless IMPORT_STATUS_DB_USER overrides it (default: sientia)
|
||||
POSTGRES_PASSWORD: password, unless IMPORT_STATUS_DB_PASSWORD overrides it (default: empty)
|
||||
|
||||
Returns:
|
||||
dict: import status database configuration dictionary
|
||||
"""
|
||||
# Written as `override or fallback or default` rather than
|
||||
# `getenv(override) or getenv(fallback, default)`: inside an `or`, the two-argument `getenv`
|
||||
# widens to `str | None` for mypy, so the `int(...)` below fails to type-check. Putting the
|
||||
# default last as its own term keeps every value a plain `str` and reads in precedence order.
|
||||
return {
|
||||
'host': getenv('IMPORT_STATUS_DB_HOST') or getenv('POSTGRES_HOST') or 'localhost',
|
||||
'port': int(getenv('IMPORT_STATUS_DB_PORT') or getenv('POSTGRES_PORT') or '5432'),
|
||||
'dbname': getenv('IMPORT_STATUS_DB_NAME') or 'sientia-core-mlops-bff',
|
||||
'user': getenv('IMPORT_STATUS_DB_USER') or getenv('POSTGRES_USER') or 'sientia',
|
||||
'password': getenv('IMPORT_STATUS_DB_PASSWORD') or getenv('POSTGRES_PASSWORD') or '',
|
||||
'schema': 'public',
|
||||
}
|
||||
|
||||
97
laborious/utils/import_password.py
Normal file
97
laborious/utils/import_password.py
Normal file
@@ -0,0 +1,97 @@
|
||||
"""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
|
||||
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 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 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
|
||||
1
laborious/utils/messages.py
Normal file
1
laborious/utils/messages.py
Normal file
@@ -0,0 +1 @@
|
||||
MISSING_MINIO_REPOSITORY = 'Minio repository not initialized'
|
||||
546
laborious/utils/models/import_errors.py
Normal file
546
laborious/utils/models/import_errors.py
Normal file
@@ -0,0 +1,546 @@
|
||||
"""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 import deque
|
||||
from collections.abc import Iterator, 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'
|
||||
|
||||
|
||||
class ImportResponsibility(StrEnum):
|
||||
"""Who can act on a step's failure: the user (re-export, re-type, re-submit) or only support."""
|
||||
|
||||
USER = 'user'
|
||||
SUPPORT = 'support'
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImportFailure:
|
||||
"""One classification of one failure: the step, its code, its sentence and whose to fix."""
|
||||
|
||||
step: ImportStep | None
|
||||
code: ImportErrorCode
|
||||
reason: str
|
||||
responsibility: ImportResponsibility = ImportResponsibility.SUPPORT
|
||||
|
||||
|
||||
_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, ImportResponsibility]]
|
||||
] = 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.',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
ImportStep.DOWNLOAD: (
|
||||
ImportErrorCode.FILE_UNREADABLE,
|
||||
'The uploaded file could not be read, or it is not the file that was sent. '
|
||||
'Upload it again.',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
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.',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
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}',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
ImportStep.EXTRACTION: (
|
||||
ImportErrorCode.BUNDLE_UNPACK_FAILED,
|
||||
f'The file could not be unpacked safely and nothing from it was kept. {_RE_EXPORT}',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
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}',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
ImportStep.CONTENT_POLICY: (
|
||||
ImportErrorCode.BUNDLE_UNEXPECTED_CONTENT,
|
||||
'The file contains something a model export should not contain, so it was '
|
||||
f'rejected. {_RE_EXPORT}',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
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}',
|
||||
ImportResponsibility.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}',
|
||||
ImportResponsibility.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}',
|
||||
ImportResponsibility.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}',
|
||||
ImportResponsibility.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}',
|
||||
)
|
||||
|
||||
# The three explicit-code failures are not pipeline steps, so `classify_import_failure` falls
|
||||
# through to `UNCLASSIFIED_FAILURE`'s reason for all of them. That sentence says "no model was
|
||||
# created", which is false for a cleanup-only failure (the model was imported) and misleading for
|
||||
# the other two (neither ever attempted the pipeline). Each gets its own reason here instead.
|
||||
_OVERRIDE_REASONS: Final[Mapping[ImportErrorCode, str]] = MappingProxyType(
|
||||
{
|
||||
ImportErrorCode.CLEANUP_INCOMPLETE: (
|
||||
'The model was imported successfully; only its temporary files could not be removed '
|
||||
'from the worker. No action needed.'
|
||||
),
|
||||
ImportErrorCode.LOG_ROW_NOT_CLAIMABLE: 'The import request could not be started.',
|
||||
ImportErrorCode.STATUS_NOT_RECORDED: 'The import outcome could not be recorded.',
|
||||
}
|
||||
)
|
||||
|
||||
# `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) -> ImportFailure:
|
||||
"""Turn a failure into the one code and the one sentence that will be recorded.
|
||||
|
||||
The step alone decides the outcome — an exception's text carries paths, table names and type
|
||||
names, and belongs in the log, never in the record.
|
||||
|
||||
Args:
|
||||
step: the step that was being attempted, or `None` when none had been entered.
|
||||
|
||||
Returns:
|
||||
ImportFailure: step, code and the reason to store.
|
||||
"""
|
||||
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, responsibility = IMPORT_FAILURE_CATALOG[step]
|
||||
return ImportFailure(step=step, code=code, reason=reason, responsibility=responsibility)
|
||||
|
||||
|
||||
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 ImportMisconfiguredError(ApplicationError):
|
||||
"""A required repository was not injected into this worker. A deploy defect, never transient."""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
super().__init__(message, type='ImportMisconfiguredError', 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
|
||||
|
||||
|
||||
_SUMMARY_LIMIT = 400
|
||||
_CHAIN_LIMIT = 10
|
||||
_UNKNOWN_ERROR_TYPE = 'unavailable'
|
||||
|
||||
|
||||
def walk_cause_chain(
|
||||
exception: BaseException, limit: int = _CHAIN_LIMIT
|
||||
) -> Iterator[BaseException]:
|
||||
"""Yield `exception`, then each `__cause__`, up to `limit` links. Bounded against a cycle."""
|
||||
candidate: BaseException | None = exception
|
||||
for _ in range(limit):
|
||||
if candidate is None:
|
||||
return
|
||||
yield candidate
|
||||
candidate = candidate.__cause__
|
||||
|
||||
|
||||
def _root_cause(exception: BaseException) -> BaseException:
|
||||
"""The deepest `__cause__` in the chain: the failure that actually happened.
|
||||
|
||||
A `MlflowException` reaches this channel wrapped — by an `ActivityError`, by an
|
||||
`ApplicationError`, by `requests`' own re-raises — and the outermost type is the wrapper's, which
|
||||
says nothing an operator can act on. The chain is walked to its end, and bounded, because a cycle
|
||||
in `__cause__` would otherwise hang the reporting of a failure.
|
||||
"""
|
||||
return deque(walk_cause_chain(exception), maxlen=1)[-1]
|
||||
|
||||
|
||||
def _type_name(exception: BaseException) -> str:
|
||||
"""The name to report a failure under: `ApplicationError.type` when there is one.
|
||||
|
||||
An activity's own exception type is the one thing worth grouping failures by, and Temporal keeps
|
||||
it in `ApplicationError.type` rather than in the Python class — a `MlflowException` raised in an
|
||||
activity reaches the workflow as an `ApplicationError` whose `type` still says `MlflowException`.
|
||||
Reporting the class name here would say `ApplicationError` for every activity failure there is.
|
||||
"""
|
||||
return str(getattr(exception, 'type', None) or type(exception).__name__)
|
||||
|
||||
|
||||
def _one_line(text: str) -> str:
|
||||
"""The first non-empty line, bounded — a log *field* has to stay one field.
|
||||
|
||||
The traceback still reaches the log in full, in the message. This is the value a dashboard groups
|
||||
by and an alert quotes, so a 40-frame `urllib3` chain would make it useless.
|
||||
"""
|
||||
for line in text.splitlines():
|
||||
stripped = line.strip()
|
||||
if stripped:
|
||||
return stripped[:_SUMMARY_LIMIT]
|
||||
return ''
|
||||
|
||||
|
||||
def summarize_failure(
|
||||
exception: BaseException | None,
|
||||
detail_text: str | None = None,
|
||||
error_type: str | None = None,
|
||||
) -> tuple[str, str]:
|
||||
"""`(error_type, error_summary)` — the two technical fields every failure log carries.
|
||||
|
||||
Args:
|
||||
exception: the failure, when the object itself is to hand.
|
||||
detail_text: the failure as text, for a caller whose exception crossed an activity boundary.
|
||||
error_type: the type name that travelled with such a text, when it is known.
|
||||
|
||||
Returns:
|
||||
tuple: the exception type name and a single-line summary of what it said.
|
||||
"""
|
||||
if exception is not None:
|
||||
cause = _root_cause(exception)
|
||||
return _type_name(cause), _one_line(str(cause))
|
||||
return error_type or _UNKNOWN_ERROR_TYPE, _one_line(detail_text or '')
|
||||
|
||||
|
||||
def describe_failure_chain(exception: BaseException) -> tuple[str, str]:
|
||||
"""`(error_type, detail)` for a failure that is about to lose its exception object.
|
||||
|
||||
The workflow is the one place that holds the whole chain — `ActivityError` wrapping the
|
||||
`ApplicationError` that carries the activity's own type and message — and it cannot hand that
|
||||
object to the channel, because the channel runs in an activity. So the chain is flattened here,
|
||||
**root cause first**, one link per line: the first line is what actually broke, which is what
|
||||
becomes `error_summary`, and the wrappers follow for whoever reads further.
|
||||
|
||||
`ApplicationError.type` is preferred over the Python class name because it is the only place the
|
||||
activity's original exception type survives the boundary: `MlflowException` arrives as an
|
||||
`ApplicationError` whose `type` says so.
|
||||
"""
|
||||
links: list[str] = []
|
||||
for candidate in walk_cause_chain(exception):
|
||||
kind = _type_name(candidate)
|
||||
message = _one_line(str(candidate))
|
||||
# `ApplicationError.__str__` already prefixes its own type, and `MlflowException:
|
||||
# MlflowException: …` reads like two failures rather than one.
|
||||
if not message:
|
||||
links.append(kind)
|
||||
elif message.startswith(f'{kind}:'):
|
||||
links.append(message)
|
||||
else:
|
||||
links.append(f'{kind}: {message}')
|
||||
links.reverse()
|
||||
root = links[0].split(':', 1)[0] if links else _UNKNOWN_ERROR_TYPE
|
||||
return root, '\n'.join(links)
|
||||
|
||||
|
||||
def _log_fields(
|
||||
*,
|
||||
failure: ImportFailure,
|
||||
where: str,
|
||||
import_run_id: Any,
|
||||
workflow_id: str,
|
||||
gate: int | None,
|
||||
error_type: str,
|
||||
error_summary: str,
|
||||
verdict: str | None,
|
||||
metadata: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""The workflow metadata plus this failure's own fields, as `logfmt` keys.
|
||||
|
||||
Every key is prefixed `import_` or `error_` so none of them can collide with the metadata the
|
||||
rest of the runtime already logs (`model_id`, `model_name`, `schedule_name`, `workflow_name`) or
|
||||
with the formatter's own (`at`, `msg`, `timestamp`). Values are single-line by construction: the
|
||||
traceback stays in the message, where a newline costs nothing.
|
||||
"""
|
||||
fields: dict[str, Any] = {
|
||||
**(metadata or {}),
|
||||
'import_run_id': import_run_id,
|
||||
'import_workflow_id': workflow_id,
|
||||
'import_step': where,
|
||||
'import_code': failure.code.value,
|
||||
'error_type': error_type,
|
||||
}
|
||||
if error_summary:
|
||||
fields['error_summary'] = error_summary
|
||||
if gate is not None:
|
||||
fields['import_gate'] = gate
|
||||
if verdict:
|
||||
fields['import_verdict_not_recorded'] = verdict
|
||||
return fields
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
if exception is None:
|
||||
return None, None
|
||||
for candidate in walk_cause_chain(exception):
|
||||
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
|
||||
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,
|
||||
error_type: str | None = None,
|
||||
gate: int | 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.
|
||||
|
||||
Both log lines this produces — the technical one here and the user-facing sentence the
|
||||
notification logs — carry the same `logfmt` fields, built once by `_log_fields`. That is what
|
||||
makes the two greppable and joinable: `import_run_id=29` finds the pair, `import_code=` groups
|
||||
them, and `error_type=` says what actually broke without reading a traceback. The message text is
|
||||
kept as it was, ids and all, so a query written against it still matches.
|
||||
|
||||
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.
|
||||
error_type: the exception type name that travelled with such a text, when it is known.
|
||||
gate: the gate that rejected the bundle, when one did — a field, never a sentence.
|
||||
|
||||
Returns:
|
||||
ImportFailure: the classification that was reported.
|
||||
"""
|
||||
failure = classify_import_failure(step)
|
||||
if code is not None:
|
||||
failure = ImportFailure(
|
||||
step=step, code=code, reason=_OVERRIDE_REASONS.get(code, 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'
|
||||
kind, summary = summarize_failure(exception, detail_text, error_type)
|
||||
fields = _log_fields(
|
||||
failure=failure,
|
||||
where=where,
|
||||
import_run_id=import_run_id,
|
||||
workflow_id=workflow_id,
|
||||
gate=gate,
|
||||
error_type=kind,
|
||||
error_summary=summary,
|
||||
verdict=verdict,
|
||||
metadata=metadata,
|
||||
)
|
||||
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}',
|
||||
fields,
|
||||
)
|
||||
|
||||
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=fields,
|
||||
notification_id=failure.code.value,
|
||||
message=failure.reason,
|
||||
block=where,
|
||||
level='ERROR',
|
||||
)
|
||||
return failure
|
||||
File diff suppressed because it is too large
Load Diff
@@ -17,6 +17,7 @@ from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.utils.models.import_errors import walk_cause_chain
|
||||
|
||||
# Requested session and secure channel lifetime (ms) before server revision; 10 minutes.
|
||||
OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS = 10 * 60 * 1000
|
||||
@@ -78,7 +79,7 @@ def _opc_authentication_token_str(client: Client | None) -> str:
|
||||
if tok is None:
|
||||
return 'unknown'
|
||||
return str(tok)
|
||||
except Exception:
|
||||
except Exception: # best-effort label for logs/metrics, must not fail loud
|
||||
return 'unknown'
|
||||
|
||||
|
||||
@@ -92,11 +93,9 @@ def _opc_status_from_exception(exc: BaseException) -> str:
|
||||
Return:
|
||||
str: Status class name or generic Python exception name.
|
||||
"""
|
||||
current: BaseException | None = exc
|
||||
while current is not None:
|
||||
if isinstance(current, UaStatusCodeError):
|
||||
return type(current).__name__
|
||||
current = current.__cause__
|
||||
for candidate in walk_cause_chain(exc):
|
||||
if isinstance(candidate, UaStatusCodeError):
|
||||
return type(candidate).__name__
|
||||
return type(exc).__name__
|
||||
|
||||
|
||||
@@ -262,12 +261,14 @@ class OpcRepository(SientiaMonitoring):
|
||||
Configure certificates and timeouts on the asyncua client.
|
||||
|
||||
Raises:
|
||||
ValueError: If cert paths or client are missing.
|
||||
ValueError: If cert paths, server URI or client are missing.
|
||||
"""
|
||||
if self.cert_path is None or self.private_key_path is None:
|
||||
raise ValueError(
|
||||
'Certificate and private key paths must be provided for secure connection.'
|
||||
)
|
||||
if self.server_uri is None:
|
||||
raise ValueError('server_uri must be provided for secure connection.')
|
||||
|
||||
cert = Path(self.cert_path)
|
||||
private_key = Path(self.private_key_path)
|
||||
@@ -304,7 +305,6 @@ class OpcRepository(SientiaMonitoring):
|
||||
|
||||
self.client = Client(self.url, timeout=10, watchdog_intervall=50) # type: ignore[attr-defined]
|
||||
self.client.name = self.pod_id
|
||||
self.client.application_name = self.pod_id
|
||||
pod_uri = self.pod_id.replace('-', ':')
|
||||
self.client.application_uri = pod_uri
|
||||
self.client.product_uri = pod_uri
|
||||
|
||||
6
laborious/utils/time.py
Normal file
6
laborious/utils/time.py
Normal file
@@ -0,0 +1,6 @@
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
def to_naive(moment: datetime) -> datetime:
|
||||
"""Strip tzinfo, if any. Postgres `TIMESTAMP WITHOUT TIME ZONE` columns need naive values."""
|
||||
return moment.replace(tzinfo=None) if moment.tzinfo else moment
|
||||
@@ -1,73 +0,0 @@
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.observability.logger import Logger
|
||||
from temporalio.client import Client
|
||||
from temporalio.worker import PollerBehaviorAutoscaling, Worker
|
||||
|
||||
parameters = [
|
||||
('MAX_CONCURRENT_WORKFLOW_TASKS', '200'),
|
||||
('MAX_CONCURRENT_ACTIVITIES', '200'),
|
||||
('MAX_CONCURRENT_LOCAL_ACTIVITIES', '200'),
|
||||
('MAX_CACHED_WORKFLOWS', '200'),
|
||||
('WORKFLOW_POLLER_BEHAVIOUR_MINIMUM', '10'),
|
||||
('WORKFLOW_POLLER_BEHAVIOUR_INITIAL', '100'),
|
||||
('WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM', '200'),
|
||||
('ACTIVITY_POLLER_BEHAVIOUR_MINIMUM', '10'),
|
||||
('ACTIVITY_POLLER_BEHAVIOUR_INITIAL', '100'),
|
||||
('ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM', '200'),
|
||||
]
|
||||
|
||||
|
||||
def camel_to_snake(text: str) -> str:
|
||||
"""Convert camelCase or PascalCase to snake_case."""
|
||||
text = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', text)
|
||||
text = re.sub('([a-z0-9])([A-Z])', r'\1_\2', text)
|
||||
return text.lower()
|
||||
|
||||
|
||||
def prepare_worker(
|
||||
main_workflow: type,
|
||||
other_workflows: Sequence[type],
|
||||
activities: Sequence[Any],
|
||||
temporal_client: Client,
|
||||
logger: Logger,
|
||||
) -> Worker:
|
||||
main_workflow_name = main_workflow.__name__.upper()
|
||||
|
||||
queue_name = f'{camel_to_snake(main_workflow.__name__)}-queue'
|
||||
|
||||
local_workflow_parameters = {}
|
||||
|
||||
for parameter in parameters:
|
||||
local_workflow_parameters[parameter[0]] = int(
|
||||
os.getenv(main_workflow_name + '_' + parameter[0], parameter[1])
|
||||
)
|
||||
|
||||
logger.info(f'Preparing worker for {main_workflow_name} with queue {queue_name}')
|
||||
logger.info(f'Worker runtime config: {local_workflow_parameters}')
|
||||
|
||||
return Worker(
|
||||
temporal_client,
|
||||
task_queue=queue_name,
|
||||
workflows=[main_workflow, *other_workflows],
|
||||
activities=[*activities],
|
||||
max_concurrent_workflow_tasks=local_workflow_parameters['MAX_CONCURRENT_WORKFLOW_TASKS'],
|
||||
max_concurrent_activities=local_workflow_parameters['MAX_CONCURRENT_ACTIVITIES'],
|
||||
max_concurrent_local_activities=local_workflow_parameters[
|
||||
'MAX_CONCURRENT_LOCAL_ACTIVITIES'
|
||||
],
|
||||
max_cached_workflows=local_workflow_parameters['MAX_CACHED_WORKFLOWS'],
|
||||
workflow_task_poller_behavior=PollerBehaviorAutoscaling(
|
||||
minimum=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIOUR_MINIMUM'],
|
||||
initial=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIOUR_INITIAL'],
|
||||
maximum=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM'],
|
||||
),
|
||||
activity_task_poller_behavior=PollerBehaviorAutoscaling(
|
||||
minimum=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIOUR_MINIMUM'],
|
||||
initial=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIOUR_INITIAL'],
|
||||
maximum=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM'],
|
||||
),
|
||||
)
|
||||
@@ -5,12 +5,16 @@ This module provides the main worker implementation for the Sientia DataOps Labo
|
||||
It orchestrates Temporal workers, manages task queues, and handles the lifecycle of
|
||||
prediction and retraining workflows.
|
||||
|
||||
The worker supports multiple task queues:
|
||||
- predictions_batch-queue: Handles batch prediction workflows (heavy workload)
|
||||
Includes activities for MLFlow, data quality gates, OPC export, PI Web API export, and PostgreSQL
|
||||
- minimal_retrain-queue: Handles model retraining workflows
|
||||
- drift-queue: Handles drift detection workflows
|
||||
- simple_metrics-queue: Handles simple metrics calculation workflows
|
||||
The worker supports multiple runtime-scoped task queues (via ``sientia_do.temporal.worker.prepare_worker``):
|
||||
- predictions_batch-{runtime}-queue: Batch prediction workflows (heavy workload)
|
||||
- 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``).
|
||||
|
||||
Key Features:
|
||||
- Resource-based scaling with WorkerTuner (CPU and memory aware)
|
||||
@@ -21,12 +25,16 @@ Key Features:
|
||||
- Multiple worker instances for different workflow types
|
||||
|
||||
Environment Variables:
|
||||
- RUNTIME: Required non-empty string; suffix for all task queue names
|
||||
- TEMPORAL_HOST: Temporal server address (default: localhost:7233)
|
||||
- TEMPORAL_NAMESPACE: Temporal namespace (default: laborious)
|
||||
- POD_ID: Kubernetes pod identifier for metrics
|
||||
- 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_NAME: import configuration, see
|
||||
``laborious.utils.connectors_config`` (the import log's server is POSTGRES_*)
|
||||
"""
|
||||
|
||||
from temporalio import client, workflow
|
||||
@@ -40,6 +48,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
from prometheus_client import start_http_server
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.logger import get_logger
|
||||
from sientia_do.temporal.worker.prepare_worker import prepare_worker
|
||||
from sientia_do.utils.connectors_config import (
|
||||
build_api_config,
|
||||
build_mongodb_config,
|
||||
@@ -49,12 +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.worker.prepare_worker import prepare_worker
|
||||
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
|
||||
@@ -66,6 +77,19 @@ with workflow.unsafe.imports_passed_through():
|
||||
POD_ID = os.getenv('HOSTNAME')
|
||||
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
|
||||
|
||||
# The values that turn the import worker on. Anything else — unset, empty, `0`, `false`, a typo —
|
||||
# leaves the four default workers alone and polls no import queue at all.
|
||||
IMPORT_ENABLED_VALUES = ('1', 'true', 'yes', 'on')
|
||||
|
||||
|
||||
def import_worker_enabled() -> bool:
|
||||
"""Whether this runtime registers the fifth (import) worker.
|
||||
|
||||
Read at worker start, not per workflow: containment (QTZPOC-19) wants the importer in its own
|
||||
runtime with its own credentials, and this flag is what makes that a deployment decision.
|
||||
"""
|
||||
return os.getenv('IMPORT_MODEL_ENABLED', '').strip().lower() in IMPORT_ENABLED_VALUES
|
||||
|
||||
|
||||
async def main():
|
||||
"""
|
||||
@@ -100,10 +124,21 @@ async def main():
|
||||
|
||||
logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
|
||||
|
||||
logger.custom_info('Starting prometheus client...', metadata)
|
||||
start_prometheus_server()
|
||||
runtime = os.getenv('RUNTIME', '').strip()
|
||||
if not runtime:
|
||||
logger.custom_critical(
|
||||
'RUNTIME environment variable is required and must be non-empty',
|
||||
metadata,
|
||||
)
|
||||
metrics.APP_UP.labels(pod_id=POD_ID).set(0)
|
||||
sys.exit(1)
|
||||
|
||||
logger.custom_info('Starting Notification Handler...', metadata)
|
||||
metadata_runtime = {**metadata, 'runtime': runtime}
|
||||
|
||||
logger.custom_info('Starting prometheus client...', metadata_runtime)
|
||||
start_prometheus_server(logger, metadata_runtime)
|
||||
|
||||
logger.custom_info('Starting Notification Handler...', metadata_runtime)
|
||||
|
||||
mongo_config = build_mongodb_config()
|
||||
notification_handler = NotificationHandler(
|
||||
@@ -113,7 +148,7 @@ async def main():
|
||||
project_name=os.getenv('PROJECT_NAME', 'laborious'),
|
||||
)
|
||||
|
||||
logger.custom_info('Starting Activities...', metadata)
|
||||
logger.custom_info('Starting Activities...', metadata_runtime)
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=build_postgres_config(),
|
||||
@@ -123,12 +158,18 @@ 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)
|
||||
logger.custom_info('Initializing OPC...', metadata_runtime)
|
||||
await activities.init_opc()
|
||||
|
||||
logger.custom_info(f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
|
||||
logger.custom_info(
|
||||
f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...',
|
||||
metadata_runtime,
|
||||
)
|
||||
|
||||
new_runtime = Runtime(
|
||||
telemetry=TelemetryConfig(
|
||||
@@ -136,7 +177,7 @@ async def main():
|
||||
)
|
||||
)
|
||||
|
||||
logger.custom_info(f'Starting Temporal Client at {host}...', metadata)
|
||||
logger.custom_info(f'Starting Temporal Client at {host}...', metadata_runtime)
|
||||
|
||||
temporal_client = await client.Client.connect(
|
||||
target_host=host,
|
||||
@@ -144,7 +185,7 @@ async def main():
|
||||
runtime=new_runtime,
|
||||
)
|
||||
|
||||
logger.custom_info('Starting Workers...', metadata)
|
||||
logger.custom_info(f'Starting Workers (runtime={runtime})...', metadata_runtime)
|
||||
|
||||
workers = [
|
||||
prepare_worker(
|
||||
@@ -159,6 +200,7 @@ async def main():
|
||||
activities.export_data_to_postgres,
|
||||
],
|
||||
logger=logger,
|
||||
runtime=runtime,
|
||||
),
|
||||
prepare_worker(
|
||||
temporal_client=temporal_client,
|
||||
@@ -170,6 +212,7 @@ async def main():
|
||||
activities.export_data_to_postgres,
|
||||
],
|
||||
logger=logger,
|
||||
runtime=runtime,
|
||||
),
|
||||
prepare_worker(
|
||||
temporal_client=temporal_client,
|
||||
@@ -182,6 +225,7 @@ async def main():
|
||||
activities.export_data_to_postgres,
|
||||
],
|
||||
logger=logger,
|
||||
runtime=runtime,
|
||||
),
|
||||
prepare_worker(
|
||||
temporal_client=temporal_client,
|
||||
@@ -211,33 +255,76 @@ async def main():
|
||||
activities.write_pi_web_api_data,
|
||||
],
|
||||
logger=logger,
|
||||
runtime=runtime,
|
||||
),
|
||||
]
|
||||
|
||||
# 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 import_worker_enabled():
|
||||
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,
|
||||
activities.report_import_pipeline_failure,
|
||||
# Bundle download/open, MLflow experiment creation and artifact upload — one
|
||||
# fused activity, so they never split across worker processes.
|
||||
activities.prepare_and_upload_import_bundle,
|
||||
# MLflow provisioning
|
||||
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())
|
||||
|
||||
logger.custom_info('Workers started successfully', metadata)
|
||||
logger.custom_info('Workers started successfully', metadata_runtime)
|
||||
|
||||
exit_code = 0
|
||||
try:
|
||||
# This will run the workers and wait for them to complete.
|
||||
# If an exception occurs in any of the worker handlers, it will be propagated here.
|
||||
await asyncio.gather(*handlers)
|
||||
except BaseException as e: # NOSONAR
|
||||
logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
|
||||
except Exception as e:
|
||||
logger.custom_error(f'An unhandled exception occurred: {e}', metadata_runtime)
|
||||
exit_code = 1
|
||||
finally:
|
||||
# No `except (SystemExit, KeyboardInterrupt)` on purpose: `finally` runs on every path,
|
||||
# caught or not, so cleanup still happens for a deliberate shutdown (Ctrl+C, a handler's
|
||||
# own sys.exit) without us having to catch and re-raise it ourselves — which would only
|
||||
# let us relabel its exit code before `sys.exit(exit_code)` below discards it anyway.
|
||||
# Letting it propagate on its own keeps the original code (e.g. a handler's sys.exit(N))
|
||||
# intact instead of collapsing it to 0.
|
||||
if notification_handler:
|
||||
notification_handler.shutdown()
|
||||
if activities:
|
||||
await activities.shutdown()
|
||||
metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN
|
||||
sys.exit(exit_code)
|
||||
|
||||
sys.exit(exit_code)
|
||||
|
||||
|
||||
def start_prometheus_server():
|
||||
def start_prometheus_server(logger, metadata: dict):
|
||||
"""
|
||||
Starts the Prometheus metrics server for monitoring and observability.
|
||||
|
||||
@@ -257,10 +344,10 @@ def start_prometheus_server():
|
||||
try:
|
||||
port = int(os.getenv('HTTP_METRICS_PORT', 9090))
|
||||
start_http_server(port)
|
||||
print(f'Prometheus server started on port {port}.')
|
||||
logger.custom_info(f'Prometheus server started on port {port}.', metadata)
|
||||
metrics.APP_UP.labels(pod_id=POD_ID).set(1) # Mark app as UP
|
||||
except Exception as e:
|
||||
print(f'Failed to start Prometheus server: {e}')
|
||||
logger.custom_critical(f'Failed to start Prometheus server: {e}', metadata)
|
||||
os._exit(1)
|
||||
|
||||
|
||||
|
||||
448
laborious/workflows/import_model.py
Normal file
448
laborious/workflows/import_model.py
Normal file
@@ -0,0 +1,448 @@
|
||||
"""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.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from datetime import timedelta
|
||||
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.temporal.policies import retry_policy
|
||||
from temporalio.exceptions import is_cancelled_exception
|
||||
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.activities.model_import import (
|
||||
IMPORT_STATUS_HINT_RETRY_POLICY,
|
||||
IMPORT_STATUS_RETRY_POLICY,
|
||||
)
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from laborious.utils.models.import_errors import (
|
||||
ImportInputError,
|
||||
describe_failure_chain,
|
||||
failure_location,
|
||||
)
|
||||
|
||||
# `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, ≈451 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=30)
|
||||
STATUS_WRITE_ENVELOPE = timedelta(seconds=900)
|
||||
# Download, decryption/extraction, experiment creation and artifact upload now run inside one
|
||||
# activity attempt (bug 8): the budget is the sum of what the four separate activities used to get,
|
||||
# since a retry of the fused activity redoes all four phases.
|
||||
PREPARE_AND_UPLOAD_TIMEOUT = (
|
||||
timedelta(seconds=1800)
|
||||
+ timedelta(seconds=1800)
|
||||
+ timedelta(seconds=300)
|
||||
+ timedelta(seconds=1800)
|
||||
)
|
||||
REGISTRATION_TIMEOUT = timedelta(seconds=600)
|
||||
DOCUMENT_TIMEOUT = timedelta(seconds=300)
|
||||
CLEANUP_TIMEOUT = timedelta(seconds=600)
|
||||
RETENTION_TIMEOUT = timedelta(seconds=300)
|
||||
|
||||
# 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,
|
||||
)
|
||||
|
||||
step: str = ImportStep.RECEIVED.value
|
||||
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)
|
||||
|
||||
# Download, decryption/extraction, experiment creation and artifact upload run inside
|
||||
# one activity attempt (bug 8), so no later phase can depend on a file a different,
|
||||
# separately-scheduled activity wrote to a different worker's disk. One hint covers the
|
||||
# whole fused step; an exception it raises with no location of its own falls back to
|
||||
# this step, since that is what the workflow was attempting.
|
||||
step = ImportStep.DOWNLOAD.value
|
||||
await self._record_hint(
|
||||
base, Activities.record_import_status, {'step': step, 'source': None}
|
||||
)
|
||||
# Achado #9: the retention check carries no data this step consumes and already
|
||||
# swallows its own failure (`_ensure_retention`), so it runs alongside the fused
|
||||
# download+decrypt+provision activity instead of blocking it. Achado #8 fused
|
||||
# `download_import_bundle` into `prepare_and_upload_import_bundle` before this change
|
||||
# landed, so this gathers against that fused activity, not the original isolated one.
|
||||
_, prepared = await asyncio.gather(
|
||||
self._ensure_retention(base, input_data),
|
||||
workflow.execute_activity_method(
|
||||
Activities.prepare_and_upload_import_bundle,
|
||||
{
|
||||
**base,
|
||||
'bucket': input_data.get('bucket'),
|
||||
'object_key': input_data['object_key'],
|
||||
'expected_digest': input_data['expected_digest'],
|
||||
'password_envelope': input_data['password_envelope'],
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=PREPARE_AND_UPLOAD_TIMEOUT,
|
||||
),
|
||||
)
|
||||
source = prepared['metadata']
|
||||
parameters = prepared['parameters']
|
||||
model_name = source['model_name']
|
||||
metadata['model_name'] = model_name
|
||||
run_name = prepared['run_name']
|
||||
|
||||
# The names are known only once the fused activity returns; they are still recorded
|
||||
# before the two provisioning steps that follow (the registered version and the model
|
||||
# document), which is the guarantee `model-import-status-record` now makes.
|
||||
await self._record_hint(
|
||||
base,
|
||||
Activities.record_import_names,
|
||||
{
|
||||
'experiment_name': source['experiment_name'],
|
||||
'run_name': run_name,
|
||||
'source': source,
|
||||
'step': ImportStep.EXPERIMENT_CREATION.value,
|
||||
},
|
||||
)
|
||||
|
||||
step = ImportStep.REGISTRATION.value
|
||||
await self._record_hint(
|
||||
base,
|
||||
Activities.record_import_status,
|
||||
{'step': step, 'source': source or None},
|
||||
)
|
||||
version = await workflow.execute_activity_method(
|
||||
Activities.register_import_model_version,
|
||||
{**base, 'model_name': model_name, 'run_id': prepared['run_id']},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=REGISTRATION_TIMEOUT,
|
||||
)
|
||||
|
||||
step = ImportStep.MODEL_DOCUMENT.value
|
||||
await self._record_hint(
|
||||
base,
|
||||
Activities.record_import_status,
|
||||
{'step': step, 'source': source or None},
|
||||
)
|
||||
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': prepared['run_id'],
|
||||
'model_name': model_name,
|
||||
'version': version['version'],
|
||||
'document_id': document.get('id'),
|
||||
}
|
||||
except BaseException as error: # classified once below, then re-raised untouched
|
||||
located_step, located_gate = failure_location(error)
|
||||
step = located_step or step
|
||||
gate = located_gate
|
||||
await self._report_failure(base, error, step=step, gate=gate)
|
||||
raise
|
||||
finally:
|
||||
cleanup = await self._cleanup(base)
|
||||
await self._record_terminal(
|
||||
base,
|
||||
succeeded=succeeded,
|
||||
step=step,
|
||||
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 _record_hint(
|
||||
self,
|
||||
base: dict[str, Any],
|
||||
activity: Callable[..., Awaitable[Any]],
|
||||
extra: dict[str, Any],
|
||||
) -> None:
|
||||
"""Record a status hint. A hint: its failure never fails the import."""
|
||||
try:
|
||||
await workflow.execute_activity_method(
|
||||
activity,
|
||||
{**base, **extra},
|
||||
retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY,
|
||||
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
|
||||
)
|
||||
except BaseException as error: # a lost hint must not fail a working import
|
||||
if is_cancelled_exception(error):
|
||||
raise
|
||||
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: # the outcome stands even when it cannot be written
|
||||
if is_cancelled_exception(error):
|
||||
raise
|
||||
await self._report_status_write_failure(
|
||||
base,
|
||||
error,
|
||||
verdict=f'succeeded={succeeded} step={step} gate={gate}',
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _report_failure(
|
||||
base: dict[str, Any],
|
||||
error: BaseException,
|
||||
*,
|
||||
step: str | None,
|
||||
gate: int | None,
|
||||
) -> None:
|
||||
"""Route the failure that ended this import through the one failure channel.
|
||||
|
||||
Runs from the `except`, before cleanup and before the terminal write, so the log reads in the
|
||||
order things happened: what failed, then what was cleaned, then what was recorded. The step
|
||||
and the gate are the ones already located from the failure, so the log line and the row can
|
||||
never name different ones.
|
||||
|
||||
Its own failure is swallowed. The channel is the thing that logs, so a report that cannot be
|
||||
made has nowhere to say so, and the failure it was reporting is re-raised either way — the
|
||||
event history and the record both still carry it.
|
||||
"""
|
||||
error_type, detail = describe_failure_chain(error)
|
||||
try:
|
||||
await workflow.execute_activity_method(
|
||||
Activities.report_import_pipeline_failure,
|
||||
{**base, 'step': step, 'gate': gate, 'detail': detail, 'error_type': error_type},
|
||||
retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY,
|
||||
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
|
||||
)
|
||||
except BaseException as error: # the report of a report is where this stops
|
||||
if is_cancelled_exception(error):
|
||||
raise
|
||||
|
||||
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 as error: # the report of a report is where this stops
|
||||
if is_cancelled_exception(error):
|
||||
raise
|
||||
# Deliberately silent otherwise: the channel this was trying to reach is the thing that
|
||||
# logs, so there is nowhere left to say it, and an import that worked must not fail
|
||||
# over it.
|
||||
|
||||
# ----------------------------------------------------------- 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 the failure is dropped here and left
|
||||
where Temporal already records it — the activity's own failure in the event history.
|
||||
"""
|
||||
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 as error: # retention is not a precondition of an import
|
||||
if is_cancelled_exception(error):
|
||||
raise
|
||||
# Dropped, not reported otherwise: the activity's own failure is already in the event
|
||||
# history, and a bucket policy the importer cannot set is an operational problem rather
|
||||
# than a bad import. Failing a valid model import over a lifecycle API would be the
|
||||
# wrong trade.
|
||||
|
||||
@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 as error: # cleanup never replaces the failure that brought us here
|
||||
if is_cancelled_exception(error):
|
||||
raise
|
||||
return {'cleaned': False}
|
||||
@@ -5,6 +5,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.temporal.policies import retry_policy
|
||||
from temporalio.exceptions import ApplicationError
|
||||
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
@@ -87,7 +88,11 @@ class MinimalRetrain:
|
||||
|
||||
storage_payload = MinioDataFramePayload.from_dict(storage_result)
|
||||
if not storage_payload.has_data():
|
||||
raise ValueError('No data returned from query')
|
||||
raise ApplicationError(
|
||||
'No data returned from query',
|
||||
type='NoDataReturnedError',
|
||||
non_retryable=True,
|
||||
)
|
||||
|
||||
experiment_response = await workflow.execute_activity_method(
|
||||
Activities.retrain_model,
|
||||
|
||||
@@ -1,106 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 23,
|
||||
"id": "e838ff21",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import csv\n",
|
||||
"\n",
|
||||
"def csv_to_tag_lists(csv_path: str) -> dict:\n",
|
||||
" read_tags = []\n",
|
||||
" write_tags = []\n",
|
||||
"\n",
|
||||
" def to_float(val):\n",
|
||||
" try:\n",
|
||||
" return float(str(val).strip())\n",
|
||||
" except Exception:\n",
|
||||
" return None\n",
|
||||
"\n",
|
||||
" with open(csv_path, newline=\"\", encoding=\"utf-8\") as f:\n",
|
||||
" reader = csv.DictReader(f)\n",
|
||||
" for row in reader:\n",
|
||||
" # Basic normalization\n",
|
||||
" op = (row.get(\"operation\") or \"\").strip()\n",
|
||||
"\n",
|
||||
" if op == \"READ\":\n",
|
||||
" # Build common tag payload with required mappings\n",
|
||||
" tag = {\n",
|
||||
" \"server_id\": \"1\",\n",
|
||||
" \"tag_address\": row.get(\"opc_tag\"),\n",
|
||||
" \"tag_name\": row.get(\"name\"),\n",
|
||||
" \"data_range\": [to_float(row.get(\"min_value\")), to_float(row.get(\"max_value\"))],\n",
|
||||
" \"aggr_func\": row.get(\"aggregation_func\").lower(),\n",
|
||||
" # keep other fields with their original names\n",
|
||||
" \"frequency\": row.get(\"frequency\"),\n",
|
||||
" \"local\": row.get(\"local\"),\n",
|
||||
" \"area\": row.get(\"area\"),\n",
|
||||
" \"description\": row.get(\"description\"),\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" read_tags.append(tag)\n",
|
||||
"\n",
|
||||
" else:\n",
|
||||
" tag = {\n",
|
||||
" \"server_id\": \"1\",\n",
|
||||
" \"addr\": row.get(\"opc_tag\"),\n",
|
||||
" \"tag_name\": row.get(\"name\"),\n",
|
||||
" \"local\": row.get(\"local\"),\n",
|
||||
" \"area\": row.get(\"area\"),\n",
|
||||
" \"description\": row.get(\"description\"),\n",
|
||||
" }\n",
|
||||
" \n",
|
||||
" if op == \"WRITE_PREDICTION\":\n",
|
||||
" tag[\"type\"] = \"prediction\"\n",
|
||||
" write_tags.append(tag)\n",
|
||||
" elif op == \"WRITE_CONFIDENCE\":\n",
|
||||
" tag[\"type\"] = \"confidence\"\n",
|
||||
" write_tags.append(tag)\n",
|
||||
" # ignore any other operation values silently\n",
|
||||
"\n",
|
||||
" return {\"read_tags\": read_tags, \"write_tags\": write_tags}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "4621cd43",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import json\n",
|
||||
"\n",
|
||||
"file_names = [\"Courier - Página1.csv\"]\n",
|
||||
"\n",
|
||||
"for file_name in file_names:\n",
|
||||
" write_file = file_name.replace(\".csv\", \".json\")\n",
|
||||
"\n",
|
||||
" with open(write_file, \"w\", encoding=\"utf-8\") as f:\n",
|
||||
" json.dump(csv_to_tag_lists(file_name), f, indent=2, ensure_ascii=False)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "venv",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "laborious"
|
||||
version = "0.0.0"
|
||||
version = "0.9.1"
|
||||
description = "Sientia DataOps Laborious - ML Model Orchestration System"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -60,6 +60,16 @@ ignore = [
|
||||
"S105", # hardcoded passwords ok in tests
|
||||
"S106", # hardcoded passwords ok in tests
|
||||
]
|
||||
"e2e/conftest.py" = [
|
||||
# Imports below the sientia/ModelAnalysis stub are load-order critical: the stub
|
||||
# must be in sys.modules before Activities is imported, so they cannot move up.
|
||||
"E402",
|
||||
]
|
||||
"scripts/**/*.py" = [
|
||||
# Notebook-style scripts split into `# %%` cells: each cell imports what it needs
|
||||
# where it needs it, so imports legitimately appear below other statements.
|
||||
"E402",
|
||||
]
|
||||
|
||||
[tool.ruff.lint.mccabe]
|
||||
max-complexity = 15
|
||||
@@ -116,12 +126,33 @@ python_functions = ["test_*"]
|
||||
addopts = [
|
||||
"-v",
|
||||
"--strict-markers",
|
||||
# pytest>=9.1 has a known bug where its unraisableexception plugin crashes
|
||||
# (tracemalloc partially-initialized AttributeError) when 2+ unraisable
|
||||
# exceptions land close together — e.g. "coroutine was never awaited" from
|
||||
# AsyncMock-mocked sync methods (metrics_controller, minio_repository) being
|
||||
# GC'd. Harmless mock artifacts turned into a hard ERROR by the plugin itself.
|
||||
"-p", "no:unraisableexception",
|
||||
]
|
||||
filterwarnings = [
|
||||
# `kafka-python`, pulled in transitively by `sientia_do`, reads its protocol schemas through the
|
||||
# legacy `importlib.resources` API that Python 3.11 deprecated. It emits ~109 warnings per run,
|
||||
# twice each (once at the call site, once inside `importlib`), which is enough noise to bury a
|
||||
# warning worth reading. Nothing here can fix it — the change belongs in `kafka-python`.
|
||||
#
|
||||
# Matched on the message rather than by module so the filter cannot quietly swallow an unrelated
|
||||
# deprecation: any *other* DeprecationWarning, including one from our own code, still shows up.
|
||||
"ignore:read_text is deprecated:DeprecationWarning",
|
||||
"ignore:open_text is deprecated:DeprecationWarning",
|
||||
]
|
||||
markers = [
|
||||
"asyncio: marks tests as async",
|
||||
"integration: marks tests as integration tests",
|
||||
"unit: marks tests as unit tests",
|
||||
"opc: marks tests that use the in-process OPC UA server (OpcRepository E2E)",
|
||||
# The .sientia import e2e suite. Outside `testpaths`, so `validate --only-tests` and the
|
||||
# coverage gate never collect it; run it on request with `pytest e2e -m import_e2e`.
|
||||
"import_e2e: marks the model-import end-to-end tests (needs Docker and an MLflow server)",
|
||||
"slow: marks tests that run on a real clock and take tens of seconds",
|
||||
]
|
||||
|
||||
[tool.coverage.run]
|
||||
@@ -148,6 +179,7 @@ exclude_lines = [
|
||||
"if TYPE_CHECKING:",
|
||||
"class .*\\bProtocol\\):",
|
||||
"@(abc\\.)?abstractmethod",
|
||||
"@overload",
|
||||
]
|
||||
|
||||
[tool.coverage.html]
|
||||
|
||||
@@ -13,9 +13,13 @@ types-requests>=2.31.0 # Type stubs for requests
|
||||
pytest>=7.4.0 # Testing framework
|
||||
pytest-cov>=4.1.0 # Coverage plugin for pytest
|
||||
pytest-asyncio>=0.21.0 # Async test support (already in main requirements)
|
||||
testcontainers[postgres,minio] # PostgreSQL and MinIO containers for E2E tests
|
||||
testcontainers[postgres,minio,mongodb]>=4.15.0 # PostgreSQL, MinIO and MongoDB containers for E2E tests
|
||||
# Floor is the `testcontainers.community.*` namespace: the flat
|
||||
# `testcontainers.minio` / `.postgres` / `.mongodb` paths the e2e suite used
|
||||
# to import are deprecated shims that warn on import.
|
||||
|
||||
# Development Tools
|
||||
ipython>=8.12.0 # Enhanced Python shell
|
||||
ipdb>=0.13.13 # IPython debugger
|
||||
ipykernel==6.30.1 # IPython kernel for Jupyter notebooks
|
||||
python-dotenv>=1.0.0 # .env loading for scripts/import_model_manual_run.py
|
||||
@@ -1,12 +1,21 @@
|
||||
temporalio
|
||||
psycopg2-binary
|
||||
sqlalchemy
|
||||
asyncua
|
||||
sqlalchemy<2.1
|
||||
asyncua==1.0.6
|
||||
redis
|
||||
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.12.0
|
||||
sientia_do>=1.12.2
|
||||
mlflow==2.10.1
|
||||
prometheus-client
|
||||
botocore
|
||||
boto3
|
||||
s3fs
|
||||
pyarrow
|
||||
mlflow
|
||||
kaleido
|
||||
hyperopt
|
||||
shap
|
||||
pycurl
|
||||
scipy<1.14.0
|
||||
scikit-learn==1.4.2
|
||||
pycaret==3.3.2
|
||||
joblib
|
||||
pynacl==1.5.0
|
||||
|
||||
26
requirements-local.txt
Normal file
26
requirements-local.txt
Normal file
@@ -0,0 +1,26 @@
|
||||
temporalio
|
||||
psycopg2-binary
|
||||
sqlalchemy<2.1
|
||||
asyncua==1.0.6
|
||||
redis
|
||||
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.12.2
|
||||
git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.41.1
|
||||
prometheus-client
|
||||
# fsspec/s3fs/gcsfs pinned as a set: nannyml (via the mlops library) pulls both s3fs and
|
||||
# gcsfs, and each pins fsspec exactly, so a free `s3fs` makes the resolver fail against
|
||||
# evidently's fsspec>=2024.2.0. Versions match the library-distribution-controller freeze.
|
||||
fsspec==2026.7.0
|
||||
s3fs==2026.7.0
|
||||
gcsfs==2026.8.0
|
||||
# aiobotocore 3.9.0 (pulled by s3fs) requires botocore<1.43.57,>=1.43.3.
|
||||
botocore==1.43.56
|
||||
boto3==1.43.56
|
||||
pyarrow
|
||||
kaleido
|
||||
hyperopt
|
||||
shap
|
||||
pycurl
|
||||
scipy<1.14.0
|
||||
scikit-learn==1.4.2
|
||||
pycaret==3.3.2
|
||||
joblib
|
||||
@@ -1,18 +1,33 @@
|
||||
temporalio
|
||||
psycopg2-binary
|
||||
sqlalchemy
|
||||
asyncua
|
||||
sqlalchemy<2.1
|
||||
asyncua==1.0.6
|
||||
redis
|
||||
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.10.4
|
||||
git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.41.0
|
||||
sientia_do>=1.12.2
|
||||
sientia>=0.41.1
|
||||
prometheus-client
|
||||
botocore
|
||||
boto3
|
||||
s3fs
|
||||
# fsspec/s3fs/gcsfs pinned as a set: nannyml (via sientia) pulls both s3fs and gcsfs, and
|
||||
# each pins fsspec exactly, so a free `s3fs` makes the resolver fail against evidently's
|
||||
# fsspec>=2024.2.0. Versions match the library-distribution-controller freeze.
|
||||
fsspec==2026.7.0
|
||||
s3fs==2026.7.0
|
||||
gcsfs==2026.8.0
|
||||
# aiobotocore 3.9.0 (pulled by s3fs) requires botocore<1.43.57,>=1.43.3.
|
||||
botocore==1.43.56
|
||||
boto3==1.43.56
|
||||
pyarrow
|
||||
kaleido
|
||||
hyperopt
|
||||
shap
|
||||
pycurl
|
||||
scipy<1.14.0
|
||||
scikit-learn==1.5.2
|
||||
scikit-learn==1.4.2
|
||||
pycaret==3.3.2
|
||||
joblib
|
||||
# .sientia bundle import (Argon2id + XChaCha20-Poly1305 secretstream): same pin as the
|
||||
# producer in sientia-projects-templates, so both sides open the same bytes.
|
||||
pynacl==1.5.0
|
||||
# AES-256-GCM password envelope; already present transitively, declared explicitly.
|
||||
cryptography
|
||||
# MongoDB models collection (the model listing); transitive via sientia_do, declared explicitly.
|
||||
pymongo
|
||||
|
||||
586
scripts/import_model_manual_run.py
Normal file
586
scripts/import_model_manual_run.py
Normal file
@@ -0,0 +1,586 @@
|
||||
"""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 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:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
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`."""
|
||||
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 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}'
|
||||
|
||||
|
||||
# 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 = 'import_model-legacy-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 = 'localhost'
|
||||
|
||||
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.
|
||||
|
||||
|
||||
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`).
|
||||
#
|
||||
# `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.
|
||||
|
||||
|
||||
envelope_key = base64.b64decode(env_required('IMPORT_PASSWORD_KEY'), validate=True)
|
||||
nonce = secrets.token_bytes(12)
|
||||
sealed = AESGCM(envelope_key).encrypt(
|
||||
nonce, env_secret_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.
|
||||
|
||||
|
||||
_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.
|
||||
|
||||
|
||||
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 = True
|
||||
|
||||
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 ---')
|
||||
|
||||
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
|
||||
# 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]):
|
||||
_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 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"run_id='{_run_id}'"):
|
||||
_registered_names.add(_version.name)
|
||||
_registered_versions.append((_version.name, str(_version.version)))
|
||||
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}')
|
||||
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 ---')
|
||||
|
||||
_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.
|
||||
#
|
||||
# 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:
|
||||
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:
|
||||
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:
|
||||
print(f' could not delete the registered model: {type(error).__name__}: {error}')
|
||||
_removed = _collection.delete_many({'name': _registered}).deleted_count
|
||||
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')
|
||||
|
||||
status_engine.dispose()
|
||||
mongo_client.close()
|
||||
@@ -1,4 +1,4 @@
|
||||
sonar.projectKey=Aignosi_sientia-dataops-laborious_temporal_beaec423-6c42-4f26-8134-b676287b499d
|
||||
sonar.projectKey=Aignosi_sientia-dataops-laborious_temporal_ca1a7039-6db9-49e5-be78-54d29bc93e4f
|
||||
sonar.projectName=sientia-dataops-laborious_temporal
|
||||
sonar.sources=laborious
|
||||
sonar.tests=tests
|
||||
@@ -8,4 +8,4 @@ sonar.qualitygate.wait=true
|
||||
sonar.qualitygate.timeout=300
|
||||
sonar.python.coverage.reportPaths=coverage.xml
|
||||
sonar.python.xunit.reportPath=pytest.xml
|
||||
sonar.python.version=3.11
|
||||
sonar.python.version=3.11
|
||||
938
tests.ipynb
938
tests.ipynb
File diff suppressed because one or more lines are too long
86
tests/fixtures/bundle/README.md
vendored
Normal file
86
tests/fixtures/bundle/README.md
vendored
Normal file
@@ -0,0 +1,86 @@
|
||||
# Golden `.sientia` fixtures
|
||||
|
||||
Binary conformance vectors for `laborious/utils/bundle/`. They are **bytes the producer's own code
|
||||
emitted**, not bytes this repository knows how to write: there is no production writer here (change
|
||||
`import-model-workflow`, design D1). If the producer's layout moves, the reader's tests fail against
|
||||
these files instead of failing in production.
|
||||
|
||||
| File | What it is |
|
||||
|---|---|
|
||||
| `golden_model_v3.sientia` | `format_version` 1 bundle, 110767 bytes, plaintext zip 110569 bytes — **two chunks**, so `MESSAGE`/`FINAL` framing is exercised by the golden fixture itself |
|
||||
| `golden_model_v3.metadata.json` | the `metadata.json` that went into that bundle, in plaintext, so a test can assert the recovered value without decrypting first |
|
||||
| `regenerate.py` | the script that produced them, importing the producer's `bundle.py`/`crypto.py` by path |
|
||||
|
||||
## Password
|
||||
|
||||
```
|
||||
Import1Model!Fixture
|
||||
```
|
||||
|
||||
Recorded deliberately: the fixture protects nothing. It satisfies the producer's export-side
|
||||
strength rule (≥ 12 characters, upper, lower, digit, symbol), which is checked by
|
||||
`crypto.validate_password_strength` at generation time.
|
||||
|
||||
## Provenance
|
||||
|
||||
- Producer repository: `sientia-projects-templates`
|
||||
- Producer commit: `8d3c1fd117c697e8b352c34f11c732809a4a129b`
|
||||
(`8d3c1fd QTZPOC-7: adapt to lint rules picked up from merging main`)
|
||||
- Producer modules used: `app/src/operations/model_export/bundle.py`,
|
||||
`app/src/operations/model_export/crypto.py` — loaded by path, so every constant, the `struct`
|
||||
layout, the zip layout, the KDF parameters and the chunk framing come from the producer and none of
|
||||
them is re-implemented here
|
||||
- Producer pin honoured: `pynacl==1.5.0` (this repository pins the same version)
|
||||
|
||||
One deviation from task 0.3, recorded rather than hidden: the bytes were generated with the
|
||||
producer's **code** under *this* repository's venv, not inside the producer's own venv, because that
|
||||
venv has neither `pynacl` nor `mlflow` installed. The dependency that matters for the bytes is
|
||||
`pynacl`, pinned identically on both sides, and `mlflow` is not reached — `regenerate.py` builds the
|
||||
artifact tree itself instead of calling `client.download_artifacts`.
|
||||
|
||||
## Contents of the bundle
|
||||
|
||||
```
|
||||
metadata.json (written first, no indent, by bundle.build_zip)
|
||||
artifacts/prediction_model/MLmodel
|
||||
artifacts/prediction_model/model.pkl (90 KiB of incompressible filler)
|
||||
artifacts/prediction_model/conda.yaml
|
||||
artifacts/prediction_model/python_env.yaml
|
||||
artifacts/prediction_model/requirements.txt
|
||||
artifacts/data_model/MLmodel
|
||||
artifacts/data_model/model.pkl (8 KiB)
|
||||
artifacts/data_model/conda.yaml
|
||||
artifacts/data_model/python_env.yaml
|
||||
artifacts/data_model/requirements.txt
|
||||
artifacts/model_card/model_card.json
|
||||
artifacts/train_data.csv
|
||||
artifacts/test_data.csv
|
||||
artifacts/raw_data.csv
|
||||
```
|
||||
|
||||
No directory entries are stored — the producer only adds paths for which `is_file()` is true.
|
||||
|
||||
`parameters` carries `target_variable: "SE"`, which gate 6 requires and which becomes
|
||||
`model_config.target` in the MongoDB model document.
|
||||
|
||||
## Digest boundary — the reason this fixture exists
|
||||
|
||||
The header's `ciphertext_digest` is `sha256(file[164:])`: the producer computes
|
||||
`hashlib.sha256(ciphertext)` over the joined chunks and only then writes
|
||||
`header + stream_header + ciphertext`, so the 24-byte `secretstream` header is **not** covered.
|
||||
Verified against this fixture and against a freshly produced bundle. Earlier prose (this change's own
|
||||
first draft, and `aig-docs/03-sientia-bundle-format.md`) said "everything after the plaintext
|
||||
header", i.e. from offset 140 — that is wrong, and `tests/laborious/utils/bundle/test_format.py`
|
||||
asserts both directions so it cannot come back.
|
||||
|
||||
## Regenerating
|
||||
|
||||
```bash
|
||||
source ./venv/bin/activate
|
||||
PRODUCER_REPO=../sientia-projects-templates python tests/fixtures/bundle/regenerate.py
|
||||
```
|
||||
|
||||
Two fields change on every run and are expected to: the 16-byte salt (`os.urandom`) and
|
||||
`metadata.metadata.export_timestamp` (`datetime.now(UTC)`). Everything else is deterministic —
|
||||
the filler bytes come from a seeded PRNG. After regenerating, update the producer commit above and
|
||||
the byte counts in the table.
|
||||
24
tests/fixtures/bundle/golden_model_v3.metadata.json
vendored
Normal file
24
tests/fixtures/bundle/golden_model_v3.metadata.json
vendored
Normal file
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"parameters": {
|
||||
"model_name": "fixture_model",
|
||||
"target_variable": "SE",
|
||||
"input_variables": "['feature_a', 'feature_b']",
|
||||
"models_params": "{'alpha': 0.1}",
|
||||
"nan_treatment": "drop",
|
||||
"scaler_name": "StandardScaler",
|
||||
"train_size": "0.8",
|
||||
"test_size": "0.2",
|
||||
"include_ar": "False",
|
||||
"retrain": "False",
|
||||
"start_date": "2026-01-01 00:00:00",
|
||||
"end_date": "2026-08-01 00:00:00"
|
||||
},
|
||||
"metadata": {
|
||||
"model_name": "fixture_model",
|
||||
"experiment_name": "fixture_experiment",
|
||||
"run_id": "8f14e45fceea167a5a36dedd4bea2543",
|
||||
"model_version": "3",
|
||||
"model_project": "fixture_project",
|
||||
"export_timestamp": "2026-08-25T12:08:59.533441+00:00"
|
||||
}
|
||||
}
|
||||
BIN
tests/fixtures/bundle/golden_model_v3.sientia
vendored
Normal file
BIN
tests/fixtures/bundle/golden_model_v3.sientia
vendored
Normal file
Binary file not shown.
185
tests/fixtures/bundle/regenerate.py
vendored
Normal file
185
tests/fixtures/bundle/regenerate.py
vendored
Normal file
@@ -0,0 +1,185 @@
|
||||
"""Regenerate the golden `.sientia` fixtures using the producer's own source.
|
||||
|
||||
This script is the auditable half of the anti-drift contract described in the change
|
||||
`import-model-workflow` (design D1): the committed fixtures are bytes produced by
|
||||
`sientia-projects-templates/app/src/operations/model_export/{bundle,crypto}.py`, loaded here by
|
||||
path so no constant is re-implemented. Regenerating is a deliberate act — see the README next to
|
||||
this file for the recorded producer commit and the fixture password.
|
||||
|
||||
Usage (from the repository root, inside the venv):
|
||||
|
||||
PRODUCER_REPO=../sientia-projects-templates ./venv/bin/python tests/fixtures/bundle/regenerate.py
|
||||
|
||||
The producer's own venv is not used: it does not have `pynacl` installed, while this repository
|
||||
pins the producer's exact version (`pynacl==1.5.0`). The bytes are therefore produced by the
|
||||
producer's code on the producer's dependency pin, which is what the fixture is for.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
FIXTURE_DIR = Path(__file__).resolve().parent
|
||||
DEFAULT_PRODUCER_REPO = FIXTURE_DIR.parents[2].parent / 'sientia-projects-templates'
|
||||
|
||||
# Recorded in the README alongside the fixtures; the password satisfies the producer's own
|
||||
# export-side strength rule (>= 12 chars, upper, lower, digit, symbol).
|
||||
FIXTURE_PASSWORD = 'Import1Model!Fixture'
|
||||
|
||||
# Deterministic filler so a regeneration differs only in the salt and the export timestamp.
|
||||
FILLER_SEED = 20260825
|
||||
|
||||
|
||||
def _load_producer_module(repo: Path, name: str) -> types.ModuleType:
|
||||
"""Import one producer module by path, without importing the Streamlit app."""
|
||||
path = repo / 'app' / 'src' / 'operations' / 'model_export' / f'{name}.py'
|
||||
if not path.is_file():
|
||||
raise SystemExit(f'producer module not found: {path}')
|
||||
spec = importlib.util.spec_from_file_location(f'producer_{name}', path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise SystemExit(f'cannot load producer module: {path}')
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _producer_commit(repo: Path) -> str:
|
||||
result = subprocess.run( # noqa: S603 - fixed argv, no shell
|
||||
['git', '-C', str(repo), 'rev-parse', 'HEAD'], # noqa: S607 - git from PATH is intended
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
return result.stdout.strip() or 'unknown'
|
||||
|
||||
|
||||
def _write_model_dir(root: Path, flavor: str, pickle_bytes: bytes) -> None:
|
||||
"""Write the files an MLflow 2.x model directory carries for this platform."""
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
(root / 'MLmodel').write_text(
|
||||
f'artifact_path: {flavor}\n'
|
||||
'flavors:\n'
|
||||
' python_function:\n'
|
||||
' loader_module: mlflow.sklearn\n'
|
||||
' model_path: model.pkl\n'
|
||||
' sklearn:\n'
|
||||
' pickled_model: model.pkl\n'
|
||||
' serialization_format: cloudpickle\n'
|
||||
' sklearn_version: 1.4.2\n'
|
||||
'mlflow_version: 2.10.1\n'
|
||||
'model_uuid: 0f9c1c1f9c1f4c0f9c1c1f9c1f4c0f9c\n'
|
||||
'utc_time_created: "2026-08-25 12:00:00.000000"\n',
|
||||
encoding='utf-8',
|
||||
)
|
||||
(root / 'model.pkl').write_bytes(pickle_bytes)
|
||||
(root / 'conda.yaml').write_text(
|
||||
'channels:\n- conda-forge\ndependencies:\n- python=3.11\n- pip\n- pip:\n'
|
||||
' - mlflow==2.10.1\n - scikit-learn==1.4.2\nname: mlflow-env\n',
|
||||
encoding='utf-8',
|
||||
)
|
||||
(root / 'python_env.yaml').write_text(
|
||||
'python: 3.11.15\nbuild_dependencies:\n- pip==24.0\ndependencies:\n- -r requirements.txt\n',
|
||||
encoding='utf-8',
|
||||
)
|
||||
(root / 'requirements.txt').write_text(
|
||||
'mlflow==2.10.1\nscikit-learn==1.4.2\ncloudpickle==3.0.0\n', encoding='utf-8'
|
||||
)
|
||||
|
||||
|
||||
def _write_csv(path: Path, rows: int, rng: random.Random) -> None:
|
||||
lines = ['datetime,feature_a,feature_b,SE']
|
||||
for i in range(rows):
|
||||
lines.append(
|
||||
f'2026-08-{(i % 27) + 1:02d}T00:00:00,'
|
||||
f'{rng.uniform(0, 100):.6f},{rng.uniform(0, 100):.6f},{rng.uniform(0, 100):.6f}'
|
||||
)
|
||||
path.write_text('\n'.join(lines) + '\n', encoding='utf-8')
|
||||
|
||||
|
||||
def _build_artifact_tree(root: Path) -> None:
|
||||
rng = random.Random(FILLER_SEED) # noqa: S311 - test filler, not key material
|
||||
# Incompressible filler keeps the zip above the 64 KiB chunk size, so the golden fixture
|
||||
# exercises multi-chunk `MESSAGE`/`FINAL` framing with real producer bytes.
|
||||
big_pickle = bytes(rng.getrandbits(8) for _ in range(90 * 1024))
|
||||
small_pickle = bytes(rng.getrandbits(8) for _ in range(8 * 1024))
|
||||
_write_model_dir(root / 'prediction_model', 'prediction_model', big_pickle)
|
||||
_write_model_dir(root / 'data_model', 'data_model', small_pickle)
|
||||
(root / 'model_card').mkdir(parents=True, exist_ok=True)
|
||||
(root / 'model_card' / 'model_card.json').write_text(
|
||||
json.dumps(
|
||||
{
|
||||
'model_name': 'fixture_model',
|
||||
'target': 'SE',
|
||||
'created_by': 'fixture',
|
||||
'creation_method': 'time_series_template',
|
||||
}
|
||||
),
|
||||
encoding='utf-8',
|
||||
)
|
||||
_write_csv(root / 'train_data.csv', 120, rng)
|
||||
_write_csv(root / 'test_data.csv', 40, rng)
|
||||
_write_csv(root / 'raw_data.csv', 200, rng)
|
||||
|
||||
|
||||
RUN_PARAMETERS = {
|
||||
'model_name': 'fixture_model',
|
||||
'target_variable': 'SE',
|
||||
'input_variables': "['feature_a', 'feature_b']",
|
||||
'models_params': "{'alpha': 0.1}",
|
||||
'nan_treatment': 'drop',
|
||||
'scaler_name': 'StandardScaler',
|
||||
'train_size': '0.8',
|
||||
'test_size': '0.2',
|
||||
'include_ar': 'False',
|
||||
'retrain': 'False',
|
||||
'start_date': '2026-01-01 00:00:00',
|
||||
'end_date': '2026-08-01 00:00:00',
|
||||
}
|
||||
|
||||
METADATA_FIELDS = {
|
||||
'model_name': 'fixture_model',
|
||||
'experiment_name': 'fixture_experiment',
|
||||
'run_id': '8f14e45fceea167a5a36dedd4bea2543',
|
||||
'model_version': '3',
|
||||
'model_project': 'fixture_project',
|
||||
}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
repo = Path(os.getenv('PRODUCER_REPO', str(DEFAULT_PRODUCER_REPO))).resolve()
|
||||
bundle = _load_producer_module(repo, 'bundle')
|
||||
crypto = _load_producer_module(repo, 'crypto')
|
||||
|
||||
crypto.validate_password_strength(FIXTURE_PASSWORD)
|
||||
|
||||
staging = Path(tempfile.mkdtemp(prefix='sientia_fixture_'))
|
||||
try:
|
||||
_build_artifact_tree(staging)
|
||||
metadata = bundle.build_metadata_json(run_parameters=RUN_PARAMETERS, **METADATA_FIELDS)
|
||||
zip_bytes = bundle.build_zip(staging, metadata)
|
||||
sientia_bytes = crypto.encrypt_archive(zip_bytes, FIXTURE_PASSWORD)
|
||||
finally:
|
||||
shutil.rmtree(staging, ignore_errors=True)
|
||||
|
||||
target = FIXTURE_DIR / 'golden_model_v3.sientia'
|
||||
target.write_bytes(sientia_bytes)
|
||||
(FIXTURE_DIR / 'golden_model_v3.metadata.json').write_text(
|
||||
json.dumps(metadata, indent=2) + '\n', encoding='utf-8'
|
||||
)
|
||||
print(f'wrote {target} ({len(sientia_bytes)} bytes, zip {len(zip_bytes)} bytes)')
|
||||
print(f'producer commit: {_producer_commit(repo)}')
|
||||
print(f'password: {FIXTURE_PASSWORD}')
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(main())
|
||||
0
tests/helpers/__init__.py
Normal file
0
tests/helpers/__init__.py
Normal file
295
tests/helpers/bundle_factory.py
Normal file
295
tests/helpers/bundle_factory.py
Normal file
@@ -0,0 +1,295 @@
|
||||
"""Test-only `.sientia` writer, mirrored field for field from the producer.
|
||||
|
||||
There is no production writer in this repository (change `import-model-workflow`, design D1): the
|
||||
producer is the Streamlit app in `sientia-projects-templates`. This factory exists for one reason —
|
||||
building the **hostile** variants the gate tests need, which the producer cannot emit. Its output is
|
||||
pinned against the golden fixture by `test_format.py::test_factory_header_matches_fixture_layout`,
|
||||
so a drift in either direction shows up as a failing test.
|
||||
|
||||
Mirrored from:
|
||||
- `app/src/operations/model_export/crypto.py` — header layout, KDF, `secretstream` framing
|
||||
- `app/src/operations/model_export/bundle.py` — zip layout and the two-top-level-key JSON
|
||||
|
||||
Deliberate deviation: the default KDF cost is a fraction of the producer's (8 MiB / 1 pass instead
|
||||
of 256 MiB / 3), because the reader honours whatever the header declares and a test suite that
|
||||
derives thirty 256 MiB Argon2id keys is a test suite nobody runs. The golden fixture is what covers
|
||||
the production cost parameters.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import stat
|
||||
import struct
|
||||
import unicodedata
|
||||
import zipfile
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import nacl.bindings as sodium
|
||||
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||||
|
||||
# --- mirrored from the producer's crypto.py ---
|
||||
MAGIC = b'SIENTIA1'
|
||||
FORMAT_VERSION = 1
|
||||
AEAD_ID = 1
|
||||
KDF_ID = 1
|
||||
SALT_SIZE = 16
|
||||
RESERVED_SIGNATURE_SIZE = 64
|
||||
CHUNK_SIZE = 64 * 1024
|
||||
HEADER_STRUCT = struct.Struct('>8sBBB16sQQB32s64s')
|
||||
HEADER_SIZE = HEADER_STRUCT.size
|
||||
PRODUCER_KDF_MEMLIMIT_BYTES = 256 * 1024 * 1024
|
||||
PRODUCER_KDF_OPSLIMIT = 3
|
||||
|
||||
# Cheap-but-valid cost parameters for the suite; see the module docstring.
|
||||
TEST_KDF_MEMLIMIT_BYTES = 8 * 1024 * 1024
|
||||
TEST_KDF_OPSLIMIT = 1
|
||||
|
||||
DEFAULT_PASSWORD = 'Factory1Model!Test'
|
||||
DEFAULT_ENVELOPE_KEY = bytes(range(32))
|
||||
|
||||
MODEL_DIRS = ('prediction_model', 'data_model')
|
||||
MODEL_FILES = ('MLmodel', 'model.pkl', 'conda.yaml', 'python_env.yaml', 'requirements.txt')
|
||||
CSV_FILES = ('train_data.csv', 'test_data.csv', 'raw_data.csv')
|
||||
|
||||
DEFAULT_METADATA = {
|
||||
'model_name': 'factory_model',
|
||||
'experiment_name': 'factory_experiment',
|
||||
'run_id': 'aa11bb22cc33dd44ee55ff6677889900',
|
||||
'model_version': '2',
|
||||
'model_project': 'factory_project',
|
||||
'export_timestamp': '2026-08-25T12:00:00+00:00',
|
||||
}
|
||||
|
||||
DEFAULT_PARAMETERS = {
|
||||
'model_name': 'factory_model',
|
||||
'target_variable': 'SE',
|
||||
'input_variables': "['feature_a', 'feature_b']",
|
||||
'retrain': 'False',
|
||||
}
|
||||
|
||||
|
||||
def default_tree() -> dict[str, bytes]:
|
||||
"""The artifact tree a well-formed bundle carries, as archive-name -> bytes."""
|
||||
files: dict[str, bytes] = {}
|
||||
for directory in MODEL_DIRS:
|
||||
for name in MODEL_FILES:
|
||||
files[f'artifacts/{directory}/{name}'] = f'{directory}/{name}\n'.encode()
|
||||
files['artifacts/model_card/model_card.json'] = json.dumps({'target': 'SE'}).encode()
|
||||
for name in CSV_FILES:
|
||||
files[f'artifacts/{name}'] = b'datetime,feature_a,SE\n2026-08-25T00:00:00,1.0,2.0\n'
|
||||
return files
|
||||
|
||||
|
||||
def build_metadata_document(
|
||||
*,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
parameters: dict[str, Any] | None = None,
|
||||
drop_metadata_field: str | None = None,
|
||||
drop_top_level_key: str | None = None,
|
||||
drop_parameter: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build `metadata.json`, with hooks for the malformed variants gate 6 must reject."""
|
||||
meta = {**DEFAULT_METADATA, **(metadata or {})}
|
||||
params = {**DEFAULT_PARAMETERS, **(parameters or {})}
|
||||
if drop_metadata_field:
|
||||
meta.pop(drop_metadata_field, None)
|
||||
if drop_parameter:
|
||||
params.pop(drop_parameter, None)
|
||||
document: dict[str, Any] = {'parameters': params, 'metadata': meta}
|
||||
if drop_top_level_key:
|
||||
document.pop(drop_top_level_key, None)
|
||||
return document
|
||||
|
||||
|
||||
@dataclass
|
||||
class ZipSpec:
|
||||
"""How to lay out the zip inside the bundle."""
|
||||
|
||||
files: dict[str, bytes] = field(default_factory=default_tree)
|
||||
document: dict[str, Any] | None = None
|
||||
include_metadata: bool = True
|
||||
symlinks: dict[str, str] = field(default_factory=dict)
|
||||
absolute_entry: str | None = None
|
||||
traversal_entry: str | None = None
|
||||
bomb_bytes: int = 0
|
||||
extra_files: dict[str, bytes] = field(default_factory=dict)
|
||||
# name -> a raw `stat.S_IF*` mode, for an entry that is neither a symlink, a regular file nor a
|
||||
# directory (e.g. `stat.S_IFIFO`) — gate 4 must reject these on their file-type bits alone.
|
||||
special_mode_entries: dict[str, int] = field(default_factory=dict)
|
||||
# Directory entries the producer never writes (`build_zip` stores none), but a gate 5 defense
|
||||
# this repository still carries has to be exercised by something.
|
||||
directory_entries: list[str] = field(default_factory=list)
|
||||
# A file name already in `files`/`extra_files`, written a second time under the same name — two
|
||||
# archive entries claiming the same path, which only gate 5's extraction can catch.
|
||||
repeat_file: str | None = None
|
||||
|
||||
|
||||
def build_zip(spec: ZipSpec | None = None) -> bytes:
|
||||
"""Zip `artifacts/` plus `metadata.json`, mirroring the producer's `build_zip`.
|
||||
|
||||
`metadata.json` is written first with `json.dumps` and no indent, entries use `/` separators,
|
||||
and no directory entries are stored — all three matter to the reader.
|
||||
"""
|
||||
spec = spec or ZipSpec()
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, 'w', zipfile.ZIP_DEFLATED) as archive:
|
||||
if spec.include_metadata:
|
||||
document = spec.document if spec.document is not None else build_metadata_document()
|
||||
archive.writestr('metadata.json', json.dumps(document))
|
||||
for name, content in {**spec.files, **spec.extra_files}.items():
|
||||
archive.writestr(_regular_file_info(name), content)
|
||||
for name, target in spec.symlinks.items():
|
||||
info = zipfile.ZipInfo(name)
|
||||
info.create_system = 3
|
||||
info.external_attr = (0o120777 << 16) | 0o200000
|
||||
archive.writestr(info, target)
|
||||
if spec.absolute_entry:
|
||||
archive.writestr(_regular_file_info(spec.absolute_entry), b'absolute\n')
|
||||
if spec.traversal_entry:
|
||||
archive.writestr(_regular_file_info(spec.traversal_entry), b'traversal\n')
|
||||
if spec.bomb_bytes:
|
||||
archive.writestr(_regular_file_info('artifacts/bomb.csv'), b'\0' * spec.bomb_bytes)
|
||||
for name, mode in spec.special_mode_entries.items():
|
||||
info = zipfile.ZipInfo(name)
|
||||
info.create_system = 3
|
||||
info.external_attr = mode << 16
|
||||
archive.writestr(info, b'')
|
||||
for name in spec.directory_entries:
|
||||
info = zipfile.ZipInfo(name if name.endswith('/') else f'{name}/')
|
||||
info.create_system = 3
|
||||
info.external_attr = (stat.S_IFDIR | 0o755) << 16
|
||||
archive.writestr(info, b'')
|
||||
if spec.repeat_file:
|
||||
content = {**spec.files, **spec.extra_files}[spec.repeat_file]
|
||||
archive.writestr(_regular_file_info(spec.repeat_file), content)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def _regular_file_info(name: str) -> zipfile.ZipInfo:
|
||||
"""A zip entry carrying a Unix regular-file mode, as `ZipFile.write` would store it."""
|
||||
info = zipfile.ZipInfo(name)
|
||||
info.create_system = 3
|
||||
info.compress_type = zipfile.ZIP_DEFLATED
|
||||
info.external_attr = 0o100644 << 16
|
||||
return info
|
||||
|
||||
|
||||
def derive_key(password: str, salt: bytes, opslimit: int, memlimit: int) -> bytes:
|
||||
"""Argon2id over the NFC-normalised password — the producer's `_derive_key`."""
|
||||
normalized = unicodedata.normalize('NFC', password)
|
||||
return sodium.crypto_pwhash_alg(
|
||||
sodium.crypto_secretstream_xchacha20poly1305_KEYBYTES,
|
||||
normalized.encode('utf-8'),
|
||||
salt,
|
||||
opslimit,
|
||||
memlimit,
|
||||
sodium.crypto_pwhash_ALG_ARGON2ID13,
|
||||
)
|
||||
|
||||
|
||||
def encrypt_archive( # noqa: PLR0913 - every knob here exists to build one hostile variant
|
||||
plaintext: bytes,
|
||||
password: str = DEFAULT_PASSWORD,
|
||||
*,
|
||||
format_version: int = FORMAT_VERSION,
|
||||
aead_id: int = AEAD_ID,
|
||||
kdf_id: int = KDF_ID,
|
||||
memlimit: int = TEST_KDF_MEMLIMIT_BYTES,
|
||||
opslimit: int = TEST_KDF_OPSLIMIT,
|
||||
declared_memlimit: int | None = None,
|
||||
declared_opslimit: int | None = None,
|
||||
parallelism: int = 1,
|
||||
reserved_signature: bytes | None = None,
|
||||
magic: bytes = MAGIC,
|
||||
digest_override: bytes | None = None,
|
||||
salt: bytes | None = None,
|
||||
drop_final_tag: bool = False,
|
||||
) -> bytes:
|
||||
"""Produce a full `.sientia` byte stream, mirroring the producer's `encrypt_archive`.
|
||||
|
||||
The `declared_*` arguments write cost parameters into the header that differ from the ones
|
||||
actually used, which is how a "512 MiB header" is built without allocating 512 MiB.
|
||||
`drop_final_tag` pushes every chunk as `MESSAGE`, so the stream is framed as if more were
|
||||
coming — the truncation the digest cannot see.
|
||||
"""
|
||||
salt = salt if salt is not None else os.urandom(SALT_SIZE)
|
||||
key = derive_key(password, salt, opslimit, memlimit)
|
||||
|
||||
state = sodium.crypto_secretstream_xchacha20poly1305_state()
|
||||
stream_header = sodium.crypto_secretstream_xchacha20poly1305_init_push(state, key)
|
||||
|
||||
chunks = []
|
||||
offset = 0
|
||||
total = len(plaintext)
|
||||
while offset < total:
|
||||
chunk = plaintext[offset : offset + CHUNK_SIZE]
|
||||
offset += len(chunk)
|
||||
is_last = offset >= total and not drop_final_tag
|
||||
tag = (
|
||||
sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL
|
||||
if is_last
|
||||
else sodium.crypto_secretstream_xchacha20poly1305_TAG_MESSAGE
|
||||
)
|
||||
chunks.append(sodium.crypto_secretstream_xchacha20poly1305_push(state, chunk, None, tag))
|
||||
ciphertext = b''.join(chunks)
|
||||
|
||||
digest = digest_override if digest_override is not None else hashlib.sha256(ciphertext).digest()
|
||||
header = HEADER_STRUCT.pack(
|
||||
magic,
|
||||
format_version,
|
||||
aead_id,
|
||||
kdf_id,
|
||||
salt,
|
||||
declared_memlimit if declared_memlimit is not None else memlimit,
|
||||
declared_opslimit if declared_opslimit is not None else opslimit,
|
||||
parallelism,
|
||||
digest,
|
||||
reserved_signature if reserved_signature is not None else bytes(RESERVED_SIGNATURE_SIZE),
|
||||
)
|
||||
return header + stream_header + ciphertext
|
||||
|
||||
|
||||
def make_bundle(
|
||||
zip_spec: ZipSpec | None = None,
|
||||
password: str = DEFAULT_PASSWORD,
|
||||
**header_kwargs: Any,
|
||||
) -> bytes:
|
||||
"""The common case: a well-formed tree, encrypted with the test cost parameters."""
|
||||
return encrypt_archive(build_zip(zip_spec), password, **header_kwargs)
|
||||
|
||||
|
||||
def tamper_ciphertext(bundle: bytes, offset_from_stream: int = 8) -> bytes:
|
||||
"""Flip one byte inside the chunk stream **and** repair the header digest.
|
||||
|
||||
Repairing the digest is what makes this a gate-3 test rather than a gate-2 one: the file is
|
||||
self-consistent and only the AEAD can tell it was altered.
|
||||
"""
|
||||
position = HEADER_SIZE + 24 + offset_from_stream
|
||||
altered = bytearray(bundle)
|
||||
altered[position] ^= 0xFF
|
||||
repaired = hashlib.sha256(bytes(altered[HEADER_SIZE + 24 :])).digest()
|
||||
altered[44:76] = repaired
|
||||
return bytes(altered)
|
||||
|
||||
|
||||
def truncate_bundle(bundle: bytes, drop_bytes: int = 32) -> bytes:
|
||||
"""Cut the tail off a bundle, leaving the header's digest describing bytes that are gone."""
|
||||
return bundle[:-drop_bytes]
|
||||
|
||||
|
||||
def build_envelope(password: str = DEFAULT_PASSWORD, key: bytes = DEFAULT_ENVELOPE_KEY) -> str:
|
||||
"""`base64(nonce || ciphertext || tag)` under AES-256-GCM, as the frontend will send it."""
|
||||
nonce = os.urandom(12)
|
||||
sealed = AESGCM(key).encrypt(nonce, password.encode('utf-8'), None)
|
||||
return base64.b64encode(nonce + sealed).decode('ascii')
|
||||
|
||||
|
||||
def envelope_key_b64(key: bytes = DEFAULT_ENVELOPE_KEY) -> str:
|
||||
"""The key as `IMPORT_PASSWORD_KEY` carries it."""
|
||||
return base64.b64encode(key).decode('ascii')
|
||||
174
tests/helpers/status_engine.py
Normal file
174
tests/helpers/status_engine.py
Normal file
@@ -0,0 +1,174 @@
|
||||
"""A recording stand-in for the import log's SQLAlchemy engine.
|
||||
|
||||
The import's status writes are the part of the feature with the strictest contract: which table they
|
||||
name, which predicate they carry, which columns they touch, how many transactions they open and
|
||||
which of them may fail without failing the import. Asserting that against a mock whose `execute`
|
||||
returns a `MagicMock` proves nothing — the assertions need the statement text and the parameters.
|
||||
|
||||
So this double records every statement it is handed, answers reads from a queue the test fills, and
|
||||
lets a test make a write fail on a chosen attempt. It deliberately understands nothing about SQL
|
||||
beyond the first word of the statement.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
class RecordedStatement:
|
||||
"""One statement as the engine saw it: its normalised text and its bound parameters."""
|
||||
|
||||
def __init__(self, statement: Any, parameters: dict[str, Any] | None) -> None:
|
||||
self.sql = ' '.join(str(statement).split())
|
||||
self.parameters: dict[str, Any] = dict(parameters or {})
|
||||
|
||||
@property
|
||||
def verb(self) -> str:
|
||||
return self.sql.split(' ', 1)[0].upper()
|
||||
|
||||
@property
|
||||
def is_write(self) -> bool:
|
||||
return self.verb in ('UPDATE', 'INSERT', 'DELETE')
|
||||
|
||||
@property
|
||||
def table(self) -> str:
|
||||
"""The relation the statement names, for the "one table only" assertions."""
|
||||
words = self.sql.split()
|
||||
if self.verb == 'UPDATE':
|
||||
return words[1]
|
||||
for index, word in enumerate(words):
|
||||
if word.upper() in ('FROM', 'INTO') and index + 1 < len(words):
|
||||
return words[index + 1]
|
||||
return ''
|
||||
|
||||
@property
|
||||
def assignments(self) -> set[str]:
|
||||
"""The column names a `SET` clause assigns, so a test can assert what is never written."""
|
||||
upper = self.sql.upper()
|
||||
if ' SET ' not in upper:
|
||||
return set()
|
||||
body = self.sql[upper.index(' SET ') + 5 :]
|
||||
where = body.upper().find(' WHERE ')
|
||||
if where != -1:
|
||||
body = body[:where]
|
||||
return {part.split('=')[0].strip() for part in body.split(',') if '=' in part}
|
||||
|
||||
def __repr__(self) -> str: # pragma: no cover - debugging aid
|
||||
return f'RecordedStatement({self.sql!r}, {self.parameters!r})'
|
||||
|
||||
|
||||
class FakeResult:
|
||||
"""What `Connection.execute` returns: a rowcount for writes, mappings for reads."""
|
||||
|
||||
def __init__(self, rowcount: int = 1, rows: list[dict[str, Any]] | None = None) -> None:
|
||||
self.rowcount = rowcount
|
||||
self._rows = rows or []
|
||||
|
||||
def mappings(self) -> FakeResult:
|
||||
return self
|
||||
|
||||
def first(self) -> dict[str, Any] | None:
|
||||
return self._rows[0] if self._rows else None
|
||||
|
||||
def all(self) -> list[dict[str, Any]]:
|
||||
return list(self._rows)
|
||||
|
||||
|
||||
class FakeConnection:
|
||||
"""A connection that records and then answers from the engine's queues."""
|
||||
|
||||
def __init__(self, engine: FakeStatusEngine) -> None:
|
||||
self._engine = engine
|
||||
|
||||
def execute(self, statement: Any, parameters: dict[str, Any] | None = None) -> FakeResult:
|
||||
recorded = RecordedStatement(statement, parameters)
|
||||
self._engine.statements.append(recorded)
|
||||
if recorded.is_write:
|
||||
return self._engine._answer_write()
|
||||
return self._engine._answer_read()
|
||||
|
||||
|
||||
class _Scope:
|
||||
"""The context manager `engine.begin()` and `engine.connect()` both hand out."""
|
||||
|
||||
def __init__(self, engine: FakeStatusEngine, transactional: bool) -> None:
|
||||
self._engine = engine
|
||||
self._transactional = transactional
|
||||
|
||||
def __enter__(self) -> FakeConnection:
|
||||
self._engine.open_scopes += 1
|
||||
if self._transactional:
|
||||
self._engine.transactions += 1
|
||||
else:
|
||||
self._engine.connections += 1
|
||||
return FakeConnection(self._engine)
|
||||
|
||||
def __exit__(self, *_exc: object) -> bool:
|
||||
self._engine.open_scopes -= 1
|
||||
return False
|
||||
|
||||
|
||||
class FakeStatusEngine:
|
||||
"""An engine double that records statements and can be told to fail a write.
|
||||
|
||||
Attributes:
|
||||
statements: every statement executed, in order.
|
||||
rowcounts: rowcounts for writes, consumed in order; the last one repeats.
|
||||
rows: rows for reads, consumed in order; the last one repeats. `None` means "no such row".
|
||||
write_errors: exceptions to raise instead of executing a write, consumed in order. A `None`
|
||||
entry means "this attempt succeeds", which is how a retry is modelled.
|
||||
transactions: how many `begin()` scopes were opened — one per write, or the writes are
|
||||
sharing a transaction.
|
||||
open_scopes: scopes currently open, so a test can assert none is left open.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rowcounts: list[int] | None = None,
|
||||
rows: list[dict[str, Any] | None] | None = None,
|
||||
write_errors: list[BaseException | None] | None = None,
|
||||
) -> None:
|
||||
self.statements: list[RecordedStatement] = []
|
||||
self.rowcounts: list[int] = list(rowcounts or [])
|
||||
self.rows: list[dict[str, Any] | None] = list(rows or [])
|
||||
self.write_errors: list[BaseException | None] = list(write_errors or [])
|
||||
self.transactions = 0
|
||||
self.connections = 0
|
||||
self.open_scopes = 0
|
||||
self.disposed = False
|
||||
|
||||
# ------------------------------------------------------------------ engine API
|
||||
def begin(self) -> _Scope:
|
||||
return _Scope(self, transactional=True)
|
||||
|
||||
def connect(self) -> _Scope:
|
||||
return _Scope(self, transactional=False)
|
||||
|
||||
def dispose(self) -> None:
|
||||
self.disposed = True
|
||||
|
||||
# ------------------------------------------------------------------- answering
|
||||
def _answer_write(self) -> FakeResult:
|
||||
if self.write_errors:
|
||||
error = self.write_errors.pop(0)
|
||||
if error is not None:
|
||||
raise error
|
||||
rowcount = self.rowcounts.pop(0) if len(self.rowcounts) > 1 else (self.rowcounts or [1])[0]
|
||||
return FakeResult(rowcount=rowcount)
|
||||
|
||||
def _answer_read(self) -> FakeResult:
|
||||
row = self.rows.pop(0) if len(self.rows) > 1 else (self.rows or [None])[0]
|
||||
return FakeResult(rows=[row] if row is not None else [])
|
||||
|
||||
# ---------------------------------------------------------------- convenience
|
||||
@property
|
||||
def writes(self) -> list[RecordedStatement]:
|
||||
return [statement for statement in self.statements if statement.is_write]
|
||||
|
||||
@property
|
||||
def reads(self) -> list[RecordedStatement]:
|
||||
return [statement for statement in self.statements if not statement.is_write]
|
||||
|
||||
@property
|
||||
def tables(self) -> set[str]:
|
||||
return {statement.table for statement in self.statements if statement.table}
|
||||
@@ -17,11 +17,13 @@ from laborious.activities.storage import Storage
|
||||
@patch('laborious.activities.activities.Gates.__init__')
|
||||
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
||||
@patch('laborious.activities.activities.API.__init__')
|
||||
@patch('laborious.activities.activities.MLFlowRepository')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@patch('laborious.activities.activities.MetricsController')
|
||||
def test___init__(
|
||||
mock_metrics_controller,
|
||||
mock_minio_repository,
|
||||
mock_mlflow_repository,
|
||||
mock_api_init,
|
||||
mock_model_metrics_init,
|
||||
mock_gates_init,
|
||||
@@ -101,16 +103,22 @@ def test___init__(
|
||||
|
||||
mock_mlflow_init.assert_called_once_with(
|
||||
ANY,
|
||||
mlflow_host=mlflow_config['host'],
|
||||
mlflow_port=mlflow_config['port'],
|
||||
mlflow_username=mlflow_config['username'],
|
||||
mlflow_password=mlflow_config['password'],
|
||||
mlflow_repository=mock_mlflow_repository.return_value,
|
||||
minio_repository=mock_minio_repository.return_value,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mock_metrics_controller.return_value,
|
||||
)
|
||||
|
||||
mock_mlflow_repository.assert_called_once_with(
|
||||
f'{mlflow_config["host"]}:{mlflow_config["port"]}',
|
||||
mlflow_config['username'],
|
||||
mlflow_config['password'],
|
||||
logger,
|
||||
notification_handler,
|
||||
mock_metrics_controller.return_value,
|
||||
)
|
||||
|
||||
mock_opc_init.assert_called_once_with(
|
||||
ANY,
|
||||
opc_servers=opc_config,
|
||||
@@ -156,6 +164,161 @@ def test___init__(
|
||||
)
|
||||
|
||||
|
||||
BASE_CONFIGS = {
|
||||
'postgres_config': {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'postgres',
|
||||
'password': 'postgres',
|
||||
'dbname': 'postgres',
|
||||
'min_connections': 1,
|
||||
'max_connections': 10,
|
||||
},
|
||||
'minio_config': {
|
||||
'endpoint_url': 'localhost:9000',
|
||||
'access_key': 'minio',
|
||||
'secret_key': 'minio123',
|
||||
'default_bucket': 'test',
|
||||
'retention_hours': 24,
|
||||
'secure': False,
|
||||
},
|
||||
'mlflow_config': {'host': 'localhost', 'port': 5000, 'username': 'm', 'password': 'm'},
|
||||
'opc_config': {},
|
||||
'pi_web_api_config': {'base_url': 'https://pi', 'auth_type': 'bearer', 'auth_token': 't'},
|
||||
}
|
||||
|
||||
IMPORT_CONFIG = {'bucket': 'sientia', 'prefix': 'imported_models/'}
|
||||
IMPORT_STATUS_CONFIG = {
|
||||
'host': 'bff-db',
|
||||
'port': 5432,
|
||||
'dbname': 'sientia-core-mlops-bff',
|
||||
'user': 'u',
|
||||
'password': 'p',
|
||||
'schema': 'public',
|
||||
}
|
||||
|
||||
|
||||
@patch('laborious.activities.activities.ModelImport.__init__')
|
||||
@patch('laborious.activities.activities.Storage.__init__')
|
||||
@patch('laborious.activities.activities.MLFlow.__init__')
|
||||
@patch('laborious.activities.activities.OPC.__init__')
|
||||
@patch('laborious.activities.activities.Gates.__init__')
|
||||
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
||||
@patch('laborious.activities.activities.API.__init__')
|
||||
@patch('laborious.activities.activities.MongoDBRepository')
|
||||
@patch('laborious.activities.activities.MLFlowRepository')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@patch('laborious.activities.activities.MetricsController')
|
||||
def test_mlflow_and_import_activities_receive_the_same_mlflow_repository_instance( # noqa: PLR0913
|
||||
_mock_metrics_controller,
|
||||
mock_minio_repository,
|
||||
mock_mlflow_repository,
|
||||
mock_mongo_repository,
|
||||
_mock_api_init,
|
||||
_mock_model_metrics_init,
|
||||
_mock_gates_init,
|
||||
_mock_opc_init,
|
||||
mock_mlflow_init,
|
||||
_mock_storage_init,
|
||||
mock_model_import_init,
|
||||
):
|
||||
"""Setting the MLflow tracking URI is a process-wide side effect, so two repository instances
|
||||
would be two chances to disagree about where MLflow is."""
|
||||
Activities(
|
||||
**BASE_CONFIGS,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
mongo_config={'connection_string': 'mongodb://mongo', 'database_name': 'sientia'},
|
||||
import_config=IMPORT_CONFIG,
|
||||
import_status_config=IMPORT_STATUS_CONFIG,
|
||||
)
|
||||
|
||||
repository = mock_mlflow_repository.return_value
|
||||
assert mock_mlflow_init.call_args.kwargs['mlflow_repository'] is repository
|
||||
|
||||
passed = mock_model_import_init.call_args.kwargs
|
||||
assert passed['mlflow_repository'] is repository
|
||||
assert passed['minio_repository'] is mock_minio_repository.return_value
|
||||
assert passed['mongo_repository'] is mock_mongo_repository.return_value
|
||||
assert passed['import_config'] is IMPORT_CONFIG
|
||||
assert passed['status_config'] is IMPORT_STATUS_CONFIG
|
||||
|
||||
|
||||
@patch('laborious.activities.activities.ModelImport.__init__')
|
||||
@patch('laborious.activities.activities.Storage.__init__')
|
||||
@patch('laborious.activities.activities.MLFlow.__init__')
|
||||
@patch('laborious.activities.activities.OPC.__init__')
|
||||
@patch('laborious.activities.activities.Gates.__init__')
|
||||
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
||||
@patch('laborious.activities.activities.API.__init__')
|
||||
@patch('laborious.activities.activities.MongoDBRepository')
|
||||
@patch('laborious.activities.activities.MLFlowRepository')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@patch('laborious.activities.activities.MetricsController')
|
||||
def test_without_a_mongo_config_the_import_gets_no_repository( # noqa: PLR0913
|
||||
_mock_metrics_controller,
|
||||
_mock_minio_repository,
|
||||
_mock_mlflow_repository,
|
||||
mock_mongo_repository,
|
||||
_mock_api_init,
|
||||
_mock_model_metrics_init,
|
||||
_mock_gates_init,
|
||||
_mock_opc_init,
|
||||
_mock_mlflow_init,
|
||||
_mock_storage_init,
|
||||
mock_model_import_init,
|
||||
):
|
||||
"""It refuses rather than guessing: the notification handler keeps its own client either way."""
|
||||
Activities(
|
||||
**BASE_CONFIGS,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
import_config=IMPORT_CONFIG,
|
||||
import_status_config=IMPORT_STATUS_CONFIG,
|
||||
)
|
||||
|
||||
assert mock_model_import_init.call_args.kwargs['mongo_repository'] is None
|
||||
mock_mongo_repository.assert_not_called()
|
||||
|
||||
|
||||
@patch('laborious.activities.activities.ModelImport.__init__')
|
||||
@patch('laborious.activities.activities.Storage.__init__')
|
||||
@patch('laborious.activities.activities.MLFlow.__init__')
|
||||
@patch('laborious.activities.activities.OPC.__init__')
|
||||
@patch('laborious.activities.activities.Gates.__init__')
|
||||
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
||||
@patch('laborious.activities.activities.API.__init__')
|
||||
@patch('laborious.activities.activities.MongoDBRepository')
|
||||
@patch('laborious.activities.activities.MLFlowRepository')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@patch('laborious.activities.activities.MetricsController')
|
||||
def test_exactly_one_mlflow_repository_is_built_even_with_import_active( # noqa: PLR0913
|
||||
_mock_metrics_controller,
|
||||
_mock_minio_repository,
|
||||
mock_mlflow_repository,
|
||||
_mock_mongo_repository,
|
||||
_mock_api_init,
|
||||
_mock_model_metrics_init,
|
||||
_mock_gates_init,
|
||||
_mock_opc_init,
|
||||
_mock_mlflow_init,
|
||||
_mock_storage_init,
|
||||
_mock_model_import_init,
|
||||
):
|
||||
"""No second `MLFlowRepository` is ever created on any path: `MLFlow` and `ModelImport` only
|
||||
receive the one instance built here, they never build their own."""
|
||||
Activities(
|
||||
**BASE_CONFIGS,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
mongo_config={'connection_string': 'mongodb://mongo', 'database_name': 'sientia'},
|
||||
import_config=IMPORT_CONFIG,
|
||||
import_status_config=IMPORT_STATUS_CONFIG,
|
||||
)
|
||||
|
||||
mock_mlflow_repository.assert_called_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.activities.Storage')
|
||||
@patch('laborious.activities.activities.MLFlow')
|
||||
@@ -163,9 +326,11 @@ def test___init__(
|
||||
@patch('laborious.activities.activities.Gates')
|
||||
@patch('laborious.activities.activities.ModelMetrics')
|
||||
@patch('laborious.activities.activities.API')
|
||||
@patch('laborious.activities.activities.MLFlowRepository')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
async def test_shutdown(
|
||||
_mock_minio_repository,
|
||||
_mock_mlflow_repository,
|
||||
mock_api_init,
|
||||
mock_model_metrics_init,
|
||||
mock_gates_init,
|
||||
@@ -173,7 +338,7 @@ async def test_shutdown(
|
||||
mock_mlflow_init,
|
||||
mock_storage_init,
|
||||
):
|
||||
mock_opc_init.close = AsyncMock()
|
||||
mock_opc_init.aclose = AsyncMock()
|
||||
postgres_config = {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
@@ -221,7 +386,7 @@ async def test_shutdown(
|
||||
)
|
||||
|
||||
await activities.shutdown()
|
||||
mock_opc_init.close.assert_called_once()
|
||||
mock_opc_init.aclose.assert_called_once()
|
||||
mock_storage_init.close.assert_called_once()
|
||||
mock_mlflow_init.close.assert_called_once()
|
||||
mock_gates_init.close.assert_called_once()
|
||||
|
||||
@@ -220,6 +220,44 @@ async def test_write_pi_web_api_data_success(mock_dataframe, api, base_input_dat
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.api.DataFrame')
|
||||
async def test_write_pi_web_api_data_partial_prediction_mismatch_downgrades_confidence(
|
||||
mock_dataframe, api, base_input_data
|
||||
):
|
||||
"""A partial write (some tags report Errors) doesn't raise — it's caught by the tag-count
|
||||
mismatch check, which still must downgrade confidence and attach the mismatch message."""
|
||||
input_data = {
|
||||
**base_input_data,
|
||||
'pi_web_api_output_config': {
|
||||
'endpoint': 'https://test-pi-server.com/piwebapi',
|
||||
'prediction_tags': {'tag1': 'web_id_1', 'tag2': 'web_id_2'},
|
||||
'confidence_tags': {'tag3': 'web_id_3'},
|
||||
},
|
||||
}
|
||||
mock_df = _create_mock_dataframe()
|
||||
mock_dataframe.return_value = mock_df
|
||||
|
||||
api.pi_web_api_client.write_value.side_effect = [
|
||||
[{'WebId': 'web_id_1', 'Errors': []}, {'WebId': 'web_id_2', 'Errors': ['bad tag']}],
|
||||
[{'WebId': 'web_id_3', 'Errors': []}],
|
||||
]
|
||||
|
||||
await api.write_pi_web_api_data(input_data)
|
||||
|
||||
mock_df.__setitem__.assert_any_call(
|
||||
'prediction_confidence', PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||
)
|
||||
mock_df.__setitem__.assert_any_call('comments', ANY)
|
||||
api.send_notification_async.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
|
||||
message=ANY,
|
||||
block='write_pi_web_api_data',
|
||||
level=NotificationLevel.ERROR,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.api.DataFrame')
|
||||
async def test_write_pi_web_api_data_prediction_error(mock_dataframe, api, base_input_data):
|
||||
|
||||
@@ -16,12 +16,12 @@ def _passthrough_from_dict():
|
||||
yield
|
||||
|
||||
|
||||
@patch('laborious.activities.mlflow.MLFlowRepository')
|
||||
@patch('laborious.activities.mlflow.MinioRepository')
|
||||
def test___init__(mock_minio_repository, mock_mlflow_repository):
|
||||
def test___init__(mock_minio_repository):
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
metrics_controller = AsyncMock()
|
||||
mock_mlflow_repository = MagicMock()
|
||||
minio_repo = mock_minio_repository(
|
||||
endpoint='localhost:9000',
|
||||
access_key='minio',
|
||||
@@ -32,24 +32,14 @@ def test___init__(mock_minio_repository, mock_mlflow_repository):
|
||||
bucket='test',
|
||||
)
|
||||
mlflow = MLFlow(
|
||||
mlflow_host='http://localhost',
|
||||
mlflow_port=5000,
|
||||
mlflow_username='admin',
|
||||
mlflow_password='admin',
|
||||
mlflow_repository=mock_mlflow_repository,
|
||||
minio_repository=minio_repo,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
assert mlflow.mlflow_host == 'http://localhost'
|
||||
assert mlflow.mlflow_port == 5000
|
||||
assert mlflow.mlflow_username == 'admin'
|
||||
assert mlflow.mlflow_password == 'admin'
|
||||
|
||||
mock_mlflow_repository.assert_called_once_with(
|
||||
'http://localhost:5000', 'admin', 'admin', ANY, ANY, ANY
|
||||
)
|
||||
assert mlflow.model_monitoring_repository is mock_mlflow_repository
|
||||
|
||||
mock_minio_repository.assert_called_once_with(
|
||||
endpoint='localhost:9000',
|
||||
@@ -62,10 +52,18 @@ def test___init__(mock_minio_repository, mock_mlflow_repository):
|
||||
)
|
||||
|
||||
|
||||
def test___init___requires_mlflow_repository():
|
||||
"""The failure moves to the signature, at construction time, instead of a `None` that only
|
||||
surfaces silently at first use."""
|
||||
minio_repo = MagicMock()
|
||||
logger = MagicMock()
|
||||
with raises(TypeError):
|
||||
MLFlow(minio_repository=minio_repo, logger=logger)
|
||||
|
||||
|
||||
@fixture
|
||||
@patch('laborious.activities.mlflow.MLFlowRepository')
|
||||
@patch('laborious.activities.mlflow.MinioRepository')
|
||||
def mlflow(mock_minio_repository, mock_mlflow_repository):
|
||||
def mlflow(mock_minio_repository):
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
metrics_controller = AsyncMock()
|
||||
@@ -79,10 +77,7 @@ def mlflow(mock_minio_repository, mock_mlflow_repository):
|
||||
bucket='test',
|
||||
)
|
||||
mlflow = MLFlow(
|
||||
mlflow_host='http://localhost:5000',
|
||||
mlflow_port=5000,
|
||||
mlflow_username='admin',
|
||||
mlflow_password='admin',
|
||||
mlflow_repository=MagicMock(),
|
||||
minio_repository=minio_repo,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
@@ -222,6 +217,39 @@ async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow):
|
||||
assert response_data == mock_from_dataframe.return_value
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch(
|
||||
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
@patch('laborious.activities.mlflow.to_datetime')
|
||||
async def test_request_predict_skips_timestamp_reinsertion_when_configured(
|
||||
mock_to_datetime, mock_from_dataframe, mlflow
|
||||
):
|
||||
"""skip_predict_timestamp=True leaves the caller's index alone — no reinsertion round-trip."""
|
||||
data_mock = MagicMock()
|
||||
payload = AsyncMock()
|
||||
payload.retrieve = AsyncMock(return_value=data_mock)
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'data': payload,
|
||||
'model_name': 'test_model',
|
||||
'model_config': {'skip_predict_timestamp': True},
|
||||
}
|
||||
|
||||
predict_response = {'success': True, 'content': MagicMock()}
|
||||
mlflow.model_monitoring_repository.predict.return_value = predict_response
|
||||
|
||||
await mlflow.request_predict(input_data)
|
||||
|
||||
data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
|
||||
mock_to_datetime.assert_not_called()
|
||||
mlflow.model_monitoring_repository.predict.assert_called_once_with(
|
||||
'test_model', data_mock, {'skip_predict_timestamp': True}, metadata['metadata']
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch(
|
||||
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
||||
@@ -491,7 +519,7 @@ async def test_retrain_model_data_error_no_minio_repository(mlflow):
|
||||
}
|
||||
)
|
||||
|
||||
assert str(e.value) == 'Minio repository not initialized'
|
||||
assert str(e.value) == 'Minio repository not initialized'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@@ -613,12 +641,12 @@ async def test_get_reference_data_exception(mlflow):
|
||||
'model_name': 'test_model',
|
||||
}
|
||||
|
||||
mlflow.model_monitoring_repository.load_artifact_dataframe.side_effect = Exception(
|
||||
mlflow.model_monitoring_repository.load_artifact_dataframe.side_effect = RuntimeError(
|
||||
'Error loading artifact'
|
||||
)
|
||||
|
||||
# Act & Assert
|
||||
with raises(Exception) as e:
|
||||
with raises(RuntimeError) as e:
|
||||
await mlflow.get_reference_data(input_data)
|
||||
|
||||
assert str(e.value) == 'Error loading artifact'
|
||||
|
||||
781
tests/laborious/activities/test_model_import_bundle.py
Normal file
781
tests/laborious/activities/test_model_import_bundle.py
Normal file
@@ -0,0 +1,781 @@
|
||||
"""Downloading the object and opening the bundle: the two activities that touch the secret.
|
||||
|
||||
Three things are asserted here:
|
||||
|
||||
- **the digest is enforced against the bytes that arrived**, before the header is parsed, with no
|
||||
path that skips the comparison — a mismatch means the object behind that key is not the file the
|
||||
user chose, and no amount of well-formedness makes it the right file;
|
||||
- **the envelope is opened inside the activity that uses the password**, so a plaintext password is
|
||||
never an activity input, an activity result, or a value the workflow holds;
|
||||
- **nothing observable carries the secret** — for every failure path, the password, the envelope key
|
||||
and the derived key are absent from the log, the notification, the metric labels, the exception and
|
||||
the return value.
|
||||
|
||||
The bundles here are real: the shared factory writes them field for field like the producer, at the
|
||||
test cost parameters, and the production reader opens them with real Argon2id and real
|
||||
`secretstream`.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from pytest import fixture, mark, raises
|
||||
from temporalio.exceptions import ApplicationError
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.activities.model_import import ModelImport, derive_run_name
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from laborious.utils.models.import_errors import ImportGateRejectedError, ImportMisconfiguredError
|
||||
from tests.helpers import bundle_factory as factory
|
||||
from tests.helpers.status_engine import FakeStatusEngine
|
||||
|
||||
# Distinctive enough that a substring search for it in a log line cannot match by accident.
|
||||
BUNDLE_PASSWORD = 'Zk9-bundle-secret-Passw0rd'
|
||||
ENVELOPE_KEY = bytes(range(100, 132))
|
||||
OTHER_ENVELOPE_KEY = bytes(range(200, 232))
|
||||
|
||||
METADATA = {
|
||||
'model_id': '77',
|
||||
'model_name': 'factory_model',
|
||||
'workflow_name': 'import_model',
|
||||
'schedule_name': '-',
|
||||
}
|
||||
WORKFLOW_ID = 'import_model-77-abc'
|
||||
OBJECT_KEY = 'imported_models/factory_model_v2.sientia'
|
||||
|
||||
|
||||
def import_config(work_dir: Path, **overrides) -> dict:
|
||||
return {
|
||||
'password_key': factory.envelope_key_b64(ENVELOPE_KEY),
|
||||
'bucket': 'sientia',
|
||||
'prefix': 'imported_models/',
|
||||
'work_dir': str(work_dir),
|
||||
'max_object_bytes': 8 * 1024 * 1024,
|
||||
'max_entries': 5000,
|
||||
'max_uncompressed_bytes': 4 * 1024 * 1024,
|
||||
'max_compression_ratio': 200.0,
|
||||
'models_collection': 'models',
|
||||
'retention_days': 7,
|
||||
**overrides,
|
||||
}
|
||||
|
||||
|
||||
@fixture
|
||||
def minio_client() -> MagicMock:
|
||||
client = MagicMock()
|
||||
client.stat_object.return_value = MagicMock(size=1024)
|
||||
return client
|
||||
|
||||
|
||||
@fixture
|
||||
def import_activity(tmp_path: Path, minio_client: MagicMock) -> ModelImport:
|
||||
repository = MagicMock()
|
||||
repository.minio_client = minio_client
|
||||
activity = ModelImport(
|
||||
status_config={
|
||||
'host': 'h',
|
||||
'port': 5432,
|
||||
'dbname': 'sientia-core-mlops-bff',
|
||||
'user': 'u',
|
||||
'password': 'p',
|
||||
'schema': 'public',
|
||||
},
|
||||
import_config=import_config(tmp_path),
|
||||
minio_repository=repository,
|
||||
mlflow_repository=MagicMock(),
|
||||
mongo_repository=AsyncMock(),
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
# `status_engine` is created lazily on first access; pre-seed it with the recording double.
|
||||
activity._status_engine = FakeStatusEngine()
|
||||
activity.emit_metric = AsyncMock()
|
||||
activity.send_notification_async = AsyncMock()
|
||||
return activity
|
||||
|
||||
|
||||
def serve_object(minio_client: MagicMock, payload: bytes) -> None:
|
||||
"""Make `fget_object` write `payload` where the activity asked for it."""
|
||||
|
||||
def write(_bucket, _key, destination):
|
||||
Path(destination).write_bytes(payload)
|
||||
|
||||
minio_client.stat_object.return_value = MagicMock(size=len(payload))
|
||||
minio_client.fget_object.side_effect = write
|
||||
|
||||
|
||||
def download_input(digest: str, **overrides) -> dict:
|
||||
return {
|
||||
'metadata': METADATA,
|
||||
'import_run_id': 77,
|
||||
'workflow_id': WORKFLOW_ID,
|
||||
'object_key': OBJECT_KEY,
|
||||
'expected_digest': digest,
|
||||
**overrides,
|
||||
}
|
||||
|
||||
|
||||
def open_input(path: str, envelope: str | None = None, **overrides) -> dict:
|
||||
return {
|
||||
'metadata': METADATA,
|
||||
'import_run_id': 77,
|
||||
'workflow_id': WORKFLOW_ID,
|
||||
'path': path,
|
||||
'password_envelope': envelope
|
||||
if envelope is not None
|
||||
else factory.build_envelope(BUNDLE_PASSWORD, ENVELOPE_KEY),
|
||||
**overrides,
|
||||
}
|
||||
|
||||
|
||||
def observable_text(activity: ModelImport) -> str:
|
||||
"""Everything this activity said out loud: logs, notifications and metric labels."""
|
||||
logger = activity.logger
|
||||
parts = [
|
||||
str(call) for call in logger.custom_info.call_args_list + logger.custom_error.call_args_list
|
||||
]
|
||||
parts += [str(call) for call in activity.send_notification_async.await_args_list]
|
||||
parts += [str(call) for call in activity.emit_metric.await_args_list]
|
||||
return '\n'.join(parts)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------- the download
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_matching_digest_yields_the_downloaded_file(
|
||||
import_activity: ModelImport, minio_client: MagicMock, tmp_path: Path
|
||||
):
|
||||
payload = factory.make_bundle(password=BUNDLE_PASSWORD)
|
||||
serve_object(minio_client, payload)
|
||||
digest = hashlib.sha256(payload).hexdigest()
|
||||
|
||||
result = await import_activity._download_bundle(download_input(digest))
|
||||
|
||||
assert Path(result['path']).read_bytes() == payload
|
||||
assert result['digest'] == digest
|
||||
assert result['size'] == len(payload)
|
||||
assert Path(result['path']).parent == tmp_path / WORKFLOW_ID
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_configured_bucket_is_used_when_the_request_names_none(
|
||||
import_activity: ModelImport, minio_client: MagicMock
|
||||
):
|
||||
payload = factory.make_bundle(password=BUNDLE_PASSWORD)
|
||||
serve_object(minio_client, payload)
|
||||
|
||||
await import_activity._download_bundle(download_input(hashlib.sha256(payload).hexdigest()))
|
||||
|
||||
assert minio_client.stat_object.call_args.args[0] == 'sientia'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_request_may_name_another_bucket(
|
||||
import_activity: ModelImport, minio_client: MagicMock
|
||||
):
|
||||
payload = factory.make_bundle(password=BUNDLE_PASSWORD)
|
||||
serve_object(minio_client, payload)
|
||||
|
||||
await import_activity._download_bundle(
|
||||
download_input(hashlib.sha256(payload).hexdigest(), bucket='other-bucket')
|
||||
)
|
||||
|
||||
assert minio_client.stat_object.call_args.args[0] == 'other-bucket'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_digest_from_another_file_fails_before_the_header_is_parsed(
|
||||
import_activity: ModelImport, minio_client: MagicMock
|
||||
):
|
||||
payload = factory.make_bundle(password=BUNDLE_PASSWORD)
|
||||
other = factory.make_bundle(password=BUNDLE_PASSWORD, salt=b'\x01' * 16)
|
||||
serve_object(minio_client, payload)
|
||||
|
||||
request = download_input(hashlib.sha256(other).hexdigest())
|
||||
|
||||
with raises(ImportGateRejectedError) as failure:
|
||||
await import_activity._download_bundle(request)
|
||||
|
||||
assert failure.value.details[0]['step'] == ImportStep.DOWNLOAD.value
|
||||
assert failure.value.details[0]['gate'] == 1
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_file_that_does_not_match_is_not_left_on_the_worker(
|
||||
import_activity: ModelImport, minio_client: MagicMock, tmp_path: Path
|
||||
):
|
||||
serve_object(minio_client, factory.make_bundle(password=BUNDLE_PASSWORD))
|
||||
|
||||
request = download_input('0' * 64)
|
||||
|
||||
with raises(ImportGateRejectedError):
|
||||
await import_activity._download_bundle(request)
|
||||
|
||||
assert not (tmp_path / WORKFLOW_ID / 'bundle.sientia').exists()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_digest_comparison_is_case_insensitive_on_the_input(
|
||||
import_activity: ModelImport, minio_client: MagicMock
|
||||
):
|
||||
payload = factory.make_bundle(password=BUNDLE_PASSWORD)
|
||||
serve_object(minio_client, payload)
|
||||
|
||||
result = await import_activity._download_bundle(
|
||||
download_input(hashlib.sha256(payload).hexdigest().upper())
|
||||
)
|
||||
|
||||
assert result['digest'] == hashlib.sha256(payload).hexdigest()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_an_object_over_the_ceiling_is_refused_before_its_body_is_fetched(
|
||||
import_activity: ModelImport, minio_client: MagicMock
|
||||
):
|
||||
minio_client.stat_object.return_value = MagicMock(size=9 * 1024 * 1024)
|
||||
request = download_input('a' * 64)
|
||||
|
||||
with raises(ImportGateRejectedError) as failure:
|
||||
await import_activity._download_bundle(request)
|
||||
|
||||
assert failure.value.details[0]['step'] == ImportStep.DOWNLOAD.value
|
||||
minio_client.fget_object.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@mark.parametrize(
|
||||
'object_key',
|
||||
['other_prefix/model.sientia', 'imported_models/model.zip'],
|
||||
)
|
||||
async def test_a_key_of_the_wrong_shape_is_refused(
|
||||
import_activity: ModelImport, minio_client: MagicMock, object_key: str
|
||||
):
|
||||
request = download_input('a' * 64, object_key=object_key)
|
||||
|
||||
with raises(ImportGateRejectedError):
|
||||
await import_activity._download_bundle(request)
|
||||
|
||||
minio_client.fget_object.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_an_unreadable_object_is_a_download_failure(
|
||||
import_activity: ModelImport, minio_client: MagicMock
|
||||
):
|
||||
minio_client.stat_object.side_effect = OSError('no such object')
|
||||
request = download_input('a' * 64)
|
||||
|
||||
with raises(ImportGateRejectedError) as failure:
|
||||
await import_activity._download_bundle(request)
|
||||
|
||||
assert failure.value.details[0]['step'] == ImportStep.DOWNLOAD.value
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_failed_transfer_is_a_download_failure(
|
||||
import_activity: ModelImport, minio_client: MagicMock
|
||||
):
|
||||
minio_client.fget_object.side_effect = OSError('connection reset')
|
||||
request = download_input('a' * 64)
|
||||
|
||||
with raises(ImportGateRejectedError) as failure:
|
||||
await import_activity._download_bundle(request)
|
||||
|
||||
assert failure.value.details[0]['step'] == ImportStep.DOWNLOAD.value
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_missing_minio_repository_is_a_configuration_error(tmp_path: Path):
|
||||
with patch('laborious.activities.model_import.create_engine', return_value=FakeStatusEngine()):
|
||||
activity = ModelImport(
|
||||
status_config={
|
||||
'host': 'h',
|
||||
'port': 5432,
|
||||
'dbname': 'd',
|
||||
'user': 'u',
|
||||
'password': 'p',
|
||||
'schema': 'public',
|
||||
},
|
||||
import_config=import_config(tmp_path),
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
|
||||
request = download_input('a' * 64)
|
||||
|
||||
with raises(ImportMisconfiguredError, match='Minio repository not initialized') as failure:
|
||||
await activity._download_bundle(request)
|
||||
|
||||
assert failure.value.non_retryable is True
|
||||
assert failure.value.type == 'ImportMisconfiguredError'
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------- opening the bundle
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_well_formed_bundle_opens_and_returns_nothing_secret(
|
||||
import_activity: ModelImport, tmp_path: Path
|
||||
):
|
||||
path = tmp_path / 'bundle.sientia'
|
||||
path.write_bytes(factory.make_bundle(password=BUNDLE_PASSWORD))
|
||||
|
||||
result = await import_activity._open_bundle(open_input(str(path)))
|
||||
|
||||
assert result['metadata']['model_name'] == 'factory_model'
|
||||
assert result['parameters']['target_variable'] == 'SE'
|
||||
assert Path(result['extracted_dir']).is_dir()
|
||||
assert (Path(result['extracted_dir']) / 'artifacts' / 'prediction_model' / 'MLmodel').is_file()
|
||||
assert set(result) == {'extracted_dir', 'metadata', 'parameters', 'digest'}
|
||||
assert BUNDLE_PASSWORD not in str(result)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_opening_a_bundle_measures_the_kdf(import_activity: ModelImport, tmp_path: Path):
|
||||
path = tmp_path / 'bundle.sientia'
|
||||
path.write_bytes(factory.make_bundle(password=BUNDLE_PASSWORD))
|
||||
|
||||
await import_activity._open_bundle(open_input(str(path)))
|
||||
|
||||
observed = [
|
||||
call.kwargs
|
||||
for call in import_activity.emit_metric.await_args_list
|
||||
if call.kwargs.get('metric_object') is metrics.MODEL_IMPORT_KDF_DURATION
|
||||
]
|
||||
assert len(observed) == 1
|
||||
assert observed[0]['method'] == 'observe'
|
||||
assert observed[0]['value'] > 0
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_wrong_envelope_key_is_not_confused_with_a_wrong_bundle_password(
|
||||
import_activity: ModelImport, tmp_path: Path
|
||||
):
|
||||
path = tmp_path / 'bundle.sientia'
|
||||
path.write_bytes(factory.make_bundle(password=BUNDLE_PASSWORD))
|
||||
envelope = factory.build_envelope(BUNDLE_PASSWORD, OTHER_ENVELOPE_KEY)
|
||||
request = open_input(str(path), envelope)
|
||||
|
||||
with raises(ImportGateRejectedError) as failure:
|
||||
await import_activity._open_bundle(request)
|
||||
|
||||
details = failure.value.details[0]
|
||||
assert details['step'] == ImportStep.DECRYPTION.value
|
||||
assert details['gate'] == 3
|
||||
# The platform's own transport failed, not the user's password. Same step — the step is where the
|
||||
# import stopped, not who is at fault — but a distinguishable cause.
|
||||
assert details['cause'] == 'envelope'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_wrong_bundle_password_is_a_decryption_failure_of_its_own(
|
||||
import_activity: ModelImport, tmp_path: Path
|
||||
):
|
||||
path = tmp_path / 'bundle.sientia'
|
||||
path.write_bytes(factory.make_bundle(password='a-different-password'))
|
||||
request = open_input(str(path))
|
||||
|
||||
with raises(ImportGateRejectedError) as failure:
|
||||
await import_activity._open_bundle(request)
|
||||
|
||||
details = failure.value.details[0]
|
||||
assert details['step'] == ImportStep.DECRYPTION.value
|
||||
assert details['cause'] != 'envelope'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_missing_password_key_stops_the_import_at_decryption(
|
||||
import_activity: ModelImport, tmp_path: Path
|
||||
):
|
||||
import_activity.import_config['password_key'] = None
|
||||
path = tmp_path / 'bundle.sientia'
|
||||
path.write_bytes(factory.make_bundle(password=BUNDLE_PASSWORD))
|
||||
request = open_input(str(path))
|
||||
|
||||
with raises(ImportGateRejectedError) as failure:
|
||||
await import_activity._open_bundle(request)
|
||||
|
||||
assert failure.value.details[0]['cause'] == 'envelope'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@mark.parametrize(
|
||||
'envelope',
|
||||
['', 'not-base64!!', 'AAAA', 'bm90LWFuLWVudmVsb3BlLWF0LWFsbA=='],
|
||||
ids=['empty', 'not-base64', 'too-short', 'wrong-bytes'],
|
||||
)
|
||||
async def test_a_tampered_or_malformed_envelope_is_refused(
|
||||
import_activity: ModelImport, tmp_path: Path, envelope: str
|
||||
):
|
||||
path = tmp_path / 'bundle.sientia'
|
||||
path.write_bytes(factory.make_bundle(password=BUNDLE_PASSWORD))
|
||||
request = open_input(str(path), envelope)
|
||||
|
||||
with raises(ImportGateRejectedError) as failure:
|
||||
await import_activity._open_bundle(request)
|
||||
|
||||
assert failure.value.details[0]['cause'] == 'envelope'
|
||||
|
||||
|
||||
def bad_magic() -> bytes:
|
||||
return factory.make_bundle(password=BUNDLE_PASSWORD, magic=b'NOTSIEN1')
|
||||
|
||||
|
||||
def wrong_digest() -> bytes:
|
||||
return factory.make_bundle(password=BUNDLE_PASSWORD, digest_override=bytes(32))
|
||||
|
||||
|
||||
def tampered() -> bytes:
|
||||
return factory.tamper_ciphertext(factory.make_bundle(password=BUNDLE_PASSWORD))
|
||||
|
||||
|
||||
def traversal() -> bytes:
|
||||
return factory.make_bundle(
|
||||
factory.ZipSpec(traversal_entry='artifacts/../../escaped.csv'), password=BUNDLE_PASSWORD
|
||||
)
|
||||
|
||||
|
||||
def absolute() -> bytes:
|
||||
return factory.make_bundle(
|
||||
factory.ZipSpec(absolute_entry='/etc/passwd'), password=BUNDLE_PASSWORD
|
||||
)
|
||||
|
||||
|
||||
def missing_data_model() -> bytes:
|
||||
files = {
|
||||
name: content
|
||||
for name, content in factory.default_tree().items()
|
||||
if not name.startswith('artifacts/data_model/')
|
||||
}
|
||||
return factory.make_bundle(factory.ZipSpec(files=files), password=BUNDLE_PASSWORD)
|
||||
|
||||
|
||||
def missing_target_variable() -> bytes:
|
||||
document = factory.build_metadata_document(drop_parameter='target_variable')
|
||||
return factory.make_bundle(factory.ZipSpec(document=document), password=BUNDLE_PASSWORD)
|
||||
|
||||
|
||||
def script_in_the_tree() -> bytes:
|
||||
return factory.make_bundle(
|
||||
factory.ZipSpec(extra_files={'artifacts/run.sh': b'#!/bin/sh\nrm -rf /\n'}),
|
||||
password=BUNDLE_PASSWORD,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@mark.parametrize(
|
||||
('build', 'step', 'gate'),
|
||||
[
|
||||
(bad_magic, ImportStep.DOWNLOAD, 2),
|
||||
(wrong_digest, ImportStep.DOWNLOAD, 2),
|
||||
(tampered, ImportStep.DECRYPTION, 3),
|
||||
(traversal, ImportStep.ARCHIVE_INSPECTION, 4),
|
||||
(absolute, ImportStep.ARCHIVE_INSPECTION, 4),
|
||||
(missing_data_model, ImportStep.STRUCTURE_VALIDATION, 6),
|
||||
(missing_target_variable, ImportStep.STRUCTURE_VALIDATION, 6),
|
||||
(script_in_the_tree, ImportStep.CONTENT_POLICY, 7),
|
||||
],
|
||||
ids=lambda value: getattr(value, '__name__', str(value)),
|
||||
)
|
||||
async def test_a_gate_failure_reaches_the_activity_as_its_own_step(
|
||||
import_activity: ModelImport, tmp_path: Path, build, step: ImportStep, gate: int
|
||||
):
|
||||
path = tmp_path / 'bundle.sientia'
|
||||
path.write_bytes(build())
|
||||
request = open_input(str(path))
|
||||
|
||||
with raises(ImportGateRejectedError) as failure:
|
||||
await import_activity._open_bundle(request)
|
||||
|
||||
assert failure.value.details[0]['step'] == step.value
|
||||
assert failure.value.details[0]['gate'] == gate
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_rejected_bundle_leaves_no_extracted_tree(
|
||||
import_activity: ModelImport, tmp_path: Path
|
||||
):
|
||||
path = tmp_path / 'bundle.sientia'
|
||||
path.write_bytes(traversal())
|
||||
request = open_input(str(path))
|
||||
|
||||
with raises(ImportGateRejectedError):
|
||||
await import_activity._open_bundle(request)
|
||||
|
||||
extract_root = tmp_path / WORKFLOW_ID / 'extract'
|
||||
assert not extract_root.exists() or list(extract_root.rglob('*')) == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------- secrecy
|
||||
|
||||
|
||||
SECRET_PATHS = [
|
||||
('happy', lambda: factory.make_bundle(password=BUNDLE_PASSWORD), None),
|
||||
('wrong-bundle-password', lambda: factory.make_bundle(password='another-password'), None),
|
||||
('wrong-envelope-key', lambda: factory.make_bundle(password=BUNDLE_PASSWORD), 'other-key'),
|
||||
('malformed-envelope', lambda: factory.make_bundle(password=BUNDLE_PASSWORD), 'garbage'),
|
||||
('gate-2', bad_magic, None),
|
||||
('gate-3', tampered, None),
|
||||
('gate-4', traversal, None),
|
||||
('gate-6', missing_data_model, None),
|
||||
('gate-7', script_in_the_tree, None),
|
||||
]
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@mark.parametrize(
|
||||
('build', 'envelope_kind'),
|
||||
[(build, kind) for _, build, kind in SECRET_PATHS],
|
||||
ids=[name for name, _, _ in SECRET_PATHS],
|
||||
)
|
||||
async def test_no_failure_path_says_the_password_out_loud(
|
||||
import_activity: ModelImport, tmp_path: Path, build, envelope_kind: str | None
|
||||
):
|
||||
"""Every path through this activity, checked against the three things that must never surface."""
|
||||
path = tmp_path / 'bundle.sientia'
|
||||
path.write_bytes(build())
|
||||
if envelope_kind == 'other-key':
|
||||
envelope = factory.build_envelope(BUNDLE_PASSWORD, OTHER_ENVELOPE_KEY)
|
||||
elif envelope_kind == 'garbage':
|
||||
envelope = 'not-base64!!'
|
||||
else:
|
||||
envelope = factory.build_envelope(BUNDLE_PASSWORD, ENVELOPE_KEY)
|
||||
|
||||
raised: BaseException | None = None
|
||||
result = None
|
||||
try:
|
||||
result = await import_activity._open_bundle(open_input(str(path), envelope))
|
||||
except BaseException as error: # noqa: BLE001 - the point is that *whatever* it is says nothing
|
||||
raised = error
|
||||
|
||||
secrets = [
|
||||
BUNDLE_PASSWORD,
|
||||
'another-password',
|
||||
factory.envelope_key_b64(ENVELOPE_KEY),
|
||||
ENVELOPE_KEY.hex(),
|
||||
]
|
||||
surfaces = [observable_text(import_activity), str(result)]
|
||||
if raised is not None:
|
||||
surfaces += [str(raised), str(getattr(raised, 'details', '')), repr(raised)]
|
||||
for secret in secrets:
|
||||
for surface in surfaces:
|
||||
assert secret not in surface
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_envelope_never_leaves_the_activity_that_opens_it(
|
||||
import_activity: ModelImport, tmp_path: Path
|
||||
):
|
||||
"""The result is what crosses back into the workflow, and therefore into the event history."""
|
||||
path = tmp_path / 'bundle.sientia'
|
||||
path.write_bytes(factory.make_bundle(password=BUNDLE_PASSWORD))
|
||||
envelope = factory.build_envelope(BUNDLE_PASSWORD, ENVELOPE_KEY)
|
||||
|
||||
result = await import_activity._open_bundle(open_input(str(path), envelope))
|
||||
|
||||
assert envelope not in str(result)
|
||||
assert 'password' not in str(result).lower()
|
||||
|
||||
|
||||
# ------------------------------------------------------- the fused activity (bug 8)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_prepare_and_upload_import_bundle_runs_all_four_phases_end_to_end(
|
||||
import_activity: ModelImport, minio_client: MagicMock
|
||||
):
|
||||
"""Download, open, experiment creation and artifact upload, one activity, one attempt."""
|
||||
payload = factory.make_bundle(password=BUNDLE_PASSWORD)
|
||||
serve_object(minio_client, payload)
|
||||
digest = hashlib.sha256(payload).hexdigest()
|
||||
|
||||
mlflow_repository = import_activity.import_mlflow_repository
|
||||
mlflow_repository.get_experiment.return_value = MagicMock(experiment_id='17')
|
||||
client = mlflow_repository.client
|
||||
client.create_run.return_value.info.run_id = 'imported-run-id'
|
||||
|
||||
request = {
|
||||
'metadata': METADATA,
|
||||
'import_run_id': 77,
|
||||
'workflow_id': WORKFLOW_ID,
|
||||
'object_key': OBJECT_KEY,
|
||||
'expected_digest': digest,
|
||||
'password_envelope': factory.build_envelope(BUNDLE_PASSWORD, ENVELOPE_KEY),
|
||||
}
|
||||
|
||||
result = await import_activity.prepare_and_upload_import_bundle(request)
|
||||
|
||||
expected_run_name = derive_run_name('factory_model', result['metadata']['model_version'])
|
||||
assert result['metadata']['model_name'] == 'factory_model'
|
||||
assert result['experiment_id'] == '17'
|
||||
assert result['run_id'] == 'imported-run-id'
|
||||
assert result['run_name'] == expected_run_name
|
||||
mlflow_repository.get_experiment.assert_called_once_with(
|
||||
'factory_experiment', create_if_not_exists=True
|
||||
)
|
||||
client.create_run.assert_called_once_with(experiment_id='17', run_name=expected_run_name)
|
||||
client.log_artifacts.assert_called_once()
|
||||
client.set_terminated.assert_called_once_with('imported-run-id', 'FINISHED')
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_gate_failure_inside_the_fused_activity_still_names_its_step_and_gate(
|
||||
import_activity: ModelImport, minio_client: MagicMock
|
||||
):
|
||||
"""A gate rejection during the fused activity's decryption/extraction phase is unaffected."""
|
||||
payload = tampered()
|
||||
serve_object(minio_client, payload)
|
||||
digest = hashlib.sha256(payload).hexdigest()
|
||||
|
||||
request = {
|
||||
'metadata': METADATA,
|
||||
'import_run_id': 77,
|
||||
'workflow_id': WORKFLOW_ID,
|
||||
'object_key': OBJECT_KEY,
|
||||
'expected_digest': digest,
|
||||
'password_envelope': factory.build_envelope(BUNDLE_PASSWORD, ENVELOPE_KEY),
|
||||
}
|
||||
|
||||
with raises(ImportGateRejectedError) as failure:
|
||||
await import_activity.prepare_and_upload_import_bundle(request)
|
||||
|
||||
assert failure.value.details[0]['step'] == ImportStep.DECRYPTION.value
|
||||
assert failure.value.details[0]['gate'] == 3
|
||||
|
||||
|
||||
async def _fake_download(_input_data: dict) -> dict:
|
||||
return {'path': 'p', 'size': 1, 'digest': 'd'}
|
||||
|
||||
|
||||
async def _fake_open(_input_data: dict) -> dict:
|
||||
return {
|
||||
'extracted_dir': 'e',
|
||||
'metadata': {
|
||||
'model_name': 'factory_model',
|
||||
'model_version': '1',
|
||||
'experiment_name': 'factory_experiment',
|
||||
},
|
||||
'parameters': {},
|
||||
'digest': 'd',
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_misconfigured_repository_at_experiment_creation_is_not_wrapped(
|
||||
import_activity: ModelImport,
|
||||
):
|
||||
"""`ImportMisconfiguredError` is a config problem, not a broken dependency — it must reach the
|
||||
workflow untouched, not get rewrapped into a generic `ApplicationError`."""
|
||||
import_activity.import_mlflow_repository = None
|
||||
|
||||
with (
|
||||
patch.object(import_activity, '_download_bundle', side_effect=_fake_download),
|
||||
patch.object(import_activity, '_open_bundle', side_effect=_fake_open),
|
||||
raises(ImportMisconfiguredError),
|
||||
):
|
||||
await import_activity.prepare_and_upload_import_bundle({'metadata': METADATA})
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_broken_dependency_at_experiment_creation_is_tagged_with_its_step(
|
||||
import_activity: ModelImport,
|
||||
):
|
||||
"""A raw MLflow/network failure carries no location of its own, so the fused activity must tag
|
||||
it with the step it actually broke at — otherwise the workflow's fallback misclassifies it."""
|
||||
|
||||
async def fake_create_experiment(_input_data: dict) -> dict:
|
||||
raise RuntimeError('mlflow unreachable')
|
||||
|
||||
with (
|
||||
patch.object(import_activity, '_download_bundle', side_effect=_fake_download),
|
||||
patch.object(import_activity, '_open_bundle', side_effect=_fake_open),
|
||||
patch.object(import_activity, '_create_experiment', side_effect=fake_create_experiment),
|
||||
raises(ApplicationError, match='mlflow unreachable') as failure,
|
||||
):
|
||||
await import_activity.prepare_and_upload_import_bundle({'metadata': METADATA})
|
||||
|
||||
assert failure.value.details[0]['step'] == ImportStep.EXPERIMENT_CREATION.value
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_misconfigured_repository_at_artifact_upload_is_not_wrapped(
|
||||
import_activity: ModelImport,
|
||||
):
|
||||
async def fake_create_experiment(_input_data: dict) -> dict:
|
||||
import_activity.import_mlflow_repository = None
|
||||
return {'experiment_id': '1', 'experiment_name': 'factory_experiment'}
|
||||
|
||||
with (
|
||||
patch.object(import_activity, '_download_bundle', side_effect=_fake_download),
|
||||
patch.object(import_activity, '_open_bundle', side_effect=_fake_open),
|
||||
patch.object(import_activity, '_create_experiment', side_effect=fake_create_experiment),
|
||||
raises(ImportMisconfiguredError),
|
||||
):
|
||||
await import_activity.prepare_and_upload_import_bundle({'metadata': METADATA})
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_retry_of_the_fused_activity_redoes_all_four_phases(import_activity: ModelImport):
|
||||
"""A retry re-runs download, open, experiment creation and upload from the start.
|
||||
|
||||
Re-sending the same artifacts to the same run is safe (MLflow overwrites rather than
|
||||
duplicates), so the cost of a retry is bandwidth and time, not corruption — but it is a real
|
||||
cost, since nothing here is skipped on a second attempt.
|
||||
"""
|
||||
calls: list[str] = []
|
||||
|
||||
async def fake_download(_input_data: dict) -> dict:
|
||||
calls.append('download')
|
||||
return {'path': 'p', 'size': 1, 'digest': 'd'}
|
||||
|
||||
async def fake_open(_input_data: dict) -> dict:
|
||||
calls.append('open')
|
||||
return {
|
||||
'extracted_dir': 'e',
|
||||
'metadata': {
|
||||
'model_name': 'factory_model',
|
||||
'model_version': '1',
|
||||
'experiment_name': 'factory_experiment',
|
||||
},
|
||||
'parameters': {},
|
||||
'digest': 'd',
|
||||
}
|
||||
|
||||
async def fake_create_experiment(_input_data: dict) -> dict:
|
||||
calls.append('experiment')
|
||||
return {'experiment_id': '1', 'experiment_name': 'factory_experiment'}
|
||||
|
||||
attempts = {'count': 0}
|
||||
|
||||
async def fake_upload(_input_data: dict) -> dict:
|
||||
calls.append('upload')
|
||||
attempts['count'] += 1
|
||||
if attempts['count'] == 1:
|
||||
raise RuntimeError('artifact upload failed on the first attempt')
|
||||
return {'run_id': 'r', 'run_name': 'import-factory_model-v1'}
|
||||
|
||||
with (
|
||||
patch.object(import_activity, '_download_bundle', side_effect=fake_download),
|
||||
patch.object(import_activity, '_open_bundle', side_effect=fake_open),
|
||||
patch.object(import_activity, '_create_experiment', side_effect=fake_create_experiment),
|
||||
patch.object(import_activity, '_upload_artifacts', side_effect=fake_upload),
|
||||
):
|
||||
with raises(ApplicationError, match='artifact upload failed on the first attempt'):
|
||||
await import_activity.prepare_and_upload_import_bundle({'metadata': METADATA})
|
||||
|
||||
await import_activity.prepare_and_upload_import_bundle({'metadata': METADATA})
|
||||
|
||||
assert calls == [
|
||||
'download',
|
||||
'open',
|
||||
'experiment',
|
||||
'upload',
|
||||
'download',
|
||||
'open',
|
||||
'experiment',
|
||||
'upload',
|
||||
]
|
||||
890
tests/laborious/activities/test_model_import_provisioning.py
Normal file
890
tests/laborious/activities/test_model_import_provisioning.py
Normal file
@@ -0,0 +1,890 @@
|
||||
"""Provisioning: the experiment, the run, the registered version, the document, the cleanup.
|
||||
|
||||
The boundary this file defends is where the import stops. A finished import leaves a version whose
|
||||
stage was never touched and a document that is inert — promoting and activating are an operator's
|
||||
decisions, owned by another card — so the tests assert not only what is written but what is *not*:
|
||||
no stage transition, no alias, no `active: true`, no update of a document already there.
|
||||
|
||||
The rest is the model listing's arithmetic (`max(int(id)) + 1`, which is not `count` and not a
|
||||
lexicographic sort) and the two operations that must never fail an import that otherwise worked:
|
||||
removing the scratch directory and putting an expiry rule on the bundle prefix.
|
||||
"""
|
||||
|
||||
import json
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import mlflow as mlflow_lib
|
||||
from minio.commonconfig import ENABLED
|
||||
from mlflow.exceptions import MlflowException
|
||||
from mlflow.protos.databricks_pb2 import INTERNAL_ERROR, RESOURCE_DOES_NOT_EXIST
|
||||
from pytest import fixture, mark, raises
|
||||
|
||||
from laborious.activities.model_import import ModelImport
|
||||
from laborious.utils.models.import_errors import ImportMisconfiguredError
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
from tests.helpers.status_engine import FakeStatusEngine
|
||||
|
||||
METADATA = {
|
||||
'model_id': '77',
|
||||
'model_name': 'factory_model',
|
||||
'workflow_name': 'import_model',
|
||||
'schedule_name': '-',
|
||||
}
|
||||
WORKFLOW_ID = 'import_model-77-abc'
|
||||
RUN_ID = '0123456789abcdef0123456789abcdef'
|
||||
# What a tracking server with proxied artifacts reports as a run's artifact root, and therefore what
|
||||
# the version's `source` is built from. The run id sits at index 2 of this URI — the whole reason the
|
||||
# registration cannot store a literal `runs:/…` source.
|
||||
ARTIFACT_URI = f'mlflow-artifacts:/17/{RUN_ID}/artifacts'
|
||||
RESOLVED_SOURCE = f'{ARTIFACT_URI}/prediction_model'
|
||||
|
||||
BUNDLE_PARAMETERS = {
|
||||
'model_name': 'factory_model',
|
||||
'target_variable': 'SE',
|
||||
'input_variables': "['feature_a', 'feature_b']",
|
||||
'retrain': 'False',
|
||||
}
|
||||
BUNDLE_METADATA = {
|
||||
'model_name': 'factory_model',
|
||||
'experiment_name': 'factory_experiment',
|
||||
'run_id': 'aa11bb22cc33dd44ee55ff6677889900',
|
||||
'model_version': '2',
|
||||
'model_project': 'factory_project',
|
||||
'export_timestamp': '2026-08-25T12:00:00+00:00',
|
||||
}
|
||||
|
||||
# The whole vocabulary of promotion. None of it may be reached by this change.
|
||||
PROMOTION_CALLS = (
|
||||
'transition_model_version_stage',
|
||||
'set_registered_model_alias',
|
||||
'set_model_version_tag',
|
||||
'update_model_version',
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
def mlflow_client() -> MagicMock:
|
||||
client = MagicMock()
|
||||
client.create_run.return_value.info.run_id = RUN_ID
|
||||
client.get_run.return_value.info.artifact_uri = ARTIFACT_URI
|
||||
client.create_model_version.return_value.version = 4
|
||||
return client
|
||||
|
||||
|
||||
@fixture
|
||||
def mongo() -> AsyncMock:
|
||||
repository = AsyncMock()
|
||||
repository.find_one.return_value = None
|
||||
repository.distinct.return_value = []
|
||||
# `_next_model_document_id` reaches Mongo through the repository's plain pymongo `database`
|
||||
# handle, not through the repository's own async methods, with a *synchronous* call
|
||||
# (`find_one_and_update` is never awaited). A sane default so tests that don't care about the
|
||||
# exact id still get one; tests that do override this return value (or its `side_effect`) per
|
||||
# collection access.
|
||||
repository.database.__getitem__.return_value.find_one_and_update.return_value = {
|
||||
'_id': 'models',
|
||||
'seq': 1,
|
||||
}
|
||||
return repository
|
||||
|
||||
|
||||
def counter(mongo: AsyncMock) -> MagicMock:
|
||||
"""The mock pymongo `Collection` behind `mongo_repository.database[collection]`."""
|
||||
return mongo.database.__getitem__.return_value.find_one_and_update
|
||||
|
||||
|
||||
@fixture
|
||||
def import_activity(tmp_path: Path, mlflow_client: MagicMock, mongo: AsyncMock) -> ModelImport:
|
||||
minio_repository = MagicMock()
|
||||
mlflow_repository = MagicMock()
|
||||
mlflow_repository.client = mlflow_client
|
||||
activity = ModelImport(
|
||||
status_config={
|
||||
'host': 'h',
|
||||
'port': 5432,
|
||||
'dbname': 'sientia-core-mlops-bff',
|
||||
'user': 'u',
|
||||
'password': 'p',
|
||||
'schema': 'public',
|
||||
},
|
||||
import_config={
|
||||
'password_key': None,
|
||||
'bucket': 'sientia',
|
||||
'prefix': 'imported_models/',
|
||||
'work_dir': str(tmp_path),
|
||||
'max_object_bytes': 1024,
|
||||
'max_entries': 10,
|
||||
'max_uncompressed_bytes': 4096,
|
||||
'max_compression_ratio': 200.0,
|
||||
'models_collection': 'models',
|
||||
'counters_collection': 'model_id_counters',
|
||||
'retention_days': 7,
|
||||
},
|
||||
minio_repository=minio_repository,
|
||||
mlflow_repository=mlflow_repository,
|
||||
mongo_repository=mongo,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
# `status_engine` is created lazily on first access; pre-seed it with the recording double.
|
||||
activity._status_engine = FakeStatusEngine()
|
||||
activity.emit_metric = AsyncMock()
|
||||
activity.send_notification_async = AsyncMock()
|
||||
return activity
|
||||
|
||||
|
||||
def base(**overrides) -> dict:
|
||||
return {
|
||||
'metadata': METADATA,
|
||||
'import_run_id': 77,
|
||||
'workflow_id': WORKFLOW_ID,
|
||||
**overrides,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------- the experiment
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_experiment_keeps_the_origin_s_name(import_activity: ModelImport):
|
||||
repository = import_activity.import_mlflow_repository
|
||||
repository.get_experiment.return_value = MagicMock(experiment_id='17')
|
||||
|
||||
result = await import_activity._create_experiment(base(experiment_name='factory_experiment'))
|
||||
|
||||
repository.get_experiment.assert_called_once_with(
|
||||
'factory_experiment', create_if_not_exists=True
|
||||
)
|
||||
assert result == {'experiment_id': '17', 'experiment_name': 'factory_experiment'}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_missing_mlflow_repository_is_a_configuration_error(import_activity: ModelImport):
|
||||
import_activity.import_mlflow_repository = None
|
||||
payload = base(experiment_name='e')
|
||||
|
||||
with raises(ImportMisconfiguredError, match='MLflow repository not initialized'):
|
||||
await import_activity._create_experiment(payload)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- the run and its tree
|
||||
|
||||
|
||||
def upload_input(extracted_dir: Path, **overrides) -> dict:
|
||||
return base(
|
||||
experiment_id='17',
|
||||
run_name='import-factory_model-v2',
|
||||
extracted_dir=str(extracted_dir),
|
||||
parameters=BUNDLE_PARAMETERS,
|
||||
source=BUNDLE_METADATA,
|
||||
**overrides,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_upload_artifacts_requires_the_mlflow_repository(
|
||||
import_activity: ModelImport, tmp_path: Path
|
||||
):
|
||||
import_activity.import_mlflow_repository = None
|
||||
payload = upload_input(tmp_path / 'extract')
|
||||
|
||||
with raises(ImportMisconfiguredError, match='MLflow repository not initialized'):
|
||||
await import_activity._upload_artifacts(payload)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_every_import_creates_a_new_run_with_the_recorded_name(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
|
||||
):
|
||||
"""The name is passed in, never recomputed, so the record and the run cannot diverge."""
|
||||
result = await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
|
||||
|
||||
mlflow_client.create_run.assert_called_once_with(
|
||||
experiment_id='17', run_name='import-factory_model-v2'
|
||||
)
|
||||
assert result == {'run_id': RUN_ID, 'run_name': 'import-factory_model-v2'}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_artifact_tree_is_logged_from_the_bundle_s_artifacts_directory(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
|
||||
):
|
||||
extracted = tmp_path / 'extract'
|
||||
|
||||
await import_activity._upload_artifacts(upload_input(extracted))
|
||||
|
||||
mlflow_client.log_artifacts.assert_called_once_with(RUN_ID, str(extracted / 'artifacts'))
|
||||
|
||||
|
||||
def _logged_params(mlflow_client: MagicMock) -> dict:
|
||||
"""Flatten every `client.log_batch(run_id, params=[...])` call into one key/value map."""
|
||||
return {
|
||||
param.key: param.value
|
||||
for call in mlflow_client.log_batch.call_args_list
|
||||
for param in call.kwargs['params']
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_bundle_s_parameters_are_logged_verbatim(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
|
||||
):
|
||||
await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
|
||||
|
||||
logged = {
|
||||
key: value
|
||||
for key, value in _logged_params(mlflow_client).items()
|
||||
if not key.startswith('source.')
|
||||
}
|
||||
assert logged == BUNDLE_PARAMETERS
|
||||
assert logged['target_variable'] == 'SE'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_origin_is_logged_as_source_parameters(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
|
||||
):
|
||||
await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
|
||||
|
||||
logged = {
|
||||
key: value
|
||||
for key, value in _logged_params(mlflow_client).items()
|
||||
if key.startswith('source.')
|
||||
}
|
||||
assert logged == {f'source.{key}': value for key, value in BUNDLE_METADATA.items()}
|
||||
# A record of where the model came from, not a link: the origin run id is never dereferenced.
|
||||
assert logged['source.run_id'] == BUNDLE_METADATA['run_id']
|
||||
mlflow_client.get_run.assert_not_called()
|
||||
mlflow_client.download_artifacts.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_parameters_are_logged_in_batches_of_at_most_100(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
|
||||
):
|
||||
"""More params than `log_batch`'s per-call limit must span multiple calls, each within the
|
||||
limit, and every key must still be logged exactly once across all of them."""
|
||||
many_parameters = {f'param_{i}': str(i) for i in range(150)}
|
||||
payload = upload_input(tmp_path / 'extract')
|
||||
payload['parameters'] = many_parameters
|
||||
|
||||
await import_activity._upload_artifacts(payload)
|
||||
|
||||
calls = mlflow_client.log_batch.call_args_list
|
||||
assert len(calls) > 1
|
||||
for call in calls:
|
||||
assert len(call.kwargs['params']) <= 100
|
||||
|
||||
logged = _logged_params(mlflow_client)
|
||||
expected = {
|
||||
**many_parameters,
|
||||
**{f'source.{key}': value for key, value in BUNDLE_METADATA.items()},
|
||||
}
|
||||
assert logged == expected
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_upload_artifacts_never_uses_log_param(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
|
||||
):
|
||||
await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
|
||||
|
||||
mlflow_client.log_param.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_run_is_closed_when_its_tree_is_in(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
|
||||
):
|
||||
await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
|
||||
|
||||
mlflow_client.set_terminated.assert_called_once_with(RUN_ID, 'FINISHED')
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------ registration
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_register_version_requires_the_mlflow_repository(import_activity: ModelImport):
|
||||
import_activity.import_mlflow_repository = None
|
||||
payload = base(model_name='factory_model', run_id=RUN_ID)
|
||||
|
||||
with raises(ImportMisconfiguredError, match='MLflow repository not initialized'):
|
||||
await import_activity.register_import_model_version(payload)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_version_source_is_the_shape_the_legacy_loader_parses(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock
|
||||
):
|
||||
result = await import_activity.register_import_model_version(
|
||||
base(model_name='factory_model', run_id=RUN_ID)
|
||||
)
|
||||
|
||||
assert result == {'model_name': 'factory_model', 'version': '4', 'source': RESOLVED_SOURCE}
|
||||
mlflow_client.create_model_version.assert_called_once_with(
|
||||
name='factory_model', source=RESOLVED_SOURCE, run_id=RUN_ID
|
||||
)
|
||||
# The one property the legacy loader depends on, asserted directly.
|
||||
assert RESOLVED_SOURCE.split('/')[2] == RUN_ID
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_literal_runs_uri_would_not_be_parseable_by_the_legacy_loader():
|
||||
"""The trap this registration avoids: `create_model_version` stores `source` verbatim, and a
|
||||
`runs:/` URI has only two slashes, so the loader's index 2 is the artifact path."""
|
||||
assert f'runs:/{RUN_ID}/prediction_model'.split('/')[2] == 'prediction_model'
|
||||
assert 'models:/m-9f1c'.split('/')[1:] == ['m-9f1c']
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_registered_model_is_created_only_when_it_is_absent(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock
|
||||
):
|
||||
"""The only exception `get_registered_model` may treat as "does not exist yet": an
|
||||
`MlflowException` whose `error_code` is literally `RESOURCE_DOES_NOT_EXIST`."""
|
||||
mlflow_client.get_registered_model.side_effect = MlflowException(
|
||||
'not found', error_code=RESOURCE_DOES_NOT_EXIST
|
||||
)
|
||||
|
||||
result = await import_activity.register_import_model_version(
|
||||
base(model_name='new_model', run_id=RUN_ID)
|
||||
)
|
||||
|
||||
mlflow_client.create_registered_model.assert_called_once_with('new_model')
|
||||
mlflow_client.create_model_version.assert_called_once()
|
||||
assert result['model_name'] == 'new_model'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_second_import_reuses_the_registered_model(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock
|
||||
):
|
||||
mlflow_client.get_registered_model.return_value = MagicMock(name='factory_model')
|
||||
|
||||
await import_activity.register_import_model_version(
|
||||
base(model_name='factory_model', run_id=RUN_ID)
|
||||
)
|
||||
|
||||
mlflow_client.create_registered_model.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_different_mlflow_error_is_not_treated_as_absent(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock
|
||||
):
|
||||
"""Only `RESOURCE_DOES_NOT_EXIST` means "create it". Any other MLflow error code — a transient
|
||||
failure, an internal error — must not be swallowed and must not trigger a spurious create."""
|
||||
error = MlflowException('backend is unavailable', error_code=INTERNAL_ERROR)
|
||||
mlflow_client.get_registered_model.side_effect = error
|
||||
payload = base(model_name='factory_model', run_id=RUN_ID)
|
||||
|
||||
with raises(MlflowException, match='backend is unavailable') as excinfo:
|
||||
await import_activity.register_import_model_version(payload)
|
||||
|
||||
assert excinfo.value is error
|
||||
mlflow_client.create_registered_model.assert_not_called()
|
||||
mlflow_client.create_model_version.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_non_mlflow_exception_is_never_treated_as_absent(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock
|
||||
):
|
||||
"""A bare `ConnectionError` (or any non-`MlflowException`) is not this function's business to
|
||||
interpret at all: it must propagate unmodified, exactly as any other unexpected error would."""
|
||||
mlflow_client.get_registered_model.side_effect = ConnectionError('connection refused')
|
||||
payload = base(model_name='factory_model', run_id=RUN_ID)
|
||||
|
||||
with raises(ConnectionError, match='connection refused'):
|
||||
await import_activity.register_import_model_version(payload)
|
||||
|
||||
mlflow_client.create_registered_model.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_import_never_promotes_anything(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
|
||||
):
|
||||
"""A `COMPLETED` import is a registered model, not a servable one."""
|
||||
await import_activity._create_experiment(base(experiment_name='factory_experiment'))
|
||||
await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
|
||||
await import_activity.register_import_model_version(
|
||||
base(model_name='factory_model', run_id=RUN_ID)
|
||||
)
|
||||
|
||||
for call_name in PROMOTION_CALLS:
|
||||
getattr(mlflow_client, call_name).assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_no_provisioning_call_passes_a_stage(
|
||||
import_activity: ModelImport, mlflow_client: MagicMock, tmp_path: Path
|
||||
):
|
||||
await import_activity._upload_artifacts(upload_input(tmp_path / 'extract'))
|
||||
result = await import_activity.register_import_model_version(
|
||||
base(model_name='factory_model', run_id=RUN_ID)
|
||||
)
|
||||
|
||||
for call in mlflow_client.method_calls:
|
||||
assert 'stage' not in call.kwargs
|
||||
assert 'alias' not in call.kwargs
|
||||
assert 'stage' not in result
|
||||
|
||||
|
||||
# ------------------------------------------------- the legacy loader reads what we wrote
|
||||
|
||||
|
||||
def legacy_repository() -> MLFlowRepository:
|
||||
"""A real `MLFlowRepository`, with the process-wide side effects of its constructor stubbed."""
|
||||
with (
|
||||
patch.object(mlflow_lib, 'set_tracking_uri'),
|
||||
patch.object(mlflow_lib.tracking, 'MlflowClient'),
|
||||
):
|
||||
return MLFlowRepository(
|
||||
host='http://localhost:5000',
|
||||
username='admin',
|
||||
password='admin',
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
|
||||
|
||||
def registry_holding(source: str) -> MagicMock:
|
||||
client = MagicMock()
|
||||
client.search_registered_models.return_value = [MagicMock()]
|
||||
version = MagicMock(version='4', source=source)
|
||||
version.current_stage = 'Production'
|
||||
client.search_model_versions.return_value = [version]
|
||||
return client
|
||||
|
||||
|
||||
def test_the_legacy_loader_resolves_the_run_the_import_created():
|
||||
"""The contract the source shape exists for: the real production code, fed a version this
|
||||
import would create, returns this import's run id."""
|
||||
repository = legacy_repository()
|
||||
repository.client = registry_holding(RESOLVED_SOURCE)
|
||||
|
||||
assert repository.get_model_run_id('factory_model') == RUN_ID
|
||||
|
||||
|
||||
def test_the_legacy_loader_would_read_a_runs_uri_as_the_artifact_path():
|
||||
"""Left as an executable warning: registering with the unresolved URI silently breaks loading,
|
||||
because nothing raises — the loader just returns the wrong string."""
|
||||
repository = legacy_repository()
|
||||
repository.client = registry_holding(f'runs:/{RUN_ID}/prediction_model')
|
||||
|
||||
assert repository.get_model_run_id('factory_model') == 'prediction_model'
|
||||
|
||||
|
||||
def test_the_legacy_loader_cannot_read_a_logged_model_uri():
|
||||
repository = legacy_repository()
|
||||
repository.client = registry_holding('models:/m-9f1c')
|
||||
|
||||
with raises(mlflow_lib.exceptions.MlflowException, match='invalid source URI'):
|
||||
repository.get_model_run_id('factory_model')
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- the model document
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_document_id_is_whatever_the_atomic_counter_returns(
|
||||
import_activity: ModelImport, mongo: AsyncMock
|
||||
):
|
||||
"""`id` is no longer computed here at all: it is `str(seq)` off the counter document
|
||||
`find_one_and_update` hands back, whatever gaps or history the collection has."""
|
||||
counter(mongo).return_value = {'_id': 'models', 'seq': 31}
|
||||
|
||||
result = await import_activity.write_import_model_document(
|
||||
base(model_name='factory_model', target='SE')
|
||||
)
|
||||
|
||||
assert result['id'] == '31'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_document_id_is_the_counter_s_seq_verbatim_not_reparsed(
|
||||
import_activity: ModelImport, mongo: AsyncMock
|
||||
):
|
||||
"""A value large enough that the old max-scan's `"9" > "30"` lexicographic bug would have
|
||||
mangled it, demonstrating the new id is never routed through any string comparison at all."""
|
||||
counter(mongo).return_value = {'_id': 'models', 'seq': 100}
|
||||
|
||||
result = await import_activity.write_import_model_document(
|
||||
base(model_name='factory_model', target='SE')
|
||||
)
|
||||
|
||||
assert result['id'] == '100'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_an_empty_collection_starts_at_one(import_activity: ModelImport, mongo: AsyncMock):
|
||||
"""`$inc` on a counter document that does not exist yet starts from zero, so the first import
|
||||
against a fresh collection gets id `1`."""
|
||||
counter(mongo).return_value = {'_id': 'models', 'seq': 1}
|
||||
|
||||
result = await import_activity.write_import_model_document(
|
||||
base(model_name='factory_model', target='SE')
|
||||
)
|
||||
|
||||
assert result['id'] == '1'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_document_is_exactly_the_contracted_shape(
|
||||
import_activity: ModelImport, mongo: AsyncMock
|
||||
):
|
||||
counter(mongo).return_value = {'_id': 'models', 'seq': 8}
|
||||
|
||||
await import_activity.write_import_model_document(base(model_name='factory_model', target='SE'))
|
||||
|
||||
collection, document, metadata = mongo.insert.await_args.args
|
||||
assert collection == 'models'
|
||||
assert metadata == METADATA
|
||||
assert document == {
|
||||
'id': '8',
|
||||
'name': 'factory_model',
|
||||
'active': False,
|
||||
'model_config': {
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'joblib',
|
||||
'retention_minutes': 60,
|
||||
'target': 'SE',
|
||||
'skip_predict_timestamp': True,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@mark.parametrize(
|
||||
'model_name, target',
|
||||
[('factory_model', 'SE'), ('another_model', 'temperature')],
|
||||
)
|
||||
async def test_the_document_always_carries_skip_predict_timestamp_true(
|
||||
import_activity: ModelImport, mongo: AsyncMock, model_name: str, target: str
|
||||
):
|
||||
"""The opt-out is decided by the importer at write time, never left for an operator to set."""
|
||||
counter(mongo).return_value = {'_id': 'models', 'seq': 8}
|
||||
|
||||
await import_activity.write_import_model_document(base(model_name=model_name, target=target))
|
||||
|
||||
_, document, _ = mongo.insert.await_args.args
|
||||
assert document['model_config']['skip_predict_timestamp'] is True
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_document_carries_no_field_the_runtime_does_not_read(
|
||||
import_activity: ModelImport, mongo: AsyncMock
|
||||
):
|
||||
await import_activity.write_import_model_document(base(model_name='factory_model', target='SE'))
|
||||
|
||||
_, document, _ = mongo.insert.await_args.args
|
||||
for absent in ('is_compressed', 'retention_target', 'stage', 'source', 'project'):
|
||||
assert absent not in document
|
||||
assert absent not in document['model_config']
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_name_already_in_the_collection_is_left_exactly_as_it_is(
|
||||
import_activity: ModelImport, mongo: AsyncMock
|
||||
):
|
||||
"""`active`, `target` and the flavors of a model already on the platform are operator-owned."""
|
||||
mongo.find_one.return_value = {'id': '12', 'name': 'factory_model', 'active': True}
|
||||
|
||||
result = await import_activity.write_import_model_document(
|
||||
base(model_name='factory_model', target='SE')
|
||||
)
|
||||
|
||||
assert result == {'written': False, 'skipped': True, 'id': '12'}
|
||||
mongo.insert.assert_not_awaited()
|
||||
# No counter increment is spent on a name that was already there: a duplicate import must
|
||||
# never move the sequence forward for a document it isn't going to write.
|
||||
counter(mongo).assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_two_imports_in_a_row_never_reuse_an_id(
|
||||
import_activity: ModelImport, mongo: AsyncMock
|
||||
):
|
||||
"""No dependency on `distinct` at all: two sequential calls simply relay whatever an atomic
|
||||
(fake, here) counter hands back, strictly increasing, with no scan of existing documents."""
|
||||
seq = {'value': 40}
|
||||
|
||||
def increment(*args, **kwargs): # noqa: ARG001 - mirrors pymongo's call signature
|
||||
seq['value'] += 1
|
||||
return {'_id': 'models', 'seq': seq['value']}
|
||||
|
||||
counter(mongo).side_effect = increment
|
||||
|
||||
first_id, _ = await import_activity._next_model_document_id('models', METADATA)
|
||||
second_id, _ = await import_activity._next_model_document_id('models', METADATA)
|
||||
|
||||
assert (first_id, second_id) == ('41', '42')
|
||||
assert int(second_id) > int(first_id)
|
||||
mongo.distinct.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_duplicate_id_fails_the_step_without_recomputing_it(
|
||||
import_activity: ModelImport, mongo: AsyncMock
|
||||
):
|
||||
"""No retry and no second guess: the unique index says the collection moved under us."""
|
||||
counter(mongo).return_value = {'_id': 'models', 'seq': 2}
|
||||
mongo.insert.side_effect = RuntimeError('E11000 duplicate key error: id_1')
|
||||
payload = base(model_name='factory_model', target='SE')
|
||||
|
||||
with raises(RuntimeError, match='duplicate key'):
|
||||
await import_activity.write_import_model_document(payload)
|
||||
|
||||
assert counter(mongo).call_count == 1
|
||||
assert mongo.insert.await_count == 1
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_import_only_reads_and_inserts_documents(
|
||||
import_activity: ModelImport, mongo: AsyncMock
|
||||
):
|
||||
"""No update, no replace, no upsert on the model document itself: an import cannot change a
|
||||
model already on the platform. The atomic counter's own `$inc` against a dedicated counters
|
||||
collection is a different concern and is explicitly allowed."""
|
||||
await import_activity.write_import_model_document(base(model_name='factory_model', target='SE'))
|
||||
|
||||
used = {name for name, _, _ in mongo.mock_calls}
|
||||
assert used <= {
|
||||
'find_one',
|
||||
'insert',
|
||||
'database.__getitem__',
|
||||
'database.__getitem__().find_one_and_update',
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_missing_mongo_repository_is_a_configuration_error(
|
||||
import_activity: ModelImport,
|
||||
):
|
||||
import_activity.mongo_repository = None
|
||||
payload = base(model_name='m', target='SE')
|
||||
|
||||
with raises(ImportMisconfiguredError, match='MongoDB repository not initialized'):
|
||||
await import_activity.write_import_model_document(payload)
|
||||
|
||||
|
||||
def test_no_import_module_writes_an_active_document():
|
||||
"""Asserted from the source too, so `active: false` cannot be flipped by a later edit.
|
||||
|
||||
`upsert` is checked against `write_import_model_document`'s own body only, not the whole file:
|
||||
`_next_model_document_id`'s atomic counter legitimately upserts a *counter* document in a
|
||||
dedicated collection — a different concern from ever updating, replacing or upserting the model
|
||||
document itself, which this test still forbids absolutely.
|
||||
"""
|
||||
source = Path('laborious/activities/model_import.py').read_text()
|
||||
start = source.index('async def write_import_model_document')
|
||||
end = source.index('async def _next_model_document_id')
|
||||
document_write = source[start:end]
|
||||
|
||||
assert "'active': False" in source
|
||||
assert "'active': True" not in source
|
||||
for forbidden in ('update_one', 'replace_one', 'update_many', 'upsert'):
|
||||
assert forbidden not in document_write
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ cleanup and retention
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_cleanup_removes_this_import_s_scratch_directory(
|
||||
import_activity: ModelImport, tmp_path: Path
|
||||
):
|
||||
work_dir = tmp_path / WORKFLOW_ID
|
||||
(work_dir / 'extract' / 'artifacts').mkdir(parents=True)
|
||||
(work_dir / 'bundle.sientia').write_bytes(b'downloaded')
|
||||
|
||||
result = await import_activity.cleanup_import_files(base())
|
||||
|
||||
assert result == {'cleaned': True, 'existed': True}
|
||||
assert not work_dir.exists()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_cleanup_of_a_run_that_never_downloaded_anything_is_not_a_failure(
|
||||
import_activity: ModelImport,
|
||||
):
|
||||
result = await import_activity.cleanup_import_files(base())
|
||||
|
||||
assert result == {'cleaned': True, 'existed': False}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_cleanup_that_fails_is_reported_and_swallowed(
|
||||
import_activity: ModelImport, tmp_path: Path
|
||||
):
|
||||
"""The user's model was imported; a leaked directory is a log line, not an ERROR."""
|
||||
(tmp_path / WORKFLOW_ID).mkdir()
|
||||
|
||||
with patch.object(shutil, 'rmtree', side_effect=OSError('device or resource busy')):
|
||||
result = await import_activity.cleanup_import_files(base())
|
||||
|
||||
assert result['cleaned'] is False
|
||||
assert result['reason'] == 'OSError'
|
||||
logged = import_activity.logger.custom_error.call_args_list[0].args[0]
|
||||
assert 'IMPORT_CLEANUP_INCOMPLETE' in logged
|
||||
import_activity.send_notification_async.assert_awaited_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_cleanup_never_touches_the_uploaded_object(
|
||||
import_activity: ModelImport, tmp_path: Path
|
||||
):
|
||||
(tmp_path / WORKFLOW_ID).mkdir()
|
||||
|
||||
await import_activity.cleanup_import_files(base())
|
||||
|
||||
client = import_activity.import_minio_repository.minio_client
|
||||
client.remove_object.assert_not_called()
|
||||
client.remove_objects.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_ensure_retention_requires_the_minio_repository(import_activity: ModelImport):
|
||||
import_activity.import_minio_repository = None
|
||||
payload = base()
|
||||
|
||||
with raises(ImportMisconfiguredError, match='Minio repository not initialized'):
|
||||
await import_activity.ensure_import_bundle_retention(payload)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_expiry_rule_is_applied_to_the_import_prefix(import_activity: ModelImport):
|
||||
client = import_activity.import_minio_repository.minio_client
|
||||
client.get_bucket_lifecycle.return_value = None
|
||||
|
||||
result = await import_activity.ensure_import_bundle_retention(base())
|
||||
|
||||
assert result == {'bucket': 'sientia', 'prefix': 'imported_models/', 'days': 7, 'rules': 1}
|
||||
config = client.set_bucket_lifecycle.call_args.args[1]
|
||||
rule = config.rules[0]
|
||||
assert rule.rule_id == 'sientia-import-bundle-expiry'
|
||||
assert rule.expiration.days == 7
|
||||
assert rule.rule_filter.prefix == 'imported_models/'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_applying_the_rule_twice_leaves_one_rule(import_activity: ModelImport):
|
||||
client = import_activity.import_minio_repository.minio_client
|
||||
client.get_bucket_lifecycle.return_value = None
|
||||
|
||||
await import_activity.ensure_import_bundle_retention(base())
|
||||
applied = client.set_bucket_lifecycle.call_args.args[1]
|
||||
client.get_bucket_lifecycle.return_value = applied
|
||||
await import_activity.ensure_import_bundle_retention(base())
|
||||
|
||||
config = client.set_bucket_lifecycle.call_args.args[1]
|
||||
assert [rule.rule_id for rule in config.rules] == ['sientia-import-bundle-expiry']
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_rule_that_belongs_to_somebody_else_is_kept(import_activity: ModelImport):
|
||||
client = import_activity.import_minio_repository.minio_client
|
||||
existing = MagicMock()
|
||||
existing.rules = [MagicMock(rule_id='predictions-offload-expiry')]
|
||||
client.get_bucket_lifecycle.return_value = existing
|
||||
|
||||
await import_activity.ensure_import_bundle_retention(base())
|
||||
|
||||
config = client.set_bucket_lifecycle.call_args.args[1]
|
||||
assert [rule.rule_id for rule in config.rules] == [
|
||||
'predictions-offload-expiry',
|
||||
'sientia-import-bundle-expiry',
|
||||
]
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_bucket_with_no_lifecycle_configuration_at_all_is_handled(
|
||||
import_activity: ModelImport,
|
||||
):
|
||||
"""Some deployments answer with an error, others with `None`; both mean "no rule yet"."""
|
||||
client = import_activity.import_minio_repository.minio_client
|
||||
client.get_bucket_lifecycle.side_effect = RuntimeError('NoSuchLifecycleConfiguration')
|
||||
|
||||
result = await import_activity.ensure_import_bundle_retention(base())
|
||||
|
||||
assert result['rules'] == 1
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_retention_rule_can_be_applied_to_a_requested_bucket(
|
||||
import_activity: ModelImport,
|
||||
):
|
||||
client = import_activity.import_minio_repository.minio_client
|
||||
client.get_bucket_lifecycle.return_value = None
|
||||
|
||||
result = await import_activity.ensure_import_bundle_retention(base(bucket='other-bucket'))
|
||||
|
||||
assert result['bucket'] == 'other-bucket'
|
||||
assert client.set_bucket_lifecycle.call_args.args[0] == 'other-bucket'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_correct_existing_rule_is_not_rewritten(import_activity: ModelImport):
|
||||
client = import_activity.import_minio_repository.minio_client
|
||||
existing = MagicMock()
|
||||
matching_rule = MagicMock(rule_id='sientia-import-bundle-expiry', status=ENABLED)
|
||||
matching_rule.expiration.days = 7
|
||||
matching_rule.rule_filter.prefix = 'imported_models/'
|
||||
existing.rules = [matching_rule]
|
||||
client.get_bucket_lifecycle.return_value = existing
|
||||
|
||||
result = await import_activity.ensure_import_bundle_retention(base())
|
||||
|
||||
client.set_bucket_lifecycle.assert_not_called()
|
||||
assert result == {'bucket': 'sientia', 'prefix': 'imported_models/', 'days': 7, 'rules': 1}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_drifted_existing_rule_is_rewritten(import_activity: ModelImport):
|
||||
client = import_activity.import_minio_repository.minio_client
|
||||
existing = MagicMock()
|
||||
drifted_rule = MagicMock(rule_id='sientia-import-bundle-expiry', status=ENABLED)
|
||||
drifted_rule.expiration.days = 3
|
||||
drifted_rule.rule_filter.prefix = 'imported_models/'
|
||||
existing.rules = [drifted_rule]
|
||||
client.get_bucket_lifecycle.return_value = existing
|
||||
|
||||
await import_activity.ensure_import_bundle_retention(base())
|
||||
|
||||
config = client.set_bucket_lifecycle.call_args.args[1]
|
||||
corrected = next(
|
||||
rule for rule in config.rules if rule.rule_id == 'sientia-import-bundle-expiry'
|
||||
)
|
||||
assert corrected.expiration.days == 7
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_an_unrelated_rule_survives_a_skipped_rewrite(import_activity: ModelImport):
|
||||
client = import_activity.import_minio_repository.minio_client
|
||||
existing = MagicMock()
|
||||
matching_rule = MagicMock(rule_id='sientia-import-bundle-expiry', status=ENABLED)
|
||||
matching_rule.expiration.days = 7
|
||||
matching_rule.rule_filter.prefix = 'imported_models/'
|
||||
unrelated_rule = MagicMock(rule_id='predictions-offload-expiry')
|
||||
existing.rules = [unrelated_rule, matching_rule]
|
||||
client.get_bucket_lifecycle.return_value = existing
|
||||
|
||||
result = await import_activity.ensure_import_bundle_retention(base())
|
||||
|
||||
client.set_bucket_lifecycle.assert_not_called()
|
||||
assert result['rules'] == 2
|
||||
|
||||
|
||||
def test_the_model_document_shape_is_documented_in_one_place():
|
||||
"""A guard against the shape drifting between the code and what the runtime reads."""
|
||||
source = Path('laborious/activities/model_import.py').read_text()
|
||||
|
||||
for key in ('transform_flavor', 'predict_flavor', 'retention_minutes', 'target'):
|
||||
assert f"'{key}'" in source
|
||||
assert json.dumps({'transform_flavor': 'sklearn'}) # sanity: the flavors are plain strings
|
||||
1046
tests/laborious/activities/test_model_import_status.py
Normal file
1046
tests/laborious/activities/test_model_import_status.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -594,6 +594,39 @@ async def test_write_opc_data_success(mock_dataframe, opc):
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.opc.DataFrame')
|
||||
async def test_write_opc_data_propagates_session_bad_and_reconnect_in_progress(mock_dataframe, opc):
|
||||
"""A single server reporting session_bad/reconnect_in_progress must flip the aggregated
|
||||
flags passed to process_confidence, even though nothing raised."""
|
||||
input_data = {
|
||||
**metadata,
|
||||
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
|
||||
'opc_output_config': {
|
||||
'server1': {
|
||||
'prediction_tags': {'tag1': {'data_type': 'float'}},
|
||||
'confidence_tags': {'tag2': {'data_type': 'float'}},
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
opc.manage_output_tags = AsyncMock(
|
||||
return_value=(False, {'tag1': 0.1, 'tag2': 0.2}, True, 'BadSessionClosed', True)
|
||||
)
|
||||
opc.process_confidence = MagicMock(return_value={'data': 'data'})
|
||||
|
||||
await opc.write_opc_data(input_data)
|
||||
|
||||
opc.process_confidence.assert_called_once_with(
|
||||
mock_dataframe.return_value,
|
||||
False,
|
||||
metadata['metadata'],
|
||||
session_bad=True,
|
||||
opc_status='BadSessionClosed',
|
||||
reconnect_in_progress=True,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_write_opc_data_empty_config(opc):
|
||||
# Arrange
|
||||
@@ -734,5 +767,5 @@ async def test_validate_server(opc):
|
||||
@mark.asyncio
|
||||
async def test_close(opc):
|
||||
opc.opc_repository['server1'].disconnect = AsyncMock(return_value=True)
|
||||
await opc.close()
|
||||
await opc.aclose()
|
||||
opc.opc_repository['server1'].disconnect.assert_called_once()
|
||||
|
||||
276
tests/laborious/test_import_absence.py
Normal file
276
tests/laborious/test_import_absence.py
Normal file
@@ -0,0 +1,276 @@
|
||||
"""What the import deliberately does **not** do.
|
||||
|
||||
Every absence here was a decision, and every one of them is the kind of thing a later well-meaning
|
||||
edit adds back: a rollback for a failed import, a resume for a half-finished one, a signal to retry,
|
||||
a `joblib.load` to "check the model is valid", a reason string built from the exception. None of
|
||||
those are missing features — a failed import stays failed, and trying again is a new upload — so the
|
||||
tests are written against the source itself, because there is no behaviour to observe when the
|
||||
behaviour is the point of not existing.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
|
||||
from pytest import mark
|
||||
|
||||
from laborious.workflows.import_model import ImportModel
|
||||
|
||||
REPOSITORY_ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
WORKFLOW_MODULE = REPOSITORY_ROOT / 'laborious/workflows/import_model.py'
|
||||
ACTIVITIES_MODULE = REPOSITORY_ROOT / 'laborious/activities/model_import.py'
|
||||
ERRORS_MODULE = REPOSITORY_ROOT / 'laborious/utils/models/import_errors.py'
|
||||
READER_MODULE = REPOSITORY_ROOT / 'laborious/utils/bundle/reader.py'
|
||||
FORMAT_MODULE = REPOSITORY_ROOT / 'laborious/utils/bundle/format.py'
|
||||
|
||||
IMPORT_PATH_MODULES = (
|
||||
WORKFLOW_MODULE,
|
||||
ACTIVITIES_MODULE,
|
||||
ERRORS_MODULE,
|
||||
READER_MODULE,
|
||||
FORMAT_MODULE,
|
||||
)
|
||||
|
||||
# The vocabulary of recovery. A partial import is not repaired; it is reported and abandoned.
|
||||
RECOVERY_WORDS = ('rollback', 'compensate', 'resume', 'reconcile', 'undo')
|
||||
|
||||
# Loading a model is the *execution* side, owned by another card. The importer moves bytes and
|
||||
# writes metadata; it never unpickles what it just accepted from a user.
|
||||
MODEL_LOADING_CALLS = (
|
||||
'mlflow.pyfunc.load_model',
|
||||
'mlflow.sklearn.load_model',
|
||||
'joblib.load',
|
||||
'pickle.load',
|
||||
'pickle.loads',
|
||||
)
|
||||
|
||||
# Decryption lives in the activity that uses the password, never in the workflow, which would put
|
||||
# the plaintext into the event history.
|
||||
DECRYPTION_PRIMITIVES = (
|
||||
'AESGCM',
|
||||
'crypto_pwhash',
|
||||
'crypto_secretstream',
|
||||
'decrypt_password_envelope',
|
||||
'load_envelope_key',
|
||||
'IMPORT_PASSWORD_KEY',
|
||||
'Argon2',
|
||||
)
|
||||
|
||||
|
||||
def parse(path: Path) -> ast.Module:
|
||||
return ast.parse(path.read_text())
|
||||
|
||||
|
||||
def called_names(tree: ast.Module) -> list[str]:
|
||||
"""Every called name, as written: `a.b.c(...)` becomes `a.b.c`."""
|
||||
names = []
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.Call):
|
||||
continue
|
||||
names.append(ast.unparse(node.func))
|
||||
return names
|
||||
|
||||
|
||||
def defined_names(tree: ast.Module) -> list[str]:
|
||||
return [
|
||||
node.name
|
||||
for node in ast.walk(tree)
|
||||
if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef)
|
||||
]
|
||||
|
||||
|
||||
# ------------------------------------------------------------- no retry beyond the record
|
||||
|
||||
|
||||
def test_the_workflow_holds_no_recovery_routine():
|
||||
tree = parse(WORKFLOW_MODULE)
|
||||
|
||||
for name in defined_names(tree) + called_names(tree):
|
||||
for word in RECOVERY_WORDS:
|
||||
assert word not in name.lower(), f'{name} looks like recovery'
|
||||
|
||||
|
||||
def test_the_activities_hold_no_recovery_routine():
|
||||
tree = parse(ACTIVITIES_MODULE)
|
||||
|
||||
for name in defined_names(tree) + called_names(tree):
|
||||
for word in RECOVERY_WORDS:
|
||||
assert word not in name.lower(), f'{name} looks like recovery'
|
||||
|
||||
|
||||
def test_the_workflow_never_continues_as_new():
|
||||
"""A single import is one workflow execution, start to finish."""
|
||||
assert not [name for name in called_names(parse(WORKFLOW_MODULE)) if 'continue_as_new' in name]
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ no handlers, no waiting
|
||||
|
||||
|
||||
def test_the_workflow_declares_no_signal_or_query_handler():
|
||||
"""Nothing outside can steer an import once it starts: no retry signal, no progress query."""
|
||||
for name in dir(ImportModel):
|
||||
member = getattr(ImportModel, name, None)
|
||||
for marker in (
|
||||
'__temporal_signal_definition',
|
||||
'__temporal_query_definition',
|
||||
'__temporal_update_definition',
|
||||
):
|
||||
assert not hasattr(member, marker), f'{name} declares {marker}'
|
||||
|
||||
|
||||
def test_the_workflow_module_carries_no_handler_decorator():
|
||||
tree = parse(WORKFLOW_MODULE)
|
||||
decorators = [
|
||||
ast.unparse(decorator)
|
||||
for node in ast.walk(tree)
|
||||
if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef)
|
||||
for decorator in node.decorator_list
|
||||
]
|
||||
|
||||
assert set(decorators) == {'workflow.run', 'staticmethod'}
|
||||
assert decorators.count('workflow.run') == 1
|
||||
|
||||
|
||||
def test_the_workflow_never_waits_for_a_condition():
|
||||
"""No `wait_condition`: there is no external event an import waits on."""
|
||||
assert not [name for name in called_names(parse(WORKFLOW_MODULE)) if 'wait_condition' in name]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------- nothing is loaded
|
||||
|
||||
|
||||
@mark.parametrize('module_path', IMPORT_PATH_MODULES, ids=lambda path: path.name)
|
||||
def test_no_module_in_the_import_path_loads_a_model(module_path: Path):
|
||||
source = module_path.read_text()
|
||||
|
||||
for call in MODEL_LOADING_CALLS:
|
||||
assert call not in source
|
||||
|
||||
|
||||
def test_the_reader_knows_nothing_about_mlflow():
|
||||
"""It stays liftable into a shared package: the reader is about the file, not the platform."""
|
||||
source = READER_MODULE.read_text()
|
||||
|
||||
assert 'mlflow' not in source
|
||||
assert 'temporalio' not in source
|
||||
|
||||
|
||||
# ----------------------------------------------------------- the workflow holds no secret
|
||||
|
||||
|
||||
@mark.parametrize('primitive', DECRYPTION_PRIMITIVES)
|
||||
def test_the_workflow_module_holds_no_decryption_primitive(primitive: str):
|
||||
assert primitive not in WORKFLOW_MODULE.read_text()
|
||||
|
||||
|
||||
def test_the_workflow_only_ever_holds_the_envelope():
|
||||
"""The envelope crosses the workflow; the password it wraps never does."""
|
||||
source = WORKFLOW_MODULE.read_text()
|
||||
|
||||
assert 'password_envelope' in source
|
||||
# The only other mention of a password is the input field the workflow refuses to accept.
|
||||
assert source.count("'password'") == 1
|
||||
|
||||
|
||||
def test_the_plaintext_password_field_is_refused_by_name():
|
||||
from laborious.workflows.import_model import REJECTED_INPUT_FIELDS
|
||||
|
||||
assert 'password' in REJECTED_INPUT_FIELDS
|
||||
|
||||
|
||||
# ------------------------------------------------------- every reason comes from the catalog
|
||||
|
||||
|
||||
def _error_message_values(tree: ast.Module) -> list[ast.expr]:
|
||||
"""Every expression written into the record's `error_message`, at an actual status write.
|
||||
|
||||
Scoped to `_execute_status_write`'s own parameter dict — not every dict in the module — because
|
||||
`_build_update`'s `columns` argument also carries an `'error_message'` key (naming the column and
|
||||
its bind parameter, not a value being written), and that is not a write this test is about.
|
||||
"""
|
||||
values = []
|
||||
for node in ast.walk(tree):
|
||||
if not (
|
||||
isinstance(node, ast.Call) and ast.unparse(node.func).endswith('_execute_status_write')
|
||||
):
|
||||
continue
|
||||
for arg in node.args:
|
||||
if not isinstance(arg, ast.Dict):
|
||||
continue
|
||||
for key, value in zip(arg.keys, arg.values, strict=True):
|
||||
if isinstance(key, ast.Constant) and key.value == 'error_message':
|
||||
values.append(value)
|
||||
return values
|
||||
|
||||
|
||||
def test_no_recorded_reason_is_built_from_an_exception():
|
||||
"""The catalog is the only source of a sentence a user reads."""
|
||||
values = _error_message_values(parse(ACTIVITIES_MODULE))
|
||||
|
||||
assert values, 'no error_message write found — has the terminal statement moved?'
|
||||
for value in values:
|
||||
assert isinstance(value, ast.Name | ast.Constant), ast.unparse(value)
|
||||
assert not isinstance(value, ast.JoinedStr | ast.BinOp)
|
||||
|
||||
|
||||
def test_the_reason_written_is_the_classified_one():
|
||||
values = _error_message_values(parse(ACTIVITIES_MODULE))
|
||||
|
||||
assert {ast.unparse(value) for value in values} == {'reason'}
|
||||
|
||||
|
||||
def test_the_classification_ignores_the_exception_by_construction():
|
||||
"""`classify_import_failure` has no exception parameter at all; nothing can leak from it."""
|
||||
from laborious.utils.models.import_errors import classify_import_failure
|
||||
|
||||
parameters = inspect.signature(classify_import_failure).parameters
|
||||
assert 'exception' not in parameters
|
||||
assert list(parameters) == ['step']
|
||||
|
||||
|
||||
@mark.parametrize('module_path', (WORKFLOW_MODULE, ACTIVITIES_MODULE), ids=lambda p: p.name)
|
||||
def test_no_status_detail_interpolates_an_exception(module_path: Path):
|
||||
"""The JSON detail is built from a code and a step, both of them enum values."""
|
||||
tree = parse(module_path)
|
||||
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.Call) or ast.unparse(node.func) != 'build_import_detail':
|
||||
continue
|
||||
for keyword in node.keywords:
|
||||
rendered = ast.unparse(keyword.value)
|
||||
assert 'error' not in rendered
|
||||
assert 'exception' not in rendered
|
||||
assert not isinstance(keyword.value, ast.JoinedStr)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ no relational target
|
||||
|
||||
|
||||
def test_the_import_writes_to_one_table_and_names_it_once():
|
||||
"""D22: no project table, no model table, no second database. One `UPDATE`, one relation."""
|
||||
statements = []
|
||||
for node in ast.walk(parse(ACTIVITIES_MODULE)):
|
||||
if not isinstance(node, ast.Call) or ast.unparse(node.func) != 'text':
|
||||
continue
|
||||
argument = node.args[0]
|
||||
if isinstance(argument, ast.Constant):
|
||||
statements.append(argument.value)
|
||||
|
||||
assert statements, 'no SQL statement found in the import activities'
|
||||
for statement in statements:
|
||||
words = statement.split()
|
||||
relations = {
|
||||
words[index + 1]
|
||||
for index, word in enumerate(words)
|
||||
if word.upper() in ('UPDATE', 'FROM', 'INTO', 'JOIN')
|
||||
}
|
||||
assert relations == {'public.experiment_run'}
|
||||
|
||||
|
||||
def test_the_import_never_inserts_a_relational_row():
|
||||
source = ACTIVITIES_MODULE.read_text().upper()
|
||||
|
||||
assert 'INSERT INTO' not in source
|
||||
assert 'DELETE FROM' not in source
|
||||
assert 'CREATE TABLE' not in source
|
||||
0
tests/laborious/utils/bundle/__init__.py
Normal file
0
tests/laborious/utils/bundle/__init__.py
Normal file
34
tests/laborious/utils/bundle/conftest.py
Normal file
34
tests/laborious/utils/bundle/conftest.py
Normal file
@@ -0,0 +1,34 @@
|
||||
"""Shared fixtures for the `.sientia` bundle tests."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
FIXTURE_DIR = Path(__file__).resolve().parents[3] / 'fixtures' / 'bundle'
|
||||
GOLDEN_BUNDLE = FIXTURE_DIR / 'golden_model_v3.sientia'
|
||||
GOLDEN_METADATA = FIXTURE_DIR / 'golden_model_v3.metadata.json'
|
||||
# Recorded in tests/fixtures/bundle/README.md; the fixture protects nothing.
|
||||
GOLDEN_PASSWORD = 'Import1Model!Fixture'
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def golden_bundle() -> Path:
|
||||
"""The `.sientia` file the real producer emitted."""
|
||||
assert GOLDEN_BUNDLE.is_file(), (
|
||||
f'golden fixture missing: {GOLDEN_BUNDLE} — regenerate it with '
|
||||
'tests/fixtures/bundle/regenerate.py'
|
||||
)
|
||||
return GOLDEN_BUNDLE
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def golden_password() -> str:
|
||||
return GOLDEN_PASSWORD
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def work_dir(tmp_path: Path) -> Path:
|
||||
"""A per-test stand-in for `IMPORT_WORK_DIR/<workflow_id>`."""
|
||||
root = tmp_path / 'workflow-id'
|
||||
root.mkdir()
|
||||
return root
|
||||
227
tests/laborious/utils/bundle/test_format.py
Normal file
227
tests/laborious/utils/bundle/test_format.py
Normal file
@@ -0,0 +1,227 @@
|
||||
"""The wire format: header parsing, ceilings, digest coverage and the fixture contract."""
|
||||
|
||||
import ast
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
|
||||
import nacl.bindings as sodium
|
||||
import pytest
|
||||
|
||||
from laborious.utils.bundle import format as fmt
|
||||
from tests.helpers import bundle_factory as factory
|
||||
|
||||
|
||||
def _write(tmp_path: Path, payload: bytes, name: str = 'bundle.sientia') -> Path:
|
||||
target = tmp_path / name
|
||||
target.write_bytes(payload)
|
||||
return target
|
||||
|
||||
|
||||
# --------------------------------------------------------------- constants
|
||||
|
||||
|
||||
def test_header_size_is_140():
|
||||
assert fmt.HEADER_SIZE == 140
|
||||
assert fmt.HEADER_STRUCT.format == '>8sBBB16sQQB32s64s'
|
||||
|
||||
|
||||
def test_libsodium_literals_match_the_binding():
|
||||
"""The literals in `format.py` keep the module import-free; these assertions keep them true."""
|
||||
assert fmt.STREAM_HEADER_SIZE == sodium.crypto_secretstream_xchacha20poly1305_HEADERBYTES
|
||||
assert fmt.ABYTES == sodium.crypto_secretstream_xchacha20poly1305_ABYTES
|
||||
assert fmt.KEY_SIZE == sodium.crypto_secretstream_xchacha20poly1305_KEYBYTES
|
||||
|
||||
|
||||
def test_producer_constants_match():
|
||||
assert fmt.MAGIC == b'SIENTIA1'
|
||||
assert fmt.CHUNK_SIZE == 65536
|
||||
assert fmt.MAX_KDF_MEMLIMIT == 268435456
|
||||
assert fmt.MAX_KDF_OPSLIMIT == 3
|
||||
assert fmt.SALT_SIZE == 16
|
||||
assert fmt.DIGEST_SIZE == 32
|
||||
assert fmt.RESERVED_SIGNATURE_SIZE == 64
|
||||
|
||||
|
||||
def test_prefix_size_is_65717():
|
||||
"""The pre-upload prefix is computed from constants, never read from a file."""
|
||||
assert fmt.PREFIX_SIZE == 140 + 24 + (65536 + 17)
|
||||
assert fmt.PREFIX_SIZE == 65717
|
||||
|
||||
|
||||
def test_digest_offset_skips_both_headers():
|
||||
assert fmt.DIGEST_OFFSET == 164
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- parse
|
||||
|
||||
|
||||
def test_parses_a_producer_emitted_header(golden_bundle):
|
||||
header = fmt.read_header(golden_bundle)
|
||||
|
||||
assert header.format_version == 1
|
||||
assert header.aead_id == fmt.AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM
|
||||
assert header.kdf_id == fmt.KDF_ID_ARGON2ID
|
||||
assert len(header.salt) == 16
|
||||
assert header.kdf_memlimit_bytes == 268435456
|
||||
assert header.kdf_opslimit == 3
|
||||
assert header.kdf_parallelism == 1
|
||||
assert len(header.ciphertext_digest) == 32
|
||||
assert header.reserved_signature == bytes(64)
|
||||
|
||||
|
||||
def test_wrong_magic_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(magic=b'NOTSIENT'))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='expected bundle marker'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_unknown_format_version_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(format_version=2))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='format version'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_unknown_aead_id_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(aead_id=7))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='encryption identifier'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_unknown_kdf_id_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(kdf_id=9))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='key derivation identifier'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_memlimit_above_the_ceiling_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(declared_memlimit=512 * 1024 * 1024))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='memory'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_opslimit_above_the_ceiling_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(declared_opslimit=4))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='passes'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_zero_cost_parameters_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(declared_opslimit=0))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='positive'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_reserved_parallelism_must_be_one(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(parallelism=2))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='parallelism'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_non_zero_reserved_signature_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(reserved_signature=b'\x01' + bytes(63)))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='reserved signature'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_file_shorter_than_header_plus_stream_header_rejected(tmp_path):
|
||||
path = _write(tmp_path, b'SIENTIA1' + bytes(100))
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='smaller than'):
|
||||
fmt.read_header(path)
|
||||
|
||||
|
||||
def test_parse_header_needs_the_whole_header():
|
||||
with pytest.raises(fmt.BundleFormatError, match='shorter than'):
|
||||
fmt.parse_header(b'SIENTIA1')
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ digest
|
||||
|
||||
|
||||
def test_digest_covers_the_chunk_stream_not_the_stream_header(golden_bundle):
|
||||
"""The producer hashes the chunks before prepending either header — offset 164, not 140."""
|
||||
header = fmt.read_header(golden_bundle)
|
||||
payload = golden_bundle.read_bytes()
|
||||
|
||||
assert hashlib.sha256(payload[164:]).digest() == header.ciphertext_digest
|
||||
assert hashlib.sha256(payload[140:]).digest() != header.ciphertext_digest
|
||||
assert fmt.compute_ciphertext_digest(golden_bundle) == header.ciphertext_digest
|
||||
|
||||
|
||||
def test_verify_ciphertext_digest_accepts_the_fixture(golden_bundle):
|
||||
fmt.verify_ciphertext_digest(golden_bundle, fmt.read_header(golden_bundle))
|
||||
|
||||
|
||||
def test_digest_mismatch_rejected(tmp_path):
|
||||
path = _write(tmp_path, factory.make_bundle(digest_override=bytes(32)))
|
||||
header = fmt.read_header(path)
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='does not match the digest'):
|
||||
fmt.verify_ciphertext_digest(path, header)
|
||||
|
||||
|
||||
def test_truncated_file_fails_the_digest(tmp_path):
|
||||
path = _write(tmp_path, factory.truncate_bundle(factory.make_bundle()))
|
||||
header = fmt.read_header(path)
|
||||
|
||||
with pytest.raises(fmt.BundleFormatError, match='does not match the digest'):
|
||||
fmt.verify_ciphertext_digest(path, header)
|
||||
|
||||
|
||||
def test_file_digest_is_the_whole_object(tmp_path):
|
||||
payload = factory.make_bundle()
|
||||
path = _write(tmp_path, payload)
|
||||
|
||||
assert fmt.file_digest(path) == hashlib.sha256(payload).hexdigest()
|
||||
|
||||
|
||||
# ------------------------------------------------------- factory conformance
|
||||
|
||||
|
||||
def test_factory_header_matches_fixture_layout(golden_bundle, tmp_path):
|
||||
"""The test-only writer must be structurally identical to the producer's output."""
|
||||
produced = fmt.read_header(_write(tmp_path, factory.make_bundle()))
|
||||
golden = fmt.read_header(golden_bundle)
|
||||
|
||||
assert (produced.format_version, produced.aead_id, produced.kdf_id) == (
|
||||
golden.format_version,
|
||||
golden.aead_id,
|
||||
golden.kdf_id,
|
||||
)
|
||||
assert produced.kdf_parallelism == golden.kdf_parallelism
|
||||
assert produced.reserved_signature == golden.reserved_signature
|
||||
assert len(produced.salt) == len(golden.salt)
|
||||
assert len(produced.ciphertext_digest) == len(golden.ciphertext_digest)
|
||||
assert factory.HEADER_STRUCT.format == fmt.HEADER_STRUCT.format
|
||||
assert factory.PRODUCER_KDF_MEMLIMIT_BYTES == golden.kdf_memlimit_bytes
|
||||
assert factory.PRODUCER_KDF_OPSLIMIT == golden.kdf_opslimit
|
||||
|
||||
|
||||
# --------------------------------------------------------------- liftability
|
||||
|
||||
FORBIDDEN_IMPORT_ROOTS = frozenset(
|
||||
{'laborious', 'temporalio', 'mlflow', 'nacl', 'cryptography', 'sqlalchemy', 'pymongo'}
|
||||
)
|
||||
|
||||
|
||||
def test_format_module_imports_nothing_platform_specific():
|
||||
"""`format.py` stays liftable into a shared package: stdlib only, no platform imports."""
|
||||
tree = ast.parse(Path(fmt.__file__).read_text(encoding='utf-8'))
|
||||
roots = set()
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Import):
|
||||
roots.update(alias.name.split('.')[0] for alias in node.names)
|
||||
elif isinstance(node, ast.ImportFrom) and node.module:
|
||||
roots.add(node.module.split('.')[0])
|
||||
|
||||
assert roots.isdisjoint(FORBIDDEN_IMPORT_ROOTS), roots & FORBIDDEN_IMPORT_ROOTS
|
||||
assert roots == {'__future__', 'hashlib', 'struct', 'dataclasses', 'pathlib', 'typing'}
|
||||
691
tests/laborious/utils/bundle/test_reader.py
Normal file
691
tests/laborious/utils/bundle/test_reader.py
Normal file
@@ -0,0 +1,691 @@
|
||||
"""The seven gates: what each one rejects, in what order, and what it leaves on disk."""
|
||||
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import stat
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
import nacl.bindings as sodium
|
||||
import pytest
|
||||
|
||||
from laborious.utils.bundle import format as fmt
|
||||
from laborious.utils.bundle.reader import BundleGateError, BundleLimits, BundleReader
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from tests.helpers import bundle_factory as factory
|
||||
|
||||
OBJECT_KEY = 'imported_models/factory_model_v2.sientia'
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def reader() -> BundleReader:
|
||||
return BundleReader()
|
||||
|
||||
|
||||
def _bundle_file(work_dir: Path, payload: bytes) -> Path:
|
||||
target = work_dir / 'download.sientia'
|
||||
target.write_bytes(payload)
|
||||
return target
|
||||
|
||||
|
||||
def _open(reader: BundleReader, work_dir: Path, payload: bytes, password: str | None = None):
|
||||
return reader.open(
|
||||
_bundle_file(work_dir, payload),
|
||||
password if password is not None else factory.DEFAULT_PASSWORD,
|
||||
work_dir / 'extract',
|
||||
work_dir,
|
||||
)
|
||||
|
||||
|
||||
def _extracted_files(root: Path) -> list[str]:
|
||||
if not root.exists():
|
||||
return []
|
||||
return sorted(str(p.relative_to(root)) for p in root.rglob('*') if p.is_file())
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 1
|
||||
|
||||
|
||||
def test_gate_1_accepts_a_well_shaped_object(reader):
|
||||
reader.check_object_limits(size=1024, key=OBJECT_KEY)
|
||||
|
||||
|
||||
def test_gate_1_rejects_an_oversized_object(reader):
|
||||
limited = BundleReader(BundleLimits(max_object_bytes=1000))
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
limited.check_object_limits(size=1001, key=OBJECT_KEY)
|
||||
|
||||
assert error.value.gate == 1
|
||||
assert error.value.step is ImportStep.DOWNLOAD
|
||||
|
||||
|
||||
def test_gate_1_rejects_an_empty_object(reader):
|
||||
with pytest.raises(BundleGateError, match='empty'):
|
||||
reader.check_object_limits(size=0, key=OBJECT_KEY)
|
||||
|
||||
|
||||
def test_gate_1_rejects_a_key_outside_the_import_prefix(reader):
|
||||
with pytest.raises(BundleGateError, match='import prefix'):
|
||||
reader.check_object_limits(size=10, key='other_place/model.sientia')
|
||||
|
||||
|
||||
def test_gate_1_rejects_a_key_with_the_wrong_suffix(reader):
|
||||
with pytest.raises(BundleGateError, match='.sientia'):
|
||||
reader.check_object_limits(size=10, key='imported_models/model.zip')
|
||||
|
||||
|
||||
# ------------------------------------------------------------- golden bundle
|
||||
|
||||
|
||||
def test_opens_the_producer_fixture(golden_bundle, golden_password, work_dir, reader):
|
||||
"""The producer's own bytes, at the producer's own KDF cost."""
|
||||
opened = reader.open(golden_bundle, golden_password, work_dir / 'extract', work_dir)
|
||||
|
||||
assert opened.metadata['model_name'] == 'fixture_model'
|
||||
assert opened.metadata['experiment_name'] == 'fixture_experiment'
|
||||
assert set(opened.metadata) == {
|
||||
'model_name',
|
||||
'experiment_name',
|
||||
'run_id',
|
||||
'model_version',
|
||||
'model_project',
|
||||
'export_timestamp',
|
||||
}
|
||||
assert opened.parameters['target_variable'] == 'SE'
|
||||
assert opened.header is not None
|
||||
assert opened.digest == opened.header.ciphertext_digest.hex()
|
||||
assert _extracted_files(opened.extracted_dir) == [
|
||||
'artifacts/data_model/MLmodel',
|
||||
'artifacts/data_model/conda.yaml',
|
||||
'artifacts/data_model/model.pkl',
|
||||
'artifacts/data_model/python_env.yaml',
|
||||
'artifacts/data_model/requirements.txt',
|
||||
'artifacts/model_card/model_card.json',
|
||||
'artifacts/prediction_model/MLmodel',
|
||||
'artifacts/prediction_model/conda.yaml',
|
||||
'artifacts/prediction_model/model.pkl',
|
||||
'artifacts/prediction_model/python_env.yaml',
|
||||
'artifacts/prediction_model/requirements.txt',
|
||||
'artifacts/raw_data.csv',
|
||||
'artifacts/test_data.csv',
|
||||
'artifacts/train_data.csv',
|
||||
'metadata.json',
|
||||
]
|
||||
assert list(work_dir.glob('*.zip')) == []
|
||||
|
||||
|
||||
def test_fixture_metadata_matches_its_committed_plaintext_copy(
|
||||
golden_bundle, golden_password, work_dir, reader
|
||||
):
|
||||
from tests.laborious.utils.bundle.conftest import GOLDEN_METADATA
|
||||
|
||||
opened = reader.open(golden_bundle, golden_password, work_dir / 'extract', work_dir)
|
||||
committed = json.loads(GOLDEN_METADATA.read_text(encoding='utf-8'))
|
||||
|
||||
assert opened.metadata == committed['metadata']
|
||||
assert opened.parameters == committed['parameters']
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 2
|
||||
|
||||
|
||||
def test_gate_2_rejects_a_wrong_magic_before_deriving(work_dir, reader, monkeypatch):
|
||||
payload = factory.make_bundle(magic=b'NOTSIENT')
|
||||
monkeypatch.setattr(
|
||||
sodium, 'crypto_pwhash_alg', _never_called('key derivation must not run at gate 2')
|
||||
)
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 2
|
||||
assert error.value.step is ImportStep.DOWNLOAD
|
||||
|
||||
|
||||
def test_gate_2_rejects_kdf_parameters_above_the_ceiling_without_allocating(
|
||||
work_dir, reader, monkeypatch
|
||||
):
|
||||
payload = factory.make_bundle(declared_memlimit=512 * 1024 * 1024)
|
||||
monkeypatch.setattr(
|
||||
sodium, 'crypto_pwhash_alg', _never_called('the ceiling must be enforced before deriving')
|
||||
)
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 2
|
||||
|
||||
|
||||
def test_gate_2_rejects_a_digest_mismatch(work_dir, reader):
|
||||
payload = factory.make_bundle(digest_override=bytes(32))
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 2
|
||||
|
||||
|
||||
def test_gate_2_rejects_a_truncated_transfer(work_dir, reader):
|
||||
payload = factory.truncate_bundle(factory.make_bundle())
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 2
|
||||
|
||||
|
||||
def test_gate_2_wraps_an_os_error_reading_the_file(work_dir, reader):
|
||||
"""A download that never landed on disk is an OS error, not a format one."""
|
||||
missing = work_dir / 'never-downloaded.sientia'
|
||||
|
||||
with pytest.raises(BundleGateError, match='could not be read') as error:
|
||||
reader.open(missing, factory.DEFAULT_PASSWORD, work_dir / 'extract', work_dir)
|
||||
|
||||
assert error.value.gate == 2
|
||||
|
||||
|
||||
def _never_called(message: str):
|
||||
def _raise(*_args, **_kwargs):
|
||||
raise AssertionError(message)
|
||||
|
||||
return _raise
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 3
|
||||
|
||||
AMBIGUOUS_REASON = 'wrong password, or the file is damaged or was altered'
|
||||
|
||||
|
||||
def test_gate_3_wrong_password(work_dir, reader):
|
||||
payload = factory.make_bundle()
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, payload, password='Wrong1Password!Here')
|
||||
|
||||
assert error.value.gate == 3
|
||||
assert error.value.step is ImportStep.DECRYPTION
|
||||
assert AMBIGUOUS_REASON in str(error.value)
|
||||
|
||||
|
||||
def test_gate_3_altered_ciphertext_reads_the_same_as_a_wrong_password(work_dir, reader):
|
||||
tampered = factory.tamper_ciphertext(factory.make_bundle())
|
||||
|
||||
with pytest.raises(BundleGateError) as altered:
|
||||
_open(reader, work_dir, tampered)
|
||||
|
||||
wrong_password_payload = factory.make_bundle()
|
||||
with pytest.raises(BundleGateError) as wrong_password:
|
||||
_open(reader, work_dir, wrong_password_payload, password='Wrong1Password!Here')
|
||||
|
||||
assert str(altered.value) == str(wrong_password.value)
|
||||
assert altered.value.gate == wrong_password.value.gate == 3
|
||||
|
||||
|
||||
def test_gate_3_stream_without_a_final_tag_is_rejected(work_dir, reader):
|
||||
payload = factory.make_bundle(drop_final_tag=True)
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 3
|
||||
assert AMBIGUOUS_REASON in str(error.value)
|
||||
|
||||
|
||||
def test_gate_3_wraps_an_os_error_while_writing_the_decrypted_zip(work_dir, reader, monkeypatch):
|
||||
"""A write failure while streaming the plaintext to disk is reported as disk trouble, not as
|
||||
a bad password — the AEAD never even got to disagree."""
|
||||
payload = factory.make_bundle()
|
||||
|
||||
def _raise_os_error(*_args, **_kwargs):
|
||||
raise OSError('disk full')
|
||||
|
||||
monkeypatch.setattr(reader, '_pull_stream', _raise_os_error)
|
||||
|
||||
with pytest.raises(BundleGateError, match='could not be written') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 3
|
||||
|
||||
|
||||
def test_gate_3_leaves_no_decrypted_zip_behind(work_dir, reader):
|
||||
payload = factory.make_bundle()
|
||||
|
||||
with pytest.raises(BundleGateError):
|
||||
_open(reader, work_dir, payload, password='Wrong1Password!Here')
|
||||
|
||||
assert list(work_dir.glob('*.zip')) == []
|
||||
|
||||
|
||||
def test_gate_3_normalises_the_password_to_nfc(work_dir, reader):
|
||||
composed = 'Café1Model!Test'
|
||||
decomposed = 'Café1Model!Test'
|
||||
payload = factory.make_bundle(password=composed)
|
||||
|
||||
opened = _open(reader, work_dir, payload, password=decomposed)
|
||||
|
||||
assert opened.metadata['model_name'] == 'factory_model'
|
||||
|
||||
|
||||
def test_gate_3_does_not_trim_the_password(work_dir, reader):
|
||||
payload = factory.make_bundle(password='Trailing1Space!Test ')
|
||||
|
||||
with pytest.raises(BundleGateError):
|
||||
_open(reader, work_dir, payload, password='Trailing1Space!Test')
|
||||
|
||||
|
||||
def test_multi_chunk_plaintext_round_trips(work_dir, reader):
|
||||
"""More than 64 KiB of plaintext, so `MESSAGE` and `FINAL` framing are both exercised."""
|
||||
spec = factory.ZipSpec()
|
||||
# Incompressible on purpose: a compressible filler would deflate below one chunk and the test
|
||||
# would silently stop covering multi-chunk framing.
|
||||
spec.files['artifacts/prediction_model/model.pkl'] = random.Random( # noqa: S311 - filler
|
||||
20260825
|
||||
).randbytes(200 * 1024)
|
||||
payload = factory.make_bundle(spec)
|
||||
assert len(payload) > fmt.CHUNK_SIZE + fmt.ABYTES + fmt.DIGEST_OFFSET
|
||||
|
||||
opened = _open(reader, work_dir, payload)
|
||||
|
||||
extracted = opened.extracted_dir / 'artifacts' / 'prediction_model' / 'model.pkl'
|
||||
assert extracted.read_bytes() == spec.files['artifacts/prediction_model/model.pkl']
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 4
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('spec_kwargs', 'match'),
|
||||
[
|
||||
({'bomb_bytes': 40 * 1024 * 1024}, 'compression ratio'),
|
||||
({'traversal_entry': 'artifacts/../../etc/passwd'}, 'escaping'),
|
||||
({'absolute_entry': '/etc/passwd'}, 'absolute'),
|
||||
({'symlinks': {'artifacts/link': '/etc/passwd'}}, 'symbolic link'),
|
||||
],
|
||||
)
|
||||
def test_gate_4_rejects_hostile_central_directories(work_dir, reader, spec_kwargs, match):
|
||||
payload = factory.make_bundle(factory.ZipSpec(**spec_kwargs))
|
||||
|
||||
with pytest.raises(BundleGateError, match=match) as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 4
|
||||
assert error.value.step is ImportStep.ARCHIVE_INSPECTION
|
||||
assert _extracted_files(work_dir / 'extract') == []
|
||||
assert not (work_dir / 'extract').exists()
|
||||
|
||||
|
||||
def test_gate_4_rejects_too_many_entries(work_dir):
|
||||
limited = BundleReader(BundleLimits(max_entries=3))
|
||||
payload = factory.make_bundle()
|
||||
|
||||
with pytest.raises(BundleGateError, match='entries') as error:
|
||||
_open(limited, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 4
|
||||
|
||||
|
||||
def test_gate_4_rejects_an_oversized_uncompressed_total(work_dir):
|
||||
limited = BundleReader(BundleLimits(max_uncompressed_bytes=64))
|
||||
payload = factory.make_bundle()
|
||||
|
||||
with pytest.raises(BundleGateError, match='uncompressed') as error:
|
||||
_open(limited, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 4
|
||||
|
||||
|
||||
def test_gate_4_rejects_content_that_is_not_an_archive(work_dir, reader):
|
||||
payload = factory.encrypt_archive(b'not a zip at all')
|
||||
|
||||
with pytest.raises(BundleGateError, match='readable archive') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 4
|
||||
|
||||
|
||||
def test_gate_4_rejects_an_empty_archive(work_dir, reader):
|
||||
payload = factory.make_bundle(factory.ZipSpec(files={}, include_metadata=False))
|
||||
|
||||
with pytest.raises(BundleGateError, match='archive is empty') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 4
|
||||
|
||||
|
||||
def test_gate_4_rejects_an_entry_that_is_neither_a_file_nor_a_directory(work_dir, reader):
|
||||
"""A FIFO's mode bits are neither a symlink's, a regular file's, nor a directory's."""
|
||||
payload = factory.make_bundle(
|
||||
factory.ZipSpec(special_mode_entries={'artifacts/pipe': stat.S_IFIFO | 0o644})
|
||||
)
|
||||
|
||||
with pytest.raises(BundleGateError, match='not a regular file') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 4
|
||||
|
||||
|
||||
def test_gate_4_and_gate_5_read_the_central_directory_once(
|
||||
golden_bundle, golden_password, work_dir, reader, monkeypatch
|
||||
):
|
||||
"""Achado #11: gate 5 extracts using the entries gate 4 already read, instead of reopening the
|
||||
archive to call `infolist()` a second time."""
|
||||
original_infolist = zipfile.ZipFile.infolist
|
||||
calls: list[None] = []
|
||||
|
||||
def _counting_infolist(self):
|
||||
calls.append(None)
|
||||
return original_infolist(self)
|
||||
|
||||
monkeypatch.setattr(zipfile.ZipFile, 'infolist', _counting_infolist)
|
||||
|
||||
reader.open(golden_bundle, golden_password, work_dir / 'extract', work_dir)
|
||||
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 5
|
||||
|
||||
|
||||
def test_gate_5_refuses_a_pre_existing_extraction_root(work_dir, reader):
|
||||
(work_dir / 'extract').mkdir()
|
||||
payload = factory.make_bundle()
|
||||
|
||||
with pytest.raises(BundleGateError, match='already exists') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 5
|
||||
assert error.value.step is ImportStep.EXTRACTION
|
||||
|
||||
|
||||
def test_gate_5_extraction_writes_plain_files(work_dir, reader):
|
||||
opened = _open(reader, work_dir, factory.make_bundle())
|
||||
|
||||
for path in opened.extracted_dir.rglob('*'):
|
||||
if path.is_file():
|
||||
assert path.stat().st_mode & 0o111 == 0, path
|
||||
|
||||
|
||||
def test_gate_5_extraction_creates_declared_directory_entries(work_dir, reader):
|
||||
"""The producer never writes directory entries, but a bundle from anywhere else might."""
|
||||
payload = factory.make_bundle(factory.ZipSpec(directory_entries=['artifacts/empty_dir']))
|
||||
|
||||
opened = _open(reader, work_dir, payload)
|
||||
|
||||
assert (opened.extracted_dir / 'artifacts' / 'empty_dir').is_dir()
|
||||
|
||||
|
||||
def test_gate_5_rejects_a_duplicate_entry_name(work_dir, reader):
|
||||
"""Two archive entries claiming the same path: gate 4 checks each name on its own and cannot
|
||||
see the collision, so extraction is where an overwrite is caught — and rejected."""
|
||||
payload = factory.make_bundle(factory.ZipSpec(repeat_file='artifacts/prediction_model/MLmodel'))
|
||||
|
||||
with pytest.raises(BundleGateError, match='overwrite') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 5
|
||||
assert not (work_dir / 'extract').exists()
|
||||
|
||||
|
||||
def test_gate_5_wraps_an_os_error_during_unpacking(work_dir, reader, monkeypatch):
|
||||
"""A write failure partway through extraction (e.g. disk full) is reported as unpacking
|
||||
trouble, not surfaced as the raw OSError."""
|
||||
payload = factory.make_bundle()
|
||||
|
||||
def _raise_os_error(*_args, **_kwargs):
|
||||
raise OSError('disk full')
|
||||
|
||||
monkeypatch.setattr(zipfile.ZipExtFile, 'read', _raise_os_error)
|
||||
|
||||
with pytest.raises(BundleGateError, match='could not be unpacked') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 5
|
||||
assert not (work_dir / 'extract').exists()
|
||||
|
||||
|
||||
def test_extract_member_rechecks_containment_on_the_resolved_path(work_dir, reader):
|
||||
"""Gate 4 already rejects a traversal by name; extraction re-checks on the *resolved* path as
|
||||
a second, independent line of defense."""
|
||||
root = work_dir / 'root'
|
||||
root.mkdir()
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, 'w') as archive:
|
||||
archive.writestr('safe.txt', b'hello')
|
||||
|
||||
with zipfile.ZipFile(io.BytesIO(buffer.getvalue())) as archive:
|
||||
entry = archive.infolist()[0]
|
||||
with pytest.raises(BundleGateError, match='resolves outside') as error:
|
||||
reader._extract_member(archive, entry, root, work_dir / 'somewhere-else')
|
||||
|
||||
assert error.value.gate == 5
|
||||
|
||||
|
||||
def test_extract_member_rejects_a_stream_larger_than_its_declared_size(
|
||||
work_dir, reader, monkeypatch
|
||||
):
|
||||
"""A decompressed stream longer than the entry's declared `file_size` is rejected mid-write,
|
||||
instead of trusting the central-directory metadata gate 4 already checked in aggregate.
|
||||
|
||||
A single crafted DEFLATE chunk can make `ZipExtFile.read()` return far more bytes than
|
||||
`file_size` declares before its own EOF/CRC bookkeeping catches up, so the stand-in here
|
||||
forces exactly that response shape rather than fighting zipfile's own CRC validation."""
|
||||
root = work_dir / 'root'
|
||||
root.mkdir()
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, 'w') as archive:
|
||||
archive.writestr('payload.bin', b'x' * 10)
|
||||
|
||||
call_count = {'n': 0}
|
||||
|
||||
def _oversized_read(self, *_args, **_kwargs):
|
||||
call_count['n'] += 1
|
||||
return b'x' * 1000 if call_count['n'] == 1 else b''
|
||||
|
||||
monkeypatch.setattr(zipfile.ZipExtFile, 'read', _oversized_read)
|
||||
|
||||
resolved_root = root.resolve()
|
||||
with zipfile.ZipFile(io.BytesIO(buffer.getvalue())) as archive:
|
||||
entry = archive.infolist()[0]
|
||||
with pytest.raises(BundleGateError, match='more than its declared') as error:
|
||||
reader._extract_member(archive, entry, root, resolved_root)
|
||||
|
||||
assert error.value.gate == 5
|
||||
|
||||
|
||||
def test_extract_member_writes_a_well_formed_member_in_full(work_dir, reader):
|
||||
"""A member whose decompressed size matches its declared `file_size` is extracted byte for
|
||||
byte, exactly as before the byte-counting check was added."""
|
||||
root = work_dir / 'root'
|
||||
root.mkdir()
|
||||
payload = b'y' * 100
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, 'w') as archive:
|
||||
archive.writestr('payload.bin', payload)
|
||||
|
||||
with zipfile.ZipFile(io.BytesIO(buffer.getvalue())) as archive:
|
||||
entry = archive.infolist()[0]
|
||||
reader._extract_member(archive, entry, root, root.resolve())
|
||||
|
||||
assert (root / 'payload.bin').read_bytes() == payload
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 6
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('spec', 'match'),
|
||||
[
|
||||
(factory.ZipSpec(include_metadata=False), 'metadata document'),
|
||||
(
|
||||
factory.ZipSpec(
|
||||
document=factory.build_metadata_document(drop_top_level_key='parameters')
|
||||
),
|
||||
'parameters and metadata blocks',
|
||||
),
|
||||
(
|
||||
factory.ZipSpec(
|
||||
document=factory.build_metadata_document(drop_metadata_field='experiment_name')
|
||||
),
|
||||
'experiment_name',
|
||||
),
|
||||
(
|
||||
factory.ZipSpec(
|
||||
document=factory.build_metadata_document(drop_parameter='target_variable')
|
||||
),
|
||||
'target_variable',
|
||||
),
|
||||
(
|
||||
factory.ZipSpec(
|
||||
document=factory.build_metadata_document(metadata={'experiment_name': 'x' * 51})
|
||||
),
|
||||
'experiment name',
|
||||
),
|
||||
(
|
||||
factory.ZipSpec(
|
||||
document=factory.build_metadata_document(metadata={'experiment_name': 'ab'})
|
||||
),
|
||||
'experiment name',
|
||||
),
|
||||
(
|
||||
factory.ZipSpec(document={'metadata': 'not-a-dict', 'parameters': {}}),
|
||||
'not objects',
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_gate_6_rejects_a_malformed_metadata_document(work_dir, reader, spec, match):
|
||||
payload = factory.make_bundle(spec)
|
||||
|
||||
with pytest.raises(BundleGateError, match=match) as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 6
|
||||
assert error.value.step is ImportStep.STRUCTURE_VALIDATION
|
||||
assert not (work_dir / 'extract').exists()
|
||||
|
||||
|
||||
def test_gate_6_rejects_unreadable_metadata_json(work_dir, reader):
|
||||
spec = factory.ZipSpec(include_metadata=False)
|
||||
spec.extra_files['metadata.json'] = b'{not json'
|
||||
payload = factory.make_bundle(spec)
|
||||
|
||||
with pytest.raises(BundleGateError, match='readable JSON') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 6
|
||||
|
||||
|
||||
def test_gate_6_rejects_a_missing_artifacts_directory(work_dir, reader):
|
||||
spec = factory.ZipSpec(files={})
|
||||
payload = factory.make_bundle(spec)
|
||||
|
||||
with pytest.raises(BundleGateError, match='artifacts directory') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 6
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('dropped', 'match'),
|
||||
[
|
||||
('artifacts/data_model/MLmodel', 'data_model artifacts are incomplete'),
|
||||
('artifacts/prediction_model/model.pkl', 'prediction_model artifacts are incomplete'),
|
||||
],
|
||||
)
|
||||
def test_gate_6_rejects_an_incomplete_model_directory(work_dir, reader, dropped, match):
|
||||
spec = factory.ZipSpec()
|
||||
spec.files.pop(dropped)
|
||||
payload = factory.make_bundle(spec)
|
||||
|
||||
with pytest.raises(BundleGateError, match=match) as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 6
|
||||
|
||||
|
||||
def test_gate_6_rejects_a_bundle_without_data_model(work_dir, reader):
|
||||
spec = factory.ZipSpec()
|
||||
for name in list(spec.files):
|
||||
if name.startswith('artifacts/data_model/'):
|
||||
spec.files.pop(name)
|
||||
payload = factory.make_bundle(spec)
|
||||
|
||||
with pytest.raises(BundleGateError, match='no data_model artifacts') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 6
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- gate 7
|
||||
|
||||
|
||||
@pytest.mark.parametrize('name', ['artifacts/run.sh', 'artifacts/tool.bin', 'artifacts/mod.so'])
|
||||
def test_gate_7_rejects_content_a_model_export_should_not_carry(work_dir, reader, name):
|
||||
spec = factory.ZipSpec()
|
||||
spec.extra_files[name] = b'#!/bin/sh\n'
|
||||
payload = factory.make_bundle(spec)
|
||||
|
||||
with pytest.raises(BundleGateError, match='should not carry') as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 7
|
||||
assert error.value.step is ImportStep.CONTENT_POLICY
|
||||
assert not (work_dir / 'extract').exists()
|
||||
|
||||
|
||||
def test_gate_7_rejects_a_non_regular_file_left_in_the_tree(work_dir, reader):
|
||||
"""Extraction itself only ever writes plain files and directories; this is the check that
|
||||
would catch anything else showing up in the tree by some other means."""
|
||||
root = work_dir / 'weird_root'
|
||||
root.mkdir()
|
||||
os.mkfifo(root / 'pipe')
|
||||
|
||||
with pytest.raises(BundleGateError, match='not a file') as error:
|
||||
reader._gate_7_content_policy(root)
|
||||
|
||||
assert error.value.gate == 7
|
||||
|
||||
|
||||
def test_gate_7_rejects_disallowed_content_regardless_of_traversal_order(work_dir, reader):
|
||||
"""Achado #12: gate 7 no longer sorts the tree before walking it, so a disallowed file that
|
||||
would not be first or last alphabetically must still be caught."""
|
||||
root = work_dir / 'unsorted_root'
|
||||
root.mkdir()
|
||||
for name in ('a_file.txt', 'b_file.txt', 'y_file.txt', 'z_file.txt'):
|
||||
(root / name).write_bytes(b'ok')
|
||||
(root / 'n_disallowed.sh').write_bytes(b'#!/bin/sh\n')
|
||||
|
||||
with pytest.raises(BundleGateError, match='should not carry') as error:
|
||||
reader._gate_7_content_policy(root)
|
||||
|
||||
assert error.value.gate == 7
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ ordering
|
||||
|
||||
|
||||
def test_the_earlier_gate_wins(work_dir, reader):
|
||||
"""A bundle that would fail gate 4 and gate 6 must report gate 4, having extracted nothing."""
|
||||
spec = factory.ZipSpec(
|
||||
traversal_entry='artifacts/../escape.csv',
|
||||
document=factory.build_metadata_document(drop_parameter='target_variable'),
|
||||
)
|
||||
payload = factory.make_bundle(spec)
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, payload)
|
||||
|
||||
assert error.value.gate == 4
|
||||
assert _extracted_files(work_dir / 'extract') == []
|
||||
|
||||
|
||||
def test_header_is_checked_before_the_password(work_dir, reader):
|
||||
"""A bundle with a bad header and a wrong password reports the header, not the password."""
|
||||
payload = factory.make_bundle(magic=b'NOTSIENT')
|
||||
|
||||
with pytest.raises(BundleGateError) as error:
|
||||
_open(reader, work_dir, payload, password='Nope1!Nope1!')
|
||||
|
||||
assert error.value.gate == 2
|
||||
705
tests/laborious/utils/models/test_import_errors.py
Normal file
705
tests/laborious/utils/models/test_import_errors.py
Normal file
@@ -0,0 +1,705 @@
|
||||
"""The failure catalog and the one channel every import failure passes through.
|
||||
|
||||
Two contracts are asserted here, and both are about what a person is allowed to see:
|
||||
|
||||
- **the catalog is closed and total** — every pipeline step that can fail has exactly one code and
|
||||
one sentence, the sentence never carries a path, a table name, an exception class or a traceback,
|
||||
and the classification is a function of the *step* alone, so rewording a sentence cannot move a
|
||||
code the frontend switches on;
|
||||
- **the channel is the only reporter** — one call reaches the log, the metric and the notification,
|
||||
the exception's text reaches nothing but the log, and no other module in the import path logs an
|
||||
error or notifies (asserted from the source, so the rule cannot rot).
|
||||
"""
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from pytest import mark, raises
|
||||
from temporalio.exceptions import ApplicationError
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.utils.bundle.steps import GATE_TO_STEP, ImportStep
|
||||
from laborious.utils.models.import_errors import (
|
||||
IMPORT_FAILURE_CATALOG,
|
||||
STEPS_WITHOUT_A_USER_FACING_REASON,
|
||||
UNCLASSIFIED_FAILURE,
|
||||
ImportErrorCode,
|
||||
ImportGateRejectedError,
|
||||
ImportInputError,
|
||||
ImportLogRowNotClaimableError,
|
||||
ImportStatusRejectedError,
|
||||
classify_import_failure,
|
||||
describe_failure_chain,
|
||||
failure_location,
|
||||
report_import_failure,
|
||||
summarize_failure,
|
||||
walk_cause_chain,
|
||||
)
|
||||
|
||||
REPOSITORY_ROOT = Path(__file__).resolve().parents[4]
|
||||
|
||||
# The modules that make up the import path. The channel lives in the third one; the other two may
|
||||
# not report a failure of their own.
|
||||
IMPORT_MODULES = (
|
||||
REPOSITORY_ROOT / 'laborious/activities/model_import.py',
|
||||
REPOSITORY_ROOT / 'laborious/workflows/import_model.py',
|
||||
)
|
||||
|
||||
# What a leaked exception message looks like: a work directory, the table, an exception class and a
|
||||
# traceback line. None of it may reach a recorded reason.
|
||||
LEAKY_MESSAGE = (
|
||||
"KeyError: 'target_variable' while writing public.experiment_run from "
|
||||
'/tmp/sientia-import/import_model-77-abc/extract\n' # noqa: S108 - the leak under test
|
||||
' File "laborious/activities/model_import.py", line 1, in write_import_model_document'
|
||||
)
|
||||
|
||||
# Every fragment of the leaky message that must not reach a recorded reason.
|
||||
LEAKS_THAT_MUST_NOT_SURFACE = (
|
||||
'sientia-import',
|
||||
'public.experiment_run',
|
||||
'KeyError',
|
||||
'File "',
|
||||
'line 1',
|
||||
)
|
||||
|
||||
STEPS_THAT_ONLY_SUPPORT_CAN_ACT_ON = (
|
||||
ImportStep.EXPERIMENT_CREATION,
|
||||
ImportStep.ARTIFACT_UPLOAD,
|
||||
ImportStep.REGISTRATION,
|
||||
ImportStep.MODEL_DOCUMENT,
|
||||
)
|
||||
|
||||
|
||||
class FakeSinks:
|
||||
"""The observability surface the channel writes to, recorded rather than performed."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.logged: list[tuple[str, dict]] = []
|
||||
self.metrics: list[tuple[Any, dict]] = []
|
||||
self.notifications: list[dict] = []
|
||||
|
||||
def error(self, message: str, metadata: dict | None = None) -> None:
|
||||
self.logged.append((message, metadata or {}))
|
||||
|
||||
def get_core_labels(self, metadata: dict | None = None, operation_type: str = '-') -> dict:
|
||||
return {
|
||||
'model_name': (metadata or {}).get('model_name', '-'),
|
||||
'operation_type': operation_type,
|
||||
}
|
||||
|
||||
async def emit_metric(
|
||||
self,
|
||||
metric_object: Any,
|
||||
tags: dict[str, Any],
|
||||
method: str = 'inc',
|
||||
value: float | None = None,
|
||||
) -> None:
|
||||
self.metrics.append((metric_object, tags))
|
||||
|
||||
async def send_notification_async(
|
||||
self,
|
||||
metadata: dict,
|
||||
notification_id: str,
|
||||
message: str,
|
||||
block: str,
|
||||
level: str = 'INFO',
|
||||
attachment_content: str | None = None,
|
||||
) -> None:
|
||||
self.notifications.append(
|
||||
{
|
||||
'metadata': metadata,
|
||||
'notification_id': notification_id,
|
||||
'message': message,
|
||||
'block': block,
|
||||
'level': level,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------- the catalog
|
||||
|
||||
|
||||
def test_every_failing_step_has_exactly_one_entry():
|
||||
"""A step added without a sentence fails here rather than reaching a user as `None`."""
|
||||
expected = set(ImportStep) - STEPS_WITHOUT_A_USER_FACING_REASON
|
||||
assert set(IMPORT_FAILURE_CATALOG) == expected
|
||||
|
||||
|
||||
def test_cleanup_is_the_only_step_without_a_sentence():
|
||||
"""Cleanup failing is a log line on a successful import, never the reported failure."""
|
||||
assert STEPS_WITHOUT_A_USER_FACING_REASON == frozenset({ImportStep.CLEANUP})
|
||||
assert ImportStep.CLEANUP not in IMPORT_FAILURE_CATALOG
|
||||
|
||||
|
||||
def test_every_gate_maps_to_a_step_that_has_a_sentence():
|
||||
for step in GATE_TO_STEP.values():
|
||||
assert step in IMPORT_FAILURE_CATALOG
|
||||
|
||||
|
||||
def test_each_step_has_its_own_code():
|
||||
codes = [code for code, _, _ in IMPORT_FAILURE_CATALOG.values()]
|
||||
assert len(codes) == len(set(codes))
|
||||
|
||||
|
||||
def test_the_codes_that_never_reach_a_record_are_absent_from_the_catalog():
|
||||
catalogued = {code for code, _, _ in IMPORT_FAILURE_CATALOG.values()}
|
||||
for code in (
|
||||
ImportErrorCode.LOG_ROW_NOT_CLAIMABLE,
|
||||
ImportErrorCode.STATUS_NOT_RECORDED,
|
||||
ImportErrorCode.CLEANUP_INCOMPLETE,
|
||||
):
|
||||
assert code not in catalogued
|
||||
|
||||
|
||||
@mark.parametrize('step', list(IMPORT_FAILURE_CATALOG))
|
||||
def test_a_reason_never_carries_the_exception(step: ImportStep):
|
||||
failure = classify_import_failure(step)
|
||||
|
||||
for leak in LEAKS_THAT_MUST_NOT_SURFACE:
|
||||
assert leak not in failure.reason
|
||||
assert failure.reason == IMPORT_FAILURE_CATALOG[step][1]
|
||||
|
||||
|
||||
@mark.parametrize('step', list(IMPORT_FAILURE_CATALOG))
|
||||
def test_the_classification_is_a_function_of_the_step_alone(step: ImportStep):
|
||||
"""Classification does not read the exception at all — it was already dropped as a parameter."""
|
||||
first = classify_import_failure(step)
|
||||
second = classify_import_failure(step)
|
||||
|
||||
assert first == second
|
||||
assert first.code is IMPORT_FAILURE_CATALOG[step][0]
|
||||
|
||||
|
||||
def test_the_decryption_reason_does_not_say_which_of_the_two_it_was():
|
||||
reason = IMPORT_FAILURE_CATALOG[ImportStep.DECRYPTION][1]
|
||||
|
||||
# Both causes are named as possibilities; neither is asserted, because the AEAD cannot tell them
|
||||
# apart and guessing would be a hint about somebody else's password.
|
||||
assert 'password is wrong' in reason
|
||||
assert 'damaged' in reason
|
||||
assert 'altered' in reason
|
||||
|
||||
|
||||
@mark.parametrize('step', STEPS_THAT_ONLY_SUPPORT_CAN_ACT_ON)
|
||||
def test_a_reason_the_user_cannot_act_on_names_the_import_identifier(step: ImportStep):
|
||||
assert "Contact support with this import's identifier." in IMPORT_FAILURE_CATALOG[step][1]
|
||||
|
||||
|
||||
def test_a_failure_with_no_step_is_unclassified():
|
||||
failure = classify_import_failure(None)
|
||||
|
||||
assert failure.step is None
|
||||
assert failure.code is ImportErrorCode.UNEXPECTED_ERROR
|
||||
assert failure.reason == UNCLASSIFIED_FAILURE[1]
|
||||
|
||||
|
||||
def test_every_reason_is_long_enough_for_the_column():
|
||||
"""`error_message` carries a `LENGTH(...) >= 3` CHECK; every sentence here clears it."""
|
||||
for _, reason, _ in IMPORT_FAILURE_CATALOG.values():
|
||||
assert len(reason) >= 3
|
||||
|
||||
|
||||
# ----------------------------------------------------------------- the typed failures
|
||||
|
||||
|
||||
def test_the_two_unrepeatable_failures_are_non_retryable():
|
||||
assert ImportLogRowNotClaimableError('x').non_retryable is True
|
||||
assert ImportStatusRejectedError('x').non_retryable is True
|
||||
assert ImportInputError('x').non_retryable is True
|
||||
|
||||
|
||||
def test_a_gate_rejection_carries_its_step_and_gate_across_the_boundary():
|
||||
"""An exception loses its Python type crossing an activity boundary but keeps its details."""
|
||||
error = ImportGateRejectedError('rejected', step=ImportStep.ARCHIVE_INSPECTION, gate=4)
|
||||
|
||||
assert failure_location(error) == ('archive_inspection', 4)
|
||||
|
||||
|
||||
def test_an_input_error_reports_the_received_step():
|
||||
assert failure_location(ImportInputError('missing object_key')) == ('received', None)
|
||||
|
||||
|
||||
def test_a_located_failure_is_found_through_the_activity_wrapper():
|
||||
wrapper = RuntimeError('activity failed')
|
||||
wrapper.__cause__ = ImportGateRejectedError('nope', step=ImportStep.EXTRACTION, gate=5)
|
||||
|
||||
assert failure_location(wrapper) == ('extraction', 5)
|
||||
|
||||
|
||||
def test_an_unlocated_failure_says_nothing_rather_than_guessing():
|
||||
assert failure_location(RuntimeError('boom')) == (None, None)
|
||||
assert failure_location(None) == (None, None)
|
||||
|
||||
|
||||
def test_a_cause_chain_cannot_loop_forever():
|
||||
first = RuntimeError('a')
|
||||
second = RuntimeError('b')
|
||||
first.__cause__ = second
|
||||
second.__cause__ = first
|
||||
|
||||
assert failure_location(first) == (None, None)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------- the channel
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_channel_hits_every_sink_once():
|
||||
sinks = FakeSinks()
|
||||
|
||||
failure = await report_import_failure(
|
||||
sinks,
|
||||
ImportStep.REGISTRATION,
|
||||
RuntimeError(LEAKY_MESSAGE),
|
||||
import_run_id=77,
|
||||
workflow_id='import_model-77-abc',
|
||||
metadata={'model_name': 'imported_model'},
|
||||
)
|
||||
|
||||
assert failure.code is ImportErrorCode.MODEL_NOT_REGISTERED
|
||||
assert len(sinks.logged) == 1
|
||||
assert len(sinks.metrics) == 1
|
||||
assert len(sinks.notifications) == 1
|
||||
assert sinks.metrics[0][0] is metrics.MODEL_IMPORT_ERROR_COUNT
|
||||
assert sinks.metrics[0][1]['operation_type'] == ImportErrorCode.MODEL_NOT_REGISTERED.value
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_exception_text_reaches_the_log_and_nothing_else():
|
||||
sinks = FakeSinks()
|
||||
try:
|
||||
raise RuntimeError(LEAKY_MESSAGE)
|
||||
except RuntimeError as raised:
|
||||
caught = raised
|
||||
|
||||
await report_import_failure(
|
||||
sinks,
|
||||
ImportStep.MODEL_DOCUMENT,
|
||||
caught,
|
||||
import_run_id=77,
|
||||
workflow_id='import_model-77-abc',
|
||||
)
|
||||
|
||||
logged, _ = sinks.logged[0]
|
||||
assert 'KeyError' in logged
|
||||
assert 'public.experiment_run' in logged
|
||||
assert 'Traceback' in logged
|
||||
notification = sinks.notifications[0]
|
||||
assert notification['message'] == IMPORT_FAILURE_CATALOG[ImportStep.MODEL_DOCUMENT][1]
|
||||
assert 'KeyError' not in notification['message']
|
||||
assert 'KeyError' not in str(sinks.metrics[0][1])
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_log_line_correlates_the_record_and_the_workflow():
|
||||
sinks = FakeSinks()
|
||||
|
||||
await report_import_failure(
|
||||
sinks,
|
||||
ImportStep.DOWNLOAD,
|
||||
RuntimeError('boom'),
|
||||
import_run_id=77,
|
||||
workflow_id='import_model-77-abc',
|
||||
)
|
||||
|
||||
logged, _ = sinks.logged[0]
|
||||
assert 'import_run_id=77' in logged
|
||||
assert 'workflow_id=import_model-77-abc' in logged
|
||||
assert 'step=download' in logged
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_failure_that_is_not_a_step_reports_under_its_marker():
|
||||
sinks = FakeSinks()
|
||||
|
||||
failure = await report_import_failure(
|
||||
sinks,
|
||||
None,
|
||||
None,
|
||||
import_run_id=77,
|
||||
workflow_id='wf',
|
||||
code=ImportErrorCode.STATUS_NOT_RECORDED,
|
||||
marker='status_write',
|
||||
verdict='succeeded=True step=model_document gate=None',
|
||||
detail_text='OperationalError: connection refused',
|
||||
)
|
||||
|
||||
logged, _ = sinks.logged[0]
|
||||
assert failure.code is ImportErrorCode.STATUS_NOT_RECORDED
|
||||
assert 'step=status_write' in logged
|
||||
assert 'verdict_not_recorded=succeeded=True step=model_document gate=None' in logged
|
||||
assert 'connection refused' in logged
|
||||
assert sinks.notifications[0]['block'] == 'status_write'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_report_without_an_exception_object_says_so():
|
||||
sinks = FakeSinks()
|
||||
|
||||
await report_import_failure(sinks, None, None, import_run_id=77, workflow_id='wf')
|
||||
|
||||
assert 'no exception object available' in sinks.logged[0][0]
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_an_explicit_code_gets_its_own_sentence_not_the_unclassified_one():
|
||||
"""The cleanup case: its own code, and a sentence a person can still read.
|
||||
|
||||
A cleanup-only failure means the model *was* imported — only its temporary files were left
|
||||
behind. `UNCLASSIFIED_FAILURE`'s sentence ("no model was created") would be false here, so the
|
||||
override in `_OVERRIDE_REASONS` must be what is reported, not the catalog's fallback.
|
||||
"""
|
||||
sinks = FakeSinks()
|
||||
|
||||
failure = await report_import_failure(
|
||||
sinks,
|
||||
ImportStep.CLEANUP,
|
||||
OSError('device busy'),
|
||||
import_run_id=77,
|
||||
workflow_id='wf',
|
||||
code=ImportErrorCode.CLEANUP_INCOMPLETE,
|
||||
marker=ImportStep.CLEANUP.value,
|
||||
)
|
||||
|
||||
assert failure.code is ImportErrorCode.CLEANUP_INCOMPLETE
|
||||
assert failure.reason != UNCLASSIFIED_FAILURE[1]
|
||||
assert 'no model was created' not in failure.reason
|
||||
assert 'unexpected reason' not in failure.reason
|
||||
assert 'imported successfully' in failure.reason
|
||||
assert sinks.notifications[0]['message'] == failure.reason
|
||||
assert sinks.notifications[0]['level'] == 'ERROR'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_an_unclaimable_row_gets_its_own_reason_not_the_generic_one():
|
||||
"""`LOG_ROW_NOT_CLAIMABLE` never attempted the pipeline, so it must not say "no model was
|
||||
created" either — it gets its own sentence about the request never starting."""
|
||||
sinks = FakeSinks()
|
||||
|
||||
failure = await report_import_failure(
|
||||
sinks,
|
||||
None,
|
||||
ImportLogRowNotClaimableError('row 29 is in status IMPORT_ERROR', import_run_id=29),
|
||||
import_run_id=29,
|
||||
workflow_id='manual-import-29',
|
||||
code=ImportErrorCode.LOG_ROW_NOT_CLAIMABLE,
|
||||
marker='not_claimable',
|
||||
)
|
||||
|
||||
assert failure.code is ImportErrorCode.LOG_ROW_NOT_CLAIMABLE
|
||||
assert failure.reason != UNCLASSIFIED_FAILURE[1]
|
||||
assert 'no model was created' not in failure.reason
|
||||
assert 'unexpected reason' not in failure.reason
|
||||
assert failure.reason == 'The import request could not be started.'
|
||||
assert sinks.notifications[0]['message'] == failure.reason
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_status_write_failure_gets_its_own_reason_not_the_generic_one():
|
||||
"""`STATUS_NOT_RECORDED` is a terminal-write failure, not a stalled pipeline — its sentence
|
||||
says the outcome could not be recorded, not that nothing was created."""
|
||||
sinks = FakeSinks()
|
||||
|
||||
failure = await report_import_failure(
|
||||
sinks,
|
||||
None,
|
||||
None,
|
||||
import_run_id=77,
|
||||
workflow_id='wf',
|
||||
code=ImportErrorCode.STATUS_NOT_RECORDED,
|
||||
marker='status_write',
|
||||
verdict='succeeded=True step=model_document gate=None',
|
||||
detail_text='OperationalError: connection refused',
|
||||
)
|
||||
|
||||
assert failure.code is ImportErrorCode.STATUS_NOT_RECORDED
|
||||
assert failure.reason != UNCLASSIFIED_FAILURE[1]
|
||||
assert 'no model was created' not in failure.reason
|
||||
assert 'unexpected reason' not in failure.reason
|
||||
assert failure.reason == 'The import outcome could not be recorded.'
|
||||
assert sinks.notifications[0]['message'] == failure.reason
|
||||
|
||||
|
||||
# --------------------------------------------- the fields the two log lines carry
|
||||
#
|
||||
# The message text is for a person reading a log; these fields are for the query that finds it. Both
|
||||
# lines the channel produces — the technical one and the user-facing sentence the notification logs —
|
||||
# carry the same ones, which is what makes `import_run_id=30` return the pair.
|
||||
|
||||
|
||||
def _wrapped_activity_failure() -> BaseException:
|
||||
"""What an activity failure looks like by the time the channel sees it.
|
||||
|
||||
A `MlflowException` raised inside an activity reaches the workflow as an `ApplicationError` whose
|
||||
`type` still names it, wrapped in the error Temporal raises for the failed activity. Built by hand
|
||||
rather than through a real activity because what is under test is the unwrapping.
|
||||
"""
|
||||
root = ApplicationError(
|
||||
'API request to http://localhost:5080/api/2.0/mlflow/experiments/get-by-name failed with '
|
||||
'exception HTTPConnectionPool(host=localhost, port=5080): Connection refused',
|
||||
type='MlflowException',
|
||||
)
|
||||
wrapper = RuntimeError('Activity task failed')
|
||||
wrapper.__cause__ = root
|
||||
return wrapper
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_log_line_carries_the_ids_the_step_and_the_code_as_fields():
|
||||
sinks = FakeSinks()
|
||||
|
||||
await report_import_failure(
|
||||
sinks,
|
||||
ImportStep.EXPERIMENT_CREATION,
|
||||
_wrapped_activity_failure(),
|
||||
import_run_id=30,
|
||||
workflow_id='manual-import-30',
|
||||
metadata={'model_name': '-', 'model_id': '30'},
|
||||
)
|
||||
|
||||
_, fields = sinks.logged[0]
|
||||
assert fields['import_run_id'] == 30
|
||||
assert fields['import_workflow_id'] == 'manual-import-30'
|
||||
assert fields['import_step'] == ImportStep.EXPERIMENT_CREATION.value
|
||||
assert fields['import_code'] == ImportErrorCode.STORAGE_PREPARATION_FAILED.value
|
||||
# The workflow metadata every other line already carries is kept, not replaced.
|
||||
assert fields['model_id'] == '30'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_fields_name_the_root_cause_not_the_wrapper():
|
||||
"""`error_type=ApplicationError` would be true of every activity failure there is."""
|
||||
sinks = FakeSinks()
|
||||
|
||||
await report_import_failure(
|
||||
sinks,
|
||||
ImportStep.EXPERIMENT_CREATION,
|
||||
_wrapped_activity_failure(),
|
||||
import_run_id=30,
|
||||
workflow_id='manual-import-30',
|
||||
)
|
||||
|
||||
_, fields = sinks.logged[0]
|
||||
assert fields['error_type'] == 'MlflowException'
|
||||
assert 'Connection refused' in fields['error_summary']
|
||||
assert '\n' not in fields['error_summary']
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_user_facing_line_is_greppable_by_the_same_fields():
|
||||
"""The sentence a person reads said nothing technical, and nothing said where to look."""
|
||||
sinks = FakeSinks()
|
||||
|
||||
await report_import_failure(
|
||||
sinks,
|
||||
None,
|
||||
ImportLogRowNotClaimableError('row 29 is in status IMPORT_ERROR', import_run_id=29),
|
||||
import_run_id=29,
|
||||
workflow_id='manual-import-29',
|
||||
code=ImportErrorCode.LOG_ROW_NOT_CLAIMABLE,
|
||||
marker='not_claimable',
|
||||
)
|
||||
|
||||
notified = sinks.notifications[0]
|
||||
assert notified['message'] == 'The import request could not be started.'
|
||||
assert notified['metadata']['import_code'] == ImportErrorCode.LOG_ROW_NOT_CLAIMABLE.value
|
||||
assert notified['metadata']['import_step'] == 'not_claimable'
|
||||
assert notified['metadata']['import_run_id'] == 29
|
||||
assert notified['metadata']['error_type'] == 'ImportLogRowNotClaimableError'
|
||||
assert notified['metadata'] == sinks.logged[0][1]
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_the_gate_and_the_unrecorded_verdict_are_fields_when_there_are_any():
|
||||
sinks = FakeSinks()
|
||||
|
||||
await report_import_failure(
|
||||
sinks,
|
||||
ImportStep.DECRYPTION,
|
||||
ImportGateRejectedError('gate 3 rejected the bundle', step=ImportStep.DECRYPTION, gate=3),
|
||||
import_run_id=77,
|
||||
workflow_id='wf',
|
||||
gate=3,
|
||||
verdict='succeeded=False step=decryption gate=3',
|
||||
)
|
||||
|
||||
_, fields = sinks.logged[0]
|
||||
assert fields['import_gate'] == 3
|
||||
assert fields['import_verdict_not_recorded'] == 'succeeded=False step=decryption gate=3'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_a_failure_with_neither_gate_nor_verdict_carries_neither_field():
|
||||
"""An empty field is noise in every line that has nothing to say in it."""
|
||||
sinks = FakeSinks()
|
||||
|
||||
await report_import_failure(
|
||||
sinks, ImportStep.DOWNLOAD, RuntimeError('boom'), import_run_id=77, workflow_id='wf'
|
||||
)
|
||||
|
||||
_, fields = sinks.logged[0]
|
||||
assert 'import_gate' not in fields
|
||||
assert 'import_verdict_not_recorded' not in fields
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_no_field_can_shadow_the_formatter_s_own_keys():
|
||||
sinks = FakeSinks()
|
||||
|
||||
await report_import_failure(
|
||||
sinks, ImportStep.DOWNLOAD, RuntimeError('boom'), import_run_id=77, workflow_id='wf'
|
||||
)
|
||||
|
||||
_, fields = sinks.logged[0]
|
||||
added = set(fields) - {'model_name', 'model_id', 'workflow_name', 'schedule_name'}
|
||||
assert all(key.startswith(('import_', 'error_')) for key in added)
|
||||
assert not {'at', 'msg', 'timestamp'} & set(fields)
|
||||
|
||||
|
||||
def test_walk_cause_chain_yields_just_the_exception_when_there_is_no_cause():
|
||||
error = RuntimeError('alone')
|
||||
|
||||
assert list(walk_cause_chain(error)) == [error]
|
||||
|
||||
|
||||
def test_walk_cause_chain_yields_the_whole_chain_when_shorter_than_the_limit():
|
||||
root = RuntimeError('root')
|
||||
middle = RuntimeError('middle')
|
||||
middle.__cause__ = root
|
||||
outer = RuntimeError('outer')
|
||||
outer.__cause__ = middle
|
||||
|
||||
assert list(walk_cause_chain(outer, limit=10)) == [outer, middle, root]
|
||||
|
||||
|
||||
def test_walk_cause_chain_stops_at_the_limit_for_a_chain_longer_than_it():
|
||||
first = RuntimeError('first')
|
||||
second = RuntimeError('second')
|
||||
first.__cause__ = second
|
||||
second.__cause__ = first
|
||||
|
||||
assert list(walk_cause_chain(first, limit=3)) == [first, second, first]
|
||||
|
||||
|
||||
def test_the_chain_is_flattened_root_cause_first():
|
||||
"""Because the first line is what becomes `error_summary`, and it has to be the cause."""
|
||||
kind, detail = describe_failure_chain(_wrapped_activity_failure())
|
||||
|
||||
assert kind == 'MlflowException'
|
||||
assert detail.splitlines()[0].startswith('MlflowException: API request to')
|
||||
assert detail.splitlines()[-1] == 'RuntimeError: Activity task failed'
|
||||
|
||||
|
||||
def test_a_cycle_in_the_chain_cannot_hang_the_reporting_of_a_failure():
|
||||
first = RuntimeError('first')
|
||||
second = RuntimeError('second')
|
||||
first.__cause__ = second
|
||||
second.__cause__ = first
|
||||
|
||||
_, detail = describe_failure_chain(first)
|
||||
|
||||
assert len(detail.splitlines()) <= 10
|
||||
|
||||
|
||||
def test_a_failure_that_arrives_as_text_keeps_the_type_that_travelled_with_it():
|
||||
kind, summary = summarize_failure(
|
||||
None,
|
||||
'MlflowException: connection refused\nRuntimeError: Activity task failed',
|
||||
'MlflowException',
|
||||
)
|
||||
|
||||
assert kind == 'MlflowException'
|
||||
assert summary == 'MlflowException: connection refused'
|
||||
|
||||
|
||||
def test_a_failure_that_arrives_as_nothing_says_so_rather_than_guessing():
|
||||
kind, summary = summarize_failure(None, None, None)
|
||||
|
||||
assert kind == 'unavailable'
|
||||
assert summary == ''
|
||||
|
||||
|
||||
def test_summarize_failure_bounds_a_cyclic_cause_chain():
|
||||
"""A cycle never hits a `None` cause, so the walk must stop by exhausting its own bound."""
|
||||
first = RuntimeError('first')
|
||||
second = RuntimeError('second')
|
||||
first.__cause__ = second
|
||||
second.__cause__ = first
|
||||
|
||||
kind, _summary = summarize_failure(first)
|
||||
|
||||
assert kind == 'RuntimeError'
|
||||
|
||||
|
||||
def test_a_leading_blank_line_is_skipped_for_the_first_real_one():
|
||||
class BlankFirstLineError(Exception):
|
||||
def __str__(self) -> str:
|
||||
return '\n \nreal message'
|
||||
|
||||
_kind, summary = summarize_failure(BlankFirstLineError())
|
||||
|
||||
assert summary == 'real message'
|
||||
|
||||
|
||||
def test_describe_failure_chain_reports_a_message_less_link_by_its_type_alone():
|
||||
"""`Exception()` has no message at all — the link still has to be readable on its own."""
|
||||
inner = Exception()
|
||||
outer = RuntimeError('outer failure')
|
||||
outer.__cause__ = inner
|
||||
|
||||
kind, detail = describe_failure_chain(outer)
|
||||
|
||||
assert kind == 'Exception'
|
||||
assert detail.splitlines() == ['Exception', 'RuntimeError: outer failure']
|
||||
|
||||
|
||||
# ------------------------------------------------- the channel is the only reporter
|
||||
|
||||
|
||||
def _reporting_calls(path: Path) -> list[str]:
|
||||
"""Every call in `path` that would log an error or send a notification."""
|
||||
forbidden = {
|
||||
'error',
|
||||
'custom_error',
|
||||
'critical',
|
||||
'custom_critical',
|
||||
'send_notification',
|
||||
'send_notification_async',
|
||||
'publish_notification',
|
||||
}
|
||||
tree = ast.parse(path.read_text())
|
||||
found = []
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.Call):
|
||||
continue
|
||||
target = node.func
|
||||
name = target.attr if isinstance(target, ast.Attribute) else getattr(target, 'id', '')
|
||||
if name in forbidden:
|
||||
found.append(f'{path.name}:{node.lineno}:{name}')
|
||||
return found
|
||||
|
||||
|
||||
@mark.parametrize('module_path', IMPORT_MODULES, ids=lambda path: path.name)
|
||||
def test_no_import_module_reports_a_failure_of_its_own(module_path: Path):
|
||||
"""So no failure is reported twice, or half-reported by whoever noticed it first."""
|
||||
assert _reporting_calls(module_path) == []
|
||||
|
||||
|
||||
def test_the_channel_itself_is_where_the_reporting_lives():
|
||||
channel = REPOSITORY_ROOT / 'laborious/utils/models/import_errors.py'
|
||||
|
||||
calls = _reporting_calls(channel)
|
||||
|
||||
assert [call.split(':')[2] for call in calls] == ['error', 'send_notification_async']
|
||||
|
||||
|
||||
def test_the_reader_neither_logs_nor_notifies():
|
||||
"""The bundle reader raises; it is the activity layer that decides what a failure means."""
|
||||
for module in ('laborious/utils/bundle/reader.py', 'laborious/utils/bundle/format.py'):
|
||||
assert _reporting_calls(REPOSITORY_ROOT / module) == []
|
||||
|
||||
|
||||
def test_an_unclassified_step_cannot_be_recorded_as_a_known_one():
|
||||
with raises(ValueError, match='not a valid ImportStep'):
|
||||
ImportStep('provisioning')
|
||||
@@ -214,6 +214,27 @@ async def test_from_dataframe_offloaded(mock_now):
|
||||
minio.upload_file.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
|
||||
async def test_from_dataframe_inline_keeps_a_provided_last_timestamp():
|
||||
"""An explicit last_timestamp must survive as-is, skipping the max() derivation."""
|
||||
minio = AsyncMock()
|
||||
logger = MagicMock()
|
||||
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
|
||||
|
||||
result = await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=df,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
last_timestamp='2030-01-01T00:00:00+00:00',
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
assert result.last_timestamp == '2030-01-01T00:00:00+00:00'
|
||||
logger.custom_debug.assert_called()
|
||||
|
||||
|
||||
def test_from_dict_inline():
|
||||
raw = {
|
||||
'last_timestamp': '2024-01-01T00:00:00+00:00',
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,5 +1,6 @@
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
@@ -8,11 +9,13 @@ from asyncua.crypto.security_policies import SecurityPolicyBasic256
|
||||
from asyncua.ua.uaerrors import BadNodeIdUnknown, BadSessionIdInvalid
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.utils.repository.opc_repository import (
|
||||
OpcClientAlreadyExistsError,
|
||||
OpcClientNotInitializedError,
|
||||
OpcRepository,
|
||||
OpcSessionAlreadyConnectedError,
|
||||
_opc_authentication_token_str,
|
||||
is_reconnectable_opcua_bad,
|
||||
)
|
||||
|
||||
@@ -80,6 +83,40 @@ def test_init(opc_repository):
|
||||
assert opc_repository.last_reconnection_time is None
|
||||
|
||||
|
||||
def test_opc_authentication_token_str_no_protocol():
|
||||
client = MagicMock()
|
||||
client.uaclient.protocol = None
|
||||
|
||||
assert _opc_authentication_token_str(client) == 'unknown'
|
||||
|
||||
|
||||
def test_opc_authentication_token_str_no_token():
|
||||
client = MagicMock()
|
||||
client.uaclient.protocol = MagicMock(authentication_token=None)
|
||||
|
||||
assert _opc_authentication_token_str(client) == 'unknown'
|
||||
|
||||
|
||||
def test_opc_authentication_token_str_swallows_errors():
|
||||
class _ExplodingClient:
|
||||
@property
|
||||
def uaclient(self):
|
||||
raise RuntimeError('boom')
|
||||
|
||||
assert _opc_authentication_token_str(_ExplodingClient()) == 'unknown'
|
||||
|
||||
|
||||
def test_is_session_open_swallows_errors(opc_repository):
|
||||
class _ExplodingClient:
|
||||
@property
|
||||
def uaclient(self):
|
||||
raise RuntimeError('boom')
|
||||
|
||||
opc_repository.client = _ExplodingClient()
|
||||
|
||||
assert opc_repository._is_session_open() is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_security(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
@@ -107,6 +144,16 @@ async def test_set_security_missing_certificates(opc_repository):
|
||||
assert str(e) == 'Certificate and private key paths must be provided for secure connection.'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_security_missing_server_uri(opc_repository):
|
||||
opc_repository.server_uri = None
|
||||
|
||||
try:
|
||||
await opc_repository.set_security()
|
||||
except ValueError as e:
|
||||
assert str(e) == 'server_uri must be provided for secure connection.'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_security_missing_client(opc_repository):
|
||||
opc_repository.client = None
|
||||
@@ -160,6 +207,46 @@ async def test_create_client_raises_when_client_exists(opc_repository, mock_clie
|
||||
await opc_repository._create_client()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_client_builds_a_secured_client(opc_repository, mock_client):
|
||||
opc_repository.pod_id = 'pod-1'
|
||||
opc_repository.set_security = AsyncMock()
|
||||
|
||||
await opc_repository._create_client()
|
||||
|
||||
assert opc_repository.client is mock_client
|
||||
assert mock_client.name == 'pod-1'
|
||||
assert mock_client.application_uri == 'pod:1'
|
||||
assert mock_client.product_uri == 'pod:1'
|
||||
opc_repository.set_security.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_client_skips_security_without_a_cert_path(opc_repository, mock_client):
|
||||
opc_repository.cert_path = None
|
||||
opc_repository.set_security = AsyncMock()
|
||||
|
||||
await opc_repository._create_client()
|
||||
|
||||
assert opc_repository.client is mock_client
|
||||
opc_repository.set_security.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_locked_reuses_an_existing_client(opc_repository, mock_client):
|
||||
"""A client already set (e.g. surviving a prior failed open) is not recreated."""
|
||||
opc_repository.client = mock_client
|
||||
opc_repository._is_session_open = MagicMock(return_value=False)
|
||||
opc_repository._create_client = AsyncMock()
|
||||
opc_repository._open_session = AsyncMock(return_value=(True, {}))
|
||||
|
||||
result = await opc_repository._connect_locked()
|
||||
|
||||
opc_repository._create_client.assert_not_called()
|
||||
opc_repository._open_session.assert_called_once()
|
||||
assert result == (True, {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_open_session_success(opc_repository):
|
||||
closed_proto = MagicMock()
|
||||
@@ -435,6 +522,82 @@ async def test_write_data_write_value_failed(opc_repository, mock_client):
|
||||
assert error_data['attachment_content'] is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_with_empty_metadata_uses_unknown_labels(opc_repository, mock_client):
|
||||
"""No model_id/model_name in the write metadata: the metric labels fall back to 'unknown'."""
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
mock_node = AsyncMock()
|
||||
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||
|
||||
result = await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', {})
|
||||
|
||||
assert result == (True, {'response_time': ANY})
|
||||
last_call = opc_repository.emit_metric.call_args_list[-1]
|
||||
assert last_call.args[0] is metrics.OPC_WRITE_ATTEMPTS_TOTAL
|
||||
assert last_call.args[1]['model_id'] == 'unknown'
|
||||
assert last_call.args[1]['model_name'] == 'unknown'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_get_node_bad_session_id_schedules_reconnect(opc_repository, mock_client):
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
opc_repository._start_reconnect = AsyncMock()
|
||||
mock_client.get_node = MagicMock(side_effect=BadSessionIdInvalid())
|
||||
|
||||
is_success, error_data = await opc_repository.write_data(
|
||||
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||
)
|
||||
|
||||
opc_repository._start_reconnect.assert_called_once()
|
||||
assert is_success is False
|
||||
assert error_data['opc_error_kind'] == 'session_bad'
|
||||
assert error_data['opc_status'] == 'BadSessionIdInvalid'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_logs_inter_arrival_and_flags_a_stale_session(opc_repository, mock_client):
|
||||
"""The second write on the same session logs the gap since the previous one, and flags it
|
||||
when the gap outlives the session's revised timeout."""
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
mock_client.session_timeout = 0.001 # 1 microsecond: any real gap will exceed it
|
||||
mock_node = AsyncMock()
|
||||
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||
|
||||
await opc_repository.write_data('ns=2;s=TestNode', 1.0, 'float', metadata['metadata'])
|
||||
await opc_repository.write_data('ns=2;s=TestNode', 2.0, 'float', metadata['metadata'])
|
||||
|
||||
logged = [c.args[0] for c in opc_repository.info.call_args_list]
|
||||
assert any('OPC write inter-arrival' in message for message in logged)
|
||||
emitted_metrics = [c.args[0] for c in opc_repository.emit_metric.call_args_list]
|
||||
assert metrics.OPC_WRITE_INTER_ARRIVAL_OVER_SESSION_TIMEOUT_TOTAL in emitted_metrics
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_write_inter_arrival_skips_the_metric_without_a_client(opc_repository):
|
||||
opc_repository._last_write_mono = time.monotonic() - 1
|
||||
opc_repository.client = None
|
||||
|
||||
await opc_repository._log_write_inter_arrival('session-1', 'ns=2;s=Node')
|
||||
|
||||
opc_repository.emit_metric.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_write_inter_arrival_skips_the_metric_within_the_session_timeout(
|
||||
opc_repository, mock_client
|
||||
):
|
||||
opc_repository._last_write_mono = time.monotonic()
|
||||
opc_repository.client = mock_client
|
||||
mock_client.session_timeout = 600_000 # 600s: comfortably longer than this test takes
|
||||
|
||||
await opc_repository._log_write_inter_arrival('session-1', 'ns=2;s=Node')
|
||||
|
||||
opc_repository.emit_metric.assert_not_called()
|
||||
|
||||
|
||||
def test_is_reconnectable_opcua_bad():
|
||||
assert is_reconnectable_opcua_bad(BadSessionIdInvalid()) is True
|
||||
assert is_reconnectable_opcua_bad(BadNodeIdUnknown()) is False
|
||||
@@ -491,6 +654,22 @@ async def test_start_reconnect_skips_within_interval(opc_repository):
|
||||
assert opc_repository._reconnect_task is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_reconnect_skips_when_already_in_progress(opc_repository):
|
||||
opc_repository.last_reconnection_time = None
|
||||
in_progress_task = asyncio.create_task(asyncio.sleep(60))
|
||||
opc_repository._reconnect_task = in_progress_task
|
||||
|
||||
await opc_repository._start_reconnect('BadSessionIdInvalid', 'tok')
|
||||
|
||||
assert opc_repository._reconnect_task is in_progress_task
|
||||
opc_repository.warning.assert_called_once()
|
||||
|
||||
in_progress_task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await in_progress_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_protocol_closed_schedules_reconnect(opc_repository):
|
||||
opc_repository.client = MagicMock()
|
||||
@@ -608,3 +787,23 @@ async def test_reconnect_locked_sets_last_reconnection_time(mock_datetime, opc_r
|
||||
opc_repository._connect_locked.assert_called_once()
|
||||
assert result == (True, {})
|
||||
assert opc_repository.last_reconnection_time == datetime(2025, 1, 1, 12, 0, 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_reconnect_failed_without_error_detail_logs_once(opc_repository):
|
||||
"""A falsy `error` (no message to add) must not log a second, empty line."""
|
||||
opc_repository._reconnect_locked = AsyncMock(return_value=(False, {}))
|
||||
|
||||
await opc_repository._run_reconnect('BadSessionIdInvalid', 'tok')
|
||||
|
||||
assert opc_repository.error.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_reconnect_swallows_an_unexpected_exception(opc_repository):
|
||||
"""A bug inside the reconnect itself must not escape the background task."""
|
||||
opc_repository._reconnect_locked = AsyncMock(side_effect=RuntimeError('boom'))
|
||||
|
||||
await opc_repository._run_reconnect('BadSessionIdInvalid', 'tok')
|
||||
|
||||
assert opc_repository.error.call_count >= 2
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
"""
|
||||
Integration test: a `sklearn`-flavor model with a custom class survives a retrain round trip.
|
||||
|
||||
Unlike the rest of `test_model_repository.py`, this test does not mock the `mlflow` module —
|
||||
`resolve_code_paths` forwarding `code_paths=` to `mlflow.sklearn.log_model` is only proven by
|
||||
letting the real MLflow file-based `code_paths` validation/copy machinery run, and by really
|
||||
dropping and reloading the custom class module afterwards. A fully mocked `mlflow.sklearn` (as
|
||||
used elsewhere in this file, and in `e2e/test_retrain_joblib_flavor.py`) would never exercise
|
||||
that machinery at all.
|
||||
|
||||
The Model Registry itself is not exercised here — MLflow's registry requires a database-backed
|
||||
store, not a bare `file:` tracking URI — so this drives the write (`log_model`) and read
|
||||
(`client.download_artifacts` + `mlflow.sklearn.load_model`) by run id directly, the same
|
||||
building blocks `download_model` composes for the `sklearn` flavor once a Production run is
|
||||
resolved.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import mlflow
|
||||
import pytest
|
||||
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
|
||||
CODE_MODULE_NAME = 'sklearn_flavor_retrain_probe'
|
||||
|
||||
CODE_MODULE_SOURCE = '''"""Custom estimator that exists only inside the artifact's code/ directory."""
|
||||
|
||||
|
||||
class SklearnFlavorProbePredictor:
|
||||
def __init__(self, offset):
|
||||
self.offset = offset
|
||||
|
||||
def fit(self, X, y=None):
|
||||
return self
|
||||
|
||||
def predict(self, X):
|
||||
return [row[0] + self.offset for row in X]
|
||||
'''
|
||||
|
||||
|
||||
def _write_code_module(code_dir: Path) -> Path:
|
||||
code_dir.mkdir(parents=True)
|
||||
(code_dir / f'{CODE_MODULE_NAME}.py').write_text(CODE_MODULE_SOURCE)
|
||||
return code_dir
|
||||
|
||||
|
||||
def _import_fresh_probe_model(code_dir: Path, offset: int):
|
||||
"""Instantiate the custom class with its module on `sys.path`, then undo both."""
|
||||
sys.path.insert(0, str(code_dir))
|
||||
try:
|
||||
module = importlib.import_module(CODE_MODULE_NAME)
|
||||
return module.SklearnFlavorProbePredictor(offset)
|
||||
finally:
|
||||
sys.path.remove(str(code_dir))
|
||||
sys.modules.pop(CODE_MODULE_NAME, None)
|
||||
importlib.invalidate_caches()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_retrain_sklearn_model_preserves_code_dir_and_reloads(tmp_path):
|
||||
"""A sklearn retrain with a custom `code/` class keeps it and can reload it afterwards."""
|
||||
tracking_uri = f'file://{tmp_path / "mlruns"}'
|
||||
mlflow.set_tracking_uri(tracking_uri)
|
||||
mlflow.set_experiment('sklearn_flavor_retrain_e2e')
|
||||
|
||||
repository = MLFlowRepository(
|
||||
host=tracking_uri,
|
||||
username='test',
|
||||
password='test',
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
repository.emit_metric = AsyncMock()
|
||||
repository.observe_lag = AsyncMock()
|
||||
|
||||
source_code_dir = _write_code_module(tmp_path / 'source_artifact' / 'code')
|
||||
model = _import_fresh_probe_model(source_code_dir, offset=10)
|
||||
model_data = {'model': model, 'artifact_path': str(tmp_path / 'source_artifact')}
|
||||
|
||||
with mlflow.start_run() as run:
|
||||
await repository.log_model(model_data, 'sklearn', 'prediction_model', {'metadata': {}})
|
||||
run_id = run.info.run_id
|
||||
|
||||
written_code_dir = Path(repository.client.download_artifacts(run_id, 'prediction_model/code'))
|
||||
assert (written_code_dir / f'{CODE_MODULE_NAME}.py').is_file()
|
||||
|
||||
# The class must resolve from the artifact's written code/, not from anything left in memory.
|
||||
sys.modules.pop(CODE_MODULE_NAME, None)
|
||||
importlib.invalidate_caches()
|
||||
|
||||
local_model_dir = repository.client.download_artifacts(run_id, 'prediction_model')
|
||||
repository.prepend_artifact_code_dir(local_model_dir, 'sklearn', {})
|
||||
|
||||
reloaded = mlflow.sklearn.load_model(local_model_dir)
|
||||
|
||||
assert reloaded.predict([[1], [2]]) == [11, 12]
|
||||
@@ -1,11 +1,51 @@
|
||||
from os import environ
|
||||
import tempfile
|
||||
from os import environ, path
|
||||
|
||||
import pytest
|
||||
|
||||
from laborious.utils.connectors_config import (
|
||||
build_import_config,
|
||||
build_import_status_config,
|
||||
build_minio_config,
|
||||
build_mlflow_config,
|
||||
build_opc_config,
|
||||
)
|
||||
|
||||
IMPORT_ENV_VARS = (
|
||||
'IMPORT_PASSWORD_KEY',
|
||||
'IMPORT_BUNDLE_BUCKET',
|
||||
'IMPORT_BUNDLE_PREFIX',
|
||||
'IMPORT_WORK_DIR',
|
||||
'IMPORT_MAX_OBJECT_BYTES',
|
||||
'IMPORT_MAX_ARCHIVE_ENTRIES',
|
||||
'IMPORT_MAX_UNCOMPRESSED_BYTES',
|
||||
'IMPORT_MAX_COMPRESSION_RATIO',
|
||||
'IMPORT_MODELS_COLLECTION',
|
||||
'IMPORT_BUNDLE_RETENTION_DAYS',
|
||||
'IMPORT_STATUS_DB_HOST',
|
||||
'IMPORT_STATUS_DB_PORT',
|
||||
'IMPORT_STATUS_DB_NAME',
|
||||
'IMPORT_STATUS_DB_USER',
|
||||
'IMPORT_STATUS_DB_PASSWORD',
|
||||
)
|
||||
|
||||
# The import log connection falls back to the POSTGRES_* server, so a developer's own POSTGRES_*
|
||||
# would otherwise decide what `build_import_status_config` returns under test.
|
||||
POSTGRES_ENV_VARS = (
|
||||
'POSTGRES_HOST',
|
||||
'POSTGRES_PORT',
|
||||
'POSTGRES_USER',
|
||||
'POSTGRES_PASSWORD',
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clean_import_env(monkeypatch):
|
||||
"""Import configuration is read from a known-empty environment."""
|
||||
for name in IMPORT_ENV_VARS + POSTGRES_ENV_VARS:
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
return monkeypatch
|
||||
|
||||
|
||||
def test_build_mlflow_config_with_env_vars():
|
||||
# Arrange
|
||||
@@ -120,3 +160,123 @@ def test_build_minio_config_with_defaults():
|
||||
'retention_hours': 24,
|
||||
'secure': False,
|
||||
}
|
||||
|
||||
|
||||
def test_build_import_config_with_defaults(clean_import_env):
|
||||
assert build_import_config() == {
|
||||
'password_key': None,
|
||||
'bucket': 'sientia',
|
||||
'prefix': 'imported_models/',
|
||||
'work_dir': path.join(tempfile.gettempdir(), 'sientia-import'),
|
||||
'max_object_bytes': 1024 * 1024 * 1024,
|
||||
'max_entries': 5000,
|
||||
'max_uncompressed_bytes': 4 * 1024 * 1024 * 1024,
|
||||
'max_compression_ratio': 200.0,
|
||||
'models_collection': 'models',
|
||||
'counters_collection': 'model_id_counters',
|
||||
'retention_days': 7,
|
||||
}
|
||||
|
||||
|
||||
def test_build_import_config_reads_the_environment(clean_import_env):
|
||||
clean_import_env.setenv('IMPORT_PASSWORD_KEY', 'a2V5')
|
||||
clean_import_env.setenv('IMPORT_BUNDLE_BUCKET', 'other-bucket')
|
||||
clean_import_env.setenv('IMPORT_BUNDLE_PREFIX', 'incoming/models')
|
||||
clean_import_env.setenv('IMPORT_WORK_DIR', '/var/lib/import')
|
||||
clean_import_env.setenv('IMPORT_MAX_OBJECT_BYTES', '2048')
|
||||
clean_import_env.setenv('IMPORT_MAX_ARCHIVE_ENTRIES', '12')
|
||||
clean_import_env.setenv('IMPORT_MAX_UNCOMPRESSED_BYTES', '4096')
|
||||
clean_import_env.setenv('IMPORT_MAX_COMPRESSION_RATIO', '12.5')
|
||||
clean_import_env.setenv('IMPORT_MODELS_COLLECTION', 'models_v2')
|
||||
clean_import_env.setenv('IMPORT_MODEL_ID_COUNTERS_COLLECTION', 'model_id_counters_v2')
|
||||
clean_import_env.setenv('IMPORT_BUNDLE_RETENTION_DAYS', '14')
|
||||
|
||||
config = build_import_config()
|
||||
|
||||
assert config['password_key'] == 'a2V5'
|
||||
assert config['bucket'] == 'other-bucket'
|
||||
# A prefix without its trailing slash would make `imported_models_backup/x` pass the gate-1
|
||||
# prefix check, so the separator is added rather than trusted.
|
||||
assert config['prefix'] == 'incoming/models/'
|
||||
assert config['work_dir'] == '/var/lib/import'
|
||||
assert config['max_object_bytes'] == 2048
|
||||
assert config['max_entries'] == 12
|
||||
assert config['max_uncompressed_bytes'] == 4096
|
||||
assert config['max_compression_ratio'] == 12.5
|
||||
assert config['models_collection'] == 'models_v2'
|
||||
assert config['counters_collection'] == 'model_id_counters_v2'
|
||||
assert config['retention_days'] == 14
|
||||
|
||||
|
||||
def test_build_import_config_does_not_validate_the_password_key(clean_import_env):
|
||||
"""A missing or malformed key must fail the import at decryption, never at worker start."""
|
||||
clean_import_env.setenv('IMPORT_PASSWORD_KEY', 'not base64 at all')
|
||||
|
||||
assert build_import_config()['password_key'] == 'not base64 at all'
|
||||
|
||||
|
||||
def test_build_import_status_config_with_defaults(clean_import_env):
|
||||
assert build_import_status_config() == {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'dbname': 'sientia-core-mlops-bff',
|
||||
'user': 'sientia',
|
||||
'password': '',
|
||||
'schema': 'public',
|
||||
}
|
||||
|
||||
|
||||
def test_build_import_status_config_uses_the_postgres_server(clean_import_env):
|
||||
"""The log is another database on the POSTGRES_* server: only its name is import configuration."""
|
||||
clean_import_env.setenv('POSTGRES_HOST', 'paradedb-rw')
|
||||
clean_import_env.setenv('POSTGRES_PORT', '5433')
|
||||
clean_import_env.setenv('POSTGRES_USER', 'sientia-user')
|
||||
clean_import_env.setenv('POSTGRES_PASSWORD', 'sientia-password')
|
||||
clean_import_env.setenv('IMPORT_STATUS_DB_NAME', 'bff-db')
|
||||
|
||||
assert build_import_status_config() == {
|
||||
'host': 'paradedb-rw',
|
||||
'port': 5433,
|
||||
'dbname': 'bff-db',
|
||||
'user': 'sientia-user',
|
||||
'password': 'sientia-password',
|
||||
'schema': 'public',
|
||||
}
|
||||
|
||||
|
||||
def test_build_import_status_config_overrides_beat_the_postgres_server(clean_import_env):
|
||||
"""A log on a server of its own: the four IMPORT_STATUS_DB_* variables win over POSTGRES_*."""
|
||||
clean_import_env.setenv('POSTGRES_HOST', 'paradedb-rw')
|
||||
clean_import_env.setenv('POSTGRES_PORT', '5433')
|
||||
clean_import_env.setenv('POSTGRES_USER', 'sientia-user')
|
||||
clean_import_env.setenv('POSTGRES_PASSWORD', 'sientia-password')
|
||||
clean_import_env.setenv('IMPORT_STATUS_DB_HOST', 'bff-postgres')
|
||||
clean_import_env.setenv('IMPORT_STATUS_DB_PORT', '6543')
|
||||
clean_import_env.setenv('IMPORT_STATUS_DB_USER', 'bff-user')
|
||||
clean_import_env.setenv('IMPORT_STATUS_DB_PASSWORD', 'bff-password')
|
||||
|
||||
config = build_import_status_config()
|
||||
|
||||
assert config['host'] == 'bff-postgres'
|
||||
assert config['port'] == 6543
|
||||
assert config['user'] == 'bff-user'
|
||||
assert config['password'] == 'bff-password'
|
||||
|
||||
|
||||
def test_build_import_status_config_reads_the_environment(clean_import_env):
|
||||
clean_import_env.setenv('IMPORT_STATUS_DB_HOST', 'bff-postgres')
|
||||
clean_import_env.setenv('IMPORT_STATUS_DB_PORT', '6543')
|
||||
clean_import_env.setenv('IMPORT_STATUS_DB_NAME', 'other-db')
|
||||
clean_import_env.setenv('IMPORT_STATUS_DB_USER', 'other-user')
|
||||
clean_import_env.setenv('IMPORT_STATUS_DB_PASSWORD', 'secret')
|
||||
|
||||
config = build_import_status_config()
|
||||
|
||||
assert config == {
|
||||
'host': 'bff-postgres',
|
||||
'port': 6543,
|
||||
'dbname': 'other-db',
|
||||
'user': 'other-user',
|
||||
'password': 'secret',
|
||||
'schema': 'public',
|
||||
}
|
||||
|
||||
40
tests/laborious/utils/test_dataframe_debug.py
Normal file
40
tests/laborious/utils/test_dataframe_debug.py
Normal file
@@ -0,0 +1,40 @@
|
||||
"""Debug message formatting: full dataframe when small, a summary when it would flood the log."""
|
||||
|
||||
from pandas import DataFrame
|
||||
|
||||
from laborious.utils.dataframe_debug import (
|
||||
DEFAULT_MAX_DEBUG_DATAFRAME_ROWS,
|
||||
build_dataframe_debug_message,
|
||||
)
|
||||
|
||||
|
||||
def test_non_dataframe_payloads_are_stringified():
|
||||
assert build_dataframe_debug_message('payload:', {'a': 1}) == "payload: {'a': 1}"
|
||||
|
||||
|
||||
def test_a_small_dataframe_is_logged_in_full():
|
||||
data = DataFrame({'a': [1, 2]})
|
||||
|
||||
message = build_dataframe_debug_message('payload:', data)
|
||||
|
||||
assert message.startswith('payload:\n')
|
||||
assert data.to_csv() in message
|
||||
|
||||
|
||||
def test_a_dataframe_over_the_row_limit_is_summarized_instead():
|
||||
data = DataFrame({'a': range(DEFAULT_MAX_DEBUG_DATAFRAME_ROWS + 1)})
|
||||
|
||||
message = build_dataframe_debug_message('payload:', data)
|
||||
|
||||
assert 'skipped because dataframe has' in message
|
||||
assert str(data.shape) in message
|
||||
assert data.to_csv() not in message
|
||||
|
||||
|
||||
def test_the_row_limit_is_configurable():
|
||||
data = DataFrame({'a': [1, 2, 3]})
|
||||
|
||||
message = build_dataframe_debug_message('payload:', data, max_rows=1)
|
||||
|
||||
assert 'skipped because dataframe has 3 rows' in message
|
||||
assert '(max: 1)' in message
|
||||
125
tests/laborious/utils/test_import_password.py
Normal file
125
tests/laborious/utils/test_import_password.py
Normal file
@@ -0,0 +1,125 @@
|
||||
"""The password envelope: AES-256-GCM in, plaintext password out, nothing in between."""
|
||||
|
||||
import base64
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||||
|
||||
from laborious.utils.import_password import (
|
||||
EnvelopeError,
|
||||
decrypt_password_envelope,
|
||||
load_envelope_key,
|
||||
)
|
||||
from tests.helpers import bundle_factory as factory
|
||||
|
||||
KEY = factory.DEFAULT_ENVELOPE_KEY
|
||||
OTHER_KEY = bytes(range(1, 33))
|
||||
|
||||
|
||||
def test_recovers_the_password_used_at_export():
|
||||
envelope = factory.build_envelope('Correct1Horse!Battery', KEY)
|
||||
|
||||
assert decrypt_password_envelope(envelope, KEY) == 'Correct1Horse!Battery'
|
||||
|
||||
|
||||
def test_wrong_key_fails_authentication():
|
||||
envelope = factory.build_envelope('Correct1Horse!Battery', KEY)
|
||||
|
||||
with pytest.raises(EnvelopeError, match='authenticated'):
|
||||
decrypt_password_envelope(envelope, OTHER_KEY)
|
||||
|
||||
|
||||
def test_tampered_envelope_fails_authentication():
|
||||
payload = bytearray(base64.b64decode(factory.build_envelope('Correct1Horse!Battery', KEY)))
|
||||
payload[-1] ^= 0xFF
|
||||
tampered_envelope = base64.b64encode(bytes(payload)).decode()
|
||||
|
||||
with pytest.raises(EnvelopeError, match='authenticated'):
|
||||
decrypt_password_envelope(tampered_envelope, KEY)
|
||||
|
||||
|
||||
def test_malformed_base64_rejected():
|
||||
with pytest.raises(EnvelopeError, match='base64'):
|
||||
decrypt_password_envelope('this is not base64!!', KEY)
|
||||
|
||||
|
||||
def test_envelope_too_short_rejected():
|
||||
short_envelope = base64.b64encode(bytes(20)).decode()
|
||||
|
||||
with pytest.raises(EnvelopeError, match='too short'):
|
||||
decrypt_password_envelope(short_envelope, KEY)
|
||||
|
||||
|
||||
def test_empty_envelope_rejected():
|
||||
with pytest.raises(EnvelopeError, match='no password envelope'):
|
||||
decrypt_password_envelope('', KEY)
|
||||
|
||||
|
||||
def test_password_is_not_trimmed():
|
||||
"""A trailing space is part of the secret; the KDF must see it."""
|
||||
envelope = factory.build_envelope('Trailing1Space!Test ', KEY)
|
||||
|
||||
assert decrypt_password_envelope(envelope, KEY) == 'Trailing1Space!Test '
|
||||
|
||||
|
||||
def test_password_keeps_its_unicode_form():
|
||||
"""Normalisation happens at key derivation, not here — the envelope round-trips verbatim."""
|
||||
decomposed = 'Café1Model!Test'
|
||||
envelope = factory.build_envelope(decomposed, KEY)
|
||||
|
||||
assert decrypt_password_envelope(envelope, KEY) == decomposed
|
||||
|
||||
|
||||
def test_loads_a_valid_key():
|
||||
assert load_envelope_key(factory.envelope_key_b64(KEY)) == KEY
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('raw_key', 'match'),
|
||||
[
|
||||
(None, 'not configured'),
|
||||
('', 'not configured'),
|
||||
('not base64 at all', 'valid base64'),
|
||||
(base64.b64encode(bytes(16)).decode(), '32 bytes'),
|
||||
(base64.b64encode(bytes(64)).decode(), '32 bytes'),
|
||||
],
|
||||
)
|
||||
def test_rejects_an_unusable_key(raw_key, match):
|
||||
"""There is no fallback: an unusable key never means the envelope is plaintext."""
|
||||
with pytest.raises(EnvelopeError, match=match):
|
||||
load_envelope_key(raw_key)
|
||||
|
||||
|
||||
@patch('laborious.utils.import_password.AESGCM')
|
||||
def test_a_generic_decrypt_failure_is_reported_as_could_not_be_opened(mock_aesgcm):
|
||||
"""Anything AESGCM raises besides InvalidTag (e.g. a malformed nonce) is not authentication."""
|
||||
mock_aesgcm.return_value.decrypt.side_effect = ValueError('bad nonce length')
|
||||
envelope = factory.build_envelope('Correct1Horse!Battery', KEY)
|
||||
|
||||
with pytest.raises(EnvelopeError, match='could not be opened'):
|
||||
decrypt_password_envelope(envelope, KEY)
|
||||
|
||||
|
||||
def test_a_non_utf8_plaintext_is_rejected():
|
||||
"""The envelope opened but what came out is not text — refuse rather than mangle it."""
|
||||
nonce = os.urandom(12)
|
||||
sealed = AESGCM(KEY).encrypt(nonce, b'\xff\xfe\xfd', None)
|
||||
envelope = base64.b64encode(nonce + sealed).decode('ascii')
|
||||
|
||||
with pytest.raises(EnvelopeError, match='not valid UTF-8'):
|
||||
decrypt_password_envelope(envelope, KEY)
|
||||
|
||||
|
||||
def test_no_secret_appears_in_the_error_text():
|
||||
password = 'Correct1Horse!Battery'
|
||||
envelope = factory.build_envelope(password, KEY)
|
||||
|
||||
with pytest.raises(EnvelopeError) as error:
|
||||
decrypt_password_envelope(envelope, OTHER_KEY)
|
||||
|
||||
message = str(error.value)
|
||||
assert password not in message
|
||||
assert envelope not in message
|
||||
assert factory.envelope_key_b64(OTHER_KEY) not in message
|
||||
19
tests/laborious/utils/test_time.py
Normal file
19
tests/laborious/utils/test_time.py
Normal file
@@ -0,0 +1,19 @@
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from laborious.utils.time import to_naive
|
||||
|
||||
|
||||
def test_to_naive_returns_naive_input_unchanged() -> None:
|
||||
moment = datetime(2024, 1, 1, 12, 30, 45)
|
||||
|
||||
assert to_naive(moment) == moment
|
||||
assert to_naive(moment).tzinfo is None
|
||||
|
||||
|
||||
def test_to_naive_strips_tzinfo_from_aware_input() -> None:
|
||||
moment = datetime(2024, 1, 1, 12, 30, 45, tzinfo=UTC)
|
||||
|
||||
result = to_naive(moment)
|
||||
|
||||
assert result.tzinfo is None
|
||||
assert result == datetime(2024, 1, 1, 12, 30, 45)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user