Coverage for src/arraybridge/types.py: 83%

1"""

2Memory type definitions for arraybridge.

4This module defines the MemoryType enum and related constants for managing

5different array/tensor frameworks.

6"""

8import importlib

9import importlib.util

10import logging

11import os

12import sys

13from collections.abc import Callable, Iterator, Mapping, MutableMapping

14from contextlib import AbstractContextManager, contextmanager, nullcontext

15from dataclasses import dataclass

16from enum import Enum

17from pathlib import Path

18from typing import Any, TypeVar, cast

20from arraybridge.array_operations import (

21 CUPY_OPERATIONS,

22 JAX_OPERATIONS,

23 NUMPY_OPERATIONS,

24 PYCLESPERANTO_OPERATIONS,

25 TENSORFLOW_OPERATIONS,

26 TORCH_OPERATIONS,

27 ArrayOperations,

28)

30T = TypeVar("T")

31logger = logging.getLogger(__name__)

32ConversionFunc = Callable[[Any], Any]

33DeviceIdsResolver = Callable[[Any], tuple[int, ...]]

34DeviceIdResolver = Callable[[Any, Any], int | None]

35DeviceScopeFactory = Callable[[Any, int], AbstractContextManager[None]]

36DeviceActivator = Callable[[Any, int], None]

37FrameworkCleanup = Callable[[Any], None]

38ActiveDeviceMover = Callable[[Any, Any, int], Any]

39CurrentDeviceResolver = Callable[[Any], int | None]

40StreamFactory = Callable[[Any], Any]

41StreamScopeFactory = Callable[[Any, Any], AbstractContextManager[None]]

42DLPackExporter = Callable[[Any, Any], Any | None]

43DLPackValidator = Callable[[Any, Any], bool]

44OOMMatcher = Callable[[BaseException, Any | None], bool]

45SubprocessEnvironmentResolver = Callable[[Mapping[str, str]], dict[str, str]]

48def _identity_subprocess_environment(

49 environment: Mapping[str, str],

50) -> dict[str, str]:

51 return dict(environment)

54def _nvidia_wheel_library_paths() -> tuple[str, ...]:

55 """Return native-library directories declared by installed NVIDIA wheels."""

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 ()

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))

78def _nvidia_wheel_subprocess_environment(

79 environment: Mapping[str, str],

80) -> dict[str, str]:

81 """Prepend installed NVIDIA wheel libraries to a child environment."""

83 prepared = dict(environment)

84 library_paths = _nvidia_wheel_library_paths()

85 if not library_paths:

86 return prepared

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

96class MemoryContractAttribute(str, Enum):

97 """Callable metadata keys owned by ArrayBridge's decorator contract."""

99 INPUT = "input_memory_type"

100 OUTPUT = "output_memory_type"

101 EXECUTION = "execution_memory_type"

102

103 def read(self, namespace: Any, default: Any = None) -> Any:

104 """Read this declaration from an attribute or mapping namespace."""

105

106 if isinstance(namespace, Mapping):

107 return namespace.get(self.value, default)

108 return getattr(namespace, self.value, default)

109

110 def write(self, namespace: Any, value: Any) -> None:

111 """Write this declaration to an attribute or mutable mapping namespace."""

112

113 if isinstance(namespace, MutableMapping):

114 namespace[self.value] = value

115 return

116 setattr(namespace, self.value, value)

117

118

119@dataclass(frozen=True, slots=True)

120class DLPackPayload:

121 """One exported capsule plus its original protocol-bearing array."""

122

123 source: Any

124 capsule: Any

125

126

127DLPackImporter = Callable[[DLPackPayload, Any], Any]

128

129

130def _no_device_ids(module: Any) -> tuple[int, ...]:

131 del module

132 return ()

133

134

135def _no_device_id(data: Any, module: Any) -> None:

136 del data, module

137 return None

138

139

140def _cupy_device_id(data: Any, module: Any) -> int:

141 del module

142 return int(data.device.id)

143

144

145def _torch_device_id(data: Any, module: Any) -> int | None:

146 del module

147 return int(data.device.index) if data.is_cuda else None

148

149

150def _tensorflow_device_id(data: Any, module: Any) -> int | None:

151 del module

152 device = data.device.lower()

153 return int(device.rsplit(":", maxsplit=1)[-1]) if "gpu" in device else None

154

155

