yappi: add --yappi-separate-threads flag

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
     )