API: HF dataset mirror for leaderboard + ckpts; fix ego_fn obs.get bug

Dockerfile

@@ -6,7 +6,9 @@ ENV PYTHONUNBUFFERED=1 \
     GUNICORN_WORKERS=2 \
     GUNICORN_TIMEOUT=180 \
     PORT=7860 \
-    BENCHMARK_UI_STORAGE=/data
+    BENCHMARK_UI_STORAGE=/data \
+    DATA_ROOT=/data \
+    BENCHMARK_HF_REPO=lainwired/jaxaht-benchmark-leaderboard
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
     curl ca-certificates build-essential git \

benchmark_ui/backend/app.py

@@ -535,6 +535,7 @@ def _register_routes(app: Flask) -> None:  # noqa: C901 (route fan-out)
         adapter = get_env(env)
         ego_factory = None
         upload_holder: dict[str, Any] = {}
+        backends = app.extensions["backends"]
 
         if payload.ego_kind == "builtin":
             if not payload.builtin_key:
@@ -563,6 +564,15 @@ def _register_routes(app: Flask) -> None:  # noqa: C901 (route fan-out)
                 idx=payload.idx,
             )
             upload_holder["uploaded"] = uploaded
+            # persist the raw upload so submissions survive container rebuilds
+            ckpt_id = backends.checkpoint.save_upload(uploaded_bytes, {
+                "agent_name": payload.agent_name,
+                "env": env,
+                "actor_type": payload.actor_type,
+                "ckpt_key": payload.ckpt_key,
+                "idx": payload.idx,
+            })
+            upload_holder["ckpt_id"] = ckpt_id
 
             def ego_factory():
                 env_obj = adapter.make_env(adapter.DEFAULT_KWARGS or {})
@@ -571,7 +581,6 @@ def _register_routes(app: Flask) -> None:  # noqa: C901 (route fan-out)
             raise ValueError(f"unknown ego_kind '{payload.ego_kind}'")
 
         jobs: JobManager = app.extensions["jobs"]
-        backends = app.extensions["backends"]
 
         def _run(progress_cb):
             try:
@@ -598,6 +607,11 @@ def _register_routes(app: Flask) -> None:  # noqa: C901 (route fan-out)
                     "per_partner": result["per_partner"],
                     "wall_clock_seconds": result["wall_clock_seconds"],
                     "notes": payload.notes,
+                    "checkpoint_id": upload_holder.get("ckpt_id"),
+                    "checkpoint_sha256": (
+                        backends.checkpoint.get_sha256(upload_holder["ckpt_id"])
+                        if "ckpt_id" in upload_holder else None
+                    ),
                 }
                 entry_id = backends.leaderboard.add_entry(env, payload.version, entry)
                 return {"entry_id": entry_id, **result}

benchmark_ui/backend/checkpoint_loader.py

@@ -188,7 +188,7 @@ def _generic_load_ego(env, uploaded: UploadedCheckpoint) -> Callable:
                 hstate_holder[0] = None
 
         if isinstance(obs, dict):
-            agent_obs = obs.get("agent_0") or next(iter(obs.values()))
+            agent_obs = obs["agent_0"] if "agent_0" in obs else next(iter(obs.values()))
         else:
             agent_obs = obs
 

benchmark_ui/backend/storage/__init__.py

@@ -1,9 +1,10 @@
 """Storage layer. Interface in ``base``; file impl in ``file``.
 
-Routes import ``build_backends`` and never touch a subclass directly.
-Swap with a Postgres impl later via STORAGE=postgres env var.
+Set BENCHMARK_HF_REPO=<dataset_id> to mirror leaderboard + ckpts to HF
+so they survive container rebuilds.
 """
 import os
