Importing jax removes logging handlers

[GeorgeGradinariu]

Hi,

I don't seem to be able to log in my project after upgrading jax to 0.3.16.

Here's a cut down example:
The following

import logging

LOGGER = logging.getLogger("run")
console = logging.StreamHandler()
file = logging.FileHandler(filename="log.log")
logging.basicConfig(level=logging.INFO, handlers=[console, file])

LOGGER.info("test")

yields, as expected:
INFO:root:test

however, if I import jax:

import logging
import jax 
LOGGER = logging.getLogger("run")
console = logging.StreamHandler()
file = logging.FileHandler(filename="log.log")
logging.basicConfig(level=logging.INFO, handlers=[console, file])

LOGGER.info("test")

Nothing is printed to the console, or the log file.

Does anyone have any insight as to why this is happening?

Thanks!

[jakevdp]

Hi - I think the issue has to do with JAX using absl.logging rather than Python's built-in logging system. I can reproduce the problem by replacing the jax import with a simple use of absl.logging, which seems to affect subsequent Python logging:

import logging
import absl.logging
absl.logging.info('message from absl')

LOGGER = logging.getLogger("run")
console = logging.StreamHandler()
file = logging.FileHandler(filename="log.log")
logging.basicConfig(level=logging.INFO, handlers=[console, file])

LOGGER.info("test")  # Not shown

I don't know enough about absl.logging to know why this is happening (though the absl.logging source does appear to modify default logging levels). For what it's worth, JAX is working on options to switch to standard Python logging; the issue is tracked in #6308 and there's a potential fix under review in #10968.