diff --git a/docs/metrics.md b/docs/metrics.md index 327412b34..86fdd1ee3 100644 --- a/docs/metrics.md +++ b/docs/metrics.md @@ -55,6 +55,7 @@ For Reinforcement Learning jobs, Tunix collects additional metrics related to the RL algorithm (e.g., PPO), reward modeling, and generation. #### Rewards & Scores + * **`rewards/sum`**: The sum of rewards for a trajectory. * **`rewards/mean`**, **`rewards/max`**, **`rewards/min`**: Statistics of the rewards across the batch. @@ -67,6 +68,7 @@ the RL algorithm (e.g., PPO), reward modeling, and generation. individual reward components are logged by name. #### Policy & Value (PPO) + * **`advantages/mean`**, **`advantages/max`**, **`advantages/min`**: Statistics of the advantages. * **`returns/mean`**, **`returns/max`**, **`returns/min`**: Statistics of @@ -81,6 +83,7 @@ the RL algorithm (e.g., PPO), reward modeling, and generation. enabled). #### Generation & Data + * **`prompts`**: The input prompts used for generation. * **`completions`**: The text completions generated by the model. * **`completions/mean_length`**, **`completions/max_length`**, @@ -210,6 +213,88 @@ depends on the execution environment. * `log_dir`: Directory where event files are written. * `flush_every_n_steps`: Frequency of flushing logs to disk (default: 100). +### Experimental: OpenTelemetry double-write + +Tunix is evaluating [OpenTelemetry](https://opentelemetry.io/) as a +vendor-neutral instrumentation layer. As an opt-in first step, `MetricsLogger` +supports a **double-write** mode: every scalar passed to `log()` is emitted both +through the existing Metrax/`jax.monitoring` backends (unchanged, still the +default) and as an OpenTelemetry gauge measurement. This lets you compare the +two pipelines side by side before any switch is made. + +Enable it with the `enable_opentelemetry` flag (requires `pip install +'google-tunix[otel]'`): + +```python +# OTLP exporter: pip install opentelemetry-exporter-otlp +from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import ( + OTLPMetricExporter, +) +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from tunix.sft import metrics_logger + +reader = PeriodicExportingMetricReader(OTLPMetricExporter()) +provider = MeterProvider(metric_readers=[reader]) + +options = metrics_logger.MetricsLoggerOptions( + log_dir="/tmp/logs", # Existing backends keep working unchanged. + enable_opentelemetry=True, # Additionally emit OpenTelemetry gauges. +) +logger = metrics_logger.MetricsLogger(options, otel_meter_provider=provider) +``` + +Key properties of the double-write path: + +* **Off by default.** With the flag unset, behavior is identical to previous + releases; the OpenTelemetry API does not even need to be installed. +* **Application-owned providers.** Tunix never configures exporters or shuts + down providers. Configure the global `MeterProvider` (or inject one via the + keyword-only `otel_meter_provider` argument) and own its flush/shutdown. + `MetricsLogger.close()` only closes the Metrax backends. +* **Process policy.** Like the Metrax backends, only JAX process 0 emits. +* **Naming.** Known metrics map to stable instrument names (`loss` → + `tunix.training.loss`, `perplexity` → `tunix.training.perplexity`, + `learning_rate` → `tunix.training.learning_rate`, `grad_norm` → + `tunix.training.gradient.norm`); other metric keys are normalized into the + `tunix.*` namespace (e.g. `rewards/score mean` → + `tunix.rewards.score.mean`). The legacy prefix and mode become the + low-cardinality attributes `tunix.metrics.prefix` and `tunix.training.mode`, + and the logical step is emitted as a separate `tunix.training.step` gauge. + +#### OpenTelemetry → Weights & Biases + +W&B ingests OpenTelemetry traces natively (via the +[Weave OTLP endpoint](https://docs.wandb.ai/weave/guides/tracking/otel)), but +has no OTLP endpoint for run metrics. Tunix therefore ships +`tunix.sft.otel_wandb.WandbMetricsExporter`, an OpenTelemetry SDK metric +exporter that forwards the gauges above to `wandb.log`, using the familiar +`{prefix}/{mode}/{name}` chart keys (e.g. `actor/train/tunix.training.loss`) and +the `tunix.training.step` gauge as the W&B step: + +```python +import wandb +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from tunix.sft import metrics_logger +from tunix.sft import otel_wandb + +run = wandb.init(project="my-project") +reader = PeriodicExportingMetricReader( + otel_wandb.WandbMetricsExporter(run), export_interval_millis=10_000 +) +provider = MeterProvider(metric_readers=[reader]) + +options = metrics_logger.MetricsLoggerOptions( + log_dir="/tmp/logs", + enable_opentelemetry=True, +) +logger = metrics_logger.MetricsLogger(options, otel_meter_provider=provider) +``` + +Because the Metrax `WandbBackend` also remains active by default, you can point +both pipelines at the same W&B project and compare the charts directly. + ### Custom metric logger You can integrate any logging service by creating a custom backend that conforms diff --git a/pyproject.toml b/pyproject.toml index b7220cea7..a6c1e6b21 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,10 +69,17 @@ test = [ "pytest-xdist", "gcsfs", "chex", + "opentelemetry-sdk>=1.25.0", + "wandb", # Offline end-to-end test for the OpenTelemetry W&B exporter ] cli = [ "tensorflow", ] +otel = [ + # Experimental OpenTelemetry double-write path for MetricsLogger. + "opentelemetry-api>=1.25.0", + "opentelemetry-sdk>=1.25.0", +] frozenlake = [ "gymnasium", ] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..ffa0e1013 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,41 @@ +"""Shared test configuration. + +The test environment installs ``wandb`` for the OpenTelemetry W&B exporter +end-to-end test. Installing it has a side effect on the legacy path: the +default Metrax ``WandbBackend`` becomes constructible, registers a +process-global ``jax.monitoring`` listener that receives JAX's own internal +compile metrics, and shares the module-global ``wandb.run`` across logger +instances — one logger's ``close()`` (``wandb.finish()``) then breaks any +other still-registered backend in the process. + +To keep the legacy default-backend behavior identical to an environment +without ``wandb`` (which is how CI has always run), the default +``WandbBackend`` is replaced with a stub that raises ``ImportError``, which +``MetricsLoggerOptions.create_backends`` already handles by skipping the +backend. Tests that exercise wandb do so explicitly: the OpenTelemetry +exporter tests pass an offline ``wandb.init`` run object directly and never +touch the Metrax backend. +""" + +import os + +import pytest + +os.environ.setdefault("WANDB_MODE", "disabled") + + +class _WandbBackendUnavailable: + + def __init__(self, *args, **kwargs): + raise ImportError( + "The default Metrax WandbBackend is disabled in tests; construct a" + " wandb run explicitly instead." + ) + + +@pytest.fixture(autouse=True) +def _disable_default_wandb_backend(monkeypatch): + monkeypatch.setattr( + "tunix.sft.metrics_logger.WandbBackend", _WandbBackendUnavailable + ) + yield diff --git a/tests/sft/metrics_logger_test.py b/tests/sft/metrics_logger_test.py index ae2ae5ec7..bcc3ab716 100644 --- a/tests/sft/metrics_logger_test.py +++ b/tests/sft/metrics_logger_test.py @@ -1,5 +1,6 @@ """Metrics logger unittest.""" +import collections import copy import os import tempfile @@ -8,6 +9,9 @@ from absl.testing import absltest import jax import metrax.logging as metrax_logging +import numpy as np +from opentelemetry.sdk import metrics as otel_sdk_metrics +from opentelemetry.sdk.metrics import export as otel_sdk_export from tunix.sft import metrics_logger from tunix.utils import env_utils @@ -229,5 +233,129 @@ def test_backend_kwargs_are_passed_to_backends(self): ) +def _gauge_points(metrics_data): + """Maps instrument name to a list of (value, attributes) data points.""" + points = collections.defaultdict(list) + if metrics_data is None: + return points + for resource_metrics in metrics_data.resource_metrics: + for scope_metrics in resource_metrics.scope_metrics: + for metric in scope_metrics.metrics: + for point in metric.data.data_points: + points[metric.name].append( + (getattr(point, "value", None), dict(point.attributes)) + ) + return points + + +class OpenTelemetryDoubleWriteTest(absltest.TestCase): + """Tests the opt-in OpenTelemetry double-write path.""" + + def setUp(self): + super().setUp() + self._temp_dir_obj = tempfile.TemporaryDirectory() + self.log_dir = self._temp_dir_obj.name + self.backend = mock.Mock(spec=metrax_logging.LoggingBackend) + self.enter_context( + mock.patch.object(jax.monitoring, "register_scalar_listener") + ) + self.mock_record_scalar = self.enter_context( + mock.patch.object(jax.monitoring, "record_scalar") + ) + self.reader = otel_sdk_export.InMemoryMetricReader() + self.meter_provider = otel_sdk_metrics.MeterProvider( + metric_readers=[self.reader] + ) + + def _make_logger(self, enable_opentelemetry=True): + options = metrics_logger.MetricsLoggerOptions( + log_dir=self.log_dir, + backend_kwargs={"custom_backend": [lambda: self.backend]}, + enable_opentelemetry=enable_opentelemetry, + ) + return metrics_logger.MetricsLogger( + metrics_logger_options=options, + otel_meter_provider=self.meter_provider, + ) + + def test_disabled_by_default_and_legacy_path_unchanged(self): + logger = self._make_logger(enable_opentelemetry=False) + logger.log("actor", "loss", 0.5, metrics_logger.Mode.TRAIN, 1) + + self.mock_record_scalar.assert_called_once_with( + "actor/train/loss", 0.5, step=1 + ) + self.assertEqual(_gauge_points(self.reader.get_metrics_data()), {}) + self.assertAlmostEqual(logger.get_metric("actor", "loss", "train"), 0.5) + + def test_double_write_emits_otel_and_legacy(self): + logger = self._make_logger() + logger.log("actor", "loss", 0.5, metrics_logger.Mode.TRAIN, 7) + + # Legacy jax.monitoring write is unchanged. + self.mock_record_scalar.assert_called_once_with( + "actor/train/loss", 0.5, step=7 + ) + # OpenTelemetry gauge is additionally emitted. + points = _gauge_points(self.reader.get_metrics_data()) + expected_attributes = { + "tunix.training.mode": "train", + "tunix.metrics.prefix": "actor", + } + self.assertEqual( + points["tunix.training.loss"], [(0.5, expected_attributes)] + ) + self.assertEqual(points["tunix.training.step"], [(7, expected_attributes)]) + # Local history is unchanged. + self.assertAlmostEqual(logger.get_metric("actor", "loss", "train"), 0.5) + + def test_dynamic_metric_names_are_normalized(self): + logger = self._make_logger() + logger.log("actor", "rewards/score mean", 1.5, "train", 1) + + points = _gauge_points(self.reader.get_metrics_data()) + self.assertIn("tunix.rewards.score.mean", points) + + def test_non_scalar_values_stay_in_history_only(self): + logger = self._make_logger() + logger.log("actor", "loss", np.array([0.5, 0.6]), "train", 1) + + points = _gauge_points(self.reader.get_metrics_data()) + self.assertNotIn("tunix.training.loss", points) + self.assertLen(logger.get_metric_history("actor", "loss", "train"), 1) + + def test_repeated_step_emits_step_gauge_once(self): + logger = self._make_logger() + logger.log("actor", "loss", 0.5, "train", 3) + logger.log("actor", "perplexity", 1.6, "train", 3) + + points = _gauge_points(self.reader.get_metrics_data()) + self.assertLen(points["tunix.training.step"], 1) + + @mock.patch.object(jax, "process_index", return_value=1) + def test_no_otel_emission_on_secondary_process(self, mock_process_index): + del mock_process_index + logger = self._make_logger() + logger.log("actor", "loss", 0.5, "train", 1) + + self.assertEqual(_gauge_points(self.reader.get_metrics_data()), {}) + + def test_missing_otel_api_raises_when_enabled(self): + with mock.patch.object(metrics_logger, "otel_metrics", new=None): + with self.assertRaisesRegex(ImportError, "google-tunix\\[otel\\]"): + self._make_logger() + + def test_close_leaves_meter_provider_running(self): + logger = self._make_logger() + logger.log("actor", "loss", 0.5, "train", 1) + logger.close() + + self.backend.close.assert_called_once() + logger_after_close = self._make_logger() + logger_after_close.log("actor", "loss", 0.25, "train", 2) + points = _gauge_points(self.reader.get_metrics_data()) + self.assertIn(0.25, [value for value, _ in points["tunix.training.loss"]]) + + if __name__ == "__main__": absltest.main() diff --git a/tests/sft/otel_wandb_test.py b/tests/sft/otel_wandb_test.py new file mode 100644 index 000000000..f30339e64 --- /dev/null +++ b/tests/sft/otel_wandb_test.py @@ -0,0 +1,256 @@ +"""Tests for the OpenTelemetry to Weights & Biases metric exporter.""" + +import glob +import json +import os +import sys +import tempfile +from unittest import mock + +from absl.testing import absltest +import jax +import metrax.logging as metrax_logging +from opentelemetry.sdk import metrics as otel_sdk_metrics +from opentelemetry.sdk.metrics import export as otel_sdk_export +from tunix.sft import metrics_logger +from tunix.sft import otel_wandb + +try: + import wandb # pytype: disable=import-error # pylint: disable=g-import-not-at-top +except ImportError: + wandb = None + + +class _FakeWandbRun: + """Captures wandb.log calls.""" + + def __init__(self): + self.calls = [] + + def log(self, data, step=None): + self.calls.append((dict(data), step)) + + +class WandbMetricsExporterTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self.run = _FakeWandbRun() + self.exporter = otel_wandb.WandbMetricsExporter(self.run) + self.reader = otel_sdk_export.PeriodicExportingMetricReader( + self.exporter, export_interval_millis=3_600_000 + ) + self.meter_provider = otel_sdk_metrics.MeterProvider( + metric_readers=[self.reader] + ) + self.enter_context( + mock.patch.object(jax.monitoring, "register_scalar_listener") + ) + self.enter_context(mock.patch.object(jax.monitoring, "record_scalar")) + + def tearDown(self): + self.meter_provider.shutdown() + super().tearDown() + + def _temp_dir(self): + temp_dir = tempfile.TemporaryDirectory() + self.addCleanup(temp_dir.cleanup) + return temp_dir.name + + def _make_logger(self): + backend = mock.Mock(spec=metrax_logging.LoggingBackend) + options = metrics_logger.MetricsLoggerOptions( + log_dir=self._temp_dir(), + backend_kwargs={"custom_backend": [lambda: backend]}, + enable_opentelemetry=True, + ) + return metrics_logger.MetricsLogger( + metrics_logger_options=options, + otel_meter_provider=self.meter_provider, + ) + + def test_double_write_reaches_wandb_run(self): + logger = self._make_logger() + logger.log("actor", "loss", 0.5, metrics_logger.Mode.TRAIN, 3) + logger.log("actor", "perplexity", 1.65, metrics_logger.Mode.TRAIN, 3) + + self.meter_provider.force_flush() + + self.assertLen(self.run.calls, 1) + values, step = self.run.calls[0] + self.assertEqual(step, 3) + self.assertAlmostEqual(values["actor/train/tunix.training.loss"], 0.5) + self.assertAlmostEqual( + values["actor/train/tunix.training.perplexity"], 1.65 + ) + + def test_groups_are_logged_in_step_order(self): + logger = self._make_logger() + logger.log("critic", "loss", 0.7, metrics_logger.Mode.EVAL, 9) + logger.log("actor", "loss", 0.5, metrics_logger.Mode.TRAIN, 2) + + self.meter_provider.force_flush() + + steps = [step for _, step in self.run.calls] + self.assertEqual(steps, sorted(steps)) + logged_keys = set() + for values, _ in self.run.calls: + logged_keys.update(values) + self.assertIn("actor/train/tunix.training.loss", logged_keys) + self.assertIn("critic/eval/tunix.training.loss", logged_keys) + + def _collect_metrics_data(self): + """Builds one metric export batch through an in-memory reader.""" + reader = otel_sdk_export.InMemoryMetricReader() + provider = otel_sdk_metrics.MeterProvider(metric_readers=[reader]) + logger = metrics_logger.MetricsLogger( + metrics_logger_options=metrics_logger.MetricsLoggerOptions( + log_dir=self._temp_dir(), + backend_kwargs={ + "custom_backend": [ + lambda: mock.Mock(spec=metrax_logging.LoggingBackend) + ] + }, + enable_opentelemetry=True, + ), + otel_meter_provider=provider, + ) + logger.log("actor", "loss", 0.5, "train", 1) + return reader.get_metrics_data() + + def test_export_failure_is_reported_not_raised(self): + failing_run = mock.Mock() + failing_run.log.side_effect = RuntimeError("wandb unavailable") + exporter = otel_wandb.WandbMetricsExporter(failing_run) + + result = exporter.export(self._collect_metrics_data()) + + self.assertEqual(result, otel_sdk_export.MetricExportResult.FAILURE) + + def test_resolves_global_wandb_run_lazily(self): + fake_wandb = mock.Mock() + fake_wandb.run = _FakeWandbRun() + exporter = otel_wandb.WandbMetricsExporter() + metrics_data = self._collect_metrics_data() + + with mock.patch.dict(sys.modules, {"wandb": fake_wandb}): + result = exporter.export(metrics_data) + + self.assertEqual(result, otel_sdk_export.MetricExportResult.SUCCESS) + self.assertLen(fake_wandb.run.calls, 1) + + +class _CountingBackend: + """Legacy Metrax-protocol backend recording jax.monitoring scalars.""" + + def __init__(self): + self.scalars = [] + + def log_scalar(self, event, value, **kwargs): + self.scalars.append((event, value, kwargs.get("step"))) + + def close(self): + pass + + +@absltest.skipIf(wandb is None, "wandb is not installed") +class WandbOfflineEndToEndTest(absltest.TestCase): + """Double-writes through the real wandb client using an offline run. + + ``wandb.init(mode="offline")`` exercises the full wandb client (run setup, + ``log`` semantics, step alignment, transaction-log serialization) without + needing an account or network access. The test reads the resulting + ``.wandb`` transaction log back with wandb's own reader and asserts the + per-step history, while a legacy backend on the same logger proves both + sides of the double-write. + """ + + def _read_history(self, wandb_dir): + """Maps step to logged values from the offline run's transaction log.""" + # pylint: disable=g-import-not-at-top + from wandb.proto import wandb_internal_pb2 # pytype: disable=import-error + from wandb.sdk.internal import datastore # pytype: disable=import-error + + # pylint: enable=g-import-not-at-top + + wandb_file = glob.glob( + os.path.join(wandb_dir, "wandb", "offline-run-*", "run-*.wandb") + )[0] + store = datastore.DataStore() + store.open_for_scan(wandb_file) + history = {} + while True: + data = store.scan_data() + if data is None: + break + record = wandb_internal_pb2.Record() + record.ParseFromString(data) + if record.WhichOneof("record_type") == "history": + items = {} + for item in record.history.item: + key = ".".join(item.nested_key) if item.nested_key else item.key + items[key] = json.loads(item.value_json) + history[items["_step"]] = items + return history + + def test_double_write_reaches_real_offline_wandb_run(self): + assert wandb is not None + wandb_dir = tempfile.TemporaryDirectory() + self.addCleanup(wandb_dir.cleanup) + self.enter_context( + mock.patch.dict( + os.environ, {"WANDB_MODE": "offline", "WANDB_SILENT": "true"} + ) + ) + run = wandb.init( + mode="offline", + project="tunix-otel-e2e", + dir=wandb_dir.name, + settings=wandb.Settings(console="off"), + ) + reader = otel_sdk_export.PeriodicExportingMetricReader( + otel_wandb.WandbMetricsExporter(run), export_interval_millis=3_600_000 + ) + provider = otel_sdk_metrics.MeterProvider(metric_readers=[reader]) + self.addCleanup(provider.shutdown) + + legacy_backend = _CountingBackend() + logger = metrics_logger.MetricsLogger( + metrics_logger_options=metrics_logger.MetricsLoggerOptions( + log_dir=os.path.join(wandb_dir.name, "logs"), + backend_kwargs={"custom_backend": [lambda: legacy_backend]}, + enable_opentelemetry=True, + ), + otel_meter_provider=provider, + ) + # close() clears the process-global scalar listener this logger registers; + # make sure that happens even when an assertion below fails first. + self.addCleanup(logger.close) + + for step, loss in enumerate([2.31, 1.87, 1.52], start=1): + logger.log("actor", "loss", loss, metrics_logger.Mode.TRAIN, step) + logger.log("actor", "rewards/score mean", 0.1 * step, "train", step) + provider.force_flush() + logger.log("actor", "loss", 1.61, metrics_logger.Mode.EVAL, 3) + provider.force_flush() + run.finish() + + history = self._read_history(wandb_dir.name) + self.assertEqual( + [ + history[step]["actor/train/tunix.training.loss"] + for step in (1, 2, 3) + ], + [2.31, 1.87, 1.52], + ) + self.assertAlmostEqual( + history[3]["actor/train/tunix.rewards.score.mean"], 0.3 + ) + self.assertEqual(history[3]["actor/eval/tunix.training.loss"], 1.61) + # The legacy jax.monitoring path received every scalar in the same run. + self.assertLen(legacy_backend.scalars, 7) + self.assertIn(("actor/train/loss", 2.31, 1), legacy_backend.scalars) + + +if __name__ == "__main__": + absltest.main() diff --git a/tunix/sft/metrics_logger.py b/tunix/sft/metrics_logger.py index abb87f1f2..50ace9bae 100644 --- a/tunix/sft/metrics_logger.py +++ b/tunix/sft/metrics_logger.py @@ -1,8 +1,23 @@ -"""Metric logger with a unified, protocol-based backend system.""" +"""Metric logger with a unified, protocol-based backend system. + +In addition to the Metrax logging backends, ``MetricsLogger`` supports an +opt-in OpenTelemetry double-write path: when +``MetricsLoggerOptions.enable_opentelemetry`` is set, every scalar logged +through ``log`` is emitted both through the existing ``jax.monitoring`` +backends and as an OpenTelemetry gauge measurement. The existing backends +remain the default and are unaffected by the flag. + +Tunix does not configure or shut down OpenTelemetry providers. Applications +configure the global meter provider (readers, exporters) or inject a provider +when constructing a ``MetricsLogger``. +""" import collections import dataclasses import enum +import importlib.metadata +import re +import threading from typing import Any, Callable from absl import logging @@ -11,6 +26,11 @@ import numpy as np from tunix.utils import env_utils +try: + from opentelemetry import metrics as otel_metrics # pylint: disable=g-import-not-at-top +except ImportError: + otel_metrics = None + LoggingBackend = metrax_logging.LoggingBackend TensorboardBackend = metrax_logging.TensorboardBackend WandbBackend = metrax_logging.WandbBackend @@ -19,6 +39,84 @@ # User backends MUST be factories (callables) to keep Options pure and copyable. BackendFactory = Callable[[], LoggingBackend] +_OTEL_INSTRUMENTATION_NAME = "tunix" +_OTEL_METRIC_NAME_PATTERN = re.compile(r"[^a-zA-Z0-9_.-]+") + +try: + _OTEL_INSTRUMENTATION_VERSION = importlib.metadata.version("google-tunix") +except importlib.metadata.PackageNotFoundError: + _OTEL_INSTRUMENTATION_VERSION = None + + +@dataclasses.dataclass(frozen=True, slots=True) +class _OtelMetricSpec: + name: str + unit: str = "1" + description: str = "" + + +_OTEL_KNOWN_METRICS = { + "loss": _OtelMetricSpec( + name="tunix.training.loss", + description="Loss reported by a Tunix training or evaluation step.", + ), + "perplexity": _OtelMetricSpec( + name="tunix.training.perplexity", + description=( + "Perplexity reported by a Tunix training or evaluation step." + ), + ), + "learning_rate": _OtelMetricSpec( + name="tunix.training.learning_rate", + description="Learning rate used by the Tunix optimizer.", + ), + "grad_norm": _OtelMetricSpec( + name="tunix.training.gradient.norm", + description="Gradient norm reported by a Tunix training step.", + ), +} + +_OTEL_STEP_METRIC = _OtelMetricSpec( + name="tunix.training.step", + unit="{step}", + description="Current logical Tunix training or evaluation step.", +) + + +def _normalize_otel_metric_name(metric_name: str) -> str: + """Builds a valid, namespaced OpenTelemetry instrument name.""" + normalized = _OTEL_METRIC_NAME_PATTERN.sub(".", metric_name).strip(".") + if not normalized: + normalized = "unnamed" + return f"tunix.{normalized.lower()}" + + +def _otel_metric_spec(metric_name: str) -> _OtelMetricSpec: + """Returns the stable specification for a Tunix metric.""" + return _OTEL_KNOWN_METRICS.get( + metric_name, + _OtelMetricSpec( + name=_normalize_otel_metric_name(metric_name), + description=f"Tunix metric derived from {metric_name!r}.", + ), + ) + + +def _as_otel_number(value: Any) -> int | float | None: + """Converts a scalar value to an OpenTelemetry-compatible number.""" + try: + array = np.asarray(value) + except Exception: # pylint: disable=broad-exception-caught + return None + if array.size != 1: + return None + scalar = array.reshape(()).item() + if isinstance(scalar, (bool, np.bool_)): + return None + if isinstance(scalar, (int, float, np.integer, np.floating)): + return scalar.item() if isinstance(scalar, np.generic) else scalar + return None + @dataclasses.dataclass class MetricsLoggerOptions: @@ -48,6 +146,12 @@ class MetricsLoggerOptions: backend_kwargs: dict[str, dict[str, Any] | list[BackendFactory]] = ( dataclasses.field(default_factory=dict) ) + # Experimental: additionally emit every logged scalar as an OpenTelemetry + # gauge measurement (double-write). The Metrax backends above keep working + # unchanged; this flag only adds a second emission path. Requires the + # OpenTelemetry API ("pip install google-tunix[otel]"). Exporters and + # provider lifecycle are owned by the application, not by Tunix. + enable_opentelemetry: bool = False def create_backends(self) -> list[LoggingBackend]: """Factory method to create a fresh set of live backends.""" @@ -91,8 +195,8 @@ def create_backends(self) -> list[LoggingBackend]: **wandb_kwargs, # pyrefly: ignore[bad-unpacking] ) ) - except ImportError: - logging.info("WandbBackend skipped: 'wandb' library not installed.") + except (ImportError, Exception) as e: + logging.info("WandbBackend skipped: %s", e) return active_backends @@ -113,13 +217,26 @@ class MetricsLogger: """Simple Metrics logger. Log metrics to multiple backends. If no backends are specified, it will log to - the default backends. + the default backends. When ``MetricsLoggerOptions.enable_opentelemetry`` is + set, every scalar is additionally emitted as an OpenTelemetry gauge + measurement (double-write) without changing the backend behavior. """ def __init__( self, metrics_logger_options: MetricsLoggerOptions | None = None, + *, + otel_meter_provider: Any = None, ): + """Initializes the metrics logger. + + Args: + metrics_logger_options: Logger configuration. ``None`` disables backend + creation and OpenTelemetry emission; local metric history still works. + otel_meter_provider: Optional OpenTelemetry ``MeterProvider`` used when + ``enable_opentelemetry`` is set. Defaults to the global provider. This + keyword-only argument exists for tests and embedding applications. + """ self._metrics = {} self._backends = ( metrics_logger_options.create_backends() @@ -130,6 +247,23 @@ def __init__( for backend in self._backends: jax.monitoring.register_scalar_listener(backend.log_scalar) + self._otel_meter = None + self._otel_instruments: dict[str, Any] = {} + self._otel_instrument_lock = threading.Lock() + self._otel_last_steps: dict[tuple[str, str], int] = {} + if metrics_logger_options and metrics_logger_options.enable_opentelemetry: + if otel_metrics is None: + raise ImportError( + "MetricsLoggerOptions.enable_opentelemetry is set but the" + " OpenTelemetry API is not installed. Install it with: pip install" + " 'google-tunix[otel]'." + ) + provider = otel_meter_provider or otel_metrics.get_meter_provider() + self._otel_meter = provider.get_meter( + _OTEL_INSTRUMENTATION_NAME, + _OTEL_INSTRUMENTATION_VERSION, + ) + def log( self, metrics_prefix: str, @@ -149,6 +283,73 @@ def log( f"{metrics_prefix}/{mode}/{metric_name}", scalar_value, step=step # pyrefly: ignore[bad-argument-type] ) + if self._otel_meter is not None: + # The OpenTelemetry double-write is best effort: emission failures must + # not interrupt training or the primary logging path. + try: + self._emit_otel_metric( + metrics_prefix, metric_name, scalar_value, mode, step + ) + except Exception: # pylint: disable=broad-exception-caught + logging.exception( + "Failed to emit OpenTelemetry metric %s/%s/%s.", + metrics_prefix, + mode, + metric_name, + ) + + def _otel_gauge(self, spec: _OtelMetricSpec): + """Returns a cached synchronous gauge for a metric specification.""" + with self._otel_instrument_lock: + gauge = self._otel_instruments.get(spec.name) + if gauge is None: + assert self._otel_meter is not None + gauge = self._otel_meter.create_gauge( + spec.name, + unit=spec.unit, + description=spec.description, + ) + self._otel_instruments[spec.name] = gauge + return gauge + + def _emit_otel_metric( + self, + metrics_prefix: str, + metric_name: str, + scalar_value: Any, + mode: Mode | str, + step: int, + ) -> None: + """Emits one OpenTelemetry metric value and its associated logical step.""" + # Match the backend policy: only the main process emits telemetry. + if jax.process_index() != 0: + return + + value = _as_otel_number(scalar_value) + if value is None: + logging.warning( + "Skipping non-scalar OpenTelemetry metric %s/%s/%s with shape %s.", + metrics_prefix, + mode, + metric_name, + np.shape(scalar_value), + ) + return + + attributes = {"tunix.training.mode": str(mode)} + if metrics_prefix: + attributes["tunix.metrics.prefix"] = metrics_prefix + self._otel_gauge(_otel_metric_spec(metric_name)).set(value, attributes) + + # The logical step is a separate gauge rather than a metric attribute: an + # attribute value per step would create unbounded time-series cardinality. + step_key = (metrics_prefix, str(mode)) + with self._otel_instrument_lock: + if self._otel_last_steps.get(step_key) == step: + return + self._otel_last_steps[step_key] = step + self._otel_gauge(_OTEL_STEP_METRIC).set(step, attributes) + def metric_exists( self, metrics_prefix, metric_name: str, mode: Mode | str ) -> bool: @@ -183,7 +384,12 @@ def get_metric_history( return np.stack(self._metrics[metrics_prefix][mode][metric_name]) def close(self): - """Closes all registered logging backends.""" + """Closes all registered logging backends. + + OpenTelemetry providers are application-owned process-wide resources; the + application that configured them owns flushing and shutdown, so ``close`` + leaves them running. + """ for backend in self._backends: backend.close() try: diff --git a/tunix/sft/otel_wandb.py b/tunix/sft/otel_wandb.py new file mode 100644 index 000000000..68fc7c4c0 --- /dev/null +++ b/tunix/sft/otel_wandb.py @@ -0,0 +1,157 @@ +"""OpenTelemetry metric exporter for Weights & Biases runs. + +Weights & Biases natively ingests OpenTelemetry traces (through the Weave OTLP +endpoint) but has no OTLP endpoint for run metrics. This module provides an +OpenTelemetry SDK ``MetricExporter`` that forwards Tunix gauge measurements to +``wandb.log`` so the OpenTelemetry double-write path of +``tunix.sft.metrics_logger.MetricsLogger`` can feed a W&B run. + +Requires the ``opentelemetry-sdk`` (``pip install 'google-tunix[otel]'``) and +``wandb`` packages. Example: + +.. code-block:: python + + import wandb + from opentelemetry.sdk.metrics import MeterProvider + from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader + from tunix.sft import metrics_logger + from tunix.sft import otel_wandb + + run = wandb.init(project="my-project") + reader = PeriodicExportingMetricReader( + otel_wandb.WandbMetricsExporter(run), export_interval_millis=10_000 + ) + provider = MeterProvider(metric_readers=[reader]) + + options = metrics_logger.MetricsLoggerOptions( + log_dir="/tmp/logs", enable_opentelemetry=True + ) + logger = metrics_logger.MetricsLogger( + options, otel_meter_provider=provider + ) +""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from typing import Any, TypedDict + +from absl import logging +from opentelemetry.sdk.metrics import export as otel_export + +_STEP_METRIC_NAME = "tunix.training.step" +_PREFIX_ATTRIBUTE = "tunix.metrics.prefix" +_MODE_ATTRIBUTE = "tunix.training.mode" + + +class _Group(TypedDict): + step: int | None + values: dict[str, Any] + + +def _wandb_key(prefix: str, mode: str, metric_name: str) -> str: + """Builds a W&B chart key mirroring the legacy `{prefix}/{mode}/{name}`.""" + return "/".join(part for part in (prefix, mode, metric_name) if part) + + +def _wandb_payloads( + metrics_data: otel_export.MetricsData, +) -> Iterator[tuple[int | None, dict[str, Any]]]: + """Groups gauge data points into per-step W&B log payloads. + + Data points are grouped by their Tunix prefix and mode attributes so each + group can be logged against the logical training step carried by the + ``tunix.training.step`` gauge of the same group. Groups are yielded in + ascending step order because W&B requires non-decreasing steps. + + Args: + metrics_data: One OpenTelemetry metric export batch. + + Yields: + ``(step, values)`` pairs, where ``step`` may be ``None`` when the batch + carries no step gauge for the group. + """ + groups: dict[tuple[str, str], _Group] = {} + for resource_metrics in metrics_data.resource_metrics: + for scope_metrics in resource_metrics.scope_metrics: + for metric in scope_metrics.metrics: + for point in getattr(metric.data, "data_points", ()): + attributes: Mapping[str, Any] = point.attributes or {} + group_key = ( + str(attributes.get(_PREFIX_ATTRIBUTE, "")), + str(attributes.get(_MODE_ATTRIBUTE, "")), + ) + group = groups.setdefault(group_key, {"step": None, "values": {}}) + if metric.name == _STEP_METRIC_NAME: + group["step"] = int(point.value) + else: + key = _wandb_key(group_key[0], group_key[1], metric.name) + group["values"][key] = point.value + ordered = sorted( + (group for group in groups.values() if group["values"]), + key=lambda group: (group["step"] is None, group["step"]), + ) + for group in ordered: + yield group["step"], group["values"] + + +class WandbMetricsExporter(otel_export.MetricExporter): + """Exports OpenTelemetry gauge data points to a Weights & Biases run. + + Chart keys mirror the legacy Metrax backend layout, + ``{prefix}/{mode}/{instrument name}`` (for example + ``actor/train/tunix.training.loss``), so double-written runs are easy to + compare side by side. The ``tunix.training.step`` gauge provides the W&B + step for the other measurements in its group. + """ + + def __init__(self, run: Any = None, **kwargs): + """Initializes the exporter. + + Args: + run: The W&B run to log to (the object returned by ``wandb.init``). + When ``None``, the active global ``wandb.run`` is resolved at export + time. + **kwargs: Forwarded to ``MetricExporter`` (for example + ``preferred_temporality``). + """ + super().__init__(**kwargs) + self._run = run + + def _resolve_run(self) -> Any: + if self._run is not None: + return self._run + import wandb # pytype: disable=import-error # pylint: disable=g-import-not-at-top + + if wandb.run is None: + raise RuntimeError( + "WandbMetricsExporter has no W&B run to log to. Pass a run to the" + " exporter or call wandb.init() before exporting." + ) + return wandb.run + + def export( + self, + metrics_data: otel_export.MetricsData, + timeout_millis: float = 10_000, + **kwargs, + ) -> otel_export.MetricExportResult: + """Logs one metric export batch to the W&B run.""" + del timeout_millis, kwargs + try: + run = self._resolve_run() + for step, values in _wandb_payloads(metrics_data): + run.log(values, step=step) + except Exception: # pylint: disable=broad-exception-caught + logging.exception("Failed to export OpenTelemetry metrics to W&B.") + return otel_export.MetricExportResult.FAILURE + return otel_export.MetricExportResult.SUCCESS + + def force_flush(self, timeout_millis: float = 10_000) -> bool: + """Nothing to flush; ``export`` logs synchronously.""" + del timeout_millis + return True + + def shutdown(self, timeout_millis: float = 30_000, **kwargs) -> None: + """Leaves the W&B run open; the application owns ``run.finish()``.""" + del timeout_millis, kwargs