Forward cuml.accel logging to subprocesses

docs/source/cuml-accel/logging-and-profiling.rst

@@ -64,6 +64,14 @@ Since the magic command doesn't accept arguments, use the programmatic installat
    # Install with desired log level before other imports
    cuml.accel.install(log_level="debug")
 
+Environment Variable
+~~~~~~~~~~~~~~~~~~~~
+
+The log level may also be configured by setting the ``CUML_ACCEL_LOG_LEVEL``
+environment variable to ``warn``, ``info`` or ``debug`` (case insensitive).
+This will be used if no explicit level is passed through either
+``cuml.accel.install`` or the CLI.
+
 Example
 -------
 

python/cuml/cuml/accel/__main__.py

@@ -128,7 +128,7 @@ def main(argv: list[str] | None = None):
     ns = parse_args(sys.argv[1:] if argv is None else argv)
 
     # Parse verbose into log_level
-    log_level = {0: "warn", 1: "info", 2: "debug"}.get(min(ns.verbose, 2))
+    log_level = {0: None, 1: "info", 2: "debug"}.get(min(ns.verbose, 2))
 
     # Enable acceleration
     install(disable_uvm=ns.disable_uvm, log_level=log_level)

python/cuml/cuml/accel/core.py

@@ -39,7 +39,7 @@ def __repr__(self):
         return "<Logger level={self.level.name!r}>"
 
     def _log(self, msg):
-        print(f"[cuml.accel] {msg}")
+        print(f"[cuml.accel] {msg}", flush=True)
 
     def set_level(self, level: str) -> None:
         """Set the logger level.
@@ -136,7 +136,7 @@ def enabled() -> bool:
 
 def install(
     disable_uvm: bool = False,
-    log_level: Literal["error", "warn", "info", "debug"] = "warn",
+    log_level: Literal["error", "warn", "info", "debug", None] = None,
 ) -> None:
     """Enable `cuml.accel`.
 
@@ -145,18 +145,24 @@ def install(
     disable_uvm : bool, optional
         Whether to disable UVM.
     log_level : {"error", "warn", "info", "debug"}, optional
-        The log level to set for the `cuml.accel` logger. Defaults to `"warn"`,
-        set to `"info"` or `"debug"` to get more information about what methods
-        `cuml.accel` accelerated for a given run.
+        The log level to set for the `cuml.accel` logger. Defaults to `None` to
+        check the value of the `CUML_ACCEL_LOG_LEVEL` environment variable,
+        falling back to `"warn"` if not configured. Set to `"info"` or
+        `"debug"` to get more information about what methods `cuml.accel`
+        accelerated for a given run.
     """
     if enabled():
         # Already enabled, no-op
         return
 
+    if log_level is None:
+        log_level = os.environ.get("CUML_ACCEL_LOG_LEVEL", "warn")
+
     logger.set_level(log_level)
     # Set the environment variable if not already set so cuml.accel will
     # be automatically enabled in subprocesses
     os.environ.setdefault("CUML_ACCEL_ENABLED", "1")
+    os.environ.setdefault("CUML_ACCEL_LOG_LEVEL", log_level)
 
     if not disable_uvm:
         if _is_concurrent_managed_access_supported():

python/cuml/cuml_accel_tests/test_cli.py

@@ -307,17 +307,28 @@ def test_cli_mix_cuml_accel_and_cudf_pandas(first, second, tmpdir):
 
 
 @pytest.mark.parametrize(
-    "args, level", [([], "warn"), (["-v"], "info"), (["-vv"], "debug")]
+    "args, env, level",
+    [
+        ([], None, "warn"),
+        (["-v"], None, "info"),
+        (["-vv"], None, "debug"),
+        # CUML_ACCEL_LOG_LEVEL used if no -v flag given
+        ([], {"CUML_ACCEL_LOG_LEVEL": "debug"}, "debug"),
+        # CUML_ACCEL_LOG_LEVEL is case insensitive
+        ([], {"CUML_ACCEL_LOG_LEVEL": "InFo"}, "info"),
+        # -v flag takes precedence over CUML_ACCEL_LOG_LEVEL
+        (["-v"], {"CUML_ACCEL_LOG_LEVEL": "warn"}, "info"),
+    ],
 )
-def test_cli_verbose(args, level):
+def test_cli_verbose(args, env, level):
     script = dedent(
         f"""
         from cuml.accel.core import logger
         level = logger.level.name.lower()
-        assert level == {level!r}
+        assert level == {level!r}, f"Got %r" % level
         """
     )
-    run(["-m", "cuml.accel", *args], stdin=script)
+    run(["-m", "cuml.accel", *args], stdin=script, env=env)
 
 
 @pytest.mark.parametrize("mode", ["script", "module", "cmd", "stdin"])

python/cuml/cuml_accel_tests/test_core.py

@@ -59,6 +59,18 @@ def test_enabled_in_loky_executor():
         assert cuml.accel.is_proxy(remote)
 
 
+def get_level():
+    return cuml.accel.core.logger.level.name.lower()
+
+
+def test_log_level_forwarded_to_subprocesses(monkeypatch):
+    monkeypatch.setenv("CUML_ACCEL_LOG_LEVEL", "debug")
+    ctx = multiprocessing.get_context("spawn")
+    with ctx.Pool(processes=1) as pool:
+        log_level = pool.apply(get_level)
+    assert log_level == "debug"
+
+
 def iter_proxy_class_methods():
     """Generate test cases of (cls, method_name) for all ProxyBase proxied methods"""
     classes = proxy_base_subclasses()