vibespatial.raster.dispatch¶

VRAM budget functions and band dispatch executors for multiband GPU processing.

Provides utilities to query available GPU memory, compute how many raster bands can be processed in a single GPU pass, and dispatch per-band operations across multiband rasters on both GPU and CPU paths.

When CuPy is unavailable, available_vram_bytes() returns 0 gracefully and dispatch_per_band_gpu will fail with a clear error at call time.

Functions¶

`analyze_raster_plan`(...)	Decide the processing strategy for a raster given VRAM constraints.
`available_vram_bytes`(→ int)	Return the effective available VRAM in bytes after headroom.
`dispatch_per_band_cpu`(...)	Apply op_fn to each band of raster on the CPU, then reassemble.
`dispatch_per_band_gpu`(...)	Apply op_fn to each band of raster on the GPU, then reassemble.
`max_bands_for_budget`(→ int)	Compute how many raster bands fit in available VRAM.
`plan_from_metadata`(→ vibespatial.raster.buffers.RasterPlan)	Convenience wrapper: extract dimensions from `RasterMetadata`.

Module Contents¶

vibespatial.raster.dispatch.analyze_raster_plan(height: int, width: int, dtype: numpy.dtype, *, band_count: int = 1, buffers_per_band: int = 2, scratch_bytes: int = 0, halo: int = 0, vram_budget: int | None = None) → vibespatial.raster.buffers.RasterPlan¶

Decide the processing strategy for a raster given VRAM constraints.

Parameters¶

height, width:: Spatial dimensions of the raster.
dtype:: NumPy dtype (e.g. np.float32).
band_count:: Number of bands.
buffers_per_band:: Device buffers required per band (default 2: input + output).
scratch_bytes:: Additional fixed scratch memory consumed by the operation.
halo:: Overlap pixels for stencil/focal operations. Each tile’s effective data area is (tile_H - 2*halo, tile_W - 2*halo).
vram_budget:: Available VRAM in bytes. None triggers auto-detection via available_vram_bytes().

Returns¶

RasterPlan: Frozen dataclass describing the strategy, tile dimensions, and estimated VRAM per tile.

vibespatial.raster.dispatch.available_vram_bytes() → int¶

Return the effective available VRAM in bytes after headroom.

When RMM is the active allocator (tiers A/B/C), the function queries rmm.mr.available_device_memory which accounts for pool-managed blocks. Otherwise it falls back to the CuPy pool query.

A 15 % headroom fraction is subtracted from the effective free memory to leave breathing room for the CUDA driver, fragmentation, and any concurrent allocations.

Returns 0 when CuPy is not importable or no CUDA device is available, making the function safe to call unconditionally on CPU-only machines.

vibespatial.raster.dispatch.dispatch_per_band_cpu(raster: vibespatial.raster.buffers.OwnedRasterArray, op_fn: collections.abc.Callable[[vibespatial.raster.buffers.OwnedRasterArray], vibespatial.raster.buffers.OwnedRasterArray]) → vibespatial.raster.buffers.OwnedRasterArray¶

Apply op_fn to each band of raster on the CPU, then reassemble.

For single-band rasters this is a zero-overhead passthrough.

For multiband rasters each band is sliced from the host numpy array, wrapped as a single-band OwnedRasterArray, passed to op_fn, and the per-band results are assembled via OwnedRasterArray.from_band_stack().

Parameters¶

rasterOwnedRasterArray: Input raster (single- or multi-band).
op_fnCallable[[OwnedRasterArray], OwnedRasterArray]: Operation to apply per band. Receives a single-band OwnedRasterArray and must return a single-band OwnedRasterArray.

Returns¶

OwnedRasterArray: Result raster with the same band count, affine, CRS, and nodata as the input.

vibespatial.raster.dispatch.dispatch_per_band_gpu(raster: vibespatial.raster.buffers.OwnedRasterArray, op_fn: collections.abc.Callable[[vibespatial.raster.buffers.OwnedRasterArray], vibespatial.raster.buffers.OwnedRasterArray], *, buffers_per_band: int = 2, scratch_bytes: int = 0) → vibespatial.raster.buffers.OwnedRasterArray¶

Apply op_fn to each band of raster on the GPU, then reassemble.

For single-band rasters this is a zero-overhead passthrough: op_fn is called once and its result is returned directly.

For multiband rasters the full raster is transferred to device once, then each band is sliced as a zero-copy 2D view and passed to op_fn. The per-band results are assembled via OwnedRasterArray.from_band_stack().

Parameters¶

rasterOwnedRasterArray: Input raster (single- or multi-band).
op_fnCallable[[OwnedRasterArray], OwnedRasterArray]: Operation to apply per band. Receives a single-band OwnedRasterArray and must return a single-band OwnedRasterArray.
buffers_per_bandint: Number of device buffers the operation needs per band (used by max_bands_for_budget for callers that want to pre-plan chunking; not consumed directly by this executor).
scratch_bytesint: Fixed scratch memory consumed by the operation (same caveat as buffers_per_band).

Returns¶

OwnedRasterArray: Result raster with the same band count, affine, CRS, and nodata as the input (metadata propagation is handled by from_band_stack).

vibespatial.raster.dispatch.max_bands_for_budget(height: int, width: int, dtype: numpy.dtype, buffers_per_band: int = 2, scratch_bytes: int = 0) → int¶

Compute how many raster bands fit in available VRAM.

Parameters¶

height, width:: Spatial dimensions of each band.
dtype:: NumPy dtype of the raster (e.g. np.float32). Used to determine per-element byte width via dtype.itemsize.
buffers_per_band:: Number of device buffers required per band (default 2 — one input and one output buffer).
scratch_bytes:: Additional fixed scratch memory consumed by the operation, subtracted from the VRAM budget before dividing by per-band cost.

Returns¶

int: Maximum number of bands that fit, but always at least 1 so that a single-band fallback is always possible.

vibespatial.raster.dispatch.plan_from_metadata(metadata: vibespatial.raster.buffers.RasterMetadata, *, buffers_per_band: int = 2, scratch_bytes: int = 0, halo: int = 0, vram_budget: int | None = None) → vibespatial.raster.buffers.RasterPlan¶

Convenience wrapper: extract dimensions from RasterMetadata.

Parameters¶

metadata:: Raster metadata (from read_raster_metadata() or raster.metadata).
buffers_per_band, scratch_bytes, halo, vram_budget:: Forwarded to analyze_raster_plan().

Returns¶

RasterPlan