Enable profiling from cli (#4437)

dvc/cli.py

@@ -137,6 +137,14 @@ def get_parent_parser():
     """
     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)
+
     log_level_group = parent_parser.add_mutually_exclusive_group()
     log_level_group.add_argument(
         "-q", "--quiet", action="count", default=0, help="Be quiet."

dvc/main.py

@@ -2,6 +2,7 @@
 
 import errno
 import logging
+from contextlib import contextmanager
 
 from dvc import analytics
 from dvc.cli import parse_args
@@ -21,6 +22,25 @@ from dvc.utils import format_link
 logger = logging.getLogger("dvc")
 
 
+@contextmanager
+def profile(enable, dump):
+    if not enable:
+        yield
+        return
+
+    import cProfile
+
+    prof = cProfile.Profile()
+    prof.enable()
+
+    yield
+
+    if not dump:
+        prof.print_stats(sort="cumtime")
+        return
+    prof.dump_stats(dump)
+
+
 def main(argv=None):  # noqa: C901
     """Run dvc CLI command.
 
@@ -50,8 +70,9 @@ def main(argv=None):  # noqa: C901
 
         logger.trace(args)
 
-        cmd = args.func(args)
-        ret = cmd.run()
+        with profile(enable=args.cprofile, dump=args.cprofile_dump):
+            cmd = args.func(args)
+            ret = cmd.run()
     except ConfigError:
         logger.exception("configuration error")
         ret = 251