sched: fence scheduled resumes at shutdown so none outlive ownership
This commit is contained in:
@@ -96,8 +96,11 @@ retrying it. Corrupt or contradictory records fail closed with
|
||||
diagnostics and block the schedule rather than clearing overlap.
|
||||
|
||||
On shutdown the server stops admission first and drains active tasks
|
||||
within the grace period; anything still running keeps its executing
|
||||
mark, and the next startup recovery abandons it truthfully.
|
||||
within the grace period; new scheduled resumes are rejected for the
|
||||
duration of the drain, and a resume still running past the deadline is
|
||||
cancelled and joined before ownership is released, keeping its marks
|
||||
for the next startup recovery. Anything else still running keeps its
|
||||
executing mark, and the next startup recovery abandons it truthfully.
|
||||
|
||||
## Occurrence inspection
|
||||
|
||||
@@ -194,8 +197,9 @@ await schedules.update_schedule(
|
||||
- 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.
|
||||
while saturated or draining, without dispatching) and releases it when the
|
||||
resumed result is persisted; a drain-cancelled resume keeps its marks
|
||||
for recovery instead.
|
||||
- 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.
|
||||
|
||||
@@ -310,7 +310,13 @@ server capacity default is deployment configuration, with deterministic tests
|
||||
using a small injected limit.
|
||||
|
||||
On shutdown stop admission first and drain active tasks within a configured
|
||||
grace period. Record cancellation/failure when possible; abrupt termination
|
||||
grace period. New scheduled resumes are rejected once shutdown begins
|
||||
(no dispatch, no ungated fallback); in-flight scheduled resumes share
|
||||
the grace window, then are cancelled and joined before ownership is
|
||||
released, so no old execution persists after a new owner takes over.
|
||||
A cancelled resume keeps its executing mark and ACTIVE attempt, and
|
||||
startup recovery fails it closed exactly like a crash mid-resume.
|
||||
Record cancellation/failure when possible; abrupt termination
|
||||
uses startup recovery. Paused/deleted schedule definitions must not prevent
|
||||
run completion or resume from updating retained occurrence history.
|
||||
|
||||
@@ -321,9 +327,11 @@ canonical ownership, recovers without executing, then ticks calendar
|
||||
polling without blocking on long workflows: each dispatch spawns exactly
|
||||
one bounded execution task behind the async-completion seam, and the
|
||||
scheduler's own capacity gate is the execution-slot bound (an executing
|
||||
run keeps its slot until it stops). Shutdown stops admission, drains
|
||||
within `drain_grace_s`, leaves unfinished work under its executing mark
|
||||
for startup recovery to abandon truthfully, and releases ownership last.
|
||||
run keeps its slot until it stops). Shutdown stops admission, rejects
|
||||
new scheduled resumes for the duration of the drain, cancels and joins
|
||||
unfinished scheduled resumes after `drain_grace_s` (cancelled work keeps
|
||||
its executing mark and ACTIVE attempt for startup recovery to abandon
|
||||
truthfully), and releases ownership last.
|
||||
A failed startup releases the lock and raises.
|
||||
|
||||
Administration (`WorkflowApi` schedules methods, `workflow.schedules.*`
|
||||
@@ -344,7 +352,9 @@ 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
|
||||
the slot when its stopped result is persisted. Shutdown drain rejects
|
||||
further scheduled resumes; a resume cancelled by the drain keeps its
|
||||
marks for recovery instead of releasing them. A resumed scheduled run
|
||||
reconciles its terminal history live through the same idempotent
|
||||
recording as dispatch; restart recovery still repairs torn boundaries.
|
||||
|
||||
|
||||
+33
-6
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from dataclasses import asdict
|
||||
from typing import Any, Protocol
|
||||
|
||||
@@ -198,14 +199,20 @@ class WorkflowRunApi:
|
||||
# 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.
|
||||
# BEFORE the ACTIVE attempt mark, so a busy or draining rejection
|
||||
# leaves no fake ACTIVE attempt for work that never dispatched.
|
||||
# The caller's task is bound atomically with acquisition, so the
|
||||
# shutdown drain can cancel and join it after the grace deadline.
|
||||
# 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
|
||||
if gate is not None:
|
||||
slot = gate.acquire(run_id, owner_task=asyncio.current_task())
|
||||
else:
|
||||
slot = None
|
||||
slot_held = slot is not None
|
||||
try:
|
||||
return await self._resume_scheduled_or_manual(
|
||||
result = await self._resume_scheduled_or_manual(
|
||||
run_id=run_id,
|
||||
resume_payload=resume_payload,
|
||||
resume_outcome=resume_outcome,
|
||||
@@ -213,9 +220,29 @@ class WorkflowRunApi:
|
||||
trace_values=trace_values,
|
||||
store=store,
|
||||
)
|
||||
finally:
|
||||
except asyncio.CancelledError:
|
||||
# Shutdown drain cancelled the execution mid-flight (or the
|
||||
# caller went away): fence, don't release. The durable
|
||||
# executing mark and the ACTIVE attempt stay exactly as a
|
||||
# crash mid-resume would leave them, so restart recovery
|
||||
# fails the run closed instead of presenting a half-resumed
|
||||
# run as safe to retry. Cancellation is never swallowed.
|
||||
if slot_held and gate is not None:
|
||||
gate.fence(run_id)
|
||||
raise
|
||||
except BaseException:
|
||||
# Pre-persist errors, torn persists, and validation failures:
|
||||
# release is truthful here because either nothing executed
|
||||
# (the marks were only ever ours) or the durable ACTIVE
|
||||
# attempt / stopped result already captures the ambiguity for
|
||||
# recovery to reconcile by attempt identity.
|
||||
if slot_held and gate is not None:
|
||||
gate.release(run_id)
|
||||
raise
|
||||
else:
|
||||
if slot_held and gate is not None:
|
||||
gate.release(run_id)
|
||||
return result
|
||||
|
||||
async def _resume_scheduled_or_manual(
|
||||
self,
|
||||
|
||||
@@ -224,7 +224,11 @@ class SchedulerService:
|
||||
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).
|
||||
(unchanged API behavior, including while draining). Once shutdown
|
||||
begins the gate rejects new scheduled resumes instead of dispatching
|
||||
them; in-flight resumes get the grace window, then are cancelled and
|
||||
joined before ownership is released, so no old execution can persist
|
||||
after a new owner takes over.
|
||||
"""
|
||||
|
||||
schedule_store: Any
|
||||
@@ -254,6 +258,9 @@ class SchedulerService:
|
||||
_live_resumes: set[str] = field(
|
||||
default_factory=set, init=False, repr=False, compare=False
|
||||
)
|
||||
_resume_tasks: dict[str, asyncio.Task[Any]] = field(
|
||||
default_factory=dict, 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)
|
||||
@@ -346,15 +353,24 @@ class SchedulerService:
|
||||
|
||||
No new tick starts after ``stop`` begins (an already-running poll
|
||||
is joined, and any dispatch it performed is covered by the
|
||||
drain). Live executions get the configured grace period; tasks
|
||||
still running afterwards are cancelled and left under their
|
||||
durable executing mark for startup recovery to abandon
|
||||
truthfully. Ownership is released only after the drain finishes.
|
||||
Safe to call when not started (returns a zero report).
|
||||
drain), and the resume gate rejects new scheduled resumes from
|
||||
that point on. Live executions and in-flight scheduled resumes
|
||||
get the configured grace period; anything still running
|
||||
afterwards is cancelled — executions are left under their durable
|
||||
executing mark, and resumes keep their executing mark plus ACTIVE
|
||||
attempt via the gate fence — for startup recovery to abandon
|
||||
truthfully. Every tracked resume task is joined before ownership
|
||||
is released, so no old execution can mutate the stores after a
|
||||
new owner takes over. Ownership is released only after the drain
|
||||
finishes. Safe to call when not started (returns a zero report).
|
||||
"""
|
||||
if not self._started:
|
||||
return DrainReport()
|
||||
self._stopping = True
|
||||
with self._lock:
|
||||
# Under the service lock, atomically with gate acquisition:
|
||||
# every resume that passed the stopping check is already
|
||||
# tracked, and every later one is rejected.
|
||||
self._stopping = True
|
||||
if self._tick_task is not None:
|
||||
self._tick_task.cancel()
|
||||
try:
|
||||
@@ -374,6 +390,26 @@ class SchedulerService:
|
||||
grace_deadline = time.monotonic() + self.config.drain_grace_s
|
||||
while self._live_resumes and time.monotonic() < grace_deadline:
|
||||
await asyncio.sleep(0.02)
|
||||
# Fence: cancel and join every still-tracked resume task BEFORE
|
||||
# releasing ownership. Joining first is what keeps a new owner
|
||||
# safe: a joined task cannot persist afterwards. Cancelled tasks
|
||||
# unwind through the gate fence (in-memory accounting dropped,
|
||||
# durable marks kept), so recovery sees the exact crash shape.
|
||||
current = asyncio.current_task()
|
||||
with self._lock:
|
||||
resume_pending = [
|
||||
task
|
||||
for task in self._resume_tasks.values()
|
||||
if task is not current and not task.done()
|
||||
]
|
||||
for task in resume_pending:
|
||||
task.cancel()
|
||||
if resume_pending:
|
||||
await asyncio.gather(*resume_pending, return_exceptions=True)
|
||||
with self._lock:
|
||||
self._resume_tasks.clear()
|
||||
self._live_resumes.clear()
|
||||
resume_cancelled = len([t for t in resume_pending if t.cancelled()])
|
||||
pending = list(self._executions)
|
||||
if pending:
|
||||
remaining = max(0.0, grace_deadline - time.monotonic())
|
||||
@@ -387,7 +423,7 @@ class SchedulerService:
|
||||
# observe the drained set deterministically.
|
||||
for task in pending:
|
||||
self._executions.discard(task)
|
||||
cancelled = len([t for t in pending if t.cancelled()])
|
||||
cancelled = len([t for t in pending if t.cancelled()]) + resume_cancelled
|
||||
with self._lock:
|
||||
report = DrainReport(
|
||||
settled=self._settled,
|
||||
|
||||
@@ -18,6 +18,7 @@ rejection leaves no fake ACTIVE attempt behind.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
|
||||
from wf_scheduling.ownership import SecondOwnerError
|
||||
@@ -27,6 +28,15 @@ class ScheduledCapacityBusyError(ValueError):
|
||||
"""A scheduled resume found every server execution slot occupied."""
|
||||
|
||||
|
||||
class ScheduledResumeShutdownError(ValueError):
|
||||
"""A scheduled resume was requested while the scheduler is draining.
|
||||
|
||||
Raised instead of acquiring once shutdown begins: a draining service
|
||||
must neither dispatch new scheduled work nor fall back to ungated
|
||||
execution, which would bypass capacity and escape drain tracking.
|
||||
"""
|
||||
|
||||
|
||||
class SchedulerResumeGate:
|
||||
"""Execution-slot gate binding scheduled resumes to scheduler capacity.
|
||||
|
||||
@@ -39,16 +49,29 @@ class SchedulerResumeGate:
|
||||
def __init__(self, service: Any) -> None:
|
||||
self._service = service
|
||||
|
||||
def acquire(self, run_id: str) -> Any | None:
|
||||
def acquire(
|
||||
self, run_id: str, *, owner_task: asyncio.Task[Any] | None = None
|
||||
) -> 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, not resumable, or the scheduler is not live — all 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.
|
||||
must :meth:`release` it, or :meth:`fence` it when cancelled), or
|
||||
``None`` when this run is genuinely manual, not resumable, or the
|
||||
scheduler is not live — all keep the legacy resume path
|
||||
unchanged. Raises :class:`ScheduledCapacityBusyError` when every
|
||||
slot is occupied, and :class:`ScheduledResumeShutdownError` once
|
||||
shutdown begins; nothing is marked in either case, in particular
|
||||
no ACTIVE resume attempt. The shutdown check sits after the
|
||||
manual/not-resumable early returns, so manual resumes stay
|
||||
independent of the drain.
|
||||
|
||||
``owner_task`` binds the caller's execution lifetime for the
|
||||
shutdown drain: the service cancels and joins tracked tasks after
|
||||
the grace deadline, before releasing ownership. Registration
|
||||
happens here, under the service lock, atomically with the
|
||||
stopping check — a resume that passes the check is always
|
||||
visible to the drain, and a resume that loses the race is
|
||||
rejected instead of running unfenced.
|
||||
"""
|
||||
service = self._service
|
||||
with service._lock:
|
||||
@@ -72,6 +95,11 @@ class SchedulerResumeGate:
|
||||
scheduler = service._scheduler
|
||||
if scheduler is None:
|
||||
return None
|
||||
if service._stopping:
|
||||
raise ScheduledResumeShutdownError(
|
||||
f"scheduler is draining for shutdown; "
|
||||
f"scheduled resume of {run_id!r} rejected without dispatch"
|
||||
)
|
||||
if scheduler._task_load() >= service.config.capacity:
|
||||
raise ScheduledCapacityBusyError(
|
||||
f"server execution capacity is saturated "
|
||||
@@ -83,14 +111,40 @@ class SchedulerResumeGate:
|
||||
# and the poll loop counts the resume as live load.
|
||||
service.run_store.mark_executing(run_id)
|
||||
service._live_resumes.add(run_id)
|
||||
if owner_task is not None:
|
||||
service._resume_tasks[run_id] = owner_task
|
||||
return admission
|
||||
|
||||
def release(self, run_id: str) -> None:
|
||||
"""Free the slot held by :meth:`acquire` after the stopped persist."""
|
||||
"""Free the slot held by :meth:`acquire` after the stopped persist.
|
||||
|
||||
Correct for every non-cancellation outcome: on success the
|
||||
stopped result plus the DONE attempt already capture the outcome;
|
||||
on pre-persist errors the marks were only ever ours (the durable
|
||||
ACTIVE attempt still records the ambiguity for recovery); on torn
|
||||
persists the durable stopped result plus ACTIVE attempt let
|
||||
recovery reconcile by attempt identity.
|
||||
"""
|
||||
service = self._service
|
||||
with service._lock:
|
||||
service.run_store.clear_executing(run_id)
|
||||
service._live_resumes.discard(run_id)
|
||||
service._resume_tasks.pop(run_id, None)
|
||||
|
||||
def fence(self, run_id: str) -> None:
|
||||
"""Forget a cancelled execution; keep its durable marks for recovery.
|
||||
|
||||
The shutdown drain cancels a resume that outlives the grace
|
||||
deadline and joins it before releasing ownership, so no old
|
||||
execution can persist afterwards. The durable executing mark and
|
||||
the ACTIVE attempt stay exactly as a crash mid-resume would leave
|
||||
them: restart recovery fails the run closed (ambiguous attempt),
|
||||
never silently resumable, never replayed.
|
||||
"""
|
||||
service = self._service
|
||||
with service._lock:
|
||||
service._live_resumes.discard(run_id)
|
||||
service._resume_tasks.pop(run_id, None)
|
||||
|
||||
def note_resumed_result(
|
||||
self,
|
||||
|
||||
@@ -43,7 +43,10 @@ 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.resume_gate import (
|
||||
ScheduledCapacityBusyError,
|
||||
ScheduledResumeShutdownError,
|
||||
)
|
||||
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
|
||||
@@ -1264,6 +1267,280 @@ async def test_service_stop_drains_inflight_scheduled_resume(tmp_path: Path) ->
|
||||
await service.stop()
|
||||
|
||||
|
||||
async def test_scheduled_resume_completes_within_grace_with_history(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""A resume finishing inside grace stops cleanly with history (shutdown).
|
||||
|
||||
Operator-driven service (no background ticks): the in-flight resume
|
||||
completes during the drain, stop reports no cancellation, and the
|
||||
resumed completion lands in occurrence history exactly once.
|
||||
"""
|
||||
_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=1, drain_grace_s=10.0, auto_tick=False)
|
||||
try:
|
||||
await service.start()
|
||||
async with asyncio.timeout(20):
|
||||
while not _run_ids(root):
|
||||
await service.poll_once(datetime.now(UTC))
|
||||
await asyncio.sleep(0.02)
|
||||
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)
|
||||
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()
|
||||
resumed = await asyncio.wait_for(resume_task, timeout=20.0)
|
||||
report = await asyncio.wait_for(stop_task, timeout=20.0)
|
||||
assert resumed["status"] == "completed"
|
||||
assert report.cancelled == 0
|
||||
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)
|
||||
assert await _kinds(root, "asker") == ["admitted", "interrupted", "completed"]
|
||||
finally:
|
||||
_gate_open.set()
|
||||
await service.stop()
|
||||
|
||||
|
||||
async def test_shutdown_timeout_fences_inflight_resume_and_restart_keeps_decision(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""A resume outliving grace cannot overwrite the next owner (shutdown).
|
||||
|
||||
Stop cancels and joins the gated resume before releasing ownership:
|
||||
the task is done-cancelled at stop return, durable state keeps the
|
||||
exact crash shape (interrupted, executing mark, ACTIVE attempt), and
|
||||
after restart recovery fails the run closed the old task cannot flip
|
||||
the decision back to completed or clear the ACTIVE attempt.
|
||||
"""
|
||||
_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=1, drain_grace_s=0.05, auto_tick=False)
|
||||
try:
|
||||
await service.start()
|
||||
async with asyncio.timeout(20):
|
||||
while not _run_ids(root):
|
||||
await service.poll_once(datetime.now(UTC))
|
||||
await asyncio.sleep(0.02)
|
||||
ask_id = _run_ids(root)[0]
|
||||
await _wait_for(
|
||||
lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted"
|
||||
)
|
||||
interrupted_checkpoint = FileRunStore(root).get_run(ask_id).latest_checkpoint_id
|
||||
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))
|
||||
# The gate stays closed: the drain must fence, not wait forever.
|
||||
report = await service.stop()
|
||||
assert report.cancelled == 1
|
||||
assert resume_task.done()
|
||||
assert resume_task.cancelled()
|
||||
# Cancellation truthfulness: crash-shaped durable state, no result.
|
||||
assert FileRunStore(root).get_run(ask_id).status.value == "interrupted"
|
||||
assert FileRunStore(root).is_executing(ask_id)
|
||||
attempt = FileRunStore(root).get_resume_attempt(ask_id)
|
||||
assert attempt is not None and attempt.state == "ACTIVE"
|
||||
assert await _kinds(root, "asker") == ["admitted", "interrupted"]
|
||||
|
||||
revived = _scheduler(server, capacity=1, auto_tick=False)
|
||||
try:
|
||||
await revived.start()
|
||||
# Restart recovery fails the ambiguous resume closed.
|
||||
assert FileRunStore(root).get_run(ask_id).status.value == "failed"
|
||||
attempt2 = FileRunStore(root).get_resume_attempt(ask_id)
|
||||
assert attempt2 is not None and attempt2.state == "ACTIVE"
|
||||
assert not FileRunStore(root).is_executing(ask_id)
|
||||
assert (
|
||||
FileRunStore(root).get_run(ask_id).latest_checkpoint_id
|
||||
== interrupted_checkpoint
|
||||
)
|
||||
assert len(_entries(root, "asker", "failed")) == 1
|
||||
# The fenced old task is already dead: opening the gate and
|
||||
# polling under the new owner changes nothing.
|
||||
_gate_open.set()
|
||||
await revived.poll_once(datetime.now(UTC))
|
||||
assert FileRunStore(root).get_run(ask_id).status.value == "failed"
|
||||
attempt3 = FileRunStore(root).get_resume_attempt(ask_id)
|
||||
assert attempt3 is not None and attempt3.state == "ACTIVE"
|
||||
assert _entries(root, "asker", "completed") == []
|
||||
assert len(_entries(root, "asker", "failed")) == 1
|
||||
finally:
|
||||
_gate_open.set()
|
||||
await revived.stop()
|
||||
finally:
|
||||
_gate_open.set()
|
||||
await service.stop()
|
||||
|
||||
|
||||
async def test_scheduled_resume_requested_during_drain_rejected(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""New scheduled resumes are rejected once shutdown begins (shutdown).
|
||||
|
||||
With a free slot available the rejection still fires (drain, not
|
||||
capacity): no dispatch, no ACTIVE attempt, no executing mark. Manual
|
||||
resumes bypass the gate and complete mid-drain.
|
||||
"""
|
||||
_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"])
|
||||
await _seed(
|
||||
server,
|
||||
"manual_ask",
|
||||
"manual_ask.default",
|
||||
_single_interrupt_plan("manual_ask"),
|
||||
["submitted"],
|
||||
)
|
||||
store = FileScheduleStore(root)
|
||||
store.create_schedule(_one_shot("blocker", "gated.default", datetime.now(UTC)))
|
||||
store.create_schedule(_one_shot("asker", "ask1.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)) == 2)
|
||||
ids = {
|
||||
FileRunStore(root).get_admission(rid).schedule_id or "": rid
|
||||
for rid in _run_ids(root)
|
||||
}
|
||||
gate_id = ids["blocker"]
|
||||
ask_id = ids["asker"]
|
||||
await _wait_for(
|
||||
lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted"
|
||||
)
|
||||
await _wait_for(lambda: FileRunStore(root).is_executing(gate_id))
|
||||
asked = await server.api.run_deployment(
|
||||
deployment_id="manual_ask.default", workflow_input={}
|
||||
)
|
||||
assert asked["status"] == "interrupted"
|
||||
manual_id = asked["run_id"]
|
||||
assert isinstance(manual_id, str)
|
||||
# Hold the drain on the gated execution; one slot stays free, so
|
||||
# only the shutdown check can reject the resume below.
|
||||
stop_task = asyncio.create_task(service.stop())
|
||||
await _wait_for(lambda: service._stopping)
|
||||
assert not stop_task.done()
|
||||
with pytest.raises(ScheduledResumeShutdownError):
|
||||
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"
|
||||
manual_resumed = await server.api.resume_run(
|
||||
run_id=manual_id, resume_payload={}, resume_outcome="submitted"
|
||||
)
|
||||
assert manual_resumed["status"] == "completed"
|
||||
_gate_open.set()
|
||||
report = await asyncio.wait_for(stop_task, timeout=20.0)
|
||||
assert report.cancelled == 0
|
||||
assert FileRunStore(root).get_run(gate_id).status.value == "completed"
|
||||
finally:
|
||||
_gate_open.set()
|
||||
await service.stop()
|
||||
|
||||
|
||||
async def test_shutdown_timeout_spares_healthy_sibling(tmp_path: Path) -> None:
|
||||
"""A fenced resume leaves a healthy sibling's outcome/history intact."""
|
||||
_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"],
|
||||
)
|
||||
await _seed(server, "const", "const.default", _constant_plan("const"), ["ok"])
|
||||
store = FileScheduleStore(root)
|
||||
store.create_schedule(_one_shot("asker", "ask_gate.default", datetime.now(UTC)))
|
||||
store.create_schedule(_one_shot("sib", "const.default", datetime.now(UTC)))
|
||||
service = _scheduler(server, capacity=1, drain_grace_s=0.05, auto_tick=False)
|
||||
try:
|
||||
await service.start()
|
||||
async with asyncio.timeout(20):
|
||||
while len(_run_ids(root)) < 2:
|
||||
await service.poll_once(datetime.now(UTC))
|
||||
await asyncio.sleep(0.02)
|
||||
ids = {
|
||||
FileRunStore(root).get_admission(rid).schedule_id or "": rid
|
||||
for rid in _run_ids(root)
|
||||
}
|
||||
ask_id = ids["asker"]
|
||||
sib_id = ids["sib"]
|
||||
await _wait_for(
|
||||
lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted"
|
||||
)
|
||||
await _wait_for(
|
||||
lambda: FileRunStore(root).get_run(sib_id).status.value == "completed"
|
||||
)
|
||||
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)
|
||||
report = await service.stop()
|
||||
assert report.cancelled == 1
|
||||
assert resume_task.cancelled()
|
||||
assert FileRunStore(root).get_run(sib_id).status.value == "completed"
|
||||
assert await _kinds(root, "sib") == ["admitted", "completed"]
|
||||
|
||||
revived = _scheduler(server, capacity=1, auto_tick=False)
|
||||
try:
|
||||
await revived.start()
|
||||
assert FileRunStore(root).get_run(ask_id).status.value == "failed"
|
||||
assert FileRunStore(root).get_run(sib_id).status.value == "completed"
|
||||
assert await _kinds(root, "sib") == ["admitted", "completed"]
|
||||
assert _entries(root, "asker", "completed") == []
|
||||
finally:
|
||||
_gate_open.set()
|
||||
await revived.stop()
|
||||
finally:
|
||||
_gate_open.set()
|
||||
await service.stop()
|
||||
|
||||
|
||||
async def _kinds(root: Path, schedule_id: str) -> list[str]:
|
||||
page = FileScheduleStore(root).list_occurrences(schedule_id, limit=100)
|
||||
rows = cast(list[dict[str, Any]], page["occurrences"])
|
||||
|
||||
Reference in New Issue
Block a user