From 0f8257443c27cdfa0b0b9343c66d32bc20d7d1d8 Mon Sep 17 00:00:00 2001
From: Pepijn <pepijn@huggingface.co>
Date: Mon, 22 Jun 2026 11:35:28 +0200
Subject: [PATCH] Retry native HTTP timeout statuses

---
 scripts/bench_episode_byte_cache.py           |  8 +++++++
 .../datasets/episode_video_streaming.py       | 22 ++++++++++++++++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/scripts/bench_episode_byte_cache.py b/scripts/bench_episode_byte_cache.py
index 86e7105f4..c7093ed5c 100644
--- a/scripts/bench_episode_byte_cache.py
+++ b/scripts/bench_episode_byte_cache.py
@@ -640,6 +640,14 @@ def _print_range_timing_summary(fetch_pool: dict[str, float]) -> None:
     if exception_counts:
         summary = ", ".join(f"{name}={count:.0f}" for name, count in sorted(exception_counts.items()))
         print(f"| http exception counts | {summary} |")
+    failed_status_counts = {
+        key.removeprefix("range_failed_status_"): value
+        for key, value in fetch_pool.items()
+        if key.startswith("range_failed_status_")
+    }
+    if failed_status_counts:
+        summary = ", ".join(f"{status}={count:.0f}" for status, count in sorted(failed_status_counts.items()))
+        print(f"| http failed status counts | {summary} |")
     status_counts = {
         key.removeprefix("range_status_"): value
         for key, value in fetch_pool.items()
diff --git a/src/lerobot/datasets/episode_video_streaming.py b/src/lerobot/datasets/episode_video_streaming.py
index b84b31e08..15ca33d78 100644
--- a/src/lerobot/datasets/episode_video_streaming.py
+++ b/src/lerobot/datasets/episode_video_streaming.py
@@ -145,6 +145,7 @@ class NativeHTTPRangeFetcher:
         httpx.RemoteProtocolError,
         httpx.PoolTimeout,
     )
+    _RETRYABLE_STATUS_CODES = {408, 425, 429, 500, 502, 503, 504}
 
     def __init__(
         self,
@@ -331,6 +332,8 @@ class NativeHTTPRangeFetcher:
                 timings[key] += value
         if status_code == 403:
             raise PermissionError(f"HTTP range request returned 403 after URL refresh: {relative_path}")
+        if status_code != 206:
+            raise RuntimeError(f"HTTP range request returned {status_code} after retries: {relative_path}")
         self._record_timing(
             range_jobs=1.0,
             range_bytes=float(len(payload)),
@@ -351,6 +354,23 @@ class NativeHTTPRangeFetcher:
             attempt_start = time.perf_counter()
             try:
                 payload, status_code, timings = self._read_range_response_once(url, headers)
+                if status_code in self._RETRYABLE_STATUS_CODES:
+                    failed_attempt_s += time.perf_counter() - attempt_start
+                    exception_attempts += 1.0
+                    status_key = f"range_failed_status_{status_code}"
+                    exception_counts[status_key] = exception_counts.get(status_key, 0.0) + 1.0
+                    if attempt >= self.max_retries:
+                        timings["range_retry_attempts"] = retry_attempts
+                        timings["range_retry_sleep_s"] = retry_sleep_s
+                        timings["range_failed_attempt_s"] = failed_attempt_s
+                        timings["range_exception_attempts"] = exception_attempts
+                        timings.update(exception_counts)
+                        return payload, status_code, timings
+                    retry_attempts += 1.0
+                    sleep_s = min(0.5 * 2**attempt, 5.0)
+                    retry_sleep_s += sleep_s
+                    time.sleep(sleep_s)
+                    continue
                 timings["range_retry_attempts"] = retry_attempts
                 timings["range_retry_sleep_s"] = retry_sleep_s
                 timings["range_failed_attempt_s"] = failed_attempt_s
@@ -387,7 +407,7 @@ class NativeHTTPRangeFetcher:
         header_start = time.perf_counter()
         with self.client.stream("GET", url, headers=headers) as response:
             header_s = time.perf_counter() - header_start
-            if response.status_code == 403:
+            if response.status_code == 403 or response.status_code in self._RETRYABLE_STATUS_CODES:
                 return (
                     b"",
                     response.status_code,