diff --git a/tests/workflow/test_resume.py b/tests/workflow/test_resume.py new file mode 100644 index 000000000..ef36cc5d5 --- /dev/null +++ b/tests/workflow/test_resume.py @@ -0,0 +1,83 @@ +"""Every resumed reader must wait for the rebuilt node-local vignette store.""" + +import os +import shutil +from collections import Counter +from pathlib import Path + +import pytest +from snakemake.io import is_flagged + + +def seed_outputs(campaign, resolve_dag): + """Materialise satisfied outputs without executing the science pipeline.""" + with resolve_dag(campaign) as dag: + for job in dag.jobs: + for output in job.output: + path = Path(output) + path.parent.mkdir(parents=True, exist_ok=True) + if is_flagged(output, "directory"): + path.mkdir(exist_ok=True) + stamp = path / ".snakemake_timestamp" + stamp.touch() + os.utime(stamp, (1, 1)) + else: + path.touch() + os.utime(path, (1, 1)) + + +@pytest.mark.parametrize("state", [ + "merged", "merged_reclaimed", "chunks_done", "partial", "missing_chunkdir", +]) +@pytest.mark.parametrize("invalidate_detect", [False, True], ids=[ + "swept_only", "swept_and_upstream_invalidated", +]) +def test_resume_after_vignets_swept( + campaign, resolve_dag, state, invalidate_detect, +): + """Sweep alone and sweep plus upstream invalidation both preserve ordering. + + The newer tile_detect input reproduces the review's counterexample: + completed chunks can still be scheduled and must retain their writer edge. + """ + seed_outputs(campaign, resolve_dag) + tile = next(iter(campaign.ready)) + root = campaign.tile_manifest(tile, "tile_vignets").parents[1] + campaign.tile_manifest(tile, "tile_vignets").unlink() + campaign.tile_manifest(tile, "tile_make_cat").unlink() + campaign.final_cat(tile).unlink() + if state == "merged_reclaimed": + # merge_cats reclaims its temp inputs in production. + for path in (root / "output").glob("run_sp_tile_ngmix_Ng*u"): + shutil.rmtree(path) + elif state != "merged": + campaign.tile_manifest(tile, "tile_merge_cats").unlink() + if state == "partial": + campaign.tile_manifest(tile, "tile_ngmix_8").unlink() + if state == "missing_chunkdir": + shutil.rmtree(root / "output/run_sp_tile_ngmix_Ng8u") + if invalidate_detect: + campaign.tile_manifest(tile, "tile_detect").touch() + with resolve_dag(campaign) as dag: + pending = list(dag.workflow.dag.needrun_jobs()) + counts = Counter(job.rule.name for job in pending) + print( + f"Resume dry-run jobs ({state}, invalidate_detect={invalidate_detect}):", + dict(sorted(counts.items())), + ) + assert counts["tile_vignets"] == 1 + assert counts["tile_ngmix"] == dag.namespace["NGMIX_CHUNKS"] + assert counts["tile_merge_cats"] == 1 + assert counts["tile_make_cat"] == 1 + writer = next(j for j in pending if j.rule.name == "tile_vignets") + readers = [ + job for job in pending + if job.rule.name in {"tile_make_cat", "tile_ngmix"} + ] + assert len(readers) == dag.namespace["NGMIX_CHUNKS"] + 1 + for job in readers: + assert job.wildcards.tile == tile + assert writer in dag.workflow.dag.dependencies[job], ( + f"pending reader {job.rule.name} {job.wildcards} has no writer " + f"dependency; inputs={list(job.input)}" + ) diff --git a/workflow/README.md b/workflow/README.md index 4f0b135c3..571728bc4 100644 --- a/workflow/README.md +++ b/workflow/README.md @@ -230,6 +230,57 @@ 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). + +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 +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. + +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_make_cat` declares that manifest as an input, so `tile_vignets` rebuilds the store even if all chunks and their merge are complete. +Every ngmix chunk also declares the vignets manifest unconditionally, so all chunks rerun after the store is rebuilt, followed by their merge and `make_cat`. +This spends ngmix time on swept tiles, but guarantees that every scheduled reader waits for the writer, even if another input or rerun trigger invalidates a previously completed chunk. +A completion manifest alone cannot prove that a chunk will not run in this invocation. +`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 cb6754108..e314d33ef 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,138 @@ 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. 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_make_cat declares it as an input, so +# tile_vignets rebuilds the store even if merge_cats is already complete. +# Every ngmix chunk retains its dependency on the writer and reruns after +# it, even if its previous output survived. This costs ngmix time on resume +# but ensures no reader can race the rebuild. Tiles with a final_cat are +# never swept. +# * 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="" +[ "$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/$(cfg run)" + 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. +# This repairs ephemeral stores, not persistent products. In particular, do not +# add an exp_maps sweep here: the map rule must declare its fragment as an +# output while the store is live, and reclamation must require it first. +# After reclamation, existing fragments must be producer-free immutable leaves; +# missing ones must fail planning with an explicit out-of-DAG backfill command. +# This prevents code/params edits from deleting products whose sources are gone. +# The same lifecycle boundary also protects non-allocation runs. +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. # @@ -227,7 +368,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: @@ -304,19 +445,24 @@ sm() { local root; root="$(code_root)" # The run config the Snakefile reads: the snapshot's copy when this run has # one, so every parse of the campaign -- head and jobs -- sees the same file. - if [ -n "${SP_RUN_CONFIG:-}" ] && [ -f "$root/run_config.yaml" ]; then - export SP_RUN_CONFIG="$root/run_config.yaml" - fi - SP_PHASE="${SP_PHASE:-passthrough}" \ - snakemake --snakefile "$root/workflow/Snakefile" \ - --profile "$root/profiles/$PROFILE" \ - --directory "$STATE_DIR" "$@" + # Scoped to this call (subshell): SP_RUN_CONFIG in the caller stays the + # external config, so a later snapshot refresh never copies the file onto itself. + ( + if [ -n "${SP_RUN_CONFIG:-}" ] && [ -f "$root/run_config.yaml" ]; then + export SP_RUN_CONFIG="$root/run_config.yaml" + fi + SP_PHASE="${SP_PHASE:-passthrough}" \ + snakemake --snakefile "$root/workflow/Snakefile" \ + --profile "$root/profiles/$PROFILE" \ + --directory "$STATE_DIR" "$@" + ) } 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 @@ -324,7 +470,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 +484,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 diff --git a/workflow/rules/tile.smk b/workflow/rules/tile.smk index 107bce4f9..d2e58bd35 100644 --- a/workflow/rules/tile.smk +++ b/workflow/rules/tile.smk @@ -626,10 +626,8 @@ rule tile_vignets: rule tile_ngmix: group: TILE_GROUP input: - # The manifest, and only the manifest. The vignette store is no longer a - # declared input because it is no longer a declared output: it is - # node-local, produced by tile_vignets earlier in THIS SAME group job - # (see TILE_LOCAL). The manifest was always the real edge. + # Unconditional: even a completed chunk can be invalidated by another + # input or rerun trigger. Every scheduled reader must wait for the store. vignets = rules.tile_vignets.output.manifest, sx = rules.tile_detect.output.manifest, output: @@ -871,9 +869,9 @@ rule tile_merge_cats: rule tile_make_cat: group: TILE_GROUP input: - # No store input: it is node-local, written by tile_vignets in this same - # group job. make_cat reads its configured PSF-interpolation output - # through $NGMIX_VIGNET_DIR (config_tile_Mc.ini). + # The store is node-local, but its manifest is the ordering edge: + # after a resume sweep, rebuild it even when merge_cats is complete. + vignets = rules.tile_vignets.output.manifest, ms = rules.tile_merge_cats.output.manifest, output: manifest = f"{TILE_DIR}/manifests/tile_make_cat.json",