156def _jax_device_id(data: Any, module: Any) -> int | None:

157 device = data.device

158 device = device() if callable(device) else device

159 if getattr(device, "platform", None) != "gpu":

160 return None

161 gpu_devices = tuple(candidate for candidate in module.devices() if candidate.platform == "gpu")

162 identity_match = next(

163 (index for index, candidate in enumerate(gpu_devices) if candidate is device),

164 None,

165 )

166 if identity_match is not None:

167 return identity_match

168 return next(

169 (index for index, candidate in enumerate(gpu_devices) if str(candidate) == str(device)),

170 None,

171 )

172

173

174def _pyclesperanto_device_id(data: Any, module: Any) -> int | None:

175 declared_device = getattr(data, "device", None)

176 if declared_device is None:

177 declared_device = module.get_device()

178 declared_selector = getattr(declared_device, "name", None)

179 if declared_selector is None:

180 declared_selector = str(declared_device)

181 devices = _pyclesperanto_devices(module)

182 return next(

183 (

184 index

185 for index, device in enumerate(devices)

186 if getattr(device, "name", str(device)) == declared_selector

187 ),

188 None,

189 )

190

191

192def _cupy_device_ids(module: Any) -> tuple[int, ...]:

193 return tuple(range(int(module.cuda.runtime.getDeviceCount())))

194

195

196def _torch_device_ids(module: Any) -> tuple[int, ...]:

197 if not module.cuda.is_available():

198 return ()

199 return tuple(range(int(module.cuda.device_count())))

200

201

202def _tensorflow_device_ids(module: Any) -> tuple[int, ...]:

203 return tuple(range(len(module.config.list_logical_devices("GPU"))))

204

205

206def _jax_device_ids(module: Any) -> tuple[int, ...]:

207 return tuple(

208 range(len(tuple(device for device in module.devices() if device.platform == "gpu")))

209 )

210

211

212def _pyclesperanto_devices(module: Any) -> tuple[Any, ...]:

213 try:

214 return tuple(module.list_available_devices("gpu"))

215 except TypeError:

216 return tuple(module.list_available_devices())

217

218

219def _pyclesperanto_select_device(

220 module: Any,

221 selector: str | int,

222 device_type: str | None = None,

223) -> None:

224 if device_type is None:

225 module.select_device(selector)

226 return

227 try:

228 module.select_device(selector, device_type)

229 except TypeError:

230 module.select_device(selector)

231

232

233def _pyclesperanto_device_ids(module: Any) -> tuple[int, ...]:

234 return tuple(range(len(_pyclesperanto_devices(module))))

235

236

237def _null_device_scope(module: Any, device_id: int) -> AbstractContextManager[None]:

238 del module, device_id

239 return nullcontext()

240

241

242def _cupy_device_scope(module: Any, device_id: int) -> AbstractContextManager[None]:

243 return cast(AbstractContextManager[None], module.cuda.Device(device_id))

244

245

246def _torch_device_scope(module: Any, device_id: int) -> AbstractContextManager[None]:

247 return cast(AbstractContextManager[None], module.cuda.device(device_id))

248

249

250def _tensorflow_device_scope(

251 module: Any,

252 device_id: int,

253) -> AbstractContextManager[None]:

254 return cast(AbstractContextManager[None], module.device(f"/device:GPU:{device_id}"))

255

256

257def _jax_device_scope(module: Any, device_id: int) -> AbstractContextManager[None]:

258 gpu_devices = tuple(device for device in module.devices() if device.platform == "gpu")

259 return cast(AbstractContextManager[None], module.default_device(gpu_devices[device_id]))

260

261

262@contextmanager

263def _pyclesperanto_device_scope(module: Any, device_id: int) -> Iterator[None]:

264 current_device = module.get_device()

265 current_id = _pyclesperanto_device_id(None, module)

266 current_selector = getattr(current_device, "name", None)

267 if current_selector is None:

268 current_selector = str(current_device)

269 _pyclesperanto_select_device(module, device_id, "gpu")

270 try:

271 yield

272 finally:

273 if current_id != device_id:

274 _pyclesperanto_select_device(module, current_selector)

275

276

277def _no_device_activation(module: Any, device_id: int) -> None:

278 del module, device_id

279

280

281def _cupy_device_activation(module: Any, device_id: int) -> None:

282 module.cuda.Device(device_id).use()

283

284

285def _torch_device_activation(module: Any, device_id: int) -> None:

286 module.cuda.set_device(device_id)

287

288

289def _pyclesperanto_device_activation(module: Any, device_id: int) -> None:

290 _pyclesperanto_select_device(module, device_id, "gpu")

291

292

293def _cupy_cleanup(module: Any) -> None:

294 module.get_default_memory_pool().free_all_blocks()

295 module.get_default_pinned_memory_pool().free_all_blocks()

296 module.cuda.runtime.deviceSynchronize()

297

298

299def _torch_cleanup(module: Any) -> None:

300 module.cuda.empty_cache()

301 module.cuda.synchronize()

302

303

304def _identity_device_move(data: Any, module: Any, device_id: int) -> Any:

305 del module, device_id

306 return data

307

308

309def _cupy_device_move(data: Any, module: Any, device_id: int) -> Any:

310 del module, device_id

311 return data.copy()

312

313

314def _torch_device_move(data: Any, module: Any, device_id: int) -> Any:

315 del module

316 return data.to(f"cuda:{device_id}")

317

318

319def _tensorflow_device_move(data: Any, module: Any, device_id: int) -> Any:

320 del device_id

321 return module.identity(data)

322

323

324def _jax_device_move(data: Any, module: Any, device_id: int) -> Any:

325 gpu_devices = tuple(device for device in module.devices() if device.platform == "gpu")

326 return module.device_put(data, gpu_devices[device_id])

327

328

329def _pyclesperanto_device_move(data: Any, module: Any, device_id: int) -> Any:

330 del device_id

331 result = module.create_like(data)

332 module.copy(data, result)

333 return result

334

335

336def _no_current_device(module: Any) -> None:

337 del module

338 return None

339

340

341def _cupy_current_device(module: Any) -> int:

342 return int(module.cuda.runtime.getDevice())

343

344

345def _torch_current_device(module: Any) -> int | None:

346 return int(module.cuda.current_device()) if module.cuda.is_available() else None

347

348

349def _cupy_stream(module: Any) -> Any:

350 return module.cuda.Stream()

351

352

353def _torch_stream(module: Any) -> Any:

354 return module.cuda.Stream()

355

356

357def _cupy_stream_scope(module: Any, stream: Any) -> AbstractContextManager[None]:

358 del module

359 return cast(AbstractContextManager[None], stream)

360

361

362def _torch_stream_scope(module: Any, stream: Any) -> AbstractContextManager[None]:

363 return cast(AbstractContextManager[None], module.cuda.stream(stream))

364

365

366def _null_stream_scope(module: Any, stream: Any) -> AbstractContextManager[None]:

367 del module, stream

368 return nullcontext()

369

370

371def _has_modern_dlpack_protocol(data: Any) -> bool:

372 return callable(getattr(data, "__dlpack__", None)) and callable(

373 getattr(data, "__dlpack_device__", None)

374 )

375

376

377def _cupy_from_dlpack(payload: DLPackPayload, module: Any) -> Any:

378 if _has_modern_dlpack_protocol(payload.source):

379 return module.from_dlpack(payload.source)

380 legacy_importer = getattr(module, "fromDlpack", None)

381 return NotImplemented if legacy_importer is None else legacy_importer(payload.capsule)

382

383

384def _torch_from_dlpack(payload: DLPackPayload, module: Any) -> Any:

385 return module.from_dlpack(payload.capsule)

386

387

388def _export_dlpack(data: Any) -> Any:

389 for attribute in ("__dlpack__", "to_dlpack", "toDlpack"):

390 exporter = getattr(data, attribute, None)

391 if callable(exporter):

392 return exporter()

393 raise TypeError(f"{type(data).__name__} does not expose a DLPack exporter")

394

395

396def _tensorflow_from_dlpack(payload: DLPackPayload, module: Any) -> Any:

397 return module.experimental.dlpack.from_dlpack(payload.capsule)

398

399

400def _dlpack_data_pointer(data: Any) -> int | None:

401 """Return a protocol source pointer when its framework exposes one."""

402

403 data_ptr = getattr(data, "data_ptr", None)

404 if callable(data_ptr):

405 return int(data_ptr())

406 cuda_interface = getattr(data, "__cuda_array_interface__", None)

407 if isinstance(cuda_interface, dict):

408 pointer, _read_only = cuda_interface.get("data", (None, False))

409 return None if pointer is None else int(pointer)

410 return None

411

412

413def _jax_aligned_dlpack_source(data: Any) -> Any:

414 """Copy a DLPack source only when JAX's 16-byte alignment is unmet."""

415

416 pointer = _dlpack_data_pointer(data)

417 if pointer is None or pointer % 16 == 0:

418 return data

419 for attribute in ("clone", "copy"):

420 copier = getattr(data, attribute, None)

421 if callable(copier):

422 return copier()

423 return data

424

425

426def _jax_from_dlpack(payload: DLPackPayload, module: Any) -> Any:

427 if not _has_modern_dlpack_protocol(payload.source):

428 return NotImplemented

429 return module.dlpack.from_dlpack(_jax_aligned_dlpack_source(payload.source))

430

431

432def _protocol_dlpack(data: Any, module: Any) -> bool:

433 del module

434 return any(

435 callable(getattr(data, attribute, None))

436 for attribute in ("__dlpack__", "toDlpack", "to_dlpack")

437 )

438

439

440def _protocol_dlpack_export(data: Any, module: Any) -> Any | None:

441 del module

442 try:

443 return _export_dlpack(data)

444 except TypeError:

445 return None

446

447

448def _tensorflow_dlpack(data: Any, module: Any) -> bool:

449 try:

450 major, minor = map(int, module.__version__.split(".")[:2])

451 except (AttributeError, TypeError, ValueError):

452 return False

453 if (major, minor) < (2, 12):

454 return False

455 if "gpu" not in str(getattr(data, "device", "")).lower():

456 return False

457 dlpack = getattr(getattr(module, "experimental", None), "dlpack", None)

458 if not callable(getattr(dlpack, "to_dlpack", None)):

459 return False

460 return True

461

462

463def _tensorflow_dlpack_export(data: Any, module: Any) -> Any | None:

464 if not _tensorflow_dlpack(data, module):

465 return None

466 return module.experimental.dlpack.to_dlpack(data)

467

468

469def _message_matches(error: BaseException, *patterns: str) -> bool:

470 message = str(error).lower()

471 return any(pattern in message for pattern in patterns)

472

473

474def _exception_type(module: Any, *path: str) -> type[BaseException] | None:

475 candidate = module

476 for attribute in path:

477 candidate = getattr(candidate, attribute, None)

478 if candidate is None:

479 return None

480 return (

481 candidate if isinstance(candidate, type) and issubclass(candidate, BaseException) else None

482 )

483

484

485def _matches_declared_exception(error: BaseException, module: Any | None, *path: str) -> bool:

486 if module is None:

487 return False

488 exception_type = _exception_type(module, *path)

489 return exception_type is not None and isinstance(error, exception_type)

490

491

492def _numpy_oom(error: BaseException, module: Any | None) -> bool:

493 del module

494 return _message_matches(error, "cannot allocate memory", "memory exhausted")

495

496

497def _never_oom(error: BaseException, module: Any | None) -> bool:

498 del error, module

499 return False

500

501

502def _cupy_oom(error: BaseException, module: Any | None) -> bool:

503 return (

504 _matches_declared_exception(error, module, "cuda", "memory", "OutOfMemoryError")

505 or _matches_declared_exception(error, module, "cuda", "runtime", "CUDARuntimeError")

506 or _message_matches(error, "out of memory", "cuda_error_out_of_memory")

507 )

508

509

510def _torch_oom(error: BaseException, module: Any | None) -> bool:

511 return _matches_declared_exception(

512 error, module, "cuda", "OutOfMemoryError"

513 ) or _message_matches(error, "out of memory", "cuda_error_out_of_memory")

514

515

516def _tensorflow_oom(error: BaseException, module: Any | None) -> bool:

517 return _matches_declared_exception(

518 error, module, "errors", "ResourceExhaustedError"

519 ) or _message_matches(error, "out of memory", "resource_exhausted")

520

521

522def _jax_oom(error: BaseException, module: Any | None) -> bool:

523 del module

524 return _message_matches(error, "out of memory", "oom when allocating", "allocation failure")

525

526

527def _pyclesperanto_oom(error: BaseException, module: Any | None) -> bool:

528 del module

529 return _message_matches(

530 error,

531 "cl_mem_object_allocation_failure",

532 "cl_out_of_resources",

533 "out of memory",

534 )

535

536

537@dataclass(frozen=True, slots=True)

