instrument environment interactions for perf metrics

PiperOrigin-RevId: 888837273

Author: s-noghabi

tests/perf/experimental/timeline_utils_test.py

@@ -188,9 +188,9 @@ def add_packet_side_effect():
     actual_descriptors = [
         {
             "name": p.track_descriptor.name,
-            "parent_uuid": p.track_descriptor.parent_uuid if i > 0 else None,
+            "parent_uuid": getattr(p.track_descriptor, "parent_uuid", None),
         }
-        for i, p in enumerate(captured_packets[:3])
+        for p in captured_packets[:3]
     ]
     expected_descriptors = [
         {"name": "overlap_timeline", "parent_uuid": None},
@@ -265,6 +265,71 @@ def add_packet_side_effect():
     self.assertEqual(actual_events, expected_events)
     mock_builder.serialize.assert_called_once()
 
+  @mock.patch.object(trace_writer_lib, "TraceProtoBuilder", autospec=True)
+  def test_write_timelines_grouping(self, mock_builder_cls):
+    mock_builder = mock_builder_cls.return_value
+    mock_builder.serialize.return_value = b""
+    captured_packets = []
+
+    def add_packet_side_effect():
+      p = mock.create_autospec(TracePacket, instance=True)
+      p.track_descriptor = mock.create_autospec(TrackDescriptor, instance=True)
+      p.track_event = mock.create_autospec(TrackEvent, instance=True)
+      captured_packets.append(p)
+      return p
+
+    mock_builder.add_packet.side_effect = add_packet_side_effect
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+      writer = trace_writer_lib.PerfettoTraceWriter(
+          trace_dir=tmp_dir, role_to_devices={"actor": ["tpu0", "tpu1"]}
+      )
+
+      t_main = tracer.Timeline("host-1", 1000.0)
+      t_main.start_span("main_span", 1001.0)
+
+      t_rollout = tracer.Timeline("host-2", 1000.0)
+      t_rollout.start_span("rollout", 1002.0)
+
+      t_tpu = tracer.Timeline("tpu0", 1000.0)
+      t_tpu.start_span("compute", 1003.0)
+
+      writer.write_timelines({
+          "host-1": t_main,
+          "host-2": t_rollout,
+          "tpu0": t_tpu,
+      })
+
+    main_group = captured_packets[0].track_descriptor
+    rollout_group = captured_packets[1].track_descriptor
+    tpu_group = captured_packets[2].track_descriptor
+    host_1 = captured_packets[3].track_descriptor
+    host_2 = captured_packets[4].track_descriptor
+    tpu0 = captured_packets[5].track_descriptor
+
+    with self.subTest("host_main_threads_group"):
+      self.assertEqual(main_group.name, "Host - Main threads")
+      self.assertEqual(main_group.uuid, 100000)
+
+    with self.subTest("host_rollout_threads_group"):
+      self.assertEqual(rollout_group.name, "Host - Rollout threads")
+      self.assertEqual(rollout_group.uuid, 100001)
+
+    with self.subTest("actor_cluster"):
+      self.assertEqual(tpu_group.name, "Actor Cluster")
+
+    with self.subTest("host_1"):
+      self.assertEqual(host_1.name, "host-1")
+      self.assertEqual(host_1.parent_uuid, 100000)
+
+    with self.subTest("host_2"):
+      self.assertEqual(host_2.name, "host-2")
+      self.assertEqual(host_2.parent_uuid, 100001)
+
+    with self.subTest("tpu0"):
+      self.assertEqual(tpu0.name, "tpu0")
+      self.assertEqual(tpu0.parent_uuid, tpu_group.uuid)
+
   def test_perfetto_trace_writer_integration(self):
     with tempfile.TemporaryDirectory() as tmp_dir:
       writer = trace_writer_lib.PerfettoTraceWriter(trace_dir=tmp_dir)
@@ -293,15 +358,16 @@ def test_perfetto_trace_writer_integration(self):
         self.assertLen(files, 1)
 
       if files:
+        file_name = files[0]
         with self.subTest("file_name_prefix"):
-          self.assertStartsWith(files[0], "perfetto_trace_v2_")
+          self.assertStartsWith(file_name, "perfetto_trace_v2_")
 
         with self.subTest("file_name_suffix"):
-          self.assertEndsWith(files[0], ".pb")
+          self.assertEndsWith(file_name, ".pb")
 
         with self.subTest("file_content"):
           self.assertGreater(
-              os.path.getsize(os.path.join(tmp_dir, files[0])), 0
+              os.path.getsize(os.path.join(tmp_dir, file_name)), 0
           )
 
   def test_perfetto_trace_writer_invalid_dir(self):

tests/perf/experimental/trace_writer_test.py

@@ -21,7 +21,6 @@
 from unittest import mock
 
 from absl.testing import absltest
-from absl.testing import parameterized
 import numpy as np
 from tunix.perf.experimental import timeline
 from tunix.perf.experimental import tracer
@@ -79,45 +78,6 @@ def trigger_all(self) -> None:
       t.join()
 
 
-class TracerUtilsTest(parameterized.TestCase):
-
-  def test_generate_host_timeline_id(self):
-    tid = tracer.generate_host_timeline_id()
-    self.assertStartsWith(tid, "host-")
-    self.assertIn(str(threading.get_ident()), tid)
-
-  @parameterized.named_parameters(
-      ("string", "tpu0", "tpu0"),
-      ("device_object", MockDevice("gpu", 7), "gpu7"),
-  )
-  def test_generate_device_timeline_id(self, device_id, expected_id):
-    self.assertEqual(tracer.generate_device_timeline_id(device_id), expected_id)
-
-  def test_generate_device_timeline_id_error(self):
-    with self.assertRaisesRegex(ValueError, "Unsupported id type"):
-      tracer.generate_device_timeline_id(123)
-
-  @parameterized.named_parameters(
-      ("none", None, []),
-      ("mixed_list", ["dev1", MockDevice("tpu", 0)], ["dev1", "tpu0"]),
-      (
-          "numpy_array",
-          np.array([MockDevice("tpu", 0), MockDevice("tpu", 1)]),
-          ["tpu0", "tpu1"],
-      ),
-      (
-          "numpy_array_2d",
-          np.array([
-              [MockDevice("tpu", 0), MockDevice("tpu", 1)],
-              [MockDevice("tpu", 2), MockDevice("tpu", 3)],
-          ]),
-          ["tpu0", "tpu1", "tpu2", "tpu3"],
-      ),
-  )
-  def test_generate_device_timeline_ids(self, devices, expected_ids):
-    self.assertEqual(tracer.generate_device_timeline_ids(devices), expected_ids)
-
-
 class PerfTracerTest(absltest.TestCase):
 
   def setUp(self):

tests/perf/experimental/tracer_test.py

@@ -22,6 +22,7 @@
 ROLE = "role"
 GROUP_ID = "group_id"
 PAIR_INDEX = "pair_index"
+QUEUED_SPAN = "queued_span"
 
 # Common Span / Event names.
 
@@ -33,3 +34,6 @@
 OLD_ACTOR_INFERENCE = "old_actor_inference"
 ADVANTAGE_COMPUTATION = "advantage_computation"
 PEFT_TRAIN = "peft_train"
+IDLE = "idle"
+QUEUE = "queue"
+ENVIRONMENT = "environment"