Image improvements

.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'
 

.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

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:

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

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

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

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

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(

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",

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":

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)

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

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:<release>-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

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

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

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

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:

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 <workdir>`, 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: