diff --git a/wsds/_timing.py b/wsds/_timing.py new file mode 100644 index 0000000..47ecf28 --- /dev/null +++ b/wsds/_timing.py @@ -0,0 +1,83 @@ +"""Env-gated phase timing for the audio-fetch hot path. + +Disabled by default (unset ``WSDS_TIMING``): ``record()`` returns a shared no-op +context manager (~130 ns/call — negligible next to the ms-scale read/decode it +wraps). When enabled, each process accumulates per-phase ``[count, total_s, +max_s]`` and periodically writes ``$WSDS_TIMING_OUT.`` as JSON so stats +survive even if DataLoader worker processes are killed. Aggregate by summing all +``.*`` files. + + from ._timing import record + with record("blob_decode"): + ... +""" +import atexit +import json +import os +import time + +_ENABLED = bool(os.environ.get("WSDS_TIMING")) +_OUT = os.environ.get("WSDS_TIMING_OUT") +_stats: dict[str, list] = {} # phase -> [count, total_s, max_s] +_n = 0 + + +class _Noop: + __slots__ = () + + def __enter__(self): + return None + + def __exit__(self, *exc): + return False + + +_NOOP = _Noop() + + +class _Timer: + __slots__ = ("phase", "_t") + + def __init__(self, phase): + self.phase = phase + + def __enter__(self): + self._t = time.perf_counter() + + def __exit__(self, *exc): + dt = time.perf_counter() - self._t + s = _stats.get(self.phase) + if s is None: + _stats[self.phase] = [1, dt, dt] + else: + s[0] += 1 + s[1] += dt + if dt > s[2]: + s[2] = dt + global _n + _n += 1 + if _n % 500 == 0: + _flush() + return False + + +def _flush(): + if _OUT and _stats: + tmp = f"{_OUT}.{os.getpid()}.tmp" + with open(tmp, "w") as f: + json.dump(_stats, f) + os.replace(tmp, f"{_OUT}.{os.getpid()}") + + +if _ENABLED: + def record(phase): + return _Timer(phase) + + atexit.register(_flush) +else: + def record(phase): # zero-work fast path: shared no-op context manager + return _NOOP + + +def get_stats(): + return _stats diff --git a/wsds/audio_codec.py b/wsds/audio_codec.py index 62eb9d4..86e1b64 100644 --- a/wsds/audio_codec.py +++ b/wsds/audio_codec.py @@ -13,6 +13,7 @@ import traceback import typing +import numpy as np import pyarrow as pa @@ -36,14 +37,25 @@ def __init__(self, reader, metadata, sample_rate, codec_delay=0): # byte-offset seek with our own index is much faster. self._use_byte_index = codec_name in ('mp3', 'mp2', 'mp1') self._packet_index = None + # On-demand demuxer-index seeding (ogg/vorbis, which has no native seek + # table). Rather than add every episode point up front (each + # av_add_index_entry is an O(n) sorted insert -> O(n*m)), we keep the + # full index and add only a small window of points around each requested + # seek target as segments are read. See set_seed_index / _seed_around. + self._seed_pts = None + self._seed_positions = None + self._seed_added = None + self._seed_window = 4 def _build_index(self): """Build a sparse packet index for byte-offset seeking.""" if self._packet_index is not None: return + from ._timing import record try: - idx = self.reader.build_packet_index( - self.reader.default_audio_stream, 128 * 1024) + with record("build_packet_index"): + idx = self.reader.build_packet_index( + self.reader.default_audio_stream, 128 * 1024) if idx and len(idx) > 1: self._packet_index = idx except Exception: @@ -97,14 +109,35 @@ def get_samples_played_in_range(self, tstart=0, tend=None, margin=.25): if index_pts is None: # Fall back to timestamp seek (or read from start) seek_target = 0.0 if read_from_start else max(0, tstart - margin) + if not read_from_start: + # Seed just the AVIndexEntry points around this target (ogg/vorbis) + # so the seek brackets in ~1 read; accumulates across seq reads. + self._seed_around(seek_target) self.reader.seek(seek_target, "key") chunks = [] - more_data = True - while more_data: - if self.reader.fill_buffer() == 1: - more_data = False + eof = False + empty_pops = 0 + while True: + if not eof and self.reader.fill_buffer() == 1: + eof = True (chunk,) = self.reader.pop_chunks() + if chunk is None: + if eof: + break # buffer fully drained + # A fill produced no decoded audio yet: mid-stream seek/decode + # hiccup (e.g. flac "read_timestamp() failed in the middle"). + # Skip it — dereferencing None.pts kills the DataLoader worker + # and with it the whole DDP run. The cap guards against a + # wedged no-progress decoder: a hung worker is worse than a + # crash. + empty_pops += 1 + if empty_pops > 65536: + raise ValueError( + f"decoder made no progress after {empty_pops} empty " + f"pops (codec={self.metadata.codec}); wedged stream?") + continue + empty_pops = 0 chunks.append(chunk) if tend is not None: chunk_end_pts = chunk.pts + chunk.shape[0] / self.sample_rate @@ -115,6 +148,13 @@ def get_samples_played_in_range(self, tstart=0, tend=None, margin=.25): if chunk_end_pts > tend + margin: break + if not chunks: + # Data-level failure (bad seek target / corrupt stream), not a bug: + # raise ValueError so sample-skipping callers can drop this read. + raise ValueError( + f"decoder produced no samples for range [{tstart:.3f}, {tend}] " + f"(codec={self.metadata.codec}, read_from_start={read_from_start})") + # Determine the reference PTS for trimming if read_from_start: chunk0_pts = 0.0 @@ -166,6 +206,37 @@ def add_seek_points(self, positions, pts_seconds): fn([int(p) for p in positions], [float(t) for t in pts_seconds]) return True + def set_seed_index(self, positions, pts_seconds): + """Store a precomputed (byte-position, pts) index for ON-DEMAND demuxer + seeding. Points are added lazily in a small window around each seek + target (see _seed_around) rather than all at once, so the per-seek cost + stays O(window) even for multi-hour episodes with tens of thousands of + points. Use only for formats WITHOUT a native seek table (ogg/vorbis); + never for mp4/mov (the moov already indexes them and av_add_index_entry + is O(n) per insert against its millions of native entries). + + `pts_seconds` MUST be ascending — the seek-index generator emits points + in blob-scan order (ascending pts) — so we store them as-is (no sort) and + binary-search with np.searchsorted in _seed_around.""" + self._seed_pts = np.asarray(pts_seconds, dtype=np.float64) + self._seed_positions = np.asarray(positions, dtype=np.int64) + self._seed_added = set() + + def _seed_around(self, target_time): + """Add the few seed-index points bracketing target_time to the demuxer's + seek index (idempotent per point, accumulates across seeks). No-op unless + a seed index was set via set_seed_index.""" + if self._seed_pts is None or self._seed_pts.size == 0: + return + i = int(np.searchsorted(self._seed_pts, target_time, side="right")) + lo = max(0, i - self._seed_window) + hi = min(self._seed_pts.size, i + self._seed_window) + sel = [k for k in range(lo, hi) if k not in self._seed_added] + if sel: + self._seed_added.update(sel) + self.add_seek_points(self._seed_positions[sel].tolist(), + self._seed_pts[sel].tolist()) + def _create_reader_humecodec(src, buffer_size): from humecodec import MediaDecoder diff --git a/wsds/ws_audio.py b/wsds/ws_audio.py index 6b22a2f..4c1794a 100644 --- a/wsds/ws_audio.py +++ b/wsds/ws_audio.py @@ -48,33 +48,52 @@ def set_seek_index(self, positions, pts_seconds): scan: `build_packet_index` for mp3/mp2/mp1 (which reads the WHOLE episode — ~200-400ms for a long spotify mp3, the dominant deep-seek cost) and the Ogg/Vorbis demuxer bisection. `positions` are byte offsets in the audio - blob (blob-relative); `pts_seconds` the matching timestamps. Applied on + blob (blob-relative); `pts_seconds` the matching timestamps. May be numpy + arrays — kept as-is for zero-copy; the per-codec consumers coerce to + Python scalars only where they must (the mp3 packet index). Applied on the next decoder (re)creation.""" - self._seek_index = ([int(p) for p in positions], [float(t) for t in pts_seconds]) + self._seek_index = (positions, pts_seconds) if self._decoder is not None: self._apply_seek_index() def _apply_seek_index(self): - if not self._seek_index or self._decoder is None: + if self._seek_index is None or self._decoder is None: return pos, pts = self._seek_index d = self._decoder if getattr(d, "_use_byte_index", False): # mp3/mp2/mp1 byte-index path from types import SimpleNamespace - d._packet_index = [SimpleNamespace(pts_seconds=t, pos=p) for p, t in zip(pos, pts)] - else: # vorbis/etc: seed the demuxer index (humecodec>=0.8) - fn = getattr(d, "add_seek_points", None) - if fn is not None: - fn(pos, pts) + # SimpleNamespace + humecodec need Python scalars, so coerce per element + # here (unavoidable for the byte index); vorbis stays fully numpy below. + d._packet_index = [SimpleNamespace(pts_seconds=float(t), pos=int(p)) for p, t in zip(pos, pts)] + else: + # ogg/vorbis has no native seek table -> seed the demuxer's + # AVIndexEntry list (humecodec>=0.8) to skip its interpolating + # bisection. Seeding is INCREMENTAL: set_seed_index just stores the + # index, and a small window of points near each requested seek target + # is added on demand (AudioDecoder._seed_around), so per-seek cost is + # O(window) regardless of episode length. + # + # mp4/mov (aac/alac) are SKIPPED: the moov already carries a full + # sample table, so seeding is redundant AND each av_add_index_entry is + # an O(n) sorted insert against its millions of native entries -> a + # 36h aac took 41s for 16k points. Native mp4 seeking is already fast. + # (The seek index is still EXTRACTED for mp4 — its absolute offsets + # feed the block-cache/backblaze prefetch — we just don't feed ffmpeg.) + codec = getattr(getattr(d, "metadata", None), "codec", "") or "" + if codec in ("vorbis", "opus") and hasattr(d, "set_seed_index"): + d.set_seed_index(pos, pts) def get_decoder(self, sample_rate=None): """Lazily creates/caches decoder via audio_codec.create_decoder().""" requested_sr = sample_rate or (self._decoder and self._decoder.metadata.sample_rate) if self._decoder is None or requested_sr != self._sample_rate: - self.src.seek(0) - self._decoder = create_decoder(self.src, sample_rate=sample_rate) - self._sample_rate = sample_rate or self._decoder.metadata.sample_rate - self._apply_seek_index() + from ._timing import record + with record("decoder_open"): # create_decoder = find_stream_info + self.src.seek(0) + self._decoder = create_decoder(self.src, sample_rate=sample_rate) + self._sample_rate = sample_rate or self._decoder.metadata.sample_rate + self._apply_seek_index() return self._decoder, self._sample_rate @property @@ -88,8 +107,10 @@ def sample_rate(self): return sr def read_segment(self, start=0, end=None, sample_rate=None): + from ._timing import record decoder, sample_rate = self.get_decoder(sample_rate) - samples = decoder.get_samples_played_in_range(start, end) + with record("audio_decode"): # seek + decode the segment + samples = decoder.get_samples_played_in_range(start, end) if hasattr(samples, "data"): samples = samples.data samples.sample_rate = sample_rate diff --git a/wsds/ws_dataset.py b/wsds/ws_dataset.py index cf96590..33ef1cb 100644 --- a/wsds/ws_dataset.py +++ b/wsds/ws_dataset.py @@ -600,12 +600,12 @@ def get_shard(self, column_dir, shard_ref): self._open_shards[shard_dir] = shard return shard - def get_sample(self, shard_ref, field, offset): + def get_sample(self, shard_ref, field, offset, raw=False): alternatives = self.fields[field] last_err = None for column_dir, column in alternatives: try: - return self.get_shard(column_dir, shard_ref).get_sample(column, offset) + return self.get_shard(column_dir, shard_ref).get_sample(column, offset, raw=raw) except WSShardMissingError as e: last_err = e continue diff --git a/wsds/ws_sample.py b/wsds/ws_sample.py index bfd7434..3859bd5 100644 --- a/wsds/ws_sample.py +++ b/wsds/ws_sample.py @@ -3,6 +3,7 @@ from .utils import WSShardMissingError, validate_shards from .ws_decode import get_audio as _get_audio +from ._timing import record if TYPE_CHECKING: from .ws_dataset import WSDataset @@ -47,7 +48,8 @@ def _verify_key_for_field(self, field: str): # Get __key__ from this column_dir try: - key = self.dataset.get_shard(column_dir, self.shard_ref).get_sample("__key__", self.offset) + with record("key_verify"): # cross-column-dir __key__ consistency check + key = self.dataset.get_shard(column_dir, self.shard_ref).get_sample("__key__", self.offset) except (WSShardMissingError, KeyError): # Can't verify if shard or key is missing self._verified_column_dirs.add(column_dir) @@ -72,6 +74,17 @@ def __getitem__(self, field): self._verify_key_for_field(field) return self.dataset.get_sample(self.shard_ref, field, self.offset) + def get_raw(self, field): + """Return the RAW pyarrow scalar for `field`, skipping the as_py()/decode + conversion that `__getitem__` applies. For trusted internal consumers that + want zero-copy numeric/struct access (e.g. `.values.to_numpy()`) instead + of Python objects — the default conversion exists mainly to keep external + callers unsurprised.""" + if field in self.overrides: + return self.overrides[field] + self._verify_key_for_field(field) + return self.dataset.get_sample(self.shard_ref, field, self.offset, raw=True) + def __setitem__(self, field, value): self.overrides[field] = value diff --git a/wsds/ws_shard.py b/wsds/ws_shard.py index 4d99b6c..c79b511 100644 --- a/wsds/ws_shard.py +++ b/wsds/ws_shard.py @@ -8,8 +8,30 @@ from .utils import WSShardMissingError from .ws_audio import WSAudioEpisode, WSAudioSegment -from .ws_decode import decode_sample, decode_arr +from .ws_decode import decode_sample, decode_arr, AUDIO_FILE_KEYS from .ws_sample import WSSample +from ._timing import record + +import struct + + +def _zero_copy_blob_reader(arr, j): + """Zero-copy seekable file-like over element ``j`` of a (large_)binary array. + + Reads only ``offsets[j:j+2]`` and slices the values buffer, WITHOUT + materializing the element into a pyarrow scalar. ``col[j]`` copies/faults the + ENTIRE blob (measured 100-300 ms on cold 1.8 GB mmap'd shards) even when the + decoder only seeks to a small region; this hands the decoder an mmap-backed + view so it faults just the pages it actually reads. Returns None if null.""" + vbuf, offbuf, valbuf = arr.buffers() # [validity, offsets, values] + idx = arr.offset + j + if vbuf is not None and not (memoryview(vbuf)[idx >> 3] & (1 << (idx & 7))): + return None + w, fmt = (8, " typing.Any: + def get_sample(self, column: str, offset: int, raw: bool = False) -> typing.Any: if self._data is None or offset < self._start or offset >= self._end: i = offset // self.batch_size if i >= self.reader.num_record_batches: raise IndexError(f"{offset} is out of range for shard {self.fname}") - self._data = self.reader.get_batch(i) + with record("batch_read"): + self._data = self.reader.get_batch(i) if i < self.reader.num_record_batches - 1: if self._data.num_rows < self.batch_size: raise ValueError( @@ -82,14 +106,29 @@ def get_sample(self, column: str, offset: int) -> typing.Any: raise IndexError(f"{offset} is out of range for shard {self.fname}") if self._data.schema.get_field_index(column) == -1: raise KeyError(f"column {column} not found in shard {self.fname}") + col_type = self._data.schema.field(column).type + ext = column.rsplit(".", 1)[-1] if "." in column else column + # Audio columns: hand the decoder a ZERO-COPY reader over just this + # element's bytes (mmap-backed) instead of materializing the whole blob + # via col[j]. Only audio benefits (it seeks; npy/pyd read fully anyway). + if not raw and ext in AUDIO_FILE_KEYS and (pa.types.is_binary(col_type) or pa.types.is_large_binary(col_type)): + with record("blob_decode"): + fd = _zero_copy_blob_reader(self._data.column(column), j) + return None if fd is None else WSAudioEpisode(fd) data = self._data[column][j] if not data.is_valid: return None # Return None for any null pyarrow scalars - col_type = self._data.schema.field(column).type + if raw: + # Internal zero-copy path: return the raw pyarrow scalar, skipping the + # as_py()/decode_sample conversion (which is there mainly to keep + # external callers unsurprised). Trusted internal consumers do their + # own .values.to_numpy() for zero-copy numeric/struct access. + return data try: if pa.types.is_binary(col_type) or pa.types.is_large_binary(col_type): - return decode_sample(column, data) - if (column.rsplit(".", 1)[-1] if "." in column else "") == "arr": + with record("blob_decode"): + return decode_sample(column, data) + if ext == "arr": return decode_arr(data, col_type) # native variable-length array return data.as_py(maps_as_pydicts="strict") except Exception as e: @@ -135,17 +174,20 @@ def from_link(cls, link, dataset, shard_ref): return cls(shard_ref, source_dataset, dataset, link["vad_column"]) def get_timestamps(self, segment_offset): - return self._source_sample[self.vad_column][segment_offset] + with record("vad_read"): + return self._source_sample[self.vad_column][segment_offset] - def get_sample(self, _column, offset): + def get_sample(self, _column, offset, raw: bool = False): file_name, segment_offset = self.derived_dataset.parse_key( WSSample(self.derived_dataset, self.shard_ref, offset)["__key__"] ) if self._source_file_name != file_name: - self._source_sample = self.source_dataset[file_name] + with record("src_key_lookup"): + self._source_sample = self.source_dataset[file_name] try: - self._source_reader = self._source_sample.get_audio() + with record("get_audio"): + self._source_reader = self._source_sample.get_audio() except KeyError: raise WSShardMissingError("no audio shards found") self._source_file_name = file_name