coco-framework 0.1.0__py3-none-any.whl

coco_graph/_internal/_queue.py

@@ -0,0 +1,124 @@
+# type: ignore
+from __future__ import annotations
+
+import asyncio
+import queue
+import threading
+import types
+from collections import deque
+from time import monotonic
+
+
+class AsyncQueue(asyncio.Queue):
+    """Async unbounded FIFO queue with a wait() method.
+
+    Subclassed from asyncio.Queue, adding a wait() method."""
+
+    async def wait(self) -> None:
+        """If queue is empty, wait until an item is available.
+
+        Copied from Queue.get(), removing the call to .get_nowait(),
+        ie. this doesn't consume the item, just waits for it.
+        """
+        while self.empty():
+            getter = self._get_loop().create_future()
+            self._getters.append(getter)
+            try:
+                await getter
+            except BaseException:
+                getter.cancel()  # Just in case getter is not done yet.
+                try:
+                    # Clean self._getters from canceled getters.
+                    self._getters.remove(getter)
+                except ValueError:
+                    # The getter could be removed from self._getters by a
+                    # previous put_nowait call.
+                    pass
+                if not self.empty() and not getter.cancelled():
+                    # We were woken up by put_nowait(), but can't take
+                    # the call.  Wake up the next in line.
+                    self._wakeup_next(self._getters)
+                raise
+
+
+class Semaphore(threading.Semaphore):
+    """Semaphore subclass with a wait() method."""
+
+    def wait(self, blocking: bool = True, timeout: float | None = None):
+        """Block until the semaphore can be acquired, but don't acquire it."""
+        if not blocking and timeout is not None:
+            raise ValueError("can't specify timeout for non-blocking acquire")
+        rc = False
+        endtime = None
+        with self._cond:
+            while self._value == 0:
+                if not blocking:
+                    break
+                if timeout is not None:
+                    if endtime is None:
+                        endtime = monotonic() + timeout
+                    else:
+                        timeout = endtime - monotonic()
+                        if timeout <= 0:
+                            break
+                self._cond.wait(timeout)
+            else:
+                rc = True
+        return rc
+
+
+class SyncQueue:
+    """Unbounded FIFO queue with a wait() method.
+    Adapted from pure Python implementation of queue.SimpleQueue.
+    """
+
+    def __init__(self):
+        self._queue = deque()
+        self._count = Semaphore(0)
+
+    def put(self, item, block=True, timeout=None):
+        """Put the item on the queue.
+
+        The optional 'block' and 'timeout' arguments are ignored, as this method
+        never blocks.  They are provided for compatibility with the Queue class.
+        """
+        self._queue.append(item)
+        self._count.release()
+
+    def get(self, block=False, timeout=None):
+        """Remove and return an item from the queue.
+
+        If optional args 'block' is true and 'timeout' is None (the default),
+        block if necessary until an item is available. If 'timeout' is
+        a non-negative number, it blocks at most 'timeout' seconds and raises
+        the Empty exception if no item was available within that time.
+        Otherwise ('block' is false), return an item if one is immediately
+        available, else raise the Empty exception ('timeout' is ignored
+        in that case).
+        """
+        if timeout is not None and timeout < 0:
+            raise ValueError("'timeout' must be a non-negative number")
+        if not self._count.acquire(block, timeout):
+            raise queue.Empty
+        try:
+            return self._queue.popleft()
+        except IndexError:
+            raise queue.Empty
+
+    def wait(self, block=True, timeout=None):
+        """If queue is empty, wait until an item maybe is available,
+        but don't consume it.
+        """
+        if timeout is not None and timeout < 0:
+            raise ValueError("'timeout' must be a non-negative number")
+        self._count.wait(block, timeout)
+
+    def empty(self):
+        """Return True if the queue is empty, False otherwise (not reliable!)."""
+        return len(self._queue) == 0
+
+    def qsize(self):
+        """Return the approximate size of the queue (not reliable!)."""
+        return len(self._queue)
+
+    __class_getitem__ = classmethod(types.GenericAlias)

coco_graph/_internal/_replay.py

@@ -0,0 +1,90 @@
+"""Replay state for subgraph checkpoint loading during time-travel."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from coco_graph._internal._constants import NS_END
+
+if TYPE_CHECKING:
+    from coco_core.runnables import RunnableConfig
+    from coco_graph.checkpoint.base import BaseCheckpointSaver, CheckpointTuple
+
+
+class ReplayState:
+    """Tracks which subgraphs have already loaded their pre-replay checkpoint.
+
+    During a parent replay, each subgraph's first invocation should restore the
+    checkpoint from before the replay point. Subsequent invocations of the same
+    subgraph (e.g. in a loop) should use normal checkpoint loading so they pick
+    up freshly created checkpoints.
+
+    The single `ReplayState` instance is shared by reference across all derived
+    configs within one parent execution.
+    """
+
+    __slots__ = ("checkpoint_id", "_visited_ns")
+
+    def __init__(self, checkpoint_id: str) -> None:
+        self.checkpoint_id = checkpoint_id
+        # DO NOT CHANGE THIS VARIABLE – it may need to be rehydrated
+        # in other runtimes
+        self._visited_ns: set[str] = set()
+
+    def _is_first_visit(self, checkpoint_ns: str) -> bool:
+        """Return True the first time a subgraph namespace is seen.
+
+        The task-id suffix is stripped so that the same logical subgraph
+        (e.g. ``"sub_node"``) is recognized across loop iterations even
+        though each iteration has a different task id.
+        """
+        # "sub_node:task_id" -> "sub_node"
+        stable_ns = (
+            checkpoint_ns.rsplit(NS_END, 1)[0]
+            if NS_END in checkpoint_ns
+            else checkpoint_ns
+        )
+        if stable_ns in self._visited_ns:
+            return False
+        self._visited_ns.add(stable_ns)
+        return True
+
+    def get_checkpoint(
+        self,
+        checkpoint_ns: str,
+        checkpointer: BaseCheckpointSaver,
+        checkpoint_config: RunnableConfig,
+    ) -> CheckpointTuple | None:
+        """Load the right checkpoint for a subgraph during replay.
+
+        On the first call for a given subgraph namespace, returns the latest
+        checkpoint created *before* the replay point. On subsequent calls
+        (e.g. the same subgraph in a later loop iteration), falls back to
+        normal latest-checkpoint loading.
+        """
+        if self._is_first_visit(checkpoint_ns):
+            for saved in checkpointer.list(
+                checkpoint_config,
+                before={"configurable": {"checkpoint_id": self.checkpoint_id}},
+                limit=1,
+            ):
+                return saved
+            return None
+        return checkpointer.get_tuple(checkpoint_config)
+
+    async def aget_checkpoint(
+        self,
+        checkpoint_ns: str,
+        checkpointer: BaseCheckpointSaver,
+        checkpoint_config: RunnableConfig,
+    ) -> CheckpointTuple | None:
+        """Async version of `get_checkpoint`."""
+        if self._is_first_visit(checkpoint_ns):
+            async for saved in checkpointer.alist(
+                checkpoint_config,
+                before={"configurable": {"checkpoint_id": self.checkpoint_id}},
+                limit=1,
+            ):
+                return saved
+            return None
+        return await checkpointer.aget_tuple(checkpoint_config)

coco_graph/_internal/_retry.py

@@ -0,0 +1,29 @@
+def default_retry_on(exc: Exception) -> bool:
+    import httpx
+    import requests
+
+    if isinstance(exc, ConnectionError):
+        return True
+    if isinstance(exc, httpx.HTTPStatusError):
+        return 500 <= exc.response.status_code < 600
+    if isinstance(exc, requests.HTTPError):
+        return 500 <= exc.response.status_code < 600 if exc.response else True
+    if isinstance(
+        exc,
+        (
+            ValueError,
+            TypeError,
+            ArithmeticError,
+            ImportError,
+            LookupError,
+            NameError,
+            SyntaxError,
+            RuntimeError,
+            ReferenceError,
+            StopIteration,
+            StopAsyncIteration,
+            OSError,
+        ),
+    ):
+        return False
+    return True