Skip to content

perf(ingest): streaming beamtime ingest + maintainability refactor - #37

Draft
HarlanHeilman wants to merge 22 commits into
mainfrom
perf/ingest-streaming-and-maintainability
Draft

perf(ingest): streaming beamtime ingest + maintainability refactor#37
HarlanHeilman wants to merge 22 commits into
mainfrom
perf/ingest-streaming-and-maintainability

Conversation

@HarlanHeilman

Copy link
Copy Markdown
Member

Summary

Executes the full plan in docs/plans/ingest-streaming-and-maintainability.md, split into two phases.

Phase A - Performance (T1-T6)

# Change Commit
T1 Bulk-read FITS pixels in read_image_i32 via unified read_bitpix16_be_bytes helper 47666a6
T2 Drop duplicate /processed zarr array in write_frame_raw (halves per-frame I/O) 726cf24
T3 Diesel RETURNING instead of SELECT-after-INSERT for catalog rows 8216275
T4 Batched catalog transactions with small-scan coalescing so short scans commit in one txn 4f5f0fa, 5470a34
T5 Streaming zarr phase with bounded-parallel read -> write pipeline 778d836
T6 Parallelize headers phase per-file (rayon) instead of per-scan b582a26

Phase B - Maintainability (T7-T10)

# Change Commits
T7 Unified bulk BITPIX=16 pixel reader in src/io/raw_pixels.rs with documented // SAFETY: mmap block; flatten_fits_error helper in src/catalog/mod.rs with Io/Validation-preserving unit tests 196d443, b44b104
T8 IngestPhase enum replaces stringly-typed phase labels; Python wire format preserved via as_str() b64e1ba, d3378bf
T9 Split ingest_beamtime_inner (~690 -> ~107 lines orchestrator) into run_headers_phase / run_catalog_phase / run_zarr_phase backed by a shared IngestContext; extracted build_scan_first_sample / insert_samples / insert_catalog_batch; added IngestContext::is_cancelled() to dedupe cancellation checks 93501c6, e2c1926
T10 Synthetic-beamtime Rust harness (tests/common/mod.rs), two self-tests, end-to-end streaming integration test with env-isolated catalog, shared scripts/_ingest_profile.py, new scripts/bench_ingest.py CLI db88dbd, fc92dda, 83d0bfe, 85b066d

Regression gates (all green)

  • cargo test --features catalog,parallel_ingest -> 86 unit + 4 ingest/harness integration tests pass
  • uv run pytest tests/test_catalog.py -> 17/17 pass
  • uv run ruff check scripts/_ingest_profile.py scripts/bench_ingest.py scripts/profile_beamtime_ingest.py tests/test_catalog.py -> clean
  • uv run ty check on branch-touched Python -> no new errors vs main (pre-existing polars/duckdb union-type errors in tests/test_catalog.py confirmed 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 expected

Backward compatibility

  • No public Python API changes on the ingest_beamtime family
  • Progress-callback wire format unchanged (phase labels remain "headers", "catalog", "zarr" via IngestPhase::as_str())
  • src/io/image_mmap.rs::load_image_pixels promoted from fn to pub fn (required for the integration-test binaries; pub(crate) does not reach separate test crates)

Test plan

  • CI green on all platforms
  • Run uv run python scripts/bench_ingest.py --scans 50 --frames-per-scan 50 --width 1024 --height 1024 locally and confirm files_per_second is higher than main baseline
  • Spot-check a real beamtime ingest via uv run python scripts/profile_beamtime_ingest.py --beamtime <path> and verify phase shares match expectation (zarr dominant, catalog small)
  • Verify zarr output on a real beamtime still round-trips through the existing downstream readers (no /processed array consumers broken)

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant