From c0f2aef4967f9bda0104ff83b79dc08db1e364a8 Mon Sep 17 00:00:00 2001 From: Michael Wang Date: Tue, 4 Aug 2026 12:04:29 -0700 Subject: [PATCH 1/2] Make static library discovery architecture-aware --- .../_static_libs/find_static_lib.py | 81 +++++++++---- .../cuda/pathfinder/_utils/binary_format.py | 110 ++++++++++++++++++ .../docs/source/release/1.6.1-notes.rst | 5 + cuda_pathfinder/tests/test_find_static_lib.py | 83 ++++++++++++- 4 files changed, 254 insertions(+), 25 deletions(-) create mode 100644 cuda_pathfinder/cuda/pathfinder/_utils/binary_format.py 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 804b1c04be7..de1eaaafd4a 100644 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py @@ -6,9 +6,15 @@ from dataclasses import dataclass from typing import NoReturn, TypedDict +from cuda.pathfinder._utils.binary_format import ( + BinaryFormat, + python_binary_format, + static_archive_matches_binary_format, +) 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_python_arch class StaticLibNotFoundError(RuntimeError): @@ -32,17 +38,29 @@ class _StaticLibInfo(TypedDict): site_packages_dirs: tuple[str, ...] +def _cudadevrt_info() -> _StaticLibInfo: + if not IS_WINDOWS: + return { + "filename": "libcudadevrt.a", + "ctk_rel_paths": ("lib64", "lib"), + "conda_rel_paths": ("lib",), + "site_packages_dirs": ("nvidia/cu13/lib", "nvidia/cuda_runtime/lib"), + } + + target_arch = windows_python_arch() + arch_dir = {"x64": "x64", "arm64": "arm64"}[target_arch] + component_wheel_dirs = ("nvidia/cuda_runtime/lib/x64",) if target_arch == "x64" else () + conda_fallback_dirs = ("lib",) if target_arch == "x64" else () + return { + "filename": "cudadevrt.lib", + "ctk_rel_paths": (os.path.join("lib", arch_dir),), + "conda_rel_paths": (os.path.join("lib", arch_dir), *conda_fallback_dirs), + "site_packages_dirs": (f"nvidia/cu13/lib/{arch_dir}", *component_wheel_dirs), + } + + _SUPPORTED_STATIC_LIBS_INFO: dict[str, _StaticLibInfo] = { - "cudadevrt": { - "filename": "cudadevrt.lib" if IS_WINDOWS else "libcudadevrt.a", - "ctk_rel_paths": (os.path.join("lib", "x64"),) if IS_WINDOWS else ("lib64", "lib"), - "conda_rel_paths": ((os.path.join("lib", "x64"), "lib") if IS_WINDOWS else ("lib",)), - "site_packages_dirs": ( - ("nvidia/cu13/lib/x64", "nvidia/cuda_runtime/lib/x64") - if IS_WINDOWS - else ("nvidia/cu13/lib", "nvidia/cuda_runtime/lib") - ), - }, + "cudadevrt": _cudadevrt_info(), } SUPPORTED_STATIC_LIBS: tuple[str, ...] = tuple(sorted(_SUPPORTED_STATIC_LIBS_INFO.keys())) @@ -68,16 +86,29 @@ def __init__(self, name: str) -> None: self.ctk_rel_paths: tuple[str, ...] = self.config["ctk_rel_paths"] self.conda_rel_paths: tuple[str, ...] = self.config["conda_rel_paths"] self.site_packages_dirs: tuple[str, ...] = self.config["site_packages_dirs"] + self.binary_format: BinaryFormat = python_binary_format() self.error_messages: list[str] = [] self.attachments: list[str] = [] + def compatible_candidate(self, file_path: str) -> str | None: + if not os.path.isfile(file_path): + return None + if static_archive_matches_binary_format(file_path, self.binary_format): + return file_path + self.error_messages.append( + f'Incompatible static library: "{file_path}" does not match ' + f"Python's {self.binary_format.kind.upper()} machine type 0x{self.binary_format.machine:04x}" + ) + return None + def try_site_packages(self) -> str | None: for rel_dir in self.site_packages_dirs: sub_dir = tuple(rel_dir.split("/")) for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): file_path = os.path.join(abs_dir, self.filename) - if os.path.isfile(file_path): - return file_path + candidate = self.compatible_candidate(file_path) + if candidate is not None: + return candidate return None def try_with_conda_prefix(self) -> str | None: @@ -88,8 +119,9 @@ def try_with_conda_prefix(self) -> str | None: anchor = os.path.join(conda_prefix, "Library") if IS_WINDOWS else conda_prefix for rel_path in self.conda_rel_paths: file_path = os.path.join(anchor, rel_path, self.filename) - if os.path.isfile(file_path): - return file_path + candidate = self.compatible_candidate(file_path) + if candidate is not None: + return candidate return None def try_with_cuda_home(self) -> str | None: @@ -100,15 +132,18 @@ def try_with_cuda_home(self) -> str | None: for rel_path in self.ctk_rel_paths: file_path = os.path.join(cuda_home, rel_path, self.filename) - if os.path.isfile(file_path): - return file_path - - _no_such_file_in_dir( - os.path.join(cuda_home, self.ctk_rel_paths[0]), - self.filename, - self.error_messages, - self.attachments, - ) + candidate = self.compatible_candidate(file_path) + if candidate is not None: + return candidate + + first_file_path = os.path.join(cuda_home, self.ctk_rel_paths[0], self.filename) + if not os.path.isfile(first_file_path): + _no_such_file_in_dir( + os.path.join(cuda_home, self.ctk_rel_paths[0]), + self.filename, + self.error_messages, + self.attachments, + ) return None def raise_not_found_error(self) -> NoReturn: diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/binary_format.py b/cuda_pathfinder/cuda/pathfinder/_utils/binary_format.py new file mode 100644 index 00000000000..110a6f70b8b --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_utils/binary_format.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import functools +import sys +from dataclasses import dataclass +from typing import Literal + +_COFF_MACHINE_TYPES = { + 0x014C, # x86 + 0x01C0, # Arm + 0x01C4, # Armv7 + 0x8664, # x64 + 0xAA64, # Arm64 +} + + +@dataclass(frozen=True, slots=True) +class BinaryFormat: + """The object format and machine type of a native binary.""" + + kind: Literal["coff", "elf"] + machine: int + + +def _binary_format_from_object(data: bytes) -> BinaryFormat | None: + if data.startswith(b"\x7fELF") and len(data) >= 20: + if data[5] == 1: + byte_order: Literal["little", "big"] = "little" + elif data[5] == 2: + byte_order = "big" + else: + return None + return BinaryFormat("elf", int.from_bytes(data[18:20], byte_order)) + + # A regular COFF object starts with Machine. Import objects and bigobj + # objects use the anonymous-object signature and store Machine at offset 6. + if data.startswith(b"\x00\x00\xff\xff") and len(data) >= 8: + machine = int.from_bytes(data[6:8], "little") + if machine in _COFF_MACHINE_TYPES: + return BinaryFormat("coff", machine) + if len(data) >= 20: + machine = int.from_bytes(data[:2], "little") + if machine in _COFF_MACHINE_TYPES: + return BinaryFormat("coff", machine) + return None + + +@functools.cache +def python_binary_format() -> BinaryFormat: + """Read the running Python executable's native object format and machine.""" + try: + with open(sys.executable, "rb") as stream: + prefix = stream.read(64) + binary_format = _binary_format_from_object(prefix) + if binary_format is not None and binary_format.kind == "elf": + return binary_format + + if prefix.startswith(b"MZ") and len(prefix) >= 64: + pe_offset = int.from_bytes(prefix[0x3C:0x40], "little") + stream.seek(pe_offset) + if stream.read(4) == b"PE\0\0": + machine_bytes = stream.read(2) + if len(machine_bytes) == 2: + return BinaryFormat("coff", int.from_bytes(machine_bytes, "little")) + except OSError as exc: + raise RuntimeError(f"Could not inspect Python executable {sys.executable!r}: {exc}") from exc + + raise RuntimeError(f"Unsupported Python executable binary format: {sys.executable!r}") + + +def static_archive_matches_binary_format(path: str, expected: BinaryFormat) -> bool: + """Return whether an ar archive contains an object matching ``expected``.""" + try: + with open(path, "rb") as stream: + if stream.read(8) != b"!\n": + return False + + while True: + header = stream.read(60) + if not header: + return False + if len(header) != 60 or header[58:60] != b"`\n": + return False + + member_name = header[:16].decode("ascii", "replace").rstrip() + member_size = int(header[48:58].decode("ascii").strip()) + member_start = stream.tell() + object_start = member_start + object_size = member_size + + if member_name.startswith("#1/"): + name_size = int(member_name[3:]) + if name_size > member_size: + return False + object_start += name_size + object_size -= name_size + + is_index = member_name in ("/", "//", "/SYM64/") or member_name.startswith("__.SYMDEF") + if not is_index: + stream.seek(object_start) + actual = _binary_format_from_object(stream.read(min(object_size, 64))) + if actual is not None: + return actual == expected + + stream.seek(member_start + member_size + member_size % 2) + except (OSError, ValueError): + return False diff --git a/cuda_pathfinder/docs/source/release/1.6.1-notes.rst b/cuda_pathfinder/docs/source/release/1.6.1-notes.rst index 29a3106e0f1..7e0bf0810d9 100644 --- a/cuda_pathfinder/docs/source/release/1.6.1-notes.rst +++ b/cuda_pathfinder/docs/source/release/1.6.1-notes.rst @@ -24,6 +24,11 @@ Highlights * Add ``UnsupportedArchError`` for unsupported Windows Python platform tags. +* Make static-library discovery binary-format aware. Candidates are now + checked against the current Python executable's object format and machine + type before they are returned, and Windows searches use the matching x64 or + Arm64 CUDA Toolkit and wheel directories. + Internal maintenance -------------------- diff --git a/cuda_pathfinder/tests/test_find_static_lib.py b/cuda_pathfinder/tests/test_find_static_lib.py index e5560dcabbf..56066cd1263 100644 --- a/cuda_pathfinder/tests/test_find_static_lib.py +++ b/cuda_pathfinder/tests/test_find_static_lib.py @@ -31,10 +31,26 @@ def clear_find_static_lib_cache(): get_cuda_path_or_home.cache_clear() -def _make_static_lib_file(dir_path: Path, filename: str) -> str: +def _archive_object(binary_format): + if binary_format.kind == "elf": + data = bytearray(64) + data[:6] = b"\x7fELF\x02\x01" + data[18:20] = binary_format.machine.to_bytes(2, "little") + return bytes(data) + + data = bytearray(20) + data[:2] = binary_format.machine.to_bytes(2, "little") + return bytes(data) + + +def _make_static_lib_file(dir_path: Path, filename: str, binary_format=None) -> str: dir_path.mkdir(parents=True, exist_ok=True) file_path = dir_path / filename - file_path.touch() + if binary_format is None: + binary_format = find_static_lib_module.python_binary_format() + member = _archive_object(binary_format) + header = b"object.o/ 0 0 0 100644 " + str(len(member)).encode("ascii").ljust(10) + b"`\n" + file_path.write_bytes(b"!\n" + header + member + (b"\n" if len(member) % 2 else b"")) return str(file_path) @@ -143,6 +159,69 @@ def test_locate_static_lib_conda_rel_path_fallback(monkeypatch, tmp_path): assert located_lib.found_via == "conda" +@pytest.mark.parametrize( + ("kind", "machine"), + (("elf", 0x003E), ("coff", 0x8664), ("coff", 0xAA64)), +) +@pytest.mark.agent_authored(model="gpt-5.6") +def test_static_archive_binary_format_matching(tmp_path, kind, machine): + expected = find_static_lib_module.BinaryFormat(kind, machine) + archive = _make_static_lib_file(tmp_path, "library.a", expected) + + assert find_static_lib_module.static_archive_matches_binary_format(archive, expected) + assert not find_static_lib_module.static_archive_matches_binary_format( + archive, + find_static_lib_module.BinaryFormat(kind, machine ^ 1), + ) + + +@pytest.mark.usefixtures("clear_find_static_lib_cache") +@pytest.mark.agent_authored(model="gpt-5.6") +def test_locate_static_lib_skips_incompatible_candidate(monkeypatch, tmp_path): + filename = CUDADEVRT_INFO["filename"] + expected = find_static_lib_module.python_binary_format() + incompatible = find_static_lib_module.BinaryFormat(expected.kind, expected.machine ^ 1) + + site_packages_lib_dir = tmp_path / "site-packages" + site_packages_path = _make_static_lib_file(site_packages_lib_dir, filename, incompatible) + conda_prefix = tmp_path / "conda-prefix" + conda_lib_dir = _conda_anchor(conda_prefix) / Path(CUDADEVRT_INFO["conda_rel_paths"][0]) + conda_path = _make_static_lib_file(conda_lib_dir, filename, expected) + + monkeypatch.setattr( + find_static_lib_module, + "find_sub_dirs_all_sitepackages", + lambda _sub_dir: [str(site_packages_lib_dir)], + ) + monkeypatch.setenv("CONDA_PREFIX", str(conda_prefix)) + monkeypatch.delenv("CUDA_HOME", raising=False) + monkeypatch.delenv("CUDA_PATH", raising=False) + + located_lib = locate_static_lib("cudadevrt") + assert located_lib.abs_path == conda_path + assert located_lib.abs_path != site_packages_path + assert located_lib.found_via == "conda" + + +@pytest.mark.parametrize( + ("target_arch", "expected_dirs"), + ( + ("x64", ("nvidia/cu13/lib/x64", "nvidia/cuda_runtime/lib/x64")), + ("arm64", ("nvidia/cu13/lib/arm64",)), + ), +) +@pytest.mark.agent_authored(model="gpt-5.6") +def test_cudadevrt_windows_paths_follow_python_arch(monkeypatch, target_arch, expected_dirs): + monkeypatch.setattr(find_static_lib_module, "IS_WINDOWS", True) + monkeypatch.setattr(find_static_lib_module, "windows_python_arch", lambda: target_arch) + + info = find_static_lib_module._cudadevrt_info() + + assert info["ctk_rel_paths"] == (os.path.join("lib", target_arch),) + assert info["site_packages_dirs"] == expected_dirs + assert all(target_arch in path for path in info["ctk_rel_paths"]) + + @pytest.mark.usefixtures("clear_find_static_lib_cache") def test_find_static_lib_not_found_error_includes_cuda_home_directory_listing(monkeypatch, tmp_path): filename = CUDADEVRT_INFO["filename"] From 3e31b2294fe54b23134d5833e6a30e62c6ea9f63 Mon Sep 17 00:00:00 2001 From: Michael Wang Date: Tue, 4 Aug 2026 14:32:51 -0700 Subject: [PATCH 2/2] Use architecture-specific static library paths --- .../_static_libs/find_static_lib.py | 55 +++------ .../cuda/pathfinder/_utils/binary_format.py | 110 ------------------ .../docs/source/release/1.6.1-notes.rst | 8 +- cuda_pathfinder/tests/test_find_static_lib.py | 94 ++++----------- 4 files changed, 45 insertions(+), 222 deletions(-) delete mode 100644 cuda_pathfinder/cuda/pathfinder/_utils/binary_format.py 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 de1eaaafd4a..ea5a740aec4 100644 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_static_lib.py @@ -6,11 +6,6 @@ from dataclasses import dataclass from typing import NoReturn, TypedDict -from cuda.pathfinder._utils.binary_format import ( - BinaryFormat, - python_binary_format, - static_archive_matches_binary_format, -) 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 @@ -47,10 +42,9 @@ def _cudadevrt_info() -> _StaticLibInfo: "site_packages_dirs": ("nvidia/cu13/lib", "nvidia/cuda_runtime/lib"), } - target_arch = windows_python_arch() - arch_dir = {"x64": "x64", "arm64": "arm64"}[target_arch] - component_wheel_dirs = ("nvidia/cuda_runtime/lib/x64",) if target_arch == "x64" else () - conda_fallback_dirs = ("lib",) if target_arch == "x64" else () + arch_dir = windows_python_arch() + component_wheel_dirs = ("nvidia/cuda_runtime/lib/x64",) if arch_dir == "x64" else () + conda_fallback_dirs = ("lib",) if arch_dir == "x64" else () return { "filename": "cudadevrt.lib", "ctk_rel_paths": (os.path.join("lib", arch_dir),), @@ -86,29 +80,16 @@ def __init__(self, name: str) -> None: self.ctk_rel_paths: tuple[str, ...] = self.config["ctk_rel_paths"] self.conda_rel_paths: tuple[str, ...] = self.config["conda_rel_paths"] self.site_packages_dirs: tuple[str, ...] = self.config["site_packages_dirs"] - self.binary_format: BinaryFormat = python_binary_format() self.error_messages: list[str] = [] self.attachments: list[str] = [] - def compatible_candidate(self, file_path: str) -> str | None: - if not os.path.isfile(file_path): - return None - if static_archive_matches_binary_format(file_path, self.binary_format): - return file_path - self.error_messages.append( - f'Incompatible static library: "{file_path}" does not match ' - f"Python's {self.binary_format.kind.upper()} machine type 0x{self.binary_format.machine:04x}" - ) - return None - def try_site_packages(self) -> str | None: for rel_dir in self.site_packages_dirs: sub_dir = tuple(rel_dir.split("/")) for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): file_path = os.path.join(abs_dir, self.filename) - candidate = self.compatible_candidate(file_path) - if candidate is not None: - return candidate + if os.path.isfile(file_path): + return file_path return None def try_with_conda_prefix(self) -> str | None: @@ -119,9 +100,8 @@ def try_with_conda_prefix(self) -> str | None: anchor = os.path.join(conda_prefix, "Library") if IS_WINDOWS else conda_prefix for rel_path in self.conda_rel_paths: file_path = os.path.join(anchor, rel_path, self.filename) - candidate = self.compatible_candidate(file_path) - if candidate is not None: - return candidate + if os.path.isfile(file_path): + return file_path return None def try_with_cuda_home(self) -> str | None: @@ -132,18 +112,15 @@ def try_with_cuda_home(self) -> str | None: for rel_path in self.ctk_rel_paths: file_path = os.path.join(cuda_home, rel_path, self.filename) - candidate = self.compatible_candidate(file_path) - if candidate is not None: - return candidate - - first_file_path = os.path.join(cuda_home, self.ctk_rel_paths[0], self.filename) - if not os.path.isfile(first_file_path): - _no_such_file_in_dir( - os.path.join(cuda_home, self.ctk_rel_paths[0]), - self.filename, - self.error_messages, - self.attachments, - ) + if os.path.isfile(file_path): + return file_path + + _no_such_file_in_dir( + os.path.join(cuda_home, self.ctk_rel_paths[0]), + self.filename, + self.error_messages, + self.attachments, + ) return None def raise_not_found_error(self) -> NoReturn: diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/binary_format.py b/cuda_pathfinder/cuda/pathfinder/_utils/binary_format.py deleted file mode 100644 index 110a6f70b8b..00000000000 --- a/cuda_pathfinder/cuda/pathfinder/_utils/binary_format.py +++ /dev/null @@ -1,110 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import functools -import sys -from dataclasses import dataclass -from typing import Literal - -_COFF_MACHINE_TYPES = { - 0x014C, # x86 - 0x01C0, # Arm - 0x01C4, # Armv7 - 0x8664, # x64 - 0xAA64, # Arm64 -} - - -@dataclass(frozen=True, slots=True) -class BinaryFormat: - """The object format and machine type of a native binary.""" - - kind: Literal["coff", "elf"] - machine: int - - -def _binary_format_from_object(data: bytes) -> BinaryFormat | None: - if data.startswith(b"\x7fELF") and len(data) >= 20: - if data[5] == 1: - byte_order: Literal["little", "big"] = "little" - elif data[5] == 2: - byte_order = "big" - else: - return None - return BinaryFormat("elf", int.from_bytes(data[18:20], byte_order)) - - # A regular COFF object starts with Machine. Import objects and bigobj - # objects use the anonymous-object signature and store Machine at offset 6. - if data.startswith(b"\x00\x00\xff\xff") and len(data) >= 8: - machine = int.from_bytes(data[6:8], "little") - if machine in _COFF_MACHINE_TYPES: - return BinaryFormat("coff", machine) - if len(data) >= 20: - machine = int.from_bytes(data[:2], "little") - if machine in _COFF_MACHINE_TYPES: - return BinaryFormat("coff", machine) - return None - - -@functools.cache -def python_binary_format() -> BinaryFormat: - """Read the running Python executable's native object format and machine.""" - try: - with open(sys.executable, "rb") as stream: - prefix = stream.read(64) - binary_format = _binary_format_from_object(prefix) - if binary_format is not None and binary_format.kind == "elf": - return binary_format - - if prefix.startswith(b"MZ") and len(prefix) >= 64: - pe_offset = int.from_bytes(prefix[0x3C:0x40], "little") - stream.seek(pe_offset) - if stream.read(4) == b"PE\0\0": - machine_bytes = stream.read(2) - if len(machine_bytes) == 2: - return BinaryFormat("coff", int.from_bytes(machine_bytes, "little")) - except OSError as exc: - raise RuntimeError(f"Could not inspect Python executable {sys.executable!r}: {exc}") from exc - - raise RuntimeError(f"Unsupported Python executable binary format: {sys.executable!r}") - - -def static_archive_matches_binary_format(path: str, expected: BinaryFormat) -> bool: - """Return whether an ar archive contains an object matching ``expected``.""" - try: - with open(path, "rb") as stream: - if stream.read(8) != b"!\n": - return False - - while True: - header = stream.read(60) - if not header: - return False - if len(header) != 60 or header[58:60] != b"`\n": - return False - - member_name = header[:16].decode("ascii", "replace").rstrip() - member_size = int(header[48:58].decode("ascii").strip()) - member_start = stream.tell() - object_start = member_start - object_size = member_size - - if member_name.startswith("#1/"): - name_size = int(member_name[3:]) - if name_size > member_size: - return False - object_start += name_size - object_size -= name_size - - is_index = member_name in ("/", "//", "/SYM64/") or member_name.startswith("__.SYMDEF") - if not is_index: - stream.seek(object_start) - actual = _binary_format_from_object(stream.read(min(object_size, 64))) - if actual is not None: - return actual == expected - - stream.seek(member_start + member_size + member_size % 2) - except (OSError, ValueError): - return False diff --git a/cuda_pathfinder/docs/source/release/1.6.1-notes.rst b/cuda_pathfinder/docs/source/release/1.6.1-notes.rst index 7e0bf0810d9..919963802ff 100644 --- a/cuda_pathfinder/docs/source/release/1.6.1-notes.rst +++ b/cuda_pathfinder/docs/source/release/1.6.1-notes.rst @@ -24,10 +24,10 @@ Highlights * Add ``UnsupportedArchError`` for unsupported Windows Python platform tags. -* Make static-library discovery binary-format aware. Candidates are now - checked against the current Python executable's object format and machine - type before they are returned, and Windows searches use the matching x64 or - Arm64 CUDA Toolkit and wheel directories. +* Make Windows static-library discovery architecture-aware. Searches now use + the current Python interpreter architecture to select the matching + ``lib/x64`` or ``lib/arm64`` CUDA Toolkit and wheel directories. CUDA 12 + component-wheel and legacy Conda fallbacks remain x64-only. Internal maintenance -------------------- diff --git a/cuda_pathfinder/tests/test_find_static_lib.py b/cuda_pathfinder/tests/test_find_static_lib.py index 56066cd1263..6d29a8def11 100644 --- a/cuda_pathfinder/tests/test_find_static_lib.py +++ b/cuda_pathfinder/tests/test_find_static_lib.py @@ -31,26 +31,10 @@ def clear_find_static_lib_cache(): get_cuda_path_or_home.cache_clear() -def _archive_object(binary_format): - if binary_format.kind == "elf": - data = bytearray(64) - data[:6] = b"\x7fELF\x02\x01" - data[18:20] = binary_format.machine.to_bytes(2, "little") - return bytes(data) - - data = bytearray(20) - data[:2] = binary_format.machine.to_bytes(2, "little") - return bytes(data) - - -def _make_static_lib_file(dir_path: Path, filename: str, binary_format=None) -> str: +def _make_static_lib_file(dir_path: Path, filename: str) -> str: dir_path.mkdir(parents=True, exist_ok=True) file_path = dir_path / filename - if binary_format is None: - binary_format = find_static_lib_module.python_binary_format() - member = _archive_object(binary_format) - header = b"object.o/ 0 0 0 100644 " + str(len(member)).encode("ascii").ljust(10) + b"`\n" - file_path.write_bytes(b"!\n" + header + member + (b"\n" if len(member) % 2 else b"")) + file_path.touch() return str(file_path) @@ -160,66 +144,38 @@ def test_locate_static_lib_conda_rel_path_fallback(monkeypatch, tmp_path): @pytest.mark.parametrize( - ("kind", "machine"), - (("elf", 0x003E), ("coff", 0x8664), ("coff", 0xAA64)), -) -@pytest.mark.agent_authored(model="gpt-5.6") -def test_static_archive_binary_format_matching(tmp_path, kind, machine): - expected = find_static_lib_module.BinaryFormat(kind, machine) - archive = _make_static_lib_file(tmp_path, "library.a", expected) - - assert find_static_lib_module.static_archive_matches_binary_format(archive, expected) - assert not find_static_lib_module.static_archive_matches_binary_format( - archive, - find_static_lib_module.BinaryFormat(kind, machine ^ 1), - ) - - -@pytest.mark.usefixtures("clear_find_static_lib_cache") -@pytest.mark.agent_authored(model="gpt-5.6") -def test_locate_static_lib_skips_incompatible_candidate(monkeypatch, tmp_path): - filename = CUDADEVRT_INFO["filename"] - expected = find_static_lib_module.python_binary_format() - incompatible = find_static_lib_module.BinaryFormat(expected.kind, expected.machine ^ 1) - - site_packages_lib_dir = tmp_path / "site-packages" - site_packages_path = _make_static_lib_file(site_packages_lib_dir, filename, incompatible) - conda_prefix = tmp_path / "conda-prefix" - conda_lib_dir = _conda_anchor(conda_prefix) / Path(CUDADEVRT_INFO["conda_rel_paths"][0]) - conda_path = _make_static_lib_file(conda_lib_dir, filename, expected) - - monkeypatch.setattr( - find_static_lib_module, - "find_sub_dirs_all_sitepackages", - lambda _sub_dir: [str(site_packages_lib_dir)], - ) - monkeypatch.setenv("CONDA_PREFIX", str(conda_prefix)) - monkeypatch.delenv("CUDA_HOME", raising=False) - monkeypatch.delenv("CUDA_PATH", raising=False) - - located_lib = locate_static_lib("cudadevrt") - assert located_lib.abs_path == conda_path - assert located_lib.abs_path != site_packages_path - assert located_lib.found_via == "conda" - - -@pytest.mark.parametrize( - ("target_arch", "expected_dirs"), + ("target_arch", "expected_ctk_dirs", "expected_conda_dirs", "expected_site_packages_dirs"), ( - ("x64", ("nvidia/cu13/lib/x64", "nvidia/cuda_runtime/lib/x64")), - ("arm64", ("nvidia/cu13/lib/arm64",)), + ( + "x64", + (os.path.join("lib", "x64"),), + (os.path.join("lib", "x64"), "lib"), + ("nvidia/cu13/lib/x64", "nvidia/cuda_runtime/lib/x64"), + ), + ( + "arm64", + (os.path.join("lib", "arm64"),), + (os.path.join("lib", "arm64"),), + ("nvidia/cu13/lib/arm64",), + ), ), ) @pytest.mark.agent_authored(model="gpt-5.6") -def test_cudadevrt_windows_paths_follow_python_arch(monkeypatch, target_arch, expected_dirs): +def test_cudadevrt_windows_paths_follow_python_arch( + monkeypatch, + target_arch, + expected_ctk_dirs, + expected_conda_dirs, + expected_site_packages_dirs, +): monkeypatch.setattr(find_static_lib_module, "IS_WINDOWS", True) monkeypatch.setattr(find_static_lib_module, "windows_python_arch", lambda: target_arch) info = find_static_lib_module._cudadevrt_info() - assert info["ctk_rel_paths"] == (os.path.join("lib", target_arch),) - assert info["site_packages_dirs"] == expected_dirs - assert all(target_arch in path for path in info["ctk_rel_paths"]) + assert info["ctk_rel_paths"] == expected_ctk_dirs + assert info["conda_rel_paths"] == expected_conda_dirs + assert info["site_packages_dirs"] == expected_site_packages_dirs @pytest.mark.usefixtures("clear_find_static_lib_cache")