Skip to content

Independent engine drain (experimental)

core.publication_mode = "engine_drain" is an experimental alternative to barrier and refresh publication. It selects rolling updates: each engine stops admitting new work, finishes the requests it already owns on their admitted weights, loads the new snapshot and reopens, while its peers keep serving. No maintained example uses it; the full-model starting profiles use refresh, and small synchronous checks use barrier.

Enabling it

Copy an example to runs/, set core.publication_mode = "engine_drain", and choose a supported MoE model and topology. Check the configuration reference, run plan and validate, and submit through the committed-image launcher.

The mode requires resident, disaggregated Core MoE with TP1 engines, one optimizer step per collection and blocking shared-engine evaluation. Automatic fault tolerance, external engines, custom generators/filters, serving TP/DP/EP/PP greater than one, and snapshot-fleet evaluation are rejected.

Field Meaning
core.engine_drain_timeout Bounds drain, request and admission waits.
core.engine_update_timeout Bounds update acknowledgement and transport setup.
core.snapshot_capacity Number of exported weight versions retained at once (default 2).

How it works

The managed async producer reserves every sibling of a prompt group on one engine and version before sending any HTTP request, including requests still waiting for the client semaphore. New groups bypass the load-balancing router and use their reserved engine's URL. Response metadata must acknowledge the reserved version. Decoding releases the engine reservation before reward computation; the existing data-source ledger retains the group until consumption or deliberate drop.

After each optimizer step, rank zero packs the exported BF16 weights into a GPU bucket, copies it through reusable pinned host staging and places one full snapshot per version in Ray's object store, shared by all receivers. The driver waits for snapshot capacity before another capture. Provision host/object-store space for snapshot_capacity × the full exported model bytes, plus staging and normal rollout data; for example, a 37 GB model needs about 74 GB at capacity two.

Each engine has a separate delivery actor on the source trainer node with its own two-rank NCCL communicator to one TP1 receiver; training communicators and live parameters are never used by background delivery. This keeps bucketed transfer but adds host copies and per-engine CUDA contexts compared with synchronous export. The initial snapshot is delivered and checked against the startup serving-weight audit before training admission opens.

Publication on each engine closes admission, waits for owned requests, loads the snapshot, flushes caches and checks the engine's version before reopening. There is no pause, retract or abort on the normal path. Admission reserves one optimizer step of lag headroom, and completed groups still pass the usual consumption-time lag check. Completed groups share the single FIFO buffer; there is no per-version queue or newest-policy priority. Groups are internally homogeneous, but an optimizer batch may contain several eligible policy versions.

Failures, checkpoints and lifecycle

Snapshot versions are delivered in order per engine. Partial updates and wrong versions make the engine unavailable and propagate a terminal error. There is no automatic per-engine retry or recovery: the run fails and a fresh process must restore a committed checkpoint.

Native checkpoints do not quiesce inference or join rolling publication; see checkpoint resume semantics. Evaluation stops producer submission, drains already-owned work and joins publication; the completion queue temporarily expands by the number of owned groups so it cannot block that join. Resume regenerates outstanding prompts at restored weights; it does not restore partial decodes, KV/KDA state, publishers or in-memory queues. Shutdown joins publishers and closes their communicators before engine disposal.

While deliveries overlap, the integration holds the inference controller's update lock and pauses its health checks. Any evaluation, offload or diagnostic operation that needs the controller lock must first quiesce the publisher. After a delivery failure, health checks stay paused so incomplete weights are never treated as a recovered engine.

Monitoring

engine_drain.jsonl records group/request ownership, actual versions, decode tokens, grade completion, per-engine drain/update/reopen transitions, snapshot retention and delivery time/bytes. Overlapping stages must not be summed as elapsed wall time. Analyze a completed run with:

python -m scripts.miles.analyze_engine_drain /path/to/run --output /tmp/drain-audit.json

The audit checks request-attempt uniqueness and actual response/group versions, and reports independent engine progress. It does not certify numerical equality, consumed lag, learning or recovery.

Freshness

Independent drain removes the fleet barrier but does not minimize policy age. It drains every engine for each publication, delivers every captured version in order and consumes completed groups FIFO, which can delay current-policy batches. Possible improvements, not implemented here, include closing admission on a subset of engines during the next optimizer step, preferring newest-policy groups in both dequeue and backpressure, and letting a slow engine skip intermediate snapshots. When comparing schedules, look at time to the first useful current-version batch, the consumed-version distribution, useful trained tokens and rejected completed work, rather than generation occupancy alone.

Relation to mixed-policy refresh

core.publication_mode = "refresh" preserves sampled tokens and behavior probabilities, rebuilds serving state under new weights and records policy-span metadata, so a single response can span several policy versions. engine_drain instead finishes each request on its admitted weights and never mixes policies inside a response. See the publication-mode guide and mixed-policy refresh.