`nemo_rl.distributed.refit_watchdog`#

Break a refit collective that a dead peer has left hanging.

WHY THIS HAS TO RUN INSIDE THE WORKER. A generation rank that dies mid-refit leaves the surviving ranks blocked in NCCL with no timeout and no error – observed directly, both policy workers stuck in packed_broadcast_producer -> cuda stream synchronize while the run sat wedged for 1801s. The controller cannot rescue them by RPC: the collective blocks the worker actor’s event loop, and the worker actors carry no max_concurrency, so an incoming abort call would queue behind the very operation it is meant to interrupt. The abort must therefore come from a thread already inside the process, which is exactly the arrangement the design’s NCCL spike validated (a survivor released 0.15s after another thread called abort()).

TWO SEMANTICS THAT SHAPE THE API, both established by that spike:

An aborted collective returns without raising. So the caller cannot detect this with try/except; it has to ask whether the abort fired. Hence :attr:fired.
The destination buffers hold partial data afterwards. A generation shard caught mid-refit holds a mix of old and new weights and must not serve until a later refit completes. Callers are responsible for propagating that – see RefitAborted.

Inert unless armed with a positive timeout, so a run that does not configure one behaves exactly as before, down to not starting a thread.

Module Contents#

Classes#

`_Abortable`
`RefitAbortWatchdog`	Abort the given group(s) if the guarded block outlives `timeout_s`.

Functions#

`is_refit_context_lost`	True when the abort orphaned GPU work and no rebuild on that device can succeed.
`is_refit_abort`	True for a RefitAborted, or for one flattened to a plain Exception in transit.
`await_off_loop`	Run a blocking call on a daemon thread so this actor’s event loop stays free.
`sync_stream_within`	Wait for `stream`’s enqueued work, giving up after `budget_s`.
`release_within`	Run a teardown that may never return, without letting it wedge the caller.
`stand_down_armed_watchdogs`	Disarm every refit deadline armed in this process, and say how many.
`hold_refit_for_fault_injection`	Block a refit receive while a test holds it open. Inert unless asked.

Data#

`REFIT_ABORTED_TOKEN`
`REFIT_CONTEXT_LOST_TOKEN`
`RELEASE_GRACE_S`
`_ARMED_LOCK`
`_ARMED`

API#

class nemo_rl.distributed.refit_watchdog._Abortable[source]#

Bases: typing.Protocol

abort() → None[source]#

nemo_rl.distributed.refit_watchdog.REFIT_ABORTED_TOKEN#: ‘[refit-aborted]’

nemo_rl.distributed.refit_watchdog.REFIT_CONTEXT_LOST_TOKEN#: ‘[refit-context-lost]’

exception nemo_rl.distributed.refit_watchdog.RefitAborted(*args: object)[source]#

Bases: RuntimeError

A refit was cut short because a peer stopped participating.

Raised by the worker that armed the watchdog, not by NCCL – the aborted call itself returns cleanly, so this is the only signal the caller gets.

The constructor prefixes :data:REFIT_ABORTED_TOKEN to the message. Idempotent, so a re-raise or an unpickle (RuntimeError.__reduce__ replays args) does not stack prefixes.

Initialization

Initialize self. See help(type(self)) for accurate signature.

nemo_rl.distributed.refit_watchdog.is_refit_context_lost(error: BaseException) → bool[source]#

True when the abort orphaned GPU work and no rebuild on that device can succeed.

Recovery after this is not merely unlikely, it is impossible in-process, and four runs established that by elimination rather than argument. Job 6521181’s py-spy dump caught both trainers in init_nccl_communicator with frame-less ncclCommAbort threads 25 minutes on. Job 6523731 killed the frozen victim before the rebuild, closing its sockets, and they wedged identically. Job 6582457 pinned the release to the caller’s CUDA device, with a SIGKILLed peer, and it still did not return. Job 6584636 stood the deadline down on confirmed death and lost the race to Ray’s own detection.

So the controller stops trying. Detecting this and ending the run in seconds is the supported behaviour; recovering from it is left to a future change – see design_vllm_fault_tolerance.md section 8.5.7.

nemo_rl.distributed.refit_watchdog.is_refit_abort(error: BaseException) → bool[source]#

True for a RefitAborted, or for one flattened to a plain Exception in transit.

The type check alone is not enough anywhere downstream of a vLLM collective_rpc: the abort is raised inside the engine core process and arrives at the Ray actor as Exception("Call to <method> method failed: ..."). A bare except RefitAborted there is dead code, which is what job 6484412 demonstrated – the deadline fired, the abort was named, and the run still wedged because the handler never matched.

async nemo_rl.distributed.refit_watchdog.await_off_loop(fn)[source]#

Run a blocking call on a daemon thread so this actor’s event loop stays free.

Ray runs a SYNC actor method directly in the event loop – sync_to_async wraps it as async def wrapper: return func(...), with no executor – so a refit that blocks in NCCL starves every other call to the same actor. max_concurrency cannot help; it interleaves coroutines, and a coroutine blocked in C never yields.

That is why the recovery could not run in job 6509685. The controller gave up on the stuck refit and called init_collective to rebuild, but that call queued behind the refit still occupying this loop. Rank 0 is the rendezvous master, so the store was never created, and the surviving generation worker timed out dialling it for 300s, twice, before the run ended.

Daemon, because asyncio.to_thread’s default executor is non-daemon and joined at interpreter exit: a thread still parked in NCCL would hang shutdown, trading a wedge in the refit for a wedge on the way out.

No timeout here on purpose. Bounding the wait is the controller’s job (_sync_weights_within); this only decides which thread blocks.

nemo_rl.distributed.refit_watchdog.sync_stream_within( stream, budget_s: Optional[float], what: str, ) → None[source]#

Wait for stream’s enqueued work, giving up after budget_s.

WHY THIS EXISTS, and why the watchdog above is not enough. Aborting a communicator does not retire work already enqueued on a CUDA stream. When a generation rank stops receiving mid-refit the sends sit on the stream forever, and torch.cuda.synchronize – which waits on the whole device – never returns. The watchdog cannot help: the abort fires, the kernels do not retire, the guarded block never exits, and

Attr:: RefitAbortWatchdog.fired is never read. No exception-translation reaches a hang.

Job 6485245 measured exactly that on 4xGB200: both policy workers parked in synchronize 1801s after their own abort had logged, while the generation workers had already unwound and gone idle.

So the wait is bounded here rather than trusted to end. The event is POLLED, not waited on, so nothing can be left holding the GIL, and the happy path still finishes with the same device-wide synchronize() – behaviour is identical when nothing is wrong.

This does NOT recover the fleet. In-flight kernels are orphaned and the caller’s CUDA context should not be trusted afterwards, so the RefitAborted raised here is expected to end the run – attributably, in seconds, rather than after a 30-minute stall. Recovering a frozen-but-alive rank on this transport stays out of scope.

budget_s of None or <= 0 keeps the original unbounded synchronize, so a run with no refit deadline configured behaves exactly as before.

nemo_rl.distributed.refit_watchdog.RELEASE_GRACE_S#: 30.0

nemo_rl.distributed.refit_watchdog.release_within(release, budget_s: float, what: str) → None[source]#

Run a teardown that may never return, without letting it wedge the caller.

THE SIXTH UNBOUNDED WAIT, and the one that only the reshard transport reaches. StatelessProcessGroup.abort() calls abort_xferdtensor_python_subcommunicators before the parent, and those split children exist ONLY on the Python reshard path – on the packed-broadcast path that call finds no cache entry and returns immediately. ncclCommAbort joins the communicator’s proxy thread, and a proxy thread blocked reading from a SIGSTOPped peer never returns: the socket is open and idle, so nothing errors and nothing times out. SIGKILL closes it and the proxy errors out at once, which is why the killed variants never see this.

Job 6518381 measured the consequence. On the rebuild, train rank 0 – the rendezvous store’s master – entered init_collective, printed its line, and never bound the port, because the only statement between the two is the old group’s release. The surviving generation rank then spent 600s (300s, an 89s backoff, 300s again) failing to connect to a store that was never created, and the run died at 700s having condemned the right shard and planned the right membership.

Bounded, on a DAEMON thread, and deliberately not joined. asyncio.to_thread and a bare ThreadPoolExecutor both use non-daemon threads that the interpreter joins at exit, which would move the wedge to shutdown rather than remove it. Nothing downstream reads a result: the release exists to stop resources accumulating across rebuilds, so a release that never finishes costs one stuck thread, while waiting on it costs the run.

nemo_rl.distributed.refit_watchdog._ARMED_LOCK#: ‘Lock(…)’

