From 3f2584d155e1b1f79d8b37bc7944cea843f81af6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?=
 =?UTF-8?q?=E0=A4=A4=29?= <suagatchhetri@outlook.com>
Date: Thu, 22 Aug 2024 14:35:23 +0545
Subject: [PATCH] bench: add --benchmark-cprofile-dump

---
 dvc/testing/benchmarks/fixtures.py | 14 +++++++++++---
 dvc/testing/benchmarks/plugin.py   |  7 +++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/dvc/testing/benchmarks/fixtures.py b/dvc/testing/benchmarks/fixtures.py
index f5170be53f..b8ebcedb87 100644
--- a/dvc/testing/benchmarks/fixtures.py
+++ b/dvc/testing/benchmarks/fixtures.py
@@ -118,9 +118,9 @@ def make_dvc_bin(
         dvc_bin = bench_config.dvc_bin
 
     def _dvc_bin(*args):
-        return check_output([dvc_bin, *args], text=True)  # noqa: S603
+        check_call([dvc_bin, *args])  # noqa: S603
 
-    _dvc_bin.version = _dvc_bin("--version")  # type: ignore[attr-defined]
+    _dvc_bin.version = check_output([dvc_bin, "--version"], text=True)  # type: ignore[attr-defined]  # noqa: S603
     return _dvc_bin
 
 
@@ -167,11 +167,19 @@ def add_suffix(_name):
 
 
 @pytest.fixture
-def bench_dvc(dvc_bin, make_bench):
+def bench_dvc(request, dvc_bin, make_bench):
     def _bench_dvc(*args, **kwargs):
         name = kwargs.pop("name", None)
         name = f"-{name}" if name else ""
         bench = make_bench(args[0] + name)
+        if request.config.getoption("--benchmark-cprofile-dump") or kwargs.pop(
+            "cprofile", False
+        ):
+            cprofile_results = request.config.invocation_params.dir / "prof"
+            cprofile_results.mkdir(exist_ok=True)
+            stats_file = cprofile_results / f"{bench.name}.prof"
+            args = (*args, "--cprofile-dump", stats_file)
+
         return bench.pedantic(dvc_bin, args=args, **kwargs)
 
     return _bench_dvc
diff --git a/dvc/testing/benchmarks/plugin.py b/dvc/testing/benchmarks/plugin.py
index 91bc716c28..34ae45bd31 100644
--- a/dvc/testing/benchmarks/plugin.py
+++ b/dvc/testing/benchmarks/plugin.py
@@ -72,6 +72,13 @@ def pytest_addoption(parser):
         help="Dataset name to use in tests (e.g. tiny/small/large/mnist/etc)",
     )
 
+    parser.addoption(
+        "--benchmark-cprofile-dump",
+        action="store_true",
+        default=False,
+        help="Save cprofile results",
+    )
+
     parser.addoption(
         "--dvc-bin",
         type=str,