From 990c3c8638f45933bad26f666d1e2e1e85d6f5e6 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 2 Oct 2026 19:41:19 -0400 Subject: [PATCH 1/3] sp run --in-allocation: run a campaign inside one SLURM allocation Runs both invocations under snakemake's local executor inside the allocation sp is called from, so a small campaign pays one queue wait instead of one per job. sp derives --cores from the cpuset and --resources mem_mb from the tightest cgroup memory limit (less a reserve), refuses to run outside an allocation, and defaults the node-local tile store to $SLURM_TMPDIR/tile-store through the existing tile_store_root rebind. README documents sizing and the resume caveat. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Dj4HnkNqqPkVCeDeUGEGuz --- workflow/README.md | 37 +++++++++++++++++++++ workflow/bin/sp | 81 +++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 114 insertions(+), 4 deletions(-) diff --git a/workflow/README.md b/workflow/README.md index 4f0b135c3..6cbad9490 100644 --- a/workflow/README.md +++ b/workflow/README.md @@ -230,6 +230,43 @@ neighbours (queue-latency amortization, per the PRD) are **not yet wired**: they require labels in `workflow/rules/*.smk`, out of scope for this profile-only pass. +## Execution mode: one allocation + +One SLURM job per rule pays one queue wait per job, and at a low fair share a +small campaign is mostly queue wait. `sp run --in-allocation` runs both +invocations under snakemake's local executor inside the allocation it is +called from, so the campaign waits once: + +```bash +#!/usr/bin/env bash +#SBATCH --account= +#SBATCH --nodes=1 --ntasks=1 +#SBATCH --cpus-per-task=96 --mem=383000M +#SBATCH --time=4:00:00 +SP_PROFILE=nibi /path/to/checkout/workflow/bin/sp run --in-allocation -c run.yaml +``` + +`sp` reads the budget from the allocation itself (cores from the cpuset, +memory from the tightest cgroup limit, less a reserve for snakemake) and +passes `--executor local --cores N --resources mem_mb=M`; later explicit +flags override them. Everything else is the profile's. `group:` labels are +inactive under local execution, so `--groups` does nothing; the tile_shape +members still share their node-local store because every job runs on this +node. That store defaults to `$SLURM_TMPDIR/tile-store`, which Slurm reclaims +when the job ends (an explicit `tile_store_root:` wins). + +A rule's `mem_mb` becomes a packing budget rather than a per-job cgroup: a job +over its share is killed only if the allocation as a whole runs out. Size the +allocation so the widest rule's last retry fits (a retry asking more than the +allocation has stops the run), at no more than 4 GB per core so memory does +not raise the billing. Keep it to one node: the tile store is node-local, so +a multi-node allocation would split a tile's vignets from its ngmix chunks. +For the same reason a time limit that ends the job mid-tile leaves that tile's +`tile_vignets.json` without its store, and a resume trips tile.smk's +"node-local vignette store is missing" guard. Delete the manifest of each tile +that has no `final_cat` yet before resuming, as the guard's message says, and +give the time limit room. + ## Layout ``` diff --git a/workflow/bin/sp b/workflow/bin/sp index cb6754108..241c920b1 100755 --- a/workflow/bin/sp +++ b/workflow/bin/sp @@ -15,6 +15,11 @@ # FILE is a user-specified config file that can be used # to add and/or overwrite entries in the default config # file workflow/config.yaml (read automatically) +# sp run --in-allocation [-c FILE] [ARGS...] +# the same, but every job runs HERE, inside the SLURM +# allocation sp is called from, under snakemake's local +# executor (see "one allocation" below). Refuses to run +# outside an allocation. # sp report [ARGS...] emit run_report.json now (mid-run is fine). # sp container VERB manage the image every job runs inside: pull, status, # sandbox, exec, resolve. `sp container --help` documents @@ -68,13 +73,17 @@ PROFILE="${SP_PROFILE:-nibi}" # the delegated command (`sp container exec python -c ...`) or for snakemake's # own escape hatch (`sp -- ...`) is left alone. `--config` is NOT sp's: keeping # it free lets snakemake's own `--config k=v` (what the Snakefile's flag() -# reads) pass through untouched. +# reads) pass through untouched. --in-allocation is sp's too (see "one +# allocation" below); snakemake never sees it. RUN_CONFIG="${SP_RUN_CONFIG:-}" +IN_ALLOCATION=0 _sp_args=() while [ $# -gt 0 ]; do case "$1" in --) shift; _sp_args+=("$@"); break ;; + --in-allocation) + IN_ALLOCATION=1; shift ;; -c|--config-file) [ $# -ge 2 ] || { echo "sp: $1 needs a file" >&2; exit 2; } RUN_CONFIG="$2"; shift 2 ;; @@ -153,6 +162,70 @@ fi STATE_DIR="${SP_STATE_DIR:-${RUN_DIR}-state}" [ "$NEEDS_CAMPAIGN" = 0 ] || mkdir -p "$STATE_DIR" +# --- one allocation (`sp run --in-allocation`) ----------------------------- +# The profile's SLURM executor submits one job per rule (or group) and so pays +# one queue wait per job; at a low fair share a small campaign is mostly queue +# wait. --in-allocation runs the same two invocations under snakemake's LOCAL +# executor instead, inside the allocation sp is called from (an sbatch script +# or salloc), so the campaign pays one wait. Everything else is the profile's: +# the apptainer wrapping, the rerun triggers, keep-going. +# +# What changes in local execution, and why each is safe: +# * `group:` labels are inactive ("Group jobs: inactive (local execution)"), +# so --groups and the rule files' groups do nothing. tile_shape's members +# still share the node-local store because they all run on THIS node. +# * The allocation's cores and memory are the scheduler's global budget, +# read here: cores from the cpuset (`nproc`), memory from the tightest +# cgroup limit above this process (an srun step's own --mem, not the +# job's), less a reserve for snakemake and the localrules. A rule's mem_mb +# is then a packing budget, not a cgroup: a job over its share is killed +# only if the whole allocation runs out. Attempt-scaled retries still +# raise the budget, and a retry asking more than the allocation has stops +# the run (snakemake's "Job needs mem_mb=..."), so size the allocation for +# the widest rule's third attempt. +# * The tile store defaults to $SLURM_TMPDIR/tile-store. The profile cannot +# name $SLURM_TMPDIR (its `--bind /local` comment says why), but sp runs +# inside the job here and knows the value, so it rebinds /local/scratch +# there exactly as `tile_store_root:` does. Slurm reclaims it at job end, +# and no other campaign on the node can share it. An explicit +# `tile_store_root:` still wins. +# Explicit --cores/--resources after it override the derived values. +ALLOC_ARGS=() +TILE_STORE_ROOT="" +[ "$NEEDS_CAMPAIGN" = 0 ] || TILE_STORE_ROOT="$(cfg tile_store_root)" +if [ "$IN_ALLOCATION" = 1 ]; then + [ "${1:-}" = run ] || { echo "sp: --in-allocation applies to \`sp run\` only" >&2; exit 2; } + [ -n "${SLURM_JOB_ID:-}" ] && [ -n "${SLURM_TMPDIR:-}" ] || { + echo "sp: --in-allocation runs every job on this machine; call it inside a SLURM" >&2 + echo " allocation (an sbatch script, or salloc), never on a login node" >&2; exit 2; } + ALLOC_CORES="$(nproc)" + ALLOC_MEM_MB="$(python - <<'PY' +import os, pathlib +limits = [] +try: + rel = next(l.split(":", 2)[2].strip() for l in open("/proc/self/cgroup") + if l.startswith("0::")) + p = pathlib.Path("/sys/fs/cgroup" + rel) + while p != p.parent: + f = p / "memory.max" + if f.exists() and f.read_text().strip() not in ("", "max"): + limits.append(int(f.read_text()) // 2**20) + p = p.parent +except (OSError, StopIteration, ValueError): + pass +if not limits and os.environ.get("SLURM_MEM_PER_NODE"): + limits.append(int(os.environ["SLURM_MEM_PER_NODE"])) +total = min(limits) if limits else 0 +print(total - max(4000, total // 20) if total else 0) +PY +)" + [ "$ALLOC_MEM_MB" -gt 0 ] || { + echo "sp: --in-allocation could not read this allocation's memory limit" >&2; exit 2; } + ALLOC_ARGS=(--executor local --cores "$ALLOC_CORES" --resources "mem_mb=$ALLOC_MEM_MB") + [ -n "$TILE_STORE_ROOT" ] || TILE_STORE_ROOT="$SLURM_TMPDIR/tile-store" + echo "sp: in allocation $SLURM_JOB_ID: ${ALLOC_ARGS[*]}" >&2 +fi + # --- the launch code snapshot ---------------------------------------------- # THE ONE HOME for this concept; everything else points here. # @@ -227,7 +300,7 @@ PY # directory (NFS) trades ngmix read latency for room; the campaign-unique store # names (LOCAL_TAG) keep concurrent campaigns apart in one root. Profiles with # no `X:/local/scratch` bind (nibi's `--bind /local`) get one appended. - python - "$SNAPSHOT/profiles/$PROFILE/config.yaml" "$(cfg tile_store_root)" <<'PY' + python - "$SNAPSHOT/profiles/$PROFILE/config.yaml" "$TILE_STORE_ROOT" <<'PY' import os, pathlib, re, sys f, root = pathlib.Path(sys.argv[1]), sys.argv[2] if root: @@ -324,7 +397,7 @@ case "$cmd" in # the compute parse's index build enforces over the WHOLE tile list — so we # record the failure and go on rather than letting `set -e` abort here. prep_rc=0 - SP_PHASE=prepare sm prepare_all_tiles "$@" || prep_rc=$? + SP_PHASE=prepare sm prepare_all_tiles ${ALLOC_ARGS[@]+"${ALLOC_ARGS[@]}"} "$@" || prep_rc=$? if [ "$prep_rc" -ne 0 ]; then echo "" >&2 echo "############################################################" >&2 @@ -338,7 +411,7 @@ case "$cmd" in fi comp_rc=0 - SP_PHASE=compute sm all "$@" || comp_rc=$? + SP_PHASE=compute sm all ${ALLOC_ARGS[@]+"${ALLOC_ARGS[@]}"} "$@" || comp_rc=$? if [ "$prep_rc" -ne 0 ] || [ "$comp_rc" -ne 0 ]; then exit "$([ "$comp_rc" -ne 0 ] && echo "$comp_rc" || echo "$prep_rc")" fi From 914510be98981b05971314162c4408e66b422e60 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 2 Oct 2026 19:42:53 -0400 Subject: [PATCH 2/3] sp --in-allocation: campaign-unique default tile store; two campaigns per allocation Data campaigns name node-local stores by bare tile id, so two campaigns over the same tiles sharing one allocation would share stores. Default to $SLURM_TMPDIR/tile-store/, and document one srun step per campaign. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Dj4HnkNqqPkVCeDeUGEGuz --- workflow/README.md | 14 ++++++++++++-- workflow/bin/sp | 15 ++++++++------- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/workflow/README.md b/workflow/README.md index 6cbad9490..1fdd9e556 100644 --- a/workflow/README.md +++ b/workflow/README.md @@ -252,8 +252,18 @@ passes `--executor local --cores N --resources mem_mb=M`; later explicit flags override them. Everything else is the profile's. `group:` labels are inactive under local execution, so `--groups` does nothing; the tile_shape members still share their node-local store because every job runs on this -node. That store defaults to `$SLURM_TMPDIR/tile-store`, which Slurm reclaims -when the job ends (an explicit `tile_store_root:` wins). +node. That store defaults to `$SLURM_TMPDIR/tile-store/`, which Slurm +reclaims when the job ends (an explicit `tile_store_root:` wins). + +Two campaigns can share one allocation, one `srun` step each. Each step's +cpuset and cgroup become that campaign's budget, and `` keeps their +stores apart: + +```bash +srun --exact -n1 -c96 --mem=380000M sp run --in-allocation -c arm_a.yaml & +srun --exact -n1 -c96 --mem=380000M sp run --in-allocation -c arm_b.yaml & +wait +``` A rule's `mem_mb` becomes a packing budget rather than a per-job cgroup: a job over its share is killed only if the allocation as a whole runs out. Size the diff --git a/workflow/bin/sp b/workflow/bin/sp index 241c920b1..8f06cf8ca 100755 --- a/workflow/bin/sp +++ b/workflow/bin/sp @@ -183,12 +183,13 @@ STATE_DIR="${SP_STATE_DIR:-${RUN_DIR}-state}" # raise the budget, and a retry asking more than the allocation has stops # the run (snakemake's "Job needs mem_mb=..."), so size the allocation for # the widest rule's third attempt. -# * The tile store defaults to $SLURM_TMPDIR/tile-store. The profile cannot -# name $SLURM_TMPDIR (its `--bind /local` comment says why), but sp runs -# inside the job here and knows the value, so it rebinds /local/scratch -# there exactly as `tile_store_root:` does. Slurm reclaims it at job end, -# and no other campaign on the node can share it. An explicit -# `tile_store_root:` still wins. +# * The tile store defaults to $SLURM_TMPDIR/tile-store/. The profile +# cannot name $SLURM_TMPDIR (its `--bind /local` comment says why), but sp +# runs inside the job here and knows the value, so it rebinds +# /local/scratch there exactly as `tile_store_root:` does. Slurm reclaims +# it at job end. keeps two campaigns over the same tiles apart when +# they share one allocation (one srun step each): data campaigns name +# their stores by bare tile id. An explicit `tile_store_root:` still wins. # Explicit --cores/--resources after it override the derived values. ALLOC_ARGS=() TILE_STORE_ROOT="" @@ -222,7 +223,7 @@ PY [ "$ALLOC_MEM_MB" -gt 0 ] || { echo "sp: --in-allocation could not read this allocation's memory limit" >&2; exit 2; } ALLOC_ARGS=(--executor local --cores "$ALLOC_CORES" --resources "mem_mb=$ALLOC_MEM_MB") - [ -n "$TILE_STORE_ROOT" ] || TILE_STORE_ROOT="$SLURM_TMPDIR/tile-store" + [ -n "$TILE_STORE_ROOT" ] || TILE_STORE_ROOT="$SLURM_TMPDIR/tile-store/$(cfg run)" echo "sp: in allocation $SLURM_JOB_ID: ${ALLOC_ARGS[*]}" >&2 fi From 5ee28fdd299e5e607ed166724dfeba3c4c4fec36 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Mon, 5 Oct 2026 09:10:19 -0400 Subject: [PATCH 3/3] sp run --in-allocation: resubmit as-is after a timeout A job cut short by its time limit leaves two things a plain resubmission trips over. Its node-local tile store is gone, so a tile whose tile_vignets.json survived without a final_cat stops at tile.smk's "vignette store is missing" guard; and a killed snakemake leaves its lock. sp now deletes those manifests before the run, and records the job that holds the state dir so that a lock whose job has left the queue is released. Any other lock stops the run before anything is swept. With -n both only report. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01D8wdTpGpBtyK6uHwU7qWQW --- workflow/README.md | 14 ++++++---- workflow/bin/sp | 64 +++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 70 insertions(+), 8 deletions(-) diff --git a/workflow/README.md b/workflow/README.md index 1fdd9e556..70c688db5 100644 --- a/workflow/README.md +++ b/workflow/README.md @@ -271,11 +271,15 @@ allocation so the widest rule's last retry fits (a retry asking more than the allocation has stops the run), at no more than 4 GB per core so memory does not raise the billing. Keep it to one node: the tile store is node-local, so a multi-node allocation would split a tile's vignets from its ngmix chunks. -For the same reason a time limit that ends the job mid-tile leaves that tile's -`tile_vignets.json` without its store, and a resume trips tile.smk's -"node-local vignette store is missing" guard. Delete the manifest of each tile -that has no `final_cat` yet before resuming, as the guard's message says, and -give the time limit room. + +A job cut short by its time limit is resumed by submitting the same script +again. The store died with the job, so before it runs `sp` deletes +`tile_vignets.json` for each listed tile that has no `final_cat` yet (tile.smk's +"node-local vignette store is missing" guard would stop its chunks otherwise); +tile_vignets rebuilds the store and the chunks rerun after it. `sp` also +records the job holding the campaign's state dir, and releases that job's lock +once the job has left the queue; any other lock stops the run. With `-n` both +only report. ## Layout diff --git a/workflow/bin/sp b/workflow/bin/sp index 8f06cf8ca..afc277a4e 100755 --- a/workflow/bin/sp +++ b/workflow/bin/sp @@ -187,9 +187,20 @@ STATE_DIR="${SP_STATE_DIR:-${RUN_DIR}-state}" # cannot name $SLURM_TMPDIR (its `--bind /local` comment says why), but sp # runs inside the job here and knows the value, so it rebinds # /local/scratch there exactly as `tile_store_root:` does. Slurm reclaims -# it at job end. keeps two campaigns over the same tiles apart when -# they share one allocation (one srun step each): data campaigns name -# their stores by bare tile id. An explicit `tile_store_root:` still wins. +# it at job end. gives each campaign sharing one allocation (one +# srun step each) its own directory. An explicit `tile_store_root:` still +# wins. +# * A resume finds the previous job's store gone. A tile whose +# tile_vignets.json survived it but whose final_cat did not would trip +# tile.smk's "node-local vignette store is missing" guard, so that +# manifest is deleted first: tile_vignets reruns and rebuilds the store, +# and the tile's chunks rerun after it. Tiles with a final_cat are never +# touched. +# * A lock left by a killed in-allocation run is released. sp records the +# job that holds the state dir; when that job is no longer in the queue, +# its lock is stale. Any other lock (a run sp did not start in an +# allocation, or one still running) stops the run instead. +# With -n both only report. # Explicit --cores/--resources after it override the derived values. ALLOC_ARGS=() TILE_STORE_ROOT="" @@ -227,6 +238,52 @@ PY echo "sp: in allocation $SLURM_JOB_ID: ${ALLOC_ARGS[*]}" >&2 fi +# The resume housekeeping (a stale lock, then stale vignette manifests; the +# lock goes first so a campaign still running is never swept). `sp run` calls +# it before the snapshot, with its own arguments, so -n only reports. +alloc_resume() { + local a alloc_dry=0 products_dir n_swept=0 tile m holder holder_file + local holder_state="" holder_rc=1 + for a in "$@"; do + case "$a" in -n|--dry-run|--dryrun|-[a-zA-Z]*n*) alloc_dry=1 ;; esac + done + holder_file="$STATE_DIR/in-allocation.job" + if [ -n "$(ls -A "$STATE_DIR/.snakemake/locks" 2>/dev/null)" ]; then + holder="$(cat "$holder_file" 2>/dev/null || true)" + if [ -n "$holder" ]; then + holder_state="$(squeue -h -j "$holder" -o %T 2>&1)" && holder_rc=0 || holder_rc=$? + fi + # Ended: gone from the controller, or listed in a final state. + if [ -n "$holder" ] && [ "$holder" != "$SLURM_JOB_ID" ] && { + [[ "$holder_state" == *"Invalid job id"* ]] || { [ "$holder_rc" = 0 ] && + [[ "$holder_state" =~ ^(|COMPLETED|FAILED|CANCELLED|TIMEOUT|OUT_OF_MEMORY|NODE_FAIL|PREEMPTED|BOOT_FAIL|DEADLINE)$ ]]; }; }; then + if [ "$alloc_dry" = 1 ]; then + echo "sp: stale lock of job $holder in $STATE_DIR (dry run, kept)" >&2 + else + echo "sp: releasing the stale lock of job $holder in $STATE_DIR" >&2 + sm --unlock >/dev/null + fi + else + echo "sp: $STATE_DIR is locked${holder:+ (job $holder: ${holder_state:-running})}." >&2 + echo " If nothing runs on this campaign, \`sp -c ... --unlock\` and resubmit." >&2 + exit 2 + fi + fi + + products_dir="$(cfg outputs.products_dir)" + while read -r tile; do + [ -n "$tile" ] || continue + m="$RUN_DIR/tiles/${tile:0:2}/$tile/manifests/tile_vignets.json" + [ -e "$m" ] || continue + [ -e "$products_dir/tiles/${tile:0:2}/$tile/final_cat-$tile.hdf5" ] && continue + if [ "$alloc_dry" = 1 ]; then echo "sp: stale (dry run, kept): $m" >&2 + else rm -f "$m"; echo "sp: deleted stale $m" >&2; fi + n_swept=$((n_swept + 1)) + done < <(sed 's/[[:space:]]//g' "$(cfg tile_list)") + echo "sp: $n_swept stale tile_vignets.json (no final_cat, store gone with its job)" >&2 + [ "$alloc_dry" = 1 ] || printf '%s\n' "$SLURM_JOB_ID" > "$holder_file" +} + # --- the launch code snapshot ---------------------------------------------- # THE ONE HOME for this concept; everything else points here. # @@ -391,6 +448,7 @@ cmd="${1:-}" case "$cmd" in run) shift + [ "$IN_ALLOCATION" = 0 ] || alloc_resume "$@" snapshot_code # PREPARE failing is NOT fatal to the run: keep-going means a failed tile # poisons only its own cone, and the tiles that lost their exposure list are