Skip to content

refactor(observability): decouple trainers from WandB - #32

Draft
leviking98z-rgb wants to merge 1 commit into
refactor/reward-ownershipfrom
refactor/observability-api
Draft

refactor(observability): decouple trainers from WandB#32
leviking98z-rgb wants to merge 1 commit into
refactor/reward-ownershipfrom
refactor/observability-api

Conversation

@leviking98z-rgb

Copy link
Copy Markdown
Owner

Summary

  • Add a provider-neutral Observer protocol, NullObserver, factory, checkpoint-state helper, and phase instrumentation under unirl.observability.
  • Make BaseTrainer, every loop program, and all trainer variants emit metrics, media, progress, and lifecycle events through self.observer; trainers no longer import or name the WandB adapter.
  • Preserve existing report_to_wandb recipes while adding explicit logging.provider: none|wandb and generic logging.enabled selection.
  • Preserve disabled-run behavior without loading the WandB adapter: console progress still prints, memory-monitor step boundaries still fire, and phase timing remains injected at the rollout-step boundary.
  • Extend the framework guard so the Observer surface remains stable and trainer code cannot reintroduce direct WandB imports, wandb_logger, or backend-named lifecycle hooks.

This is a boundary-stability PR, not a LOC-reduction PR: 630 insertions / 364 deletions (net +266), including the new contract, null implementation, static guard, and documentation. It removes the duplicated timing/progress policy from the concrete adapter and makes another telemetry backend possible without changing trainer programs.

This is stacked on refactor/reward-ownership (#31 in this fork), not on main.

Related Issue

N/A

Test Plan

  • pre-commit run --all-files — passed, including all 120 recipe compositions, target resolution, Ruff, architecture boundaries, experimental boundaries, and framework contracts.
  • python -m compileall -q unirl/observability unirl/trainer unirl/utils/wandb_logger.py unirl/utils/memory_monitor.py unirl/utils/media.py unirl/types/media_preview.py lint/check_framework_contracts.py — passed.
  • python lint/check_framework_contracts.py — passed; reports one Observer contract and verifies no trainer-facing WandB dependency remains.
  • python lint/check_architecture_boundaries.py — passed (476 files / 4,279 internal import edges).
  • python lint/check_experimental_boundaries.py — passed.
  • One-off stdlib CPU harnesses exercised structural Observer conformance, old/new checkpoint run-id state, disabled factory selection, invalid-provider rejection, enabled-adapter argument forwarding, phase-time injection across consecutive steps, console progress, and memory-monitor step boundaries — passed.
  • Not run: GPU training or live WandB upload. Metric extraction, WandB SDK calls, and training numerics are unchanged; this PR changes their ownership and call boundary.

Compatibility / Risk

  • Existing logging.report_to_wandb, project, tags, entity, directory, and media fields remain accepted with the same enabled behavior.
  • New explicit configuration uses provider: wandb plus enabled: true; provider: none selects the dependency-light null observer. Unknown providers fail before training starts.
  • Checkpoints now write observer_run_id and continue writing wandb_run_id; resume accepts either field, preferring the provider-neutral one. Older and newer checkouts can therefore exchange trainer state.
  • BaseTrainer.wandb_logger, _init_wandb, _finish_wandb, and _loop_wandb_extra are replaced by provider-neutral names. Out-of-tree trainer subclasses using those internal seams must migrate to observer, _init_observability, _finish_observability, and _loop_observability_metadata.
  • Phase and memory instrumentation still wrap collaborators lazily on the first training step. No recipe, model, reward, dataset, tensor wire format, or GPU placement changes.

Reviewer Notes

Checklist

  • I reviewed the changed code and removed unrelated/generated artifacts.
  • I updated tests, docs, and configs where needed, or explained why not.

@leviking98z-rgb
leviking98z-rgb force-pushed the refactor/observability-api branch from fe46af1 to 926fabd Compare August 2, 2026 17:28
@leviking98z-rgb
leviking98z-rgb force-pushed the refactor/reward-ownership branch from 3cdf886 to 39165d1 Compare August 2, 2026 17:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant