diff --git a/src/sp_validation/image_sims.py b/src/sp_validation/image_sims.py index e259109d..fa46e1d7 100644 --- a/src/sp_validation/image_sims.py +++ b/src/sp_validation/image_sims.py @@ -152,11 +152,17 @@ def load_catalogs(self, verbose=True): grids_dir = self.cfg["grids_dir"] num = self.cfg["num"] cat_name = self.cfg.get("catalog_name", "shape_catalog_cut_ngmix.fits") + # Sim-directory naming mirrors the workflow's contract: grid sims live in + # ``{branch}_grid_{num}``, all other families (e.g. blended) in + # ``{branch}_{num}``. Defaults to ``grid`` so pre-``sims_type`` configs + # and the synthetic-recovery tests keep the historical layout. + sims_type = self.cfg.get("sims_type", "grid") + suffix = f"_grid_{num}" if sims_type == "grid" else f"_{num}" # ``sim_names`` (incl. the unsheared reference) comes from the config's # branch map. The +g/-g pool estimator pairs the sheared sims directly; # the reference is loaded for completeness and null-test diagnostics. for name in self.sim_names: - path = f"{grids_dir}/{name}_grid_{num}/{cat_name}" + path = f"{grids_dir}/{name}{suffix}/{cat_name}" if verbose: print(f" Loading {path}") self.cats[name] = _load_cat(path, self.w_cols) diff --git a/workflow/image_sims/Snakefile b/workflow/image_sims/Snakefile index d7a89065..02b16288 100644 --- a/workflow/image_sims/Snakefile +++ b/workflow/image_sims/Snakefile @@ -33,9 +33,10 @@ there under ``if "image_sims" in config``. configfile: "workflow/image_sims/config.yaml" -# The image-sims rules own their container invocation explicitly, so no -# top-level container is needed here. -container: None +# Each rule declares its own image via ``container:`` (the ShapePipe vs +# sp_validation split); with ``--software-deployment-method apptainer`` in the +# candide profile, Snakemake wraps every job in ``apptainer exec``. No +# top-level container here -- the per-rule directives own the choice. include: "../rules/image_sims.smk" diff --git a/workflow/image_sims/config.yaml b/workflow/image_sims/config.yaml index 8f719c68..2ade5f7d 100644 --- a/workflow/image_sims/config.yaml +++ b/workflow/image_sims/config.yaml @@ -51,6 +51,13 @@ image_sims: # --- simulation grid -------------------------------------------------- sims_type: grid # 'grid' -> *_grid_{num}; anything else -> *_{num} [operational] num: 1 + # Optional: pull exposures from a *different* realization than tiles. Unset -> + # exposures follow `num` (prior behaviour, bit-for-bit). Only meaningful for + # non-grid sims: tiles stay on `{branch}_{num}`, exposures come from + # `{branch}_{exp_num}`. Reproduces Fabian's blended production split -- the + # blended-family realizations _1.._4 carry stale cross-linked SP_exp (grid-sims + # exposures symlinked in, signal-free); _5 is the clean all-real exposure set. + # exp_num: 5 # Branches this run requests: the unsheared reference plus the four +/- # sheared branches. Injected shear is NOT set here -- it is parsed from each # branch's basic_info.txt by im_manifest. [operational] @@ -63,6 +70,13 @@ image_sims: # ShapePipe cfis configs (final_cat.param etc.); default is # {shapepipe_repo}/example/cfis_image_sims once #766 lands. [operational] config_dir: /n17data/cdaley/unions/scratch-wf/imsims-run/grids/_cfis_image_sims + # ShapePipe config tree the *pipeline* stage runs against (run_job's -c + # override), the one that owns the ngmix ini and hence METACAL_PSF. Empty -> + # run_job self-computes {shapepipe_repo}/example/cfis_image_sims (the + # historical default, byte-exact). Point it at a per-arm copy carrying a + # different ngmix ini to A/B the metacal PSF without editing the worktree. + # [operational] + pipeline_config_dir: "" psf_model: psfex # [operational] n_smp: -1 # [operational] # Extract/calibrate scripts run from the sp_validation repo checkout (branch diff --git a/workflow/profiles/candide/config.yaml b/workflow/profiles/candide/config.yaml index 0316b909..dde1e295 100644 --- a/workflow/profiles/candide/config.yaml +++ b/workflow/profiles/candide/config.yaml @@ -11,29 +11,30 @@ # cluster policy (executor, account, partition, node excludes, per-job # defaults); it carries no science and no workflow logic. # -# What is deliberately NOT here: -# -# * Container / apptainer settings. The image-sims rules set -# ``container: None`` and own their ``apptainer exec`` call through the -# shared ``EXEC`` prefix (one image for every stage, with PYTHONPATH / -# PSF_DICT / OMP_NUM_THREADS injected there). So no -# ``software-deployment-method: apptainer`` / ``apptainer-args`` -- those -# would wrap a *second*, redundant container around jobs that already run -# inside one. +# Containers: the image-sims rules declare their image with Snakemake's own +# per-rule ``container:`` directive (the ShapePipe vs sp_validation split), and +# ``software-deployment-method: apptainer`` here turns that into an +# ``apptainer exec `` around every job. ``apptainer-args`` carries the +# static bind mounts -- the candide paths every job needs mounted. The per-run +# env injections (PYTHONPATH shadowing the branch build, PSF_DICT, OMP_NUM_THREADS +# =1) are NOT here: they are per-run values derived from the run config's repo and +# PSF paths, so they live in committed workflow code (``os.environ`` in +# workflow/rules/image_sims.smk) and reach the container natively -- the slurm +# executor submits ``sbatch --export=ALL`` and apptainer inherits the host env by +# default. Setting them in the .smk, not by hand in the launching shell, keeps +# the "one run command, no implicit uncommitted state" property. # -# * OMP_NUM_THREADS. It is pinned to 1 on the ``apptainer exec`` line in -# workflow/rules/image_sims.smk, not here. The slurm executor submits with -# ``--export=ALL``, which propagates the *driver's* ambient environment; a -# profile only sets CLI flags, never the driver's own env, so an -# ``OMP_NUM_THREADS`` set here would silently depend on the operator having -# exported it by hand. Injecting it at the container boundary puts it where -# the compute runs, committed and independent of the launching shell. +# What is deliberately NOT here: # # * Per-rule resources (mem_mb, runtime). Those live on each rule in the # .smk; the ``default-resources`` below are only the floor for rules that # set none. executor: slurm +software-deployment-method: apptainer +# Bind the candide filesystems every job's inputs/outputs/repos live on. +# /automnt is required when repos or data resolve through the automounter. +apptainer-args: "--bind /n17data,/n09data,/home,/automnt" # Cluster policy applied to every job unless a rule overrides it. The excludes # are the flaky/no-internet candide nodes (n17 mount issues, n09 no internet, diff --git a/workflow/rules/image_sims.smk b/workflow/rules/image_sims.smk index a4eb92d0..36425fb1 100644 --- a/workflow/rules/image_sims.smk +++ b/workflow/rules/image_sims.smk @@ -25,10 +25,12 @@ same split the gate766 baseline ran: ``calibrate`` (-> cut cat) and ``m_bias`` (-> ``m_bias_results.yaml``) run in ``sif`` (the sp_validation image). -Every rule sets ``container: None`` and calls ``apptainer exec`` explicitly -through a shared prefix template (``EXEC_PIPELINE`` / ``EXEC`` -- identical -env injections, different image), because the images are not the workflow's -top-level container. Everything is parameterised under +Every rule declares its image with Snakemake's own ``container:`` directive +(``sif_pipeline`` for the ShapePipe stages, ``sif`` for the sp_validation ones); +``--software-deployment-method apptainer`` in the candide profile turns that +into an ``apptainer exec`` around each job, with the bind mounts and env +injections carried once on ``--apptainer-args`` (see the container block below +and workflow/profiles/candide/config.yaml). Everything is parameterised under ``config["image_sims"]`` -- the two ``sif`` keys, repository roots, data roots, the PSF dictionary, the explicit ``tile_ids`` list and the sim/calibration knobs -- so a fresh user drives it from config alone, with no hard-coded clone @@ -86,11 +88,18 @@ _OPERATIONAL_KEYS = { "branches", "shape", "config_dir", + "pipeline_config_dir", "psf_model", "n_smp", "extract_script", "calibrate_script", } +# Optional override: not required, so an unset key defaults in code (below) to +# NUM -- reproducing prior behaviour bit-for-bit. Listed here only to keep it +# out of the unknown-key error. +_OPTIONAL_KEYS = { + "exp_num", +} # Structural keys: paths/identifiers the run must supply (no sensible default). _STRUCTURAL_KEYS = { "sif", @@ -104,7 +113,11 @@ _STRUCTURAL_KEYS = { "tile_ids", } _ALLOWED_KEYS = ( - _SCIENCE_KEYS | _DEPRECATED_KEYS | _OPERATIONAL_KEYS | _STRUCTURAL_KEYS + _SCIENCE_KEYS + | _DEPRECATED_KEYS + | _OPERATIONAL_KEYS + | _OPTIONAL_KEYS + | _STRUCTURAL_KEYS ) _unknown = set(IMSIM) - _ALLOWED_KEYS @@ -155,6 +168,22 @@ SIMS_TYPE = IMSIM["sims_type"] _SUFFIX = f"_{SIMS_TYPE}_{NUM}" if SIMS_TYPE == "grid" else f"_{NUM}" SIM_BASES = list(IMSIM["branches"]) SIMS = [f"{base}{_SUFFIX}" for base in SIM_BASES] +# Exposure source may be decoupled from the tile source. Tiles stay on +# ``{branch}_{num}`` (SP_tiles is correct in every realization); exposures can be +# pulled from a different realization ``{branch}_{exp_num}`` via the optional +# ``exp_num`` key. Motivation: the blended-family realizations _1.._4 carry +# stale cross-linked SP_exp (grid-sims exposures symlinked in, signal-free); _5 +# is the clean all-real exposure set -- Fabian's production split (exposures +# from _5, tiles from _1). Unset ``exp_num`` -> NUM, reproducing prior behaviour +# bit-for-bit. Only meaningful for non-grid sims (grid names carry SIMS_TYPE, and +# the blended cross-link pathology this addresses is a non-grid concern), so the +# override is guarded to the non-grid branch. +EXP_NUM = IMSIM.get("exp_num", NUM) +if SIMS_TYPE == "grid": + # Grid names embed sims_type; exposure decoupling is not defined there. + SIM_EXP_BASE = {sim: sim for sim in SIMS} +else: + SIM_EXP_BASE = {f"{base}{_SUFFIX}": f"{base}_{EXP_NUM}" for base in SIM_BASES} MANIFEST = f"{GRIDS_BASE}/manifest.yaml" BUILD_MANIFEST = f"{SPV_REPO}/workflow/scripts/im_build_manifest.py" @@ -168,6 +197,13 @@ MASK_CONFIG = IMSIM["mask_config"] # e.g. config/calibration/mask_v1.X.9_im_sim PARAMS_TEMPLATE = f"{SPV_REPO}/workflow/image_sims/params_im_sim.py" # ShapePipe cfis_image_sims config dir (per-tile/exposure configs + final_cat.param). CONFIG_DIR = IMSIM["config_dir"] +# ShapePipe config dir the *pipeline* stage's run_job uses (its ``-c`` override, +# the tree that owns the ngmix ini and thus METACAL_PSF). Empty (the default) +# means "let run_job self-compute {shapepipe_repo}/example/cfis_image_sims" -- +# byte-for-byte the historical behaviour, so an unset key reproduces prior runs +# exactly. Set it (e.g. per A/B arm) to point the ShapePipe stage at a config +# tree carrying a different ngmix ini without touching the shared worktree. +PIPELINE_CONFIG_DIR = IMSIM["pipeline_config_dir"] # ShapePipe scripts live in the ShapePipe repo (also baked into its image). CREATE_FINAL_CAT = f"{SHAPEPIPE_REPO}/scripts/python/create_final_cat.py" @@ -181,10 +217,24 @@ CALIBRATE = IMSIM["calibrate_script"] # m-bias is *this branch's* extracted core, injected on PYTHONPATH. COMPUTE_M_BIAS = f"{SPV_REPO}/scripts/compute_m_bias_image_sims.py" -# --- container exec prefixes ---------------------------------------------- -# One prefix *shape* for every stage -- two instances, one per image. Three -# env injections make the on-disk branch -# code and the sim PSF win over the image's baked copies: +# --- containers ----------------------------------------------------------- +# Each stage runs inside its repo's image via Snakemake's own ``container:`` +# directive (below on each rule): ``sif_pipeline`` (the ShapePipe image) for the +# pipeline/merge stages, ``sif`` (the sp_validation image) for manifest/extract/ +# calibrate/m-bias. With ``--software-deployment-method apptainer`` (set in the +# candide profile) Snakemake wraps every job in ``apptainer exec ``; the +# static bind mounts travel once on ``--apptainer-args`` in that profile. +# +# The three per-run env injections that used to be spelled on every hand-rolled +# ``apptainer exec`` line are set here, in ``os.environ``, at workflow parse +# time. They then reach the container by the native path: the slurm executor +# submits ``sbatch --export=ALL`` (the driver env rides to each job) and +# apptainer inherits the host environment by default. Because they are set in +# committed workflow code -- not exported by hand in the launching shell -- this +# keeps the "one run command, no implicit uncommitted state" property the +# hand-rolled prefix was protecting, while letting Snakemake own the exec. +# ``--apptainer-args`` cannot carry them: it is a single static profile string, +# but these are per-run values derived from the run config's repo/PSF paths. # # * PYTHONPATH prepends BOTH repos' ``src`` (ShapePipe first, then # sp_validation), so Python resolves the worktree build before @@ -196,33 +246,28 @@ COMPUTE_M_BIAS = f"{SPV_REPO}/scripts/compute_m_bias_image_sims.py" # shadowed by PYTHONPATH. # * PSF_DICT points the fake_psf module (PSF_DICT_PATH = $PSF_DICT, expanded # via getexpanded) at this run's PSF dictionary. +# * OMP_NUM_THREADS=1 pins the OpenMP/BLAS pool: the chain is MPI-free (one +# SLURM job per branch x tile, parallelism is ShapePipe's own ``-N n_smp``), +# so the thread pool inside the container must be 1 to avoid oversubscription. +os.environ["PYTHONPATH"] = f"{SHAPEPIPE_REPO}/src:{SPV_REPO}/src" +os.environ["PSF_DICT"] = str(PSF_DICT) +os.environ["OMP_NUM_THREADS"] = "1" + +# No SLURM-env scrubbing is needed. The old prefix ran ``env -u SLURM_*`` to +# keep the ShapePipe stage's OpenMPI from attaching to the host SLURM launcher, +# but the pipeline image's OpenMPI is built ``--with-pmix=internal`` and carries +# no SLURM RAS plugin (verified: ``ompi_info --param ras slurm`` reports the +# framework absent), so it sees a singleton regardless of the host SLURM vars -- +# and ShapePipe drives its own multiprocessing, not MPI-across-tasks. The strip +# was defensive, not load-bearing, so it drops cleanly. # -# The SLURM env vars are stripped (``env -u ...``) so that when the ShapePipe -# pipeline stage's OpenMPI initialises inside the image it does not try to -# attach to the host SLURM launcher (cf. apptainer_noslurm.sh). The strip is -# harmless for the pure-Python sp_validation stages, so one prefix serves all. -# -# ``OMP_NUM_THREADS=1`` is injected here, at the ``apptainer exec`` call, and -# not left to the SLURM profile. The chain is MPI-free: Snakemake fans out one -# job per branch x tile and each job's parallelism is ShapePipe's own internal -# multiprocessing (``-N n_smp``), so the OpenMP/BLAS thread pool inside the -# container must be pinned to 1 to avoid oversubscription. The SLURM profile -# cannot pin it reliably: the slurm executor submits with ``--export=ALL``, -# which propagates the *driver's* ambient environment -- but a Snakemake -# profile only sets CLI flags, never the driver's own env, so an -# ``OMP_NUM_THREADS`` there would depend on the operator having exported it by -# hand (the implicit, uncommitted state the "one run command" is meant to -# retire). Injecting it on the ``apptainer exec`` line puts it where the -# compute actually runs -- inside the container, independent of the driver's -# env -- the same lever this prefix already uses for PYTHONPATH/PSF_DICT. -_EXEC_PREFIX = ( - "env -u SLURM_JOBID -u SLURM_JOB_ID -u SLURM_PROCID " - f"apptainer exec --bind {BINDS} " - f"--env PYTHONPATH={SHAPEPIPE_REPO}/src:{SPV_REPO}/src " - f"--env PSF_DICT={PSF_DICT} --env OMP_NUM_THREADS=1 " -) -EXEC = _EXEC_PREFIX + SIF # sp_validation stages -EXEC_PIPELINE = _EXEC_PREFIX + SIF_PIPELINE # ShapePipe stages +# The m-bias rule is the one exception to the ``container:`` mechanism: it is a +# ``run:`` block, which Snakemake executes in the driver process where the +# ``container:`` directive does not reach. Its single ``shell()`` call keeps an +# explicit ``apptainer exec``, built from the same config (binds + image) so it +# stays in lockstep with the rest; the env injections above serve it too (same +# driver process). +M_BIAS_EXEC = f"apptainer exec --bind {BINDS} {SIF}" JOB_MASK = sum([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048]) @@ -300,8 +345,10 @@ rule im_manifest: input_sims_base=INPUT_SIMS_BASE, sims_type=SIMS_TYPE, num=NUM, + container: + SIF shell: - "{EXEC} python {BUILD_MANIFEST} " + "python {BUILD_MANIFEST} " "--input-sims-base {params.input_sims_base} " "--sims-type {params.sims_type} --num {params.num} " "{params.branch_args} -o {output.manifest}" @@ -335,7 +382,10 @@ rule im_init: run_dir=lambda wc: f"{GRIDS_BASE}/{wc.sim}", cfis=lambda wc: f"{GRIDS_BASE}/{wc.sim}/cfis", sim_tiles=lambda wc: f"{INPUT_SIMS_BASE}/{wc.sim}/images/SP_tiles", - sim_exp=lambda wc: f"{INPUT_SIMS_BASE}/{wc.sim}/images/SP_exp", + # Exposures may come from a different realization than tiles (see EXP_NUM + # / SIM_EXP_BASE above): SIM_EXP_BASE maps the tile sim to its exposure + # sim, which is ``{branch}_{exp_num}`` for non-grid, ``wc.sim`` otherwise. + sim_exp=lambda wc: f"{INPUT_SIMS_BASE}/{SIM_EXP_BASE[wc.sim]}/images/SP_exp", shell: # cfis / input_tiles / input_exp are stable read-only symlinks (used by # get_images, merge, extract); created here but not tracked as outputs, @@ -372,14 +422,22 @@ rule im_pipeline: run_dir=lambda wc: f"{GRIDS_BASE}/{wc.sim}", psf=IMSIM["psf_model"], n_smp=IMSIM["n_smp"], + # ``-c DIR`` only when pipeline_config_dir is set; empty -> run_job + # self-computes its default, so the flag (and its output) is absent for + # a bit-exact reproduction of a run that predates this key. + config_flag=f"-c {PIPELINE_CONFIG_DIR}" if PIPELINE_CONFIG_DIR else "", resources: - mem_mb=16000, + # measured: job uses ~1 core (eff 0.9); post-#843 MaxRSS ~72MB accounting / 2.3GB live-peak; 2cpu/4GB packs ~512 jobs on the usable partition + mem_mb=4000, + cpus_per_task=2, runtime=720, + container: + SIF_PIPELINE shell: "cd {params.run_dir} && " - "{EXEC_PIPELINE} bash {RUN_JOB} " + "bash {RUN_JOB} " "-e {wildcards.tile} -t image_sims -j {JOB_MASK} " - "-p {params.psf} -N {params.n_smp}" + "-p {params.psf} -N {params.n_smp} {params.config_flag}" rule im_merge: @@ -397,9 +455,11 @@ rule im_merge: cat=f"{GRIDS_BASE}/{{sim}}/final_cat_{{sim}}.hdf5", params: run_dir=lambda wc: f"{GRIDS_BASE}/{wc.sim}", + container: + SIF_PIPELINE shell: "cd {params.run_dir} && " - "{EXEC_PIPELINE} python {CREATE_FINAL_CAT} " + "python {CREATE_FINAL_CAT} " "-I -m final_cat_{wildcards.sim}.hdf5 -i .. " "-p cfis/final_cat.param -P {wildcards.sim} " "-o n_tiles_final.txt -v" @@ -418,8 +478,10 @@ rule im_extract: cat=f"{GRIDS_BASE}/{{sim}}/shape_catalog_comprehensive_{SHAPE}.fits", params: run_dir=lambda wc: f"{GRIDS_BASE}/{wc.sim}", + container: + SIF shell: - "cd {params.run_dir} && {EXEC} python {EXTRACT_INFO}" + "cd {params.run_dir} && python {EXTRACT_INFO}" rule im_calibrate: @@ -436,9 +498,11 @@ rule im_calibrate: cat=f"{GRIDS_BASE}/{{sim}}/shape_catalog_cut_{SHAPE}.fits", params: run_dir=lambda wc: f"{GRIDS_BASE}/{wc.sim}", + container: + SIF shell: "cd {params.run_dir} && " - "{EXEC} python {CALIBRATE} -s calibrate" + "python {CALIBRATE} -s calibrate" rule im_mbias: @@ -458,10 +522,14 @@ rule im_mbias: ), output: results=f"{GRIDS_BASE}/results/m_bias_results.yaml", + resources: + # Blended catalogues are ~5x grid size; the snakemake 1000M default OOMs. + mem_mb=8000, params: cfg=f"{GRIDS_BASE}/results/m_bias_config.yaml", grids_base=GRIDS_BASE, num=NUM, + sims_type=SIMS_TYPE, cat_name=f"shape_catalog_cut_{SHAPE}.fits", sif=SIF, sif_pipeline=SIF_PIPELINE, @@ -547,6 +615,7 @@ rule im_mbias: mbias_cfg = { "grids_dir": params.grids_base, "num": params.num, + "sims_type": params.sims_type, "catalog_name": params.cat_name, # Injected shear: from the manifest, the single source of truth. "shear_amplitude": manifest["shear_amplitude"], @@ -564,5 +633,5 @@ rule im_mbias: with open(params.cfg, "w") as fh: yaml.safe_dump(mbias_cfg, fh) shell( - "{EXEC} python {COMPUTE_M_BIAS} -c {params.cfg} -v" + "{M_BIAS_EXEC} python {COMPUTE_M_BIAS} -c {params.cfg} -v" )