Validate checkpoint credentials once; log throughput

app.py

@@ -21,6 +21,7 @@ ENV_NAME = os.environ.get("ENV_NAME", "G1JoystickRoughTerrain")
 NUM_TIMESTEPS = int(os.environ.get("NUM_TIMESTEPS", 200_000_000))
 SEED = int(os.environ.get("SEED", 0))
 HF_REPO = os.environ.get("HF_REPO", "").strip()
+HF_TOKEN = os.environ.get("HF_TOKEN", "").strip() or None
 AUTO_START = os.environ.get("AUTO_START", "1") == "1"
 
 # /data exists only when persistent storage is attached; fall back to /tmp.
@@ -39,18 +40,49 @@ class Stopped(Exception):
     pass
 
 
-def _upload(path: Path) -> None:
+UPLOAD = {"enabled": False, "fails": 0}
+
+
+def init_upload() -> None:
+    """Validate credentials once, up front, instead of failing on every eval."""
     if not HF_REPO:
+        log(f"HF_REPO unset -- checkpoints stay in {OUT} and are lost on restart")
+        return
+    if not HF_TOKEN:
+        log("HF_TOKEN secret unset -- cannot push checkpoints")
         return
     try:
         from huggingface_hub import HfApi
-        api = HfApi()
+        api = HfApi(token=HF_TOKEN)
+        who = api.whoami().get("name")
         api.create_repo(HF_REPO, repo_type="model", exist_ok=True)
-        api.upload_file(path_or_fileobj=str(path), path_in_repo=path.name,
-                        repo_id=HF_REPO, repo_type="model")
-        log(f"uploaded {path.name} -> {HF_REPO}")
+        UPLOAD["enabled"] = True
+        log(f"checkpoints -> https://huggingface.co/{HF_REPO} (as {who})")
+    except Exception as e:
+        log(f"cannot write to HF_REPO={HF_REPO}: {type(e).__name__}: "
+            f"{str(e).splitlines()[0][:200]}")
+        log("fix: HF_TOKEN needs write scope on that namespace. A token scoped to "
+            "your own user cannot write to an org repo -- point HF_REPO at a "
+            "namespace the token owns, e.g. <your-user>/g1-rough-terrain.")
+
+
+def _upload(path: Path) -> None:
+    if not UPLOAD["enabled"]:
+        return
+    try:
+        from huggingface_hub import HfApi
+        HfApi(token=HF_TOKEN).upload_file(
+            path_or_fileobj=str(path), path_in_repo=path.name,
+            repo_id=HF_REPO, repo_type="model")
+        log(f"uploaded {path.name}")
+        UPLOAD["fails"] = 0
     except Exception as e:                       # never let upload kill training
-        log(f"upload failed ({type(e).__name__}): {e}")
+        UPLOAD["fails"] += 1
+        log(f"upload failed ({type(e).__name__}): {str(e).splitlines()[0][:160]}")
+        if UPLOAD["fails"] >= 3:
+            UPLOAD["enabled"] = False
+            log("disabling uploads after 3 failures; training continues, "
+                f"checkpoints remain in {OUT}")
 
 
 def install_jax_pmap_shims() -> None:
@@ -127,14 +159,24 @@ def train_worker() -> None:
             f"timesteps={NUM_TIMESTEPS:,}")
 
         OUT.mkdir(parents=True, exist_ok=True)
+        init_upload()
+        t_start = [None]
 
         def progress(step, metrics):
             if STATE["stop"]:
                 raise Stopped()
+            import time
             STATE["step"] = int(step)
             rew = metrics.get("eval/episode_reward", float("nan"))
             ln = metrics.get("eval/avg_episode_length", float("nan"))
-            log(f"step {int(step):>12,}  reward {rew:8.2f}  ep_len {ln:7.1f}")
+            if t_start[0] is None:
+                t_start[0] = (time.time(), int(step))
+                sps = float("nan")
+            else:
+                t0, s0 = t_start[0]
+                sps = (int(step) - s0) / max(time.time() - t0, 1e-9)
+            log(f"step {int(step):>12,}  reward {rew:8.2f}  ep_len {ln:7.1f}  "
+                f"{sps:,.0f} steps/s")
 
         def policy_params_fn(step, make_policy, params):
             p = OUT / f"ckpt_{int(step)}.pkl"