add viztracer async option (#8856)

* add viztracer async option

* Update dvc/_debug.py

Co-authored-by: skshetry <[email protected]>
This commit is contained in:
Dave Berenbaum
2023-01-20 22:29:11 +05:45
committed by GitHub
co-authored by skshetry
parent f1c9e9365b
commit 18f039d7c7
+18 -4
View File
@@ -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