slurm: stop cancelled jobs from leaking whole GPUs by RyanPCo · Pull Request #682 · GaTech-RL2/EgoVerse

RyanPCo · 2026-09-20T16:51:17Z

Stops cancelled DDP jobs from leaking whole GPUs. Dataloader workers now die with their rank, ranks die with their launcher under Slurm, and each job first clears what an earlier job left on the node.

Behaviour changes

Nothing a run computes changes. Under Slurm, trainHydra.main now SIGKILLs your own processes on the node that are reparented to init and hold an nvidia fd, and ties each rank to its parent. EGOMIMIC_REAP_GPU_ORPHANS=0 turns the reaper off.

The problem

A cancelled DDP job left GPUs allocated, with nvidia-smi blaming PIDs absent from /proc. The driver is fine: the memory is held by live processes reparented to init. One node had 126 pt_data_worker survivors across three cancellations. Each held /dev/nvidia0-7, and together they pinned 7 of 8 GPUs at ~39 GiB each and ~1.7 TB of RAM.

Why it happens

Workers inherit the GPUs. Workers fork after CUDA init, so they inherit the rank's device fds. Ranks open every visible GPU, so one survivor pins the node.
Workers can't exit on their own. A worker inherits both ends of its result pipe, so a dead rank never sends EPIPE. The worker blocks in pipe_write and never reaches PyTorch's ppid check.
Whole ranks leak the same way. A rank wedged in torch.cuda.synchronize() inside inductor's cudagraph_trees ignores SIGTERM until Slurm gives up (UnkillableStepTimeout). Slurm then drops the job with its ranks still spinning.
Slurm doesn't sweep them. With ProctrackType=proctrack/linuxproc Slurm walks the PPID tree, so it never sweeps anything reparented to init. proctrack/cgroup would, but that is a cluster admin change.

The fix

egomimic/utils/gpu_orphans.py:

orphan_guarded / die_with_parent set PR_SET_PDEATHSIG on every worker, so the kernel kills it with its rank whatever it is blocked on. They are wired into the train and val loaders (pl_data_utils.py) and the norm-stats pass (zarr_dataset_multi.py).
arm_rank_pdeathsig does the same for the rank process, under Slurm only.
- A process whose parent is already gone when it arms exits 1. It first writes one line to stderr saying why, since os._exit skips logging's flush.
- SLURM_JOB_ID is also set inside a salloc shell, so a run started there under nohup (SIGHUP ignored) is left unarmed. Armed, it would die with the shell it was meant to outlive.
reap_orphans runs from trainHydra.main before CUDA init and clears what an earlier job left, so a leak cannot outlive the job that made it.
- It only signals processes that are yours, reparented to init and holding an nvidia fd.
- It is Slurm-gated and runs once per node (SLURM_LOCALID 0).
- EGOMIMIC_REAP_GPU_ORPHANS is parsed with str2bool: 0/false/no disable it, any truthy spelling keeps it on, and anything else raises.
- It reads PPid from /proc/<pid>/status, not /proc/<pid>/stat, where a comm with spaces shifts the fields.

AGENTS.md gains a section on this, including that any new DataLoader on a training path needs orphan_guarded or die_with_parent.

A manual reaper script for nodes you are not submitting to is kept out of this PR.

Also folded in: DDP devices in the trainer configs and gres in the submitit launchers are now per node (Lightning multiplies by num_nodes), and eval-mode val subsampling is single-rank.

Tests

tests/unit/test_worker_orphan_guard.py:

Two tests reproduce the leak end to end. A worker inherits both pipe ends and blocks writing, then its rank is SIGKILLed. The unguarded worker must survive (a strict baseline that fails if the leak stops reproducing), and the guarded one must die.
- The wait channel is pipe_write on Linux >= 5.5 and pipe_wait on 5.4. The tests accept either.
- On any other wait channel they kill the worker first and then skip. Failing there used to leave it running under init.
The rest pin down:
- the reaper's selection rules against a fake /proc (yours vs. other users, orphaned vs. live parent, GPU fd vs. not, a process exiting mid-scan);
- the Slurm and once-per-node gates;
- the kill-switch spellings and a misspelling raising (without scanning the real /proc);
- rank arming under Slurm, and leaving a nohup run alone.

Verified in practice

Running the reaper against the node above took all 8 GPUs from ~39 GiB each to 0 MiB, with 0 compute apps left (observed when this PR was written).

Folded in (to keep the stack under Graphite's 50-branch cap):

trainer: DDP devices are per node; eval-mode val subsampling is single-rank #683 — trainer: DDP devices are per node; eval-mode val subsampling is single-rank

🤖 Generated with Claude Code

https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw

RyanPCo · 2026-09-20T18:52:53Z

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

This stack of pull requests is managed by Graphite. Learn more about stacking.

RyanPCo

Automated adversarial review (read-only; findings were checked by a second skeptical pass before posting): 1 medium, 3 low, 1 nit.

RyanPCo · 2026-10-30T20:55:45Z

+                continue
+        except OSError:
+            continue
+        if _ppid(proc_root, pid) != 1:


medium (correctness): find_orphans treats every same-uid, PPid-1 process holding /dev/nvidia* as wreckage. That includes the case arm_rank_pdeathsig and AGENTS.md say this PR supports: a nohup run in a salloc shell that kept running after the shell exited. That run is reparented to init and still holds the GPU, so the next trainHydra launch in the same allocation SIGKILLs it, and its workers die with it through pdeathsig. A detached inference server or nvidia-smi dmon gets the same treatment. Restrict the reap to processes that are clearly leftovers. One way is to match pt_data_worker in /proc/<pid>/comm. Another is to read SLURM_JOB_ID from /proc/<pid>/environ and skip it when it equals the current job or names a job that is still running.

RyanPCo · 2026-10-30T20:55:45Z

+        pytest.fail("worker survived its rank despite PR_SET_PDEATHSIG")
+
+
+def test_orphan_guard_runs_before_a_caller_init():


low (test-quality): This test runs the real die_with_parent inside pytest. That arms PDEATHSIG on the test runner and never disarms it, and if the runner's ppid is 1 (setsid/nohup), the os._exit(1) kills the whole session. The assertion also passes whether or not the guard ran, and in either order. Monkeypatch gpu_orphans.die_with_parent to append ("guard", i) and assert calls == [("guard", 3), ("user", 3)]. _chained looks the name up as a global at call time, so the patch takes effect.

RyanPCo · 2026-10-30T20:55:45Z

+    times wider than it looks from one rank. Eval mode runs on one device
+    (``train`` forces ``devices=1`` after the loaders are built), so it is 1.
    """
+    legacy = "train" if cfg.get("train") else "eval" if cfg.get("eval") else None


low (duplication): This re-implements _resolve_mode; the comment says so. It has already drifted: _resolve_mode tests mode is not None, this uses truthiness. Give _resolve_mode a required: bool = True flag that returns None instead of raising, and call _resolve_mode(cfg, required=False) == "eval". This eval world-size fix and the devices-per-node yaml changes are unrelated to orphan reaping and would review more easily as a separate PR.

RyanPCo · 2026-10-30T20:55:45Z

+        logger.debug("PR_SET_PDEATHSIG unavailable; orphaned workers may leak GPUs")
+        return
+    # pdeathsig does not fire retroactively if the rank died during the fork.
+    if os.getppid() == 1:


low (correctness): os.getppid() == 1 is used here to mean 'the rank died before pdeathsig was armed'. When the training process is PID 1 (a container started as python egomimic/trainHydra.py with no init shim), every worker's parent is PID 1 by design, so each worker calls os._exit(1) at startup and the DataLoader fails. Nothing gates this on Slurm: orphan_guarded wraps every train/val loader, and the norm-stats loader always installs die_with_parent. Record the parent pid before the fork, for example functools.partial(die_with_parent, parent=os.getpid()) in orphan_guarded, and exit only when os.getppid() != parent.

RyanPCo · 2026-10-30T20:55:45Z

+
+logger = logging.getLogger(__name__)
+
+_DISABLE_ENV = "EGOMIMIC_REAP_GPU_ORPHANS"


nit (naming): EGOMIMIC_REAP_GPU_ORPHANS is an enable switch that defaults to on, so _DISABLE_ENV reads backwards in if not str2bool(os.environ.get(_DISABLE_ENV, "1")). Rename it _ENABLE_ENV or _REAP_ENV.

A cancelled DDP job left GPUs allocated with nvidia-smi blaming PIDs absent from /proc. The driver is fine: the memory is held by live processes reparented to init. gpu19 had 126 pt_data_worker survivors across three cancellations, each holding /dev/nvidia0-7, pinning 7 GPUs and ~1.7 TB of RAM. Workers fork after CUDA init, so they inherit the rank's device fds, and ranks open every visible GPU -- one survivor pins the node. A worker also inherits both ends of its result pipe, so a dead rank never sends EPIPE; it blocks in pipe_write and never reaches PyTorch's ppid check. Whole ranks leak the same way: one wedged in torch.cuda.synchronize() inside inductor cudagraph_trees ignores SIGTERM until Slurm gives up (UnkillableStepTimeout) and drops the job with its ranks still spinning. ProctrackType=proctrack/linuxproc walks the PPID tree, so it never sweeps anything reparented to init. egomimic/utils/gpu_orphans.py: * orphan_guarded/die_with_parent put PR_SET_PDEATHSIG on every worker, so the kernel kills it with its rank whatever it is blocked on. Wired into the train and val loaders and the norm-stats pass. * arm_rank_pdeathsig does the same for the rank process. A process whose parent is already gone when it arms exits 1, after one line on stderr saying why (os._exit skips logging's flush). SLURM_JOB_ID is also set inside a salloc shell, so a run started there with nohup is left unarmed (SIGHUP ignored): armed, it died with the shell it was meant to outlive. * reap_orphans runs from trainHydra.main before CUDA init and clears what an earlier job left, so a leak cannot outlive the job that made it. It only signals processes that are ours, reparented to init and holding an nvidia fd. Slurm-gated; EGOMIMIC_REAP_GPU_ORPHANS=0 (or false/no) disables it, any truthy spelling keeps it on, and anything else raises. The module reads PPid from /proc/<pid>/status: in /proc/<pid>/stat, field 2 is comm, and a comm with spaces shifts field 4 far enough to read a live process's PPid as 1. The PDEATHSIG test waits for the unguarded worker to block on its pipe, which Linux >= 5.5 reports as pipe_write and 5.4 as pipe_wait. It accepts either and skips on anything else, killing the worker first: failing there used to leave it running under init. The kill-switch test never scans the real /proc. Co-Authored-By: Claude Opus 5 <[email protected]> Claude-Session: https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw

…e-rank trainer/ddp.yaml and ddp_pi.yaml set devices to gpus_per_node * nodes next to num_nodes, but Lightning reads devices per node and multiplies by num_nodes itself, so a 2-node run asked each node for twice the GPUs it has. The submitit launchers made the same mistake in gres, which Slurm also reads per node. Both now use gpus_per_node; single-node runs are unchanged. trainHydra scales each val head's frames-per-episode K by devices * num_nodes, but eval mode forces devices=1 only after the loaders are built. An eval run of a multi-GPU recipe therefore kept W times the frames per episode that its single rank's limit_val_batches window holds, and the window ran out before the last episodes. _trainer_world_size now returns 1 in eval mode. Both predate this stack (the ddp devices expression is from 2025-01). Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01UkAZwERXcFKjueZQKfRF6J

RyanPCo · 2026-10-30T23:10:01Z

Review findings: status after the fixup pass. The fixes are squashed into this PR's commits, and the whole stack was force-pushed.

Fixed

orphan-guard test arms PDEATHSIG on pytest and doesn't check order: die_with_parent stubbed; asserts [('guard',3),('user',3)]
_trainer_world_size re-implements _resolve_mode: _trainer_world_size is gone: the trainHydra: metric_frames_per_episode: auto, and warn on empty prompts #671 fixup stopped scaling K by the world size, which removed the function. This PR's eval-mode branch of it and its single-rank test went with it, so nothing copies _resolve_mode any more. _resolve_mode is unchanged. The per-node devices config fix and its test are kept.
_DISABLE_ENV names an enable switch: renamed to _REAP_ENV
reap_orphans SIGKILLs live nohup runs: reaps only a PPID-1 GPU holder whose SLURM_JOB_ID is another job that squeue reports as over (squeue failure is treated as alive); AGENTS.md updated; tests added
die_with_parent kills every worker when the rank is PID 1: os.register_at_fork records the forking pid in the child, and die_with_parent checks getppid() against it instead of 1; no call-site changes; tests added

RyanPCo force-pushed the ryanco/pi-scratch-init branch from 211db25 to 6efcdae Compare September 20, 2026 19:09

RyanPCo force-pushed the ryanco/gpu-orphan-reap branch 2 times, most recently from 319a6f3 to 58f8d05 Compare September 21, 2026 18:45

RyanPCo force-pushed the ryanco/pi-scratch-init branch 2 times, most recently from 2fb9f92 to 663b77b Compare September 22, 2026 13:18

RyanPCo force-pushed the ryanco/gpu-orphan-reap branch 3 times, most recently from eb7ebf3 to 61e90ed Compare September 22, 2026 14:06

RyanPCo changed the base branch from ryanco/pi-scratch-init to graphite-base/682 September 22, 2026 14:07

RyanPCo changed the base branch from graphite-base/682 to ryanco/rdt September 22, 2026 14:07

This was referenced Sep 22, 2026

trainer: DDP devices are per node; eval-mode val subsampling is single-rank #683

Closed

eval: log validation videos to wandb #684

Draft

RyanPCo force-pushed the ryanco/gpu-orphan-reap branch from 946d600 to 7bf4804 Compare October 12, 2026 21:50

RyanPCo force-pushed the ryanco/rdt branch from 6213e4d to 6a8a7b9 Compare October 12, 2026 21:50

RyanPCo force-pushed the ryanco/gpu-orphan-reap branch from 7bf4804 to 7246780 Compare October 12, 2026 23:19

RyanPCo mentioned this pull request Oct 19, 2026

lance migration #716

Open

RyanPCo force-pushed the ryanco/rdt branch from 23b0df3 to eef4d01 Compare October 19, 2026 19:51

RyanPCo force-pushed the ryanco/gpu-orphan-reap branch from 7246780 to 222aa11 Compare October 19, 2026 19:51

RyanPCo force-pushed the ryanco/rdt branch from eef4d01 to a7ad86d Compare October 19, 2026 20:22

RyanPCo force-pushed the ryanco/gpu-orphan-reap branch from 222aa11 to 9c51f34 Compare October 19, 2026 20:22

RyanPCo force-pushed the ryanco/rdt branch from a7ad86d to 96c6bbd Compare October 30, 2026 18:27

RyanPCo force-pushed the ryanco/gpu-orphan-reap branch from 9c51f34 to 26995d8 Compare October 30, 2026 18:27

RyanPCo commented Oct 30, 2026

View reviewed changes

RyanPCo and others added 2 commits October 30, 2026 18:51

RyanPCo force-pushed the ryanco/gpu-orphan-reap branch from 26995d8 to 863ec5e Compare October 30, 2026 23:05

Navigation Menu

Repository navigation

slurm: stop cancelled jobs from leaking whole GPUs - #682

slurm: stop cancelled jobs from leaking whole GPUs#682
RyanPCo wants to merge 2 commits into
ryanco/rdtfrom
ryanco/gpu-orphan-reap

RyanPCo commented Sep 20, 2026 •

edited

Loading

Uh oh!

RyanPCo commented Sep 20, 2026 •

edited

Loading

Uh oh!

RyanPCo left a comment

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo commented Oct 30, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

		pytest.fail("worker survived its rank despite PR_SET_PDEATHSIG")


		def test_orphan_guard_runs_before_a_caller_init():


		logger = logging.getLogger(__name__)

		_DISABLE_ENV = "EGOMIMIC_REAP_GPU_ORPHANS"

Conversation

RyanPCo commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Behaviour changes

The problem

Why it happens

The fix

Tests

Verified in practice

Uh oh!

RyanPCo commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

RyanPCo left a comment

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo commented Oct 30, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

RyanPCo commented Sep 20, 2026 •

edited

Loading

RyanPCo commented Sep 20, 2026 •

edited

Loading