diff --git a/dvc/_debug.py b/dvc/_debug.py index db358e299..ce3030075 100644 --- a/dvc/_debug.py +++ b/dvc/_debug.py @@ -12,13 +12,13 @@ def viztracer_profile( max_stack_depth: int = -1, ): try: - from viztracer import VizTracer # pylint: disable=import-error + import viztracer # pylint: disable=import-error except ImportError: print("Failed to run profiler, viztracer is not installed") yield return - tracer = VizTracer(max_stack_depth=max_stack_depth) + tracer = viztracer.VizTracer(max_stack_depth=max_stack_depth) tracer.start() yield @@ -126,12 +126,12 @@ def debugtools(args: "Namespace" = None, **kwargs): stack.enter_context( yappi_profile(path=lambda: output.format(datetime.now())) ) - if kw.get("viztracer"): + if kw.get("viztracer") or kw.get("viztracer_depth"): output = "viztracer.dvc-{0:%Y%m%d}_{0:%H%M%S}.json" stack.enter_context( viztracer_profile( path=lambda: output.format(datetime.now()), - max_stack_depth=kw.get("max_stack_depth", -1), + max_stack_depth=kw.get("viztracer_depth") or -1, ) ) yield @@ -149,7 +149,9 @@ def add_debugging_flags(parser): parser.add_argument( "--viztracer", action="store_true", default=False, help=SUPPRESS ) - parser.add_argument("--max-stack-depth", default=-1, help=SUPPRESS) + parser.add_argument( + "--viztracer-depth", type=int, default=-1, help=SUPPRESS + ) parser.add_argument("--cprofile-dump", help=SUPPRESS) parser.add_argument( "--pdb", action="store_true", default=False, help=SUPPRESS diff --git a/tests/unit/command/test_debug.py b/tests/unit/command/test_debug.py new file mode 100644 index 000000000..6de3c24f2 --- /dev/null +++ b/tests/unit/command/test_debug.py @@ -0,0 +1,38 @@ +from dvc.cli import main + + +def test_viztracer(tmp_dir, dvc, mocker): + viztracer_profile = mocker.patch("dvc._debug.viztracer_profile") + + assert main(["status", "--viztracer"]) == 0 + + args = viztracer_profile.call_args[1] + assert callable(args["path"]) + assert args["max_stack_depth"] == -1 + + assert main(["status", "--viztracer", "--viztracer-depth", "5"]) == 0 + args = viztracer_profile.call_args[1] + assert args["max_stack_depth"] == 5 + + assert main(["status", "--viztracer-depth", "2"]) == 0 + args = viztracer_profile.call_args[1] + assert args["max_stack_depth"] == 2 + + +def test_viztracer_internal(tmp_dir, dvc, mocker): + import sys + + viztracer = mocker.MagicMock() + sys.modules["viztracer"] = viztracer + from dvc import _debug + + profile = mocker.spy(_debug, "viztracer_profile") + tracer = viztracer.VizTracer.return_value + assert main(["status", "--viztracer"]) == 0 + + args = viztracer.VizTracer.call_args[1] + assert args["max_stack_depth"] == -1 + + tracer.start.assert_called_once() + tracer.stop.assert_called_once() + tracer.save.assert_called_once_with(profile.call_args[1]["path"]())