Repository: jax-ml/jax
Pull request #23468 (2024-09-06 01:23:08)

Adding `JAX_LOGGING_LEVEL` configuration option
For example, setting `JAX_LOGGING_LEVEL=DEBUG` for

```python
jax.jit(lambda x: x)(jnp.ones((10,)))
```

gives

```
DEBUG:2024-09-16 10:11:49,929:jax._src.path:45: etils.epath was not found. Using pathlib for file I/O.                                                            
DEBUG:2024-09-16 10:11:50,045:jax._src.dispatch:178: Finished tracing + transforming convert_element_type for pjit in 0.000222683 sec                             
DEBUG:2024-09-16 10:11:50,048:jax._src.xla_bridge:579: Discovered path based JAX plugin: jax_plugins.cuda                                                         
DEBUG:2024-09-16 10:11:50,048:jax._src.xla_bridge:579: Discovered path based JAX plugin: jax_plugins.rocm                                                         
DEBUG:2024-09-16 10:11:50,054:jax._src.xla_bridge:594: Loading plugin module jax_plugins.rocm                                                                     
DEBUG:2024-09-16 10:11:50,055:jax._src.xla_bridge:594: Loading plugin module jax_plugins.cuda                                                                     
DEBUG:2024-09-16 10:11:50,055:jax._src.xla_bridge:970: Initializing backend 'cpu'                                                                                 
DEBUG:2024-09-16 10:11:50,057:jax._src.xla_bridge:982: Backend 'cpu' initialized                                                                                  
DEBUG:2024-09-16 10:11:50,057:jax._src.xla_bridge:970: Initializing backend 'rocm'                                                                                
INFO:2024-09-16 10:11:50,057:jax._src.xla_bridge:895: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'    
DEBUG:2024-09-16 10:11:50,057:jax._src.xla_bridge:970: Initializing backend 'tpu'                                                                                 
INFO:2024-09-16 10:11:50,058:jax._src.xla_bridge:895: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object
 file: No such file or directory                                                                                                                                  
DEBUG:2024-09-16 10:11:50,059:jax._src.interpreters.pxla:1903: Compiling convert_element_type with global shapes and types [ShapedArray(float32[])]. Argument mapp
ing: (UnspecifiedValue,).                                                                                                                                         
DEBUG:2024-09-16 10:11:50,066:jax._src.dispatch:178: Finished jaxpr to MLIR module conversion jit(convert_element_type) in 0.007055998 sec                        
DEBUG:2024-09-16 10:11:50,066:jax._src.compiler:168: get_compile_options: num_replicas=1 num_partitions=1 device_assignment=[[CpuDevice(id=0)]]                   
DEBUG:2024-09-16 10:11:50,067:jax._src.compiler:227: get_compile_options XLA-AutoFDO profile: using XLA-AutoFDO profile version -1                                
DEBUG:2024-09-16 10:11:50,067:jax._src.cache_key:120: get_cache_key hash of serialized computation: 96f6c2c849611cb65ca14596725417384a55335a17aaf471c2f18031b69bb7
36                                                                                                                                                                
DEBUG:2024-09-16 10:11:50,067:jax._src.cache_key:126: get_cache_key hash after serializing computation: 96f6c2c849611cb65ca14596725417384a55335a17aaf471c2f18031b69bb736
DEBUG:2024-09-16 10:11:50,067:jax._src.cache_key:120: get_cache_key hash of serialized jax_lib version: 5f45bbcdc1f605c0aa3cce55583a208fba43af4faefecfdf157d982dac286c25
DEBUG:2024-09-16 10:11:50,067:jax._src.cache_key:126: get_cache_key hash after serializing jax_lib version: 321fe1b73386c60fa57bb696231a28dbda25a62a329efc2c26d72fee523ea5fd
...
```

in the formatting style of already existing JAX module debugging: `jax_debug_log_modules` option.


Patch:
diff --git a/docs/persistent_compilation_cache.md b/docs/persistent_compilation_cache.md
index 47a7587b6..c49e18394 100644
--- a/docs/persistent_compilation_cache.md
+++ b/docs/persistent_compilation_cache.md
@@ -70,11 +70,11 @@ cc.set_cache_dir("/tmp/jax_cache")
 * `jax_persistent_cache_min_entry_size_bytes`: The minimum size (in bytes)
    of an entry that will be cached in the persistent compilation cache:
 
