From 261ddc947b7c57b2ad56781f60232b32f6281039 Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Mon, 4 Jan 2021 15:23:40 +0545
Subject: [PATCH] Add pyinstrument: a statistical profiler (#5196)

* Add pyinstrument: a statistical profiler

It is easier than cprofile, in that it can show stacktrace
of where things are slow.

Also refactored/moved debugging related utils to a separate
module

* use ipdb first, then fallback to pdb

* use ipdb first, then fallback to pdb
---
 dvc/_debug.py | 90 +++++++++++++++++++++++++++++++++++++++++++++++++++
 dvc/cli.py    | 15 ++-------
 dvc/main.py   | 42 +++---------------------
 3 files changed, 96 insertions(+), 51 deletions(-)
 create mode 100644 dvc/_debug.py

diff --git a/dvc/_debug.py b/dvc/_debug.py
new file mode 100644
index 000000000..b3d776a99
--- /dev/null
+++ b/dvc/_debug.py
@@ -0,0 +1,90 @@
+import argparse
+from contextlib import ExitStack, contextmanager
+
+
+@contextmanager
+def instrument(html_output):
+    """Run a statistical profiler"""
+    try:
+        from pyinstrument import Profiler  # pylint: disable=import-error
+    except ImportError:
+        print("Failed to run profiler, pyinstrument is not installed")
+        yield
+        return
+
+    profiler = Profiler()
+
+    profiler.start()
+    yield
+    profiler.stop()
+
+    if html_output:
+        profiler.open_in_browser()
+        return
+    print(profiler.output_text(unicode=True, color=True))
+
+
+@contextmanager
+def profile(dump):
+    """Run a cprofile"""
+    import cProfile
+
+    prof = cProfile.Profile()
+    prof.enable()
+
+    yield
+
+    prof.disable()
+    if not dump:
+        prof.print_stats(sort="cumtime")
+        return
+    prof.dump_stats(dump)
+
+
+@contextmanager
+def debug():
+    try:
+        yield
+    except Exception:  # pylint: disable=broad-except
+        try:
+            import ipdb as pdb  # noqa: T100, pylint: disable=import-error
+        except ImportError:
+            import pdb  # noqa: T100
+        pdb.post_mortem()
+
+
+@contextmanager
+def debugtools(args):
+    with ExitStack() as stack:
+        if args.pdb:
+            stack.enter_context(debug())
+        if args.cprofile:
+            stack.enter_context(profile(args.cprofile_dump))
+        if args.instrument:
+            stack.enter_context(instrument(args.instrument_open))
+        yield
+
+
+def add_debugging_flags(parser):
+    parser.add_argument(
+        "--cprofile",
+        action="store_true",
+        default=False,
+        help=argparse.SUPPRESS,
+    )
+    parser.add_argument("--cprofile-dump", help=argparse.SUPPRESS)
+    parser.add_argument(
+        "--pdb", action="store_true", default=False, help=argparse.SUPPRESS,
+    )
+    parser.add_argument(
+        "--instrument",
+        action="store_true",
+        default=False,
+        help=argparse.SUPPRESS,
+    )
+    parser.add_argument(
+        "--instrument-open",
+        action="store_true",
+        default=False,
+        help=argparse.SUPPRESS,
+    )
diff --git a/dvc/cli.py b/dvc/cli.py
index 48ebbc292..2a346da96 100644
--- a/dvc/cli.py
+++ b/dvc/cli.py
@@ -4,6 +4,7 @@ import logging
 import os
 import sys
 
+from ._debug import add_debugging_flags
 from .command import (
     add,
     cache,
@@ -139,19 +140,7 @@ def get_parent_parser():
     in order to prevent some weird behavior.
     """
     parent_parser = argparse.ArgumentParser(add_help=False)
-
-    parent_parser.add_argument(
-        "--cprofile",
-        action="store_true",
-        default=False,
-        help=argparse.SUPPRESS,
-    )
-    parent_parser.add_argument("--cprofile-dump", help=argparse.SUPPRESS)
-
-    parent_parser.add_argument(
-        "--pdb", action="store_true", default=False, help=argparse.SUPPRESS,
-    )
-
+    add_debugging_flags(parent_parser)
     log_level_group = parent_parser.add_mutually_exclusive_group()
     log_level_group.add_argument(
         "-q", "--quiet", action="count", default=0, help="Be quiet."
diff --git a/dvc/main.py b/dvc/main.py
index c4c11421d..da0225f6e 100644
--- a/dvc/main.py
+++ b/dvc/main.py
@@ -2,9 +2,9 @@
 
 import errno
 import logging
-from contextlib import contextmanager
 
 from dvc import analytics
+from dvc._debug import debugtools
 from dvc.cli import parse_args
 from dvc.config import ConfigError
 from dvc.exceptions import DvcException, DvcParserError, NotDvcRepoError
@@ -22,39 +22,6 @@ from dvc.utils import error_link
 logger = logging.getLogger("dvc")
 
 
-@contextmanager
-def profile(enable, dump):
-    if not enable:
-        yield
-        return
-
-    import cProfile
-
-    prof = cProfile.Profile()
-    prof.enable()
-
-    yield
-
-    prof.disable()
-    if not dump:
-        prof.print_stats(sort="cumtime")
-        return
-    prof.dump_stats(dump)
-
-
-@contextmanager
-def debug(enable):
-    try:
-        yield
-        return
-    except Exception:
-        if enable:
-            import pdb  # noqa: T100
-
-            pdb.post_mortem()
-        raise
-
-
 def main(argv=None):  # noqa: C901
     """Run dvc CLI command.
 
@@ -84,10 +51,9 @@ def main(argv=None):  # noqa: C901
 
         logger.trace(args)
 
-        with profile(enable=args.cprofile, dump=args.cprofile_dump):
-            with debug(args.pdb):
-                cmd = args.func(args)
-                ret = cmd.run()
+        with debugtools(args):
+            cmd = args.func(args)
+            ret = cmd.run()
     except ConfigError:
         logger.exception("configuration error")
         ret = 251