pyodide-pool 0.1.0__tar.gz

pyodide_pool-0.1.0/.gitignore

@@ -0,0 +1,12 @@
+node_modules/
+dist/
+.venv/
+test-results/
+_output/
+bench/results/node-smoke*
+playwright-report/
+.jupyterlite.doit.db
+demos/jupyterlite/files/assets/
+demos/jupyterlite/files/wheels/
+demos/jupyterlite/pypi/
+demos/jupyterlite/.cache/

pyodide_pool-0.1.0/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: pyodide-pool
+Version: 0.1.0
+Summary: Run Python callables and dask graphs on a pool of Pyodide Web Workers (driver-side API)
+License-Expression: MIT
+Requires-Python: >=3.11
+Requires-Dist: cloudpickle
+Provides-Extra: dask
+Requires-Dist: dask; extra == 'dask'

pyodide_pool-0.1.0/pyodide_pool/__init__.py

@@ -0,0 +1,83 @@
+"""Run Python callables and dask graphs on a pool of Pyodide Web Workers.
+
+This package runs in the MAIN Pyodide instance (the driver). The embedding
+JS creates a ``PyodidePool`` and registers it before use::
+
+    pyodide.registerJsModule('js_pyodide_pool', { pool })
+
+Inside a Pyodide kernel with no JS embedder (JupyterLite notebooks), let
+``loader`` pull in the self-contained JS bundle instead::
+
+    pool = await pyodide_pool.create_pool(pool_size=4)
+
+Driver usage::
+
+    import pyodide_pool
+
+    value = await pyodide_pool.submit(fn, *args, **kwargs)  # one remote call
+    result = await pyodide_pool.compute(delayed_obj)        # dask collections
+    values = await pyodide_pool.get(dsk, keys)              # raw dask graph
+
+Design: docs/architecture/dask-scheduler-design.md. The package is pure
+Python with no hard dask dependency — ``compute``/``get`` import dask lazily
+via ``scheduler`` — so it can also back the Phase 06 multiprocessing shim.
+"""
+
+from __future__ import annotations
+
+import sys
+from typing import Any
+
+import cloudpickle
+
+from ._bridge import (
+    RemoteExecutionError,
+    RemoteTraceback,
+    WorkerPool,
+    default_pool,
+    set_default_pool,
+    submit,
+)
+from ._packages import EXCLUDED_FROM_MIRROR, PackageSnapshot, snapshot_packages
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "EXCLUDED_FROM_MIRROR",
+    "PackageSnapshot",
+    "RemoteExecutionError",
+    "RemoteTraceback",
+    "WorkerPool",
+    "compute",
+    "create_pool",
+    "default_pool",
+    "get",
+    "set_default_pool",
+    "snapshot_packages",
+    "submit",
+]
+
+# Helpers from this package that end up inside task payloads (the
+# scheduler's nested-expression evaluator, user wrappers) must travel BY
+# VALUE so workers never need pyodide_pool installed. Registering the
+# package covers all submodules — cloudpickle walks parent packages.
+cloudpickle.register_pickle_by_value(sys.modules[__name__])
+
+
+def __getattr__(name: str) -> Any:
+    # compute/get live in .scheduler, resolved lazily: scheduler pulls in
+    # dask, which must not become an import-time dependency of the package.
+    if name in ("compute", "get"):
+        from . import scheduler
+
+        value = getattr(scheduler, name)
+        globals()[name] = value
+        return value
+    # create_pool lives in .loader, resolved lazily for symmetry: it is only
+    # meaningful where the JS bundle is reachable (browser kernels, Node).
+    if name == "create_pool":
+        from .loader import create_pool
+
+        globals()[name] = create_pool
+        return create_pool
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

pyodide_pool-0.1.0/pyodide_pool/_bridge.py

