Snakemake orchestration for the real-data pipeline

.gitignore

@@ -139,3 +139,4 @@ code
 # .felt here is a machine-local symlink into it. Never track it in this repo.
 /.felt/
 /.felt
+.snakemake/

profiles/nibi/config.yaml

@@ -0,0 +1,103 @@
+# Snakemake profile for the Nibi cluster (Digital Research Alliance).
+#
+# ONE-ALLOCATION mode: sbatch one node up front, then run snakemake with the
+# LOCAL scheduler inside it (no per-job sbatch; the local scheduler bin-packs the
+# DAG into the allocation against `cores` + `resources: mem_mb`). This is the
+# proven nibi path (job_p3_batch1) and the single honest first-pass mode;
+# multi-node scaling = multiple tile-disjoint invocations, deferred.
+#
+# Launch via `workflow/bin/sp` (loads apptainer/1.4.5, uses the /project venv).
+# software-deployment-method wraps every job's shell in `apptainer exec` — the
+# user never types apptainer; the container is set in the Snakefile (local .sif).
+
+# Local scheduler (no `executor:` -> built-in local). Size to the allocation.
+cores: 192                       # node cores available to the DAG
+local-cores: 4                   # cap localrule/aggregation concurrency
+resources:
+  mem_mb: 720000                 # node RAM cap the local scheduler bin-packs under
+
+software-deployment-method: [apptainer]
+# Explicit container environment (finding 2/3): the apptainer SDM otherwise drops
+# the proven-recipe env and hardcodes --home <workdir>, hiding ~/.ssl/cadcproxy.pem.
+#  --cleanenv          strict host-env isolation (APPTAINERENV_*/SINGULARITYENV_* survive)
+#  OMP_NUM_THREADS=1   caps OpenBLAS fork-explosion (verified: pool 32->1)
+#  MALLOC_ARENA_MAX=2  bounded allocator (im_sims/nibi lesson)
+#  --home /home/cdaley wins over the SDM's --home <cwd>; restores cadcproxy.pem for vos/vcp
+#  PYTHONPATH          SETTLED CALL 3 — P0 pins THIS branch's src/ (identical to
+#                      develop@97e16d50: orchestration commits never touch src/).
+#                      NOT shapepipe-prod (drifted to a PR branch mid-run, and the
+#                      live p3-batch1 job reads it) and not the sif default
+#                      (frozen pre-#843). Production later rebuilds the sif at the
+#                      validated commit and DROPS this --env line.
+apptainer-args: "--cleanenv --env OMP_NUM_THREADS=1 --env MALLOC_ARENA_MAX=2 --env PYTHONPATH=/project/def-mjhudson/cdaley/shapepipe-snakemake/src --home /home/cdaley --bind /project --bind /scratch"
+
+latency-wait: 60                 # NFS: wait for outputs to appear after a job
+keep-going: true                 # a failed job poisons only its cone; siblings run on
+rerun-incomplete: true           # re-do jobs left incomplete by an unclean death
+show-failed-logs: true
+printshellcmds: true
+
+# rerun-triggers left at the full Snakemake v9 default
+# {mtime,params,input,code,software-env}: params/code fixes SHOULD propagate.
+# No override — the wrapper writes outputs only on change so mtimes move only
+# when reality moves; protected() guards finished science products.
+
+default-resources:
+  mem_mb: 4000
+  runtime: 120                   # minutes
+  cpus_per_task: 1
+
+# Per-stage resources — measured footprints from the nibi P3 runs (#808).
+# threads == cpus_per_task == ShapePipe SMP fork width (sp_rule sets
+# SMP_BATCH_SIZE={threads}); mem_mb == threads x measured per-worker footprint.
+#
+# set-threads is REQUIRED for {threads} to resolve above 1: the local scheduler
+# derives cpus_per_task from threads, NOT the reverse (verified snakemake 9.23 —
+# set-resources cpus_per_task alone leaves {threads}==1, i.e. SMP_BATCH_SIZE=1,
+# serial). Keep the two blocks in lockstep.
+set-threads:
+  exp_split: 4
+  exp_mask: 4
+  exp_psf: 8
+  tile_mask: 16
+  tile_detect: 16
+  tile_merge_headers: 8
+  tile_vignets: 8
+  tile_ngmix: 8
+  tile_merge_cats: 4
+  tile_make_cat: 4
+
+set-resources:
+  exp_split:
+    cpus_per_task: 4
+    mem_mb: 8000
+  exp_mask:
+    cpus_per_task: 4
+    mem_mb: 8000
+  exp_psf:
+    cpus_per_task: 8
+    mem_mb: 16000
+    runtime: 240
+  tile_mask:
+    cpus_per_task: 16
+    mem_mb: 16000
+  tile_detect:
+    cpus_per_task: 16
+    mem_mb: 16000
+  tile_merge_headers:
+    cpus_per_task: 8
+    mem_mb: 8000
+  tile_vignets:
+    cpus_per_task: 8
+    mem_mb: 24000
+    runtime: 240
+  tile_ngmix:
+    cpus_per_task: 8
+    mem_mb: 24000
+    runtime: 360
+  tile_merge_cats:
+    cpus_per_task: 4
+    mem_mb: 8000
+  tile_make_cat:
+    cpus_per_task: 4
+    mem_mb: 8000

