mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #18724 from apaszke:downgrade-logging

Browse Source 
PiperOrigin-RevId: 586304793
This commit is contained in:
jax authors 2023-11-29 05:06:18 -08:00
commit 528a90e97d
2 changed files with 10 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
jax/_src/compilation_cache.py
View File

@ -63,13 +63,13 @@ def _initialize_cache() -> None:
global _cache_initialized
with _cache_initialized_mutex:
if _cache_initialized:
logger.info("_initialize_cache: cache has already been initialized!")
logger.debug("_initialize_cache: cache has already been initialized!")
return
_cache_initialized = True
# Nothing to do if the cache is disabled.
if not _is_cache_enabled():
logger.warning("_initialize_cache: cache is disabled!")
logger.debug("_initialize_cache: cache is disabled!")
return
global _cache
@ -80,7 +80,7 @@ def _initialize_cache() -> None:
return
_cache = get_file_cache(path)
logger.warning("Initialized persistent compilation cache at %s", path)
logger.debug("Initialized persistent compilation cache at %s", path)
def _get_cache() -> Optional[CacheInterface]:
@ -100,7 +100,7 @@ def get_executable_and_time(
"""
cache = _get_cache()
if cache is None:
logger.info("get_executable_and_time: cache is disabled/not initialized")
logger.debug("get_executable_and_time: cache is disabled/not initialized")
return None, None
executable_and_time = cache.get(cache_key)
if not executable_and_time:
@ -129,9 +129,9 @@ def put_executable_and_time(
"""
cache = _get_cache()
if cache is None:
logger.info("put_executable_and_time: cache is disabled/not initialized")
logger.debug("put_executable_and_time: cache is disabled/not initialized")
return
logger.info(
logger.debug(
"Writing %s to persistent compilation cache with key %s.",
module_name,
cache_key,
@ -167,8 +167,8 @@ def reset_cache() -> None:
"""Get back to pristine, uninitialized state."""
global _cache
global _cache_initialized
logger.info("Resetting cache at %s.",
_cache._path if _cache is not None else "<empty>")
logger.debug("Resetting cache at %s.",
_cache._path if _cache is not None else "<empty>")
_cache = None
_cache_initialized = False

4
jax/_src/compiler.py
View File

@ -326,7 +326,7 @@ def compile_or_get_cached(
if retrieved_executable is not None:
assert retrieved_compile_time is not None
logger.info("Persistent compilation cache hit for '%s'", module_name)
logger.debug("Persistent compilation cache hit for '%s'", module_name)
if config.use_original_compilation_cache_key_generation.value:
# TODO(b/293308239) Remove metrics for the original cache after the new
@ -389,7 +389,7 @@ def _cache_write(cache_key: str,
return
if host_callbacks:
logger.info(
logger.debug(
"Not writing persistent cache entry for '%s' because it uses host "
"callbacks (e.g. from jax.debug.print or breakpoint)", module_name)
return