Repository navigation
assemble: guard direct-shard-encode with shard-aligned band asserts - #18
Open
mark-a-potts wants to merge 21 commits into
Open
mark-a-potts wants to merge 21 commits into
mark-a-potts wants to merge 21 commits into
Conversation
offset() accepts per-axis overlap (overlap_x, overlap_y) and derives the tile pitch from the correct axis (corr_x from width shape[-1], corr_y from height shape[-2]); the prior code used the wrong dimension, which only mattered for non-square tiles. assemble.stitch() swaps the stored tile_shape Y/X dims on odd rot90 so canvas/placement/EDT weights match the rotated pixels. Scalar overlap and square tiles are unaffected (backward-compatible). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Weight each edge in the global tile-position solve by its phase-correlation confidence, so low-confidence seams pull less on the fit. Off by default (equal-weight edges, behavior byte-for-byte unchanged). Applies to the IRLS (L1, default), LSMR (L2), and CPU fallback paths; modes: linear / squared / threshold:<cut>, with a small weight floor so no tile is left under-constrained. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…loader
Adds the LiveScreen scale-out plumbing to `assemble_streaming`:
- **Pinned host memory for D2H**: `_d2h_and_submit_writes` now writes into a
reusable pinned buffer allocated via `cupyx.empty_pinned`. Cupy's default
`.get()` returned unpinned numpy arrays that capped PCIe TX at ~1.5 GB/s on
H100 gen5; pinned unlocks full-line DMA (~50 GB/s).
- **On-GPU uint16 cast before D2H**: halves the transferred payload when the
output dtype is uint16 (default LiveScreen output). Peak per-band GPU
memory drops proportionally.
- **Direct chunk writes (bypass zarr V3 sync API)**: new `_write_zarr_block_direct`
encodes each chunk with numcodecs.Blosc and writes the file directly to
`<arr_root>/c/0/{C}/0/{Y}/{X}`. A probe on gpu-f-6 showed zarr V3's sync
API (arr[slice] = data) caps at ~500 MB/s regardless of thread count —
serial bottleneck in the shared asyncio event loop. Direct writes scaled
to ~2.8 GB/s at 16 workers. Opt-in via `use_direct_writes=True` and
requires `divide_tile_size` to match the array's chunk shape.
- **`_write_executor` bumped 14 → 32 workers**; matches the concurrency
LiveScreen convert uses to saturate NFS.
- **`RawShardTileSource` (new file)**: raw acquire-zarr shard reader with a
per-shard `_loading` Event to serialize duplicate reads (was reading each
1 GB shard 18x under the previous naive multi-thread pattern).
Together these turn write drain (the previous wall) into a hidden background
operation — per-band GPU compute is now the dominant cost, not I/O.
Adds `_process_y_band_gpu_xblock` as a generator that allocates per-X-block numer/denom instead of full-band (peak GPU memory drops ~500× — from ~19 GB per band to ~34 MB per X-block). Enables running multiple bands concurrently per GPU (multi-band CUDA streams) and, via a multi-process wrapper on top, lets N workers share a single GPU without the GIL bottleneck that limited in-process concurrency. Supporting changes: - `assemble_streaming` gets `xblock_scoped_gpu`, `n_concurrent_bands`, `y_range` parameters. `y_range` is what the multi-process runner uses to partition the canvas across worker processes. - `_write_zarr_block_direct` extended with a matching `_d2h_and_write_xblock` path that allocates per-call pinned host buffers (cupyx.empty_pinned). - `_prefetch_depth` is now env-configurable via `OPS_PREFETCH_DEPTH`. - Fix subtle bug in the tile-meta build: was calling `xp.asarray(shift)` per tile — 167k cupy roundtrips + D2H syncs (~2 min extra per worker when N processes serialize on shared CUDA driver context). Now uses numpy for the CPU-only slice math. - Fix shared-dict bug in multi-band path: `tile_cache.clear()` was wiping the dict a worker was still reading. Guard the clear call with the multi-band flag. Measured on LiveScreen 06-26 (167k tiles, converted HCS input): serial single-process: ~35 min (baseline direct writes + pinned D2H) 8 workers × 8 loader × 2 prefetch: **14.2 min** (best sweep result) 6 workers × 10 loader: 15.8 min 4 workers × 8 loader: 20.1 min (write bottleneck reappears) Full pipeline convert + assemble now ~49 min vs 8-16 h projected on the original direct-shard reader.
PyYAML's pure-Python `SafeLoader` was taking ~30 s to parse the 35 MB LiveScreen shifts YAML (167k `total_translation` entries), single-threaded in every process that calls `read_shifts_biahub`. In the multi-process assemble path that's 1 parent + 8 workers = 9 parses = ~4.5 min of wasted main-thread setup before workers actually start GPU work. CSafeLoader (backed by libyaml C bindings) parses the same file in ~2-3 s. Falls back to SafeLoader if PyYAML wasn't built against libyaml.
The direct chunk-write path wrote trailing edge blocks at their truncated size,
but zarr v3 stores every chunk at the full chunk shape (fill-padded) and crops on
read using the array's shape metadata. A short edge chunk was therefore
undecodable ("cannot reshape array of size ... into shape ..."). Pad the
bottom/right of edge blocks with zeros up to the full chunk shape before encoding
so the bytes on disk match what the reader expects. The zarr-API write path was
never affected because it pads internally.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…streaming fetch - assemble.py: OPS_ASSEMBLE_BATCH_UPLOAD uploads each unique tile in a band to GPU once; X-block loop then slices D2D instead of firing ~1200-1600 per-touch H2Ds - assemble.py: OPS_ASSEMBLE_FUSED_ACCUM collapses the 8 per-touch launches (slice, cast, ne, mult, iadd_num, iadd_den) into one ElementwiseKernel; the tile stays uint16 through the slice and the kernel casts inline; tile_weights cached on GPU once per band - converted_stream.py: dragon_fetch_server streams DRAGON_FETCH_MSG_TILES tiles per put() with a (host, None) sentinel — the prior single-put of the full batch dict (~4.5 GB) overflowed Dragon's 4 GB default dynheap and deadlocked the loader On the ops0113 convert store (2 nodes, 4 workers total across 4 H100s), assemble wall drops from ~22 min baseline to 14:47 (-33 %). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…(K=1 stable) Adds a CuPy RawKernel (env-gated OPS_ASSEMBLE_BATCHED_KERNEL=1) that launches ONCE per X-block and internally iterates all K tile-touches, replacing the Python for-loop over touches. Requires the batch-upload path so tiles are already on GPU as a stackable (N_unique, TCZ, H, W) uint16 tensor. Design: - Module-level _BATCHED_ACCUM_EDT_KERNEL compiled lazily - Per-band setup: xp.stack([_gpu_tiles[n]... for n in names]) plus shape-validation of tile_weights against (tile_H, tile_W) - Per-X-block: gather (src_y, src_x, dst_y, dst_x, oh, ow, ti) into int32 cupy arrays, launch kernel with grid=(K, ceil(max_oh/16), ceil(max_ow/16)) block=(1, 16, 16); atomicAdd into numer/denom - Bounds-safe: int64 offset arithmetic inside the kernel (int32 overflowed at tile_i*TCZ*tile_H*tile_W > 2.1 G) - Fall-back to fused-per-touch path if shapes don't match K=1 verified end-to-end: 516 bands, wall 18:51, per-band p50 1.19 s mean 3.43 s (vs fused-only p50 3.99 s mean 4.27 s). Known limitation: at K=2 workers/node, peer workers hit NVRM Xid 31 MMU fault after band 1 and hang. Root cause unresolved — likely CUDA-context race between the RawKernel and Dragon HSTA on the same node. Keep the flag off (or use K=1) until debugged. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
ElementwiseKernel that combines xp.maximum(denom, 1e-12) + xp.divide + xp.nan_to_num + .astype(uint16) into ONE launch per X-block, down from 4. Env-gated OPS_ASSEMBLE_FUSED_NORMCAST=1; only applies to uint16 output (LiveScreen case), float32 falls through to legacy path. Measured impact at K=2 batched: wall 14:51 → 14:41 (marginal). Bottleneck at K=2 is loader-bound (blosc decode CPU-limited at ~0.7 GB/s local), not GPU. Kernel wins buried under I/O wait. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The streaming fetch server (converted_stream.py:210) now emits multiple result chunks per batch_get, terminated by (host, None). Both consumers must drain until sentinel, but only the assemble loader had been updated. Estimate's _fetch_one() and load_many() were still doing a single resp_q.get() and crashed on KeyError once the response arrived in >1 chunk. _fetch_one: while-loop drain until (host, None), then read from merged dict. load_many: count sentinels seen (one per outstanding request) instead of message count. Estimate now runs cleanly under the streaming protocol. Measured wall: 231 s (3:51) at K=2, matching the pre-streaming baseline. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Thread an optional t_out (default 0, unchanged behavior) through
assemble_streaming and the write helpers (_write_zarr_block,
_write_zarr_block_direct, _d2h_and_submit_writes, _d2h_and_write_xblock)
so a multi-timepoint (N,C,Z,Y,X) store can be assembled one timepoint per
T index — direct-chunk key c/{t_out}/... and zarr T-slice. Counterpart to
the livescreen multi-timepoint assemble.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
assemble_streaming accepts an optional block_hook(y0, x0, norm_block), called per level-0 block right before the async D2H (so the hook can read the GPU block), plus block_hook.finalize() at the end. Enables downsampling each block on-GPU as it is produced — the livescreen level-1 pyramid fusion. No-op when block_hook is None: the level-0-only path is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
When each node's /tmp holds the full set of tiles its Y-workers will
need (via Y-partitioned convert), the shard_map / cross-node fetch
plumbing is dead weight. Setting OPS_ASSEMBLE_FORCE_LOCAL=1 makes
_split_by_host return (all_names, {}) so every tile is read from
local disk. Removes the last dependency on Dragon HSTA for
multi-node assemble.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
`_d2h_and_write_xblock` used to enqueue an async D2H on transfer_stream via `norm_gpu_x.get(out=pinned_buf)`, then immediately `del norm_gpu_x` and `xp.get_default_memory_pool().free_all_blocks()` — which cudaFrees the source pointer while the DMA is still in flight. Under K=1 (single worker per GPU) the GPU is idle enough that the copy always completes before free runs; under K=2 workers/GPU, compute contention delays the copy, the cudaFree wins the race, and pinned_buf.copy() reads back the freshly-allocated (all-zero) pinned buffer. Reproduced as ~5–7 empty output chunks per full LiveScreen assemble at K=2. Added a `transfer_stream.synchronize()` between the enqueued D2H and the source release. Verified: with fix + K=2 the empty-chunk symptom disappears (35 → 0 empty chunks in a controlled diff comparison of cold vs warm assembled outputs on the same shifts).
Previously ConvertedTileSource hardcoded the ``A/1`` well path when building tile-file paths and the sample glob. Multi-well acquisitions (6-well plates on the new LiveScreen instrument) can't be read without per-well parameterization. Add ``well_id: str = "A/1"`` kwarg (backwards-compatible default) to ``ConvertedTileSource`` and ``dragon_fetch_server``, thread it into the store-root path assembly. Callers now pass the target well when reading a per-well converted store. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…lation Three unrelated but overlapping fixes needed to run stitch reliably from inside our ops uv venv: - Silence arbol's per-call "Memory pool clearing not enabled!" print, which floods stdout with 334k lines on a full plate run (fires every dexp CupyBackend.__exit__). - CuPy runtime kernel compilation needs ``CUDA_PATH`` pointing at a toolkit with ``bin/nvcc``. Our uv .venv doesn't set it; resolve via nvidia.cuda_nvcc wheel then fall back to ``/hpc/apps/x86_64/cuda/<torch-cuda-ver>``. Must run before ``import cupy`` so the first kernel compile doesn't crash. - Add ``batched_phase_correlation`` for the GPU-batched estimate path in ``ops_process.processes.remap_estimate_gpu``. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…y' into feat/livescreen-multi-well-support
The xblock_scoped_gpu path casts float32 numer → int32 masks on the caller's compute stream (`with stream_b:` in _run_one_band_xblock), but _d2h_and_write_xblock reads the tensor on a separate transfer_stream. Without an inter-stream barrier, under GPU contention (K=2+ mask_band processes sharing one GPU via MPS) the D2H can begin BEFORE the astype kernel drains, copying pre-cast float32 numer bytes into the int32 pinned buffer and thence into the int32 chunk file. Result: ~1e9 "labels" from float32 bit-patterns of small-integer labels (e.g. 0x40800002 = 4.000001f = label 4). Verified on B/3 t5 at K=4 workers/GPU: 0/400 corrupt chunks (was ~9% chunks corrupt before). Cost: sub-ms per x-block (astype is fast). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…3 sharding_indexed
Sharded stitched-zarr output was slow (129s per A/1 t=42) AND corrupt (per-block
zarr.__setitem__ writes race on the same shard file — 32-thread writer + zarr's
read-modify-write per shard = last-writer-wins). This adds two coupled changes
so sharded output is correct AND ~2.3x faster than the initial sharded path,
closing to within ~40% of the un-sharded direct-file-write baseline.
1. Shard-batching (`_SHARD_BATCH`): allocate one in-memory buffer covering the
worker's full Y-range (one shard row tall) when the output array is
sharding_indexed. Each block-write copies into the buffer; nothing hits
zarr per-block. At end of Y-range, flush.
2. Direct-shard-encode (`_encode_shard_direct` + `_encode_chunks_to_bytes` +
`_assemble_and_write_shard`): bypass zarr's synchronous shard writer at
flush time. For each shard file:
- encode inner chunks in parallel via numcodecs (Blosc or Zstd auto-
detected from zarr.json)
- skip all-fill_value chunks with the (2^64-1, 2^64-1) empty sentinel
(huge win on sparse LiveScreen tissue: >50% chunks background)
- assemble shard bytes = concatenated encoded chunks || uint64 index ||
CRC32C over the index (matches zarr v3 sharding_indexed spec)
- single open("wb").write() per shard file
Correctness: 10/10 crops byte-identical vs the non-sharded (per-chunk-file,
direct-write) ground truth at same tp.
Perf (gpu-i-1, VAST-RDMA, A/1 t=42 assemble wall):
* sharded per-block writes: 210 s (corrupt)
* sharded batched + zarr flush: 129 s
* sharded batched + direct-enc: 90 s ← this change
* non-sharded (reference): 64 s
The remaining 26s gap vs non-sharded is fundamental to sharding_indexed's
chunk-granular design: 50k encode calls per Y-range on 0.5MB chunks vs 1500
calls on 4MB tiles. Documented in memory notes; nvCOMP GPU-side zstd is the
next lever (~2-3s encode/Y-range would put us at parity).
Also (unrelated cleanup on this branch):
* fix(assemble): sync compute stream before yield to avoid mask-junk race
Env knobs for tuning: OPS_SHARD_BATCH (default 1), OPS_SHARD_DIRECT_ENCODE
(default 1), OPS_ENCODE_WORKERS (default 32), OPS_FLUSH_WORKERS (default 16).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The direct-encode flush maps a band to shard row y_lo//sh_h and places it at intra-shard rows [0:y_hi-y_lo]. That is correct only when the band starts on a shard boundary and spans at most one shard row. A mis-aligned band makes two workers target the same shard file (last writer wins, dropping chunks) with the CRC still valid -> silently wrong data; a >1-row band broadcast-crashes cryptically. Assert both preconditions loudly. The ops_process partition guarantees them; this makes the cross-repo coupling explicit.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacks on #16 — please review that first. Only the final commit
(
46f9e8d, +17 lines institch/stitch/assemble.py) is new here; everythingelse in the diff is inherited from #16 and will disappear once that merges.
What this adds
assemble_streaming's direct-shard-encode path writes a whole shard file with asingle
open("wb").write(). It maps a band to shard rowy_lo // sh_handplaces the data at intra-shard rows
[0 : y_hi - y_lo].That is only correct when the band starts on a shard boundary and spans at
most one shard row. If the caller's Y-partition violates either:
drops the other's chunks, and the CRC still validates, so it reads back
as plausible-but-wrong data;
The
ops_processpartition (assemble_gpu_multi) does guarantee both, but thatcoupling is implicit and across repos. This asserts the two preconditions with
messages that point at the actual cause, so a future partition change fails
loudly instead of corrupting output.
No behavioural change when the preconditions hold.
Provenance
Commit authored by Gav Sturm; opened on his behalf from the biohub-hpc fork.
🤖 Generated with Claude Code