Adopt OpenTelemetry behind the MetricsLogger interface
- Dominant language
- Python
- Stars
- 2.5k
- Forks
- 345
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 240
Description
## Problem
Tunix currently couples `MetricsLogger` to Metrax backends through process-global JAX monitoring listeners. Backend configuration is backend-specific and lives inside a training-library object, and there is no standard path for correlated metrics, traces, and structured events. [OpenTelemetry](https://opentelemetry.io/) would give Tunix a vendor-neutral instrumentation layer.
## Updated plan: phased adoption
*(Edited: the original proposal was a full replacement of the Metrax backends. Per maintainer feedback, we are not ready to fully commit to OpenTelemetry yet — we want to try it first. The plan is therefore split into phases so each step can be evaluated before the next.)*
### Phase 1 — opt-in double-write (#1686)
Add an experimental, **off-by-default** flag `MetricsLoggerOptions.enable_opentelemetry`:
- The existing Metrax/`jax.monitoring` backends stay the default and are completely unchanged. With the flag unset, behavior is identical to today and OpenTelemetry does not even need to be installed (it is a new optional extra: `pip install 'google-tunix[otel]'`, not a core dependency).
- With the flag set, every scalar passed to `MetricsLogger.log()` is **double-written**: once through the existing backends, and once as an OpenTelemetry gauge with stable `tunix.*` instrument names, low-cardinality `tunix.metrics.prefix` / `tunix.training.mode` attributes, and a separate `tunix.training.step` gauge for the logical step.
- OpenTelemetry providers and exporters are application-owned. Tunix never configures exporters and never flushes or shuts down providers; a provider can also be injected for tests/embedding.
- **Weights & Biases compatibility**: W&B natively ingests OTLP traces (Weave endpoint) but has no OTLP endpoint for run metrics, so Phase 1 ships `tunix.sft.otel_wandb.WandbMetricsExporter` — an OpenTelemetry SDK metric exporter that forwards the gauges to `wandb.log` using the familiar `{prefix}/{mode}/{name}` chart keys. Since the Metrax `WandbBackend` also keeps running, both pipelines can be compared side by side in the same W&B project.
### Phase 2 — spans and structured events (follow-up)
Once double-write metrics look good: add `span()` / `event()` helpers to the same facade, instrument high-value operations (training step, RL rollout), and fix shared-logger ownership so a logger injected into actor and critic trainers is not closed by each trainer.
### Phase 3 — switch (only after validation)
Only after double-write has been validated in real runs (including the W&B path): decide whether to make OpenTelemetry the primary emission path and deprecate the backend-specific option fields. Not in scope until Phases 1–2 have been evaluated.
## Validation for Phase 1
- Existing `MetricsLogger` tests pass unchanged (default-off path).
- In-memory OpenTelemetry SDK tests for double-write: instrument naming/normalization, attributes, step gauge, non-scalar handling, process-zero policy, missing-API error, provider lifecycle across `close()`.
- `WandbMetricsExporter` tests with a fake run: key layout, step ordering, failure isolation, global-run resolution.
- Documentation of the flag, the naming scheme, and the OTLP + W&B setups.
Contributor guide
Assessment
This issue has not been assessed yet.