When a restart runs, what it waits for, and how a shard re-enters

Companion illustration for NVIDIA-NeMo/RL #3592. Permalinks pin the PR head 5b4a732.

Seven pages on this PR, one question each — this one: the order things happen in.

Short answer up front: the restart itself is fully asynchronous and is queued behind nothing. It starts on the next fleet-health tick after the shard is marked DEAD, and it runs on a thread, so training and generation carry on underneath it. What is queued is everything after the engine comes back: re-entering the collective waits for the next refit, and on the NeMo-Gym path taking traffic again waits for one more fleet-health tick after that. On the default path — native Ray calls, no router — traffic resumes the moment the refit promotes it.

  1. The four loops, and where a restart actually runs
  2. Two ways a shard is marked DEAD, and why they take different times
  3. Scenario A — it dies between refits
  4. Scenario B — it dies during a refit
  5. Scenario C — it finishes restarting during a refit
  6. What waits for what

Every timeline below uses the same run: 4 generation shards, shard 2 dies, with the shipped defaults probe_interval_s: 5, unhealthy_threshold: 3 and refit_timeout_s: 300. Time axes are not to scale — a model reload takes minutes while a probe takes milliseconds — so each tick is labelled.

The four loops, and where a restart actually runs

The SingleController runs four asyncio tasks on one event loop. One loop means they never truly run at the same instant — they take turns, handing control over at each await. Anything that would block for a long time is pushed onto a thread so the other three keep their turns. Note what that does and does not buy: moving work to a thread keeps the loop free, but the task that awaits that work is still stuck until it finishes.

one event loop, four tasks taking turns rollout pump keeps generating throughout train pump trains, and drives the refit watchdog pump stall detection, metrics fleet probe pump probe, then tick, then push to the router threads, off the loop the shard restart nobody awaits it — fire and forget the weight transfer the train pump DOES await it, so training waits — the other pumps do not The probe pump does three things in a fixed order every 5s, and the order matters: probe the shards → supervisor.tick() → push serving URLs to the router single_controller.py L2275 — a shard condemned by the probe starts restarting on this tick, not the next

"On a thread" and "asynchronous" are not the same thing, and the two boxes on the right differ exactly there. Both the restart and the weight transfer are pushed off the loop, so neither freezes the other three pumps. But the weight transfer is awaited by the train pump: the refit is step 6 of a training step, right after step 5 trains the policy, and the pump cannot start the next step until it returns. So the weight transfer is synchronous with respect to the train pump — it is moved off the loop to protect the watchdog and the probe, not to let training run ahead of it. The restart is the opposite: it is started and then nobody waits for it at all.

Two more things worth stating plainly. First, nothing takes a lock: there is no mutual exclusion between a restart and a refit, and no queue that one waits in for the other. Second, generation is not paused during a refit on this path: nothing waits for in-flight rollouts to finish before weights start being written, and the code marks adding that as a known TODO.

Two ways a shard is marked DEAD, and why they take different times

Before any restart can start, the shard has to reach DEAD. How long that takes depends on how it failed, and the difference is large.

How it failsWhat the probe seesTime to DEADWhy
Crash — the process is gone RayActorError one probe, so ≤ 5s Ray only reports this once the process is really gone, so it is conclusive. It goes straight to DEAD with no threshold.
Hang — alive but not answering timeout three probes, so ≈ 15s A timeout is ambiguous — it could be a blip. It counts one failure, so the shard sits in SUSPECT until unhealthy_threshold is reached.

A crash also does one extra thing, and it needs a sentence of background first. Every refit arms a timer inside each trainer worker. If the weight transfer has not finished within refit_timeout_s (300s by default), that timer breaks the NCCL communicator so the trainers stop waiting on a rank that is never going to answer. It is the safety net for a generation shard that goes silent mid-transfer.

On a confirmed death, the controller cancels that timer — the code calls this "standing the deadline down". Cancelling the safety net sounds backwards, so here is why it is right. A process that has genuinely died closes its sockets, and NCCL sees that and unblocks the survivors by itself, which is the clean way out. If the timer fired first it would break the communicator instead, and breaking it strands GPU work already queued on the trainers' own streams — leaving their devices in a state no rebuild can use. So on a confirmed death, doing nothing is strictly better than intervening.

A frozen-but-alive rank never produces a confirmed death, so its timer is never cancelled and still fires. That is the case the timer exists for.

