Issue 1:
Description

Hello team - First of all, not sure if this is a bug.

I'm having trouble setting up Python logging with Jax without running into log spam issues.

The following code works as expected:

import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)-5.5s] [%(name)-12.12s]: %(message)s')
logger = logging.getLogger(__name__)

logger.info('hello world')

Output

2022-11-24 21:26:11,974 [INFO ] [__main__ ]: hello world

But not when also importing jax

import logging
import jax

logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)-5.5s] [%(name)-12.12s]: %(message)s')
logger = logging.getLogger(__name__)

logger.info('hello world') # no visible output

It only works if one passes the force=True flag to the logging config

import logging
import jax

logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)-5.5s] [%(name)-12.12s]: %(message)s', force=True)
logger = logging.getLogger(__name__)

logger.info('hello world') # works as expected now

However, if one checks the loggers, this overwrote now the config of all previous loggers

print(logging.root.manager.loggerDict)
# {'urllib3.util.retry': <Logger urllib3.util.retry (INFO)>, 'urllib3.util': <logging.PlaceHolder object at 0x7fadd2974a90>, 'urllib3': <Logger urllib3 (INFO)>, 'urllib3.connection': <Logger urllib3.connection (INFO)>, 'urllib3.response': <Logger urllib3.response (INFO)>, 'urllib3.connectionpool': <Logger urllib3.connectionpool (INFO)>, 'urllib3.poolmanager': <L
ogger urllib3.poolmanager (INFO)>, 'charset_normalizer': <Logger charset_normalizer (INFO)>, 'requests': <Logger requests (INFO)>, 'jaxlib.tpu_client': <Logger jaxlib.tpu_client (INFO)>, 'jaxlib': <logging.PlaceHolder object at 0x7fadc6f0cac0>, 'jax._src.config': <Logger jax._src.config (INFO)>, 'jax._src': <logging.PlaceHolder object at 0x7fadd24a8d90>, 'jax'
: <logging.PlaceHolder object at 0x7fadd24a8df0>, 'jax._src.util': <Logger jax._src.util (INFO)>, 'jax._src.clusters.cluster': <Logger jax._src.clusters.cluster (INFO)>, 'jax._src.clusters': <logging.PlaceHolder object at 0x7fadc6bc0850>, 'jax._src.distributed': <Logger jax._src.distributed (INFO)>, 'jax._src.lib.xla_bridge': <Logger jax._src.lib.xla_bridge (I
NFO)>, 'jax._src.lib': <logging.PlaceHolder object at 0x7fadc6bc0220>, 'jax._src.profiler': <Logger jax._src.profiler (INFO)>, 'jax._src.path': <Logger jax._src.path (INFO)>, 'concurrent.futures': <Logger concurrent.futures (INFO)>, 'concurrent': <logging.PlaceHolder object at 0x7fadc679d580>, 'asyncio': <Logger asyncio (INFO)>, 'jax._src.dispatch': <Logger ja
x._src.dispatch (INFO)>, 'jax.interpreters.pxla': <Logger jax.interpreters.pxla (INFO)>, 'jax.interpreters': <logging.PlaceHolder object at 0x7fadc6552f10>, 'rich': <Logger rich (INFO)>, '__main__': <Logger __main__ (INFO)>}

This means that certain loggers (e.g. 'jax.experimental.compilation_cache.compilation_cache': <Logger jax.experimental.compilation_cache.compilation_cache (INFO)>) will now throw a lot of log spam à la

2022-11-24 20:58:53,231 [INFO ] [root ]: Not writing persistent cache entry for 'jit_prim_fun' because it took < 1.00 seconds to compile (0.32s)
2022-11-24 20:58:53,247 [INFO ] [root ]: Not writing persistent cache entry for 'jit_prim_fun' because it took < 1.00 seconds to compile (0.01s)
2022-11-24 20:58:53,321 [INFO ] [root ]: Not writing persistent cache entry for 'jit_prim_fun' because it took < 1.00 seconds to compile (0.07s)
2022-11-24 20:58:53,360 [INFO ] [root ]: Not writing persistent cache entry for 'jit_true_divide' because it took < 1.00 seconds to compile (0.03s)
2022-11-24 20:58:53,840 [INFO ] [root ]: Not writing persistent cache entry for 'jit_prim_fun' because it took < 1.00 seconds to compile (0.47s)
2022-11-24 20:58:53,889 [INFO ] [root ]: Not writing persistent cache entry for 'jit_prim_fun' because it took < 1.00 seconds to compile (0.04s)
2022-11-24 20:58:54,514 [INFO ] [root ]: Not writing persistent cache entry for 'jit_prim_fun' because it took < 1.00 seconds to compile (0.62s)
2022-11-24 20:58:55,159 [INFO ] [root ]: Not writing persistent cache entry for 'jit_prim_fun' because it took < 1.00 seconds to compile (0.63s)
2022-11-24 20:58:55,950 [INFO ] [root ]: Not writing persistent cache entry for 'jit_prim_fun' because it took < 1.00 seconds to compile (0.76s)
2022-11-24 20:58:56,251 [INFO ] [root ]: Not writing persistent cache entry for 'jit_prim_fun' because it took < 1.00 seconds to compile (0.26s)

because their config has changed. Any solutions to this?

I could reduce log spam when jax was using abseil logging, e.g. with

logging.getLogger('absl').setLevel(logging.WARNING)

but this has no effect now (I assume jax switched to Python logging).

Thanks for your help on this!

What jax/jaxlib version are you using?

jax==0.3.25 jaxlib==0.3.25

Which accelerator(s) are you using?

TPU

Additional system info

Linux

NVIDIA GPU info

No response

Issue 2:
Description

Importing jaxlib.mlir._mlir_libs._mlir (and by extension any modules depending on these, including jaxlib and jax), seems to render python's logging module unable to work.

I have a script test.py:

import logging
import importlib
import sys

logger = logging.getLogger(__name__)

import jax
logging.basicConfig(level=logging.INFO)
# import jax

logger.info('Test')

Running python -m test results in no logging output. If the second import statement, after the logging.basicConfig, is used instead of the first, then the expected INFO:__main__:Test line appears. If another module e.g. numpy is used instead of jax, the log line appears. I also do see the log line when using jax==0.3.14, jaxlib==0.3.14 - it disappears on version 0.3.15-0.3.18.

By looking at sys.modules before and after the import, the module that replicates this behaviour with the smallest number of dependencies appears to be jaxlib.mlir._mlir_libs._mlir.

What jax/jaxlib version are you using?

jax >=0.3.14, jaxlib >=0.3.14

Which accelerator(s) are you using?

CPU

Additional system info

Linux

NVIDIA GPU info

No response