<h3 dir="auto">Description</h3>
<p dir="auto">Hello team - First of all, not sure if this is a bug.</p>
<p dir="auto">I'm having trouble setting up Python logging with Jax without running into log spam issues.</p>
<p dir="auto">The following code works as expected:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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')"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">logging</span>

<span class="pl-s1">logging</span>.<span class="pl-en">basicConfig</span>(<span class="pl-s1">level</span><span class="pl-c1">=</span><span class="pl-s1">logging</span>.<span class="pl-v">INFO</span>, <span class="pl-s1">format</span><span class="pl-c1">=</span><span class="pl-s">'%(asctime)s [%(levelname)-5.5s] [%(name)-12.12s]: %(message)s'</span>)
<span class="pl-s1">logger</span> <span class="pl-c1">=</span> <span class="pl-s1">logging</span>.<span class="pl-en">getLogger</span>(<span class="pl-s1">__name__</span>)

<span class="pl-s1">logger</span>.<span class="pl-en">info</span>(<span class="pl-s">'hello world'</span>)</pre></div>
<p dir="auto">Output</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="2022-11-24 21:26:11,974 [INFO ] [__main__    ]: hello world"><pre class="notranslate"><code class="notranslate">2022-11-24 21:26:11,974 [INFO ] [__main__    ]: hello world
</code></pre></div>
<p dir="auto">But not when also importing jax</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">logging</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>

<span class="pl-s1">logging</span>.<span class="pl-en">basicConfig</span>(<span class="pl-s1">level</span><span class="pl-c1">=</span><span class="pl-s1">logging</span>.<span class="pl-v">INFO</span>, <span class="pl-s1">format</span><span class="pl-c1">=</span><span class="pl-s">'%(asctime)s [%(levelname)-5.5s] [%(name)-12.12s]: %(message)s'</span>)
<span class="pl-s1">logger</span> <span class="pl-c1">=</span> <span class="pl-s1">logging</span>.<span class="pl-en">getLogger</span>(<span class="pl-s1">__name__</span>)

<span class="pl-s1">logger</span>.<span class="pl-en">info</span>(<span class="pl-s">'hello world'</span>)  <span class="pl-c"># no visible output</span></pre></div>
<p dir="auto">It only works if one passes the <code class="notranslate">force=True</code> flag to the logging config</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">logging</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>

<span class="pl-s1">logging</span>.<span class="pl-en">basicConfig</span>(<span class="pl-s1">level</span><span class="pl-c1">=</span><span class="pl-s1">logging</span>.<span class="pl-v">INFO</span>, <span class="pl-s1">format</span><span class="pl-c1">=</span><span class="pl-s">'%(asctime)s [%(levelname)-5.5s] [%(name)-12.12s]: %(message)s'</span>, <span class="pl-s1">force</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)
<span class="pl-s1">logger</span> <span class="pl-c1">=</span> <span class="pl-s1">logging</span>.<span class="pl-en">getLogger</span>(<span class="pl-s1">__name__</span>)

