Quickstart

uv venv --python 3.12
source .venv/bin/activate

uv pip install --torch-backend=cpu belfort-ml belfort-ml-client huggingface_hub safetensors pyarrow pillow

pip install --extra-index-url https://download.pytorch.org/whl/cpu belfort-ml belfort-ml-client huggingface_hub safetensors pyarrow pillow

export BELFORT_API_KEY=blf_live_...

"""The model and the test images, both downloaded from Hugging Face."""

import io

import numpy as np
import pyarrow.parquet as pq
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from safetensors.torch import load_file
from torch import nn


class MLP(nn.Module):
    """dacorvo/mnist-mlp: a 28x28 image, flattened, to ten digit scores."""

    def __init__(self) -> None:
        super().__init__()
        self.input_layer = nn.Linear(784, 256)
        self.mid_layer = nn.Linear(256, 256)
        self.output_layer = nn.Linear(256, 10)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = torch.relu(self.input_layer(x))
        x = torch.relu(self.mid_layer(x))
        return self.output_layer(x)


def load_model() -> MLP:
    weights = hf_hub_download(
        "dacorvo/mnist-mlp",
        "model.safetensors",
        revision="7ed1e96d182395fd7380ea8495b9eb2364ce9dc1",
    )
    net = MLP()
    net.load_state_dict(load_file(weights))
    return net.eval()


def test_images() -> tuple[torch.Tensor, list[int]]:
    """The 10,000 MNIST test images, normalized and flattened as the model expects."""
    path = hf_hub_download(
        "ylecun/mnist",
        "mnist/test-00000-of-00001.parquet",
        repo_type="dataset",
        revision="77f3279092a1c1579b2250db8eafed0ad422088c",
    )
    table = pq.read_table(path)
    pngs = table["image"].to_pylist()
    pixels = np.stack([np.array(Image.open(io.BytesIO(png["bytes"]))) for png in pngs])
    pixels = pixels.astype(np.float32) / 255
    images = torch.from_numpy((pixels - 0.1307) / 0.3081).flatten(1)
    return images, table["label"].to_pylist()

"""The model provider: compile the model, then serve it to one end client.

    uv run provider.py compile   # compile the model on Belfort
    uv run provider.py serve     # mint a client token and start a GPU deploy
    uv run provider.py serve --dev   # testing only: the gateway chooses the seed

Both are safe to run again: a model already compiled from the same MLIR is
reused, not compiled again.
"""

import hashlib
import sys
from datetime import timedelta
from pathlib import Path

import belfort_ml
from mnist import load_model, test_images

NAME = "quickstart_mnist"
# client.py on this machine reads the token from here; it exists while serve
# runs.
TOKEN_FILE = Path(".client-token")


def compiled_model(provider: belfort_ml.Provider) -> belfort_ml.Model:
    images, _ = test_images()
    # The last 512 test images; client.py sends the first ones.
    calibration = list(images[-512:].split(1))
    artifacts = belfort_ml.compile(
        load_model(), calibration[0], NAME, "./build", calibration_inputs=calibration
    )
    return provider.models.get(
        name=NAME, create_from=artifacts, timeout=1800, use_composite_relu=False
    )


def hand_over(token: str, seed: str | None = None) -> None:
    """Write the token, and a developer-mode seed where client.py looks for it."""
    TOKEN_FILE.touch(mode=0o600)
    TOKEN_FILE.write_text(token)
    export = f"  export BELFORT_CLIENT_TOKEN={token}\n"
    if seed is not None:
        digest = hashlib.sha256(token.encode()).hexdigest()[:12]
        seed_file = Path(f".secret-seed-{digest}")
        seed_file.touch(mode=0o600)
        seed_file.write_text(seed)
        export += f"  export BELFORT_CLIENT_SECRET_SEED={seed}\n"
    print(
        f"\nStart client.py now. On this machine it reads the token from"
        f" {TOKEN_FILE}.\nOn another machine, run this before client.py:\n{export}",
        flush=True,
    )


def serve(model: belfort_ml.Model, dev: bool = False) -> None:
    TOKEN_FILE.unlink(missing_ok=True)  # a killed earlier run may have left one
    try:
        token = None if dev else model.create_client(label="quickstart")
        if token is not None:
            hand_over(token.token)
            # The GPU bills from here, so start it only once the client's keys
            # are in.
            model.wait_for_keys(token)
            print("Keys uploaded. Starting the deploy; Ctrl-C stops it.", flush=True)
        with model.deploy(
            key_set=token,
            max_lifetime=timedelta(minutes=30),
            timeout=1800,
            dev_mode=dev,
        ) as deploy:
            if dev:
                # The gateway supplies the developer seed and a client token.
                assert deploy.client_token is not None
                hand_over(deploy.client_token, deploy.secret_seed)
            input("Deploy ready. Press Enter to stop it...")
    finally:
        TOKEN_FILE.unlink(missing_ok=True)


if __name__ == "__main__":
    command = sys.argv[1] if len(sys.argv) > 1 else ""
    dev = sys.argv[2:] == ["--dev"]
    if command not in ("compile", "serve") or (sys.argv[2:] and not dev):
        sys.exit(__doc__)
    with belfort_ml.Provider() as provider:  # reads BELFORT_API_KEY
        model = compiled_model(provider)
        if command == "compile":
            print(f"{NAME} is compiled ({model.id})")
        else:
            serve(model, dev)

uv run provider.py compile

uv run provider.py serve

"""The end client: encrypt MNIST test images, have the deploy classify them,
and decrypt the scores.

    uv run client.py [COUNT]

The secret key derives from a seed this script keeps in a local file per
token, so running it again on the same token reuses the key pair.
"""

import hashlib
import os
import sys
from pathlib import Path

from belfort_ml_client import Client
from mnist import test_images

token_file = Path(".client-token")
if token_file.exists():
    token = token_file.read_text().strip()
elif not (token := os.environ.get("BELFORT_CLIENT_TOKEN")):
    sys.exit("Run `uv run provider.py serve` here first, or set BELFORT_CLIENT_TOKEN.")

seed_file = Path(f".secret-seed-{hashlib.sha256(token.encode()).hexdigest()[:12]}")
seed = bytes.fromhex(seed_file.read_text()) if seed_file.exists() else None
images, labels = test_images()
count = int(sys.argv[1]) if len(sys.argv) > 1 else 10

with Client(token, secret_seed=seed) as client:
    client.prepare()
    if seed is None:
        seed_file.touch(mode=0o600)
        seed_file.write_text(client.secret_seed.hex())
    correct = 0
    for image, label in zip(images[:count], labels[:count]):
        answer = client.infer([image.tolist()])
        scores = answer["values"]
        predicted = scores.index(max(scores))
        correct += predicted == label
        print(f"digit {label}: predicted {predicted} in {answer['latency_ms']} ms")
    print(f"{correct} of {count} correct")

uv run client.py

digit 7: predicted 7 in 287 ms
digit 2: predicted 2 in 189 ms
...

uv run provider.py serve --dev