perf(streaming): sub-range parallel fetch + non-blocking pool replacement

The 64-vs-128-worker benchmark pair proved a per-host throughput ceiling
(~270 MiB/s) on the HF bucket path: doubling connections exactly halved
per-connection speed (4.8 -> 2.2 MiB/s) and left the aggregate flat,
while per-episode latency doubled (5.7s -> 12s) and keep-up worsened.
Steady-state demand (148 MiB/s) is well below the ceiling; the keep-up
misses come entirely from consumer stalls (refill_wait 14-19s of ~84s):
the sim blocks the training hot path on ensure_ready() for the FIFO-head
replacement while episodes take 5.7-12s to arrive.

Two fixes:

- Non-blocking replacements: EpisodeByteCache.is_ready() (all cameras
  cached or futures done, no blocking) and the stream sim now swaps a
  replacement only when it is already resident, deferring otherwise;
  fetch capacity (~2x demand) repays the debt on later batches. A
  deferred_swaps metric is reported.
- Sub-range parallel fetch (native-http): --range-subranges N splits one
  camera GET into N concurrent sub-range GETs. Under a per-host ceiling
  this adds no bandwidth but divides per-episode latency by ~N. Keep
  workers x subranges near the ~64-connection saturation point (e.g.
  --workers 16 --range-subranges 4).

Verified: sub-range span math + order-preserving concat and is_ready
semantics (unit-level, network stubbed).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

scripts/bench_episode_byte_cache.py

@@ -70,7 +70,20 @@ def parse_args() -> argparse.Namespace:
         help="Limit manifest construction to the first N episodes for local smoke tests.",
     )
     parser.add_argument("--pool-size", type=int, default=16)
