hpt: build the policy on CPU and let Lightning place it

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
+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] <node>` 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.

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,

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:

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)

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

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> [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 <srun args...>; 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"