From be64ded80f7b03d68a382e03c3925a75d18bdb18 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Fri, 3 Jul 2026 14:05:46 +0200 Subject: [PATCH] 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 --- scripts/bench_episode_byte_cache.py | 28 ++++++++++- .../datasets/episode_video_streaming.py | 50 +++++++++++++++++++ 2 files changed, 76 insertions(+), 2 deletions(-) diff --git a/scripts/bench_episode_byte_cache.py b/scripts/bench_episode_byte_cache.py index 919a7a830..c8dc65e74 100644 --- a/scripts/bench_episode_byte_cache.py +++ b/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) diff --git a/src/lerobot/datasets/episode_video_streaming.py b/src/lerobot/datasets/episode_video_streaming.py index 79800ea51..f48d578ec 100644 --- a/src/lerobot/datasets/episode_video_streaming.py +++ b/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"]