Add time info through user warning for debugging b/496201097

Author: Charles Li

tests/integration/aot_identical_test.py

@@ -18,12 +18,14 @@
 training run (using train.py).
 """
 
+import datetime
 import tempfile
 import unittest
 import pytest
 import os
 import shutil
 import hashlib
+import warnings
 import re
 import jax
 from tests.utils.test_helpers import get_test_config_path
@@ -194,7 +196,11 @@ def assert_compile_and_real_match_jaxpr(self, test_name, *extra_args):
         get_test_config_path(),
         f"dump_jaxpr_local_dir={train_dump_dir}",
     ) + tuple(shared_args)
+    start_time = datetime.datetime.now()
     train.main(train_argv)
+    seconds_elapsed = (datetime.datetime.now() - start_time).total_seconds()
+    print(f"train cost {seconds_elapsed} secs")
+    warnings.warn(f"DEUBG: train cost {seconds_elapsed} secs")
     jax.clear_caches()
 
     # Run train_compile.py and dump jaxpr
@@ -206,7 +212,11 @@ def assert_compile_and_real_match_jaxpr(self, test_name, *extra_args):
         f"compile_topology={topology}",
         "compile_topology_num_slices=1",
     ) + tuple(shared_args)
+    start_time = datetime.datetime.now()
     train_compile.main(compile_argv)
+    seconds_elapsed = (datetime.datetime.now() - start_time).total_seconds()
+    print(f"train_compile cost {seconds_elapsed} secs")
+    warnings.warn(f"DEUBG: train_compile cost {seconds_elapsed} secs")
     jax.clear_caches()
 
     # Compare results
@@ -218,5 +228,6 @@ def assert_compile_and_real_match_jaxpr(self, test_name, *extra_args):
     )
 
   @pytest.mark.tpu_only
+  @pytest.mark.filterwarnings("always::UserWarning")
   def test_default_jaxpr_match(self):
     self.assert_compile_and_real_match_jaxpr("default_run")