Optimize modelarrayio to-modelarray runtime

[tien-tong]

Summary

Rewrites the modelarrayio to-modelarray conversion path for speed without changing the output schema or breaking ModelArray reads. On real RBC freesurferpost data on CUBIC, one full 8,693-subject CIFTI scalar drops from about 3.3 h to about 7 min at one worker (27.5×), and to 69–91 s at eight workers. All 70 production scalars now convert in ≈22 min of job-array wall-clock, where the original code took 112 h for its first scalar and stalled at 15/70. Every output is bitwise-equal to the deployed reference.

The design, measurements and rejected alternatives are in design/TO_MODELARRAY_RUNTIME_OPTIMIZATION_PLAN.md.

What was slow

• HDF5 writes were capped by single-threaded chunk compression (~70 MB/s), not disk.
• TileDB writes reopened the array once per spatial stripe: one scalar became ~92 opens and ~92 fragments across hundreds of files.
• CIFTI loading parsed every file's XML header up front, GIL-bound, so extra workers made it slower.
• Every loader built the whole cohort across all scalars in memory before writing anything.
• NIfTI row assembly decoded full-volume float64 temporaries for the scalar and its mask on every row.
• S3 loading created a new client per file.

What changed

Loader (utils/loading.py, utils/{cifti,nifti,mif}.py, utils/s3_utils.py)

• One scalar at a time is staged into a typed, C-contiguous matrix through a bounded ThreadPoolExecutor.
• CIFTI rows are read directly from the data block after the first header is validated, skipping the per-file XML parse; rows served this way are counted as fast_path_rows.
• NIfTI masks are cached and selected with flat indices; decoding happens in the storage dtype.
• CIFTI axis validation compares a cheap spatial signature instead of full label arrays.
• One shared S3 connection pool sized to the loader concurrency; MODELARRAYIO_S3_ANON=1 for public buckets. Broken git-annex symlinks report an actionable error instead of being read.

Storage (storage/{h5_storage,tiledb_storage,utils}.py)

• HDF5 chunks are gzip-compressed on several threads and written with direct chunk writes (MODELARRAYIO_HDF5_COMPRESSION_WORKERS, default all CPUs).
• Each TileDB scalar matrix is written in a single dense submission: 1 fragment instead of one per stripe (350 → 8 files for MIF, 596 → 38 for CIFTI on the CUBIC comparison).
• TileDB's internal thread pools are recorded in telemetry but never capped.
• Backend compression defaults resolved in one place: HDF5 gzip level 4, TileDB zstd level 5, byte shuffle on; --compression none disables the filter pipeline.

CLI (cli/runtime.py, cli/to_modelarray.py, cli/{cifti,nifti,mif}_to_h5.py, cli/parser_utils.py)

• --workers N now means one thing everywhere: the maximum number of input sources loaded concurrently, for local and S3 paths, every modality and both backends. This changes the meaning of an existing option: it previously controlled concurrent TileDB scalar writes. The resolved meaning is logged on the first run with --workers > 1.
• --s3-workers is a deprecated alias for --workers for one minor release. Matching values are accepted, conflicting values rejected.
• --dtype float32 stays the default; --dtype float64 is the documented opt-in that keeps source precision end to end.
• New --timing-jsonl PATH appends structured per-scalar and total runtime records (load/write phases, peak RSS, fragment counts, fast-path rows, TileDB pool settings).
• Outputs are written to a same-filesystem temporary path, validated, then published atomically. Existing TileDB directories use a rename-and-rollback sequence.
• Group mask and cohort inputs accept s3:// URIs.
• tqdm dependency removed.

Docs

• README gains a "Fast conversion and input sources" section covering the worker semantics, dtype trade-off, S3/annex behaviour and timing output.
• Sphinx: autosectionlabel prefixing fixes duplicate-label warnings; the sphinx-argparse patch tolerates versions without ArgParseDomain.

Benchmarks (benchmarks/)

