MemoryManager

class cubie.memory.mem_manager.MemoryManager(registry: dict[int, InstanceMemorySettings] = NOTHING, stream_groups: StreamGroups = NOTHING, mode: str = 'passive', auto_pool: list[int] = NOTHING, manual_pool: list[int] = NOTHING, queued_allocations: Dict[str, Dict] = NOTHING, group_chunk_parameters: Dict[Tuple[str, int], Tuple[int, int]] = NOTHING, pinned_max_bytes: int | None = None, allocation_granule_bytes: int = 33554432)

Bases: object

Singleton interface coordinating GPU memory allocation and stream usage.

Parameters:

• registry (dict[int, cubie.memory.mem_manager.InstanceMemorySettings]) – Registry mapping instance identifiers to their memory settings.

• stream_groups (cubie.memory.stream_groups.StreamGroups) – Manager for organising instances into stream groups.

totalmem

Total GPU memory in bytes as of the last probe_device(), or None when no device answered.

Type:

int | None

Notes

The manager accepts ArrayRequest objects and returns ArrayResponse instances that reference allocated arrays and chunking information. Active mode enforces per-instance VRAM proportions while passive mode mirrors standard allocation behaviour using chunking only when necessary.

Construction succeeds without a device: the manager stays unsized and every decision needing a byte figure reprobes, then raises NoCudaDeviceError if no device answers.

See also

ArrayRequest

Describes a single allocation request.

ArrayResponse

Reports allocation outcomes.

StreamGroups

Manages CUDA stream groups.

_add_auto_proportion(instance: object) → float

Add an instance to the auto allocation pool with equal share.

Parameters:

instance – Instance to add to auto allocation pool.

Returns:

Proportion assigned to this instance.

Return type:

float

Raises:

ValueError – If available auto-allocation pool is less than minimum required size.

Notes

Splits the non-manually-allocated portion of VRAM equally among all auto-allocated instances. Triggers rebalancing of the auto pool.

_add_manual_proportion(instance: object, proportion: float) → None

Add an instance to the manual allocation pool with the specified proportion.

Parameters:

• instance – Instance to add to manual allocation pool.

• proportion – Memory proportion to assign (0.0 to 1.0).

Raises:

ValueError – If manual proportion would exceed total available memory or leave insufficient memory for auto-allocated processes.

Warning

UserWarning

If manual proportion leaves less than 5% of memory for auto allocation.

Notes

Updates the instance’s proportion and cap, then rebalances the auto pool. Enforces minimum auto pool size constraints.

_apply_pinned_releases() → None

Move queued finalizer releases from live to retained.

_cap_bytes(proportion: float) → int | None

Bytes for proportion of VRAM; None with no device.

_check_requests(requests: dict[str, ArrayRequest]) → None

Validate that all requests are properly formatted.

Parameters:

requests – Dictionary of requests to validate.

Raises:

TypeError – If requests is not a dict or contains invalid ArrayRequest objects.

_create_spill_array(shape: tuple[int, ...], dtype: DTypeLike) → memmap

Create a disk-backed array in the cache root.

_drop_instance(instance_id: int) → None

Free and deregister one instance.

_evict_idle_owners(exclude_ids: set[int], required_bytes: int) → int

Free the least-recently-used idle owners.

Whole owners are evicted, oldest first, until required_bytes are released. Returns the number of bytes actually released.

_on_pinned_released(nbytes: int) → None

Queue a collected pinned array’s bytes without locking.

_owned_by(instance_id: int, owner_id: int) → bool

Return whether a registered instance belongs to owner_id.

_owner_settings(owner_id: int) → list[InstanceMemorySettings]

Return registry entries owned by one client.

_purge_dead_instances() → None

Drop registry entries whose instance has been collected.

Registered instances are held weakly. Once an instance is garbage collected, its manual or auto reservation is released, its stream-group membership is removed, and any allocations it still had queued are discarded. Teardowns recorded by GC finalizers run first, here, in normal call context.

_rebalance_auto_pool() → None

Redistribute available memory equally among auto-allocated instances.

Notes

Calculates the available proportion after manual allocations and divides it equally among all instances in the auto pool. Updates both proportion and cap for each auto-allocated instance.

_require_device() → None

Reprobe an unsized manager, then require a device size.

Raises:

NoCudaDeviceError – If the device probe failed, chained to its error.

