feat(aorlta): two runs on one node: engine ports chosen free, GPU and directory claims between jobs

aorlta/jobs/aorlta-run.sbatch

@@ -2,8 +2,21 @@
 # AORLTA: one run of the GRPO loop (`aorlta.orchestrator.run`) on one node, such as configs/standin.yaml's 3-iteration
 # stand-in run. The TokenSpeed image is built in-job if the node lacks it, then the trainer image from
 # docker/Dockerfile.trainer; an EXIT/TERM trap removes this job's containers, its /dev/shm slots and every image it
-# built or pulled. docker runs in the background with `wait`, so a TERM (scancel) runs the trap at once: bash defers
-# traps while a foreground command runs, and a killed docker client leaves its container running.
+# built or pulled that no other job on the node may be using. docker runs in the background with `wait`, so a TERM
+# (scancel) runs the trap at once: bash defers traps while a foreground command runs, and a killed docker client
+# leaves its container running.
+#
+# Two runs on one node. The job asks for the whole node (the --exclusive directive below); a command-line option
+# overrides it, so `sbatch --exclusive=user -w <node> ...` lets your own jobs share a node, on a partition that allows
+# it (one with OverSubscribe=EXCLUSIVE never puts two jobs on a node). Their configs must name
+# disjoint GPUs (rollout.engines, trainer.devices, reward.remote's GPUs, and any GPUs a task's own checks use), and
+# each its own run id and reward.remote.socket_dir. The job claims the GPUs and its two /dev/shm directories
+# (<sync.dir>/<run id> and the socket directory) in files under /dev/shm/aorlta-claims, and ends with exit 89 when a
+# job Slurm still runs holds one, or a directory inside or around one (a GPU a task's own checks use is not claimed:
+# the job cannot see it). Engine ports
+# need nothing with rollout.port_base null (the default: each engine start picks a free window); two runs that both
+# fix a port_base must give windows that do not overlap. A container or claim of a job Slurm no longer runs is what
+# a cancelled or killed job left: the job removes or takes it over, and never touches another live job's.
 #
 # With `reward.remote` in the config (the standin-gpu profile's ConSan and ASan legs), the checkers run in
 # rocm/dev-ubuntu-24.04:7.0.2 (glibc 2.39 for the ConSan hook, the ROCm 7.0.2 ASan runtime): first aorta's ASan
@@ -57,6 +70,8 @@
 #   AORLTA_EXCLUDE_NODES   default none. Nodes this job refuses to run on (exit 90), as a Slurm host list. #SBATCH
 #                          lines cannot read the environment, so give Slurm the same list when submitting:
 #                            sbatch --exclude="$AORLTA_EXCLUDE_NODES" jobs/aorlta-run.sbatch
+#   AORLTA_CLAIMS_DIR      default /dev/shm/aorlta-claims. Where jobs on this node claim GPUs and directories: node-
+#                          local, and the same for every job of the node, or two jobs cannot see each other's claims
 # A container sees each directory at its host path. Mounted besides the directories above: the package root, /dev/shm
 # (reward.remote.socket_dir must be under it, and sync.dir is by default), and from the config run.root,
 # model.base_dir and model.adapter_dir (the adapter read-only; for a Hugging Face snapshot, the model's whole cache
@@ -114,7 +129,118 @@ BUILT_TRAINER=0
 PULLED_CHK=0
 SOCK_DIR=
 VRAM_PID=
+CLAIMS=${AORLTA_CLAIMS_DIR:-/dev/shm/aorlta-claims}
+CLAIMED=()
+OWN_DIRS=0  # this job holds the claims on <sync.dir>/<run id> and the socket directory, so its trap may remove them
 stamp() { echo "=== [$(date -u +%FT%TZ)] $*"; }