New self-contained suite under benchmarks/to_modelarray/: a harness that runs the full CLI across NIfTI/CIFTI/MIF × HDF5/TileDB, cohort generators (public ABIDE from s3://fcp-indi, the MIF walkthrough demo, a reproducible CIFTI synthesizer), an old-commit launcher pinned to 9ab13dd, output-equality and R read gates, a comparison-report generator, and the CUBIC SLURM ladder (container env, sbatch templates, scalar-sharded job array). Pinned results live in baselines/, the rendered CUBIC report in reports/.

Cross-modality, 100 subjects, original 9ab13dd vs this branch, one worker, CUBIC (baselines/cubic_cross_modality_9ab13dd_vs_optimized_20260922.csv):

Modality	Backend	Original (s)	Optimized (s)	Speedup	Output files	Values equal
NIfTI	HDF5	3.02	2.17	1.40×	1 → 1	yes
NIfTI	TileDB	3.41	2.33	1.46×	44 → 8	yes
MIF	HDF5	10.76	2.91	3.70×	1 → 1	yes
MIF	TileDB	14.52	3.53	4.11×	350 → 8	yes
CIFTI (3 scalars)	HDF5	385.78	12.12	31.8×	1 → 1	yes
CIFTI (3 scalars)	TileDB	396.47	12.56	31.6×	596 → 38	yes

Production scale, one full CIFTI scalar (area, 8,693 subjects), CUBIC (baselines/cubic_production_scale_9ab13dd_vs_optimized_20260923.csv):

Backend	Original, 1 worker (s)	Optimized, 1 worker (s)	4 workers (s)	8 workers (s)	Peak RSS
HDF5	11,854.8	429.9	120.9	90.6	~12 GB
TileDB	11,501.6	418.1	83.2	68.6	~20 GB

Peak memory is flat across workers because the staging matrix dominates. Other modalities were not swept at production scale, so the CLI keeps the conservative default of one worker.

Deferred

• Overlapping load and write: loading dominates (≈1.7 s load vs ≈0.1 s write per example scalar), so hiding the write buys about five percent; its real value would be hiding S3 latency.
• --only-scalar: per-scalar cohort CSVs select the scalar for the sharded HPC recipe instead.

Tests

New: test_runtime.py, test_runtime_integration_matrix.py, test_loading.py, test_storage_utils.py, test_s3_utils.py, test_benchmark_harness.py, test_comparison_report.py, test_cubic_tools.py, test_make_cifti_cohort.py. Extended: CIFTI/NIfTI utils, HDF5/TileDB storage, CLI utils and parser tests.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 84.53815% with 231 lines in your changes missing coverage. Please review.
✅ Project coverage is 87.70%. Comparing base (9ab13dd) to head (6632ffe).

Files with missing lines	Patch %	Lines
src/modelarrayio/utils/cifti.py	75.50%	49 Missing and 24 partials ⚠️
src/modelarrayio/cli/runtime.py	84.37%	27 Missing and 23 partials ⚠️
src/modelarrayio/utils/loading.py	88.46%	17 Missing and 10 partials ⚠️
src/modelarrayio/storage/h5_storage.py	85.13%	14 Missing and 8 partials ⚠️
src/modelarrayio/storage/utils.py	75.00%	9 Missing and 9 partials ⚠️
src/modelarrayio/utils/nifti.py	86.61%	11 Missing and 6 partials ⚠️
src/modelarrayio/utils/s3_utils.py	84.00%	7 Missing and 5 partials ⚠️
src/modelarrayio/storage/tiledb_storage.py	91.66%	4 Missing and 2 partials ⚠️
src/modelarrayio/utils/mif.py	81.25%	4 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #110      +/-   ##
==========================================
- Coverage   90.81%   87.70%   -3.12%     
==========================================
  Files          21       23       +2     
  Lines        1470     2602    +1132     
  Branches      305      523     +218     
==========================================
+ Hits         1335     2282     +947     
- Misses         76      198     +122     
- Partials       59      122      +63     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.