Part III - Cluster and platform layer

The ML Platform Bridge

Lesson 12 made the whole thing run on a real cluster: a head node, autoscaling worker groups, KubeRay turning a YAML file into pods. But we still have five libraries that we have only ever met one at a time — Data (09), Train (10), Tune (08), Serve (11). A production ML system is never one of them; it is all of them in a line, and the line is where systems actually break. This lesson reads the book's "Ray AI Runtime" chapter not as a product to learn but as the answer to one question: how do these libraries compose into a platform, and what is the object that crosses every boundary between them? The answer is the checkpoint — and the sharpest hazard in the whole pipeline lives in the handoff it travels through.

1 · A platform is four stages around one runtime

Each library we have met answers a different question about the same model. Put them in the order a model actually moves through and the platform appears:

The book calls the umbrella over these "AIR." The durable idea is not the brand — it is that all four run on one Ray runtime, sharing one object store (lesson 04), so a value produced by one stage can be handed to the next without a glue service, a separate scheduler, or a round-trip through external storage you have to wire up by hand. That shared substrate is the platform. And the value that crosses every boundary is a single kind of object.

2 · The checkpoint is the universal handoff

Watch one checkpoint cross every boundary. The same object that Train writes is the object Tune resumes from, the object batch prediction loads, and the artifact a Serve deployment opens. Nothing else travels the full length of the pipeline.

This is why the platform is not five systems bolted together: a single object type, written once, flows the whole length. The handoff contract is "produce a checkpoint; consume a checkpoint." Everything else — schedulers, replicas, blocks — is local to a stage. But notice the quiet assumption hiding in that contract: the checkpoint must contain everything the next stage needs. If it carries the weights but not the feature transform, the boundary leaks. That leak has a name.

3 · Train/serve skew — the hazard at the feature boundary

The model never sees raw input. Between raw data and the model sits a preprocessor: a learned feature transform — a normalization (subtract mean, divide by standard deviation), a tokenizer, a categorical encoder, a learned vocabulary. It is fitted on data: the normalizer learns the training mean and standard deviation; the encoder learns which categories exist. The model's weights were tuned assuming inputs arrive already passed through that specific fitted transform.

The bug, worked in numbers

Suppose one feature is a user's account age in days. On the training set it has mean = 50, standard deviation = 30. The normalizer fitted at training time computes z = (x − 50) / 30. The model learned its weights against z-scores in roughly the [−2, +3] range — a typical account is age 50, which maps to z = 0; a 6-month account, age 180, maps to z = (180 − 50)/30 ≈ 4.3, a meaningful "this is an unusually old account" signal.

Now the serving code makes the classic mistake: instead of loading the saved training mean and standard deviation, it normalizes each incoming request batch by that batch's own statistics. A burst of long-tenured "power users" arrives — a request batch with mean = 200, standard deviation = 40. Trace the same age-180 user through the wrong transform:

The fix: fit on training data, save with the checkpoint, reuse byte-for-byte

The transform must be fitted once on the training data and then frozen and saved alongside the model in the checkpoint, so Serve loads the identical mean = 50 / std = 30 and applies it to every request, regardless of what the live batch looks like. The checkpoint's job is to be the single source of truth for both the weights and the transform that the weights expect. If both ride in the same checkpoint, the boundary cannot leak — there is only one transform in existence and every stage reads it.

import ray
from ray import serve
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig, Checkpoint
from ray.data.preprocessors import StandardScaler

# --- STAGE 1+2: Data fits the preprocessor on TRAINING data, Train consumes it ---
train_ds = ray.data.read_parquet("s3://data/train")

scaler = StandardScaler(columns=["account_age_days"])   # the feature transform
scaler.fit(train_ds)                  # learns mean=50, std=30 from TRAINING data ONLY
train_ds = scaler.transform(train_ds) # apply the FITTED transform to training data

def train_loop(config):
    shard = ray.train.get_dataset_shard("train")
    model = build_and_fit(shard)      # weights learn against the scaled features
    # CRITICAL: save the fitted scaler INSIDE the checkpoint, with the weights.
    ckpt = save_model_and_preprocessor(model, scaler)   # -> Checkpoint directory
    ray.train.report({"val_acc": evaluate(model)}, checkpoint=ckpt)

trainer = TorchTrainer(
    train_loop,
    scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
    datasets={"train": train_ds},
)
result = trainer.fit()
best_ckpt = result.checkpoint            # the universal handoff object

# --- STAGE 4: Serve loads the SAME checkpoint AND the SAME fitted preprocessor ---
@serve.deployment(num_replicas=2, ray_actor_options={"num_gpus": 1})
class OnlineModel:
    def __init__(self, ckpt: Checkpoint):
        self.model, self.scaler = load_model_and_preprocessor(ckpt)  # mean=50, std=30
    async def __call__(self, request):
        row = await request.json()
        row = self.scaler.transform_batch(row)   # SAME transform as training — no skew
        return {"score": self.model(row)}

serve.run(OnlineModel.bind(best_ckpt))

The whole defense is two lines of discipline: scaler.fit runs on training data only, and the fitted scaler is serialized into the checkpoint so the deployment's __init__ reconstructs the exact same transform. Never call fit on request data; never recompute statistics online. The preprocessor is part of the model, not part of the request handler.

4 · The platform trade-off: one runtime vs best-of-breed

Why build the whole pipeline on one Ray runtime at all? The alternative is a best-of-breed stack: a dedicated feature store, a separate training service, and a Kubernetes-native serving system (KServe, Triton) — each the best tool for its stage, wired together with glue. Both are real, defensible architectures. The choice turns on a single tension: a shared substrate buys you frictionless handoffs but concentrates your operational risk.

Neither answer is universally right. The honest reading of AIR is that it optimizes the handoff-and-skew axis hard, and asks you to accept a concentrated blast radius — which is why lesson 12's head-node SPOF and lesson 14's failure model are not side topics but the bill that comes due for unification.

5 · Memory, failure, and autoscaling of the multi-stage pipeline

A pipeline that chains four stages has behavior no single stage shows. At altitude — the deep failure model is lesson 14 — three properties matter:

6 · Failure modes & checklist

Checkpoint exercise

Where this points next

We have leaned on one phrase repeatedly — "a worker dies and resumes from its checkpoint" — as if that recovery were automatic and free. It is neither. How does Ray know an object was lost, and when can it rebuild it versus when is it gone for good? What does "at-least-once" execution mean for a stage with side effects? Lesson 14 opens the failure model the platform quietly relies on: lineage-based reconstruction of lost objects, where it works and where it can't (actor state, side effects, a dead owner's ray.put), the retry-vs-idempotency contract, and the observability — dashboard, timeline, ray memory — you need to see the object-store contention and stragglers this lesson only named. The platform composes cleanly; lesson 14 is what happens when a piece of it dies.

Interview prompts

• What single object joins Data, Train, Tune, and Serve into a platform, and how does it cross each boundary? (§1–2 — the checkpoint: Train writes it, Tune resumes from and selects among checkpoints, batch prediction loads it in map_batches, and a Serve replica loads it in __init__ — the same durable directory the whole way.)
• What is train/serve skew, why is it dangerous, and where does it live? (§3 — the serving-time feature transform differs from the training-time one, so the model sees a different input distribution online; dangerous because it is silent (200 OK, confident wrong answers, no exception); it lives at the Data→Serve preprocessor boundary.)
• Walk the numeric skew bug: training mean 50/std 30, a live batch mean 200/std 40, raw input 180. (§3 — correct: (180−50)/30 ≈ +4.3; live-fitted: (180−200)/40 = −0.5; a ~5σ swing on one feature into a region the weights never trained on, collapsing slice accuracy ~0.93 → ~0.50.)
• What is the fix for train/serve skew, and why does it belong in the checkpoint? (§3 — fit the preprocessor on training data only, freeze it, and serialize it inside the checkpoint so Serve loads the identical transform; if weights and transform ride together there is only one transform in existence, so the boundary can't leak.)
• Compare one unified runtime to a best-of-breed stack. (§4 — unified: one checkpoint handoff, in-cluster data sharing, skew preventable by construction, but concentrated blast radius and tighter coupling; best-of-breed: contained per-stage failure and independent scaling, but glue at every boundary and skew risk at each. Unify when small/Python-native; isolate the SLO-bound online tier first.)
• What is blast radius, and how does it argue for isolating an online serving tier? (§4 — the set of things that fail when one component fails; sharing the cluster means a head-node fault takes down live serving too (100% of stages), so an SLO-bound tier often gets its own cluster to shrink the radius — at the cost of a new boundary the preprocessor must cross.)
• Why is the cluster under an AIR-style pipeline heterogeneous, and how do you size it? (§5 — stages have opposite resource shapes: CPU-heavy preprocessing vs a small GPU training gang vs traffic-scaled GPU serving; size the CPU:GPU ratio so GPUs never starve, e.g. 8000 rows/s ÷ 500 rows/s per CPU worker = 16 CPU : 4 GPU, encoded as per-worker-group min/max.)