Problem
Tunix currently couples MetricsLogger to Metrax backends through process-global JAX monitoring listeners. The same class also owns in-memory metric history used by trainers and progress bars. This makes telemetry configuration backend-specific, gives individual logger instances process-global cleanup behavior, and does not provide a standard path for correlated metrics, traces, and structured events.
Proposed design
Keep MetricsLogger as the public compatibility facade while making OpenTelemetry the emission implementation:
- preserve
log, metric_exists, get_metric, get_metric_history, and close;
- retain local history because OpenTelemetry is an emission API rather than a query store;
- emit numeric measurements with OpenTelemetry instruments using stable
tunix.* names and low-cardinality attributes for mode and prefix;
- expose OTel-backed span and structured-event helpers through the same facade;
- use global OTel providers by default while allowing provider injection for tests and embedding;
- keep provider/exporter lifecycle owned by the application, never by
MetricsLogger;
- emit metrics from JAX process 0 by default, with an opt-in all-process policy;
- accept legacy
MetricsLoggerOptions construction during migration while deprecating exporter-specific backend fields;
- fix trainer ownership so a shared logger injected into actor and critic trainers is not closed by each trainer.
Compatibility and migration
The initial change should avoid modifying trainer metric call sites. Existing backend-specific configuration will need to migrate to standard OpenTelemetry SDK/exporter configuration. Performance tracing can then move incrementally onto the same facade without requiring a second public telemetry object.
Validation
Add in-memory OpenTelemetry tests for metric names, attributes, process policy, provider lifecycle, spans, events, scalar conversion, history compatibility, and shared logger ownership.
Problem
Tunix currently couples
MetricsLoggerto Metrax backends through process-global JAX monitoring listeners. The same class also owns in-memory metric history used by trainers and progress bars. This makes telemetry configuration backend-specific, gives individual logger instances process-global cleanup behavior, and does not provide a standard path for correlated metrics, traces, and structured events.Proposed design
Keep
MetricsLoggeras the public compatibility facade while making OpenTelemetry the emission implementation:log,metric_exists,get_metric,get_metric_history, andclose;tunix.*names and low-cardinality attributes for mode and prefix;MetricsLogger;MetricsLoggerOptionsconstruction during migration while deprecating exporter-specific backend fields;Compatibility and migration
The initial change should avoid modifying trainer metric call sites. Existing backend-specific configuration will need to migrate to standard OpenTelemetry SDK/exporter configuration. Performance tracing can then move incrementally onto the same facade without requiring a second public telemetry object.
Validation
Add in-memory OpenTelemetry tests for metric names, attributes, process policy, provider lifecycle, spans, events, scalar conversion, history compatibility, and shared logger ownership.