diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 913717c..9f0ea64 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -105,7 +105,7 @@ jobs: # 1. Framework unavailable tests (always run) # 2. GPU unavailable fallback paths (will run in CI) # 3. Mocked GPU tests (test cleanup code with mocked GPU state) - pytest -v tests/test_gpu_cleanup.py \ + pytest -v tests/test_libraries.py \ --cov=arraybridge \ --cov-report=term-missing \ --cov-report=html \ diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index e665667..2e49d1d 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -46,7 +46,7 @@ jobs: run: | # Run all GPU cleanup tests # Tests will use frameworks if available, skip gracefully if not - pytest -v tests/test_gpu_cleanup.py \ + pytest -v tests/test_libraries.py \ --cov=arraybridge \ --cov-report=term-missing \ --cov-report=html \ diff --git a/ADDING_NEW_FRAMEWORKS.md b/ADDING_NEW_FRAMEWORKS.md index 82a03c7..a2aad93 100644 --- a/ADDING_NEW_FRAMEWORKS.md +++ b/ADDING_NEW_FRAMEWORKS.md @@ -1,55 +1,180 @@ -# Adding a framework +# Adding a library -This is a maintainer how-to for extending ArrayBridge's closed framework -taxonomy. Applications should consume `MemoryType`; they should not add their -own framework registry. +A library is one class: a subclass of `arraybridge.libraries.ArrayLibrary` +that inherits the role parents it plays, states `name = "..."`, and defines +the abstract methods those parents ask for. Nothing else is edited. The class +can live in `src/arraybridge/libraries.py` or in any other module; it +registers when its module is imported. -## 1. Define typed leaves +## What follows from the class -Add an `ArrayOperations` leaf in `src/arraybridge/array_operations.py`: a -subclass declaring `memory_type = ""`. It registers itself, and -`ArrayOperations.for_memory(MemoryType.)` returns it. It must provide: +- Lookup: `ArrayLibrary.named(name)`, `ArrayLibrary.of(array)`, + `ArrayLibrary.for_module_name(module_name)`, `ArrayLibrary.all()`. +- A `MemoryType` member: `MemoryType.` with value `name` and + `.library` returning the library object. +- The decorator `arraybridge.`. +- Conversion to and from every other library through `convert_memory`, + `stack_slices` and `unstack_slices`. `ConversionRoute.planned` picks the + route from the role parents: same library is a device move, a DLPack + importer paired with a DLPack exporter is a DLPack exchange, anything else + is a copy through NumPy. +- Device answers (`host_devices`, `accelerator_devices`, `available_devices`, + `require_device`) and the array operations (`max`, `reshape`, + `stack_arrays`, ...) with NumPy-style defaults that call the library's + module. -- project an array to NumPy; -- create an array from NumPy on the requested framework-local device; -- stack prepared 2D arrays into one 3D array; -- apply the framework's dtype-scaling semantics; -- an override of each namespace primitive (`max`, `mean`, `assign`, - `linspace`, ...) whose library call or numeric rule differs from the - NumPy-compatible default. +## A host-only example -Define device, stream, DLPack, cleanup, and OOM leaves in -`src/arraybridge/types.py`. Use the no-op defaults in `FrameworkRuntime` only -when the framework genuinely lacks that capability. Do not encode behavior as -strings or add a parallel `MemoryType`-keyed table. +This file runs as it is with only NumPy installed. -## 2. Add the declaration +```python +import numpy as np -Add one `MemoryType` member carrying: +from arraybridge import HOST, ConversionRoute, MemoryType +from arraybridge import convert_memory, detect_memory_type, stack_slices +from arraybridge.libraries import ArrayLibrary, HostDevices -1. its stable value and import name; -2. its display name and module aliases; -3. whether it is GPU-backed; -4. any environment defaults that must be set before first import; -5. its `FrameworkRuntime` leaves. -The converter classes, target conversion methods, compatibility operation view, -dtype-scaling view, cleanup adapters, and framework decorator are generated -from the enum. Do not wire them manually. +class FrozenArray: + """Values in main memory that cannot be changed after creation.""" -## 3. Prove the extension + def __init__(self, values): + self.values = np.array(values) + self.values.flags.writeable = False + self.shape = self.values.shape + self.dtype = self.values.dtype -Add focused tests for: + def __getitem__(self, index): + return FrozenArray(self.values[index]) -- module-name detection and optional-dependency failure; -- NumPy round trips and every supported DLPack pair; -- CPU-only installations of a nominally GPU framework; -- invalid, negative, and multiple framework-local device IDs; -- same-framework moves between two devices; -- one-, two-, and multi-slice stack/unstack behavior; -- immutable array types; -- cleanup and OOM behavior without importing an absent framework; -- generated registry and export completeness. + +class FrozenLibrary(HostDevices[FrozenArray]): + name = "frozen" + # The module whose classes are this library's arrays. It defaults to + # ``name``; here the array class lives in this file. + import_name = __name__ + + # asked for by ArrayLibrary + def device_of(self, data): + return HOST + + def _to_host(self, data): + return np.array(data.values) + + def _scale_dtype(self, result, target_dtype): + return FrozenArray(result.values.astype(target_dtype)) + + # asked for by HostDevices + def _wrap_host_array(self, data): + return FrozenArray(data) + + def _move_to_host(self, data): + return data + + # The inherited default calls ``self.module.stack``; this module has none. + def _stack(self, values): + return FrozenArray(np.stack([value.values for value in values])) + + +# Lookup, by name and by array. +library = ArrayLibrary.named("frozen") +assert library in ArrayLibrary.all() +assert ArrayLibrary.of(FrozenArray([1, 2])) is library + +# The MemoryType member. +assert MemoryType.FROZEN.value == "frozen" +assert MemoryType.FROZEN.library is library + +# Devices. +assert library.available_devices() == (HOST,) + +# Conversion in both directions, by host copy. +numpy_library = ArrayLibrary.named("numpy") +assert ConversionRoute.planned(numpy_library, HOST, library, HOST).mechanism == "host copy" +frozen = convert_memory(np.arange(6).reshape(2, 3), "numpy", "frozen", HOST) +assert detect_memory_type(frozen) == "frozen" +assert convert_memory(frozen, "frozen", "numpy", HOST).tolist() == [[0, 1, 2], [3, 4, 5]] + +# Stacking. +stack = stack_slices([frozen, np.ones((2, 3))], "frozen", HOST) +assert stack.shape == (2, 2, 3) + +# The decorator named after the library. +from arraybridge import frozen as frozen_function + + +@frozen_function +def first_row(image): + return image[0] + + +assert first_row.execution_memory_type == "frozen" +assert first_row(frozen).shape == (3,) +``` + +## What the class states + +| Attribute | Meaning | Default | +|---|---|---| +| `name` | The registered name, the `MemoryType` value and the decorator name. | required | +| `import_name` | Top-level module to import; objects from it belong to the library. | `name` | +| `display_name` | Name used in messages. | `name` | +| `module_aliases` | Other top-level module names whose objects belong to the library (JAX: `jaxlib`). | `()` | +| `import_environment` | Environment variables set, if unset, before the first import and in child processes. | `()` | +| `oom_messages` | Lower-case text that marks an out-of-memory error message. | `()` | +| `oom_exception_paths` | Attribute paths from the module to its out-of-memory exception classes. | `()` | + +Every library defines three methods asked for by `ArrayLibrary` itself: + +- `device_of(data)`: the device holding the array. It never returns `None`. +- `_to_host(data)`: the array as a NumPy array. +- `_scale_dtype(result, target_dtype)`: convert the dtype, rescaling float + intensities into the integer range where the library does that. + +## Role parents + +Inherit each role the library plays. The abstract methods are the ones the +class must define. + +| Parent | The library... | Abstract methods | +|---|---|---| +| `HostDevices` | keeps arrays in main memory. | `_wrap_host_array(data)`, `_move_to_host(data)` | +| `AcceleratorDevices` | can keep arrays on numbered accelerators. | `_accelerator_count()`, `_accelerator_scope(device)`, `_move_to_accelerator(data, device)`, `_upload(data, device)` | +| `CudaLibrary` | is an `AcceleratorDevices` library with its own CUDA streams and a memory cache to release. | those of `AcceleratorDevices`, plus `stream_scope(stream)`, `current_device()`, `_free_cached_memory()` | +| `DLPackImport` | can wrap another library's buffer through DLPack. | `_import_dlpack(payload)` | +| `DLPackProtocolExport` | has arrays with the standard DLPack export methods on every device. | none | +| `NvidiaWheelLibraries` | loads CUDA from the `nvidia-*` wheels; child processes get their library directories on the search path. | none | +| `MutableArrays` | has the NumPy array API and arrays that can be changed in place; supplies `_scale_dtype` and in-place plane normalisation. | `_rescales_into(target_dtype)` | + +Methods with a default worth knowing: + +- `HostDevices._host_scope()`: the block in which new arrays are created in + main memory. The default does nothing; TensorFlow and JAX override it. +- `AcceleratorDevices._activate_accelerator(device)`: select a device for the + whole process. The default does nothing. +- `cleanup_loaded(device=None)`: release cached device memory. `CudaLibrary` + implements it; elsewhere the default does nothing. +- `exports_dlpack_from`, `exports_dlpack_data`, `export_dlpack`: override + these directly when export depends on the device or the array, as + TensorFlow does. +- `_stack`, `_cast`, `_reshape`, `_broadcast_to`, `_ones_like`, + `_logical_and`, `max`, `min`, `mean`, `sum`, `linspace`, `ones`, `assign`, + `outer`, `floor`, `accepts`, `nbytes_of`, `dtype_name`: the defaults call + NumPy-named functions on the library's module or methods on the array. + Override each one whose call or numeric rule differs. + +The six built-in classes at the end of `src/arraybridge/libraries.py` are the +reference for each combination of roles. + +## Tests to add + +- detection by array and by module name, and behaviour when the package is + not installed; +- NumPy round trips, and the planned route against every other library; +- each device the library accepts, and `ValueError` for one it does not have; +- moves between two devices of the library; +- stacking and unstacking one, two and several planes; +- cleanup and out-of-memory recognition without importing an absent library. Run: @@ -61,6 +186,6 @@ pytest python -m sphinx -E -W --keep-going -b html docs/source /tmp/arraybridge-docs ``` -When GPU hardware is available, also exercise every installed framework pair in -one process. This catches import-order, allocator-coexistence, DLPack, and real -API-shape defects that isolated mocks cannot prove. +When accelerator hardware is available, also convert between every installed +pair of libraries in one process. Import order, allocators sharing a device +and DLPack differences only show there. diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..1c92d91 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,122 @@ +# Changelog + +## 0.5.0 + +Each supported array library is now one class of one family, +`arraybridge.libraries.ArrayLibrary`. The class states the library's identity +and owns everything that library does its own way: devices, moving data, +streams, DLPack, out-of-memory recognition, child-process environment and +native array operations. Adding a library is writing one class; lookup, +conversion, the `MemoryType` member and the decorator follow from it. + +This release removes names without aliases. The table says what to call +instead. + +### The family + +| Class | Parents (roles) | +|---|---| +| `NumpyLibrary` | `HostDevices`, `MutableArrays` | +| `CupyLibrary` | `CudaLibrary`, `DLPackImport`, `DLPackProtocolExport`, `NvidiaWheelLibraries`, `MutableArrays` | +| `TorchLibrary` | `HostDevices`, `CudaLibrary`, `DLPackImport`, `DLPackProtocolExport` | +| `TensorflowLibrary` | `HostDevices`, `AcceleratorDevices`, `DLPackImport` | +| `JaxLibrary` | `HostDevices`, `AcceleratorDevices`, `DLPackImport`, `DLPackProtocolExport` | +| `PyclesperantoLibrary` | `AcceleratorDevices` | + +Lookup: `ArrayLibrary.named("torch")`, `ArrayLibrary.of(array)` (which library +made this), `ArrayLibrary.for_module_name("jaxlib")`, `ArrayLibrary.all()`. + +### `MemoryType` + +`MemoryType` is kept as a view: an `Enum` whose members are generated from the +registered library classes (`MemoryType.TORCH.value == "torch"`). It carries no +behaviour. `MemoryType.X.library` is the library object. A library class +defined later appears as a member. + +### Devices + +`arraybridge.devices`: `Device` with leaves `HostDevice()` (the value `HOST`), +`AcceleratorDevice(index)` and `NamedDevice(name)` (a device the library +reports that is not one of its numbered accelerators, such as pyclesperanto on +a CPU OpenCL device). Each has a `label`. + +`is_gpu` is removed. It answered "is this NumPy" and was read as "this needs a +GPU", which is false for torch, JAX and TensorFlow. Ask instead: + +| Question | Call | +|---|---| +| Does the library keep arrays in main memory? | `isinstance(library, HostDevices)`; `library.host_devices()` | +| Can the library use accelerators at all? | `isinstance(library, AcceleratorDevices)` | +| Which accelerators exist on this machine now? | `library.accelerator_devices()` | +| Every device it can use now | `library.available_devices()` | +| Where is this array? | `library.device_of(array)` (never `None`) | + +Every `gpu_id` / `device_id` integer parameter is now a `Device`. A device the +library does not have raises `ValueError` (`convert_memory` wraps it in +`MemoryConversionError`, as it wraps every failure); in particular a NumPy target with an +accelerator device raises (0.4.1 ignored the number), and torch, JAX and +TensorFlow accept `HOST` (0.4.1 always moved torch data to CUDA). +`process_slices` keeps a torch CPU tensor on the host (0.4.1 sent it to GPU 0). + +### Conversion + +`arraybridge.conversion`: + +- `ConversionRoute`: the mechanisms, tried in declaration order: `DeviceMove` + (`"device move"`), `DLPackExchange` (`"dlpack"`), `HostCopy` (`"host copy"`). + `ConversionRoute.planned(source_library, source_device, target_library, + target_device)` returns the route that will be tried first. It reads only + the library declarations and the device values: no device, no import. The + conversion itself selects its route with the same test. +- `ConversionRecord`: source and target library and device label, mechanism, + `nbytes`, `seconds`, and `declined_mechanisms` (routes that applied by + declaration but could not carry that array). Plain fields; `to_json()`, + `from_json()`. +- `ConversionObserver`: one hook, `ConversionObserver.current`. The default + ignores conversions and nothing is timed. `observer.install()` or + `with observer.installed():`. `convert_to`, `to_numpy` and `from_numpy` + report; a call that returns its input unchanged does not. +- `ConversionCounters`: an observer keeping totals per (source, target, + mechanism); `merge`, `rows`, `to_json`, `from_json`, `format_table`. + +DLPack is no longer attempted from a main-memory buffer into CuPy, which +cannot wrap one; the result (a host copy) is unchanged and the warning is gone. + +### Removed, and what to call instead + +| Removed | Use | +|---|---| +| `MemoryType.X.(...)` (any method) | `MemoryType.X.library.(...)` | +| the `module=` argument of those methods | gone; the library uses its imported module | +| `MemoryType.is_gpu` | see Devices | +| `CPU_MEMORY_TYPES`, `GPU_MEMORY_TYPES` | filter `MemoryType` by `isinstance(m.library, HostDevices / AcceleratorDevices)` | +| `SUPPORTED_MEMORY_TYPES`, `VALID_MEMORY_TYPES`, `VALID_GPU_MEMORY_TYPES`, `MEMORY_TYPE_*` | `MemoryType` itself; `m.value` | +| `device_id_of(data)` | `library.device_of(data)` | +| `available_device_ids()` | `library.accelerator_devices()` | +| `require_device(id)` returning the module | `library.require_device(device)` returns `None`; `library.module` | +| `canonical_dtype_name(dtype)` | `library.dtype_name(dtype)` | +| `supports_dlpack` | `library.imports_dlpack`, `library.imports_dlpack_from(device)` | +| `supports_dlpack_data(data)` | `library.exports_dlpack_data(data)` | +| `current_device_id()`, `create_stream()` | on `CudaLibrary` only: `current_device()`, `create_stream()`, `thread_stream(device)` | +| `MemoryType.subprocess_environment(env)` | `ArrayLibrary.child_environment(env)` | +| `arraybridge.types.FrameworkRuntime` and the `__` functions | methods of the library classes | +| `arraybridge.types.DLPackPayload` | `arraybridge.libraries.DLPackPayload` | +| `arraybridge.array_operations.ArrayOperations`, `ArrayOperations.for_memory(m)` | `ArrayLibrary`, `m.library` (same primitive names) | +| `ArrayOperations.to_numpy(data, module)` and the other `module`-taking leaves | the library method without `module`: `to_numpy(data)`, `from_numpy(data, device)`, `scale_dtype`, `stack_arrays(values, device)`, `astype`, `logical_and`, `reshape`, `broadcast_to`, `ones_like`, `normalize_planes` | +| `arraybridge.framework_config`, `framework_ops`, `_FRAMEWORK_CONFIG`, `_FRAMEWORK_OPS` | nothing used them | +| `arraybridge.converters_registry` (`ConverterBase`, `get_converter`) | `library.convert_to(data, target_library, device)` | +| `arraybridge.dtype_scaling.SCALING_FUNCTIONS` | `library.scale_dtype(data, dtype)` | +| `arraybridge.gpu_cleanup.cleanup_all_gpu_frameworks(device_id)` | `ArrayLibrary.cleanup_all_loaded(device)` | +| `cleanup__gpu`, `MEMORY_TYPE_CLEANUP_REGISTRY` | `library.cleanup_loaded(device)` | +| `arraybridge.oom_recovery._execute_with_oom_recovery(f, name, retries, device_id)` | `library.run_with_oom_recovery(f, retries, device)` | +| `_is_oom_error(e, name)` | `library.is_oom_error(e)` | +| `arraybridge.utils._get_device_id`, `_set_device`, `_move_to_device` | `library.device_of`, `activate_device`, `move_to_device` | +| `arraybridge.decorators.ThreadGPUContext` | `CudaLibrary.thread_stream(device)`; per-call wrapping is `library.execution_wrapper(func, oom_recovery)` | +| `convert_memory(..., gpu_id=)`, `stack_slices(..., gpu_id=)`, `unstack_slices(..., gpu_id=)`, `process_slices(..., gpu_id=)` | the same functions with `device` | + +### Typing + +`typing.Any` is not used in the package. `arraybridge.arrays.Array` is the +nominal type of an array of any supported library; `library.array_class` is an +array of that library. Native array classes are registered with them when a +library first meets one. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index cd3a403..80455c2 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -117,29 +117,34 @@ git checkout -b fix/issue-description Use Google-style docstrings: ```python -def convert_memory(data: Any, source_type: str, target_type: str, gpu_id: int) -> Any: +def convert_memory( + data: Array, + source_type: str | MemoryType, + target_type: str | MemoryType, + device: Device, +) -> Array: """ - Convert data between memory types. + Convert data from the source library to the target library on a device. Args: - data: The data to convert + data: The array to convert source_type: The source memory type (e.g., "numpy", "torch") target_type: The target memory type (e.g., "cupy", "jax") - gpu_id: The target GPU device ID + device: The device of the target library to put the result on Returns: - The converted data in the target memory type + The converted array in the target memory type Raises: - ValueError: If source_type or target_type is invalid + ValueError: If source_type or target_type is not a registered library MemoryConversionError: If conversion fails Example: import numpy as np - from arraybridge import convert_memory + from arraybridge import HOST, convert_memory data = np.array([1, 2, 3]) - result = convert_memory(data, "numpy", "torch", gpu_id=0) + result = convert_memory(data, "numpy", "torch", HOST) """ ``` @@ -159,22 +164,22 @@ def convert_memory(data: Any, source_type: str, target_type: str, gpu_id: int) - ```python import pytest import numpy as np -from arraybridge import convert_memory +from arraybridge import HOST, convert_memory class TestConvertMemory: """Tests for convert_memory function.""" def test_convert_numpy_to_numpy(self): - """Test converting NumPy to NumPy (no-op).""" + """Converting NumPy to NumPy on the host returns the same array.""" arr = np.array([1, 2, 3]) - result = convert_memory(arr, "numpy", "numpy", gpu_id=0) - np.testing.assert_array_equal(result, arr) + result = convert_memory(arr, "numpy", "numpy", HOST) + assert result is arr def test_invalid_source_type_raises_error(self): """Test that invalid source type raises ValueError.""" arr = np.array([1, 2, 3]) with pytest.raises(ValueError): - convert_memory(arr, "invalid", "numpy", gpu_id=0) + convert_memory(arr, "invalid", "numpy", HOST) ``` ### Running Tests @@ -400,18 +405,13 @@ GitHub Actions will automatically: 4. **Test thoroughly**: Include unit, integration, and edge case tests 5. **Maintain compatibility**: Avoid breaking changes when possible -### Adding Framework Support +### Adding Library Support -To add support for a new framework: - -1. Add one `ArrayOperations` leaf in `array_operations.py` declaring its `memory_type`. -2. Add one `MemoryType` declaration carrying its runtime leaves. -3. Add detection, conversion-pair, stack, cross-device, and optional-dependency tests. -4. Update the extension and reference documentation. - -Do not add framework behavior to `_FRAMEWORK_CONFIG`, a consumer registry, or -generated decorator exports. See `ADDING_NEW_FRAMEWORKS.md` for the complete -declaration-owned extension procedure. +A new library is one class: a subclass of `ArrayLibrary` with the role parents +it plays, `name = "..."`, and the abstract methods of those parents. Lookup, +the `MemoryType` member, the decorator and the conversion routes follow from +the class. See `ADDING_NEW_FRAMEWORKS.md` for a runnable example, the methods +each role parent asks for, and the tests to add. ## Getting Help @@ -427,4 +427,4 @@ Contributors will be recognized in: - Release notes (for significant contributions) - Project README (for major features) -Thank you for contributing to arraybridge! 🎉 +Thank you for contributing to arraybridge! diff --git a/README.md b/README.md index eeb298c..d149fc5 100644 --- a/README.md +++ b/README.md @@ -1,36 +1,66 @@ # arraybridge -ArrayBridge provides explicit conversion and shared lifecycle utilities for -NumPy, CuPy, PyTorch, TensorFlow, JAX, and pyclesperanto arrays. +ArrayBridge converts arrays between NumPy, CuPy, PyTorch, TensorFlow, JAX and +pyclesperanto, and gives code one way to ask each library about devices, +memory and array operations. -Core dependencies are NumPy and metaclass-registry. Other frameworks are -optional. +Core dependencies are NumPy and metaclass-registry. The other libraries are +optional and are imported only when used. + +Version 0.5.0 removed names without aliases. [CHANGELOG.md](CHANGELOG.md) +lists each removed name and what to call instead. ## Quick start ```python import numpy as np -from arraybridge import convert_memory, detect_memory_type +from arraybridge import HOST, convert_memory, detect_memory_type value = np.arange(6).reshape(2, 3) assert detect_memory_type(value) == "numpy" -copy = convert_memory( - value, - source_type="numpy", - target_type="numpy", - gpu_id=0, -) +result = convert_memory(value, source_type="numpy", target_type="numpy", device=HOST) +assert result is value ``` -`convert_memory` requires the declared source type, target type, and device id. -The source and target `MemoryType` declarations perform conversion directly. -The device id is required even for CPU conversions so call sites have one -stable signature. A GPU target must declare that identifier as available; -ArrayBridge does not silently place the value on the CPU. +`convert_memory(data, source_type, target_type, device)` takes all four +arguments. `device` is a device of the target library: `HOST` for main memory, +`AcceleratorDevice(0)` for the library's first accelerator. A conversion to a +device the target library does not have fails; nothing is placed somewhere +else instead. + +## Devices + +A device is a value: `HOST` (main memory), `AcceleratorDevice(index)` (an +accelerator as its library numbers it) or `NamedDevice(name)` (a device a +library reports that is not one of its numbered accelerators, such as +pyclesperanto on a CPU OpenCL device). Indices belong to one library; the same +index in two libraries need not be the same hardware. + +| Target library | Devices it accepts | +|---|---| +| NumPy | `HOST` | +| CuPy, pyclesperanto | `AcceleratorDevice(i)` | +| PyTorch, JAX, TensorFlow | `HOST` or `AcceleratorDevice(i)` | + +Each library answers what it has on this machine: + +```python +import numpy as np + +from arraybridge import HOST, AcceleratorDevice, ArrayLibrary + +numpy_library = ArrayLibrary.named("numpy") +assert numpy_library.host_devices() == (HOST,) +assert numpy_library.accelerator_devices() == () +assert numpy_library.available_devices() == (HOST,) +assert numpy_library.device_of(np.zeros(3)) == HOST + +assert AcceleratorDevice(0).label == "accelerator:0" +``` -## Declarative decorators +## Decorators ```python from arraybridge import numpy @@ -38,74 +68,120 @@ from arraybridge import numpy @numpy def normalize(image): return image / max(float(image.max()), 1.0) + +assert normalize.input_memory_type == "numpy" ``` -The framework decorators attach `input_memory_type`, `output_memory_type`, and -`execution_memory_type` metadata, provide dtype/slice runtime parameters, and -add framework-specific stream/OOM handling where supported. They do **not** -convert inputs or outputs between frameworks and do not accept a `gpu_id` -argument. A host runtime must call `convert_memory` at the boundary it plans and -scope execution using the execution declaration. +The decorator named after a library sets `input_memory_type`, +`output_memory_type` and `execution_memory_type` on the function, adds the +`slice_by_slice` and `dtype_config` keyword arguments, and wraps the call with +what that library's devices need (per-thread CUDA streams for CuPy and +PyTorch, out-of-memory retries for libraries with accelerators). It does not +convert inputs or outputs between libraries and does not choose a device. The +caller converts with `convert_memory`. ## Stack utilities ```python import numpy as np -from arraybridge import stack_slices, unstack_slices +from arraybridge import HOST, stack_slices, unstack_slices slices = [np.zeros((8, 8)), np.ones((8, 8))] -stack = stack_slices(slices, memory_type="numpy", gpu_id=0) -restored = unstack_slices(stack, memory_type="numpy", gpu_id=0) +stack = stack_slices(slices, memory_type="numpy", device=HOST) +restored = unstack_slices(stack, memory_type="numpy", device=HOST) +assert stack.shape == (2, 8, 8) and len(restored) == 2 ``` -`stack_slices` requires non-empty 2D inputs. `unstack_slices` requires a 3D -array. Both validate shape and use explicit target memory/device declarations. +`stack_slices` takes a non-empty list of 2D arrays. `unstack_slices` takes a +3D array. Both raise `ValueError` for a wrong shape or a device the target +library does not have. -## Installation +## Which route a conversion takes -```bash -pip install arraybridge -pip install "arraybridge[torch]" -pip install "arraybridge[cupy]" +A conversion tries three mechanisms in order: `"device move"` (same library, +other device), `"dlpack"` (the target wraps the source's buffer) and +`"host copy"` (download to NumPy, then upload). `ConversionRoute.planned` +tells which one will be tried first. It reads only the library classes and the +device values, so it works without the libraries or the hardware: + +```python +from arraybridge import HOST, AcceleratorDevice, ArrayLibrary, ConversionRoute + +torch = ArrayLibrary.named("torch") +cupy = ArrayLibrary.named("cupy") +gpu = AcceleratorDevice(0) + +assert ConversionRoute.planned(torch, gpu, cupy, gpu).mechanism == "dlpack" +assert ConversionRoute.planned(torch, HOST, cupy, gpu).mechanism == "host copy" +assert ConversionRoute.planned(torch, HOST, torch, gpu).mechanism == "device move" ``` -Documentation: +To see what conversions actually ran, install a `ConversionCounters` observer. +It keeps a count, bytes and seconds per source, target and mechanism. This +example needs JAX: + +```python +import numpy as np + +from arraybridge import HOST, ConversionCounters, convert_memory -### Native array geometry +counters = ConversionCounters() +with counters.installed(): + convert_memory(np.zeros((4, 4), dtype=np.float32), "numpy", "jax", HOST) -`MemoryType.reshape(array, shape)` and `MemoryType.broadcast_to(array, shape)` -operate in the array's framework and device. `MemoryType.ones_like(reference, -shape=..., dtype=...)` allocates on the reference's device; omitting `shape` -retains its shape. Reshape and broadcast preserve views where the backend can -represent them. Broadcasting is not a promise of a writable, independent buffer. +row = counters.rows()[0] +assert (row["source_library"], row["target_library"]) == ("numpy", "jax") +assert (row["mechanism"], row["count"], row["nbytes"]) == ("host copy", 1, 64) +print(counters.format_table()) +``` -These operations never silently project pixels to NumPy. Pyclesperanto cannot -provide native changed-shape reshape or broadcasting, so those requests raise -`NotImplementedError`. Its native ones allocation supports one to three -dimensions and its supported integer/float32 dtypes; boolean allocation is -explicitly unsupported. Existing conversion and disk boundaries remain separate. +A call that returns its input unchanged is not counted. With no observer +installed nothing is timed. -### Writing an operation once for every framework +## Array operations in the array's own library -`ArrayOperations.for_memory(memory_type)` returns that framework's operations -leaf: one registered subclass per `MemoryType`, holding only the primitives its -library does differently (`max`, `min`, `mean`, `sum`, `astype`, -`prepend_axis`, `linspace`, `ones`, `assign`, `outer`, `floor`, `accepts`). -Code written against the leaf runs natively in each framework: +`ArrayLibrary.named("torch")`, `ArrayLibrary.of(array)` and +`MemoryType.TORCH.library` all return the library object. Its methods run in +that library, on the device of their input: `max`, `min`, `mean`, `sum`, +`astype`, `reshape`, `broadcast_to`, `ones_like`, `ones`, `prepend_axis`, +`linspace`, `assign`, `outer`, `floor`, `logical_and`, `stack_arrays`, +`normalize_planes`, `scale_dtype`, `to_numpy`, `from_numpy`. ```python -from arraybridge import MemoryType -from arraybridge.array_operations import ArrayOperations +import numpy as np + +from arraybridge import ArrayLibrary -def max_projection(ops, stack): - return ops.prepend_axis(ops.max(stack, 0)) +def max_projection(stack): + library = ArrayLibrary.of(stack) + return library.prepend_axis(library.max(stack, 0)) -ops = ArrayOperations.for_memory(MemoryType.TORCH) +assert max_projection(np.ones((4, 8, 8))).shape == (1, 8, 8) ``` -`ops.module` is the framework, imported on first use. Each leaf keeps its -library's numeric rules: torch, JAX and TensorFlow take means in float32, and -pyclesperanto reduces only along the leading axis. +`library.module` is the library's module, imported on first use. Each library +keeps its own numeric rules: PyTorch, JAX and TensorFlow take means in +float32, and pyclesperanto reduces only along the leading axis. + +`reshape`, `broadcast_to` and `ones_like` do not download pixels to NumPy. +pyclesperanto cannot change shape or broadcast natively, so those calls raise +`NotImplementedError` unless the shape is unchanged; its `ones_like` supports +one to three dimensions and its float32 and integer dtypes, not bool. + +## Adding a library + +A library is one class. See +[ADDING_NEW_FRAMEWORKS.md](ADDING_NEW_FRAMEWORKS.md). + +## Installation + +```bash +pip install arraybridge +pip install "arraybridge[torch]" +pip install "arraybridge[cupy]" +``` + +Documentation: diff --git a/TESTING.md b/TESTING.md index 0a465ff..752c3b6 100644 --- a/TESTING.md +++ b/TESTING.md @@ -24,7 +24,7 @@ exclude framework-specific tests. On a GPU host, validate more than a single NumPy round trip: - import all installed frameworks in one fresh process; -- enumerate framework-local devices through each `MemoryType`; +- enumerate each library's devices through `ArrayLibrary.all()` (`available_devices()`); - convert a small exact array through every available source-target pair; - stack and unstack one, two, and several planes in every framework; - verify same-framework movement across multiple local devices when present; diff --git a/docs/source/README.md b/docs/source/README.md index 4d4c711..f8b0317 100644 --- a/docs/source/README.md +++ b/docs/source/README.md @@ -9,7 +9,7 @@ python -m pip install -e ".[docs]" python -m sphinx -E -W --keep-going -b html docs/source docs/source/_build/html ``` -Narrative pages document public behavior and ownership. `api_reference.rst` -tracks the root export surface. Keep code examples valid under +The pages describe public behaviour. `api_reference.rst` lists the root +exports, the device model and conversion reporting. Keep code examples valid under `scripts/validate_docs.py` in the OpenHCS superproject and verify signatures against `src/arraybridge/` when public APIs change. diff --git a/docs/source/advanced_topics.rst b/docs/source/advanced_topics.rst index 8d95f7e..8fe2a28 100644 --- a/docs/source/advanced_topics.rst +++ b/docs/source/advanced_topics.rst @@ -1,24 +1,28 @@ -Advanced extension topics -========================= +Adding a library +================ -Framework behavior is carried by each ``MemoryType`` member. Add a framework by -extending that owning declaration: +Each supported library is one class of the family +``arraybridge.libraries.ArrayLibrary``. The class states the library's name +and owns what that library does its own way: devices, moving data, streams, +DLPack, out-of-memory recognition, child-process environment and array +operations. What several libraries share is a role parent they inherit: +``HostDevices``, ``AcceleratorDevices``, ``CudaLibrary``, ``DLPackImport``, +``DLPackProtocolExport``, ``NvidiaWheelLibraries``, ``MutableArrays``. -1. add the ``MemoryType`` member; -2. provide its typed runtime and array-operation leaves; -3. define conversion, stack, scaling, and device behavior; -4. add cleanup and OOM behavior only when the framework supports them; -5. exercise detection and every supported conversion pair. +To add a library, write one subclass with ``name = "..."``, the role parents +it plays and their abstract methods. Lookup, the ``MemoryType`` member, the +decorator and the conversion routes follow from the class. The repository's +`ADDING_NEW_FRAMEWORKS.md +`_ +has a runnable example and the list of methods each role parent asks for. -Do not add a parallel framework-name table in a consumer. Generic consumers use -``MemoryType`` directly. The converter registry and historical framework config -are generated projections, not authorities. +Do not keep a second table of library names in a consumer. Iterate +``MemoryType`` or ``ArrayLibrary.all()``. -Version 0.3 keeps the private ``_FRAMEWORK_CONFIG`` and ``_FRAMEWORK_OPS`` names -only as read-only operation views. Their historical mapping shape is not a -compatibility contract. Migrate callers to ``MemoryType`` methods. +Caching and batching +-------------------- -Conversion caching and batching belong to the caller because value lifetime and -semantic identity are application concerns. ArrayBridge functions remain -stateless at the conversion boundary except for declared thread-local GPU -context. +ArrayBridge does not cache or batch conversions. How long a converted value +stays valid is the caller's knowledge. The state ArrayBridge keeps is the +per-thread CUDA stream of each CuPy and PyTorch device and the installed +``ConversionObserver``. diff --git a/docs/source/api_reference.rst b/docs/source/api_reference.rst index 1a6650f..f5a3d53 100644 --- a/docs/source/api_reference.rst +++ b/docs/source/api_reference.rst @@ -1,45 +1,261 @@ API reference ============= +Version 0.5.0 removed names without aliases. The repository's `CHANGELOG.md +`_ lists +each removed name and what to call instead. + Public imports -------------- -``MemoryType``, ``CPU_MEMORY_TYPES``, ``GPU_MEMORY_TYPES``, ``SUPPORTED_MEMORY_TYPES`` - Framework declarations. Each member owns import identity, conversion, - stacking, dtype scaling, device discovery, scoped selection, movement, - DLPack import/export, cleanup, and pre-import coexistence defaults. +``ArrayLibrary`` + The family of supported libraries. ``ArrayLibrary.named("torch")``, + ``ArrayLibrary.of(array)``, ``ArrayLibrary.for_module_name("jaxlib")`` and + ``ArrayLibrary.all()`` return library objects. A library object answers for + its devices, moves and converts arrays, and runs array operations in that + library. + +``MemoryType`` + The registered library names as an ``Enum`` + (``MemoryType.TORCH.value == "torch"``). It carries no behaviour; + ``MemoryType.TORCH.library`` is the library object. A library class defined + later appears as a member. + +``Device``, ``HostDevice``, ``AcceleratorDevice``, ``NamedDevice``, ``HOST`` + Where an array lives. See `Devices`_. ``detect_memory_type`` and ``convert_memory`` - Detection and explicit conversion. + Name the library that made an array; convert an array to another library on + a device. + +``ConversionRoute``, ``ConversionRecord``, ``ConversionObserver``, ``ConversionCounters`` + Which mechanism a conversion uses, and a hook that receives each conversion. + See `Conversion reporting`_. + +``Array`` + The nominal type of an array of any supported library. + ``library.array_class`` is an array of that library. ``ArrayGeometry`` - Framework-neutral shape inspection for arrays and nominal ``ArrayPayload`` - values. Geometry inspection reads declared shape metadata and does not move - device data to host memory. + Shape inspection for arrays and ``ArrayPayload`` values. It reads the + declared shape and does not move device data to main memory. + +``ArrayPayload`` + Abstract container that keeps its context around array data. ``memory_types``, ``numpy``, ``cupy``, ``torch``, ``tensorflow``, ``jax``, ``pyclesperanto`` - Callable memory declarations and wrappers. + Decorators that declare a function's memory types. There is one for every + registered library, made on first access. + +``MemoryContractAttribute`` + Names of the attributes the decorators set on a function. ``DtypeConversion`` Output dtype policy values. -``stack_slices``, ``unstack_slices``, ``process_slices`` - Validated plane/stack operations. +``SliceBySliceRuntimeParameter``, ``wrap_dtype_preserving_callable`` + The slice control and the dtype and slice wrapper used by the decorators. -``cleanup_all_gpu_frameworks`` - Installed-framework cache cleanup. +``stack_slices``, ``unstack_slices``, ``process_slices`` + Plane and stack operations. ``MemoryConversionError`` - Conversion failure boundary. + Raised by ``convert_memory`` when a conversion fails. + +The export list is ``arraybridge.__all__``. Names beginning with an underscore +are not part of the public interface. + +Devices +------- + +A device is a value of ``arraybridge.devices``. Equal fields mean equal +devices, and each has a ``label``. + +``HOST`` + Main memory, the one ``HostDevice()``. Label ``"host"``. + +``AcceleratorDevice(index)`` + One accelerator as its library numbers it (a CUDA or an OpenCL device). + Label ``"accelerator:"``. The index belongs to one library; the same + index in two libraries need not be the same hardware. + +``NamedDevice(name)`` + A device a library reports that is not one of its numbered accelerators, + such as pyclesperanto running on a CPU OpenCL device. Label + ``"device:"``. + +Every ``device`` parameter takes one of these. NumPy accepts ``HOST``; CuPy +and pyclesperanto accept ``AcceleratorDevice(i)``; PyTorch, JAX and TensorFlow +accept either. A device the library does not have on this machine raises +``ValueError`` from the library methods and from ``stack_slices`` and +``unstack_slices``; ``convert_memory`` raises ``MemoryConversionError`` +carrying that message. + +A library object answers these questions: + +``library.host_devices()`` + ``(HOST,)`` if the library keeps arrays in main memory, else ``()``. + +``library.accelerator_devices()`` + The accelerators the library can use on this machine now. Empty when the + package is not installed or sees no accelerator. + +``library.available_devices()`` + Both of the above, host first. + +``library.device_of(array)`` + The device holding the array. Never ``None``. + +``library.require_device(device)`` + Raises ``ValueError`` unless the library has the device. + +``isinstance(library, HostDevices)`` and ``isinstance(library, AcceleratorDevices)`` +(both in ``arraybridge.libraries``) answer whether the library can use main +memory or accelerators at all, without importing it. + +.. code-block:: python + + import numpy as np + + from arraybridge import HOST, AcceleratorDevice, ArrayLibrary, MemoryType + from arraybridge.libraries import AcceleratorDevices, HostDevices + + numpy_library = ArrayLibrary.named("numpy") + assert numpy_library.host_devices() == (HOST,) + assert numpy_library.accelerator_devices() == () + assert numpy_library.available_devices() == (HOST,) + assert numpy_library.device_of(np.zeros(3)) == HOST + + assert (HOST.label, AcceleratorDevice(1).label) == ("host", "accelerator:1") + + on_host = [m.value for m in MemoryType if isinstance(m.library, HostDevices)] + on_accelerators = [m.value for m in MemoryType if isinstance(m.library, AcceleratorDevices)] + assert on_host == ["numpy", "torch", "tensorflow", "jax"] + assert on_accelerators == ["cupy", "torch", "tensorflow", "jax", "pyclesperanto"] + +Conversion reporting +-------------------- + +``ConversionRoute`` has one subclass per mechanism, tried in this order: -The canonical export list is ``arraybridge.__all__``. Names beginning with an -underscore are implementation surfaces even when temporarily re-exported for a -host migration. +``DeviceMove`` (``"device move"``) + Source and target are the same library; the array changes device only. -Documented public surface -------------------------- +``DLPackExchange`` (``"dlpack"``) + The target wraps the source's buffer. Applies when the target imports + DLPack from the source's device and the source exports from it. + +``HostCopy`` (``"host copy"``) + Download to a NumPy array, then upload to the target. Always applies. + +``ConversionRoute.planned(source_library, source_device, target_library, +target_device)`` returns the route class that a conversion of that pair tries +first. It reads only the library classes and the device values: no library is +imported and no device is touched. The conversion selects its route with the +same test. A route that applies but cannot carry a particular array is +skipped and the next one is tried. + +.. code-block:: python + + from arraybridge import HOST, AcceleratorDevice, ArrayLibrary, ConversionRoute + + numpy_library = ArrayLibrary.named("numpy") + torch = ArrayLibrary.named("torch") + cupy = ArrayLibrary.named("cupy") + gpu = AcceleratorDevice(0) + + assert ConversionRoute.planned(torch, gpu, cupy, gpu).mechanism == "dlpack" + assert ConversionRoute.planned(torch, HOST, cupy, gpu).mechanism == "host copy" + assert ConversionRoute.planned(numpy_library, HOST, torch, HOST).mechanism == "host copy" + assert ConversionRoute.planned(torch, HOST, torch, gpu).mechanism == "device move" + +``ConversionObserver`` + Receives one ``ConversionRecord`` per conversion while it is installed. + Subclass it and define ``record(record)``. ``observer.install()`` makes it + the process's observer and returns the one it replaced; + ``with observer.installed():`` does so for a block. + ``ConversionObserver.current`` is the installed observer. The default + ignores conversions, and then nothing is timed. ``library.convert_to`` + (which ``convert_memory`` calls), ``library.to_numpy`` and + ``library.from_numpy`` report. A call that returns its input unchanged does + not. + +``ConversionRecord`` + Fields: ``source_library``, ``source_device``, ``target_library``, + ``target_device`` (device labels), ``mechanism``, ``nbytes``, ``seconds`` + and ``declined_mechanisms`` (routes that applied but could not carry that + array). ``to_json()`` and ``from_json()`` convert to and from a plain + mapping. + +``ConversionCounters`` + An observer that keeps count, bytes and seconds per source library and + device, target library and device, and mechanism. ``rows()`` (also named + ``to_json()``) returns one row per kind, largest byte total first; + ``from_json(rows)`` rebuilds counters; ``merge(other)`` adds another + process's totals; ``format_table()`` returns aligned text. + +.. code-block:: python + + from arraybridge import ConversionCounters, ConversionObserver, ConversionRecord + + record = ConversionRecord( + source_library="numpy", + source_device="host", + target_library="torch", + target_device="accelerator:0", + mechanism="host copy", + nbytes=64, + seconds=0.25, + ) + assert ConversionRecord.from_json(record.to_json()) == record + + counters = ConversionCounters() + with counters.installed(): + # Conversions made in this block are reported to ``counters``. + ConversionObserver.current.record(record) + + worker = ConversionCounters.from_json(counters.to_json()) + counters.merge(worker) + row = counters.rows()[0] + assert (row["mechanism"], row["count"], row["nbytes"]) == ("host copy", 2, 128) + print(counters.format_table()) + +With JAX installed, a real conversion is counted the same way: + +.. code-block:: python + + import numpy as np + + from arraybridge import HOST, ConversionCounters, convert_memory + + counters = ConversionCounters() + with counters.installed(): + convert_memory(np.zeros((4, 4), dtype=np.float32), "numpy", "jax", HOST) + + row = counters.rows()[0] + assert (row["source_library"], row["target_library"]) == ("numpy", "jax") + assert (row["mechanism"], row["count"], row["nbytes"]) == ("host copy", 1, 64) + +Modules +------- + +The root package exports the names listed above. The role parents and the +library classes are in ``arraybridge.libraries``, and the route classes in +``arraybridge.conversion``. .. automodule:: arraybridge :members: - :exclude-members: _FRAMEWORK_CONFIG, _FRAMEWORK_OPS, _ensure_module, _execute_with_oom_recovery, _get_device_id, _supports_dlpack + :exclude-members: _ensure_module, _supports_dlpack + :member-order: bysource + +.. automodule:: arraybridge.libraries + :members: + :member-order: bysource + +.. automodule:: arraybridge.devices + :members: + :member-order: bysource + +.. automodule:: arraybridge.conversion + :members: :member-order: bysource diff --git a/docs/source/ci-cd.rst b/docs/source/ci-cd.rst index d0b6fcc..e24d080 100644 --- a/docs/source/ci-cd.rst +++ b/docs/source/ci-cd.rst @@ -1,7 +1,7 @@ Continuous integration ====================== -The repository workflows are authoritative. The main test matrix currently +The workflow files in the repository define what CI runs. The main test matrix currently covers Python 3.10 through 3.13 on Linux, Windows, and macOS. Separate jobs exercise optional GPU paths and code quality. diff --git a/docs/source/contributing.rst b/docs/source/contributing.rst index e2410b9..ce61bb6 100644 --- a/docs/source/contributing.rst +++ b/docs/source/contributing.rst @@ -13,15 +13,15 @@ by the repository: python -m pytest python -m sphinx -E -W --keep-going -b html docs/source docs/source/_build/html -Framework extensions --------------------- +Adding a library +---------------- -A framework extension must update the nominal ``MemoryType`` declaration and -its typed operation leaves. Add detection, conversion-pair, stack, -cross-device, and optional-dependency tests. Consumers should not maintain a -second framework-name registry. +A new library is one ``ArrayLibrary`` subclass. Add tests for detection, +conversion against the other libraries, stacking, its devices and the case +where the package is not installed. -See :doc:`advanced_topics` for the ownership model and the repository's +See :doc:`advanced_topics` for how the class is put together and the +repository's `contribution guide `_ for the review workflow. diff --git a/docs/source/converters.rst b/docs/source/converters.rst index ae65d72..f6211d6 100644 --- a/docs/source/converters.rst +++ b/docs/source/converters.rst @@ -1,47 +1,61 @@ -Conversion system -================= +Conversion +========== -``convert_memory(data, source_type, target_type, gpu_id)`` resolves the source -and target ``MemoryType`` declarations and lets those declarations perform the -conversion. +``convert_memory(data, source_type, target_type, device)`` converts ``data`` +from the source library to the target library on ``device``. It looks up the +two library objects and calls ``source.convert_to(data, target, device)``. .. code-block:: python import numpy as np - from arraybridge import MemoryType, convert_memory + from arraybridge import HOST, MemoryType, convert_memory value = np.arange(4) - result = convert_memory( - value, - MemoryType.NUMPY, - MemoryType.NUMPY, - 0, - ) + result = convert_memory(value, MemoryType.NUMPY, MemoryType.NUMPY, HOST) + assert result is value -Both ``source_type`` and ``target_type`` accept either the canonical -``MemoryType`` member or its string value. Decorators normalize the same -members to string metadata at declaration time, so callable contracts and -conversion plans share one taxonomy without carrying duplicate enum classes. +``source_type`` and ``target_type`` accept a ``MemoryType`` member or its +string value. Call ``detect_memory_type`` when the source library is not +already known. Passing the wrong source type is a caller error. -Call ``detect_memory_type`` when the source type is not already known. Passing a -wrong source declaration is a caller error; conversion planning should keep the -declared type aligned with the actual value. +Devices +------- -Device semantics ----------------- +``device`` is required and is a device of the target library: ``HOST`` for +NumPy; ``AcceleratorDevice(i)`` for CuPy and pyclesperanto; either for +PyTorch, JAX and TensorFlow. Moving an array between two devices of one +library is a conversion with the same source and target type. See +:doc:`api_reference` for the device values and the questions a library +answers about them. -``gpu_id`` is required. GPU converters use it for allocation or device -selection. CPU conversion accepts the same argument for a uniform API. Moving a -value between two devices is expressed as conversion with the target device id. +.. code-block:: python + + import numpy as np + + from arraybridge import AcceleratorDevice, MemoryConversionError, convert_memory + + try: + convert_memory(np.arange(4), "numpy", "numpy", AcceleratorDevice(0)) + except MemoryConversionError as error: + assert "accelerator:0 is unavailable" in str(error) + else: + raise AssertionError("NumPy has no accelerator device") + +Routes +------ -Declarations prefer framework-native or DLPack paths when both endpoints -support them and otherwise use an explicit NumPy round trip. No zero-copy -guarantee applies to every pair. +A conversion tries, in order: a device move when source and target are the +same library; a DLPack exchange when the target can wrap the source's buffer +on its device; a copy through a NumPy array in main memory. +``ConversionRoute.planned`` tells which is tried first for a pair, and a +``ConversionObserver`` receives what actually ran; both are described in +:doc:`api_reference`. No pair is promised to be zero-copy. Failures -------- -Invalid framework names raise ``ValueError``. Conversion failures raise -``MemoryConversionError`` with source/target context. Optional frameworks are -loaded only when their declaration is used. +An unregistered library name raises ``ValueError``. Any failure during the +conversion, including a missing optional package or a device the target does +not have, raises ``MemoryConversionError`` with the source and target names. +An optional library is imported only when a conversion uses it. diff --git a/docs/source/decorators.rst b/docs/source/decorators.rst index c844d43..870393c 100644 --- a/docs/source/decorators.rst +++ b/docs/source/decorators.rst @@ -1,8 +1,10 @@ Memory declaration decorators ============================= -``memory_types`` and the six framework helpers declare callable memory -contracts. +``memory_types`` and the decorator named after each library (``numpy``, +``cupy``, ``torch``, ``tensorflow``, ``jax``, ``pyclesperanto``) declare which +library a function's input and output arrays belong to. A decorator exists for +every registered library and is made on first access. .. code-block:: python @@ -16,34 +18,34 @@ contracts. def offset(image, value=1): return image + value -Framework helpers accept ``input_type``, ``output_type``, ``oom_recovery``, and -``contract``. They may be used as ``@numpy`` or ``@numpy()``. +The library decorators accept ``input_type``, ``output_type``, +``oom_recovery``, ``contract``, ``slice_by_slice_default`` and +``dtype_config_default``. They may be used as ``@numpy`` or ``@numpy()``. Direct-call behavior -------------------- -Decorators do not convert the input or output between frameworks. They attach -metadata, add keyword-only ``slice_by_slice`` and ``dtype_config`` runtime -parameters, apply the selected dtype policy, and wrap supported GPU frameworks -with thread-local stream/OOM handling. +Decorators do not convert the input or output between libraries. They set +the memory type attributes, add the keyword-only ``slice_by_slice`` and +``dtype_config`` parameters, apply the selected dtype policy, and wrap the +call with ``library.execution_wrapper(func, oom_recovery)``: per-thread CUDA +streams for CuPy and PyTorch and out-of-memory retries for libraries with +accelerators (see :doc:`gpu_features`). For NumPy the wrapper adds nothing. -Framework helpers also attach ``execution_memory_type``. It identifies the -framework that owns the function body even when ``input_type`` or -``output_type`` is overridden. Host runtimes use that declaration when scoping -callable execution; boundary types alone do not prove execution ownership. +The library decorators also set ``execution_memory_type``: the library the +function body is written for, even when ``input_type`` or ``output_type`` +names another one. -Metadata key API ----------------- +Attribute names +--------------- -``MemoryContractAttribute`` is the public owner of the callable metadata keys: -``INPUT`` maps to ``input_memory_type``, ``OUTPUT`` maps to -``output_memory_type``, and ``EXECUTION`` maps to ``execution_memory_type``. -Each member provides ``read(namespace, default=None)`` and -``write(namespace, value)`` for object and mapping namespaces. Host libraries -can consume these members without maintaining a duplicate key registry. +``MemoryContractAttribute`` names the attributes: ``INPUT`` is +``input_memory_type``, ``OUTPUT`` is ``output_memory_type`` and ``EXECUTION`` +is ``execution_memory_type``. Each member has ``read(namespace, default=None)`` +and ``write(namespace, value)`` for objects and mappings. -The decorator does not accept ``gpu_id`` or ``clear_cuda_cache``. Device -selection belongs to explicit conversion or the host runtime. A non-callable +The decorators take no device argument. The device is chosen by the +conversion that produces the input or by the application. A non-callable ``contract`` is stored as declarative metadata; a callable contract validates the returned value. @@ -54,6 +56,6 @@ Direct calls default to preserving the input dtype. Hosts can pass an object implementing ``DtypeConversionConfig`` to select native or explicit output dtype behavior. -``wrap_dtype_preserving_callable(func, memory_type)`` exposes the dtype and -slice-control wrapper as a public integration boundary for host registries that -already own their memory and processing contracts. +``wrap_dtype_preserving_callable(func, mem_type)`` applies only the dtype and +slice wrapper, for applications that declare memory types themselves. +``mem_type`` is a ``MemoryType`` member. diff --git a/docs/source/examples/basic_conversion.rst b/docs/source/examples/basic_conversion.rst index efabe2c..cf47bde 100644 --- a/docs/source/examples/basic_conversion.rst +++ b/docs/source/examples/basic_conversion.rst @@ -5,12 +5,13 @@ Basic conversion import numpy as np - from arraybridge import convert_memory, detect_memory_type + from arraybridge import HOST, convert_memory, detect_memory_type source = np.arange(9).reshape(3, 3) source_type = detect_memory_type(source) - result = convert_memory(source, source_type, "numpy", 0) + result = convert_memory(source, source_type, "numpy", HOST) -For an optional target such as ``torch`` or ``cupy``, install that framework and -replace the target string. The returned value is allocated according to the -registered converter and requested device. +For an optional target, install that library, replace the target string and +pass a device the target has: ``convert_memory(source, source_type, "torch", +HOST)`` or ``convert_memory(source, source_type, "cupy", +AcceleratorDevice(0))``. diff --git a/docs/source/examples/multi_framework.rst b/docs/source/examples/multi_framework.rst index 8f4918a..cc1da17 100644 --- a/docs/source/examples/multi_framework.rst +++ b/docs/source/examples/multi_framework.rst @@ -1,18 +1,23 @@ -Multi-framework boundary -======================== +Converting between libraries +============================ -Keep conversion visible at the point that owns the boundary: +Convert at the place that knows the target library and device: .. code-block:: python - from arraybridge import convert_memory, detect_memory_type + import numpy as np - def as_framework(value, target, gpu_id): - source = detect_memory_type(value) - if source == target: - return value - return convert_memory(value, source, target, gpu_id) + from arraybridge import HOST, convert_memory, detect_memory_type -For a sequence of framework-specific functions, a host compiler should read -their declared input/output memory types and create one conversion plan. The -functions themselves should not each guess the incoming framework. + def as_library(value, target, device): + return convert_memory(value, detect_memory_type(value), target, device) + + value = np.arange(4) + assert as_library(value, "numpy", HOST) is value + +A conversion to the library and device the value is already on returns the +same object, so no check is needed first. + +For a sequence of functions written for different libraries, read their +declared input and output memory types and convert between steps. The +functions should not each guess the library of what they receive. diff --git a/docs/source/gpu_features.rst b/docs/source/gpu_features.rst index eeefd8b..8b87ebe 100644 --- a/docs/source/gpu_features.rst +++ b/docs/source/gpu_features.rst @@ -1,39 +1,72 @@ -GPU lifecycle -============= - -GPU-backed conversion uses the explicit ``gpu_id`` passed to -``convert_memory``, ``stack_slices``, or ``unstack_slices``. Framework-specific -``MemoryType`` declarations own device discovery, validation, scopes, movement, -and cache cleanup. Device identifiers are local to one framework. Equal numeric -identifiers from different frameworks do not establish physical-device -identity. - -Before importing TensorFlow or JAX, their declarations apply coexistence-safe -defaults with ``setdefault``: TensorFlow memory growth is enabled and JAX GPU -preallocation is disabled. An environment value supplied by the host is never -overwritten. Hosts that import optional frameworks themselves should call -``MemoryType.prepare_import()`` first. ArrayBridge warns when it encounters an -already-loaded framework whose import-time defaults were absent. - -``MemoryType.subprocess_environment()`` projects the same member-owned import -defaults into a child-process environment. The CuPy member also derives native -library search paths from installed NVIDIA wheels, so a fresh child interpreter -does not depend on another framework having loaded CUDA libraries first. - -JAX float64 output requires x64 mode to be enabled before import. ArrayBridge -raises instead of silently returning float32 when a caller requests float64 -while that capability is disabled. - -Decorated GPU callables may execute in a thread-local framework stream and may -retry supported OOM failures when ``oom_recovery=True``. Decorators do not -select a device and do not convert an incompatible input. - -``cleanup_all_gpu_frameworks(device_id=None)`` invokes allocator cleanup only -for loaded frameworks that declare a real cleanup leaf. Supplying ``device_id`` -targets that framework-local identifier. Cleanup never imports an absent -framework. JAX compilation-cache clearing is not presented as GPU-memory -cleanup. The application decides when values are no longer live and cleanup is -appropriate. - -ArrayBridge exposes mechanisms, not scheduling policy. Worker counts, device -assignment, retry scope, and concurrency limits belong to the host runtime. +Accelerators and memory +======================= + +Conversion and stacking use the ``device`` passed to ``convert_memory``, +``stack_slices`` or ``unstack_slices``. Each library class owns its device +discovery, checks, scopes, moves and cache cleanup. An accelerator index +belongs to one library: equal indices in two libraries need not be the same +hardware. + +Import defaults +--------------- + +Before importing TensorFlow or JAX, ArrayBridge sets defaults with +``setdefault`` so the libraries can share a GPU: TensorFlow memory growth is +enabled (``TF_FORCE_GPU_ALLOW_GROWTH=true``) and JAX preallocation is disabled +(``XLA_PYTHON_CLIENT_PREALLOCATE=false``). A value already in the environment +is not overwritten. An application that imports an optional library itself +should call ``library.prepare_import()`` first, for example +``ArrayLibrary.named("tensorflow").prepare_import()``. ArrayBridge logs a +warning when it meets an already imported library whose defaults were not set. + +``ArrayLibrary.child_environment(environment=None)`` returns an environment +with the same defaults for a child process. For CuPy it also adds the library +directories of installed NVIDIA wheels to the search path, so a new +interpreter does not depend on another library having loaded CUDA first. + +.. code-block:: python + + from arraybridge import ArrayLibrary + + environment = ArrayLibrary.child_environment({}) + assert environment["TF_FORCE_GPU_ALLOW_GROWTH"] == "true" + assert environment["XLA_PYTHON_CLIENT_PREALLOCATE"] == "false" + +JAX float64 output requires x64 mode to be enabled before import. Scaling a +JAX array to float64 raises ``ValueError`` when it is not. + +Decorated functions +------------------- + +A function decorated for a library with accelerators is wrapped only when +that library is installed and sees an accelerator. For CuPy and PyTorch the +call runs on a per-thread CUDA stream of the device the library currently has +selected. With ``oom_recovery=True`` (the default) a call that fails with the +library's out-of-memory error is retried up to two times; before each retry +CuPy and PyTorch release their cached memory on that device and Python's +garbage collector runs. Decorators do not choose a device and do not convert +an input of another library. + +The same pieces are methods of the library object: +``library.run_with_oom_recovery(func, max_retries=2, device=None)``, +``library.is_oom_error(error)`` and, on CuPy and PyTorch, +``library.thread_stream(device)``. + +Cleanup +------- + +``ArrayLibrary.cleanup_all_loaded(device=None)`` asks every library that is +already imported to release its cached device memory, on every accelerator or +on the given ``AcceleratorDevice``. ``library.cleanup_loaded(device=None)`` +does it for one library. CuPy and PyTorch release caches; the other libraries +have nothing to release. Cleanup never imports a library. The application +decides when values are no longer needed. + +.. code-block:: python + + from arraybridge import ArrayLibrary + + ArrayLibrary.cleanup_all_loaded() + +ArrayBridge provides these mechanisms and no scheduling. Worker counts, device +assignment and concurrency limits belong to the application. diff --git a/docs/source/index.rst b/docs/source/index.rst index 575ebf3..e83ab60 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -1,9 +1,11 @@ arraybridge =========== -ArrayBridge owns explicit array-framework conversion, memory-type declarations, -stack utilities, dtype policy, GPU stream/OOM helpers, and cleanup for scientific -applications. +ArrayBridge converts arrays between NumPy, CuPy, PyTorch, TensorFlow, JAX and +pyclesperanto. Each library is one class that answers for its devices, moves +and converts its arrays, and runs array operations natively. The package also +provides memory-type decorators, stack utilities, a dtype policy, per-thread +CUDA streams, out-of-memory retries and cache cleanup. .. toctree:: :maxdepth: 2 @@ -23,13 +25,12 @@ applications. release_035_validation release_036_validation -Boundary rule -------------- +What converts and what does not +------------------------------- -Conversion is explicit. ``convert_memory`` moves values between frameworks. -Decorators declare callable input/output memory types and runtime policies; they -do not perform cross-framework conversion on direct calls. +``convert_memory`` moves values between libraries. Decorators declare a +function's input and output memory types; calling a decorated function does +not convert anything. -ArrayBridge does not own a host application's compiler, function catalog, or -resource scheduler. Hosts consume the declarations and plan conversions/devices -at their own boundary. +ArrayBridge does not plan pipelines or schedule devices. An application reads +the declared types and decides where to convert and on which device. diff --git a/docs/source/installation.rst b/docs/source/installation.rst index ec3800f..4883a91 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -45,14 +45,14 @@ Verify the required NumPy path without optional frameworks: import numpy as np - from arraybridge import convert_memory, detect_memory_type + from arraybridge import HOST, convert_memory, detect_memory_type value = np.arange(3) assert detect_memory_type(value) == "numpy" - converted = convert_memory(value, "numpy", "numpy", gpu_id=0) + converted = convert_memory(value, "numpy", "numpy", HOST) -An optional-framework conversion fails when that framework is not installed. -ArrayBridge does not fall back to another target. +A conversion to an optional library raises ``MemoryConversionError`` when that +library is not installed. ArrayBridge does not substitute another target. See the `repository `_ for supported version constraints and current workflows. diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst index 32a6580..f4c8e34 100644 --- a/docs/source/quickstart.rst +++ b/docs/source/quickstart.rst @@ -8,7 +8,7 @@ Detection and conversion import numpy as np - from arraybridge import convert_memory, detect_memory_type + from arraybridge import HOST, convert_memory, detect_memory_type value = np.arange(6).reshape(2, 3) source = detect_memory_type(value) @@ -16,12 +16,14 @@ Detection and conversion value, source_type=source, target_type="numpy", - gpu_id=0, + device=HOST, ) ``detect_memory_type`` returns a string such as ``"numpy"`` or ``"torch"``. -Unknown array classes fail loudly. ``convert_memory`` requires all four -arguments; use a real GPU id when the target framework is GPU-backed. +An array of an unknown class raises ``ValueError``. ``convert_memory`` +requires all four arguments. ``device`` is a device of the target library: +``HOST`` for NumPy, ``AcceleratorDevice(i)`` for CuPy and pyclesperanto, either +for PyTorch, JAX and TensorFlow. Callable declarations --------------------- @@ -37,18 +39,20 @@ Callable declarations assert double.input_memory_type == "numpy" assert double.output_memory_type == "numpy" -The decorator expects the caller to supply a compatible array. A host runtime -may read the metadata and call ``convert_memory`` before invocation. +The decorator expects the caller to pass an array of the declared library. An +application may read the attributes and call ``convert_memory`` first. Stacks ------ .. code-block:: python - from arraybridge import stack_slices, unstack_slices + import numpy as np + + from arraybridge import HOST, stack_slices, unstack_slices slices = [np.zeros((4, 4)), np.ones((4, 4))] - volume = stack_slices(slices, memory_type="numpy", gpu_id=0) - restored = unstack_slices(volume, memory_type="numpy", gpu_id=0) + volume = stack_slices(slices, memory_type="numpy", device=HOST) + restored = unstack_slices(volume, memory_type="numpy", device=HOST) See :doc:`decorators`, :doc:`converters`, and :doc:`stack_utils`. diff --git a/docs/source/stack_utils.rst b/docs/source/stack_utils.rst index 7c83ef7..a75117e 100644 --- a/docs/source/stack_utils.rst +++ b/docs/source/stack_utils.rst @@ -11,15 +11,19 @@ and returns its 2D planes. import numpy as np - from arraybridge import stack_slices, unstack_slices + from arraybridge import HOST, stack_slices, unstack_slices planes = [np.zeros((3, 5)), np.ones((3, 5))] - stack = stack_slices(planes, memory_type="numpy", gpu_id=0) - planes_again = unstack_slices(stack, memory_type="numpy", gpu_id=0) - -Both functions require ``memory_type`` and ``gpu_id``. They validate dimensions -and fail on empty/malformed inputs. A host should declare what the stack axis -means; ArrayBridge owns only array shape, conversion, and device mechanics. - -``process_slices`` is the lower-level helper used by decorated functions when -``slice_by_slice=True``. + stack = stack_slices(planes, memory_type="numpy", device=HOST) + planes_again = unstack_slices(stack, memory_type="numpy", device=HOST) + +Both functions require ``memory_type`` and ``device``, a device of the target +library (see :doc:`converters`). They raise ``ValueError`` for an empty list, +a wrong number of dimensions or a device the target library does not have. +What the stack axis means is the application's knowledge. + +``process_slices(image, func, args, kwargs, device=None)`` is the helper +decorated functions use when ``slice_by_slice=True``. It unstacks on the +device holding ``image`` unless ``device`` is given, calls ``func`` on each +plane, and stacks the results in the library and on the device of the first +result. diff --git a/docs/source/user_guide.rst b/docs/source/user_guide.rst index c4c0d8a..a166a81 100644 --- a/docs/source/user_guide.rst +++ b/docs/source/user_guide.rst @@ -1,35 +1,47 @@ User guide ========== -Choose the operation by ownership: +Choose the operation by what it does: ``detect_memory_type`` - Identify the framework for an existing supported value. + Name the library that made an array. ``convert_memory`` - Convert one value from an explicit source type to an explicit target type on - the requested device. + Convert an array from a stated source library to a stated target library on + a stated device of the target. -Framework decorator - Declare a callable's input/output memory contract and apply dtype, slice, and - supported stream/OOM wrappers. It does not move the value between frameworks. +Library decorator (``numpy``, ``torch``, ...) + Declare a function's input and output memory types and apply the dtype, + slice, stream and out-of-memory wrappers. It does not move the value between + libraries. ``stack_slices`` / ``unstack_slices`` - Validate and convert plane collections at an explicit stack boundary. + Check, convert and stack 2D planes, or split a 3D array into planes. -``cleanup_all_gpu_frameworks`` - Ask every loaded supported GPU framework to release caches after an - application-owned lifecycle boundary. +``ArrayLibrary.named(name)`` / ``ArrayLibrary.of(array)`` + Get the library object: its devices, its array operations, its cleanup. -Framework names ---------------- +``ArrayLibrary.cleanup_all_loaded(device=None)`` + Ask every imported library to release its cached device memory. -The public conversion and stack APIs use string values: ``numpy``, ``cupy``, -``torch``, ``tensorflow``, ``jax``, and ``pyclesperanto``. ``convert_memory`` -accepts the canonical ``MemoryType`` member at both its source and target -boundaries. Availability depends on optional dependencies and hardware. -ArrayBridge never silently substitutes a different framework for an invalid -target. +``ConversionRoute.planned`` / ``ConversionCounters`` + Tell which mechanism a conversion will use; count the conversions that ran. -Applications should carry these values through typed configuration or plans, -not infer them from function names. +Library names +------------- + +The conversion and stack functions take the string values ``numpy``, ``cupy``, +``torch``, ``tensorflow``, ``jax`` and ``pyclesperanto``. ``convert_memory`` +also accepts the ``MemoryType`` member for source and target. What is usable +depends on the installed packages and the hardware. ArrayBridge does not +substitute another library for a target that is not usable. + +``MemoryType`` is the set of registered names as an ``Enum``. Use it where a +name is stored or offered as a choice, and ``MemoryType.X.library`` for +anything the library does. + +Devices +------- + +Every ``device`` argument is ``HOST``, ``AcceleratorDevice(index)`` or +``NamedDevice(name)`` from ``arraybridge``. See :doc:`api_reference`. diff --git a/paper.md b/paper.md index 08b8400..7fa58ab 100644 --- a/paper.md +++ b/paper.md @@ -73,25 +73,33 @@ Writing correct multi-framework code requires understanding all these difference # Software Design -The architecture uses data-driven metaprogramming. All framework-specific behavior is declared in `_FRAMEWORK_CONFIG`: +The architecture uses data-driven metaprogramming. All framework-specific behavior is declared in one `ArrayLibrary` subclass per framework (abridged): ```python -_FRAMEWORK_CONFIG = { - MemoryType.CUPY: { - "conversion_ops": { - "to_numpy": "data.get()", - "from_numpy": "({mod}.cuda.Device(gpu_id), {mod}.array(data))[1]", - "from_dlpack": "{mod}.from_dlpack(data)", - }, - "oom_exception_types": ["{mod}.cuda.memory.OutOfMemoryError"], - "oom_clear_cache": "{mod}.get_default_memory_pool().free_all_blocks()", - "stream_context": "{mod}.cuda.Stream()", - }, - # Similar entries for TORCH, TENSORFLOW, JAX, PYCLESPERANTO, NUMPY -} +class CupyLibrary(CudaLibrary, DLPackImport, DLPackProtocolExport, + NvidiaWheelLibraries, MutableArrays): + name = "cupy" + oom_exception_paths = ( + ("cuda", "memory", "OutOfMemoryError"), + ("cuda", "runtime", "CUDARuntimeError"), + ) + + def _to_host(self, data): + return data.get() + + def _upload(self, data, device): + return self.module.array(data) + + def _import_dlpack(self, payload): + return self.module.from_dlpack(payload.source) + + def _free_cached_memory(self): + self.module.get_default_memory_pool().free_all_blocks() + +# Similar classes for TORCH, TENSORFLOW, JAX, PYCLESPERANTO, NUMPY ``` -At import time, converter classes are generated dynamically via `AutoRegisterMeta` from the `metaclass-registry` library [@metaclassregistry]. Each converter implements `to_numpy()`, `from_numpy()`, `from_dlpack()`, and `to_X()` methods for all target frameworks. Adding a seventh framework requires only adding its entry to `_FRAMEWORK_CONFIG`—no new code paths. +At import time, the library classes are registered via `AutoRegisterMeta` from the `metaclass-registry` library [@metaclassregistry]. Each library implements `to_numpy()`, `from_numpy()`, `from_dlpack()`, and `convert_to()` for all target frameworks. Adding a seventh framework requires only adding its `ArrayLibrary` subclass—no new code paths. **Thread-local GPU streams** are managed via `threading.local()`. Without this, multiple threads sharing a GPU would serialize on the default stream or corrupt each other's operations. With per-thread streams, true parallel GPU execution is possible: @@ -104,14 +112,17 @@ def segment_image(image): **OOM recovery** (enabled by default) unifies detection across frameworks. The library checks both exception types and error string patterns (e.g., "out of memory", "resource_exhausted"), clears framework-specific caches, and retries: ```python -def _execute_with_oom_recovery(func, memory_type, max_retries=2): - for attempt in range(max_retries + 1): +def run_with_oom_recovery(self, func, max_retries=2, device=None): + attempt = 0 + while True: try: return func() - except Exception as e: - if not _is_oom_error(e, memory_type) or attempt == max_retries: + except Exception as error: + if not self.is_oom_error(error) or attempt == max_retries: raise - _clear_cache_for_memory_type(memory_type) + attempt += 1 + self.cleanup_loaded(device) + gc.collect() ``` This is critical for microscopy pipelines where image sizes are unpredictable. A 4K image might fit in GPU memory, but a 16K image might not. With OOM recovery, the pipeline automatically falls back to CPU processing without user intervention. @@ -125,7 +136,7 @@ This is critical for microscopy pipelines where image sizes are unpredictable. A result = gaussian_filter(image) # Still uint16 ``` -The `SCALING_FUNCTIONS` registry applies framework-specific scaling and clamping to convert back to the original dtype without overflow. This preserves the user's mental model—they think in terms of their original data type, not the intermediate float32 representation. +Each library's `scale_dtype` method applies framework-specific scaling and clamping to convert back to the original dtype without overflow. This preserves the user's mental model—they think in terms of their original data type, not the intermediate float32 representation. # Research Impact Statement diff --git a/pyproject.toml b/pyproject.toml index 61d0176..d53a490 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ keywords = ["numpy", "cupy", "pytorch", "tensorflow", "jax", "gpu", "array", "te dependencies = [ "numpy>=1.20", - "metaclass-registry>=0.1.5", + "metaclass-registry>=0.4.0", "python-introspect>=0.1.11", ] @@ -133,7 +133,10 @@ select = ["E", "F", "I", "N", "W", "UP"] python_version = "3.10" follow_imports_for_stubs = true ignore_missing_imports = true -warn_return_any = true +# Optional array libraries are not importable by the checker, so what their +# modules return is untyped; those values are admitted by Array registration +# at runtime (see arraybridge.libraries). +warn_return_any = false warn_unused_configs = true disallow_untyped_defs = false # Too strict for scientific code diff --git a/src/arraybridge/__init__.py b/src/arraybridge/__init__.py index bf3496c..99d79d9 100644 --- a/src/arraybridge/__init__.py +++ b/src/arraybridge/__init__.py @@ -10,16 +10,24 @@ NumPy/numcodecs import cost. """ -__version__ = "0.4.1" +__version__ = "0.5.0" _LAZY_EXPORTS: dict[str, str] = { "MemoryType": ".types", + "ArrayLibrary": ".libraries", + "Array": ".libraries", + "Device": ".devices", + "HostDevice": ".devices", + "AcceleratorDevice": ".devices", + "NamedDevice": ".devices", + "HOST": ".devices", + "ConversionRoute": ".conversion", + "ConversionRecord": ".conversion", + "ConversionObserver": ".conversion", + "ConversionCounters": ".conversion", "MemoryContractAttribute": ".types", "ArrayPayload": ".array_payload", "ArrayGeometry": ".array_geometry", - "CPU_MEMORY_TYPES": ".types", - "GPU_MEMORY_TYPES": ".types", - "SUPPORTED_MEMORY_TYPES": ".types", "convert_memory": ".converters", "detect_memory_type": ".converters", "memory_types": ".decorators", @@ -29,81 +37,34 @@ "stack_slices": ".stack_utils", "unstack_slices": ".stack_utils", "process_slices": ".slice_processing", - "cleanup_all_gpu_frameworks": ".gpu_cleanup", "MemoryConversionError": ".exceptions", - "SCALING_FUNCTIONS": ".dtype_scaling", - "_FRAMEWORK_CONFIG": ".framework_config", - "_FRAMEWORK_OPS": ".framework_ops", - "_execute_with_oom_recovery": ".oom_recovery", "_ensure_module": ".utils", "_supports_dlpack": ".utils", - "_get_device_id": ".utils", - # Decorator exports named after each memory type (numpy, cupy, torch, ...). - "numpy": ".decorators", - "cupy": ".decorators", - "torch": ".decorators", - "tensorflow": ".decorators", - "jax": ".decorators", - "pyclesperanto": ".decorators", } -__all__ = [ - # Types - "MemoryType", - "MemoryContractAttribute", - "ArrayPayload", - "ArrayGeometry", - "CPU_MEMORY_TYPES", - "GPU_MEMORY_TYPES", - "SUPPORTED_MEMORY_TYPES", - # Converters - "convert_memory", - "detect_memory_type", - # Decorators - "memory_types", - "DtypeConversion", - "SliceBySliceRuntimeParameter", - "wrap_dtype_preserving_callable", - # Stack utilities - "stack_slices", - "unstack_slices", - # Slice processing - "process_slices", - # GPU cleanup - "cleanup_all_gpu_frameworks", - # Exceptions - "MemoryConversionError", - # Scaling - "SCALING_FUNCTIONS", - # Framework config (internal but needed by some consumers) - "_FRAMEWORK_CONFIG", - "_FRAMEWORK_OPS", - # OOM recovery - "_execute_with_oom_recovery", - # Utils - "_ensure_module", - "_supports_dlpack", - "_get_device_id", - "numpy", - "cupy", - "torch", - "tensorflow", - "jax", - "pyclesperanto", -] +def _library_names() -> list[str]: + """The registered library names; each is also the name of its decorator.""" + from arraybridge.libraries import ArrayLibrary + + return list(ArrayLibrary.__registry__) -def __getattr__(name: str): - module_name = _LAZY_EXPORTS.get(name) - if module_name is None: - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") +def __getattr__(name: str): import importlib + if name == "__all__": + # Read on demand so that a library registered later is exported too. + return [*_LAZY_EXPORTS, *_library_names()] + # A name that is not a fixed export is the decorator of a registered library. + module_name = _LAZY_EXPORTS.get(name, ".decorators") + if name.startswith("_") and name not in _LAZY_EXPORTS: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(importlib.import_module(module_name, __name__), name) globals()[name] = value return value def __dir__() -> list[str]: - return sorted(set(globals()) | set(_LAZY_EXPORTS)) + return sorted(set(globals()) | set(_LAZY_EXPORTS) | set(_library_names())) diff --git a/src/arraybridge/array_geometry.py b/src/arraybridge/array_geometry.py index 6014c96..caa153b 100644 --- a/src/arraybridge/array_geometry.py +++ b/src/arraybridge/array_geometry.py @@ -3,7 +3,6 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any import numpy as np @@ -23,7 +22,7 @@ def ndim(self) -> int: return len(self.shape) @classmethod - def from_value(cls, value: Any) -> ArrayGeometry | None: + def from_value(cls, value: object) -> ArrayGeometry | None: """Inspect an array or nominal payload without forcing host conversion.""" data = value.array_payload_data() if isinstance(value, ArrayPayload) else value @@ -43,7 +42,7 @@ def from_value(cls, value: Any) -> ArrayGeometry | None: @classmethod def require_from_value( cls, - value: Any, + value: object, *, value_name: str = "Value", ) -> ArrayGeometry: diff --git a/src/arraybridge/array_operations.py b/src/arraybridge/array_operations.py deleted file mode 100644 index a628c0f..0000000 --- a/src/arraybridge/array_operations.py +++ /dev/null @@ -1,697 +0,0 @@ -"""Native array operations, one family leaf per MemoryType declaration.""" - -from __future__ import annotations - -from abc import abstractmethod -from collections.abc import Mapping, Sequence -from functools import cache, cached_property -from typing import TYPE_CHECKING, Any, ClassVar - -import numpy as np -from metaclass_registry import AutoRegisterMeta - -from arraybridge.array_payload import ArrayPayload - -if TYPE_CHECKING: - from arraybridge.types import MemoryType - -_SCALING_RANGES: dict[str, float | tuple[float, float]] = { - "uint8": 255.0, - "uint16": 65535.0, - "uint32": 4294967295.0, - "int16": (65535.0, 32768.0), - "int32": (4294967295.0, 2147483648.0), -} - - -def _dtype_name(dtype: Any) -> str: - declared_name = getattr(dtype, "name", None) - if declared_name is not None: - return str(declared_name) - return getattr(dtype, "__name__", str(dtype).rsplit(".", maxsplit=1)[-1]) - - -def _numpy_dtype_name(dtype: Any) -> str: - return str(np.dtype(dtype).name) - - -def _torch_dtype_name(dtype: Any) -> str: - return str(dtype).rsplit(".", maxsplit=1)[-1] - - -def _tensorflow_dtype_name(dtype: Any) -> str: - numpy_dtype = getattr(dtype, "as_numpy_dtype", dtype) - return str(np.dtype(numpy_dtype).name) - - -def _scaled_values(result: Any, result_min: Any, result_max: Any, target_dtype: Any) -> Any: - normalized = (result - result_min) / (result_max - result_min) - range_info = _SCALING_RANGES.get(_dtype_name(target_dtype)) - if range_info is None: - return normalized - if isinstance(range_info, tuple): - scale, offset = range_info - return normalized * scale - offset - return normalized * range_info - - -def _clamp_bounds(target_dtype: Any) -> tuple[float, float] | None: - range_info = _SCALING_RANGES.get(_dtype_name(target_dtype)) - if range_info is None: - return None - if isinstance(range_info, tuple): - scale, offset = range_info - return -offset, scale - offset - 128 - return 0, range_info - - -def _mapped_dtype(target_dtype: Any, module: Any) -> Any: - try: - dtype_name = np.dtype(target_dtype).name - except TypeError as error: - raise TypeError(f"Unsupported target dtype {target_dtype!r}") from error - mapped = getattr(module, dtype_name, None) - if mapped is None: - module_name = getattr(module, "__name__", type(module).__name__) - raise TypeError(f"{module_name} does not expose dtype {dtype_name}") - return mapped - - -@cache -def _leaf_instance(memory_type_value: str) -> ArrayOperations: - return ArrayOperations.__registry__[memory_type_value]() - - -class ArrayOperations(metaclass=AutoRegisterMeta): - """Native array semantics, one registered leaf per ``MemoryType`` value. - - Static leaves take the framework module from their caller (``MemoryType``). - Instance primitives are the namespace array-library-neutral code is written - against: each leaf supplies only what its library genuinely does differently. - """ - - __registry_key__ = "memory_type" - # Leaves are declared in this module, without lazy discovery. - __registry__: ClassVar[Mapping[str, type[ArrayOperations]]] = {} - memory_type: ClassVar[str | None] = None - - @classmethod - def for_memory(cls, memory_type: MemoryType) -> ArrayOperations: - """Return the operations leaf declared for ``memory_type``.""" - return _leaf_instance(memory_type.value) - - @cached_property - def module(self) -> Any: - """The framework module, imported on first use.""" - from arraybridge.types import MemoryType - - # Only registered leaves are instantiated, and each declares memory_type. - assert self.memory_type is not None - return MemoryType(self.memory_type).import_module() - - @property - def array_module(self) -> Any: - """The module holding the framework's array functions.""" - return self.module - - # ----- primitives for array-library-neutral operations ----- - - def accepts(self, value: Any) -> bool: - """Whether ``value`` is a native array of this framework.""" - return isinstance(value, self.module.ndarray) - - def astype(self, data: Any, dtype: Any) -> Any: - return self.cast(data, dtype, self.module) - - def max(self, data: Any, axis: int | tuple[int, ...]) -> Any: - return self.array_module.max(data, axis=axis) - - def min(self, data: Any, axis: int | tuple[int, ...]) -> Any: - return self.array_module.min(data, axis=axis) - - def mean(self, data: Any, axis: int | tuple[int, ...]) -> Any: - return self.array_module.mean(data, axis=axis) - - def sum(self, data: Any, axis: int | tuple[int, ...]) -> Any: - return self.array_module.sum(data, axis=axis) - - def prepend_axis(self, data: Any) -> Any: - """Add a leading axis of length one.""" - return self.reshape(data, (1, *data.shape), self.module) - - def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> Any: - return self.array_module.linspace(start, stop, num, endpoint=endpoint) - - def ones(self, shape: int | tuple[int, ...], dtype: Any) -> Any: - """Allocate ones on the framework's default device.""" - return self.array_module.ones(shape, dtype=_mapped_dtype(dtype, self.array_module)) - - def assign(self, data: Any, index: slice, values: Any) -> Any: - """Return ``data`` with ``data[index]`` replaced by ``values``.""" - data[index] = values - return data - - def outer(self, left: Any, right: Any) -> Any: - return self.array_module.outer(left, right) - - def floor(self, value: float) -> Any: - """Floor a host scalar in the framework's default precision.""" - return self.array_module.floor(value) - - @staticmethod - @abstractmethod - def to_numpy(data: Any, module: Any) -> Any: - """Project pixels to the explicit host boundary.""" - - @staticmethod - @abstractmethod - def from_numpy(data: Any, module: Any, device_id: int) -> Any: - """Admit host pixels on the selected framework device.""" - - @staticmethod - @abstractmethod - def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any: - """Apply the framework's existing intensity conversion law.""" - - @staticmethod - def dtype_name(dtype: Any) -> str: - return _numpy_dtype_name(dtype) - - @staticmethod - def stack(values: Sequence[Any], module: Any) -> Any: - return module.stack(tuple(values), axis=0) - - @staticmethod - def cast(data: Any, dtype: Any, module: Any) -> Any: - return data.astype(dtype, copy=False) - - @staticmethod - def logical_and(left: Any, right: Any, module: Any) -> Any: - return module.logical_and(left, right) - - @staticmethod - def reshape(data: Any, shape: tuple[int, ...], module: Any) -> Any: - return data.reshape(shape) - - @staticmethod - def broadcast_to(data: Any, shape: tuple[int, ...], module: Any) -> Any: - return module.broadcast_to(data, shape) - - @staticmethod - def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) -> Any: - return module.ones(shape, dtype=dtype) - - @classmethod - def normalize_planes( - cls, - data: Any, - dtype: Any, - scales: Sequence[float | None], - module: Any, - ) -> Any: - """Cast and scale each plane before assembling its owned output.""" - return cls.stack( - tuple( - ( - cls.cast(plane, dtype, module) - if scale is None - else cls.cast(plane, dtype, module) / float(scale) - ) - for plane, scale in zip(data, scales, strict=True) - ), - module, - ) - - -class MutableArrayOperations(ArrayOperations): - """Arrays whose allocated output supports native in-place arithmetic.""" - - @classmethod - def normalize_planes( - cls, - data: Any, - dtype: Any, - scales: Sequence[float | None], - module: Any, - ) -> Any: - # Integer division promotes the output dtype in the original recipe. - if not np.issubdtype(np.dtype(dtype), np.inexact) or len(data) == 0: - return super().normalize_planes(data, dtype, scales, module) - normalized = module.array(data, dtype=dtype, copy=True) - for index, scale in enumerate(scales): - if scale is not None: - normalized[index] /= float(scale) - return normalized - - -class NumpyArrayOperations(MutableArrayOperations): - """Numpy native operation leaves.""" - - memory_type = "numpy" - - def accepts(self, value: Any) -> bool: - return isinstance(value, (np.ndarray, ArrayPayload)) - - @staticmethod - def to_numpy(data: Any, module: Any) -> Any: - del module - return data - - @staticmethod - def from_numpy(data: Any, module: Any, device_id: int) -> Any: - del module, device_id - return data - - @staticmethod - def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any: - if not hasattr(result, "dtype"): - return result - if not ( - module.issubdtype(result.dtype, module.floating) - and module.issubdtype(target_dtype, module.integer) - ): - return result.astype(target_dtype) - result_min = result.min() - result_max = result.max() - if result_max <= result_min: - return result.astype(target_dtype) - scaled = _scaled_values(result, result_min, result_max, target_dtype) - bounds = _clamp_bounds(target_dtype) - if bounds is not None: - scaled = module.clip(scaled, *bounds) - return scaled.astype(target_dtype) - - -class CupyArrayOperations(MutableArrayOperations): - """Cupy native operation leaves.""" - - memory_type = "cupy" - - @staticmethod - def to_numpy(data: Any, module: Any) -> Any: - del module - return data.get() - - @staticmethod - def from_numpy(data: Any, module: Any, device_id: int) -> Any: - del device_id - return module.array(data) - - @staticmethod - def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any: - if not hasattr(result, "dtype"): - return result - if not ( - module.issubdtype(result.dtype, module.floating) - and not module.issubdtype(target_dtype, module.floating) - ): - return result.astype(target_dtype) - result_min = module.min(result) - result_max = module.max(result) - if result_max <= result_min: - return result.astype(target_dtype) - scaled = _scaled_values(result, result_min, result_max, target_dtype) - bounds = _clamp_bounds(target_dtype) - if bounds is not None: - scaled = module.clip(scaled, *bounds) - return scaled.astype(target_dtype) - - -class TorchArrayOperations(ArrayOperations): - """Torch native operation leaves.""" - - memory_type = "torch" - - def accepts(self, value: Any) -> bool: - return isinstance(value, self.module.Tensor) - - def max(self, data: Any, axis: int | tuple[int, ...]) -> Any: - # CUDA has no uint16 reduction; float32 holds every uint16 exactly. - values = data.float() if data.dtype == self.module.uint16 else data - return self.module.amax(values, dim=axis).to(data.dtype) - - def min(self, data: Any, axis: int | tuple[int, ...]) -> Any: - values = data.float() if data.dtype == self.module.uint16 else data - return self.module.amin(values, dim=axis).to(data.dtype) - - def mean(self, data: Any, axis: int | tuple[int, ...]) -> Any: - return self.module.mean(data.float(), dim=axis) - - def sum(self, data: Any, axis: int | tuple[int, ...]) -> Any: - return self.module.sum(data, dim=axis) - - def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> Any: - float32 = self.module.float32 - if endpoint: - return self.module.linspace(start, stop, num, dtype=float32) - return self.module.linspace(start, stop, num + 1, dtype=float32)[:-1] - - def floor(self, value: float) -> Any: - return self.module.floor(self.module.tensor(value)) - - @staticmethod - def to_numpy(data: Any, module: Any) -> Any: - del module - return data.cpu().numpy() - - @staticmethod - def from_numpy(data: Any, module: Any, device_id: int) -> Any: - host_data = ( - np.ascontiguousarray(data) - if any(stride < 0 for stride in getattr(data, "strides", ())) - else data - ) - return module.from_numpy(host_data).to(f"cuda:{device_id}") - - @staticmethod - def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any: - if not hasattr(result, "dtype"): - return result - mapped = _mapped_dtype(target_dtype, module) - floats = (module.float16, module.float32, module.float64) - if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)): - return result.to(mapped) - result_min = result.min() - result_max = result.max() - if result_max <= result_min: - return result.to(mapped) - scaled = _scaled_values(result, result_min, result_max, target_dtype) - bounds = _clamp_bounds(target_dtype) - if bounds is not None: - scaled = module.clamp(scaled, min=bounds[0], max=bounds[1]) - return scaled.to(mapped) - - @staticmethod - def stack(values: Sequence[Any], module: Any) -> Any: - return module.stack(tuple(values), dim=0) - - @staticmethod - def cast(data: Any, dtype: Any, module: Any) -> Any: - native = dtype if isinstance(dtype, module.dtype) else _mapped_dtype(dtype, module) - return data.to(dtype=native) - - @staticmethod - def dtype_name(dtype: Any) -> str: - return _torch_dtype_name(dtype) - - @staticmethod - def broadcast_to(data: Any, shape: tuple[int, ...], module: Any) -> Any: - return data.expand(shape) - - @staticmethod - def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) -> Any: - return module.ones(shape, dtype=_mapped_dtype(dtype, module), device=reference.device) - - -class TensorflowArrayOperations(ArrayOperations): - """Tensorflow native operation leaves.""" - - memory_type = "tensorflow" - - def accepts(self, value: Any) -> bool: - return isinstance(value, self.module.Tensor) - - def max(self, data: Any, axis: int | tuple[int, ...]) -> Any: - return self.module.reduce_max(data, axis=axis) - - def min(self, data: Any, axis: int | tuple[int, ...]) -> Any: - return self.module.reduce_min(data, axis=axis) - - def mean(self, data: Any, axis: int | tuple[int, ...]) -> Any: - return self.module.reduce_mean(self.module.cast(data, self.module.float32), axis=axis) - - def sum(self, data: Any, axis: int | tuple[int, ...]) -> Any: - return self.module.reduce_sum(data, axis=axis) - - def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> Any: - if endpoint: - return self.module.linspace(float(start), float(stop), num) - return self.module.linspace(float(start), float(stop), num + 1)[:-1] - - def assign(self, data: Any, index: slice, values: Any) -> Any: - rows = self.module.range(*index.indices(int(data.shape[0]))) - return self.module.tensor_scatter_nd_update( - data, self.module.reshape(rows, [-1, 1]), values - ) - - def outer(self, left: Any, right: Any) -> Any: - return self.module.tensordot(left, right, axes=0) - - def floor(self, value: float) -> Any: - return self.module.math.floor(value) - - @staticmethod - def dtype_name(dtype: Any) -> str: - return _tensorflow_dtype_name(dtype) - - @staticmethod - def to_numpy(data: Any, module: Any) -> Any: - del module - return data.numpy() - - @staticmethod - def from_numpy(data: Any, module: Any, device_id: int) -> Any: - del device_id - return module.convert_to_tensor(data) - - @staticmethod - def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any: - if not hasattr(result, "dtype"): - return result - mapped = _mapped_dtype(target_dtype, module) - floats = (module.float16, module.float32, module.float64) - if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)): - return module.cast(result, mapped) - result_min = module.reduce_min(result) - result_max = module.reduce_max(result) - if result_max <= result_min: - return module.cast(result, mapped) - scaled = _scaled_values(result, result_min, result_max, target_dtype) - bounds = _clamp_bounds(target_dtype) - if bounds is not None: - scaled = module.clip_by_value(scaled, *bounds) - return module.cast(scaled, mapped) - - @staticmethod - def cast(data: Any, dtype: Any, module: Any) -> Any: - native = dtype if isinstance(dtype, module.dtypes.DType) else _mapped_dtype(dtype, module) - return module.cast(data, native) - - @staticmethod - def reshape(data: Any, shape: tuple[int, ...], module: Any) -> Any: - return module.reshape(data, shape) - - @staticmethod - def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) -> Any: - with module.device(reference.device): - return module.ones(shape, dtype=_mapped_dtype(dtype, module)) - - -class JaxArrayOperations(ArrayOperations): - """Jax native operation leaves.""" - - memory_type = "jax" - - @property - def array_module(self) -> Any: - return self.module.numpy - - def accepts(self, value: Any) -> bool: - return isinstance(value, self.module.numpy.ndarray) - - def mean(self, data: Any, axis: int | tuple[int, ...]) -> Any: - jnp = self.module.numpy - return jnp.mean(data.astype(jnp.float32), axis=axis) - - def assign(self, data: Any, index: slice, values: Any) -> Any: - return data.at[index].set(values) - - @staticmethod - def to_numpy(data: Any, module: Any) -> Any: - del module - return np.asarray(data) - - @staticmethod - def from_numpy(data: Any, module: Any, device_id: int) -> Any: - devices = tuple(device for device in module.devices() if device.platform == "gpu") - return module.device_put(data, devices[device_id]) - - @staticmethod - def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any: - if not hasattr(result, "dtype"): - return result - if np.dtype(target_dtype) == np.dtype(np.float64): - x64_enabled = getattr(module.config, "x64_enabled", None) - if x64_enabled is None: - x64_enabled = module.config.read("jax_enable_x64") - if not x64_enabled: - raise ValueError( - "JAX float64 output requires x64 mode; set JAX_ENABLE_X64=true before import" - ) - jnp = module.numpy - mapped = _mapped_dtype(target_dtype, jnp) - floats = (jnp.float16, jnp.float32, jnp.float64) - if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)): - return result.astype(mapped) - result_min = jnp.min(result) - result_max = jnp.max(result) - if result_max <= result_min: - return result.astype(mapped) - scaled = _scaled_values(result, result_min, result_max, target_dtype) - bounds = _clamp_bounds(target_dtype) - if bounds is not None: - scaled = jnp.clip(scaled, *bounds) - return scaled.astype(mapped) - - @staticmethod - def stack(values: Sequence[Any], module: Any) -> Any: - return module.numpy.stack(tuple(values), axis=0) - - @staticmethod - def cast(data: Any, dtype: Any, module: Any) -> Any: - return data.astype(_mapped_dtype(dtype, module.numpy)) - - @staticmethod - def logical_and(left: Any, right: Any, module: Any) -> Any: - return module.numpy.logical_and(left, right) - - @staticmethod - def broadcast_to(data: Any, shape: tuple[int, ...], module: Any) -> Any: - return module.numpy.broadcast_to(data, shape) - - @staticmethod - def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) -> Any: - return module.numpy.ones(shape, dtype=dtype, device=reference.device) - - -class PyclesperantoArrayOperations(ArrayOperations): - """Pyclesperanto native operation leaves.""" - - memory_type = "pyclesperanto" - - def accepts(self, value: Any) -> bool: - # pyclesperanto kernels push any host array they receive. - return hasattr(value, "shape") - - def max(self, data: Any, axis: int | tuple[int, ...]) -> Any: - _require_z_axis(axis) - return self.module.maximum_z_projection(data) - - def min(self, data: Any, axis: int | tuple[int, ...]) -> Any: - _require_z_axis(axis) - return self.module.minimum_z_projection(data) - - def mean(self, data: Any, axis: int | tuple[int, ...]) -> Any: - _require_z_axis(axis) - return self.module.mean_z_projection(data) - - def sum(self, data: Any, axis: int | tuple[int, ...]) -> Any: - _require_z_axis(axis) - return self.module.sum_z_projection(data) - - def prepend_axis(self, data: Any) -> Any: - result = self.module.create((1, *data.shape), dtype=data.dtype) - result[0] = data - return result - - @staticmethod - def to_numpy(data: Any, module: Any) -> Any: - return module.pull(data) - - @staticmethod - def from_numpy(data: Any, module: Any, device_id: int) -> Any: - del device_id - return module.push(data) - - @staticmethod - def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any: - if not hasattr(result, "dtype"): - return result - target_is_int = np.issubdtype(np.dtype(target_dtype), np.integer) - if not (np.issubdtype(result.dtype, np.floating) and target_is_int): - return module.push(module.pull(result).astype(target_dtype)) - result_min = float(module.minimum_of_all_pixels(result)) - result_max = float(module.maximum_of_all_pixels(result)) - if result_max <= result_min: - return module.push(module.pull(result).astype(target_dtype)) - normalized = module.subtract_image_from_scalar(result, scalar=result_min) - normalized = module.multiply_image_and_scalar( - normalized, - scalar=1.0 / (result_max - result_min), - ) - range_info = _SCALING_RANGES.get(_dtype_name(target_dtype)) - if isinstance(range_info, tuple): - scale, offset = range_info - scaled = module.multiply_image_and_scalar(normalized, scalar=scale) - scaled = module.subtract_image_from_scalar(scaled, scalar=offset) - elif range_info is not None: - scaled = module.multiply_image_and_scalar(normalized, scalar=range_info) - else: - scaled = normalized - host_values = module.pull(scaled) - bounds = _clamp_bounds(target_dtype) - if bounds is not None: - host_values = np.clip(host_values, *bounds) - return module.push(host_values.astype(target_dtype)) - - @staticmethod - def stack(values: Sequence[Any], module: Any) -> Any: - if not values: - raise ValueError("Cannot stack an empty pyclesperanto sequence") - if len(values) == 1: - source = values[0] - result = module.create((1, *source.shape), dtype=source.dtype) - return module.copy_slice(source, result, 0) - result = values[0] - for value in values[1:]: - result = module.concatenate_along_z(result, value) - return result - - @staticmethod - def cast(data: Any, dtype: Any, module: Any) -> Any: - return module.push(module.pull(data).astype(dtype, copy=False)) - - @staticmethod - def logical_and(left: Any, right: Any, module: Any) -> Any: - return module.push(np.logical_and(module.pull(left), module.pull(right))) - - @staticmethod - def reshape(data: Any, shape: tuple[int, ...], module: Any) -> Any: - if tuple(data.shape) == shape: - return data - raise NotImplementedError( - "pyclesperanto does not provide native array reshape; its reshape downloads pixels" - ) - - @staticmethod - def broadcast_to(data: Any, shape: tuple[int, ...], module: Any) -> Any: - if tuple(data.shape) == shape: - return data - raise NotImplementedError("pyclesperanto does not provide native array broadcasting") - - @staticmethod - def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) -> Any: - if np.dtype(dtype).name not in { - "float32", - "int8", - "int16", - "int32", - "uint8", - "uint16", - "uint32", - }: - raise NotImplementedError(f"pyclesperanto cannot allocate native dtype {dtype!r}") - if not 1 <= len(shape) <= 3: - raise NotImplementedError( - "pyclesperanto native allocation requires one to three dimensions" - ) - result = module.create(shape, dtype=dtype, device=reference.device) - module.set(result, scalar=1, device=reference.device) - return result - - -def _require_z_axis(axis: int | tuple[int, ...]) -> None: - if axis != 0: - raise NotImplementedError( - f"pyclesperanto reduces only along the leading (z) axis, not axis {axis!r}" - ) diff --git a/src/arraybridge/array_payload.py b/src/arraybridge/array_payload.py index 6655d4c..5267de5 100644 --- a/src/arraybridge/array_payload.py +++ b/src/arraybridge/array_payload.py @@ -2,7 +2,9 @@ from abc import ABC, abstractmethod from collections.abc import Callable -from typing import Any, TypeVar +from typing import TypeVar + +from arraybridge.arrays import Array ArrayPayloadT = TypeVar("ArrayPayloadT", bound="ArrayPayload") @@ -11,16 +13,16 @@ class ArrayPayload(ABC): """Container whose array data can be transformed without losing its context.""" @abstractmethod - def array_payload_data(self) -> Any: + def array_payload_data(self) -> Array: """Return the concrete array governed by this payload.""" @abstractmethod - def with_data(self: ArrayPayloadT, data: Any) -> ArrayPayloadT: + def with_data(self: ArrayPayloadT, data: Array) -> ArrayPayloadT: """Return this payload's context attached to replacement array data.""" def map_array_payload( self: ArrayPayloadT, - transform: Callable[[Any], Any], + transform: Callable[[Array], Array], ) -> ArrayPayloadT: """Apply one array transform while preserving the payload container.""" diff --git a/src/arraybridge/arrays.py b/src/arraybridge/arrays.py new file mode 100644 index 0000000..3037728 --- /dev/null +++ b/src/arraybridge/arrays.py @@ -0,0 +1,12 @@ +"""The nominal type of an array of any supported library.""" + +from abc import ABC + + +class Array(ABC): + """An array of any supported library. + + Membership is by registration: each library's ``array_class`` is a + subclass, and a library's native array classes are registered with it when + the library first meets one (``ArrayLibrary.of``). + """ diff --git a/src/arraybridge/conversion.py b/src/arraybridge/conversion.py new file mode 100644 index 0000000..442758c --- /dev/null +++ b/src/arraybridge/conversion.py @@ -0,0 +1,369 @@ +"""How an array gets from one library and device to another, and what was done. + +``ConversionRoute`` leaves are tried in the order they are declared in this +module. ``applies`` reads only the two library declarations and the two device +descriptions, so the same test answers "which mechanism will be used" on a +machine without the device and selects the route when the conversion runs. +""" + +from __future__ import annotations + +import logging +import time +from abc import ABC, abstractmethod +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from dataclasses import asdict, dataclass, field +from typing import TYPE_CHECKING, ClassVar, TypeVar + +from metaclass_registry import AutoRegisterMeta + +from arraybridge.devices import Device + +if TYPE_CHECKING: + from arraybridge.libraries import ArrayLibrary + +logger = logging.getLogger(__name__) + +S = TypeVar("S") +T = TypeVar("T") + + +class RouteDeclinedError(Exception): + """A route that applied by declaration could not carry this particular array.""" + + +class ConversionRoute(ABC, metaclass=AutoRegisterMeta): + """One mechanism for moving an array between libraries and devices.""" + + __registry_key__ = "mechanism" + __registry__: ClassVar[Mapping[str, type[ConversionRoute]]] = {} + mechanism: ClassVar[str | None] = None + + @classmethod + @abstractmethod + def applies( + cls, + source: ArrayLibrary[S], + source_device: Device, + target: ArrayLibrary[T], + target_device: Device, + ) -> bool: + """Whether this route is declared for the pair. Touches no device.""" + + @classmethod + @abstractmethod + def convert( + cls, + data: S, + source: ArrayLibrary[S], + target: ArrayLibrary[T], + target_device: Device, + ) -> T: + """Carry ``data`` to ``target`` on ``target_device`` or raise ``RouteDeclinedError``.""" + + @classmethod + def candidates( + cls, + source: ArrayLibrary[S], + source_device: Device, + target: ArrayLibrary[T], + target_device: Device, + ) -> tuple[type[ConversionRoute], ...]: + """Every declared route for the pair, first choice first.""" + + return tuple( + route + for route in cls.__registry__.values() + if route.applies(source, source_device, target, target_device) + ) + + @classmethod + def planned( + cls, + source: ArrayLibrary[S], + source_device: Device, + target: ArrayLibrary[T], + target_device: Device, + ) -> type[ConversionRoute]: + """The route a conversion of this pair will try first.""" + + return cls.candidates(source, source_device, target, target_device)[0] + + @classmethod + def run( + cls, + data: S, + source: ArrayLibrary[S], + target: ArrayLibrary[T], + target_device: Device, + ) -> T: + """Convert through the first candidate route that carries ``data``.""" + + observer = ConversionObserver.current + started = time.perf_counter() if observer.wants_records else 0.0 + source_device = source.device_of(data) + declined: list[str] = [] + for route in cls.candidates(source, source_device, target, target_device): + try: + result = route.convert(data, source, target, target_device) + except RouteDeclinedError as reason: + logger.debug("%s declined %s to %s: %s", route.mechanism, source, target, reason) + declined.append(str(route.mechanism)) + continue + if observer.wants_records and result is not data: + observer.record( + ConversionRecord( + source_library=source.name, + source_device=source_device.label, + target_library=target.name, + target_device=target_device.label, + mechanism=str(route.mechanism), + nbytes=target.nbytes_of(result), + seconds=time.perf_counter() - started, + declined_mechanisms=tuple(declined), + ) + ) + return result + raise LookupError(f"No conversion route from {source} to {target}") + + +class DeviceMove(ConversionRoute): + """Same library: the array stays native and only changes device.""" + + mechanism = "device move" + + @classmethod + def applies(cls, source, source_device, target, target_device) -> bool: + return source is target + + @classmethod + def convert(cls, data, source, target, target_device): + return target.move_to_device(data, target_device) + + +class DLPackExchange(ConversionRoute): + """Both libraries speak DLPack: the target wraps the source's buffer.""" + + mechanism = "dlpack" + + @classmethod + def applies(cls, source, source_device, target, target_device) -> bool: + return target.imports_dlpack_from(source_device) and source.exports_dlpack_from( + source_device + ) + + @classmethod + def convert(cls, data, source, target, target_device): + try: + payload = source.export_dlpack(data) + if payload is None: + raise RouteDeclinedError("the source array has no DLPack export") + with target.device_scope(target_device): + result = target.from_dlpack(payload) + if result is NotImplemented: + raise RouteDeclinedError("the target cannot import this capsule") + return target.move_to_device(result, target_device) + except RouteDeclinedError: + raise + except Exception as error: + logger.warning( + "DLPack conversion from %s to %s failed: %s. Using CPU roundtrip.", + source.name, + target.name, + error, + ) + raise RouteDeclinedError(str(error)) from error + + +class HostCopy(ConversionRoute): + """Download to a NumPy array in main memory, then upload to the target.""" + + mechanism = "host copy" + + @classmethod + def applies(cls, source, source_device, target, target_device) -> bool: + return True + + @classmethod + def convert(cls, data, source, target, target_device): + return target.unrecorded_from_numpy(source.unrecorded_to_numpy(data), target_device) + + +@dataclass(frozen=True, slots=True) +class ConversionRecord: + """One conversion that produced a new array. Plain fields: pickles and dumps to JSON.""" + + source_library: str + source_device: str + target_library: str + target_device: str + mechanism: str + nbytes: int + seconds: float + declined_mechanisms: tuple[str, ...] = () + + def to_json(self) -> dict[str, str | int | float | list[str]]: + return {**asdict(self), "declined_mechanisms": list(self.declined_mechanisms)} + + @classmethod + def from_json(cls, row: Mapping[str, str | int | float | list[str]]) -> ConversionRecord: + declined = row["declined_mechanisms"] + assert isinstance(declined, list) + return cls( + source_library=str(row["source_library"]), + source_device=str(row["source_device"]), + target_library=str(row["target_library"]), + target_device=str(row["target_device"]), + mechanism=str(row["mechanism"]), + nbytes=int(str(row["nbytes"])), + seconds=float(str(row["seconds"])), + declined_mechanisms=tuple(declined), + ) + + +class ConversionObserver(ABC): + """Receives one ``ConversionRecord`` per conversion while it is installed. + + ``ConversionObserver.current`` is the one hook. Its default ignores + conversions and asks for no records, so nothing is timed or sized. + """ + + current: ClassVar[ConversionObserver] + wants_records: ClassVar[bool] = True + + @abstractmethod + def record(self, record: ConversionRecord) -> None: + """Take one finished conversion.""" + + @contextmanager + def installed(self) -> Iterator[ConversionObserver]: + """Make this the process's observer for the duration of a block.""" + + previous = ConversionObserver.current + ConversionObserver.current = self + try: + yield self + finally: + ConversionObserver.current = previous + + def install(self) -> ConversionObserver: + """Make this the process's observer and return the one it replaced.""" + + previous = ConversionObserver.current + ConversionObserver.current = self + return previous + + +class IgnoreConversions(ConversionObserver): + """The default observer: conversions are neither measured nor kept.""" + + wants_records = False + + def record(self, record: ConversionRecord) -> None: + del record + + +ConversionObserver.current = IgnoreConversions() + +ConversionKey = tuple[str, str, str, str, str] + + +@dataclass(slots=True) +class ConversionTotal: + """How many conversions of one kind happened, and their bytes and seconds.""" + + count: int = 0 + nbytes: int = 0 + seconds: float = 0.0 + + +@dataclass(slots=True) +class ConversionCounters(ConversionObserver): + """Per-process totals keyed by source, target and mechanism. + + Counters from worker processes are plain data: pickle them or send + ``to_json()`` back, then ``merge`` them into the parent's. + """ + + totals: dict[ConversionKey, ConversionTotal] = field(default_factory=dict) + + def record(self, record: ConversionRecord) -> None: + self._add( + ( + record.source_library, + record.source_device, + record.target_library, + record.target_device, + record.mechanism, + ), + 1, + record.nbytes, + record.seconds, + ) + + def _add(self, key: ConversionKey, count: int, nbytes: int, seconds: float) -> None: + total = self.totals.setdefault(key, ConversionTotal()) + total.count += count + total.nbytes += nbytes + total.seconds += seconds + + def merge(self, other: ConversionCounters) -> ConversionCounters: + """Add another process's totals into this one.""" + + for key, total in other.totals.items(): + self._add(key, total.count, total.nbytes, total.seconds) + return self + + def rows(self) -> list[dict[str, str | int | float]]: + """One JSON-ready row per kind of conversion, largest byte total first.""" + + return [ + { + "source_library": key[0], + "source_device": key[1], + "target_library": key[2], + "target_device": key[3], + "mechanism": key[4], + "count": total.count, + "nbytes": total.nbytes, + "seconds": total.seconds, + } + for key, total in sorted( + self.totals.items(), key=lambda item: (-item[1].nbytes, item[0]) + ) + ] + + to_json = rows + + @classmethod + def from_json(cls, rows: list[dict[str, str | int | float]]) -> ConversionCounters: + counters = cls() + for row in rows: + counters._add( + ( + str(row["source_library"]), + str(row["source_device"]), + str(row["target_library"]), + str(row["target_device"]), + str(row["mechanism"]), + ), + int(row["count"]), + int(row["nbytes"]), + float(row["seconds"]), + ) + return counters + + def format_table(self) -> str: + """The totals as aligned text, one line per kind of conversion.""" + + lines = [ + f"{'source':<28} {'target':<28} {'mechanism':<12} " + f"{'count':>7} {'MB':>10} {'seconds':>9}" + ] + for key, total in sorted(self.totals.items(), key=lambda item: (-item[1].nbytes, item[0])): + lines.append( + f"{key[0] + ' ' + key[1]:<28} {key[2] + ' ' + key[3]:<28} {key[4]:<12} " + f"{total.count:>7} {total.nbytes / 1e6:>10.1f} {total.seconds:>9.3f}" + ) + return "\n".join(lines) diff --git a/src/arraybridge/converters.py b/src/arraybridge/converters.py index 2e482ea..4bf399b 100644 --- a/src/arraybridge/converters.py +++ b/src/arraybridge/converters.py @@ -1,85 +1,54 @@ -"""Memory conversion public API for OpenHCS.""" +"""Convert an array from one library to another, and tell which library made one.""" -from typing import Any, cast - -import numpy as np +from __future__ import annotations +from arraybridge.devices import Device from arraybridge.exceptions import MemoryConversionError -from arraybridge.types import VALID_MEMORY_TYPES, MemoryType +from arraybridge.libraries import Array, ArrayLibrary +from arraybridge.types import MemoryType + + +def _library(memory_type: str | MemoryType, role: str) -> ArrayLibrary[Array]: + name = memory_type.value if isinstance(memory_type, MemoryType) else memory_type + try: + return ArrayLibrary.named(name) + except ValueError: + raise ValueError( + f"Invalid {role} '{name}'. Available types: {sorted(ArrayLibrary.__registry__)}" + ) from None def convert_memory( - data: Any, + data: Array, source_type: str | MemoryType, target_type: str | MemoryType, - gpu_id: int, -) -> Any: - """ - Convert data between memory types using the unified converter infrastructure. - - Args: - data: The data to convert - source_type: The source memory type (e.g., "numpy", "torch") - target_type: The target memory type (e.g., "cupy", "jax") - gpu_id: The target GPU device ID - - Returns: - The converted data in the target memory type + device: Device, +) -> Array: + """Convert ``data`` from the source library to the target library on ``device``. Raises: - ValueError: If source_type or target_type is invalid + ValueError: If source_type or target_type is not a registered library MemoryConversionError: If conversion fails """ - source_name = source_type.value if isinstance(source_type, MemoryType) else source_type - target_name = target_type.value if isinstance(target_type, MemoryType) else target_type - if source_name not in VALID_MEMORY_TYPES: - raise ValueError( - f"Invalid source_type '{source_name}'. Available types: {sorted(VALID_MEMORY_TYPES)}" - ) - if target_name not in VALID_MEMORY_TYPES: - raise ValueError( - f"Invalid target_type '{target_name}'. Available types: {sorted(VALID_MEMORY_TYPES)}" - ) - - source = MemoryType(source_name) - target = MemoryType(target_name) + source = _library(source_type, "source_type") + target = _library(target_type, "target_type") try: - return source.convert_to(data, target, gpu_id) + return source.convert_to(data, target, device) except MemoryConversionError: raise except Exception as error: raise MemoryConversionError( - source_type=source_name, - target_type=target_name, - method="MemoryType.convert_to", + source_type=str(source.name), + target_type=str(target.name), + method="ArrayLibrary.convert_to", reason=str(error), ) from error -def detect_memory_type(data: Any) -> str: - """ - Detect the memory type of data using framework config. - - Args: - data: The data to detect - - Returns: - The detected memory type string (e.g., "numpy", "torch") +def detect_memory_type(data: object) -> str: + """The name of the library that made ``data`` (``"numpy"``, ``"torch"``, ...). Raises: - ValueError: If memory type cannot be detected + ValueError: If no registered library made it """ - # NumPy special case (most common, check first) - if isinstance(data, np.ndarray): - return cast(str, MemoryType.NUMPY.value) - - # Check all frameworks using their module names from config - module_name = type(data).__module__ - - top_level = module_name.split(".")[0] - - for mem_type in MemoryType: - if top_level in mem_type.recognized_module_names: - return cast(str, mem_type.value) - - raise ValueError(f"Unknown memory type for {type(data)} (module: {module_name})") + return str(ArrayLibrary.of(data).name) diff --git a/src/arraybridge/converters_registry.py b/src/arraybridge/converters_registry.py deleted file mode 100644 index 1b82c3e..0000000 --- a/src/arraybridge/converters_registry.py +++ /dev/null @@ -1,174 +0,0 @@ -""" -Registry-based converter infrastructure using metaclass-registry. - -This module provides the ConverterBase class using AutoRegisterMeta, -concrete converter implementations for each framework, and a helper -function for registry lookups. -""" - -from abc import abstractmethod -from collections.abc import Mapping -from types import MappingProxyType -from typing import ClassVar - -from metaclass_registry import AutoRegisterMeta - -from arraybridge.types import MemoryType - - -class ConverterBase(metaclass=AutoRegisterMeta): - """Base class for memory type converters using auto-registration. - - Each concrete converter sets memory_type to register itself in the registry. - The registry key is the memory_type attribute (e.g., "numpy", "torch"). - """ - - __registry_key__ = "memory_type" - # Simple dict: converters are created in this module, without lazy discovery. - __registry__: ClassVar[Mapping[str, type["ConverterBase"]]] = {} - memory_type: str | None = None - - @abstractmethod - def to_numpy(self, data, gpu_id): - """Extract to NumPy (type-specific implementation).""" - pass - - @abstractmethod - def from_numpy(self, data, gpu_id): - """Create from NumPy (type-specific implementation).""" - pass - - @abstractmethod - def from_dlpack(self, data, gpu_id): - """Create from DLPack capsule (type-specific implementation).""" - pass - - @abstractmethod - def move_to_device(self, data, gpu_id): - """Move data to specified GPU device if needed (type-specific implementation).""" - pass - - -def _make_to_numpy(mem_type: MemoryType): - def to_numpy(self, data, gpu_id): - del self, gpu_id - return mem_type.to_numpy(data) - - to_numpy.__qualname__ = f"{mem_type.value.capitalize()}Converter.to_numpy" - return to_numpy - - -def _make_from_numpy(mem_type: MemoryType): - def from_numpy(self, data, gpu_id): - del self - return mem_type.from_numpy(data, gpu_id) - - from_numpy.__qualname__ = f"{mem_type.value.capitalize()}Converter.from_numpy" - return from_numpy - - -def _make_device_mover(mem_type: MemoryType): - """Create a converter adapter over declaration-owned device movement.""" - - def move_to_device(self, data, gpu_id): - del self - return mem_type.move_to_device(data, gpu_id) - - move_to_device.__qualname__ = f"{mem_type.value.capitalize()}Converter.move_to_device" - return move_to_device - - -def _make_dlpack_importer(mem_type: MemoryType): - """Create a converter adapter over declaration-owned DLPack import.""" - - def from_dlpack(self, data, gpu_id): - del self - module = mem_type.import_module() - with mem_type.device_scope(gpu_id, module): - return mem_type.from_dlpack(data, module) - - from_dlpack.__qualname__ = f"{mem_type.value.capitalize()}Converter.from_dlpack" - return from_dlpack - - -# Auto-generate converter classes for each memory type -def _create_converter_classes(): - """Create concrete converter classes for each memory type.""" - for mem_type in MemoryType: - class_attrs = { - "memory_type": mem_type.value, - "to_numpy": _make_to_numpy(mem_type), - "from_numpy": _make_from_numpy(mem_type), - "move_to_device": _make_device_mover(mem_type), - "from_dlpack": _make_dlpack_importer(mem_type), - } - class_name = f"{mem_type.value.capitalize()}Converter" - type(class_name, (ConverterBase,), class_attrs) - - -# Create all converter classes at module load time -_create_converter_classes() - - -def get_converter(memory_type: str): - """Get a converter instance for the given memory type. - - Args: - memory_type: The memory type string (e.g., "numpy", "torch") - - Returns: - A converter instance for the memory type - - Raises: - ValueError: If memory type is not registered - """ - converter_class = ConverterBase.__registry__.get(memory_type) - if converter_class is None: - raise ValueError( - f"No converter registered for memory type '{memory_type}'. " - f"Available types: {sorted(ConverterBase.__registry__.keys())}" - ) - return converter_class() - - -def _add_converter_methods(): - """Add to_X() methods to ConverterBase. - - For each target memory type, generates a method like to_cupy(), to_torch(), etc. - that tries GPU-to-GPU conversion via DLPack first, then falls back to CPU roundtrip. - """ - for target_type in MemoryType: - method_name = f"to_{target_type.value}" - - def make_method(tgt): - def method(self, data, gpu_id): - source_type = MemoryType(self.memory_type) - return source_type.convert_to(data, tgt, gpu_id) - - return method - - setattr(ConverterBase, method_name, make_method(target_type)) - - -def _validate_registry(): - """Validate that all memory types are registered.""" - required_types = {mt.value for mt in MemoryType} - registered_types = set(ConverterBase.__registry__.keys()) - - if required_types != registered_types: - missing = required_types - registered_types - extra = registered_types - required_types - msg_parts = [] - if missing: - msg_parts.append(f"Missing: {missing}") - if extra: - msg_parts.append(f"Extra: {extra}") - raise RuntimeError(f"Registry validation failed. {', '.join(msg_parts)}") - - -# Add to_X() conversion methods after converter classes are created -_add_converter_methods() - -# Run validation at module load time -_validate_registry() -ConverterBase.__registry__ = MappingProxyType(dict(ConverterBase.__registry__)) diff --git a/src/arraybridge/decorators.py b/src/arraybridge/decorators.py index 28d96ce..8aa1852 100644 --- a/src/arraybridge/decorators.py +++ b/src/arraybridge/decorators.py @@ -8,31 +8,33 @@ attributes and provide automatic thread-local CUDA stream management for GPU frameworks to enable true parallelization across multiple threads. -Framework-specific runtime capabilities are delegated to ``MemoryType`` members. +What a library needs around a call (streams, out-of-memory retries) comes from +its ``ArrayLibrary`` class. The decorator named after a library (``numpy``, +``torch``, ...) is made on first access for any registered library. """ import functools import inspect import logging -import threading from abc import ABC, abstractmethod from collections.abc import Callable from dataclasses import dataclass from enum import Enum -from typing import Any, ClassVar, TypeVar, cast +from typing import ClassVar, ParamSpec, TypeVar import numpy as np from metaclass_registry import AutoRegisterMeta, RegistryFamily, RegistryKeyAttribute from python_introspect import RuntimeParameterDeclarationABC from arraybridge.array_payload import ArrayPayload -from arraybridge.oom_recovery import _execute_with_oom_recovery +from arraybridge.libraries import Array, ArrayLibrary, DType from arraybridge.slice_processing import process_slices from arraybridge.types import MemoryContractAttribute, MemoryType logger = logging.getLogger(__name__) -F = TypeVar("F", bound=Callable[..., Any]) +P = ParamSpec("P") +R = TypeVar("R") class DtypeConversion(Enum): @@ -136,7 +138,7 @@ class EnumValueRegistryKeyMixin: strategy_label: ClassVar[str | None] = None - def __init_subclass__(cls, **kwargs: Any) -> None: + def __init_subclass__(cls, **kwargs: object) -> None: super().__init_subclass__(**kwargs) member = cls.registry_enum_member() if isinstance(member, Enum) and cls.__dict__.get("strategy_label") is None: @@ -152,10 +154,10 @@ def registry_enum_member(cls) -> Enum | None: class DtypeConversionRequest: """Runtime data needed to convert one decorated function output.""" - array: Any + array: Array original_dtype_name: str | None array_dtype_name: str | None - scale_func: Callable[[Any, Any], Any] + scale_func: Callable[[Array, DType], Array] class DtypeConversionRunner( @@ -178,10 +180,11 @@ def for_dtype_conversion( cls, dtype_conversion: DtypeConversion, ) -> "DtypeConversionRunner": - return cast(DtypeConversionRunner, cls.__registry__[dtype_conversion.value]()) + runner: DtypeConversionRunner = cls.__registry__[dtype_conversion.value]() + return runner @abstractmethod - def apply(self, request: DtypeConversionRequest) -> Any: + def apply(self, request: DtypeConversionRequest) -> Array: """Return output converted according to the configured dtype policy.""" @@ -190,7 +193,7 @@ class PreserveInputDtypeConversionRunner(DtypeConversionRunner): dtype_conversion = DtypeConversion.PRESERVE_INPUT - def apply(self, request: DtypeConversionRequest) -> Any: + def apply(self, request: DtypeConversionRequest) -> Array: if ( request.original_dtype_name is not None and request.array_dtype_name != request.original_dtype_name @@ -204,14 +207,14 @@ class NativeOutputDtypeConversionRunner(DtypeConversionRunner): dtype_conversion = DtypeConversion.NATIVE_OUTPUT - def apply(self, request: DtypeConversionRequest) -> Any: + def apply(self, request: DtypeConversionRequest) -> Array: return request.array class FixedDtypeConversionRunner(DtypeConversionRunner): """Scale output to the dtype declared by a fixed DtypeConversion member.""" - def apply(self, request: DtypeConversionRequest) -> Any: + def apply(self, request: DtypeConversionRequest) -> Array: if self.dtype_conversion is None: raise TypeError("FixedDtypeConversionRunner requires dtype_conversion.") target_dtype = self.dtype_conversion.numpy_dtype @@ -266,66 +269,11 @@ def _insertion_index(parameters: list[inspect.Parameter]) -> int: return len(parameters) -class ThreadGPUContextStorage(threading.local): - """Declare the existing native thread-local context slot and its type.""" - - context: "ThreadGPUContext | None" = None - - -class ThreadGPUContext: - """Runtime-owned thread-local streams keyed by framework/device identity. - - Keep the runtime handle on its importable owner, not in decorator function - globals. A retained, unpublished decorated callable is serialized by value; - its durable closure must not pull a ``threading.local`` into history. - """ - - _contexts: ClassVar[ThreadGPUContextStorage] = ThreadGPUContextStorage() - - @classmethod - def current(cls) -> "ThreadGPUContext": - """Return this thread's runtime context without serializing its handle.""" - context = cls._contexts.context - if context is None: - context = cls() - cls._contexts.context = context - return context - - def __init__(self): - self._streams: dict[tuple[MemoryType, int], Any] = {} - - def stream_for( - self, - memory_type: MemoryType, - module: Any, - ) -> tuple[int | None, Any | None]: - """Return the current device and its stable thread-local stream.""" - - device_id = memory_type.current_device_id(module) - if device_id is None: - return None, None - memory_type.require_device(device_id, module) - key = (memory_type, device_id) - if key not in self._streams: - with memory_type.device_scope(device_id, module): - stream = memory_type.create_stream(module) - if stream is None: - return device_id, None - self._streams[key] = stream - logger.debug( - "Created %s stream for device %d in thread %s", - memory_type.display_name, - device_id, - threading.current_thread().name, - ) - return device_id, self._streams[key] - - def memory_types( input_type: str | MemoryType, output_type: str | MemoryType, - contract: Any | None = None, -) -> Callable[[F], F]: + contract: object | None = None, +) -> Callable[[Callable[P, R]], Callable[P, R]]: """ Base decorator for declaring memory types of a function. @@ -337,9 +285,9 @@ def memory_types( input_memory_type = input_member.value output_memory_type = output_member.value - def decorator(func: F) -> F: + def decorator(func: Callable[P, R]) -> Callable[P, R]: @functools.wraps(func) - def wrapper(*args, **kwargs): + def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: result = func(*args, **kwargs) # Apply output validation only when a callable contract was provided. @@ -355,7 +303,7 @@ def wrapper(*args, **kwargs): if contract is not None and not callable(contract): setattr(wrapper, "__processing_contract__", contract) - return cast(F, wrapper) + return wrapper return decorator @@ -374,9 +322,9 @@ def wrap_dtype_preserving_callable( complete framework decorator or ArrayBridge internals. """ func_name = func.__name__ - input_memory_type = MemoryType(MemoryContractAttribute.INPUT.read(func, mem_type.value)) - output_memory_type = MemoryType(MemoryContractAttribute.OUTPUT.read(func, mem_type.value)) - scale_func = output_memory_type.scale_dtype + input_library = MemoryType(MemoryContractAttribute.INPUT.read(func, mem_type.value)).library + output_library = MemoryType(MemoryContractAttribute.OUTPUT.read(func, mem_type.value)).library + scale_func = output_library.scale_dtype default_dtype_config = ( DtypeConversionConfig.default_value() if dtype_config_default is None @@ -402,9 +350,7 @@ def dtype_wrapper(image, *args, **kwargs): # Store original dtype original_dtype = getattr(image, "dtype", None) original_dtype_name = ( - None - if original_dtype is None - else input_memory_type.canonical_dtype_name(original_dtype) + None if original_dtype is None else input_library.dtype_name(original_dtype) ) # Handle slice_by_slice processing for 3D arrays @@ -423,7 +369,7 @@ def _apply_dtype_conversion(array): DtypeConversionRequest( array=array, original_dtype_name=original_dtype_name, - array_dtype_name=output_memory_type.canonical_dtype_name(array.dtype), + array_dtype_name=output_library.dtype_name(array.dtype), scale_func=scale_func, ) ) @@ -468,59 +414,10 @@ def _apply_dtype_conversion(array): return dtype_wrapper -def _create_gpu_wrapper(func, mem_type: MemoryType, oom_recovery: bool): - """ - Auto-generate GPU stream/device wrapper for any GPU memory type. - - This function creates the GPU-specific wrapper with stream management and OOM recovery. - """ - - @functools.wraps(func) - def gpu_wrapper(*args, **kwargs): - framework = mem_type.import_if_installed() - - # Check if GPU is available for this framework - if framework is not None and mem_type.available_device_ids(framework): - # Get thread-local context - ctx = ThreadGPUContext.current() - - device_id, stream = ctx.stream_for(mem_type, framework) - - # Define execution function that captures args/kwargs - def execute_with_stream(): - with mem_type.stream_scope(stream, framework): - return func(*args, **kwargs) - - # Execute with OOM recovery if enabled - if oom_recovery: - return _execute_with_oom_recovery( - execute_with_stream, - mem_type.value, - device_id=device_id, - ) - return execute_with_stream() - - # CPU fallback or framework not available - return func(*args, **kwargs) - - # Preserve memory type attributes - MemoryContractAttribute.INPUT.write( - gpu_wrapper, - MemoryContractAttribute.INPUT.read(func), - ) - MemoryContractAttribute.OUTPUT.write( - gpu_wrapper, - MemoryContractAttribute.OUTPUT.read(func), - ) - - return gpu_wrapper - - def _create_memory_decorator(mem_type: MemoryType): """ Factory function that creates a decorator for a specific memory type. - This single factory replaces 6 nearly-identical decorator functions. """ def decorator( @@ -565,9 +462,8 @@ def inner_decorator(func): dtype_config_default=dtype_config_default, ) - # Apply GPU wrapper if this is a GPU memory type - if mem_type.is_gpu: - func = _create_gpu_wrapper(func, mem_type, oom_recovery) + # What the library's devices need around the call (streams, OOM retries). + func = mem_type.library.execution_wrapper(func, oom_recovery) MemoryContractAttribute.EXECUTION.write(func, mem_type.value) @@ -580,18 +476,25 @@ def inner_decorator(func): # Set proper function name and docstring decorator.__name__ = mem_type.value - decorator.__doc__ = (decorator.__doc__ or "").format(mem_type=mem_type.display_name) + decorator.__doc__ = (decorator.__doc__ or "").format(mem_type=mem_type.library.display_name) + + return decorator + + +def __getattr__(name: str): + """The decorator named after a registered library, made on first access.""" + if name.startswith("_") or name not in ArrayLibrary.__registry__: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + decorator = _create_memory_decorator(MemoryType(name)) + globals()[name] = decorator return decorator -# Auto-generate all 6 memory type decorators -for mem_type in MemoryType: - decorator_func = _create_memory_decorator(mem_type) - globals()[mem_type.value] = decorator_func +def __dir__() -> list[str]: + return sorted(set(globals()) | set(ArrayLibrary.__registry__)) -# Export the fixed decorator infrastructure plus declaration-derived helpers. __all__ = [ "memory_types", "DtypeConversion", @@ -600,4 +503,4 @@ def inner_decorator(func): "PRESERVE_INPUT_DTYPE_CONFIG", "SliceBySliceRuntimeParameter", "wrap_dtype_preserving_callable", -] + [memory_type.value for memory_type in MemoryType] +] diff --git a/src/arraybridge/devices.py b/src/arraybridge/devices.py new file mode 100644 index 0000000..e1ae6b0 --- /dev/null +++ b/src/arraybridge/devices.py @@ -0,0 +1,54 @@ +"""Where an array lives: the host processor or one accelerator of a library.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass + + +class Device(ABC): + """One place array data can live. Leaves are values: equal fields, equal device.""" + + @property + @abstractmethod + def label(self) -> str: + """Short text naming this device in records and messages.""" + + +@dataclass(frozen=True, slots=True) +class HostDevice(Device): + """Main memory, computed on by the host processor.""" + + @property + def label(self) -> str: + return "host" + + +@dataclass(frozen=True, slots=True) +class AcceleratorDevice(Device): + """One accelerator as its library numbers it (a CUDA or an OpenCL device).""" + + index: int + + @property + def label(self) -> str: + return f"accelerator:{self.index}" + + +@dataclass(frozen=True, slots=True) +class NamedDevice(Device): + """A device a library reports that is not one of its numbered accelerators. + + pyclesperanto running on a CPU OpenCL device is the case that exists: the + array belongs to that OpenCL device, which is neither NumPy's main memory + nor a listed GPU. It is identified by the name the library gives it. + """ + + name: str + + @property + def label(self) -> str: + return f"device:{self.name}" + + +HOST = HostDevice() diff --git a/src/arraybridge/dtype_scaling.py b/src/arraybridge/dtype_scaling.py deleted file mode 100644 index b3f3a76..0000000 --- a/src/arraybridge/dtype_scaling.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Declaration-owned dtype scaling and compatibility projections.""" - -from functools import partial -from types import MappingProxyType -from typing import Any - -from arraybridge.array_operations import _SCALING_RANGES as _SCALING_RANGES -from arraybridge.types import MemoryType - - -def _scale_generic(result: Any, target_dtype: Any, mem_type: MemoryType) -> Any: - """Scale through the operation leaf carried by ``mem_type``.""" - - return mem_type.scale_dtype(result, target_dtype) - - -def _scale_pyclesperanto(result: Any, target_dtype: Any) -> Any: - """Compatibility adapter for the pyclesperanto scaling leaf.""" - - return MemoryType.PYCLESPERANTO.scale_dtype(result, target_dtype) - - -SCALING_FUNCTIONS = MappingProxyType( - {memory_type.value: partial(_scale_generic, mem_type=memory_type) for memory_type in MemoryType} -) diff --git a/src/arraybridge/framework_config.py b/src/arraybridge/framework_config.py deleted file mode 100644 index 0eb4f19..0000000 --- a/src/arraybridge/framework_config.py +++ /dev/null @@ -1,26 +0,0 @@ -"""Deprecated compatibility view of declaration-owned array operations. - -Runtime code consumes :class:`arraybridge.types.MemoryType` directly. This -read-only mapping exists only for callers that still import the historical -private name. -""" - -from types import MappingProxyType - -from arraybridge.types import MemoryType - - -def _operation_view(memory_type: MemoryType): - return MappingProxyType( - { - "to_numpy": memory_type.to_numpy, - "from_numpy": memory_type.from_numpy, - "stack_arrays": memory_type.stack_arrays, - "scale_dtype": memory_type.scale_dtype, - } - ) - - -_FRAMEWORK_CONFIG = MappingProxyType( - {memory_type: _operation_view(memory_type) for memory_type in MemoryType} -) diff --git a/src/arraybridge/framework_ops.py b/src/arraybridge/framework_ops.py deleted file mode 100644 index 668fc6f..0000000 --- a/src/arraybridge/framework_ops.py +++ /dev/null @@ -1,7 +0,0 @@ -""" -Deprecated alias for the read-only, declaration-derived operation view. -""" - -from arraybridge.framework_config import _FRAMEWORK_CONFIG - -_FRAMEWORK_OPS = _FRAMEWORK_CONFIG diff --git a/src/arraybridge/gpu_cleanup.py b/src/arraybridge/gpu_cleanup.py deleted file mode 100644 index 136e618..0000000 --- a/src/arraybridge/gpu_cleanup.py +++ /dev/null @@ -1,78 +0,0 @@ -""" -GPU memory cleanup utilities for different frameworks. - -This module provides unified GPU memory cleanup functions for PyTorch, CuPy, -TensorFlow, JAX, and pyclesperanto. The cleanup functions are designed to be called -after processing steps to free up GPU memory that's no longer needed. - -Framework-specific cleanup behavior belongs to each ``MemoryType`` declaration. -""" - -import logging -from types import MappingProxyType - -from arraybridge.types import MemoryType - -logger = logging.getLogger(__name__) - - -def _cleanup_declared(mem_type: MemoryType, device_id: int | None = None) -> None: - try: - mem_type.cleanup_loaded(device_id) - except Exception as error: - logger.warning( - "Failed to cleanup %s GPU memory: %s", - mem_type.display_name, - error, - ) - - -def _create_cleanup_function(mem_type: MemoryType): - """Create one compatibility cleanup function from its enum declaration.""" - - def cleanup(device_id: int | None = None) -> None: - """Clean an already-loaded framework without importing it.""" - _cleanup_declared(mem_type, device_id) - - cleanup.__name__ = f"cleanup_{mem_type.import_name}_gpu" - cleanup.__doc__ = f"Clean already-loaded {mem_type.display_name} GPU resources." - - return cleanup - - -# Auto-generate all cleanup functions -for mem_type in MemoryType: - cleanup_func = _create_cleanup_function(mem_type) - globals()[cleanup_func.__name__] = cleanup_func - - -# Auto-generate cleanup registry -MEMORY_TYPE_CLEANUP_REGISTRY = MappingProxyType( - {mem_type.value: globals()[f"cleanup_{mem_type.import_name}_gpu"] for mem_type in MemoryType} -) - - -def cleanup_all_gpu_frameworks(device_id: int | None = None) -> None: - """ - Clean up GPU memory for all available frameworks. - - This function calls cleanup for all GPU frameworks that are currently loaded. - It's safe to call even if some frameworks aren't available. - - Args: - device_id: Optional GPU device ID. If None, cleans all devices. - """ - logger.debug(f"🔥 GPU CLEANUP: Starting cleanup for all GPU frameworks (device_id={device_id})") - - for mem_type in MemoryType: - if mem_type.is_gpu: - _cleanup_declared(mem_type, device_id) - - logger.debug("🔥 GPU CLEANUP: Completed cleanup for all GPU frameworks") - - -# Export all cleanup functions and utilities -__all__ = [ - "cleanup_all_gpu_frameworks", - "MEMORY_TYPE_CLEANUP_REGISTRY", -] + [f"cleanup_{mem_type.import_name}_gpu" for mem_type in MemoryType] diff --git a/src/arraybridge/libraries.py b/src/arraybridge/libraries.py new file mode 100644 index 0000000..2f16fe2 --- /dev/null +++ b/src/arraybridge/libraries.py @@ -0,0 +1,1696 @@ +"""One class per supported array library. + +``ArrayLibrary`` is the family. A leaf states the library's identity and owns +everything that library does differently: devices, moving data, DLPack, +out-of-memory recognition, child-process environment and native array +operations. What several libraries share is a parent they inherit from +(``HostDevices``, ``AcceleratorDevices``, ``CudaLibrary``, ``DLPackImport``, +``DLPackProtocolExport``, ``NvidiaWheelLibraries``, ``MutableArrays``). + +Adding a library is writing one class here or anywhere else: registration, +lookup, the ``MemoryType`` member, the decorator and conversion follow from it. + +Typing. ``Array`` is the nominal type of "an array of any supported library"; +each library's ``array_class`` is "an array of that library". Third-party +array classes are admitted with ``ABC.register`` when the library first meets +one, so no optional library is imported to define them. Inside a leaf the +native class is the type parameter, named through a ``TYPE_CHECKING`` import +because the optional package cannot be imported at runtime here; a checker +without that package installed sees it as unknown. A library's module is a +``types.ModuleType``: what it returns has no static type, so those values are +checked by ``Array`` registration at runtime, not by the checker. + +Importing this module imports no array library, NumPy included. +""" + +from __future__ import annotations + +import functools +import gc +import importlib +import importlib.util +import logging +import operator +import os +import sys +import threading +import time +from abc import ABCMeta, abstractmethod +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import AbstractContextManager, contextmanager, nullcontext +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType +from typing import TYPE_CHECKING, ClassVar, Generic, ParamSpec, TypeVar + +from metaclass_registry import AutoRegisterMeta + +from arraybridge.array_payload import ArrayPayload +from arraybridge.arrays import Array +from arraybridge.conversion import ConversionObserver, ConversionRecord, ConversionRoute, HostCopy +from arraybridge.devices import HOST, AcceleratorDevice, Device, HostDevice, NamedDevice + +if TYPE_CHECKING: + # Optional packages: importable by a type checker that has them, never at runtime here. + import cupy + import jax + import numpy as np + import pyclesperanto + import tensorflow + import torch + +logger = logging.getLogger(__name__) + +P = ParamSpec("P") +R = TypeVar("R") + +# A dtype as callers spell it: a NumPy dtype, a Python type, a name, or a library's own dtype. +DType = object +Axis = int | tuple[int, ...] + +_SCALING_RANGES: dict[str, float | tuple[float, float]] = { + "uint8": 255.0, + "uint16": 65535.0, + "uint32": 4294967295.0, + "int16": (65535.0, 32768.0), + "int32": (4294967295.0, 2147483648.0), +} + + +def _numpy() -> ModuleType: + """NumPy, imported when an operation first needs it.""" + + return importlib.import_module("numpy") + + +NativeT = TypeVar("NativeT") +TargetT = TypeVar("TargetT") + + +@dataclass(frozen=True, slots=True) +class DLPackPayload: + """One exported capsule plus the array it was exported from.""" + + source: object + capsule: object + + +def _dtype_name(dtype: DType) -> str: + declared_name = getattr(dtype, "name", None) + if declared_name is not None: + return str(declared_name) + return str(getattr(dtype, "__name__", str(dtype).rsplit(".", maxsplit=1)[-1])) + + +def _scaled_values(result, result_min, result_max, target_dtype: DType): + normalized = (result - result_min) / (result_max - result_min) + range_info = _SCALING_RANGES.get(_dtype_name(target_dtype)) + if range_info is None: + return normalized + if isinstance(range_info, tuple): + scale, offset = range_info + return normalized * scale - offset + return normalized * range_info + + +def _clamp_bounds(target_dtype: DType) -> tuple[float, float] | None: + range_info = _SCALING_RANGES.get(_dtype_name(target_dtype)) + if range_info is None: + return None + if isinstance(range_info, tuple): + scale, offset = range_info + return -offset, scale - offset - 128 + return 0, range_info + + +def _mapped_dtype(target_dtype: DType, module: ModuleType): + try: + dtype_name = _numpy().dtype(target_dtype).name + except TypeError as error: + raise TypeError(f"Unsupported target dtype {target_dtype!r}") from error + mapped = getattr(module, dtype_name, None) + if mapped is None: + raise TypeError(f"{module.__name__} does not expose dtype {dtype_name}") + return mapped + + +def _has_dlpack_protocol(data: object) -> bool: + return callable(getattr(data, "__dlpack__", None)) and callable( + getattr(data, "__dlpack_device__", None) + ) + + +class ArrayLibrary(Generic[NativeT], metaclass=AutoRegisterMeta): + """One array library: its identity and everything it does its own way.""" + + __registry_key__ = "name" + __registry__: ClassVar[Mapping[str, type[ArrayLibrary[Array]]]] = {} + _instances: ClassVar[dict[str, ArrayLibrary[Array]]] = {} + + # ----- identity, stated by each leaf ----- + name: ClassVar[str] + display_name: ClassVar[str] + module_aliases: ClassVar[tuple[str, ...]] = () + import_environment: ClassVar[tuple[tuple[str, str], ...]] = () + oom_messages: ClassVar[tuple[str, ...]] = () + oom_exception_paths: ClassVar[tuple[tuple[str, ...], ...]] = () + array_class: ClassVar[ABCMeta] + + def __init_subclass__(cls, **kwargs: object) -> None: + super().__init_subclass__(**kwargs) + declared = cls.__dict__.get("name") + if declared is None: + return + if "import_name" not in cls.__dict__: + cls.import_name = declared + if "display_name" not in cls.__dict__: + cls.display_name = declared + cls.array_class = ABCMeta(f"{cls.__name__}.Array", (Array,), {}) + + import_name: ClassVar[str] + + def __repr__(self) -> str: + return f"<{type(self).__name__}>" + + def __str__(self) -> str: + return str(self.name) + + def __reduce__(self) -> tuple[Callable[[str], ArrayLibrary[Array]], tuple[str]]: + return ArrayLibrary.named, (str(self.name),) + + # ----- lookup ----- + + @classmethod + def all(cls) -> tuple[ArrayLibrary[Array], ...]: + """Every registered library, in declaration order.""" + + return tuple(cls.named(name) for name in ArrayLibrary.__registry__) + + @classmethod + def named(cls, name: str) -> ArrayLibrary[Array]: + """The library registered under ``name`` (``"numpy"``, ``"torch"``, ...).""" + + try: + return ArrayLibrary._instances[name] + except KeyError: + pass + try: + leaf = ArrayLibrary.__registry__[name] + except KeyError: + raise ValueError( + f"'{name}' is not a registered array library. " + f"Available: {sorted(ArrayLibrary.__registry__)}" + ) from None + return ArrayLibrary._instances.setdefault(name, leaf()) + + @classmethod + def for_module_name(cls, module_name: str) -> ArrayLibrary[Array]: + """The library owning a top-level module name or one of its aliases.""" + + top_level = module_name.split(".")[0] + for library in cls.all(): + if top_level in library.recognized_module_names: + return library + raise ValueError(f"No array library owns module '{module_name}'") + + @classmethod + def of(cls, value: object) -> ArrayLibrary[Array]: + """Which library made ``value``. Imports nothing.""" + + libraries = cls.all() + for library in libraries: + if isinstance(value, library.array_class): + return library + # First sight of this class: the library is the owner of the nearest + # ancestor's module, so a subclass defined elsewhere is still recognised. + for ancestor in type(value).__mro__: + top_level = ancestor.__module__.split(".")[0] + for library in libraries: + if top_level in library.recognized_module_names: + library.array_class.register(type(value)) + return library + raise ValueError( + f"Unknown memory type for {type(value)} (module: {type(value).__module__})" + ) + + # ----- the module ----- + + @property + def recognized_module_names(self) -> frozenset[str]: + """Top-level module names whose objects belong to this library.""" + + return frozenset((self.import_name, *self.module_aliases)) + + def is_installed(self) -> bool: + """Whether the package is present, without importing it.""" + + if self.loaded_module() is not None: + return True + try: + return importlib.util.find_spec(self.import_name) is not None + except (ImportError, ModuleNotFoundError, ValueError): + return False + + def prepare_import(self) -> None: + """Set this library's import-time environment defaults before it is imported.""" + + missing_names = tuple( + name for name, _value in self.import_environment if name not in os.environ + ) + if missing_names and self.loaded_module() is not None: + logger.warning( + "%s was loaded before its ArrayBridge import defaults were set; " + "the current process cannot guarantee those import-time settings: %s", + self.display_name, + ", ".join(missing_names), + ) + for name, value in self.import_environment: + os.environ.setdefault(name, value) + + def loaded_module(self) -> ModuleType | None: + """The library's module if something already imported it.""" + + return sys.modules.get(self.import_name) + + def import_module(self) -> ModuleType: + """Import the library's module.""" + + self.prepare_import() + return importlib.import_module(self.import_name) + + def import_if_installed(self) -> ModuleType | None: + """Import the library only when its package is present.""" + + loaded = self.loaded_module() + if loaded is not None: + self.prepare_import() + return loaded + if not self.is_installed(): + return None + return self.import_module() + + @property + def module(self) -> ModuleType: + """The library's module, imported on first use.""" + + return self.loaded_module() or self.import_module() + + @property + def array_module(self) -> ModuleType: + """The module holding the library's array functions.""" + + return self.module + + def extend_child_environment(self, environment: dict[str, str]) -> dict[str, str]: + """Add what a child process needs to import this library.""" + + for name, value in self.import_environment: + environment.setdefault(name, value) + return environment + + @classmethod + def child_environment(cls, environment: Mapping[str, str] | None = None) -> dict[str, str]: + """An environment in which a child process can import every library.""" + + prepared = dict(os.environ if environment is None else environment) + for library in cls.all(): + prepared = library.extend_child_environment(prepared) + return prepared + + # ----- devices ----- + + def host_devices(self) -> tuple[HostDevice, ...]: + """The host device, if this library keeps arrays in main memory.""" + + return () + + def accelerator_devices(self) -> tuple[AcceleratorDevice, ...]: + """The accelerator devices this library can use on this machine now.""" + + return () + + def available_devices(self) -> tuple[Device, ...]: + return (*self.host_devices(), *self.accelerator_devices()) + + @abstractmethod + def device_of(self, data: NativeT) -> Device: + """The device holding ``data``.""" + + def require_device(self, device: Device) -> None: + """Raise unless this library has ``device`` on this machine.""" + + available = self.available_devices() + if device not in available: + raise ValueError( + f"{self.display_name} device {device.label} is unavailable; " + f"available devices are {tuple(each.label for each in available)}" + ) + + def device_scope(self, device: Device) -> AbstractContextManager[None]: + """A block in which new arrays of this library are created on ``device``.""" + + self.require_device(device) + return self._scope(device) + + def activate_device(self, device: Device) -> None: + """Select ``device`` for the whole process where the library has such a setting.""" + + self.require_device(device) + self._activate(device) + + def move_to_device(self, data: NativeT, device: Device) -> NativeT: + """Return ``data`` on ``device``; the same object when it is already there.""" + + self.require_device(device) + if self.device_of(data) == device: + return data + with self._scope(device): + return self._move(data, device) + + def _array_scope(self, reference: NativeT) -> AbstractContextManager[None]: + """The scope of the accelerator holding ``reference``; host data needs none.""" + + device = self.device_of(reference) + if isinstance(device, AcceleratorDevice): + return self.device_scope(device) + return nullcontext() + + # Each device parent answers for its own kind of device and passes the rest on. + + def _scope(self, device: Device) -> AbstractContextManager[None]: + raise ValueError(f"{self.display_name} has no {device.label} device") + + def _activate(self, device: Device) -> None: + raise ValueError(f"{self.display_name} has no {device.label} device") + + def _move(self, data: NativeT, device: Device) -> NativeT: + raise ValueError(f"{self.display_name} has no {device.label} device") + + def _place(self, data: np.ndarray, device: Device) -> NativeT: + raise ValueError(f"{self.display_name} has no {device.label} device") + + # ----- running a function of this library ----- + + def execution_wrapper(self, func: Callable[P, R], oom_recovery: bool) -> Callable[P, R]: + """Wrap a function written for this library with what its devices need.""" + + del oom_recovery + return func + + def is_oom_error(self, error: BaseException) -> bool: + """Whether ``error`` is this library running out of memory. Imports nothing.""" + + module = self.loaded_module() + if module is not None: + for path in self.oom_exception_paths: + candidate: object = module + for attribute in path: + candidate = getattr(candidate, attribute, None) + if ( + isinstance(candidate, type) + and issubclass(candidate, BaseException) + and isinstance(error, candidate) + ): + return True + message = str(error).lower() + return any(pattern in message for pattern in self.oom_messages) + + def cleanup_loaded(self, device: AcceleratorDevice | None = None) -> None: + """Release cached device memory if the library is loaded; every device by default.""" + + del device + + @classmethod + def cleanup_all_loaded(cls, device: AcceleratorDevice | None = None) -> None: + """Release cached device memory of every loaded library; failures are logged.""" + + for library in cls.all(): + try: + library.cleanup_loaded(device) + except Exception as error: + logger.warning( + "Failed to cleanup %s GPU memory: %s", library.display_name, error + ) + + def run_with_oom_recovery( + self, + func: Callable[[], R], + max_retries: int = 2, + device: AcceleratorDevice | None = None, + ) -> R: + """Call ``func``; on this library's out-of-memory error release caches and retry.""" + + attempt = 0 + while True: + try: + return func() + except Exception as error: + if not self.is_oom_error(error) or attempt == max_retries: + raise + attempt += 1 + try: + self.cleanup_loaded(device) + except Exception as error: + logger.warning("Failed to clear cache for %s: %s", self.name, error) + gc.collect() + + # ----- DLPack ----- + + @property + def imports_dlpack(self) -> bool: + """Whether this library can wrap another library's buffer through DLPack.""" + + return isinstance(self, DLPackImport) + + def imports_dlpack_from(self, device: Device) -> bool: + """Whether this library can wrap a buffer that lives on ``device``.""" + + del device + return False + + def exports_dlpack_from(self, device: Device) -> bool: + """Whether arrays of this library on ``device`` are offered through DLPack.""" + + del device + return False + + def exports_dlpack_data(self, data: NativeT) -> bool: + """Whether this particular array can be exported through DLPack.""" + + del data + return False + + def export_dlpack(self, data: NativeT) -> DLPackPayload | None: + """Export one array through DLPack, or ``None`` when this array cannot be.""" + + del data + return None + + # ----- conversion ----- + + def unrecorded_to_numpy(self, data: NativeT) -> np.ndarray: + """Download to NumPy without writing a conversion record (for routes).""" + + return self._to_host(data) + + def unrecorded_from_numpy(self, data: np.ndarray, device: Device) -> NativeT: + """Upload from NumPy without writing a conversion record (for routes).""" + + with self.device_scope(device): + return self._place(data, device) + + def to_numpy(self, data: NativeT) -> np.ndarray: + """Download ``data`` to a NumPy array in main memory.""" + + observer = ConversionObserver.current + if not observer.wants_records: + return self._to_host(data) + started = time.perf_counter() + source_device = self.device_of(data) + result = self._to_host(data) + if result is not data: + observer.record( + ConversionRecord( + source_library=str(self.name), + source_device=source_device.label, + target_library=str(NumpyLibrary.name), + target_device=HOST.label, + mechanism=str(HostCopy.mechanism), + nbytes=int(result.nbytes), + seconds=time.perf_counter() - started, + ) + ) + return result + + def from_numpy(self, data: np.ndarray, device: Device) -> NativeT: + """Create an array of this library on ``device`` from a NumPy array.""" + + observer = ConversionObserver.current + if not observer.wants_records: + return self.unrecorded_from_numpy(data, device) + started = time.perf_counter() + result = self.unrecorded_from_numpy(data, device) + if result is not data: + observer.record( + ConversionRecord( + source_library=str(NumpyLibrary.name), + source_device=HOST.label, + target_library=str(self.name), + target_device=device.label, + mechanism=str(HostCopy.mechanism), + nbytes=int(data.nbytes), + seconds=time.perf_counter() - started, + ) + ) + return result + + def convert_to(self, data: NativeT, target: ArrayLibrary[TargetT], device: Device) -> TargetT: + """Convert ``data`` to ``target`` on ``device`` by the first route that carries it.""" + + return ConversionRoute.run(data, self, target, device) + + def nbytes_of(self, data: NativeT) -> int: + """Size of the array's elements in bytes.""" + + np = _numpy() + return int(np.prod(tuple(getattr(data, "shape")), dtype=np.int64)) * int( + np.dtype(getattr(data, "dtype")).itemsize + ) + + # ----- array operations, each on the device of its input ----- + + def accepts(self, value: object) -> bool: + """Whether ``value`` is a native array of this library.""" + + return isinstance(value, self.module.ndarray) + + def stack_arrays(self, arrays: Sequence[NativeT], device: Device) -> NativeT: + """Stack arrays along a new leading axis on ``device``.""" + + with self.device_scope(device): + return self._stack(arrays) + + def normalize_planes( + self, data: NativeT, dtype: DType, scales: Sequence[float | None] + ) -> NativeT: + """Cast every plane to ``dtype`` and divide it by its scale.""" + + if len(scales) != int(getattr(data, "shape")[0]): + raise ValueError("Normalization scales must match the leading plane axis.") + with self._array_scope(data): + return self._normalize_planes(data, dtype, scales) + + def scale_dtype(self, data: NativeT, target_dtype: DType) -> NativeT: + """Convert to ``target_dtype``, rescaling float intensities into an integer range.""" + + if self.import_if_installed() is None or not hasattr(data, "dtype"): + return data + with self._array_scope(data): + return self._scale_dtype(data, target_dtype) + + def dtype_name(self, dtype: DType) -> str: + """The dtype's name as NumPy spells it (``"uint16"``, ``"float32"``).""" + + return str(_numpy().dtype(dtype).name) + + def astype(self, data: NativeT, dtype: DType) -> NativeT: + with self._array_scope(data): + return self._cast(data, dtype) + + def reshape(self, data: NativeT, shape: Sequence[int]) -> NativeT: + with self._array_scope(data): + return self._reshape(data, tuple(shape)) + + def broadcast_to(self, data: NativeT, shape: Sequence[int]) -> NativeT: + with self._array_scope(data): + return self._broadcast_to(data, tuple(shape)) + + def ones_like( + self, + reference: NativeT, + *, + shape: Sequence[int] | None = None, + dtype: DType = bool, + ) -> NativeT: + """Allocate ones on the device of ``reference``.""" + + with self._array_scope(reference): + return self._ones_like( + reference, + tuple(getattr(reference, "shape") if shape is None else shape), + dtype, + ) + + def logical_and(self, left: NativeT, right: NativeT) -> NativeT: + with self._array_scope(left): + return self._logical_and(left, right) + + def max(self, data: NativeT, axis: Axis) -> NativeT: + return self.array_module.max(data, axis=axis) + + def min(self, data: NativeT, axis: Axis) -> NativeT: + return self.array_module.min(data, axis=axis) + + def mean(self, data: NativeT, axis: Axis) -> NativeT: + return self.array_module.mean(data, axis=axis) + + def sum(self, data: NativeT, axis: Axis) -> NativeT: + return self.array_module.sum(data, axis=axis) + + def prepend_axis(self, data: NativeT) -> NativeT: + """Add a leading axis of length one.""" + + return self._reshape(data, (1, *getattr(data, "shape"))) + + def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> NativeT: + return self.array_module.linspace(start, stop, num, endpoint=endpoint) + + def ones(self, shape: int | tuple[int, ...], dtype: DType) -> NativeT: + """Allocate ones on the library's default device.""" + + return self.array_module.ones(shape, dtype=_mapped_dtype(dtype, self.array_module)) + + def assign(self, data: NativeT, index: slice, values: NativeT) -> NativeT: + """Return ``data`` with ``data[index]`` replaced by ``values``.""" + + getattr(data, "__setitem__")(index, values) + return data + + def outer(self, left: NativeT, right: NativeT) -> NativeT: + return self.array_module.outer(left, right) + + def floor(self, value: float) -> NativeT: + """Floor a host scalar in the library's default precision.""" + + return self.array_module.floor(value) + + # ----- what a leaf supplies; these run without choosing a device ----- + + @abstractmethod + def _to_host(self, data: NativeT) -> np.ndarray: + """Download to NumPy.""" + + @abstractmethod + def _scale_dtype(self, result: NativeT, target_dtype: DType) -> NativeT: + """The library's intensity-preserving dtype conversion.""" + + def _stack(self, values: Sequence[NativeT]) -> NativeT: + return self.module.stack(tuple(values), axis=0) + + def _cast(self, data: NativeT, dtype: DType) -> NativeT: + return getattr(data, "astype")(dtype, copy=False) + + def _logical_and(self, left: NativeT, right: NativeT) -> NativeT: + return self.module.logical_and(left, right) + + def _reshape(self, data: NativeT, shape: tuple[int, ...]) -> NativeT: + return getattr(data, "reshape")(shape) + + def _broadcast_to(self, data: NativeT, shape: tuple[int, ...]) -> NativeT: + return self.module.broadcast_to(data, shape) + + def _ones_like(self, reference: NativeT, shape: tuple[int, ...], dtype: DType) -> NativeT: + del reference + return self.module.ones(shape, dtype=dtype) + + def _normalize_planes( + self, data: NativeT, dtype: DType, scales: Sequence[float | None] + ) -> NativeT: + return self._stack( + tuple( + ( + self._cast(plane, dtype) + if scale is None + else operator.truediv(self._cast(plane, dtype), float(scale)) + ) + for plane, scale in ( + (getattr(data, "__getitem__")(index), scale) + for index, scale in enumerate(scales) + ) + ) + ) + + +# --------------------------------------------------------------------------- +# Roles several libraries share. A library inherits each role it plays. +# --------------------------------------------------------------------------- + + +class HostDevices(ArrayLibrary[NativeT]): + """Libraries that keep arrays in main memory (NumPy, torch, TensorFlow, JAX).""" + + def host_devices(self) -> tuple[HostDevice, ...]: + return (HOST,) + + def _scope(self, device: Device) -> AbstractContextManager[None]: + if isinstance(device, HostDevice): + return self._host_scope() + return super()._scope(device) + + def _activate(self, device: Device) -> None: + if not isinstance(device, HostDevice): + super()._activate(device) + + def _move(self, data: NativeT, device: Device) -> NativeT: + if isinstance(device, HostDevice): + return self._move_to_host(data) + return super()._move(data, device) + + def _place(self, data: np.ndarray, device: Device) -> NativeT: + if isinstance(device, HostDevice): + return self._wrap_host_array(data) + return super()._place(data, device) + + def _host_scope(self) -> AbstractContextManager[None]: + return nullcontext() + + @abstractmethod + def _wrap_host_array(self, data: np.ndarray) -> NativeT: + """A native array in main memory holding the NumPy array's values.""" + + @abstractmethod + def _move_to_host(self, data: NativeT) -> NativeT: + """The same native array in main memory.""" + + +class AcceleratorDevices(ArrayLibrary[NativeT]): + """Libraries that can keep arrays on an accelerator (every one but NumPy).""" + + def accelerator_devices(self) -> tuple[AcceleratorDevice, ...]: + if self.import_if_installed() is None: + return () + return tuple(AcceleratorDevice(index) for index in range(self._accelerator_count())) + + def _scope(self, device: Device) -> AbstractContextManager[None]: + if isinstance(device, AcceleratorDevice): + return self._accelerator_scope(device) + return super()._scope(device) + + def _activate(self, device: Device) -> None: + if isinstance(device, AcceleratorDevice): + self._activate_accelerator(device) + else: + super()._activate(device) + + def _move(self, data: NativeT, device: Device) -> NativeT: + if isinstance(device, AcceleratorDevice): + return self._move_to_accelerator(data, device) + return super()._move(data, device) + + def _place(self, data: np.ndarray, device: Device) -> NativeT: + if isinstance(device, AcceleratorDevice): + return self._upload(data, device) + return super()._place(data, device) + + @abstractmethod + def _accelerator_count(self) -> int: + """How many accelerators the imported library sees.""" + + @abstractmethod + def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[None]: + """A block in which new arrays are created on ``device``.""" + + def _activate_accelerator(self, device: AcceleratorDevice) -> None: + """Select ``device`` for the process; libraries without such a setting do nothing.""" + + del device + + @abstractmethod + def _move_to_accelerator(self, data: NativeT, device: AcceleratorDevice) -> NativeT: + """Copy ``data`` to ``device``; called inside that device's scope.""" + + @abstractmethod + def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> NativeT: + """Create a native array on ``device``; called inside that device's scope.""" + + def _thread_execution( + self, + ) -> tuple[AcceleratorDevice | None, AbstractContextManager[None]]: + """The device a decorated function runs on, if one is selected, and its scope.""" + + return None, nullcontext() + + def execution_wrapper(self, func: Callable[P, R], oom_recovery: bool) -> Callable[P, R]: + @functools.wraps(func) + def on_accelerator(*args: P.args, **kwargs: P.kwargs) -> R: + if self.import_if_installed() is None or not self.accelerator_devices(): + return func(*args, **kwargs) + device, scope = self._thread_execution() + + def execute() -> R: + with scope: + return func(*args, **kwargs) + + if oom_recovery: + return self.run_with_oom_recovery(execute, device=device) + return execute() + + return on_accelerator + + +class CudaLibrary(AcceleratorDevices[NativeT]): + """CUDA libraries with their own streams and a memory cache to release (CuPy, torch).""" + + _thread_streams: ClassVar[threading.local] = threading.local() + + def create_stream(self) -> object: + """A new CUDA stream on the current device.""" + + return self.module.cuda.Stream() + + @abstractmethod + def stream_scope(self, stream: object) -> AbstractContextManager[None]: + """A block whose work is queued on ``stream``.""" + + @abstractmethod + def current_device(self) -> AcceleratorDevice | None: + """The accelerator the process currently has selected, if any.""" + + @abstractmethod + def _free_cached_memory(self) -> None: + """Release the library's cached memory on the device in scope.""" + + def thread_stream(self, device: AcceleratorDevice) -> object: + """This thread's stream for ``device``, created on first use.""" + + streams: dict[tuple[str, int], object] | None = getattr( + CudaLibrary._thread_streams, "streams", None + ) + if streams is None: + streams = {} + CudaLibrary._thread_streams.streams = streams + key = (str(self.name), device.index) + if key not in streams: + with self.device_scope(device): + streams[key] = self.create_stream() + logger.debug( + "Created %s stream for device %d in thread %s", + self.display_name, + device.index, + threading.current_thread().name, + ) + return streams[key] + + def _thread_execution( + self, + ) -> tuple[AcceleratorDevice | None, AbstractContextManager[None]]: + device = self.current_device() + if device is None: + return None, nullcontext() + self.require_device(device) + return device, self.stream_scope(self.thread_stream(device)) + + def cleanup_loaded(self, device: AcceleratorDevice | None = None) -> None: + if self.loaded_module() is None: + return + available = self.accelerator_devices() + if not available: + return + for target in available if device is None else (device,): + with self.device_scope(target): + self._free_cached_memory() + + +class DLPackImport(ArrayLibrary[NativeT]): + """Libraries that can wrap another library's buffer handed over through DLPack.""" + + def imports_dlpack_from(self, device: Device) -> bool: + # A library without main-memory arrays cannot wrap a main-memory buffer. + return isinstance(self, HostDevices) or not isinstance(device, HostDevice) + + def from_dlpack(self, data: DLPackPayload | object) -> NativeT: + """Wrap an exported buffer; ``NotImplemented`` when this capsule cannot be read.""" + + payload = data if isinstance(data, DLPackPayload) else DLPackPayload(data, data) + return self._import_dlpack(payload) + + @abstractmethod + def _import_dlpack(self, payload: DLPackPayload) -> NativeT: + """Read one payload with the library's importer.""" + + +class DLPackProtocolExport(ArrayLibrary[NativeT]): + """Libraries whose arrays carry the standard DLPack export methods on every device.""" + + _export_methods: ClassVar[tuple[str, ...]] = ("__dlpack__", "to_dlpack", "toDlpack") + + def exports_dlpack_from(self, device: Device) -> bool: + del device + return True + + def exports_dlpack_data(self, data: NativeT) -> bool: + return any( + callable(getattr(data, attribute, None)) for attribute in self._export_methods + ) + + def export_dlpack(self, data: NativeT) -> DLPackPayload | None: + for attribute in self._export_methods: + exporter = getattr(data, attribute, None) + if callable(exporter): + return DLPackPayload(source=data, capsule=exporter()) + return None + + +class NvidiaWheelLibraries(ArrayLibrary[NativeT]): + """Libraries that load CUDA from the ``nvidia-*`` wheels' library directories.""" + + @staticmethod + def nvidia_wheel_library_paths() -> tuple[str, ...]: + """Native-library directories of the installed NVIDIA wheels.""" + + try: + spec = importlib.util.find_spec("nvidia") + except (ImportError, ModuleNotFoundError, ValueError): + return () + if spec is None or spec.submodule_search_locations is None: + return () + paths: set[str] = set() + for package_root in map(Path, spec.submodule_search_locations): + try: + components = tuple(package_root.iterdir()) + except OSError: + continue + for component in components: + for library_directory_name in ("lib", "bin"): + candidate = component / library_directory_name + if candidate.is_dir(): + paths.add(str(candidate)) + return tuple(sorted(paths)) + + def extend_child_environment(self, environment: dict[str, str]) -> dict[str, str]: + prepared = super().extend_child_environment(environment) + library_paths = self.nvidia_wheel_library_paths() + if not library_paths: + return prepared + search_variable = "PATH" if os.name == "nt" else "LD_LIBRARY_PATH" + existing_paths = tuple( + path for path in prepared.get(search_variable, "").split(os.pathsep) if path + ) + prepared[search_variable] = os.pathsep.join( + dict.fromkeys((*library_paths, *existing_paths)) + ) + return prepared + + +class MutableArrays(ArrayLibrary[NativeT]): + """Libraries with the NumPy array API whose arrays can be changed in place.""" + + def _normalize_planes( + self, data: NativeT, dtype: DType, scales: Sequence[float | None] + ) -> NativeT: + np = _numpy() + # Integer division promotes the output dtype, so only floats divide in place. + if not np.issubdtype(np.dtype(dtype), np.inexact) or len(scales) == 0: + return super()._normalize_planes(data, dtype, scales) + normalized = self.module.array(data, dtype=dtype, copy=True) + for index, scale in enumerate(scales): + if scale is not None: + normalized[index] /= float(scale) + return normalized + + def _scale_dtype(self, result: NativeT, target_dtype: DType) -> NativeT: + module = self.module + values = getattr(result, "astype") + if not ( + module.issubdtype(getattr(result, "dtype"), module.floating) + and self._rescales_into(target_dtype) + ): + return values(target_dtype) + result_min = module.min(result) + result_max = module.max(result) + if result_max <= result_min: + return values(target_dtype) + scaled = _scaled_values(result, result_min, result_max, target_dtype) + bounds = _clamp_bounds(target_dtype) + if bounds is not None: + scaled = module.clip(scaled, *bounds) + return scaled.astype(target_dtype) + + @abstractmethod + def _rescales_into(self, target_dtype: DType) -> bool: + """Whether a float array converted to ``target_dtype`` is rescaled first.""" + + +# --------------------------------------------------------------------------- +# The libraries. +# --------------------------------------------------------------------------- + + +class NumpyLibrary(HostDevices["np.ndarray"], MutableArrays["np.ndarray"]): + name = "numpy" + display_name = "NumPy" + oom_messages = ("cannot allocate memory", "memory exhausted") + + def device_of(self, data: np.ndarray) -> Device: + del data + return HOST + + def accepts(self, value: object) -> bool: + return isinstance(value, (_numpy().ndarray, ArrayPayload)) + + def _to_host(self, data: np.ndarray) -> np.ndarray: + return data + + def _wrap_host_array(self, data: np.ndarray) -> np.ndarray: + return data + + def _move_to_host(self, data: np.ndarray) -> np.ndarray: + return data + + def _rescales_into(self, target_dtype: DType) -> bool: + return bool(self.module.issubdtype(target_dtype, self.module.integer)) + + +class CupyLibrary( + CudaLibrary["cupy.ndarray"], + DLPackImport["cupy.ndarray"], + DLPackProtocolExport["cupy.ndarray"], + NvidiaWheelLibraries["cupy.ndarray"], + MutableArrays["cupy.ndarray"], +): + name = "cupy" + display_name = "CuPy" + oom_messages = ("out of memory", "cuda_error_out_of_memory") + oom_exception_paths = ( + ("cuda", "memory", "OutOfMemoryError"), + ("cuda", "runtime", "CUDARuntimeError"), + ) + + def device_of(self, data: cupy.ndarray) -> Device: + return AcceleratorDevice(int(data.device.id)) + + def _accelerator_count(self) -> int: + return int(self.module.cuda.runtime.getDeviceCount()) + + def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[None]: + return self.module.cuda.Device(device.index) + + def _activate_accelerator(self, device: AcceleratorDevice) -> None: + self.module.cuda.Device(device.index).use() + + def _move_to_accelerator(self, data: cupy.ndarray, device: AcceleratorDevice) -> cupy.ndarray: + del device + return data.copy() + + def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> cupy.ndarray: + del device + return self.module.array(data) + + def current_device(self) -> AcceleratorDevice | None: + return AcceleratorDevice(int(self.module.cuda.runtime.getDevice())) + + def stream_scope(self, stream: object) -> AbstractContextManager[None]: + if not isinstance(stream, AbstractContextManager): + raise TypeError(f"{stream!r} is not a CuPy stream") + # A CuPy stream is its own context manager. + return stream + + def _free_cached_memory(self) -> None: + module = self.module + module.get_default_memory_pool().free_all_blocks() + module.get_default_pinned_memory_pool().free_all_blocks() + module.cuda.runtime.deviceSynchronize() + + def _import_dlpack(self, payload: DLPackPayload) -> cupy.ndarray: + if _has_dlpack_protocol(payload.source): + return self.module.from_dlpack(payload.source) + capsule_importer = getattr(self.module, "fromDlpack", None) + return NotImplemented if capsule_importer is None else capsule_importer(payload.capsule) + + def _to_host(self, data: cupy.ndarray) -> np.ndarray: + return data.get() + + def _rescales_into(self, target_dtype: DType) -> bool: + return not self.module.issubdtype(target_dtype, self.module.floating) + + +class TorchLibrary( + HostDevices["torch.Tensor"], + CudaLibrary["torch.Tensor"], + DLPackImport["torch.Tensor"], + DLPackProtocolExport["torch.Tensor"], +): + name = "torch" + display_name = "PyTorch" + oom_messages = ("out of memory", "cuda_error_out_of_memory") + oom_exception_paths = (("cuda", "OutOfMemoryError"),) + + def device_of(self, data: torch.Tensor) -> Device: + return AcceleratorDevice(int(data.device.index)) if data.is_cuda else HOST + + def _accelerator_count(self) -> int: + cuda = self.module.cuda + return int(cuda.device_count()) if cuda.is_available() else 0 + + def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[None]: + return self.module.cuda.device(device.index) + + def _activate_accelerator(self, device: AcceleratorDevice) -> None: + self.module.cuda.set_device(device.index) + + def _move_to_accelerator(self, data: torch.Tensor, device: AcceleratorDevice) -> torch.Tensor: + return data.to(f"cuda:{device.index}") + + def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> torch.Tensor: + return self._wrap_host_array(data).to(f"cuda:{device.index}") + + def _wrap_host_array(self, data: np.ndarray) -> torch.Tensor: + host_data = ( + _numpy().ascontiguousarray(data) + if any(stride < 0 for stride in getattr(data, "strides", ())) + else data + ) + return self.module.from_numpy(host_data) + + def _move_to_host(self, data: torch.Tensor) -> torch.Tensor: + return data.cpu() + + def current_device(self) -> AcceleratorDevice | None: + cuda = self.module.cuda + return AcceleratorDevice(int(cuda.current_device())) if cuda.is_available() else None + + def stream_scope(self, stream: object) -> AbstractContextManager[None]: + return self.module.cuda.stream(stream) + + def _free_cached_memory(self) -> None: + self.module.cuda.empty_cache() + self.module.cuda.synchronize() + + def _import_dlpack(self, payload: DLPackPayload) -> torch.Tensor: + return self.module.from_dlpack(payload.capsule) + + def nbytes_of(self, data: torch.Tensor) -> int: + return int(data.element_size()) * int(data.nelement()) + + def accepts(self, value: object) -> bool: + return isinstance(value, self.module.Tensor) + + def dtype_name(self, dtype: DType) -> str: + return str(dtype).rsplit(".", maxsplit=1)[-1] + + def max(self, data: torch.Tensor, axis: Axis) -> torch.Tensor: + # CUDA has no uint16 reduction; float32 holds every uint16 exactly. + values = data.float() if data.dtype == self.module.uint16 else data + return self.module.amax(values, dim=axis).to(data.dtype) + + def min(self, data: torch.Tensor, axis: Axis) -> torch.Tensor: + values = data.float() if data.dtype == self.module.uint16 else data + return self.module.amin(values, dim=axis).to(data.dtype) + + def mean(self, data: torch.Tensor, axis: Axis) -> torch.Tensor: + return self.module.mean(data.float(), dim=axis) + + def sum(self, data: torch.Tensor, axis: Axis) -> torch.Tensor: + return self.module.sum(data, dim=axis) + + def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> torch.Tensor: + float32 = self.module.float32 + if endpoint: + return self.module.linspace(start, stop, num, dtype=float32) + return self.module.linspace(start, stop, num + 1, dtype=float32)[:-1] + + def floor(self, value: float) -> torch.Tensor: + return self.module.floor(self.module.tensor(value)) + + def _to_host(self, data: torch.Tensor) -> np.ndarray: + return data.cpu().numpy() + + def _scale_dtype(self, result: torch.Tensor, target_dtype: DType) -> torch.Tensor: + module, np = self.module, _numpy() + mapped = _mapped_dtype(target_dtype, module) + floats = (module.float16, module.float32, module.float64) + if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)): + return result.to(mapped) + result_min = result.min() + result_max = result.max() + if result_max <= result_min: + return result.to(mapped) + scaled = _scaled_values(result, result_min, result_max, target_dtype) + bounds = _clamp_bounds(target_dtype) + if bounds is not None: + scaled = module.clamp(scaled, min=bounds[0], max=bounds[1]) + return scaled.to(mapped) + + def _stack(self, values: Sequence[torch.Tensor]) -> torch.Tensor: + return self.module.stack(tuple(values), dim=0) + + def _cast(self, data: torch.Tensor, dtype: DType) -> torch.Tensor: + module = self.module + native = dtype if isinstance(dtype, module.dtype) else _mapped_dtype(dtype, module) + return data.to(dtype=native) + + def _broadcast_to(self, data: torch.Tensor, shape: tuple[int, ...]) -> torch.Tensor: + return data.expand(shape) + + def _ones_like( + self, reference: torch.Tensor, shape: tuple[int, ...], dtype: DType + ) -> torch.Tensor: + module = self.module + return module.ones(shape, dtype=_mapped_dtype(dtype, module), device=reference.device) + + +class TensorflowLibrary( + HostDevices["tensorflow.Tensor"], + AcceleratorDevices["tensorflow.Tensor"], + DLPackImport["tensorflow.Tensor"], +): + name = "tensorflow" + display_name = "TensorFlow" + import_environment = (("TF_FORCE_GPU_ALLOW_GROWTH", "true"),) + oom_messages = ("out of memory", "resource_exhausted") + oom_exception_paths = (("errors", "ResourceExhaustedError"),) + + def device_of(self, data: tensorflow.Tensor) -> Device: + device = data.device.lower() + if "gpu" not in device: + return HOST + return AcceleratorDevice(int(device.rsplit(":", maxsplit=1)[-1])) + + def _accelerator_count(self) -> int: + return len(self.module.config.list_logical_devices("GPU")) + + def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[None]: + return self.module.device(f"/device:GPU:{device.index}") + + def _host_scope(self) -> AbstractContextManager[None]: + return self.module.device("/device:CPU:0") + + def _move_to_accelerator( + self, data: tensorflow.Tensor, device: AcceleratorDevice + ) -> tensorflow.Tensor: + del device + return self.module.identity(data) + + def _move_to_host(self, data: tensorflow.Tensor) -> tensorflow.Tensor: + return self.module.identity(data) + + def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> tensorflow.Tensor: + del device + return self.module.convert_to_tensor(data) + + def _wrap_host_array(self, data: np.ndarray) -> tensorflow.Tensor: + return self.module.convert_to_tensor(data) + + def exports_dlpack_from(self, device: Device) -> bool: + # TensorFlow hands over accelerator tensors only. + return isinstance(device, AcceleratorDevice) + + def exports_dlpack_data(self, data: tensorflow.Tensor) -> bool: + module = self.module + try: + major, minor = map(int, module.__version__.split(".")[:2]) + except (AttributeError, TypeError, ValueError): + return False + if (major, minor) < (2, 12): + return False + if "gpu" not in str(getattr(data, "device", "")).lower(): + return False + dlpack = getattr(getattr(module, "experimental", None), "dlpack", None) + return callable(getattr(dlpack, "to_dlpack", None)) + + def export_dlpack(self, data: tensorflow.Tensor) -> DLPackPayload | None: + if not self.exports_dlpack_data(data): + return None + return DLPackPayload( + source=data, capsule=self.module.experimental.dlpack.to_dlpack(data) + ) + + def _import_dlpack(self, payload: DLPackPayload) -> tensorflow.Tensor: + return self.module.experimental.dlpack.from_dlpack(payload.capsule) + + def nbytes_of(self, data: tensorflow.Tensor) -> int: + return int(data.shape.num_elements()) * int(data.dtype.size) + + def accepts(self, value: object) -> bool: + return isinstance(value, self.module.Tensor) + + def dtype_name(self, dtype: DType) -> str: + return str(_numpy().dtype(getattr(dtype, "as_numpy_dtype", dtype)).name) + + def max(self, data: tensorflow.Tensor, axis: Axis) -> tensorflow.Tensor: + return self.module.reduce_max(data, axis=axis) + + def min(self, data: tensorflow.Tensor, axis: Axis) -> tensorflow.Tensor: + return self.module.reduce_min(data, axis=axis) + + def mean(self, data: tensorflow.Tensor, axis: Axis) -> tensorflow.Tensor: + module = self.module + return module.reduce_mean(module.cast(data, module.float32), axis=axis) + + def sum(self, data: tensorflow.Tensor, axis: Axis) -> tensorflow.Tensor: + return self.module.reduce_sum(data, axis=axis) + + def linspace( + self, start: float, stop: float, num: int, *, endpoint: bool + ) -> tensorflow.Tensor: + if endpoint: + return self.module.linspace(float(start), float(stop), num) + return self.module.linspace(float(start), float(stop), num + 1)[:-1] + + def assign( + self, data: tensorflow.Tensor, index: slice, values: tensorflow.Tensor + ) -> tensorflow.Tensor: + module = self.module + rows = module.range(*index.indices(int(data.shape[0]))) + return module.tensor_scatter_nd_update(data, module.reshape(rows, [-1, 1]), values) + + def outer(self, left: tensorflow.Tensor, right: tensorflow.Tensor) -> tensorflow.Tensor: + return self.module.tensordot(left, right, axes=0) + + def floor(self, value: float) -> tensorflow.Tensor: + return self.module.math.floor(value) + + def _to_host(self, data: tensorflow.Tensor) -> np.ndarray: + return data.numpy() + + def _scale_dtype(self, result: tensorflow.Tensor, target_dtype: DType) -> tensorflow.Tensor: + module, np = self.module, _numpy() + mapped = _mapped_dtype(target_dtype, module) + floats = (module.float16, module.float32, module.float64) + if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)): + return module.cast(result, mapped) + result_min = module.reduce_min(result) + result_max = module.reduce_max(result) + if result_max <= result_min: + return module.cast(result, mapped) + scaled = _scaled_values(result, result_min, result_max, target_dtype) + bounds = _clamp_bounds(target_dtype) + if bounds is not None: + scaled = module.clip_by_value(scaled, *bounds) + return module.cast(scaled, mapped) + + def _cast(self, data: tensorflow.Tensor, dtype: DType) -> tensorflow.Tensor: + module = self.module + native = dtype if isinstance(dtype, module.dtypes.DType) else _mapped_dtype(dtype, module) + return module.cast(data, native) + + def _reshape(self, data: tensorflow.Tensor, shape: tuple[int, ...]) -> tensorflow.Tensor: + return self.module.reshape(data, shape) + + def _ones_like( + self, reference: tensorflow.Tensor, shape: tuple[int, ...], dtype: DType + ) -> tensorflow.Tensor: + module = self.module + with module.device(reference.device): + return module.ones(shape, dtype=_mapped_dtype(dtype, module)) + + +class JaxLibrary( + HostDevices["jax.Array"], + AcceleratorDevices["jax.Array"], + DLPackImport["jax.Array"], + DLPackProtocolExport["jax.Array"], +): + name = "jax" + display_name = "JAX" + module_aliases = ("jaxlib",) + import_environment = (("XLA_PYTHON_CLIENT_PREALLOCATE", "false"),) + oom_messages = ("out of memory", "oom when allocating", "allocation failure") + + @property + def array_module(self) -> ModuleType: + return self.module.numpy + + def _platform_devices(self, platform: str) -> tuple[object, ...]: + return tuple(device for device in self.module.devices() if device.platform == platform) + + def device_of(self, data: jax.Array) -> Device: + device = data.device + device = device() if callable(device) else device + if getattr(device, "platform", None) != "gpu": + return HOST + gpu_devices = self._platform_devices("gpu") + for index, candidate in enumerate(gpu_devices): + if candidate is device: + return AcceleratorDevice(index) + for index, candidate in enumerate(gpu_devices): + if str(candidate) == str(device): + return AcceleratorDevice(index) + return NamedDevice(str(device)) + + def _accelerator_count(self) -> int: + return len(self._platform_devices("gpu")) + + def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[None]: + return self.module.default_device(self._platform_devices("gpu")[device.index]) + + def _host_scope(self) -> AbstractContextManager[None]: + return self.module.default_device(self.module.devices("cpu")[0]) + + def _move_to_accelerator(self, data: jax.Array, device: AcceleratorDevice) -> jax.Array: + return self.module.device_put(data, self._platform_devices("gpu")[device.index]) + + def _move_to_host(self, data: jax.Array) -> jax.Array: + return self.module.device_put(data, self.module.devices("cpu")[0]) + + def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> jax.Array: + return self.module.device_put(data, self._platform_devices("gpu")[device.index]) + + def _wrap_host_array(self, data: np.ndarray) -> jax.Array: + return self.module.device_put(data, self.module.devices("cpu")[0]) + + def _import_dlpack(self, payload: DLPackPayload) -> jax.Array: + if not _has_dlpack_protocol(payload.source): + return NotImplemented + return self.module.dlpack.from_dlpack(self._aligned(payload.source)) + + @staticmethod + def _aligned(data: object) -> object: + """Copy a DLPack source only when JAX's 16-byte alignment is unmet.""" + + pointer: int | None = None + data_ptr = getattr(data, "data_ptr", None) + cuda_interface = getattr(data, "__cuda_array_interface__", None) + if callable(data_ptr): + pointer = int(data_ptr()) + elif isinstance(cuda_interface, dict): + declared, _read_only = cuda_interface.get("data", (None, False)) + pointer = None if declared is None else int(declared) + if pointer is None or pointer % 16 == 0: + return data + for attribute in ("clone", "copy"): + copier = getattr(data, attribute, None) + if callable(copier): + return copier() + return data + + def accepts(self, value: object) -> bool: + return isinstance(value, self.module.numpy.ndarray) + + def mean(self, data: jax.Array, axis: Axis) -> jax.Array: + jnp = self.module.numpy + return jnp.mean(data.astype(jnp.float32), axis=axis) + + def assign(self, data: jax.Array, index: slice, values: jax.Array) -> jax.Array: + return data.at[index].set(values) + + def _to_host(self, data: jax.Array) -> np.ndarray: + return _numpy().asarray(data) + + def _scale_dtype(self, result: jax.Array, target_dtype: DType) -> jax.Array: + module, np = self.module, _numpy() + if np.dtype(target_dtype) == np.dtype(np.float64): + x64_enabled = getattr(module.config, "x64_enabled", None) + if x64_enabled is None: + x64_enabled = module.config.read("jax_enable_x64") + if not x64_enabled: + raise ValueError( + "JAX float64 output requires x64 mode; set JAX_ENABLE_X64=true before import" + ) + jnp = module.numpy + mapped = _mapped_dtype(target_dtype, jnp) + floats = (jnp.float16, jnp.float32, jnp.float64) + if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)): + return result.astype(mapped) + result_min = jnp.min(result) + result_max = jnp.max(result) + if result_max <= result_min: + return result.astype(mapped) + scaled = _scaled_values(result, result_min, result_max, target_dtype) + bounds = _clamp_bounds(target_dtype) + if bounds is not None: + scaled = jnp.clip(scaled, *bounds) + return scaled.astype(mapped) + + def _stack(self, values: Sequence[jax.Array]) -> jax.Array: + return self.module.numpy.stack(tuple(values), axis=0) + + def _cast(self, data: jax.Array, dtype: DType) -> jax.Array: + return data.astype(_mapped_dtype(dtype, self.module.numpy)) + + def _logical_and(self, left: jax.Array, right: jax.Array) -> jax.Array: + return self.module.numpy.logical_and(left, right) + + def _broadcast_to(self, data: jax.Array, shape: tuple[int, ...]) -> jax.Array: + return self.module.numpy.broadcast_to(data, shape) + + def _ones_like(self, reference: jax.Array, shape: tuple[int, ...], dtype: DType) -> jax.Array: + return self.module.numpy.ones(shape, dtype=dtype, device=reference.device) + + +class PyclesperantoLibrary(AcceleratorDevices["pyclesperanto.Array"]): + name = "pyclesperanto" + oom_messages = ("cl_mem_object_allocation_failure", "cl_out_of_resources", "out of memory") + + _native_dtype_names: ClassVar[frozenset[str]] = frozenset( + {"float32", "int8", "int16", "int32", "uint8", "uint16", "uint32"} + ) + + def _gpu_devices(self) -> tuple[object, ...]: + try: + return tuple(self.module.list_available_devices("gpu")) + except TypeError: + return tuple(self.module.list_available_devices()) + + @staticmethod + def _device_name(device: object) -> str: + return str(getattr(device, "name", None) or device) + + def _select(self, selector: str | int, device_type: str | None = None) -> None: + if device_type is None: + self.module.select_device(selector) + return + try: + self.module.select_device(selector, device_type) + except TypeError: + self.module.select_device(selector) + + def _device_named(self, name: str) -> Device: + for index, device in enumerate(self._gpu_devices()): + if self._device_name(device) == name: + return AcceleratorDevice(index) + # An OpenCL device that is not a listed GPU, such as a CPU OpenCL device. + return NamedDevice(name) + + def device_of(self, data: pyclesperanto.Array) -> Device: + declared_device = getattr(data, "device", None) + if declared_device is None: + declared_device = self.module.get_device() + return self._device_named(self._device_name(declared_device)) + + def _accelerator_count(self) -> int: + return len(self._gpu_devices()) + + @contextmanager + def _accelerator_scope(self, device: AcceleratorDevice) -> Iterator[None]: + previous_name = self._device_name(self.module.get_device()) + previous = self._device_named(previous_name) + self._select(device.index, "gpu") + try: + yield + finally: + if previous != device: + self._select(previous_name) + + def _activate_accelerator(self, device: AcceleratorDevice) -> None: + self._select(device.index, "gpu") + + def _move_to_accelerator( + self, data: pyclesperanto.Array, device: AcceleratorDevice + ) -> pyclesperanto.Array: + del device + module = self.module + result = module.create_like(data) + module.copy(data, result) + return result + + def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> pyclesperanto.Array: + del device + return self.module.push(data) + + def accepts(self, value: object) -> bool: + # pyclesperanto kernels push any host array they receive. + return hasattr(value, "shape") + + @staticmethod + def _require_z_axis(axis: Axis) -> None: + if axis != 0: + raise NotImplementedError( + f"pyclesperanto reduces only along the leading (z) axis, not axis {axis!r}" + ) + + def max(self, data: pyclesperanto.Array, axis: Axis) -> pyclesperanto.Array: + self._require_z_axis(axis) + return self.module.maximum_z_projection(data) + + def min(self, data: pyclesperanto.Array, axis: Axis) -> pyclesperanto.Array: + self._require_z_axis(axis) + return self.module.minimum_z_projection(data) + + def mean(self, data: pyclesperanto.Array, axis: Axis) -> pyclesperanto.Array: + self._require_z_axis(axis) + return self.module.mean_z_projection(data) + + def sum(self, data: pyclesperanto.Array, axis: Axis) -> pyclesperanto.Array: + self._require_z_axis(axis) + return self.module.sum_z_projection(data) + + def prepend_axis(self, data: pyclesperanto.Array) -> pyclesperanto.Array: + result = self.module.create((1, *data.shape), dtype=data.dtype) + result[0] = data + return result + + def _to_host(self, data: pyclesperanto.Array) -> np.ndarray: + return self.module.pull(data) + + def _scale_dtype( + self, result: pyclesperanto.Array, target_dtype: DType + ) -> pyclesperanto.Array: + module, np = self.module, _numpy() + target_is_int = np.issubdtype(np.dtype(target_dtype), np.integer) + if not (np.issubdtype(result.dtype, np.floating) and target_is_int): + return module.push(module.pull(result).astype(target_dtype)) + result_min = float(module.minimum_of_all_pixels(result)) + result_max = float(module.maximum_of_all_pixels(result)) + if result_max <= result_min: + return module.push(module.pull(result).astype(target_dtype)) + normalized = module.subtract_image_from_scalar(result, scalar=result_min) + normalized = module.multiply_image_and_scalar( + normalized, + scalar=1.0 / (result_max - result_min), + ) + range_info = _SCALING_RANGES.get(_dtype_name(target_dtype)) + if isinstance(range_info, tuple): + scale, offset = range_info + scaled = module.multiply_image_and_scalar(normalized, scalar=scale) + scaled = module.subtract_image_from_scalar(scaled, scalar=offset) + elif range_info is not None: + scaled = module.multiply_image_and_scalar(normalized, scalar=range_info) + else: + scaled = normalized + host_values = module.pull(scaled) + bounds = _clamp_bounds(target_dtype) + if bounds is not None: + host_values = np.clip(host_values, *bounds) + return module.push(host_values.astype(target_dtype)) + + def _stack(self, values: Sequence[pyclesperanto.Array]) -> pyclesperanto.Array: + module = self.module + if not values: + raise ValueError("Cannot stack an empty pyclesperanto sequence") + if len(values) == 1: + source = values[0] + result = module.create((1, *source.shape), dtype=source.dtype) + return module.copy_slice(source, result, 0) + result = values[0] + for value in values[1:]: + result = module.concatenate_along_z(result, value) + return result + + def _cast(self, data: pyclesperanto.Array, dtype: DType) -> pyclesperanto.Array: + module = self.module + return module.push(module.pull(data).astype(dtype, copy=False)) + + def _logical_and( + self, left: pyclesperanto.Array, right: pyclesperanto.Array + ) -> pyclesperanto.Array: + module = self.module + return module.push(_numpy().logical_and(module.pull(left), module.pull(right))) + + def _reshape(self, data: pyclesperanto.Array, shape: tuple[int, ...]) -> pyclesperanto.Array: + if tuple(data.shape) == shape: + return data + raise NotImplementedError( + "pyclesperanto does not provide native array reshape; its reshape downloads pixels" + ) + + def _broadcast_to( + self, data: pyclesperanto.Array, shape: tuple[int, ...] + ) -> pyclesperanto.Array: + if tuple(data.shape) == shape: + return data + raise NotImplementedError("pyclesperanto does not provide native array broadcasting") + + def _ones_like( + self, reference: pyclesperanto.Array, shape: tuple[int, ...], dtype: DType + ) -> pyclesperanto.Array: + module = self.module + if _numpy().dtype(dtype).name not in self._native_dtype_names: + raise NotImplementedError(f"pyclesperanto cannot allocate native dtype {dtype!r}") + if not 1 <= len(shape) <= 3: + raise NotImplementedError( + "pyclesperanto native allocation requires one to three dimensions" + ) + result = module.create(shape, dtype=dtype, device=reference.device) + module.set(result, scalar=1, device=reference.device) + return result diff --git a/src/arraybridge/oom_recovery.py b/src/arraybridge/oom_recovery.py deleted file mode 100644 index 2b95889..0000000 --- a/src/arraybridge/oom_recovery.py +++ /dev/null @@ -1,90 +0,0 @@ -""" -GPU Out of Memory (OOM) recovery utilities. - -Provides comprehensive OOM detection and cache clearing for all supported -GPU frameworks in OpenHCS. - -OOM classification and cache cleanup delegate to each ``MemoryType`` declaration. -""" - -import gc -import logging - -from arraybridge.types import MemoryType - -logger = logging.getLogger(__name__) - - -def _is_oom_error(e: Exception, memory_type: str) -> bool: - """ - Detect Out of Memory errors for all GPU frameworks. - - Args: - e: Exception to check - memory_type: Memory type string (e.g., 'torch', 'cupy') - - Returns: - True if exception is an OOM error for the given framework - """ - try: - mem_type_enum = MemoryType(memory_type) - except ValueError: - return False - - return mem_type_enum.is_oom_error(e) - - -def _clear_cache_for_memory_type(memory_type: str, device_id: int | None = None): - """ - Clear GPU cache for specific memory type. - - Args: - memory_type: Memory type string (e.g., 'torch', 'cupy') - device_id: Optional framework-local GPU device ID. ``None`` cleans all. - """ - try: - mem_type_enum = MemoryType(memory_type) - except ValueError: - logger.warning(f"Unknown memory type for cache clearing: {memory_type}") - gc.collect() - return - - try: - mem_type_enum.cleanup_loaded(device_id) - except Exception as e: - logger.warning(f"Failed to clear cache for {memory_type}: {e}") - - # Always trigger Python garbage collection - gc.collect() - - -def _execute_with_oom_recovery( - func_callable, - memory_type: str, - max_retries: int = 2, - device_id: int | None = None, -): - """ - Execute function with automatic OOM recovery. - - Args: - func_callable: Function to execute - memory_type: Memory type from MemoryType enum - max_retries: Maximum number of retry attempts - device_id: Optional framework-local device whose cache should be cleared - - Returns: - Function result - - Raises: - Original exception if not OOM or retries exhausted - """ - for attempt in range(max_retries + 1): - try: - return func_callable() - except Exception as e: - if not _is_oom_error(e, memory_type) or attempt == max_retries: - raise - - # Clear cache and retry - _clear_cache_for_memory_type(memory_type, device_id) diff --git a/src/arraybridge/slice_processing.py b/src/arraybridge/slice_processing.py index b07df4e..6f5f596 100644 --- a/src/arraybridge/slice_processing.py +++ b/src/arraybridge/slice_processing.py @@ -5,12 +5,11 @@ that works for all memory types, eliminating duplication across dtype wrappers. """ -from arraybridge.converters import detect_memory_type +from arraybridge.libraries import ArrayLibrary from arraybridge.stack_utils import stack_slices, unstack_slices -from arraybridge.utils import _get_device_id -def process_slices(image, func, args, kwargs, gpu_id=None): +def process_slices(image, func, args, kwargs, device=None): """ Process a 3D array slice-by-slice using the provided function. @@ -26,21 +25,17 @@ def process_slices(image, func, args, kwargs, gpu_id=None): func: Function to apply to each slice args: Positional arguments to pass to func kwargs: Keyword arguments to pass to func - gpu_id: Optional GPU device ID override. If not provided, attempts - to derive from the input image and falls back to 0. + device: The device to unstack on. By default the device holding the image. Returns: Processed 3D array, or tuple of (processed_3d_array, special_outputs...) if func returns tuples """ - # Detect memory type and use proper OpenHCS utilities - memory_type = detect_memory_type(image) - if gpu_id is None: - detected_gpu_id = _get_device_id(image, memory_type) - gpu_id = 0 if detected_gpu_id is None else detected_gpu_id + library = ArrayLibrary.of(image) + if device is None: + device = library.device_of(image) - # Unstack 3D array into 2D slices - slices_2d = unstack_slices(image, memory_type, gpu_id) + slices_2d = unstack_slices(image, library.name, device) # Process each slice and handle special outputs main_outputs = [] @@ -82,12 +77,11 @@ def process_slices(image, func, args, kwargs, gpu_id=None): # framework for the per-slice function. if not main_outputs: raise ValueError("Slice processing produced no main outputs to stack") - output_memory_type = detect_memory_type(main_outputs[0]) - output_gpu_id = _get_device_id(main_outputs[0], output_memory_type) + output_library = ArrayLibrary.of(main_outputs[0]) result = stack_slices( main_outputs, - output_memory_type, - 0 if output_gpu_id is None else output_gpu_id, + output_library.name, + output_library.device_of(main_outputs[0]), ) # If we have special outputs, combine them and return tuple diff --git a/src/arraybridge/stack_utils.py b/src/arraybridge/stack_utils.py index 0967cdb..e810f96 100644 --- a/src/arraybridge/stack_utils.py +++ b/src/arraybridge/stack_utils.py @@ -10,10 +10,10 @@ """ import logging -from typing import Any -from arraybridge.converters import detect_memory_type -from arraybridge.types import MemoryType +from arraybridge.converters import convert_memory +from arraybridge.devices import Device +from arraybridge.libraries import Array, ArrayLibrary logger = logging.getLogger(__name__) @@ -21,7 +21,7 @@ logger.debug("🔄 STACK_UTILS: Module loaded - memory conversion logging enabled") -def _is_2d(data: Any) -> bool: +def _is_2d(data: object) -> bool: """ Check if data is a 2D array. @@ -39,7 +39,7 @@ def _is_2d(data: Any) -> bool: return len(data.shape) == 2 -def _is_3d(data: Any) -> bool: +def _is_3d(data: object) -> bool: """ Check if data is a 3D array. @@ -57,23 +57,7 @@ def _is_3d(data: Any) -> bool: return len(data.shape) == 3 -def _enforce_gpu_device_requirements(memory_type: str, gpu_id: int) -> None: - """ - Enforce GPU device requirements. - - Args: - memory_type: The memory type - gpu_id: The GPU device ID - - Raises: - ValueError: If gpu_id is negative - """ - mem_type = MemoryType(memory_type) - if mem_type.is_gpu: - mem_type.require_device(gpu_id) - - -def stack_slices(slices: list[Any], memory_type: str, gpu_id: int) -> Any: +def stack_slices(slices: list[Array], memory_type: str, device: Device) -> Array: """ Stack 2D slices into a 3D array with the specified memory type. @@ -83,14 +67,14 @@ def stack_slices(slices: list[Any], memory_type: str, gpu_id: int) -> Any: Args: slices: List of 2D slices (numpy arrays, cupy arrays, torch tensors, etc.) memory_type: The memory type to use for the stacked array (REQUIRED) - gpu_id: The target GPU device ID (REQUIRED) + device: The device of the target library to put the result on (REQUIRED) Returns: A 3D array with the specified memory type of shape [Z, Y, X] Raises: ValueError: If memory_type is not supported or slices is empty - ValueError: If gpu_id is negative for GPU memory types + ValueError: If the target library does not have the device ValueError: If slices are not 2D arrays MemoryConversionError: If conversion fails """ @@ -102,30 +86,22 @@ def stack_slices(slices: list[Any], memory_type: str, gpu_id: int) -> Any: if not _is_2d(slice_data): raise ValueError(f"Slice at index {i} is not a 2D array. All slices must be 2D.") - # Check GPU requirements - _enforce_gpu_device_requirements(memory_type, gpu_id) + library = ArrayLibrary.named(memory_type) + library.require_device(device) - # Convert each slice and enforce the requested framework-local device. conversion_count = 0 - mem_type = MemoryType(memory_type) converted_slices = [] for slice_data in slices: - source_type = detect_memory_type(slice_data) - if source_type == memory_type: - converted_data = mem_type.move_to_device(slice_data, gpu_id) - else: - from arraybridge.converters import convert_memory - - converted_data = convert_memory( - data=slice_data, - source_type=source_type, - target_type=memory_type, - gpu_id=gpu_id, - ) + source = ArrayLibrary.of(slice_data) + if source is not library: conversion_count += 1 - converted_slices.append(converted_data) + converted_slices.append( + convert_memory(slice_data, str(source.name), memory_type, device) + ) + else: + converted_slices.append(library.move_to_device(slice_data, device)) - result = mem_type.stack_arrays(converted_slices, gpu_id) + result = library.stack_arrays(converted_slices, device) # 🔍 MEMORY CONVERSION LOGGING: Only log when conversions happen or issues occur if conversion_count > 0: @@ -139,8 +115,8 @@ def stack_slices(slices: list[Any], memory_type: str, gpu_id: int) -> Any: def unstack_slices( - array: Any, memory_type: str, gpu_id: int, validate_slices: bool = True -) -> list[Any]: + array: Array, memory_type: str, device: Device, validate_slices: bool = True +) -> list[Array]: """ Split a 3D array into 2D slices along axis 0 and convert to the specified memory type. @@ -149,7 +125,7 @@ def unstack_slices( Args: array: 3D array to split - MUST BE 3D memory_type: The memory type to use for the output slices (REQUIRED) - gpu_id: The target GPU device ID (REQUIRED) + device: The device of the target library to put the result on (REQUIRED) validate_slices: If True, validates that each extracted slice is 2D Returns: @@ -158,38 +134,29 @@ def unstack_slices( Raises: ValueError: If array is not 3D ValueError: If validate_slices is True and any extracted slice is not 2D - ValueError: If gpu_id is negative for GPU memory types + ValueError: If the target library does not have the device ValueError: If memory_type is not supported MemoryConversionError: If conversion fails """ - # Detect input type and check if conversion is needed - input_type = detect_memory_type(array) - getattr(array, "shape", "unknown") + source = ArrayLibrary.of(array) + source_type = source.name # Verify the array is 3D - fail loudly if not if not _is_3d(array): raise ValueError(f"Array must be 3D, got shape {getattr(array, 'shape', 'unknown')}") - # Check GPU requirements - _enforce_gpu_device_requirements(memory_type, gpu_id) - - # Convert to target memory type - source_type = input_type # Reuse already detected type - - # Direct conversion - if source_type == memory_type: - array = MemoryType(memory_type).move_to_device(array, gpu_id) - else: - # Convert and log the conversion - from arraybridge.converters import convert_memory - + library = ArrayLibrary.named(memory_type) + library.require_device(device) + if source is not library: logger.debug(f"🔄 UNSTACK_SLICES: Converting array - {source_type} → {memory_type}") - array = convert_memory( - data=array, source_type=source_type, target_type=memory_type, gpu_id=gpu_id - ) + array = convert_memory(array, str(source_type), memory_type, device) + else: + array = library.move_to_device(array, device) # Extract slices along axis 0 (already in the target memory type) - slices = [array[i] for i in range(array.shape[0])] + slices = [ + getattr(array, "__getitem__")(index) for index in range(getattr(array, "shape")[0]) + ] # Validate that all extracted slices are 2D if requested if validate_slices: diff --git a/src/arraybridge/types.py b/src/arraybridge/types.py index 36268fc..a317bc6 100644 --- a/src/arraybridge/types.py +++ b/src/arraybridge/types.py @@ -1,1089 +1,121 @@ -""" -Memory type definitions for arraybridge. +"""``MemoryType``: the names of the registered array libraries as an enum. + +The enum is a view. Its members are generated from the classes registered in +``arraybridge.libraries.ArrayLibrary`` and carry no behaviour; use it where a +name is stored, serialised or offered as a choice, and ask +``MemoryType.X.library`` for anything the library does. A library class +defined after this module was imported appears as a member the next time the +enum is looked at. -This module defines the MemoryType enum and related constants for managing -different array/tensor frameworks. +Importing this module imports no array library, NumPy included. """ from __future__ import annotations -import importlib -import importlib.util -import logging -import os -import sys -from collections.abc import Callable, Iterator, Mapping, MutableMapping, Sequence -from contextlib import AbstractContextManager, contextmanager, nullcontext -from dataclasses import dataclass -from enum import Enum -from pathlib import Path -from typing import TYPE_CHECKING, Any, TypeVar, cast - -if TYPE_CHECKING: - from arraybridge.array_operations import ArrayOperations - -T = TypeVar("T") -logger = logging.getLogger(__name__) -ConversionFunc = Callable[[Any], Any] -DeviceIdsResolver = Callable[[Any], tuple[int, ...]] -DeviceIdResolver = Callable[[Any, Any], int | None] -DeviceScopeFactory = Callable[[Any, int], AbstractContextManager[None]] -DeviceActivator = Callable[[Any, int], None] -FrameworkCleanup = Callable[[Any], None] -ActiveDeviceMover = Callable[[Any, Any, int], Any] -CurrentDeviceResolver = Callable[[Any], int | None] -StreamFactory = Callable[[Any], Any] -StreamScopeFactory = Callable[[Any, Any], AbstractContextManager[None]] -DLPackExporter = Callable[[Any, Any], Any | None] -DLPackValidator = Callable[[Any, Any], bool] -OOMMatcher = Callable[[BaseException, Any | None], bool] -SubprocessEnvironmentResolver = Callable[[Mapping[str, str]], dict[str, str]] - - -def _identity_subprocess_environment( - environment: Mapping[str, str], -) -> dict[str, str]: - return dict(environment) - - -def _nvidia_wheel_library_paths() -> tuple[str, ...]: - """Return native-library directories declared by installed NVIDIA wheels.""" - - try: - spec = importlib.util.find_spec("nvidia") - except (ImportError, ModuleNotFoundError, ValueError): - return () - if spec is None or spec.submodule_search_locations is None: - return () - - paths: set[str] = set() - for package_root in map(Path, spec.submodule_search_locations): - try: - components = tuple(package_root.iterdir()) - except OSError: - continue - for component in components: - for library_directory_name in ("lib", "bin"): - candidate = component / library_directory_name - if candidate.is_dir(): - paths.add(str(candidate)) - return tuple(sorted(paths)) - - -def _nvidia_wheel_subprocess_environment( - environment: Mapping[str, str], -) -> dict[str, str]: - """Prepend installed NVIDIA wheel libraries to a child environment.""" - - prepared = dict(environment) - library_paths = _nvidia_wheel_library_paths() - if not library_paths: - return prepared - - search_variable = "PATH" if os.name == "nt" else "LD_LIBRARY_PATH" - existing_paths = tuple( - path for path in prepared.get(search_variable, "").split(os.pathsep) if path - ) - prepared[search_variable] = os.pathsep.join(dict.fromkeys((*library_paths, *existing_paths))) - return prepared - - -class MemoryContractAttribute(str, Enum): - """Callable metadata keys owned by ArrayBridge's decorator contract.""" - - INPUT = "input_memory_type" - OUTPUT = "output_memory_type" - EXECUTION = "execution_memory_type" - - def read(self, namespace: Any, default: Any = None) -> Any: - """Read this declaration from an attribute or mapping namespace.""" - - if isinstance(namespace, Mapping): - return namespace.get(self.value, default) - return getattr(namespace, self.value, default) - - def write(self, namespace: Any, value: Any) -> None: - """Write this declaration to an attribute or mutable mapping namespace.""" - - if isinstance(namespace, MutableMapping): - namespace[self.value] = value - return - setattr(namespace, self.value, value) - - -@dataclass(frozen=True, slots=True) -class DLPackPayload: - """One exported capsule plus its original protocol-bearing array.""" - - source: Any - capsule: Any - - -DLPackImporter = Callable[[DLPackPayload, Any], Any] - - -def _no_device_ids(module: Any) -> tuple[int, ...]: - del module - return () - - -def _no_device_id(data: Any, module: Any) -> None: - del data, module - return None - - -def _cupy_device_id(data: Any, module: Any) -> int: - del module - return int(data.device.id) - - -def _torch_device_id(data: Any, module: Any) -> int | None: - del module - return int(data.device.index) if data.is_cuda else None - - -def _tensorflow_device_id(data: Any, module: Any) -> int | None: - del module - device = data.device.lower() - return int(device.rsplit(":", maxsplit=1)[-1]) if "gpu" in device else None - - -def _jax_device_id(data: Any, module: Any) -> int | None: - device = data.device - device = device() if callable(device) else device - if getattr(device, "platform", None) != "gpu": - return None - gpu_devices = tuple(candidate for candidate in module.devices() if candidate.platform == "gpu") - identity_match = next( - (index for index, candidate in enumerate(gpu_devices) if candidate is device), - None, - ) - if identity_match is not None: - return identity_match - return next( - (index for index, candidate in enumerate(gpu_devices) if str(candidate) == str(device)), - None, - ) - - -def _pyclesperanto_device_id(data: Any, module: Any) -> int | None: - declared_device = getattr(data, "device", None) - if declared_device is None: - declared_device = module.get_device() - declared_selector = getattr(declared_device, "name", None) - if declared_selector is None: - declared_selector = str(declared_device) - devices = _pyclesperanto_devices(module) - return next( - ( - index - for index, device in enumerate(devices) - if getattr(device, "name", str(device)) == declared_selector - ), - None, - ) - - -def _cupy_device_ids(module: Any) -> tuple[int, ...]: - return tuple(range(int(module.cuda.runtime.getDeviceCount()))) - - -def _torch_device_ids(module: Any) -> tuple[int, ...]: - if not module.cuda.is_available(): - return () - return tuple(range(int(module.cuda.device_count()))) - - -def _tensorflow_device_ids(module: Any) -> tuple[int, ...]: - return tuple(range(len(module.config.list_logical_devices("GPU")))) - - -def _jax_device_ids(module: Any) -> tuple[int, ...]: - return tuple( - range(len(tuple(device for device in module.devices() if device.platform == "gpu"))) - ) - - -def _pyclesperanto_devices(module: Any) -> tuple[Any, ...]: - try: - return tuple(module.list_available_devices("gpu")) - except TypeError: - return tuple(module.list_available_devices()) - - -def _pyclesperanto_select_device( - module: Any, - selector: str | int, - device_type: str | None = None, -) -> None: - if device_type is None: - module.select_device(selector) - return - try: - module.select_device(selector, device_type) - except TypeError: - module.select_device(selector) - - -def _pyclesperanto_device_ids(module: Any) -> tuple[int, ...]: - return tuple(range(len(_pyclesperanto_devices(module)))) - - -def _null_device_scope(module: Any, device_id: int) -> AbstractContextManager[None]: - del module, device_id - return nullcontext() - - -def _cupy_device_scope(module: Any, device_id: int) -> AbstractContextManager[None]: - return cast(AbstractContextManager[None], module.cuda.Device(device_id)) - - -def _torch_device_scope(module: Any, device_id: int) -> AbstractContextManager[None]: - return cast(AbstractContextManager[None], module.cuda.device(device_id)) - - -def _tensorflow_device_scope( - module: Any, - device_id: int, -) -> AbstractContextManager[None]: - return cast(AbstractContextManager[None], module.device(f"/device:GPU:{device_id}")) - - -def _jax_device_scope(module: Any, device_id: int) -> AbstractContextManager[None]: - gpu_devices = tuple(device for device in module.devices() if device.platform == "gpu") - return cast(AbstractContextManager[None], module.default_device(gpu_devices[device_id])) - - -@contextmanager -def _pyclesperanto_device_scope(module: Any, device_id: int) -> Iterator[None]: - current_device = module.get_device() - current_id = _pyclesperanto_device_id(None, module) - current_selector = getattr(current_device, "name", None) - if current_selector is None: - current_selector = str(current_device) - _pyclesperanto_select_device(module, device_id, "gpu") - try: - yield - finally: - if current_id != device_id: - _pyclesperanto_select_device(module, current_selector) - - -def _no_device_activation(module: Any, device_id: int) -> None: - del module, device_id - - -def _cupy_device_activation(module: Any, device_id: int) -> None: - module.cuda.Device(device_id).use() - - -def _torch_device_activation(module: Any, device_id: int) -> None: - module.cuda.set_device(device_id) - - -def _pyclesperanto_device_activation(module: Any, device_id: int) -> None: - _pyclesperanto_select_device(module, device_id, "gpu") - - -def _cupy_cleanup(module: Any) -> None: - module.get_default_memory_pool().free_all_blocks() - module.get_default_pinned_memory_pool().free_all_blocks() - module.cuda.runtime.deviceSynchronize() - - -def _torch_cleanup(module: Any) -> None: - module.cuda.empty_cache() - module.cuda.synchronize() - - -def _identity_device_move(data: Any, module: Any, device_id: int) -> Any: - del module, device_id - return data - - -def _cupy_device_move(data: Any, module: Any, device_id: int) -> Any: - del module, device_id - return data.copy() - - -def _torch_device_move(data: Any, module: Any, device_id: int) -> Any: - del module - return data.to(f"cuda:{device_id}") - - -def _tensorflow_device_move(data: Any, module: Any, device_id: int) -> Any: - del device_id - return module.identity(data) - - -def _jax_device_move(data: Any, module: Any, device_id: int) -> Any: - gpu_devices = tuple(device for device in module.devices() if device.platform == "gpu") - return module.device_put(data, gpu_devices[device_id]) - - -def _pyclesperanto_device_move(data: Any, module: Any, device_id: int) -> Any: - del device_id - result = module.create_like(data) - module.copy(data, result) - return result - - -def _no_current_device(module: Any) -> None: - del module - return None - - -def _cupy_current_device(module: Any) -> int: - return int(module.cuda.runtime.getDevice()) - - -def _torch_current_device(module: Any) -> int | None: - return int(module.cuda.current_device()) if module.cuda.is_available() else None - - -def _cupy_stream(module: Any) -> Any: - return module.cuda.Stream() - - -def _torch_stream(module: Any) -> Any: - return module.cuda.Stream() - - -def _cupy_stream_scope(module: Any, stream: Any) -> AbstractContextManager[None]: - del module - return cast(AbstractContextManager[None], stream) - - -def _torch_stream_scope(module: Any, stream: Any) -> AbstractContextManager[None]: - return cast(AbstractContextManager[None], module.cuda.stream(stream)) - - -def _null_stream_scope(module: Any, stream: Any) -> AbstractContextManager[None]: - del module, stream - return nullcontext() +from collections.abc import Mapping, MutableMapping +from enum import Enum, EnumMeta +from types import MappingProxyType +from arraybridge.libraries import Array, ArrayLibrary -def _has_modern_dlpack_protocol(data: Any) -> bool: - return callable(getattr(data, "__dlpack__", None)) and callable( - getattr(data, "__dlpack_device__", None) - ) +class _RegisteredLibraryNames(EnumMeta): + """Keeps the members equal to the registered library names. -def _cupy_from_dlpack(payload: DLPackPayload, module: Any) -> Any: - if _has_modern_dlpack_protocol(payload.source): - return module.from_dlpack(payload.source) - legacy_importer = getattr(module, "fromDlpack", None) - return NotImplemented if legacy_importer is None else legacy_importer(payload.capsule) + The standard library offers no way to add a member to an existing enum, so + a missing one is entered into the enum's own member tables here. The + overrides below keep ``EnumMeta``'s signatures and add only that step. + """ + def _add_missing_members(cls): + for name in ArrayLibrary.__registry__: + if name not in cls._value2member_map_: + member = object.__new__(cls) + member._name_ = name.upper() + member._value_ = name + cls._member_map_[member._name_] = member + cls._member_names_.append(member._name_) + cls._value2member_map_[name] = member -def _torch_from_dlpack(payload: DLPackPayload, module: Any) -> Any: - return module.from_dlpack(payload.capsule) + def __call__(cls, *args, **kwargs): + cls._add_missing_members() + return super().__call__(*args, **kwargs) + def __getitem__(cls, name): + cls._add_missing_members() + return super().__getitem__(name) -def _export_dlpack(data: Any) -> Any: - for attribute in ("__dlpack__", "to_dlpack", "toDlpack"): - exporter = getattr(data, attribute, None) - if callable(exporter): - return exporter() - raise TypeError(f"{type(data).__name__} does not expose a DLPack exporter") + def __reversed__(cls): + cls._add_missing_members() + return super().__reversed__() + def __dir__(cls): + cls._add_missing_members() + return super().__dir__() -def _tensorflow_from_dlpack(payload: DLPackPayload, module: Any) -> Any: - return module.experimental.dlpack.from_dlpack(payload.capsule) + def __iter__(cls): + cls._add_missing_members() + return super().__iter__() + def __len__(cls): + cls._add_missing_members() + return super().__len__() -def _dlpack_data_pointer(data: Any) -> int | None: - """Return a protocol source pointer when its framework exposes one.""" - - data_ptr = getattr(data, "data_ptr", None) - if callable(data_ptr): - return int(data_ptr()) - cuda_interface = getattr(data, "__cuda_array_interface__", None) - if isinstance(cuda_interface, dict): - pointer, _read_only = cuda_interface.get("data", (None, False)) - return None if pointer is None else int(pointer) - return None - - -def _jax_aligned_dlpack_source(data: Any) -> Any: - """Copy a DLPack source only when JAX's 16-byte alignment is unmet.""" - - pointer = _dlpack_data_pointer(data) - if pointer is None or pointer % 16 == 0: - return data - for attribute in ("clone", "copy"): - copier = getattr(data, attribute, None) - if callable(copier): - return copier() - return data - - -def _jax_from_dlpack(payload: DLPackPayload, module: Any) -> Any: - if not _has_modern_dlpack_protocol(payload.source): - return NotImplemented - return module.dlpack.from_dlpack(_jax_aligned_dlpack_source(payload.source)) - - -def _protocol_dlpack(data: Any, module: Any) -> bool: - del module - return any( - callable(getattr(data, attribute, None)) - for attribute in ("__dlpack__", "toDlpack", "to_dlpack") - ) - - -def _protocol_dlpack_export(data: Any, module: Any) -> Any | None: - del module - try: - return _export_dlpack(data) - except TypeError: - return None - - -def _tensorflow_dlpack(data: Any, module: Any) -> bool: - try: - major, minor = map(int, module.__version__.split(".")[:2]) - except (AttributeError, TypeError, ValueError): - return False - if (major, minor) < (2, 12): - return False - if "gpu" not in str(getattr(data, "device", "")).lower(): - return False - dlpack = getattr(getattr(module, "experimental", None), "dlpack", None) - if not callable(getattr(dlpack, "to_dlpack", None)): - return False - return True - - -def _tensorflow_dlpack_export(data: Any, module: Any) -> Any | None: - if not _tensorflow_dlpack(data, module): - return None - return module.experimental.dlpack.to_dlpack(data) - - -def _message_matches(error: BaseException, *patterns: str) -> bool: - message = str(error).lower() - return any(pattern in message for pattern in patterns) - - -def _exception_type(module: Any, *path: str) -> type[BaseException] | None: - candidate = module - for attribute in path: - candidate = getattr(candidate, attribute, None) - if candidate is None: - return None - return ( - candidate if isinstance(candidate, type) and issubclass(candidate, BaseException) else None - ) - - -def _matches_declared_exception(error: BaseException, module: Any | None, *path: str) -> bool: - if module is None: - return False - exception_type = _exception_type(module, *path) - return exception_type is not None and isinstance(error, exception_type) - - -def _numpy_oom(error: BaseException, module: Any | None) -> bool: - del module - return _message_matches(error, "cannot allocate memory", "memory exhausted") - - -def _never_oom(error: BaseException, module: Any | None) -> bool: - del error, module - return False - - -def _cupy_oom(error: BaseException, module: Any | None) -> bool: - return ( - _matches_declared_exception(error, module, "cuda", "memory", "OutOfMemoryError") - or _matches_declared_exception(error, module, "cuda", "runtime", "CUDARuntimeError") - or _message_matches(error, "out of memory", "cuda_error_out_of_memory") - ) - - -def _torch_oom(error: BaseException, module: Any | None) -> bool: - return _matches_declared_exception( - error, module, "cuda", "OutOfMemoryError" - ) or _message_matches(error, "out of memory", "cuda_error_out_of_memory") - - -def _tensorflow_oom(error: BaseException, module: Any | None) -> bool: - return _matches_declared_exception( - error, module, "errors", "ResourceExhaustedError" - ) or _message_matches(error, "out of memory", "resource_exhausted") - - -def _jax_oom(error: BaseException, module: Any | None) -> bool: - del module - return _message_matches(error, "out of memory", "oom when allocating", "allocation failure") - - -def _pyclesperanto_oom(error: BaseException, module: Any | None) -> bool: - del module - return _message_matches( - error, - "cl_mem_object_allocation_failure", - "cl_out_of_resources", - "out of memory", - ) - - -@dataclass(frozen=True, slots=True) -class FrameworkRuntime: - """Typed execution leaves carried by one ``MemoryType`` declaration.""" - - device_ids: DeviceIdsResolver = _no_device_ids - device_id: DeviceIdResolver = _no_device_id - device_scope: DeviceScopeFactory = _null_device_scope - activate_device: DeviceActivator = _no_device_activation - cleanup: FrameworkCleanup | None = None - move_to_active_device: ActiveDeviceMover = _identity_device_move - current_device: CurrentDeviceResolver = _no_current_device - stream_factory: StreamFactory | None = None - stream_scope: StreamScopeFactory = _null_stream_scope - dlpack_importer: DLPackImporter | None = None - dlpack_exporter: DLPackExporter | None = None - dlpack_validator: DLPackValidator = _protocol_dlpack - oom_matcher: OOMMatcher = _never_oom - subprocess_environment: SubprocessEnvironmentResolver = _identity_subprocess_environment - - -class _MemoryTypeFields: - import_name: str - display_name: str - is_gpu: bool - module_aliases: tuple[str, ...] - import_environment: tuple[tuple[str, str], ...] - _runtime: FrameworkRuntime - - -class MemoryType(_MemoryTypeFields, Enum): - """Array-framework declarations with member-owned runtime capability leaves.""" - - def __new__( - cls, - value: str, - *declaration: Any, - ) -> MemoryType: - ( - import_name, - display_name, - is_gpu, - module_aliases, - import_environment, - runtime, - ) = declaration - member = object.__new__(cls) - member._value_ = value - member.import_name = cast(str, import_name) - member.display_name = cast(str, display_name) - member.is_gpu = cast(bool, is_gpu) - member.module_aliases = cast(tuple[str, ...], module_aliases) - member.import_environment = cast(tuple[tuple[str, str], ...], import_environment) - member._runtime = cast(FrameworkRuntime, runtime) - return member + def __contains__(cls, value): + cls._add_missing_members() + return super().__contains__(value) @property - def _operations(self) -> ArrayOperations: - """Resolve this declaration's operations leaf on first access.""" - - from arraybridge.array_operations import ArrayOperations + def __members__(cls): + cls._add_missing_members() + return MappingProxyType(cls._member_map_) - return ArrayOperations.for_memory(self) + def __getattr__(cls, name): + if not name.startswith("_"): + cls._add_missing_members() + member = cls._member_map_.get(name) + if member is not None: + return member + raise AttributeError(name) - NUMPY = ( - "numpy", - "numpy", - "NumPy", - False, - (), - (), - FrameworkRuntime(oom_matcher=_numpy_oom), - ) - CUPY = ( - "cupy", - "cupy", - "CuPy", - True, - (), - (), - FrameworkRuntime( - device_ids=_cupy_device_ids, - device_id=_cupy_device_id, - device_scope=_cupy_device_scope, - activate_device=_cupy_device_activation, - cleanup=_cupy_cleanup, - move_to_active_device=_cupy_device_move, - current_device=_cupy_current_device, - stream_factory=_cupy_stream, - stream_scope=_cupy_stream_scope, - dlpack_importer=_cupy_from_dlpack, - dlpack_exporter=_protocol_dlpack_export, - oom_matcher=_cupy_oom, - subprocess_environment=_nvidia_wheel_subprocess_environment, - ), - ) - TORCH = ( - "torch", - "torch", - "PyTorch", - True, - (), - (), - FrameworkRuntime( - device_ids=_torch_device_ids, - device_id=_torch_device_id, - device_scope=_torch_device_scope, - activate_device=_torch_device_activation, - cleanup=_torch_cleanup, - move_to_active_device=_torch_device_move, - current_device=_torch_current_device, - stream_factory=_torch_stream, - stream_scope=_torch_stream_scope, - dlpack_importer=_torch_from_dlpack, - dlpack_exporter=_protocol_dlpack_export, - oom_matcher=_torch_oom, - ), - ) - TENSORFLOW = ( - "tensorflow", - "tensorflow", - "TensorFlow", - True, - (), - (("TF_FORCE_GPU_ALLOW_GROWTH", "true"),), - FrameworkRuntime( - device_ids=_tensorflow_device_ids, - device_id=_tensorflow_device_id, - device_scope=_tensorflow_device_scope, - move_to_active_device=_tensorflow_device_move, - dlpack_importer=_tensorflow_from_dlpack, - dlpack_exporter=_tensorflow_dlpack_export, - dlpack_validator=_tensorflow_dlpack, - oom_matcher=_tensorflow_oom, - ), - ) - JAX = ( - "jax", - "jax", - "JAX", - True, - ("jaxlib",), - (("XLA_PYTHON_CLIENT_PREALLOCATE", "false"),), - FrameworkRuntime( - device_ids=_jax_device_ids, - device_id=_jax_device_id, - device_scope=_jax_device_scope, - move_to_active_device=_jax_device_move, - dlpack_importer=_jax_from_dlpack, - dlpack_exporter=_protocol_dlpack_export, - oom_matcher=_jax_oom, - ), - ) - PYCLESPERANTO = ( - "pyclesperanto", - "pyclesperanto", - "pyclesperanto", - True, - (), - (), - FrameworkRuntime( - device_ids=_pyclesperanto_device_ids, - device_id=_pyclesperanto_device_id, - device_scope=_pyclesperanto_device_scope, - activate_device=_pyclesperanto_device_activation, - move_to_active_device=_pyclesperanto_device_move, - oom_matcher=_pyclesperanto_oom, - ), - ) - @property - def recognized_module_names(self) -> frozenset[str]: - """Return top-level module names owned by this framework declaration.""" - - return frozenset((self.import_name, *self.module_aliases)) - - def is_installed(self) -> bool: - """Check package presence without importing the optional framework.""" - - try: - return importlib.util.find_spec(self.import_name) is not None - except (ImportError, ModuleNotFoundError, ValueError): - return False - - def prepare_import(self) -> None: - """Apply declaration-owned coexistence defaults before framework import.""" - - missing_names = tuple( - name for name, _value in self.import_environment if name not in os.environ - ) - if missing_names and self.loaded_module() is not None: - logger.warning( - "%s was loaded before its ArrayBridge import defaults were set; " - "the current process cannot guarantee those import-time settings: %s", - self.display_name, - ", ".join(missing_names), - ) - for name, value in self.import_environment: - os.environ.setdefault(name, value) +class MemoryType(Enum, metaclass=_RegisteredLibraryNames): + """The name of one registered array library.""" @classmethod - def subprocess_environment( - cls, - environment: Mapping[str, str] | None = None, - ) -> dict[str, str]: - """Project framework import requirements into a child environment.""" - - prepared = dict(os.environ if environment is None else environment) - for memory_type in cls: - for name, value in memory_type.import_environment: - prepared.setdefault(name, value) - prepared = memory_type._runtime.subprocess_environment(prepared) - return prepared - - def loaded_module(self) -> Any | None: - """Return an already-loaded framework without causing an import.""" - - return sys.modules.get(self.import_name) - - def import_module(self) -> Any: - """Import this declaration's framework module.""" - - self.prepare_import() - return importlib.import_module(self.import_name) - - def import_if_installed(self) -> Any | None: - """Import this framework only when its package is present.""" - - loaded = self.loaded_module() - if loaded is not None: - self.prepare_import() - return loaded - if not self.is_installed(): - return None - return self.import_module() - - def to_numpy(self, data: Any, module: Any | None = None) -> Any: - """Project one array to NumPy through this declaration's leaf.""" - - framework = module if module is not None else self.import_module() - return self._operations.to_numpy(data, framework) - - def from_numpy( - self, - data: Any, - device_id: int, - module: Any | None = None, - ) -> Any: - """Create one array from NumPy on a declared framework-local device.""" - - framework = module if module is not None else self.import_module() - with self.device_scope(device_id, framework): - return self._operations.from_numpy(data, framework, device_id) - - def stack_arrays( - self, - arrays: list[Any], - device_id: int, - module: Any | None = None, - ) -> Any: - """Stack prepared arrays on a declared framework-local device.""" - - framework = module if module is not None else self.import_module() - with self.device_scope(device_id, framework): - return self._operations.stack(arrays, framework) - - def normalize_planes( - self, - data: Any, - dtype: Any, - scales: Sequence[float | None], - module: Any | None = None, - ) -> Any: - """Normalize one owned plane stack through its native operation owner.""" - if len(scales) != len(data): - raise ValueError("Normalization scales must match the leading plane axis.") - framework = module if module is not None else self.import_module() - with self._array_device_scope(data, framework): - return self._operations.normalize_planes(data, dtype, scales, framework) - - def scale_dtype( - self, - data: Any, - target_dtype: Any, - module: Any | None = None, - ) -> Any: - """Scale an array through this declaration's typed operation leaf.""" - - framework = module if module is not None else self.import_if_installed() - if framework is None: - return data - device_id = self.device_id_of(data, framework) if hasattr(data, "dtype") else None - scope = nullcontext() if device_id is None else self.device_scope(device_id, framework) - with scope: - return self._operations.scale_dtype(data, target_dtype, framework) - - def canonical_dtype_name(self, dtype: Any) -> str: - """Return this framework member's portable dtype identity.""" - - return self._operations.dtype_name(dtype) - - def astype( - self, - data: Any, - dtype: Any, - module: Any | None = None, - ) -> Any: - """Cast one array through this framework member's operation leaf.""" - - framework = module if module is not None else self.import_module() - with self._array_device_scope(data, framework): - return self._operations.cast(data, dtype, framework) - - def _array_device_scope(self, reference: Any, module: Any) -> AbstractContextManager[None]: - """Derive every native operation scope from its actual input array.""" - device_id = self.device_id_of(reference, module) - return nullcontext() if device_id is None else self.device_scope(device_id, module) - - def reshape(self, data: Any, shape: Sequence[int], module: Any | None = None) -> Any: - """Reshape on the input device without an implicit host projection.""" - framework = module if module is not None else self.import_module() - with self._array_device_scope(data, framework): - return self._operations.reshape(data, tuple(shape), framework) - - def broadcast_to(self, data: Any, shape: Sequence[int], module: Any | None = None) -> Any: - """Broadcast on the input device without an implicit host projection.""" - framework = module if module is not None else self.import_module() - with self._array_device_scope(data, framework): - return self._operations.broadcast_to(data, tuple(shape), framework) - - def ones_like( - self, - reference: Any, - *, - shape: Sequence[int] | None = None, - dtype: Any = bool, - module: Any | None = None, - ) -> Any: - """Allocate ones on the supplied reference's framework-local device.""" - framework = module if module is not None else self.import_module() - with self._array_device_scope(reference, framework): - return self._operations.ones_like( - reference, - tuple(reference.shape if shape is None else shape), - dtype, - framework, - ) - - def logical_and( - self, - left: Any, - right: Any, - module: Any | None = None, - ) -> Any: - """Intersect two arrays through this framework member's operation leaf.""" - - framework = module if module is not None else self.import_module() - with self._array_device_scope(left, framework): - return self._operations.logical_and(left, right, framework) - - def available_device_ids(self, module: Any | None = None) -> tuple[int, ...]: - """Return every framework-local GPU device identifier.""" - - if not self.is_gpu: - return () - framework = module if module is not None else self.import_if_installed() - return () if framework is None else self._runtime.device_ids(framework) - - def require_device(self, device_id: int, module: Any | None = None) -> Any: - """Return the framework after proving that its local device exists.""" - - if not self.is_gpu: - return module - framework = module if module is not None else self.import_module() - available = self.available_device_ids(framework) - if device_id not in available: - raise ValueError( - f"{self.display_name} device {device_id} is unavailable; " - f"available device IDs are {available}" - ) - return framework - - def device_id_of(self, data: Any, module: Any | None = None) -> int | None: - """Return this framework's local device identifier for an array.""" - - if not self.is_gpu: - return None - framework = module if module is not None else self.import_module() - return self._runtime.device_id(data, framework) - - def device_scope( - self, - device_id: int, - module: Any | None = None, - ) -> AbstractContextManager[None]: - """Return this framework member's scoped device activation leaf.""" - - if not self.is_gpu: - return nullcontext() - framework = self.require_device(device_id, module) - return self._runtime.device_scope(framework, device_id) - - def activate_device(self, device_id: int, module: Any | None = None) -> None: - """Activate a process-global device where the framework supports it.""" - - if not self.is_gpu: - return - framework = self.require_device(device_id, module) - self._runtime.activate_device(framework, device_id) - - def move_to_device( - self, - data: Any, - device_id: int, - module: Any | None = None, - ) -> Any: - """Move an array to one framework-local device without leaking selection.""" - - if not self.is_gpu: - return data - framework = self.require_device(device_id, module) - if self.device_id_of(data, framework) == device_id: - return data - with self.device_scope(device_id, framework): - return self._runtime.move_to_active_device(data, framework, device_id) + def _missing_(cls, value: object) -> MemoryType | None: + cls._add_missing_members() + member = cls._value2member_map_.get(value) + return member if isinstance(member, cls) else None @property - def supports_dlpack(self) -> bool: - """Whether this framework declares a DLPack import leaf.""" - - return self._runtime.dlpack_importer is not None - - def supports_dlpack_data(self, data: Any, module: Any | None = None) -> bool: - """Validate DLPack export for an array owned by this framework.""" - - if self._runtime.dlpack_exporter is None: - return False - framework = module if module is not None else self.import_module() - return self._runtime.dlpack_validator(data, framework) - - def export_dlpack(self, data: Any, module: Any | None = None) -> DLPackPayload | None: - """Export one array through this framework's declaration, if supported.""" - - exporter = self._runtime.dlpack_exporter - if exporter is None: - return None - framework = module if module is not None else self.import_module() - capsule = exporter(data, framework) - return None if capsule is None else DLPackPayload(source=data, capsule=capsule) - - def from_dlpack( - self, - data: DLPackPayload | Any, - module: Any | None = None, - ) -> Any: - """Import DLPack data through this framework's declared leaf.""" - - importer = self._runtime.dlpack_importer - if importer is None: - raise NotImplementedError(f"DLPack not supported for {self.value}") - framework = module if module is not None else self.import_module() - payload = data if isinstance(data, DLPackPayload) else DLPackPayload(data, data) - return importer(payload, framework) + def library(self) -> ArrayLibrary[Array]: + """The library this name stands for.""" - def convert_to(self, data: Any, target: MemoryType, device_id: int) -> Any: - """Convert one array through the source and target declarations.""" + return ArrayLibrary.named(self.value) - if self is target: - return target.move_to_device(data, device_id) - if target.supports_dlpack: - try: - payload = self.export_dlpack(data) - if payload is not None: - module = target.import_module() - with target.device_scope(device_id, module): - result = target.from_dlpack(payload, module) - if result is not NotImplemented: - return target.move_to_device(result, device_id, module) - except Exception as error: - logger.warning( - "DLPack conversion from %s to %s failed: %s. Using CPU roundtrip.", - self.value, - target.value, - error, - ) +MemoryType._add_missing_members() - return target.from_numpy(self.to_numpy(data), device_id) - def current_device_id(self, module: Any | None = None) -> int | None: - """Return the framework-local process device used for new streams.""" - - if not self.is_gpu: - return None - framework = module if module is not None else self.import_module() - return self._runtime.current_device(framework) - - def create_stream(self, module: Any | None = None) -> Any | None: - """Create a stream on the framework's current device when supported.""" - - factory = self._runtime.stream_factory - if factory is None: - return None - framework = module if module is not None else self.import_module() - return factory(framework) - - def stream_scope( - self, - stream: Any, - module: Any | None = None, - ) -> AbstractContextManager[None]: - """Return this framework member's execution scope for one stream.""" +class MemoryContractAttribute(str, Enum): + """Names of the attributes the memory decorators set on a function.""" - if stream is None: - return nullcontext() - framework = module if module is not None else self.import_module() - return self._runtime.stream_scope(framework, stream) + INPUT = "input_memory_type" + OUTPUT = "output_memory_type" + EXECUTION = "execution_memory_type" - def is_oom_error(self, error: BaseException) -> bool: - """Classify an error without importing an optional framework.""" + def read(self, namespace: object, default: str | None = None) -> str | None: + """Read this attribute from an object or a mapping.""" - return self._runtime.oom_matcher(error, self.loaded_module()) + if isinstance(namespace, Mapping): + return namespace.get(self.value, default) + return getattr(namespace, self.value, default) - def cleanup_loaded(self, device_id: int | None = None) -> None: - """Clean an already-loaded GPU framework without importing absent modules.""" + def write(self, namespace: object, value: str) -> None: + """Write this attribute to an object or a mutable mapping.""" - if not self.is_gpu: - return - framework = self.loaded_module() - if framework is None: - return - cleanup = self._runtime.cleanup - if cleanup is None: - return - available = self.available_device_ids(framework) - if not available: + if isinstance(namespace, MutableMapping): + namespace[self.value] = value return - targets = available if device_id is None else (device_id,) - for target in targets: - with self.device_scope(target, framework): - cleanup(framework) - - -# Memory type sets -CPU_MEMORY_TYPES: frozenset[MemoryType] = frozenset( - memory_type for memory_type in MemoryType if not memory_type.is_gpu -) -GPU_MEMORY_TYPES: frozenset[MemoryType] = frozenset( - memory_type for memory_type in MemoryType if memory_type.is_gpu -) -SUPPORTED_MEMORY_TYPES: frozenset[MemoryType] = CPU_MEMORY_TYPES | GPU_MEMORY_TYPES - -# String value sets for validation -VALID_MEMORY_TYPES = frozenset(mt.value for mt in MemoryType) -VALID_GPU_MEMORY_TYPES = frozenset(mt.value for mt in GPU_MEMORY_TYPES) - -# Compatibility constants are generated projections of the enum declaration. -for _memory_type in MemoryType: - globals()[f"MEMORY_TYPE_{_memory_type.name}"] = _memory_type.value + setattr(namespace, self.value, value) diff --git a/src/arraybridge/utils.py b/src/arraybridge/utils.py index 22e85b9..cd802e3 100644 --- a/src/arraybridge/utils.py +++ b/src/arraybridge/utils.py @@ -8,16 +8,12 @@ import importlib import logging -from typing import Any - -from arraybridge.types import MemoryType - -from .exceptions import MemoryConversionError +from types import ModuleType logger = logging.getLogger(__name__) -def optional_import(module_name: str) -> Any | None: +def optional_import(module_name: str) -> ModuleType | None: """Import an optional module, returning ``None`` when it is unavailable. Args: @@ -48,7 +44,7 @@ def optional_import(module_name: str) -> Any | None: raise -def _ensure_module(module_name: str) -> Any: +def _ensure_module(module_name: str) -> ModuleType: """ Ensure a module is imported and meets version requirements. @@ -72,7 +68,7 @@ def _ensure_module(module_name: str) -> Any: return module -def _supports_cuda_array_interface(obj: Any) -> bool: +def _supports_cuda_array_interface(obj: object) -> bool: """ Check if an object supports the CUDA Array Interface. @@ -85,89 +81,10 @@ def _supports_cuda_array_interface(obj: Any) -> bool: return hasattr(obj, "__cuda_array_interface__") -def _supports_dlpack(obj: Any) -> bool: +def _supports_dlpack(obj: object) -> bool: """Return whether an object exposes a standard DLPack export protocol.""" return any( callable(getattr(obj, attribute, None)) for attribute in ("__dlpack__", "toDlpack", "to_dlpack") ) - - -# Compatibility adapters over declaration-owned device operations. - - -def _get_device_id(data: Any, memory_type: str) -> int | None: - """ - Get the declaration-owned GPU device ID from a data object. - - Args: - data: The data object - memory_type: The memory type - - Returns: - The GPU device ID or None if not applicable - - Raises: - MemoryConversionError: If the device ID cannot be determined for a GPU memory type - """ - mem_type = MemoryType(memory_type) - try: - return mem_type.device_id_of(data) - except Exception as e: - raise MemoryConversionError( - source_type=memory_type, - target_type=memory_type, - method="device_identification", - reason=f"Failed to identify the {mem_type.value} device: {e}", - ) from e - - -def _set_device(memory_type: str, device_id: int) -> None: - """ - Set the current device through its memory-type declaration. - - Args: - memory_type: The memory type - device_id: The GPU device ID - - Raises: - MemoryConversionError: If the device cannot be set - """ - mem_type = MemoryType(memory_type) - try: - mem_type.activate_device(device_id) - except Exception as e: - raise MemoryConversionError( - source_type=memory_type, - target_type=memory_type, - method="device_selection", - reason=f"Failed to set {mem_type.value} device to {device_id}: {e}", - ) from e - - -def _move_to_device(data: Any, memory_type: str, device_id: int) -> Any: - """ - Move data through its memory-type declaration. - - Args: - data: The data to move - memory_type: The memory type - device_id: The target GPU device ID - - Returns: - The data on the target device - - Raises: - MemoryConversionError: If the data cannot be moved to the specified device - """ - mem_type = MemoryType(memory_type) - try: - return mem_type.move_to_device(data, device_id) - except Exception as e: - raise MemoryConversionError( - source_type=memory_type, - target_type=memory_type, - method="device_movement", - reason=f"Failed to move {mem_type.value} array to device {device_id}: {e}", - ) from e diff --git a/tests/conftest.py b/tests/conftest.py index de0944c..9ebdac3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,160 +1,84 @@ -"""Pytest configuration and fixtures for arraybridge tests.""" +"""Shared fixtures: every test starts with clean process-wide library state.""" + +import importlib.util +import os +import subprocess +import sys +import threading -import numpy as np import pytest +from arraybridge.conversion import ConversionObserver, ConversionRecord +from arraybridge.libraries import ArrayLibrary, CudaLibrary -# Helper functions for safe module checking -def _module_available(module_name): - """Check if a module is available without triggering ImportError.""" - try: - __import__(module_name) - return True - except ImportError: - return False - - -def _module_has_attribute(module_name, attribute): - """Check if a module is available and has an attribute.""" - try: - module = __import__(module_name) - return hasattr(module, attribute) - except ImportError: - return False - - -def _can_import_and_has_cuda(module_name): - """Check if module is available and has CUDA.""" - try: - module = __import__(module_name) - if module_name == "cupy": - return hasattr(module, "cuda") - elif module_name == "torch": - return hasattr(module, "cuda") - elif module_name == "tensorflow": - return hasattr(module, "config") - elif module_name == "jax": - return hasattr(module, "numpy") - elif module_name == "pyclesperanto": - return hasattr(module, "get_device") - return False - except ImportError: - return False - - -def pytest_configure(config): - """Configure pytest with custom markers.""" - config.addinivalue_line("markers", "gpu: tests that require actual GPU hardware") - config.addinivalue_line("markers", "cupy: tests that require CuPy") - config.addinivalue_line("markers", "torch: tests that require PyTorch") - config.addinivalue_line("markers", "tensorflow: tests that require TensorFlow") - config.addinivalue_line("markers", "jax: tests that require JAX") - config.addinivalue_line("markers", "pyclesperanto: tests that require pyclesperanto") - config.addinivalue_line("markers", "slow: tests that take a long time to run") +@pytest.fixture(autouse=True) +def process_state_restored(monkeypatch): + """Undo what library calls leave behind in the process. -@pytest.fixture -def sample_2d_array(): - """Create a sample 2D NumPy array for testing.""" - return np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32) + Import defaults written to ``os.environ`` are removed again, and the CUDA + streams a test's fake module created for this thread are forgotten. + """ + monkeypatch.setattr(CudaLibrary, "_thread_streams", threading.local()) + environment = dict(os.environ) + yield + for name in set(os.environ) - set(environment): + del os.environ[name] + os.environ.update(environment) @pytest.fixture -def sample_3d_array(): - """Create a sample 3D NumPy array for testing.""" - return np.random.rand(5, 10, 10).astype(np.float32) +def absent_libraries(monkeypatch): + """Every array library but NumPy is neither loaded nor installed; imports fail the test.""" + for library in ArrayLibrary.all(): + if library.name != "numpy": + monkeypatch.delitem(sys.modules, library.import_name, raising=False) + monkeypatch.setattr(importlib.util, "find_spec", lambda name: None) -@pytest.fixture -def sample_slices(): - """Create a list of 2D slices for testing.""" - return [np.random.rand(10, 10).astype(np.float32) for _ in range(5)] + import_module = importlib.import_module + + def fail_import(name, package=None): + if name == "numpy": + return import_module(name, package) + raise AssertionError(f"unexpected import of {name}") + + monkeypatch.setattr(importlib, "import_module", fail_import) @pytest.fixture -def sample_uint8_array(): - """Create a sample uint8 array for dtype testing.""" - return np.random.randint(0, 256, size=(10, 10), dtype=np.uint8) +def loaded_libraries_after(): + """Run code in a fresh interpreter; the array-library modules it left loaded.""" + + def run(code): + report = ( + "import sys\n" + "from arraybridge.libraries import ArrayLibrary\n" + "names = {library.import_name for library in ArrayLibrary.all()}\n" + "print(sorted(names & set(sys.modules)))\n" + ) + finished = subprocess.run( + [sys.executable, "-c", code + report], capture_output=True, text=True, check=False + ) + assert finished.returncode == 0, finished.stderr + return finished.stdout.strip() + + return run + + +class RecordedConversions(ConversionObserver): + """Keeps every record it is given.""" + + def __init__(self): + self.records: list[ConversionRecord] = [] + + def record(self, record): + self.records.append(record) @pytest.fixture -def sample_uint16_array(): - """Create a sample uint16 array for dtype testing.""" - return np.random.randint(0, 65536, size=(10, 10), dtype=np.uint16) - - -# Framework availability fixtures -@pytest.fixture(scope="session") -def torch_available(): - """Check if PyTorch is available.""" - try: - import torch # noqa: F401 - return True - except ImportError: - return False - - -@pytest.fixture(scope="session") -def cupy_available(): - """Check if CuPy is available and has GPU access.""" - try: - import cupy as cp - - # Try to create a small array to verify GPU access - _ = cp.array([1, 2, 3]) - return True - except (ImportError, Exception): - # CuPy not installed or no GPU available - return False - - -@pytest.fixture(scope="session") -def tensorflow_available(): - """Check if TensorFlow is available.""" - try: - import tensorflow as tf # noqa: F401 - return True - except ImportError: - return False - - -@pytest.fixture(scope="session") -def jax_available(): - """Check if JAX is available.""" - try: - import jax # noqa: F401 - return True - except ImportError: - return False - - -@pytest.fixture(scope="session") -def pyclesperanto_available(): - """Check if pyclesperanto is available.""" - try: - import pyclesperanto_prototype # noqa: F401 - return True - except ImportError: - return False - - -@pytest.fixture(scope="session") -def gpu_available(): - """Check if a GPU is available (CUDA or similar).""" - try: - import torch - - if torch.cuda.is_available(): - return True - except ImportError: - pass - - try: - import cupy as cp - - _ = cp.array([1]) - return True - except (ImportError, Exception): - pass - - return False +def recorded(): + """The records of the conversions made during the test.""" + + with RecordedConversions().installed() as observer: + yield observer.records diff --git a/tests/fakes.py b/tests/fakes.py new file mode 100644 index 0000000..e29e0ab --- /dev/null +++ b/tests/fakes.py @@ -0,0 +1,291 @@ +"""Stand-ins for array libraries that are not installed. + +A library class always works on ``sys.modules[import_name]``, so a test installs +one of these with ``monkeypatch.setitem(sys.modules, "cupy", FakeCupy())``. Each +fake implements only what the library class calls and writes what was done to +``events``. Fake arrays set ``__module__`` to the library's module name, which is +how ``ArrayLibrary.of`` recognises them. +""" + +from dataclasses import dataclass +from types import SimpleNamespace + +import numpy as np + + +class DeviceScope: + """Makes ``index`` the fake module's current device for a block.""" + + def __init__(self, module, index): + self.module = module + self.index = index + self.previous = None + + def __enter__(self): + self.previous = self.module.current + self.module.current = self.index + self.module.events.append(("enter", self.index)) + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.module.current = self.previous + self.module.events.append(("exit", self.index)) + return False + + +class FakeStream: + """A CUDA stream: remembers the device it was created on and is its own scope.""" + + def __init__(self, module): + self.module = module + self.device_index = module.current + + def __enter__(self): + self.module.events.append(("stream enter", self)) + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.module.events.append(("stream exit", self)) + return False + + +class _Pool: + def __init__(self, name, events): + self.name = name + self.events = events + + def free_all_blocks(self): + self.events.append(("free", self.name)) + + +@dataclass(frozen=True) +class FakeCapsule: + """What a fake array hands over through DLPack: its buffer and where it lives.""" + + values: np.ndarray + device_index: int | None + + +class _FakeArray: + def __init__(self, values): + self.values = np.asarray(values) + + @property + def shape(self): + return self.values.shape + + @property + def dtype(self): + return self.values.dtype + + +class FakeCupyArray(_FakeArray): + __module__ = "cupy" + + def __init__(self, values, device_index, module): + super().__init__(values) + self.device = SimpleNamespace(id=device_index) + self.module = module + + def copy(self): + return type(self)(self.values.copy(), self.module.current, self.module) + + def get(self): + return self.values.copy() + + def __getitem__(self, index): + return type(self)(self.values[index], self.device.id, self.module) + + def __dlpack__(self): + return FakeCapsule(self.values, self.device.id) + + def __dlpack_device__(self): + return (2, self.device.id) + + +class FakeCupy: + def __init__(self, device_count=2): + self.events = [] + self.current = 0 + self.cuda = SimpleNamespace( + runtime=SimpleNamespace( + getDeviceCount=lambda: device_count, + getDevice=lambda: self.current, + deviceSynchronize=lambda: self.events.append(("synchronize", None)), + ), + Device=lambda index: DeviceScope(self, index), + Stream=lambda: FakeStream(self), + ) + self._pools = (_Pool("regular", self.events), _Pool("pinned", self.events)) + + def get_default_memory_pool(self): + return self._pools[0] + + def get_default_pinned_memory_pool(self): + return self._pools[1] + + def array(self, data): + return FakeCupyArray(data, self.current, self) + + def stack(self, values, axis): + stacked = np.stack([value.values for value in values], axis=axis) + return FakeCupyArray(stacked, self.current, self) + + +class FakeTorchTensor(_FakeArray): + __module__ = "torch" + + def __init__(self, values, cuda_index=None): + super().__init__(values) + self.device = SimpleNamespace(index=cuda_index) + + @property + def is_cuda(self): + return self.device.index is not None + + def cpu(self): + return type(self)(self.values) if self.is_cuda else self + + def to(self, device): + return type(self)(self.values, int(device.removeprefix("cuda:"))) + + def numpy(self): + return self.values.copy() + + def element_size(self): + return self.values.dtype.itemsize + + def nelement(self): + return self.values.size + + +class FakeTorch: + """torch with main memory and ``device_count`` CUDA devices (none by default).""" + + def __init__(self, device_count=0): + self.events = [] + self.current = 0 + self.cuda = SimpleNamespace( + is_available=lambda: device_count > 0, + device_count=lambda: device_count, + current_device=lambda: self.current, + set_device=lambda index: self.events.append(("set_device", index)), + device=lambda index: DeviceScope(self, index), + Stream=lambda: FakeStream(self), + stream=lambda stream: stream, + empty_cache=lambda: self.events.append(("empty_cache", self.current)), + synchronize=lambda: self.events.append(("synchronize", self.current)), + ) + + def from_numpy(self, data): + self.events.append(("from_numpy", data)) + return FakeTorchTensor(data) + + def from_dlpack(self, capsule): + return FakeTorchTensor(capsule.values, capsule.device_index) + + +class FakeTensorflowTensor(_FakeArray): + __module__ = "tensorflow" + + def __init__(self, values, gpu_index): + super().__init__(values) + self.device = f"/device:GPU:{gpu_index}" + + def __getitem__(self, index): + return type(self)(self.values[index], int(self.device.rsplit(":", 1)[-1])) + + def numpy(self): + return self.values + + +class FakeTensorflow: + """TensorFlow with two GPUs; new tensors land on the device in scope.""" + + def __init__(self): + self.events = [] + self.current = 0 + self.config = SimpleNamespace( + list_logical_devices=lambda kind: [object(), object()] if kind == "GPU" else [] + ) + + def device(self, name): + return DeviceScope(self, int(name.rsplit(":", 1)[-1])) + + def identity(self, data): + return FakeTensorflowTensor(data.values, self.current) + + def stack(self, values, axis): + stacked = np.stack([value.values for value in values], axis=axis) + return FakeTensorflowTensor(stacked, self.current) + + +class FakePyclesperantoArray(_FakeArray): + __module__ = "pyclesperanto" + + def __init__(self, values, device_name): + super().__init__(values) + self.device = SimpleNamespace(name=device_name) + + +class FakePyclesperanto: + """pyclesperanto with GPUs ``gpu0`` and ``gpu1``; ``current`` may be any device.""" + + gpu_devices = ("gpu0", "gpu1") + + def __init__(self, current="gpu0"): + self.current = current + self.events = [] + + def list_available_devices(self, device_type=None): + assert device_type in (None, "gpu") + return self.gpu_devices + + def get_device(self): + return self.current + + def select_device(self, selector, device_type=None): + if isinstance(selector, int): + assert device_type in (None, "gpu") + self.current = self.gpu_devices[selector] + else: + self.current = selector + + def create_like(self, data): + return ("copy", data) + + def copy(self, source, target): + self.events.append(("copy", source, target)) + + def pull(self, data): + return data.values + + def push(self, data): + return FakePyclesperantoArray(data, self.current) + + def create(self, shape, dtype): + self.events.append(("create", shape)) + return FakePyclesperantoArray(np.empty(shape, dtype=dtype), self.current) + + def copy_slice(self, source, target, index): + self.events.append(("copy_slice", index)) + target.values[index] = source.values + return target + + def concatenate_along_z(self, left, right): + self.events.append(("concatenate", left.shape, right.shape)) + planes = [ + each.values if each.values.ndim == 3 else each.values[None] for each in (left, right) + ] + return FakePyclesperantoArray(np.concatenate(planes, axis=0), self.current) + + +def fake_jax(*platforms): + """A jax module that lists one device per given platform (``"cpu"``, ``"gpu"``).""" + + devices = [SimpleNamespace(platform=platform) for platform in platforms] + + def list_devices(platform=None): + return [each for each in devices if platform in (None, each.platform)] + + return SimpleNamespace(devices=list_devices) diff --git a/tests/test_array_operations_family.py b/tests/test_array_operations_family.py deleted file mode 100644 index 803cf69..0000000 --- a/tests/test_array_operations_family.py +++ /dev/null @@ -1,74 +0,0 @@ -"""The ArrayOperations family: one leaf per MemoryType, found by its declaration.""" - -import ast -from pathlib import Path - -import numpy as np -import pytest - -from arraybridge import MemoryType -from arraybridge.array_operations import ArrayOperations - -SOURCE = Path(__file__).parents[1] / "src" / "arraybridge" - - -def test_every_memory_type_has_exactly_one_leaf(): - assert set(ArrayOperations.__registry__) == {member.value for member in MemoryType} - for member in MemoryType: - leaf = ArrayOperations.for_memory(member) - assert leaf.memory_type == member.value - assert ArrayOperations.for_memory(member) is leaf - assert member._operations is leaf - - -def test_leaves_are_not_named_by_strings_or_module_singletons(): - types_source = (SOURCE / "types.py").read_text() - assert "_OPERATIONS" not in types_source - tree = ast.parse((SOURCE / "array_operations.py").read_text()) - module_level_names = { - target.id - for node in tree.body - if isinstance(node, ast.Assign) - for target in node.targets - if isinstance(target, ast.Name) - } - assert not {name for name in module_level_names if name.endswith("_OPERATIONS")} - - -def test_numpy_primitives(): - ops = ArrayOperations.for_memory(MemoryType.NUMPY) - stack = np.arange(24, dtype=np.uint16).reshape(2, 3, 4) - assert ops.module is np - assert ops.accepts(stack) and not ops.accepts([1, 2]) - np.testing.assert_array_equal(ops.max(stack, 0), stack[1]) - np.testing.assert_array_equal(ops.min(stack, 0), stack[0]) - np.testing.assert_array_equal(ops.mean(stack, 0), stack.mean(axis=0)) - np.testing.assert_array_equal(ops.sum(stack, (1, 2)), stack.sum(axis=(1, 2))) - assert ops.prepend_axis(stack[0]).shape == (1, 3, 4) - assert ops.astype(stack, np.float32).dtype == np.float32 - ramp = ops.linspace(0, 1, 4, endpoint=False) - np.testing.assert_array_equal(ramp, [0, 0.25, 0.5, 0.75]) - weights = ops.assign(ops.ones(5, np.float32), slice(-2, None), ramp[:2]) - np.testing.assert_array_equal(weights, [1, 1, 1, 0, 0.25]) - np.testing.assert_array_equal(ops.outer(ramp, ramp), np.outer(ramp, ramp)) - assert ops.floor(2.7) == 2 - - -def test_jax_primitives_are_functional(): - jnp = pytest.importorskip("jax.numpy") - ops = ArrayOperations.for_memory(MemoryType.JAX) - values = jnp.ones(4, dtype=jnp.float32) - updated = ops.assign(values, slice(None, 2), jnp.zeros(2, dtype=jnp.float32)) - np.testing.assert_array_equal(np.asarray(updated), [0, 0, 1, 1]) - np.testing.assert_array_equal(np.asarray(values), [1, 1, 1, 1]) - stack = jnp.arange(24, dtype=jnp.uint16).reshape(2, 3, 4) - assert ops.mean(stack, 0).dtype == jnp.float32 - - -def test_tensorflow_dtype_names_come_from_the_numpy_dtype(): - class TensorflowDtype: - as_numpy_dtype = np.uint16 - - ops = ArrayOperations.for_memory(MemoryType.TENSORFLOW) - assert ops.dtype_name(TensorflowDtype()) == "uint16" - assert MemoryType.TENSORFLOW.canonical_dtype_name(TensorflowDtype()) == "uint16" diff --git a/tests/test_conversion.py b/tests/test_conversion.py new file mode 100644 index 0000000..1187e4d --- /dev/null +++ b/tests/test_conversion.py @@ -0,0 +1,494 @@ +"""Conversion routes: which mechanism is planned, which one ran, and what was recorded.""" + +import json +import pickle +import sys +from types import SimpleNamespace + +import numpy as np +import pytest + +from arraybridge import conversion, libraries +from arraybridge.conversion import ( + ConversionCounters, + ConversionObserver, + ConversionRecord, + ConversionRoute, + DeviceMove, + DLPackExchange, + HostCopy, +) +from arraybridge.devices import HOST, AcceleratorDevice +from arraybridge.libraries import ( + AcceleratorDevices, + ArrayLibrary, + DLPackImport, + DLPackPayload, + HostDevices, +) +from tests.fakes import FakeCupy, FakeTorch, FakeTorchTensor + +LIBRARIES = ArrayLibrary.all() +NUMPY, CUPY, TORCH, TENSORFLOW, JAX, PYCLESPERANTO = ( + ArrayLibrary.named(name) + for name in ("numpy", "cupy", "torch", "tensorflow", "jax", "pyclesperanto") +) +GPU = AcceleratorDevice(0) + + +def placements(): + """Every (library, device) a conversion can start from or end at, by declaration.""" + + return [ + (library, device) + for library in LIBRARIES + for role, device in ((HostDevices, HOST), (AcceleratorDevices, GPU)) + if isinstance(library, role) + ] + + +def planned(source, source_device, target, target_device): + return ConversionRoute.planned(source, source_device, target, target_device).mechanism + + +class TestPlanning: + def test_routes_are_tried_in_declaration_order(self): + assert list(ConversionRoute.__registry__.items()) == [ + ("device move", DeviceMove), + ("dlpack", DLPackExchange), + ("host copy", HostCopy), + ] + + @pytest.mark.parametrize(("source", "source_device"), placements(), ids=str) + def test_every_pair_has_a_route_ending_in_a_host_copy(self, source, source_device): + for target, target_device in placements(): + candidates = ConversionRoute.candidates(source, source_device, target, target_device) + mechanism = planned(source, source_device, target, target_device) + + assert candidates[0].mechanism == mechanism + assert candidates[-1] is HostCopy + if source is target: + assert mechanism == "device move" + elif source is NUMPY or target in (NUMPY, PYCLESPERANTO): + assert mechanism == "host copy" + else: + assert mechanism in ("dlpack", "host copy") + + @pytest.mark.parametrize( + ("source", "source_device", "target", "target_device", "mechanism"), + [ + (CUPY, GPU, TORCH, GPU, "dlpack"), + (CUPY, GPU, TORCH, HOST, "dlpack"), + (TORCH, GPU, CUPY, GPU, "dlpack"), + (TORCH, HOST, CUPY, GPU, "host copy"), + (TORCH, HOST, JAX, HOST, "dlpack"), + (TORCH, GPU, JAX, GPU, "dlpack"), + (JAX, HOST, TORCH, HOST, "dlpack"), + (JAX, GPU, TORCH, GPU, "dlpack"), + (JAX, HOST, CUPY, GPU, "host copy"), + (TENSORFLOW, GPU, TORCH, GPU, "dlpack"), + (TENSORFLOW, HOST, TORCH, HOST, "host copy"), + (TORCH, HOST, TENSORFLOW, HOST, "dlpack"), + (PYCLESPERANTO, GPU, TORCH, GPU, "host copy"), + (TORCH, GPU, PYCLESPERANTO, GPU, "host copy"), + (TORCH, HOST, TORCH, GPU, "device move"), + ], + ids=str, + ) + def test_expected_mechanism(self, source, source_device, target, target_device, mechanism): + assert planned(source, source_device, target, target_device) == mechanism + + def test_planning_every_pair_imports_no_array_library(self, loaded_libraries_after): + loaded = loaded_libraries_after( + "from arraybridge.conversion import ConversionRoute\n" + "from arraybridge.devices import HOST, AcceleratorDevice\n" + "from arraybridge.libraries import AcceleratorDevices, ArrayLibrary, HostDevices\n" + "placements = [\n" + " (library, device)\n" + " for library in ArrayLibrary.all()\n" + " for role, device in (\n" + " (HostDevices, HOST), (AcceleratorDevices, AcceleratorDevice(0))\n" + " )\n" + " if isinstance(library, role)\n" + "]\n" + "assert len(placements) == 9, placements\n" + "for source in placements:\n" + " for target in placements:\n" + " route = ConversionRoute.planned(*source, *target)\n" + " assert issubclass(route, ConversionRoute) and route.mechanism\n" + ) + assert loaded == "[]" + + +class TestDLPack: + def test_libraries_that_wrap_foreign_buffers(self): + importers = {library.name for library in LIBRARIES if library.imports_dlpack} + + assert importers == {"cupy", "torch", "tensorflow", "jax"} + assert importers == { + library.name for library in LIBRARIES if isinstance(library, DLPackImport) + } + assert not CUPY.imports_dlpack_from(HOST) and CUPY.imports_dlpack_from(GPU) + assert TORCH.imports_dlpack_from(HOST) and TORCH.imports_dlpack_from(GPU) + + @pytest.mark.parametrize( + ("version", "device", "exportable"), + [ + ("2.15.0", "GPU:0", True), + ("2.10.0", "GPU:0", False), + ("2.15.0", "CPU:0", False), + ("unknown", "GPU:0", False), + ], + ) + def test_tensorflow_exports_accelerator_tensors_of_recent_versions( + self, version, device, exportable, monkeypatch + ): + capsule = object() + module = SimpleNamespace( + __version__=version, + experimental=SimpleNamespace(dlpack=SimpleNamespace(to_dlpack=lambda data: capsule)), + ) + monkeypatch.setitem(sys.modules, "tensorflow", module) + data = SimpleNamespace(device=device) + + payload = TENSORFLOW.export_dlpack(data) + + assert TENSORFLOW.exports_dlpack_data(data) is exportable + if exportable: + assert (payload.source, payload.capsule) == (data, capsule) + else: + assert payload is None + + def test_tensorflow_without_an_exporter_exports_nothing(self, monkeypatch): + module = SimpleNamespace( + __version__="2.15.0", experimental=SimpleNamespace(dlpack=object()) + ) + monkeypatch.setitem(sys.modules, "tensorflow", module) + + assert not TENSORFLOW.exports_dlpack_data(SimpleNamespace(device="GPU:0")) + + @pytest.mark.parametrize( + ("library", "importer_path", "received"), + [ + (TENSORFLOW, ("experimental", "dlpack", "from_dlpack"), "capsule"), + (TORCH, ("from_dlpack",), "capsule"), + (CUPY, ("from_dlpack",), "source"), + ], + ids=str, + ) + def test_importer_receives_what_the_library_reads( + self, library, importer_path, received, monkeypatch + ): + capsule = object() + source = SimpleNamespace(__dlpack__=lambda: capsule, __dlpack_device__=lambda: (2, 0)) + imported = object() + calls = [] + module = lambda value: calls.append(value) or imported # noqa: E731 + for attribute in reversed(importer_path): + module = SimpleNamespace(**{attribute: module}) + monkeypatch.setitem(sys.modules, library.import_name, module) + + assert library.from_dlpack(DLPackPayload(source, capsule)) is imported + assert calls == [{"capsule": capsule, "source": source}[received]] + + def test_bare_capsule_is_its_own_source(self, monkeypatch): + capsule = object() + calls = [] + monkeypatch.setitem(sys.modules, "torch", SimpleNamespace(from_dlpack=calls.append)) + + TORCH.from_dlpack(capsule) + + assert calls == [capsule] + + def test_cupy_reads_a_capsule_only_with_its_capsule_importer(self, monkeypatch): + payload = DLPackPayload(source=object(), capsule=object()) + imported = object() + monkeypatch.setitem(sys.modules, "cupy", SimpleNamespace()) + assert CUPY.from_dlpack(payload) is NotImplemented + + monkeypatch.setitem(sys.modules, "cupy", SimpleNamespace(fromDlpack=lambda _: imported)) + assert CUPY.from_dlpack(payload) is imported + + def test_jax_declines_a_capsule_without_its_source_array(self): + payload = DLPackPayload(source=object(), capsule=object()) + + assert JAX.from_dlpack(payload) is NotImplemented + + @pytest.mark.parametrize("exporter_name", ["__dlpack__", "to_dlpack", "toDlpack"]) + def test_protocol_export_accepts_each_exporter_name(self, exporter_name): + capsule = object() + data = SimpleNamespace(**{exporter_name: lambda: capsule}) + + payload = CUPY.export_dlpack(data) + + assert CUPY.exports_dlpack_data(data) + assert (payload.source, payload.capsule) == (data, capsule) + assert CUPY.export_dlpack(object()) is None + assert not CUPY.exports_dlpack_data(object()) + + @pytest.mark.torch + @pytest.mark.jax + def test_jax_import_copies_a_misaligned_torch_view(self): + torch = pytest.importorskip("torch") + pytest.importorskip("jax") + source = torch.arange(9, dtype=torch.int32).to(torch.uint16)[1:] + assert source.data_ptr() % 16 != 0 + + result = TORCH.convert_to(source, JAX, HOST) + + assert np.array_equal(np.asarray(result), np.arange(1, 9, dtype=np.uint16)) + + +@pytest.fixture +def cupy_module(monkeypatch): + module = FakeCupy() + monkeypatch.setitem(sys.modules, "cupy", module) + return module + + +@pytest.fixture +def torch_module(monkeypatch): + module = FakeTorch(device_count=1) + monkeypatch.setitem(sys.modules, "torch", module) + return module + + +class TestExecution: + """The route that ran is the planned one, or the next candidate when it declined.""" + + def test_dlpack_shares_the_buffer(self, cupy_module, torch_module, recorded): + source = cupy_module.array(np.arange(6, dtype=np.float32)) + assert planned(CUPY, GPU, TORCH, GPU) == "dlpack" + + result = CUPY.convert_to(source, TORCH, GPU) + + assert result.values is source.values + assert TORCH.device_of(result) == GPU + (record,) = recorded + assert record == ConversionRecord( + "cupy", "accelerator:0", "torch", "accelerator:0", "dlpack", 24, record.seconds + ) + assert record.declined_mechanisms == () + + def test_failed_dlpack_import_falls_back_to_a_host_copy( + self, cupy_module, torch_module, recorded, monkeypatch, caplog + ): + source = cupy_module.array(np.arange(6, dtype=np.float32)) + + def cannot_import(capsule): + raise RuntimeError("capsule rejected") + + monkeypatch.setattr(torch_module, "from_dlpack", cannot_import) + assert planned(CUPY, GPU, TORCH, HOST) == "dlpack" + + result = CUPY.convert_to(source, TORCH, HOST) + + assert isinstance(result, FakeTorchTensor) + assert TORCH.device_of(result) == HOST + assert result.values is not source.values + np.testing.assert_array_equal(result.values, source.values) + (record,) = recorded + assert record.mechanism == "host copy" + assert record.declined_mechanisms == ("dlpack",) + assert (record.source_library, record.source_device) == ("cupy", "accelerator:0") + assert (record.target_library, record.target_device) == ("torch", "host") + assert "capsule rejected" in caplog.text + + def test_array_without_an_exporter_falls_back_to_a_host_copy( + self, cupy_module, torch_module, recorded + ): + source = SimpleNamespace( + device=SimpleNamespace(id=0), get=lambda: np.arange(3, dtype=np.uint8) + ) + + result = CUPY.convert_to(source, TORCH, HOST) + + np.testing.assert_array_equal(result.values, [0, 1, 2]) + assert [(each.mechanism, each.declined_mechanisms) for each in recorded] == [ + ("host copy", ("dlpack",)) + ] + + def test_no_route_left_is_an_error(self, cupy_module, torch_module, monkeypatch): + monkeypatch.setattr(ConversionRoute, "candidates", lambda *pair: ()) + + with pytest.raises(LookupError, match="No conversion route from cupy to torch"): + ConversionRoute.run(cupy_module.array([1]), CUPY, TORCH, HOST) + + +class TestRecords: + values = np.arange(12, dtype=np.uint16).reshape(3, 4) + + def test_default_observer_asks_for_no_records(self, torch_module, monkeypatch): + def no_record(**fields): + pytest.fail(f"a record was built: {fields}") + + monkeypatch.setattr(conversion, "ConversionRecord", no_record) + monkeypatch.setattr(libraries, "ConversionRecord", no_record) + + assert type(ConversionObserver.current) is conversion.IgnoreConversions + assert ConversionObserver.current.wants_records is False + tensor = NUMPY.convert_to(self.values, TORCH, HOST) + np.testing.assert_array_equal(TORCH.to_numpy(tensor), self.values) + TORCH.from_numpy(self.values, HOST) + + def test_install_replaces_the_observer_until_it_is_put_back(self): + default = ConversionObserver.current + counters = ConversionCounters() + + assert counters.install() is default + assert ConversionObserver.current is counters + assert default.install() is counters + with counters.installed() as installed: + assert installed is counters is ConversionObserver.current + assert ConversionObserver.current is default + + @pytest.mark.parametrize( + ("convert", "source", "target"), + [ + (lambda: NUMPY.convert_to(TestRecords.values, TORCH, HOST), "numpy", "torch"), + (lambda: TORCH.from_numpy(TestRecords.values, HOST), "numpy", "torch"), + (lambda: TORCH.to_numpy(FakeTorchTensor(TestRecords.values)), "torch", "numpy"), + ( + lambda: TORCH.convert_to(FakeTorchTensor(TestRecords.values), NUMPY, HOST), + "torch", + "numpy", + ), + ], + ids=["convert_to", "from_numpy", "to_numpy", "convert_to numpy"], + ) + def test_each_conversion_writes_one_record( + self, convert, source, target, torch_module, recorded + ): + convert() + + (record,) = recorded + assert record == ConversionRecord( + source_library=source, + source_device="host", + target_library=target, + target_device="host", + mechanism="host copy", + nbytes=self.values.nbytes, + seconds=record.seconds, + ) + assert record.seconds >= 0 + + def test_upload_to_an_accelerator_names_the_device(self, torch_module, recorded): + TORCH.from_numpy(self.values, GPU) + + assert [(each.target_library, each.target_device) for each in recorded] == [ + ("torch", "accelerator:0") + ] + + def test_returning_the_input_writes_no_record(self, torch_module, recorded): + tensor = FakeTorchTensor(self.values) + + assert NUMPY.to_numpy(self.values) is self.values + assert NUMPY.from_numpy(self.values, HOST) is self.values + assert NUMPY.convert_to(self.values, NUMPY, HOST) is self.values + assert TORCH.convert_to(tensor, TORCH, HOST) is tensor + + assert recorded == [] + + def test_device_move_is_recorded(self, torch_module, recorded): + moved = TORCH.convert_to(FakeTorchTensor(self.values), TORCH, GPU) + + assert TORCH.device_of(moved) == GPU + assert [(each.mechanism, each.source_device, each.target_device) for each in recorded] == [ + ("device move", "host", "accelerator:0") + ] + + def test_record_is_plain_data(self): + record = ConversionRecord( + "cupy", "accelerator:0", "torch", "host", "host copy", 24, 0.5, ("dlpack",) + ) + + assert pickle.loads(pickle.dumps(record)) == record + assert ConversionRecord.from_json(json.loads(json.dumps(record.to_json()))) == record + assert hash(record) == hash(pickle.loads(pickle.dumps(record))) + + +class TestCounters: + @staticmethod + def record(source="numpy", target="torch", nbytes=10, seconds=0.25, mechanism="host copy"): + return ConversionRecord(source, "host", target, "accelerator:0", mechanism, nbytes, seconds) + + def counters(self, *records): + counters = ConversionCounters() + for record in records: + counters.record(record) + return counters + + def test_totals_are_kept_per_kind_of_conversion(self): + counters = self.counters( + self.record(), self.record(nbytes=30), self.record("cupy", mechanism="dlpack") + ) + + assert counters.rows() == [ + { + "source_library": "numpy", + "source_device": "host", + "target_library": "torch", + "target_device": "accelerator:0", + "mechanism": "host copy", + "count": 2, + "nbytes": 40, + "seconds": 0.5, + }, + { + "source_library": "cupy", + "source_device": "host", + "target_library": "torch", + "target_device": "accelerator:0", + "mechanism": "dlpack", + "count": 1, + "nbytes": 10, + "seconds": 0.25, + }, + ] + + def test_installed_counters_count_conversions(self, torch_module): + values = np.zeros(8, dtype=np.uint8) + + with ConversionCounters().installed() as counters: + for _ in range(3): + TORCH.from_numpy(values, HOST) + + (row,) = counters.rows() + assert (row["count"], row["nbytes"], row["mechanism"]) == (3, 24, "host copy") + + def test_merge_sums_counts_bytes_and_seconds(self): + parent = self.counters(self.record(), self.record("cupy", nbytes=5)) + worker = self.counters(self.record(nbytes=30, seconds=1.0), self.record("jax", nbytes=7)) + + assert parent.merge(worker) is parent + + totals = { + row["source_library"]: (row["count"], row["nbytes"], row["seconds"]) + for row in parent.rows() + } + assert totals == {"numpy": (2, 40, 1.25), "cupy": (1, 5, 0.25), "jax": (1, 7, 0.25)} + assert len(worker.rows()) == 2 + + def test_counters_are_plain_data(self): + counters = self.counters(self.record(), self.record("cupy", mechanism="dlpack")) + + assert pickle.loads(pickle.dumps(counters)) == counters + assert ConversionCounters.from_json(json.loads(json.dumps(counters.to_json()))) == counters + + def test_table_has_one_line_per_row(self): + counters = self.counters( + self.record(nbytes=2_500_000), self.record("cupy", mechanism="dlpack") + ) + + header, *lines = counters.format_table().splitlines() + + assert header.split() == ["source", "target", "mechanism", "count", "MB", "seconds"] + assert len(lines) == len(counters.rows()) + for line, row in zip(lines, counters.rows(), strict=True): + for field in ("source_library", "source_device", "target_library", "target_device"): + assert row[field] in line + assert row["mechanism"] in line + assert f"{row['count']:>7}" in line + assert f"{row['nbytes'] / 1e6:.1f}" in line + assert f"{row['seconds']:.3f}" in line diff --git a/tests/test_converters.py b/tests/test_converters.py index 78378e8..858fba3 100644 --- a/tests/test_converters.py +++ b/tests/test_converters.py @@ -1,230 +1,118 @@ -"""Tests for arraybridge.converters module.""" +"""``convert_memory`` and ``detect_memory_type``: the name-based conversion functions.""" import sys -from types import SimpleNamespace import numpy as np import pytest from arraybridge.converters import convert_memory, detect_memory_type +from arraybridge.devices import HOST, AcceleratorDevice from arraybridge.exceptions import MemoryConversionError from arraybridge.types import MemoryType +from tests.fakes import FakeTorch, FakeTorchTensor, fake_jax class TestDetectMemoryType: - """Tests for detect_memory_type function.""" - - def test_detect_numpy_array(self): - """Test detecting NumPy array.""" - arr = np.array([1, 2, 3]) - detected = detect_memory_type(arr) - assert detected == "numpy" - assert detected == MemoryType.NUMPY.value - - def test_detect_numpy_2d_array(self): - """Test detecting 2D NumPy array.""" - arr = np.array([[1, 2], [3, 4]]) - detected = detect_memory_type(arr) - assert detected == "numpy" - - def test_detect_numpy_3d_array(self): - """Test detecting 3D NumPy array.""" - arr = np.zeros((5, 10, 10)) - detected = detect_memory_type(arr) - assert detected == "numpy" + @pytest.mark.parametrize("shape", [(3,), (2, 2), (5, 10, 10), (2, 3, 4, 5)]) + def test_numpy_array_of_any_rank(self, shape): + assert detect_memory_type(np.zeros(shape)) == "numpy" == MemoryType.NUMPY.value - @pytest.mark.torch - def test_detect_torch_tensor(self, torch_available): - """Test detecting PyTorch tensor.""" - if not torch_available: - pytest.skip("PyTorch not available") - - import torch - - tensor = torch.tensor([1, 2, 3]) - detected = detect_memory_type(tensor) - assert detected == "torch" - assert detected == MemoryType.TORCH.value + def test_array_of_a_loaded_library(self, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", FakeTorch()) - def test_detect_unknown_type_raises_error(self): - """Test that unknown types raise ValueError.""" - with pytest.raises(ValueError) as exc_info: - detect_memory_type([1, 2, 3]) # Plain list + assert detect_memory_type(FakeTorchTensor([1, 2, 3])) == "torch" - assert "Unknown memory type" in str(exc_info.value) + @pytest.mark.torch + def test_installed_torch_tensor(self): + torch = pytest.importorskip("torch") - def test_detect_none_raises_error(self): - """Test that None raises ValueError.""" - with pytest.raises(ValueError) as exc_info: - detect_memory_type(None) + assert detect_memory_type(torch.tensor([1, 2, 3])) == "torch" - assert "Unknown memory type" in str(exc_info.value) + @pytest.mark.parametrize("value", [[1, 2, 3], None, "not an array"]) + def test_value_of_no_library_is_rejected(self, value): + with pytest.raises(ValueError, match="Unknown memory type"): + detect_memory_type(value) class TestConvertMemory: - """Tests for convert_memory function.""" - - def test_convert_numpy_to_numpy(self): - """Test converting NumPy to NumPy (no-op).""" - arr = np.array([1, 2, 3, 4, 5]) - result = convert_memory(arr, source_type="numpy", target_type="numpy", gpu_id=0) - - # Should return the same or equivalent array - assert isinstance(result, np.ndarray) - np.testing.assert_array_equal(result, arr) - - def test_convert_accepts_canonical_memory_type_members(self): - """Both conversion boundaries accept the owner taxonomy directly.""" - - arr = np.array([1, 2, 3, 4, 5]) - - result = convert_memory( - arr, - source_type=MemoryType.NUMPY, - target_type=MemoryType.NUMPY, - gpu_id=0, - ) - - np.testing.assert_array_equal(result, arr) - - def test_convert_preserves_data(self): - """Test that conversion preserves data values.""" - arr = np.array([[1.5, 2.5], [3.5, 4.5]], dtype=np.float32) - result = convert_memory(arr, source_type="numpy", target_type="numpy", gpu_id=0) + @pytest.mark.parametrize("dtype", [np.uint8, np.uint16, np.int32, np.float32, np.float64]) + @pytest.mark.parametrize("shape", [(10,), (10, 10), (5, 10, 10), (2, 3, 4, 5)]) + def test_numpy_to_numpy_returns_the_array(self, dtype, shape): + values = np.random.rand(*shape).astype(dtype) + + assert convert_memory(values, "numpy", "numpy", HOST) is values + assert convert_memory(values, MemoryType.NUMPY, MemoryType.NUMPY, HOST) is values + + @pytest.mark.parametrize( + ("source_type", "target_type", "role"), + [("invalid_type", "numpy", "source_type"), ("numpy", "invalid_type", "target_type")], + ) + def test_unregistered_name_is_rejected(self, source_type, target_type, role): + with pytest.raises(ValueError, match=f"Invalid {role} 'invalid_type'"): + convert_memory(np.array([1, 2, 3]), source_type, target_type, HOST) + + def test_device_the_target_lacks_is_a_conversion_error(self, monkeypatch): + monkeypatch.setitem(sys.modules, "jax", fake_jax("cpu")) + + with pytest.raises( + MemoryConversionError, match="JAX device accelerator:0 is unavailable" + ) as raised: + convert_memory( + np.array([1, 2, 3]), MemoryType.NUMPY, MemoryType.JAX, AcceleratorDevice(0) + ) - np.testing.assert_array_almost_equal(result, arr) + assert (raised.value.source_type, raised.value.target_type) == ("numpy", "jax") + assert isinstance(raised.value.__cause__, ValueError) - def test_convert_invalid_source_type(self): - """Test that invalid source type raises ValueError.""" - arr = np.array([1, 2, 3]) - with pytest.raises(ValueError): - convert_memory(arr, source_type="invalid_type", target_type="numpy", gpu_id=0) + def test_numpy_has_no_accelerator(self): + with pytest.raises(MemoryConversionError, match="NumPy device accelerator:0"): + convert_memory(np.array([1, 2, 3]), "numpy", "numpy", AcceleratorDevice(0)) - def test_convert_invalid_target_type(self): - """Test that invalid target type raises error.""" - arr = np.array([1, 2, 3]) - with pytest.raises(ValueError): - convert_memory(arr, source_type="numpy", target_type="invalid_type", gpu_id=0) + def test_round_trip_through_a_loaded_library(self, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", FakeTorch()) + original = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) - def test_convert_to_unavailable_declared_gpu_is_a_conversion_error(self, monkeypatch): - cpu_device = SimpleNamespace(platform="cpu") - monkeypatch.setitem(sys.modules, "jax", SimpleNamespace(devices=lambda: [cpu_device])) + tensor = convert_memory(original, "numpy", "torch", HOST) + result = convert_memory(tensor, detect_memory_type(tensor), "numpy", HOST) - with pytest.raises(MemoryConversionError, match="JAX device 0 is unavailable"): - convert_memory( - np.array([1, 2, 3]), - source_type=MemoryType.NUMPY, - target_type=MemoryType.JAX, - gpu_id=0, - ) + assert isinstance(tensor, FakeTorchTensor) + assert result.dtype == original.dtype + np.testing.assert_array_equal(result, original) @pytest.mark.torch - def test_convert_numpy_to_torch(self, torch_available): - """Test converting NumPy to PyTorch.""" - if not torch_available: - pytest.skip("PyTorch not available") + def test_round_trip_through_installed_torch_in_main_memory(self): + torch = pytest.importorskip("torch") + original = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) - import torch + tensor = convert_memory(original, "numpy", "torch", HOST) + result = convert_memory(tensor, "torch", "numpy", HOST) - arr = np.array([1.0, 2.0, 3.0], dtype=np.float32) + assert isinstance(tensor, torch.Tensor) and not tensor.is_cuda + np.testing.assert_array_equal(result, original) if not torch.cuda.is_available(): - with pytest.raises(MemoryConversionError, match="device 0 is unavailable"): - convert_memory(arr, source_type="numpy", target_type="torch", gpu_id=0) - return - result = convert_memory(arr, source_type="numpy", target_type="torch", gpu_id=0) - - assert isinstance(result, torch.Tensor) - np.testing.assert_array_almost_equal(result.cpu().numpy(), arr) + with pytest.raises(MemoryConversionError, match="accelerator:0 is unavailable"): + convert_memory(original, "numpy", "torch", AcceleratorDevice(0)) @pytest.mark.torch - def test_convert_torch_to_numpy(self, torch_available): - """Test converting PyTorch to NumPy.""" - if not torch_available: - pytest.skip("PyTorch not available") - - import torch + @pytest.mark.gpu + def test_round_trip_through_installed_torch_on_a_gpu(self): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("PyTorch CUDA device not available") + original = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) - tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) - result = convert_memory(tensor, source_type="torch", target_type="numpy", gpu_id=0) + tensor = convert_memory(original, "numpy", "torch", AcceleratorDevice(0)) + result = convert_memory(tensor, "torch", "numpy", HOST) - assert isinstance(result, np.ndarray) - np.testing.assert_array_almost_equal(result, tensor.cpu().numpy()) + assert tensor.is_cuda + np.testing.assert_array_equal(result, original) @pytest.mark.jax - def test_convert_jax_to_numpy(self, jax_available): - """JAX conversion resolves NumPy through the converter environment.""" - if not jax_available: - pytest.skip("JAX not available") - - import jax.numpy as jnp - + def test_installed_jax_to_numpy(self): + jnp = pytest.importorskip("jax.numpy") array = jnp.asarray([1.0, 2.0, 3.0], dtype=jnp.float32) - result = convert_memory( - array, - source_type=MemoryType.JAX, - target_type=MemoryType.NUMPY, - gpu_id=0, - ) + + result = convert_memory(array, MemoryType.JAX, MemoryType.NUMPY, HOST) assert isinstance(result, np.ndarray) np.testing.assert_array_equal(result, np.asarray(array)) - - def test_convert_2d_arrays(self): - """Test converting 2D arrays.""" - arr = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32) - result = convert_memory(arr, source_type="numpy", target_type="numpy", gpu_id=0) - - assert result.shape == arr.shape - np.testing.assert_array_equal(result, arr) - - def test_convert_3d_arrays(self): - """Test converting 3D arrays.""" - arr = np.random.rand(5, 10, 10).astype(np.float32) - result = convert_memory(arr, source_type="numpy", target_type="numpy", gpu_id=0) - - assert result.shape == arr.shape - np.testing.assert_array_almost_equal(result, arr) - - def test_convert_different_dtypes(self): - """Test converting arrays with different dtypes.""" - for dtype in [np.uint8, np.uint16, np.int32, np.float32, np.float64]: - arr = np.array([1, 2, 3], dtype=dtype) - result = convert_memory(arr, source_type="numpy", target_type="numpy", gpu_id=0) - assert result.dtype == dtype - - -class TestConversionIntegration: - """Integration tests for memory conversion.""" - - def test_detect_and_convert_workflow(self): - """Test detecting type and then converting.""" - arr = np.array([1, 2, 3, 4, 5]) - - # Detect the type - detected_type = detect_memory_type(arr) - assert detected_type == "numpy" - - # Convert using detected type - result = convert_memory(arr, source_type=detected_type, target_type="numpy", gpu_id=0) - np.testing.assert_array_equal(result, arr) - - @pytest.mark.torch - def test_round_trip_conversion_numpy_torch(self, torch_available): - """Test round-trip conversion: NumPy -> PyTorch -> NumPy.""" - if not torch_available: - pytest.skip("PyTorch not available") - - import torch - - if not torch.cuda.is_available(): - pytest.skip("PyTorch CUDA device not available") - - original = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) - - # NumPy -> PyTorch - torch_tensor = convert_memory(original, source_type="numpy", target_type="torch", gpu_id=0) - - # PyTorch -> NumPy - result = convert_memory(torch_tensor, source_type="torch", target_type="numpy", gpu_id=0) - - np.testing.assert_array_almost_equal(result, original) diff --git a/tests/test_converters_registry.py b/tests/test_converters_registry.py deleted file mode 100644 index 87ad215..0000000 --- a/tests/test_converters_registry.py +++ /dev/null @@ -1,155 +0,0 @@ -"""Tests for arraybridge.converters_registry module.""" - -from types import MappingProxyType - -import pytest - - -class TestConverterRegistry: - """Tests for converter registry functionality.""" - - def test_registry_contains_all_memory_types(self): - """Test that registry contains converters for all memory types.""" - from arraybridge.converters_registry import ConverterBase - from arraybridge.types import MemoryType - - expected_types = {mt.value for mt in MemoryType} - registered_types = set(ConverterBase.__registry__.keys()) - - assert ( - expected_types == registered_types - ), f"Registry mismatch. Expected: {expected_types}, Got: {registered_types}" - - def test_get_converter_returns_valid_converter(self): - """Test that get_converter returns a valid converter instance.""" - from arraybridge.converters_registry import get_converter - - converter = get_converter("numpy") - assert converter is not None - assert hasattr(converter, "to_numpy") - assert hasattr(converter, "from_numpy") - assert hasattr(converter, "from_dlpack") - assert hasattr(converter, "move_to_device") - - def test_get_converter_for_all_types(self): - """Test that get_converter works for all memory types.""" - from arraybridge.converters_registry import get_converter - from arraybridge.types import MemoryType - - for mem_type in MemoryType: - converter = get_converter(mem_type.value) - assert converter is not None - assert converter.memory_type == mem_type.value - - def test_get_converter_invalid_type_raises_error(self): - """Test that get_converter raises ValueError for invalid types.""" - from arraybridge.converters_registry import get_converter - - with pytest.raises(ValueError) as exc_info: - get_converter("invalid_type") - - assert "No converter registered" in str(exc_info.value) - assert "invalid_type" in str(exc_info.value) - - def test_registry_validation_passes(self): - """Test that registry validation passes with current setup.""" - from arraybridge.converters_registry import _validate_registry - - # Should not raise any exception - _validate_registry() - - def test_registry_validation_fails_on_missing_type(self, monkeypatch): - """Test that registry validation fails if a memory type is missing.""" - from arraybridge.converters_registry import ConverterBase, _validate_registry - - removed_type = "numpy" - incomplete = { - name: converter - for name, converter in ConverterBase.__registry__.items() - if name != removed_type - } - monkeypatch.setattr(ConverterBase, "__registry__", MappingProxyType(incomplete)) - - with pytest.raises(RuntimeError) as exc_info: - _validate_registry() - assert "Missing" in str(exc_info.value) - assert removed_type in str(exc_info.value) - - def test_registry_validation_fails_on_extra_type(self, monkeypatch): - """Test that registry validation fails if there's an extra type.""" - from arraybridge.converters_registry import ConverterBase, _validate_registry - - extended = dict(ConverterBase.__registry__) - extended["extra_type"] = type("ExtraConverter", (), {}) - monkeypatch.setattr(ConverterBase, "__registry__", MappingProxyType(extended)) - - with pytest.raises(RuntimeError) as exc_info: - _validate_registry() - assert "Extra" in str(exc_info.value) - assert "extra_type" in str(exc_info.value) - - def test_registry_projection_is_read_only(self): - from arraybridge.converters_registry import ConverterBase - - assert isinstance(ConverterBase.__registry__, MappingProxyType) - with pytest.raises(TypeError): - ConverterBase.__registry__["extra_type"] = type("ExtraConverter", (), {}) - - def test_converter_has_to_x_methods(self): - """Test that converters have to_X() methods for all memory types.""" - from arraybridge.converters_registry import get_converter - from arraybridge.types import MemoryType - - numpy_converter = get_converter("numpy") - - # Check that it has to_X() methods for all memory types - for target_type in MemoryType: - method_name = f"to_{target_type.value}" - assert hasattr(numpy_converter, method_name), f"Converter missing method: {method_name}" - - def test_converter_classes_registered_with_correct_names(self): - """Test that converter classes are registered with expected names.""" - from arraybridge.converters_registry import ConverterBase - - # Check numpy converter - numpy_class = ConverterBase.__registry__["numpy"] - assert numpy_class.__name__ == "NumpyConverter" - - # Check torch converter - torch_class = ConverterBase.__registry__["torch"] - assert torch_class.__name__ == "TorchConverter" - - def test_multiple_get_converter_calls_return_new_instances(self): - """Test that get_converter returns new instances each time.""" - from arraybridge.converters_registry import get_converter - - converter1 = get_converter("numpy") - converter2 = get_converter("numpy") - - # They should be different instances - assert converter1 is not converter2 - # But same type - assert isinstance(converter1, type(converter2)) - - -class TestMemoryTypeRegistryResolution: - """Tests that nominal memory types resolve through the registry owner.""" - - def test_memory_type_value_resolves_through_registry(self): - """Registry lookup accepts the nominal enum value.""" - from arraybridge.converters_registry import get_converter - from arraybridge.types import MemoryType - - numpy_converter = get_converter(MemoryType.NUMPY.value) - assert numpy_converter is not None - assert numpy_converter.memory_type == "numpy" - - def test_registry_resolves_all_nominal_memory_types(self): - """Every declared memory type has one registered converter owner.""" - from arraybridge.converters_registry import get_converter - from arraybridge.types import MemoryType - - for mem_type in MemoryType: - converter = get_converter(mem_type.value) - assert converter is not None - assert converter.memory_type == mem_type.value diff --git a/tests/test_decorators.py b/tests/test_decorators.py index d0442f6..7aad9ec 100644 --- a/tests/test_decorators.py +++ b/tests/test_decorators.py @@ -1,27 +1,29 @@ """Tests for arraybridge.decorators module.""" -from contextlib import nullcontext +import sys +import threading from dataclasses import dataclass from typing import Any +import dill import numpy as np import pytest from python_introspect import RuntimeParameterDeclarationABC -from arraybridge import ArrayPayload +import arraybridge +from arraybridge import ArrayPayload, decorators from arraybridge.decorators import ( DtypeConversion, DtypeConversionConfig, PreserveInputDtypeConfig, SliceBySliceRuntimeParameter, - ThreadGPUContext, memory_types, wrap_dtype_preserving_callable, ) -from arraybridge.decorators import ( - torch as torch_memory, -) +from arraybridge.devices import AcceleratorDevice +from arraybridge.libraries import ArrayLibrary from arraybridge.types import MemoryType +from tests.fakes import FakeCupy, FakeTorch class TestDtypeConversion: @@ -57,130 +59,81 @@ def test_numpy_dtype_property(self): assert DtypeConversion.NATIVE_OUTPUT.numpy_dtype is None -class TestThreadGPUContext: - def test_streams_are_cached_by_framework_local_device(self): - state = {"device": 0} - created = [] - - class DeviceScope: - def __init__(self, device_id): - self.device_id = device_id - self.previous = None - - def __enter__(self): - self.previous = state["device"] - state["device"] = self.device_id - - def __exit__(self, exc_type, exc_value, traceback): - state["device"] = self.previous - return False - - class Stream: - def __init__(self): - self.device_id = state["device"] - created.append(self) - - module = type( - "FakeCupy", - (), - { - "cuda": type( - "Cuda", - (), - { - "runtime": type( - "Runtime", - (), - { - "getDeviceCount": staticmethod(lambda: 2), - "getDevice": staticmethod(lambda: state["device"]), - }, - )(), - "Device": staticmethod(DeviceScope), - "Stream": Stream, - }, - )() - }, - )() - context = ThreadGPUContext() - - _, first = context.stream_for(MemoryType.CUPY, module) - _, repeated = context.stream_for(MemoryType.CUPY, module) - state["device"] = 1 - _, second = context.stream_for(MemoryType.CUPY, module) - state["device"] = 0 - _, first_again = context.stream_for(MemoryType.CUPY, module) - - assert first is repeated is first_again - assert second is not first - assert [stream.device_id for stream in created] == [0, 1] - - def test_framework_without_stream_leaf_does_not_invent_one(self): - module = type( - "FakeTensorFlow", - (), - {"config": type("Config", (), {"list_logical_devices": lambda self, _: []})()}, - )() - - assert ThreadGPUContext().stream_for(MemoryType.TENSORFLOW, module) == ( - None, - None, - ) +class TestExecutionOnAccelerators: + """What a library's devices need around a decorated call.""" - def test_torch_stream_execution_uses_framework_owned_scope(self, monkeypatch): - events = [] - - class Stream: - pass - - class StreamScope: - def __init__(self, stream): - self.stream = stream - - def __enter__(self): - events.append(("enter", self.stream)) - - def __exit__(self, exc_type, exc_value, traceback): - events.append(("exit", self.stream)) - return False - - stream = Stream() - module = type( - "FakeTorch", - (), - { - "cuda": type( - "Cuda", - (), - { - "is_available": staticmethod(lambda: True), - "device_count": staticmethod(lambda: 1), - "current_device": staticmethod(lambda: 0), - "device": staticmethod(lambda _device_id: nullcontext()), - "Stream": staticmethod(lambda: stream), - "stream": staticmethod(StreamScope), - }, - )() - }, - )() - monkeypatch.setattr(MemoryType, "import_if_installed", lambda self: module) - - @torch_memory + @pytest.mark.parametrize( + ("name", "fake"), [("torch", FakeTorch), ("cupy", FakeCupy)], ids=["torch", "cupy"] + ) + def test_cuda_call_runs_on_this_threads_stream(self, name, fake, monkeypatch): + module = fake(device_count=1) + monkeypatch.setitem(sys.modules, name, module) + + @getattr(decorators, name) def identity(value): - events.append(("call", value)) + module.events.append(("call", value)) return value assert identity("payload") == "payload" - assert events == [ - ("enter", stream), + assert identity("again") == "again" + stream = ArrayLibrary.named(name).thread_stream(AcceleratorDevice(0)) + assert module.events == [ + ("enter", 0), + ("exit", 0), + ("stream enter", stream), ("call", "payload"), - ("exit", stream), + ("stream exit", stream), + ("stream enter", stream), + ("call", "again"), + ("stream exit", stream), ] - def test_cross_framework_dtype_preservation_uses_portable_identity( - self, - monkeypatch, + def test_torch_without_cuda_calls_the_function_directly(self, monkeypatch): + module = FakeTorch() + monkeypatch.setitem(sys.modules, "torch", module) + + @decorators.torch + def identity(value): + return value + + assert identity("payload") == "payload" + assert module.events == [] + + @pytest.mark.parametrize("name", [library.name for library in ArrayLibrary.all()]) + def test_library_that_is_not_installed_calls_the_function_directly( + self, name, absent_libraries ): + @getattr(decorators, name) + def add_one(values): + return values + 1 + + np.testing.assert_array_equal(add_one(np.array([1, 2])), [2, 3]) + assert add_one.execution_memory_type == name + + def test_out_of_memory_releases_the_cache_and_retries(self, monkeypatch): + module = FakeCupy(device_count=1) + monkeypatch.setitem(sys.modules, "cupy", module) + calls = [] + + def fails_once(value): + calls.append(value) + if len(calls) == 1: + raise MemoryError("out of memory") + return value + + assert decorators.cupy(fails_once)("payload") == "payload" + assert calls == ["payload", "payload"] + assert [event for event in module.events if event[0] == "free"] == [ + ("free", "regular"), + ("free", "pinned"), + ] + + calls.clear() + with pytest.raises(MemoryError, match="out of memory"): + decorators.cupy(fails_once, oom_recovery=False)("payload") + assert calls == ["payload"] + + def test_dtype_is_restored_by_name_across_libraries(self, monkeypatch): class TorchDtype: def __str__(self): return "torch.uint16" @@ -191,12 +144,11 @@ def __init__(self, dtype): converted = [] - def scale_dtype(memory_type, array, target_dtype, module=None): - del memory_type, module - converted.append(target_dtype) + def scale_dtype(library, array, target_dtype): + converted.append((library.name, target_dtype)) return Array(np.dtype(target_dtype)) - monkeypatch.setattr(MemoryType, "scale_dtype", scale_dtype) + monkeypatch.setattr(ArrayLibrary, "scale_dtype", scale_dtype) @memory_types("torch", "jax") def convert(_image): @@ -205,9 +157,33 @@ def convert(_image): wrapped = wrap_dtype_preserving_callable(convert, MemoryType.JAX) result = wrapped(Array(TorchDtype())) - assert converted == ["uint16"] + assert converted == [("jax", "uint16")] assert result.dtype == np.dtype("uint16") + @pytest.mark.parametrize("name", ["numpy", "cupy"]) + def test_decorated_function_pickles_without_this_threads_streams(self, name, monkeypatch): + # A function with no importable name, as kept by an undo snapshot: dill + # stores its body, and the stream of the thread that made it stays behind. + module = FakeCupy(device_count=1) + monkeypatch.setitem(sys.modules, "cupy", module) + stream = ArrayLibrary.named("cupy").thread_stream(AcceleratorDevice(0)) + stream.unpicklable = threading.Lock() + + @getattr(decorators, name) + def declared(image, scale=3): + return image * scale + + restored = dill.loads(dill.dumps(declared)) + + assert restored is not declared + assert restored.__name__ == declared.__name__ + assert restored.__wrapped__.__name__ == declared.__wrapped__.__name__ + assert restored.execution_memory_type == name + np.testing.assert_array_equal( + restored(np.array([1, 2], dtype=np.uint16), scale=4), [4, 8] + ) + assert ArrayLibrary.named("cupy").thread_stream(AcceleratorDevice(0)) is stream + class TestMemoryTypesDecorator: """Tests for memory_types decorator.""" @@ -287,9 +263,6 @@ def test_slice_stacking_uses_declared_output_framework(self): torch = pytest.importorskip("torch") from arraybridge.decorators import numpy - if not MemoryType.TORCH.available_device_ids(torch): - pytest.skip("PyTorch GPU stacking requires an available device") - @numpy(output_type="torch", slice_by_slice_default=True) def planes_to_torch(value): return torch.as_tensor(value) @@ -301,15 +274,34 @@ def planes_to_torch(value): assert result.dtype is torch.uint16 assert tuple(result.shape) == source.shape - def test_decorator_exports_are_declaration_derived(self): - import arraybridge - from arraybridge import decorators + def test_every_registered_library_has_a_decorator(self): + names = {library.name for library in ArrayLibrary.all()} + + assert names <= set(dir(decorators)) + assert names <= set(arraybridge.__all__) + for name in names: + decorator = getattr(decorators, name) + assert decorator is getattr(arraybridge, name) + assert decorator.__name__ == name + assert ArrayLibrary.named(name).display_name in decorator.__doc__ + with pytest.raises(AttributeError): + decorators.not_a_library + + @pytest.mark.torch + def test_torch_function_runs_in_main_memory_without_cuda(self): + torch = pytest.importorskip("torch") + + @arraybridge.torch + def double(tensor): + return tensor * 2 - expected = {memory_type.value for memory_type in MemoryType} + source = torch.arange(6, dtype=torch.int32).reshape(2, 3) + result = double(source) - assert expected <= set(decorators.__all__) - assert expected <= set(arraybridge.__all__) - assert all(callable(getattr(arraybridge, name)) for name in expected) + assert isinstance(result, torch.Tensor) and not result.is_cuda + assert result.dtype is torch.int32 + assert double.execution_memory_type == "torch" + np.testing.assert_array_equal(result.numpy(), np.arange(6).reshape(2, 3) * 2) def test_memory_types_preserves_function_metadata(self): """Test that memory_types preserves function name, docstring, etc.""" @@ -328,12 +320,6 @@ def test_func(x, y=10): class TestFrameworkDecorators: """Tests for auto-generated framework-specific decorators.""" - def test_numpy_decorator_exists(self): - """Test that numpy decorator is available.""" - from arraybridge.decorators import numpy - - assert callable(numpy) - def test_numpy_decorator_basic(self): """Test basic numpy decorator functionality.""" from arraybridge.decorators import numpy @@ -442,36 +428,6 @@ def unchanged(_image): assert unchanged(np.zeros(3, dtype=np.float32)) is unchanged_payload - def test_cupy_decorator_exists(self): - """Test that cupy decorator is available.""" - from arraybridge.decorators import cupy - - assert callable(cupy) - - def test_torch_decorator_exists(self): - """Test that torch decorator is available.""" - from arraybridge.decorators import torch - - assert callable(torch) - - def test_tensorflow_decorator_exists(self): - """Test that tensorflow decorator is available.""" - from arraybridge.decorators import tensorflow - - assert callable(tensorflow) - - def test_jax_decorator_exists(self): - """Test that jax decorator is available.""" - from arraybridge.decorators import jax - - assert callable(jax) - - def test_pyclesperanto_decorator_exists(self): - """Test that pyclesperanto decorator is available.""" - from arraybridge.decorators import pyclesperanto - - assert callable(pyclesperanto) - class TestDecoratorParameters: """Tests for decorator parameter handling.""" diff --git a/tests/test_devices.py b/tests/test_devices.py new file mode 100644 index 0000000..622b2a6 --- /dev/null +++ b/tests/test_devices.py @@ -0,0 +1,264 @@ +"""Devices are values, and each library selects, reports and moves between its own.""" + +import importlib +import pickle +import sys +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +import numpy as np +import pytest + +from arraybridge.devices import HOST, AcceleratorDevice, Device, HostDevice, NamedDevice +from arraybridge.libraries import ArrayLibrary, HostDevices +from tests.fakes import ( + FakeCupy, + FakePyclesperanto, + FakePyclesperantoArray, + FakeTorch, + FakeTorchTensor, + fake_jax, +) + +LIBRARIES = ArrayLibrary.all() +NUMPY = ArrayLibrary.named("numpy") +TORCH = ArrayLibrary.named("torch") +JAX = ArrayLibrary.named("jax") +PYCLESPERANTO = ArrayLibrary.named("pyclesperanto") + + +class TestDeviceValues: + @pytest.mark.parametrize( + ("device", "equal", "label"), + [ + (HOST, HostDevice(), "host"), + (AcceleratorDevice(1), AcceleratorDevice(1), "accelerator:1"), + (NamedDevice("cpu0"), NamedDevice("cpu0"), "device:cpu0"), + ], + ) + def test_leaves_are_equal_by_value_hashable_and_picklable(self, device, equal, label): + assert isinstance(device, Device) + assert device == equal + assert device is not equal + assert hash(device) == hash(equal) + assert {device: "found"}[equal] == "found" + assert pickle.loads(pickle.dumps(device)) == device + assert device.label == label + + def test_different_devices_are_not_equal(self): + devices = [HOST, AcceleratorDevice(0), AcceleratorDevice(1), NamedDevice("0")] + + assert len(set(devices)) == len(devices) + + def test_a_device_cannot_be_changed(self): + with pytest.raises(AttributeError): + AcceleratorDevice(0).index = 1 + + +class TestEveryLibrary: + @pytest.mark.parametrize("library", LIBRARIES, ids=str) + def test_without_its_package_only_main_memory_is_available(self, library, absent_libraries): + host = (HOST,) if isinstance(library, HostDevices) else () + + assert library.host_devices() == host + assert library.accelerator_devices() == () + assert library.available_devices() == host + for device in host: + assert library.require_device(device) is None + for missing in (AcceleratorDevice(0), NamedDevice("cpu0"), *({HOST} - set(host))): + with pytest.raises(ValueError, match=f"device {missing.label} is unavailable"): + library.require_device(missing) + with pytest.raises(ValueError, match="is unavailable"): + library.device_scope(missing) + + +class TestNumpy: + def test_arrays_live_in_main_memory(self): + values = np.array([1, 2, 3]) + + assert NUMPY.device_of(values) == HOST + assert NUMPY.move_to_device(values, HOST) is values + assert NUMPY.activate_device(HOST) is None + with pytest.raises(ValueError, match="NumPy device accelerator:0 is unavailable"): + NUMPY.move_to_device(values, AcceleratorDevice(0)) + + def test_host_scope_does_not_import_numpy(self, monkeypatch): + monkeypatch.setattr( + importlib, "import_module", lambda name: pytest.fail(f"imported {name}") + ) + + with NUMPY.device_scope(HOST): + pass + + +class TestTorch: + def test_device_of_reads_the_tensor(self, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", FakeTorch()) + + assert TORCH.device_of(FakeTorchTensor([1, 2])) == HOST + assert TORCH.device_of(FakeTorchTensor([1, 2], cuda_index=1)) == AcceleratorDevice(1) + with pytest.raises(AttributeError): + TORCH.device_of(SimpleNamespace()) + + def test_cuda_devices_are_numbered_as_torch_numbers_them(self, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", FakeTorch(device_count=2)) + + assert TORCH.accelerator_devices() == (AcceleratorDevice(0), AcceleratorDevice(1)) + assert TORCH.available_devices() == (HOST, AcceleratorDevice(0), AcceleratorDevice(1)) + for index in (-1, 2): + with pytest.raises(ValueError, match=f"device accelerator:{index} is unavailable"): + TORCH.require_device(AcceleratorDevice(index)) + + @pytest.mark.parametrize("index", [0, 1, 2]) + def test_activate_selects_the_device_for_the_process(self, index, monkeypatch): + module = FakeTorch(device_count=3) + monkeypatch.setitem(sys.modules, "torch", module) + + TORCH.activate_device(AcceleratorDevice(index)) + + assert module.events == [("set_device", index)] + + def test_move_copies_only_when_the_device_differs(self, monkeypatch): + module = FakeTorch(device_count=2) + monkeypatch.setitem(sys.modules, "torch", module) + on_first = FakeTorchTensor([1, 2], cuda_index=0) + + on_second = TORCH.move_to_device(on_first, AcceleratorDevice(1)) + on_host = TORCH.move_to_device(on_second, HOST) + + assert TORCH.move_to_device(on_first, AcceleratorDevice(0)) is on_first + assert TORCH.device_of(on_second) == AcceleratorDevice(1) + assert TORCH.device_of(on_host) == HOST + assert module.events == [("enter", 1), ("exit", 1)] + + @pytest.mark.torch + def test_installed_torch_works_in_main_memory(self): + torch = pytest.importorskip("torch") + values = np.arange(6, dtype=np.float32).reshape(2, 3) + + tensor = TORCH.from_numpy(values, HOST) + + assert isinstance(tensor, torch.Tensor) and not tensor.is_cuda + assert TORCH.device_of(tensor) == HOST + assert TORCH.move_to_device(tensor, HOST) is tensor + np.testing.assert_array_equal(TORCH.to_numpy(tensor), values) + assert TORCH.available_devices()[0] == HOST + if not torch.cuda.is_available(): + assert TORCH.accelerator_devices() == () + assert TORCH.current_device() is None + with pytest.raises(ValueError, match="accelerator:0 is unavailable"): + TORCH.require_device(AcceleratorDevice(0)) + + +class TestJax: + def test_cpu_install_has_no_accelerator(self, monkeypatch): + monkeypatch.setitem(sys.modules, "jax", fake_jax("cpu")) + + assert JAX.accelerator_devices() == () + assert JAX.available_devices() == (HOST,) + with pytest.raises(ValueError, match=r"available devices are \('host',\)"): + JAX.device_scope(AcceleratorDevice(0)) + + @pytest.mark.parametrize("index", [-1, 2]) + def test_scope_rejects_an_accelerator_it_does_not_list(self, index, monkeypatch): + monkeypatch.setitem(sys.modules, "jax", fake_jax("gpu", "gpu")) + + with pytest.raises(ValueError, match=f"device accelerator:{index} is unavailable"): + JAX.device_scope(AcceleratorDevice(index)) + + def test_accelerators_are_numbered_among_the_gpu_devices(self, monkeypatch): + module = fake_jax("cpu", "gpu", "gpu") + monkeypatch.setitem(sys.modules, "jax", module) + cpu, _first_gpu, second_gpu = module.devices() + + assert JAX.accelerator_devices() == (AcceleratorDevice(0), AcceleratorDevice(1)) + assert JAX.device_of(SimpleNamespace(device=second_gpu)) == AcceleratorDevice(1) + assert JAX.device_of(SimpleNamespace(device=cpu)) == HOST + assert JAX.device_of(SimpleNamespace(device=lambda: second_gpu)) == AcceleratorDevice(1) + + +class TestPyclesperanto: + def test_move_restores_the_selected_gpu(self, monkeypatch): + module = FakePyclesperanto(current="gpu0") + monkeypatch.setitem(sys.modules, "pyclesperanto", module) + source = object() + + result = PYCLESPERANTO.move_to_device(source, AcceleratorDevice(1)) + + assert result == ("copy", source) + assert module.current == "gpu0" + assert module.events == [("copy", source, result)] + + def test_scope_restores_a_device_that_is_not_a_listed_gpu(self, monkeypatch): + module = FakePyclesperanto(current="cpu0") + monkeypatch.setitem(sys.modules, "pyclesperanto", module) + + with PYCLESPERANTO.device_scope(AcceleratorDevice(1)): + assert module.current == "gpu1" + + assert module.current == "cpu0" + + def test_array_on_an_unlisted_device_is_not_taken_for_gpu_zero(self, monkeypatch): + module = FakePyclesperanto(current="cpu0") + monkeypatch.setitem(sys.modules, "pyclesperanto", module) + source = object() + + result = PYCLESPERANTO.move_to_device(source, AcceleratorDevice(0)) + + assert result == ("copy", source) + assert module.current == "cpu0" + assert module.events == [("copy", source, result)] + + @pytest.mark.parametrize( + ("current", "array_device", "expected"), + [ + (SimpleNamespace(id=None, name="gpu1"), None, AcceleratorDevice(1)), + ("gpu0", "gpu1", AcceleratorDevice(1)), + (SimpleNamespace(id=0, name="cpu0"), None, NamedDevice("cpu0")), + ("gpu0", "cpu0", NamedDevice("cpu0")), + ], + ) + def test_device_of_goes_by_the_device_name(self, current, array_device, expected, monkeypatch): + monkeypatch.setitem(sys.modules, "pyclesperanto", FakePyclesperanto(current=current)) + data = object() if array_device is None else FakePyclesperantoArray([1], array_device) + + assert PYCLESPERANTO.device_of(data) == expected + + +class TestCudaStreams: + def test_only_cuda_libraries_have_streams(self): + assert {library.name for library in LIBRARIES if hasattr(library, "thread_stream")} == { + "cupy", + "torch", + } + + def test_streams_are_kept_per_thread_and_device(self, monkeypatch): + module = FakeCupy() + monkeypatch.setitem(sys.modules, "cupy", module) + cupy = ArrayLibrary.named("cupy") + + first = cupy.thread_stream(AcceleratorDevice(0)) + second = cupy.thread_stream(AcceleratorDevice(1)) + with ThreadPoolExecutor(max_workers=1) as executor: + other_thread, repeated = executor.submit( + lambda: ( + cupy.thread_stream(AcceleratorDevice(0)), + cupy.thread_stream(AcceleratorDevice(0)), + ) + ).result() + + assert cupy.thread_stream(AcceleratorDevice(0)) is first + assert cupy.thread_stream(AcceleratorDevice(1)) is second + assert (first.device_index, second.device_index) == (0, 1) + assert other_thread is repeated + assert other_thread is not first + assert module.current == 0 + + def test_current_device_is_the_one_the_library_has_selected(self, monkeypatch): + cupy_module = FakeCupy() + cupy_module.current = 1 + monkeypatch.setitem(sys.modules, "cupy", cupy_module) + monkeypatch.setitem(sys.modules, "torch", FakeTorch()) + + assert ArrayLibrary.named("cupy").current_device() == AcceleratorDevice(1) + assert TORCH.current_device() is None diff --git a/tests/test_dtype_scaling.py b/tests/test_dtype_scaling.py deleted file mode 100644 index 9c76a06..0000000 --- a/tests/test_dtype_scaling.py +++ /dev/null @@ -1,357 +0,0 @@ -"""Tests for arraybridge.dtype_scaling module.""" - -import numpy as np -import pytest - -from arraybridge.dtype_scaling import SCALING_FUNCTIONS -from arraybridge.types import MemoryType - - -class TestScalingRanges: - """Tests for scaling range constants.""" - - def test_scaling_ranges_uint8(self): - """Test uint8 scaling range.""" - from arraybridge.dtype_scaling import _SCALING_RANGES - - assert _SCALING_RANGES["uint8"] == 255.0 - - def test_scaling_ranges_uint16(self): - """Test uint16 scaling range.""" - from arraybridge.dtype_scaling import _SCALING_RANGES - - assert _SCALING_RANGES["uint16"] == 65535.0 - - def test_scaling_ranges_int16(self): - """Test int16 scaling range (tuple format).""" - from arraybridge.dtype_scaling import _SCALING_RANGES - - scale_val, offset_val = _SCALING_RANGES["int16"] - assert scale_val == 65535.0 - assert offset_val == 32768.0 - - -class TestScalingFunctions: - """Tests for scaling functions.""" - - def test_scaling_functions_registry(self): - """Test that all memory types have scaling functions.""" - for mem_type in MemoryType: - assert mem_type.value in SCALING_FUNCTIONS - assert callable(SCALING_FUNCTIONS[mem_type.value]) - - def test_numpy_scaling_no_conversion_needed(self): - """Test numpy scaling when no conversion is needed.""" - scale_func = SCALING_FUNCTIONS["numpy"] - - # int to int - no scaling needed - arr = np.array([1, 2, 3], dtype=np.uint8) - result = scale_func(arr, np.uint16) - assert result.dtype == np.uint16 - np.testing.assert_array_equal(result, [1, 2, 3]) - - def test_numpy_scaling_float_to_int(self): - """Test numpy scaling from float to int.""" - scale_func = SCALING_FUNCTIONS["numpy"] - - # float64 [0, 1] to uint8 [0, 255] - arr = np.array([0.0, 0.5, 1.0], dtype=np.float64) - result = scale_func(arr, np.uint8) - assert result.dtype == np.uint8 - np.testing.assert_array_equal(result, [0, 127, 255]) - - def test_numpy_scaling_float_to_int16(self): - """Test numpy scaling from float to int16.""" - scale_func = SCALING_FUNCTIONS["numpy"] - - # float32 [0, 1] to uint16 [0, 65535] - arr = np.array([0.0, 0.5, 1.0], dtype=np.float32) - result = scale_func(arr, np.uint16) - assert result.dtype == np.uint16 - np.testing.assert_array_equal(result, [0, 32767, 65535]) - - def test_numpy_scaling_int16_range(self): - """Test numpy scaling to int16 range.""" - scale_func = SCALING_FUNCTIONS["numpy"] - - # float64 [0, 1] to int16 [-32768, 32767] - arr = np.array([0.0, 0.5, 1.0], dtype=np.float64) - result = scale_func(arr, np.int16) - assert result.dtype == np.int16 - # Check that values are in expected range - assert np.all(result >= -32768) - assert np.all(result <= 32767) - - def test_numpy_scaling_constant_image(self): - """Test numpy scaling with constant image.""" - scale_func = SCALING_FUNCTIONS["numpy"] - - # Constant float image should convert without error - arr = np.full((10, 10), 0.5, dtype=np.float32) - result = scale_func(arr, np.uint8) - assert result.dtype == np.uint8 - # For constant images, all values should be the same - unique_vals = np.unique(result) - assert len(unique_vals) == 1 # All values should be identical - - def test_numpy_scaling_edge_cases(self): - """Test numpy scaling edge cases.""" - scale_func = SCALING_FUNCTIONS["numpy"] - - # Test with very small range - arr = np.array([0.499, 0.501], dtype=np.float64) - result = scale_func(arr, np.uint8) - assert result.dtype == np.uint8 - - # Test with single value - arr = np.array([0.7], dtype=np.float32) - result = scale_func(arr, np.uint16) - assert result.dtype == np.uint16 - - @pytest.mark.skipif(not hasattr(np, "float16"), reason="float16 not available") - def test_numpy_scaling_float16(self): - """Test numpy scaling with float16.""" - scale_func = SCALING_FUNCTIONS["numpy"] - - arr = np.array([0.0, 1.0], dtype=np.float16) - result = scale_func(arr, np.uint8) - assert result.dtype == np.uint8 - - @pytest.mark.parametrize( - "memory_type", - [MemoryType.TORCH, MemoryType.CUPY, MemoryType.PYCLESPERANTO], - ) - def test_scaling_unavailable_framework_returns_input(self, memory_type, monkeypatch): - monkeypatch.setattr(MemoryType, "import_if_installed", lambda self: None) - arr = np.array([1, 2, 3]) - - result = SCALING_FUNCTIONS[memory_type.value](arr, np.float32) - - assert result is arr - - def test_scaling_non_array_input(self): - """Test scaling with non-array input.""" - scale_func = SCALING_FUNCTIONS["numpy"] - - # Should return input unchanged - result = scale_func("not an array", np.uint8) - assert result == "not an array" - - def test_scaling_empty_array(self): - """Test scaling with empty array.""" - scale_func = SCALING_FUNCTIONS["numpy"] - - arr = np.array([], dtype=np.float32) - # Empty arrays may not be handled by the scaling function due to min/max operations - # This is acceptable as empty arrays are edge cases - try: - result = scale_func(arr, np.uint8) - assert result.dtype == np.uint8 - assert result.size == 0 - except ValueError: - # Expected for empty arrays due to min/max operations - pytest.skip("Empty arrays not supported by scaling function (expected)") - - def test_generic_scaling_delegates_to_member_leaf(self): - from arraybridge.dtype_scaling import _scale_generic - - values = np.array([0.0, 1.0], dtype=np.float32) - - result = _scale_generic(values, np.uint8, MemoryType.NUMPY) - - np.testing.assert_array_equal(result, [0, 255]) - - def test_scaling_ranges_comprehensive(self): - """Test all scaling ranges are properly defined.""" - from arraybridge.dtype_scaling import _SCALING_RANGES - - # Test all expected dtypes - expected_ranges = { - "uint8": 255.0, - "uint16": 65535.0, - "uint32": 4294967295.0, - "int16": (65535.0, 32768.0), - "int32": (4294967295.0, 2147483648.0), - } - - for dtype_name, expected_range in expected_ranges.items(): - assert dtype_name in _SCALING_RANGES - assert _SCALING_RANGES[dtype_name] == expected_range - - def test_torch_scaling_with_gpu_array(self): - """Test torch scaling with actual GPU array (when torch available).""" - torch = pytest.importorskip("torch") - scale_func = SCALING_FUNCTIONS["torch"] - - # Create torch tensor (float to int conversion should trigger scaling) - arr = torch.tensor([0.0, 0.5, 1.0], dtype=torch.float32) - result = scale_func(arr, np.int32) # Use numpy dtype for scaling - - # Should return torch tensor - assert isinstance(result, torch.Tensor) - # Check that scaling occurred (should be scaled to int32 range) - assert result.dtype == torch.int32 - # With clamping fix, values should be correctly scaled to int32 range - # 0.0 -> INT32_MIN, 0.5 -> ~0, 1.0 -> close to INT32_MAX - int32_min = -(2**31) - int32_max = 2**31 - 1 - # Due to float32 precision limits, we clamp to INT32_MAX - 128 to avoid overflow - # Allow tolerance of up to 150 from the bounds - assert int32_min <= result[0].item() <= int32_min + 150 - assert abs(result[1].item()) <= 150 # Close to 0 - assert int32_max - 150 <= result[2].item() <= int32_max # Close to max, not overflowed - - def test_cupy_scaling_with_gpu_array(self): - """Test cupy scaling with actual GPU array (when cupy available).""" - cupy = pytest.importorskip("cupy") - scale_func = SCALING_FUNCTIONS["cupy"] - - # Create cupy array - arr = cupy.array([0.0, 0.5, 1.0], dtype=cupy.float32) - result = scale_func(arr, np.int32) - - # Should return cupy array - assert isinstance(result, cupy.ndarray) - assert result.dtype == cupy.int32 - - def test_jax_scaling_with_gpu_array(self): - """Test jax scaling with actual GPU array (when jax available).""" - jax = pytest.importorskip("jax") - jnp = jax.numpy - scale_func = SCALING_FUNCTIONS["jax"] - - # Create jax array - JAX uses numpy dtypes - arr = jnp.array([0.0, 0.5, 1.0], dtype=np.float32) - result = scale_func(arr, np.int32) - - # Should return jax array - assert hasattr(result, "dtype") - assert str(result.dtype) == "int32" - - def test_tensorflow_scaling_with_gpu_array(self): - """Test tensorflow scaling with actual GPU array (when tensorflow available).""" - tf = pytest.importorskip("tensorflow") - scale_func = SCALING_FUNCTIONS["tensorflow"] - - # Create tensorflow tensor - arr = tf.constant([0.0, 0.5, 1.0], dtype=tf.float32) - result = scale_func(arr, np.int32) - - # Should return tensorflow tensor - assert isinstance(result, tf.Tensor) - assert result.dtype == tf.int32 - - @pytest.mark.parametrize( - ("memory_type", "module_name"), - [ - (MemoryType.TORCH, "torch"), - (MemoryType.TENSORFLOW, "tensorflow"), - (MemoryType.JAX, "jax"), - ], - ) - def test_uint16_policy_is_honored_by_typed_frameworks(self, memory_type, module_name): - module = pytest.importorskip(module_name) - if memory_type is MemoryType.TORCH: - values = module.tensor([0.0, 0.5, 1.0], dtype=module.float32) - elif memory_type is MemoryType.TENSORFLOW: - values = module.constant([0.0, 0.5, 1.0], dtype=module.float32) - else: - values = module.numpy.asarray([0.0, 0.5, 1.0], dtype=module.numpy.float32) - - result = memory_type.scale_dtype(values, np.uint16, module) - - expected_dtype = module.numpy.uint16 if memory_type is MemoryType.JAX else module.uint16 - assert result.dtype == expected_dtype - - def test_jax_float64_fails_when_x64_capability_is_disabled(self): - module = type( - "FakeJax", - (), - { - "config": type("Config", (), {"x64_enabled": False})(), - "numpy": object(), - }, - )() - values = type("FakeArray", (), {"dtype": np.dtype(np.float32), "device": None})() - - with pytest.raises(ValueError, match="JAX_ENABLE_X64"): - MemoryType.JAX.scale_dtype(values, np.float64, module) - - def test_pyclesperanto_scaling_preserves_payload_device(self): - class FakeArray: - def __init__(self, values, device): - self.values = np.asarray(values) - self.dtype = self.values.dtype - self.device = type("Device", (), {"name": device})() - - class FakePyclesperanto: - def __init__(self): - self.current = "gpu0" - - def list_available_devices(self, device_type=None): - assert device_type in (None, "gpu") - return ("gpu0", "gpu1") - - def get_device(self): - return self.current - - def select_device(self, selector, device_type=None): - assert device_type in (None, "gpu") - self.current = ("gpu0", "gpu1")[selector] if isinstance(selector, int) else selector - - def pull(self, data): - return data.values - - def push(self, data): - return FakeArray(data, self.current) - - module = FakePyclesperanto() - values = FakeArray([1, 2, 3], "gpu1") - - result = MemoryType.PYCLESPERANTO.scale_dtype(values, np.int32, module) - - assert result.device.name == "gpu1" - assert module.current == "gpu0" - - def test_pyclesperanto_scaling_with_gpu_array(self): - """Test pyclesperanto scaling with actual GPU array (when pyclesperanto available).""" - pytest.importorskip("pyclesperanto") - scale_func = SCALING_FUNCTIONS["pyclesperanto"] - - # Create numpy array (pyclesperanto works with numpy arrays pushed to GPU) - arr = np.array([[0.0, 0.5], [0.25, 1.0]], dtype=np.float32) - result = scale_func(arr, np.int32) - - # pyclesperanto returns its own array type (can be converted via cle.pull()) - # Check it has correct dtype - assert hasattr(result, "dtype") - assert result.dtype == np.int32 or str(result.dtype) == "int32" - - def test_pyclesperanto_scaling_constant_image(self): - """Test pyclesperanto scaling with constant image.""" - pytest.importorskip("pyclesperanto") - scale_func = SCALING_FUNCTIONS["pyclesperanto"] - - # Create constant image - arr = np.full((10, 10), 0.5, dtype=np.float32) - result = scale_func(arr, np.int32) - - # pyclesperanto returns its own array type - assert hasattr(result, "dtype") - assert result.dtype == np.int32 or str(result.dtype) == "int32" - - def test_pyclesperanto_scaling_no_conversion_needed(self): - """Test pyclesperanto scaling when no conversion needed.""" - cle = pytest.importorskip("pyclesperanto") - scale_func = SCALING_FUNCTIONS["pyclesperanto"] - - # int to int - no scaling needed - arr = np.array([1, 2, 3], dtype=np.int32) - result = scale_func(arr, np.int32) - - # pyclesperanto returns its own array type - assert hasattr(result, "dtype") - assert result.dtype == np.int32 or str(result.dtype) == "int32" - # Convert to numpy for value comparison - result_np = cle.pull(result) if hasattr(cle, "pull") else np.asarray(result) - np.testing.assert_array_equal(result_np, [1, 2, 3]) diff --git a/tests/test_durable_decorator_context.py b/tests/test_durable_decorator_context.py deleted file mode 100644 index 24385d5..0000000 --- a/tests/test_durable_decorator_context.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Retained callable declarations exclude thread-local GPU lifecycle handles.""" - -import threading -from concurrent.futures import ThreadPoolExecutor - -import dill -import numpy as np -import pytest - -from arraybridge.decorators import ThreadGPUContext, cupy, numpy -from arraybridge.types import MemoryType - - -def test_context_is_stable_per_thread_and_never_shared_between_threads(): - current = ThreadGPUContext.current() - with ThreadPoolExecutor(max_workers=1) as executor: - other, repeated = executor.submit( - lambda: (ThreadGPUContext.current(), ThreadGPUContext.current()) - ).result() - assert ThreadGPUContext.current() is current - assert other is repeated - assert other is not current - - -@pytest.mark.parametrize("decorator", [numpy, cupy]) -def test_unpublished_decorated_callable_serializes_without_runtime_context(decorator): - # Like a replaced custom function retained by an undo snapshot, this - # declaration has no importable public alias. Dill must persist its body. - @decorator - def declared(image, scale=3): - return image * scale - - context = ThreadGPUContext.current() - key = (MemoryType.CUPY, -1) - handle = threading.local() - context._streams[key] = handle - try: - restored = dill.loads(dill.dumps(declared)) - assert restored.__name__ == declared.__name__ - assert restored.__wrapped__.__name__ == declared.__wrapped__.__name__ - # CPU NumPy invocation checks actual restored behavior. GPU behavior - # keeps its existing focused tests; this receipt starts no GPU runtime. - if decorator is numpy: - np.testing.assert_array_equal( - restored(np.array([1, 2], dtype=np.uint16), scale=4), [4, 8] - ) - assert ThreadGPUContext.current() is context - assert context._streams[key] is handle - finally: - del context._streams[key] diff --git a/tests/test_framework_config.py b/tests/test_framework_config.py deleted file mode 100644 index 203f144..0000000 --- a/tests/test_framework_config.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Tests for the deprecated declaration-derived operation view.""" - -from types import MappingProxyType - -from arraybridge.framework_config import _FRAMEWORK_CONFIG -from arraybridge.framework_ops import _FRAMEWORK_OPS -from arraybridge.types import MemoryType - - -def test_compatibility_view_is_read_only_and_exhaustive(): - assert isinstance(_FRAMEWORK_CONFIG, MappingProxyType) - assert set(_FRAMEWORK_CONFIG) == set(MemoryType) - assert _FRAMEWORK_OPS is _FRAMEWORK_CONFIG - - -def test_compatibility_view_projects_member_owned_operations(): - expected_names = {"to_numpy", "from_numpy", "stack_arrays", "scale_dtype"} - for memory_type, operations in _FRAMEWORK_CONFIG.items(): - assert isinstance(operations, MappingProxyType) - assert set(operations) == expected_names - for name in expected_names: - operation = operations[name] - assert callable(operation) - assert operation.__self__ is memory_type - - -def test_compatibility_view_contains_no_executable_strings(): - for operations in _FRAMEWORK_CONFIG.values(): - assert all(not isinstance(operation, str) for operation in operations.values()) diff --git a/tests/test_gpu_cleanup.py b/tests/test_gpu_cleanup.py deleted file mode 100644 index fd8379a..0000000 --- a/tests/test_gpu_cleanup.py +++ /dev/null @@ -1,156 +0,0 @@ -"""Behavioral tests for declaration-owned GPU cleanup.""" - -import sys -from contextlib import AbstractContextManager -from types import MappingProxyType, SimpleNamespace - -import pytest - -from arraybridge.gpu_cleanup import MEMORY_TYPE_CLEANUP_REGISTRY, cleanup_all_gpu_frameworks -from arraybridge.types import GPU_MEMORY_TYPES, MemoryType - - -class _DeviceScope(AbstractContextManager): - def __init__(self, events, device_id): - self.events = events - self.device_id = device_id - - def __enter__(self): - self.events.append(("enter", self.device_id)) - return self - - def __exit__(self, exc_type, exc_value, traceback): - self.events.append(("exit", self.device_id)) - return False - - -class _Pool: - def __init__(self, name, events): - self.name = name - self.events = events - - def free_all_blocks(self): - self.events.append(("free", self.name)) - - -def _fake_cupy(events, device_count=2): - regular_pool = _Pool("regular", events) - pinned_pool = _Pool("pinned", events) - runtime = SimpleNamespace( - getDeviceCount=lambda: device_count, - deviceSynchronize=lambda: events.append(("synchronize", None)), - ) - cuda = SimpleNamespace( - runtime=runtime, - Device=lambda device_id: _DeviceScope(events, device_id), - ) - return SimpleNamespace( - cuda=cuda, - get_default_memory_pool=lambda: regular_pool, - get_default_pinned_memory_pool=lambda: pinned_pool, - ) - - -class TestCleanupProjection: - def test_compatibility_registry_is_derived_for_every_declaration(self): - assert isinstance(MEMORY_TYPE_CLEANUP_REGISTRY, MappingProxyType) - assert set(MEMORY_TYPE_CLEANUP_REGISTRY) == { - memory_type.value for memory_type in MemoryType - } - for memory_type in MemoryType: - cleanup = MEMORY_TYPE_CLEANUP_REGISTRY[memory_type.value] - assert cleanup.__name__ == f"cleanup_{memory_type.import_name}_gpu" - - def test_cleanup_functions_remain_importable(self): - from arraybridge import gpu_cleanup - - for memory_type in MemoryType: - assert callable(getattr(gpu_cleanup, f"cleanup_{memory_type.import_name}_gpu")) - - -class TestCleanupBehavior: - def test_absent_frameworks_are_never_imported(self, monkeypatch): - from arraybridge import types - - for memory_type in GPU_MEMORY_TYPES: - monkeypatch.delitem(sys.modules, memory_type.import_name, raising=False) - - def fail_import(name): - raise AssertionError(f"cleanup imported absent framework {name}") - - monkeypatch.setattr(types.importlib, "import_module", fail_import) - cleanup_all_gpu_frameworks() - - def test_loaded_cupy_cleanup_executes_declared_leaf(self, monkeypatch): - events = [] - monkeypatch.setitem(sys.modules, "cupy", _fake_cupy(events)) - - MEMORY_TYPE_CLEANUP_REGISTRY[MemoryType.CUPY.value]() - - assert events == [ - ("enter", 0), - ("free", "regular"), - ("free", "pinned"), - ("synchronize", None), - ("exit", 0), - ("enter", 1), - ("free", "regular"), - ("free", "pinned"), - ("synchronize", None), - ("exit", 1), - ] - - def test_device_specific_cleanup_is_scoped(self, monkeypatch): - events = [] - monkeypatch.setitem(sys.modules, "cupy", _fake_cupy(events)) - - MEMORY_TYPE_CLEANUP_REGISTRY[MemoryType.CUPY.value](1) - - assert events == [ - ("enter", 1), - ("free", "regular"), - ("free", "pinned"), - ("synchronize", None), - ("exit", 1), - ] - - def test_declaration_rejects_undeclared_cleanup_device(self, monkeypatch): - events = [] - monkeypatch.setitem(sys.modules, "cupy", _fake_cupy(events, device_count=1)) - - with pytest.raises(ValueError, match="device 3 is unavailable"): - MemoryType.CUPY.cleanup_loaded(3) - - assert events == [] - - def test_compatibility_cleanup_contains_undeclared_device(self, monkeypatch, caplog): - events = [] - monkeypatch.setitem(sys.modules, "cupy", _fake_cupy(events, device_count=1)) - - MEMORY_TYPE_CLEANUP_REGISTRY[MemoryType.CUPY.value](3) - - assert events == [] - assert "device 3 is unavailable" in caplog.text - - def test_jax_compilation_cache_is_not_misrepresented_as_gpu_cleanup(self, monkeypatch): - events = [] - module = SimpleNamespace( - clear_caches=lambda: events.append("clear"), - devices=lambda: [SimpleNamespace(platform="gpu")], - ) - monkeypatch.setitem(sys.modules, "jax", module) - - MemoryType.JAX.cleanup_loaded(0) - - assert events == [] - - def test_numpy_cleanup_is_a_noop_without_import(self, monkeypatch): - from arraybridge import types - - monkeypatch.delitem(sys.modules, "numpy", raising=False) - - def fail_import(name): - raise AssertionError(f"cleanup imported {name}") - - monkeypatch.setattr(types.importlib, "import_module", fail_import) - MEMORY_TYPE_CLEANUP_REGISTRY[MemoryType.NUMPY.value](0) diff --git a/tests/test_integration.py b/tests/test_integration.py deleted file mode 100644 index 869f9f6..0000000 --- a/tests/test_integration.py +++ /dev/null @@ -1,195 +0,0 @@ -"""Integration tests for arraybridge.""" - -import numpy as np -import pytest - -from arraybridge import ( - CPU_MEMORY_TYPES, - GPU_MEMORY_TYPES, - MemoryType, - convert_memory, - detect_memory_type, -) - - -class TestBasicWorkflow: - """Test basic arraybridge workflows.""" - - def test_import_all_exports(self): - """Test that all main exports are importable.""" - from arraybridge import ( - CPU_MEMORY_TYPES, - GPU_MEMORY_TYPES, - SUPPORTED_MEMORY_TYPES, - MemoryConversionError, - MemoryType, - convert_memory, - detect_memory_type, - memory_types, - ) - - # Verify types exist - assert MemoryType is not None - assert CPU_MEMORY_TYPES is not None - assert GPU_MEMORY_TYPES is not None - assert SUPPORTED_MEMORY_TYPES is not None - - # Verify functions exist - assert callable(convert_memory) - assert callable(detect_memory_type) - assert callable(memory_types) - - # Verify exception exists - assert issubclass(MemoryConversionError, Exception) - - def test_simple_numpy_workflow(self): - """Test a simple workflow with NumPy arrays.""" - # Create array - data = np.array([[1, 2], [3, 4]]) - - # Detect type - mem_type = detect_memory_type(data) - assert mem_type == "numpy" - - # Convert (no-op for numpy to numpy) - result = convert_memory(data, source_type="numpy", target_type="numpy", gpu_id=0) - np.testing.assert_array_equal(result, data) - - def test_readme_quick_start_example(self): - """Test the Quick Start example from README.""" - # Create NumPy array - data = np.array([[1, 2], [3, 4]]) - - # Detect memory type - mem_type = detect_memory_type(data) - assert mem_type == "numpy" - - def test_multiple_conversions(self): - """Test multiple sequential conversions.""" - original = np.array([1.0, 2.0, 3.0], dtype=np.float32) - - # Multiple numpy->numpy conversions - result = original - for _ in range(5): - result = convert_memory(result, source_type="numpy", target_type="numpy", gpu_id=0) - - np.testing.assert_array_almost_equal(result, original) - - def test_different_array_shapes(self): - """Test conversion with different array shapes.""" - shapes = [ - (10,), # 1D - (10, 10), # 2D - (5, 10, 10), # 3D - (2, 3, 4, 5), # 4D - ] - - for shape in shapes: - arr = np.random.rand(*shape).astype(np.float32) - detected = detect_memory_type(arr) - assert detected == "numpy" - - result = convert_memory(arr, source_type="numpy", target_type="numpy", gpu_id=0) - assert result.shape == shape - - -class TestFrameworkAvailability: - """Test framework availability detection.""" - - def test_numpy_always_available(self): - """Test that NumPy is always available.""" - import numpy - - assert numpy is not None - - def test_optional_framework_import(self): - """Test optional framework imports.""" - from arraybridge.utils import optional_import - - # NumPy should always work - np = optional_import("numpy") - assert np - assert hasattr(np, "array") - - # Non-existent module should return falsy placeholder - fake = optional_import("nonexistent_module_12345") - assert not fake - - -class TestMemoryTypeConstants: - """Test memory type constants and sets.""" - - def test_cpu_memory_types_contain_only_numpy(self): - """Test CPU memory types.""" - assert len(CPU_MEMORY_TYPES) == 1 - assert MemoryType.NUMPY in CPU_MEMORY_TYPES - - def test_gpu_memory_types_contain_all_gpu_frameworks(self): - """Test GPU memory types.""" - expected = { - MemoryType.CUPY, - MemoryType.TORCH, - MemoryType.TENSORFLOW, - MemoryType.JAX, - MemoryType.PYCLESPERANTO, - } - assert GPU_MEMORY_TYPES == expected - - def test_no_overlap_between_cpu_and_gpu(self): - """Test that CPU and GPU memory types don't overlap.""" - overlap = CPU_MEMORY_TYPES & GPU_MEMORY_TYPES - assert len(overlap) == 0 - - -class TestErrorHandling: - """Test error handling in arraybridge.""" - - def test_invalid_source_type_raises_error(self): - """Test that invalid source type raises error.""" - arr = np.array([1, 2, 3]) - with pytest.raises(ValueError): - convert_memory(arr, source_type="invalid", target_type="numpy", gpu_id=0) - - def test_invalid_target_type_raises_error(self): - """Test that invalid target type raises error.""" - arr = np.array([1, 2, 3]) - with pytest.raises((ValueError, AttributeError)): - convert_memory(arr, source_type="numpy", target_type="invalid", gpu_id=0) - - def test_detect_invalid_type_raises_error(self): - """Test that detecting invalid type raises error.""" - with pytest.raises(ValueError): - detect_memory_type("not an array") - - def test_detect_none_raises_error(self): - """Test that detecting None raises error.""" - with pytest.raises(ValueError): - detect_memory_type(None) - - -class TestDtypePreservation: - """Test that dtypes are preserved during conversions.""" - - def test_uint8_dtype_preserved(self): - """Test uint8 dtype preservation.""" - arr = np.array([0, 128, 255], dtype=np.uint8) - result = convert_memory(arr, source_type="numpy", target_type="numpy", gpu_id=0) - assert result.dtype == np.uint8 - - def test_uint16_dtype_preserved(self): - """Test uint16 dtype preservation.""" - arr = np.array([0, 1000, 65535], dtype=np.uint16) - result = convert_memory(arr, source_type="numpy", target_type="numpy", gpu_id=0) - assert result.dtype == np.uint16 - - def test_float32_dtype_preserved(self): - """Test float32 dtype preservation.""" - arr = np.array([1.5, 2.5, 3.5], dtype=np.float32) - result = convert_memory(arr, source_type="numpy", target_type="numpy", gpu_id=0) - assert result.dtype == np.float32 - - def test_float64_dtype_preserved(self): - """Test float64 dtype preservation.""" - arr = np.array([1.5, 2.5, 3.5], dtype=np.float64) - result = convert_memory(arr, source_type="numpy", target_type="numpy", gpu_id=0) - assert result.dtype == np.float64 diff --git a/tests/test_libraries.py b/tests/test_libraries.py new file mode 100644 index 0000000..0b0fb56 --- /dev/null +++ b/tests/test_libraries.py @@ -0,0 +1,519 @@ +"""The ``ArrayLibrary`` family: lookup, identity, import environment, memory and operations.""" + +import importlib +import importlib.util +import os +import pickle +import sys +from types import SimpleNamespace + +import numpy as np +import pytest +from metaclass_registry.exceptions import DuplicateRegistryKeyError + +from arraybridge.devices import HOST, AcceleratorDevice +from arraybridge.libraries import ( + AcceleratorDevices, + ArrayLibrary, + CudaLibrary, + HostDevices, + NumpyLibrary, +) +from tests.fakes import FakeCupy, FakePyclesperanto, FakePyclesperantoArray, FakeTorch, fake_jax + +LIBRARIES = ArrayLibrary.all() +NUMPY = ArrayLibrary.named("numpy") + + +def names(role): + return {library.name for library in LIBRARIES if isinstance(library, role)} + + +class TestLookup: + def test_one_object_per_registered_class(self): + assert [library.name for library in LIBRARIES] == list(ArrayLibrary.__registry__) + for library in LIBRARIES: + assert type(library) is ArrayLibrary.__registry__[library.name] + assert ArrayLibrary.named(library.name) is library + assert ArrayLibrary.for_module_name(f"{library.import_name}.submodule") is library + assert pickle.loads(pickle.dumps(library)) is library + assert str(library) == library.name + + def test_roles_are_inherited(self): + assert names(HostDevices) == {"numpy", "torch", "tensorflow", "jax"} + assert names(AcceleratorDevices) == { + "cupy", + "torch", + "tensorflow", + "jax", + "pyclesperanto", + } + assert names(CudaLibrary) == {"cupy", "torch"} + + def test_identity_is_stated_by_each_class(self): + assert NUMPY.import_name == "numpy" + assert NUMPY.display_name == "NumPy" + assert ArrayLibrary.named("pyclesperanto").display_name == "pyclesperanto" + assert ArrayLibrary.named("jax").recognized_module_names == frozenset({"jax", "jaxlib"}) + assert ArrayLibrary.for_module_name("jaxlib") is ArrayLibrary.named("jax") + + def test_unknown_names_are_rejected(self): + with pytest.raises(ValueError, match="not a registered array library"): + ArrayLibrary.named("unknown_type") + with pytest.raises(ValueError, match="No array library owns module"): + ArrayLibrary.for_module_name("scipy.ndimage") + + def test_of_names_the_library_that_made_a_value(self): + assert ArrayLibrary.of(np.zeros((2, 2))) is NUMPY + for value in ([1, 2, 3], None, "not an array"): + with pytest.raises(ValueError, match="Unknown memory type"): + ArrayLibrary.of(value) + + def test_a_name_is_registered_to_one_class(self): + with pytest.raises(DuplicateRegistryKeyError, match="key 'numpy' is already registered"): + + class SecondNumpy(NumpyLibrary): + name = "numpy" + + assert ArrayLibrary.__registry__["numpy"] is NumpyLibrary + assert ArrayLibrary.named("numpy") is NUMPY + + +class TestImporting: + @pytest.mark.parametrize("library", LIBRARIES[1:], ids=str) + def test_absent_library_is_never_imported(self, library, absent_libraries, monkeypatch): + monkeypatch.delenv("XLA_PYTHON_CLIENT_PREALLOCATE", raising=False) + monkeypatch.delenv("TF_FORCE_GPU_ALLOW_GROWTH", raising=False) + values = np.array([1, 2, 3]) + + assert library.import_if_installed() is None + assert library.accelerator_devices() == () + assert library.scale_dtype(values, np.float32) is values + library.cleanup_loaded() + assert not library.is_oom_error(RuntimeError("unrelated")) + assert not {"XLA_PYTHON_CLIENT_PREALLOCATE", "TF_FORCE_GPU_ALLOW_GROWTH"} & set(os.environ) + + def test_cleanup_of_every_library_imports_nothing(self, absent_libraries): + ArrayLibrary.cleanup_all_loaded() + ArrayLibrary.cleanup_all_loaded(AcceleratorDevice(0)) + + def test_import_preparation_sets_only_missing_defaults(self, monkeypatch): + tensorflow, jax = ArrayLibrary.named("tensorflow"), ArrayLibrary.named("jax") + monkeypatch.delenv("TF_FORCE_GPU_ALLOW_GROWTH", raising=False) + monkeypatch.setenv("XLA_PYTHON_CLIENT_PREALLOCATE", "custom") + + tensorflow.prepare_import() + jax.prepare_import() + + assert tensorflow.import_environment == (("TF_FORCE_GPU_ALLOW_GROWTH", "true"),) + assert jax.import_environment == (("XLA_PYTHON_CLIENT_PREALLOCATE", "false"),) + assert os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] == "true" + assert os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] == "custom" + + def test_library_loaded_before_its_defaults_is_reported(self, monkeypatch, caplog): + module = fake_jax("cpu") + monkeypatch.delenv("XLA_PYTHON_CLIENT_PREALLOCATE", raising=False) + monkeypatch.setitem(sys.modules, "jax", module) + jax = ArrayLibrary.named("jax") + + assert jax.import_if_installed() is module + assert jax.module is module + assert os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] == "false" + assert "loaded before its ArrayBridge import defaults" in caplog.text + + def test_child_environment_adds_nvidia_wheel_directories(self, monkeypatch, tmp_path): + nvidia_root = tmp_path / "nvidia" + cublas_lib = nvidia_root / "cublas" / "lib" + runtime_bin = nvidia_root / "cuda_runtime" / "bin" + cublas_lib.mkdir(parents=True) + runtime_bin.mkdir(parents=True) + monkeypatch.setattr( + importlib.util, + "find_spec", + lambda name: ( + SimpleNamespace(submodule_search_locations=(str(nvidia_root),)) + if name == "nvidia" + else None + ), + ) + search_variable = "PATH" if os.name == "nt" else "LD_LIBRARY_PATH" + original = {search_variable: "/existing", "CUSTOM": "value"} + + prepared = ArrayLibrary.child_environment(original) + + assert original == {search_variable: "/existing", "CUSTOM": "value"} + assert prepared[search_variable].split(os.pathsep) == [ + *sorted((str(cublas_lib), str(runtime_bin))), + "/existing", + ] + assert prepared["CUSTOM"] == "value" + assert prepared["TF_FORCE_GPU_ALLOW_GROWTH"] == "true" + assert prepared["XLA_PYTHON_CLIENT_PREALLOCATE"] == "false" + + def test_child_environment_keeps_the_caller_values(self, monkeypatch): + monkeypatch.setattr(importlib.util, "find_spec", lambda name: None) + original = { + "TF_FORCE_GPU_ALLOW_GROWTH": "host", + "XLA_PYTHON_CLIENT_PREALLOCATE": "host", + } + + prepared = ArrayLibrary.child_environment(original) + + assert prepared == original + assert prepared is not original + + +class TestOutOfMemory: + @pytest.mark.parametrize( + ("name", "messages"), + [ + ("numpy", ("cannot allocate memory", "Memory exhausted")), + ("cupy", ("Out of memory", "CUDA_ERROR_OUT_OF_MEMORY")), + ("torch", ("CUDA out of memory", "cuda_error_out_of_memory")), + ("tensorflow", ("out of memory", "RESOURCE_EXHAUSTED: failed to allocate")), + ("jax", ("out of memory", "OOM when allocating", "allocation failure")), + ("pyclesperanto", ("CL_MEM_OBJECT_ALLOCATION_FAILURE", "CL_OUT_OF_RESOURCES")), + ], + ) + def test_messages_are_recognised_per_library(self, name, messages, absent_libraries): + library = ArrayLibrary.named(name) + for message in messages: + assert library.is_oom_error(RuntimeError(message)), message + assert not library.is_oom_error(RuntimeError("some random error")) + + @pytest.mark.parametrize( + ("name", "path"), + [ + ("torch", ("cuda", "OutOfMemoryError")), + ("cupy", ("cuda", "memory", "OutOfMemoryError")), + ("cupy", ("cuda", "runtime", "CUDARuntimeError")), + ("tensorflow", ("errors", "ResourceExhaustedError")), + ], + ) + def test_loaded_module_exception_classes_are_recognised(self, name, path, monkeypatch): + class LibraryOutOfMemoryError(Exception): + pass + + module = LibraryOutOfMemoryError + for attribute in reversed(path): + module = SimpleNamespace(**{attribute: module}) + monkeypatch.setitem(sys.modules, name, module) + monkeypatch.setattr( + importlib, "import_module", lambda name: pytest.fail(f"imported {name}") + ) + library = ArrayLibrary.named(name) + + assert library.is_oom_error(LibraryOutOfMemoryError("allocation failed")) + assert not library.is_oom_error(RuntimeError("allocation failed")) + + @pytest.mark.parametrize("library", LIBRARIES, ids=str) + def test_retries_after_an_out_of_memory_error(self, library, absent_libraries): + calls = [] + + def fails_once(): + calls.append(len(calls)) + if len(calls) == 1: + raise MemoryError(library.oom_messages[0]) + return "success" + + assert library.run_with_oom_recovery(fails_once, max_retries=1) == "success" + assert calls == [0, 1] + + @pytest.mark.parametrize(("max_retries", "expected_calls"), [(0, 1), (2, 3)]) + def test_gives_up_after_the_allowed_retries(self, max_retries, expected_calls): + calls = [] + + def always_fails(): + calls.append(None) + raise MemoryError("out of memory") + + with pytest.raises(MemoryError, match="out of memory"): + ArrayLibrary.named("torch").run_with_oom_recovery( + always_fails, max_retries=max_retries + ) + assert len(calls) == expected_calls + + def test_other_errors_are_not_retried(self): + calls = [] + + def raises_value_error(): + calls.append(None) + raise ValueError("not an OOM error") + + with pytest.raises(ValueError, match="not an OOM error"): + ArrayLibrary.named("torch").run_with_oom_recovery(raises_value_error) + assert len(calls) == 1 + + def test_retry_releases_the_cache_of_the_given_device(self, monkeypatch): + module = FakeCupy() + monkeypatch.setitem(sys.modules, "cupy", module) + outcomes = [MemoryError("out of memory"), "success"] + + def fails_once(): + outcome = outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + return outcome + + result = ArrayLibrary.named("cupy").run_with_oom_recovery( + fails_once, device=AcceleratorDevice(1) + ) + + assert result == "success" + assert module.events == [ + ("enter", 1), + ("free", "regular"), + ("free", "pinned"), + ("synchronize", None), + ("exit", 1), + ] + + +class TestCleanup: + def test_cupy_pools_are_freed_on_every_device(self, monkeypatch): + module = FakeCupy() + monkeypatch.setitem(sys.modules, "cupy", module) + + ArrayLibrary.named("cupy").cleanup_loaded() + + freed = [("free", "regular"), ("free", "pinned"), ("synchronize", None)] + assert module.events == [ + ("enter", 0), + *freed, + ("exit", 0), + ("enter", 1), + *freed, + ("exit", 1), + ] + + def test_torch_cache_is_emptied_on_the_given_device(self, monkeypatch): + module = FakeTorch(device_count=2) + monkeypatch.setitem(sys.modules, "torch", module) + + ArrayLibrary.cleanup_all_loaded(AcceleratorDevice(1)) + + assert module.events == [ + ("enter", 1), + ("empty_cache", 1), + ("synchronize", 1), + ("exit", 1), + ] + + def test_torch_without_cuda_has_nothing_to_release(self, monkeypatch): + module = FakeTorch() + monkeypatch.setitem(sys.modules, "torch", module) + + ArrayLibrary.named("torch").cleanup_loaded() + + assert module.events == [] + + def test_unavailable_device_is_rejected_by_the_library(self, monkeypatch): + module = FakeCupy(device_count=1) + monkeypatch.setitem(sys.modules, "cupy", module) + + with pytest.raises(ValueError, match="device accelerator:3 is unavailable"): + ArrayLibrary.named("cupy").cleanup_loaded(AcceleratorDevice(3)) + + assert module.events == [] + + def test_cleanup_of_every_library_logs_a_failing_one(self, monkeypatch, caplog): + module = FakeCupy(device_count=1) + monkeypatch.setitem(sys.modules, "cupy", module) + + ArrayLibrary.cleanup_all_loaded(AcceleratorDevice(3)) + + assert module.events == [] + assert "Failed to cleanup CuPy GPU memory" in caplog.text + assert "device accelerator:3 is unavailable" in caplog.text + + def test_jax_compilation_cache_is_not_cleared_as_device_memory(self, monkeypatch): + events = [] + module = fake_jax("gpu") + module.clear_caches = lambda: events.append("clear") + monkeypatch.setitem(sys.modules, "jax", module) + + ArrayLibrary.named("jax").cleanup_loaded(AcceleratorDevice(0)) + + assert events == [] + + +class TestScaleDtype: + @pytest.mark.parametrize( + ("values", "dtype", "expected"), + [ + (np.array([1, 2, 3], dtype=np.uint8), np.uint16, [1, 2, 3]), + (np.array([0.0, 0.5, 1.0], dtype=np.float64), np.uint8, [0, 127, 255]), + (np.array([0.0, 0.5, 1.0], dtype=np.float32), np.uint16, [0, 32767, 65535]), + (np.array([0.0, 1.0], dtype=np.float16), np.uint8, [0, 255]), + (np.array([0.0, 1.0], dtype=np.float64), np.int16, [-32768, 32639]), + (np.array([0.0, 1.0], dtype=np.float64), np.uint32, [0, 4294967295]), + (np.array([0.25, 0.75], dtype=np.float64), np.float32, [0.25, 0.75]), + ], + ) + def test_numpy_values(self, values, dtype, expected): + result = NUMPY.scale_dtype(values, dtype) + + assert result.dtype == dtype + np.testing.assert_array_equal(result, expected) + + def test_numpy_int32_stays_inside_the_range(self): + result = NUMPY.scale_dtype(np.array([0.0, 0.5, 1.0], dtype=np.float64), np.int32) + + assert result.dtype == np.int32 + assert result[0] == np.iinfo(np.int32).min + assert abs(int(result[1])) <= 150 + assert np.iinfo(np.int32).max - 150 <= result[2] <= np.iinfo(np.int32).max + + @pytest.mark.parametrize( + "values", + [ + np.full((10, 10), 0.5, dtype=np.float32), + np.array([0.7], dtype=np.float32), + ], + ) + def test_numpy_constant_image_is_cast_without_rescaling(self, values): + result = NUMPY.scale_dtype(values, np.uint8) + + assert result.dtype == np.uint8 + np.testing.assert_array_equal(result, np.zeros(values.shape, dtype=np.uint8)) + + def test_value_without_a_dtype_is_returned_unchanged(self): + assert NUMPY.scale_dtype("not an array", np.uint8) == "not an array" + + def test_numpy_empty_array_has_no_range_to_scale(self): + with pytest.raises(ValueError): + NUMPY.scale_dtype(np.array([], dtype=np.float32), np.uint8) + + @pytest.mark.parametrize("name", ["torch", "cupy", "jax", "tensorflow"]) + def test_installed_library_scales_into_the_integer_range(self, name): + module = pytest.importorskip(name) + library = ArrayLibrary.named(name) + host_values = np.array([0.0, 0.5, 1.0], dtype=np.float32) + values = library.from_numpy(host_values, _first_device(library)) + + as_uint16 = library.scale_dtype(values, np.uint16) + as_int32 = library.to_numpy(library.scale_dtype(values, np.int32)) + + assert library.accepts(as_uint16) + assert library.dtype_name(as_uint16.dtype) == "uint16" + np.testing.assert_array_equal(library.to_numpy(as_uint16), [0, 32767, 65535]) + assert as_int32.dtype == np.int32 + assert as_int32[0] <= np.iinfo(np.int32).min + 150 + assert abs(int(as_int32[1])) <= 150 + assert np.iinfo(np.int32).max - 150 <= as_int32[2] + assert module is library.module + + def test_jax_float64_fails_when_x64_is_disabled(self, monkeypatch): + module = fake_jax("cpu") + module.config = SimpleNamespace(x64_enabled=False) + module.numpy = object() + monkeypatch.setitem(sys.modules, "jax", module) + values = SimpleNamespace(dtype=np.dtype(np.float32), device=None) + + with pytest.raises(ValueError, match="JAX_ENABLE_X64"): + ArrayLibrary.named("jax").scale_dtype(values, np.float64) + + def test_pyclesperanto_scales_on_the_device_of_its_input(self, monkeypatch): + module = FakePyclesperanto(current="gpu0") + monkeypatch.setitem(sys.modules, "pyclesperanto", module) + values = FakePyclesperantoArray([1, 2, 3], "gpu1") + + result = ArrayLibrary.named("pyclesperanto").scale_dtype(values, np.int32) + + assert result.device.name == "gpu1" + assert result.dtype == np.int32 + np.testing.assert_array_equal(result.values, [1, 2, 3]) + assert module.current == "gpu0" + + @pytest.mark.pyclesperanto + @pytest.mark.parametrize( + "values", + [ + np.array([[0.0, 0.5], [0.25, 1.0]], dtype=np.float32), + np.full((10, 10), 0.5, dtype=np.float32), + np.array([[1, 2, 3]], dtype=np.int32), + ], + ) + def test_installed_pyclesperanto_scales_to_int32(self, values): + pytest.importorskip("pyclesperanto") + library = ArrayLibrary.named("pyclesperanto") + if not library.accelerator_devices(): + pytest.skip("No OpenCL GPU device") + + result = library.scale_dtype(library.from_numpy(values, AcceleratorDevice(0)), np.int32) + + assert str(result.dtype) == "int32" + if values.dtype == np.int32: + np.testing.assert_array_equal(library.to_numpy(result), values) + + +def _first_device(library): + """Main memory where the library has it, otherwise its first accelerator.""" + + devices = library.available_devices() + if not devices: + pytest.skip(f"{library.display_name} has no device on this machine") + return devices[0] + + +class TestOperations: + def test_numpy_operations(self): + stack = np.arange(24, dtype=np.uint16).reshape(2, 3, 4) + + assert NUMPY.module is np + assert NUMPY.accepts(stack) and not NUMPY.accepts([1, 2]) + np.testing.assert_array_equal(NUMPY.max(stack, 0), stack[1]) + np.testing.assert_array_equal(NUMPY.min(stack, 0), stack[0]) + np.testing.assert_array_equal(NUMPY.mean(stack, 0), stack.mean(axis=0)) + np.testing.assert_array_equal(NUMPY.sum(stack, (1, 2)), stack.sum(axis=(1, 2))) + assert NUMPY.prepend_axis(stack[0]).shape == (1, 3, 4) + assert NUMPY.astype(stack, np.float32).dtype == np.float32 + ramp = NUMPY.linspace(0, 1, 4, endpoint=False) + np.testing.assert_array_equal(ramp, [0, 0.25, 0.5, 0.75]) + weights = NUMPY.assign(NUMPY.ones(5, np.float32), slice(-2, None), ramp[:2]) + np.testing.assert_array_equal(weights, [1, 1, 1, 0, 0.25]) + np.testing.assert_array_equal(NUMPY.outer(ramp, ramp), np.outer(ramp, ramp)) + assert NUMPY.floor(2.7) == 2 + assert NUMPY.nbytes_of(stack) == stack.nbytes + + @pytest.mark.jax + def test_jax_assignment_returns_a_new_array(self): + jnp = pytest.importorskip("jax.numpy") + jax = ArrayLibrary.named("jax") + values = jnp.ones(4, dtype=jnp.float32) + + updated = jax.assign(values, slice(None, 2), jnp.zeros(2, dtype=jnp.float32)) + + np.testing.assert_array_equal(np.asarray(updated), [0, 0, 1, 1]) + np.testing.assert_array_equal(np.asarray(values), [1, 1, 1, 1]) + stack = jnp.arange(24, dtype=jnp.uint16).reshape(2, 3, 4) + assert jax.mean(stack, 0).dtype == jnp.float32 + + def test_dtype_names_are_spelled_as_numpy_spells_them(self): + class TensorflowDtype: + as_numpy_dtype = np.uint16 + + class TorchDtype: + def __str__(self): + return "torch.uint16" + + assert NUMPY.dtype_name(np.uint16) == "uint16" + assert ArrayLibrary.named("tensorflow").dtype_name(TensorflowDtype()) == "uint16" + assert ArrayLibrary.named("torch").dtype_name(TorchDtype()) == "uint16" + + def test_torch_wraps_a_reversed_view_as_a_contiguous_copy(self, monkeypatch): + module = FakeTorch(device_count=1) + monkeypatch.setitem(sys.modules, "torch", module) + values = np.arange(12, dtype=np.float32).reshape(3, 4)[:, ::-1] + torch = ArrayLibrary.named("torch") + + on_host = torch.from_numpy(values, HOST) + on_accelerator = torch.from_numpy(values, AcceleratorDevice(0)) + + wrapped = [data for event, data in module.events if event == "from_numpy"] + assert len(wrapped) == 2 + for host_array in wrapped: + assert all(stride >= 0 for stride in host_array.strides) + np.testing.assert_array_equal(host_array, values) + assert torch.device_of(on_host) == HOST + assert torch.device_of(on_accelerator) == AcceleratorDevice(0) diff --git a/tests/test_memory_type.py b/tests/test_memory_type.py new file mode 100644 index 0000000..949c672 --- /dev/null +++ b/tests/test_memory_type.py @@ -0,0 +1,77 @@ +"""``MemoryType`` is the registered library names as an enum, and nothing more.""" + +import pickle +from types import SimpleNamespace + +import pytest + +import arraybridge +from arraybridge.libraries import ArrayLibrary +from arraybridge.types import MemoryContractAttribute, MemoryType + +LIBRARY_NAMES = {"numpy", "cupy", "torch", "tensorflow", "jax", "pyclesperanto"} + + +def test_members_are_the_registered_libraries_in_order(): + assert [member.value for member in MemoryType] == list(ArrayLibrary.__registry__) + assert {member.value for member in MemoryType} == LIBRARY_NAMES + assert len(MemoryType) == len(ArrayLibrary.all()) + for member in MemoryType: + assert member.name == member.value.upper() + assert MemoryType(member.value) is member + assert getattr(MemoryType, member.name) is member + assert MemoryType.__members__[member.name] is member + assert member in MemoryType + assert member.library is ArrayLibrary.named(member.value) + assert pickle.loads(pickle.dumps(member)) is member + + +def test_unregistered_name_is_not_a_member(): + with pytest.raises(ValueError): + MemoryType("invalid_type") + with pytest.raises(AttributeError): + MemoryType.INVALID_TYPE + + +def test_members_carry_no_library_behaviour(): + removed = ("is_gpu", "to_numpy", "from_numpy", "stack_arrays", "scale_dtype", "import_name") + for member in MemoryType: + assert not any(hasattr(member, name) for name in removed) + for constant in ("CPU_MEMORY_TYPES", "GPU_MEMORY_TYPES", "SUPPORTED_MEMORY_TYPES"): + assert not hasattr(arraybridge, constant) + + +def test_importing_the_enum_imports_no_array_library(loaded_libraries_after): + loaded = loaded_libraries_after( + "from arraybridge.types import MemoryType\n" + "names = [member.value for member in MemoryType]\n" + "assert len(names) == 6 and 'numpy' in names, names\n" + "assert MemoryType.TORCH.value == 'torch'\n" + ) + assert loaded == "[]" + + +def test_package_root_exports_resolve(): + for name in arraybridge.__all__: + assert getattr(arraybridge, name) is not None + assert arraybridge.MemoryType is MemoryType + assert LIBRARY_NAMES <= set(arraybridge.__all__) + with pytest.raises(AttributeError): + arraybridge.not_a_library + + +class TestMemoryContractAttribute: + def test_reads_and_writes_object_namespaces(self): + namespace = SimpleNamespace() + + MemoryContractAttribute.EXECUTION.write(namespace, "torch") + + assert MemoryContractAttribute.EXECUTION.read(namespace) == "torch" + + def test_reads_and_writes_mapping_namespaces(self): + namespace = {} + + MemoryContractAttribute.INPUT.write(namespace, "numpy") + + assert MemoryContractAttribute.INPUT.read(namespace) == "numpy" + assert MemoryContractAttribute.OUTPUT.read(namespace) is None diff --git a/tests/test_native_geometry.py b/tests/test_native_geometry.py index 5aa0ba8..2aa4720 100644 --- a/tests/test_native_geometry.py +++ b/tests/test_native_geometry.py @@ -1,22 +1,26 @@ """Native geometry must preserve pixels, sharing and the reference device.""" +import sys + import numpy as np import pytest -from arraybridge.array_operations import ArrayOperations -from arraybridge.types import MemoryType +from arraybridge.libraries import ArrayLibrary +from tests.fakes import FakePyclesperanto -PYCLESPERANTO_OPERATIONS = ArrayOperations.for_memory(MemoryType.PYCLESPERANTO) +NUMPY = ArrayLibrary.named("numpy") +CUPY = ArrayLibrary.named("cupy") +PYCLESPERANTO = ArrayLibrary.named("pyclesperanto") def test_numpy_views_and_independent_reference_allocation(): source = np.arange(12, dtype=np.float32).reshape(3, 4)[:, ::-1] - shaped = MemoryType.NUMPY.reshape(source, (3, 1, 4)) - broadcast = MemoryType.NUMPY.broadcast_to(shaped, (3, 2, 4)) + shaped = NUMPY.reshape(source, (3, 1, 4)) + broadcast = NUMPY.broadcast_to(shaped, (3, 2, 4)) assert np.shares_memory(shaped, source) assert np.shares_memory(broadcast, source) np.testing.assert_array_equal(broadcast, np.broadcast_to(source[:, None], (3, 2, 4))) - ones = MemoryType.NUMPY.ones_like(source, shape=(2, 4), dtype=bool) + ones = NUMPY.ones_like(source, shape=(2, 4), dtype=bool) assert ones.dtype == np.bool_ assert not np.shares_memory(ones, source) np.testing.assert_array_equal(ones, np.ones((2, 4), dtype=bool)) @@ -25,31 +29,33 @@ def test_numpy_views_and_independent_reference_allocation(): def test_numpy_geometry_rejects_invalid_shapes(): source = np.arange(6).reshape(2, 3) with pytest.raises(ValueError): - MemoryType.NUMPY.reshape(source, (7,)) + NUMPY.reshape(source, (7,)) with pytest.raises(ValueError): - MemoryType.NUMPY.broadcast_to(source, (2, 4)) + NUMPY.broadcast_to(source, (2, 4)) def test_cupy_geometry_and_allocation_stay_native(): cp = pytest.importorskip("cupy") - if not MemoryType.CUPY.available_device_ids(cp): + if not CUPY.accelerator_devices(): pytest.skip("No CUDA device") source = cp.asarray(np.arange(12, dtype=np.float32).reshape(3, 4))[:, ::-1] - shaped = MemoryType.CUPY.reshape(source, (3, 1, 4)) - broadcast = MemoryType.CUPY.broadcast_to(shaped, (3, 2, 4)) + shaped = CUPY.reshape(source, (3, 1, 4)) + broadcast = CUPY.broadcast_to(shaped, (3, 2, 4)) assert shaped.data.ptr == source.data.ptr assert broadcast.data.ptr == source.data.ptr assert broadcast.device.id == source.device.id np.testing.assert_array_equal( broadcast.get(), np.broadcast_to(source.get()[:, None], (3, 2, 4)) ) - ones = MemoryType.CUPY.ones_like(source, shape=(2, 4), dtype=bool) + ones = CUPY.ones_like(source, shape=(2, 4), dtype=bool) assert ones.device.id == source.device.id assert ones.dtype == cp.bool_ np.testing.assert_array_equal(ones.get(), np.ones((2, 4), dtype=bool)) -def test_pycles_geometry_rejects_host_fallback(): +def test_pycles_geometry_rejects_host_fallback(monkeypatch): + monkeypatch.setitem(sys.modules, "pyclesperanto", FakePyclesperanto(current="cpu0")) + class DeviceOnlyArray: shape = (2, 3) @@ -57,14 +63,14 @@ def __array__(self): pytest.fail("Native geometry must never download pixels") source = DeviceOnlyArray() - assert PYCLESPERANTO_OPERATIONS.reshape(source, (2, 3), None) is source - assert PYCLESPERANTO_OPERATIONS.broadcast_to(source, (2, 3), None) is source + assert PYCLESPERANTO.reshape(source, (2, 3)) is source + assert PYCLESPERANTO.broadcast_to(source, (2, 3)) is source with pytest.raises(NotImplementedError, match="reshape"): - PYCLESPERANTO_OPERATIONS.reshape(source, (1, 2, 3), None) + PYCLESPERANTO.reshape(source, (1, 2, 3)) with pytest.raises(NotImplementedError, match="broadcasting"): - PYCLESPERANTO_OPERATIONS.broadcast_to(source, (4, 2, 3), None) + PYCLESPERANTO.broadcast_to(source, (4, 2, 3)) with pytest.raises(NotImplementedError, match="dtype"): - PYCLESPERANTO_OPERATIONS.ones_like(source, (2, 3), bool, None) + PYCLESPERANTO.ones_like(source, shape=(2, 3), dtype=bool) @pytest.mark.parametrize("dtype", (np.float32, np.complex64, np.int32)) @@ -72,11 +78,11 @@ def test_normalize_planes_preserves_source_and_division_dtype(dtype): source = np.arange(24, dtype=dtype).reshape(2, 3, 4)[:, :, ::-1] original = source.copy() expected = np.stack((source[0] / 3.0, source[1])) - actual = MemoryType.NUMPY.normalize_planes(source, dtype, (3.0, None)) + actual = NUMPY.normalize_planes(source, dtype, (3.0, None)) np.testing.assert_array_equal(actual, expected) np.testing.assert_array_equal(source, original) assert actual.dtype == expected.dtype assert not np.shares_memory(actual, source) for scales in ((3.0,), (3.0, None, 2.0)): with pytest.raises(ValueError, match="leading plane axis"): - MemoryType.NUMPY.normalize_planes(source, dtype, scales) + NUMPY.normalize_planes(source, dtype, scales) diff --git a/tests/test_oom_recovery.py b/tests/test_oom_recovery.py deleted file mode 100644 index dc33058..0000000 --- a/tests/test_oom_recovery.py +++ /dev/null @@ -1,141 +0,0 @@ -"""Tests for arraybridge.oom_recovery module.""" - -import pytest - - -class TestOOMRecovery: - """Tests for OOM recovery functions.""" - - def test_is_oom_error_none_memory_type(self): - """Test _is_oom_error with None/unknown memory type.""" - from arraybridge.oom_recovery import _is_oom_error - - e = Exception("some error") - assert _is_oom_error(e, "unknown_type") is False - - def test_is_oom_error_generic_exception(self): - """Test _is_oom_error with generic exception.""" - from arraybridge.oom_recovery import _is_oom_error - - e = Exception("some random error") - assert _is_oom_error(e, "torch") is False - - def test_is_oom_error_memory_error(self): - """Test _is_oom_error with MemoryError.""" - from arraybridge.oom_recovery import _is_oom_error - - e = MemoryError("out of memory") - # Should detect based on string patterns - assert _is_oom_error(e, "torch") is True - - def test_is_oom_error_string_patterns(self): - """Test _is_oom_error with various OOM string patterns.""" - from arraybridge.oom_recovery import _is_oom_error - - # Test torch patterns - torch_oom_messages = [ - "out of memory", - "cuda_error_out_of_memory", - ] - for msg in torch_oom_messages: - e = Exception(msg) - assert _is_oom_error(e, "torch") is True, f"Failed to detect OOM in torch: {msg}" - - # Test numpy patterns - numpy_oom_messages = ["memory exhausted", "cannot allocate memory"] - for msg in numpy_oom_messages: - e = Exception(msg) - assert _is_oom_error(e, "numpy") is True, f"Failed to detect OOM in numpy: {msg}" - - def test_clear_cache_for_memory_type_unknown(self): - """Test _clear_cache_for_memory_type with unknown memory type.""" - from arraybridge.oom_recovery import _clear_cache_for_memory_type - - # Should not raise, just log warning and do gc.collect() - _clear_cache_for_memory_type("unknown_type") - - def test_clear_cache_for_memory_type_numpy(self): - """Test _clear_cache_for_memory_type with numpy (CPU).""" - from arraybridge.oom_recovery import _clear_cache_for_memory_type - - # Should just do gc.collect() - _clear_cache_for_memory_type("numpy") - - def test_execute_with_oom_recovery_no_oom(self): - """Test _execute_with_oom_recovery when no OOM occurs.""" - from arraybridge.oom_recovery import _execute_with_oom_recovery - - def successful_func(): - return "success" - - result = _execute_with_oom_recovery(successful_func, "torch", max_retries=2) - assert result == "success" - - def test_execute_with_oom_recovery_oom_retry_success(self): - """Test _execute_with_oom_recovery with OOM that succeeds on retry.""" - from arraybridge.oom_recovery import _execute_with_oom_recovery - - call_count = {"count": 0} - - def failing_then_success_func(): - call_count["count"] += 1 - if call_count["count"] == 1: - raise MemoryError("out of memory") - return "success" - - result = _execute_with_oom_recovery(failing_then_success_func, "torch", max_retries=2) - assert result == "success" - assert call_count["count"] == 2 - - def test_execute_with_oom_recovery_oom_exhausted_retries(self): - """Test _execute_with_oom_recovery with OOM that exhausts retries.""" - from arraybridge.oom_recovery import _execute_with_oom_recovery - - def always_fails_func(): - raise MemoryError("out of memory") - - with pytest.raises(MemoryError) as exc_info: - _execute_with_oom_recovery(always_fails_func, "torch", max_retries=2) - assert "out of memory" in str(exc_info.value) - - def test_execute_with_oom_recovery_non_oom_exception(self): - """Test _execute_with_oom_recovery with non-OOM exception (should not retry).""" - from arraybridge.oom_recovery import _execute_with_oom_recovery - - def raises_value_error(): - raise ValueError("not an OOM error") - - with pytest.raises(ValueError) as exc_info: - _execute_with_oom_recovery(raises_value_error, "torch", max_retries=2) - assert "not an OOM error" in str(exc_info.value) - - def test_execute_with_oom_recovery_max_retries_zero(self): - """Test _execute_with_oom_recovery with max_retries=0.""" - from arraybridge.oom_recovery import _execute_with_oom_recovery - - def always_fails_func(): - raise MemoryError("out of memory") - - with pytest.raises(MemoryError): - _execute_with_oom_recovery(always_fails_func, "torch", max_retries=0) - - @pytest.mark.parametrize("memory_type", ["torch", "cupy", "tensorflow"]) - def test_execute_with_oom_recovery_different_frameworks(self, memory_type): - """Test _execute_with_oom_recovery with different GPU frameworks.""" - from arraybridge.oom_recovery import _execute_with_oom_recovery - - call_count = {"count": 0} - - def failing_then_success_func(): - call_count["count"] += 1 - if call_count["count"] == 1: - raise MemoryError("out of memory") - return f"success_{memory_type}" - - result = _execute_with_oom_recovery(failing_then_success_func, memory_type, max_retries=1) - assert result == f"success_{memory_type}" - assert call_count["count"] == 2 - - -def test_torch_cache_clear_mock(): - """Test torch cache clear with mocked torch.""" diff --git a/tests/test_registry_integration.py b/tests/test_registry_integration.py deleted file mode 100644 index c0737f7..0000000 --- a/tests/test_registry_integration.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Integration tests demonstrating metaclass-registry benefits.""" - -import numpy as np - - -class TestRegistryIntegration: - """Integration tests showing how the registry simplifies converter management.""" - - def test_registry_discoverability(self): - """Test that all converters are discoverable via the registry.""" - from arraybridge.converters_registry import ConverterBase - - # Registry makes it easy to discover all available converters - available_converters = sorted(ConverterBase.__registry__.keys()) - - assert len(available_converters) == 6 - assert available_converters == [ - "cupy", - "jax", - "numpy", - "pyclesperanto", - "tensorflow", - "torch", - ] - - def test_registry_enables_programmatic_access(self): - """Test that registry enables programmatic access to all converters.""" - from arraybridge.converters_registry import ConverterBase, get_converter - - # Can iterate over all registered converters - for memory_type, converter_class in ConverterBase.__registry__.items(): - converter = get_converter(memory_type) - - # Verify each converter has the expected interface - assert hasattr(converter, "to_numpy") - assert hasattr(converter, "from_numpy") - assert hasattr(converter, "from_dlpack") - assert hasattr(converter, "move_to_device") - - # Verify memory_type matches - assert converter.memory_type == memory_type - - def test_memory_type_enum_integration(self): - """Test that MemoryType enum integrates seamlessly with registry.""" - from arraybridge.converters_registry import get_converter - from arraybridge.types import MemoryType - - np.array([1, 2, 3, 4, 5]) - - # The enum declares identity; the registry owns converter resolution. - for mem_type in MemoryType: - converter = get_converter(mem_type.value) - assert converter.memory_type == mem_type.value - - def test_convert_memory_uses_registry(self): - """Test that convert_memory function uses registry-based converters.""" - from arraybridge.converters import convert_memory - - arr = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) - - # convert_memory should work with registry - result = convert_memory(arr, source_type="numpy", target_type="numpy", gpu_id=0) - - assert isinstance(result, np.ndarray) - np.testing.assert_array_almost_equal(result, arr) - - def test_registry_validation_on_import(self): - """Test that registry validates all memory types are registered on import.""" - from arraybridge.converters_registry import ConverterBase - from arraybridge.types import MemoryType - - # Registry should contain exactly the memory types defined in MemoryType enum - expected = {mt.value for mt in MemoryType} - actual = set(ConverterBase.__registry__.keys()) - - assert ( - expected == actual - ), f"Registry validation failed. Expected: {expected}, Got: {actual}" - - def test_declarations_generate_converter_registration(self): - """Every declaration automatically generates one registered converter.""" - from arraybridge.converters_registry import ConverterBase - from arraybridge.types import MemoryType - - current_count = len(ConverterBase.__registry__) - assert current_count == len(MemoryType) - - -class TestConverterIsolation: - """Test that converters are properly isolated and independent.""" - - def test_converters_are_independent_instances(self): - """Test that multiple calls to get_converter return independent instances.""" - from arraybridge.converters_registry import get_converter - - # Each call should return a new instance - conv1 = get_converter("numpy") - conv2 = get_converter("numpy") - - assert conv1 is not conv2 - assert isinstance(conv1, type(conv2)) - assert conv1.memory_type == conv2.memory_type - - def test_converter_classes_are_registered_not_instances(self): - """Test that registry stores classes, not instances.""" - from arraybridge.converters_registry import ConverterBase, get_converter - - # Registry should contain classes - numpy_class = ConverterBase.__registry__["numpy"] - assert isinstance(numpy_class, type) - - # get_converter creates instances - instance = get_converter("numpy") - assert isinstance(instance, numpy_class) diff --git a/tests/test_slice_processing.py b/tests/test_slice_processing.py index 174e890..bb52bab 100644 --- a/tests/test_slice_processing.py +++ b/tests/test_slice_processing.py @@ -1,16 +1,20 @@ """Tests for arraybridge.slice_processing module.""" +import sys + import numpy as np import pytest +from arraybridge.devices import AcceleratorDevice +from arraybridge.slice_processing import process_slices +from tests.fakes import FakeCupy, FakeCupyArray + class TestProcessSlices: """Tests for process_slices function.""" def test_process_slices_single_output(self): """Test process_slices with function returning single output.""" - from arraybridge.slice_processing import process_slices - # Create a 3D array (2 slices of 2x2) image_3d = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) @@ -25,8 +29,6 @@ def double_func(slice_2d): def test_process_slices_tuple_output(self): """Test process_slices with function returning tuple (main + special outputs).""" - from arraybridge.slice_processing import process_slices - # Create a 3D array image_3d = np.array([[[1, 2]], [[3, 4]]]) @@ -49,8 +51,6 @@ def func_with_special(slice_2d): def test_process_slices_multiple_special_outputs(self): """Test process_slices with function returning multiple special outputs.""" - from arraybridge.slice_processing import process_slices - image_3d = np.array([[[1, 2]], [[3, 4]]]) def func_multiple_special(slice_2d): @@ -70,8 +70,6 @@ def func_multiple_special(slice_2d): def test_process_slices_with_args_kwargs(self): """Test process_slices passing additional args and kwargs to function.""" - from arraybridge.slice_processing import process_slices - image_3d = np.array([[[1]], [[2]]]) def func_with_args_kwargs(slice_2d, multiplier, offset=0): @@ -85,8 +83,6 @@ def func_with_args_kwargs(slice_2d, multiplier, offset=0): def test_process_slices_empty_special_outputs(self): """Reject a tuple/non-tuple result boundary between slices.""" - from arraybridge.slice_processing import process_slices - image_3d = np.array([[[1]], [[2]]]) # Mix of single output and tuple output @@ -102,8 +98,6 @@ def mixed_func(slice_2d): def test_process_slices_rejects_result_tuple_arity_drift(self): """Reject side-output cardinality changes between slices.""" - from arraybridge.slice_processing import process_slices - image_3d = np.array([[[1]], [[2]]]) def drifting_func(slice_2d): @@ -117,8 +111,6 @@ def drifting_func(slice_2d): def test_process_slices_rejects_empty_result_tuple(self): """Reject a tuple that has no declared main output.""" - from arraybridge.slice_processing import process_slices - with pytest.raises(ValueError, match="cannot be empty"): process_slices(np.ones((1, 1, 1)), lambda _slice: (), (), {}) @@ -132,8 +124,6 @@ def test_process_slices_rejects_empty_result_tuple(self): ) def test_process_slices_different_shapes(self, shape): """Test process_slices with different 3D array shapes.""" - from arraybridge.slice_processing import process_slices - image_3d = np.random.rand(*shape) def identity_func(slice_2d): @@ -143,3 +133,24 @@ def identity_func(slice_2d): # Should return the same array np.testing.assert_array_equal(result, image_3d) + + @pytest.mark.parametrize( + ("device", "seen", "result_index"), + [(None, [0, 0], 0), (AcceleratorDevice(1), [1, 1], 1)], + ids=["device of the image", "given device"], + ) + def test_process_slices_unstacks_on_the_device(self, device, seen, result_index, monkeypatch): + module = FakeCupy() + monkeypatch.setitem(sys.modules, "cupy", module) + image = FakeCupyArray(np.arange(8).reshape(2, 2, 2), 0, module) + slice_devices = [] + + def record_device(slice_2d): + slice_devices.append(slice_2d.device.id) + return slice_2d + + result = process_slices(image, record_device, (), {}, device=device) + + assert slice_devices == seen + assert result.device.id == result_index + np.testing.assert_array_equal(result.values, image.values) diff --git a/tests/test_stack_utils.py b/tests/test_stack_utils.py index 863b8db..3893f64 100644 --- a/tests/test_stack_utils.py +++ b/tests/test_stack_utils.py @@ -1,381 +1,132 @@ -"""Tests for arraybridge.stack_utils module.""" +"""Stacking 2D slices into a 3D array of a named library on a device, and back.""" import sys -from types import SimpleNamespace import numpy as np import pytest - -class _FakeDeviceScope: - def __init__(self, state, device_id): - self._state = state - self._device_id = device_id - self._previous = None - - def __enter__(self): - self._previous = self._state["device"] - self._state["device"] = self._device_id - return self - - def __exit__(self, exc_type, exc_value, traceback): - self._state["device"] = self._previous - return False - - -class _FakeCupyArray: - __module__ = "cupy" - - def __init__(self, values, device_id, state): - self.values = np.asarray(values) - self.device = SimpleNamespace(id=device_id) - self._state = state - - @property - def shape(self): - return self.values.shape - - @property - def dtype(self): - return self.values.dtype - - def copy(self): - return type(self)(self.values.copy(), self._state["device"], self._state) - - def __getitem__(self, index): - return type(self)(self.values[index], self.device.id, self._state) - - def __setitem__(self, index, value): - self.values[index] = value.values - - -def _fake_cupy_module(state): - runtime = SimpleNamespace( - getDeviceCount=lambda: 2, - getDevice=lambda: state["device"], - ) - cuda = SimpleNamespace( - runtime=runtime, - Device=lambda device_id: _FakeDeviceScope(state, device_id), - ) - return SimpleNamespace( - cuda=cuda, - stack=lambda values, axis: _FakeCupyArray( - np.stack([value.values for value in values], axis=axis), - state["device"], - state, - ), - ) - - -class _FakeTensorFlowTensor: - __module__ = "tensorflow" - - def __init__(self, values, device_id, state): - self.values = np.asarray(values) - self.device = f"/device:GPU:{device_id}" - self._state = state - - @property - def shape(self): - return self.values.shape - - @property - def dtype(self): - return self.values.dtype - - def __getitem__(self, index): - return type(self)(self.values[index], int(self.device.rsplit(":", 1)[-1]), self._state) - - def numpy(self): - return self.values - - -def _fake_tensorflow_module(state): - def device_scope(name): - return _FakeDeviceScope(state, int(name.rsplit(":", 1)[-1])) - - return SimpleNamespace( - config=SimpleNamespace( - list_logical_devices=lambda kind: [object(), object()] if kind == "GPU" else [] - ), - device=device_scope, - identity=lambda data: _FakeTensorFlowTensor(data.values, state["device"], state), - stack=lambda values, axis: _FakeTensorFlowTensor( - np.stack([value.values for value in values], axis=axis), - state["device"], - state, - ), - ) - - -class TestStackUtils: - """Tests for stack utilities functions.""" - - def test_is_2d_numpy(self): - """Test _is_2d with numpy arrays.""" - from arraybridge.stack_utils import _is_2d - - # 2D array - arr_2d = np.array([[1, 2], [3, 4]]) - assert _is_2d(arr_2d) is True - - # 1D array - arr_1d = np.array([1, 2, 3]) - assert _is_2d(arr_1d) is False - - # 3D array - arr_3d = np.array([[[1, 2]], [[3, 4]]]) - assert _is_2d(arr_3d) is False - - def test_is_3d_numpy(self): - """Test _is_3d with numpy arrays.""" - from arraybridge.stack_utils import _is_3d - - # 3D array - arr_3d = np.array([[[1, 2]], [[3, 4]]]) - assert _is_3d(arr_3d) is True - - # 2D array - arr_2d = np.array([[1, 2], [3, 4]]) - assert _is_3d(arr_2d) is False - - # 1D array - arr_1d = np.array([1, 2, 3]) - assert _is_3d(arr_1d) is False - - def test_enforce_gpu_device_requirements_valid(self, monkeypatch): - """Test _enforce_gpu_device_requirements with valid inputs.""" - from arraybridge.stack_utils import _enforce_gpu_device_requirements - - # CPU memory type should not raise - _enforce_gpu_device_requirements("numpy", 0) - - torch = SimpleNamespace( - cuda=SimpleNamespace(is_available=lambda: True, device_count=lambda: 2) - ) - monkeypatch.setitem(sys.modules, "torch", torch) - - _enforce_gpu_device_requirements("torch", 0) - _enforce_gpu_device_requirements("torch", 1) - - def test_enforce_gpu_device_requirements_invalid_gpu_id(self, monkeypatch): - """Test _enforce_gpu_device_requirements with invalid GPU device ID.""" - from arraybridge.stack_utils import _enforce_gpu_device_requirements - - torch = SimpleNamespace( - cuda=SimpleNamespace(is_available=lambda: True, device_count=lambda: 2) - ) - monkeypatch.setitem(sys.modules, "torch", torch) - - with pytest.raises(ValueError, match="device -1 is unavailable"): - _enforce_gpu_device_requirements("torch", -1) - - with pytest.raises(ValueError, match="device 2 is unavailable"): - _enforce_gpu_device_requirements("torch", 2) - - def test_cpu_only_jax_cannot_satisfy_stack_device_requirement(self, monkeypatch): - from arraybridge.stack_utils import _enforce_gpu_device_requirements - - cpu_device = SimpleNamespace(platform="cpu") - monkeypatch.setitem(sys.modules, "jax", SimpleNamespace(devices=lambda: [cpu_device])) - - with pytest.raises(ValueError, match="available device IDs are \\(\\)"): - _enforce_gpu_device_requirements("jax", 0) - - def test_stack_slices_empty_list(self): - """Test stack_slices with empty list raises error.""" - from arraybridge.stack_utils import stack_slices - - with pytest.raises(ValueError) as exc_info: - stack_slices([], "numpy", 0) - assert "Cannot stack empty list" in str(exc_info.value) - - def test_stack_slices_not_2d(self): - """Test stack_slices with non-2D slices raises error.""" - from arraybridge.stack_utils import stack_slices - - slices = [np.array([1, 2, 3]), np.array([4, 5, 6])] # 1D arrays - - with pytest.raises(ValueError) as exc_info: - stack_slices(slices, "numpy", 0) - assert "not a 2D array" in str(exc_info.value) - - def test_stack_slices_numpy(self): - """Test stack_slices with numpy arrays.""" - from arraybridge.stack_utils import stack_slices - - slice1 = np.array([[1, 2], [3, 4]]) - slice2 = np.array([[5, 6], [7, 8]]) - slices = [slice1, slice2] - - result = stack_slices(slices, "numpy", 0) - - assert result.shape == (2, 2, 2) - expected = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) - np.testing.assert_array_equal(result, expected) - - def test_stack_slices_single_slice(self): - """Test stack_slices with single slice.""" - from arraybridge.stack_utils import stack_slices - - slice1 = np.array([[1, 2, 3], [4, 5, 6]]) - result = stack_slices([slice1], "numpy", 0) - - assert result.shape == (1, 2, 3) - expected = np.array([[[1, 2, 3], [4, 5, 6]]]) - np.testing.assert_array_equal(result, expected) - - def test_unstack_slices_not_3d(self): - """Test unstack_slices with non-3D array raises error.""" - from arraybridge.stack_utils import unstack_slices - - arr_2d = np.array([[1, 2], [3, 4]]) - - with pytest.raises(ValueError) as exc_info: - unstack_slices(arr_2d, "numpy", 0) - assert "Array must be 3D" in str(exc_info.value) - - def test_unstack_slices_numpy(self): - """Test unstack_slices with numpy array.""" - from arraybridge.stack_utils import unstack_slices - - arr_3d = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) - result = unstack_slices(arr_3d, "numpy", 0) - - assert len(result) == 2 - assert result[0].shape == (2, 2) - assert result[1].shape == (2, 2) - - np.testing.assert_array_equal(result[0], np.array([[1, 2], [3, 4]])) - np.testing.assert_array_equal(result[1], np.array([[5, 6], [7, 8]])) - - def test_unstack_slices_single_slice(self): - """Test unstack_slices with single slice.""" - from arraybridge.stack_utils import unstack_slices - - arr_3d = np.array([[[1, 2, 3], [4, 5, 6]]]) - result = unstack_slices(arr_3d, "numpy", 0) - - assert len(result) == 1 - assert result[0].shape == (2, 3) - np.testing.assert_array_equal(result[0], np.array([[1, 2, 3], [4, 5, 6]])) - - def test_unstack_slices_validate_slices_false(self): - """Test unstack_slices with validate_slices=False.""" - from arraybridge.stack_utils import unstack_slices - - arr_3d = np.array([[[1, 2]], [[3, 4]]]) # Shape: (2, 1, 2) - result = unstack_slices(arr_3d, "numpy", 0, validate_slices=False) - - assert len(result) == 2 - assert result[0].shape == (1, 2) # Each slice has shape (1, 2) - assert result[1].shape == (1, 2) - - def test_numpy_stack_unstack_roundtrip(self): - """Test roundtrip: stack_slices -> unstack_slices.""" - from arraybridge.stack_utils import stack_slices, unstack_slices - - # Create test slices - slice1 = np.array([[1, 2], [3, 4]]) - slice2 = np.array([[5, 6], [7, 8]]) - original_slices = [slice1, slice2] - - # Stack them - stacked = stack_slices(original_slices, "numpy", 0) - - # Unstack them - unstaked = unstack_slices(stacked, "numpy", 0) - - # Verify roundtrip - assert len(unstaked) == len(original_slices) - for original, result in zip(original_slices, unstaked): - np.testing.assert_array_equal(original, result) - - def test_same_framework_stack_and_unstack_honor_requested_device(self, monkeypatch): - from arraybridge.stack_utils import stack_slices, unstack_slices - - state = {"device": 0} - monkeypatch.setitem(sys.modules, "cupy", _fake_cupy_module(state)) - inputs = [ - _FakeCupyArray([[1, 2], [3, 4]], 0, state), - _FakeCupyArray([[5, 6], [7, 8]], 0, state), - ] - - stacked = stack_slices(inputs, "cupy", 1) - slices = unstack_slices(stacked, "cupy", 1) +from arraybridge.devices import HOST, AcceleratorDevice +from arraybridge.libraries import ArrayLibrary +from arraybridge.stack_utils import _is_2d, _is_3d, stack_slices, unstack_slices +from tests.fakes import ( + FakeCupy, + FakeCupyArray, + FakePyclesperanto, + FakePyclesperantoArray, + FakeTensorflow, + FakeTensorflowTensor, + FakeTorch, + FakeTorchTensor, +) + +PLANES = [np.array([[1, 2], [3, 4]]), np.array([[5, 6], [7, 8]])] + + +@pytest.mark.parametrize("rank", [1, 2, 3]) +def test_rank_checks(rank): + values = np.zeros((2,) * rank) + + assert _is_2d(values) is (rank == 2) + assert _is_3d(values) is (rank == 3) + assert not _is_2d("no shape") and not _is_3d("no shape") + + +class TestNumpy: + @pytest.mark.parametrize("planes", [PLANES, PLANES[:1]], ids=["two", "one"]) + def test_stack_and_unstack_round_trip(self, planes): + stacked = stack_slices(planes, "numpy", HOST) + unstacked = unstack_slices(stacked, "numpy", HOST) + + assert stacked.shape == (len(planes), 2, 2) + np.testing.assert_array_equal(stacked, np.stack(planes)) + assert len(unstacked) == len(planes) + for original, result in zip(planes, unstacked, strict=True): + assert result.shape == (2, 2) + np.testing.assert_array_equal(result, original) + + def test_unstack_can_skip_the_slice_check(self): + slices = unstack_slices(np.zeros((2, 1, 2)), "numpy", HOST, validate_slices=False) + + assert [each.shape for each in slices] == [(1, 2), (1, 2)] + + def test_malformed_input_is_rejected(self): + with pytest.raises(ValueError, match="Cannot stack empty list"): + stack_slices([], "numpy", HOST) + with pytest.raises(ValueError, match="Slice at index 0 is not a 2D array"): + stack_slices([np.array([1, 2, 3])], "numpy", HOST) + with pytest.raises(ValueError, match="Array must be 3D"): + unstack_slices(PLANES[0], "numpy", HOST) + + def test_unregistered_library_and_missing_device_are_rejected(self): + with pytest.raises(ValueError, match="not a registered array library"): + stack_slices(PLANES, "invalid_type", HOST) + with pytest.raises(ValueError, match="NumPy device accelerator:0 is unavailable"): + stack_slices(PLANES, "numpy", AcceleratorDevice(0)) + with pytest.raises(ValueError, match="NumPy device accelerator:0 is unavailable"): + unstack_slices(np.stack(PLANES), "numpy", AcceleratorDevice(0)) + + +class TestOtherLibraries: + def test_stack_and_unstack_move_to_the_requested_device(self, monkeypatch): + module = FakeCupy() + monkeypatch.setitem(sys.modules, "cupy", module) + inputs = [FakeCupyArray(plane, 0, module) for plane in PLANES] + + stacked = stack_slices(inputs, "cupy", AcceleratorDevice(1)) + slices = unstack_slices(stacked, "cupy", AcceleratorDevice(1)) assert stacked.device.id == 1 - assert [slice_.device.id for slice_ in slices] == [1, 1] - np.testing.assert_array_equal(stacked.values, np.stack([item.values for item in inputs])) - assert state["device"] == 0 - - def test_tensorflow_stack_uses_immutable_framework_leaf(self, monkeypatch): - from arraybridge.stack_utils import stack_slices, unstack_slices + assert [each.device.id for each in slices] == [1, 1] + np.testing.assert_array_equal(stacked.values, np.stack(PLANES)) + assert module.current == 0 - state = {"device": 0} - monkeypatch.setitem(sys.modules, "tensorflow", _fake_tensorflow_module(state)) - inputs = [ - _FakeTensorFlowTensor([[1, 2], [3, 4]], 0, state), - _FakeTensorFlowTensor([[5, 6], [7, 8]], 0, state), - ] + def test_tensorflow_tensors_are_copied_to_the_device_not_changed(self, monkeypatch): + module = FakeTensorflow() + monkeypatch.setitem(sys.modules, "tensorflow", module) + inputs = [FakeTensorflowTensor(plane, 0) for plane in PLANES] - stacked = stack_slices(inputs, "tensorflow", 1) - slices = unstack_slices(stacked, "tensorflow", 1) + stacked = stack_slices(inputs, "tensorflow", AcceleratorDevice(1)) + slices = unstack_slices(stacked, "tensorflow", AcceleratorDevice(1)) assert stacked.device == "/device:GPU:1" - assert [slice_.device for slice_ in slices] == [ - "/device:GPU:1", - "/device:GPU:1", - ] - np.testing.assert_array_equal(stacked.values, np.stack([item.values for item in inputs])) - assert state["device"] == 0 - - def test_pyclesperanto_stack_uses_binary_api_and_preserves_singleton_axis(self): - from arraybridge.array_operations import ArrayOperations - from arraybridge.types import MemoryType - - operations = ArrayOperations.for_memory(MemoryType.PYCLESPERANTO) - - class Array: - def __init__(self, values): - self.values = np.asarray(values) - self.shape = self.values.shape - self.dtype = self.values.dtype - - class Module: - calls = [] - - @classmethod - def create(cls, shape, dtype): - cls.calls.append(("create", shape)) - return Array(np.empty(shape, dtype=dtype)) + assert [each.device for each in slices] == ["/device:GPU:1", "/device:GPU:1"] + assert [each.device for each in inputs] == ["/device:GPU:0", "/device:GPU:0"] + np.testing.assert_array_equal(stacked.values, np.stack(PLANES)) + assert module.current == 0 + + def test_slices_of_another_library_are_converted_first(self, monkeypatch): + module = FakeTorch() + module.stack = lambda values, dim: FakeTorchTensor( + np.stack([each.values for each in values], axis=dim) + ) + monkeypatch.setitem(sys.modules, "torch", module) - @classmethod - def copy_slice(cls, source, target, index): - cls.calls.append(("copy_slice", index)) - target.values[index] = source.values - return target + stacked = stack_slices([PLANES[0], FakeTorchTensor(PLANES[1])], "torch", HOST) + as_numpy = unstack_slices(stacked, "numpy", HOST) - @classmethod - def concatenate_along_z(cls, left, right): - cls.calls.append(("concatenate", left.shape, right.shape)) - left_values = left.values if left.values.ndim == 3 else left.values[None] - right_values = right.values if right.values.ndim == 3 else right.values[None] - return Array(np.concatenate((left_values, right_values), axis=0)) + assert isinstance(stacked, FakeTorchTensor) + np.testing.assert_array_equal(stacked.values, np.stack(PLANES)) + np.testing.assert_array_equal(as_numpy, PLANES) - first = Array([[1, 2], [3, 4]]) - second = Array([[5, 6], [7, 8]]) - third = Array([[9, 10], [11, 12]]) + def test_pyclesperanto_stacks_pairwise_and_keeps_a_single_plane_3d(self, monkeypatch): + module = FakePyclesperanto() + monkeypatch.setitem(sys.modules, "pyclesperanto", module) + library = ArrayLibrary.named("pyclesperanto") + first, second, third = ( + FakePyclesperantoArray(plane, "gpu0") for plane in (*PLANES, PLANES[0] + 8) + ) - singleton = operations.stack([first], Module) - stack = operations.stack([first, second, third], Module) + single = library.stack_arrays([first], AcceleratorDevice(0)) + stack = library.stack_arrays([first, second, third], AcceleratorDevice(0)) - assert singleton.shape == (1, 2, 2) + assert single.shape == (1, 2, 2) assert stack.shape == (3, 2, 2) - assert Module.calls == [ + np.testing.assert_array_equal(stack.values[2], third.values) + assert module.events == [ ("create", (1, 2, 2)), ("copy_slice", 0), ("concatenate", (2, 2), (2, 2)), ("concatenate", (2, 2, 2), (2, 2)), ] + with pytest.raises(ValueError, match="empty pyclesperanto sequence"): + library.stack_arrays([], AcceleratorDevice(0)) diff --git a/tests/test_toy_library.py b/tests/test_toy_library.py new file mode 100644 index 0000000..6d35e1e --- /dev/null +++ b/tests/test_toy_library.py @@ -0,0 +1,149 @@ +"""A library is one class: writing it is all that adding a library takes.""" + +from contextlib import contextmanager + +import numpy as np +import pytest +from metaclass_registry.testing import restored_registry + +import arraybridge +from arraybridge import decorators +from arraybridge.conversion import ConversionRoute +from arraybridge.converters import convert_memory, detect_memory_type +from arraybridge.devices import HOST +from arraybridge.libraries import ArrayLibrary, HostDevices +from arraybridge.types import MemoryType + + +class ToyArray: + """The smallest array: NumPy values behind a class of this module.""" + + def __init__(self, values): + self.values = np.asarray(values) + + @property + def shape(self): + return self.values.shape + + @property + def dtype(self): + return self.values.dtype + + +@contextmanager +def toy_library_defined(): + """Define ``ToyLibrary`` for a block, then remove every trace of it. + + A registered class stays registered for the life of the process, and each + thing derived from it is kept where it was first asked for, so all of + those are undone here. + """ + with restored_registry(ArrayLibrary): + + class ToyLibrary(HostDevices[ToyArray]): + name = "toy" + import_name = __name__.split(".")[0] + + def device_of(self, data): + return HOST + + def _to_host(self, data): + return data.values + + def _wrap_host_array(self, data): + return ToyArray(data) + + def _move_to_host(self, data): + return data + + def _scale_dtype(self, result, target_dtype): + return ToyArray(result.values.astype(target_dtype)) + + try: + yield ToyLibrary + finally: + ArrayLibrary._instances.pop("toy", None) + MemoryType._value2member_map_.pop("toy", None) + if MemoryType._member_map_.pop("TOY", None) is not None: + MemoryType._member_names_.remove("TOY") + vars(decorators).pop("toy", None) + vars(arraybridge).pop("toy", None) + + +@pytest.fixture +def toy_library_class(): + with toy_library_defined() as library_class: + yield library_class + + +def test_defining_the_class_is_the_whole_registration(toy_library_class): + from arraybridge.decorators import toy + + library = ArrayLibrary.named("toy") + host_values = np.arange(6, dtype=np.uint16).reshape(2, 3) + + assert type(library) is toy_library_class + assert ArrayLibrary.all()[-1] is library + assert ArrayLibrary.for_module_name(__name__) is library + assert library.display_name == "toy" + assert library.available_devices() == (HOST,) + + assert MemoryType("toy") is MemoryType.TOY + assert MemoryType.TOY in list(MemoryType) + assert (MemoryType.TOY.name, MemoryType.TOY.value) == ("TOY", "toy") + assert MemoryType.TOY.library is library + + toy_values = convert_memory(host_values, "numpy", "toy", HOST) + assert isinstance(toy_values, ToyArray) + assert ArrayLibrary.of(toy_values) is library + assert isinstance(toy_values, library.array_class) + assert detect_memory_type(toy_values) == "toy" + np.testing.assert_array_equal(toy_values.values, host_values) + np.testing.assert_array_equal( + convert_memory(toy_values, MemoryType.TOY, MemoryType.NUMPY, HOST), host_values + ) + + numpy_library = ArrayLibrary.named("numpy") + assert ConversionRoute.planned(numpy_library, HOST, library, HOST).mechanism == "host copy" + assert ConversionRoute.planned(library, HOST, numpy_library, HOST).mechanism == "host copy" + assert ConversionRoute.planned(library, HOST, library, HOST).mechanism == "device move" + + @toy + def halve(image): + return ToyArray(image.values / 2) + + assert toy is arraybridge.toy + assert toy.__name__ == "toy" + assert halve.input_memory_type == "toy" + assert halve.output_memory_type == "toy" + assert halve.execution_memory_type == "toy" + halved = halve(ToyArray(np.array([0, 2, 4], dtype=np.uint16))) + assert isinstance(halved, ToyArray) + assert halved.dtype == np.uint16 + np.testing.assert_array_equal(halved.values, [0, 1, 2]) + + +def test_name_lookup_sees_a_library_defined_after_import(toy_library_class): + # The first look at the enum after the class statement is by member name. + assert MemoryType["TOY"].value == "toy" + + +def test_nothing_of_the_class_is_left_afterwards(): + registered = list(ArrayLibrary.__registry__) + members = list(MemoryType) + + with toy_library_defined(): + assert MemoryType.TOY.library is ArrayLibrary.named("toy") + assert callable(arraybridge.toy) and callable(decorators.toy) + + assert list(ArrayLibrary.__registry__) == registered + assert [library.name for library in ArrayLibrary.all()] == registered + assert list(MemoryType) == members + assert list(MemoryType.__members__) == [member.name for member in members] + for lookup in (lambda: MemoryType("toy"), lambda: ArrayLibrary.named("toy")): + with pytest.raises(ValueError): + lookup() + for owner in (MemoryType, decorators, arraybridge): + assert not hasattr(owner, "TOY") and not hasattr(owner, "toy") + with pytest.raises(ValueError, match="Unknown memory type"): + ArrayLibrary.of(ToyArray([1])) diff --git a/tests/test_types.py b/tests/test_types.py deleted file mode 100644 index c277fe4..0000000 --- a/tests/test_types.py +++ /dev/null @@ -1,498 +0,0 @@ -"""Tests for arraybridge.types module.""" - -import os -from contextlib import contextmanager -from types import SimpleNamespace - -import numpy as np -import pytest - -from arraybridge.types import ( - CPU_MEMORY_TYPES, - GPU_MEMORY_TYPES, - SUPPORTED_MEMORY_TYPES, - VALID_MEMORY_TYPES, - DLPackPayload, - MemoryContractAttribute, - MemoryType, -) - - -class TestMemoryContractAttribute: - """Tests for the declaration-owned callable metadata protocol.""" - - def test_reads_and_writes_object_namespaces(self): - namespace = SimpleNamespace() - - MemoryContractAttribute.EXECUTION.write(namespace, "torch") - - assert MemoryContractAttribute.EXECUTION.read(namespace) == "torch" - - def test_reads_and_writes_mapping_namespaces(self): - namespace = {} - - MemoryContractAttribute.INPUT.write(namespace, "numpy") - - assert MemoryContractAttribute.INPUT.read(namespace) == "numpy" - assert MemoryContractAttribute.OUTPUT.read(namespace) is None - - -class TestMemoryType: - """Tests for MemoryType enum.""" - - def test_memory_type_values(self): - """Test that MemoryType enum has expected values.""" - assert MemoryType.NUMPY.value == "numpy" - assert MemoryType.CUPY.value == "cupy" - assert MemoryType.TORCH.value == "torch" - assert MemoryType.TENSORFLOW.value == "tensorflow" - assert MemoryType.JAX.value == "jax" - assert MemoryType.PYCLESPERANTO.value == "pyclesperanto" - - def test_memory_type_from_string(self): - """Test creating MemoryType from string.""" - assert MemoryType("numpy") == MemoryType.NUMPY - assert MemoryType("torch") == MemoryType.TORCH - assert MemoryType("cupy") == MemoryType.CUPY - - def test_invalid_memory_type_raises_error(self): - """Test that invalid memory type raises ValueError.""" - with pytest.raises(ValueError): - MemoryType("invalid_type") - - def test_cpu_memory_types(self): - """Test that CPU_MEMORY_TYPES contains only NumPy.""" - assert CPU_MEMORY_TYPES == {MemoryType.NUMPY} - assert isinstance(CPU_MEMORY_TYPES, frozenset) - assert len(CPU_MEMORY_TYPES) == 1 - - def test_gpu_memory_types(self): - """Test that GPU_MEMORY_TYPES contains all GPU frameworks.""" - expected_gpu_types = { - MemoryType.CUPY, - MemoryType.TORCH, - MemoryType.TENSORFLOW, - MemoryType.JAX, - MemoryType.PYCLESPERANTO, - } - assert GPU_MEMORY_TYPES == expected_gpu_types - assert isinstance(GPU_MEMORY_TYPES, frozenset) - assert len(GPU_MEMORY_TYPES) == 5 - - def test_supported_memory_types(self): - """Test that SUPPORTED_MEMORY_TYPES contains all types.""" - assert SUPPORTED_MEMORY_TYPES == CPU_MEMORY_TYPES | GPU_MEMORY_TYPES - assert len(SUPPORTED_MEMORY_TYPES) == 6 - - def test_valid_memory_types_strings(self): - """Test that VALID_MEMORY_TYPES contains string values.""" - expected_strings = {"numpy", "cupy", "torch", "tensorflow", "jax", "pyclesperanto"} - assert VALID_MEMORY_TYPES == expected_strings - assert isinstance(VALID_MEMORY_TYPES, frozenset) - - -class TestMemoryTypeOwnership: - """Tests for declaration-owned framework and device capabilities.""" - - def test_array_operations_are_owned_by_enum_members(self): - assert not hasattr(MemoryType.NUMPY, "converter") - assert callable(MemoryType.NUMPY.to_numpy) - assert callable(MemoryType.NUMPY.from_numpy) - assert callable(MemoryType.NUMPY.stack_arrays) - assert callable(MemoryType.NUMPY.scale_dtype) - - def test_framework_identity_is_owned_by_members(self): - assert MemoryType.NUMPY.import_name == "numpy" - assert MemoryType.NUMPY.display_name == "NumPy" - assert MemoryType.NUMPY.is_gpu is False - assert MemoryType.JAX.recognized_module_names == frozenset({"jax", "jaxlib"}) - assert all(memory_type.is_gpu for memory_type in GPU_MEMORY_TYPES) - - def test_absent_optional_framework_is_not_imported(self, monkeypatch): - from arraybridge import types - - monkeypatch.delitem(types.sys.modules, "cupy", raising=False) - monkeypatch.setattr(types.importlib.util, "find_spec", lambda name: None) - - def fail_import(name): - raise AssertionError(f"unexpected import of {name}") - - monkeypatch.setattr(types.importlib, "import_module", fail_import) - - assert MemoryType.CUPY.import_if_installed() is None - assert MemoryType.CUPY.available_device_ids() == () - MemoryType.CUPY.cleanup_loaded() - - def test_absent_framework_does_not_apply_import_environment(self, monkeypatch): - from arraybridge import types - - monkeypatch.delenv("XLA_PYTHON_CLIENT_PREALLOCATE", raising=False) - monkeypatch.delitem(types.sys.modules, "jax", raising=False) - monkeypatch.setattr(types.importlib.util, "find_spec", lambda name: None) - - assert MemoryType.JAX.import_if_installed() is None - assert "XLA_PYTHON_CLIENT_PREALLOCATE" not in os.environ - - def test_import_preparation_sets_only_declaration_owned_defaults(self, monkeypatch): - monkeypatch.delenv("TF_FORCE_GPU_ALLOW_GROWTH", raising=False) - monkeypatch.setenv("XLA_PYTHON_CLIENT_PREALLOCATE", "custom") - - MemoryType.TENSORFLOW.prepare_import() - MemoryType.JAX.prepare_import() - - assert MemoryType.TENSORFLOW.import_environment == (("TF_FORCE_GPU_ALLOW_GROWTH", "true"),) - assert MemoryType.JAX.import_environment == (("XLA_PYTHON_CLIENT_PREALLOCATE", "false"),) - assert os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] == "true" - assert os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] == "custom" - - def test_subprocess_environment_derives_nvidia_wheel_paths_from_cupy_owner( - self, - monkeypatch, - tmp_path, - ): - from arraybridge import types - - nvidia_root = tmp_path / "nvidia" - cublas_lib = nvidia_root / "cublas" / "lib" - runtime_bin = nvidia_root / "cuda_runtime" / "bin" - cublas_lib.mkdir(parents=True) - runtime_bin.mkdir(parents=True) - monkeypatch.setattr( - types.importlib.util, - "find_spec", - lambda name: ( - SimpleNamespace(submodule_search_locations=(str(nvidia_root),)) - if name == "nvidia" - else None - ), - ) - search_variable = "PATH" if os.name == "nt" else "LD_LIBRARY_PATH" - original = {search_variable: "/existing", "CUSTOM": "value"} - - prepared = MemoryType.subprocess_environment(original) - - assert original == {search_variable: "/existing", "CUSTOM": "value"} - assert prepared[search_variable].split(os.pathsep) == [ - *sorted((str(cublas_lib), str(runtime_bin))), - "/existing", - ] - assert prepared["CUSTOM"] == "value" - assert prepared["TF_FORCE_GPU_ALLOW_GROWTH"] == "true" - assert prepared["XLA_PYTHON_CLIENT_PREALLOCATE"] == "false" - - def test_subprocess_environment_preserves_host_values_without_nvidia_wheels( - self, - monkeypatch, - ): - from arraybridge import types - - monkeypatch.setattr(types.importlib.util, "find_spec", lambda name: None) - original = { - "TF_FORCE_GPU_ALLOW_GROWTH": "host", - "XLA_PYTHON_CLIENT_PREALLOCATE": "host", - } - - prepared = MemoryType.subprocess_environment(original) - - assert prepared == original - assert prepared is not original - - def test_preloaded_framework_reports_late_import_preparation( - self, - monkeypatch, - caplog, - ): - from arraybridge import types - - module = SimpleNamespace() - monkeypatch.delenv("XLA_PYTHON_CLIENT_PREALLOCATE", raising=False) - monkeypatch.setitem(types.sys.modules, "jax", module) - - assert MemoryType.JAX.import_if_installed() is module - assert os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] == "false" - assert "loaded before its ArrayBridge import defaults" in caplog.text - - def test_jax_dlpack_import_repairs_misaligned_torch_views(self): - torch = pytest.importorskip("torch") - pytest.importorskip("jax") - - source = torch.arange(9, dtype=torch.int32).to(torch.uint16)[1:] - assert source.data_ptr() % 16 != 0 - - result = MemoryType.TORCH.convert_to(source, MemoryType.JAX, device_id=0) - - assert np.array_equal(np.asarray(result), np.arange(1, 9, dtype=np.uint16)) - - def test_dlpack_capability_is_owned_by_members(self): - assert {memory_type for memory_type in MemoryType if memory_type.supports_dlpack} == { - MemoryType.CUPY, - MemoryType.TORCH, - MemoryType.TENSORFLOW, - MemoryType.JAX, - } - - @pytest.mark.parametrize( - ("version", "device"), - [ - ("2.10.0", "GPU:0"), - ("2.15.0", "CPU:0"), - ("unknown", "GPU:0"), - ], - ) - def test_tensorflow_dlpack_validation_is_declaration_owned( - self, - version, - device, - ): - module = SimpleNamespace( - __version__=version, - experimental=SimpleNamespace(dlpack=object()), - ) - data = SimpleNamespace(device=device, __dlpack__=lambda: object()) - - assert not MemoryType.TENSORFLOW.supports_dlpack_data(data, module) - - def test_tensorflow_dlpack_validation_accepts_supported_gpu_tensor(self): - def exporter(data): - return object() - - module = SimpleNamespace( - __version__="2.15.0", - experimental=SimpleNamespace(dlpack=SimpleNamespace(to_dlpack=exporter)), - ) - data = SimpleNamespace(device="GPU:0") - - assert MemoryType.TENSORFLOW.supports_dlpack_data(data, module) - - def test_tensorflow_dlpack_export_uses_module_owned_leaf(self): - capsule = object() - received = [] - module = SimpleNamespace( - __version__="2.15.0", - experimental=SimpleNamespace( - dlpack=SimpleNamespace(to_dlpack=lambda data: received.append(data) or capsule) - ), - ) - data = SimpleNamespace(device="GPU:0") - - payload = MemoryType.TENSORFLOW.export_dlpack(data, module) - - assert payload is not None - assert payload.source is data - assert payload.capsule is capsule - assert received == [data] - - def test_tensorflow_dlpack_import_uses_exported_capsule(self): - capsule = object() - imported = object() - received = [] - module = SimpleNamespace( - experimental=SimpleNamespace( - dlpack=SimpleNamespace( - from_dlpack=lambda value: received.append(value) or imported, - ) - ) - ) - assert MemoryType.TENSORFLOW.from_dlpack(capsule, module) is imported - assert received == [capsule] - - def test_modern_cupy_dlpack_import_uses_protocol_source(self): - capsule = object() - source = SimpleNamespace(__dlpack__=lambda: capsule, __dlpack_device__=lambda: (2, 0)) - imported = object() - received = [] - module = SimpleNamespace(from_dlpack=lambda value: received.append(value) or imported) - - result = MemoryType.CUPY.from_dlpack(DLPackPayload(source, capsule), module) - - assert result is imported - assert received == [source] - - def test_torch_dlpack_import_uses_capsule(self): - capsule = object() - source = SimpleNamespace(__dlpack__=lambda: capsule, __dlpack_device__=lambda: (2, 0)) - imported = object() - received = [] - module = SimpleNamespace(from_dlpack=lambda value: received.append(value) or imported) - - result = MemoryType.TORCH.from_dlpack(DLPackPayload(source, capsule), module) - - assert result is imported - assert received == [capsule] - - def test_jax_dlpack_import_declines_capsule_only_source(self): - payload = DLPackPayload(source=object(), capsule=object()) - - assert MemoryType.JAX.from_dlpack(payload, SimpleNamespace()) is NotImplemented - - @pytest.mark.parametrize("exporter_name", ["to_dlpack", "toDlpack"]) - def test_protocol_dlpack_export_accepts_legacy_exporters(self, exporter_name): - capsule = object() - data = SimpleNamespace(**{exporter_name: lambda: capsule}) - - payload = MemoryType.CUPY.export_dlpack(data, object()) - - assert payload is not None - assert payload.source is data - assert payload.capsule is capsule - - def test_oom_classification_uses_loaded_module_without_import(self, monkeypatch): - from arraybridge import types - - class TorchOOMError(Exception): - pass - - module = SimpleNamespace(cuda=SimpleNamespace(OutOfMemoryError=TorchOOMError)) - monkeypatch.setitem(types.sys.modules, "torch", module) - - def fail_import(name): - raise AssertionError(f"unexpected import of {name}") - - monkeypatch.setattr(types.importlib, "import_module", fail_import) - - assert MemoryType.TORCH.is_oom_error(TorchOOMError("allocation failed")) - assert not MemoryType.TORCH.is_oom_error(RuntimeError("unrelated")) - - def test_cpu_device_scope_does_not_import_numpy(self, monkeypatch): - from arraybridge import types - - def fail_import(name): - raise AssertionError(f"unexpected import of {name}") - - monkeypatch.setattr(types.importlib, "import_module", fail_import) - with MemoryType.NUMPY.device_scope(0): - pass - - def test_jax_cpu_install_does_not_declare_a_gpu(self): - cpu_device = SimpleNamespace(platform="cpu") - module = SimpleNamespace(devices=lambda: [cpu_device]) - - assert MemoryType.JAX.available_device_ids(module) == () - with pytest.raises(ValueError, match="available device IDs are \\(\\)"): - MemoryType.JAX.device_scope(0, module) - - @pytest.mark.parametrize("device_id", [-1, 2]) - def test_device_scope_rejects_undeclared_ids(self, device_id): - devices = [SimpleNamespace(platform="gpu") for _ in range(2)] - - @contextmanager - def default_device(device): - yield device - - module = SimpleNamespace(devices=lambda: devices, default_device=default_device) - - with pytest.raises(ValueError, match=f"device {device_id} is unavailable"): - MemoryType.JAX.device_scope(device_id, module) - - def test_jax_device_identity_uses_gpu_local_index(self): - cpu = SimpleNamespace(platform="cpu") - first_gpu = SimpleNamespace(platform="gpu") - second_gpu = SimpleNamespace(platform="gpu") - module = SimpleNamespace(devices=lambda: [cpu, first_gpu, second_gpu]) - data = SimpleNamespace(device=second_gpu) - - assert MemoryType.JAX.available_device_ids(module) == (0, 1) - assert MemoryType.JAX.device_id_of(data, module) == 1 - - def test_pyclesperanto_move_restores_recognized_origin(self): - module = _FakePyclesperanto(current="gpu0") - source = object() - - result = MemoryType.PYCLESPERANTO.move_to_device(source, 1, module) - - assert result == ("copy", source) - assert module.current == "gpu0" - assert module.copies == [(source, result)] - - def test_pyclesperanto_scope_restores_origin_outside_gpu_inventory(self): - module = _FakePyclesperanto(current="cpu0") - - with MemoryType.PYCLESPERANTO.device_scope(1, module): - assert module.current == "gpu1" - - assert module.current == "cpu0" - - def test_pyclesperanto_unknown_origin_is_not_guessed_as_gpu_zero(self): - module = _FakePyclesperanto(current="cpu0") - source = object() - - result = MemoryType.PYCLESPERANTO.move_to_device(source, 0, module) - - assert result == ("copy", source) - assert module.current == "cpu0" - assert module.copies == [(source, result)] - - def test_pyclesperanto_device_identity_uses_device_name(self): - module = _FakePyclesperanto(current=SimpleNamespace(id=None, name="gpu1")) - - assert MemoryType.PYCLESPERANTO.device_id_of(object(), module) == 1 - - def test_pyclesperanto_device_identity_prefers_payload_evidence(self): - module = _FakePyclesperanto(current="gpu0") - data = SimpleNamespace(device=SimpleNamespace(id=None, name="gpu1")) - - assert MemoryType.PYCLESPERANTO.device_id_of(data, module) == 1 - - def test_pyclesperanto_non_gpu_numeric_id_is_not_accepted(self): - module = _FakePyclesperanto(current=SimpleNamespace(id=0, name="cpu0")) - - assert MemoryType.PYCLESPERANTO.device_id_of(object(), module) is None - - def test_torch_from_numpy_materializes_negative_strides(self): - received = [] - - class FakeTensor: - def to(self, device): - assert device == "cuda:0" - return self - - @contextmanager - def device_scope(device_id): - assert device_id == 0 - yield - - def from_numpy(data): - received.append(data) - assert all(stride >= 0 for stride in data.strides) - return FakeTensor() - - module = SimpleNamespace( - cuda=SimpleNamespace( - is_available=lambda: True, - device_count=lambda: 1, - device=device_scope, - ), - from_numpy=from_numpy, - ) - values = np.arange(12, dtype=np.float32).reshape(3, 4)[:, ::-1] - - MemoryType.TORCH.from_numpy(values, 0, module) - - np.testing.assert_array_equal(received[0], values) - - -class _FakePyclesperanto: - def __init__(self, current: str) -> None: - self.current = current - self.gpu_devices = ("gpu0", "gpu1") - self.copies = [] - - def list_available_devices(self, device_type=None): - assert device_type in (None, "gpu") - return self.gpu_devices - - def get_device(self): - return self.current - - def select_device(self, selector, device_type=None): - if isinstance(selector, int): - assert device_type in (None, "gpu") - self.current = self.gpu_devices[selector] - else: - self.current = selector - - def create_like(self, data): - return ("copy", data) - - def copy(self, source, target): - self.copies.append((source, target)) diff --git a/tests/test_utils.py b/tests/test_utils.py index e27adfb..b205ce2 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -107,99 +107,6 @@ def test_supports_dlpack_object_without_it(self): assert not _supports_dlpack(obj) -class TestDeviceOperations: - """Tests for device-related utility functions.""" - - def test_get_device_id_numpy(self): - """Test getting device ID for NumPy arrays.""" - import numpy as np - - from arraybridge.utils import _get_device_id - - arr = np.array([1, 2, 3]) - device_id = _get_device_id(arr, "numpy") - assert device_id is None # NumPy is CPU-only - - def test_set_device_numpy(self): - """Test setting device for NumPy (should be no-op).""" - from arraybridge.utils import _set_device - - # Should not raise - _set_device("numpy", 0) - - def test_move_to_device_numpy(self): - """Test moving NumPy array to device (should return same array).""" - import numpy as np - - from arraybridge.utils import _move_to_device - - arr = np.array([1, 2, 3]) - result = _move_to_device(arr, "numpy", 0) - assert result is arr # Should return same object - - @pytest.mark.parametrize("device_id", [0, 1, 2]) - def test_set_device_torch_mock(self, device_id, monkeypatch): - """Test setting device for torch with mock.""" - import types - - selected = [] - mock_torch = types.SimpleNamespace( - cuda=types.SimpleNamespace( - is_available=lambda: True, - device_count=lambda: 3, - set_device=selected.append, - ) - ) - monkeypatch.setitem(sys.modules, "torch", mock_torch) - - from arraybridge.utils import _set_device - - _set_device("torch", device_id) - assert selected == [device_id] - - def test_get_device_id_torch_mock(self, monkeypatch): - """Test getting device ID for torch tensor with mock.""" - import types - - mock_device = types.SimpleNamespace(index=1) - mock_tensor = types.SimpleNamespace(is_cuda=True, device=mock_device) - mock_torch = types.SimpleNamespace(cuda=types.SimpleNamespace(current_device=lambda: 1)) - monkeypatch.setitem(sys.modules, "torch", mock_torch) - - from arraybridge.utils import _get_device_id - - device_id = _get_device_id(mock_tensor, "torch") - assert device_id == 1 - - def test_move_to_device_torch_mock(self, monkeypatch): - """Test moving torch tensor to device with mock.""" - import types - - mock_tensor = types.SimpleNamespace( - is_cuda=True, device=types.SimpleNamespace(index=0), to=lambda device: "moved_tensor" - ) - - class DeviceScope: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, traceback): - return False - - mock_torch = types.SimpleNamespace( - cuda=types.SimpleNamespace( - is_available=lambda: True, - device_count=lambda: 2, - device=lambda device_id: DeviceScope(), - ) - ) - monkeypatch.setitem(sys.modules, "torch", mock_torch) - - from arraybridge.utils import _move_to_device - - assert _move_to_device(mock_tensor, "torch", 1) == "moved_tensor" - - class TestEnsureModuleTensorFlowVersion: def test_ordinary_tensorflow_import_respects_supported_baseline(self, monkeypatch): import types @@ -210,39 +117,3 @@ def test_ordinary_tensorflow_import_respects_supported_baseline(self, monkeypatc from arraybridge.utils import _ensure_module assert _ensure_module("tensorflow") is mock_tf - - -class TestGetDeviceIdCallableHandler: - """Tests for _get_device_id with callable handlers.""" - - def test_get_device_id_with_callable_handler(self, monkeypatch): - """Test _get_device_id with a callable handler (pyclesperanto).""" - import types - - from arraybridge.utils import _get_device_id - - mock_cle = types.SimpleNamespace( - get_device=lambda: "gpu1", - list_available_devices=lambda device_type=None: ["gpu0", "gpu1"], - ) - monkeypatch.setitem(sys.modules, "pyclesperanto", mock_cle) - - # Create mock data - mock_data = types.SimpleNamespace() - - assert _get_device_id(mock_data, "pyclesperanto") == 1 - - def test_get_device_id_fallback_on_error(self, monkeypatch): - """Invalid GPU objects fail rather than inventing a fallback device.""" - import types - - from arraybridge.exceptions import MemoryConversionError - from arraybridge.utils import _get_device_id - - # Create a mock torch tensor that will fail device ID extraction - mock_tensor = types.SimpleNamespace() # Missing device attribute - mock_torch = types.SimpleNamespace() - monkeypatch.setitem(sys.modules, "torch", mock_torch) - - with pytest.raises(MemoryConversionError, match="device_identification"): - _get_device_id(mock_tensor, "torch")