From d33d68e3263868c48e065f4371ccdad9259b68a0 Mon Sep 17 00:00:00 2001
From: Shivam Raikundalia <sraikund@meta.com>
Date: Sat, 24 Aug 2024 23:03:19 +0000
Subject: [PATCH] [Profiler] Add test to make sure FunctionEvents are processed
 lazily (#134359)

Summary: Create simple test that checks that FunctionEvent build tree happens lazily by checking that the metrics for it changes before and after call.

Test Plan: Make sure test passes in CI

Reviewed By: briancoutinho

Differential Revision: D61685429

Pull Request resolved: https://github.com/pytorch/pytorch/pull/134359
Approved by: https://github.com/briancoutinho
---
 test/profiler/test_profiler.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/test/profiler/test_profiler.py b/test/profiler/test_profiler.py
index 22baadba7b4e..cf92a08d836e 100644
--- a/test/profiler/test_profiler.py
+++ b/test/profiler/test_profiler.py
@@ -1974,6 +1974,21 @@ assert KinetoStepTracker.current_step() == initial_step + 2 * niters
                 torch.add(x, y)
         self.assertTrue(len(p2.events()) == 0)
 
+    @skipIfTorchDynamo("profiler gets ignored if dynamo activated")
+    def test_lazy_build_tree(self):
+        with profile() as p:
+            self.payload()
+
+        stats = p._stats()
+        # Test that the tree is not built
+        self.assertEqual(stats.function_events_build_tree_call_duration_us, 0)
+        self.assertEqual(stats.number_of_events, 0)
+
+        # Test that the tree is built on demand
+        p.events()
+        self.assertGreater(stats.function_events_build_tree_call_duration_us, 0)
+        self.assertGreater(stats.number_of_events, 0)
+
 
 class SimpleNet(nn.Module):
     def __init__(self) -> None: