|
14 | 14 | from contextlib import AbstractContextManager, contextmanager, nullcontext |
15 | 15 | from dataclasses import dataclass |
16 | 16 | from enum import Enum |
| 17 | +from pathlib import Path |
17 | 18 | from typing import Any, TypeVar, cast |
18 | 19 |
|
19 | 20 | from arraybridge.array_operations import ( |
|
41 | 42 | DLPackExporter = Callable[[Any, Any], Any | None] |
42 | 43 | DLPackValidator = Callable[[Any, Any], bool] |
43 | 44 | OOMMatcher = Callable[[BaseException, Any | None], bool] |
| 45 | +SubprocessEnvironmentResolver = Callable[[Mapping[str, str]], dict[str, str]] |
| 46 | + |
| 47 | + |
| 48 | +def _identity_subprocess_environment( |
| 49 | + environment: Mapping[str, str], |
| 50 | +) -> dict[str, str]: |
| 51 | + return dict(environment) |
| 52 | + |
| 53 | + |
| 54 | +def _nvidia_wheel_library_paths() -> tuple[str, ...]: |
| 55 | + """Return native-library directories declared by installed NVIDIA wheels.""" |
| 56 | + |
| 57 | + try: |
| 58 | + spec = importlib.util.find_spec("nvidia") |
| 59 | + except (ImportError, ModuleNotFoundError, ValueError): |
| 60 | + return () |
| 61 | + if spec is None or spec.submodule_search_locations is None: |
| 62 | + return () |
| 63 | + |
| 64 | + paths: set[str] = set() |
| 65 | + for package_root in map(Path, spec.submodule_search_locations): |
| 66 | + try: |
| 67 | + components = tuple(package_root.iterdir()) |
| 68 | + except OSError: |
| 69 | + continue |
| 70 | + for component in components: |
| 71 | + for library_directory_name in ("lib", "bin"): |
| 72 | + candidate = component / library_directory_name |
| 73 | + if candidate.is_dir(): |
| 74 | + paths.add(str(candidate)) |
| 75 | + return tuple(sorted(paths)) |
| 76 | + |
| 77 | + |
| 78 | +def _nvidia_wheel_subprocess_environment( |
| 79 | + environment: Mapping[str, str], |
| 80 | +) -> dict[str, str]: |
| 81 | + """Prepend installed NVIDIA wheel libraries to a child environment.""" |
| 82 | + |
| 83 | + prepared = dict(environment) |
| 84 | + library_paths = _nvidia_wheel_library_paths() |
| 85 | + if not library_paths: |
| 86 | + return prepared |
| 87 | + |
| 88 | + search_variable = "PATH" if os.name == "nt" else "LD_LIBRARY_PATH" |
| 89 | + existing_paths = tuple( |
| 90 | + path for path in prepared.get(search_variable, "").split(os.pathsep) if path |
| 91 | + ) |
| 92 | + prepared[search_variable] = os.pathsep.join(dict.fromkeys((*library_paths, *existing_paths))) |
| 93 | + return prepared |
44 | 94 |
|
45 | 95 |
|
46 | 96 | class MemoryContractAttribute(str, Enum): |
@@ -501,6 +551,7 @@ class FrameworkRuntime: |
501 | 551 | dlpack_exporter: DLPackExporter | None = None |
502 | 552 | dlpack_validator: DLPackValidator = _protocol_dlpack |
503 | 553 | oom_matcher: OOMMatcher = _never_oom |
| 554 | + subprocess_environment: SubprocessEnvironmentResolver = _identity_subprocess_environment |
504 | 555 |
|
505 | 556 |
|
506 | 557 | class _MemoryTypeFields: |
@@ -571,6 +622,7 @@ def __new__( |
571 | 622 | dlpack_importer=_cupy_from_dlpack, |
572 | 623 | dlpack_exporter=_protocol_dlpack_export, |
573 | 624 | oom_matcher=_cupy_oom, |
| 625 | + subprocess_environment=_nvidia_wheel_subprocess_environment, |
574 | 626 | ), |
575 | 627 | CUPY_OPERATIONS, |
576 | 628 | ) |
@@ -682,6 +734,20 @@ def prepare_import(self) -> None: |
682 | 734 | for name, value in self.import_environment: |
683 | 735 | os.environ.setdefault(name, value) |
684 | 736 |
|
| 737 | + @classmethod |
| 738 | + def subprocess_environment( |
| 739 | + cls, |
| 740 | + environment: Mapping[str, str] | None = None, |
| 741 | + ) -> dict[str, str]: |
| 742 | + """Project framework import requirements into a child environment.""" |
| 743 | + |
| 744 | + prepared = dict(os.environ if environment is None else environment) |
| 745 | + for memory_type in cls: |
| 746 | + for name, value in memory_type.import_environment: |
| 747 | + prepared.setdefault(name, value) |
| 748 | + prepared = memory_type._runtime.subprocess_environment(prepared) |
| 749 | + return prepared |
| 750 | + |
685 | 751 | def loaded_module(self) -> Any | None: |
686 | 752 | """Return an already-loaded framework without causing an import.""" |
687 | 753 |
|
|
0 commit comments