@@ -0,0 +1,236 @@
+"""Driver-side bridge to the JS ``PyodidePool``.
+
+This module runs in the MAIN Pyodide instance (the driver), never in
+workers: worker-executed code travels inside cloudpickle payloads (the
+package is registered pickle-by-value in ``__init__``), so workers never
+import this package and it is never installed there.
+
+``submit`` (one call over ``runPickled``) and ``submit_batch`` (many calls
+as one ``mapPickled`` run) are the only choke points between driver Python
+and the pool: cloudpickle each ``(func, args, kwargs)`` call triple, hand
+the bytes to the JS pool as transferable buffers together with the current
+package snapshot, await the JS promise (directly awaitable on Pyodide's
+event loop), and unpickle the results — or re-raise the remote exception
+with its original traceback attached.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any, NoReturn
+
+import cloudpickle
+from js import Object, Uint8Array
+from pyodide.ffi import JsException, to_js
+
+from ._packages import snapshot_packages
+
+__all__ = [
+    "BatchRun",
+    "RemoteExecutionError",
+    "RemoteTraceback",
+    "WorkerPool",
+    "default_pool",
+    "set_default_pool",
+    "submit",
+]
+
+#: One driver call shipped to a worker: ``(func, args, kwargs)``.
+Call = tuple[Any, tuple[Any, ...], dict[str, Any]]
+
+
+class RemoteTraceback(Exception):
+    """Formatted traceback of an exception raised on a pool worker.
+
+    Attached as ``__cause__`` when the original exception re-raises on the
+    driver (the concurrent.futures pattern), so the user sees where their
+    code actually failed, not where the bridge noticed.
+    """
+
+    def __init__(self, formatted: str) -> None:
+        super().__init__(formatted)
+        self.formatted = formatted
+
+    def __str__(self) -> str:
+        return f'\n"""\n{self.formatted}"""'
+
+
+class RemoteExecutionError(Exception):
+    """A worker task failed and the original exception could not be
+    reconstructed (it was unpicklable); carries the remote traceback text."""
+
+    def __init__(self, message: str, remote_traceback: str | None = None) -> None:
+        super().__init__(message)
+        self.remote_traceback = remote_traceback
+
+
+def _resolve_js_pool() -> Any:
+    try:
+        import js_pyodide_pool  # registered by the embedding JS
+    except ImportError as exc:
+        raise RuntimeError(
+            "No JS worker pool is registered. Create a PyodidePool in JS and "
+            "expose it with pyodide.registerJsModule('js_pyodide_pool', "
+            "{ pool }) before using pyodide_pool."
+        ) from exc
+    pool = getattr(js_pyodide_pool, "pool", None)
+    if pool is None:
+        raise RuntimeError(
+            "The js_pyodide_pool module has no 'pool' attribute; register the "
+            "JS pool as pyodide.registerJsModule('js_pyodide_pool', { pool })."
+        )
+    return pool
+
+
+def _to_js_buffer(data: bytes) -> Any:
+    """bytes -> plain JS ArrayBuffer (transferable; runPickled detaches it)."""
+    view = Uint8Array.new(len(data))
+    view.assign(data)
+    return view.buffer
+
+
+def _to_py_bytes(js_buffer: Any) -> bytes:
+    """JS ArrayBuffer (or any view over one) -> bytes."""
+    return Uint8Array.new(js_buffer).to_bytes()
+
+
+def _raise_remote(exc: JsException) -> NoReturn:
+    """Re-raise a failed ``runPickled`` as the original Python exception.
+
+    ``JsException`` is itself a JsProxy of the JS error (``PyodideTaskError``),
+    so the worker's ``pythonTraceback`` and cloudpickled ``exceptionPayload``
+    read straight off the caught exception. Undefined properties surface as
+    AttributeError, which ``getattr`` defaults cover.
+    """
+    formatted = getattr(exc, "pythonTraceback", None)
+    payload = getattr(exc, "exceptionPayload", None)
+    if payload is not None:
+        try:
+            remote = cloudpickle.loads(_to_py_bytes(payload))
+        except BaseException:
+            remote = None
+        if isinstance(remote, BaseException):
+            if isinstance(formatted, str) and formatted:
+                raise remote from RemoteTraceback(formatted)
+            raise remote from exc
+    message = getattr(exc, "message", None) or str(exc)
+    raise RemoteExecutionError(message, remote_traceback=formatted) from exc
+
+
+class BatchRun:
+    """Driver-side handle over one JS ``mapPickled`` run.
+
+    Returned by :meth:`WorkerPool.start_batch` so callers (the Phase 06
+    multiprocessing shim's ``terminate``) can cancel queued-not-started
+    payloads while the run is in flight; plain callers use
+    :meth:`WorkerPool.submit_batch` and never see this class.
+    """
+
+    def __init__(self, js_run: Any) -> None:
+        self._js_run = js_run
+
+    @property
+    def run_id(self) -> int:
+        """worker-pool run id; ``-1`` for an empty batch (nothing to cancel)."""
+        return int(self._js_run.runId)
+
+    def cancel(self) -> None:
+        """Cancel payloads not yet started; :meth:`results` then raises the
+        worker-pool cancellation error. In-flight payloads finish."""
+        self._js_run.cancel()
+
+    async def results(self) -> list[Any]:
+        """Await the whole run; unpickled results in input order. A failed
+        payload re-raises its original remote exception — ``mapPickled``
+        rejects with the first failure once the run settles, matching stdlib
+        ``Pool.map`` surfacing one exception."""
+        try:
+            buffers = await self._js_run.promise
+        except JsException as exc:
+            _raise_remote(exc)
+        return [cloudpickle.loads(_to_py_bytes(buffer)) for buffer in buffers]
+
+
+class WorkerPool:
+    """Thin Python handle over the JS ``PyodidePool``.
+
+    Wraps the pool registered as ``js_pyodide_pool.pool`` by default; pass a
+    JsProxy of another pool for explicit wiring (tests, multiple pools).
+    """
+
+    def __init__(self, js_pool: Any = None) -> None:
+        self._js_pool = js_pool if js_pool is not None else _resolve_js_pool()
+
+    @property
+    def pool_size(self) -> int:
+        """Maximum number of concurrent workers in the JS pool."""
+        return int(self._js_pool.poolSize)
+
+    async def submit(self, func: Any, /, *args: Any, **kwargs: Any) -> Any:
+        """Run ``func(*args, **kwargs)`` on a pool worker; return its result.
+
+        The call triple ships as a cloudpickle payload together with the
+        driver's current package snapshot, so the worker mirrors the
+        driver's installed packages before executing.
+        """
+        payload = _to_js_buffer(cloudpickle.dumps((func, args, kwargs)))
+        packages, wheels = snapshot_packages()
+        options = to_js(
+            {"packages": packages, "wheels": wheels},
+            dict_converter=Object.fromEntries,
+        )
+        try:
+            result = await self._js_pool.runPickled(payload, options)
+        except JsException as exc:
+            _raise_remote(exc)
+        return cloudpickle.loads(_to_py_bytes(result))
+
+    def start_batch(self, calls: Sequence[Call]) -> BatchRun:
+        """Dispatch ``calls`` as ONE JS ``mapPickled`` run and return its
+        handle without awaiting.
+
+        One run means one FIFO queue, input-order results, and one
+        ``cancel`` handle for the whole batch, at one task message per call
+        — the batching substrate of the multiprocessing shim's ``chunksize``.
+        """
+        payloads = [_to_js_buffer(cloudpickle.dumps(call)) for call in calls]
+        packages, wheels = snapshot_packages()
+        options = to_js(
+            {"packages": packages, "wheels": wheels},
+            dict_converter=Object.fromEntries,
+        )
+        return BatchRun(self._js_pool.mapPickled(to_js(payloads), options))
+
+    async def submit_batch(self, calls: Sequence[Call]) -> list[Any]:
+        """Run every ``(func, args, kwargs)`` triple on the pool as one
+        batched run; return the results in input order."""
+        return await self.start_batch(calls).results()
+
+    def terminate(self) -> None:
+        """Terminate idle workers; the pool stays usable (next task re-boots)."""
+        self._js_pool.terminate()
+
+
+_default_pool: WorkerPool | None = None
+
+
+def default_pool() -> WorkerPool:
+    """The process-wide :class:`WorkerPool` over ``js_pyodide_pool.pool``."""
+    global _default_pool
+    if _default_pool is None:
+        _default_pool = WorkerPool()
+    return _default_pool
+
+
+def set_default_pool(pool: WorkerPool | None) -> None:
+    """Install ``pool`` as the process-wide default — the target of
+    module-level :func:`submit` and the dask scheduler. ``loader.create_pool``
+    calls this so notebooks never wire JS themselves; ``None`` resets to lazy
+    resolution via ``js_pyodide_pool``."""
+    global _default_pool
+    _default_pool = pool
+
+
+async def submit(func: Any, /, *args: Any, **kwargs: Any) -> Any:
+    """Module-level :meth:`WorkerPool.submit` on the default pool."""
+    return await default_pool().submit(func, *args, **kwargs)

pyodide_pool-0.1.0/pyodide_pool/_packages.py

@@ -0,0 +1,99 @@
+"""Package-mirroring snapshot helpers (driver side).
+
+Workers are stateless and must satisfy the imports of whatever user code
+they receive. The driver therefore snapshots what its own Pyodide instance
+has installed and rides the snapshot on every submitted task
+(``_bridge.WorkerPool.submit``); the worker's module-scope installed-set
+makes replaying the same snapshot cheap and idempotent, so no separate
+driver/worker sync protocol exists.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import NamedTuple
+
+#: Never mirrored to workers: ``pyodide_pool`` and the Phase 06
+#: ``wasm_multiprocessing`` shim travel by value inside task payloads (they
+#: are never installed on workers, see each package's ``__init__``), the
+#: worker's execPickled shim loads cloudpickle itself, and micropip is the
+#: install mechanism, not a task dependency. The second group is JupyterLite kernel
+#: machinery — never a task dependency, and piplite installs it from its
+#: bundled index while reporting ``source == "pypi"``, so name-based
+#: mirroring would resolve against REAL PyPI where ``pyodide-kernel`` does
+#: not even exist (and ``ipykernel`` pulls ``tornado``, which has no pure
+#: wheel there).
+EXCLUDED_FROM_MIRROR = frozenset(
+    {"pyodide_pool", "wasm_multiprocessing", "cloudpickle", "micropip"}
+    | {"piplite", "pyodide_kernel", "ipykernel", "comm", "widgetsnbextension"}
+)
+
+#: ``pyodide.loadedPackages`` value for packages from the distribution.
+_DEFAULT_CHANNEL = "default channel"
+
+
+class PackageSnapshot(NamedTuple):
+    """What the driver has installed, split by worker install mechanism."""
+
+    packages: list[str]
+    """Pyodide-distribution names, installed via ``pyodide.loadPackage``."""
+
+    wheels: list[str]
+    """micropip targets: PyPI names or wheel URLs."""
+
+
+def _is_url(value: str) -> bool:
+    return "://" in value or value.endswith(".whl")
+
+
+def _canonical(name: str) -> str:
+    """PEP 503 name normalization — micropip lists ``fake_pkg`` as
+    ``fake-pkg``, so exclusion and dedup must compare canonical forms."""
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
+_EXCLUDED_CANONICAL = frozenset(map(_canonical, EXCLUDED_FROM_MIRROR))
+
+
+def snapshot_packages() -> PackageSnapshot:
+    """Snapshot the driver's installed packages for mirroring to workers.
+
+    Combines the JS side's ``pyodide.loadedPackages`` (covers both
+    ``pyodide.loadPackage`` and micropip installs; values are
+    ``"default channel"`` for distribution packages, otherwise the install
+    source) with ``micropip.list()`` when micropip is present. Outside a
+    Pyodide driver both sources are absent and the snapshot is empty.
+    """
+    packages: dict[str, str] = {}  # canonical -> name as first seen
+    wheels: set[str] = set()
+
+    def add(name: str, source: str) -> None:
+        canonical = _canonical(name)
+        if canonical in _EXCLUDED_CANONICAL:
+            return
+        if source in (_DEFAULT_CHANNEL, "pyodide"):
+            packages.setdefault(canonical, name)
+        elif _is_url(source):
+            wheels.add(source)
+        elif canonical not in packages:
+            # Unknown source marker (e.g. "pypi"): let the worker's micropip
+            # resolve the package by name.
+            wheels.add(name)
+
+    try:
+        import pyodide_js
+    except ImportError:
+        pyodide_js = None
+    if pyodide_js is not None:
+        for name, channel in pyodide_js.loadedPackages.to_py().items():
+            add(name, channel)
+
+    try:
+        import micropip
+    except ImportError:
+        micropip = None
+    if micropip is not None:
+        for name, meta in micropip.list().items():
+            add(name, str(getattr(meta, "source", "") or ""))
+
+    return PackageSnapshot(sorted(packages.values()), sorted(wheels))

pyodide_pool-0.1.0/pyodide_pool/loader.py

@@ -0,0 +1,68 @@
+"""Bootstrap the JS worker pool from inside a Pyodide kernel (JupyterLite).
+
+Runs in the MAIN Pyodide instance (the driver), like the rest of the
+package. An embedding web app wires the JS side itself (create a
+``PyodidePool``, ``pyodide.registerJsModule('js_pyodide_pool', { pool })``),
+but notebook code inside JupyterLite's Pyodide kernel has no JS embedder to
+lean on — the kernel itself runs in a Web Worker. :func:`create_pool` pulls
+the JS side in from Python instead: dynamically ``import()`` the
+self-contained browser bundle (``pyodide-pool.browser.js``, built by
+``npm run build`` and copied into the site's ``files/assets/`` by
+``npm run build:lite``), construct the JS ``PyodidePool`` — its workers
+spawn as NESTED Web Workers from a Blob URL inlined in the bundle — and
+install the Python :class:`WorkerPool` wrapper as the default pool, so
+``pyodide_pool.submit``/``compute`` work with no further setup.
+
+Environment requirements: nested-worker support (all Chromium-based
+browsers; Firefox and Safari in current versions) and network access to
+the Pyodide CDN. The pool is postMessage-based, so it does NOT need a
+cross-origin-isolated page — it runs on plain static hosts like GitHub
+Pages (verified by the Pages deployment of demos/jupyterlite). COOP/COEP
+headers (scripts/serve-lite.mjs sends them) only add SharedArrayBuffer
+extras such as the JupyterLite kernel's coincident transport.
+"""
+
+from __future__ import annotations
+
+from . import _bridge
+from ._bridge import WorkerPool
+
+__all__ = ["DEFAULT_JS_URL", "create_pool"]
+
+#: Where the JupyterLite site serves the bundle when hosted at the domain
+#: root (scripts/serve-lite.mjs serves ``_output`` that way). Deployments
+#: under a sub-path must pass ``js_url`` explicitly.
+DEFAULT_JS_URL = "/files/assets/pyodide-pool.browser.js"
+
+
+async def create_pool(pool_size: int = 4, js_url: str | None = None) -> WorkerPool:
+    """Import the JS bundle, start a worker pool, and make it the default.
+
+    Parameters
+    ----------
+    pool_size:
+        Maximum number of concurrent workers, one Pyodide interpreter each.
+    js_url:
+        URL of the self-contained browser bundle, defaulting to
+        :data:`DEFAULT_JS_URL`. Anything the JS ``import()`` accepts in the
+        current context works: an absolute path, a full URL, or a ``file://``
+        URL under Node.
+    """
+    from js import Object
+    from pyodide.code import run_js
+    from pyodide.ffi import to_js
+
+    url = DEFAULT_JS_URL if js_url is None else js_url
+    # Dynamic import() resolves relative and path-absolute URLs against the
+    # REFERRER MODULE (pyodide.asm.js) — in JupyterLite that is the jsDelivr
+    # CDN, so "/files/assets/..." would 404 on cdn.jsdelivr.net. Resolve
+    # against the worker/page location instead, which is the site origin.
+    # Fully-qualified URLs (http://, file://) pass through new URL() as-is,
+    # and Node (no global `location`) keeps the raw url.
+    module = await run_js(
+        f"(u => import(typeof location === 'undefined' ? u : new URL(u, location.href).href))({url!r})"
+    )
+    options = to_js({"poolSize": pool_size}, dict_converter=Object.fromEntries)
+    pool = WorkerPool(module.createPool(options))
+    _bridge.set_default_pool(pool)
+    return pool

pyodide_pool-0.1.0/pyodide_pool/scheduler.py

@@ -0,0 +1,209 @@
+"""Async dask scheduler executing task graphs on the Pyodide worker pool.
+
+This module runs in the MAIN Pyodide instance (the driver). It implements
+the design-doc algorithm (docs/architecture/dask-scheduler-design.md): a
+Kahn-style in-degree executor over the raw graph dict that dispatches every
+ready task to the worker pool at once and, as each completes, decrements its
+dependents and dispatches the newly ready ones. Concurrency is bounded by
+the JS pool's size, never by the scheduler — over-dispatching just fills the
+pool queue, which is exactly what utilization wants.
+
+dask is imported lazily inside :func:`get`/:func:`compute` so the package
+keeps no import-time dask dependency (the Phase 06 multiprocessing shim
+shares the substrate). Both graph dialects execute:
+
+- **Legacy tuple tasks** (``(func, arg, ...)``) ship as a self-contained
+  recursive evaluator plus resolved dependency values, so workers never need
+  dask installed.
+- **Modern task-spec nodes** (``dask._task_spec``, dask >= 2024.12 — what
+  ``dask.delayed``/``dask.bag`` emit today) are callable with a mapping of
+  dependency values and pickle by reference to dask; they execute on workers
+  because package mirroring ships the driver's dask there.
+
+Literal graph values, key aliases, and ``Alias``/``DataNode`` nodes resolve
+locally without a worker round-trip.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from . import _bridge
+
+__all__ = ["compute", "get"]
+
+
+def _evaluate_node(node: Any, values: dict[Any, Any]) -> Any:
+    """Recursively evaluate a legacy-dialect graph node against resolved
+    dependency ``values`` (the dask graph spec: tasks are tuples with a
+    callable head, lists are traversed, keys substitute, the rest are
+    literals).
+
+    Runs on workers inside cloudpickle payloads — the package is registered
+    pickle-by-value, so this function travels with each task. It must stay
+    self-contained: no dask, no imports, no references to other globals.
+    """
+    if isinstance(node, list):
+        return [_evaluate_node(item, values) for item in node]
+    if type(node) is tuple and node and callable(node[0]):
+        return node[0](*(_evaluate_node(arg, values) for arg in node[1:]))
+    try:
+        if node in values:
+            return values[node]
+    except TypeError:  # unhashable literal (dict, set, ...)
+        pass
+    return node
+
+
+def _nested_get(keys: Any, cache: dict[Any, Any]) -> Any:
+    """Assemble results following the nesting of ``keys``, like dask's own
+    ``get``: nested key lists become tuples of results."""
+    if isinstance(keys, list):
+        return tuple(_nested_get(k, cache) for k in keys)
+    return cache[keys]
+
+
+async def get(dsk: Any, keys: Any, pool: Any = None, **kwargs: Any) -> Any:
+    """Execute a dask graph on the worker pool; return the values of ``keys``.
+
+    Parameters
+    ----------
+    dsk:
+        A dask graph — a mapping of key to node, or anything with a
+        ``__dask_graph__()`` (HighLevelGraph, expression) to materialize one.
+    keys:
+        A key or arbitrarily nested lists of keys, like dask's own ``get``;
+        the result mirrors the nesting (lists come back as tuples).
+    pool:
+        A :class:`~pyodide_pool._bridge.WorkerPool` to dispatch through;
+        defaults to the pool registered as ``js_pyodide_pool.pool``.
+    kwargs:
+        Accepted for scheduler-call compatibility (``num_workers`` etc.) and
+        ignored: the JS pool bounds concurrency, the scheduler dispatches
+        everything that is ready.
+    """
+    from dask.core import flatten, get_dependencies, istask, reverse_dict, toposort
+
+    try:
+        from dask._task_spec import Alias, DataNode, GraphNode
+    except ImportError:  # dask predates the task-spec dialect
+        Alias = DataNode = GraphNode = ()  # type: ignore[assignment]
+
+    if not isinstance(dsk, dict) and hasattr(dsk, "__dask_graph__"):
+        dsk = dsk.__dask_graph__()
+    dsk = dict(dsk)
+
+    requested = list(flatten(keys)) if isinstance(keys, list) else [keys]
+    missing = [k for k in requested if k not in dsk]
+    if missing:
+        raise KeyError(f"requested keys not in the graph: {missing!r}")
+
+    dependencies = {k: get_dependencies(dsk, k) for k in dsk}
+    toposort(dsk, dependencies=dependencies)  # raises on cycles
+    dependents = reverse_dict(dependencies)
+    indegree = {k: len(deps) for k, deps in dependencies.items()}
+
+    submit = _bridge.submit if pool is None else pool.submit
+    cache: dict[Any, Any] = {}
+    pending: dict[Any, asyncio.Future[Any]] = {}
+    ready = [k for k in dsk if indegree[k] == 0]
+
+    def finish(key: Any) -> None:
+        for dependent in dependents.get(key, ()):
+            indegree[dependent] -= 1
+            if indegree[dependent] == 0:
+                ready.append(dependent)
+
+    def dispatch_ready() -> None:
+        # Drain the ready list: local nodes resolve immediately (readying
+        # further keys, hence the loop), remote tasks become pending futures.
+        while ready:
+            key = ready.pop()
+            node = dsk[key]
+            values = {dep: cache[dep] for dep in dependencies[key]}
+            if isinstance(node, GraphNode):
+                if isinstance(node, Alias) or isinstance(node, DataNode):
+                    cache[key] = node(values)  # cheap: no worker round-trip
+                    finish(key)
+                else:
+                    pending[key] = asyncio.ensure_future(submit(node, values))
+            elif istask(node):
+                pending[key] = asyncio.ensure_future(
+                    submit(_evaluate_node, node, values)
+                )
+            else:
+                # Literal, key alias, or container of keys/literals.
+                cache[key] = _evaluate_node(node, values)
+                finish(key)
+
+    try:
+        dispatch_ready()
+        while pending:
+            done, _ = await asyncio.wait(
+                pending.values(), return_when=asyncio.FIRST_COMPLETED
+            )
+            for key in [k for k, fut in pending.items() if fut in done]:
+                cache[key] = pending.pop(key).result()
+                finish(key)
+            dispatch_ready()
+    except BaseException:
+        # Fail fast: drop the outstanding work and let the first error
+        # (already carrying its remote traceback, see _bridge) propagate.
+        if pending:
+            for fut in pending.values():
+                fut.cancel()
+            await asyncio.gather(*pending.values(), return_exceptions=True)
+        raise
+
+    return _nested_get(keys, cache)
+
+
+async def compute(
+    *args: Any,
+    traverse: bool = True,
+    optimize_graph: bool = True,
+    pool: Any = None,
+    **kwargs: Any,
+) -> Any:
+    """Compute dask collections on the worker pool.
+
+    The async counterpart of :func:`dask.compute`: accepts ``dask.delayed``
+    objects, bags, arrays, plain values, and (with ``traverse=True``)
+    builtin containers of any of these. A single argument returns its bare
+    result — ``await compute(delayed_obj)`` matches ``delayed_obj.compute()``
+    — while multiple arguments return a tuple like ``dask.compute``.
+
+    On expression-based dask (2025+) this follows ``dask.base.compute``:
+    ``collections_to_expr`` + ``FinalizeCompute`` compile each collection's
+    finalization into the graph itself. On older dask it falls back to the
+    classic ``collections_to_dsk`` + ``__dask_postcompute__`` protocol.
+    """
+    from dask.base import unpack_collections
+    from dask.core import flatten
+
+    collections, repack = unpack_collections(*args, traverse=traverse)
+    if not collections:
+        return args[0] if len(args) == 1 else args
+
+    try:
+        from dask._expr import FinalizeCompute
+        from dask.base import collections_to_expr
+    except ImportError:  # pre-expression dask: classic collection protocol
+        from dask.base import collections_to_dsk
+
+        dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
+        nested_keys = [c.__dask_keys__() for c in collections]
+        postcomputes = [c.__dask_postcompute__() for c in collections]
+        results = await get(dsk, nested_keys, pool=pool)
+        out = repack(
+            [f(r, *extra) for r, (f, extra) in zip(results, postcomputes)]
+        )
+    else:
+        expr = FinalizeCompute(collections_to_expr(collections, optimize_graph))
+        expr = expr.optimize()
+        keys = list(flatten(expr.__dask_keys__()))
+        results = await get(expr.__dask_graph__(), keys, pool=pool)
+        out = repack(results)
+
+    return out[0] if len(args) == 1 else out

pyodide_pool-0.1.0/pyproject.toml

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "pyodide-pool"
+description = "Run Python callables and dask graphs on a pool of Pyodide Web Workers (driver-side API)"
+license = "MIT"
+requires-python = ">=3.11"
+# cloudpickle is the wire format of every submitted task (_bridge); dask is
+# only needed for compute/get and stays an extra. Both resolve in Pyodide:
+# cloudpickle ships in the distribution, dask installs from PyPI via
+# micropip/piplite.
+dependencies = ["cloudpickle"]
+dynamic = ["version"]
+
+[project.optional-dependencies]
+dask = ["dask"]
+
+[tool.hatch.version]
+path = "pyodide_pool/__init__.py"
+
+[tool.hatch.build.targets.wheel]
+packages = ["pyodide_pool"]

pyodide_pool-0.1.0/wasm_multiprocessing/pyproject.toml

@@ -0,0 +1,21 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "wasm-multiprocessing"
+description = "multiprocessing.Pool-compatible API over a pool of Pyodide Web Workers (driver-side shim)"
+license = "MIT"
+requires-python = ">=3.11"
+# pyodide-pool is the single choke point for cloudpickle + JS bridging
+# (_bridge.submit / submit_batch); the shim adds no second copy — it only
+# needs cloudpickle itself for the pickle-by-value registration. pyodide-pool
+# is not on PyPI: install both wheels together (JupyterLite ships both).
+dependencies = ["pyodide-pool", "cloudpickle"]
+dynamic = ["version"]
+
+[tool.hatch.version]
+path = "wasm_multiprocessing/__init__.py"
+
+[tool.hatch.build.targets.wheel]
+packages = ["wasm_multiprocessing"]

pyodide_pool-0.1.0/wasm_multiprocessing/wasm_multiprocessing/__init__.py

@@ -0,0 +1,94 @@
+"""``multiprocessing.Pool`` for Pyodide: same API, Web Workers underneath.
+
+Change one import line and existing ``multiprocessing.Pool`` code runs on
+the Pyodide worker pool::
+
+    # was: from multiprocessing import Pool
+    from wasm_multiprocessing import Pool
+
+    with Pool(4) as pool:
+        results = await pool.amap(count_primes, ranges)  # portable everywhere
+        results = pool.map(count_primes, ranges)         # blocks only under JSPI
+
+Like ``pyodide_pool``, this package runs ONLY in the driver: the chunk
+runners it ships inside task payloads travel by value (pickle-by-value
+registration below), so workers never install it.
+
+Sync methods (``map``/``starmap``/``apply``/``join``, ``AsyncResult.get``/
+``wait``) truly block only where JSPI stack switching is available
+(``pyodide.ffi.can_run_sync()``, detected per call); everywhere else they
+raise a ``RuntimeError`` naming the exact async replacement. Anything
+fork-shaped (``Process``, ``Queue``, ``Manager``, shared memory, locks)
+raises ``NotImplementedError`` — WebAssembly workers share nothing.
+
+Design + full API mapping: docs/architecture/multiprocessing-shim-design.md.
+"""
+
+from __future__ import annotations
+
+import sys
+from typing import Any
+
+import cloudpickle
+
+from .pool import AsyncResult, Pool, TimeoutError, cpu_count
+
+__version__ = "0.1.0"
+
+__all__ = ["AsyncResult", "Pool", "TimeoutError", "cpu_count"]
+
+# The chunk runners (pool._run_chunk / _run_star_chunk) ship inside
+# cloudpickle payloads; registering the package makes them travel BY VALUE
+# so workers never need wasm_multiprocessing installed (registering the
+# parent package covers submodules — cloudpickle walks parent packages).
+cloudpickle.register_pickle_by_value(sys.modules[__name__])
+
+#: stdlib ``multiprocessing`` names that cannot exist on WebAssembly — no
+#: processes, no fork, no shared memory, no channels between workers.
+#: Attribute access raises immediately instead of failing obscurely
+#: mid-task; see the design doc's "out of scope" section.
+_UNSUPPORTED = frozenset(
+    {
+        "Process",
+        "Queue",
+        "SimpleQueue",
+        "JoinableQueue",
+        "Pipe",
+        "Manager",
+        "Value",
+        "Array",
+        "RawValue",
+        "RawArray",
+        "Lock",
+        "RLock",
+        "Semaphore",
+        "BoundedSemaphore",
+        "Event",
+        "Condition",
+        "Barrier",
+        "shared_memory",
+        "current_process",
+        "parent_process",
+        "active_children",
+        "freeze_support",
+        "get_context",
+        "get_start_method",
+        "set_start_method",
+        "get_all_start_methods",
+        "set_executable",
+        "set_forkserver_preload",
+        "log_to_stderr",
+        "get_logger",
+    }
+)
+
+
+def __getattr__(name: str) -> Any:
+    if name in _UNSUPPORTED:
+        raise NotImplementedError(
+            f"wasm_multiprocessing.{name} is intentionally unsupported: "
+            "WebAssembly workers share nothing and cannot fork — tasks "
+            "communicate only through arguments and results. See "
+            "docs/architecture/multiprocessing-shim-design.md (out of scope)."
+        )
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

pyodide_pool-0.1.0/wasm_multiprocessing/wasm_multiprocessing/pool.py

@@ -0,0 +1,483 @@
+"""The ``multiprocessing.Pool`` shim over ``pyodide_pool``.
+
+Every remote execution goes through ``pyodide_pool`` — batched chunks as
+one JS ``mapPickled`` run (``WorkerPool.start_batch``), streamed chunks as
+single ``submit`` calls — so cloudpickle/JS bridging and remote-exception
+re-raising (original type with ``RemoteTraceback`` chained) live in exactly
+one place.
+
+The async-native core is the portable surface: ``amap`` / ``astarmap`` /
+``aapply`` / ``imap`` / ``imap_unordered`` / ``ajoin`` plus
+``AsyncResult.aget`` / ``await_ready``. The stdlib-shaped sync methods are
+thin wrappers over it via ``_block_on``: JSPI ``run_sync`` where the
+runtime supports stack switching, a ``RuntimeError`` naming the exact
+async replacement everywhere else.
+
+Design + API mapping: docs/architecture/multiprocessing-shim-design.md.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Mapping, Sequence
+from typing import Any
+
+import pyodide_pool
+from pyodide_pool._bridge import BatchRun, Call
+
+try:  # the class users already catch; importable on Emscripten
+    from multiprocessing import TimeoutError
+except Exception:  # pragma: no cover — stdlib multiprocessing unavailable
+
+    class TimeoutError(Exception):  # type: ignore[no-redef]
+        """Raised when a blocking ``AsyncResult`` call times out."""
+
+
+__all__ = ["AsyncResult", "Pool", "TimeoutError", "cpu_count"]
+
+# Pool lifecycle states (stdlib names).
+_RUN, _CLOSE, _TERMINATE = "RUN", "CLOSE", "TERMINATE"
+
+_JSPI_HINT = (
+    "Synchronous blocking needs WebAssembly JSPI stack switching "
+    "(pyodide.ffi.run_sync): on by default in Chrome 137+, in Node 24 via "
+    "--experimental-wasm-jspi, behind a flag in Firefox, not yet in Safari."
+)
+
+
+def _block_on(awaitable: Coroutine[Any, Any, Any], replacement: str) -> Any:
+    """Run ``awaitable`` to completion synchronously where the runtime can,
+    else raise the guidance error naming the exact async replacement.
+
+    Capability is detected PER CALL, not at import: ``can_run_sync()``
+    depends on how Python was entered (async entry via ``runPythonAsync`` /
+    ``callPromising`` allows stack switching; plain ``runPython`` does not).
+    """
+    try:
+        from pyodide.ffi import can_run_sync, run_sync
+
+        capable = can_run_sync()
+    except ImportError:  # non-Pyodide Python (CI helpers, docs builds)
+        capable = False
+    if capable:
+        return run_sync(awaitable)
+    awaitable.close()  # don't leak a never-awaited coroutine
+    raise RuntimeError(
+        f"this runtime cannot block on pool results — use the async form: "
+        f"{replacement}\n{_JSPI_HINT}"
+    )
+
+
+def cpu_count() -> int:
+    """``navigator.hardwareConcurrency`` (browsers, workers, Node >= 21);
+    1 where unavailable."""
+    try:
+        import js
+
+        count = int(js.navigator.hardwareConcurrency)
+        if count >= 1:
+            return count
+    except (ImportError, AttributeError, TypeError, ValueError):
+        pass
+    return 1
+
+
+def _run_chunk(func: Callable[..., Any], items: list[Any]) -> list[Any]:
+    """Executes ON WORKERS (ships by value): one ``map`` chunk."""
+    return [func(item) for item in items]
+
+
+def _run_star_chunk(func: Callable[..., Any], items: list[Any]) -> list[Any]:
+    """Executes ON WORKERS (ships by value): one ``starmap`` chunk."""
+    return [func(*item) for item in items]
+
+
+def _default_chunksize(n_items: int, processes: int) -> int:
+    """CPython's own ``Pool._map_async`` heuristic: ~4 chunks per (soft-cap)
+    worker, so per-message overhead amortizes over coarse batches."""
+    chunksize, extra = divmod(n_items, processes * 4)
+    if extra:
+        chunksize += 1
+    return max(chunksize, 1)
+
+
+def _split(items: list[Any], chunksize: int) -> list[list[Any]]:
+    return [items[i : i + chunksize] for i in range(0, len(items), chunksize)]
+
+
+class AsyncResult:
+    """Result handle of the ``*_async`` methods, backed by an asyncio task.
+
+    ``ready``/``successful`` are non-blocking and stdlib-shaped. The
+    portable completion forms are awaitables — ``await result`` (sugar for
+    ``aget()``) and ``await_ready()``; the blocking ``get``/``wait`` work
+    only under JSPI (see ``_block_on``).
+    """
+
+    def __init__(
+        self,
+        task: asyncio.Task[Any],
+        callback: Callable[[Any], object] | None = None,
+        error_callback: Callable[[BaseException], object] | None = None,
+    ) -> None:
+        self._task = task
+        self._callback = callback
+        self._error_callback = error_callback
+        task.add_done_callback(self._fire_callbacks)
+
+    def _fire_callbacks(self, task: asyncio.Task[Any]) -> None:
+        if task.cancelled():
+            return
+        # .exception() also marks the failure as retrieved, so an abandoned
+        # failed result never trips asyncio's never-retrieved warning.
+        exc = task.exception()
+        if exc is None:
+            if self._callback is not None:
+                self._callback(task.result())
+        elif self._error_callback is not None:
+            self._error_callback(exc)
+
+    def ready(self) -> bool:
+        return self._task.done()
+
+    def successful(self) -> bool:
+        if not self._task.done():
+            raise ValueError(f"{self!r} not ready")
+        return not self._task.cancelled() and self._task.exception() is None
+
+    async def aget(self, timeout: float | None = None) -> Any:
+        """Await the result; a failed task re-raises its original exception
+        (worker traceback chained as ``RemoteTraceback``). On ``timeout``
+        raises :class:`TimeoutError` and leaves the computation running —
+        a later ``aget()`` can still collect it (stdlib ``get`` semantics,
+        hence the ``shield``)."""
+        if timeout is None:
+            return await asyncio.shield(self._task)
+        try:
+            return await asyncio.wait_for(asyncio.shield(self._task), timeout)
+        except asyncio.TimeoutError:
+            raise TimeoutError from None
+
+    async def await_ready(self, timeout: float | None = None) -> None:
+        """Awaitable form of ``wait``: returns once ready or after
+        ``timeout`` seconds, never raising on timeout (stdlib ``wait``)."""
+        await asyncio.wait({self._task}, timeout=timeout)
+
+    def get(self, timeout: float | None = None) -> Any:
+        return _block_on(self.aget(timeout), "await result.aget(timeout)")
+
+    def wait(self, timeout: float | None = None) -> None:
+        return _block_on(
+            self.await_ready(timeout), "await result.await_ready(timeout)"
+        )
+
+    def __await__(self) -> Any:
+        return self.aget().__await__()
+
+
+class Pool:
+    """``multiprocessing.Pool`` over the shared Pyodide worker pool.
+
+    The JS pool's ``poolSize`` is the HARD concurrency cap; ``processes``
+    is a driver-side soft cap (``asyncio.Semaphore``) on in-flight chunks,
+    which is what makes ``Pool(1)``/``Pool(2)``/``Pool(4)`` behave
+    differently on one shared 4-worker JS pool.
+    """
+
+    def __init__(
+        self,
+        processes: int | None = None,
+        initializer: Callable[..., object] | None = None,
+        initargs: Sequence[Any] = (),
+        maxtasksperchild: int | None = None,
+        context: Any = None,
+        *,
+        pool: pyodide_pool.WorkerPool | None = None,
+    ) -> None:
+        if initializer is not None:
+            raise NotImplementedError(
+                "Pool(initializer=...) is not supported yet (planned as a "
+                "cloudpickled per-worker prelude); ship state inside func "
+                "instead. See docs/architecture/multiprocessing-shim-design.md."
+            )
+        if maxtasksperchild is not None or context is not None:
+            raise NotImplementedError(
+                "Pool(maxtasksperchild=/context=) has no meaning here — "
+                "worker recycling is the JS pool's policy. See "
+                "docs/architecture/multiprocessing-shim-design.md."
+            )
+        if processes is None:
+            processes = cpu_count()
+        if not isinstance(processes, int) or processes < 1:
+            raise ValueError("Number of processes must be at least 1")
+        self._processes = processes
+        self._pool = pool if pool is not None else pyodide_pool.default_pool()
+        self._semaphore = asyncio.Semaphore(processes)
+        self._state = _RUN
+        self._active_batches: set[BatchRun] = set()
+        self._pending: set[asyncio.Task[Any]] = set()
+
+    def __repr__(self) -> str:
+        return f"<wasm_multiprocessing.Pool processes={self._processes} state={self._state}>"
+
+    # -- internals ----------------------------------------------------------
+
+    def _check_running(self) -> None:
+        if self._state != _RUN:
+            raise ValueError("Pool not running")
+
+    def _spawn(self, coro: Coroutine[Any, Any, Any]) -> asyncio.Task[Any]:
+        """Create a tracked asyncio task (``join`` awaits the tracked set;
+        ``terminate`` cancels it). Works from sync driver code too:
+        Pyodide's WebLoop accepts task creation outside a coroutine."""
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = asyncio.get_event_loop()
+        task = loop.create_task(coro)
+        self._pending.add(task)
+        task.add_done_callback(self._forget)
+        return task
+
+    def _forget(self, task: asyncio.Task[Any]) -> None:
+        self._pending.discard(task)
+        if not task.cancelled():
+            task.exception()  # abandoned failures stay quiet (see AsyncResult)
+
+    def _map_calls(
+        self,
+        runner: Callable[..., Any],
+        func: Callable[..., Any],
+        iterable: Iterable[Any],
+        chunksize: int | None,
+    ) -> list[Call]:
+        """Chunk ``iterable`` into ``(runner, (func, chunk), {})`` call
+        triples — one chunk = one task message. Consumed eagerly (v1)."""
+        items = list(iterable)
+        if chunksize is None:
+            chunksize = _default_chunksize(len(items), self._processes) if items else 1
+        elif chunksize < 1:
+            raise ValueError("Chunksize must be 1+")
+        return [(runner, (func, chunk), {}) for chunk in _split(items, chunksize)]
+
+    async def _submit_call(self, call: Call) -> Any:
+        """One chunk as a single ``submit``, gated by the soft cap."""
+        func, args, kwargs = call
+        async with self._semaphore:
+            return await self._pool.submit(func, *args, **kwargs)
+
+    async def _run_chunks(self, calls: list[Call]) -> list[Any]:
+        """Chunk results in input order.
+
+        Fast path: when the soft cap cannot bind tighter than the JS pool's
+        own worker cap, the whole call is ONE ``mapPickled`` run (single
+        FIFO queue, one cancel handle). Otherwise chunks go through
+        semaphore-gated single submits so at most ``processes`` are ever in
+        flight.
+        """
+        if self._processes >= self._pool.pool_size:
+            batch = self._pool.start_batch(calls)
+            self._active_batches.add(batch)
+            try:
+                return await batch.results()
+            finally:
+                self._active_batches.discard(batch)
+        return list(await asyncio.gather(*(self._submit_call(call) for call in calls)))
+
+    async def _gather_map(self, calls: list[Call]) -> list[Any]:
+        if not calls:
+            return []
+        chunk_results = await self._run_chunks(calls)
+        return [result for chunk in chunk_results for result in chunk]
+
+    def _stream_chunks(
+        self,
+        runner: Callable[..., Any],
+        func: Callable[..., Any],
+        iterable: Iterable[Any],
+        chunksize: int,
+    ) -> list[asyncio.Task[Any]]:
+        """Per-chunk submit futures for ``imap*`` — a ``mapPickled`` run
+        settles only as a whole, so streaming needs one future per chunk."""
+        if chunksize < 1:
+            raise ValueError("Chunksize must be 1+")
+        calls = self._map_calls(runner, func, iterable, chunksize)
+        return [self._spawn(self._submit_call(call)) for call in calls]
+
+    # -- async-native core --------------------------------------------------
+
+    async def amap(
+        self,
+        func: Callable[[Any], Any],
+        iterable: Iterable[Any],
+        chunksize: int | None = None,
+    ) -> list[Any]:
+        """Async ``map``: results in input order."""
+        self._check_running()
+        return await self._spawn(self._gather_map(self._map_calls(_run_chunk, func, iterable, chunksize)))
+
+    async def astarmap(
+        self,
+        func: Callable[..., Any],
+        iterable: Iterable[Sequence[Any]],
+        chunksize: int | None = None,
+    ) -> list[Any]:
+        """Async ``starmap``: each item is an argument tuple."""
+        self._check_running()
+        return await self._spawn(self._gather_map(self._map_calls(_run_star_chunk, func, iterable, chunksize)))
+
+    async def aapply(
+        self,
+        func: Callable[..., Any],
+        args: Sequence[Any] = (),
+        kwds: Mapping[str, Any] | None = None,
+    ) -> Any:
+        """Async ``apply``: one unbatched call."""
+        self._check_running()
+        return await self._spawn(self._pool.submit(func, *args, **dict(kwds or {})))
+
+    def imap(
+        self,
+        func: Callable[[Any], Any],
+        iterable: Iterable[Any],
+        chunksize: int = 1,
+    ) -> AsyncIterator[Any]:
+        """Async generator over results IN INPUT ORDER, yielding as chunks
+        complete. Chunks dispatch eagerly at call time (v1), bounded in
+        flight by the soft cap."""
+        self._check_running()
+        futures = self._stream_chunks(_run_chunk, func, iterable, chunksize)
+
+        async def iterate() -> AsyncIterator[Any]:
+            for future in futures:
+                for result in await future:
+                    yield result
+
+        return iterate()
+
+    def imap_unordered(
+        self,
+        func: Callable[[Any], Any],
+        iterable: Iterable[Any],
+        chunksize: int = 1,
+    ) -> AsyncIterator[Any]:
+        """Like ``imap`` but yields whole chunks in COMPLETION order."""
+        self._check_running()
+        futures = self._stream_chunks(_run_chunk, func, iterable, chunksize)
+
+        async def iterate() -> AsyncIterator[Any]:
+            for next_done in asyncio.as_completed(futures):
+                for result in await next_done:
+                    yield result
+
+        return iterate()
+
+    async def ajoin(self) -> None:
+        """Await every outstanding result. Stdlib rule kept: ``ValueError``
+        unless the pool was closed/terminated first."""
+        if self._state == _RUN:
+            raise ValueError("Pool is still running")
+        while self._pending:
+            await asyncio.wait(set(self._pending))
+
+    # -- stdlib-shaped API --------------------------------------------------
+
+    def map(
+        self,
+        func: Callable[[Any], Any],
+        iterable: Iterable[Any],
+        chunksize: int | None = None,
+    ) -> list[Any]:
+        self._check_running()
+        return _block_on(self.amap(func, iterable, chunksize), "await pool.amap(func, iterable)")
+
+    def starmap(
+        self,
+        func: Callable[..., Any],
+        iterable: Iterable[Sequence[Any]],
+        chunksize: int | None = None,
+    ) -> list[Any]:
+        self._check_running()
+        return _block_on(self.astarmap(func, iterable, chunksize), "await pool.astarmap(func, iterable)")
+
+    def apply(
+        self,
+        func: Callable[..., Any],
+        args: Sequence[Any] = (),
+        kwds: Mapping[str, Any] | None = None,
+    ) -> Any:
+        self._check_running()
+        return _block_on(self.aapply(func, args, kwds), "await pool.aapply(func, args, kwds)")
+
+    def map_async(
+        self,
+        func: Callable[[Any], Any],
+        iterable: Iterable[Any],
+        chunksize: int | None = None,
+        callback: Callable[[Any], object] | None = None,
+        error_callback: Callable[[BaseException], object] | None = None,
+    ) -> AsyncResult:
+        self._check_running()
+        calls = self._map_calls(_run_chunk, func, iterable, chunksize)
+        return AsyncResult(self._spawn(self._gather_map(calls)), callback, error_callback)
+
+    def starmap_async(
+        self,
+        func: Callable[..., Any],
+        iterable: Iterable[Sequence[Any]],
+        chunksize: int | None = None,
+        callback: Callable[[Any], object] | None = None,
+        error_callback: Callable[[BaseException], object] | None = None,
+    ) -> AsyncResult:
+        self._check_running()
+        calls = self._map_calls(_run_star_chunk, func, iterable, chunksize)
+        return AsyncResult(self._spawn(self._gather_map(calls)), callback, error_callback)
+
+    def apply_async(
+        self,
+        func: Callable[..., Any],
+        args: Sequence[Any] = (),
+        kwds: Mapping[str, Any] | None = None,
+        callback: Callable[[Any], object] | None = None,
+        error_callback: Callable[[BaseException], object] | None = None,
+    ) -> AsyncResult:
+        self._check_running()
+        task = self._spawn(self._pool.submit(func, *args, **dict(kwds or {})))
+        return AsyncResult(task, callback, error_callback)
+
+    def close(self) -> None:
+        """Stop accepting submissions; outstanding work keeps running —
+        drain it with ``await pool.ajoin()`` (or ``join()`` under JSPI)."""
+        if self._state == _RUN:
+            self._state = _CLOSE
+
+    def join(self) -> None:
+        if self._state == _RUN:
+            raise ValueError("Pool is still running")
+        return _block_on(self.ajoin(), "await pool.ajoin()")
+
+    def terminate(self) -> None:
+        """Stop without draining, approximately (documented): cancel batched
+        chunks not yet started plus every tracked pending future, then
+        terminate the JS pool's idle workers. In-flight worker calls run to
+        completion — no preemption without an interrupt buffer."""
+        self._state = _TERMINATE
+        for batch in list(self._active_batches):
+            batch.cancel()
+        for task in list(self._pending):
+            task.cancel()
+        self._pool.terminate()
+
+    def __enter__(self) -> Pool:
+        self._check_running()
+        return self
+
+    def __exit__(self, *exc_info: object) -> None:
+        self.terminate()
+
+    async def __aenter__(self) -> Pool:
+        self._check_running()
+        return self
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        self.terminate()