+job_ended() {  # job_ended ID: Slurm no longer runs job ID: it does not know the job, or names a final state. Any other
+  # answer, and no answer at all (squeue missing, the controller unreachable), counts as running, so a live job's
+  # containers, images and claims are never taken for left-overs.
+  local out
+  if ! out=$(squeue -h -j "$1" -o %T 2>&1); then
+    grep -qi 'invalid job id' <<< "$out"
+    return
+  fi
+  case "$out" in
+    ''|BOOT_FAIL|CANCELLED|COMPLETED|DEADLINE|FAILED|NODE_FAIL|OUT_OF_MEMORY|PREEMPTED|REVOKED|SPECIAL_EXIT|TIMEOUT)
+      return 0 ;;
+  esac
+  return 1
+}
+stale_containers() {  # the aorlta-run-<job id>[-part] containers on this node whose job Slurm no longer runs
+  local name id
+  local -A ended=()
+  while read -r name; do
+    id=${name#aorlta-run-}
+    id=${id%%-*}
+    [ "$id" = "$SLURM_JOB_ID" ] && continue
+    if [ -z "${ended[$id]:-}" ]; then
+      if job_ended "$id"; then ended[$id]=1; else ended[$id]=0; fi
+    fi
+    [ "${ended[$id]}" = 1 ] && echo "$name"
+  done < <(docker ps -a --format '{{.Names}}' | grep -xE 'aorlta-run-[0-9]+(-[a-z]+)?')
+}
+other_jobs() {  # the other Slurm jobs running on this node, one per line; "unknown" if Slurm cannot say
+  local out
+  out=$(squeue -h -w "${SLURMD_NODENAME:-$(hostname -s)}" -o %i 2>/dev/null) || { echo unknown; return; }
+  grep -vxF -- "$SLURM_JOB_ID" <<< "$out" | grep .
+}
+remove_image() {  # remove_image IMAGE: remove an image this job built or pulled, unless another job may use it: a
+  # container of another job was made from it (or from an image built on it), or another job runs on this node (it
+  # may have found the image when it started and not have run it yet). A kept image is recorded under
+  # $CLAIMS/images, and every job's cleanup tries it again (retry_kept_images), so the last job to leave removes it.
+  local users jobs rec
+  rec=$CLAIMS/images/$(printf %s "$1" | sha256sum | cut -c1-16)
+  users=$(docker ps -a --filter "ancestor=$1" --format '{{.Names}}' 2>/dev/null | grep -vxE "$CNAME(-[a-z]+)?")
+  jobs=$(other_jobs)
+  if [ -n "$users$jobs" ]; then
+    echo "kept $1: another job may use it (containers: ${users:-none}; jobs on this node: ${jobs:-none})" | tr '\n' ' '
+    echo
+    mkdir -p "$CLAIMS/images" 2>/dev/null && chmod 1777 "$CLAIMS/images" 2>/dev/null
+    printf '%s\n' "$1" > "$rec" 2>/dev/null
+    return
+  fi
+  if docker rmi -f "$1" >/dev/null 2>&1; then
+    echo "removed $1"
+    rm -f "$rec"
+  elif ! docker image inspect "$1" >/dev/null 2>&1; then
+    rm -f "$rec"  # gone already
+  fi
+}
+retry_kept_images() {  # the images an earlier job built or pulled and kept for another job (remove_image)
+  local f img
+  for f in "$CLAIMS"/images/*; do
+    [ -f "$f" ] || continue
+    img=$(head -n 1 "$f")
+    if [ -z "$img" ] || ! docker image inspect "$img" >/dev/null 2>&1; then rm -f "$f"; continue; fi
+    remove_image "$img"
+  done
+}
+claim() {  # claim WHAT KEY...: hold each KEY for this job (a file under $CLAIMS naming its job id). A key held by a job
+  # Slurm still runs ends this job, exit 89; one whose job ended is taken over. A directory (WHAT "directory PATH",
+  # PATH canonical: canon_dir) also ends the job when a live job holds one inside it or one it is inside, since
+  # either job's cleanup removes its directory recursively. Under flock, so two jobs taking over the same left-over
+  # claim cannot both get it.
+  local what=$1 key f owner fd held mine other
+  shift
+  mkdir -p "$CLAIMS" && chmod 1777 "$CLAIMS" 2>/dev/null  # other users' jobs on a shared node claim here too
+  touch "$CLAIMS/.lock" 2>/dev/null && chmod 666 "$CLAIMS/.lock" 2>/dev/null
+  exec {fd}>>"$CLAIMS/.lock" && flock -w 60 "$fd" || { echo "[guard] cannot lock $CLAIMS/.lock" >&2; exit 89; }
+  for key in "$@"; do
+    f=$CLAIMS/$key
+    if [ "${what#directory /}" != "$what" ]; then
+      mine=${what#directory }
+      for held in "$CLAIMS"/dir-*; do
+        [ -f "$held" ] && [ "$held" != "$f" ] || continue
+        owner=$(cut -d' ' -f1 "$held" 2>/dev/null)
+        [ -n "$owner" ] && [ "$owner" != "$SLURM_JOB_ID" ] || continue
+        other=$(cut -d' ' -f3- "$held" 2>/dev/null)
+        case "$mine/" in "$other/"*) ;; *) case "$other/" in "$mine/"*) ;; *) continue ;; esac ;; esac
+        if ! job_ended "$owner"; then
+          echo "[guard] job $owner, which Slurm still runs (or Slurm cannot say), holds directory $other on this node," \
+            "which overlaps $mine: give this run another run id and socket directory" >&2
+          exec {fd}>&-
+          exit 89
+        fi
+      done
+    fi
+    owner=$(cut -d' ' -f1 "$f" 2>/dev/null)
+    if [ -n "$owner" ] && [ "$owner" != "$SLURM_JOB_ID" ]; then
+      if ! job_ended "$owner"; then
+        echo "[guard] job $owner, which Slurm still runs (or Slurm cannot say), holds $what on this node ($f):" \
+          "give this run other GPUs, or another run id and socket directory" >&2
+        exec {fd}>&-
+        exit 89
+      fi
+      stamp "taking over $key from job $owner, which has ended"
+    fi
+    echo "$SLURM_JOB_ID $what" > "$f" && CLAIMED+=("$f") || { echo "[guard] cannot write $f" >&2; exec {fd}>&-; exit 89; }
+  done
+  exec {fd}>&-
+}
+canon_dir() { realpath -m -- "$1"; }  # one spelling per directory: symlinks resolved, no trailing slash, no . or ..
+dir_key() { echo "dir-$(canon_dir "$1" | tr -d '\n' | sha256sum | cut -c1-16)"; }
+claim_dir() { local d; d=$(canon_dir "$1"); claim "directory $d" "$(dir_key "$d")"; }
 vram_sampler() {  # every 10 s, from sysfs (no GPU context): UTC seconds, PCI address, VRAM used and total in bytes
   local c
   while :; do
@@ -217,7 +343,11 @@ clean_img() {  # an image this node has, for removing root-owned files
   local i
   for i in "$BASE" "$CHK_IMG"; do docker image inspect "$i" >/dev/null 2>&1 && { echo "$i"; return; }; done
 }
-cleanup() {
+CLEANED=0
+cleanup() {  # once: the TERM/INT trap's exit runs the EXIT trap too, and by then this job's claims are released, so a
+  # second pass could remove the directories of a job that claimed them in between
+  [ "$CLEANED" = 1 ] && return
+  CLEANED=1
   [ -n "$VRAM_PID" ] && kill "$VRAM_PID" 2>/dev/null
   docker ps -a --format '{{.Names}}' | grep -xE "$CNAME(-[a-z]+)?" | xargs -r docker rm -f >/dev/null 2>&1
   find "$OUT" -maxdepth 3 -name 'gpucore.*' -type f -delete 2>/dev/null
@@ -226,25 +356,52 @@ cleanup() {
   add_mounts mounts rw /dev/shm "$OUT" "$RUN_ROOT" 2>/dev/null
   # The GPU legs' logs, harnesses and the candidates' parsed sources (stand-in data) are kept, small, before the
   # node directory goes: they are how an unsettled ConSan or ASan leg is read afterwards.
+  # (a job that did not get its claims leaves the directories alone: another live job holds them)
   [ -n "$img" ] && docker run --rm --network none --ulimit core=0 "${mounts[@]}" --entrypoint bash "$img" -c \
-    "if cd '$SYNC_DIR/$RUN_ID/scratch/reward' 2>/dev/null; then
+    "if [ $OWN_DIRS = 1 ] && cd '$SYNC_DIR/$RUN_ID/scratch/reward' 2>/dev/null; then
        find pkg \( -path '*/san.*' \( -name '*.log' -o -name aorlta_harness.hip -o -name '*.json' \) \
          -o -path '*/gate.parse-*' -name source.hip \) -size -2M -print0 | tar --null -czf '$OUT/gpu-leg-evidence.tgz' -T -;
        cd /; fi
-     rm -rf '$SYNC_DIR/$RUN_ID' ${SOCK_DIR:+'$SOCK_DIR'}; chown -R $(id -u):$(id -g) '$OUT';
+     [ $OWN_DIRS = 1 ] && rm -rf '$SYNC_DIR/$RUN_ID' ${SOCK_DIR:+'$SOCK_DIR'}; chown -R $(id -u):$(id -g) '$OUT';
      [ -d '$RUN_DIR' ] && chown -R $(id -u):$(id -g) '$RUN_DIR'" >/dev/null 2>&1
-  [ "$BUILT_TRAINER" = 1 ] && docker rmi -f "$IMG" >/dev/null 2>&1 && echo "removed $IMG"
-  [ "$BUILT_BASE" = 1 ] && docker rmi -f "$BASE" >/dev/null 2>&1 && echo "removed $BASE"
-  [ "$PULLED_CHK" = 1 ] && docker rmi -f "$CHK_IMG" >/dev/null 2>&1 && echo "removed $CHK_IMG (pulled by this job)"
+  [ "$BUILT_TRAINER" = 1 ] && remove_image "$IMG"
+  [ "$BUILT_BASE" = 1 ] && remove_image "$BASE"
+  [ "$PULLED_CHK" = 1 ] && remove_image "$CHK_IMG"
+  retry_kept_images
+  local f
+  for f in "${CLAIMED[@]}"; do
+    [ "$(cut -d' ' -f1 "$f" 2>/dev/null)" = "$SLURM_JOB_ID" ] && rm -f "$f"
+  done
   echo "containers left: $(docker ps -a --format '{{.Names}}' | grep -cxE "$CNAME(-[a-z]+)?")"
 }
 trap cleanup EXIT
 trap 'cleanup; exit 143' TERM INT
 
 stamp "node $(hostname -s) job $SLURM_JOB_ID run $RUN_ID config $CONFIG components $COMPONENTS resume $RESUME commit $COMMIT"
-# The node is exclusive: a container of this run left by an earlier, cancelled job is removed before starting.
-stale=$(docker ps -a --format '{{.Names}}' | grep -xE 'aorlta-run-[0-9]+(-[a-z]+)?')
-[ -n "$stale" ] && { stamp "removing stale containers"; echo "$stale" | xargs -r docker rm -f >/dev/null 2>&1; }
+# Before anything uses them (the checker canary empties the socket directory): this run's /dev/shm directories, and
+# the GPUs it uses on this node, which are the engines', the trainer ranks' own (with trainer.devices null they share
+# the engines') and the checker servers'. The fakes use no GPU.
+claim_dir "$SYNC_DIR/$RUN_ID"
+[ -n "$SOCK_DIR" ] && claim_dir "$SOCK_DIR"
+OWN_DIRS=1
+RUN_GPUS=()
+if [ "$COMPONENTS" = real ]; then
+  read -ra RUN_GPUS < <(PYTHONPATH="$CODE/src$XPATH" "$PY" -c "
+from aorlta.config import load_config
+c = load_config('$CODE/$CONFIG')
+gpus = set(c.rollout.engines)
+if c.orchestrator.mode != 'score':
+    gpus |= set(c.trainer.devices or ())
+if c.reward.remote is not None:
+    for name in c.reward.remote.adapters:
+        gpus |= set(c.reward.remote.gpus(name))
+print(*sorted(gpus))") || { echo "[guard] $CONFIG does not load" >&2; exit 94; }
+fi
+for g in "${RUN_GPUS[@]}"; do claim "GPU $g" "gpu$g"; done
+stamp "claimed GPUs ${RUN_GPUS[*]:-none} and the run's /dev/shm directories"
+# A container left by a job Slurm no longer runs (cancelled, killed) is removed; another live job's are left alone.
+stale=$(stale_containers)
+[ -n "$stale" ] && { stamp "removing stale containers: $(echo $stale)"; echo "$stale" | xargs -r docker rm -f >/dev/null 2>&1; }
 df -h /dev/shm | tail -1
 
 # --- the remote checkers, first: a node whose canary or pre-flight fails ends the job before anything is built ---

aorlta/src/aorlta/config.py

@@ -33,6 +33,7 @@
 
 from .atomic import atomic_write_json, is_partial
 from .ids import IdError, sha256_file, sha256_json, validate_run_id
+from .rollout.ports import PORT_STRIDE
 
 Source = Literal["standin"]  # the stand-in data source; a new source adds a value here
 RUN_PLAN_SCHEMA = "aorlta.run_plan/1"
@@ -439,9 +440,19 @@ class RolloutSection(_Section):
     # with near-unique answers and few prompts may raise it: one such group in 8 is 12.5%. The engines' start-up
     # self-test checks the sampler on its own.
     identical_groups_stop: float = 0.05
+    # the engines' host ports. null: each engine start picks a free window of 2000 ports at random, outside the
+    # kernel's ephemeral range, and tries another if one is taken before the engine binds it (rollout/ports.py), so
+    # two runs can share a node. An int in [1024, 65536 - 2000 x engines]: engine i on port_base + 2000 i, started
+    # once; 20000 is the layout every run used before this key. The ports used are in <run_dir>/engines.jsonl.
+    port_base: int | None = None
 
     def _validate(self) -> None:
         _check(len(self.engines) >= 1, "engines: at least one engine GPU is needed")
+        if self.port_base is not None:
+            top = 65536 - PORT_STRIDE * len(self.engines)
+            _check(1024 <= self.port_base <= top,
+                   f"port_base: must be null or in [1024, {top}] ({len(self.engines)} engines, {PORT_STRIDE} ports "
+                   f"each), got {self.port_base}")
         _check(all(g >= 0 for g in self.engines), f"engines: GPU indices must be >= 0, got {list(self.engines)}")
         _check(len(set(self.engines)) == len(self.engines), f"engines: a GPU is listed twice: {list(self.engines)}")
         _check(self.K >= 2, f"K: must be >= 2 so a group can have an advantage, got {self.K}")