A detail that surprises people: the probe only walks serving_shards(), which is HEALTHY plus SUSPECT. So a shard is not probed while it is DEAD, RESTARTING or STALE. That is deliberate for the first two, but it does mean a freshly restarted STALE shard is unwatched in the window between its engine coming up and the refit that promotes it.

Scenario A — it dies between refits

The common case, and the one the feature is built for. Shard 2 crashes while the run is generating. Nothing is mid-collective, so nothing is stuck.

crash next probe (≤5s) minutes later next refit next probe shard 2 HEALTHY — the ledger has not noticed yet RESTARTING STALE — in the group, but takes no traffic HEALTHY probe RayActorError → DEAD → RESTARTING all in one pass: the probe records the death, then tick() starts the restart. DEAD is momentary. pushes the new URL to the router gap: ≤ 5s Gym path only restart on a thread — reloads the model mark_loaded → STALE, with a NEW port refit rebuild group, send weights promoted to HEALTHY here the model was already loaded during the restart; this step rebuilds the NCCL group to include shard 2, then sends it current weights between the crash and the probe the process is gone but the ledger still lists it as serving — that window is up to one probe interval rollout generation never stops — it just runs on the 3 surviving shards in the middle HEALTHY: serving DEAD: process gone RESTARTING: engine coming up STALE: in the group, no traffic

Two things about the right-hand end. First, the refit step is not reloading anything — the model was loaded from disk minutes earlier, during the restart. What the refit does is rebuild the NCCL group so it includes shard 2 again, and then send it the current weights. That is why the shard has to wait for a refit at all: it has a model, but the wrong version of it.

Second, the shard becomes HEALTHY at the end of that refit, when _promote_refit_shards runs — not at the start. The little gap after it is a real wait, and whether it costs you anything depends on how rollouts are routed:

RoutingWhen the shard gets traffic againWhy
Native Ray calls
(generation_router.enabled: false, the default)		 immediately, on the next request _next_dp_shard_idx asks the selector per request, and it reads serving_shards() live. Nothing is cached.
Via the NeMo-Gym router at the next probe tick, so ≤ 5s later Gym holds one fixed URL and never re-resolves it, so the router has to be told. _push_router_membership sends it serving_base_urls() on the probe pump's clock.

So your guess about the router was right, but only for the Gym path — and that push matters there for a specific reason: the restarted engine binds a new port, so without it Gym keeps talking to a socket nobody is listening on.

Scenario B — it dies during a refit

Now the timing is different, because the trainers are inside a collective that needs every rank. This is where the refit deadline earns its keep.

refit starts shard 2 hangs deadline fires second attempt it completes trainers sending weights blocked in NCCL same weights, sent again watchdog a per-refit thread in each worker is armed with refit_timeout_s (300s); when it expires it aborts the communicator, which raises RefitAborted recovery 1 2 3 all three finish before the second attempt starts — 1 and 2 are near-instant, 3 is the slow one 1 probe every shard · 2 mark the survivors STALE · 3 rebuild the refit communicator without shard 2 shard 2 HEALTHY SUSP DEAD RESTARTING — begins ~15s in, long before the refit gives up the others HEALTHY — still serving rollouts STALE — no traffic, right through the second attempt promoted back to HEALTHY only when the transfer completes an abort leaves every engine that was receiving with a mix of old and new weights, so the survivors are pulled out of service too the recovery probes on the spot rather than waiting for the 5s tick: the abort is an event, the probe is a clock, and the clock used to lose that race — leaving an empty absent set and no rebuild HEALTHY SUSPECT: failing probes, still serving RESTARTING STALE: holds partial weights, no traffic

What gets retried is the weight transfer, once. Not the request, not the training step. The controller calls _sync_weights_within(..., "first"); when that raises it recovers and calls the same function again with "retry" — same weights, same training step, smaller group — then _promote_refit_shards puts the survivors back in service. A second failure ends the run: at that point it is a fault, not a membership problem.

The three numbered steps in the figure are _recover_from_failed_refit: probe every shard on the spot rather than on the 5s clock, mark the survivors STALE if the failure was an abort, then rebuild the communicator without the dead shard. If nothing is absent but exactly one shard is already SUSPECT, that one is condemned as the silent participant — with several suspects, or none, it gives up rather than condemn a healthy shard by guessing.

Why are the survivors STALE when they are perfectly healthy? Because STALE is a verdict on the weights, not the process. Every other state describes the process; STALE says "this engine's weights are not the current ones". An abort stops a transfer mid-flight, so every receiving engine holds a mix of old and new — generating from that is silently wrong. mark_weights_partial puts it in one line: "the process is alive and refits normally, it just must not serve until a refit completes."

It is the same state a restarted engine sits in, for the same reason and a different cause. Both stay in the group so the refit can reach them, and out of the serving set so nothing reads from them meanwhile.

Two things about the timing. The ledger marks shard 2 DEAD after three failed probes — about 15s — but the trainers stay blocked for the full 300s deadline, because nothing connects the ledger's verdict to an in-flight collective. Only a confirmed process death ends the wait early, and it does so by cancelling the deadline and letting NCCL unblock the survivors itself. So a hang costs the full 300s; a crash costs seconds.

And the restart begins at that 15s mark, running on a thread the whole time the trainers are parked — the clearest "queued behind nothing" on this page. Its first act is to kill the hung worker, which recreate_worker expects ("a worker can be unresponsive rather than gone"). The replacement then has to load onto the GPU that worker was holding, so if its EngineCore is orphaned rather than reaped the restart fails on memory — which is where RAY_process_group_cleanup_enabled matters most.

One case does not recover. On the nccl_reshard transport, a fault while the bulk transfer is in flight cannot be recovered in-process: aborting the communicator does not retire CUDA work already queued on the trainers' own streams, so no rebuild can bootstrap on those devices. The run detects this and ends in seconds with a named cause rather than wedging. On the collective transport the same fault recovers.

Scenario C — it finishes restarting during a refit

This one is a bug rather than a timing story. A page of its own covers it start to finish, including why the metric meant to catch it cannot.

Shard 2 did not fail during this refit. It died in an earlier step and has been restarting ever since, so when this refit settled its membership shard 2 was already RESTARTING — absent, and correctly left out. All of that works as intended. The only new thing is when the reload finishes: it takes minutes and the loop is free throughout, so it can land at any moment, including in the middle of someone else's refit.

membership settled transfer in flight restart finishes refit ends the refit sending weights to shards 0, 1 and 3 — shard 2 is not in the group shard 2 RESTARTING since an earlier step — absent, correctly left out STALE HEALTHY mark_loaded fires here promote every STALE shard what each shard actually got out of this refit shard 0 received current weights HEALTHY shard 1 received current weights HEALTHY shard 2 received nothing at all HEALTHY ← now serving rollouts from the checkpoint it loaded off disk shard 3 received current weights HEALTHY the promotion asks "is this shard STALE?", never "was it in the group I just sent weights to?" — and until this PR nothing could turn a shard STALE in the middle of a refit, so the two questions had the same answer

It should stay STALE — that is exactly the bug. Shard 2 took no part in this refit, so nothing gave it current weights, and the design is explicit that a completed refit is the only way back into the serving set. But _promote_refit_shards promotes every shard that is STALE when it runs, and report_refit takes the shard index on faith. Neither one checks participation, because until this PR nothing could turn a shard STALE in the middle of a refit — so "is STALE" and "was refitted" described the same set, and no check was needed.

This is a live bug, raised in review and not yet fixed. The shard serves one step of rollouts generated from the checkpoint it loaded off disk, mixed into the training batch with nothing marking it. It self-corrects at the next refit, so the damage is bounded to one step — but it is silent, and the metric meant to catch it is written by the same call, so it reports a healthy-looking weight version. The fix is to promote only the shards the refit actually wrote to.

What waits for what

StepRuns onWaits forHow long
Notice a crashprobe pumpthe next probe tick≤ 5s
Notice a hangprobe pumpthree consecutive failed probes≈ 15s
Start the restartprobe pumpnothing — same tick as the verdictimmediate
Reload the modela threadnothing — nobody awaits itminutes
Transfer the weightsa thread—, but the train pump waits on this oneseconds
Rejoin the collectivetrain pumpthe next refitone training step
Take traffic again — native callsrollout pumpnothing; the selector reads the ledger liveimmediate
Take traffic again — via the Gym routerprobe pumpthe next tick after it is HEALTHY≤ 5s

So: is the restart synchronous with anything? No. It is started from the fleet-health tick and immediately handed to a thread, and nothing waits on it — not the refit, not training, not generation. The supervisor's tick() returns as soon as it has started the work, and a shard already being restarted is skipped on later ticks so a slow reload cannot pile up.

The waiting is all on the way back in, and it is deliberate: a restarted engine holds whatever weights it loaded from disk, so it must sit out until a refit gives it current ones. That is what STALE means — in the group, so the refit reaches it, but not in the serving set.