diff --git a/container_templates/base.def b/container_templates/base.def new file mode 100644 index 000000000..eeb218a18 --- /dev/null +++ b/container_templates/base.def @@ -0,0 +1,58 @@ +# Written by scripts/build_containers.py from container_templates/base.def. +# The base image and install steps come from uv.def, or micromamba.def for +# projects with a conda lock file. +Bootstrap: docker +@@FROM@@ + +%files +@@FILES_BLOCK@@ + +%post +# Targets for the data directories bound in at run time +mkdir -p /cvmfs /hdfs /gpfs /ceph /hadoop + +apt-get update +apt-get install -y --no-install-recommends build-essential git +rm -rf /var/lib/apt/lists/* + +@@EXTRA_POST@@ + +# Installing with `--no-cache` unpacks packages into a temporary directory +# under TMPDIR. Apptainer bind-mounts the host's /tmp into %post, so the +# default would use the host's /tmp rather than the build filesystem, which +# can easily run out of space. Keep it inside the image instead, where it +# follows APPTAINER_TMPDIR, and remove it before the image is finalized. +export TMPDIR=/opt/build-tmp +mkdir -p $TMPDIR + +@@INSTALL@@ + +rm -rf /opt/build-tmp + +# Record the environment hash (scripts/env_hash.py) so the pipeline can +# tell whether this image matches the repo it's run from. +# Also record commit it was built from, with "-dirty" if its environment +# files had uncommitted changes. +echo @@ENV_HASH@@ > /opt/env_hash +echo @@BUILD_COMMIT@@ > /opt/build_commit + +@@PURGE_BUILD_TOOLS@@ + +%environment +# Put the environment on PATH. The snakemake apptainer integration uses +# `apptainer exec`, which ignores %runscript, so this is what makes the +# project commands findable. +export PATH="/opt/env/bin:$PATH" + +# Condor jobs may have no usable $HOME, so keep caches (astropy, gwpy, +# pycbc, matplotlib, torch inductor, triton) in the job's scratch directory. +# Outside condor, nothing changes. +if [ -n "${_CONDOR_SCRATCH_DIR:-}" ]; then + export HOME="$_CONDOR_SCRATCH_DIR" + export XDG_CACHE_HOME="$HOME/.cache" + export MPLCONFIGDIR="$XDG_CACHE_HOME/matplotlib" + export TORCHINDUCTOR_CACHE_DIR="$XDG_CACHE_HOME/torchinductor" + export TRITON_CACHE_DIR="$XDG_CACHE_HOME/triton" +fi + +@@EXTRA_ENV@@ diff --git a/container_templates/micromamba.def b/container_templates/micromamba.def index 01d77b39e..48f4cc38a 100644 --- a/container_templates/micromamba.def +++ b/container_templates/micromamba.def @@ -1,28 +1,7 @@ -Bootstrap: docker From: mambaorg/micromamba:2 -Stage: build - -%files -@@FILES_BLOCK@@ - -%environment -# Put the conda env on PATH. The snakemake apptainer integration -# uses `apptainer exec`, which ignores %runscript; prepending -# to PATH here makes the project commands findable. -export PATH=/opt/env/bin:$PATH - -@@EXTRA_ENV@@ %post -mkdir -p /cvmfs /hdfs /gpfs /ceph /hadoop - -apt-get update -apt-get install -y --no-install-recommends git build-essential -rm -rf /var/lib/apt/lists/* - -@@EXTRA_POST@@ - -# activate micromamba and create environment from lockfile +# create the environment from the project's lockfile micromamba create -p /opt/env -f /opt/aframe/projects/@@PROJECT@@/@@PROJECT@@.conda-lock.yml # install uv so we can install local deps of deps editably @@ -32,15 +11,7 @@ micromamba run -p /opt/env python -m \ # Export rather than `uv sync`, so that uv installs into the existing conda # env instead of managing its own. Export to pylock.toml rather than # requirements.txt so that the index each package was locked from is -# recorded. This allows torch to be installed from the PyTorch CPU index. -# `--no-cache` makes uv unpack wheels into a temporary directory under -# TMPDIR. Apptainer bind-mounts the host's /tmp into %post, so the default -# would use the host's /tmp rather than the build filesystem, which -# can easily run out of space. Keep it inside the image instead, where it -# follows APPTAINER_TMPDIR, and remove it before the image is finalized. -export TMPDIR=/opt/build-tmp -mkdir -p $TMPDIR - +# recorded. This allows torch to be installed from the PyTorch CPU index. cd /opt/aframe micromamba run -p /opt/env \ @@UV_CMD@@ --no-cache -o pylock.toml @@ -49,29 +20,4 @@ micromamba run -p /opt/env \ uv pip install --no-cache -r pylock.toml rm pylock.toml -rm -rf /opt/build-tmp micromamba clean -ay - -# Record the environment hash (scripts/env_hash.py) so the pipeline can -# tell whether this image matches the repo it's run from. -# Also record commit it was built from, with "-dirty" if its environment -# files had uncommitted changes. -echo @@ENV_HASH@@ > /opt/env_hash -echo @@BUILD_COMMIT@@ > /opt/build_commit - -@@PURGE_BUILD_TOOLS@@ - -# initialize our shell so that we can execute -# commands in our environment at run time -# set path, and add it to /etc/profile -# so that it will be set if login shell -# is invoked -export PATH=/opt/env/bin:$PATH -echo export PATH=$PATH >> /etc/profile - - -%runscript -#!/bin/bash -eval "$(micromamba shell hook --shell bash)" -micromamba activate /opt/env -exec "$@" diff --git a/container_templates/uv.def b/container_templates/uv.def index 2475d45d7..046c217aa 100644 --- a/container_templates/uv.def +++ b/container_templates/uv.def @@ -1,48 +1,11 @@ -Bootstrap: docker # Important: 0.9.30 is the last uv image built on Debian 12 (bookworm). # Every tag afterwards is Debian 13 (trixie), and NVIDIA's Debian 13 repo # doesn't have CUDA 12 support, which precludes running on V100s. From: ghcr.io/astral-sh/uv:0.9.30-python3.12-bookworm-slim -%files -@@FILES_BLOCK@@ - %post -apt-get update -apt-get install -y --no-install-recommends build-essential git -rm -rf /var/lib/apt/lists/* - -@@EXTRA_POST@@ - cd /opt/aframe/projects/@@PROJECT@@ # Set venv dir outside of project for # when binding the repo into the container export UV_PROJECT_ENVIRONMENT=/opt/env - -# `--no-cache` makes uv unpack wheels into a temporary directory under -# TMPDIR. Apptainer bind-mounts the host's /tmp into %post, so the default -# would use the host's /tmp rather than the build filesystem, which -# can easily run out of space. Keep it inside the image instead, where it -# follows APPTAINER_TMPDIR, and remove it before the image is finalized. -export TMPDIR=/opt/build-tmp -mkdir -p $TMPDIR - @@UV_CMD@@ --no-cache - -rm -rf /opt/build-tmp - -# Record the environment hash (scripts/env_hash.py) so the pipeline can -# tell whether this image matches the repo it's run from. -# Also record commit it was built from, with "-dirty" if its environment -# files had uncommitted changes. -echo @@ENV_HASH@@ > /opt/env_hash -echo @@BUILD_COMMIT@@ > /opt/build_commit - -@@PURGE_BUILD_TOOLS@@ - -%environment -# Append venv dir to PATH so the -# environment is active by default -export PATH="/opt/env/bin:$PATH" - -@@EXTRA_ENV@@ diff --git a/pipeline/config/config.yaml b/pipeline/config/config.yaml index ba8cb9716..795c3f075 100644 --- a/pipeline/config/config.yaml +++ b/pipeline/config/config.yaml @@ -23,6 +23,9 @@ run_dir: /path/to/run # Can be shared across experiments with identical waveform parameters. # waveforms_dir: /path/to/shared/waveforms +# Snakemake and job logs. Defaults to {run_dir}/logs if omitted. +# log_dir: /path/to/logs + # --- Interferometers --------------------------------------------------------- ifos: - H1 @@ -101,6 +104,16 @@ container_source: local # own (/igwn/cit/staging/). Anyone's can be read. osdf_staging_dir: null +# --- Rule placement ---------------------------------------------------------- +# Whether rules run on the node snakemake runs on rather than as batch jobs. +# On LDG, both should be true, though the second is allowed to be false. +# On Delta, both should be false. +# +# Whether train and export, which need a GPU, run locally +gpu_rules_local: true +# Whether the aggregation rules run locally +aggregate_rules_local: true + # --- Resources --------------------------------------------------------------- # Memory (MB) and walltime (minutes) for rules submitted as batch jobs, # under slurm or condor. Rules and keys not listed here use the profile's @@ -205,6 +218,7 @@ integration_window_length: 1.5 # seconds cluster_window_length: 8.0 # seconds Tb: 31536000.0 # target background livetime (seconds) zero_lag: false # also analyze unshifted (zero-lag) data +return_timeseries: false # also write the network's output timeseries # Triton server. model_name: aframe-stream diff --git a/pipeline/profiles/ldg/config.yaml b/pipeline/profiles/ldg/config.yaml index d936d226a..23dbe864e 100644 --- a/pipeline/profiles/ldg/config.yaml +++ b/pipeline/profiles/ldg/config.yaml @@ -66,8 +66,9 @@ apptainer-args: >- default-resources: # Forward environment variables from the submit node. # HOME is not accessible on the execute node, and the bearer - # token is delivered by condor. - getenv: AFRAME_*, PATH, KRB5* + # token is delivered by condor. Fetching proprietary data needs the + # datafind server. + getenv: AFRAME_*, PATH, KRB5*, GWDATAFIND_SERVER # Pass through authentication and account info as raw job ad attributes # using the classad_ prefix. The plugin supports passing attributes diff --git a/pipeline/resources.smk b/pipeline/resources.smk index 042050b32..b9d8481ef 100644 --- a/pipeline/resources.smk +++ b/pipeline/resources.smk @@ -28,9 +28,9 @@ def container(project): `osdf_staging_dir` (your own staging directory by default) through the AP's `/osdf` mount. """ - if config.get("container_source", "local") == "osdf" and project in OSDF_PROJECTS: + if config["container_source"] == "osdf" and project in OSDF_PROJECTS: name = image_name(project, env_hash(project)) - source = config.get("osdf_staging_dir") or staging_dir() + source = config["osdf_staging_dir"] or staging_dir() return f"/osdf{source}/{name}" return os.path.join(os.getenv("AFRAME_CONTAINER_ROOT", ""), f"{project}.sif") @@ -78,9 +78,9 @@ def rule_resources(name): default-resources. With `epnfs`, condor jobs only match execute points that mount the AP's /home. """ - res = config.get("resources", {}).get(name, {}) + res = config["resources"].get(name, {}) out = {} - if config.get("epnfs"): + if config["epnfs"]: out["requirements"] = "TARGET.EPNFS =?= True" if "mem_mb" in res: out["mem_mb"] = out["htcondor_request_mem_mb"] = res["mem_mb"] @@ -99,14 +99,14 @@ def gpu_resources(): instead. """ res = { - "slurm_partition": config.get("inference_partition", "gpuA40x4"), + "slurm_partition": config["inference_partition"], "gpu": 1, "request_gpus": 1, } - if config.get("inference_backend", "export") == "aoti": + if config["inference_backend"] == "aoti": res["require_gpus"] = f"Capability == {config['aoti_gpu_capability']}" else: res["gpus_minimum_capability"] = config["gpu_min_capability"] - if config.get("gpu_min_memory_mb"): + if config["gpu_min_memory_mb"]: res["gpus_minimum_memory"] = f"{config['gpu_min_memory_mb']}M" return res diff --git a/projects/data/data.smk b/projects/data/data.smk index 606e1023c..dcfef8d3e 100644 --- a/projects/data/data.smk +++ b/projects/data/data.smk @@ -6,6 +6,7 @@ Rules: fetch_train_background download one chunk of train strain data fetch_test_background download one chunk of test strain data compute_waveform_branches checkpoint: enumerate (segment, shifts) combos + compute_psd PSDs that waveform jobs rejection-sample against val_waveforms_branch validation waveforms for one branch aggregate_val_waveforms merge per-branch validation waveforms testing_waveforms_branch testing waveforms for one branch @@ -17,8 +18,8 @@ Directory layout: {background_dir}/{train,test}/segments.txt {background_dir}/{train,test}/background-{start}-{duration}.hdf5 - {waveforms_dir}/train/{val_waveforms,training_waveforms}.hdf5 - {waveforms_dir}/test/{waveforms,rejected_parameters}.hdf5 + {waveforms_dir}/train/{val_waveforms,training_waveforms,psd}.hdf5 + {waveforms_dir}/test/{waveforms,rejected_parameters,psd}.hdf5 The fetch rules wildcard over {start} and {duration}. Testing waveforms wildcard over {wbranch_id}, which are indices into the @@ -41,15 +42,16 @@ data_log_dir = log_dir / "data" DATA_CONTAINER = container("data") -num_validation_jobs = int(config.get("num_validation_jobs", 200)) +num_validation_jobs = int(config["num_validation_jobs"]) validation_branch_ids = [str(i) for i in range(num_validation_jobs)] -num_train_waveform_jobs = int(config.get("num_train_waveform_jobs", 10)) +num_train_waveform_jobs = int(config["num_train_waveform_jobs"]) training_branch_ids = [str(i) for i in range(num_train_waveform_jobs)] localrules: compute_waveform_branches, + compute_psd, # DQSegDB seems to be unreachable from execute points generate_train_segments, generate_test_segments, @@ -93,7 +95,7 @@ def _read_segments(segments_file): def _segment_chunks(segments_file): """Split segments into (start, duration) chunks of <= max_duration.""" - max_duration = float(config.get("max_duration", -1)) + max_duration = float(config["max_duration"]) chunks = [] for start, duration in _read_segments(segments_file): step = duration if max_duration == -1 else max_duration @@ -121,14 +123,27 @@ def get_test_background_files(wildcards): ] -def _train_psd_file(wildcards): - """PSD reference for validation waveforms.""" - return get_train_background_files(wildcards)[-1] +# Shortest background file to compute waveform PSDs from +PSD_MIN_DURATION = 2048 -def _test_psd_file(wildcards): - """PSD reference file for testing waveforms""" - return get_test_background_files(wildcards)[-1] +def _psd_background_file(wildcards): + """ + The background file from which to compute PSDs that waveforms are + rejection-sampled against. We take the last file of the split that + is at least PSD_MIN_DURATION long. + """ + if wildcards.split == "train": + files = get_train_background_files(wildcards) + else: + files = get_test_background_files(wildcards) + for fname in reversed(files): + if int(Path(fname).stem.split("-")[-1]) >= PSD_MIN_DURATION: + return fname + raise WorkflowError( + f"No {wildcards.split} background file is at least " + f"{PSD_MIN_DURATION} s long" + ) def _branch_params(wildcards, input): @@ -328,16 +343,38 @@ Runs locally on the submit node. json.dump(branch_map, f, indent=2) +rule compute_psd: + """Compute the PSDs that a split's waveform jobs rejection-sample +against so that each job reads a small file instead of a full +background file. +""" + input: + _psd_background_file, + output: + str(waveform_dir / "{split}" / "psd.hdf5"), + log: + str(data_log_dir / "compute_psd-{split}.log"), + wildcard_constraints: + split="train|test", + container: + DATA_CONTAINER + params: + ifos=config["ifos"], + df=1 / config["waveform_duration"], + script: + "scripts/compute_psd.py" + + rule testing_waveforms_branch: """Generate testing waveforms for one (segment, shifts) branch. -Rejection-samples waveforms against the PSD of the last fetched +Rejection-samples waveforms against the PSDs of the last fetched test-background chunk and writes the accepted injection set and the rejected parameters for this branch. """ input: branch_map=str(test_waveforms / "waveform_branch_map.json"), - psd_file=_test_psd_file, + psd_file=str(test_waveforms / "psd.hdf5"), output: waveforms=str(test_waveforms / "branches" / "{wbranch_id}" / "waveforms.hdf5"), rejected=str( @@ -352,7 +389,6 @@ the rejected parameters for this branch. params: branch=_branch_params, ifos=_fmt_list(config["ifos"]), - output_dir=lambda wc, output: str(Path(output.waveforms).parent), prior=config["prior"], minimum_frequency=config["minimum_frequency"], reference_frequency=config["reference_frequency"], @@ -388,7 +424,8 @@ the rejected parameters for this branch. " --psd_file {input.psd_file}" " --max_num_samples {params.max_num_samples}" " --seed {params.seed}" - " --output_dir {params.output_dir}" + " --waveforms_file {output.waveforms}" + " --rejected_file {output.rejected}" " &> {log}" @@ -401,7 +438,7 @@ rule aggregate_testing_waveforms: rejected=str(test_waveforms / "rejected_parameters.hdf5"), log: str(data_log_dir / "aggregate_testing_waveforms.log"), - localrule: config.get("aggregate_rules_local", False) + localrule: config["aggregate_rules_local"] container: DATA_CONTAINER resources: @@ -416,10 +453,10 @@ rule val_waveforms_branch: """Generate one branch of validation waveforms via rejection sampling. num_validation_signals is split evenly across num_validation_jobs branches. -The PSD reference is the last fetched train-background chunk. +The PSDs are those of the last fetched train-background chunk. """ input: - psd_file=_train_psd_file, + psd_file=str(train_waveforms / "psd.hdf5"), output: str(train_waveforms / "validation_tmp" / "waveforms-{vbranch_id}.hdf5"), log: @@ -473,7 +510,7 @@ rule aggregate_val_waveforms: str(train_waveforms / "val_waveforms.hdf5"), log: str(data_log_dir / "aggregate_val_waveforms.log"), - localrule: config.get("aggregate_rules_local", False) + localrule: config["aggregate_rules_local"] container: DATA_CONTAINER resources: @@ -484,7 +521,7 @@ rule aggregate_val_waveforms: "scripts/aggregate_val_waveforms.py" -if config.get("pregenerate_training_waveforms", False): +if config["pregenerate_training_waveforms"]: rule training_waveforms_branch: """Generate one branch of training waveform polarizations. @@ -534,7 +571,7 @@ if config.get("pregenerate_training_waveforms", False): str(train_waveforms / "training_waveforms.hdf5"), log: str(data_log_dir / "aggregate_training_waveforms.log"), - localrule: config.get("aggregate_rules_local", False) + localrule: config["aggregate_rules_local"] container: DATA_CONTAINER resources: diff --git a/projects/data/data/waveforms/testing.py b/projects/data/data/waveforms/testing.py index 17f59e840..cfa72802b 100644 --- a/projects/data/data/waveforms/testing.py +++ b/projects/data/data/waveforms/testing.py @@ -27,7 +27,8 @@ def testing_waveforms( snr_threshold: float, psd_file: Path, max_num_samples: int, - output_dir: Path, + waveforms_file: Path, + rejected_file: Path, jitter: float = 0.1, seed: int | None = None, ): @@ -85,14 +86,15 @@ def testing_waveforms( that result in an SNR below this threshold will be rejected, but saved for later use psd_file: - Background file from which to calculate PSDs used for - estimating waveforms SNR + PSDs written by `utils.write_psds`, or a background file from + which to calculate them, used for estimating waveforms SNR max_num_samples: Maximum number of samples to generate at once in the rejection sampling process. - output_dir: - Directory to which the waveform file and rejected parameter - file will be written + waveforms_file: + File to which the accepted waveforms will be written + rejected_file: + File to which the rejected parameters will be written jitter: Scale of random jitter to add to injection times seed: @@ -133,7 +135,7 @@ def testing_waveforms( # calculate psd that will be used for snr calculation df = 1 / waveform_duration - logging.info(f"Using background file {psd_file} for psd calculation") + logging.info(f"Using {psd_file} for PSDs") psds = utils.load_psds(psd_file, ifos, df=df) # perform the rejection sampling @@ -166,17 +168,15 @@ def testing_waveforms( parameters["injection_time"] = injection_times parameters["shift"] = np.array([shifts for _ in range(num_signals)]) - output_dir.mkdir(parents=True, exist_ok=True) + for fname in (waveforms_file, rejected_file): + fname.parent.mkdir(parents=True, exist_ok=True) response_set = ResponseSet(**parameters) - waveform_fname = output_dir / "waveforms.hdf5" - utils.io_with_blocking(response_set.write, waveform_fname) - - rejected_fname = output_dir / "rejected_parameters.hdf5" - utils.io_with_blocking(rejected_params.write, rejected_fname) + utils.io_with_blocking(response_set.write, waveforms_file) + utils.io_with_blocking(rejected_params.write, rejected_file) # TODO: compute probability of all parameters against # source and all target priors here then save them somehow - return waveform_fname, rejected_fname + return waveforms_file, rejected_file def main(args): diff --git a/projects/data/data/waveforms/utils.py b/projects/data/data/waveforms/utils.py index d9624c708..8c155f300 100644 --- a/projects/data/data/waveforms/utils.py +++ b/projects/data/data/waveforms/utils.py @@ -4,6 +4,7 @@ from pathlib import Path from zlib import adler32 +import h5py import numpy as np import torch from gwpy.timeseries import TimeSeriesDict @@ -65,8 +66,8 @@ def calc_segment_injection_times( return injection_times -def load_psds(background: Path, ifos: list[str], df: float) -> torch.Tensor: - """Calculate PSDs from background generated by `background.py`""" +def compute_psds(background: Path, ifos: list[str], df: float) -> torch.Tensor: + """Calculate PSDs from a background file written by `fetch-data`""" background = TimeSeriesDict.read(background, path=ifos) psds = [] for ifo in ifos: @@ -76,6 +77,30 @@ def load_psds(background: Path, ifos: list[str], df: float) -> torch.Tensor: return psds +def write_psds(psds: torch.Tensor, fname: Path, ifos: list[str], df: float): + """Write PSDs from `compute_psds`, one dataset per ifo, for `load_psds`""" + with h5py.File(fname, "w") as f: + f.attrs["df"] = df + for ifo, psd in zip(ifos, psds, strict=True): + f[ifo] = psd.numpy() + + +def load_psds(fname: Path, ifos: list[str], df: float) -> torch.Tensor: + """ + Load PSDs written by `write_psds`, or calculate them from a background + file if that's what `fname` is + """ + with h5py.File(fname, "r") as f: + if "df" in f.attrs: + if f.attrs["df"] != df: + raise ValueError( + f"PSDs in {fname} have df={f.attrs['df']}, not {df}" + ) + psds = np.stack([f[ifo][:] for ifo in ifos]) + return torch.tensor(psds, dtype=torch.float64) + return compute_psds(fname, ifos, df) + + def io_with_blocking(f, fname, timeout=10): """ Function that assists with multiple processes writing to the same file diff --git a/projects/data/scripts/compute_psd.py b/projects/data/scripts/compute_psd.py new file mode 100644 index 000000000..25e4642ac --- /dev/null +++ b/projects/data/scripts/compute_psd.py @@ -0,0 +1,18 @@ +# ruff: noqa: F821 +"""Compute the PSDs that waveform jobs rejection-sample against. + +Executed via the snakemake `script:` directive. +The `snakemake` object is injected by snakemake. +""" + +import sys + +from data.waveforms.utils import compute_psds, write_psds + +# `script:` directives are not auto-redirected to the rule's log, so +# send stdout/stderr there to capture tracebacks on failure. +sys.stdout = sys.stderr = open(snakemake.log[0], "w", buffering=1) + +ifos, df = snakemake.params.ifos, snakemake.params.df +psds = compute_psds(snakemake.input[0], ifos, df) +write_psds(psds, snakemake.output[0], ifos, df) diff --git a/projects/export/export.smk b/projects/export/export.smk index e927a85c8..4085948d2 100644 --- a/projects/export/export.smk +++ b/projects/export/export.smk @@ -9,7 +9,7 @@ export_log_dir = log_dir / "export" EXPORT_CONTAINER = container("export") -remote_train = config.get("remote_train", False) +remote_train = config["remote_train"] def _train_artifacts(wildcards): @@ -30,13 +30,13 @@ snakemake is invoked. model_repo=directory(str(export_out / "model_repo")), log: str(export_log_dir / "export.log"), - localrule: config.get("gpu_rules_local", True) + localrule: config["gpu_rules_local"] container: EXPORT_CONTAINER # never reaches condor, so slurm GPU keys only resources: **rule_resources("export"), - slurm_partition=config.get("inference_partition", "gpuA40x4"), + slurm_partition=config["inference_partition"], gpu=1, params: preprocessor=config["export_preprocessor"], @@ -46,7 +46,7 @@ snakemake is invoked. inference_sampling_rate=config["inference_sampling_rate"], batch_size=config["inference_batch_size"], fduration=config["fduration"], - fftlength=config.get("fftlength") or "null", + fftlength=config["fftlength"] or "null", psd_length=config["psd_length"], highpass=config["highpass"], streams_per_gpu=config["streams_per_gpu"], diff --git a/projects/infer/infer.smk b/projects/infer/infer.smk index 3388d96c4..040bf1fc1 100644 --- a/projects/infer/infer.smk +++ b/projects/infer/infer.smk @@ -25,11 +25,11 @@ from pathlib import Path triton_dir = run_dir / "triton" infer_dir = run_dir / "infer" infer_log_dir = log_dir / "infer" -zero_lag = config.get("zero_lag", False) -return_timeseries = config.get("return_timeseries", False) +zero_lag = config["zero_lag"] +return_timeseries = config["return_timeseries"] -INFERENCE_MODE = config.get("inference_mode", "triton") -INFERENCE_BACKEND = config.get("inference_backend", "export") +INFERENCE_MODE = config["inference_mode"] +INFERENCE_BACKEND = config["inference_backend"] if INFERENCE_MODE not in ("triton", "inprocess"): raise WorkflowError( f"inference_mode must be 'triton' or 'inprocess', got {INFERENCE_MODE}" @@ -39,7 +39,7 @@ if INFERENCE_BACKEND not in ("export", "compile", "aoti"): f"inference_backend must be 'export', 'compile', or 'aoti', got {INFERENCE_BACKEND}" ) -ANALYSIS_TYPE = config.get("analysis_type", "hdf5") +ANALYSIS_TYPE = config["analysis_type"] if ANALYSIS_TYPE not in ("hdf5", "rnp"): raise WorkflowError(f"analysis_type must be 'hdf5' or 'rnp', got {ANALYSIS_TYPE}") if ANALYSIS_TYPE == "rnp" and INFERENCE_MODE == "triton": @@ -63,7 +63,7 @@ def check_triton_image(): AOTI_PKG = str(export_out / "model_aoti.pt2") -BRANCHES_PER_JOB = config.get("branches_per_job", 1) +BRANCHES_PER_JOB = config["branches_per_job"] if ANALYSIS_TYPE == "rnp": @@ -111,13 +111,15 @@ def _assign_groups(branch_map, split_at_file_changes): last = branch["fname"] -def _group_branches(branch_map, group_id): - return [b for b in branch_map.values() if b["group"] == group_id] +def _group_branches(group_id): + """The branches for one group""" + branch_map = _load_branch_map() + return {i: b for i, b in branch_map.items() if b["group"] == int(group_id)} def get_infer_group_inputs(wildcards): """Strain files, plus their testing waveforms, for one group.""" - branches = _group_branches(_load_branch_map(), int(wildcards.group_id)) + branches = _group_branches(wildcards.group_id).values() inputs = {"background": sorted({b["fname"] for b in branches})} if ANALYSIS_TYPE == "hdf5": inputs["waveforms"] = [b["waveforms"] for b in branches if b["waveforms"]] @@ -243,6 +245,38 @@ else: json.dump(branch_map, f, indent=2) +rule infer_group_branches: + """Write one group's branches for infer_group. Contains the ID, shifts, +and indices into infer_group's --background and --waveforms lists, so +that the job finds its files only through its declared inputs. + +Both rules take their file lists from get_infer_group_inputs, so the +indices match infer_group's inputs. +""" + input: + str(infer_dir / "branch_map.json"), + output: + str(group_dir / "branches.json"), + localrule: True + run: + inputs = get_infer_group_inputs(wildcards) + background = inputs["background"] + waveforms = inputs.get("waveforms", []) + branches = [ + { + "id": i, + "shifts": b["shifts"], + "background": background.index(b["fname"]), + "waveforms": ( + waveforms.index(b["waveforms"]) if b.get("waveforms") else None + ), + } + for i, b in _group_branches(wildcards.group_id).items() + ] + with open(output[0], "w") as f: + json.dump(branches, f, indent=2) + + _group_common_params = dict( analysis_type=ANALYSIS_TYPE, ifos="[" + ",".join(config["ifos"]) + "]", @@ -255,8 +289,9 @@ _group_common_params = dict( ) _GROUP_SHELL_SUFFIX = ( - " --branch_map {input.branch_map}" - " --group_id {wildcards.group_id}" + " --background {input.background}" + + (" --waveforms {input.waveforms}" if ANALYSIS_TYPE == "hdf5" else "") + + " --branches {input.branches}" " --analysis_type {params.analysis_type}" " --background_out {output.background}" " --foreground_out {output.foreground}" @@ -280,7 +315,7 @@ if INFERENCE_MODE == "triton": streams_per_gpu = config["streams_per_gpu"] workflow.global_resources["triton_streams"] = streams_per_gpu * num_gpus - rate_per_gpu = config.get("rate_per_gpu") + rate_per_gpu = config["rate_per_gpu"] infer_rate = 2 * rate_per_gpu / streams_per_gpu if rate_per_gpu else "null" # The clients run where the server does because we can't @@ -309,7 +344,7 @@ if INFERENCE_MODE == "triton": gpus=config["gpus"], batch_size=config["inference_batch_size"], triton_image=TRITON_IMAGE, - idle_timeout=config.get("triton_idle_timeout", 3600), + idle_timeout=config["triton_idle_timeout"], script: "scripts/start_triton.py" @@ -317,7 +352,7 @@ if INFERENCE_MODE == "triton": """Stream a group of branches to the Triton server from a local client.""" input: unpack(get_infer_group_inputs), - branch_map=str(infer_dir / "branch_map.json"), + branches=str(group_dir / "branches.json"), triton_started=str(triton_dir / "triton.started"), output: **_group_outputs, @@ -385,7 +420,7 @@ else: """Run a group of branches in-process on one GPU.""" input: unpack(get_infer_group_inputs), - branch_map=str(infer_dir / "branch_map.json"), + branches=str(group_dir / "branches.json"), artifact=_artifact, output: **_group_outputs, @@ -398,18 +433,17 @@ else: **gpu_resources(), params: **_group_common_params, - weights=_artifact, backend=INFERENCE_BACKEND, - aoti_arg=(f" --aoti_path {AOTI_PKG}" if INFERENCE_BACKEND == "aoti" else ""), sample_rate=config["sample_rate"], kernel_length=config["kernel_length"], highpass=config["highpass"], - fftlength=config.get("fftlength") or "null", + fftlength=config["fftlength"] or "null", shell: "infer-local" - " --weights {params.weights}" - " --backend {params.backend}{params.aoti_arg}" - " --sample_rate {params.sample_rate}" + " --weights {input.artifact}" + " --backend {params.backend}" + + (" --aoti_path {input.artifact}" if INFERENCE_BACKEND == "aoti" else "") + + " --sample_rate {params.sample_rate}" " --kernel_length {params.kernel_length}" " --highpass {params.highpass}" " --fftlength {params.fftlength}" diff --git a/projects/infer/infer/cli.py b/projects/infer/infer/cli.py index 4f47888ba..9eedb7066 100644 --- a/projects/infer/infer/cli.py +++ b/projects/infer/infer/cli.py @@ -79,7 +79,27 @@ def _run_branch(client, cfg, branch_id, branch, outdir, rate=None): _write_outputs(outdir, branch_id, results, seq, postproc, cfg) -def _merge_group(cfg, branch_map, group, scratch): +def _resolve_branches(cfg): + """The group's branches, by ID, with the files they read. + + The `--branches` file, written by infer_group_branches in infer.smk, + gives each branch's ID, shifts, and indices into `--background` and + `--waveforms` (null for none). + """ + with open(cfg.branches) as f: + branch_list = json.load(f) + branches = {} + for branch in branch_list: + index = branch["waveforms"] + branches[str(branch["id"])] = { + "fname": cfg.background[branch["background"]], + "waveforms": None if index is None else cfg.waveforms[index], + "shifts": branch["shifts"], + } + return branches + + +def _merge_group(cfg, branches, scratch): """Merge the group's branch outputs into the group's declared outputs. With `zero_lag_out`, branches with all-zero shifts go there rather @@ -89,9 +109,9 @@ def _merge_group(cfg, branch_map, group, scratch): """ split_zero_lag = cfg.zero_lag_out is not None background, zero_lag = [], [] - for branch_id in group: + for branch_id, branch in branches.items(): fname = scratch / branch_id / "background.hdf5" - shifts = branch_map[branch_id]["shifts"] + shifts = branch["shifts"] if split_zero_lag and all(s == 0 for s in shifts): zero_lag.append(fname) else: @@ -105,7 +125,7 @@ def _merge_group(cfg, branch_map, group, scratch): "background": (EventSet, background, cfg.background_out), "foreground": ( foreground_cls, - [scratch / i / "foreground.hdf5" for i in group], + [scratch / i / "foreground.hdf5" for i in branches], cfg.foreground_out, ), } @@ -122,16 +142,13 @@ def _merge_group(cfg, branch_map, group, scratch): if cfg.timeseries_out is not None: merge_timeseries( - [scratch / i / "timeseries.hdf5" for i in group], + [scratch / i / "timeseries.hdf5" for i in branches], cfg.timeseries_out, ) def _run_group(client, cfg, rate=None, reset=None): - with open(cfg.branch_map) as f: - branch_map = json.load(f) - # groups are assigned by compute_branch_map in infer.smk - group = [i for i, b in branch_map.items() if b["group"] == cfg.group_id] + branches = _resolve_branches(cfg) # Per-branch outputs only exist until they're merged. Keep them next # to the group's outputs rather than in /tmp, which may be small. @@ -140,26 +157,27 @@ def _run_group(client, cfg, rate=None, reset=None): with tempfile.TemporaryDirectory(dir=outdir) as scratch: scratch = Path(scratch) with client: - for branch_id in group: + for branch_id, branch in branches.items(): if reset: reset() _run_branch( client, cfg, branch_id=branch_id, - branch=branch_map[branch_id], + branch=branch, outdir=scratch / branch_id, rate=rate, ) - _merge_group(cfg, branch_map, group, scratch) + _merge_group(cfg, branches, scratch) def _shared_args(p): p.add_argument("--config", action=jsonargparse.ActionConfigFile) p.add_argument("--verbose", type=bool, default=False) p.add_argument("--logfile", type=str, default=None) - p.add_argument("--branch_map", type=str) - p.add_argument("--group_id", type=int) + p.add_argument("--background", type=str, nargs="+") + p.add_argument("--waveforms", type=str, nargs="*", default=[]) + p.add_argument("--branches", type=str) p.add_argument("--analysis_type", type=str, default="hdf5") p.add_argument("--background_out", type=str) p.add_argument("--foreground_out", type=str) diff --git a/projects/plots/plots.smk b/projects/plots/plots.smk index 2fecb739d..5f6b32088 100644 --- a/projects/plots/plots.smk +++ b/projects/plots/plots.smk @@ -16,7 +16,7 @@ plots_log_dir = log_dir / "plots" PLOTS_CONTAINER = container("plots") -VETOS = config.get("vetos") +VETOS = config["vetos"] localrules: @@ -76,15 +76,8 @@ rule sensitive_volume: ifos=_fmt_list(config["ifos"]), mass_combos=json.dumps(config["mass_combos"]), source_prior=config["source_prior"], - output_dir=lambda wc, output: str(Path(output.sv_data).parent), - dt=config.get("dt") or "null", - # Omitted entirely when unset. If we pass an empty list, - # we still do a query. - vetos=lambda wc, input: ( - f" --vetos '{_fmt_list(VETOS)}' --veto_segments {input.veto_segments}" - if VETOS - else "" - ), + dt=config["dt"] or "null", + vetos=_fmt_list(VETOS or []), shell: "sensitive-volume" " --background {input.background}" @@ -94,7 +87,15 @@ rule sensitive_volume: " --ifos '{params.ifos}'" " --mass_combos '{params.mass_combos}'" " --source_prior {params.source_prior}" - " --output_dir {params.output_dir}" + " --output_file {output.sv_data}" + " --plot_file {output.sv_plot}" + " --gwtc3_file {output.gwtc3_sv}" " --dt {params.dt}" - "{params.vetos}" - " &> {log}" + # Omitted entirely when unset. If we pass an empty list, + # we still do a query. + + ( + " --vetos '{params.vetos}' --veto_segments {input.veto_segments}" + if VETOS + else "" + ) + + " &> {log}" diff --git a/projects/plots/plots/core/gwtc3.py b/projects/plots/plots/core/gwtc3.py index 61d535ed6..65baa95fd 100644 --- a/projects/plots/plots/core/gwtc3.py +++ b/projects/plots/plots/core/gwtc3.py @@ -309,7 +309,7 @@ def get_logdNs( def _write_result_file( - output_dir: Path, + output_file: Path, detection_criterion: str, detection_thresholds: np.ndarray, pipelines: list[str], @@ -317,10 +317,9 @@ def _write_result_file( sv: dict, err: dict, ) -> None: - """Write the `gwtc-3_pipeline_sv.hdf5` data file.""" - output_dir.mkdir(parents=True, exist_ok=True) - outfile = output_dir / "gwtc-3_pipeline_sv.hdf5" - with h5py.File(outfile, "w") as f: + """Write the pipelines' sensitive volumes to `output_file`.""" + output_file.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(output_file, "w") as f: f.create_dataset(f"{detection_criterion}", data=detection_thresholds) for p in pipelines: g = f.create_group(p) @@ -334,7 +333,7 @@ def main( mass_combos: list[float], detection_criterion: str, detection_thresholds: list[float], - output_dir: Path, + output_file: Path, injection_file: Path | None = None, pipelines: list[str] = None, sig_lognorm: float = 0.1, @@ -387,7 +386,7 @@ def main( ) _write_result_file( - output_dir, + output_file, detection_criterion, detection_thresholds, pipelines, diff --git a/projects/plots/plots/core/sv.py b/projects/plots/plots/core/sv.py index 9f4bd0ae8..321c18e77 100644 --- a/projects/plots/plots/core/sv.py +++ b/projects/plots/plots/core/sv.py @@ -321,6 +321,6 @@ def update(self, result: SensitiveVolumeResult) -> None: def layout(self): return self.grid - def save(self, output_dir: Path) -> None: - self.result.write(output_dir / "sensitive_volume.hdf5") - io.save(self.grid, filename=output_dir / "sensitive_volume.html") + def save(self, data_file: Path, plot_file: Path) -> None: + self.result.write(data_file) + io.save(self.grid, filename=plot_file) diff --git a/projects/plots/plots/main.py b/projects/plots/plots/main.py index c732b1626..21b119189 100644 --- a/projects/plots/plots/main.py +++ b/projects/plots/plots/main.py @@ -92,7 +92,9 @@ def sensitive_volume( ifos: list[str], mass_combos: list[tuple], source_prior: Callable, - output_dir: Path, + output_file: Path, + plot_file: Path, + gwtc3_file: Path, log_file: Path | None = None, dt: float | None = None, max_far: float = 365, @@ -117,8 +119,12 @@ def sensitive_volume( rejected_params: Path to the rejected parameter set. Should be an HDF5 file readable by `ledger.injections.InjectionParameterSet.read` - output_dir: - Path to the directory to save the output plots and data + output_file: + Path to write the sensitive volume data to + plot_file: + Path to write the sensitive volume plot to + gwtc3_file: + Path to write the GWTC-3 pipelines' sensitive volumes to log_file: Path to the log file. If not provided, will log to stdout dt: @@ -191,9 +197,9 @@ def sensitive_volume( injection_file=injection_file, detection_criterion="far", detection_thresholds=result.fars, - output_dir=output_dir, + output_file=gwtc3_file, ) comparisons = comparisons_from_gwtc3_curves( gwtc3_sv, gwtc3_err, mass_combos ) - SensitiveVolumePlot(result, comparisons).save(output_dir) + SensitiveVolumePlot(result, comparisons).save(output_file, plot_file) diff --git a/projects/plots/plots/vizapp/pages/summary/page.py b/projects/plots/plots/vizapp/pages/summary/page.py index a43921998..59a5ed30e 100644 --- a/projects/plots/plots/vizapp/pages/summary/page.py +++ b/projects/plots/plots/vizapp/pages/summary/page.py @@ -35,7 +35,7 @@ def __init__(self, *args, **kwargs): mass_combos=self.app.mass_combos, detection_criterion="far", detection_thresholds=result.fars, - output_dir=self.app.results_dir, + output_file=self.app.results_dir / "gwtc-3_pipeline_sv.hdf5", ) comparisons = comparisons_from_gwtc3_curves( gwtc3_sv, gwtc3_err, self.app.mass_combos diff --git a/projects/plots/tests/test_cli.py b/projects/plots/tests/test_cli.py index 04058f6de..f674a146c 100644 --- a/projects/plots/tests/test_cli.py +++ b/projects/plots/tests/test_cli.py @@ -45,8 +45,12 @@ def _run(paths, output_dir): str([list(c) for c in MASS_COMBOS]), "--source_prior", "priors.priors.end_o3_ratesandpops", - "--output_dir", - str(output_dir), + "--output_file", + str(output_dir / "sensitive_volume.hdf5"), + "--plot_file", + str(output_dir / "sensitive_volume.html"), + "--gwtc3_file", + str(output_dir / "gwtc-3_pipeline_sv.hdf5"), ] cli.main(args) diff --git a/projects/plots/tests/test_gwtc3.py b/projects/plots/tests/test_gwtc3.py index 04facfb4c..811c4019e 100644 --- a/projects/plots/tests/test_gwtc3.py +++ b/projects/plots/tests/test_gwtc3.py @@ -80,7 +80,7 @@ def test_gwtc3_vectorized_matches_reference_loop( mass_combos=MASS_COMBOS, detection_criterion=criterion, detection_thresholds=thresholds, - output_dir=tmp_path / "run", + output_file=tmp_path / "run" / "gwtc-3_pipeline_sv.hdf5", injection_file=injection_file, pipelines=PIPELINES, ) diff --git a/projects/train/train.smk b/projects/train/train.smk index 1d56d2b3b..ba8ee92d3 100644 --- a/projects/train/train.smk +++ b/projects/train/train.smk @@ -24,17 +24,17 @@ TRAIN_CONTAINER = container("train") # GPUs for local training. `train_gpus` pins specific devices on a shared # node; otherwise use `train_num_gpus` of whatever is visible, which under # slurm is the allocation. -if config.get("train_gpus") is not None: +if config["train_gpus"] is not None: TRAIN_GPU_ENV = f"CUDA_VISIBLE_DEVICES={config['train_gpus']} " TRAIN_NUM_GPUS = len(str(config["train_gpus"]).split(",")) else: TRAIN_GPU_ENV = "" - TRAIN_NUM_GPUS = config.get("train_num_gpus", 1) + TRAIN_NUM_GPUS = config["train_num_gpus"] def _train_waveform_inputs(wildcards): """Pre-generated training waveforms, when enabled.""" - if config.get("pregenerate_training_waveforms", False): + if config["pregenerate_training_waveforms"]: return [str(train_waveforms / "training_waveforms.hdf5")] return [] @@ -47,7 +47,7 @@ train_data_params = dict( sample_rate=config["sample_rate"], kernel_length=config["kernel_length"], fduration=config["fduration"], - fftlength=config.get("fftlength") or "null", + fftlength=config["fftlength"] or "null", highpass=config["highpass"], lowpass=config["lowpass"] or "null", ) @@ -68,7 +68,7 @@ train_cli_args = ( ) -if config.get("remote_train", False): +if config["remote_train"]: rule train_remote: """Submit training to Nautilus and wait for the pod. @@ -119,13 +119,13 @@ else: batch=str(train_out / "batch.hdf5"), log: str(train_log_dir / "train.log"), - localrule: config.get("gpu_rules_local", True) + localrule: config["gpu_rules_local"] container: TRAIN_CONTAINER # never reaches condor, so slurm GPU keys only resources: **rule_resources("train"), - slurm_partition=config.get("train_partition", "gpuA40x4"), + slurm_partition=config["train_partition"], gpu=TRAIN_NUM_GPUS, params: **train_data_params, diff --git a/scripts/build_containers.py b/scripts/build_containers.py index 07d41ac40..b3ee198f5 100644 --- a/scripts/build_containers.py +++ b/scripts/build_containers.py @@ -91,16 +91,27 @@ def _get_uv_command(project_name: str, subcommand: str) -> str: def create_definition_file(project_name: str) -> Path: """ - Create the apptainer definition file for a project from the appropriate - template and write it to projects//apptainer.def. + Create the apptainer definition file for a project and write it to + projects//apptainer.def. - Projects with a .conda-lock.yml use the micromamba template; all - others use the uv template. + Every definition is base.def, with its base image and install steps + from micromamba.def for projects with a .conda-lock.yml, and + from uv.def otherwise. """ project_dir = BASE_DIR / project_name is_micromamba = (project_dir / f"{project_name}.conda-lock.yml").exists() template_name = "micromamba.def" if is_micromamba else "uv.def" - template_text = (TEMPLATES_DIR / template_name).read_text() + header, sep, install = ( + (TEMPLATES_DIR / template_name).read_text().partition("\n%post\n") + ) + if not sep: + raise ValueError(f"{template_name} has no %post line") + template_text = ( + (TEMPLATES_DIR / "base.def") + .read_text() + .replace("@@FROM@@", header.strip()) + .replace("@@INSTALL@@", install.strip()) + ) files_block = _get_files_block(project_name) # Optional per-project hooks diff --git a/scripts/env_hash.py b/scripts/env_hash.py index ca6c16dae..e9f487b31 100644 --- a/scripts/env_hash.py +++ b/scripts/env_hash.py @@ -100,6 +100,7 @@ def env_files(project: str) -> list[Path]: template = "micromamba.def" if conda_lock.exists() else "uv.def" candidates = [ ROOT_DIR / "scripts" / "build_containers.py", + TEMPLATES_DIR / "base.def", TEMPLATES_DIR / template, project_dir / "pyproject.toml", project_dir / "apptainer.post",