diff --git a/workflow/Snakefile b/workflow/Snakefile index ac9d0760c..0143ce5c5 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -1177,8 +1177,21 @@ def unit_pre(stage, unit, *, exp_name=None, forest=None, env=None, # Also set via apptainer-args in the profile; kept here so a hand-run of # this same line outside snakemake behaves identically. f"export {_THREAD_CAPS}", - 'mkdir -p "$SP_RUN/output" "$SP_RUN/manifests" "$SP_RUN/logs"', ] + if level == "exp": + # Under `sp run --in-allocation`, SP_EXP_STORE names a node-local root + # and the exposure's output/ is a link into it (sp's header says why). + # The link is made once, by whichever stage runs first; a dangling one + # is swept by sp before the run, with the manifests that attested it. + # An output/ already on /scratch (a store begun outside an allocation) + # stays where it is. + lines.append( + 'if [ -n "${SP_EXP_STORE:-}" ] && [ ! -e "$SP_RUN/output" ] ' + '&& [ ! -L "$SP_RUN/output" ]; then ' + f'mkdir -p "$SP_EXP_STORE/{unit}" && ' + f'ln -s "$SP_EXP_STORE/{unit}" "$SP_RUN/output"; fi') + # Unchanged text for the tile side: unit_pre is a params rerun trigger. + lines.append('mkdir -p "$SP_RUN/output" "$SP_RUN/manifests" "$SP_RUN/logs"') if PSF_DICT: lines.append(f"export PSF_DICT='{PSF_DICT}'") if forest: diff --git a/workflow/bin/sp b/workflow/bin/sp index afc277a4e..ad8c2f905 100755 --- a/workflow/bin/sp +++ b/workflow/bin/sp @@ -190,6 +190,15 @@ STATE_DIR="${SP_STATE_DIR:-${RUN_DIR}-state}" # it at job end. gives each campaign sharing one allocation (one # srun step each) its own directory. An explicit `tile_store_root:` still # wins. +# * The exposure stores' bulk (each exposure's output/, ~7 GB and ~1,000 +# files) goes node-local too, under $SLURM_TMPDIR/exp-store/: sp +# exports SP_EXP_STORE, and the exposure jobs link /output to +# $SP_EXP_STORE/ (sp_shell, Snakefile). The link names the real +# path, which /local's bind makes the same on the host and in the +# container. Manifests, logs and the tombstone stay on /scratch beside +# the link, so the DAG reads them as before; clean_exposure removes the +# node-local tree with the link. What downstream reads after the store +# goes is on the products root (exp_persist's tar, exp_maps' fragment). # * 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 @@ -235,6 +244,7 @@ PY 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)" + export SP_EXP_STORE="$(realpath -m "$SLURM_TMPDIR")/exp-store/$(cfg run)" echo "sp: in allocation $SLURM_JOB_ID: ${ALLOC_ARGS[*]}" >&2 fi @@ -281,6 +291,20 @@ alloc_resume() { 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 + + # An exposure whose node-local output went with an earlier job: its link + # dangles and its scratch manifests attest work that is gone. Dropping the + # link and the manifests reruns its chain; a cleaned exposure (tombstone) + # is reclaimed, not stale, and is never touched. + n_swept=0 + for m in "$RUN_DIR"/exp/*/*/output; do + [ -L "$m" ] && [ ! -e "$m" ] || continue + [ -e "$(dirname "$m")/cleaned.json" ] && continue + if [ "$alloc_dry" = 1 ]; then echo "sp: stale (dry run, kept): $m" >&2 + else rm -f "$m"; rm -rf "$(dirname "$m")/manifests"; echo "sp: deleted stale $m and its manifests" >&2; fi + n_swept=$((n_swept + 1)) + done + echo "sp: $n_swept stale exposure store(s) (no tombstone, store gone with its job)" >&2 [ "$alloc_dry" = 1 ] || printf '%s\n' "$SLURM_JOB_ID" > "$holder_file" } @@ -379,6 +403,24 @@ if root: print(f"sp: tile store root {root} (bound to /local/scratch)", file=sys.stderr) PY + # The exposure jobs read SP_EXP_STORE (sp_shell); --cleanenv drops it unless + # the profile passes it. Like the bind above, an apptainer arg is no rerun + # trigger, so a new job's path never reschedules finished exposures. + if [ -n "${SP_EXP_STORE:-}" ]; then + mkdir -p "$SP_EXP_STORE" + python - "$SNAPSHOT/profiles/$PROFILE/config.yaml" "$SP_EXP_STORE" <<'PY' +import pathlib, re, sys +f, store = pathlib.Path(sys.argv[1]), sys.argv[2] +text, n = re.subn(r'(apptainer-args:\s*"[^"]*)"', + lambda m: m.group(1) + f" --env SP_EXP_STORE={store}" + '"', + f.read_text(), count=1) +if not n: + sys.exit("sp: SP_EXP_STORE is set but the profile has no quoted apptainer-args") +f.write_text(text) +print(f"sp: exposure store {store} (node-local)", file=sys.stderr) +PY + fi + python - "$SNAPSHOT/snapshot.json" "$REPO" <<'PY' import json, pathlib, subprocess, sys, time out, repo = pathlib.Path(sys.argv[1]), sys.argv[2] diff --git a/workflow/scripts/clean_exposure.py b/workflow/scripts/clean_exposure.py index edff0f06f..0fd2ecb60 100644 --- a/workflow/scripts/clean_exposure.py +++ b/workflow/scripts/clean_exposure.py @@ -66,6 +66,7 @@ import argparse import json +import os import shutil import time from pathlib import Path @@ -114,6 +115,12 @@ def main() -> None: for target in targets: # NEVER rmtree a symlink (see the module docstring). if target.is_symlink(): + # The one link we own: output/ into the node-local exposure store + # (SP_EXP_STORE, sp --in-allocation). Its tree goes with it. + store = os.environ.get("SP_EXP_STORE") + real = target.resolve() + if store and real.is_dir() and real.parent == Path(store).resolve(): + shutil.rmtree(real) target.unlink() else: shutil.rmtree(target)