From 8f655052a86417a3c51869310081447457eb4f37 Mon Sep 17 00:00:00 2001 From: Tian Gao Date: Thu, 10 Oct 2024 09:22:08 -0700 Subject: [PATCH] Fix logsparse on multiprocess spawn (#480) --- src/viztracer/patch.py | 3 ++- tests/test_logsparse.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/viztracer/patch.py b/src/viztracer/patch.py index 394b1b3d..3eb21a14 100644 --- a/src/viztracer/patch.py +++ b/src/viztracer/patch.py @@ -157,7 +157,8 @@ def run(self) -> None: tracer = viztracer.VizTracer(**self._viztracer_kwargs) install_all_hooks(tracer, self._cmdline_args) tracer.register_exit() - tracer.start() + if not self._viztracer_kwargs.get("log_sparse"): + tracer.start() self._run() diff --git a/tests/test_logsparse.py b/tests/test_logsparse.py index d85705f2..e7e1366c 100644 --- a/tests/test_logsparse.py +++ b/tests/test_logsparse.py @@ -76,6 +76,27 @@ def f(x): time.sleep(0.1) """ +file_multiprocess_spawn = """ +import multiprocessing +from viztracer import get_tracer + +def bar(): + pass + +def child(it): + for idx in it: + with get_tracer().log_event(f"custom_event"): + bar() + +def main(): + p = multiprocessing.get_context("spawn").Process(target=child, args=(range(3), )) + p.start() + p.join() + +if __name__ == "__main__": + main() +""" + file_context_manager = """ from viztracer import VizTracer, log_sparse @@ -199,6 +220,14 @@ def test_multiprocess(self): if not os.getenv("COVERAGE_RUN"): raise e + def test_multiprocess_spawn(self): + self.template(["viztracer", "-o", "result.json", "--log_sparse", "cmdline_test.py"], + script=file_multiprocess_spawn, + expected_output_file="result.json", + expected_entries=3, + check_func=functools.partial(self.check_func, target=['custom_event'] * 3), + concurrency="multiprocessing") + def test_context_manager(self): self.template(["python", "cmdline_test.py"], script=file_context_manager, expected_output_file="result.json", expected_entries=4,