diff --git a/.gitignore b/.gitignore index 43e869c..9793443 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ /cuvis/cuvis_il.py /cuvis/_cuvis_pyil.pyd /venv +/__pycache__ /cuvis/__pycache__ /cuvis/git-hash.txt diff --git a/CHANGELOG.md b/CHANGELOG.md index 414bfb0..3f70df0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,13 @@ Pre-releases (`b*`, `rc*`) are not listed. - `CI` - `scripts/check_changelog.py` validates this file's structure (header format, allowed section names, descending versions) and the tag/version/changelog agreement at release time. - `CONTRIBUTING.md` - documents the branch model, the version scheme, the changelog conventions and the release checklist. - `cuvis.BindingInfo` - new frozen dataclass with the fields `built_against: str`, `library_version: str`, `library_path: str` and `missing_symbols: Tuple[str, ...]`, the read-only property `is_complete: bool`, and a `__str__` rendering a report fit for a bug report. +- `cuvis.CudaImageData` - new class, device-resident image data backed by a shareable CUDA buffer, exposed zero-copy through DLPack or `__cuda_array_interface__` with no host copy. +- `cuvis.CudaImageData.export_payload` - new method, returns `bytes`. +- `cuvis.CudaImageData.make_ipc` - new method, returns `bytes`. +- `cuvis.CudaImageData.to_torch` - new method, returns `torch.Tensor`. +- `cuvis.Measurement.get_cube` - new method, returns `Union[ImageData, CudaImageData]` depending on whether `cuvis.cuda.enable` was called. +- `cuvis.Measurement.get_cube_cuda` - new method, returns `CudaImageData`. +- `cuvis.Measurement.get_cube_cuda_ipc` - new method, returns `CudaImageData`. - `cuvis.SdkSettings` - new class, a `MutableMapping` of setting id to value that writes the SDK's `cuvis.settings` file, so the SDK configuration can be built in Python instead of maintained by hand. Values are stored as strings: `bool` becomes `true`/`false`, an `Enum` becomes its value, anything else goes through `str()`, and `None` drops the entry. - `cuvis.SdkSettings.__enter__`, `cuvis.SdkSettings.__exit__` - new methods; entering the context serializes the settings into a temporary directory and returns its path as `str`, leaving the context removes the directory. @@ -32,6 +39,25 @@ Pre-releases (`b*`, `rc*`) are not listed. - `cuvis.binding.info` - new function, returns `BindingInfo`. - `cuvis.binding.missing_symbols` - new function, returns `FrozenSet[str]`. - `cuvis.binding.require` - new function, raises `UnavailableSDKFunction` naming whichever of the given functions the installed cuvis library does not provide. +- `cuvis.binding.unavailable` - new function, returns `Tuple[str, ...]`. + A function the binding never exposed is unusable as well, so availability cannot be answered from the reported-missing list alone. +- `cuvis.cuda` - new module gating the optional CUDA feature surface; CUDA stays off until `cuvis.cuda.enable` is called. +- `cuvis.cuda.BACKEND_NONE`, `cuvis.cuda.BACKEND_POOL`, `cuvis.cuda.BACKEND_LEGACY`, `cuvis.cuda.BACKEND_VMM` - new constants, the IPC backend codes from `cuvis.h`. +- `cuvis.cuda.CudaCapabilities` - new `NamedTuple` with the `bool` fields `same_process`, `ipc_pool`, `ipc_legacy`, `ipc_vmm`, `torch` and `cuda_python`, and the read-only property `any_ipc: bool`. +- `cuvis.cuda.capabilities` - new function, returns `CudaCapabilities`. +- `cuvis.cuda.disable` - new function. +- `cuvis.cuda.enable` - new function. +- `cuvis.cuda.is_enabled` - new function, returns `bool`. +- `cuvis.cuda.require_device` - new function, raises `UnavailableSDKFunction` unless the installed library provides the same-process device path. +- `cuvis.cuda.require_ipc` - new function, raises `UnavailableSDKFunction` unless the installed library provides the cross-process export path. +- `cuvis_ipc` - new top-level module, the consumer side of cross-process CUDA IPC. + It sits outside the `cuvis` package because importing `cuvis` requires the SDK and a consumer process does not have one; its only import is `struct`. +- `cuvis_ipc.ImportedCube` - new class, a mapped IPC buffer in the consumer process, usable as a context manager. +- `cuvis_ipc.open` - new function, returns `ImportedCube`. +- `cuvis_ipc.open_descriptor` - new function, returns `ImportedCube`. +- `cuvis_ipc.pack_payload` - new function, returns `bytes`. +- `pyproject.toml` - `py-modules` declaring the top-level `cuvis_ipc`. +- `tests/` - `test_binding.py`, `test_cuda.py` and `test_cuvis_ipc.py`. ### Changed diff --git a/cuvis/Measurement.py b/cuvis/Measurement.py index a7ff1b3..11d305c 100644 --- a/cuvis/Measurement.py +++ b/cuvis/Measurement.py @@ -14,7 +14,8 @@ _utc_from_epoch_ms, ) from .cuvis_types import DataFormat, ProcessingMode, ReferenceType -from .cube_utils import ImageData +from .cube_utils import ImageData, CudaImageData +from . import cuda import cuvis.cuvis_types as internal @@ -37,6 +38,13 @@ class Measurement(object): session_info: SessionData # read-only frame_id: int # read-only + # When False, refresh() skips fetching image data to the host via + # cuvis_measurement_get_data_image. That host fetch moves a GPU-processed cube + # into host memory and frees the device copy, which makes the same-process CUDA + # path (get_cube_cuda) unavailable. Set False before processing when you intend + # to read the cube as CUDA device memory. Default True preserves normal behaviour. + _refresh_images = True + def __init__(self, base: Union[int, str, Path]): self._handle = None self._session = None @@ -107,6 +115,10 @@ def refresh(self) -> None: ) cdtype = cuvis_il.p_cuvis_data_type_t_value(pType) if cdtype == cuvis_il.data_type_image: + if not Measurement._refresh_images: + # Skip the host fetch so a GPU-processed cube stays in device + # memory and remains reachable via get_cube_cuda. + continue data = cuvis_il.cuvis_imbuffer_t() cuvis_il.cuvis_measurement_get_data_image(self._handle, key, data) # t0 = datetime.datetime.now() @@ -275,6 +287,93 @@ def cube(self) -> ImageData: "This Measurement does not have a cube saved. Consider reprocessing with a Processing Context." ) + def get_cube_cuda(self, key: str = "cube") -> CudaImageData: + """Image data as a device-resident CUDA buffer for same-process, zero-copy use. + + .. code-block:: python3 + + from cuvis import cuda + + if cuda.capabilities().same_process: + cuda.enable() # BEFORE loading or processing + tensor = mesu.get_cube_cuda().to_torch() + + Returns a :class:`cuvis.CudaImageData` wrapping a CUVIS_CUDA_MEM handle; read it + with ``.to_torch()`` (DLPack, which ties the buffer lifetime to the tensor) or + through ``__cuda_array_interface__`` (which does not, so keep the CudaImageData + alive). No host copy is made. + + The cube must still be on the device, which it is only when `cuda.enable` was + called before it was processed; the host fetch in `refresh` otherwise moves it to + host memory and frees the device copy. + + :param key: which image entry to read, `"cube"` unless the measurement carries + several. + :raises cuvis.UnavailableSDKFunction: the installed cuvis library provides no + CUDA support; call `cuda.capabilities` first to avoid this. + :raises cuvis.cuvis_aux.SDKException: the image data is not device-backed. + """ + cuda.require_device() + buf = cuvis_il.cuvis_cuda_imbuffer_t() + if cuvis_il.status_ok != cuvis_il.cuvis_measurement_get_data_image_cuda( + self._handle, key, buf + ): + raise SDKException() + return CudaImageData(buf) + + def get_cube_cuda_ipc(self, key: str = "cube", backend: int = 0) -> CudaImageData: + """Image data as a shareable CUDA buffer for cross-process use. + + Producer side; the consumer opens the payload with :mod:`cuvis_ipc`, which needs + no SDK of its own. + + .. code-block:: python3 + + cimg = mesu.get_cube_cuda_ipc() # keep alive until the consumer is done + send(cimg.export_payload()) # descriptor + geometry, one blob + + # ... in the consumer process, no cuvis installed ... + import cuvis_ipc + with cuvis_ipc.open(payload) as cube: + tensor = cube.to_torch() + + Fetches the device buffer with `get_cube_cuda` and creates an IPC export on it, + filling `.descriptor` with the transportable bytes. The returned object is the + in-process pin, since legacy IPC carries no cross-process refcount: drop it and + the consumer is reading freed device memory. + + :param key: which image entry to read. + :param backend: which mechanism to export with, one of `cuvis.cuda.BACKEND_NONE` + (auto), `BACKEND_POOL`, `BACKEND_LEGACY` or `BACKEND_VMM`. `cuda.capabilities` + reports which of them this device supports. + :raises cuvis.UnavailableSDKFunction: the installed cuvis library provides no + CUDA IPC support. + :raises cuvis.cuvis_aux.SDKException: the requested backend is unavailable on + this device. + """ + cimg = self.get_cube_cuda(key) + cimg.make_ipc(backend) + return cimg + + def get_cube(self, key: str = "cube") -> Union[ImageData, CudaImageData]: + """Cube through whichever mode is active, so one call site serves both. + + .. code-block:: python3 + + cube = mesu.get_cube() # CudaImageData after cuda.enable(), else ImageData + + With CUDA mode on (`cuvis.cuda.enable`) this is `get_cube_cuda` and raises when + the device path is unavailable. There is deliberately no silent fallback to the + host: a zero-copy path that quietly degrades to two copies is worse than an error, + because the cost is invisible. + + :param key: which image entry to read. + :return: `CudaImageData` in CUDA mode, otherwise the host `ImageData`. + """ + if cuda.is_enabled(): + return self.get_cube_cuda(key) + return self.cube + @property def thumbnail(self): thumb = [val for key, val in self.data.items() if "view" in key] diff --git a/cuvis/__init__.py b/cuvis/__init__.py index 065de6c..23388e6 100644 --- a/cuvis/__init__.py +++ b/cuvis/__init__.py @@ -41,7 +41,8 @@ from .Export import CubeExporter, EnviExporter, TiffExporter, ViewExporter from .Calibration import Calibration from .AcquisitionContext import AcquisitionContext -from .cube_utils import ImageData +from .cube_utils import ImageData, CudaImageData +from . import cuda import os import platform import sys diff --git a/cuvis/_dlpack.py b/cuvis/_dlpack.py new file mode 100644 index 0000000..1084e7b --- /dev/null +++ b/cuvis/_dlpack.py @@ -0,0 +1,104 @@ +"""Minimal DLPack producer over a raw CUDA device pointer. + +Builds a DLManagedTensor PyCapsule via ctypes so torch.from_dlpack can consume a +foreign device buffer zero-copy. The capsule deleter calls a supplied on_delete +callback when torch releases the tensor, which is how buffer lifetime is tied to +the tensor (the callback drops the SDK reference that pins the buffer). + +Ported from utils_data_cuda/examples/torch_local.py; the only change is that the +deleter calls an injected callback instead of a hard-coded ctypes SDK. +""" + +import ctypes + +_kDLCUDA = 2 +_kDLInt = 0 +_kDLUInt = 1 +_kDLFloat = 2 + + +class _DLDevice(ctypes.Structure): + _fields_ = [("device_type", ctypes.c_int), ("device_id", ctypes.c_int)] + + +class _DLDataType(ctypes.Structure): + _fields_ = [ + ("code", ctypes.c_uint8), + ("bits", ctypes.c_uint8), + ("lanes", ctypes.c_uint16), + ] + + +class _DLTensor(ctypes.Structure): + _fields_ = [ + ("data", ctypes.c_void_p), + ("device", _DLDevice), + ("ndim", ctypes.c_int), + ("dtype", _DLDataType), + ("shape", ctypes.POINTER(ctypes.c_int64)), + ("strides", ctypes.POINTER(ctypes.c_int64)), + ("byte_offset", ctypes.c_uint64), + ] + + +class _DLManagedTensor(ctypes.Structure): + pass + + +_DELETER = ctypes.CFUNCTYPE(None, ctypes.POINTER(_DLManagedTensor)) +_DLManagedTensor._fields_ = [ + ("dl_tensor", _DLTensor), + ("manager_ctx", ctypes.c_void_p), + ("deleter", _DELETER), +] + +_pycapsule_new = ctypes.pythonapi.PyCapsule_New +_pycapsule_new.restype = ctypes.py_object +_pycapsule_new.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + +# Keep ctypes structs / deleters / callbacks alive until torch invokes the deleter. +_LIVE = {} + + +class CudaDlpack: + """DLPack producer over (ptr, nbytes) on a CUDA device. on_delete() runs when + torch releases the tensor. Exposes a flat uint8 buffer; reshape/retype in torch.""" + + def __init__(self, ptr, nbytes, device, on_delete): + self._ptr = int(ptr) + self._n = int(nbytes) + self._dev = int(device) + self._on_delete = on_delete + + def __dlpack_device__(self): + return (_kDLCUDA, self._dev) + + def __dlpack__(self, stream=None, max_version=None, dl_device=None, copy=None): + shape = (ctypes.c_int64 * 1)(self._n) + mt = _DLManagedTensor() + mt.dl_tensor.data = ctypes.c_void_p(self._ptr) + mt.dl_tensor.device = _DLDevice(_kDLCUDA, self._dev) + mt.dl_tensor.ndim = 1 + mt.dl_tensor.dtype = _DLDataType(_kDLUInt, 8, 1) + mt.dl_tensor.shape = shape + mt.dl_tensor.strides = None + mt.dl_tensor.byte_offset = 0 + + on_delete = self._on_delete + key = id(mt) + + def _del(_p): + try: + on_delete() + finally: + _LIVE.pop(key, None) + + deleter = _DELETER(_del) + mt.deleter = deleter + _LIVE[key] = (mt, shape, deleter) + return _pycapsule_new(ctypes.byref(mt), b"dltensor", None) + + +def make_cuda_dlpack(ptr, nbytes, device, on_delete): + """Return an object that torch.from_dlpack consumes into a zero-copy CUDA tensor.""" + return CudaDlpack(ptr, nbytes, device, on_delete) diff --git a/cuvis/binding.py b/cuvis/binding.py index 8a89a1c..e550029 100644 --- a/cuvis/binding.py +++ b/cuvis/binding.py @@ -29,9 +29,11 @@ derives from :class:`RuntimeError`, so a single ``except RuntimeError`` covers both, while ``except SDKException`` still catches it as an ordinary cuvis error. -Against a binding too old to report any of this (an older ``cuvis_il`` wheel), every -query answers empty: :func:`missing_symbols` is empty, :func:`available` is ``True`` -and :func:`require` never raises. Absence of evidence, not evidence of absence. +Against a binding too old to report any of this (an older ``cuvis_il`` wheel), +:func:`missing_symbols` is empty and :func:`info` reports unknown throughout: absence of +evidence, not evidence of absence. :func:`available` and :func:`require` stay meaningful +there, because a function such an old binding never exposed is unusable whether or not +anything reports it missing. """ from dataclasses import dataclass, field @@ -161,8 +163,25 @@ def missing_symbols() -> FrozenSet[str]: return frozenset(getattr(cuvis_il, "missing_symbols", ())) +def unavailable(*names: str) -> Tuple[str, ...]: + """Which of the named functions cannot be called, in the order given. + + A function is unusable for either of two reasons, and callers care about neither: + the binding never exposed it, which is what an older ``cuvis_il`` wheel looks like, + or the binding exposes it but the loaded library does not export it. Checking only + the second would report a function the binding does not even have as available. + + :param names: C function names as they appear in ``cuvis.h``. + :return: the subset that is unusable, empty when all of them can be called. + """ + absent = missing_symbols() + return tuple( + name for name in names if name in absent or not hasattr(cuvis_il, name) + ) + + def available(*names: str) -> bool: - """Whether every named function is provided by the installed cuvis library. + """Whether every named function can actually be called. .. code-block:: python3 @@ -170,16 +189,14 @@ def available(*names: str) -> bool: cube = mesu.get_cube_cuda() :param names: C function names as they appear in ``cuvis.h``. - :return: ``True`` if none of them is reported missing. With a binding too old to - report anything this is always ``True``, so treat it as "nothing known to be - missing" rather than a guarantee. + :return: ``True`` if the binding exposes every one of them and none is reported + missing from the loaded library. """ - absent = missing_symbols() - return not any(name in absent for name in names) + return not unavailable(*names) def require(*names: str) -> None: - """Raise unless every named function is provided by the installed cuvis library. + """Raise unless every named function can actually be called. Use it at the start of an operation to fail with a clear explanation, instead of letting a call fail deeper in with less context. @@ -189,13 +206,12 @@ def require(*names: str) -> None: binding.require("cuvis_cuda_mem_get_view", "cuvis_cuda_mem_free") :param names: C function names as they appear in ``cuvis.h``. - :raises UnavailableSDKFunction: naming whichever of them are missing; the message + :raises UnavailableSDKFunction: naming whichever of them are unusable; the message also states the loaded SDK version and the one the binding expects. """ - absent = missing_symbols() - unavailable = tuple(name for name in names if name in absent) - if unavailable: - raise UnavailableSDKFunction(*unavailable) + missing = unavailable(*names) + if missing: + raise UnavailableSDKFunction(*missing) __all__ = [ @@ -203,6 +219,7 @@ def require(*names: str) -> None: "UnavailableSDKFunction", "info", "missing_symbols", + "unavailable", "available", "require", ] diff --git a/cuvis/cube_utils.py b/cuvis/cube_utils.py index 9ecf9f5..2fb71db 100644 --- a/cuvis/cube_utils.py +++ b/cuvis/cube_utils.py @@ -3,6 +3,8 @@ import numpy as np import operator from .cuvis_aux import SDKException +from .cuvis_types import DataFormat +from . import cuda _IMBUF_READERS = { 1: cuvis_il.cuvis_read_imbuf_uint8, @@ -417,3 +419,160 @@ def apply(self, other): ImageData.__abs__ = lambda self: self._wrap(abs(self.array)) del _op, _reflected, _method + + +class CudaImageData(object): + """Device-resident image data backed by a shareable CUDA buffer. + + Wraps a CUVIS_CUDA_MEM handle plus geometry, exposing the device memory as a + zero-copy CUDA tensor (torch via DLPack, or any consumer via + __cuda_array_interface__). Unlike ImageData, no host copy is made. + + Lifetime: this object owns the CUVIS_CUDA_MEM handle and frees it in __del__. + to_torch() takes its own SDK reference so the returned tensor can outlive this + object; the __cuda_array_interface__ path does not, so keep this object alive + until such a consumer is done with it. + + For cross-process sharing, make_ipc() (called by Measurement.get_cube_cuda_ipc) + creates a CUVIS_CUDA_IPC handle and fills .descriptor. That IPC handle is an + independent reference to the same buffer, freed in __del__; freeing the mem handle + while it is open would not release the memory. + """ + + _TYPESTR = {1: "|u1", 2: " bytes: + """A single transportable blob (IPC descriptor + geometry) for a consumer process. + + Send these bytes out-of-band; the consumer opens them with cuvis_ipc.open(payload) + and gets a correctly shaped/typed tensor. Calls make_ipc(backend) if not already done + (backend: 0=auto, 1=pool, 2=legacy, 3=VMM). Keep this CudaImageData alive until the + consumer is finished (legacy IPC has no cross-process refcount). + """ + if self.descriptor is None: + self.make_ipc(backend) + # Imported here, not at module scope: cuvis_ipc is the consumer half and lives + # outside the package, so a producer that never exports must not fail to import + # cuvis because it is absent. + import cuvis_ipc + + return cuvis_ipc.pack_payload( + self.descriptor, self.width, self.height, self.channels, self._format + ) + + def __del__(self): + try: + if self._ipc_handle is not None: + cuvis_il.cuvis_cuda_ipc_handle_free(self._ipc_handle) + except Exception: + pass + try: + cuvis_il.cuvis_cuda_mem_free(self._handle) + except Exception: + pass diff --git a/cuvis/cuda.py b/cuvis/cuda.py new file mode 100644 index 0000000..60c8ce3 --- /dev/null +++ b/cuvis/cuda.py @@ -0,0 +1,186 @@ +"""Opt-in CUDA support for cuvis. + +CUDA is off by default. Reach it explicitly: + + from cuvis import cuda + + caps = cuda.capabilities() # what this SDK, this device and this env support + if caps.same_process: + cuda.enable() # BEFORE loading/processing measurements + ... + t = mesu.get_cube().to_torch() # zero-copy device tensor (DLPack, lifecycle-managed) + +`enable()` also disables the host auto-refresh (`Measurement._refresh_images = False`) so a +GPU-processed cube stays on the device instead of being copied to host and freed. + +Three unrelated things can each deny CUDA, and they are answered separately rather than +collapsed into one boolean: the installed cuvis library may not provide the functions, which +`cuvis.binding` reports without calling anything; the device or driver may not support a +backend, which only the SDK can answer; and the optional consumer packages may be absent. +Keeping them apart is what lets `enable()` say which of the three went wrong. +""" + +import importlib.util +from typing import NamedTuple + +from . import binding +from ._cuvis_il import cuvis_il + +# Backend codes, matching CUVIS_CUDA_IPC_BACKEND_* in cuvis.h. BACKEND_NONE doubles as the +# same-process probe and as "auto" where an export picks a backend itself. +BACKEND_NONE = 0 +BACKEND_POOL = 1 +BACKEND_LEGACY = 2 +BACKEND_VMM = 3 + +# Functions the same-process device path calls, as named in cuvis.h. BACKEND_PROBE answers +# for a backend and so gates every capability query, including same_process. +BACKEND_PROBE = "cuvis_cuda_ipc_backend_available" +DEVICE_FUNCTIONS = ( + "cuvis_measurement_get_data_image_cuda", + "cuvis_cuda_mem_get_view", + "cuvis_cuda_mem_copy_handle", + "cuvis_cuda_mem_free", +) + +# Needed on top of those to export a device buffer to another process. +IPC_FUNCTIONS = ( + "cuvis_cuda_ipc_handle_create", + "cuvis_cuda_ipc_get_descriptor", + "cuvis_cuda_ipc_handle_free", +) + +_enabled = False + + +class CudaCapabilities(NamedTuple): + """What the installed cuvis library, the current device and this environment support. + + same_process: the CUDA boundary responds, so same-process device sharing works. + ipc_pool: exportable memory pool backend (zero-copy cross-process, needs an exportable pool). + ipc_legacy: legacy cudaIpc backend (copy-on-export; available on most WDDM/Linux GPUs). + ipc_vmm: driver-API VMM backend (copy-on-export, with an exportable handle type). + torch / cuda_python: optional consumer packages importable (checked, not imported). + """ + + same_process: bool + ipc_pool: bool + ipc_legacy: bool + ipc_vmm: bool + torch: bool + cuda_python: bool + + @property + def any_ipc(self) -> bool: + return self.ipc_pool or self.ipc_legacy or self.ipc_vmm + + +def _installed(package: str) -> bool: + """Whether an optional consumer package is installed, without importing it. + + find_spec imports parent packages to reach a submodule, and cuvis puts its own + directory on sys.path, so probing `cuda.bindings` can resolve `cuda` to this very + module whenever cuda-python is absent. A probe that cannot resolve is an answer: + the package is not installed. + """ + try: + return importlib.util.find_spec(package) is not None + except (ImportError, ValueError): + return False + + +def _backend_supported(code: int) -> bool: + """Ask the SDK whether this device and driver support one backend. + + Reached only once the function is known to exist, so a false answer here is the SDK's + verdict on the hardware rather than a missing symbol wearing the same disguise. + """ + out = cuvis_il.new_p_int() + if cuvis_il.status_ok != cuvis_il.cuvis_cuda_ipc_backend_available(code, out): + return False + return bool(cuvis_il.p_int_value(out)) + + +def capabilities() -> CudaCapabilities: + """Probe what is supported, here and now. Safe to call before init() or enable().""" + device = binding.available(BACKEND_PROBE, *DEVICE_FUNCTIONS) + ipc = device and binding.available(*IPC_FUNCTIONS) + return CudaCapabilities( + same_process=device and _backend_supported(BACKEND_NONE), + ipc_pool=ipc and _backend_supported(BACKEND_POOL), + ipc_legacy=ipc and _backend_supported(BACKEND_LEGACY), + ipc_vmm=ipc and _backend_supported(BACKEND_VMM), + torch=_installed("torch"), + cuda_python=_installed("cuda.bindings"), + ) + + +def require_device() -> None: + """Raise unless the installed cuvis library provides the same-process device path. + + Guards the entry points so a library without CUDA fails by naming the functions it + lacks, rather than as an AttributeError from deep inside the binding. + + :raises cuvis.UnavailableSDKFunction: naming the functions that are unavailable. + """ + binding.require(*DEVICE_FUNCTIONS) + + +def require_ipc() -> None: + """Raise unless the installed cuvis library provides the cross-process export path. + + :raises cuvis.UnavailableSDKFunction: naming the functions that are unavailable. + """ + binding.require(*DEVICE_FUNCTIONS, *IPC_FUNCTIONS) + + +def enable() -> None: + """Turn on CUDA mode. Call this BEFORE loading or processing measurements. + + Routes `Measurement.get_cube()` through the device path and disables the host + auto-refresh so the GPU cube is kept on the device. + + :raises cuvis.UnavailableSDKFunction: the installed cuvis library does not provide the + CUDA functions; the message names them and both library versions. + :raises RuntimeError: the library provides them, but this device or driver reports no + CUDA support. + """ + global _enabled + binding.require(BACKEND_PROBE, *DEVICE_FUNCTIONS) + if not _backend_supported(BACKEND_NONE): + raise RuntimeError( + "the installed CUVIS SDK provides the CUDA functions, but this device " + "reports no CUDA support\n{}".format(binding.info()) + ) + from .Measurement import Measurement + + Measurement._refresh_images = False + _enabled = True + + +def disable() -> None: + """Turn off CUDA mode and restore the host auto-refresh.""" + global _enabled + from .Measurement import Measurement + + Measurement._refresh_images = True + _enabled = False + + +def is_enabled() -> bool: + return _enabled + + +__all__ = [ + "CudaCapabilities", + "capabilities", + "require_device", + "require_ipc", + "enable", + "disable", + "is_enabled", + "BACKEND_NONE", + "BACKEND_POOL", + "BACKEND_LEGACY", + "BACKEND_VMM", +] diff --git a/cuvis_ipc.py b/cuvis_ipc.py new file mode 100644 index 0000000..0bb7ec0 --- /dev/null +++ b/cuvis_ipc.py @@ -0,0 +1,317 @@ +"""Import-safe cross-process CUDA IPC consumer utilities. + +Use this in a SEPARATE process that receives a cuvis IPC payload - it does NOT initialize +or link the cuvis SDK. It needs only cuda-python (`cuda.bindings`) and torch, imported lazily. + +It sits outside the `cuvis` package on purpose. Importing `cuvis` requires the SDK to be +installed and the CUVIS environment variable to be set, which a consumer process by +definition does not have; `import cuvis_ipc` needs neither, and the only import here is +`struct`. + +Producer (in the cuvis process): + cimg = mesu.get_cube_cuda_ipc() # keep this alive until the consumer is done + payload = cimg.export_payload() # a single transportable bytes blob (descriptor + geometry) + +Consumer (this module, any process): + import cuvis_ipc + with cuvis_ipc.open(payload) as cube: + t = cube.to_torch() # correctly shaped/typed zero-copy CUDA tensor + ... # use t inside the block + # leaving the block releases this process's mapping (does not free the exporter's memory) + +The exporting process must outlive this importer: legacy IPC has no cross-process refcount. +""" + +import struct + +# --- IPC descriptor wire format (locked; mirrors cuvis_cuda_ipc_descriptor_t) --- +# 48-byte header + 64-byte blob (pool OS handle) at offset 48, then ptr_blob_len(+pad) and a +# 64-byte ptr_blob (cudaMemPoolPtrExportData) at offset 120. Total 184. +_HEAD = struct.Struct( + " torch dtype name / __cuda_array_interface__ typestr (1/2/3/4 = u8/u16/u32/f32) +_TORCH_DTYPE = {1: "uint8", 2: "uint16", 3: "uint32", 4: "float32"} +_TYPESTR = {1: "|u1", 2: " bytes: + """Bundle an IPC descriptor and cube geometry into one transportable blob.""" + if len(descriptor) != _DESC_LEN: + raise ValueError(f"descriptor must be {_DESC_LEN} bytes, got {len(descriptor)}") + return _PAYLOAD_HDR.pack( + _MAGIC, _VERSION, int(width), int(height), int(channels), int(format_code) + ) + bytes(descriptor) + + +def _unpack_payload(payload: bytes): + magic, version, width, height, channels, fmt = _PAYLOAD_HDR.unpack_from(payload, 0) + if magic != _MAGIC: + raise ValueError("not a cuvis IPC payload (bad magic)") + if version != _VERSION: + raise ValueError(f"unsupported cuvis IPC payload version {version}") + descriptor = bytes(payload[_PAYLOAD_HDR.size :]) + return (width, height, channels, fmt), descriptor + + +def _ck(ret, what): + err = ret[0] + if int(err) != 0: + raise RuntimeError(f"{what} failed: {err}") + return ret[1:] if len(ret) > 1 else None + + +class _CudaArray: + def __init__(self, ptr, nbytes): + self.__cuda_array_interface__ = { + "shape": (nbytes,), + "typestr": "|u1", + "data": (int(ptr), False), + "version": 3, + } + + +class ImportedCube: + """A mapped IPC buffer in the consumer process. Use as a context manager. + + open()/open_descriptor() return this; leaving the `with` block releases the mapping. + to_torch() (preferred) and __cuda_array_interface__ produce zero-copy views that are + valid only while the block is open. + """ + + def __init__(self, descriptor_bytes: bytes, shape=None, format_code=None): + ( + self.backend, + self.device, + self.htype, + blob_len, + self.size, + self.alloc, + self.offset, + self.pid, + ) = _HEAD.unpack_from(descriptor_bytes, 0) + if blob_len > _BLOB_MAX: + raise ValueError(f"blob_len {blob_len} exceeds {_BLOB_MAX}") + self._blob = bytes(descriptor_bytes[_BLOB_OFF : _BLOB_OFF + blob_len]) + (ptr_blob_len,) = struct.unpack_from(" _PTR_BLOB_MAX: + raise ValueError(f"ptr_blob_len {ptr_blob_len} exceeds {_PTR_BLOB_MAX}") + self._ptr_blob = bytes( + descriptor_bytes[_PTR_BLOB_OFF : _PTR_BLOB_OFF + ptr_blob_len] + ) + self._shape = tuple(shape) if shape is not None else None + self._format = format_code + self._close = None + + from cuda.bindings import runtime + + runtime.cudaSetDevice(self.device) + + if self.backend == BACKEND_POOL: + self._ptr, self._close = self._open_pool() + elif self.backend == BACKEND_LEGACY: + self._ptr, self._close = self._open_legacy() + elif self.backend == BACKEND_VMM: + self._ptr, self._close = self._open_vmm() + else: + raise NotImplementedError( + f"backend {self.backend} is not importable cross-process" + ) + self._ptr += ( + self.offset + ) # 0 for pool/legacy/vmm (import returns the exact base pointer) + + @property + def shape(self): + return self._shape + + @property + def device_ptr(self): + return self._ptr + + def _open_pool(self): + # IPC-capable memory pool: import the pool from its OS shareable handle, grant this + # device access, then import the exact pointer from the per-allocation export data. + from cuda.bindings import runtime + + if self.htype == H_WIN32_KMT: + htype = runtime.cudaMemAllocationHandleType.cudaMemHandleTypeWin32Kmt + elif self.htype == H_POSIX_FD: + htype = ( + runtime.cudaMemAllocationHandleType.cudaMemHandleTypePosixFileDescriptor + ) + else: + raise NotImplementedError( + f"pool handle_type {self.htype} needs out-of-band duplication (DuplicateHandle / SCM_RIGHTS)" + ) + + handle_val = int.from_bytes(self._blob, "little") + (pool,) = _ck( + runtime.cudaMemPoolImportFromShareableHandle(handle_val, htype, 0), + "cudaMemPoolImportFromShareableHandle", + ) + + acc = runtime.cudaMemAccessDesc() + acc.location.type = runtime.cudaMemLocationType.cudaMemLocationTypeDevice + acc.location.id = self.device + acc.flags = runtime.cudaMemAccessFlags.cudaMemAccessFlagsProtReadWrite + _ck(runtime.cudaMemPoolSetAccess(pool, [acc], 1), "cudaMemPoolSetAccess") + + export_data = runtime.cudaMemPoolPtrExportData() + export_data.reserved = self._ptr_blob.ljust(_PTR_BLOB_MAX, b"\x00") + (ptr,) = _ck( + runtime.cudaMemPoolImportPointer(pool, export_data), + "cudaMemPoolImportPointer", + ) + + def close(): + runtime.cudaFree(ptr) # release this process's imported pointer + runtime.cudaMemPoolDestroy(pool) # release the imported pool handle + + return int(ptr), close + + def _open_legacy(self): + # Legacy cudaIpc: the descriptor blob is a self-contained 64-byte cudaIpcMemHandle_t. + from cuda.bindings import runtime + + h = runtime.cudaIpcMemHandle_t() + h.reserved = self._blob.ljust(_BLOB_MAX, b"\x00") + (ptr,) = _ck( + runtime.cudaIpcOpenMemHandle(h, runtime.cudaIpcMemLazyEnablePeerAccess), + "cudaIpcOpenMemHandle", + ) + + def close(): + runtime.cudaIpcCloseMemHandle( + ptr + ) # release this process's mapping (not the exporter's) + + return int(ptr), close + + def _open_vmm(self): + # VMM: import the generic handle from its OS shareable handle, reserve + map + grant access. + from cuda.bindings import driver + + driver.cuInit(0) + if self.htype == H_WIN32_KMT: + htype = driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_WIN32_KMT + elif self.htype == H_POSIX_FD: + htype = driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + elif self.htype == H_WIN32: + htype = driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_WIN32 + else: + raise NotImplementedError(f"vmm handle_type {self.htype} not supported") + + handle_val = int.from_bytes(self._blob, "little") + (gen,) = _ck( + driver.cuMemImportFromShareableHandle(handle_val, htype), + "cuMemImportFromShareableHandle", + ) + size = self.alloc + (ptr,) = _ck(driver.cuMemAddressReserve(size, 0, 0, 0), "cuMemAddressReserve") + _ck(driver.cuMemMap(ptr, size, 0, gen, 0), "cuMemMap") + + acc = driver.CUmemAccessDesc() + acc.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + acc.location.id = self.device + acc.flags = driver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + _ck(driver.cuMemSetAccess(ptr, size, [acc], 1), "cuMemSetAccess") + + def close(): + driver.cuMemUnmap(ptr, size) + driver.cuMemAddressFree(ptr, size) + driver.cuMemRelease(gen) + + return int(ptr), close + + def _torch_dtype(self, torch): + return ( + None if self._format is None else getattr(torch, _TORCH_DTYPE[self._format]) + ) + + def to_torch(self, dtype=None, shape=None): + """Zero-copy CUDA torch tensor over the mapped buffer, valid inside the with-block. + + With open(payload), dtype and shape are taken from the payload geometry; overrides + may be passed explicitly (needed after open_descriptor without geometry).""" + import torch + + n = self.size - self.offset + t = torch.as_tensor(_CudaArray(self._ptr, n), device=f"cuda:{self.device}") + dt = dtype if dtype is not None else self._torch_dtype(torch) + if dt is not None and dt != torch.uint8: + t = t.view(dt) + sh = shape if shape is not None else self._shape + if sh is not None: + t = t.reshape(*sh) + return t + + tensor = to_torch # back-compat alias + + @property + def __cuda_array_interface__(self): + if self._shape is None or self._format is None: + raise RuntimeError( + "__cuda_array_interface__ needs geometry; open via open(payload) or use to_torch(dtype, shape)" + ) + return { + "shape": self._shape, + "typestr": _TYPESTR[self._format], + "data": (int(self._ptr), False), + "version": 3, + } + + def close(self): + if self._close is not None: + self._close() + self._close = None + + def __enter__(self): + return self + + def __exit__(self, *exc): + self.close() + + +def open(payload: bytes) -> ImportedCube: + """Open a payload from CudaImageData.export_payload(): maps the buffer and carries geometry.""" + (width, height, channels, fmt), descriptor = _unpack_payload(payload) + return ImportedCube(descriptor, shape=(height, width, channels), format_code=fmt) + + +def open_descriptor(descriptor_bytes: bytes, shape=None) -> ImportedCube: + """Advanced: open a raw 184-byte descriptor when you transport geometry yourself. + + Pass shape here (or dtype/shape to .to_torch()); prefer open(payload) for the easy path.""" + return ImportedCube(descriptor_bytes, shape=shape) + + +__all__ = [ + "ImportedCube", + "open", + "open_descriptor", + "pack_payload", + "BACKEND_NONE", + "BACKEND_POOL", + "BACKEND_LEGACY", + "BACKEND_VMM", +] diff --git a/pyproject.toml b/pyproject.toml index b438a4d..484cd40 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,9 @@ Issues = "https://github.com/cubert-hyperspectral/cuvis.python/issues" [tool.setuptools] packages = ["cuvis"] +# cuvis_ipc sits outside the package because a consumer process opening a CUDA IPC +# payload has no SDK, and importing `cuvis` requires one. +py-modules = ["cuvis_ipc"] include-package-data = true [tool.setuptools.package-data] diff --git a/tests/test_binding.py b/tests/test_binding.py new file mode 100644 index 0000000..9eaeaa6 --- /dev/null +++ b/tests/test_binding.py @@ -0,0 +1,65 @@ +""" +Tests for cuvis.binding. + +Covers availability reporting. The installed library differs between CI and developer +machines, so the tests pin the two ends that are always true: a function the binding +does expose is available, and one nothing exposes is not. +""" + +import pytest + +import cuvis +from cuvis import binding + +PRESENT = "cuvis_measurement_load" +ABSENT = "cuvis_a_function_that_does_not_exist" + + +def test_info_reports_without_an_initialised_sdk(): + """Test info() is answerable before cuvis.init, which is what gates a feature check.""" + current = binding.info() + assert isinstance(current, cuvis.BindingInfo) + assert isinstance(current.is_complete, bool) + assert "cuvis binding" in str(current) + + +def test_a_function_the_binding_exposes_is_available(): + """Test a function present in the binding and not reported missing is available.""" + assert binding.available(PRESENT) + assert binding.unavailable(PRESENT) == () + binding.require(PRESENT) + + +def test_a_function_no_binding_exposes_is_unavailable(): + """Test availability accounts for absent symbols, not only reported-missing ones. + + A binding too old to report missing symbols reports none, so consulting that list + alone would call a function the binding never had and fail with an AttributeError. + """ + assert not binding.available(ABSENT) + assert binding.unavailable(ABSENT) == (ABSENT,) + + +def test_unavailable_preserves_the_requested_order(): + """Test the report names only the unusable functions, in the order asked for.""" + assert binding.unavailable(PRESENT, ABSENT, PRESENT) == (ABSENT,) + + +def test_require_names_every_unavailable_function(): + """Test the raised error carries the names, so the message says what to install.""" + with pytest.raises(cuvis.UnavailableSDKFunction) as excinfo: + binding.require(PRESENT, ABSENT) + assert excinfo.value.names == (ABSENT,) + assert ABSENT in str(excinfo.value) + + +@pytest.mark.parametrize("expected", [RuntimeError, cuvis.cuvis_aux.SDKException]) +def test_unavailable_function_is_catchable_by_either_base(expected): + """Test the exception satisfies both except clauses its docstring advertises.""" + with pytest.raises(expected): + binding.require(ABSENT) + + +def test_missing_symbols_is_a_frozenset(): + """Test the reported set is immutable, so a caller cannot corrupt it.""" + assert isinstance(binding.missing_symbols(), frozenset) diff --git a/tests/test_cuda.py b/tests/test_cuda.py new file mode 100644 index 0000000..7ef719f --- /dev/null +++ b/tests/test_cuda.py @@ -0,0 +1,74 @@ +""" +Tests for cuvis.cuda capability reporting. + +The assertions hold whether or not the installed cuvis library provides the CUDA +functions, since CI and developer machines differ on that. What is pinned is that +capabilities() answers rather than raises, and that an unavailable function is +reported by name instead of surfacing as an AttributeError from the binding. +""" + +import pytest + +import cuvis +from cuvis import binding, cuda + +_HAS_DEVICE = binding.available(cuda.BACKEND_PROBE, *cuda.DEVICE_FUNCTIONS) + + +def test_capabilities_answers_without_raising(): + """Test capabilities() reports booleans instead of raising on a CUDA-less SDK.""" + caps = cuda.capabilities() + assert isinstance(caps, cuda.CudaCapabilities) + assert all(isinstance(value, bool) for value in caps) + assert isinstance(caps.any_ipc, bool) + + +def test_capabilities_is_safe_before_init(): + """Test capabilities() needs no initialised SDK, so a caller can gate on it first.""" + assert cuda.capabilities() == cuda.capabilities() + + +def test_cuda_mode_is_off_by_default(): + """Test CUDA stays opt-in, leaving the host refresh path in place.""" + assert cuda.is_enabled() is False + assert cuvis.Measurement._refresh_images is True + + +def test_disable_restores_the_host_refresh(): + """Test disable() is safe to call unconditionally and restores the host path.""" + cuda.disable() + assert cuda.is_enabled() is False + assert cuvis.Measurement._refresh_images is True + + +@pytest.mark.skipif(_HAS_DEVICE, reason="this SDK provides the CUDA functions") +def test_capabilities_are_false_without_the_sdk_functions(): + """Test a library lacking the CUDA functions reports no CUDA support.""" + caps = cuda.capabilities() + assert caps.same_process is False + assert caps.any_ipc is False + + +@pytest.mark.skipif(_HAS_DEVICE, reason="this SDK provides the CUDA functions") +@pytest.mark.parametrize("guard", [cuda.require_device, cuda.require_ipc, cuda.enable]) +def test_missing_functions_are_reported_by_name(guard): + """Test the CUDA entry points name what the SDK lacks (see cuvis.binding).""" + with pytest.raises(cuvis.UnavailableSDKFunction) as excinfo: + guard() + assert excinfo.value.names + assert all(name in str(excinfo.value) for name in excinfo.value.names) + assert cuda.is_enabled() is False + + +@pytest.mark.skipif(_HAS_DEVICE, reason="this SDK provides the CUDA functions") +def test_get_cube_cuda_reports_the_missing_functions(test_measurement): + """Test the device path fails with the diagnostic, not an AttributeError.""" + with pytest.raises(cuvis.UnavailableSDKFunction): + test_measurement.get_cube_cuda() + + +def test_unavailable_function_is_catchable_either_way(): + """Test UnavailableSDKFunction satisfies both except clauses it advertises.""" + for expected in (RuntimeError, cuvis.cuvis_aux.SDKException): + with pytest.raises(expected): + binding.require("cuvis_a_function_that_does_not_exist") diff --git a/tests/test_cuvis_ipc.py b/tests/test_cuvis_ipc.py new file mode 100644 index 0000000..0168f9b --- /dev/null +++ b/tests/test_cuvis_ipc.py @@ -0,0 +1,63 @@ +""" +Tests for the cuvis_ipc cross-process consumer module. + +Covers the payload codec and the property the module exists for: it is importable in +a process that has no cuvis SDK. Mapping a buffer needs a CUDA device and a live +exporting process, so that is not covered here. +""" + +import os +import subprocess +import sys + +import pytest + +import cuvis_ipc + +DESCRIPTOR = bytes(range(184)) +GEOMETRY = (290, 275, 51, 2) # width, height, channels, format code + + +def test_payload_round_trip(): + """Test a packed payload decodes back to the same descriptor and geometry.""" + payload = cuvis_ipc.pack_payload(DESCRIPTOR, *GEOMETRY) + geometry, descriptor = cuvis_ipc._unpack_payload(payload) + assert geometry == GEOMETRY + assert descriptor == DESCRIPTOR + + +def test_pack_payload_rejects_a_wrong_sized_descriptor(): + """Test the descriptor length is checked, since the wire format is fixed.""" + with pytest.raises(ValueError, match="184 bytes"): + cuvis_ipc.pack_payload(DESCRIPTOR[:-1], *GEOMETRY) + + +def test_unpack_rejects_foreign_bytes(): + """Test the magic guards against anything that is not a cuvis payload.""" + with pytest.raises(ValueError, match="magic"): + cuvis_ipc._unpack_payload(b"XXXX" + bytes(200)) + + +def test_unpack_rejects_a_future_version(): + """Test a payload from a newer wire format is refused rather than misread.""" + payload = bytearray(cuvis_ipc.pack_payload(DESCRIPTOR, *GEOMETRY)) + payload[4:8] = (cuvis_ipc._VERSION + 1).to_bytes(4, "little") + with pytest.raises(ValueError, match="version"): + cuvis_ipc._unpack_payload(bytes(payload)) + + +def test_importable_without_the_sdk(): + """Test the consumer module imports with no CUVIS environment variable set. + + This is the whole reason it sits outside the cuvis package: importing cuvis + requires the SDK, and a consumer process opening a payload does not have one. + """ + env = {k: v for k, v in os.environ.items() if k != "CUVIS"} + result = subprocess.run( + [sys.executable, "-c", "import cuvis_ipc, sys; print('cuvis' in sys.modules)"], + capture_output=True, + text=True, + env=env, + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "False"