Skip to content

Commit 8104dba

Browse files
committed
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.
1 parent 3bd069a commit 8104dba

2 files changed

Lines changed: 77 additions & 4 deletions

File tree

cuda_bindings/cuda/bindings/_example_helpers/helper_string.py

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,20 @@
55

66

77
def check_cmd_line_flag(string_ref):
8-
return any(string_ref == i and k < len(sys.argv) - 1 for i, k in enumerate(sys.argv))
8+
"""Return whether ``string_ref`` was passed on the command line.
9+
10+
``sys.argv[0]`` is the program name and is never considered a flag.
11+
"""
12+
return string_ref in sys.argv[1:]
913

1014

1115
def get_cmd_line_argument_int(string_ref):
12-
for i, k in enumerate(sys.argv):
13-
if string_ref == i and k < len(sys.argv) - 1:
14-
return sys.argv[k + 1]
16+
"""Return the integer that follows ``string_ref`` on the command line.
17+
18+
Returns 0 if ``string_ref`` was not passed, or if nothing follows it.
19+
"""
20+
args = sys.argv[1:]
21+
for idx, arg in enumerate(args):
22+
if arg == string_ref and idx + 1 < len(args):
23+
return int(args[idx + 1])
1524
return 0
Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
import sys
5+
6+
import pytest
7+
8+
from cuda.bindings._example_helpers import check_cmd_line_flag, get_cmd_line_argument_int
9+
10+
11+
@pytest.fixture
12+
def argv(monkeypatch):
13+
"""Replace sys.argv, keeping a realistic program name in argv[0]."""
14+
15+
def _argv(*args, prog="example.py"):
16+
monkeypatch.setattr(sys, "argv", [prog, *args])
17+
18+
return _argv
19+
20+
21+
@pytest.mark.agent_authored(model="claude-opus-5")
22+
@pytest.mark.parametrize(
23+
("args", "flag", "expected"),
24+
[
25+
((), "device=", False),
26+
(("device=", "0"), "device=", True),
27+
(("help",), "help", True), # a boolean flag has nothing after it
28+
(("wA=", "128", "hA=", "256"), "hA=", True),
29+
(("wA=", "128"), "hA=", False),
30+
],
31+
)
32+
def test_check_cmd_line_flag(argv, args, flag, expected):
33+
argv(*args)
34+
assert check_cmd_line_flag(flag) is expected
35+
36+
37+
@pytest.mark.agent_authored(model="claude-opus-5")
38+
def test_check_cmd_line_flag_ignores_the_program_name(argv):
39+
argv(prog="help")
40+
assert check_cmd_line_flag("help") is False
41+
42+
43+
@pytest.mark.agent_authored(model="claude-opus-5")
44+
@pytest.mark.parametrize(
45+
("args", "expected"),
46+
[
47+
((), 0),
48+
(("device=", "3"), 3),
49+
(("wA=", "128", "device=", "2"), 2),
50+
(("device=",), 0), # nothing follows the flag
51+
(("nomatch", "7"), 0),
52+
],
53+
)
54+
def test_get_cmd_line_argument_int(argv, args, expected):
55+
argv(*args)
56+
value = get_cmd_line_argument_int("device=")
57+
assert value == expected
58+
assert isinstance(value, int)
59+
60+
61+
@pytest.mark.agent_authored(model="claude-opus-5")
62+
def test_get_cmd_line_argument_int_ignores_the_program_name(argv):
63+
argv("3", prog="device=")
64+
assert get_cmd_line_argument_int("device=") == 0

0 commit comments

Comments
 (0)