Zero-copy audio blob reads + mp4/ogg seek-index fixes - #84
Open
jpc wants to merge 6 commits into
Open
Conversation
pyarrow's `col[j]` on a binary array materializes the element into a BinaryScalar: Array::GetScalar -> GetString -> std::string(GetView(i)) COPIES the bytes into a fresh heap allocation (confirmed in arrow C++ source and by address check). On cold mmap'd 1.8 GB shards this faults + copies the ENTIRE episode blob (measured 100-300 ms, ~1 GB/s memcpy even warm) just to decode an 8 s crop. Hand the audio decoder a zero-copy `pa.BufferReader` over `values.slice(offsets[j], offsets[j+1])` instead, so it faults only the pages it actually seeks/reads. Byte-identical output (verified); get_audio 248 ms -> 7.85 ms (32x), dataloader throughput +7.5x. Applies only to audio columns (npy/pyd read fully anyway); other columns keep the scalar path. Mirrors pupyarrow LazyBinaryArray._get_single used by S3/Modal shards. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
_apply_seek_index used to call add_seek_points(all_points) for every non-mp3 codec. For mp4/mov (aac/alac) the moov already carries a full sample table, and av_add_index_entry is an O(n) sorted insert against its millions of native entries -> O(n*m): a 36 h aac audiobook took 41 s to seed 16 k points (and could hang the dataloader). Skip mp4 entirely (native seeking is already fast; the index is still EXTRACTED for the block-cache/backblaze prefetch, just not fed to ffmpeg). For ogg/vorbis (no native seek table) seed INCREMENTALLY: set_seed_index stores the index and AudioDecoder._seed_around adds only a small window of points bracketing each requested seek target as segments are read, so per-seek cost is O(window) regardless of episode length and accumulates across a seq_per_seek run. Byte-identical decode verified. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Wraps the hot-path phases (shard_open, batch_read, blob_decode, vad_read, src_key_lookup, get_audio, decoder_open, audio_decode, build_packet_index, key_verify) in a `record()` context manager that accumulates per-phase [count, total_s, max_s] per process and periodically writes $WSDS_TIMING_OUT.<pid> as JSON (survives DataLoader worker kills). This is how the two fixes in this PR were located. Disabled by default: with WSDS_TIMING unset, `record()` returns a shared no-op context manager. Measured overhead: - micro: 132 ns/call (vs 847 ns for a naive @contextmanager) - end-to-end: 27.8 crops/s (disabled) vs 27.3 crops/s (code removed) — within noise; the ~130 ns/call is nothing next to the ms-scale read/ decode each call wraps. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ted lookup Review feedback: avoid the allocations in set_seed_index/_seed_around. The seek-index generator emits points in blob-scan order, i.e. ascending pts (verified across podcasts-en + spotify episodes), so the per-episode `sorted(range(n), key=lambda)` was redundant — store `pts`/`positions` directly as float64/int64 ndarrays (one asarray each, no Python sort or list comps). Replace the per-seek `bisect` with `np.searchsorted`. Behavior is identical (same 8-point windows accumulate; decode bit-identical). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…hrough attach The seek index is a pyarrow struct<... seek_pts: list<u32>, seek_pos: list<u64>>, but reading it went struct -> as_py() (Python dict of 15k boxed ints) -> list-comp in the consumer -> back to numpy: ~11 ms/episode on a 15k-point index. Add `WSSample.get_raw(field)` / `get_sample(..., raw=True)` that returns the raw pyarrow scalar, skipping the as_py()/decode conversion (which exists mainly to keep external callers unsurprised). Trusted internal consumers pull list values straight to numpy zero-copy. `WSAudioEpisode.set_seek_index` now keeps arrays as-is (no `[int(p) for p in ...]`); only the mp3 packet-index path coerces per element (SimpleNamespace + humecodec need Python scalars), vorbis stays numpy. Read path measured 323x faster (10863 -> 34 us) on a 15k-point index; decode bit-identical, seek index verified attaching on 505/505 mp3 crops. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
pop_chunks() returns None when a fill produced no decoded audio (e.g. a mid-stream flac seek failure: "read_timestamp() failed in the middle"). get_samples_played_in_range dereferenced chunk.pts unconditionally, so a single bad seek killed the DataLoader worker with AttributeError and, via NCCL, the whole DDP training run (two 1B runs died at ~3k steps; the peers' "invalid access of peer GPU memory over nvlink" errors were the symptom of the dead rank, not a hardware fault). - skip None chunks while more input remains, with a 65536-empty-pop cap so a wedged decoder raises instead of hanging the worker - keep popping after fill_buffer() reports EOF until the buffer is drained: chunks are 1s each, so the old exit-after-one-pop could truncate the tail of reads that reach end of stream - raise ValueError (a data error callers already skip) when a range decodes to zero chunks Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
1 participant
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.
Three localized changes to the audio-fetch hot path, found while profiling a random-seek dataloader over multi-GB podcast/audiobook shards. Stacked on #83.
1. Zero-copy audio blob reads (
ws_shard)pyarrow
col[j]on a binary column materializes the element into aBinaryScalar—Array::GetScalar → GetString → std::string(GetView(i))copies the bytes into a fresh heap allocation (confirmed in the arrow C++ source and by an address check; a real ~1 GB/s memcpy even warm). On cold mmap'd 1.8 GB shards this faults + copies the entire episode blob (100–300 ms) just to decode an 8 s crop.Fix: hand the decoder a zero-copy
pa.BufferReaderovervalues.slice(offsets[j], offsets[j+1]), so it faults only the pages it seeks/reads. Audio-only (npy/pyd read fully anyway); mirrors pupyarrow'sLazyBinaryArray._get_singleused by the S3/Modal shards.get_audio248 ms → 7.85 ms (32×), dataloader throughput +7.5×2. Seek-index attach: skip mp4, seed ogg incrementally (
ws_audio,audio_codec)_apply_seek_indexcalledadd_seek_points(all_points)for every non-mp3 codec. For mp4/mov (aac/alac) the moov already has a full sample table, andav_add_index_entryis an O(n) sorted insert against its millions of native entries → O(n·m): a 36 h aac audiobook took 41 s to seed 16k points. Now mp4 is skipped (native seeking is already fast; the index is still extracted for the block-cache/backblaze prefetch, just not fed to ffmpeg).For ogg/vorbis (no native seek table) seeding is now incremental:
set_seed_indexstores the index andAudioDecoder._seed_aroundadds only a small window of points bracketing each requested seek target as segments are read — O(window) per seek regardless of episode length. Byte-identical decode verified.3.
WSDS_TIMINGphase profiler (_timing.py)Env-gated per-phase profiler for the audio-fetch hot path — how the two fixes above were located. Disabled by default:
record()returns a shared no-op context manager. Measured overhead: 132 ns/call (micro) and 27.8 vs 27.3 crops/s end-to-end with vs without the code (within noise).Notes
WSAudioSegment's underlyingWSAudioEpisode, not the segment — the bug that silently dropped 99.9% of crops) lives in the downstreambestrqrepo.🤖 Generated with Claude Code