sched: gate scheduled resumes on the shared execution slot, with drain tracking (B2)

This commit is contained in:
lda
2026-09-09 17:23:28 +07:00 Verified
parent fd2f85ef96
commit ada0239a1c
8 changed files with 459 additions and 15 deletions
+5 -2
View File
@@ -191,8 +191,11 @@ await schedules.update_schedule(
- One composition's stores nested inside another live composition's - One composition's stores nested inside another live composition's
store subtree (without sharing its identical roots) is unsupported store subtree (without sharing its identical roots) is unsupported
operator error; shared-store cross layouts are rejected outright. operator error; shared-store cross layouts are rejected outright.
- Manual runs and resumes bypass scheduler capacity by design; - Manual runs bypass scheduler capacity by design; capacity governs
capacity governs scheduled dispatch only. 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 - A set `max_steps` budget cannot be cleared back to unset through
update (recreate the schedule for an unbounded budget). update (recreate the schedule for an unbounded budget).
- MCP-backed servers reject scheduler enablement for now. - MCP-backed servers reject scheduler enablement for now.
@@ -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 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 crash can only leave the op unapplied (retryable), never a new revision
that backfills. Occurrence pages carry the stored history plus a live that backfills. Occurrence pages carry the stored history plus a live
held-candidate `pending` synthesis on the first page. Manual runs and held-candidate `pending` synthesis on the first page. Manual runs bypass
resumes bypass scheduler capacity by design (unchanged API behavior); scheduler capacity by design (unchanged API behavior); capacity governs
scheduler capacity governs scheduled dispatch only, and a scheduled scheduled dispatch plus scheduled resumes. A scheduled interrupted run
interrupted run resumed manually reconciles its terminal history through resumed through the run API acquires a server execution slot through the
recovery. 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 Known limitations: pointing one composition's stores inside another live
composition's store subtree (without sharing its identical roots) is composition's store subtree (without sharing its identical roots) is
+38
View File
@@ -66,10 +66,16 @@ class WorkflowRunApi:
context: WorkflowOperationContext, context: WorkflowOperationContext,
*, *,
resume_locks: AsyncKeyedLock | None = None, resume_locks: AsyncKeyedLock | None = None,
resume_slot_gate: Any | None = None,
) -> None: ) -> None:
self.context = context self.context = context
self.deployments = WorkflowDeploymentApi(context) self.deployments = WorkflowDeploymentApi(context)
self._resume_locks = resume_locks or AsyncKeyedLock() 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: def _run_store(self) -> RunStore:
if self.context.run_store is None: if self.context.run_store is None:
@@ -189,6 +195,38 @@ class WorkflowRunApi:
) -> RunResult: ) -> RunResult:
trace_values = _trace_range_values(trace_range) trace_values = _trace_range_values(trace_range)
store = self._run_store() 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) pre_attempt = store.get_resume_attempt(run_id)
if pre_attempt is not None and pre_attempt.state == "ACTIVE": if pre_attempt is not None and pre_attempt.state == "ACTIVE":
raise ValueError( raise ValueError(
+20 -6
View File
@@ -22,6 +22,7 @@ from __future__ import annotations
import asyncio import asyncio
import threading import threading
import time
from collections.abc import Callable, Coroutine from collections.abc import Callable, Coroutine
from dataclasses import dataclass, field from dataclasses import dataclass, field
from datetime import UTC, datetime from datetime import UTC, datetime
@@ -219,9 +220,11 @@ class SchedulerService:
dispatch happens only with a provably free slot, an executing run dispatch happens only with a provably free slot, an executing run
keeps its slot until it stops, and every dispatch spawns exactly one keeps its slot until it stops, and every dispatch spawns exactly one
task that settles or abandons exactly once — so live execution tasks task that settles or abandons exactly once — so live execution tasks
can never exceed capacity. Manual runs and resumes bypass the can never exceed capacity. Scheduled resumes of schedule-owned runs
scheduler entirely (unchanged API behavior): scheduler capacity acquire the same slot through the resume gate (same load count, same
governs scheduled dispatch only. durable executing mark, tracked for drain); genuinely manual runs —
those with no schedule admission — bypass the scheduler entirely
(unchanged API behavior).
""" """
schedule_store: Any schedule_store: Any
@@ -248,6 +251,9 @@ class SchedulerService:
_executions: set[asyncio.Future[Any]] = field( _executions: set[asyncio.Future[Any]] = field(
default_factory=set, init=False, repr=False, compare=False 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) _stopping: bool = field(default=False, init=False, repr=False, compare=False)
_started: 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) _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 # Join an in-flight poll (a cancelled to_thread wait keeps running
# in its worker): its dispatches are registered tasks by now. # in its worker): its dispatches are registered tasks by now.
await asyncio.to_thread(self._join_tick) 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) pending = list(self._executions)
if pending: if pending:
_, still_pending = await asyncio.wait( remaining = max(0.0, grace_deadline - time.monotonic())
pending, timeout=self.config.drain_grace_s _, still_pending = await asyncio.wait(pending, timeout=remaining)
)
for task in still_pending: for task in still_pending:
task.cancel() task.cancel()
if still_pending: if still_pending:
+9 -1
View File
@@ -190,11 +190,19 @@ class Scheduler:
return admission.schedule_id return admission.schedule_id
def _task_load(self) -> int: def _task_load(self) -> int:
from wf_scheduling.recovery import _is_pending from wf_scheduling.recovery import _is_pending, is_executing
count = 0 count = 0
for run in self.run_store.list_runs(): for run in self.run_store.list_runs():
if self._status_value(run) != "admitted": 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 continue
if _is_pending(self.run_store, run.id): if _is_pending(self.run_store, run.id):
continue continue
+91
View File
@@ -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)
+10 -1
View File
@@ -15,6 +15,7 @@ from contextlib import asynccontextmanager
from wf_config import WorkflowConfigFile from wf_config import WorkflowConfigFile
from wf_scheduling.lifecycle import SchedulerService, SchedulerServiceConfig from wf_scheduling.lifecycle import SchedulerService, SchedulerServiceConfig
from wf_scheduling.ownership import SchedulerOwnership from wf_scheduling.ownership import SchedulerOwnership
from wf_scheduling.resume_gate import SchedulerResumeGate
from wf_scheduling.store import FileScheduleStore from wf_scheduling.store import FileScheduleStore
from .context import WorkflowServer from .context import WorkflowServer
@@ -31,8 +32,14 @@ def build_scheduler_service(
composition root (server run data lives at ``<store_root>/runs``, composition root (server run data lives at ``<store_root>/runs``,
schedules at ``<store_root>/schedules``), so one lock at schedules at ``<store_root>/schedules``), so one lock at
``<store_root>/scheduler.lock`` covers both stores. ``<store_root>/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), schedule_store=FileScheduleStore(server.config.store_root),
run_store=server.stores.run_store, run_store=server.stores.run_store,
runtime=server.context.runtime, runtime=server.context.runtime,
@@ -40,6 +47,8 @@ def build_scheduler_service(
ownership=SchedulerOwnership(server.config.store_root, owner=SCHEDULER_OWNER), ownership=SchedulerOwnership(server.config.store_root, owner=SCHEDULER_OWNER),
config=config, config=config,
) )
server.api.runs.resume_slot_gate = SchedulerResumeGate(service)
return service
def server_scheduler_config( def server_scheduler_config(
@@ -43,6 +43,7 @@ from wf_platform import (
) )
from wf_scheduling.lifecycle import SchedulerServiceConfig, SchedulerStartupError from wf_scheduling.lifecycle import SchedulerServiceConfig, SchedulerStartupError
from wf_scheduling.models import Schedule from wf_scheduling.models import Schedule
from wf_scheduling.resume_gate import ScheduledCapacityBusyError
from wf_scheduling.store import FileScheduleStore, ScheduleExistsError from wf_scheduling.store import FileScheduleStore, ScheduleExistsError
from wf_server import WorkflowServer, build_local_static_workflow_server from wf_server import WorkflowServer, build_local_static_workflow_server
from wf_server.scheduling import build_scheduler_service from wf_server.scheduling import build_scheduler_service
@@ -984,3 +985,278 @@ async def test_paused_deleted_run_completion(
finally: finally:
_gate_open.set() _gate_open.set()
await deleter.stop() 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()