Zero-copy audio blob reads + mp4/ogg seek-index fixes

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.<pid>`` as JSON so stats
+survive even if DataLoader worker processes are killed. Aggregate by summing all
+``<out>.*`` 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

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

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

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

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