Function benchmark_script

bench/python/benchmark.py:912–1132

Main function responsible for running an arbitrary python script and benchmarking it. :param process_idx: The 0-based index of this process. :param device_id: The GPU device id to use. :param output_dir: The output directory to use to store artifacts. :param warmup_batches: The

(
    process_idx,
    device_id,
    output_dir,
    warmup_batches,
    script,
    args,
)

def benchmark_script(
    process_idx,
    device_id,
    output_dir,
    warmup_batches,
    script,
    args,
):
    """
    Main function responsible for running an arbitrary python script and benchmarking it.
    :param process_idx: The 0-based index of this process.
    :param device_id: The GPU device id to use.
    :param output_dir: The output directory to use to store artifacts.
    :param warmup_batches: The numbers of batches that should be ignored from benchmarking.
    :param script: The python script to execute.
    :param args: Any optional command line arguments that should be passed to the script.
    """

    # Make a copy of the environment variables and add our own env-vars to it.
    my_env = os.environ.copy()
    # Change the CUDA visible devices for this process.
    my_env["CUDA_VISIBLE_DEVICES"] = str(device_id)
    # Add the benchmark flag so that perf_utils knows that this is a benchmark run.
    my_env["BENCHMARK_PY"] = "1"

    # Set a path to store the SQLITE report created by NSYS.
    out_sqlite_path = os.path.join(output_dir, "perf_report")
    # Set a path to the benchmark.json created by the script.
    benchmark_json_path = os.path.join(output_dir, "benchmark.json")

    # Remove any existing benchmark.json files.
    if os.path.isfile(benchmark_json_path):
        os.remove(benchmark_json_path)

    # Setup the command that will launch nsys and ask it to benchmark the script
    # that we were interested in.
    nsys_root_path = "/opt/nvidia/nsight-systems/2025.5.1/"
    nsys_binary_path = os.path.join(nsys_root_path, "bin/nsys")
    nsys_reports_path = os.path.join(nsys_root_path, "target-linux-x64/reports")
    nsys_gpu_proj_trace_report_path = os.path.join(
        nsys_reports_path, "nvtx_gpu_proj_trace"
    )
    nsys_pushpop_trace_report_path = os.path.join(
        nsys_reports_path, "nvtx_pushpop_trace"
    )

    if not os.path.isfile(nsys_binary_path):
        raise ValueError(
            "Unable to locate nsys binary at %s. Make sure you have nsight-systems installed."
            % nsys_binary_path
        )

    cmd = [
        nsys_binary_path,
        "profile",
        "--export",
        "sqlite",
        "-o",