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

QTZPOC-20

Browse Source 
QTZPOC-20: Implement explicit joblib model flavor and enhance model import workflow with extensive E2E tests.

This commit introduces a first-class joblib model flavor for direct loading and saving of Python pickle artifacts, including custom code within the artifact's code/ directory, bypassing MLflow's flavor loaders. The model import workflow's status database configuration is now streamlined to default to the POSTGRES_* server, with explicit IMPORT_STATUS_DB_* overrides.

Comprehensive end-to-end tests have been added for the model import workflow, covering success paths, gate rejections, digest enforcement, provisioning failures, status record integrity, cleanup, and retry/recovery scenarios using a real-clock environment. New E2
This commit is contained in:
Vitor Pimentel dos Santos
2026-08-26 15:45:33 -03:00
parent dac4f1718f
commit e95e65b72b
32 changed files with 4247 additions and 1040 deletions
Download Patch File Download Diff File Expand all files Collapse all files

112
docs/model-flavors.md Normal file
View File

@@ -0,0 +1,112 @@
# 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.
## Accepted values
| Flavor | Read | Write (retrain) | How it loads |
|---|---|---|---|
| `sklearn` | yes | yes | `mlflow.sklearn.load_model(<uri>)` |
| `pyfunc` | yes | yes | `mlflow.pyfunc.load_model(<uri>)`; the wrapper path also unwraps `._model_impl.python_model` |
| `pytorch` | yes | yes | `mlflow.pytorch.load_model(<uri>)` |
| `joblib` | yes | yes | downloads the artifact and `joblib.load`s its pickle — **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 four values govern reading and writing.
**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.
## The `joblib` flavor
Some models the pipeline produces do not load through an MLflow flavor loader even though the
artifact is intact: the `.pkl` is in the compressed form `joblib.dump` writes, which MLflow's
`pickle.load` does not read. Declaring `joblib` is how such a model is served.
**Reading.** Resolve the Production run, `client.download_artifacts(run_id, <artifact>)`, take the
**first top-level `.pkl` in name order**, and `joblib.load` it. The listing does not recurse, so a
nested wrapper artifact (`artifacts/*.pkl`) is never deserialized in place of the model. No
`MLmodel` is parsed and no `models:/<name>/production` URI is built.
**The artifact's `code/` goes on `sys.path` first.** A model whose class lives inside the artifact
would otherwise raise `ModuleNotFoundError` at unpickle time. Before `joblib.load`, the flavor
prepends `<artifact>/code` to `sys.path` when that directory exists — the directory itself, matching
MLflow's own convention where `code/utils/…` imports as the package `utils`. The entry is not
duplicated and is **not** removed after the load, because the deserialized object may import
lazily. A missing `code/` is not an error, and no extra download is performed: `download_artifacts`
already brings `code/` along.
> `sys.path` is process-global. Two different models that embed a package with the same name
> (`utils/`, typically) and are served by the same worker resolve to whichever entered first — the
> second silently runs the first one's code. This is pre-existing behaviour of MLflow's own pyfunc
> path with the same artifact layout; the mitigation, if it ever bites, is one worker per model.
**Writing / retraining.** `joblib` is a write flavor too, so a joblib model retrains. `log_model`
writes `joblib.dump(model, 'model.pkl', compress=3)` under the artifact path and re-logs the source
artifact's `code/` to `<artifact>/code`, so the retrained model stays deserializable. Registration
and promotion to `Production` are unchanged.
**The written artifact has no `MLmodel` manifest.** Only this runtime's `joblib` flavor reads it:
`mlflow.sklearn.load_model` and `models:/<name>/production` both fail on it. A model retrained in
joblib **stays** joblib — `predict_flavor` needs no edit after a retrain, and no consumer outside
this runtime (the MLflow UI, MLflow serving, another service) recognizes that artifact as a model.
**What `joblib` is for.** A pickled object exposing `predict(data)` — and `fit(data)` if it is
retrained. It is **not** for a bare pyfunc `PythonModel` wrapper: `get_prediction_data` and
`get_cached_operation` call `predict(data)` for every non-pyfunc flavor, so a wrapper expecting
`predict(context, model_input)` raises `TypeError` at predict and retrain time.
**Observability.** A joblib load is an ordinary load: `MODEL_READ_LAG` / `MODEL_READ_COUNT` on
success and `MODEL_READ_ERROR_COUNT` on failure, under the base `operation_type`
(`load_predict_model` / `load_transform_model`). A joblib write emits the usual `MODEL_WRITE_*`
under `log_model`. No suffixed operation type, no new metric, no new label.
## Switching a model to `joblib`
The runtime does not migrate documents. Editing the `models` collection is an operator step, and the
order matters:
1. **Inventory first, before the deploy.** Collect the models that fail in their configured flavor
loader — run predict/transform and note which ones raise. There is no automatic discovery
afterwards: once deployed, such a model simply fails, and the failure names the flavor loader,
not the fix.
2. **Edit the document.** Set `model_config.predict_flavor: 'joblib'` and/or
`model_config.transform_flavor: 'joblib'` for each model in the inventory. Nothing else changes;
models already on a working flavor need no edit and behave identically.
3. **Imported models arrive ready.** The `import_model` workflow writes
`predict_flavor: 'joblib'` itself — see [`model-import.md`](model-import.md) § 12.
4. **Check the first retrain.** In MLflow, confirm the new run carries
`prediction_model/model.pkl` **and** `prediction_model/code/` when the source model had a
`code/` directory, and that the promoted version loads on the next predict.
5. **Rolling back is not just a revert.** Reverting the code leaves any artifact a joblib retrain
already wrote in MLflow: that Production version stops being loadable, so the rollback also means
promoting the previous version back in the registry, and setting the documents' flavors back —
manual, in reverse order.
## Known gap: how the run id is resolved
`get_model_run_id` resolves the Production run with `source.split('/')[2]` on the registered
version's `source`. **This has not been confirmed against the real tracking server.** No artifact
scheme reproduced locally puts the run id in that position (a file store yields `''`, S3 yields the
bucket, `mlflow-artifacts:/` yields the experiment id), yet predict works in production — so the
real `source` has a shape none of the local probes produced.
This is not specific to `joblib`: every flavor that resolves a run goes through the same line, and
the joblib read path depends on it. Confirm it with a **read-only** query against the tracking
server before relying on it, ideally before the next deploy.

59
docs/model-import.md
View File

@@ -88,14 +88,18 @@ is at fault.
| `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_HOST` | `localhost` | the BFF database holding the import log |
| `IMPORT_STATUS_DB_PORT` | `5432` | |
| `IMPORT_STATUS_DB_NAME` | `sientia-core-mlops-bff` | |
| `IMPORT_STATUS_DB_USER` | `mlops_bff_user` | |
| `IMPORT_STATUS_DB_PASSWORD` | empty | |
| `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 |
`IMPORT_STATUS_DB_*` is a **second** Postgres connection, independent of the one Laborious already
holds for predictions. The import writes to that database and no other, and to one table in it.
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.
@@ -324,10 +328,11 @@ A `COMPLETED` import is **not** a servable model. Stated plainly:
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 a **read-only**
flavor. It is what the artifact actually is, and it is enough to load and predict, but nothing in this
platform writes that flavor — so an imported model cannot be retrained until a write flavor is
configured for it.
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
@@ -347,3 +352,35 @@ 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. 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).