[Tunix] Adopt OpenTelemetry behind MetricsLogger

[erain]

Resolves #1685

Summary

This change keeps MetricsLogger as the public Tunix telemetry facade while replacing the Metrax/JAX-monitoring backend implementation with OpenTelemetry.

The facade now:

• preserves the existing log(), metric_exists(), get_metric(), get_metric_history(), and close() call patterns;
• retains local metric history because OpenTelemetry is an emission API, not a query store;
• emits numeric measurements through an application-owned OpenTelemetry MeterProvider;
• exposes span() for correlated tracing and event() for structured Python log records that an OpenTelemetry logging bridge can export;
• uses global OpenTelemetry providers by default and supports keyword-only provider injection for tests and embedding;
• does not configure exporters or shut down global/injected providers;
• preserves process-zero emission by default, with an opt-in all-process policy;
• removes the runtime Metrax dependency and adds only opentelemetry-api to the core package. The SDK remains a test/application dependency.

Motivation

The previous implementation coupled each logger to Metrax backends through process-global JAX monitoring listeners. That caused three architectural problems:

1. Backend configuration lived inside a training-library object instead of at the application/observability boundary.
2. A logger's close() could clear process-global listeners and close resources shared by other trainers.
3. Metrics, traces, and structured events had no common correlation or export model.

OpenTelemetry gives Tunix a vendor-neutral instrumentation layer. Applications can choose OTLP, a collector, or a vendor distribution without changing Tunix training call sites.

Design

1. Compatibility facade

MetricsLogger remains the interface used by SFT, RL, performance metrics, and progress bars. Existing metric logging and local history behavior is preserved, including geometric-mean aggregation for perplexity.

MetricsLoggerOptions remains copy-safe and continues to accept these legacy fields so existing Python and YAML configuration still parses:

• log_dir
• project_name
• run_name
• flush_every_n_steps
• backend_kwargs

Those fields no longer configure telemetry backends. backend_kwargs emits a deprecation warning because readers, processors, and exporters now belong on OpenTelemetry providers. The legacy log_dir remains temporarily meaningful to agentic trajectory CSV logging; that artifact path is explicitly documented as separate from telemetry.

Two signal-level options are added:

• enabled: disables signal emission while retaining local metric history.
• emit_on_all_processes: emits on every JAX process and adds tunix.jax.process.index; otherwise process zero emits.

attributes supplies common low-cardinality signal attributes. Service, deployment, and run metadata are documented as OpenTelemetry Resource attributes instead.

2. Provider and lifecycle ownership

The core package depends on opentelemetry-api, not the SDK. A process without an SDK receives the standard no-op meter/tracer behavior and training continues normally.

Applications may configure global providers or inject meter_provider and tracer_provider when constructing a logger. Injection is keyword-only so the existing positional constructor remains unchanged.

MetricsLogger.close() deliberately does not flush or shut down providers. Provider lifecycle is owned by the application that configured the readers, processors, and exporters.

Trainer ownership now follows the same rule:

• PeftTrainer closes a logger only when it created that logger.
• A logger injected into actor/critic trainers remains open.
• RLCluster, which creates and shares the RL logger, closes that facade once after its trainers close.

This prevents one trainer from invalidating telemetry used by another trainer or by the host application.

3. Metric data model

Synchronous gauges represent the latest training values because loss, perplexity, learning rate, and gradient norm are non-monotonic current measurements.

Tunix metric key	OpenTelemetry instrument
loss	tunix.training.loss
perplexity	tunix.training.perplexity
learning_rate	tunix.training.learning_rate
grad_norm	tunix.training.gradient.norm

Unknown user metric names remain supported. They are normalized into the tunix.* namespace, so a key such as rewards/score mean becomes tunix.rewards.score.mean.

The legacy prefix and mode become attributes:

• tunix.metrics.prefix
• tunix.training.mode

The logical step is emitted as a separate tunix.training.step gauge. It is intentionally not a metric attribute, because a distinct attribute value per step would create an unbounded time-series cardinality.

Only numeric scalar values are emitted. Non-scalar values remain in local history and produce a warning rather than interrupting training. Telemetry conversion/export setup failures are best effort and cannot prevent history updates.

The instrumentation scope is the stable name tunix, with the installed google-tunix package version when available.

4. Traces and structured events

The same facade adds:

• span(name, attributes=...), implemented with start_as_current_span;
• event(name, body=..., attributes=..., severity=...), implemented as a structured standard-library log record.

Using standard Python logging avoids depending on private/SDK logging APIs in the core library. Applications that install an OpenTelemetry logging handler can export these records and correlate them with the current span.

This change instruments two high-value operations without replacing Tunix's existing performance tracing:

• tunix.model.train around a PEFT training step;
• tunix.rollout around RL generation.

The existing Perfetto-oriented tracers continue to run alongside these OpenTelemetry spans. Injected legacy-compatible logger objects that do not implement span() continue to work through a no-op context fallback.

5. Concurrency and process behavior

Metric history and instrument creation use locks so shared logger instances are safe against concurrent logging. Process-zero emission remains the default for all three signal types to avoid duplicate data. All-process emission is explicit and identifies the source JAX process.

Migration impact

This is source-compatible for normal trainer and logger call sites, but intentionally changes backend ownership:

• Tunix no longer creates TensorBoard, W&B, CLU, or custom Metrax backends.
• Existing backend-specific option fields parse during migration but do not create exporters.
• Applications configure an OpenTelemetry SDK, Resource, readers/processors, and exporters before creating trainers.
• Destinations can be reached through OTLP collectors, OpenTelemetry distributions, or destination-specific exporters.
• Applications are responsible for provider flush and shutdown.

The documentation includes an OTLP metric/trace setup example and explains the no-SDK behavior, multi-process policy, structured logging bridge, lifecycle boundary, and legacy trajectory path.

Non-goals

• Configuring a global SDK or exporter from inside Tunix.
• Making the OpenTelemetry SDK a core runtime dependency.
• Replacing the existing detailed Perfetto performance trace format in this change.
• Moving the agentic trajectory artifact path out of MetricsLoggerOptions; that is called out as a follow-up.
• Preserving the Metrax custom-backend factory architecture, which would retain the coupling this design removes.

Validation

Installed the package normally with the test extra:

python -m pip install -e '.[test]'

Focused package integration:

python -m pytest   tests/sft/metrics_logger_test.py   tests/sft/progress_bar_test.py   tests/sft/peft_trainer_test.py::PeftTrainerTest::test_close_does_not_close_injected_metrics_logger   -q
# 21 passed, 3 subtests passed

Complete SFT trainer integration:

python -m pytest tests/sft/peft_trainer_test.py -q
# 23 passed

Applicable RL cluster integration:

python -m pytest tests/rl/rl_cluster_test.py -q   -k 'not test_init_vllm_rollout_engine and not test_init_sglang_jax_rollout_engine'
# 25 passed, 6 deselected

The deselected cases require optional vllm or sgl_jax installations that are not included in .[test]. An unfiltered run passed the other 26 cases and failed only while importing those optional engines.

Additional checks:

• 15 OpenTelemetry SDK tests cover known and dynamic metric naming, attributes, step handling, history compatibility, perplexity aggregation, scalar conversion, no-op providers, disabled emission, process policy, span exception recording, structured events, and provider lifecycle.
• OpenTelemetry API 1.25 compatibility was checked for synchronous gauges, no-op providers, and invalid-span support.
• Progress-bar integration: 5 tests and 3 subtests passed.
• Python compilation passed for every modified Python file.
• Pyink, isort, Pylint, and git diff --check passed for the modified sources. Pylint reports only pre-existing non-blocking warnings in trainer/cluster modules.
• The repository's aggregate pre-commit invocation currently cannot initialize its Markdown hook because mdformat 1.0.0 is pinned with mdformat-gfm 0.3.5, which requires mdformat<0.8; Python hooks were run individually.

Reference

• Issue: Adopt OpenTelemetry behind the MetricsLogger interface #1685
• OpenTelemetry Python documentation: https://opentelemetry.io/docs/languages/python/
• OpenTelemetry metrics semantic conventions: https://opentelemetry.io/docs/specs/semconv/general/metrics/

Colab Notebook

Not applicable. The existing MetricsLogger and trainer entry points remain the usage surface; SDK/exporter configuration and the additive span/event methods are documented and covered with in-memory SDK tests.

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and all applicable unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch.
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.