diff --git a/benchmarks/cuda_bindings/runner/main.py b/benchmarks/cuda_bindings/runner/main.py index 9c984c340d6..eb2bcdacaf0 100644 --- a/benchmarks/cuda_bindings/runner/main.py +++ b/benchmarks/cuda_bindings/runner/main.py @@ -232,11 +232,18 @@ def parse_args(argv: list[str], default_output: Path = DEFAULT_OUTPUT) -> tuple[ def main( *, - bench_dir: Path = BENCH_DIR, - default_output: Path = DEFAULT_OUTPUT, + bench_dir: Path | None = None, + default_output: Path | None = None, module_name_prefix: str = DEFAULT_MODULE_NAME_PREFIX, bench_filter_env_var: str = DEFAULT_BENCH_FILTER_ENV_VAR, ) -> None: + # Resolve the defaults inside the call, for the same reason + # discover_benchmarks() does: a literal default would be bound at def-time + # and would ignore a later monkeypatch of the module-level constant. + if bench_dir is None: + bench_dir = BENCH_DIR + if default_output is None: + default_output = DEFAULT_OUTPUT parsed, remaining_argv = parse_args(sys.argv[1:], default_output=default_output) registry = discover_benchmarks(bench_dir=bench_dir, module_name_prefix=module_name_prefix) diff --git a/benchmarks/cuda_bindings/tests/test_runner.py b/benchmarks/cuda_bindings/tests/test_runner.py index 56d88444c9e..836653522a1 100644 --- a/benchmarks/cuda_bindings/tests/test_runner.py +++ b/benchmarks/cuda_bindings/tests/test_runner.py @@ -164,3 +164,24 @@ def test_bench_launch_initializes_on_first_use(monkeypatch): assert len(compile_calls) == 1 assert len(launch_calls) == 2 + + +def test_main_honors_a_monkeypatched_bench_dir(monkeypatch, tmp_path, capsys): + """main() must resolve BENCH_DIR at call time, like discover_benchmarks() does. + + A literal default would be bound at def-time and would silently ignore a + later patch of the module-level constant. + """ + runner_main = load_runner_main(monkeypatch) + + (tmp_path / "bench_patched.py").write_text( + "def bench_only_here(loops: int) -> float:\n return loops + 0.5\n", + encoding="utf-8", + ) + monkeypatch.setattr(runner_main, "BENCH_DIR", tmp_path) + runner_main._MODULE_CACHE.clear() + monkeypatch.setattr(sys, "argv", ["run_pyperf.py", "--list"]) + + runner_main.main() + + assert capsys.readouterr().out.split() == ["patched.only_here"]