-    parser.add_argument("--workers", type=int, default=8)
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=8,
+        help="Concurrent camera-fetch jobs. Total connections ~= workers x range-subranges; "
+        "the HF bucket path saturates around 64 connections per host, so keep the product near 64.",
+    )
+    parser.add_argument(
+        "--range-subranges",
+        type=int,
+        default=1,
+        help="Split each camera byte-range GET into N concurrent sub-range GETs (native-http only). "
+        "Divides per-episode latency by ~N under the per-host throughput ceiling.",
+    )
     parser.add_argument(
         "--native-http-connections",
         type=int,
@@ -392,10 +405,19 @@ def run_pool_stream_simulation(
     decoded_samples: list[tuple[int, float]] = []
     start = time.perf_counter()
 
+    deferred_swaps = 0
+
     def consume_ready_replacement() -> bool:
-        nonlocal refill_wait_s, replacement_count
+        nonlocal refill_wait_s, replacement_count, deferred_swaps
         if not pending:
             return False
+        # Non-blocking: only swap when the head replacement is fully resident. Blocking here
+        # stalls the training hot path on remote fetch latency (head-of-line); deferring lets
+        # the fetch pipeline (capacity ~2x demand) catch up while training continues on the
+        # current pool. The replacement debt is repaid on subsequent batches.
+        if not cache.is_ready(pending[0]):
+            deferred_swaps += 1
+            return False
         new_ep = pending.pop(0)
         wait_start = time.perf_counter()
         cache.ensure_ready(new_ep)
@@ -463,6 +485,7 @@ def run_pool_stream_simulation(
         "deadline_miss_s": deadline_miss_s,
         "replacements": float(replacement_count),
         "replacement_episodes_s": replacement_count / elapsed if elapsed > 0 else 0.0,
+        "deferred_swaps": float(deferred_swaps),
         "samples_per_episode": float(samples_per_episode),
         "prefetch_episodes": float(prefetch_episodes),
         "batch_size": float(batch_size),
@@ -722,6 +745,7 @@ def run_fetch_pool(
         native_http_connections=args.native_http_connections,
         native_http_timeout=args.native_http_timeout,
         native_http_retries=args.native_http_retries,
+        native_http_subranges=args.range_subranges,
         open_decoders=False,
     ) as cache:
         elapsed = _fill_cache(cache, episodes, progress_interval=args.progress_interval)

src/lerobot/datasets/episode_video_streaming.py

@@ -366,11 +366,24 @@ class NativeHTTPRangeFetcher:
         max_connections: int = 32,
         timeout: float = 60.0,
         max_retries: int = 4,
+        subrange_parts: int = 1,
+        subrange_min_bytes: int = 8 * 1024 * 1024,
     ):
         self.data_root = str(data_root).rstrip("/")
         if not self.data_root.startswith("hf://"):
             raise ValueError("NativeHTTPRangeFetcher only supports hf:// roots")
         self.max_retries = max_retries
+        # Sub-range parallelism: split one large GET into `subrange_parts` concurrent GETs.
+        # Under a per-host throughput ceiling this adds no aggregate bandwidth, but divides
+        # per-request latency by ~parts - keep (in-flight jobs x parts) near the ceiling's
+        # connection sweet spot (~64 on the observed HF bucket path) rather than raising both.
+        self.subrange_parts = max(1, subrange_parts)
+        self.subrange_min_bytes = max(1, subrange_min_bytes)
+        self._subrange_pool = (
+            ThreadPoolExecutor(max_workers=max_connections, thread_name_prefix="subrange")
+            if self.subrange_parts > 1
+            else None
+        )
         self.api = HfApi()
         self.fs: HfFileSystem | None = None
         self._bucket_id: str | None = None
@@ -526,6 +539,21 @@ class NativeHTTPRangeFetcher:
             response.close()
 
     def read_range(self, relative_path: str, offset: int, length: int) -> bytes:
+        parts = self.subrange_parts
+        if self._subrange_pool is None or parts <= 1 or length < 2 * self.subrange_min_bytes:
+            return self._read_range_single(relative_path, offset, length)
+        parts = min(parts, max(1, length // self.subrange_min_bytes))
+        if parts <= 1:
+            return self._read_range_single(relative_path, offset, length)
+        step = (length + parts - 1) // parts
+        spans = [(offset + i * step, min(step, length - i * step)) for i in range(parts)]
+        futures = [
+            self._subrange_pool.submit(self._read_range_single, relative_path, span_off, span_len)
+            for span_off, span_len in spans
+        ]
+        return b"".join(future.result() for future in futures)
+
+    def _read_range_single(self, relative_path: str, offset: int, length: int) -> bytes:
         resolve_start = time.perf_counter()
         resolved = self._resolve_url(relative_path)
         source = self._source_url(relative_path)
@@ -693,6 +721,8 @@ class NativeHTTPRangeFetcher:
             return dict(self._timing_totals)
 
     def close(self) -> None:
+        if self._subrange_pool is not None:
+            self._subrange_pool.shutdown(wait=False, cancel_futures=True)
         self.client.close()
 
 
@@ -704,6 +734,7 @@ def make_range_fetcher(
     native_http_connections: int | None = None,
     native_http_timeout: float = 60.0,
     native_http_retries: int = 4,
+    native_http_subranges: int = 1,
 ):
     if range_backend == "fsspec":
         return ThreadLocalRangeFetcher(data_root)
@@ -714,6 +745,7 @@ def make_range_fetcher(
             max_connections=max_connections,
             timeout=native_http_timeout,
             max_retries=native_http_retries,
+            subrange_parts=native_http_subranges,
         )
     raise ValueError(f"Unknown range backend: {range_backend}")
 
@@ -961,6 +993,7 @@ class EpisodeByteCache:
         native_http_connections: int | None = None,
         native_http_timeout: float = 60.0,
         native_http_retries: int = 4,
+        native_http_subranges: int = 1,
         open_decoders: bool = True,
     ):
         self.manifest = manifest
@@ -971,6 +1004,7 @@ class EpisodeByteCache:
             native_http_connections=native_http_connections,
             native_http_timeout=native_http_timeout,
             native_http_retries=native_http_retries,
+            native_http_subranges=native_http_subranges,
         )
         self.byte_budget = byte_budget
         self.open_decoders = open_decoders
@@ -1009,6 +1043,22 @@ class EpisodeByteCache:
         for camera_key in self.manifest.video_keys:
             self.get_bytes(episode_index, camera_key)
 
+    def is_ready(self, episode_index: int) -> bool:
+        """Non-blocking: True when every camera of the episode is fetched (cached or future done).
+
+        Lets a consumer swap in replacements only when they are already resident, instead of
+        blocking the training hot path on a remote fetch (head-of-line stall).
+        """
+        for camera_key in self.manifest.video_keys:
+            key = (episode_index, camera_key)
+            with self._lock:
+                if key in self._cache:
+                    continue
+                future = self._futures.get(key)
+            if future is None or not future.done():
+                return False
+        return True
+
     def get_bytes(self, episode_index: int, camera_key: str) -> bytes:
         return self._get_entry(episode_index, camera_key)["bytes"]