nemo_rl.models.generation.engine_supervisor

Brings dead generation shards back.

Without this the fleet only shrinks: a shard lost at hour one is gone for the rest of the run, and a job that sheds a few transient failures ends up permanently smaller. The recovery path itself already handles both directions – a rebuild can name more shards than the last one – so all that is missing is something to restart the engine and say when it is ready.

The restart is deliberately not awaited by the caller. Reloading a model takes minutes, and the control loop this runs from also drives the rollout pump, the watchdog and the refit; blocking it for a restart would stall the training that the surviving shards are still perfectly able to do.

Handover to the rest of the system is through fleet-health states, not through this class:

DEAD --(restart starts)--> RESTARTING --(engine up)--> STALE --(next refit)--> HEALTHY

RESTARTING is absent from collectives, so a rebuild that happens mid-restart correctly leaves the shard out. STALE is present but not serving, which is what lets the next refit write current weights into it before it takes traffic again.

Module Contents

Classes

EngineSupervisor

Drives restarts of dead generation shards, one background task per shard.

Data

DEFAULT_RESTART_TIMEOUT_S
DEFAULT_RESTART_BACKOFF_S

API

nemo_rl.models.generation.engine_supervisor.DEFAULT_RESTART_TIMEOUT_S

1800.0

nemo_rl.models.generation.engine_supervisor.DEFAULT_RESTART_BACKOFF_S

60.0

class nemo_rl.models.generation.engine_supervisor.EngineSupervisor(

generation: nemo_rl.models.generation.interfaces.GenerationInterface,

monitor: nemo_rl.models.generation.fleet_health.GenerationFleetHealth,

restart_timeout_s: float = DEFAULT_RESTART_TIMEOUT_S,

restart_backoff_s: float = DEFAULT_RESTART_BACKOFF_S,

clock: Any = time.monotonic,

)

Drives restarts of dead generation shards, one background task per shard.

Initialization

as_metrics() → dict[str, float]

Restart counters, merged into the per-step metrics dict.

gen_fleet/, matching GenerationFleetHealth.as_metrics – the two are merged two lines apart and gen_fleet/restart_attempts already carries restart state, so a second prefix would split one story across two namespaces.

tick() → None

Start a restart for any shard that needs one. Returns immediately.

Safe to call on every watchdog tick: a shard already being restarted is skipped, and a shard whose attempts are exhausted has been RETIRED by the monitor and is no longer DEAD, so it is never picked up again.

_restartable_shards() → list[int]

DEAD, not already being restarted, and past its cooldown.

The cooldown is the whole reason this is not just a state test. A failed restart returns the shard to DEAD and the next probe tick picks it straight back up, so with a 5s probe interval the entire five-attempt budget can be spent inside 25s – on one cause, none of the attempts having waited for it to clear. The cause that motivated the feature is exactly this shape: an orphaned EngineCore held its GPU for 370s, so every attempt inside that window was guaranteed to fail on placement.

_begin_restart(shard_idx: int) → None

_forget(shard_idx: int, task: asyncio.Task) → None

Drop the finished task so a later tick can retry this shard.

async _restart(shard_idx: int, attempt: int) → None

async _restart_off_loop(

shard_idx: int,

attempt: int,

) → Optional[str]

Run one restart on a dedicated daemon thread, and stop waiting after the budget.

A DEDICATED DAEMON THREAD, not asyncio.to_thread, for the same reason _sync_weights_within gives: to_thread runs on the default ThreadPoolExecutor, whose workers are not daemons and are joined at interpreter exit with no timeout, and a worker only checks for the stop sentinel between items. A thread parked inside a reload never sees it, so the join never returns and the process cannot exit. A model reload takes minutes, which makes this the call in the feature most likely to be that thread. drain cannot help: the join happens whether or not anyone calls it.

BOUNDED, because nothing else in the chain is. restart_shard ends in ray.get calls with no timeout, and create_worker returns before the actor is scheduled – so a placement-group bundle that can never be filled, which is what a lost node looks like from here, does not raise: it blocks in post_init. Then the except above never fires and the shard sits in RESTARTING for the rest of the run, never retried because it is no longer DEAD and never retired because retirement is driven by restart attempts. Timing out converts that silent park into a failed restart that spends an attempt and eventually retires the shard.

The orphaned thread is a real cost, not a free win: the reload may still be running, and if it later succeeds it does so into a fleet that has written the shard off. It is bounded by the attempt budget, and the alternative is a shard that disappears.

async drain(timeout_s: Optional[float] = None) → None

Wait briefly for in-flight restarts so shutdown does not abandon one.

Bounded on purpose: restart_shard can block on a placement-group bundle that will never be filled again, so giving up here is what lets the process exit. The thread it gave up on is a daemon (see _restart_off_loop), which is what makes giving up harmless rather than a leak.