diff --git a/.claude/skills/aetherscan-repo-context/SKILL.md b/.claude/skills/aetherscan-repo-context/SKILL.md index 76d6e80c..0c8fb950 100644 --- a/.claude/skills/aetherscan-repo-context/SKILL.md +++ b/.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 ``` --- diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 00000000..5c6842db --- /dev/null +++ b/.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 diff --git a/CLAUDE.md b/CLAUDE.md index 2f651694..fd01916d 100644 --- a/CLAUDE.md +++ b/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) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b3d6e38b..4ae9c149 100644 --- a/CONTRIBUTING.md +++ b/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). --- diff --git a/README.md b/README.md index 4a410837..3012261a 100644 --- a/README.md +++ b/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 /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 diff --git a/pyproject.toml b/pyproject.toml index 109eb90f..386c214d 100644 --- a/pyproject.toml +++ b/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 diff --git a/src/aetherscan/cli.py b/src/aetherscan/cli.py index 03cc30da..89a394f3 100644 --- a/src/aetherscan/cli.py +++ b/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 /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) diff --git a/src/aetherscan/config.py b/src/aetherscan/config.py index 0e686a20..fe4c3102 100644 --- a/src/aetherscan/config.py +++ b/src/aetherscan/config.py @@ -183,6 +183,16 @@ class TrainingConfig: signal_injection_chunk_size: int = ( 50000 # Maximum cadences to process at once during data generation ) + data_gen_task_size: int = 256 # Cadences per batched worker task (workers write results straight into the round memmap) + + # Round data pipeline params (disk-backed per-round datasets, see round_data.py) + round_data_dir: str | None = None # Defaults to /round_data at runtime + overlap_data_generation: bool = ( + True # Generate round k+1's data in a background process while round k trains + ) + keep_round_data: bool = ( + False # Retain each round's on-disk data after its training completes (debugging) + ) # TODO: tune to include sufficient info without bottlenecking training plot_injection_subsampling_count: int = ( 100000 # Max points per series in injection_stats scatter plots @@ -526,6 +536,10 @@ def to_dict(self) -> dict: "effective_batch_size": self.training.effective_batch_size, "per_replica_val_batch_size": self.training.per_replica_val_batch_size, "signal_injection_chunk_size": self.training.signal_injection_chunk_size, + "data_gen_task_size": self.training.data_gen_task_size, + "round_data_dir": self.training.round_data_dir, + "overlap_data_generation": self.training.overlap_data_generation, + "keep_round_data": self.training.keep_round_data, "plot_injection_subsampling_count": self.training.plot_injection_subsampling_count, "plot_injection_outlier_percentile": self.training.plot_injection_outlier_percentile, "latent_viz_num_cadences_per_type": self.training.latent_viz_num_cadences_per_type, diff --git a/src/aetherscan/data_generation.py b/src/aetherscan/data_generation.py index 241c622c..324564b7 100644 --- a/src/aetherscan/data_generation.py +++ b/src/aetherscan/data_generation.py @@ -2,6 +2,11 @@ Synthetic data generation for Aetherscan Pipeline Handles signal injection & log-normalization for training Uses multiprocessing and shared memory to process data in parallel + +Generated rounds are written straight into disk-backed .npy memmaps (see +aetherscan.round_data): each pool task generates a batch of cadences and writes them +into its slice of the round's memmap in-place, returning only the small per-sample +stats dicts — no more per-sample IPC pickling or ~294 GB in-RAM round arrays. """ from __future__ import annotations @@ -11,9 +16,11 @@ import logging import os import random +import shutil import signal import time -from multiprocessing import Pool, cpu_count +from dataclasses import dataclass +from multiprocessing import Pool from multiprocessing.shared_memory import SharedMemory import numpy as np @@ -25,11 +32,11 @@ from aetherscan.db import get_db from aetherscan.logger import init_worker_logging from aetherscan.manager import get_manager +from aetherscan.round_data import RoundDataPaths, build_manifest, write_done_manifest logger = logging.getLogger(__name__) # NOTE: find a way to avoid using global refs (store under manager.py maybe?) -# NOTE: is there any room to use asyncio & load all chunks simultaneously? # Global variables to store background data for multiprocessing workers # This avoids serialization overhead when passing data between workers _GLOBAL_SHM = None @@ -37,47 +44,25 @@ _GLOBAL_SHAPE = None _GLOBAL_DTYPE = None -# NOTE: come back to this later -# # NEW: Output shared memory for zero-copy results -# _GLOBAL_OUTPUT_SHM = None -# _GLOBAL_OUTPUT_ARRAY = None -# _GLOBAL_OUTPUT_SHAPE = None - -def _init_worker(shm_name, shape, dtype): +def _init_worker(shm_name, shape, dtype, log_queue=None): """ Worker pool initializer: attach to the named shared-memory block holding the background plates, seed numpy/random from the worker PID so each process gets a distinct RNG state, and set up logging. Passing shm_name/shape/dtype through the pool initializer (rather than per-task args) avoids - re-serializing the array on every map() call. The worker installs a SIGTERM handler that - closes its shared-memory file descriptor before letting the signal kill the process; the - main process is responsible for unlinking shared memory afterwards (handled by - ResourceManager). + re-serializing the array on every map() call. `log_queue` must be passed explicitly for + pools whose parent process was spawn-started (the RoundDataProducer's — no inherited Logger + singleton there); fork-started pools omit it and inherit the singleton's queue. The worker + installs a SIGTERM handler that closes its shared-memory file descriptor before letting the + signal kill the process; the main process is responsible for unlinking shared memory + afterwards (handled by ResourceManager). """ - # NOTE: come back to this later - # def _init_worker(shm_name, shape, dtype, output_shm_name=None, output_shape=None): - # """ - # Initialize worker process with shared memory references. - # - # Args: - # shm_name: Name of the shared memory block containing background plates - # shape: Shape of the background array - # dtype: Data type of the background array - # output_shm_name: Name of shared memory for output results (optional) - # output_shape: Shape of the output array (optional) - # - # Why this change: - # Adding output_shm_name and output_shape parameters allows workers to write - # results directly to shared memory instead of returning them through IPC. - # This eliminates ~23GB of pickle serialization per training round. - # """ global _GLOBAL_SHM, _GLOBAL_BACKGROUNDS, _GLOBAL_SHAPE, _GLOBAL_DTYPE - # global _GLOBAL_OUTPUT_SHM, _GLOBAL_OUTPUT_ARRAY, _GLOBAL_OUTPUT_SHAPE # Initialize worker logging - init_worker_logging() + init_worker_logging(log_queue) # Seed processes with process IDs so each worker gets a different random state random.seed(os.getpid()) @@ -91,15 +76,6 @@ def _init_worker(shm_name, shape, dtype): _GLOBAL_SHAPE = shape _GLOBAL_DTYPE = dtype - # NOTE: come back to this later - # # NEW: Attach to output shared memory if provided - # if output_shm_name is not None and output_shape is not None: - # _GLOBAL_OUTPUT_SHM = SharedMemory(name=output_shm_name) - # _GLOBAL_OUTPUT_ARRAY = np.ndarray( - # output_shape, dtype=np.float32, buffer=_GLOBAL_OUTPUT_SHM.buf - # ) - # _GLOBAL_OUTPUT_SHAPE = output_shape - # Ignore SIGINT (Ctrl+C) in workers - let manager from parent handle cleanup coordination signal.signal(signal.SIGINT, signal.SIG_IGN) @@ -565,18 +541,138 @@ def create_true_double( return final, sample_info -def _single_cadence_wrapper(args): +# Cadence generator functions addressable by name, so batched tasks stay picklable-cheap +# (workers resolve the callable locally instead of unpickling a function reference per task) +_CREATE_FUNCTIONS = { + "create_false": create_false, + "create_true_single": create_true_single, + "create_true_double": create_true_double, +} + + +@dataclass(frozen=True) +class ChunkSegment: + """ + One contiguous class-segment of a chunk: `count` cadences of a single signal type written + at rows [start_idx, start_idx + count) of the round's `array_name` memmap. Each chunk is + made of 8 segments (4 quarters for main, 2 halves each for false/true) — the same + contiguous per-chunk layout the in-RAM generator used, so the labels array preserves the + row -> signal-type mapping and the stratified split downstream keeps working unchanged. + """ + + array_name: str # "main" | "false" | "true" + signal_class: str # "main" | "false" | "true" (DB signal_class) + signal_type: str # e.g. "false_no_signal" + create_fn_name: str # key into _CREATE_FUNCTIONS + inject: bool | None + dynamic_range: float | None + start_idx: int # absolute row offset into the round array + count: int + + +# The 4 signal types in main-array (and labels) order within each chunk +_SIGNAL_TYPES = ("false_no_signal", "false_with_rfi", "true_only_eti", "true_eti_rfi") + + +def build_chunk_segments(chunk_start: int, chunk_size: int) -> list[ChunkSegment]: + """ + Lay out the 8 class-segments for one chunk starting at absolute row `chunk_start`. + + Requires chunk_size % 4 == 0 (validated for signal_injection_chunk_size in cli.py), so + quarter = chunk_size // 4 and half = chunk_size // 2 are exact — no oversample/subsample + dance and no post-hoc stats filtering. """ - Multiprocessing worker: unpack the per-task args tuple (function, snr_base, snr_range, - width_bin, freq_resolution, time_resolution, inject, dynamic_range) and call the chosen - create_* generator against _GLOBAL_BACKGROUNDS. Returns the (cadence, sample_info) pair from - the generator. + if chunk_size % 4 != 0: + raise ValueError(f"chunk_size must be divisible by 4, got {chunk_size}") + + quarter = chunk_size // 4 + half = chunk_size // 2 + + def seg(array_name, signal_class, signal_type, fn, inject, dyn, offset, count): + return ChunkSegment( + array_name=array_name, + signal_class=signal_class, + signal_type=signal_type, + create_fn_name=fn, + inject=inject, + dynamic_range=dyn, + start_idx=chunk_start + offset, + count=count, + ) + + return [ + # main: 4-way balanced quarters (order matches _SIGNAL_TYPES / the labels array) + seg("main", "main", "false_no_signal", "create_false", False, None, 0, quarter), + seg("main", "main", "false_with_rfi", "create_false", True, None, quarter, quarter), + seg( + "main", "main", "true_only_eti", "create_true_single", None, None, 2 * quarter, quarter + ), + seg("main", "main", "true_eti_rfi", "create_true_double", None, 1.0, 3 * quarter, quarter), + # false: 2-way balanced halves + seg("false", "false", "false_no_signal", "create_false", False, None, 0, half), + seg("false", "false", "false_with_rfi", "create_false", True, None, half, half), + # true: 2-way balanced halves + seg("true", "true", "true_only_eti", "create_true_single", None, None, 0, half), + seg("true", "true", "true_eti_rfi", "create_true_double", None, 1.0, half, half), + ] + + +def build_segment_tasks( + segment: ChunkSegment, + array_path: str, + task_size: int, + snr_base: float, + snr_range: float, + width_bin: int, + freq_resolution: float, + time_resolution: float, + seed_rng: np.random.Generator, +) -> list[tuple]: + """ + Partition one segment into batched worker tasks of at most `task_size` cadences each. + Together the tasks cover rows [segment.start_idx, segment.start_idx + segment.count) + exactly once. Each task carries a fresh RNG seed drawn from `seed_rng` so results don't + depend on which persistent worker picks the task up. + """ + if task_size < 1: + raise ValueError(f"task_size must be >= 1, got {task_size}") + + tasks = [] + for offset in range(0, segment.count, task_size): + count = min(task_size, segment.count - offset) + seed = int(seed_rng.integers(0, 2**31 - 1)) + tasks.append( + ( + array_path, + segment.start_idx + offset, + count, + segment.create_fn_name, + snr_base, + snr_range, + width_bin, + freq_resolution, + time_resolution, + segment.inject, + segment.dynamic_range, + seed, + ) + ) + return tasks + - Pulling backgrounds from the global rather than passing them per-task avoids pickling the - full plate on every dispatch. +def _run_memmap_task(args: tuple, backgrounds: np.ndarray) -> tuple[float, list[dict]]: + """ + Execute one batched generation task against `backgrounds`: open the target .npy r+ (like + preprocessing._extract_stamps_worker), generate `count` cadences with the chosen create_* + function, and write each result row straight into the memmap. Tasks address disjoint row + ranges, so concurrent pool writes never collide. Returns (elapsed_seconds, + [sample_info, ...]) — the only data that crosses the IPC boundary. """ ( - function, + array_path, + start_idx, + count, + create_fn_name, snr_base, snr_range, width_bin, @@ -584,147 +680,30 @@ def _single_cadence_wrapper(args): time_resolution, inject, dynamic_range, + seed, ) = args - return function( - # NOTE: - # is _GLOBAL_BACKGROUNDS a reference to the shared memory, or the entire data array itself? - # is this method of doing things slower due to pickling costs? - # should we move the index selection into _single_cadence_wrapper so only one background plate is sent at a time? - # how can we benchmark this? - # # # Select random background from plate - # # background_index = int(plate.shape[0] * random.random()) - # # base = plate[background_index, :, :, :] - _GLOBAL_BACKGROUNDS, - snr_base=snr_base, - snr_range=snr_range, - width_bin=width_bin, - freq_resolution=freq_resolution, - time_resolution=time_resolution, - inject=inject, - dynamic_range=dynamic_range, - ) + task_start = time.time() -# NOTE: come back to this later -# def _batch_cadence_worker(args): -# """ -# Process a batch of cadences and write results directly to shared memory. -# -# Args: -# args: Tuple of (start_idx, end_idx, function, snr_base, snr_range, -# width_bin, freq_resolution, time_resolution, inject, dynamic_range) -# -# Returns: -# None - results written directly to _GLOBAL_OUTPUT_ARRAY -# -# Why this approach: -# Instead of returning 192KB per cadence through IPC (which requires pickle -# serialization), workers write directly to pre-allocated shared memory. -# This eliminates the dominant bottleneck: ~23GB of IPC data transfer per round. -# -# Processing cadences in batches (e.g., 500-1000 per task) also reduces -# task scheduling overhead from 120,000 dispatches to ~120-240 dispatches. -# """ -# ( -# start_idx, -# end_idx, -# function, -# snr_base, -# snr_range, -# width_bin, -# freq_resolution, -# time_resolution, -# inject, -# dynamic_range, -# ) = args -# -# # Process each cadence in this batch -# for i in range(start_idx, end_idx): -# result = function( -# _GLOBAL_BACKGROUNDS, -# snr_base=snr_base, -# snr_range=snr_range, -# width_bin=width_bin, -# freq_resolution=freq_resolution, -# time_resolution=time_resolution, -# inject=inject, -# dynamic_range=dynamic_range, -# ) -# -# # Write directly to shared memory output array - NO IPC! -# _GLOBAL_OUTPUT_ARRAY[i, :, :, :] = result -# -# # Return None - no data through IPC - - -def batch_create_cadence( - function, - samples: int, - plate: np.ndarray, - snr_base: int = 10, - snr_range: float = 40, - width_bin: int = 512, - freq_resolution: float = 2.7939677238464355, - time_resolution: float = 18.25361108, - inject: bool | None = None, - dynamic_range: float | None = None, - pool: Pool | None = None, - n_processes: int | None = cpu_count(), - chunks_per_worker: int | None = 4, -) -> tuple[np.ndarray, list[dict]]: - """ - Batch-generate `samples` cadences with the chosen `function` (create_false / - create_true_single / create_true_double), returning (cadence_array of shape - (samples, 6, 16, width_bin), all_sample_info: list of per-sample dicts). - - When `pool` is provided, work is dispatched via _single_cadence_wrapper against the worker - pool's shared-memory background plate; otherwise the loop runs sequentially against `plate` - in-process. chunks_per_worker controls map() chunksize to balance scheduling overhead vs. - parallelism. dynamic_range only applies to create_true_double; inject only to create_false. - """ - # Pre-allocate output array - cadence = np.zeros((samples, 6, 16, width_bin)) + # Per-task seeding keeps results independent of worker scheduling (workers are + # persistent, so PID-only seeding from _init_worker would tie the stream to which + # worker happens to pick the task up). + # NOTE: this seeds the LEGACY global RNGs because new_cadence/create_* draw from + # random.random() and the np.random.* module API. The determinism therefore holds only + # as long as nothing else mutates the global RNG state between cadences within a task; + # threading an explicit np.random.Generator through the create_* functions would remove + # that coupling if it ever matters. + random.seed(seed) + np.random.seed(seed % (2**32)) + + create_fn = _CREATE_FUNCTIONS[create_fn_name] all_sample_info = [] - if pool: - # Parallel execution using provided pool - # Prepare arguments for each parallel task (no plate - uses global) - args_list = [ - ( - function, - snr_base, - snr_range, - width_bin, - freq_resolution, - time_resolution, - inject, - dynamic_range, - ) - for _ in range(samples) - ] - - # Calculate optimal chunksize for load balancing - try: - n_workers = pool._processes - except AttributeError: - n_workers = n_processes - # NOTE: should we use separate chunks_per_worker? how to benchmark? - chunksize = max(1, samples // (n_workers * chunks_per_worker)) - - # Use pool to generate cadences in parallel - for i, (result, sample_info) in enumerate( - # TEST: does return order matter? - pool.map(_single_cadence_wrapper, args_list, chunksize=chunksize) - # pool.imap(_single_cadence_wrapper, args_list, chunksize=chunksize) - # pool.imap_unordered(_single_cadence_wrapper, args_list, chunksize=chunksize) - ): - cadence[i, :, :, :] = result - all_sample_info.append(sample_info) - else: - # Fallback to sequential execution - for i in range(samples): - result, sample_info = function( - plate, + out = np.lib.format.open_memmap(array_path, mode="r+") + try: + for i in range(count): + cadence, sample_info = create_fn( + backgrounds, snr_base=snr_base, snr_range=snr_range, width_bin=width_bin, @@ -733,161 +712,286 @@ def batch_create_cadence( inject=inject, dynamic_range=dynamic_range, ) - cadence[i, :, :, :] = result + out[start_idx + i] = cadence all_sample_info.append(sample_info) + finally: + # Flush in the finally so rows written before a mid-task exception still reach disk + # deterministically (a failed round has no .done manifest and reads as garbage either + # way, but not all filesystems are guaranteed to write back dirty pages on munmap) + out.flush() + del out - return cadence, all_sample_info - - -# NOTE: come back to this later -# def batch_create_cadence( -# function, -# samples: int, -# plate: np.ndarray, -# snr_base: int = 10, -# snr_range: float = 40, -# width_bin: int = 512, -# freq_resolution: float = 2.7939677238464355, -# time_resolution: float = 18.25361108, -# inject: bool | None = None, -# dynamic_range: float | None = None, -# pool: Pool | None = None, -# output_shm: SharedMemory | None = None, -# output_array: np.ndarray | None = None, -# output_offset: int = 0, -# n_processes: int | None = cpu_count(), -# chunks_per_worker: int | None = 4, -# # NOTE: parametrize this in config -# batch_size: int = 500, # cadences per task -# ) -> np.ndarray: -# """ -# Batch wrapper for creating multiple cadences using multiprocessing. -# -# Key changes from original: -# 1. Workers write directly to shared memory (output_array) instead of returning results -# 2. Tasks are batched (batch_size cadences per task) to reduce scheduling overhead -# 3. output_offset allows writing to a specific slice of the output array -# -# Args: -# function: Cadence generation function (create_false, create_true_single, create_true_double) -# samples: Number of cadences to generate -# plate: Background plate array (only used if pool is None) -# snr_base: Base SNR value -# snr_range: SNR range for randomization -# width_bin: Number of frequency bins -# freq_resolution: Frequency resolution in Hz -# time_resolution: Time resolution in seconds -# inject: Whether to inject signals (for create_false) -# dynamic_range: Dynamic range for signal injection (for create_true_double) -# pool: Pre-initialized multiprocessing Pool with shared memory -# output_shm: SharedMemory object for output (used for cleanup tracking) -# output_array: numpy array view of output shared memory -# output_offset: Starting index in output_array for this batch -# n_processes: Number of processes in multiprocessing Pool -# chunks_per_worker: Used to calculate optimal chunksize for load balancing -# batch_size: Number of cadences per worker task (default 500) -# -# Returns: -# Array of shape (samples, 6, 16, width_bin) containing generated cadences -# (either from shared memory or newly allocated) -# """ -# if pool and output_array is not None: -# # OPTIMIZED PATH: Use shared memory output -# -# # Create batched task arguments -# # Instead of 1 task per cadence, create 1 task per batch_size cadences -# tasks = [] -# for batch_start in range(0, samples, batch_size): -# batch_end = min(batch_start + batch_size, samples) -# # Indices are relative to output_offset in the shared memory -# tasks.append( -# ( -# output_offset + batch_start, # start_idx in output array -# output_offset + batch_end, # end_idx in output array -# function, -# snr_base, -# snr_range, -# width_bin, -# freq_resolution, -# time_resolution, -# inject, -# dynamic_range, -# ) -# ) -# -# # Calculate chunksize for pool.map -# # With batched tasks, we have far fewer tasks, so chunksize can be smaller -# n_tasks = len(tasks) -# try: -# n_workers = pool._processes -# except AttributeError: -# n_workers = n_processes -# -# # Aim for ~4 chunks per worker for load balancing -# chunksize = max(1, n_tasks // (n_workers * 4)) -# -# logger.debug( -# f"Dispatching {n_tasks} batched tasks (batch_size={batch_size}, chunksize={chunksize})" -# ) -# -# # Execute - results are written directly to shared memory -# # We use pool.map which blocks until complete, but returns None for each task -# list(pool.map(_batch_cadence_worker, tasks, chunksize=chunksize)) -# -# # Return view of the shared memory output for this batch -# return output_array[output_offset : output_offset + samples] -# -# elif pool: -# # LEGACY PATH: Pool exists but no shared memory output -# # Fall back to original IPC-based approach (for backward compatibility) -# cadence = np.zeros((samples, 6, 16, width_bin), dtype=np.float32) -# -# args_list = [ -# ( -# function, -# snr_base, -# snr_range, -# width_bin, -# freq_resolution, -# time_resolution, -# inject, -# dynamic_range, -# ) -# for _ in range(samples) -# ] -# -# try: -# n_workers = pool._processes -# except AttributeError: -# n_workers = n_processes -# -# # FIX: Use reasonable chunksize instead of always 1 -# chunksize = max(50, samples // (n_workers * 4)) -# -# for i, result in enumerate( -# pool.imap(_single_cadence_wrapper, args_list, chunksize=chunksize) -# ): -# cadence[i, :, :, :] = result -# -# return cadence -# -# else: -# # Sequential execution (no pool) -# cadence = np.zeros((samples, 6, 16, width_bin), dtype=np.float32) -# -# for i in range(samples): -# cadence[i, :, :, :] = function( -# plate, -# snr_base=snr_base, -# snr_range=snr_range, -# width_bin=width_bin, -# freq_resolution=freq_resolution, -# time_resolution=time_resolution, -# inject=inject, -# dynamic_range=dynamic_range, -# ) -# -# return cadence + return time.time() - task_start, all_sample_info + + +def _memmap_task_worker(args: tuple) -> tuple[float, list[dict]]: + """Pool entry point for batched memmap tasks: run against the shared-memory backgrounds.""" + return _run_memmap_task(args, _GLOBAL_BACKGROUNDS) + + +def write_segment_stats(db, tag: str, segment: dict) -> None: + """ + Write one class-segment's injection stats to the DB (main-process only — the DB queue is + a thread queue.Queue, not process-safe). `segment` is the dict emitted by + generate_round_to_memmap's stats_cb. + + Per-sample writes: 6 intensity stats x 3 stages = 18 rows, plus 0-12 signal-characteristic + rows depending on signal_type. Segment-level writes: 4 metadata rows. + """ + if db is None: + raise RuntimeError("No database instance detected - cannot write injection stats") + + round_number = segment["round_number"] + chunk_number = segment["chunk_number"] + signal_class = segment["signal_class"] + signal_type = segment["signal_type"] + timestamp = segment["timestamp"] + + for sample_idx, sample_info in enumerate(segment["stats_list"]): + background_index = sample_info["background_index"] + intensity_stats = sample_info["intensity_stats"] + signal_info = sample_info["signal_info"] + slope_was_clamped = sample_info.get("slope_was_clamped", False) + + # Intensity stats for each stage + for stage in ["A", "B", "C"]: + stage_stats = intensity_stats[stage] + + for stat_name in [ + "global_mean", + "global_median", + "global_std", + "global_mad", + "global_skew", + "global_kurtosis", + ]: + db.write_injection_stat( + stat_name=stat_name, + value=stage_stats[stat_name], + round_number=round_number, + chunk_number=chunk_number, + sample_index=sample_idx, + background_index=background_index, + signal_class=signal_class, + signal_type=signal_type, + injection_stage=stage, + slope_clamped=slope_was_clamped, + tag=tag, + timestamp=timestamp, + ) + + # NOTE: what happens when signal_info is empty for false_no_rfi signal types? + # Signal characteristics (eti_snr, rfi_drift_rate, etc.) + # injection_stage=None since these describe the injection itself + for stat_name, value in signal_info.items(): + db.write_injection_stat( + stat_name=stat_name, + value=float(value), + round_number=round_number, + chunk_number=chunk_number, + sample_index=sample_idx, + background_index=background_index, + signal_class=signal_class, + signal_type=signal_type, + injection_stage=None, + slope_clamped=slope_was_clamped, + tag=tag, + timestamp=timestamp, + ) + + # Segment-level metadata stats (once per segment, not per sample) + metadata_stats = [ + ("snr_range_floor", segment["snr_range_floor"]), + ("snr_range_ceil", segment["snr_range_ceil"]), + ("num_samples", float(segment["num_samples"])), + ("inject_duration", segment["inject_duration"]), + ] + + for stat_name, value in metadata_stats: + db.write_injection_stat( + stat_name=stat_name, + value=value, + round_number=round_number, + chunk_number=chunk_number, + sample_index=None, + background_index=None, + signal_class=signal_class, + signal_type=signal_type, + injection_stage=None, + tag=tag, + timestamp=timestamp, + ) + + +def generate_round_to_memmap( + paths: RoundDataPaths, + n_samples: int, + snr_base: float, + snr_range: float, + *, + width_bin: int, + num_observations: int, + time_bins: int, + chunk_size: int, + task_size: int, + freq_resolution: float, + time_resolution: float, + pool: Pool | None = None, + backgrounds: np.ndarray | None = None, + round_num: int | None = None, + stats_cb=None, + progress_cb=None, +) -> dict: + """ + Generate one round's triplet dataset straight into disk-backed .npy memmaps. + + main: collapsed cadences, 1/4 balanced across the 4 signal types (labels array tracks the + per-row type); false: 1/2 false_no_signal + 1/2 false_with_rfi; true: 1/2 true_only_eti + + 1/2 true_eti_rfi — each of shape (n_samples, num_observations, time_bins, width_bin) + float32, laid out contiguously per chunk (see build_chunk_segments). + + Work is dispatched as one unified batched task list per chunk through a single pool.map + barrier (instead of the old 8 sequential per-class barriers); workers write rows in-place + and return only stats dicts. With `pool` None, tasks run sequentially in-process against + `backgrounds`. On success the labels array is saved and an atomic .done manifest is + written (a crash mid-generation leaves no manifest, so the dir reads as garbage). + + stats_cb(segment_dict) fires once per class-segment per chunk; progress_cb(chunk, + n_chunks) once per chunk. Returns the manifest dict. + """ + if n_samples % 4 != 0: + raise ValueError(f"n_samples must be divisible by 4, got {n_samples}") + if chunk_size % 4 != 0: + raise ValueError(f"chunk_size must be divisible by 4, got {chunk_size}") + if pool is None and backgrounds is None: + raise ValueError("backgrounds must be provided when no pool is given") + + wall_start = time.time() + + # Start from a clean slate: a previous partial generation (no .done manifest) may have + # left stale arrays behind. + shutil.rmtree(paths.round_dir, ignore_errors=True) + os.makedirs(paths.round_dir, exist_ok=True) + + # Pre-create the three destination memmaps; workers reopen them r+ per task + shape = (n_samples, num_observations, time_bins, width_bin) + array_paths = paths.array_paths + for path in array_paths.values(): + mm = np.lib.format.open_memmap(path, mode="w+", dtype=np.float32, shape=shape) + del mm # Close immediately: creation only reserves the file; workers do the writing + + labels = np.empty(n_samples, dtype="U20") + + n_chunks = max(1, (n_samples + chunk_size - 1) // chunk_size) + seed_rng = np.random.default_rng() # OS entropy; per-task seeds are drawn from this + + logger.info( + f"Generating {n_samples} samples into {paths.round_dir} " + f"({n_chunks} chunks of max {chunk_size}, task size {task_size})" + ) + + for chunk_idx in range(n_chunks): + chunk_start = chunk_idx * chunk_size + this_chunk_size = min(chunk_size, n_samples - chunk_start) + if this_chunk_size <= 0: + break + + # Capture single timestamp for all stats in this chunk + chunk_timestamp = time.time() + + segments = build_chunk_segments(chunk_start, this_chunk_size) + + # One unified task list across all 8 class-segments -> one pool.map barrier per chunk + tasks: list[tuple] = [] + task_owners: list[int] = [] # task index -> segment index (map preserves order) + for segment_idx, segment in enumerate(segments): + segment_tasks = build_segment_tasks( + segment, + array_paths[segment.array_name], + task_size, + snr_base, + snr_range, + width_bin, + freq_resolution, + time_resolution, + seed_rng, + ) + tasks.extend(segment_tasks) + task_owners.extend([segment_idx] * len(segment_tasks)) + + # Labels mirror the main array's contiguous per-chunk layout + quarter = this_chunk_size // 4 + chunk_labels: list[str] = [] + for signal_type in _SIGNAL_TYPES: + chunk_labels.extend([signal_type] * quarter) + labels[chunk_start : chunk_start + this_chunk_size] = chunk_labels + + logger.info( + f"Generating chunk {chunk_idx + 1}/{n_chunks} " + f"({this_chunk_size} samples, {len(tasks)} tasks)" + ) + + if pool is not None: + # chunksize=1: each task is already a coarse unit of work (task_size cadences), + # so per-task scheduling overhead is negligible and load balance is best + results = pool.map(_memmap_task_worker, tasks, chunksize=1) + else: + results = [_run_memmap_task(task, backgrounds) for task in tasks] + + # Group returned stats by segment. inject_duration is the sum of the segment's + # per-task wall times (aggregate worker time, not the barrier wall time — tasks from + # all segments run interleaved across the pool, so per-segment barrier walls no + # longer exist). + segment_stats: list[list[dict]] = [[] for _ in segments] + segment_durations = [0.0] * len(segments) + for owner, (task_duration, task_stats) in zip(task_owners, results, strict=True): + segment_stats[owner].extend(task_stats) + segment_durations[owner] += task_duration + + if stats_cb is not None: + for segment, stats_list, duration in zip( + segments, segment_stats, segment_durations, strict=True + ): + stats_cb( + { + "round_number": round_num, + "chunk_number": chunk_idx + 1, + "signal_class": segment.signal_class, + "signal_type": segment.signal_type, + "snr_range_floor": snr_base, + "snr_range_ceil": snr_base + snr_range, + "num_samples": len(stats_list), + "inject_duration": duration, + "timestamp": chunk_timestamp, + "stats_list": stats_list, + } + ) + + del results, segment_stats + gc.collect() + + if progress_cb is not None: + progress_cb(chunk_idx + 1, n_chunks) + logger.info(f"Chunk {chunk_idx + 1}/{n_chunks} complete") + + np.save(paths.labels_path, labels) + + # The .done manifest is only written after every chunk has finished — the same atomicity + # idea as preprocessing's .tmp -> os.replace stamp extraction + manifest = build_manifest( + paths, + n_samples=n_samples, + snr_base=snr_base, + snr_range=snr_range, + wall_time_s=time.time() - wall_start, + chunk_count=n_chunks, + ) + write_done_manifest(paths, manifest) + + logger.info( + f"Round data generation complete: {paths.round_dir} ({manifest['wall_time_s']:.1f}s wall)" + ) + return manifest class DataGenerator: @@ -900,7 +1004,10 @@ def __init__( """ Initialize the generator. background_plates is an array of preprocessed background observations with shape (n_backgrounds, 6, 16, 512). Plates are copied into shared memory - for worker access and a persistent multiprocessing pool is set up here. + for worker access; the worker pool itself is created lazily on the first in-process + generate_round() call (when overlap_data_generation is on, the RoundDataProducer process + owns its own pool against the same shared memory, and this one is never needed for the + beta-VAE rounds). """ self.config = get_config() if self.config is None: @@ -914,12 +1021,12 @@ def __init__( if self.manager is None: raise ValueError("get_manager() returned None") + # Pool is created on demand by _ensure_pool() + self.pool = None + # Load background plates into shared memory self._load_backgrounds(background_plates) - # Setup persistent process pool for efficient parallel execution - self._setup_managed_pool() - def _load_backgrounds(self, background_plates: np.ndarray): """Load background plates into shared memory""" # Sanity check: verify no NaN or Inf values in background plates @@ -945,7 +1052,6 @@ def _load_backgrounds(self, background_plates: np.ndarray): # Get multiprocessing params from config self.n_processes = self.config.manager.n_processes - self.chunks_per_worker = self.config.manager.chunks_per_worker # Setup shared memory to avoid duplicating background data across workers if self.n_processes > 1: @@ -972,53 +1078,25 @@ def _load_backgrounds(self, background_plates: np.ndarray): logger.info(f" Background shape: {self._background_shape}") logger.info(f" Background dtype: {self._background_dtype}") - def _setup_managed_pool(self): + def _ensure_pool(self): """ - Stand up a persistent ResourceManager-owned multiprocessing pool whose workers attach to - the shared-memory background block at init time, so per-task dispatches don't have to - re-serialize the plates. Falls back to sequential mode (self.pool = None) when - n_processes == 1. + Lazily stand up the persistent ResourceManager-owned multiprocessing pool whose workers + attach to the shared-memory background block at init time, so per-task dispatches don't + have to re-serialize the plates. No-op in sequential mode (n_processes == 1, no shared + memory) — generate_round() then runs tasks in-process. The pool must be released via _free_managed_pool() or close() — the ResourceManager won't reap it automatically. """ # NOTE: should we explicitly guarantee only 1 shm & 1 pool can exist at a time? - # If shared memory exists, then create pool using shared memory reference - if self.shm: - self.pool = self.manager.create_pool( - n_processes=self.n_processes, - name=f"DataGen_pool_{id(self)}", # NOTE: come back to this later - initializer=_init_worker, - initargs=(self.shm.name, self._background_shape, self._background_dtype), - ) - # Else run in sequential mode (no pool) - else: - self.pool = None - logger.info("DataGenerator running in sequential mode (n_processes=1)") - - # NOTE: come back to this later - # def _setup_managed_pool(self): - # """ - # Setup managed multiprocessing pool with shared memory. - # - # Why we don't create output shared memory here: - # Output shared memory size depends on the batch size requested in generate_triplet_batch(), - # which varies between calls. We create output shared memory on-demand in generate_triplet_batch() - # and pass references to workers via the task arguments. - # - # The pool is initialized with background shared memory only. Workers will attach to - # output shared memory when provided via _batch_cadence_worker args. - # """ - # if self.shm: - # self.pool = self.manager.create_pool( - # n_processes=self.n_processes, - # name=f"DataGen_pool_{id(self)}", - # initializer=_init_worker, - # initargs=(self.shm.name, self._background_shape, self._background_dtype), - # ) - # else: - # self.pool = None - # logger.info("DataGenerator running in sequential mode (n_processes=1)") + if self.pool is not None or not self.shm: + return + self.pool = self.manager.create_pool( + n_processes=self.n_processes, + name=f"DataGen_pool_{id(self)}", # NOTE: come back to this later + initializer=_init_worker, + initargs=(self.shm.name, self._background_shape, self._background_dtype), + ) def _free_managed_pool(self): """Close multiprocessing pool""" @@ -1032,13 +1110,13 @@ def reset_managed_pool(self): Should be called between training rounds, since workers can accumulate memory through memory fragmentation in long-lived processes, python's reference counter leaking in workers, - and caches / global state accumulating in workers. + and caches / global state accumulating in workers. The pool is re-created lazily by + _ensure_pool() on the next generate_round() call. """ if hasattr(self, "pool") and self.pool is not None: try: self._free_managed_pool() gc.collect() # Garbage collect between resets - self._setup_managed_pool() except Exception as e: logger.warning(f"Error resetting DataGenerator pool: {e}") @@ -1054,674 +1132,41 @@ def close(self): self._free_managed_shared_memory() logger.info("DataGenerator closed") - # NOTE: update docstring values when more stats are added - def _write_batch_stats( + def generate_round( self, - stats_list: list[dict], - round_number: int | None, - chunk_number: int, - signal_class: str, - signal_type: str, - snr_range_floor: float, - snr_range_ceil: float, - num_samples: int, - inject_duration: float, - timestamp: float, - ): - """ - Write per-sample stats to DB. - - Per-sample writes: - - Intensity stats: 6 stats × 3 stages = 18 writes per sample - - Signal characteristics: 0-12 writes depending on signal_type - - background_index is included with each write - - Batch-level writes: - - 4 metadata stats written once per batch - """ - tag = self.config.checkpoint.save_tag - - if self.db is None: - raise RuntimeError( - "No database instance detected - cannot generate training progress plot" - ) - - # Write per-sample stats - for sample_idx, sample_info in enumerate(stats_list): - background_index = sample_info["background_index"] - intensity_stats = sample_info["intensity_stats"] - signal_info = sample_info["signal_info"] - slope_was_clamped = sample_info.get("slope_was_clamped", False) - - # Write intensity stats for each stage - for stage in ["A", "B", "C"]: - stage_stats = intensity_stats[stage] - - for stat_name in [ - "global_mean", - "global_median", - "global_std", - "global_mad", - "global_skew", - "global_kurtosis", - ]: - self.db.write_injection_stat( - stat_name=stat_name, - value=stage_stats[stat_name], - round_number=round_number, - chunk_number=chunk_number, - sample_index=sample_idx, - background_index=background_index, - signal_class=signal_class, - signal_type=signal_type, - injection_stage=stage, - slope_clamped=slope_was_clamped, - tag=tag, - timestamp=timestamp, - ) - - # NOTE: what happens when signal_info is empty for false_no_rfi signal types? - # Write signal characteristics (eti_snr, rfi_drift_rate, etc.) - # injection_stage=None since these describe the injection itself - for stat_name, value in signal_info.items(): - self.db.write_injection_stat( - stat_name=stat_name, - value=float(value), - round_number=round_number, - chunk_number=chunk_number, - sample_index=sample_idx, - background_index=background_index, - signal_class=signal_class, - signal_type=signal_type, - injection_stage=None, - slope_clamped=slope_was_clamped, - tag=tag, - timestamp=timestamp, - ) - - # Write batch-level metadata stats (once per batch, not per sample) - metadata_stats = [ - ("snr_range_floor", snr_range_floor), - ("snr_range_ceil", snr_range_ceil), - ("num_samples", float(num_samples)), - ("inject_duration", inject_duration), - ] - - for stat_name, value in metadata_stats: - self.db.write_injection_stat( - stat_name=stat_name, - value=value, - round_number=round_number, - chunk_number=chunk_number, - sample_index=None, - background_index=None, - signal_class=signal_class, - signal_type=signal_type, - injection_stage=None, - tag=tag, - timestamp=timestamp, - ) - - def generate_triplet_batch( - self, n_samples: int, snr_base: int, snr_range: int, round_num: int | None = None - ) -> dict[str, np.ndarray]: + paths, + n_samples: int, + snr_base: float, + snr_range: float, + round_num: int | None = None, + ) -> dict: """ - Generate triplet batch using chunking & multiprocessing - - main: collapsed cadences - - total: n_samples - - split: 1/4 balanced between false-no-signal, false-with-rfi, true-single, true-double - false: non-collapsed false cadences - - total: n_samples - - split: 1/2 balanced between false-no-signal, false-with-rfi - true: non-collapsed true cadences - - total: n_samples - - split: 1/2 balanced between true-single, true-double + Generate one round's triplet dataset into the disk-backed memmaps at `paths` + (a round_data.RoundDataPaths), writing injection stats to the DB as chunks complete. + Used for the sequential (non-overlapped) generation path — the RF dataset, and beta-VAE + rounds when overlap_data_generation is disabled. Returns the .done manifest dict. """ - max_chunk_size = self.config.training.signal_injection_chunk_size - n_chunks = max(1, (n_samples + max_chunk_size - 1) // max_chunk_size) - - logger.info(f"Generating {n_samples} samples in {n_chunks} chunks of max {max_chunk_size}") - - # NOTE: freq & time dims hard-coded here - # Pre-allocate output arrays - all_main = np.empty((n_samples, 6, 16, self.width_bin), dtype=np.float32) - all_labels = np.empty(n_samples, dtype="U20") - all_false = np.empty((n_samples, 6, 16, self.width_bin), dtype=np.float32) - all_true = np.empty((n_samples, 6, 16, self.width_bin), dtype=np.float32) - - for chunk_idx in range(n_chunks): - chunk_size = min(max_chunk_size, n_samples - chunk_idx * max_chunk_size) - if chunk_size <= 0: - break - - start_idx = chunk_idx * max_chunk_size - end_idx = start_idx + chunk_size - - logger.info(f"Generating chunk {chunk_idx + 1}/{n_chunks} with {chunk_size} samples") - - # Capture single timestamp for all stats in this chunk - chunk_timestamp = time.time() - - # Split chunk into equal partitions (for balanced classes) - # Ceiling division ensures 4*quarter >= chunk_size and 2*half >= chunk_size, - # so we always generate enough samples. Excess is randomly discarded after concat. - quarter = max(1, -(-chunk_size // 4)) - half = max(1, -(-chunk_size // 2)) - - # Pure background (main) - batch_start = time.time() - quarter_false_no_signal, stats_main_false_no_signal = batch_create_cadence( - create_false, - quarter, - self.backgrounds, - snr_base=snr_base, - snr_range=snr_range, - width_bin=self.width_bin, - freq_resolution=self.freq_resolution, - time_resolution=self.time_resolution, - inject=False, - pool=self.pool, - n_processes=self.n_processes, - chunks_per_worker=self.chunks_per_worker, - ) - duration_main_false_no_signal = time.time() - batch_start - - # RFI only (main) - batch_start = time.time() - quarter_false_with_rfi, stats_main_false_with_rfi = batch_create_cadence( - create_false, - quarter, - self.backgrounds, - snr_base=snr_base, - snr_range=snr_range, - width_bin=self.width_bin, - freq_resolution=self.freq_resolution, - time_resolution=self.time_resolution, - inject=True, - pool=self.pool, - n_processes=self.n_processes, - chunks_per_worker=self.chunks_per_worker, - ) - duration_main_false_with_rfi = time.time() - batch_start - - # ETI only (main) - batch_start = time.time() - quarter_true_single, stats_main_true_only_eti = batch_create_cadence( - create_true_single, - quarter, - self.backgrounds, - snr_base=snr_base, - snr_range=snr_range, - width_bin=self.width_bin, - freq_resolution=self.freq_resolution, - time_resolution=self.time_resolution, - pool=self.pool, - n_processes=self.n_processes, - chunks_per_worker=self.chunks_per_worker, - ) - duration_main_true_only_eti = time.time() - batch_start - - # ETI + RFI (main) - batch_start = time.time() - quarter_true_double, stats_main_true_eti_rfi = batch_create_cadence( - create_true_double, - quarter, - self.backgrounds, - snr_base=snr_base, - snr_range=snr_range, - width_bin=self.width_bin, - freq_resolution=self.freq_resolution, - time_resolution=self.time_resolution, - dynamic_range=1, - pool=self.pool, - n_processes=self.n_processes, - chunks_per_worker=self.chunks_per_worker, - ) - duration_main_true_eti_rfi = time.time() - batch_start - - # Concatenate for main training data (collapsed cadences) - chunk_main = np.concatenate( - [ - quarter_false_no_signal, - quarter_false_with_rfi, - quarter_true_single, - quarter_true_double, - ], - axis=0, - ) + self._ensure_pool() - # Build per-cadence labels for this chunk - chunk_labels = np.array( - ["false_no_signal"] * quarter - + ["false_with_rfi"] * quarter - + ["true_only_eti"] * quarter - + ["true_eti_rfi"] * quarter, - dtype="U20", - ) - - # Subsample to chunk_size if we generated more than needed (ceiling division remainder) - n_generated_main = 4 * quarter - if n_generated_main > chunk_size: - keep_main = np.sort( - np.random.choice(n_generated_main, size=chunk_size, replace=False) - ) - chunk_main = chunk_main[keep_main] - chunk_labels = chunk_labels[keep_main] - - # Filter stats to only include kept samples - main_stats_groups = [ - (stats_main_false_no_signal, 0), - (stats_main_false_with_rfi, quarter), - (stats_main_true_only_eti, 2 * quarter), - (stats_main_true_eti_rfi, 3 * quarter), - ] - for stats_ref, offset in main_stats_groups: - type_mask = (keep_main >= offset) & (keep_main < offset + quarter) - kept_indices = keep_main[type_mask] - offset - stats_ref[:] = [stats_ref[i] for i in kept_indices] - - # Write main batch stats (deferred to after subsampling so only kept samples are written) - for stats_list, signal_type, duration in [ - (stats_main_false_no_signal, "false_no_signal", duration_main_false_no_signal), - (stats_main_false_with_rfi, "false_with_rfi", duration_main_false_with_rfi), - (stats_main_true_only_eti, "true_only_eti", duration_main_true_only_eti), - (stats_main_true_eti_rfi, "true_eti_rfi", duration_main_true_eti_rfi), - ]: - self._write_batch_stats( - stats_list=stats_list, - round_number=round_num, - chunk_number=chunk_idx + 1, - signal_class="main", - signal_type=signal_type, - snr_range_floor=snr_base, - snr_range_ceil=snr_base + snr_range, - num_samples=len(stats_list), - inject_duration=duration, - timestamp=chunk_timestamp, - ) - - # Generate separate true/false non-collapsed cadences for training set diversity - # Used to calculate clustering loss & train RF - - # Pure background (false) - batch_start = time.time() - half_false_no_signal, stats_false_no_signal = batch_create_cadence( - create_false, - half, - self.backgrounds, - snr_base=snr_base, - snr_range=snr_range, - width_bin=self.width_bin, - freq_resolution=self.freq_resolution, - time_resolution=self.time_resolution, - inject=False, - pool=self.pool, - n_processes=self.n_processes, - chunks_per_worker=self.chunks_per_worker, - ) - duration_false_no_signal = time.time() - batch_start - - # RFI only (false) - batch_start = time.time() - half_false_with_rfi, stats_false_with_rfi = batch_create_cadence( - create_false, - half, - self.backgrounds, - snr_base=snr_base, - snr_range=snr_range, - width_bin=self.width_bin, - freq_resolution=self.freq_resolution, - time_resolution=self.time_resolution, - inject=True, - pool=self.pool, - n_processes=self.n_processes, - chunks_per_worker=self.chunks_per_worker, - ) - duration_false_with_rfi = time.time() - batch_start - - # ETI only (true) - batch_start = time.time() - half_true_single, stats_true_only_eti = batch_create_cadence( - create_true_single, - half, - self.backgrounds, - snr_base=snr_base, - snr_range=snr_range, - width_bin=self.width_bin, - freq_resolution=self.freq_resolution, - time_resolution=self.time_resolution, - pool=self.pool, - n_processes=self.n_processes, - chunks_per_worker=self.chunks_per_worker, - ) - duration_true_only_eti = time.time() - batch_start - - # ETI + RFI (true) - batch_start = time.time() - half_true_double, stats_true_eti_rfi = batch_create_cadence( - create_true_double, - half, - self.backgrounds, - snr_base=snr_base, - snr_range=snr_range, - width_bin=self.width_bin, - freq_resolution=self.freq_resolution, - time_resolution=self.time_resolution, - dynamic_range=1, - pool=self.pool, - n_processes=self.n_processes, - chunks_per_worker=self.chunks_per_worker, - ) - duration_true_eti_rfi = time.time() - batch_start - - chunk_false = np.concatenate([half_false_no_signal, half_false_with_rfi], axis=0) - - chunk_true = np.concatenate([half_true_single, half_true_double], axis=0) - - # Subsample to chunk_size if we generated more than needed (ceiling division remainder) - n_generated_half = 2 * half - if n_generated_half > chunk_size: - keep_false = np.sort( - np.random.choice(n_generated_half, size=chunk_size, replace=False) - ) - keep_true = np.sort( - np.random.choice(n_generated_half, size=chunk_size, replace=False) - ) - chunk_false = chunk_false[keep_false] - chunk_true = chunk_true[keep_true] - - # Filter stats to only include kept samples - # chunk_false layout: [0, half) = no_signal, [half, 2*half) = with_rfi - stats_false_no_signal[:] = [ - stats_false_no_signal[i] for i in keep_false[keep_false < half] - ] - stats_false_with_rfi[:] = [ - stats_false_with_rfi[i - half] for i in keep_false[keep_false >= half] - ] - # chunk_true layout: [0, half) = only_eti, [half, 2*half) = eti_rfi - stats_true_only_eti[:] = [ - stats_true_only_eti[i] for i in keep_true[keep_true < half] - ] - stats_true_eti_rfi[:] = [ - stats_true_eti_rfi[i - half] for i in keep_true[keep_true >= half] - ] - - # Write false/true batch stats (deferred to after subsampling) - for stats_list, signal_class, signal_type, duration in [ - (stats_false_no_signal, "false", "false_no_signal", duration_false_no_signal), - (stats_false_with_rfi, "false", "false_with_rfi", duration_false_with_rfi), - (stats_true_only_eti, "true", "true_only_eti", duration_true_only_eti), - (stats_true_eti_rfi, "true", "true_eti_rfi", duration_true_eti_rfi), - ]: - self._write_batch_stats( - stats_list=stats_list, - round_number=round_num, - chunk_number=chunk_idx + 1, - signal_class=signal_class, - signal_type=signal_type, - snr_range_floor=snr_base, - snr_range_ceil=snr_base + snr_range, - num_samples=len(stats_list), - inject_duration=duration, - timestamp=chunk_timestamp, - ) + tag = self.config.checkpoint.save_tag - # Store chunks directly into output array - all_main[start_idx:end_idx] = chunk_main - all_labels[start_idx:end_idx] = chunk_labels - all_false[start_idx:end_idx] = chunk_false - all_true[start_idx:end_idx] = chunk_true - - # Clean up chunk data immediately - del ( - quarter_false_no_signal, - quarter_false_with_rfi, - quarter_true_single, - quarter_true_double, - ) - del half_false_no_signal, half_false_with_rfi, half_true_single, half_true_double - del chunk_main, chunk_labels, chunk_false, chunk_true - del stats_main_false_no_signal, stats_main_false_with_rfi - del stats_main_true_only_eti, stats_main_true_eti_rfi - del stats_false_no_signal, stats_false_with_rfi - del stats_true_only_eti, stats_true_eti_rfi - gc.collect() - - logger.info(f"Chunk {chunk_idx + 1} complete, memory cleared") - - # Create result dictionary with references to pre-allocated arrays - result = { - "concatenated": all_main, - "labels": all_labels, - "false": all_false, - "true": all_true, - } - - # NOTE: is there a more efficient way to do this? these checks currently take a few minutes to complete. should we comment this portion out? - # # Sanity check: verify post-injection data normalization - # for key in ["concatenated", "false", "true"]: - # min_val = np.min(result[key]) - # max_val = np.max(result[key]) - # mean_val = np.mean(result[key]) - # logger.info( - # f"Post-injection {key} stats: min={min_val:.6f}, max={max_val:.6f}, mean={mean_val:.6f}" - # ) - # if max_val > 1.0: - # logger.error(f"Post-injection {key} values too large! Max: {max_val}") - # raise ValueError(f"Post-injection {key} normalization check failed") - # elif min_val < 0.0: - # logger.error(f"Post-injection {key} values too small! Min: {min_val}") - # raise ValueError(f"Post-injection {key} normalization check failed") - # elif np.isnan(result[key]).any(): - # logger.error(f"Post-injection {key} contains NaN values!") - # raise ValueError(f"Post-injection {key} normalization check failed") - # elif np.isinf(result[key]).any(): - # logger.error(f"Post-injection {key} contains Inf values!") - # raise ValueError(f"Post-injection {key} normalization check failed") - # else: - # logger.info(f"Post-injection {key} data properly normalized") - - return result - - # NOTE: come back to this later - # def generate_triplet_batch( - # self, n_samples: int, snr_base: int, snr_range: int - # ) -> dict[str, np.ndarray]: - # """ - # Generate triplet batch using unified task submission and shared memory outputs. - # - # Key optimizations: - # 1. Single shared memory allocation for all outputs - # 2. All 8 batch types submitted as unified task queue - # 3. Workers write directly to shared memory (no IPC returns) - # 4. Single synchronization point instead of 8 - # - # Output structure: - # main: collapsed cadences (n_samples total) - # - 1/4 false-no-signal, 1/4 false-with-rfi, 1/4 true-single, 1/4 true-double - # false: non-collapsed false cadences (n_samples total) - # - 1/2 false-no-signal, 1/2 false-with-rfi - # true: non-collapsed true cadences (n_samples total) - # - 1/2 true-single, 1/2 true-double - # """ - # max_chunk_size = self.config.training.signal_injection_chunk_size - # n_chunks = max(1, (n_samples + max_chunk_size - 1) // max_chunk_size) - # - # logger.info(f"Generating {n_samples} samples in {n_chunks} chunks of max {max_chunk_size}") - # - # # Pre-allocate output arrays - # all_main = np.empty((n_samples, 6, 16, self.width_bin), dtype=np.float32) - # all_false = np.empty((n_samples, 6, 16, self.width_bin), dtype=np.float32) - # all_true = np.empty((n_samples, 6, 16, self.width_bin), dtype=np.float32) - # - # for chunk_idx in range(n_chunks): - # chunk_size = min(max_chunk_size, n_samples - chunk_idx * max_chunk_size) - # if chunk_size <= 0: - # break - # - # start_idx = chunk_idx * max_chunk_size - # end_idx = start_idx + chunk_size - # - # logger.info(f"Generating chunk {chunk_idx + 1}/{n_chunks} with {chunk_size} samples") - # - # # Calculate sample counts for this chunk - # quarter = max(1, chunk_size // 4) - # half = max(1, chunk_size // 2) - # - # # Total outputs needed for this chunk: - # # - main: 4 * quarter = chunk_size - # # - false: 2 * half = chunk_size - # # - true: 2 * half = chunk_size - # # Total: 3 * chunk_size - # total_outputs = 3 * chunk_size - # - # # Create output shared memory for this chunk - # output_shape = (total_outputs, 6, 16, self.width_bin) - # output_nbytes = int(np.prod(output_shape) * np.float32().nbytes) - # - # output_shm = self.manager.create_shared_memory( - # size=output_nbytes, - # name=f"DataGen_output_chunk_{chunk_idx}_{id(self)}", - # ) - # - # # Create numpy view of output shared memory - # output_array = np.ndarray(output_shape, dtype=np.float32, buffer=output_shm.buf) - # - # # Reinitialize pool with output shared memory reference - # # Workers need to attach to the new output shared memory - # if self.pool is not None: - # self._free_managed_pool() - # - # self.pool = self.manager.create_pool( - # n_processes=self.n_processes, - # name=f"DataGen_pool_chunk_{chunk_idx}_{id(self)}", - # initializer=_init_worker, - # initargs=( - # self.shm.name, - # self._background_shape, - # self._background_dtype, - # output_shm.name, # NEW: output shared memory - # output_shape, # NEW: output shape - # ), - # ) - # - # # Define output layout in shared memory: - # # [0:quarter] -> main: false_no_signal - # # [quarter:2*quarter] -> main: false_with_rfi - # # [2*quarter:3*quarter] -> main: true_single - # # [3*quarter:4*quarter] -> main: true_double - # # [chunk_size:chunk_size+half] -> false: no_signal - # # [chunk_size+half:chunk_size+2*half] -> false: with_rfi - # # [2*chunk_size:2*chunk_size+half] -> true: single - # # [2*chunk_size+half:3*chunk_size] -> true: double - # - # # Build unified task list for all 8 batch types - # batch_size = 500 # Cadences per task - # all_tasks = [] - # - # # Helper to create batched tasks for a given output range - # def add_tasks( - # output_start, - # count, - # function, - # inject=None, - # dynamic_range=None, - # _batch_size=batch_size, - # _all_tasks=all_tasks, - # ): - # for batch_start in range(0, count, _batch_size): - # batch_end = min(batch_start + _batch_size, count) - # _all_tasks.append( - # ( - # output_start + batch_start, - # output_start + batch_end, - # function, - # snr_base, - # snr_range, - # self.width_bin, - # self.freq_resolution, - # self.time_resolution, - # inject, - # dynamic_range, - # ) - # ) - # - # # Main outputs (quarters) - # add_tasks(0, quarter, create_false, inject=False) - # add_tasks(quarter, quarter, create_false, inject=True) - # add_tasks(2 * quarter, quarter, create_true_single) - # add_tasks(3 * quarter, quarter, create_true_double, dynamic_range=1) - # - # # False outputs (halves) - # add_tasks(chunk_size, half, create_false, inject=False) - # add_tasks(chunk_size + half, half, create_false, inject=True) - # - # # True outputs (halves) - # add_tasks(2 * chunk_size, half, create_true_single) - # add_tasks(2 * chunk_size + half, half, create_true_double, dynamic_range=1) - # - # logger.info(f"Submitting {len(all_tasks)} unified tasks to pool") - # - # # Execute ALL tasks in single pool.map call - # # This is the key optimization: one sync barrier instead of 8 - # n_workers = self.n_processes - # chunksize = max(1, len(all_tasks) // (n_workers * 4)) - # - # list(self.pool.map(_batch_cadence_worker, all_tasks, chunksize=chunksize)) - # - # logger.info("All tasks complete, extracting results from shared memory") - # - # # Extract results from shared memory into output arrays - # # Main outputs - # chunk_main = np.concatenate( - # [ - # output_array[0:quarter], - # output_array[quarter : 2 * quarter], - # output_array[2 * quarter : 3 * quarter], - # output_array[3 * quarter : 4 * quarter], - # ], - # axis=0, - # ) - # - # # False outputs - # chunk_false = np.concatenate( - # [ - # output_array[chunk_size : chunk_size + half], - # output_array[chunk_size + half : chunk_size + 2 * half], - # ], - # axis=0, - # ) - # - # # True outputs - # chunk_true = np.concatenate( - # [ - # output_array[2 * chunk_size : 2 * chunk_size + half], - # output_array[2 * chunk_size + half : 3 * chunk_size], - # ], - # axis=0, - # ) - # - # # Store chunks in pre-allocated arrays - # all_main[start_idx:end_idx] = chunk_main - # all_false[start_idx:end_idx] = chunk_false - # all_true[start_idx:end_idx] = chunk_true - # - # # Cleanup chunk resources - # del chunk_main, chunk_false, chunk_true - # del output_array - # - # # Close pool before closing shared memory (workers have references) - # self._free_managed_pool() - # self.manager.close_shared_memory(output_shm) - # - # gc.collect() - # - # logger.info(f"Chunk {chunk_idx + 1} complete, memory cleared") - # - # # Recreate pool for next generate_triplet_batch call - # self._setup_managed_pool() - # - # result = {"concatenated": all_main, "false": all_false, "true": all_true} - # - # return result + def _stats_cb(segment: dict) -> None: + write_segment_stats(self.db, tag, segment) + + return generate_round_to_memmap( + paths=paths, + n_samples=n_samples, + snr_base=snr_base, + snr_range=snr_range, + width_bin=self.width_bin, + num_observations=self._background_shape[1], + time_bins=self._background_shape[2], + chunk_size=self.config.training.signal_injection_chunk_size, + task_size=self.config.training.data_gen_task_size, + freq_resolution=self.freq_resolution, + time_resolution=self.time_resolution, + pool=self.pool, + backgrounds=self.backgrounds if self.pool is None else None, + round_num=round_num, + stats_cb=_stats_cb, + ) diff --git a/src/aetherscan/logger/logger.py b/src/aetherscan/logger/logger.py index 32762368..747fc917 100644 --- a/src/aetherscan/logger/logger.py +++ b/src/aetherscan/logger/logger.py @@ -293,24 +293,30 @@ def init_logger() -> Logger: return logger_instance -def init_worker_logging(): +def init_worker_logging(log_queue=None): """ Initialize logging in a multiprocessing worker: reset stdout/stderr to the original streams (so the inherited StreamToLogger from the parent doesn't fire), then attach a QueueHandler - pointing at the main process's shared log queue. If no Logger singleton has been initialized - in the parent, falls back to a NullHandler to avoid noise. - """ - logger_instance = Logger._instance + pointing at the main process's shared log queue. - if logger_instance is None: - logger.warning( - "No logger instance initialized - disabling worker logging to avoid conflicts" - ) - logging.getLogger().handlers.clear() - logging.getLogger().addHandler(logging.NullHandler()) - return + `log_queue` defaults to the Logger singleton's queue, which fork-started workers inherit. + Spawn-started processes (e.g. the RoundDataProducer and its pool workers) run a fresh + interpreter where the singleton doesn't exist — they must pass the queue explicitly + (a multiprocessing.Queue survives the spawn pickling boundary). Without either, falls back + to a NullHandler to avoid noise. + """ + if log_queue is None: + logger_instance = Logger._instance + + if logger_instance is None: + logger.warning( + "No logger instance initialized - disabling worker logging to avoid conflicts" + ) + logging.getLogger().handlers.clear() + logging.getLogger().addHandler(logging.NullHandler()) + return - log_queue = logger_instance.log_queue + log_queue = logger_instance.log_queue # Reset stdout/stderr to avoid inherited StreamToLogger from parent sys.stdout = sys.__stdout__ diff --git a/src/aetherscan/manager/manager.py b/src/aetherscan/manager/manager.py index 8a220e5e..afc38e7f 100644 --- a/src/aetherscan/manager/manager.py +++ b/src/aetherscan/manager/manager.py @@ -10,6 +10,7 @@ import atexit import contextlib import logging +import multiprocessing import os import signal import sys @@ -35,12 +36,66 @@ class ResourceStats: pools_active: int = 0 pools_closed: int = 0 + processes_active: int = 0 + processes_closed: int = 0 shared_memories_active: int = 0 shared_memories_cleaned: int = 0 total_memory_freed_gb: float = 0.0 cleanup_time_seconds: float = 0.0 +@dataclass +class ManagedProcess: + """Wrapper for a tracked multiprocessing.Process (e.g. the RoundDataProducer)""" + + process: multiprocessing.Process + name: str + created_at: float + closed: bool = False + + def close(self, timeout): + """Close the process with terminate -> join -> kill escalation (mirrors ManagedPool.close)""" + if self.closed: + return + + try: + if self.process.is_alive(): + logger.info(f"Terminating process '{self.name}' (PID {self.process.pid})") + + # SIGTERM first: the process's handler gets a chance to clean up (e.g. the + # producer terminates its own worker pool before dying) + self.process.terminate() + self.process.join(timeout) + + if self.process.is_alive(): + logger.warning(f"Process '{self.name}' survived SIGTERM, escalating to SIGKILL") + # SIGKILL gives the process no chance to reap its own children (e.g. the + # producer's pool workers), so kill any survivors in its subtree first — + # otherwise they'd be orphaned and keep running + with contextlib.suppress(Exception): + for child in psutil.Process(self.process.pid).children(recursive=True): + with contextlib.suppress(psutil.NoSuchProcess): + child.kill() + self.process.kill() + self.process.join(timeout) + + if self.process.is_alive(): + # Surviving SIGKILL is extremely rare (uninterruptible sleep state 'D'). + # Log and proceed — the OS will clean up on exit + logger.error( + f"Process '{self.name}' survived SIGKILL (uninterruptible state?)" + ) + + self.closed = True + logger.info(f"Process '{self.name}' closed") + + except Exception as e: + logger.warning(f"Error terminating process '{self.name}': {e}") + with contextlib.suppress(Exception): + self.process.kill() + self.closed = True + + @dataclass class ManagedPool: """Wrapper for tracked multiprocessing Pool""" @@ -305,6 +360,7 @@ def __init__(self): # NOTE: should these be strong or weak references? import weakref ... # Track resources self._pools: list[ManagedPool] = [] + self._processes: list[ManagedProcess] = [] self._shared_memories: list[ManagedSharedMemory] = [] # Track threads @@ -409,11 +465,13 @@ def cleanup_all(self): """ Unified cleanup of all resources. Strict order: - 1. Pools - 2. SharedMemory - 3. Monitor - 4. DB - 5. Logger + 1. Processes (before shared memory — e.g. the producer's workers attach to the + background-plate block) + 2. Pools + 3. SharedMemory + 4. Monitor + 5. DB + 6. Logger """ # Skip if not main process if os.getpid() != self._main_process_pid: @@ -439,6 +497,12 @@ def cleanup_all(self): start_time = time.time() initial_memory = self._get_memory_usage() + # Close managed processes + logger.info("Closing managed processes...") + for managed in list(self._processes): + if not managed.closed: + self._close_managed_process(managed) + # Close managed pools logger.info("Closing multiprocessing pools...") for managed in self._pools: @@ -486,6 +550,7 @@ def cleanup_all(self): logger.info("=" * 60) logger.info("ResourceManager cleanup complete:") logger.info(f" Pools closed: {self.stats.pools_closed}") + logger.info(f" Processes closed: {self.stats.processes_closed}") logger.info(f" Shared memories cleaned: {self.stats.shared_memories_cleaned}") logger.info(f" Memory freed: {self.stats.total_memory_freed_gb:.2f} GB") logger.info(f" Cleanup time: {self.stats.cleanup_time_seconds:.2f} seconds") @@ -548,6 +613,35 @@ def _close_managed_pool(self, managed: ManagedPool): self.stats.pools_active -= 1 self.stats.pools_closed += 1 + def register_process(self, process: multiprocessing.Process, name: str = "unnamed"): + """ + Register an externally-created multiprocessing.Process (e.g. the RoundDataProducer) + for lifecycle tracking, so cleanup_all() can escalate terminate -> join -> kill on it. + """ + managed = ManagedProcess(process=process, name=name, created_at=time.time()) + + self._processes.append(managed) + self.stats.processes_active += 1 + + logger.info(f"Registered process '{name}' (PID {process.pid})") + logger.info(f" Current total active: {self.stats.processes_active}") + + def close_process(self, process: multiprocessing.Process): + """Explicitly close a specific tracked process by reference""" + for managed in self._processes: + if managed.process is process and not managed.closed: + self._close_managed_process(managed) + return + logger.warning("ResourceManager.close_process(): Process not found in managed processes") + + def _close_managed_process(self, managed: ManagedProcess): + """Internal method to close a ManagedProcess""" + managed.close(timeout=self.config.manager.pool_terminate_timeout) + # Remove from tracking list to allow garbage collection of the Process object + self._processes.remove(managed) + self.stats.processes_active -= 1 + self.stats.processes_closed += 1 + def create_shared_memory(self, size: int, name: str = "unnamed") -> SharedMemory: """ Allocate a POSIX shared-memory block of `size` bytes and register it for lifecycle diff --git a/src/aetherscan/round_data.py b/src/aetherscan/round_data.py new file mode 100644 index 00000000..01c4bbc6 --- /dev/null +++ b/src/aetherscan/round_data.py @@ -0,0 +1,640 @@ +""" +Disk-backed (memmap) round datasets for Aetherscan training. + +Each training round's data lives on disk as .npy memmaps under +{round_data_dir}/{save_tag}/round_{k:02d}/ instead of ~294 GB of main-process RAM: +workers write generated cadences straight into the memmaps, training reads them back +through np.load(mmap_mode="r"), and the OS page cache keeps steady-state reads at +RAM speed while remaining evictable under memory pressure (no more OOM kills). + +This module owns: +- RoundDataPaths: the on-disk layout of one round's dataset. +- The atomic `.done` manifest protocol (written only after all workers finish, via + .tmp -> os.replace like preprocessing's _extract_stamps_worker). +- Startup archive/reuse/delete semantics for a tag's round-data directory. +- RoundDataProducer: a dedicated process that generates round k+1 while round k + trains, streaming stats back to the main process (DB writes stay in main — the + DB queue is a thread queue.Queue, not process-safe). +""" + +from __future__ import annotations + +import contextlib +import json +import logging +import math +import multiprocessing +import os +import queue +import re +import shutil +import signal +import threading +import time +import traceback +from dataclasses import dataclass +from logging.handlers import QueueListener + +import numpy as np + +from aetherscan.logger import init_worker_logging +from aetherscan.manager import get_manager + +logger = logging.getLogger(__name__) + +# The producer process is started with the SPAWN method, not fork: the training parent holds +# deep TF / NCCL / gRPC / CUDA thread state whose locks a forked child can inherit in a locked +# state — an early producer prototype deadlocked on a futex before reaching its own code. A +# spawned child runs a fresh interpreter with none of that baggage (its own pool workers are +# then forked from the clean single-threaded producer, which is safe). +_MP_CONTEXT = multiprocessing.get_context("spawn") + +# Directory-name pattern for one round's dataset (1-based, mirrors checkpoint round_XX tags) +_ROUND_DIR_PATTERN = re.compile(r"^round_(\d+)$") + +# Number of values sampled per array for the manifest's cheap corruption check +_MANIFEST_SAMPLE_COUNT = 8 + + +@dataclass(frozen=True) +class RoundDataPaths: + """On-disk layout of one round's dataset: {main,true,false,labels}.npy + a .done manifest.""" + + round_dir: str + round_idx: int + + @classmethod + def for_round(cls, base_dir: str, round_idx: int) -> RoundDataPaths: + """Build the paths for 1-based round `round_idx` under `base_dir` (round_{k:02d}).""" + return cls(round_dir=os.path.join(base_dir, f"round_{round_idx:02d}"), round_idx=round_idx) + + @property + def main_path(self) -> str: + return os.path.join(self.round_dir, "main.npy") + + @property + def true_path(self) -> str: + return os.path.join(self.round_dir, "true.npy") + + @property + def false_path(self) -> str: + return os.path.join(self.round_dir, "false.npy") + + @property + def labels_path(self) -> str: + return os.path.join(self.round_dir, "labels.npy") + + @property + def done_path(self) -> str: + return os.path.join(self.round_dir, f"round_{self.round_idx:02d}.done") + + @property + def array_paths(self) -> dict[str, str]: + """Mapping of array name -> .npy path for the three cadence arrays.""" + return {"main": self.main_path, "true": self.true_path, "false": self.false_path} + + +def _array_checksum(arr: np.ndarray) -> dict: + """ + Cheap, sha-less checksum for a (possibly huge) memmap: total element count plus a few + deterministically-sampled flat-index values. Positions are derived from the element count, + so validation re-samples the same spots without storing an RNG state. + """ + # NOTE: reshape(-1) is a zero-copy view here (and in validate_done_manifest) because + # open_memmap/np.save arrays are always C-contiguous; on a non-contiguous array it would + # silently materialize a full in-RAM copy — switch to .ravel()/.flat if the layout ever changes + flat = arr.reshape(-1) + n = int(flat.size) + rng = np.random.default_rng(n) + positions = sorted({int(i) for i in rng.integers(0, n, size=min(_MANIFEST_SAMPLE_COUNT, n))}) + if arr.dtype.kind in ("U", "S"): + samples = [[i, str(flat[i])] for i in positions] + else: + samples = [[i, float(flat[i])] for i in positions] + return {"elements": n, "samples": samples} + + +def _checksum_value_matches(recorded, actual) -> bool: + """Compare one manifest sample against the on-disk value (NaN == NaN counts as a match).""" + if isinstance(recorded, str): + return recorded == str(actual) + actual = float(actual) + return recorded == actual or (math.isnan(recorded) and math.isnan(actual)) + + +def build_manifest( + paths: RoundDataPaths, + n_samples: int, + snr_base: float, + snr_range: float, + wall_time_s: float, + chunk_count: int, +) -> dict: + """Assemble the .done manifest dict by re-opening the finished arrays read-only.""" + shapes: dict[str, list[int]] = {} + checksums: dict[str, dict] = {} + for name, path in {**paths.array_paths, "labels": paths.labels_path}.items(): + arr = np.load(path, mmap_mode="r") + shapes[name] = list(arr.shape) + checksums[name] = _array_checksum(arr) + del arr + return { + "round_idx": paths.round_idx, + "n_samples": int(n_samples), + "shapes": shapes, + "snr_base": float(snr_base), + "snr_range": float(snr_range), + "checksums": checksums, + "wall_time_s": float(wall_time_s), + "chunk_count": int(chunk_count), + "created_at": time.time(), + } + + +def write_done_manifest(paths: RoundDataPaths, manifest: dict) -> None: + """Write the manifest atomically (.tmp -> os.replace) so a crash can't leave a partial + .done file — a round dir either has a complete manifest or is treated as garbage.""" + tmp_path = paths.done_path + ".tmp" + with open(tmp_path, "w") as f: + json.dump(manifest, f, indent=2) + os.replace(tmp_path, paths.done_path) + logger.info(f"Wrote round-data manifest: {paths.done_path}") + + +def validate_done_manifest( + paths: RoundDataPaths, expected_n_samples: int | None = None +) -> dict | None: + """ + Validate a round dir's .done manifest against the arrays on disk. Returns the manifest + dict when everything checks out (round index, optional expected sample count, per-array + existence, shape, element count, and the sampled checksum values), else None. + """ + try: + if not os.path.isfile(paths.done_path): + return None + with open(paths.done_path) as f: + manifest = json.load(f) + + if manifest.get("round_idx") != paths.round_idx: + logger.warning(f"Manifest round_idx mismatch in {paths.done_path}") + return None + if expected_n_samples is not None and manifest.get("n_samples") != expected_n_samples: + logger.warning( + f"Manifest n_samples ({manifest.get('n_samples')}) != expected " + f"({expected_n_samples}) in {paths.done_path}" + ) + return None + + all_paths = {**paths.array_paths, "labels": paths.labels_path} + for name, path in all_paths.items(): + if not os.path.isfile(path): + logger.warning(f"Manifest array missing on disk: {path}") + return None + arr = np.load(path, mmap_mode="r") + try: + if list(arr.shape) != manifest["shapes"][name]: + logger.warning(f"Shape mismatch for {path}") + return None + checksum = manifest["checksums"][name] + if int(arr.size) != checksum["elements"]: + logger.warning(f"Element-count mismatch for {path}") + return None + flat = arr.reshape(-1) + for position, recorded in checksum["samples"]: + if not _checksum_value_matches(recorded, flat[position]): + logger.warning(f"Sampled-value mismatch for {path} at {position}") + return None + finally: + del arr + + return manifest + except Exception as e: + logger.warning(f"Failed to validate round-data manifest {paths.done_path}: {e}") + return None + + +def load_round_arrays(paths: RoundDataPaths) -> dict[str, np.ndarray]: + """ + Open one round's arrays for training. The three cadence arrays come back as read-only + memmaps (np.load(mmap_mode="r")) so nothing is pulled into RAM until batches gather it — + the OS page cache keeps hot pages resident and evicts them under memory pressure instead + of OOM-killing the process. Labels are tiny and loaded eagerly. + """ + return { + "concatenated": np.load(paths.main_path, mmap_mode="r"), + "true": np.load(paths.true_path, mmap_mode="r"), + "false": np.load(paths.false_path, mmap_mode="r"), + "labels": np.load(paths.labels_path), + } + + +def prepare_round_data_dir(base_dir: str, start_round: int) -> None: + """ + Startup cleanup of a tag's round-data directory, mirroring checkpoint-archiving semantics + (train.archive_directory) for resumable runs — except entries are deleted rather than + archived (a round is ~295 GB; archiving would silently double the disk budget): + + - round dirs with index >= start_round are deleted (regenerated by this run), + - round dirs with index < start_round are kept only if their .done manifest validates + (reusable instead of regenerated); .done-less or corrupt dirs are deleted, + - any non-round entry (e.g. a stale RF dataset dir) is deleted. + """ + os.makedirs(base_dir, exist_ok=True) + + for entry in sorted(os.listdir(base_dir)): + entry_path = os.path.join(base_dir, entry) + if not os.path.isdir(entry_path): + continue + match = _ROUND_DIR_PATTERN.match(entry) + if match is None: + logger.info(f"Deleting stale round-data entry: {entry_path}") + shutil.rmtree(entry_path, ignore_errors=True) + continue + round_idx = int(match.group(1)) + if round_idx >= start_round: + logger.info(f"Deleting round-data dir for round >= {start_round}: {entry_path}") + shutil.rmtree(entry_path, ignore_errors=True) + elif validate_done_manifest(RoundDataPaths(entry_path, round_idx)) is None: + logger.info(f"Deleting round-data dir with missing/invalid manifest: {entry_path}") + shutil.rmtree(entry_path, ignore_errors=True) + else: + logger.info(f"Keeping validated round-data dir: {entry_path}") + + +# --------------------------------------------------------------------------- +# RoundDataProducer — background generation process +# --------------------------------------------------------------------------- + +# Producer-process-local handle to its worker pool, so the SIGTERM handler can +# terminate the pool's children before the producer itself dies. +_PRODUCER_POOL = None + + +def _default_generate(paths, round_idx, snr_base, snr_range, pool, params, stats_cb, progress_cb): + """Real generation entry point for the producer process (test hooks can replace it).""" + # Deferred import: data_generation pulls in setigen/scipy, which the parent may not need + # at round_data import time (cli.py must stay stdlib-importable via config/cli only). + from aetherscan.data_generation import generate_round_to_memmap # noqa: PLC0415 + + return generate_round_to_memmap( + paths=paths, + n_samples=params["n_samples"], + snr_base=snr_base, + snr_range=snr_range, + width_bin=params["width_bin"], + num_observations=params["num_observations"], + time_bins=params["time_bins"], + chunk_size=params["chunk_size"], + task_size=params["task_size"], + freq_resolution=params["freq_resolution"], + time_resolution=params["time_resolution"], + pool=pool, + round_num=round_idx, + stats_cb=stats_cb, + progress_cb=progress_cb, + ) + + +def _producer_main( + request_queue, result_queue, params: dict, generate_fn=None, log_queue=None +) -> None: + """ + Producer process entry point (spawn-started — see _MP_CONTEXT). Owns a private worker pool + (workers attach to the background-plate shared memory created by the main process) and + generates rounds on request, isolated from the main process's GIL (TF's prefetch threads no + longer slow generation down, and generation no longer steals cycles from training). + + Protocol (multiprocessing.Queues): + - in: ("generate", round_idx, snr_base, snr_range) | ("shutdown",) + - out: ("stats", round_idx, segment_dict) per class-segment per chunk + ("progress", round_idx, chunk, n_chunks) per chunk + ("done", round_idx, manifest) on success (or valid reuse) + ("error", round_idx, traceback_str) on failure (producer keeps serving) + ("shutdown_ack",) right before a graceful exit + + `log_queue` is a spawn-context multiprocessing queue relayed into the main process's + logging pipeline by RoundDataProducer.start() — a spawned child has no inherited Logger + singleton, so it (and its pool workers) attach QueueHandlers to this queue explicitly. + `generate_fn` is a test seam: when provided, no worker pool is created and the stub is + called in place of the real memmap generation (see tests/unit/test_round_data.py). + """ + global _PRODUCER_POOL + + is_main_thread = threading.current_thread() is threading.main_thread() + if is_main_thread: + # The log queue is a multiprocessing.Queue — safe to use from this process. + init_worker_logging(log_queue) + + # Ignore SIGINT: the main process's ResourceManager coordinates Ctrl-C cleanup. + signal.signal(signal.SIGINT, signal.SIG_IGN) + + # SIGTERM handler mirrors data_generation._init_worker's: + # WARN: DO NOT LOG ANYTHING IN THIS HANDLER — the log QueueHandler's feeder thread + # needs the GIL, which may be held elsewhere, and a blocked handler deadlocks + # termination (see data_generation.py's canonical worker handler). + def _cleanup_on_sigterm(signum, frame): + with contextlib.suppress(Exception): + if _PRODUCER_POOL is not None: + # Forward termination to the pool's children so they don't outlive us + # (SIGKILL on this process would otherwise orphan them). + _PRODUCER_POOL.terminate() + signal.signal(signal.SIGTERM, signal.SIG_DFL) + os.kill(os.getpid(), signal.SIGTERM) + + signal.signal(signal.SIGTERM, _cleanup_on_sigterm) + + pool = None + if generate_fn is None: + generate_fn = _default_generate + # Deferred import (setigen/scipy) — only the real producer needs it. + from aetherscan.data_generation import _init_worker # noqa: PLC0415 + + # Plain fork-context Pool: this spawned process is single-threaded, so forking its + # workers is safe — and they attach to the parent-created background SHM by name. + pool = multiprocessing.Pool( + processes=max(1, params["n_processes"]), + initializer=_init_worker, + initargs=( + params["shm_name"], + tuple(params["background_shape"]), + np.dtype(params["background_dtype"]), + log_queue, + ), + ) + _PRODUCER_POOL = pool + + logger.info(f"RoundDataProducer started (PID {os.getpid()})") + + try: + while True: + message = request_queue.get() + if message[0] == "shutdown": + logger.info("RoundDataProducer received shutdown request") + result_queue.put(("shutdown_ack",)) + break + if message[0] != "generate": + logger.warning(f"RoundDataProducer ignoring unknown message: {message[0]!r}") + continue + + _, round_idx, snr_base, snr_range = message + paths = RoundDataPaths.for_round(params["base_dir"], round_idx) + + existing = validate_done_manifest(paths, expected_n_samples=params["n_samples"]) + if existing is not None: + logger.info(f"RoundDataProducer: reusing validated round {round_idx} data") + result_queue.put(("done", round_idx, existing)) + continue + + try: + logger.info( + f"RoundDataProducer: generating round {round_idx} data " + f"(SNR {snr_base}-{snr_base + snr_range})" + ) + manifest = generate_fn( + paths, + round_idx, + snr_base, + snr_range, + pool, + params, + lambda segment, _r=round_idx: result_queue.put(("stats", _r, segment)), + lambda chunk, n_chunks, _r=round_idx: result_queue.put( + ("progress", _r, chunk, n_chunks) + ), + ) + logger.info(f"RoundDataProducer: round {round_idx} data complete") + result_queue.put(("done", round_idx, manifest)) + except Exception: + result_queue.put(("error", round_idx, traceback.format_exc())) + finally: + _PRODUCER_POOL = None + if pool is not None: + with contextlib.suppress(Exception): + pool.terminate() + pool.join() + + +class RoundDataProducer: + """ + Main-process handle for the background round-data generation process. + + Lifecycle: start() spawns the producer process (registered with ResourceManager for + terminate -> join -> kill cleanup) plus a drainer thread that consumes the result queue — + writing streamed injection stats to the DB from the main process (the DB writer queue is + a thread queue.Queue, not process-safe) while the GPUs train. request_generation() is + fire-and-forget; await_round() blocks until that round's done/error message arrives. + """ + + def __init__( + self, + *, + base_dir: str, + n_samples: int, + shm_name: str, + background_shape: tuple, + background_dtype: str, + n_processes: int, + width_bin: int, + num_observations: int, + time_bins: int, + chunk_size: int, + task_size: int, + freq_resolution: float, + time_resolution: float, + db, + tag: str, + ): + self._params = { + "base_dir": base_dir, + "n_samples": n_samples, + "shm_name": shm_name, + "background_shape": tuple(background_shape), + "background_dtype": str(background_dtype), + "n_processes": n_processes, + "width_bin": width_bin, + "num_observations": num_observations, + "time_bins": time_bins, + "chunk_size": chunk_size, + "task_size": task_size, + "freq_resolution": freq_resolution, + "time_resolution": time_resolution, + } + self._db = db + self._tag = tag + # Queues come from the spawn context so they can cross the spawn pickling boundary + # (mixing contexts raises "A SemLock created in a fork context is being shared with a + # process in a spawn context" — which also rules out handing the child the Logger + # singleton's fork-context queue directly; see the relay in start()) + self._request_queue: multiprocessing.Queue = _MP_CONTEXT.Queue() + self._result_queue: multiprocessing.Queue = _MP_CONTEXT.Queue() + self._producer_log_queue: multiprocessing.Queue = _MP_CONTEXT.Queue() + self._log_relay: QueueListener | None = None + self._process: multiprocessing.Process | None = None + self._drainer: threading.Thread | None = None + self._drainer_done = False + self._results: dict[int, tuple[str, object]] = {} + self._condition = threading.Condition() + + def start(self) -> None: + """Spawn the producer process and the main-side result drainer thread.""" + # A spawned child has no inherited Logger singleton, and the singleton's fork-context + # queue can't cross the spawn boundary — so the producer tree logs into its own + # spawn-context queue, and this main-side QueueListener relays each record into the + # main process's root handlers (i.e. into the normal file/console/Slack pipeline). + self._log_relay = QueueListener( + self._producer_log_queue, + *logging.getLogger().handlers, + respect_handler_level=False, + ) + self._log_relay.start() + + self._process = _MP_CONTEXT.Process( + target=_producer_main, + args=( + self._request_queue, + self._result_queue, + self._params, + None, + self._producer_log_queue, + ), + name="RoundDataProducer", + ) + + # The spawned child re-imports the aetherscan module chain, which includes TF; blank + # out GPU visibility for the child's entire process tree so nothing in the producer + # can ever initialize CUDA (generation is pure CPU). + # NOTE: another thread reading os.environ during this brief window would see the + # blanked value — intentionally acceptable ONLY because the parent's TF read the var + # once at its own GPU init (long before this point) and nothing else in the pipeline + # spawns GPU-visible subprocesses concurrently with training startup. + original_cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") + os.environ["CUDA_VISIBLE_DEVICES"] = "" + try: + self._process.start() + finally: + if original_cuda_visible is None: + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + else: + os.environ["CUDA_VISIBLE_DEVICES"] = original_cuda_visible + + manager = get_manager() + if manager is not None: + manager.register_process(self._process, name="RoundDataProducer") + + self._drainer = threading.Thread( + target=self._drain_results, name="RoundDataDrainer", daemon=True + ) + self._drainer.start() + logger.info(f"RoundDataProducer process started (PID {self._process.pid})") + + def request_generation(self, round_idx: int, snr_base: float, snr_range: float) -> None: + """Queue generation of 1-based round `round_idx` (non-blocking).""" + logger.info(f"Requesting background generation of round {round_idx} data") + self._request_queue.put(("generate", round_idx, snr_base, snr_range)) + + def await_round(self, round_idx: int) -> dict: + """ + Block until round `round_idx`'s done/error message arrives; return the manifest on + success, raise RuntimeError (carrying the producer traceback) on generation failure + or if the producer died without answering. + """ + while True: + with self._condition: + if round_idx in self._results: + kind, payload = self._results[round_idx] + break + if self._drainer_done: + raise RuntimeError( + f"RoundDataProducer exited before producing round {round_idx} data" + ) + # NOTE: the drainer's notify_all() wakes this immediately on done/error; the + # 1 s timeout is not a polling latency, just a liveness re-check so a drainer + # that died without notifying can't strand this wait forever + self._condition.wait(timeout=1.0) + + if kind == "error": + raise RuntimeError(f"Round {round_idx} data generation failed in producer:\n{payload}") + return payload + + def shutdown(self, timeout: float = 10.0) -> None: + """Request a graceful producer exit, escalating to terminate/kill via the manager's + ManagedProcess if the producer doesn't wind down in time (e.g. mid-generation).""" + if self._process is None: + return + + with contextlib.suppress(Exception): + self._request_queue.put(("shutdown",)) + self._process.join(timeout) + + manager = get_manager() + if manager is not None: + # No-op join if already exited; terminate -> join -> kill escalation otherwise. + manager.close_process(self._process) + elif self._process.is_alive(): + self._process.terminate() + self._process.join(timeout) + + if self._drainer is not None: + self._drainer.join(timeout=timeout) + if self._log_relay is not None: + with contextlib.suppress(Exception): + self._log_relay.stop() + self._log_relay = None + self._process = None + logger.info("RoundDataProducer shut down") + + def _handle_message(self, message: tuple) -> None: + """Dispatch one producer message (runs on the drainer thread).""" + kind = message[0] + if kind == "stats": + # Deferred import: keep round_data importable without the setigen stack. + from aetherscan.data_generation import write_segment_stats # noqa: PLC0415 + + _, _round_idx, segment = message + write_segment_stats(self._db, self._tag, segment) + elif kind == "progress": + _, round_idx, chunk, n_chunks = message + logger.info(f"Round {round_idx} data generation: chunk {chunk}/{n_chunks} complete") + elif kind in ("done", "error"): + _, round_idx, payload = message + with self._condition: + self._results[round_idx] = (kind, payload) + self._condition.notify_all() + else: + logger.warning(f"RoundDataDrainer ignoring unknown message: {kind!r}") + + def _drain_results(self) -> None: + """Drainer thread body: consume producer messages until shutdown-ack or producer death, + then sweep any messages still buffered in the queue. A failure handling one message + (e.g. a DB hiccup on a stats write) must not kill the drainer — done/error messages + for in-flight rounds still need to unblock await_round().""" + + def _handle_safely(message: tuple) -> None: + try: + self._handle_message(message) + except Exception as e: + logger.error(f"RoundDataDrainer failed to handle {message[0]!r} message: {e}") + + while True: + try: + message = self._result_queue.get(timeout=1.0) + except queue.Empty: + if self._process is None or not self._process.is_alive(): + break + continue + if message[0] == "shutdown_ack": + break + _handle_safely(message) + + # Final sweep: the producer may have exited with messages still in the pipe. + while True: + try: + message = self._result_queue.get_nowait() + except queue.Empty: + break + if message[0] != "shutdown_ack": + _handle_safely(message) + + with self._condition: + self._drainer_done = True + self._condition.notify_all() diff --git a/src/aetherscan/train.py b/src/aetherscan/train.py index 56c2e989..59055a8c 100644 --- a/src/aetherscan/train.py +++ b/src/aetherscan/train.py @@ -53,6 +53,13 @@ create_beta_vae_model, prepare_latent_features, ) +from aetherscan.round_data import ( + RoundDataPaths, + RoundDataProducer, + load_round_arrays, + prepare_round_data_dir, + validate_done_manifest, +) logger = logging.getLogger(__name__) @@ -335,24 +342,22 @@ def check_encoder_trained(encoder, threshold=0.2): # Create data holder objects, to be paired with data generators, for TF's distributed datasets -# Allows for explicit dereferencing of large arrays using holder.clear(), which lets +# Allows for explicit dereferencing of the backing arrays using holder.clear(), which lets # Python's garbage collector free up memory on-demand -# Note, holder.clear() is only useful at the end of an epoch, once indices have been exhausted, -# since the data generators' local caches maintain references to the data until then -# This is not an issue in our current implementation, where we only clear & reset resources at the -# end of a round. However, if you require early exit behavior, you may want to remove the _lock and -# use explicit _cleared() checks instead, which negates the need for local caches (see commit hash -# 2a404a4). The trade-off being that you're at risk of race conditions if multiple threads attempt -# to access/clear the holder simultaneously. While this is not the case in our current -# implementation, we opted for a more defensive approach rather than accomodating future design -# patterns. As well, the data should not be modified once the holder has been initialized to -# prevent corrupted state in the holder -# Note, there's a potential deadlock issue with holder lock contention -# Since the generators acquire locks at the start of every loop iteration, if TF's prefetch threads -# (.prefetch(tf.data.AUTOTUNE)) are blocked waiting on this lock while the main thread is trying to -# call holder.clear() (which also needs the lock), there could be contention. -# This has not been an issue so far, but if you encounter this in the future, pls update this -# comment with your findings +# The holders now hold np.load(mmap_mode="r") memmap references rather than ~294 GB of in-RAM +# arrays, so clear() drops file mappings (letting the round's directory be deleted and its page +# cache reclaimed) rather than freeing huge heap allocations — but the semantics are unchanged: +# clear() is only fully effective at the end of an epoch, once the generators' local caches have +# been dropped, and the data must not be modified once the holder has been initialized +# Note, if you require early exit behavior, you may want to remove the _lock and use explicit +# _cleared() checks instead, which negates the need for local caches (see commit hash 2a404a4). +# The trade-off being that you're at risk of race conditions if multiple threads attempt to +# access/clear the holder simultaneously — we opted for the defensive approach +# Lock contention with TF's prefetch threads is a non-issue in the batched design: each generator +# acquires the lock exactly once per epoch pass (to snapshot the array references), then yields +# whole global batches without touching the lock, so a clear() on the main thread can no longer +# race hundreds of thousands of per-sample lock acquisitions (which the pre-memmap, per-sample +# generators were at least theoretically exposed to) class TrainDataHolder: def __init__(self, concat, true, false): self._cleared = False @@ -396,26 +401,42 @@ def prepare_distributed_train_dataset( shuffle: bool = True, ) -> dict: """ - Build distributed training and validation datasets from the in-memory `data` dict, returning - a dict with the two tf.data datasets, sample/step counts, the shared TrainDataHolder, and the + Build distributed training and validation datasets from the `data` dict, returning a dict + with the two tf.data datasets, sample/step counts, the shared TrainDataHolder, and the stratified train/val indices into the original arrays. - `data` must contain 'concatenated', 'true', 'false', and 'labels' (numpy arrays). The split - is stratified across the 4 signal types (false_no_signal, false_with_rfi, true_only_eti, - true_eti_rfi) — generate_triplet_batch arranges labels sequentially within chunks, so a naive - positional split would over-represent later signal types in val. Each generated batch has the - signature ((concat, true, false), concat). Sample counts are trimmed to the global / effective - batch size to keep all replicas evenly fed; the holder is shared by both generators so neither - pays a memory cost beyond index subsets. + `data` must contain 'concatenated', 'true', 'false', and 'labels' — typically the read-only + memmaps returned by round_data.load_round_arrays(), though plain in-RAM ndarrays work too. + The split is stratified across the 4 signal types (false_no_signal, false_with_rfi, + true_only_eti, true_eti_rfi) — generation lays labels out sequentially within chunks, so a + naive positional split would over-represent later signal types in val. + + The generators yield whole GLOBAL batches (leading batch dim in the output signature; no + .batch() call downstream): each batch gathers sorted fancy indices from the memmaps, cutting + the per-sample Python/tf.data boundary crossings by a factor of the global batch size — the + per-sample yields were the main source of the 0-14 % GPU utilization. Randomness lives at + the epoch level (train_indices are reshuffled each pass); sorting *within* a batch only + improves memmap read locality and the model is order-invariant within a batch. + + Page-cache framing: gathering from the memmaps pulls pages through the OS page cache, so + after the first epoch a round's ~294 GB (at full-scale defaults) is served at RAM speed from + otherwise-free memory on the 503 GB training nodes — but under memory pressure the kernel + evicts pages instead of OOM-killing the process, which is exactly the failure mode the old + in-RAM arrays hit. + + Each generated batch has the signature ((concat, true, false), concat). Sample counts are + trimmed to the global / effective batch size to keep all replicas evenly fed (so every epoch + pass yields whole batches exactly); the holder is shared by both generators so neither pays + a memory cost beyond index subsets. """ global_train_batch_size = per_replica_batch_size * num_replicas global_val_batch_size = per_replica_val_batch_size * num_replicas # Stratified train/val split to ensure both sets contain proportional representation # of all 4 signal types (false_no_signal, false_with_rfi, true_only_eti, true_eti_rfi). - # This is necessary because generate_triplet_batch() arranges labels sequentially within + # This is necessary because generation arranges labels sequentially within # chunks, so a naive positional split would over-represent later signal types in val. - labels = data["labels"] + labels = np.asarray(data["labels"]) unique_labels = np.unique(labels) train_indices = [] @@ -450,15 +471,25 @@ def prepare_distributed_train_dataset( if n_val_trimmed < n_val: val_indices = np.random.choice(val_indices, size=n_val_trimmed, replace=False) + # Sort both index sets ascending. For shuffle=False this pins the generators' yield order + # to the returned train_indices/val_indices arrays (the alignment contract + # train_random_forest depends on) while giving monotone memmap reads; for shuffle=True the + # epoch-level reshuffle below supplies the randomness anyway. Stratification is a property + # of index *membership*, not order, so sorting doesn't affect it. + train_indices = np.sort(train_indices) + val_indices = np.sort(val_indices) + logger.info(f"Data alignment: Train {n_train}→{n_train_trimmed}, Val {n_val}→{n_val_trimmed}") # Share the original arrays between train and val generators via a single data holder. # The stratified split requires non-contiguous indices, which would force numpy to create - # full copies via fancy indexing (~2x peak memory). Instead, both generators read from the - # same original arrays using their respective index subsets — zero extra copies. + # full copies via fancy indexing (~2x peak memory). Instead, both generators gather + # per-batch slices from the same original arrays using their respective index subsets — + # only one global batch is materialized at a time. train_holder = TrainDataHolder(data["concatenated"], data["true"], data["false"]) - # Create generator functions for memory-efficient data loading + # Create generator functions yielding whole global batches gathered from the (memmap) + # arrays — see the docstring for the batching/locality/page-cache rationale def train_generator(): while True: # Make generators infinite to reset state between epochs # Acquire lock to check cleared status and capture data references @@ -471,14 +502,20 @@ def train_generator(): true = train_holder.true false = train_holder.false - # Work with local references (safe from clearing, no per-sample lock needed) + # Work with local references (safe from clearing, no per-batch lock needed) # Copy train_indices because np.random.shuffle mutates in-place indices = train_indices.copy() if shuffle: # Perform global shuffle on each epoch so each pass through the data is unique np.random.shuffle(indices) - for idx in indices: - yield (concat[idx], true[idx], false[idx]), concat[idx] + for start in range(0, len(indices), global_train_batch_size): + batch_indices = indices[start : start + global_train_batch_size] + if shuffle: + # Within-batch sorted order improves memmap read locality; random batch + # membership is already guaranteed by the epoch-level shuffle above + batch_indices = np.sort(batch_indices) + concat_batch = concat[batch_indices] + yield (concat_batch, true[batch_indices], false[batch_indices]), concat_batch # Remove cache references to ensure garbage collection in future del concat, true, false @@ -495,44 +532,40 @@ def val_generator(): true = train_holder.true false = train_holder.false - # Maintain order on each epoch since shuffling provides no benefits (no gradients - # are calculated during validation) - for idx in val_indices: - yield (concat[idx], true[idx], false[idx]), concat[idx] + # Maintain val_indices order on each epoch (already sorted above): no gradients are + # calculated during validation, and train_random_forest relies on the i-th encoded + # val cadence corresponding to val_indices[i] + for start in range(0, len(val_indices), global_val_batch_size): + batch_indices = val_indices[start : start + global_val_batch_size] + concat_batch = concat[batch_indices] + yield (concat_batch, true[batch_indices], false[batch_indices]), concat_batch # Remove cache references to ensure garbage collection in future del concat, true, false - # Determine dataset output signature + # Determine dataset output signature: the generators yield whole global batches, so the + # specs carry a leading (None) batch dimension and no .batch() call is applied downstream sample_shape = data["concatenated"].shape[1:] - output_signature = ( - ( - tf.TensorSpec(shape=sample_shape, dtype=tf.float32), - tf.TensorSpec(shape=sample_shape, dtype=tf.float32), - tf.TensorSpec(shape=sample_shape, dtype=tf.float32), - ), - tf.TensorSpec(shape=sample_shape, dtype=tf.float32), - ) + batch_spec = tf.TensorSpec(shape=(None, *sample_shape), dtype=tf.float32) + output_signature = ((batch_spec, batch_spec, batch_spec), batch_spec) # Create datasets using generators to reduce GPU memory pressure # Data is kept on CPU & transferred to GPU in batches on-demand - # Note that the datasets yield data in batches before being sharded (distributed) across replicas - # Hence, we use global batch sizes here to ensure per replica batch sizes match expectations + # Note that the generators yield data in global batches before being sharded (distributed) + # across replicas, ensuring per replica batch sizes match expectations logger.info( - f"Creating infinite datasets from generators with global batch size - " + f"Creating infinite batched datasets from generators with global batch size - " f"Train: {global_train_batch_size}, Val: {global_val_batch_size}" ) train_dataset = ( tf.data.Dataset.from_generator(train_generator, output_signature=output_signature) - .batch(global_train_batch_size, drop_remainder=True) .repeat() .prefetch(tf.data.AUTOTUNE) ) val_dataset = ( tf.data.Dataset.from_generator(val_generator, output_signature=output_signature) - .batch(global_val_batch_size, drop_remainder=True) # NOTE: do we need repeat for val dataset? run test without repeat & see if anything breaks? .repeat() .prefetch(tf.data.AUTOTUNE) @@ -615,7 +648,9 @@ def prepare_distributed_viz_dataset( viz_holder = VizDataHolder(padded_data) - # Create generator function for memory-efficient data loading + # Create generator function yielding whole global batches — this feeds + # _capture_latent_snapshot every latent_viz_step_interval training steps, so per-sample + # yields here used to tax every capture during the epoch loop def viz_generator(): while True: # Make generator infinite to reset state between passes # Acquire lock to check cleared status and capture data references @@ -627,29 +662,29 @@ def viz_generator(): concat = viz_holder.concat # WARN: DO NOT SHUFFLE viz_generator(), OR ELSE YOU'LL BREAK plot_latent_space_gif() - # Maintain order on each epoch since shuffling provides no benefits (no gradients - # are calculated during inference) - for idx in range(len(concat)): - yield concat[idx] + # Contiguous in-order slices preserve the original cadence order on every pass + # (n_padded is an exact multiple of the global batch size, so slices are whole) + for start in range(0, len(concat), global_viz_batch_size): + yield concat[start : start + global_viz_batch_size] # Remove cache references for future garbage collection del concat - # Determine dataset output signature + # Determine dataset output signature: the generator yields whole global batches, so the + # spec carries a leading (None) batch dimension and no .batch() call is applied downstream sample_shape = padded_data.shape[1:] - output_signature = tf.TensorSpec(shape=sample_shape, dtype=tf.float32) + output_signature = tf.TensorSpec(shape=(None, *sample_shape), dtype=tf.float32) # Create dataset using generator to reduce GPU memory pressure # Data is kept on CPU & transferred to GPU in batches on-demand - # Note that the dataset yields data in batches before being sharded (distributed) across replicas - # Hence, we use global batch sizes here to ensure per replica batch sizes match expectations + # Note that the generator yields data in global batches before being sharded (distributed) + # across replicas, ensuring per replica batch sizes match expectations logger.info( - f"Creating infinite dataset from generator with global batch size: {global_viz_batch_size}" + f"Creating infinite batched dataset from generator with global batch size: {global_viz_batch_size}" ) viz_dataset = ( tf.data.Dataset.from_generator(viz_generator, output_signature=output_signature) - .batch(global_viz_batch_size, drop_remainder=True) # NOTE: do we need repeat for viz dataset? run test without repeat & see if anything breaks? .repeat() .prefetch(tf.data.AUTOTUNE) @@ -780,6 +815,10 @@ def __init__(self, background_data, strategy: tf.distribute.Strategy = None): # Initialize RF model as None self.rf_model = None + # Background round-data producer (created in train_beta_vae when + # overlap_data_generation is enabled; None otherwise) + self._round_producer: RoundDataProducer | None = None + # In-memory caches for RF eval artifacts and SHAP values, keyed by tag. # All ten RF plots consume the same eval-artifact joblib (and the five SHAP # plots additionally share a SHAP-values joblib); without these caches each @@ -866,6 +905,15 @@ def _setup_directories(self): plot_checkpoints_dir = os.path.join(self.config.output_path, "plots", "checkpoints") archive_directory(plot_checkpoints_dir, target_dirs=None, round_num=start_round) + # Disk-backed round-data directory for this tag: delete round dirs >= start_round, + # keep earlier ones only if their .done manifest validates (round-data mirror of the + # checkpoint archiving above, minus the archiving — a round is ~295 GB) + round_data_root = self.config.training.round_data_dir or os.path.join( + self.config.output_path, "round_data" + ) + self._round_data_base_dir = os.path.join(round_data_root, self.config.checkpoint.save_tag) + prepare_round_data_dir(self._round_data_base_dir, start_round) + logger.info("Setup directories complete") # COMMENTED OUT: Removing TensorBoard support @@ -920,6 +968,46 @@ def train_beta_vae(self): # NOTE: this approach doesn't play well with fault tolerance. rethink later self.start_time = time.time() + # Stand up the background round-data producer (a dedicated process owning its own + # worker pool against the shared-memory background plates), so round k+1's data + # generates while round k trains AND generation escapes this process's GIL — TF's + # prefetch/callback threads made round 2+ generation far slower than round 1's. + # Falls back to sequential in-process generation when disabled + # (--no-overlap-data-generation) or when the DataGenerator has no shared memory for + # producer workers to attach to (n_processes == 1). + if self.config.training.overlap_data_generation: + if self.data_generator.shm is not None: + self._round_producer = RoundDataProducer( + base_dir=self._round_data_base_dir, + n_samples=self.config.training.num_samples_beta_vae, + shm_name=self.data_generator.shm.name, + background_shape=self.data_generator._background_shape, + background_dtype=str(self.data_generator._background_dtype), + n_processes=self.data_generator.n_processes, + width_bin=self.data_generator.width_bin, + num_observations=self.config.data.num_observations, + time_bins=self.config.data.time_bins, + chunk_size=self.config.training.signal_injection_chunk_size, + task_size=self.config.training.data_gen_task_size, + freq_resolution=self.data_generator.freq_resolution, + time_resolution=self.data_generator.time_resolution, + db=self.db, + tag=self.config.checkpoint.save_tag, + ) + self._round_producer.start() + # Kick off the first round's data right away (nothing to overlap with yet — + # the unavoidable serial start) + first_snr_base, first_snr_range = self._calculate_curriculum_snr(start_round - 1) + self._round_producer.request_generation( + start_round, first_snr_base, first_snr_range + ) + else: + logger.warning( + "overlap_data_generation is enabled but DataGenerator is in sequential " + "mode (n_processes=1, no shared memory) — falling back to in-process " + "round data generation" + ) + try: for round_idx in range(start_round - 1, n_rounds): snr_base, snr_range = self._calculate_curriculum_snr(round_idx) @@ -945,6 +1033,12 @@ def train_beta_vae(self): round_idx=round_idx, epochs=epochs, snr_base=snr_base, snr_range=snr_range ) finally: + # Wind down the producer (graceful shutdown message, escalating to + # terminate -> kill through the ResourceManager if it's mid-generation) + if self._round_producer is not None: + self._round_producer.shutdown() + self._round_producer = None + # NOTE: this approach doesn't play well with fault tolerance. rethink later # Free the latent viz batch once all rounds are complete (or on failure) del self._latent_viz_batch, self._latent_viz_labels @@ -960,10 +1054,37 @@ def train_round(self, round_idx: int, epochs: int, snr_base: int, snr_range: int f"Training round {round_idx + 1} - Epochs: {epochs}, SNR: {snr_base}-{snr_base + snr_range}" ) - # Generate training data - train_data = self.data_generator.generate_triplet_batch( - self.config.training.num_samples_beta_vae, snr_base, snr_range, round_idx + 1 - ) + round_number = round_idx + 1 + n_samples = self.config.training.num_samples_beta_vae + paths = RoundDataPaths.for_round(self._round_data_base_dir, round_number) + round_trained = False # Set True once the round fully completes (drives dir deletion) + + # Obtain this round's disk-backed data: reuse a validated on-disk dataset if one + # exists, otherwise wait on the background producer (which was asked to generate it + # while the previous round trained) or generate in-process (overlap disabled) + if validate_done_manifest(paths, expected_n_samples=n_samples) is not None: + logger.info(f"Reusing validated round {round_number} data at {paths.round_dir}") + elif self._round_producer is not None: + logger.info(f"Waiting for round {round_number} data from the background producer") + wait_start = time.time() + self._round_producer.await_round(round_number) + logger.info(f"Round {round_number} data ready (waited {time.time() - wait_start:.1f}s)") + else: + self.data_generator.generate_round(paths, n_samples, snr_base, snr_range, round_number) + + # Immediately queue generation of the next round's data so it runs in the producer + # process while this round's epochs train (curriculum SNR for round k+1 is + # deterministic, so it can be computed ahead of time) + if ( + self._round_producer is not None + and round_number < self.config.training.num_training_rounds + ): + next_snr_base, next_snr_range = self._calculate_curriculum_snr(round_idx + 1) + self._round_producer.request_generation(round_number + 1, next_snr_base, next_snr_range) + + # Open the round's arrays as read-only memmaps (nothing is loaded into RAM here; the + # batched generators gather from the OS page cache during training) + train_data = load_round_arrays(paths) # Extract labels before distributing (prepare_distributed_train_dataset keeps the # original arrays alive via a shared train_holder — no copies — so we can free the @@ -1227,6 +1348,8 @@ def train_round(self, round_idx: int, epochs: int, snr_base: int, snr_range: int # Save checkpoint self.save_models(tag=f"round_{round_idx + 1:02d}", dir="checkpoints") + round_trained = True + except Exception as e: logger.error(f"Error in train_round(): {e}") raise # Re-raise to propagate error @@ -1260,6 +1383,16 @@ def train_round(self, round_idx: int, epochs: int, snr_base: int, snr_range: int gc.collect() + # Delete the round's on-disk data as soon as its training completed (keeps the + # disk footprint at ~2 rounds max with overlap). A failed round leaves its data + # in place; the retry's _setup_directories() -> prepare_round_data_dir() decides + # what survives (dirs >= the resume round are regenerated). Deleting after the + # holder.clear()/clear_session() above is safe even if stray memmap handles + # linger — POSIX keeps the inodes alive until the mappings drop. + if round_trained and not self.config.training.keep_round_data: + shutil.rmtree(paths.round_dir, ignore_errors=True) + logger.info(f"Deleted round {round_number} data directory: {paths.round_dir}") + def _train_epoch( self, round_idx, @@ -1827,9 +1960,20 @@ def train_random_forest(self): time_bins = self.config.data.time_bins width_bin = self.config.data.width_bin // self.config.data.downsample_factor - # Generate training data (concatenated is 4-way balanced; labels track per-sample subtype) + # Generate training data (concatenated is 4-way balanced; labels track per-sample + # subtype) into a disk-backed dataset alongside the per-round dirs. Generation is + # in-process (sequential with training) — there is nothing left to overlap with, the + # beta-VAE producer has already been shut down by train_beta_vae() logger.info(f"Preparing training set with SNR: {snr_base}-{snr_base + snr_range}") - rf_data = self.data_generator.generate_triplet_batch(n_samples, snr_base, snr_range) + rf_paths = RoundDataPaths( + round_dir=os.path.join(self._round_data_base_dir, "rf"), round_idx=0 + ) + rf_trained = False # Set True once RF training fully completes (drives dir deletion) + if validate_done_manifest(rf_paths, expected_n_samples=n_samples) is not None: + logger.info(f"Reusing validated RF dataset at {rf_paths.round_dir}") + else: + self.data_generator.generate_round(rf_paths, n_samples, snr_base, snr_range) + rf_data = load_round_arrays(rf_paths) # Prepare distributed train+val datasets (stratified split). shuffle=False so the # train generator yields in train_indices order, letting us align encoded features @@ -1983,6 +2127,8 @@ def encode_fn(data): joblib.dump(artifacts, artifact_path) logger.info(f"Saved RF eval artifacts to {artifact_path}") + rf_trained = True + # NOTE: come back to this later (are we dereferencing the correct things? can we instead write things to db instead of storing in memory?) del ( artifacts, @@ -2018,6 +2164,13 @@ def encode_fn(data): self.data_generator.reset_managed_pool() logger.info("Reset managed pools") + # Delete the RF dataset once RF training fully completed. On failure it stays on + # disk for post-mortem; the next run's prepare_round_data_dir() treats it as + # stale and regenerates (matching the pre-memmap per-attempt regeneration) + if rf_trained and not self.config.training.keep_round_data: + shutil.rmtree(rf_paths.round_dir, ignore_errors=True) + logger.info(f"Deleted RF data directory: {rf_paths.round_dir}") + # NOTE: is this the right way to check if arrays exist before dereferencing? if train_latents is not None: del train_latents @@ -2734,7 +2887,7 @@ def _add_snr_range_shading( ) # TODO: reorder plot methods (def & call sites): train -> latent -> injection - # TODO: move injection plots to data_generation.py & call at end of generate_triplet_batch() (instead of at the end of train_round() & run_training_pipeline()) + # TODO: move injection plots to data_generation.py & call at end of generate_round_to_memmap() (instead of at the end of train_round() & run_training_pipeline()) # NOTE: there's a ton of improvements we could make to this function (and subsequent _plot functions), but i just care that it works well enough for now def plot_injection_stats(self, tag: str | None = None, dir: str | None = None): """ diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..5a4c9e7c --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,226 @@ +""" +Shared pytest fixtures for the Aetherscan test suite. + +Provides three things: +1. An autouse fixture that isolates every (non-integration) test: AETHERSCAN_{DATA,MODEL,OUTPUT}_PATH + point at tmp_path, all singletons are reset via their _reset() hooks, and a fresh Config is + initialized. Teardown stops any started background threads/pools and resets everything again. +2. Synthetic data factories: tiny .npy background plates, tiny .h5 observation files, and tiny + inference CSVs — sized for speed, shaped like the real thing. +3. A headless matplotlib backend for CI (set before any aetherscan module imports pyplot). +""" + +# TODO: integration tests that verify clean shutdown under adverse conditions +# (SIGTERM/SIGINT mid-run, resource cleanup) — carried over from the old tests/placeholder +# TODO: build tests that tagged releases deploy properly to various environments +# (PyPI, HuggingFace, etc.) once release engineering lands — carried over from placeholder + +from __future__ import annotations + +import contextlib +import csv +import os +import signal +import sys + +import numpy as np +import pytest + +# Force a non-interactive backend before any test imports matplotlib.pyplot +# (train.py imports it at module level; CI runners have no display). +os.environ.setdefault("MPLBACKEND", "Agg") +# Quiet TF's C++ INFO/WARNING chatter in test output. +os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2") + + +def _reset_all_singletons(): + """Reset every singleton class via its _reset() teardown hook. + + Imports are deferred so integration runs (which skip the isolation fixture entirely) + never pull TensorFlow into the pytest parent process. + """ + from aetherscan.config import Config # noqa: PLC0415 + from aetherscan.db.db import Database # noqa: PLC0415 + from aetherscan.logger.logger import Logger # noqa: PLC0415 + from aetherscan.manager.manager import ResourceManager # noqa: PLC0415 + from aetherscan.monitor.monitor import ResourceMonitor # noqa: PLC0415 + + for cls in (ResourceMonitor, Database, Logger, ResourceManager, Config): + cls._reset() + + +def _teardown_singletons(): + """Stop background resources still held by singletons, then reset them all. + + Mirrors ResourceManager.cleanup_all()'s ordering (pools -> shared memory -> monitor -> + db -> logger) without its logging side effects, so a test that leaked a resource can't + poison the next test. + """ + from aetherscan.db.db import Database # noqa: PLC0415 + from aetherscan.logger.logger import Logger # noqa: PLC0415 + from aetherscan.manager.manager import ResourceManager # noqa: PLC0415 + from aetherscan.monitor.monitor import ResourceMonitor # noqa: PLC0415 + + manager = ResourceManager._instance + if manager is not None: + for managed_process in list(manager._processes): + with contextlib.suppress(Exception): + managed_process.close(timeout=5.0) + for managed_pool in list(manager._pools): + with contextlib.suppress(Exception): + managed_pool.close(timeout=5.0) + for managed_shm in list(manager._shared_memories): + with contextlib.suppress(Exception): + managed_shm.close() + # Each ResourceManager construction registers an atexit callback; unregister it so + # per-test instances don't accumulate stale cleanup hooks over the session. + import atexit # noqa: PLC0415 + + atexit.unregister(manager.cleanup_all) + + monitor = ResourceMonitor._instance + if monitor is not None: + with contextlib.suppress(Exception): + monitor.stop() + + db = Database._instance + if db is not None: + with contextlib.suppress(Exception): + db.stop() + + logger_instance = Logger._instance + if logger_instance is not None: + with contextlib.suppress(Exception): + logger_instance.stop() + + _reset_all_singletons() + + +@pytest.fixture(autouse=True) +def aetherscan_isolated_env(request, tmp_path, monkeypatch): + """Autouse: fresh singletons + tmp-path-scoped env for every test. + + Integration tests are exempt (marker `integration`): they exercise the real pipeline as a + subprocess and must inherit the real AETHERSCAN_* environment, and their pytest parent + process should never import TensorFlow. + """ + if request.node.get_closest_marker("integration"): + yield + return + + data_path = tmp_path / "data" + model_path = tmp_path / "models" + output_path = tmp_path / "outputs" + for sub in ("training", "testing", "inference"): + (data_path / sub).mkdir(parents=True, exist_ok=True) + model_path.mkdir(exist_ok=True) + output_path.mkdir(exist_ok=True) + + monkeypatch.setenv("AETHERSCAN_DATA_PATH", str(data_path)) + monkeypatch.setenv("AETHERSCAN_MODEL_PATH", str(model_path)) + monkeypatch.setenv("AETHERSCAN_OUTPUT_PATH", str(output_path)) + # Keep unit tests from ever talking to Slack, whatever the host env holds. + monkeypatch.delenv("SLACK_BOT_TOKEN", raising=False) + monkeypatch.delenv("SLACK_CHANNEL", raising=False) + + # ResourceManager's constructor installs SIGINT/SIGTERM handlers; snapshot and restore + # them so per-test instances don't leave stale bound-method handlers behind. + original_sigint = signal.getsignal(signal.SIGINT) + original_sigterm = signal.getsignal(signal.SIGTERM) + original_stdout, original_stderr = sys.stdout, sys.stderr + + from aetherscan.config import init_config # noqa: PLC0415 + + _reset_all_singletons() + init_config() + + yield + + _teardown_singletons() + signal.signal(signal.SIGINT, original_sigint) + signal.signal(signal.SIGTERM, original_sigterm) + sys.stdout, sys.stderr = original_stdout, original_stderr + + +@pytest.fixture +def make_background_npy(tmp_path): + """Factory for tiny .npy background plates shaped like real training backgrounds. + + Returns a callable(filename, n_cadences=8, width_bin=512) -> Path that writes a positive + float32 array of shape (n_cadences, 6, 16, width_bin) under the tmp data/training dir. + """ + + def _make(filename="backgrounds.npy", n_cadences=8, width_bin=512): + rng = np.random.default_rng(11) + # Chi-squared-ish positive noise, loosely resembling detected power spectra. + plate = rng.chisquare(df=4, size=(n_cadences, 6, 16, width_bin)).astype(np.float32) + path = tmp_path / "data" / "training" / filename + path.parent.mkdir(parents=True, exist_ok=True) + np.save(path, plate) + return path + + return _make + + +@pytest.fixture +def make_h5_observation(tmp_path): + """Factory for tiny .h5 observation files matching the filterbank-style layout. + + Returns a callable(filename, n_chans=2048, time_bins=16) -> Path that writes an .h5 with a + 'data' dataset of shape (time_bins, 1, n_chans) plus fch1/foff/nchans attrs (plain h5py, + no bitshuffle). + """ + import h5py # noqa: PLC0415 + + def _make(filename="obs.h5", n_chans=2048, time_bins=16): + rng = np.random.default_rng(7) + data = rng.chisquare(df=4, size=(time_bins, 1, n_chans)).astype(np.float32) + path = tmp_path / filename + with h5py.File(path, "w") as hf: + dset = hf.create_dataset("data", data=data) + dset.attrs["fch1"] = 8421.386717353016 + dset.attrs["foff"] = -2.7939677238464355e-06 + dset.attrs["nchans"] = n_chans + return path + + return _make + + +@pytest.fixture +def make_inference_csv(tmp_path): + """Factory for tiny inference CSVs in the cadence-grouping layout. + + Returns a callable(filename, groups) -> Path. `groups` is a list of (key_dict, h5_paths) + pairs; each h5 path becomes one row carrying the group's key columns. The h5-path column + name is read from InferenceConfig.cadence_h5_path_col on the initialized singleton, and + the default group keys mirror InferenceConfig.cadence_group_by_cols. + """ + from aetherscan.config import get_config # noqa: PLC0415 + + def _make(filename="subset.csv", groups=None): + if groups is None: + groups = [ + ( + { + "Target": "HIP110750", + "Session": "AGBT21B_999_31", + "Band": "L", + "Cadence ID": "0", + "Frequency": "1400", + }, + [f"/data/obs_{i}.h5" for i in range(6)], + ) + ] + h5_col = get_config().inference.cadence_h5_path_col + fieldnames = list(groups[0][0].keys()) + [h5_col] + path = tmp_path / "data" / "inference" / filename + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for key_dict, h5_paths in groups: + for h5 in h5_paths: + writer.writerow({**key_dict, h5_col: h5}) + return path + + return _make diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py new file mode 100644 index 00000000..775fe117 --- /dev/null +++ b/tests/integration/conftest.py @@ -0,0 +1,53 @@ +"""Shared fixtures for the cluster integration smokes. + +Both smoke tests launch `python -m aetherscan.main ...` as a subprocess against the real +cluster environment; the repo root, the AETHERSCAN_* path resolution (env var or the cluster +defaults baked into utils/run_container.sh), and the PYTHONPATH-injected subprocess launcher +live here so the two tests can't drift apart. +""" + +from __future__ import annotations + +import os +import subprocess +import sys + +import pytest + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +_SUBPROCESS_TIMEOUT_SECONDS = 7200 + + +@pytest.fixture +def cluster_paths(): + """(data_path, model_path, output_path) from AETHERSCAN_* env, falling back to the same + cluster defaults utils/run_container.sh uses.""" + return ( + os.environ.get("AETHERSCAN_DATA_PATH", "/datax/scratch/zachy/data/aetherscan"), + os.environ.get("AETHERSCAN_MODEL_PATH", "/datax/scratch/zachy/models/aetherscan"), + os.environ.get("AETHERSCAN_OUTPUT_PATH", "/datax/scratch/zachy/outputs/aetherscan"), + ) + + +@pytest.fixture +def run_pipeline(): + """Launcher for `python -m aetherscan.main ` from the repo root, with /src + prepended to PYTHONPATH so the subprocess can `import aetherscan` both inside the NGC + container and in a bare env. Returns the CompletedProcess (check=False — callers assert + on returncode so failures surface the log tail, not a bare CalledProcessError).""" + + def _run(args: list[str]) -> subprocess.CompletedProcess: + env = dict(os.environ) + src = os.path.join(_REPO_ROOT, "src") + env["PYTHONPATH"] = src + os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else src + return subprocess.run( + [sys.executable, "-m", "aetherscan.main", *args], + cwd=_REPO_ROOT, + env=env, + capture_output=True, + text=True, + timeout=_SUBPROCESS_TIMEOUT_SECONDS, + check=False, + ) + + return _run diff --git a/tests/integration/test_inference_smoke.py b/tests/integration/test_inference_smoke.py new file mode 100644 index 00000000..c526c38f --- /dev/null +++ b/tests/integration/test_inference_smoke.py @@ -0,0 +1,70 @@ +"""End-to-end CSV-inference smoke test (cluster-only, blpc3). + +Runs subset CSV inference against the persisted dummy model (test_v17) as a subprocess of +`python -m aetherscan.main inference` and asserts a clean exit plus a config snapshot. The +subset CSV's raw .h5 files live under /datag (blpc3 only), which run_container.sh does not +bind by default — either bind it explicitly: + + SINGULARITY_BIND=/datag ./utils/run_container.sh python -m pytest tests/ -m "gpu or cluster" -q + +or rely on the resume path: preprocessing skips any cadence whose stamp .npy already exists +under /preprocessed, so a previously-preprocessed subset runs without /datag. +""" + +from __future__ import annotations + +import glob +import os +import shutil +from datetime import datetime + +import pytest + +pytestmark = [pytest.mark.integration, pytest.mark.gpu, pytest.mark.cluster] + +_MODEL_TAG = "test_v17" # persisted dummy model on blpc3 +_CSV_NAME = "subset_test.csv" # 2 complete 6-observation cadences + + +def test_inference_smoke(cluster_paths, run_pipeline): + if shutil.which("nvidia-smi") is None: + pytest.skip("requires a GPU host (nvidia-smi not found)") + + data_path, model_path, output_path = cluster_paths + + encoder = os.path.join(model_path, f"vae_encoder_{_MODEL_TAG}.keras") + rf = os.path.join(model_path, f"random_forest_{_MODEL_TAG}.joblib") + saved_config = os.path.join(model_path, f"config_{_MODEL_TAG}.json") + csv_path = os.path.join(data_path, "inference", _CSV_NAME) + for required in (encoder, rf, saved_config, csv_path): + if not os.path.exists(required): + pytest.skip(f"required cluster artifact missing: {required}") + + # Raw .h5 reads need /datag; already-preprocessed stamps make it optional. + preprocessed = glob.glob(os.path.join(output_path, "preprocessed", "subset_test_*.npy")) + if not os.path.exists("/datag") and not preprocessed: + pytest.skip("/datag not mounted and no preprocessed subset stamps to resume from") + + tag = datetime.now().strftime("%Y%m%d_%H%M%S") + proc = run_pipeline( + [ + "inference", + "--encoder-path", + encoder, + "--rf-path", + rf, + "--config-path", + saved_config, + "--inference-files", + _CSV_NAME, + "--save-tag", + tag, + "--max-retries", + "1", + ] + ) + + tail = "\n".join(proc.stdout.splitlines()[-40:]) + assert proc.returncode == 0, f"inference smoke run failed (tag={tag}); last output:\n{tail}" + assert "Inference completed successfully!" in proc.stdout + assert os.path.exists(os.path.join(output_path, f"config_{tag}.json")) diff --git a/tests/integration/test_train_smoke.py b/tests/integration/test_train_smoke.py new file mode 100644 index 00000000..81d6bf86 --- /dev/null +++ b/tests/integration/test_train_smoke.py @@ -0,0 +1,87 @@ +"""End-to-end training smoke test (cluster-only). + +Runs the known-good blpc3 5-GPU smoke config from the repo runbook as a subprocess of +`python -m aetherscan.main train` and asserts the run exits cleanly with all final model +artifacts on disk. Batch/sample sizes here divide cleanly for exactly 5 replicas — this test +is sized for blpc3 (5x RTX PRO 6000). Run inside the NGC container: + + ./utils/run_container.sh python -m pytest tests/ -m "gpu or cluster" -q +""" + +from __future__ import annotations + +import os +import re +import shutil + +import pytest + +pytestmark = [pytest.mark.integration, pytest.mark.gpu, pytest.mark.cluster] + +# The known-good blpc3 smoke config: repo defaults are only divisible for 4 or 6 replicas, so +# the batch/sample geometry is overridden to divide for 5. Mirrored (as the subject under +# test) by _SMOKE_FLAGS_5_REPLICAS in tests/unit/test_cli_validation.py — keep in sync. +_SMOKE_FLAGS = [ + "--max-retries", + "1", + "--num-training-rounds", + "2", + "--epochs-per-round", + "2", + "--per-replica-batch-size", + "4", + "--per-replica-val-batch-size", + "4", + "--effective-batch-size", + "20", + "--num-samples-beta-vae", + "200", + "--num-samples-rf", + "200", + "--latent-viz-num-cadences-per-type", + "5", +] + + +def _next_test_tag(model_path: str) -> str: + """Scan existing artifacts for test_vNN tags and return the next unused one — reusing a + tag causes stale-artifact confusion.""" + versions = [0] + for root in (model_path, os.path.join(model_path, "checkpoints")): + if not os.path.isdir(root): + continue + for entry in os.listdir(root): + match = re.search(r"test_v(\d+)", entry) + if match: + versions.append(int(match.group(1))) + return f"test_v{max(versions) + 1}" + + +def test_train_smoke(cluster_paths, run_pipeline): + if shutil.which("nvidia-smi") is None: + pytest.skip("requires a GPU host (nvidia-smi not found)") + + data_path, model_path, output_path = cluster_paths + if not os.path.isdir(os.path.join(data_path, "training")): + pytest.skip(f"training data not found under {data_path}") + + tag = _next_test_tag(model_path) + proc = run_pipeline(["train", "--save-tag", tag, *_SMOKE_FLAGS]) + + tail = "\n".join(proc.stdout.splitlines()[-40:]) + assert proc.returncode == 0, f"train smoke run failed (tag={tag}); last output:\n{tail}" + assert "Training completed successfully!" in proc.stdout + + # Final artifacts: encoder/decoder/RF under model_path, config snapshot under output_path. + for artifact in ( + f"vae_encoder_{tag}.keras", + f"vae_decoder_{tag}.keras", + f"random_forest_{tag}.joblib", + ): + assert os.path.exists(os.path.join(model_path, artifact)), f"missing {artifact}" + assert os.path.exists(os.path.join(output_path, f"config_{tag}.json")) + + # Per-round checkpoints for both rounds. + checkpoints = os.path.join(model_path, "checkpoints") + for round_tag in ("round_01", "round_02"): + assert os.path.exists(os.path.join(checkpoints, f"vae_encoder_{round_tag}.keras")) diff --git a/tests/placeholder b/tests/placeholder deleted file mode 100644 index 3f1c5d6a..00000000 --- a/tests/placeholder +++ /dev/null @@ -1,3 +0,0 @@ -# TODO: -# Integration tests to verify clean shutdown under various conditions -# Build tests to ensure tagged releases can be deployed properly to various environments (PyPI, HuggingFace, etc.) diff --git a/tests/unit/test_cli_validation.py b/tests/unit/test_cli_validation.py new file mode 100644 index 00000000..7d723263 --- /dev/null +++ b/tests/unit/test_cli_validation.py @@ -0,0 +1,461 @@ +"""Unit tests for aetherscan.cli: tag pattern, validation matrix, cross-param solver, and +apply_saved_config precedence (defaults < saved config < CLI args).""" + +from __future__ import annotations + +import collections +import json + +import pytest + +from aetherscan import cli +from aetherscan.cli import ( + _TAG_PATTERN, + _check_cross_constraints, + _solve_cross_param_constraints, + apply_args_to_config, + apply_saved_config, + collect_validation_errors, + setup_argument_parser, +) +from aetherscan.config import get_config + +# The blpc3 smoke config from the repo's known-good runbook: divisible for exactly 5 replicas. +# Mirrored by _SMOKE_FLAGS in tests/integration/test_train_smoke.py (which runs it for real on +# the cluster) — keep the two in sync. +_SMOKE_FLAGS_5_REPLICAS = [ + "--per-replica-batch-size", + "4", + "--per-replica-val-batch-size", + "4", + "--effective-batch-size", + "20", + "--num-samples-beta-vae", + "200", + "--num-samples-rf", + "200", + "--latent-viz-num-cadences-per-type", + "5", +] + + +def _parse(argv): + return setup_argument_parser().parse_args(argv) + + +def _cross_param_errors(args, num_replicas): + errors = collect_validation_errors(args, num_replicas) + return [e for e in errors if e.fix_kind == "cross_param"] + + +@pytest.fixture(autouse=True) +def _ample_disk_space(monkeypatch): + """The round-data disk-budget check reads the real filesystem via shutil.disk_usage (the + full-scale default config needs ~650 GB free — more than CI runners or dev boxes have). + Pin it to a huge value so validation tests are machine-independent; tests that exercise + the check itself (TestRoundDataFlags) patch their own values on top.""" + usage = collections.namedtuple("usage", ["total", "used", "free"]) + monkeypatch.setattr(cli.shutil, "disk_usage", lambda path: usage(2**60, 0, 2**60)) + + +class TestTagPattern: + @pytest.mark.parametrize( + "tag", + ["20260712_123456", "final_v1", "final_v12", "round_01", "round_5", "test_v1", "test_v17"], + ) + def test_accepted_formats(self, tag): + assert _TAG_PATTERN.match(tag) + + @pytest.mark.parametrize( + "tag", + [ + "smoke_blackwell", # free-form slug + "final_1", # missing v + "final_v", # missing version number + "test_17", # missing v + "TEST_V1", # wrong case + "2026_0712", # malformed timestamp + "20260712-123456", # wrong separator + "20260712_12345", # HHMMS (5 digits) + "round_", # missing round number + " test_v1", # leading whitespace + "test_v1 ", # trailing whitespace + "", + ], + ) + def test_rejected_formats(self, tag): + assert not _TAG_PATTERN.match(tag) + + +class TestCrossReplicaDivisibilityMatrix: + """Default config divides cleanly across 4 or 6 replicas but NOT 5; the blpc3 smoke config + is the inverse. collect_validation_errors must reproduce that matrix exactly.""" + + @pytest.mark.parametrize("num_replicas", [4, 6]) + def test_defaults_valid_for_4_and_6_replicas(self, num_replicas): + assert _cross_param_errors(_parse(["train"]), num_replicas) == [] + + def test_defaults_invalid_for_5_replicas(self): + errors = _cross_param_errors(_parse(["train"]), 5) + violated = {e.field for e in errors} + assert violated == { + "training.effective_batch_size", + "training.num_samples_beta_vae", + "training.num_samples_rf", + "training.latent_viz_num_cadences_per_type", + } + + def test_smoke_config_valid_for_5_replicas(self): + assert _cross_param_errors(_parse(["train", *_SMOKE_FLAGS_5_REPLICAS]), 5) == [] + + @pytest.mark.parametrize("num_replicas", [4, 6]) + def test_smoke_config_invalid_for_4_and_6_replicas(self, num_replicas): + assert _cross_param_errors(_parse(["train", *_SMOKE_FLAGS_5_REPLICAS]), num_replicas) + + def test_unknown_replica_count_skips_cross_checks(self): + # num_replicas=None (no GPUs detectable) must skip the divisibility section entirely, + # even for a config that would fail on any replica count. + assert _cross_param_errors(_parse(["train", *_SMOKE_FLAGS_5_REPLICAS]), None) == [] + + +class TestSemanticChecks: + def test_save_tag_format_error(self): + errors = collect_validation_errors(_parse(["train", "--save-tag", "bogus"]), None) + assert any(e.field == "checkpoint.save_tag" and e.fix_kind == "format" for e in errors) + + def test_load_tag_format_error(self): + errors = collect_validation_errors(_parse(["train", "--load-tag", "bogus"]), None) + assert any(e.field == "checkpoint.load_tag" and e.fix_kind == "format" for e in errors) + + def test_num_samples_divisible_by_4(self): + errors = collect_validation_errors(_parse(["train", "--num-samples-beta-vae", "202"]), None) + assert any(e.field == "training.num_samples_beta_vae" and e.divisor == 4 for e in errors) + + def test_curriculum_schedule_enum(self): + errors = collect_validation_errors( + _parse(["train", "--curriculum-schedule", "sigmoid"]), None + ) + assert any( + e.field == "training.curriculum_schedule" and e.fix_kind == "enum" for e in errors + ) + + def test_step_schedule_sum_constraint(self): + argv = [ + "train", + "--curriculum-schedule", + "step", + "--num-training-rounds", + "10", + "--step-easy-rounds", + "3", + "--step-hard-rounds", + "4", + ] + errors = collect_validation_errors(_parse(argv), None) + assert any( + e.field == "training.step_easy_rounds" and e.fix_kind == "cross_param" for e in errors + ) + + def test_snr_curriculum_ordering(self): + argv = ["train", "--initial-snr-range", "5", "--final-snr-range", "10"] + errors = collect_validation_errors(_parse(argv), None) + assert any(e.field == "training.initial_snr_range" for e in errors) + + def test_missing_train_files_reported(self): + errors = collect_validation_errors(_parse(["train"]), None) + # Default train_files don't exist under the tmp data path. + file_errors = [e for e in errors if e.fix_kind == "file_exists"] + assert {e.field for e in file_errors} == {"data.train_files"} + assert len(file_errors) == len(get_config().data.train_files) + + def test_existing_train_files_pass(self, tmp_path): + config = get_config() + for filename in config.data.train_files: + (tmp_path / "data" / "training" / filename).touch() + errors = collect_validation_errors(_parse(["train"]), None) + assert [e for e in errors if e.fix_kind == "file_exists"] == [] + + def test_inference_requires_model_artifacts(self): + errors = collect_validation_errors(_parse(["inference"]), None) + fields = {e.field for e in errors if e.fix_kind == "file_exists"} + assert { + "inference.encoder_path", + "inference.rf_path", + "inference.config_path", + } <= fields + + def test_inference_stamp_width_must_match_width_bin(self, make_inference_csv): + csv_path = make_inference_csv("subset.csv") + argv = [ + "inference", + "--inference-files", + csv_path.name, + "--stamp-width", + "2048", # width_bin default is 4096 + ] + errors = collect_validation_errors(_parse(argv), None) + assert any(e.field == "inference.stamp_width" for e in errors) + + +class TestCrossParamSolver: + _VALID_BASE = { + # The repo defaults: valid for 4 and 6 replicas. + "num_samples_beta_vae": 499200, + "num_samples_rf": 99840, + "train_val_split": 0.8, + "per_replica_batch_size": 128, + "effective_batch_size": 3072, + "per_replica_val_batch_size": 80, + } + + def test_check_cross_constraints_matrix(self): + assert _check_cross_constraints(**self._VALID_BASE, num_replicas_list=[4]) + assert _check_cross_constraints(**self._VALID_BASE, num_replicas_list=[6]) + assert not _check_cross_constraints(**self._VALID_BASE, num_replicas_list=[5]) + + def test_solver_returns_base_when_already_valid(self): + assert _solve_cross_param_constraints(self._VALID_BASE, [4, 6]) == self._VALID_BASE + + def test_solver_respects_candidate_budget(self): + # The default search ranges exceed a tiny budget, so the solver must bail with None + # rather than grind through the grid. + base = dict(self._VALID_BASE) + base["effective_batch_size"] = 3070 # invalidate so the grid search is attempted + assert _solve_cross_param_constraints(base, [4], max_candidates=10) is None + + def test_solver_finds_nearest_valid_config(self, monkeypatch): + # Shrink the search ranges to a tractable grid and verify the solver picks the + # L1-nearest satisfying combination. + monkeypatch.setattr( + cli, + "_SEARCH_RANGES", + { + "num_samples_beta_vae": (160, 320, 80), + "num_samples_rf": (40, 80, 20), + "per_replica_batch_size": (4, 8, 4), + "effective_batch_size": (16, 64, 16), + "per_replica_val_batch_size": (4, 8, 4), + }, + ) + base = { + "num_samples_beta_vae": 250, + "num_samples_rf": 50, + "train_val_split": 0.8, + "per_replica_batch_size": 5, + "effective_batch_size": 20, + "per_replica_val_batch_size": 5, + } + solution = _solve_cross_param_constraints(base, [4]) + assert solution is not None + assert solution["train_val_split"] == base["train_val_split"] # held fixed + assert _check_cross_constraints(**solution, num_replicas_list=[4]) + # Every solved field must come from the (patched) search grid. + for field_name, (lo, hi, step) in cli._SEARCH_RANGES.items(): + assert solution[field_name] in range(lo, hi + 1, step) + # And the solution must be L1-minimal among all valid grid points. + assert self._l1(solution, base) == min( + self._l1(candidate, base) for candidate in self._valid_grid_points(base) + ) + + @staticmethod + def _l1(candidate, base): + return sum( + abs(candidate[f] - base[f]) + for f in ( + "num_samples_beta_vae", + "num_samples_rf", + "per_replica_batch_size", + "effective_batch_size", + "per_replica_val_batch_size", + ) + ) + + @staticmethod + def _valid_grid_points(base): + from itertools import product # noqa: PLC0415 + + ranges = {f: range(lo, hi + 1, step) for f, (lo, hi, step) in cli._SEARCH_RANGES.items()} + for nsb, nsr, prb, eb, prvb in product( + ranges["num_samples_beta_vae"], + ranges["num_samples_rf"], + ranges["per_replica_batch_size"], + ranges["effective_batch_size"], + ranges["per_replica_val_batch_size"], + ): + candidate = { + "num_samples_beta_vae": nsb, + "num_samples_rf": nsr, + "train_val_split": base["train_val_split"], + "per_replica_batch_size": prb, + "effective_batch_size": eb, + "per_replica_val_batch_size": prvb, + } + if _check_cross_constraints(**candidate, num_replicas_list=[4]): + yield candidate + + +class TestApplySavedConfigPrecedence: + def test_saved_config_overrides_defaults(self, tmp_path): + saved = { + "training": {"num_training_rounds": 7, "snr_base": 33}, + "checkpoint": {"save_tag": "final_v9"}, + "data_path": "/saved/data/path", + } + path = tmp_path / "saved_config.json" + path.write_text(json.dumps(saved)) + + apply_saved_config(str(path)) + config = get_config() + assert config.training.num_training_rounds == 7 + assert config.training.snr_base == 33 + # Documented sharp edge: the saved file clobbers checkpoint.save_tag too. + assert config.checkpoint.save_tag == "final_v9" + assert config.data_path == "/saved/data/path" + + def test_cli_args_override_saved_config(self, tmp_path): + saved = {"training": {"snr_base": 33, "num_training_rounds": 7}} + path = tmp_path / "saved_config.json" + path.write_text(json.dumps(saved)) + apply_saved_config(str(path)) + + args = _parse(["train", "--snr-base", "44"]) + apply_args_to_config(args) + config = get_config() + assert config.training.snr_base == 44 # CLI wins over saved + assert config.training.num_training_rounds == 7 # saved wins over default (20) + + def test_unknown_keys_and_fields_skipped(self, tmp_path): + saved = { + "not_a_section": {"whatever": 1}, + "training": {"not_a_field": 123, "snr_base": 21}, + } + path = tmp_path / "saved_config.json" + path.write_text(json.dumps(saved)) + apply_saved_config(str(path)) + config = get_config() + assert config.training.snr_base == 21 + assert not hasattr(config.training, "not_a_field") + assert not hasattr(config, "not_a_section") + + def test_missing_file_raises(self): + with pytest.raises(ValueError, match="does not exist"): + apply_saved_config("/nonexistent/config.json") + + def test_malformed_json_raises(self, tmp_path): + path = tmp_path / "bad.json" + path.write_text("{not json") + with pytest.raises(ValueError, match="not valid JSON"): + apply_saved_config(str(path)) + + +class TestApplyArgsToConfig: + def test_none_args_leave_defaults(self): + config = get_config() + default_rounds = config.training.num_training_rounds + apply_args_to_config(_parse(["train"])) + assert config.training.num_training_rounds == default_rounds + + def test_mode_scoped_flags_route_to_correct_section(self): + # --per-replica-batch-size / --max-retries exist in both subparsers; the command + # gates which config section receives them. + config = get_config() + apply_args_to_config(_parse(["train", "--per-replica-batch-size", "64"])) + assert config.training.per_replica_batch_size == 64 + + inference_default = config.inference.per_replica_batch_size + apply_args_to_config(_parse(["inference", "--per-replica-batch-size", "32"])) + assert config.inference.per_replica_batch_size == 32 + assert config.training.per_replica_batch_size == 64 # untouched by inference flag + assert inference_default != 32 + + def test_load_tag_infers_start_round(self): + config = get_config() + apply_args_to_config(_parse(["train", "--load-tag", "round_03"])) + assert config.checkpoint.start_round == 4 + + +class TestRoundDataFlags: + """Flags and validation for the disk-backed round-data pipeline (round_data.py).""" + + def test_flags_apply_to_config(self): + apply_args_to_config( + _parse( + [ + "train", + "--round-data-dir", + "/scratch/rounds", + "--no-overlap-data-generation", + "--keep-round-data", + "--data-gen-task-size", + "128", + ] + ) + ) + config = get_config() + assert config.training.round_data_dir == "/scratch/rounds" + assert config.training.overlap_data_generation is False + assert config.training.keep_round_data is True + assert config.training.data_gen_task_size == 128 + + def test_defaults_preserved_when_omitted(self): + apply_args_to_config(_parse(["train"])) + config = get_config() + assert config.training.round_data_dir is None + assert config.training.overlap_data_generation is True + assert config.training.keep_round_data is False + assert config.training.data_gen_task_size == 256 + + def test_data_gen_task_size_below_one_rejected(self): + errors = collect_validation_errors(_parse(["train", "--data-gen-task-size", "0"]), None) + assert any( + e.field == "training.data_gen_task_size" and e.fix_kind == "clamp_low" for e in errors + ) + + def _patch_free_bytes(self, monkeypatch, free_bytes): + usage = collections.namedtuple("usage", ["total", "used", "free"]) + monkeypatch.setattr( + cli.shutil, "disk_usage", lambda path: usage(free_bytes * 2, free_bytes, free_bytes) + ) + + def test_disk_budget_error_when_insufficient(self, monkeypatch): + config = get_config() + round_nbytes = cli._estimate_round_data_nbytes( + config.training.num_samples_beta_vae, + config.data.num_observations, + config.data.time_bins, + config.data.width_bin // config.data.downsample_factor, + ) + # Between 1.1x and 2.2x one round: fails with overlap (default), passes without + self._patch_free_bytes(monkeypatch, int(1.5 * round_nbytes)) + + errors = collect_validation_errors(_parse(["train"]), None) + disk_errors = [e for e in errors if e.field == "training.round_data_dir"] + assert len(disk_errors) == 1 + assert "GB free" in disk_errors[0].message + + errors = collect_validation_errors(_parse(["train", "--no-overlap-data-generation"]), None) + assert not any(e.field == "training.round_data_dir" for e in errors) + + def test_disk_budget_ok_when_sufficient(self, monkeypatch): + config = get_config() + round_nbytes = cli._estimate_round_data_nbytes( + config.training.num_samples_beta_vae, + config.data.num_observations, + config.data.time_bins, + config.data.width_bin // config.data.downsample_factor, + ) + self._patch_free_bytes(monkeypatch, int(10 * round_nbytes)) + errors = collect_validation_errors(_parse(["train"]), None) + assert not any(e.field == "training.round_data_dir" for e in errors) + + def test_estimate_scales_with_sample_count(self): + one = cli._estimate_round_data_nbytes(4, 6, 16, 512) + two = cli._estimate_round_data_nbytes(8, 6, 16, 512) + assert two == 2 * one + # 3 arrays x n x 6 x 16 x 512 float32 + n x U20 labels + assert one == 3 * 4 * 6 * 16 * 512 * 4 + 4 * 80 + + def test_nearest_existing_ancestor(self, tmp_path): + missing = tmp_path / "a" / "b" / "c" + assert cli._nearest_existing_ancestor(str(missing)) == str(tmp_path) + assert cli._nearest_existing_ancestor(str(tmp_path)) == str(tmp_path) diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py new file mode 100644 index 00000000..42ecf565 --- /dev/null +++ b/tests/unit/test_config.py @@ -0,0 +1,107 @@ +"""Unit tests for aetherscan.config: singleton semantics and to_dict() field coverage.""" + +from __future__ import annotations + +import dataclasses +import json +import os + +from aetherscan.config import Config, get_config, init_config + +# Sub-dataclass sections serialized by Config.to_dict(), keyed by attribute name. +_CONFIG_SECTIONS = ( + "db", + "manager", + "monitor", + "logger", + "beta_vae", + "rf", + "gpu", + "data", + "training", + "inference", + "checkpoint", +) + + +class TestSingletonSemantics: + def test_constructor_returns_same_instance(self): + assert Config() is Config() + + def test_get_config_returns_singleton(self): + assert get_config() is Config() + + def test_init_config_returns_singleton(self): + assert init_config() is get_config() + + def test_mutation_visible_across_accessors(self): + Config().training.num_training_rounds = 3 + assert get_config().training.num_training_rounds == 3 + + def test_reset_produces_fresh_instance(self): + first = Config() + first.training.num_training_rounds = 3 + Config._reset() + second = init_config() + assert second is not first + assert second.training.num_training_rounds != 3 + + def test_env_var_paths_respected(self, tmp_path): + # The autouse fixture points AETHERSCAN_* at tmp_path before init_config(). + config = get_config() + assert config.data_path == os.environ["AETHERSCAN_DATA_PATH"] + assert config.model_path == os.environ["AETHERSCAN_MODEL_PATH"] + assert config.output_path == os.environ["AETHERSCAN_OUTPUT_PATH"] + assert config.data_path.startswith(str(tmp_path)) + + def test_file_path_helpers(self): + config = get_config() + assert config.get_training_file_path("a.npy") == os.path.join( + config.data_path, "training", "a.npy" + ) + assert config.get_test_file_path("b.npy") == os.path.join( + config.data_path, "testing", "b.npy" + ) + assert config.get_inference_file_path("c.csv") == os.path.join( + config.data_path, "inference", "c.csv" + ) + # base_path override (used by validate_args before --data-path is applied) + assert config.get_training_file_path("a.npy", base_path="/elsewhere") == os.path.join( + "/elsewhere", "training", "a.npy" + ) + + +class TestToDict: + def test_sections_cover_every_dataclass_field(self): + """to_dict() must be updated by hand for every new config field — catch drift.""" + config = get_config() + serialized = config.to_dict() + for section in _CONFIG_SECTIONS: + expected = {f.name for f in dataclasses.fields(type(getattr(config, section)))} + actual = set(serialized[section].keys()) + assert actual == expected, ( + f"Config.to_dict()['{section}'] is out of sync with " + f"{type(getattr(config, section)).__name__}: " + f"missing={expected - actual}, extra={actual - expected}" + ) + + def test_values_round_trip(self): + config = get_config() + config.training.num_training_rounds = 5 + config.gpu.num_replicas = 4 + serialized = config.to_dict() + assert serialized["training"]["num_training_rounds"] == 5 + assert serialized["gpu"]["num_replicas"] == 4 + assert serialized["paths"]["data_path"] == config.data_path + for section in _CONFIG_SECTIONS: + for field_name, value in serialized[section].items(): + stored = getattr(getattr(config, section), field_name) + if isinstance(stored, tuple): + # json-level equality: tuples serialize as lists downstream + assert list(stored) == list(value) + else: + assert stored == value + + def test_json_serializable(self): + # Saved run configs are json.dump'ed — the dict must be JSON-native throughout. + json.dumps(get_config().to_dict()) diff --git a/tests/unit/test_data_generation.py b/tests/unit/test_data_generation.py new file mode 100644 index 00000000..aca15aa9 --- /dev/null +++ b/tests/unit/test_data_generation.py @@ -0,0 +1,266 @@ +"""Unit tests for aetherscan.data_generation: log-norm, intersection checks, signal injection, +create_* cadence generators, and intensity statistics.""" + +from __future__ import annotations + +import math +import random + +import numpy as np +import pytest + +from aetherscan.data_generation import ( + _compute_intensity_stats, + check_valid_intersection, + create_false, + create_true_double, + create_true_single, + log_norm, + new_cadence, +) + +# Keep injection fast: small frequency axis, real-ish resolutions. +_WIDTH_BIN = 128 +_FREQ_RES = 2.7939677238464355 # Hz +_TIME_RES = 18.25361108 # seconds + +_SIGNAL_INFO_KEYS = { + "snr", + "drift_rate", + "signal_width", + "starting_bin", + "slope_pixel", + "y_intercept", +} + +_STAT_KEYS = { + "global_mean", + "global_median", + "global_std", + "global_mad", + "global_skew", + "global_kurtosis", +} + + +@pytest.fixture(autouse=True) +def _seed_rngs(): + random.seed(11) + np.random.seed(11) + + +@pytest.fixture +def plate(make_background_npy): + """A tiny background plate loaded from a factory-written .npy (round-trips the factory).""" + path = make_background_npy("plate.npy", n_cadences=4, width_bin=_WIDTH_BIN) + loaded = np.load(path) + assert loaded.shape == (4, 6, 16, _WIDTH_BIN) + assert loaded.dtype == np.float32 + return loaded + + +class TestLogNorm: + def test_output_bounded_in_unit_interval(self): + rng = np.random.default_rng(0) + data = rng.chisquare(df=4, size=(16, 64)) + result = log_norm(data) + assert result.shape == data.shape + assert result.min() == 0.0 + assert result.max() == pytest.approx(1.0) + + def test_preserves_ordering(self): + data = np.array([[1.0, 10.0, 100.0, 1000.0]]) + result = log_norm(data) + assert np.all(np.diff(result[0]) > 0) + + def test_constant_input_maps_to_zeros(self): + result = log_norm(np.full((4, 4), 3.0)) + assert np.all(result == 0.0) + + def test_idempotence_on_normalized_scale(self): + # Re-normalizing an already [0, 1] array must stay within [0, 1] with the same + # endpoints — the transform is a monotone squash onto the unit interval. + rng = np.random.default_rng(1) + once = log_norm(rng.chisquare(df=4, size=(8, 32))) + twice = log_norm(once) + assert twice.min() == 0.0 + assert twice.max() == pytest.approx(1.0) + assert np.array_equal(np.argsort(once, axis=None), np.argsort(twice, axis=None)) + + +class TestCheckValidIntersection: + def test_parallel_lines_are_valid(self): + assert check_valid_intersection(2.0, 2.0, 0.0, 50.0) is True + + def test_intersection_inside_on_region_invalid(self): + # y = x and y = -x + 20 intersect at (10, 10): inside the first ON block [0, 16]. + assert check_valid_intersection(1.0, -1.0, 0.0, 20.0) is False + + def test_intersection_between_on_regions_valid(self): + # y = x and y = -x + 40 intersect at (20, 20): inside the first OFF block (16, 32). + assert check_valid_intersection(1.0, -1.0, 0.0, 40.0) is True + + @pytest.mark.parametrize("y_target", [0.0, 16.0, 32.0, 48.0, 64.0, 80.0]) + def test_on_region_boundaries_are_inclusive(self, y_target): + # Intersection exactly on an ON boundary counts as inside (invalid). + # y = x and y = -x + 2*y_target intersect at (y_target, y_target). + assert check_valid_intersection(1.0, -1.0, 0.0, 2 * y_target) is False + + +class TestNewCadence: + def _background(self, rng=None): + rng = rng or np.random.default_rng(3) + return rng.chisquare(df=4, size=(96, _WIDTH_BIN)) + + def test_shapes_and_signal_info(self): + data = self._background() + modified, signal_info, clamped = new_cadence( + data.copy(), + snr=20.0, + width_bin=_WIDTH_BIN, + freq_resolution=_FREQ_RES, + time_resolution=_TIME_RES, + ) + assert modified.shape == data.shape + assert set(signal_info.keys()) == _SIGNAL_INFO_KEYS + assert all(isinstance(v, float) for v in signal_info.values()) + assert signal_info["snr"] == 20.0 + assert 1 <= signal_info["starting_bin"] <= _WIDTH_BIN - 1 + assert isinstance(clamped, bool) + + def test_injection_adds_power(self): + data = self._background() + modified, _, _ = new_cadence( + data.copy(), + snr=50.0, + width_bin=_WIDTH_BIN, + freq_resolution=_FREQ_RES, + time_resolution=_TIME_RES, + ) + # A signal only adds intensity — total power must strictly increase. + assert modified.sum() > data.sum() + + def test_drift_rate_inverse_of_slope(self): + _, signal_info, clamped = new_cadence( + self._background(), + snr=20.0, + width_bin=_WIDTH_BIN, + freq_resolution=_FREQ_RES, + time_resolution=_TIME_RES, + ) + if not clamped: + slope_physical = signal_info["slope_pixel"] * (_TIME_RES / _FREQ_RES) + # drift_rate = -1/slope_physical up to the additive noise term on the slope, + # so only the sign relationship is exact. + assert math.copysign(1, signal_info["drift_rate"]) == -math.copysign(1, slope_physical) + + def test_near_zero_slope_is_clamped(self, monkeypatch): + # Force the degenerate geometry: time/freq resolution ratio ~0 makes the physical + # slope collapse below the clamp threshold when the additive noise term is 0. + # new_cadence draws random.random() in order: starting_bin, slope noise, signal width. + draws = iter([0.5, 0.0]) + monkeypatch.setattr(random, "random", lambda: next(draws, 0.25)) + monkeypatch.setattr(np.random, "choice", lambda *a, **k: 1) + _, signal_info, clamped = new_cadence( + self._background(), + snr=20.0, + width_bin=_WIDTH_BIN, + freq_resolution=1.0, + time_resolution=1e-12, + ) + assert clamped is True + # Clamped slope preserves direction but pins magnitude at the floor: |drift| = 1e6. + assert abs(signal_info["drift_rate"]) == pytest.approx(1e6) + + +class TestCreateCadences: + def _kwargs(self): + return { + "snr_base": 10.0, + "snr_range": 5.0, + "width_bin": _WIDTH_BIN, + "freq_resolution": _FREQ_RES, + "time_resolution": _TIME_RES, + } + + def _assert_common(self, final, sample_info, plate): + assert final.shape == (6, 16, _WIDTH_BIN) + assert 0 <= sample_info["background_index"] < plate.shape[0] + assert set(sample_info["intensity_stats"].keys()) == {"A", "B", "C"} + for stage_stats in sample_info["intensity_stats"].values(): + assert set(stage_stats.keys()) == _STAT_KEYS + assert isinstance(sample_info["slope_was_clamped"], bool) + # Output is log-normalized per observation. + assert final.min() >= 0.0 + assert final.max() <= 1.0 + + def test_create_false_injected(self, plate): + final, sample_info = create_false(plate, inject=True, **self._kwargs()) + self._assert_common(final, sample_info, plate) + keys = set(sample_info["signal_info"].keys()) + assert keys == {f"rfi_{k}" for k in _SIGNAL_INFO_KEYS} + + def test_create_false_no_injection(self, plate): + final, sample_info = create_false(plate, inject=False, **self._kwargs()) + self._assert_common(final, sample_info, plate) + assert sample_info["signal_info"] == {} + # Without injection, Stage B mirrors Stage A exactly. + assert sample_info["intensity_stats"]["B"] == sample_info["intensity_stats"]["A"] + # And every observation is just the log-normalized background (float32 path: the + # generator normalizes the raw float32 plate slices directly). + base = plate[sample_info["background_index"]] + for obs in range(6): + np.testing.assert_allclose(final[obs], log_norm(base[obs])) + + def test_create_true_single_injects_on_only(self, plate): + final, sample_info = create_true_single(plate, **self._kwargs()) + self._assert_common(final, sample_info, plate) + keys = set(sample_info["signal_info"].keys()) + assert keys == {f"eti_{k}" for k in _SIGNAL_INFO_KEYS} + # OFF observations (1, 3, 5) carry no injection: they equal the log-normed background. + base = plate[sample_info["background_index"]] + for obs in (1, 3, 5): + np.testing.assert_allclose(final[obs], log_norm(base[obs].astype(np.float64))) + + def test_create_true_double_injects_both(self, plate): + final, sample_info = create_true_double(plate, **self._kwargs()) + self._assert_common(final, sample_info, plate) + keys = set(sample_info["signal_info"].keys()) + expected = {f"rfi_{k}" for k in _SIGNAL_INFO_KEYS} | {f"eti_{k}" for k in _SIGNAL_INFO_KEYS} + assert keys == expected + # The two injected trajectories must not intersect inside an ON region. + assert check_valid_intersection( + sample_info["signal_info"]["rfi_slope_pixel"], + sample_info["signal_info"]["eti_slope_pixel"], + sample_info["signal_info"]["rfi_y_intercept"], + sample_info["signal_info"]["eti_y_intercept"], + ) + + +class TestComputeIntensityStats: + def test_normal_case_matches_numpy(self): + rng = np.random.default_rng(5) + data = rng.chisquare(df=4, size=(6, 16, 32)).astype(np.float32) + stats = _compute_intensity_stats(data) + assert set(stats.keys()) == _STAT_KEYS + flat = data.ravel().astype(np.float64) + assert stats["global_mean"] == pytest.approx(np.mean(flat)) + assert stats["global_median"] == pytest.approx(np.median(flat)) + assert stats["global_std"] == pytest.approx(np.std(flat)) + assert stats["global_mad"] == pytest.approx(np.median(np.abs(flat - np.median(flat)))) + assert all(math.isfinite(v) for v in stats.values()) + + def test_empty_array_returns_nan_for_every_key(self): + stats = _compute_intensity_stats(np.array([])) + assert set(stats.keys()) == _STAT_KEYS + assert all(math.isnan(v) for v in stats.values()) + + def test_constant_array_higher_moments_degenerate(self): + stats = _compute_intensity_stats(np.full((4, 4), 7.0)) + assert stats["global_mean"] == 7.0 + assert stats["global_std"] == 0.0 + assert stats["global_mad"] == 0.0 + # Skew/kurtosis of a zero-variance array are NaN — the DB layer records these with + # is_finite=0 rather than rejecting the write. + assert math.isnan(stats["global_skew"]) + assert math.isnan(stats["global_kurtosis"]) diff --git a/tests/unit/test_db.py b/tests/unit/test_db.py new file mode 100644 index 00000000..486d67bb --- /dev/null +++ b/tests/unit/test_db.py @@ -0,0 +1,227 @@ +"""Unit tests for aetherscan.db: writer thread lifecycle, flush sentinel protocol, executemany +batching across tables, is_finite sanitization, and query filters / column whitelists — all +against a tmp-path SQLite file.""" + +from __future__ import annotations + +import json +import time + +import numpy as np +import pytest + +from aetherscan.config import get_config +from aetherscan.db.db import Database + + +@pytest.fixture +def db(tmp_path): + """A started Database against the tmp output path, with the periodic flush effectively + disabled (long write_interval / huge buffer) so only the flush sentinel can drain writes — + making flush() the thing under test rather than a timing accident.""" + config = get_config() + config.db.write_interval = 300.0 + config.db.write_buffer_max_size = 100_000 + database = Database() + assert database.db_path.startswith(str(tmp_path)) + database.start() + yield database + database.stop() + + +class TestWriterThreadLifecycle: + def test_start_spawns_writer_thread(self, db): + assert db.writer_thread is not None + assert db.writer_thread.is_alive() + + def test_start_is_idempotent(self, db): + thread = db.writer_thread + db.start() + assert db.writer_thread is thread # no second thread spawned + + def test_stop_terminates_thread(self, db): + db.stop() + assert not db.writer_thread.is_alive() + + def test_restart_after_stop(self, db): + db.stop() + db.start() + assert db.writer_thread.is_alive() + db.write_training_stat("beta_vae", "total_loss", 1.0, tag="test_v1") + assert db.flush(timeout=10) is True + assert len(db.query_training_stat(tag="test_v1")) == 1 + + def test_singleton_semantics(self, db): + assert Database() is db + + +class TestFlushSentinel: + def test_flush_drains_queued_writes(self, db): + for i in range(5): + db.write_training_stat( + "beta_vae", "total_loss", float(i), round_number=1, epoch_number=i, tag="test_v1" + ) + # write_interval is 300 s, buffer cap 100k: only the sentinel can have flushed these. + assert db.flush(timeout=10) is True + rows = db.query_training_stat(tag="test_v1") + assert len(rows) == 5 + assert sorted(r["value"] for r in rows) == [0.0, 1.0, 2.0, 3.0, 4.0] + + def test_flush_without_writer_is_noop_success(self): + database = Database() # never started + assert database.flush(timeout=1) is True + + def test_flush_during_shutdown_returns_false(self, db): + # Deterministic shutdown-in-progress simulation: setting stop_event on the *live* + # writer races (the thread may exit before flush() checks it, flipping the result + # to True via the not-running branch). Stop the real writer first, then stand in an + # always-alive stub so flush() deterministically hits the stop_event check. + class _AliveStub: + @staticmethod + def is_alive(): + return True + + db.stop() + db.writer_thread = _AliveStub() + db.stop_event.set() # stop() already set it; explicit for the reader + try: + assert db.flush(timeout=1) is False + finally: + db.writer_thread = None # let the fixture's stop() no-op cleanly + + +class TestExecutemanyBatching: + def test_mixed_table_buffer_lands_in_all_tables(self, db): + """One flush must route grouped executemany() inserts to every destination table.""" + tag = "test_v1" + db.write_system_resource("cpu", "system_total", 42.0, unit="percent", tag=tag) + db.write_training_stat("beta_vae", "total_loss", 0.5, 1, 1, tag=tag) + db.write_injection_stat("eti_snr", 12.5, round_number=1, signal_class="true", tag=tag) + db.write_latent_snapshot("beta_vae", 1, 1, 10, 0, "true_only_eti", [[0.1] * 4] * 6, tag=tag) + db.write_inference_result("/tmp/x.npy", 0, 1, 0.99, latent_vector=np.arange(4.0), tag=tag) + assert db.flush(timeout=10) is True + + assert len(db.query_system_resource(tag=tag)) == 1 + assert len(db.query_training_stat(tag=tag)) == 1 + assert len(db.query_injection_stat(tag=tag)) == 1 + snapshots = db.query_latent_snapshots(tag=tag) + assert len(snapshots) == 1 + assert json.loads(snapshots[0]["latent_vector"]) == [[0.1] * 4] * 6 + results = db.query_inference_result(tag=tag) + assert len(results) == 1 + assert json.loads(results[0]["latent_vector"]) == [0.0, 1.0, 2.0, 3.0] + + def test_multiple_rows_per_table_batch(self, db): + for i in range(20): + db.write_training_stat("beta_vae", "kl_loss", float(i), 1, i, tag="test_v2") + assert db.flush(timeout=10) is True + assert len(db.query_training_stat(tag="test_v2", stat_name="kl_loss")) == 20 + + +class TestIsFiniteSanitization: + def test_non_finite_values_stored_as_zero_with_flag(self, db): + tag = "test_v3" + db.write_injection_stat("global_skew", float("nan"), tag=tag) + db.write_injection_stat("global_skew", float("inf"), tag=tag) + db.write_injection_stat("global_skew", float("-inf"), tag=tag) + db.write_injection_stat("global_skew", 1.25, tag=tag) + assert db.flush(timeout=10) is True + + finite_only = db.query_injection_stat(tag=tag) # only_finite defaults to True + assert [r["value"] for r in finite_only] == [1.25] + + everything = db.query_injection_stat(tag=tag, only_finite=False) + assert len(everything) == 4 + sanitized = [r for r in everything if r["is_finite"] == 0] + assert len(sanitized) == 3 + assert all(r["value"] == 0.0 for r in sanitized) + + def test_slope_clamped_flag_round_trip(self, db): + tag = "test_v4" + db.write_injection_stat("rfi_drift_rate", 1.0, slope_clamped=True, tag=tag) + db.write_injection_stat("rfi_drift_rate", 2.0, slope_clamped=False, tag=tag) + db.write_injection_stat("rfi_drift_rate", 3.0, slope_clamped=None, tag=tag) + assert db.flush(timeout=10) is True + + clamped = db.query_injection_stat(tag=tag, only_slope_clamped=True) + assert [r["value"] for r in clamped] == [1.0] + unclamped = db.query_injection_stat(tag=tag, only_slope_clamped=False) + assert sorted(r["value"] for r in unclamped) == [2.0, 3.0] # None defaults to 0 + + def test_stability_aggregation(self, db): + tag = "test_v5" + db.write_injection_stat("global_skew", float("nan"), round_number=1, tag=tag) + db.write_injection_stat("global_skew", 1.0, round_number=1, slope_clamped=True, tag=tag) + db.write_injection_stat("global_skew", 2.0, round_number=2, tag=tag) + assert db.flush(timeout=10) is True + + rows = db.query_injection_stat_stability(stat_name="global_skew", tag=tag) + by_round = {r["round_number"]: r for r in rows} + assert by_round[1]["total_count"] == 2 + assert by_round[1]["non_finite_count"] == 1 + assert by_round[1]["clamped_count"] == 1 + assert by_round[2]["total_count"] == 1 + assert by_round[2]["non_finite_count"] == 0 + + +class TestQueryFiltersAndWhitelists: + @pytest.fixture(autouse=True) + def _rows(self, db): + self.db = db + now = time.time() + for round_number in (1, 2, 3): + for tag in ("test_v1", "test_v2"): + db.write_training_stat( + "beta_vae", + "total_loss", + round_number * 1.0, + round_number, + 1, + tag=tag, + timestamp=now + round_number, + ) + assert db.flush(timeout=10) is True + + def test_scalar_filter_uses_equality(self): + rows = self.db.query_training_stat(tag="test_v1") + assert len(rows) == 3 + assert all(r["tag"] == "test_v1" for r in rows) + + def test_list_filter_uses_in(self): + rows = self.db.query_training_stat(tag=["test_v1", "test_v2"]) + assert len(rows) == 6 + + def test_range_filters_inclusive(self): + rows = self.db.query_training_stat(tag="test_v1", start_round_number=2, end_round_number=3) + assert sorted(r["round_number"] for r in rows) == [2, 3] + + def test_column_projection(self): + rows = self.db.query_training_stat(tag="test_v1", columns=["value", "round_number"]) + assert rows + assert all(set(r.keys()) == {"value", "round_number"} for r in rows) + + def test_invalid_column_rejected(self): + with pytest.raises(ValueError, match="Invalid column"): + self.db.query_training_stat(columns=["value; DROP TABLE training_stats;--"]) + + def test_invalid_column_rejected_per_table(self): + # round_number is valid for training_stats but not system_resources. + with pytest.raises(ValueError, match="Invalid column"): + self.db.query_system_resource(columns=["round_number"]) + + def test_empty_list_filter_matches_everything(self): + # NOTE: this documents CURRENT behavior, not desired behavior — an empty IN-list is + # treated as "no filter" (the `if tag:` gate in every query_* skips falsy values), so + # a caller passing a filtered-to-empty list gets the whole table instead of no rows. + # If query_* semantics are ever fixed to "empty list matches nothing", update this + # test deliberately rather than treating the failure as a regression. + assert len(self.db.query_training_stat(tag=[])) == 6 + + def test_confidence_bounds_on_inference_results(self): + self.db.write_inference_result("/a.npy", 0, 1, 0.99, tag="test_v9") + self.db.write_inference_result("/a.npy", 1, 0, 0.42, tag="test_v9") + assert self.db.flush(timeout=10) is True + rows = self.db.query_inference_result(tag="test_v9", min_confidence=0.9) + assert [r["confidence"] for r in rows] == [0.99] + rows = self.db.query_inference_result(tag="test_v9", prediction=0) + assert [r["confidence"] for r in rows] == [0.42] diff --git a/tests/unit/test_manager.py b/tests/unit/test_manager.py new file mode 100644 index 00000000..7eb10f27 --- /dev/null +++ b/tests/unit/test_manager.py @@ -0,0 +1,199 @@ +"""Unit tests for aetherscan.manager: pool / shared-memory tracking and cleanup idempotence. + +ResourceManager registers atexit + signal handlers on construction; the conftest autouse +fixture snapshots/restores signal handlers and unregisters the atexit hook after each test, +so per-test instances don't accumulate process-global state. +""" + +from __future__ import annotations + +import atexit +import multiprocessing +import signal +import sys +import time +from multiprocessing.shared_memory import SharedMemory + +import pytest + +from aetherscan.manager import get_manager, init_manager +from aetherscan.manager.manager import ResourceManager + + +@pytest.fixture +def manager(): + return init_manager() + + +class TestSingletonSemantics: + def test_init_returns_singleton(self, manager): + assert ResourceManager() is manager + assert get_manager() is manager + + def test_reset_produces_fresh_instance(self, manager): + ResourceManager._reset() + try: + assert init_manager() is not manager + finally: + # The conftest teardown only unregisters the *current* instance's atexit hook; + # drop the first instance's hook here so it can't fire at interpreter exit. + atexit.unregister(manager.cleanup_all) + + +class TestPoolTracking: + def test_create_pool_tracks_and_works(self, manager): + pool = manager.create_pool(n_processes=2, name="test-pool") + assert manager.stats.pools_active == 1 + # abs is importable-by-reference in spawned workers — keeps the test light. + assert pool.map(abs, [-2, -3, 4]) == [2, 3, 4] + + def test_close_pool_updates_stats(self, manager): + pool = manager.create_pool(n_processes=2, name="test-pool") + manager.close_pool(pool) + assert manager.stats.pools_active == 0 + assert manager.stats.pools_closed == 1 + assert manager._pools == [] + + def test_close_pool_twice_is_safe(self, manager): + pool = manager.create_pool(n_processes=2, name="test-pool") + manager.close_pool(pool) + manager.close_pool(pool) # logs a warning; must not raise or double-count + assert manager.stats.pools_closed == 1 + + def test_managed_pool_close_is_idempotent(self, manager): + manager.create_pool(n_processes=2, name="test-pool") + managed = manager._pools[0] + managed.close(timeout=10.0) + assert managed.closed is True + managed.close(timeout=10.0) # second call short-circuits on the closed flag + assert managed.closed is True + + def test_multiple_pools_tracked_independently(self, manager): + pool_a = manager.create_pool(n_processes=1, name="a") + pool_b = manager.create_pool(n_processes=1, name="b") + assert manager.stats.pools_active == 2 + manager.close_pool(pool_a) + assert manager.stats.pools_active == 1 + assert manager._pools[0].pool is pool_b + manager.close_pool(pool_b) + assert manager.stats.pools_active == 0 + + +class TestSharedMemoryTracking: + def test_create_tracks_and_is_usable(self, manager): + shm = manager.create_shared_memory(size=1024, name="test-shm") + assert manager.stats.shared_memories_active == 1 + shm.buf[:4] = b"seti" + # Another handle attached by name sees the same bytes (i.e. it's real POSIX shm). + other = SharedMemory(name=shm.name) + try: + assert bytes(other.buf[:4]) == b"seti" + finally: + other.close() + + def test_close_unlinks_and_updates_stats(self, manager): + shm = manager.create_shared_memory(size=1024, name="test-shm") + name = shm.name + manager.close_shared_memory(shm) + assert manager.stats.shared_memories_active == 0 + assert manager.stats.shared_memories_cleaned == 1 + # The segment must be gone from the system namespace. + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + + def test_close_twice_is_safe(self, manager): + shm = manager.create_shared_memory(size=1024, name="test-shm") + manager.close_shared_memory(shm) + manager.close_shared_memory(shm) # logs a warning; must not raise or double-count + assert manager.stats.shared_memories_cleaned == 1 + + +@pytest.mark.skipif( + sys.platform != "linux", + reason="cleanup_all reports process-tree memory via PSS, which psutil only exposes on Linux", +) +class TestCleanupAll: + def test_cleanup_all_closes_everything(self, manager): + manager.create_pool(n_processes=1, name="pool") + shm = manager.create_shared_memory(size=512, name="shm") + name = shm.name + + manager.cleanup_all() + + assert manager.stats.pools_active == 0 + assert manager.stats.pools_closed == 1 + assert manager.stats.shared_memories_active == 0 + assert manager.stats.shared_memories_cleaned == 1 + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + + def test_cleanup_all_is_idempotent(self, manager): + manager.create_pool(n_processes=1, name="pool") + manager.cleanup_all() + closed = manager.stats.pools_closed + manager.cleanup_all() # second call must be a guarded no-op + assert manager.stats.pools_closed == closed + + def test_cleanup_time_recorded(self, manager): + manager.create_shared_memory(size=512, name="shm") + manager.cleanup_all() + assert manager.stats.cleanup_time_seconds >= 0.0 + + +def _sleep_target(ready_event): + """Process target that just idles (module-level so spawn-based platforms can pickle it).""" + ready_event.set() + time.sleep(60) + + +def _sigterm_ignoring_target(ready_event): + """Process target that ignores SIGTERM, forcing the kill escalation path.""" + signal.signal(signal.SIGTERM, signal.SIG_IGN) + ready_event.set() + time.sleep(60) + + +def _start_process(target): + ready = multiprocessing.Event() + process = multiprocessing.Process(target=target, args=(ready,)) + process.start() + assert ready.wait(timeout=30) # don't race the target's signal-handler setup + return process + + +class TestProcessTracking: + def test_register_and_close_process(self, manager): + process = _start_process(_sleep_target) + manager.register_process(process, name="test-process") + assert manager.stats.processes_active == 1 + + manager.close_process(process) + assert not process.is_alive() + assert manager.stats.processes_active == 0 + assert manager.stats.processes_closed == 1 + assert manager._processes == [] + + def test_close_process_twice_is_safe(self, manager): + process = _start_process(_sleep_target) + manager.register_process(process, name="test-process") + manager.close_process(process) + manager.close_process(process) # logs a warning; must not raise or double-count + assert manager.stats.processes_closed == 1 + + def test_close_escalates_to_kill_when_sigterm_ignored(self, manager): + process = _start_process(_sigterm_ignoring_target) + manager.register_process(process, name="stubborn-process") + managed = manager._processes[0] + managed.close(timeout=1.0) # SIGTERM is ignored -> join times out -> SIGKILL + assert managed.closed is True + process.join(timeout=10) + assert not process.is_alive() + + def test_managed_process_close_on_dead_process_is_clean(self, manager): + process = _start_process(_sleep_target) + manager.register_process(process, name="test-process") + process.terminate() + process.join(timeout=10) + # Closing an already-dead process must succeed without escalation + manager.close_process(process) + assert manager.stats.processes_closed == 1 diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py new file mode 100644 index 00000000..24f3a33a --- /dev/null +++ b/tests/unit/test_models.py @@ -0,0 +1,179 @@ +"""Unit tests for aetherscan.models: latent feature layout, RandomForestModel behavior, the +Sampling layer, and encoder/decoder symmetry.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from aetherscan.config import get_config +from aetherscan.models import RandomForestModel, Sampling, prepare_latent_features + + +class TestPrepareLatentFeatures: + def test_row_major_layout(self): + """features[i] must be the row-major concatenation of cadence i's 6 latents.""" + n_cadences, num_obs, latent_dim = 3, 6, 4 + # Encode (cadence, obs, dim) into each value so misordering is unambiguous. + latents = np.zeros((n_cadences * num_obs, latent_dim)) + for cad in range(n_cadences): + for obs in range(num_obs): + for dim in range(latent_dim): + latents[cad * num_obs + obs, dim] = cad * 1000 + obs * 10 + dim + + features = prepare_latent_features(latents, num_observations=num_obs) + + assert features.shape == (n_cadences, num_obs * latent_dim) + for cad in range(n_cadences): + expected = np.concatenate([latents[cad * num_obs + obs] for obs in range(num_obs)]) + np.testing.assert_array_equal(features[cad], expected) + # Spot-check the row-major invariant: feature column obs*latent_dim + dim. + assert features[cad, 1 * latent_dim + 2] == cad * 1000 + 10 + 2 + + def test_indivisible_row_count_raises(self): + with pytest.raises(ValueError, match="Not divisible"): + prepare_latent_features(np.zeros((7, 8)), num_observations=6) + + def test_single_cadence(self): + latents = np.arange(6 * 8, dtype=float).reshape(6, 8) + features = prepare_latent_features(latents, num_observations=6) + np.testing.assert_array_equal(features, latents.ravel()[None, :]) + + +def _toy_latents(n_per_class=16, latent_dim=8, num_obs=6, seed=0): + """Separable toy data: class-1 latents cluster at +2, class-0 at -2.""" + rng = np.random.default_rng(seed) + neg = rng.normal(-2.0, 0.3, size=(n_per_class * num_obs, latent_dim)) + pos = rng.normal(+2.0, 0.3, size=(n_per_class * num_obs, latent_dim)) + latents = np.concatenate([neg, pos]) + labels = np.concatenate([np.zeros(n_per_class), np.ones(n_per_class)]).astype(int) + return latents, labels + + +class TestRandomForestModel: + @pytest.fixture(autouse=True) + def _small_forest(self): + # 1000 trees is overkill for toy data; shrink for speed. + get_config().rf.n_estimators = 10 + + def test_train_and_predict_separable_data(self): + latents, labels = _toy_latents() + model = RandomForestModel() + assert model.is_trained is False + model.train(latents, labels) + assert model.is_trained is True + + predictions = model.predict(latents) + np.testing.assert_array_equal(predictions, labels) + + def test_predict_proba_shape_and_normalization(self): + latents, labels = _toy_latents() + model = RandomForestModel() + model.train(latents, labels) + probas = model.predict_proba(latents) + assert probas.shape == (len(labels), 2) + np.testing.assert_allclose(probas.sum(axis=1), 1.0) + + def test_predict_verbose_confidence_of_predicted_class(self): + latents, labels = _toy_latents() + model = RandomForestModel() + model.train(latents, labels) + predictions, confidences = model.predict_verbose(latents) + probas = model.predict_proba(latents) + np.testing.assert_array_equal(predictions, labels) + # Confidence is P(predicted class) — high for confident negatives too. + expected = np.where(predictions, probas[:, 1], probas[:, 0]) + np.testing.assert_array_equal(confidences, expected) + assert confidences.min() >= 0.5 + + def test_threshold_semantics(self): + latents, labels = _toy_latents() + model = RandomForestModel() + model.train(latents, labels) + # Strict inequality: P(1) > 1.0 is impossible, so everything is class 0. + assert model.predict(latents, threshold=1.0).sum() == 0 + # Threshold 0 flags anything with nonzero positive probability. + assert model.predict(latents, threshold=0.0).sum() >= labels.sum() + + def test_feature_label_mismatch_raises(self): + latents, _ = _toy_latents(n_per_class=4) + wrong_labels = np.zeros(3, dtype=int) + model = RandomForestModel() + with pytest.raises(ValueError, match="mismatch"): + model.train(latents, wrong_labels) + + def test_save_load_round_trip(self, tmp_path): + latents, labels = _toy_latents() + model = RandomForestModel() + model.train(latents, labels) + path = str(tmp_path / "rf.joblib") + model.save(path) + + restored = RandomForestModel() + restored.load(path) + assert restored.is_trained is True + np.testing.assert_array_equal(restored.predict_proba(latents), model.predict_proba(latents)) + + +class TestSamplingLayer: + def test_output_shape(self): + import tensorflow as tf # noqa: PLC0415 + + z_mean = tf.zeros((5, 8)) + z_log_var = tf.zeros((5, 8)) + z = Sampling()([z_mean, z_log_var]) + assert z.shape == (5, 8) + + def test_collapses_to_mean_at_negligible_variance(self): + import tensorflow as tf # noqa: PLC0415 + + z_mean = tf.constant(np.arange(10, dtype=np.float32).reshape(2, 5)) + z_log_var = tf.fill((2, 5), -100.0) # std = exp(-50) ~ 0 + z = Sampling()([z_mean, z_log_var]) + np.testing.assert_allclose(z.numpy(), z_mean.numpy(), atol=1e-6) + + def test_variance_injects_noise(self): + import tensorflow as tf # noqa: PLC0415 + + z_mean = tf.zeros((4, 8)) + z_log_var = tf.zeros((4, 8)) # std = 1 + z1 = Sampling()([z_mean, z_log_var]).numpy() + z2 = Sampling()([z_mean, z_log_var]).numpy() + assert not np.allclose(z1, z2) + assert np.std(z1) > 0.1 + + +@pytest.mark.slow +class TestEncoderDecoderSymmetry: + """Builds the full Beta-VAE graph on CPU — slow but CI-safe.""" + + def test_symmetry_and_forward_pass(self): + import tensorflow as tf # noqa: PLC0415 + + from aetherscan.models import create_beta_vae_model # noqa: PLC0415 + + vae = create_beta_vae_model() + config = get_config() + latent_dim = config.beta_vae.latent_dim + dense_size = config.beta_vae.dense_layer_size + + # Encoder consumes single observations; decoder must emit the exact mirror shape. + assert vae.encoder.input_shape == (None, 16, dense_size, 1) + z_mean_shape, z_log_var_shape, z_shape = vae.encoder.output_shape + assert z_mean_shape == (None, latent_dim) + assert z_log_var_shape == (None, latent_dim) + assert z_shape == (None, latent_dim) + assert vae.decoder.input_shape == (None, latent_dim) + assert vae.decoder.output_shape == vae.encoder.input_shape + + # Forward pass round-trips the cadence shape (same graph — built once per test run). + batch = tf.random.uniform((2, 6, 16, dense_size)) + reconstruction, z_mean, z_log_var, z = vae(batch, training=False) + assert reconstruction.shape == (2, 6, 16, dense_size) + # The encoder operates per observation: 2 cadences * 6 observations. + assert z_mean.shape == (12, latent_dim) + assert z.shape == (12, latent_dim) + # Sigmoid output stays within [0, 1] like the log-normed inputs. + recon = reconstruction.numpy() + assert recon.min() >= 0.0 + assert recon.max() <= 1.0 diff --git a/tests/unit/test_preprocessing.py b/tests/unit/test_preprocessing.py new file mode 100644 index 00000000..d6acd2e5 --- /dev/null +++ b/tests/unit/test_preprocessing.py @@ -0,0 +1,235 @@ +"""Unit tests for aetherscan.preprocessing: hit deduplication, CSV cadence grouping, filename +sanitization, JSON coercion, DC-spike removal, and spline bandpass fitting.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from aetherscan.config import get_config +from aetherscan.preprocessing import ( + DataPreprocessor, + _fit_channel_bandpass, + _read_coarse_channel_worker, + _remove_dc_spike, + group_observations_from_csv, +) + + +@pytest.fixture +def group_cols(): + """The real grouping columns from InferenceConfig — read from the config rather than + duplicated, so these tests track the defaults the pipeline actually uses.""" + return list(get_config().inference.cadence_group_by_cols) + + +@pytest.fixture +def h5_col(): + return get_config().inference.cadence_h5_path_col + + +class TestDeduplicateHits: + def test_empty_input(self): + assert DataPreprocessor._deduplicate_hits([], stamp_width=512) == [] + + def test_far_apart_hits_survive(self): + hits = [(0, 5.0, 0.01), (300, 4.0, 0.01)] + assert DataPreprocessor._deduplicate_hits(hits, stamp_width=512) == hits + + def test_nearby_hits_merge_keeping_higher_statistic(self): + hits = [(0, 5.0, 0.01), (100, 9.0, 0.02)] + assert DataPreprocessor._deduplicate_hits(hits, stamp_width=512) == [(100, 9.0, 0.02)] + + def test_merge_is_greedy_left_to_right(self): + # 0 and 100 merge (keep 100 @ 9.0); 300 is then within 256 of 100 but weaker, + # so it merges away too — one survivor. + hits = [(0, 5.0, 0.01), (100, 9.0, 0.02), (300, 2.0, 0.03)] + assert DataPreprocessor._deduplicate_hits(hits, stamp_width=512) == [(100, 9.0, 0.02)] + + def test_input_order_does_not_matter(self): + hits = [(300, 2.0, 0.03), (0, 5.0, 0.01), (100, 9.0, 0.02)] + assert DataPreprocessor._deduplicate_hits(hits, stamp_width=512) == [(100, 9.0, 0.02)] + + def test_boundary_distance_exactly_half_width_kept(self): + # Merge condition is strict (<): a gap of exactly stamp_width // 2 keeps both. + hits = [(0, 5.0, 0.01), (256, 4.0, 0.02)] + assert DataPreprocessor._deduplicate_hits(hits, stamp_width=512) == hits + + +class TestGroupObservationsFromCsv: + def test_valid_and_flagged_groups(self, make_inference_csv, group_cols, h5_col): + key_a = { + "Target": "HIP110750", + "Session": "AGBT21B_999_31", + "Band": "L", + "Cadence ID": "0", + "Frequency": "1400", + } + key_b = {**key_a, "Cadence ID": "1"} + csv_path = make_inference_csv( + "two_groups.csv", + groups=[ + (key_a, [f"/data/a_{i}.h5" for i in range(6)]), + (key_b, [f"/data/b_{i}.h5" for i in range(3)]), # short cadence -> flagged + ], + ) + valid, flagged = group_observations_from_csv( + str(csv_path), group_cols, h5_col, expected_obs=6 + ) + assert len(valid) == 1 + assert len(flagged) == 1 + assert valid[0].is_valid is True + assert valid[0].key == tuple(key_a[c] for c in group_cols) + # Row order within the group is preserved. + assert valid[0].h5_paths == [f"/data/a_{i}.h5" for i in range(6)] + assert flagged[0].is_valid is False + assert len(flagged[0].h5_paths) == 3 + + def test_missing_column_raises_keyerror(self, make_inference_csv, group_cols, h5_col): + csv_path = make_inference_csv("ok.csv") + with pytest.raises(KeyError, match="missing required column"): + group_observations_from_csv( + str(csv_path), [*group_cols, "Nonexistent"], h5_col, expected_obs=6 + ) + + def test_missing_file_raises(self, group_cols, h5_col): + with pytest.raises(FileNotFoundError): + group_observations_from_csv("/nope.csv", group_cols, h5_col) + + def test_expected_obs_parameter(self, make_inference_csv, group_cols, h5_col): + key = { + "Target": "T", + "Session": "S", + "Band": "L", + "Cadence ID": "0", + "Frequency": "1", + } + csv_path = make_inference_csv("three.csv", groups=[(key, ["/a.h5", "/b.h5", "/c.h5"])]) + valid, flagged = group_observations_from_csv( + str(csv_path), group_cols, h5_col, expected_obs=3 + ) + assert len(valid) == 1 + assert flagged == [] + + +class TestCadenceNpyFilename: + def test_clean_key_passes_through(self): + name = DataPreprocessor._cadence_npy_filename("catalog", ("HIP110750", "L", "0")) + assert name == "catalog_HIP110750_L_0.npy" + + def test_hostile_characters_collapse_to_underscore(self): + name = DataPreprocessor._cadence_npy_filename( + "catalog", ("HIP 110750", "AGBT/18A", "a;b|c", "$(rm -rf)") + ) + assert "/" not in name + assert " " not in name + assert ";" not in name + assert "$" not in name + assert name.startswith("catalog_") + assert name.endswith(".npy") + + def test_dash_and_dot_survive(self): + name = DataPreprocessor._cadence_npy_filename("catalog", ("L-band", "1.5")) + assert name == "catalog_L-band_1.5.npy" + + def test_non_string_key_components(self): + name = DataPreprocessor._cadence_npy_filename("catalog", (0, 1400.5)) + assert name == "catalog_0_1400.5.npy" + + +class TestToJsonSafe: + def test_bytes_decode(self): + assert DataPreprocessor._to_json_safe(b"HIP110750") == "HIP110750" + + def test_numpy_scalars(self): + assert DataPreprocessor._to_json_safe(np.float32(1.5)) == 1.5 + assert DataPreprocessor._to_json_safe(np.int64(7)) == 7 + assert DataPreprocessor._to_json_safe(np.bool_(True)) is True + + def test_ndarray_to_list(self): + result = DataPreprocessor._to_json_safe(np.arange(3, dtype=np.int32)) + assert result == [0, 1, 2] + assert all(isinstance(v, int) for v in result) + + def test_nested_structures_and_keys(self): + obj = { + 1: {"vals": np.array([1.0, 2.0], dtype=np.float64)}, + "meta": (b"x", np.int16(2)), + } + result = DataPreprocessor._to_json_safe(obj) + assert result == {"1": {"vals": [1.0, 2.0]}, "meta": ["x", 2]} + + def test_plain_values_pass_through(self): + assert DataPreprocessor._to_json_safe("s") == "s" + assert DataPreprocessor._to_json_safe(3.5) == 3.5 + assert DataPreprocessor._to_json_safe(None) is None + + +class TestRemoveDcSpike: + def test_spike_replaced_by_neighbor_interpolation(self): + rng = np.random.default_rng(9) + coarse_width, n_coarse = 32, 3 + block = rng.chisquare(df=4, size=(4, n_coarse * coarse_width)) + original = block.copy() + + _remove_dc_spike(block, coarse_width, n_coarse) + + half = coarse_width // 2 + for i in range(n_coarse): + dc = i * coarse_width + half + np.testing.assert_allclose( + block[:, dc], (original[:, dc + 1] + original[:, dc - 3]) / 2 + ) + np.testing.assert_allclose( + block[:, dc - 1], (original[:, dc + 2] + original[:, dc - 2]) / 2 + ) + # Everything else in the channel is untouched. + untouched = [ + j for j in range(i * coarse_width, (i + 1) * coarse_width) if j not in (dc, dc - 1) + ] + np.testing.assert_array_equal(block[:, untouched], original[:, untouched]) + + def test_visible_spike_is_flattened(self): + coarse_width = 64 + block = np.ones((4, coarse_width)) + dc = coarse_width // 2 + block[:, dc] = 100.0 + block[:, dc - 1] = 100.0 + _remove_dc_spike(block, coarse_width, 1) + np.testing.assert_allclose(block, np.ones((4, coarse_width))) + + +class TestFitChannelBandpass: + def test_smooth_bandpass_recovered(self): + channel_width, spl_order = 1024, 16 + x = np.arange(channel_width) + bandpass = 100.0 + 10.0 * np.sin(2 * np.pi * x / channel_width) + fit = _fit_channel_bandpass(bandpass, channel_width, spl_order) + assert fit.shape == bandpass.shape + # A smooth curve must be recovered nearly exactly by the spline. + assert np.max(np.abs(fit - bandpass)) < 0.1 + + def test_subtracting_fit_flattens_channel(self): + rng = np.random.default_rng(2) + channel_width, spl_order = 1024, 16 + x = np.arange(channel_width) + bandpass = 100.0 + 20.0 * np.cos(2 * np.pi * x / channel_width) + noisy = bandpass + rng.normal(0, 0.5, size=channel_width) + residual = noisy - _fit_channel_bandpass(noisy, channel_width, spl_order) + # Residuals lose the 20-unit bandpass structure and keep only ~noise-scale variance. + assert np.abs(residual.mean()) < 0.5 + assert residual.std() < 1.0 + + +class TestReadCoarseChannelWorker: + def test_reads_one_coarse_channel(self, make_h5_observation): + n_chans, coarse_width = 2048, 512 + h5_path = make_h5_observation("obs.h5", n_chans=n_chans) + channel = _read_coarse_channel_worker((str(h5_path), 2, coarse_width)) + assert channel.shape == (16, coarse_width) + + import h5py # noqa: PLC0415 + + with h5py.File(h5_path, "r") as hf: + expected = hf["data"][:, 0, 2 * coarse_width : 3 * coarse_width] + np.testing.assert_array_equal(channel, expected) diff --git a/tests/unit/test_round_data.py b/tests/unit/test_round_data.py new file mode 100644 index 00000000..fae1afb4 --- /dev/null +++ b/tests/unit/test_round_data.py @@ -0,0 +1,622 @@ +"""Unit tests for aetherscan.round_data (paths, manifests, startup dir cleanup, producer +protocol) and the batched memmap generation in aetherscan.data_generation.""" + +from __future__ import annotations + +import os +import queue +import random +import threading +import time +from multiprocessing.shared_memory import SharedMemory + +import numpy as np +import pytest + +from aetherscan.data_generation import ( + build_chunk_segments, + build_segment_tasks, + generate_round_to_memmap, +) +from aetherscan.round_data import ( + RoundDataPaths, + RoundDataProducer, + _producer_main, + build_manifest, + load_round_arrays, + prepare_round_data_dir, + validate_done_manifest, + write_done_manifest, +) + +# Keep injection fast: small frequency axis, real-ish resolutions (mirrors +# tests/unit/test_data_generation.py). +_WIDTH_BIN = 128 +_FREQ_RES = 2.7939677238464355 # Hz +_TIME_RES = 18.25361108 # seconds + +_SIGNAL_TYPES = ["false_no_signal", "false_with_rfi", "true_only_eti", "true_eti_rfi"] + + +@pytest.fixture(autouse=True) +def _seed_rngs(): + random.seed(11) + np.random.seed(11) + + +def _write_small_round(paths: RoundDataPaths, n_samples=8, width_bin=16, snr_base=10.0): + """Write a tiny but structurally-complete round dataset + validated manifest.""" + os.makedirs(paths.round_dir, exist_ok=True) + rng = np.random.default_rng(paths.round_idx + 1) + shape = (n_samples, 6, 4, width_bin) + for path in paths.array_paths.values(): + arr = np.lib.format.open_memmap(path, mode="w+", dtype=np.float32, shape=shape) + arr[:] = rng.random(shape, dtype=np.float32) + del arr + labels = np.array( + [_SIGNAL_TYPES[i % 4] for i in range(n_samples)], + dtype="U20", + ) + np.save(paths.labels_path, labels) + manifest = build_manifest( + paths, + n_samples=n_samples, + snr_base=snr_base, + snr_range=40.0, + wall_time_s=1.0, + chunk_count=1, + ) + write_done_manifest(paths, manifest) + return manifest + + +class TestRoundDataPaths: + def test_for_round_naming(self, tmp_path): + paths = RoundDataPaths.for_round(str(tmp_path), 3) + assert paths.round_dir == str(tmp_path / "round_03") + assert paths.round_idx == 3 + assert paths.done_path.endswith("round_03.done") + + def test_array_and_label_paths(self, tmp_path): + paths = RoundDataPaths.for_round(str(tmp_path), 12) + assert set(paths.array_paths.keys()) == {"main", "true", "false"} + for name, path in paths.array_paths.items(): + assert path == os.path.join(paths.round_dir, f"{name}.npy") + assert paths.labels_path == os.path.join(paths.round_dir, "labels.npy") + + +class TestManifest: + def test_roundtrip_validates(self, tmp_path): + paths = RoundDataPaths.for_round(str(tmp_path), 1) + written = _write_small_round(paths) + manifest = validate_done_manifest(paths, expected_n_samples=8) + assert manifest is not None + assert manifest["n_samples"] == written["n_samples"] == 8 + assert manifest["round_idx"] == 1 + assert set(manifest["checksums"].keys()) == {"main", "true", "false", "labels"} + + def test_missing_done_file_invalid(self, tmp_path): + paths = RoundDataPaths.for_round(str(tmp_path), 1) + _write_small_round(paths) + os.remove(paths.done_path) + assert validate_done_manifest(paths) is None + + def test_expected_n_samples_mismatch_invalid(self, tmp_path): + paths = RoundDataPaths.for_round(str(tmp_path), 1) + _write_small_round(paths, n_samples=8) + assert validate_done_manifest(paths, expected_n_samples=16) is None + + def test_round_idx_mismatch_invalid(self, tmp_path): + # A manifest written for round 1 dropped into a round-2 dir must not validate. + paths_1 = RoundDataPaths.for_round(str(tmp_path), 1) + manifest = _write_small_round(paths_1) + paths_2 = RoundDataPaths.for_round(str(tmp_path), 2) + _write_small_round(paths_2) + write_done_manifest(paths_2, manifest) + assert validate_done_manifest(paths_2) is None + + def test_missing_array_invalid(self, tmp_path): + paths = RoundDataPaths.for_round(str(tmp_path), 1) + _write_small_round(paths) + os.remove(paths.true_path) + assert validate_done_manifest(paths) is None + + def test_shape_mismatch_invalid(self, tmp_path): + paths = RoundDataPaths.for_round(str(tmp_path), 1) + _write_small_round(paths) + np.save(paths.main_path, np.zeros((2, 2), dtype=np.float32)) + assert validate_done_manifest(paths) is None + + def test_corrupted_values_invalid(self, tmp_path): + paths = RoundDataPaths.for_round(str(tmp_path), 1) + _write_small_round(paths) + # Same shape, different content: the sampled-value checksum must catch it. + arr = np.load(paths.false_path, mmap_mode="r+") + arr[:] = arr[:] + 7.0 + arr.flush() + del arr + assert validate_done_manifest(paths) is None + + def test_corrupted_manifest_json_invalid(self, tmp_path): + paths = RoundDataPaths.for_round(str(tmp_path), 1) + _write_small_round(paths) + with open(paths.done_path, "w") as f: + f.write("{not json") + assert validate_done_manifest(paths) is None + + def test_load_round_arrays_returns_memmaps(self, tmp_path): + paths = RoundDataPaths.for_round(str(tmp_path), 1) + _write_small_round(paths, n_samples=8, width_bin=16) + data = load_round_arrays(paths) + assert set(data.keys()) == {"concatenated", "true", "false", "labels"} + for key in ("concatenated", "true", "false"): + assert isinstance(data[key], np.memmap) + assert data[key].shape == (8, 6, 4, 16) + assert data[key].dtype == np.float32 + assert data["labels"].shape == (8,) + assert not isinstance(data["labels"], np.memmap) + + +class TestPrepareRoundDataDir: + def test_creates_missing_base_dir(self, tmp_path): + base = tmp_path / "round_data" / "test_v1" + prepare_round_data_dir(str(base), start_round=1) + assert base.is_dir() + + def test_resume_semantics(self, tmp_path): + base = str(tmp_path) + # round 1: valid manifest (< start_round -> kept) + _write_small_round(RoundDataPaths.for_round(base, 1)) + # round 2: no .done (< start_round -> deleted) + paths_2 = RoundDataPaths.for_round(base, 2) + _write_small_round(paths_2) + os.remove(paths_2.done_path) + # round 3: valid manifest but >= start_round -> deleted (regenerated by this run) + _write_small_round(RoundDataPaths.for_round(base, 3)) + # non-round entry (stale RF dataset) -> deleted + os.makedirs(os.path.join(base, "rf")) + # loose file -> untouched + with open(os.path.join(base, "stray.txt"), "w") as f: + f.write("x") + + prepare_round_data_dir(base, start_round=3) + + assert os.path.isdir(os.path.join(base, "round_01")) + assert not os.path.exists(os.path.join(base, "round_02")) + assert not os.path.exists(os.path.join(base, "round_03")) + assert not os.path.exists(os.path.join(base, "rf")) + assert os.path.isfile(os.path.join(base, "stray.txt")) + + def test_fresh_run_deletes_everything(self, tmp_path): + base = str(tmp_path) + _write_small_round(RoundDataPaths.for_round(base, 1)) + _write_small_round(RoundDataPaths.for_round(base, 2)) + prepare_round_data_dir(base, start_round=1) + assert not os.path.exists(os.path.join(base, "round_01")) + assert not os.path.exists(os.path.join(base, "round_02")) + + +def _segment_coverage(segments, array_name): + """Row indices covered by `array_name`'s segments (with multiplicity).""" + covered = [] + for segment in segments: + if segment.array_name == array_name: + covered.extend(range(segment.start_idx, segment.start_idx + segment.count)) + return covered + + +class TestChunkSegmentsAndTasks: + def test_segments_cover_each_array_exactly_once(self): + segments = build_chunk_segments(chunk_start=8, chunk_size=8) + assert len(segments) == 8 + for array_name in ("main", "false", "true"): + covered = _segment_coverage(segments, array_name) + assert sorted(covered) == list(range(8, 16)) # exactly once, no gaps/overlaps + + def test_main_segment_order_matches_labels(self): + segments = build_chunk_segments(chunk_start=0, chunk_size=8) + main_types = [s.signal_type for s in segments if s.array_name == "main"] + assert main_types == _SIGNAL_TYPES + + def test_chunk_size_not_divisible_by_4_raises(self): + with pytest.raises(ValueError, match="divisible by 4"): + build_chunk_segments(chunk_start=0, chunk_size=6) + + def test_task_partitioning_covers_every_row_exactly_once(self): + segments = build_chunk_segments(chunk_start=0, chunk_size=40) + seed_rng = np.random.default_rng(0) + for segment in segments: + tasks = build_segment_tasks( + segment, "arr.npy", 3, 10.0, 40.0, _WIDTH_BIN, _FREQ_RES, _TIME_RES, seed_rng + ) + covered = [] + for _, start_idx, count, *_rest in tasks: + covered.extend(range(start_idx, start_idx + count)) + assert sorted(covered) == list( + range(segment.start_idx, segment.start_idx + segment.count) + ) + # Every task respects the max size and carries the segment's generation params + assert all(task[2] <= 3 for task in tasks) + assert all(task[3] == segment.create_fn_name for task in tasks) + + def test_task_size_below_one_raises(self): + segment = build_chunk_segments(0, 4)[0] + with pytest.raises(ValueError, match="task_size"): + build_segment_tasks( + segment, + "arr.npy", + 0, + 10.0, + 40.0, + _WIDTH_BIN, + _FREQ_RES, + _TIME_RES, + np.random.default_rng(0), + ) + + +class TestGenerateRoundToMemmap: + @pytest.fixture + def plate(self, make_background_npy): + path = make_background_npy("plate.npy", n_cadences=4, width_bin=_WIDTH_BIN) + return np.load(path) + + def test_sequential_generation_end_to_end(self, tmp_path, plate): + paths = RoundDataPaths.for_round(str(tmp_path), 1) + stats_segments = [] + progress = [] + + manifest = generate_round_to_memmap( + paths, + n_samples=8, + snr_base=10.0, + snr_range=5.0, + width_bin=_WIDTH_BIN, + num_observations=6, + time_bins=16, + chunk_size=4, + task_size=3, + freq_resolution=_FREQ_RES, + time_resolution=_TIME_RES, + pool=None, + backgrounds=plate, + round_num=1, + stats_cb=stats_segments.append, + progress_cb=lambda chunk, n_chunks: progress.append((chunk, n_chunks)), + ) + + # Arrays on disk with the right shape/dtype, every row populated & log-normalized + data = load_round_arrays(paths) + for key in ("concatenated", "true", "false"): + arr = data[key] + assert arr.shape == (8, 6, 16, _WIDTH_BIN) + assert arr.dtype == np.float32 + row_maxes = arr.max(axis=(1, 2, 3)) + assert np.all(row_maxes > 0) # no row was left unwritten + assert float(arr.max()) <= 1.0 + assert float(arr.min()) >= 0.0 + + # Labels mirror the contiguous per-chunk layout (chunk_size=4 -> quarter=1) + assert list(data["labels"]) == _SIGNAL_TYPES + _SIGNAL_TYPES + + # Manifest validates and matches the generation request + assert validate_done_manifest(paths, expected_n_samples=8) is not None + assert manifest["chunk_count"] == 2 + assert manifest["snr_base"] == 10.0 + + # Stats: 8 class-segments per chunk x 2 chunks; per-segment sample counts add up + assert len(stats_segments) == 16 + for segment in stats_segments: + assert segment["num_samples"] == len(segment["stats_list"]) + assert segment["round_number"] == 1 + assert segment["snr_range_ceil"] == 15.0 + main_total = sum(s["num_samples"] for s in stats_segments if s["signal_class"] == "main") + false_total = sum(s["num_samples"] for s in stats_segments if s["signal_class"] == "false") + true_total = sum(s["num_samples"] for s in stats_segments if s["signal_class"] == "true") + assert main_total == false_total == true_total == 8 + + # Signal-info keys match the signal type + by_type = {(s["signal_class"], s["signal_type"]): s for s in stats_segments} + assert by_type[("main", "false_no_signal")]["stats_list"][0]["signal_info"] == {} + double_keys = set(by_type[("true", "true_eti_rfi")]["stats_list"][0]["signal_info"]) + assert any(k.startswith("eti_") for k in double_keys) + assert any(k.startswith("rfi_") for k in double_keys) + + assert progress == [(1, 2), (2, 2)] + + def test_regeneration_clears_stale_dir(self, tmp_path, plate): + paths = RoundDataPaths.for_round(str(tmp_path), 1) + os.makedirs(paths.round_dir) + stale = os.path.join(paths.round_dir, "leftover.npy") + with open(stale, "w") as f: + f.write("stale") + generate_round_to_memmap( + paths, + n_samples=4, + snr_base=10.0, + snr_range=5.0, + width_bin=_WIDTH_BIN, + num_observations=6, + time_bins=16, + chunk_size=4, + task_size=2, + freq_resolution=_FREQ_RES, + time_resolution=_TIME_RES, + backgrounds=plate, + ) + assert not os.path.exists(stale) + assert validate_done_manifest(paths, expected_n_samples=4) is not None + + def test_invalid_inputs_raise(self, tmp_path, plate): + paths = RoundDataPaths.for_round(str(tmp_path), 1) + common = { + "width_bin": _WIDTH_BIN, + "num_observations": 6, + "time_bins": 16, + "task_size": 2, + "freq_resolution": _FREQ_RES, + "time_resolution": _TIME_RES, + "backgrounds": plate, + } + with pytest.raises(ValueError, match="n_samples"): + generate_round_to_memmap(paths, 6, 10.0, 5.0, chunk_size=4, **common) + with pytest.raises(ValueError, match="chunk_size"): + generate_round_to_memmap(paths, 8, 10.0, 5.0, chunk_size=6, **common) + with pytest.raises(ValueError, match="backgrounds"): + generate_round_to_memmap( + paths, + 8, + 10.0, + 5.0, + chunk_size=4, + **{**common, "backgrounds": None}, + ) + + +def _stub_generate_ok(paths, round_idx, snr_base, snr_range, pool, params, stats_cb, progress_cb): + """Producer-protocol stub: emits one stats + one progress message, returns a manifest.""" + stats_cb({"signal_class": "main", "stats_list": []}) + progress_cb(1, 1) + return {"round_idx": round_idx, "n_samples": params["n_samples"], "stub": True} + + +def _stub_generate_boom(paths, round_idx, snr_base, snr_range, pool, params, stats_cb, progress_cb): + raise RuntimeError("stub generation exploded") + + +class TestProducerProtocol: + """Drive _producer_main in a thread with plain queues and a stub generate_fn — validates + the message protocol without child processes or the setigen generation stack.""" + + def _run(self, tmp_path, generate_fn, requests): + request_queue: queue.Queue = queue.Queue() + result_queue: queue.Queue = queue.Queue() + params = {"base_dir": str(tmp_path), "n_samples": 8} + for request in requests: + request_queue.put(request) + request_queue.put(("shutdown",)) + thread = threading.Thread( + target=_producer_main, + args=(request_queue, result_queue, params), + kwargs={"generate_fn": generate_fn}, + daemon=True, + ) + thread.start() + thread.join(timeout=30) + assert not thread.is_alive() + messages = [] + while True: + try: + messages.append(result_queue.get_nowait()) + except queue.Empty: + break + return messages + + def test_generate_emits_stats_progress_done(self, tmp_path): + messages = self._run(tmp_path, _stub_generate_ok, [("generate", 1, 10, 40)]) + kinds = [m[0] for m in messages] + assert kinds == ["stats", "progress", "done", "shutdown_ack"] + done = messages[2] + assert done[1] == 1 + assert done[2]["stub"] is True + + def test_error_is_reported_and_producer_keeps_serving(self, tmp_path): + messages = self._run( + tmp_path, + _stub_generate_boom, + [("generate", 1, 10, 40), ("generate", 2, 10, 40)], + ) + kinds = [m[0] for m in messages] + assert kinds == ["error", "error", "shutdown_ack"] + assert messages[0][1] == 1 + assert "stub generation exploded" in messages[0][2] + assert messages[1][1] == 2 + + def test_existing_valid_round_short_circuits(self, tmp_path): + # A validated on-disk round must be reused: done comes back without generate_fn firing. + _write_small_round(RoundDataPaths.for_round(str(tmp_path), 1), n_samples=8) + messages = self._run(tmp_path, _stub_generate_boom, [("generate", 1, 10, 40)]) + kinds = [m[0] for m in messages] + assert kinds == ["done", "shutdown_ack"] + assert messages[0][2]["n_samples"] == 8 + + def test_unknown_message_ignored(self, tmp_path): + messages = self._run(tmp_path, _stub_generate_ok, [("bogus",)]) + assert [m[0] for m in messages] == ["shutdown_ack"] + + +class _FakeProcess: + """Stand-in for the producer multiprocessing.Process on the main-side handle.""" + + def __init__(self): + self.alive = True + self.pid = 4242 + + def is_alive(self): + return self.alive + + +class _FakeDB: + def __init__(self): + self.writes = [] + + def write_injection_stat(self, **kwargs): + self.writes.append(kwargs) + + +def _minimal_sample_info(): + stage = dict.fromkeys( + [ + "global_mean", + "global_median", + "global_std", + "global_mad", + "global_skew", + "global_kurtosis", + ], + 0.5, + ) + return { + "background_index": 0, + "intensity_stats": {"A": dict(stage), "B": dict(stage), "C": dict(stage)}, + "signal_info": {}, + "slope_was_clamped": False, + } + + +def _minimal_segment(): + return { + "round_number": 1, + "chunk_number": 1, + "signal_class": "main", + "signal_type": "false_no_signal", + "snr_range_floor": 10, + "snr_range_ceil": 50, + "num_samples": 1, + "inject_duration": 0.1, + "timestamp": time.time(), + "stats_list": [_minimal_sample_info()], + } + + +class TestRoundDataProducerDrainer: + """Exercise the main-side drainer thread + await_round against a fake process, feeding + the result queue directly (no child process involved).""" + + @pytest.fixture + def producer(self, tmp_path): + producer = RoundDataProducer( + base_dir=str(tmp_path), + n_samples=8, + shm_name="unused", + background_shape=(1, 6, 4, 16), + background_dtype="float32", + n_processes=1, + width_bin=16, + num_observations=6, + time_bins=4, + chunk_size=4, + task_size=2, + freq_resolution=_FREQ_RES, + time_resolution=_TIME_RES, + db=_FakeDB(), + tag="test_v1", + ) + producer._process = _FakeProcess() + producer._drainer = threading.Thread(target=producer._drain_results, daemon=True) + producer._drainer.start() + yield producer + producer._process.alive = False + producer._drainer.join(timeout=10) + + def test_done_unblocks_await_round(self, producer): + producer._result_queue.put(("done", 1, {"n_samples": 8})) + assert producer.await_round(1) == {"n_samples": 8} + + def test_error_raises_with_producer_traceback(self, producer): + producer._result_queue.put(("error", 2, "Traceback: boom")) + with pytest.raises(RuntimeError, match="boom"): + producer.await_round(2) + + def test_stats_are_written_to_db_from_main_process(self, producer): + producer._result_queue.put(("stats", 1, _minimal_segment())) + producer._result_queue.put(("done", 1, {"n_samples": 8})) + producer.await_round(1) + # 18 per-sample intensity rows (6 stats x 3 stages) + 0 signal rows + 4 metadata rows + deadline = time.time() + 10 + while len(producer._db.writes) < 22 and time.time() < deadline: + time.sleep(0.05) + assert len(producer._db.writes) == 22 + metadata_names = {w["stat_name"] for w in producer._db.writes if w["sample_index"] is None} + assert metadata_names == { + "snr_range_floor", + "snr_range_ceil", + "num_samples", + "inject_duration", + } + assert all(w["tag"] == "test_v1" for w in producer._db.writes) + + def test_producer_death_unblocks_await_round(self, producer): + producer._process.alive = False + with pytest.raises(RuntimeError, match="exited before producing"): + producer.await_round(5) + + +@pytest.mark.slow +class TestRoundDataProducerSpawnEndToEnd: + """Real spawn-started producer process + real (tiny) generation against real shared + memory — exercises the spawn pickling boundary, the child's import chain, its pool + creation, and the stats/done/shutdown protocol for real. Both cluster-smoke failures + this PR hit (a fork-inherited deadlocked lock; a fork-context SemLock crossing the + spawn boundary) were only reachable through a real child process.""" + + def test_spawned_producer_generates_round(self, tmp_path): + rng = np.random.default_rng(11) + plate = rng.chisquare(df=4, size=(4, 6, 16, _WIDTH_BIN)).astype(np.float32) + + shm = SharedMemory(create=True, size=plate.nbytes) + producer = None + try: + shared = np.ndarray(plate.shape, dtype=plate.dtype, buffer=shm.buf) + shared[:] = plate + + db = _FakeDB() + producer = RoundDataProducer( + base_dir=str(tmp_path), + n_samples=8, + shm_name=shm.name, + background_shape=plate.shape, + background_dtype=str(plate.dtype), + n_processes=2, + width_bin=_WIDTH_BIN, + num_observations=6, + time_bins=16, + chunk_size=4, + task_size=3, + freq_resolution=_FREQ_RES, + time_resolution=_TIME_RES, + db=db, + tag="test_v1", + ) + producer.start() + producer.request_generation(1, 10, 5) + manifest = producer.await_round(1) + assert manifest["n_samples"] == 8 + + paths = RoundDataPaths.for_round(str(tmp_path), 1) + assert validate_done_manifest(paths, expected_n_samples=8) is not None + + # Streamed stats land in the main-process drainer (DB writes stay in main) + deadline = time.time() + 30 + while not db.writes and time.time() < deadline: + time.sleep(0.2) + assert db.writes + + # A second request for the already-generated round short-circuits via the manifest + producer.request_generation(1, 10, 5) + assert producer.await_round(1)["n_samples"] == 8 + finally: + if producer is not None: + producer.shutdown() + shm.close() + shm.unlink() diff --git a/tests/unit/test_train_datasets.py b/tests/unit/test_train_datasets.py new file mode 100644 index 00000000..5399ee03 --- /dev/null +++ b/tests/unit/test_train_datasets.py @@ -0,0 +1,193 @@ +"""Unit tests for train.py's batched distributed dataset builders: whole-global-batch yields, +exact index coverage, stratified splits, the shuffle=False alignment contract that +train_random_forest depends on, and the viz dataset's order guarantee.""" + +from __future__ import annotations + +import numpy as np +import pytest +import tensorflow as tf + +from aetherscan.train import ( + TrainDataHolder, + prepare_distributed_train_dataset, + prepare_distributed_viz_dataset, +) + +# Tiny sample shape — the builders never hardcode (6, 16, 512) +_SAMPLE_SHAPE = (2, 3, 4) +_SIGNAL_TYPES = ["false_no_signal", "false_with_rfi", "true_only_eti", "true_eti_rfi"] + +pytestmark = pytest.mark.slow # builds TF graphs on CPU + + +def _make_data(n_samples=40): + """Arrays whose every sample is a constant equal to (offset + row index), so a yielded + batch reveals exactly which rows it gathered.""" + base = np.arange(n_samples, dtype=np.float32)[:, None, None, None] + ones = np.ones((n_samples, *_SAMPLE_SHAPE), dtype=np.float32) + labels = np.array([_SIGNAL_TYPES[(4 * i) // n_samples] for i in range(n_samples)], dtype="U20") + return { + "concatenated": base * ones, + "true": (base + 1000.0) * ones, + "false": (base + 2000.0) * ones, + "labels": labels, + } + + +def _build(data, shuffle, train_val_split=0.8, prb=4, eb=8, prvb=4): + return prepare_distributed_train_dataset( + data=data, + train_val_split=train_val_split, + per_replica_batch_size=prb, + effective_batch_size=eb, + per_replica_val_batch_size=prvb, + num_replicas=1, + strategy=tf.distribute.get_strategy(), + shuffle=shuffle, + ) + + +def _batch_row_ids(batch): + """Recover the source row indices a yielded ((c, t, f), y) batch gathered.""" + (concat, true, false), y = batch + concat_ids = concat.numpy()[:, 0, 0, 0] + np.testing.assert_array_equal(concat_ids, y.numpy()[:, 0, 0, 0]) + np.testing.assert_array_equal(concat_ids + 1000.0, true.numpy()[:, 0, 0, 0]) + np.testing.assert_array_equal(concat_ids + 2000.0, false.numpy()[:, 0, 0, 0]) + return concat_ids.astype(np.int64) + + +class TestPrepareDistributedTrainDataset: + def test_step_counts_and_batched_shapes(self): + results = _build(_make_data(), shuffle=True) + assert results["n_train_trimmed"] == 32 + assert results["n_val_trimmed"] == 8 + assert results["train_steps"] == 4 # 32 / effective(8) + assert results["accumulation_steps"] == 2 # effective(8) / global(4) + assert results["val_steps"] == 2 # 8 / global val(4) + + iterator = iter(results["train_dataset"]) + batch = next(iterator) + (concat, true, false), y = batch + # Whole global batches: leading batch dim is the global batch size (no .batch() op) + assert concat.shape == (4, *_SAMPLE_SHAPE) + assert true.shape == (4, *_SAMPLE_SHAPE) + assert false.shape == (4, *_SAMPLE_SHAPE) + assert y.shape == (4, *_SAMPLE_SHAPE) + results["_train_holder"].clear() + + def test_epoch_covers_train_indices_exactly_once(self): + results = _build(_make_data(), shuffle=True) + n_batches_per_epoch = results["train_steps"] * results["accumulation_steps"] + iterator = iter(results["train_dataset"]) + seen = [] + for _ in range(n_batches_per_epoch): + batch_ids = _batch_row_ids(next(iterator)) + # Within-batch sorted gathers (read locality); membership is epoch-shuffled + assert np.all(np.diff(batch_ids) >= 0) + seen.extend(batch_ids.tolist()) + assert sorted(seen) == sorted(results["train_indices"].tolist()) + results["_train_holder"].clear() + + def test_shuffle_false_yields_train_indices_order(self): + """The alignment contract train_random_forest depends on: with shuffle=False the i-th + yielded cadence is row train_indices[i], across exactly train_steps x accumulation_steps + batches per epoch pass.""" + results = _build(_make_data(), shuffle=False) + n_batches_per_epoch = results["train_steps"] * results["accumulation_steps"] + iterator = iter(results["train_dataset"]) + seen = [] + for _ in range(n_batches_per_epoch): + seen.extend(_batch_row_ids(next(iterator)).tolist()) + assert seen == results["train_indices"].tolist() + # A second epoch pass repeats the same order (.repeat() on an ordered generator) + second = [] + for _ in range(n_batches_per_epoch): + second.extend(_batch_row_ids(next(iterator)).tolist()) + assert second == seen + results["_train_holder"].clear() + + def test_val_yields_val_indices_order(self): + results = _build(_make_data(), shuffle=True) + iterator = iter(results["val_dataset"]) + seen = [] + for _ in range(results["val_steps"]): + seen.extend(_batch_row_ids(next(iterator)).tolist()) + assert seen == results["val_indices"].tolist() + results["_train_holder"].clear() + + def test_split_is_stratified_and_disjoint(self): + data = _make_data() + results = _build(data, shuffle=True) + train_indices = results["train_indices"] + val_indices = results["val_indices"] + assert set(train_indices).isdisjoint(set(val_indices)) + labels = data["labels"] + for signal_type in _SIGNAL_TYPES: + # 40 samples, 10 per type, 0.8 split -> 8 train / 2 val of each type + assert int(np.sum(labels[train_indices] == signal_type)) == 8 + assert int(np.sum(labels[val_indices] == signal_type)) == 2 + results["_train_holder"].clear() + + def test_memmap_inputs_supported(self, tmp_path): + """Round data arrives as np.load(mmap_mode='r') memmaps — gathers must produce plain + in-RAM batches from them.""" + data = _make_data() + mmap_data = {"labels": data["labels"]} + for key in ("concatenated", "true", "false"): + path = tmp_path / f"{key}.npy" + np.save(path, data[key]) + mmap_data[key] = np.load(path, mmap_mode="r") + results = _build(mmap_data, shuffle=False) + iterator = iter(results["train_dataset"]) + batch_ids = _batch_row_ids(next(iterator)) + assert batch_ids.tolist() == results["train_indices"][:4].tolist() + results["_train_holder"].clear() + + def test_holder_clear_drops_references(self): + holder = TrainDataHolder(np.zeros(2), np.zeros(2), np.zeros(2)) + holder.clear() + assert holder.concat is None and holder.true is None and holder.false is None + holder.clear() # idempotent + assert holder._cleared + + +class TestPrepareDistributedVizDataset: + def _build_viz(self, concat, prib=4): + return prepare_distributed_viz_dataset( + concat_data=concat, + per_replica_inf_batch_size=prib, + num_replicas=1, + strategy=tf.distribute.get_strategy(), + ) + + def test_order_preserved_with_padding(self): + n = 10 + base = np.arange(n, dtype=np.float32)[:, None, None, None] + concat = base * np.ones((n, *_SAMPLE_SHAPE), dtype=np.float32) + results = self._build_viz(concat) + assert results["n_samples"] == 10 + assert results["n_padded"] == 12 + assert results["viz_steps"] == 3 + + iterator = iter(results["viz_dataset"]) + seen = [] + for _ in range(results["viz_steps"]): + batch = next(iterator) + assert batch.shape == (4, *_SAMPLE_SHAPE) # whole global batches + seen.extend(batch.numpy()[:, 0, 0, 0].astype(np.int64).tolist()) + # WARN-comment contract: the first n_samples yields are the original cadence order — + # plot_latent_space_gif and _capture_latent_snapshot truncate the padded tail + assert seen[:n] == list(range(n)) + # Padding rows duplicate real cadences + assert all(0 <= idx < n for idx in seen[n:]) + results["_viz_holder"].clear() + + def test_no_padding_when_divisible(self): + n = 8 + concat = np.ones((n, *_SAMPLE_SHAPE), dtype=np.float32) + results = self._build_viz(concat) + assert results["n_padded"] == n + assert results["viz_steps"] == 2 + results["_viz_holder"].clear() diff --git a/tests/unit/test_train_utils.py b/tests/unit/test_train_utils.py new file mode 100644 index 00000000..1150e982 --- /dev/null +++ b/tests/unit/test_train_utils.py @@ -0,0 +1,320 @@ +"""Unit tests for aetherscan.train pure-logic helpers: checkpoint tag resolution, curriculum +schedules, directory archiving, encoder-trained heuristics, and SHAP output normalization.""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +from aetherscan.config import get_config +from aetherscan.train import ( + TrainingPipeline, + _select_positive_class_shap, + archive_directory, + check_encoder_trained, + compute_expected_std, + get_latest_tag, +) + + +def _touch_pair(checkpoints_dir, tag): + """Create a matching encoder/decoder checkpoint pair for `tag`.""" + os.makedirs(checkpoints_dir, exist_ok=True) + for prefix in ("vae_encoder", "vae_decoder"): + with open(os.path.join(checkpoints_dir, f"{prefix}_{tag}.keras"), "w") as f: + f.write("stub") + + +class TestGetLatestTag: + def test_priority_ladder(self, tmp_path): + d = str(tmp_path / "ckpt") + for tag in ("test_v3", "20240101_000000", "round_02", "final_v1"): + _touch_pair(d, tag) + assert get_latest_tag(d) == "final_v1" + + def test_round_beats_timestamp_and_test(self, tmp_path): + d = str(tmp_path / "ckpt") + for tag in ("test_v9", "20991231_235959", "round_02", "round_10"): + _touch_pair(d, tag) + # Numeric compare, not lexicographic: round_10 > round_02. + assert get_latest_tag(d) == "round_10" + + def test_timestamp_beats_test(self, tmp_path): + d = str(tmp_path / "ckpt") + for tag in ("test_v9", "20240101_000000", "20250101_000000"): + _touch_pair(d, tag) + assert get_latest_tag(d) == "20250101_000000" + + def test_test_tags_ranked_by_version(self, tmp_path): + d = str(tmp_path / "ckpt") + for tag in ("test_v2", "test_v17", "test_v9"): + _touch_pair(d, tag) + assert get_latest_tag(d) == "test_v17" + + def test_final_ranked_by_version(self, tmp_path): + d = str(tmp_path / "ckpt") + for tag in ("final_v1", "final_v12", "final_v3"): + _touch_pair(d, tag) + assert get_latest_tag(d) == "final_v12" + + def test_encoder_without_decoder_ignored(self, tmp_path): + d = str(tmp_path / "ckpt") + _touch_pair(d, "round_01") + # Higher-priority final_v2 lacks its decoder — must not win. + with open(os.path.join(d, "vae_encoder_final_v2.keras"), "w") as f: + f.write("stub") + assert get_latest_tag(d) == "round_01" + + def test_missing_directory_raises(self, tmp_path): + with pytest.raises(FileNotFoundError, match="doesn't exist"): + get_latest_tag(str(tmp_path / "nope")) + + def test_empty_directory_raises(self, tmp_path): + d = tmp_path / "empty" + d.mkdir() + with pytest.raises(FileNotFoundError, match="No encoder files"): + get_latest_tag(str(d)) + + def test_no_complete_pair_raises(self, tmp_path): + d = tmp_path / "orphans" + d.mkdir() + with open(d / "vae_encoder_round_01.keras", "w") as f: + f.write("stub") + with pytest.raises(FileNotFoundError, match="No valid model pairs"): + get_latest_tag(str(d)) + + +class _PipelineStub: + """Just enough of TrainingPipeline to drive _calculate_curriculum_snr.""" + + def __init__(self, config): + self.config = config + + +def _curriculum(round_idx): + return TrainingPipeline._calculate_curriculum_snr(_PipelineStub(get_config()), round_idx) + + +class TestCalculateCurriculumSnr: + @pytest.fixture(autouse=True) + def _setup_config(self): + config = get_config() + config.training.num_training_rounds = 5 + config.training.snr_base = 10 + config.training.initial_snr_range = 40 + config.training.final_snr_range = 10 + + def test_linear_endpoints_and_monotonicity(self): + get_config().training.curriculum_schedule = "linear" + ranges = [_curriculum(i)[1] for i in range(5)] + assert ranges[0] == 40 + assert ranges[-1] == 10 + assert all(a >= b for a, b in zip(ranges, ranges[1:], strict=False)) + assert all(base == 10 for base, _ in (_curriculum(i) for i in range(5))) + + def test_linear_midpoint(self): + get_config().training.curriculum_schedule = "linear" + # progress = 2/4 = 0.5 -> 40 - 0.5 * 30 = 25 + assert _curriculum(2)[1] == 25 + + def test_exponential_endpoints_exact(self): + config = get_config() + config.training.curriculum_schedule = "exponential" + config.training.exponential_decay_rate = -3.0 + assert _curriculum(0)[1] == 40 + assert _curriculum(4)[1] == 10 + + def test_exponential_decays_faster_than_linear(self): + config = get_config() + config.training.curriculum_schedule = "exponential" + config.training.exponential_decay_rate = -3.0 + ranges = [_curriculum(i)[1] for i in range(5)] + assert all(a >= b for a, b in zip(ranges, ranges[1:], strict=False)) + # Exponential front-loads the difficulty ramp: below linear at the midpoint. + assert ranges[2] < 25 + + def test_exponential_rejects_nonnegative_decay_rate(self): + config = get_config() + config.training.curriculum_schedule = "exponential" + config.training.exponential_decay_rate = 0.5 + with pytest.raises(ValueError, match="must be < 0"): + _curriculum(1) + + def test_step_schedule(self): + config = get_config() + config.training.curriculum_schedule = "step" + config.training.step_easy_rounds = 2 + config.training.step_hard_rounds = 3 + assert [_curriculum(i)[1] for i in range(5)] == [40, 40, 10, 10, 10] + + def test_step_schedule_rejects_bad_sum(self): + config = get_config() + config.training.curriculum_schedule = "step" + config.training.step_easy_rounds = 2 + config.training.step_hard_rounds = 2 # 2 + 2 != 5 + with pytest.raises(ValueError, match="must equal total_rounds"): + _curriculum(0) + + def test_single_round_returns_initial_range(self): + config = get_config() + config.training.num_training_rounds = 1 + config.training.curriculum_schedule = "linear" + assert _curriculum(0) == (10, 40) + + def test_unknown_schedule_raises(self): + get_config().training.curriculum_schedule = "sigmoid" + with pytest.raises(ValueError, match="invalid"): + _curriculum(0) + + +class TestArchiveDirectory: + def test_empty_directory_is_noop(self, tmp_path): + base = tmp_path / "plots" + archive_directory(str(base)) + assert base.exists() + assert list(base.iterdir()) == [] + + def test_fresh_run_moves_files_to_archive(self, tmp_path): + base = tmp_path / "plots" + base.mkdir() + (base / "a.png").write_text("a") + (base / "b.png").write_text("b") + (base / "subdir").mkdir() # not in target_dirs -> untouched + + archive_directory(str(base), round_num=1) + + remaining = {p.name for p in base.iterdir()} + assert remaining == {"archive", "subdir"} + archived = list((base / "archive").iterdir()) + assert len(archived) == 1 # one timestamped snapshot + assert {p.name for p in archived[0].iterdir()} == {"a.png", "b.png"} + + def test_resume_copies_then_deletes_rounds_geq(self, tmp_path): + base = tmp_path / "checkpoints" + base.mkdir() + for tag in ("round_01", "round_02", "round_03"): + (base / f"vae_encoder_{tag}.keras").write_text("stub") + (base / "notes.txt").write_text("keep me") + + archive_directory(str(base), round_num=2) + + remaining = {p.name for p in base.iterdir()} + # Rounds >= 2 deleted; round_01 and non-round files kept. + assert remaining == {"archive", "vae_encoder_round_01.keras", "notes.txt"} + # Everything (including the later-deleted files) was backed up first. + snapshot = next((base / "archive").iterdir()) + assert {p.name for p in snapshot.iterdir()} == { + "vae_encoder_round_01.keras", + "vae_encoder_round_02.keras", + "vae_encoder_round_03.keras", + "notes.txt", + } + + def test_target_dirs_moved_and_recreated_empty(self, tmp_path): + base = tmp_path / "tb" + (base / "train").mkdir(parents=True) + (base / "train" / "events.1").write_text("x") + (base / "validation").mkdir() + + archive_directory(str(base), target_dirs=["train"], round_num=1) + + assert (base / "train").exists() + assert list((base / "train").iterdir()) == [] # replaced with an empty dir + assert (base / "validation").exists() # not a target -> untouched + snapshot = next((base / "archive").iterdir()) + assert (snapshot / "train" / "events.1").exists() + + +class TestEncoderTrainedHeuristics: + @staticmethod + def _dense_model(initializer, units=256, input_dim=256): + import tensorflow as tf # noqa: PLC0415 + + model = tf.keras.Sequential( + [ + tf.keras.layers.Input(shape=(input_dim,)), + tf.keras.layers.Dense(units, kernel_initializer=initializer), + ] + ) + return model + + def test_expected_std_he_normal(self): + from tensorflow.keras.initializers import HeNormal # noqa: PLC0415 + + model = self._dense_model(HeNormal(seed=0), units=4, input_dim=8) + expected = compute_expected_std(model.layers[-1]) + assert expected == pytest.approx(np.sqrt(2.0 / 8)) + + def test_expected_std_glorot_normal(self): + from tensorflow.keras.initializers import GlorotNormal # noqa: PLC0415 + + model = self._dense_model(GlorotNormal(seed=0), units=4, input_dim=8) + expected = compute_expected_std(model.layers[-1]) + assert expected == pytest.approx(np.sqrt(2.0 / (8 + 4))) + + def test_expected_std_unknown_initializer_returns_none(self): + model = self._dense_model("glorot_uniform", units=4, input_dim=8) + assert compute_expected_std(model.layers[-1]) is None + + def test_fresh_encoder_reports_untrained(self): + from tensorflow.keras.initializers import HeNormal # noqa: PLC0415 + + # 256x256 kernel: sampling noise on the std is ~0.3%, far under the 20% threshold. + model = self._dense_model(HeNormal(seed=0)) + assert check_encoder_trained(model) is False + + def test_scaled_weights_report_trained(self): + from tensorflow.keras.initializers import HeNormal # noqa: PLC0415 + + model = self._dense_model(HeNormal(seed=0)) + layer = model.layers[-1] + kernel, bias = layer.get_weights() + layer.set_weights([kernel * 3.0, bias]) # 200% deviation from expected std + assert check_encoder_trained(model) is True + + +class TestSelectPositiveClassShap: + N, F = 5, 8 + + def test_list_of_class_arrays_selects_positive(self): + neg = np.zeros((self.N, self.F)) + pos = np.ones((self.N, self.F)) + result = _select_positive_class_shap([neg, pos]) + np.testing.assert_array_equal(result, pos) + + def test_trailing_class_axis_values(self): + values = np.stack( + [np.zeros((self.N, self.F)), np.ones((self.N, self.F))], axis=-1 + ) # (N, F, 2) + result = _select_positive_class_shap(values) + assert result.shape == (self.N, self.F) + assert np.all(result == 1.0) + + def test_trailing_class_axis_interactions(self): + values = np.stack( + [np.zeros((self.N, self.F, self.F)), np.ones((self.N, self.F, self.F))], axis=-1 + ) # (N, F, F, 2) + result = _select_positive_class_shap(values) + assert result.shape == (self.N, self.F, self.F) + assert np.all(result == 1.0) + + def test_single_output_passthrough(self): + values = np.arange(self.N * self.F, dtype=float).reshape(self.N, self.F) + np.testing.assert_array_equal(_select_positive_class_shap(values), values) + + def test_log_loss_list_selects_first(self): + first = np.ones((self.N, self.F)) + result = _select_positive_class_shap([first, np.zeros((self.N, self.F))], log_loss=True) + np.testing.assert_array_equal(result, first) + + def test_log_loss_trailing_class_axis(self): + values = np.stack([np.zeros((self.N, self.F)), np.ones((self.N, self.F))], axis=-1) + result = _select_positive_class_shap(values, log_loss=True) + assert result.shape == (self.N, self.F) + assert np.all(result == 1.0) + + def test_log_loss_passthrough(self): + values = np.arange(self.N * self.F, dtype=float).reshape(self.N, self.F) + np.testing.assert_array_equal(_select_positive_class_shap(values, log_loss=True), values)