diff --git a/AGENTS.md b/AGENTS.md index c081b5e36..17bed48f8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,4 +25,39 @@ use plan mode for anything except extremely simple tasks ## Slurm rules If you're on a slurm cluster, request a GPU before running or testing training. -On sky1/sky2: salloc -p rl2-lab -A rl2-lab --gres=gpu:a40:1 -c 12 --mem=30G \ No newline at end of file +On sky1/sky2: salloc -p rl2-lab -A rl2-lab --gres=gpu:a40:1 -c 12 --mem=30G + +## Cancelled jobs and leaked GPU memory +A cancelled DDP job can leave whole GPUs allocated, with `nvidia-smi` blaming PIDs that are not in +`/proc`. Nothing is wrong with the driver: the memory is held by **live** `pt_data_worker` processes +reparented to init. + +Dataloader workers are forked after CUDA init, so each inherits the rank's `/dev/nvidia*` fds — and +ranks open *every* visible GPU, so one survivor pins the whole node. The worker cannot exit on its +own: it inherited both ends of its result pipe, so a dead rank never gives it an `EPIPE`, it blocks +in `pipe_write`, and PyTorch's parent check (which only runs between `index_queue` timeouts) is +never reached. `ProctrackType=proctrack/linuxproc` then walks the PPID tree and never sees it. + +Whole ranks leak the same way. A rank wedged in `torch.cuda.synchronize()` inside inductor's +`cudagraph_trees` does not answer SIGTERM, Slurm hits `UnkillableStepTimeout` (60s) and gives up +(`srun: error: Timed out waiting for job step to complete`), and the job leaves the queue with its +ranks still spinning on every GPU. Enabling `torch.compile` makes this more likely, not less. + +`egomimic/utils/gpu_orphans.py` owns the fix: + +- `orphan_guarded(params)` / `die_with_parent` set `PR_SET_PDEATHSIG` on every worker, so the kernel + kills it with its rank regardless of what it is blocked on. **Any new `DataLoader` on a training + path needs one of them** — they are already on the train/val loaders and the norm-stats pass. +- `arm_rank_pdeathsig()` does the same for the rank process itself, so a launcher that gives up + takes its ranks with it. Slurm-gated, so a detached interactive run is not shot when its shell + exits. +- `reap_orphans()` runs from `trainHydra.main` before CUDA init and clears what an earlier job left + on the node, so a leak cannot outlive the job that made it. It only ever signals processes that + are yours, reparented to init and holding an nvidia fd; a live job's workers have a live parent, + so they are never touched. Off outside Slurm and on non-zero `SLURM_LOCALID`; disable with + `EGOMIMIC_REAP_GPU_ORPHANS=0`. +- `./reap_gpu_orphans.sh [--kill] ` is the manual version, for nodes you are not about to + submit to. It attaches to your own job on a busy node rather than queueing behind it. + +The cluster-level fix is `ProctrackType=proctrack/cgroup`, which sweeps by cgroup instead of by PPID +tree and would make all of this unnecessary. That needs admin; `/etc/slurm` is read-only to us. diff --git a/egomimic/pl_utils/pl_data_utils.py b/egomimic/pl_utils/pl_data_utils.py index c85cc3098..ac74ccbb3 100644 --- a/egomimic/pl_utils/pl_data_utils.py +++ b/egomimic/pl_utils/pl_data_utils.py @@ -4,6 +4,8 @@ from lightning.pytorch.utilities.combined_loader import CombinedLoader from torch.utils.data import DataLoader, default_collate +from egomimic.utils.gpu_orphans import orphan_guarded + logger = logging.getLogger(__name__) # Val heads that can carry a metric loader, in val_dataloader() order. @@ -181,7 +183,7 @@ def train_dataloader(self): dataset, shuffle=True, collate_fn=self.collate_fn, - **dataset_params, + **orphan_guarded(dataset_params), ) return CombinedLoader(iterables, "max_size_cycle") @@ -196,7 +198,7 @@ def _build_val_style_loader( raise ValueError( f"No dataloader params found for dataset {dataset_name}. Please add {dataset_name} into your data config {kind}_dataloader_params." ) - dataset_params = dict(dataset_params) + dataset_params = orphan_guarded(dataset_params) shuffle = dataset_params.pop("shuffle", False) iterables[dataset_name] = DataLoader( dataset, diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index ed4244424..a76758c5b 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -53,6 +53,7 @@ _apply_unnorm_one, ) from egomimic.utils.env import load_env +from egomimic.utils.gpu_orphans import die_with_parent from egomimic.utils.pose_utils import ( bimanual_cartesian_layout, bimanual_keypoint_layout, @@ -1501,6 +1502,7 @@ def infer_norm_from_dataset( num_workers=num_workers, shuffle=True, generator=torch.Generator().manual_seed(seed), + worker_init_fn=die_with_parent, ) N = len(dataset) if N <= 0: diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 2b94b51b1..30ac91405 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -31,6 +31,7 @@ from egomimic.utils.checkpoint_utils import load_checkpoint_weights from egomimic.utils.compile_cache import set_per_job_compile_cache_dir from egomimic.utils.env import load_env +from egomimic.utils.gpu_orphans import arm_rank_pdeathsig, reap_orphans from egomimic.utils.instantiators import instantiate_callbacks, instantiate_loggers from egomimic.utils.logging_utils import log_hyperparameters from egomimic.utils.pylogger import RankedLogger @@ -726,6 +727,11 @@ def main(cfg: DictConfig) -> Optional[float]: # the job runs main(), so each job keys the cache on its own SLURM_JOB_ID. set_per_job_compile_cache_dir() + # Before CUDA init: a cancelled job can leave dataloader workers -- or whole + # ranks -- holding GPU contexts, and nothing else on this cluster sweeps them. + reap_orphans() + arm_rank_pdeathsig() + # apply extra utilities # (e.g. ask for tags if none are provided in cfg, print cfg tree, etc.) extras(cfg) diff --git a/egomimic/utils/gpu_orphans.py b/egomimic/utils/gpu_orphans.py new file mode 100644 index 000000000..8cdf0287e --- /dev/null +++ b/egomimic/utils/gpu_orphans.py @@ -0,0 +1,163 @@ +"""Reclaim GPUs pinned by dataloader workers that outlived a cancelled job. + +A worker forked after CUDA init holds its rank's ``/dev/nvidia*`` fds, so the +rank's whole context stays resident after the rank is gone -- ``nvidia-smi`` then +reports the memory against a PID that is no longer in ``/proc``. The worker +cannot exit on its own (it blocks in ``pipe_write`` on a result pipe nobody +drains), and ``ProctrackType=proctrack/linuxproc`` walks the PPID tree, so +Slurm never sweeps it once it is reparented to init. + +``pl_data_utils._orphan_guarded`` stops new leaks. This is the other half: a job +clears what an earlier job left on its node, before it touches CUDA itself. + +Nothing is imported from torch here -- this has to run before CUDA init. +""" + +import ctypes +import logging +import os +import signal + +logger = logging.getLogger(__name__) + +_DISABLE_ENV = "EGOMIMIC_REAP_GPU_ORPHANS" +_PR_SET_PDEATHSIG = 1 + + +def die_with_parent(_worker_id: int = 0) -> None: + """Ask the kernel to SIGKILL this dataloader worker when its rank dies. + + PyTorch's own parent check only runs between ``index_queue.get()`` timeouts, + which a worker blocked writing to a full result pipe never reaches. The + kernel does not care what the worker is blocked on. + """ + try: + ctypes.CDLL("libc.so.6", use_errno=True).prctl( + _PR_SET_PDEATHSIG, signal.SIGKILL + ) + except (OSError, AttributeError): # non-Linux, or no libc + 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: + os._exit(1) + + +def orphan_guarded(params: dict) -> dict: + """``params`` with ``die_with_parent`` run before any caller ``worker_init_fn``.""" + params = dict(params) + if params.get("num_workers", 0) == 0: + return params + user_init = params.get("worker_init_fn") + if user_init is None: + params["worker_init_fn"] = die_with_parent + else: + + def _chained(worker_id: int, _user_init=user_init) -> None: + die_with_parent(worker_id) + _user_init(worker_id) + + params["worker_init_fn"] = _chained + return params + + +def _ppid(proc_root: str, pid: str) -> int | None: + try: + with open(f"{proc_root}/{pid}/status") as f: + for line in f: + if line.startswith("PPid:"): + return int(line.split()[1]) + except (OSError, ValueError): + pass + return None + + +def _holds_gpu(proc_root: str, pid: str) -> bool: + fd_dir = f"{proc_root}/{pid}/fd" + try: + fds = os.listdir(fd_dir) + except OSError: # gone, or not ours + return False + for fd in fds: + try: + if os.readlink(f"{fd_dir}/{fd}").startswith("/dev/nvidia"): + return True + except OSError: + continue + return False + + +def arm_rank_pdeathsig() -> bool: + """Have the kernel kill this process if its launcher dies first. + + Workers are not the only thing that leaks. A rank wedged in + ``torch.cuda.synchronize()`` inside inductor's cudagraph_trees does not + answer SIGTERM; Slurm hits ``UnkillableStepTimeout``, gives up, and the job + leaves the queue with its ranks still spinning on every GPU of the node. + + Only armed under Slurm, where the parent is slurmstepd or the launcher that + spawned this rank and never legitimately exits first. Off elsewhere, so a + detached interactive run is not shot when its shell goes away. + """ + if "SLURM_JOB_ID" not in os.environ: + return False + die_with_parent() + return True + + +def find_orphans(proc_root: str = "/proc", uid: int | None = None) -> list[int]: + """PIDs that are ours, reparented to init, and holding a GPU device open. + + A live job's workers always have a live parent, so the PPID-1 test alone + keeps this off anything still running -- ours or anyone else's. + """ + uid = os.getuid() if uid is None else uid + mine = {os.getpid(), os.getppid()} + orphans = [] + for pid in os.listdir(proc_root): + if not pid.isdigit() or int(pid) in mine: + continue + try: + if os.stat(f"{proc_root}/{pid}").st_uid != uid: + continue + except OSError: + continue + if _ppid(proc_root, pid) != 1: + continue + if _holds_gpu(proc_root, pid): + orphans.append(int(pid)) + return sorted(orphans) + + +def reap_orphans(dry_run: bool = False) -> list[int]: + """Kill this node's leaked GPU holders. Returns the PIDs signalled. + + Off outside Slurm: only there do we know every process of ours on the node + belongs to a job, so that a PPID-1 GPU holder can only be wreckage. + """ + if os.environ.get(_DISABLE_ENV, "1") != "1": + return [] + if "SLURM_JOB_ID" not in os.environ: + return [] + if os.environ.get("SLURM_LOCALID", "0") != "0": # once per node + return [] + + orphans = find_orphans() + if not orphans: + return [] + logger.warning( + "Reaping %d orphaned GPU holder(s) left by an earlier job on %s: %s " + "(set %s=0 to disable)", + len(orphans), + os.uname().nodename, + orphans, + _DISABLE_ENV, + ) + if dry_run: + return orphans + for pid in orphans: + try: + os.kill(pid, signal.SIGKILL) + except OSError as e: + logger.warning("could not kill orphan %d: %s", pid, e) + return orphans diff --git a/reap_gpu_orphans.sh b/reap_gpu_orphans.sh new file mode 100755 index 000000000..264ab52d6 --- /dev/null +++ b/reap_gpu_orphans.sh @@ -0,0 +1,76 @@ +#!/usr/bin/env bash +# Find (and optionally kill) your own dataloader workers that outlived a cancelled +# job and are still pinning GPU memory. +# +# A worker forked after CUDA init holds its rank's /dev/nvidia* fds, so the rank's +# whole context stays resident even after the rank is gone -- nvidia-smi then +# reports memory against a PID that is no longer in /proc. ProctrackType on this +# cluster is proctrack/linuxproc, which tracks the PPID tree, so workers that got +# reparented to init are never swept. +# +# ./reap_gpu_orphans.sh gpu19 # report only +# ./reap_gpu_orphans.sh --kill gpu19 # reclaim +# ./reap_gpu_orphans.sh --kill gpu11 gpu19 +# +# Only processes that are yours, orphaned (PPID 1) and holding an nvidia fd are +# ever signalled, so this is safe to run against a node with live jobs on it. +set -uo pipefail + +KILL=0 +[[ ${1:-} == --kill ]] && { KILL=1; shift; } +if [[ $# -eq 0 ]]; then + echo "usage: $0 [--kill] [node...]" >&2 + exit 2 +fi + +PROBE=' +ME=$(id -u); victims=() +for d in /proc/[0-9]*; do + p=${d#/proc/} + [[ $(stat -c %u "$d" 2>/dev/null) == "$ME" ]] || continue + [[ $(awk "{print \$4}" "$d/stat" 2>/dev/null) == 1 ]] || continue + ls -l "$d/fd" 2>/dev/null | grep -q /dev/nvidia || continue + victims+=("$p") +done +echo "node=$(hostname -s) orphans=${#victims[@]}" +[[ ${#victims[@]} -eq 0 ]] && exit 0 +for p in "${victims[@]}"; do + echo " $p $(cat /proc/$p/comm 2>/dev/null) rss=$(awk "/VmRSS/{print \$2}" /proc/$p/status 2>/dev/null)kB age=$(ps -o etime= -p $p 2>/dev/null | tr -d " ")" +done +echo "before: $(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader | tr "\n" " ")" +if [[ $KILL == 1 ]]; then + kill -9 "${victims[@]}" 2>/dev/null + for i in 1 2 3 4 5 6; do + sleep 4 + n=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader | wc -l) + [[ $n -eq 0 ]] && break + done + echo "after: $(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader | tr "\n" " ")" + echo "compute apps left: $n" +fi +' + +probe() { # probe ; returns srun's status, not grep's + "$@" bash -c "KILL=$KILL; $PROBE" 2>&1 | grep -v '^srun: job' + return "${PIPESTATUS[0]}" +} + +rc=0 +for node in "$@"; do + echo "=== $node ===" + # --immediate everywhere: the partition is EXCLUSIVE, so a fresh allocation on a + # busy node queues forever, and a step inside a job whose training step already + # holds every CPU never starts. Better to say so than to hang. + job=$(squeue -h -u "$USER" -w "$node" -t R -o %i | head -1) + if [[ -n $job ]]; then + echo "(node held by your job $job; attaching to it)" + probe srun --overlap --immediate=20 --jobid="$job" -n1 -c 1 && continue + echo "(could not get a step in $job -- it is probably mid-training)" + fi + probe srun -p loaner -A loaner --nodelist="$node" -n1 -c 2 --mem=4G --overlap \ + --immediate=30 -t 5 && continue + echo "$node is busy; rerun when it is idle. A job starting there cleans it" >&2 + echo "automatically anyway (egomimic.utils.gpu_orphans.reap_orphans)." >&2 + rc=1 +done +exit "$rc" diff --git a/tests/unit/test_worker_orphan_guard.py b/tests/unit/test_worker_orphan_guard.py new file mode 100644 index 000000000..468087434 --- /dev/null +++ b/tests/unit/test_worker_orphan_guard.py @@ -0,0 +1,179 @@ +"""Dataloader workers must die with their rank. + +Workers are forked after CUDA init, so they inherit the rank's ``/dev/nvidia*`` +fds. One survivor keeps every GPU context on the node alive, and +``proctrack/linuxproc`` will not sweep it because it has been reparented to init. +""" + +import os +import subprocess +import sys +import textwrap +import time + +import pytest + +from egomimic.utils.gpu_orphans import ( + arm_rank_pdeathsig, + die_with_parent, + find_orphans, + orphan_guarded, + reap_orphans, +) + +# A worker that inherits both ends of its result pipe never gets an EPIPE when +# the rank dies; it blocks in pipe_write, so PyTorch's own ppid check -- which +# only runs between index_queue timeouts -- is never reached. +RANK = textwrap.dedent( + """ + import os, sys, time + sys.path.insert(0, {root!r}) + from egomimic.utils.gpu_orphans import die_with_parent + guard = sys.argv[1] == "guard" + r, w = os.pipe() + if os.fork() == 0: + if guard: + die_with_parent(0) + with open(sys.argv[2], "w") as f: + f.write(str(os.getpid())) + while True: + os.write(w, b"x" * 65536) + time.sleep(600) + """ +) + + +def _alive(pid: int) -> bool: + return os.path.isdir(f"/proc/{pid}") + + +def _pid_written(pidfile) -> bool: + return pidfile.exists() and pidfile.read_text().strip().isdigit() + + +def _run_rank(tmp_path, mode: str) -> int: + """Start a rank, SIGKILL it once its worker is wedged, return the worker pid.""" + root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + pidfile = tmp_path / f"pid_{mode}" + rank = subprocess.Popen( + [sys.executable, "-c", RANK.format(root=root), mode, str(pidfile)] + ) + try: + deadline = time.time() + 180 # the rank imports torch first + while time.time() < deadline and not _pid_written(pidfile): + time.sleep(0.2) + assert _pid_written(pidfile), "worker never started" + worker = int(pidfile.read_text()) + deadline = time.time() + 30 + while time.time() < deadline: + if open(f"/proc/{worker}/wchan").read() == "pipe_write": + break + time.sleep(0.2) + else: + pytest.fail("worker never blocked in pipe_write") + finally: + rank.kill() + rank.wait() + return worker + + +@pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only") +def test_unguarded_worker_outlives_its_rank(tmp_path): + worker = _run_rank(tmp_path, "noguard") + time.sleep(3) + try: + assert _alive(worker), "baseline no longer reproduces the leak" + finally: + os.kill(worker, 9) + + +@pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only") +def test_guarded_worker_dies_with_its_rank(tmp_path): + worker = _run_rank(tmp_path, "guard") + deadline = time.time() + 15 + while time.time() < deadline and _alive(worker): + time.sleep(0.2) + if _alive(worker): + os.kill(worker, 9) + pytest.fail("worker survived its rank despite PR_SET_PDEATHSIG") + + +def test_orphan_guard_runs_before_a_caller_init(): + calls = [] + params = orphan_guarded( + {"num_workers": 2, "worker_init_fn": lambda i: calls.append(("user", i))} + ) + params["worker_init_fn"](3) + assert calls == [("user", 3)] + + +def test_orphan_guard_is_a_noop_without_workers(): + assert orphan_guarded({"num_workers": 0}) == {"num_workers": 0} + + +def test_orphan_guard_installs_the_default(): + assert orphan_guarded({"num_workers": 4})["worker_init_fn"] is die_with_parent + + +def _fake_proc(root, pid: int, *, ppid: int, fds: list[str]) -> None: + d = root / str(pid) + (d / "fd").mkdir(parents=True) + (d / "status").write_text(f"Name:\tpt_data_worker\nPPid:\t{ppid}\n") + for i, target in enumerate(fds): + os.symlink(target, d / "fd" / str(i)) + + +def test_find_orphans_picks_only_ours_orphaned_and_on_a_gpu(tmp_path): + uid = os.getuid() + gpu = ["/dev/nvidia0", "/dev/nvidiactl"] + _fake_proc(tmp_path, 100, ppid=1, fds=gpu) # the leak + _fake_proc(tmp_path, 101, ppid=42, fds=gpu) # live job's worker + _fake_proc(tmp_path, 102, ppid=1, fds=["/dev/null"]) # no GPU + assert find_orphans(str(tmp_path), uid=uid) == [100] + + +def test_find_orphans_skips_other_users(tmp_path): + uid = os.getuid() + _fake_proc(tmp_path, 200, ppid=1, fds=["/dev/nvidia0"]) + assert find_orphans(str(tmp_path), uid=uid + 1) == [] + + +def test_find_orphans_tolerates_a_process_exiting_mid_scan(tmp_path): + uid = os.getuid() + (tmp_path / "300").mkdir() # no status/fd: raced away + _fake_proc(tmp_path, 301, ppid=1, fds=["/dev/nvidia0"]) + assert find_orphans(str(tmp_path), uid=uid) == [301] + + +def test_reap_is_off_outside_slurm(monkeypatch): + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + assert reap_orphans() == [] + + +def test_reap_runs_once_per_node(monkeypatch): + monkeypatch.setenv("SLURM_JOB_ID", "1") + monkeypatch.setenv("SLURM_LOCALID", "3") + assert reap_orphans() == [] + + +def test_reap_honours_the_kill_switch(monkeypatch): + monkeypatch.setenv("SLURM_JOB_ID", "1") + monkeypatch.setenv("SLURM_LOCALID", "0") + monkeypatch.setenv("EGOMIMIC_REAP_GPU_ORPHANS", "0") + assert reap_orphans() == [] + + +def test_rank_pdeathsig_is_off_outside_slurm(monkeypatch): + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + assert arm_rank_pdeathsig() is False + + +def test_rank_pdeathsig_arms_under_slurm(monkeypatch): + import ctypes + + monkeypatch.setenv("SLURM_JOB_ID", "1") + try: + assert arm_rank_pdeathsig() is True + finally: + # do not leave pytest itself armed + ctypes.CDLL("libc.so.6").prctl(1, 0)