<span class="pl-s1">logger</span>.<span class="pl-en">info</span>(<span class="pl-s">'hello world'</span>)  <span class="pl-c"># works as expected now</span></pre></div>
<p dir="auto">However, if one checks the loggers, this overwrote now the config of all previous loggers</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="print(logging.root.manager.loggerDict)
# {'urllib3.util.retry': &lt;Logger urllib3.util.retry (INFO)&gt;, 'urllib3.util': &lt;logging.PlaceHolder object at 0x7fadd2974a90&gt;, 'urllib3': &lt;Logger urllib3 (INFO)&gt;, 'urllib3.connection': &lt;Logger urllib3.connection (INFO)&gt;, 'urllib3.response': &lt;Logger urllib3.response (INFO)&gt;, 'urllib3.connectionpool': &lt;Logger urllib3.connectionpool (INFO)&gt;, 'urllib3.poolmanager': &lt;L
ogger urllib3.poolmanager (INFO)&gt;, 'charset_normalizer': &lt;Logger charset_normalizer (INFO)&gt;, 'requests': &lt;Logger requests (INFO)&gt;, 'jaxlib.tpu_client': &lt;Logger jaxlib.tpu_client (INFO)&gt;, 'jaxlib': &lt;logging.PlaceHolder object at 0x7fadc6f0cac0&gt;, 'jax._src.config': &lt;Logger jax._src.config (INFO)&gt;, 'jax._src': &lt;logging.PlaceHolder object at 0x7fadd24a8d90&gt;, 'jax'
: &lt;logging.PlaceHolder object at 0x7fadd24a8df0&gt;, 'jax._src.util': &lt;Logger jax._src.util (INFO)&gt;, 'jax._src.clusters.cluster': &lt;Logger jax._src.clusters.cluster (INFO)&gt;, 'jax._src.clusters': &lt;logging.PlaceHolder object at 0x7fadc6bc0850&gt;, 'jax._src.distributed': &lt;Logger jax._src.distributed (INFO)&gt;, 'jax._src.lib.xla_bridge': &lt;Logger jax._src.lib.xla_bridge (I
NFO)&gt;, 'jax._src.lib': &lt;logging.PlaceHolder object at 0x7fadc6bc0220&gt;, 'jax._src.profiler': &lt;Logger jax._src.profiler (INFO)&gt;, 'jax._src.path': &lt;Logger jax._src.path (INFO)&gt;, 'concurrent.futures': &lt;Logger concurrent.futures (INFO)&gt;, 'concurrent': &lt;logging.PlaceHolder object at 0x7fadc679d580&gt;, 'asyncio': &lt;Logger asyncio (INFO)&gt;, 'jax._src.dispatch': &lt;Logger ja
x._src.dispatch (INFO)&gt;, 'jax.interpreters.pxla': &lt;Logger jax.interpreters.pxla (INFO)&gt;, 'jax.interpreters': &lt;logging.PlaceHolder object at 0x7fadc6552f10&gt;, 'rich': &lt;Logger rich (INFO)&gt;, '__main__': &lt;Logger __main__ (INFO)&gt;}"><pre class="notranslate"><code class="notranslate">print(logging.root.manager.loggerDict)
# {'urllib3.util.retry': &lt;Logger urllib3.util.retry (INFO)&gt;, 'urllib3.util': &lt;logging.PlaceHolder object at 0x7fadd2974a90&gt;, 'urllib3': &lt;Logger urllib3 (INFO)&gt;, 'urllib3.connection': &lt;Logger urllib3.connection (INFO)&gt;, 'urllib3.response': &lt;Logger urllib3.response (INFO)&gt;, 'urllib3.connectionpool': &lt;Logger urllib3.connectionpool (INFO)&gt;, 'urllib3.poolmanager': &lt;L
ogger urllib3.poolmanager (INFO)&gt;, 'charset_normalizer': &lt;Logger charset_normalizer (INFO)&gt;, 'requests': &lt;Logger requests (INFO)&gt;, 'jaxlib.tpu_client': &lt;Logger jaxlib.tpu_client (INFO)&gt;, 'jaxlib': &lt;logging.PlaceHolder object at 0x7fadc6f0cac0&gt;, 'jax._src.config': &lt;Logger jax._src.config (INFO)&gt;, 'jax._src': &lt;logging.PlaceHolder object at 0x7fadd24a8d90&gt;, 'jax'
: &lt;logging.PlaceHolder object at 0x7fadd24a8df0&gt;, 'jax._src.util': &lt;Logger jax._src.util (INFO)&gt;, 'jax._src.clusters.cluster': &lt;Logger jax._src.clusters.cluster (INFO)&gt;, 'jax._src.clusters': &lt;logging.PlaceHolder object at 0x7fadc6bc0850&gt;, 'jax._src.distributed': &lt;Logger jax._src.distributed (INFO)&gt;, 'jax._src.lib.xla_bridge': &lt;Logger jax._src.lib.xla_bridge (I
NFO)&gt;, 'jax._src.lib': &lt;logging.PlaceHolder object at 0x7fadc6bc0220&gt;, 'jax._src.profiler': &lt;Logger jax._src.profiler (INFO)&gt;, 'jax._src.path': &lt;Logger jax._src.path (INFO)&gt;, 'concurrent.futures': &lt;Logger concurrent.futures (INFO)&gt;, 'concurrent': &lt;logging.PlaceHolder object at 0x7fadc679d580&gt;, 'asyncio': &lt;Logger asyncio (INFO)&gt;, 'jax._src.dispatch': &lt;Logger ja
x._src.dispatch (INFO)&gt;, 'jax.interpreters.pxla': &lt;Logger jax.interpreters.pxla (INFO)&gt;, 'jax.interpreters': &lt;logging.PlaceHolder object at 0x7fadc6552f10&gt;, 'rich': &lt;Logger rich (INFO)&gt;, '__main__': &lt;Logger __main__ (INFO)&gt;}
</code></pre></div>
<p dir="auto">This means that certain loggers (e.g. <code class="notranslate">'jax.experimental.compilation_cache.compilation_cache': &lt;Logger jax.experimental.compilation_cache.compilation_cache (INFO)&gt;</code>) will now throw a lot of log spam à la</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="2022-11-24 20:58:53,231 [INFO ] [root        ]: Not writing persistent cache entry for 'jit_prim_fun' because it took &lt; 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 &lt; 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 &lt; 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 &lt; 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 &lt; 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 &lt; 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 &lt; 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 &lt; 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 &lt; 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 &lt; 1.00 seconds to compile (0.26s)"><pre class="notranslate"><code class="notranslate">2022-11-24 20:58:53,231 [INFO ] [root        ]: Not writing persistent cache entry for 'jit_prim_fun' because it took &lt; 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 &lt; 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 &lt; 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 &lt; 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 &lt; 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 &lt; 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 &lt; 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 &lt; 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 &lt; 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 &lt; 1.00 seconds to compile (0.26s)
</code></pre></div>
<p dir="auto">because their config has changed. Any solutions to this?</p>
<p dir="auto">I could reduce log spam when jax was using abseil logging, e.g. with</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="logging.getLogger('absl').setLevel(logging.WARNING)"><pre class="notranslate"><code class="notranslate">logging.getLogger('absl').setLevel(logging.WARNING)
</code></pre></div>
<p dir="auto">but this has no effect now (I assume jax switched to Python logging).</p>
<p dir="auto">Thanks for your help on this!</p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">jax==0.3.25 jaxlib==0.3.25</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">TPU</p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto">Linux</p>
<h3 dir="auto">NVIDIA GPU info</h3>
<p dir="auto"><em>No response</em></p>
<h3 dir="auto">Description</h3>
<p dir="auto">Importing <code class="notranslate">jaxlib.mlir._mlir_libs._mlir</code> (and by extension any modules depending on these, including <code class="notranslate">jaxlib</code> and <code class="notranslate">jax</code>), seems to render python's logging module unable to work.</p>
<p dir="auto">I have a script <code class="notranslate">test.py</code>:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import logging
import importlib
import sys

logger = logging.getLogger(__name__)

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

logger.info('Test')"><pre class="notranslate"><code class="notranslate">import logging
import importlib
import sys

logger = logging.getLogger(__name__)

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

logger.info('Test')
</code></pre></div>
<p dir="auto">Running <code class="notranslate">python -m test</code> results in no logging output. If the second import statement, after the <code class="notranslate">logging.basicConfig</code>, is used instead of the first, then the expected <code class="notranslate">INFO:__main__:Test</code> line appears. If another module e.g. <code class="notranslate">numpy</code> is used instead of <code class="notranslate">jax</code>, the log line appears. I also do see the log line when using <code class="notranslate">jax==0.3.14, jaxlib==0.3.14</code> - it disappears on version 0.3.15-0.3.18.</p>
<p dir="auto">By looking at <code class="notranslate">sys.modules</code> before and after the import, the module that replicates this behaviour with the smallest number of dependencies appears to be <code class="notranslate">jaxlib.mlir._mlir_libs._mlir.</code></p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">jax &gt;=0.3.14, jaxlib &gt;=0.3.14</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">CPU</p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto">Linux</p>
<h3 dir="auto">NVIDIA GPU info</h3>
<p dir="auto"><em>No response</em></p>