Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions cuda_pathfinder/cuda/pathfinder/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,9 @@
find_nvidia_binary_utility as find_nvidia_binary_utility,
)
from cuda.pathfinder._binaries.supported_nvidia_binaries import SUPPORTED_BINARIES as _SUPPORTED_BINARIES
from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import (
find_nvidia_dynamic_lib as find_nvidia_dynamic_lib,
)
from cuda.pathfinder._dynamic_libs.load_dl_common import (
DynamicLibNotAvailableError as DynamicLibNotAvailableError,
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,19 +13,17 @@
from cuda.pathfinder._dynamic_libs.platform_loader import LOADER
from cuda.pathfinder._dynamic_libs.subprocess_protocol import (
MODE_CANARY,
MODE_FIND,
MODE_LOAD,
STATUS_NOT_FOUND,
STATUS_OK,
VALID_MODES,
format_dynamic_lib_subprocess_payload,
)

# NOTE: The main entrypoint (below) serves both production (canary probe)
# and tests (full loader). Keeping them together ensures a single subprocess
# protocol and CLI surface, so the test subprocess stays aligned with the
# production flow while avoiding a separate test-only module.
# Any production-code impact is negligible since the extra logic only runs
# in the subprocess entrypoint and only in test mode.
# The main entrypoint serves three modes — canary probe, find-without-load,
# and (test-only) full-loader exercise — behind a single subprocess protocol
# so test and production flows stay aligned.


def _probe_canary_abs_path(libname: str) -> str | None:
Expand Down Expand Up @@ -94,6 +92,22 @@ def probe_dynamic_lib_and_print_json(libname: str, mode: str) -> None:
print(format_dynamic_lib_subprocess_payload(status, abs_path))
return

if mode == MODE_FIND:
from cuda.pathfinder import load_nvidia_dynamic_lib

try:
loaded = load_nvidia_dynamic_lib(libname)
except DynamicLibNotFoundError as exc:
error = {"type": exc.__class__.__name__, "message": str(exc)}
print(format_dynamic_lib_subprocess_payload(STATUS_NOT_FOUND, None, error=error))
return
abs_path = loaded.abs_path
if not isinstance(abs_path, str):
raise RuntimeError(f"loaded.abs_path is not a string: {abs_path!r}")
_validate_abs_path(abs_path)
print(format_dynamic_lib_subprocess_payload(STATUS_OK, abs_path))
return

if mode == MODE_LOAD:
# Test-only path: exercises full loader behavior in isolation.
try:
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Locate an NVIDIA dynamic library on disk without loading it in this process.

Resolution is delegated to ``load_nvidia_dynamic_lib`` running in a fresh
Python subprocess. The full loader runs (including ``dlopen`` /
``LoadLibraryExW``) but only inside the child, so the caller's process is left
untouched.
"""

from __future__ import annotations

import functools

from cuda.pathfinder._dynamic_libs import load_nvidia_dynamic_lib as _load_module
from cuda.pathfinder._dynamic_libs.load_dl_common import (
DynamicLibNotAvailableError,
DynamicLibNotFoundError,
DynamicLibUnknownError,
)
from cuda.pathfinder._dynamic_libs.subprocess_protocol import (
MODE_FIND,
STATUS_OK,
run_dynamic_lib_subprocess,
)
from cuda.pathfinder._utils.platform_aware import IS_WINDOWS

# The subprocess runs the full loader (site-packages / conda / CUDA_PATH /
# canary cascade), which can be substantially slower than a single canary
# probe. Bound it so a wedged child cannot hang the caller indefinitely.
_FIND_SUBPROCESS_TIMEOUT_SECONDS = 120.0 if IS_WINDOWS else 30.0


@functools.cache
def find_nvidia_dynamic_lib(libname: str) -> str:
"""Return the absolute path to an NVIDIA dynamic library without loading it.

Resolution is performed by running :func:`load_nvidia_dynamic_lib` in a
fresh Python subprocess and reporting back the resolved absolute path.
The caller's process does **not** dlopen / LoadLibrary the library.

Args:
libname: Short name of the library (e.g., ``"cufile"``,
``"nvJitLink"``, ``"cudart"``).

Returns:
The absolute path the loader would have used in the caller's process.

Raises:
DynamicLibUnknownError: If ``libname`` is not a recognized library.
DynamicLibNotAvailableError: If ``libname`` is recognized but not
supported on this platform.
DynamicLibNotFoundError: If the library cannot be located.

Notes:
Because resolution happens in a separate process, results may differ
from an in-process ``load_nvidia_dynamic_lib`` if the caller's process
has DSOs loaded with custom ``RPATH``s or has already loaded a matching
library by some other mechanism. The intent is to report the path the
loader would pick when not influenced by other DSOs in the caller.
"""
# Indirect attribute access (not `from ... import`) so tests can
# monkeypatch the source-of-truth tables in `load_nvidia_dynamic_lib`.
if libname not in _load_module._ALL_KNOWN_LIBNAMES:
raise DynamicLibUnknownError(
f"Unknown library name: {libname!r}. Known names: {sorted(_load_module._ALL_KNOWN_LIBNAMES)}"
)
if libname not in _load_module._ALL_SUPPORTED_LIBNAMES:
raise DynamicLibNotAvailableError(
f"Library name {libname!r} is known but not available on {_load_module._PLATFORM_NAME}. "
f"Supported names on {_load_module._PLATFORM_NAME}: {sorted(_load_module._ALL_SUPPORTED_LIBNAMES)}"
)

payload = run_dynamic_lib_subprocess(
MODE_FIND,
libname,
timeout=_FIND_SUBPROCESS_TIMEOUT_SECONDS,
error_label=f"find_nvidia_dynamic_lib subprocess for {libname!r}",
)
if payload.status == STATUS_OK:
abs_path: str | None = payload.abs_path
assert abs_path is not None
return abs_path

error = payload.error
message = error["message"] if error and "message" in error else f"could not locate {libname!r}"
raise DynamicLibNotFoundError(message)
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@

import functools
import struct
import subprocess
import sys
from typing import TYPE_CHECKING

Expand All @@ -27,12 +26,9 @@
run_find_steps,
)
from cuda.pathfinder._dynamic_libs.subprocess_protocol import (
DYNAMIC_LIB_SUBPROCESS_CWD,
MODE_CANARY,
STATUS_OK,
DynamicLibSubprocessPayload,
build_dynamic_lib_subprocess_command,
parse_dynamic_lib_subprocess_payload,
run_dynamic_lib_subprocess,
)
from cuda.pathfinder._utils.platform_aware import IS_WINDOWS

Expand Down Expand Up @@ -74,61 +70,17 @@ def _load_driver_lib_no_cache(desc: LibDescriptor) -> LoadedDL:
)


def _coerce_subprocess_output(output: str | bytes | None) -> str:
if isinstance(output, bytes):
return output.decode(errors="replace")
return "" if output is None else output


def _raise_canary_probe_child_process_error(
*,
returncode: int | None = None,
timeout: float | None = None,
stderr: str | bytes | None = None,
) -> None:
if timeout is None:
error_line = f"Canary probe child process exited with code {returncode}."
else:
error_line = f"Canary probe child process timed out after {timeout} seconds."
raise ChildProcessError(
f"{error_line}\n"
"--- stderr-from-child-process ---\n"
f"{_coerce_subprocess_output(stderr)}"
"<end-of-stderr-from-child-process>\n"
)


@functools.cache
def _resolve_system_loaded_abs_path_in_subprocess(
libname: str,
*,
timeout: float = _CANARY_PROBE_TIMEOUT_SECONDS,
) -> str | None:
"""Resolve a canary library's absolute path in a fresh Python subprocess."""
try:
result = subprocess.run( # noqa: S603 - trusted argv: current interpreter + internal probe module
build_dynamic_lib_subprocess_command(MODE_CANARY, libname),
capture_output=True,
text=True,
timeout=timeout,
check=False,
cwd=DYNAMIC_LIB_SUBPROCESS_CWD,
)
except subprocess.TimeoutExpired as exc:
_raise_canary_probe_child_process_error(timeout=exc.timeout, stderr=exc.stderr)

if result.returncode != 0:
_raise_canary_probe_child_process_error(returncode=result.returncode, stderr=result.stderr)

payload: DynamicLibSubprocessPayload = parse_dynamic_lib_subprocess_payload(
result.stdout,
libname=libname,
error_label="Canary probe child process",
payload = run_dynamic_lib_subprocess(
MODE_CANARY, libname, timeout=timeout, error_label="Canary probe child process"
)
abs_path: str | None = payload.abs_path
if payload.status == STATUS_OK:
return abs_path
return None
return payload.abs_path if payload.status == STATUS_OK else None


def _loadable_via_canary_subprocess(libname: str, *, timeout: float = _CANARY_PROBE_TIMEOUT_SECONDS) -> bool:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,16 @@
from __future__ import annotations

import json
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Literal
from typing import Literal, NoReturn

MODE_CANARY: Literal["canary"] = "canary"
MODE_LOAD: Literal["load"] = "load"
VALID_MODES: tuple[Literal["canary"], Literal["load"]] = (MODE_CANARY, MODE_LOAD)
MODE_FIND: Literal["find"] = "find"
VALID_MODES: tuple[Literal["canary"], Literal["load"], Literal["find"]] = (MODE_CANARY, MODE_LOAD, MODE_FIND)

STATUS_OK: Literal["ok"] = "ok"
STATUS_NOT_FOUND: Literal["not-found"] = "not-found"
Expand All @@ -24,6 +26,7 @@
class DynamicLibSubprocessPayload:
status: Literal["ok", "not-found"]
abs_path: str | None
error: dict[str, str] | None = None


def format_dynamic_lib_subprocess_payload(
Expand Down Expand Up @@ -60,12 +63,78 @@ def parse_dynamic_lib_subprocess_payload(
raise RuntimeError(f"{error_label} emitted unexpected payload for {libname!r}: {payload!r}")
status = payload.get("status")
abs_path = payload.get("abs_path")
if status == STATUS_OK:
if not isinstance(abs_path, str):
raise RuntimeError(f"{error_label} emitted unexpected payload for {libname!r}: {payload!r}")
return DynamicLibSubprocessPayload(status=STATUS_OK, abs_path=abs_path)
if status == STATUS_NOT_FOUND:
if abs_path is not None:
raise RuntimeError(f"{error_label} emitted unexpected payload for {libname!r}: {payload!r}")
return DynamicLibSubprocessPayload(status=STATUS_NOT_FOUND, abs_path=None)
raise RuntimeError(f"{error_label} emitted unexpected payload for {libname!r}: {payload!r}")
error = payload.get("error")

def reject() -> NoReturn:
raise RuntimeError(f"{error_label} emitted unexpected payload for {libname!r}: {payload!r}")

if error is not None and not (
isinstance(error, dict) and all(isinstance(k, str) and isinstance(v, str) for k, v in error.items())
):
reject()
# STATUS_OK carries a path and never an error; STATUS_NOT_FOUND has no
# path but may carry a structured error from the child loader.
if status == STATUS_OK and isinstance(abs_path, str) and error is None:
return DynamicLibSubprocessPayload(status=STATUS_OK, abs_path=abs_path, error=None)
if status == STATUS_NOT_FOUND and abs_path is None:
return DynamicLibSubprocessPayload(status=STATUS_NOT_FOUND, abs_path=None, error=error)
reject()


def _coerce_subprocess_output(output: str | bytes | None) -> str:
if isinstance(output, bytes):
return output.decode(errors="replace")
return "" if output is None else output


def raise_subprocess_child_process_error(
error_label: str,
*,
returncode: int | None = None,
timeout: float | None = None,
stdout: str | bytes | None = None,
stderr: str | bytes | None = None,
) -> NoReturn:
if timeout is not None:
first_line = f"{error_label} timed out after {timeout} seconds."
else:
first_line = f"{error_label} exited with code {returncode}."
raise ChildProcessError(
f"{first_line}\n"
"--- stdout-from-child-process ---\n"
f"{_coerce_subprocess_output(stdout)}<end-of-stdout-from-child-process>\n"
"--- stderr-from-child-process ---\n"
f"{_coerce_subprocess_output(stderr)}<end-of-stderr-from-child-process>\n"
)


def run_dynamic_lib_subprocess(
mode: str,
libname: str,
*,
timeout: float,
error_label: str,
) -> DynamicLibSubprocessPayload:
"""Run the dynamic-lib subprocess and parse its payload.

Raises ``ChildProcessError`` if the child times out or exits non-zero;
otherwise returns the parsed payload (which may itself be ``STATUS_NOT_FOUND``).
"""
try:
result = subprocess.run( # noqa: S603 - trusted argv: current interpreter + internal probe module
build_dynamic_lib_subprocess_command(mode, libname),
capture_output=True,
text=True,
timeout=timeout,
check=False,
cwd=DYNAMIC_LIB_SUBPROCESS_CWD,
)
except subprocess.TimeoutExpired as exc:
raise_subprocess_child_process_error(error_label, timeout=exc.timeout, stdout=exc.stdout, stderr=exc.stderr)

if result.returncode != 0:
raise_subprocess_child_process_error(
error_label, returncode=result.returncode, stdout=result.stdout, stderr=result.stderr
)

return parse_dynamic_lib_subprocess_payload(result.stdout, libname=libname, error_label=error_label)
1 change: 1 addition & 0 deletions cuda_pathfinder/docs/source/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ CUDA bitcode and static libraries.

SUPPORTED_NVIDIA_LIBNAMES
load_nvidia_dynamic_lib
find_nvidia_dynamic_lib
LoadedDL
DynamicLibNotFoundError
DynamicLibUnknownError
Expand Down
Loading
Loading