Which process owns each piece of the refit recovery

PR #3591 · all links pinned at ab0ffe2f

Six pages, one question each — this one: which process owns what.

The recovery machinery is spread across three kinds of process, and the piece that does the actual rescuing lives in the least obvious one. The ledger of shard health and the decision to rebuild are both in the driver, with no GPU. The communicators are in the GPU workers. The watchdog that breaks the hang is a thread inside each GPU worker process, created fresh for every refit and owned by nothing — which is why you cannot find it by looking for a long-lived object. The router, despite being the thing that routes traffic away from a dead shard, takes no part in the refit at all.

Where each object lives · 2 trainer ranks, 3 generation shards SingleControllerActor — one process, no GPU Everything here runs in one Python process. Nothing below writes to it. GenerationFleetHealth the ledger: state per shard + absent_shards() new here rest came with #3590 WeightSynchronizer issues the refit + reconcile_communicator() new here — decides to rebuild VllmGeneration handle to the gen actors + _refit_membership new — who to dispatch to plan_refit_membership() pure function — no Ray, no torch membership.py:55 new file — the rank layout Trainer — Ray actor, GPU (x2) self.model_update_group StatelessProcessGroup base_policy_worker.py:60 self.pp_comm_group second group, reshard only + abort() — new here RefitAbortWatchdog — a thread, per refit created by a `with` block, joined on exit. No attribute holds it. VllmGenerationWorker — Ray actor, GPU (x3) vLLM TP worker process — reached by collective_rpc VllmInternalWorkerExtension self.model_update_group self.pp_comm_groups {stage: g} RefitAbortWatchdog — a thread, per refit two sites, inside this innermost process .remote() .remote() GenerationRouter — its own Ray actor, no GPU Routes traffic away from a dead shard; reports failures to the ledger. Nothing in the refit path calls it. The watchdog exists in six places and is owned by none of them It is never stored on an object, so there is no "who owns it". Each refit builds one, arms it, and joins it. Four sites on the training side, two on the generation side: megatron_policy_worker.py:2501 megatron_policy_worker.py:2879 dtensor_policy_worker.py:1922 dtensor_policy_worker_v2.py:1203 vllm_backend.py:785 vllm_backend.py:1167 Reading the box colors Solid outline = a real OS process. Dashed = a process inside a process, or a component the refit never touches. Blue fill = added by this PR. Green text = a method or field this PR adds to something that already existed. Green arrow = a Ray call from the driver into a worker. The ledger is written only in the driver; no GPU worker writes to it.
Three process kinds, and a fourth thing that is not a process at all — the watchdog thread.

Why the watchdog cannot live in the driver

This is the design choice that makes the rest work, and it is worth being explicit about because the driver is where you would naturally put a timeout.

The driver cannot tell a slow refit from a wedged one. From outside, both look like ray.get not returning yet, and a large model's refit legitimately takes seconds. The driver also cannot break the hang even if it guessed right: freeing a rank stuck in NCCL means calling abort() on that rank's own communicator object, which only exists inside that worker's process.

And the stuck thread cannot rescue itself — it is blocked inside a NCCL call and will not run another line of Python until that call returns. So the only place a clock can usefully sit is a second thread in the same process: close enough to touch the communicator, separate enough to still be running. That is what RefitAbortWatchdog is, and why it is created at the call site rather than owned by anything.

QuestionAnswer
Does it live in the driver?No. Six worker-side sites, four training and two generation.
Does VllmGeneration own it?No. That object is driver-side and holds the membership, not the watchdog.
Does the generation Ray actor own it?Not quite. It is one level deeper — in the vLLM TP worker process, inside VllmInternalWorkerExtension.
Does the router own it?No. The router is not in the refit path at all.
What holds a reference?Nothing. It is a with-block local, joined when the block exits.

Does it work?

The one thing this arrangement depends on is that a second thread calling ncclCommAbort can actually free a thread already blocked inside a collective. That was measured: on 2×A6000, with a peer SIGKILLed mid-broadcast, the blocked survivor was released 0.15 s after another thread called abort(). If that had not worked, the in-process thread design would have bought nothing.

End to end, no CI job has watched this recover. The recovery harness needs 3 GPUs and both lanes have 2, so on the PR's own green run every registered variant printed [recovery] SKIP: … found 2. and exited. The feature was exercised by hand on larger hardware, so the evidence exists — it is just not in CI, and a green lane says nothing about it.

One consequence of the watchdog being per-process. It can only abort communicators it was handed. On the reshard path the transfer also creates split child communicators, and those are never passed to it — so a rank blocked on a child is not freed, and the run wedges even though the watchdog fired. That is raised on the PR; the fix is to release the children in abort() before the parent goes.