03. Reproducible training & registry

Training and evaluation as one-shot containers from pinned images: they log to the self-hosted MLflow container, register a model version with reconstructable lineage, and leave a results-DB record — the same images become Azure Container Apps Jobs in Part II.

Outcome

Training and evaluation run as one-shot containers built from src/train_job/Dockerfile, log to the self-hosted MLflow, and produce a registered model version whose lineage (code digest, tracked dataset, params, metrics) can be reconstructed from MLflow alone. Re-running the same inputs yields an equivalent registered artifact, and every run is also visible in the results DB.

Locally these containers are Compose services triggered through the dashboard and runner APIs (chapter 02); in Part II the same images run as Azure Container Apps Jobs (chapter 11). The execution shape is identical in both worlds: ephemeral, image-pinned, logging to MLflow and writing a results-DB row.

What “model” means here: anything with learned weights — a scikit-learn/ XGBoost estimator and a fine-tuned transformer classifier — travels the identical registry path. An LLM app (prompt + config, no weights trained) is packaged differently but registered in the same registry (chapter 07).

Design: training and evaluation are one-shot jobs

A training run executes a pinned image exactly once: the container starts, runs its entrypoint to completion, and exits. Locally Compose owns that lifecycle and the dashboard/runner APIs trigger it; in Part II an ACA Job definition does, started manually, on a schedule, or through the same trigger API (chapter 11). The training script:

  1. Builds an MLflow dataset from its source with mlflow.data, capturing the source location, a content digest, and the schema, not just “the latest table.”
  2. Starts an MLflow run, logging params, the code image digest, and the dataset via mlflow.log_input(dataset, context="training").
  3. Trains the candidate and logs training-run metrics and artifacts.
  4. Registers the model as a new version in the registry served by the same self-hosted server.
  5. Writes a results-DB record (name='train:<model>', status, output carrying the MLflow run ID and registered version).
import mlflow

dataset = tracked_dataset(raw_df, source=data_source, name="wine-quality-white")
with mlflow.start_run() as run:
    mlflow.log_input(dataset, context="training")
    mlflow.log_params(params)
    mlflow.set_tag("code.image_digest", image_digest)
    model = ElasticNet(alpha=alpha, l1_ratio=l1_ratio).fit(x_train, y_train)
    mlflow.log_metrics(evaluate(model, x_test, y_test))
    mlflow.sklearn.log_model(model, name="model",
                             registered_model_name="wine-quality")

The candidate then runs through a separate eval job from the same image. evaluate.py loads models:/<name>/<version>, scores a held-out dataset, writes its own MLflow and results-DB records, and tags the evaluation run with the exact model URI and pass/fail decision. A threshold rejection raises inside both run contexts, so MLflow and the results DB record FAILURE and the process exits non-zero.

A version can therefore exist in the registry yet remain unpromotable. demo/promote.py searches the evaluation experiment for a passing record tied to that exact model URI before it moves the production alias. The alias flip and exact-version serving redeploy are chapter 05’s subject, with the cross-environment rules written in chapter 08.

Build in projects/ml-platform/

projects/ml-platform/
├── src/train_job/
│   ├── Dockerfile              # pinned base + training deps + ml_platform
│   ├── requirements.txt        # mlflow pinned to the server's version
│   ├── train.py                # entrypoint: dataset → run → register
│   ├── evaluate.py             # held-out metrics; separate eval job
│   └── register_llm.py         # LLM artifacts into the same registry (Ch 07)
└── src/ml_platform/common/
    ├── mlflow_client.py        # configure_mlflow(): tracking + registry URI
    ├── datasets.py             # load_csv + tracked_dataset (mlflow.data)
    ├── schemas.py              # Pandera contract, validated before training
    └── results.py              # record_run(): results-DB row per job

Locally the training job is the train Compose service, built once at stack startup and configured entirely by its environment block in demo/docker-compose.yml: MLFLOW_TRACKING_URI=http://mlflow:5000 reaches the MLflow container by service-name DNS, PGHOST/PGUSER/RESULTS_DB point at the Postgres container’s results database, and TRIGGERED_BY names the caller. IMAGE_DIGEST carries the honest placeholder demo-local, because an image built on a laptop never acquires a registry digest.

Part II runs the identical image as an ACA Job under the id-jobs-train managed identity. Same variables, different answers: the tracking URI becomes the MLflow App’s URL, Postgres auth switches from the demo password to a short-lived Entra access token fetched at runtime by DefaultAzureCredential, and Terraform sets IMAGE_DIGEST to the digest CI pinned, so the tag on each run names the exact code that produced it. The port itself is chapter 11’s work.

results.record_run(...) degrades gracefully rather than hard-wiring the database: with PGHOST unset it is a no-op wrapper, so bare python train.py on a laptop still trains and registers; inside the demo stack PGHOST is set, so every bootstrap and triggered run lands a results row from day one. Chapter 04 formalizes that table and adds the query helpers.

How the pieces connect

The entrypoint is thin because the reusable machinery lives in common/, shared with every later job (eval, batch, LLM registration):

  • common/mlflow_client.py: configure_mlflow(experiment) points both the tracking and registry URIs at the one self-hosted server and raises if MLFLOW_TRACKING_URI is unset. One call, and mlflow.* and the registry agree.
  • common/datasets.py: load_csv(source, delimiter=…) reads the source (URL or path); tracked_dataset(df, source=…, name=…, targets=…) wraps mlflow.data.from_pandas so the run captures source + content digest + schema, not “whatever the table held today.”
  • common/schemas.py: validate(df) runs the Pandera wine_quality_schema (11 physicochemical floats + integer quality) at the boundary, before training. Bad data fails the run early instead of poisoning a registered version.
  • common/results.py: record_run("train:<model>") is a context manager that opens a results-DB row and closes it SUCCESS, or marks it FAILURE and re-raises on error. It yields a mutable output dict; train.py fills it with the MLflow run id and registered version, so the operational record points straight at the lineage. The local runner passes its execution id through RESULTS_RUN_ID, making trigger response and results row share one id.

train.py sequences configure_mlflow → load_csv → validate → tracked_dataset → train_test_split → inside record_run(...) and mlflow.start_run(): log the dataset, parameters, code digest, metrics, and registered model. evaluate.py loads the exact registered version and records eval.model_uri, eval.max_rmse, and eval.passed. A failed threshold raises an ordinary exception before the context managers close, which makes both run records fail honestly instead of returning a non-zero code from an otherwise successful context.

Both execution adapters expose evaluation as a first-class job. The local runner accepts POST /api/runs/eval/trigger; Terraform creates a separate manual ACA Job from the same train image with python evaluate.py as its command. Train and eval executions may overlap because each request creates an independent runner subprocess or ACA Job execution.

register_llm.py shares this image and registry path for LLM artifacts; chapter 07 builds it.

Golden-path position & acceptance evidence

This chapter builds the tracked dataset → train job → registered version → eval job segment of the golden path (chapter 01).

The registered version number is the canonical model identity: every later stage refers to a model only through it. Reading the platform’s records in order gives the provenance chain registered version → evaluation record → Git tag → deployed image digest, and each link is written down somewhere queryable rather than reconstructed from memory.

Acceptance evidence:

  • A training run appears in MLflow with its dataset digest, code image digest, params, and metrics; the model is a registered version (models:/wine-quality/<n>).
  • The same inputs re-run produce an equivalent registered artifact, and the run is queryable in the results DB by name and status.
  • A passing eval run names the exact models:/<name>/<version> candidate and records its threshold; a rejected candidate ends FAILURE in both execution and operational state.
  • Promotion of a version with no passing evaluation is rejected before the registry alias or serving deployment changes.

Extensions (deferred from the MVP)

Deferred capability Current baseline
Rich stage-identity chain (many lineage fields) Dataset digest + code image digest + version
Distributed / multi-GPU training Single-node job (see chapter 14)
Full evaluator evidence package Held-out metrics + threshold check
Scheduled retrains as a workflow Manual triggers: dashboard button and runner API now, ACA Job schedules in Part II (see chapter 04)

Next: 04 — Results DB & batch workflows turns run-recording into the operational backbone and adds scheduling and batch fan-out.

Back to top