Training data pipeline: disk-backed memmap rounds, background producer, batched tf.data input#117
Training data pipeline: disk-backed memmap rounds, background producer, batched tf.data input#117zachtheyek wants to merge 9 commits into
Conversation
New src/aetherscan/round_data.py:
- RoundDataPaths: on-disk layout of one round's dataset
({main,true,false,labels}.npy + round_XX.done manifest)
- Atomic .done manifest protocol (.tmp -> os.replace) with cheap
element-count + sampled-value checksums; validate_done_manifest()
gates reuse of on-disk rounds
- prepare_round_data_dir(): startup archive/reuse/delete semantics
mirroring checkpoint archiving (delete rounds >= start_round, keep
earlier rounds only with a validating manifest)
- RoundDataProducer: dedicated multiprocessing.Process that generates
round k+1 while round k trains (generate/shutdown in;
stats/progress/done/error out); a main-process drainer thread writes
streamed injection stats to the DB (the DB queue is thread-only);
producer installs the canonical no-log SIGTERM handler and uses
init_worker_logging()
ResourceManager gains a tracked ManagedProcess (terminate -> join ->
kill escalation mirroring ManagedPool.close, including killing the
producer's orphaned pool children before SIGKILL); processes are closed
first in cleanup_all().
Part of the training-data-pipeline overhaul (#114).
…mmaps Replaces the 1-task-per-sample pool dispatch (per-sample IPC pickling + main-process copies were the GIL bottleneck that made round 2+ data-gen far slower than round 1) with batched tasks: each task generates up to training.data_gen_task_size cadences and writes them directly into its disjoint slice of the round's on-disk memmap (open_memmap r+, like preprocessing's _extract_stamps_worker), returning only the small per-sample stats dicts. - The 8 class-segments per chunk are submitted as ONE unified task list through a single pool.map barrier (instead of 8 sequential barriers) - Requires chunk_size % 4 == 0 (already CLI-validated), which drops the ceil-division oversample/subsample dance and its stats-filtering code - generate_round_to_memmap() replaces generate_triplet_batch(): labels written per chunk, atomic .done manifest written only after all chunks finish; per-task seeds keep results independent of worker scheduling - write_segment_stats() (module-level) replaces DataGenerator._write_batch_stats so the producer's main-process drainer can reuse it - DataGenerator pool creation is now lazy (_ensure_pool): with overlap enabled the producer process owns the only worker pool, and the in-process pool exists only for the sequential path / RF dataset - Remove the commented-out shared-memory-output redesign blocks: this change is that redesign Part of the training-data-pipeline overhaul (#114).
prepare_distributed_train_dataset() now yields whole GLOBAL batches gathered from np.load(mmap_mode="r") arrays (sorted fancy indices per batch, epoch-level shuffle preserved), drops the .batch() op, and keeps .repeat().prefetch(AUTOTUNE) — removing the per-sample Python/tf.data boundary crossings behind the 0-14 % GPU utilization. The stratified split and the shuffle=False alignment contract used by train_random_forest are preserved (train/val indices are sorted once so the yield order matches the returned index arrays); val and viz datasets get the same batched treatment (viz order guarantee kept). TrainDataHolder semantics are unchanged, now holding memmap refs; the stale lock-contention comment is refreshed. train_beta_vae() schedules generation through the RoundDataProducer when overlap_data_generation is on: request round k+1 as soon as round k's data is ready, train round k, delete round k's directory after it completes (keep_round_data retains it). Sequential in-process generation remains for --no-overlap-data-generation and pool-less (n_processes=1) runs. _setup_directories() prepares the tag's round-data dir with checkpoint-style resume semantics. train_random_forest() generates its dataset through the same disk-backed builder (per-round dir sibling "rf") — the train_steps x accumulation_steps encode-count contract still holds and remains guarded by _distributed_encode's drift check. RAM no longer scales with round size: batches gather through the OS page cache, which evicts under pressure instead of OOM-killing the run. Part of the training-data-pipeline overhaul (#114).
New TrainingConfig fields (+ Config.to_dict entries): - round_data_dir (None -> <output_path>/round_data at runtime) - overlap_data_generation (True), keep_round_data (False) — both BooleanOptionalAction flags - data_gen_task_size (256), validated >= 1 collect_validation_errors() gains a train-mode disk-budget check: --round-data-dir must have >= 2.2x one round's estimated size free with overlap enabled (two rounds coexist on disk), 1.1x without — a hard ValidationError carrying the computed numbers. The estimate is stdlib- only so utils/print_cli_help.py keeps importing cli.py without the scientific stack. README CLI Reference blocks regenerated with `PYTHONPATH=src python utils/print_cli_help.py all`. Part of the training-data-pipeline overhaul (#114).
- tests/unit/test_round_data.py: RoundDataPaths layout; .done manifest roundtrip + corruption/shape/count mismatches; prepare_round_data_dir resume semantics; chunk-segment and task partitioning covering every row exactly once; sequential end-to-end generate_round_to_memmap (labels layout, stats segments, manifest, progress); producer protocol driven in-thread with a stub generator fn (stats/progress/ done, error reporting, valid-round short-circuit); main-side drainer + await_round against a fake process (DB writes land in main) - tests/unit/test_train_datasets.py (slow): whole-global-batch yields, exact per-epoch index coverage, stratified split, the shuffle=False train_random_forest alignment contract, memmap inputs, viz order guarantee with padding - tests/unit/test_manager.py: ManagedProcess register/close, idempotent close, SIGTERM-ignoring kill escalation, dead-process close - tests/unit/test_cli_validation.py: new round-data flags routing + defaults, data_gen_task_size bound, disk-budget check (deterministic via an autouse shutil.disk_usage stub — the full-scale default needs more free disk than CI runners have) - tests/conftest.py: teardown also closes leaked ManagedProcesses Part of the training-data-pipeline overhaul (#114).
- RoundDataProducer's drainer thread now survives a failure while handling a single message (e.g. a DB hiccup on a stats write) — otherwise every later await_round() would raise a misleading "producer exited" error while the producer kept generating fine - Update the --num-samples-rf divisibility message that referenced the removed generate_triplet_batch() Part of the training-data-pipeline overhaul (#114).
The first cluster smoke hung: the fork-started producer froze on a futex before reaching its own code — a lock inherited in a locked state from the training parent's deep TF/NCCL/gRPC thread stack (single thread, 0% CPU, no pool children, nothing logged from the child). Fix: start the producer with the multiprocessing "spawn" context — a fresh interpreter inherits no locks (its own pool workers are then forked from the clean single-threaded producer, which is safe). Plumbing that spawn requires: - init_worker_logging() accepts an explicit log_queue (a spawned child has no inherited Logger singleton); _init_worker forwards it to the producer's pool workers via initargs - Producer queues come from the spawn context so they cross the spawn pickling boundary - CUDA_VISIBLE_DEVICES is blanked around Process.start() so the spawned child's TF import chain can never initialize CUDA (generation is pure CPU) Part of the training-data-pipeline overhaul (#114).
The spawn-started producer cannot receive the Logger singleton's
fork-context multiprocessing queue ("A SemLock created in a fork
context is being shared with a process in a spawn context" at
Process.start()). The producer tree now logs into its own spawn-context
queue, and a main-side logging.handlers.QueueListener relays each
record into the main process's root handlers (the normal
file/console/Slack pipeline). Relay lifecycle is tied to
start()/shutdown().
Adds a slow-marked end-to-end test that runs a REAL spawn-started
producer against real shared memory with a real tiny generation — the
two cluster-smoke failures this PR hit (a fork-inherited deadlocked
lock, and the fork/spawn SemLock mix above) were only reachable through
a real child process, which the in-thread protocol tests can't cover.
Part of the training-data-pipeline overhaul (#114).
|
Claude finished @zachtheyek's task in 2m 39s —— View job Code Review: Training Data Pipeline
SummaryThis is an impressive, well-engineered overhaul that simultaneously solves four coupled problems (OOM, GIL contention, serial phases, GPU starvation) with a cohesive design. The architecture decisions are sound — spawn-started producer, atomic manifests, batched tf.data yields — and the cluster smoke results validate the approach convincingly. Strengths
Issues & SuggestionsI've posted 6 inline comments on specific code locations. The most actionable one:
The remaining inline comments are informational/defensive notes rather than blocking issues:
Architecture Observations (non-blocking)
VerdictThe PR is well-structured, well-tested, and addresses a real cluster-scale problem with appropriate engineering rigor. The one suggested code change (flush in finally) is minor. Ready to merge after the bla0 full-scale validation confirms the OOM fix at production geometry. |
|
Contribution compliance note: This PR's |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| deterministically-sampled flat-index values. Positions are derived from the element count, | ||
| so validation re-samples the same spots without storing an RNG state. | ||
| """ | ||
| flat = arr.reshape(-1) |
There was a problem hiding this comment.
Nit: arr.reshape(-1) on a memmap will create a full copy in RAM if the array isn't C-contiguous. Since these are always C-contiguous float32 memmaps created by open_memmap, this is fine in practice, but it'd be defensive to use arr.ravel() (which avoids a copy when contiguous) or .flat[positions] directly.
Not a bug at current usage, just a note for future-proofing if the layout ever changes.
There was a problem hiding this comment.
Agreed — added a NOTE at the checksum flatten (2bfa118) documenting the C-contiguity assumption and pointing at .ravel()/.flat if the layout ever changes. Kept reshape(-1) since open_memmap/np.save arrays are always C-contiguous today.
| 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" | ||
| ) | ||
| self._condition.wait(timeout=1.0) |
There was a problem hiding this comment.
Minor: the while True / wait(timeout=1.0) polling loop means await_round can take up to 1 second longer than strictly necessary after the result arrives (if notify_all() races with the timeout expiry). In practice this is negligible compared to generation wall times, but calling this out in case latency ever matters here. The Condition pattern is correct.
There was a problem hiding this comment.
Right — one clarification: notify_all() from the drainer wakes the wait immediately, so the 1 s timeout isn't added latency on the happy path; it's only a liveness re-check so a drainer that died without notifying can't strand the wait. Added a NOTE saying exactly that (2bfa118).
|
|
||
| create_fn = _CREATE_FUNCTIONS[create_fn_name] |
There was a problem hiding this comment.
np.random.seed(seed % (2**32)) seeds the legacy global RNG. This is needed because new_cadence and the create_* functions use both random.random() and np.random.choice([-1, 1]) (via the legacy API). Worth noting that this creates implicit coupling between per-task determinism and legacy RNG state — if any library call mutates the global RNG between cadences in the same task, results become scheduling-dependent.
Not actionable right now (the create_* functions would need to accept an explicit rng parameter), but flagging for awareness — especially since the PR description calls out that results don't depend on worker scheduling.
There was a problem hiding this comment.
Good flag. Added a NOTE (2bfa118) spelling out the legacy-global-RNG coupling and that determinism holds only while nothing else mutates global RNG state within a task; threading an explicit np.random.Generator through the create_* functions is the eventual fix if it ever matters.
| @@ -733,161 +707,283 @@ def batch_create_cadence( | |||
| inject=inject, | |||
| dynamic_range=dynamic_range, | |||
| ) | |||
| cadence[i, :, :, :] = result | |||
| out[start_idx + i] = cadence | |||
| all_sample_info.append(sample_info) | |||
| out.flush() | |||
| finally: | |||
| del out | |||
There was a problem hiding this comment.
The memmap is opened r+ without a lock, relying on disjoint row ranges for safety. This is correct and well-documented, but I want to flag one edge case: if create_fn raises partway through the for i in range(count) loop, out.flush() is never called for the rows that were written before the exception (the finally only dels the reference).
On Linux this is fine (the kernel will flush dirty pages eventually and the half-written round has no manifest so it's treated as garbage), but explicitly calling out.flush() in the finally block would make the intent clearer and help on exotic filesystems that might not auto-flush on munmap.
| 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, | |
| freq_resolution=freq_resolution, | |
| time_resolution=time_resolution, | |
| inject=inject, | |
| dynamic_range=dynamic_range, | |
| ) | |
| out[start_idx + i] = cadence | |
| all_sample_info.append(sample_info) | |
| finally: | |
| out.flush() | |
| del out |
There was a problem hiding this comment.
Applied in 2bfa118 — out.flush() now runs in the finally (plus a comment on why: deterministic write-back for rows written before a mid-task exception; the missing .done manifest already guards correctness either way).
| 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): |
There was a problem hiding this comment.
Good design — keeping this stdlib-only so utils/print_cli_help.py doesn't need numpy.
One note: the U20 label size calculation (n_samples * 20 * 4) is correct for numpy's UCS-4 encoding on disk, but could drift if np.save changes encoding or if labels grow beyond U20. A brief inline comment noting the assumption (e.g. # numpy U20 on-disk: 20 UCS-4 codepoints × 4 bytes each) would help future maintainers.
There was a problem hiding this comment.
The assumption is already annotated inline on that line: labels_bytes = n_samples * 20 * 4 # numpy "U20" = 20 UCS-4 code points per label — so skipping an additional comment to avoid duplication. If labels ever outgrow U20, the dtype literal and this estimate live one grep apart.
| # its own GPU init, so temporarily mutating it here is safe. | ||
| 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 |
There was a problem hiding this comment.
The CUDA_VISIBLE_DEVICES blanking is clever and well-motivated, but there's a small TOCTOU window: if another thread in the parent reads os.environ["CUDA_VISIBLE_DEVICES"] between lines 505 and 509-511, it would see the blanked value.
In practice this is safe because:
- TF reads this env var only at GPU initialization (already done by this point)
- No other thread should be spawning GPU-visible subprocesses concurrently
But a brief comment noting this is intentionally safe only because TF's GPU init is already complete would help a future reader who wonders about the race.
There was a problem hiding this comment.
Added a NOTE (2bfa118) making the safety argument explicit: the brief env-mutation window is acceptable only because the parent's TF read CUDA_VISIBLE_DEVICES once at GPU init (long before this point) and nothing else in the pipeline spawns GPU-visible subprocesses concurrently with training startup.
- _run_memmap_task: move out.flush() into the finally block so rows written before a mid-task exception still reach disk deterministically (review suggestion on #117) - Add clarifying NOTE comments from the review's informational points: reshape(-1) contiguity assumption in the manifest checksums, legacy global-RNG coupling of per-task seeding, the CUDA_VISIBLE_DEVICES env-mutation window, and the await_round wait-timeout being a liveness re-check rather than polling latency - Follow-up for the create_true_double retry-loop straggler filed as #118 Part of the training-data-pipeline overhaul (#114).
|
Review addressed in 2bfa118:
Each inline thread has an individual reply. Full local suite re-run after the change: 235 passed, 3 skipped, 2 deselected. |
|
Re: the contribution compliance note — the missing formal issue link is an artifact of this being a stacked PR, not a missing reference:
No action needed otherwise; the body already carries |
|
@claude please perform a second-pass code review of this PR. Re-review the current diff for any remaining correctness bugs, security issues, or significant quality problems, and post your findings. If the diff is clean, please say so explicitly. (Final verification round before merge.) |
Second-pass code review — ✅ CleanI've performed a thorough second-pass review of the full diff (round_data.py, data_generation.py, train.py, manager.py, config.py, cli.py, logger.py, and the test files). Checked for correctness bugs, race conditions, security issues, resource leaks, and significant quality problems. Verdict: No blocking issues found.The core patterns are sound:
Minor observations (non-blocking, informational):
LGTM for merge. 🚀 |
1 participant
Closes #114
Stacked on #116 (
feature/test-suite) — this PR's base is that branch, so the diff shows only this unit's changes.Summary
Overhauls the training data pipeline to fix four coupled problems that share one design (full details in #114): the ~294 GB/round in-RAM arrays that OOM-killed full-scale runs, GIL contention that made round 2+ data generation far slower than round 1, strictly serial generate-then-train phases, and per-sample
tf.datagenerators that starved the GPUs (0–14 % utilization).(a) Memmap-backed round datasets — new
src/aetherscan/round_data.py{main,true,false,labels}.npylive on disk under{round_data_dir}/{save_tag}/round_XX/(round_data_dirdefaults to<output_path>/round_data), created withnp.lib.format.open_memmapand read back vianp.load(mmap_mode="r").round_XX.donemanifest (sample counts, shapes, SNR params, element-count + sampled-value checksums, wall time, chunk count) is written atomically (.tmp→os.replace) only after all workers finish — a crash mid-generation leaves no manifest and the dir reads as garbage._setup_directories()prepares the tag's round-data dir with checkpoint-archiving-style resume semantics: delete round dirs ≥start_round, keep earlier dirs only if their manifest validates, delete.done-less dirs (deleted rather than archived — a round is ~295 GB).--keep-round-dataretains for debugging), so the steady-state disk footprint is ≤ 2 rounds with overlap.(b) Workers write straight into the memmap —
data_generation.pyrewrite--data-gen-task-size(default 256) cadences; each worker opens the round memmapr+and writes its disjoint row range in place, returning only the small per-sample stats dicts.pool.mapbarrier instead of 8 sequential ones.signal_injection_chunk_size % 4 == 0(already CLI-validated) makes quarters/halves exact, removing the ceil-division oversample/subsample dance and its stats-filtering code.(c) Background producer process — overlap + GIL isolation
RoundDataProducerprocess (owning its own worker pool attached to the background-plate SHM by name) generates round k+1 while round k trains. Queue protocol:generate/shutdownin;stats/progress/done/errorout.db.write_injection_stat(the DB queue is a threadqueue.Queue, not process-safe) — also moving the per-sample DB-write GIL load off the training thread's critical path.QueueListener(the Logger singleton's fork-context queue can't cross the spawn boundary), andCUDA_VISIBLE_DEVICESis blanked aroundProcess.start()so nothing in the producer tree can initialize CUDA.init_worker_logging()and the canonical no-log SIGTERM handler; the producer is registered with the ResourceManager via a new trackedManagedProcess(terminate → join → kill escalation mirroringManagedPool.close, including killing orphaned pool children before SIGKILL).--no-overlap-data-generationdrops to sequential in-process generation for debugging (also the automatic fallback whenn_processes=1leaves no SHM to attach to).(d) Batched
tf.datainput —prepare_distributed_train_datasetrewrite.batch()and keeping.repeat().prefetch(AUTOTUNE)— ~global-batch-size× fewer Python/tf.data boundary crossings.shuffle=Falsealignment contract used bytrain_random_forestare preserved (train/val indices are sorted once so yield order matches the returned index arrays); thetrain_steps × accumulation_stepsencode-count contract still holds and stays guarded by_distributed_encode's drift check.plot_latent_space_gifand_capture_latent_snapshotdepend on it).TrainDataHoldernow holds memmap references;clear()semantics unchanged; the stale lock-contention comment is refreshed. Page-cache framing documented in code comments: steady-state reads come from otherwise-free RAM, and the kernel evicts pages under pressure instead of OOM-killing the run.Config / CLI
New
TrainingConfigfields + train-only Pattern A flags (per docs/CONFIG_AND_CLI.md), serialized viaConfig.to_dict():--round-data-dir,--overlap-data-generation/--no-overlap-data-generation,--keep-round-data/--no-keep-round-data,--data-gen-task-size(validated ≥ 1).collect_validation_errorsgains a train-mode disk-budget check: ≥ 2.2× one round's estimated size free with overlap (two rounds coexist on disk), 1.1× without — a hard ValidationError with the computed numbers. README CLI Reference blocks regenerated withPYTHONPATH=src python utils/print_cli_help.py all.Verification
Local (macOS arm64, CPU-only venv: Python 3.12, tensorflow 2.17.1, pins from
requirements-container.txt):New tests:
tests/unit/test_round_data.py— paths/manifest/dir-prep/task-partitioning, end-to-end sequential generation, producer protocol driven with a stub generator fn, drainer +await_round, plus a real spawn-started producer end-to-end test (which would have caught both cluster failures below);tests/unit/test_train_datasets.py— batched yields, exact index coverage, stratification, theshuffle=Falsealignment contract, memmap inputs, viz order;ManagedProcesstests intest_manager.py; round-data flag/validation tests intest_cli_validation.py.Cluster smoke (blpc3, NGC 25.02 container, 5× RTX PRO 6000, tag
test_v21, branch @ 291334f):(Divisibility pre-validated for 5 replicas with
utils/find_optimal_configs.py --num-gpus 5; GPUs confirmed idle before launch.) Result: clean completion, 0 ERROR/Traceback lines, all artifacts saved,/dev/shmclean, round-data dirs auto-deleted.Acceptance criteria (PLAN §5.4, adapted to blpc3):
16:16:35 RoundDataProducer: generating round 3 dataruns through round 2's epochs (16:17:29 Epoch 1 … 16:18:35 Epoch 3) and completes at16:21:16.test_v22: batch 128/replica, effective 640, 3200 samples) sustained 5120 samples/epoch in 7.8 s once warm (~660 samples/s end-to-end) — the input pipeline is no longer the bottleneck at smoke scale, but per-step GPU compute for this VAE at these sample counts sits far below the 1 Hz sampling interval, so a comparable full-scale GPU-% number must come from the bla0 rehearsal (deferred below).random_forest_test_v21.joblib+ eval/SHAP artifacts); 56 tagged plots rendered incl. 24 latent-space UMAP GIF variants and the 10 RF diagnostics.The smoke also exercised: manifest-gated reuse short-circuiting, per-round dir deletion, RF's sequential in-process generation path, producer graceful shutdown, and (via two interrupted runs) the Ctrl-C cleanup path through the new
ManagedProcess.Observation (pre-existing, not a regression): the RF dataset generation (
600 samples, 1022.5 s) was dominated by one worker spinning ~10 min increate_true_double's non-intersection retry loop — the loop predates this PR (see itsNote, small but nonzero probability for "infinite" (long-running) loops). Batched tasks make such stragglers more visible since one unlucky sample stalls its whole task; a retry cap is worth a follow-up issue.Deferred validation
AI assistance disclosure
Implemented with Claude Code (implementation, unit tests, docs/README regeneration, and cluster smoke verification), per AI_POLICY.md.
🤖 Generated with Claude Code