feat: add W&B worker GPU system metrics#1336
Closed
EazyReal wants to merge 1 commit into
Closed
Conversation
In single-controller training the controller may run on a CPU node while
actor/rollout workers run on GPU nodes. With ``wandb.mode=online`` only the
controller's process metrics are recorded — worker GPU utilisation is missing.
This change introduces an opt-in worker-side W&B system metrics client using
W&B shared mode:
- Controller initialises the primary run with ``x_primary=True``.
- Selected worker processes attach a non-primary client with a role/rank label,
``x_primary=False`` and ``x_update_finish_state=False`` so they contribute
system metrics only and never finalise the run.
- Worker init is locked, exception-tolerant, and silently skips when disabled.
- ``stats_logger.wandb.system_metrics`` exposes ``enabled``, ``roles`` and
``gpu_device_ids``; the WandBConfig validator requires ``mode='shared'``.
Key changes:
- areal/api/cli_args.py: WandBSystemMetricsConfig dataclass + WandBConfig wiring
- areal/utils/wandb_system_metrics.py: worker init/finish + configure hook
- areal/utils/stats_logger.py: controller becomes the primary shared-mode client
- areal/utils/logging.py: register WandBSystemMetrics logger color
- areal/infra/rpc/{rpc_server,ray_rpc_server}.py: register worker hooks
- areal/infra/scheduler/{local,slurm,ray}.py: pre-resolve wandb run id
- areal/trainer/*: init StatsLogger before worker /configure
- docs/{en,zh}: metrics_tracking guide + regenerated CLI reference
- tests/test_wandb_system_metrics.py: unit tests for config and init/finish
EazyReal
requested review from
HwVanICI,
fishcrap,
garrett4wade,
guozhihao-224,
nuzant,
rchardx and
sitabulaixizawaluduo
as code owners
May 14, 2026 07:41
This was referenced May 14, 2026
3 tasks
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
In single-controller training the process that calls
wandb.init()runs on aCPU/controller host while actor and rollout workers run on GPU nodes. W&B
collects system metrics from the initialized W&B client, so
wandb.mode=onlineon only the controller records controller-process memory but not worker GPU
utilisation.
This PR opts workers in to W&B system metrics via shared-mode:
x_primary=True.x_label,x_primary=False, andx_update_finish_state=False, so theycontribute system metrics only and never finalise the primary run.
Supersedes #1335 (the previous attempt was closed before the head branch could
be reopened).
Usage
Hardening since #1335
WandBSystemMetricsConfig.rolesis nowlist[str] | None, normalized fromtuples in
__post_init__. Empty roles and negative GPU ids are rejected.init_worker_wandb_system_metricsis protected by athreading.Lockandcatches
wandb.initfailures so worker telemetry can never crash training.finish_worker_wandb_system_metricsis exception-tolerant and idempotent.init-failure path. Logger color registered for
WandBSystemMetrics.Validation
uv run pre-commit run --from-ref main --to-ref HEAD— all hooks pass(mdformat, ruff lint+format, conventional-pre-commit, CLI doc regen, SPDX).
uv run python -m pytest tests/test_wandb_system_metrics.py tests/test_trackio_backend.py— 20 passed.
Test plan
safe-to-testCI on a 2-GPU runner.utilisation appears alongside the controller run.