make trace writing a configurable option

PiperOrigin-RevId: 869369502

Author: s-noghabi

tests/perf/export_test.py

@@ -15,6 +15,7 @@
 from unittest import mock
 
 from absl.testing import absltest
+from absl.testing import parameterized
 import jax
 import numpy as np
 from tunix.perf import export
@@ -31,12 +32,18 @@
 
 
 def _create_mock_cluster_config_with_perf_metrics(
-    perf_metrics_log_dir: str | None,
+    *,
+    enable_perf_metrics: bool = True,
+    enable_trace_writer: bool = True,
+    perf_metrics_log_dir: str = "",
 ) -> mock.Mock:
   """Creates a mock ClusterConfig object for testing.
 
   Args:
-    perf_metrics_log_dir: the log directory for perf metrics.
+    enable_perf_metrics: If True, enables perf metrics.
+    enable_trace_writer: If True, enables trace writer to write out the trace
+      timeline.
+    perf_metrics_log_dir: The log directory for perf metrics.
 
   Returns:
     A mock ClusterConfig object.
@@ -60,48 +67,89 @@ def _create_mock_cluster_config_with_perf_metrics(
   )
   cluster_config.training_config = mock_training_config
 
-  if perf_metrics_log_dir:
+  if enable_perf_metrics:
     mock_options = mock.create_autospec(
         metrics.PerfMetricsOptions, instance=True
     )
     mock_options.log_dir = perf_metrics_log_dir
+    mock_options.enable_trace_writer = enable_trace_writer
     mock_training_config.perf_metrics_options = mock_options
   else:
     mock_training_config.perf_metrics_options = None
   return cluster_config
 
 
-class ExportTest(absltest.TestCase):
+class ExportTest(parameterized.TestCase):
 
-  def test_from_cluster_config_with_export_dir(self):
-    log_dir = "/tmp/test_log_dir"
-    cluster_config = _create_mock_cluster_config_with_perf_metrics(log_dir)
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="with_export_dir",
+          perf_metrics_log_dir="test_log_dir",
+          expected_log_dir="test_log_dir",
+      ),
+      dict(
+          testcase_name="without_export_dir",
+          perf_metrics_log_dir="",
+          expected_log_dir=None,
+      ),
+  )
+  def test_from_cluster_config_trace_writer_enabled(
+      self, perf_metrics_log_dir, expected_log_dir
+  ):
+    cluster_config = _create_mock_cluster_config_with_perf_metrics(
+        perf_metrics_log_dir=perf_metrics_log_dir
+    )
     with mock.patch.object(
-        export, "PerfettoTraceWriter", autospec=True
+        export, "PerfettoTraceWriter", autospec=True, spec_set=True
     ) as mock_writer:
       PerfMetricsExport.from_cluster_config(cluster_config)
-      mock_writer.assert_called_with(log_dir)
-
-  def test_from_cluster_config_without_export_dir(self):
+      mock_writer.assert_called_with(expected_log_dir)
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="disabled_trace_writer",
+          enable_perf_metrics=True,
+          enable_trace_writer=False,
+      ),
+      dict(
+          testcase_name="disabled_perf_metrics",
+          enable_perf_metrics=False,
+          enable_trace_writer=True,
+      ),
+  )
+  def test_from_cluster_config_trace_writer_disabled(
+      self, enable_perf_metrics, enable_trace_writer
+  ):
+    log_dir = "test_log_dir"
     cluster_config = _create_mock_cluster_config_with_perf_metrics(
-        perf_metrics_log_dir=None
+        perf_metrics_log_dir=log_dir,
+        enable_perf_metrics=enable_perf_metrics,
+        enable_trace_writer=enable_trace_writer,
     )
     with mock.patch.object(
-        export, "PerfettoTraceWriter", autospec=True
+        export, "PerfettoTraceWriter", autospec=True, spec_set=True
     ) as mock_writer:
       PerfMetricsExport.from_cluster_config(cluster_config)
-      mock_writer.assert_called_with(None)
+      mock_writer.assert_not_called()
 
   @patch("time.perf_counter")
-  def test_export_grpo_metrics_colocated(self, mock_perf_counter):
+  def test_export_grpo_metrics_colocated_with_trace_writer(
+      self, mock_perf_counter
+  ):
     # tpu0 span end times
     mock_perf_counter.side_effect = [0.41, 0.61, 1.21]
+    mock_trace_writer = mock.create_autospec(
+        export.PerfettoTraceWriter, instance=True, spec_set=True
+    )
 
-    export_fn = PerfMetricsExport.from_role_to_devices({
-        "rollout": ["tpu0"],
-        "refer": ["tpu0"],
-        "actor": ["tpu0"],
-    })
+    export_fn = PerfMetricsExport.from_role_to_devices(
+        {
+            "rollout": ["tpu0"],
+            "refer": ["tpu0"],
+            "actor": ["tpu0"],
+        },
+        trace_writer=mock_trace_writer,
+    )
     host_timeline = ThreadTimeline("host", 0.0)
     tpu0_timeline = DeviceTimeline("tpu0", 0.0)
     timelines = {
@@ -165,7 +213,11 @@ def test_export_grpo_metrics_colocated(self, mock_perf_counter):
     for k, v in export_fn(PerfSpanQuery(timelines, "host")).items():
       actual_metrics[k] = float(v[0])
 
-    self.assertDictAlmostEqual(actual_metrics, expected_metrics)
+    with self.subTest("metrics"):
+      self.assertDictAlmostEqual(actual_metrics, expected_metrics)
+
+    with self.subTest("trace_logging"):
+      mock_trace_writer.log_trace.assert_called_once()
 
   @patch("time.perf_counter")
   def test_export_grpo_metrics_rollout_1_actor_2_reference_2(

tunix/cli/base_config.yaml

@@ -147,6 +147,7 @@ training_config: &base_training_config
   # Currently only supported for grpo_main. Configs you can specify are:
     # log_dir = the directory to save the trace files.
     # custom_export_fn_path = the path to the custom export function.
+    # enable_trace_writer = whether to enable writing out the trace for visualization (enabled by default).
   perf_metrics_options: {}
   profiler_options:
     log_dir: "/tmp/profiling"

tunix/perf/export.py

@@ -83,7 +83,6 @@ def from_role_to_devices(
       role_to_devices: A dictionary mapping role names to a list of device
         identifiers.
       trace_writer: An optional PerfettoTraceWriter to log performance traces.
-        If None, a default writer is created.
       log_rollout_time_at_micro_batch_level: Whether to log rollout time at the
         micro batch level. This is a temporary flag. It will be removed once
         metrics are exported at the micro batch.
@@ -95,11 +94,6 @@ def from_role_to_devices(
       A callable function that takes a PerfSpanQuery and returns MetricsT.
     """
 
-    if trace_writer is None:
-      # If no trace writer is provided, create a default one.
-      logging.info("Creating a default trace writer for metrics export.")
-      trace_writer = PerfettoTraceWriter(None)
-
     r2d = role_to_devices
     if r2d["rollout"] == r2d["actor"] == r2d["refer"]:
       logging.info(
@@ -169,19 +163,25 @@ def from_cluster_config(
         trace.create_device_timeline_id, refer_mesh.devices.flatten().tolist()
     )
 
-    export_dir = (
-        cluster_config.training_config.perf_metrics_options.log_dir
-        if cluster_config.training_config.perf_metrics_options
-        else None  # A default directory will be used in this case.
-    )
+    perf_metrics_options = cluster_config.training_config.perf_metrics_options
+    if (
+        perf_metrics_options is not None
+        and perf_metrics_options.enable_trace_writer
+    ):
+      # Setting export_dir to None will cause the trace writer to use a
+      # default directory.
+      export_dir = perf_metrics_options.log_dir or None
+      trace_writer = PerfettoTraceWriter(export_dir)
+    else:
+      trace_writer = None
 
     return PerfMetricsExport.from_role_to_devices(
         role_to_devices={
             "rollout": list(rollo_devices),
             "actor": list(actor_devices),
             "refer": list(refer_devices),
         },
-        trace_writer=PerfettoTraceWriter(export_dir),
+        trace_writer=trace_writer,
         log_rollout_time_at_micro_batch_level=log_rollout_time_at_micro_batch_level,
         log_actor_train_time_at_micro_batch_level=log_actor_train_time_at_micro_batch_level,
     )
@@ -196,7 +196,7 @@ def create_metrics_export_fn(
   @staticmethod
   def _grpo_metrics_colocated(
       extract_spans_fn: _GrpoExtractSpansFn,
-      trace_writer: PerfettoTraceWriter,
+      trace_writer: PerfettoTraceWriter | None,
       query: PerfSpanQuery,
   ) -> MetricsT:
     """GRPO workflow: rollout, actor and reference are colocated on the same mesh.
@@ -251,12 +251,13 @@ def _grpo_metrics_colocated(
         span.duration for span in actor_train_step_spans
     ]
 
-    trace_writer.log_trace(
-        global_step_groups,
-        rollout_spans,
-        refer_inference_spans,
-        actor_train_groups,
-    )
+    if trace_writer is not None:
+      trace_writer.log_trace(
+          global_step_groups,
+          rollout_spans,
+          refer_inference_spans,
+          actor_train_groups,
+      )
 
     # pyformat: disable
     return {
@@ -276,7 +277,7 @@ def _grpo_metrics_colocated(
   @staticmethod
   def _grpo_metrics_rollout_1_actor_2_reference_2(
       extract_spans_fn: _GrpoExtractSpansFn,
-      trace_writer: PerfettoTraceWriter,
+      trace_writer: PerfettoTraceWriter | None,
       query: PerfSpanQuery,
   ) -> MetricsT:
     """GRPO workflow: actor and reference are on the same mesh,rollout is on a different mesh.
@@ -343,12 +344,13 @@ def _grpo_metrics_rollout_1_actor_2_reference_2(
         for a, b in zip(actor_train_groups[:-1], refer_inference_spans[1:])
     ] + [0.0]
 
-    trace_writer.log_trace(
-        global_step_groups,
-        rollout_spans,
-        refer_inference_spans,
-        actor_train_groups,
-    )
+    if trace_writer is not None:
+      trace_writer.log_trace(
+          global_step_groups,
+          rollout_spans,
+          refer_inference_spans,
+          actor_train_groups,
+      )
 
     # pyformat: disable
     return {
@@ -372,7 +374,7 @@ def _grpo_metrics_rollout_1_actor_2_reference_2(
   @staticmethod
   def _grpo_metrics_fully_disaggregated(
       extract_spans_fn: _GrpoExtractSpansFn,
-      trace_writer: PerfettoTraceWriter,
+      trace_writer: PerfettoTraceWriter | None,
       query: PerfSpanQuery,
   ) -> MetricsT:
     """GRPO workflow: rollout, actor and reference are all on different meshes.
@@ -442,12 +444,13 @@ def _grpo_metrics_fully_disaggregated(
         for a, b in zip(actor_train_groups[:-1], actor_train_groups[1:])
     ] + [0.0]
 
-    trace_writer.log_trace(
-        global_step_groups,
-        rollout_spans,
-        refer_inference_spans,
-        actor_train_groups,
-    )
+    if trace_writer is not None:
+      trace_writer.log_trace(
+          global_step_groups,
+          rollout_spans,
+          refer_inference_spans,
+          actor_train_groups,
+      )
 
     # pyformat: disable
     return {

tunix/perf/metrics.py

@@ -73,10 +73,19 @@ class MetricsBuffer:
 
 @dataclasses.dataclass(frozen=True)
 class PerfMetricsOptions:
-  # Directory to write the raw metrics/events to.
+  """Options for configuring performance metrics.
+
+  Attributes:
+    enable_trace_writer: Whether to enable the trace writer. By default, it is
+      enabled when perf metrics are enabled. If False, the trace will not be
+      written out.
+    log_dir: Directory to write the raw metrics/events to.
+    custom_export_fn_path: Path to the custom export function. If set, the
+      custom export function will be loaded from the path instead of being
+      created by PerfMetricsExport.
+  """
+  enable_trace_writer: bool = True
   log_dir: str = ""
-  # Path to the custom export function. If set, the custom export function will
-  # be loaded from the path instead of being created by PerfMetricsExport.
   custom_export_fn_path: str = ""