_reserve_pinned_bytes(nbytes: int, force: bool = False) → bool

Atomically reserve nbytes against the pinned budget.

Live and retained bytes count together; pressure empties the CuPy pinned pool before a refusal. force reserves past the budget.

allocate(shape: tuple[int, ...], dtype: Callable, memory_type: str, stream: cuda.cudadrv.driver.Stream = 0) → object

Allocate a single C-contiguous array with specified parameters.

Parameters:

• shape – Shape of the array to allocate.

• dtype – Constructor returning the precision object for the array elements.

• memory_type – Type of memory: “device” or “pinned”.

• stream – CUDA stream for the allocation. Defaults to 0.

Returns:

Allocated GPU array.

Return type:

object

Raises:

ValueError – If memory_type is not “device” or “pinned”.

allocate_all(requests: dict[str, ArrayRequest], instance_id: int, stream: cuda.cudadrv.driver.Stream, settings: InstanceMemorySettings | None = None) → dict[str, object]

Allocate multiple arrays based on a dictionary of requests.

Parameters:

• requests – Dictionary mapping labels to array requests.

• instance_id – ID of the requesting instance.

• stream – CUDA stream for the allocations.

• settings – Registration to record the allocations in. Passing it avoids a registry lookup that can fail when the client is garbage collected between queueing and allocation.

Returns:

Dictionary mapping labels to allocated arrays.

Return type:

dict of str to object

allocate_pinned_array(shape: Tuple[int, ...], dtype: DTypeLike, force: bool = False) → ndarray | None

Allocate one budget-accounted pinned host array.

Reserves bytes, attempts the driver allocation, and attaches a finalizer that releases the bytes once the array and every view of it are collected.

Parameters:

• shape – Shape of the array to allocate.

• dtype – Data type for the array elements.

• force – Reserve even past the budget. A driver failure then propagates instead of returning None.

Returns:

The pinned array, or None when the budget or the driver refuses.

Return type:

numpy.ndarray or None

Raises:

NoCudaDeviceError – If the last device probe failed.

allocate_queue(triggering_instance: object, stream: FakeStream | None = None) → None

Process all queued requests for a stream group with coordinated chunking.

Chunking is always performed along the run axis when memory constraints require splitting the batch.

Parameters:

triggering_instance – The instance that triggered queue processing.

Notes

Processes the pending requests of the triggering instance’s owner within its stream group, applying coordinated chunking based on available memory, and calls allocation_ready_hook for each of those instances with their results. Requests queued by other owners in the group stay queued until their own owner triggers. The owner’s instances with no queued requests receive an empty response carrying the owner’s chunk parameters. When nothing of the owner’s is queued the call returns at once.

allocation_granule_bytes: int

property auto_pool_proportion

Total proportion of VRAM currently distributed automatically.

begin_work(owner: object) → None

Mark an owner as submitting CUDA work.

cap(instance: object) → int | None

Get the maximum allocatable bytes for an instance.

Parameters:

instance – Instance to query.

Returns:

Maximum allocatable bytes for this instance.

Return type:

int or None

change_stream_group(instance: object, new_group: str) → None

Move the owner’s registrations, queued requests and partition.

Parameters:

• instance – Instance to move.

• new_group – Name of the new stream group.

choose_host_memory_type(nbytes: int, allow_pinned: bool = True) → str

Pick the backing for a host array of nbytes bytes.

Memmap above HOST_SPILL_FRACTION of RAM, pinned within pinned_max_bytes, else pageable.

Parameters:

• nbytes – Size of the array in bytes.

• allow_pinned – Permit the "pinned" choice; chunked staging passes False.

Returns:

"pinned", "host", or "memmap".

Return type:

str

Raises:

NoCudaDeviceError – If a pinned choice is reachable but no device answers the probe.

compute_chunked_shapes(requests: dict[str, ArrayRequest], chunk_size: int) → dict[str, Tuple[int, ...]]

Compute per-array chunked shapes based on available memory.

Parameters:

• requests – Dictionary mapping labels to array requests.

• chunk_size – Length of chunked arrays along run axis

Returns:

Mapping from array labels to their per-chunk shapes.

Return type:

dict[str, tuple[int, …]]

Notes

Unchunkable arrays retain their original shape.

create_host_array(shape: tuple[int, ...], dtype: DTypeLike, memory_type: str = 'pinned', like: ndarray | None = None) → ndarray

