pyodide-pool 0.1.0__py3-none-any.whl

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/_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/_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/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/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.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,8 @@
+pyodide_pool/__init__.py,sha256=mIl2RE3dlASJYBJwVsr7UGZQKRg-HUbUXC_m0xS4UyQ,2616
+pyodide_pool/_bridge.py,sha256=i7BH39YXr4yY3bToEdMFInH7xoDfAknLKgtFe1OrG3U,8913
+pyodide_pool/_packages.py,sha256=zf0fOQaatYv7-mGqny9BYppPyB_s3m-WytoymrokM1w,3770
+pyodide_pool/loader.py,sha256=zCZhBQkZL5FD_L0308xsDSfsiba3j5epbVVX8kWHO3Q,3264
+pyodide_pool/scheduler.py,sha256=ItS14NLiV5nQeLTO5SnQrzWiC2rhNoRzQehmvP8CwHo,8598
+pyodide_pool-0.1.0.dist-info/METADATA,sha256=Esy-HZEL5-d0MBsQJQKAhoKMWSpKK9_GZxzVKmi22m8,286
+pyodide_pool-0.1.0.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+pyodide_pool-0.1.0.dist-info/RECORD,,

pyodide_pool-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.31.0
+Root-Is-Purelib: true
+Tag: py3-none-any