refactor(pathfinder): unify dynamic-lib subprocess probing (production, test code)

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/dynamic_lib_subprocess.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+import sys
+from collections.abc import Sequence
+
+from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS
+from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError, LoadedDL
+from cuda.pathfinder._dynamic_libs.platform_loader import LOADER
+from cuda.pathfinder._dynamic_libs.subprocess_protocol import (
+    MODE_CANARY,
+    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.
+
+
+def _probe_canary_abs_path(libname: str) -> str | None:
+    desc = LIB_DESCRIPTORS.get(libname)
+    if desc is None:
+        raise ValueError(f"Unsupported canary library name: {libname!r}")
+    try:
+        loaded: LoadedDL | None = LOADER.load_with_system_search(desc)
+    except DynamicLibNotFoundError:
+        return None
+    if loaded is None:
+        return None
+    abs_path: str | None = loaded.abs_path
+    return abs_path
+
+
+def _validate_abs_path(abs_path: str) -> None:
+    assert abs_path, f"empty path: {abs_path=!r}"
+    assert os.path.isabs(abs_path), f"not absolute: {abs_path=!r}"
+    assert os.path.isfile(abs_path), f"not a file: {abs_path=!r}"
+
+
+def _load_nvidia_dynamic_lib_for_test(libname: str) -> str:
+    """Test-only loader used by the subprocess entrypoint."""
+    # Keep imports inside the subprocess body so startup stays focused on the
+    # code under test rather than the parent test module.
+    from cuda.pathfinder import load_nvidia_dynamic_lib
+    from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import _load_lib_no_cache
+    from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import (
+        SUPPORTED_LINUX_SONAMES,
+        SUPPORTED_WINDOWS_DLLS,
+    )
+    from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
+
+    loaded_dl_fresh = load_nvidia_dynamic_lib(libname)
+    if loaded_dl_fresh.was_already_loaded_from_elsewhere:
+        raise RuntimeError("loaded_dl_fresh.was_already_loaded_from_elsewhere")
+
+    abs_path = loaded_dl_fresh.abs_path
+    if not isinstance(abs_path, str):
+        raise RuntimeError(f"loaded_dl_fresh.abs_path is not a string: {abs_path!r}")
+    _validate_abs_path(abs_path)
+    assert loaded_dl_fresh.found_via is not None
+
+    loaded_dl_from_cache = load_nvidia_dynamic_lib(libname)
+    if loaded_dl_from_cache is not loaded_dl_fresh:
+        raise RuntimeError("loaded_dl_from_cache is not loaded_dl_fresh")
+
+    loaded_dl_no_cache = _load_lib_no_cache(libname)
+    supported_libs = SUPPORTED_WINDOWS_DLLS if IS_WINDOWS else SUPPORTED_LINUX_SONAMES
+    if not loaded_dl_no_cache.was_already_loaded_from_elsewhere and libname in supported_libs:
+        raise RuntimeError("not loaded_dl_no_cache.was_already_loaded_from_elsewhere")
+    abs_path_no_cache = loaded_dl_no_cache.abs_path
+    if not isinstance(abs_path_no_cache, str):
+        raise RuntimeError(f"loaded_dl_no_cache.abs_path is not a string: {abs_path_no_cache!r}")
+    if not os.path.samefile(abs_path_no_cache, abs_path):
+        raise RuntimeError(f"not os.path.samefile({abs_path_no_cache=!r}, {abs_path=!r})")
+    _validate_abs_path(abs_path_no_cache)
+    return abs_path
+
+
+def probe_dynamic_lib_and_print_json(libname: str, mode: str) -> None:
+    if mode == MODE_CANARY:
+        abs_path = _probe_canary_abs_path(libname)
+        status = STATUS_OK if abs_path is not None else STATUS_NOT_FOUND
+        print(format_dynamic_lib_subprocess_payload(status, abs_path))
+        return
+
+    if mode == MODE_LOAD:
+        # Test-only path: exercises full loader behavior in isolation.
+        try:
+            abs_path = _load_nvidia_dynamic_lib_for_test(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
+        print(format_dynamic_lib_subprocess_payload(STATUS_OK, abs_path))
+        return
+
+    raise ValueError(f"Unsupported subprocess probe mode: {mode!r}")
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    args = list(sys.argv[1:] if argv is None else argv)
+    if len(args) != 2 or args[0] not in VALID_MODES:
+        modes = ", ".join(VALID_MODES)
+        raise SystemExit(
+            f"Usage: python -m cuda.pathfinder._dynamic_libs.dynamic_lib_subprocess <mode> <libname>\nModes: {modes}"
+        )
+    mode, libname = args
+    probe_dynamic_lib_and_print_json(libname, mode)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py

@@ -4,11 +4,9 @@
 from __future__ import annotations
 
 import functools
-import json
 import struct
 import subprocess
 import sys
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS
@@ -28,6 +26,14 @@
     find_via_ctk_root,
     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,
+)
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 
 if TYPE_CHECKING:
@@ -40,9 +46,7 @@
     name for name, desc in LIB_DESCRIPTORS.items() if (desc.windows_dlls if IS_WINDOWS else desc.linux_sonames)
 )
 _PLATFORM_NAME = "Windows" if IS_WINDOWS else "Linux"
