From 8104dbaec93dccf74b4711f7ae0d39f9f025fe44 Mon Sep 17 00:00:00 2001 From: LeSingh1 Date: Sat, 8 Aug 2026 18:34:27 -0700 Subject: [PATCH] Fix command line parsing in the cuda.bindings example helpers Both helpers unpack enumerate() backwards: def check_cmd_line_flag(string_ref): return any(string_ref == i and k < len(sys.argv) - 1 for i, k in enumerate(sys.argv)) enumerate() yields (index, value), so `i` is an int and `k` is a str, but the body uses `i` as the argument text and `k` as the index. `string_ref == i` compares str to int and is therefore always False: $ python -c "import sys; from cuda.bindings._example_helpers import *; \ print(check_cmd_line_flag('device='), get_cmd_line_argument_int('device='))" device= 3 False 0 check_cmd_line_flag() always returns False and get_cmd_line_argument_int() always returns 0, so every command line option in the examples is silently ignored: device=, wA=, hA=, wB=, hB=, kernel=, help, ? and use_generic_memory, via helper_cuda.find_cuda_device(), find_cuda_device_drv(), global_to_shmem_async_copy.py, simple_zero_copy.py and stream_ordered_allocation.py. The dead branch would not have worked either: `k < len(sys.argv) - 1` is str < int (TypeError) and `sys.argv[k + 1]` indexes with a str. Alongside the unpacking: - check_cmd_line_flag() no longer requires a following argument. That condition belongs to the value lookup; requiring it would keep `help` and `?` broken whenever they are the last argument, which is the normal way to pass them. - Both helpers skip sys.argv[0], matching the C samples' helper_string.h, which scans from argv[1]. - get_cmd_line_argument_int() returns an int, as its name says and as its callers require: helper_cuda.find_cuda_device() passes the result straight to cudaSetDevice(), and find_cuda_device_drv() to cuDeviceGet(). Returning sys.argv[k + 1] unchanged would hand those APIs a str. This is also the only value the function has ever actually returned, since the literal 0 fallback was the sole reachable path. Adds cuda_bindings/tests/test_example_helpers.py. Five of its assertions fail against main. --- .../_example_helpers/helper_string.py | 17 +++-- cuda_bindings/tests/test_example_helpers.py | 64 +++++++++++++++++++ 2 files changed, 77 insertions(+), 4 deletions(-) create mode 100644 cuda_bindings/tests/test_example_helpers.py diff --git a/cuda_bindings/cuda/bindings/_example_helpers/helper_string.py b/cuda_bindings/cuda/bindings/_example_helpers/helper_string.py index ad15cecc9ca..af53a8b3697 100644 --- a/cuda_bindings/cuda/bindings/_example_helpers/helper_string.py +++ b/cuda_bindings/cuda/bindings/_example_helpers/helper_string.py @@ -5,11 +5,20 @@ def check_cmd_line_flag(string_ref): - return any(string_ref == i and k < len(sys.argv) - 1 for i, k in enumerate(sys.argv)) + """Return whether ``string_ref`` was passed on the command line. + + ``sys.argv[0]`` is the program name and is never considered a flag. + """ + return string_ref in sys.argv[1:] def get_cmd_line_argument_int(string_ref): - for i, k in enumerate(sys.argv): - if string_ref == i and k < len(sys.argv) - 1: - return sys.argv[k + 1] + """Return the integer that follows ``string_ref`` on the command line. + + Returns 0 if ``string_ref`` was not passed, or if nothing follows it. + """ + args = sys.argv[1:] + for idx, arg in enumerate(args): + if arg == string_ref and idx + 1 < len(args): + return int(args[idx + 1]) return 0 diff --git a/cuda_bindings/tests/test_example_helpers.py b/cuda_bindings/tests/test_example_helpers.py new file mode 100644 index 00000000000..c7d50acb65a --- /dev/null +++ b/cuda_bindings/tests/test_example_helpers.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import sys + +import pytest + +from cuda.bindings._example_helpers import check_cmd_line_flag, get_cmd_line_argument_int + + +@pytest.fixture +def argv(monkeypatch): + """Replace sys.argv, keeping a realistic program name in argv[0].""" + + def _argv(*args, prog="example.py"): + monkeypatch.setattr(sys, "argv", [prog, *args]) + + return _argv + + +@pytest.mark.agent_authored(model="claude-opus-5") +@pytest.mark.parametrize( + ("args", "flag", "expected"), + [ + ((), "device=", False), + (("device=", "0"), "device=", True), + (("help",), "help", True), # a boolean flag has nothing after it + (("wA=", "128", "hA=", "256"), "hA=", True), + (("wA=", "128"), "hA=", False), + ], +) +def test_check_cmd_line_flag(argv, args, flag, expected): + argv(*args) + assert check_cmd_line_flag(flag) is expected + + +@pytest.mark.agent_authored(model="claude-opus-5") +def test_check_cmd_line_flag_ignores_the_program_name(argv): + argv(prog="help") + assert check_cmd_line_flag("help") is False + + +@pytest.mark.agent_authored(model="claude-opus-5") +@pytest.mark.parametrize( + ("args", "expected"), + [ + ((), 0), + (("device=", "3"), 3), + (("wA=", "128", "device=", "2"), 2), + (("device=",), 0), # nothing follows the flag + (("nomatch", "7"), 0), + ], +) +def test_get_cmd_line_argument_int(argv, args, expected): + argv(*args) + value = get_cmd_line_argument_int("device=") + assert value == expected + assert isinstance(value, int) + + +@pytest.mark.agent_authored(model="claude-opus-5") +def test_get_cmd_line_argument_int_ignores_the_program_name(argv): + argv("3", prog="device=") + assert get_cmd_line_argument_int("device=") == 0