diff --git a/dvc/_debug.py b/dvc/_debug.py index 7a94d7a8d..10cc63cc0 100644 --- a/dvc/_debug.py +++ b/dvc/_debug.py @@ -27,7 +27,9 @@ def viztracer_profile(path: Union[Callable[[], str], str], depth: int = -1): @contextmanager def yappi_profile( - path: Union[Callable[[], str], str] = None, wall_clock: bool = True + path: Union[Callable[[], str], str] = None, + wall_clock: Optional[bool] = True, + separate_threads: Optional[bool] = False, ): try: import yappi # pylint: disable=import-error @@ -42,14 +44,26 @@ def yappi_profile( yield yappi.stop() - # pylint:disable=no-member - if path: - stats = yappi.get_func_stats() - fpath = path() if callable(path) else path - stats.save(fpath, "callgrind") + threads = yappi.get_thread_stats() + stats = {} + if separate_threads: + for thread in threads: + ctx_id = thread.id + stats[ctx_id] = yappi.get_func_stats(ctx_id=ctx_id) else: - yappi.get_func_stats().print_all() - yappi.get_thread_stats().print_all() + stats[None] = yappi.get_func_stats() + + fpath = path() if callable(path) else path + for ctx_id, st in stats.items(): + if fpath: + out = f"{fpath}-{ctx_id}" if ctx_id is not None else fpath + st.save(out, type="callgrind") + else: + if ctx_id is not None: + print(f"\nThread {ctx_id}") + st.print_all() # pylint:disable=no-member + if ctx_id is None: + threads.print_all() # pylint:disable=no-member yappi.clear_stats() @@ -155,7 +169,12 @@ def debugtools(args: "Namespace" = None, **kwargs): stack.enter_context(show_stack()) if kw.get("yappi"): path_func = _get_path_func("callgrind", "out") - stack.enter_context(yappi_profile(path=path_func)) + stack.enter_context( + yappi_profile( + path=path_func, + separate_threads=kw.get("yappi_separate_threads"), + ) + ) if kw.get("viztracer") or kw.get("viztracer_depth"): path_func = _get_path_func("viztracer", "json") depth = kw.get("viztracer_depth") or -1 @@ -173,6 +192,12 @@ def add_debugging_flags(parser): parser.add_argument( "--yappi", action="store_true", default=False, help=SUPPRESS ) + parser.add_argument( + "--yappi-separate-threads", + action="store_true", + default=False, + help=SUPPRESS, + ) parser.add_argument( "--viztracer", action="store_true", default=False, help=SUPPRESS )