-_CANARY_PROBE_MODULE = "cuda.pathfinder._dynamic_libs.canary_probe_subprocess"
 _CANARY_PROBE_TIMEOUT_SECONDS = 10.0
-_CANARY_PROBE_IMPORT_ROOT = Path(__file__).resolve().parents[3]
 
 # Driver libraries: shipped with the NVIDIA display driver, always on the
 # system linker path.  These skip all CTK search steps (site-packages,
@@ -99,34 +103,28 @@ def _resolve_system_loaded_abs_path_in_subprocess(libname: str) -> 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
-            [sys.executable, "-m", _CANARY_PROBE_MODULE, libname],
+            build_dynamic_lib_subprocess_command(MODE_CANARY, libname),
             capture_output=True,
             text=True,
             timeout=_CANARY_PROBE_TIMEOUT_SECONDS,
             check=False,
-            cwd=_CANARY_PROBE_IMPORT_ROOT,
+            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)
 
-    # Use the final non-empty line in case earlier output lines are emitted.
-    lines = [line for line in result.stdout.splitlines() if line.strip()]
-    if not lines:
-        raise RuntimeError(f"Canary probe child process produced no stdout payload for {libname!r}")
-    try:
-        payload = json.loads(lines[-1])
-    except json.JSONDecodeError:
-        raise RuntimeError(
-            f"Canary probe child process emitted invalid JSON payload for {libname!r}: {lines[-1]!r}"
-        ) from None
-    if isinstance(payload, str):
-        return payload
-    if payload is None:
-        return None
-    raise RuntimeError(f"Canary probe child process emitted unexpected payload for {libname!r}: {payload!r}")
+    payload: DynamicLibSubprocessPayload = parse_dynamic_lib_subprocess_payload(
+        result.stdout,
+        libname=libname,
+        error_label="Canary probe child process",
+    )
+    abs_path: str | None = payload.abs_path
+    if payload.status == STATUS_OK:
+        return abs_path
+    return None
 
 
 def _try_ctk_root_canary(ctx: SearchContext) -> str | None:

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/subprocess_protocol.py

@@ -0,0 +1,71 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import json
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+MODE_CANARY: Literal["canary"] = "canary"
+MODE_LOAD: Literal["load"] = "load"
+VALID_MODES: tuple[Literal["canary"], Literal["load"]] = (MODE_CANARY, MODE_LOAD)
+
+STATUS_OK: Literal["ok"] = "ok"
+STATUS_NOT_FOUND: Literal["not-found"] = "not-found"
+
+DYNAMIC_LIB_SUBPROCESS_MODULE = "cuda.pathfinder._dynamic_libs.dynamic_lib_subprocess"
+DYNAMIC_LIB_SUBPROCESS_CWD = Path(__file__).resolve().parents[3]
+
+
+@dataclass(frozen=True)
+class DynamicLibSubprocessPayload:
+    status: Literal["ok", "not-found"]
+    abs_path: str | None
+
+
+def format_dynamic_lib_subprocess_payload(
+    status: Literal["ok", "not-found"],
+    abs_path: str | None,
+    *,
+    error: dict[str, str] | None = None,
+) -> str:
+    payload: dict[str, object] = {"status": status, "abs_path": abs_path}
+    if error is not None:
+        payload["error"] = error
+    return json.dumps(payload)
+
+
+def build_dynamic_lib_subprocess_command(mode: str, libname: str) -> list[str]:
+    return [sys.executable, "-m", DYNAMIC_LIB_SUBPROCESS_MODULE, mode, libname]
+
+
+def parse_dynamic_lib_subprocess_payload(
+    stdout: str,
+    *,
+    libname: str,
+    error_label: str,
+) -> DynamicLibSubprocessPayload:
+    # Use the final non-empty line in case earlier output lines are emitted.
+    lines = [line for line in stdout.splitlines() if line.strip()]
+    if not lines:
+        raise RuntimeError(f"{error_label} produced no stdout payload for {libname!r}")
+    try:
+        payload = json.loads(lines[-1])
+    except json.JSONDecodeError:
+        raise RuntimeError(f"{error_label} emitted invalid JSON payload for {libname!r}: {lines[-1]!r}") from None
+    if not isinstance(payload, dict):
+        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}")