Enable passing a dict of module names: log level to set_logs python a…

…pi (#98989)

Adds "module" kwarg to set_logs to allow a user to pass a dict of module qualified names to log level to the API.

Pull Request resolved: #98989
Approved by: https://github.com/ezyang

test/dynamo/test_logging.py

@@ -151,6 +151,14 @@ def test_open_registration(self, records):
         logger.info("hi")
         self.assertEqual(len(records), 1)
 
+    # check logging to a random log that is not a child log of a registered
+    # logger registers it and sets handlers properly
+    @make_logging_test(modules={"torch.utils": logging.INFO})
+    def test_open_registration_python_api(self, records):
+        logger = logging.getLogger("torch.utils")
+        logger.info("hi")
+        self.assertEqual(len(records), 1)
+
 
 # single record tests
 exclusions = {"bytecode", "output_code", "schedule", "aot_graphs", "recompiles"}

torch/_logging/_internal.py

@@ -132,11 +132,12 @@ def set_logs(
     recompiles=False,
     output_code=False,
     schedule=False,
+    modules=None,
 ):
     """
     Enable setting the log level of individual components through kwargs.
     Args are set using the following format:
-        set_logs(<log_name>=<log_level>,...<artifact_name>=<True or False>)
+        set_logs(<log_name>=<log_level>,...<artifact_name>=<True or False>, ...,modules={"module.qualified.name":<log_level>})
     """
     # ignore if env var is set
     if LOG_ENV_VAR in os.environ:
@@ -147,8 +148,11 @@ def set_logs(
 
     log_state.clear()
 
+    if modules is None:
+        modules = {}
+
     def _set_logs(**kwargs):
-        for alias, val in kwargs.items():
+        for alias, val in itertools.chain(kwargs.items(), modules.items()):
             if log_registry.is_artifact(alias):
                 if val:
                     log_state.enable_artifact(alias)

torch/testing/_internal/logging_utils.py

@@ -35,12 +35,22 @@ def kwargs_to_settings(**kwargs):
     INT_TO_VERBOSITY = {10: "+", 20: "", 40: "-"}
 
     settings = []
+
+    def append_setting(name, level):
+        if isinstance(name, str) and isinstance(level, int) and level in INT_TO_VERBOSITY:
+            settings.append(INT_TO_VERBOSITY[level] + name)
+            return
+        else:
+            raise ValueError("Invalid value for setting")
+
     for name, val in kwargs.items():
         if isinstance(val, bool):
             settings.append(name)
         elif isinstance(val, int):
-            if val in INT_TO_VERBOSITY:
-                settings.append(INT_TO_VERBOSITY[val] + name)
+            append_setting(name, val)
+        elif isinstance(val, dict) and name == "modules":
+            for module_qname, level in val.items():
+                append_setting(module_qname, level)
         else:
             raise ValueError("Invalid value for setting")