Training data pipeline: disk-backed memmap rounds, background producer, batched tf.data input

.claude/skills/aetherscan-repo-context/SKILL.md

@@ -114,7 +114,10 @@ utils/                   # fetch_run_outputs.sh, find_optimal_configs.py,
                          # run_container.sh, start_tmux_session.sh,
                          # verify_train_test_files.py
 docs/                    # GPU_RUNTIME_GUIDE.md, CONFIG_AND_CLI.md, README.md, assets/
-tests/                   # Placeholder — no test suite yet
+tests/                   # Pytest suite: conftest.py (singleton-reset fixtures, synthetic
+                         # data factories), unit/, integration/ (gpu+cluster-marked smokes).
+                         # Default selection: pytest -m "not gpu and not cluster" -q
+                         # (runs in CI via .github/workflows/tests.yml); see CONTRIBUTING.md
 ```
 
 ---

.github/workflows/tests.yml

@@ -0,0 +1,57 @@
+# Unit test suite (pytest)
+name: Tests
+
+# Run on every PR and on pushes to master
+on:
+  push:
+    branches: ["master"]
+  pull_request:
+    branches: ["**"]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  tests:
+    runs-on: ubuntu-latest
+    strategy:
+      # Let both Python versions report rather than cancelling the sibling on first failure
+      fail-fast: false
+      matrix:
+        # The two supported runtimes: conda env (3.10) and NGC container (3.12)
+        python-version: ["3.10", "3.12"]
+
+    steps:
+      # Pull repo contents into GitHub Actions runner
+      - name: Checkout code
+        uses: actions/checkout@v5
+
+      # Install the matrix Python version, with built-in pip caching keyed on the
+      # dependency pins (invalidates when requirements-container.txt changes)
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: "pip"
+          # Include the workflow itself so the inline pins below also invalidate the cache
+          cache-dependency-path: |
+            requirements-container.txt
+            .github/workflows/tests.yml
+
+      # Mirror the NGC container's dependency surface on a CPU-only runner:
+      # tensorflow-cpu stands in for the container's GPU TF 2.17 build, the pinned
+      # extras come from requirements-container.txt, and h5py/hdf5plugin/pandas/
+      # psutil/pytest are installed explicitly because the NGC base image ships
+      # them (so the requirements file intentionally omits them; hdf5plugin is
+      # pinned explicitly in environment.yml too — don't rely on setigen pulling
+      # it transitively)
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install "tensorflow-cpu==2.17.*" -r requirements-container.txt h5py hdf5plugin pandas psutil pytest
+
+      # GPU- and cluster-marked tests need physical GPUs / cluster-resident data;
+      # everything else (including slow CPU TF graph tests) runs here
+      - name: Run test suite
+        run: pytest -m "not gpu and not cluster" -q

CLAUDE.md

@@ -4,7 +4,7 @@ Always-on, bare-essentials rules for any coding agent in this repo. The full dee
 
 ## Project
 
-Aetherscan: Breakthrough Listen's deep-learning SETI pipeline. Two-stage ML (Beta-VAE → Random Forest), single-node multi-GPU distributed train/inference. Sole entry point: `src/aetherscan/main.py`. `tests/` is a placeholder — **no test suite yet**, so don't claim `pytest` passes; verify changes another way and say how.
+Aetherscan: Breakthrough Listen's deep-learning SETI pipeline. Two-stage ML (Beta-VAE → Random Forest), single-node multi-GPU distributed train/inference. Sole entry point: `src/aetherscan/main.py`. Pytest suite lives in `tests/` (unit + gpu/cluster-marked integration; see [CONTRIBUTING.md](CONTRIBUTING.md#testing)) — run it before claiming a change works, and ship unit tests with new logic. Most test modules import TensorFlow at collection; if the full dependency stack isn't available in your environment, run the subset you can and say exactly what you ran.
 
 ## Run / lint
 
@@ -15,6 +15,8 @@ Aetherscan: Breakthrough Listen's deep-learning SETI pipeline. Two-stage ML (Bet
 PYTHONPATH=src python -m aetherscan.main {train|inference} --save-tag final_v1
 # Lint + format (also enforced via pre-commit)
 ruff check src/ && ruff format src/
+# Tests (default selection = what CI runs; gpu/cluster-marked tests need a cluster)
+pytest -m "not gpu and not cluster" -q
 ```
 
 ## Hard rules (don't break these)

CONTRIBUTING.md

@@ -190,7 +190,10 @@ Aetherscan/
 │       ├── __init__.py         # Manager exports
 │       └── manager.py          # Resource lifecycle management
 ├── docs/                       # Documentation (placeholder; no docs yet)
-├── tests/                      # Test suite (placeholder; no tests yet)
+├── tests/                      # Pytest suite (see Testing section below)
+│   ├── conftest.py                  # Singleton-reset fixtures, tmp paths, synthetic data factories
+│   ├── unit/                        # Fast, hardware-independent unit tests (run in CI)
+│   └── integration/                 # gpu+cluster-marked end-to-end smoke tests
 ├── utils/                      # Utility scripts
 │   ├── fetch_run_outputs.sh         # rsync a run's outputs from a remote node
 │   ├── find_optimal_configs.py      # Per-host config helper
@@ -510,9 +513,43 @@ Prefer descriptive, full-word names: `num_training_rounds`, `per_replica_batch_s
 
 ## Testing
 
-> [!WARNING]
->
-> # TODO: update this section once test suite is operational
+The test suite lives in `tests/` and runs on [pytest](https://docs.pytest.org/) (configured in `pyproject.toml` under `[tool.pytest.ini_options]` — `pythonpath = ["src"]` makes `import aetherscan` work without a `PYTHONPATH` prefix).
+
+### Running the suite
+
+```bash
+# Default selection — everything that doesn't need GPUs or cluster-resident data.
+# This is exactly what CI runs (.github/workflows/tests.yml).
+pytest -m "not gpu and not cluster" -q
+
+# Cluster integration smokes (end-to-end train + CSV inference), inside the NGC container:
+./utils/run_container.sh python -m pytest tests/ -m "gpu or cluster" -q
+
+# A single file / test while iterating:
+pytest tests/unit/test_config.py -q
+pytest tests/unit/test_db.py::TestFlushSentinel -q
+```
+
+### Markers
+
+| Marker        | Meaning                                                                                     |
+| ------------- | ------------------------------------------------------------------------------------------- |
+| `slow`        | Slower tests (e.g. builds TF graphs on CPU); still included in the default CI selection      |
+| `gpu`         | Requires one or more physical GPUs; excluded from CI                                         |
+| `cluster`     | Requires cluster-resident data/models (blpc3/bla0); excluded from CI                         |
+| `integration` | End-to-end pipeline runs; skips the per-test singleton/env isolation fixture in `conftest.py` |
+
+### Test environment & fixtures
+
+- Every non-integration test runs isolated: `tests/conftest.py` points `AETHERSCAN_{DATA,MODEL,OUTPUT}_PATH` at a pytest `tmp_path`, resets all singletons (`Config`, `ResourceManager`, `Database`, `Logger`, `ResourceMonitor`) via their `_reset()` hooks, and re-runs `init_config()` — so tests never touch real data or leak state into each other.
+- Synthetic data factories are provided as fixtures: `make_background_npy` (tiny background plates), `make_h5_observation` (tiny filterbank-style `.h5` files), and `make_inference_csv` (tiny cadence-grouping CSVs).
+- CI (`.github/workflows/tests.yml`) runs the default selection on Ubuntu with `tensorflow-cpu==2.17.*` + the pins from `requirements-container.txt`, on Python 3.10 and 3.12 (the conda and container runtimes respectively). A few tests assert Linux-only behavior (e.g. PSS-based memory accounting) and self-skip elsewhere.
+
+### Writing tests
+
+- Every PR that adds or changes logic should ship unit tests for it (`tests/unit/`), reusing the conftest fixtures rather than instantiating singletons directly.
+- Mark anything that needs hardware or cluster data with the appropriate marker so the default selection stays runnable on a laptop and in CI.
+- Test style follows the same ruff lint+format rules as `src/` (pre-commit runs on `tests/` too).
 
 ---
 

README.md

@@ -417,6 +417,10 @@ usage: train [-h] [--data-path DATA_PATH] [--model-path MODEL_PATH]
              [--effective-batch-size EFFECTIVE_BATCH_SIZE]
              [--per-replica-val-batch-size PER_REPLICA_VAL_BATCH_SIZE]
              [--signal-injection-chunk-size SIGNAL_INJECTION_CHUNK_SIZE]
+             [--data-gen-task-size DATA_GEN_TASK_SIZE]
+             [--round-data-dir ROUND_DATA_DIR]
+             [--overlap-data-generation | --no-overlap-data-generation]
+             [--keep-round-data | --no-keep-round-data]
              [--plot-injection-subsampling-count PLOT_INJECTION_SUBSAMPLING_COUNT]
              [--plot-injection-outlier-percentile PLOT_INJECTION_OUTLIER_PERCENTILE]
              [--latent-viz-num-cadences-per-type LATENT_VIZ_NUM_CADENCES_PER_TYPE]
@@ -551,6 +555,25 @@ options:
   --signal-injection-chunk-size SIGNAL_INJECTION_CHUNK_SIZE
                         Maximum cadences to process at once during synthetic
                         signal injection (must be divisible by 4)
+  --data-gen-task-size DATA_GEN_TASK_SIZE
+                        Cadences per batched signal-injection worker task
+                        (workers write results straight into the round's on-
+                        disk memmap; must be >= 1)
+  --round-data-dir ROUND_DATA_DIR
+                        Directory for disk-backed per-round training datasets
+                        (defaults to <output-path>/round_data; needs ~2.2x one
+                        round's size free when data-generation overlap is
+                        enabled, ~1.1x otherwise)
+  --overlap-data-generation, --no-overlap-data-generation
+                        Generate round k+1's training data in a background
+                        producer process while round k trains (default:
+                        enabled). Pass --no-overlap-data-generation to fall
+                        back to sequential in-process generation for debugging
+  --keep-round-data, --no-keep-round-data
+                        Retain each round's on-disk training data after that
+                        round finishes (default: disabled — round k's data
+                        directory is deleted as soon as round k's training
+                        completes). Enable for debugging
   --plot-injection-subsampling-count PLOT_INJECTION_SUBSAMPLING_COUNT
                         Max points per stat name, per signal type, for A→B
                         intensity bias scatter plots. Outliers are

pyproject.toml

@@ -55,6 +55,19 @@ Issues = "https://github.com/zachtheyek/Aetherscan/issues"
 [project.optional-dependencies]
 dev = ["ruff"]  # Python linter
 
+[tool.pytest.ini_options]
+# Make `import aetherscan` work without a PYTHONPATH=src prefix
+pythonpath = ["src"]
+testpaths = ["tests"]
+# Reject undeclared markers so a typo'd @pytest.mark.gpu can't silently slip into CI
+addopts = "--strict-markers"
+markers = [
+    "slow: slower tests (e.g. builds TF graphs on CPU); included in the default CI selection",
+    "gpu: requires one or more physical GPUs; excluded from CI (-m 'not gpu and not cluster')",
+    "cluster: requires cluster-resident data/models (blpc3/bla0); excluded from CI",
+    "integration: end-to-end pipeline runs (skips the per-test singleton/env isolation fixture)",
+]
+
 [tool.ruff]
 # Configure linter
 target-version = "py310"   # Target Python version

src/aetherscan/cli.py

@@ -9,6 +9,7 @@
 import logging
 import os
 import re
+import shutil
 from dataclasses import dataclass, field, is_dataclass
 from itertools import product
 from typing import Any
@@ -55,6 +56,32 @@ def _resolve(args: argparse.Namespace, arg_name: str, default: Any) -> Any:
     return val if val is not None else default
 
 
+def _estimate_round_data_nbytes(
+    n_samples: int, num_observations: int, time_bins: int, width_bin_downsampled: int
+) -> int:
+    """
+    Estimate the on-disk size of one round's disk-backed dataset (see round_data.py): three
+    float32 arrays of shape (n_samples, num_observations, time_bins, width_bin_downsampled)
+    plus a tiny U20 labels array. Kept stdlib-only (no numpy) so utils/print_cli_help.py can
+    keep importing cli.py without the scientific stack.
+    """
+    per_sample_bytes = num_observations * time_bins * width_bin_downsampled * 4  # float32
+    labels_bytes = n_samples * 20 * 4  # numpy "U20" = 20 UCS-4 code points per label
+    return 3 * n_samples * per_sample_bytes + labels_bytes
+
+
+def _nearest_existing_ancestor(path: str) -> str:
+    """Walk up from `path` to the closest directory that exists (for shutil.disk_usage on a
+    round-data dir that hasn't been created yet)."""
+    probe = os.path.abspath(path)
+    while probe and not os.path.exists(probe):
+        parent = os.path.dirname(probe)
+        if parent == probe:
+            break
+        probe = parent
+    return probe
+
+
 def _resolve_num_replicas(args: argparse.Namespace) -> int | None:
     """Resolve the replica count to validate cross-replica constraints against, and
     fail fast if an explicitly-requested count is unusable on this host.
@@ -394,6 +421,30 @@ def _add_train_flags_to(parser):
         # NOTE: divisible by 4 or num_replicas?
         help="Maximum cadences to process at once during synthetic signal injection (must be divisible by 4)",
     )
+    parser.add_argument(
+        "--data-gen-task-size",
+        type=int,
+        default=None,
+        help="Cadences per batched signal-injection worker task (workers write results straight into the round's on-disk memmap; must be >= 1)",
+    )
+    parser.add_argument(
+        "--round-data-dir",
+        type=str,
+        default=None,
+        help="Directory for disk-backed per-round training datasets (defaults to <output-path>/round_data; needs ~2.2x one round's size free when data-generation overlap is enabled, ~1.1x otherwise)",
+    )
+    parser.add_argument(
+        "--overlap-data-generation",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+        help="Generate round k+1's training data in a background producer process while round k trains (default: enabled). Pass --no-overlap-data-generation to fall back to sequential in-process generation for debugging",
+    )
+    parser.add_argument(
+        "--keep-round-data",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+        help="Retain each round's on-disk training data after that round finishes (default: disabled — round k's data directory is deleted as soon as round k's training completes). Enable for debugging",
+    )
     parser.add_argument(
         "--plot-injection-subsampling-count",
         type=int,
@@ -920,6 +971,18 @@ def apply_args_to_config(args: argparse.Namespace) -> None:
         and args.signal_injection_chunk_size is not None
     ):
         config.training.signal_injection_chunk_size = args.signal_injection_chunk_size
+    if hasattr(args, "data_gen_task_size") and args.data_gen_task_size is not None:
+        config.training.data_gen_task_size = args.data_gen_task_size
+    if hasattr(args, "round_data_dir") and args.round_data_dir is not None:
+        config.training.round_data_dir = args.round_data_dir
+    # overlap_data_generation / keep_round_data use argparse.BooleanOptionalAction with
+    # default=None so that the CLI can express "leave the config default" (omit), "force on"
+    # (--overlap-data-generation), and "force off" (--no-overlap-data-generation). The
+    # `is not None` guard preserves the config default when the user passes neither
+    if hasattr(args, "overlap_data_generation") and args.overlap_data_generation is not None:
+        config.training.overlap_data_generation = args.overlap_data_generation
+    if hasattr(args, "keep_round_data") and args.keep_round_data is not None:
+        config.training.keep_round_data = args.keep_round_data
     if (
         hasattr(args, "plot_injection_subsampling_count")
         and args.plot_injection_subsampling_count is not None
@@ -1216,7 +1279,7 @@ def collect_validation_errors(
                 ValidationError(
                     field="training.num_samples_rf",
                     current=nsr,
-                    message=f"--num-samples-rf must be divisible by 2 for generate_triplet_batch, got {nsr}",
+                    message=f"--num-samples-rf must be divisible by 2 for the balanced true/false halves in data generation, got {nsr}",
                     fix_kind="divisibility",
                     divisor=2,
                 )
@@ -1247,6 +1310,62 @@ def collect_validation_errors(
                 )
             )
 
+        # data_gen_task_size >= 1 (batched memmap worker tasks)
+        dgts = _resolve(args, "data_gen_task_size", config.training.data_gen_task_size)
+        if dgts is not None and dgts < 1:
+            errors.append(
+                ValidationError(
+                    field="training.data_gen_task_size",
+                    current=dgts,
+                    message=f"--data-gen-task-size must be >= 1, got {dgts}",
+                    fix_kind="clamp_low",
+                    min_val=1,
+                )
+            )
+
+        # Disk budget for the disk-backed round datasets (round_data.py): with overlap enabled
+        # two rounds coexist on disk (round k trains while round k+1 generates), so require
+        # 2.2x one round's estimated size free; 1.1x when overlap is disabled. Estimated from
+        # sample counts only — actual usage tracks the estimate closely since the arrays are
+        # fixed-shape float32.
+        nob = _resolve(args, "num_observations", config.data.num_observations)
+        tb = _resolve(args, "time_bins", config.data.time_bins)
+        overlap = _resolve(args, "overlap_data_generation", config.training.overlap_data_generation)
+        output_path = _resolve(args, "output_path", config.output_path)
+        round_data_dir = _resolve(args, "round_data_dir", config.training.round_data_dir)
+        if round_data_dir is None:
+            round_data_dir = os.path.join(output_path, "round_data")
+        if all(v is not None for v in (nsb, nob, tb, wb)) and df:
+            round_nbytes = _estimate_round_data_nbytes(nsb, nob, tb, wb // df)
+            required_factor = 2.2 if overlap else 1.1
+            required_bytes = required_factor * round_nbytes
+            try:
+                free_bytes = shutil.disk_usage(_nearest_existing_ancestor(round_data_dir)).free
+            except OSError as e:
+                logger.warning(
+                    f"Could not check free disk space for --round-data-dir "
+                    f"({round_data_dir}): {e} — skipping the disk-budget check"
+                )
+            else:
+                if free_bytes < required_bytes:
+                    errors.append(
+                        ValidationError(
+                            field="training.round_data_dir",
+                            current=round_data_dir,
+                            message=(
+                                f"--round-data-dir ({round_data_dir}) needs >= "
+                                f"{required_bytes / 1e9:.1f} GB free ({required_factor}x one "
+                                f"round's ~{round_nbytes / 1e9:.1f} GB"
+                                f"{' with data-generation overlap enabled' if overlap else ''}), "
+                                f"but only {free_bytes / 1e9:.1f} GB is available. Free up disk "
+                                f"space, point --round-data-dir at a larger volume, reduce "
+                                f"--num-samples-beta-vae, or pass --no-overlap-data-generation "
+                                f"to halve the requirement"
+                            ),
+                            fix_kind="file_exists",
+                        )
+                    )
+
         # SNR sanity (positivity + curriculum ordering)
         snr_base = _resolve(args, "snr_base", config.training.snr_base)
         snr_init = _resolve(args, "initial_snr_range", config.training.initial_snr_range)