Why the restart needs its own thread, not asyncio.to_thread

Companion illustration for NVIDIA-NeMo/RL #3592. Permalinks pin the PR head 5b4a732; stdlib quotes are CPython 3.13.

Seven pages on this PR, one question each — this one: why the choice of thread matters.

Restarting an engine means blocking Ray calls that know nothing about asyncio, so the control loop cannot simply await them. Both ways of handling that put the work on a real thread, and both leave the event loop completely free — the rollout pump and the watchdog keep running either way. They differ only in what happens to that thread when the restart never finishes. Today the supervisor uses asyncio.to_thread; the trainer side of the same codebase already refuses to, and starts its own daemon thread instead.

the run ends asyncio.run waits 300s, warns, gives up _python_exit joins process exits to_thread pooled, not a daemon restart_shard still running 300s stall, then a RuntimeWarning joined with no timeout — never returns the process never exits, and the GPUs stay held dedicated thread daemon=True restart_shard still running — the wait already timed out and marked the shard failed exits not joined — daemon threads are skipped, so the process exits and the GPUs come back Both rows keep the event loop free the whole time. The only difference is the thread each one runs on, and whether the interpreter waits for it. A stuck restart never returns to the pool's queue, so the sentinel that asks it to stop is never seen. the restart, still running waiting on it stuck forever clean exit

The same function, both ways

# today — engine_supervisor.py async def _restart(self, shard_idx: int) -> None: try: # Borrows a worker from the event loop's DEFAULT ThreadPoolExecutor, # shared by every to_thread in the process and capped at min(32, cpus+4). # The loop is free while we wait. No timeout, and no way to add one # that helps -- see below. url = await asyncio.to_thread(self._generation.restart_shard, shard_idx) except Exception as e: self._monitor.mark_restart_failed(shard_idx) return self._monitor.mark_loaded(shard_idx, base_url=url)
# the pattern already used by _sync_weights_within in this repo async def _restart(self, shard_idx: int, budget_s: float) -> None: loop = asyncio.get_running_loop() settled: asyncio.Future = loop.create_future() def _settle(setter, value) -> None: # wait_for cancels `settled` on timeout, and setting a result on a # cancelled future raises InvalidStateError inside the loop callback. if not settled.done(): setter(value) def _run() -> None: # Runs on OUR thread. Results cross back via call_soon_threadsafe, # because a Future may only be touched from the loop's own thread. try: url = self._generation.restart_shard(shard_idx) except BaseException as exc: loop.call_soon_threadsafe(_settle, settled.set_exception, exc) else: loop.call_soon_threadsafe(_settle, settled.set_result, url) threading.Thread(target=_run, name=f"gen-restart-{shard_idx}", daemon=True).start() try: url = await asyncio.wait_for(settled, budget_s) # bounded except asyncio.TimeoutError: self._monitor.mark_restart_failed(shard_idx) # counts as a failure return self._monitor.mark_loaded(shard_idx, base_url=url)
to_threaddedicated daemon thread
Event loop blocked?NoNo — identical
Thread comes fromthe shared default pool, capped at min(32, cpus+4)created for this one call
Daemon?No — the process waits for it at exitYes — the process does not
Wait bounded?not as writtenwait_for(budget_s)
A hang costsa pool slot forever, and shutdownone thread, and a recorded failure

Why a timeout on to_thread is not enough

The obvious repair is await asyncio.wait_for(asyncio.to_thread(...), 900). It does bound the wait. It does not fix anything else, because a thread in Python cannot be cancelled. The timeout abandons the await; the thread keeps running restart_shard for as long as it wants. So one pool slot is gone for the rest of the run, and the thread is still not a daemon, so shutdown still waits for it.

That is the whole reason the fix is three things at once: dedicated (does not eat a shared slot), daemon (cannot hold up exit), and bounded (the shard is marked failed instead of vanishing).

What "daemon" actually means. Python does not exit when your code finishes. It first runs threading._shutdown, whose own comment reads "Wait for all non-daemon threads to exit." So the flag answers exactly one question: does the process wait for this thread before exiting? daemon=True means no — the interpreter exits and the thread dies wherever it happens to be.

And the join really is unbounded. concurrent.futures.thread registers an exit hook that puts a stop-sentinel on each worker's queue and then calls t.join() with no timeout. A worker only checks that queue between work items, so a thread stuck inside restart_shard never sees the sentinel and the join never returns.

Why this matters for a restart specifically. A model reload takes minutes, and the restart is aimed at a shard that may be on a node that is never coming back — so "this call might never return" is the expected case, not the unlucky one. That makes it the worst call in the feature to run on a shared, non-daemon, unbounded thread. The cost is a training job that has already given up on a shard but cannot put its GPUs down.