538class FrameworkRuntime:

539 """Typed execution leaves carried by one ``MemoryType`` declaration."""

540

541 device_ids: DeviceIdsResolver = _no_device_ids

542 device_id: DeviceIdResolver = _no_device_id

543 device_scope: DeviceScopeFactory = _null_device_scope

544 activate_device: DeviceActivator = _no_device_activation

545 cleanup: FrameworkCleanup | None = None

546 move_to_active_device: ActiveDeviceMover = _identity_device_move

547 current_device: CurrentDeviceResolver = _no_current_device

548 stream_factory: StreamFactory | None = None

549 stream_scope: StreamScopeFactory = _null_stream_scope

550 dlpack_importer: DLPackImporter | None = None

551 dlpack_exporter: DLPackExporter | None = None

552 dlpack_validator: DLPackValidator = _protocol_dlpack

553 oom_matcher: OOMMatcher = _never_oom

554 subprocess_environment: SubprocessEnvironmentResolver = _identity_subprocess_environment

555

556

557class _MemoryTypeFields:

558 import_name: str

559 display_name: str

560 is_gpu: bool

561 module_aliases: tuple[str, ...]

562 import_environment: tuple[tuple[str, str], ...]

563 _runtime: FrameworkRuntime

564 _operations: ArrayOperations

565

566

567class MemoryType(_MemoryTypeFields, Enum):

568 """Array-framework declarations with member-owned runtime capability leaves."""

569

570 def __new__(

571 cls,

572 value: str,

573 *declaration: Any,

574 ) -> "MemoryType":

575 (

576 import_name,

577 display_name,

578 is_gpu,

579 module_aliases,

580 import_environment,

581 runtime,

582 operations,

583 ) = declaration

584 member = object.__new__(cls)

585 member._value_ = value

586 member.import_name = cast(str, import_name)

587 member.display_name = cast(str, display_name)

588 member.is_gpu = cast(bool, is_gpu)

589 member.module_aliases = cast(tuple[str, ...], module_aliases)

590 member.import_environment = cast(tuple[tuple[str, str], ...], import_environment)

591 member._runtime = cast(FrameworkRuntime, runtime)

592 member._operations = cast(ArrayOperations, operations)

593 return member

594

595 NUMPY = (

596 "numpy",

597 "numpy",

598 "NumPy",

599 False,

600 (),

601 (),

602 FrameworkRuntime(oom_matcher=_numpy_oom),

603 NUMPY_OPERATIONS,

604 )

605 CUPY = (

606 "cupy",

607 "cupy",

608 "CuPy",

609 True,

610 (),

611 (),

612 FrameworkRuntime(

613 device_ids=_cupy_device_ids,

614 device_id=_cupy_device_id,

615 device_scope=_cupy_device_scope,

616 activate_device=_cupy_device_activation,

617 cleanup=_cupy_cleanup,

618 move_to_active_device=_cupy_device_move,

619 current_device=_cupy_current_device,

620 stream_factory=_cupy_stream,

621 stream_scope=_cupy_stream_scope,

622 dlpack_importer=_cupy_from_dlpack,

623 dlpack_exporter=_protocol_dlpack_export,

624 oom_matcher=_cupy_oom,

625 subprocess_environment=_nvidia_wheel_subprocess_environment,

626 ),

627 CUPY_OPERATIONS,

628 )

629 TORCH = (

630 "torch",

631 "torch",

632 "PyTorch",

633 True,

634 (),

635 (),

636 FrameworkRuntime(

637 device_ids=_torch_device_ids,

638 device_id=_torch_device_id,

639 device_scope=_torch_device_scope,

640 activate_device=_torch_device_activation,

641 cleanup=_torch_cleanup,

642 move_to_active_device=_torch_device_move,

643 current_device=_torch_current_device,

644 stream_factory=_torch_stream,

645 stream_scope=_torch_stream_scope,

646 dlpack_importer=_torch_from_dlpack,

647 dlpack_exporter=_protocol_dlpack_export,

648 oom_matcher=_torch_oom,

649 ),

650 TORCH_OPERATIONS,

651 )

