diff --git a/adapters/python/coflux/checkpoint.py b/adapters/python/coflux/checkpoint.py index 28d97564..46f6a15d 100644 --- a/adapters/python/coflux/checkpoint.py +++ b/adapters/python/coflux/checkpoint.py @@ -2,12 +2,20 @@ from __future__ import annotations +from collections.abc import Callable from typing import Any, Generic, TypeVar, overload from .state import get_context T = TypeVar("T") +# Checkpoint names starting with this are the adapter's own. Reserved as a +# namespace rather than name by name, so later internal state doesn't need +# another round of this. Enforced here in ``Checkpoint`` rather than in the +# context's checkpoint_get/set/reset, which are the path the adapter's own +# cursors go through. +RESERVED_PREFIX = "_" + class Checkpoint(Generic[T]): """A named value that survives across executions of a step. @@ -33,7 +41,13 @@ def poll_orders(): 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. + returns or fails is delivered. + + Writes are cut into deltas only where the execution could resume from + them, so a checkpoint is always read back as part of a state the step was + actually in: one written while a stream item is in a loop body's hands is + published with the cursor advance that consumes it, and dropped if that + iteration never finishes. 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. @@ -48,7 +62,8 @@ def poll_orders(): nothing is enforced at runtime. Args: - name: Checkpoint name, unique within the step. + name: Checkpoint name, unique within the step. Can't start with + ``_`` — that prefix is reserved for adapter-managed state. default: Value returned when the checkpoint has never been set, or has been reset. Client-side only — the server never sees it. """ @@ -63,6 +78,12 @@ def __init__(self, name: str) -> None: ... # ``-> 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: + if name.startswith(RESERVED_PREFIX): + raise ValueError( + f"checkpoint name {name!r} is reserved: names starting with" + f" {RESERVED_PREFIX!r} are used for adapter-managed state," + " such as the cursors behind stream suspension" + ) self._name = name self._default = default @@ -93,6 +114,27 @@ def set(self, value: T) -> None: """Set the value, replacing anything already there.""" get_context().checkpoint_set(self._name, value) + def update(self, fn: Callable[[T], T]) -> T: + """Set the value to ``fn(current)``, and return what was stored. + + The read-modify-write that most checkpoints do — advancing a + cursor, accumulating a total — without naming the old value:: + + n = count.update(lambda x: x + 1) + + ``fn`` receives the declared default when the checkpoint isn't set, + exactly as ``get()`` would return it. + + This is a read followed by a write, not an atomic swap: two threads + of one execution updating the same checkpoint can still lose one of + the updates. That only arises if you share a checkpoint across + threads — a task body and a ``cf.stream`` generator, say — in which + case guard it yourself. + """ + value = fn(self.get()) + self.set(value) + return value + def reset(self) -> None: """Clear the checkpoint, so ``get()`` returns the declared default. @@ -126,7 +168,9 @@ def flush() -> None: send_notification() Not needed before suspending, returning or raising — those are flushed - automatically. + automatically. Inside a stream loop body it also publishes what is being + held for the current item, ahead of the cursor advance that would + normally carry it. """ get_context().flush() diff --git a/adapters/python/coflux/context.py b/adapters/python/coflux/context.py index a090f30e..3ddf8e3b 100644 --- a/adapters/python/coflux/context.py +++ b/adapters/python/coflux/context.py @@ -11,7 +11,7 @@ from collections.abc import Callable, Iterator from contextlib import contextmanager from pathlib import Path -from typing import Any +from typing import Any, NoReturn from . import protocol from .dispatcher import get_dispatcher @@ -21,6 +21,7 @@ ExecutionCrashed, ExecutionTimeout, InputDismissed, + Suspending, create_execution_error, ) from .models import Asset, AssetEntry, AssetMetadata, Execution, Input @@ -83,7 +84,10 @@ def _timeout_to_ms(timeout: float | dt.timedelta | None) -> int | None: _group_id: contextvars.ContextVar[int | None] = contextvars.ContextVar( "_group_id", default=None ) -# Context variable for timeout tracking (not yet enforced) +# Enclosing `cf.suspense` timeout. Read by `select` when deciding how long +# to wait before suspending, and by stream subscriptions, where it also +# switches on cursor tracking so a resumed execution carries on rather than +# re-reading from the start. _timeout: contextvars.ContextVar[float | None] = contextvars.ContextVar( "_timeout", default=None ) @@ -126,6 +130,23 @@ def __init__(self, execution_id: str, working_dir: Path | None = None): # guarded by ``self._lock``. self._checkpoint_wire: dict[str, Any] = {} self._checkpoint_values: dict[str, Any] = {} + # Writes held back until the execution is somewhere it could resume + # from. ``_checkpoint_holds`` counts the non-replayable inputs + # currently in the body's hands (see ``hold_checkpoints``); while it + # is non-zero, writes accumulate here instead of going on the wire. + # A name is either set or reset, never both — the later write + # replaces the earlier one, so the delta only describes the net + # effect, exactly as the worker-side throttle coalesces them. + self._checkpoint_holds = 0 + self._checkpoint_pending_set: dict[str, Any] = {} + self._checkpoint_pending_reset: set[str] = set() + # Occurrence counts for auto-named stream cursors, so two loops + # over an identical view of the same stream get distinct + # checkpoints. Keyed by the content-addressed base name; the + # count is deterministic across attempts as long as subscriptions + # are opened in the same order, which is the determinism suspend + # already requires. + self._cursor_occurrences: dict[str, int] = {} def set_default_streams(self, streams: Streams | None) -> None: """Record the decorator's stream config so ``cf.stream(...)`` can @@ -624,25 +645,136 @@ 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)} - ) + # Serialised here rather than at publication time: the value is the + # one the caller passed, and holding a reference to a mutable object + # would record whatever it became later instead. + self._record_checkpoint_delta(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]) + self._record_checkpoint_delta(reset=[name]) + + def hold_checkpoints(self) -> None: + """Note that a non-replayable input is in the body's hands. + + Checkpoint state is one snapshot of a step's progress rather than a + set of independent cells — the server stores an execution's row-set + as a complete snapshot and applies each delta in a transaction. What + decides whether that snapshot is *coherent* is where the deltas get + cut, and one cut at an arbitrary point describes a state the + execution was never in. + + That only matters for state derived from something a replay can't + re-read. A result resolves again; a stream item does not — its + position lives in a cursor, and if the cursor and whatever the body + derived from the item reach the server separately, a crash in + between leaves the successor counting on from a position it never + actually reached. + + So writes made while an item is in hand are held, and published in + the same delta as the cursor advance that retires it. The pair moves + together or not at all, and a replay repeats whole items rather than + fractions of one. + + Balanced by ``release_checkpoints``. Nested holds — a body iterating + two streams — publish at the outermost release, the only point at + which every cursor involved is up to date. + """ + with self._lock: + self._checkpoint_holds += 1 + + def release_checkpoints(self, *, publish: bool) -> None: + """Retire a hold taken by ``hold_checkpoints``. + + ``publish`` says whether the item the hold covered was consumed. On + the way out of a completed iteration it is true, and the held writes + go out with the cursor advance. Where the item is abandoned instead + — ``break``, an exception, a dropped iterator — it is false: the + cursor was never advanced, so the item will be delivered again, and + anything derived from it must not be recorded or the replay counts + it twice. + + Discarding drops the whole pending delta, including writes made + under an enclosing hold. That is not over-eager: an enclosing hold + means that iteration has not advanced its own cursor either, so + everything pending derives from an item that is still unconsumed. + """ + with self._lock: + if not self._checkpoint_holds: + return + self._checkpoint_holds -= 1 + held = self._checkpoint_holds + if not publish: + self._checkpoint_pending_set.clear() + self._checkpoint_pending_reset.clear() + if publish and not held: + self._publish_checkpoints() + + def _record_checkpoint_delta( + self, + set_: dict[str, Any] | None = None, + reset: list[str] | None = None, + ) -> None: + """Put a write on the wire, or hold it for the next safe point.""" + with self._lock: + if self._checkpoint_holds: + for name, value in (set_ or {}).items(): + self._checkpoint_pending_set[name] = value + self._checkpoint_pending_reset.discard(name) + for name in reset or []: + self._checkpoint_pending_reset.add(name) + self._checkpoint_pending_set.pop(name, None) + return + protocol.send_checkpoint_update(self.execution_id, set_=set_, reset=reset) + + def _publish_checkpoints(self) -> None: + """Send whatever is being held, as a single delta.""" + with self._lock: + set_ = self._checkpoint_pending_set + reset = self._checkpoint_pending_reset + self._checkpoint_pending_set = {} + self._checkpoint_pending_reset = set() + if set_ or reset: + protocol.send_checkpoint_update( + self.execution_id, + set_=set_ or None, + # Sorted only so the delta is deterministic; the server + # applies the whole thing at once either way. + reset=sorted(reset) or None, + ) def flush(self) -> None: - """Block until buffered state has reached the server.""" + """Block until buffered state has reached the server. + + Publishes anything currently held first. An explicit flush is the + caller declaring this point consistent, which is what makes it the + escape hatch for state that has to be durable before a side effect — + including inside a loop body, where the runtime would otherwise wait + for the iteration to end. The cursor advance is still to come at + that point, so a flush there deliberately records derived state + ahead of the position it came from. + """ + self._publish_checkpoints() 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: - """Suspend the current execution, optionally resuming after a delay.""" + self, + delay: float | dt.timedelta | dt.datetime | None = None, + stream_wait: tuple[str, int] | None = None, + ) -> NoReturn: + """Signal that this execution should suspend. + + Raises rather than performing the handshake here. The server + records a suspension as a completion, and a completed execution's + checkpoint writes are rejected — so everything that runs while the + body unwinds (``finally`` blocks, cancelled tasks, generator + cleanup) has to happen *before* the request is sent, or its state + is silently dropped. ``finish_suspension`` completes it once the + body is done. + """ execute_after = None if isinstance(delay, dt.datetime): execute_after = int(delay.timestamp() * 1000) @@ -657,12 +789,84 @@ def suspend_execution( ).timestamp() * 1000 ) - request_id = protocol.request_suspend(self.execution_id, execute_after) + raise Suspending(execute_after, stream_wait) + + def finish_suspension( + self, + execute_after: int | None, + stream_wait: tuple[str, int] | None = None, + ) -> None: + """Complete a suspension once the body has unwound. Never returns. + + Stops any in-flight stream producers and joins their driver threads + first, so their cleanup runs while the execution is still live. + + Winding the generators down does *not* close their streams: the + driver skips ``send_stream_close`` on ``GeneratorExit``, and the + server leaves a suspended execution's streams paused rather than + closing them, so the execution that resumes the step continues + them. + + ``stream_wait`` gates the successor on a stream reaching a + sequence, for a consumer that suspended partway through iterating. + """ + try: + self.close_streams() + self.wait_streams() + except Exception: # noqa: BLE001, S110 + # Best-effort teardown — the suspension below is what matters. + pass + request_id = protocol.request_suspend( + self.execution_id, execute_after, stream_wait + ) self._wait_response(request_id) # Suspension confirmed. Block until the server aborts this execution. get_dispatcher().wait_closed() raise SystemExit(0) + def take_stream_suspension( + self, + ) -> tuple[int | None, tuple[str, int] | None] | None: + """Claim a suspension requested from inside a generator body. + + Returns ``(execute_after, stream_wait)`` — either of which may + itself be ``None`` — or ``None`` when no generator asked to + suspend. The executor checks this after its streams have drained. + """ + return self._stream_driver.take_suspension() + + def stream_available(self, stream_id: str, sequence: int) -> bool: + """Whether ``stream_id`` has reached ``sequence``, or has closed. + + The consumer can't answer this itself. Its queue is fed + asynchronously, so an empty one means "nothing has arrived yet", + not "the stream has nothing" — checking locally right after + subscribing always finds it empty, whatever the stream holds. + + A poll, never a suspension: the server reports what it knows and + the decision of what to do about it stays here, so a suspension + still unwinds the body before the handshake. + """ + request_id = protocol.request_select( + self.execution_id, + [{"type": "stream", "id": stream_id, "sequence": sequence}], + timeout_ms=0, + suspend=False, + ) + return self._wait_response(request_id) is not None + + @property + def suspense_timeout(self) -> float | None: + """The enclosing ``cf.suspense`` timeout, or ``None`` outside one.""" + return _timeout.get() + + def next_cursor_occurrence(self, name: str) -> int: + """Count of prior subscriptions in this execution sharing ``name``.""" + with self._lock: + occurrence = self._cursor_occurrences.get(name, 0) + self._cursor_occurrences[name] = occurrence + 1 + return occurrence + def _parse_response(self, msg: dict) -> Any: """Extract the result from a response message, raising on error.""" if msg.get("error"): diff --git a/adapters/python/coflux/errors.py b/adapters/python/coflux/errors.py index 804a3a73..d432942b 100644 --- a/adapters/python/coflux/errors.py +++ b/adapters/python/coflux/errors.py @@ -5,6 +5,37 @@ import importlib +class Suspending(BaseException): + """Internal control-flow signal: this execution wants to suspend. + + Raised at the suspend point — ``cf.suspend()``, or a wait expiring + inside a ``cf.suspense`` scope — and caught by the executor, or by the + stream driver when the suspend point is inside a generator body. + + Derives from ``BaseException`` so a bare ``except Exception`` in user + code can't swallow a suspension, the same protection ``SystemExit`` + provided before. + + The handshake with the server deliberately happens *after* this has + unwound the body rather than at the suspend point. The server records + a suspension as a completion, and a completed execution's checkpoint + writes are rejected, so anything running during unwinding — ``finally`` + blocks, cancelled tasks — has to happen while the execution is still + live. See ``ExecutorContext.finish_suspension``. + """ + + def __init__( + self, + execute_after: int | None = None, + stream_wait: tuple[str, int] | None = None, + ): + self.execute_after = execute_after + # (stream_id, sequence) when the suspension is waiting on a stream + # rather than the clock. + self.stream_wait = stream_wait + super().__init__("execution suspending") + + class ExecutionError(Exception): """Raised when a child execution failed. diff --git a/adapters/python/coflux/executor.py b/adapters/python/coflux/executor.py index 8dcac6fd..b8f42d1f 100644 --- a/adapters/python/coflux/executor.py +++ b/adapters/python/coflux/executor.py @@ -14,6 +14,7 @@ from . import protocol from .context import ExecutorContext from .dispatcher import start_dispatcher +from .errors import Suspending from .models import Input from .output import capture_output from .serialization import deserialize_value, serialize_value @@ -194,6 +195,23 @@ def execute_target( # serialised by Protocol._write_lock. ctx.wait_streams() + # A generator body may have asked to suspend while draining. The + # driver only records the request and stops its siblings, so that + # every generator's cleanup runs before the execution is finalised + # — the handshake happens here, on the executor thread. + pending = ctx.take_stream_suspension() + if pending is not None: + execute_after, stream_wait = pending + ctx.finish_suspension(execute_after, stream_wait) + + except Suspending as suspending: + # Raised at the suspend point and allowed to unwind the whole body + # first, so `finally` blocks and cancelled tasks run while the + # execution is still live and its checkpoint writes still land. + if ctx is None: + raise + ctx.finish_suspension(suspending.execute_after, suspending.stream_wait) + except Exception as e: # noqa: BLE001 # Any failure in user code is reported back to the server as an # execution error rather than crashing the adapter. diff --git a/adapters/python/coflux/protocol.py b/adapters/python/coflux/protocol.py index 7072fafa..23932b38 100644 --- a/adapters/python/coflux/protocol.py +++ b/adapters/python/coflux/protocol.py @@ -340,11 +340,24 @@ def request_upload_blob( ) -def request_suspend(execution_id: str, execute_after: int | None = None) -> int: - """Request to suspend execution.""" +def request_suspend( + execution_id: str, + execute_after: int | None = None, + stream_wait: tuple[str, int] | None = None, +) -> int: + """Request to suspend execution. + + ``stream_wait`` is a ``(stream_id, sequence)`` pair for a consumer that + suspended partway through iterating: the successor is held until the + stream reaches that sequence, or closes. Without it the successor is + scheduled on ``execute_after`` alone. + """ params: dict[str, Any] = {"execution_id": execution_id} if execute_after is not None: params["execute_after"] = execute_after + if stream_wait is not None: + stream_id, sequence = stream_wait + params["stream_wait"] = {"stream_id": stream_id, "sequence": sequence} return get_protocol().send_request("suspend", params) diff --git a/adapters/python/coflux/streams.py b/adapters/python/coflux/streams.py index e7f121fb..4b59af06 100644 --- a/adapters/python/coflux/streams.py +++ b/adapters/python/coflux/streams.py @@ -43,8 +43,8 @@ from . import protocol from .dispatcher import get_dispatcher -from .errors import raise_for_close -from .models import Stream +from .errors import Suspending, raise_for_close +from .models import Stream, Stride from .serialization import deserialize_value, serialize_value from .state import get_context from .target import Streams, _validate_buffer, _validate_timeout @@ -152,6 +152,12 @@ def __init__(self, execution_id: str) -> None: # index, then applied. self._pending_demand: dict[int, int] = {} self._closing = False + # Set when a generator body suspends. The handshake belongs on the + # executor thread, once every driver has wound down, so a driver + # only records the request and stops. A one-tuple, so an + # ``execute_after`` of ``None`` stays distinguishable from "no + # request". + self._suspend_request: tuple[int | None, tuple[str, int] | None] | None = None self._demand_handler_registered = False self._force_close_handler_registered = False # Indexes of streams the worker (CLI) has force-closed — typically @@ -348,6 +354,10 @@ def _is_force_closed(self, index: int) -> bool: with self._demand_cv: return index in self._force_closed + def _is_closing(self) -> bool: + with self._demand_cv: + return self._closing + def _run(self, index: int, generator: Any, start_sequence: int) -> None: """Run one sync generator to exhaustion (or error). @@ -382,11 +392,16 @@ def _run(self, index: int, generator: Any, start_sequence: int) -> None: # lifecycle closure on execution-end, or has already recorded # the force-close reason (e.g. "timeout"). return - except SystemExit: + except Suspending as suspending: # The generator body suspended (``cf.suspend()`` / implicit # suspense). The stream stays open — paused — for the # execution that resumes the step to continue, so nothing is # sent: a close here would end it for every consumer. + # + # The handshake happens on the executor thread once every + # driver has stopped, so that cleanup here (and in the other + # generators) lands before the execution is finalised. + self._record_suspension(suspending.execute_after, suspending.stream_wait) return except BaseException as e: # noqa: BLE001 - we propagate all if self._is_force_closed(index): @@ -402,7 +417,15 @@ def _run(self, index: int, generator: Any, start_sequence: int) -> None: traceback=tb, ) else: - if self._is_force_closed(index): + if self._is_force_closed(index) or self._is_closing(): + # Either the worker already recorded a close, or the driver + # is shutting down — in which case this StopIteration came + # from ``close_all`` closing the generator out from under + # us between the demand check and ``next()``, not from the + # body finishing. Reporting a normal close here would end + # the stream for every consumer, including the suspend + # case, where it has to stay paused for the execution that + # resumes the step. return protocol.send_stream_close(self._execution_id, index) @@ -444,8 +467,9 @@ async def iterate() -> None: loop.run_until_complete(iterate()) except (GeneratorExit, asyncio.CancelledError): return - except SystemExit: + except Suspending as suspending: # Suspended from inside the generator — see ``_run``. + self._record_suspension(suspending.execute_after, suspending.stream_wait) return except BaseException as e: # noqa: BLE001 - we propagate all if self._is_force_closed(index): @@ -460,7 +484,15 @@ async def iterate() -> None: traceback=tb, ) else: - if self._is_force_closed(index): + if self._is_force_closed(index) or self._is_closing(): + # Either the worker already recorded a close, or the driver + # is shutting down — in which case this StopIteration came + # from ``close_all`` closing the generator out from under + # us between the demand check and ``next()``, not from the + # body finishing. Reporting a normal close here would end + # the stream for every consumer, including the suspend + # case, where it has to stay paused for the execution that + # resumes the step. return protocol.send_stream_close(self._execution_id, index) finally: @@ -482,6 +514,35 @@ def _record_loop(self, generator: Any, loop: asyncio.AbstractEventLoop) -> None: entry["loop"] = loop return + def _record_suspension( + self, + execute_after: int | None, + stream_wait: tuple[str, int] | None = None, + ) -> None: + """Note that a generator body asked to suspend, and stop the rest. + + Closing the siblings is what lets the executor's ``wait_all()`` + return: they would otherwise keep producing, and nothing else is + going to interrupt them. Their streams stay open — paused — because + the driver skips ``send_stream_close`` on ``GeneratorExit`` and the + server leaves a suspended execution's streams alone. + + First request wins; a second generator suspending while we tear + down is the same suspension. + """ + with self._lock: + if self._suspend_request is None: + self._suspend_request = (execute_after, stream_wait) + # Outside the lock: close_all takes both _demand_cv and _lock. + self.close_all() + + def take_suspension(self) -> tuple[int | None, tuple[str, int] | None] | None: + """Claim any recorded suspension request. See ``_record_suspension``.""" + with self._lock: + request = self._suspend_request + self._suspend_request = None + return request + def wait_all(self) -> None: """Block until every worker thread has finished.""" with self._lock: @@ -561,6 +622,71 @@ def __init__(self, reason: str, error: dict[str, Any] | None) -> None: _ACK_BATCH = max(1, _PREFETCH // 2) +_CURSOR_PREFIX = "_cursor" + + +class _Resume: + """Cursor state for a subscription that is allowed to suspend. + + Only created inside a ``cf.suspense`` scope; outside one there is no + cursor and iteration blocks for as long as it takes, exactly as it + always has. + + ``position`` counts items of *this view*, not raw sequences — a + partition consumer counts its own items — which is what makes + resuming a plain slice on the same view. + """ + + __slots__ = ("name", "position", "stream_id", "stride", "timeout") + + def __init__( + self, + name: str, + position: int, + timeout: float, + stream_id: str, + stride: Stride, + ) -> None: + self.name = name + self.position = position + self.timeout = timeout + self.stream_id = stream_id + # The view's *original* stride, so the mapping from position to + # sequence stays fixed as the position advances. + self.stride = stride + + def next_sequence(self) -> int: + """Absolute sequence of the item this view wants next. + + A composed stride is a linear map, so item k of any view sits at + ``start + k*step``. That's both where a resumed subscription starts + and what the server gates a suspension on — no stride needed at the + far end, just the one number. + """ + start, _stop, step = self.stride + return start + self.position * step + + +def _cursor_name(ctx: Any, stream_id: str, stride: Stride) -> str: + """Content-addressed checkpoint name for a view's cursor. + + Derived from what is being consumed rather than from call order, so it + survives branching, and a producer re-run — which opens a new stream, + hence a new id — starts a fresh cursor rather than resuming a + brand-new stream at a stale offset. The stride is part of the name + because two partitions of one stream are different views with + different positions. + + Two loops over an *identical* view in one body are the one case + content addressing can't separate, so they fall back to an occurrence + counter. + """ + start, stop, step = stride + name = f"{_CURSOR_PREFIX}/{stream_id}/{start},{stop},{step}" + occurrence = ctx.next_cursor_occurrence(name) + return name if not occurrence else f"{name}#{occurrence}" + + class _Subscription: """Bookkeeping shared by the sync and async subscription readers. @@ -581,10 +707,19 @@ def __init__(self, subscription_id: int, execution_id: str) -> None: self._acked_sequence = -1 # Sequence of the item handed to the caller by the previous # ``__next__`` and not yet retired. It only counts as processed - # once the caller comes back for another one. + # once the caller comes back for another one, and where there is a + # cursor it also holds the execution's checkpoint writes back — + # see ``_take_in_hand``. self._in_hand: int | None = None # Retired items not yet reported. self._unreported = 0 + # Cursor state, when this subscription was opened somewhere it is + # allowed to suspend. ``None`` everywhere else. + self._resume: _Resume | None = None + + def attach_resume(self, resume: _Resume | None) -> None: + """Give this subscription its cursor. Called before subscribing.""" + self._resume = resume def on_items(self, items: list[list[Any]]) -> None: """Called by the registry when the server pushes items for this @@ -623,6 +758,17 @@ def close(self) -> None: if self._done: return self._done = True + if self._in_hand is not None and self._resume is not None: + # Abandoned rather than consumed: the cursor is deliberately not + # advanced (see the ack note above), so the item will be + # delivered again and anything the body derived from it must not + # be recorded, or the replay counts it twice. + try: + get_context().release_checkpoints(publish=False) + except Exception: # noqa: BLE001, S110 + # `close` also runs from `__del__`, by which point the + # execution context may be gone. Nothing to publish then. + pass self._in_hand = None self._unreported = 0 _stream_registry().drop(self._subscription_id) @@ -649,6 +795,20 @@ def __del__(self) -> None: # `__del__` must never raise. pass + def _take_in_hand(self, sequence: int) -> None: + """Hand an item to the caller, pending its retirement. + + Where this subscription keeps a cursor, it also holds the + execution's checkpoint writes (``hold_checkpoints``) until the item + retires, so that whatever the loop body derives from it is + published in the same delta as the cursor advance that consumes it. + Outside a suspense scope there is no cursor, no recorded position + for derived state to disagree with, and so nothing to hold. + """ + self._in_hand = sequence + if self._resume is not None: + get_context().hold_checkpoints() + def _retire_in_hand(self) -> None: """Count the previously-yielded item as processed. @@ -663,9 +823,39 @@ def _retire_in_hand(self) -> None: self._acked_sequence = max(self._acked_sequence, self._in_hand) self._in_hand = None self._unreported += 1 + if self._resume is not None: + # Same boundary the acknowledgement uses: the item counts as + # processed once the caller comes back for another. A + # suspension flushes checkpoints before it is recorded, so a + # pause never loses or repeats an item, and a crash before this + # point replays the whole item — the advance and whatever the + # body derived from it are published together, below. + self._resume.position += 1 + get_context().checkpoint_set(self._resume.name, self._resume.position) + # Releases the hold taken in ``_take_in_hand``, publishing the + # cursor advance and whatever the loop body wrote while holding + # the item as one delta. + get_context().release_checkpoints(publish=True) if self._unreported >= _ACK_BATCH: self._flush_ack() + def _suspend(self, resume: _Resume) -> None: + """Give up the worker slot until there is more to read. Never returns. + + The subscription is released first. An abandoned one keeps holding + the producer's backpressure watermark down until the server tears + it down, and this execution is about to go away regardless. + + No delay: the successor is gated on the stream instead, and the + server releases it when the item lands or the stream closes. Every + way an execution can end closes its open streams, so a producer + that dies still wakes us rather than leaving us gated forever. + """ + self.close() + get_context().suspend_execution( + None, stream_wait=(resume.stream_id, resume.next_sequence()) + ) + def _flush_ack(self) -> None: if self._unreported == 0: return @@ -712,6 +902,31 @@ def on_items(self, items: list[list[Any]]) -> None: def on_closed(self, reason: str, error: dict[str, Any] | None) -> None: self._queue.put(_Closed(reason, error)) + def _wait_for_item(self) -> Any: + """Block for the next item, suspending instead if that's allowed. + + Outside a ``cf.suspense`` scope this waits indefinitely, as it + always has. Inside one, a gap longer than the scope's timeout means + the execution gives up its worker slot rather than holding it + through the wait. + """ + resume = self._resume + if resume is None: + return self._queue.get() + if resume.timeout: + try: + return self._queue.get(timeout=resume.timeout) + except queue.Empty: + pass + # An empty queue says nothing about the stream — items arrive + # asynchronously — so the server decides, exactly as it does for a + # result. If the item is already there its push is in flight, and + # blocking for it is bounded: a dispatcher EOF wakes every + # iterator with a synthetic close. + if get_context().stream_available(resume.stream_id, resume.next_sequence()): + return self._queue.get() + self._suspend(resume) + def __iter__(self) -> _StreamIterator: return self @@ -734,7 +949,7 @@ def __next__(self) -> Any: # the producer may be waiting on precisely the acknowledgement # we're batching, so holding it back would deadlock. self._flush_ack() - item = self._queue.get() + item = self._wait_for_item() if isinstance(item, _Closed): # Same release path as an early exit — the server has already @@ -745,7 +960,7 @@ def __next__(self) -> Any: raise StopIteration sequence, value = item - self._in_hand = sequence + self._take_in_hand(sequence) return deserialize_value(value) @@ -808,6 +1023,25 @@ async def aclose(self) -> None: """ self.close() + async def _wait_for_item(self) -> Any: + """Await the next item, suspending instead if that's allowed. + + The sync path's counterpart. Cancelling a ``Queue.get`` leaves the + item in the queue for the next reader, so timing out here doesn't + drop anything. + """ + resume = self._resume + if resume is None: + return await self._queue.get() + if resume.timeout: + try: + return await asyncio.wait_for(self._queue.get(), resume.timeout) + except asyncio.TimeoutError: + pass + if get_context().stream_available(resume.stream_id, resume.next_sequence()): + return await self._queue.get() + self._suspend(resume) + async def __aenter__(self) -> _AsyncStreamIterator: return self @@ -825,7 +1059,7 @@ async def __anext__(self) -> Any: # the sync path does: a lockstep producer may be waiting on # precisely the acknowledgement we're batching. self._flush_ack() - item = await self._queue.get() + item = await self._wait_for_item() if isinstance(item, _Closed): self.close() @@ -833,7 +1067,7 @@ async def __anext__(self) -> Any: raise StopAsyncIteration sequence, value = item - self._in_hand = sequence + self._take_in_hand(sequence) return deserialize_value(value) @@ -956,9 +1190,28 @@ def _open_subscription( """ ctx = get_context() execution_id = ctx.execution_id + + # A cursor is kept only where iteration is allowed to suspend. Outside + # a suspense scope nothing is written and behaviour is unchanged. + resume: _Resume | None = None + timeout = ctx.suspense_timeout + if timeout is not None: + name = _cursor_name(ctx, stream_id, stride) + try: + position = ctx.checkpoint_get(name) + except KeyError: + position = 0 + resume = _Resume(name, position, timeout, stream_id, stride) + subscription_id, iterator = _stream_registry().allocate(execution_id, factory) + iterator.attach_resume(resume) start, stop, step = stride + if resume is not None and resume.position: + # Resume where the last attempt stopped, so the server begins + # delivery there rather than the consumer reading a backlog in + # order to discard it. + start = resume.next_sequence() wire_stride = {"start": start, "stop": stop, "step": step} protocol.send_stream_subscribe( diff --git a/adapters/python/coflux/target.py b/adapters/python/coflux/target.py index 7a59876d..d74b771a 100644 --- a/adapters/python/coflux/target.py +++ b/adapters/python/coflux/target.py @@ -571,6 +571,23 @@ def definition(self) -> TargetDefinition: def fn(self) -> t.Callable[P, T]: return self._fn + def _bind_arguments(self, args: tuple[t.Any, ...]) -> tuple[t.Any, ...]: + """Fill in the defaults for arguments the caller left out. + + Arguments are always sent in full, so that two call sites differing + only in whether they spelled a default out submit the same list — + otherwise they'd derive different cache, memo and defer keys for what + is the same call. Binding against the signature uses the real default + objects, so nothing is lost the way it would be by reconstructing them + from the manifest's JSON. + + Every parameter is positional-or-keyword (enforced when the target is + defined), so the bound arguments are always a plain positional tuple. + """ + bound = inspect.signature(self._fn).bind(*args) + bound.apply_defaults() + return bound.args + def submit(self, *args: P.args, **kwargs: P.kwargs) -> Execution[T]: """Submit this target for execution and return a handle.""" if kwargs: @@ -582,7 +599,7 @@ def submit(self, *args: P.args, **kwargs: P.kwargs) -> Execution[T]: # been registered via cf.stream(...) — the caller becomes the # producer, the callee gets a Stream handle. Bare generators # raise; the user should wrap them explicitly. - serialized_args = [serialize_value(arg) for arg in args] + serialized_args = [serialize_value(arg) for arg in self._bind_arguments(args)] # Use only the declared wait_for from the decorator wait_for_val = ( diff --git a/adapters/python/tests/test_argument_binding.py b/adapters/python/tests/test_argument_binding.py new file mode 100644 index 00000000..84fa8e16 --- /dev/null +++ b/adapters/python/tests/test_argument_binding.py @@ -0,0 +1,92 @@ +"""Defaults are filled in at the call site, not left to the callee. + +Cache, memo and defer keys are derived from the arguments a call was +submitted with, so ``expensive(5)`` and ``expensive(5, 1)`` — the same call, +one of them spelling out the default — have to reach the server as the same +list or they key differently and never share a cached result. + +Binding against the signature rather than reconstructing defaults from the +manifest also keeps the real objects: a tuple default stays a tuple, where a +JSON round-trip would have made it a list. +""" + +from __future__ import annotations + +import pytest + +import coflux as cf + + +def bind(target, *args): + return target._bind_arguments(args) + + +def test_omitted_defaults_are_filled_in(): + @cf.task() + def expensive(x: int, y: int = 1, z: int = 2) -> int: + return x + y + z + + assert bind(expensive, 5) == (5, 1, 2) + + +def test_omitted_and_explicit_defaults_agree(): + @cf.task() + def expensive(x: int, y: int = 1, z: int = 2) -> int: + return x + y + z + + assert bind(expensive, 5) == bind(expensive, 5, 1, 2) + assert bind(expensive, 5, 1) == bind(expensive, 5, 1, 2) + + +def test_overridden_values_are_kept(): + @cf.task() + def expensive(x: int, y: int = 1, z: int = 2) -> int: + return x + y + z + + assert bind(expensive, 5, 9) == (5, 9, 2) + assert bind(expensive, 5) != bind(expensive, 5, 9) + + +def test_no_defaults_is_unchanged(): + @cf.task() + def add(x: int, y: int) -> int: + return x + y + + assert bind(add, 1, 2) == (1, 2) + + +def test_no_parameters_is_unchanged(): + @cf.task() + def tick() -> None: + return None + + assert bind(tick) == () + + +def test_default_objects_survive_intact(): + # Reconstructing this from the manifest's JSON would yield a list. + @cf.task() + def batched(x: int, shape: tuple = (1, 2)) -> int: + return x + + (_, shape) = bind(batched, 5) + assert shape == (1, 2) + assert isinstance(shape, tuple) + + +def test_too_many_arguments_is_rejected_at_submit(): + @cf.task() + def add(x: int, y: int) -> int: + return x + y + + with pytest.raises(TypeError, match="too many positional arguments"): + bind(add, 1, 2, 3) + + +def test_missing_required_argument_is_rejected_at_submit(): + @cf.task() + def add(x: int, y: int) -> int: + return x + y + + with pytest.raises(TypeError, match="missing a required argument"): + bind(add, 1) diff --git a/adapters/python/tests/test_checkpoint.py b/adapters/python/tests/test_checkpoint.py new file mode 100644 index 00000000..c13e6b6f --- /dev/null +++ b/adapters/python/tests/test_checkpoint.py @@ -0,0 +1,79 @@ +"""The reserved checkpoint namespace. + +Names starting with ``_`` belong to the adapter — currently the cursors +behind stream suspension, which are named after the stream and view being +read. Reserving the whole prefix rather than individual names means later +internal state doesn't need another round of this. + +The check lives in ``Checkpoint`` rather than in the context's +``checkpoint_get``/``checkpoint_set``/``checkpoint_reset``, because those +are the path the adapter's own cursors go through — validating there would +reject the very thing the prefix exists for. +""" + +from __future__ import annotations + +import pytest + +from coflux.checkpoint import RESERVED_PREFIX, Checkpoint + + +def test_ordinary_names_are_accepted(): + assert Checkpoint("cursor").name == "cursor" + assert Checkpoint("total", default=0.0).default == 0.0 + # Only the *leading* character is reserved. + assert Checkpoint("last_seen").name == "last_seen" + + +@pytest.mark.parametrize("name", ["_cursor", "_", "_cursor/Erun:2_0/0,None,1"]) +def test_reserved_names_are_rejected(name): + with pytest.raises(ValueError, match="reserved"): + Checkpoint(name) + + +def test_the_prefix_is_what_the_adapter_uses(): + """The cursor names the stream subscription generates have to be + exactly what this rejects, or the two could drift apart.""" + from coflux.streams import _CURSOR_PREFIX + + assert _CURSOR_PREFIX.startswith(RESERVED_PREFIX) + + +class _FakeContext: + """Just enough context for get/set — the storage semantics are the + server's business and are covered by the e2e suite.""" + + def __init__(self, values=None): + self.values = dict(values or {}) + + def checkpoint_get(self, name): + if name not in self.values: + raise KeyError(name) + return self.values[name] + + def checkpoint_set(self, name, value): + self.values[name] = value + + +@pytest.fixture +def context(monkeypatch): + ctx = _FakeContext() + monkeypatch.setattr("coflux.checkpoint.get_context", lambda: ctx) + return ctx + + +def test_update_applies_to_the_current_value(context): + context.values["count"] = 4 + count = Checkpoint("count", default=0) + + assert count.update(lambda n: n + 1) == 5 + assert context.values["count"] == 5 + + +def test_update_starts_from_the_default_when_unset(context): + """`fn` sees what `get()` would have returned, so an unset checkpoint + doesn't need special-casing at the call site.""" + count = Checkpoint("count", default=10) + + assert count.update(lambda n: n + 1) == 11 + assert context.values["count"] == 11 diff --git a/adapters/python/tests/test_checkpoint_publication.py b/adapters/python/tests/test_checkpoint_publication.py new file mode 100644 index 00000000..d659b337 --- /dev/null +++ b/adapters/python/tests/test_checkpoint_publication.py @@ -0,0 +1,177 @@ +"""Where checkpoint deltas get cut. + +Checkpoint state is one snapshot of a step's progress rather than a set of +independent cells: the server stores an execution's row-set as a complete +snapshot and applies each delta in a transaction. So what decides whether +the successor of a crashed or suspended execution reads a *coherent* state +is where the deltas are cut — and one cut at an arbitrary point describes a +state the execution was never in. + +That only matters for state derived from something a replay can't re-read. +A result resolves again; a stream item doesn't, so its position is recorded +in a cursor, and the cursor and whatever the body derived from the item have +to reach the server together. The context holds writes back while such an +item is in hand and publishes them with the advance that consumes it. + +These cover the context's half of that — the holding, the coalescing, and +what becomes of held writes when the item is abandoned instead. Which +subscription takes a hold, and when, is in ``test_stream_iteration.py``. +""" + +from __future__ import annotations + +import pytest + +from coflux import protocol +from coflux.context import ExecutorContext + +CURSOR = "_cursor/Eproducer_0/0,None,1" + + +class _Wire: + """Records the deltas reaching the worker, values unwrapped.""" + + def __init__(self): + self.deltas = [] + + def send(self, execution_id, set_=None, reset=None): + self.deltas.append( + ( + {name: value["value"] for name, value in (set_ or {}).items()}, + list(reset or []), + ) + ) + + +@pytest.fixture +def wire(monkeypatch): + w = _Wire() + monkeypatch.setattr(protocol, "send_checkpoint_update", w.send) + return w + + +@pytest.fixture +def ctx(): + return ExecutorContext("Econsumer") + + +def test_writes_go_out_immediately_when_nothing_is_held(ctx, wire): + """The ordinary case — a body not consuming a stream is always at a + point it could resume from, so nothing is deferred.""" + ctx.checkpoint_set("cursor", 4) + ctx.checkpoint_set("cursor", 5) + + assert wire.deltas == [({"cursor": 4}, []), ({"cursor": 5}, [])] + + +def test_a_hold_publishes_one_delta_at_the_release(ctx, wire): + """The whole point: derived state and the cursor advance land together + or not at all.""" + ctx.hold_checkpoints() + ctx.checkpoint_set("total", 7) + ctx.checkpoint_set(CURSOR, 1) + assert wire.deltas == [] + + ctx.release_checkpoints(publish=True) + + assert wire.deltas == [({"total": 7, CURSOR: 1}, [])] + + +def test_the_last_write_to_a_name_wins(ctx, wire): + """Held writes coalesce, the way the worker-side throttle does: the + delta describes the net effect, not the history.""" + ctx.hold_checkpoints() + ctx.checkpoint_set("total", 1) + ctx.checkpoint_set("total", 2) + ctx.checkpoint_reset("total") + ctx.checkpoint_set("total", 3) + ctx.release_checkpoints(publish=True) + + assert wire.deltas == [({"total": 3}, [])] + + +def test_a_reset_supersedes_a_held_write(ctx, wire): + """A name is either set or reset in the delta, never both.""" + ctx.hold_checkpoints() + ctx.checkpoint_set("total", 1) + ctx.checkpoint_reset("total") + ctx.release_checkpoints(publish=True) + + assert wire.deltas == [({}, ["total"])] + + +def test_an_abandoned_item_discards_what_the_body_derived(ctx, wire): + """Nothing is recorded, because the item it came from was never + consumed and will be delivered again.""" + ctx.hold_checkpoints() + ctx.checkpoint_set("total", 7) + ctx.release_checkpoints(publish=False) + + assert wire.deltas == [] + # The execution still sees its own writes; only the record is dropped. + assert ctx.checkpoint_get("total") == 7 + + +def test_nested_holds_publish_at_the_outermost_release(ctx, wire): + """A body iterating two streams is only somewhere it could resume from + once every cursor involved is up to date.""" + ctx.hold_checkpoints() + ctx.checkpoint_set("outer", 1) + ctx.hold_checkpoints() + ctx.checkpoint_set("inner", 2) + + ctx.release_checkpoints(publish=True) + assert wire.deltas == [] + + ctx.release_checkpoints(publish=True) + assert wire.deltas == [({"outer": 1, "inner": 2}, [])] + + +def test_an_abandoned_inner_item_discards_the_enclosing_writes_too(ctx, wire): + """Not over-eager: an enclosing hold means that iteration hasn't + advanced its own cursor either, so everything pending derives from an + item that is still unconsumed.""" + ctx.hold_checkpoints() + ctx.checkpoint_set("outer", 1) + ctx.hold_checkpoints() + ctx.checkpoint_set("inner", 2) + + ctx.release_checkpoints(publish=False) + ctx.release_checkpoints(publish=True) + + assert wire.deltas == [] + + +def test_writes_are_immediate_again_after_the_release(ctx, wire): + ctx.hold_checkpoints() + ctx.release_checkpoints(publish=True) + ctx.checkpoint_set("cursor", 1) + + assert wire.deltas == [({"cursor": 1}, [])] + + +def test_an_unbalanced_release_is_a_no_op(ctx, wire): + """The count can't go negative — one that did would leave every later + write held for the rest of the execution.""" + ctx.release_checkpoints(publish=True) + ctx.checkpoint_set("cursor", 1) + + assert wire.deltas == [({"cursor": 1}, [])] + + +def test_flush_publishes_what_is_held(ctx, wire, monkeypatch): + """An explicit flush is the caller declaring the point consistent — + the escape hatch for state that has to be durable before a side + effect, including inside a loop body.""" + monkeypatch.setattr(protocol, "request_flush", lambda execution_id: "R1") + monkeypatch.setattr(ctx, "_wait_response", lambda request_id: None) + + ctx.hold_checkpoints() + ctx.checkpoint_set("sent", True) + ctx.flush() + + assert wire.deltas == [({"sent": True}, [])] + + # The flush cut a delta; it didn't end the item. + ctx.checkpoint_set("total", 1) + assert len(wire.deltas) == 1 diff --git a/adapters/python/tests/test_stream_driver.py b/adapters/python/tests/test_stream_driver.py index ffbbfbbc..8473e647 100644 --- a/adapters/python/tests/test_stream_driver.py +++ b/adapters/python/tests/test_stream_driver.py @@ -14,6 +14,7 @@ import pytest from coflux import protocol, streams +from coflux.errors import Suspending from coflux.streams import StreamDriver @@ -129,16 +130,49 @@ def test_suspend_inside_the_generator_sends_no_close(harness): def gen(): yield "a" - # ``cf.suspend()`` ends the calling thread with SystemExit once the - # server has confirmed the suspension. The stream must be left - # paused for the resumed execution, not closed. - raise SystemExit(0) + # ``cf.suspend()`` raises ``Suspending`` at the suspend point; the + # handshake happens later, on the executor thread. The stream must + # be left paused for the resumed execution, not closed. + raise Suspending(1234) driver.register(gen(), buffer=None) driver.wait_all() assert h.appends == [(0, 0, "a")] assert h.closes == [] + # Recorded for the executor to complete once every driver has stopped. + assert driver.take_suspension() == (1234, None) + # Claimed exactly once. + assert driver.take_suspension() is None + + +def test_suspend_stops_sibling_generators(harness): + """One generator suspending winds the others down, so ``wait_all`` + returns rather than waiting on producers that would never stop.""" + h = harness({"id": "run1:2_0", "index": 0, "head": -1}) + driver = StreamDriver("run1:2:1") + + started = threading.Event() + + def sibling(): + yield "sibling" + started.set() + while True: + # Would run forever if nothing closed it. + yield "more" + + def suspending(): + started.wait(timeout=5) + raise Suspending(None) + yield # pragma: no cover - unreachable, makes this a generator + + driver.register(sibling(), buffer=None) + driver.register(suspending(), buffer=None) + driver.wait_all() + + assert driver.take_suspension() == (None, None) + # Neither stream was closed — both are paused for the resumed execution. + assert h.closes == [] def test_generator_error_closes_with_the_error(harness): diff --git a/adapters/python/tests/test_stream_iteration.py b/adapters/python/tests/test_stream_iteration.py index 3213cf4e..5692c87a 100644 --- a/adapters/python/tests/test_stream_iteration.py +++ b/adapters/python/tests/test_stream_iteration.py @@ -24,16 +24,70 @@ import asyncio import gc -from types import SimpleNamespace import pytest from coflux import protocol, streams +from coflux.errors import Suspending from coflux.models import Stream PRODUCER_STREAM_ID = "Eproducer_0" +class _FakeContext: + """Stands in for ``ExecutorContext`` — only what subscriptions reach for. + + ``suspense_timeout`` is ``None`` by default, which is the no-cursor + path: no checkpoint is written and iteration blocks indefinitely. The + resume tests set it to opt into cursor tracking. + """ + + def __init__(self): + self.execution_id = "Econsumer" + self.suspense_timeout = None + self.checkpoints = {} + self.suspended = [] + self.probes = [] + self.available = False + # Ordered log of everything the subscription does to checkpoint + # state, so tests can assert on the bracket itself. What the real + # context does with held writes is tested against the real context, + # in ``test_checkpoint_publication.py``. + self.events = [] + self._occurrences = {} + + def next_cursor_occurrence(self, name): + occurrence = self._occurrences.get(name, 0) + self._occurrences[name] = occurrence + 1 + return occurrence + + def checkpoint_get(self, name): + if name not in self.checkpoints: + raise KeyError(name) + return self.checkpoints[name] + + def checkpoint_set(self, name, value): + self.checkpoints[name] = value + self.events.append(("set", name, value)) + + def hold_checkpoints(self): + self.events.append(("hold",)) + + def release_checkpoints(self, *, publish): + self.events.append(("release", publish)) + + def stream_available(self, stream_id, sequence): + """Stands in for the server's answer. Tests set ``available`` to + say what it should report; the default is "nothing there", which + is what makes a consumer suspend.""" + self.probes.append((stream_id, sequence)) + return self.available + + def suspend_execution(self, delay=None, stream_wait=None): + self.suspended.append((delay, stream_wait)) + raise Suspending(None, stream_wait) + + class _FakeDispatcher: def __init__(self): self.closed = False @@ -60,8 +114,10 @@ def __init__(self, registry, dispatcher): self.registry = registry self.dispatcher = dispatcher self.subscribes = [] + self.strides = [] self.unsubscribes = [] self.acks = [] + self.context = _FakeContext() self._items = [] self._close = None @@ -81,6 +137,7 @@ def on_subscribe( stride=None, ): self.subscribes.append(subscription_id) + self.strides.append(stride) if self._items: self.registry._on_items( {"subscription_id": subscription_id, "items": list(self._items)} @@ -108,9 +165,7 @@ def harness(monkeypatch): monkeypatch.setattr(streams, "_registry_instance", registry) monkeypatch.setattr(streams, "get_dispatcher", lambda: dispatcher) - monkeypatch.setattr( - streams, "get_context", lambda: SimpleNamespace(execution_id="Econsumer") - ) + monkeypatch.setattr(streams, "get_context", lambda: h.context) # Values go on the wire as tagged envelopes; the lifecycle is what's # under test, so keep them opaque. monkeypatch.setattr(streams, "deserialize_value", lambda value: value) @@ -377,3 +432,267 @@ def test_async_iterator_requires_a_running_loop(harness): than failing somewhere further in.""" with pytest.raises(RuntimeError, match="running event loop"): Stream(PRODUCER_STREAM_ID).__aiter__() + + +# --- Resuming from a checkpoint cursor --------------------------------------- +# +# A subscription opened inside a `cf.suspense` scope keeps its position in +# an adapter-managed checkpoint, so the execution that resumes the step +# carries on rather than re-reading from sequence 0. Outside such a scope +# none of this engages and iteration behaves exactly as it always has. + +CURSOR = f"_cursor/{PRODUCER_STREAM_ID}/0,None,1" + + +def test_no_cursor_outside_a_suspense_scope(harness): + """The default path writes no checkpoint and never suspends.""" + harness.serve([[0, "a"], [1, "b"]], close="complete") + + assert list(Stream(PRODUCER_STREAM_ID)) == ["a", "b"] + assert harness.context.checkpoints == {} + assert harness.context.suspended == [] + + +def test_cursor_advances_as_items_are_consumed(harness): + harness.context.suspense_timeout = 30 + harness.serve([[0, "a"], [1, "b"], [2, "c"]], close="complete") + + assert list(Stream(PRODUCER_STREAM_ID)) == ["a", "b", "c"] + assert harness.context.checkpoints == {CURSOR: 3} + + +def test_cursor_lags_the_item_in_hand(harness): + """An item counts as consumed only once the caller comes back for the + next one — the same boundary the acknowledgement uses, so a suspension + mid-body replays that item rather than skipping it.""" + harness.context.suspense_timeout = 30 + harness.serve([[0, "a"], [1, "b"]]) + + iterator = iter(Stream(PRODUCER_STREAM_ID)) + assert next(iterator) == "a" + # "a" is in hand, not yet retired. + assert harness.context.checkpoints == {} + assert next(iterator) == "b" + assert harness.context.checkpoints == {CURSOR: 1} + + +def test_resume_subscribes_at_the_cursor(harness): + """The server starts delivery at the cursor, rather than the consumer + reading a backlog in order to discard it.""" + harness.context.suspense_timeout = 30 + harness.context.checkpoints[CURSOR] = 2 + harness.serve([[2, "c"]], close="complete") + + assert list(Stream(PRODUCER_STREAM_ID)) == ["c"] + assert harness.strides == [{"start": 2, "stop": None, "step": 1}] + # Counting continues from where it resumed. + assert harness.context.checkpoints == {CURSOR: 3} + + +def test_resume_of_a_partition_counts_its_own_items(harness): + """A partition consumer's cursor counts items of its view, not raw + sequences, so resuming is a plain slice on the same view.""" + harness.context.suspense_timeout = 30 + name = f"_cursor/{PRODUCER_STREAM_ID}/1,None,4" + harness.context.checkpoints[name] = 3 + harness.serve([], close="complete") + + assert list(Stream(PRODUCER_STREAM_ID).partition(4, 1)) == [] + # Item 3 of the view sits at sequence 1 + 3*4. + assert harness.strides == [{"start": 13, "stop": None, "step": 4}] + + +def test_identical_views_get_distinct_cursors(harness): + """Content addressing can't separate two loops over the same view, so + they fall back to an occurrence counter.""" + harness.context.suspense_timeout = 30 + harness.serve([[0, "a"]], close="complete") + + assert list(Stream(PRODUCER_STREAM_ID)) == ["a"] + assert list(Stream(PRODUCER_STREAM_ID)) == ["a"] + + assert harness.context.checkpoints == {CURSOR: 1, f"{CURSOR}#1": 1} + + +def test_idle_stream_suspends_and_releases_the_subscription(harness): + """Nothing arrives within the timeout, so the execution gives up its + worker slot — after unsubscribing, since an abandoned subscription + would pin the producer's backpressure watermark.""" + harness.context.suspense_timeout = 0.01 + harness.serve([]) + + with pytest.raises(Suspending): + list(Stream(PRODUCER_STREAM_ID)) + + # The local wait expired, then the server confirmed there was nothing. + assert harness.context.probes == [(PRODUCER_STREAM_ID, 0)] + # No delay — the successor is gated on the stream instead, and the + # server releases it when the next item lands or the stream closes. + assert harness.context.suspended == [(None, (PRODUCER_STREAM_ID, 0))] + assert harness.unsubscribes == harness.subscribes + + +def test_the_server_decides_whether_to_suspend(harness): + """An empty queue is not evidence that the stream is empty — items + arrive asynchronously, so a consumer that checked locally would + suspend before hearing anything, wake at once because the item had + been there all along, and repeat forever. The server is asked + instead, exactly as it is for a result.""" + harness.context.suspense_timeout = 0 + # Only the first item is delivered on subscribe; the second arrives + # after the server has been asked, which is the ordering that broke a + # consumer deciding for itself — it would have suspended here. + harness.serve([[0, "a"]]) + + def available(stream_id, sequence): + harness.context.probes.append((stream_id, sequence)) + subscription_id = harness.subscribes[-1] + harness.registry._on_items( + {"subscription_id": subscription_id, "items": [[1, "b"]]} + ) + harness.registry._on_closed( + {"subscription_id": subscription_id, "reason": "complete"} + ) + return True + + harness.context.stream_available = available + + assert list(Stream(PRODUCER_STREAM_ID)) == ["a", "b"] + assert harness.context.probes == [(PRODUCER_STREAM_ID, 1)] + assert harness.context.suspended == [] + + +def test_zero_timeout_suspends_once_caught_up(harness): + """With nothing left, the server says so and the consumer suspends — + once, gated on the sequence it is waiting for.""" + harness.context.suspense_timeout = 0 + harness.context.available = False + harness.serve([[0, "a"]]) + + iterator = iter(Stream(PRODUCER_STREAM_ID)) + assert next(iterator) == "a" + with pytest.raises(Suspending): + next(iterator) + + assert harness.context.probes == [(PRODUCER_STREAM_ID, 1)] + assert harness.context.suspended == [(None, (PRODUCER_STREAM_ID, 1))] + assert harness.context.checkpoints == {CURSOR: 1} + + +# --- holding writes while an item is in hand --------------------------------- +# +# Where there's a cursor, whatever the loop body derives from an item has to +# reach the server in the same delta as the cursor advance that consumes it, +# or a crash between the two leaves the successor counting on from a position +# it never reached. These cover the iterator's half of that — when it takes a +# hold and how it releases one. + + +def test_body_writes_are_held_until_the_item_retires(harness): + """The hold spans the loop body, and is released by the retire that + writes the cursor — so the pair leaves as one delta.""" + harness.context.suspense_timeout = 30 + harness.serve([[0, "a"], [1, "b"]], close="complete") + + for total, _value in enumerate(Stream(PRODUCER_STREAM_ID), start=1): + harness.context.checkpoint_set("total", total) + + assert harness.context.events == [ + ("hold",), + ("set", "total", 1), + ("set", CURSOR, 1), + ("release", True), + ("hold",), + ("set", "total", 2), + ("set", CURSOR, 2), + ("release", True), + ] + + +def test_async_iteration_holds_the_same_way(harness): + """`async for` shares the accounting, including the bracket.""" + harness.context.suspense_timeout = 30 + harness.serve([[0, "a"]], close="complete") + + async def consume(): + async for _value in Stream(PRODUCER_STREAM_ID): + harness.context.checkpoint_set("total", 1) + + asyncio.run(consume()) + + assert harness.context.events == [ + ("hold",), + ("set", "total", 1), + ("set", CURSOR, 1), + ("release", True), + ] + + +def test_breaking_mid_item_releases_without_publishing(harness): + """The cursor is deliberately not advanced for an abandoned item, so + the item will be delivered again — and anything derived from it must + not be recorded, or the replay counts it twice.""" + harness.context.suspense_timeout = 30 + harness.serve([[0, "a"], [1, "b"]]) + + for _value in Stream(PRODUCER_STREAM_ID): + harness.context.checkpoint_set("total", 1) + break + gc.collect() + + assert harness.context.events == [ + ("hold",), + ("set", "total", 1), + ("release", False), + ] + assert CURSOR not in harness.context.checkpoints + + +def test_an_exception_in_the_body_discards_the_same_way(harness): + """Unwinding through the loop body abandons the item just as `break` + does; the retry re-reads it.""" + harness.context.suspense_timeout = 30 + harness.serve([[0, "a"]]) + + with pytest.raises(RuntimeError): + for _value in Stream(PRODUCER_STREAM_ID): + harness.context.checkpoint_set("total", 1) + raise RuntimeError("boom") + gc.collect() + + assert ("release", False) in harness.context.events + assert CURSOR not in harness.context.checkpoints + + +def test_suspending_leaves_nothing_held(harness): + """The suspension happens after the retire, so the cursor advance and + the body's writes are published before the handshake goes out. That's + what makes a pause neither lose nor repeat an item.""" + harness.context.suspense_timeout = 0 + harness.context.available = False + harness.serve([[0, "a"]]) + + iterator = iter(Stream(PRODUCER_STREAM_ID)) + assert next(iterator) == "a" + harness.context.checkpoint_set("total", 1) + with pytest.raises(Suspending): + next(iterator) + + assert harness.context.events == [ + ("hold",), + ("set", "total", 1), + ("set", CURSOR, 1), + ("release", True), + ] + + +def test_nothing_is_held_without_a_cursor(harness): + """Outside a suspense scope there's no recorded position for derived + state to disagree with, so iteration doesn't take a hold at all and + writes go out as they always have.""" + harness.serve([[0, "a"], [1, "b"]], close="complete") + + for _value in Stream(PRODUCER_STREAM_ID): + harness.context.checkpoint_set("total", 1) + + assert [event for event in harness.context.events if event[0] != "set"] == [] diff --git a/cli/cmd/coflux/submit.go b/cli/cmd/coflux/submit.go index 014cb8f2..267c25cc 100644 --- a/cli/cmd/coflux/submit.go +++ b/cli/cmd/coflux/submit.go @@ -59,6 +59,36 @@ func init() { submitCmd.MarkFlagsMutuallyExclusive("memo", "no-memo") } +// padArguments fills in the defaults for arguments that weren't given on the +// command line. +// +// Arguments are positional, so a submission that omits its trailing defaults +// and one that spells them out would otherwise derive different cache, memo +// and defer keys for the same call. Adapters bind their own defaults against +// the signature before submitting; this is the equivalent here, using the +// defaults the workflow was registered with (already JSON-encoded). +// +// Padding stops at a parameter with no default, since nothing after it could +// be positioned without it. +func padArguments(arguments [][]any, parameters any) [][]any { + params, ok := parameters.([]any) + if !ok { + return arguments + } + for _, p := range params[min(len(arguments), len(params)):] { + param, ok := p.(map[string]any) + if !ok { + return arguments + } + def, ok := param["default"].(string) + if !ok { + return arguments + } + arguments = append(arguments, []any{"json", def}) + } + return arguments +} + func runSubmit(cmd *cobra.Command, args []string) error { target := args[0] arguments := args[1:] @@ -94,6 +124,7 @@ func runSubmit(cmd *cobra.Command, args []string) error { for i, arg := range arguments { submitArgs[i] = []any{"json", arg} } + submitArgs = padArguments(submitArgs, workflow["parameters"]) // Build options from workflow definition options := make(map[string]any) diff --git a/cli/internal/adapter/protocol.go b/cli/internal/adapter/protocol.go index 7e60fe75..126cccdb 100644 --- a/cli/internal/adapter/protocol.go +++ b/cli/internal/adapter/protocol.go @@ -117,8 +117,12 @@ type SelectResult struct { // SelectHandle identifies a single handle in a select call. type SelectHandle struct { - Type string `json:"type"` // "execution" or "input" + Type string `json:"type"` // "execution", "input" or "stream" ID string `json:"id"` + // Sequence is only meaningful for a "stream" handle: the absolute + // sequence the consumer is waiting for. Omitted for the others, which + // resolve on their own terms. + Sequence *int64 `json:"sequence,omitempty"` } // ReadyMessage is sent by executor when it's ready for work @@ -256,6 +260,17 @@ type GetAssetResult struct { type SuspendParams struct { ExecutionID string `json:"execution_id"` ExecuteAfter *int64 `json:"execute_after,omitempty"` // timestamp in ms + // StreamWait, when set, holds the successor until the stream reaches + // the given sequence (or closes) — a consumer that suspended partway + // through iterating, rather than one waiting on the clock. + StreamWait *StreamWait `json:"stream_wait,omitempty"` +} + +// StreamWait names the stream and absolute sequence a suspended consumer +// is waiting for. +type StreamWait struct { + StreamID string `json:"stream_id"` + Sequence int64 `json:"sequence"` } // CancelParams for cancel request diff --git a/cli/internal/pool/pool.go b/cli/internal/pool/pool.go index c3ffdec1..f2fa9d7a 100644 --- a/cli/internal/pool/pool.go +++ b/cli/internal/pool/pool.go @@ -27,7 +27,7 @@ type ExecutionHandler interface { // UploadBlob uploads a local file as a blob UploadBlob(ctx context.Context, executionID, sourcePath string) (string, error) // Suspend suspends an execution - Suspend(ctx context.Context, executionID string, executeAfter *int64) error + Suspend(ctx context.Context, executionID string, executeAfter *int64, streamWait *adapter.StreamWait) error // Cancel cancels one or more handles (executions and/or inputs) Cancel(ctx context.Context, executionID string, handles []adapter.SelectHandle) error // RegisterGroup registers a group for organizing child executions @@ -196,6 +196,14 @@ func (p *Pool) spawnExecutor(ctx context.Context) (*adapter.Executor, error) { return exec, nil } +// SetStreamTimerPaused stops or restarts a producer stream's idle +// countdown. The server pauses it while every consumer of the stream is +// suspended waiting on it, so a consumer's nap can't time out the +// producer it is waiting for. +func (p *Pool) SetStreamTimerPaused(executionID string, index int, paused bool) { + p.streamTimers.SetPaused(streamKey{executionID, index}, paused) +} + // Execute runs a target. Uses a warm executor if available, otherwise spawns // one on demand. Returns an error if spawning fails (caller should report to server). // timeoutMs, if > 0, enforces a wall-clock timeout on the execution. @@ -828,7 +836,7 @@ func (p *Pool) handleRequest(ctx context.Context, exec *adapter.Executor, method // 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 { + if err := p.handler.Suspend(ctx, req.ExecutionID, req.ExecuteAfter, req.StreamWait); err != nil { errInfo = &adapter.ErrorInfo{Code: "suspend_error", Message: err.Error()} } else { result = map[string]any{} diff --git a/cli/internal/pool/stream_timers.go b/cli/internal/pool/stream_timers.go index 5b45d7c8..9874542a 100644 --- a/cli/internal/pool/stream_timers.go +++ b/cli/internal/pool/stream_timers.go @@ -27,6 +27,12 @@ type streamKey struct { type streamTimer struct { timeout time.Duration timer *time.Timer + // paused while every consumer of this stream is suspended waiting on + // it. Their nap is not the producer being idle — the mirror of the + // rule that a suspended producer's own pause doesn't count — and + // without this a lockstep producer would be timed out by the very + // consumers waiting for it. + paused bool } // streamTimers is a concurrency-safe registry of active stream timers @@ -66,14 +72,35 @@ func (s *streamTimers) Register(key streamKey, timeoutMs int) { s.mu.Unlock() } +// SetPaused stops or restarts a stream's countdown. Pausing leaves the +// entry in place so a later append or close still finds it; resuming +// starts a fresh full-length window rather than resuming a partial one, +// which is the same thing a resumed producer gets. +func (s *streamTimers) SetPaused(key streamKey, paused bool) { + s.mu.Lock() + st, ok := s.timers[key] + if ok { + st.paused = paused + } + s.mu.Unlock() + if !ok { + return + } + st.timer.Stop() + if !paused { + st.timer.Reset(st.timeout) + } +} + // Reset restarts the countdown for a stream. No-op if no timer was // registered (stream has no timeout configured, or was already -// cleared). +// cleared), or while it is paused. func (s *streamTimers) Reset(key streamKey) { s.mu.Lock() st, ok := s.timers[key] + paused := ok && st.paused s.mu.Unlock() - if !ok { + if !ok || paused { return } // time.Timer.Reset is safe to call on a timer that has already diff --git a/cli/internal/worker/worker.go b/cli/internal/worker/worker.go index 23cab7d2..4cb3f2a3 100644 --- a/cli/internal/worker/worker.go +++ b/cli/internal/worker/worker.go @@ -425,6 +425,7 @@ func (w *Worker) runConnection(ctx context.Context, targets map[string]map[strin conn.RegisterHandler("stream_items", w.handleStreamItems) conn.RegisterHandler("stream_closed", w.handleStreamClosed) conn.RegisterHandler("stream_demand", w.handleStreamDemand) + conn.RegisterHandler("stream_timer_pause", w.handleStreamTimerPause) conn.SetOnSession(w.handleSession) if err := conn.Connect(ctx); err != nil { @@ -874,6 +875,31 @@ func (w *Worker) handleStreamDemand(params []any) error { }) } +// handleStreamTimerPause stops or restarts a producer stream's idle +// countdown. Params: [execution_id, index, paused]. The server pauses it +// while every consumer of the stream is suspended waiting on it — their +// nap is not the producer being idle, and without this a lockstep +// producer would be timed out by the consumers waiting for it. +func (w *Worker) handleStreamTimerPause(params []any) error { + if len(params) < 3 { + return fmt.Errorf("stream_timer_pause: insufficient params") + } + executionID, ok := params[0].(string) + if !ok { + return fmt.Errorf("stream_timer_pause: execution_id is not a string (got %T)", params[0]) + } + index, ok := params[1].(float64) + if !ok { + return fmt.Errorf("stream_timer_pause: index is not a number (got %T)", params[1]) + } + paused, ok := params[2].(bool) + if !ok { + return fmt.Errorf("stream_timer_pause: paused is not a bool (got %T)", params[2]) + } + w.pool.SetStreamTimerPaused(executionID, int(index), paused) + return nil +} + func (w *Worker) heartbeatLoop(ctx context.Context) { ticker := time.NewTicker(heartbeatInterval) defer ticker.Stop() @@ -1146,6 +1172,12 @@ func (w *Worker) Select(ctx context.Context, params *adapter.SelectParams) (*ada case "ok": valueArr, ok := resultMap["value"].([]any) if !ok { + if _, present := resultMap["value"]; !present { + // A stream handle resolves without one: the answer is only + // "there is something, stop waiting", and the item itself + // reaches the consumer over its own subscription. + break + } return nil, fmt.Errorf("ok status missing value tuple: %v", resultMap) } value, err := api.ParseValue(valueArr) @@ -1323,12 +1355,20 @@ func (w *Worker) GetAsset(ctx context.Context, executionID string, assetID strin return entriesMap, nil } -func (w *Worker) Suspend(ctx context.Context, executionID string, executeAfter *int64) error { +func (w *Worker) Suspend(ctx context.Context, executionID string, executeAfter *int64, streamWait *adapter.StreamWait) error { conn, err := w.requireConn() if err != nil { return err } - // Python params: (execution_id, execute_after_ms) + // Params: (execution_id, execute_after_ms[, stream_wait]). The third is + // only sent when there is one, so the message stays the shape older + // servers expect. + if streamWait != nil { + return conn.Notify("suspend", executionID, executeAfter, map[string]any{ + "stream_id": streamWait.StreamID, + "sequence": streamWait.Sequence, + }) + } return conn.Notify("suspend", executionID, executeAfter) } diff --git a/docs/docs/checkpoints.md b/docs/docs/checkpoints.md index d8533dc7..eb5cc581 100644 --- a/docs/docs/checkpoints.md +++ b/docs/docs/checkpoints.md @@ -38,12 +38,21 @@ 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 +cursor.get() # the current value, or the default if unset +cursor.set(value) # replace the value +cursor.update(fn) # set it to fn(current), and return that +cursor.reset() # clear it, so get() returns the default again +cursor.is_set() # whether it has a value ``` +`update` is the read-modify-write most checkpoints do — advancing a cursor, accumulating a total — without naming the old value: + +```python +n = count.update(lambda x: x + 1) +``` + +`fn` receives the declared default when the checkpoint isn't set, so an unset checkpoint needs no special case at the call site. It's a read then a write rather than an atomic swap, so if you share one checkpoint between a task body and a `cf.stream` generator, guard it yourself. + 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. @@ -64,7 +73,7 @@ Writes are throttled and delivered in the background, so a crash can lose up to 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: +Whatever has been written when an execution suspends, returns, or fails is delivered before the next attempt starts — with one deliberate exception, described under [consistency](#consistency) below. 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) @@ -74,6 +83,14 @@ send_notification() `cf.flush()` returns once the server has acknowledged the write. +## Consistency + +A step's checkpoints are one snapshot of its progress rather than a set of independent cells: whatever has been written is delivered as a single delta and applied at once. Deltas are only cut where the execution could resume from — never part-way through consuming something a replay can't re-read. + +That's what keeps state derived from a [stream](./streams.md) honest. A checkpoint written in the loop body is published in the same delta as the cursor advance that consumes the item, so a running total never counts an item the cursor says was never read. If the iteration doesn't finish — a `break`, an exception, a lost worker — the item stays unconsumed and the writes derived from it are dropped, because the next attempt reads that item again. + +A replay therefore repeats whole items rather than fractions of one. That isn't exactly-once: the item *is* delivered again, so anything else the loop body did happens again too. [Memoize](./memoizing.md) what it calls. + ## 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. @@ -84,6 +101,10 @@ A checkpoint belongs to the step that actually executes. A step resolved from th 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. ::: +## Reserved names + +Names starting with an underscore belong to the adapter, and `cf.Checkpoint("_...")` raises. They're currently used for the cursors behind [stream](./streams.md) suspension, which track how far a consumer has read; you'll see them alongside your own in Studio. The Python variable name is unrestricted — only the checkpoint's name matters, so `_cursor = cf.Checkpoint("cursor")` is fine. + ## 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 559e06dc..a61e58fc 100644 --- a/docs/docs/python_reference.md +++ b/docs/docs/python_reference.md @@ -341,7 +341,7 @@ cf.Checkpoint[T | None]( | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `name` | `str` | required | Checkpoint name, unique within the step | +| `name` | `str` | required | Checkpoint name, unique within the step. Can't start with `_`, which is reserved for adapter-managed state | | `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. @@ -354,6 +354,10 @@ The current value, or the declared default if it isn't set. A checkpoint explici Sets the value, replacing anything already there. +#### `checkpoint.update(fn: Callable[[T], T]) -> T` + +Sets the value to `fn(current)` and returns what was stored. `fn` receives the declared default when the checkpoint isn't set. A read followed by a write, not an atomic swap. + #### `checkpoint.reset() -> None` Clears the checkpoint, so `get()` returns the declared default again. Distinct from `set(None)`, which stores `None`. diff --git a/docs/docs/streams.md b/docs/docs/streams.md index 1108d3d6..982c0545 100644 --- a/docs/docs/streams.md +++ b/docs/docs/streams.md @@ -115,7 +115,49 @@ def fetch_pages(url: str): ... If the producer doesn't append an item within that window, the stream is closed as timed out and the generator is stopped. The window counts time spent waiting for consumer demand too, so with a lockstep buffer a slow or absent consumer can time out the producer. The producing execution still completes with its value, but is excluded from [caching](./caching.md) and [memoizing](./memoizing.md). -The timeout is measured from the last item appended, per execution. A suspended step is not idle, so the pause between a suspend and its resumption doesn't count; the timer starts again when the resumed execution registers. There is no consumer-side timeout: iterating a stream waits as long as it takes. +The timeout is measured from the last item appended, per execution. A suspended step is not idle, so the pause between a suspend and its resumption doesn't count; the timer starts again when the resumed execution registers. The same applies from the other side: while every consumer of a stream is suspended waiting on it (below), the producer's countdown is paused — a consumer's nap isn't the producer being idle, and otherwise a lockstep producer could be timed out by the very consumers waiting for it. + +This timeout is the producer's. For the consumer's, see below. + +## Suspending while consuming + +By default, iterating a stream waits as long as it takes — the consumer holds its worker slot through every gap. Iterating inside a [`cf.suspense`](./suspense.md) scope instead gives the slot up when the stream goes quiet: + +```python +@cf.task() +def handle_events(events: cf.Stream[dict]): + with cf.suspense(30): + for event in events: + store.submit(event) +``` + +If thirty seconds pass with nothing arriving, the execution suspends. It is resumed when the stream next holds the item it was waiting for — or when the stream closes, since it will never hold it then. The resumed execution runs the body from the top, as always — but it doesn't re-read the stream. The adapter keeps the consumed position in a checkpoint of its own, named after the stream and the view being read, and re-subscribes there, so each item is delivered to exactly one attempt. + +Two things follow from the body restarting: + +- **Anything derived from the stream needs its own [checkpoint](./checkpoints.md).** The position survives; your running total doesn't, unless you keep it somewhere durable. Keeping the two in step is the runtime's job: a checkpoint written in the loop body is published in the same delta as the cursor advance that consumes the item, so the pair moves together or not at all. +- **The scope covers the loop body too.** A `.result()` inside it inherits the same timeout, so the usual care about re-execution applies — [memoize](./memoizing.md) what the body calls. + +A bare `cf.suspense()` means what it means for a result: don't wait at all. The consumer asks the server whether the next item is there, and suspends only if it isn't — so a zero timeout costs a round trip whenever the local queue is momentarily empty, rather than a wasted restart. Whether to suspend is always the server's answer, never a guess from the consumer's own queue, which is fed asynchronously and so says nothing about what the stream holds. + +Pick the threshold with `buffer` in mind. Under the default lockstep budget the producer is never more than one item ahead, so a short threshold makes the consumer suspend on nearly every item; give the producer a `buffer` (or `buffer=None`) when the consumer is going to nap. + +The cursor is never cleared. Re-running a consumer that already drained its stream therefore does nothing — it resumes at the end. Clear the step's checkpoints to make it read again. + +### An idle pipeline + +The two halves compose. Have the workflow *submit* the consumer rather than wait on it, and a whole pipeline can sit idle holding no worker slots at all: + +```python +@cf.workflow() +def events_pipeline(): + events = tail_events() # returns once the stream is registered + return handle_events.submit(events) # a handle, not a result +``` + +Calling the producer doesn't block — a generator task's result *is* the stream reference, recorded before the first item — and submitting the consumer doesn't either, so the workflow step finishes immediately. Between bursts the producer is suspended, the consumer is suspended, and the workflow is done: nothing is running, and the consumer is only scheduled again once there is something for it to read. + +The trade is that the run's result is a handle rather than a value. It still resolves — anything reading it waits for the consumer's result the usual way — but to see the value directly you look at the consumer's step. ## Workspaces diff --git a/docs/docs/suspense.md b/docs/docs/suspense.md index 3b923f5a..7ba3adc1 100644 --- a/docs/docs/suspense.md +++ b/docs/docs/suspense.md @@ -53,3 +53,7 @@ with cf.suspense(10): :::warning It's important that any tasks called within the suspense block _or before it_ are [memoized](/memoizing) (or cached). Otherwise the task is likely to keep suspending as a new task will be spawned on each execution. ::: + +## Suspense and streams + +A suspense scope also applies to iterating a [stream](/streams): a gap longer than the timeout suspends the consumer, and the resumed execution carries on from where it stopped rather than re-reading. Unlike a result — which the re-run simply resolves again — a stream's position isn't naturally replayable, so the adapter keeps it in a [checkpoint](/checkpoints) of its own. Anything _derived_ from the stream is still yours to checkpoint. See [suspending while consuming](/streams#suspending-while-consuming). diff --git a/server/lib/coflux/handlers/worker.ex b/server/lib/coflux/handlers/worker.ex index c1b24817..ba98ff3e 100644 --- a/server/lib/coflux/handlers/worker.ex +++ b/server/lib/coflux/handlers/worker.ex @@ -593,15 +593,28 @@ defmodule Coflux.Handlers.Worker do end "suspend" -> - [execution_id, execute_after] = message["params"] + [execution_id, execute_after | rest] = message["params"] # TODO: validate execute_after + # An optional stream to wait on: a consumer that suspended + # mid-iteration wants its successor held until the stream reaches + # the sequence it stopped at (or closes). Older adapters send two + # params and get the unconditional behaviour. + dependencies = + case rest do + [%{"stream_id" => stream_id, "sequence" => sequence}] -> + [{:stream, stream_id, sequence}] + + _ -> + [] + end + if is_recognised_execution?(execution_id, state) do :ok = Orchestration.record_result( state.project_id, execution_id, - {:suspended, execute_after, []} + {:suspended, execute_after, dependencies} ) {[], state} @@ -800,6 +813,10 @@ defmodule Coflux.Handlers.Worker do {[command_message("stream_items", [execution_external_id, subscription_id, encoded])], state} end + def websocket_info({:stream_timer_pause, execution_external_id, index, paused}, state) do + {[command_message("stream_timer_pause", [execution_external_id, index, paused])], state} + end + def websocket_info({:stream_demand, execution_external_id, index, n}, state) do {[command_message("stream_demand", [execution_external_id, index, n])], state} end @@ -1005,6 +1022,12 @@ defmodule Coflux.Handlers.Worker do base = %{"winner" => idx} case detail do + # A stream handle resolves with no value — the item reaches the + # consumer through its own subscription; this only says "stop + # waiting". + :available -> + Map.put(base, "status", "ok") + {:value, value} -> Map.merge(base, %{"status" => "ok", "value" => compose_value(value)}) diff --git a/server/lib/coflux/orchestration/epoch.ex b/server/lib/coflux/orchestration/epoch.ex index 3703bfe1..74ee2987 100644 --- a/server/lib/coflux/orchestration/epoch.ex +++ b/server/lib/coflux/orchestration/epoch.ex @@ -627,18 +627,23 @@ defmodule Coflux.Orchestration.Epoch do {:ok, stream_deps} = query( source_db, - "SELECT stream_ref_id, created_at FROM stream_dependencies WHERE execution_id = ?1", + "SELECT stream_ref_id, sequence, created_at FROM stream_dependencies WHERE execution_id = ?1", {old_exec_id} ) - Enum.each(stream_deps, fn {old_ref_id, created_at} -> + Enum.each(stream_deps, fn {old_ref_id, sequence, created_at} -> new_ref_id = ensure_stream_ref(source_db, target_db, old_ref_id) {:ok, _} = insert_one( target_db, :stream_dependencies, - %{execution_id: new_exec_id, stream_ref_id: new_ref_id, created_at: created_at}, + %{ + execution_id: new_exec_id, + stream_ref_id: new_ref_id, + sequence: sequence, + created_at: created_at + }, on_conflict: "DO NOTHING" ) end) diff --git a/server/lib/coflux/orchestration/runs.ex b/server/lib/coflux/orchestration/runs.ex index 393590f2..dc65cdec 100644 --- a/server/lib/coflux/orchestration/runs.ex +++ b/server/lib/coflux/orchestration/runs.ex @@ -1,5 +1,5 @@ defmodule Coflux.Orchestration.Runs do - alias Coflux.Orchestration.{Models, Results, Values, TagSets, CacheConfigs, Utils} + alias Coflux.Orchestration.{Models, Results, Values, TagSets, CacheConfigs, Utils, Streams} import Coflux.Store @@ -548,6 +548,7 @@ defmodule Coflux.Orchestration.Runs do execute_after, dependency_ref_ids, input_dependency_ids \\ [], + stream_waits \\ [], created_by \\ nil ) do with_transaction(db, fn -> @@ -573,6 +574,15 @@ defmodule Coflux.Orchestration.Runs do Enum.map(input_dependency_ids, &{execution_id, &1, now}) ) + # A consumer that suspended mid-stream: gate the successor until the + # stream reaches the sequence it was waiting for. Written through + # Streams so the on-conflict rule stays in one place — a lineage row + # for this execution can already exist if it subscribed before + # suspending again. + Enum.each(stream_waits, fn {stream_ref_id, sequence} -> + :ok = Streams.record_wait(db, execution_id, stream_ref_id, sequence) + end) + {:ok, execution_id, attempt, now} end) end diff --git a/server/lib/coflux/orchestration/server.ex b/server/lib/coflux/orchestration/server.ex index 3ab20e7f..c6c3bb1e 100644 --- a/server/lib/coflux/orchestration/server.ex +++ b/server/lib/coflux/orchestration/server.ex @@ -115,6 +115,10 @@ defmodule Coflux.Orchestration.Server do # execution_id -> MapSet of execution_ids that this execution is waiting on pending_dependencies: %{}, + # Stream waits indexed by stream, so an append can check for + # waiters with one lookup instead of scanning every pending + # dependency. Derived from pending_dependencies; rebuilt with it. + stream_dependency_keys: %{}, # execution_id -> MapSet of execution_ids that are waiting on this execution dependency_waiters: %{}, @@ -2046,6 +2050,7 @@ defmodule Coflux.Orchestration.Server do |> ensure_stream_producer(stream_id, producer_session_id) |> push_stream_item(stream_id, sequence, value) |> notify_stream_item_appended(stream_id, execution_id, sequence, value, created_at) + |> update_dependencies_on_stream(stream_id, sequence) |> flush_notifications() {:reply, :ok, state} @@ -2078,6 +2083,7 @@ defmodule Coflux.Orchestration.Server do state |> push_stream_closed(stream_id, reason, error) |> notify_stream_closed(stream_id, execution_id, reason, error, closed_at) + |> update_dependencies_on_stream(stream_id, :closed) |> drop_stream_producer(stream_id) |> flush_notifications() @@ -2173,26 +2179,24 @@ defmodule Coflux.Orchestration.Server do # being read. Uses stream_refs so the edge survives epoch rotation. {:ok, stream_ref_id} = Streams.create_stream_ref_for(state.db, stream_id) - {:ok, inserted_id} = - Streams.record_dependency(state.db, consumer_execution_id, stream_ref_id) + {:ok, _} = Streams.record_dependency(state.db, consumer_execution_id, stream_ref_id) - state = - if inserted_id do - {:ok, {run_external_id}} = - Runs.get_external_run_id_for_execution(state.db, consumer_execution_id) + # Announced whether or not the insert was new: the row may already + # exist as a wait recorded against this execution before it ran. The + # topic merges, so re-announcing an edge it already holds is harmless. + {:ok, {run_external_id}} = + Runs.get_external_run_id_for_execution(state.db, consumer_execution_id) - {:ok, {stream_run_ext_id, step_number, index, module, target}} = - Streams.get_stream_ref(state.db, stream_ref_id) + {:ok, {stream_run_ext_id, step_number, index, module, target}} = + Streams.get_stream_ref(state.db, stream_ref_id) - notify_listeners( - state, - {:run, run_external_id}, - {:stream_dependency, consumer_execution_external_id, - stream_external_id(stream_run_ext_id, step_number, index), module, target} - ) - else - state - end + state = + notify_listeners( + state, + {:run, run_external_id}, + {:stream_dependency, consumer_execution_external_id, + stream_external_id(stream_run_ext_id, step_number, index), module, target} + ) state = flush_notifications(state) @@ -4668,11 +4672,32 @@ defmodule Coflux.Orchestration.Server do dependency_keys = Keyword.get(opts, :dependency_keys, []) created_by = Keyword.get(opts, :created_by) - # Separate execution and input dependencies - {exec_deps, input_deps} = - Enum.reduce(dependency_keys, {[], []}, fn - {:execution, id}, {execs, inputs} -> {[id | execs], inputs} - {:input, id}, {execs, inputs} -> {execs, [id | inputs]} + # Separate execution, input and stream dependencies. + # + # A stream wait arrives naming the stream externally, and is persisted + # against a *stream ref* — the same indirection subscription lineage + # uses, so the edge survives epoch rotation. A stream that can't be + # resolved is dropped rather than recorded: gating on it would strand + # the successor forever. + {exec_deps, input_deps, stream_waits} = + Enum.reduce(dependency_keys, {[], [], []}, fn + {:execution, id}, {execs, inputs, streams} -> + {[id | execs], inputs, streams} + + {:input, id}, {execs, inputs, streams} -> + {execs, [id | inputs], streams} + + {:stream, external_id, sequence}, {execs, inputs, streams} -> + case resolve_stream_id(state, external_id) do + {:ok, stream_id} -> + case Streams.create_stream_ref_for(state.db, stream_id) do + {:ok, ref_id} -> {execs, inputs, [{ref_id, sequence} | streams]} + {:error, :not_found} -> {execs, inputs, streams} + end + + {:error, :not_found} -> + {execs, inputs, streams} + end end) # Convert internal dependency execution IDs to execution_ref IDs @@ -4692,6 +4717,7 @@ defmodule Coflux.Orchestration.Server do execute_after, dependency_ref_ids, input_deps, + stream_waits, created_by ) do {:ok, execution_id, attempt, created_at} -> @@ -4790,6 +4816,23 @@ defmodule Coflux.Orchestration.Server do {:step, step.module, step.target, step.type, run.external_id, step.number, attempt} ) + # Announce the stream waits `rerun_step` recorded. They are lineage + # edges against an execution that hasn't run yet, so without this + # nothing announces them until it subscribes — and it may never get + # that far. A topic opened later reads them from the snapshot. + state = + Enum.reduce(stream_waits, state, fn {stream_ref_id, _sequence}, state -> + {:ok, {stream_run_ext_id, stream_step_number, index, module, target}} = + Streams.get_stream_ref(state.db, stream_ref_id) + + notify_listeners( + state, + {:run, run.external_id}, + {:stream_dependency, execution_external_id, + stream_external_id(stream_run_ext_id, stream_step_number, index), module, target} + ) + end) + # Notify run topic about input dependencies for this execution state = Enum.reduce(input_deps, state, fn input_id, state -> @@ -4973,6 +5016,7 @@ defmodule Coflux.Orchestration.Server do |> remap_config_ids(id_mappings) |> copy_in_flight_runs() |> Map.put(:pending_dependencies, %{}) + |> Map.put(:stream_dependency_keys, %{}) |> Map.put(:dependency_waiters, %{}) |> initialize_pending_dependencies() |> maybe_start_index_build() @@ -6865,6 +6909,7 @@ defmodule Coflux.Orchestration.Server do state |> push_stream_closed(stream_id, push_reason, push_error) |> notify_stream_closed(stream_id, execution_id, push_reason, push_error, closed_at) + |> update_dependencies_on_stream(stream_id, :closed) |> drop_stream_producer(stream_id) {:error, :already_closed} -> @@ -7562,6 +7607,10 @@ defmodule Coflux.Orchestration.Server do {:input, _input_id} -> # Input dependencies are not shown in the queue nil + + {:stream, _stream_id, _sequence} -> + # Nor stream waits — the queue lists executions being waited on. + nil end) |> Enum.reject(&is_nil/1) end @@ -7688,9 +7737,53 @@ defmodule Coflux.Orchestration.Server do end) end + # Collect unmet stream waits. Only rows with a sequence are waits; + # the rest of the table is subscription lineage. This runs solely for + # executions that have not been assigned yet, which is what makes it + # safe for the sequence to stay on the row after the gate clears — a + # completed execution's row is never read back here. + stream_dependencies = + case Streams.get_wait_dependencies(db, execution_id) do + {:ok, waits} -> + Enum.reduce(waits, MapSet.new(), fn {stream_ref_id, sequence}, acc -> + case resolve_stream_ref_id(db, stream_ref_id) do + {:ok, stream_id} -> + if stream_reached?(db, stream_id, sequence) do + acc + else + MapSet.put(acc, {:stream, stream_id, sequence}) + end + + {:error, :not_found} -> + # The stream is gone (a pruned epoch, say). Waiting on it + # forever would strand the execution, so treat it as met. + acc + end + end) + end + argument_dependencies |> MapSet.union(result_dependencies) |> MapSet.union(input_dependencies) + |> MapSet.union(stream_dependencies) + end + + # A stream wait is met once the stream holds the sequence, or can never + # hold it because it closed. + defp stream_reached?(db, stream_id, sequence) do + case Streams.get_head(db, stream_id) do + {:ok, head} -> head >= sequence || Streams.closed?(db, stream_id) + end + end + + defp resolve_stream_ref_id(db, stream_ref_id) do + case Streams.get_stream_ref(db, stream_ref_id) do + {:ok, {run_external_id, step_number, index, _module, _target}} -> + Streams.get_stream_id_by_key(db, run_external_id, step_number, index) + + {:error, :not_found} -> + {:error, :not_found} + end end # Walk references and collect tagged dependency keys that are still pending. @@ -7753,15 +7846,86 @@ defmodule Coflux.Orchestration.Server do put_in(state, [Access.key(:pending_dependencies), execution_id], dependencies) Enum.reduce(dependencies, state, fn dependency_id, state -> - update_in( - state, + state + |> update_in( [Access.key(:dependency_waiters), Access.key(dependency_id, MapSet.new())], &MapSet.put(&1, execution_id) ) + |> index_stream_dependency(dependency_id) end) end end + # Stream waits get a secondary index, keyed by stream. Appends are hot, + # and without it every appended item would have to scan the whole + # dependency_waiters map to find out whether anything was waiting; with + # it the check is one map lookup that almost always misses. + defp index_stream_dependency(state, {:stream, stream_id, _sequence} = key) do + was_waiting = stream_has_waiters?(state, stream_id) + + state = + update_in( + state, + [Access.key(:stream_dependency_keys), Access.key(stream_id, MapSet.new())], + &MapSet.put(&1, key) + ) + + # First waiter: the producer's idle countdown stops. A consumer's nap + # is not the producer being idle — the mirror of the existing rule + # that a suspended producer's own pause doesn't count against it. + if was_waiting, do: state, else: set_stream_timer_paused(state, stream_id, true) + end + + defp index_stream_dependency(state, _key), do: state + + defp unindex_stream_dependency(state, {:stream, stream_id, _sequence} = key) do + state = + update_in( + state, + [Access.key(:stream_dependency_keys), Access.key(stream_id, MapSet.new())], + &MapSet.delete(&1, key) + ) + + if MapSet.size(state.stream_dependency_keys[stream_id] || MapSet.new()) == 0 do + state + |> update_in([Access.key(:stream_dependency_keys)], &Map.delete(&1, stream_id)) + |> set_stream_timer_paused(stream_id, false) + else + state + end + end + + defp unindex_stream_dependency(state, _key), do: state + + defp stream_has_waiters?(state, stream_id) do + MapSet.size(Map.get(state.stream_dependency_keys, stream_id, MapSet.new())) > 0 + end + + # Tell the producer's worker to stop or restart the stream's idle + # countdown. Enforcement is worker-side, so this is the only way to say + # it. A producer with no live session has no timer to pause. + # + # Deliberately resolved from the database rather than from + # `stream_producers`: that map only exists to track demand, so a stream + # with `buffer=nil` has no entry at all — and an unbuffered producer is + # exactly what you pair with a suspending consumer, so it is the case + # that most needs this. Infrequent enough for the lookup not to matter: + # once when the first waiter arrives, once when the last one clears. + defp set_stream_timer_paused(state, stream_id, paused) do + with execution_external_id when is_binary(execution_external_id) <- + producer_external_id(state.db, stream_id), + {:ok, session_id} <- find_session_for_execution(state, execution_external_id), + {:ok, stream} <- Streams.get_stream(state.db, stream_id) do + send_session( + state, + session_id, + {:stream_timer_pause, execution_external_id, stream.index, paused} + ) + else + _ -> state + end + end + # Remove an execution from the dependency tracking (when assigned or completed). defp unregister_pending_dependencies(state, execution_id) do case Map.fetch(state.pending_dependencies, execution_id) do @@ -7777,11 +7941,12 @@ defmodule Coflux.Orchestration.Server do # Clean up empty waiter entries if MapSet.size(state.dependency_waiters[dependency_id] || MapSet.new()) == 0 do - update_in( - state, + state + |> update_in( [Access.key(:dependency_waiters)], &Map.delete(&1, dependency_id) ) + |> unindex_stream_dependency(dependency_id) else state end @@ -7885,19 +8050,46 @@ defmodule Coflux.Orchestration.Server do end end + # Called when a stream gains an item, or closes. Wakes any execution that + # suspended mid-iteration and is gated on this stream. + # + # `head` is the highest sequence now available, or `:closed` — a closed + # stream will never reach the sequence anyone is still waiting for, so + # every waiter on it is released rather than stranded. The consumer + # re-subscribes at its checkpoint cursor and sees the closure. + defp update_dependencies_on_stream(state, stream_id, head) do + case Map.fetch(state.stream_dependency_keys, stream_id) do + {:ok, keys} -> + keys + |> Enum.filter(fn {:stream, _stream_id, sequence} -> + head == :closed || head >= sequence + end) + |> Enum.reduce(state, &clear_dependency_key(&2, &1)) + + :error -> + state + end + end + # Called when an input response is recorded. Resolves the {:input, id} # dependency for any executions that were waiting on this input. defp update_dependencies_on_input(state, input_id) do - dependency_key = {:input, input_id} + clear_dependency_key(state, {:input, input_id}) + end + # Drop one dependency key: forget its waiter set, and take the key out of + # each waiter's pending set, scheduling any execution that has nothing + # left to wait for. + defp clear_dependency_key(state, dependency_key) do case Map.fetch(state.dependency_waiters, dependency_key) do {:ok, waiters} -> state = - update_in( - state, + state + |> update_in( [Access.key(:dependency_waiters)], &Map.delete(&1, dependency_key) ) + |> unindex_stream_dependency(dependency_key) Enum.reduce(waiters, state, fn waiter_id, state -> case Map.fetch(state.pending_dependencies, waiter_id) do @@ -8102,6 +8294,37 @@ defmodule Coflux.Orchestration.Server do end end + # A stream handle asks one question: has the stream reached this + # sequence (or closed, so it never will)? A consumer can't answer it + # itself — its queue is fed asynchronously, so an empty one means + # "nothing has arrived yet", not "the stream has nothing". + # + # Resolving carries no value: the item reaches the consumer through the + # subscription it already holds. The answer only says "there is + # something, stop waiting". + defp process_select_handle( + state, + %{"type" => "stream", "id" => stream_external_id} = handle, + _from_execution_id, + _from_execution_external_id + ) do + sequence = Map.get(handle, "sequence", 0) + + case resolve_stream_id(state, stream_external_id) do + {:error, :not_found} -> + {{:error, :not_found}, state} + + {:ok, stream_id} -> + if stream_reached?(state.db, stream_id, sequence) do + {{:ok, {:resolved, :available}}, state} + else + {{:ok, + {:pending, {:stream, stream_external_id}, {:stream, stream_external_id, sequence}}}, + state} + end + end + end + defp process_select_handle( state, %{"type" => "input", "id" => input_external_id}, diff --git a/server/lib/coflux/orchestration/streams.ex b/server/lib/coflux/orchestration/streams.ex index d8f91179..1e97cf45 100644 --- a/server/lib/coflux/orchestration/streams.ex +++ b/server/lib/coflux/orchestration/streams.ex @@ -718,8 +718,10 @@ defmodule Coflux.Orchestration.Streams do end end - # Records that `execution_id` subscribed to the stream. Returns - # `{:ok, id}` for a new edge, `{:ok, nil}` if it already existed. + # Records that `execution_id` subscribed to the stream. An edge may + # already exist — `record_wait` writes one against an execution before it + # runs — so this is idempotent, and the subscribe is announced to the run + # topic either way rather than only when a row is written. def record_dependency(db, execution_id, stream_ref_id) do with_transaction(db, fn -> insert_one( @@ -735,6 +737,72 @@ defmodule Coflux.Orchestration.Streams do end) end + # Records that `execution_id` is gated on the stream reaching `sequence`. + # + # Written when a consumer suspends mid-iteration, against the *successor* + # execution — which has not run, so it has no lineage row of its own yet. + # If it later subscribes, `record_dependency`'s `DO NOTHING` leaves this + # sequence in place. + # No transaction of its own: this is a single insert, and its only caller + # runs inside `Runs.rerun_step`'s transaction — SQLite has no nested + # transactions. + def record_wait(db, execution_id, stream_ref_id, sequence) do + {:ok, _} = + insert_one( + db, + :stream_dependencies, + %{ + execution_id: execution_id, + stream_ref_id: stream_ref_id, + sequence: sequence, + created_at: current_timestamp() + }, + on_conflict: "(execution_id, stream_ref_id) DO UPDATE SET sequence = excluded.sequence" + ) + + :ok + end + + # `[{stream_ref_id, sequence}, ...]` for the waits an execution is gated + # on. Lineage-only edges (a NULL sequence) are not waits and are excluded. + def get_wait_dependencies(db, execution_id) do + query( + db, + """ + SELECT stream_ref_id, sequence + FROM stream_dependencies + WHERE execution_id = ?1 AND sequence IS NOT NULL + """, + {execution_id} + ) + end + + # Highest sequence in the stream, or -1 when it holds no items. + def get_head(db, stream_id) do + case query_one( + db, + "SELECT MAX(sequence) FROM stream_items WHERE stream_id = ?1", + {stream_id} + ) do + {:ok, {nil}} -> {:ok, -1} + {:ok, {head}} -> {:ok, head} + end + end + + # Whether the stream has recorded a closure. A closed stream never gains + # another item, so a consumer waiting on one has to be woken rather than + # left gated forever. + def closed?(db, stream_id) do + case query_one( + db, + "SELECT 1 FROM stream_closures WHERE stream_id = ?1", + {stream_id} + ) do + {:ok, nil} -> false + {:ok, _} -> true + end + end + # `%{execution_id => [stream_ref_id, ...]}` for every consumer execution # in the run. def get_run_dependencies(db, run_id) do diff --git a/server/priv/migrations/orchestration/4.sql b/server/priv/migrations/orchestration/4.sql index 3dc8458f..2b4a0ee7 100644 --- a/server/priv/migrations/orchestration/4.sql +++ b/server/priv/migrations/orchestration/4.sql @@ -254,9 +254,19 @@ CREATE TABLE stream_refs ( -- is written when a consumer subscribes (regardless of whether items are -- read), so data lineage is preserved even for subscriptions that yield no -- values. +-- +-- `sequence` annotates that edge for a consumer that suspended while +-- reading: it is the sequence the execution was gated on, and the +-- execution is not scheduled until the stream reaches it (or closes). It +-- stays populated afterwards — it is a record of where this attempt +-- resumed from, not live state needing cleanup, and only +-- `compute_pending_dependencies` reads it, which runs solely for +-- executions that have not been assigned yet. NULL means a plain lineage +-- edge, which is every row written by subscribing. CREATE TABLE stream_dependencies ( execution_id INTEGER NOT NULL, stream_ref_id INTEGER NOT NULL, + sequence INTEGER, created_at INTEGER NOT NULL, PRIMARY KEY (execution_id, stream_ref_id), FOREIGN KEY (execution_id) REFERENCES executions ON DELETE CASCADE, diff --git a/tests/support/executor.py b/tests/support/executor.py index 6f3da300..31dbd3f7 100644 --- a/tests/support/executor.py +++ b/tests/support/executor.py @@ -236,9 +236,13 @@ def cancel(self, execution_id, target_execution_id): ) return self._request(msg) - def suspend(self, execution_id, execute_after=None): - """Suspend the current execution.""" - msg = protocol.suspend_request(None, execution_id, execute_after) + def suspend(self, execution_id, execute_after=None, stream_wait=None): + """Suspend the current execution. + + ``stream_wait`` is a ``(stream_id, sequence)`` pair gating the + successor on that stream reaching the sequence, or closing. + """ + msg = protocol.suspend_request(None, execution_id, execute_after, stream_wait) return self._request(msg) def checkpoint_set(self, execution_id, **values): diff --git a/tests/support/manifest.py b/tests/support/manifest.py index 1f3aa50c..4eeed46b 100644 --- a/tests/support/manifest.py +++ b/tests/support/manifest.py @@ -18,7 +18,11 @@ def _target( "module": module, "name": name, "type": type, - "parameters": [{"name": p} for p in (parameters or [])], + # A parameter is either a name, or a dict for one carrying a default + # (JSON-encoded, as the adapter reports it). + "parameters": [ + {"name": p} if isinstance(p, str) else dict(p) for p in (parameters or []) + ], } if retries is not None: target["retries"] = retries diff --git a/tests/support/protocol.py b/tests/support/protocol.py index c09edd24..586fca01 100644 --- a/tests/support/protocol.py +++ b/tests/support/protocol.py @@ -196,10 +196,16 @@ def cancel_request(request_id, execution_id, handles): } -def suspend_request(request_id, execution_id, execute_after=None): +def suspend_request(request_id, execution_id, execute_after=None, stream_wait=None): + """``stream_wait`` is a ``(stream_id, sequence)`` pair for a consumer + that suspended mid-iteration: the successor is held until the stream + reaches that sequence, or closes.""" params = {"execution_id": execution_id} if execute_after is not None: params["execute_after"] = execute_after + if stream_wait is not None: + stream_id, sequence = stream_wait + params["stream_wait"] = {"stream_id": stream_id, "sequence": sequence} return {"id": request_id, "method": "suspend", "params": params} diff --git a/tests/test_caching.py b/tests/test_caching.py index d39377a6..13c54075 100644 --- a/tests/test_caching.py +++ b/tests/test_caching.py @@ -282,3 +282,49 @@ def test_run_level_memo_inherited_by_child_tasks(worker): ex0.conn.complete(ex0.execution_id, value="done") assert ctx.result(run_id)["value"]["data"] == "done" + + +def test_submit_fills_in_omitted_defaults(worker): + """A workflow submitted without its trailing defaults records them anyway. + + Adapters bind their own defaults before submitting, so a submission that + arrives without a signature to bind against - the CLI here - is padded + from the manifest. Otherwise the same call keys differently depending on + whether the default was spelled out. + """ + targets = [ + workflow( + "test", + "main", + parameters=["x", {"name": "y", "default": "1"}, {"name": "z", "default": "2"}], + ), + ] + + with worker(targets) as ctx: + resp = ctx.submit("test", "main", "42") + + ex = ctx.executor.next_execute() + assert [a["value"] for a in ex.arguments] == [42, 1, 2] + + ex.conn.complete(ex.execution_id, value="done") + assert ctx.result(resp["runId"])["value"]["data"] == "done" + + +def test_submit_keeps_explicit_values(worker): + """Padding only fills in what was left out.""" + targets = [ + workflow( + "test", + "main", + parameters=["x", {"name": "y", "default": "1"}, {"name": "z", "default": "2"}], + ), + ] + + with worker(targets) as ctx: + resp = ctx.submit("test", "main", "42", "9") + + ex = ctx.executor.next_execute() + assert [a["value"] for a in ex.arguments] == [42, 9, 2] + + ex.conn.complete(ex.execution_id, value="done") + assert ctx.result(resp["runId"])["value"]["data"] == "done" diff --git a/tests/test_streams.py b/tests/test_streams.py index c6257925..61de02d5 100644 --- a/tests/test_streams.py +++ b/tests/test_streams.py @@ -1928,3 +1928,265 @@ def test_backpressure_partition_consumer_advances_demand_past_unmatched(worker): prod_ex.conn.complete(prod_ex.execution_id) cons_ex.conn.complete(cons_ex.execution_id) ctx.result(prod_resp["runId"]) + + +# --- Consumer-side suspension: gating the successor on the stream ------------ +# +# A consumer that suspends mid-iteration names the stream and the absolute +# sequence it stopped at. The successor isn't dispatched until the stream +# reaches that sequence — or closes, which is the only other way the wait +# could ever end. + + +def _suspended_consumer(ctx, stream_id, sequence): + """Run a consumer up to the point of suspending on ``stream_id``.""" + ctx.submit("test", "consumer") + cons_ex = ctx.executor.next_execute() + cons_ex.conn.stream_subscribe( + cons_ex.execution_id, subscription_id=1, stream_id=stream_id + ) + # No delay: the stream gate is the only thing that can hold the + # successor back, so a dispatch here would mean the gate isn't working. + cons_ex.conn.suspend( + cons_ex.execution_id, + stream_wait=(stream_id, sequence), + ) + return cons_ex + + +def test_stream_wait_holds_successor_until_the_item_lands(worker): + """The successor is held with no delay set, and appending the sequence + it was waiting for is what dispatches it.""" + targets = [workflow("test", "producer"), workflow("test", "consumer")] + + with worker(targets, concurrency=3) as ctx: + ctx.submit("test", "producer") + prod_ex = ctx.executor.next_execute() + stream = prod_ex.conn.stream_register(prod_ex.execution_id, 0) + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 0, "v0") + + # Consumed item 0, so it is waiting for sequence 1. + _suspended_consumer(ctx, stream["id"], sequence=1) + + # Nothing to wake it yet. + with pytest.raises(TimeoutError): + ctx.executor.next_execute(timeout=1) + + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 1, "v1") + + cons2 = ctx.executor.next_execute() + assert cons2.target == "consumer" + cons2.conn.complete(cons2.execution_id) + + prod_ex.conn.stream_close(prod_ex.execution_id, 0) + prod_ex.conn.complete(prod_ex.execution_id, value="done") + + +def test_stream_wait_is_released_by_the_stream_closing(worker): + """A closed stream will never reach the sequence, so the waiter is + released rather than stranded — it re-subscribes and sees the closure.""" + targets = [workflow("test", "producer"), workflow("test", "consumer")] + + with worker(targets, concurrency=3) as ctx: + ctx.submit("test", "producer") + prod_ex = ctx.executor.next_execute() + stream = prod_ex.conn.stream_register(prod_ex.execution_id, 0) + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 0, "v0") + + _suspended_consumer(ctx, stream["id"], sequence=1) + + with pytest.raises(TimeoutError): + ctx.executor.next_execute(timeout=1) + + # Sequence 1 never arrives; the stream ends instead. + prod_ex.conn.stream_close(prod_ex.execution_id, 0) + + cons2 = ctx.executor.next_execute() + assert cons2.target == "consumer" + cons2.conn.complete(cons2.execution_id) + + prod_ex.conn.complete(prod_ex.execution_id, value="done") + + +def test_stream_wait_already_satisfied_does_not_hold_the_successor(worker): + """The item can arrive between the consumer deciding to suspend and the + server recording it. The gate is a condition, not an edge, so an + already-met wait schedules straight away.""" + targets = [workflow("test", "producer"), workflow("test", "consumer")] + + with worker(targets, concurrency=3) as ctx: + ctx.submit("test", "producer") + prod_ex = ctx.executor.next_execute() + stream = prod_ex.conn.stream_register(prod_ex.execution_id, 0) + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 0, "v0") + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 1, "v1") + + # Waiting for a sequence the stream already holds. + _suspended_consumer(ctx, stream["id"], sequence=1) + + cons2 = ctx.executor.next_execute() + assert cons2.target == "consumer" + cons2.conn.complete(cons2.execution_id) + + prod_ex.conn.stream_close(prod_ex.execution_id, 0) + prod_ex.conn.complete(prod_ex.execution_id, value="done") + + +def test_suspended_consumer_does_not_time_out_the_producer(worker): + """A consumer's nap is not the producer being idle. + + The mirror of the rule that a suspended *producer* isn't idle: while + every consumer of a stream is suspended waiting on it, the producer's + idle countdown is paused. Without this, a lockstep producer would be + force-closed by the very consumers waiting for it — they stop + acknowledging when they suspend, so it can't emit, so its idle window + runs out. + """ + targets = [workflow("test", "producer"), workflow("test", "consumer")] + + with worker(targets, concurrency=3) as ctx: + ctx.submit("test", "producer") + prod_ex = ctx.executor.next_execute() + stream = prod_ex.conn.stream_register( + prod_ex.execution_id, 0, buffer=None, timeout_ms=150 + ) + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 0, "v0") + + _suspended_consumer(ctx, stream["id"], sequence=1) + + # Well past the 150ms window. The countdown is paused, so no + # force-close arrives and the stream is still open to append. + with pytest.raises(TimeoutError): + prod_ex.conn.recv_push("stream_force_close", timeout=1.5) + + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 1, "v1") + + # Appending released the waiter, which also restarts the countdown. + cons2 = ctx.executor.next_execute() + assert cons2.target == "consumer" + cons2.conn.complete(cons2.execution_id) + + force = prod_ex.conn.recv_push("stream_force_close", timeout=2) + assert force["reason"] == "timeout" + prod_ex.conn.complete(prod_ex.execution_id) + + +def test_stream_select_reports_whether_the_sequence_is_available(worker): + """The consumer asks the server rather than trusting its own queue. + + Items arrive asynchronously, so an empty queue means "nothing has + arrived yet", not "the stream has nothing" — a consumer that decided + for itself would suspend before hearing anything, be rescheduled at + once because the item was there all along, and never make progress. + """ + targets = [workflow("test", "producer"), workflow("test", "consumer")] + + with worker(targets, concurrency=3) as ctx: + ctx.submit("test", "producer") + prod_ex = ctx.executor.next_execute() + stream = prod_ex.conn.stream_register(prod_ex.execution_id, 0) + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 0, "v0") + + ctx.submit("test", "consumer") + cons_ex = ctx.executor.next_execute() + cons_ex.conn.stream_subscribe( + cons_ex.execution_id, subscription_id=1, stream_id=stream["id"] + ) + # Appends are fire-and-forget; the push confirms it landed. + cons_ex.conn.recv_push("stream_items", subscription_id=1, timeout=2) + + # Sequence 0 exists, so the poll resolves — no value, just "stop + # waiting"; the item reaches the consumer via its subscription. + resolved = cons_ex.conn.select( + cons_ex.execution_id, + [{"type": "stream", "id": stream["id"], "sequence": 0}], + timeout_ms=0, + suspend=False, + ) + assert resolved["winner"] == 0 + assert resolved["status"] == "ok" + + # Sequence 1 doesn't yet, and a poll never suspends — it reports + # nothing and leaves the decision to the consumer. + assert ( + cons_ex.conn.select( + cons_ex.execution_id, + [{"type": "stream", "id": stream["id"], "sequence": 1}], + timeout_ms=0, + suspend=False, + ) + is None + ) + + # A closed stream will never reach it, so that resolves too. + prod_ex.conn.stream_close(prod_ex.execution_id, 0) + cons_ex.conn.recv_push("stream_closed", subscription_id=1, timeout=2) + closed = cons_ex.conn.select( + cons_ex.execution_id, + [{"type": "stream", "id": stream["id"], "sequence": 1}], + timeout_ms=0, + suspend=False, + ) + assert closed["status"] == "ok" + + cons_ex.conn.complete(cons_ex.execution_id) + prod_ex.conn.complete(prod_ex.execution_id, value="done") + + +def test_stream_dependency_reported_for_each_consumer_attempt(worker): + """The lineage edge reaches an already-open topic on every attempt. + + A consumer that suspends mid-stream has its wait recorded against the + successor before that successor runs, so the successor's subscribe + finds the row already there. With neither the wait nor that subscribe + announced, a topic open for the run showed the dependency against the + first attempt only — while a reload, built from the snapshot, showed + it against all of them. + """ + targets = [workflow("test", "producer"), workflow("test", "consumer")] + + with worker(targets, concurrency=3) as ctx: + ctx.submit("test", "producer") + prod_ex = ctx.executor.next_execute() + stream = prod_ex.conn.stream_register(prod_ex.execution_id, 0) + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 0, "v0") + + cons_resp = ctx.submit("test", "consumer") + cons_ex = ctx.executor.next_execute() + + # Open the topic before anything subscribes, so every edge below has + # to arrive as an update rather than being read from the snapshot. + ctx.inspect(cons_resp["runId"]) + + cons_ex.conn.stream_subscribe( + cons_ex.execution_id, subscription_id=1, stream_id=stream["id"] + ) + cons_ex.conn.suspend(cons_ex.execution_id, stream_wait=(stream["id"], 1)) + + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 1, "v1") + + cons2 = ctx.executor.next_execute() + assert cons2.target == "consumer" + cons2.conn.stream_subscribe( + cons2.execution_id, subscription_id=1, stream_id=stream["id"] + ) + cons2.conn.complete(cons2.execution_id) + # Ordered on the one connection, so a result means the subscribe + # before it has been processed — and its notification sent. + ctx.result(cons_resp["runId"]) + + prod_ex.conn.stream_close(prod_ex.execution_id, 0) + prod_ex.conn.complete(prod_ex.execution_id, value="done") + + expected = { + stream["id"]: { + "type": "stream", + "streamId": stream["id"], + "module": "test", + "target": "producer", + } + } + _, step = next(iter(ctx.inspect(cons_resp["runId"])["steps"].items())) + assert sorted(step["executions"]) == ["1", "2"] + for attempt, execution in step["executions"].items(): + assert execution["dependencies"] == expected, f"attempt {attempt}"