Tweak setup_logging to use config.log_level (#82)

* Tweak `setup_logging` to use `config.log_level` Signed-off-by: Fabrice Normandin <[email protected]> * Only create ConfigLoader once (faster tests) Signed-off-by: Fabrice Normandin <[email protected]> --------- Signed-off-by: Fabrice Normandin <[email protected]>
mila-iqia · Nov 5, 2024 · e8514f6 · e8514f6
1 parent 727fa67
commit e8514f6
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 21 deletions.
diff --git a/docs/profiling_test.py b/docs/profiling_test.py
@@ -118,7 +118,7 @@ def test_notebook_commands_dont_cause_errors(experiment_dictconfig: DictConfig):
     config = resolve_dictconfig(experiment_dictconfig)
     # check for any errors when actually instantiating the components.
     # _experiment = _setup_experiment(config)
-    setup_logging(config)
+    setup_logging(log_level=config.log_level)
     lightning.seed_everything(config.seed, workers=True)
     _trainer = instantiate_trainer(config)
     datamodule = instantiate_datamodule(config.datamodule)

diff --git a/project/conftest.py b/project/conftest.py
@@ -289,7 +289,7 @@ def algorithm(
 def trainer(
     experiment_config: Config,
 ) -> pl.Trainer:
-    setup_logging(experiment_config)
+    setup_logging(log_level=experiment_config.log_level)
     lightning.seed_everything(experiment_config.seed, workers=True)
     return instantiate_trainer(experiment_config)
 

diff --git a/project/experiment.py b/project/experiment.py
@@ -14,7 +14,6 @@
 import copy
 import functools
 import logging
-import os
 from logging import getLogger as get_logger
 from typing import Any
 
@@ -44,10 +43,11 @@
 instantiate = hydra_zen.instantiate
 
 
-def setup_logging(experiment_config: Config) -> None:
-    LOGLEVEL = os.environ.get("LOGLEVEL", "info").upper()
+def setup_logging(log_level: str, global_log_level: str = "WARNING") -> None:
+    from project.main import PROJECT_NAME
+
     logging.basicConfig(
-        level=LOGLEVEL,
+        level=global_log_level.upper(),
         # format="%(asctime)s - %(levelname)s - %(message)s",
         format="%(message)s",
         datefmt="[%X]",
@@ -62,12 +62,8 @@ def setup_logging(experiment_config: Config) -> None:
         ],
     )
 
-    root_logger = logging.getLogger("project")
-
-    if experiment_config.debug:
-        root_logger.setLevel(logging.INFO)
-    elif experiment_config.verbose:
-        root_logger.setLevel(logging.DEBUG)
+    project_logger = logging.getLogger(PROJECT_NAME)
+    project_logger.setLevel(log_level.upper())
 
 
 def instantiate_trainer(experiment_config: Config) -> Trainer:

diff --git a/project/main.py b/project/main.py
@@ -45,14 +45,16 @@
 logger = get_logger(__name__)
 
 PROJECT_NAME = Path(__file__).parent.name
-
 add_configs_to_hydra_store()
+setup_logging(log_level="INFO", global_log_level="ERROR")
+
 
 auto_schema_plugin.config = auto_schema_plugin.AutoSchemaPluginConfig(
     schemas_dir=REPO_ROOTDIR / ".schemas",
     regen_schemas=False,
     stop_on_error=False,
     quiet=True,
+    verbose=False,
     add_headers=False,  # don't fallback to adding headers if we can't use vscode settings file.
 )
 
@@ -81,12 +83,13 @@ def main(dict_config: DictConfig) -> dict:
     # Resolve all the interpolations in the configs.
     config: Config = resolve_dictconfig(dict_config)
 
-    # Now we instantiate the components.
+    setup_logging(
+        log_level=config.log_level,
+        global_log_level="DEBUG" if config.debug else "INFO" if config.verbose else "WARNING",
+    )
 
     # seed the random number generators, so the weights that are
     # constructed are deterministic and reproducible.
-
-    setup_logging(config)
     lightning.seed_everything(seed=config.seed, workers=True)
 
     # Create the Trainer

diff --git a/project/utils/hydra_config_utils.py b/project/utils/hydra_config_utils.py
@@ -1,3 +1,4 @@
+import functools
 import importlib
 import inspect
 import typing
@@ -12,15 +13,31 @@
 logger = get_logger(__name__)
 
 
+@functools.cache
 def get_config_loader():
     from hydra._internal.config_loader_impl import ConfigLoaderImpl
     from hydra._internal.utils import create_automatic_config_search_path
 
     from project.main import PROJECT_NAME
 
+    # TODO: This (loading a config) is actually taking a long time, in part because this is
+    # triggering the hydra-auto-schema plugin to add schemas to all the yaml files.
+    AutoSchemaPlugin = None
+    backup = None
+    try:
+        from hydra_plugins.hydra_auto_schema.auto_schema_plugin import (  # type: ignore
+            AutoSchemaPlugin,
+        )
+
+        backup = AutoSchemaPlugin._ALREADY_DID
+        AutoSchemaPlugin._ALREADY_DID = True
+    except ImportError:
+        pass
     search_path = create_automatic_config_search_path(
         calling_file=None, calling_module=None, config_path=f"pkg://{PROJECT_NAME}.configs"
     )
+    if AutoSchemaPlugin is not None:
+        AutoSchemaPlugin._ALREADY_DID = backup
     config_loader = ConfigLoaderImpl(config_search_path=search_path)
     return config_loader
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -40,7 +40,7 @@ dependencies = [
     "mkdocstrings[python]>=0.26.2",
     "remote-slurm-executor",
     "watchdog>=5.0.3",
-    "hydra-auto-schema>=0.0.3",
+    "hydra-auto-schema>=0.0.4",
 ]
 readme = "README.md"
 requires-python = ">= 3.10"

diff --git a/uv.lock b/uv.lock