Make nested multiprocess work (#264)

* Make nested multiprocess work

Author: gaogaotiantian

src/viztracer/main.py

@@ -10,7 +10,6 @@
 import multiprocessing.util  # type: ignore
 import os
 import platform
-import re
 import shutil
 import signal
 import sys
@@ -23,7 +22,7 @@
 from . import __version__
 from .attach_process.add_code_to_python_process import run_python_code  # type: ignore
 from .code_monkey import CodeMonkey
-from .patch import patch_multiprocessing, patch_subprocess
+from .patch import install_all_hooks
 from .report_builder import ReportBuilder
 from .util import time_str_to_us, color_print, same_line_print, pid_exists
 from .viztracer import VizTracer
@@ -153,13 +152,6 @@ def create_parser(self) -> argparse.ArgumentParser:
                             help="time you want to trace the process")
         return parser
 
-    @property
-    def is_main_process(self) -> bool:
-        options = self.options
-        if options.subprocess_child:
-            return False  # pragma: no cover
-        return self.parent_pid == os.getpid()
-
     def load_config_file(self, filename: str = ".viztracerrc") -> argparse.Namespace:
         ret = argparse.Namespace()
         if os.path.exists(filename):
@@ -262,6 +254,7 @@ def parse(self, argv: List[str]) -> VizProcedureResult:
             "log_gc": options.log_gc,
             "log_sparse": options.log_sparse,
             "log_async": options.log_async,
+            "log_audit": options.log_audit,
             "vdb": options.vdb,
             "pid_suffix": True,
             "file_info": False,
@@ -296,45 +289,6 @@ def search_file(self, file_name: str) -> Optional[str]:
 
         return None
 
-    def install_all_hooks(self, tracer: VizTracer) -> None:
-        options = self.options
-
-        # multiprocess hook
-        if not options.ignore_multiprocess:
-            patch_multiprocessing(tracer)
-            if not options.subprocess_child:
-                patch_subprocess(self.args + ["--subprocess_child", "--dump_raw", "-o", tracer.output_file])
-
-            # If we want to hook fork correctly with file waiter, we need to
-            # use os.register_at_fork to write the file, and make sure
-            # os.exec won't clear viztracer so that the file lives forever.
-            # This is basically equivalent to py3.8 + Linux
-            if hasattr(sys, "addaudithook"):
-                if hasattr(os, "register_at_fork"):
-                    def audit_hook(event, args):  # pragma: no cover
-                        if event == "os.exec":
-                            tracer.exit_routine()
-                    sys.addaudithook(audit_hook)  # type: ignore
-                    os.register_at_fork(after_in_child=lambda: tracer.label_file_to_write())  # type: ignore
-                if options.log_audit is not None:
-                    audit_regex_list = [re.compile(regex) for regex in options.log_audit]
-
-                    def audit_hook(event, args):  # pragma: no cover
-                        if len(audit_regex_list) == 0 or any((regex.fullmatch(event) for regex in audit_regex_list)):
-                            tracer.log_instant(event, args={"args": [str(arg) for arg in args]})
-                    sys.addaudithook(audit_hook)  # type: ignore
-
-        # SIGTERM hook
-        def term_handler(signalnum, frame):
-            sys.exit(0)
-
-        if self.is_main_process:
-            signal.signal(signal.SIGTERM, term_handler)
-            multiprocessing.util.Finalize(self, self.exit_routine, exitpriority=-1)
-        else:
-            signal.signal(signal.SIGTERM, term_handler)
-            multiprocessing.util.Finalize(tracer, tracer.exit_routine, exitpriority=-1)
-
     def run(self) -> VizProcedureResult:
         if self.options.version:
             return self.show_version()
@@ -365,7 +319,19 @@ def run_code(self, code: Any, global_dict: Dict[str, Any]) -> VizProcedureResult
         tracer = VizTracer(**self.init_kwargs)
         self.tracer = tracer
 
-        self.install_all_hooks(tracer)
+        install_all_hooks(tracer,
+                          self.args,
+                          patch_multiprocess=not options.ignore_multiprocess)
+
+        def term_handler(signalnum, frame):
+            sys.exit(0)
+
+        signal.signal(signal.SIGTERM, term_handler)
+
+        if options.subprocess_child:
+            multiprocessing.util.Finalize(tracer, tracer.exit_routine, exitpriority=-1)
+        else:
+            multiprocessing.util.Finalize(self, self.exit_routine, exitpriority=-1)
 
         if not options.log_sparse:
             tracer.start()

src/viztracer/patch.py

@@ -5,6 +5,7 @@
 import functools
 from multiprocessing import Process
 import os
+import re
 import sys
 import textwrap
 from typing import Any, Callable, Dict, List, Sequence, Union
@@ -53,7 +54,7 @@ def subprocess_init(self, args: Union[str, Sequence], **kwargs) -> None:
     setattr(subprocess.Popen, "__init__", subprocess_init)
 
 
-def patch_multiprocessing(tracer: VizTracer) -> None:
+def patch_multiprocessing(tracer: VizTracer, args: List[str]) -> None:
 
     # For fork process
     def func_after_fork(tracer: VizTracer):
@@ -83,7 +84,7 @@ def get_command_line(**kwds) -> List[str]:
             prog = textwrap.dedent(f"""
                     from multiprocessing.spawn import spawn_main;
                     from viztracer.patch import patch_spawned_process;
-                    patch_spawned_process({tracer.init_kwargs});
+                    patch_spawned_process({tracer.init_kwargs}, {args});
                     spawn_main(%s)
                     """)
             prog %= ', '.join('%s=%r' % item for item in kwds.items())
@@ -100,24 +101,27 @@ def __init__(
             run: Callable,
             target: Callable,
             args: List[Any],
-            kwargs: Dict[str, Any]):
+            kwargs: Dict[str, Any],
+            cmdline_args: List[str]):
         self._viztracer_kwargs = viztracer_kwargs
         self._run = run
         self._target = target
         self._args = args
         self._kwargs = kwargs
+        self._cmdline_args = cmdline_args
         self._exiting = False
 
     def run(self) -> None:
         import viztracer
 
         tracer = viztracer.VizTracer(**self._viztracer_kwargs)
+        install_all_hooks(tracer, self._cmdline_args)
         tracer.register_exit()
         tracer.start()
         self._run()
 
 
-def patch_spawned_process(viztracer_kwargs: Dict[str, Any]):
+def patch_spawned_process(viztracer_kwargs: Dict[str, Any], cmdline_args: List[str]):
     from multiprocessing import reduction, process  # type: ignore
     from multiprocessing.spawn import prepare
     import multiprocessing.spawn
@@ -130,7 +134,7 @@ def _main_3839(fd, parent_sentinel):
                 preparation_data = reduction.pickle.load(from_parent)
                 prepare(preparation_data)
                 self: Process = reduction.pickle.load(from_parent)
-                sp = SpawnProcess(viztracer_kwargs, self.run, self._target, self._args, self._kwargs)
+                sp = SpawnProcess(viztracer_kwargs, self.run, self._target, self._args, self._kwargs, cmdline_args)
                 self.run = sp.run
             finally:
                 del process.current_process()._inheriting
@@ -144,7 +148,7 @@ def _main_3637(fd):
                 preparation_data = reduction.pickle.load(from_parent)
                 prepare(preparation_data)
                 self: Process = reduction.pickle.load(from_parent)
-                sp = SpawnProcess(viztracer_kwargs, self.run, self._target, self._args, self._kwargs)
+                sp = SpawnProcess(viztracer_kwargs, self.run, self._target, self._args, self._kwargs, cmdline_args)
                 self.run = sp.run
             finally:
                 del process.current_process()._inheriting
@@ -154,3 +158,33 @@ def _main_3637(fd):
         multiprocessing.spawn._main = _main_3839  # type: ignore
     else:
         multiprocessing.spawn._main = _main_3637  # type: ignore
+
+
+def install_all_hooks(
+        tracer: VizTracer,
+        args: List[str],
+        patch_multiprocess: bool = True) -> None:
+
+    # multiprocess hook
+    if patch_multiprocess:
+        patch_multiprocessing(tracer, args)
+        patch_subprocess(args + ["--subprocess_child", "--dump_raw", "-o", tracer.output_file])
+
+    # If we want to hook fork correctly with file waiter, we need to
+    # use os.register_at_fork to write the file, and make sure
+    # os.exec won't clear viztracer so that the file lives forever.
+    # This is basically equivalent to py3.8 + Linux
+    if hasattr(sys, "addaudithook"):
+        if hasattr(os, "register_at_fork") and patch_multiprocess:
+            def audit_hook(event, _):  # pragma: no cover
+                if event == "os.exec":
+                    tracer.exit_routine()
+            sys.addaudithook(audit_hook)  # type: ignore
+            os.register_at_fork(after_in_child=lambda: tracer.label_file_to_write())  # type: ignore
+        if tracer.log_audit is not None:
+            audit_regex_list = [re.compile(regex) for regex in tracer.log_audit]
+
+            def audit_hook(event, _):  # pragma: no cover
+                if len(audit_regex_list) == 0 or any((regex.fullmatch(event) for regex in audit_regex_list)):
+                    tracer.log_instant(event, args={"args": [str(arg) for arg in args]})
+            sys.addaudithook(audit_hook)  # type: ignore

src/viztracer/viztracer.py

