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..fa28347a399 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,71 @@ 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: + with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, product_key_path, 0, access) as product_key: + current_version, _ = winreg.QueryValueEx(product_key, "CurrentVersion") + if not isinstance(current_version, str) or not current_version: + raise RuntimeError(f"Invalid CurrentVersion value in {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: + return None + + if not isinstance(install_root, str) or not install_root: + raise RuntimeError(f"Invalid installation directory for {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 +150,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. @@ -88,6 +183,16 @@ def find_nvidia_binary_utility(utility_name: str) -> str | None: UnsupportedBinaryError: If ``utility_name`` is not in the supported set (see ``SUPPORTED_BINARY_UTILITIES``). + 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 +205,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 +261,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..d7b6264d16b 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 = { @@ -35,6 +36,20 @@ def windows_python_arch() -> str: raise UnsupportedArchError(raw_platform_tag) +def windows_machine_arch() -> str: + """Return the native Windows machine architecture.""" + 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_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/tests/test_find_nvidia_binaries.py b/cuda_pathfinder/tests/test_find_nvidia_binaries.py index 2784633ff38..f2c424088ad 100644 --- a/cuda_pathfinder/tests/test_find_nvidia_binaries.py +++ b/cuda_pathfinder/tests/test_find_nvidia_binaries.py @@ -128,6 +128,326 @@ 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 = mocker.MagicMock() + winreg.HKEY_LOCAL_MACHINE = object() + winreg.KEY_READ = 0x20019 + winreg.KEY_WOW64_64KEY = 0x0100 + winreg.OpenKey.side_effect = (product_context, version_context) + winreg.QueryValueEx.side_effect = (("2026.1.3", 1), (install_root, 1)) + mocker.patch.object(binary_finder_module.importlib, "import_module", return_value=winreg) + + 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.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..99ce416f1d6 100644 --- a/cuda_pathfinder/tests/test_search_steps.py +++ b/cuda_pathfinder/tests/test_search_steps.py @@ -147,6 +147,25 @@ def test_rejects_unknown_sysconfig_tag(self, mocker): assert exc_info.value.platform_tag == "custom-win" +class TestWindowsMachineArch: + @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_normalizes_platform_machine(self, mocker, reported_machine, expected): + 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_rejects_unknown_machine(self, mocker): + mocker.patch.object(windows_arch_mod.platform, "machine", return_value="mips64") + + with pytest.raises(RuntimeError, match=r"Unsupported Windows machine architecture: 'mips64'"): + windows_arch_mod.windows_machine_arch() + + @pytest.mark.parametrize( ("machine", "target_arch", "expected"), (