diff --git a/.github/project-filters.yaml b/.github/project-filters.yaml index add6e2827..e594982b6 100644 --- a/.github/project-filters.yaml +++ b/.github/project-filters.yaml @@ -3,6 +3,7 @@ workflow: &workflow - '.github/workflows/project-build-test.yaml' - 'container_templates/**' - 'scripts/build_containers.py' + - 'scripts/env_hash.py' - 'pyproject.toml' - 'uv.lock' diff --git a/.github/workflows/project-build-test.yaml b/.github/workflows/project-build-test.yaml index b943e8fbd..de0d21fba 100644 --- a/.github/workflows/project-build-test.yaml +++ b/.github/workflows/project-build-test.yaml @@ -81,7 +81,7 @@ jobs: - name: Write apptainer definition - run: python scripts/build_containers.py --definition-only ${{ matrix.project }} + run: python -m scripts.build_containers --definition-only ${{ matrix.project }} # build the singularity image as a sandbox directory # inside a docker container that has singularity diff --git a/Snakefile b/Snakefile index b25287cab..fa22badd8 100644 --- a/Snakefile +++ b/Snakefile @@ -40,6 +40,11 @@ include: "projects/infer/infer.smk" include: "projects/plots/plots.smk" +onstart: + check_images() + check_triton_image() + + rule all: default_target: True input: diff --git a/container_templates/micromamba.def b/container_templates/micromamba.def index 6a178954b..d368ee53d 100644 --- a/container_templates/micromamba.def +++ b/container_templates/micromamba.def @@ -52,6 +52,10 @@ 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 +echo @@ENV_HASH@@ > /opt/env_hash + @@PURGE_BUILD_TOOLS@@ # initialize our shell so that we can execute diff --git a/container_templates/uv.def b/container_templates/uv.def index f053fa882..dd6469d6b 100644 --- a/container_templates/uv.def +++ b/container_templates/uv.def @@ -31,6 +31,10 @@ mkdir -p $TMPDIR 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 +echo @@ENV_HASH@@ > /opt/env_hash + @@PURGE_BUILD_TOOLS@@ %environment diff --git a/docs/first_pipeline.md b/docs/first_pipeline.md index b46d99bac..db4e6f679 100644 --- a/docs/first_pipeline.md +++ b/docs/first_pipeline.md @@ -80,7 +80,6 @@ In the experiment directory a `run.sh` file will be created that looks like ```bash #!/bin/bash -# Set AFRAME_DEV=1 to bind the working tree into containers. cd /home/albert.einstein/projects/aframe source pipeline/.env snakemake --configfile /home/albert.einstein/aframe/my-first-run/config.yaml --profile pipeline/profiles/ldg diff --git a/libs/architectures/architectures/networks/wavenet.py b/libs/architectures/architectures/networks/wavenet.py index e51091b56..b6d510598 100644 --- a/libs/architectures/architectures/networks/wavenet.py +++ b/libs/architectures/architectures/networks/wavenet.py @@ -3,7 +3,7 @@ import numpy as np import torch import torch.nn as nn -from ml4gw.nn.norm import GroupNorm1DGetter, NormLayer +from ml4gw.nn.norm import GroupNorm1D, NormLayer class GatedActivation(torch.nn.Module): @@ -95,7 +95,7 @@ def __init__( kernel_size: int = 2, norm_layer: NormLayer | None = None, ): - self.norm_layer = norm_layer or GroupNorm1DGetter() + self.norm_layer = norm_layer or GroupNorm1D super().__init__() self.init_conv = nn.Conv1d( in_channels, res_channels, kernel_size=2, dilation=1 diff --git a/libs/architectures/architectures/networks/xylophone.py b/libs/architectures/architectures/networks/xylophone.py index 191ae4904..4a9c5013d 100644 --- a/libs/architectures/architectures/networks/xylophone.py +++ b/libs/architectures/architectures/networks/xylophone.py @@ -1,7 +1,8 @@ from collections import namedtuple import torch -from ml4gw.nn.resnet.resnet_1d import GroupNorm1DGetter, NormLayer, convN +from ml4gw.nn.norm import GroupNorm1D, NormLayer +from ml4gw.nn.resnet.resnet_1d import convN class XylophoneResidualBlock(torch.nn.Module): @@ -87,7 +88,7 @@ def __init__( base_channels: int = 8, ): super().__init__() - self._norm_layer = norm_layer or GroupNorm1DGetter() + self._norm_layer = norm_layer or GroupNorm1D self.base_channels = base_channels self.initial = torch.nn.Sequential( diff --git a/libs/architectures/pyproject.toml b/libs/architectures/pyproject.toml index 2564527fb..76dc65e91 100644 --- a/libs/architectures/pyproject.toml +++ b/libs/architectures/pyproject.toml @@ -8,7 +8,7 @@ license = "MIT" dependencies = [ "einops>=0.8,<0.9", "jaxtyping>=0.2", - "ml4gw>=0.7.2", + "ml4gw>=0.8.4", "h5py>=3.10,<4", "numpy>=2", "torch", diff --git a/libs/ledger/ledger/events.py b/libs/ledger/ledger/events.py index f9bea6957..d20148224 100644 --- a/libs/ledger/ledger/events.py +++ b/libs/ledger/ledger/events.py @@ -62,9 +62,11 @@ class EventSet(Ledger): @classmethod def compare_metadata(cls, key, ours, theirs): - # accumulate background time when merging or appending + # accumulate background time when merging or appending. Files + # aggregated from no sources, e.g. a group with no zero-lag + # branches, have no Tb, which means no livetime. if key == "Tb": - return ours + theirs + return (ours or 0) + (theirs or 0) return super().compare_metadata(key, ours, theirs) def get_shift(self, shift: np.ndarray) -> "EventSet": diff --git a/libs/ledger/tests/test_events.py b/libs/ledger/tests/test_events.py index 04d583453..fe4fb8141 100644 --- a/libs/ledger/tests/test_events.py +++ b/libs/ledger/tests/test_events.py @@ -33,6 +33,22 @@ def test_append(self): assert (obj2.detection_time == times).all() assert (obj2.shift == shifts).all() + def test_aggregate_no_sources(self, tmp_path): + # a file aggregated from no sources has no Tb, and merging it + # with others counts it as no livetime + empty = tmp_path / "empty.hdf5" + events.EventSet.aggregate([], empty) + full = tmp_path / "full.hdf5" + events.EventSet( + np.arange(3.0), np.arange(3.0), np.zeros((3, 2)), 100 + ).write(full) + + merged = tmp_path / "merged.hdf5" + events.EventSet.aggregate([empty, full], merged, clean=False) + obj = events.EventSet.read(merged) + assert obj.Tb == 100 + assert len(obj) == 3 + def test_sorting(self): det_stats = np.arange(10)[::-1] times = np.arange(10) diff --git a/pipeline/.env.example b/pipeline/.env.example index 5bb72bb83..706de9d44 100644 --- a/pipeline/.env.example +++ b/pipeline/.env.example @@ -22,8 +22,3 @@ # --- Weights & Biases --------------------------------------------------------- # Required when a WandbLogger is enabled in the train config # export WANDB_API_KEY= - -# --- Development -------------------------------------------------------------- -# Uncomment to run rules against your current repo instead of the code -# in the images without rebuilding (analog of law's --dev). -# export AFRAME_DEV=1 diff --git a/pipeline/README.md b/pipeline/README.md index 93f651052..b65a605a6 100644 --- a/pipeline/README.md +++ b/pipeline/README.md @@ -13,6 +13,11 @@ cp pipeline/.env.example pipeline/.env # then edit pipeline/.env (git-ignored) # One-time: build the per-project containers uv run build-containers # may need to build one at a time +# One-time, for inference_mode: triton: pull the server image named by +# triton_image in config.yaml +apptainer pull $AFRAME_CONTAINER_ROOT/tritonserver_25.06.sif \ + docker://nvcr.io/nvidia/tritonserver:25.06-py3 + # Initialize a run directory with a config.yaml, train.yaml, and run.sh uv run aframe-init snakemake -d /path/to/my-run # Edit /path/to/my-run/config.yaml to override paramerers in pipeline/config/config.yaml @@ -181,7 +186,9 @@ If these conditions are met, then we can implement various upgrades: must currently name a local image (a bare filename resolved against `$AFRAME_CONTAINER_ROOT`, or an absolute path) because the more modern Triton containers have not been added to CVMFS. We could instead -have a shared cache directory to auto-pull from `ghcr.io/ml4gw/hermes`. +have a shared cache directory to auto-pull from +`nvcr.io/nvidia/tritonserver:-py3`. (hermes' images on +`ghcr.io/ml4gw/hermes` add only a label and stop at 24.12.) **Hyperparameter tuning currently has no Snakemake equivalent.** The law pipeline had a `TuneTask` that stood up a Ray cluster on Kubernetes via diff --git a/pipeline/config/config.yaml b/pipeline/config/config.yaml index 046a51ce0..720fbb4f7 100644 --- a/pipeline/config/config.yaml +++ b/pipeline/config/config.yaml @@ -136,6 +136,9 @@ gpu_min_memory_mb: null # 8.6 is the A10. aoti_gpu_capability: 8.6 +# Only match condor execute points that mount /home +epnfs: false + # --- Training ---------------------------------------------------------------- # Lightning CLI YAML config for `train fit`. train_config: projects/train/train.yaml diff --git a/pipeline/config/small.yaml b/pipeline/config/small.yaml new file mode 100644 index 000000000..30890bab8 --- /dev/null +++ b/pipeline/config/small.yaml @@ -0,0 +1,28 @@ +# A small-scale run that exercises every stage, for checking that the +# pipeline runs. Layer it over config.yaml: +# +# snakemake --configfile pipeline/config/small.yaml --config run_dir=... + +# Continuous H1+L1 open data in O3a, four background files per split +train_start: 1250917000 +train_end: 1250925192 +test_start: 1250925192 +test_end: 1250933384 +max_duration: 2048.0 + +train_config: pipeline/config/small_train.yaml + +num_training_signals: 2000 +num_train_waveform_jobs: 2 + +num_validation_signals: 500 +num_validation_jobs: 2 + +num_testing_signals: 500 + +# Three shifts plus zero lag: 16 inference branches +Tb: 20000.0 +zero_lag: true + +# Groups of more than one branch, still within one file +branches_per_job: 2 diff --git a/pipeline/config/small_train.yaml b/pipeline/config/small_train.yaml new file mode 100644 index 000000000..c6a5a34f5 --- /dev/null +++ b/pipeline/config/small_train.yaml @@ -0,0 +1,86 @@ +# projects/train/train.yaml shrunk for pipeline/config/small.yaml +model: + class_path: train.model.SupervisedAframe + init_args: + arch: + class_path: architectures.supervised.SupervisedTimeDomainResNet + init_args: + layers: [1, 1, 1, 1] + norm_layer: + class_path: ml4gw.nn.norm.GroupNorm1D + init_args: + num_groups: 16 + metric: + class_path: train.metrics.TimeSlideAUROC + init_args: + max_fpr: 1e-3 + pool_length: 8 + weight_decay: 0.0 + learning_rate: 0.000585 + pct_lr_ramp: 0.115 +data: + class_path: train.data.supervised.TimeDomainSupervisedAframeDataset + init_args: + batch_size: 32 + batches_per_epoch: 4 + num_files_per_batch: 2 + chunks_per_epoch: 1 + chunk_size: 1000 + psd_length: 8 + fftlength: null + waveform_prob: 0.277 + swap_prob: 0.014 + mute_prob: 0.055 + left_pad: 0.25 + right_pad: 0.05 + snr_sampler: + class_path: ml4gw.distributions.PowerLaw + init_args: + minimum: 4 + maximum: 100 + index: -3 + waveform_sampler: + class_path: train.data.waveforms.WaveformLoader + init_args: + val_waveform_file: ${oc.env:AFRAME_TRAIN_WAVEFORMS_DIR}/val_waveforms.hdf5 + training_waveform_path: ${oc.env:AFRAME_TRAIN_WAVEFORMS_DIR}/training_waveforms.hdf5 + dec: + class_path: ml4gw.distributions.Cosine + psi: + class_path: torch.distributions.Uniform + init_args: + low: 0 + high: 3.14159 + validate_args: false + phi: + class_path: torch.distributions.Uniform + init_args: + low: 0 + high: 6.28318 + validate_args: false + valid_stride: 0.5 + num_valid_views: 5 + # one background file + min_valid_duration: 2048 + valid_livetime: 3600 +trainer: + logger: + class_path: lightning.pytorch.loggers.CSVLogger + init_args: + name: lightning_logs + flush_logs_every_n_steps: 10 + callbacks: + - class_path: train.callbacks.ModelCheckpoint + init_args: + monitor: "valid_auroc" + mode: "max" + save_top_k: 1 + save_last: true + auto_insert_metric_name: false + - class_path: train.callbacks.SaveAugmentedBatch + # devices: set by the pipeline from train_gpus / train_num_gpus + accelerator: auto + max_epochs: 2 + check_val_every_n_epoch: 1 + log_every_n_steps: 1 + benchmark: false diff --git a/pipeline/profiles/delta/config.yaml b/pipeline/profiles/delta/config.yaml index 4f25e1d50..d8bf3822e 100644 --- a/pipeline/profiles/delta/config.yaml +++ b/pipeline/profiles/delta/config.yaml @@ -30,10 +30,11 @@ local-cores: 4 software-deployment-method: apptainer +# The local repo is bound to /opt/aframe apptainer-args: >- --nv - --bind $AFRAME_DATA_DIRS - ${AFRAME_DEV:+--bind $PWD:/opt/aframe} + --bind ${AFRAME_DATA_DIRS:?must be set, see pipeline/.env.example} + --bind $PWD:/opt/aframe --home $HOME default-resources: diff --git a/pipeline/profiles/ldg/config.yaml b/pipeline/profiles/ldg/config.yaml index fbd6e730d..d936d226a 100644 --- a/pipeline/profiles/ldg/config.yaml +++ b/pipeline/profiles/ldg/config.yaml @@ -50,15 +50,17 @@ local-cores: 16 software-deployment-method: apptainer # $AFRAME_DATA_DIRS binds the run/data/waveform dirs into the container. -# Setting $AFRAME_DEV to any non-empty value binds the $PWD to /opt/aframe. -# Because snakemake should be run from the repo's root directory, this -# is the equivalent of the --dev flag for law. +# If it's unset, the rule fails with a message saying so, rather than +# apptainer misreading the next flag as the bind path. +# The local repo ($PWD, since snakemake runs from the repo root) is bound to +# /opt/aframe, where the images install the projects in editable mode, so +# code always comes from the local repo. # Snakemake invokes apptainer with `--home `, which sets $HOME # to the repo root and creates caches there. A `--home $HOME` overrides it. apptainer-args: >- --nv - --bind $AFRAME_DATA_DIRS - ${AFRAME_DEV:+--bind $PWD:/opt/aframe} + --bind ${AFRAME_DATA_DIRS:?must be set, see pipeline/.env.example} + --bind $PWD:/opt/aframe --home $HOME default-resources: diff --git a/pipeline/resources.smk b/pipeline/resources.smk index 2cc092f44..8c996a407 100644 --- a/pipeline/resources.smk +++ b/pipeline/resources.smk @@ -5,17 +5,65 @@ Slurm uses `mem_mb`, `runtime` and `gpu`, while htcondor ignores those and reads `htcondor_request_mem_mb`, `allowed_execute_duration`, and `request_gpus` plus its GPU matchmaking keys. Each helper returns both slurm and htcondor sets. + +Also resolves each project's container image. """ +import os +import shutil +import subprocess + +from snakemake.logging import logger + +from scripts.env_hash import env_hash + + +def container(project): + """The image to run `project`'s rules in: + `$AFRAME_CONTAINER_ROOT/.sif`. + """ + return os.path.join(os.getenv("AFRAME_CONTAINER_ROOT", ""), f"{project}.sif") + + +def check_images(projects=("data", "train", "export", "infer", "plots")): + """Warn about images built for a different environment than the local + repo's. + + Images record their environment hash (scripts/env_hash.py) at build time. + """ + exe = shutil.which("apptainer") or shutil.which("singularity") + if exe is None: + return + for project in projects: + image = container(project) + if not os.path.exists(image): + continue + result = subprocess.run( + [exe, "exec", image, "cat", "/opt/env_hash"], + capture_output=True, + text=True, + ) + built = result.stdout.strip() if result.returncode == 0 else "unknown" + expected = env_hash(project) + if built != expected: + logger.warning( + f"{image} was built for environment {built}, but the " + f"local repo's is {expected}. Rebuild it with " + f"`uv run build-containers {project}`." + ) + def rule_resources(name): """Memory and walltime for rule `name`, from the config's `resources`. Anything a rule doesn't set there fall back to the profile's - default-resources. + default-resources. With `epnfs`, condor jobs only match execute points + that mount the AP's /home. """ res = config.get("resources", {}).get(name, {}) out = {} + if config.get("epnfs"): + out["requirements"] = "TARGET.EPNFS =?= True" if "mem_mb" in res: out["mem_mb"] = out["htcondor_request_mem_mb"] = res["mem_mb"] if "runtime" in res: diff --git a/projects/data/data.smk b/projects/data/data.smk index 9c2967f28..606e1023c 100644 --- a/projects/data/data.smk +++ b/projects/data/data.smk @@ -29,7 +29,6 @@ split num_validation_signals between them. import json import math -import os from pathlib import Path bg_dir = Path(config["background_dir"]) @@ -40,7 +39,7 @@ train_waveforms = waveform_dir / "train" test_waveforms = waveform_dir / "test" data_log_dir = log_dir / "data" -DATA_CONTAINER = os.path.join(os.getenv("AFRAME_CONTAINER_ROOT", ""), "data.sif") +DATA_CONTAINER = container("data") num_validation_jobs = int(config.get("num_validation_jobs", 200)) validation_branch_ids = [str(i) for i in range(num_validation_jobs)] diff --git a/projects/export/export.smk b/projects/export/export.smk index a43cc2d44..e927a85c8 100644 --- a/projects/export/export.smk +++ b/projects/export/export.smk @@ -4,12 +4,10 @@ Rules: export: compile the trained model into an accelerated format """ -import os - export_out = run_dir / "export" export_log_dir = log_dir / "export" -EXPORT_CONTAINER = os.path.join(os.getenv("AFRAME_CONTAINER_ROOT", ""), "export.sif") +EXPORT_CONTAINER = container("export") remote_train = config.get("remote_train", False) diff --git a/projects/infer/infer.smk b/projects/infer/infer.smk index c3cd96ba9..3388d96c4 100644 --- a/projects/infer/infer.smk +++ b/projects/infer/infer.smk @@ -45,12 +45,22 @@ if ANALYSIS_TYPE not in ("hdf5", "rnp"): if ANALYSIS_TYPE == "rnp" and INFERENCE_MODE == "triton": raise WorkflowError("analysis_type 'rnp' requires inference_mode 'inprocess'") -INFER_CONTAINER = os.path.join(os.getenv("AFRAME_CONTAINER_ROOT", ""), "infer.sif") +INFER_CONTAINER = container("infer") TRITON_IMAGE = os.path.join( os.getenv("AFRAME_CONTAINER_ROOT", ""), config["triton_image"] ) + +def check_triton_image(): + if INFERENCE_MODE == "triton" and not os.path.exists(TRITON_IMAGE): + raise WorkflowError( + f"Triton image {TRITON_IMAGE} doesn't exist. Pull the release " + f"in its name, e.g. `apptainer pull {TRITON_IMAGE} " + "docker://nvcr.io/nvidia/tritonserver:25.06-py3`." + ) + + AOTI_PKG = str(export_out / "model_aoti.pt2") BRANCHES_PER_JOB = config.get("branches_per_job", 1) diff --git a/projects/online/config.yaml b/projects/online/config.yaml index 7c7054b02..acecb9b29 100644 --- a/projects/online/config.yaml +++ b/projects/online/config.yaml @@ -19,9 +19,9 @@ amplfi_hl_architecture: time_layers: [5, 3, 3] freq_layers: [5, 3, 3] norm_layer: - class_path: ml4gw.nn.norm.GroupNorm1DGetter + class_path: ml4gw.nn.norm.GroupNorm1D init_args: - groups: 8 + num_groups: 8 amplfi_hlv_architecture: class_path: amplfi.train.architectures.flows.NSF init_args: @@ -37,9 +37,9 @@ amplfi_hlv_architecture: time_layers: [5, 3, 3] freq_layers: [5, 3, 3] norm_layer: - class_path: ml4gw.nn.norm.GroupNorm1DGetter + class_path: ml4gw.nn.norm.GroupNorm1D init_args: - groups: 8 + num_groups: 8 amplfi_hl_weights: ${oc.env:AMPLFI_HL_WEIGHTS} amplfi_hlv_weights: ${oc.env:AMPLFI_HLV_WEIGHTS} amplfi_parameter_sampler: ./prior.yaml diff --git a/projects/online/pyproject.toml b/projects/online/pyproject.toml index a918486f9..96e1ebca0 100644 --- a/projects/online/pyproject.toml +++ b/projects/online/pyproject.toml @@ -16,7 +16,7 @@ dependencies = [ "architectures", "arrakis>=0.20,<0.21", "amplfi", - "ml4gw>=0.7.4", + "ml4gw>=0.8.4", "torch==2.10.0", "omegaconf>=2.3.0,<3", "numpy>=2", diff --git a/projects/plots/plots.smk b/projects/plots/plots.smk index 8e4815dac..2fecb739d 100644 --- a/projects/plots/plots.smk +++ b/projects/plots/plots.smk @@ -10,12 +10,11 @@ inputs and needs no network access. """ import json -import os plots_dir = run_dir / "plots" plots_log_dir = log_dir / "plots" -PLOTS_CONTAINER = os.path.join(os.getenv("AFRAME_CONTAINER_ROOT", ""), "plots.sif") +PLOTS_CONTAINER = container("plots") VETOS = config.get("vetos") diff --git a/projects/train/configs/bns.yaml b/projects/train/configs/bns.yaml index f740e3089..cfab56c1d 100644 --- a/projects/train/configs/bns.yaml +++ b/projects/train/configs/bns.yaml @@ -14,9 +14,9 @@ model: init_args: layers: [3, 4, 6, 3] norm_layer: - class_path: ml4gw.nn.norm.GroupNorm2DGetter + class_path: ml4gw.nn.norm.GroupNorm2D init_args: - groups: 16 + num_groups: 16 metric: class_path: train.metrics.TimeSlideAUROC init_args: diff --git a/projects/train/configs/bns_time_spectrogram.yaml b/projects/train/configs/bns_time_spectrogram.yaml index a73287b97..2a632f870 100644 --- a/projects/train/configs/bns_time_spectrogram.yaml +++ b/projects/train/configs/bns_time_spectrogram.yaml @@ -23,9 +23,9 @@ model: spec_kernel_size: 3 spec_classes: 1 spec_norm_layer: - class_path: ml4gw.nn.norm.GroupNorm2DGetter + class_path: ml4gw.nn.norm.GroupNorm2D init_args: - groups: 16 + num_groups: 16 train_X_coeff: 0.3 train_X_spec_coeff: 0.7 val_X_coeff: 0.5 diff --git a/projects/train/pyproject.toml b/projects/train/pyproject.toml index c48755de6..9c527c816 100644 --- a/projects/train/pyproject.toml +++ b/projects/train/pyproject.toml @@ -23,7 +23,7 @@ dependencies = [ "fsspec[s3]>=2024", "s3fs>=2024", "utils[torch]", - "ml4gw>=0.8.0", + "ml4gw>=0.8.4", "ledger", "architectures", "filelock>=3.13.1,<5", diff --git a/projects/train/train.smk b/projects/train/train.smk index 4eadd4a55..1d56d2b3b 100644 --- a/projects/train/train.smk +++ b/projects/train/train.smk @@ -15,12 +15,10 @@ The train_config YAML defines all the model/data/trainer hyperparameters. This rule adds the preprocessing args shared with inference. """ -import os - train_out = run_dir / "train" train_log_dir = log_dir / "train" -TRAIN_CONTAINER = os.path.join(os.getenv("AFRAME_CONTAINER_ROOT", ""), "train.sif") +TRAIN_CONTAINER = container("train") # GPUs for local training. `train_gpus` pins specific devices on a shared diff --git a/pyproject.toml b/pyproject.toml index ef3500a9f..ce961f874 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,13 +78,12 @@ packages = ["scripts"] [tool.ruff] line-length = 79 -extend-include = ["T201"] extend-exclude = ["docs"] [tool.ruff.lint] extend-safe-fixes = ["C408"] mccabe = {"max-complexity" = 18} -select = ["B", "C", "E", "F", "W", "I", "UP", "B9"] +select = ["B", "C", "E", "F", "W", "I", "UP", "B9", "T201"] # ignore W605 and E203 for consistency with black # ignore F722 for compatibility with jaxtyping ignore = ["W605", "E203", "F722"] diff --git a/scripts/aframe_init.py b/scripts/aframe_init.py index afa8c77b1..428537729 100644 --- a/scripts/aframe_init.py +++ b/scripts/aframe_init.py @@ -45,7 +45,6 @@ def create_snakemake_runfile(path: Path, profile: str): cmd = f"snakemake --configfile {config} --profile {profile}" content = f""" #!/bin/bash - # Set AFRAME_DEV=1 to bind the working tree into containers. cd {root} source pipeline/.env {cmd} diff --git a/scripts/build_containers.py b/scripts/build_containers.py index 7374d58f0..e9b699e6b 100644 --- a/scripts/build_containers.py +++ b/scripts/build_containers.py @@ -1,12 +1,13 @@ import logging import os -import tomllib from concurrent.futures import ProcessPoolExecutor, as_completed from pathlib import Path from jsonargparse import ArgumentParser from spython.main import Client +from scripts.env_hash import env_hash, local_libs + # Define the directory where the projects are located ROOT_DIR: Path = Path(__file__).resolve().parent.parent BASE_DIR: Path = ROOT_DIR / "projects" @@ -39,32 +40,6 @@ } -def _local_libs(project_name: str) -> list[str]: - """ - Collect a list of the `libs/` that a project depends on, including - those that enter via a library's own dependencies. - """ - seen: set[str] = set() - pyproject_files: list[Path] = [BASE_DIR / project_name / "pyproject.toml"] - - while len(pyproject_files) > 0: - with open(pyproject_files.pop(), "rb") as f: - data = tomllib.load(f) - - sources = data.get("tool", {}).get("uv", {}).get("sources", {}) - for name, source in sources.items(): - # index pins like torch are lists, not dicts. - # git sources are dicts but don't have a `workspace` key - if not isinstance(source, dict) or not source.get("workspace"): - continue - if name in seen: - continue - seen.add(name) - pyproject_files.append(LIBS_DIR / name / "pyproject.toml") - - return sorted(seen) - - def _copy_entries( source_prefix: str, directory: Path, destination: str ) -> list[str]: @@ -90,7 +65,7 @@ def _get_files_block(project_name: str) -> str: BASE_DIR / project_name, f"/opt/aframe/projects/{project_name}", ) - for lib in _local_libs(project_name): + for lib in local_libs(project_name): lines.extend( _copy_entries( f"../../libs/{lib}/", @@ -154,6 +129,7 @@ def create_definition_file(project_name: str) -> Path: .replace("@@EXTRA_POST@@", extra_post) .replace("@@EXTRA_ENV@@", extra_env) .replace("@@PURGE_BUILD_TOOLS@@", PURGE_BUILD_TOOLS) + .replace("@@ENV_HASH@@", env_hash(project_name)) ) output_path = project_dir / "apptainer.def" diff --git a/scripts/env_hash.py b/scripts/env_hash.py new file mode 100644 index 000000000..abc7b2949 --- /dev/null +++ b/scripts/env_hash.py @@ -0,0 +1,84 @@ +"""Hash the inputs that determine a project's container environment. + +Project and library sources are installed in editable mode, so an +image only has to be updated when the environment changes. Images +record this hash at build time and the pipeline compares it with +the local repo's to warn about stale images. +""" + +import hashlib +import sys +import tomllib +from pathlib import Path + +ROOT_DIR: Path = Path(__file__).resolve().parent.parent +PROJECTS_DIR: Path = ROOT_DIR / "projects" +LIBS_DIR: Path = ROOT_DIR / "libs" +TEMPLATES_DIR: Path = ROOT_DIR / "container_templates" + + +def local_libs(project: str) -> list[str]: + """ + The `libs/` that a project depends on, including those that enter via + a library's own dependencies. + """ + seen: set[str] = set() + pyproject_files: list[Path] = [PROJECTS_DIR / project / "pyproject.toml"] + + while len(pyproject_files) > 0: + with open(pyproject_files.pop(), "rb") as f: + data = tomllib.load(f) + + sources = data.get("tool", {}).get("uv", {}).get("sources", {}) + for name, source in sources.items(): + # index pins like torch are lists, not dicts. + # git sources are dicts but don't have a `workspace` key + if not isinstance(source, dict) or not source.get("workspace"): + continue + if name in seen: + continue + seen.add(name) + pyproject_files.append(LIBS_DIR / name / "pyproject.toml") + + return sorted(seen) + + +def env_files(project: str) -> list[Path]: + """Every file whose contents impact the project's environment.""" + project_dir = PROJECTS_DIR / project + conda_lock = project_dir / f"{project}.conda-lock.yml" + template = "micromamba.def" if conda_lock.exists() else "uv.def" + candidates = [ + ROOT_DIR / "uv.lock", + ROOT_DIR / "pyproject.toml", + ROOT_DIR / "scripts" / "build_containers.py", + TEMPLATES_DIR / template, + project_dir / "pyproject.toml", + project_dir / "apptainer.post", + project_dir / "apptainer.env", + conda_lock, + *(LIBS_DIR / lib / "pyproject.toml" for lib in local_libs(project)), + ] + return [path for path in candidates if path.exists()] + + +def manifest(project: str) -> str: + """`sha256sum`-style lines for each environment file, sorted by path.""" + paths = sorted(str(p.relative_to(ROOT_DIR)) for p in env_files(project)) + return "".join( + f"{hashlib.sha256((ROOT_DIR / p).read_bytes()).hexdigest()} {p}\n" + for p in paths + ) + + +def env_hash(project: str) -> str: + """A short hash of the files that impact the project's environment. + + Hashes the manifest, so from the repo root this equals the first 12 + characters of `sha256sum | sha256sum`. + """ + return hashlib.sha256(manifest(project).encode()).hexdigest()[:12] + + +if __name__ == "__main__": + print(env_hash(sys.argv[1])) # noqa: T201 diff --git a/uv.lock b/uv.lock index 745d5beaa..36256a3ef 100644 --- a/uv.lock +++ b/uv.lock @@ -448,7 +448,7 @@ requires-dist = [ { name = "einops", specifier = ">=0.8,<0.9" }, { name = "h5py", specifier = ">=3.10,<4" }, { name = "jaxtyping", specifier = ">=0.2" }, - { name = "ml4gw", specifier = ">=0.7.2" }, + { name = "ml4gw", specifier = ">=0.8.4" }, { name = "numpy", specifier = ">=2" }, { name = "torch" }, ] @@ -3442,7 +3442,7 @@ wheels = [ [[package]] name = "ml4gw" -version = "0.8.0" +version = "0.8.4" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jaxtyping" }, @@ -3455,9 +3455,9 @@ dependencies = [ { name = "torchaudio", version = "2.10.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-data-cpu' and extra == 'extra-4-data-cuda') or (extra != 'extra-4-data-cpu' and extra == 'extra-5-plots-vizapp-cuda') or (extra != 'extra-4-data-cpu' and extra != 'extra-5-plots-vizapp') or (extra != 'extra-4-data-cuda' and extra == 'extra-5-plots-vizapp' and extra == 'extra-5-plots-vizapp-cuda')" }, { name = "torchaudio", version = "2.10.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-data-cpu') or (sys_platform != 'darwin' and extra == 'extra-5-plots-vizapp') or (extra == 'extra-4-data-cpu' and extra == 'extra-4-data-cuda') or (extra == 'extra-5-plots-vizapp' and extra == 'extra-5-plots-vizapp-cuda')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/2a/56/abb490d353f989802f918ee92cf6c9a37336483aa20c7f17df7730d81744/ml4gw-0.8.0.tar.gz", hash = "sha256:43a2411ae348f8f911fdc0e2defd4fa54370414fa8b51c63518de3cb805754ba", size = 121709, upload-time = "2026-04-17T13:15:20.347Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/51/8a4f2c8061fb1dfa9a8bc6cbb3d0f79ddfa10121eb8e1f6ee910ba29ed7e/ml4gw-0.8.4.tar.gz", hash = "sha256:3a801759d38a4aaab64030a810505a72ba73acfb0a5f286b2bda534cc23e47a8", size = 131414, upload-time = "2026-09-25T13:53:41.624Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/22/64102f10ad7f9043083d8bafcf84b8f2e8abc84bfc214fd04dd8d243ece9/ml4gw-0.8.0-py3-none-any.whl", hash = "sha256:0b4377541d5a90dcf9c728efa4008b4d57b942452acf564856ac58f599273070", size = 132926, upload-time = "2026-04-17T13:15:18.751Z" }, + { url = "https://files.pythonhosted.org/packages/ce/ce/f65f38f5db0a37160c1e171a0471b61727709e4008f69879991d41026239/ml4gw-0.8.4-py3-none-any.whl", hash = "sha256:bac66cbed02b2082ad1c2d845a47bedf67cfb8cce0928d2157658e8788eb7bb3", size = 143589, upload-time = "2026-09-25T13:53:40.554Z" }, ] [[package]] @@ -4123,7 +4123,7 @@ requires-dist = [ { name = "ligo-gracedb", extras = ["kafka"], specifier = ">=2.15.4" }, { name = "ligo-skymap", specifier = ">=2.4.0,<3" }, { name = "matplotlib", specifier = ">=3.9.4" }, - { name = "ml4gw", specifier = ">=0.7.4" }, + { name = "ml4gw", specifier = ">=0.8.4" }, { name = "numpy", specifier = ">=2" }, { name = "omegaconf", specifier = ">=2.3.0,<3" }, { name = "p-astro", editable = "libs/p_astro" }, @@ -6811,7 +6811,7 @@ requires-dist = [ { name = "ledger", editable = "libs/ledger" }, { name = "lightning", specifier = ">=2.2.1" }, { name = "lightray", specifier = ">=0.2.3" }, - { name = "ml4gw", specifier = ">=0.8.0" }, + { name = "ml4gw", specifier = ">=0.8.4" }, { name = "numpy", specifier = ">=2" }, { name = "ray", extras = ["default", "tune"], specifier = ">=2.8.0,<3" }, { name = "s3fs", specifier = ">=2024" },