Skip to content

Snakemake orchestration for the real-data run (design proposal) #849

Description

@cailmdaley

Snakemake orchestration for the real-data ShapePipe run

Status: design proposal — for discussion before any implementation.

This issue proposes replacing ShapePipe's shell-script orchestration for real-data (UNIONS/CFIS) runs with a Snakemake workflow, leaving the ShapePipe module code untouched. It is the real-data sibling of the image-sims Snakemake work (#766, whose orchestration later moved to sp_validation/workflow/), grounded in a survey of both existing orchestration regimes and the Snakemake v8/v9 feature set, and it has been through an adversarial review pass (four independent reviewers + verification against the code and Snakemake source) before posting. Diagrams below; comments very welcome — nothing here is implemented yet.

1. What the orchestration does today

Two regimes exist side by side:

(A) CANFAR (canonical, scripts/sh/): three levels — curl_canfar_local.sh (dispatch: loops image IDs, spawns one headless Skaha container per ID, throttles via batch/sleep polling), init_run_exclusive_canfar.sh (per-ID run-dir surgery: symlinks global run dirs into a per-ID output/, dedup by rm -rf run_sp_*, the -F repair path, sqlite size sanity checks, the ngmix _prev keep-the-best-partial scheme), and job_sp_canfar.bash (bit-coded module driver: each bit runs shapepipe_run -c <config>.ini, with per-ID restriction emulated by rewriting NUMBER_LIST into config copies since -e was retired in #746). A v2.0 generation (*_v2.0.bash, init_run_v2.0.py) partially coexists with v1.1.

(B) Nibi (production, monolithic): one sbatch script (job_p3_batch1.sh) runs the full module chain over all tiles in one shared run dir on one exclusive node; ShapePipe's SMP engine provides within-step parallelism (SMP_BATCH_SIZE); resume is whole-step (START_STEP); success is judged by "did the step produce any output files" because shapepipe_run exits nonzero if any of its ~2000 sub-processes failed (per-CCD star-selection attrition of ~0.2% is normal in the setools step). Star catalogues for masking are pre-generated out-of-band on a login node (gen_star_cats.sh — network access to VizieR/GSC, plus a 40-per-exposure per-CCD symlink fan-out), and a helper (make_exp_forest.sh) fakes the per-exposure directory hierarchy ($SP_EXP/<2-digit-prefix>/<exp>/output/...) that tile-level multi-epoch modules expect, by symlinking the monolithic run dirs into per-exposure stubs.

What both regimes share, and what Snakemake replaces: input discovery (tile → exposures via image headers), bookkeeping (which units are done — currently directory existence and log files), dedup (an exposure shared by several tiles is processed once), config mutation as control flow (sed/perl rewriting SMP_BATCH_SIZE, NUMBER_LIST, ngmix ID_OBJ_MIN/MAX into config copies), scatter/gather sequencing, retry/repair, and cluster-specific job submission.

What stays: every shapepipe_run -c config.ini module invocation, the module code, the config-file grammar, the image-number conventions.

2. Design at a glance

flowchart TD
    subgraph tile_front ["per tile {T}"]
        GitT["stage_tile: image+weight"] --> Uz["uncompress_weight"]
        GitT --> Fe["find_exposures → exp_numbers-{T}.txt"]
        GitT --> ScT["star_cat_tile (GSC query — network)"]
    end
    Fe --> IDX{{"checkpoint build_index:
    manifest.sqlite (tile↔exp↔ccd, immutable)"}}
    IDX -.->|"SELECT DISTINCT exp_id"| GitE
    subgraph exp_chain ["per exposure {E} (deduplicated across tiles)"]
        GitE["stage_exposure: image+weight+flag"] --> Sp["split_exp → 40 CCDs + headers-{E}.npy"]
        GitE --> ScE["star_cat_exp (network) + per-CCD fan-out"]
        Sp --> MaE["mask_exp (per CCD)"]
        ScE --> MaE
        MaE --> Psf["sextractor+setools+psfex+interp (per CCD)
        → PSF models + verdict-{E}"]
    end
    subgraph tile_back ["per tile {T}"]
        Mh["merge_headers → log_exp_headers-{T}.sqlite"]
        MaT["mask_tile"] --> SxT["sextractor_tile → objects"]
        SxT --> PiVi["psfex_interp ME + vignetmaker ×2"]
        Mh --> PiVi
        PiVi --> Ng["ngmix chunks {T}.{k} (fixed k per tile)"]
        Ng --> Merge["merge_sep_cats"]
        Merge --> Cat["make_cat → final_cat-{T}.fits"]
    end
    Uz --> MaT
    ScT --> MaT
    Fe -->|"exp list"| Mh
    Sp -->|"headers-{E}.npy, all E of T"| Mh
    Psf -.->|"per-CCD products at rendered paths,
    gated by verdicts"| PiVi
    Cat --> All(["rule all: expand(final_cat, T in config tiles)"])
Loading

Principles, each argued below:

  1. One immutable SQLite index for membership; per-unit verdict files for mutable status. The manifest (tile↔exposure↔CCD) is built once per run scope by a checkpoint and is read-only thereafter, read only by the scheduler process — never opened for write by a job. This matters because scratch/project on nibi are NFS, where concurrent-writer SQLite is a corruption trap. All mutable state (done/bad/quarantined) lives in per-unit verdict files — which are declared Snakemake outputs, not log-files-as-markers; a queryable rollup is generated from them serially on demand.
  2. Targets are real outputs with deterministic paths. ShapePipe's native behavior — timestamped run_sp_<module>_<datetime> dirs discovered by globbing, with cross-step references resolved through a run-log file — is incompatible with Snakemake's declared-output model. The generated-config layer (below) therefore pins RUN_DATETIME=False + per-unit RUN_NAME, renders OUTPUT_DIR per unit, and rewrites every last:/run_sp_X:runner INPUT_DIR token to the explicit path of the upstream unit's rule output. ShapePipe's run-log name resolution is deliberately not used.
  3. Rules run at tile and exposure granularity; CCDs stay inside shapepipe_run's SMP engine. One Snakemake job = one shapepipe_run invocation restricted to one unit.
  4. Generated configs are tracked artifacts, not control flow. The sed/perl mutation becomes a render step whose outputs are declared inputs of the module rules (write-if-unchanged, so re-rendering identical content never triggers downstream reruns).
  5. Expected failure is data, not error. Verdict outputs + filter-at-gather for expected-by-data failures; --retries and --keep-going for bad luck; the two never mix (§5 spells out the output contract that makes this actually true in Snakemake).
  6. Machine specifics live in profiles. The same Snakefile runs on a laptop subset, inside one nibi allocation, across many nibi nodes, and (via cluster-generic) on CANFAR.

3. The manifest (tile ↔ exposure ↔ CCD index)

A checkpoint build_index rule consumes the per-tile find_exposures outputs (exp_numbers-{T}.txt — so find_exposures survives as a rule; the index doesn't re-parse headers) and writes manifest.sqlite:

  • tiles(tile_id, ra, dec)
  • exposures(exp_id, n_ccd)
  • tile_exp(tile_id, exp_id) — the many-to-many map

Everything downstream resolves membership through it: per-exposure rules exist because some tile references them (dedup is a SELECT DISTINCT exp_id, not a filesystem check); the tile gather rules enumerate their exposure inputs from it.

Design points, sharpened by review:

  • Checkpoint semantics, precisely. Tile→exposure membership is only knowable after tile images are staged, so build_index is a Snakemake checkpoint. Downstream input functions call checkpoints.build_index.get() — which forces the checkpoint to run and the DAG to re-evaluate — and then read the sqlite once into a module-level memoized dict (keyed on checkpoint completion; it cannot be loaded at parse time, since on a fresh run the file doesn't exist yet). One checkpoint, one re-evaluation, no per-job queries.
  • Immutable, and NFS-safe by construction. The manifest is written once by the checkpoint job and never again. Status/quarantine columns do not live here (see principle 1); nothing concurrent ever writes it.
  • v1.4 lesson (index-first serialization). The old global index forced "finish job 1 everywhere before job 2 anywhere." Here the index is per run scope (the configured tile list), built in seconds from staged headers, and the DAG means a 4-tile run touches only 4 tiles' worth of the graph: snakemake final_cat-196.302.fits works. A 1268-tile P3 run builds one index once.
  • DAG scale. Benchmarked (Snakemake 9.x, synthetic workflow of this shape): ~180k-job DAG dry-runs in ~3 min / 2.5 GB, scaling linearly — full DR6 (~2–4×10⁵ jobs) is workable in one invocation. If campaign partitioning is ever wanted, --batch applies to the terminal aggregation rule (it partitions that rule's inputs); note --batch + checkpoint is permitted but untested territory, and per-patch tile lists achieve the same thing more simply.

4. Rules, granularity, and the generated-config layer

One Snakemake job = one shapepipe_run on one unit (one tile, or one exposure). Granularity rationale:

  • Per-CCD Snakemake jobs would mean ~40× job count for seconds-to-minutes tasks — scheduler poison. The SMP engine already parallelizes CCDs within an exposure job; SMP_BATCH_SIZE is rendered from the rule's cpus_per_task.
  • Per-exposure jobs give dedup where it matters: the exposure chain from split_exp onward is one wildcard family {E}, requested by every tile that needs it via the manifest, executed once.
  • Per-tile jobs for the tile front and back.

Config → rule mapping (the current production chain, enumerated — PSFEx path; each rule's config is rendered per unit into generated/{rule}/{unit}.ini):

# Current config / script Snakemake rule Unit Notes
1 config_tile_Git_vos.ini stage_tile tile retrieve/symlink image+weight
2 config_tile_Uz.ini uncompress_weight tile
3 config_tile_Fe.ini find_exposures tile reads tile image header
(new) build_index (checkpoint) run scope consumes all exp_numbers-{T}.txt
4 config_exp_Gie_vos.ini stage_exposure exposure see restructuring note below
5 config_exp_Sp.ini (split half of SpMh) split_exp exposure 40 CCDs + headers-{E}.npy
gen_star_cats.sh / create_star_cat.py star_cat_tile, star_cat_exp tile / exposure network step — see below
6 config_exp_Ma_onthefly.ini mask_exp exposure (per-CCD inside) consumes star_cat_exp
7 config_exp_psfex.ini psf_exp exposure (per-CCD inside) sextractor+setools+psfex+interp; verdict output
8 config_tile_Mh_exp.ini merge_headers tile per-tile gather of its exposures' headers-{E}.npylog_exp_headers-{T}.sqlite
9 config_tile_Ma_onthefly.ini mask_tile tile consumes star_cat_tile
10 config_tile_Sx.ini sextractor_tile tile object detection
11 config_tile_PiViVi_canfar_sx.ini psf_vignets_tile tile psfex_interp ME + vignetmaker ×2
12 config_tile_Ng{k}u.ini (from template) ngmix_chunk tile × chunk fixed chunk count; per-tile boundaries
13 config_merge_sep_cats.ini merge_sep_cats tile
14 config_make_cat_psfex_nosm.ini make_cat tile final_cat-{T}.fits

The generated-config layer is a full I/O rewrite, not a three-field tweak. Review against the code showed the tile-back configs wire inputs by named run references (INPUT_DIR = run_sp_tile_Sx:sextractor_runner, last:...), resolved by substring-searching a run-log file in the step's output dir — which breaks under both per-unit isolation (upstream line absent → "no runs found") and shared dirs with reruns (two matching lines → ambiguous). And RUN_DATETIME defaults to True, so output dirs are timestamp-named — undeclarable as Snakemake outputs, invisible to incomplete-run cleanup. So render_config rewrites, per unit: NUMBER_LIST (unit restriction), SMP_BATCH_SIZE (from cpus_per_task), RUN_DATETIME=False + RUN_NAME (deterministic dirs), OUTPUT_DIR (per-unit path; for exposures, the $SP_EXP/<prefix>/<base>/output forest path with the strip-trailing-letter base convention), every INPUT_DIR token (explicit upstream paths), and for ngmix ID_OBJ_MIN/MAX. Rendering is write-if-unchanged. The rule wrapper clears the unit's own stale run dir before invoking (ShapePipe errors if a fixed-name run dir already exists — this is also the retry path).

The exposure forest becomes real as a consequence: each exposure job's OUTPUT_DIR is its forest entry, in exactly the layout the multi-epoch tile modules glob. make_exp_forest.sh is deleted, not ported.

Exposure staging is restructured, honestly. get_images at the exposure level is today tile-indexed: it ingests a tile's exp_numbers-{T}.txt and stages that tile's whole exposure list (so NUMBER_LIST cannot restrict it to one exposure, and it re-visits shared exposures per tile). The workflow instead has the manifest emit one-line list files generated/exp_list/{E}.txt, and stage_exposure renders get_images to read that — a genuinely per-exposure staging rule; dedup lives in the manifest, not the get step. (On nibi, staging is retrieve=symlink off the local store, so this is cheap; with real VOSpace downloads it's also where --retries earns its keep.)

Star catalogues become explicit rules. Both mask steps consume pre-generated GSC star catalogues (USE_EXT_STAR=True), produced today by an out-of-band login-node script (network access to VizieR; per-CCD symlink fan-out for exposures; container cert workarounds). The design makes them rules (star_cat_tile, star_cat_exp) with the fan-out inside — and flags them as the one stage that needs network, which the execution model must accommodate (§7): on nibi, run them under a login/DTN-side profile or as a pre-staging invocation; their outputs are ordinary declared files the mask rules depend on, so Snakemake sequences the two phases naturally.

ngmix chunking. Today: NSH_JOBS=8 fixed object-ID ranges computed from the average object count across tiles, last chunk open-ended — so above-average tiles pile the overflow into chunk 8 (the observed straggler flaw). The workflow keeps a fixed chunk count (static DAG — a variable count would require a second, per-tile checkpoint after sextractor_tile; explicitly not worth it) but computes per-tile equal boundaries: the ngmix render_config job takes the tile's SExtractor catalog as an input and computes that tile's ID_OBJ_MIN/MAX ranges at run time. No open-ended absorber chunk. ⚠️ Note: this is deliberately not enabled in the parity validation (§9) — chunk boundaries feed ngmix's RNG stream, so re-chunking changes per-object noise realizations; parity runs replicate the current scheme exactly, and the per-tile fix flips on immediately after.

5. Failure handling

Taxonomy and mechanism:

Class Example Mechanism
Transient / infrastructural network hiccup at staging, node OOM, filesystem blip rule fails → --retries 2 (global, with attempt-scaled mem_mb/runtime where useful); --keep-going lets everything not downstream proceed
Expected-by-data too few stars on a CCD (setools), too few objects on a tile verdict outputs (contract below): quarantine recorded, run continues, nothing retried
Novel / real bugs crash, corrupt output rule fails loudly; ensure(non_empty=True) on always-present products; --show-failed-logs prints the culprit log at end of run

The verdict output contract, precisely. Snakemake fails any job whose declared outputs are missing — even at exit 0 — so "product may legitimately not exist" and "product is the declared output" cannot coexist. The design commits to the following split:

  • Rules with no expected-by-data failure mode (staging, split, merge_headers, sextractor_tile, vignets, merge, make_cat): declared outputs are the science products, guarded by ensure(non_empty=True). Missing/empty product = real failure = retry/report.
  • Rules with an expected-by-data failure mode (today: psf_exp via setools star selection; tile-level object-count guards if adopted): the declared output is the verdict file (always written: ok or bad: <reason> + per-CCD detail), which the rule produces on every clean execution. The science products (PSF models per CCD) are written at rendered, deterministic paths but are not declared outputs — Snakemake tracks the step through its verdict. Consequence, stated plainly: those products lose per-file Snakemake tracking (no MissingOutputException safety net, no temp()); the verdict file's content is the tracking. Downstream gathers (psf_vignets_tile) declare all relevant verdict files as inputs and filter bad CCDs inside the job — exactly the attrition tolerance the pipeline already has, now with the quarantine list as a queryable artifact instead of folklore.
  • Bad data is never retried (a verdict is a successful job); bad luck never quarantines (a failed job produces no verdict). The taxonomy is enforced by the output contract, not by convention.

Consequences we care about:

  • 1% of jobs failing does not halt the run. --keep-going + per-unit granularity: a failed exposure blocks only tiles that need it; everything else completes. The end-of-run report lists every failure with rule, wildcards, and log path; --log-handler-script (v8) / a logger plugin (v9) receives structured per-job error records, from which a failed_jobs.jsonl is a few lines of code — debuggable in parallel while the run continues.
  • Fix upstream, downstream heals. Re-running after fixing an input marks exactly the affected descendants outdated — native behavior, no hand bookkeeping. Killed runs resume with --rerun-incomplete; the deterministic run-dir contract (§4) plus pre-invocation stale-dir cleanup makes that actually safe.
  • Fail-early checks live where the knowledge is. Predictable expected-failures (star count below threshold) are checked in the verdict step before burning compute; failures that only manifest downstream keep their verdict at the step that detects them.
  • shapepipe_run's unreliable exit code is contained: each invocation covers one unit, and the wrapper judges success by the unit's declared outputs and verdict content, not the raw exit code.

6. Debuggability

Treated as a requirement (it is the main pain point with Snakemake in practice):

  • Every rule declares log: "logs/{rule}/{unit}.log" with explicit > {log} 2>&1 — one hop from "rule X failed for tile Y" to the actual stderr; the path is printed in the failure block, and --show-failed-logs dumps failed logs to the console at end of run.
  • SLURM-level logs: the profile pins --slurm-logdir somewhere findable; rule log: files are the durable record.
  • benchmark: on the heavy rules (ngmix, psf_exp) gives per-unit wall/RSS/CPU TSVs for free — the throughput model stops being hand-assembled.
  • Hand-running one step stays possible, by construction: the generated config for any unit sits in generated/, so apptainer exec <sif> shapepipe_run -c generated/ngmix_chunk/196.302.1.ini reproduces exactly what the rule ran. Hard requirement, not a convenience.
  • snakemake --dry-run (with --quiet at scale) answers "what would run and why" before anything launches.

7. Execution modes and profiles

The Snakefile is machine-agnostic; profiles carry the rest:

  • Mode A — one allocation (the current nibi model): sbatch one exclusive node → apptainer exec <sif> snakemake --cores $SLURM_CPUS_ON_NODE --resources mem_mb=<node> with the local executor. Snakemake packs jobs against the core/memory budget using honest per-rule threads/mem_mb (measured: ngmix ~2.3 GB/worker at SAVE_BATCH=250; benchmarks accumulate thereafter) instead of hand-tuned SMP_BATCH_SIZE tables. Nested apptainer exec (rules exec-ing inside the outer container) verified working on nibi.
  • Mode B — many nodes: snakemake on a login node with --executor slurm (the executor's docs endorse login-node controllers; ~1% CPU); per-rule resources: map to sbatch; group:/--slurm-array-jobs batch small jobs; --jobs N throttles against nibi's per-user submit limits. This is the P3/DR6 scale path.
  • The network phase: star_cat_tile/star_cat_exp (and real-VOSpace staging, when not pre-staged) need network, which nibi compute nodes lack. They run as a separate profile invocation on the login side — same Snakefile, same DAG; Snakemake simply finds their outputs present when the compute profile runs. One documented seam, not hidden state.
  • CANFAR: no skaha executor plugin exists, and the kubernetes plugin doesn't apply (CANFAR exposes the headless-job REST API, not the k8s API). Honest answer: the cluster-generic executor with three small scripts wrapping the skaha API (submit/status/cancel) gets CANFAR working without plugin authoring; a proper snakemake-executor-plugin-skaha is officially supported to write if CANFAR becomes a primary target. The workflow itself doesn't change either way.
  • Container: rules use an explicit apptainer exec prefix (the sp_validation pattern: container: None, SLURM env vars stripped, OMP_NUM_THREADS=1 pinned at the exec line, binds from config) — identical under both modes, keeps the host snakemake thin, sidesteps sbatch-from-inside-container fragility.
  • Profiles: workflow/profiles/nibi-node/ (mode A), workflow/profiles/nibi-slurm/ (mode B), workflow/profiles/candide/, … each ~15 lines of yaml.

8. Where it lives, and the sims relationship

Proposal: workflow/ at the shapepipe repo root (standard Snakemake layout: workflow/Snakefile, workflow/rules/*.smk, workflow/profiles/<machine>/, workflow/scripts/), configs staying in example/cfis/ as templates.

Naming the tension: the image-sims orchestration deliberately migrated out of shapepipe into sp_validation/workflow/. The split that makes sense: orchestration that produces the shear catalogue lives with the pipeline (this proposal, in shapepipe); orchestration that consumes/validates it stays in sp_validation. The sims workflow currently drives ShapePipe through run_job_sp_canfar_v2.0.bash — once the real-data rules exist, the sims im_pipeline rule can call the same per-unit rules (the pipelines are deliberately near-identical up to the final catalogue), and the bash runner retires from that seam too. Not in scope for the first pass, but the convergence point is designed in.

Reused from the sp_validation precedent (battle-tested): the profile-driven single-command UX, the apptainer exec prefix pattern, config schema fail-fast at parse time, the operational/science/structural config split, retries+latency-wait conventions, provenance capture (git commits + container revision) in the terminal rule.

9. What the first pass does and doesn't do

In scope:

  1. workflow/ skeleton, manifest checkpoint, generated-config layer (the full I/O rewrite of §4).
  2. Rules for the full production chain (table in §4, PSFEx path), tile+exposure granularity, verdict contract on psf_exp.
  3. Explicit star_cat_* rules with the network-profile seam.
  4. nibi-node and nibi-slurm profiles; container exec prefix.
  5. Parity validation: a P3 subset (e.g. the 10-tile set) end-to-end under the workflow with the current ngmix chunking replicated exactly, compared against the existing monolithic run by joining on object ID and comparing science columns (e1/e2, fluxes, PSF quantities) to stated tolerances, ignoring row order and provenance/timestamp headers, reporting matched/dropped counts. (Byte identity is impossible by construction — timestamped provenance, parallel ordering; and chunk boundaries feed ngmix's RNG stream, which is why the chunking fix is fenced off from parity.)
  6. Per-tile ngmix chunk boundaries, enabled immediately after parity is established.

Out of scope (first pass): touching any ShapePipe module; MCCD path; CANFAR profile (documented route, not built); sims-workflow convergence; retiring the CANFAR bash scripts (they keep working untouched).

Success criteria: parity per above on the validation subset; a killed-and-resumed run continues where it stopped with no hand bookkeeping (--rerun-incomplete + deterministic run dirs); an injected expected-failure quarantines that unit and completes the run with the quarantine queryable; an injected real failure blocks exactly its descendants and appears in failed_jobs.jsonl; node utilization on mode A ≥ the current hand-tuned script.

10. Open questions for discussion

  1. Snakemake version pin: ≥9 proposed (current plugin ecosystem); note the structured failed-job feed exists in v8 too (--log-handler-script), so this is about currency, not capability.
  2. Verdict granularity: per-exposure verdict carrying per-CCD detail (proposed) vs 40 per-CCD verdict files per exposure. Proposal keeps file counts sane; the per-CCD detail is inside the verdict.
  3. Manifest scope: per-run-directory index (proposed) vs one shared campaign-level index for dedup across runs (e.g. P3 batches sharing exposures). Schema is identical either way; promoting later is cheap. Related: should exposure products live in a campaign-level forest so runs share them too (today's exp store), and if so, who owns cleanup?
  4. Rerun triggers in production: full triggers (code/params/software-env) rerun aggressively on container upgrades or Snakefile edits; proposal pins --rerun-triggers mtime in production profiles, full triggers during development.
  5. SpMh split: today split_exp and merge_headers are bundled in one early config; the design splits them (split per exposure, merge per tile — matching what the nibi run already does). Any reason to keep them fused?
  6. Anything about the CANFAR timeline that should pull the cluster-generic wrapper into scope earlier?

Drafted from a survey of scripts/sh/ (origin/develop), the nibi production runs (job_p3_batch1.sh et al.), the sp_validation workflow, and the Snakemake docs/source; adversarially reviewed (4 independent review lenses, 30 findings raised, 20 confirmed against code and integrated, 10 refuted) before posting. — Claude (Fable) on behalf of Cail

Metadata

Metadata

Labels

No labels
No labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions