coco-framework 0.1.0__py3-none-any.whl

coco_graph/pregel/_draw.py

@@ -0,0 +1,294 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Mapping, Sequence
+from typing import Any, NamedTuple, cast
+
+from coco_core.runnables.config import RunnableConfig
+from coco_core.runnables.graph import Graph, Node
+from coco_graph.checkpoint.base import BaseCheckpointSaver
+
+from coco_graph._internal._constants import CONF, CONFIG_KEY_SEND, INPUT
+from coco_graph.channels.base import BaseChannel
+from coco_graph.channels.last_value import LastValueAfterFinish
+from coco_graph.constants import END, START
+from coco_graph.managed.base import ManagedValueSpec
+from coco_graph.pregel._algo import (
+    PregelTaskWrites,
+    apply_writes,
+    increment,
+    prepare_next_tasks,
+)
+from coco_graph.pregel._checkpoint import channels_from_checkpoint, empty_checkpoint
+from coco_graph.pregel._io import map_input
+from coco_graph.pregel._read import PregelNode
+from coco_graph.pregel._write import ChannelWrite
+from coco_graph.types import All, Checkpointer
+
+
+class Edge(NamedTuple):
+    source: str
+    target: str
+    conditional: bool
+    data: str | None
+
+
+class TriggerEdge(NamedTuple):
+    source: str
+    conditional: bool
+    data: str | None
+
+
+def draw_graph(
+    config: RunnableConfig,
+    *,
+    nodes: dict[str, PregelNode],
+    specs: dict[str, BaseChannel | ManagedValueSpec],
+    input_channels: str | Sequence[str],
+    interrupt_after_nodes: All | Sequence[str],
+    interrupt_before_nodes: All | Sequence[str],
+    trigger_to_nodes: Mapping[str, Sequence[str]],
+    checkpointer: Checkpointer,
+    subgraphs: dict[str, Graph],
+    limit: int = 250,
+) -> Graph:
+    """Get the graph for this Pregel instance.
+
+    Args:
+        config: The configuration to use for the graph.
+        subgraphs: The subgraphs to include in the graph.
+        checkpointer: The checkpointer to use for the graph.
+
+    Returns:
+        The graph for this Pregel instance.
+    """
+    # (src, dest, is_conditional, label)
+    edges: set[Edge] = set()
+
+    step = -1
+    checkpoint = empty_checkpoint()
+    get_next_version = (
+        checkpointer.get_next_version
+        if isinstance(checkpointer, BaseCheckpointSaver)
+        else increment
+    )
+    channels, managed = channels_from_checkpoint(
+        specs,
+        checkpoint,
+    )
+    static_seen: set[Any] = set()
+    sources: dict[str, set[TriggerEdge]] = {}
+    step_sources: dict[str, set[TriggerEdge]] = {}
+    static_declared_writes: dict[str, set[TriggerEdge]] = defaultdict(set)
+    # remove node mappers
+    nodes = {
+        k: v.copy(update={"mapper": None}) if v.mapper is not None else v
+        for k, v in nodes.items()
+    }
+    # apply input writes
+    input_writes = list(map_input(input_channels, {}))
+    updated_channels = apply_writes(
+        checkpoint,
+        channels,
+        [
+            PregelTaskWrites((), INPUT, input_writes, []),
+        ],
+        get_next_version,
+        trigger_to_nodes,
+    )
+    # prepare first tasks
+    tasks = prepare_next_tasks(
+        checkpoint,
+        [],
+        nodes,
+        channels,
+        managed,
+        config,
+        step,
+        -1,
+        for_execution=True,
+        store=None,
+        checkpointer=None,
+        manager=None,
+        trigger_to_nodes=trigger_to_nodes,
+        updated_channels=updated_channels,
+    )
+    start_tasks = tasks
+    # run the pregel loop
+    for step in range(step, limit):
+        if not tasks:
+            break
+        conditionals: dict[tuple[str, str, Any], str | None] = {}
+        # run task writers
+        for task in tasks.values():
+            for w in task.writers:
+                # apply regular writes
+                if isinstance(w, ChannelWrite):
+                    empty_input = (
+                        cast(BaseChannel, specs["__root__"]).ValueType()
+                        if "__root__" in specs
+                        else None
+                    )
+                    w.invoke(empty_input, task.config)
+                # apply conditional writes declared for static analysis, only once
+                if w not in static_seen:
+                    static_seen.add(w)
+                    # apply static writes
+                    if writes := ChannelWrite.get_static_writes(w):
+                        # END writes are not written, but become edges directly
+                        for t in writes:
+                            if t[0] == END:
+                                edges.add(Edge(task.name, t[0], True, t[2]))
+                        writes = [t for t in writes if t[0] != END]
+                        conditionals.update(
+                            {(task.name, t[0], t[1] or None): t[2] for t in writes}
+                        )
+                        # record static writes for edge creation
+                        for t in writes:
+                            static_declared_writes[task.name].add(
+                                TriggerEdge(t[0], True, t[2])
+                            )
+                        task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes])
+        # collect sources
+        step_sources = {}
+        for task in tasks.values():
+            task_edges = {
+                TriggerEdge(
+                    w[0],
+                    (task.name, w[0], w[1] or None) in conditionals,
+                    conditionals.get((task.name, w[0], w[1] or None)),
+                )
+                for w in task.writes
+            }
+            task_edges |= static_declared_writes.get(task.name, set())
+            step_sources[task.name] = task_edges
+        sources.update(step_sources)
+        # invert triggers
+        trigger_to_sources: dict[str, set[TriggerEdge]] = defaultdict(set)
+        for src, triggers in sources.items():
+            for trigger, cond, label in triggers:
+                trigger_to_sources[trigger].add(TriggerEdge(src, cond, label))
+        # apply writes
+        updated_channels = apply_writes(
+            checkpoint, channels, tasks.values(), get_next_version, trigger_to_nodes
+        )
+        # prepare next tasks
+        tasks = prepare_next_tasks(
+            checkpoint,
+            [],
+            nodes,
+            channels,
+            managed,
+            config,
+            step,
+            limit,
+            for_execution=True,
+            store=None,
+            checkpointer=None,
+            manager=None,
+            trigger_to_nodes=trigger_to_nodes,
+            updated_channels=updated_channels,
+        )
+        # collect deferred nodes
+        deferred_nodes: set[str] = set()
+        edges_to_deferred_nodes: set[Edge] = set()
+        for channel, item in channels.items():
+            if isinstance(item, LastValueAfterFinish):
+                deferred_node = channel.split(":", 2)[-1]
+                deferred_nodes.add(deferred_node)
+        # collect edges
+        for task in tasks.values():
+            added = False
+            for trigger in task.triggers:
+                for src, cond, label in sorted(trigger_to_sources[trigger]):
+                    # record edge to be reviewed later
+                    if task.name in deferred_nodes:
+                        edges_to_deferred_nodes.add(Edge(src, task.name, cond, label))
+                    edges.add(Edge(src, task.name, cond, label))
+                    # if the edge is from this step, skip adding the implicit edges
+                    if (trigger, cond, label) in step_sources.get(src, set()):
+                        added = True
+                    else:
+                        sources[src].discard(TriggerEdge(trigger, cond, label))
+            # if no edges from this step, add implicit edges from all previous tasks
+            if not added:
+                for src in step_sources:
+                    edges.add(Edge(src, task.name, True, None))
+
+    # assemble the graph
+    graph = Graph()
+    # add nodes
+    for name, node in nodes.items():
+        metadata = dict(node.metadata or {})
+        if name in deferred_nodes:
+            metadata["defer"] = True
+        if name in interrupt_before_nodes and name in interrupt_after_nodes:
+            metadata["__interrupt"] = "before,after"
+        elif name in interrupt_before_nodes:
+            metadata["__interrupt"] = "before"
+        elif name in interrupt_after_nodes:
+            metadata["__interrupt"] = "after"
+        graph.add_node(node.bound, name, metadata=metadata or None)
+    # add start node
+    if START not in nodes:
+        graph.add_node(None, START)
+        for task in start_tasks.values():
+            add_edge(graph, START, task.name)
+    # add discovered edges
+    for src, dest, is_conditional, label in sorted(edges):
+        add_edge(
+            graph,
+            src,
+            dest,
+            data=label if label != dest else None,
+            conditional=is_conditional,
+        )
+    # add end edges
+    termini = {d for _, d, _, _ in edges if d != END}.difference(
+        s for s, _, _, _ in edges
+    )
+    end_edge_exists = any(d == END for _, d, _, _ in edges)
+    if termini:
+        for src in sorted(termini):
+            add_edge(graph, src, END)
+    elif len(step_sources) == 1 and not end_edge_exists:
+        for src in sorted(step_sources):
+            add_edge(graph, src, END, conditional=True)
+    # replace subgraphs
+    for name, subgraph in subgraphs.items():
+        if (
+            len(subgraph.nodes) > 1
+            and name in graph.nodes
+            and subgraph.first_node()
+            and subgraph.last_node()
+        ):
+            subgraph.trim_first_node()
+            subgraph.trim_last_node()
+            # replace the node with the subgraph
+            graph.nodes.pop(name)
+            first, last = graph.extend(subgraph, prefix=name)
+            for idx, edge in enumerate(graph.edges):
+                if edge.source == name:
+                    edge = edge.copy(source=cast(Node, last).id)
+                if edge.target == name:
+                    edge = edge.copy(target=cast(Node, first).id)
+                graph.edges[idx] = edge
+
+    return graph
+
+
+def add_edge(
+    graph: Graph,
+    source: str,
+    target: str,
+    *,
+    data: Any | None = None,
+    conditional: bool = False,
+) -> None:
+    """Add an edge to the graph."""
+    for edge in graph.edges:
+        if edge.source == source and edge.target == target:
+            return
+    if target not in graph.nodes and target == END:
+        graph.add_node(None, END)
+    graph.add_edge(graph.nodes[source], graph.nodes[target], data, conditional)

coco_graph/pregel/_executor.py

@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import time
+from collections.abc import Awaitable, Callable, Coroutine
+from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack
+from contextvars import copy_context
+from types import TracebackType
+from typing import (
+    Protocol,
+    TypeVar,
+    cast,
+)
+
+from coco_core.runnables import RunnableConfig
+from coco_core.runnables.config import get_executor_for_config
+from typing_extensions import ParamSpec
+
+from coco_graph._internal._future import CONTEXT_NOT_SUPPORTED, run_coroutine_threadsafe
+from coco_graph.errors import GraphBubbleUp
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+class Submit(Protocol[P, T]):
+    def __call__(  # type: ignore[valid-type]
+        self,
+        fn: Callable[P, T],
+        *args: P.args,
+        __name__: str | None = None,
+        __cancel_on_exit__: bool = False,
+        __reraise_on_exit__: bool = True,
+        __next_tick__: bool = False,
+        **kwargs: P.kwargs,
+    ) -> concurrent.futures.Future[T]: ...
+
+
+class BackgroundExecutor(AbstractContextManager):
+    """A context manager that runs sync tasks in the background.
+    Uses a thread pool executor to delegate tasks to separate threads.
+    On exit,
+    - cancels any (not yet started) tasks with `__cancel_on_exit__=True`
+    - waits for all tasks to finish
+    - re-raises the first exception from tasks with `__reraise_on_exit__=True`"""
+
+    def __init__(self, config: RunnableConfig) -> None:
+        self.stack = ExitStack()
+        self.executor = self.stack.enter_context(get_executor_for_config(config))
+        # mapping of Future to (__cancel_on_exit__, __reraise_on_exit__) flags
+        self.tasks: dict[concurrent.futures.Future, tuple[bool, bool]] = {}
+
+    def submit(  # type: ignore[valid-type]
+        self,
+        fn: Callable[P, T],
+        *args: P.args,
+        __name__: str | None = None,  # currently not used in sync version
+        __cancel_on_exit__: bool = False,  # for sync, can cancel only if not started
+        __reraise_on_exit__: bool = True,
+        __next_tick__: bool = False,
+        **kwargs: P.kwargs,
+    ) -> concurrent.futures.Future[T]:
+        ctx = copy_context()
+        if __next_tick__:
+            task = cast(
+                concurrent.futures.Future[T],
+                self.executor.submit(next_tick, ctx.run, fn, *args, **kwargs),  # type: ignore[arg-type]
+            )
+        else:
+            task = self.executor.submit(ctx.run, fn, *args, **kwargs)
+        self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__)
+        # add a callback to remove the task from the tasks dict when it's done
+        task.add_done_callback(self.done)
+        return task
+
+    def done(self, task: concurrent.futures.Future) -> None:
+        """Remove the task from the tasks dict when it's done."""
+        try:
+            task.result()
+        except GraphBubbleUp:
+            # This exception is an interruption signal, not an error
+            # so we don't want to re-raise it on exit
+            self.tasks.pop(task)
+        except BaseException:
+            pass
+        else:
+            self.tasks.pop(task)
+
+    def __enter__(self) -> Submit:
+        return self.submit
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool | None:
+        # copy the tasks as done() callback may modify the dict
+        tasks = self.tasks.copy()
+        # cancel all tasks that should be cancelled
+        for task, (cancel, _) in tasks.items():
+            if cancel:
+                task.cancel()
+        # wait for all tasks to finish
+        if pending := {t for t in tasks if not t.done()}:
+            concurrent.futures.wait(pending)
+        # shutdown the executor
+        self.stack.__exit__(exc_type, exc_value, traceback)
+        # if there's already an exception being raised, don't raise another one
+        if exc_type is None:
+            # re-raise the first exception that occurred in a task
+            for task, (_, reraise) in tasks.items():
+                if not reraise:
+                    continue
+                try:
+                    task.result()
+                except concurrent.futures.CancelledError:
+                    pass
+
+
+class AsyncBackgroundExecutor(AbstractAsyncContextManager):
+    """A context manager that runs async tasks in the background.
+    Uses the current event loop to delegate tasks to asyncio tasks.
+    On exit,
+    - cancels any tasks with `__cancel_on_exit__=True`
+    - waits for all tasks to finish
+    - re-raises the first exception from tasks with `__reraise_on_exit__=True`
+      ignoring CancelledError"""
+
+    def __init__(self, config: RunnableConfig) -> None:
+        self.tasks: dict[asyncio.Future, tuple[bool, bool]] = {}
+        self.sentinel = object()
+        self.loop = asyncio.get_running_loop()
+        if max_concurrency := config.get("max_concurrency"):
+            self.semaphore: asyncio.Semaphore | None = asyncio.Semaphore(
+                max_concurrency
+            )
+        else:
+            self.semaphore = None
+
+    def submit(  # type: ignore[valid-type]
+        self,
+        fn: Callable[P, Awaitable[T]],
+        *args: P.args,
+        __name__: str | None = None,
+        __cancel_on_exit__: bool = False,
+        __reraise_on_exit__: bool = True,
+        __next_tick__: bool = False,  # noop in async (always True)
+        **kwargs: P.kwargs,
+    ) -> asyncio.Future[T]:
+        coro = cast(Coroutine[None, None, T], fn(*args, **kwargs))
+        if self.semaphore:
+            coro = gated(self.semaphore, coro)
+        if CONTEXT_NOT_SUPPORTED:
+            task = run_coroutine_threadsafe(
+                coro, self.loop, name=__name__, lazy=__next_tick__
+            )
+        else:
+            task = run_coroutine_threadsafe(
+                coro,
+                self.loop,
+                name=__name__,
+                context=copy_context(),
+                lazy=__next_tick__,
+            )
+        self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__)
+        task.add_done_callback(self.done)
+        return task
+
+    def done(self, task: asyncio.Future) -> None:
+        try:
+            if exc := task.exception():
+                # This exception is an interruption signal, not an error
+                # so we don't want to re-raise it on exit
+                if isinstance(exc, GraphBubbleUp):
+                    self.tasks.pop(task)
+            else:
+                self.tasks.pop(task)
+        except asyncio.CancelledError:
+            self.tasks.pop(task)
+
+    async def __aenter__(self) -> Submit:
+        return self.submit
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        # copy the tasks as done() callback may modify the dict
+        tasks = self.tasks.copy()
+        # cancel all tasks that should be cancelled
+        for task, (cancel, _) in tasks.items():
+            if cancel:
+                task.cancel(self.sentinel)
+        # wait for all tasks to finish
+        if tasks:
+            await asyncio.wait(tasks)
+        # if there's already an exception being raised, don't raise another one
+        if exc_type is None:
+            # re-raise the first exception that occurred in a task
+            for task, (_, reraise) in tasks.items():
+                if not reraise:
+                    continue
+                try:
+                    if exc := task.exception():
+                        raise exc
+                except asyncio.CancelledError:
+                    pass
+
+
+async def gated(semaphore: asyncio.Semaphore, coro: Coroutine[None, None, T]) -> T:
+    """A coroutine that waits for a semaphore before running another coroutine."""
+    async with semaphore:
+        return await coro
+
+
+def next_tick(fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
+    """A function that yields control to other threads before running another function."""
+    time.sleep(0)
+    return fn(*args, **kwargs)

coco_graph/pregel/_io.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+from collections import Counter
+from collections.abc import Iterator, Mapping, Sequence
+from typing import Any, Literal
+
+from coco_graph._internal._constants import (
+    ERROR,
+    INTERRUPT,
+    NULL_TASK_ID,
+    RESUME,
+    RETURN,
+    TASKS,
+)
+from coco_graph._internal._typing import EMPTY_SEQ, MISSING
+from coco_graph.channels.base import BaseChannel, EmptyChannelError
+from coco_graph.constants import START, TAG_HIDDEN
+from coco_graph.errors import InvalidUpdateError
+from coco_graph.pregel._log import logger
+from coco_graph.types import Command, PregelExecutableTask, Send
+
+
+def read_channel(
+    channels: Mapping[str, BaseChannel],
+    chan: str,
+    *,
+    catch: bool = True,
+) -> Any:
+    try:
+        return channels[chan].get()
+    except EmptyChannelError:
+        if catch:
+            return None
+        else:
+            raise
+
+
+def read_channels(
+    channels: Mapping[str, BaseChannel],
+    select: Sequence[str] | str,
+    *,
+    skip_empty: bool = True,
+) -> dict[str, Any] | Any:
+    if isinstance(select, str):
+        return read_channel(channels, select)
+    else:
+        values: dict[str, Any] = {}
+        for k in select:
+            try:
+                values[k] = read_channel(channels, k, catch=not skip_empty)
+            except EmptyChannelError:
+                pass
+        return values
+
+
+def map_command(cmd: Command) -> Iterator[tuple[str, str, Any]]:
+    """Map input chunk to a sequence of pending writes in the form (channel, value)."""
+    if cmd.graph == Command.PARENT:
+        raise InvalidUpdateError("There is no parent graph")
+    if cmd.goto:
+        if isinstance(cmd.goto, (tuple, list)):
+            sends = cmd.goto
+        else:
+            sends = [cmd.goto]
+        for send in sends:
+            if isinstance(send, Send):
+                yield (NULL_TASK_ID, TASKS, send)
+            elif isinstance(send, str):
+                yield (NULL_TASK_ID, f"branch:to:{send}", START)
+            else:
+                raise TypeError(
+                    f"In Command.goto, expected Send/str, got {type(send).__name__}"
+                )
+    if cmd.resume is not None:
+        yield (NULL_TASK_ID, RESUME, cmd.resume)
+    if cmd.update:
+        for k, v in cmd._update_as_tuples():
+            yield (NULL_TASK_ID, k, v)
+
+
+def map_input(
+    input_channels: str | Sequence[str],
+    chunk: dict[str, Any] | Any | None,
+) -> Iterator[tuple[str, Any]]:
+    """Map input chunk to a sequence of pending writes in the form (channel, value)."""
+    if chunk is None:
+        return
+    elif isinstance(input_channels, str):
+        yield (input_channels, chunk)
+    else:
+        if not isinstance(chunk, dict):
+            raise TypeError(f"Expected chunk to be a dict, got {type(chunk).__name__}")
+        for k in chunk:
+            if k in input_channels:
+                yield (k, chunk[k])
+            else:
+                logger.warning(f"Input channel {k} not found in {input_channels}")
+
+
+def map_output_values(
+    output_channels: str | Sequence[str],
+    pending_writes: Literal[True] | Sequence[tuple[str, Any]],
+    channels: Mapping[str, BaseChannel],
+) -> Iterator[dict[str, Any] | Any]:
+    """Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
+    if isinstance(output_channels, str):
+        if pending_writes is True or any(
+            chan == output_channels for chan, _ in pending_writes
+        ):
+            yield read_channel(channels, output_channels)
+    else:
+        if pending_writes is True or {
+            c for c, _ in pending_writes if c in output_channels
+        }:
+            yield read_channels(channels, output_channels)
+
+
+def map_output_updates(
+    output_channels: str | Sequence[str],
+    tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]],
+    cached: bool = False,
+) -> Iterator[dict[str, Any | dict[str, Any]]]:
+    """Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
+    output_tasks = [
+        (t, ww)
+        for t, ww in tasks
+        if (not t.config or TAG_HIDDEN not in t.config.get("tags", EMPTY_SEQ))
+        and ww[0][0] != ERROR
+        and ww[0][0] != INTERRUPT
+    ]
+    if not output_tasks:
+        return
+    updated: list[tuple[str, Any]] = []
+    for task, writes in output_tasks:
+        rtn = next((value for chan, value in writes if chan == RETURN), MISSING)
+        if rtn is not MISSING:
+            updated.append((task.name, rtn))
+        elif isinstance(output_channels, str):
+            updated.extend(
+                (task.name, value) for chan, value in writes if chan == output_channels
+            )
+        elif any(chan in output_channels for chan, _ in writes):
+            counts = Counter(chan for chan, _ in writes)
+            if any(counts[chan] > 1 for chan in output_channels):
+                updated.extend(
+                    (
+                        task.name,
+                        {chan: value},
+                    )
+                    for chan, value in writes
+                    if chan in output_channels
+                )
+            else:
+                updated.append(
+                    (
+                        task.name,
+                        {
+                            chan: value
+                            for chan, value in writes
+                            if chan in output_channels
+                        },
+                    )
+                )
+    grouped: dict[str, Any] = {t.name: [] for t, _ in output_tasks}
+    for node, value in updated:
+        grouped[node].append(value)
+    for node, value in grouped.items():
+        if len(value) == 0:
+            grouped[node] = None
+        if len(value) == 1:
+            grouped[node] = value[0]
+    if cached:
+        grouped["__metadata__"] = {"cached": cached}
+    yield grouped

coco_graph/pregel/_log.py

@@ -0,0 +1,3 @@
+import logging
+
+logger = logging.getLogger("coco_graph")