coco-framework 0.1.0__py3-none-any.whl

coco_graph/channels/base.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import Any, Generic, TypeVar
+
+from typing_extensions import Self
+
+from coco_graph._internal._typing import MISSING
+from coco_graph.errors import EmptyChannelError
+
+Value = TypeVar("Value")
+Update = TypeVar("Update")
+Checkpoint = TypeVar("Checkpoint")
+
+__all__ = ("BaseChannel",)
+
+
+class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
+    """Base class for all channels."""
+
+    __slots__ = ("key", "typ")
+
+    def __init__(self, typ: Any, key: str = "") -> None:
+        self.typ = typ
+        self.key = key
+
+    @property
+    @abstractmethod
+    def ValueType(self) -> Any:
+        """The type of the value stored in the channel."""
+
+    @property
+    @abstractmethod
+    def UpdateType(self) -> Any:
+        """The type of the update received by the channel."""
+
+    # serialize/deserialize methods
+
+    def copy(self) -> Self:
+        """Return a copy of the channel.
+
+        By default, delegates to `checkpoint()` and `from_checkpoint()`.
+
+        Subclasses can override this method with a more efficient implementation.
+        """
+        return self.from_checkpoint(self.checkpoint())
+
+    def checkpoint(self) -> Checkpoint | Any:
+        """Return a serializable representation of the channel's current state.
+
+        Raises `EmptyChannelError` if the channel is empty (never updated yet),
+        or doesn't support checkpoints.
+        """
+        try:
+            return self.get()
+        except EmptyChannelError:
+            return MISSING
+
+    @abstractmethod
+    def from_checkpoint(self, checkpoint: Checkpoint | Any) -> Self:
+        """Return a new identical channel, optionally initialized from a checkpoint.
+
+        If the checkpoint contains complex data structures, they should be copied.
+        """
+
+    # read methods
+
+    @abstractmethod
+    def get(self) -> Value:
+        """Return the current value of the channel.
+
+        Raises `EmptyChannelError` if the channel is empty (never updated yet)."""
+
+    def is_available(self) -> bool:
+        """Return `True` if the channel is available (not empty), `False` otherwise.
+
+        Subclasses should override this method to provide a more efficient
+        implementation than calling `get()` and catching `EmptyChannelError`.
+        """
+        try:
+            self.get()
+            return True
+        except EmptyChannelError:
+            return False
+
+    # write methods
+
+    @abstractmethod
+    def update(self, values: Sequence[Update]) -> bool:
+        """Update the channel's value with the given sequence of updates.
+        The order of the updates in the sequence is arbitrary.
+        This method is called by Pregel for all channels at the end of each step.
+
+        If there are no updates, it is called with an empty sequence.
+
+        Raises `InvalidUpdateError` if the sequence of updates is invalid.
+
+        Returns `True` if the channel was updated, `False` otherwise."""
+
+    def consume(self) -> bool:
+        """Notify the channel that a subscribed task ran.
+
+        By default, no-op.
+
+        A channel can use this method to modify its state, preventing the value from being consumed again.
+
+        Returns `True` if the channel was updated, `False` otherwise.
+        """
+        return False
+
+    def finish(self) -> bool:
+        """Notify the channel that the Pregel run is finishing.
+
+        By default, no-op.
+
+        A channel can use this method to modify its state, preventing finish.
+
+        Returns `True` if the channel was updated, `False` otherwise.
+        """
+        return False

coco_graph/channels/binop.py