nemo_rl.distributed.refit_watchdog._ARMED: set[RefitAbortWatchdog]#: ‘set(…)’

nemo_rl.distributed.refit_watchdog.stand_down_armed_watchdogs() → int[source]#

Disarm every refit deadline armed in this process, and say how many.

THE DEADLINE IS FOR A SILENT PEER, NOT A DEAD ONE, and firing it on a dead one is strictly harmful. When a generation rank’s process is gone its sockets close, NCCL’s own error path unblocks the survivors, and the run recovers off the pre-existing actor-death route – job 6405953 passed the reshard kill variant that way with RefitAborted appearing zero times, before any deadline existed.

Once the deadline was added it started winning that race. It aborts at its timeout, sync_stream_within gives up on kernels already enqueued on the trainers’ streams, and the CUDA context cannot be trusted afterwards – so the rebuild that would have succeeded now cannot. recovery-reshard-refit has failed continuously since job 6512153, which is the run where the deadline first began firing on that path, and jobs 6521181/6523731/6582457 each confirmed the abort never retires: not with the peer SIGKILLed, not with the release pinned to the caller’s device.

So the controller stands the deadline down the moment a probe reports an actor DEATH, which is conclusive in a way a timeout is not. The frozen case is untouched: a frozen rank is alive, no death is ever recorded, and the deadline still fires and still ends the run attributably.

The controller can reach this at all only because the refit runs off the actor’s event loop (see await_off_loop); a worker blocked in the loop could not service the call.

Idempotent, and safe to call when nothing is armed.

class nemo_rl.distributed.refit_watchdog.RefitAbortWatchdog( group: Optional[Union[nemo_rl.distributed.refit_watchdog._Abortable, collections.abc.Sequence[Optional[nemo_rl.distributed.refit_watchdog._Abortable]]]], timeout_s: Optional[float], )[source]#

Abort the given group(s) if the guarded block outlives timeout_s.

Use as a context manager around the collective::

with RefitAbortWatchdog(self.model_update_group, timeout_s) as guard:
    ...collective...
if guard.fired:
    raise RefitAborted(...)

A sequence may be passed instead of one group, and the nccl_reshard transport needs that: it moves weights over per-PP-stage bulk groups and then broadcasts the remainder over the shared model_update_group, so a hang can be in either family and nothing at this level can tell which. Aborting all of them costs nothing – abort() is idempotent and safe on a group that never built a communicator – and the recovery rebuilds every family regardless.

timeout_s of None or <= 0 disarms it entirely: no thread is started and fired stays False, so the default configuration is bit-for-bit the old behaviour.

Initialization

property armed: bool#

stand_down() → None[source]#

Cancel this deadline without firing it; see stand_down_armed_watchdogs.

Sets the same event the guarded block sets on a clean exit, so the watch thread returns without aborting anything and fired stays False. Racing a watch thread that is already past its wait is harmless: _done is only ever set, never cleared, and the abort it performs is idempotent.

property fired: bool#: True if the deadline passed and abort() was called.

_watch() → None[source]#

__enter__() → nemo_rl.distributed.refit_watchdog.RefitAbortWatchdog[source]#

__exit__( exc_type: Optional[type[BaseException]], exc: Optional[BaseException], tb: Optional[types.TracebackType], ) → None[source]#

nemo_rl.distributed.refit_watchdog.hold_refit_for_fault_injection() → None[source]#

Block a refit receive while a test holds it open. Inert unless asked.

Does nothing unless NRL_REFIT_HOLD_FILE names a path that exists, so a real run pays one os.path.exists per refit and behaves no differently.

It exists because “kill a shard during the refit” is otherwise untestable. A refit on the functional test’s model takes ~0.10s, and the harness has to notice one started and then find and kill a process: job 5925668 aimed at the collective and landed in the RPC epilogue instead. That is a real failure mode and worth handling, but it is not the one the test claimed to cover, so the abort-and-rebuild path went unexercised while the run still reported a result.

A file rather than a fixed delay because the harness has to hold one specific refit – the one after the step it kills at. A delay on every refit would slow the whole run for the sake of one moment and still not be aimed at it.

Bounded by NRL_REFIT_HOLD_MAX_S so a harness that dies mid-test cannot wedge the worker it was holding.

nemo_rl.distributed.refit_watchdog#