(tvm) branch main updated: [Disco] Propagate structlog/logging config to workers (#16715)

[lunderberg]

This is an automated email from the ASF dual-hosted git repository.

lunderberg pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 1891b4db49 [Disco] Propagate structlog/logging config to workers 
(#16715)
1891b4db49 is described below

commit 1891b4db4933388fd67e8fbececc263930a148d1
Author: Eric Lunderberg <lunderb...@users.noreply.github.com>
AuthorDate: Wed Mar 27 09:11:08 2024 -0500

    [Disco] Propagate structlog/logging config to workers (#16715)
    
    This is a follow-up to #16618, which propagates the `structlog`
    configuration to disco worker processes.  For configurations that
    use `structlog.stdlib` to integrate `structlog` with the stdlib
    `logging` module, this integration must also be forwarded.
---
 python/tvm/runtime/disco/session.py | 29 ++++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/python/tvm/runtime/disco/session.py 
b/python/tvm/runtime/disco/session.py
index 344212a2f6..b8f74bacb0 100644
--- a/python/tvm/runtime/disco/session.py
+++ b/python/tvm/runtime/disco/session.py
@@ -18,6 +18,7 @@
 with the distributed runtime.
 """
 
+import logging
 import os
 import pickle
 from typing import Any, Callable, Optional, Sequence, Union
@@ -402,7 +403,19 @@ class ProcessSession(Session):
         except ImportError:
             return
 
-        config = pickle.dumps(structlog.get_config())
+        root_logger = logging.getLogger()
+        if len(root_logger.handlers) == 1 and isinstance(
+            root_logger.handlers[0].formatter, 
structlog.stdlib.ProcessorFormatter
+        ):
+            stdlib_formatter = root_logger.handlers[0].formatter
+        else:
+            stdlib_formatter = None
+
+        stdlib_level = root_logger.level
+
+        full_config = (structlog.get_config(), stdlib_formatter, stdlib_level)
+
+        config = pickle.dumps(full_config)
         func = self.get_global_func("runtime.disco._configure_structlog")
         func(config, os.getpid())
 
@@ -428,8 +441,18 @@ def _configure_structlog(pickled_config: bytes, 
parent_pid: int) -> None:
 
     import structlog  # pylint: disable=import-outside-toplevel
 
-    config = pickle.loads(pickled_config)
-    structlog.configure(**config)
+    full_config = pickle.loads(pickled_config)
+    structlog_config, stdlib_formatter, stdlib_level = full_config
+
+    root_logger = logging.getLogger()
+
+    root_logger.setLevel(stdlib_level)
+    if stdlib_formatter is not None:
+        handler = logging.StreamHandler()
+        handler.setFormatter(stdlib_formatter)
+        root_logger.addHandler(handler)
+
+    structlog.configure(**structlog_config)
 
 
 @register_func("runtime.disco._import_python_module")