viztracer 1.1.1__cp313-cp313-win32.whl

viztracer/patch.py

@@ -0,0 +1,317 @@
+# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
+# For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
+
+from __future__ import annotations
+
+import functools
+import os
+import re
+import shutil
+import sys
+import textwrap
+from multiprocessing import Process
+from typing import Any, Callable, Sequence, no_type_check
+
+from .viztracer import VizTracer
+
+
+def patch_subprocess(viz_args: list[str]) -> None:
+    import shlex
+    import subprocess
+
+    # Try to detect the end of the python argument list and parse out various invocation patterns:
+    # `file.py args` | - args | `-- file.py args` | `-cprint(5) args` | `-Esm mod args`
+    py_arg_pat = re.compile("([^-].+)|-$|(--)$|-([a-z]+)?(c|m)(.+)?", re.IGNORECASE)
+    # Note: viztracer doesn't really work in interactive mode and arg handling is weird.
+    # Unlikely to be used in practice anyway so we just skip wrapping interactive python processes.
+    interactive_pat = re.compile("-[A-Za-z]*?i[A-Za-z]*$")
+
+    def build_command(args: Sequence[str]) -> list[str] | None:
+        py_args: list[str] = []
+        mode: list[str] | None = []
+        script = None
+        args_iter = iter(args[1:])
+        for arg in args_iter:
+            if interactive_pat.match(arg):
+                return None
+
+            match = py_arg_pat.match(arg)
+            if match:
+                file, ddash, cm_py_args, cm, cm_arg = match.groups()
+                if file:
+                    # file.py [script args]
+                    script = file
+                elif ddash:
+                    # -- file.py [script args]
+                    script = next(args_iter, None)
+                elif cm:
+                    # -m mod [script args]
+                    if cm_py_args:
+                        # "-[pyopts]m"
+                        py_args.append(f"-{cm_py_args}")
+                    mode = [f"-{cm}"]
+                    # -m mod | -mmod
+                    cm_arg = cm_arg or next(args_iter, None)
+                    if cm_arg is not None:
+                        if cm_arg.split('.')[0] == "viztracer":
+                            # Avoid tracing viztracer subprocess
+                            # This is mainly used to avoid tracing --open
+                            return None
+                        mode.append(cm_arg)
+                    else:
+                        mode = None
+                break
+
+            # -pyopts
+            py_args.append(arg)
+            if arg in ("-X", "-W", "--check-hash-based-pycs"):
+                arg_next = next(args_iter, None)
+                if arg_next is not None:
+                    py_args.append(arg_next)
+                else:
+                    return None
+
+        if script:
+            return [sys.executable, *py_args, "-m", "viztracer", "--quiet", *viz_args, "--", script, *args_iter]
+        elif mode:
+            return [sys.executable, *py_args, "-m", "viztracer", "--quiet", *viz_args, *mode, "--", *args_iter]
+        return None
+
+    def is_python_entry(path: str) -> bool:
+        real_path = shutil.which(path)
+        if real_path is None:
+            return False
+        try:
+            with open(real_path, "rb") as f:
+                if f.read(2) == b"#!":
+                    executable = f.readline().decode("utf-8").strip()
+                    if "python" in executable.split('/')[-1]:
+                        return True
+        except Exception:  # pragma: no cover
+            pass
+        return False
+
+    @functools.wraps(subprocess.Popen.__init__)
+    def subprocess_init(self: subprocess.Popen[Any], args: str | Sequence[Any] | Any, **kwargs: Any) -> None:
+        new_args = args
+        if isinstance(new_args, str):
+            new_args = shlex.split(new_args, posix=sys.platform != "win32")
+        if isinstance(new_args, Sequence):
+            if "python" in os.path.basename(new_args[0]):
+                new_args = build_command(new_args)
+            elif is_python_entry(new_args[0]):
+                new_args = ["python", "-m", "viztracer", "--quiet", *viz_args, "--", *new_args]
+            else:
+                new_args = None
+            if new_args is not None and kwargs.get("shell") and isinstance(args, str):
+                # For shell=True, we should convert the commands back to string
+                # if it was passed as string
+                # This is mostly for Unix shell
+                new_args = " ".join(new_args)
+
+        if new_args is None:
+            new_args = args
+        assert hasattr(subprocess_init, "__wrapped__")  # for mypy
+        subprocess_init.__wrapped__(self, new_args, **kwargs)
+
+    # We need to filter the arguments as there are something we may not want
+    if "-m" in viz_args:
+        # If it's a module run, we don't want to use that module for subprocess
+        idx = viz_args.index("-m")
+        viz_args.pop(idx)
+        viz_args.pop(idx)
+
+    setattr(subprocess.Popen, "__originit__", subprocess.Popen.__init__)
+    setattr(subprocess.Popen, "__init__", subprocess_init)
+
+
+def patch_multiprocessing(tracer: VizTracer, viz_args: list[str]) -> None:
+
+    # For fork process
+    def func_after_fork(tracer: VizTracer):
+        tracer.register_exit()
+
+        tracer.clear()
+        tracer.reset_stack()
+
+        if tracer._afterfork_cb:
+            tracer._afterfork_cb(tracer, *tracer._afterfork_args, **tracer._afterfork_kwargs)
+
+    import multiprocessing.spawn
+    import multiprocessing.util
+    from multiprocessing.util import register_after_fork  # type: ignore
+
+    register_after_fork(tracer, func_after_fork)
+
+    if sys.platform == "win32":
+        # For spawn process on Windows
+        @functools.wraps(multiprocessing.spawn.get_command_line)
+        def get_command_line(**kwds) -> list[str]:
+            """
+            Returns prefix of command line used for spawning a child process
+            """
+            if getattr(sys, 'frozen', False):  # pragma: no cover
+                return ([sys.executable, '--multiprocessing-fork']
+                        + ['%s=%r' % item for item in kwds.items()])
+            else:
+                prog = textwrap.dedent(f"""
+                        from multiprocessing.spawn import spawn_main;
+                        from viztracer.patch import patch_spawned_process;
+                        patch_spawned_process({tracer.init_kwargs}, {viz_args});
+                        spawn_main(%s)
+                        """)
+                prog %= ', '.join('%s=%r' % item for item in kwds.items())
+                opts = multiprocessing.util._args_from_interpreter_flags()  # type: ignore
+                return [multiprocessing.spawn._python_exe] + opts + ['-c', prog, '--multiprocessing-fork']  # type: ignore
+
+        multiprocessing.spawn.get_command_line = get_command_line
+    else:
+        # POSIX
+        # For forkserver process and spawned process
+        # We patch spawnv_passfds to trace forkserver parent process so the forked
+        # children can be traced
+        _spawnv_passfds = multiprocessing.util.spawnv_passfds
+
+        @functools.wraps(_spawnv_passfds)
+        def spawnv_passfds(path, args, passfds):
+            if "-c" in args:
+                idx = args.index("-c")
+                cmd = args[idx + 1]
+                if "forkserver" in cmd:
+                    # forkserver will not end before main process, avoid deadlock by --patch_only
+                    args = (
+                        args[:idx]
+                        + ["-m", "viztracer", "--patch_only", *viz_args]
+                        + ["--subprocess_child", "--dump_raw", "-o", tracer.output_file]
+                        + args[idx:]
+                    )
+                elif "resource_tracker" not in cmd:
+                    # We don't trace resource_tracker as it does not quit before the main process
+                    # This is a normal spawned process. Only one of spawnv_passfds and spawn._main
+                    # can be patched. forkserver process will use spawn._main after forking a child,
+                    # so on POSIX we patch spawnv_passfds which has a similar effect on spawned processes.
+                    args = (
+                        args[:idx]
+                        + ["-m", "viztracer", *viz_args]
+                        + ["--subprocess_child", "--dump_raw", "-o", tracer.output_file]
+                        + args[idx:]
+                    )
+            ret = _spawnv_passfds(path, args, passfds)
+            return ret
+
+        multiprocessing.util.spawnv_passfds = spawnv_passfds  # type: ignore
+
+
+class SpawnProcess:
+    def __init__(
+            self,
+            viztracer_kwargs: dict[str, Any],
+            run: Callable,
+            target: Callable,
+            args: list[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()
+        if not self._viztracer_kwargs.get("log_sparse"):
+            tracer.start()
+        self._run()
+
+
+def patch_spawned_process(viztracer_kwargs: dict[str, Any], cmdline_args: list[str]):
+    import multiprocessing.spawn
+    from multiprocessing import process, reduction  # type: ignore
+    from multiprocessing.spawn import prepare
+
+    @no_type_check
+    @functools.wraps(multiprocessing.spawn._main)
+    def _main(fd, parent_sentinel) -> Any:
+        with os.fdopen(fd, 'rb', closefd=True) as from_parent:
+            process.current_process()._inheriting = True
+            try:
+                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, cmdline_args)
+                self.run = sp.run
+            finally:
+                del process.current_process()._inheriting
+        return self._bootstrap(parent_sentinel)
+
+    multiprocessing.spawn._main = _main  # type: ignore
+
+
+def filter_args(args: list[str]) -> list[str]:
+    new_args = []
+    i = 0
+    while i < len(args):
+        arg = args[i]
+        if arg == "-u" or arg == "--unique_output_file":
+            i += 1
+            continue
+        elif arg == "-o" or arg == "--output_file":
+            i += 2
+            continue
+        new_args.append(arg)
+        i += 1
+    return new_args
+
+
+def install_all_hooks(
+        tracer: VizTracer,
+        args: list[str],
+        patch_multiprocess: bool = True) -> None:
+
+    args = filter_args(args)
+
+    # 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
+
+            def callback():
+                if "--patch_only" in args:
+                    # We use --patch_only for forkserver process so we need to
+                    # turn on tracer in the forked child process and register
+                    # for exit routine
+                    tracer.register_exit()
+                    tracer.start()
+                else:
+                    # otherwise we need to add a new exit_routine callback because the one
+                    # from parent won't be executed as it has a different pid.
+                    # Also make sure to label the file because it's a new process
+                    import multiprocessing.util
+                    multiprocessing.util.Finalize(tracer, tracer.exit_routine, exitpriority=-1)
+                    tracer.label_file_to_write()
+            os.register_at_fork(after_in_child=callback)  # 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

viztracer/report_builder.py

@@ -0,0 +1,311 @@
+# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
+# For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
+
+try:
+    import orjson as json  # type: ignore
+except ImportError:
+    import json  # type: ignore
+
+import gzip
+import importlib
+import os
+import re
+import tokenize
+from string import Template
+from typing import Any, Sequence, TextIO
+
+from . import __version__
+from .util import color_print, same_line_print
+
+
+def get_json(data: dict[str, Any] | str | tuple[str, dict]) -> dict[str, Any]:
+    # This function will return a json object if data is already json object
+    # or a opened file or a file path
+    if isinstance(data, dict):
+        # This is an object already
+        return data
+    elif isinstance(data, str):
+        with open(data, encoding="utf-8") as f:
+            json_str = f.read()
+    elif isinstance(data, tuple):
+        path, args = data
+        if args['type'] == 'torch':
+            with open(path, encoding="utf-8") as f:
+                json_str = f.read()
+            ret = json.loads(json_str)
+            base_offset = args['base_offset']
+            # torch 2.4.0+ uses baseTimeNanoseconds to store the offset
+            # before that they simply use the absolute timestamp which is
+            # equivalent to baseTimeNanoseconds = 0
+            torch_offset = ret.get('baseTimeNanoseconds', 0)
+            # convert to us
+            offset_diff = (torch_offset - base_offset) / 1000
+
+            for event in ret['traceEvents']:
+                if 'ts' in event:
+                    event['ts'] += offset_diff
+                if event['ph'] == 'M':
+                    # Pop metadata timestamp so it won't overwrite
+                    # process and thread names
+                    event.pop('ts', None)
+
+            ret.pop("baseTimeNanoseconds", None)
+            ret.pop("displayTimeUnit", None)
+            ret.pop("traceName")
+            ret.pop("deviceProperties")
+            ret.pop("schemaVersion")
+
+            ret['viztracer_metadata'] = {}
+
+            return ret
+
+    return json.loads(json_str)
+
+
+class ReportBuilder:
+    def __init__(
+            self,
+            data: Sequence[str | dict | tuple[str, dict]] | dict[str, Any],
+            verbose: int = 1,
+            align: bool = False,
+            minimize_memory: bool = False,
+            base_time: int | None = None) -> None:
+        self.data = data
+        self.verbose = verbose
+        self.combined_json: dict = {}
+        self.entry_number_threshold = 4000000
+        self.align = align
+        self.minimize_memory = minimize_memory
+        self.jsons: list[dict] = []
+        self.invalid_json_paths: list[str] = []
+        self.json_loaded = False
+        self.base_time = base_time
+        self.final_messages: list[tuple[str, dict]] = []
+        if not isinstance(data, (dict, list, tuple)):
+            raise TypeError("Invalid data type for ReportBuilder")
+        if isinstance(data, (list, tuple)):
+            for path in data:
+                if isinstance(path, dict):
+                    continue
+                if isinstance(path, tuple):
+                    path = path[0]
+                if not isinstance(path, str):
+                    raise TypeError("Path should be a string")
+                if not os.path.exists(path):
+                    raise ValueError(f"{path} does not exist")
+                if not path.endswith(".json"):
+                    raise ValueError(f"{path} is not a json file")
+
+    def load_jsons(self) -> None:
+        if not self.json_loaded:
+            self.json_loaded = True
+            if isinstance(self.data, dict):
+                self.jsons = [get_json(self.data)]
+            elif isinstance(self.data, (list, tuple)):
+                self.jsons = []
+                self.invalid_json_paths = []
+                for idx, j in enumerate(self.data):
+                    if self.verbose > 0:
+                        same_line_print(f"Loading trace data from processes {idx}/{len(self.data)}")
+                    try:
+                        self.jsons.append(get_json(j))
+                    except json.JSONDecodeError:
+                        assert isinstance(j, str)
+                        self.invalid_json_paths.append(j)
+                if len(self.invalid_json_paths) > 0:
+                    self.final_messages.append(("invalid_json", {"paths": self.invalid_json_paths}))
+
+    def combine_json(self) -> None:
+        if self.verbose > 0:
+            same_line_print("Combining trace data")
+        if self.combined_json:
+            return
+        if not self.jsons:
+            if self.invalid_json_paths:
+                raise ValueError("No valid json files found")
+            else:
+                raise ValueError("Can't get report of nothing")
+        if self.align:
+            for one in self.jsons:
+                self.align_events(one["traceEvents"], one['viztracer_metadata'].get('sync_marker', None))
+        self.combined_json = self.jsons[0]
+        if "viztracer_metadata" not in self.combined_json:
+            self.combined_json["viztracer_metadata"] = {}
+        for one in self.jsons[1:]:
+            if "traceEvents" in one:
+                self.combined_json["traceEvents"].extend(one["traceEvents"])
+            if one.get("viztracer_metadata", {}).get("overflow", False):
+                self.combined_json["viztracer_metadata"]["overflow"] = True
+            if one.get("viztracer_metadata", {}).get("baseTimeNanoseconds") is not None:
+                self.combined_json["viztracer_metadata"]["baseTimeNanoseconds"] = \
+                    one["viztracer_metadata"]["baseTimeNanoseconds"]
+            if "file_info" in one:
+                if "file_info" not in self.combined_json:
+                    self.combined_json["file_info"] = {"files": {}, "functions": {}}
+                self.combined_json["file_info"]["files"].update(one["file_info"]["files"])
+                self.combined_json["file_info"]["functions"].update(one["file_info"]["functions"])
+
+    def align_events(self, original_events: list[dict[str, Any]], sync_marker: float | None = None) -> list[dict[str, Any]]:
+        """
+        Apply an offset to all the trace events, making the start timestamp 0
+        This is useful when comparing multiple runs of the same script
+
+        If sync_marker is not None then sync_marker  be used as an offset
+
+        This function will change the timestamp in place, and return the original list
+        """
+        if sync_marker is None:
+            offset_ts = min((event["ts"] for event in original_events if "ts" in event))
+        else:
+            offset_ts = sync_marker
+
+        for event in original_events:
+            if "ts" in event:
+                event["ts"] -= offset_ts
+        return original_events
+
+    def prepare_json(self, file_info: bool = True, display_time_unit: str | None = None) -> None:
+        # This will prepare self.combined_json to be ready to output
+        self.load_jsons()
+        self.combine_json()
+        if self.verbose > 0:
+            entries = len(self.combined_json["traceEvents"])
+            same_line_print(f"Dumping trace data, total entries: {entries}")
+            self.final_messages.append(("total_entries", {"total_entries": entries}))
+            if self.combined_json["viztracer_metadata"].get("overflow", False):
+                self.final_messages.append(("overflow", {}))
+
+        if display_time_unit is not None:
+            self.combined_json["displayTimeUnit"] = display_time_unit
+
+        self.combined_json["viztracer_metadata"]["version"] = __version__
+
+        if self.base_time is not None:
+            self.combined_json["viztracer_metadata"]["baseTimeNanoseconds"] = self.base_time
+
+        if file_info:
+            if "file_info" not in self.combined_json:
+                self.combined_json["file_info"] = {"files": {}, "functions": {}}
+            pattern = re.compile(r".*\((.*):([0-9]*)\)")
+            file_dict = self.combined_json["file_info"]["files"]
+            func_dict = self.combined_json["file_info"]["functions"]
+            for event in self.combined_json["traceEvents"]:
+                if event["ph"] == 'X':
+                    if event["name"] not in func_dict:
+                        func_dict[event["name"]] = None
+                        m = pattern.match(event["name"])
+                        if m is not None:
+                            file_name = m.group(1)
+                            lineno = int(m.group(2))
+                            if file_name not in file_dict:
+                                content = self.get_source_from_filename(file_name)
+                                if content is None:
+                                    continue
+                                file_dict[file_name] = [content, content.count("\n")]
+                            func_dict[event["name"]] = [file_name, lineno]
+            unknown_func_dict = set(func for func in func_dict if func_dict[func] is None)
+            for func in unknown_func_dict:
+                del func_dict[func]
+
+    @classmethod
+    def get_source_from_filename(cls, filename: str) -> str | None:
+        if filename.startswith("<frozen "):
+            m = re.match(r"<frozen (.*)>", filename)
+            if not m:
+                return None
+            module_name = m.group(1)
+            try:
+                module = importlib.import_module(module_name)
+            except ImportError:
+                return None
+            if hasattr(module, "__file__") and module.__file__ is not None:
+                filename = module.__file__
+            else:
+                return None
+        try:
+            with tokenize.open(filename) as f:
+                return f.read()
+        except Exception:
+            return None
+
+    def generate_report(
+            self,
+            output_file: TextIO,
+            output_format: str,
+            file_info: bool = True) -> None:
+        sub = {}
+        if output_format == "html":
+            self.prepare_json(file_info=file_info, display_time_unit="ns")
+            with open(os.path.join(os.path.dirname(__file__), "html/trace_viewer_embedder.html"), encoding="utf-8") as f:
+                tmpl = f.read()
+            with open(os.path.join(os.path.dirname(__file__), "html/trace_viewer_full.html"), encoding="utf-8") as f:
+                sub["trace_viewer_full"] = f.read()
+            if json.__name__ == "orjson":
+                sub["json_data"] = json.dumps(self.combined_json) \
+                                       .decode("utf-8") \
+                                       .replace("</script>", "<\\/script>")
+            else:
+                sub["json_data"] = json.dumps(self.combined_json) \
+                                       .replace("</script>", "<\\/script>")  # type: ignore
+            output_file.write(Template(tmpl).substitute(sub))
+        elif output_format == "json":
+            self.prepare_json(file_info=file_info)
+            if json.__name__ == "orjson":
+                output_file.write(json.dumps(self.combined_json).decode("utf-8"))
+            else:
+                if self.minimize_memory:
+                    json.dump(self.combined_json, output_file)  # type: ignore
+                else:
+                    output_file.write(json.dumps(self.combined_json))  # type: ignore
+
+    def save(self, output_file: str | TextIO = "result.html", file_info: bool = True) -> None:
+        if isinstance(output_file, str):
+            file_type = output_file.split(".")[-1]
+
+            if file_type == "html":
+                with open(output_file, "w", encoding="utf-8") as f:
+                    self.generate_report(f, output_format="html", file_info=file_info)
+            elif file_type == "json":
+                with open(output_file, "w", encoding="utf-8") as f:
+                    self.generate_report(f, output_format="json", file_info=file_info)
+            elif file_type == "gz":
+                with gzip.open(output_file, "wt") as f:
+                    self.generate_report(f, output_format="json", file_info=file_info)
+            else:
+                raise Exception("Only html, json and gz are supported")
+        else:
+            self.generate_report(output_file, output_format="json", file_info=file_info)
+
+        if isinstance(output_file, str):
+            self.final_messages.append(("view_command", {"output_file": os.path.abspath(output_file)}))
+
+        self.print_messages()
+
+    def print_messages(self):
+        if self.verbose > 0:
+            same_line_print("")
+            for msg_type, msg_args in self.final_messages:
+                if msg_type == "overflow":
+                    print("")
+                    color_print("WARNING", ("Circular buffer is full, you lost some early data, "
+                                            "but you still have the most recent data."))
+                    color_print("WARNING", ("    If you need more buffer, use \"viztracer --tracer_entries <entry_number>\""))
+                    color_print("WARNING", "    Or, you can try the filter options to filter out some data you don't need")
+                    color_print("WARNING", "    use --quiet to shut me up")
+                    print("")
+                elif msg_type == "total_entries":
+                    print(f"Total Entries: {msg_args['total_entries']}")
+                elif msg_type == "view_command":
+                    report_abspath = os.path.abspath(msg_args["output_file"])
+                    print("Use the following command to open the report:")
+                    if " " in report_abspath:
+                        color_print("OKGREEN", f"vizviewer \"{report_abspath}\"")
+                    else:
+                        color_print("OKGREEN", f"vizviewer {report_abspath}")
+                elif msg_type == "invalid_json":
+                    print("")
+                    color_print("WARNING", "Found and ignored invalid json file, you may lost some process data.")
+                    color_print("WARNING", "Invalid json file:")
+                    for msg in msg_args["paths"]:
+                        color_print("WARNING", f"    {msg}")
+                    print("")