Repository navigation
fetch_asset progress hook + total --num-envs semantics - #36
Merged
Merged
Conversation
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).
1 task done
KraHsu
added a commit
that referenced
this pull request
Jun 5, 2026
…envs fetch_asset progress hook + total --num-envs semantics
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
utils/download.pykeeps its stdlib-only contract but now exposes aProgressCallbackProtocol and aPROGRESS_CALLBACKContextVar. The_download_toloop 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 ininfo/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-envsis total, new--num-envs-per-gpuis per-rank: passing both errors out, and--num-envsmust 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-gpuso every worker sees the same per-rank value verbatim._strip_gpus_flagis renamed to_strip_distributed_flagsto reflect the broader scope.Notes for review
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.utils, the Rich rendering lives incli. Asset-zoo factories do not change._resolve_per_rank_num_envsis 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 checkuv run pyright(scoped to touched files — 0 errors)uv run pytest— 169 passeduv run mkdocs build --strictuv run genelab train --helpshows both--num_envsand--num_envs_per_gpuwith semantic notes.cache/assets/g1/to visually confirm the Rich download bar (cannot exercise headless in CI; needs a fresh fetch)