@@ -0,0 +1,155 @@
+import collections.abc
+from collections.abc import Callable, Sequence
+from typing import Any, Generic
+
+from typing_extensions import NotRequired, Required, Self
+
+from coco_graph._internal._constants import OVERWRITE
+from coco_graph._internal._typing import MISSING
+from coco_graph.channels.base import BaseChannel, Value
+from coco_graph.errors import (
+    EmptyChannelError,
+    ErrorCode,
+    InvalidUpdateError,
+    create_error_message,
+)
+from coco_graph.types import Overwrite
+
+__all__ = ("BinaryOperatorAggregate",)
+
+
+# Adapted from typing_extensions
+def _strip_extras(t):  # type: ignore[no-untyped-def]
+    """Strips Annotated, Required and NotRequired from a given type."""
+    if hasattr(t, "__origin__"):
+        if t.__origin__ in (Required, NotRequired):
+            return _strip_extras(t.__args__[0])
+        return _strip_extras(t.__origin__)
+    return t
+
+
+def _get_overwrite(value: Any) -> tuple[bool, Any]:
+    """Inspects the given value and returns (is_overwrite, overwrite_value).
+
+    Recognises three forms:
+
+    * The typed `Overwrite` dataclass instance.
+    * The sentinel-keyed `{"__overwrite__": value}` dict form.
+    * The dataclass-erased `{"value": ..., "type": "__overwrite__"}` form that
+      results from JSON-serialising an `Overwrite` (e.g. an `orjson`-encoded
+      state update routed through the CocoGraph API server). This keeps the
+      `Overwrite` semantics intact across JSON boundaries that strip dataclass
+      types.
+    """
+    if isinstance(value, Overwrite):
+        return True, value.value
+    if isinstance(value, dict):
+        if len(value) == 1 and OVERWRITE in value:
+            return True, value[OVERWRITE]
+        if value.get("type") == OVERWRITE and "value" in value:
+            return True, value["value"]
+    return False, None
+
+
+def _operators_equal(a: Callable, b: Callable) -> bool:
+    """Return True if two reducer operators should be considered equal.
+
+    Lambdas all share the name '<lambda>' so identity comparison is
+    unreliable; treat any pairing that includes a lambda as equal.
+    """
+    if a.__name__ == "<lambda>" or b.__name__ == "<lambda>":
+        return True
+    return a is b
+
+
+class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
+    """Stores the result of applying a binary operator to the current value and each new value.
+
+    ```python
+    import operator
+
+    total = Channels.BinaryOperatorAggregate(int, operator.add)
+    ```
+    """
+
+    __slots__ = ("value", "operator")
+
+    def __init__(self, typ: type[Value], operator: Callable[[Value, Value], Value]):
+        super().__init__(typ)
+        self.operator = operator
+        # special forms from typing or collections.abc are not instantiable
+        # so we need to replace them with their concrete counterparts
+        typ = _strip_extras(typ)
+        if typ in (collections.abc.Sequence, collections.abc.MutableSequence):
+            typ = list
+        if typ in (collections.abc.Set, collections.abc.MutableSet):
+            typ = set
+        if typ in (collections.abc.Mapping, collections.abc.MutableMapping):
+            typ = dict
+        try:
+            self.value = typ()
+        except Exception:
+            self.value = MISSING
+
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, BinaryOperatorAggregate) and _operators_equal(
+            self.operator, value.operator
+        )
+
+    @property
+    def ValueType(self) -> type[Value]:
+        """The type of the value stored in the channel."""
+        return self.typ
+
+    @property
+    def UpdateType(self) -> type[Value]:
+        """The type of the update received by the channel."""
+        return self.typ
+
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.operator)
+        empty.key = self.key
+        empty.value = self.value
+        return empty
+
+    def from_checkpoint(self, checkpoint: Value) -> Self:
+        empty = self.__class__(self.typ, self.operator)
+        empty.key = self.key
+        if checkpoint is not MISSING:
+            empty.value = checkpoint
+        return empty
+
+    def update(self, values: Sequence[Value]) -> bool:
+        if not values:
+            return False
+        if self.value is MISSING:
+            self.value = values[0]
+            values = values[1:]
+        seen_overwrite: bool = False
+        for value in values:
+            is_overwrite, overwrite_value = _get_overwrite(value)
+            if is_overwrite:
+                if seen_overwrite:
+                    msg = create_error_message(
+                        message="Can receive only one Overwrite value per super-step.",
+                        error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
+                    )
+                    raise InvalidUpdateError(msg)
+                self.value = overwrite_value
+                seen_overwrite = True
+                continue
+            if not seen_overwrite:
+                self.value = self.operator(self.value, value)
+        return True
+
+    def get(self) -> Value:
+        if self.value is MISSING:
+            raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING
+
+    def checkpoint(self) -> Value:
+        return self.value

coco_graph/channels/delta.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+import collections.abc
+import copy as _copy
+from collections.abc import Callable, Sequence
+from typing import Any, Generic
+
+from coco_graph.checkpoint.base import PendingWrite
+from coco_graph.checkpoint.serde.types import _DeltaSnapshot
+from typing_extensions import Self
+
+from coco_graph._internal._typing import MISSING
+from coco_graph.channels.base import BaseChannel, Value
+from coco_graph.channels.binop import _get_overwrite, _operators_equal, _strip_extras
+from coco_graph.errors import (
+    EmptyChannelError,
+    ErrorCode,
+    InvalidUpdateError,
+    create_error_message,
+)
+
+__all__ = ("DeltaChannel",)
+
+
+class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
+    """Reducer channel that stores only a sentinel in checkpoint blobs and
+    reconstructs state by replaying ancestor writes through the reducer.
+
+    !!! warning "Beta"
+
+        `DeltaChannel` is in beta. The API and on-disk representation may
+        change in future releases. Threads written with `DeltaChannel` today
+        are expected to remain readable, but the surrounding contract
+        (`BaseCheckpointSaver.get_delta_channel_history`, the
+        `_DeltaSnapshot` blob shape, the `counters_since_delta_snapshot`
+        metadata field) is not yet stable.
+
+    The reducer receives the current accumulated value and a batch of writes
+    in one call: `reducer(state, [write1, write2, ...]) -> new_state`.
+
+    Reducers must be deterministic and batching-invariant (associative across
+    folds): applying two consecutive write batches separately must produce the
+    same state as applying their concatenation once:
+
+        reducer(reducer(state, xs), ys) == reducer(state, xs + ys)
+
+    This lets CocoGraph replay checkpointed writes in larger batches than they
+    were originally produced without changing reconstructed state.
+
+    Snapshot cadence is driven by two counters: per-channel update count and
+    total supersteps since last snapshot. `create_checkpoint` writes a full
+    `_DeltaSnapshot` blob when EITHER the update count reaches
+    `snapshot_frequency` OR the supersteps count reaches the system-wide
+    `DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT` bound (default 5000), bounding
+    replay depth even for channels that stop receiving writes.
+
+    Parameters:
+        reducer: `(state, list[writes]) -> new_state`. Must be deterministic
+            and batching-invariant as described above.
+        typ: The value type (e.g. `list`, `dict`). Inferred automatically
+            from the outer type when used inside `Annotated[T, DeltaChannel(...)]`.
+        snapshot_frequency: Every Nth update to this channel writes a snapshot
+            blob (default `1000`). Must be a positive int.
+    """
+
+    __slots__ = ("value", "reducer", "snapshot_frequency")
+    value: Value | Any
+
+    def __init__(
+        self,
+        reducer: Callable[[Any, Sequence[Any]], Any],
+        typ: type[Value] | None = None,
+        *,
+        snapshot_frequency: int = 1000,
+    ) -> None:
+        if snapshot_frequency <= 0:
+            raise ValueError(
+                f"snapshot_frequency must be a positive int, got {snapshot_frequency}"
+            )
+        if typ is None:
+            typ = list  # type: ignore[assignment]  # placeholder; overridden by _is_field_channel
+        super().__init__(typ)
+        self.reducer = reducer
+        self.snapshot_frequency = snapshot_frequency
+        typ = _strip_extras(typ)
+        if typ in (collections.abc.Sequence, collections.abc.MutableSequence):
+            typ = list
+        if typ in (collections.abc.Set, collections.abc.MutableSet):
+            typ = set
+        if typ in (collections.abc.Mapping, collections.abc.MutableMapping):
+            typ = dict
+        self.typ = typ
+        self.value: Any = MISSING
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, DeltaChannel):
+            return False
+        if self.snapshot_frequency != other.snapshot_frequency:
+            return False
+        return _operators_equal(self.reducer, other.reducer)
+
+    @property
+    def ValueType(self) -> Any:
+        return self.typ
+
+    @property
+    def UpdateType(self) -> Any:
+        return self.typ
+
+    def copy(self) -> Self:
+        new = self.__class__(
+            self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency
+        )
+        new.key = self.key
+        new.value = self.value if self.value is MISSING else _copy.copy(self.value)
+        return new
+
+    def from_checkpoint(self, checkpoint: Any) -> Self:
+        """Initialize from a stored blob.
+
+        Blob types:
+          * `MISSING`: start empty; caller replays writes.
+          * `_DeltaSnapshot(value)`: restore value directly from snapshot.
+          * plain value (migration from old `BinaryOperatorAggregate` blobs):
+            use directly.
+        """
+        new = self.__class__(
+            self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency
+        )
+        new.key = self.key
+        if checkpoint is MISSING:
+            new.value = self.typ()
+        elif isinstance(checkpoint, _DeltaSnapshot):
+            new.value = checkpoint.value
+        else:
+            new.value = checkpoint
+        return new
+
+    def replay_writes(self, writes: Sequence[PendingWrite]) -> None:
+        """Apply ancestor writes oldest-to-newest via a single reducer call.
+
+        If any write is an Overwrite, the last one in the sequence acts as
+        the reset point: its value becomes the new base and only writes
+        after it are passed to the reducer.
+        """
+        values = [v for _, _, v in writes]
+        if not values:
+            return
+        base = self.value
+        start = 0
+        for i, v in enumerate(values):
+            is_ow, ow_value = _get_overwrite(v)
+            if is_ow:
+                base = _copy.copy(ow_value) if ow_value is not None else self.typ()
+                start = i + 1
+        remaining = values[start:]
+        self.value = self.reducer(base, remaining) if remaining else base
+
+    def update(self, values: Sequence[Any]) -> bool:
+        if not values:
+            return False
+        overwrite_idx: int | None = None
+        for i, v in enumerate(values):
+            is_ow, _ = _get_overwrite(v)
+            if is_ow:
+                if overwrite_idx is not None:
+                    msg = create_error_message(
+                        message="Can receive only one Overwrite value per super-step.",
+                        error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
+                    )
+                    raise InvalidUpdateError(msg)
+                overwrite_idx = i
+        if overwrite_idx is not None:
+            _, overwrite_value = _get_overwrite(values[overwrite_idx])
+            self.value = (
+                _copy.copy(overwrite_value)
+                if overwrite_value is not None
+                else self.typ()
+            )
+            return True
+        base = self.typ() if self.value is MISSING else self.value
+        self.value = self.reducer(base, list(values))
+        return True
+
+    def get(self) -> Any:
+        if self.value is MISSING:
+            raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING
+
+    def checkpoint(self) -> Any:
+        """Return stored representation: always `MISSING`.
+
+        Snapshot decisions live in `create_checkpoint` (which has the channel
+        version) and write `_DeltaSnapshot(ch.get())` directly into
+        `channel_values`. For non-snapshot steps the channel does not appear
+        in `channel_values`; reconstruction walks ancestor writes via the
+        saver's `get_delta_channel_history`.
+        """
+        return MISSING

