PreparedMap

A map staged once into a persistent region, run many times.

Usage

Source

PreparedMap(
    pool,
    fn,
    x,
    *,
    args=(),
    kwargs=None,
    n_chunks=None,
    seed=None,
    template=None,
    collect=None,
    stream=False,
)

Create with Pool.map_prepare(); run with run(), which returns the run’s map handle (the collect happens on it, as for Pool.map()). Re-arming is O(1) — a generation bump, a cursor reset, a cancel-word clear — and workers reuse their cached contexts, so a re-run pays neither the descriptor pickle nor the region create nor the worker-side re-attach. A view-collected run hands its region to the view; the next run restages into a fresh one (the mirror of mizu). Close with close() (or a with block) to unlink the region.

Parameters

pool: pymizu.Pool

The Pool to run on.

fn: _Callable[…, _Any]

The function to map, or a call spec — as for Pool.map().

x: _Any

The data to map over.

args: _Any = ()

Constant positional arguments appended to every call.

kwargs: dict | None = None

Constant keyword arguments for every call.

n_chunks: int | None = None

Overrides the morsel count (the scheduling granularity).

seed: int | bytes | bytearray | tuple[int | bytes | bytearray, int] | None = None

Deterministic per-element seeding — as for Pool.map().

template: _Any = None

An exemplar buffer declaring each result’s dtype and length — as for Pool.map().

collect: str | None = None

‘copy’ or ‘view’ with template; ‘list’ without.

stream: bool = False
Stream slices of x instead of staging it whole — as for Pool.map().

Examples

Stage the map once, then re-run it over replacement data at memcpy cost (same dtype and length swap in place):

import numpy as np
import pymizu

with pymizu.Pool.create(2) as pool:
    with pool.map_prepare(np.square, np.arange(4.0)) as pm:
        first = pm.run().collect(timeout=10)
        second = pm.run(np.array([4.0, -5.0, 6.0, 7.0])).collect(
            timeout=10
        )
second
[np.float64(16.0), np.float64(25.0), np.float64(36.0), np.float64(49.0)]

Methods

Name Description
close() Unlink the staged region (idempotent; the GC backstop is the
run() Run the prepared map once; return the run’s map handle.

close()

Unlink the staged region (idempotent; the GC backstop is the

Usage

Source

close()

capsule destructor).

run()

Run the prepared map once; return the run’s map handle.

Usage

Source

run(
    x=None,
    timeout=None,
)

Parameters

x: _Any = None

Replaces the staged data for this and later runs. When both the staged and the replacement x are raw-buffer eligible with the same dtype and length, the swap is an in-place memcpy over the region’s x section — the iterate-over-same-shape loop (optimizer steps, simulation sweeps) runs at memcpy cost, skipping the region create and the worker-side re-attach. Any other x restages transparently on this run.

timeout: float | None = None
Seconds before the run’s submission gives up; on expiry the tasks already in are cancelled and the TIMEOUT Sentinel is returned (no handle). The collection deadline is the handle’s collect().

Returns

The run’s map handle — collected exactly once, by
collect() (blocking) or collect_async() (event-driven). A clean value collect keeps the stage for the O(1) re-arm; a view collect hands the region to the view, and an unclean run (timeout, an error, a worker death) drops it — the next run restages into a fresh one either way.