Import upstream snapshot d424a0c899f6b33cbd3daeb279913134349c0b63
Upstream: https://github.com/pollen-robotics/microduck_rl Upstream-Commit: d424a0c899f6b33cbd3daeb279913134349c0b63 Upstream-Branch: develop
This commit is contained in:
Upstream Snapshot
@@ -0,0 +1,79 @@
|
||||
# HF Jobs training
|
||||
|
||||
Train mjlab-microduck on Hugging Face's managed GPUs. Auth is the cached HF
|
||||
token (`hf auth login` or `HF_TOKEN`); everything goes through the
|
||||
`huggingface_hub` Python API — the standalone `hf` CLI is not required.
|
||||
|
||||
## One-time setup
|
||||
|
||||
```fish
|
||||
hf auth login # or export HF_TOKEN (any tool that caches the token works)
|
||||
wandb login # auto-detected from ~/.netrc and forwarded
|
||||
```
|
||||
|
||||
## Submit a run
|
||||
|
||||
Your normal train command, plus `--hf-jobs`:
|
||||
|
||||
```fish
|
||||
uv run train Mjlab-Kick-Flat-MicroDuck \
|
||||
--env.scene.num-envs 4096 --agent.max_iterations 4000 --hf-jobs
|
||||
```
|
||||
|
||||
You'll be asked which namespace to run under — your personal account or one
|
||||
of your orgs. Repos, uv-cache bucket, billing and the job itself all live in
|
||||
the chosen namespace. Pass `--namespace <name>` to skip the prompt
|
||||
(non-interactive runs default to personal).
|
||||
|
||||
Without `--hf-jobs` the command behaves exactly as before (local training).
|
||||
Submission flags are consumed locally; everything else is forwarded to
|
||||
`uv run train` inside the job.
|
||||
|
||||
Useful flags:
|
||||
- `--namespace <name>` — account/org to run under; skips the prompt
|
||||
- `--flavor l4x1` (default) / `a10g-large` / `a100-large`
|
||||
- `--timeout 12h` (default) — job is killed past this
|
||||
- `--detach` — submit and return immediately (default streams logs; Ctrl-C detaches without killing the job)
|
||||
- `--dry-run` — build tarball, print the job spec, do not submit
|
||||
- `--run-name <tag>` — overrides the auto-generated `<task>-<timestamp>` name
|
||||
- `--no-uv-cache` — disable the persistent `uv` cache bucket (first-run cost on every run)
|
||||
- `--no-wandb` — don't forward a wandb key
|
||||
|
||||
(`uv run scripts/hf/train_hf.py <task> ...` still works — it's a shim to the
|
||||
same code, which lives in `src/mjlab_microduck/hf_jobs.py`.)
|
||||
|
||||
## What happens under the hood
|
||||
|
||||
1. `git ls-files` snapshots tracked + uncommitted files of the repo you run
|
||||
from (worktree-aware) → `src-<stamp>.tar.gz`.
|
||||
2. Tarball is uploaded to private dataset `<namespace>/mjlab-microduck-src`.
|
||||
3. Private model repo `<namespace>/<run-name>` is created for checkpoints.
|
||||
4. A private HF bucket `<namespace>/mjlab-uv-cache` is mounted at `/uv-cache`
|
||||
and used as `UV_CACHE_DIR` so wheel downloads persist across runs (first
|
||||
run cold, subsequent runs fast).
|
||||
5. `HfApi.run_job` launches a container that:
|
||||
- installs `uv`, extracts the tarball, runs `uv sync` (warm-cached),
|
||||
- starts `scripts/hf/uploader.py` in background (watches `logs/rsl_rl/**/model_*.pt`, pushes every 60s),
|
||||
- runs `uv run train <task> <args>`,
|
||||
- does a final one-shot upload on exit.
|
||||
6. wandb credentials are forwarded as a secret — runs show up live in your
|
||||
wandb project.
|
||||
|
||||
## Browsing checkpoints
|
||||
|
||||
The submitter prints `https://huggingface.co/<namespace>/<run-name>` at
|
||||
start; new `.pt` files appear there during training.
|
||||
|
||||
## Managing jobs
|
||||
|
||||
The job id and URL are printed at submission. From Python:
|
||||
|
||||
```python
|
||||
from huggingface_hub import HfApi
|
||||
api = HfApi()
|
||||
api.list_jobs() # or namespace="pollen-robotics"
|
||||
for l in api.fetch_job_logs(job_id="...", follow=True): print(l)
|
||||
api.cancel_job(job_id="...")
|
||||
```
|
||||
|
||||
(or the `hf jobs ps/logs/cancel` CLI if you have it installed.)
|
||||
@@ -0,0 +1,15 @@
|
||||
"""Back-compat shim: the submission logic moved to mjlab_microduck.hf_jobs.
|
||||
|
||||
Prefer the integrated flag:
|
||||
uv run train <task> <train args...> --hf-jobs [--namespace <ns>] [...]
|
||||
|
||||
This script keeps the old invocation working:
|
||||
uv run scripts/hf/train_hf.py <task> [submission flags] <train args...>
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
from mjlab_microduck.hf_jobs import submit
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(submit(sys.argv[1:]))
|
||||
@@ -0,0 +1,75 @@
|
||||
"""Checkpoint uploader run inside an HF Job.
|
||||
|
||||
Watches `logs/rsl_rl/**/model_*.pt` and uploads new/updated files to the
|
||||
target HF Model repo. Designed to be `nohup uv run`-launched from the job
|
||||
bootstrap, with auth coming from the HF_TOKEN secret injected by `hf jobs run`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from huggingface_hub import HfApi, CommitOperationAdd
|
||||
|
||||
|
||||
def main() -> int:
|
||||
repo_id = os.environ.get("CKPT_REPO")
|
||||
if not repo_id:
|
||||
print("[uploader] CKPT_REPO not set, exiting", flush=True)
|
||||
return 1
|
||||
|
||||
poll_interval = float(os.environ.get("CKPT_POLL_INTERVAL", "60"))
|
||||
root = Path(os.environ.get("CKPT_ROOT", "logs/rsl_rl"))
|
||||
|
||||
one_shot = os.environ.get("CKPT_ONE_SHOT") == "1"
|
||||
|
||||
api = HfApi()
|
||||
api.create_repo(repo_id, repo_type="model", private=True, exist_ok=True)
|
||||
mode = "one-shot" if one_shot else f"every {poll_interval}s"
|
||||
print(f"[uploader] watching {root} -> {repo_id} ({mode})", flush=True)
|
||||
|
||||
sent: dict[Path, float] = {}
|
||||
while True:
|
||||
try:
|
||||
files = list(root.glob("**/model_*.pt"))
|
||||
# also pick up the dumped configs once
|
||||
files += [p for p in root.glob("**/params/*.yaml")]
|
||||
files += [p for p in root.glob("**/params/*.json")]
|
||||
|
||||
to_upload: list[CommitOperationAdd] = []
|
||||
for f in files:
|
||||
try:
|
||||
mtime = f.stat().st_mtime
|
||||
except FileNotFoundError:
|
||||
continue
|
||||
if sent.get(f) == mtime:
|
||||
continue
|
||||
# use path-in-repo relative to logs/rsl_rl so the repo mirrors run dirs
|
||||
rel = f.relative_to(root)
|
||||
to_upload.append(
|
||||
CommitOperationAdd(path_in_repo=str(rel), path_or_fileobj=str(f))
|
||||
)
|
||||
sent[f] = mtime
|
||||
|
||||
if to_upload:
|
||||
msg = f"upload {len(to_upload)} file(s)"
|
||||
api.create_commit(
|
||||
repo_id=repo_id,
|
||||
repo_type="model",
|
||||
operations=to_upload,
|
||||
commit_message=msg,
|
||||
)
|
||||
print(f"[uploader] pushed {len(to_upload)} file(s)", flush=True)
|
||||
except Exception as e:
|
||||
print(f"[uploader] error: {e}", flush=True)
|
||||
|
||||
if one_shot:
|
||||
return 0
|
||||
time.sleep(poll_interval)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user