perf(ingest): streaming beamtime ingest + maintainability refactor - #37
Draft
HarlanHeilman wants to merge 22 commits into
Draft
perf(ingest): streaming beamtime ingest + maintainability refactor#37HarlanHeilman wants to merge 22 commits into
HarlanHeilman wants to merge 22 commits into
Conversation
Records the six performance fixes (T1-T6) and four maintainability refactors (T7-T10) driving the perf/ingest-streaming-and-maintainability branch. Serves as the authoritative reference for per-task reviews.
Replaces the per-pixel read_exact(2) loop with a single read_exact into a preallocated buffer, then decodes via chunks_exact(2). Cuts ~170k syscalls per file on network-mounted FITS to 1, which dominates wall-clock during the zarr phase on NAS-backed beamtimes. Part of docs/plans/ingest-streaming-and-maintainability.md (T1).
Removes the second zarr array that write_frame_raw wrote at
/{scan}/{frame}/processed with identical bytes to /raw. Nothing in the
codebase reads /processed; downstream processing produces its own
outputs. Halves the filesystem metadata and chunk-write work per frame
during the zarr phase.
Updates module/function docstrings and the schema.rs comment that
described the old two-array layout.
Part of docs/plans/ingest-streaming-and-maintainability.md (T2).
… phase Replaces the N+1 SELECT pattern for samples/scans/files/tags with `.returning(id).get_result(conn)` inside ingest_beamtime_inner. Cuts SQLite round-trips in the catalog transaction roughly in half for the rows inserted this ingest run. Diesel's returning_clauses_for_sqlite_3_35 feature is already enabled and the bundled libsqlite3-sys is >= 3.44. Part of docs/plans/ingest-streaming-and-maintainability.md (T3).
Splits the monolithic catalog-phase SQLite transaction into batches of MIN_BATCH_ROWS=200 to MAX_BATCH_ROWS=1000 rows. Scans are never split across batches. Small scans coalesce until the batch reaches MIN or the next scan would push it over MAX; very large scans (>MAX) get their own batch. Samples insert stays in a separate one-shot transaction up front. Shorter lock hold times let concurrent SQLite readers query during long ingests and make committed rows visible incrementally. Catalog progress events remain unchanged. Part of docs/plans/ingest-streaming-and-maintainability.md (T4).
- Remove unused MIN_BATCH_ROWS constant. The greedy planner only enforces MAX; MIN was aspirational, never enforced, and the associated `plan_catalog_batches_fills_to_min_before_sealing` test did not actually exercise a MIN boundary. - Rename `ingest_produces_same_row_counts_as_single_transaction` to `ingest_produces_expected_row_counts_for_minimal_fixture` since the fixture is too small to cross a batch boundary. - Add module-level doc comment explaining ingest phases and partial-failure semantics: prior committed batches remain on error; re-invoking ingest deletes the beamtime row and cascades stale data. - Check cancel inside the samples pre-transaction so long samples-heavy ingests respond to cancellation. Addresses Important findings from T4 code-quality review. Part of docs/plans/ingest-streaming-and-maintainability.md (T4).
Replaces the memory-buffering zarr phase (~9 GB peak resident on large beamtimes) with a crossbeam-channel pipeline. Reader pool pushes decoded images onto a bounded channel of capacity n_workers*2; writer drains the channel and writes each frame via write_frame_raw, emitting FileComplete as each write lands. Peak image residency is now bounded by the channel capacity (typically O(16 MB)). FileComplete events fire incrementally rather than bursting after the final read, giving live progress bars during zarr writes. Order of FileComplete emission switches from row-index order to write completion order; global_done still reaches rows.len() exactly. Part of docs/plans/ingest-streaming-and-maintainability.md (T5).
Replaces `scan_groups.par_iter().map(read_multiple_fits_headers_only_rows)` with a flat `paths_only.par_iter().map(read_fits_headers_only_row)` over the Rayon pool. Evens out worker utilization when scan sizes are skewed (e.g., one 3000-file scan plus many 50-file scans), where the old per-scan parallelism left N-1 workers idle once small scans completed. No change to the sort order of the resulting `rows` Vec - the same `(scan_number, frame_number, file_path)` sort follows the flat collect. Part of docs/plans/ingest-streaming-and-maintainability.md (T6).
…th IngestContext Split the 420-line ingest_beamtime_inner into a thin orchestrator and three phase functions with a shared IngestContext struct: - IngestContext: carries beamtime_id, progress, cancel, pool, scan_total_map - run_headers_phase (28 lines): parallel FITS header parsing + sort - run_catalog_phase (231 lines): batched SQLite transactions - run_zarr_phase (208 lines): streaming crossbeam read/write pipeline - ingest_beamtime_inner (107 lines): thin orchestrator Cancel handling changed from Arc<AtomicBool> threaded through to Option<&AtomicBool> borrowed into IngestContext. Progress event ordering preserved: Layout in orchestrator, Phase(Headers)/Phase(Catalog)/CatalogRow/ Phase(Zarr)/FileComplete in respective phase functions. Strict refactor: no behavior change, no wire-format change, all existing tests pass unchanged.
…helper
Follow-up to T9 addressing code-review items I1, I2, M1, M5.
- I1: Add `IngestContext::is_cancelled()` helper using `is_some_and` and
replace the three `cancel.map(...).unwrap_or(false)` occurrences (two in
run_catalog_phase, one in the zarr reader closure via cancel_ref).
- I2: Extract three helpers from run_catalog_phase so the driver shrinks
from 231 to 38 lines:
- build_scan_first_sample(rows): pure map from scan -> first sample name
- insert_samples(conn, beamtime_id, rows, cancel): returns sample_cache
- insert_catalog_batch(conn, ctx, rows, start, end, caches, ...): runs
ONE batch transaction and emits CatalogRow per row
- M1: Add short doc comments on IngestContext and all phase/helper fns.
- M5: Remove dead `let _ = layout_label(layout);` and the unused
`layout_label` helper (no other callers) and its `BeamtimeLayout` import.
No behavior change, no wire-format change, no public API change.
Adds reproducible synthetic-beamtime fixtures so the streaming ingest can be exercised end-to-end without a real ALS dataset. Rust: - tests/common/mod.rs: build_synthetic_beamtime() generates N valid, ingestible BITPIX=16 FITS files in the root/CCD/ layout. Headers and pixel data are padded to 2880-byte blocks; required cards (NAXIS, NAXIS1/2, BZERO=32768, DATE, Beamline Energy, Sample/CCD Theta, EPU Polarization, Higher Order Suppressor) are emitted with the HIERARCH long-keyword convention. - tests/synthetic_harness.rs: round-trips fixtures through pyref::loader::read_fits_headers_only_row and verifies pixel layout via io::image_mmap::load_image_pixels (now pub) using a positional sentinel. - tests/ingest_streaming.rs: ingests 10 scans x 10 frames at 16x16 px into isolated PYREF_CATALOG_DB / PYREF_CACHE_ROOT tempdirs and asserts 100 files / 100 frames / 10 scans plus catalog_row -> file_complete ordering through the IngestProgressSink callback. - src/io/image_mmap.rs: expose load_image_pixels for harness use. Python: - scripts/_ingest_profile.py: shared helpers (IngestProfile, isolated_catalog_env, run_ingest_with_profile, render_markdown_table) factored out so both the real-beamtime profiler and the synthetic bench print the same markdown timing report. - scripts/bench_ingest.py: CI-friendly benchmark that materialises synthetic FITS via astropy.io.fits, ingests them under an isolated catalog/cache, and prints the timing table plus a files_per_second summary. - scripts/profile_beamtime_ingest.py: refactored to reuse the shared helpers; behaviour and output are unchanged.
- emit SIMPLE=T (standard FITS logical) instead of SIMPLE=1 in the Rust synthetic writer - mirror the Rust writer's per-scan/per-frame header variation (Beamline Energy, Sample Theta) in the Python astropy writer so the two generators stay in sync - add return type Iterator[Path | None] to isolated_catalog_env
…g catalog path - Add shape-bucketed per-scan 3D uint16 Zarr with shuffle+Zstd; frames zarr index columns - SQLite migration; ingest, query, and image load paths (Zarr-first, FITS fallback) - Add migrate-zarr-3d binary with dry-run compression/ETA sampling; optional zstd dep - Default catalog.db under ~/.config/pyref (XDG_CONFIG_HOME, PYREF_HOME overrides) - Python IO/catalog_path updates; notebooks and tests
…, bench CLI - Pipelined ingest, scan scheduler, zarr and catalog refactors - Rename frame EAV to header_values; metadata module; remove duplicate io schema - Migrations for zarr bucket and header_values/frames prune - Move ingest profiling into python/pyref; add pyref bench; remove old scripts - Update AGENTS, readers, and beamtime/catalog Python surface
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Executes the full plan in
docs/plans/ingest-streaming-and-maintainability.md, split into two phases.Phase A - Performance (T1-T6)
read_image_i32via unifiedread_bitpix16_be_byteshelper47666a6/processedzarr array inwrite_frame_raw(halves per-frame I/O)726cf24RETURNINGinstead ofSELECT-after-INSERTfor catalog rows82162754f5f0fa,5470a34778d836b582a26Phase B - Maintainability (T7-T10)
src/io/raw_pixels.rswith documented// SAFETY:mmap block;flatten_fits_errorhelper insrc/catalog/mod.rswith Io/Validation-preserving unit tests196d443,b44b104IngestPhaseenum replaces stringly-typed phase labels; Python wire format preserved viaas_str()b64e1ba,d3378bfingest_beamtime_inner(~690 -> ~107 lines orchestrator) intorun_headers_phase/run_catalog_phase/run_zarr_phasebacked by a sharedIngestContext; extractedbuild_scan_first_sample/insert_samples/insert_catalog_batch; addedIngestContext::is_cancelled()to dedupe cancellation checks93501c6,e2c1926tests/common/mod.rs), two self-tests, end-to-end streaming integration test with env-isolated catalog, sharedscripts/_ingest_profile.py, newscripts/bench_ingest.pyCLIdb88dbd,fc92dda,83d0bfe,85b066dRegression gates (all green)
cargo test --features catalog,parallel_ingest-> 86 unit + 4 ingest/harness integration tests passuv run pytest tests/test_catalog.py-> 17/17 passuv run ruff check scripts/_ingest_profile.py scripts/bench_ingest.py scripts/profile_beamtime_ingest.py tests/test_catalog.py-> cleanuv run ty checkon branch-touched Python -> no new errors vs main (pre-existing polars/duckdb union-type errors intests/test_catalog.pyconfirmed present on main)cargo clippy --features catalog,parallel_ingest-> 17 warnings vs 18 on main (net -1; all remaining are pre-existing)scripts/bench_ingest.py --scans 8 --frames-per-scan 25 --width 128 --height 128-> 200 files at ~1080 files/s with streaming zarr dominating wall time as expectedBackward compatibility
ingest_beamtimefamily"headers","catalog","zarr"viaIngestPhase::as_str())src/io/image_mmap.rs::load_image_pixelspromoted fromfntopub fn(required for the integration-test binaries;pub(crate)does not reach separate test crates)Test plan
uv run python scripts/bench_ingest.py --scans 50 --frames-per-scan 50 --width 1024 --height 1024locally and confirm files_per_second is higher thanmainbaselineuv run python scripts/profile_beamtime_ingest.py --beamtime <path>and verify phase shares match expectation (zarr dominant, catalog small)/processedarray consumers broken)