@@ -32,6 +32,7 @@ def __init__(self,
                  log_gc: bool = False,
                  log_sparse: bool = False,
                  log_async: bool = False,
+                 log_audit: Optional[Sequence[str]] = None,
                  vdb: bool = False,
                  pid_suffix: bool = False,
                  file_info: bool = True,
@@ -66,6 +67,7 @@ def __init__(self,
         self.output_file = output_file
         self.system_print = None
         self.log_sparse = log_sparse
+        self.log_audit = log_audit
         self.dump_raw = dump_raw
         self.sanitize_function_name = sanitize_function_name
         self.minimize_memory = minimize_memory
@@ -112,6 +114,7 @@ def init_kwargs(self) -> Dict:
             "log_gc": self.log_gc,
             "log_sparse": self.log_sparse,
             "log_async": self.log_async,
+            "log_audit": self.log_audit,
             "vdb": self.vdb,
             "pid_suffix": self.pid_suffix,
             "min_duration": self.min_duration,

tests/test_multiprocess.py

@@ -11,6 +11,12 @@
 from .cmdline_tmpl import CmdlineTmpl
 
 
+file_grandparent = """
+import subprocess
+subprocess.run(["python", "parent.py"])
+"""
+
+
 file_parent = """
 import subprocess
 subprocess.run(["python", "child.py"])
@@ -82,6 +88,33 @@ def f():
     time.sleep(0.1)
 """
 
+file_nested_multiprocessing = """
+import multiprocessing
+from multiprocessing import Process
+import time
+
+
+def fib(n):
+    if n < 2:
+        return 1
+    return fib(n-1) + fib(n-2)
+
+def f():
+    fib(5)
+
+def spawn():
+    p = Process(target=f)
+    p.start()
+    p.join()
+
+if __name__ == "__main__":
+    fib(2)
+    p = Process(target=spawn)
+    p.start()
+    p.join()
+    time.sleep(0.1)
+"""
+
 file_multiprocessing_overload_run = """
 import multiprocessing
 from multiprocessing import Process
@@ -196,6 +229,30 @@ def test_child_process(self):
                           expected_output_file=None)
             self.assertEqual(len(os.listdir(tmpdir)), 1)
 
+    def test_nested(self):
+        def check_func(data):
+            pids = set()
+            for entry in data["traceEvents"]:
+                pids.add(entry["pid"])
+            self.assertEqual(len(pids), 5)
+        with open("parent.py", "w") as f:
+            f.write(file_parent)
+        self.template(["viztracer", "-o", "result.json", "cmdline_test.py"],
+                      expected_output_file="result.json", script=file_grandparent, check_func=check_func)
+        os.remove("parent.py")
+
+    def test_nested_multiproessing(self):
+        def check_func(data):
+            pids = set()
+            for entry in data["traceEvents"]:
+                pids.add(entry["pid"])
+            self.assertEqual(len(pids), 3)
+        with open("parent.py", "w") as f:
+            f.write(file_multiprocessing)
+        self.template(["viztracer", "-o", "result.json", "cmdline_test.py"],
+                      expected_output_file="result.json", script=file_grandparent, check_func=check_func)
+        os.remove("parent.py")
+
     @unittest.skipIf(sys.platform == "win32", "Can't get anything on Windows with SIGTERM")
     def test_term(self):
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -248,6 +305,19 @@ def check_func(data):
                       check_func=check_func,
                       concurrency="multiprocessing")
 
+    def test_nested_multiprosessing(self):
+        def check_func(data):
+            pids = set()
+            for entry in data["traceEvents"]:
+                pids.add(entry["pid"])
+            self.assertEqual(len(pids), 3)
+
+        self.template(["viztracer", "-o", "result.json", "cmdline_test.py"],
+                      expected_output_file="result.json",
+                      script=file_nested_multiprocessing,
+                      check_func=check_func,
+                      concurrency="multiprocessing")
+
     def test_multiprocessing_entry_limit(self):
         result = self.template(["viztracer", "-o", "result.json", "--tracer_entries", "10", "cmdline_test.py"],
                                expected_output_file="result.json",
@@ -263,7 +333,7 @@ def check_func(data):
                 pids.add(entry["pid"])
             self.assertEqual(len(pids), 1)
 
-        self.template(["viztracer", "-o", "result.json", "--ignore_multiproces", "cmdline_test.py"],
+        self.template(["viztracer", "-o", "result.json", "--ignore_multiprocess", "cmdline_test.py"],
                       expected_output_file="result.json",
                       script=file_multiprocessing,
                       check_func=check_func,

tests/test_patch.py

@@ -32,7 +32,7 @@
 set_spawning_popen(None)
 child_r, parent_w = os.pipe()
 
-patch_spawned_process({'output_file': "$tmpdir/result.json", 'pid_suffix': True})
+patch_spawned_process({'output_file': "$tmpdir/result.json", 'pid_suffix': True}, [])
 pid = os.getpid()
 
 assert multiprocessing.spawn._main.__qualname__ == "_main"