use current and commited step in perf tracer

PiperOrigin-RevId: 904696205

Author: s-noghabi

tests/perf/experimental/timeline_test.py

@@ -134,42 +134,6 @@ def test_stop_span_error_cases(self):
     with self.assertRaisesRegex(ValueError, "ended at .* before it began"):
       t.stop_span(1.0)
 
-  def test_snapshot(self):
-    t = timeline.Timeline("test_tl", 0.0)
-
-    snap0 = t.snapshot()
-    with self.subTest("initially_empty"):
-      self.assertEmpty(snap0.spans)
-      self.assertIsNot(snap0, t)
-
-    s1 = t.start_span("span1", 1.0)
-    snap1 = t.snapshot()
-    with self.subTest("active_span_not_in_snapshot"):
-      self.assertEmpty(snap1.spans)
-
-    t.stop_span(2.0)
-    snap2 = t.snapshot()
-    with self.subTest("completed_span_in_snapshot"):
-      self.assertLen(snap2.spans, 1)
-      self.assertEqual(snap2.spans[0].name, "span1")
-      self.assertEqual(snap2.spans[0].end, 2.0)
-
-    s2 = t.start_span("span2", 3.0)
-    s3 = t.start_span("span3", 4.0)
-    t.stop_span(5.0)  # stops s3
-    snap3 = t.snapshot()
-    with self.subTest("nested_active_span_not_in_snapshot"):
-      self.assertLen(snap3.spans, 2)
-      self.assertIn(s1.id, snap3.spans)
-      self.assertIn(s3.id, snap3.spans)
-      self.assertNotIn(s2.id, snap3.spans)
-
-    t.stop_span(6.0)  # stops s2
-    snap4 = t.snapshot()
-    with self.subTest("all_spans_completed"):
-      self.assertLen(snap4.spans, 3)
-      self.assertIn(s2.id, snap4.spans)
-
   def test_nested_timeline_with_tags_repr(self):
     born = 1000.0
     t = timeline.Timeline("test_tl", born)
@@ -195,8 +159,9 @@ def test_nested_timeline_with_tags_repr(self):
     # Check full repr string
     expected_repr = (
         f"Timeline(test_tl, {born:.6f})\n"
-        "[0] root: 1.000000, 4.000000, tags={'type': 'root_span'}\n"
-        "[1] child: 2.000000, 3.000000 (parent=0), tags={'iter': 1}\n"
+        "Current Step -0:\n"
+        "  [0] root: 1.000000, 4.000000, tags={'type': 'root_span'}\n"
+        "  [1] child: 2.000000, 3.000000 (parent=0), tags={'iter': 1}\n"
     )
     self.assertEqual(repr(t), expected_repr)
 
@@ -224,8 +189,8 @@ def test_span_success(self):
     t.span("async_op", 1.0, waitlist)
 
     self.mock_async_wait.assert_called_once()
-    self.assertLen(t.spans, 1)
-    s = t.spans[0]
+    self.assertLen(t.cur_step, 1)
+    s = t.cur_step[0]
     self.assertEqual(s.name, "async_op")
     self.assertEqual(s.begin, 1.0)
     self.assertTrue(s.ended)  # Ended because mock calls success immediately
@@ -234,8 +199,8 @@ def test_span_with_no_waitlist(self):
     t = timeline.AsyncTimeline("dev", 0.0)
     t.span("immediate", 1.0, [])
     self.mock_async_wait.assert_not_called()
-    self.assertLen(t.spans, 1)
-    self.assertTrue(t.spans[0].ended)
+    self.assertLen(t.cur_step, 1)
+    self.assertTrue(t.cur_step[0].ended)
 
   def test_delayed_completion(self):
     t = timeline.AsyncTimeline("dev", 0.0)
@@ -251,17 +216,40 @@ def capture_wait(waitlist, success, failure):
     self.mock_async_wait.side_effect = capture_wait
 
     t.span("delayed", 1.0, ["wait"])
-
-    self.assertEmpty(t.spans)  # Not yet recorded
+    self.assertEmpty(t.cur_step)  # Not yet recorded
 
     # Simulate completion
     with mock.patch.object(time, "perf_counter", return_value=5.0):
       callbacks["success"]()
 
-    self.assertLen(t.spans, 1)
-    s = t.spans[0]
+    self.assertLen(t.cur_step, 1)
+    s = t.cur_step[0]
     self.assertEqual(s.end, 5.0)
 
+  def test_nested_async_span_parent(self):
+    t = timeline.AsyncTimeline("dev", 0.0)
+
+    # Start sync spans to populate the stack with multiple items
+    s0 = t.start_span("root", 1.0)
+    s1 = t.start_span("child1", 2.0)
+    s2 = t.start_span("child2", 3.0)
+
+    # Now create an async span
+    t.span("async_op", 4.0, ["wait"])
+
+    # Find the async span.
+    async_s = None
+    for span in t.cur_step.values():
+      if span.name == "async_op":
+        async_s = span
+        break
+
+    self.assertIsNotNone(async_s)
+    self.assertEqual(async_s.parent_id, s2.id)
+
+    # Protect against index hardcoding bugs (e.g. self._spans_stack[1])
+    self.assertNotEqual(async_s.parent_id, s1.id)
+
   def test_failure(self):
     t = timeline.AsyncTimeline("dev", 0.0)
 

tests/perf/experimental/timeline_utils_test.py

