Skip to content

Commit 49745fb

Browse files
committed
Declare child framework environments on memory types
1 parent a29bf9f commit 49745fb

9 files changed

Lines changed: 158 additions & 64 deletions

File tree

‎.github/workflows/publish.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ jobs:
4242
4343
- name: Verify release tag and main ancestry
4444
run: |
45-
declared_version=$(python -c 'import tomllib; print(tomllib.load(open("pyproject.toml", "rb"))["project"]["version"])')
45+
declared_version=$(python -c 'import arraybridge; print(arraybridge.__version__)')
4646
test "$GITHUB_REF_NAME" = "v$declared_version"
4747
git fetch origin main
4848
git merge-base --is-ancestor "$GITHUB_SHA" origin/main

‎docs/source/conf.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,15 @@
1616
sys.path.insert(0, str(EXTERNAL_ROOT / "metaclass-registry" / "src"))
1717
sys.path.insert(0, str(PACKAGE_ROOT / "src"))
1818

19+
from arraybridge import __version__
20+
1921
# -- Project information -----------------------------------------------------
2022
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
2123

2224
project = 'arraybridge'
2325
copyright = '2025, Tristan Simas'
2426
author = 'Tristan Simas'
25-
release = '0.3.1'
27+
release = __version__
2628

2729
# -- General configuration ---------------------------------------------------
2830
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

‎docs/source/gpu_features.rst‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,11 @@ overwritten. Hosts that import optional frameworks themselves should call
1515
``MemoryType.prepare_import()`` first. ArrayBridge warns when it encounters an
1616
already-loaded framework whose import-time defaults were absent.
1717

18+
``MemoryType.subprocess_environment()`` projects the same member-owned import
19+
defaults into a child-process environment. The CuPy member also derives native
20+
library search paths from installed NVIDIA wheels, so a fresh child interpreter
21+
does not depend on another framework having loaded CUDA libraries first.
22+
1823
JAX float64 output requires x64 mode to be enabled before import. ArrayBridge
1924
raises instead of silently returning float32 when a caller requests float64
2025
while that capability is disabled.

‎pyproject.toml‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
55

66
[project]
77
name = "arraybridge"
8-
version = "0.3.2"
8+
dynamic = ["version"]
99
description = "Unified API for NumPy, CuPy, PyTorch, TensorFlow, JAX, and pyclesperanto with automatic memory type conversion"
1010
authors = [{name = "Tristan Simas", email = "tristan.simas@mail.mcgill.ca"}]
1111
license = {text = "MIT"}
@@ -98,6 +98,9 @@ Issues = "https://github.com/OpenHCSDev/arraybridge/issues"
9898
[tool.hatch.build.targets.wheel]
9999
packages = ["src/arraybridge"]
100100

101+
[tool.hatch.version]
102+
path = "src/arraybridge/__init__.py"
103+
101104
[tool.pytest.ini_options]
102105
testpaths = ["tests"]
103106
python_files = ["test_*.py"]

‎scripts/update_version.py‎

Lines changed: 24 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#!/usr/bin/env python3
2-
"""Update ArrayBridge's three version projections without Git side effects."""
2+
"""Update ArrayBridge's canonical package version without Git side effects."""
33

44
import argparse
55
import re
@@ -8,39 +8,21 @@
88
from packaging.version import InvalidVersion, Version
99

1010
PROJECT_ROOT = Path(__file__).resolve().parents[1]
11-
VERSION_FILES = {
12-
PROJECT_ROOT
13-
/ "src/arraybridge/__init__.py": (
14-
r"(__version__\s*=\s*['\"])([^'\"]+)(['\"])",
15-
"package",
16-
),
17-
PROJECT_ROOT
18-
/ "pyproject.toml": (
19-
r"(^version\s*=\s*['\"])([^'\"]+)(['\"])",
20-
"project",
21-
),
22-
PROJECT_ROOT
23-
/ "docs/source/conf.py": (
24-
r"(^release\s*=\s*['\"])([^'\"]+)(['\"])",
25-
"documentation",
26-
),
27-
}
11+
VERSION_FILE = PROJECT_ROOT / "src/arraybridge/__init__.py"
12+
VERSION_PATTERN = r"(__version__\s*=\s*['\"])([^'\"]+)(['\"])"
2813

2914

3015
def current_version() -> Version:
31-
"""Return the agreed version, failing when a projection has drifted."""
32-
declared_versions: dict[str, Version] = {}
33-
for path, (pattern, projection_name) in VERSION_FILES.items():
34-
match = re.search(pattern, path.read_text(encoding="utf-8"), flags=re.MULTILINE)
35-
if match is None:
36-
raise RuntimeError(f"Version declaration is missing from {path}")
37-
declared_versions[projection_name] = Version(match.group(2))
16+
"""Return the package-owned canonical version."""
3817

39-
unique_versions = set(declared_versions.values())
40-
if len(unique_versions) != 1:
41-
declarations = ", ".join(f"{name}={version}" for name, version in declared_versions.items())
42-
raise RuntimeError(f"Version projections have drifted: {declarations}")
43-
return unique_versions.pop()
18+
match = re.search(
19+
VERSION_PATTERN,
20+
VERSION_FILE.read_text(encoding="utf-8"),
21+
flags=re.MULTILINE,
22+
)
23+
if match is None:
24+
raise RuntimeError(f"Version declaration is missing from {VERSION_FILE}")
25+
return Version(match.group(2))
4426

4527

4628
def update_version(version_text: str) -> None:
@@ -52,19 +34,18 @@ def update_version(version_text: str) -> None:
5234
if requested <= current:
5335
raise ValueError(f"New version {requested} must be greater than {current}")
5436

55-
for path, (pattern, _projection_name) in VERSION_FILES.items():
56-
content = path.read_text(encoding="utf-8")
57-
updated, count = re.subn(
58-
pattern,
59-
rf"\g<1>{requested}\g<3>",
60-
content,
61-
flags=re.MULTILINE,
62-
)
63-
if count != 1:
64-
raise RuntimeError(f"Expected one version declaration in {path}, found {count}")
65-
path.write_text(updated, encoding="utf-8")
66-
67-
print(f"Updated ArrayBridge version projections to {requested}")
37+
content = VERSION_FILE.read_text(encoding="utf-8")
38+
updated, count = re.subn(
39+
VERSION_PATTERN,
40+
rf"\g<1>{requested}\g<3>",
41+
content,
42+
flags=re.MULTILINE,
43+
)
44+
if count != 1:
45+
raise RuntimeError(f"Expected one version declaration in {VERSION_FILE}, found {count}")
46+
VERSION_FILE.write_text(updated, encoding="utf-8")
47+
48+
print(f"Updated ArrayBridge version to {requested}")
6849
print("Run scripts/verify_release_ready.py --allow-dirty, then review and commit.")
6950

7051

‎scripts/verify_release_ready.py‎

Lines changed: 2 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,6 @@
1616
import tempfile
1717
from pathlib import Path
1818

19-
import tomllib
20-
2119
PROJECT_ROOT = Path(__file__).resolve().parents[1]
2220

2321

@@ -41,20 +39,6 @@ def check_version():
4139
print(f" ❌ Version '{version}' doesn't follow semantic versioning (MAJOR.MINOR.PATCH)")
4240
return False
4341

44-
pyproject_version = tomllib.loads(
45-
(PROJECT_ROOT / "pyproject.toml").read_text(encoding="utf-8")
46-
)["project"]["version"]
47-
docs_match = re.search(
48-
r"release\s*=\s*['\"]([^'\"]+)",
49-
(PROJECT_ROOT / "docs/source/conf.py").read_text(encoding="utf-8"),
50-
)
51-
docs_version = docs_match.group(1) if docs_match else None
52-
if len({version, pyproject_version, docs_version}) != 1:
53-
print(
54-
" ❌ Version authorities disagree: "
55-
f"package={version}, pyproject={pyproject_version}, docs={docs_version}"
56-
)
57-
return False
5842
print(f" ✅ Version: {version}")
5943
return True
6044

