Shared Memory Manager

Cross-process shared memory for multi-worker data loading.

See Also

• Memory Overview - Memory tools
• Memory Overview - Shared-memory tools
• Performance - Optimization
• Distributed - Multi-process training

---

datarax.memory.shared_memory_manager

Shared memory manager for multi-worker data pipeline scenarios.

logger module-attribute

logger = logging.getLogger(__name__)

SharedMemoryManager

SharedMemoryManager()

Manage shared memory arrays for multi-worker scenarios.

Automatically converts large numpy arrays to shared memory to avoid duplication across worker processes.

This is a plain resource manager, not a Flax NNX module: it owns multiprocessing.shared_memory blocks and plain metadata, none of which are traced JAX state. (An earlier version subclassed nnx.Module and stored numpy arrays / Python dicts inside nnx.Variable; that broke nnx.split/checkpointing and is deliberately avoided here.) Use it as a context manager, or call cleanup() explicitly, to release blocks.

shared_blocks instance-attribute

shared_blocks: dict[str, SharedMemory] = {}

array_metadata instance-attribute

array_metadata: dict[str, dict[str, Any]] = {}

make_shared

make_shared(name: str, array: Array, force: bool = False) -> Array

Convert array to shared memory.

Parameters:

Name	Type	Description	Default
name	str	Name for the shared memory block.	required
array	Array	Array to store in shared memory.	required
force	bool	If True, always use shared memory regardless of size.	False

Returns:

Type	Description
Array	The original array (shared memory is accessed via get_shared).

get_shared

get_shared(name: str) -> Array | None

Get shared array by name, or None if it was never stored.

cleanup

cleanup() -> None

Close and unlink all shared memory blocks, then clear tracking.