652 TENSORFLOW = (

653 "tensorflow",

654 "tensorflow",

655 "TensorFlow",

656 True,

657 (),

658 (("TF_FORCE_GPU_ALLOW_GROWTH", "true"),),

659 FrameworkRuntime(

660 device_ids=_tensorflow_device_ids,

661 device_id=_tensorflow_device_id,

662 device_scope=_tensorflow_device_scope,

663 move_to_active_device=_tensorflow_device_move,

664 dlpack_importer=_tensorflow_from_dlpack,

665 dlpack_exporter=_tensorflow_dlpack_export,

666 dlpack_validator=_tensorflow_dlpack,

667 oom_matcher=_tensorflow_oom,

668 ),

669 TENSORFLOW_OPERATIONS,

670 )

671 JAX = (

672 "jax",

673 "jax",

674 "JAX",

675 True,

676 ("jaxlib",),

677 (("XLA_PYTHON_CLIENT_PREALLOCATE", "false"),),

678 FrameworkRuntime(

679 device_ids=_jax_device_ids,

680 device_id=_jax_device_id,

681 device_scope=_jax_device_scope,

682 move_to_active_device=_jax_device_move,

683 dlpack_importer=_jax_from_dlpack,

684 dlpack_exporter=_protocol_dlpack_export,

685 oom_matcher=_jax_oom,

686 ),

687 JAX_OPERATIONS,

688 )

689 PYCLESPERANTO = (

690 "pyclesperanto",

691 "pyclesperanto",

692 "pyclesperanto",

693 True,

694 (),

695 (),

696 FrameworkRuntime(

697 device_ids=_pyclesperanto_device_ids,

698 device_id=_pyclesperanto_device_id,

699 device_scope=_pyclesperanto_device_scope,

700 activate_device=_pyclesperanto_device_activation,

701 move_to_active_device=_pyclesperanto_device_move,

702 oom_matcher=_pyclesperanto_oom,

703 ),

704 PYCLESPERANTO_OPERATIONS,

705 )

706

707 @property

708 def recognized_module_names(self) -> frozenset[str]:

709 """Return top-level module names owned by this framework declaration."""

710

711 return frozenset((self.import_name, *self.module_aliases))

712

713 def is_installed(self) -> bool:

714 """Check package presence without importing the optional framework."""

715

716 try:

717 return importlib.util.find_spec(self.import_name) is not None

718 except (ImportError, ModuleNotFoundError, ValueError):

719 return False

720

721 def prepare_import(self) -> None:

722 """Apply declaration-owned coexistence defaults before framework import."""

723

724 missing_names = tuple(

725 name for name, _value in self.import_environment if name not in os.environ

726 )

727 if missing_names and self.loaded_module() is not None:

728 logger.warning(

729 "%s was loaded before its ArrayBridge import defaults were set; "

730 "the current process cannot guarantee those import-time settings: %s",

731 self.display_name,

732 ", ".join(missing_names),

733 )

734 for name, value in self.import_environment:

735 os.environ.setdefault(name, value)

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

751 def loaded_module(self) -> Any | None:

752 """Return an already-loaded framework without causing an import."""

753

754 return sys.modules.get(self.import_name)

755

756 def import_module(self) -> Any:

757 """Import this declaration's framework module."""

758

759 self.prepare_import()

760 return importlib.import_module(self.import_name)

761

762 def import_if_installed(self) -> Any | None:

763 """Import this framework only when its package is present."""

764

765 loaded = self.loaded_module()

766 if loaded is not None:

767 self.prepare_import()

768 return loaded

769 if not self.is_installed():

770 return None

771 return self.import_module()

772

773 def to_numpy(self, data: Any, module: Any | None = None) -> Any:

774 """Project one array to NumPy through this declaration's leaf."""

775

776 framework = module if module is not None else self.import_module()

777 return self._operations.to_numpy(data, framework)

778

779 def from_numpy(

780 self,

781 data: Any,

782 device_id: int,

783 module: Any | None = None,

784 ) -> Any:

785 """Create one array from NumPy on a declared framework-local device."""

786

787 framework = module if module is not None else self.import_module()

788 with self.device_scope(device_id, framework):

789 return self._operations.from_numpy(data, framework, device_id)

790

791 def stack_arrays(

792 self,

793 arrays: list[Any],

794 device_id: int,

795 module: Any | None = None,

796 ) -> Any:

797 """Stack prepared arrays on a declared framework-local device."""

798

799 framework = module if module is not None else self.import_module()

800 with self.device_scope(device_id, framework):

801 return self._operations.stack(arrays, framework)

802

