Skip to content

feat: add W&B worker GPU system metrics#1336

Closed
EazyReal wants to merge 1 commit into
areal-project:mainfrom
VmaxAI:codex/wandb-worker-system-metrics
Closed

feat: add W&B worker GPU system metrics#1336
EazyReal wants to merge 1 commit into
areal-project:mainfrom
VmaxAI:codex/wandb-worker-system-metrics

Conversation

@EazyReal

Copy link
Copy Markdown
Contributor

Summary

In single-controller training the process that calls wandb.init() runs on a
CPU/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=online
on 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:

  • Controller initialises the primary run with x_primary=True.
  • Selected worker processes attach a non-primary client with a role/rank
    x_label, x_primary=False, and x_update_finish_state=False, so they
    contribute system metrics only and never finalise the primary run.
  • Worker init is locked, exception-tolerant, and a no-op when disabled.

Supersedes #1335 (the previous attempt was closed before the head branch could
be reopened).

Usage

stats_logger:
  wandb:
    mode: shared
    system_metrics:
      enabled: true
      roles: [actor, rollout]   # null = every configured role
      gpu_device_ids: null      # let W&B use the worker's visible GPUs

Hardening since #1335

  • WandBSystemMetricsConfig.roles is now list[str] | None, normalized from
    tuples in __post_init__. Empty roles and negative GPU ids are rejected.
  • init_worker_wandb_system_metrics is protected by a threading.Lock and
    catches wandb.init failures so worker telemetry can never crash training.
  • finish_worker_wandb_system_metrics is exception-tolerant and idempotent.
  • New unit tests cover the validators, tuple to list normalization, and the
    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-test CI on a 2-GPU runner.
  • Optional: run an end-to-end shared-mode job and confirm worker GPU
    utilisation appears alongside the controller run.

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
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