-   *  `-1`: disable the size restriction and prevent overrides. 
+   *  `-1`: disable the size restriction and prevent overrides.
 
    *  Leave at default (`0`) to allow for overrides. The override will
       typically ensure that the minimum size is optimal for the file system
-      being used for the cache. 
+      being used for the cache.
 
    *  `> 0`: the actual minimum size desired; no overrides.
 
@@ -155,7 +155,14 @@ import os
 os.environ["JAX_DEBUG_LOG_MODULES"] = "jax._src.compiler,jax._src.lru_cache"
 ```
 
-on the top of the script.
+on the top of the script. Alternatively, you can change the global jax logging level with
+
+```python
+import os
+os.environ["JAX_LOGGING_LEVEL"] = "DEBUG"
+# or locally with
+jax.config.update("jax_logging_level", "DEBUG")
+```
 
 ### Examining cache misses
 
diff --git a/jax/_src/config.py b/jax/_src/config.py
index 0e113c894..215ef443c 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -1647,6 +1647,7 @@ default_matmul_precision = optional_enum_state(
           'the name of a :class:`~jax.lax.DotAlgorithmPreset`.\n\n'),
     include_in_jit_key=True)
 
+
 traceback_filtering = enum_state(
     name = 'jax_traceback_filtering',
     enum_values=["off", "tracebackhide", "remove_frames", "quiet_remove_frames",
@@ -1913,14 +1914,6 @@ array_garbage_collection_guard = optional_enum_state(
     ),
 )
 
-def _update_debug_log_modules(module_names_str: str | None):
-  logging_config.disable_all_debug_logging()
-  if not module_names_str:
-    return
-  module_names = module_names_str.split(',')
-  for module_name in module_names:
-    logging_config.enable_debug_logging(module_name)
-
 # Don't define a context manager since this isn't threadsafe.
 string_state(
     name='jax_debug_log_modules',
@@ -1928,7 +1921,20 @@ string_state(
     help=('Comma-separated list of module names (e.g. "jax" or '
           '"jax._src.xla_bridge,jax._src.dispatch") to enable debug logging '
           'for.'),
-    update_global_hook=_update_debug_log_modules)
+    update_global_hook=logging_config.update_debug_log_modules)
+
+# Don't define a context manager since this isn't threadsafe.
+optional_enum_state(
+    name='jax_logging_level',
+    enum_values=['NOTSET', 'DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
+    default=logging.getLevelName(logging.getLogger("jax").level),
+    help=('Set the corresponding logging level on all jax loggers. Only string'
+          ' values from ["NOTSET", "DEBUG", "INFO", "WARNING", "ERROR",'
+          ' "CRITICAL"] are accepted. If None, the logging level will not be'
+          ' set. Includes C++ logging.'),
+    update_global_hook=lambda logging_level: \
+      logging_config.update_logging_level_global(logging_level=logging_level)
+)
 
 pmap_no_rank_reduction = bool_state(
     name='jax_pmap_no_rank_reduction',
diff --git a/jax/_src/logging_config.py b/jax/_src/logging_config.py
index d2f9d9c8f..bdf588d20 100644
--- a/jax/_src/logging_config.py
+++ b/jax/_src/logging_config.py
@@ -13,19 +13,92 @@
 # limitations under the License.
 
 import logging
+import os
 import sys
 
-_debug_handler = logging.StreamHandler(sys.stderr)
-_debug_handler.setLevel(logging.DEBUG)
 # Example log message:
 # DEBUG:2023-06-07 00:14:40,280:jax._src.xla_bridge:590: Initializing backend 'cpu'
-_debug_handler.setFormatter(logging.Formatter(
-    "{levelname}:{asctime}:{name}:{lineno}: {message}", style='{'))
+logging_formatter = logging.Formatter(
+    "{levelname}:{asctime}:{name}:{lineno}: {message}", style='{')
 
-_debug_enabled_loggers = []
+_logging_level_set: dict[str, int] = {}
+_default_TF_CPP_MIN_LOG_LEVEL = os.environ.get("TF_CPP_MIN_LOG_LEVEL", "1")
+
+_jax_logger_handler = logging.StreamHandler(sys.stderr)
+_jax_logger_handler.setFormatter(logging_formatter)
+
+_nameToLevel = {
+    'CRITICAL': logging.CRITICAL,
+    'FATAL': logging.FATAL,
+    'ERROR': logging.ERROR,
+    'WARN': logging.WARNING,
+    'WARNING': logging.WARNING,
+    'INFO': logging.INFO,
+    'DEBUG': logging.DEBUG,
+    'NOTSET': logging.NOTSET,
+}
+
+_tf_cpp_map = {
+    'CRITICAL': 3,
+    'FATAL': 3,
+    'ERROR': 2,
+    'WARN': 1,
+    'WARNING': 1,
+    'INFO': 0,
+    'DEBUG': 0,
+}
+
+def _set_TF_CPP_MIN_LOG_LEVEL(logging_level: str | None = None):
+  if logging_level in (None, "NOTSET"):
+    # resetting to user-default TF_CPP_MIN_LOG_LEVEL
+    # this is typically "1", but if the user overrode it, it can be != "1"
+    os.environ["TF_CPP_MIN_LOG_LEVEL"] = _default_TF_CPP_MIN_LOG_LEVEL
+  else:
+    # set cpp runtime logging level if the level is anything but NOTSET
+    if logging_level not in _tf_cpp_map:
+      raise ValueError(f"Attempting to set log level \"{logging_level}\" which"
+                       f" isn't one of the supported:"
+                       f" {list(_tf_cpp_map.keys())}.")
+    # config the CPP logging level 0 - debug, 1 - info, 2 - warning, 3 - error
+    os.environ["TF_CPP_MIN_LOG_LEVEL"] = str(_tf_cpp_map[logging_level])
+
+def update_logging_level_global(logging_level: str | None) -> None:
+  # remove previous handlers
+  for logger_name, level in _logging_level_set.items():
+    logger = logging.getLogger(logger_name)
+    logger.removeHandler(_jax_logger_handler)
+    logger.setLevel(level)
+  _logging_level_set.clear()
+  _set_TF_CPP_MIN_LOG_LEVEL(logging_level)
+
+  if logging_level is None:
+    return
+
+  logging_level_num = _nameToLevel[logging_level]
 
+  # update jax and jaxlib root loggers for propagation
+  root_loggers = [logging.getLogger("jax"), logging.getLogger("jaxlib")]
+  for logger in root_loggers:
+    logger.setLevel(logging_level_num)
+    logger.addHandler(_jax_logger_handler)
+    _logging_level_set[logger.name] = logger.level
 
-def enable_debug_logging(logger_name):
+# per-module debug logging
+
+_jax_logger = logging.getLogger("jax")
+
+class _DebugHandlerFilter(logging.Filter):
+  def filter(self, _):
+    return _jax_logger.level > logging.DEBUG
+
+_debug_handler = logging.StreamHandler(sys.stderr)
+_debug_handler.setLevel(logging.DEBUG)
+_debug_handler.setFormatter(logging_formatter)
+_debug_handler.addFilter(_DebugHandlerFilter())
+
+_debug_enabled_loggers = []
+
+def _enable_debug_logging(logger_name):
   """Makes the specified logger log everything to stderr.
 
   Also adds more useful debug information to the log messages, e.g. the time.
@@ -34,21 +107,28 @@ def enable_debug_logging(logger_name):
     logger_name: the name of the logger, e.g. "jax._src.xla_bridge".
   """
   logger = logging.getLogger(logger_name)
+  _debug_enabled_loggers.append((logger, logger.level))
+
   logger.addHandler(_debug_handler)
   logger.setLevel(logging.DEBUG)
-  _debug_enabled_loggers.append(logger)
 
 
-def disable_all_debug_logging():
+def _disable_all_debug_logging():
   """Disables all debug logging enabled via `enable_debug_logging`.
 
   The default logging behavior will still be in effect, i.e. WARNING and above
   will be logged to stderr without extra message formatting.
   """
-  for logger in _debug_enabled_loggers:
+  for logger, prev_level in _debug_enabled_loggers:
+    logger: logging.Logger
     logger.removeHandler(_debug_handler)
-    # Assume that the default non-debug log level is always WARNING. In theory
-    # we could keep track of what it was set to before. This shouldn't make a
-    # difference if not other handlers are attached, but set it back in case
-    # something else gets attached (e.g. absl logger) and for consistency.
-    logger.setLevel(logging.WARNING)
+    logger.setLevel(prev_level)
+  _debug_enabled_loggers.clear()
+
+def update_debug_log_modules(module_names_str: str | None):
+  _disable_all_debug_logging()
+  if not module_names_str:
+    return
+  module_names = module_names_str.split(',')
+  for module_name in module_names:
+    _enable_debug_logging(module_name)


Test patch:
diff --git a/tests/logging_test.py b/tests/logging_test.py
index a1d6695a1..05d874b73 100644
--- a/tests/logging_test.py
+++ b/tests/logging_test.py
@@ -15,8 +15,9 @@
 import contextlib
 import io
 import logging
-import os
 import platform
+import re
+import shlex
 import subprocess
 import sys
 import tempfile
@@ -49,10 +50,23 @@ def jax_debug_log_modules(value):
   finally:
     jax.config.update("jax_debug_log_modules", original_value)
 
+@contextlib.contextmanager
+def jax_logging_level(value):
+  # jax_logging_level doesn't have a context manager, because it's
+  # not thread-safe. But since tests are always single-threaded, we
+  # can define one here.
+  original_value = jax.config.jax_logging_level
+  jax.config.update("jax_logging_level", value)
+  try:
+    yield
+  finally:
+    jax.config.update("jax_logging_level", original_value)
+
 
 @contextlib.contextmanager
 def capture_jax_logs():
   log_output = io.StringIO()
+
   handler = logging.StreamHandler(log_output)
   logger = logging.getLogger("jax")
 
@@ -91,21 +105,8 @@ class LoggingTest(jtu.JaxTestCase):
     """))
       python = sys.executable
       assert "python" in python
-      env_variables = {"TF_CPP_MIN_LOG_LEVEL": "1"}
-      if os.getenv("ASAN_OPTIONS"):
-        env_variables["ASAN_OPTIONS"] = os.getenv("ASAN_OPTIONS")
-      if os.getenv("PYTHONPATH"):
-        env_variables["PYTHONPATH"] = os.getenv("PYTHONPATH")
-      if os.getenv("LD_LIBRARY_PATH"):
-        env_variables["LD_LIBRARY_PATH"] = os.getenv("LD_LIBRARY_PATH")
-      if os.getenv("LD_PRELOAD"):
-        env_variables["LD_PRELOAD"] = os.getenv("LD_PRELOAD")
       # Make sure C++ logging is at default level for the test process.
-      proc = subprocess.run(
-          [python, f.name],
-          capture_output=True,
-          env=env_variables,
-      )
+      proc = subprocess.run([python, f.name], capture_output=True)
 
       lines = proc.stdout.split(b"\n")
       lines.extend(proc.stderr.split(b"\n"))
@@ -155,6 +156,153 @@ class LoggingTest(jtu.JaxTestCase):
         jax.jit(lambda x: x + 1)(1)
       self.assertEmpty(log_output.getvalue())
 
+  @unittest.skipIf(platform.system() == "Windows",
+                   "Subprocess test doesn't work on Windows")
+  def test_subprocess_stderr_info_logging(self):
+    if sys.executable is None:
+      raise self.skipTest("test requires access to python binary")
+
+    program = """
+    import jax  # this prints INFO logging from backend imports
+    jax.jit(lambda x: x)(1)  # this prints logs to DEBUG (from compilation)
+    """
+
+    # strip the leading whitespace from the program script
+    program = re.sub(r"^\s+", "", program, flags=re.MULTILINE)
+
+    # test INFO
+    cmd = shlex.split(f"env JAX_LOGGING_LEVEL=INFO {sys.executable} -c"
+                      f" '{program}'")
+    p = subprocess.run(cmd, capture_output=True, text=True)
+    log_output = p.stderr
+    info_lines = log_output.split("\n")
+    self.assertGreater(len(info_lines), 0)
+    self.assertIn("INFO", log_output)
+    self.assertNotIn("DEBUG", log_output)
+
+  @unittest.skipIf(platform.system() == "Windows",
+                   "Subprocess test doesn't work on Windows")
+  def test_subprocess_stderr_debug_logging(self):
+    if sys.executable is None:
+      raise self.skipTest("test requires access to python binary")
+
+    program = """
+    import jax  # this prints INFO logging from backend imports
+    jax.jit(lambda x: x)(1)  # this prints logs to DEBUG (from compilation)
+    """
+
+    # strip the leading whitespace from the program script
+    program = re.sub(r"^\s+", "", program, flags=re.MULTILINE)
+
+    # test DEBUG
+    cmd = shlex.split(f"env JAX_LOGGING_LEVEL=DEBUG {sys.executable} -c"
+                      f" '{program}'")
+    p = subprocess.run(cmd, capture_output=True, text=True)
+    log_output = p.stderr
+    self.assertIn("INFO", log_output)
+    self.assertIn("DEBUG", log_output)
+
+    # test JAX_DEBUG_MODULES
+    cmd = shlex.split(f"env JAX_DEBUG_LOG_MODULES=jax {sys.executable} -c"
+                      f" '{program}'")
+    p = subprocess.run(cmd, capture_output=True, text=True)
+    log_output = p.stderr
+    self.assertIn("DEBUG", log_output)
+
+  @unittest.skipIf(platform.system() == "Windows",
+                   "Subprocess test doesn't work on Windows")
+  def test_subprocess_toggling_logging_level(self):
+    if sys.executable is None:
+      raise self.skipTest("test requires access to python binary")
+
+    _separator = "---------------------------"
+    program = f"""
+    import sys
+    import jax  # this prints INFO logging from backend imports
+    jax.jit(lambda x: x)(1)  # this prints logs to DEBUG (from compilation)
+    jax.config.update("jax_logging_level", None)
+    sys.stderr.write("{_separator}")
+    jax.jit(lambda x: x)(1)  # should not log anything now
+    """
+
+    # strip the leading whitespace from the program script
+    program = re.sub(r"^\s+", "", program, flags=re.MULTILINE)
+
+    cmd = shlex.split(f"env JAX_LOGGING_LEVEL=DEBUG {sys.executable} -c"
+                      f" '{program}'")
+    p = subprocess.run(cmd, capture_output=True, text=True)
+    log_output = p.stderr
+    m = re.search(_separator, log_output)
+    self.assertTrue(m is not None)
+    log_output_verbose = log_output[:m.start()]
+    log_output_silent = log_output[m.end():]
+
+    self.assertIn("Finished tracing + transforming <lambda> for pjit",
+                  log_output_verbose)
+    self.assertEqual(log_output_silent, "")
+
+  @unittest.skipIf(platform.system() == "Windows",
+                   "Subprocess test doesn't work on Windows")
+  def test_subprocess_double_logging_absent(self):
+    if sys.executable is None:
+      raise self.skipTest("test requires access to python binary")
+
+    program = """
+    import jax  # this prints INFO logging from backend imports
+    jax.config.update("jax_debug_log_modules", "jax._src.compiler,jax._src.dispatch")
+    jax.jit(lambda x: x)(1)  # this prints logs to DEBUG (from compilation)
+    """
+
+    # strip the leading whitespace from the program script
+    program = re.sub(r"^\s+", "", program, flags=re.MULTILINE)
+
+    cmd = shlex.split(f"env JAX_LOGGING_LEVEL=DEBUG {sys.executable} -c"
+                      f" '{program}'")
+    p = subprocess.run(cmd, capture_output=True, text=True)
+    log_output = p.stderr
+    self.assertNotEmpty(log_output)
+    log_lines = log_output.strip().split("\n")
+    # 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)
+
+  @unittest.skipIf(platform.system() == "Windows",
+                   "Subprocess test doesn't work on Windows")
+  def test_subprocess_cpp_logging_level(self):
+    if sys.executable is None:
+      raise self.skipTest("test requires access to python binary")
+
+    program = """
+    import sys
+    import jax  # this prints INFO logging from backend imports
+    jax.distributed.initialize("127.0.0.1:12345", num_processes=1, process_id=0)
+    """
+
+    # strip the leading whitespace from the program script
+    program = re.sub(r"^\s+", "", program, flags=re.MULTILINE)
+
+    # verbose logging: DEBUG, VERBOSE
+    cmd = shlex.split(f"env JAX_LOGGING_LEVEL=DEBUG {sys.executable} -c"
+                      f" '{program}'")
+    p = subprocess.run(cmd, capture_output=True, text=True)
+    self.assertIn("Initializing CoordinationService", p.stderr)
+
+    cmd = shlex.split(f"env JAX_LOGGING_LEVEL=INFO {sys.executable} -c"
+                      f" '{program}'")
+    p = subprocess.run(cmd, capture_output=True, text=True)
+    self.assertIn("Initializing CoordinationService", p.stderr)
+
+    # verbose logging: WARNING, None
+    cmd = shlex.split(f"env JAX_LOGGING_LEVEL=WARNING {sys.executable} -c"
+                      f" '{program}'")
+    p = subprocess.run(cmd, capture_output=True, text=True)
+    self.assertNotIn("Initializing CoordinationService", p.stderr)
+
+    cmd = shlex.split(f"{sys.executable} -c"
+                      f" '{program}'")
+    p = subprocess.run(cmd, capture_output=True, text=True)
+    self.assertNotIn("Initializing CoordinationService", p.stderr)
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())