Deployment Guide

This guide covers the full deployment story for AITune-tuned models: saving a tuned model to a checkpoint, loading it in production, and optionally serving it as an OpenAI-compatible HTTP endpoint via NVIDIA Dynamo.

Save a Tuned Model

Basic Save

1
import aitune.torch as ait
2
3
# After tuning
4
ait.save(model, "checkpoints/model.ait")

This creates:

• checkpoints/model.ait: Compressed checkpoint with tuned modules
• checkpoints/model_sha256_sums.txt: SHA256 checksums
• checkpoints/model/: Decompressed artifacts (after first load)

With Custom Storage

1
from aitune.torch import LocalTorchStorage
2
3
storage = LocalTorchStorage(
4
    base_folder="production/models",
5
    remove_checkpoint_after_tune=False,
6
)
7
8
ait.save(model, "model_v2.ait", storage=storage)

Load in Production

Basic Load

1
import aitune.torch as ait
2
3
model = YourModel()
4
model.eval()
5
model.to("cuda")
6
7
ait.load(model, "checkpoints/model.ait")
8
9
output = model(input_data)

With Custom Storage

1
from aitune.torch import LocalTorchStorage
2
3
storage = LocalTorchStorage(base_folder="production/models")
4
ait.load(model, "model.ait", storage=storage)

Loading Process

1. First load — decompresses .ait file, extracts artifacts, verifies checksums, loads backend and weights. Slower due to decompression.
2. Subsequent loads — uses decompressed files from checkpoints/, skips decompression. Faster startup.

Serve with Dynamo Worker

After loading a tuned model, you can expose it as an OpenAI-compatible HTTP endpoint using AITune’s Dynamo integration. The worker registers the model with the Dynamo HTTP frontend, deserializes incoming requests, packs inference results into the Dynamo wire format, and blocks until SIGTERM/SIGINT.

Prerequisites

Install the Dynamo extra:

$
uv pip install "aitune[dynamo]"

For local development without etcd/NATS, set DYN_DISCOVERY_BACKEND=file before starting any Dynamo process.

Quick Start — Embedding Model

1
import numpy as np
2
import aitune.dynamo as dyn
3
import aitune.torch as ait  # for ait.config, ait.load, ait.save
4
from sentence_transformers import SentenceTransformer
5
6
model = SentenceTransformer("intfloat/e5-large-v2")
7
ait.config.device_after_tuning = "cpu"
8
ait.load(model, "checkpoints/e5large.ait")
9
10
11
def mapping(req) -> dict:
12
    sentences = req.input if isinstance(req.input, list) else [req.input]
13
    return {"sentences": sentences}
14
15
16
def embed(sentences: list[str]) -> np.ndarray:
17
    return model.encode(sentences, normalize_embeddings=True, device="cuda")
18
19
20
config = dyn.DynamoWorkerConfig(
21
    type="embedding",
22
    model_path="intfloat/e5-large-v2",
23
    mapping=mapping,
24
)
25
dyn.dynamo_worker(embed, config)  # blocks until shutdown

API Reference

Import from aitune.dynamo:

1
import aitune.dynamo as dyn

DynamoWorkerConfig

dynamo_worker(model_or_fn, config)

Functional API. Validates config, starts the Dynamo runtime, and blocks until shutdown.

• model_or_fn: any callable, or a torch.nn.Module (requires config.mapping)
• config: DynamoWorkerConfig

DynamoWorker (class-based API)

For more control, subclass DynamoWorker and override setup() and serve():

1
import aitune.dynamo as dyn
2
3
4
class MyEmbeddingWorker(dyn.DynamoWorker):
5
    def setup(self) -> None:
6
        # called once at startup — load or tune your model here
7
        self.model = load_my_model()
8
9
    async def serve(self, request):
10
        sentences = request.input if isinstance(request.input, list) else [request.input]
11
        embeddings = self.model.encode(sentences)
12
        yield embeddings
13
14
15
MyEmbeddingWorker().run()

Override on_ready(runtime, endpoint) for post-startup work such as custom register_model calls.

Modality Types

If your function returns a plain dict, it is forwarded to the runtime as-is (no auto-packing).

Serving with run_dynamo.sh

The recommended way to start all processes locally is a run_dynamo.sh script that:

1. Starts the Dynamo HTTP frontend in the background
2. Starts the backend worker in the background
3. Polls /health until the endpoint is registered
4. Runs a smoke-test client request

$
#!/bin/bash
$
export DYN_DISCOVERY_BACKEND=file
$
$
python -m dynamo.frontend --http-port 8000 &
$
FRONTEND_PID=$!
$
$
python -m myapp.dynamo.backend &
$
BACKEND_PID=$!
$
$
trap "kill -9 $FRONTEND_PID; kill -9 $BACKEND_PID" EXIT
$
$
for i in {1..10}; do
$
  curl -s http://localhost:8000/health | grep -q '"dyn://aitune.backend.generate"' && break
$
  echo "Waiting for endpoint... (attempt $i)"
$
  sleep 5
$
done
$
$
python -m myapp.dynamo.client

See the E5Large example for a complete working version.

Next Steps

• AOT Tuning Guide — tuning a model before saving
• Tune Strategies — selecting the right optimization strategy
• Backend Guides — backend-specific deployment notes
• E5Large example — end-to-end embedding worker
• FLUX example — end-to-end image generation worker