diff --git a/.gitignore b/.gitignore index b429e48..0d277c9 100644 --- a/.gitignore +++ b/.gitignore @@ -73,6 +73,12 @@ instance/ docs/_build/ docs/generated/ docs/auto_examples/ +docs/sg_execution_times.rst + +# Downloaded/generated benchmark data and results (one set per suite) +benchmarks/*/cohorts/ +benchmarks/*/data/ +benchmarks/*/results/ # PyBuilder target/ diff --git a/README.rst b/README.rst index 0db98e4..7769dfe 100644 --- a/README.rst +++ b/README.rst @@ -71,6 +71,39 @@ Both backends expose a similar API: Notes and minor differences: * Chunking vs tiling: HDF5 uses chunks; TileDB uses tiles. We compute tile sizes analogous to chunk sizes to keep write/read patterns similar. -* Compression: HDF5 uses ``gzip`` by default; TileDB defaults to ``zstd`` with shuffle for better speed/ratio. You can switch to ``gzip`` for parity. -* Metadata: HDF5 stores ``column_names`` as a dataset attribute; TileDB stores names as JSON metadata on the array/group. +* Compression: HDF5 defaults to ``gzip`` level 4; TileDB defaults to ``zstd`` level 5. Compressed outputs use byte shuffle; ``--compression none`` disables the filter pipeline. You can switch TileDB to ``gzip`` for parity. +* Metadata: HDF5 stores ``column_names`` as a sibling dataset under each scalar group. TileDB stores names as JSON metadata, and CIFTI outputs also retain their dedicated column-name arrays. * Layout: Both backends keep dimensions in the same order and use zero-based indices. + +Fast conversion and input sources +================================= + +``modelarrayio to-modelarray`` stages one scalar at a time into a typed, +C-contiguous matrix, then writes that matrix in one backend submission. Use +``--workers N`` to bound how many input sources are loaded concurrently. This +meaning applies to local paths and ``s3://`` URLs, every modality, and both +backends; it does not cap HDF5 compression threads or TileDB's internal pools. + +``--workers`` previously controlled concurrent TileDB scalar writes. +``--s3-workers`` is accepted as a deprecated alias for one minor release. +Matching values are accepted, conflicting values are rejected, and the alias +will be removed in the next breaking release. + +Local files and materialized DataLad/git-annex symlinks are read directly. +ModelArrayIO never runs ``datalad get`` implicitly; a broken annex symlink +reports an actionable error. S3 sources use one shared connection pool sized +for the loader concurrency. Set ``MODELARRAYIO_S3_ANON=1`` for public buckets. + +``--dtype float32`` is the faster, smaller default. ``--dtype float64`` keeps +source precision through decoding, staging, and storage, at twice the staging +memory. On the committed 20-subject fixture, three warm runs showed HDF5's +median write phase increasing from 0.00090 s to 0.00138 s (1.54x) and compressed +output from 37,460 to 54,355 bytes (1.45x). That fixture is too small for a +throughput claim; use the reproducible `benchmark harness +`_ to measure the medium ABIDE cohort on deployment +hardware. + +Pass ``--timing-jsonl PATH`` to append structured per-scalar and total runtime +records. Outputs are written to a same-filesystem temporary path and validated +before publication. File replacement is atomic; existing TileDB directories +use a recoverable rename-and-rollback sequence. diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 0000000..a011e06 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,32 @@ +# ModelArrayIO benchmarks + +Each benchmark suite lives in its own folder, named after the `modelarrayio` subcommand it +measures (with underscores). A suite is self-contained: its scripts, committed fixtures, pinned +results (`baselines/`), rendered findings (`reports/`), its own `README.md`, and the gitignored +inputs and outputs it produces. + +| Suite | Measures | Guide | +| --- | --- | --- | +| [`to_modelarray/`](to_modelarray/) | `modelarrayio to-modelarray` runtime and output correctness across NIfTI, CIFTI and MIF on HDF5 and TileDB, including the original-vs-optimized comparison and the CUBIC (HPC) validation ladder | [`to_modelarray/README.md`](to_modelarray/README.md) | + +## Suite layout + +Each suite keeps the same three workspace folders, ignored by Git through the +`benchmarks/*/{data,cohorts,results}/` patterns in `.gitignore`: + +- `data/`: downloaded or synthesized input files. +- `cohorts/`: generated cohort CSVs. Locally materialized cohorts contain absolute paths into + `data/`, so regenerate them (or rewrite the path prefix) after moving the checkout. +- `results//`: per-run outputs, logs, timings and `run_meta` files. Set + `MODELARRAYIO_BENCHMARK_RESULTS_DIR` to write them elsewhere (for example cluster scratch). + +Run every command from the repository root. + +## Adding a suite + +Create `benchmarks//` with its own `README.md`, `fixtures/`, `baselines/` and +`reports/`, and add a row to the table above. When a second suite needs a tool or input that +currently lives in `to_modelarray/` (the shared modules `harness_common.py` and +`output_compare.py`, the cohort generators, `source_identity.py`, the old-commit launcher in +`baselines/`, or a downloaded cohort), move it to `benchmarks/common/` and update both suites; +until then, keep each tool and its data with the suite that uses it. diff --git a/benchmarks/to_modelarray/README.md b/benchmarks/to_modelarray/README.md new file mode 100644 index 0000000..a17cae6 --- /dev/null +++ b/benchmarks/to_modelarray/README.md @@ -0,0 +1,472 @@ +# `to-modelarray` benchmarks + +This folder benchmarks the complete CLI conversion. Results stay under +`benchmarks/to_modelarray/results//` (or `$MODELARRAYIO_BENCHMARK_RESULTS_DIR//`) +and are intentionally not deleted: correctness comparisons and the R read gate need the +produced arrays. + +## Cohorts + +- `fixture`: the committed NIfTI and CIFTI test data, used only for correctness and smoke runs. +- `small`: 100 public ABIDE subjects and one derivative. +- `medium`: 1,102 ABIDE subjects and one derivative (`alff`). +- `wide`: 1,102 subjects and nine derivatives (`alff`, `falff`, `reho`, `vmhc`, `lfcd`, + `degree_binarize`, `degree_weighted`, `eigenvector_binarize`, and + `eigenvector_weighted`). +- `production confirmation`: internal data, run only for the final candidates. + +The tracked fixture manifests are `fixtures/nifti-20.csv` and `fixtures/cifti.csv`. ABIDE is read +anonymously from `s3://fcp-indi`. Generate every public ABIDE cohort as follows: + +```bash +# Small: 100 subjects x 1 derivative +python benchmarks/to_modelarray/make_cohort.py \ + --subjects 100 \ + --derivatives alff \ + --repeat-factor 1 \ + --output benchmarks/to_modelarray/cohorts/abide-small.csv + +# Medium: 1,102 subjects x 1 derivative +python benchmarks/to_modelarray/make_cohort.py \ + --subjects 1102 \ + --derivatives alff \ + --repeat-factor 1 \ + --output benchmarks/to_modelarray/cohorts/abide-medium.csv + +# Wide: 1,102 subjects x 9 derivatives +python benchmarks/to_modelarray/make_cohort.py \ + --subjects 1102 \ + --derivatives alff falff reho vmhc lfcd degree_binarize degree_weighted \ + eigenvector_binarize eigenvector_weighted \ + --repeat-factor 1 \ + --output benchmarks/to_modelarray/cohorts/abide-wide.csv +``` + +The generator lists public keys, intersects subjects having every requested derivative and +`func_mask`, sorts the subject IDs, and takes the first requested count. Pass +`--subject-manifest subjects.txt` to use a frozen newline-delimited subject list without an S3 +listing. Repeats reuse the real source URIs and receive unique cohort IDs; they are valid for +decode-throughput tests but not compression-ratio conclusions. + +Use one downloaded mask as the group mask. For example, download the first mask in the generated +small cohort with the AWS CLI: + +```bash +mkdir -p benchmarks/to_modelarray/data +aws s3 cp \ + s3://fcp-indi/data/Projects/ABIDE_Initiative/Outputs/cpac/filt_global/func_mask/CMU_a_0050642_func_mask.nii.gz \ + benchmarks/to_modelarray/data/group_mask.nii.gz \ + --no-sign-request +``` + +`cohorts/`, `data/` and `results/` hold generated or downloaded files and are ignored by Git (see +[`../README.md`](../README.md#suite-layout)). Set `MODELARRAYIO_S3_ANON=1` for manual CLI runs; +the harness sets it when absent. + +### MIF (fixel) cohorts + +MIF has a fully public, reproducible cohort: the ModelArray walkthrough FDC demo (100 fixel +`.mif` subjects with 602,229 fixels each, plus `index.mif` and `directions.mif`), hosted on OSF. +`make_mif_cohort.py` downloads and extracts it (once) and writes a normalized long-format +`subject_id,scalar_name,source_file` cohort with absolute paths: + +```bash +# Downloads ~212 MiB from https://osf.io/tce9d/download into benchmarks/to_modelarray/data/mif-walkthrough/ +python benchmarks/to_modelarray/make_mif_cohort.py \ + --data-dir benchmarks/to_modelarray/data/mif-walkthrough \ + --output benchmarks/to_modelarray/cohorts/mif-walkthrough.csv +``` + +Run it on both backends with the R gate; MIF needs `--index-file` and `--directions-file` and no +`--mask`: + +```bash +python benchmarks/to_modelarray/run_benchmarks.py \ + --cohort benchmarks/to_modelarray/cohorts/mif-walkthrough.csv \ + --index-file benchmarks/to_modelarray/data/mif-walkthrough/FDC/index.mif \ + --directions-file benchmarks/to_modelarray/data/mif-walkthrough/FDC/directions.mif \ + --tier quick --workers 1 4 --dtypes float32 \ + --backends hdf5 tiledb --repetitions 3 --cache-state warm \ + --r-gate-command 'Rscript benchmarks/to_modelarray/r_read_gate.R {output} {backend}' +``` + +This is real diffusion fixel data, so its runtime, `--workers` scaling and compression ratios are +all quotable — the MIF analogue of the public ABIDE NIfTI cohort. + +### CIFTI cohorts + +ABIDE has no public CIFTI equivalent, so the CIFTI path is benchmarked with `make_cifti_cohort.py`, +which emits a long-format `subject_id,scalar_name,source_file` cohort the harness runs with combined +output (no `--mask`, no `--scalar-columns`). + +CIFTI matters because its bottleneck is unlike NIfTI's: uncompressed `.dscalar.nii` files carry a +large CIFTI XML header that nibabel parses at roughly one second per file while holding the GIL, so +thread workers cannot speed it up. The loader's raw greyordinate fast path +(`MODELARRAYIO_CIFTI_RAW`, on by default) reads the data block directly and sidesteps that parse; +benchmark both settings to measure it. + +The generator has three modes: + +```bash +# 1. Default synthesize: reproducible anywhere with no real data. Uses scalar-appropriate marginal +# ranges and a medial-wall zero block, so it is valid for load/runtime timing but NOT compression +# ratios (its per-vertex values carry no real spatial autocorrelation). +python benchmarks/to_modelarray/make_cifti_cohort.py synthesize \ + --output-dir benchmarks/to_modelarray/data/cifti-synth \ + --subjects 40 --scalars thickness area sulc \ + --grayordinates 327684 --seed 0 \ + --cohort benchmarks/to_modelarray/cohorts/cifti-synth.csv + +# 2. Reference-profile synthesize: derives the medial-wall mask and per-vertex mean/std from a +# small set of real dscalars and draws fresh noise around the real mean map, so compression is +# realistic (measured 1.38x vs 1.41x for the real reference) while no real subject row is copied. +python benchmarks/to_modelarray/make_cifti_cohort.py synthesize \ + --output-dir benchmarks/to_modelarray/data/cifti-synth-real \ + --subjects 40 --scalars thickness area sulc \ + --reference-dir /path/to/real/dscalars --seed 0 \ + --cohort benchmarks/to_modelarray/cohorts/cifti-synth-real.csv + +# 3. Subset: slice an existing long-format internal CIFTI cohort without copying files. +python benchmarks/to_modelarray/make_cifti_cohort.py subset \ + --input /path/to/rbc_freesurferpost_cohort.csv \ + --output benchmarks/to_modelarray/cohorts/cifti-real-small.csv \ + --subjects 100 --scalars thickness area sulc + +# Run it (fast path on vs. off) to quantify the raw-read win. +python benchmarks/to_modelarray/run_benchmarks.py \ + --cohort benchmarks/to_modelarray/cohorts/cifti-synth-real.csv \ + --tier quick --backends tiledb hdf5 \ + --cache-state warm --run-id cifti-synth-raw \ + --modelarrayio .venv/bin/modelarrayio +MODELARRAYIO_CIFTI_RAW=0 python benchmarks/to_modelarray/run_benchmarks.py \ + --cohort benchmarks/to_modelarray/cohorts/cifti-synth-real.csv \ + --tier quick --backends tiledb hdf5 \ + --cache-state warm --run-id cifti-synth-nibabel \ + --modelarrayio .venv/bin/modelarrayio +``` + +The `subset` mode reads any long-format CIFTI cohort, keeps the first `--subjects` subjects (by +first appearance, using `subject_id` and `session_id` when present) for the requested scalars, and +never lists or copies files, so it is safe to point at a large internal cohort CSV. The +`--reference-dir` mode reads only aggregate per-vertex statistics, never an individual subject row. + +## Running + +The cohort controls input size and scalar count; `--tier` independently controls the parameter +grid. The complete grids are: + +| Tier | Loader workers | Dtypes | Chunk targets (MiB) | Codecs | Cases/repetition | +| --- | --- | --- | --- | --- | ---: | +| `quick` | 1 | float32 | 2 | HDF5 gzip-4; TileDB zstd-5 | 2 | +| `medium` | 1, 4, 8 | float32, float64 | 2 | HDF5 gzip-4; TileDB zstd-5 | 12 | +| `full` | 1, 2, 4, 8, allocated CPUs | float32, float64 | 8, 32, 64, 128 | HDF5 gzip-1/gzip-4/none; TileDB zstd-1/zstd-3/zstd-5/none | hardware-dependent | + +The full tier also sweeps HDF5 gzip compression workers over 1, 4, and allocated CPUs. If `W` is +the number of distinct loader-worker settings and `P` is the number of distinct compression-worker +settings, it runs `8 * W * (2 * P + 5)` cases per repetition: 440 cases on a 14-CPU host. Three +repetitions therefore launch 1,320 conversions before accounting for multiple scalars. Use it only +for an intentional overnight or cluster sweep. + +### Full validation suite + +Run the complete Python regression and runtime-integration suite before benchmarking: + +```bash +python -m pytest -q +ruff check src test benchmarks docs/examples +ruff format --check src test benchmarks docs/examples +Rscript -e "invisible(parse(file='benchmarks/to_modelarray/r_read_gate.R')); invisible(parse(file='benchmarks/to_modelarray/plot_results.R'))" +``` + +Run every benchmark parameter combination once on the committed NIfTI fixture. This is an +exhaustive correctness/configuration test, not a throughput result: + +```bash +python benchmarks/to_modelarray/run_benchmarks.py \ + --cohort benchmarks/to_modelarray/fixtures/nifti-20.csv \ + --mask test/data_voxel_toy/group_mask_FA.nii.gz \ + --scalar-columns FA \ + --tier full \ + --repetitions 1 \ + --cache-state warm \ + --run-id fixture-nifti-full \ + --modelarrayio .venv/bin/modelarrayio \ + --r-gate-command 'Rscript benchmarks/to_modelarray/r_read_gate.R {output} {backend}' +``` + +On the 14-CPU reference host this is 440 cases: 280 HDF5 and 160 TileDB. The 2026-09-19 run +passed all 440 converter and R-gate checks; its tracked audit is +[`baselines/full_fixture_validation_20260919.csv`](baselines/full_fixture_validation_20260919.csv). +The run was split into worker-specific shards to execute independent R gates concurrently. To do +the same, launch the command once per worker in `1 2 4 8 14`, add `--workers WORKER`, and use a +unique `--result-shard worker-WORKER` with a shared run ID. + +Wide cohorts split by default and long-format cohorts combine, matching the converter. The +explicit `--split-files` and `--no-split-files` flags remain available for compatibility tests. +Every reportable command below includes the R read gate. + +### Fixture smoke and correctness + +Run both committed fixture modalities before public or production sweeps: + +```bash +python benchmarks/to_modelarray/run_benchmarks.py \ + --cohort benchmarks/to_modelarray/fixtures/nifti-20.csv \ + --mask test/data_voxel_toy/group_mask_FA.nii.gz \ + --scalar-columns FA \ + --tier quick \ + --repetitions 3 \ + --cache-state warm \ + --r-gate-command 'Rscript benchmarks/to_modelarray/r_read_gate.R {output} {backend}' +``` + +```bash +python benchmarks/to_modelarray/run_benchmarks.py \ + --cohort benchmarks/to_modelarray/fixtures/cifti.csv \ + --tier quick \ + --repetitions 3 \ + --cache-state warm \ + --r-gate-command 'Rscript benchmarks/to_modelarray/r_read_gate.R {output} {backend}' +``` + +### Small iteration case + +```bash +python benchmarks/to_modelarray/run_benchmarks.py \ + --cohort benchmarks/to_modelarray/cohorts/abide-small.csv \ + --mask benchmarks/to_modelarray/data/group_mask.nii.gz \ + --scalar-columns alff \ + --tier quick \ + --repetitions 3 \ + --cache-state warm \ + --r-gate-command 'Rscript benchmarks/to_modelarray/r_read_gate.R {output} {backend}' +``` + +### Medium comparison case + +The medium tier is the default public performance comparison. It runs both backends at loader +worker counts 1, 4, and 8 and both storage dtypes: + +```bash +python benchmarks/to_modelarray/run_benchmarks.py \ + --cohort benchmarks/to_modelarray/cohorts/abide-medium.csv \ + --mask benchmarks/to_modelarray/data/group_mask.nii.gz \ + --scalar-columns alff \ + --tier medium \ + --repetitions 3 \ + --cache-state warm \ + --r-gate-command 'Rscript benchmarks/to_modelarray/r_read_gate.R {output} {backend}' +``` + +### Wide scalar case + +First validate the nine-scalar path with the quick tier: + +```bash +python benchmarks/to_modelarray/run_benchmarks.py \ + --cohort benchmarks/to_modelarray/cohorts/abide-wide.csv \ + --mask benchmarks/to_modelarray/data/group_mask.nii.gz \ + --scalar-columns alff falff reho vmhc lfcd degree_binarize degree_weighted \ + eigenvector_binarize eigenvector_weighted \ + --tier quick \ + --repetitions 3 \ + --cache-state warm \ + --r-gate-command 'Rscript benchmarks/to_modelarray/r_read_gate.R {output} {backend}' +``` + +Run the exhaustive grid only when its storage and scheduling cost is acceptable: + +```bash +python benchmarks/to_modelarray/run_benchmarks.py \ + --cohort benchmarks/to_modelarray/cohorts/abide-wide.csv \ + --mask benchmarks/to_modelarray/data/group_mask.nii.gz \ + --scalar-columns alff falff reho vmhc lfcd degree_binarize degree_weighted \ + eigenvector_binarize eigenvector_weighted \ + --tier full \ + --repetitions 3 \ + --cache-state warm \ + --r-gate-command 'Rscript benchmarks/to_modelarray/r_read_gate.R {output} {backend}' +``` + +Use `--backends`, `--workers`, and `--result-shard` to divide a full sweep into explicit, +non-overlapping jobs. A shard label isolates artifacts but does not partition the grid for you. + +### Production confirmation case + +Production data are site-specific. Set these two paths, then run the final public-cohort settings +on the internal cohort. This template uses the quick-tier codecs with an explicit finalist worker +count and dtype: + +```bash +COHORT_CSV=/absolute/path/to/production-cohort.csv +GROUP_MASK=/absolute/path/to/production-group-mask.nii.gz +python benchmarks/to_modelarray/run_benchmarks.py \ + --cohort "$COHORT_CSV" \ + --mask "$GROUP_MASK" \ + --scalar-columns alff \ + --tier quick \ + --workers 8 \ + --dtypes float32 \ + --repetitions 3 \ + --cache-state warm \ + --r-gate-command 'Rscript benchmarks/to_modelarray/r_read_gate.R {output} {backend}' +``` + +For CIFTI cohorts, omit `--mask`. For MIF cohorts, omit `--mask` and pass +`--index-file` plus `--directions-file` instead. Long-format cohorts record +their actual `scalar_name` values in result rows and R-gate commands. + +Important comparisons must use the same cohort, cache state, filesystem, and empty case output. +The harness records `--cache-state` as metadata but does not warm or clear caches itself. +Run at least three repetitions and compare medians. `--result-shard node-01` writes +`results.node-01.csv`; its outputs, logs, timing, and metadata are also stored below a private +`shard-node-01` directory, so parallel workers cannot overwrite one another. Reusing the same +run ID with the same configuration skips completed case keys and writes an interrupted case to a +new retry directory, preserving all prior artifacts. A changed configuration requires a new run +ID. The plotting script merges all shard result files. `run_meta..json` records commands, +versions, the cohort hash, CPU/RAM allocation, and TileDB concurrency settings. Each result row includes phase counters, output/file/fragment +counts, input bytes, dtype, codec, loader workers, compression workers, and resource usage. +Before a production sweep, set an absolute stopping rule with +`--acceptance-files-per-minute`; a case below that target makes the harness exit nonzero. Use +`--compression-workers 1 2 4 8` to sweep the internal HDF5 gzip pool without adding another +public converter option. + +Runs without an R gate are recorded as `unverified`, remain useful as smoke tests, and are +excluded from the plotting script's candidate set. A failed gate marks the case failed and makes +the sweep exit nonzero. The default two-backend commands require the R packages `rhdf5`, +`HDF5Array`, `ModelArray`, and `tiledb`. + +For HDF5, the gate constructs a `ModelArray` through `rhdf5`/`HDF5Array` and reads a representative +slice. For TileDB, it opens and reads the array through the R `tiledb` package. A failed gate makes +the sweep fail. Deployment environments should extend the TileDB branch to construct the local +ModelArray TileDB adapter when one is available. + +Create plots with: + +```bash +RUN_DIR=$(ls -dt benchmarks/to_modelarray/results/*/ | head -n 1) +Rscript benchmarks/to_modelarray/plot_results.R "$RUN_DIR" +``` + +## CUBIC (HPC) validation + +The cluster ladder — container-based environment, SLURM job templates, the old-code reference +comparison, and the scalar-sharded job-array recipe — lives in [`cubic/`](cubic/README.md). + +Pinned CUBIC results (original `9ab13dd` vs optimized, real RBC freesurferpost data): +[`baselines/cubic_cross_modality_9ab13dd_vs_optimized_20260922.csv`](baselines/cubic_cross_modality_9ab13dd_vs_optimized_20260922.csv) +(step 2: NIfTI 1.40×/1.46×, MIF 3.70×/4.11×, CIFTI 31.8×/31.6×) and +[`baselines/cubic_production_scale_9ab13dd_vs_optimized_20260923.csv`](baselines/cubic_production_scale_9ab13dd_vs_optimized_20260923.csv) +(one full 8,693-subject scalar 27.6×/27.5× at one worker, 69–91 s at eight; 15-scalar job array in +21 min; all 70 production scalars in ≈22 min of array wall-clock, every output bitwise-equal to the +deployed reference). The CUBIC PDF is `benchmarks/to_modelarray/reports/to-modelarray-cubic-benchmark.pdf`. + +## Original versus optimized comparison + +`compare_revisions.py` runs two `modelarrayio` executables against the same materialized cohort, +inputs, environment, backend settings, and warm-cache policy. It infers the modality from the +cohort (so it drives NIfTI, CIFTI and MIF), discards one warm-up per implementation and backend, +counterbalances the remaining runs, applies the R gate, and compares every optimized output with +the baseline for shape, dtype, source order, NaN placement, and values. + +The original code is pinned to commit `9ab13dd` (the last commit before the optimization branch), +exported with `git archive` into an isolated cache and run under the current checkout's Python +environment by `benchmarks/to_modelarray/baselines/modelarrayio-baseline-9ab13dd`. Override the cache or +interpreter with `MODELARRAYIO_BASELINE_ROOT` or `MODELARRAYIO_PYTHON`. + +The reportable result is the CUBIC run above (step 2 of the ladder, on real CIFTI). The commands +below reproduce the same comparison from public data on any machine; CIFTI uses the realistic +synthesizer there because no real CIFTI cohort is public. Run one comparison per modality (each +writes its own results directory): + +```bash +BASELINE=benchmarks/to_modelarray/baselines/modelarrayio-baseline-9ab13dd +OPTIMIZED=.venv/bin/modelarrayio +GATE=benchmarks/to_modelarray/r_read_gate.R + +# NIfTI (100 ABIDE subjects) +python benchmarks/to_modelarray/compare_revisions.py \ + --cohort benchmarks/to_modelarray/cohorts/abide-nifti100-local.csv \ + --mask benchmarks/to_modelarray/data/group_mask.nii.gz --scalar-columns alff \ + --baseline-executable "$BASELINE" --optimized-executable "$OPTIMIZED" \ + --results-dir benchmarks/to_modelarray/results/original-vs-optimized-9ab13dd-nifti \ + --repetitions 3 --warmups 1 --r-gate-script "$GATE" --baseline-revision 9ab13dd + +# MIF (100 walkthrough FDC subjects) +python benchmarks/to_modelarray/compare_revisions.py \ + --cohort benchmarks/to_modelarray/cohorts/mif-walkthrough.csv \ + --index-file benchmarks/to_modelarray/data/mif-walkthrough/FDC/index.mif \ + --directions-file benchmarks/to_modelarray/data/mif-walkthrough/FDC/directions.mif \ + --baseline-executable "$BASELINE" --optimized-executable "$OPTIMIZED" \ + --results-dir benchmarks/to_modelarray/results/original-vs-optimized-9ab13dd-mif \ + --repetitions 3 --warmups 1 --r-gate-script "$GATE" --baseline-revision 9ab13dd + +# CIFTI (20 realistic-synth subjects x 3 scalars) +python benchmarks/to_modelarray/compare_revisions.py \ + --cohort benchmarks/to_modelarray/cohorts/cifti-cmp20.csv \ + --baseline-executable "$BASELINE" --optimized-executable "$OPTIMIZED" \ + --results-dir benchmarks/to_modelarray/results/original-vs-optimized-9ab13dd-cifti \ + --repetitions 3 --warmups 1 --r-gate-script "$GATE" --baseline-revision 9ab13dd +``` + +Then render the cross-modality findings PDF from the three run directories. It is written to +`benchmarks/to_modelarray/results/` by default; only a run that supersedes the pinned CUBIC report +belongs in `reports/`: + +```bash +python benchmarks/to_modelarray/create_comparison_report.py \ + --run-dir benchmarks/to_modelarray/results/original-vs-optimized-9ab13dd-nifti \ + --run-dir benchmarks/to_modelarray/results/original-vs-optimized-9ab13dd-mif \ + --run-dir benchmarks/to_modelarray/results/original-vs-optimized-9ab13dd-cifti \ + --cohort-source 'nifti=ABIDE alff, public S3 (make_cohort.py, materialize_cohort.py)' \ + --cohort-source 'mif=ModelArray walkthrough FDC, public OSF (make_mif_cohort.py)' \ + --cohort-source 'cifti=synthetic, real reference profile (make_cifti_cohort.py --reference-dir)' +``` + +The report shows the original-vs-optimized median wall time, speedup, downstream R-read time, +output file count and bitwise correctness for every modality on both backends. Its text comes from +the run metadata (date, platform, repetitions, warm-ups, workers, chunk target, cohort file and +hash); `--cohort-source` adds where each cohort came from, which the metadata cannot tell. It needs ReportLab, +which is not a declared dependency. + +## What the results show about the design + +The optimized design and the alternatives it rejects are described in +[`design/TO_MODELARRAY_RUNTIME_OPTIMIZATION_PLAN.md`](../../design/TO_MODELARRAY_RUNTIME_OPTIMIZATION_PLAN.md). +The benchmarks measure its two main choices: + +- **One dense TileDB write per scalar** instead of one fragment per stripe: on the CUBIC comparison + this reduced the output from 350 files to 8 for MIF and from 596 to 38 for CIFTI while preserving + values and metadata. +- **Concurrent loader workers.** Loader-worker scaling is pinned by the CUBIC step-3 sweep in + [`baselines/cubic_production_scale_9ab13dd_vs_optimized_20260923.csv`](baselines/cubic_production_scale_9ab13dd_vs_optimized_20260923.csv): + one full production CIFTI scalar (8,693 subjects), one repetition per setting. Workers 1/4/8 + measured 403.9/120.9/90.6 seconds for HDF5 and 186.3/83.2/68.6 seconds for TileDB, while peak + memory stayed flat (about 12 GB for HDF5 and 20 GB for TileDB) because it is dominated by the + staging matrix. Other modalities were not swept at production scale, so the CLI keeps the + conservative default of one worker; each production workload should validate `--workers 4` or + `--workers 8` with the medium tier and multiple repetitions. + +## Interpretation and limits + +The optimization target is end-to-end elapsed time while preserving values, NaN locations, +source order, metadata, and R readability. `writer_idle_seconds` is the time a sequential +writer awaits staging; `loader_blocked_seconds` and queue occupancy remain zero until the +optional loader/writer overlap pipeline exists. `scheduler_wait_seconds` and the in-flight +source counters describe the bounded loader itself. HDF5 `compression_workers` is an internal runtime policy; +`--workers` controls only concurrent input loads. TileDB's compute and I/O pools are recorded and +never capped to `--workers`. + +`baselines/baseline_fixture.csv` pins three-run medians from the committed 20-subject fixture. +It is a schema/size and gross-regression reference, not a production throughput target. The +reportable original-versus-optimized results are the CUBIC CSVs above: real data, the original +source tree at commit `9ab13dd`, float32 and warm local data. Step 2 uses one loader worker and +step 3 sweeps workers; remote S3 throughput remains a separate measurement. + +ABIDE's roughly 65k selected voxels and small compressed files emphasize per-file overhead more +than production CIFTI/fixel inputs. MIF has a public equivalent cohort (the walkthrough FDC demo, +via `make_mif_cohort.py`); CIFTI has no public real cohort, so it uses the reproducible synthesizer +(`make_cifti_cohort.py`), optionally seeded from a small real reference for realistic compression. +Repeated local rows measure a warm-cache decode ceiling, and repeated values compress +unrealistically well. Never compare warm and cold cache rows or quote synthetic/layout-sweep rows +as end-to-end results. diff --git a/benchmarks/to_modelarray/baselines/baseline_fixture.csv b/benchmarks/to_modelarray/baselines/baseline_fixture.csv new file mode 100644 index 0000000..96dd153 --- /dev/null +++ b/benchmarks/to_modelarray/baselines/baseline_fixture.csv @@ -0,0 +1,5 @@ +cohort,cohort_sha256,measured_at,platform,repetitions,aggregation,backend,workers,dtype,compression,compression_level,target_chunk_mb,wall_seconds,write_seconds,output_bytes,output_file_count,fragment_count,r_gate +fixture-20-nifti,867c3eb9194d51b7f20ce511a308fa818d1bbe75124f6fc66f76bcc7c85508fb,2026-09-18,macOS-arm64-14cpu,3,median,hdf5,1,float32,gzip,4,2.0,0.4263049580,0.0008974160,37460,1,,passed +fixture-20-nifti,867c3eb9194d51b7f20ce511a308fa818d1bbe75124f6fc66f76bcc7c85508fb,2026-09-18,macOS-arm64-14cpu,3,median,hdf5,1,float64,gzip,4,2.0,0.4020678340,0.0013775420,54355,1,,passed +fixture-20-nifti,867c3eb9194d51b7f20ce511a308fa818d1bbe75124f6fc66f76bcc7c85508fb,2026-09-18,macOS-arm64-14cpu,3,median,tiledb,1,float32,zstd,5,2.0,0.4307647080,0.0062273330,20018,8,1,passed +fixture-20-nifti,867c3eb9194d51b7f20ce511a308fa818d1bbe75124f6fc66f76bcc7c85508fb,2026-09-18,macOS-arm64-14cpu,3,median,tiledb,1,float64,zstd,5,2.0,0.4465746670,0.0061256250,37116,8,1,passed diff --git a/benchmarks/to_modelarray/baselines/cubic_cross_modality_9ab13dd_vs_optimized_20260922.csv b/benchmarks/to_modelarray/baselines/cubic_cross_modality_9ab13dd_vs_optimized_20260922.csv new file mode 100644 index 0000000..6ad8b5f --- /dev/null +++ b/benchmarks/to_modelarray/baselines/cubic_cross_modality_9ab13dd_vs_optimized_20260922.csv @@ -0,0 +1,7 @@ +modality,backend,n_subjects,n_elements,n_scalars,baseline_revision,baseline_median_wall_seconds,optimized_median_wall_seconds,speedup,wall_time_reduction_percent,baseline_median_r_read_seconds,optimized_median_r_read_seconds,baseline_output_bytes,optimized_output_bytes,baseline_output_file_count,optimized_output_file_count,values_equal,max_abs_difference +nifti,hdf5,100,66502,1,9ab13dd,3.0238,2.1674,1.3951,28.32,4.0079,4.0271,19965668,19967444,1,1,True,0.0 +nifti,tiledb,100,66502,1,9ab13dd,3.4056,2.332,1.4604,31.52,0.2337,0.2294,18351911,18304134,44,8,True,0.0 +mif,hdf5,100,602229,1,9ab13dd,10.7605,2.9104,3.6972,72.95,4.0573,3.9861,229198658,229200434,1,1,True,0.0 +mif,tiledb,100,602229,1,9ab13dd,14.5217,3.5334,4.1098,75.67,0.2643,0.2358,194154640,193700782,350,8,True,0.0 +cifti,hdf5,100,327684,3,9ab13dd,385.7822,12.124,31.8197,96.86,4.0078,3.9943,289790692,289791564,1,1,True,0.0 +cifti,tiledb,100,327684,3,9ab13dd,396.466,12.5594,31.5673,96.83,0.4726,0.2439,288509910,287768750,596,38,True,0.0 diff --git a/benchmarks/to_modelarray/baselines/cubic_production_scale_9ab13dd_vs_optimized_20260923.csv b/benchmarks/to_modelarray/baselines/cubic_production_scale_9ab13dd_vs_optimized_20260923.csv new file mode 100644 index 0000000..1320ea1 --- /dev/null +++ b/benchmarks/to_modelarray/baselines/cubic_production_scale_9ab13dd_vs_optimized_20260923.csv @@ -0,0 +1,109 @@ +step,scalars,subjects,backend,workers,baseline_median_wall_seconds,optimized_median_wall_seconds,speedup,input_files_per_minute,peak_memory_gb,fragment_count,fast_path_rows,passed +3_full_scalar_original_vs_optimized,area,8693,hdf5,1,11854.8,429.9,27.58,,,,, +3_full_scalar_original_vs_optimized,area,8693,tiledb,1,11501.6,418.1,27.51,,,,, +3_full_scalar_worker_sweep,area,8693,hdf5,1,,403.9,,1291,12.1,,8692, +3_full_scalar_worker_sweep,area,8693,hdf5,4,,120.9,,4313,12.0,,8692, +3_full_scalar_worker_sweep,area,8693,hdf5,8,,90.6,,5756,12.2,,8692, +3_full_scalar_worker_sweep,area,8693,tiledb,1,,186.3,,2800,20.3,1,8692, +3_full_scalar_worker_sweep,area,8693,tiledb,4,,83.2,,6271,20.3,1,8692, +3_full_scalar_worker_sweep,area,8693,tiledb,8,,68.6,,7603,20.4,1,8692, +4_scalar_sharded_array,area,8693,hdf5,8,,137.4,,,,,,True +4_scalar_sharded_array,area,8693,tiledb,8,,85.1,,,,,,True +4_scalar_sharded_array,areaPial,8693,hdf5,8,,140.6,,,,,,True +4_scalar_sharded_array,areaPial,8693,tiledb,8,,106.8,,,,,,True +4_scalar_sharded_array,curv,8693,hdf5,8,,134.5,,,,,,True +4_scalar_sharded_array,curv,8693,tiledb,8,,102.2,,,,,,True +4_scalar_sharded_array,desc_fwhm0_area,8693,hdf5,8,,136.5,,,,,,True +4_scalar_sharded_array,desc_fwhm0_area,8693,tiledb,8,,105.3,,,,,,True +4_scalar_sharded_array,desc_fwhm0_areaPial,8693,hdf5,8,,141.3,,,,,,True +4_scalar_sharded_array,desc_fwhm0_areaPial,8693,tiledb,8,,99.8,,,,,,True +4_scalar_sharded_array,desc_fwhm0_curv,8693,hdf5,8,,140.5,,,,,,True +4_scalar_sharded_array,desc_fwhm0_curv,8693,tiledb,8,,95.7,,,,,,True +4_scalar_sharded_array,desc_fwhm0_jacobianWhite,8693,hdf5,8,,112.8,,,,,,True +4_scalar_sharded_array,desc_fwhm0_jacobianWhite,8693,tiledb,8,,81.0,,,,,,True +4_scalar_sharded_array,desc_fwhm0_sulc,8693,hdf5,8,,140.7,,,,,,True +4_scalar_sharded_array,desc_fwhm0_sulc,8693,tiledb,8,,118.4,,,,,,True +4_scalar_sharded_array,desc_fwhm0_thickness,8693,hdf5,8,,139.9,,,,,,True +4_scalar_sharded_array,desc_fwhm0_thickness,8693,tiledb,8,,117.0,,,,,,True +4_scalar_sharded_array,desc_fwhm0_volume,8693,hdf5,8,,118.7,,,,,,True +4_scalar_sharded_array,desc_fwhm0_volume,8693,tiledb,8,,93.3,,,,,,True +4_scalar_sharded_array,desc_fwhm0_wgPct,8693,hdf5,8,,114.2,,,,,,True +4_scalar_sharded_array,desc_fwhm0_wgPct,8693,tiledb,8,,109.7,,,,,,True +4_scalar_sharded_array,desc_fwhm0_whiteH,8693,hdf5,8,,139.0,,,,,,True +4_scalar_sharded_array,desc_fwhm0_whiteH,8693,tiledb,8,,109.9,,,,,,True +4_scalar_sharded_array,desc_fwhm0_whiteK,8693,hdf5,8,,124.1,,,,,,True +4_scalar_sharded_array,desc_fwhm0_whiteK,8693,tiledb,8,,97.8,,,,,,True +4_scalar_sharded_array,desc_fwhm10_area,8693,hdf5,8,,111.9,,,,,,True +4_scalar_sharded_array,desc_fwhm10_area,8693,tiledb,8,,74.7,,,,,,True +4_scalar_sharded_array,desc_fwhm10_areaPial,8693,hdf5,8,,115.0,,,,,,True +4_scalar_sharded_array,desc_fwhm10_areaPial,8693,tiledb,8,,69.7,,,,,,True +5_production_build,area,8693,hdf5,8,,113.0,,,,,,True +5_production_build,areaPial,8693,hdf5,8,,114.8,,,,,,True +5_production_build,curv,8693,hdf5,8,,139.7,,,,,,True +5_production_build,desc_fwhm0_area,8693,hdf5,8,,138.2,,,,,,True +5_production_build,desc_fwhm0_areaPial,8693,hdf5,8,,138.7,,,,,,True +5_production_build,desc_fwhm0_curv,8693,hdf5,8,,144.1,,,,,,True +5_production_build,desc_fwhm0_jacobianWhite,8693,hdf5,8,,186.3,,,,,,True +5_production_build,desc_fwhm0_thickness,8693,hdf5,8,,181.3,,,,,,True +5_production_build,desc_fwhm0_sulc,8693,hdf5,8,,184.3,,,,,,True +5_production_build,desc_fwhm0_volume,8693,hdf5,8,,192.1,,,,,,True +5_production_build,desc_fwhm0_wgPct,8693,hdf5,8,,111.9,,,,,,True +5_production_build,desc_fwhm0_whiteH,8693,hdf5,8,,114.5,,,,,,True +5_production_build,desc_fwhm0_whiteK,8693,hdf5,8,,148.8,,,,,,True +5_production_build,desc_fwhm10_area,8693,hdf5,8,,130.5,,,,,,True +5_production_build,desc_fwhm10_areaPial,8693,hdf5,8,,135.0,,,,,,True +5_production_build,desc_fwhm10_curv,8693,hdf5,8,,139.8,,,,,,True +5_production_build,desc_fwhm10_wgPct,8693,hdf5,8,,108.7,,,,,,True +5_production_build,desc_fwhm10_volume,8693,hdf5,8,,132.4,,,,,,True +5_production_build,desc_fwhm10_whiteH,8693,hdf5,8,,113.5,,,,,,True +5_production_build,desc_fwhm10_thickness,8693,hdf5,8,,165.1,,,,,,True +5_production_build,desc_fwhm10_jacobianWhite,8693,hdf5,8,,165.1,,,,,,True +5_production_build,desc_fwhm10_sulc,8693,hdf5,8,,170.6,,,,,,True +5_production_build,desc_fwhm10_whiteK,8693,hdf5,8,,139.8,,,,,,True +5_production_build,desc_fwhm15_area,8693,hdf5,8,,132.2,,,,,,True +5_production_build,desc_fwhm15_curv,8693,hdf5,8,,112.6,,,,,,True +5_production_build,desc_fwhm15_sulc,8693,hdf5,8,,113.6,,,,,,True +5_production_build,desc_fwhm15_areaPial,8693,hdf5,8,,136.0,,,,,,True +5_production_build,desc_fwhm15_jacobianWhite,8693,hdf5,8,,137.0,,,,,,True +5_production_build,desc_fwhm15_wgPct,8693,hdf5,8,,130.0,,,,,,True +5_production_build,desc_fwhm15_thickness,8693,hdf5,8,,144.3,,,,,,True +5_production_build,desc_fwhm15_volume,8693,hdf5,8,,148.7,,,,,,True +5_production_build,desc_fwhm15_whiteH,8693,hdf5,8,,145.9,,,,,,True +5_production_build,desc_fwhm15_whiteK,8693,hdf5,8,,140.2,,,,,,True +5_production_build,desc_fwhm20_areaPial,8693,hdf5,8,,110.7,,,,,,True +5_production_build,desc_fwhm20_area,8693,hdf5,8,,130.5,,,,,,True +5_production_build,desc_fwhm20_curv,8693,hdf5,8,,113.7,,,,,,True +5_production_build,desc_fwhm20_jacobianWhite,8693,hdf5,8,,132.0,,,,,,True +5_production_build,desc_fwhm20_thickness,8693,hdf5,8,,130.9,,,,,,True +5_production_build,desc_fwhm20_sulc,8693,hdf5,8,,166.8,,,,,,True +5_production_build,desc_fwhm20_whiteK,8693,hdf5,8,,112.9,,,,,,True +5_production_build,desc_fwhm20_wgPct,8693,hdf5,8,,162.1,,,,,,True +5_production_build,desc_fwhm20_volume,8693,hdf5,8,,168.2,,,,,,True +5_production_build,desc_fwhm20_whiteH,8693,hdf5,8,,168.9,,,,,,True +5_production_build,desc_fwhm25_curv,8693,hdf5,8,,112.8,,,,,,True +5_production_build,desc_fwhm25_area,8693,hdf5,8,,129.5,,,,,,True +5_production_build,desc_fwhm25_areaPial,8693,hdf5,8,,133.5,,,,,,True +5_production_build,desc_fwhm25_jacobianWhite,8693,hdf5,8,,127.8,,,,,,True +5_production_build,desc_fwhm25_sulc,8693,hdf5,8,,135.7,,,,,,True +5_production_build,desc_fwhm25_thickness,8693,hdf5,8,,127.9,,,,,,True +5_production_build,desc_fwhm25_volume,8693,hdf5,8,,107.2,,,,,,True +5_production_build,desc_fwhm5_area,8693,hdf5,8,,112.8,,,,,,True +5_production_build,desc_fwhm25_wgPct,8693,hdf5,8,,155.1,,,,,,True +5_production_build,desc_fwhm5_areaPial,8693,hdf5,8,,141.2,,,,,,True +5_production_build,desc_fwhm25_whiteH,8693,hdf5,8,,163.9,,,,,,True +5_production_build,desc_fwhm5_curv,8693,hdf5,8,,144.2,,,,,,True +5_production_build,desc_fwhm25_whiteK,8693,hdf5,8,,168.1,,,,,,True +5_production_build,desc_fwhm5_jacobianWhite,8693,hdf5,8,,133.8,,,,,,True +5_production_build,desc_fwhm5_volume,8693,hdf5,8,,111.3,,,,,,True +5_production_build,desc_fwhm5_wgPct,8693,hdf5,8,,111.7,,,,,,True +5_production_build,desc_fwhm5_thickness,8693,hdf5,8,,137.5,,,,,,True +5_production_build,desc_fwhm5_sulc,8693,hdf5,8,,166.3,,,,,,True +5_production_build,desc_fwhm5_whiteK,8693,hdf5,8,,144.7,,,,,,True +5_production_build,sulc,8693,hdf5,8,,143.8,,,,,,True +5_production_build,volume,8693,hdf5,8,,142.5,,,,,,True +5_production_build,wgPct,8693,hdf5,8,,111.2,,,,,,True +5_production_build,desc_fwhm5_whiteH,8693,hdf5,8,,177.7,,,,,,True +5_production_build,jacobianWhite,8693,hdf5,8,,175.0,,,,,,True +5_production_build,thickness,8693,hdf5,8,,172.5,,,,,,True +5_production_build,whiteH,8693,hdf5,8,,118.5,,,,,,True +5_production_build,whiteK,8693,hdf5,8,,149.1,,,,,,True diff --git a/benchmarks/to_modelarray/baselines/full_fixture_validation_20260919.csv b/benchmarks/to_modelarray/baselines/full_fixture_validation_20260919.csv new file mode 100644 index 0000000..d951b40 --- /dev/null +++ b/benchmarks/to_modelarray/baselines/full_fixture_validation_20260919.csv @@ -0,0 +1,2 @@ +measured_at,run_id,cohort,cohort_sha256,platform,allocated_cpus,repetitions,total_cases,hdf5_cases,tiledb_cases,loader_workers,dtypes,target_chunk_mb,hdf5_codecs,tiledb_codecs,hdf5_compression_workers,passed_cases,failed_cases,r_gate_passed_cases,status +2026-09-19,fixture-nifti-full-sharded-20260919,nifti-20,23abf80d6356e0ebb2f8e1c0320bb254735c3654364d8df78a9db497a3191c52,macOS-arm64-M3-Max,14,1,440,280,160,"1,2,4,8,14","float32,float64","8,32,64,128","gzip-1,gzip-4,none","zstd-1,zstd-3,zstd-5,none","1,4,14",440,0,440,passed diff --git a/benchmarks/to_modelarray/baselines/modelarrayio-baseline-9ab13dd b/benchmarks/to_modelarray/baselines/modelarrayio-baseline-9ab13dd new file mode 100755 index 0000000..1b69ddb --- /dev/null +++ b/benchmarks/to_modelarray/baselines/modelarrayio-baseline-9ab13dd @@ -0,0 +1,48 @@ +#!/bin/sh +set -eu + +# Launch the unmodified ModelArrayIO sources from the commit used by the +# published benchmark while reusing this checkout's Python environment. +BASELINE_REVISION=9ab13dd7a4efaed5b86c5d5a184515c4a6968a67 +SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) +REPOSITORY_ROOT=$(CDPATH= cd -- "$SCRIPT_DIR/../../.." && pwd) +BASELINE_ROOT=${MODELARRAYIO_BASELINE_ROOT:-${TMPDIR:-/tmp}/modelarrayio-baseline-9ab13dd} + +prepare_baseline() { + if [ -f "$BASELINE_ROOT/src/modelarrayio/__init__.py" ]; then + return + fi + if [ -e "$BASELINE_ROOT" ]; then + echo "Baseline cache exists but is incomplete: $BASELINE_ROOT" >&2 + echo "Remove it or set MODELARRAYIO_BASELINE_ROOT to an empty path." >&2 + exit 2 + fi + + baseline_parent=$(dirname -- "$BASELINE_ROOT") + mkdir -p "$baseline_parent" + baseline_stage=$(mktemp -d "$baseline_parent/.modelarrayio-baseline-stage.XXXXXX") + trap 'rm -rf "$baseline_stage"' EXIT HUP INT TERM + git -C "$REPOSITORY_ROOT" archive "$BASELINE_REVISION" | tar -x -C "$baseline_stage" + mv "$baseline_stage" "$BASELINE_ROOT" + trap - EXIT HUP INT TERM +} + +prepare_baseline + +if [ "${1:-}" = "--prepare-only" ]; then + printf '%s\n' "$BASELINE_ROOT" + exit 0 +fi + +if [ -n "${MODELARRAYIO_PYTHON:-}" ]; then + PYTHON_EXECUTABLE=$MODELARRAYIO_PYTHON +elif [ -x "$REPOSITORY_ROOT/.venv/bin/python" ]; then + PYTHON_EXECUTABLE=$REPOSITORY_ROOT/.venv/bin/python +else + PYTHON_EXECUTABLE=$(command -v python3) +fi + +export PYTHONPATH="$BASELINE_ROOT/src${PYTHONPATH:+:$PYTHONPATH}" +exec "$PYTHON_EXECUTABLE" -c \ + 'import sys; from modelarrayio.cli.main import main; raise SystemExit(main(sys.argv[1:]))' \ + "$@" diff --git a/benchmarks/to_modelarray/compare_revisions.py b/benchmarks/to_modelarray/compare_revisions.py new file mode 100644 index 0000000..b7d16b4 --- /dev/null +++ b/benchmarks/to_modelarray/compare_revisions.py @@ -0,0 +1,447 @@ +"""Compare original and optimized ``to-modelarray`` executables on identical inputs.""" + +from __future__ import annotations + +import argparse +import csv +import functools +import json +import os +import platform +import resource +import shlex +import shutil +import subprocess +import tempfile +import time +from datetime import UTC, datetime +from pathlib import Path + +import harness_common as common +import output_compare +import source_identity + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--baseline-executable', type=Path, required=True) + parser.add_argument('--optimized-executable', type=Path, required=True) + parser.add_argument('--cohort', type=Path, required=True) + parser.add_argument('--mask', type=Path, help='Group mask for NIfTI cohorts.') + parser.add_argument('--index-file', type=Path, help='Index MIF for fixel cohorts.') + parser.add_argument('--directions-file', type=Path, help='Directions MIF for fixel cohorts.') + parser.add_argument('--results-dir', type=Path, required=True) + parser.add_argument( + '--scalar-columns', + nargs='+', + help='Wide-format NIfTI scalar columns; omit for long-format CIFTI/MIF cohorts.', + ) + parser.add_argument( + '--backends', nargs='+', choices=['hdf5', 'tiledb'], default=['hdf5', 'tiledb'] + ) + parser.add_argument('--workers', type=int, default=1) + parser.add_argument('--repetitions', type=int, default=3) + parser.add_argument('--warmups', type=int, default=1) + parser.add_argument( + '--target-chunk-mb', + type=float, + default=2.0, + help='Chunk/tile target passed to both executables (production CIFTI used 32).', + ) + parser.add_argument('--r-gate-script', type=Path, required=True) + parser.add_argument( + '--gate-command', + help=( + 'Optional gate command template with {output} and {backend} fields, used instead ' + 'of "Rscript " (for example a dispatcher that gates TileDB in ' + 'Python where the R tiledb package is unavailable).' + ), + ) + parser.add_argument('--baseline-revision', default='HEAD') + parser.add_argument('--optimized-revision', default='working-tree') + parser.add_argument( + '--verify-only', + action='store_true', + help='Reuse completed case directories and only run correctness checks.', + ) + return parser + + +def _source_identity(repo: Path) -> dict[str, object]: + """Pin the exact (possibly uncommitted or rsynced) source tree that ran.""" + identity = source_identity.source_identity(repo) + return { + 'source_commit': identity['commit'], + 'source_dirty': identity['dirty'], + 'source_tree_sha256': identity['source_tree_sha256'], + } + + +def _case_order(repetitions: int) -> list[tuple[str, int]]: + """Counterbalance adjacent runs while retaining stable repetition identifiers.""" + result = [] + for repetition in range(1, repetitions + 1): + implementations = ( + ('baseline', 'optimized') if repetition % 2 else ('optimized', 'baseline') + ) + result.extend((implementation, repetition) for implementation in implementations) + return result + + +def _command( + executable: Path, + implementation: str, + backend: str, + cohort: Path, + output: Path, + inputs: dict[str, object], + workers: int, +) -> list[str]: + compression, level = ('gzip', 4) if backend == 'hdf5' else ('zstd', 5) + modality = inputs['modality'] + split = modality == 'nifti' + command = [ + str(executable), + 'to-modelarray', + '--cohort-file', + str(cohort), + '--output', + str(output), + '--backend', + backend, + '--workers', + str(workers), + '--dtype', + 'float32', + '--compression', + compression, + '--compression-level', + str(level), + '--target-chunk-mb', + str(inputs.get('target_chunk_mb', 2.0)), + '--split-files' if split else '--no-split-files', + '--log-level', + 'WARNING', + ] + if modality == 'nifti': + command.extend(['--mask', str(inputs['mask'])]) + command.extend(['--scalar-columns', *inputs['scalar_columns']]) + elif modality == 'mif': + command.extend(['--index-file', str(inputs['index_file'])]) + command.extend(['--directions-file', str(inputs['directions_file'])]) + if implementation == 'baseline': + command.extend(['--s3-workers', str(workers)]) + return command + + +def _expected_outputs(case_dir: Path, backend: str, inputs: dict[str, object]) -> list[Path]: + if inputs['modality'] == 'nifti': + return [ + common.output_path(case_dir, backend, scalar) for scalar in inputs['scalar_columns'] + ] + return [common.output_path(case_dir, backend)] + + +def _run_gate( + rscript: str | None, + script: Path, + output: Path, + backend: str, + log, + template: str | None = None, +) -> tuple[bool, float]: + if template: + command = shlex.split(template.format(output=str(output), backend=backend)) + else: + command = [str(rscript), str(script), str(output), backend] + started = time.perf_counter() + try: + completed = subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=False) + except (OSError, ValueError) as exc: + log.write(f'Could not execute gate command {command!r}: {exc}\n') + return False, time.perf_counter() - started + return completed.returncode == 0, time.perf_counter() - started + + +def _run_case( + *, + implementation: str, + executable: Path, + backend: str, + repetition: int, + case_dir: Path, + cohort: Path, + inputs: dict[str, object], + workers: int, + rscript: str | None, + r_gate_script: Path, + gate_template: str | None = None, +) -> dict[str, object]: + case_dir.mkdir(parents=True, exist_ok=False) + requested_output = common.output_path(case_dir, backend) + command = _command( + executable, + implementation, + backend, + cohort, + requested_output, + inputs, + workers, + ) + stdout_path = case_dir / 'stdout.log' + stderr_path = case_dir / 'stderr.log' + before = resource.getrusage(resource.RUSAGE_CHILDREN) + started = time.perf_counter() + environment = os.environ.copy() + environment.setdefault('MODELARRAYIO_S3_ANON', '1') + with ( + stdout_path.open('w', encoding='utf-8') as stdout, + stderr_path.open('w', encoding='utf-8') as stderr, + ): + # Resolve relative cohort paths the same way run_benchmarks.py and the + # converter's own convention do: from the cohort file's directory. + completed = subprocess.run( + command, + cwd=cohort.parent, + env=environment, + stdout=stdout, + stderr=stderr, + check=False, + ) + wall_seconds = time.perf_counter() - started + after = resource.getrusage(resource.RUSAGE_CHILDREN) + outputs = _expected_outputs(case_dir, backend, inputs) + outputs_present = all(path.exists() for path in outputs) + gate_passed = False + r_read_seconds = 0.0 + gate_path = case_dir / 'r_gate.log' + with gate_path.open('w', encoding='utf-8') as gate_log: + if completed.returncode == 0 and outputs_present: + gate_results = [ + _run_gate(rscript, r_gate_script, output, backend, gate_log, gate_template) + for output in outputs + ] + gate_passed = all(passed for passed, _ in gate_results) + r_read_seconds = sum(seconds for _, seconds in gate_results) + output_bytes = 0 + output_files = 0 + if outputs_present: + for output in outputs: + size, count = common.output_stats(output) + output_bytes += size + output_files += count + status = ( + 'passed' if completed.returncode == 0 and outputs_present and gate_passed else 'failed' + ) + return { + 'implementation': implementation, + 'backend': backend, + 'repetition': repetition, + 'status': status, + 'returncode': completed.returncode, + 'wall_seconds': wall_seconds, + 'child_user_cpu_seconds': after.ru_utime - before.ru_utime, + 'child_system_cpu_seconds': after.ru_stime - before.ru_stime, + 'r_gate': 'passed' if gate_passed else 'failed', + 'r_read_seconds': r_read_seconds, + 'output_bytes': output_bytes, + 'output_file_count': output_files, + 'output_paths': json.dumps([str(path) for path in outputs]), + 'stdout_log': str(stdout_path), + 'stderr_log': str(stderr_path), + 'r_gate_log': str(gate_path), + 'command': shlex.join(command), + } + + +def _compare_outputs( + reference: Path, candidate: Path, backend: str, inputs: dict[str, object] +) -> dict[str, object]: + comparisons = [] + split = inputs['modality'] == 'nifti' + for scalar in inputs['scalar_names']: + output_scalar = scalar if split else None + comparisons.append( + output_compare.compare_outputs( + common.output_path(reference, backend, output_scalar), + common.output_path(candidate, backend, output_scalar), + scalar, + ) + ) + result = { + key: all(bool(comparison[key]) for comparison in comparisons) + for key in ('shape_equal', 'dtype_equal', 'sources_equal', 'nan_equal', 'values_equal') + } + result['max_abs_difference'] = max( + float(comparison['max_abs_difference']) for comparison in comparisons + ) + return result + + +def main() -> int: + args = _parser().parse_args() + if args.repetitions < 1 or args.warmups < 0 or args.workers < 1: + raise ValueError('repetitions/workers must be positive and warmups must be non-negative') + baseline = args.baseline_executable.expanduser().resolve() + optimized = args.optimized_executable.expanduser().resolve() + cohort = args.cohort.expanduser().resolve() + gate = args.r_gate_script.expanduser().resolve() + rscript = shutil.which('Rscript') + if rscript is None and not args.gate_command: + raise FileNotFoundError('Rscript was not found on PATH (or pass --gate-command)') + + cohort_rows = common.read_cohort(cohort) + scalar_columns = args.scalar_columns or common.infer_scalar_columns(cohort_rows) + modality = common.source_modality(common.scalar_source_paths(cohort_rows, scalar_columns)) + scalar_names = common.scalar_names(cohort_rows, scalar_columns) + mask = args.mask.expanduser().resolve() if args.mask is not None else None + index_file = args.index_file.expanduser().resolve() if args.index_file is not None else None + directions_file = ( + args.directions_file.expanduser().resolve() if args.directions_file is not None else None + ) + if modality == 'nifti' and mask is None: + raise ValueError('--mask is required for a NIfTI cohort.') + if modality == 'mif' and (index_file is None or directions_file is None): + raise ValueError('--index-file and --directions-file are required for a MIF cohort.') + inputs: dict[str, object] = { + 'modality': modality, + 'mask': mask, + 'index_file': index_file, + 'directions_file': directions_file, + 'scalar_columns': scalar_columns, + 'scalar_names': scalar_names, + 'target_chunk_mb': args.target_chunk_mb, + } + required = [baseline, optimized, cohort, gate] + required += [path for path in (mask, index_file, directions_file) if path is not None] + for path in required: + if not path.exists(): + raise FileNotFoundError(path) + run_dir = args.results_dir.expanduser().resolve() + executables = {'baseline': baseline, 'optimized': optimized} + if args.verify_only: + results_path = run_dir / 'comparison_results.csv' + if not results_path.is_file(): + raise FileNotFoundError(results_path) + with results_path.open(newline='', encoding='utf-8') as stream: + rows = list(csv.DictReader(stream)) + case_dirs = { + (implementation, backend, repetition): run_dir + / f'{implementation}-{backend}-rep-{repetition:02d}' + for backend in args.backends + for implementation in ('baseline', 'optimized') + for repetition in range(1, args.repetitions + 1) + } + else: + run_dir.mkdir(parents=True, exist_ok=False) + # Warm-up and measured runs share every setting except the case itself. + run_case = functools.partial( + _run_case, + cohort=cohort, + inputs=inputs, + workers=args.workers, + rscript=rscript, + r_gate_script=gate, + gate_template=args.gate_command, + ) + for backend in args.backends: + for implementation in ('baseline', 'optimized'): + for warmup in range(1, args.warmups + 1): + with tempfile.TemporaryDirectory( + prefix=f'warmup-{implementation}-{backend}-{warmup:02d}-', dir=run_dir + ) as temporary: + temporary_path = Path(temporary) + row = run_case( + implementation=implementation, + executable=executables[implementation], + backend=backend, + repetition=0, + case_dir=temporary_path / 'case', + ) + if row['status'] != 'passed': + raise RuntimeError( + f'Warm-up failed for {implementation}/{backend}: ' + f'{row["stderr_log"]}' + ) + + rows = [] + case_dirs: dict[tuple[str, str, int], Path] = {} + for backend in args.backends: + for implementation, repetition in _case_order(args.repetitions): + case_dir = run_dir / f'{implementation}-{backend}-rep-{repetition:02d}' + row = run_case( + implementation=implementation, + executable=executables[implementation], + backend=backend, + repetition=repetition, + case_dir=case_dir, + ) + rows.append(row) + case_dirs[(implementation, backend, repetition)] = case_dir + common.write_csv(run_dir / 'comparison_results.csv', rows) + print( + f'{implementation}/{backend}/rep-{repetition:02d}: ' + f'status={row["status"]} wall={row["wall_seconds"]:.3f}s', + flush=True, + ) + if row['status'] != 'passed': + return 1 + + correctness_rows = [] + for backend in args.backends: + reference = case_dirs[('baseline', backend, 1)] + for repetition in range(1, args.repetitions + 1): + candidate = case_dirs[('optimized', backend, repetition)] + correctness_rows.append( + { + 'backend': backend, + 'optimized_repetition': repetition, + **_compare_outputs(reference, candidate, backend, inputs), + } + ) + common.write_csv(run_dir / 'correctness_results.csv', correctness_rows) + + metadata = { + 'created_at': datetime.now(UTC).isoformat(), + 'platform': platform.platform(), + 'python': platform.python_version(), + 'baseline_executable': str(baseline), + 'optimized_executable': str(optimized), + 'baseline_revision': args.baseline_revision, + 'optimized_revision': args.optimized_revision, + 'cohort': str(cohort), + 'cohort_sha256': common.sha256_file(cohort), + 'modality': modality, + 'mask': str(mask) if mask is not None else None, + 'mask_sha256': common.sha256_file(mask) if mask is not None else None, + 'index_file': str(index_file) if index_file is not None else None, + 'directions_file': str(directions_file) if directions_file is not None else None, + 'scalar_columns': scalar_columns, + 'scalar_names': scalar_names, + 'backends': args.backends, + 'workers': args.workers, + 'target_chunk_mb': args.target_chunk_mb, + 'gate_command': args.gate_command, + 'repetitions': args.repetitions, + 'warmups': args.warmups, + 'versions': common.package_versions(), + **_source_identity(Path(__file__).resolve().parents[2]), + 'case_order': _case_order(args.repetitions), + 'comparison_results': 'comparison_results.csv', + 'correctness_results': 'correctness_results.csv', + } + (run_dir / 'comparison_meta.json').write_text( + json.dumps(metadata, indent=2, sort_keys=True), encoding='utf-8' + ) + correctness_passed = all( + all( + bool(row[key]) + for key in ('shape_equal', 'dtype_equal', 'sources_equal', 'nan_equal', 'values_equal') + ) + for row in correctness_rows + ) + return int(not correctness_passed) + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/benchmarks/to_modelarray/compare_to_reference.py b/benchmarks/to_modelarray/compare_to_reference.py new file mode 100644 index 0000000..683e4e0 --- /dev/null +++ b/benchmarks/to_modelarray/compare_to_reference.py @@ -0,0 +1,195 @@ +"""Bitwise-compare a converter output against a reference TileDB conversion. + +On CUBIC the original converter's TileDB output of the production CIFTI cohort +(``rbc_freesurferpost.tdb``, written by the stripe writer with one fragment per +column stripe) is the ground truth for the scalars it completed. This script +checks an optimized output — a per-scalar ``.h5`` file or a ``.tdb`` root — +against that reference scalar by scalar, without ever holding an 11 GB matrix +twice: it streams column blocks (old fragments are column stripes, so each +block touches one or two of them), aligns rows by ``column_names`` rather than +by position, and reports shape/dtype/NaN-map/value equality plus the maximum +absolute difference. It also flags reference scalars whose non-empty domain does +not cover the full matrix (the original run stalled part-way through a scalar). + +Examples:: + + # Which reference scalars are complete? + python benchmarks/to_modelarray/compare_to_reference.py --reference REF.tdb --list + + # Check an optimized output for one scalar (rows may be a subset). + python benchmarks/to_modelarray/compare_to_reference.py --reference REF.tdb \ + --candidate out/area_modelarray.h5 --scalar area --output-csv check.csv +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np + +import harness_common as common +import output_compare + +FILELOCK_CONFIG = {'vfs.file.enable_filelocks': 'false'} + + +def _tiledb(): + import tiledb + + return tiledb + + +def _reference_names(root: Path, scalar: str, array) -> list[str]: + names = output_compare.names_from_meta(array) + if names: + return names + tiledb = _tiledb() + column_uri = root / 'scalars' / scalar / 'column_names' + if tiledb.object_type(str(column_uri)) == 'array': + with tiledb.open(str(column_uri), 'r') as column_array: + values = np.asarray(column_array[:]['values']).reshape(-1) + return [v.decode() if isinstance(v, bytes) else str(v) for v in values] + return [] + + +def reference_completeness(root: Path) -> list[dict[str, object]]: + """Report every reference scalar's shape and whether its domain is fully written.""" + tiledb = _tiledb() + rows = [] + with tiledb.scope_ctx(tiledb.Config(FILELOCK_CONFIG)): + for scalar_dir in sorted(p for p in (root / 'scalars').glob('*') if p.is_dir()): + uri = scalar_dir / 'values' + if tiledb.object_type(str(uri)) != 'array': + continue + with tiledb.open(str(uri), 'r') as array: + n_rows, n_cols = array.shape + nonempty = array.nonempty_domain() + complete = nonempty is not None and all( + int(low) == 0 and int(high) == extent - 1 + for (low, high), extent in zip(nonempty, (n_rows, n_cols), strict=True) + ) + rows.append( + { + 'scalar': scalar_dir.name, + 'rows': n_rows, + 'columns': n_cols, + 'dtype': str(array.dtype), + 'complete': complete, + 'nonempty_domain': str(nonempty), + } + ) + return rows + + +def compare_scalar( + reference_root: Path, + candidate_path: Path, + scalar: str, + *, + block_columns: int = 4096, +) -> dict[str, object]: + """Compare one scalar; rows are aligned by column_names, columns must match.""" + tiledb = _tiledb() + result: dict[str, object] = {'scalar': scalar, 'candidate': str(candidate_path)} + reference_uri = reference_root / 'scalars' / scalar / 'values' + with tiledb.scope_ctx(tiledb.Config(FILELOCK_CONFIG)): + if tiledb.object_type(str(reference_uri)) != 'array': + # The original run completed only some scalars; say so instead of crashing. + result.update( + { + 'reference_complete': False, + 'reference_missing': True, + 'values_equal': False, + 'nan_equal': False, + 'max_abs_difference': float('nan'), + 'passed': False, + } + ) + return result + result['reference_missing'] = False + reference = tiledb.open(str(reference_uri), 'r') + candidate = output_compare.OutputReader(candidate_path, scalar) + try: + ref_rows, ref_cols = (int(v) for v in reference.shape) + ref_names = _reference_names(reference_root, scalar, reference) + nonempty = reference.nonempty_domain() + result['reference_complete'] = nonempty is not None and all( + int(low) == 0 and int(high) == extent - 1 + for (low, high), extent in zip(nonempty, (ref_rows, ref_cols), strict=True) + ) + result['columns_equal'] = candidate.shape[1] == ref_cols + result['dtype_equal'] = candidate.dtype == np.dtype(reference.dtype) + index_of = {name: i for i, name in enumerate(ref_names)} + missing = [name for name in candidate.names if name not in index_of] + result['rows_compared'] = len(candidate.names) - len(missing) + result['rows_missing_in_reference'] = len(missing) + aligned = bool(candidate.names) and not missing + result['sources_aligned'] = aligned + values_equal = aligned and result['columns_equal'] and result['dtype_equal'] + nan_equal = values_equal + max_abs = 0.0 + if values_equal: + # Reference rows are re-ordered to the candidate's column_names. + ref_index = np.asarray([index_of[name] for name in candidate.names]) + + def reference_block(start: int, stop: int) -> np.ndarray: + return np.asarray(reference[:, start:stop]['values'])[ref_index] + + nan_equal, values_equal, max_abs = output_compare.compare_blocks( + reference_block, candidate.block, ref_cols, block_columns + ) + result['nan_equal'] = nan_equal + result['values_equal'] = values_equal + result['max_abs_difference'] = max_abs + result['passed'] = bool(values_equal and nan_equal and result['reference_complete']) + finally: + candidate.close() + reference.close() + return result + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--reference', type=Path, required=True, help='Reference .tdb root') + parser.add_argument('--list', action='store_true', help='List reference scalar completeness') + parser.add_argument('--candidate', type=Path, help='Optimized .h5 file or .tdb root') + parser.add_argument('--scalar', action='append', default=[], dest='scalars') + parser.add_argument('--block-columns', type=int, default=4096) + parser.add_argument('--output-csv', type=Path) + return parser + + +def main() -> int: + args = _parser().parse_args() + reference = args.reference.expanduser().resolve() + if args.list: + rows = reference_completeness(reference) + for row in rows: + print(json.dumps(row, sort_keys=True)) + if args.output_csv and rows: + common.write_csv(args.output_csv, rows) + return 0 + if args.candidate is None or not args.scalars: + print('--candidate and at least one --scalar are required', file=sys.stderr) + return 2 + results = [ + compare_scalar( + reference, + args.candidate.expanduser().resolve(), + scalar, + block_columns=args.block_columns, + ) + for scalar in args.scalars + ] + for row in results: + print(json.dumps(row, sort_keys=True)) + if args.output_csv: + common.write_csv(args.output_csv, results) + return int(not all(row['passed'] for row in results)) + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/benchmarks/to_modelarray/comparison_runs.py b/benchmarks/to_modelarray/comparison_runs.py new file mode 100644 index 0000000..28aeef8 --- /dev/null +++ b/benchmarks/to_modelarray/comparison_runs.py @@ -0,0 +1,166 @@ +"""Read ``compare_revisions.py`` run directories into per-modality/backend groups. + +Shared by ``create_comparison_report.py`` (the findings PDF) and ``summarize_comparison.py`` +(the pinned summary CSV), so both aggregate a comparison the same way. +""" + +from __future__ import annotations + +import csv +import json +import statistics +from dataclasses import dataclass, field +from pathlib import Path + +import harness_common as common +from output_compare import OutputReader + +CORRECTNESS_KEYS = ('shape_equal', 'dtype_equal', 'sources_equal', 'nan_equal', 'values_equal') + + +@dataclass +class Group: + """Aggregated original-vs-optimized medians for one modality/backend.""" + + modality: str + backend: str + n_subjects: int + n_elements: int + n_scalars: int + baseline_revision: str = '' + baseline_wall: list[float] = field(default_factory=list) + optimized_wall: list[float] = field(default_factory=list) + baseline_r: list[float] = field(default_factory=list) + optimized_r: list[float] = field(default_factory=list) + baseline_bytes: int = 0 + optimized_bytes: int = 0 + baseline_files: int = 0 + optimized_files: int = 0 + values_equal: bool = True + shape_equal: bool = True + nan_equal: bool = True + sources_equal: bool = True + dtype_equal: bool = True + max_abs_difference: float = 0.0 + + @staticmethod + def _median(values: list[float]) -> float: + return statistics.median(values) if values else 0.0 + + @property + def baseline_median(self) -> float: + return self._median(self.baseline_wall) + + @property + def optimized_median(self) -> float: + return self._median(self.optimized_wall) + + @property + def baseline_r_median(self) -> float: + return self._median(self.baseline_r) + + @property + def optimized_r_median(self) -> float: + return self._median(self.optimized_r) + + @property + def speedup(self) -> float: + return self.baseline_median / self.optimized_median if self.optimized_median else 0.0 + + @property + def reduction(self) -> float: + base = self.baseline_median + return 100.0 * (1.0 - self.optimized_median / base) if base else 0.0 + + @property + def correctness_ok(self) -> bool: + return all( + ( + self.shape_equal, + self.dtype_equal, + self.sources_equal, + self.nan_equal, + self.values_equal, + ) + ) + + +def to_float(value: str | None) -> float | None: + if value in (None, '', 'None'): + return None + try: + return float(value) + except ValueError: + return None + + +def as_bool(value: object) -> bool: + return str(value).strip().lower() in {'1', 'true', 'yes'} + + +def output_shape(run_dir: Path, backend: str, scalar_names: list[str], modality: str) -> tuple: + """Read (n_elements, n_subjects) from a representative optimized output.""" + scalar = scalar_names[0] + split = modality == 'nifti' # wide NIfTI cohorts split into one file per scalar + output = common.output_path( + run_dir / f'optimized-{backend}-rep-01', backend, scalar if split else None + ) + with OutputReader(output, scalar) as reader: + return reader.shape[1], reader.shape[0] + + +def load_groups(run_dirs: list[Path]) -> dict[str, list[Group]]: + """Aggregate every comparison run directory into per-modality groups.""" + groups: dict[tuple[str, str], Group] = {} + for run_dir in run_dirs: + meta = json.loads((run_dir / 'comparison_meta.json').read_text(encoding='utf-8')) + modality = meta['modality'] + scalar_names = meta.get('scalar_names') or meta.get('scalar_columns') or [] + n_scalars = len(scalar_names) + with (run_dir / 'comparison_results.csv').open(newline='', encoding='utf-8') as stream: + comparison_rows = list(csv.DictReader(stream)) + correctness_path = run_dir / 'correctness_results.csv' + correctness_rows = [] + if correctness_path.exists(): + with correctness_path.open(newline='', encoding='utf-8') as stream: + correctness_rows = list(csv.DictReader(stream)) + backends = sorted({row['backend'] for row in comparison_rows}) + for backend in backends: + n_elements, n_subjects = output_shape(run_dir, backend, scalar_names, modality) + group = Group( + modality=modality, + backend=backend, + n_subjects=n_subjects, + n_elements=n_elements, + n_scalars=n_scalars, + baseline_revision=meta.get('baseline_revision') or '', + ) + for row in comparison_rows: + if row['backend'] != backend: + continue + wall = to_float(row.get('wall_seconds')) + r_read = to_float(row.get('r_read_seconds')) + if row['implementation'] == 'baseline': + group.baseline_wall.append(wall) + group.baseline_r.append(r_read) + group.baseline_bytes = int(row.get('output_bytes') or 0) + group.baseline_files = int(row.get('output_file_count') or 0) + else: + group.optimized_wall.append(wall) + group.optimized_r.append(r_read) + group.optimized_bytes = int(row.get('output_bytes') or 0) + group.optimized_files = int(row.get('output_file_count') or 0) + backend_correctness = [row for row in correctness_rows if row['backend'] == backend] + for row in backend_correctness: + for key in CORRECTNESS_KEYS: + setattr(group, key, getattr(group, key) and as_bool(row.get(key))) + group.max_abs_difference = max( + group.max_abs_difference, to_float(row.get('max_abs_difference')) or 0.0 + ) + groups[(modality, backend)] = group + grouped: dict[str, list[Group]] = {} + for group in groups.values(): + grouped.setdefault(group.modality, []).append(group) + for modality_groups in grouped.values(): + modality_groups.sort(key=lambda g: g.backend) + return grouped diff --git a/benchmarks/to_modelarray/create_comparison_report.py b/benchmarks/to_modelarray/create_comparison_report.py new file mode 100644 index 0000000..2324fbb --- /dev/null +++ b/benchmarks/to_modelarray/create_comparison_report.py @@ -0,0 +1,862 @@ +"""Render the original-versus-optimized ``to-modelarray`` benchmark PDF. + +Reads one or more ``compare_revisions.py`` run directories (each holding +``comparison_results.csv``, ``correctness_results.csv`` and +``comparison_meta.json``) and renders a findings report that compares the +original committed code (a pinned baseline revision) against the optimized +working tree for every modality (NIfTI, CIFTI, MIF) on both backends +(HDF5, TileDB): runtime speedup, downstream R-read time, output size and +bitwise correctness. + +Example:: + + python benchmarks/to_modelarray/create_comparison_report.py \ + --run-dir benchmarks/to_modelarray/results/original-vs-optimized-9ab13dd-nifti \ + --run-dir benchmarks/to_modelarray/results/original-vs-optimized-9ab13dd-mif \ + --run-dir benchmarks/to_modelarray/results/original-vs-optimized-9ab13dd-cifti \ + --output benchmarks/to_modelarray/results/to-modelarray-cross-modality-benchmark.pdf +""" + +from __future__ import annotations + +import argparse +import csv +import json +import statistics +from pathlib import Path +from xml.sax.saxutils import escape + +from reportlab.lib import colors +from reportlab.lib.colors import HexColor +from reportlab.lib.enums import TA_CENTER, TA_LEFT +from reportlab.lib.pagesizes import letter +from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet +from reportlab.lib.units import inch +from reportlab.platypus import ( + Flowable, + PageBreak, + Paragraph, + SimpleDocTemplate, + Spacer, + Table, + TableStyle, +) + +from comparison_runs import Group, load_groups + +NAVY = HexColor('#172A3A') +TEAL = HexColor('#0F766E') +TEAL_LIGHT = HexColor('#DDF3F0') +CORAL = HexColor('#E76F51') +CORAL_LIGHT = HexColor('#FCE8E2') +SLATE = HexColor('#48606F') +PALE = HexColor('#F3F6F8') +GRID = HexColor('#D3DDE3') +INK = HexColor('#172A3A') +WHITE = colors.white + +MODALITY_LABEL = {'nifti': 'NIfTI (voxel)', 'cifti': 'CIFTI (surface)', 'mif': 'MIF (fixel)'} +MODALITY_ORDER = ['nifti', 'cifti', 'mif'] +BACKEND_LABEL = {'hdf5': 'HDF5 gzip-4', 'tiledb': 'TileDB zstd-5'} + + +MODALITY_NAME = {'nifti': 'NIfTI', 'cifti': 'CIFTI', 'mif': 'MIF'} +BACKEND_NAME = {'hdf5': 'HDF5', 'tiledb': 'TileDB'} +NUMBER_WORDS = ( + 'no', + 'one', + 'two', + 'three', + 'four', + 'five', + 'six', + 'seven', + 'eight', + 'nine', + 'ten', +) +RUN_SETTINGS = ('repetitions', 'warmups', 'workers', 'target_chunk_mb') +# What explains a configuration's gain, for the executive finding. +GAIN_REASON = { + 'cifti': 'where the raw-read fast path removes the per-file XML header parse', + 'hdf5': 'where chunks are now compressed on a thread pool instead of one writer thread', + 'tiledb': 'where each scalar is now written as one dense fragment instead of one per stripe', +} + + +def _count(value: str, noun: str) -> str: + """Spell a run setting for prose (``'3'`` -> ``'three s'``); mixed values stay digits.""" + if not value.isdigit(): + return f'{value} {noun}s' + number = int(value) + word = NUMBER_WORDS[number] if number < len(NUMBER_WORDS) else value + return f'{word} {noun}' if number == 1 else f'{word} {noun}s' + + +def _plural(number: int, noun: str) -> str: + return f'{number:,} {noun}' if number == 1 else f'{number:,} {noun}s' + + +def _join(items: list[str]) -> str: + if len(items) <= 2: + return ' and '.join(items) + return f'{", ".join(items[:-1])}, and {items[-1]}' + + +def _scope(ordered: list[Group]) -> str: + """Name the modalities and backends actually measured, e.g. 'NIfTI and MIF on HDF5'.""" + modalities = _join(list(dict.fromkeys(MODALITY_NAME[group.modality] for group in ordered))) + backends = list(dict.fromkeys(BACKEND_NAME[group.backend] for group in ordered)) + return f'{modalities} on {"both HDF5 and TileDB" if len(backends) == 2 else backends[0]}' + + +def _largest_gain(ordered: list[Group]) -> str: + best = max(ordered, key=lambda group: group.speedup) + reason = GAIN_REASON['cifti' if best.modality == 'cifti' else best.backend] + return ( + f'The largest gain is {MODALITY_NAME[best.modality]} on {BACKEND_NAME[best.backend]} ' + f'({best.speedup:.1f}x), {reason}.' + ) + + +def _run_setting(metas: list[dict], key: str) -> str: + """One recorded run setting across every run directory; differing values are joined by '/'.""" + values = {f'{float(meta[key]):g}' for meta in metas if meta.get(key) is not None} + return '/'.join(sorted(values, key=float)) or '?' + + +def _cohort_sources(values: list[str]) -> dict[str, str]: + sources = {} + for value in values: + modality, separator, text = value.partition('=') + if not separator or modality not in MODALITY_LABEL: + raise SystemExit( + f'--cohort-source expects MODALITY=TEXT with MODALITY one of ' + f'{sorted(MODALITY_LABEL)}, got {value!r}' + ) + sources[modality] = text + return sources + + +def _median_cells(group: Group) -> list[str]: + """Original median, optimized median and speedup, as every results table shows them.""" + return [ + f'{group.baseline_median:.2f} s', + f'{group.optimized_median:.2f} s', + f'{group.speedup:.2f}x', + ] + + +def _equal_cell(group: Group) -> str: + return 'yes' if group.correctness_ok else 'NO' + + +def _ordered_groups(grouped: dict[str, list[Group]]) -> list[Group]: + ordered: list[Group] = [] + for modality in MODALITY_ORDER: + for backend in ('hdf5', 'tiledb'): + ordered.extend( + group for group in grouped.get(modality, []) if group.backend == backend + ) + return ordered + + +def _styles() -> dict[str, ParagraphStyle]: + sample = getSampleStyleSheet() + return { + 'title': ParagraphStyle( + 'Title', + parent=sample['Title'], + fontName='Helvetica-Bold', + fontSize=22, + leading=25, + textColor=NAVY, + alignment=TA_LEFT, + spaceAfter=4, + ), + 'subtitle': ParagraphStyle( + 'Subtitle', + parent=sample['Normal'], + fontName='Helvetica', + fontSize=10, + leading=14, + textColor=SLATE, + spaceAfter=12, + ), + 'h1': ParagraphStyle( + 'Heading1', + parent=sample['Heading1'], + fontName='Helvetica-Bold', + fontSize=15, + leading=18, + textColor=NAVY, + spaceBefore=2, + spaceAfter=7, + ), + 'h2': ParagraphStyle( + 'Heading2', + parent=sample['Heading2'], + fontName='Helvetica-Bold', + fontSize=11, + leading=14, + textColor=TEAL, + spaceBefore=4, + spaceAfter=5, + ), + 'body': ParagraphStyle( + 'Body', + parent=sample['BodyText'], + fontName='Helvetica', + fontSize=8.5, + leading=12, + textColor=INK, + spaceAfter=5, + ), + 'small': ParagraphStyle( + 'Small', + parent=sample['BodyText'], + fontName='Helvetica', + fontSize=7.4, + leading=10, + textColor=SLATE, + ), + 'card_number': ParagraphStyle( + 'CardNumber', + parent=sample['Normal'], + fontName='Helvetica-Bold', + fontSize=16, + leading=17, + textColor=NAVY, + alignment=TA_CENTER, + ), + 'card_label': ParagraphStyle( + 'CardLabel', + parent=sample['Normal'], + fontName='Helvetica', + fontSize=6.7, + leading=8.4, + textColor=SLATE, + alignment=TA_CENTER, + ), + 'table_header': ParagraphStyle( + 'TableHeader', + parent=sample['Normal'], + fontName='Helvetica-Bold', + fontSize=7.2, + leading=9, + textColor=WHITE, + ), + 'table': ParagraphStyle( + 'Table', + parent=sample['Normal'], + fontName='Helvetica', + fontSize=7.2, + leading=9.1, + textColor=INK, + ), + } + + +def _paragraph(value: object, style: ParagraphStyle) -> Paragraph: + return Paragraph(str(value), style) + + +def _table(rows, widths, styles, *, header=True) -> Table: + prepared = [] + for row_index, row in enumerate(rows): + style = styles['table_header'] if header and row_index == 0 else styles['table'] + prepared.append([_paragraph(value, style) for value in row]) + table = Table(prepared, colWidths=widths, repeatRows=1 if header else 0, hAlign='LEFT') + commands = [ + ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'), + ('LEFTPADDING', (0, 0), (-1, -1), 5), + ('RIGHTPADDING', (0, 0), (-1, -1), 5), + ('TOPPADDING', (0, 0), (-1, -1), 3.5), + ('BOTTOMPADDING', (0, 0), (-1, -1), 3.5), + ('GRID', (0, 0), (-1, -1), 0.35, GRID), + ] + if header: + commands.append(('BACKGROUND', (0, 0), (-1, 0), NAVY)) + commands.append(('ROWBACKGROUNDS', (0, 1), (-1, -1), [WHITE, PALE])) + table.setStyle(TableStyle(commands)) + return table + + +def _cards(cards, styles, backgrounds) -> Table: + content = [ + [ + [_paragraph(number, styles['card_number']), _paragraph(label, styles['card_label'])] + for number, label in cards + ] + ] + width = 504 / len(cards) + table = Table(content, colWidths=[width] * len(cards), rowHeights=[54]) + style = [ + ('BOX', (0, 0), (-1, -1), 0.5, GRID), + ('INNERGRID', (0, 0), (-1, -1), 0.5, WHITE), + ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'), + ('TOPPADDING', (0, 0), (-1, -1), 6), + ('BOTTOMPADDING', (0, 0), (-1, -1), 6), + ] + for index, background in enumerate(backgrounds): + style.append(('BACKGROUND', (index, 0), (index, 0), background)) + table.setStyle(TableStyle(style)) + return table + + +def _callout(text, styles, background=TEAL_LIGHT, border=TEAL) -> Table: + table = Table([[_paragraph(text, styles['body'])]], colWidths=[504]) + table.setStyle( + TableStyle( + [ + ('BACKGROUND', (0, 0), (-1, -1), background), + ('BOX', (0, 0), (-1, -1), 0.6, border), + ('LEFTPADDING', (0, 0), (-1, -1), 10), + ('RIGHTPADDING', (0, 0), (-1, -1), 10), + ('TOPPADDING', (0, 0), (-1, -1), 7), + ('BOTTOMPADDING', (0, 0), (-1, -1), 7), + ] + ) + ) + return table + + +class SpeedupBars(Flowable): + """Grouped horizontal original-vs-optimized median wall-time bars.""" + + def __init__(self, groups: list[Group], width=504, height=250): + super().__init__() + self.groups = groups + self.width = width + self.height = height + + def draw(self) -> None: + canvas = self.canv + left = 118 + right = 44 + chart_width = self.width - left - right + maximum = max((g.baseline_median for g in self.groups), default=1.0) * 1.16 or 1.0 + top = self.height - 10 + bottom = 26 + lane = (top - bottom) / max(len(self.groups), 1) + canvas.setFont('Helvetica', 6.5) + for tick in range(6): + value = maximum * tick / 5 + x = left + chart_width * tick / 5 + canvas.setStrokeColor(GRID) + canvas.setLineWidth(0.35) + canvas.line(x, bottom - 2, x, top) + canvas.setFillColor(SLATE) + canvas.drawCentredString(x, 12, f'{value:.1f}') + canvas.setFillColor(SLATE) + canvas.drawCentredString( + left + chart_width / 2, 2, 'Median conversion wall time (seconds)' + ) + for index, group in enumerate(self.groups): + y = top - (index + 0.5) * lane + label = f'{MODALITY_LABEL[group.modality].split(" ")[0]} · {group.backend.upper()}' + canvas.setFont('Helvetica-Bold', 7) + canvas.setFillColor(NAVY) + canvas.drawRightString(left - 8, y - 2, label) + for offset, value, fill in ( + (2.4, group.baseline_median, CORAL), + (-8.4, group.optimized_median, TEAL), + ): + bar = chart_width * value / maximum + canvas.setFillColor(fill) + canvas.roundRect(left, y + offset - 2.2, max(bar, 0.4), 9, 1.5, fill=1, stroke=0) + canvas.setFillColor(NAVY) + canvas.setFont('Helvetica-Bold', 6.3) + canvas.drawString(left + max(bar, 0.4) + 3, y + offset - 0.5, f'{value:.2f}s') + canvas.setFillColor(TEAL) + canvas.setFont('Helvetica-Bold', 6.8) + canvas.drawString(left + chart_width + 6, y - 2, f'{group.speedup:.2f}x') + + +def _fmt_bytes(value: float) -> str: + if value >= 2**30: + return f'{value / 2**30:.2f} GiB' + return f'{value / 2**20:.1f} MiB' + + +def _header_footer(canvas, document, measured_at: str, baseline_revision: str, cache: str) -> None: + canvas.saveState() + canvas.setStrokeColor(GRID) + canvas.setLineWidth(0.5) + canvas.line(42, 764, 570, 764) + canvas.setFillColor(SLATE) + canvas.setFont('Helvetica', 6.8) + canvas.drawString( + 42, + 771, + f'Measured {measured_at} | Original {baseline_revision} vs optimized | {cache}', + ) + canvas.drawRightString(570, 771, f'Page {document.page}') + canvas.line(42, 28, 570, 28) + canvas.drawString(42, 18, 'ModelArrayIO original-vs-optimized runtime benchmark') + canvas.restoreState() + + +def _report_story(grouped, meta, styles) -> list[Flowable]: + ordered = _ordered_groups(grouped) + total = len(ordered) + faster = sum(group.speedup > 1.0 for group in ordered) + correct = sum(group.correctness_ok for group in ordered) + best = max((group.speedup for group in ordered), default=0.0) + max_diff = max((group.max_abs_difference for group in ordered), default=0.0) + all_good = faster == total and correct == total + scope = _scope(ordered) + settings = meta['settings'] + revision = meta['baseline_revision'] + repetitions = settings['repetitions'] + warmups = settings['warmups'] + speed = ( + f'in every one of the {total} modality x backend configurations, by up to {best:.1f}x' + if faster == total + else f'in {faster} of the {total} modality x backend configurations (best {best:.1f}x)' + ) + correctness = ( + f'while producing byte-for-byte identical output: all {correct} of {total} optimized ' + f'outputs were bitwise equal to the original (maximum absolute difference {max_diff:.1f}), ' + 'with source order and NaN positions preserved' + if correct == total + else f'but only {correct} of {total} optimized outputs were bitwise equal to the original ' + f'(maximum absolute difference {max_diff:.3g})' + ) + + cards = [ + (f'{faster}/{total}', 'configurations faster than original'), + (f'{best:.1f}x', 'best median speedup'), + (f'{correct}/{total}', 'outputs bitwise equal to original'), + (f'{max_diff:.1f}', 'max abs value difference'), + ] + card_backgrounds = [TEAL_LIGHT, TEAL_LIGHT, TEAL_LIGHT, PALE] + + summary_rows = [ + ['Modality', 'Backend', 'Original', 'Optimized', 'Speedup', 'Reduction', 'Equal'], + *( + [ + MODALITY_LABEL[group.modality], + group.backend.upper(), + *_median_cells(group), + f'{group.reduction:.0f}%', + _equal_cell(group), + ] + for group in ordered + ), + ] + + story: list[Flowable] = [ + Spacer(1, 8), + _paragraph('to-modelarray: original vs optimized', styles['title']), + _paragraph( + f'End-to-end conversion runtime of the original committed code (commit {revision}) ' + f'against the optimized working tree for {scope}, on identical inputs', + styles['subtitle'], + ), + _cards(cards, styles, card_backgrounds), + Spacer(1, 12), + _paragraph('Executive finding', styles['h1']), + _paragraph( + f'The optimized converter is faster than the original ({revision}) {speed}, ' + f'{correctness}, and every R read gate passed. {_largest_gain(ordered)}', + styles['body'], + ), + _paragraph( + f'Table 1. Medians of {_count(repetitions, "counterbalanced repetition")} after ' + f'{_count(warmups, "discarded warm-up")} per implementation and backend. float32, ' + f'gzip-4 HDF5 / zstd-5 TileDB, {settings["target_chunk_mb"]} MiB target, ' + f'{_count(settings["workers"], "loader worker")}. Converter wall time excludes the ' + 'separately timed R read gate.', + styles['small'], + ), + _table(summary_rows, [92, 74, 66, 68, 58, 62, 44], styles), + Spacer(1, 8), + SpeedupBars(ordered), + _paragraph( + 'Figure 1. Original (coral) vs optimized (teal) median wall time; the right-hand label ' + 'is the speedup. Element counts differ by modality, so bars compare each pipeline ' + 'against its own original, not across modalities.', + styles['small'], + ), + Spacer(1, 6), + _callout( + f'Decision: The optimization is a clear, correctness-preserving runtime win for ' + f'{scope}. It is safe to adopt: outputs are identical and remain readable through ' + 'ModelArray in R.' + if all_good + else f'Decision: Not ready to adopt as measured: {total - faster} of {total} ' + f'configurations were not faster and {total - correct} of {total} outputs differed ' + 'from the original. See Table 1 and the correctness gates.', + styles, + ), + PageBreak(), + Spacer(1, 8), + _paragraph('Method: controlled before-and-after comparison', styles['h1']), + _paragraph( + f'The baseline is the unmodified source at commit {revision}, exported with git archive ' + 'into an isolated tree and run under the same Python environment as the optimized ' + 'working tree. For each modality and backend the harness discards ' + f'{_count(warmups, "warm-up")} per implementation, then runs ' + f'{_count(repetitions, "counterbalanced measured repetition")} (original/optimized ' + 'order alternates), applies the R read gate, and compares every optimized output ' + 'against the original for shape, dtype, source order, NaN placement and values.', + styles['body'], + ), + _paragraph('Per-modality detail', styles['h2']), + ] + + for modality in MODALITY_ORDER: + groups = grouped.get(modality) + if not groups: + continue + first = groups[0] + story.append( + _paragraph( + f'{MODALITY_LABEL[modality]} — {_plural(first.n_subjects, "subject")} x ' + f'{_plural(first.n_scalars, "scalar")}, {_plural(first.n_elements, "element")}', + styles['h2'], + ) + ) + rows = [ + [ + 'Backend', + 'Orig median', + 'Opt median', + 'Speedup', + 'Orig R-read', + 'Opt R-read', + 'Files (orig→opt)', + 'Bitwise equal', + ], + *( + [ + group.backend.upper(), + *_median_cells(group), + f'{group.baseline_r_median:.2f} s', + f'{group.optimized_r_median:.2f} s', + f'{group.baseline_files} → {group.optimized_files}', + _equal_cell(group), + ] + for group in groups + ), + ] + story.append(_table(rows, [50, 62, 62, 52, 62, 62, 90, 64], styles)) + story.append(Spacer(1, 7)) + + story.extend( + [ + _callout( + 'What changed operationally. HDF5: chunks are byte-shuffled and compressed on ' + 'a thread pool, then handed to one writer thread, so the same standard gzip file is ' + 'written far faster. TileDB: each scalar is written as one dense fragment instead of ' + 'one fragment per spatial stripe, cutting the file count sharply. CIFTI: the loader ' + 'reads greyordinate data blocks directly and parses the large XML header only once ' + 'per scalar. All three keep the exact output schema ModelArray reads.', + styles, + background=PALE, + border=SLATE, + ), + PageBreak(), + Spacer(1, 8), + _paragraph('Correctness and reproducibility', styles['h1']), + _paragraph( + 'Correctness gates. Every optimized output was compared against the original ' + 'output for the same inputs. ' + + ( + 'All checks passed:' + if correct == total + else f'{correct} of {total} configurations passed every check:' + ), + styles['body'], + ), + _table( + [ + [ + 'Modality', + 'Backend', + 'Shape', + 'Dtype', + 'Sources', + 'NaN map', + 'Values', + 'Max abs diff', + ], + *( + [ + MODALITY_LABEL[group.modality], + group.backend.upper(), + 'eq' if group.shape_equal else 'NO', + 'eq' if group.dtype_equal else 'NO', + 'eq' if group.sources_equal else 'NO', + 'eq' if group.nan_equal else 'NO', + 'eq' if group.values_equal else 'NO', + f'{group.max_abs_difference:.1f}', + ] + for group in ordered + ), + ], + [86, 66, 44, 44, 52, 54, 50, 78], + styles, + ), + Spacer(1, 8), + _paragraph('Reproduce', styles['h2']), + _table(_cohort_table(grouped, meta), [52, 150, 130, 172], styles), + Spacer(1, 6), + _paragraph( + f'Baseline: git commit {meta["baseline_revision"]}, exported and run under the ' + f'current environment via benchmarks/to_modelarray/baselines/modelarrayio-baseline-' + f'{meta["baseline_revision"]}. Environment: {meta["platform"]}; Python ' + f'{meta["versions"]["python"]}, NumPy {meta["versions"]["numpy"]}, nibabel ' + f'{meta["versions"]["nibabel"]}, h5py {meta["versions"]["h5py"]}, TileDB-Py ' + f'{meta["versions"]["tiledb"]}. Medians of {_count(repetitions, "repetition")} after ' + f'{_count(warmups, "discarded warm-up")}. Harness: ' + 'benchmarks/to_modelarray/compare_revisions.py; cohort generators and commands: ' + 'benchmarks/to_modelarray/README.md.', + styles['small'], + ), + Spacer(1, 8), + *_scale_section(meta.get('scale_rows', []), meta.get('original_run_note', ''), styles), + _callout( + f'Bottom line: Against the original code at {revision}, the optimized ' + f'to-modelarray is faster for {scope}, with byte-for-byte identical, R-readable ' + 'output.' + if all_good + else f'Bottom line: Against the original code at {revision}, the optimized ' + f'to-modelarray is faster in {faster} of {total} configurations and bitwise equal ' + f'in {correct} of {total}.', + styles, + ), + ] + ) + return story + + +def _load_scale(path: Path | None) -> list[dict[str, str]]: + if path is None or not path.exists(): + return [] + with path.open(newline='', encoding='utf-8') as stream: + return list(csv.DictReader(stream)) + + +def _scale_section(rows: list[dict[str, str]], note: str, styles) -> list[Flowable]: + """Production-scale page: full-scalar old-vs-new, worker sweep, sharded array, build.""" + if not rows: + return [] + by_step: dict[str, list[dict[str, str]]] = {} + for row in rows: + by_step.setdefault(row['step'], []).append(row) + story: list[Flowable] = [ + PageBreak(), + Spacer(1, 8), + _paragraph('Production scale', styles['h1']), + ] + if note: + story.append(_paragraph(note, styles['body'])) + compare = by_step.get('3_full_scalar_original_vs_optimized', []) + if compare: + story.append(_paragraph('One full scalar, original vs optimized', styles['h2'])) + story.append( + _table( + [ + ['Scalar', 'Subjects', 'Backend', 'Original', 'Optimized', 'Speedup'], + *( + [ + r['scalars'], + r['subjects'], + r['backend'].upper(), + f'{float(r["baseline_median_wall_seconds"]) / 3600:.2f} h', + f'{float(r["optimized_median_wall_seconds"]) / 60:.1f} min', + f'{float(r["speedup"]):.1f}x', + ] + for r in compare + ), + ], + [90, 70, 70, 90, 90, 70], + styles, + ) + ) + sweep = by_step.get('3_full_scalar_worker_sweep', []) + if sweep: + story.append(Spacer(1, 6)) + story.append(_paragraph('Optimized worker sweep on the same scalar', styles['h2'])) + story.append( + _table( + [ + [ + 'Backend', + 'Workers', + 'Wall', + 'Files/min', + 'Peak RSS', + 'Fragments', + 'Fast-path rows', + ], + *( + [ + r['backend'].upper(), + r['workers'], + f'{float(r["optimized_median_wall_seconds"]):.1f} s', + r.get('input_files_per_minute', ''), + f'{r.get("peak_memory_gb", "")} GB', + r.get('fragment_count', '') or '—', + r.get('fast_path_rows', ''), + ] + for r in sweep + ), + ], + [64, 56, 66, 70, 70, 70, 90], + styles, + ) + ) + for step, title in ( + ('4_scalar_sharded_array', 'Scalar-sharded SLURM job array (one task per scalar)'), + ('5_production_build', 'Full production build (per-scalar HDF5 files)'), + ): + items = by_step.get(step, []) + if not items: + continue + walls = [float(r['optimized_median_wall_seconds']) for r in items] + passed = sum(r.get('passed', '').lower() == 'true' for r in items) + backends = sorted({r['backend'] for r in items}) + story.append(Spacer(1, 6)) + story.append(_paragraph(title, styles['h2'])) + story.append( + _paragraph( + f'{len(items)} conversions ({", ".join(b.upper() for b in backends)}), ' + f'{passed} passed their gates; per-scalar wall {min(walls):.0f}-{max(walls):.0f} s, ' + f'median {statistics.median(walls):.0f} s, serial-equivalent sum ' + f'{sum(walls) / 60:.0f} min.', + styles['body'], + ) + ) + return story + + +def _cohort_table(grouped: dict[str, list[Group]], meta: dict) -> list[list[str]]: + """Describe each modality's cohort from what its comparison recorded.""" + rows = [['Modality', 'Cohort', 'Cohort file', 'Source']] + for modality in MODALITY_ORDER: + groups = grouped.get(modality) + if not groups: + continue + recorded = meta['cohorts'][modality] + scalars = recorded.get('scalar_names') or [] + rows.append( + [ + MODALITY_NAME[modality], + ( + f'{_plural(groups[0].n_subjects, "subject")} x ' + f'{_plural(len(scalars), "scalar")}: {escape(", ".join(scalars))}' + ), + ( + f'{escape(Path(recorded["cohort"]).name)}
' + f'sha256 {recorded.get("cohort_sha256", "")[:12]}' + ), + escape(meta['cohort_sources'].get(modality, 'not recorded')), + ] + ) + return rows + + +def _parser() -> argparse.ArgumentParser: + repo = Path(__file__).resolve().parents[2] + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--run-dir', type=Path, action='append', required=True, dest='run_dirs') + parser.add_argument( + '--measured-at', help='Measurement date (default: date of the earliest run directory).' + ) + parser.add_argument( + '--scale-csv', + type=Path, + default=None, + help='Optional summarize_scale.py CSV; adds a production-scale section (full-scalar ' + 'old-vs-new, worker sweep, scalar-sharded array, production build).', + ) + parser.add_argument( + '--original-run-note', + default='', + help='One sentence describing the original production run for the scale section.', + ) + parser.add_argument( + '--platform', help='Hardware description (default: the platform the runs recorded).' + ) + parser.add_argument( + '--cohort-source', + action='append', + default=[], + metavar='MODALITY=TEXT', + help='Where a modality\'s cohort came from, e.g. "cifti=RBC subset, internal"; ' + 'repeat per modality. Unset modalities read "not recorded".', + ) + parser.add_argument( + '--output', + type=Path, + default=repo + / 'benchmarks' + / 'to_modelarray' + / 'results' + / 'to-modelarray-cross-modality-benchmark.pdf', + ) + return parser + + +def main() -> int: + args = _parser().parse_args() + run_dirs = [run_dir.expanduser().resolve() for run_dir in args.run_dirs] + grouped = load_groups(run_dirs) + if not grouped: + raise ValueError('No comparison run directories were loaded.') + metas = [ + json.loads((run_dir / 'comparison_meta.json').read_text(encoding='utf-8')) + for run_dir in run_dirs + ] + first_meta = metas[0] + created = sorted(recorded['created_at'] for recorded in metas if recorded.get('created_at')) + meta = { + 'baseline_revision': first_meta.get('baseline_revision', 'baseline'), + 'measured_at': args.measured_at or (created[0][:10] if created else 'date not recorded'), + 'platform': args.platform or first_meta.get('platform', 'not recorded'), + 'settings': {key: _run_setting(metas, key) for key in RUN_SETTINGS}, + # Discarded warm-up runs read every input first, so measured runs see a warm cache. + 'cache': ( + 'Warm cache' + if all(int(recorded.get('warmups') or 0) > 0 for recorded in metas) + else 'Cache not pre-warmed' + ), + 'cohorts': {recorded['modality']: recorded for recorded in metas}, + 'cohort_sources': _cohort_sources(args.cohort_source), + # The versions that ran the comparison, not the environment rendering the report. + 'versions': {name: version or 'n/a' for name, version in first_meta['versions'].items()}, + 'scale_rows': _load_scale(args.scale_csv), + 'original_run_note': args.original_run_note, + } + styles = _styles() + output = args.output.expanduser().resolve() + output.parent.mkdir(parents=True, exist_ok=True) + document = SimpleDocTemplate( + str(output), + pagesize=letter, + leftMargin=0.58 * inch, + rightMargin=0.58 * inch, + topMargin=0.48 * inch, + bottomMargin=0.48 * inch, + title='to-modelarray original vs optimized runtime benchmark', + author='ModelArrayIO benchmark harness', + subject=f'Original (committed) vs optimized to-modelarray for {_scope(_ordered_groups(grouped))}', + ) + document.build( + _report_story(grouped, meta, styles), + onFirstPage=lambda canvas, doc: _header_footer( + canvas, doc, meta['measured_at'], meta['baseline_revision'], meta['cache'] + ), + onLaterPages=lambda canvas, doc: _header_footer( + canvas, doc, meta['measured_at'], meta['baseline_revision'], meta['cache'] + ), + ) + print(output) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/benchmarks/to_modelarray/cubic/README.md b/benchmarks/to_modelarray/cubic/README.md new file mode 100644 index 0000000..9d35516 --- /dev/null +++ b/benchmarks/to_modelarray/cubic/README.md @@ -0,0 +1,126 @@ +# CUBIC validation ladder for `to-modelarray` + +Runs the optimized converter on the CUBIC cluster against the original code (commit +`9ab13dd`) and against the production CIFTI cohort, with a correctness gate at every step. +Everything writable lives under `/cbica/home/tongt/modelarray_opt`; the production data and +its old-code reference output are read in place from `/cbica/projects/pennlinc_rbc/...`. + +## Environment route + +The ModelArray container (`~/software/modelarray_confixel_0.1.5.sif`, Ubuntu 20.04, glibc 2.31) +runs *everything*: a uv-managed Python 3.12 venv created inside it installs the exact local +library versions as `manylinux_2_28` wheels (the CentOS 7 host, glibc 2.17, cannot), and its +R ModelArray/rhdf5/HDF5Array gate HDF5 outputs natively. TileDB outputs are gated in Python +(`benchmarks/to_modelarray/tiledb_read_gate.py`) because the container has no R `tiledb`. `env.sh` defines +`in_sif`, which runs a command in the container — or directly when already inside one, so the +gate dispatcher works from within a containerised job without nesting. + +## One-time setup (login node; it has internet, compute nodes may not) + +```bash +# local machine: copy the working tree (with .git, for the baseline export) — no commit needed +rsync -a --delete --exclude .venv --exclude 'benchmarks/*/data' --exclude 'benchmarks/*/results' \ + --exclude 'benchmarks/*/cohorts' --exclude .DS_Store \ + ModelArrayIO/ tongt@cubic-login.uphs.upenn.edu:/cbica/home/tongt/modelarray_opt/ModelArrayIO/ +python benchmarks/to_modelarray/source_identity.py # note source_tree_sha256 + +# local machine: the current ModelArray R source, for the HDF5 gate +rsync -a --delete --exclude .git --exclude docs ModelArray/ \ + tongt@cubic-login.uphs.upenn.edu:/cbica/home/tongt/modelarray_opt/ModelArray/ + +# cubic login node +bash /cbica/home/tongt/modelarray_opt/ModelArrayIO/benchmarks/to_modelarray/cubic/setup_env.sh +sbatch /cbica/home/tongt/modelarray_opt/ModelArrayIO/benchmarks/to_modelarray/cubic/sbatch/00_smoke.sbatch # also writes reference_scalars.csv +bash /cbica/home/tongt/modelarray_opt/ModelArrayIO/benchmarks/to_modelarray/cubic/make_scalar_cohorts.sh # after the smoke job +``` + +`setup_env.sh` prints the container's `source_tree_sha256`; it must equal the local one — this +is the rsync-fidelity gate, and every result file records it. + +## The ladder + +| Step | Job | Gate | +|---|---|---| +| 0 | `sbatch/00_smoke.sbatch` (short) | imports on a compute node, offline pytest, quick-tier fixtures through both gate branches, fixture old-vs-new bitwise | +| 2 | `sbatch/02_compare_public.sbatch` (all, exclusive) | old-vs-new on ABIDE-100, MIF walkthrough-100, RBC CIFTI 100×3: bitwise + gates; raw fast path vs nibabel both equal the deployed reference; CUBIC PDF | +| 3 | `sbatch/03_single_scalar.sbatch` (all, exclusive) | one full scalar, 8,693 subjects: old-vs-new (1 rep), worker sweep, every output bitwise vs `rbc_freesurferpost.tdb` | +| 4 | `sbatch/04_scalar_array.sbatch` `--array` | one task per complete reference scalar, per-scalar file, bitwise vs reference, no aggregation | +| 5 | `submit_production.sh` → `05_production.sbatch` `--array` | all scalars → kept per-scalar HDF5; runs only when every step-4 shard PASSED | + +`submit_ladder.sh` chains steps 0→2→3→4 with `--dependency=afterok`. + +## Pinning the results + +Step 2 writes the cross-modality summary CSV and an interim PDF without the production-scale +section. Once steps 3–5 have passed, build the production-scale CSV and the final report. +`summarize_scale.py` reads the step-3 comparison and worker sweep, the step-4 shard walls and the +step-5 production walls (it skips any step not yet run). The report opens one TileDB output per +backend, so run both commands on a compute node, not the login node: + +```bash +# cubic login node: open a shell on a compute node +srun -p short -c 2 --mem=8G -t 1:00:00 --propagate=NONE --pty bash + +source /cbica/home/tongt/modelarray_opt/ModelArrayIO/benchmarks/to_modelarray/cubic/env.sh +cd "$REPO" +in_sif python benchmarks/to_modelarray/summarize_scale.py --results-root "$OPT_ROOT/results" \ + --scalar area --output "$OPT_ROOT/results/cubic_production_scale_9ab13dd_vs_optimized.csv" +# The note is the report's hand-written "Production scale" summary; this is the text of the pinned +# 2026-09-22 report. Rewrite it from the new numbers whenever the ladder is re-run. +NOTE='The original production run (code/02_create_tiledb.sh, --workers 15) needed 112 h 23 m for its first scalar and stalled at 15 of 70 scalars after about 119 h. On the same data, the optimized converter builds one full 8,693-subject scalar in about 7 minutes at one loader worker and in 69-91 seconds at eight, and the scalar-sharded job array converted the 15 reference scalars in 21 minutes of wall-clock time, every output bitwise-equal to the deployed reference. The full production build - all 70 scalars, 608,510 files, 484 GB of per-scalar HDF5 - completed in about 22 minutes of job-array wall-clock (median 138 s per scalar), 70 of 70 gated.' +in_sif python benchmarks/to_modelarray/create_comparison_report.py \ + --run-dir "$OPT_ROOT/results/cubic-02-nifti" --run-dir "$OPT_ROOT/results/cubic-02-mif" \ + --run-dir "$OPT_ROOT/results/cubic-02-cifti" \ + --scale-csv "$OPT_ROOT/results/cubic_production_scale_9ab13dd_vs_optimized.csv" \ + --original-run-note "$NOTE" --measured-at 2026-09-22 \ + --platform 'CUBIC compute nodes (Xeon Gold 6130, 8-CPU allocations, GPFS)' \ + --cohort-source 'nifti=ABIDE alff, public S3 (make_cohort.py, materialize_cohort.py)' \ + --cohort-source 'mif=ModelArray walkthrough FDC, public OSF (make_mif_cohort.py)' \ + --cohort-source 'cifti=RBC freesurferpost subset, internal (make_cifti_cohort.py subset)' \ + --output "$OPT_ROOT/results/cubic-cross-modality-benchmark.pdf" +``` + +Then copy the three files into the repository, naming each CSV with its measurement date. +Replace the previous pinned files rather than keeping both, and update the links and numbers in +[`../README.md`](../README.md#cubic-hpc-validation): + +```bash +# local machine, from the ModelArrayIO root +R=tongt@cubic-login.uphs.upenn.edu:/cbica/home/tongt/modelarray_opt/results +scp "$R/cubic_cross_modality_9ab13dd_vs_optimized.csv" \ + benchmarks/to_modelarray/baselines/cubic_cross_modality_9ab13dd_vs_optimized_YYYYMMDD.csv +scp "$R/cubic_production_scale_9ab13dd_vs_optimized.csv" \ + benchmarks/to_modelarray/baselines/cubic_production_scale_9ab13dd_vs_optimized_YYYYMMDD.csv +scp "$R/cubic-cross-modality-benchmark.pdf" benchmarks/to_modelarray/reports/to-modelarray-cubic-benchmark.pdf +``` + +## Why this is the HPC pattern + +ModelArray's `element-splitting` shards the *element* axis across jobs and stitches chunk files +back together because HDF5 cannot be written concurrently. Building the matrix inverts that: +inputs are whole per-subject files (gzip is not seekable, so element shards would re-decode every +subject), while each **scalar's** output is already a separate file. So the job array shards by +scalar and there is nothing to aggregate. Within a task, `--workers` threads decode subjects +concurrently and the HDF5 compression pool / TileDB internal pool compress in parallel. + +Production-scale runs use `--target-chunk-mb 32` (the original run's tile geometry) and `--mem=48G` +(an 11.4 GB float32 matrix plus TileDB write buffers); calibrate from `peak_memory_bytes` in step 3. + +## Defects the ladder caught (and how the scripts now avoid them) + +- **Stale reader in the container.** ModelArray 0.1.5 reads `column_names` from an HDF5 + *attribute*; every converter version writes the dataset, so 0.1.5 rejects old and new outputs + alike. The HDF5 gate uses the current ModelArray installed into `$OPT_ROOT/Rlib`. +- **CRLF cohort files.** `csv.DictWriter`'s default `\r\n` broke shell consumers such as the + cache-priming `xargs cat` (`path\r` → "No such file"). All cohort writers now emit `\n`; the + priming step also strips `\r` defensively. +- **CPU-time rlimit.** The login node's 2 h `ulimit -t` propagates into jobs + (`PropagateResourceLimits=ALL`) and killed the original converter with SIGXCPU two-thirds of the + way through a full scalar (it needs ~3 h CPU). Every submission uses `sbatch --propagate=NONE` + (probe-verified unlimited inside jobs); `env.sh` also raises the soft limit to the hard limit. +- **Login node cannot open TileDB arrays** (Ivy Bridge, no AVX2 → SIGILL in libtiledb). Anything + that touches an array runs in a SLURM job; compute nodes are Xeon Gold with AVX-512. +- **Silent fast-path fallback.** Production dscalars carry per-subject provenance in the CIFTI + XML, shifting `vox_offset`; the exact-size check rejected 91/100 files. The fast path now + verifies dims/dtype/scaling plus a digest of the brain-model `MatrixIndicesMap` bytes; watch + `fast_path_rows` in `timing.jsonl`. diff --git a/benchmarks/to_modelarray/cubic/env.sh b/benchmarks/to_modelarray/cubic/env.sh new file mode 100755 index 0000000..af2a3dd --- /dev/null +++ b/benchmarks/to_modelarray/cubic/env.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Shared environment for the CUBIC validation ladder. Source this file. +# Everything writable lives under /cbica/home/tongt (never comp_space, never the project dir). +export OPT_ROOT=/cbica/home/tongt/modelarray_opt +export REPO=$OPT_ROOT/ModelArrayIO +export VENV=$OPT_ROOT/.venv +export SIF=/cbica/home/tongt/software/modelarray_confixel_0.1.5.sif + +# uv / apptainer caches must not land in the (small, purged) default locations. +export UV_CACHE_DIR=$OPT_ROOT/.uv-cache +export UV_PYTHON_INSTALL_DIR=$OPT_ROOT/.uv-python +export APPTAINER_CACHEDIR=$OPT_ROOT/.apptainer-cache +# The production cohort CSV uses /gpfs/fs001/... absolute paths; bind both spellings. +export APPTAINER_BIND=/cbica/home/tongt,/cbica/projects/pennlinc_rbc,/gpfs/fs001 + +# The login node sets a 2 h CPU-time ulimit that SLURM propagates into jobs; the +# original converter needs ~3 h CPU per full CIFTI scalar and was killed by SIGXCPU. +# Lift it (soft -> hard) for every job and gate; submit with --propagate=NONE too. +ulimit -t unlimited 2>/dev/null || ulimit -t "$(ulimit -Ht)" 2>/dev/null || true +export MODELARRAYIO_BENCHMARK_RESULTS_DIR=$OPT_ROOT/results +export MODELARRAYIO_BASELINE_ROOT=$OPT_ROOT/baseline-9ab13dd +export MODELARRAYIO_PYTHON=$VENV/bin/python +export MODELARRAYIO_S3_ANON=1 + +export REF_TDB=/cbica/projects/pennlinc_rbc/modelarray_rbc_tiledb/derivatives/rbc_freesurferpost.tdb +export RBC_COHORT=/cbica/projects/pennlinc_rbc/modelarray_rbc_tiledb/derivatives/rbc_freesurferpost_cohort.csv + +# Current ModelArray (installed from source into this user library) is the HDF5 gate; +# the container ships ModelArray 0.1.5, which reads column_names from an HDF5 attribute +# and cannot parse the current format. +export MODELARRAYIO_R_LIBS=$OPT_ROOT/Rlib +export R_LIBS_USER=$OPT_ROOT/Rlib +export GATE="$REPO/benchmarks/to_modelarray/cubic/gate.sh {output} {backend}" + +# Run a command inside the ModelArray container (Ubuntu 20.04, glibc 2.31) with the +# venv first on PATH. When already inside the container, run it directly. +in_sif() { + if [ -n "${APPTAINER_CONTAINER:-}${SINGULARITY_CONTAINER:-}" ]; then + PATH="$VENV/bin:/cbica/home/tongt/.local/bin:$PATH" "$@" + else + apptainer exec "$SIF" env PATH="$VENV/bin:/cbica/home/tongt/.local/bin:$PATH" "$@" + fi +} +export -f in_sif diff --git a/benchmarks/to_modelarray/cubic/gate.sh b/benchmarks/to_modelarray/cubic/gate.sh new file mode 100755 index 0000000..b3495b2 --- /dev/null +++ b/benchmarks/to_modelarray/cubic/gate.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# Read-gate dispatcher: `gate.sh OUTPUT BACKEND`. +# hdf5 -> R ModelArray/rhdf5/HDF5Array inside the container (r_read_gate.R) +# tiledb -> Python tiledb-py gate (the container has no R tiledb package) +set -euo pipefail +source "$(dirname "$0")/env.sh" +output=$1; backend=$2 +case "$backend" in + hdf5) in_sif Rscript "$REPO/benchmarks/to_modelarray/r_read_gate.R" "$output" hdf5 ;; + tiledb) in_sif python "$REPO/benchmarks/to_modelarray/tiledb_read_gate.py" "$output" tiledb ;; + *) echo "unknown backend: $backend" >&2; exit 2 ;; +esac diff --git a/benchmarks/to_modelarray/cubic/make_scalar_cohorts.sh b/benchmarks/to_modelarray/cubic/make_scalar_cohorts.sh new file mode 100755 index 0000000..1f73211 --- /dev/null +++ b/benchmarks/to_modelarray/cubic/make_scalar_cohorts.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# One pre-filtered long-format cohort CSV per scalar (all 8,693 subject-sessions), +# which is how a job array selects a scalar without any converter change. +# Usage: make_scalar_cohorts.sh complete reference scalars; also writes scalars.txt +# make_scalar_cohorts.sh SCALAR ... just these cohorts; scalars.txt is left unchanged +# Cohorts are deterministic subsets of $RBC_COHORT, so existing ones are kept; delete one to rebuild it. +set -euo pipefail +source "$(dirname "$0")/env.sh" +mkdir -p "$OPT_ROOT/cohorts/per-scalar" +if [ $# -gt 0 ]; then scalars=("$@"); write_list=0; else + mapfile -t scalars < <(python3 -c "import csv,sys;[print(r['scalar']) for r in csv.DictReader(open('$OPT_ROOT/reference_scalars.csv')) if r['complete']=='True']") + write_list=1 +fi +for s in "${scalars[@]}"; do + out="$OPT_ROOT/cohorts/per-scalar/$s.csv" + [ -f "$out" ] && continue + in_sif python "$REPO/benchmarks/to_modelarray/make_cifti_cohort.py" subset --input "$RBC_COHORT" \ + --output "$out" --subjects 8693 --scalars "$s" +done +if [ "$write_list" = 1 ]; then + printf '%s\n' "${scalars[@]}" > "$OPT_ROOT/cohorts/per-scalar/scalars.txt" +fi +echo "per-scalar cohorts ready for ${#scalars[@]} scalars" diff --git a/benchmarks/to_modelarray/cubic/sbatch/00_smoke.sbatch b/benchmarks/to_modelarray/cubic/sbatch/00_smoke.sbatch new file mode 100755 index 0000000..efeaa1b --- /dev/null +++ b/benchmarks/to_modelarray/cubic/sbatch/00_smoke.sbatch @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +#SBATCH -J ma-smoke +#SBATCH -p short +#SBATCH -c 4 +#SBATCH --mem=16G +#SBATCH -t 1:00:00 +#SBATCH -o /cbica/home/tongt/modelarray_opt/logs/%x-%j.out +set -euo pipefail +source /cbica/home/tongt/modelarray_opt/ModelArrayIO/benchmarks/to_modelarray/cubic/env.sh +cd "$REPO" +echo "== reference scalar completeness (needs a compute-node CPU for libtiledb) ==" +in_sif python benchmarks/to_modelarray/compare_to_reference.py --reference "$REF_TDB" --list --output-csv "$OPT_ROOT/reference_scalars.csv" +echo "== compute-node imports =="; in_sif python -c "import tiledb,h5py,numpy,nibabel,modelarrayio;print('ok',tiledb.__version__,h5py.__version__,numpy.__version__)" +echo "== source identity =="; in_sif python benchmarks/to_modelarray/source_identity.py --repo "$REPO" +echo "== pytest (offline subset) =="; in_sif python -m pytest -q -m "not s3 and not downloaded_data" -p no:cacheprovider +echo "== fixture quick tier with both gate branches ==" +in_sif python benchmarks/to_modelarray/run_benchmarks.py --cohort benchmarks/to_modelarray/fixtures/nifti-20.csv \ + --mask test/data_voxel_toy/group_mask_FA.nii.gz --scalar-columns FA --tier quick --repetitions 1 \ + --cache-state warm --run-id cubic-00-smoke --result-shard nifti --modelarrayio "$VENV/bin/modelarrayio" --r-gate-command "$GATE" +in_sif python benchmarks/to_modelarray/run_benchmarks.py --cohort benchmarks/to_modelarray/fixtures/cifti.csv --tier quick --repetitions 1 \ + --cache-state warm --run-id cubic-00-smoke --result-shard cifti --modelarrayio "$VENV/bin/modelarrayio" --r-gate-command "$GATE" +echo "== fixture original-vs-optimized (bitwise) ==" +in_sif python benchmarks/to_modelarray/compare_revisions.py --cohort benchmarks/to_modelarray/fixtures/cifti.csv \ + --baseline-executable benchmarks/to_modelarray/baselines/modelarrayio-baseline-9ab13dd --optimized-executable "$VENV/bin/modelarrayio" \ + --results-dir "$OPT_ROOT/results/cubic-00-fixture-compare" --repetitions 1 --warmups 0 \ + --r-gate-script benchmarks/to_modelarray/r_read_gate.R --gate-command "$GATE" --baseline-revision 9ab13dd +echo "SMOKE PASSED" diff --git a/benchmarks/to_modelarray/cubic/sbatch/02_compare_public.sbatch b/benchmarks/to_modelarray/cubic/sbatch/02_compare_public.sbatch new file mode 100755 index 0000000..c9cfe67 --- /dev/null +++ b/benchmarks/to_modelarray/cubic/sbatch/02_compare_public.sbatch @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +#SBATCH -J ma-compare +#SBATCH -p all +#SBATCH -c 8 +#SBATCH --mem=32G +#SBATCH -t 4:00:00 +#SBATCH -o /cbica/home/tongt/modelarray_opt/logs/%x-%j.out +set -euo pipefail +source /cbica/home/tongt/modelarray_opt/ModelArrayIO/benchmarks/to_modelarray/cubic/env.sh +cd "$REPO" +BASE=benchmarks/to_modelarray/baselines/modelarrayio-baseline-9ab13dd; OPT=$VENV/bin/modelarrayio +# area areaPial curv: scalars the original production run completed, so every CIFTI output can +# also be checked bitwise against the deployed reference. +CIFTI_COHORT=$OPT_ROOT/cohorts/cifti-rbc-100x3ref.csv +CIFTI_SCALARS=(--scalar area --scalar areaPial --scalar curv) +# Cohort provenance for the report (it cannot tell public, internal and synthetic data apart). +SOURCES=(--cohort-source 'nifti=ABIDE alff, public S3 (make_cohort.py, materialize_cohort.py)' + --cohort-source 'mif=ModelArray walkthrough FDC, public OSF (make_mif_cohort.py)' + --cohort-source 'cifti=RBC freesurferpost subset, internal (make_cifti_cohort.py subset)') +COMMON=(--baseline-executable $BASE --optimized-executable $OPT --repetitions 3 --warmups 1 --workers 1 --r-gate-script benchmarks/to_modelarray/r_read_gate.R --gate-command "$GATE" --baseline-revision 9ab13dd) +in_sif python benchmarks/to_modelarray/compare_revisions.py --cohort "$OPT_ROOT/cohorts/abide-100-local.csv" --mask "$OPT_ROOT/data/group_mask.nii.gz" --scalar-columns alff --results-dir "$OPT_ROOT/results/cubic-02-nifti" "${COMMON[@]}" +in_sif python benchmarks/to_modelarray/compare_revisions.py --cohort "$OPT_ROOT/cohorts/mif-walkthrough.csv" --index-file "$OPT_ROOT/data/mif-walkthrough/FDC/index.mif" --directions-file "$OPT_ROOT/data/mif-walkthrough/FDC/directions.mif" --results-dir "$OPT_ROOT/results/cubic-02-mif" "${COMMON[@]}" +in_sif python benchmarks/to_modelarray/compare_revisions.py --cohort "$CIFTI_COHORT" --results-dir "$OPT_ROOT/results/cubic-02-cifti" "${COMMON[@]}" +for d in "$OPT_ROOT"/results/cubic-02-cifti/optimized-*-rep-01; do + in_sif python benchmarks/to_modelarray/compare_to_reference.py --reference "$REF_TDB" --candidate "$d"/modelarray.* "${CIFTI_SCALARS[@]}" --output-csv "$d/reference_check.csv" +done +echo "== raw fast path vs nibabel on production headers (both must equal the reference) ==" +for raw in 1 0; do + MODELARRAYIO_CIFTI_RAW=$raw in_sif python benchmarks/to_modelarray/run_benchmarks.py --cohort "$CIFTI_COHORT" --tier quick --repetitions 1 --cache-state warm \ + --run-id cubic-02-cifti-raw$raw --backends hdf5 --modelarrayio "$OPT" --r-gate-command "$GATE" + out=$(ls -d "$OPT_ROOT"/results/cubic-02-cifti-raw$raw/outputs/*/modelarray.h5 | head -1) + in_sif python benchmarks/to_modelarray/compare_to_reference.py --reference "$REF_TDB" --candidate "$out" "${CIFTI_SCALARS[@]}" --output-csv "$OPT_ROOT/results/cubic-02-cifti-raw$raw/reference_check.csv" + # The raw run must read every row through the fast path; a silent fallback shows up here. + in_sif python - "$OPT_ROOT/results/cubic-02-cifti-raw$raw" <<'PY' +import json, glob, sys +recs = [json.loads(l) for p in glob.glob(sys.argv[1] + '/outputs/*/timing.jsonl') for l in open(p) if l.strip()] +print('fast_path_rows:', [(r['scalar_name'], r.get('fast_path_rows'), r['matrix_shape'][0]) for r in recs if r.get('event') == 'scalar_complete']) +PY +done +echo "== report ==" +in_sif python benchmarks/to_modelarray/summarize_comparison.py --run-dir "$OPT_ROOT/results/cubic-02-nifti" --run-dir "$OPT_ROOT/results/cubic-02-mif" --run-dir "$OPT_ROOT/results/cubic-02-cifti" --output "$OPT_ROOT/results/cubic_cross_modality_9ab13dd_vs_optimized.csv" +in_sif python benchmarks/to_modelarray/create_comparison_report.py --run-dir "$OPT_ROOT/results/cubic-02-nifti" --run-dir "$OPT_ROOT/results/cubic-02-mif" --run-dir "$OPT_ROOT/results/cubic-02-cifti" --platform "CUBIC compute node ($(nproc) CPU allocation)" "${SOURCES[@]}" --output "$OPT_ROOT/results/cubic-cross-modality-benchmark.pdf" +echo "COMPARE PASSED" diff --git a/benchmarks/to_modelarray/cubic/sbatch/03_single_scalar.sbatch b/benchmarks/to_modelarray/cubic/sbatch/03_single_scalar.sbatch new file mode 100755 index 0000000..f1a41cd --- /dev/null +++ b/benchmarks/to_modelarray/cubic/sbatch/03_single_scalar.sbatch @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +#SBATCH -J ma-scalar +#SBATCH -p all +#SBATCH -c 8 +#SBATCH --mem=48G +#SBATCH -t 12:00:00 +#SBATCH -o /cbica/home/tongt/modelarray_opt/logs/%x-%j.out +# One full production scalar (8,693 subjects): old vs new (1 rep), then a worker sweep. +set -euo pipefail +source /cbica/home/tongt/modelarray_opt/ModelArrayIO/benchmarks/to_modelarray/cubic/env.sh +cd "$REPO" +SCALAR=${SCALAR:-area}; COHORT="$OPT_ROOT/cohorts/per-scalar/$SCALAR.csv" +BASE=benchmarks/to_modelarray/baselines/modelarrayio-baseline-9ab13dd; OPT=$VENV/bin/modelarrayio +echo "== prime page cache ($SCALAR) =="; cut -d, -f3 "$COHORT" | tail -n +2 | tr -d "\r" | xargs -P 8 -n 200 cat > /dev/null +echo "== old vs new, 1 rep, chunk target 32 MiB ==" +in_sif python benchmarks/to_modelarray/compare_revisions.py --cohort "$COHORT" --results-dir "$OPT_ROOT/results/cubic-03-$SCALAR-compare" \ + --baseline-executable $BASE --optimized-executable $OPT --repetitions 1 --warmups 0 --workers 1 --target-chunk-mb 32 \ + --r-gate-script benchmarks/to_modelarray/r_read_gate.R --gate-command "$GATE" --baseline-revision 9ab13dd +echo "== optimized worker sweep ==" +in_sif python benchmarks/to_modelarray/run_benchmarks.py --cohort "$COHORT" --tier quick --workers 1 4 8 --target-chunk-mb 32 \ + --backends hdf5 tiledb --repetitions 1 --cache-state warm --run-id cubic-03-$SCALAR-sweep --modelarrayio "$OPT" --r-gate-command "$GATE" +echo "== bitwise vs the deployed reference for every optimized output ==" +for out in "$OPT_ROOT"/results/cubic-03-$SCALAR-compare/optimized-*-rep-01/modelarray.* "$OPT_ROOT"/results/cubic-03-$SCALAR-sweep/outputs/*/modelarray.*; do + in_sif python benchmarks/to_modelarray/compare_to_reference.py --reference "$REF_TDB" --candidate "$out" --scalar "$SCALAR" --output-csv "$out.reference_check.csv" +done +echo "== free space: drop bulky outputs, keep logs/timing/checks ==" +find "$OPT_ROOT/results/cubic-03-$SCALAR-sweep/outputs" -maxdepth 2 \( -name 'modelarray.h5' -o -name 'modelarray.tdb' \) -exec rm -rf {} + +echo "SCALAR PASSED" diff --git a/benchmarks/to_modelarray/cubic/sbatch/04_scalar_array.sbatch b/benchmarks/to_modelarray/cubic/sbatch/04_scalar_array.sbatch new file mode 100755 index 0000000..7bb4f59 --- /dev/null +++ b/benchmarks/to_modelarray/cubic/sbatch/04_scalar_array.sbatch @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +#SBATCH -J ma-shard +#SBATCH -p all +#SBATCH -c 8 +#SBATCH --mem=48G +#SBATCH -t 4:00:00 +#SBATCH -o /cbica/home/tongt/modelarray_opt/logs/%x-%A_%a.out +# Scalar-sharded conversion: one task per scalar, its own per-scalar file, no aggregation. +# Submit with: sbatch --array=1-$(wc -l < .../per-scalar/scalars.txt) 04_scalar_array.sbatch +set -euo pipefail +source /cbica/home/tongt/modelarray_opt/ModelArrayIO/benchmarks/to_modelarray/cubic/env.sh +cd "$REPO" +SCALAR=$(sed -n "${SLURM_ARRAY_TASK_ID}p" "$OPT_ROOT/cohorts/per-scalar/scalars.txt") +COHORT="$OPT_ROOT/cohorts/per-scalar/$SCALAR.csv"; OUT="$OPT_ROOT/results/cubic-04-shard/$SCALAR"; mkdir -p "$OUT" +KEEP=${KEEP_OUTPUTS:-0} +for backend in hdf5 tiledb; do + suffix=$([ $backend = hdf5 ] && echo h5 || echo tdb) + comp=$([ $backend = hdf5 ] && echo "gzip --compression-level 4" || echo "zstd --compression-level 1") + start=$(date +%s.%N) + in_sif "$VENV/bin/modelarrayio" to-modelarray --cohort-file "$COHORT" --output "$OUT/${SCALAR}_modelarray.$suffix" \ + --backend $backend --workers 8 --dtype float32 --compression $comp --target-chunk-mb 32 --no-split-files \ + --timing-jsonl "$OUT/timing_$backend.jsonl" --log-level WARNING + echo "wall_seconds,$backend,$(echo "$(date +%s.%N) - $start" | bc)" >> "$OUT/wall.csv" + bash benchmarks/to_modelarray/cubic/gate.sh "$OUT/${SCALAR}_modelarray.$suffix" $backend | tee -a "$OUT/gate_$backend.log" + in_sif python benchmarks/to_modelarray/compare_to_reference.py --reference "$REF_TDB" --candidate "$OUT/${SCALAR}_modelarray.$suffix" --scalar "$SCALAR" --output-csv "$OUT/reference_check_$backend.csv" + ls "$OUT"/.*modelarrayio-tmp-* "$OUT"/.*modelarrayio-backup-* 2>/dev/null && { echo "leftover temporaries!"; exit 1; } || true + [ "$KEEP" = 1 ] || rm -rf "$OUT/${SCALAR}_modelarray.$suffix" +done +echo "SHARD $SCALAR PASSED" | tee "$OUT/PASSED" diff --git a/benchmarks/to_modelarray/cubic/sbatch/05_production.sbatch b/benchmarks/to_modelarray/cubic/sbatch/05_production.sbatch new file mode 100755 index 0000000..dfed6cf --- /dev/null +++ b/benchmarks/to_modelarray/cubic/sbatch/05_production.sbatch @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +#SBATCH -J ma-prod +#SBATCH -p all +#SBATCH -c 8 +#SBATCH --mem=48G +#SBATCH -t 4:00:00 +#SBATCH -o /cbica/home/tongt/modelarray_opt/logs/%x-%A_%a.out +# Full production build: every scalar -> its own HDF5 file (kept). Submit via submit_production.sh. +set -euo pipefail +source /cbica/home/tongt/modelarray_opt/ModelArrayIO/benchmarks/to_modelarray/cubic/env.sh +cd "$REPO" +SCALAR=$(sed -n "${SLURM_ARRAY_TASK_ID}p" "$OPT_ROOT/cohorts/per-scalar/all_scalars.txt") +COHORT="$OPT_ROOT/cohorts/per-scalar/$SCALAR.csv"; OUT="$OPT_ROOT/results/production"; mkdir -p "$OUT/checks" +start=$(date +%s.%N) +in_sif "$VENV/bin/modelarrayio" to-modelarray --cohort-file "$COHORT" --output "$OUT/${SCALAR}_modelarray.h5" \ + --backend hdf5 --workers 8 --dtype float32 --compression gzip --compression-level 4 --target-chunk-mb 32 --no-split-files \ + --timing-jsonl "$OUT/checks/${SCALAR}_timing.jsonl" --log-level WARNING +echo "$SCALAR,$(echo "$(date +%s.%N) - $start" | bc)" >> "$OUT/checks/wall_seconds.csv" +bash benchmarks/to_modelarray/cubic/gate.sh "$OUT/${SCALAR}_modelarray.h5" hdf5 | tee "$OUT/checks/${SCALAR}_gate.log" +if grep -qx "$SCALAR" "$OPT_ROOT/cohorts/per-scalar/scalars.txt"; then + in_sif python benchmarks/to_modelarray/compare_to_reference.py --reference "$REF_TDB" --candidate "$OUT/${SCALAR}_modelarray.h5" --scalar "$SCALAR" --output-csv "$OUT/checks/${SCALAR}_reference_check.csv" +fi +echo "PROD $SCALAR PASSED" | tee "$OUT/checks/${SCALAR}.PASSED" diff --git a/benchmarks/to_modelarray/cubic/setup_env.sh b/benchmarks/to_modelarray/cubic/setup_env.sh new file mode 100755 index 0000000..ebe7000 --- /dev/null +++ b/benchmarks/to_modelarray/cubic/setup_env.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +# Step 0 (login node): environment, baseline export, public data, cohorts. +# Route: run everything inside the ModelArray container (glibc 2.31) with a +# uv-managed Python 3.12 so the exact local library versions install as wheels. +set -euo pipefail +source "$(dirname "$0")/env.sh" +UV=/cbica/home/tongt/.local/bin/uv +mkdir -p "$OPT_ROOT"/{data,cohorts,results,logs} "$UV_CACHE_DIR" "$UV_PYTHON_INSTALL_DIR" "$APPTAINER_CACHEDIR" + +echo "== quota =="; df -h /cbica/home/tongt | tail -1 + +echo "== python 3.12 (uv-managed; runs on glibc >= 2.17, host and container) ==" +"$UV" python install 3.12 +PY312=$("$UV" python find 3.12) +echo "python: $PY312" + +echo "== venv inside the container so wheels resolve against its glibc ==" +in_sif "$UV" venv --python "$PY312" "$VENV" +in_sif "$UV" pip install --python "$VENV/bin/python" -e "$REPO[s3]" pytest reportlab +in_sif python - <<'PY' +import h5py, nibabel, numpy, tiledb, modelarrayio, sys +print('in-container:', sys.version.split()[0], 'numpy', numpy.__version__, 'h5py', h5py.__version__, + 'tiledb', tiledb.__version__, 'nibabel', nibabel.__version__) +PY + +echo "== current ModelArray R package into a user library (the container ships 0.1.5) ==" +# ModelArray 0.1.5 reads column_names from an HDF5 *attribute*; the current format (and the +# converter) store it as a dataset, so the HDF5 gate must use the current ModelArray source, +# rsynced to $OPT_ROOT/ModelArray. Deps come from CRAN (login node has internet). +mkdir -p "$OPT_ROOT/Rlib" +in_sif Rscript -e 'lib<-Sys.getenv("R_LIBS_USER"); .libPaths(c(lib,.libPaths())); deps<-c("DelayedArray","HDF5Array","broom","crayon","dplyr","glue","hdf5r","magrittr","mgcv","pbapply","pbmcapply","rhdf5","tibble","tidyr","cheapr"); miss<-deps[!vapply(deps, requireNamespace, TRUE, quietly=TRUE)]; cat("missing:", miss, "\n"); if (length(miss)) install.packages(miss, lib=lib, repos="https://cloud.r-project.org", Ncpus=4)' +in_sif R CMD INSTALL --library="$OPT_ROOT/Rlib" "$OPT_ROOT/ModelArray" +in_sif Rscript -e '.libPaths(c(Sys.getenv("R_LIBS_USER"), .libPaths())); cat("ModelArray", as.character(packageVersion("ModelArray")), "from", find.package("ModelArray"), "\n")' + +echo "== baseline 9ab13dd export (host git; cache under home so compute nodes see it) ==" +"$REPO/benchmarks/to_modelarray/baselines/modelarrayio-baseline-9ab13dd" --prepare-only + +echo "== source identity (must equal the local tree) ==" +in_sif python "$REPO/benchmarks/to_modelarray/source_identity.py" --repo "$REPO" | tee "$OPT_ROOT/source_identity.json" + +echo "== public NIfTI (ABIDE-100) materialized locally ==" +in_sif python "$REPO/benchmarks/to_modelarray/make_cohort.py" --subjects 100 --derivatives alff \ + --output "$OPT_ROOT/cohorts/abide-100-s3.csv" +GROUP_MASK_KEY=s3://fcp-indi/data/Projects/ABIDE_Initiative/Outputs/cpac/filt_global/func_mask/CMU_a_0050642_func_mask.nii.gz +in_sif python "$REPO/benchmarks/to_modelarray/materialize_cohort.py" --cohort "$OPT_ROOT/cohorts/abide-100-s3.csv" \ + --data-dir "$OPT_ROOT/data/abide" --output "$OPT_ROOT/cohorts/abide-100-local.csv" \ + --extra-key "$GROUP_MASK_KEY" +cp "$OPT_ROOT/data/abide/func_mask/CMU_a_0050642_func_mask.nii.gz" "$OPT_ROOT/data/group_mask.nii.gz" + +echo "== public MIF (walkthrough FDC) ==" +in_sif python "$REPO/benchmarks/to_modelarray/make_mif_cohort.py" --data-dir "$OPT_ROOT/data/mif-walkthrough" \ + --output "$OPT_ROOT/cohorts/mif-walkthrough.csv" + +echo "== real CIFTI subset cohorts (read-only project CSV, no copying) ==" +in_sif python "$REPO/benchmarks/to_modelarray/make_cifti_cohort.py" subset --input "$RBC_COHORT" \ + --output "$OPT_ROOT/cohorts/cifti-rbc-100x3ref.csv" --subjects 100 --scalars area areaPial curv + +# The reference listing opens TileDB arrays; libtiledb needs a compute-node CPU +# (the login node raised SIGILL), so it runs in sbatch/00_smoke.sbatch instead. +echo "setup complete (run sbatch/00_smoke.sbatch next; then make_scalar_cohorts.sh)" diff --git a/benchmarks/to_modelarray/cubic/submit_ladder.sh b/benchmarks/to_modelarray/cubic/submit_ladder.sh new file mode 100755 index 0000000..cbd3b60 --- /dev/null +++ b/benchmarks/to_modelarray/cubic/submit_ladder.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# Submit steps 0 -> 2 -> 3 -> 4 as a dependency chain (each waits for the previous to succeed). +set -euo pipefail +source "$(dirname "$0")/env.sh" +D="$(dirname "$0")/sbatch" +j0=$(sbatch --propagate=NONE --parsable "$D/00_smoke.sbatch") +j2=$(sbatch --propagate=NONE --parsable --dependency=afterok:$j0 "$D/02_compare_public.sbatch") +j3=$(sbatch --propagate=NONE --parsable --dependency=afterok:$j2 "$D/03_single_scalar.sbatch") +n=$(wc -l < "$OPT_ROOT/cohorts/per-scalar/scalars.txt") +j4=$(sbatch --propagate=NONE --parsable --dependency=afterok:$j3 --array=1-${n}%6 "$D/04_scalar_array.sbatch") +echo "submitted: smoke=$j0 compare=$j2 scalar=$j3 shard=$j4 (array 1-$n)" diff --git a/benchmarks/to_modelarray/cubic/submit_production.sh b/benchmarks/to_modelarray/cubic/submit_production.sh new file mode 100755 index 0000000..ba148c7 --- /dev/null +++ b/benchmarks/to_modelarray/cubic/submit_production.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# Guarded step-5 submission: refuse unless every step-4 shard PASSED, then build all scalars. +# scalars.txt stays the list of *complete reference* scalars (used for bitwise checks); +# all_scalars.txt drives the production array. +set -euo pipefail +source "$(dirname "$0")/env.sh" +n_expected=$(wc -l < "$OPT_ROOT/cohorts/per-scalar/scalars.txt") +n_passed=$(ls "$OPT_ROOT"/results/cubic-04-shard/*/PASSED 2>/dev/null | wc -l) +[ "$n_passed" -eq "$n_expected" ] || { echo "step 4: $n_passed/$n_expected shards passed; refusing to submit step 5"; exit 1; } +python3 -c "import csv;print('\n'.join(sorted({r['scalar_name'] for r in csv.DictReader(open('$RBC_COHORT'))})))" \ + > "$OPT_ROOT/cohorts/per-scalar/all_scalars.txt" +# Per-scalar cohorts for every scalar (the step-4 reference scalars already have theirs). +mapfile -t all_scalars < "$OPT_ROOT/cohorts/per-scalar/all_scalars.txt" +bash "$(dirname "$0")/make_scalar_cohorts.sh" "${all_scalars[@]}" +n=$(wc -l < "$OPT_ROOT/cohorts/per-scalar/all_scalars.txt") +echo "free space before production build:"; df -h /cbica/home/tongt | tail -1 +sbatch --propagate=NONE --array=1-${n}%10 "$(dirname "$0")/sbatch/05_production.sbatch" diff --git a/benchmarks/to_modelarray/fixtures/cifti.csv b/benchmarks/to_modelarray/fixtures/cifti.csv new file mode 100644 index 0000000..93d92e6 --- /dev/null +++ b/benchmarks/to_modelarray/fixtures/cifti.csv @@ -0,0 +1,2 @@ +subject_id,scalar_name,source_file +example,example,../../../test/data_cifti_toy/example.dscalar.nii diff --git a/benchmarks/to_modelarray/fixtures/nifti-20.csv b/benchmarks/to_modelarray/fixtures/nifti-20.csv new file mode 100644 index 0000000..ce39287 --- /dev/null +++ b/benchmarks/to_modelarray/fixtures/nifti-20.csv @@ -0,0 +1,21 @@ +subject_id,source_mask_file,FA +sub-01,../../../test/data_voxel_toy/individual_masks/sub-01_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-01_FA.nii.gz +sub-02,../../../test/data_voxel_toy/individual_masks/sub-02_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-02_FA.nii.gz +sub-03,../../../test/data_voxel_toy/individual_masks/sub-03_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-03_FA.nii.gz +sub-04,../../../test/data_voxel_toy/individual_masks/sub-04_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-04_FA.nii.gz +sub-05,../../../test/data_voxel_toy/individual_masks/sub-05_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-05_FA.nii.gz +sub-06,../../../test/data_voxel_toy/individual_masks/sub-06_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-06_FA.nii.gz +sub-07,../../../test/data_voxel_toy/individual_masks/sub-07_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-07_FA.nii.gz +sub-08,../../../test/data_voxel_toy/individual_masks/sub-08_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-08_FA.nii.gz +sub-09,../../../test/data_voxel_toy/individual_masks/sub-09_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-09_FA.nii.gz +sub-10,../../../test/data_voxel_toy/individual_masks/sub-10_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-10_FA.nii.gz +sub-11,../../../test/data_voxel_toy/individual_masks/sub-11_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-11_FA.nii.gz +sub-12,../../../test/data_voxel_toy/individual_masks/sub-12_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-12_FA.nii.gz +sub-13,../../../test/data_voxel_toy/individual_masks/sub-13_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-13_FA.nii.gz +sub-14,../../../test/data_voxel_toy/individual_masks/sub-14_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-14_FA.nii.gz +sub-15,../../../test/data_voxel_toy/individual_masks/sub-15_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-15_FA.nii.gz +sub-16,../../../test/data_voxel_toy/individual_masks/sub-16_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-16_FA.nii.gz +sub-17,../../../test/data_voxel_toy/individual_masks/sub-17_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-17_FA.nii.gz +sub-18,../../../test/data_voxel_toy/individual_masks/sub-18_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-18_FA.nii.gz +sub-19,../../../test/data_voxel_toy/individual_masks/sub-19_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-19_FA.nii.gz +sub-20,../../../test/data_voxel_toy/individual_masks/sub-20_FA_mask.nii.gz,../../../test/data_voxel_toy/FA/sub-20_FA.nii.gz diff --git a/benchmarks/to_modelarray/harness_common.py b/benchmarks/to_modelarray/harness_common.py new file mode 100644 index 0000000..bdbb1f5 --- /dev/null +++ b/benchmarks/to_modelarray/harness_common.py @@ -0,0 +1,126 @@ +"""Cohort, CSV and provenance helpers shared by the ``to-modelarray`` benchmark scripts. + +The scripts import this module as a sibling: a script run by path has its own directory on +``sys.path``, and the test suite adds this directory through pytest's ``pythonpath`` setting. +""" + +from __future__ import annotations + +import csv +import hashlib +import importlib.metadata +import platform +from collections.abc import Sequence +from pathlib import Path +from urllib.parse import urlparse + +PACKAGE_NAMES = ('modelarrayio', 'numpy', 'nibabel', 'h5py', 'tiledb') +# Header of a long-format cohort, which the converter combines into one output. +LONG_COHORT_FIELDS = ('subject_id', 'scalar_name', 'source_file') + + +def read_cohort(path: Path) -> list[dict[str, str]]: + with path.open(newline='', encoding='utf-8') as stream: + return list(csv.DictReader(stream)) + + +def write_csv( + path: Path, rows: list[dict[str, object]], fieldnames: Sequence[str] | None = None +) -> None: + """Write ``rows`` with ``\\n`` line endings, which shell consumers such as ``cut`` need. + + Without ``fieldnames`` the header is every key in first-seen order across the rows. + """ + if fieldnames is None: + fieldnames = list(dict.fromkeys(key for row in rows for key in row)) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open('w', newline='', encoding='utf-8') as stream: + writer = csv.DictWriter(stream, fieldnames=fieldnames, lineterminator='\n') + writer.writeheader() + writer.writerows(rows) + + +def output_path(case_dir: Path, backend: str, scalar: str | None = None) -> Path: + """Return the converter output in ``case_dir``: the requested ``modelarray.`` path, or + the ``_modelarray.`` file it writes for one scalar of a split (wide) cohort.""" + suffix = '.h5' if backend == 'hdf5' else '.tdb' + return case_dir / (f'modelarray{suffix}' if scalar is None else f'{scalar}_modelarray{suffix}') + + +def source_modality(paths: list[str]) -> str: + """Return the single modality (``nifti``, ``cifti`` or ``mif``) shared by ``paths``.""" + modalities = set() + for source in paths: + path = urlparse(source).path.lower() + if path.endswith(('.dscalar.nii', '.pscalar.nii', '.pconn.nii')): + modalities.add('cifti') + elif path.endswith(('.mif', '.mif.gz')): + modalities.add('mif') + elif path.endswith(('.nii', '.nii.gz')): + modalities.add('nifti') + else: + raise ValueError(f'Cannot infer benchmark modality from {source!r}.') + if len(modalities) != 1: + raise ValueError(f'Benchmark cohort must contain one modality, got {sorted(modalities)}.') + return modalities.pop() + + +def infer_scalar_columns(rows: list[dict[str, str]]) -> list[str]: + """Return wide-format scalar columns, or [] for a long-format cohort.""" + if not rows: + raise ValueError('Cohort CSV has no rows.') + if {'scalar_name', 'source_file'} <= rows[0].keys(): + return [] + excluded = {'subject_id', 'source_mask_file'} + columns = [] + for column in rows[0]: + if column in excluded: + continue + values = [row[column] for row in rows if row.get(column)] + if not values: + continue + try: + source_modality(values) + except ValueError: + continue + columns.append(column) + if not columns: + raise ValueError('Could not infer wide-format scalar columns; pass --scalar-columns.') + return columns + + +def scalar_source_paths(rows: list[dict[str, str]], scalar_columns: list[str]) -> list[str]: + """Return scalar inputs without companion mask paths.""" + if scalar_columns: + return [row[column] for row in rows for column in scalar_columns if row.get(column)] + return [row['source_file'] for row in rows if row.get('source_file')] + + +def scalar_names(rows: list[dict[str, str]], scalar_columns: list[str]) -> list[str]: + if scalar_columns: + return list(scalar_columns) + return list(dict.fromkeys(row['scalar_name'] for row in rows if row.get('scalar_name'))) + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open('rb') as stream: + for block in iter(lambda: stream.read(1024 * 1024), b''): + digest.update(block) + return digest.hexdigest() + + +def output_stats(path: Path) -> tuple[int, int]: + """Return (total bytes, file count) for an output file or directory tree.""" + files = [path] if path.is_file() else [child for child in path.rglob('*') if child.is_file()] + return sum(child.stat().st_size for child in files), len(files) + + +def package_versions() -> dict[str, str | None]: + result: dict[str, str | None] = {'python': platform.python_version()} + for package in PACKAGE_NAMES: + try: + result[package] = importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + result[package] = None + return result diff --git a/benchmarks/to_modelarray/make_cifti_cohort.py b/benchmarks/to_modelarray/make_cifti_cohort.py new file mode 100644 index 0000000..8340acc --- /dev/null +++ b/benchmarks/to_modelarray/make_cifti_cohort.py @@ -0,0 +1,354 @@ +"""Create CIFTI dscalar benchmark cohorts. + +Two modes: + +``synthesize`` + Write real, seeded ``.dscalar.nii`` files (surface greyordinates at a + configurable density) plus a long-format cohort CSV. Reproducible on any + machine with no internal data, so the CIFTI load/runtime path is publicly + testable. + + By default, values are drawn from scalar-appropriate marginal + distributions (``thickness``/``area``/``sulc`` ranges) with a deterministic + medial-wall zero block, which is realistic enough for load/runtime timing + but still carries no real spatial autocorrelation, so it is not valid for + compression-ratio conclusions. + + Pass ``--reference-dir`` to derive, per scalar, the shared medial-wall zero + mask and the per-vertex mean and standard deviation across a small set of + real ``.dscalar.nii`` files, then generate each synthetic subject as + ``clip(mean + std * noise, observed_min, observed_max)`` with the medial + wall zeroed. Only aggregate per-vertex statistics are read and fresh noise + is drawn, so no real subject's row is reproduced, while the real mean map's + spatial smoothness makes compression ratios realistic. The reference axes + are copied verbatim so the header (and therefore ``vox_offset``) matches the + real files exactly. + +``subset`` + Take an existing long-format CIFTI cohort (for example the internal + freesurferpost cohort) and emit a smaller deterministic slice — the first + ``--subjects`` subjects for the selected scalars — so the real data can be + benchmarked with the same harness without editing CSVs by hand. + +Both modes emit ``subject_id,scalar_name,source_file`` rows, which +``run_benchmarks.py`` drives as a long-format (combined-output) CIFTI cohort. +""" + +from __future__ import annotations + +import argparse +import csv +from collections import OrderedDict +from pathlib import Path + +import harness_common as common + +# fs_LR den-164k: 163,842 vertices per hemisphere, both hemispheres, no medial +# wall removed, which is what the production freesurferpost dscalars carry. +DEFAULT_GRAYORDINATES = 327684 +DEFAULT_SCALARS = ('thickness', 'area', 'sulc') + +# Marginal value profiles for the no-reference synthesizer, chosen to match the +# ranges of real fs_LR surface metrics: (mean, std, min, max). These reproduce +# realistic value ranges but not real spatial autocorrelation. +SCALAR_MARGINALS: dict[str, tuple[float, float, float, float]] = { + 'thickness': (2.24, 1.02, 0.0, 5.0), + 'area': (0.52, 0.32, 0.0, 13.5), + 'sulc': (0.09, 5.35, -14.0, 17.0), + 'curv': (0.0, 0.13, -1.0, 1.0), + 'myelin': (1.6, 0.5, 0.0, 8.0), +} +# Fraction of greyordinates that fall on the medial wall (constant zero block) +# in real fs_LR den-164k dscalars. +DEFAULT_MEDIAL_WALL_FRACTION = 0.1017 + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest='mode', required=True) + + synth = subparsers.add_parser( + 'synthesize', help='Write seeded dscalar files and a long-format cohort CSV.' + ) + synth.add_argument( + '--output-dir', + type=Path, + required=True, + help='Directory to hold the generated .dscalar.nii files and cohort.csv.', + ) + synth.add_argument('--subjects', type=int, default=20) + synth.add_argument('--scalars', nargs='+', default=list(DEFAULT_SCALARS)) + synth.add_argument( + '--grayordinates', + type=int, + default=DEFAULT_GRAYORDINATES, + help='Total surface greyordinates (split evenly across two hemispheres).', + ) + synth.add_argument('--seed', type=int, default=0) + synth.add_argument( + '--reference-dir', + type=Path, + default=None, + help=( + 'Directory of real .dscalar.nii files to derive realistic per-vertex mean/std and ' + 'the medial-wall zero mask from (matched per scalar by filename suffix). When set, ' + '--grayordinates is taken from the reference header and generated compression ratios ' + 'are realistic; only aggregate statistics are read, never a real subject row.' + ), + ) + synth.add_argument( + '--cohort', + type=Path, + help='Cohort CSV path (default: /cohort.csv).', + ) + synth.set_defaults(func=synthesize) + + subset = subparsers.add_parser( + 'subset', help='Slice an existing long-format CIFTI cohort into a smaller one.' + ) + subset.add_argument( + '--input', type=Path, required=True, help='Existing long-format cohort CSV.' + ) + subset.add_argument('--output', type=Path, required=True) + subset.add_argument('--subjects', type=int, default=100) + subset.add_argument( + '--scalars', + nargs='+', + help='Scalars to keep (default: every scalar present in the input).', + ) + subset.set_defaults(func=subset_cohort) + + return parser + + +def _validate_common(subjects: int) -> None: + if subjects < 1: + raise ValueError('--subjects must be positive') + + +def _build_brain_model_axis(grayordinates: int): + """Return a two-hemisphere surface greyordinate axis of the requested length.""" + import numpy as np + from nibabel.cifti2.cifti2_axes import BrainModelAxis + + if grayordinates < 2: + raise ValueError('--grayordinates must be at least 2') + left_vertices = grayordinates // 2 + right_vertices = grayordinates - left_vertices + left = BrainModelAxis.from_surface(np.arange(left_vertices), left_vertices, name='CortexLeft') + right = BrainModelAxis.from_surface( + np.arange(right_vertices), right_vertices, name='CortexRight' + ) + return left + right + + +def _reference_files(reference_dir: Path, scalar: str) -> list[Path]: + """Return the real .dscalar.nii files for one scalar, matched by suffix.""" + matches = sorted(reference_dir.glob(f'**/*_{scalar}.dscalar.nii')) + if not matches: + raise ValueError( + f'No reference files matching *_{scalar}.dscalar.nii under {reference_dir}.' + ) + return matches + + +def _reference_profile(reference_dir: Path, scalar: str): + """Derive the aggregate per-vertex generation profile for one scalar. + + Returns the reference brain-model axis plus per-vertex mean and standard + deviation, the shared medial-wall zero mask, and the observed value range. + Only these aggregates are read; individual subject rows are never emitted. + """ + import nibabel as nb + import numpy as np + + files = _reference_files(reference_dir, scalar) + stack = [] + brain_axis = None + for path in files: + image = nb.load(str(path)) + if brain_axis is None: + brain_axis = image.header.get_axis(1) + stack.append(np.asarray(image.dataobj).reshape(-1).astype(np.float64)) + matrix = np.vstack(stack) + finite = np.isfinite(matrix) + # A vertex is medial (constant zero) when it is zero in every reference subject. + zero_mask = np.all(matrix == 0.0, axis=0) + with np.errstate(invalid='ignore'): + mean = np.where(finite, matrix, np.nan) + per_vertex_mean = np.nanmean(mean, axis=0) + per_vertex_std = np.nan_to_num(np.nanstd(mean, axis=0), nan=0.0) + per_vertex_mean = np.nan_to_num(per_vertex_mean, nan=0.0) + observed = matrix[finite] + return { + 'brain_axis': brain_axis, + 'mean': per_vertex_mean.astype(np.float32), + 'std': per_vertex_std.astype(np.float32), + 'zero_mask': zero_mask, + 'vmin': float(observed.min()), + 'vmax': float(observed.max()), + 'reference_count': len(files), + } + + +def _marginal_row(scalar: str, n_elements: int, medial_mask, rng) -> object: + """Draw one synthetic row from a scalar's marginal profile (no reference).""" + import numpy as np + + mean, std, vmin, vmax = SCALAR_MARGINALS.get(scalar, (0.0, 1.0, -6.0, 6.0)) + values = rng.normal(mean, std, size=n_elements).astype(np.float32) + np.clip(values, vmin, vmax, out=values) + values[medial_mask] = 0.0 + return values.reshape(1, n_elements) + + +def _reference_row(profile, rng) -> object: + """Draw one synthetic row from a reference profile; never a real row.""" + import numpy as np + + mean = profile['mean'] + std = profile['std'] + noise = rng.standard_normal(mean.shape[0]).astype(np.float32) + values = mean + std * noise + np.clip(values, profile['vmin'], profile['vmax'], out=values) + values[profile['zero_mask']] = 0.0 + return values.reshape(1, mean.shape[0]).astype(np.float32) + + +def synthesize(args: argparse.Namespace) -> int: + """Write the seeded cohort and return its row count.""" + import nibabel as nb + import numpy as np + from nibabel.cifti2.cifti2_axes import ScalarAxis + + _validate_common(args.subjects) + if not args.scalars: + raise ValueError('--scalars must name at least one scalar') + + output_dir = args.output_dir + output_dir.mkdir(parents=True, exist_ok=True) + cohort_path = args.cohort or (output_dir / 'cohort.csv') + + reference_dir = getattr(args, 'reference_dir', None) + profiles: dict[str, dict] = {} + if reference_dir is not None: + for scalar in args.scalars: + profiles[scalar] = _reference_profile(reference_dir, scalar) + n_elements = profiles[args.scalars[0]]['mean'].shape[0] + if any(profile['mean'].shape[0] != n_elements for profile in profiles.values()): + raise ValueError('Reference scalars have differing greyordinate counts.') + medial_mask = None + density = 'reference' + summary = ', '.join(f'{s}:{profiles[s]["reference_count"]}' for s in args.scalars) + print(f'Derived realistic profiles from {reference_dir} ({summary} files).') + else: + brain_axis = _build_brain_model_axis(args.grayordinates) + n_elements = len(brain_axis) + # A deterministic contiguous medial-wall block, matching the real zero fraction. + medial_count = round(n_elements * DEFAULT_MEDIAL_WALL_FRACTION) + medial_mask = np.zeros(n_elements, dtype=bool) + medial_mask[:medial_count] = True + density = 'den-164k' if args.grayordinates == DEFAULT_GRAYORDINATES else 'synthetic' + + rng = np.random.default_rng(args.seed) + + rows: list[dict[str, str]] = [] + for subject_index in range(args.subjects): + subject_id = f'sub-synth{subject_index + 1:04d}' + subject_dir = output_dir / subject_id + subject_dir.mkdir(parents=True, exist_ok=True) + for scalar in args.scalars: + # The scalar becomes the CIFTI map name, whose length shifts vox_offset + # exactly as real per-metric dscalars do, exercising the per-scalar + # raw-read plan rather than one shared offset. + if reference_dir is not None: + brain_axis = profiles[scalar]['brain_axis'] + values = _reference_row(profiles[scalar], rng) + else: + values = _marginal_row(scalar, n_elements, medial_mask, rng) + header = nb.cifti2.Cifti2Header.from_axes((ScalarAxis([scalar]), brain_axis)) + path = subject_dir / f'{subject_id}_space-fsLR_den-164k_{scalar}.dscalar.nii' + nb.Cifti2Image(values, header=header).to_filename(str(path)) + rows.append( + { + 'subject_id': subject_id, + 'scalar_name': scalar, + 'source_file': str(path.resolve()), + } + ) + + common.write_csv(cohort_path, rows, common.LONG_COHORT_FIELDS) + print( + f'Wrote {len(rows)} rows ({args.subjects} subjects x {len(args.scalars)} scalars, ' + f'{n_elements} greyordinates, values={density}) to {cohort_path}' + ) + return len(rows) + + +def _read_long_cohort(path: Path) -> list[dict[str, str]]: + with path.open(newline='', encoding='utf-8') as stream: + reader = csv.DictReader(stream) + fields = set(reader.fieldnames or []) + if not {'scalar_name', 'source_file'} <= fields: + raise ValueError( + f'{path} is not a long-format CIFTI cohort; it must have ' + "'scalar_name' and 'source_file' columns." + ) + return list(reader) + + +def _subject_key(row: dict[str, str]) -> str: + """Return a stable subject identifier, tolerating cohorts without a session column.""" + subject = row.get('subject_id') or '' + session = row.get('session_id') or '' + return f'{subject}|{session}' if session else subject + + +def subset_cohort(args: argparse.Namespace) -> int: + """Write the deterministic subset and return its row count.""" + _validate_common(args.subjects) + rows = _read_long_cohort(args.input) + if not rows: + raise ValueError(f'{args.input} has no rows.') + + available_scalars = list(dict.fromkeys(row['scalar_name'] for row in rows)) + scalars = args.scalars or available_scalars + missing = [scalar for scalar in scalars if scalar not in available_scalars] + if missing: + raise ValueError(f'Scalars not present in {args.input}: {", ".join(missing)}') + + # Subjects in first-seen order, then keep the first --subjects of them. + ordered_subjects = list(OrderedDict.fromkeys(_subject_key(row) for row in rows)) + if len(ordered_subjects) < args.subjects: + raise ValueError( + f'Requested {args.subjects} subjects but {args.input} has only ' + f'{len(ordered_subjects)}.' + ) + keep_subjects = set(ordered_subjects[: args.subjects]) + keep_scalars = set(scalars) + + selected = [ + { + 'subject_id': row.get('subject_id', ''), + 'scalar_name': row['scalar_name'], + 'source_file': row['source_file'], + } + for row in rows + if row['scalar_name'] in keep_scalars and _subject_key(row) in keep_subjects + ] + common.write_csv(args.output, selected, common.LONG_COHORT_FIELDS) + print( + f'Wrote {len(selected)} rows ({args.subjects} subjects x {len(scalars)} scalars) ' + f'to {args.output}' + ) + return len(selected) + + +def main() -> int: + args = _parser().parse_args() + args.func(args) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/benchmarks/to_modelarray/make_cohort.py b/benchmarks/to_modelarray/make_cohort.py new file mode 100644 index 0000000..af4e61f --- /dev/null +++ b/benchmarks/to_modelarray/make_cohort.py @@ -0,0 +1,129 @@ +"""Create deterministic wide-format ABIDE cohort CSVs from public S3 keys.""" + +from __future__ import annotations + +import argparse +import re +from pathlib import Path + +import harness_common as common + +DEFAULT_BUCKET = 'fcp-indi' +DEFAULT_PREFIX = 'data/Projects/ABIDE_Initiative/Outputs/cpac/filt_global' +DEFAULT_DERIVATIVES = ('alff',) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--output', type=Path, required=True) + parser.add_argument('--subjects', type=int, default=100) + parser.add_argument('--derivatives', nargs='+', default=list(DEFAULT_DERIVATIVES)) + parser.add_argument('--repeat-factor', type=int, default=1) + parser.add_argument('--bucket', default=DEFAULT_BUCKET) + parser.add_argument('--prefix', default=DEFAULT_PREFIX) + parser.add_argument( + '--subject-manifest', + type=Path, + help='Optional newline-delimited subject IDs; avoids listing S3 and is useful offline.', + ) + return parser + + +def _validate_args(args: argparse.Namespace) -> None: + if args.subjects < 1: + raise ValueError('--subjects must be positive') + if args.repeat_factor < 1: + raise ValueError('--repeat-factor must be positive') + for derivative in args.derivatives: + if not re.fullmatch(r'[A-Za-z0-9_]+', derivative): + raise ValueError(f'Invalid derivative name: {derivative!r}') + + +def _unsigned_s3_client(): + try: + import boto3 + from botocore import UNSIGNED + from botocore.config import Config + except ImportError as exc: + raise RuntimeError( + 'Listing the public ABIDE bucket requires boto3; install modelarrayio[s3] ' + 'or pass --subject-manifest.' + ) from exc + return boto3.client('s3', config=Config(signature_version=UNSIGNED)) + + +def _subjects_for_derivative(client, bucket: str, prefix: str, derivative: str) -> set[str]: + key_prefix = f'{prefix.rstrip("/")}/{derivative}/' + suffix = f'_{derivative}.nii.gz' + subjects: set[str] = set() + paginator = client.get_paginator('list_objects_v2') + for page in paginator.paginate(Bucket=bucket, Prefix=key_prefix): + for item in page.get('Contents', []): + filename = str(item['Key']).rsplit('/', 1)[-1] + if filename.endswith(suffix): + subjects.add(filename[: -len(suffix)]) + return subjects + + +def _available_subjects(args: argparse.Namespace) -> list[str]: + if args.subject_manifest is not None: + subjects = { + line.strip() + for line in args.subject_manifest.read_text(encoding='utf-8').splitlines() + if line.strip() and not line.lstrip().startswith('#') + } + return sorted(subjects) + + client = _unsigned_s3_client() + required = [ + _subjects_for_derivative(client, args.bucket, args.prefix, derivative) + for derivative in [*args.derivatives, 'func_mask'] + ] + return sorted(set.intersection(*required)) if required else [] + + +def _s3_uri(bucket: str, prefix: str, derivative: str, subject: str) -> str: + return f's3://{bucket}/{prefix.rstrip("/")}/{derivative}/{subject}_{derivative}.nii.gz' + + +def make_cohort(args: argparse.Namespace) -> int: + """Write the requested deterministic cohort and return its row count.""" + _validate_args(args) + available = _available_subjects(args) + if len(available) < args.subjects: + raise ValueError( + f'Requested {args.subjects} subjects, but only {len(available)} have every ' + f'requested derivative and func_mask.' + ) + selected = available[: args.subjects] + rows = [] + for repeat_index in range(args.repeat_factor): + for subject in selected: + row = { + 'subject_id': ( + subject + if args.repeat_factor == 1 + else f'{subject}__repeat-{repeat_index + 1:04d}' + ), + 'source_mask_file': _s3_uri(args.bucket, args.prefix, 'func_mask', subject), + } + row.update( + { + derivative: _s3_uri(args.bucket, args.prefix, derivative, subject) + for derivative in args.derivatives + } + ) + rows.append(row) + common.write_csv(args.output, rows, ['subject_id', 'source_mask_file', *args.derivatives]) + return len(rows) + + +def main() -> int: + args = _parser().parse_args() + row_count = make_cohort(args) + print(f'Wrote {row_count} rows to {args.output}') + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/benchmarks/to_modelarray/make_mif_cohort.py b/benchmarks/to_modelarray/make_mif_cohort.py new file mode 100644 index 0000000..3f66783 --- /dev/null +++ b/benchmarks/to_modelarray/make_mif_cohort.py @@ -0,0 +1,113 @@ +"""Fetch and normalize the public MIF (fixel) benchmark cohort. + +The ModelArray walkthrough ships a fully public fixel demo — 100 FDC ``.mif`` +subjects plus ``index.mif`` and ``directions.mif`` — hosted on OSF. It is the +MIF analogue of the public ABIDE NIfTI cohort: real diffusion fixel data whose +runtime, ``--workers`` scaling and compression ratios are all quotable. + +This helper downloads and extracts that archive (unless it is already present) +and writes a normalized long-format ``subject_id,scalar_name,source_file`` +cohort CSV with absolute paths, which ``run_benchmarks.py`` drives directly with +``--index-file`` and ``--directions-file``. The archive already contains a +long-format ``cohort_FDC_n100.csv``; normalizing it here just resolves the +relative source paths and drops the extra demographic columns so the harness and +report see a clean cohort. + +Reference: https://pennlinc.github.io/ModelArray/articles/walkthrough.html +""" + +from __future__ import annotations + +import argparse +import tarfile +import urllib.request +from pathlib import Path + +import harness_common as common + +DEFAULT_URL = 'https://osf.io/download/tce9d' +SCALAR_NAME = 'FDC' + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + '--data-dir', + type=Path, + default=Path('benchmarks/to_modelarray/data/mif-walkthrough'), + help='Directory holding (or to receive) the extracted FDC archive.', + ) + parser.add_argument( + '--output', + type=Path, + default=Path('benchmarks/to_modelarray/cohorts/mif-walkthrough.csv'), + help='Normalized long-format cohort CSV to write.', + ) + parser.add_argument('--url', default=DEFAULT_URL, help='OSF download URL for the FDC archive.') + parser.add_argument( + '--skip-download', + action='store_true', + help='Use the already-extracted FDC directory without downloading.', + ) + return parser + + +def _download_and_extract(url: str, data_dir: Path) -> None: + data_dir.mkdir(parents=True, exist_ok=True) + archive = data_dir / 'download.tar.gz' + print(f'Downloading FDC demo archive from {url} ...') + with urllib.request.urlopen(url) as response, archive.open('wb') as handle: # noqa: S310 + while chunk := response.read(1024 * 1024): + handle.write(chunk) + print(f'Extracting {archive} ...') + with tarfile.open(archive) as tar: + tar.extractall(data_dir, filter='data') + archive.unlink() + + +def _find_fdc_dir(data_dir: Path) -> Path: + for candidate in (data_dir / 'FDC', *sorted(data_dir.glob('**/FDC'))): + if (candidate / 'index.mif').exists() and (candidate / 'directions.mif').exists(): + return candidate + raise FileNotFoundError( + f'Could not find an FDC directory with index.mif and directions.mif under {data_dir}.' + ) + + +def _subject_rows(fdc_dir: Path) -> list[dict[str, str]]: + rows = [ + { + 'subject_id': mif_file.stem, + 'scalar_name': SCALAR_NAME, + 'source_file': str(mif_file.resolve()), + } + for mif_file in sorted(fdc_dir.glob('sub-*.mif')) + ] + if not rows: + raise FileNotFoundError(f'No sub-*.mif files found in {fdc_dir}.') + return rows + + +def make_mif_cohort(args: argparse.Namespace) -> int: + """Write the normalized MIF cohort and return its row count.""" + if not args.skip_download and not (args.data_dir / 'FDC' / 'index.mif').exists(): + _download_and_extract(args.url, args.data_dir) + fdc_dir = _find_fdc_dir(args.data_dir) + rows = _subject_rows(fdc_dir) + common.write_csv(args.output, rows, common.LONG_COHORT_FIELDS) + print( + f'Wrote {len(rows)} rows to {args.output}\n' + f' --index-file {fdc_dir / "index.mif"}\n' + f' --directions-file {fdc_dir / "directions.mif"}' + ) + return len(rows) + + +def main() -> int: + args = _parser().parse_args() + make_mif_cohort(args) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/benchmarks/to_modelarray/materialize_cohort.py b/benchmarks/to_modelarray/materialize_cohort.py new file mode 100644 index 0000000..dbae63d --- /dev/null +++ b/benchmarks/to_modelarray/materialize_cohort.py @@ -0,0 +1,97 @@ +"""Download every ``s3://`` source in a cohort CSV and write a local-path cohort. + +Public benchmark cohorts (``make_cohort.py``) reference ``s3://fcp-indi`` keys. +Cluster compute nodes may have no internet, so the files are fetched once on a +login node with unsigned requests and the cohort is rewritten to absolute local +paths. Existing files are kept, so re-running is cheap. Every ``s3://`` cell in +any column is materialized, which covers scalar columns and +``source_mask_file`` alike. + +Example:: + + python benchmarks/to_modelarray/materialize_cohort.py \ + --cohort benchmarks/to_modelarray/cohorts/abide-small.csv \ + --data-dir /cbica/home/tongt/modelarray_opt/data/abide \ + --output /cbica/home/tongt/modelarray_opt/cohorts/abide-small-local.csv +""" + +from __future__ import annotations + +import argparse +import csv +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from urllib.parse import urlparse + +import harness_common as common + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--cohort', type=Path, required=True) + parser.add_argument('--data-dir', type=Path, required=True) + parser.add_argument('--output', type=Path, required=True) + parser.add_argument('--workers', type=int, default=8) + parser.add_argument( + '--extra-key', + action='append', + default=[], + help='Additional s3:// object to download (for example the group mask).', + ) + return parser + + +def local_path_for(uri: str, data_dir: Path) -> Path: + """Map ``s3://bucket/a/b/derivative/file`` to ``data_dir/derivative/file``.""" + parsed = urlparse(uri) + key = Path(parsed.path.lstrip('/')) + return data_dir / key.parent.name / key.name + + +def _download(client, uri: str, target: Path) -> None: + if target.exists() and target.stat().st_size > 0: + return + parsed = urlparse(uri) + target.parent.mkdir(parents=True, exist_ok=True) + temporary = target.with_suffix(target.suffix + '.part') + client.download_file(parsed.netloc, parsed.path.lstrip('/'), str(temporary)) + temporary.replace(target) + + +def materialize(args: argparse.Namespace, client=None) -> int: + """Download all s3 cells (and extra keys), write the local cohort, return file count.""" + with args.cohort.open(newline='', encoding='utf-8') as stream: + reader = csv.DictReader(stream) + fieldnames = list(reader.fieldnames or []) + rows = list(reader) + + uris = {value for row in rows for value in row.values() if value and value.startswith('s3://')} + uris.update(args.extra_key) + if client is None and uris: + import boto3 + from botocore import UNSIGNED + from botocore.config import Config + + client = boto3.client('s3', config=Config(signature_version=UNSIGNED)) + + targets = {uri: local_path_for(uri, args.data_dir) for uri in sorted(uris)} + with ThreadPoolExecutor(max_workers=max(1, args.workers)) as pool: + list(pool.map(lambda item: _download(client, item[0], item[1]), targets.items())) + + for row in rows: + for column, value in row.items(): + if value and value.startswith('s3://'): + row[column] = str(targets[value].resolve()) + common.write_csv(args.output, rows, fieldnames) + return len(targets) + + +def main() -> int: + args = _parser().parse_args() + count = materialize(args) + print(f'Materialized {count} objects under {args.data_dir}; wrote {args.output}') + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/benchmarks/to_modelarray/output_compare.py b/benchmarks/to_modelarray/output_compare.py new file mode 100644 index 0000000..7fd2ef9 --- /dev/null +++ b/benchmarks/to_modelarray/output_compare.py @@ -0,0 +1,118 @@ +"""Stream-compare ModelArray outputs (an HDF5 file or a TileDB root) one column block at a time. + +Shared by ``compare_revisions.py`` (original vs optimized output), ``compare_to_reference.py`` +(optimized output vs a deployed reference) and ``create_comparison_report.py`` (output shapes). +``h5py`` and ``tiledb`` are imported only when an output is opened, because libtiledb cannot +load on CPUs without AVX2 (the CUBIC login node) and the scripts' ``--help`` must still work there. +""" + +from __future__ import annotations + +import json +from collections.abc import Callable +from pathlib import Path +from typing import Self + +import numpy as np + +BlockReader = Callable[[int, int], np.ndarray] + + +def names_from_meta(array) -> list[str]: + """Return the ``column_names`` a TileDB values array stores in its metadata.""" + raw = array.meta.get('column_names') + return [str(value) for value in json.loads(raw)] if raw is not None else [] + + +def _hdf5_column_names(handle, scalar: str) -> list[str]: + group = handle[f'scalars/{scalar}'] + if 'column_names' in group: + return [str(value) for value in group['column_names'].asstr()[:]] + # Older outputs stored the names as an attribute of the values dataset. + return [ + value.decode() if isinstance(value, bytes) else str(value) + for value in group['values'].attrs.get('column_names', []) + ] + + +class OutputReader: + """Column-block reader over one scalar of an HDF5 file or a TileDB root.""" + + def __init__(self, path: Path, scalar: str): + self._h5 = None + self._array = None + if path.is_file(): + import h5py + + self._h5 = h5py.File(path, 'r') + self._dataset = self._h5[f'scalars/{scalar}/values'] + self.shape = tuple(int(value) for value in self._dataset.shape) + self.dtype = np.dtype(self._dataset.dtype) + self.names = _hdf5_column_names(self._h5, scalar) + else: + import tiledb + + self._array = tiledb.open(str(path / 'scalars' / scalar / 'values'), 'r') + self.shape = tuple(int(value) for value in self._array.shape) + self.dtype = np.dtype(self._array.dtype) + self.names = names_from_meta(self._array) + + def block(self, start: int, stop: int) -> np.ndarray: + if self._h5 is not None: + return np.asarray(self._dataset[:, start:stop]) + return np.asarray(self._array[:, start:stop]['values']) + + def close(self) -> None: + if self._h5 is not None: + self._h5.close() + if self._array is not None: + self._array.close() + + def __enter__(self) -> Self: + return self + + def __exit__(self, *exc_info: object) -> None: + self.close() + + +def compare_blocks( + left: BlockReader, right: BlockReader, n_columns: int, block_columns: int = 4096 +) -> tuple[bool, bool, float]: + """Return ``(nan_equal, values_equal, max_abs_difference)`` over equal-shape matrices.""" + nan_equal = values_equal = True + max_abs_difference = 0.0 + for start in range(0, n_columns, block_columns): + stop = min(start + block_columns, n_columns) + left_block = left(start, stop) + right_block = right(start, stop) + left_nan = np.isnan(left_block) + right_nan = np.isnan(right_block) + nan_equal = nan_equal and np.array_equal(left_nan, right_nan) + finite = ~left_nan & ~right_nan + if np.any(finite): + difference = np.max(np.abs(left_block[finite] - right_block[finite])) + max_abs_difference = max(max_abs_difference, float(difference)) + values_equal = values_equal and np.array_equal(left_block, right_block, equal_nan=True) + return nan_equal, values_equal, max_abs_difference + + +def compare_outputs( + left: Path, right: Path, scalar: str, block_columns: int = 4096 +) -> dict[str, object]: + """Compare one scalar of two outputs written by the same backend.""" + with OutputReader(left, scalar) as left_reader, OutputReader(right, scalar) as right_reader: + shape_equal = left_reader.shape == right_reader.shape + nan_equal = values_equal = False + max_abs_difference = 0.0 + if shape_equal: + nan_equal, values_equal, max_abs_difference = compare_blocks( + left_reader.block, right_reader.block, left_reader.shape[1], block_columns + ) + return { + 'shape_equal': shape_equal, + 'dtype_equal': left_reader.dtype == right_reader.dtype, + 'sources_equal': left_reader.names == right_reader.names, + 'nan_equal': nan_equal, + 'values_equal': values_equal, + 'max_abs_difference': max_abs_difference, + } diff --git a/benchmarks/to_modelarray/plot_results.R b/benchmarks/to_modelarray/plot_results.R new file mode 100644 index 0000000..7207c1b --- /dev/null +++ b/benchmarks/to_modelarray/plot_results.R @@ -0,0 +1,51 @@ +#!/usr/bin/env Rscript + +args <- commandArgs(trailingOnly = TRUE) +results_root <- if (length(args) >= 1L) args[[1L]] else Sys.getenv( + "MODELARRAYIO_BENCHMARK_RESULTS_DIR", "benchmarks/to_modelarray/results" +) +plot_dir <- if (length(args) >= 2L) args[[2L]] else file.path(results_root, "plots") +dir.create(plot_dir, recursive = TRUE, showWarnings = FALSE) + +files <- list.files(results_root, pattern = "^results(?:\\.[^.]+)?\\.csv$", + recursive = TRUE, full.names = TRUE) +if (!length(files)) stop("No benchmark result CSVs found") + +tables <- lapply(files, function(path) read.csv(path, stringsAsFactors = FALSE)) +all_names <- unique(unlist(lapply(tables, names))) +tables <- lapply(tables, function(table) { + for (name in setdiff(all_names, names(table))) table[[name]] <- NA + table[all_names] +}) +results <- do.call(rbind, tables) +results <- results[results$status == "passed" & results$run_kind == "end_to_end", , drop = FALSE] +if ("r_gate" %in% names(results)) { + gate_failed <- !is.na(results$r_gate) & grepl("^failed", results$r_gate) + results <- results[!gate_failed, , drop = FALSE] +} +if (!nrow(results)) stop("No passing end-to-end benchmark rows found") + +png(file.path(plot_dir, "runtime_by_workers.png"), width = 1200, height = 800, res = 130) +groups <- interaction(results$backend, results$dtype, results$compression, drop = TRUE) +colours <- seq_along(levels(groups)) +plot(results$workers, results$wall_seconds, log = "xy", pch = 19, + col = colours[groups], xlab = "Loader workers", ylab = "Wall time (seconds)", + main = "to-modelarray end-to-end runtime") +legend("topright", legend = levels(groups), col = colours, pch = 19, cex = 0.8) +dev.off() + +phase_names <- intersect( + c("source_load_seconds", "validation_seconds", "staging_allocation_seconds", + "compression_seconds", "submission_seconds"), + names(results) +) +if (length(phase_names)) { + phase_values <- sapply(phase_names, function(name) { + values <- suppressWarnings(as.numeric(results[[name]])) + mean(values, na.rm = TRUE) + }) + phase_values[!is.finite(phase_values)] <- 0 + png(file.path(plot_dir, "mean_phase_seconds.png"), width = 1200, height = 800, res = 130) + barplot(phase_values, las = 2, ylab = "Mean seconds", main = "Mean recorded phase time") + dev.off() +} diff --git a/benchmarks/to_modelarray/r_read_gate.R b/benchmarks/to_modelarray/r_read_gate.R new file mode 100644 index 0000000..31a59d7 --- /dev/null +++ b/benchmarks/to_modelarray/r_read_gate.R @@ -0,0 +1,43 @@ +#!/usr/bin/env Rscript + +user_lib <- Sys.getenv("MODELARRAYIO_R_LIBS", "") +if (nzchar(user_lib) && dir.exists(user_lib)) .libPaths(c(user_lib, .libPaths())) +args <- commandArgs(trailingOnly = TRUE) +if (length(args) != 2L) { + stop("usage: r_read_gate.R OUTPUT BACKEND") +} + +output <- args[[1L]] +backend <- args[[2L]] +started <- proc.time()[["elapsed"]] + +if (backend == "hdf5") { + if (!requireNamespace("rhdf5", quietly = TRUE) || + !requireNamespace("HDF5Array", quietly = TRUE) || + !requireNamespace("ModelArray", quietly = TRUE)) { + stop("The HDF5 gate requires rhdf5, HDF5Array, and ModelArray") + } + listing <- rhdf5::h5ls(output, recursive = TRUE) + scalar_rows <- listing$name == "values" & grepl("^/scalars/[^/]+$", listing$group) + scalar_names <- sub("^/scalars/", "", listing$group[scalar_rows]) + if (!length(scalar_names)) stop("No scalar matrices found") + modelarray <- ModelArray::ModelArray(output, scalar_types = scalar_names) + scalar <- ModelArray::scalars(modelarray)[[scalar_names[[1L]]]] + invisible(as.matrix(scalar[seq_len(min(64L, nrow(scalar))), + seq_len(min(8L, ncol(scalar))), drop = FALSE])) +} else if (backend == "tiledb") { + if (!requireNamespace("tiledb", quietly = TRUE)) { + stop("The TileDB gate requires the tiledb R package") + } + scalar_roots <- list.dirs(file.path(output, "scalars"), recursive = FALSE) + value_uris <- file.path(scalar_roots, "values") + value_uris <- value_uris[file.exists(value_uris)] + if (!length(value_uris)) stop("No TileDB scalar matrices found") + array <- tiledb::tiledb_array(value_uris[[1L]], query_type = "READ") + invisible(array[0L, 0L]) +} else { + stop(sprintf("Unknown backend: %s", backend)) +} + +elapsed <- proc.time()[["elapsed"]] - started +cat(sprintf("read_seconds=%.6f\n", elapsed)) diff --git a/benchmarks/to_modelarray/reports/to-modelarray-cubic-benchmark.pdf b/benchmarks/to_modelarray/reports/to-modelarray-cubic-benchmark.pdf new file mode 100644 index 0000000..47a8733 Binary files /dev/null and b/benchmarks/to_modelarray/reports/to-modelarray-cubic-benchmark.pdf differ diff --git a/benchmarks/to_modelarray/run_benchmarks.py b/benchmarks/to_modelarray/run_benchmarks.py new file mode 100644 index 0000000..5d7ce66 --- /dev/null +++ b/benchmarks/to_modelarray/run_benchmarks.py @@ -0,0 +1,733 @@ +"""Run reproducible end-to-end ``to-modelarray`` benchmark sweeps.""" + +from __future__ import annotations + +import argparse +import csv +import itertools +import json +import os +import platform +import re +import resource +import shlex +import subprocess +import time +from datetime import UTC, datetime +from pathlib import Path +from urllib.parse import urlparse + +import harness_common as common +import source_identity + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--cohort', type=Path, required=True) + parser.add_argument('--mask', help='Group mask for NIfTI cohorts (local path or S3 URI).') + parser.add_argument('--index-file', help='Index MIF for fixel cohorts (local path or S3 URI).') + parser.add_argument( + '--directions-file', help='Directions MIF for fixel cohorts (local path or S3 URI).' + ) + parser.add_argument('--scalar-columns', nargs='+') + parser.add_argument('--tier', choices=['quick', 'medium', 'full'], default='quick') + parser.add_argument( + '--backends', nargs='+', choices=['hdf5', 'tiledb'], default=['hdf5', 'tiledb'] + ) + parser.add_argument('--workers', nargs='+', type=int) + parser.add_argument('--dtypes', nargs='+', choices=['float32', 'float64']) + parser.add_argument('--repetitions', type=int, default=1) + parser.add_argument('--target-chunk-mb', nargs='+', type=float) + parser.add_argument( + '--compression-workers', + nargs='+', + type=int, + help='Internal HDF5 compression-pool sizes to sweep (not a public CLI option).', + ) + parser.add_argument('--acceptance-files-per-minute', type=float) + parser.add_argument('--run-id', default=datetime.now(UTC).strftime('%Y%m%dT%H%M%SZ')) + parser.add_argument('--results-dir', type=Path) + parser.add_argument('--modelarrayio', default='modelarrayio') + parser.add_argument( + '--cache-state', choices=['warm', 'cold', 'unspecified'], default='unspecified' + ) + parser.add_argument('--repeat-factor', type=int, default=1) + split_group = parser.add_mutually_exclusive_group() + split_group.add_argument('--split-files', dest='split_files', action='store_true') + split_group.add_argument('--no-split-files', dest='split_files', action='store_false') + parser.set_defaults(split_files=None) + parser.add_argument('--skip-input-bytes', action='store_true') + parser.add_argument('--result-shard', help='Write results..csv for parallel sweeps.') + parser.add_argument( + '--r-gate-command', + help=( + 'Optional command template run for every output. Available fields: ' + '{output}, {backend}, and {scalars}.' + ), + ) + return parser + + +def _validated_shard(value: str | None) -> str | None: + """Return a path-safe shard identifier.""" + if value is None: + return None + if not re.fullmatch(r'[A-Za-z0-9][A-Za-z0-9_.-]*', value): + raise ValueError( + '--result-shard must start with an alphanumeric character and contain only ' + 'letters, digits, underscores, dots, or hyphens.' + ) + return value + + +def _validated_run_id(value: str) -> str: + """Return a path-safe benchmark run identifier.""" + if not re.fullmatch(r'[A-Za-z0-9][A-Za-z0-9_.-]*', value): + raise ValueError( + '--run-id must start with an alphanumeric character and contain only ' + 'letters, digits, underscores, dots, or hyphens.' + ) + return value + + +def _run_artifact_paths(run_dir: Path, shard: str | None) -> tuple[Path, Path, Path]: + """Return isolated artifact, result, and metadata paths for one shard.""" + artifact_dir = run_dir / f'shard-{shard}' if shard else run_dir + result_name = f'results.{shard}.csv' if shard else 'results.csv' + meta_name = f'run_meta.{shard}.json' if shard else 'run_meta.json' + return artifact_dir, run_dir / result_name, run_dir / meta_name + + +def _available_cpu_count() -> int: + """Return CPUs allocated to this process rather than the host total.""" + process_cpu_count = getattr(os, 'process_cpu_count', None) + if process_cpu_count is not None: + count = process_cpu_count() + if count: + return max(1, int(count)) + + sched_getaffinity = getattr(os, 'sched_getaffinity', None) + if sched_getaffinity is not None: + try: + return max(1, len(sched_getaffinity(0))) + except OSError: + pass + + return max(1, os.cpu_count() or 1) + + +def _machine_metadata() -> dict[str, object]: + metadata: dict[str, object] = { + 'platform': platform.platform(), + 'allocated_cpus': _available_cpu_count(), + 'versions': common.package_versions(), + } + try: + page_size = os.sysconf('SC_PAGE_SIZE') + page_count = os.sysconf('SC_PHYS_PAGES') + metadata['physical_memory_bytes'] = int(page_size * page_count) + except (AttributeError, OSError, ValueError): + metadata['physical_memory_bytes'] = None + try: + import tiledb + + config = tiledb.Config() + metadata['tiledb_compute_concurrency_level'] = int(config['sm.compute_concurrency_level']) + metadata['tiledb_io_concurrency_level'] = int(config['sm.io_concurrency_level']) + except (ImportError, KeyError, TypeError, ValueError): + metadata['tiledb_compute_concurrency_level'] = None + metadata['tiledb_io_concurrency_level'] = None + return metadata + + +def _source_paths(rows: list[dict[str, str]], scalar_columns: list[str]) -> list[str]: + paths: list[str] = [] + for row in rows: + if scalar_columns: + paths.extend(row[column] for column in scalar_columns if row.get(column)) + elif row.get('source_file'): + paths.append(row['source_file']) + if row.get('source_mask_file'): + paths.append(row['source_mask_file']) + return paths + + +def _resolve_split_files(requested: bool | None, scalar_columns: list[str]) -> bool: + """Mirror the converter's wide-split/long-combined default.""" + return bool(scalar_columns) if requested is None else requested + + +def _resolved_input(value: str | None) -> str | None: + if value is None or value.startswith('s3://'): + return value + return str(Path(value).resolve()) + + +def _resolved_executable(value: str) -> str: + """Resolve an explicitly pathed executable before converter cwd changes.""" + candidate = Path(value).expanduser() + if candidate.exists() or candidate.parent != Path('.'): + return str(candidate.resolve()) + return value + + +def _input_bytes( + paths: list[str], auxiliary_paths: list[str], *, source_root: Path +) -> tuple[int | None, int]: + sizes: dict[str, int] = {} + s3_client = None + total = 0 + all_paths = [*auxiliary_paths, *paths] + for source in all_paths: + if source not in sizes: + if source.startswith('s3://'): + if s3_client is None: + import boto3 + from botocore import UNSIGNED + from botocore.config import Config + + s3_client = boto3.client('s3', config=Config(signature_version=UNSIGNED)) + parsed = urlparse(source) + response = s3_client.head_object(Bucket=parsed.netloc, Key=parsed.path.lstrip('/')) + sizes[source] = int(response['ContentLength']) + else: + try: + local_source = Path(source) + if not local_source.is_absolute(): + local_source = source_root / local_source + sizes[source] = local_source.stat().st_size + except OSError: + return None, len(all_paths) + total += sizes[source] + return total, len(all_paths) + + +def _case_grid(args: argparse.Namespace) -> list[dict[str, object]]: + available_cpus = _available_cpu_count() + default_workers = { + 'quick': [1], + 'medium': [1, 4, 8], + 'full': sorted({1, 2, 4, 8, available_cpus}), + }[args.tier] + workers = args.workers or default_workers + dtypes = args.dtypes or (['float32', 'float64'] if args.tier != 'quick' else ['float32']) + chunk_targets = args.target_chunk_mb or ( + [8.0, 32.0, 64.0, 128.0] if args.tier == 'full' else [2.0] + ) + cases = [] + for backend in args.backends: + if args.tier == 'full': + codecs = ( + [('gzip', 1), ('gzip', 4), ('none', 0)] + if backend == 'hdf5' + else [ + ('zstd', 1), + ('zstd', 3), + ('zstd', 5), + ('none', 0), + ] + ) + else: + codecs = [('gzip', 4)] if backend == 'hdf5' else [('zstd', 5)] + for worker_count, dtype, chunk_target, codec in itertools.product( + workers, dtypes, chunk_targets, codecs + ): + compression, level = codec + compression_workers = [None] + if backend == 'hdf5' and compression != 'none': + compression_workers = args.compression_workers or ( + sorted({1, 4, available_cpus}) if args.tier == 'full' else [None] + ) + cases.extend( + { + 'backend': backend, + 'workers': worker_count, + 'dtype': dtype, + 'target_chunk_mb': chunk_target, + 'compression': compression, + 'compression_level': level, + 'compression_workers_requested': compression_worker_count, + } + for compression_worker_count in compression_workers + ) + return cases + + +def _load_timing(path: Path) -> dict[str, object]: + if not path.exists(): + return {} + records = [json.loads(line) for line in path.read_text(encoding='utf-8').splitlines()] + scalars = [record for record in records if record.get('event') == 'scalar_complete'] + complete = next( + (record for record in reversed(records) if record.get('event') == 'conversion_complete'), + {}, + ) + aggregate: dict[str, object] = dict(complete) + for field in ( + 'source_load_seconds', + 'source_cpu_seconds', + 'validation_seconds', + 'staging_allocation_seconds', + 'writer_idle_seconds', + 'loader_blocked_seconds', + 'write_seconds', + 'compression_seconds', + 'submission_seconds', + 'fragment_count', + 'fast_path_rows', + ): + values = [ + record[field] for record in scalars if isinstance(record.get(field), (int, float)) + ] + aggregate[field] = sum(values) if values else None + aggregate['compression_workers'] = max( + (int(record.get('compression_workers', 0)) for record in scalars), default=0 + ) + return aggregate + + +def _append_csv(path: Path, row: dict[str, object]) -> None: + existing: list[dict[str, str]] = [] + fieldnames: list[str] = [] + if path.exists(): + with path.open(newline='', encoding='utf-8') as stream: + reader = csv.DictReader(stream) + fieldnames = list(reader.fieldnames or []) + existing = list(reader) + for key in row: + if key not in fieldnames: + fieldnames.append(key) + temporary = path.with_suffix(path.suffix + '.tmp') + common.write_csv(temporary, [*existing, row], fieldnames) + os.replace(temporary, path) + + +def _completed_case_keys(path: Path) -> set[str]: + """Return stable case keys already recorded by an interrupted sweep.""" + if not path.exists(): + return set() + with path.open(newline='', encoding='utf-8') as stream: + return { + row.get('case_key') or row['case_id'] + for row in csv.DictReader(stream) + if row.get('case_key') or row.get('case_id') + } + + +def _recorded_failure_count(path: Path) -> int: + """Count failed rows so a resumed run preserves its nonzero status.""" + if not path.exists(): + return 0 + with path.open(newline='', encoding='utf-8') as stream: + return sum(row.get('status') == 'failed' for row in csv.DictReader(stream)) + + +def _write_json(path: Path, value: object) -> None: + """Atomically replace a JSON metadata file.""" + temporary = path.with_suffix(path.suffix + '.tmp') + temporary.write_text(json.dumps(value, indent=2, sort_keys=True), encoding='utf-8') + os.replace(temporary, path) + + +def _run_r_gate( + template: str | None, + outputs: list[Path], + backend: str, + scalar_columns: list[str], + log_path: Path, +) -> tuple[str, float | None]: + if not template: + return 'not_run', None + started = time.perf_counter() + with log_path.open('w', encoding='utf-8') as log: + for output in outputs: + command = shlex.split( + template.format( + output=str(output), + backend=backend, + scalars=','.join(scalar_columns), + ) + ) + try: + completed = subprocess.run( + command, + stdout=log, + stderr=subprocess.STDOUT, + check=False, + ) + except (OSError, ValueError) as exc: + log.write(f'Could not execute R gate command: {exc}\n') + return 'failed', time.perf_counter() - started + if completed.returncode != 0: + return 'failed', time.perf_counter() - started + return 'passed', time.perf_counter() - started + + +def _case_failure_reasons( + *, + returncode: int, + outputs_present: bool, + r_gate: str, + meets_acceptance_target: bool | None, +) -> list[str]: + """Return every correctness or performance gate that failed for a case.""" + reasons = [] + if returncode != 0: + reasons.append('converter') + if not outputs_present: + reasons.append('missing_output') + if r_gate.startswith('failed'): + reasons.append('r_gate') + if meets_acceptance_target is False: + reasons.append('acceptance_target') + return reasons + + +def _case_status(failure_reasons: list[str], r_gate: str) -> str: + """Distinguish failed, R-verified, and smoke-only benchmark rows.""" + if failure_reasons: + return 'failed' + if r_gate == 'not_run': + return 'unverified' + return 'passed' + + +def main() -> int: + args = _parser().parse_args() + if args.repetitions < 1 or args.repeat_factor < 1: + raise ValueError('--repetitions and --repeat-factor must be positive') + result_shard = _validated_shard(args.result_shard) + run_id = _validated_run_id(args.run_id) + repo = Path(__file__).resolve().parents[2] + root = args.results_dir or Path( + os.environ.get( + 'MODELARRAYIO_BENCHMARK_RESULTS_DIR', repo / 'benchmarks' / 'to_modelarray' / 'results' + ) + ) + root = root.expanduser().resolve() + run_dir = root / run_id + run_dir.mkdir(parents=True, exist_ok=True) + artifact_dir, results_csv, meta_path = _run_artifact_paths(run_dir, result_shard) + artifact_dir.mkdir(parents=True, exist_ok=True) + + rows = common.read_cohort(args.cohort) + scalar_columns = args.scalar_columns or common.infer_scalar_columns(rows) + scalar_names = common.scalar_names(rows, scalar_columns) + split_files = _resolve_split_files(args.split_files, scalar_columns) + scalar_sources = common.scalar_source_paths(rows, scalar_columns) + modality = common.source_modality(scalar_sources) + mask = _resolved_input(args.mask) + index_file = _resolved_input(args.index_file) + directions_file = _resolved_input(args.directions_file) + if modality == 'nifti' and mask is None: + raise ValueError('--mask is required for a NIfTI benchmark cohort.') + if modality == 'mif' and (index_file is None or directions_file is None): + raise ValueError( + '--index-file and --directions-file are required for a MIF benchmark cohort.' + ) + auxiliary_paths = ( + [mask] + if modality == 'nifti' + else [index_file, directions_file] + if modality == 'mif' + else [] + ) + paths = _source_paths(rows, scalar_columns) + input_bytes, input_file_count = ( + (None, len(paths) + len(auxiliary_paths)) + if args.skip_input_bytes + else _input_bytes(paths, auxiliary_paths, source_root=args.cohort.resolve().parent) + ) + machine = _machine_metadata() + identity = source_identity.source_identity(repo) + cases = _case_grid(args) + modelarrayio_command = _resolved_executable(args.modelarrayio) + configuration = { + 'cases': cases, + 'repetitions': args.repetitions, + 'cache_state': args.cache_state, + 'acceptance_files_per_minute': args.acceptance_files_per_minute, + 'split_files': split_files, + 'skip_input_bytes': args.skip_input_bytes, + 'mask': mask, + 'index_file': index_file, + 'directions_file': directions_file, + 'r_gate_command': args.r_gate_command, + 'modelarrayio': modelarrayio_command, + } + fresh_metadata = { + 'run_id': run_id, + 'created_at': datetime.now(UTC).isoformat(), + 'run_kind': 'end_to_end', + 'result_shard': result_shard, + 'tier': args.tier, + 'modality': modality, + 'cohort': str(args.cohort.resolve()), + 'cohort_sha256': common.sha256_file(args.cohort), + 'repeat_factor': args.repeat_factor, + 'acceptance_files_per_minute': args.acceptance_files_per_minute, + 'scalar_columns': scalar_columns, + 'scalar_names': scalar_names, + 'modelarrayio_commit': identity['commit'], + 'source_tree_sha256': identity['source_tree_sha256'], + 'source_dirty': identity['dirty'], + **machine, + 'configuration': configuration, + 'commands': [], + } + if meta_path.exists(): + metadata = json.loads(meta_path.read_text(encoding='utf-8')) + identity_keys = ( + 'run_id', + 'run_kind', + 'result_shard', + 'tier', + 'modality', + 'cohort', + 'cohort_sha256', + 'repeat_factor', + 'scalar_columns', + 'scalar_names', + 'modelarrayio_commit', + 'source_tree_sha256', + ) + mismatches = [key for key in identity_keys if metadata.get(key) != fresh_metadata.get(key)] + previous_configuration = metadata.get('configuration') + if previous_configuration is not None and previous_configuration != configuration: + mismatches.append('configuration') + if mismatches: + raise ValueError( + f'Cannot resume benchmark run {run_id!r}: metadata differs for ' + f'{", ".join(mismatches)}. Choose a new --run-id.' + ) + metadata['configuration'] = configuration + metadata.setdefault('commands', []) + metadata.setdefault('resume_events', []).append( + { + 'resumed_at': datetime.now(UTC).isoformat(), + 'machine': machine, + } + ) + else: + metadata = fresh_metadata + _write_json(meta_path, metadata) + + completed_case_keys = _completed_case_keys(results_csv) + environment = os.environ.copy() + environment.setdefault('MODELARRAYIO_S3_ANON', '1') + failures = _recorded_failure_count(results_csv) + for case_index, case in enumerate(cases, start=1): + for repetition in range(1, args.repetitions + 1): + case_prefix = f'{result_shard}-' if result_shard else '' + case_key = f'{case_prefix}case-{case_index:04d}-rep-{repetition:02d}' + if case_key in completed_case_keys: + print(f'{case_key}: already recorded; skipping') + continue + + case_id = case_key + case_dir = artifact_dir / 'outputs' / case_id + retry = 1 + while case_dir.exists(): + case_id = f'{case_key}-retry-{retry:02d}' + case_dir = artifact_dir / 'outputs' / case_id + retry += 1 + case_dir.mkdir(parents=True, exist_ok=True) + output = common.output_path(case_dir, case['backend']) + timing_path = case_dir / 'timing.jsonl' + command = [ + modelarrayio_command, + 'to-modelarray', + '--cohort-file', + str(args.cohort.resolve()), + '--output', + str(output), + '--backend', + str(case['backend']), + '--workers', + str(case['workers']), + '--dtype', + str(case['dtype']), + '--compression', + str(case['compression']), + '--compression-level', + str(case['compression_level']), + '--target-chunk-mb', + str(case['target_chunk_mb']), + '--timing-jsonl', + str(timing_path), + '--split-files' if split_files else '--no-split-files', + ] + if modality == 'nifti': + command.extend(['--mask', str(mask)]) + elif modality == 'mif': + command.extend( + [ + '--index-file', + str(index_file), + '--directions-file', + str(directions_file), + ] + ) + if scalar_columns: + command.extend(['--scalar-columns', *scalar_columns]) + + stdout_path = case_dir / 'stdout.log' + stderr_path = case_dir / 'stderr.log' + started = time.perf_counter() + before_usage = resource.getrusage(resource.RUSAGE_CHILDREN) + with ( + stdout_path.open('w', encoding='utf-8') as stdout, + stderr_path.open('w', encoding='utf-8') as stderr, + ): + case_environment = environment.copy() + if case['compression_workers_requested'] is None: + case_environment.pop('MODELARRAYIO_HDF5_COMPRESSION_WORKERS', None) + else: + case_environment['MODELARRAYIO_HDF5_COMPRESSION_WORKERS'] = str( + case['compression_workers_requested'] + ) + completed = subprocess.run( + command, + cwd=args.cohort.resolve().parent, + env=case_environment, + stdout=stdout, + stderr=stderr, + check=False, + ) + wall_seconds = time.perf_counter() - started + after_usage = resource.getrusage(resource.RUSAGE_CHILDREN) + + outputs = [output] if output.exists() else sorted(case_dir.glob(f'*_{output.name}')) + output_bytes = 0 + output_files = 0 + for candidate in outputs: + candidate_bytes, candidate_files = common.output_stats(candidate) + output_bytes += candidate_bytes + output_files += candidate_files + if completed.returncode != 0: + r_gate, r_read_seconds = 'not_run_converter_failed', None + elif not outputs: + r_gate, r_read_seconds = 'failed_missing_output', None + else: + r_gate, r_read_seconds = _run_r_gate( + args.r_gate_command, + outputs, + str(case['backend']), + scalar_names, + case_dir / 'r_gate.log', + ) + timing = _load_timing(timing_path) + files_per_minute = 60.0 * input_file_count / wall_seconds + meets_acceptance_target = ( + None + if args.acceptance_files_per_minute is None + else files_per_minute >= args.acceptance_files_per_minute + ) + failure_reasons = _case_failure_reasons( + returncode=completed.returncode, + outputs_present=bool(outputs), + r_gate=r_gate, + meets_acceptance_target=meets_acceptance_target, + ) + case_status = _case_status(failure_reasons, r_gate) + case_failed = case_status == 'failed' + row = { + 'run_id': run_id, + 'case_key': case_key, + 'case_id': case_id, + 'run_kind': 'end_to_end', + 'tier': args.tier, + 'status': case_status, + 'failure_reason': ','.join(failure_reasons), + 'returncode': completed.returncode, + 'backend': case['backend'], + 'workers': case['workers'], + 'dtype': case['dtype'], + 'compression': case['compression'], + 'compression_level': case['compression_level'], + 'compression_workers': timing.get('compression_workers'), + 'compression_workers_requested': case['compression_workers_requested'], + 'target_chunk_mb': case['target_chunk_mb'], + 'cache_state': args.cache_state, + 'repeat_factor': args.repeat_factor, + 'repetition': repetition, + 'scalar_names': ','.join(scalar_names), + 'cohort_sha256': metadata['cohort_sha256'], + 'input_file_count': input_file_count, + 'input_bytes': input_bytes, + 'input_files_per_minute': files_per_minute, + 'acceptance_files_per_minute': args.acceptance_files_per_minute, + 'meets_acceptance_target': meets_acceptance_target, + 'wall_seconds': wall_seconds, + 'child_user_cpu_seconds': after_usage.ru_utime - before_usage.ru_utime, + 'child_system_cpu_seconds': after_usage.ru_stime - before_usage.ru_stime, + 'children_peak_rss_platform_units_cumulative': after_usage.ru_maxrss, + 'peak_memory_bytes': timing.get('peak_memory_bytes'), + 'output_bytes': output_bytes, + 'output_file_count': output_files, + 'r_gate': r_gate, + 'r_read_seconds': r_read_seconds, + 'modelarrayio_commit': metadata['modelarrayio_commit'], + 'source_filesystem': ( + 's3' + if paths and all(path.startswith('s3://') for path in paths) + else 'mixed' + if any(path.startswith('s3://') for path in paths) + else 'local' + ), + 'output_filesystem': 'local', + 'tiledb_compute_concurrency_level': machine.get( + 'tiledb_compute_concurrency_level' + ), + 'tiledb_io_concurrency_level': machine.get('tiledb_io_concurrency_level'), + **{ + key: timing.get(key) + for key in ( + 'source_load_seconds', + 'source_cpu_seconds', + 'validation_seconds', + 'staging_allocation_seconds', + 'writer_idle_seconds', + 'loader_blocked_seconds', + 'write_seconds', + 'compression_seconds', + 'submission_seconds', + 'fragment_count', + 'fast_path_rows', + 'peak_open_file_descriptors_sampled', + ) + }, + 'source_tree_sha256': metadata.get('source_tree_sha256'), + 'command': shlex.join(command), + } + _append_csv(results_csv, row) + metadata['commands'].append( + { + 'case_key': case_key, + 'case_id': case_id, + 'command': command, + 'environment_overrides': { + 'MODELARRAYIO_S3_ANON': case_environment.get('MODELARRAYIO_S3_ANON'), + 'MODELARRAYIO_HDF5_COMPRESSION_WORKERS': case_environment.get( + 'MODELARRAYIO_HDF5_COMPRESSION_WORKERS' + ), + }, + 'returncode': completed.returncode, + 'status': row['status'], + 'r_gate': r_gate, + 'stdout': str(stdout_path.relative_to(run_dir)), + 'stderr': str(stderr_path.relative_to(run_dir)), + } + ) + _write_json(meta_path, metadata) + completed_case_keys.add(case_key) + failures += int(case_failed) + print( + f'{case_id}: backend={case["backend"]} workers={case["workers"]} ' + f'dtype={case["dtype"]} status={row["status"]} wall={wall_seconds:.3f}s' + ) + return int(failures > 0) + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/benchmarks/to_modelarray/source_identity.py b/benchmarks/to_modelarray/source_identity.py new file mode 100644 index 0000000..afef374 --- /dev/null +++ b/benchmarks/to_modelarray/source_identity.py @@ -0,0 +1,86 @@ +"""Identify the exact source tree a benchmark ran against, committed or not. + +``git rev-parse HEAD`` is misleading for an uncommitted working tree (and for a +tree copied to a cluster with rsync): it names the last commit, not the code +that ran. This module hashes the content of every tracked-or-untracked, +non-ignored file under the paths that affect conversion results, so two hosts +holding the same working tree report the same ``source_tree_sha256`` and any +drift (a forgotten rsync, a local edit) is visible in the recorded metadata. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shutil +import subprocess +from pathlib import Path + +IDENTITY_PATHS = ('src', 'pyproject.toml', 'benchmarks') +_SKIP_DIRS = {'.venv', '__pycache__', 'data', 'results', 'cohorts', '.git'} + + +def _git(repo: Path, *args: str) -> str | None: + executable = shutil.which('git') + if executable is None or not (repo / '.git').exists(): + return None + completed = subprocess.run( + [executable, *args], cwd=repo, capture_output=True, text=True, check=False + ) + return completed.stdout if completed.returncode == 0 else None + + +def _listed_files(repo: Path) -> list[Path]: + """Return tracked + untracked (not ignored) files under the identity paths.""" + listing = _git(repo, 'ls-files', '-co', '--exclude-standard', '--', *IDENTITY_PATHS) + if listing is not None: + return sorted( + repo / line for line in listing.splitlines() if line and (repo / line).is_file() + ) + files: list[Path] = [] + for root_name in IDENTITY_PATHS: + root = repo / root_name + if root.is_file(): + files.append(root) + continue + for directory, dirnames, filenames in os.walk(root): + dirnames[:] = sorted(d for d in dirnames if d not in _SKIP_DIRS) + files.extend(Path(directory) / name for name in sorted(filenames)) + return sorted(files) + + +def source_tree_sha256(repo: Path) -> str: + """Hash sorted (relative path, bytes) pairs of every identity file.""" + digest = hashlib.sha256() + for path in _listed_files(repo): + digest.update(path.relative_to(repo).as_posix().encode('utf-8')) + digest.update(b'\0') + digest.update(path.read_bytes()) + digest.update(b'\0') + return digest.hexdigest() + + +def source_identity(repo: str | Path) -> dict[str, object]: + """Return ``commit``, ``dirty`` and ``source_tree_sha256`` for a repository.""" + root = Path(repo).resolve() + commit = _git(root, 'rev-parse', 'HEAD') + status = _git(root, 'status', '--porcelain', '--', *IDENTITY_PATHS) + return { + 'commit': commit.strip() if commit else None, + 'dirty': bool(status and status.strip()), + 'source_tree_sha256': source_tree_sha256(root), + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--repo', type=Path, default=Path(__file__).resolve().parents[2]) + args = parser.parse_args() + print(json.dumps(source_identity(args.repo), indent=2, sort_keys=True)) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/benchmarks/to_modelarray/summarize_comparison.py b/benchmarks/to_modelarray/summarize_comparison.py new file mode 100644 index 0000000..f3fab10 --- /dev/null +++ b/benchmarks/to_modelarray/summarize_comparison.py @@ -0,0 +1,54 @@ +"""Aggregate ``compare_revisions.py`` run directories into one pinned summary CSV. + +Reads the per-modality original-vs-optimized run directories and writes a single +committed summary row per modality/backend, so a comparison remains auditable and +regressions are detectable by rerunning one command. The full findings PDF is +rendered separately by ``create_comparison_report.py``. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import harness_common as common +from comparison_runs import Group, load_groups + + +def _summary_row(group: Group) -> dict[str, object]: + return { + 'modality': group.modality, + 'backend': group.backend, + 'n_subjects': group.n_subjects, + 'n_elements': group.n_elements, + 'n_scalars': group.n_scalars, + 'baseline_revision': group.baseline_revision, + 'baseline_median_wall_seconds': round(group.baseline_median, 4), + 'optimized_median_wall_seconds': round(group.optimized_median, 4), + 'speedup': round(group.speedup, 4), + 'wall_time_reduction_percent': round(group.reduction, 2), + 'baseline_median_r_read_seconds': round(group.baseline_r_median, 4), + 'optimized_median_r_read_seconds': round(group.optimized_r_median, 4), + 'baseline_output_bytes': group.baseline_bytes, + 'optimized_output_bytes': group.optimized_bytes, + 'baseline_output_file_count': group.baseline_files, + 'optimized_output_file_count': group.optimized_files, + 'values_equal': group.correctness_ok, + 'max_abs_difference': group.max_abs_difference, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--run-dir', type=Path, action='append', required=True, dest='run_dirs') + parser.add_argument('--output', type=Path, required=True) + args = parser.parse_args() + grouped = load_groups([run_dir.expanduser().resolve() for run_dir in args.run_dirs]) + rows = [_summary_row(group) for groups in grouped.values() for group in groups] + common.write_csv(args.output, rows) + print(f'Wrote {len(rows)} rows to {args.output}') + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/benchmarks/to_modelarray/summarize_scale.py b/benchmarks/to_modelarray/summarize_scale.py new file mode 100644 index 0000000..1e568dc --- /dev/null +++ b/benchmarks/to_modelarray/summarize_scale.py @@ -0,0 +1,139 @@ +"""Pin the production-scale CUBIC results (full-scalar and scalar-sharded runs) to one CSV. + +Reads the step-3 single-scalar comparison and worker sweep, the step-4 scalar-sharded +job-array walls, and (when present) the step-5 production-build walls, and writes one +row per measurement so the numbers are committed and reproducible alongside the +cross-modality summary. Missing inputs are skipped, so the CSV can be regenerated as +the ladder progresses. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import statistics +from pathlib import Path + +import harness_common as common + + +def _rows_compare(run_dir: Path) -> list[dict[str, object]]: + path = run_dir / 'comparison_results.csv' + if not path.exists(): + return [] + with path.open(newline='', encoding='utf-8') as stream: + rows = list(csv.DictReader(stream)) + meta = json.loads((run_dir / 'comparison_meta.json').read_text(encoding='utf-8')) + n_rows = len(common.read_cohort(Path(meta['cohort']))) + out = [] + for backend in sorted({r['backend'] for r in rows}): + walls = { + impl: statistics.median( + float(r['wall_seconds']) + for r in rows + if r['backend'] == backend and r['implementation'] == impl + ) + for impl in ('baseline', 'optimized') + } + out.append( + { + 'step': '3_full_scalar_original_vs_optimized', + 'scalars': ','.join(meta.get('scalar_names') or []), + 'subjects': n_rows, + 'backend': backend, + 'workers': meta.get('workers'), + 'baseline_median_wall_seconds': round(walls['baseline'], 1), + 'optimized_median_wall_seconds': round(walls['optimized'], 1), + 'speedup': round(walls['baseline'] / walls['optimized'], 2), + } + ) + return out + + +def _rows_sweep(run_dir: Path) -> list[dict[str, object]]: + path = run_dir / 'results.csv' + if not path.exists(): + return [] + with path.open(newline='', encoding='utf-8') as stream: + return [ + { + 'step': '3_full_scalar_worker_sweep', + 'scalars': r['scalar_names'], + 'subjects': int(r['input_file_count']), + 'backend': r['backend'], + 'workers': int(r['workers']), + 'baseline_median_wall_seconds': '', + 'optimized_median_wall_seconds': round(float(r['wall_seconds']), 1), + 'speedup': '', + 'input_files_per_minute': round(float(r['input_files_per_minute'])), + 'peak_memory_gb': round(float(r['peak_memory_bytes'] or 0) / 2**30, 1), + 'fragment_count': r.get('fragment_count'), + 'fast_path_rows': r.get('fast_path_rows'), + } + for r in csv.DictReader(stream) + ] + + +def _wall_row( + step: str, scalar: str, backend: str, seconds: str, passed: bool +) -> dict[str, object]: + """One job-array conversion; steps 4 and 5 convert every subject with 8 loader workers.""" + return { + 'step': step, + 'scalars': scalar, + 'subjects': 8693, + 'backend': backend, + 'workers': 8, + 'baseline_median_wall_seconds': '', + 'optimized_median_wall_seconds': round(float(seconds), 1), + 'speedup': '', + 'passed': passed, + } + + +def _rows_shards(shard_root: Path, step: str) -> list[dict[str, object]]: + if not shard_root.exists(): + return [] + out = [] + for wall in sorted(shard_root.glob('*/wall.csv')): + for line in wall.read_text(encoding='utf-8').splitlines(): + _, backend, seconds = line.split(',') + passed = (wall.parent / 'PASSED').exists() + out.append(_wall_row(step, wall.parent.name, backend, seconds, passed)) + return out + + +def _rows_production(prod_root: Path) -> list[dict[str, object]]: + wall = prod_root / 'checks' / 'wall_seconds.csv' + if not wall.exists(): + return [] + out = [] + for line in wall.read_text(encoding='utf-8').splitlines(): + scalar, seconds = line.split(',') + passed = (prod_root / 'checks' / f'{scalar}.PASSED').exists() + out.append(_wall_row('5_production_build', scalar, 'hdf5', seconds, passed)) + return out + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--results-root', type=Path, required=True) + parser.add_argument('--scalar', default='area') + parser.add_argument('--output', type=Path, required=True) + args = parser.parse_args() + root = args.results_root + rows: list[dict[str, object]] = [] + rows += _rows_compare(root / f'cubic-03-{args.scalar}-compare') + rows += _rows_sweep(root / f'cubic-03-{args.scalar}-sweep') + rows += _rows_shards(root / 'cubic-04-shard', '4_scalar_sharded_array') + rows += _rows_production(root / 'production') + if not rows: + raise SystemExit('no production-scale results found') + common.write_csv(args.output, rows) + print(f'Wrote {len(rows)} rows to {args.output}') + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/benchmarks/to_modelarray/tiledb_read_gate.py b/benchmarks/to_modelarray/tiledb_read_gate.py new file mode 100644 index 0000000..a85932d --- /dev/null +++ b/benchmarks/to_modelarray/tiledb_read_gate.py @@ -0,0 +1,90 @@ +"""Python read gate for converter outputs, mirroring ``r_read_gate.R``. + +The R gate needs the R ``tiledb`` package for TileDB outputs, which is not +available in every deployment (for example the only ModelArray container on +CUBIC). This gate performs the same checks from Python: every scalar matrix +opens, its ``column_names`` metadata matches the row count, and a +representative element range reads back. It handles both backends so a +dispatcher can fall back to it for HDF5 too. It prints ``read_seconds=`` like +the R gate and exits non-zero on any failure. + +Usage: ``python tiledb_read_gate.py OUTPUT BACKEND [--n-columns N]`` +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + + +def _fail(message: str) -> int: + print(f'gate failed: {message}', file=sys.stderr) + return 1 + + +def _gate_tiledb(output: Path, n_columns: int) -> int: + import numpy as np + import tiledb + + roots = sorted(p for p in (output / 'scalars').glob('*') if p.is_dir()) + uris = [root / 'values' for root in roots if (root / 'values').exists()] + if not uris: + return _fail(f'no TileDB scalar matrices under {output}') + config = tiledb.Config({'vfs.file.enable_filelocks': 'false'}) + with tiledb.scope_ctx(config): + for uri in uris: + with tiledb.open(str(uri), 'r') as array: + rows, cols = array.shape + names = json.loads(array.meta.get('column_names', '[]')) + if len(names) != rows: + return _fail(f'{uri}: {len(names)} column names for {rows} rows') + first = array[0, 0]['values'] + block = array[0 : min(64, rows), 0 : min(n_columns, cols)]['values'] + if not np.isfinite(np.asarray(first)).size or block.shape[0] == 0: + return _fail(f'{uri}: representative read returned nothing') + return 0 + + +def _gate_hdf5(output: Path, n_columns: int) -> int: + import h5py + + with h5py.File(output, 'r') as handle: + if 'scalars' not in handle: + return _fail(f'no scalars group in {output}') + names = list(handle['scalars']) + if not names: + return _fail(f'no scalar matrices in {output}') + for name in names: + group = handle[f'scalars/{name}'] + if 'values' not in group or 'column_names' not in group: + return _fail(f'scalars/{name}: missing values or column_names') + values = group['values'] + rows, cols = values.shape + if group['column_names'].shape[0] != rows: + return _fail(f'scalars/{name}: column_names length != rows') + values[0, 0] + block = values[0 : min(64, rows), 0 : min(n_columns, cols)] + if block.shape[0] == 0: + return _fail(f'scalars/{name}: representative read returned nothing') + return 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('output', type=Path) + parser.add_argument('backend', choices=['hdf5', 'tiledb']) + parser.add_argument('--n-columns', type=int, default=8) + args = parser.parse_args() + started = time.perf_counter() + gate = _gate_tiledb if args.backend == 'tiledb' else _gate_hdf5 + status = gate(args.output, args.n_columns) + if status == 0: + print(f'read_seconds={time.perf_counter() - started:.6f}') + return status + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/design/TO_MODELARRAY_RUNTIME_OPTIMIZATION_PLAN.md b/design/TO_MODELARRAY_RUNTIME_OPTIMIZATION_PLAN.md new file mode 100644 index 0000000..a243b94 --- /dev/null +++ b/design/TO_MODELARRAY_RUNTIME_OPTIMIZATION_PLAN.md @@ -0,0 +1,851 @@ +# `to-modelarray` runtime optimization plan + +## Goal + +Make `modelarrayio to-modelarray` finish as fast as possible without changing +its output schema or breaking ModelArray reads. Memory use, output size and +downstream read speed are things to measure and check, not things to optimize. + +Where a choice trades speed against precision or output size, the fast option +is the default and the slower, more precise option is an explicit, documented +CLI flag. + +## Terms used in this document + +- **chunk** (HDF5) / **tile** (TileDB) — the fixed-size block a file format + stores and compresses as a unit. Its shape decides how much data gets + compressed at once, and therefore how finely compression can be split across + threads. +- **fragment** (TileDB) — one write commit. Many small commits mean many files + and slower reads. +- **staging matrix** — the in-memory `subjects × elements` array we fill from + the input images and then hand to the backend. +- **ceiling** — the fastest a step could possibly go given the hardware, e.g. + how fast the machine can decode `.nii.gz` files at all. Useful as a stopping + rule: once we are near the ceiling there is nothing left to win. + +## Measured starting point + +**Production baseline (CUBIC, real run).** The original converter was run on the +full RBC freesurferpost cohort — 8,693 subject-sessions × 70 scalars, 608,510 +`fs_LR den-164k` dscalars — with `--backend tiledb --compression zstd +--compression-level 1 --target-chunk-mb 32 --workers 15` +(`modelarray_rbc_tiledb/code/02_create_tiledb.sh`). Its log shows the **first +scalar took 112 h 23 m** (the old loader parsed all 608,510 CIFTI XML headers +up front, GIL-bound, so extra workers made it slower), and the run stalled at +15/70 scalars after ~119 h; 18 scalars were completed. That output is now the +bitwise correctness reference for the optimized converter. + +The scouting numbers below come from a 14-core macOS workstation with a warm +page cache; they rank the work, and the CUBIC ladder in the benchmark plan +replaces them with production measurements. + +```text +HDF5 write, gzip-4 + shuffle + h5py dataset assignment 73 MB/s + parallel direct chunk write, 14 threads 695 MB/s +HDF5 write, lzf + shuffle 336 MB/s +HDF5 write, no compression 1611 MB/s + +TileDB write, zstd + shuffle + current column-stripe path 299 MB/s 92 fragments, 278 files + one full-matrix query 1992 MB/s 1 fragment, 5 files + 100-subject blocked writes 922 MB/s 4 fragments + +nibabel .nii.gz decode, 1 thread 108 MB/s decoded +nibabel .nii.gz decode, 14 threads 1011 MB/s decoded +get_fdata(float32) vs float64 default no measurable difference + +NIfTI row assembly, current flattened_image 52 ms/subject +NIfTI row assembly, cached mask + float32 34 ms/subject +``` + +## What was slow in the original code + +The measurements above explain each bottleneck; the fix for each is in +*Decisions*. + +1. **HDF5 writes were limited by compression, not disk speed.** `h5py` + compresses each chunk one at a time on the single thread that writes the + file, capping output at roughly 70 MB/s no matter how fast the loader feeds + it. Compressing chunks on several threads removes the cap (695 MB/s) and is + the single largest write-side win. +2. **TileDB writes reopened the array once per spatial stripe**, and + `create_empty_scalar_matrix_array` opened it one more time just to write + `column_names`: one scalar became 92 array opens and 92 write commits across + 278 files. One full-matrix query is 6.7 times faster and leaves 1 commit + across 5 files. +3. `load_cohort_cifti`, `load_cohort_voxels` and `load_cohort_mif` each built + the **entire cohort across all scalars** in memory — as dictionaries of + per-subject row lists — before any writing started, and submitted one task + per cohort item when workers > 1. +4. CIFTI-to-TileDB started several scalar jobs at once under its own + `--workers`, so a run could hold several full scalar matrices in memory. +5. `flattened_image` decoded each scalar and its individual mask into + full-volume `float64` temporaries, thresholded the mask at `> 0` into another + float64 array, and mutated the decoded scalar volume in place before + fancy-indexing it — several full-volume temporaries per row at double the + precision the output needs. +6. CIFTI axis validation built and compared large label arrays for every input + file, and S3 loading created a new client for every file. + +**Loading dominates, so overlapping load and write buys little.** With 8 loader +threads the example scalar takes about 1.7 s to load and about 0.10 s to write +to TileDB. Even hiding the write completely behind the load would save about +five percent, and once compression runs on several threads the writer and the +loader compete for the same cores anyway. Overlap is therefore deferred +(implementation step 4); its real value would be hiding S3 network latency. + +Relevant code: `src/modelarrayio/cli/{to_modelarray,cifti_to_h5}.py`, +`src/modelarrayio/utils/{cifti,nifti,mif}.py`, +`src/modelarrayio/storage/{h5_storage,tiledb_storage}.py`. + +## Decisions + +### One public worker option + +*Status: implemented — cli/runtime.py:resolve_loader_workers, cli/to_modelarray.py.* + +`--workers N` means one thing: *the maximum number of input sources being +loaded at the same time*. A source is one scalar image plus any companion input +it needs, such as its NIfTI mask. It applies to local paths and `s3://` URLs, +every modality, and both backends. It does not control HDF5 writers, TileDB +writers, or how many scalars are held in memory. + +This changes the meaning of an existing option. `--workers` is currently +documented as "maximum number of parallel TileDB write workers ... has no +effect when `--backend=hdf5`". Someone passing `--workers 8 --backend tiledb` +today is asking for write concurrency and will now get load concurrency. An +alias cannot cover this because the option name does not change, so: say it in +the release notes and the CLI help, and log the resolved meaning ("loading N +sources concurrently") on the first run with `--workers > 1`, so the change +shows up in existing job logs. + +`--s3-workers` is replaced by `--workers` everywhere. Accept it as a deprecated +alias for one minor release with a warning; give both options `default=None` so +argparse can tell whether the user supplied both; reject conflicting values and +accept matching ones; remove the alias in the next breaking release. + +Do not add separate public options for source, scalar and write workers. +Writer lanes stay an internal runtime policy. + +### Loader concurrency uses threads + +*Status: implemented — utils/loading.py:stage_source_rows (bounded ThreadPoolExecutor).* + +Use a `ThreadPoolExecutor`. This is settled: nibabel `.nii.gz` decoding scaled +from 108 to 1011 MB/s between 1 and 14 threads, because zlib lets other threads +run while it decompresses. Do not build a process pool or a shared-memory ring +buffer. + +### CIFTI loading bypasses the XML parse instead of adding threads + +*Status: implemented — utils/cifti.py:_CiftiRawReadPlan / _read_cifti_row_raw; rows served this way are counted in `fast_path_rows`.* + +The thread scaling above is a NIfTI/gzip result and does **not** transfer to +uncompressed CIFTI. On real `fs_LR den-164k` `.dscalar.nii` files (327,684 +greyordinates, ~3.4 MB, of which ~2 MB is the CIFTI XML header) nibabel spends +about 0.78 s parsing the header and 0.48 s in `get_axis` while reading the data +block takes 0.5 ms. That parse holds the GIL, so it does not thread: measured on +cubic, the loader ran 0.8 files/s at one worker and *0.3 files/s at eight*. A +608,510-file freesurferpost conversion under this path took nearly seven days +with `--workers 15` pinned to roughly one core. + +Because every source in a cohort shares one surface template, its values are a +contiguous floating block after the XML; the XML only needs parsing once. The +loader therefore decodes each scalar's first source with nibabel to fix the +reference row, brain-model axis and layout, then reads later sources straight +from their data block (`np.fromfile` at each file's own `vox_offset`) whenever a +cheap check passes — falling back to a full nibabel decode otherwise. The plan +self-validates against the decoded reference row, so any layout it cannot +reproduce (parcellated-connectivity ordering, on-disk scaling, byte order, +remote or compressed inputs) disables it rather than risking a wrong row. +`MODELARRAYIO_CIFTI_RAW=0` forces the nibabel path. + +**What the cheap check must be (learned on CUBIC).** The first version required +the file's `vox_offset` and size to equal the reference's. On the production RBC +files that rejected 91 of 100 `area` files: the pipeline embeds per-subject +provenance strings (`wb_command` paths with job/subject/session IDs) in the CIFTI +metadata, so the XML length — and `vox_offset` — differs by a few bytes between +subjects while the brain-model `MatrixIndicesMap` is byte-identical +(2,072,200 bytes across every size class). Those rows silently fell back to the +1 s nibabel parse and the CUBIC step-2 CIFTI comparison showed no speedup at +all. The check is now: NIfTI-2 dims, dtype and scaling equal the reference, the +file is large enough to hold the block at its own offset, and a SHA-256 of the +raw XML bytes of the dimension-1 (brain-model) `MatrixIndicesMap` equals the +reference's — a few milliseconds per file, and strictly safer than the size +check, since identical map bytes are what guarantee identical greyordinate +order. `fast_path_rows` in the per-scalar telemetry counts rows served this +way; a value below the row count means silent fallbacks and must be treated as +a performance bug, not noise. + +Measured effect on real data (16 subjects × 3 scalars, warm cache): output +bitwise-identical to nibabel, and 448 s → 6.3 s at `--workers 8` (71×). Because +the raw read releases the GIL, threads help again; the per-scalar reference +decode (~70 nibabel loads for the full cohort) is negligible. This is the +CIFTI-specific answer to "loading dominates" — for uncompressed CIFTI the fix is +to stop parsing the header, not to add loader threads or a process pool. + +### Compress HDF5 chunks on several threads + +*Status: implemented — storage/h5_storage.py:_write_direct_gzip_chunks (pool size `MODELARRAYIO_HDF5_COMPRESSION_WORKERS`, default all CPUs).* + +Instead of letting `h5py` compress chunks on the writer thread, compress them +in a thread pool and hand the finished bytes straight to HDF5: + +```text +chunk-compression pool (N threads) writer thread + byte-shuffle + zlib.compress(chunk) ──→ write_direct_chunk(offset, raw) +``` + +Measured at gzip-4 with shuffle: 70, 139, 269, 507 and 695 MB/s at 1, 2, 4, 8 +and 14 threads. The results read back correctly through ordinary `h5py`. + +Two properties make this safe. The dataset is still created with the standard +`compression='gzip'`, `shuffle=True` settings, and each worker produces bytes +that this standard pipeline can decode — so `rhdf5` and `HDF5Array` read the +file with no plugin and no schema change. And there is still exactly one writer +thread, so a shared output file still has exactly one h5py handle. + +Requirements: + +1. Apply the byte-shuffle step exactly as HDF5's shuffle filter defines it for + the element size in use, covered by a round-trip test against a dataset + written the normal way. +2. Pad the last partial chunk out to the full chunk size before compressing, + because `write_direct_chunk` writes whole chunks only, and check that the + padded region reads back correctly. +3. Skip the direct-chunk path entirely for `--compression none`; uncompressed + writes already run at memory speed. +4. Fall back to ordinary dataset assignment if a worker cannot reproduce the + requested filter settings, and log the fallback instead of silently writing + a slow file. + +### Write each scalar matrix in one submission + +*Status: implemented — storage/tiledb_storage.py (one dense write, fragment_count 1); storage/h5_storage.py:write_full_matrix.* + +Write each staged scalar matrix in a single backend call instead of as a list +of rows. Keep every source-name and CIFTI metadata item exactly where and how +it lives today — in particular, do **not** replace the existing +`scalars//column_names` dataset with an HDF5 attribute, and keep the +TileDB column-name arrays the CIFTI outputs rely on. The current +`create_scalar_matrix_dataset` writes an attribute, so fix or wrap it first. + +For TileDB, use one dense write query per scalar. Keeping the array open does +not merge the stripe queries into one commit, and TileDB-Py 0.36 offers only +`write_direct` and `write_subarray` with no incremental global-order write, so +a single-commit write has no streaming alternative from Python. + +Writing in one call means the whole scalar matrix has to be in memory, which +rules out overlapping loader and writer *within* one scalar. That is +acceptable: one matrix is about 1 GiB even at 500 subjects by 500,000 elements. +Allocate it once per scalar as `np.empty((n_subjects, n_elements), +dtype=storage_dtype, order='C')`, in the layout the backend wants, so neither +side needs a second full-size cast or stack copy. Because of this the plan +needs no peak-memory admission checks, no `--staging-dir`, and no +memory-mapped staging. + +If some future workload cannot hold one matrix, the documented fallback is +writing in subject blocks with the subject tile size shrunk to the block size — +measured at 922 MB/s for 100-subject blocks, three times better than today and +two times worse than one big write. + +Treat the per-subject extent of the chunk or tile as a tunable, since it sets +both the compression unit and how finely compression can be parallelized. Make +whichever compression and layout wins the new default, rather than keeping +gzip-4 because it happened to be the old default. + +### Default to the fast dtype, with precision as an explicit opt-in + +*Status: implemented — cli/to_modelarray.py `--dtype` (float32 default); readers decode directly in the storage dtype.* + +`--dtype float32` stays the default and is the fast path. `--dtype float64` is +the supported opt-in for workloads that need double precision end to end. + +Be precise about where float32 actually wins. *Decoding* in float32 is not +faster: `get_fdata(dtype=np.float32, caching='unchanged')` measured no faster +than the float64 default, because decoding dominates. The win is everything +after decoding — half the bytes to shuffle, half to compress, half to write, +and half the staging memory and cache pressure. Since compression is the HDF5 +bottleneck, halving what it has to compress saves real runtime, not just +memory. + +Implementation: + +- On the default path, read numerical data directly as `np.float32` through the + nibabel API without filling the image cache, and allocate the staging matrix + as C-contiguous `float32`. The CIFTI and NIfTI readers currently call + `get_fdata()` at its float64 default and cast afterwards, which creates a + large avoidable temporary array. +- On `--dtype float64`, load and stage as float64 throughout so precision is + exactly preserved. Do not silently downcast. +- Cast MIF data once, at staging, and only when its requested storage type + differs. +- Benchmark both dtypes at least once on the medium cohort so the CLI help can + state the real cost of `--dtype float64` instead of an estimate. Expect + roughly double the write time and double the output size; confirm it. +- CLI help must say that float32 is the default because it is faster and + smaller, and that float64 preserves full source precision at a measured cost. + +### Cache masks and select with flat indices + +*Status: implemented — utils/nifti.py:_flat_group_indices, ReducedMaskCache, _selected_mask_values.* + +`flattened_image` builds two full-volume temporaries per row: a `float64` mask +thresholded at `> 0`, and a `float64` scalar volume that it modifies in place +and then fancy-indexes. + +The group mask fixes the output columns — the set of voxels every subject's row +is reported at. Each subject also has its **own** individual mask, defined by +that subject's anatomy, which may cover only part of the group mask and differs +from subject to subject. We cannot assume any subject's mask matches the group +mask or another subject's. Voxels that are inside the group mask but outside a +given subject's individual mask are written as `NaN` for that subject; ModelArray +reads those per-subject `NaN`s and handles the differing coverage downstream. So +the per-subject mask is real, per-subject work, not a shared file that can be +loaded once and reused. + +Compute `np.flatnonzero(group_mask)` once; fill the typed staging row by flat +selection at those group indices; reduce each subject's individual mask to the +same group indices; and set the out-of-mask positions of the already-selected +row to `NaN`. The scalar is decoded directly in the storage dtype (`float32` by +default) via `get_fdata(dtype=storage_dtype, caching='unchanged')`, and the mask +is read from its native `dataobj` and thresholded `> 0` — neither is ever +promoted to `float64`. Cache reduced masks by file path in a bounded dictionary +so cohorts that *do* repeat a mask (for example repeated rows, or sessions that +share one mask) avoid re-decoding it, but never skip applying a subject's mask on +the assumption that the group mask covers it. The one safe short-circuit is +when a subject's mask path is literally the group-mask file: every selected voxel +is then in-mask by construction, so applying it is a no-op and is skipped. This +dropped row assembly from 52 to 34 ms per subject with bitwise-identical output, +including `NaN` positions. + +Cache the reduced boolean selection — not the nibabel image and not its +`get_fdata()` array. Each row is a fresh `np.take` result that gets its +out-of-mask positions overwritten with NaN, which is only safe because every +call decodes and selects a fresh copy; a cache that handed back a reused +scalar array would let one subject's NaNs corrupt the next subject's data. The +cached reduced mask is read-only and never mutated. Test that loading the same +path twice through the cache leaves the second row unaffected by the first. + +Do not build a float64 mask just to test `> 0`, do not keep a second full +scalar volume after selection, and stage in `float32` unless the user asks for +`--dtype float64`. + +### Reuse CIFTI setup work and validate cheaply + +*Status: implemented — utils/cifti.py (first row reused; cifti_spatial_signature instead of label compare).* + +Load, decode and validate the first CIFTI source once. Reuse its decoded row as +the first row of the first staging matrix, and reuse its metadata to create the +greyordinate/parcel output metadata without reopening the file. + +Build a short, stable signature of the spatial axis from that reference header +and check later sources against the signature instead of comparing full label +arrays element by element. The signature must cover every spatial property the +current comparison protects, and must leave out scalar labels that legitimately +differ between inputs. Cache signatures that pass, and build full label arrays +only for the reference metadata or for a failure message. + +### Do not throttle TileDB's internal thread pools + +*Status: implemented — concurrency levels are only read into telemetry (cli/runtime.py:_tiledb_runtime_settings), never set.* + +TileDB writes measured 791 MB/s against `h5py`'s 73 at the same codec and +level, and the reason is that TileDB compresses on many threads internally. +Capping `sm.compute_concurrency_level` or `sm.io_concurrency_level` to +`--workers` would give that away. Leave the compute pool at the allocated core +count; if contention shows up in measurements, reduce *loader* threads while a +write is in flight instead. Record TileDB's concurrency settings in every +benchmark run. + +For S3, reuse one configured client whose connection pool can serve `--workers` +loads at once. + +### Per-scalar output is the primary target + +*Status: implemented — cli/to_modelarray.py `--split-files`/`--no-split-files`; per-scalar atomic publish in every CLI.* + +Optimize for **one output file per scalar** — one `.h5` file, or one `.tdb` +root, per scalar. This is what we want in practice: ModelArray reads per-scalar +files directly, and separate files make a cohort easier to share, copy, and +reprocess one scalar at a time. `--split-files` produces per-scalar outputs and +is the default for wide `--scalar-columns` input; `--no-split-files` keeps every +scalar under one combined output root and stays supported for callers that want +a single file. Split is a data-layout choice, not a performance knob — but +because per-scalar files are fully independent, the scalar is also the unit the +converter parallelizes and shards on (next section), so the fast path is built +around it. + +Each scalar is staged lazily and written to its own atomically published output +with no cross-scalar coupling, so one scalar can be built, validated, and +published without touching another. Preserve the greyordinate/parcel and +fixel/voxel metadata each per-scalar file needs. + +### Scalar-sharded conversion for HPC + +*Status: implemented — recipe and job templates in benchmarks/to_modelarray/cubic/ (README, sbatch/04_scalar_array.sbatch); `--only-scalar` deferred — per-scalar cohort CSVs select the scalar.* + +Because each scalar's file is independent, the scalar is the natural unit of +coarse parallelism, and on a cluster the productive pattern is **one job per +scalar** rather than more threads inside one job. This mirrors ModelArray's HPC +sharding (`vignette("element-splitting")` and its `run_chunk.R`), with one +important difference that comes from building the matrix rather than analysing +one. + +ModelArray shards the **element** axis: each array task analyses a disjoint +range of fixels/voxels, writes its own `result_chunk_k.h5`, and a final +single-process step concatenates the chunk files into one output. That works +because it reads element slices lazily from an already-built matrix, and it needs +the combine step only because the shards are disjoint element ranges and HDF5 +forbids concurrent writes. + +Building the matrix is the opposite shape, so the axes map differently: + +- **Element-range sharding does not transfer.** To read elements *j..k* of a + subject you must decode that subject's whole `.nii.gz`/`.mif` — gzip is not + seekable — so *E* element-shards re-decode all *N* subjects (*E×N* decodes). + This is the one piece of ModelArray's approach that does not carry over to the + build. +- **Scalar sharding is the win, with no aggregation.** One array task per scalar, + each decoding its subjects once and writing its own final per-scalar file. The + shards are already the desired output granularity, so there is nothing to + combine and no concurrent-write hazard — "shard and aggregate" collapses to + just "shard." The CLI already supports this (one `--scalar-columns X`, or a + scalar-filtered cohort, per `$SLURM_ARRAY_TASK_ID`); an explicit + `--only-scalar` selector would let a long-format job array map cleanly to the + task index without pre-filtering the CSV. Document this array recipe as the + build-side analogue of `element-splitting`. +- **Subject-block sharding is a niche add-on** for a single scalar too large for + one node. Decode each subject once inside its block, then recombine along the + subject (row) axis. The backend decides the cost: HDF5 inherits ModelArray's + no-concurrent-write rule, so it needs a single-process combine that copies the + row blocks into the final dataset — one extra full pass, worth it only when + decode dominates a plain copy, which it does. TileDB does not inherit it — + subject-block jobs can write disjoint row ranges as concurrent fragments into + one array with no copy, trading the single-fragment write for a bounded number + of fragments (one per block, not per stripe). For the sharded build this is a + real TileDB advantage over HDF5. + +Scalar-sharded jobs and the in-node thread parallelism (`--workers` decode +threads plus the HDF5 compression pool / TileDB internal pool) are orthogonal +and compose: shard scalars across nodes, thread within each job. Neither adds a +public worker knob; `--only-scalar` selects work, it does not size a pool. + +## Implementation plan + +### 1. Measurement before behaviour + +*Status: implemented — cli/runtime.py:TimingRecorder, `--timing-jsonl`.* + +Per-scalar structured timing at INFO level plus optional JSON Lines output for +benchmark automation records: source load time, CIFTI validation time, staging +allocation time, write time with compression separated from submission, +per-scalar and total time, matrix shape and dtype, chunk/tile shape, +compression settings, peak memory, TileDB fragment counts and output bytes, +CPU utilization, and the bounded-loader counters that actually exist — +`scheduler_wait_seconds`, `in_flight_sources_mean/max`, and `fast_path_rows` +(rows served by the CIFTI raw read; a value below the source count means some +rows silently fell back to nibabel). Loader/writer *overlap* counters +(`loader_blocked_seconds`, queue occupancy) are not measured because the +pipeline is sequential per scalar; they are emitted as zeros and must not be +used as evidence. Keep instrumentation cheap: never per-file at production scale. + +### 2. Shared bounded source loader + +Build one modality-neutral bounded-executor helper that: + +- takes the ordered list of sources for the scalar being staged; +- keeps at most `workers` loads submitted or running at any time, rather than + one task per file; +- writes each row straight into its own index in the preallocated staging + matrix, so completion order does not matter and no results dictionary is + needed; +- includes the source path and scalar name in errors; and +- never builds a list of rows or a second full-size `stack`/cast copy. + +Use it for local files and S3 alike; only the scheduling becomes shared, and +each loader still picks its own nibabel/S3/MRtrix reader. Avoid pools inside +pools — CIFTI-to-TileDB must use this loader instead of launching one +long-running outer task per scalar. Loader threads hold no HDF5 handles: for a +combined HDF5 output, every h5py write goes through its single writer lane. For +S3, release the response body and any compressed or decompressed byte buffers +as soon as the typed row is copied into staging. + +### 3. Replace the stripe writers + +Implement the full-matrix writers, the parallel HDF5 chunk-compression path, +and the removal of the redundant array opens. Steps 2 and 3 are the bulk of the +expected win and must be benchmarked before any overlap work begins. + +Use one writer lane for a combined HDF5 file. For TileDB, benchmark one lane +first and add a small number of lanes only for scalar-split outputs that can be +written independently; concurrent writers must never modify shared +output-group metadata unless that is proven safe. Expect the compression pool, +not extra writer lanes, to be the effective lever. + +In-process writer lanes are for building several per-scalar files within one +invocation. On HPC, prefer the coarser and simpler split: one process (one job +array task) per scalar, each writing its own per-scalar file, as in +*Scalar-sharded conversion for HPC*. That needs no shared-handle coordination +and no aggregation, and it uses the scheduler's cores instead of competing with +the in-node loader and compression threads. + +### 4. Deferred: overlap load and write + +Not implemented, and not justified by measurement: on local and GPFS input the +write phase is a small fraction of wall time and the two phases run +sequentially per scalar. Revive it only for S3-latency-bound input, with a +one-item queue (a deeper queue only multiplies resident matrices). Subject-block +sharding is likewise deferred until a single scalar no longer fits one node. + +### 5. Recoverable output completion + +For scalar-split output, write each scalar to a temporary name in the +destination parent and publish it only after its data and metadata validate. +For a combined HDF5 file or TileDB root, write the whole conversion to one +temporary root on the same filesystem and publish after every scalar validates. +Use atomic rename where supported; never turn publication into a hidden copy +whose duration is left out of the benchmark. + +Published scalar outputs let a user target missing work manually, but the +converter must not implicitly skip them on a later invocation. Resume behaviour +stays an explicit opt-in, decided separately. + +## Benchmark plan + +### Benchmark data: real ABIDE images + +Benchmark on real images, not generated arrays. The primary cohort is the open +ABIDE derivatives already used by `test/test_nifti_s3.py`: + +```text +s3://fcp-indi/data/Projects/ABIDE_Initiative/Outputs/cpac/filt_global/ + alff/_alff.nii.gz + func_mask/_func_mask.nii.gz + falff/, reho/, vmhc/, lfcd/, degree_binarize/, degree_weighted/, + eigenvector_binarize/, eigenvector_weighted/ (same 1102 subjects) +``` + +The bucket is public, so `MODELARRAYIO_S3_ANON=1` with unsigned requests is +enough and no credentials appear in any recorded command. The properties that +matter: 1102 subjects per derivative; 61×73×61 volumes with about 65,193 +in-group-mask voxels, so 0.249 MiB per `float32` row and 274 MiB for a full +1102-subject matrix; about 285 KB per gzipped file and 0.29 GiB per derivative; +and the files store **`float64` on disk**, so the float32 change gets tested +against a genuinely double-precision source rather than data that was already +float32. + +Cohorts: + +- **fixture** — the 20 committed `test/data_voxel_toy/FA` volumes plus + `test/data_cifti_toy`, for correctness tests only. +- **small** — 100 ABIDE subjects, 1 scalar, for iteration. +- **medium** — 1102 ABIDE subjects, `alff`. The default comparison cohort. +- **wide** — 1102 ABIDE subjects × 9 derivatives, about 2.4 GiB staged. The + only cohort that exercises the per-scalar loop, `--scalar-columns` split + outputs, and writer lanes. +- **production confirmation** — the real internal workload, run once on the + final two or three candidates. + +To go beyond 1102 input files, repeat real subjects as extra cohort rows rather +than synthesizing values: each repeated row is a genuine open, decode and +mask-apply, so loader throughput, `--workers` scaling and S3 latency stay real. +Record the repeat factor. Two caveats: repeats read warm from the page cache +locally, so a repeated-row local run measures the decode ceiling rather than +the filesystem; and duplicate rows compress unrealistically well, so never draw +output-size or compression-ratio conclusions from them. + +Two limits to state wherever ABIDE numbers are reported. First, 65 k voxels and +285 KB files make per-file overhead matter more than it does in production +fixel and CIFTI workloads — which is exactly what the production confirmation +run exists to catch. Second, ABIDE is NIfTI-only, so the other two modalities +use their own reproducible cohorts, described next. + +### MIF and CIFTI benchmark data + +MIF has a fully public, reproducible cohort: the ModelArray walkthrough FDC +demo — 100 fixel `.mif` subjects (602,229 fixels each) plus `index.mif` and +`directions.mif`, downloaded from OSF (`https://osf.io/tce9d/download`, mirrored +by `benchmarks/to_modelarray/make_mif_cohort.py`). This is real diffusion fixel data, so its +runtime, `--workers` scaling and compression ratios are all quotable, and the +shipped `cohort_FDC_n100.csv` is already long-format and drives the harness +directly. It is the MIF analogue of ABIDE. + +CIFTI has no public `.dscalar.nii` cohort, so it is benchmarked with +`benchmarks/to_modelarray/make_cifti_cohort.py`, which has two reproducible modes. The default +`synthesize` writes seeded `fs_LR den-164k` dscalars with scalar-appropriate +marginal ranges and a medial-wall zero block — reproducible anywhere and valid +for load/runtime timing, but its per-vertex values are independent, so it is not +quoted for compression ratios. Passing `--reference-dir` derives, per scalar, +the shared medial-wall mask and the per-vertex mean and standard deviation from +a small set of real dscalars and generates each subject as +`clip(mean + std·noise, min, max)`; only aggregate statistics are read and fresh +noise is drawn, so no real subject row is reproduced, while the real mean map's +spatial smoothness makes compression realistic (measured 1.38× versus 1.41× for +the real reference at gzip-4, against 1.17× for the independent-noise default). +The reference axes are copied verbatim so the header and `vox_offset` match the +real files exactly, which is what the raw-read fast path depends on. + +Use the same cohort file and an empty output location for every comparison. +Label warm-cache and cold-cache runs separately and never compare across that +line. Repeat important candidates at least three times and compare medians. + +### Benchmarking lives in its own folder + +Everything about benchmarking goes in one new self-contained folder so it can +be added, run and updated without touching the library or the main test suite: + +```text +benchmarks/ + README.md index of suites and their shared conventions + to_modelarray/ this suite; a future export-results suite gets a sibling folder + README.md what each cohort is, how to run, how to read results + make_cohort.py deterministic ABIDE (NIfTI) cohort-CSV generator + make_mif_cohort.py downloads/normalizes the public walkthrough FDC cohort + make_cifti_cohort.py seeded/realistic CIFTI dscalar cohort generator + run_benchmarks.py builds the CLI commands, times them, appends result rows + r_read_gate.R per-candidate R read gate (rhdf5/HDF5Array/tiledb) + plot_results.R figures from the results CSV + compare_revisions.py original (pinned commit) vs optimized, per modality, with a bitwise check + harness_common.py cohort, CSV, output-naming and provenance helpers shared by the scripts + output_compare.py block-wise HDF5/TileDB output comparison (revisions and reference) + comparison_runs.py reads compare_revisions run directories for the report and summary + create_comparison_report.py renders the cross-modality original-vs-optimized findings PDF + fixtures/ committed fixture cohort CSVs + baselines/ + baseline_.csv pinned reference numbers, committed + reports/ rendered findings PDFs, committed + cubic/ HPC validation ladder (SLURM templates, container env) + data/ cohorts/ downloaded/generated inputs, gitignored + results/ gitignored + / + results.csv one row per run + run_meta.json exact commands, environment, versions, cohort hash + plots/ generated figures +``` + +Benchmark output is written under `benchmarks/to_modelarray/results//` and nowhere +else — not into the package, not into the test data folders, not into `docs/`. +An environment variable (`MODELARRAYIO_BENCHMARK_RESULTS_DIR`) overrides the +results root so a cluster job can write to scratch. The harness documentation +lives in `benchmarks/to_modelarray/README.md`; the main docs link to it rather than +duplicating it. + +### Reuse the conventions, not the code, from earlier benchmarking + +An earlier round of HDF5 benchmarking in this project produced result CSVs, +plots and a pytest-based runner. Do not port that code forward: it was written +before `to-modelarray` existed, it benchmarks the writer rather than the CLI, +and adapting it costs more than writing the small harness above. Its +conventions are worth keeping, though, because they worked: + +- Append one row per run to a CSV, so an interrupted sweep keeps the runs it + already finished and later runs accumulate instead of overwriting. +- Keep a `run_meta.json` beside the CSV with the exact commands and environment + for every run, so any row can be reproduced later. +- Select size tiers by flag (quick / medium / full) so the same harness serves + both a 30-second sanity check and an overnight sweep. +- Add columns over time instead of renaming them, and fill missing columns when + loading, so old rows stay plottable by the current plotting script. +- Keep a `run_kind` column that separates end-to-end runs from synthetic layout + sweeps, so the two are never quoted interchangeably. +- When the sweep runs in parallel, write one result file per parallel worker and + merge them at plot time. + +The earlier *numbers* are reusable too, as hypotheses rather than results. Its +chunk-geometry and gzip-level sweeps point the layout and codec sweep below at +the promising region instead of a blind grid. But they came from a writer-only +harness fed synthesized values, so no figure from them counts as a result here; +every number this plan reports comes from a run of the new harness. + +### What the harness must do differently + +1. **Benchmark the CLI end to end, not the writer.** Drive `to-modelarray` + against an ABIDE cohort CSV and record the phase counters from step 1. + Timing a single call to `create_empty_scalar_matrix_dataset` filled with + generated values covers about five percent of what this plan optimizes. +2. **Keep the output.** Do not delete the converted file after measuring its + size; the R read gate and the correctness comparison both need it. +3. **Cover both backends.** Record the backend, TileDB fragment count, output + file count and TileDB concurrency settings. +4. **Record the new knobs**: `--workers`, dtype, compression-pool size, decode + time, writer-idle and loader-blocked time, peak memory and repeat factor. + +Keep a small seeded data generator for the one thing synthetic data is good at: +sweeping chunk geometry and codecs at 10,000 and 40,000 rows, where holding +that many real decoded subjects is not the point. Mark those rows in `run_kind` +and never quote them as end-to-end runtimes. + +### Sequence: the CUBIC validation ladder + +The local original-vs-optimized comparison (2026-09-21, superseded by the CUBIC step-2 results +pinned in `benchmarks/to_modelarray/baselines/cubic_cross_modality_9ab13dd_vs_optimized_20260922.csv`) +fixes the architecture; what remains is validating it where it runs. Each step is +gated on correctness before the next starts; scripts live in `benchmarks/to_modelarray/cubic/`. + +0. **Environment and smoke** — container venv (glibc 2.31) with the exact local + library versions, offline pytest on a compute node, quick-tier fixtures + through both gate branches (R ModelArray for HDF5, Python for TileDB), and a + fixture old-vs-new bitwise comparison. `source_tree_sha256` must equal the + local tree's. +1. **Production-data preflight** — reference-scalar completeness + (`compare_to_reference.py --list`), per-scalar cohort CSVs, header-size + uniformity of one scalar's files (the raw fast path needs it). +2. **Original vs optimized on the public and RBC-subset cohorts** — ABIDE-100, + MIF walkthrough-100, RBC CIFTI 100 × 3; both backends; 3 counterbalanced + repetitions after a warm-up; bitwise equality, both gates; the raw fast path + and the nibabel path must both equal the deployed reference. +3. **One full production scalar** (8,693 subjects) — old vs new on the same + node, a `--workers` sweep, every output bitwise against the deployed + reference (`rbc_freesurferpost.tdb`), `fragment_count == 1`, no leftover + temporaries, an R read of a representative element range. +4. **Scalar-sharded job array** — one task per complete reference scalar, its + own per-scalar file, no aggregation; every task bitwise against its + reference scalar. Aggregate wall versus the serial sum is the HPC result. +5. **Full production build** — every scalar as a job array, outputs kept, the 18 + cross-checkable scalars bitwise, every file through the HDF5 gate. Submitted + only when every step-4 shard passed. + +Every run records: ModelArrayIO commit *and* `source_tree_sha256`, Python, +nibabel, h5py and TileDB versions, allocated CPUs and RAM, TileDB concurrency +settings, source and output filesystem types, cohort hash, scalar names, input +byte total, and the exact CLI command; plus wall time, CPU utilization, peak +memory, output bytes, phase times, `fast_path_rows`, TileDB fragment and file +counts, and peak open file descriptors. + +### Every HDF5 codec candidate must pass an R read gate + +ModelArray reads HDF5 through `rhdf5` and `HDF5Array`. LZF ships with `h5py`; +zstd and blosc need a filter plugin registered against whatever libhdf5 those +Bioconductor packages were built with. A codec that Python writes quickly but R +cannot open is not a candidate at all, however good its number looks: LZF +measured 336 MB/s against gzip-4's 73, which is exactly the kind of result that +gets adopted by accident. + +Read every candidate output back through ModelArray in R, on the deployment +environment, as a pass/fail gate before its timing counts. This is why the +harness must keep its output. The parallel direct-chunk path makes the question +largely moot for HDF5, since it reaches 695 MB/s while writing a standard +gzip+shuffle file that needs no plugin. + +Record R-side read time for a representative element range alongside the gate. +A candidate that halves conversion time by turning off compression and then +doubles every downstream read is a bad trade that a write-only benchmark cannot +see. + +## CUBIC results (ladder completed 2026-09-23) + +Environment: CUBIC compute nodes (Xeon Gold 6130), everything inside the ModelArray +container with a uv-managed Python 3.12 venv (numpy 2.5.3, h5py 3.16.0, tiledb +0.36.1, nibabel 5.4.2); HDF5 gated by the current ModelArray R package installed +into a user library, TileDB gated in Python; `source_tree_sha256` recorded in every +run. Warm GPFS cache; medians of three counterbalanced repetitions after one +warm-up; one loader worker; identical inputs for both converters. + +| Step | Result | +|---|---| +| 0 Env + smoke | 415 tests pass on a compute node; quick-tier NIfTI and CIFTI fixtures pass both gate branches; fixture original-vs-optimized bitwise on both backends. | +| 2 Original (`9ab13dd`) vs optimized | NIfTI ABIDE-100: **1.40× HDF5, 1.46× TileDB**. MIF walkthrough-100: **3.70× HDF5, 4.11× TileDB**. RBC CIFTI 100 × 3 (area, areaPial, curv): **31.8× HDF5 (385.8 → 12.1 s), 31.6× TileDB (396.5 → 12.6 s)**. All 36 runs passed their gate; every optimized output bitwise-equal to the original (max abs diff 0.0). | +| 2b Raw fast path on production headers | After the digest fix, `fast_path_rows` = 99/100 per scalar (was 8/100); raw path 27.2 s vs nibabel path 659.2 s for 100 × 3; **both** paths bitwise-equal to the deployed old-code TileDB reference (rows aligned by source name, NaN maps equal). | +| 3 One full scalar (8,693 subjects), old vs new | `area`, 8,693 × 327,684, `--target-chunk-mb 32`, one loader worker, 1 rep each: **HDF5 11,854.8 s (3 h 17 m) → 429.9 s (7.2 min), 27.6×; TileDB 11,501.6 s (3 h 12 m) → 418.1 s (7.0 min), 27.5×** (output files 1,031 → 14). Both gates passed; optimized output bitwise-equal to the original on the full matrix (max abs diff 0.0). Optimized worker sweep on the same scalar (`fast_path_rows` = 8,692, i.e. every non-reference row): HDF5 403.9 / 120.9 / **90.6 s** at 1 / 4 / 8 workers (5,756 files/min), peak RSS ≈ 12 GB; TileDB 186.3 / 83.2 / **68.6 s** (7,603 files/min), `fragment_count` 1, peak RSS ≈ 20 GB — about 130× (HDF5) and 170× (TileDB) the original at 8 workers. All eight optimized outputs bitwise-equal to the deployed old-code reference for all 8,693 rows. | +| 4 Scalar-sharded job array (15 reference scalars) | `--array=1-15%6`, 8 CPUs / 48 GB per task, one scalar per task, HDF5 then TileDB, no aggregation: **15/15 shards passed**; per-scalar wall 112–141 s (HDF5) and 70–118 s (TileDB) at 8 workers; serial-equivalent sum **1,947 s (32 min) HDF5 / 1,466 s (24 min) TileDB** for 15 scalars that took the original run ~119 h; **array wall-clock 21.0 min** (first task start to last task end, rerun included). All 30 per-scalar outputs bitwise-equal to the deployed old-code reference. One task was cancelled by the node (signal, 3 min in, on the node step 3 had just vacated) and passed on rerun. | +| 5 Full production build (70 scalars) | `--array=1-70%10`, 8 CPUs / 48 GB per task, per-scalar HDF5 gzip-4, `--target-chunk-mb 32`, 8 workers: **70/70 scalars passed** (converter validation + current-ModelArray R gate on every file; the 15 cross-checkable scalars bitwise-equal to the deployed reference). Per-scalar wall 107–192 s, median 138 s; serial-equivalent 9,800 s (2.7 h); **array wall-clock 22.2 min** for the entire 608,510-file dataset (484 GB of per-scalar files). The original run needed 112 h for its first scalar and never finished. | + +Three harness/cluster defects surfaced and were fixed by the ladder before any +production number was trusted: the container's ModelArray 0.1.5 cannot read the +current HDF5 format (attribute- vs dataset-stored `column_names`); the cohort +writers emitted CRLF line endings that broke shell consumers; and the login node's +2 h CPU-time `ulimit` propagates into SLURM jobs, killing the original converter +with SIGXCPU two-thirds of the way through a full scalar (jobs now submit with +`--propagate=NONE`). All are documented in `benchmarks/to_modelarray/cubic/README.md`. + +## Acceptance criteria + +Compare every optimized output against the current converter on the same +inputs: + +- scalar names, source ordering, shapes, dtypes, and the candidate's declared + chunk/tile metadata — not equality with the baseline layout, when the sweep + picked a different compatible one; +- CIFTI greyordinate and parcel metadata; +- NaN positions, and values bitwise equal where the source transform permits; + otherwise an explicitly documented float32 rounding tolerance; +- successful ModelArray construction and representative reads in R for the + selected codec and layout; and +- no unexpected TileDB fragment growth and no leftover temporary output. + +The CUBIC ladder above is the acceptance path: each step's gate must pass on +production data before the next runs. Ship the fastest candidate that produces +correct, compatible output, clears the read gates, and fits the allocated +resources. Report output size and +downstream read observations, but do not use them to reject a faster correct +conversion unless the R-side read regression is large enough to be recorded as +a deliberate trade. + +## Deliverables + +Code and tests: + +- Unit tests for bounded scheduling, ordered rows, worker validation, and + float32/float64 reader behaviour. +- A round-trip test proving the parallel direct-chunk path produces a dataset + whose declared filter settings and read-back values match one written the + normal way, including the padded final partial chunk and the + `--compression none` bypass. +- A test that the direct-chunk fallback triggers and logs when a worker cannot + reproduce the filter settings. +- Mask-caching and flat-index tests covering the group-mask case, the + differing-mask case, and the cached-row mutation hazard. +- Integration tests for CIFTI, NIfTI and MIF at `--workers 1` and `>1`, on both + backends, and at both dtypes. +- Tests that `--s3-workers` warns and aliases `--workers`, including the + conflicting-value case. +- Regression tests that scalar-split versus long-cohort behaviour is unchanged, + that full-matrix writers preserve every existing HDF5 and TileDB metadata + location including CIFTI column-name arrays, and that one scalar commits as + one TileDB fragment. +- CIFTI tests that reuse the first decoded row, accept an equivalent later + header without building label arrays, and still reject every spatial-axis + mismatch the current validation catches. + +Documentation and benchmarking: + +- CLI help, README and usage docs covering the single worker meaning, the + redefinition of `--workers`, the `--s3-workers` removal schedule, and the + measured cost of `--dtype float64`. +- The `benchmarks/to_modelarray/` folder described above, containing the harness, the + plotting script, `benchmarks/to_modelarray/README.md`, and a deterministic ABIDE cohort-CSV + generator that takes a subject count, derivative list and repeat factor, so + anyone can reproduce the small, medium and wide cohorts from the public + bucket with no credentials. +- A pinned baseline results CSV committed under `benchmarks/to_modelarray/baselines/`, so a + future regression is detectable by rerunning one command. +- An HPC scalar-sharding recipe — the build-side analogue of ModelArray's + `element-splitting` — showing a SLURM job array that converts one scalar per + task into its own per-scalar file with no aggregation step, plus a note on the + TileDB concurrent-fragment option for a single scalar too large for one node. + +Optional, decided separately: + +- An `--only-scalar ` selector so a long-format cohort maps one scalar to + one job-array task without pre-filtering the CSV, with a regression test that + it produces exactly the same per-scalar output as converting that scalar + alone. + +## Non-goals + +- No independent public knobs for source, scalar and write workers. (A + `--only-scalar` selector for scalar-sharded HPC job arrays selects work, not + worker counts, and is a separate, optional addition.) +- No output schema change; chunk/tile layout is tunable only while the output + stays compatible and readable from R. +- No process-based or shared-memory loader *inside one conversion*. A + scalar-sharded HPC job array — one ordinary threaded process per scalar, + writing its own file — is a supported deployment pattern, not a process-pool + loader, and does not violate this. +- No peak-memory admission checks, memory-mapped staging, or `--staging-dir`. +- No mandatory node-local scratch. +- No exhaustive sweep of every combination of formats, backends and settings + before the target architecture is fixed. diff --git a/docs/conf.py b/docs/conf.py index c7894f4..b373b7c 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -40,7 +40,13 @@ source_suffix = {'.rst': 'restructuredtext'} master_doc = 'index' -exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] +exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store', 'examples/README.rst'] + +# Gallery examples intentionally reuse headings such as "Prepare data". Prefix +# their automatic labels by document and avoid labels for deep argparse-generated +# subsections, which can repeat within a single usage page. +autosectionlabel_prefix_document = True +autosectionlabel_maxdepth = 2 html_theme = 'sphinx_rtd_theme' html_static_path = ['_static'] @@ -81,7 +87,12 @@ def _sync_overview_figure(app) -> None: def _patch_sphinxarg_parallel_merge() -> None: """Teach sphinx-argparse's domain how to merge parallel-read data.""" - from sphinxarg.ext import ArgParseDomain + try: + from sphinxarg.ext import ArgParseDomain + except ImportError: + # sphinx-argparse versions without a domain do not need this legacy + # parallel-data merge patch. + return def merge_domaindata(self, docnames, otherdata): self.data['commands'].extend(otherdata.get('commands', [])) diff --git a/docs/examples/plot_cifti_workflow.py b/docs/examples/plot_cifti_workflow.py index 397aa4a..5f6c7b6 100644 --- a/docs/examples/plot_cifti_workflow.py +++ b/docs/examples/plot_cifti_workflow.py @@ -5,8 +5,8 @@ For imaging data in CIFTI format, use the ``modelarrayio to-modelarray`` command to convert the CIFTI files to the HDF5 format (``.h5``) used by **ModelArray**, and ``modelarrayio export-results`` to export results back to CIFTI. -The CIFTI workflow is very similar to the MIF workflow -(:ref:`sphx_glr_auto_examples_plot_mif_workflow.py`). +The CIFTI workflow is very similar to the fixel workflow +(:ref:`sphx_glr_auto_examples_plot_fixel_workflow.py`). """ # %% diff --git a/docs/examples/plot_voxel_workflow.py b/docs/examples/plot_voxel_workflow.py index 290ddd0..8671f7f 100644 --- a/docs/examples/plot_voxel_workflow.py +++ b/docs/examples/plot_voxel_workflow.py @@ -6,7 +6,7 @@ the NIfTI files to the HDF5 format (``.h5``) used by **ModelArray**, and ``modelarrayio export-results`` to export results back to NIfTI. The voxel workflow is very similar to the fixel workflow -(:ref:`sphx_glr_auto_examples_plot_mif_workflow.py`). +(:ref:`sphx_glr_auto_examples_plot_fixel_workflow.py`). """ # %% diff --git a/pyproject.toml b/pyproject.toml index 6ccb90c..2ce3d78 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,6 @@ dependencies = [ "nibabel>=5.0", "numpy>=1.26", "pandas>=2.0", - "tqdm>=4.65", "h5py>=3.10", "tiledb>=0.23", ] @@ -99,7 +98,8 @@ exclude = ".*" [tool.ruff] line-length = 99 -src = ["src"] +# The benchmark scripts import each other as top-level modules from their own directory. +src = ["src", "benchmarks/to_modelarray"] extend-exclude = ["src/modelarrayio/__about__.py"] [tool.ruff.lint] @@ -180,6 +180,8 @@ exclude_lines = [ ] [tool.pytest.ini_options] +# Benchmark scripts and their tests import the shared benchmark modules by name. +pythonpath = ["benchmarks/to_modelarray"] markers = [ "s3: tests that require network access to public S3 (deselect with '-m not s3')", "downloaded_data: tests that download external fixture data before running", diff --git a/src/modelarrayio/cli/cifti_to_h5.py b/src/modelarrayio/cli/cifti_to_h5.py index d091a96..77f28c3 100644 --- a/src/modelarrayio/cli/cifti_to_h5.py +++ b/src/modelarrayio/cli/cifti_to_h5.py @@ -3,20 +3,24 @@ from __future__ import annotations import logging -from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path -import h5py -from tqdm import tqdm - from modelarrayio.cli import utils as cli_utils +from modelarrayio.cli.runtime import ( + TimingRecorder, + publish_scalar_outputs, + resolve_loader_workers, + validate_timing_jsonl_path, +) +from modelarrayio.storage import utils as storage_utils from modelarrayio.utils.cifti import ( _get_cifti_parcel_info, brain_names_to_dataframe, - extract_cifti_scalar_data, load_cohort_cifti, ) +from modelarrayio.utils.loading import LazyScalarMatrices from modelarrayio.utils.misc import build_scalar_sources +from modelarrayio.utils.s3_utils import configure_s3_client_pool logger = logging.getLogger(__name__) @@ -31,9 +35,10 @@ def cifti_to_h5( shuffle=True, chunk_voxels=0, target_chunk_mb=2.0, - workers=1, - s3_workers=1, + workers=None, + s3_workers=None, split_outputs=False, + timing_jsonl=None, ): """Load all CIFTI data and write to an HDF5 or TileDB file. @@ -61,12 +66,13 @@ def cifti_to_h5( target_chunk_mb : :obj:`float` Target chunk/tile size in MiB when auto-computing the spatial axis length workers : :obj:`int` - Maximum number of parallel TileDB write workers. Default 1. - Has no effect when ``backend='hdf5'``. + Maximum number of input sources loaded concurrently, for every backend. s3_workers : :obj:`int` - Number of workers for parallel S3 downloads + Deprecated alias for ``workers``. split_outputs : :obj:`bool` If True, write one output file per scalar. Default False. + timing_jsonl : path-like, optional + Append structured per-scalar and total runtime records to this file. Returns ------- @@ -76,127 +82,87 @@ def cifti_to_h5( scalar_sources = build_scalar_sources(cohort_long) if not scalar_sources: raise ValueError('Unable to derive scalar sources from cohort file.') - scalar_names = list(scalar_sources.keys()) - split_scalar_outputs = split_outputs - - _first_scalar, first_sources = next(iter(scalar_sources.items())) - first_path = first_sources[0] - cifti_type, parcel_arrays = _get_cifti_parcel_info(first_path) - - if backend == 'hdf5': - if split_scalar_outputs: - scalars, last_brain_names = load_cohort_cifti(cohort_long, s3_workers) - if cifti_type == 'dscalar': - greyordinate_table, structure_names = brain_names_to_dataframe(last_brain_names) - outputs: list[Path] = [] - for scalar_name in scalar_names: - scalar_output = cli_utils.prepare_output_parent( - cli_utils.prefixed_output_path(output, scalar_name) - ) - with h5py.File(scalar_output, 'w') as h5_file: - if cifti_type == 'dscalar': - cli_utils.write_table_dataset( - h5_file, - 'greyordinates', - greyordinate_table, - extra_attrs={'structure_names': structure_names}, - ) - else: - cli_utils.write_hdf5_parcel_arrays(h5_file, parcel_arrays) - cli_utils.write_hdf5_scalar_matrices( - h5_file, - {scalar_name: scalars[scalar_name]}, - {scalar_name: scalar_sources[scalar_name]}, - storage_dtype=storage_dtype, - compression=compression, - compression_level=compression_level, - shuffle=shuffle, - chunk_voxels=chunk_voxels, - target_chunk_mb=target_chunk_mb, - ) - outputs.append(scalar_output) - return int(not all(path.exists() for path in outputs)) - - scalars, last_brain_names = load_cohort_cifti(cohort_long, s3_workers) - output = cli_utils.prepare_output_parent(output) - with h5py.File(output, 'w') as h5_file: - if cifti_type == 'dscalar': - greyordinate_table, structure_names = brain_names_to_dataframe(last_brain_names) - cli_utils.write_table_dataset( - h5_file, - 'greyordinates', - greyordinate_table, - extra_attrs={'structure_names': structure_names}, - ) - else: - cli_utils.write_hdf5_parcel_arrays(h5_file, parcel_arrays) - cli_utils.write_hdf5_scalar_matrices( + scalar_names = list(scalar_sources) + worker_count = resolve_loader_workers(workers, s3_workers) + configure_s3_client_pool(worker_count) + compression, compression_level = storage_utils.resolve_backend_compression( + backend, compression, compression_level + ) + planned_outputs = cli_utils.conversion_output_paths( + output, + scalar_names, + split_outputs=split_outputs, + ) + validate_timing_jsonl_path( + timing_jsonl, + planned_outputs, + directory_outputs=backend == 'tiledb', + ) + recorder = TimingRecorder( + timing_jsonl, + modality='cifti', + backend=backend, + workers=worker_count, + compression=compression, + compression_level=compression_level, + shuffle=shuffle, + ) + + scalars, reference_brain_names = load_cohort_cifti( + cohort_long, + workers=worker_count, + storage_dtype=storage_dtype, + ) + # The loader hands back the decoded first image so the metadata can be + # derived without reopening the file; plain mappings (tests) reopen it. + reference_image = scalars.reference_image if isinstance(scalars, LazyScalarMatrices) else None + first_path = scalar_sources[scalar_names[0]][0] + cifti_type, parcel_arrays = _get_cifti_parcel_info( + reference_image if reference_image is not None else first_path + ) + + if cifti_type == 'dscalar': + greyordinate_table, structure_names = brain_names_to_dataframe(reference_brain_names) + + def _hdf5_metadata(h5_file): + cli_utils.write_table_dataset( h5_file, - scalars, - scalar_sources, - storage_dtype=storage_dtype, - compression=compression, - compression_level=compression_level, - shuffle=shuffle, - chunk_voxels=chunk_voxels, - target_chunk_mb=target_chunk_mb, - ) - return int(not output.exists()) - - output = Path(output) - output.mkdir(parents=True, exist_ok=True) - - _, reference_brain_names = extract_cifti_scalar_data(first_path) - - def _process_scalar_job(scalar_name, source_files): - rows = [] - for source_file in source_files: - cifti_data, _ = extract_cifti_scalar_data( - source_file, reference_brain_names=reference_brain_names + 'greyordinates', + greyordinate_table, + extra_attrs={'structure_names': structure_names}, ) - rows.append(cifti_data) - if rows: - scalar_output = ( - cli_utils.prefixed_output_path(output, scalar_name) - if split_scalar_outputs - else output - ) - cli_utils.write_tiledb_scalar_matrices( - scalar_output, - {scalar_name: rows}, - {scalar_name: source_files}, - storage_dtype=storage_dtype, - compression=compression, - compression_level=compression_level, - shuffle=shuffle, - chunk_voxels=chunk_voxels, - target_chunk_mb=target_chunk_mb, - write_column_name_arrays=True, - ) - return scalar_name - - worker_count = min(len(scalar_names), workers) - - if worker_count <= 1: - for scalar_name in scalar_names: - _process_scalar_job(scalar_name, scalar_sources[scalar_name]) + hdf5_required_paths = ('greyordinates',) else: - desc = 'TileDB scalars' - with ThreadPoolExecutor(max_workers=worker_count) as executor: - futures = { - executor.submit(_process_scalar_job, scalar_name, scalar_sources[scalar_name]): ( - scalar_name - ) - for scalar_name in scalar_names - } - for future in tqdm(as_completed(futures), total=len(futures), desc=desc): - future.result() - if parcel_arrays: - if split_scalar_outputs: - for scalar_name in scalar_names: - scalar_output = cli_utils.prefixed_output_path(output, scalar_name) - cli_utils.write_tiledb_parcel_arrays(scalar_output, parcel_arrays) - else: - cli_utils.write_tiledb_parcel_arrays(output, parcel_arrays) + + def _hdf5_metadata(h5_file): + cli_utils.write_hdf5_parcel_arrays(h5_file, parcel_arrays) + + hdf5_required_paths = tuple(f'parcels/{name}' for name in parcel_arrays) + + def _tiledb_metadata(root): + if parcel_arrays: + cli_utils.write_tiledb_parcel_arrays(root, parcel_arrays) + + publish_scalar_outputs( + recorder, + backend=backend, + output=output, + split_outputs=split_outputs, + scalar_names=scalar_names, + scalars=scalars, + sources_by_scalar=scalar_sources, + storage_dtype=storage_dtype, + compression=compression, + compression_level=compression_level, + shuffle=shuffle, + chunk_voxels=chunk_voxels, + target_chunk_mb=target_chunk_mb, + expected_elements=len(reference_brain_names), + hdf5_metadata=_hdf5_metadata, + hdf5_required_paths=hdf5_required_paths, + tiledb_metadata=_tiledb_metadata, + tiledb_required_paths=tuple(f'parcels/{name}' for name in parcel_arrays), + column_name_arrays=True, + ) return 0 diff --git a/src/modelarrayio/cli/mif_to_h5.py b/src/modelarrayio/cli/mif_to_h5.py index bcd41da..850c411 100644 --- a/src/modelarrayio/cli/mif_to_h5.py +++ b/src/modelarrayio/cli/mif_to_h5.py @@ -3,14 +3,18 @@ from __future__ import annotations import logging -from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path -import h5py -from tqdm import tqdm - from modelarrayio.cli import utils as cli_utils +from modelarrayio.cli.runtime import ( + TimingRecorder, + publish_scalar_outputs, + resolve_loader_workers, + validate_timing_jsonl_path, +) +from modelarrayio.storage import utils as storage_utils from modelarrayio.utils.mif import gather_fixels, load_cohort_mif +from modelarrayio.utils.s3_utils import configure_s3_client_pool logger = logging.getLogger(__name__) @@ -27,9 +31,10 @@ def mif_to_h5( shuffle=True, chunk_voxels=0, target_chunk_mb=2.0, - workers=1, - s3_workers=1, + workers=None, + s3_workers=None, split_outputs=False, + timing_jsonl=None, ): """Load all fixeldb data and write to an HDF5 or TileDB file. @@ -61,97 +66,76 @@ def mif_to_h5( target_chunk_mb : :obj:`float` Target chunk/tile size in MiB when auto-computing the spatial axis length workers : :obj:`int` - Maximum number of parallel TileDB write workers. Default 1. - Has no effect when ``backend='hdf5'``. + Maximum number of input sources loaded concurrently, for every backend. s3_workers : :obj:`int` - Number of parallel workers for S3 downloads. Default 1. + Deprecated alias for ``workers``. split_outputs : :obj:`bool` If True, write one output file per scalar. Default False. + timing_jsonl : path-like, optional + Append structured per-scalar and total runtime records to this file. Returns ------- status : :obj:`int` 0 if successful, 1 if failed. """ + worker_count = resolve_loader_workers(workers, s3_workers) + configure_s3_client_pool(worker_count) + compression, compression_level = storage_utils.resolve_backend_compression( + backend, compression, compression_level + ) + planned_scalar_names = list(dict.fromkeys(str(value) for value in cohort_long['scalar_name'])) + planned_outputs = cli_utils.conversion_output_paths( + output, + planned_scalar_names, + split_outputs=split_outputs, + ) + validate_timing_jsonl_path( + timing_jsonl, + planned_outputs, + directory_outputs=backend == 'tiledb', + ) + recorder = TimingRecorder( + timing_jsonl, + modality='mif', + backend=backend, + workers=worker_count, + compression=compression, + compression_level=compression_level, + shuffle=shuffle, + ) # gather fixel data fixel_table, voxel_table = gather_fixels(index_file, directions_file) logger.info('Extracting .mif data...') - scalars, sources_lists = load_cohort_mif(cohort_long, s3_workers) + scalars, sources_lists = load_cohort_mif( + cohort_long, + workers=worker_count, + storage_dtype=storage_dtype, + ) if not sources_lists: raise ValueError('Unable to derive scalar sources from cohort file.') - scalar_names = list(sources_lists.keys()) - split_scalar_outputs = split_outputs - - if backend == 'hdf5': - if split_scalar_outputs: - outputs: list[Path] = [] - for scalar_name in scalar_names: - scalar_output = cli_utils.prepare_output_parent( - cli_utils.prefixed_output_path(output, scalar_name) - ) - with h5py.File(scalar_output, 'w') as h5_file: - cli_utils.write_table_dataset(h5_file, 'fixels', fixel_table) - cli_utils.write_table_dataset(h5_file, 'voxels', voxel_table) - cli_utils.write_hdf5_scalar_matrices( - h5_file, - {scalar_name: scalars[scalar_name]}, - {scalar_name: sources_lists[scalar_name]}, - storage_dtype=storage_dtype, - compression=compression, - compression_level=compression_level, - shuffle=shuffle, - chunk_voxels=chunk_voxels, - target_chunk_mb=target_chunk_mb, - ) - outputs.append(scalar_output) - return int(not all(path.exists() for path in outputs)) - - output = cli_utils.prepare_output_parent(output) - with h5py.File(output, 'w') as h5_file: - cli_utils.write_table_dataset(h5_file, 'fixels', fixel_table) - cli_utils.write_table_dataset(h5_file, 'voxels', voxel_table) - cli_utils.write_hdf5_scalar_matrices( - h5_file, - scalars, - sources_lists, - storage_dtype=storage_dtype, - compression=compression, - compression_level=compression_level, - shuffle=shuffle, - chunk_voxels=chunk_voxels, - target_chunk_mb=target_chunk_mb, - ) - return int(not output.exists()) - - worker_count = min(len(scalar_names), workers) - - def _write_scalar_job(scalar_name): - scalar_output = ( - cli_utils.prefixed_output_path(output, scalar_name) if split_scalar_outputs else output - ) - cli_utils.write_tiledb_scalar_matrices( - scalar_output, - {scalar_name: scalars[scalar_name]}, - {scalar_name: sources_lists[scalar_name]}, - storage_dtype=storage_dtype, - compression=compression, - compression_level=compression_level, - shuffle=shuffle, - chunk_voxels=chunk_voxels, - target_chunk_mb=target_chunk_mb, - ) - - if worker_count <= 1: - for scalar_name in scalar_names: - _write_scalar_job(scalar_name) - else: - with ThreadPoolExecutor(max_workers=worker_count) as executor: - futures = { - executor.submit(_write_scalar_job, scalar_name): scalar_name - for scalar_name in scalar_names - } - for future in tqdm(as_completed(futures), total=len(futures), desc='TileDB scalars'): - future.result() + def _hdf5_metadata(h5_file): + cli_utils.write_table_dataset(h5_file, 'fixels', fixel_table) + cli_utils.write_table_dataset(h5_file, 'voxels', voxel_table) + + publish_scalar_outputs( + recorder, + backend=backend, + output=output, + split_outputs=split_outputs, + scalar_names=list(sources_lists), + scalars=scalars, + sources_by_scalar=sources_lists, + storage_dtype=storage_dtype, + compression=compression, + compression_level=compression_level, + shuffle=shuffle, + chunk_voxels=chunk_voxels, + target_chunk_mb=target_chunk_mb, + expected_elements=len(fixel_table), + hdf5_metadata=_hdf5_metadata, + hdf5_required_paths=('fixels', 'voxels'), + ) return 0 diff --git a/src/modelarrayio/cli/nifti_to_h5.py b/src/modelarrayio/cli/nifti_to_h5.py index d113676..1cfc699 100644 --- a/src/modelarrayio/cli/nifti_to_h5.py +++ b/src/modelarrayio/cli/nifti_to_h5.py @@ -3,17 +3,25 @@ from __future__ import annotations import logging -from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path -import h5py -import nibabel as nb import numpy as np import pandas as pd -from tqdm import tqdm from modelarrayio.cli import utils as cli_utils +from modelarrayio.cli.runtime import ( + TimingRecorder, + publish_scalar_outputs, + resolve_loader_workers, + validate_timing_jsonl_path, +) +from modelarrayio.storage import utils as storage_utils from modelarrayio.utils.nifti import load_cohort_voxels +from modelarrayio.utils.s3_utils import ( + configure_s3_client_pool, + load_nibabel, + release_image_buffers, +) logger = logging.getLogger(__name__) @@ -29,9 +37,10 @@ def nifti_to_h5( shuffle=True, chunk_voxels=0, target_chunk_mb=2.0, - workers=1, - s3_workers=1, + workers=None, + s3_workers=None, split_outputs=False, + timing_jsonl=None, ): """Load all volume data and write to an HDF5 or TileDB file. @@ -61,15 +70,44 @@ def nifti_to_h5( target_chunk_mb : :obj:`float` Target chunk/tile size in MiB when auto-computing. Default 2.0. workers : :obj:`int` - Maximum number of parallel TileDB write workers. Default 1. - Has no effect when ``backend='hdf5'``. + Maximum number of input sources loaded concurrently, for every backend. s3_workers : :obj:`int` - Number of parallel workers for S3 downloads. Default 1. + Deprecated alias for ``workers``. split_outputs : :obj:`bool` If True, write one output file per scalar. Default False. + timing_jsonl : path-like, optional + Append structured per-scalar and total runtime records to this file. """ - group_mask_img = nb.load(group_mask_file) - group_mask_matrix = group_mask_img.get_fdata() > 0 + worker_count = resolve_loader_workers(workers, s3_workers) + configure_s3_client_pool(worker_count) + compression, compression_level = storage_utils.resolve_backend_compression( + backend, compression, compression_level + ) + planned_scalar_names = list(dict.fromkeys(str(value) for value in cohort_long['scalar_name'])) + planned_outputs = cli_utils.conversion_output_paths( + output, + planned_scalar_names, + split_outputs=split_outputs, + ) + validate_timing_jsonl_path( + timing_jsonl, + planned_outputs, + directory_outputs=backend == 'tiledb', + ) + recorder = TimingRecorder( + timing_jsonl, + modality='nifti', + backend=backend, + workers=worker_count, + compression=compression, + compression_level=compression_level, + shuffle=shuffle, + ) + group_mask_img = load_nibabel(group_mask_file) + try: + group_mask_matrix = np.asanyarray(group_mask_img.dataobj) > 0 + finally: + release_image_buffers(group_mask_img) voxel_coords = np.column_stack(np.nonzero(group_mask_matrix)) voxel_table = pd.DataFrame( @@ -82,79 +120,35 @@ def nifti_to_h5( ) logger.info('Extracting NIfTI data...') - scalars, sources_lists = load_cohort_voxels(cohort_long, group_mask_matrix, s3_workers) + scalars, sources_lists = load_cohort_voxels( + cohort_long, + group_mask_matrix, + workers=worker_count, + storage_dtype=storage_dtype, + group_mask_path=group_mask_file, + ) if not sources_lists: raise ValueError('Unable to derive scalar sources from cohort file.') - scalar_names = list(sources_lists.keys()) - split_scalar_outputs = split_outputs - - if backend == 'hdf5': - if split_scalar_outputs: - outputs: list[Path] = [] - for scalar_name in scalar_names: - scalar_output = cli_utils.prepare_output_parent( - cli_utils.prefixed_output_path(output, scalar_name) - ) - with h5py.File(scalar_output, 'w') as h5_file: - cli_utils.write_table_dataset(h5_file, 'voxels', voxel_table) - cli_utils.write_hdf5_scalar_matrices( - h5_file, - {scalar_name: scalars[scalar_name]}, - {scalar_name: sources_lists[scalar_name]}, - storage_dtype=storage_dtype, - compression=compression, - compression_level=compression_level, - shuffle=shuffle, - chunk_voxels=chunk_voxels, - target_chunk_mb=target_chunk_mb, - ) - outputs.append(scalar_output) - return int(not all(path.exists() for path in outputs)) - - output = cli_utils.prepare_output_parent(output) - with h5py.File(output, 'w') as h5_file: - cli_utils.write_table_dataset(h5_file, 'voxels', voxel_table) - cli_utils.write_hdf5_scalar_matrices( - h5_file, - scalars, - sources_lists, - storage_dtype=storage_dtype, - compression=compression, - compression_level=compression_level, - shuffle=shuffle, - chunk_voxels=chunk_voxels, - target_chunk_mb=target_chunk_mb, - ) - return int(not output.exists()) - - worker_count = min(len(scalar_names), workers) - - def _write_scalar_job(scalar_name): - scalar_output = ( - cli_utils.prefixed_output_path(output, scalar_name) if split_scalar_outputs else output - ) - cli_utils.write_tiledb_scalar_matrices( - scalar_output, - {scalar_name: scalars[scalar_name]}, - {scalar_name: sources_lists[scalar_name]}, - storage_dtype=storage_dtype, - compression=compression, - compression_level=compression_level, - shuffle=shuffle, - chunk_voxels=chunk_voxels, - target_chunk_mb=target_chunk_mb, - ) - - if worker_count <= 1: - for scalar_name in scalar_names: - _write_scalar_job(scalar_name) - else: - with ThreadPoolExecutor(max_workers=worker_count) as executor: - futures = { - executor.submit(_write_scalar_job, scalar_name): scalar_name - for scalar_name in scalar_names - } - for future in tqdm(as_completed(futures), total=len(futures), desc='TileDB scalars'): - future.result() + def _hdf5_metadata(h5_file): + cli_utils.write_table_dataset(h5_file, 'voxels', voxel_table) + + publish_scalar_outputs( + recorder, + backend=backend, + output=output, + split_outputs=split_outputs, + scalar_names=list(sources_lists), + scalars=scalars, + sources_by_scalar=sources_lists, + storage_dtype=storage_dtype, + compression=compression, + compression_level=compression_level, + shuffle=shuffle, + chunk_voxels=chunk_voxels, + target_chunk_mb=target_chunk_mb, + expected_elements=len(voxel_table), + hdf5_metadata=_hdf5_metadata, + hdf5_required_paths=('voxels',), + ) return 0 diff --git a/src/modelarrayio/cli/parser_utils.py b/src/modelarrayio/cli/parser_utils.py index 21cfeb9..b72076a 100644 --- a/src/modelarrayio/cli/parser_utils.py +++ b/src/modelarrayio/cli/parser_utils.py @@ -1,6 +1,9 @@ from __future__ import annotations from pathlib import Path +from urllib.parse import urlparse + +from modelarrayio.utils.s3_utils import is_s3_path def add_log_level_arg(parser): @@ -28,3 +31,13 @@ def _is_file(path: str | Path | None, parser) -> Path: if not path.is_file(): raise parser.error(f'Path should point to a file (or symlink of file): <{path}>.') return path + + +def _is_file_or_s3(path: str | Path | None, parser) -> Path | str: + """Accept an existing local file or a syntactically valid S3 object URI.""" + if path is not None and is_s3_path(path): + parsed = urlparse(str(path)) + if parsed.netloc and parsed.path.lstrip('/'): + return str(path) + raise parser.error(f'Invalid S3 URI: <{path}>.') + return _is_file(path, parser) diff --git a/src/modelarrayio/cli/runtime.py b/src/modelarrayio/cli/runtime.py new file mode 100644 index 0000000..00412bc --- /dev/null +++ b/src/modelarrayio/cli/runtime.py @@ -0,0 +1,629 @@ +"""Runtime policy, transactional publication, and cheap conversion metrics.""" + +from __future__ import annotations + +import json +import logging +import os +import shutil +import sys +import time +import uuid +from collections.abc import Callable, Mapping, Sequence +from contextlib import contextmanager +from pathlib import Path + +import h5py +import numpy as np +import tiledb + +try: + import resource +except ImportError: # pragma: no cover - resource is available on supported benchmark hosts + resource = None + +from modelarrayio.cli import utils as cli_utils +from modelarrayio.utils.loading import resolve_loader_workers, validate_workers + +logger = logging.getLogger(__name__) + +__all__ = [ + 'TimingRecorder', + 'atomic_output', + 'load_timing_metrics', + 'path_stats', + 'peak_rss_bytes', + 'publish_scalar_outputs', + 'resolve_loader_workers', + 'scalar_output_stats', + 'validate_hdf5_output', + 'validate_tiledb_output', + 'validate_timing_jsonl_path', + 'validate_workers', +] + + +def load_timing_metrics(load_stats) -> dict[str, float | int]: + """Translate staging counters into consistently named runtime metrics.""" + values = load_stats.to_dict() if load_stats is not None else {} + source_wall = float(values.get('source_wall_seconds', 0.0)) + staging_wall = float(values.get('wall_seconds', 0.0)) + staging_cpu = float(values.get('cpu_seconds', 0.0)) + return { + 'source_load_seconds': source_wall, + # Retain the original keys for benchmark-file compatibility while also + # exposing their precise staging semantics. + 'source_cpu_seconds': staging_cpu, + 'source_cpu_utilization_percent': ( + 100.0 * staging_cpu / staging_wall if staging_wall else 0.0 + ), + 'source_worker_seconds': float(values.get('source_seconds', 0.0)), + 'validation_seconds': float(values.get('validation_seconds', 0.0)), + 'staging_seconds': staging_wall, + 'staging_cpu_seconds': staging_cpu, + 'staging_cpu_utilization_percent': ( + 100.0 * staging_cpu / staging_wall if staging_wall else 0.0 + ), + 'staging_allocation_seconds': float(values.get('allocation_seconds', 0.0)), + 'writer_idle_seconds': staging_wall, + 'loader_blocked_seconds': 0.0, + 'scheduler_wait_seconds': float(values.get('scheduler_wait_seconds', 0.0)), + 'queue_occupancy_mean': 0.0, + 'queue_occupancy_max': 0, + 'in_flight_sources_mean': float(values.get('in_flight_sources_mean', 0.0)), + 'in_flight_sources_max': int(values.get('in_flight_sources_max', 0)), + # Rows served by a reader fast path (CIFTI raw greyordinate read); a + # value below source_count means some rows silently fell back. + 'fast_path_rows': int(values.get('fast_path_rows', 0)), + } + + +def _tiledb_runtime_settings() -> dict[str, int | str]: + """Return the TileDB thread settings without modifying its global pools.""" + config = tiledb.Config() + settings: dict[str, int | str] = {} + for metric_name, config_name in ( + ('tiledb_compute_concurrency_level', 'sm.compute_concurrency_level'), + ('tiledb_io_concurrency_level', 'sm.io_concurrency_level'), + ): + value = config[config_name] + try: + settings[metric_name] = int(value) + except (TypeError, ValueError): + settings[metric_name] = str(value) + return settings + + +def _open_file_descriptor_count() -> int | None: + """Return the current descriptor count on Linux/macOS when available.""" + for descriptor_dir in (Path('/proc/self/fd'), Path('/dev/fd')): + try: + return sum(1 for _ in descriptor_dir.iterdir()) + except OSError: + continue + return None + + +def _temporary_sibling(output: Path) -> Path: + token = uuid.uuid4().hex + return output.with_name(f'.{output.name}.modelarrayio-tmp-{token}') + + +def validate_timing_jsonl_path( + jsonl_path: str | Path | None, + output_paths: Sequence[str | Path], + *, + directory_outputs: bool, +) -> None: + """Reject telemetry paths that would be replaced by or corrupt an output.""" + if jsonl_path is None: + return + timing_path = Path(jsonl_path).resolve(strict=False) + for output_path in output_paths: + resolved_output = Path(output_path).resolve(strict=False) + conflicts = timing_path == resolved_output + if directory_outputs: + conflicts = conflicts or resolved_output in timing_path.parents + if conflicts: + raise ValueError( + f'Timing JSONL path {str(jsonl_path)!r} conflicts with conversion output ' + f'{str(output_path)!r}; choose a path outside the output.' + ) + + +def _remove_path(path: Path) -> None: + if path.is_dir() and not path.is_symlink(): + shutil.rmtree(path) + elif path.exists() or path.is_symlink(): + path.unlink() + + +def _publish_directory(temporary: Path, output: Path) -> None: + """Publish a directory with rollback when replacing a non-empty root.""" + if not output.exists(): + temporary.replace(output) + return + + backup = output.with_name(f'.{output.name}.modelarrayio-backup-{uuid.uuid4().hex}') + storage_logger = logging.getLogger('modelarrayio.storage.tiledb_storage') + storage_logger.warning('Removing existing array tree %s during atomic publication', output) + output.replace(backup) + try: + temporary.replace(output) + except BaseException: + if not output.exists() and backup.exists(): + backup.replace(output) + raise + else: + _remove_path(backup) + + +@contextmanager +def atomic_output(output: str | Path, *, directory: bool): + """Yield a same-parent temporary path and publish it only on success.""" + final_path = Path(output) + final_path.parent.mkdir(parents=True, exist_ok=True) + temporary = _temporary_sibling(final_path) + _remove_path(temporary) + try: + yield temporary + if not temporary.exists(): + raise RuntimeError(f'Conversion did not create temporary output {temporary}.') + if directory: + _publish_directory(temporary, final_path) + else: + os.replace(temporary, final_path) + except BaseException: + _remove_path(temporary) + raise + + +def _normalized_names(values) -> list[str]: + """Return storage strings as ordinary Python text.""" + array = np.asarray(values).reshape(-1) + return [ + value.decode('utf-8') if isinstance(value, (bytes, bytearray, np.bytes_)) else str(value) + for value in array + ] + + +def _expected_element_count(expected_elements, scalar_name: str) -> int | None: + if expected_elements is None: + return None + if isinstance(expected_elements, Mapping): + value = expected_elements.get(scalar_name) + return int(value) if value is not None else None + return int(expected_elements) + + +def validate_hdf5_output( + path: str | Path, + scalar_names: Sequence[str], + *, + expected_sources: Mapping[str, Sequence[str]] | None = None, + expected_dtype: str | np.dtype | None = None, + expected_elements: int | Mapping[str, int] | None = None, + required_paths: Sequence[str] = (), +) -> None: + """Validate scalar shape, dtype, source order, and required metadata.""" + dtype = np.dtype(expected_dtype) if expected_dtype is not None else None + with h5py.File(path, 'r') as h5_file: + for required_path in required_paths: + if required_path not in h5_file: + raise RuntimeError( + f'Output validation failed: missing required HDF5 path {required_path!r}.' + ) + for scalar_name in scalar_names: + values_path = f'scalars/{scalar_name}/values' + columns_path = f'scalars/{scalar_name}/column_names' + if values_path not in h5_file or columns_path not in h5_file: + raise RuntimeError( + f'Output validation failed: missing {values_path!r} or {columns_path!r}.' + ) + values = h5_file[values_path] + columns = h5_file[columns_path] + if not isinstance(values, h5py.Dataset) or values.ndim != 2: + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + f'{values_path!r} must be a two-dimensional dataset.' + ) + if not isinstance(columns, h5py.Dataset) or columns.ndim != 1: + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + f'{columns_path!r} must be a one-dimensional dataset.' + ) + if values.shape[0] != columns.shape[0]: + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + f'{values.shape[0]} rows but {columns.shape[0]} source names.' + ) + if dtype is not None and values.dtype != dtype: + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + f'dtype {values.dtype} != expected {dtype}.' + ) + element_count = _expected_element_count(expected_elements, scalar_name) + if element_count is not None and values.shape[1] != element_count: + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + f'{values.shape[1]} elements != expected {element_count}.' + ) + stored_sources = _normalized_names(columns[...]) + if expected_sources is not None: + expected = [str(value) for value in expected_sources[scalar_name]] + if stored_sources != expected: + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + 'source names or order do not match the cohort.' + ) + if values.shape[0] and values.shape[1]: + # Force representative first/last chunks through the normal + # decoder before publication without reading the whole matrix. + values[0, 0] + values[-1, -1] + + +def validate_tiledb_output( + path: str | Path, + scalar_names: Sequence[str], + *, + require_column_name_arrays: bool = False, + expected_sources: Mapping[str, Sequence[str]] | None = None, + expected_dtype: str | np.dtype | None = None, + expected_elements: int | Mapping[str, int] | None = None, + required_paths: Sequence[str] = (), +) -> None: + """Validate TileDB scalar shape, dtype, source order, and metadata.""" + root = Path(path) + dtype = np.dtype(expected_dtype) if expected_dtype is not None else None + for required_path in required_paths: + required_uri = str(root / required_path) + if tiledb.object_type(required_uri) is None: + raise RuntimeError( + f'Output validation failed: missing required TileDB path {required_uri!r}.' + ) + for scalar_name in scalar_names: + uri = str(root / 'scalars' / scalar_name / 'values') + if tiledb.object_type(uri) != 'array': + raise RuntimeError(f'Output validation failed: missing TileDB array {uri!r}.') + with tiledb.open(uri, 'r') as array: + schema = array.schema + if schema.sparse or schema.domain.ndim != 2: + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + 'values must be a two-dimensional dense TileDB array.' + ) + if schema.nattr != 1 or not schema.has_attr('values'): + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + "array must contain exactly one 'values' attribute." + ) + if dtype is not None and schema.attr('values').dtype != dtype: + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + f'dtype {schema.attr("values").dtype} != expected {dtype}.' + ) + element_count = _expected_element_count(expected_elements, scalar_name) + if element_count is not None and array.shape[1] != element_count: + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + f'{array.shape[1]} elements != expected {element_count}.' + ) + encoded_names = array.meta.get('column_names') + try: + column_names = json.loads(encoded_names) if encoded_names is not None else None + except (TypeError, ValueError) as exc: + raise RuntimeError( + f'Output validation failed: invalid column_names metadata in {uri!r}.' + ) from exc + if not isinstance(column_names, list) or len(column_names) != array.shape[0]: + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + f'{array.shape[0]} rows but ' + f'{len(column_names) if isinstance(column_names, list) else 0} source names.' + ) + stored_sources = [str(value) for value in column_names] + if expected_sources is not None: + expected = [str(value) for value in expected_sources[scalar_name]] + if stored_sources != expected: + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + 'source names or order do not match the cohort.' + ) + if array.shape[0] and array.shape[1]: + domain = schema.domain + first = tuple(int(domain.dim(index).domain[0]) for index in range(2)) + last = tuple(int(domain.dim(index).domain[1]) for index in range(2)) + array[first[0], first[1]]['values'] + array[last[0], last[1]]['values'] + if require_column_name_arrays: + column_uri = str(root / 'scalars' / scalar_name / 'column_names') + if tiledb.object_type(column_uri) != 'array': + raise RuntimeError( + f'Output validation failed: missing TileDB array {column_uri!r}.' + ) + with tiledb.open(column_uri, 'r') as column_array: + array_sources = _normalized_names(column_array[:]['values']) + if array_sources != stored_sources: + raise RuntimeError( + f'Output validation failed for scalar {scalar_name!r}: ' + 'column-name array does not match values metadata.' + ) + + +def path_stats(path: str | Path) -> tuple[int, int]: + """Return total bytes and regular-file count below a path.""" + output = Path(path) + if output.is_file(): + return output.stat().st_size, 1 + total_bytes = 0 + file_count = 0 + if output.exists(): + for child in output.rglob('*'): + if child.is_file(): + file_count += 1 + total_bytes += child.stat().st_size + return total_bytes, file_count + + +def scalar_output_stats( + writer_result, + scalar_name: str, + output_path: str | Path, + *, + split_output: bool, +) -> tuple[int | None, int | None]: + """Return non-cumulative bytes/files for one scalar write when measurable.""" + raw = writer_result.get(scalar_name, {}) if isinstance(writer_result, dict) else {} + if isinstance(raw, dict) and isinstance(raw.get('output_bytes'), int): + file_count = raw.get('output_file_count') + return raw['output_bytes'], file_count if isinstance(file_count, int) else None + if split_output: + return path_stats(output_path) + # A combined HDF5 file is still open when per-scalar telemetry is emitted. + # Its current file size is buffered and cumulative, so report only the final + # conversion-level size rather than a misleading scalar value. + return None, None + + +def peak_rss_bytes() -> int: + """Return process peak RSS normalized across macOS and Linux.""" + if resource is None: + return 0 + rss = int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) + return rss if sys.platform == 'darwin' else rss * 1024 + + +class TimingRecorder: + """Emit structured INFO timing records and optional JSON Lines output.""" + + def __init__(self, jsonl_path: str | Path | None, **context) -> None: + from modelarrayio.utils.s3_utils import s3_client_pool_connections + + self.jsonl_path = Path(jsonl_path) if jsonl_path is not None else None + context.setdefault('s3_pool_connections', s3_client_pool_connections()) + if context.get('backend') == 'tiledb': + for key, value in _tiledb_runtime_settings().items(): + context.setdefault(key, value) + self.context = context + self.started_wall = time.perf_counter() + self.started_cpu = time.process_time() + self.peak_open_file_descriptors = _open_file_descriptor_count() + if self.jsonl_path is not None: + self.jsonl_path.parent.mkdir(parents=True, exist_ok=True) + + def emit(self, event: str, **values) -> dict: + open_descriptors = _open_file_descriptor_count() + if open_descriptors is not None: + if self.peak_open_file_descriptors is None: + self.peak_open_file_descriptors = open_descriptors + else: + self.peak_open_file_descriptors = max( + self.peak_open_file_descriptors, open_descriptors + ) + record = { + 'event': event, + 'timestamp_unix': time.time(), + **self.context, + **values, + } + record.setdefault('open_file_descriptors', open_descriptors) + record.setdefault('peak_open_file_descriptors_sampled', self.peak_open_file_descriptors) + logger.info('runtime %s', json.dumps(record, sort_keys=True, default=str)) + if self.jsonl_path is not None: + try: + with self.jsonl_path.open('a', encoding='utf-8') as stream: + stream.write(json.dumps(record, sort_keys=True, default=str) + '\n') + except OSError: + logger.warning( + 'Could not append runtime metrics to %s; continuing without JSONL output.', + self.jsonl_path, + exc_info=True, + ) + self.jsonl_path = None + return record + + def finish(self, **values) -> dict: + wall_seconds = time.perf_counter() - self.started_wall + cpu_seconds = time.process_time() - self.started_cpu + return self.emit( + 'conversion_complete', + total_seconds=wall_seconds, + cpu_seconds=cpu_seconds, + cpu_utilization_percent=(100.0 * cpu_seconds / wall_seconds if wall_seconds else 0.0), + peak_memory_bytes=peak_rss_bytes(), + **values, + ) + + +def _writer_metrics(result, scalar_name: str) -> dict[str, object]: + """Return a writer's per-scalar metrics minus keys the scalar record sets itself.""" + if not isinstance(result, dict): + return {} + value = result.get(scalar_name) + metrics = value if isinstance(value, dict) else result + reserved = {'scalar_name', 'matrix_shape', 'dtype', 'output_bytes', 'output_file_count'} + return {key: item for key, item in metrics.items() if key not in reserved} + + +def _matrix_description(matrix, result, scalar_name: str, storage_dtype) -> tuple[list, str]: + raw = result.get(scalar_name, {}) if isinstance(result, dict) else {} + if isinstance(raw, dict) and raw.get('matrix_shape') is not None: + return list(raw['matrix_shape']), str(raw.get('dtype', storage_dtype)) + array = matrix if hasattr(matrix, 'shape') else np.asarray(matrix) + return list(array.shape), str(array.dtype) + + +def publish_scalar_outputs( + recorder: TimingRecorder, + *, + backend: str, + output: str | Path, + split_outputs: bool, + scalar_names: Sequence[str], + scalars: Mapping[str, object], + sources_by_scalar: Mapping[str, Sequence[str]], + storage_dtype: str, + compression: str, + compression_level: int | None, + shuffle: bool, + chunk_voxels: int, + target_chunk_mb: float, + expected_elements: int, + hdf5_metadata: Callable[[h5py.File], None] | None = None, + hdf5_required_paths: Sequence[str] = (), + tiledb_metadata: Callable[[Path], None] | None = None, + tiledb_required_paths: Sequence[str] = (), + column_name_arrays: bool = False, +) -> list[Path]: + """Stage, write, validate and atomically publish every scalar; emit telemetry. + + This is the one per-scalar loop shared by the CIFTI, NIfTI and MIF + converters. ``scalars`` is indexed once per scalar (normally a + :class:`~modelarrayio.utils.loading.LazyScalarMatrices`, so each access + stages that scalar) and the matrix is dropped right after its write. + ``hdf5_metadata`` runs once per HDF5 file right after it is created; + ``tiledb_metadata`` runs once per TileDB root after its scalar arrays are + written. Validation happens on the temporary path before publication. + """ + directory = backend == 'tiledb' + # Stats are filled in as scalars are accessed, so keep the live mapping + # itself rather than a snapshot of it (an empty dict is falsy). + load_stats = getattr(scalars, 'stats', None) + if load_stats is None: + load_stats = {} + writer_kwargs = { + 'storage_dtype': storage_dtype, + 'compression': compression, + 'compression_level': compression_level, + 'shuffle': shuffle, + 'chunk_voxels': chunk_voxels, + 'target_chunk_mb': target_chunk_mb, + } + + def _validate(path, names): + if directory: + validate_tiledb_output( + path, + names, + require_column_name_arrays=column_name_arrays, + expected_sources=sources_by_scalar, + expected_dtype=storage_dtype, + expected_elements=expected_elements, + required_paths=tiledb_required_paths, + ) + else: + validate_hdf5_output( + path, + names, + expected_sources=sources_by_scalar, + expected_dtype=storage_dtype, + expected_elements=expected_elements, + required_paths=hdf5_required_paths, + ) + + def _write(target, scalar_name): + started = time.perf_counter() + matrix = scalars[scalar_name] + one_scalar = {scalar_name: matrix} + one_source = {scalar_name: sources_by_scalar[scalar_name]} + if directory: + result = cli_utils.write_tiledb_scalar_matrices( + target, + one_scalar, + one_source, + write_column_name_arrays=column_name_arrays, + **writer_kwargs, + ) + else: + result = cli_utils.write_hdf5_scalar_matrices( + target, one_scalar, one_source, **writer_kwargs + ) + return matrix, result, started + + def _emit(scalar_name, matrix, result, started, output_path): + output_bytes, output_files = scalar_output_stats( + result, scalar_name, output_path, split_output=split_outputs + ) + matrix_shape, matrix_dtype = _matrix_description( + matrix, result, scalar_name, storage_dtype + ) + recorder.emit( + 'scalar_complete', + scalar_name=scalar_name, + scalar_seconds=time.perf_counter() - started, + matrix_shape=matrix_shape, + dtype=matrix_dtype, + **load_timing_metrics(load_stats.get(scalar_name)), + peak_memory_bytes=peak_rss_bytes(), + output_bytes=output_bytes, + output_file_count=output_files, + **_writer_metrics(result, scalar_name), + ) + + published: list[Path] = [] + if split_outputs: + for scalar_name in scalar_names: + scalar_output = Path(cli_utils.prefixed_output_path(output, scalar_name)) + with atomic_output(scalar_output, directory=directory) as temporary: + if directory: + matrix, result, started = _write(temporary, scalar_name) + if tiledb_metadata is not None: + tiledb_metadata(temporary) + else: + with h5py.File(temporary, 'w') as h5_file: + if hdf5_metadata is not None: + hdf5_metadata(h5_file) + matrix, result, started = _write(h5_file, scalar_name) + _validate(temporary, [scalar_name]) + _emit(scalar_name, matrix, result, started, temporary) + del matrix + published.append(scalar_output) + else: + final_output = Path(output) + with atomic_output(final_output, directory=directory) as temporary: + if directory: + for scalar_name in scalar_names: + matrix, result, started = _write(temporary, scalar_name) + _emit(scalar_name, matrix, result, started, temporary) + del matrix + if tiledb_metadata is not None: + tiledb_metadata(temporary) + else: + with h5py.File(temporary, 'w') as h5_file: + if hdf5_metadata is not None: + hdf5_metadata(h5_file) + for scalar_name in scalar_names: + matrix, result, started = _write(h5_file, scalar_name) + _emit(scalar_name, matrix, result, started, temporary) + del matrix + _validate(temporary, list(scalar_names)) + published.append(final_output) + + output_stats = [path_stats(path) for path in published] + recorder.finish( + scalar_count=len(scalar_names), + output_bytes=sum(total_bytes for total_bytes, _ in output_stats), + output_file_count=sum(file_count for _, file_count in output_stats), + ) + return published diff --git a/src/modelarrayio/cli/to_modelarray.py b/src/modelarrayio/cli/to_modelarray.py index 6b39e83..8590b08 100644 --- a/src/modelarrayio/cli/to_modelarray.py +++ b/src/modelarrayio/cli/to_modelarray.py @@ -11,7 +11,9 @@ from modelarrayio.cli.cifti_to_h5 import cifti_to_h5 from modelarrayio.cli.mif_to_h5 import mif_to_h5 from modelarrayio.cli.nifti_to_h5 import nifti_to_h5 -from modelarrayio.cli.parser_utils import _is_file, add_log_level_arg +from modelarrayio.cli.parser_utils import _is_file, _is_file_or_s3, add_log_level_arg +from modelarrayio.cli.runtime import resolve_loader_workers +from modelarrayio.storage import utils as storage_utils from modelarrayio.utils.misc import load_and_normalize_cohort logger = logging.getLogger(__name__) @@ -22,18 +24,19 @@ def to_modelarray( backend='hdf5', output=Path('modelarray.h5'), storage_dtype='float32', - compression='gzip', - compression_level=4, + compression=None, + compression_level=None, shuffle=True, chunk_voxels=0, target_chunk_mb=2.0, - workers=1, - s3_workers=1, + workers=None, + s3_workers=None, scalar_columns=None, split_outputs=None, group_mask_file=None, index_file=None, directions_file=None, + timing_jsonl=None, ): """Load neuroimaging data and write to an HDF5 or TileDB modelarray file. @@ -57,6 +60,15 @@ def to_modelarray( cohort_long, modality = load_and_normalize_cohort(cohort_file, scalar_columns=scalar_columns) logger.info('Detected modality: %s', modality) + workers = resolve_loader_workers(workers, s3_workers) + if workers > 1: + logger.info('Loading %d sources concurrently.', workers) + compression, compression_level = storage_utils.resolve_backend_compression( + backend, + compression, + compression_level, + ) + if split_outputs is None: split_outputs = bool(scalar_columns) @@ -71,8 +83,9 @@ def to_modelarray( 'chunk_voxels': chunk_voxels, 'target_chunk_mb': target_chunk_mb, 'workers': workers, - 's3_workers': s3_workers, + 's3_workers': None, 'split_outputs': split_outputs, + 'timing_jsonl': timing_jsonl, } if modality == 'nifti': @@ -111,7 +124,7 @@ def _parse_to_modelarray(): ), formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) - IsFile = partial(_is_file, parser=parser) + IsFileOrS3 = partial(_is_file_or_s3, parser=parser) parser.add_argument( '--cohort-file', @@ -162,7 +175,13 @@ def _parse_to_modelarray(): ) parser.add_argument( '--dtype', - help='Floating dtype for storing values: float32 (default) or float64', + help=( + 'Floating dtype for loading and storing values. float32 is the faster, smaller ' + 'default; float64 preserves full source precision and generally doubles staging ' + 'memory and bytes processed. On the bundled 20-subject fixture, float64 made the ' + 'HDF5 write phase 1.54x slower and the compressed output 1.45x larger; real cohorts ' + 'vary.' + ), choices=['float32', 'float64'], default='float32', dest='storage_dtype', @@ -170,20 +189,20 @@ def _parse_to_modelarray(): parser.add_argument( '--compression', help=( - 'Compression filter (default gzip). ' + 'Compression filter. The default is gzip for HDF5 and zstd for TileDB. ' 'gzip works for both backends; ' 'lzf is HDF5-only; ' 'zstd is TileDB-only.' ), choices=['gzip', 'zstd', 'lzf', 'none'], - default='gzip', + default=None, ) parser.add_argument( '--compression-level', '--compression_level', type=int, - help='Compression level (codec-dependent). Default 4.', - default=4, + help='Compression level. Effective defaults: gzip 4; TileDB zstd 5.', + default=None, ) parser.add_argument( '--no-shuffle', @@ -212,16 +231,22 @@ def _parse_to_modelarray(): default=2.0, ) - tiledb_group = parser.add_argument_group('TileDB arguments') - tiledb_group.add_argument( + runtime_group = parser.add_argument_group('Runtime arguments') + runtime_group.add_argument( '--workers', type=int, help=( - 'Maximum number of parallel TileDB write workers. ' - 'Default 1. ' - 'Has no effect when --backend=hdf5.' + 'Maximum number of input sources loaded concurrently, for local or S3 input, ' + 'every modality, and either backend. Effective default: 1.' ), - default=1, + default=None, + ) + runtime_group.add_argument( + '--timing-jsonl', + '--timing_jsonl', + type=Path, + default=None, + help='Append structured per-scalar and total runtime metrics to this JSON Lines file.', ) s3_group = parser.add_argument_group('S3 arguments') @@ -229,11 +254,10 @@ def _parse_to_modelarray(): '--s3-workers', '--s3_workers', type=int, - default=1, + default=None, help=( - 'Number of parallel worker processes for loading image files. ' - 'Set > 1 to enable parallel downloads when cohort paths begin with s3://. ' - 'Default 1 (serial).' + 'Deprecated alias for --workers. Matching values are accepted for one minor ' + 'release; conflicting values are rejected.' ), ) @@ -241,7 +265,7 @@ def _parse_to_modelarray(): nifti_group.add_argument( '--mask', help='Path to a NIfTI binary group mask file.', - type=IsFile, + type=IsFileOrS3, default=None, dest='group_mask_file', ) @@ -251,14 +275,14 @@ def _parse_to_modelarray(): '--index-file', '--index_file', help='Nifti2 index file.', - type=IsFile, + type=IsFileOrS3, default=None, ) mif_group.add_argument( '--directions-file', '--directions_file', help='Nifti2 directions file.', - type=IsFile, + type=IsFileOrS3, default=None, ) diff --git a/src/modelarrayio/cli/utils.py b/src/modelarrayio/cli/utils.py index 69b5da0..61f62ef 100644 --- a/src/modelarrayio/cli/utils.py +++ b/src/modelarrayio/cli/utils.py @@ -12,6 +12,7 @@ import pandas as pd from modelarrayio.storage import h5_storage, tiledb_storage +from modelarrayio.storage import utils as storage_utils def configure_logging(log_level: str) -> None: @@ -45,6 +46,33 @@ def prefixed_output_path(output_path: str | Path, prefix: str) -> Path: return path.with_name(f'{safe_prefix}_{path.name}') +def conversion_output_paths( + output_path: str | Path, + scalar_names: Sequence[str], + *, + split_outputs: bool, +) -> list[Path]: + """Return final output paths and reject colliding split filenames.""" + if not split_outputs: + for scalar_name in scalar_names: + storage_utils.validate_scalar_name(scalar_name) + return [Path(output_path)] + + validated_names = [storage_utils.validate_scalar_name(name) for name in scalar_names] + paths = [prefixed_output_path(output_path, scalar_name) for scalar_name in validated_names] + collision_keys: dict[str, tuple[str, Path]] = {} + for scalar_name, path in zip(validated_names, paths, strict=True): + key = str(path.resolve(strict=False)).casefold() + if key in collision_keys: + previous_name, previous_path = collision_keys[key] + raise ValueError( + f'Scalar names {previous_name!r} and {str(scalar_name)!r} map to the same ' + f'split output {str(previous_path)!r}.' + ) + collision_keys[key] = (str(scalar_name), path) + return paths + + def write_table_dataset( h5_file: h5py.File, dataset_name: str, @@ -62,7 +90,7 @@ def write_table_dataset( def write_hdf5_scalar_matrices( h5_file: h5py.File, - scalars: Mapping[str, Sequence[np.ndarray]], + scalars: Mapping[str, np.ndarray | Sequence[np.ndarray]], sources_by_scalar: Mapping[str, Sequence[str]], *, storage_dtype: str, @@ -71,32 +99,37 @@ def write_hdf5_scalar_matrices( shuffle: bool, chunk_voxels: int, target_chunk_mb: float, -) -> None: - """Write per-scalar matrices into an open HDF5 file.""" - for scalar_name, rows in scalars.items(): - n_files = len(rows) +) -> dict[str, dict[str, object]]: + """Write each scalar as one full matrix into an open HDF5 file.""" + metrics: dict[str, dict[str, object]] = {} + for scalar_name, matrix in scalars.items(): + scalar_name = storage_utils.validate_scalar_name(scalar_name) + n_files = len(matrix) if n_files == 0: continue - n_elements = rows[0].shape[0] - dataset = h5_storage.create_empty_scalar_matrix_dataset( + scalar_metrics: dict[str, object] = {} + h5_storage.create_scalar_matrix_dataset( h5_file, f'scalars/{scalar_name}/values', - n_files, - n_elements, + matrix, + sources_by_scalar[scalar_name], storage_dtype=storage_dtype, compression=compression, compression_level=compression_level, shuffle=shuffle, chunk_voxels=chunk_voxels, target_chunk_mb=target_chunk_mb, - sources_list=sources_by_scalar[scalar_name], + metrics=scalar_metrics, ) - h5_storage.write_rows_in_column_stripes(dataset, rows) + scalar_metrics['backend'] = 'hdf5' + scalar_metrics['scalar_name'] = str(scalar_name) + metrics[str(scalar_name)] = scalar_metrics + return metrics def write_tiledb_scalar_matrices( output_dir: str | Path, - scalars: Mapping[str, Sequence[np.ndarray]], + scalars: Mapping[str, np.ndarray | Sequence[np.ndarray]], sources_by_scalar: Mapping[str, Sequence[str]], *, storage_dtype: str, @@ -106,35 +139,40 @@ def write_tiledb_scalar_matrices( chunk_voxels: int, target_chunk_mb: float, write_column_name_arrays: bool = False, -) -> None: - """Write per-scalar matrices into a TileDB directory.""" +) -> dict[str, dict[str, object]]: + """Write each scalar in one dense submission into a TileDB directory.""" output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) + metrics: dict[str, dict[str, object]] = {} - for scalar_name, rows in scalars.items(): - n_files = len(rows) + for scalar_name, matrix in scalars.items(): + scalar_name = storage_utils.validate_scalar_name(scalar_name) + n_files = len(matrix) if n_files == 0: continue - n_elements = rows[0].shape[0] dataset_path = f'scalars/{scalar_name}/values' - tiledb_storage.create_empty_scalar_matrix_array( + scalar_metrics: dict[str, object] = {} + tiledb_storage.create_scalar_matrix_array( str(output_path), dataset_path, - n_files, - n_elements, + matrix, + sources_by_scalar[scalar_name], storage_dtype=storage_dtype, compression=compression, compression_level=compression_level, shuffle=shuffle, tile_voxels=chunk_voxels, target_tile_mb=target_chunk_mb, - sources_list=sources_by_scalar[scalar_name], + metrics=scalar_metrics, ) if write_column_name_arrays: tiledb_storage.write_column_names( str(output_path), scalar_name, sources_by_scalar[scalar_name] ) - tiledb_storage.write_rows_in_column_stripes(str(output_path / dataset_path), rows) + scalar_metrics['backend'] = 'tiledb' + scalar_metrics['scalar_name'] = str(scalar_name) + metrics[str(scalar_name)] = scalar_metrics + return metrics def write_hdf5_parcel_arrays( diff --git a/src/modelarrayio/storage/h5_storage.py b/src/modelarrayio/storage/h5_storage.py index 33a0d3c..c07f8e7 100644 --- a/src/modelarrayio/storage/h5_storage.py +++ b/src/modelarrayio/storage/h5_storage.py @@ -3,17 +3,43 @@ from __future__ import annotations import logging +import operator +import os +import time +import zlib +from collections.abc import MutableMapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from itertools import islice, product import h5py import numpy as np import pandas as pd -from tqdm import tqdm -from tqdm.contrib.logging import logging_redirect_tqdm from modelarrayio.storage import utils as storage_utils logger = logging.getLogger(__name__) +_H5_FILTER_DEFLATE = 1 +_H5_FILTER_SHUFFLE = 2 + + +class _DirectChunkUnsupported(RuntimeError): + """Raised when a dataset's filter pipeline cannot be reproduced safely.""" + + +def _write_column_names_to_group( + group: h5py.Group, + sources: pd.Series | Sequence[str], +) -> h5py.Dataset: + """Write source names beside the values dataset in *group*.""" + values = np.array(storage_utils.normalize_column_names(sources), dtype=object) + return group.create_dataset( + 'column_names', + data=values, + dtype=h5py.string_dtype(encoding='utf-8'), + shape=(len(values),), + ) + def resolve_dtype(storage_dtype): """Resolve a storage dtype to a supported NumPy floating type. @@ -116,6 +142,7 @@ def create_scalar_matrix_dataset( shuffle=True, chunk_voxels=0, target_chunk_mb=2.0, + metrics: MutableMapping[str, object] | None = None, ): """Create a scalar matrix dataset in an HDF5 file. @@ -150,8 +177,7 @@ def create_scalar_matrix_dataset( storage_np_dtype = resolve_dtype(storage_dtype) comp, comp_opts, use_shuffle = resolve_compression(compression, compression_level, shuffle) - if stacked_values.dtype != storage_np_dtype: - stacked_values = stacked_values.astype(storage_np_dtype) + stacked_values = _coerce_full_matrix(stacked_values, dtype=storage_np_dtype) n_files, n_elements = stacked_values.shape chunk_shape = compute_chunk_shape_full_subjects( @@ -176,10 +202,13 @@ def create_scalar_matrix_dataset( shuffle=use_shuffle, ) logger.info('Writing full dataset %s to HDF5 (this may take a while)...', dataset_path) - dset[...] = stacked_values + write_metrics = write_full_matrix(dset, stacked_values) logger.info('Finished writing dataset %s', dataset_path) if sources_list is not None: - dset.attrs['column_names'] = list(sources_list) + _write_column_names_to_group(dset.parent, sources_list) + if metrics is not None: + metrics.update(write_metrics) + metrics['dataset_path'] = str(dataset_path) return dset @@ -194,7 +223,7 @@ def create_empty_scalar_matrix_dataset( shuffle=True, chunk_voxels=0, target_chunk_mb=2.0, - sources_list=None | pd.Series | list, + sources_list: pd.Series | list | None = None, ): """Create an empty scalar matrix dataset in an HDF5 file. @@ -253,9 +282,7 @@ def create_empty_scalar_matrix_dataset( shuffle=use_shuffle, ) if sources_list is not None: - # dataset_path is e.g. 'scalars/FA/values'; extract the scalar name segment - scalar_name = dataset_path.split('/')[1] if dataset_path.count('/') >= 2 else dataset_path - write_column_names(h5file, scalar_name, sources_list) + _write_column_names_to_group(dset.parent, sources_list) return dset @@ -271,24 +298,265 @@ def write_column_names(h5_file: h5py.File, scalar: str, sources: pd.Series | lis sources : :obj:`list` Sources list. """ - values = np.array(storage_utils.normalize_column_names(sources), dtype=object) + scalar = storage_utils.validate_scalar_name(scalar) grp = h5_file.require_group(f'scalars/{scalar}') + _write_column_names_to_group(grp, sources) - # Variable-length UTF-8 string dtype - vlen_str = h5py.string_dtype(encoding='utf-8') - # Create 1-D dataset of strings - grp.create_dataset( - 'column_names', - data=values, - dtype=vlen_str, - shape=(len(values),), +# Shared with the other backend; kept under the private name for callers and tests. +_coerce_full_matrix = storage_utils.coerce_full_matrix + + +def _hdf5_byte_shuffle(raw: bytes, element_size: int) -> bytes: + """Apply the HDF5 shuffle filter's byte transpose exactly.""" + if element_size <= 0: + raise ValueError('element_size must be positive') + byte_values = np.frombuffer(raw, dtype=np.uint8) + if byte_values.size % element_size: + raise ValueError('raw chunk size must be divisible by element_size') + return byte_values.reshape(-1, element_size).T.tobytes(order='C') + + +def _chunk_offsets(shape: tuple[int, int], chunk_shape: tuple[int, int]) -> list[tuple[int, int]]: + return list( + product( + range(0, shape[0], chunk_shape[0]), + range(0, shape[1], chunk_shape[1]), + ) + ) + + +def _encode_gzip_chunk( + matrix: np.ndarray, + offset: tuple[int, int], + chunk_shape: tuple[int, int], + fill_value: object, + compression_level: int, + shuffle: bool, +) -> tuple[tuple[int, int], bytes, float, float]: + """Pad, optionally byte-shuffle, and gzip one HDF5 chunk.""" + wall_started = time.perf_counter() + cpu_started = time.thread_time() + row_start, column_start = offset + row_stop = min(row_start + chunk_shape[0], matrix.shape[0]) + column_stop = min(column_start + chunk_shape[1], matrix.shape[1]) + source = matrix[row_start:row_stop, column_start:column_stop] + + if source.shape == chunk_shape: + chunk = np.ascontiguousarray(source) + else: + chunk = np.empty(chunk_shape, dtype=matrix.dtype, order='C') + chunk.fill(fill_value) + chunk[: source.shape[0], : source.shape[1]] = source + + raw = chunk.tobytes(order='C') + if shuffle: + raw = _hdf5_byte_shuffle(raw, matrix.dtype.itemsize) + compressed = zlib.compress(raw, level=compression_level) + return ( + offset, + compressed, + time.perf_counter() - wall_started, + time.thread_time() - cpu_started, ) +def _available_cpu_count() -> int: + """Return CPUs available to this process, respecting affinity/cgroup limits.""" + process_cpu_count = getattr(os, 'process_cpu_count', None) + if process_cpu_count is not None: + count = process_cpu_count() + if count: + return max(1, int(count)) + + sched_getaffinity = getattr(os, 'sched_getaffinity', None) + if sched_getaffinity is not None: + try: + return max(1, len(sched_getaffinity(0))) + except OSError: + pass + + return max(1, os.cpu_count() or 1) + + +def _resolve_compression_workers(compression_workers: int | None) -> int: + """Resolve and validate the direct-chunk compression worker count.""" + configured_workers: object = compression_workers + if configured_workers is None: + configured_workers = os.environ.get('MODELARRAYIO_HDF5_COMPRESSION_WORKERS') + if configured_workers is None: + return _available_cpu_count() + try: + requested_workers = int(configured_workers) + except ValueError as exc: + raise ValueError( + 'MODELARRAYIO_HDF5_COMPRESSION_WORKERS must be a positive integer' + ) from exc + else: + if isinstance(configured_workers, bool): + raise ValueError('HDF5 compression workers must be a positive integer') + try: + requested_workers = operator.index(configured_workers) + except TypeError as exc: + raise ValueError('HDF5 compression workers must be a positive integer') from exc + + if requested_workers < 1: + raise ValueError('HDF5 compression workers must be a positive integer') + return requested_workers + + +def _validate_direct_chunk_pipeline(dset: h5py.Dataset) -> None: + """Verify that direct bytes can reproduce the dataset's full filter pipeline.""" + if dset.chunks is None: + raise _DirectChunkUnsupported('the dataset is not chunked') + if dset.compression != 'gzip': + raise _DirectChunkUnsupported( + f'compression {dset.compression!r} is not supported by the direct-chunk writer' + ) + if not hasattr(dset.id, 'write_direct_chunk'): + raise _DirectChunkUnsupported('h5py does not expose write_direct_chunk') + + expected: list[tuple[int, tuple[int, ...]]] = [] + if dset.shuffle: + expected.append((_H5_FILTER_SHUFFLE, (dset.dtype.itemsize,))) + expected.append((_H5_FILTER_DEFLATE, (int(dset.compression_opts),))) + + creation_properties = dset.id.get_create_plist() + actual = [] + for index in range(creation_properties.get_nfilters()): + filter_id, _flags, client_data, _name = creation_properties.get_filter(index) + actual.append((int(filter_id), tuple(int(value) for value in client_data))) + if actual != expected: + raise _DirectChunkUnsupported( + f'the declared HDF5 filter pipeline {actual!r} does not match {expected!r}' + ) + + +def _write_direct_gzip_chunks( + dset: h5py.Dataset, + matrix: np.ndarray, + *, + compression_workers: int | None, +) -> dict[str, object]: + """Compress chunks in worker threads and submit bytes from the calling thread.""" + _validate_direct_chunk_pipeline(dset) + chunk_shape = tuple(int(value) for value in dset.chunks) + offsets = _chunk_offsets(matrix.shape, chunk_shape) + requested_workers = _resolve_compression_workers(compression_workers) + worker_count = max(1, min(requested_workers, len(offsets))) + compression_seconds = 0.0 + compression_cpu_seconds = 0.0 + submission_seconds = 0.0 + + offset_iterator = iter(offsets) + with ThreadPoolExecutor( + max_workers=worker_count, + thread_name_prefix='modelarrayio-hdf5', + ) as executor: + while batch := tuple(islice(offset_iterator, worker_count)): + compression_started = time.perf_counter() + futures = [ + executor.submit( + _encode_gzip_chunk, + matrix, + offset, + chunk_shape, + dset.fillvalue, + int(dset.compression_opts), + bool(dset.shuffle), + ) + for offset in batch + ] + encoded_chunks = [future.result() for future in futures] + compression_seconds += time.perf_counter() - compression_started + compression_cpu_seconds += sum(result[3] for result in encoded_chunks) + + for offset, compressed, _worker_wall_seconds, _worker_cpu_seconds in encoded_chunks: + submission_started = time.perf_counter() + dset.id.write_direct_chunk(offset, compressed, filter_mask=0) + submission_seconds += time.perf_counter() - submission_started + + return { + 'compression_seconds': compression_seconds, + 'compression_cpu_seconds': compression_cpu_seconds, + 'submission_seconds': submission_seconds, + 'compression_workers': worker_count, + 'chunk_count': len(offsets), + } + + +def write_full_matrix( + dset: h5py.Dataset, + values: np.ndarray | Sequence[np.ndarray], + *, + compression_workers: int | None = None, +) -> dict[str, object]: + """Write a complete matrix, using parallel direct-chunk gzip when safe. + + The HDF5 handle is only accessed by the calling thread. Worker threads receive + NumPy slices and return fully filtered byte buffers. Unsupported filters and + worker failures are logged and safely rewritten through ordinary h5py assignment. + """ + matrix = _coerce_full_matrix(values, dtype=dset.dtype, shape=dset.shape) + started = time.perf_counter() + metrics: dict[str, object] = { + 'matrix_shape': [int(value) for value in matrix.shape], + 'dtype': matrix.dtype.name, + 'chunk_shape': [int(value) for value in dset.chunks] if dset.chunks else None, + 'compression': dset.compression or 'none', + 'compression_level': (int(dset.compression_opts) if dset.compression == 'gzip' else None), + 'shuffle': bool(dset.shuffle), + 'direct_chunk': False, + 'fallback': False, + 'fallback_reason': None, + 'compression_seconds': 0.0, + 'compression_cpu_seconds': 0.0, + 'submission_seconds': 0.0, + 'compression_workers': 0, + 'chunk_count': 0, + } + + if dset.compression is None: + submission_started = time.perf_counter() + dset[...] = matrix + metrics['submission_seconds'] = time.perf_counter() - submission_started + else: + # Configuration errors are user errors, not direct-writer failures. Resolve + # them before entering the fallback boundary so they cannot be hidden by a + # successful but unexpectedly serial h5py write. + resolved_workers = _resolve_compression_workers(compression_workers) + try: + direct_metrics = _write_direct_gzip_chunks( + dset, + matrix, + compression_workers=resolved_workers, + ) + # Any worker or filter failure must fall back safely; the debug log keeps the traceback. + except Exception as exc: + reason = str(exc) or type(exc).__name__ + logger.warning( + 'Falling back to ordinary HDF5 dataset assignment for %s: %s', + dset.name, + reason, + ) + logger.debug('Direct-chunk write failure details', exc_info=True) + submission_started = time.perf_counter() + dset[...] = matrix + metrics['submission_seconds'] = time.perf_counter() - submission_started + metrics['compression_seconds'] = None + metrics['compression_cpu_seconds'] = None + metrics['fallback'] = True + metrics['fallback_reason'] = reason + else: + metrics.update(direct_metrics) + metrics['direct_chunk'] = True + + metrics['write_seconds'] = time.perf_counter() - started + return metrics + + def write_rows_in_column_stripes(dset, rows): - """Fill a 2D HDF5 dataset by buffering column-aligned stripes to minimize - chunk recompression, using about one chunk's worth of memory. + """Compatibility wrapper that now writes rows as one full matrix. Parameters ---------- @@ -296,39 +564,11 @@ def write_rows_in_column_stripes(dset, rows): Target dataset with shape (n_files, n_elements) and chunking set. rows : Sequence[np.ndarray] List/sequence of 1D arrays, one per subject, length == n_elements. - Each will be cast on write to dset.dtype if needed. - """ - n_files, n_elements = dset.shape - if len(rows) != n_files: - raise ValueError('rows length does not match dataset subjects dimension') - stripe_width = dset.chunks[1] if dset.chunks is not None else max(1, n_elements // 8) - logger.info( - 'Stripe-writing dataset %s with stripe width=%d (chunks=%s)', - dset.name, - stripe_width, - str(dset.chunks), - ) + Each will be cast into one C-contiguous staging matrix if needed. - buf = np.empty((n_files, stripe_width), dtype=dset.dtype) - with logging_redirect_tqdm(): - for start in tqdm( - range(0, n_elements, stripe_width), - bar_format=( - '{percentage:3.0f}% {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]' - ), - ascii=True, - mininterval=max(1, (n_elements / stripe_width) // 200), - ): - end = min(start + stripe_width, n_elements) - width = end - start - if width != stripe_width: - # resize buffer view on last partial stripe - buf_view = buf[:, :width] - else: - buf_view = buf - for i, row in enumerate(rows): - # slice is contiguous; cast on assignment if needed - buf_view[i, :] = row[start:end] - logger.debug('Writing stripe [%d:%d] to %s', start, end, dset.name) - dset[:, start:end] = buf_view - logger.info('Finished stripe-writing dataset %s', dset.name) + Returns + ------- + dict + JSON-serializable write metrics from :func:`write_full_matrix`. + """ + return write_full_matrix(dset, rows) diff --git a/src/modelarrayio/storage/tiledb_storage.py b/src/modelarrayio/storage/tiledb_storage.py index 7ca212b..b0a035a 100644 --- a/src/modelarrayio/storage/tiledb_storage.py +++ b/src/modelarrayio/storage/tiledb_storage.py @@ -5,7 +5,8 @@ import json import logging import os -from collections.abc import Sequence +import time +from collections.abc import MutableMapping, Sequence import numpy as np import tiledb @@ -34,27 +35,39 @@ def resolve_dtype(storage_dtype): return storage_utils.resolve_dtype(storage_dtype) +def _resolve_compression_settings( + compression: str | None, + compression_level: int | None, +) -> tuple[str, int | None]: + """Return the effective TileDB codec and level used by the schema.""" + comp = 'none' if compression is None else str(compression).lower() + if comp == 'none': + return 'none', None + if comp not in {'zstd', 'gzip'}: + logger.warning("Unknown compression '%s' for TileDB; disabling compression.", comp) + return 'none', None + + default_level = 5 if comp == 'zstd' else 4 + try: + level = int(compression_level) if compression_level is not None else default_level + except (TypeError, ValueError): + level = default_level + return comp, level + + def _build_filter_list(compression: str | None, compression_level: int | None, shuffle: bool): + compression, compression_level = _resolve_compression_settings(compression, compression_level) filters = [] - if shuffle: + # A standalone byte-shuffle still consumes CPU even though it cannot + # reduce output size. Keep ``--compression none`` as the true no-filter + # path, matching HDF5 and the CLI description. + if shuffle and compression != 'none': # ByteShuffle works well for float data; BitShuffle is also available filters.append(tiledb.ByteShuffleFilter()) - if compression is None or str(compression).lower() == 'none': - pass - else: - comp = str(compression).lower() - level = None - try: - level = int(compression_level) if compression_level is not None else None - except (TypeError, ValueError): - level = None - if comp == 'zstd': - filters.append(tiledb.ZstdFilter(level=level if level is not None else 5)) - elif comp == 'gzip': - filters.append(tiledb.GzipFilter(level=level if level is not None else 4)) - else: - # Fallback: no compression if an unknown codec is provided - logger.warning("Unknown compression '%s' for TileDB; disabling compression.", comp) + if compression == 'zstd': + filters.append(tiledb.ZstdFilter(level=compression_level)) + elif compression == 'gzip': + filters.append(tiledb.GzipFilter(level=compression_level)) return tiledb.FilterList(filters) @@ -120,6 +133,7 @@ def create_scalar_matrix_array( shuffle=True, tile_voxels=0, target_tile_mb=2.0, + metrics: MutableMapping[str, object] | None = None, ): """Create a scalar matrix array in a TileDB directory. @@ -152,8 +166,7 @@ def create_scalar_matrix_array( URI of the created array. """ storage_np_dtype = resolve_dtype(storage_dtype) - if stacked_values.dtype != storage_np_dtype: - stacked_values = stacked_values.astype(storage_np_dtype) + stacked_values = _coerce_full_matrix(stacked_values, dtype=storage_np_dtype) n_files, n_elements = stacked_values.shape tile_shape = compute_tile_shape_full_subjects( @@ -171,7 +184,11 @@ def create_scalar_matrix_array( name='items', domain=(0, n_elements - 1), tile=tile_shape[1], dtype=_DIM_DTYPE ) dom = tiledb.Domain(dim_subjects, dim_items) - attr_filters = _build_filter_list(compression, compression_level, shuffle) + effective_compression, effective_level = _resolve_compression_settings( + compression, compression_level + ) + effective_shuffle = bool(shuffle and effective_compression != 'none') + attr_filters = _build_filter_list(effective_compression, effective_level, shuffle) attr_values = tiledb.Attr(name='values', dtype=storage_np_dtype, filters=attr_filters) schema = tiledb.ArraySchema(domain=dom, attrs=[attr_values], sparse=False) @@ -188,9 +205,12 @@ def create_scalar_matrix_array( tiledb.remove(uri) tiledb.Array.create(uri, schema) - logger.info('Writing full array %s to TileDB (this may take a while)...', uri) + logger.info('Writing full array %s to TileDB in one dense submission...', uri) + write_started = time.perf_counter() with tiledb.open(uri, 'w') as A: + submission_started = time.perf_counter() A[:] = {'values': stacked_values} + submission_seconds = time.perf_counter() - submission_started if sources_list is not None: try: A.meta['column_names'] = json.dumps( @@ -199,6 +219,28 @@ def create_scalar_matrix_array( except (TypeError, ValueError, tiledb.TileDBError): # Fallback without metadata if serialization fails logger.warning('Failed to write column_names metadata for %s', uri) + write_seconds = time.perf_counter() - write_started + if metrics is not None: + fragment_count = len(tiledb.array_fragments(uri)) + output_file_count, output_bytes = _local_output_stats(uri) + metrics.update( + { + 'dataset_path': str(dataset_path), + 'matrix_shape': [int(value) for value in stacked_values.shape], + 'dtype': stacked_values.dtype.name, + 'tile_shape': [int(value) for value in tile_shape], + 'compression': effective_compression, + 'compression_level': effective_level, + 'shuffle': effective_shuffle, + 'compression_seconds': None, + 'submission_seconds': submission_seconds, + 'write_seconds': write_seconds, + 'dense_submissions': 1, + 'fragment_count': fragment_count, + 'output_file_count': output_file_count, + 'output_bytes': output_bytes, + } + ) logger.info('Finished writing array %s', uri) return uri @@ -291,9 +333,73 @@ def create_empty_scalar_matrix_array( return uri +# Shared with the other backend; kept under the private name for callers and tests. +_coerce_full_matrix = storage_utils.coerce_full_matrix + + +def _local_output_stats(uri: str) -> tuple[int | None, int | None]: + """Return local file count and bytes without treating remote URIs as paths.""" + if '://' in uri or not os.path.exists(uri): + return None, None + file_count = 0 + byte_count = 0 + for root, _directories, files in os.walk(uri): + for filename in files: + file_count += 1 + try: + byte_count += os.path.getsize(os.path.join(root, filename)) + except OSError: + logger.debug('Could not stat TileDB output file %s', filename, exc_info=True) + return file_count, byte_count + + +def write_full_matrix(uri: str, values: np.ndarray | Sequence[np.ndarray]) -> dict[str, object]: + """Write an existing dense TileDB array in exactly one submission.""" + with tiledb.open(uri, 'r') as array_info: + schema = array_info.schema + if schema.sparse: + raise ValueError(f"TileDB array '{uri}' must be dense") + if schema.domain.ndim != 2: + raise ValueError( + f"TileDB array '{uri}' must have two dimensions, got {schema.domain.ndim}" + ) + if schema.nattr != 1 or not schema.has_attr('values'): + raise ValueError( + f"TileDB array '{uri}' must have exactly one attribute named 'values'" + ) + + domain = schema.domain + shape = tuple( + int(domain.dim(index).domain[1] - domain.dim(index).domain[0] + 1) + for index in range(2) + ) + attr_dtype = schema.attr('values').dtype + tile_shape = tuple(int(domain.dim(index).tile) for index in range(2)) + + matrix = _coerce_full_matrix(values, dtype=attr_dtype, shape=shape) + write_started = time.perf_counter() + with tiledb.open(uri, 'w') as array: + submission_started = time.perf_counter() + array[:] = {'values': matrix} + submission_seconds = time.perf_counter() - submission_started + write_seconds = time.perf_counter() - write_started + output_file_count, output_bytes = _local_output_stats(uri) + return { + 'matrix_shape': [int(value) for value in matrix.shape], + 'dtype': matrix.dtype.name, + 'tile_shape': [int(value) for value in tile_shape], + 'compression_seconds': None, + 'submission_seconds': submission_seconds, + 'write_seconds': write_seconds, + 'dense_submissions': 1, + 'fragment_count': len(tiledb.array_fragments(uri)), + 'output_file_count': output_file_count, + 'output_bytes': output_bytes, + } + + def write_rows_in_column_stripes(uri: str, rows: Sequence[np.ndarray]): - """Fill a 2D TileDB dense array by buffering column-aligned stripes to minimize - tile writes, using about one tile's worth of memory. + """Compatibility wrapper that now writes rows in one dense submission. Parameters ---------- @@ -301,34 +407,14 @@ def write_rows_in_column_stripes(uri: str, rows: Sequence[np.ndarray]): Target array URI with shape (n_files, n_elements). rows : Sequence[np.ndarray] List/sequence of 1D arrays, one per subject, length == n_elements. - Each will be cast on write to array attr dtype if needed. + Each will be cast into one C-contiguous staging matrix if needed. + + Returns + ------- + dict + JSON-serializable write metrics from :func:`write_full_matrix`. """ - with tiledb.open(uri, 'r') as Ainfo: - dom = Ainfo.schema.domain - n_files = dom.dim(0).domain[1] - dom.dim(0).domain[0] + 1 - n_elements = dom.dim(1).domain[1] - dom.dim(1).domain[0] + 1 - attr_dtype = Ainfo.schema.attr(0).dtype - - if len(rows) != n_files: - raise ValueError('rows length does not match array subjects dimension') - - # Try to align stripe width to the items tile for best throughput - with tiledb.open(uri, 'r') as Ainfo2: - items_tile = Ainfo2.schema.domain.dim(1).tile - stripe_width = items_tile if items_tile is not None else max(1, n_elements // 8) - - buf = np.empty((n_files, stripe_width), dtype=attr_dtype) - for start in range(0, n_elements, stripe_width): - end = min(start + stripe_width, n_elements) - width = end - start - if width != stripe_width: - buf_view = buf[:, :width] - else: - buf_view = buf - for i, row in enumerate(rows): - buf_view[i, :] = row[start:end] - with tiledb.open(uri, 'w') as A: - A[:, start:end] = {'values': buf_view} + return write_full_matrix(uri, rows) def write_parcel_names(base_uri: str, array_path: str, names: Sequence[str]): @@ -380,6 +466,7 @@ def write_column_names(base_uri: str, scalar: str, sources: Sequence[str]): sources : :obj:`list` Sources list. """ + scalar = storage_utils.validate_scalar_name(scalar) sources = storage_utils.normalize_column_names(sources) uri = os.path.join(base_uri, 'scalars', scalar, 'column_names') _ensure_parent_group(uri) diff --git a/src/modelarrayio/storage/utils.py b/src/modelarrayio/storage/utils.py index a2afbed..df09e40 100644 --- a/src/modelarrayio/storage/utils.py +++ b/src/modelarrayio/storage/utils.py @@ -2,6 +2,7 @@ from __future__ import annotations +import operator from collections.abc import Sequence import numpy as np @@ -51,12 +52,25 @@ def compute_full_subject_chunk_shape( f'n_files={n_files}, n_elements={n_elements}' ) + try: + if isinstance(item_chunk, bool): + raise TypeError('a boolean is not an item chunk') # bool passes operator.index + resolved_item_chunk = operator.index(item_chunk) + except TypeError as exc: + raise ValueError(f'item chunk must be a non-negative integer, got {item_chunk!r}') from exc + if resolved_item_chunk < 0: + raise ValueError(f'item chunk must be non-negative, got {resolved_item_chunk}') + subjects_per_chunk = n_files - if int(item_chunk) > 0: - items_per_chunk = min(int(item_chunk), n_elements) + if resolved_item_chunk > 0: + items_per_chunk = min(resolved_item_chunk, n_elements) else: bytes_per_value = np.dtype(resolve_dtype(storage_np_dtype)).itemsize target_bytes = float(target_chunk_mb) * 1024.0 * 1024.0 + if not np.isfinite(target_bytes) or target_bytes <= 0: + raise ValueError( + f'target chunk size must be a positive finite value, got {target_chunk_mb!r}' + ) items_per_chunk = max(1, int(target_bytes / (bytes_per_value * subjects_per_chunk))) items_per_chunk = min(items_per_chunk, n_elements) @@ -68,3 +82,104 @@ def normalize_column_names(sources: Sequence[str] | pd.Series) -> list[str]: if isinstance(sources, pd.Series): return sources.astype(str).tolist() return [str(source) for source in sources] + + +def validate_scalar_name(scalar_name: object) -> str: + """Return a scalar identifier that is safe as one storage path segment.""" + name = str(scalar_name) + if not name.strip(): + raise ValueError('Scalar names must not be empty or whitespace-only.') + if '\x00' in name or '/' in name or '\\' in name or name in {'.', '..'}: + raise ValueError( + f'Unsafe scalar name {name!r}: names must be one path segment without ' + "'/', '\\', NUL, '.' or '..'." + ) + return name + + +def resolve_backend_compression( + backend: str, + compression: str | None, + compression_level: int | None, +) -> tuple[str, int | None]: + """Resolve a supported backend codec and its effective level.""" + backend_name = str(backend).lower() + supported = { + 'hdf5': {'gzip', 'lzf', 'none'}, + 'tiledb': {'gzip', 'zstd', 'none'}, + } + if backend_name not in supported: + raise ValueError(f'Unsupported storage backend {backend!r}.') + + codec = ( + ('gzip' if backend_name == 'hdf5' else 'zstd') + if compression is None + else str(compression).lower() + ) + if codec not in supported[backend_name]: + choices = ', '.join(sorted(supported[backend_name])) + raise ValueError( + f'Compression {codec!r} is not supported by {backend_name}; choose one of: {choices}.' + ) + if codec in {'none', 'lzf'}: + return codec, None + + default_level = 4 if codec == 'gzip' else 5 + if compression_level is None: + return codec, default_level + try: + if isinstance(compression_level, bool): + raise TypeError('a boolean is not a compression level') # bool passes operator.index + level = operator.index(compression_level) + except TypeError as exc: + raise ValueError( + f'compression level must be an integer, got {compression_level!r}' + ) from exc + if codec == 'gzip' and not 0 <= level <= 9: + raise ValueError(f'gzip compression level must be between 0 and 9, got {level}.') + return codec, level + + +def coerce_full_matrix( + values: np.ndarray | Sequence[np.ndarray], + *, + dtype: str | np.dtype | type[np.floating], + shape: tuple[int, int] | None = None, +) -> np.ndarray: + """Return one C-contiguous matrix, preserving an already-staged matrix. + + A 2-D array of the requested dtype is returned as-is (no copy). A sequence + of 1-D rows is stacked into a new matrix. ``shape`` is enforced when given. + """ + target_dtype = np.dtype(dtype) + if isinstance(values, np.ndarray): + if values.ndim != 2: + raise ValueError(f'Expected a 2D scalar matrix, got shape {values.shape!r}') + matrix = values + if matrix.dtype != target_dtype: + matrix = matrix.astype(target_dtype, copy=False) + if not matrix.flags.c_contiguous: + matrix = np.ascontiguousarray(matrix) + else: + rows = list(values) + if shape is None: + if not rows: + raise ValueError('Cannot build a scalar matrix from an empty row sequence') + first = np.asarray(rows[0]) + if first.ndim != 1: + raise ValueError(f'Expected 1D scalar rows, got shape {first.shape!r}') + shape = (len(rows), first.shape[0]) + if len(rows) != shape[0]: + raise ValueError('rows length does not match the subjects dimension') + matrix = np.empty(shape, dtype=target_dtype, order='C') + for index, row in enumerate(rows): + row_array = np.asarray(row) + if row_array.ndim != 1 or row_array.shape[0] != shape[1]: + raise ValueError( + f'Row {index} has shape {row_array.shape!r}; expected ({shape[1]},)' + ) + matrix[index] = row_array + + if shape is not None and matrix.shape != shape: + raise ValueError(f'Matrix shape {matrix.shape!r} does not match expected shape {shape!r}') + return matrix diff --git a/src/modelarrayio/utils/cifti.py b/src/modelarrayio/utils/cifti.py index 6bf6dc7..636f02c 100644 --- a/src/modelarrayio/utils/cifti.py +++ b/src/modelarrayio/utils/cifti.py @@ -2,19 +2,111 @@ from __future__ import annotations -from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor, as_completed -from pathlib import Path +import hashlib +import logging +import os +import time +from dataclasses import dataclass +from io import BytesIO import nibabel as nb import numpy as np import pandas as pd -from tqdm import tqdm -from modelarrayio.utils.s3_utils import load_nibabel - - -def extract_cifti_scalar_data(cifti_file, reference_brain_names=None): +from modelarrayio.utils.loading import ( + LazyScalarMatrices, + LoadedRow, + SourceLoadError, + resolve_loader_workers, + stage_source_rows, +) +from modelarrayio.utils.misc import build_scalar_sources +from modelarrayio.utils.s3_utils import is_s3_path, load_nibabel, release_image_buffers + +logger = logging.getLogger(__name__) + +# NIfTI-2 fixed header length in bytes. The CIFTI XML extension and the data +# block follow it, and ``vox_offset`` in this header points past both. +_NIFTI2_HEADER_SIZE = 540 + +# Environment switch to disable the raw greyordinate fast path (diagnostics +# and strict-validation runs). Any of 0/false/no/off turns it off. +_RAW_READ_ENV = 'MODELARRAYIO_CIFTI_RAW' + + +def _hash_signature_value(hasher, value): + """Feed nested axis state into a stable digest without materializing labels.""" + if value is None: + hasher.update(b'none;') + elif isinstance(value, dict): + hasher.update(b'dict{') + for key in sorted(value, key=str): + _hash_signature_value(hasher, str(key)) + _hash_signature_value(hasher, value[key]) + hasher.update(b'}') + elif isinstance(value, np.ndarray): + hasher.update(f'array:{value.dtype.str}:{value.shape};'.encode()) + if value.dtype.hasobject: + for item in value.flat: + _hash_signature_value(hasher, item) + elif value.dtype.kind in {'U', 'S'}: + for item in value.flat: + _hash_signature_value(hasher, str(item)) + else: + hasher.update(np.ascontiguousarray(value).tobytes()) + elif isinstance(value, (list, tuple)): + hasher.update(f'seq:{len(value)}['.encode()) + for item in value: + _hash_signature_value(hasher, item) + hasher.update(b']') + elif isinstance(value, bytes): + hasher.update(len(value).to_bytes(8, 'little')) + hasher.update(value) + else: + encoded = str(value).encode('utf-8') + hasher.update(len(encoded).to_bytes(8, 'little')) + hasher.update(encoded) + + +def cifti_spatial_signature(cifti_file): + """Return a short signature of all non-scalar CIFTI axis properties.""" + if not hasattr(cifti_file, 'get_fdata'): + cifti = load_nibabel(cifti_file, cifti=True) + try: + return cifti_spatial_signature(cifti) + finally: + release_image_buffers(cifti) + cifti = cifti_file + hasher = hashlib.blake2b(digest_size=20) + spatial_count = 0 + for axis_index in range(cifti.ndim): + axis = cifti.header.get_axis(axis_index) + if isinstance(axis, nb.cifti2.cifti2_axes.ScalarAxis): + continue + spatial_count += 1 + _hash_signature_value(hasher, axis_index) + _hash_signature_value(hasher, type(axis).__name__) + if isinstance(axis, nb.cifti2.cifti2_axes.BrainModelAxis): + for attr in ('voxel', 'vertex', '_name', 'nvertices', '_affine', '_volume_shape'): + _hash_signature_value(hasher, getattr(axis, attr, None)) + elif isinstance(axis, nb.cifti2.cifti2_axes.ParcelsAxis): + for attr in ('name', 'voxels', 'vertices', '_affine', '_volume_shape', 'nvertices'): + _hash_signature_value(hasher, getattr(axis, attr, None)) + else: + _hash_signature_value(hasher, vars(axis)) + if spatial_count == 0: + raise ValueError('CIFTI image does not contain a supported spatial axis.') + return spatial_count, hasher.hexdigest() + + +def extract_cifti_scalar_data( + cifti_file, + reference_brain_names=None, + *, + storage_dtype='float32', + reference_signature=None, + spatial_signature=None, +): """Load a scalar cifti file and get its data and mapping. Supports dscalar (.dscalar.nii), parcellated scalar (.pscalar.nii), and @@ -43,7 +135,23 @@ def extract_cifti_scalar_data(cifti_file, reference_brain_names=None): are the parcel names. For pconn files these are the row-parcel names repeated ``n_col_parcels`` times (one entry per flattened element). """ - cifti = cifti_file if hasattr(cifti_file, 'get_fdata') else nb.load(Path(cifti_file)) + if not hasattr(cifti_file, 'get_fdata'): + cifti = load_nibabel(cifti_file, cifti=True) + try: + return extract_cifti_scalar_data( + cifti, + reference_brain_names, + storage_dtype=storage_dtype, + reference_signature=reference_signature, + spatial_signature=spatial_signature, + ) + finally: + release_image_buffers(cifti) + cifti = cifti_file + if reference_signature is not None: + signature = spatial_signature or cifti_spatial_signature(cifti) + if signature != reference_signature: + raise ValueError(f'Inconsistent spatial axis in CIFTI file {cifti_file!r}.') cifti_hdr = cifti.header axes = [cifti_hdr.get_axis(i) for i in range(cifti.ndim)] if len(axes) != 2: @@ -58,62 +166,79 @@ def extract_cifti_scalar_data(cifti_file, reference_brain_names=None): if len(scalar_axes) == 1 and len(brain_axes) == 1: # dscalar.nii: ScalarAxis (axis 0) + BrainModelAxis (axis 1) brain_axis = brain_axes[0] - cifti_data = cifti.get_fdata().squeeze().astype(np.float32) - if cifti_data.ndim != 1: - raise ValueError(f'Expected 1-D scalar CIFTI data in {cifti_file!r}.') - brain_names = brain_axis.name + cifti_data = cifti.get_fdata( + dtype=np.dtype(storage_dtype).type, caching='unchanged' + ).reshape(-1) + brain_names = ( + reference_brain_names + if reference_signature is not None and reference_brain_names is not None + else brain_axis.name + ) if cifti_data.shape[0] != brain_names.shape[0]: raise ValueError( f'Mismatch between brain names and data array in {cifti_file!r}: ' f'{brain_names.shape[0]} names vs {cifti_data.shape[0]} values.' ) - if reference_brain_names is not None: - if not np.array_equal(brain_names, reference_brain_names): - raise ValueError(f'Inconsistent greyordinate names in CIFTI file {cifti_file!r}.') + if ( + reference_brain_names is not None + and reference_signature is None + and not np.array_equal(brain_names, reference_brain_names) + ): + raise ValueError(f'Inconsistent greyordinate names in CIFTI file {cifti_file!r}.') return cifti_data, brain_names elif len(scalar_axes) == 1 and len(parcel_axes) == 1: # pscalar.nii: ScalarAxis (axis 0) + ParcelsAxis (axis 1) parcel_axis = parcel_axes[0] - cifti_data = cifti.get_fdata().squeeze().astype(np.float32) - if cifti_data.ndim != 1: - raise ValueError(f'Expected 1-D parcellated scalar CIFTI data in {cifti_file!r}.') - parcel_names = parcel_axis.name + cifti_data = cifti.get_fdata( + dtype=np.dtype(storage_dtype).type, caching='unchanged' + ).reshape(-1) + parcel_names = ( + reference_brain_names + if reference_signature is not None and reference_brain_names is not None + else parcel_axis.name + ) if cifti_data.shape[0] != parcel_names.shape[0]: raise ValueError( f'Mismatch between parcel names and data array in {cifti_file!r}: ' f'{parcel_names.shape[0]} names vs {cifti_data.shape[0]} values.' ) - if reference_brain_names is not None: - if not np.array_equal(parcel_names, reference_brain_names): - raise ValueError(f'Inconsistent parcel names in CIFTI file {cifti_file!r}.') + if ( + reference_brain_names is not None + and reference_signature is None + and not np.array_equal(parcel_names, reference_brain_names) + ): + raise ValueError(f'Inconsistent parcel names in CIFTI file {cifti_file!r}.') return cifti_data, parcel_names elif len(parcel_axes) == 2: # pconn.nii: ParcelsAxis (axis 0, rows) + ParcelsAxis (axis 1, cols) row_axis = parcel_axes[0] col_axis = parcel_axes[1] - cifti_data = cifti.get_fdata().astype(np.float32) + cifti_data = cifti.get_fdata(dtype=np.dtype(storage_dtype).type, caching='unchanged') if cifti_data.ndim != 2: raise ValueError( f'Expected 2-D parcellated connectivity CIFTI data in {cifti_file!r}.' ) n_rows, n_cols = cifti_data.shape - row_parcel_names = row_axis.name - col_parcel_names = col_axis.name - # Validate that parcel axes lengths match the data shape - if len(row_parcel_names) != n_rows: + if len(row_axis) != n_rows: raise ValueError( f'Mismatch between row parcel names and data array in {cifti_file!r}: ' - f'{len(row_parcel_names)} names vs {n_rows} rows.' + f'{len(row_axis)} names vs {n_rows} rows.' ) - if len(col_parcel_names) != n_cols: + if len(col_axis) != n_cols: raise ValueError( f'Mismatch between column parcel names and data array in {cifti_file!r}: ' - f'{len(col_parcel_names)} names vs {n_cols} columns.' + f'{len(col_axis)} names vs {n_cols} columns.' ) - # One parcel name per flattened entry: repeat each row parcel name n_cols times - parcel_names = np.repeat(row_parcel_names, n_cols) + if reference_signature is not None and reference_brain_names is not None: + # The compact signature already covers both parcel axes, so avoid + # rebuilding or comparing their expanded labels on every source. + parcel_names = reference_brain_names + else: + # One parcel name per flattened entry: repeat each row parcel name + # once per column for the reference metadata. + parcel_names = np.repeat(row_axis.name, n_cols) cifti_data_flat = cifti_data.ravel() if reference_brain_names is not None: if reference_brain_names.ndim != 1: @@ -127,7 +252,9 @@ def extract_cifti_scalar_data(cifti_file, reference_brain_names=None): f'{reference_brain_names.shape[0]} names vs ' f'{cifti_data_flat.shape[0]} flattened connectivity values.' ) - if not np.array_equal(parcel_names, reference_brain_names): + if reference_signature is None and not np.array_equal( + parcel_names, reference_brain_names + ): raise ValueError(f'Inconsistent parcel names in CIFTI file {cifti_file!r}.') return cifti_data_flat, parcel_names @@ -157,9 +284,13 @@ def _get_cifti_parcel_info(cifti_file): For ``'pconn'``: ``{'parcel_id_from': row_names, 'parcel_id_to': col_names}`` — one entry per unique row/column parcel axis. """ - cifti = ( - cifti_file if hasattr(cifti_file, 'get_fdata') else load_nibabel(cifti_file, cifti=True) - ) + if not hasattr(cifti_file, 'get_fdata'): + cifti = load_nibabel(cifti_file, cifti=True) + try: + return _get_cifti_parcel_info(cifti) + finally: + release_image_buffers(cifti) + cifti = cifti_file axes = [cifti.header.get_axis(i) for i in range(cifti.ndim)] scalar_axes = [ax for ax in axes if isinstance(ax, nb.cifti2.cifti2_axes.ScalarAxis)] @@ -209,71 +340,396 @@ def brain_names_to_dataframe(brain_names): return greyordinate_df, structure_name_strings -def load_cohort_cifti(cohort_long, s3_workers): - """Load all CIFTI scalar rows from the cohort, optionally in parallel. +@dataclass(frozen=True) +class _CiftiRawReadPlan: + """On-disk layout for reading a CIFTI scalar row without parsing its XML. - The first file is always loaded serially to obtain the reference brain - structure axis used for validation. When s3_workers > 1, remaining rows - are submitted to a ThreadPoolExecutor and collected via as_completed. - Threads share memory so reference_brain_names is accessed directly with - no copying overhead. + Every source in a cohort produced by one pipeline on one surface template + stores its values as a contiguous floating block at a fixed ``vox_offset``. + Recovering that block through nibabel means parsing the CIFTI XML header and + rebuilding the brain-model axis, which measures ~1 s per file; reading the + block directly measures well under a millisecond. A plan captures one + reference file's layout so later sources with an identical cheap header can + be read straight from disk. + """ + + offset: int + disk_dtype: np.dtype + n_elements: int + slope: float + inter: float + expected_file_size: int + storage_dtype: np.dtype + # Cheap header facts that must match the reference before a data block is + # trusted: the NIfTI-2 dims (so the block holds exactly n_elements) and a + # digest of the raw XML bytes of the brain-model MatrixIndicesMap (so the + # greyordinate order is byte-identical). vox_offset may differ between + # files — production pipelines embed per-subject provenance strings in the + # CIFTI metadata, which shifts the data block without changing its layout. + dims: tuple[int, ...] = () + brain_model_digest: bytes = b'' + + +_BRAIN_MODEL_DIM_MARKER = b'AppliesToMatrixDimension="1"' + + +def _brain_model_map_digest_bytes(head: bytes) -> bytes | None: + """Hash the XML bytes of the dimension-1 (brain-model) MatrixIndicesMap in *head*. + + *head* is the file prefix ``[0, vox_offset)``. The XML is never parsed, so + this costs a few milliseconds where a nibabel decode costs about a second. + Returns None when the map cannot be located, which disables the fast path + for that file. + """ + xml_start = head.find(b'', marker) + if start < 0 or end < 0: + return None + end += len(b'') + return hashlib.sha256(head[start:end]).digest() + + +def _brain_model_map_digest(path: str, data_offset: int) -> bytes | None: + """Hash the brain-model map of the file at *path*, reading only ``[0, vox_offset)``.""" + try: + with open(path, 'rb') as handle: + head = handle.read(int(data_offset)) + except (OSError, ValueError, OverflowError): + return None + return _brain_model_map_digest_bytes(head) + + +def _raw_read_enabled() -> bool: + """Return True unless the raw greyordinate fast path is disabled by env.""" + return os.environ.get(_RAW_READ_ENV, '').strip().lower() not in {'0', 'false', 'no', 'off'} + + +def _parse_nifti2_header(header_bytes: bytes): + """Parse a 540-byte NIfTI-2 header, or return None when it cannot be trusted.""" + if len(header_bytes) != _NIFTI2_HEADER_SIZE: + return None + try: + return nb.Nifti2Header.from_fileobj(BytesIO(header_bytes), check=False) + except Exception: # noqa: BLE001 - any parse failure must fall back to nibabel + return None + + +def _read_nifti2_header(path: str): + """Parse only the 540-byte NIfTI-2 header, skipping the CIFTI XML extension.""" + try: + with open(path, 'rb') as handle: + header_bytes = handle.read(_NIFTI2_HEADER_SIZE) + except OSError: + return None + return _parse_nifti2_header(header_bytes) + + +def _normalize_slope_inter(header) -> tuple[float, float]: + """Return (slope, inter) as plain floats, treating 'no scaling' as (1, 0).""" + try: + slope, inter = header.get_slope_inter() + except Exception: # noqa: BLE001 - unreadable scaling must fall back to nibabel + return 1.0, 0.0 + slope = 1.0 if slope is None else float(slope) + inter = 0.0 if inter is None else float(inter) + return slope, inter + + +def _read_cifti_row_raw(source, plan: _CiftiRawReadPlan) -> np.ndarray | None: + """Read one scalar row straight from its data block, or None to fall back. + + Returns None unless *source* is a local, uncompressed file whose NIfTI-2 + header matches *plan* in dims, dtype and scaling, whose brain-model + MatrixIndicesMap is byte-identical to the reference's, and whose file is + large enough to hold the data block at its own ``vox_offset``. The offset + itself may differ from the reference (per-subject provenance metadata + lengthens the XML); the digest is what guarantees the same greyordinate + order. Any mismatch means the caller decodes the source with nibabel. + + The file is opened once and read sequentially: header, XML extension and + data block in a single read, which matters on network filesystems where + every open is a metadata round trip. + """ + if is_s3_path(source): + return None + path = os.fspath(source) + if path.endswith('.gz'): + return None + data_bytes = plan.n_elements * plan.disk_dtype.itemsize + try: + with open(path, 'rb') as handle: + header_bytes = handle.read(_NIFTI2_HEADER_SIZE) + header = _parse_nifti2_header(header_bytes) + if header is None: + return None + try: + offset = int(header.get_data_offset()) + if tuple(int(v) for v in header.get_data_shape()) != plan.dims: + return None + if np.dtype(header.get_data_dtype()) != plan.disk_dtype: + return None + except Exception: # noqa: BLE001 - any header mismatch must fall back to nibabel + return None + if _normalize_slope_inter(header) != (plan.slope, plan.inter): + return None + if offset < _NIFTI2_HEADER_SIZE: + return None + xml_length = offset - _NIFTI2_HEADER_SIZE + tail = handle.read(xml_length + data_bytes) + except OSError: + return None + if len(tail) < xml_length + data_bytes: + # Truncated file: the data block does not fit at its own offset. + return None + if _brain_model_map_digest_bytes(header_bytes + tail[:xml_length]) != plan.brain_model_digest: + return None + values = np.frombuffer(tail, dtype=plan.disk_dtype, count=plan.n_elements, offset=xml_length) + if plan.slope != 1.0 or plan.inter != 0.0: + values = values.astype(np.float64) * plan.slope + plan.inter + return np.array(values, dtype=plan.storage_dtype, order='C') + + +def _build_cifti_raw_read_plan( + reference_image, reference_row: np.ndarray, source, storage_dtype +) -> _CiftiRawReadPlan | None: + """Capture a raw-read plan from a reference source, or None if ineligible. + + A plan is returned only when reading the reference file's data block + reproduces *reference_row* exactly, so any layout the fast path cannot + handle — parcellated-connectivity ordering, on-disk scaling, unusual byte + order, or remote/compressed inputs — disables it rather than risking a + silently wrong row. + """ + if is_s3_path(source): + return None + path = os.fspath(source) + if path.endswith('.gz'): + return None + nifti_header = getattr(reference_image, 'nifti_header', None) + dataobj = getattr(reference_image, 'dataobj', None) + if nifti_header is None or dataobj is None: + return None + try: + disk_dtype = np.dtype(nifti_header.get_data_dtype()) + offset = int(dataobj.offset) + except (AttributeError, TypeError, ValueError): + return None + if disk_dtype.kind != 'f': + return None + reference_row = np.asarray(reference_row) + if reference_row.ndim != 1: + return None + n_elements = int(reference_row.shape[0]) + slope, inter = _normalize_slope_inter(nifti_header) + try: + dims = tuple(int(v) for v in nifti_header.get_data_shape()) + except Exception: # noqa: BLE001 - unreadable dims must disable the fast path + return None + digest = _brain_model_map_digest(path, offset) + if digest is None: + return None + plan = _CiftiRawReadPlan( + offset=offset, + disk_dtype=disk_dtype, + n_elements=n_elements, + slope=slope, + inter=inter, + expected_file_size=offset + n_elements * disk_dtype.itemsize, + storage_dtype=np.dtype(storage_dtype), + dims=dims, + brain_model_digest=digest, + ) + probe = _read_cifti_row_raw(path, plan) + if probe is None: + return None + expected = reference_row.astype(plan.storage_dtype, copy=False) + if not np.array_equal(probe, expected, equal_nan=True): + logger.debug('Raw CIFTI fast path disabled for %s: reference probe mismatch.', path) + return None + return plan + + +def load_cohort_cifti( + cohort_long, + s3_workers=None, + *, + workers=None, + storage_dtype='float32', +): + """Return lazy per-scalar CIFTI matrices and reference spatial metadata. + + Each scalar's first source is opened and decoded with nibabel once, to fix + its reference row and brain-model axis. Later sources are read straight from + their data block when their cheap NIfTI-2 header matches that reference + exactly (see :func:`_read_cifti_row_raw`), which avoids re-parsing the CIFTI + XML per file; any source that does not match falls back to a full nibabel + decode validated against a compact spatial signature. Set + ``MODELARRAYIO_CIFTI_RAW=0`` to force the nibabel path for every source. + + Rows are derived with :func:`~modelarrayio.utils.misc.build_scalar_sources`, + the same helper the CLI uses for ``column_names``, so the staged rows and + the stored source names always agree. Parameters ---------- cohort_long : :obj:`pandas.DataFrame` Long-format cohort dataframe s3_workers : :obj:`int` - Number of workers to use for parallel loading + Deprecated positional alias for the loader worker count. + workers : :obj:`int` + Maximum number of concurrent source loads. + storage_dtype : {'float32', 'float64'} + Dtype used for decoding and staging. Returns ------- - scalars : :obj:`dict` - Per-scalar ordered list of 1-D subject arrays, ready for stripe-write. + scalars : :class:`modelarrayio.utils.loading.LazyScalarMatrices` + Lazy mapping that stages one ordered, typed matrix per scalar. Its + ``reference_image``, ``reference_brain_names`` and + ``reference_signature`` attributes carry the decoded first source. reference_brain_names : :obj:`numpy.ndarray` Brain structure names from the first file, for building greyordinate table. """ - # Assign stable per-scalar subject indices in cohort order - scalar_subj_counter = defaultdict(int) - rows_with_idx = [] - for row in cohort_long.itertuples(index=False): - subj_idx = scalar_subj_counter[row.scalar_name] - scalar_subj_counter[row.scalar_name] += 1 - rows_with_idx.append((row.scalar_name, subj_idx, row.source_file)) - - # Load the first file serially to get the reference brain axis - first_sn, _, first_src = rows_with_idx[0] - first_data, reference_brain_names = extract_cifti_scalar_data( - load_nibabel(first_src, cifti=True) + worker_count = resolve_loader_workers(workers, s3_workers) + + sources_by_scalar = build_scalar_sources(cohort_long) + if not sources_by_scalar: + raise ValueError('Cannot load CIFTI cohort without any sources.') + + first_scalar = next(iter(sources_by_scalar)) + first_src = sources_by_scalar[first_scalar][0] + first_started = time.perf_counter() + first_cpu_started = time.process_time() + first_image = None + try: + first_image = load_nibabel(first_src, cifti=True) + first_data, reference_brain_names = extract_cifti_scalar_data( + first_image, + storage_dtype=storage_dtype, + ) + validation_started = time.perf_counter() + reference_signature = cifti_spatial_signature(first_image) + first_validation_seconds = time.perf_counter() - validation_started + except Exception as exc: + raise SourceLoadError( + f'Failed to load source {str(first_src)!r} for scalar {first_scalar!r}: {exc}' + ) from exc + finally: + if first_image is not None: + release_image_buffers(first_image) + first_loaded_row = LoadedRow( + first_data, + validation_seconds=first_validation_seconds, + source_seconds=time.perf_counter() - first_started, + cpu_seconds=time.process_time() - first_cpu_started, ) - def _worker(job): - sn, subj_idx, src = job - arr, _ = extract_cifti_scalar_data( - load_nibabel(src, cifti=True), reference_brain_names=reference_brain_names + raw_enabled = _raw_read_enabled() + first_plan = ( + _build_cifti_raw_read_plan(first_image, first_data, first_src, storage_dtype) + if raw_enabled + else None + ) + + def _validated_signature(image, source): + """Return the image's spatial signature, rejecting any axis mismatch.""" + validation_started = time.perf_counter() + signature = cifti_spatial_signature(image) + if signature != reference_signature: + raise ValueError(f'Inconsistent spatial axis in CIFTI file {source!r}.') + return signature, time.perf_counter() - validation_started + + def _decode_reference_row(source): + """Full nibabel decode of a scalar's first source, validated against the cohort.""" + started = time.perf_counter() + cpu_started = time.process_time() + image = load_nibabel(source, cifti=True) + try: + signature, validation_seconds = _validated_signature(image, source) + values, _ = extract_cifti_scalar_data( + image, + reference_brain_names=reference_brain_names, + storage_dtype=storage_dtype, + reference_signature=reference_signature, + spatial_signature=signature, + ) + except Exception: + release_image_buffers(image) + raise + row = LoadedRow( + values, + validation_seconds=validation_seconds, + source_seconds=time.perf_counter() - started, + cpu_seconds=time.process_time() - cpu_started, + ) + return row, image + + def _load(source, plan=None): + if plan is not None: + raw_values = _read_cifti_row_raw(source, plan) + if raw_values is not None: + return LoadedRow(raw_values, fast_path=True) + image = load_nibabel(source, cifti=True) + try: + signature, validation_seconds = _validated_signature(image, source) + values, _ = extract_cifti_scalar_data( + image, + reference_brain_names=reference_brain_names, + storage_dtype=storage_dtype, + reference_signature=reference_signature, + spatial_signature=signature, + ) + finally: + release_image_buffers(image) + return LoadedRow(values, validation_seconds=validation_seconds) + + def _stage(scalar_name): + sources = sources_by_scalar[scalar_name] + if scalar_name == first_scalar: + reusable_first = first_loaded_row + plan = first_plan + else: + try: + reusable_first, reference_image = _decode_reference_row(sources[0]) + except Exception as exc: + raise SourceLoadError( + f'Failed to load source {str(sources[0])!r} for scalar {scalar_name!r}: {exc}' + ) from exc + try: + plan = ( + _build_cifti_raw_read_plan( + reference_image, reusable_first.values, sources[0], storage_dtype + ) + if raw_enabled + else None + ) + finally: + release_image_buffers(reference_image) + + def _load_source(source): + return _load(source, plan) + + return stage_source_rows( + sources, + _load_source, + scalar_name=scalar_name, + storage_dtype=storage_dtype, + workers=worker_count, + first_row=reusable_first, ) - return sn, subj_idx, arr - - if s3_workers > 1 and len(rows_with_idx) > 1: - results = {first_sn: {0: first_data}} - jobs = [(sn, subj_idx, src) for sn, subj_idx, src in rows_with_idx[1:]] - with ThreadPoolExecutor(max_workers=s3_workers) as pool: - futures = {pool.submit(_worker, job): job for job in jobs} - for future in tqdm( - as_completed(futures), - total=len(futures), - desc='Loading CIFTI data', - ): - sn, subj_idx, arr = future.result() - results.setdefault(sn, {})[subj_idx] = arr - scalars = { - sn: [results[sn][i] for i in range(cnt)] for sn, cnt in scalar_subj_counter.items() - } - else: - scalars = defaultdict(list) - scalars[first_sn].append(first_data) - remaining = [(sn, subj_idx, src) for sn, subj_idx, src in rows_with_idx[1:]] - for job in tqdm(remaining, desc='Loading CIFTI data'): - sn, subj_idx, arr = _worker(job) - scalars[sn].append(arr) + scalars = LazyScalarMatrices( + list(sources_by_scalar), + _stage, + sources_by_scalar=sources_by_scalar, + reference_image=first_image, + reference_brain_names=reference_brain_names, + reference_signature=reference_signature, + ) return scalars, reference_brain_names diff --git a/src/modelarrayio/utils/loading.py b/src/modelarrayio/utils/loading.py new file mode 100644 index 0000000..4210f11 --- /dev/null +++ b/src/modelarrayio/utils/loading.py @@ -0,0 +1,402 @@ +"""Bounded, ordered source loading into typed staging matrices.""" + +from __future__ import annotations + +import logging +import operator +import time +import warnings +from collections.abc import Callable, Iterator, Mapping, Sequence +from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait +from dataclasses import asdict, dataclass +from typing import Any, Generic, TypeVar + +import numpy as np + +logger = logging.getLogger(__name__) + +SourceT = TypeVar('SourceT') + +# Seconds between progress lines while a scalar is being staged. +DEFAULT_PROGRESS_INTERVAL_SECONDS = 30.0 + + +class SourceLoadError(RuntimeError): + """A source could not be loaded into its staging-matrix row.""" + + +@dataclass(frozen=True) +class LoadedRow: + """A decoded row plus optional validation time recorded by its reader.""" + + values: np.ndarray + validation_seconds: float = 0.0 + source_seconds: float = 0.0 + cpu_seconds: float = 0.0 + # True when a reader took a fast path (for example the CIFTI raw greyordinate + # read) instead of a full decode, so silent fallbacks are visible in telemetry. + fast_path: bool = False + + +@dataclass +class LoadStats: + """Cheap aggregate counters for one scalar's bounded load.""" + + allocation_seconds: float = 0.0 + wall_seconds: float = 0.0 + source_wall_seconds: float = 0.0 + cpu_seconds: float = 0.0 + source_seconds: float = 0.0 + validation_seconds: float = 0.0 + scheduler_wait_seconds: float = 0.0 + in_flight_sources_mean: float = 0.0 + in_flight_sources_max: int = 0 + # Reserved for a future loader/writer overlap queue. The bounded source + # scheduler has in-flight tasks, but no producer/consumer queue. + queue_occupancy_mean: float = 0.0 + queue_occupancy_max: int = 0 + # Rows whose reader reported a fast path (see LoadedRow.fast_path). The + # first row of a scalar is its reference decode, so a fully served scalar + # reports ``source_count - 1`` here. + fast_path_rows: int = 0 + source_count: int = 0 + workers: int = 1 + + def to_dict(self) -> dict[str, float | int]: + """Return JSON-serializable counters.""" + return asdict(self) + + +def validate_workers(workers: int | None) -> int: + """Return a positive loader-worker count.""" + if workers is None: + return 1 + try: + if isinstance(workers, bool): + raise TypeError('a boolean is not a worker count') # bool passes operator.index + resolved = operator.index(workers) + except TypeError as exc: + raise ValueError(f'workers must be a positive integer, got {workers!r}') from exc + if resolved < 1: + raise ValueError(f'workers must be a positive integer, got {workers!r}') + return resolved + + +def resolve_loader_workers(workers: int | None, s3_workers: int | None) -> int: + """Resolve ``workers`` and its deprecated ``s3_workers`` alias to one count. + + This is the single place the alias is interpreted, for the CLI + (``--workers`` / ``--s3-workers``) and for the modality loaders alike. + """ + if workers is None and s3_workers is None: + return 1 + + resolved_workers = validate_workers(workers) if workers is not None else None + resolved_s3 = validate_workers(s3_workers) if s3_workers is not None else None + if resolved_s3 is not None: + message = ( + '--s3-workers is deprecated and will be removed in the next breaking release; ' + 'use --workers instead. --workers now controls source loading for local and S3 input.' + ) + warnings.warn(message, DeprecationWarning, stacklevel=3) + logger.warning(message) + if ( + resolved_workers is not None + and resolved_s3 is not None + and resolved_workers != resolved_s3 + ): + raise ValueError( + f'Conflicting loader worker counts: --workers={resolved_workers} and ' + f'--s3-workers={resolved_s3}. Use one option or pass matching values.' + ) + return resolved_workers if resolved_workers is not None else int(resolved_s3) + + +def _source_label(source: Any) -> str: + if isinstance(source, Mapping): + for key in ('source_file', 'path', 'uri'): + if key in source: + return str(source[key]) + if isinstance(source, tuple) and source: + return str(source[0]) + return str(source) + + +def _unwrap_loaded_row( + value: np.ndarray | LoadedRow, +) -> tuple[np.ndarray, float, float, float, bool]: + if isinstance(value, LoadedRow): + return ( + np.asarray(value.values), + float(value.validation_seconds), + float(value.source_seconds), + float(value.cpu_seconds), + bool(value.fast_path), + ) + return np.asarray(value), 0.0, 0.0, 0.0, False + + +def stage_source_rows( + sources: Sequence[SourceT], + load_row: Callable[[SourceT], np.ndarray | LoadedRow], + *, + scalar_name: str, + storage_dtype: str | np.dtype | type[np.floating] = np.float32, + workers: int = 1, + n_elements: int | None = None, + first_row: np.ndarray | LoadedRow | None = None, + progress_interval_seconds: float | None = DEFAULT_PROGRESS_INTERVAL_SECONDS, +) -> tuple[np.ndarray, LoadStats]: + """Load ordered sources into one C-contiguous staging matrix. + + Only ``workers`` calls are submitted or running at once. Each task copies + directly into its assigned row, so completion order never changes cohort + order and no second list/stack copy is built. While the load runs, an INFO + progress line is logged at most every ``progress_interval_seconds`` (``None`` + or a non-positive value disables it), so a long scalar is never silent. + """ + worker_count = validate_workers(workers) + source_count = len(sources) + if source_count == 0: + raise ValueError(f'Cannot stage scalar {scalar_name!r}: no sources were provided.') + + dtype = np.dtype(storage_dtype) + stats = LoadStats(source_count=source_count, workers=worker_count) + started = time.perf_counter() + started_cpu = time.process_time() + + loaded_first: np.ndarray | None = None + start_index = 0 + preloaded_source_seconds = 0.0 + preloaded_cpu_seconds = 0.0 + in_flight_total = 0 + in_flight_samples = 0 + in_flight_max = 0 + + def _sample_in_flight(count: int) -> None: + nonlocal in_flight_total, in_flight_samples, in_flight_max + if count <= 0: + return + in_flight_total += count + in_flight_samples += 1 + in_flight_max = max(in_flight_max, count) + + if first_row is not None: + ( + loaded_first, + validation_seconds, + preloaded_source_seconds, + preloaded_cpu_seconds, + first_fast_path, + ) = _unwrap_loaded_row(first_row) + stats.validation_seconds += validation_seconds + stats.fast_path_rows += int(first_fast_path) + stats.source_seconds += preloaded_source_seconds + stats.source_wall_seconds += preloaded_source_seconds + if preloaded_source_seconds: + _sample_in_flight(1) + start_index = 1 + elif n_elements is None: + source_started = time.perf_counter() + try: + loaded_first, validation_seconds, _, _, first_fast_path = _unwrap_loaded_row( + load_row(sources[0]) + ) + except Exception as exc: + raise SourceLoadError( + f'Failed to load source {_source_label(sources[0])!r} ' + f'for scalar {scalar_name!r}: {exc}' + ) from exc + first_source_seconds = time.perf_counter() - source_started + stats.source_seconds += first_source_seconds + stats.source_wall_seconds += first_source_seconds + stats.validation_seconds += validation_seconds + stats.fast_path_rows += int(first_fast_path) + _sample_in_flight(1) + start_index = 1 + + if loaded_first is not None: + if loaded_first.ndim != 1: + raise SourceLoadError( + f'Loaded source {_source_label(sources[0])!r} for scalar {scalar_name!r} ' + f'produced shape {loaded_first.shape}; expected one dimension.' + ) + inferred_elements = int(loaded_first.shape[0]) + if n_elements is None: + n_elements = inferred_elements + elif inferred_elements != int(n_elements): + raise SourceLoadError( + f'Loaded source {_source_label(sources[0])!r} for scalar {scalar_name!r} ' + f'has {inferred_elements} elements; expected {n_elements}.' + ) + + if n_elements is None or int(n_elements) <= 0: + raise ValueError( + f'Cannot stage scalar {scalar_name!r}: n_elements must be positive, got {n_elements!r}.' + ) + n_elements = int(n_elements) + + allocation_started = time.perf_counter() + matrix = np.empty((source_count, n_elements), dtype=dtype, order='C') + stats.allocation_seconds = time.perf_counter() - allocation_started + if loaded_first is not None: + np.copyto(matrix[0], loaded_first, casting='unsafe') + del loaded_first + + def _load_into(index: int, source: SourceT) -> tuple[float, float, bool]: + row_started = time.perf_counter() + try: + values, validation_seconds, _, _, fast_path = _unwrap_loaded_row(load_row(source)) + if values.ndim != 1 or values.shape[0] != n_elements: + raise ValueError(f'loaded row has shape {values.shape}; expected ({n_elements},)') + np.copyto(matrix[index], values, casting='unsafe') + except Exception as exc: + raise SourceLoadError( + f'Failed to load source {_source_label(source)!r} ' + f'for scalar {scalar_name!r}: {exc}' + ) from exc + return time.perf_counter() - row_started, validation_seconds, fast_path + + progress_interval = ( + float(progress_interval_seconds) if progress_interval_seconds is not None else 0.0 + ) + progress_started = time.perf_counter() + last_progress = progress_started + completed = start_index + + def _record_completion(in_flight: int) -> None: + nonlocal completed, last_progress + completed += 1 + if progress_interval <= 0 or completed >= source_count: + return + now = time.perf_counter() + if now - last_progress < progress_interval: + return + last_progress = now + elapsed = now - progress_started + logger.info( + 'Staging %s: %d/%d sources loaded (%.1f sources/s, %d in flight)', + scalar_name, + completed, + source_count, + completed / elapsed if elapsed > 0 else 0.0, + in_flight, + ) + + remaining = range(start_index, source_count) + remaining_started = time.perf_counter() + loaded_remaining = False + if worker_count == 1: + for index in remaining: + loaded_remaining = True + _sample_in_flight(1) + elapsed, validation_seconds, fast_path = _load_into(index, sources[index]) + stats.source_seconds += elapsed + stats.validation_seconds += validation_seconds + stats.fast_path_rows += int(fast_path) + _record_completion(0) + else: + source_iter: Iterator[int] = iter(remaining) + futures: dict[Future[tuple[float, float, bool]], int] = {} + with ThreadPoolExecutor( + max_workers=worker_count, + thread_name_prefix='modelarray-loader', + ) as executor: + for _ in range(worker_count): + try: + index = next(source_iter) + except StopIteration: + break + futures[executor.submit(_load_into, index, sources[index])] = index + + try: + while futures: + loaded_remaining = True + _sample_in_flight(len(futures)) + wait_started = time.perf_counter() + done, _ = wait(futures, return_when=FIRST_COMPLETED) + stats.scheduler_wait_seconds += time.perf_counter() - wait_started + for future in done: + futures.pop(future) + elapsed, validation_seconds, fast_path = future.result() + stats.source_seconds += elapsed + stats.validation_seconds += validation_seconds + stats.fast_path_rows += int(fast_path) + try: + index = next(source_iter) + except StopIteration: + pass + else: + futures[executor.submit(_load_into, index, sources[index])] = index + _record_completion(len(futures)) + except Exception: + for future in futures: + future.cancel() + raise + + if loaded_remaining: + stats.source_wall_seconds += time.perf_counter() - remaining_started + if in_flight_samples: + stats.in_flight_sources_mean = in_flight_total / in_flight_samples + stats.in_flight_sources_max = in_flight_max + stats.wall_seconds = time.perf_counter() - started + if preloaded_source_seconds: + # The preloaded row was decoded before this function could infer and + # allocate the matrix. Include it in end-to-end staging telemetry. + stats.wall_seconds += preloaded_source_seconds + stats.cpu_seconds = time.process_time() - started_cpu + preloaded_cpu_seconds + return matrix, stats + + +class LazyScalarMatrices(Mapping[str, np.ndarray], Generic[SourceT]): + """Mapping that stages one scalar matrix at a time on access. + + Every access re-runs the staging function: there is no cache, so index each + scalar once, write it, and drop it. Loader-specific setup products travel + with the mapping as explicit attributes instead of being attached ad hoc: + + ``stats`` + Per-scalar :class:`LoadStats`, filled in as scalars are accessed. + ``sources_by_scalar`` + Ordered source paths per scalar, exactly as the rows were staged. + ``reference_image``, ``reference_brain_names``, ``reference_signature`` + CIFTI only: the decoded first image, its element names and spatial + signature, so callers can build metadata without reopening the file. + ``flat_indices``, ``mask_cache`` + NIfTI only: the group-mask flat selection and the reduced-mask cache. + """ + + def __init__( + self, + scalar_names: Sequence[str], + stage_scalar: Callable[[str], tuple[np.ndarray, LoadStats]], + *, + sources_by_scalar: Mapping[str, Sequence[Any]] | None = None, + reference_image: Any = None, + reference_brain_names: np.ndarray | None = None, + reference_signature: tuple[int, str] | None = None, + flat_indices: np.ndarray | None = None, + mask_cache: Any = None, + ) -> None: + self._scalar_names = tuple(scalar_names) + self._stage_scalar = stage_scalar + self.stats: dict[str, LoadStats] = {} + self.sources_by_scalar = sources_by_scalar + self.reference_image = reference_image + self.reference_brain_names = reference_brain_names + self.reference_signature = reference_signature + self.flat_indices = flat_indices + self.mask_cache = mask_cache + + def __getitem__(self, scalar_name: str) -> np.ndarray: + if scalar_name not in self._scalar_names: + raise KeyError(scalar_name) + matrix, stats = self._stage_scalar(scalar_name) + self.stats[scalar_name] = stats + return matrix + + def __iter__(self) -> Iterator[str]: + return iter(self._scalar_names) + + def __len__(self) -> int: + return len(self._scalar_names) diff --git a/src/modelarrayio/utils/mif.py b/src/modelarrayio/utils/mif.py index 5a93fc0..1ca411e 100644 --- a/src/modelarrayio/utils/mif.py +++ b/src/modelarrayio/utils/mif.py @@ -1,19 +1,40 @@ """Utility functions for MIF data.""" -from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor, as_completed +from collections import OrderedDict +from io import BytesIO from pathlib import Path import numpy as np import pandas as pd -from tqdm import tqdm +from nibabel.filebasedimages import FileHolder +from modelarrayio.utils.loading import ( + LazyScalarMatrices, + resolve_loader_workers, + stage_source_rows, +) from modelarrayio.utils.mif_image import MifHeader, MifImage +from modelarrayio.utils.s3_utils import ( + _check_local_source, + is_s3_path, + read_s3_bytes, + release_image_buffers, +) __all__ = ['MifHeader', 'MifImage', 'gather_fixels', 'load_cohort_mif', 'mif_to_image'] -def mif_to_image(mif_file): +def _load_mif_image(mif_file): + """Load a local, materialized-annex, or S3 MIF image.""" + if not is_s3_path(mif_file): + return MifImage.from_filename(str(_check_local_source(Path(mif_file)))) + + data = read_s3_bytes(mif_file, decompress_gzip=True) + file_map = {'image': FileHolder(fileobj=BytesIO(data))} + return MifImage.from_file_map(file_map) + + +def mif_to_image(mif_file, dtype=np.float32): """Load a `.mif` file into a :class:`MifImage`. Parameters @@ -28,72 +49,73 @@ def mif_to_image(mif_file): data : :obj:`numpy.ndarray` Data from the MIF image """ - input_path = Path(mif_file) - mif_img = MifImage.from_filename(str(input_path)) - data = mif_img.get_fdata(dtype=np.float32).squeeze() + mif_img = _load_mif_image(mif_file) + release_after_decode = is_s3_path(mif_file) + try: + if hasattr(mif_img, 'dataobj'): + data = np.asanyarray(mif_img.dataobj).squeeze() + if dtype is not None and data.dtype != np.dtype(dtype): + data = data.astype(dtype, copy=False) + else: + # Compatibility with nibabel-like test doubles and older readers. + data = mif_img.get_fdata(dtype=dtype).squeeze() + finally: + if release_after_decode: + release_image_buffers(mif_img) return mif_img, data -def load_cohort_mif(cohort_long, s3_workers): - """Load all MIF scalar rows from the cohort, optionally in parallel. +def load_cohort_mif( + cohort_long, + s3_workers=None, + *, + workers=None, + storage_dtype='float32', +): + """Return lazy per-scalar MIF staging matrices and ordered sources. - When s3_workers > 1, a ThreadPoolExecutor is used to run mrconvert - calls concurrently (subprocess calls release the GIL). Results arrive - via as_completed and are indexed by (scalar_name, subj_idx) so the - final ordered lists are reconstructed correctly regardless of completion - order. + Only the matrix requested by the writer is resident. Its rows are loaded + with the shared bounded thread scheduler and cast once as they enter the + requested staging dtype. Parameters ---------- cohort_long : :obj:`pandas.DataFrame` Long-format cohort dataframe with columns 'scalar_name' and 'source_file'. s3_workers : :obj:`int` - Number of parallel workers for loading. + Deprecated positional alias for the loader worker count. + workers : :obj:`int` + Maximum number of concurrent source loads. + storage_dtype : {'float32', 'float64'} + Dtype of the staging matrix. Returns ------- - scalars : dict[str, list[np.ndarray]] - Per-scalar ordered list of 1-D subject arrays, ready for stripe-write. + scalars : modelarrayio.utils.loading.LazyScalarMatrices + Lazy mapping that stages one ordered, typed matrix per scalar. sources_lists : dict[str, list[str]] Per-scalar ordered list of source file paths (for HDF5 metadata). """ - scalar_subj_counter = defaultdict(int) - jobs = [] - sources_lists = defaultdict(list) + worker_count = resolve_loader_workers(workers, s3_workers) + sources_lists: OrderedDict[str, list[str]] = OrderedDict() for row in cohort_long.itertuples(index=False): - sn = row.scalar_name - subj_idx = scalar_subj_counter[sn] - scalar_subj_counter[sn] += 1 - src = row.source_file - jobs.append((sn, subj_idx, src)) - sources_lists[sn].append(src) - - def _worker(job): - sn, subj_idx, src = job - _img, data = mif_to_image(src) - return sn, subj_idx, data - - if s3_workers > 1: - results = defaultdict(dict) - with ThreadPoolExecutor(max_workers=s3_workers) as pool: - futures = {pool.submit(_worker, job): job for job in jobs} - for future in tqdm( - as_completed(futures), - total=len(futures), - desc='Loading MIF data', - ): - sn, subj_idx, data = future.result() - results[sn][subj_idx] = data - scalars = { - sn: [results[sn][i] for i in range(cnt)] for sn, cnt in scalar_subj_counter.items() - } - else: - scalars = defaultdict(list) - for job in tqdm(jobs, desc='Loading MIF data'): - sn, subj_idx, data = _worker(job) - scalars[sn].append(data) - + sources_lists.setdefault(str(row.scalar_name), []).append(str(row.source_file)) + + def _load_native(source): + _image, data = mif_to_image(source, dtype=None) + return np.asarray(data).squeeze() + + def _stage(scalar_name): + return stage_source_rows( + sources_lists[scalar_name], + _load_native, + scalar_name=scalar_name, + storage_dtype=storage_dtype, + workers=worker_count, + ) + + scalars = LazyScalarMatrices(list(sources_lists), _stage, sources_by_scalar=sources_lists) return scalars, sources_lists diff --git a/src/modelarrayio/utils/nifti.py b/src/modelarrayio/utils/nifti.py index 9e18ff7..416643d 100644 --- a/src/modelarrayio/utils/nifti.py +++ b/src/modelarrayio/utils/nifti.py @@ -1,124 +1,321 @@ -"""Utility functions for NIfTI data.""" +"""Utility functions for loading NIfTI data into staging matrices.""" from __future__ import annotations -from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor, as_completed +import os +import threading +from collections import OrderedDict +from collections.abc import Callable, Hashable +from concurrent.futures import Future from pathlib import Path +from typing import Any -import nibabel as nb import numpy as np -from tqdm import tqdm -from modelarrayio.utils.s3_utils import load_nibabel +from modelarrayio.utils.loading import ( + LazyScalarMatrices, + resolve_loader_workers, + stage_source_rows, +) +from modelarrayio.utils.s3_utils import is_s3_path, load_nibabel, release_image_buffers +_DEFAULT_MASK_CACHE_SIZE = 128 -def load_cohort_voxels(cohort_long, group_mask_matrix, s3_workers): - """Load all voxel rows from the cohort, optionally in parallel. - When s3_workers > 1, a ThreadPoolExecutor is used. Threads share memory so - group_mask_matrix is accessed directly with no copying overhead. Results - arrive via as_completed and are indexed by (scalar_name, subj_idx) so the - final ordered lists are reconstructed correctly regardless of completion order. +def _storage_dtype(dtype: str | np.dtype | type[np.floating]) -> np.dtype: + """Return a nibabel-compatible floating-point dtype.""" + resolved = np.dtype(dtype) + if resolved not in (np.dtype(np.float32), np.dtype(np.float64)): + raise ValueError(f'storage_dtype must be float32 or float64, got {resolved.name!r}.') + return resolved + + +def _path_cache_key(path: str | os.PathLike[str]) -> Hashable: + """Return a stable key while preserving remote URI semantics.""" + if is_s3_path(path): + return str(path) + return Path(path).expanduser().resolve(strict=False) + + +class ReducedMaskCache: + """Bounded, thread-safe LRU cache of group-selected boolean masks. + + Concurrent requests for the same path share one in-flight load. Cached + arrays are read-only so callers cannot make one subject's mask affect a + later subject. + """ + + def __init__(self, max_size: int = _DEFAULT_MASK_CACHE_SIZE) -> None: + if isinstance(max_size, bool) or not isinstance(max_size, int) or max_size < 1: + raise ValueError(f'mask cache size must be a positive integer, got {max_size!r}.') + self.max_size = max_size + self._values: OrderedDict[Hashable, np.ndarray] = OrderedDict() + self._in_flight: dict[Hashable, Future[np.ndarray]] = {} + self._lock = threading.Lock() + + def get_or_load( + self, + path: str | os.PathLike[str], + loader: Callable[[], np.ndarray], + ) -> np.ndarray: + """Return the reduced mask for ``path``, loading it at most once concurrently.""" + key = _path_cache_key(path) + owns_load = False + with self._lock: + cached = self._values.get(key) + if cached is not None: + self._values.move_to_end(key) + return cached + + pending = self._in_flight.get(key) + if pending is None: + pending = Future() + self._in_flight[key] = pending + owns_load = True + + if not owns_load: + return pending.result() + + try: + reduced = np.array(loader(), dtype=bool, copy=True, order='C') + if reduced.ndim != 1: + raise ValueError( + f'Reduced mask for {path!r} has shape {reduced.shape}; expected one dimension.' + ) + reduced.setflags(write=False) + except BaseException as exc: + with self._lock: + self._in_flight.pop(key, None) + pending.set_exception(exc) + raise + + with self._lock: + self._values[key] = reduced + self._values.move_to_end(key) + while len(self._values) > self.max_size: + self._values.popitem(last=False) + self._in_flight.pop(key, None) + pending.set_result(reduced) + return reduced + + def __len__(self) -> int: + """Return the current number of cached masks.""" + with self._lock: + return len(self._values) + + +def _as_nifti_image(source: Any): + """Return an already-loaded image or load a local/S3 source.""" + return source if hasattr(source, 'get_fdata') else load_nibabel(source) + + +def _flat_group_indices(group_mask_matrix: np.ndarray) -> tuple[tuple[int, ...], np.ndarray]: + """Return the group-mask shape and its C-order selected flat indices.""" + group_mask = np.asarray(group_mask_matrix, dtype=bool) + flat_indices = np.flatnonzero(group_mask) + flat_indices.setflags(write=False) + return group_mask.shape, flat_indices + + +def _selected_scalar_values( + scalar_image: Any, + *, + group_shape: tuple[int, ...], + flat_indices: np.ndarray, + storage_dtype: np.dtype, +) -> np.ndarray: + """Decode one scalar directly in its storage dtype and select group voxels.""" + loaded_here = not hasattr(scalar_image, 'get_fdata') + image = _as_nifti_image(scalar_image) + try: + scalar_matrix = image.get_fdata(dtype=storage_dtype, caching='unchanged') + if scalar_matrix.shape != group_shape: + raise ValueError( + 'Scalar image and group mask must have matching shapes: ' + f'{scalar_matrix.shape} != {group_shape}' + ) + + # np.take with axis=None uses C-order flat indices and returns only the + # selected row. It never exposes the decoded volume for in-place mutation. + return np.take(scalar_matrix, flat_indices) + finally: + if loaded_here: + release_image_buffers(image) + + +def _selected_mask_values( + scalar_mask: Any, + *, + group_shape: tuple[int, ...], + flat_indices: np.ndarray, +) -> np.ndarray: + """Decode one mask without forcing float64 and reduce it to group voxels.""" + loaded_here = not hasattr(scalar_mask, 'get_fdata') + image = _as_nifti_image(scalar_mask) + try: + mask_matrix = np.asanyarray(image.dataobj) + if mask_matrix.shape != group_shape: + raise ValueError( + 'Scalar mask and group mask must have matching shapes: ' + f'{mask_matrix.shape} != {group_shape}' + ) + return np.take(mask_matrix, flat_indices) > 0 + finally: + if loaded_here: + release_image_buffers(image) + + +def flattened_image( + scalar_image, + scalar_mask, + group_mask_matrix, + *, + storage_dtype: str | np.dtype | type[np.floating] = np.float32, +): + """Flatten a scalar image to a typed 1-D group-mask selection. + + The decoded scalar volume is never modified. NaNs for an individual mask + are applied only to the selected row, which keeps repeated image objects + and repeated paths independent. Parameters ---------- - cohort_long : :obj:`pandas.DataFrame` - Long-format cohort dataframe with columns 'scalar_name', 'source_file', - and 'source_mask_file'. - group_mask_matrix : :obj:`numpy.ndarray` - Boolean group mask array. - s3_workers : :obj:`int` - Number of parallel workers for loading. + scalar_image : nibabel image or path-like + Scalar image. + scalar_mask : nibabel image or path-like + Individual scalar mask. + group_mask_matrix : numpy.ndarray + Boolean group mask matrix. + storage_dtype : {'float32', 'float64'} + Floating-point dtype to request directly from nibabel. Returns ------- - scalars : dict[str, list[np.ndarray]] - Per-scalar ordered list of 1-D subject arrays, ready for stripe-write. - sources_lists : dict[str, list[str]] - Per-scalar ordered list of source file paths (for HDF5 metadata). + numpy.ndarray + One-dimensional selected scalar row. """ - scalar_subj_counter = defaultdict(int) - jobs = [] - sources_lists = defaultdict(list) - - for row in cohort_long.itertuples(index=False): - sn = row.scalar_name - subj_idx = scalar_subj_counter[sn] - scalar_subj_counter[sn] += 1 - src = row.source_file - msk = row.source_mask_file - jobs.append((sn, subj_idx, src, msk)) - sources_lists[sn].append(src) - - def _worker(job): - sn, subj_idx, scalar_path, mask_path = job - scalar_img = load_nibabel(scalar_path) - mask_img = load_nibabel(mask_path) - arr = flattened_image(scalar_img, mask_img, group_mask_matrix) - return sn, subj_idx, arr - - if s3_workers > 1: - results = defaultdict(dict) - with ThreadPoolExecutor(max_workers=s3_workers) as pool: - futures = {pool.submit(_worker, job): job for job in jobs} - for future in tqdm( - as_completed(futures), - total=len(futures), - desc='Loading voxel data', - ): - sn, subj_idx, arr = future.result() - results[sn][subj_idx] = arr - scalars = { - sn: [results[sn][i] for i in range(cnt)] for sn, cnt in scalar_subj_counter.items() - } - else: - scalars = defaultdict(list) - for job in tqdm(jobs, desc='Loading voxel data'): - sn, subj_idx, arr = _worker(job) - scalars[sn].append(arr) - - return scalars, sources_lists - - -def flattened_image(scalar_image, scalar_mask, group_mask_matrix): - """Flatten a scalar image to a 1-D array. + dtype = _storage_dtype(storage_dtype) + group_shape, flat_indices = _flat_group_indices(group_mask_matrix) + row = _selected_scalar_values( + scalar_image, + group_shape=group_shape, + flat_indices=flat_indices, + storage_dtype=dtype, + ) + reduced_mask = _selected_mask_values( + scalar_mask, + group_shape=group_shape, + flat_indices=flat_indices, + ) + row[~reduced_mask] = np.nan + return row + + +def load_cohort_voxels( + cohort_long, + group_mask_matrix, + s3_workers=None, + *, + workers=None, + storage_dtype='float32', + group_mask_path=None, + mask_cache_size=None, +): + """Return lazily staged per-scalar voxel matrices and ordered source paths. + + The third positional argument remains the legacy ``s3_workers`` count. + New callers should pass ``workers``. Accessing one key in the returned + mapping allocates and fills only that scalar's typed C-contiguous staging + matrix, with no more than ``workers`` source loads submitted at once. Parameters ---------- - scalar_image : :obj:`nibabel.Nifti1Image` - Scalar image. - scalar_mask : :obj:`nibabel.Nifti1Image` - Scalar mask. - group_mask_matrix : :obj:`numpy.ndarray` - Group mask matrix. + cohort_long : pandas.DataFrame + Long-format cohort dataframe with ``scalar_name``, ``source_file``, + and ``source_mask_file`` columns. + group_mask_matrix : numpy.ndarray + Boolean group mask array. + s3_workers : int, optional + Deprecated positional name for the loader worker count. + workers : int, optional + Maximum concurrent source loads. + storage_dtype : {'float32', 'float64'} + Dtype used for both nibabel decoding and the staging matrix. + group_mask_path : path-like, optional + Source path of ``group_mask_matrix``. A matching individual-mask path + is not opened because it cannot exclude any selected group voxel. + mask_cache_size : int, optional + Maximum number of reduced individual masks retained. Defaults to 128. Returns ------- - :obj:`numpy.ndarray` - Flattened scalar image. + scalars : modelarrayio.utils.loading.LazyScalarMatrices + Mapping from scalar name to a lazily constructed 2-D numpy array. + sources_by_scalar : dict[str, list[str]] + Source paths in exact cohort order for each scalar. """ - scalar_mask_img = ( - scalar_mask if hasattr(scalar_mask, 'get_fdata') else nb.load(Path(scalar_mask)) - ) - scalar_mask_matrix = scalar_mask_img.get_fdata() > 0 + worker_count = resolve_loader_workers(workers, s3_workers) + dtype = _storage_dtype(storage_dtype) + group_shape, flat_indices = _flat_group_indices(group_mask_matrix) + if flat_indices.size == 0: + raise ValueError('Cannot load NIfTI cohort: the group mask selects no voxels.') - scalar_img = ( - scalar_image if hasattr(scalar_image, 'get_fdata') else nb.load(Path(scalar_image)) - ) - scalar_matrix = scalar_img.get_fdata() + cache_size = _DEFAULT_MASK_CACHE_SIZE if mask_cache_size is None else mask_cache_size + mask_cache = ReducedMaskCache(cache_size) + group_mask_key = _path_cache_key(group_mask_path) if group_mask_path is not None else None + + sources_by_scalar: dict[str, list[str]] = {} + source_pairs_by_scalar: dict[str, list[tuple[str, str]]] = {} + for cohort_row in cohort_long.itertuples(index=False): + scalar_name = cohort_row.scalar_name + source_file = cohort_row.source_file + source_mask_file = cohort_row.source_mask_file + sources_by_scalar.setdefault(scalar_name, []).append(source_file) + source_pairs_by_scalar.setdefault(scalar_name, []).append((source_file, source_mask_file)) + + def _load_row(source_pair: tuple[str, str]) -> np.ndarray: + source_file, source_mask_file = source_pair + row = _selected_scalar_values( + source_file, + group_shape=group_shape, + flat_indices=flat_indices, + storage_dtype=dtype, + ) + + mask_key = _path_cache_key(source_mask_file) + if group_mask_key is not None and mask_key == group_mask_key: + return row - if scalar_matrix.shape != scalar_mask_matrix.shape: - raise ValueError( - 'Scalar image and scalar mask must have matching shapes: ' - f'{scalar_matrix.shape} != {scalar_mask_matrix.shape}' + reduced_mask = mask_cache.get_or_load( + source_mask_file, + lambda: _selected_mask_values( + source_mask_file, + group_shape=group_shape, + flat_indices=flat_indices, + ), ) - if scalar_matrix.shape != group_mask_matrix.shape: - raise ValueError( - 'Scalar image and group mask must have matching shapes: ' - f'{scalar_matrix.shape} != {group_mask_matrix.shape}' + row[~reduced_mask] = np.nan + return row + + scalar_names = tuple(sources_by_scalar) + + def _stage_scalar(scalar_name: str): + return stage_source_rows( + source_pairs_by_scalar[scalar_name], + _load_row, + scalar_name=scalar_name, + storage_dtype=dtype, + workers=worker_count, + n_elements=int(flat_indices.size), ) - scalar_matrix[np.logical_not(scalar_mask_matrix)] = np.nan - # .shape = (#voxels,) # squeeze() is to remove the 2nd dimension which is not necessary - return scalar_matrix[group_mask_matrix].squeeze() + # The flat selection and the reduced-mask cache ride along explicitly, so + # callers and tests need neither a second group-mask scan nor the volumes. + scalars = LazyScalarMatrices( + scalar_names, + _stage_scalar, + sources_by_scalar=sources_by_scalar, + flat_indices=flat_indices, + mask_cache=mask_cache, + ) + return scalars, sources_by_scalar diff --git a/src/modelarrayio/utils/s3_utils.py b/src/modelarrayio/utils/s3_utils.py index b37d436..383ca7b 100644 --- a/src/modelarrayio/utils/s3_utils.py +++ b/src/modelarrayio/utils/s3_utils.py @@ -5,18 +5,48 @@ import gzip import logging import os +import threading from io import BytesIO from os import PathLike +from pathlib import Path from urllib.parse import urlparse logger = logging.getLogger(__name__) +_S3_CLIENT_CACHE = {'client': None, 'key': None} +_S3_CLIENT_LOCK = threading.Lock() +_S3_SETTINGS = {'pool_connections': 10} + def is_s3_path(path: str | PathLike[str]) -> bool: """Return True if path is an S3 URI (s3://).""" return str(path).startswith('s3://') +def configure_s3_client_pool(workers: int) -> None: + """Ensure new shared S3 clients have enough connections for loaders.""" + pool_connections = max(10, int(workers)) + if pool_connections == _S3_SETTINGS['pool_connections']: + return + _S3_SETTINGS['pool_connections'] = pool_connections + clear_s3_client_cache() + + +def s3_client_pool_connections() -> int: + """Return the configured shared-client connection-pool size.""" + return _S3_SETTINGS['pool_connections'] + + +def clear_s3_client_cache() -> None: + """Clear cached clients (primarily for tests and credential changes).""" + with _S3_CLIENT_LOCK: + client = _S3_CLIENT_CACHE['client'] + _S3_CLIENT_CACHE.update(client=None, key=None) + close = getattr(client, 'close', None) + if close is not None: + close() + + def _make_s3_client(): """Create a boto3 S3 client. @@ -34,12 +64,83 @@ def _make_s3_client(): 'boto3 is required for s3:// paths. Install with: pip install modelarrayio[s3]' ) from None anon = os.environ.get('MODELARRAYIO_S3_ANON', '').lower() in ('1', 'true', 'yes') - if anon: - from botocore import UNSIGNED + pool_connections = s3_client_pool_connections() + cache_key = (anon, pool_connections) + with _S3_CLIENT_LOCK: + if _S3_CLIENT_CACHE['client'] is not None and _S3_CLIENT_CACHE['key'] == cache_key: + return _S3_CLIENT_CACHE['client'] + + previous_client = _S3_CLIENT_CACHE['client'] + from botocore.config import Config - return boto3.client('s3', config=Config(signature_version=UNSIGNED)) - return boto3.client('s3') + config_kwargs = {'max_pool_connections': pool_connections} + if anon: + from botocore import UNSIGNED + + config_kwargs['signature_version'] = UNSIGNED + client = boto3.client('s3', config=Config(**config_kwargs)) + _S3_CLIENT_CACHE.update(client=client, key=cache_key) + close = getattr(previous_client, 'close', None) + if close is not None: + close() + return client + + +def read_s3_bytes(path: str | PathLike[str], *, decompress_gzip: bool = False) -> bytes: + """Read an S3 object with the shared client and always close its body.""" + if not is_s3_path(path): + raise ValueError(f'Expected an s3:// URI, got {path!r}.') + parsed = urlparse(str(path)) + bucket = parsed.netloc + key = parsed.path.lstrip('/') + if not bucket or not key: + raise ValueError(f'Invalid S3 URI: {path!r}.') + + logger.debug('Loading s3://%s/%s into memory', bucket, key) + response = _make_s3_client().get_object(Bucket=bucket, Key=key) + body = response['Body'] + try: + data = body.read() + finally: + close = getattr(body, 'close', None) + if close is not None: + close() + if decompress_gzip and os.path.basename(key).endswith('.gz'): + return gzip.decompress(data) + return data + + +def _check_local_source(path: str | PathLike[str]) -> Path: + """Validate a local source and give actionable git-annex diagnostics.""" + local_path = Path(path) + if local_path.exists(): + return local_path + hint = '' + if local_path.is_symlink(): + hint = ( + ' The path is a broken symlink; if it belongs to DataLad/git-annex, ' + 'materialize it first (for example, with `datalad get`).' + ) + raise FileNotFoundError(f'Image source does not exist: {local_path}.{hint}') + + +def release_image_buffers(image) -> None: + """Close in-memory file objects after a nibabel image has been decoded. + + Images loaded from S3 retain their downloaded bytes through ``FileHolder`` + objects. Callers that have already copied the decoded data can release + those buffers while continuing to use the parsed image header. + """ + seen: set[int] = set() + for holder in getattr(image, 'file_map', {}).values(): + fileobj = getattr(holder, 'fileobj', None) + if fileobj is None or id(fileobj) in seen: + continue + seen.add(id(fileobj)) + close = getattr(fileobj, 'close', None) + if close is not None: + close() def load_nibabel(path: str | PathLike[str], *, cifti: bool = False): @@ -66,17 +167,9 @@ def load_nibabel(path: str | PathLike[str], *, cifti: bool = False): import nibabel as nb if not is_s3_path(path): - return nb.load(path) - - parsed = urlparse(str(path)) - bucket = parsed.netloc - key = parsed.path.lstrip('/') - - logger.debug('Loading s3://%s/%s into memory', bucket, key) - data = _make_s3_client().get_object(Bucket=bucket, Key=key)['Body'].read() + return nb.load(_check_local_source(path)) - if os.path.basename(key).endswith('.gz'): - data = gzip.decompress(data) + data = read_s3_bytes(path, decompress_gzip=True) from nibabel.filebasedimages import FileHolder diff --git a/test/test_benchmark_harness.py b/test/test_benchmark_harness.py new file mode 100644 index 0000000..0620a08 --- /dev/null +++ b/test/test_benchmark_harness.py @@ -0,0 +1,240 @@ +"""Regression tests for benchmark result gating and artifact isolation.""" + +from __future__ import annotations + +import csv +from pathlib import Path +from types import SimpleNamespace + +import h5py +import numpy as np +import pytest + +import harness_common +import output_compare +import run_benchmarks as harness + + +def test_r_gate_and_missing_output_make_case_fail() -> None: + assert harness._case_failure_reasons( + returncode=0, + outputs_present=True, + r_gate='failed', + meets_acceptance_target=True, + ) == ['r_gate'] + assert harness._case_failure_reasons( + returncode=0, + outputs_present=False, + r_gate='failed_missing_output', + meets_acceptance_target=None, + ) == ['missing_output', 'r_gate'] + assert harness._case_status([], 'not_run') == 'unverified' + assert harness._case_status([], 'passed') == 'passed' + assert harness._case_status(['r_gate'], 'failed') == 'failed' + + +def test_missing_r_gate_executable_is_recorded_as_failure(tmp_path: Path) -> None: + output = tmp_path / 'output.h5' + output.touch() + + status, _seconds = harness._run_r_gate( + 'definitely-missing-r-command {output}', + [output], + 'hdf5', + ['FA'], + tmp_path / 'r-gate.log', + ) + + assert status == 'failed' + assert 'Could not execute R gate command' in (tmp_path / 'r-gate.log').read_text() + + +def test_shard_artifacts_are_isolated(tmp_path: Path) -> None: + first = harness._run_artifact_paths(tmp_path, 'node-01') + second = harness._run_artifact_paths(tmp_path, 'node-02') + + assert first[0] == tmp_path / 'shard-node-01' + assert first[1] == tmp_path / 'results.node-01.csv' + assert first[2] == tmp_path / 'run_meta.node-01.json' + assert set(first).isdisjoint(second) + + +@pytest.mark.parametrize('value', ['../escape', '/absolute', 'contains space', '.hidden']) +def test_shard_identifier_rejects_unsafe_values(value: str) -> None: + with pytest.raises(ValueError, match='result-shard'): + harness._validated_shard(value) + + +@pytest.mark.parametrize('value', ['../escape', '/absolute', 'contains space', '.hidden']) +def test_run_identifier_rejects_unsafe_values(value: str) -> None: + with pytest.raises(ValueError, match='run-id'): + harness._validated_run_id(value) + + +def test_completed_case_keys_support_resume_and_legacy_rows(tmp_path: Path) -> None: + results = tmp_path / 'results.csv' + with results.open('w', newline='', encoding='utf-8') as stream: + writer = csv.DictWriter(stream, fieldnames=['case_key', 'case_id', 'status']) + writer.writeheader() + writer.writerow( + { + 'case_key': 'stable-case', + 'case_id': 'stable-case-retry-01', + 'status': 'passed', + } + ) + writer.writerow({'case_key': '', 'case_id': 'legacy-case', 'status': 'failed'}) + + assert harness._completed_case_keys(results) == {'stable-case', 'legacy-case'} + assert harness._recorded_failure_count(results) == 1 + + +def test_available_cpu_count_prefers_process_allocation(monkeypatch) -> None: + monkeypatch.setattr(harness.os, 'process_cpu_count', lambda: 3, raising=False) + monkeypatch.setattr(harness.os, 'cpu_count', lambda: 64) + + assert harness._available_cpu_count() == 3 + + +def test_full_tier_covers_complete_reference_grid(monkeypatch) -> None: + monkeypatch.setattr(harness, '_available_cpu_count', lambda: 14) + args = SimpleNamespace( + tier='full', + workers=None, + dtypes=None, + target_chunk_mb=None, + backends=['hdf5', 'tiledb'], + compression_workers=None, + ) + + cases = harness._case_grid(args) + + assert len(cases) == 440 + assert sum(case['backend'] == 'hdf5' for case in cases) == 280 + assert sum(case['backend'] == 'tiledb' for case in cases) == 160 + assert {case['workers'] for case in cases} == {1, 2, 4, 8, 14} + assert {case['dtype'] for case in cases} == {'float32', 'float64'} + assert {case['target_chunk_mb'] for case in cases} == {8.0, 32.0, 64.0, 128.0} + assert { + (case['compression'], case['compression_level']) + for case in cases + if case['backend'] == 'hdf5' + } == {('gzip', 1), ('gzip', 4), ('none', 0)} + assert { + (case['compression'], case['compression_level']) + for case in cases + if case['backend'] == 'tiledb' + } == {('zstd', 1), ('zstd', 3), ('zstd', 5), ('none', 0)} + + +def test_resolved_executable_keeps_path_lookup_and_absolutizes_paths(tmp_path: Path) -> None: + executable = tmp_path / 'bin' / 'modelarrayio' + executable.parent.mkdir() + executable.touch() + + assert harness._resolved_executable('modelarrayio') == 'modelarrayio' + assert harness._resolved_executable(str(executable)) == str(executable.resolve()) + + +def test_long_cohort_scalar_names_and_modality_are_recorded() -> None: + rows = [ + {'scalar_name': 'FA', 'source_file': 'first.nii.gz'}, + {'scalar_name': 'FA', 'source_file': 'second.nii.gz'}, + {'scalar_name': 'MD', 'source_file': 'third.nii.gz'}, + ] + + assert harness_common.scalar_names(rows, []) == ['FA', 'MD'] + assert harness_common.source_modality(harness_common.scalar_source_paths(rows, [])) == 'nifti' + assert harness._resolve_split_files(None, []) is False + assert harness._resolve_split_files(None, ['FA', 'MD']) is True + assert harness._resolve_split_files(False, ['FA', 'MD']) is False + + +def test_wide_scalar_inference_excludes_demographic_columns() -> None: + rows = [ + { + 'subject_id': 'sub-01', + 'age': '12', + 'source_mask_file': 'mask.nii.gz', + 'FA': 'fa.nii.gz', + 'MD': 'md.nii.gz', + } + ] + + assert harness_common.infer_scalar_columns(rows) == ['FA', 'MD'] + + +def test_scalar_inference_rejects_empty_and_uninferable_cohorts() -> None: + # Both harnesses share this check, so compare_revisions.py also fails loudly instead of + # treating a wide cohort with no recognizable scalar columns as long-format. + with pytest.raises(ValueError, match='no rows'): + harness_common.infer_scalar_columns([]) + with pytest.raises(ValueError, match='pass --scalar-columns'): + harness_common.infer_scalar_columns([{'subject_id': 'sub-01', 'age': '12'}]) + + +@pytest.mark.parametrize( + ('path', 'expected'), + [ + ('s3://bucket/data.dscalar.nii', 'cifti'), + ('input.pscalar.nii', 'cifti'), + ('input.pconn.nii', 'cifti'), + ('input.mif.gz', 'mif'), + ], +) +def test_benchmark_modality_supports_cifti_and_mif(path: str, expected: str) -> None: + assert harness_common.source_modality([path]) == expected + + +def test_write_csv_uses_lf_and_unions_columns_in_first_seen_order(tmp_path: Path) -> None: + path = tmp_path / 'nested' / 'rows.csv' + harness_common.write_csv(path, [{'a': 1, 'b': 2}, {'b': 3, 'c': 4}]) + + assert path.read_bytes() == b'a,b,c\n1,2,\n,3,4\n' + + +def test_output_path_follows_the_converter_naming() -> None: + case_dir = Path('case') + + assert harness_common.output_path(case_dir, 'hdf5') == case_dir / 'modelarray.h5' + assert harness_common.output_path(case_dir, 'tiledb', 'FA') == case_dir / 'FA_modelarray.tdb' + + +def _write_hdf5_output(path: Path, values: np.ndarray, names: list[str]) -> Path: + with h5py.File(path, 'w') as handle: + handle['scalars/FA/values'] = values + handle['scalars/FA/column_names'] = np.asarray(names, dtype=h5py.string_dtype()) + return path + + +def test_compare_outputs_reports_value_nan_and_source_differences(tmp_path: Path) -> None: + values = np.asarray([[1.0, np.nan, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32) + names = ['sub-01', 'sub-02'] + reference = _write_hdf5_output(tmp_path / 'reference.h5', values, names) + + same = output_compare.compare_outputs( + reference, _write_hdf5_output(tmp_path / 'same.h5', values.copy(), names), 'FA', 2 + ) + assert same == { + 'shape_equal': True, + 'dtype_equal': True, + 'sources_equal': True, + 'nan_equal': True, + 'values_equal': True, + 'max_abs_difference': 0.0, + } + + changed = values.copy() + changed[1, 2] = 6.5 + changed[0, 1] = 2.0 + different = output_compare.compare_outputs( + reference, _write_hdf5_output(tmp_path / 'diff.h5', changed, names[::-1]), 'FA', 2 + ) + assert different == { + 'shape_equal': True, + 'dtype_equal': True, + 'sources_equal': False, + 'nan_equal': False, + 'values_equal': False, + 'max_abs_difference': pytest.approx(0.5), + } diff --git a/test/test_cifti_to_h5_unit.py b/test/test_cifti_to_h5_unit.py index a449c30..2ed3caf 100644 --- a/test/test_cifti_to_h5_unit.py +++ b/test/test_cifti_to_h5_unit.py @@ -7,6 +7,7 @@ import numpy as np import pandas as pd import pytest +import tiledb from modelarrayio.cli import cifti_to_h5 @@ -19,35 +20,23 @@ def test_cifti_to_h5_raises_when_scalar_sources_missing(monkeypatch) -> None: def test_cifti_to_h5_tiledb_split_outputs_and_parcels(monkeypatch, tmp_path: Path) -> None: scalar_sources = {'FA': ['fa1.nii'], 'MD': ['md1.nii']} - write_calls = [] - parcel_calls = [] - monkeypatch.setattr(cifti_to_h5, 'build_scalar_sources', lambda _cohort: scalar_sources) monkeypatch.setattr( cifti_to_h5, '_get_cifti_parcel_info', - lambda _first: ('pscalar', {'parcel_id': np.array(['P1'])}), + lambda _first: ('pscalar', {'parcel_id': np.array(['P1', 'P2'])}), ) monkeypatch.setattr( cifti_to_h5, - 'extract_cifti_scalar_data', - lambda source_file, reference_brain_names=None: ( - np.array([1.0, 2.0], dtype=np.float32), - ['brain-a'], - ), - ) - monkeypatch.setattr( - cifti_to_h5.cli_utils, - 'write_tiledb_scalar_matrices', - lambda output, scalars, sources, **kwargs: write_calls.append( - (Path(output), scalars, sources) + 'load_cohort_cifti', + lambda *_args, **_kwargs: ( + { + 'FA': np.array([[1.0, 2.0]], dtype=np.float32), + 'MD': np.array([[3.0, 4.0]], dtype=np.float32), + }, + ['brain-a', 'brain-b'], ), ) - monkeypatch.setattr( - cifti_to_h5.cli_utils, - 'write_tiledb_parcel_arrays', - lambda output, parcels: parcel_calls.append((Path(output), parcels)), - ) status = cifti_to_h5.cifti_to_h5( cohort_long=pd.DataFrame( @@ -60,10 +49,16 @@ def test_cifti_to_h5_tiledb_split_outputs_and_parcels(monkeypatch, tmp_path: Pat ) assert status == 0 - assert len(write_calls) == 2 - assert sorted(path.name for path, _, _ in write_calls) == ['FA_store.tdb', 'MD_store.tdb'] - assert len(parcel_calls) == 2 - assert sorted(path.name for path, _ in parcel_calls) == ['FA_store.tdb', 'MD_store.tdb'] + for scalar_name in scalar_sources: + scalar_output = tmp_path / f'{scalar_name}_store.tdb' + assert ( + tiledb.object_type(str(scalar_output / 'scalars' / scalar_name / 'values')) == 'array' + ) + assert tiledb.object_type(str(scalar_output / 'parcels' / 'parcel_id')) == 'array' + assert ( + tiledb.object_type(str(scalar_output / 'scalars' / scalar_name / 'column_names')) + == 'array' + ) def test_cifti_to_h5_hdf5_split_outputs_for_dscalar(monkeypatch, tmp_path: Path) -> None: @@ -73,7 +68,9 @@ def test_cifti_to_h5_hdf5_split_outputs_for_dscalar(monkeypatch, tmp_path: Path) monkeypatch.setattr(cifti_to_h5, 'build_scalar_sources', lambda _cohort: scalar_sources) monkeypatch.setattr(cifti_to_h5, '_get_cifti_parcel_info', lambda _first: ('dscalar', {})) monkeypatch.setattr( - cifti_to_h5, 'load_cohort_cifti', lambda _cohort, _workers: (scalars, ['Left', 'Right']) + cifti_to_h5, + 'load_cohort_cifti', + lambda *_args, **_kwargs: (scalars, ['Left', 'Right']), ) monkeypatch.setattr( cifti_to_h5, diff --git a/test/test_cifti_utils.py b/test/test_cifti_utils.py index d2f354e..91f8a5e 100644 --- a/test/test_cifti_utils.py +++ b/test/test_cifti_utils.py @@ -2,12 +2,51 @@ from __future__ import annotations +import nibabel as nb import numpy as np +import pandas as pd import pytest -from nibabel.cifti2.cifti2_axes import ScalarAxis +from nibabel.cifti2.cifti2_axes import BrainModelAxis, ScalarAxis from utils import make_dscalar, make_parcels_axis, make_pconn, make_pscalar -from modelarrayio.utils.cifti import brain_names_to_dataframe, extract_cifti_scalar_data +from modelarrayio.utils import cifti as cifti_utils +from modelarrayio.utils.cifti import ( + _build_cifti_raw_read_plan, + _read_cifti_row_raw, + brain_names_to_dataframe, + cifti_spatial_signature, + extract_cifti_scalar_data, + load_cohort_cifti, +) + + +def _write_dscalar(path, mask, values, *, map_name='synthetic', metadata=None): + """Write a real dscalar ``.nii`` file to disk and return its path string.""" + header = nb.cifti2.Cifti2Header.from_axes( + (ScalarAxis([map_name]), BrainModelAxis.from_mask(mask)) + ) + if metadata: + # Provenance-style metadata lengthens the CIFTI XML and shifts vox_offset + # without touching the brain-model axis, as production pipelines do. + header.matrix.metadata = nb.cifti2.Cifti2MetaData(metadata) + image = nb.Cifti2Image(np.asarray(values, dtype=np.float32).reshape(1, -1), header=header) + image.to_filename(str(path)) + return str(path) + + +def _spy_raw_reads(monkeypatch): + """Record which sources the raw fast path successfully read from disk.""" + successes: list[str] = [] + real = cifti_utils._read_cifti_row_raw + + def _wrapped(source, plan): + row = real(source, plan) + if row is not None: + successes.append(str(source)) + return row + + monkeypatch.setattr(cifti_utils, '_read_cifti_row_raw', _wrapped) + return successes class _FakeHeader: @@ -24,8 +63,8 @@ def __init__(self, axes, data: np.ndarray): self.ndim = len(axes) self._data = data - def get_fdata(self) -> np.ndarray: - return self._data + def get_fdata(self, dtype=np.float64, caching='fill') -> np.ndarray: + return self._data.astype(dtype, copy=False) def test_extract_cifti_scalar_data_returns_data_and_names() -> None: @@ -40,6 +79,20 @@ def test_extract_cifti_scalar_data_returns_data_and_names() -> None: assert names.shape == (2,) +def test_extract_cifti_path_releases_owned_image_buffers(monkeypatch) -> None: + mask = np.zeros((2, 2, 2), dtype=bool) + mask[0, 0, 0] = True + image = make_dscalar(mask, np.array([1.0], dtype=np.float32)) + released = [] + monkeypatch.setattr(cifti_utils, 'load_nibabel', lambda *_args, **_kwargs: image) + monkeypatch.setattr(cifti_utils, 'release_image_buffers', released.append) + + data, _ = extract_cifti_scalar_data('source.dscalar.nii') + + np.testing.assert_array_equal(data, np.array([1.0], dtype=np.float32)) + assert released == [image] + + def test_extract_cifti_scalar_data_rejects_wrong_axis_count() -> None: image = _FakeCiftiImage([ScalarAxis(['a'])], np.array([1.0], dtype=np.float32)) @@ -136,3 +189,436 @@ def test_brain_names_to_dataframe() -> None: assert 'structure_id' in gdf.columns assert gdf['vertex_id'].tolist() == [0, 1, 2] assert len(struct_strings) == 2 # factorize unique structures + + +@pytest.mark.parametrize('storage_dtype', ['float32', 'float64']) +def test_extract_cifti_requests_storage_dtype_end_to_end(storage_dtype: str) -> None: + mask = np.zeros((2, 2, 2), dtype=bool) + mask[0, 0, 0] = True + mask[1, 1, 1] = True + header = nb.cifti2.Cifti2Header.from_axes( + (ScalarAxis(['precision']), BrainModelAxis.from_mask(mask)) + ) + image = nb.Cifti2Image( + np.array([[1.123456789, 2.987654321]], dtype=np.float64), + header=header, + ) + + data, _ = extract_cifti_scalar_data(image, storage_dtype=storage_dtype) + + assert data.dtype == np.dtype(storage_dtype) + + +def test_cifti_signature_ignores_scalar_labels_but_rejects_spatial_changes() -> None: + mask = np.zeros((2, 2, 2), dtype=bool) + mask[0, 0, 0] = True + mask[1, 1, 1] = True + brain_axis = BrainModelAxis.from_mask(mask) + image_a = nb.Cifti2Image( + np.array([[1.0, 2.0]], dtype=np.float32), + header=nb.cifti2.Cifti2Header.from_axes((ScalarAxis(['FA']), brain_axis)), + ) + image_b = nb.Cifti2Image( + np.array([[3.0, 4.0]], dtype=np.float32), + header=nb.cifti2.Cifti2Header.from_axes((ScalarAxis(['MD']), brain_axis)), + ) + + reference_signature = cifti_spatial_signature(image_a) + assert cifti_spatial_signature(image_b) == reference_signature + + changed_mask = np.zeros_like(mask) + changed_mask[0, 0, 1] = True + changed_mask[1, 1, 0] = True + changed = make_dscalar(changed_mask, np.array([1.0, 2.0])) + with pytest.raises(ValueError, match='Inconsistent spatial axis'): + extract_cifti_scalar_data(changed, reference_signature=reference_signature) + + +def test_cifti_signature_covers_both_pconn_axes() -> None: + row_axis = make_parcels_axis(['row-a', 'row-b']) + column_axis = make_parcels_axis(['column-a', 'column-b']) + changed_column_axis = make_parcels_axis(['column-a', 'changed-column']) + values = np.arange(4, dtype=np.float32).reshape(2, 2) + reference = nb.Cifti2Image( + values, + header=nb.cifti2.Cifti2Header.from_axes((row_axis, column_axis)), + ) + changed = nb.Cifti2Image( + values, + header=nb.cifti2.Cifti2Header.from_axes((row_axis, changed_column_axis)), + ) + + assert cifti_spatial_signature(reference) != cifti_spatial_signature(changed) + + +def test_load_cohort_cifti_reuses_first_decoded_row_and_cheap_validation(monkeypatch) -> None: + mask = np.zeros((2, 2, 2), dtype=bool) + mask[0, 0, 0] = True + mask[1, 1, 1] = True + brain_axis = BrainModelAxis.from_mask(mask) + + class _CountingImage: + def __init__(self, scalar_label: str, values: list[float]) -> None: + self.header = nb.cifti2.Cifti2Header.from_axes( + (ScalarAxis([scalar_label]), brain_axis) + ) + self.ndim = 2 + self._values = np.asarray([values], dtype=np.float64) + self.decode_count = 0 + + def get_fdata(self, dtype=np.float64, caching='fill'): + self.decode_count += 1 + return self._values.astype(dtype, copy=True) + + images = { + 'first.dscalar.nii': _CountingImage('FA', [1.0, 2.0]), + 'second.dscalar.nii': _CountingImage('different-scalar-label', [3.0, 4.0]), + } + open_counts = dict.fromkeys(images, 0) + + def _load(source, *, cifti=False): + assert cifti is True + open_counts[source] += 1 + return images[source] + + validation_calls = [] + real_extract = cifti_utils.extract_cifti_scalar_data + + def _extract(*args, **kwargs): + validation_calls.append(kwargs.get('reference_signature')) + return real_extract(*args, **kwargs) + + monkeypatch.setattr(cifti_utils, 'load_nibabel', _load) + monkeypatch.setattr(cifti_utils, 'extract_cifti_scalar_data', _extract) + cohort = pd.DataFrame( + { + 'scalar_name': ['FA', 'FA'], + 'source_file': ['first.dscalar.nii', 'second.dscalar.nii'], + } + ) + + scalars, _ = load_cohort_cifti(cohort, workers=2, storage_dtype='float32') + matrix = scalars['FA'] + + np.testing.assert_array_equal(matrix, np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)) + assert matrix.flags.c_contiguous + assert open_counts == {'first.dscalar.nii': 1, 'second.dscalar.nii': 1} + assert images['first.dscalar.nii'].decode_count == 1 + assert images['second.dscalar.nii'].decode_count == 1 + assert validation_calls[0] is None + assert validation_calls[1] == scalars.reference_signature + + +def test_signature_validated_extract_skips_full_name_comparison(monkeypatch) -> None: + mask = np.zeros((2, 2, 2), dtype=bool) + mask[0, 0, 0] = True + mask[1, 1, 1] = True + image = make_dscalar(mask, np.array([1.0, 2.0], dtype=np.float32)) + signature = cifti_spatial_signature(image) + _, names = extract_cifti_scalar_data(image) + + def _unexpected_comparison(*_args, **_kwargs): + raise AssertionError('compact signature should replace full label comparison') + + monkeypatch.setattr(cifti_utils.np, 'array_equal', _unexpected_comparison) + values, returned_names = extract_cifti_scalar_data( + image, + reference_brain_names=names, + reference_signature=signature, + spatial_signature=signature, + ) + + np.testing.assert_allclose(values, [1.0, 2.0]) + assert returned_names is names + + +def test_load_cohort_cifti_releases_image_when_signature_validation_fails(monkeypatch) -> None: + first_mask = np.zeros((2, 2, 2), dtype=bool) + first_mask[0, 0, 0] = True + first_mask[0, 1, 1] = True + changed_mask = np.zeros_like(first_mask) + changed_mask[0, 0, 0] = True + changed_mask[1, 1, 1] = True + images = { + 'first.dscalar.nii': make_dscalar(first_mask, np.array([1.0, 2.0], dtype=np.float32)), + 'changed.dscalar.nii': make_dscalar(changed_mask, np.array([3.0, 4.0], dtype=np.float32)), + } + released = [] + monkeypatch.setattr(cifti_utils, 'load_nibabel', lambda source, **_kwargs: images[source]) + monkeypatch.setattr(cifti_utils, 'release_image_buffers', released.append) + cohort = pd.DataFrame( + { + 'scalar_name': ['FA', 'FA'], + 'source_file': ['first.dscalar.nii', 'changed.dscalar.nii'], + } + ) + + scalars, _ = load_cohort_cifti(cohort) + with pytest.raises(cifti_utils.SourceLoadError, match=r'changed\.dscalar\.nii'): + _ = scalars['FA'] + + assert released == [images['first.dscalar.nii'], images['changed.dscalar.nii']] + + +def _toy_mask() -> np.ndarray: + mask = np.zeros((4, 4, 4), dtype=bool) + mask[0, 0, 0] = mask[1, 1, 1] = mask[2, 2, 2] = mask[3, 3, 3] = True + return mask + + +@pytest.mark.parametrize('workers', [1, 2]) +def test_load_cohort_cifti_raw_fast_path_matches_nibabel(tmp_path, monkeypatch, workers) -> None: + mask = _toy_mask() + n = int(mask.sum()) + truth = [] + rows = [] + for i in range(4): + values = np.arange(n, dtype=np.float32) + i * 10.0 + if i == 2: + values[1] = np.nan # NaN positions must survive the raw read + path = _write_dscalar(tmp_path / f'sub{i}.dscalar.nii', mask, values) + rows.append({'scalar_name': 'FA', 'source_file': path}) + truth.append(values) + cohort = pd.DataFrame(rows) + + raw_reads = _spy_raw_reads(monkeypatch) + scalars, _ = load_cohort_cifti(cohort, workers=workers, storage_dtype='float32') + matrix = scalars['FA'] + + np.testing.assert_array_equal(matrix, np.stack(truth)) + assert matrix.flags.c_contiguous + # The three non-reference subjects are read straight from their data block. + assert {f'sub{i}.dscalar.nii' for i in (1, 2, 3)} <= { + source.split('/')[-1] for source in raw_reads + } + + +def test_load_cohort_cifti_raw_path_handles_per_scalar_offsets(tmp_path, monkeypatch) -> None: + # A longer map name lengthens the CIFTI XML extension and shifts vox_offset, + # so the raw plan must be rebuilt per scalar rather than shared globally. + mask = _toy_mask() + n = int(mask.sum()) + rows = [] + truth = {} + for scalar, map_name in (('thickness', 'thickness'), ('area', 'desc-fwhm25_area_long_map')): + truth[scalar] = [] + for i in range(3): + values = np.arange(n, dtype=np.float32) + hash(scalar) % 7 + i + path = _write_dscalar( + tmp_path / f'{scalar}_sub{i}.dscalar.nii', mask, values, map_name=map_name + ) + rows.append({'scalar_name': scalar, 'source_file': path}) + truth[scalar].append(values) + cohort = pd.DataFrame(rows) + + offsets = {nb.load(row['source_file']).dataobj.offset for row in rows} + assert len(offsets) >= 2 # the two scalars really do have different offsets + + raw_reads = _spy_raw_reads(monkeypatch) + scalars, _ = load_cohort_cifti(cohort, workers=2, storage_dtype='float32') + for scalar in ('thickness', 'area'): + np.testing.assert_array_equal(scalars[scalar], np.stack(truth[scalar])) + # Both scalars used the fast path despite differing offsets. + assert any('thickness_sub' in source for source in raw_reads) + assert any('area_sub' in source for source in raw_reads) + + +def test_load_cohort_cifti_raw_path_tolerates_metadata_only_offset_shift( + tmp_path, monkeypatch +) -> None: + # Production dscalars embed per-subject provenance in the CIFTI metadata, so + # vox_offset differs between files that share one brain-model axis. The raw + # path must still serve them (verified by the brain-model map digest) and + # reproduce the nibabel values exactly. + mask = _toy_mask() + n = int(mask.sum()) + rows = [] + truth = [] + for i in range(3): + values = np.arange(n, dtype=np.float32) + i + metadata = None if i != 1 else {'Provenance': f'job-{i}-sub-{i:05d}-' + 'x' * 40} + path = _write_dscalar(tmp_path / f'sub{i}.dscalar.nii', mask, values, metadata=metadata) + rows.append({'scalar_name': 'FA', 'source_file': path}) + truth.append(values) + cohort = pd.DataFrame(rows) + + offsets = [nb.load(row['source_file']).dataobj.offset for row in rows] + assert offsets[1] != offsets[0] # the offset shift is real + + raw_reads = _spy_raw_reads(monkeypatch) + scalars, _ = load_cohort_cifti(cohort, workers=2, storage_dtype='float32') + + np.testing.assert_array_equal(scalars['FA'], np.stack(truth)) + # Every non-reference source went through the fast path despite the shift. + assert any('sub1.dscalar.nii' in source for source in raw_reads) + assert any('sub2.dscalar.nii' in source for source in raw_reads) + assert scalars.stats['FA'].fast_path_rows == 2 + + +def test_load_cohort_cifti_raw_path_declines_brain_model_mismatch(tmp_path, monkeypatch) -> None: + # A source with the same element count but a different brain-model axis has + # a different MatrixIndicesMap, so the raw path must decline it; the nibabel + # path then validates the spatial axis and rejects the inconsistent file. + mask = _toy_mask() + other = np.zeros_like(mask) + other[0, 0, 1] = other[1, 1, 2] = other[2, 2, 3] = other[3, 3, 0] = True # same count + n = int(mask.sum()) + rows = [] + for i in range(3): + values = np.arange(n, dtype=np.float32) + i + path = _write_dscalar(tmp_path / f'sub{i}.dscalar.nii', other if i == 1 else mask, values) + rows.append({'scalar_name': 'FA', 'source_file': path}) + cohort = pd.DataFrame(rows) + + raw_reads = _spy_raw_reads(monkeypatch) + scalars, _ = load_cohort_cifti(cohort, workers=1, storage_dtype='float32') + with pytest.raises(cifti_utils.SourceLoadError, match='Inconsistent spatial axis'): + scalars['FA'] + assert not any('sub1.dscalar.nii' in source for source in raw_reads) + + +def test_load_cohort_cifti_raw_path_disabled_by_env(tmp_path, monkeypatch) -> None: + mask = _toy_mask() + n = int(mask.sum()) + rows = [] + truth = [] + for i in range(3): + values = np.arange(n, dtype=np.float32) + i + path = _write_dscalar(tmp_path / f'sub{i}.dscalar.nii', mask, values) + rows.append({'scalar_name': 'FA', 'source_file': path}) + truth.append(values) + cohort = pd.DataFrame(rows) + + monkeypatch.setenv('MODELARRAYIO_CIFTI_RAW', '0') + raw_reads = _spy_raw_reads(monkeypatch) + scalars, _ = load_cohort_cifti(cohort, workers=2, storage_dtype='float32') + + np.testing.assert_array_equal(scalars['FA'], np.stack(truth)) + assert raw_reads == [] # every source decoded through nibabel + + +def test_load_cohort_cifti_pconn_row_major_values_are_correct(tmp_path) -> None: + # pconn stores a 2-D matrix flattened row-major. The raw plan self-validates + # against the decoded reference row, so it only engages when the on-disk + # order reproduces ``get_fdata().ravel()`` exactly; either way the staged + # values must match the flattened matrices. + parcel_names = ['a', 'b', 'c'] + rows = [] + truth = [] + for i in range(2): + matrix = np.arange(9, dtype=np.float32).reshape(3, 3) + i + path = str(tmp_path / f'sub{i}.pconn.nii') + make_pconn(parcel_names, matrix).to_filename(path) + rows.append({'scalar_name': 'conn', 'source_file': path}) + truth.append(matrix.ravel()) + cohort = pd.DataFrame(rows) + + scalars, _ = load_cohort_cifti(cohort, workers=2, storage_dtype='float32') + + np.testing.assert_array_equal(scalars['conn'], np.stack(truth)) + + +def test_read_cifti_row_raw_rejects_non_matching_sources(tmp_path) -> None: + mask = _toy_mask() + n = int(mask.sum()) + reference = _write_dscalar(tmp_path / 'ref.dscalar.nii', mask, np.arange(n, dtype=np.float32)) + ref_image = nb.load(reference) + ref_row = ref_image.get_fdata(dtype=np.float32, caching='unchanged').reshape(-1) + plan = _build_cifti_raw_read_plan(ref_image, ref_row, reference, 'float32') + assert plan is not None + + # An s3:// URI and a .gz path are declined outright. + assert _read_cifti_row_raw('s3://bucket/key.dscalar.nii', plan) is None + assert _read_cifti_row_raw(f'{reference}.gz', plan) is None + + # A truncated copy (wrong size) is declined. + truncated = tmp_path / 'short.dscalar.nii' + truncated.write_bytes(b'\x00' * (plan.expected_file_size - 8)) + assert _read_cifti_row_raw(str(truncated), plan) is None + + # A missing file is declined. + assert _read_cifti_row_raw(str(tmp_path / 'missing.dscalar.nii'), plan) is None + + # The reference itself round-trips through the fast path. + np.testing.assert_array_equal(_read_cifti_row_raw(reference, plan), ref_row) + + +def test_build_cifti_raw_read_plan_rejects_in_memory_image() -> None: + # An in-memory image has no on-disk data offset, so no plan can be built. + image = make_dscalar(_toy_mask(), np.arange(int(_toy_mask().sum()), dtype=np.float32)) + row = image.get_fdata(dtype=np.float32).reshape(-1) + assert _build_cifti_raw_read_plan(image, row, 'memory.dscalar.nii', 'float32') is None + + +def test_load_cohort_cifti_raw_fast_path_serves_pscalar(tmp_path, monkeypatch) -> None: + # Parcellated scalars store a ParcelsAxis map at dimension 1, so the raw + # plan must engage for them exactly as it does for dense greyordinates. + parcel_names = [f'p{i}' for i in range(6)] + truth = [] + rows = [] + for i in range(4): + values = np.arange(len(parcel_names), dtype=np.float32) + i * 10.0 + path = str(tmp_path / f'sub{i}.pscalar.nii') + make_pscalar(parcel_names, values).to_filename(path) + rows.append({'scalar_name': 'thickness', 'source_file': path}) + truth.append(values) + cohort = pd.DataFrame(rows) + + raw_reads = _spy_raw_reads(monkeypatch) + scalars, names = load_cohort_cifti(cohort, workers=2, storage_dtype='float32') + matrix = scalars['thickness'] + + np.testing.assert_array_equal(matrix, np.stack(truth)) + assert list(names) == parcel_names + assert {f'sub{i}.pscalar.nii' for i in (1, 2, 3)} <= { + source.split('/')[-1] for source in raw_reads + } + assert scalars.stats['thickness'].fast_path_rows == 3 + + +def test_load_cohort_cifti_rows_match_build_scalar_sources(tmp_path) -> None: + # The loader and the CLI's column_names must be built from the same rows, + # so a blank source cell is skipped by both rather than failing late. + from modelarrayio.utils.misc import build_scalar_sources + + mask = _toy_mask() + n = int(mask.sum()) + rows = [] + for i in range(2): + path = _write_dscalar(tmp_path / f'sub{i}.dscalar.nii', mask, np.full(n, float(i))) + rows.append({'scalar_name': 'FA', 'source_file': path}) + rows.append({'scalar_name': 'FA', 'source_file': ''}) + cohort = pd.DataFrame(rows) + + scalars, _ = load_cohort_cifti(cohort, workers=1, storage_dtype='float32') + matrix = scalars['FA'] + + expected_sources = build_scalar_sources(cohort)['FA'] + assert matrix.shape[0] == len(expected_sources) == 2 + assert list(scalars.sources_by_scalar['FA']) == expected_sources + + +def test_read_cifti_row_raw_reads_file_once(tmp_path, monkeypatch) -> None: + mask = _toy_mask() + n = int(mask.sum()) + reference = _write_dscalar(tmp_path / 'ref.dscalar.nii', mask, np.arange(n, dtype=np.float32)) + other = _write_dscalar( + tmp_path / 'other.dscalar.nii', mask, np.arange(n, dtype=np.float32) + 1 + ) + ref_image = nb.load(reference) + ref_row = ref_image.get_fdata(dtype=np.float32).reshape(-1) + plan = _build_cifti_raw_read_plan(ref_image, ref_row, reference, 'float32') + assert plan is not None + + opened: list[str] = [] + real_open = open + + def _counting_open(file, *args, **kwargs): + if isinstance(file, (str, bytes)) or hasattr(file, '__fspath__'): + opened.append(str(file)) + return real_open(file, *args, **kwargs) + + monkeypatch.setattr('builtins.open', _counting_open) + row = _read_cifti_row_raw(other, plan) + np.testing.assert_array_equal(row, np.arange(n, dtype=np.float32) + 1) + assert opened.count(other) == 1 diff --git a/test/test_cli_main_and_s3_utils.py b/test/test_cli_main_and_s3_utils.py index 16805d1..16682ab 100644 --- a/test/test_cli_main_and_s3_utils.py +++ b/test/test_cli_main_and_s3_utils.py @@ -70,7 +70,7 @@ def client(service, **kwargs): monkeypatch.setenv('MODELARRAYIO_S3_ANON', '0') s3_utils._make_s3_client() - assert calls[1] == ('s3', {}) + assert calls[1] == ('s3', {'config': ('cfg', {'max_pool_connections': 10})}) def test_make_s3_client_requires_boto3(monkeypatch) -> None: @@ -96,6 +96,19 @@ def test_load_nibabel_local_path(monkeypatch, tmp_path) -> None: np.testing.assert_array_equal(loaded.get_fdata(), data) +def test_load_nibabel_materialized_datalad_symlink(tmp_path) -> None: + annex_object = tmp_path / '.git' / 'annex' / 'objects' / 'image.nii.gz' + annex_object.parent.mkdir(parents=True) + data = np.arange(8, dtype=np.float32).reshape(2, 2, 2) + nb.Nifti1Image(data, np.eye(4)).to_filename(annex_object) + worktree_link = tmp_path / 'image.nii.gz' + worktree_link.symlink_to(annex_object) + + loaded = s3_utils.load_nibabel(worktree_link) + + np.testing.assert_array_equal(loaded.get_fdata(), data) + + def test_load_nibabel_from_s3_bytes(monkeypatch, tmp_path) -> None: data = np.arange(8, dtype=np.float32).reshape(2, 2, 2) image = nb.Nifti1Image(data, np.eye(4)) diff --git a/test/test_cli_utils.py b/test/test_cli_utils.py index 7d7aeee..e12a875 100644 --- a/test/test_cli_utils.py +++ b/test/test_cli_utils.py @@ -8,6 +8,7 @@ import h5py import numpy as np import pandas as pd +import pytest from modelarrayio.cli import utils as cli_utils @@ -31,6 +32,27 @@ def test_prefixed_output_path_sanitizes_prefix(tmp_path: Path) -> None: assert prefixed.name == 'p.value_result_stats.h5' +def test_conversion_output_paths_rejects_sanitized_collisions(tmp_path: Path) -> None: + with pytest.raises(ValueError, match='map to the same split output'): + cli_utils.conversion_output_paths( + tmp_path / 'output.h5', + ['a b', 'a_b'], + split_outputs=True, + ) + + +@pytest.mark.parametrize('scalar_name', ['a/b', r'a\b', '..']) +def test_conversion_output_paths_rejects_unsafe_scalar_names( + tmp_path: Path, scalar_name: str +) -> None: + with pytest.raises(ValueError, match='Unsafe scalar name'): + cli_utils.conversion_output_paths( + tmp_path / 'output.h5', + [scalar_name], + split_outputs=False, + ) + + def test_write_table_dataset_writes_transposed_data_and_attrs(tmp_path: Path) -> None: table = pd.DataFrame({'a': [1, 2], 'b': [3, 4]}) h5_path = tmp_path / 'table.h5' @@ -50,7 +72,7 @@ def test_write_table_dataset_writes_transposed_data_and_attrs(tmp_path: Path) -> def test_write_hdf5_scalar_matrices_skips_empty_rows_and_writes_values(tmp_path: Path) -> None: h5_path = tmp_path / 'scalars.h5' with h5py.File(h5_path, 'w') as h5_file: - cli_utils.write_hdf5_scalar_matrices( + metrics = cli_utils.write_hdf5_scalar_matrices( h5_file, scalars={ 'FA': [ @@ -73,27 +95,28 @@ def test_write_hdf5_scalar_matrices_skips_empty_rows_and_writes_values(tmp_path: assert 'scalars/MD/values' not in h5_file values = h5_file['scalars/FA/values'][...] np.testing.assert_array_equal(values, np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)) + assert h5_file['scalars/FA/column_names'].asstr()[...].tolist() == ['sub-1', 'sub-2'] + assert set(metrics) == {'FA'} + assert metrics['FA']['backend'] == 'hdf5' + assert metrics['FA']['matrix_shape'] == [2, 2] def test_write_tiledb_scalar_matrices_calls_column_name_writer( monkeypatch, tmp_path: Path ) -> None: - called = {'create': [], 'columns': [], 'write': []} + called = {'create': [], 'columns': []} def _fake_create(*args, **kwargs): called['create'].append((args, kwargs)) + kwargs['metrics'].update({'dense_submissions': 1}) def _fake_write_columns(*args, **kwargs): called['columns'].append((args, kwargs)) - def _fake_write_rows(*args, **kwargs): - called['write'].append((args, kwargs)) - - monkeypatch.setattr(cli_utils.tiledb_storage, 'create_empty_scalar_matrix_array', _fake_create) + monkeypatch.setattr(cli_utils.tiledb_storage, 'create_scalar_matrix_array', _fake_create) monkeypatch.setattr(cli_utils.tiledb_storage, 'write_column_names', _fake_write_columns) - monkeypatch.setattr(cli_utils.tiledb_storage, 'write_rows_in_column_stripes', _fake_write_rows) - cli_utils.write_tiledb_scalar_matrices( + metrics = cli_utils.write_tiledb_scalar_matrices( tmp_path / 'out.tdb', scalars={'FA': [np.array([1.0, 2.0], dtype=np.float32)], 'MD': []}, sources_by_scalar={'FA': ['sub-1.nii.gz'], 'MD': []}, @@ -108,7 +131,18 @@ def _fake_write_rows(*args, **kwargs): assert len(called['create']) == 1 assert len(called['columns']) == 1 - assert len(called['write']) == 1 + create_args, create_kwargs = called['create'][0] + assert create_args[1] == 'scalars/FA/values' + np.testing.assert_array_equal(create_args[2], [np.array([1.0, 2.0], dtype=np.float32)]) + assert create_args[3] == ['sub-1.nii.gz'] + assert create_kwargs['tile_voxels'] == 16 + assert metrics == { + 'FA': { + 'dense_submissions': 1, + 'backend': 'tiledb', + 'scalar_name': 'FA', + } + } def test_write_hdf5_and_tiledb_parcel_arrays(monkeypatch, tmp_path: Path) -> None: diff --git a/test/test_comparison_report.py b/test/test_comparison_report.py new file mode 100644 index 0000000..5105c1d --- /dev/null +++ b/test/test_comparison_report.py @@ -0,0 +1,188 @@ +"""Tests for the original-vs-optimized findings report (``create_comparison_report.py``).""" + +from __future__ import annotations + +import csv +from pathlib import Path + +import pytest + +# The report renders PDFs with ReportLab, an optional benchmark tool that is not a package +# dependency; skip these tests (rather than fail the whole session) when it is not installed. +pytest.importorskip('reportlab') + +import comparison_runs +import create_comparison_report as comparison_report + +# --- comparison report: production-scale section ------------------------------ + + +def test_scale_section_renders_every_step_and_is_empty_without_rows(tmp_path: Path) -> None: + styles = comparison_report._styles() + assert comparison_report._scale_section([], 'note', styles) == [] + + scale = tmp_path / 'scale.csv' + with scale.open('w', newline='') as stream: + fields = [ + 'step', + 'scalars', + 'subjects', + 'backend', + 'workers', + 'baseline_median_wall_seconds', + 'optimized_median_wall_seconds', + 'speedup', + 'input_files_per_minute', + 'peak_memory_gb', + 'fragment_count', + 'fast_path_rows', + 'passed', + ] + writer = csv.DictWriter(stream, fieldnames=fields) + writer.writeheader() + writer.writerow( + { + 'step': '3_full_scalar_original_vs_optimized', + 'scalars': 'area', + 'subjects': '8693', + 'backend': 'hdf5', + 'workers': '1', + 'baseline_median_wall_seconds': '11854.8', + 'optimized_median_wall_seconds': '429.9', + 'speedup': '27.57', + } + ) + writer.writerow( + { + 'step': '3_full_scalar_worker_sweep', + 'scalars': 'area', + 'subjects': '8693', + 'backend': 'tiledb', + 'workers': '8', + 'optimized_median_wall_seconds': '68.6', + 'input_files_per_minute': '7603', + 'peak_memory_gb': '20.4', + 'fragment_count': '1', + 'fast_path_rows': '8692', + } + ) + writer.writerow( + { + 'step': '4_scalar_sharded_array', + 'scalars': 'curv', + 'subjects': '8693', + 'backend': 'hdf5', + 'workers': '8', + 'optimized_median_wall_seconds': '134.5', + 'passed': 'True', + } + ) + rows = comparison_report._load_scale(scale) + assert len(rows) == 3 + flowables = comparison_report._scale_section(rows, 'Original run: 112 h first scalar.', styles) + # A page break, heading, note, and one heading+table per populated step. + assert len(flowables) >= 8 + assert comparison_report._load_scale(None) == [] + + +# --- comparison report: text follows the measured runs ---------------------------------- + + +def _group(modality: str, backend: str, baseline: float, optimized: float, *, equal: bool = True): + return comparison_runs.Group( + modality=modality, + backend=backend, + n_subjects=100, + n_elements=10, + n_scalars=1, + baseline_wall=[baseline], + optimized_wall=[optimized], + values_equal=equal, + ) + + +def _story_text(flowables) -> str: + """Flatten the report story, including paragraphs nested in table cells.""" + parts = [] + for flowable in flowables: + if hasattr(flowable, 'getPlainText'): + parts.append(flowable.getPlainText()) + parts.extend( + _story_text(cell for cell in row if not isinstance(cell, str)) + for row in getattr(flowable, '_cellvalues', []) + ) + return ' '.join(parts) + + +def _report_meta(grouped) -> dict[str, object]: + return { + 'baseline_revision': '9ab13dd', + 'platform': 'CUBIC compute nodes', + 'versions': dict.fromkeys(('python', 'numpy', 'nibabel', 'h5py', 'tiledb'), '1.0'), + 'settings': {'repetitions': '3', 'warmups': '1', 'workers': '1', 'target_chunk_mb': '2'}, + 'cohorts': { + modality: { + 'cohort': f'/c/{modality}.csv', + 'cohort_sha256': 'ab' * 32, + 'scalar_names': ['x'], + } + for modality in grouped + }, + 'cohort_sources': {'cifti': 'RBC subset, internal & real'}, + } + + +def test_report_text_names_what_was_measured() -> None: + # The pinned CUBIC medians: TileDB beats HDF5 for NIfTI and MIF, so no HDF5 claim is made. + cubic = [ + _group('nifti', 'hdf5', 3.02, 2.17), + _group('nifti', 'tiledb', 3.41, 2.33), + _group('cifti', 'hdf5', 385.78, 12.12), + _group('cifti', 'tiledb', 396.47, 12.56), + _group('mif', 'hdf5', 10.76, 2.91), + _group('mif', 'tiledb', 14.52, 3.53), + ] + assert comparison_report._scope(cubic) == 'NIfTI, CIFTI, and MIF on both HDF5 and TileDB' + assert comparison_report._largest_gain(cubic) == ( + 'The largest gain is CIFTI on HDF5 (31.8x), where the raw-read fast path removes the ' + 'per-file XML header parse.' + ) + mif_only = [_group('mif', 'hdf5', 2.0, 1.0), _group('mif', 'tiledb', 4.0, 1.0)] + assert comparison_report._scope(mif_only) == 'MIF on both HDF5 and TileDB' + assert 'one dense fragment' in comparison_report._largest_gain(mif_only) + assert comparison_report._scope([_group('nifti', 'hdf5', 2.0, 1.0)]) == 'NIfTI on HDF5' + + assert comparison_report._count('3', 'repetition') == 'three repetitions' + assert comparison_report._count('1', 'warm-up') == 'one warm-up' + assert comparison_report._count('0', 'warm-up') == 'no warm-ups' + assert comparison_report._count('1/3', 'repetition') == '1/3 repetitions' + settings = [{'target_chunk_mb': 32}, {'target_chunk_mb': 2.0}, {}] + assert comparison_report._run_setting(settings, 'target_chunk_mb') == '2/32' + with pytest.raises(SystemExit, match='MODALITY=TEXT'): + comparison_report._cohort_sources(['dwi=somewhere']) + + +def test_report_verdict_follows_the_results() -> None: + styles = comparison_report._styles() + passing = {'cifti': [_group('cifti', 'hdf5', 30.0, 3.0), _group('cifti', 'tiledb', 30.0, 3.0)]} + text = _story_text(comparison_report._report_story(passing, _report_meta(passing), styles)) + assert 'It is safe to adopt' in text + assert 'for CIFTI on both HDF5 and TileDB' in text + assert 'three counterbalanced repetitions after one discarded warm-up' in text + assert 'RBC subset, internal & real' in text + assert 'synth' not in text + assert 'Local' not in text + + failing = { + 'nifti': [ + _group('nifti', 'hdf5', 1.0, 2.0), + _group('nifti', 'tiledb', 2.0, 1.0, equal=False), + ] + } + text = _story_text(comparison_report._report_story(failing, _report_meta(failing), styles)) + assert 'safe to adopt' not in text + assert 'in 1 of the 2 modality x backend configurations' in text + assert 'but only 1 of 2 optimized outputs were bitwise equal' in text + assert 'faster in 1 of 2 configurations and bitwise equal in 1 of 2' in text + assert '1 of 2 configurations passed every check' in text + assert '100 subjects x 1 scalar:' in text diff --git a/test/test_cubic_tools.py b/test/test_cubic_tools.py new file mode 100644 index 0000000..102b818 --- /dev/null +++ b/test/test_cubic_tools.py @@ -0,0 +1,327 @@ +"""Tests for the CUBIC validation tools: source identity, gates, reference compare.""" + +from __future__ import annotations + +import csv +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import h5py +import numpy as np +import pytest +import tiledb + +import compare_to_reference as compare_reference +import materialize_cohort as materialize +import source_identity +import summarize_scale +import tiledb_read_gate as read_gate + +REPO = Path(__file__).resolve().parents[1] + + +# --- source identity ------------------------------------------------------- + + +def test_source_identity_is_deterministic_and_pins_content(tmp_path: Path) -> None: + repo = tmp_path / 'repo' + (repo / 'src').mkdir(parents=True) + (repo / 'src' / 'a.py').write_text('x = 1\n') + (repo / 'pyproject.toml').write_text('[project]\nname = "x"\n') + suite = repo / 'benchmarks' / 'to_modelarray' + suite.mkdir(parents=True) + (suite / 'run.py').write_text('print()\n') + # Ignored-style directories inside a suite must not affect the hash even without git. + (suite / 'results').mkdir() + (suite / 'results' / 'big.csv').write_text('1,2\n') + + first = source_identity.source_identity(repo) + assert first['commit'] is None # no git repository + assert len(first['source_tree_sha256']) == 64 + assert source_identity.source_identity(repo) == first + + (suite / 'results' / 'big.csv').write_text('changed\n') + assert ( + source_identity.source_identity(repo)['source_tree_sha256'] == first['source_tree_sha256'] + ) + + (repo / 'src' / 'a.py').write_text('x = 2\n') + assert ( + source_identity.source_identity(repo)['source_tree_sha256'] != first['source_tree_sha256'] + ) + + +def test_source_identity_on_this_repository_reports_commit() -> None: + identity = source_identity.source_identity(REPO) + assert identity['commit'] is None or len(identity['commit']) == 40 + assert len(identity['source_tree_sha256']) == 64 + + +# --- helpers to build outputs ---------------------------------------------- + + +def _write_tiledb(root: Path, scalar: str, values: np.ndarray, names: list[str], *, rows=None): + uri = root / 'scalars' / scalar / 'values' + uri.parent.mkdir(parents=True, exist_ok=True) + n_rows = rows if rows is not None else values.shape[0] + domain = tiledb.Domain( + tiledb.Dim(name='subject', domain=(0, n_rows - 1), tile=n_rows, dtype=np.int32), + tiledb.Dim(name='element', domain=(0, values.shape[1] - 1), tile=4, dtype=np.int32), + ) + schema = tiledb.ArraySchema( + domain=domain, attrs=[tiledb.Attr(name='values', dtype=np.float32)], sparse=False + ) + tiledb.Array.create(str(uri), schema) + with tiledb.open(str(uri), 'w') as array: + array[0 : values.shape[0], 0 : values.shape[1]] = {'values': values} + array.meta['column_names'] = json.dumps(names) + return uri + + +def _write_hdf5(path: Path, scalar: str, values: np.ndarray, names: list[str]) -> Path: + with h5py.File(path, 'w') as handle: + group = handle.create_group(f'scalars/{scalar}') + group.create_dataset('values', data=values) + group.create_dataset('column_names', data=np.array(names, dtype='S')) + return path + + +# --- read gate --------------------------------------------------------------- + + +def test_read_gate_passes_and_fails_on_both_backends(tmp_path: Path, capsys) -> None: + values = np.arange(12, dtype=np.float32).reshape(3, 4) + names = ['s1', 's2', 's3'] + _write_tiledb(tmp_path / 'ok.tdb', 'x', values, names) + assert read_gate._gate_tiledb(tmp_path / 'ok.tdb', 8) == 0 + h5 = _write_hdf5(tmp_path / 'ok.h5', 'x', values, names) + assert read_gate._gate_hdf5(h5, 8) == 0 + + _write_tiledb(tmp_path / 'bad.tdb', 'x', values, ['only-one']) + assert read_gate._gate_tiledb(tmp_path / 'bad.tdb', 8) == 1 + assert 'column names' in capsys.readouterr().err + + +# --- compare to reference ---------------------------------------------------- + + +def _reference(tmp_path: Path): + ref = tmp_path / 'ref.tdb' + values = np.arange(32, dtype=np.float32).reshape(4, 8) + values[1, 3] = np.nan + names = ['a', 'b', 'c', 'd'] + _write_tiledb(ref, 'area', values, names) + return ref, values, names + + +def test_compare_aligns_rows_by_name_and_accepts_subsets(tmp_path: Path) -> None: + ref, values, _names = _reference(tmp_path) + # Candidate holds rows d and b, in that order — alignment must use names. + candidate = _write_hdf5(tmp_path / 'cand.h5', 'area', values[[3, 1]], ['d', 'b']) + result = compare_reference.compare_scalar(ref, candidate, 'area', block_columns=3) + assert result['passed'] + assert result['values_equal'] + assert result['nan_equal'] + assert result['rows_compared'] == 2 + assert result['max_abs_difference'] == 0.0 + assert result['reference_complete'] is True + + +def test_compare_detects_value_change_and_misaligned_names(tmp_path: Path) -> None: + ref, values, names = _reference(tmp_path) + changed = values.copy() + changed[2, 5] += 0.5 + candidate = _write_hdf5(tmp_path / 'changed.h5', 'area', changed, names) + result = compare_reference.compare_scalar(ref, candidate, 'area') + assert not result['values_equal'] + assert result['max_abs_difference'] == pytest.approx(0.5) + assert not result['passed'] + + bad_names = _write_hdf5(tmp_path / 'names.h5', 'area', values, ['a', 'b', 'c', 'zzz']) + result = compare_reference.compare_scalar(ref, bad_names, 'area') + assert result['sources_aligned'] is False + assert result['rows_missing_in_reference'] == 1 + assert not result['passed'] + + +def test_compare_flags_partial_reference_and_lists_completeness(tmp_path: Path) -> None: + ref = tmp_path / 'partial.tdb' + values = np.arange(16, dtype=np.float32).reshape(2, 8) + # Domain has 4 rows but only 2 were written (a stalled original run). + _write_tiledb(ref, 'curv', values, ['a', 'b'], rows=4) + listing = compare_reference.reference_completeness(ref) + assert listing == [ + { + 'scalar': 'curv', + 'rows': 4, + 'columns': 8, + 'dtype': 'float32', + 'complete': False, + 'nonempty_domain': listing[0]['nonempty_domain'], + } + ] + candidate = _write_hdf5(tmp_path / 'c.h5', 'curv', values, ['a', 'b']) + result = compare_reference.compare_scalar(ref, candidate, 'curv') + assert result['values_equal'] + assert result['reference_complete'] is False + assert not result['passed'] + + +def test_compare_accepts_tiledb_candidate(tmp_path: Path) -> None: + ref, values, names = _reference(tmp_path) + cand = tmp_path / 'cand.tdb' + _write_tiledb(cand, 'area', values, names) + result = compare_reference.compare_scalar(ref, cand, 'area', block_columns=5) + assert result['passed'] + assert result['rows_compared'] == 4 + + +# --- materialize cohort ------------------------------------------------------- + + +class _FakeClient: + def __init__(self) -> None: + self.calls: list[tuple[str, str]] = [] + + def download_file(self, bucket: str, key: str, target: str) -> None: + self.calls.append((bucket, key)) + Path(target).write_bytes(b'data:' + key.encode()) + + +def test_materialize_rewrites_s3_cells_and_downloads_once(tmp_path: Path) -> None: + cohort = tmp_path / 'cohort.csv' + with cohort.open('w', newline='') as stream: + writer = csv.DictWriter(stream, fieldnames=['subject_id', 'source_mask_file', 'alff']) + writer.writeheader() + writer.writerow( + { + 'subject_id': 's1', + 'source_mask_file': 's3://bkt/p/func_mask/s1_func_mask.nii.gz', + 'alff': 's3://bkt/p/alff/s1_alff.nii.gz', + } + ) + writer.writerow( + { + 'subject_id': 's2', + 'source_mask_file': 's3://bkt/p/func_mask/s1_func_mask.nii.gz', + 'alff': 's3://bkt/p/alff/s2_alff.nii.gz', + } + ) + client = _FakeClient() + args = SimpleNamespace( + cohort=cohort, + data_dir=tmp_path / 'data', + output=tmp_path / 'local.csv', + workers=2, + extra_key=['s3://bkt/p/func_mask/group.nii.gz'], + ) + count = materialize.materialize(args, client=client) + assert count == 4 # 2 alff + 1 shared mask + 1 extra key + assert sorted(client.calls) == sorted( + [ + ('bkt', 'p/alff/s1_alff.nii.gz'), + ('bkt', 'p/alff/s2_alff.nii.gz'), + ('bkt', 'p/func_mask/s1_func_mask.nii.gz'), + ('bkt', 'p/func_mask/group.nii.gz'), + ] + ) + with (tmp_path / 'local.csv').open(newline='') as stream: + rows = list(csv.DictReader(stream)) + assert rows[0]['alff'].endswith('data/alff/s1_alff.nii.gz') + assert Path(rows[0]['alff']).is_absolute() + assert Path(rows[0]['alff']).exists() + assert rows[1]['source_mask_file'] == rows[0]['source_mask_file'] + + +def test_compare_reports_missing_reference_scalar(tmp_path: Path) -> None: + ref, values, names = _reference(tmp_path) + candidate = _write_hdf5(tmp_path / 'c.h5', 'thickness', values, names) + result = compare_reference.compare_scalar(ref, candidate, 'thickness') + assert result['reference_missing'] is True + assert result['passed'] is False + + +# --- summarize_scale ----------------------------------------------------------- + + +def test_summarize_scale_collects_every_ladder_step(tmp_path: Path) -> None: + root = tmp_path / 'results' + compare = root / 'cubic-03-area-compare' + compare.mkdir(parents=True) + cohort = tmp_path / 'area.csv' + cohort.write_text('subject_id,scalar_name,source_file\ns1,area,/a\ns2,area,/b\n') + (compare / 'comparison_meta.json').write_text( + json.dumps({'cohort': str(cohort), 'scalar_names': ['area'], 'workers': 1}) + ) + with (compare / 'comparison_results.csv').open('w', newline='') as stream: + writer = csv.DictWriter(stream, fieldnames=['implementation', 'backend', 'wall_seconds']) + writer.writeheader() + writer.writerows( + [ + {'implementation': 'baseline', 'backend': 'hdf5', 'wall_seconds': '100'}, + {'implementation': 'optimized', 'backend': 'hdf5', 'wall_seconds': '4'}, + ] + ) + sweep = root / 'cubic-03-area-sweep' + sweep.mkdir() + with (sweep / 'results.csv').open('w', newline='') as stream: + fields = [ + 'scalar_names', + 'input_file_count', + 'backend', + 'workers', + 'wall_seconds', + 'input_files_per_minute', + 'peak_memory_bytes', + 'fragment_count', + 'fast_path_rows', + ] + writer = csv.DictWriter(stream, fieldnames=fields) + writer.writeheader() + writer.writerow( + { + 'scalar_names': 'area', + 'input_file_count': '2', + 'backend': 'tiledb', + 'workers': '8', + 'wall_seconds': '3.5', + 'input_files_per_minute': '34.3', + 'peak_memory_bytes': str(2**30), + 'fragment_count': '1', + 'fast_path_rows': '1', + } + ) + shard = root / 'cubic-04-shard' / 'curv' + shard.mkdir(parents=True) + (shard / 'wall.csv').write_text('wall_seconds,hdf5,60.5\nwall_seconds,tiledb,50.0\n') + (shard / 'PASSED').write_text('ok') + prod = root / 'production' / 'checks' + prod.mkdir(parents=True) + (prod / 'wall_seconds.csv').write_text('sulc,70.25\n') + + output = tmp_path / 'scale.csv' + args = ['--results-root', str(root), '--output', str(output)] + + old_argv = sys.argv + sys.argv = ['summarize_scale.py', *args] + try: + assert summarize_scale.main() == 0 + finally: + sys.argv = old_argv + with output.open(newline='') as stream: + rows = list(csv.DictReader(stream)) + steps = [row['step'] for row in rows] + assert steps == [ + '3_full_scalar_original_vs_optimized', + '3_full_scalar_worker_sweep', + '4_scalar_sharded_array', + '4_scalar_sharded_array', + '5_production_build', + ] + assert rows[0]['speedup'] == '25.0' + assert rows[0]['subjects'] == '2' + assert rows[1]['peak_memory_gb'] == '1.0' + assert rows[2]['passed'] == 'True' + assert rows[4]['passed'] == 'False' diff --git a/test/test_h5_storage.py b/test/test_h5_storage.py index 9a03df6..8ec535f 100644 --- a/test/test_h5_storage.py +++ b/test/test_h5_storage.py @@ -2,15 +2,21 @@ from __future__ import annotations +import json +import logging +import zlib + import h5py import numpy as np import pytest +from modelarrayio.storage import h5_storage from modelarrayio.storage.h5_storage import ( compute_chunk_shape_full_subjects, create_scalar_matrix_dataset, resolve_compression, resolve_dtype, + write_full_matrix, write_rows_in_column_stripes, ) @@ -82,7 +88,7 @@ def test_compute_chunk_shape_rejects_zero_dimension(n_files: int, n_elements: in ) -def test_create_scalar_matrix_dataset_writes_data_and_attrs(tmp_path) -> None: +def test_create_scalar_matrix_dataset_writes_data_and_column_names_dataset(tmp_path) -> None: path = tmp_path / 'out.h5' data = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32) with h5py.File(path, 'w') as h5: @@ -97,7 +103,219 @@ def test_create_scalar_matrix_dataset_writes_data_and_attrs(tmp_path) -> None: with h5py.File(path, 'r') as h5: dset = h5['scalars/FA/values'] np.testing.assert_array_equal(dset[...], data) - assert list(dset.attrs['column_names']) == ['sub-01', 'sub-02'] + assert 'column_names' not in dset.attrs + assert h5['scalars/FA/column_names'].asstr()[...].tolist() == ['sub-01', 'sub-02'] + + +@pytest.mark.parametrize('dtype', [np.float32, np.float64]) +def test_parallel_direct_chunks_match_normal_hdf5_and_pad_partial_chunk(tmp_path, dtype) -> None: + path = tmp_path / 'direct.h5' + data = np.arange(21, dtype=dtype).reshape(3, 7) + + with h5py.File(path, 'w') as h5: + normal = h5.create_dataset( + 'normal', + shape=data.shape, + dtype=data.dtype, + chunks=(2, 4), + compression='gzip', + compression_opts=4, + shuffle=True, + ) + normal[...] = data + direct = h5.create_dataset( + 'direct', + shape=data.shape, + dtype=data.dtype, + chunks=(2, 4), + compression='gzip', + compression_opts=4, + shuffle=True, + ) + metrics = write_full_matrix(direct, data, compression_workers=2) + + np.testing.assert_array_equal(direct[...], normal[...]) + assert direct.compression == normal.compression == 'gzip' + assert direct.compression_opts == normal.compression_opts == 4 + assert direct.shuffle is normal.shuffle is True + assert direct.chunks == normal.chunks == (2, 4) + assert metrics['direct_chunk'] is True + assert metrics['fallback'] is False + assert metrics['chunk_count'] == 4 + assert metrics['compression_cpu_seconds'] >= 0.0 + json.dumps(metrics) + + for offset in ((0, 0), (0, 4), (2, 0), (2, 4)): + direct_mask, direct_bytes = direct.id.read_direct_chunk(offset) + normal_mask, normal_bytes = normal.id.read_direct_chunk(offset) + assert direct_mask == normal_mask == 0 + # Deflate streams are not guaranteed byte-for-byte stable across zlib + # builds. Their decoded shuffled bytes are the portable equivalence. + assert zlib.decompress(direct_bytes) == zlib.decompress(normal_bytes) + + _filter_mask, compressed = direct.id.read_direct_chunk((2, 4)) + shuffled = zlib.decompress(compressed) + unshuffled = ( + np.frombuffer(shuffled, dtype=np.uint8).reshape(data.dtype.itemsize, -1).T.reshape(-1) + ) + padded_chunk = unshuffled.view(data.dtype).reshape(2, 4) + np.testing.assert_array_equal(padded_chunk[0, :3], data[2, 4:]) + np.testing.assert_array_equal(padded_chunk[0, 3], np.zeros((), dtype=data.dtype)) + np.testing.assert_array_equal(padded_chunk[1], np.zeros(4, dtype=data.dtype)) + + +def test_no_compression_bypasses_direct_chunk_path(tmp_path, monkeypatch) -> None: + path = tmp_path / 'uncompressed.h5' + data = np.arange(10, dtype=np.float32).reshape(2, 5) + + def _unexpected_direct_write(*_args, **_kwargs): + raise AssertionError('direct chunk path should not be called') + + monkeypatch.setattr(h5_storage, '_write_direct_gzip_chunks', _unexpected_direct_write) + with h5py.File(path, 'w') as h5: + dataset = h5.create_dataset('values', shape=data.shape, dtype=data.dtype, chunks=(2, 3)) + metrics = write_full_matrix(dataset, data) + np.testing.assert_array_equal(dataset[...], data) + + assert metrics['direct_chunk'] is False + assert metrics['fallback'] is False + assert metrics['compression_seconds'] == 0.0 + + +def test_internal_compression_worker_environment_controls_pool(tmp_path, monkeypatch) -> None: + path = tmp_path / 'workers.h5' + data = np.arange(64, dtype=np.float32).reshape(8, 8) + monkeypatch.setenv('MODELARRAYIO_HDF5_COMPRESSION_WORKERS', '2') + + with h5py.File(path, 'w') as h5: + dataset = h5.create_dataset( + 'values', + shape=data.shape, + dtype=data.dtype, + chunks=(2, 2), + compression='gzip', + compression_opts=1, + shuffle=True, + ) + metrics = write_full_matrix(dataset, data) + + assert metrics['compression_workers'] == 2 + with h5py.File(path, 'r') as h5: + np.testing.assert_array_equal(data, h5['values'][...]) + + +@pytest.mark.parametrize('configured_workers', ['invalid', '0', '-1']) +def test_invalid_compression_worker_environment_is_not_silently_ignored( + tmp_path, monkeypatch, configured_workers +) -> None: + path = tmp_path / 'invalid-workers.h5' + data = np.arange(4, dtype=np.float32).reshape(2, 2) + monkeypatch.setenv('MODELARRAYIO_HDF5_COMPRESSION_WORKERS', configured_workers) + + with h5py.File(path, 'w') as h5: + dataset = h5.create_dataset( + 'values', + shape=data.shape, + dtype=data.dtype, + chunks=(1, 2), + compression='gzip', + ) + with pytest.raises(ValueError, match='positive integer'): + write_full_matrix(dataset, data) + + +def test_default_compression_workers_respect_available_cpu_count(tmp_path, monkeypatch) -> None: + path = tmp_path / 'available-cpus.h5' + data = np.arange(16, dtype=np.float32).reshape(4, 4) + monkeypatch.delenv('MODELARRAYIO_HDF5_COMPRESSION_WORKERS', raising=False) + monkeypatch.setattr(h5_storage, '_available_cpu_count', lambda: 2) + + with h5py.File(path, 'w') as h5: + dataset = h5.create_dataset( + 'values', + shape=data.shape, + dtype=data.dtype, + chunks=(1, 2), + compression='gzip', + ) + metrics = write_full_matrix(dataset, data) + + assert metrics['compression_workers'] == 2 + + +def test_column_names_are_written_beside_nested_values_dataset(tmp_path) -> None: + path = tmp_path / 'nested.h5' + data = np.arange(6, dtype=np.float32).reshape(2, 3) + + with h5py.File(path, 'w') as h5: + create_scalar_matrix_dataset( + h5, + 'scalars/group/FA/values', + data, + sources_list=['sub-01', 'sub-02'], + compression=None, + ) + + with h5py.File(path, 'r') as h5: + assert h5['scalars/group/FA/column_names'].asstr()[...].tolist() == [ + 'sub-01', + 'sub-02', + ] + + +def test_direct_chunk_worker_failure_logs_and_falls_back(tmp_path, monkeypatch, caplog) -> None: + path = tmp_path / 'fallback.h5' + data = np.arange(10, dtype=np.float64).reshape(2, 5) + + def _failed_worker(*_args, **_kwargs): + raise RuntimeError('injected compression failure') + + monkeypatch.setattr(h5_storage, '_encode_gzip_chunk', _failed_worker) + with h5py.File(path, 'w') as h5, caplog.at_level(logging.WARNING): + dataset = h5.create_dataset( + 'values', + shape=data.shape, + dtype=data.dtype, + chunks=(2, 3), + compression='gzip', + compression_opts=1, + shuffle=True, + ) + metrics = write_full_matrix(dataset, data, compression_workers=2) + np.testing.assert_array_equal(dataset[...], data) + + assert metrics['direct_chunk'] is False + assert metrics['fallback'] is True + assert metrics['compression_seconds'] is None + assert metrics['fallback_reason'] == 'injected compression failure' + assert any('Falling back to ordinary HDF5' in record.message for record in caplog.records) + + +def test_unsupported_direct_filter_logs_and_falls_back(tmp_path, caplog) -> None: + path = tmp_path / 'lzf.h5' + data = np.arange(10, dtype=np.float32).reshape(2, 5) + + with h5py.File(path, 'w') as h5, caplog.at_level(logging.WARNING): + dataset = h5.create_dataset( + 'values', + shape=data.shape, + dtype=data.dtype, + chunks=(2, 3), + compression='lzf', + shuffle=True, + ) + metrics = write_full_matrix(dataset, data) + np.testing.assert_array_equal(dataset[...], data) + + assert metrics['direct_chunk'] is False + assert metrics['fallback'] is True + assert "compression 'lzf' is not supported" in metrics['fallback_reason'] + assert any('Falling back to ordinary HDF5' in record.message for record in caplog.records) + + +def test_full_matrix_coercion_preserves_staged_array() -> None: + matrix = np.empty((2, 3), dtype=np.float32, order='C') + assert h5_storage._coerce_full_matrix(matrix, dtype=np.float32) is matrix def test_write_rows_in_column_stripes_matches_dense_write(tmp_path) -> None: diff --git a/test/test_loading.py b/test/test_loading.py new file mode 100644 index 0000000..6e49a0b --- /dev/null +++ b/test/test_loading.py @@ -0,0 +1,152 @@ +"""Tests for bounded scalar staging.""" + +from __future__ import annotations + +import threading +import time + +import numpy as np +import pytest + +from modelarrayio.utils.loading import LoadedRow, SourceLoadError, stage_source_rows + + +def test_stage_source_rows_is_bounded_ordered_and_typed() -> None: + lock = threading.Lock() + active = 0 + maximum = 0 + + def load(source: int) -> np.ndarray: + nonlocal active, maximum + with lock: + active += 1 + maximum = max(maximum, active) + try: + # Force completion in reverse order. + time.sleep((5 - source) * 0.002) + return np.array([source + 0.25, source + 0.5], dtype=np.float64) + finally: + with lock: + active -= 1 + + matrix, stats = stage_source_rows( + list(range(6)), + load, + scalar_name='FA', + storage_dtype='float32', + workers=2, + n_elements=2, + ) + + assert maximum <= 2 + assert matrix.flags.c_contiguous + assert matrix.dtype == np.float32 + np.testing.assert_array_equal(matrix[:, 0], np.arange(6, dtype=np.float32) + 0.25) + assert 1 <= stats.in_flight_sources_max <= 2 + assert stats.in_flight_sources_mean >= 1 + assert stats.queue_occupancy_max == 0 + assert stats.source_wall_seconds > 0 + assert stats.source_count == 6 + + +def test_stage_source_rows_reuses_first_row_and_records_validation() -> None: + calls: list[str] = [] + + def load(source: str): + calls.append(source) + return LoadedRow(np.array([2.0, 3.0]), validation_seconds=0.125) + + matrix, stats = stage_source_rows( + ['first', 'second'], + load, + scalar_name='MD', + first_row=LoadedRow( + np.array([0.0, 1.0]), + validation_seconds=0.25, + source_seconds=0.5, + cpu_seconds=0.125, + ), + storage_dtype='float64', + workers=1, + ) + + assert calls == ['second'] + np.testing.assert_array_equal(matrix, [[0.0, 1.0], [2.0, 3.0]]) + assert stats.validation_seconds == pytest.approx(0.375) + assert stats.source_wall_seconds >= 0.5 + assert stats.source_seconds >= 0.5 + assert stats.wall_seconds >= 0.5 + assert stats.cpu_seconds >= 0.125 + + +@pytest.mark.parametrize('workers', [0, -1, True, 2.0, '2', 'bad']) +def test_stage_source_rows_rejects_invalid_workers(workers) -> None: + with pytest.raises(ValueError, match='positive integer'): + stage_source_rows( + ['source'], + lambda _source: np.ones(1), + scalar_name='FA', + workers=workers, + n_elements=1, + ) + + +def test_stage_source_rows_adds_scalar_and_source_to_errors() -> None: + def fail(_source): + raise OSError('unavailable') + + with pytest.raises(SourceLoadError, match=r"source 'broken\.nii\.gz'.*scalar 'ALFF'"): + stage_source_rows( + ['broken.nii.gz'], + fail, + scalar_name='ALFF', + workers=1, + n_elements=3, + ) + + +def test_stage_source_rows_rejects_wrong_row_shape() -> None: + with pytest.raises(SourceLoadError, match=r'expected \(3,\)'): + stage_source_rows( + ['source'], + lambda _source: np.ones(2), + scalar_name='FA', + workers=1, + n_elements=3, + ) + + +def test_stage_source_rows_logs_progress_while_loading(caplog) -> None: + import logging + + def load(source: int) -> np.ndarray: + return np.array([float(source)], dtype=np.float32) + + with caplog.at_level(logging.INFO, logger='modelarrayio.utils.loading'): + matrix, stats = stage_source_rows( + list(range(6)), + load, + scalar_name='FA', + workers=1, + progress_interval_seconds=0.0000001, + ) + np.testing.assert_array_equal(matrix[:, 0], np.arange(6, dtype=np.float32)) + progress = [record for record in caplog.records if 'Staging FA:' in record.getMessage()] + assert progress, 'expected at least one progress line' + # The last source is reported by the completion record, not a progress line. + assert all('6/6' not in record.getMessage() for record in progress) + assert stats.source_count == 6 + + +def test_stage_source_rows_progress_can_be_disabled(caplog) -> None: + import logging + + with caplog.at_level(logging.INFO, logger='modelarrayio.utils.loading'): + stage_source_rows( + [0, 1, 2], + lambda source: np.array([1.0], dtype=np.float32), + scalar_name='FA', + workers=2, + progress_interval_seconds=None, + ) + assert not [record for record in caplog.records if 'Staging FA:' in record.getMessage()] diff --git a/test/test_make_cifti_cohort.py b/test/test_make_cifti_cohort.py new file mode 100644 index 0000000..ae22a94 --- /dev/null +++ b/test/test_make_cifti_cohort.py @@ -0,0 +1,291 @@ +"""Tests for the CIFTI benchmark cohort generator.""" + +from __future__ import annotations + +import csv +from pathlib import Path + +import nibabel as nb +import numpy as np +import pytest + +import harness_common +import make_cifti_cohort as generator + +_ONE_THICKNESS_ROW = { + 'subject_id': 'sub-a', + 'scalar_name': 'thickness', + 'source_file': '/data/a.dscalar.nii', +} + + +def _read(path: Path) -> list[dict[str, str]]: + with path.open(newline='', encoding='utf-8') as stream: + return list(csv.DictReader(stream)) + + +def test_synthesize_writes_real_dscalars_and_long_cohort(tmp_path) -> None: + args = generator._parser().parse_args( + [ + 'synthesize', + '--output-dir', + str(tmp_path / 'data'), + '--subjects', + '3', + '--scalars', + 'thickness', + 'area', + '--grayordinates', + '400', + '--seed', + '7', + ] + ) + row_count = args.func(args) + assert row_count == 6 # 3 subjects x 2 scalars + + rows = _read(tmp_path / 'data' / 'cohort.csv') + assert list(rows[0].keys()) == ['subject_id', 'scalar_name', 'source_file'] + assert {row['scalar_name'] for row in rows} == {'thickness', 'area'} + + # Every source is a real dscalar with the requested greyordinate count. + for row in rows: + image = nb.load(row['source_file']) + assert isinstance(image, nb.Cifti2Image) + assert image.shape == (1, 400) + + +def test_synthesize_is_deterministic_for_a_seed(tmp_path) -> None: + def _run(target: Path) -> np.ndarray: + args = generator._parser().parse_args( + [ + 'synthesize', + '--output-dir', + str(target), + '--subjects', + '1', + '--scalars', + 'thickness', + '--grayordinates', + '256', + '--seed', + '11', + ] + ) + args.func(args) + source = _read(target / 'cohort.csv')[0]['source_file'] + return nb.load(source).get_fdata(dtype=np.float32).reshape(-1) + + np.testing.assert_array_equal(_run(tmp_path / 'a'), _run(tmp_path / 'b')) + + +def test_synthesize_varies_offset_per_scalar(tmp_path) -> None: + # Different map-name lengths must shift vox_offset, matching real per-metric + # dscalars and exercising the per-scalar raw-read plan. + args = generator._parser().parse_args( + [ + 'synthesize', + '--output-dir', + str(tmp_path / 'data'), + '--subjects', + '1', + '--scalars', + 'curv', + 'desc-fwhm25_thickness', + '--grayordinates', + '256', + ] + ) + args.func(args) + rows = _read(tmp_path / 'data' / 'cohort.csv') + offsets = {row['scalar_name']: nb.load(row['source_file']).dataobj.offset for row in rows} + assert offsets['curv'] != offsets['desc-fwhm25_thickness'] + + +def test_synthesize_default_has_medial_wall_zeros_and_realistic_range(tmp_path) -> None: + # The no-reference default must carry a constant zero block (medial wall) + # and scalar-appropriate value ranges rather than raw standard-normal noise. + args = generator._parser().parse_args( + [ + 'synthesize', + '--output-dir', + str(tmp_path / 'data'), + '--subjects', + '2', + '--scalars', + 'thickness', + '--grayordinates', + '1000', + '--seed', + '3', + ] + ) + args.func(args) + rows = _read(tmp_path / 'data' / 'cohort.csv') + values = nb.load(rows[0]['source_file']).get_fdata(dtype=np.float32).reshape(-1) + assert float(values.min()) >= 0.0 # thickness is non-negative + assert float(values.max()) <= 5.0 + assert np.mean(values == 0.0) > 0.05 # a real medial-wall-sized zero block + + +def _write_reference_dscalar(path: Path, values: np.ndarray, scalar: str) -> None: + from nibabel.cifti2.cifti2_axes import ScalarAxis + + n_elements = values.shape[-1] + brain_axis = generator._build_brain_model_axis(n_elements) + header = nb.cifti2.Cifti2Header.from_axes((ScalarAxis([scalar]), brain_axis)) + path.parent.mkdir(parents=True, exist_ok=True) + nb.Cifti2Image(values.reshape(1, n_elements).astype(np.float32), header=header).to_filename( + str(path) + ) + + +def test_synthesize_reference_mode_matches_stats_without_copying_rows(tmp_path) -> None: + rng = np.random.default_rng(0) + n_elements = 800 + zero_mask = np.zeros(n_elements, dtype=bool) + zero_mask[:100] = True # a known medial-wall block shared across subjects + reference_dir = tmp_path / 'reference' + reference_rows = [] + for index in range(3): + row = rng.normal(2.2, 1.0, size=n_elements).astype(np.float32) + np.clip(row, 0.0, 5.0, out=row) + row[zero_mask] = 0.0 + reference_rows.append(row) + _write_reference_dscalar( + reference_dir / f'sub-real{index}' / f'sub-real{index}_thickness.dscalar.nii', + row, + 'thickness', + ) + + args = generator._parser().parse_args( + [ + 'synthesize', + '--output-dir', + str(tmp_path / 'out'), + '--subjects', + '5', + '--scalars', + 'thickness', + '--reference-dir', + str(reference_dir), + '--seed', + '1', + ] + ) + args.func(args) + + rows = _read(tmp_path / 'out' / 'cohort.csv') + assert len(rows) == 5 + generated = np.vstack( + [nb.load(row['source_file']).get_fdata(dtype=np.float32).reshape(-1) for row in rows] + ) + # The medial-wall block is preserved exactly. + assert np.all(generated[:, zero_mask] == 0.0) + # Values stay within the observed reference range. + assert float(generated.min()) >= 0.0 + assert float(generated.max()) <= 5.0 + # Fresh noise is drawn: no generated row equals any reference row. + reference = np.vstack(reference_rows) + for generated_row in generated: + assert not any(np.array_equal(generated_row, real_row) for real_row in reference) + + +def test_subset_selects_first_subjects_and_requested_scalars(tmp_path) -> None: + source = tmp_path / 'full.csv' + harness_common.write_csv( + source, + [ + { + 'subject_id': subject, + 'session_id': 'ses-1', + 'scalar_name': scalar, + 'source_file': f'/data/{subject}_{scalar}.dscalar.nii', + } + for subject in ('sub-a', 'sub-b', 'sub-c') + for scalar in ('thickness', 'area', 'sulc') + ], + ) + + output = tmp_path / 'subset.csv' + args = generator._parser().parse_args( + [ + 'subset', + '--input', + str(source), + '--output', + str(output), + '--subjects', + '2', + '--scalars', + 'thickness', + 'sulc', + ] + ) + row_count = args.func(args) + + rows = _read(output) + assert row_count == 4 # 2 subjects x 2 scalars + assert {row['subject_id'] for row in rows} == {'sub-a', 'sub-b'} + assert {row['scalar_name'] for row in rows} == {'thickness', 'sulc'} + + +def test_subset_rejects_more_subjects_than_available(tmp_path) -> None: + source = tmp_path / 'full.csv' + harness_common.write_csv(source, [_ONE_THICKNESS_ROW]) + + args = generator._parser().parse_args( + [ + 'subset', + '--input', + str(source), + '--output', + str(tmp_path / 'out.csv'), + '--subjects', + '5', + ] + ) + with pytest.raises(ValueError, match='only 1'): + args.func(args) + + +def test_subset_rejects_unknown_scalar(tmp_path) -> None: + source = tmp_path / 'full.csv' + harness_common.write_csv(source, [_ONE_THICKNESS_ROW]) + + args = generator._parser().parse_args( + [ + 'subset', + '--input', + str(source), + '--output', + str(tmp_path / 'out.csv'), + '--subjects', + '1', + '--scalars', + 'nope', + ] + ) + with pytest.raises(ValueError, match='not present'): + args.func(args) + + +def test_subset_rejects_wide_format_cohort(tmp_path) -> None: + source = tmp_path / 'wide.csv' + harness_common.write_csv( + source, [{'subject_id': 'sub-a', 'alff': 's3://bucket/a_alff.nii.gz'}] + ) + + args = generator._parser().parse_args( + [ + 'subset', + '--input', + str(source), + '--output', + str(tmp_path / 'out.csv'), + '--subjects', + '1', + ] + ) + with pytest.raises(ValueError, match='long-format'): + args.func(args) diff --git a/test/test_mif_cli.py b/test/test_mif_cli.py index d2b3735..0e0fec1 100644 --- a/test/test_mif_cli.py +++ b/test/test_mif_cli.py @@ -33,7 +33,7 @@ def test_mif_to_h5_scalar_columns_writes_prefixed_outputs(tmp_path, monkeypatch) def fake_gather_fixels(_index_file, _directions_file): return fixel_table, voxel_table - def fake_load_cohort_mif(cohort_long, _s3_workers): + def fake_load_cohort_mif(cohort_long, **_kwargs): scalars = {} sources = {} for scalar_name, group in cohort_long.groupby('scalar_name'): diff --git a/test/test_mif_helpers.py b/test/test_mif_helpers.py index 2f70ba5..94c1149 100644 --- a/test/test_mif_helpers.py +++ b/test/test_mif_helpers.py @@ -23,6 +23,7 @@ def _fake_from_filename(path): assert path == str(tmp_path / 'input.mif') return _FakeImage() + (tmp_path / 'input.mif').touch() monkeypatch.setattr(mif.MifImage, 'from_filename', staticmethod(_fake_from_filename)) img, data = mif.mif_to_image(tmp_path / 'input.mif') assert isinstance(img, _FakeImage) @@ -37,7 +38,7 @@ def test_load_cohort_mif_sequential_and_parallel(monkeypatch) -> None: } ) - def _fake_mif_to_image(path): + def _fake_mif_to_image(path, dtype=None): mapping = { 'a': np.array([1.0, 2.0], dtype=np.float32), 'b': np.array([3.0, 4.0], dtype=np.float32), @@ -70,7 +71,7 @@ def test_gather_fixels_builds_tables(monkeypatch) -> None: dtype=np.float32, ) - def _fake_mif_to_image(path): + def _fake_mif_to_image(path, dtype=np.float32): if str(path).endswith('index.mif'): return object(), index_data return object(), directions @@ -89,7 +90,7 @@ def test_gather_fixels_raises_when_terminal_count_missing(monkeypatch) -> None: index_data[..., 1] = np.nan directions = np.zeros((1, 3), dtype=np.float32) - def _fake_mif_to_image(path): + def _fake_mif_to_image(path, dtype=np.float32): return object(), index_data if str(path).endswith('index.mif') else directions monkeypatch.setattr(mif, 'mif_to_image', _fake_mif_to_image) diff --git a/test/test_mif_to_h5_unit.py b/test/test_mif_to_h5_unit.py index 23b9784..06851f5 100644 --- a/test/test_mif_to_h5_unit.py +++ b/test/test_mif_to_h5_unit.py @@ -7,6 +7,7 @@ import numpy as np import pandas as pd import pytest +import tiledb from modelarrayio.cli import mif_to_h5 @@ -26,7 +27,7 @@ def test_mif_to_h5_raises_when_sources_missing(monkeypatch, tmp_path: Path) -> N 'gather_fixels', lambda index_file, directions_file: (pd.DataFrame(), pd.DataFrame()), ) - monkeypatch.setattr(mif_to_h5, 'load_cohort_mif', lambda cohort_long, s3_workers: ({}, {})) + monkeypatch.setattr(mif_to_h5, 'load_cohort_mif', lambda *_args, **_kwargs: ({}, {})) with pytest.raises(ValueError, match='Unable to derive scalar sources'): mif_to_h5.mif_to_h5('index.mif', 'directions.mif', _cohort(), output=tmp_path / 'out.h5') @@ -67,18 +68,10 @@ def test_mif_to_h5_tiledb_parallel_and_split(monkeypatch, tmp_path: Path) -> Non 'MD': [np.array([2.0], dtype=np.float32)], } sources = {'FA': ['fa.mif'], 'MD': ['md.mif']} - calls = [] - monkeypatch.setattr( mif_to_h5, 'gather_fixels', lambda *_args, **_kwargs: (fixel_table, voxel_table) ) monkeypatch.setattr(mif_to_h5, 'load_cohort_mif', lambda *_args, **_kwargs: (scalars, sources)) - monkeypatch.setattr( - mif_to_h5.cli_utils, - 'write_tiledb_scalar_matrices', - lambda output, scalars, sources, **kwargs: calls.append((Path(output), scalars, sources)), - ) - status = mif_to_h5.mif_to_h5( 'index.mif', 'directions.mif', @@ -89,4 +82,9 @@ def test_mif_to_h5_tiledb_parallel_and_split(monkeypatch, tmp_path: Path) -> Non split_outputs=True, ) assert status == 0 - assert sorted(path.name for path, _, _ in calls) == ['FA_fixels.tdb', 'MD_fixels.tdb'] + for scalar_name, scalar_values in scalars.items(): + scalar_output = tmp_path / f'{scalar_name}_fixels.tdb' + values_uri = scalar_output / 'scalars' / scalar_name / 'values' + assert tiledb.object_type(str(values_uri)) == 'array' + with tiledb.open(str(values_uri), 'r') as array: + np.testing.assert_array_equal(array[:]['values'], np.asarray(scalar_values)) diff --git a/test/test_nifti_s3.py b/test/test_nifti_s3.py index 08a2e78..025165f 100644 --- a/test/test_nifti_s3.py +++ b/test/test_nifti_s3.py @@ -170,9 +170,5 @@ def test_nifti_s3_serial_matches_parallel(tmp_path, group_mask_path, monkeypatch serial_data = s['scalars/alff/values'][...] parallel_data = p['scalars/alff/values'][...] - # Row order in the parallel result may differ from cohort order, so sort both - # by their row fingerprint before comparing assert serial_data.shape == parallel_data.shape - serial_sorted = serial_data[np.lexsort(serial_data.T)] - parallel_sorted = parallel_data[np.lexsort(parallel_data.T)] - np.testing.assert_array_equal(serial_sorted, parallel_sorted) + np.testing.assert_array_equal(serial_data, parallel_data) diff --git a/test/test_nifti_utils.py b/test/test_nifti_utils.py index 7ccbb05..a5227b5 100644 --- a/test/test_nifti_utils.py +++ b/test/test_nifti_utils.py @@ -2,11 +2,32 @@ from __future__ import annotations +import threading +import time + import nibabel as nb import numpy as np +import pandas as pd import pytest -from modelarrayio.utils.nifti import flattened_image +from modelarrayio.utils import nifti +from modelarrayio.utils.loading import SourceLoadError +from modelarrayio.utils.nifti import ReducedMaskCache, flattened_image, load_cohort_voxels + + +class _TrackingImage: + def __init__(self, data: np.ndarray, *, return_shared: bool = False) -> None: + self.dataobj = data + self._data = data + self._return_shared = return_shared + self.get_fdata_calls: list[tuple[np.dtype, str]] = [] + + def get_fdata(self, *, dtype, caching): + resolved = np.dtype(dtype) + self.get_fdata_calls.append((resolved, caching)) + if self._return_shared and self._data.dtype == resolved: + return self._data + return np.asarray(self._data, dtype=resolved) def test_flattened_image_extracts_group_masked_values() -> None: @@ -25,6 +46,7 @@ def test_flattened_image_extracts_group_masked_values() -> None: flat = flattened_image(scalar_img, mask_img, group_mask) assert flat.shape == (2,) + assert flat.dtype == np.float32 assert np.allclose(flat, np.array([1.5, 2.5], dtype=np.float32)) @@ -46,3 +68,248 @@ def test_flattened_image_nan_outside_individual_mask() -> None: assert flat.shape == (2,) assert np.isnan(flat[1]) assert flat[0] == pytest.approx(3.0) + + +def test_flattened_image_requests_dtype_without_caching_and_preserves_precision() -> None: + precise = np.array([[[1.0 + 2.0**-40]]], dtype=np.float64) + scalar_img = _TrackingImage(precise) + mask_img = _TrackingImage(np.ones((1, 1, 1), dtype=np.uint8)) + + flat = flattened_image( + scalar_img, + mask_img, + np.ones((1, 1, 1), dtype=bool), + storage_dtype='float64', + ) + + assert flat.dtype == np.float64 + assert flat[0] == precise.item() + assert scalar_img.get_fdata_calls == [(np.dtype(np.float64), 'unchanged')] + assert mask_img.get_fdata_calls == [] + + +def test_load_cohort_voxels_is_lazy_ordered_typed_and_skips_group_mask( + monkeypatch, tmp_path +) -> None: + group_mask = np.array([[[True, False, True]]]) + group_mask_path = tmp_path / 'group.nii.gz' + equivalent_group_path = tmp_path / 'unused' / '..' / group_mask_path.name + images = { + 'late.nii.gz': _TrackingImage(np.array([[[30.0, 0.0, 31.0]]])), + 'early.nii.gz': _TrackingImage(np.array([[[10.0, 0.0, 11.0]]])), + 'other.nii.gz': _TrackingImage(np.array([[[20.0, 0.0, 21.0]]])), + } + loaded_paths: list[str] = [] + lock = threading.Lock() + + def fake_load(path): + path = str(path) + with lock: + loaded_paths.append(path) + if path == 'late.nii.gz': + time.sleep(0.02) + return images[path] + + monkeypatch.setattr(nifti, 'load_nibabel', fake_load) + cohort = pd.DataFrame( + [ + { + 'scalar_name': 'fa', + 'source_file': 'late.nii.gz', + 'source_mask_file': str(equivalent_group_path), + }, + { + 'scalar_name': 'fa', + 'source_file': 'early.nii.gz', + 'source_mask_file': str(group_mask_path), + }, + { + 'scalar_name': 'md', + 'source_file': 'other.nii.gz', + 'source_mask_file': str(group_mask_path), + }, + ] + ) + + scalars, sources = load_cohort_voxels( + cohort, + group_mask, + workers=2, + storage_dtype='float64', + group_mask_path=group_mask_path, + ) + + assert loaded_paths == [] + assert list(scalars) == ['fa', 'md'] + assert sources == { + 'fa': ['late.nii.gz', 'early.nii.gz'], + 'md': ['other.nii.gz'], + } + + fa = scalars['fa'] + assert sorted(loaded_paths) == ['early.nii.gz', 'late.nii.gz'] + assert np.array_equal(fa, [[30.0, 31.0], [10.0, 11.0]]) + assert fa.dtype == np.float64 + assert fa.flags.c_contiguous + assert all( + image.get_fdata_calls == [(np.dtype(np.float64), 'unchanged')] + for name, image in images.items() + if name != 'other.nii.gz' + ) + + +def test_mask_cache_deduplicates_concurrent_loads_and_preserves_order( + monkeypatch, tmp_path +) -> None: + group_mask = np.ones((1, 1, 2), dtype=bool) + mask_path = str(tmp_path / 'individual-mask.nii.gz') + images = { + f'sub-{index}.nii.gz': _TrackingImage( + np.array([[[float(index), float(index) + 0.5]]], dtype=np.float32) + ) + for index in range(6) + } + images[mask_path] = _TrackingImage(np.array([[[1, 0]]], dtype=np.uint8)) + mask_load_count = 0 + lock = threading.Lock() + + def fake_load(path): + nonlocal mask_load_count + path = str(path) + if path == mask_path: + with lock: + mask_load_count += 1 + time.sleep(0.02) + return images[path] + + monkeypatch.setattr(nifti, 'load_nibabel', fake_load) + cohort = pd.DataFrame( + [ + { + 'scalar_name': 'fa', + 'source_file': f'sub-{index}.nii.gz', + 'source_mask_file': mask_path, + } + for index in range(6) + ] + ) + + scalars, _ = load_cohort_voxels(cohort, group_mask, workers=4) + matrix = scalars['fa'] + + assert mask_load_count == 1 + assert np.array_equal(matrix[:, 0], np.arange(6, dtype=np.float32)) + assert np.isnan(matrix[:, 1]).all() + assert len(scalars.mask_cache) == 1 + + +def test_repeated_scalar_decode_is_never_mutated_by_subject_mask(monkeypatch, tmp_path) -> None: + group_mask = np.ones((1, 1, 2), dtype=bool) + group_path = str(tmp_path / 'group.nii.gz') + individual_path = str(tmp_path / 'individual.nii.gz') + shared_data = np.array([[[1.0, 2.0]]], dtype=np.float32) + scalar_image = _TrackingImage(shared_data, return_shared=True) + images = { + 'shared.nii.gz': scalar_image, + individual_path: _TrackingImage(np.array([[[1, 0]]], dtype=np.uint8)), + } + monkeypatch.setattr(nifti, 'load_nibabel', lambda path: images[str(path)]) + cohort = pd.DataFrame( + [ + { + 'scalar_name': 'fa', + 'source_file': 'shared.nii.gz', + 'source_mask_file': individual_path, + }, + { + 'scalar_name': 'fa', + 'source_file': 'shared.nii.gz', + 'source_mask_file': group_path, + }, + ] + ) + + scalars, _ = load_cohort_voxels( + cohort, + group_mask, + workers=1, + group_mask_path=group_path, + ) + matrix = scalars['fa'] + + assert np.isnan(matrix[0, 1]) + assert np.array_equal(matrix[1], [1.0, 2.0]) + assert np.array_equal(shared_data, [[[1.0, 2.0]]]) + + +def test_reduced_mask_cache_is_bounded_and_protects_cached_arrays(tmp_path) -> None: + cache = ReducedMaskCache(max_size=2) + loads = 0 + + def load(value): + nonlocal loads + loads += 1 + return np.array([value], dtype=bool) + + first = cache.get_or_load(tmp_path / 'a.nii.gz', lambda: load(True)) + cache.get_or_load(tmp_path / 'b.nii.gz', lambda: load(False)) + cache.get_or_load(tmp_path / 'c.nii.gz', lambda: load(True)) + + assert len(cache) == 2 + assert not first.flags.writeable + with pytest.raises(ValueError, match='read-only'): + first[0] = False + + cache.get_or_load(tmp_path / 'a.nii.gz', lambda: load(True)) + assert loads == 4 + + +def test_load_cohort_voxels_adds_scalar_and_source_context(monkeypatch, tmp_path) -> None: + bad_source = str(tmp_path / 'bad.nii.gz') + group_path = str(tmp_path / 'group.nii.gz') + monkeypatch.setattr( + nifti, + 'load_nibabel', + lambda _path: _TrackingImage(np.ones((2, 2, 2), dtype=np.float32)), + ) + cohort = pd.DataFrame( + [ + { + 'scalar_name': 'fa', + 'source_file': bad_source, + 'source_mask_file': group_path, + } + ] + ) + scalars, _ = load_cohort_voxels( + cohort, + np.ones((1, 1, 1), dtype=bool), + workers=1, + group_mask_path=group_path, + ) + + with pytest.raises(SourceLoadError) as exc_info: + scalars['fa'] + message = str(exc_info.value) + assert bad_source in message + assert "scalar 'fa'" in message + assert 'matching shapes' in message + + +def test_load_cohort_voxels_legacy_worker_argument_and_conflict() -> None: + empty_cohort = pd.DataFrame(columns=['scalar_name', 'source_file', 'source_mask_file']) + scalars, sources = load_cohort_voxels( + empty_cohort, + np.ones((1, 1, 1), dtype=bool), + 2, + ) + assert len(scalars) == 0 + assert sources == {} + + with pytest.raises(ValueError, match='Conflicting loader worker counts'): + load_cohort_voxels( + empty_cohort, + np.ones((1, 1, 1), dtype=bool), + 2, + workers=3, + ) diff --git a/test/test_parser_utils.py b/test/test_parser_utils.py index 957d867..eb82e82 100644 --- a/test/test_parser_utils.py +++ b/test/test_parser_utils.py @@ -30,13 +30,14 @@ def test_parse_to_modelarray_minimal_defaults(tmp_path): assert args.output == Path('modelarray.h5') assert args.backend == 'hdf5' assert args.storage_dtype == 'float32' - assert args.compression == 'gzip' - assert args.compression_level == 4 + assert args.compression is None + assert args.compression_level is None assert args.shuffle is True assert args.chunk_voxels == 0 assert args.target_chunk_mb == 2.0 - assert args.workers == 1 - assert args.s3_workers == 1 + assert args.workers is None + assert args.s3_workers is None + assert args.timing_jsonl is None assert args.scalar_columns is None assert args.split_outputs is None assert args.group_mask_file is None @@ -113,6 +114,36 @@ def test_parse_to_modelarray_optional_paths_and_flags(tmp_path): assert args.log_level == 'WARNING' +def test_parse_to_modelarray_accepts_s3_auxiliary_inputs(tmp_path): + cohort = tmp_path / 'cohort.csv' + cohort.touch() + parser = _parse_to_modelarray() + args = parser.parse_args( + [ + '--cohort-file', + str(cohort), + '--mask', + 's3://bucket/group-mask.nii.gz', + '--index-file', + 's3://bucket/index.mif', + '--directions-file', + 's3://bucket/directions.mif', + ] + ) + + assert args.group_mask_file == 's3://bucket/group-mask.nii.gz' + assert args.index_file == 's3://bucket/index.mif' + assert args.directions_file == 's3://bucket/directions.mif' + + +def test_parse_to_modelarray_rejects_invalid_s3_auxiliary_uri(tmp_path): + cohort = tmp_path / 'cohort.csv' + cohort.touch() + parser = _parse_to_modelarray() + with pytest.raises(SystemExit): + parser.parse_args(['--cohort-file', str(cohort), '--mask', 's3://bucket']) + + def test_parse_to_modelarray_target_chunk_mb_branch(tmp_path): cohort = tmp_path / 'cohort.csv' cohort.touch() diff --git a/test/test_runtime.py b/test/test_runtime.py new file mode 100644 index 0000000..6db6829 --- /dev/null +++ b/test/test_runtime.py @@ -0,0 +1,362 @@ +"""Tests for conversion runtime policy and transactional publication.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import h5py +import numpy as np +import pytest +import tiledb + +from modelarrayio.cli import runtime +from modelarrayio.cli.runtime import ( + TimingRecorder, + atomic_output, + load_timing_metrics, + path_stats, + resolve_loader_workers, + scalar_output_stats, + validate_hdf5_output, + validate_tiledb_output, + validate_timing_jsonl_path, +) +from modelarrayio.storage import tiledb_storage +from modelarrayio.storage.tiledb_storage import create_scalar_matrix_array +from modelarrayio.utils.loading import LoadStats + + +def test_resolve_loader_workers_defaults_and_deprecated_alias() -> None: + assert resolve_loader_workers(None, None) == 1 + assert resolve_loader_workers(3, None) == 3 + + with pytest.warns(DeprecationWarning, match='--s3-workers is deprecated'): + assert resolve_loader_workers(None, 4) == 4 + with pytest.warns(DeprecationWarning, match='--s3-workers is deprecated'): + assert resolve_loader_workers(4, 4) == 4 + + +def test_resolve_loader_workers_rejects_conflicts_and_nonpositive_values() -> None: + with ( + pytest.warns(DeprecationWarning, match='--s3-workers is deprecated'), + pytest.raises(ValueError, match='Conflicting loader worker counts'), + ): + resolve_loader_workers(2, 3) + with pytest.raises(ValueError, match='positive integer'): + resolve_loader_workers(0, None) + + +def test_load_timing_metrics_distinguishes_scheduler_and_pipeline_waiting() -> None: + metrics = load_timing_metrics( + LoadStats( + allocation_seconds=0.25, + wall_seconds=3.0, + source_wall_seconds=2.0, + cpu_seconds=1.5, + source_seconds=4.0, + validation_seconds=0.5, + scheduler_wait_seconds=1.75, + in_flight_sources_mean=1.5, + in_flight_sources_max=2, + ) + ) + + assert metrics['source_load_seconds'] == 2.0 + assert metrics['staging_seconds'] == 3.0 + assert metrics['writer_idle_seconds'] == 3.0 + assert metrics['scheduler_wait_seconds'] == 1.75 + assert metrics['loader_blocked_seconds'] == 0.0 + assert metrics['queue_occupancy_mean'] == 0.0 + assert metrics['in_flight_sources_max'] == 2 + + +def test_scalar_output_stats_avoids_cumulative_combined_file_size(tmp_path: Path) -> None: + output = tmp_path / 'combined.h5' + output.write_bytes(b'cumulative') + + assert scalar_output_stats({}, 'FA', output, split_output=False) == (None, None) + assert scalar_output_stats({}, 'FA', output, split_output=True) == ( + len(b'cumulative'), + 1, + ) + assert scalar_output_stats( + {'FA': {'output_bytes': 123, 'output_file_count': 4}}, + 'FA', + output, + split_output=False, + ) == (123, 4) + + +def test_atomic_file_publication_replaces_only_after_success(tmp_path: Path) -> None: + output = tmp_path / 'result.h5' + output.write_bytes(b'old') + + with atomic_output(output, directory=False) as temporary: + assert output.read_bytes() == b'old' + temporary.write_bytes(b'new') + + assert output.read_bytes() == b'new' + assert not list(tmp_path.glob('.*.modelarrayio-tmp-*')) + + +def test_atomic_directory_failure_preserves_published_output(tmp_path: Path) -> None: + output = tmp_path / 'result.tdb' + output.mkdir() + (output / 'old.txt').write_text('old', encoding='utf-8') + + def _fail_during_conversion() -> None: + with atomic_output(output, directory=True) as temporary: + temporary.mkdir() + (temporary / 'new.txt').write_text('new', encoding='utf-8') + raise RuntimeError('injected failure') + + with pytest.raises(RuntimeError, match='injected failure'): + _fail_during_conversion() + + assert (output / 'old.txt').read_text(encoding='utf-8') == 'old' + assert not list(tmp_path.glob('.*.modelarrayio-tmp-*')) + assert not list(tmp_path.glob('.*.modelarrayio-backup-*')) + + +@pytest.mark.parametrize('directory', [False, True]) +def test_atomic_output_cleans_temporary_path_on_keyboard_interrupt( + tmp_path: Path, directory: bool +) -> None: + output = tmp_path / ('result.tdb' if directory else 'result.h5') + if directory: + output.mkdir() + (output / 'old.txt').write_text('old', encoding='utf-8') + else: + output.write_bytes(b'old') + + def _interrupt_conversion() -> None: + with atomic_output(output, directory=directory) as temporary: + if directory: + temporary.mkdir() + (temporary / 'new.txt').write_text('new', encoding='utf-8') + else: + temporary.write_bytes(b'new') + raise KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + _interrupt_conversion() + + if directory: + assert (output / 'old.txt').read_text(encoding='utf-8') == 'old' + else: + assert output.read_bytes() == b'old' + assert not list(tmp_path.glob('.*.modelarrayio-tmp-*')) + + +def test_atomic_directory_publish_rolls_back_on_keyboard_interrupt( + tmp_path: Path, monkeypatch +) -> None: + output = tmp_path / 'result.tdb' + output.mkdir() + (output / 'old.txt').write_text('old', encoding='utf-8') + real_replace = Path.replace + + def _interrupt_temporary_publish(self, target): + if '.modelarrayio-tmp-' in self.name and Path(target) == output: + raise KeyboardInterrupt + return real_replace(self, target) + + monkeypatch.setattr(Path, 'replace', _interrupt_temporary_publish) + + def _interrupt_publication() -> None: + with atomic_output(output, directory=True) as temporary: + temporary.mkdir() + (temporary / 'new.txt').write_text('new', encoding='utf-8') + + with pytest.raises(KeyboardInterrupt): + _interrupt_publication() + + assert (output / 'old.txt').read_text(encoding='utf-8') == 'old' + assert not list(tmp_path.glob('.*.modelarrayio-tmp-*')) + assert not list(tmp_path.glob('.*.modelarrayio-backup-*')) + + +def test_validate_hdf5_output_checks_source_name_count(tmp_path: Path) -> None: + output = tmp_path / 'result.h5' + with h5py.File(output, 'w') as h5_file: + h5_file.create_dataset('scalars/FA/values', data=np.zeros((2, 3))) + h5_file.create_dataset('scalars/FA/column_names', data=np.array([b'one'])) + + with pytest.raises(RuntimeError, match='2 rows but 1 source names'): + validate_hdf5_output(output, ['FA']) + + +def test_validate_hdf5_output_checks_full_contract(tmp_path: Path) -> None: + output = tmp_path / 'result.h5' + with h5py.File(output, 'w') as h5_file: + h5_file.create_dataset('voxels', data=np.zeros((4, 3))) + h5_file.create_dataset('scalars/FA/values', data=np.zeros((2, 3), dtype=np.float32)) + h5_file.create_dataset( + 'scalars/FA/column_names', + data=np.array(['two', 'one'], dtype=h5py.string_dtype('utf-8')), + ) + + with pytest.raises(RuntimeError, match='source names or order'): + validate_hdf5_output( + output, + ['FA'], + expected_sources={'FA': ['one', 'two']}, + expected_dtype='float32', + expected_elements=3, + required_paths=('voxels',), + ) + with pytest.raises(RuntimeError, match='elements != expected'): + validate_hdf5_output(output, ['FA'], expected_elements=4) + with pytest.raises(RuntimeError, match='missing required HDF5 path'): + validate_hdf5_output(output, ['FA'], required_paths=('missing',)) + + +def test_validate_tiledb_output_checks_metadata_and_cifti_name_array(tmp_path: Path) -> None: + output = tmp_path / 'result.tdb' + uri = create_scalar_matrix_array( + str(output), + 'scalars/FA/values', + np.zeros((2, 3), dtype=np.float32), + ['one', 'two'], + compression='none', + ) + validate_tiledb_output(output, ['FA']) + + with pytest.raises(RuntimeError, match='missing TileDB array'): + validate_tiledb_output(output, ['FA'], require_column_name_arrays=True) + + with tiledb.open(uri, 'w') as array: + del array.meta['column_names'] + with pytest.raises(RuntimeError, match='0 source names'): + validate_tiledb_output(output, ['FA']) + + +def test_validate_tiledb_output_checks_full_contract(tmp_path: Path) -> None: + output = tmp_path / 'result.tdb' + tiledb_storage.create_scalar_matrix_array( + str(output), + 'scalars/FA/values', + np.zeros((2, 3), dtype=np.float32), + ['one', 'two'], + compression='none', + ) + tiledb_storage.write_column_names(str(output), 'FA', ['one', 'two']) + validate_tiledb_output( + output, + ['FA'], + require_column_name_arrays=True, + expected_sources={'FA': ['one', 'two']}, + expected_dtype='float32', + expected_elements=3, + ) + + tiledb_storage.write_column_names(str(output), 'FA', ['two', 'one']) + with pytest.raises(RuntimeError, match='column-name array does not match'): + validate_tiledb_output(output, ['FA'], require_column_name_arrays=True) + with pytest.raises(RuntimeError, match='elements != expected'): + validate_tiledb_output(output, ['FA'], expected_elements=4) + + +def test_validate_timing_jsonl_path_rejects_output_collisions(tmp_path: Path) -> None: + hdf5_output = tmp_path / 'result.h5' + tiledb_output = tmp_path / 'result.tdb' + + with pytest.raises(ValueError, match='conflicts with conversion output'): + validate_timing_jsonl_path( + hdf5_output, + [hdf5_output], + directory_outputs=False, + ) + with pytest.raises(ValueError, match='conflicts with conversion output'): + validate_timing_jsonl_path( + tiledb_output / 'timing.jsonl', + [tiledb_output], + directory_outputs=True, + ) + validate_timing_jsonl_path( + tmp_path / 'timing.jsonl', + [hdf5_output, tiledb_output], + directory_outputs=False, + ) + + +def test_timing_recorder_appends_json_lines_and_path_stats(tmp_path: Path) -> None: + metrics_path = tmp_path / 'metrics' / 'runtime.jsonl' + recorder = TimingRecorder(metrics_path, modality='nifti', workers=2) + recorder.emit('scalar_complete', scalar_name='FA', matrix_shape=[2, 3]) + recorder.finish(scalar_count=1) + + records = [json.loads(line) for line in metrics_path.read_text().splitlines()] + assert [record['event'] for record in records] == [ + 'scalar_complete', + 'conversion_complete', + ] + assert all(record['modality'] == 'nifti' for record in records) + assert records[-1]['cpu_utilization_percent'] >= 0 + assert records[-1]['s3_pool_connections'] >= 10 + assert 'peak_open_file_descriptors_sampled' in records[-1] + assert path_stats(metrics_path) == (metrics_path.stat().st_size, 1) + + +def test_timing_recorder_reports_tiledb_concurrency(tmp_path: Path) -> None: + recorder = TimingRecorder(tmp_path / 'tiledb.jsonl', backend='tiledb') + record = recorder.finish() + + assert record['tiledb_compute_concurrency_level'] >= 1 + assert record['tiledb_io_concurrency_level'] >= 1 + + +def test_timing_recorder_jsonl_failure_does_not_fail_conversion_metrics( + tmp_path: Path, caplog +) -> None: + jsonl_directory = tmp_path / 'metrics.jsonl' + jsonl_directory.mkdir() + recorder = TimingRecorder(jsonl_directory) + + with caplog.at_level('WARNING', logger=runtime.__name__): + record = recorder.emit('scalar_complete') + + assert record['event'] == 'scalar_complete' + assert recorder.jsonl_path is None + assert any('continuing without JSONL output' in item.message for item in caplog.records) + + +def test_publish_scalar_outputs_reports_loader_stats_filled_on_access(tmp_path) -> None: + import numpy as np + + from modelarrayio.cli.runtime import TimingRecorder, publish_scalar_outputs + from modelarrayio.utils.loading import LazyScalarMatrices, LoadStats + + def _stage(scalar_name): + stats = LoadStats(source_count=2, workers=1, fast_path_rows=1) + return np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32), stats + + scalars = LazyScalarMatrices(['FA'], _stage) + assert scalars.stats == {} # nothing staged yet: the loop must not snapshot this + jsonl = tmp_path / 'timing.jsonl' + recorder = TimingRecorder(jsonl, modality='test', backend='hdf5') + + published = publish_scalar_outputs( + recorder, + backend='hdf5', + output=tmp_path / 'out.h5', + split_outputs=False, + scalar_names=['FA'], + scalars=scalars, + sources_by_scalar={'FA': ['a.nii', 'b.nii']}, + storage_dtype='float32', + compression='gzip', + compression_level=4, + shuffle=True, + chunk_voxels=0, + target_chunk_mb=2.0, + expected_elements=2, + ) + + assert published == [tmp_path / 'out.h5'] + records = [json.loads(line) for line in jsonl.read_text().splitlines() if line.strip()] + scalar_records = [record for record in records if record['event'] == 'scalar_complete'] + assert len(scalar_records) == 1 + assert scalar_records[0]['fast_path_rows'] == 1 + assert scalar_records[0]['matrix_shape'] == [2, 2] diff --git a/test/test_runtime_integration_matrix.py b/test/test_runtime_integration_matrix.py new file mode 100644 index 0000000..1dc6653 --- /dev/null +++ b/test/test_runtime_integration_matrix.py @@ -0,0 +1,186 @@ +"""Small end-to-end runtime matrix across modalities, backends, workers, and dtypes.""" + +from __future__ import annotations + +from pathlib import Path + +import h5py +import nibabel as nb +import numpy as np +import pandas as pd +import pytest +import tiledb +from nibabel.cifti2.cifti2_axes import BrainModelAxis, ScalarAxis + +from modelarrayio.cli.main import main as modelarrayio_main +from modelarrayio.utils.mif_image import MifImage + +RUNTIME_CASES = [ + pytest.param(backend, dtype, workers, id=f'{backend}-{dtype}-workers{workers}') + for backend in ('hdf5', 'tiledb') + for dtype in ('float32', 'float64') + for workers in (1, 2) +] + + +def _read_values(output: Path, backend: str, scalar_name: str) -> np.ndarray: + if backend == 'hdf5': + with h5py.File(output, 'r') as h5_file: + return h5_file[f'scalars/{scalar_name}/values'][...] + with tiledb.open(str(output / 'scalars' / scalar_name / 'values'), 'r') as array: + return array[:]['values'] + + +def _common_args( + cohort: Path, + output: Path, + backend: str, + dtype: str, + workers: int, +) -> list[str]: + return [ + 'to-modelarray', + '--cohort-file', + str(cohort), + '--output', + str(output), + '--backend', + backend, + '--dtype', + dtype, + '--workers', + str(workers), + '--compression', + 'none', + ] + + +@pytest.mark.parametrize(('backend', 'dtype', 'workers'), RUNTIME_CASES) +def test_nifti_runtime_matrix( + tmp_path: Path, + backend: str, + dtype: str, + workers: int, +) -> None: + mask_path = tmp_path / 'mask.nii.gz' + mask = np.ones((2, 2, 1), dtype=np.uint8) + nb.Nifti1Image(mask, np.eye(4)).to_filename(mask_path) + values = [ + np.array([[[1.125], [2.25]], [[3.5], [4.75]]], dtype=np.float64), + np.array([[[5.125], [6.25]], [[7.5], [8.75]]], dtype=np.float64), + ] + sources = [] + for index, data in enumerate(values): + source = tmp_path / f'sub-{index}.nii.gz' + nb.Nifti1Image(data, np.eye(4)).to_filename(source) + sources.append(source) + cohort = tmp_path / 'cohort.csv' + pd.DataFrame( + { + 'scalar_name': ['FA', 'FA'], + 'source_file': [str(source) for source in sources], + 'source_mask_file': [str(mask_path), str(mask_path)], + } + ).to_csv(cohort, index=False) + output = tmp_path / ('output.h5' if backend == 'hdf5' else 'output.tdb') + + status = modelarrayio_main( + [*_common_args(cohort, output, backend, dtype, workers), '--mask', str(mask_path)] + ) + + assert status == 0 + actual = _read_values(output, backend, 'FA') + expected = np.stack([value.ravel() for value in values]).astype(dtype) + assert actual.dtype == np.dtype(dtype) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize(('backend', 'dtype', 'workers'), RUNTIME_CASES) +def test_cifti_runtime_matrix( + tmp_path: Path, + backend: str, + dtype: str, + workers: int, +) -> None: + mask = np.ones((1, 1, 2), dtype=bool) + brain_axis = BrainModelAxis.from_mask(mask) + values = [ + np.array([1.125, 2.25], dtype=np.float64), + np.array([3.5, 4.75], dtype=np.float64), + ] + sources = [] + for index, data in enumerate(values): + source = tmp_path / f'sub-{index}.dscalar.nii' + header = nb.cifti2.Cifti2Header.from_axes( + (ScalarAxis([f'scalar-label-{index}']), brain_axis) + ) + nb.Cifti2Image(data.reshape(1, -1), header=header).to_filename(source) + sources.append(source) + cohort = tmp_path / 'cohort.csv' + pd.DataFrame( + { + 'scalar_name': ['FA', 'FA'], + 'source_file': [str(source) for source in sources], + } + ).to_csv(cohort, index=False) + output = tmp_path / ('output.h5' if backend == 'hdf5' else 'output.tdb') + + status = modelarrayio_main(_common_args(cohort, output, backend, dtype, workers)) + + assert status == 0 + actual = _read_values(output, backend, 'FA') + expected = np.stack(values).astype(dtype) + assert actual.dtype == np.dtype(dtype) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize(('backend', 'dtype', 'workers'), RUNTIME_CASES) +def test_mif_runtime_matrix( + tmp_path: Path, + backend: str, + dtype: str, + workers: int, +) -> None: + index_path = tmp_path / 'index.mif' + directions_path = tmp_path / 'directions.mif' + index_data = np.zeros((2, 2, 2, 2), dtype=np.float32) + index_data[..., 0] = 1 + index_data[..., 1] = np.arange(8, dtype=np.float32).reshape(2, 2, 2) + MifImage(index_data, np.eye(4)).to_filename(index_path) + MifImage( + np.tile(np.array([[1.0, 0.0, 0.0]], dtype=np.float32), (8, 1)), + np.eye(4), + ).to_filename(directions_path) + + values = [ + np.arange(8, dtype=np.float64) + 0.125, + np.arange(8, dtype=np.float64) + 10.25, + ] + sources = [] + for index, data in enumerate(values): + source = tmp_path / f'sub-{index}.mif' + MifImage(data, np.eye(4)).to_filename(source) + sources.append(source) + cohort = tmp_path / 'cohort.csv' + pd.DataFrame( + { + 'scalar_name': ['FD', 'FD'], + 'source_file': [str(source) for source in sources], + } + ).to_csv(cohort, index=False) + output = tmp_path / ('output.h5' if backend == 'hdf5' else 'output.tdb') + args = [ + *_common_args(cohort, output, backend, dtype, workers), + '--index-file', + str(index_path), + '--directions-file', + str(directions_path), + ] + + status = modelarrayio_main(args) + + assert status == 0 + actual = _read_values(output, backend, 'FD') + expected = np.stack(values).astype(dtype) + assert actual.dtype == np.dtype(dtype) + np.testing.assert_array_equal(actual, expected) diff --git a/test/test_s3_utils.py b/test/test_s3_utils.py index ef03ff4..0c03bea 100644 --- a/test/test_s3_utils.py +++ b/test/test_s3_utils.py @@ -2,8 +2,12 @@ from __future__ import annotations +import sys +import types + import pytest +from modelarrayio.utils import s3_utils from modelarrayio.utils.s3_utils import is_s3_path @@ -21,3 +25,126 @@ ) def test_is_s3_path(path: str, expected: bool) -> None: assert is_s3_path(path) is expected + + +def test_s3_client_is_shared_and_pool_matches_workers(monkeypatch) -> None: + calls = [] + client = object() + + class _FakeBoto3: + @staticmethod + def client(service, **kwargs): + calls.append((service, kwargs)) + return client + + monkeypatch.setitem(sys.modules, 'boto3', _FakeBoto3) + monkeypatch.setitem(sys.modules, 'botocore', types.SimpleNamespace(UNSIGNED='unsigned')) + monkeypatch.setitem( + sys.modules, + 'botocore.config', + types.SimpleNamespace(Config=lambda **kwargs: ('config', kwargs)), + ) + monkeypatch.setenv('MODELARRAYIO_S3_ANON', '1') + s3_utils.clear_s3_client_cache() + s3_utils.configure_s3_client_pool(23) + + assert s3_utils._make_s3_client() is client + assert s3_utils._make_s3_client() is client + assert calls == [ + ( + 's3', + { + 'config': ( + 'config', + {'max_pool_connections': 23, 'signature_version': 'unsigned'}, + ) + }, + ) + ] + + s3_utils.clear_s3_client_cache() + s3_utils.configure_s3_client_pool(1) + + +def test_s3_client_cache_closes_replaced_and_cleared_clients(monkeypatch) -> None: + clients = [] + + class _Client: + def __init__(self): + self.closed = False + + def close(self): + self.closed = True + + class _FakeBoto3: + @staticmethod + def client(_service, **_kwargs): + client = _Client() + clients.append(client) + return client + + monkeypatch.setitem(sys.modules, 'boto3', _FakeBoto3) + monkeypatch.setitem(sys.modules, 'botocore', types.SimpleNamespace(UNSIGNED='unsigned')) + monkeypatch.setitem( + sys.modules, + 'botocore.config', + types.SimpleNamespace(Config=lambda **kwargs: kwargs), + ) + monkeypatch.setenv('MODELARRAYIO_S3_ANON', '0') + s3_utils.clear_s3_client_cache() + s3_utils.configure_s3_client_pool(10) + + first = s3_utils._make_s3_client() + s3_utils.configure_s3_client_pool(24) + assert first.closed is True + second = s3_utils._make_s3_client() + assert second is not first + + s3_utils.clear_s3_client_cache() + assert second.closed is True + s3_utils.configure_s3_client_pool(1) + + +def test_read_s3_bytes_closes_response_body(monkeypatch) -> None: + class _Body: + closed = False + + def read(self): + return b'payload' + + def close(self): + self.closed = True + + body = _Body() + client = types.SimpleNamespace(get_object=lambda **_kwargs: {'Body': body}) + monkeypatch.setattr(s3_utils, '_make_s3_client', lambda: client) + + assert s3_utils.read_s3_bytes('s3://bucket/key') == b'payload' + assert body.closed is True + + +def test_read_s3_bytes_closes_body_when_read_fails(monkeypatch) -> None: + class _Body: + closed = False + + def read(self): + raise OSError('network interrupted') + + def close(self): + self.closed = True + + body = _Body() + client = types.SimpleNamespace(get_object=lambda **_kwargs: {'Body': body}) + monkeypatch.setattr(s3_utils, '_make_s3_client', lambda: client) + + with pytest.raises(OSError, match='network interrupted'): + s3_utils.read_s3_bytes('s3://bucket/key') + assert body.closed is True + + +def test_broken_datalad_symlink_has_actionable_error(tmp_path) -> None: + source = tmp_path / 'image.nii.gz' + source.symlink_to(tmp_path / '.git' / 'annex' / 'objects' / 'missing') + + with pytest.raises(FileNotFoundError, match='datalad get'): + s3_utils._check_local_source(source) diff --git a/test/test_storage_utils.py b/test/test_storage_utils.py index 65354b2..59d2dd1 100644 --- a/test/test_storage_utils.py +++ b/test/test_storage_utils.py @@ -28,6 +28,65 @@ def test_compute_full_subject_chunk_shape_auto() -> None: assert 1 <= chunk[1] <= 128 +@pytest.mark.parametrize( + ('item_chunk', 'target_chunk_mb'), + [(-1, 1.0), (1.5, 1.0), (True, 1.0), (0, 0.0), (0, -1.0), (0, np.nan)], +) +def test_compute_full_subject_chunk_shape_rejects_invalid_layout( + item_chunk, target_chunk_mb +) -> None: + with pytest.raises(ValueError, match=r'item chunk|target chunk size'): + storage_utils.compute_full_subject_chunk_shape( + n_files=4, + n_elements=128, + item_chunk=item_chunk, + target_chunk_mb=target_chunk_mb, + storage_np_dtype='float32', + ) + + def test_normalize_column_names_accepts_series_and_sequences() -> None: assert storage_utils.normalize_column_names(['a', 2]) == ['a', '2'] assert storage_utils.normalize_column_names(pd.Series(['x', 'y'])).__class__ is list + + +@pytest.mark.parametrize('name', ['', ' ', '.', '..', 'a/b', r'a\b', 'a\x00b']) +def test_validate_scalar_name_rejects_unsafe_path_segments(name: str) -> None: + with pytest.raises(ValueError, match=r'Scalar names|Unsafe scalar name'): + storage_utils.validate_scalar_name(name) + + +def test_validate_scalar_name_preserves_valid_identifier() -> None: + assert storage_utils.validate_scalar_name('mean FA') == 'mean FA' + + +@pytest.mark.parametrize( + ('backend', 'compression', 'level', 'expected'), + [ + ('hdf5', None, None, ('gzip', 4)), + ('tiledb', None, None, ('zstd', 5)), + ('hdf5', 'lzf', 9, ('lzf', None)), + ('tiledb', 'none', 9, ('none', None)), + ('tiledb', 'gzip', 2, ('gzip', 2)), + ], +) +def test_resolve_backend_compression_defaults_and_levels( + backend, compression, level, expected +) -> None: + assert storage_utils.resolve_backend_compression(backend, compression, level) == expected + + +@pytest.mark.parametrize( + ('backend', 'compression', 'level'), + [ + ('hdf5', 'zstd', None), + ('tiledb', 'lzf', None), + ('hdf5', 'gzip', 10), + ('hdf5', 'gzip', 2.0), + ], +) +def test_resolve_backend_compression_rejects_invalid_combinations( + backend, compression, level +) -> None: + with pytest.raises(ValueError, match=r'not supported|between 0 and 9|must be an integer'): + storage_utils.resolve_backend_compression(backend, compression, level) diff --git a/test/test_tiledb_storage.py b/test/test_tiledb_storage.py index 5ed4d31..f1c2027 100644 --- a/test/test_tiledb_storage.py +++ b/test/test_tiledb_storage.py @@ -22,6 +22,10 @@ def test_build_filter_list_variants() -> None: assert isinstance(no_filters, tiledb.FilterList) assert len(no_filters) == 0 + # ``none`` is a true no-filter fast path even when the CLI's default + # shuffle flag remains enabled. + assert len(tiledb_storage._build_filter_list('none', None, shuffle=True)) == 0 + zstd = tiledb_storage._build_filter_list('zstd', 9, shuffle=True) assert len(zstd) >= 2 @@ -29,6 +33,22 @@ def test_build_filter_list_variants() -> None: assert len(fallback) == 0 +@pytest.mark.parametrize( + ('compression', 'level', 'expected'), + [ + ('zstd', None, ('zstd', 5)), + ('gzip', None, ('gzip', 4)), + ('zstd', 'bad', ('zstd', 5)), + (None, 9, ('none', None)), + ('unknown', 9, ('none', None)), + ], +) +def test_resolve_compression_settings_reports_effective_values( + compression, level, expected +) -> None: + assert tiledb_storage._resolve_compression_settings(compression, level) == expected + + def test_create_empty_scalar_matrix_array_writes_metadata_and_overwrites(tmp_path: Path) -> None: base = tmp_path / 'store.tdb' uri = tiledb_storage.create_empty_scalar_matrix_array( @@ -75,10 +95,14 @@ def test_write_rows_in_column_stripes_round_trip(tmp_path: Path) -> None: np.array([6, 7, 8, 9, 10], dtype=np.float32), np.array([11, 12, 13, 14, 15], dtype=np.float32), ] - tiledb_storage.write_rows_in_column_stripes(uri, rows) + metrics = tiledb_storage.write_rows_in_column_stripes(uri, rows) with tiledb.open(uri, 'r') as array: np.testing.assert_array_equal(array[:]['values'], np.vstack(rows)) + assert len(tiledb.array_fragments(uri)) == 1 + assert metrics['dense_submissions'] == 1 + assert metrics['fragment_count'] == 1 + json.dumps(metrics) def test_write_rows_in_column_stripes_rejects_wrong_row_count(tmp_path: Path) -> None: @@ -122,6 +146,7 @@ def test_write_parcel_names_and_column_names(tmp_path: Path) -> None: def test_create_scalar_matrix_array_writes_values_and_metadata(tmp_path: Path) -> None: base = tmp_path / 'store.tdb' values = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) + metrics = {} uri = tiledb_storage.create_scalar_matrix_array( str(base), 'scalars/MD/values', @@ -130,8 +155,106 @@ def test_create_scalar_matrix_array_writes_values_and_metadata(tmp_path: Path) - storage_dtype='float32', compression='zstd', compression_level=3, + metrics=metrics, ) with tiledb.open(uri, 'r') as array: assert_domain_dtypes(array, ['int32', 'int32']) np.testing.assert_array_equal(array[:]['values'], values) assert json.loads(array.meta['column_names']) == ['first', 'second'] + assert len(tiledb.array_fragments(uri)) == 1 + assert metrics['dense_submissions'] == 1 + assert metrics['fragment_count'] == 1 + assert metrics['matrix_shape'] == [2, 2] + json.dumps(metrics) + + +def test_create_scalar_matrix_array_reports_default_compression_level(tmp_path: Path) -> None: + metrics = {} + tiledb_storage.create_scalar_matrix_array( + str(tmp_path / 'store.tdb'), + 'scalars/MD/values', + np.ones((2, 2), dtype=np.float32), + ['first', 'second'], + compression='zstd', + compression_level=None, + metrics=metrics, + ) + + assert metrics['compression'] == 'zstd' + assert metrics['compression_level'] == 5 + + +def test_create_scalar_matrix_array_reports_effective_shuffle_for_no_compression( + tmp_path: Path, +) -> None: + metrics = {} + uri = tiledb_storage.create_scalar_matrix_array( + str(tmp_path / 'store.tdb'), + 'scalars/MD/values', + np.ones((2, 2), dtype=np.float32), + ['first', 'second'], + compression='none', + compression_level=None, + shuffle=True, + metrics=metrics, + ) + + with tiledb.open(uri, 'r') as array: + assert len(array.schema.attr('values').filters) == 0 + assert metrics['compression'] == 'none' + assert metrics['shuffle'] is False + + +def test_create_scalar_matrix_array_skips_output_stats_without_metrics( + tmp_path: Path, monkeypatch +) -> None: + def _unexpected_stats(*_args, **_kwargs): + raise AssertionError('output statistics should only be collected for telemetry') + + monkeypatch.setattr(tiledb_storage, '_local_output_stats', _unexpected_stats) + monkeypatch.setattr(tiledb_storage.tiledb, 'array_fragments', _unexpected_stats) + + uri = tiledb_storage.create_scalar_matrix_array( + str(tmp_path / 'store.tdb'), + 'scalars/MD/values', + np.ones((2, 2), dtype=np.float32), + ['first', 'second'], + ) + with tiledb.open(uri, 'r') as array: + np.testing.assert_array_equal(array[:]['values'], np.ones((2, 2), dtype=np.float32)) + + +def test_write_full_matrix_rejects_non_matrix_schema(tmp_path: Path) -> None: + uri = str(tmp_path / 'one-dimensional.tdb') + dimension = tiledb.Dim(name='items', domain=(0, 2), tile=3, dtype=np.int32) + schema = tiledb.ArraySchema( + domain=tiledb.Domain(dimension), + attrs=[tiledb.Attr(name='values', dtype=np.float32)], + sparse=False, + ) + tiledb.Array.create(uri, schema) + + with pytest.raises(ValueError, match='must have two dimensions'): + tiledb_storage.write_full_matrix(uri, np.ones((1, 3), dtype=np.float32)) + + +def test_write_full_matrix_rejects_unexpected_attributes(tmp_path: Path) -> None: + uri = str(tmp_path / 'wrong-attributes.tdb') + dimensions = [ + tiledb.Dim(name='subjects', domain=(0, 1), tile=2, dtype=np.int32), + tiledb.Dim(name='items', domain=(0, 2), tile=3, dtype=np.int32), + ] + schema = tiledb.ArraySchema( + domain=tiledb.Domain(*dimensions), + attrs=[tiledb.Attr(name='other', dtype=np.float32)], + sparse=False, + ) + tiledb.Array.create(uri, schema) + + with pytest.raises(ValueError, match="attribute named 'values'"): + tiledb_storage.write_full_matrix(uri, np.ones((2, 3), dtype=np.float32)) + + +def test_full_matrix_coercion_preserves_staged_array() -> None: + matrix = np.empty((2, 3), dtype=np.float32, order='C') + assert tiledb_storage._coerce_full_matrix(matrix, dtype=np.float32) is matrix diff --git a/uv.lock b/uv.lock index 700f0e8..aa62b88 100644 --- a/uv.lock +++ b/uv.lock @@ -557,7 +557,6 @@ dependencies = [ { name = "numpy" }, { name = "pandas" }, { name = "tiledb" }, - { name = "tqdm" }, ] [package.optional-dependencies] @@ -643,7 +642,6 @@ requires-dist = [ { name = "sphinxcontrib-bibtex", marker = "extra == 'all'" }, { name = "sphinxcontrib-bibtex", marker = "extra == 'doc'" }, { name = "tiledb", specifier = ">=0.23" }, - { name = "tqdm", specifier = ">=4.65" }, ] provides-extras = ["all", "doc", "s3", "test"] @@ -1176,27 +1174,27 @@ wheels = [ [[package]] name = "ruff" -version = "0.15.20" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/43/dc/35b341fc554ba02f217fc10da57d1a75168cfbcf75b0ef2202176d4c4f2d/ruff-0.15.20.tar.gz", hash = "sha256:1416eb04349192646b54de98f146c4f59afe37d0decfc02c3cbbf396f3a28566", size = 4755489, upload-time = "2026-06-25T17:20:37.578Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/94/d9/2d5014f0253ba541d2061d9fa7193f48e941c8b21bb88a7ff9bbe0bd0596/ruff-0.15.20-py3-none-linux_armv6l.whl", hash = "sha256:00e188c53e499c3c1637f73c91dcf2fb56d576cab76ce1be50a27c4e80e37078", size = 10839665, upload-time = "2026-06-25T17:19:44.702Z" }, - { url = "https://files.pythonhosted.org/packages/c6/d3/ac1798ba64f670698867fcfc591d50e7e421bef137db564858f619a30fcf/ruff-0.15.20-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:9ebd1fd9b9c95fc0bd7b2761aebec1f030013d2e193a2901b224af68fe47251b", size = 11208649, upload-time = "2026-06-25T17:19:48.787Z" }, - { url = "https://files.pythonhosted.org/packages/47/47/d3ac899991202095dfcf3d5176be4272642be3cf981a2f1a30f72a2afb95/ruff-0.15.20-py3-none-macosx_11_0_arm64.whl", hash = "sha256:c5b16cdd67ca108185cd36dce98c576350c03b1660a751de725fb049193a0632", size = 10622638, upload-time = "2026-06-25T17:19:51.354Z" }, - { url = "https://files.pythonhosted.org/packages/33/13/4e043fe30aa94d4ff5213a9881fc296d12960f5971b234a5263fdc225312/ruff-0.15.20-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3413bb3c3d2ca6a8208f1f4809cd2dca3c6de6d0b491c0e70847672bde6e6efd", size = 10984227, upload-time = "2026-06-25T17:19:54.044Z" }, - { url = "https://files.pythonhosted.org/packages/76/e6/92e7bf40388bc5800073b96564f56264f7e48bfd1a498f5ced6ae6d5a769/ruff-0.15.20-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:bd7ec42b3bb3da066488db093308a69c4ac5ee6d2af333a86ba6e2eb2e7dd44b", size = 10622882, upload-time = "2026-06-25T17:19:57.037Z" }, - { url = "https://files.pythonhosted.org/packages/13/7a/43460be3f24495a3aa46d4b16873e2c4941b3b5f0b00cf88c03b7b94b339/ruff-0.15.20-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e1a36ad0eb77fba9aabfb69ede54de6f376d04ac18ebea022847046d340a8267", size = 11474808, upload-time = "2026-06-25T17:20:00.357Z" }, - { url = "https://files.pythonhosted.org/packages/27/a0/f37077884873221c6b33b4ab49eb18f9f88e54a16a25a5bca59bef46dd66/ruff-0.15.20-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b6df3b1e4610432f0386dba04d853b5f08cbbc903410c6fcc02f620f05aff53c", size = 12293094, upload-time = "2026-06-25T17:20:03.446Z" }, - { url = "https://files.pythonhosted.org/packages/a6/74/165545b60256a9704c21ac0ec4a0d07933b320812f9584836c9f4aca4292/ruff-0.15.20-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e89f198a1ea6ef0d727c1cf16088bc91a6cb0ab947dedc966715691647186eae", size = 11526176, upload-time = "2026-06-25T17:20:06.301Z" }, - { url = "https://files.pythonhosted.org/packages/86/b1/a976a136d40ade83ce743578399865f57001003a409acadc0ecbb3051082/ruff-0.15.20-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:309809086c2acb67624950a3c8133e80f32d0d3e27106c0cd60ff26657c9f24b", size = 11520767, upload-time = "2026-06-25T17:20:09.191Z" }, - { url = "https://files.pythonhosted.org/packages/19/0f/f032696cb01c9b54c0263fa393474d7758f1cdc021a01b04e3cbc2500999/ruff-0.15.20-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:2d2374caa2f2c2f9e2b7da0a50802cfb8b79f55a9b5e49379f564544fbf56487", size = 11500132, upload-time = "2026-06-25T17:20:13.602Z" }, - { url = "https://files.pythonhosted.org/packages/4b/f4/51b1a14bc69e8c224b15dab9cce8e99b425e0455d462caa2b3c9be2b6a8e/ruff-0.15.20-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:a1ed17b65293e0c2f22fc387bc13198a5de94bf4429589b0ff6946b0feaf21a3", size = 10943828, upload-time = "2026-06-25T17:20:16.635Z" }, - { url = "https://files.pythonhosted.org/packages/71/4b/fe267640783cd02bf6c5cc290b1df1051be2ec294c678b5c15fe19e52343/ruff-0.15.20-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:f701305e66b38ea6c91882490eb73459796808e4c6362a1b765255e0cdcd4053", size = 10645418, upload-time = "2026-06-25T17:20:19.4Z" }, - { url = "https://files.pythonhosted.org/packages/b0/c0/a65aa4ec2f5e87a1df32dc3ec1fede434fe3dfd5cbcf3b503cafc676ab54/ruff-0.15.20-py3-none-musllinux_1_2_i686.whl", hash = "sha256:5b9c0c367ad8e5d0d5b5b8537864c469a0a0e55417aadfbeca41fa61333be9f4", size = 11211770, upload-time = "2026-06-25T17:20:22.033Z" }, - { url = "https://files.pythonhosted.org/packages/5a/a4/0caa331d954ae2723d729d351c989cb4ca8b6077d5c6c2cb6de75e98c041/ruff-0.15.20-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:01cc00dd58f0df339d0e902219dd53990ea99996a0344e5d9cc8d45d5307e460", size = 11618698, upload-time = "2026-06-25T17:20:25.259Z" }, - { url = "https://files.pythonhosted.org/packages/10/9b/5f14927848d2fd4aa891fd88d883788c5a7baba561c7874732364045708c/ruff-0.15.20-py3-none-win32.whl", hash = "sha256:ed65ef510e43a137207e0f01cfcf998aeddb1aeeda5c9d35023e910284d7cf21", size = 10857322, upload-time = "2026-06-25T17:20:28.612Z" }, - { url = "https://files.pythonhosted.org/packages/fa/f0/fe47c501f9dea92a26d788ff98bb5d92ed4cb4c88792c5c88af6b697dc8e/ruff-0.15.20-py3-none-win_amd64.whl", hash = "sha256:a525c81c70fb0380344dd1d8745d8cc1c890b7fc94a58d5a07bd8eb9557b8415", size = 11993274, upload-time = "2026-06-25T17:20:31.871Z" }, - { url = "https://files.pythonhosted.org/packages/d7/2b/9555445e1201d92b3195f45cdb153a0b68f24e0a4273f6e3d5ab46e212bb/ruff-0.15.20-py3-none-win_arm64.whl", hash = "sha256:2f5b2a6d614e8700388806a14996c40fab2c47b819ef57d790a34878858ed9ca", size = 11343498, upload-time = "2026-06-25T17:20:35.03Z" }, +version = "0.17.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e7/e6d5749983b020e6cb1055abec0a6f361a8e94f0184cdcf56d4760995f02/ruff-0.17.1.tar.gz", hash = "sha256:5bb796c5112e9fb9527f2ef2faa0a130d6312d19de92bcfaddd70dbb2d41020f", size = 5323349, upload-time = "2026-10-15T15:23:18.64Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/78/90c46ca6250709f7c37e397b60d06061adee8d0f7d13959ed6a8e49b14fd/ruff-0.17.1-py3-none-linux_armv6l.whl", hash = "sha256:1796c41eafec66f4635125dd1e074a7aa27b8d20c187d484360acf11d247fc4a", size = 10147528, upload-time = "2026-10-15T15:22:33.667Z" }, + { url = "https://files.pythonhosted.org/packages/aa/d2/a2e28bc35b083ae51fa07dfba5280a667c8b01ecaee7802fc8849dcd5c46/ruff-0.17.1-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:8eadb09a788010be29c895418bdfa513e4e1c65e64b94e5b19e4cd381f6cc7fb", size = 10225974, upload-time = "2026-10-15T15:22:36.328Z" }, + { url = "https://files.pythonhosted.org/packages/8e/71/7a9029a81c8206ee01a98f6cfeeb56e6c095af86857581a37045c46d7e94/ruff-0.17.1-py3-none-macosx_11_0_arm64.whl", hash = "sha256:29b454f18ca3eeaf37fbf140841dcf27d36109e29dee9aae8acecb68c26c6f74", size = 9270361, upload-time = "2026-10-15T15:22:38.651Z" }, + { url = "https://files.pythonhosted.org/packages/c0/8c/e075e9c917cb0575297880c96768304e881240171c870cbd8d9b7b54bec2/ruff-0.17.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:872e03dd231744d6c5ec5c8af6be52b71323fc6ee096e4c8b34d9f844ee3b978", size = 9938300, upload-time = "2026-10-15T15:22:41.284Z" }, + { url = "https://files.pythonhosted.org/packages/87/27/0b4bb926a49de2383f161272b4d863571a4c7ecdbf595a6bc58813a13367/ruff-0.17.1-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4c0d6cb5f8fbac77eec9af49f4fc0eaf856b05545e372b0694d24b4e76b1c640", size = 9948464, upload-time = "2026-10-15T15:22:43.681Z" }, + { url = "https://files.pythonhosted.org/packages/14/27/518d57e0e02aecb8ea4acac6b57b51a53173d25fa734952bfdd7d531e88e/ruff-0.17.1-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:78beced4a3c94dfedbee4b999b95b6a27b35820fdc0afbacc6d37a8da26c1cd1", size = 10638639, upload-time = "2026-10-15T15:22:46.081Z" }, + { url = "https://files.pythonhosted.org/packages/01/82/d89a4e498ade59e6cee52d67f8da569019372212b87468c1fdeae884601c/ruff-0.17.1-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7feefc63921202894f0e158b0e32f45a5950d70a7fb693ccb9264816c25b1c05", size = 11337238, upload-time = "2026-10-15T15:22:48.894Z" }, + { url = "https://files.pythonhosted.org/packages/a2/7c/3b0fdef9352656317c977086fe1d802902406e4f2a0d8578da90c0e6d1c5/ruff-0.17.1-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:684154861a819bf1f433669d2a0ce9c3e54677c65b8ebb866cc2a7ec86345785", size = 10896327, upload-time = "2026-10-15T15:22:51.45Z" }, + { url = "https://files.pythonhosted.org/packages/62/86/5566290219708ce2928f3a29c4f843a1185b341ec6111c4631a33c5ceced/ruff-0.17.1-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a6ef98e51de8ef38082f90653ceeb84944bd7673783eea13e1647c5b9b064189", size = 10350811, upload-time = "2026-10-15T15:22:54.401Z" }, + { url = "https://files.pythonhosted.org/packages/ad/d2/ca8462dabb3c413990d49703ba306c30c862195a75448d080fa8fcce9775/ruff-0.17.1-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:8f9aa5b5bd5fa1937803ea609c8f153b8de39bf3e0e149897ae34b0c5b52f467", size = 10569994, upload-time = "2026-10-15T15:22:57.142Z" }, + { url = "https://files.pythonhosted.org/packages/e3/ad/dd1864c1f1db17c8cfb297ce70b352388e6977620f66f7182c52855407df/ruff-0.17.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:80e53b0e5faa3c5667f035daaa032893ae42b7cfe708c886d39ed131e54688ab", size = 10059672, upload-time = "2026-10-15T15:23:00.084Z" }, + { url = "https://files.pythonhosted.org/packages/65/45/3f56cf1d2168c2165cdd5ac6f9668aee66a4cfdd4caf2b2cc5ccb2b1d568/ruff-0.17.1-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:056c5d5fa0ac790051ae165d7035c7bf956d02de2db557f2e9b31d877a505fca", size = 9947803, upload-time = "2026-10-15T15:23:03.132Z" }, + { url = "https://files.pythonhosted.org/packages/aa/4b/5457cb2be4c1d1ea5e919558131a0f197a4a454a0530d9e30820e8206e4c/ruff-0.17.1-py3-none-musllinux_1_2_i686.whl", hash = "sha256:bd8f5a55b10f371642eacd0e8665a42a8b97da15c9a7187c79e697b77b5e0ce4", size = 10351807, upload-time = "2026-10-15T15:23:05.628Z" }, + { url = "https://files.pythonhosted.org/packages/49/fb/db15b4f8fa24a1e36b99e1401f42ef6df16ee299a48670d5d262b771e7d7/ruff-0.17.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:ce05e6f3973611d9550a600e04bb531061e368baba31610eb57c1509a8f21a0a", size = 10765716, upload-time = "2026-10-15T15:23:08.092Z" }, + { url = "https://files.pythonhosted.org/packages/0a/e8/1e186578076aff7a80735f625005fde678412ab609cd323ed182eb7d6b87/ruff-0.17.1-py3-none-win32.whl", hash = "sha256:1de0030a8aa78c65c5fecb84b2e5c390ac684c0452875a02afc7c083d49a496a", size = 9895093, upload-time = "2026-10-15T15:23:10.582Z" }, + { url = "https://files.pythonhosted.org/packages/6a/2e/da3d9589bae8e8b58f3a48bd5bc0b980232fb14a83bf02c3be7660f87c1b/ruff-0.17.1-py3-none-win_amd64.whl", hash = "sha256:5df76f16580fab58d0f6c98a3ac254a4abceb7b2489b8b678335869cff77b3d7", size = 10097955, upload-time = "2026-10-15T15:23:13.309Z" }, + { url = "https://files.pythonhosted.org/packages/df/c9/0f2a4dea06a5e12c1191ee3d96b1e1233797d0d5e18f97a2b909a95a9836/ruff-0.17.1-py3-none-win_arm64.whl", hash = "sha256:42400b9a6ff8515ff27796f07f92e0195be011c86e81e381606f58ea5502fef4", size = 9853542, upload-time = "2026-10-15T15:23:16.05Z" }, ] [[package]] @@ -1546,18 +1544,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/23/d1/136eb2cb77520a31e1f64cbae9d33ec6df0d78bdf4160398e86eec8a8754/tomli-2.4.0-py3-none-any.whl", hash = "sha256:1f776e7d669ebceb01dee46484485f43a4048746235e683bcdffacdf1fb4785a", size = 14477, upload-time = "2026-01-11T11:22:37.446Z" }, ] -[[package]] -name = "tqdm" -version = "4.68.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ae/5f/57ff8b434839e70dab45601284ea413e947a63799891b7553e5960a793a8/tqdm-4.68.4.tar.gz", hash = "sha256:19829c9673638f2a0b8617da4cdcb927e831cd88bcfcb6e78d42a4d1af131520", size = 792418, upload-time = "2026-07-07T09:58:18.369Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/22/2a/5e5e750890ada51017d18d0d4c30da696e5b5bd3180947729927628fc3cb/tqdm-4.68.4-py3-none-any.whl", hash = "sha256:5168118b2368f48c561afda8020fd79195b1bdb0bdf8086b88442c267a315dc2", size = 676612, upload-time = "2026-07-07T09:58:16.256Z" }, -] - [[package]] name = "typing-extensions" version = "4.15.0"