@@ -81,11 +81,59 @@ def test_is_timeline_only_of_allowed_type(
     tl = timeline.Timeline("test", 0)
     for i, span_name in enumerate(spans_to_add):
       tl.start_span(span_name, float(i))
-
-    self.assertEqual(
-        timeline_utils.is_timeline_only_of_allowed_type(tl, allowed_types),
-        expected,
-    )
+      tl.stop_span(float(i) + 0.5)
+
+    with self.subTest("without_commit_no_cur_step"):
+      # without current step it shouldn't see anything without a commit
+      self.assertFalse(
+          timeline_utils.is_timeline_only_of_allowed_type(tl, allowed_types)
+      )
+
+    with self.subTest("without_commit_with_cur_step"):
+      # Using include_current_step should match expected
+      self.assertEqual(
+          timeline_utils.is_timeline_only_of_allowed_type(
+              tl, allowed_types, include_cur_step=True
+          ),
+          expected,
+      )
+
+    tl.commit_step()
+
+    with self.subTest("after_commit_no_cur_step"):
+      # After commit, default should match expected with or without
+      # include_cur_step
+      self.assertEqual(
+          timeline_utils.is_timeline_only_of_allowed_type(tl, allowed_types),
+          expected,
+      )
+
+    with self.subTest("after_commit_with_cur_step"):
+      self.assertEqual(
+          timeline_utils.is_timeline_only_of_allowed_type(
+              tl, allowed_types, include_cur_step=True
+          ),
+          expected,
+      )
+
+    if expected:
+      # Add an unallowed span to the current step
+      tl.start_span("unallowed_span_xyz", 10.0)
+      tl.stop_span(10.5)
+
+      with self.subTest("unallowed_span_in_cur_step_default"):
+        # Default (include_cur_step=False) ignores cur_step
+        self.assertTrue(
+            timeline_utils.is_timeline_only_of_allowed_type(tl, allowed_types)
+        )
+
+      with self.subTest("unallowed_span_in_cur_step_included"):
+        # With include_cur_step=True, the unallowed span causes a False return
+        self.assertFalse(
+            timeline_utils.is_timeline_only_of_allowed_type(
+                tl, allowed_types, include_cur_step=True
+            )
+        )
 
   @parameterized.named_parameters(
       dict(testcase_name="host_with_id", tl_id="host-12345", expected=True),

tests/perf/experimental/trace_writer_test.py

@@ -121,6 +121,7 @@ def add_packet_side_effect():
           },
       )
       t.stop_span(1002.0)
+      t.commit_step()
 
       writer.write_timelines({"timeline_test": t})
 
@@ -193,6 +194,8 @@ def add_packet_side_effect():
       span3.end = 1010.0
       t._cur_step[3] = span3
 
+      t.commit_step()
+
       writer.write_timelines({"overlap_timeline": t})
 
     self.assertLen(captured_packets, 9)
@@ -307,22 +310,29 @@ def add_packet_side_effect():
 
       t_main = tracer.Timeline("host-1", 1000.0)
       t_main.start_span("main_span", 1001.0)
+      t_main.stop_span(1002.0)
 
       t_rollout = tracer.Timeline("host-2", 1000.0)
       t_rollout.start_span("rollout", 1002.0)
+      t_rollout.stop_span(1003.0)
 
       t_tpu = tracer.Timeline("tpu0", 1000.0)
       t_tpu.start_span("compute", 1003.0)
+      t_tpu.stop_span(1004.0)
 
       t_tpu1 = tracer.Timeline("tpu1", 1000.0)
       t_tpu1.start_span("compute2", 1004.0)
+      t_tpu1.stop_span(1005.0)
 
-      writer.write_timelines({
+      timelines = {
           "host-1": t_main,
           "host-2": t_rollout,
           "tpu0": t_tpu,
           "tpu1": t_tpu1,
-      })
+      }
+      for tl in timelines.values():
+        tl.commit_step()
+      writer.write_timelines(timelines)
 
     main_group = captured_packets[0].track_descriptor
     rollout_group = captured_packets[1].track_descriptor
@@ -382,6 +392,9 @@ def test_perfetto_trace_writer_integration(self):
 
       timelines = {"timeline1": t1, "timeline_tags": t2}
 
+      for tl in timelines.values():
+        tl.commit_step()
+
       writer.write_timelines(timelines)
 
       # Check if file was created and has content
@@ -425,13 +438,23 @@ def test_perfetto_trace_writer_empty_timelines(self):
       # No content should be written.
       self.assertEmpty(files)
 
+  def test_perfetto_trace_writer_timeline_with_empty_committed_steps(self):
+    with tempfile.TemporaryDirectory() as tmp_dir:
+      writer = trace_writer_lib.PerfettoTraceWriter(trace_dir=tmp_dir)
+      t = tracer.Timeline("timeline_test", 1000.0)
+      writer.write_timelines({"timeline_test": t})
+      files = os.listdir(tmp_dir)
+      self.assertEmpty(files)
+
 
 class NoopTraceWriterTest(absltest.TestCase):
 
   def test_noop_trace_writer_write_timelines(self):
     writer = trace_writer_lib.NoopTraceWriter()
     t = tracer.Timeline("timeline", 1000.0)
     t.start_span("span1", 1001.0)
+    t.stop_span(1002.0)
+    t.commit_step()
     # Should not crash and do nothing.
     writer.write_timelines({"timeline": t})