coco_graph/channels/ephemeral_value.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any, Generic
+
+from typing_extensions import Self
+
+from coco_graph._internal._typing import MISSING
+from coco_graph.channels.base import BaseChannel, Value
+from coco_graph.errors import EmptyChannelError, InvalidUpdateError
+
+__all__ = ("EphemeralValue",)
+
+
+class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
+    """Stores the value received in the step immediately preceding, clears after."""
+
+    __slots__ = ("value", "guard")
+
+    value: Value | Any
+    guard: bool
+
+    def __init__(self, typ: Any, guard: bool = True) -> None:
+        super().__init__(typ)
+        self.guard = guard
+        self.value = MISSING
+
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, EphemeralValue) and value.guard == self.guard
+
+    @property
+    def ValueType(self) -> type[Value]:
+        """The type of the value stored in the channel."""
+        return self.typ
+
+    @property
+    def UpdateType(self) -> type[Value]:
+        """The type of the update received by the channel."""
+        return self.typ
+
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.guard)
+        empty.key = self.key
+        empty.value = self.value
+        return empty
+
+    def from_checkpoint(self, checkpoint: Value) -> Self:
+        empty = self.__class__(self.typ, self.guard)
+        empty.key = self.key
+        if checkpoint is not MISSING:
+            empty.value = checkpoint
+        return empty
+
+    def update(self, values: Sequence[Value]) -> bool:
+        if len(values) == 0:
+            if self.value is not MISSING:
+                self.value = MISSING
+                return True
+            else:
+                return False
+        if len(values) != 1 and self.guard:
+            raise InvalidUpdateError(
+                f"At key '{self.key}': EphemeralValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values."
+            )
+
+        self.value = values[-1]
+        return True
+
+    def get(self) -> Value:
+        if self.value is MISSING:
+            raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING
+
+    def checkpoint(self) -> Value:
+        return self.value