803 def scale_dtype(

804 self,

805 data: Any,

806 target_dtype: Any,

807 module: Any | None = None,

808 ) -> Any:

809 """Scale an array through this declaration's typed operation leaf."""

810

811 framework = module if module is not None else self.import_if_installed()

812 if framework is None:

813 return data

814 device_id = self.device_id_of(data, framework) if hasattr(data, "dtype") else None

815 scope = nullcontext() if device_id is None else self.device_scope(device_id, framework)

816 with scope:

817 return self._operations.scale_dtype(data, target_dtype, framework)

818

819 def canonical_dtype_name(self, dtype: Any) -> str:

820 """Return this framework member's portable dtype identity."""

821

822 return self._operations.dtype_name(dtype)

823

824 def astype(

825 self,

826 data: Any,

827 dtype: Any,

828 module: Any | None = None,

829 ) -> Any:

830 """Cast one array through this framework member's operation leaf."""

831

832 framework = module if module is not None else self.import_module()

833 device_id = self.device_id_of(data, framework)

834 scope = nullcontext() if device_id is None else self.device_scope(device_id, framework)

835 with scope:

836 return self._operations.cast(data, dtype, framework)

837

838 def logical_and(

839 self,

840 left: Any,

841 right: Any,

842 module: Any | None = None,

843 ) -> Any:

844 """Intersect two arrays through this framework member's operation leaf."""

845

846 framework = module if module is not None else self.import_module()

847 device_id = self.device_id_of(left, framework)

848 scope = nullcontext() if device_id is None else self.device_scope(device_id, framework)

849 with scope:

850 return self._operations.logical_and(left, right, framework)

851

852 def available_device_ids(self, module: Any | None = None) -> tuple[int, ...]:

853 """Return every framework-local GPU device identifier."""

854

855 if not self.is_gpu:

856 return ()

857 framework = module if module is not None else self.import_if_installed()

858 return () if framework is None else self._runtime.device_ids(framework)

859

860 def require_device(self, device_id: int, module: Any | None = None) -> Any:

861 """Return the framework after proving that its local device exists."""

862

863 if not self.is_gpu:

864 return module

865 framework = module if module is not None else self.import_module()

866 available = self.available_device_ids(framework)

867 if device_id not in available:

868 raise ValueError(

869 f"{self.display_name} device {device_id} is unavailable; "

870 f"available device IDs are {available}"

871 )

872 return framework

873

874 def device_id_of(self, data: Any, module: Any | None = None) -> int | None:

875 """Return this framework's local device identifier for an array."""

876

877 if not self.is_gpu:

878 return None

879 framework = module if module is not None else self.import_module()

880 return self._runtime.device_id(data, framework)

881

882 def device_scope(

883 self,

884 device_id: int,

885 module: Any | None = None,

886 ) -> AbstractContextManager[None]:

887 """Return this framework member's scoped device activation leaf."""

888

889 if not self.is_gpu:

890 return nullcontext()

891 framework = self.require_device(device_id, module)

892 return self._runtime.device_scope(framework, device_id)

893

894 def activate_device(self, device_id: int, module: Any | None = None) -> None:

895 """Activate a process-global device where the framework supports it."""

896

897 if not self.is_gpu:

898 return

899 framework = self.require_device(device_id, module)

900 self._runtime.activate_device(framework, device_id)

901

902 def move_to_device(

903 self,

904 data: Any,

905 device_id: int,

906 module: Any | None = None,

907 ) -> Any:

908 """Move an array to one framework-local device without leaking selection."""

909

910 if not self.is_gpu:

911 return data

912 framework = self.require_device(device_id, module)

913 if self.device_id_of(data, framework) == device_id:

914 return data

915 with self.device_scope(device_id, framework):

916 return self._runtime.move_to_active_device(data, framework, device_id)

917

918 @property

919 def supports_dlpack(self) -> bool:

920 """Whether this framework declares a DLPack import leaf."""

921

922 return self._runtime.dlpack_importer is not None

923

924 def supports_dlpack_data(self, data: Any, module: Any | None = None) -> bool:

925 """Validate DLPack export for an array owned by this framework."""

926

927 if self._runtime.dlpack_exporter is None:

928 return False

929 framework = module if module is not None else self.import_module()

930 return self._runtime.dlpack_validator(data, framework)

931

932 def export_dlpack(self, data: Any, module: Any | None = None) -> DLPackPayload | None:

