Add support for checkpoints

adapters/python/CHANGELOG.md

@@ -1,6 +1,9 @@
 ## 0.12.0
 
-No changes.
+Enhancements:
+
+- Adds `cf.Checkpoint` for state that survives across executions of a step — retries, suspensions, recurrences and re-runs. Supports `get`, `set` and `reset`, with a declared default.
+- Adds `cf.flush` for synchronously flushing buffered state to the server.
 
 ## 0.11.0
 

adapters/python/coflux/__init__.py

@@ -12,6 +12,7 @@
 from pathlib import Path
 
 from ._version import __version__
+from .checkpoint import Checkpoint, flush
 from .decorators import stub, task, workflow
 from .errors import (
     ExecutionAbandoned,
@@ -63,6 +64,7 @@
     "MetricScale",
     "Prompt",
     "Cache",
+    "Checkpoint",
     "Defer",
     "Retries",
     "Streams",
@@ -86,6 +88,7 @@
     "log_error",
     "progress",
     "asset",
+    "flush",
 ]
 
 

adapters/python/coflux/checkpoint.py

@@ -0,0 +1,134 @@
+"""Step-scoped durable state."""
+
+from __future__ import annotations
+
+from typing import Any, Generic, TypeVar, overload
+
+from .state import get_context
+
+T = TypeVar("T")
+
+
+class Checkpoint(Generic[T]):
+    """A named value that survives across executions of a step.
+
+    Created via ``cf.Checkpoint(...)``, usually at module level. The name
+    identifies storage scoped to the current step and workspace, so the value
+    written by one attempt is what the next attempt reads — across retries,
+    suspends, recurrences and manual re-runs.
+
+    ::
+
+        cursor = cf.Checkpoint("cursor", default=0)
+
+        @cf.workflow(recurrent=True, delay=60)
+        def poll_orders():
+            since = cursor.get()
+            orders, next_since = fetch_orders(since)
+            for order in orders:
+                process_order.submit(order)
+            cursor.set(next_since)
+
+    Writes are throttled, so a crash can lose up to one throttle window: treat
+    a checkpoint as at-least-once and make the code that follows a read safe
+    to re-run from it. ``cf.flush()`` gives an explicit boundary where that
+    isn't good enough. Whatever is written before an execution suspends,
+    returns or fails is always delivered.
+
+    A checkpoint is not part of any cache, memo or defer key, and a step that
+    resolves from the cache never runs and never sees one.
+
+    ``T`` is the type ``get()`` returns. It's inferred from ``default`` when
+    one is given, so ``cursor`` above is a ``Checkpoint[int]``. Without a
+    default the checkpoint can also read as ``None``, so spell that out::
+
+        cursor = cf.Checkpoint[int | None]("cursor")
+
+    As with task arguments and results, this only informs type checkers —
+    nothing is enforced at runtime.
+
+    Args:
+        name: Checkpoint name, unique within the step.
+        default: Value returned when the checkpoint has never been set, or has
+            been reset. Client-side only — the server never sees it.
+    """
+
+    @overload
+    def __init__(self, name: str, *, default: T) -> None: ...
+
+    @overload
+    def __init__(self, name: str) -> None: ...
+
+    # ``Any`` rather than ``T | None``: the stored default has to satisfy the
+    # ``-> T`` on ``default`` and ``get()``, which it can't when ``T`` is
+    # non-optional and no default was given.
+    def __init__(self, name: str, *, default: Any = None) -> None:
+        self._name = name
+        self._default = default
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def default(self) -> T:
+        return self._default
+
+    def get(self) -> T:
+        """The current value, or the declared default if it isn't set.
+
+        A checkpoint explicitly set to ``None`` reads back as ``None``; only
+        an unset or reset checkpoint falls back to the default.
+        """
+        try:
+            return get_context().checkpoint_get(self._name)
+        except KeyError:
+            return self._default
+
+    def is_set(self) -> bool:
+        """Whether the checkpoint has a value (including an explicit ``None``)."""
+        return get_context().checkpoint_has(self._name)
+
+    def set(self, value: T) -> None:
+        """Set the value, replacing anything already there."""
+        get_context().checkpoint_set(self._name, value)
+
+    def reset(self) -> None:
+        """Clear the checkpoint, so ``get()`` returns the declared default.
+
+        Distinct from ``set(None)``, which stores ``None`` as a value.
+        """
+        get_context().checkpoint_reset(self._name)
+
+    def __repr__(self) -> str:
+        return f"Checkpoint({self._name!r})"
+
+    def __reduce__(self):
+        # Unlike cf.Metric, a checkpoint handle names step-scoped storage
+        # rather than describing itself. Serialising one and passing it to
+        # another execution would silently rebind it to that execution's step.
+        raise TypeError(
+            "Checkpoint handles can't be passed between executions — "
+            "declare cf.Checkpoint(...) in the target that uses it"
+        )
+
+
+def flush() -> None:
+    """Block until buffered state has reached the server.
+
+    Checkpoint writes are throttled and metrics and logs are batched; this
+    delivers whatever is outstanding and returns once the server has
+    acknowledged it. Useful for pinning a checkpoint before a side effect that
+    shouldn't be repeated::
+
+        cursor.set(next_cursor)
+        cf.flush()
+        send_notification()
+
+    Not needed before suspending, returning or raising — those are flushed
+    automatically.
+    """
+    get_context().flush()
+
+
+__all__ = ["Checkpoint", "flush"]

adapters/python/coflux/context.py

@@ -118,6 +118,14 @@ def __init__(self, execution_id: str, working_dir: Path | None = None):
         # executor from the target's ``@cf.task(streams=...)`` setting.
         # Used by ``cf.stream(...)`` to fill in unspecified options.
         self._default_streams: Streams | None = None
+        # Checkpoint state, split by whether it has been materialised yet.
+        # ``_checkpoint_wire`` holds what arrived with the execute message,
+        # still in protocol form — a checkpoint that's never read is never
+        # deserialised. ``_checkpoint_values`` holds materialised reads and
+        # anything written by this execution, and always wins. Both are
+        # guarded by ``self._lock``.
+        self._checkpoint_wire: dict[str, Any] = {}
+        self._checkpoint_values: dict[str, Any] = {}
 
     def set_default_streams(self, streams: Streams | None) -> None:
         """Record the decorator's stream config so ``cf.stream(...)`` can
@@ -576,6 +584,61 @@ def suspense(self, timeout: float | None = None) -> Iterator[None]:
         finally:
             _timeout.reset(token)
 
+    def set_checkpoints(self, checkpoints: dict[str, Any] | None) -> None:
+        """Seed the effective checkpoint state from the execute message."""
+        with self._lock:
+            self._checkpoint_wire = dict(checkpoints or {})
+            self._checkpoint_values = {}
+
+    def checkpoint_get(self, name: str) -> Any:
+        """Read a checkpoint value, raising ``KeyError`` if it isn't set.
+
+        Reads are served entirely from local state — the effective checkpoint
+        arrives with the execute message, and this execution's own writes are
+        applied to it — so this never round-trips to the server.
+        """
+        with self._lock:
+            if name in self._checkpoint_values:
+                return self._checkpoint_values[name]
+            if name not in self._checkpoint_wire:
+                raise KeyError(name)
+            wire = self._checkpoint_wire[name]
+
+        # Materialise outside the lock: a blob-backed value is read off disk
+        # here, and the lock is shared with the stream driver threads.
+        value = deserialize_value(wire)
+
+        with self._lock:
+            self._checkpoint_wire.pop(name, None)
+            # A write may have landed while this was materialising, and a
+            # write always wins over the inherited value.
+            if name not in self._checkpoint_values:
+                self._checkpoint_values[name] = value
+            return self._checkpoint_values[name]
+
+    def checkpoint_has(self, name: str) -> bool:
+        with self._lock:
+            return name in self._checkpoint_values or name in self._checkpoint_wire
+
+    def checkpoint_set(self, name: str, value: Any) -> None:
+        with self._lock:
+            self._checkpoint_values[name] = value
+            self._checkpoint_wire.pop(name, None)
+        protocol.send_checkpoint_update(
+            self.execution_id, set_={name: serialize_value(value)}
+        )
+
+    def checkpoint_reset(self, name: str) -> None:
+        with self._lock:
+            self._checkpoint_values.pop(name, None)
+            self._checkpoint_wire.pop(name, None)
+        protocol.send_checkpoint_update(self.execution_id, reset=[name])
+
+    def flush(self) -> None:
+        """Block until buffered state has reached the server."""
+        request_id = protocol.request_flush(self.execution_id)
+        self._wait_response(request_id)
+
     def suspend_execution(
         self, delay: float | dt.timedelta | dt.datetime | None = None
     ) -> None:

adapters/python/coflux/executor.py

@@ -101,6 +101,7 @@ def execute_target(
     arguments: list[dict[str, Any]],
     working_dir: str | None = None,
     streams: dict[str, Any] | None = None,
+    checkpoints: dict[str, Any] | None = None,
 ) -> None:
     """Execute a target with the given arguments.
 
@@ -110,6 +111,10 @@ def execute_target(
     decorator's static config; it's applied both to the auto-registered
     stream for generator-bodied tasks and to ``cf.stream(...)`` calls
     inside the body.
+
+    ``checkpoints`` is the step's effective checkpoint state, resolved by the
+    server. It's kept in protocol form until read, so a target that never
+    touches a checkpoint doesn't pay to deserialise it.
     """
     original_dir = os.getcwd()
     # Start the stdin dispatcher. From here on, all incoming messages flow
@@ -148,6 +153,7 @@ def execute_target(
         effective_streams = _resolve_execute_streams(target_obj, streams)
         if effective_streams is not None or hasattr(target_obj, "definition"):
             ctx.set_default_streams(effective_streams)
+        ctx.set_checkpoints(checkpoints)
         set_context(ctx)
 
         with capture_output(execution_id):
@@ -256,6 +262,7 @@ def run_executor() -> int:
             arguments=params.get("arguments", []),
             working_dir=params.get("working_dir"),
             streams=params.get("streams"),
+            checkpoints=params.get("checkpoints"),
         )
         return 0
 

adapters/python/coflux/protocol.py

@@ -471,6 +471,40 @@ def send_metric(
     get_protocol().send_message("metric", params)
 
 
+def send_checkpoint_update(
+    execution_id: str,
+    set_: dict[str, Any] | None = None,
+    reset: list[str] | None = None,
+) -> None:
+    """Send a checkpoint delta.
+
+    ``set_`` maps names to serialized values; ``reset`` lists names to clear
+    back to their declared default. Delivery is throttled by the worker, so
+    this returning does not mean the delta has reached the server — use
+    ``request_flush`` for that.
+
+    Args:
+        execution_id: The execution recording the delta.
+        set_: Names assigned a new value (serialized).
+        reset: Names to clear.
+    """
+    params: dict[str, Any] = {"execution_id": execution_id}
+    if set_:
+        params["set"] = set_
+    if reset:
+        params["reset"] = reset
+    get_protocol().send_message("checkpoint_update", params)
+
+
+def request_flush(execution_id: str) -> int:
+    """Request a flush of buffered state — checkpoints, metrics and logs.
+
+    The response is withheld until the buffered state has reached the server,
+    which is what makes this a durability boundary rather than a hint.
+    """
+    return get_protocol().send_request("flush", {"execution_id": execution_id})
+
+
 def send_stream_register(
     execution_id: str,
     index: int,

cli/internal/adapter/adapter.go

@@ -246,7 +246,7 @@ func (e *Executor) Send(msg any) error {
 }
 
 // SendExecute sends an execute command to the executor
-func (e *Executor) SendExecute(executionID, module, target string, arguments []Argument, workingDir string, streams *StreamsConfig) error {
+func (e *Executor) SendExecute(executionID, module, target string, arguments []Argument, workingDir string, streams *StreamsConfig, checkpoints map[string]*Value) error {
 	req := ExecuteRequest{
 		Method: "execute",
 		Params: ExecuteRequestParams{
@@ -256,6 +256,7 @@ func (e *Executor) SendExecute(executionID, module, target string, arguments []A
 			Arguments:   arguments,
 			WorkingDir:  workingDir,
 			Streams:     streams,
+			Checkpoints: checkpoints,
 		},
 	}
 	return e.Send(req)

cli/internal/adapter/protocol.go

@@ -80,6 +80,11 @@ type ExecuteRequestParams struct {
 	Arguments   []Argument     `json:"arguments"`
 	WorkingDir  string         `json:"working_dir,omitempty"`
 	Streams     *StreamsConfig `json:"streams,omitempty"`
+	// Checkpoints is the step's effective checkpoint state, resolved by the
+	// server and delivered eagerly so reads never need a round-trip. Values
+	// take the same form as arguments — a blob-backed one has already been
+	// downloaded by the worker and arrives as a local file path.
+	Checkpoints map[string]*Value `json:"checkpoints,omitempty"`
 }
 
 // Argument is the same structure as Value (used for arguments to distinguish context)
@@ -259,6 +264,25 @@ type CancelParams struct {
 	Handles     []SelectHandle `json:"handles"`
 }
 
+// CheckpointUpdateParams for the checkpoint_update notification.
+// Set carries names assigned a new value; Reset carries names cleared back to
+// their declared default. A name appears in at most one of the two — the
+// adapter tracks the net effect, so a set followed by a reset arrives only as
+// a reset.
+type CheckpointUpdateParams struct {
+	ExecutionID string            `json:"execution_id"`
+	Set         map[string]*Value `json:"set,omitempty"`
+	Reset       []string          `json:"reset,omitempty"`
+}
+
+// FlushParams for the flush request. Flushes everything buffered on the
+// execution's behalf — checkpoints, metrics and logs — and only responds once
+// it has reached the server, giving the adapter an explicit durability
+// boundary.
+type FlushParams struct {
+	ExecutionID string `json:"execution_id"`
+}
+
 // RegisterGroupParams for register_group notification
 type RegisterGroupParams struct {
 	ExecutionID string  `json:"execution_id"`