workflow/README.md

@@ -0,0 +1,105 @@
+# ShapePipe Snakemake orchestration
+
+Snakemake workflow that orchestrates real-data ShapePipe runs, replacing the
+`curl_canfar_local.sh → run_job_sp_canfar_v2.0.bash → job_sp_canfar_v2.0.bash`
+bash layers (and per-site sbatch reimplementations). **Module code is untouched**:
+rules call `shapepipe_run -c <config>` on the existing config chains. Design and
+rationale: [CosmoStat/shapepipe#848](https://github.com/CosmoStat/shapepipe/issues/848)
+(the living PRD).
+
+## Quick start (nibi)
+
+**One-allocation mode.** `sbatch` one node, then run Snakemake with the *local*
+scheduler inside it: it bin-packs the DAG into the allocation (no per-job
+sbatch). Each job's shell runs `shapepipe_run` **inside the container** (the
+profile's apptainer software-deployment — you never type `apptainer`).
+
+```bash
+# One-time: a snakemake env on a SHARED filesystem (/project — NOT /tmp, which is
+# node-local; the executor re-invokes this python inside jobs).
+uv venv /project/def-mjhudson/cdaley/snakemake-env --python 3.12
+source /project/def-mjhudson/cdaley/snakemake-env/bin/activate
+uv pip install 'snakemake>=9,<10'
+
+# Edit workflow/config.yaml: tile_list, run_dir, config_src, container, star_cats.
+
+# The committed launcher loads apptainer/1.4.5 + the /project venv, so a fresh
+# shell always has the right state. Inside your allocation:
+workflow/bin/sp prepare   # prepare_tiles -> build_index.py -> prepare_exposures
+workflow/bin/sp run       # the main compute DAG
+```
+
+`sp` subcommands: `prepare` (chains all three prepare steps), `prepare_tiles`,
+`index`, `prepare_exposures`, `run`, `rerun <target>` (forced recompute of
+protected products), `cancel <run>` (scancel sweep before `--unlock`),
+`clean-exposures [--apply]`, `report`.
+
+## Execution is three static invocations
+
+The exposure job set is data-derived from the tiles' `find_exposures` output, so
+it cannot be scheduled in the same static DAG that produces it. Hence:
+
+1. `snakemake prepare_tiles` — per-tile static DAG (Git_vos → Uz → Fe),
+   `keep-going` so tile failures are independent.
+2. `build_index.py` — a **plain script**, not a DAG node. It iterates the
+   declared tile list, checks each tile's Fe output at its deterministic path (no
+   glob), writes `run_index.sqlite` + `missing.json`, and exits nonzero only if
+   the missing *fraction* exceeds a threshold. (There is no `--allow-missing`
+   Snakemake flag; that mechanism does not exist.)
+3. `snakemake prepare_exposures` — per-exposure static DAG (Gie_vos + per-unit
+   star cats), parseable now the index exists.
+4. `snakemake` (main DAG) — the tile/exposure compute chain.
+
+`sp prepare` chains 1–3 so the UX stays one command.
+
+## Layout
+
+```
+workflow/
+  Snakefile              parse-time index load; global container:; onsuccess/onerror report hooks
+  config.yaml            the run: tile list, paths, container, chunk count
+  bin/sp                 committed launcher (module load + /project venv + subcommands)
+  rules/
+    prepare.smk          tile get_images/uncompress/find_exposures + per-unit star cats
+    exposure.smk         per-exposure: get_images, split, mask, psf (no temp())
+    tile.smk             per-tile: exp forest, merge_headers, mask, detect, vignets, ngmix, merge, make_cat
+  scripts/
+    sp_rule.py           the thin per-unit wrapper (isolation furniture, config copy, log-sync, count floor)
+    build_index.py       prepare-phase run_index.sqlite builder (plain script)
+    build_forest.py      per-tile exposure symlink forest (group-compatible shell)
+    completeness.py      the ported count-floor table (shared by sp_rule + run_report)
+    run_report.py        standalone report (NOT a DAG node; run_report hooks call it)
+    clean_exposures.py   index-driven exposure-store reclaim (replaces temp())
+profiles/nibi/config.yaml  local scheduler; apptainer SDM; resources; keep-going
+```
+
+## How it works
+
+- **Completeness = the count floor.** Each stage declares its deterministic
+  `output/run_sp_<prefix>/` directory as `directory()` output (the per-unit
+  config copy forces `RUN_DATETIME = False`, so the path is known at DAG time).
+  After the run, `sp_rule.py` counts products per mandatory runner against
+  `completeness.py`'s floor and exits nonzero below it; per-CCD attrition between
+  floor and `expect` is tolerated (the directory simply lacks that product).
+  This *is* the failure policy — no 3-class taxonomy, no error-signature
+  whitelist. `--keep-going` isolates a failure to its own cone.
+- **Per-unit isolation is by work-dir content, not `-e/--exclusive`.** Every
+  tile/exposure runs its own `shapepipe_run` in `tiles/<ID>/` or `exp/<base>/`,
+  with a `cfis` config symlink, `star_cat_{exp,tiles}` symlinks, a
+  `tile_numbers.txt` (tiles) or fabricated `exp_numbers-000-000.txt` pseudo-Fe
+  (exposures), and a config copy setting `RUN_DATETIME=False` and
+  `SMP_BATCH_SIZE={threads}`. `NUMBER_LIST=-<ID>` is injected **only** for
+  stages whose numbering scheme is the unit ID (tile-scheme stages + exp split),
+  never for get_images / exp_mask / exp_psf.
+- **No `temp()` on shared exposure dirs.** Exposures overlap tiles, so temp()
+  would cascade destructive reruns when a tile is appended. The store is
+  persistent; `sp clean-exposures` reclaims space only for exposures all of
+  whose consuming tiles have a final_cat.
+- **The index is parse-time data, never a rule input.** Appending tiles changes
+  which jobs exist without invalidating completed work. Star cats are re-keyed
+  per unit for the same reason. `final_cat` is `protected()`.
+- **Failure is a report, not a gate.** `run_report.py` disk-scans the trees
+  against the count table and enumerates shortfalls (whole-unit absence vs
+  per-CCD attrition). It is a standalone script — a DAG report node would be
+  poisoned by the very failures it must enumerate — emitted automatically by the
+  `onsuccess`/`onerror` hooks and runnable any time via `sp report`.

workflow/Snakefile

@@ -0,0 +1,141 @@
+"""ShapePipe real-data orchestration — Snakemake workflow.
+
+Design / rationale: CosmoStat/shapepipe#848 (the living PRD).
+
+Execution is honestly THREE static invocations, chained by ``workflow/bin/sp``
+so the UX stays one command (``sp prepare`` / ``sp run``):
+
+    snakemake --profile profiles/nibi prepare_tiles      # Git_vos + Uz + Fe (keep-going)
+    python workflow/scripts/build_index.py ...            # PLAIN SCRIPT, not a DAG node
+    snakemake --profile profiles/nibi prepare_exposures  # Gie_vos + star cats
+    snakemake --profile profiles/nibi                     # the main compute DAG
+
+The run index is loaded ONCE here at parse time into plain dicts; it is never a
+rule input, so appending tiles and rebuilding it changes which jobs exist without
+invalidating completed work. build_index.py is a plain script *between*
+prepare_tiles and prepare_exposures (the exposure job set is data-derived from
+the tiles' find_exposures output, so it cannot be scheduled in the same static
+DAG that produces it).
+"""
+
+import sqlite3
+from pathlib import Path
+
+# Resolved relative to THIS file, not the working directory: snakemake runs with
+# --directory on /scratch (bin/sp) so .snakemake/ state never lands on /project
+# (group quota is a hard 27/27 TiB — a metadata write mid-run died on it live).
+configfile: str(Path(workflow.snakefile).parent / "config.yaml")
+
+# Every job's shell runs inside this container (apptainer software-deployment in
+# the profile); the user never types apptainer.
+container: config["container"]
+
+# --- paths -----------------------------------------------------------------
+RUN_DIR    = Path(config["run_dir"])
+CONFIG_SRC = Path(config["config_src"])
+STAR_CATS  = Path(config["star_cats"])
+INDEX_DB   = Path(config["index_db"])
+SCRIPTS    = Path(workflow.basedir) / "scripts"
+NGMIX_CHUNKS = config.get("ngmix_chunks", 8)
+
+with open(config["tile_list"]) as f:
+    TILES = [ln.strip() for ln in f if ln.strip()]
+
+# --- parse-time index load (compute phase only) ----------------------------
+# EXP: exposure base-id -> original name (2605805 -> 2605805p; the name goes
+# verbatim into the fabricated per-unit exp_numbers list so get_images matches
+# <name>.fits.fz in the store).  TILE_EXP: tile -> [exp_ids].
+EXP, TILE_EXP = {}, {}
+if INDEX_DB.exists():
+    _con = sqlite3.connect(INDEX_DB)
+    EXP = dict(_con.execute("SELECT exp_id, name FROM exposures"))
+    for tile, exp in _con.execute("SELECT tile_id, exp_id FROM tile_exposures"):
+        TILE_EXP.setdefault(tile, []).append(exp)
+    _con.close()
+
+wildcard_constraints:
+    tile = r"\d{3}\.\d{3}",
+    exp  = r"\d{6,7}",
+    chunk = r"\d+",
+
+# --- shared wrapper invocation ---------------------------------------------
+def sp_rule(stage, config_name, level, unit, *, isolate=True, exp_forest=None,
+            extra=""):
+    """Build the shell command that runs one unit through sp_rule.py.
+
+    NUMBER_LIST is injected by default (tile-scheme stages + exp split); pass
+    ``isolate=False`` for get_images / exp_mask / exp_psf, whose per-CCD/download
+    numbering must not be constrained. ``--threads {threads}`` is appended in
+    each rule's ``shell:`` string (NOT here: Snakemake formats the shell string
+    once, so a ``{threads}`` inside params.cmd would survive literally) and
+    becomes SMP_BATCH_SIZE (fork width == cpus_per_task).
+    """
+    cmd = (
+        f"python {SCRIPTS}/sp_rule.py --stage {stage} --config {config_name} "
+        f"--unit {unit} --level {level} --run-dir {RUN_DIR} "
+        f"--config-src {CONFIG_SRC} --star-cats {STAR_CATS} "
+        f"--script-hash {SCRIPT_HASH}"
+    )
+    if not isolate:
+        cmd += " --no-isolate"
+    if exp_forest:
+        cmd += f" --exp-forest {exp_forest}"
+    if extra:
+        cmd += " " + extra
+    return cmd
+
+def tile_dir(tile):
+    return RUN_DIR / "tiles" / tile
+
+def exp_dir(exp):
+    return RUN_DIR / "exp" / exp
+
+# Content hash of the wrapper scripts, computed once at parse time. Exposed as a
+# param on every rule (via sp_rule() and tile_exp_forest) so the default
+# rerun-triggers' *params* trigger covers script edits: the `code` trigger only
+# hashes the rule's own shell string, NOT external scripts it calls — without
+# this, a wrapper/forest-script fix silently leaves stale outputs in place
+# (bitten live: the forest-shard fix rebuilt nothing on rerun).
+import hashlib as _hashlib
+SCRIPT_HASH = _hashlib.md5(b"".join(
+    (SCRIPTS / n).read_bytes()
+    for n in ("sp_rule.py", "completeness.py", "build_forest.py"))).hexdigest()[:12]
+
+include: "rules/prepare.smk"
+include: "rules/exposure.smk"
+include: "rules/tile.smk"
+
+# --- top-level targets ------------------------------------------------------
+# Only aggregation targets are localrules (no-ops under the local executor, but
+# they keep the compute chain group-compatible: no mid-chain localrules).
+localrules: all, prepare_tiles, prepare_exposures
+
+rule all:
+    input:
+        expand(str(RUN_DIR / "tiles/{tile}/final_cat-{tile}.fits"), tile=TILES)
+
+# Phase A.1 — per-tile static DAG (download + find_exposures). keep-going so
+# tile failures are independent; build_index.py runs next as a plain script.
+rule prepare_tiles:
+    input:
+        expand(str(RUN_DIR / "tiles/{tile}/output/run_sp_tile_Fe"), tile=TILES)
+
+# Phase A.2 — per-exposure static DAG (download + star cats), parseable now the
+# index exists. Star cats are re-keyed per unit (see prepare.smk).
+rule prepare_exposures:
+    input:
+        expand(str(RUN_DIR / "exp/{exp}/output/run_sp_exp_Gie"), exp=sorted(EXP))
+
+# --- report hooks -----------------------------------------------------------
+# run_report is NOT a DAG node (a descendant of every job would be poisoned by
+# any hard failure — the exact case it exists for). It is a standalone script,
+# emitted automatically at the end of every invocation, runnable any time.
+def _report(status):
+    shell(f"python {SCRIPTS}/run_report.py --run-dir {RUN_DIR} "
+          f"--index {INDEX_DB} --status {status} || true")
+
+onsuccess:
+    _report("success")
+
+onerror:
+    _report("error")