NVIDIA/nvbench
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvbench.git synced 2026-03-14 20:27:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Replace main_arg_run_benchmarks with run_interriptible

Browse Source 
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.
This commit is contained in:
Oleksandr Pavlyk
2025-12-08 14:28:26 -06:00
parent 8e6154511e
commit 7e9a9a8983
1 changed files with 24 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
python/src/py_nvbench.cpp
View File

@@ -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);