Create a C-contiguous host array.

Parameters:

• shape – Shape of the array to create.

• dtype – Data type for the array elements.

• memory_type – "pinned", "host", or "memmap".

• like – Optional source data.

Returns:

C-contiguous host array; a numpy.memmap in the cache root when spilled. A "pinned" request whose reservation or driver allocation fails lands pageable.

Return type:

numpy.ndarray

Raises:

ValueError – If memory_type is not "pinned", "host", or "memmap".

defer_teardown(teardown: Callable[[], None]) → None

Record a teardown to run at the manager’s next entry point.

Garbage-collection finalizers call this instead of deregistering directly: GC runs inside whatever allocation triggered it — possibly while this manager iterates its registry, possibly on the transfer watcher’s own thread — so the callback only appends. _purge_dead_instances runs the recorded teardowns in normal call context.

end_work(owner: object, stream: FakeStream) → None

Record completion of an owner’s submitted CUDA work.

free(array_label: str) → None

Free an allocation by label across all instances.

Parameters:

array_label – Label of the allocation to free.

free_all() → None

Free all allocations across all registered instances.

from_device(instance: object, from_arrays: list[object], to_arrays: list[object], stream: FakeStream | None = None) → None

Copy data from device arrays using the instance’s stream.

Parameters:

• instance – Instance whose stream to use for copying.

• from_arrays – Source device arrays to copy from.

• to_arrays – Destination arrays to copy to.

get_available_memory(group: str) → int

Get available memory for an entire stream group.

Parameters:

group – Name of the stream group.

Returns:

Available memory in bytes for the group.

Return type:

int

Warning

UserWarning

If group has used more than 95% of allocated memory.

Raises:

NoCudaDeviceError – If no device answers the probe.

get_chunk_parameters(requests: Dict[str, Dict], axis_length: int, stream_group: str) → Tuple[int, int]

Calculate number of chunks and chunk size for a dict of array requests.

Chunking is performed along the run axis only.

Parameters:

• requests – Dictionary mapping instance IDs to their array requests.

• axis_length – Unchunked length of the chunking axis.

• stream_group – Name of the stream group making the request.

Returns:

Length of chunked axis (ceil(axis_length / num_chunks)) and number of chunks needed to fit the request.

Return type:

int, int

Warning

UserWarning

If request exceeds available VRAM by more than 20x.

Raises:

NoCudaDeviceError – If no device answers the probe.

Notes

The request is offered min((1 - CHUNK_HEADROOM_FRACTION) × available, free - allocation_granule_bytes) bytes.

get_group_stream(group: str = 'default') → FakeStream | int

Get the dedicated stream for a named group, creating it if needed.

Parameters:

group – Name of the stream group.

Returns:

The group’s dedicated CUDA stream. Because each process owns its own manager, this stream is private to the process; synchronizing it never orders against the device-wide default stream.

Return type:

Stream

get_memory_info() → tuple[int, int]

Get free and total GPU memory information.

Returns:

(free_memory, total_memory) in bytes.

Return type:

tuple of int

get_registration(instance: object) → InstanceMemorySettings

Return the registry entry for a registered instance.

Parameters:

instance – Registered instance to look up.

Raises:

KeyError – If the instance is not registered.

get_stream(instance: object) → object

Get the CUDA stream associated with an instance.

Parameters:

instance – Instance to retrieve the stream for.

Returns:

CUDA stream associated with the instance.

Return type:

object

get_stream_group(instance: object) → str

Get the name of the stream group for an instance.

Parameters:

instance – Instance to query.

Returns:

Name of the stream group.

Return type:

str

invalidate_all() → None

Call each invalidate hook and release all allocations.

is_grouped(instance: object) → bool

Check if instance is grouped with others in a named stream.

Parameters:

instance – Instance to check.

Returns:

True if instance shares a stream group with other instances.

Return type:

bool

property manual_pool_proportion

Total proportion of VRAM currently assigned manually.

manual_proportion(instance: object) → float | None

Return an instance’s reserved proportion; None if auto.

Parameters:

instance – Instance to query.

Returns:

Manual proportion, or None for auto-pool instances.

Return type:

float or None

property pinned_budget_bytes: int

pinned_max_bytes capped at the spill fraction of RAM.

Raises:

NoCudaDeviceError – If the last device probe failed.

property pinned_live_bytes: int

Pinned bytes currently backing reachable arrays.

pinned_max_bytes: int | None

property pinned_retained_bytes: int

Pinned bytes held only by CuPy’s pinned pool.

probe_device() → None

Read the device’s memory size.

An absent device stores its error and leaves the sizes None.

Raises:

Exception – The probe’s own error, when it does not signal an absent device.

proportion(instance: object) → float

Get the maximum proportion of VRAM allocated to an instance.

Parameters:

instance – Instance to query.

Returns:

Proportion of VRAM allocated to this instance.

Return type:

float

queue_request(instance: object, requests: dict[str, ArrayRequest]) → None

Queue allocation requests for batched stream group processing.

Parameters:

• instance – The instance making the request.

• requests – Dictionary mapping labels to array requests.

Raises:

ValueError – If the instance registered without a live invalidate hook. Allocation holders are eviction candidates, so they must be able to drop handles when their buffers are freed.

Notes

Requests are queued per stream group, allowing multiple components to contribute to a single coordinated allocation that can be optimally chunked together.

register(instance: object, proportion: float | None = None, invalidate_cache_hook: Callable = <function placeholder_invalidate>, allocation_ready_hook: Callable = <function placeholder_dataready>, stream_group: str = 'default', owner: object | None = None) → None

Register an instance and configure its memory allocation settings.

Parameters:

• instance – Instance to register for memory management.

• proportion – Proportion of VRAM to allocate (0.0 to 1.0). When omitted, the instance joins the automatic allocation pool.

• invalidate_cache_hook – Function to call when CUDA memory system changes occur.

• allocation_ready_hook – Function to call when allocations are ready.

• stream_group – Name of the stream group to assign the instance to.

• owner – Eviction unit for this registration. A solver kernel passes itself when registering its input and output array managers, so pressure evicts the whole solver or nothing. Defaults to the instance itself.

Raises:

ValueError – If instance is already registered or proportion is not between 0 and 1.

registry: dict[int, InstanceMemorySettings]

reinit_streams() → None

Reinitialise all streams after a CUDA context reset.

release_host_array(array: ndarray) → None

Release a spill-backed host array.

release_instance(instance_id: int, settings: InstanceMemorySettings) → None

Release a matching registered instance.

set_auto_limit_mode(instance: object) → None

Convert a manual-limited instance to auto allocation mode.

Parameters:

instance – Instance to convert to auto mode.

Notes

If the instance is already in the auto pool, this is a no-op.

set_limit_mode(mode: str) → None

Set the memory allocation limiting mode.

Parameters:

mode – Either "passive" or "active" memory management mode.

Raises:

ValueError – If mode is not “passive” or “active”.

set_manual_limit_mode(instance: object, proportion: float) → None

Convert an auto-limited instance to manual allocation mode.

Parameters:

• instance – Instance to convert to manual mode.

• proportion – Memory proportion to assign (0.0 to 1.0).

Notes

If the instance is already in the manual pool, this is a no-op.

set_manual_proportion(instance: object, proportion: float) → None

Set manual allocation proportion for an instance.

If instance is currently in the auto-allocation pool, shift it to manual.

Parameters:

• instance – Instance to update proportion for.

• proportion – New proportion between 0 and 1.

Raises:

ValueError – If proportion is not between 0 and 1.

stream_groups: StreamGroups

sync_stream(instance: object, stream: FakeStream | None = None) → None

Synchronize the CUDA stream for an instance.

Parameters:

instance – Instance whose stream to synchronize.

to_device(instance: object, from_arrays: list[object], to_arrays: list[object], stream: FakeStream | None = None) → None

Copy data to device arrays using the instance’s stream.

Parameters:

• instance – Instance whose stream to use for copying.

• from_arrays – Source arrays to copy from.

• to_arrays – Destination device arrays to copy to.

totalmem: int | None

update(instance: object, updates_dict: Dict[str, Any] | None = None, silent: bool = False, **kwargs: Any) → Set[str]

Update a registered instance’s memory settings.

Parameters:

• instance – Registered instance to update.

• updates_dict – Setting names to new values: stream_group, mem_proportion (None selects the automatic limit) and memory_manager (this manager only).

• silent – Ignore unknown names instead of raising.

• **kwargs – Further updates.

Returns:

The recognised names.

Return type:

Set[str]

Raises:

• KeyError – Unknown names when not silent.

• ValueError – A different memory manager.