From 158a88e9bde453aeaf16fdd584243456aa286d42 Mon Sep 17 00:00:00 2001 From: Joe Freeman Date: Tue, 1 Sep 2026 12:11:44 +0100 Subject: [PATCH] Add support for checkpoints --- adapters/python/CHANGELOG.md | 5 +- adapters/python/coflux/__init__.py | 3 + adapters/python/coflux/checkpoint.py | 134 ++++ adapters/python/coflux/context.py | 63 ++ adapters/python/coflux/executor.py | 7 + adapters/python/coflux/protocol.py | 34 + cli/internal/adapter/adapter.go | 3 +- cli/internal/adapter/protocol.go | 24 + cli/internal/checkpoint/throttle.go | 184 ++++++ cli/internal/pool/pool.go | 101 ++- cli/internal/worker/worker.go | 105 +++- docs/docs/checkpoints.md | 89 +++ docs/docs/python_reference.md | 44 ++ docs/docs/recurring.md | 2 + docs/docs/suspense.md | 2 + docs/sidebars.ts | 1 + server/CHANGELOG.md | 4 +- server/lib/coflux/handlers/worker.ex | 42 +- server/lib/coflux/orchestration.ex | 7 + .../lib/coflux/orchestration/checkpoints.ex | 316 ++++++++++ server/lib/coflux/orchestration/epoch.ex | 34 +- server/lib/coflux/orchestration/runs.ex | 14 + server/lib/coflux/orchestration/server.ex | 146 ++++- server/lib/coflux/orchestration/workspaces.ex | 18 + server/lib/coflux/topics/run.ex | 34 +- server/priv/migrations/orchestration/5.sql | 51 ++ server/test/coflux/checkpoints_test.exs | 322 ++++++++++ tests/support/executor.py | 30 +- tests/support/protocol.py | 26 + tests/test_checkpoints.py | 582 ++++++++++++++++++ 30 files changed, 2402 insertions(+), 25 deletions(-) create mode 100644 adapters/python/coflux/checkpoint.py create mode 100644 cli/internal/checkpoint/throttle.go create mode 100644 docs/docs/checkpoints.md create mode 100644 server/lib/coflux/orchestration/checkpoints.ex create mode 100644 server/priv/migrations/orchestration/5.sql create mode 100644 server/test/coflux/checkpoints_test.exs create mode 100644 tests/test_checkpoints.py diff --git a/adapters/python/CHANGELOG.md b/adapters/python/CHANGELOG.md index f2b5d5fc..b90221ce 100644 --- a/adapters/python/CHANGELOG.md +++ b/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 diff --git a/adapters/python/coflux/__init__.py b/adapters/python/coflux/__init__.py index ceb30f5e..02d5893b 100644 --- a/adapters/python/coflux/__init__.py +++ b/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", ] diff --git a/adapters/python/coflux/checkpoint.py b/adapters/python/coflux/checkpoint.py new file mode 100644 index 00000000..28d97564 --- /dev/null +++ b/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"] diff --git a/adapters/python/coflux/context.py b/adapters/python/coflux/context.py index a9aaf8d8..a090f30e 100644 --- a/adapters/python/coflux/context.py +++ b/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: diff --git a/adapters/python/coflux/executor.py b/adapters/python/coflux/executor.py index aeb43941..8dcac6fd 100644 --- a/adapters/python/coflux/executor.py +++ b/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 diff --git a/adapters/python/coflux/protocol.py b/adapters/python/coflux/protocol.py index deea2b02..e7dcf802 100644 --- a/adapters/python/coflux/protocol.py +++ b/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, diff --git a/cli/internal/adapter/adapter.go b/cli/internal/adapter/adapter.go index 1430c1df..ef7ab7c9 100644 --- a/cli/internal/adapter/adapter.go +++ b/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) diff --git a/cli/internal/adapter/protocol.go b/cli/internal/adapter/protocol.go index 6cde284a..88bc1789 100644 --- a/cli/internal/adapter/protocol.go +++ b/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"` diff --git a/cli/internal/checkpoint/throttle.go b/cli/internal/checkpoint/throttle.go new file mode 100644 index 00000000..5ec036e7 --- /dev/null +++ b/cli/internal/checkpoint/throttle.go @@ -0,0 +1,184 @@ +// Package checkpoint buffers and coalesces checkpoint writes on their way +// from an adapter to the server. +// +// Unlike metrics, checkpoints are semantic state rather than telemetry: the +// last write before an execution suspends or terminates has to reach the +// server before the successor starts, or the successor resumes from a stale +// value. Buffering therefore comes with synchronous, acknowledged flush +// points rather than a best-effort background drain. +package checkpoint + +import ( + "context" + "sync" + "time" + + "github.com/bitroot/coflux/cli/internal/adapter" +) + +// DefaultInterval is the minimum gap between checkpoint deltas reaching the +// server for a given execution. Writes arriving within a window are coalesced +// (last write wins per name) and delivered when it closes. +// +// Slower than the metric default because a checkpoint value can be +// blob-backed: coalescing here means a value superseded before anything could +// read it is never uploaded at all. +const DefaultInterval = 500 * time.Millisecond + +// Sink delivers a coalesced delta to the server. +type Sink interface { + SetCheckpoints(ctx context.Context, executionID string, set map[string]*adapter.Value, reset []string) error +} + +// entry is the most recent operation buffered for a name. A name is either +// set to a value or reset, never both — a later operation replaces an earlier +// one, so the delta only ever describes the net effect. +type entry struct { + value *adapter.Value + reset bool +} + +type state struct { + pending map[string]entry + lastSent time.Time + timer *time.Timer + + // Held across delivery so a timer-driven flush and an explicit Flush + // can't interleave and land deltas out of order — the later snapshot + // must not be overtaken by the earlier one, or a superseded value + // becomes the stored one. + sending sync.Mutex +} + +// Throttle coalesces per-execution checkpoint deltas. +type Throttle struct { + sink Sink + interval time.Duration + + mu sync.Mutex + execs map[string]*state +} + +func NewThrottle(sink Sink, interval time.Duration) *Throttle { + if interval <= 0 { + interval = DefaultInterval + } + return &Throttle{ + sink: sink, + interval: interval, + execs: make(map[string]*state), + } +} + +// Record buffers a delta. The first write for an execution (or the first +// after a quiet period longer than the interval) is delivered immediately; +// subsequent writes are coalesced and delivered when the window closes. +func (t *Throttle) Record(ctx context.Context, executionID string, set map[string]*adapter.Value, reset []string) error { + t.mu.Lock() + st, existed := t.execs[executionID] + if !existed { + st = &state{pending: make(map[string]entry)} + t.execs[executionID] = st + } + + for name, value := range set { + st.pending[name] = entry{value: value} + } + for _, name := range reset { + st.pending[name] = entry{reset: true} + } + + // Leading edge: deliver immediately if this is the first write or the + // previous one is already older than the window. + immediate := !existed || time.Since(st.lastSent) >= t.interval + if !immediate && st.timer == nil { + remaining := t.interval - time.Since(st.lastSent) + st.timer = time.AfterFunc(remaining, func() { + _ = t.deliver(context.WithoutCancel(ctx), executionID, st) + }) + } + t.mu.Unlock() + + if immediate { + return t.deliver(ctx, executionID, st) + } + return nil +} + +// Flush delivers anything buffered for the execution and returns once the +// server has acknowledged it. Safe to call when nothing is pending. +func (t *Throttle) Flush(ctx context.Context, executionID string) error { + t.mu.Lock() + st, ok := t.execs[executionID] + t.mu.Unlock() + + if !ok { + return nil + } + return t.deliver(ctx, executionID, st) +} + +// Remove drops buffered state for a finished execution. Callers are expected +// to have flushed first — anything still pending here can no longer be read +// by the server, since the execution has terminated. +func (t *Throttle) Remove(executionID string) { + t.mu.Lock() + st, ok := t.execs[executionID] + if ok { + if st.timer != nil { + st.timer.Stop() + st.timer = nil + } + delete(t.execs, executionID) + } + t.mu.Unlock() +} + +// deliver snapshots the pending delta and sends it. +// +// The send lock is taken *before* t.mu, and t.mu is never held while +// acquiring it, so concurrent callers serialise into a consistent order: +// whoever holds the send lock snapshots last and therefore sends last. +func (t *Throttle) deliver(ctx context.Context, executionID string, st *state) error { + st.sending.Lock() + defer st.sending.Unlock() + + t.mu.Lock() + if st.timer != nil { + st.timer.Stop() + st.timer = nil + } + + if len(st.pending) == 0 { + t.mu.Unlock() + return nil + } + + set := make(map[string]*adapter.Value, len(st.pending)) + var reset []string + for name, e := range st.pending { + if e.reset { + reset = append(reset, name) + } else { + set[name] = e.value + } + } + sent := st.pending + st.pending = make(map[string]entry) + st.lastSent = time.Now() + t.mu.Unlock() + + err := t.sink.SetCheckpoints(ctx, executionID, set, reset) + if err != nil { + // Put the delta back so a later write or an explicit flush retries + // it. Names written while this send was in flight are newer and win. + t.mu.Lock() + for name, e := range sent { + if _, superseded := st.pending[name]; !superseded { + st.pending[name] = e + } + } + t.mu.Unlock() + } + return err +} diff --git a/cli/internal/pool/pool.go b/cli/internal/pool/pool.go index bbe316fb..7ef102c0 100644 --- a/cli/internal/pool/pool.go +++ b/cli/internal/pool/pool.go @@ -39,6 +39,16 @@ type ExecutionHandler interface { DefineMetric(ctx context.Context, executionID string, key string, definition map[string]any) error // RecordMetric records a metric data point RecordMetric(ctx context.Context, executionID string, key string, value float64, at *float64) error + // SetCheckpoints records a checkpoint delta. Writes are throttled, so this + // returns before the delta necessarily reaches the server. + SetCheckpoints(ctx context.Context, executionID string, set map[string]*adapter.Value, reset []string) error + // FlushCheckpoints delivers anything buffered for the execution and + // returns once the server has acknowledged it. + FlushCheckpoints(ctx context.Context, executionID string) error + // FlushBuffers delivers everything buffered on the execution's behalf — + // checkpoints, metrics and logs — and returns once the server has + // acknowledged it. + FlushBuffers(ctx context.Context, executionID string) error // ReportResult reports execution completion ReportResult(ctx context.Context, executionID string, result *adapter.Value) error // ReportError reports execution failure @@ -189,7 +199,7 @@ func (p *Pool) spawnExecutor(ctx context.Context) (*adapter.Executor, error) { // timeoutMs, if > 0, enforces a wall-clock timeout on the execution. // streams (if non-nil) is the default stream config — forwarded to the // adapter so generator-bodied tasks and cf.stream(...) calls pick it up. -func (p *Pool) Execute(ctx context.Context, executionID, module, target string, arguments []adapter.Argument, timeoutMs int64, streams *adapter.StreamsConfig) error { +func (p *Pool) Execute(ctx context.Context, executionID, module, target string, arguments []adapter.Argument, timeoutMs int64, streams *adapter.StreamsConfig, checkpoints map[string]*adapter.Value) error { p.mu.Lock() if p.shutdown { p.mu.Unlock() @@ -220,12 +230,12 @@ func (p *Pool) Execute(ctx context.Context, executionID, module, target string, p.wg.Add(1) p.mu.Unlock() - go p.runExecution(ctx, exec, executionID, module, target, arguments, timeoutMs, streams) + go p.runExecution(ctx, exec, executionID, module, target, arguments, timeoutMs, streams, checkpoints) return nil } -func (p *Pool) runExecution(ctx context.Context, exec *adapter.Executor, executionID, module, target string, arguments []adapter.Argument, timeoutMs int64, streams *adapter.StreamsConfig) { +func (p *Pool) runExecution(ctx context.Context, exec *adapter.Executor, executionID, module, target string, arguments []adapter.Argument, timeoutMs int64, streams *adapter.StreamsConfig, checkpoints map[string]*adapter.Value) { defer p.wg.Done() // Create a temporary directory for this execution @@ -241,7 +251,7 @@ func (p *Pool) runExecution(ctx context.Context, exec *adapter.Executor, executi logger := p.logger.With("execution_id", executionID, "module", module, "target", target) // Send execute command - if err := exec.SendExecute(executionID, module, target, arguments, workingDir, streams); err != nil { + if err := exec.SendExecute(executionID, module, target, arguments, workingDir, streams, checkpoints); err != nil { logger.Error("failed to send execute command", "error", err) p.handler.ReportError(ctx, executionID, "internal", err.Error(), "", nil) os.RemoveAll(workingDir) @@ -295,6 +305,11 @@ loop: logger.Debug("executor exited after result", "error", err) } else { logger.Error("failed to receive message", "error", err) + // The adapter died without reporting anything — the case + // checkpoints exist for. Whatever it managed to write is what + // the retry should resume from, so land it before the error + // schedules that retry. + p.flushCheckpoints(ctx, executionID, "crash", logger) p.handler.ReportError(ctx, executionID, "internal", err.Error(), "", nil) } break @@ -331,7 +346,16 @@ loop: case "metric": p.handleMetric(execCtx, executionID, params, logger) - case "submit_execution", "select", "persist_asset", "get_asset", "suspend", "cancel", "download_blob", "upload_blob", "submit_input": + case "checkpoint_update": + // Handled inline rather than dispatched async like the requests + // below, even though the leading-edge write does hit the server: + // deltas are last-write-wins per name, so two goroutines racing + // to merge them could leave the older value as the stored one. + // The server call is a short-lived orchestration call, and the + // throttle means at most one per window per execution. + p.handleCheckpointUpdate(execCtx, executionID, params, logger) + + case "submit_execution", "select", "persist_asset", "get_asset", "suspend", "cancel", "download_blob", "upload_blob", "submit_input", "flush": // Dispatch async: these can block on the server (e.g. a // `select` that waits for a child execution). Blocking the // message loop here would stop us reading the adapter's @@ -365,13 +389,16 @@ loop: default: err := fmt.Errorf("unknown message method: %s", method) logger.Error("unknown message method", "method", method) + p.flushCheckpoints(ctx, executionID, "error", logger) p.handler.ReportError(ctx, executionID, "internal", err.Error(), "", nil) break loop } } - // Report timeout to the server if the execution timed out + // Report timeout to the server if the execution timed out. execCtx is + // already past its deadline, so the flush goes via ctx. if timedOut { + p.flushCheckpoints(ctx, executionID, "timeout", logger) if err := p.handler.ReportTimeout(ctx, executionID); err != nil { logger.Error("failed to report timeout", "error", err) } @@ -387,6 +414,13 @@ loop: // Clean up temp dir os.RemoveAll(workingDir) + // Backstop for the paths that reach here without having reported anything + // — principally a server-side abort, where nothing above ran a flush. Every + // path that does report has already flushed, so this is normally a no-op + // against an empty buffer. Uses ctx rather than execCtx, which may already + // be cancelled by a timeout. + p.flushCheckpoints(ctx, executionID, "termination", logger) + // Remove from busy tracking and notify server that the process has // exited. This frees the concurrency slot on the server side. p.finishExecution(executionID, nil) @@ -451,6 +485,11 @@ func (p *Pool) handleExecutionResult(ctx context.Context, executionID string, pa return } + // A recurrent step's successor is created off the back of this, so a + // trailing write that hasn't reached the server yet would be lost to the + // next iteration. + p.flushCheckpoints(ctx, executionID, "result", logger) + if err := p.handler.ReportResult(ctx, executionID, result.Result); err != nil { logger.Error("failed to report result", "error", err) } @@ -471,11 +510,31 @@ func (p *Pool) handleExecutionError(ctx context.Context, executionID string, par return } + // A failing execution's checkpoint is exactly what its retry needs, and + // the retry is scheduled off this. + p.flushCheckpoints(ctx, executionID, "error", logger) + if err := p.handler.ReportError(ctx, executionID, result.Error.Type, result.Error.Message, result.Error.Traceback, result.Error.Retryable); err != nil { logger.Error("failed to report error", "error", err) } } +// flushCheckpoints lands anything buffered for the execution, and must precede +// every report that puts the execution into a terminal state. +// +// The report is what schedules the successor — retry, recurrence or +// resumption — and the successor reads the checkpoint as it starts, so a +// delta still sitting in the buffer at that point is one the successor may +// miss. Where the report also records a completion (a timeout, say) it's worse +// than a race: the server refuses checkpoint writes against a completed +// execution, so a later flush is rejected outright and the delta is stranded +// in a buffer nothing will drain. +func (p *Pool) flushCheckpoints(ctx context.Context, executionID, before string, logger *slog.Logger) { + if err := p.handler.FlushCheckpoints(ctx, executionID); err != nil { + logger.Error("failed to flush checkpoints", "before", before, "error", err) + } +} + func (p *Pool) handleLog(ctx context.Context, executionID string, params json.RawMessage, logger *slog.Logger) { var logMsg struct { ExecutionID string `json:"execution_id"` @@ -526,6 +585,18 @@ func (p *Pool) handleMetric(ctx context.Context, executionID string, params json } } +func (p *Pool) handleCheckpointUpdate(ctx context.Context, executionID string, params json.RawMessage, logger *slog.Logger) { + var req adapter.CheckpointUpdateParams + if err := json.Unmarshal(params, &req); err != nil { + logger.Error("failed to parse checkpoint_update message", "error", err) + return + } + + if err := p.handler.SetCheckpoints(ctx, req.ExecutionID, req.Set, req.Reset); err != nil { + logger.Error("failed to record checkpoints", "error", err) + } +} + func (p *Pool) handleRegisterGroup(ctx context.Context, executionID string, params json.RawMessage, logger *slog.Logger) { var req adapter.RegisterGroupParams if err := json.Unmarshal(params, &req); err != nil { @@ -726,12 +797,30 @@ func (p *Pool) handleRequest(ctx context.Context, exec *adapter.Executor, method result = map[string]any{"blob_key": blobKey} } + case "flush": + var req adapter.FlushParams + if err := json.Unmarshal(params, &req); err != nil { + errInfo = &adapter.ErrorInfo{Code: "parse_error", Message: err.Error()} + break + } + // Everything the adapter can have buffered, not just checkpoints — + // `cf.flush()` is documented as a boundary for all of it. + if err := p.handler.FlushBuffers(ctx, req.ExecutionID); err != nil { + errInfo = &adapter.ErrorInfo{Code: "flush_error", Message: err.Error()} + } else { + result = map[string]any{} + } + case "suspend": var req adapter.SuspendParams if err := json.Unmarshal(params, &req); err != nil { errInfo = &adapter.ErrorInfo{Code: "parse_error", Message: err.Error()} break } + // The successor may be scheduled the moment the server records the + // suspension, so anything still buffered has to land first — otherwise + // it resumes from a stale checkpoint. + p.flushCheckpoints(ctx, req.ExecutionID, "suspend", logger) if err := p.handler.Suspend(ctx, req.ExecutionID, req.ExecuteAfter); err != nil { errInfo = &adapter.ErrorInfo{Code: "suspend_error", Message: err.Error()} } else { diff --git a/cli/internal/worker/worker.go b/cli/internal/worker/worker.go index cf01653c..99393f6c 100644 --- a/cli/internal/worker/worker.go +++ b/cli/internal/worker/worker.go @@ -20,6 +20,7 @@ import ( "github.com/bitroot/coflux/cli/internal/adapter" "github.com/bitroot/coflux/cli/internal/api" "github.com/bitroot/coflux/cli/internal/blob" + "github.com/bitroot/coflux/cli/internal/checkpoint" "github.com/bitroot/coflux/cli/internal/config" logstore "github.com/bitroot/coflux/cli/internal/log" "github.com/bitroot/coflux/cli/internal/metric" @@ -58,6 +59,7 @@ type Worker struct { metrics metric.Store tracker *metric.Tracker throttle *metric.Throttle + checkpoints *checkpoint.Throttle connMu sync.RWMutex conn *api.Connection @@ -326,6 +328,12 @@ func (w *Worker) Run(ctx context.Context, modules []string, register bool) error w.tracker = metric.NewTracker(w.logger) defer func() { _ = w.metrics.Close() }() + // Checkpoint writes go over the worker connection rather than the metrics + // HTTP path: they're semantic state, and the pool needs acknowledged + // flushes at suspend / result / termination rather than a best-effort + // background drain. + w.checkpoints = checkpoint.NewThrottle(checkpointSink{w}, checkpoint.DefaultInterval) + // Determine pool size (default to CPU count + 4) poolSize := w.cfg.Worker.Concurrency if poolSize <= 0 { @@ -580,6 +588,24 @@ func (w *Worker) handleExecute(params []any) error { } } + // Optional checkpoints (9th param). The step's effective checkpoint + // state, resolved server-side and delivered eagerly so the adapter can + // answer reads without a round-trip. Values are in the same wire form as + // arguments. + var checkpoints map[string]*adapter.Value + if len(params) > 8 && params[8] != nil { + if m, ok := params[8].(map[string]any); ok && len(m) > 0 { + checkpoints = make(map[string]*adapter.Value, len(m)) + for name, raw := range m { + value, err := w.convertValueFromServer(raw) + if err != nil { + return fmt.Errorf("checkpoint %q: %w", name, err) + } + checkpoints[name] = value + } + } + } + w.logger.Debug("executing", "execution_id", executionID, "module", moduleName, "target", targetName, "run_id", runID, "timeout_ms", timeoutMs) // Track execution @@ -619,7 +645,7 @@ func (w *Worker) handleExecute(params []any) error { w.mu.Unlock() // Execute on pool - if err := w.pool.Execute(context.Background(), executionID, moduleName, targetName, args, timeoutMs, streams); err != nil { + if err := w.pool.Execute(context.Background(), executionID, moduleName, targetName, args, timeoutMs, streams, checkpoints); err != nil { w.logger.Error("failed to execute", "error", err, "run_id", runID) w.ReportError(context.Background(), executionID, "internal", err.Error(), "", nil) } @@ -1307,6 +1333,80 @@ func (w *Worker) Suspend(ctx context.Context, executionID string, executeAfter * return conn.Notify("suspend", executionID, executeAfter) } +// SetCheckpoints buffers a checkpoint delta. Delivery is throttled, so this +// returns before the delta necessarily reaches the server — the pool calls +// FlushCheckpoints at the points where that matters. +func (w *Worker) SetCheckpoints(ctx context.Context, executionID string, set map[string]*adapter.Value, reset []string) error { + return w.checkpoints.Record(ctx, executionID, set, reset) +} + +// FlushCheckpoints delivers anything buffered for the execution, returning +// once the server has acknowledged it. +func (w *Worker) FlushCheckpoints(ctx context.Context, executionID string) error { + return w.checkpoints.Flush(ctx, executionID) +} + +// FlushBuffers delivers everything buffered on the execution's behalf — +// checkpoints, metrics and logs — returning once the server has acknowledged +// it. This is what `cf.flush()` reaches, so it has to cover every buffer the +// adapter can fill, not just the one that motivated the call. +// +// All three are attempted regardless of failures: a caller flushing before a +// side effect wants as much delivered as possible, and reporting only the +// first failure would hide the rest. +func (w *Worker) FlushBuffers(ctx context.Context, executionID string) error { + var errs []error + + if err := w.checkpoints.Flush(ctx, executionID); err != nil { + errs = append(errs, fmt.Errorf("checkpoints: %w", err)) + } + + // Metrics and logs are batched per worker rather than per execution, so + // these deliver a superset — harmless, and cheaper than threading + // per-execution flushing through both stores. + if err := w.metrics.Flush(); err != nil { + errs = append(errs, fmt.Errorf("metrics: %w", err)) + } + + if err := w.logs.Flush(); err != nil { + errs = append(errs, fmt.Errorf("logs: %w", err)) + } + + return errors.Join(errs...) +} + +// checkpointSink is the throttle's outbound half. Kept separate from +// Worker.SetCheckpoints — which is the inbound half, feeding the throttle — +// so the two directions don't share a method name. +type checkpointSink struct{ w *Worker } + +func (s checkpointSink) SetCheckpoints(ctx context.Context, executionID string, set map[string]*adapter.Value, reset []string) error { + conn, err := s.w.requireConn() + if err != nil { + return err + } + + // Values are composed here rather than in the adapter, so a value + // superseded within a throttle window is never uploaded to the blob store. + composed := make(map[string]any, len(set)) + for name, value := range set { + serverValue, err := s.w.convertValueToServerFormat(value) + if err != nil { + return fmt.Errorf("failed to convert checkpoint %q: %w", name, err) + } + composed[name] = serverValue + } + + if reset == nil { + reset = []string{} + } + + // A request rather than a notification: the response is what makes a + // flush an actual barrier. + _, err = conn.Request(ctx, "checkpoint_update", executionID, composed, reset) + return err +} + func (w *Worker) DownloadBlob(ctx context.Context, executionID, blobKey, targetPath string) error { // Download blob to the target path return w.blobs.DownloadTo(blobKey, targetPath) @@ -1942,6 +2042,9 @@ func (w *Worker) NotifyTerminated(ctx context.Context, executionID string) error // Clean up metric tracking for this execution w.tracker.UnregisterExecution(executionID) w.throttle.RemoveExecution(executionID) + // The pool flushes checkpoints before getting here, so anything still + // buffered couldn't be read anyway — the execution has terminated. + w.checkpoints.Remove(executionID) w.mu.Lock() state, ok := w.executions[executionID] diff --git a/docs/docs/checkpoints.md b/docs/docs/checkpoints.md new file mode 100644 index 00000000..d8533dc7 --- /dev/null +++ b/docs/docs/checkpoints.md @@ -0,0 +1,89 @@ +# Checkpoints + +A checkpoint is a named value that survives across executions of a step. When a step [suspends](./suspense.md), [retries](./retries.md), [recurs](./recurring.md), or is re-run from Studio, the code starts again from the top — a checkpoint is how it picks up where it left off. + +```python +import coflux as cf + +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) +``` + +Each iteration reads what the previous one wrote, so the poller only ever fetches what it hasn't seen. + +## Checkpoints vs memoizing + +These solve overlapping problems, and it's worth being deliberate about which you reach for. + +[Memoizing](./memoizing.md) is the right tool when the state you want to keep is *the result of some work*. A memoized child task runs once and its result is reused by every subsequent attempt, so the work isn't repeated. + +Checkpoints are for state that isn't naturally a task result — a cursor, a page token, a partially-built accumulator. There's no work to memoize, just a value to carry forward. + +## Reading and writing + +Create a `Checkpoint` with a name, and optionally a default: + +```python +cursor = cf.Checkpoint("cursor", default=0) +``` + +The name identifies storage scoped to the step, so declaring the handle at module level is fine — it isn't module state. + +```python +cursor.get() # the current value, or the default if unset +cursor.set(value) # replace the value +cursor.reset() # clear it, so get() returns the default again +cursor.is_set() # whether it has a value +``` + +Reads are served locally: the effective state arrives with the execution, and an execution always sees its own writes. Nothing round-trips to the server. + +`set(None)` and `reset()` are different. `set(None)` stores `None`, and `get()` returns `None`. `reset()` removes the checkpoint, so `get()` falls back to the declared default. + +A checkpoint is typed by what `get()` returns. That's inferred from the default when there is one, so `cursor` above is a `Checkpoint[int]`. Without a default, `get()` can return `None`, so declare that: + +```python +cursor = cf.Checkpoint[int | None]("cursor") +``` + +Like argument and result types, this only informs type checkers — nothing is enforced at runtime. + +A checkpoint handle can't be passed to another execution — it names storage belonging to a particular step, so declare it in the target that uses it. + +## Durability + +Writes are throttled and delivered in the background, so a crash can lose up to a fraction of a second of them. **Treat a checkpoint as at-least-once**: the code after a read has to be safe to run again from the value it read. + +In the polling example above, that means a crash may cause some orders to be fetched twice — which is fine, because `process_order` is submitted with the same arguments and can be memoized. + +Whatever has been written when an execution suspends, returns, or fails is always delivered before the next attempt starts. You only need to think about this for a side effect *within* an execution that must not be repeated. `cf.flush()` gives an explicit boundary: + +```python +cursor.set(next_cursor) +cf.flush() +send_notification() +``` + +`cf.flush()` returns once the server has acknowledged the write. + +## Scope + +Checkpoints are scoped to a step within a [workspace](./concepts.md). Reads fall back through the workspace's bases, so re-running a step in a derived workspace reads the base's real state — useful for debugging against production values — while writes only ever land in the workspace doing the writing. A derived workspace can't corrupt the state its base is using, and once it has written its own value it reads that instead. + +A checkpoint belongs to the step that actually executes. A step resolved from the [cache](./caching.md) or by [memoization](./memoizing.md) never runs, so it never sees one, and a checkpoint is never part of a cache, memo or defer key. + +:::warning +Checkpoints are scoped to a step within a run, so a recurring workflow keeps its checkpoints for as long as its run is alive — across every recurrence, retry and suspension. But if recurrence stops (retries are exhausted, the task returns a non-`None` value, or the run is cancelled), submitting the workflow again creates a new run with a fresh step, which starts from the declared defaults. +::: + +## Size + +A checkpoint value is serialized like any other, so a large one is stored as a [blob](./blobs.md) and downloaded at the start of every execution of the step. That's cheap for a cursor and expensive for a large dataframe — prefer keeping checkpoints small, and use an [asset](./assets.md) for anything substantial. diff --git a/docs/docs/python_reference.md b/docs/docs/python_reference.md index 18f94a35..4b0471f2 100644 --- a/docs/docs/python_reference.md +++ b/docs/docs/python_reference.md @@ -284,6 +284,46 @@ cf.Retries( | `backoff` | `tuple` | `(1, 60)` | Backoff range (min, max) in seconds | | `when` | type, tuple, callable, or `None` | `None` | Exception filter (`None` = retry on any error) | +## Checkpoints + +State that survives across executions of a step — retries, suspensions, recurrences and re-runs. See [checkpoints](./checkpoints.md). + +### `Checkpoint` + +```python +cf.Checkpoint( + name: str, + *, + default: T, +) +cf.Checkpoint[T | None]( + name: str, +) +``` + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `name` | `str` | required | Checkpoint name, unique within the step | +| `default` | `T` | — | Returned when the checkpoint is unset or has been reset. Omit it and `get()` may return `None` | + +`T` is the type `get()` returns. It's inferred from `default` when one is given, so `cf.Checkpoint("cursor", default=0)` is a `Checkpoint[int]`. Without a default the checkpoint can read as `None`, so spell the type out: `cf.Checkpoint[int | None]("cursor")`. Types only inform type checkers — nothing is enforced at runtime. + +#### `checkpoint.get() -> T` + +The current value, or the declared default if it isn't set. A checkpoint explicitly set to `None` reads back as `None`. + +#### `checkpoint.set(value) -> None` + +Sets the value, replacing anything already there. + +#### `checkpoint.reset() -> None` + +Clears the checkpoint, so `get()` returns the declared default again. Distinct from `set(None)`, which stores `None`. + +#### `checkpoint.is_set() -> bool` + +Whether the checkpoint has a value (including an explicit `None`). + ## Metrics Record numeric values from executions, streamed in real-time and rendered as charts in Studio. See [metrics](./metrics.md). @@ -366,6 +406,10 @@ Context manager that sets a timeout on `.result()` calls within its scope. If th Explicitly suspends the current execution. It will be re-run after the specified delay. `delay` can be `float` (seconds), `timedelta`, or `datetime`. +### `flush()` + +Blocks until buffered state (checkpoints, metrics, logs) has reached the server. Not needed before suspending, returning or raising — those flush automatically. See [checkpoints](./checkpoints.md). + ### `select(handles, *, cancel_remaining=False)` Wait for the first of one or more handles (`Execution` and/or `Input`) to resolve. Returns `(winner, remaining)` — call `winner.result()` to get the value (or to raise the exception that resolved it). Picks up its timeout from any enclosing `cf.suspense(timeout=...)` scope; raises `TimeoutError` if the wait expires. See [select](./select.md). diff --git a/docs/docs/recurring.md b/docs/docs/recurring.md index e9fd6e05..b80a68c0 100644 --- a/docs/docs/recurring.md +++ b/docs/docs/recurring.md @@ -12,6 +12,8 @@ def poll_for_updates(): The task recurs as long as it returns `None`. Returning any other value completes the cycle and stops recurrence. The run can also be stopped by cancelling it, or if an error occurs (without a successful retry). +Each iteration is a fresh execution starting from the top, so anything that needs to carry forward between them — a cursor, say — belongs in a [checkpoint](./checkpoints.md). + ## Delay By default, recurring tasks restart immediately. Use `delay` to wait between executions: diff --git a/docs/docs/suspense.md b/docs/docs/suspense.md index 14392b0e..6f198126 100644 --- a/docs/docs/suspense.md +++ b/docs/docs/suspense.md @@ -4,6 +4,8 @@ Suspense is a way of putting a task to sleep — the current execution will be s The suspense can be either _explicit_ or _implicit_. In either case, it's important that the code up to the point of suspense is safe to re-execute — i.e., any side-effects need to be idempotent. (An easy way to achieve this is to ensure that any tasks called by the execution are [memoized](/memoizing).) +State that needs to survive the suspension, but isn't the result of a task, can be kept in a [checkpoint](/checkpoints). + Suspense is useful as a way of freeing up resources used by a waiting execution. ## Explicit suspense diff --git a/docs/sidebars.ts b/docs/sidebars.ts index 412e9d39..6c65a353 100644 --- a/docs/sidebars.ts +++ b/docs/sidebars.ts @@ -38,6 +38,7 @@ const sidebars: SidebarsConfig = { "memoizing", "deferring", "suspense", + "checkpoints", "select", ], }, diff --git a/server/CHANGELOG.md b/server/CHANGELOG.md index 7b971541..66544efe 100644 --- a/server/CHANGELOG.md +++ b/server/CHANGELOG.md @@ -1,6 +1,8 @@ ## 0.12.0 -No changes. +Enhancements: + +- Adds checkpoints — named values scoped to a step within a workspace, carried across retries, suspensions, recurrences and re-runs. Reads fall back through the workspace's bases; writes stay in the writing workspace. History is compacted to the effective state at epoch rotation. ## 0.11.0 diff --git a/server/lib/coflux/handlers/worker.ex b/server/lib/coflux/handlers/worker.ex index 3e5fc0a2..8550f395 100644 --- a/server/lib/coflux/handlers/worker.ex +++ b/server/lib/coflux/handlers/worker.ex @@ -196,6 +196,39 @@ defmodule Coflux.Handlers.Worker do {[], state} + # Shaped as a request rather than a notification so the worker gets a + # definite server ack. Writes are throttled worker-side, so the + # round-trip is cheap — and it makes "checkpoints flushed before the + # execution suspends or terminates" an awaited barrier rather than an + # assumption about message ordering. + "checkpoint_update" -> + [execution_id, set, reset] = message["params"] + + if is_recognised_execution?(execution_id, state) do + set = Map.new(set, fn {name, value} -> {name, parse_value(value)} end) + + case Orchestration.set_checkpoints( + state.project_id, + execution_id, + set, + reset + ) do + :ok -> + {[success_message(message["id"], nil)], state} + + # The execution has already been finalised — its writes can no + # longer be read by anything, so this is reported rather than + # retried. + {:error, :completed} -> + {[error_message(message["id"], "execution_completed")], state} + + {:error, :not_found} -> + {[{:close, 4000, "execution_invalid"}], nil} + end + else + {[{:close, 4000, "execution_invalid"}], nil} + end + "define_metric" -> [execution_id, key, definition] = message["params"] @@ -702,11 +735,15 @@ defmodule Coflux.Handlers.Worker do def websocket_info( {:execute, execution_external_id, module, target, arguments, run_id, - workspace_external_id, timeout, streams}, + workspace_external_id, timeout, streams, checkpoints}, state ) do arguments = Enum.map(arguments, &compose_value/1) + # Checkpoints ride along in the same wire form as arguments, and follow + # the same path from here on. + checkpoints = Map.new(checkpoints, fn {name, value} -> {name, compose_value(value)} end) + state = Map.update!(state, :execution_ids, &MapSet.put(&1, execution_external_id)) {[ @@ -718,7 +755,8 @@ defmodule Coflux.Handlers.Worker do run_id, workspace_external_id, timeout, - compose_streams(streams) + compose_streams(streams), + checkpoints ]) ], state} end diff --git a/server/lib/coflux/orchestration.ex b/server/lib/coflux/orchestration.ex index faa1225d..6ad74d22 100644 --- a/server/lib/coflux/orchestration.ex +++ b/server/lib/coflux/orchestration.ex @@ -181,6 +181,13 @@ defmodule Coflux.Orchestration do call_server(project_id, {:record_result, execution_id, result}) end + # Applies a checkpoint delta recorded by an execution. `set` is a map of + # name => value; `reset` is a list of names to clear. Both are applied + # together, on top of the previous effective state. + def set_checkpoints(project_id, execution_id, set, reset) do + call_server(project_id, {:set_checkpoints, execution_id, set, reset}) + end + # Stream producer messages — worker registers a stream, appends items, # and closes the stream. `index` identifies the stream within its # producer execution; `sequence` identifies an item within the stream. diff --git a/server/lib/coflux/orchestration/checkpoints.ex b/server/lib/coflux/orchestration/checkpoints.ex new file mode 100644 index 00000000..e6cfdfe3 --- /dev/null +++ b/server/lib/coflux/orchestration/checkpoints.ex @@ -0,0 +1,316 @@ +defmodule Coflux.Orchestration.Checkpoints do + @moduledoc """ + Storage for step-scoped durable state ("checkpoints"). + + A checkpoint is a named value that survives across executions of a step — + retries, suspends, recurrences and manual re-runs. It exists for the state + an execution needs to carry forward but that isn't naturally a task result + (a cursor, a page token, an accumulator); anything that *is* a task result + should be memoised instead. + + Scope is `(step, workspace)`. Reads walk the workspace chain nearest-first + and take the latest attempt that touched checkpoints in the *nearest* + workspace that has any — a workspace that has written its own state never + falls back to its base. Writes always land on the writing execution's own + attempt, so inheritance is read-only and one-directional: a re-run in a + descendant workspace reads the base's state without being able to corrupt + it. + + Invariants: + + * Each execution's row-set is a complete snapshot of the effective + checkpoint, not a delta. `apply_delta/7` materialises the previous + snapshot on an execution's first write, then applies the delta, in one + transaction. + * A `NULL` `value_id` is a tombstone — the name was explicitly reset at + that attempt. Distinct from a stored null *value*, which is an ordinary + `values_` row. + * `reset` always writes a tombstone, even for a name that was never set. + That keeps an execution's row-set non-empty whenever it touched + checkpoints, so "reset everything" can't be mistaken for "wrote + nothing" and resurrect the pre-reset state from an earlier attempt. + * Rows are updated in place within an execution — a deliberate exception + to the append-only convention, documented in `5.sql`. Nothing can + observe an intra-execution overwrite: the execution reads its own + writes from memory, and the worker-side throttle discards intermediate + values before they reach the server. + * Attempts are globally ordered per step (`executions` has + `UNIQUE (step_id, attempt)`), so ordering by attempt is well-defined + even across workspaces. Ordering by attempt rather than wall-clock is + what makes a stale worker harmless: its writes land on its own, older + attempt and are never read. + + History across attempts is preserved by the `(execution_id, name)` key and + compacted to the effective snapshot at epoch rotation. + """ + + import Coflux.Store + + alias Coflux.Orchestration.Values + + @doc """ + Resolves the effective checkpoint for a step. + + `workspace_chain` is a list of workspace ids ordered nearest-first (the + execution's own workspace, then its bases). `before_attempt` bounds the + lookup to attempts strictly below it — pass the reading execution's own + attempt so it sees what it started with, or `nil` for the current state. + + Returns `{:ok, %{name => value}}` with tombstones dropped and values + resolved. + """ + def get_effective(db, step_id, workspace_chain, before_attempt \\ nil) do + {:ok, rows} = get_snapshot_rows(db, step_id, workspace_chain, before_attempt) + + entries = + rows + |> Enum.reject(fn {_name, value_id} -> is_nil(value_id) end) + |> Map.new(fn {name, value_id} -> + {:ok, value} = Values.get_value_by_id(db, value_id) + {name, value} + end) + + {:ok, entries} + end + + @doc """ + Applies a checkpoint delta recorded by `execution_id`. + + `set` is a map of `name => value` (values in the same `{:raw, ...}` / + `{:blob, ...}` form as step arguments); `reset` is a list of names. Both are + applied in a single transaction, preceded by carry-forward of the previous + snapshot if this is the execution's first checkpoint write. Resets are + applied before sets, so a name appearing in both ends up set. + + Returns `{:ok, timestamp}`. + """ + def apply_delta(db, execution_id, step_id, workspace_chain, attempt, set, reset) do + with_transaction(db, fn -> + now = current_timestamp() + + :ok = ensure_snapshot(db, execution_id, step_id, workspace_chain, attempt, now) + + Enum.each(reset, fn name -> + :ok = put(db, execution_id, name, nil, now) + end) + + Enum.each(set, fn {name, value} -> + {:ok, value_id} = Values.get_or_create_value(db, value) + :ok = put(db, execution_id, name, value_id, now) + end) + + {:ok, now} + end) + end + + @doc """ + Every checkpoint row recorded by `execution_id`, tombstones included. + + Returns `{:ok, [{name, value_id | nil, updated_at}]}` in name order. Used by + epoch rotation and by the run topic. + """ + def get_rows_for_execution(db, execution_id) do + query( + db, + """ + SELECT name, value_id, updated_at + FROM checkpoints + WHERE execution_id = ?1 + ORDER BY name + """, + {execution_id} + ) + end + + @doc """ + The checkpoint an execution ended up holding, with values resolved and + tombstones dropped. + + Since each execution's row-set is a complete snapshot, this is the whole + effective checkpoint as of that attempt — not just what it changed. An + execution that never wrote has no rows and returns an empty map. + + Returns `{:ok, %{name => value}}`. + """ + def get_effective_for_execution(db, execution_id) do + {:ok, rows} = get_rows_for_execution(db, execution_id) + {:ok, resolve_execution_rows(db, rows)} + end + + @doc """ + What an execution started from, and what it ended up holding. + + The first element is the effective checkpoint as of the execution's own + attempt — what it was handed when it started. The second is its own + snapshot, or the first again when it has no rows: an execution that never + wrote left the state exactly as it found it, which isn't the same as + holding nothing. + + Returns `{:ok, {before, after}}`, both `%{name => value}`. + """ + def get_execution_snapshots(db, execution_id, step_id, workspace_chain, attempt) do + {:ok, before} = get_effective(db, step_id, workspace_chain, attempt) + {:ok, rows} = get_rows_for_execution(db, execution_id) + + # Only an empty row-set means "didn't write" — a row-set that resolves to + # nothing is an execution that reset everything, which is a real change. + after_ = if Enum.empty?(rows), do: before, else: resolve_execution_rows(db, rows) + + {:ok, {before, after_}} + end + + defp resolve_execution_rows(db, rows) do + rows + |> Enum.reject(fn {_name, value_id, _updated_at} -> is_nil(value_id) end) + |> Map.new(fn {name, value_id, _updated_at} -> + {:ok, value} = Values.get_value_by_id(db, value_id) + {name, value} + end) + end + + @doc """ + The execution ids holding the effective snapshot for `step_id` — one per + workspace that has any checkpoint rows. + + Epoch rotation copies only these, discarding the rest of the history. + """ + def get_snapshot_execution_ids(db, step_id) do + case query( + db, + """ + SELECT e.id + FROM executions AS e + WHERE e.step_id = ?1 + AND EXISTS (SELECT 1 FROM checkpoints AS c WHERE c.execution_id = e.id) + AND e.attempt = ( + SELECT MAX(e2.attempt) + FROM executions AS e2 + WHERE e2.step_id = e.step_id + AND e2.workspace_id = e.workspace_id + AND EXISTS (SELECT 1 FROM checkpoints AS c2 WHERE c2.execution_id = e2.id) + ) + """, + {step_id} + ) do + {:ok, rows} -> + {:ok, Enum.map(rows, fn {id} -> id end)} + end + end + + # The raw row-set of the effective snapshot — tombstones included, values + # unresolved. Carry-forward needs the tombstones (so a reset stays reset), + # `get_effective/4` drops them. + defp get_snapshot_rows(db, step_id, workspace_chain, before_attempt) do + location = + Enum.find_value(workspace_chain, fn workspace_id -> + case get_latest_attempt(db, step_id, workspace_id, before_attempt) do + {:ok, nil} -> nil + {:ok, attempt} -> {workspace_id, attempt} + end + end) + + case location do + nil -> {:ok, []} + {workspace_id, attempt} -> get_attempt_rows(db, step_id, workspace_id, attempt) + end + end + + # The highest attempt of `step_id` in `workspace_id` that recorded any + # checkpoint rows, below `before_attempt` if given. + # + # Written as ORDER BY ... LIMIT 1 rather than MAX(...) so SQLite walks + # idx_executions_step_workspace descending and stops at the first hit, + # rather than scanning every attempt of the step. + defp get_latest_attempt(db, step_id, workspace_id, before_attempt) do + {sql, args} = + if before_attempt do + {""" + SELECT e.attempt + FROM executions AS e + WHERE e.step_id = ?1 AND e.workspace_id = ?2 AND e.attempt < ?3 + AND EXISTS (SELECT 1 FROM checkpoints AS c WHERE c.execution_id = e.id) + ORDER BY e.attempt DESC + LIMIT 1 + """, {step_id, workspace_id, before_attempt}} + else + {""" + SELECT e.attempt + FROM executions AS e + WHERE e.step_id = ?1 AND e.workspace_id = ?2 + AND EXISTS (SELECT 1 FROM checkpoints AS c WHERE c.execution_id = e.id) + ORDER BY e.attempt DESC + LIMIT 1 + """, {step_id, workspace_id}} + end + + case query_one(db, sql, args) do + {:ok, nil} -> {:ok, nil} + {:ok, {attempt}} -> {:ok, attempt} + end + end + + defp get_attempt_rows(db, step_id, workspace_id, attempt) do + query( + db, + """ + SELECT c.name, c.value_id + FROM checkpoints AS c + INNER JOIN executions AS e ON e.id = c.execution_id + WHERE e.step_id = ?1 AND e.workspace_id = ?2 AND e.attempt = ?3 + """, + {step_id, workspace_id, attempt} + ) + end + + # Materialise the previous effective snapshot into this execution, if it + # hasn't written anything yet. Bounded by the execution's own attempt so the + # stored snapshot matches what the execution actually started with. + defp ensure_snapshot(db, execution_id, step_id, workspace_chain, attempt, now) do + case has_rows?(db, execution_id) do + {:ok, true} -> + :ok + + {:ok, false} -> + {:ok, rows} = get_snapshot_rows(db, step_id, workspace_chain, attempt) + + Enum.each(rows, fn {name, value_id} -> + :ok = put(db, execution_id, name, value_id, now) + end) + + :ok + end + end + + defp has_rows?(db, execution_id) do + case query_one( + db, + "SELECT 1 FROM checkpoints WHERE execution_id = ?1 LIMIT 1", + {execution_id} + ) do + {:ok, nil} -> {:ok, false} + {:ok, {1}} -> {:ok, true} + end + end + + defp put(db, execution_id, name, value_id, now) do + {:ok, _} = + insert_one( + db, + :checkpoints, + %{ + execution_id: execution_id, + name: name, + value_id: value_id, + updated_at: now + }, + on_conflict: + "(execution_id, name) DO UPDATE SET value_id = excluded.value_id, updated_at = excluded.updated_at" + ) + + :ok + end + + defp current_timestamp() do + System.os_time(:millisecond) + end +end diff --git a/server/lib/coflux/orchestration/epoch.ex b/server/lib/coflux/orchestration/epoch.ex index cb137d7d..a9092186 100644 --- a/server/lib/coflux/orchestration/epoch.ex +++ b/server/lib/coflux/orchestration/epoch.ex @@ -3,7 +3,7 @@ defmodule Coflux.Orchestration.Epoch do Copies data between epoch databases during rotation and copy-on-reference. """ - alias Coflux.Orchestration.Runs + alias Coflux.Orchestration.{Checkpoints, Runs} import Coflux.Store @@ -231,6 +231,8 @@ defmodule Coflux.Orchestration.Epoch do Map.put(acc, old_exec_id, new_exec_id) end) + copy_step_checkpoints(source_db, target_db, old_step_id, exec_acc) + {Map.put(step_acc, old_step_id, new_step_id), exec_acc} end) @@ -473,6 +475,36 @@ defmodule Coflux.Orchestration.Epoch do # Per-execution copy helpers + # Rotation is the compaction point for checkpoint history: only the + # effective snapshot is carried over — the latest checkpoint-bearing attempt + # per workspace — and everything older is discarded. + # + # Tombstones are copied along with values. Discarding the history they were + # masking isn't enough to make them redundant: a tombstone in a derived + # workspace also masks the *base* workspace's value, which survives rotation + # in its own right. Dropping it would empty that workspace's row-set, and an + # empty row-set falls back to the base — resurrecting the value the reset + # removed. + defp copy_step_checkpoints(source_db, target_db, old_step_id, execution_ids) do + {:ok, snapshot_execution_ids} = + Checkpoints.get_snapshot_execution_ids(source_db, old_step_id) + + Enum.each(snapshot_execution_ids, fn old_exec_id -> + new_exec_id = Map.fetch!(execution_ids, old_exec_id) + {:ok, rows} = Checkpoints.get_rows_for_execution(source_db, old_exec_id) + + Enum.each(rows, fn {name, value_id, updated_at} -> + {:ok, _} = + insert_one(target_db, :checkpoints, %{ + execution_id: new_exec_id, + name: name, + value_id: value_id && ensure_value(source_db, target_db, value_id), + updated_at: updated_at + }) + end) + end) + end + defp copy_execution_groups(source_db, target_db, old_exec_id, new_exec_id) do {:ok, groups} = query(source_db, "SELECT group_id, name FROM groups WHERE execution_id = ?1", {old_exec_id}) diff --git a/server/lib/coflux/orchestration/runs.ex b/server/lib/coflux/orchestration/runs.ex index 9c7f487d..102315c5 100644 --- a/server/lib/coflux/orchestration/runs.ex +++ b/server/lib/coflux/orchestration/runs.ex @@ -148,6 +148,20 @@ defmodule Coflux.Orchestration.Runs do end end + # The step, workspace and attempt an execution belongs to. Checkpoint reads + # and writes need all three: the step and workspace define the scope, and + # the attempt orders it against the step's other executions. + def get_execution_location(db, execution_id) do + case query_one( + db, + "SELECT step_id, workspace_id, attempt FROM executions WHERE id = ?1", + {execution_id} + ) do + {:ok, {step_id, workspace_id, attempt}} -> {:ok, {step_id, workspace_id, attempt}} + {:ok, nil} -> {:error, :not_found} + end + end + def get_run_id_for_execution(db, execution_id) do case query_one( db, diff --git a/server/lib/coflux/orchestration/server.ex b/server/lib/coflux/orchestration/server.ex index 494074e3..fdec9545 100644 --- a/server/lib/coflux/orchestration/server.ex +++ b/server/lib/coflux/orchestration/server.ex @@ -11,6 +11,7 @@ defmodule Coflux.Orchestration.Server do Runs, Results, Streams, + Checkpoints, Assets, Inputs, Values, @@ -1442,6 +1443,14 @@ defmodule Coflux.Orchestration.Server do recurrent = Keyword.get(opts, :recurrent, false) + {:ok, checkpoints} = + Checkpoints.get_effective( + state.db, + step_id, + get_workspace_chain(state, workspace_id), + attempt + ) + state |> notify_listeners( {:run, run.external_id}, @@ -1465,7 +1474,7 @@ defmodule Coflux.Orchestration.Server do |> notify_listeners( {:run, run.external_id}, {:execution, step_number, attempt, execution_external_id, ws_ext_id, created_at, - execute_after, %{}, nil} + execute_after, %{}, nil, enrich_checkpoints(checkpoints, state.db)} ) else state @@ -1867,6 +1876,63 @@ defmodule Coflux.Orchestration.Server do end end + def handle_call({:set_checkpoints, execution_external_id, set, reset}, _from, state) do + case Map.fetch(state.execution_ids, execution_external_id) do + {:ok, execution_id} -> + # Reject writes from an execution the server has already finalised. A + # worker that was declared abandoned while still alive flushes its + # buffer on reconnect; per-execution keying already makes those writes + # unreadable (they land on an older attempt), but there's no reason to + # mutate a terminated execution's history. + case Results.has_completion?(state.db, execution_id) do + {:ok, true} -> + {:reply, {:error, :completed}, state} + + {:ok, false} -> + {:ok, {step_id, workspace_id, attempt}} = + Runs.get_execution_location(state.db, execution_id) + + set = Map.new(set, fn {name, value} -> {name, normalize_value(value)} end) + + {:ok, _updated_at} = + Checkpoints.apply_delta( + state.db, + execution_id, + step_id, + get_workspace_chain(state, workspace_id), + attempt, + set, + reset + ) + + # Push the whole resolved snapshot rather than the delta — it's + # what Studio renders, and carry-forward means the first write of + # an execution materialises names it never touched. Only the + # "after" side moves; what the execution started from is fixed. + {:ok, checkpoints} = + Checkpoints.get_effective_for_execution(state.db, execution_id) + + checkpoints = enrich_checkpoints(checkpoints, state.db) + + {:ok, {run_external_id, _step_number, _attempt}} = + Runs.get_execution_key(state.db, execution_id) + + state = + state + |> notify_listeners( + {:run, run_external_id}, + {:checkpoints, execution_external_id, checkpoints} + ) + |> flush_notifications() + + {:reply, :ok, state} + end + + :error -> + {:reply, {:error, :not_found}, state} + end + end + def handle_call( {:register_stream, execution_external_id, index, buffer, timeout_ms, session_external_id}, _from, @@ -3349,6 +3415,22 @@ defmodule Coflux.Orchestration.Server do # Enrich arguments with resolved references (asset/execution metadata) enriched_arguments = Enum.map(arguments, &build_value(&1, state.db)) + # Checkpoints travel with the execute message in the same + # wire form as arguments, and are handled the same way + # end-to-end — including the worker downloading any + # blob-backed value before the adapter starts. Bounded by + # this execution's own attempt, so it sees what it started + # with rather than anything a stale writer lands later. + {:ok, checkpoints} = + Checkpoints.get_effective( + state.db, + execution.step_id, + get_workspace_chain(state, execution.workspace_id), + execution.attempt + ) + + enriched_checkpoints = enrich_checkpoints(checkpoints, state.db) + workspace_external_id = state.workspaces[execution.workspace_id].external_id execution_external_id = @@ -3380,7 +3462,7 @@ defmodule Coflux.Orchestration.Server do build_streams_config( execution.streams_buffer, execution.streams_timeout_ms - )} + ), enriched_checkpoints} ) # Notify sessions topic of updated total @@ -4303,6 +4385,20 @@ defmodule Coflux.Orchestration.Server do end end + # The workspace inheritance chain ordered nearest-first (the workspace + # itself, then its bases). Checkpoint reads walk this and stop at the first + # workspace that has any state, so the nearest scope wins outright. + defp get_workspace_chain(state, workspace_id, ids \\ []) do + workspace = Map.fetch!(state.workspaces, workspace_id) + ids = [workspace_id | ids] + + if workspace.base_id do + get_workspace_chain(state, workspace.base_id, ids) + else + Enum.reverse(ids) + end + end + defp get_cache_workspace_ids(state, workspace_id, ids \\ []) do workspace = Map.fetch!(state.workspaces, workspace_id) @@ -4642,13 +4738,21 @@ defmodule Coflux.Orchestration.Server do ws_ext_id = workspace_external_id(state, workspace_id) + {:ok, checkpoints} = + Checkpoints.get_effective( + state.db, + step.id, + get_workspace_chain(state, workspace_id), + attempt + ) + state = state |> put_in([Access.key(:execution_ids), execution_external_id], execution_id) |> notify_listeners( {:run, run.external_id}, {:execution, step.number, attempt, execution_external_id, ws_ext_id, created_at, - execute_after, dependencies, principal} + execute_after, dependencies, principal, enrich_checkpoints(checkpoints, state.db)} ) |> notify_listeners( {:modules, ws_ext_id}, @@ -5307,6 +5411,20 @@ defmodule Coflux.Orchestration.Server do {:ok, run_submitted_inputs} = Inputs.get_submitted_inputs_for_run(db, run.id) {:ok, run_asset_deps} = Runs.get_asset_dependencies_for_run(db, run.id) + # Resolving a checkpoint needs the workspace chain of the execution + # reading it. Resolved from `db` rather than `state` because this also + # runs against archived epochs, which remap workspace ids. + workspace_chains = + run_executions + |> Enum.map(fn {_execution_id, _step_id, _attempt, workspace_id, _, _, _, _, _} -> + workspace_id + end) + |> Enum.uniq() + |> Map.new(fn workspace_id -> + {:ok, chain} = Workspaces.get_workspace_chain(db, workspace_id) + {workspace_id, chain} + end) + submitted_inputs_by_execution = run_submitted_inputs |> Enum.group_by( @@ -5524,6 +5642,15 @@ defmodule Coflux.Orchestration.Server do streams = streams_with_resolved_reasons(db, execution_id) + {:ok, {checkpoints_before, checkpoints_after}} = + Checkpoints.get_execution_snapshots( + db, + execution_id, + step.id, + Map.fetch!(workspace_chains, workspace_id), + attempt + ) + {attempt, %{ execution_id: exec_external_id, @@ -5543,7 +5670,11 @@ defmodule Coflux.Orchestration.Server do result_created_by: result_created_by, children: Map.get(run_children, execution_id, []), metric_definitions: Map.get(metric_definitions_by_execution, execution_id, %{}), - streams: streams + streams: streams, + checkpoints: %{ + before: enrich_checkpoints(checkpoints_before, db), + after: enrich_checkpoints(checkpoints_after, db) + } }} end) }} @@ -6125,6 +6256,13 @@ defmodule Coflux.Orchestration.Server do defp normalize_value({:blob, key, size, refs}), do: {:blob, key, size, normalize_references(refs)} + # Checkpoint values carry references (assets, executions, inputs) in the + # same form as arguments, so they need the same resolution before going out + # to a worker or a topic. + defp enrich_checkpoints(checkpoints, db) do + Map.new(checkpoints, fn {name, value} -> {name, build_value(value, db)} end) + end + defp build_value(value, db) do case value do {:raw, data, references} -> diff --git a/server/lib/coflux/orchestration/workspaces.ex b/server/lib/coflux/orchestration/workspaces.ex index a155f369..6a1bbb7f 100644 --- a/server/lib/coflux/orchestration/workspaces.ex +++ b/server/lib/coflux/orchestration/workspaces.ex @@ -17,6 +17,24 @@ defmodule Coflux.Orchestration.Workspaces do end end + @doc """ + The workspace inheritance chain, nearest-first (the workspace itself, then + its bases). + + The orchestration server keeps the same thing in memory, but reads from an + archived epoch can't use that — rotation remaps workspace ids — so this + resolves it against whichever database the ids belong to. + """ + def get_workspace_chain(db, workspace_id, ids \\ []) do + case get_workspace_by_id(db, workspace_id) do + {:ok, %{base_id: nil}} -> + {:ok, Enum.reverse([workspace_id | ids])} + + {:ok, %{base_id: base_id}} -> + get_workspace_chain(db, base_id, [workspace_id | ids]) + end + end + def get_all_workspaces(db) do case query( db, diff --git a/server/lib/coflux/topics/run.ex b/server/lib/coflux/topics/run.ex index 2051f486..0d64ae0f 100644 --- a/server/lib/coflux/topics/run.ex +++ b/server/lib/coflux/topics/run.ex @@ -75,7 +75,7 @@ defmodule Coflux.Topics.Run do defp process_notification( topic, {:execution, step_number, attempt, execution_external_id, workspace_external_id, - created_at, execute_after, dependencies, created_by} + created_at, execute_after, dependencies, created_by, checkpoints} ) do if workspace_external_id in topic.state.workspace_ids do Topic.set( @@ -105,7 +105,10 @@ defmodule Coflux.Topics.Run do inputs: %{}, result: nil, metrics: %{}, - streams: %{} + streams: %{}, + # Nothing has run yet, so what the execution will start from is also + # what it currently holds. + checkpoints: build_checkpoints(checkpoints, checkpoints) } ) else @@ -237,6 +240,15 @@ defmodule Coflux.Topics.Run do end) end + defp process_notification(topic, {:checkpoints, execution_external_id, checkpoints}) do + # Only the "after" side moves — what the execution started from was fixed + # when it was created. The whole snapshot is replaced rather than merged: a + # reset drops a name entirely, so merging would leave it behind. + update_execution(topic, execution_external_id, fn topic, base_path -> + Topic.set(topic, base_path ++ [:checkpoints, :after], build_values(checkpoints)) + end) + end + defp process_notification( topic, {:stream_opened, execution_external_id, index, buffer, timeout_ms, created_at} @@ -443,7 +455,12 @@ defmodule Coflux.Topics.Run do upper: def_data.upper }} end), - streams: build_streams(execution.streams) + streams: build_streams(execution.streams), + checkpoints: + build_checkpoints( + execution.checkpoints.before, + execution.checkpoints.after + ) }} end) }} @@ -590,6 +607,17 @@ defmodule Coflux.Topics.Run do end end + # Checkpoints are reported as a pair so an attempt can be read as a + # transition: `before` is what it was handed when it started, `after` what it + # ended up holding. They're equal for an attempt that didn't write. + defp build_checkpoints(before, after_) do + %{before: build_values(before), after: build_values(after_)} + end + + defp build_values(values) do + Map.new(values, fn {name, value} -> {name, build_value(value)} end) + end + defp build_streams(streams) do Map.new(streams, fn {index, buffer, timeout_ms, opened_at, nil, nil, nil} -> diff --git a/server/priv/migrations/orchestration/5.sql b/server/priv/migrations/orchestration/5.sql new file mode 100644 index 00000000..e375deef --- /dev/null +++ b/server/priv/migrations/orchestration/5.sql @@ -0,0 +1,51 @@ +-- Checkpoints — durable, named state scoped to a step within a workspace. +-- +-- Motivation: an execution that suspends, retries or recurs re-runs from the +-- beginning, so anything it needs to carry forward (a cursor, a page token, +-- an accumulator) has nowhere to live. Memoised child tasks cover values that +-- are naturally task results; checkpoints cover the rest. +-- +-- Scope is (step, workspace). Reads walk the workspace chain nearest-first +-- and take the latest attempt that touched checkpoints in the nearest +-- workspace that has any. A re-run in a descendant workspace therefore reads +-- the base's state but writes only to its own overlay — inheritance is +-- read-only and one-directional. +-- +-- Each execution's row-set is a complete snapshot, not a delta: on an +-- execution's first checkpoint write the server materialises the previous +-- effective set and then applies the delta, in one transaction. The worker +-- only ever sends what changed, so unchanged values are carried forward as +-- values_ FK copies and are never re-serialised or re-uploaded. +-- +-- DELIBERATE EXCEPTION TO THE APPEND-ONLY CONVENTION: rows are updated in +-- place within an execution. Nothing can observe an intra-execution +-- overwrite — the execution reads its own writes from memory, and the +-- worker-side throttle discards intermediate values before they reach the +-- server. History across attempts is preserved by the (execution_id, name) +-- key, and is compacted to the effective snapshot at epoch rotation. +-- +-- Three states per name: +-- * row with value_id — set to a value (possibly the null value) +-- * row with NULL value_id — explicitly reset at this attempt (tombstone) +-- * no row — never existed +-- +-- reset() ALWAYS writes a tombstone, even for a name that was never set. +-- That is what keeps an execution's row-set non-empty whenever it touched +-- checkpoints at all, so "reset everything" can never be mistaken for "wrote +-- nothing" and resurrect the pre-reset state from an earlier attempt. + +CREATE TABLE checkpoints ( + execution_id INTEGER NOT NULL, + name TEXT NOT NULL, + value_id INTEGER, + updated_at INTEGER NOT NULL, + PRIMARY KEY (execution_id, name), + FOREIGN KEY (execution_id) REFERENCES executions ON DELETE CASCADE, + FOREIGN KEY (value_id) REFERENCES values_ ON DELETE RESTRICT +) STRICT; + +-- Resolving "the latest attempt of this step in this workspace" is the read +-- path for every execution start. executions has UNIQUE (step_id, attempt) +-- and a separate idx_executions_workspace_id, neither of which can satisfy +-- that in one seek. +CREATE INDEX idx_executions_step_workspace ON executions(step_id, workspace_id, attempt); diff --git a/server/test/coflux/checkpoints_test.exs b/server/test/coflux/checkpoints_test.exs new file mode 100644 index 00000000..9344bacc --- /dev/null +++ b/server/test/coflux/checkpoints_test.exs @@ -0,0 +1,322 @@ +defmodule Coflux.CheckpointsTest do + use ExUnit.Case, async: true + + alias Coflux.Orchestration.Checkpoints + alias Coflux.Store.Migrations + alias Exqlite.Sqlite3 + + @base_ws 1 + @child_ws 2 + + setup do + {:ok, db} = Sqlite3.open(":memory:") + :ok = Migrations.run(db, "orchestration") + + create_workspace(db, @base_ws, "base") + create_workspace(db, @child_ws, "child") + create_run(db, 1, "r1") + create_step(db, 1, 1, 0) + + {:ok, db: db} + end + + # The chain a base-workspace execution sees, and the chain a child sees. + defp base_chain, do: [@base_ws] + defp child_chain, do: [@child_ws, @base_ws] + + defp val(data), do: {:raw, data, []} + + describe "get_effective/4" do + test "is empty when nothing has been written", %{db: db} do + assert {:ok, %{}} = Checkpoints.get_effective(db, 1, base_chain()) + end + + test "returns what the latest attempt recorded", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + set(db, 1, 1, %{"cursor" => val(5)}) + + assert {:ok, %{"cursor" => {:raw, 5, []}}} = Checkpoints.get_effective(db, 1, base_chain()) + end + + test "bounds the lookup to attempts below before_attempt", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + create_execution(db, 2, 1, 2, @base_ws) + set(db, 1, 1, %{"cursor" => val(5)}) + set(db, 2, 2, %{"cursor" => val(9)}) + + # What attempt 2 started with, rather than what it ended with. + assert {:ok, %{"cursor" => {:raw, 5, []}}} = + Checkpoints.get_effective(db, 1, base_chain(), 2) + + assert {:ok, %{"cursor" => {:raw, 9, []}}} = Checkpoints.get_effective(db, 1, base_chain()) + end + + test "distinguishes a stored null from a reset", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + set(db, 1, 1, %{"cursor" => val(nil)}) + + assert {:ok, entries} = Checkpoints.get_effective(db, 1, base_chain()) + assert Map.has_key?(entries, "cursor") + assert entries["cursor"] == {:raw, nil, []} + + create_execution(db, 2, 1, 2, @base_ws) + reset(db, 2, 2, ["cursor"]) + + assert {:ok, entries} = Checkpoints.get_effective(db, 1, base_chain()) + refute Map.has_key?(entries, "cursor") + end + end + + describe "apply_delta/7 carry-forward" do + test "preserves names the latest attempt didn't write", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + set(db, 1, 1, %{"cursor" => val(5), "batch" => val("a")}) + + create_execution(db, 2, 1, 2, @base_ws) + set(db, 2, 2, %{"cursor" => val(9)}) + + assert {:ok, entries} = Checkpoints.get_effective(db, 1, base_chain()) + assert entries["cursor"] == {:raw, 9, []} + assert entries["batch"] == {:raw, "a", []} + end + + test "writes a complete snapshot for each attempt", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + set(db, 1, 1, %{"cursor" => val(5), "batch" => val("a")}) + + create_execution(db, 2, 1, 2, @base_ws) + set(db, 2, 2, %{"cursor" => val(9)}) + + # Attempt 2 carries `batch` forward rather than leaving it to a merge. + assert {:ok, rows} = Checkpoints.get_rows_for_execution(db, 2) + assert Enum.map(rows, fn {name, _, _} -> name end) == ["batch", "cursor"] + end + + test "only carries forward once per execution", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + set(db, 1, 1, %{"cursor" => val(5)}) + + create_execution(db, 2, 1, 2, @base_ws) + reset(db, 2, 2, ["cursor"]) + # A later write in the same execution must not re-materialise the + # previous snapshot and undo the reset. + set(db, 2, 2, %{"batch" => val("a")}) + + assert {:ok, entries} = Checkpoints.get_effective(db, 1, base_chain()) + refute Map.has_key?(entries, "cursor") + assert entries["batch"] == {:raw, "a", []} + end + end + + describe "reset" do + test "a reset name falls back to absent in the next attempt", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + set(db, 1, 1, %{"cursor" => val(5)}) + + create_execution(db, 2, 1, 2, @base_ws) + reset(db, 2, 2, ["cursor"]) + + create_execution(db, 3, 1, 3, @base_ws) + set(db, 3, 3, %{"other" => val(1)}) + + assert {:ok, entries} = Checkpoints.get_effective(db, 1, base_chain()) + refute Map.has_key?(entries, "cursor") + end + + test "resetting every name does not resurrect the pre-reset state", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + set(db, 1, 1, %{"cursor" => val(5), "batch" => val("a")}) + + create_execution(db, 2, 1, 2, @base_ws) + reset(db, 2, 2, ["cursor", "batch"]) + + # Attempt 2's row-set is all tombstones, which must still count as + # "touched checkpoints" — otherwise the read falls back to attempt 1. + assert {:ok, %{}} = Checkpoints.get_effective(db, 1, base_chain()) + end + + test "resetting a name that was never set is recorded", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + reset(db, 1, 1, ["cursor"]) + + assert {:ok, [{"cursor", nil, _}]} = Checkpoints.get_rows_for_execution(db, 1) + assert {:ok, %{}} = Checkpoints.get_effective(db, 1, base_chain()) + end + end + + describe "workspace scoping" do + test "a child reads the base's state", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + set(db, 1, 1, %{"cursor" => val(5)}) + + assert {:ok, %{"cursor" => {:raw, 5, []}}} = Checkpoints.get_effective(db, 1, child_chain()) + end + + test "a child's writes do not affect the base", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + set(db, 1, 1, %{"cursor" => val(5)}) + + create_execution(db, 2, 1, 2, @child_ws) + set(db, 2, 2, %{"cursor" => val(9)}, child_chain()) + + assert {:ok, %{"cursor" => {:raw, 5, []}}} = Checkpoints.get_effective(db, 1, base_chain()) + assert {:ok, %{"cursor" => {:raw, 9, []}}} = Checkpoints.get_effective(db, 1, child_chain()) + end + + test "the nearest workspace wins outright", %{db: db} do + create_execution(db, 1, 1, 1, @child_ws) + set(db, 1, 1, %{"cursor" => val(9)}, child_chain()) + + # A later base-workspace attempt must not override the child's own + # state — inheritance is a fallback, not a merge. + create_execution(db, 2, 1, 2, @base_ws) + set(db, 2, 2, %{"cursor" => val(5)}) + + assert {:ok, %{"cursor" => {:raw, 9, []}}} = Checkpoints.get_effective(db, 1, child_chain()) + end + end + + describe "ordering" do + test "a write from an older attempt does not affect the current state", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + create_execution(db, 2, 1, 2, @base_ws) + + set(db, 2, 2, %{"cursor" => val(9)}) + # A stale worker flushing after being superseded. + set(db, 1, 1, %{"cursor" => val(5)}) + + assert {:ok, %{"cursor" => {:raw, 9, []}}} = Checkpoints.get_effective(db, 1, base_chain()) + end + end + + describe "get_execution_snapshots/5" do + test "reports what an attempt started with and what it left", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + create_execution(db, 2, 1, 2, @base_ws) + set(db, 1, 1, %{"cursor" => val(5)}) + set(db, 2, 2, %{"cursor" => val(9)}) + + assert {:ok, {before, after_}} = + Checkpoints.get_execution_snapshots(db, 2, 1, base_chain(), 2) + + assert before == %{"cursor" => {:raw, 5, []}} + assert after_ == %{"cursor" => {:raw, 9, []}} + end + + test "an attempt that didn't write leaves the state as it found it", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + create_execution(db, 2, 1, 2, @base_ws) + set(db, 1, 1, %{"cursor" => val(5)}) + + assert {:ok, {before, after_}} = + Checkpoints.get_execution_snapshots(db, 2, 1, base_chain(), 2) + + assert before == %{"cursor" => {:raw, 5, []}} + assert after_ == before + end + + test "an attempt that reset everything is distinct from one that didn't write", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + create_execution(db, 2, 1, 2, @base_ws) + set(db, 1, 1, %{"cursor" => val(5)}) + reset(db, 2, 2, ["cursor"]) + + assert {:ok, {before, after_}} = + Checkpoints.get_execution_snapshots(db, 2, 1, base_chain(), 2) + + assert before == %{"cursor" => {:raw, 5, []}} + assert after_ == %{} + end + + test "the first attempt starts from nothing", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + set(db, 1, 1, %{"cursor" => val(5)}) + + assert {:ok, {%{}, after_}} = + Checkpoints.get_execution_snapshots(db, 1, 1, base_chain(), 1) + + assert after_ == %{"cursor" => {:raw, 5, []}} + end + + test "a child workspace attempt starts from the base's state", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + create_execution(db, 2, 1, 2, @child_ws) + set(db, 1, 1, %{"cursor" => val(5)}) + set(db, 2, 2, %{"cursor" => val(9)}, child_chain()) + + assert {:ok, {before, after_}} = + Checkpoints.get_execution_snapshots(db, 2, 1, child_chain(), 2) + + assert before == %{"cursor" => {:raw, 5, []}} + assert after_ == %{"cursor" => {:raw, 9, []}} + end + end + + describe "get_snapshot_execution_ids/2" do + test "returns the latest checkpoint-bearing execution per workspace", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + create_execution(db, 2, 1, 2, @base_ws) + create_execution(db, 3, 1, 3, @child_ws) + + set(db, 1, 1, %{"cursor" => val(1)}) + set(db, 2, 2, %{"cursor" => val(2)}) + set(db, 3, 3, %{"cursor" => val(3)}, child_chain()) + + assert {:ok, ids} = Checkpoints.get_snapshot_execution_ids(db, 1) + assert Enum.sort(ids) == [2, 3] + end + + test "is empty when no checkpoints exist", %{db: db} do + create_execution(db, 1, 1, 1, @base_ws) + + assert {:ok, []} = Checkpoints.get_snapshot_execution_ids(db, 1) + end + end + + ## Helpers + + defp set(db, execution_id, attempt, entries, chain \\ nil) do + {:ok, _} = + Checkpoints.apply_delta(db, execution_id, 1, chain || base_chain(), attempt, entries, []) + end + + defp reset(db, execution_id, attempt, names, chain \\ nil) do + {:ok, _} = + Checkpoints.apply_delta(db, execution_id, 1, chain || base_chain(), attempt, %{}, names) + end + + defp create_workspace(db, id, external_id) do + :ok = + Sqlite3.execute( + db, + "INSERT INTO workspaces (id, external_id) VALUES (#{id}, '#{external_id}')" + ) + end + + defp create_run(db, id, external_id) do + :ok = + Sqlite3.execute( + db, + "INSERT INTO runs (id, external_id, created_at) VALUES (#{id}, '#{external_id}', 0)" + ) + end + + defp create_step(db, id, run_id, number) do + :ok = + Sqlite3.execute(db, """ + INSERT INTO steps ( + id, number, run_id, module, target, type, priority, wait_for, + retry_limit, retry_backoff_min, retry_backoff_max, created_at + ) + VALUES (#{id}, #{number}, #{run_id}, 'module', 'target', 0, 0, 0, 0, 0, 0, 0) + """) + end + + defp create_execution(db, id, step_id, attempt, workspace_id) do + :ok = + Sqlite3.execute(db, """ + INSERT INTO executions (id, step_id, attempt, workspace_id, created_at) + VALUES (#{id}, #{step_id}, #{attempt}, #{workspace_id}, 0) + """) + end +end diff --git a/tests/support/executor.py b/tests/support/executor.py index 1ac02173..aff1493f 100644 --- a/tests/support/executor.py +++ b/tests/support/executor.py @@ -42,8 +42,8 @@ def _unwrap_select_result(result): Execution = namedtuple( "Execution", - ["conn", "execution_id", "module", "target", "arguments", "streams"], - defaults=[None], + ["conn", "execution_id", "module", "target", "arguments", "streams", "checkpoints"], + defaults=[None, None], ) @@ -104,7 +104,10 @@ def run_one(self, handler): self.send(response) def recv_execute(self, **kwargs): - """Receive an execute message, return (execution_id, module, target, arguments, streams).""" + """Receive an execute message. + + Returns (execution_id, module, target, arguments, streams, checkpoints). + """ msg = self.recv(**kwargs) assert msg["method"] == "execute", f"expected execute, got {msg['method']}" p = msg["params"] @@ -114,6 +117,7 @@ def recv_execute(self, **kwargs): p["target"], p.get("arguments", []), p.get("streams"), + p.get("checkpoints") or {}, ) def _request(self, msg): @@ -225,6 +229,18 @@ def suspend(self, execution_id, execute_after=None): msg = protocol.suspend_request(None, execution_id, execute_after) return self._request(msg) + def checkpoint_set(self, execution_id, **values): + """Set one or more checkpoints (plain JSON values).""" + self.send(protocol.checkpoint_update_notification(execution_id, set_=values)) + + def checkpoint_reset(self, execution_id, *names): + """Reset one or more checkpoints back to their declared default.""" + self.send(protocol.checkpoint_update_notification(execution_id, reset=names)) + + def flush(self, execution_id): + """Flush buffered state, returning once the server has acknowledged it.""" + return self._request(protocol.flush_request(None, execution_id)) + def persist_asset(self, execution_id, paths, metadata=None): """Persist files as an asset and return the result.""" msg = protocol.persist_asset_request(None, execution_id, paths, metadata) @@ -538,9 +554,13 @@ def next_execute(self, timeout=10): if idx in self._consumed: continue try: - eid, module, target, args, streams = conn.recv_execute(timeout=0.1) + eid, module, target, args, streams, checkpoints = conn.recv_execute( + timeout=0.1 + ) self._consumed.add(idx) - return Execution(conn, eid, module, target, args, streams) + return Execution( + conn, eid, module, target, args, streams, checkpoints + ) except TimeoutError: continue except (ConnectionError, OSError): diff --git a/tests/support/protocol.py b/tests/support/protocol.py index 54fe19ab..1d95dcaa 100644 --- a/tests/support/protocol.py +++ b/tests/support/protocol.py @@ -224,6 +224,32 @@ def get_asset_request(request_id, execution_id, asset_id): } +def checkpoint_update_notification(execution_id, set_=None, reset=None): + """Build a checkpoint delta notification. + + ``set_`` maps names to plain JSON values (wrapped here); ``reset`` lists + names to clear. + """ + params = {"execution_id": execution_id} + if set_: + params["set"] = { + name: {"type": "inline", "format": "json", "value": value} + for name, value in set_.items() + } + if reset: + params["reset"] = list(reset) + return {"method": "checkpoint_update", "params": params} + + +def flush_request(request_id, execution_id): + """Build a flush request — responds once buffered state reaches the server.""" + return { + "id": request_id, + "method": "flush", + "params": {"execution_id": execution_id}, + } + + def register_group_notification(execution_id, group_id, name=None): params = {"execution_id": execution_id, "group_id": group_id} if name is not None: diff --git a/tests/test_checkpoints.py b/tests/test_checkpoints.py new file mode 100644 index 00000000..9e658b1d --- /dev/null +++ b/tests/test_checkpoints.py @@ -0,0 +1,582 @@ +"""Tests for checkpoints — step-scoped durable state.""" + +import json + +import pytest +from support import cli +from support.helpers import api_post, managed_worker, poll_result +from support.manifest import workflow + + +def values(ex): + """Plain {name: value} from an execute message's checkpoint payload. + + Blob-backed values arrive as a local file path (the CLI downloads them, + exactly as it does for arguments), so those are read back off disk. + + A stored null has no "value" key at all — it's omitted on the wire, the + same as a null argument — so presence of the name is what distinguishes it + from a reset. + """ + result = {} + for name, value in (ex.checkpoints or {}).items(): + if value.get("type") == "file": + with open(value["path"]) as f: + result[name] = json.load(f) + else: + result[name] = value.get("value") + return result + + +def test_carried_across_suspend(worker): + """A checkpoint written before suspending is visible to the successor.""" + targets = [workflow("test", "poller")] + + with worker(targets) as ctx: + resp = ctx.submit("test", "poller") + run_id = resp["runId"] + + ex0 = ctx.executor.next_execute() + assert values(ex0) == {} + ex0.conn.checkpoint_set(ex0.execution_id, cursor=5) + ex0.conn.suspend(ex0.execution_id) + + ex1 = ctx.executor.next_execute() + assert ex1.execution_id != ex0.execution_id + assert values(ex1) == {"cursor": 5} + + ex1.conn.complete(ex1.execution_id, value="done") + assert ctx.result(run_id)["value"]["data"] == "done" + + +def test_carried_across_retry(worker): + """A failing attempt's checkpoint is visible to its retry.""" + targets = [ + workflow( + "test", + "flaky", + retries={"limit": 1, "backoff_min_ms": 0, "backoff_max_ms": 0}, + ) + ] + + with worker(targets) as ctx: + resp = ctx.submit("test", "flaky") + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, progress=17) + ex0.conn.fail(ex0.execution_id, "RuntimeError", "transient failure") + + ex1 = ctx.executor.next_execute() + assert values(ex1) == {"progress": 17} + + ex1.conn.complete(ex1.execution_id, value="recovered") + assert ctx.result(resp["runId"])["value"]["data"] == "recovered" + + +def test_carried_across_recurrence(worker): + """Each recurrence sees what the previous iteration wrote.""" + targets = [workflow("test", "ticker", recurrent=True)] + + with worker(targets) as ctx: + ctx.submit("test", "ticker") + + seen = [] + for i in range(3): + ex = ctx.executor.next_execute() + seen.append(values(ex).get("count")) + ex.conn.checkpoint_set(ex.execution_id, count=i) + # Recurrent targets recur while they return None. + ex.conn.complete(ex.execution_id, value=None) + + assert seen == [None, 0, 1] + + +def test_carried_across_manual_rerun(worker): + """Re-running a step from the CLI preserves its checkpoint.""" + targets = [workflow("test", "main")] + + with worker(targets) as ctx: + resp = ctx.submit("test", "main") + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor="a") + ex0.conn.complete(ex0.execution_id, value="first") + ctx.result(resp["runId"]) + + ctx.rerun(resp["stepId"]) + + ex1 = ctx.executor.next_execute() + assert values(ex1) == {"cursor": "a"} + ex1.conn.complete(ex1.execution_id, value="second") + + +def test_scoped_to_step(worker): + """A checkpoint is not visible to another step, or to another run.""" + targets = [workflow("test", "main"), workflow("test", "other")] + + with worker(targets, concurrency=2) as ctx: + resp = ctx.submit("test", "main") + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=1) + ex0.conn.complete(ex0.execution_id, value="done") + ctx.result(resp["runId"]) + + # A different target — different step, different run. + resp2 = ctx.submit("test", "other") + ex1 = ctx.executor.next_execute() + assert values(ex1) == {} + ex1.conn.complete(ex1.execution_id, value="done") + ctx.result(resp2["runId"]) + + # The same target again — still a new run, so still a new step. + resp3 = ctx.submit("test", "main") + ex2 = ctx.executor.next_execute() + assert values(ex2) == {} + ex2.conn.complete(ex2.execution_id, value="done") + ctx.result(resp3["runId"]) + + +def test_names_carried_forward_independently(worker): + """A name the latest attempt didn't write is still carried forward.""" + targets = [workflow("test", "poller")] + + with worker(targets) as ctx: + ctx.submit("test", "poller") + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=1, batch="a") + ex0.conn.suspend(ex0.execution_id) + + ex1 = ctx.executor.next_execute() + assert values(ex1) == {"cursor": 1, "batch": "a"} + ex1.conn.checkpoint_set(ex1.execution_id, cursor=2) + ex1.conn.suspend(ex1.execution_id) + + ex2 = ctx.executor.next_execute() + assert values(ex2) == {"cursor": 2, "batch": "a"} + ex2.conn.complete(ex2.execution_id, value="done") + + +def test_reset(worker): + """A reset name is gone for the next attempt, not restored from an older one.""" + targets = [workflow("test", "poller")] + + with worker(targets) as ctx: + ctx.submit("test", "poller") + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=1, batch="a") + ex0.conn.suspend(ex0.execution_id) + + ex1 = ctx.executor.next_execute() + ex1.conn.checkpoint_reset(ex1.execution_id, "cursor") + ex1.conn.suspend(ex1.execution_id) + + ex2 = ctx.executor.next_execute() + assert values(ex2) == {"batch": "a"} + + # And it stays gone through an attempt that touches nothing. + ex2.conn.suspend(ex2.execution_id) + ex3 = ctx.executor.next_execute() + assert values(ex3) == {"batch": "a"} + ex3.conn.complete(ex3.execution_id, value="done") + + +def test_reset_every_name(worker): + """Resetting everything leaves an empty checkpoint, not the pre-reset state. + + An attempt whose entire row-set is tombstones still has to count as having + touched checkpoints — otherwise the read falls back to the previous attempt + and resurrects exactly what was just cleared. + """ + targets = [workflow("test", "poller")] + + with worker(targets) as ctx: + ctx.submit("test", "poller") + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=1, batch="a") + ex0.conn.suspend(ex0.execution_id) + + ex1 = ctx.executor.next_execute() + assert values(ex1) == {"cursor": 1, "batch": "a"} + ex1.conn.checkpoint_reset(ex1.execution_id, "cursor", "batch") + ex1.conn.suspend(ex1.execution_id) + + ex2 = ctx.executor.next_execute() + assert values(ex2) == {} + ex2.conn.complete(ex2.execution_id, value="done") + + +def test_null_value_is_not_a_reset(worker): + """A checkpoint set to null is stored; only reset clears it.""" + targets = [workflow("test", "poller")] + + with worker(targets) as ctx: + ctx.submit("test", "poller") + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=None) + ex0.conn.suspend(ex0.execution_id) + + ex1 = ctx.executor.next_execute() + assert "cursor" in (ex1.checkpoints or {}) + assert values(ex1) == {"cursor": None} + + ex1.conn.checkpoint_reset(ex1.execution_id, "cursor") + ex1.conn.suspend(ex1.execution_id) + + ex2 = ctx.executor.next_execute() + assert "cursor" not in (ex2.checkpoints or {}) + ex2.conn.complete(ex2.execution_id, value="done") + + +def test_repeated_writes_coalesce(worker): + """Rapid writes are throttled, but the final value always lands.""" + targets = [workflow("test", "poller")] + + with worker(targets) as ctx: + ctx.submit("test", "poller") + + ex0 = ctx.executor.next_execute() + for i in range(50): + ex0.conn.checkpoint_set(ex0.execution_id, cursor=i) + ex0.conn.suspend(ex0.execution_id) + + ex1 = ctx.executor.next_execute() + assert values(ex1) == {"cursor": 49} + ex1.conn.complete(ex1.execution_id, value="done") + + +def test_flush_is_acknowledged(worker): + """flush returns only once buffered state has reached the server.""" + targets = [workflow("test", "poller")] + + with worker(targets) as ctx: + ctx.submit("test", "poller") + + ex0 = ctx.executor.next_execute() + # Two writes inside one throttle window: the second is buffered, and + # the flush is what gets it to the server. + ex0.conn.checkpoint_set(ex0.execution_id, cursor=1) + ex0.conn.checkpoint_set(ex0.execution_id, cursor=2) + resp = ex0.conn.flush(ex0.execution_id) + assert "error" not in resp + + ex0.conn.suspend(ex0.execution_id) + ex1 = ctx.executor.next_execute() + assert values(ex1) == {"cursor": 2} + ex1.conn.complete(ex1.execution_id, value="done") + + +def test_buffered_write_survives_a_crash(worker): + """A write still in the throttle buffer lands when the adapter dies. + + The first write of an execution goes out on the leading edge, so it's the + second — coalesced into the open window — that has to survive a process + that never reports anything. Covers the crash path end to end; it doesn't + pin down *which* flush delivers it, since the retry is not in practice + assigned early enough to distinguish them. + """ + targets = [ + workflow( + "test", + "poller", + retries={"limit": 1, "backoff_min_ms": 0, "backoff_max_ms": 0}, + ) + ] + + with worker(targets) as ctx: + ctx.submit("test", "poller") + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=1) + ex0.conn.checkpoint_set(ex0.execution_id, cursor=2) + # Die without reporting a result or an error. + ex0.conn.close() + + ex1 = ctx.executor.next_execute() + assert values(ex1) == {"cursor": 2} + ex1.conn.complete(ex1.execution_id, value="done") + + +def test_buffered_write_survives_a_timeout(worker): + """A write still in the throttle buffer lands when the execution times out. + + Same shape as the crash case: the timeout is reported by the worker rather + than the adapter, so it needs its own flush ahead of the report. + """ + targets = [ + workflow( + "test", + "poller", + timeout=300, + retries={"limit": 1, "backoff_min_ms": 0, "backoff_max_ms": 0}, + ) + ] + + with worker(targets) as ctx: + ctx.submit("test", "poller") + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=1) + ex0.conn.checkpoint_set(ex0.execution_id, cursor=2) + # Never report anything — let the timeout fire. It's shorter than the + # throttle window, so the second write is still buffered when it does. + + ex1 = ctx.executor.next_execute() + assert values(ex1) == {"cursor": 2} + ex1.conn.complete(ex1.execution_id, value="done") + + +def test_write_after_completion_ignored(worker): + """A write from a superseded attempt can't affect the current one. + + The old connection may already be torn down by the abort that follows a + suspension, in which case nothing is sent at all — either way the + successor must see the pre-suspend value. + """ + targets = [workflow("test", "poller")] + + with worker(targets) as ctx: + ctx.submit("test", "poller") + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=1) + ex0.conn.suspend(ex0.execution_id) + + try: + ex0.conn.checkpoint_set(ex0.execution_id, cursor=999) + except (ConnectionError, OSError): + pass + + ex1 = ctx.executor.next_execute() + assert values(ex1) == {"cursor": 1} + ex1.conn.complete(ex1.execution_id, value="done") + + +def test_blob_backed_value(worker): + """A checkpoint too large to inline round-trips via the blob store.""" + targets = [workflow("test", "poller")] + payload = "x" * 5000 + + with worker(targets) as ctx: + ctx.submit("test", "poller") + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, blob=payload) + ex0.conn.suspend(ex0.execution_id) + + ex1 = ctx.executor.next_execute() + assert values(ex1) == {"blob": payload} + ex1.conn.complete(ex1.execution_id, value="done") + + +def test_workspace_inheritance(worker): + """A derived workspace reads the base's state but writes only its own.""" + targets = [workflow("test", "poller")] + + with worker(targets, workspace="base") as ctx_base: + resp = ctx_base.submit("test", "poller") + run_id = resp["runId"] + step_id = resp["stepId"] + + ex0 = ctx_base.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor="base-1") + ex0.conn.complete(ex0.execution_id, value="done") + ctx_base.result(run_id) + + saved_host = ctx_base.host + + cli.workspaces_create("derived", base="base", host=saved_host, workspace="derived") + + with worker(targets, workspace="derived") as ctx_derived: + # Inherited from the base, since the derived workspace has no state. + ctx_derived.rerun(step_id) + ex1 = ctx_derived.executor.next_execute() + assert values(ex1) == {"cursor": "base-1"} + + ex1.conn.checkpoint_set(ex1.execution_id, cursor="derived-1") + ex1.conn.complete(ex1.execution_id, value="done") + + # Now the derived workspace has its own, which wins outright. + ctx_derived.rerun(step_id) + ex2 = ctx_derived.executor.next_execute() + assert values(ex2) == {"cursor": "derived-1"} + ex2.conn.complete(ex2.execution_id, value="done") + + # The base is untouched by any of that. + with worker(targets, workspace="base") as ctx_base: + ctx_base.rerun(step_id) + ex3 = ctx_base.executor.next_execute() + assert values(ex3) == {"cursor": "base-1"} + ex3.conn.complete(ex3.execution_id, value="done") + + +def test_survives_epoch_rotation(isolated_server, tmp_path): + """Rotation carries the effective checkpoint into the new epoch.""" + server, host, project_id = isolated_server + targets = [workflow("test", "poller")] + + with managed_worker(targets, host, tmp_path) as executor: + resp = cli.submit("test/poller", host=host) + run_id = resp["runId"] + step_id = resp["stepId"] + + ex0 = executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=7, batch="a") + ex0.conn.suspend(ex0.execution_id) + + ex1 = executor.next_execute() + assert values(ex1) == {"cursor": 7, "batch": "a"} + # Reset one name, so the snapshot carried over is the post-reset state + # rather than everything ever written. + ex1.conn.checkpoint_reset(ex1.execution_id, "batch") + ex1.conn.complete(ex1.execution_id, value="done") + poll_result(run_id, host) + + api_post(server.port, project_id, "rotate_epoch") + + cli.runs_rerun(step_id, host=host) + ex2 = executor.next_execute() + assert values(ex2) == {"cursor": 7} + ex2.conn.complete(ex2.execution_id, value="after-rotation") + + result = poll_result(run_id, host) + assert result["value"]["data"] == "after-rotation" + + +def test_reset_in_derived_workspace_survives_epoch_rotation(isolated_server, tmp_path): + """A reset in a derived workspace keeps masking the base after rotation. + + The tombstone is the only thing keeping the derived workspace's row-set + non-empty, and an empty row-set falls back to the base — so discarding it + at rotation would resurrect the value the reset removed. + """ + server, host, project_id = isolated_server + targets = [workflow("test", "poller")] + + with managed_worker(targets, host, tmp_path / "base") as executor: + resp = cli.submit("test/poller", host=host) + run_id = resp["runId"] + step_id = resp["stepId"] + + ex0 = executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=7) + ex0.conn.complete(ex0.execution_id, value="done") + poll_result(run_id, host) + + cli.workspaces_create("derived", base="default", host=host, workspace="derived") + + with managed_worker( + targets, host, tmp_path / "derived", workspace="derived" + ) as executor: + cli.runs_rerun(step_id, host=host, workspace="derived") + ex1 = executor.next_execute() + assert values(ex1) == {"cursor": 7} + + # Reset the only name it holds, so its entire row-set is tombstones. + ex1.conn.checkpoint_reset(ex1.execution_id, "cursor") + ex1.conn.complete(ex1.execution_id, value="done") + poll_result(run_id, host, workspace="derived") + + api_post(server.port, project_id, "rotate_epoch") + + cli.runs_rerun(step_id, host=host, workspace="derived") + ex2 = executor.next_execute() + assert values(ex2) == {} + ex2.conn.complete(ex2.execution_id, value="done") + poll_result(run_id, host, workspace="derived") + + # The base's own state is untouched by any of that. + with managed_worker(targets, host, tmp_path / "base-again") as executor: + cli.runs_rerun(step_id, host=host) + ex3 = executor.next_execute() + assert values(ex3) == {"cursor": 7} + ex3.conn.complete(ex3.execution_id, value="done") + + +def test_exposed_in_run_topic(worker): + """The run topic carries each execution's checkpoint, for Studio.""" + targets = [workflow("test", "poller")] + + with worker(targets) as ctx: + resp = ctx.submit("test", "poller") + run_id = resp["runId"] + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=5, batch="a") + ex0.conn.suspend(ex0.execution_id) + + ex1 = ctx.executor.next_execute() + ex1.conn.checkpoint_reset(ex1.execution_id, "batch") + ex1.conn.complete(ex1.execution_id, value="done") + ctx.result(run_id) + + executions = next(iter(ctx.inspect(run_id)["steps"].values()))["executions"] + + # Each attempt is reported as a transition: what it was handed when it + # started, and what it ended up holding. The first attempt started from + # nothing; the second inherited both names and reset one of them, so the + # reset shows up as an absence on the "after" side only. + assert executions["1"]["checkpoints"]["before"] == {} + assert executions["1"]["checkpoints"]["after"]["cursor"]["data"] == 5 + assert executions["1"]["checkpoints"]["after"]["batch"]["data"] == "a" + + assert executions["2"]["checkpoints"]["before"]["cursor"]["data"] == 5 + assert executions["2"]["checkpoints"]["before"]["batch"]["data"] == "a" + assert executions["2"]["checkpoints"]["after"]["cursor"]["data"] == 5 + assert "batch" not in executions["2"]["checkpoints"]["after"] + + +def test_run_topic_reports_untouched_state(worker): + """An attempt that didn't write reports the state it left untouched.""" + targets = [workflow("test", "poller")] + + with worker(targets) as ctx: + resp = ctx.submit("test", "poller") + run_id = resp["runId"] + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=5) + ex0.conn.suspend(ex0.execution_id) + + ex1 = ctx.executor.next_execute() + ex1.conn.complete(ex1.execution_id, value="done") + ctx.result(run_id) + + executions = next(iter(ctx.inspect(run_id)["steps"].values()))["executions"] + + # No rows of its own — but it still held the checkpoint, rather than + # nothing. + checkpoints = executions["2"]["checkpoints"] + assert checkpoints["before"]["cursor"]["data"] == 5 + assert checkpoints["after"]["cursor"]["data"] == 5 + + +def test_no_checkpoints_sent_when_empty(worker): + """An execution with no checkpoint state gets no payload at all.""" + targets = [workflow("test", "main")] + + with worker(targets) as ctx: + ctx.submit("test", "main") + ex = ctx.executor.next_execute() + assert ex.checkpoints == {} + ex.conn.complete(ex.execution_id, value="done") + + +@pytest.mark.parametrize("value", [0, "", False, []]) +def test_falsy_values_round_trip(worker, value): + """Falsy checkpoint values survive — absence is distinct from emptiness.""" + targets = [workflow("test", "poller")] + + with worker(targets) as ctx: + ctx.submit("test", "poller") + + ex0 = ctx.executor.next_execute() + ex0.conn.checkpoint_set(ex0.execution_id, cursor=value) + ex0.conn.suspend(ex0.execution_id) + + ex1 = ctx.executor.next_execute() + assert values(ex1) == {"cursor": value} + ex1.conn.complete(ex1.execution_id, value="done")