grain.multiprocessing module

grain.multiprocessing module#

APIs to work with multiprocessing.

List of Members#

class grain.multiprocessing.MultiprocessingOptions(num_workers=0, per_worker_buffer_size=1, enable_profiling=False)[source]#

Options for using Python multiprocessing.

Example

With MultiprocessingOptions.num_workers=2, there will be 2 worker processes:

import grain

# Create a source with 4 elements.
source = grain.sources.RangeDataSource(0, 4, 1)

# Configure multiprocessing options.
multiprocessing_options = grain.MultiprocessingOptions(
    num_workers=2,
    per_worker_buffer_size=10,
)

# Use the multiprocessing options to create an iterator dataset.
ds = (
    grain.MapDataset.source(source)
    .to_iter_dataset()
    .batch(2)
    .mp_prefetch(multiprocessing_options)
)
print(list(ds))
# [SharedMemoryArray([0, 2]), SharedMemoryArray([1, 3])]
Parameters:

  • num_workers (int)

  • per_worker_buffer_size (int)

  • enable_profiling (bool)

num_workers#

Number of Python worker processes. More processes can speed up the pipeline if it’s compute bound and bottlenecked on the CPython’s GIL. The default value of 0 means no Python multiprocessing, and as a result all data loading and transformation will run in the main Python process.

Type:

int

per_worker_buffer_size#

Size of the buffer for preprocessed elements that each worker maintains. These are elements after all transformations. If your transformations include batching this means a single element is a batch.

Type:

int

enable_profiling#

If True, profiling info is logged. This is only available when num_workers >= 1.

Type:

bool

class grain.multiprocessing.SharedMemoryArray(shape, dtype=<class 'float'>)[source]#

A NumPy array subclass which is backed by shared memory.

This should be used in combination with Python multiprocessing. Compared with the normal NumPy ndarray it avoids expensive serialization when sending a root array to another Python process (on the same machine). It also doesn’t require a copy on the receiving side. Views/slices of a SharedMemoryArray do not carry offset or stride metadata and are therefore serialized by value as standard NumPy ndarrays (or copied to a new shared memory block by copy_to_shm).

Exactly one owning process (typically the consumer opening the array via open_from_shm) must call unlink_on_del() on the root array, or call metadata.close_and_unlink_shm(). Otherwise the shared memory block will not be freed.

Parameters:

  • shape (Iterable[int])

  • dtype (npt.DTypeLike)

classmethod enable_async_del(num_threads=1, max_outstanding_requests=50)[source]#

Enables asynchronous deletion of shared memory arrays.

Parameters:

  • num_threads (int) – The number of threads to use for deletion.

  • max_outstanding_requests (int) – The maximum number of outstanding requests to close/unlink shared memory. A larger value may make the __del__ method faster, but it may also lead to OOM errors or hitting file descriptor limits, since max_outstanding_requests shared memory objects and their associated file descriptors may be buffered before deletion.

Return type:

None

Mark this object responsible for unlinking the shared memory.

Return type:

None

class grain.multiprocessing.SharedMemoryArrayMetadata(name, shape, dtype)[source]#

Metadata for a SharedMemoryArray.

Parameters:

  • name (str)

  • shape (Iterable[int])

  • dtype (numpy.typing.DTypeLike)

Closes and unlinks the shared memory. No-op if already unlinked.

Return type:

None

Contents