Instrument HfFileSystem range requests

scripts/bench_episode_byte_cache.py

@@ -661,6 +661,63 @@ def _print_range_timing_summary(fetch_pool: dict[str, float]) -> None:
     if status_counts:
         summary = ", ".join(f"{status}={count:.0f}" for status, count in sorted(status_counts.items()))
         print(f"| http status counts | {summary} |")
+    for method in ("head", "get"):
+        request_count = fetch_pool.get(f"range_hffs_{method}_requests", 0.0)
+        if request_count <= 0:
+            continue
+        print(f"| hffs {method.upper()} requests/range | {request_count / range_jobs:.3f} |")
+        print(
+            f"| hffs {method.upper()} total | {fetch_pool[f'range_hffs_{method}_s'] * 1000 / range_jobs:.3f} |"
+        )
+        retries = fetch_pool.get(f"range_hffs_{method}_retries", 0.0)
+        exceptions = fetch_pool.get(f"range_hffs_{method}_exception_attempts", 0.0)
+        if retries:
+            print(f"| hffs {method.upper()} retries/range | {retries / range_jobs:.3f} |")
+            print(
+                f"| hffs {method.upper()} retry sleep | "
+                f"{fetch_pool.get(f'range_hffs_{method}_retry_sleep_s', 0.0) * 1000 / range_jobs:.3f} |"
+            )
+        if exceptions:
+            print(f"| hffs {method.upper()} exceptions/range | {exceptions / range_jobs:.3f} |")
+            print(
+                f"| hffs {method.upper()} failed attempts | "
+                f"{fetch_pool.get(f'range_hffs_{method}_failed_attempt_s', 0.0) * 1000 / range_jobs:.3f} |"
+            )
+        bytes_read = fetch_pool.get(f"range_hffs_{method}_bytes", 0.0)
+        total_s = fetch_pool.get(f"range_hffs_{method}_s", 0.0)
+        if bytes_read > 0 and total_s > 0:
+            print(f"| hffs {method.upper()} MiB/s | {bytes_read / total_s / 1024**2:.1f} |")
+        hffs_status_counts = {
+            key.removeprefix(f"range_hffs_{method}_status_"): value
+            for key, value in fetch_pool.items()
+            if key.startswith(f"range_hffs_{method}_status_")
+        }
+        if hffs_status_counts:
+            summary = ", ".join(
+                f"{status}={count:.0f}" for status, count in sorted(hffs_status_counts.items())
+            )
+            print(f"| hffs {method.upper()} status counts | {summary} |")
+        hffs_failed_status_counts = {
+            key.removeprefix(f"range_hffs_{method}_failed_status_"): value
+            for key, value in fetch_pool.items()
+            if key.startswith(f"range_hffs_{method}_failed_status_")
+        }
+        if hffs_failed_status_counts:
+            summary = ", ".join(
+                f"{status}={count:.0f}" for status, count in sorted(hffs_failed_status_counts.items())
+            )
+            print(f"| hffs {method.upper()} failed status counts | {summary} |")
+        hffs_exception_counts = {
+            key.removeprefix(f"range_hffs_{method}_exception_"): value
+            for key, value in fetch_pool.items()
+            if key.startswith(f"range_hffs_{method}_exception_")
+            and key != f"range_hffs_{method}_exception_attempts"
+        }
+        if hffs_exception_counts:
+            summary = ", ".join(
+                f"{name}={count:.0f}" for name, count in sorted(hffs_exception_counts.items())
+            )
+            print(f"| hffs {method.upper()} exception counts | {summary} |")
     chunks = fetch_pool.get("range_chunks", 0.0)
     if chunks > 0:
         bytes_read = fetch_pool.get("range_bytes", 0.0)

src/lerobot/datasets/episode_video_streaming.py

@@ -18,6 +18,7 @@ from concurrent.futures import Future, ThreadPoolExecutor
 from dataclasses import dataclass
 from importlib import metadata
 from pathlib import Path
+from types import MethodType
 from typing import Any
 from urllib.parse import quote, urljoin, urlparse
 
@@ -25,7 +26,7 @@ import fsspec
 import httpx
 import numpy as np
 from huggingface_hub import HfApi, HfFileSystem, constants
-from huggingface_hub.utils import hf_raise_for_status
+from huggingface_hub.utils import get_session, hf_raise_for_status
 
 from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
 from lerobot.datasets.mp4 import Mp4Index, Mp4SampleSlice, fetch_mp4_index, synthesize_mp4
@@ -85,6 +86,7 @@ class ThreadLocalRangeFetcher:
             handle = self.fs.open(
                 self._url(relative_path), "rb", block_size=self.block_size, cache_type=self.cache_type
             )
+            self._instrument_hf_handle(handle)
             handles[relative_path] = handle
         return handle
 
@@ -115,6 +117,117 @@ class ThreadLocalRangeFetcher:
             for key, value in kwargs.items():
                 self._timing_totals[key] = self._timing_totals.get(key, 0.0) + value
 
+    def _instrument_hf_handle(self, handle: Any) -> None:
+        if getattr(handle, "_lerobot_range_timing", False):
+            return
+        if not hasattr(handle, "_request_with_retry"):
+            return
+
+        def request_with_retry(
+            handle_self,
+            method: str,
+            url: str,
+            *,
+            headers: dict[str, str],
+            follow_redirects: bool | None = None,
+            max_retries: int = 5,
+        ) -> httpx.Response:
+            from huggingface_hub.hf_file_system import _RANGE_RETRY_EXCEPTIONS, _RANGE_RETRY_STATUS_CODES
+
+            method_key = method.lower()
+            sleep_time = 1.0
+            retry_attempts = 0.0
+            retry_sleep_s = 0.0
+            failed_attempt_s = 0.0
+            exception_attempts = 0.0
+            extra_counts: dict[str, float] = {}
+            call_start = time.perf_counter()
+            request_kwargs: dict[str, Any] = {
+                "headers": headers,
+                "timeout": constants.HF_HUB_DOWNLOAD_TIMEOUT,
+            }
+            if follow_redirects is not None:
+                request_kwargs["follow_redirects"] = follow_redirects
+
+            for attempt in range(max_retries + 1):
+                attempt_start = time.perf_counter()
+                try:
+                    response = get_session().request(method, url, **request_kwargs)
+                except _RANGE_RETRY_EXCEPTIONS as exc:
+                    failed_attempt_s += time.perf_counter() - attempt_start
+                    exception_attempts += 1.0
+                    key = f"range_hffs_{method_key}_exception_{type(exc).__name__}"
+                    extra_counts[key] = extra_counts.get(key, 0.0) + 1.0
+                    if attempt == max_retries:
+                        self._record_hffs_request_timing(
+                            method_key,
+                            time.perf_counter() - call_start,
+                            retry_attempts,
+                            retry_sleep_s,
+                            failed_attempt_s,
+                            exception_attempts,
+                            None,
+                            0,
+                            extra_counts,
+                        )
+                        raise
+                else:
+                    elapsed = time.perf_counter() - attempt_start
+                    if response.status_code not in _RANGE_RETRY_STATUS_CODES or attempt == max_retries:
+                        self._record_hffs_request_timing(
+                            method_key,
+                            time.perf_counter() - call_start,
+                            retry_attempts,
+                            retry_sleep_s,
+                            failed_attempt_s,
+                            exception_attempts,
+                            response.status_code,
+                            len(response.content),
+                            extra_counts,
+                        )
+                        return response
+
+                    failed_attempt_s += elapsed
+                    key = f"range_hffs_{method_key}_failed_status_{response.status_code}"
+                    extra_counts[key] = extra_counts.get(key, 0.0) + 1.0
+                    response.close()
+
+                time.sleep(sleep_time)
+                retry_attempts += 1.0
+                retry_sleep_s += sleep_time
+                sleep_time = min(8.0, sleep_time * 2)
+
+            raise RuntimeError("unreachable")
+
+        handle._request_with_retry = MethodType(request_with_retry, handle)
+        handle._lerobot_range_timing = True
+
+    def _record_hffs_request_timing(
+        self,
+        method_key: str,
+        total_s: float,
+        retry_attempts: float,
+        retry_sleep_s: float,
+        failed_attempt_s: float,
+        exception_attempts: float,
+        status_code: int | None,
+        byte_count: int,
+        extra_counts: dict[str, float],
+    ) -> None:
+        timings = {
+            f"range_hffs_{method_key}_requests": 1.0,
+            f"range_hffs_{method_key}_s": total_s,
+            f"range_hffs_{method_key}_retries": retry_attempts,
+            f"range_hffs_{method_key}_retry_sleep_s": retry_sleep_s,
+            f"range_hffs_{method_key}_failed_attempt_s": failed_attempt_s,
+            f"range_hffs_{method_key}_exception_attempts": exception_attempts,
+            f"range_hffs_{method_key}_bytes": float(byte_count),
+        }
+        if status_code is not None:
+            timings[f"range_hffs_{method_key}_status_{status_code}"] = 1.0
+        timings.update(extra_counts)
+        self._record_timing(**timings)
+
     def timing_summary(self) -> dict[str, float]:
         with self._timing_lock:
             return dict(self._timing_totals)