Don't use root logger at import time

[hajarmitage]

Summary

At import time, these calls to logging.debug() implicitly call logging.basicConfig (https://docs.python.org/3/library/logging.html#logging.basicConfig), setting logging config for the whole project which cannot then be overwritten later. For instance, consider the following test script:

import logging
import jax

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
logger.info('info')

This should log out 'info', but because when import jax is called, this _mlir_lib/__init__.py file is run and a logging.debug is called, calling logging.basicConfig, my logging.basicConfig(level=logging.INFO) does nothing.

Fix: instead of using root logger, use a module level logger.

Found in this issue: https://github.com/google/jax/issues/12526

[stellaraccident]

Good catch - thanks.

[stellaraccident]

Also note https://reviews.llvm.org/D133450 in flight. I think yours should go in first

[hajarmitage]

I don't think I can commit this - can you commit this for me.

[stellaraccident]

Sorry for the delay - I attempted to rebase/land but the patch conflicts. I don't have time to actually update it right now - can you rebase/fix? I'll be back from traveling later in the week and can do the rebase/fix if you don't get to it but am not in a position to be making actual edits right now.

[hajarmitage]

Have updated the diff, and replaced a new call to logging.warning with logger.warning

[phawkins]

Gentle ping: can this be merged? The same bug got reported again.

[stellaraccident]

In D134812#3955348, @phawkins wrote:

Gentle ping: can this be merged? The same bug got reported again.

Landed. Very sorry for the delay getting this merged (do feel free to ask on discord if delays like this happen with individuals in the future: any committer can land).

mlir/python/mlir/_mlir_libs/__init__.py

 #
 # This facility allows downstreams to customize Context creation to their
 # needs.
 def _site_initialize():
   import importlib
   import itertools
   import logging
   from ._mlir import ir
+  logger = logging.getLogger(__name__)
   registry = ir.DialectRegistry()
   post_init_hooks = []
 
   def process_initializer_module(module_name):
     try:
       m = importlib.import_module(f".{module_name}", __name__)
     except ModuleNotFoundError:
       return False
     except ImportError:
       message = (f"Error importing mlir initializer {module_name}. This may "
       "happen in unclean incremental builds but is likely a real bug if "
       "encountered otherwise and the MLIR Python API may not function.")
-      logging.warning(message, exc_info=True)
+      logger.warning(message, exc_info=True)
 
-    logging.debug("Initializing MLIR with module: %s", module_name)
+    logger.debug("Initializing MLIR with module: %s", module_name)
     if hasattr(m, "register_dialects"):
-      logging.debug("Registering dialects from initializer %r", m)
+      logger.debug("Registering dialects from initializer %r", m)
       m.register_dialects(registry)
     if hasattr(m, "context_init_hook"):
-      logging.debug("Adding context init hook from %r", m)
+      logger.debug("Adding context init hook from %r", m)
       post_init_hooks.append(m.context_init_hook)
     return True
 
 
   # If _mlirRegisterEverything is built, then include it as an initializer
   # module.
   process_initializer_module("_mlirRegisterEverything")