dvc/dvc/_debug.py

from contextlib import ExitStack, contextmanager
from datetime import datetime
from typing import TYPE_CHECKING, Callable, Optional, Union

if TYPE_CHECKING:
    from argparse import Namespace
    from types import FrameType


@contextmanager
def viztracer_profile(
    path: Union[Callable[[], str], str],
    depth: int = -1,
    log_async: bool = False,
):
    try:
        import viztracer  # ty: ignore[unresolved-import]
    except ImportError:
        print("Failed to run profiler, viztracer is not installed")  # noqa: T201
        yield
        return

    tracer = viztracer.VizTracer(max_stack_depth=depth, log_async=log_async)

    tracer.start()
    try:
        yield
    finally:
        tracer.stop()
        tracer.save(path() if callable(path) else path)


@contextmanager
def yappi_profile(
    path: Optional[Union[Callable[[], str], str]] = None,
    wall_clock: Optional[bool] = True,
    separate_threads: Optional[bool] = False,
):
    try:
        import yappi  # ty: ignore[unresolved-import]
    except ImportError:
        print("Failed to run profiler, yappi is not installed")  # noqa: T201
        yield
        return

    yappi.set_clock_type("wall" if wall_clock else "cpu")

    yappi.start()
    try:
        yield
    finally:
        yappi.stop()

        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:
            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}")  # noqa: T201
                st.print_all()
                if ctx_id is None:
                    threads.print_all()

        yappi.clear_stats()


@contextmanager
def instrument(html_output=False):
    """Run a statistical profiler"""
    try:
        from pyinstrument import Profiler  # ty: ignore[unresolved-import]
    except ImportError:
        print("Failed to run profiler, pyinstrument is not installed")  # noqa: T201
        yield
        return

    profiler = Profiler()

    profiler.start()
    try:
        yield
    finally:
        profiler.stop()

        if html_output:
            profiler.open_in_browser()
        else:
            print(profiler.output_text(unicode=True, color=True))  # noqa: T201


@contextmanager
def profile(dump_path: Optional[str] = None):
    """Run a cprofile"""
    import cProfile

    prof = cProfile.Profile()
    prof.enable()

    try:
        yield
    finally:
        prof.disable()
        if dump_path:
            prof.dump_stats(dump_path)
        else:
            prof.print_stats(sort="cumtime")


@contextmanager
def debug():
    try:
        yield
    except Exception:
        try:
            import ipdb as pdb  # noqa: T100  # ty: ignore[unresolved-import]
        except ImportError:
            import pdb  # type: ignore[no-redef]  # noqa: T100
        pdb.post_mortem()

        raise  # prevent from jumping ahead


def _sigshow(_, frame: Optional["FrameType"]) -> None:
    import sys
    from shutil import get_terminal_size
    from traceback import format_stack

    lines = "\u2015" * get_terminal_size().columns
    stack = format_stack(frame)
    print(lines, "\n", *stack, lines, sep="", file=sys.stderr)  # noqa: T201


@contextmanager
def show_stack():
    r"""Show stack trace on SIGQUIT (Ctrl-\) or SIGINFO (Ctrl-T on macOS)."""
    import signal
    import sys

    if sys.platform != "win32":
        signal.signal(signal.SIGQUIT, _sigshow)

    try:
        # only available on macOS
        signal.signal(signal.SIGINFO, _sigshow)  # type: ignore[attr-defined]
    except AttributeError:
        pass
    yield


def _get_path_func(tool: str, ext: str):
    fmt = f"{tool}.dvc-{{now:%Y%m%d}}_{{now:%H%M%S}}.{ext}"

    def func(now: Optional["datetime"] = None) -> str:
        return fmt.format(now=now or datetime.now())  # noqa: DTZ005

    return func


@contextmanager
def debugtools(args: Optional["Namespace"] = None, **kwargs):
    kw = vars(args) if args else {}
    kw.update(kwargs)

    with ExitStack() as stack:
        if kw.get("pdb"):
            stack.enter_context(debug())
        if kw.get("cprofile") or kw.get("cprofile_dump"):
            stack.enter_context(profile(kw.get("cprofile_dump")))
        if kw.get("instrument") or kw.get("instrument_open"):
            stack.enter_context(instrument(kw.get("instrument_open", False)))
        if kw.get("show_stack", False):
            stack.enter_context(show_stack())
        if kw.get("yappi"):
            path_func = _get_path_func("callgrind", "out")
            stack.enter_context(
                yappi_profile(
                    path=path_func,
                    separate_threads=kw.get("yappi_separate_threads"),
                )
            )
        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
            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


def add_debugging_flags(parser):
    from argparse import SUPPRESS

    # For detailed info see:
    # https://github.com/treeverse/dvc/wiki/Debugging,-Profiling-and-Benchmarking-DVC
    args, _ = parser.parse_known_args()
    verbose = args.verbose

    def debug_help(msg):
        if verbose:
            return msg
        return SUPPRESS

    parser = parser.add_argument_group("debug options")

    parser.add_argument(
        "--cprofile",
        action="store_true",
        default=False,
        help=debug_help("Generate cprofile data for tools like snakeviz / tuna"),
    )
    parser.add_argument(
        "--cprofile-dump", help=debug_help("Location to dump cprofile file")
    )
    parser.add_argument(
        "--yappi",
        action="store_true",
        default=False,
        help=debug_help(
            "Generate a callgrind file for use with tools like "
            "kcachegrind / qcachegrind"
        ),
    )
    parser.add_argument(
        "--yappi-separate-threads",
        action="store_true",
        default=False,
        help=debug_help("Generate one callgrind file per thread"),
    )
    parser.add_argument(
        "--viztracer",
        action="store_true",
        default=False,
        help=debug_help("Generate a viztracer file for use with vizviewer"),
    )
    parser.add_argument(
        "--viztracer-depth",
        type=int,
        help=debug_help("Set viztracer maximum stack depth"),
    )
    parser.add_argument(
        "--viztracer-async",
        action="store_true",
        default=False,
        help=debug_help("Treat async tasks as threads"),
    )
    parser.add_argument(
        "--pdb",
        action="store_true",
        default=False,
        help=debug_help("Drop into the pdb/ipdb debugger on any exception"),
    )
    parser.add_argument(
        "--instrument",
        action="store_true",
        default=False,
        help=debug_help("Use pyinstrument CLI profiler"),
    )
    parser.add_argument(
        "--instrument-open",
        action="store_true",
        default=False,
        help=debug_help("Use pyinstrument web profiler"),
    )
    parser.add_argument(
        "--show-stack",
        "--ss",
        action="store_true",
        default=False,
        help=debug_help(
            r"Use Ctrl+T on macOS or Ctrl+\ on Linux to print the stack "
            "frame currently executing. Unavailable on Windows."
        ),
    )