diff --git a/dvc/_debug.py b/dvc/_debug.py index 20c4098c1..2cf877d76 100644 --- a/dvc/_debug.py +++ b/dvc/_debug.py @@ -8,7 +8,11 @@ if TYPE_CHECKING: @contextmanager -def viztracer_profile(path: Union[Callable[[], str], str], depth: int = -1): +def viztracer_profile( + path: Union[Callable[[], str], str], + depth: int = -1, + log_async: bool = False, +): try: import viztracer # pylint: disable=import-error except ImportError: @@ -18,7 +22,7 @@ def viztracer_profile(path: Union[Callable[[], str], str], depth: int = -1): yield return - tracer = viztracer.VizTracer(max_stack_depth=depth) + tracer = viztracer.VizTracer(max_stack_depth=depth, log_async=log_async) tracer.start() yield @@ -182,10 +186,17 @@ def debugtools(args: Optional["Namespace"] = None, **kwargs): separate_threads=kw.get("yappi_separate_threads"), ) ) - if kw.get("viztracer") or kw.get("viztracer_depth"): + if ( + kw.get("viztracer") + or kw.get("viztracer_depth") + or kw.get("viztracer_async") + ): path_func = _get_path_func("viztracer", "json") depth = kw.get("viztracer_depth") or -1 - prof = viztracer_profile(path=path_func, depth=depth) + log_async = kw.get("viztracer_async") or False + prof = viztracer_profile( + path=path_func, depth=depth, log_async=log_async + ) stack.enter_context(prof) yield @@ -209,6 +220,9 @@ def add_debugging_flags(parser): "--viztracer", action="store_true", default=False, help=SUPPRESS ) parser.add_argument("--viztracer-depth", type=int, help=SUPPRESS) + parser.add_argument( + "--viztracer-async", action="store_true", default=False, help=SUPPRESS + ) parser.add_argument("--cprofile-dump", help=SUPPRESS) parser.add_argument( "--pdb", action="store_true", default=False, help=SUPPRESS