force-upgrade huggingface_hub during runtime install (datasets needs BucketNotFoundError)

handler.py

@@ -1,14 +1,18 @@
-"""HF Inference Endpoint handler for the lerobot folding_final policy.
-
-Strategy: lerobot 0.5.1 is the version that matches the model's config
-schema and doesn't need a patched transformers, but its pyproject pins
-`requires_python>=3.12`. HFIE's default container is Python 3.11. The
-lerobot/pi0 source itself doesn't use any 3.12-only syntax — the version
-pin is metadata-only. So we install lerobot at runtime inside __init__
-with `--ignore-requires-python`, then import.
-
-This adds ~90–180 s to the first cold-start (the install + ~600 MB of
-deps). Subsequent warm requests are normal.
+"""HF Inference Endpoint handler for lerobot folding_final (pi0.5).
+
+Why this is fiddly:
+  - The model was saved by lerobot >=0.5.1 (config has fields added there).
+  - lerobot 0.5.x's pyproject pins `requires_python>=3.12`, but HFIE's
+    default container is Python 3.11.
+  - lerobot 0.5.x's pi0 source itself uses no Python-3.12-only syntax —
+    the version pin is metadata-only. So we install it at runtime with
+    `--ignore-requires-python`.
+  - The container pre-installs an OLDER `huggingface_hub`; lerobot 0.5.1's
+    transitive `datasets` dep needs a newer one (uses `BucketNotFoundError`).
+    Forcing `--upgrade` so pip overwrites the pre-installed copy.
+
+Cold-start: ~120–180 s extra for the pip install (lerobot + ~600 MB deps).
+Warm requests are normal.
 """
 from __future__ import annotations
 
@@ -40,35 +44,34 @@ ACTION_FEATURE_NAMES = [
 
 
 def _ensure_lerobot() -> None:
-    """Install lerobot 0.5.1 ignoring its >=3.12 pip metadata check.
-    The pi0 module itself is 3.11-compatible (no PEP 695 syntax)."""
+    """Install lerobot 0.5.1 + its modern deps, overriding the pre-installed
+    huggingface_hub (too old for `datasets` to import)."""
     try:
         import lerobot.policies.pi0.modeling_pi0  # noqa: F401
-        log.info("lerobot already installed")
+        from huggingface_hub.errors import BucketNotFoundError  # noqa: F401
+        log.info("lerobot + modern hub already present")
         return
     except ImportError:
         pass
-    log.info("installing lerobot[pi0]==0.5.1 (ignoring requires-python) ...")
+    log.info("installing lerobot 0.5.1 + upgraded deps ...")
     subprocess.check_call(
         [
-            sys.executable,
-            "-m",
-            "pip",
-            "install",
+            sys.executable, "-m", "pip", "install",
             "--no-cache-dir",
-            "--ignore-requires-python",
+            "--upgrade",                       # overwrite pre-installed older copies
+            "--ignore-requires-python",        # pip metadata gate; pi0 source is 3.11-compat
             "lerobot[pi0]==0.5.1",
+            "huggingface_hub>=0.30",           # for BucketNotFoundError
         ],
         env={**os.environ, "PIP_IGNORE_REQUIRES_PYTHON": "1"},
     )
-    log.info("lerobot install done")
+    log.info("install done")
 
 
 class EndpointHandler:
     def __init__(self, path: str = "") -> None:
         _ensure_lerobot()
 
-        # Imports must be after _ensure_lerobot.
         from huggingface_hub import snapshot_download
         from lerobot.policies.pi0.modeling_pi0 import PI0Policy
 
@@ -98,7 +101,7 @@ class EndpointHandler:
         images_in = inputs.get("images") or {}
 
         batch = self._build_batch(images_in, state, prompt)
-        actions = self.policy.predict_action_chunk(batch)  # (1, T, A_padded)
+        actions = self.policy.predict_action_chunk(batch)
         arr = actions.squeeze(0).detach().to("cpu").float().numpy()
         if arr.shape[1] > 16:
             arr = arr[:, :16]
@@ -111,12 +114,7 @@ class EndpointHandler:
 
     # ── helpers ──────────────────────────────────────────────────────
 
-    def _build_batch(
-        self,
-        images_in: dict[str, str],
-        state: np.ndarray,
-        prompt: str,
-    ) -> dict[str, Any]:
+    def _build_batch(self, images_in, state, prompt):
         try:
             from lerobot.constants import OBS_STATE
         except ImportError:
@@ -135,7 +133,7 @@ class EndpointHandler:
         return batch
 
     @staticmethod
-    def _decode_image(b64: str | None) -> np.ndarray | None:
+    def _decode_image(b64):
         if not b64:
             return None
         try:
@@ -146,7 +144,7 @@ class EndpointHandler:
             log.warning("image decode failed: %s", e)
             return None
 
-    def _warmup(self) -> None:
+    def _warmup(self):
         dummy = base64.b64encode(_dummy_jpeg(224, 224)).decode("ascii")
         self({"inputs": {
             "images": {"left_wrist": dummy, "right_wrist": dummy, "base": dummy},
@@ -155,7 +153,7 @@ class EndpointHandler:
         }})
 
 
-def _dummy_jpeg(w: int, h: int) -> bytes:
+def _dummy_jpeg(w, h):
     buf = io.BytesIO()
     Image.fromarray(np.zeros((h, w, 3), dtype=np.uint8)).save(buf, format="JPEG")
     return buf.getvalue()