Only log elapsed time during tracing for top level jit

PiperOrigin-RevId: 890643915

Author: yashk2810

jax/_src/pjit.py

@@ -16,6 +16,7 @@
 
 from collections import defaultdict
 from collections.abc import Callable, Sequence, Iterable
+import contextlib
 from dataclasses import dataclass, replace
 from functools import partial
 import inspect
@@ -511,9 +512,12 @@ def _trace_for_jit(
 
   qdd_token = _qdd_cache_index(fun, in_type.vals)  # represents qdd state context
 
-  with dispatch.log_elapsed_time(
-      "Finished tracing + transforming {fun_name} for pjit in {elapsed_time:.9f} sec",
-      fun_name(fun), event=dispatch.JAXPR_TRACE_EVENT):
+  elapsed_time_ctx = (
+      dispatch.log_elapsed_time(
+          "Finished tracing {fun_name} for jit in {elapsed_time:.9f} sec",
+          fun_name(fun), event=dispatch.JAXPR_TRACE_EVENT)
+      if core.trace_state_clean() else contextlib.nullcontext())
+  with elapsed_time_ctx:
     if ji.use_resource_env:  # pjit
       with (_internal_use_concrete_mesh(ctx_mesh),
             mesh_lib.use_abstract_mesh(ctx_mesh.abstract_mesh)):

tests/api_test.py

@@ -4745,7 +4745,7 @@ def add(x):
         add(inp)
       tracing_add_count = 0
       for m in cm.output:
-        if 'Finished tracing + transforming add for pjit' in m:
+        if 'Finished tracing add for jit' in m:
           tracing_add_count += 1
       self.assertEqual(tracing_add_count, 2)
 

tests/logging_test.py

@@ -149,7 +149,7 @@ def test_debug_logging(self):
     with jax_debug_log_modules("jax"):
       with capture_jax_logs() as log_output:
         jax.jit(lambda x: x + 1)(1)
-      self.assertIn("Finished tracing + transforming", log_output.getvalue())
+      self.assertIn("Finished tracing", log_output.getvalue())
       self.assertIn("Compiling jit(<lambda>)", log_output.getvalue())
 
     # Turn off all debug logging.
@@ -162,7 +162,7 @@ def test_debug_logging(self):
     with jax_debug_log_modules("jax._src.dispatch"):
       with capture_jax_logs() as log_output:
         jax.jit(lambda x: x + 1)(1)
-      self.assertIn("Finished tracing + transforming", log_output.getvalue())
+      self.assertIn("Finished tracing", log_output.getvalue())
       self.assertNotIn("Compiling jit(<lambda>)", log_output.getvalue())
 
     # Turn everything off again.
@@ -233,7 +233,7 @@ def test_subprocess_toggling_logging_level(self):
     log_output_verbose = log_output[:m.start()]
     log_output_silent = log_output[m.end():]
 
-    self.assertIn("Finished tracing + transforming <lambda> for pjit",
+    self.assertIn("Finished tracing <lambda> for jit",
                   log_output_verbose)
     self.assertEqual(log_output_silent, "")
 
@@ -256,7 +256,7 @@ def test_subprocess_double_logging_absent(self):
     # only one tracing line should be printed, if there's more than one
     # then logs are printing duplicated
     self.assertLen([line for line in log_lines
-                    if "Finished tracing + transforming" in line], 1)
+                    if "Finished tracing" in line], 1)
 
   @jtu.skip_on_devices("tpu")
   @unittest.skipIf(platform.system() == "Windows",