Skip to content

Adopt OpenTelemetry behind the MetricsLogger interface #1685

Description

@erain

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.

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions