diff --git a/triton_kernel_agent/opt_worker_component/benchmarking/benchmark.py b/triton_kernel_agent/opt_worker_component/benchmarking/benchmark.py index 9f8314ac..46de3b34 100644 --- a/triton_kernel_agent/opt_worker_component/benchmarking/benchmark.py +++ b/triton_kernel_agent/opt_worker_component/benchmarking/benchmark.py @@ -20,6 +20,7 @@ import json import logging +import os import subprocess import sys import traceback @@ -129,8 +130,14 @@ def benchmark_kernel( results_json = self.artifacts_dir / "benchmark_results.json" benchmark_script = Path(__file__).parent / "kernel_subprocess.py" + # Use KERNEL_PROFILER_PYTHON (the PAR bootstrap) when set, like + # ncu_profiler.py; bare sys.executable is un-bootstrapped in a PAR. + bench_python = ( + os.environ.get("KERNEL_PROFILER_PYTHON") or sys.executable + ) + cmd = [ - sys.executable, + bench_python, str(benchmark_script), "--problem", str(problem_file),