933 """Export one array through this framework's declaration, if supported."""

934

935 exporter = self._runtime.dlpack_exporter

936 if exporter is None:

937 return None

938 framework = module if module is not None else self.import_module()

939 capsule = exporter(data, framework)

940 return None if capsule is None else DLPackPayload(source=data, capsule=capsule)

941

942 def from_dlpack(

943 self,

944 data: DLPackPayload | Any,

945 module: Any | None = None,

946 ) -> Any:

947 """Import DLPack data through this framework's declared leaf."""

948

949 importer = self._runtime.dlpack_importer

950 if importer is None:

951 raise NotImplementedError(f"DLPack not supported for {self.value}")

952 framework = module if module is not None else self.import_module()

953 payload = data if isinstance(data, DLPackPayload) else DLPackPayload(data, data)

954 return importer(payload, framework)

955

956 def convert_to(self, data: Any, target: "MemoryType", device_id: int) -> Any:

957 """Convert one array through the source and target declarations."""

958

959 if self is target:

960 return target.move_to_device(data, device_id)

961

962 if target.supports_dlpack:

963 try:

964 payload = self.export_dlpack(data)

965 if payload is not None:

966 module = target.import_module()

967 with target.device_scope(device_id, module):

968 result = target.from_dlpack(payload, module)

969 if result is not NotImplemented:

970 return target.move_to_device(result, device_id, module)

971 except Exception as error:

972 logger.warning(

973 "DLPack conversion from %s to %s failed: %s. Using CPU roundtrip.",

974 self.value,

975 target.value,

976 error,

977 )

978

979 return target.from_numpy(self.to_numpy(data), device_id)

980

981 def current_device_id(self, module: Any | None = None) -> int | None:

982 """Return the framework-local process device used for new streams."""

983

984 if not self.is_gpu:

985 return None

986 framework = module if module is not None else self.import_module()

987 return self._runtime.current_device(framework)

988

989 def create_stream(self, module: Any | None = None) -> Any | None:

990 """Create a stream on the framework's current device when supported."""

991

992 factory = self._runtime.stream_factory

993 if factory is None:

994 return None

995 framework = module if module is not None else self.import_module()

996 return factory(framework)

997

998 def stream_scope(

999 self,

1000 stream: Any,

1001 module: Any | None = None,

1002 ) -> AbstractContextManager[None]:

1003 """Return this framework member's execution scope for one stream."""

1004

1005 if stream is None:

1006 return nullcontext()

1007 framework = module if module is not None else self.import_module()

1008 return self._runtime.stream_scope(framework, stream)

1009

1010 def is_oom_error(self, error: BaseException) -> bool:

1011 """Classify an error without importing an optional framework."""

1012

1013 return self._runtime.oom_matcher(error, self.loaded_module())

1014

1015 def cleanup_loaded(self, device_id: int | None = None) -> None:

1016 """Clean an already-loaded GPU framework without importing absent modules."""

1017

1018 if not self.is_gpu:

1019 return

1020 framework = self.loaded_module()

1021 if framework is None:

1022 return

1023 cleanup = self._runtime.cleanup

1024 if cleanup is None:

1025 return

1026 available = self.available_device_ids(framework)

1027 if not available:

1028 return

1029 targets = available if device_id is None else (device_id,)

1030 for target in targets:

1031 with self.device_scope(target, framework):

1032 cleanup(framework)

1033

1034

1035# Memory type sets

1036CPU_MEMORY_TYPES: frozenset[MemoryType] = frozenset(

1037 memory_type for memory_type in MemoryType if not memory_type.is_gpu

1038)

1039GPU_MEMORY_TYPES: frozenset[MemoryType] = frozenset(

1040 memory_type for memory_type in MemoryType if memory_type.is_gpu

1041)

1042SUPPORTED_MEMORY_TYPES: frozenset[MemoryType] = CPU_MEMORY_TYPES | GPU_MEMORY_TYPES

1043

1044# String value sets for validation

1045VALID_MEMORY_TYPES = frozenset(mt.value for mt in MemoryType)

1046VALID_GPU_MEMORY_TYPES = frozenset(mt.value for mt in GPU_MEMORY_TYPES)

1047

1048# Compatibility constants are generated projections of the enum declaration.

1049for _memory_type in MemoryType:

1050 globals()[f"MEMORY_TYPE_{_memory_type.name}"] = _memory_type.value