Skip to content

Commit

Permalink
[Disco] Propagate structlog/logging config to workers
Browse files Browse the repository at this point in the history
This is a follow-up to apache#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.
  • Loading branch information
Lunderberg committed Mar 13, 2024
1 parent 831d769 commit 9d5786d
Showing 1 changed file with 26 additions and 3 deletions.
29 changes: 26 additions & 3 deletions python/tvm/runtime/disco/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
with the distributed runtime.
"""

import logging
import os
import pickle

Expand Down Expand Up @@ -397,7 +398,19 @@ def _configure_structlog(self) -> None:
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())

Expand All @@ -423,8 +436,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")
Expand Down

0 comments on commit 9d5786d

Please sign in to comment.