How a dead generation shard gets back into the fleet

Companion illustration for NVIDIA-NeMo/RL #3592 (part 4/4 of the SingleController elastic-recovery stack). Permalinks pin the PR head 5b4a732; "before" code is pinned at the base commit d0f8c07, where it still exists. Earlier in this stack: which failures survive.

Seven pages on this PR, one question each — this one: why a returning shard was shut out.

A run keeps several generation shards. The earlier PRs in this stack made a dead one survivable: it is dropped, the refit communicator is rebuilt over the survivors, and training continues. But the fleet could only ever get smaller.

Two separate things were missing, and this PR adds both. Nothing restarted a dead engine — there was no supervisor and no restart_shard method to call. And even if you had restarted one by hand, it could not have rejoined, because the code that decides whether to rebuild the communicator asked only "is anything absent right now?" A returning shard makes the absent set empty, which looks exactly like a healthy fleet with nothing to do.

Three stories, one dead shard

Three shards, numbered 0, 1 and 2. Shard 1 dies. The top row is what actually happened before this PR. The middle row is the thought experiment that shows why restarting alone would not have been enough. The bottom row is this PR.

steady state shard 1 dies minutes later next refit before what happened 0 1 2 0 1 2 0 1 2 0 1 2 nothing restarts it smaller fleet, for good restart only a thought experiment 0 1 2 0 1 2 0 1 2 0 1 2 engine up, so nothing is absent the check returns early: no rebuild, still out with this PR 0 1 2 0 1 2 0 1 2 0 1 2 the supervisor restarts it wanted and built differ, rebuilt — it serves again serving dead engine up, old weights, no traffic alive but shut out of the refit the check that decided the middle row, before: if not absent_shards: return False

How we know the top row is what really happened

At the base commit, three things are true and each is easy to check. engine_supervisor.py is added by this PR, so no file existed to drive a restart. restart_shard is also added by this PR — it had been written earlier and then deleted by a merge, and the comment above it records that. And the ledger's restart transitions had no caller at all — the old code says so itself, in the comment this PR deletes:

because absent_shards() never empties again -- nothing in production calls mark_restarting or mark_loaded

collective_weight_synchronizer.py at the base commit. mark_restarting is what moves a shard out of DEAD, and mark_loaded is what says its engine is back. With neither ever called, a shard that died stayed dead, and the absent set it was in never emptied.

What the fix compares

The new code stops asking "is anything absent" and compares two sets of shard numbers: the set the refit wants in the communicator now, and the set the communicator was actually last built over.

The wanted set is not a wish list. It is every shard minus the ones currently marked absent, and absent means one of three states: DEAD, RESTARTING or RETIRED. So it is smaller than the full fleet exactly when one of those is true, and equal to the full fleet the rest of the time. Same three shards, same moments:

moment wants in the group last built over same? decision --------------------------------------------------------------------------- steady state {0,1,2} {0,1,2} yes skip shard 1 dies {0,2} {0,1,2} no rebuild over {0,2} engine back up {0,1,2} {0,2} no rebuild over {0,1,2} <- the new one next refit {0,1,2} {0,1,2} yes skip

That set is computed by desired_membership(), which is phrased as "what the membership should be" rather than "what changed" — that is the whole trick, because it reads the same in both directions. The built set is remembered in _built_membership, and the comparison is one line. Notice the third row: that is the case the old check could never see, because its absent set is empty and so was a healthy fleet's.

The other half is the thing that makes the engine come back at all: EngineSupervisor.tick() picks up dead shards and calls restart_shard in the background, handing the shard along four states:

DEAD --(restart starts)--> RESTARTING --(engine up)--> STALE --(next refit)--> HEALTHY ^^^^^^^^^^ ^^^^^ not in the collective in the collective, at all but takes no traffic

RESTARTING counts as absent, so a rebuild that happens while the engine is still loading correctly leaves it out. STALE is the opposite: the shard is in the group, so the refit can write weights into it, but it takes no traffic until that refit lands. A restarted engine only has whatever it loaded from disk, and that ordering is what keeps those old weights out of a rollout.

What this buys you. A long run that sheds a few transient failures no longer ends permanently smaller — a shard lost at hour one can be serving again a few minutes later. It is off by default: set async_rl.generation_fleet_health.restart_dead_shards: true on top of generation_fleet_health.enabled.

Two things to know before you turn it on. Restarting does not raise the min_healthy_shards floor — a restarting shard is not serving, so a run can still stop while a restart is in flight. And a killed vLLM worker leaves its EngineCore holding the GPU unless the raylet is told to reap it, which needs RAY_process_group_cleanup_enabled=1 set before ray start, not from the driver. Without that the replacement has nowhere to load.