(tvm) branch main updated: [Disco] Propagate structlog configuration to disco workers (#16618)

[ma...@apache.org]

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

masahi 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 1278c3544d [Disco] Propagate structlog configuration to disco workers (#16618)
1278c3544d is described below

commit 1278c3544d15d6b743073ba1058e5efbfc2894f4
Author: Eric Lunderberg <Lu...@users.noreply.github.com>
AuthorDate: Tue Mar 12 05:09:20 2024 -0500

    [Disco] Propagate structlog configuration to disco workers (#16618)
    
    Prior to this commit, while `structlog.configure(...)` would only
    impact log statements generated in the main process.  Any workers
    started with `disco.session.ProcessSession` do not inherit the
    `structlog` configuration.  While `disco.session.ThreadedSession` would
    inherit the `structlog` configuration, it would also inherit
    process-specific CUDA variables.
    
    This commit updates `disco.session.ProcessSession` to explicitly
    propagate any `structlog` configuration to child processes.  This
    implementation intentionally avoids introducing a new dependency for
    TVM.  If the `structlog` package is not available, the config
    propagation is skipped.
---
 python/tvm/runtime/disco/session.py | 41 +++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/python/tvm/runtime/disco/session.py b/python/tvm/runtime/disco/session.py
index 1013d14a89..53b362f579 100644
--- a/python/tvm/runtime/disco/session.py
+++ b/python/tvm/runtime/disco/session.py
@@ -17,6 +17,11 @@
 """This module defines a Session in Disco. Session is the primary interface that users interact
 with the distributed runtime.
 """
+
+import os
+import pickle
+
+
 from typing import Any, Callable, Optional, Sequence, Union
 
 import numpy as np
@@ -384,6 +389,42 @@ class ProcessSession(Session):
             "runtime.disco.create_process_pool",
             entrypoint,
         )
+        self._configure_structlog()
+
+    def _configure_structlog(self) -> None:
+        try:
+            import structlog  # pylint: disable=import-outside-toplevel
+        except ImportError:
+            return
+
+        config = pickle.dumps(structlog.get_config())
+        func = self.get_global_func("runtime.disco._configure_structlog")
+        func(config, os.getpid())
+
+
+@register_func("runtime.disco._configure_structlog")
+def _configure_structlog(pickled_config: bytes, parent_pid: int) -> None:
+    """Configure structlog for all disco workers
+
+    The child processes
+
+    Parameters
+    ----------
+    pickled_config: bytes
+
+        The pickled configuration for structlog
+
+    parent_pid: int
+
+        The PID of the main process.  This is used to restrict the
+    """
+    if os.getpid() == parent_pid:
+        return
+
+    import structlog  # pylint: disable=import-outside-toplevel
+
+    config = pickle.loads(pickled_config)
+    structlog.configure(**config)
 
 
 @register_func("runtime.disco._import_python_module")