@@ -70,7 +54,8 @@ def check_pyproject_toml():
7054
content = pyproject_file.read_text(encoding="utf-8")
7155
required_fields = {
7256
"name": r'name\s*=\s*["\']arraybridge["\']',
73-
"version": r"version\s*=",
57+
"dynamic version": r'dynamic\s*=\s*\[["\']version["\']\]',
58+
"version authority": r'path\s*=\s*["\']src/arraybridge/__init__\.py["\']',
7459
"description": r"description\s*=",
7560
"authors": r"authors\s*=",
7661
"build-backend": r'build-backend\s*=\s*["\']hatchling\.build["\']',

‎src/arraybridge/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
and unified utilities for working with multiple array/tensor frameworks.
66
"""
77

8-
__version__ = "0.3.2"
8+
__version__ = "0.3.3"
99

1010
from . import decorators as _decorators
1111
from .array_geometry import ArrayGeometry

‎src/arraybridge/types.py‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
from contextlib import AbstractContextManager, contextmanager, nullcontext
1515
from dataclasses import dataclass
1616
from enum import Enum
17+
from pathlib import Path
1718
from typing import Any, TypeVar, cast
1819

1920
from arraybridge.array_operations import (
@@ -41,6 +42,55 @@
4142
DLPackExporter = Callable[[Any, Any], Any | None]
4243
DLPackValidator = Callable[[Any, Any], bool]
4344
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
4494

4595

4696
class MemoryContractAttribute(str, Enum):
@@ -501,6 +551,7 @@ class FrameworkRuntime:
501551
dlpack_exporter: DLPackExporter | None = None
502552
dlpack_validator: DLPackValidator = _protocol_dlpack
503553
oom_matcher: OOMMatcher = _never_oom
554+
subprocess_environment: SubprocessEnvironmentResolver = _identity_subprocess_environment
504555

505556

506557
class _MemoryTypeFields:
@@ -571,6 +622,7 @@ def __new__(
571622
dlpack_importer=_cupy_from_dlpack,
572623
dlpack_exporter=_protocol_dlpack_export,
573624
oom_matcher=_cupy_oom,
625+
subprocess_environment=_nvidia_wheel_subprocess_environment,
574626
),
575627
CUPY_OPERATIONS,
576628
)
@@ -682,6 +734,20 @@ def prepare_import(self) -> None:
682734
for name, value in self.import_environment:
683735
os.environ.setdefault(name, value)
684736

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+
685751
def loaded_module(self) -> Any | None:
686752
"""Return an already-loaded framework without causing an import."""
687753

‎tests/test_types.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,58 @@ def test_import_preparation_sets_only_declaration_owned_defaults(self, monkeypat
145145
assert os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] == "true"
146146
assert os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] == "custom"
147147

148+
def test_subprocess_environment_derives_nvidia_wheel_paths_from_cupy_owner(
149+
self,
150+
monkeypatch,
151+
tmp_path,
152+
):
153+
from arraybridge import types
154+
155+
nvidia_root = tmp_path / "nvidia"
156+
cublas_lib = nvidia_root / "cublas" / "lib"
157+
runtime_bin = nvidia_root / "cuda_runtime" / "bin"
158+
cublas_lib.mkdir(parents=True)
159+
runtime_bin.mkdir(parents=True)
160+
monkeypatch.setattr(
161+
types.importlib.util,
162+
"find_spec",
163+
lambda name: (
164+
SimpleNamespace(submodule_search_locations=(str(nvidia_root),))
165+
if name == "nvidia"
166+
else None
167+
),
168+
)
169+
search_variable = "PATH" if os.name == "nt" else "LD_LIBRARY_PATH"
170+
original = {search_variable: "/existing", "CUSTOM": "value"}
171+
172+
prepared = MemoryType.subprocess_environment(original)
173+
174+
assert original == {search_variable: "/existing", "CUSTOM": "value"}
175+
assert prepared[search_variable].split(os.pathsep) == [
176+
*sorted((str(cublas_lib), str(runtime_bin))),
177+
"/existing",
178+
]
179+
assert prepared["CUSTOM"] == "value"
180+
assert prepared["TF_FORCE_GPU_ALLOW_GROWTH"] == "true"
181+
assert prepared["XLA_PYTHON_CLIENT_PREALLOCATE"] == "false"
182+
183+
def test_subprocess_environment_preserves_host_values_without_nvidia_wheels(
184+
self,
185+
monkeypatch,
186+
):
187+
from arraybridge import types
188+
189+
monkeypatch.setattr(types.importlib.util, "find_spec", lambda name: None)
190+
original = {
191+
"TF_FORCE_GPU_ALLOW_GROWTH": "host",
192+
"XLA_PYTHON_CLIENT_PREALLOCATE": "host",
193+
}
194+
195+
prepared = MemoryType.subprocess_environment(original)
196+
197+
assert prepared == original
198+
assert prepared is not original
199+
148200
def test_preloaded_framework_reports_late_import_preparation(
149201
self,
150202
monkeypatch,

0 commit comments

Comments
 (0)