Skip to content

Commit

Permalink
[Disco] Propagate structlog/logging config to workers (apache#16715)
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 authored and thaisacs committed Apr 3, 2024
1 parent 1e44b01 commit 243f49c
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
from typing import Any, Callable, Optional, Sequence, Union
Expand Down Expand Up @@ -402,7 +403,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 @@ -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")
Expand Down

0 comments on commit 243f49c

Please sign in to comment.