Skip to content

fetch_asset progress hook + total --num-envs semantics - #36

Merged
KraHsu merged 3 commits into
mainfrom
feat/fetch-progress-and-total-num-envs
May 15, 2026
Merged

KraHsu merged 3 commits into
mainfrom
feat/fetch-progress-and-total-num-envs

Conversation

@KraHsu

@KraHsu KraHsu commented May 15, 2026

Copy link
Copy Markdown
Owner

Summary

  • fetch_asset progress hook: utils/download.py keeps its stdlib-only contract but now exposes a ProgressCallback Protocol and a PROGRESS_CALLBACK ContextVar. The _download_to loop streams chunks manually and reports cumulative bytes + Content-Length. CLI side, cli/_progress.fetch_progress() is a context manager that installs a Rich-backed callback for the duration of cfg construction in info / play / train — so first-time downloads of the ~19 MB G1 (or any other asset-zoo) tar.gz now render a live download bar instead of a silent stall.
  • --num-envs is total, new --num-envs-per-gpu is per-rank: passing both errors out, and --num-envs must divide evenly by --gpus. The torchrun relaunch strips both env-count tokens from the inner argv and re-injects the parent-resolved --num-envs-per-gpu so every worker sees the same per-rank value verbatim. _strip_gpus_flag is renamed to _strip_distributed_flags to reflect the broader scope.

Notes for review

  • Stacked on Update unitree tests for asset zoo migration #35 (fix/unitree-test-asset-zoo) via a cherry-pick of the test update. Once Update unitree tests for asset zoo migration #35 merges, a rebase here will drop the duplicate commit cleanly.
  • The download progress callback design is intentionally stdlib-only — the Protocol type and ContextVar live in utils, the Rich rendering lives in cli. Asset-zoo factories do not change.
  • _resolve_per_rank_num_envs is the single source of truth for the CLI's env-count math; both _dispatch_play (gpus=1) and _dispatch_train (real gpus) call it.

Test plan

  • uv run ruff check
  • uv run pyright (scoped to touched files — 0 errors)
  • uv run pytest — 169 passed
  • uv run mkdocs build --strict
  • uv run genelab train --help shows both --num_envs and --num_envs_per_gpu with semantic notes
  • Manual: re-run G1 velocity train with a cleared .cache/assets/g1/ to visually confirm the Rich download bar (cannot exercise headless in CI; needs a fresh fetch)

KraHsu added 3 commits May 15, 2026 21:47
Asserts that the g1 robot is provided by genelab.asset_zoo (not the example),
matching the change in the previous commit. Renames the MJCF-resolution test
since the path is now resolved through fetch_asset rather than the vendored
directory.
* utils/download.py grows a stdlib-only ProgressCallback Protocol and a
  PROGRESS_CALLBACK ContextVar. _download_to now streams chunks manually
  and reports cumulative bytes + Content-Length when available, so a UI
  can render a live download bar without making utils depend on Rich.
* New cli/_progress.fetch_progress() context manager installs a Rich
  Progress-backed callback for the duration of cfg construction in
  info / play / train command bodies.
* --num-envs is now treated as the total env count across ranks. The new
  --num-envs-per-gpu flag is the per-rank knob. Passing both errors, and
  --num-envs must divide evenly by --gpus. The torchrun relaunch strips
  both env-count tokens from the inner argv and re-injects the resolved
  --num-envs-per-gpu so each worker sees the parent-decided per-rank
  value verbatim.
* Tests cover the resolver (mutual exclusion, divisibility, fallthrough),
  the relaunch argv rewrite, and progress callback dispatch (via kwarg,
  via ContextVar, kwarg overrides ContextVar).
@KraHsu
KraHsu merged commit 1708b9d into main May 15, 2026
3 checks passed
@KraHsu
KraHsu deleted the feat/fetch-progress-and-total-num-envs branch May 15, 2026 13:55
@KraHsu KraHsu mentioned this pull request May 15, 2026
1 task done
KraHsu added a commit that referenced this pull request Jun 5, 2026
…envs

fetch_asset progress hook + total --num-envs semantics
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant