diff --git a/docs/deployment_scheduling.md b/docs/deployment_scheduling.md index 06de1972..35ee77d7 100644 --- a/docs/deployment_scheduling.md +++ b/docs/deployment_scheduling.md @@ -191,8 +191,11 @@ await schedules.update_schedule( - One composition's stores nested inside another live composition's store subtree (without sharing its identical roots) is unsupported operator error; shared-store cross layouts are rejected outright. -- Manual runs and resumes bypass scheduler capacity by design; - capacity governs scheduled dispatch only. +- Manual runs bypass scheduler capacity by design; capacity governs + scheduled dispatch plus scheduled resumes. Resuming a scheduled + interrupted run acquires a server execution slot first (rejected + while saturated, without dispatching) and releases it when the + resumed result is persisted. - A set `max_steps` budget cannot be cleared back to unset through update (recreate the schedule for an unbounded budget). - MCP-backed servers reject scheduler enablement for now. diff --git a/docs/superpowers/specs/2026-09-08-deployment-scheduling-design.md b/docs/superpowers/specs/2026-09-08-deployment-scheduling-design.md index dccdfd7f..f06609df 100644 --- a/docs/superpowers/specs/2026-09-08-deployment-scheduling-design.md +++ b/docs/superpowers/specs/2026-09-08-deployment-scheduling-design.md @@ -337,11 +337,16 @@ all mutating admin ops clear the old revision's unadmitted work and advance the watermark BEFORE the revision bump or flag flip lands, so a crash can only leave the op unapplied (retryable), never a new revision that backfills. Occurrence pages carry the stored history plus a live -held-candidate `pending` synthesis on the first page. Manual runs and -resumes bypass scheduler capacity by design (unchanged API behavior); -scheduler capacity governs scheduled dispatch only, and a scheduled -interrupted run resumed manually reconciles its terminal history through -recovery. +held-candidate `pending` synthesis on the first page. Manual runs bypass +scheduler capacity by design (unchanged API behavior); capacity governs +scheduled dispatch plus scheduled resumes. A scheduled interrupted run +resumed through the run API acquires a server execution slot through the +scheduler's own accounting before dispatch — rejection leaves no resume +attempt behind — holds the durable executing mark for the re-execution +(visible to capacity and drain like any live execution), and releases +the slot when its stopped result is persisted. A resumed scheduled run +reconciles its terminal history live through the same idempotent +recording as dispatch; restart recovery still repairs torn boundaries. Known limitations: pointing one composition's stores inside another live composition's store subtree (without sharing its identical roots) is diff --git a/src/wf_api/runs.py b/src/wf_api/runs.py index cfd5030d..97c99705 100644 --- a/src/wf_api/runs.py +++ b/src/wf_api/runs.py @@ -66,10 +66,16 @@ class WorkflowRunApi: context: WorkflowOperationContext, *, resume_locks: AsyncKeyedLock | None = None, + resume_slot_gate: Any | None = None, ) -> None: self.context = context self.deployments = WorkflowDeploymentApi(context) self._resume_locks = resume_locks or AsyncKeyedLock() + # Optional scheduler resume gate (installed by the opt-in server + # scheduler composition): schedule-owned resumes acquire a shared + # execution slot through it. Genuinely manual runs never consult + # it, and a missing gate keeps the legacy path unchanged. + self.resume_slot_gate = resume_slot_gate def _run_store(self) -> RunStore: if self.context.run_store is None: @@ -189,6 +195,38 @@ class WorkflowRunApi: ) -> RunResult: trace_values = _trace_range_values(trace_range) store = self._run_store() + # Shared execution slot for schedule-owned resumes: the resume + # gate holds the scheduler's own capacity accounting (no second + # semaphore) behind the durable executing mark. Acquisition runs + # BEFORE the ACTIVE attempt mark, so a busy rejection leaves no + # fake ACTIVE attempt for work that never dispatched. Manual runs + # (no schedule admission, or no live scheduler) skip the gate. + gate = self.resume_slot_gate + slot = gate.acquire(run_id) if gate is not None else None + slot_held = slot is not None + try: + return await self._resume_scheduled_or_manual( + run_id=run_id, + resume_payload=resume_payload, + resume_outcome=resume_outcome, + trace_range=trace_range, + trace_values=trace_values, + store=store, + ) + finally: + if slot_held and gate is not None: + gate.release(run_id) + + async def _resume_scheduled_or_manual( + self, + *, + run_id: str, + resume_payload: dict[str, Any], + resume_outcome: str, + trace_range: TraceRangeLike | None, + trace_values: tuple[int, int] | None, + store: RunStore, + ) -> RunResult: pre_attempt = store.get_resume_attempt(run_id) if pre_attempt is not None and pre_attempt.state == "ACTIVE": raise ValueError( diff --git a/src/wf_scheduling/lifecycle.py b/src/wf_scheduling/lifecycle.py index 2cde31cd..0fbbc81a 100644 --- a/src/wf_scheduling/lifecycle.py +++ b/src/wf_scheduling/lifecycle.py @@ -22,6 +22,7 @@ from __future__ import annotations import asyncio import threading +import time from collections.abc import Callable, Coroutine from dataclasses import dataclass, field from datetime import UTC, datetime @@ -219,9 +220,11 @@ class SchedulerService: dispatch happens only with a provably free slot, an executing run keeps its slot until it stops, and every dispatch spawns exactly one task that settles or abandons exactly once — so live execution tasks - can never exceed capacity. Manual runs and resumes bypass the - scheduler entirely (unchanged API behavior): scheduler capacity - governs scheduled dispatch only. + can never exceed capacity. Scheduled resumes of schedule-owned runs + acquire the same slot through the resume gate (same load count, same + durable executing mark, tracked for drain); genuinely manual runs — + those with no schedule admission — bypass the scheduler entirely + (unchanged API behavior). """ schedule_store: Any @@ -248,6 +251,9 @@ class SchedulerService: _executions: set[asyncio.Future[Any]] = field( default_factory=set, init=False, repr=False, compare=False ) + _live_resumes: set[str] = field( + default_factory=set, init=False, repr=False, compare=False + ) _stopping: bool = field(default=False, init=False, repr=False, compare=False) _started: bool = field(default=False, init=False, repr=False, compare=False) _tick_count: int = field(default=0, init=False, repr=False, compare=False) @@ -359,11 +365,19 @@ class SchedulerService: # Join an in-flight poll (a cancelled to_thread wait keeps running # in its worker): its dispatches are registered tasks by now. await asyncio.to_thread(self._join_tick) + # In-flight scheduled resumes (API-driven, slot held) join the + # same grace window: they release their executing mark on their + # stopped persist, so the drain usually just observes them + # finish. Anything still holding a slot afterwards keeps its + # durable executing mark for startup recovery to abandon + # truthfully. + grace_deadline = time.monotonic() + self.config.drain_grace_s + while self._live_resumes and time.monotonic() < grace_deadline: + await asyncio.sleep(0.02) pending = list(self._executions) if pending: - _, still_pending = await asyncio.wait( - pending, timeout=self.config.drain_grace_s - ) + remaining = max(0.0, grace_deadline - time.monotonic()) + _, still_pending = await asyncio.wait(pending, timeout=remaining) for task in still_pending: task.cancel() if still_pending: diff --git a/src/wf_scheduling/poll.py b/src/wf_scheduling/poll.py index 0705450e..92ab7e91 100644 --- a/src/wf_scheduling/poll.py +++ b/src/wf_scheduling/poll.py @@ -190,11 +190,19 @@ class Scheduler: return admission.schedule_id def _task_load(self) -> int: - from wf_scheduling.recovery import _is_pending + from wf_scheduling.recovery import _is_pending, is_executing count = 0 for run in self.run_store.list_runs(): if self._status_value(run) != "admitted": + # A mid-resume scheduled run is interrupted but holds the + # durable executing mark: it occupies a live execution + # slot exactly like a dispatched run (B2), so the shared + # capacity gate must count it. + if is_executing(self.run_store, run.id) and not _is_pending( + self.run_store, run.id + ): + count += 1 continue if _is_pending(self.run_store, run.id): continue diff --git a/src/wf_scheduling/resume_gate.py b/src/wf_scheduling/resume_gate.py new file mode 100644 index 00000000..77ff35f2 --- /dev/null +++ b/src/wf_scheduling/resume_gate.py @@ -0,0 +1,91 @@ +"""Capacity-gated resume slots for schedule-owned runs (B2). + +A scheduled interrupted run occupies its schedule's overlap slot while +waiting, but the re-execution itself needs a live server execution slot — +the same slot scheduled dispatch consumes. Without a gate, a manual +``resume_run`` call on a schedule-owned run would execute even when every +slot is occupied, oversubscribing the bound the scheduler promises. + +:class:`SchedulerResumeGate` closes that hole without a second semaphore: +acquisition checks the scheduler's own load (executing-marked runs, +including mid-resume work) against the same capacity the poll loop uses, +and holds the durable executing mark for the whole re-execution, so the +poll loop observes the resume exactly like any other live execution. +Genuinely manual runs (no schedule admission) never touch the gate, and a +rejected acquisition happens before any resume-attempt mark, so a busy +rejection leaves no fake ACTIVE attempt behind. +""" + +from __future__ import annotations + +from typing import Any + + +class ScheduledCapacityBusyError(ValueError): + """A scheduled resume found every server execution slot occupied.""" + + +class SchedulerResumeGate: + """Execution-slot gate binding scheduled resumes to scheduler capacity. + + Bound to one live :class:`SchedulerService`; the workflow run API + consults it when the run being resumed carries a schedule admission. + All state changes happen under the service lock, so the + check-then-mark sequence is atomic with scheduler ticks. + """ + + def __init__(self, service: Any) -> None: + self._service = service + + def acquire(self, run_id: str) -> Any | None: + """Take the shared execution slot for a scheduled resume. + + Returns the schedule admission when the slot is held (the caller + must :meth:`release` it), or ``None`` when this run is genuinely + manual or the scheduler is not live — both keep the legacy + resume path unchanged. Raises + :class:`ScheduledCapacityBusyError` when every slot is occupied; + nothing is marked in that case, in particular no ACTIVE resume + attempt. + """ + service = self._service + with service._lock: + if not service.running: + return None + try: + admission = service.run_store.get_admission(run_id) + except KeyError: + return None + if admission.schedule_id is None: + return None + try: + record = service.run_store.get_run(run_id) + except KeyError: + return None + status = getattr(record, "status", None) + if getattr(status, "value", status) != "interrupted": + # Not resumable: leave slot accounting alone so the normal + # path reports the proper validation error, not a busy slot. + return None + scheduler = service._scheduler + if scheduler is None: + return None + if scheduler._task_load() >= service.config.capacity: + raise ScheduledCapacityBusyError( + f"server execution capacity is saturated " + f"({service.config.capacity} slot(s) busy); " + f"scheduled resume of {run_id!r} rejected without dispatch" + ) + # Durable executing transition, same as dispatch: a crash from + # here on abandons truthfully through the normal marker rules, + # and the poll loop counts the resume as live load. + service.run_store.mark_executing(run_id) + service._live_resumes.add(run_id) + return admission + + def release(self, run_id: str) -> None: + """Free the slot held by :meth:`acquire` after the stopped persist.""" + service = self._service + with service._lock: + service.run_store.clear_executing(run_id) + service._live_resumes.discard(run_id) diff --git a/src/wf_server/scheduling.py b/src/wf_server/scheduling.py index fab873ea..feb5703d 100644 --- a/src/wf_server/scheduling.py +++ b/src/wf_server/scheduling.py @@ -15,6 +15,7 @@ from contextlib import asynccontextmanager from wf_config import WorkflowConfigFile from wf_scheduling.lifecycle import SchedulerService, SchedulerServiceConfig from wf_scheduling.ownership import SchedulerOwnership +from wf_scheduling.resume_gate import SchedulerResumeGate from wf_scheduling.store import FileScheduleStore from .context import WorkflowServer @@ -31,8 +32,14 @@ def build_scheduler_service( composition root (server run data lives at ``/runs``, schedules at ``/schedules``), so one lock at ``/scheduler.lock`` covers both stores. + + Also installs the resume gate on the server's run API, so + schedule-owned resumes acquire the same execution slot scheduled + dispatch uses. The gate is inert until the service runs (and only + ever applies to schedule-owned runs); manual runs keep the legacy + path either way. """ - return SchedulerService( + service = SchedulerService( schedule_store=FileScheduleStore(server.config.store_root), run_store=server.stores.run_store, runtime=server.context.runtime, @@ -40,6 +47,8 @@ def build_scheduler_service( ownership=SchedulerOwnership(server.config.store_root, owner=SCHEDULER_OWNER), config=config, ) + server.api.runs.resume_slot_gate = SchedulerResumeGate(service) + return service def server_scheduler_config( diff --git a/tests/wf_server/test_scheduler_integration.py b/tests/wf_server/test_scheduler_integration.py index a4abb2c2..0449f015 100644 --- a/tests/wf_server/test_scheduler_integration.py +++ b/tests/wf_server/test_scheduler_integration.py @@ -43,6 +43,7 @@ from wf_platform import ( ) from wf_scheduling.lifecycle import SchedulerServiceConfig, SchedulerStartupError from wf_scheduling.models import Schedule +from wf_scheduling.resume_gate import ScheduledCapacityBusyError from wf_scheduling.store import FileScheduleStore, ScheduleExistsError from wf_server import WorkflowServer, build_local_static_workflow_server from wf_server.scheduling import build_scheduler_service @@ -984,3 +985,278 @@ async def test_paused_deleted_run_completion( finally: _gate_open.set() await deleter.stop() + + +def _interrupt_then_fail_plan(name: str) -> RawWorkflowPlan: + """Approval interrupt whose resumed continuation raises (resume failure).""" + return RawWorkflowPlan.model_validate( + { + "name": name, + "input_schema": {"type": "object", "properties": {}}, + "state_schema": {"type": "object", "properties": {}}, + "output_schema": {"type": "object", "properties": {}}, + "outcomes": ["submitted"], + "start": "ask", + "nodes": [ + {"id": "ask", "type": "interrupt", "kind": "approval"}, + { + "id": "fail", + "type": "node", + "node": "wf.std.runtime_error", + "input": [ + { + "value": "boom on resume", + "target": {"root": "local", "parts": ["message"]}, + } + ], + "output": [], + }, + {"id": "end_submitted", "type": "end", "outcome": "submitted"}, + ], + "edges": [ + {"from": "ask", "outcome": "submitted", "to": "fail"}, + {"from": "fail", "outcome": "ok", "to": "end_submitted"}, + ], + "output": [], + } + ) + + +def _interrupt_then_gate_plan(name: str) -> RawWorkflowPlan: + """Approval interrupt whose resumed continuation blocks on the test gate.""" + return RawWorkflowPlan.model_validate( + { + "name": name, + "input_schema": {"type": "object", "properties": {}}, + "state_schema": { + "type": "object", + "properties": {"result": {"type": "string"}}, + }, + "output_schema": { + "type": "object", + "properties": {"result": {"type": "string"}}, + }, + "outcomes": ["submitted"], + "start": "ask", + "nodes": [ + {"id": "ask", "type": "interrupt", "kind": "approval"}, + { + "id": "gate", + "type": "node", + "node": "test.gate.wait", + "input": [ + { + "value": "hello", + "target": {"root": "local", "parts": ["value"]}, + } + ], + "output": [ + { + "source": {"root": "local", "parts": ["value"]}, + "target": {"root": "state", "parts": ["result"]}, + } + ], + }, + {"id": "end_submitted", "type": "end", "outcome": "submitted"}, + ], + "edges": [ + {"from": "ask", "outcome": "submitted", "to": "gate"}, + {"from": "gate", "outcome": "ok", "to": "end_submitted"}, + ], + "output": [], + } + ) + + +async def test_scheduled_resume_rejected_when_capacity_saturated( + tmp_path: Path, +) -> None: + """A scheduled resume acquires the shared execution slot (B2). + + With the only slot occupied by a gated scheduled run, resuming the + interrupted scheduled run is rejected BEFORE any dispatch: no ACTIVE + attempt, no executing mark, both runs untouched. Freeing the slot + lets the same resume through, releasing the slot afterwards. + """ + _gate_open.clear() + root = tmp_path / "store" + server = build_local_static_workflow_server(root, extra_sources=_gate_sources()) + await _seed(server, "ask1", "ask1.default", _single_interrupt_plan("ask1"), ["submitted"]) + await _seed(server, "gated", "gated.default", _gate_plan("gated"), ["ok"]) + store = FileScheduleStore(root) + store.create_schedule(_one_shot("asker", "ask1.default", datetime.now(UTC))) + service = _scheduler(server, capacity=1) + try: + await service.start() + await _wait_for(lambda: len(_run_ids(root)) == 1) + ask_id = _run_ids(root)[0] + await _wait_for( + lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted" + ) + store.create_schedule(_one_shot("blocker", "gated.default", datetime.now(UTC))) + await _wait_for(lambda: len(_run_ids(root)) == 2) + gate_id = [rid for rid in _run_ids(root) if rid != ask_id][0] + await _wait_for(lambda: FileRunStore(root).is_executing(gate_id)) + + with pytest.raises(ScheduledCapacityBusyError): + await server.api.resume_run( + run_id=ask_id, resume_payload={}, resume_outcome="submitted" + ) + assert FileRunStore(root).get_resume_attempt(ask_id) is None + assert not FileRunStore(root).is_executing(ask_id) + assert FileRunStore(root).get_run(ask_id).status.value == "interrupted" + assert FileRunStore(root).is_executing(gate_id) + assert len(_entries(root, "asker", "completed")) == 0 + + _gate_open.set() + await _wait_for( + lambda: FileRunStore(root).get_run(gate_id).status.value == "completed" + ) + resumed = await server.api.resume_run( + run_id=ask_id, resume_payload={}, resume_outcome="submitted" + ) + assert resumed["status"] == "completed" + attempt = FileRunStore(root).get_resume_attempt(ask_id) + assert attempt is not None and attempt.state == "DONE" + assert not FileRunStore(root).is_executing(ask_id) + finally: + _gate_open.set() + await service.stop() + + +async def test_scheduled_resume_releases_slot_on_reinterruption( + tmp_path: Path, +) -> None: + """A scheduled resume that interrupts again frees its slot (B2). + + The re-interruption stays durably resumable (attempt DONE, new + checkpoint) and a second resume of the same run completes. + """ + root = tmp_path / "store" + server = build_local_static_workflow_server(root) + await _seed( + server, + "ask_twice", + "ask_twice.default", + _double_interrupt_plan("ask_twice"), + ["submitted"], + ) + store = FileScheduleStore(root) + store.create_schedule(_one_shot("asker", "ask_twice.default", datetime.now(UTC))) + service = _scheduler(server, capacity=1) + try: + await service.start() + await _wait_for(lambda: len(_run_ids(root)) == 1) + ask_id = _run_ids(root)[0] + await _wait_for( + lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted" + ) + first = await server.api.resume_run( + run_id=ask_id, resume_payload={}, resume_outcome="submitted" + ) + assert first["status"] == "interrupted" + attempt = FileRunStore(root).get_resume_attempt(ask_id) + assert attempt is not None and attempt.state == "DONE" + assert not FileRunStore(root).is_executing(ask_id) + second = await server.api.resume_run( + run_id=ask_id, resume_payload={}, resume_outcome="submitted" + ) + assert second["status"] == "completed" + assert not FileRunStore(root).is_executing(ask_id) + finally: + await service.stop() + + +async def test_scheduled_resume_releases_slot_on_failure(tmp_path: Path) -> None: + """A scheduled resume that fails releases its slot (B2). + + The failed stopped result is durably persisted with its attempt + cleared, and the freed slot admits new scheduled work. + """ + root = tmp_path / "store" + server = build_local_static_workflow_server(root) + await _seed( + server, + "ask_fail", + "ask_fail.default", + _interrupt_then_fail_plan("ask_fail"), + ["submitted"], + ) + await _seed( + server, + "const", + "const.default", + _constant_plan("const"), + ["ok"], + ) + store = FileScheduleStore(root) + store.create_schedule(_one_shot("asker", "ask_fail.default", datetime.now(UTC))) + service = _scheduler(server, capacity=1) + try: + await service.start() + await _wait_for(lambda: len(_run_ids(root)) == 1) + ask_id = _run_ids(root)[0] + await _wait_for( + lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted" + ) + resumed = await server.api.resume_run( + run_id=ask_id, resume_payload={}, resume_outcome="submitted" + ) + assert resumed["status"] == "failed" + attempt = FileRunStore(root).get_resume_attempt(ask_id) + assert attempt is not None and attempt.state == "DONE" + assert not FileRunStore(root).is_executing(ask_id) + # The freed slot admits new scheduled work. + store.create_schedule(_one_shot("after", "const.default", datetime.now(UTC))) + await _wait_for(lambda: len(_run_ids(root)) == 2) + after_id = [rid for rid in _run_ids(root) if rid != ask_id][0] + await _wait_for( + lambda: FileRunStore(root).get_run(after_id).status.value == "completed" + ) + finally: + await service.stop() + + +async def test_service_stop_drains_inflight_scheduled_resume(tmp_path: Path) -> None: + """Shutdown waits for an in-flight scheduled resume (B2 drain).""" + _gate_open.clear() + root = tmp_path / "store" + server = build_local_static_workflow_server(root, extra_sources=_gate_sources()) + await _seed( + server, + "ask_gate", + "ask_gate.default", + _interrupt_then_gate_plan("ask_gate"), + ["submitted"], + ) + store = FileScheduleStore(root) + store.create_schedule(_one_shot("asker", "ask_gate.default", datetime.now(UTC))) + service = _scheduler(server, capacity=2, drain_grace_s=10.0) + try: + await service.start() + await _wait_for(lambda: len(_run_ids(root)) == 1) + ask_id = _run_ids(root)[0] + await _wait_for( + lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted" + ) + resume_task = asyncio.create_task( + server.api.resume_run( + run_id=ask_id, resume_payload={}, resume_outcome="submitted" + ) + ) + await _wait_for(lambda: ask_id in service._live_resumes) + await _wait_for(lambda: FileRunStore(root).is_executing(ask_id)) + stop_task = asyncio.create_task(service.stop()) + await asyncio.sleep(0.3) + # The gate is still closed: the drain must still be waiting. + assert not stop_task.done() + assert not resume_task.done() + _gate_open.set() + report = await asyncio.wait_for(stop_task, timeout=20.0) + resumed = await asyncio.wait_for(resume_task, timeout=20.0) + assert resumed["status"] == "completed" + assert not FileRunStore(root).is_executing(ask_id) + assert report.cancelled == 0 + finally: + _gate_open.set() + await service.stop()