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 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.
Problem
Tunix currently couples
MetricsLoggerto 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 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:jax.monitoringbackends 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).MetricsLogger.log()is double-written: once through the existing backends, and once as an OpenTelemetry gauge with stabletunix.*instrument names, low-cardinalitytunix.metrics.prefix/tunix.training.modeattributes, and a separatetunix.training.stepgauge for the logical step.tunix.sft.otel_wandb.WandbMetricsExporter— an OpenTelemetry SDK metric exporter that forwards the gauges towandb.logusing the familiar{prefix}/{mode}/{name}chart keys. Since the MetraxWandbBackendalso 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
MetricsLoggertests pass unchanged (default-off path).close().WandbMetricsExportertests with a fake run: key layout, step ordering, failure isolation, global-run resolution.