diff --git a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binary_utility.py b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binary_utility.py index 834db8fe8fa..2c784675cc5 100644 --- a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binary_utility.py +++ b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binary_utility.py @@ -2,13 +2,29 @@ # SPDX-License-Identifier: Apache-2.0 import functools +import importlib import os +from collections.abc import Iterable +from typing import Any from cuda.pathfinder._binaries import supported_nvidia_binaries from cuda.pathfinder._utils.ctk_root_canary import CTK_ROOT_CANARY_ANCHOR_LIBNAMES from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages from cuda.pathfinder._utils.platform_aware import IS_WINDOWS +from cuda.pathfinder._utils.windows_arch import windows_machine_arch + +_NSIGHT_REGISTRY_ROOT = r"SOFTWARE\NVIDIA Corporation\Installed Products\Nsight" + +_NSYS_TARGET_DIR_BY_ARCH = { + "x64": "target-windows-x64", + "arm64": "target-windows-armv8", +} + +_NCU_TARGET_DIR_BY_ARCH = { + "x64": os.path.join("target", "windows-desktop-win7-x64"), + "arm64": os.path.join("target", "windows-desktop-win10-t23x-a64"), +} class UnsupportedBinaryError(Exception): @@ -46,6 +62,84 @@ def _ctk_bin_subdirs(root: str) -> list[str]: return [os.path.join(root, "bin")] +def _resolve_candidate_paths(candidates: Iterable[str]) -> str | None: + """Return the first executable candidate, preserving candidate order.""" + seen: set[str] = set() + for candidate in candidates: + if candidate in seen: + continue + seen.add(candidate) + if _is_executable_candidate(candidate): + return os.path.abspath(candidate) + return None + + +def _find_windows_compute_sanitizer(ctk_root: str) -> str | None: + return _resolve_candidate_paths( + ( + os.path.join(ctk_root, "bin", "compute-sanitizer.bat"), + os.path.join(ctk_root, "compute-sanitizer", "compute-sanitizer.exe"), + ) + ) + + +def _windows_installed_nsight_root(product: str) -> str | None: + """Return the active Nsight product installation recorded by its MSI.""" + # ``winreg`` attributes are absent from the type stubs on non-Windows hosts. + winreg: Any = importlib.import_module("winreg") + + access = winreg.KEY_READ | winreg.KEY_WOW64_64KEY + product_key_path = rf"{_NSIGHT_REGISTRY_ROOT}\{product}" + try: + product_context = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, product_key_path, 0, access) + except FileNotFoundError: + return None + + try: + with product_context as product_key: + current_version, _ = winreg.QueryValueEx(product_key, "CurrentVersion") + if not isinstance(current_version, str) or not current_version.strip(): + raise RuntimeError( + f"Invalid CurrentVersion value {current_version!r} in " + f"Nsight {product!r} registry registration at {product_key_path!r}" + ) + with winreg.OpenKey(product_key, current_version, 0, access) as version_key: + install_root, _ = winreg.QueryValueEx(version_key, None) + except FileNotFoundError as exc: + raise RuntimeError( + f"Incomplete Nsight {product!r} registry registration at {product_key_path!r}" + ) from exc + + if not isinstance(install_root, str) or not install_root.strip(): + raise RuntimeError( + f"Invalid installation directory {install_root!r} in Nsight {product!r} " + f"registry registration at {product_key_path!r} version {current_version!r}" + ) + return install_root + + +def _find_windows_nsys() -> str | None: + install_root = _windows_installed_nsight_root("Systems") + if install_root is None: + return None + + target_dir = _NSYS_TARGET_DIR_BY_ARCH[windows_machine_arch()] + return _resolve_candidate_paths((os.path.join(install_root, target_dir, "nsys.exe"),)) + + +def _find_windows_ncu() -> str | None: + install_root = _windows_installed_nsight_root("Compute") + if install_root is None: + return None + + launcher = os.path.join(install_root, "ncu.bat") + if (found := _resolve_candidate_paths((launcher,))) is not None: + return found + + target_dir = _NCU_TARGET_DIR_BY_ARCH[windows_machine_arch()] + return _resolve_candidate_paths((os.path.join(install_root, target_dir, "ncu.exe"),)) + + def _resolve_ctk_root_via_canary() -> str | None: from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import resolve_ctk_root_via_canary @@ -69,6 +163,20 @@ def _resolve_in_trusted_dirs(normalized_name: str, dirs: list[str]) -> str | Non return None +def _resolve_names_in_trusted_dirs(candidate_names: tuple[str, ...], dirs: list[str]) -> str | None: + """Resolve ordered candidate names within each trusted directory.""" + seen: set[str] = set() + for directory in dirs: + if directory in seen: + continue + assert directory + seen.add(directory) + found = _resolve_candidate_paths(os.path.join(directory, name) for name in candidate_names) + if found is not None: + return found + return None + + @functools.cache def find_nvidia_binary_utility(utility_name: str) -> str | None: """Locate a CUDA binary utility executable. @@ -87,6 +195,19 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None: Raises: UnsupportedBinaryError: If ``utility_name`` is not in the supported set (see ``SUPPORTED_BINARY_UTILITIES``). + RuntimeError: If a native Windows architecture needed for an + architecture-specific utility layout cannot be determined, or an + installed Nsight product has incomplete or invalid registry data. + + Windows on ARM (WoA) Note: + Binary utilities execute in separate processes and do not need to match + the Python process architecture. When choosing among architecture-specific + Windows layouts, this API deliberately targets the native machine + architecture rather than the Python interpreter architecture. For + example, standalone ``nsys`` and ``ncu`` discovery under x64 Python on an + Arm64 machine selects the Arm64 target. This differs from + ``load_nvidia_dynamic_lib`` and ``find_static_lib``, which target the + Python interpreter architecture. Search order: 1. **NVIDIA Python wheels** @@ -100,17 +221,27 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None: environment variable, which use platform-specific bin directory layouts (``Library/bin`` on Windows, ``bin`` on Linux). - 3. **CUDA Toolkit environment variables** + 3. **Library-specific standalone installations** + + - Search the installation paths for the CUDA Toolkit, Nsight Systems, + and Nsight Compute. - - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order), searching - ``bin/x64``, ``bin/x86_64``, and ``bin`` subdirectories on Windows, - or just ``bin`` on Linux. + 3.1. **Nsight installations**: On Windows, locate Nsight Systems and + Nsight Compute from their installer registry entries. Select + architecture-specific binaries using the native machine + architecture, independent of Python. Lookup of the standalone + ``nsys`` and ``ncu`` CLIs is terminal; a miss does not fall + through to CUDA Toolkit locations. + + 3.2. **CUDA Toolkit installation**: Use ``CUDA_HOME`` or ``CUDA_PATH`` + (in that order), searching ``bin/x64``, ``bin/x86_64``, and + ``bin`` subdirectories on Windows, or just ``bin`` on Linux. 4. **CTK-root canary fallback** - - Only when steps 1-3 miss: resolve the ``cudart`` library through the - OS dynamic loader, derive the CUDA Toolkit root from it, and search - that root's bin layout. + - For utilities that reach this step after the earlier searches miss, + resolve the ``cudart`` library through the OS dynamic loader, derive + the CUDA Toolkit root from it, and search that root's bin layout. Note: Results are cached using ``@functools.cache`` for performance. The cache @@ -146,17 +277,35 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None: else: dirs.append(os.path.join(conda_prefix, "bin")) - # 3. Search in CUDA Toolkit (CUDA_HOME/CUDA_PATH) - if (cuda_home := get_cuda_path_or_home()) is not None: - dirs.extend(_ctk_bin_subdirs(cuda_home)) - normalized_name = _normalize_utility_name(utility_name) - found = _resolve_in_trusted_dirs(normalized_name, dirs) + if IS_WINDOWS and utility_name in ("compute-sanitizer", "ncu"): + candidate_names = (f"{utility_name}.bat", normalized_name) + found = _resolve_names_in_trusted_dirs(candidate_names, dirs) + else: + found = _resolve_in_trusted_dirs(normalized_name, dirs) if found is not None: return found + # 3. Search library-specific standalone installations. + # 3.1. Standalone Nsight CLI lookup is terminal; CTK does not contain nsys/ncu. + if IS_WINDOWS and utility_name == "nsys": + return _find_windows_nsys() + if IS_WINDOWS and utility_name == "ncu": + return _find_windows_ncu() + + # 3.2. Search in CUDA Toolkit (CUDA_HOME/CUDA_PATH). + if (cuda_home := get_cuda_path_or_home()) is not None: + if IS_WINDOWS and utility_name == "compute-sanitizer": + found = _find_windows_compute_sanitizer(cuda_home) + else: + found = _resolve_in_trusted_dirs(normalized_name, _ctk_bin_subdirs(cuda_home)) + if found is not None: + return found + # 4. CTK-root canary fallback. ctk_root = _resolve_ctk_root_via_canary() if ctk_root is not None: + if IS_WINDOWS and utility_name == "compute-sanitizer": + return _find_windows_compute_sanitizer(ctk_root) return _resolve_in_trusted_dirs(normalized_name, _ctk_bin_subdirs(ctk_root)) return None diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py index 53446107da3..61bf31720d3 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py @@ -231,6 +231,15 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL: DynamicLibNotFoundError: If the library cannot be found or loaded. RuntimeError: If Python is not 64-bit. + Windows on ARM (WoA) Note: + On Windows, this API aims to load a dynamic library whose architecture + matches the Python interpreter architecture. For example, x64 Python + running on an Arm64 machine targets an x64 DLL, while native Arm64 Python + targets an Arm64 DLL. A library loaded into the Python process must be + compatible with that process. This differs from + ``find_nvidia_binary_utility``, which targets the native machine + architecture when selecting architecture-specific executables. + Search order: 0. **Already loaded in the current process** diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py index ea5a740aec4..8e0d81cab01 100644 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py @@ -175,5 +175,13 @@ def find_static_lib(name: str) -> str: Raises: ValueError: If ``name`` is not a supported static library. StaticLibNotFoundError: If the static library cannot be found. + + Windows on ARM (WoA) Note: + On Windows, this API aims to return the path to a static library whose + architecture matches the Python interpreter architecture. For example, + x64 Python running on an Arm64 machine targets the x64 library, while + native Arm64 Python targets the Arm64 library. This differs from + ``find_nvidia_binary_utility``, which targets the native machine + architecture when selecting architecture-specific executables. """ return locate_static_lib(name).abs_path diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/windows_arch.py b/cuda_pathfinder/cuda/pathfinder/_utils/windows_arch.py index 9313f3a9f17..96eb53a4e69 100644 --- a/cuda_pathfinder/cuda/pathfinder/_utils/windows_arch.py +++ b/cuda_pathfinder/cuda/pathfinder/_utils/windows_arch.py @@ -3,6 +3,7 @@ from __future__ import annotations +import platform import sysconfig WINDOWS_PE_MACHINE_BY_ARCH = { @@ -10,6 +11,8 @@ "arm64": 0xAA64, } +_WINDOWS_ARCH_BY_PE_MACHINE = {machine: arch for arch, machine in WINDOWS_PE_MACHINE_BY_ARCH.items()} + class UnsupportedArchError(RuntimeError): """Raised when Python reports an unsupported Windows architecture.""" @@ -35,6 +38,76 @@ def windows_python_arch() -> str: raise UnsupportedArchError(raw_platform_tag) +def _windows_machine_arch_from_platform() -> str: + """Return the Windows architecture reported by Python's platform module.""" + raw_machine = platform.machine() + machine = raw_machine.lower().replace("_", "-") + + if machine in ("amd64", "x86-64"): + return "x64" + + if machine in ("arm64", "aarch64"): + return "arm64" + + raise RuntimeError(f"Unsupported Windows machine architecture: {raw_machine!r}") + + +def _windows_native_machine() -> int | None: + """Return the native Windows PE machine type, or None on older Windows.""" + import ctypes + from ctypes import wintypes + + try: + # These ctypes attributes are absent from the type stubs on non-Windows hosts. + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) # type: ignore[attr-defined] + except OSError as exc: + raise RuntimeError("Failed to load kernel32 while detecting the native Windows architecture") from exc + + get_current_process = kernel32.GetCurrentProcess + try: + is_wow64_process2 = kernel32.IsWow64Process2 + except AttributeError: + return None + + get_current_process.argtypes = () + get_current_process.restype = wintypes.HANDLE + is_wow64_process2.argtypes = ( + wintypes.HANDLE, + ctypes.POINTER(wintypes.USHORT), + ctypes.POINTER(wintypes.USHORT), + ) + is_wow64_process2.restype = wintypes.BOOL + + process_machine = wintypes.USHORT() + native_machine = wintypes.USHORT() + if not is_wow64_process2( + get_current_process(), + ctypes.byref(process_machine), + ctypes.byref(native_machine), + ): + error_code = ctypes.get_last_error() # type: ignore[attr-defined] + error = ctypes.WinError(error_code) # type: ignore[attr-defined] + raise RuntimeError( + f"IsWow64Process2 failed while detecting the native Windows architecture " + f"(Windows error {error_code}): {error}" + ) from error + return native_machine.value + + +def windows_machine_arch() -> str: + """Return the native Windows machine architecture, ignoring process emulation.""" + native_machine = _windows_native_machine() + if native_machine is None: + # IsWow64Process2 predates x64-on-Arm emulation, so this fallback is only + # needed on older Windows versions where platform.machine() is sufficient. + return _windows_machine_arch_from_platform() + + try: + return _WINDOWS_ARCH_BY_PE_MACHINE[native_machine] + except KeyError: + raise RuntimeError(f"Unsupported native Windows PE machine type: 0x{native_machine:04x}") from None + + def windows_pe_matches_arch(path: str, target_arch: str) -> bool: """Return whether a Windows Portable Executable (PE) targets the requested architecture. diff --git a/cuda_pathfinder/docs/source/install.rst b/cuda_pathfinder/docs/source/install.rst index 53f11ebbf18..078abf47ee3 100644 --- a/cuda_pathfinder/docs/source/install.rst +++ b/cuda_pathfinder/docs/source/install.rst @@ -9,7 +9,7 @@ Runtime Requirements ``cuda.pathfinder`` is a pure-Python package with no runtime dependencies: -* Linux (x86-64, arm64) and Windows (x86-64) +* Linux (x86-64, arm64) and Windows (x86-64, arm64) * Python 3.10 - 3.14 Installing from PyPI diff --git a/cuda_pathfinder/tests/test_find_nvidia_binaries.py b/cuda_pathfinder/tests/test_find_nvidia_binaries.py index 2784633ff38..668be893863 100644 --- a/cuda_pathfinder/tests/test_find_nvidia_binaries.py +++ b/cuda_pathfinder/tests/test_find_nvidia_binaries.py @@ -58,6 +58,15 @@ def fake_is_executable_candidate(path): return checked +def _patch_winreg(mocker): + winreg = mocker.MagicMock() + winreg.HKEY_LOCAL_MACHINE = object() + winreg.KEY_READ = 0x20019 + winreg.KEY_WOW64_64KEY = 0x0100 + mocker.patch.object(binary_finder_module.importlib, "import_module", return_value=winreg) + return winreg + + @pytest.mark.usefixtures("clear_find_binary_cache") def test_find_binary_search_path_includes_site_packages_conda_cuda(monkeypatch, mocker): conda_prefix = os.path.join(os.sep, "conda") @@ -128,6 +137,415 @@ def test_find_binary_windows_extension_and_search_dirs(monkeypatch, mocker): assert checked == [os.path.join(d, "nvcc.exe") for d in expected_dirs] +@pytest.mark.parametrize( + ("launcher_exists", "expected_rel", "checked_rels"), + ( + (True, os.path.join("bin", "compute-sanitizer.bat"), (os.path.join("bin", "compute-sanitizer.bat"),)), + ( + False, + os.path.join("compute-sanitizer", "compute-sanitizer.exe"), + ( + os.path.join("bin", "compute-sanitizer.bat"), + os.path.join("compute-sanitizer", "compute-sanitizer.exe"), + ), + ), + ), +) +@pytest.mark.usefixtures("clear_find_binary_cache") +@pytest.mark.agent_authored(model="gpt-5.6") +def test_find_compute_sanitizer_prefers_ctk_launcher_with_executable_fallback( + monkeypatch, mocker, launcher_exists, expected_rel, checked_rels +): + cuda_home = os.path.join(os.sep, "cuda") + launcher = os.path.join(cuda_home, "bin", "compute-sanitizer.bat") + executable = os.path.join(cuda_home, "compute-sanitizer", "compute-sanitizer.exe") + + mocker.patch.object(binary_finder_module, "IS_WINDOWS", new=True) + mocker.patch.object(binary_finder_module.supported_nvidia_binaries, "SITE_PACKAGES_BINDIRS", {}) + monkeypatch.delenv("CONDA_PREFIX", raising=False) + mocker.patch.object(binary_finder_module, "get_cuda_path_or_home", return_value=cuda_home) + canary_mock = mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary") + existing = [executable] + if launcher_exists: + existing.append(launcher) + checked = _patch_exec_probe(mocker, existing=existing) + + assert find_nvidia_binary_utility("compute-sanitizer") == os.path.abspath(os.path.join(cuda_home, expected_rel)) + assert checked == [os.path.join(cuda_home, rel) for rel in checked_rels] + canary_mock.assert_not_called() + + +@pytest.mark.usefixtures("clear_find_binary_cache") +@pytest.mark.agent_authored(model="gpt-5.6") +def test_find_compute_sanitizer_uses_canary_ctk_root(monkeypatch, mocker): + ctk_root = os.path.join(os.sep, "cuda") + launcher = os.path.join(ctk_root, "bin", "compute-sanitizer.bat") + + mocker.patch.object(binary_finder_module, "IS_WINDOWS", new=True) + mocker.patch.object(binary_finder_module.supported_nvidia_binaries, "SITE_PACKAGES_BINDIRS", {}) + monkeypatch.delenv("CONDA_PREFIX", raising=False) + mocker.patch.object(binary_finder_module, "get_cuda_path_or_home", return_value=None) + canary = mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary", return_value=ctk_root) + checked = _patch_exec_probe(mocker, existing=[launcher]) + + assert find_nvidia_binary_utility("compute-sanitizer") == os.path.abspath(launcher) + assert checked == [launcher] + canary.assert_called_once_with() + + +@pytest.mark.parametrize( + ("machine_arch", "target_dir"), + ( + ("x64", "target-windows-x64"), + ("arm64", "target-windows-armv8"), + ), +) +@pytest.mark.agent_authored(model="gpt-5.6") +def test_find_windows_nsys_uses_machine_arch(mocker, machine_arch, target_dir): + install_root = os.path.join(os.sep, "Program Files", "Nsight Systems") + expected = os.path.join(install_root, target_dir, "nsys.exe") + mocker.patch.object(binary_finder_module, "_windows_installed_nsight_root", return_value=install_root) + mocker.patch.object(binary_finder_module, "windows_machine_arch", return_value=machine_arch) + checked = _patch_exec_probe(mocker, existing=[expected]) + + assert binary_finder_module._find_windows_nsys() == os.path.abspath(expected) + assert checked == [expected] + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_find_windows_nsys_does_not_fallback_to_other_arch(mocker): + install_root = os.path.join(os.sep, "Program Files", "Nsight Systems") + arm64 = os.path.join(install_root, "target-windows-armv8", "nsys.exe") + x64 = os.path.join(install_root, "target-windows-x64", "nsys.exe") + mocker.patch.object(binary_finder_module, "_windows_installed_nsight_root", return_value=install_root) + mocker.patch.object(binary_finder_module, "windows_machine_arch", return_value="arm64") + checked = _patch_exec_probe(mocker, existing=[x64]) + + assert binary_finder_module._find_windows_nsys() is None + assert checked == [arm64] + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_find_windows_ncu_prefers_launcher(mocker): + install_root = os.path.join(os.sep, "Program Files", "Nsight Compute") + launcher = os.path.join(install_root, "ncu.bat") + mocker.patch.object(binary_finder_module, "_windows_installed_nsight_root", return_value=install_root) + machine_arch_mock = mocker.patch.object(binary_finder_module, "windows_machine_arch") + checked = _patch_exec_probe(mocker, existing=[launcher]) + + assert binary_finder_module._find_windows_ncu() == os.path.abspath(launcher) + assert checked == [launcher] + machine_arch_mock.assert_not_called() + + +@pytest.mark.parametrize( + ("machine_arch", "target_dir"), + ( + ("x64", os.path.join("target", "windows-desktop-win7-x64")), + ("arm64", os.path.join("target", "windows-desktop-win10-t23x-a64")), + ), +) +@pytest.mark.agent_authored(model="gpt-5.6") +def test_find_windows_ncu_falls_back_to_machine_binary(mocker, machine_arch, target_dir): + install_root = os.path.join(os.sep, "Program Files", "Nsight Compute") + launcher = os.path.join(install_root, "ncu.bat") + expected = os.path.join(install_root, target_dir, "ncu.exe") + mocker.patch.object(binary_finder_module, "_windows_installed_nsight_root", return_value=install_root) + mocker.patch.object(binary_finder_module, "windows_machine_arch", return_value=machine_arch) + checked = _patch_exec_probe(mocker, existing=[expected]) + + assert binary_finder_module._find_windows_ncu() == os.path.abspath(expected) + assert checked == [launcher, expected] + + +@pytest.mark.parametrize( + ("utility_name", "candidate_names"), + ( + ("nsys", ("nsys.exe",)), + ("ncu", ("ncu.bat", "ncu.exe")), + ), +) +@pytest.mark.usefixtures("clear_find_binary_cache") +@pytest.mark.agent_authored(model="gpt-5.6") +def test_find_binary_windows_nsight_conda_precedes_registry(monkeypatch, mocker, utility_name, candidate_names): + site_dir = os.path.join(os.sep, "site-packages", utility_name, "bin") + conda_prefix = os.path.join(os.sep, "conda") + conda_bin = os.path.join(conda_prefix, "Library", "bin") + expected = os.path.join(conda_bin, candidate_names[0]) + + mocker.patch.object(binary_finder_module, "IS_WINDOWS", new=True) + mocker.patch.object(binary_finder_module, "find_sub_dirs_all_sitepackages", return_value=[site_dir]) + monkeypatch.setenv("CONDA_PREFIX", conda_prefix) + registry_root = mocker.patch.object(binary_finder_module, "_windows_installed_nsight_root") + machine_arch = mocker.patch.object(binary_finder_module, "windows_machine_arch") + get_cuda_path = mocker.patch.object(binary_finder_module, "get_cuda_path_or_home") + canary = mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary") + checked = _patch_exec_probe(mocker, existing=[expected]) + + assert find_nvidia_binary_utility(utility_name) == os.path.abspath(expected) + assert checked == [ + *(os.path.join(site_dir, name) for name in candidate_names), + os.path.join(conda_bin, candidate_names[0]), + ] + registry_root.assert_not_called() + machine_arch.assert_not_called() + get_cuda_path.assert_not_called() + canary.assert_not_called() + + +@pytest.mark.parametrize( + ("utility_name", "product", "machine_arch", "target_rel", "candidate_names"), + ( + ("nsys", "Systems", "x64", os.path.join("target-windows-x64", "nsys.exe"), ("nsys.exe",)), + ("nsys", "Systems", "arm64", os.path.join("target-windows-armv8", "nsys.exe"), ("nsys.exe",)), + ( + "ncu", + "Compute", + "x64", + os.path.join("target", "windows-desktop-win7-x64", "ncu.exe"), + ("ncu.bat", "ncu.exe"), + ), + ( + "ncu", + "Compute", + "arm64", + os.path.join("target", "windows-desktop-win10-t23x-a64", "ncu.exe"), + ("ncu.bat", "ncu.exe"), + ), + ), +) +@pytest.mark.usefixtures("clear_find_binary_cache") +@pytest.mark.agent_authored(model="gpt-5.6") +def test_find_binary_windows_nsight_composes_registry_and_native_target( + monkeypatch, mocker, utility_name, product, machine_arch, target_rel, candidate_names +): + site_dir = os.path.join(os.sep, "site-packages", utility_name, "bin") + conda_prefix = os.path.join(os.sep, "conda") + conda_bin = os.path.join(conda_prefix, "Library", "bin") + install_root = os.path.join(os.sep, "Program Files", utility_name) + expected = os.path.join(install_root, target_rel) + + mocker.patch.object(binary_finder_module, "IS_WINDOWS", new=True) + mocker.patch.object(binary_finder_module, "find_sub_dirs_all_sitepackages", return_value=[site_dir]) + monkeypatch.setenv("CONDA_PREFIX", conda_prefix) + registry_root = mocker.patch.object( + binary_finder_module, "_windows_installed_nsight_root", return_value=install_root + ) + machine_arch_mock = mocker.patch.object(binary_finder_module, "windows_machine_arch", return_value=machine_arch) + get_cuda_path = mocker.patch.object(binary_finder_module, "get_cuda_path_or_home") + canary = mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary") + checked = _patch_exec_probe(mocker, existing=[expected]) + + assert find_nvidia_binary_utility(utility_name) == os.path.abspath(expected) + assert checked == [ + *(os.path.join(directory, name) for directory in (site_dir, conda_bin) for name in candidate_names), + *((os.path.join(install_root, "ncu.bat"),) if utility_name == "ncu" else ()), + expected, + ] + registry_root.assert_called_once_with(product) + machine_arch_mock.assert_called_once_with() + get_cuda_path.assert_not_called() + canary.assert_not_called() + + +@pytest.mark.parametrize(("utility_name", "product"), (("nsys", "Systems"), ("ncu", "Compute"))) +@pytest.mark.usefixtures("clear_find_binary_cache") +@pytest.mark.agent_authored(model="gpt-5.6") +def test_find_binary_windows_nsight_registry_miss_is_terminal(monkeypatch, mocker, utility_name, product): + mocker.patch.object(binary_finder_module, "IS_WINDOWS", new=True) + mocker.patch.object(binary_finder_module.supported_nvidia_binaries, "SITE_PACKAGES_BINDIRS", {}) + monkeypatch.delenv("CONDA_PREFIX", raising=False) + registry_root = mocker.patch.object(binary_finder_module, "_windows_installed_nsight_root", return_value=None) + machine_arch = mocker.patch.object(binary_finder_module, "windows_machine_arch") + get_cuda_path = mocker.patch.object(binary_finder_module, "get_cuda_path_or_home") + canary = mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary") + + assert find_nvidia_binary_utility(utility_name) is None + registry_root.assert_called_once_with(product) + machine_arch.assert_not_called() + get_cuda_path.assert_not_called() + canary.assert_not_called() + + +@pytest.mark.parametrize("utility_name", ("nsight-sys", "nsight-compute")) +@pytest.mark.usefixtures("clear_find_binary_cache") +@pytest.mark.agent_authored(model="gpt-5.6") +def test_find_windows_nsight_legacy_names_remain_literal_in_early_search(monkeypatch, mocker, utility_name): + site_key = os.path.join("nvidia", utility_name, "bin") + site_dir = os.path.join(os.sep, "site-packages", utility_name, "bin") + conda_prefix = os.path.join(os.sep, "conda") + conda_bin = os.path.join(conda_prefix, "Library", "bin") + expected = os.path.join(conda_bin, f"{utility_name}.exe") + + mocker.patch.object(binary_finder_module, "IS_WINDOWS", new=True) + mocker.patch.object( + binary_finder_module.supported_nvidia_binaries, + "SITE_PACKAGES_BINDIRS", + {utility_name: (site_key,)}, + ) + find_sub_dirs = mocker.patch.object(binary_finder_module, "find_sub_dirs_all_sitepackages", return_value=[site_dir]) + monkeypatch.setenv("CONDA_PREFIX", conda_prefix) + get_cuda_path = mocker.patch.object(binary_finder_module, "get_cuda_path_or_home") + nsys_finder = mocker.patch.object(binary_finder_module, "_find_windows_nsys") + ncu_finder = mocker.patch.object(binary_finder_module, "_find_windows_ncu") + checked = _patch_exec_probe(mocker, existing=[expected]) + + assert find_nvidia_binary_utility(utility_name) == os.path.abspath(expected) + assert checked == [os.path.join(site_dir, f"{utility_name}.exe"), expected] + find_sub_dirs.assert_called_once_with(site_key.split(os.sep)) + get_cuda_path.assert_not_called() + nsys_finder.assert_not_called() + ncu_finder.assert_not_called() + + +@pytest.mark.parametrize("utility_name", ("nsight-sys", "nsight-compute")) +@pytest.mark.usefixtures("clear_find_binary_cache") +@pytest.mark.agent_authored(model="gpt-5.6") +def test_find_windows_nsight_legacy_names_remain_literal_in_ctk(monkeypatch, mocker, utility_name): + cuda_home = os.path.join(os.sep, "cuda") + expected = os.path.join(cuda_home, "bin", f"{utility_name}.exe") + + mocker.patch.object(binary_finder_module, "IS_WINDOWS", new=True) + mocker.patch.object(binary_finder_module.supported_nvidia_binaries, "SITE_PACKAGES_BINDIRS", {}) + monkeypatch.delenv("CONDA_PREFIX", raising=False) + mocker.patch.object(binary_finder_module, "get_cuda_path_or_home", return_value=cuda_home) + nsys_finder = mocker.patch.object(binary_finder_module, "_find_windows_nsys") + ncu_finder = mocker.patch.object(binary_finder_module, "_find_windows_ncu") + canary = mocker.patch.object(binary_finder_module, "_resolve_ctk_root_via_canary") + checked = _patch_exec_probe(mocker, existing=[expected]) + + assert find_nvidia_binary_utility(utility_name) == os.path.abspath(expected) + assert checked == [ + os.path.join(cuda_home, "bin", "x64", f"{utility_name}.exe"), + os.path.join(cuda_home, "bin", "x86_64", f"{utility_name}.exe"), + expected, + ] + nsys_finder.assert_not_called() + ncu_finder.assert_not_called() + canary.assert_not_called() + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_windows_installed_nsight_root_reads_64_bit_registry(mocker): + install_root = os.path.join(os.sep, "Program Files", "Nsight Systems") + product_key = mocker.MagicMock() + version_key = mocker.MagicMock() + product_context = mocker.MagicMock() + product_context.__enter__.return_value = product_key + version_context = mocker.MagicMock() + version_context.__enter__.return_value = version_key + winreg = _patch_winreg(mocker) + winreg.OpenKey.side_effect = (product_context, version_context) + winreg.QueryValueEx.side_effect = (("2026.1.3", 1), (install_root, 1)) + + assert binary_finder_module._windows_installed_nsight_root("Systems") == install_root + access = winreg.KEY_READ | winreg.KEY_WOW64_64KEY + winreg.OpenKey.assert_has_calls( + ( + mocker.call( + winreg.HKEY_LOCAL_MACHINE, + rf"{binary_finder_module._NSIGHT_REGISTRY_ROOT}\Systems", + 0, + access, + ), + mocker.call(product_key, "2026.1.3", 0, access), + ) + ) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_windows_installed_nsight_root_returns_none_when_product_key_is_absent(mocker): + winreg = _patch_winreg(mocker) + winreg.OpenKey.side_effect = FileNotFoundError("Nsight Systems is not installed") + + assert binary_finder_module._windows_installed_nsight_root("Systems") is None + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_windows_installed_nsight_root_rejects_missing_current_version(mocker): + product_context = mocker.MagicMock() + product_context.__enter__.return_value = mocker.MagicMock() + winreg = _patch_winreg(mocker) + winreg.OpenKey.return_value = product_context + winreg.QueryValueEx.side_effect = FileNotFoundError("CurrentVersion is missing") + + with pytest.raises(RuntimeError, match=r"Incomplete Nsight 'Systems' registry registration") as exc_info: + binary_finder_module._windows_installed_nsight_root("Systems") + + assert isinstance(exc_info.value.__cause__, FileNotFoundError) + + +@pytest.mark.parametrize("current_version", (None, "", " ", 2026)) +@pytest.mark.agent_authored(model="gpt-5.6") +def test_windows_installed_nsight_root_rejects_invalid_current_version(mocker, current_version): + product_context = mocker.MagicMock() + product_context.__enter__.return_value = mocker.MagicMock() + winreg = _patch_winreg(mocker) + winreg.OpenKey.return_value = product_context + winreg.QueryValueEx.return_value = (current_version, 1) + + with pytest.raises(RuntimeError, match=r"Invalid CurrentVersion value .*Nsight 'Systems' registry registration"): + binary_finder_module._windows_installed_nsight_root("Systems") + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_windows_installed_nsight_root_rejects_missing_version_key(mocker): + product_key = mocker.MagicMock() + product_context = mocker.MagicMock() + product_context.__enter__.return_value = product_key + winreg = _patch_winreg(mocker) + winreg.OpenKey.side_effect = (product_context, FileNotFoundError("Version key is missing")) + winreg.QueryValueEx.return_value = ("2026.1.3", 1) + + with pytest.raises(RuntimeError, match=r"Incomplete Nsight 'Systems' registry registration") as exc_info: + binary_finder_module._windows_installed_nsight_root("Systems") + + assert isinstance(exc_info.value.__cause__, FileNotFoundError) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_windows_installed_nsight_root_rejects_missing_installation_directory(mocker): + product_context = mocker.MagicMock() + product_context.__enter__.return_value = mocker.MagicMock() + version_context = mocker.MagicMock() + version_context.__enter__.return_value = mocker.MagicMock() + winreg = _patch_winreg(mocker) + winreg.OpenKey.side_effect = (product_context, version_context) + winreg.QueryValueEx.side_effect = (("2026.1.3", 1), FileNotFoundError("Installation directory is missing")) + + with pytest.raises(RuntimeError, match=r"Incomplete Nsight 'Systems' registry registration") as exc_info: + binary_finder_module._windows_installed_nsight_root("Systems") + + assert isinstance(exc_info.value.__cause__, FileNotFoundError) + + +@pytest.mark.parametrize("install_root", (None, "", " ", 2026)) +@pytest.mark.agent_authored(model="gpt-5.6") +def test_windows_installed_nsight_root_rejects_invalid_installation_directory(mocker, install_root): + product_context = mocker.MagicMock() + product_context.__enter__.return_value = mocker.MagicMock() + version_context = mocker.MagicMock() + version_context.__enter__.return_value = mocker.MagicMock() + winreg = _patch_winreg(mocker) + winreg.OpenKey.side_effect = (product_context, version_context) + winreg.QueryValueEx.side_effect = (("2026.1.3", 1), (install_root, 1)) + + with pytest.raises( + RuntimeError, + match=r"Invalid installation directory .*Nsight 'Systems' registry registration.*version '2026.1.3'", + ): + binary_finder_module._windows_installed_nsight_root("Systems") + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_windows_installed_nsight_root_propagates_access_errors(mocker): + winreg = _patch_winreg(mocker) + winreg.OpenKey.side_effect = PermissionError("Registry access denied") + + with pytest.raises(PermissionError, match="Registry access denied"): + binary_finder_module._windows_installed_nsight_root("Systems") + + @pytest.mark.usefixtures("clear_find_binary_cache") def test_find_binary_first_matching_dir_wins(monkeypatch, mocker): conda_prefix = os.path.join(os.sep, "conda") diff --git a/cuda_pathfinder/tests/test_search_steps.py b/cuda_pathfinder/tests/test_search_steps.py index 54136dc34e1..fc78e22c708 100644 --- a/cuda_pathfinder/tests/test_search_steps.py +++ b/cuda_pathfinder/tests/test_search_steps.py @@ -5,7 +5,9 @@ from __future__ import annotations +import ctypes import os +from ctypes import wintypes import pytest @@ -147,6 +149,85 @@ def test_rejects_unknown_sysconfig_tag(self, mocker): assert exc_info.value.platform_tag == "custom-win" +class TestWindowsMachineArch: + @pytest.mark.parametrize( + ("native_machine", "expected"), + ((0x8664, "x64"), (0xAA64, "arm64")), + ) + @pytest.mark.agent_authored(model="gpt-5.6") + def test_uses_native_pe_machine(self, mocker, native_machine, expected): + mocker.patch.object(windows_arch_mod, "_windows_native_machine", return_value=native_machine) + platform_machine = mocker.patch.object(windows_arch_mod.platform, "machine", return_value="AMD64") + + assert windows_arch_mod.windows_machine_arch() == expected + platform_machine.assert_not_called() + + @pytest.mark.agent_authored(model="gpt-5.6") + def test_rejects_unknown_native_pe_machine(self, mocker): + mocker.patch.object(windows_arch_mod, "_windows_native_machine", return_value=0x014C) + + with pytest.raises(RuntimeError, match=r"Unsupported native Windows PE machine type: 0x014c"): + windows_arch_mod.windows_machine_arch() + + @pytest.mark.parametrize( + ("reported_machine", "expected"), + (("AMD64", "x64"), ("x86_64", "x64"), ("ARM64", "arm64"), ("aarch64", "arm64")), + ) + @pytest.mark.agent_authored(model="gpt-5.6") + def test_old_windows_fallback_normalizes_platform_machine(self, mocker, reported_machine, expected): + mocker.patch.object(windows_arch_mod, "_windows_native_machine", return_value=None) + mocker.patch.object(windows_arch_mod.platform, "machine", return_value=reported_machine) + + assert windows_arch_mod.windows_machine_arch() == expected + + @pytest.mark.agent_authored(model="gpt-5.6") + def test_native_machine_returns_none_when_is_wow64_process2_is_unavailable(self, mocker): + kernel32 = mocker.Mock(spec=["GetCurrentProcess"]) + mocker.patch.object(ctypes, "WinDLL", create=True, return_value=kernel32) + + assert windows_arch_mod._windows_native_machine() is None + + @pytest.mark.agent_authored(model="gpt-5.6") + def test_native_machine_configures_api_and_returns_native_machine(self, mocker): + kernel32 = mocker.Mock() + kernel32.GetCurrentProcess.return_value = wintypes.HANDLE(1) + + def report_native_machine(_process, _process_machine, native_machine): + native_machine._obj.value = 0xAA64 + return True + + kernel32.IsWow64Process2.side_effect = report_native_machine + mocker.patch.object(ctypes, "WinDLL", create=True, return_value=kernel32) + + assert windows_arch_mod._windows_native_machine() == 0xAA64 + assert kernel32.GetCurrentProcess.argtypes == () + assert kernel32.GetCurrentProcess.restype is wintypes.HANDLE + assert kernel32.IsWow64Process2.argtypes == ( + wintypes.HANDLE, + ctypes.POINTER(wintypes.USHORT), + ctypes.POINTER(wintypes.USHORT), + ) + assert kernel32.IsWow64Process2.restype is wintypes.BOOL + + @pytest.mark.agent_authored(model="gpt-5.6") + def test_native_machine_raises_contextual_error_when_api_call_fails(self, mocker): + kernel32 = mocker.Mock() + kernel32.GetCurrentProcess.return_value = wintypes.HANDLE(1) + kernel32.IsWow64Process2.return_value = False + mocker.patch.object(ctypes, "WinDLL", create=True, return_value=kernel32) + mocker.patch.object(ctypes, "get_last_error", create=True, return_value=87) + windows_error = OSError(87, "The parameter is incorrect") + mocker.patch.object(ctypes, "WinError", create=True, return_value=windows_error) + + with pytest.raises( + RuntimeError, + match=r"IsWow64Process2 failed while detecting the native Windows architecture \(Windows error 87\)", + ) as exc_info: + windows_arch_mod._windows_native_machine() + + assert exc_info.value.__cause__ is windows_error + + @pytest.mark.parametrize( ("machine", "target_arch", "expected"), (