+from pathlib import Path
 
 from .base import (
     Backends,
@@ -17,19 +18,46 @@ from .file import build_backends as _build_file
 
 
 def build_backends() -> Backends:
-    """Construct the storage bundle according to STORAGE env var.
-
-    Default: file-backed under ./benchmark_ui/data.
-    """
     backend = os.environ.get("STORAGE", "file").lower()
-    if backend == "file":
-        root = os.environ.get("DATA_ROOT", "./benchmark_ui/data")
-        return _build_file(root)
-    if backend == "postgres":
-        raise NotImplementedError(
-            "Postgres backend not implemented yet. Set STORAGE=file."
+    if backend != "file":
+        if backend == "postgres":
+            raise NotImplementedError("Postgres backend not implemented yet.")
+        raise ValueError(f"Unknown STORAGE: {backend}")
+
+    root = os.environ.get("DATA_ROOT", "./benchmark_ui/data")
+    hf_repo = os.environ.get("BENCHMARK_HF_REPO")
+
+    if hf_repo:
+        from .hf_mirror import bootstrap_from_hf, HFMirrorCheckpointStore, HFMirrorLeaderboardStore
+        bootstrap_root = Path(root)
+        bootstrap_from_hf(hf_repo, bootstrap_root)
+        _migrate_bootstrap_paths(bootstrap_root)
+
+    backends = _build_file(root)
+
+    if hf_repo:
+        backends = Backends(
+            leaderboard=HFMirrorLeaderboardStore(backends.leaderboard, hf_repo),
+            trajectory=backends.trajectory,
+            checkpoint=HFMirrorCheckpointStore(backends.checkpoint, hf_repo),
+            session=backends.session,
+            job=backends.job,
         )
-    raise ValueError(f"Unknown STORAGE: {backend}")
+    return backends
+
+
+def _migrate_bootstrap_paths(root: Path) -> None:
+    # the HF dataset uses "leaderboard/" while FileLeaderboardStore reads from "leaderboards/"
+    src = root / "leaderboard"
+    dst = root / "leaderboards"
+    if not src.exists():
+        return
+    dst.mkdir(exist_ok=True)
+    for f in src.iterdir():
+        if f.is_file():
+            target = dst / f.name
+            if not target.exists():
+                f.replace(target)
 
 
 __all__ = [

benchmark_ui/backend/storage/hf_mirror.py

@@ -0,0 +1,110 @@
+# HF Hub mirror for leaderboard + checkpoint stores.
+# Local file backs writes; after each successful write we push to HF dataset.
+# Bootstrap: snapshot_download pulls the dataset into the local root before serving.
+from __future__ import annotations
+
+import logging
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+from .base import CheckpointStore, LeaderboardStore
+from .file import FileCheckpointStore, FileLeaderboardStore
+
+log = logging.getLogger("hf_mirror")
+
+
+def bootstrap_from_hf(repo_id: str, local_root: Path) -> None:
+    # pull existing leaderboard + ckpts from the dataset on container start.
+    # idempotent: snapshot_download skips files that already exist locally.
+    from huggingface_hub import snapshot_download
+    local_root = Path(local_root)
+    local_root.mkdir(parents=True, exist_ok=True)
+    try:
+        snapshot_download(
+            repo_id=repo_id,
+            repo_type="dataset",
+            local_dir=str(local_root),
+            allow_patterns=["leaderboard/*", "checkpoints/**"],
+        )
+        log.info("bootstrapped storage from %s into %s", repo_id, local_root)
+    except Exception as exc:
+        log.warning("failed to bootstrap from %s (%s); starting empty", repo_id, exc)
+
+
+class HFMirrorLeaderboardStore(LeaderboardStore):
+    def __init__(self, inner: FileLeaderboardStore, repo_id: str):
+        self.inner = inner
+        self.repo_id = repo_id
+        self._executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="hf-push")
+        self._lock = threading.Lock()
+
+    def _push(self, env: str, version: str) -> None:
+        from huggingface_hub import HfApi, CommitOperationAdd
+        local = self.inner._path(env, version)
+        if not local.exists():
+            return
+        rel = f"leaderboard/{env}_{version}.json"
+        try:
+            HfApi().create_commit(
+                repo_id=self.repo_id, repo_type="dataset",
+                operations=[CommitOperationAdd(path_in_repo=rel, path_or_fileobj=str(local))],
+                commit_message=f"update {rel}",
+            )
+            log.info("pushed %s to HF", rel)
+        except Exception as exc:
+            log.exception("HF push failed for %s: %s", rel, exc)
+
+    def add_entry(self, env: str, version: str, entry: dict) -> str:
+        entry_id = self.inner.add_entry(env, version, entry)
+        # fire-and-forget push so the API response isn't blocked
+        self._executor.submit(self._push, env, version)
+        return entry_id
+
+    def list_entries(self, env: str, version: str) -> list[dict]:
+        return self.inner.list_entries(env, version)
+
+    def get_entry(self, env: str, version: str, entry_id: str):
+        return self.inner.get_entry(env, version, entry_id)
+
+    def clear(self, env: str, version: str) -> None:
+        self.inner.clear(env, version)
+        self._executor.submit(self._push, env, version)
+
+
+class HFMirrorCheckpointStore(CheckpointStore):
+    # mirrors user-uploaded ckpts to HF so they survive container restarts.
+    # the local FileCheckpointStore stays the source of truth for in-process reads.
+    def __init__(self, inner: FileCheckpointStore, repo_id: str):
+        self.inner = inner
+        self.repo_id = repo_id
+        self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="hf-ckpt-push")
+
+    def _push(self, ckpt_id: str) -> None:
+        # Push the whole ckpt dir (zip + extracted + meta) so a future bootstrap can rehydrate it fully.
+        from huggingface_hub import HfApi
+        ckpt_dir = self.inner.root / ckpt_id
+        if not ckpt_dir.exists():
+            return
+        try:
+            HfApi().upload_folder(
+                folder_path=str(ckpt_dir),
+                path_in_repo=f"checkpoints/{ckpt_id}",
+                repo_id=self.repo_id,
+                repo_type="dataset",
+                commit_message=f"upload ckpt {ckpt_id}",
+            )
+            log.info("pushed ckpt %s to HF", ckpt_id)
+        except Exception as exc:
+            log.exception("HF ckpt push failed for %s: %s", ckpt_id, exc)
+
+    def save_upload(self, zip_bytes: bytes, metadata: dict | None = None) -> str:
+        ckpt_id = self.inner.save_upload(zip_bytes, metadata)
+        self._executor.submit(self._push, ckpt_id)
+        return ckpt_id
+
+    def get_path(self, ckpt_id: str):
+        return self.inner.get_path(ckpt_id)
+
+    def get_sha256(self, ckpt_id: str):
+        return self.inner.get_sha256(ckpt_id)