From 7e9a9a8983e018890f5a16a5629eaee5b9ea261a Mon Sep 17 00:00:00 2001
From: Oleksandr Pavlyk <21087696+oleksandr-pavlyk@users.noreply.github.com>
Date: Mon, 8 Dec 2025 14:28:26 -0600
Subject: [PATCH] Replace main_arg_run_benchmarks with run_interriptible

This loop uses benchmark.run_or_skip to resolve #284 even
for scripts that contain more than one benchmark, or when
a script with a single benchmark is executed when more than
one device is available.
---
 python/src/py_nvbench.cpp | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/python/src/py_nvbench.cpp b/python/src/py_nvbench.cpp
index e1b86ea..159936f 100644
--- a/python/src/py_nvbench.cpp
+++ b/python/src/py_nvbench.cpp
@@ -119,6 +119,29 @@ struct nvbench_run_error : std::runtime_error
 };
 py::handle benchmark_exc{};
 
+void run_interruptible(nvbench::option_parser &parser)
+{
+  auto &printer    = parser.get_printer();
+  auto &benchmarks = parser.get_benchmarks();
+
+  std::size_t total_states = 0;
+  for (auto &bench_ptr : benchmarks)
+  {
+    total_states += bench_ptr->get_config_count();
+  }
+
+  printer.set_completed_state_count(0);
+  printer.set_total_state_count(total_states);
+
+  bool skip_remaining_flag = false;
+  for (auto &bench_ptr : benchmarks)
+  {
+    bench_ptr->set_printer(printer);
+    bench_ptr->run_or_skip(skip_remaining_flag);
+    bench_ptr->clear_printer();
+  }
+}
+
 class GlobalBenchmarkRegistry
 {
   bool m_finalized;
@@ -185,7 +208,7 @@ public:
         parser.parse(argv);
 
         NVBENCH_MAIN_PRINT_PREAMBLE(parser);
-        NVBENCH_MAIN_RUN_BENCHMARKS(parser);
+        run_interruptible(parser);
         NVBENCH_MAIN_PRINT_EPILOGUE(parser);
 
         NVBENCH_MAIN_PRINT_RESULTS(parser);