tiny improvement to the cprofile wrapper (#120100)

shunting314 · pytorchmergebot · commit becfda005e52 · 2024-02-17T02:10:25.000Z
1. right now we double increment the profile counter. The PR avoid that so we don't end up with profile_0, profile_2, profile_4 ... 2. log the latency to run the passed in function with profiling on so we can easily skip those _compile call which returns quickly. Pull Request resolved: #120100 Approved by: https://github.com/eellison
diff --git a/torch/_dynamo/utils.py b/torch/_dynamo/utils.py
@@ -131,12 +131,17 @@ def cprofile_wrapper(func):
     @wraps(func)
     def profile_wrapper(*args, **kwargs):
         global timer_counter
-        profile_path = Path(func.__name__ + f"{next(timer_counter)}.profile")
+        profile_cnt = next(timer_counter)
+        profile_path = Path(func.__name__ + f"{profile_cnt}.profile")
         prof = cProfile.Profile()
         prof.enable()
+        start_ts = time.time()
         retval = prof.runcall(func, *args, **kwargs)
+        profile_latency = time.time() - start_ts
         prof.disable()
-        print(f"### Cprofile for {func.__name__} iter {next(timer_counter)} ###")
+        print(
+            f"### Cprofile for {func.__name__} iter {profile_cnt} took {profile_latency:.3f} seconds ###"
+        )
         ps = pstats.Stats(prof)
         prof.dump_stats(profile_path)
         svg_path = profile_path.with_suffix(".svg")
