The full MNIST example

This is a minimal working example that implements the contract as is. Copy the first four files (the last one is for local verification only) and swap in your training code. The code below was actually built into an image and run against both a local MLflow and the platform.

mnist/
├── Dockerfile            # Image definition. ENTRYPOINT is just train.py
├── train.py              # Training entry point. Implements the three contract items (metrics · model · exit code)
├── predictor.py          # Serving wrapper + model structure. The file registered in the registry and started by mlflow models serve
├── experiment.json       # Example of the settings file the platform provides (for local verification)
└── make_sample_data.py   # Local verification only — imitates the data directory the platform prepares

How the contract maps to the code#

What it does not do: set_tracking_uri() · set_experiment() · creating a run · downloading data · register_model(). The platform does all of these.

Why predictor.py holds the model structure#

The serving container has no training code. mlflow models serve loads only the registered model, so model structure, preprocessing and inference must stand on their own in the single file predictor.py. That is why the dependency runs train.py → predictor.py. If you reverse it, training works but serving alone fails with ModuleNotFoundError.

Keeping preprocessing (preprocess) in one place has the same reason. If training and serving preprocess differently, you get the hardest bug to find: "training accuracy is high, but serving predictions are wrong."

Dockerfile#

# MNIST training container — contract sample
#
# The point is that there is nothing special: no platform-supplied script is added,
# and ENTRYPOINT simply points at the training script. The platform runs this ENTRYPOINT as is.

FROM python:3.11-slim

# CPU-only wheel — keeps the image at ~1.5GB.
# For GPU training, use `pip install torch mlflow==3.13.0 pillow pandas` instead of this line (6-8GB).
RUN pip install --no-cache-dir \
        --index-url https://download.pytorch.org/whl/cpu torch==2.12.0 \
    && pip install --no-cache-dir mlflow==3.13.0 pillow pandas

# train.py imports predictor.py, so the two must be in the same directory.
COPY train.py predictor.py /app/

# Without this, training progress logs are buffered and do not appear on screen in real time.
ENV PYTHONUNBUFFERED=1

# The location the platform mounts. If you ever use relative paths inside the container, this is the base.
WORKDIR /geo

ENTRYPOINT ["python", "/app/train.py"]

train.py#

# -*- coding: utf-8 -*-

"""MNIST training — sample implementation of the training container contract.

For the contract, this file does **only three things**:

1. Records metrics to MLflow (`step` = epoch, total epochs in the param `epochs`)
2. Logs the model once with `mlflow.pyfunc.log_model()`
3. Exits with 0 on success and non-zero on failure

What it does **not** do: set the MLflow address or token, create a run, name the experiment,
download the dataset, register in the model registry — the platform handles all of these.
"""

import json
import os
import random
import signal
import sys
from pathlib import Path

import mlflow
import pandas as pd
import torch
import torch.nn.functional as F
from mlflow.models import infer_signature
from PIL import Image
from torch.utils.data import DataLoader, Dataset

# Imported from the serving wrapper — model structure and preprocessing must be the same for training and serving.
from predictor import CLASS_NAMES, SmallCNN, preprocess

# ── Provided by the platform ───────────────────────────────────────────────
DATA_DIR = Path(os.environ.get("GEO_DATA_DIR", "/geo/dataset"))
WORK_DIR = Path(os.environ.get("GEO_WORK_DIR", "/geo/work"))
CONFIG_PATH = Path(os.environ.get("GEO_CONFIG", "/geo/experiment.json"))
DEVICE = "cpu" if os.environ.get("GEO_PARAM_DEVICE", "cpu") == "cpu" else "cuda:0"

# SIGTERM is a "stop request". Do not exit in the handler; only set a flag, then
# save a checkpoint at a safe point (epoch boundary) and leave. Any exit code is fine.
stop_requested = False


def _on_sigterm(*_):
    global stop_requested
    stop_requested = True
    print("-> SIGTERM: finishing this epoch, then saving and exiting", flush=True)


def hyperparameter(name: str, default):
    """Reads the ``GEO_HP_<name>`` environment variable as the type of the default."""

    raw = os.environ.get(f"GEO_HP_{name.upper()}")
    if raw is None or raw == "":
        return default
    return type(default)(raw)


def load_config() -> dict:
    """``experiment.json`` — must work without it (convenient for local verification)."""

    try:
        return json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return {}


class ManifestDataset(Dataset):
    """A dataset that reads the image files ``manifest.json`` points to, as is.

    The platform places the data in a local directory **before** training starts.
    No download and no storage credentials are needed — just open the files.

    Labels come from each entry's ``meta.label`` (for classification datasets).
    """

    def __init__(self, items: list[dict]) -> None:
        self.items = items

    def __len__(self) -> int:
        return len(self.items)

    def __getitem__(self, index: int):
        item = self.items[index]
        image = Image.open(DATA_DIR / item["path"])
        # Use only predictor.preprocess for preprocessing — if training and serving differ,
        # you get the hardest bug to find: good training accuracy, wrong serving predictions.
        tensor = preprocess(image)[0]
        return tensor, CLASS_NAMES.index(str(item["meta"]["label"]))


def load_manifest() -> list[dict]:
    """List of image files (only those with labels)."""

    manifest = json.loads((DATA_DIR / "manifest.json").read_text(encoding="utf-8"))
    items = [
        f
        for f in manifest.get("files", [])
        if f.get("kind") == "image" and (f.get("meta") or {}).get("label") is not None
    ]
    if not items:
        raise SystemExit("no labeled images in manifest.json")
    return items


def split_items(items: list[dict], config: dict) -> tuple[list[dict], list[dict]]:
    """Splits train/val as ``experiment.json`` instructs.

    **Splitting is the image's responsibility.** The platform does not split the files; it only gives the ratio.
    """

    spec = config.get("split") or {}
    percent = int(spec.get("train_percent", 80))
    shuffled = list(items)
    random.Random(int(spec.get("seed", 0))).shuffle(shuffled)
    cut = max(1, len(shuffled) * percent // 100)
    return shuffled[:cut], shuffled[cut:] or shuffled[cut - 1 :]


def run_epoch(model, loader, optimizer=None) -> tuple[float, float]:
    """(mean loss, accuracy). Evaluation mode when there is no ``optimizer``."""

    training = optimizer is not None
    model.train(training)
    total_loss, correct, seen = 0.0, 0, 0
    with torch.set_grad_enabled(training):
        for images, labels in loader:
            images, labels = images.to(DEVICE), labels.to(DEVICE)
            logits = model(images)
            loss = F.cross_entropy(logits, labels)
            if training:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
            total_loss += loss.item() * len(labels)
            correct += int((logits.argmax(dim=-1) == labels).sum())
            seen += len(labels)
    return total_loss / seen, correct / seen


def build_signature(model, sample_item: dict):
    """(signature, input_example) — builds the serving API contract from one real input.

    The column names and types become the inference request format as is, so build
    exactly the shape serving will receive (a base64 image string) here. ``input_example``
    makes MLflow run one prediction at logging time, catching errors during training
    that would otherwise surface only in serving.
    """

    import base64

    encoded = base64.b64encode((DATA_DIR / sample_item["path"]).read_bytes()).decode()
    model_input = pd.DataFrame({"image_b64": [encoded]})
    with torch.no_grad():
        probs = model.cpu()(preprocess(Image.open(DATA_DIR / sample_item["path"])))
        probs = probs.softmax(dim=-1)[0]
    index = int(probs.argmax())
    model_output = pd.DataFrame(
        [{"label": CLASS_NAMES[index], "confidence": float(probs[index])}]
    )
    return infer_signature(model_input, model_output), model_input


def main() -> int:
    signal.signal(signal.SIGTERM, _on_sigterm)
    WORK_DIR.mkdir(parents=True, exist_ok=True)

    config = load_config()
    epochs = hyperparameter("epochs", 5)
    batch_size = hyperparameter("batch", 64)
    learning_rate = hyperparameter("lr", 0.001)

    train_items, val_items = split_items(load_manifest(), config)
    print(f"-> train {len(train_items)} / val {len(val_items)} @ {DEVICE}", flush=True)

    train_loader = DataLoader(
        ManifestDataset(train_items), batch_size=batch_size, shuffle=True
    )
    val_loader = DataLoader(ManifestDataset(val_items), batch_size=batch_size)

    model = SmallCNN(len(CLASS_NAMES)).to(DEVICE)
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    checkpoint = WORK_DIR / "best.pt"
    best_accuracy = 0.0

    # Connection details (URI, token) and the run are already set via environment variables.
    # Do not call set_tracking_uri() or set_experiment() — the experiment mismatches and the run dies at once.
    with mlflow.start_run():
        # Total epochs as a param. Progress on the training screen is calculated from this value.
        mlflow.log_params(
            {"epochs": epochs, "batch": batch_size, "lr": learning_rate}
        )

        for epoch in range(epochs):
            train_loss, train_accuracy = run_epoch(model, train_loader, optimizer)
            _, val_accuracy = run_epoch(model, val_loader)
            # MLflow is the single source of metrics. step is the epoch — the x axis on screen.
            mlflow.log_metrics(
                {
                    "train/loss": train_loss,
                    "train/accuracy": train_accuracy,
                    "eval/accuracy": val_accuracy,
                },
                step=epoch,
            )
            print(
                f"-> epoch {epoch + 1}/{epochs} "
                f"loss={train_loss:.4f} val_acc={val_accuracy:.4f}",
                flush=True,
            )

            if val_accuracy >= best_accuracy:
                best_accuracy = val_accuracy
                torch.save(model.state_dict(), checkpoint)

            if stop_requested:  # stop request — saving is already done
                break

        if not checkpoint.exists():  # if not even one epoch ran, there is no model to keep
            raise SystemExit("no checkpoint — training did not make progress")

        model.load_state_dict(torch.load(checkpoint, map_location=DEVICE))
        signature, input_example = build_signature(model, val_items[0])

        # Always keep the model as an MLflow model. With only a raw .pt there is no MLmodel,
        # so it does not start with `mlflow models serve` and no serving image can be built.
        mlflow.pyfunc.log_model(
            name="model",
            python_model=str(Path(__file__).resolve().parent / "predictor.py"),
            artifacts={"weights": str(checkpoint)},
            signature=signature,
            input_example=input_example,
            pip_requirements=[
                f"torch=={torch.__version__.split('+')[0]}",
                "pillow",
                "pandas",
                "mlflow==3.13.0",  # platform-pinned version
            ],
            metadata={
                "input_kind": "image_b64",  # decides the inference console's input widget
                "class_names": {str(i): name for i, name in enumerate(CLASS_NAMES)},
            },
        )
        # Do not call registry registration (register_model) — the platform does it.

        mlflow.log_metric("eval/best_accuracy", best_accuracy)
        print(f"-> done. best val accuracy = {best_accuracy:.4f}", flush=True)

    return 0


if __name__ == "__main__":
    sys.exit(main())

predictor.py#

# -*- coding: utf-8 -*-

"""Serving wrapper — the file registered in the model registry and started by `mlflow models serve`.

**The serving container has no training code (`train.py`).** So this one file must
handle the model structure definition + weight loading + inference on its own (models-from-code).
That is also why `train.py` imports ``SmallCNN`` from this file and not the other way
around — the reverse would make serving need the training code.

Input/output contract (= `signature`):

- Input: a single ``image_b64`` column. PNG/JPEG bytes encoded as a base64 string.
  Spelling out a 640x640 tensor as JSON numbers makes a request tens of MB, so image models use base64.
- Output: two columns, ``label`` (string) and ``confidence`` (float between 0 and 1).
"""

import base64
import io

import mlflow
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image

#: Index → display name. For MNIST the index equals the digit, but in general
#: list the classes in exactly the order used in training.
CLASS_NAMES = [str(i) for i in range(10)]

#: Training and serving must use the same preprocessing. Changing these constants changes train.py too.
IMAGE_SIZE = 28
NORM_MEAN = 0.1307
NORM_STD = 0.3081


class SmallCNN(nn.Module):
    """Small CNN for 28x28 grayscale input (it is a sample, so the structure itself does not matter)."""

    def __init__(self, num_classes: int = 10) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.fc1 = nn.Linear(32 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.flatten(1)
        x = F.relu(self.fc1(x))
        return self.fc2(x)


def preprocess(image: Image.Image) -> torch.Tensor:
    """PIL image → (1, 1, 28, 28) tensor. **Identical** preprocessing to training."""

    image = image.convert("L").resize((IMAGE_SIZE, IMAGE_SIZE))
    buffer = bytearray(image.tobytes())  # writable copy — required by frombuffer
    tensor = torch.frombuffer(buffer, dtype=torch.uint8).float() / 255.0
    tensor = tensor.reshape(1, 1, IMAGE_SIZE, IMAGE_SIZE)
    return (tensor - NORM_MEAN) / NORM_STD


class MnistPredictor(mlflow.pyfunc.PythonModel):
    """The model the scoring server loads."""

    def load_context(self, context):
        """Once at container start. The weights come along via ``artifacts``."""

        self.model = SmallCNN(len(CLASS_NAMES))
        state = torch.load(context.artifacts["weights"], map_location="cpu")
        self.model.load_state_dict(state)
        self.model.eval()

    def predict(self, context, model_input, params=None):
        """Each request brings in one DataFrame and gets one back.

        The return value must be **JSON-serializable** — DataFrame / ndarray / list / dict.

        At logging time an "Add type hints to the `predict` method" warning appears; you can ignore it.
        MLflow's type-hint-based validation supports only the ``list[...]`` form, and this model
        uses a DataFrame contract (`signature`), so adding hints would only add more warnings.
        """

        rows = []
        for encoded in model_input["image_b64"]:
            image = Image.open(io.BytesIO(base64.b64decode(encoded)))
            with torch.no_grad():
                probs = self.model(preprocess(image)).softmax(dim=-1)[0]
            index = int(probs.argmax())
            rows.append(
                {"label": CLASS_NAMES[index], "confidence": float(probs[index])}
            )
        return pd.DataFrame(rows)


# The last line of models-from-code. Without this call the model does not load.
mlflow.models.set_model(MnistPredictor())

The order of CLASS_NAMES defines what each model output index means. It must match the order used in training.

experiment.json#

On the platform this file is created automatically. For local verification, mount it at /geo/experiment.json and the script reads split and the rest from it (without it, the script runs on defaults).

{
  "experiment_id": "exp-local-mnist",
  "name": "mnist-local-20260801-1030",
  "tenant": "DEMO",
  "created_at": "2026-08-01T10:30:12Z",

  "task": "classification",
  "framework": "mnist-example",
  "model": "small-cnn",
  "classes": ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"],

  "split": { "method": "random", "train_percent": 80, "seed": 0 },
  "hyperparameters": { "epochs": 5, "batch": 64, "lr": 0.001 },

  "device": "cpu",
  "gpu": 0,

  "evaluation": { "benchmark": true, "speed_test": false },
  "export": { "onnx": false, "tensorrt": false },
  "serving": { "runtime": "cpu" },
  "register": { "enabled": true, "model_name": "mnist" },
  "pretrained": { "source": "catalog" },

  "dataset": {
    "id": "ds-local-mnist",
    "name": "MNIST (local verification)",
    "file_count": 2000,
    "path": "/geo/dataset",
    "manifest": "/geo/dataset/manifest.json"
  },
  "paths": { "data_dir": "/geo/dataset", "work_dir": "/geo/work" },
  "mlflow": { "experiment_name": "mnist-local-20260801-1030", "run_id": "0000000000000000000000000000abcd" }
}

make_sample_data.py (local verification only)#

# -*- coding: utf-8 -*-

"""Data generator for local verification — **not needed on the platform**.

In a real run, the platform places the dataset in `GEO_DATA_DIR` beforehand
and writes `manifest.json` as well. This script imitates that state locally
so you can check your image yourself before handing it over.

    python make_sample_data.py ./sample-data              # torchvision MNIST (recommended)
    python make_sample_data.py ./sample-data --synthetic  # synthetic images, no network needed
"""

import argparse
import json
import random
from pathlib import Path

from PIL import Image, ImageDraw


def write_manifest(root: Path, records: list[tuple[str, str]]) -> None:
    """Writes a ``manifest.json`` in the same format as the platform."""

    manifest = {
        "dataset": {"id": "ds-local-mnist", "name": "MNIST (local verification)"},
        "files": [
            {"path": path, "kind": "image", "meta": {"label": label}}
            for path, label in records
        ],
    }
    (root / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(f"-> {root/'manifest.json'} ({len(records)} files)")


def from_torchvision(root: Path, count: int) -> list[tuple[str, str]]:
    from torchvision import datasets

    dataset = datasets.MNIST(root=str(root / ".cache"), train=True, download=True)
    records = []
    for index in range(min(count, len(dataset))):
        image, label = dataset[index]
        name = f"images/{index:05d}.png"
        (root / "images").mkdir(parents=True, exist_ok=True)
        image.save(root / name)
        records.append((name, str(label)))
    return records


def synthetic(root: Path, count: int) -> list[tuple[str, str]]:
    """28x28 images with a digit drawn in. For checking that training runs."""

    (root / "images").mkdir(parents=True, exist_ok=True)
    rng = random.Random(0)
    records = []
    for index in range(count):
        label = rng.randrange(10)
        image = Image.new("L", (28, 28), color=0)
        draw = ImageDraw.Draw(image)
        draw.text((9, 8), str(label), fill=255)
        name = f"images/{index:05d}.png"
        image.save(root / name)
        records.append((name, str(label)))
    return records


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("output", type=Path, help="directory to use as GEO_DATA_DIR")
    parser.add_argument("--count", type=int, default=2000)
    parser.add_argument(
        "--synthetic", action="store_true", help="use synthetic images without torchvision"
    )
    args = parser.parse_args()

    args.output.mkdir(parents=True, exist_ok=True)
    if args.synthetic:
        records = synthetic(args.output, args.count)
    else:
        records = from_torchvision(args.output, args.count)
    write_manifest(args.output, records)


if __name__ == "__main__":
    main()

Next: Verify locally first