From 42deba6399aa2f63f1d62b0a9dcbc4ac6af55ea0 Mon Sep 17 00:00:00 2001 From: lda Date: Wed, 9 Sep 2026 10:53:48 +0700 Subject: [PATCH] sched: server lifecycle service with bounded real execution (T12 core) --- src/wf_scheduling/lifecycle.py | 522 +++++++++++++++++++++++ src/wf_scheduling/ownership.py | 94 +++-- src/wf_scheduling/prepare.py | 8 +- tests/scheduling/test_lifecycle.py | 558 +++++++++++++++++++++++++ tests/scheduling/test_lock_identity.py | 49 ++- tests/scheduling/test_preparation.py | 16 + 6 files changed, 1195 insertions(+), 52 deletions(-) create mode 100644 src/wf_scheduling/lifecycle.py create mode 100644 tests/scheduling/test_lifecycle.py diff --git a/src/wf_scheduling/lifecycle.py b/src/wf_scheduling/lifecycle.py new file mode 100644 index 00000000..0965e6cb --- /dev/null +++ b/src/wf_scheduling/lifecycle.py @@ -0,0 +1,522 @@ +"""Opt-in same-server scheduler lifecycle (T12). + +The :class:`Scheduler` core is synchronous and never owns threads, tasks, +or clocks: this module adds the explicitly enabled server composition +around it. :class:`SchedulerService` acquires the canonical store +ownership, runs startup recovery (which never executes), then ticks +calendar polling on the server event loop without blocking it on +long-running workflows. Real execution runs in bounded asyncio tasks +behind :class:`RuntimeDispatcher`; settlement flows back through the +scheduler's guarded async-completion seam, so capacity accounting, +overlap slots, and exactly-once settlement keep working exactly as the +poll-level tests pin them. + +Shutdown stops admission first (no new ticks), drains live executions +within the configured grace period, leaves anything still running under +its durable executing mark for startup recovery to abandon truthfully, +and releases ownership only after settlement is finished. A failed +startup releases the lock and raises instead of running unprotected. +""" + +from __future__ import annotations + +import asyncio +import threading +from collections.abc import Callable, Coroutine +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import Any + + +def _utcnow() -> datetime: + return datetime.now(UTC) + + +@dataclass(frozen=True, slots=True) +class SchedulerServiceConfig: + """Tuning for the opt-in scheduler service (deployment configuration).""" + + poll_interval_s: float = 1.0 + capacity: int = 4 + drain_grace_s: float = 30.0 + auto_tick: bool = True + + def __post_init__(self) -> None: + if not self.poll_interval_s > 0: + raise ValueError("poll_interval_s must be positive") + if self.capacity < 1: + raise ValueError("capacity must be at least 1") + if not self.drain_grace_s >= 0: + raise ValueError("drain_grace_s must be non-negative") + + +class SchedulerStartupError(Exception): + """Scheduler startup failed; no lock is held and nothing is running.""" + + +@dataclass(frozen=True, slots=True) +class DrainReport: + """Outcome of :meth:`SchedulerService.stop`.""" + + settled: int = 0 + abandoned: int = 0 + cancelled: int = 0 + + +class StoreDeploymentDirectory: + """Production deployment contract source backed by the artifact store. + + ``deployment_revision`` reuses the store-backed deployment revision + (T07), so admission rechecks observe edits; unknown ids raise + ``KeyError``, which preparation reports as a ``deployment-deleted`` + preflight rejection without inventing a run. ``required_inputs`` + reuses the required keys of the pinned root artifact's input schema; + the schema validator remains the hard gate, so an unshaped schema + simply contributes no required keys here. + """ + + def __init__(self, artifact_store: Any) -> None: + self._artifact_store = artifact_store + + def _root_artifact(self, deployment_id: str) -> Any: + deployment = self._artifact_store.get_deployment(deployment_id) + return self._artifact_store.get_artifact( + deployment.artifact_id, deployment.artifact_version + ) + + def deployment_revision(self, deployment_id: str) -> int: + return int(self._artifact_store.get_deployment(deployment_id).revision) + + def required_inputs(self, deployment_id: str) -> list[str]: + artifact = self._root_artifact(deployment_id) + schema = getattr(artifact, "input_schema", None) + if not isinstance(schema, dict): + return [] + required = schema.get("required", []) + if not isinstance(required, list): + return [] + return [key for key in required if isinstance(key, str)] + + +def build_pinned_environment( + artifact_store: Any, +) -> Callable[[Any], Any]: + """Build a production environment pinning callable for a schedule. + + Loads the schedule's deployment, resolves its root artifact and saved + subgraph tree, and freezes them with + :func:`wf_api.run_lifecycle.create_pinned_environment`. Unknown + deployments raise ``KeyError`` (mapped to ``deployment-deleted`` by + preparation). Live dependency validation is deliberately NOT part of + admission: the admission recheck list covers revision, inputs, and + schema, while broken live dependencies surface as failed runs at + execution, exactly like manual runs. + """ + + def build(sched: Any) -> Any: + from wf_api.run_lifecycle import create_pinned_environment + from wf_api.saved_subgraphs import resolve_saved_subgraph_tree + + deployment = artifact_store.get_deployment(sched.deployment_id) + artifact = artifact_store.get_artifact( + deployment.artifact_id, deployment.artifact_version + ) + tree = resolve_saved_subgraph_tree( + root_artifact=artifact, artifact_store=artifact_store + ) + return create_pinned_environment( + deployment=deployment, artifact=artifact, tree=tree + ) + + return build + + +class RuntimeDispatcher: + """Execution collaborator that runs admitted invocations for real. + + ``dispatch`` never blocks the polling tick on workflow execution: it + submits the execution coroutine through the service-provided + ``submit`` hook and immediately returns :class:`StillRunning`. The + coroutine runs the pinned invocation through the server runtime and + settles the genuine stopped state through the service-provided + ``settle`` hook (the scheduler's guarded async-completion seam, which + persists through the shared lifecycle boundary and records terminal + history exactly once). Anything unexpected — an execution failure or + a settlement failure such as a torn write — goes through the + service-provided ``abandon`` hook, which reuses startup recovery to + fail the run closed without replay. Cancellation (shutdown drain) is + re-raised unsettled: the durable executing mark is left for startup + recovery, which abandons it truthfully. + """ + + def __init__( + self, + runtime: Any, + *, + submit: Callable[[Coroutine[Any, Any, None]], Any], + settle: Callable[[str, Any], None], + abandon: Callable[[str, BaseException], None], + ) -> None: + self._runtime = runtime + self._submit = submit + self._settle = settle + self._abandon = abandon + + def dispatch(self, *, admission: Any, now: datetime) -> Any: + from wf_scheduling.dispatch import StillRunning + + self._submit(self._run(admission)) + return StillRunning() + + async def _run(self, admission: Any) -> None: + from wf_api.artifact_plans import raw_plan_from_artifact + from wf_api.saved_subgraphs import saved_subgraph_tree_from_snapshots + from wf_core import RunLimits + + try: + plan = raw_plan_from_artifact(admission.environment.root_artifact) + tree = saved_subgraph_tree_from_snapshots( + admission.environment.child_artifacts + ) + max_steps = admission.max_steps + limits = ( + RunLimits(max_steps=max_steps) if max_steps is not None else RunLimits() + ) + state = await self._runtime.run_workflow_from_plan( + plan, + dict(admission.resolved_input), + deployment=admission.environment.deployment, + artifact=admission.environment.root_artifact, + saved_subgraph_tree=tree, + limits=limits, + ) + except asyncio.CancelledError: + raise + except Exception as exc: + self._abandon(admission.id, exc) + return + try: + self._settle(admission.id, state) + except asyncio.CancelledError: + raise + except Exception as exc: + self._abandon(admission.id, exc) + + +@dataclass(slots=True) +class SchedulerService: + """Explicitly enabled same-server scheduler composition. + + Owns exactly one :class:`Scheduler` built over the server's schedule + and run stores, the canonical :class:`SchedulerOwnership` handle, one + periodic tick task, and the set of live execution tasks. All + scheduler mutations (ticks and settlements) serialize on one lock; the + workflow coroutines themselves run without it, so a long-running + workflow never blocks calendar polling. + + The execution-slot bound is the scheduler's own capacity gate: + 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. + """ + + schedule_store: Any + run_store: Any + runtime: Any + artifact_store: Any + ownership: Any + config: SchedulerServiceConfig = field(default_factory=SchedulerServiceConfig) + clock: Callable[[], datetime] = _utcnow + on_tick_error: Callable[[BaseException], None] | None = None + _lock: threading.Lock = field( + default_factory=threading.Lock, init=False, repr=False, compare=False + ) + _scheduler: Any = field(default=None, init=False, repr=False, compare=False) + _sources: dict[str, Any] = field( + default_factory=dict, init=False, repr=False, compare=False + ) + _loop: asyncio.AbstractEventLoop | None = field( + default=None, init=False, repr=False, compare=False + ) + _tick_task: asyncio.Task[None] | None = field( + default=None, init=False, repr=False, compare=False + ) + _executions: set[asyncio.Future[Any]] = 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) + _settled: int = field(default=0, init=False, repr=False, compare=False) + _abandoned: int = field(default=0, init=False, repr=False, compare=False) + _errors: list[str] = field( + default_factory=list, init=False, repr=False, compare=False + ) + _last_tick_error: BaseException | None = field( + default=None, init=False, repr=False, compare=False + ) + _dispatcher: RuntimeDispatcher = field(init=False, repr=False, compare=False) + _preparer: Any = field(default=None, init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + from wf_scheduling.prepare import SchedulePreparer + + self._preparer = SchedulePreparer( + StoreDeploymentDirectory(self.artifact_store), + build_pinned_environment(self.artifact_store), + ) + self._dispatcher = RuntimeDispatcher( + self.runtime, + submit=self._submit, + settle=self._settle, + abandon=self._abandon, + ) + + # -- observability ------------------------------------------------- + @property + def running(self) -> bool: + """Whether the tick loop is active (started and not stopped).""" + return self._started and not self._stopping + + @property + def tick_count(self) -> int: + """Number of poll ticks attempted since start.""" + return self._tick_count + + @property + def live_executions(self) -> int: + """Number of execution tasks currently tracked for drain.""" + return len(self._executions) + + @property + def last_tick_error(self) -> BaseException | None: + """The most recent tick failure, if any (ticks continue).""" + return self._last_tick_error + + @property + def errors(self) -> list[str]: + """Recent execution/settlement error summaries (capped).""" + return list(self._errors) + + # -- lifecycle ----------------------------------------------------- + async def start(self) -> None: + """Acquire ownership, recover without executing, start polling. + + Raises :class:`SchedulerStartupError` (wrapping second-owner, + unsupported-locking, or corrupt-store failures) without holding + the lock and without leaving tasks behind. + """ + if self._started: + raise SchedulerStartupError("scheduler service already started") + try: + self.ownership.acquire() + except Exception as exc: + raise SchedulerStartupError( + f"scheduler startup cannot prove store ownership: {exc}" + ) from exc + try: + self._build_scheduler() + await asyncio.to_thread(self._recover) + except Exception as exc: + self.ownership.release() + raise SchedulerStartupError( + f"scheduler startup recovery failed: {exc}" + ) from exc + self._loop = asyncio.get_running_loop() + self._stopping = False + self._started = True + # Operator-driven services (auto_tick=False) poll only through + # explicit poll_once calls; the drain and ownership rules are + # identical either way. + if self.config.auto_tick: + self._tick_task = asyncio.get_running_loop().create_task(self._run()) + + async def stop(self) -> DrainReport: + """Stop admission, drain live executions, release ownership. + + 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). + """ + if not self._started: + return DrainReport() + self._stopping = True + if self._tick_task is not None: + self._tick_task.cancel() + try: + await self._tick_task + except asyncio.CancelledError: + pass + self._tick_task = None + # 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) + pending = list(self._executions) + if pending: + _, still_pending = await asyncio.wait( + pending, timeout=self.config.drain_grace_s + ) + for task in still_pending: + task.cancel() + if still_pending: + await asyncio.gather(*still_pending, return_exceptions=True) + # Drop drained tasks explicitly: done-callbacks flush on a + # later loop pass, but the report and release below must + # observe the drained set deterministically. + for task in pending: + self._executions.discard(task) + cancelled = len([t for t in pending if t.cancelled()]) + with self._lock: + report = DrainReport( + settled=self._settled, + abandoned=self._abandoned, + cancelled=cancelled, + ) + self.ownership.release() + self._started = False + return report + + async def poll_once(self, now: datetime | None = None) -> dict[str, str]: + """Run one guarded poll tick (started services only). + + Tick failures are recorded like background-tick failures and + re-raised to the direct caller. + """ + if not self._started: + raise SchedulerStartupError("scheduler service is not started") + instant = now if now is not None else self.clock() + try: + return await asyncio.to_thread(self._tick, instant) + except Exception as exc: + self._note_tick_error(exc) + raise + + def _note_tick_error(self, exc: BaseException) -> None: + self._last_tick_error = exc + if self.on_tick_error is not None: + try: + self.on_tick_error(exc) + except Exception: + pass + + # -- internals ----------------------------------------------------- + def _build_scheduler(self) -> None: + from wf_scheduling.poll import Scheduler + + self._scheduler = Scheduler( + schedule_store=self.schedule_store, + run_store=self.run_store, + sources=self._sources, + capacity=self.config.capacity, + preparer=self._preparer, + dispatcher=self._dispatcher, + ownership=self.ownership, + ) + + def _recover(self) -> list[str]: + from wf_scheduling import recovery as sched_recovery + + with self._lock: + return sched_recovery.recover( + schedule_store=self.schedule_store, + run_store=self.run_store, + now=self.clock(), + ownership=self.ownership, + ) + + def _join_tick(self) -> None: + with self._lock: + return None + + def _refresh_sources(self) -> None: + from wf_scheduling.poll import source_for_trigger + + try: + schedules = self.schedule_store.list_schedules(include_deleted=True) + except Exception: + schedules = [] + fresh: dict[str, Any] = {} + for sched in schedules: + if getattr(sched, "deleted", False): + continue + try: + fresh[sched.id] = source_for_trigger(sched.trigger) + except Exception: + # No source entry: the poll raises a loud per-schedule + # definition error instead of ticking a stale calendar. + continue + self._sources.clear() + self._sources.update(fresh) + + def _tick(self, now: datetime) -> dict[str, str]: + with self._lock: + self._refresh_sources() + return self._scheduler.poll(now) + + async def _run(self) -> None: + try: + while not self._stopping: + instant = self.clock() + try: + await asyncio.to_thread(self._tick, instant) + except asyncio.CancelledError: + raise + except Exception as exc: + self._note_tick_error(exc) + finally: + self._tick_count += 1 + await asyncio.sleep(self.config.poll_interval_s) + except asyncio.CancelledError: + raise + + def _submit(self, coro: Coroutine[Any, Any, None]) -> asyncio.Future[Any]: + """Register one execution task; dispatch always spawns exactly one.""" + if self._loop is None: # pragma: no cover - start() sets this first + raise SchedulerStartupError("scheduler service is not started") + # Submitted from the poll worker thread: wrap explicitly against + # the server loop, which has no current-loop binding out there. + tracked = asyncio.wrap_future( + asyncio.run_coroutine_threadsafe(coro, self._loop), loop=self._loop + ) + self._executions.add(tracked) + tracked.add_done_callback(self._executions.discard) + return tracked + + def _settle(self, run_id: str, state: Any) -> None: + with self._lock: + self._scheduler.record_stopped_execution(run_id, state, self.clock()) + self._settled += 1 + + def _abandon(self, run_id: str, error: BaseException) -> None: + """Fail a dispatcher-broken run closed through startup recovery. + + The run holds a durable executing mark with no genuine stopped + result — exactly the crashed-process shape — so the same recovery + that owns restart abandonment owns it here: no replay, no + fabricated outcome, slot freed. Recovery failures are recorded; + the run then waits for the next restart recovery. + """ + with self._lock: + self._errors.append(f"{run_id}: {error}") + del self._errors[: max(0, len(self._errors) - 20)] + try: + from wf_scheduling import recovery as sched_recovery + + sched_recovery.recover( + schedule_store=self.schedule_store, + run_store=self.run_store, + now=self.clock(), + ownership=self.ownership, + ) + except Exception as exc: + self._errors.append(f"{run_id}: recovery failed: {exc}") + del self._errors[: max(0, len(self._errors) - 20)] + return + self._abandoned += 1 diff --git a/src/wf_scheduling/ownership.py b/src/wf_scheduling/ownership.py index c95d6794..0cac7589 100644 --- a/src/wf_scheduling/ownership.py +++ b/src/wf_scheduling/ownership.py @@ -11,18 +11,17 @@ Locking design (Windows-tested, per the store transaction boundary): one new module holding ``msvcrt.locking`` (Windows) / ``fcntl.flock`` (POSIX) on ``/scheduler.lock``. No ad-hoc per-run lock files. -One store composition has exactly one lock identity: the composition root -whose distinct schedule and run store directories are its direct children -(see :func:`canonical_lock_root`). Any ancestor is NOT good enough — -different ancestors yield different lock files and therefore independent -OS locks, so ancestor containment can never prove exclusivity. Nor is a -mere common ancestor enough: cross pairs that reuse one protected store -with a different partner (shared schedule store, shared run store), -nested pairs, and same-directory dual use would map to different lock -files while covering the same store files, so they have no lock identity -at all. Guards require the held lock's frozen identity to equal the -composition identity; compositions that share no supported layout are -rejected instead of claimed safe. +One store composition has exactly one lock identity: the deepest +composition root that contains each store root as itself or a direct +child (see :func:`canonical_lock_root`). Any ancestor is NOT good enough +— different ancestors yield different lock files and therefore +independent OS locks, so ancestor containment can never prove +exclusivity. Nor is a mere common ancestor enough: cross pairs that +reuse one protected store with a different partner (shared schedule +store, shared run store) have no single lock file and no identity at +all. Guards require the held lock's frozen identity to equal the +composition identity; compositions with no single lock file are rejected +instead of claimed safe. The lock is acquired on the composition root that contains the protected schedule and run stores. Entry-point guards do not trust the caller to @@ -52,19 +51,19 @@ def describe_unsupported_layout(*store_roots: Path | str | None) -> str | None: """Explain why store roots form no supportable lock composition. Returns None when :func:`canonical_lock_root` yields an identity; - otherwise returns a message naming the supported layout (distinct - store directories under one composition root) so operators learn the - layout is rejected, not merely the lock. Guards raise this as - :class:`SecondOwnerError`: without one provable lock file, no owner - can claim exclusive ownership of the stores. + otherwise returns a message naming the supported layout (one + composition root containing every store as itself or a direct child) + so operators learn the layout is rejected, not merely the lock. + Guards raise this as :class:`SecondOwnerError`: without one provable + lock file, no owner can claim exclusive ownership of the stores. """ if canonical_lock_root(*store_roots) is not None: return None return ( - "scheduler stores are not a supported composition: the schedule " - "and run stores must be distinct directories under one composition " - "root; shared-store cross pairs, nested pairs, and same-directory " - "dual use cannot prove exclusive ownership and are rejected" + "scheduler stores are not a supported composition: one composition " + "root must contain every store as itself or a direct child; " + "shared-store cross pairs and split layouts cannot prove exclusive " + "ownership and are rejected" ) @@ -86,22 +85,25 @@ def canonical_store_path(path: Path | str) -> str: def canonical_lock_root(*store_roots: Path | str | None) -> str | None: """Return the one canonical lock identity for a store composition. - The supported layout is two or more distinct store directories that - are direct children of one composition root; the identity is that - parent and the lock file lives at ``/scheduler.lock``. Both - the held guard's frozen root and the guarded stores map through this - function, so exactly one lock file can authorize one composition, and - any two compositions covering the same store files map to the same - lock file (identical pairs) or are rejected (shared-store cross - pairs, which are not siblings). A single root maps to itself, so a + The identity is the deepest composition root that contains every + store root as itself or a direct child; the lock file lives at + ``/scheduler.lock``. Both the held guard's frozen root and + the guarded stores map through this function. Identical pairs share + the identity of their root (so the server layout, where the schedule + and run stores flank the same composition root, needs exactly one + lock file), distinct siblings share their parent, and a nested pair + shares the outer root — any two compositions covering the same store + files through identical, sibling, or nested roots map to the same + lock file and exclude each other. A single root maps to itself, so a lock still proves the directory it was acquired on. - Returns None when a root is missing, when the roots are not distinct - siblings under one parent (nested pairs, same-directory dual use, or - ancestor-less roots on different drives), or when no roots are given: - such compositions are unsupported and guards reject them instead of - claiming safety. Nesting one composition's stores inside another - composition's store subtree without sharing the identical root stays - operator error outside the supported layout. + Returns None when a root is missing, when no roots are given, or + when no single root contains every store root as itself or a direct + child (shared-store cross pairs, split layouts, ancestor-less roots + on different drives): such compositions have no single lock file + that could prove exclusivity, so guards reject them instead of + claiming safety. Pointing one composition's stores inside another + live composition's store subtree without sharing its identical roots + stays operator error outside the supported layout. """ canonical: list[str] = [] for store_root in store_roots: @@ -112,10 +114,20 @@ def canonical_lock_root(*store_roots: Path | str | None) -> str | None: return None if len(canonical) == 1: return canonical[0] - parents = {os.path.dirname(path) for path in canonical} - if len(parents) != 1 or len(set(canonical)) != len(canonical): - return None - return parents.pop() + candidates = {path: os.path.dirname(path) for path in canonical} + for path in canonical: + identity = path + if all( + other == identity or os.path.dirname(other) == identity + for other in canonical + ): + return identity + parent = candidates[path] + if all( + other == parent or os.path.dirname(other) == parent for other in canonical + ): + return parent + return None class SchedulerOwnership: @@ -164,7 +176,7 @@ class SchedulerOwnership: nothing, and a cross pair reusing one protected store with a different partner has no lock identity at all. An unrelated (even held) lock, a released lock, a mutated handle, and rootless - or non-sibling stores are never covered: guards reject those + or composition-less stores are never covered: guards reject those before any write or side effect instead of trusting the caller's promise. """ diff --git a/src/wf_scheduling/prepare.py b/src/wf_scheduling/prepare.py index 2496235c..7a45a86a 100644 --- a/src/wf_scheduling/prepare.py +++ b/src/wf_scheduling/prepare.py @@ -86,7 +86,13 @@ class SchedulePreparer: def prepare( self, *, sched: Any, intended: datetime, now: datetime ) -> PreparedInvocation | PreparationRejected: - environment = self._build_environment(sched) + try: + environment = self._build_environment(sched) + except KeyError: + # The deployment vanished between schedule creation and + # admission: a preflight rejection, never a tick failure, and + # no run is invented for the occurrence. + return PreparationRejected(reason="deployment-deleted") try: revision = self._deployments.deployment_revision(environment.deployment.id) required = list( diff --git a/tests/scheduling/test_lifecycle.py b/tests/scheduling/test_lifecycle.py new file mode 100644 index 00000000..f5a05f3a --- /dev/null +++ b/tests/scheduling/test_lifecycle.py @@ -0,0 +1,558 @@ +"""Opt-in server scheduler lifecycle (T12): service behavior over real stores. + +The service owns canonical ownership, startup recovery (never executes), +periodic polling that never blocks on long workflows, bounded execution +tasks behind the async-completion seam, and stop-admission-first drain +with truthful abandonment. Execution itself is scripted here through a +stub async runtime; real-workflow and real-server integration arrives in +dedicated integration tests. +""" + +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import Any, cast + +import pytest + +from tests.artifacts.test_run_store import artifact as _artifact +from tests.artifacts.test_run_store import deployment as _deployment +from wf_artifacts.runs.store import FileRunStore +from wf_artifacts.store import FileWorkflowArtifactStore +from wf_core import END, RunState, RunStatus +from wf_scheduling.lifecycle import ( + DrainReport, + SchedulerService, + SchedulerServiceConfig, + SchedulerStartupError, + StoreDeploymentDirectory, +) +from wf_scheduling.models import Schedule +from wf_scheduling.ownership import SchedulerOwnership, SecondOwnerError +from wf_scheduling.store import FileScheduleStore + + +def ts(y: int, mo: int, d: int, h: int = 0, mi: int = 0) -> datetime: + return datetime(y, mo, d, h, mi, tzinfo=UTC) + + +def _sched_model(sid: str, intended: datetime, **kw: Any) -> Schedule: + base: dict[str, Any] = { + "id": sid, + "deployment_id": "parent.personal", + "trigger": {"kind": "oneshot", "at": intended.isoformat()}, + "input_bindings": [], + "created_at": intended.isoformat(), + "updated_at": intended.isoformat(), + } + base.update(kw) + return Schedule.model_validate(base) + + +def _stopped(status: RunStatus) -> RunState: + return RunState( + workflow_name="sched", + status=status, + workflow_input={}, + state={}, + ) + + +class ScriptedRuntime: + """Stub async runtime with per-test outcomes (never touches stores).""" + + def __init__(self, outcome: Any = "complete") -> None: + self.outcome = outcome + self.calls: list[dict[str, Any]] = [] + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def run_workflow_from_plan( + self, + plan: Any, + workflow_input: dict[str, Any], + deployment: Any = None, + artifact: Any = None, + saved_subgraph_tree: Any = None, + limits: Any = None, + ) -> RunState: + self.calls.append( + { + "plan": plan, + "input": workflow_input, + "deployment": deployment, + "artifact": artifact, + "limits": limits, + } + ) + spec = self.outcome + if callable(spec): + spec = spec(len(self.calls)) + if spec == "hang": + self.started.set() + await self.release.wait() + return _stopped(RunStatus.COMPLETED) + if spec == "raise": + raise RuntimeError("injected execution failure") + assert isinstance(spec, str), f"unknown outcome spec {spec!r}" + return _stopped( + { + "complete": RunStatus.COMPLETED, + "fail": RunStatus.FAILED, + "interrupt": RunStatus.INTERRUPTED, + }[spec] + ) + + async def resume_workflow_from_plan(self, *args: Any, **kwargs: Any) -> RunState: + raise AssertionError("scheduler never resumes through the runtime") + + +def _artifact_store(tmp_path: Path) -> FileWorkflowArtifactStore: + store = FileWorkflowArtifactStore(tmp_path / "wf") + store.save_artifact(_artifact().model_copy(update={"plan": _plan_dict()})) + store.save_deployment(_deployment()) + return store + + +def _plan_dict() -> dict[str, Any]: + """Minimal executable plan (constant workflow, proven shape).""" + return { + "name": "parent", + "input_schema": {"type": "object", "properties": {}}, + "state_schema": { + "type": "object", + "properties": {"result": {"type": "string"}}, + }, + "output_schema": { + "type": "object", + "properties": {"result": {"type": "string"}}, + }, + "outcomes": ["ok"], + "start": "constant", + "nodes": [ + { + "id": "constant", + "type": "node", + "node": "wf.std.constant", + "input": [ + { + "value": "hello from scheduler", + "target": {"root": "local", "parts": ["value"]}, + } + ], + "output": [ + { + "source": {"root": "local", "parts": ["value"]}, + "target": {"root": "state", "parts": ["result"]}, + } + ], + } + ], + "edges": [{"from": "constant", "outcome": "ok", "to": END}], + "output": [ + { + "path": {"root": "state", "parts": ["result"]}, + "target": {"root": "local", "parts": ["result"]}, + } + ], + } + + +def _service( + tmp_path: Path, + runtime: Any, + config: SchedulerServiceConfig | None = None, +) -> SchedulerService: + if config is None: + # Deterministic tests drive ticks explicitly through poll_once; + # the background loop has its own test below. + config = SchedulerServiceConfig(poll_interval_s=0.01, auto_tick=False) + return SchedulerService( + schedule_store=FileScheduleStore(tmp_path), + run_store=FileRunStore(tmp_path), + runtime=runtime, + artifact_store=_artifact_store(tmp_path), + ownership=SchedulerOwnership(tmp_path, owner="test"), + config=config, + ) + + +async def _wait_for(cond: Any, timeout: float = 10.0) -> None: + async with asyncio.timeout(timeout): + while not cond(): + await asyncio.sleep(0.01) + + +def _entries(store: FileScheduleStore, sid: str, kind: str) -> list[dict[str, Any]]: + page = store.list_occurrences(sid, limit=100) + rows = cast(list[dict[str, Any]], page["occurrences"]) + return [r for r in rows if r["kind"] == kind] + + +def _only_run_id(run_store: Any) -> str: + runs = list(run_store.list_runs()) + assert len(runs) == 1 + return runs[0].id + + +async def test_scheduled_completion_through_lifecycle(tmp_path: Path) -> None: + intended = ts(2026, 9, 8, 12, 0) + service = _service(tmp_path, ScriptedRuntime("complete")) + try: + await service.start() + service.schedule_store.create_schedule(_sched_model("a", intended)) + result = await service.poll_once(intended + timedelta(seconds=1)) + run_id = _only_run_id(service.run_store) + assert result == {"a": f"admit:{run_id}"} + await _wait_for(lambda: service.live_executions == 0) + record = service.run_store.get_run(run_id) + assert record.status.value == "completed" + assert len(_entries(service.schedule_store, "a", "completed")) == 1 + # The dispatcher ran the pinned deployment input, not a fixture. + assert len(service.runtime.calls) == 1 + assert service.runtime.calls[0]["input"] == {} + assert service.runtime.calls[0]["deployment"].id == "parent.personal" + assert service.runtime.calls[0]["limits"].max_steps == 10_000 + finally: + report = await service.stop() + assert report.settled == 1 + assert report.cancelled == 0 + + +async def test_scheduled_failure_records_failed_history(tmp_path: Path) -> None: + intended = ts(2026, 9, 8, 12, 0) + service = _service(tmp_path, ScriptedRuntime("fail")) + try: + await service.start() + service.schedule_store.create_schedule(_sched_model("a", intended)) + result = await service.poll_once(intended + timedelta(seconds=1)) + run_id = _only_run_id(service.run_store) + assert result == {"a": f"admit:{run_id}"} + await _wait_for(lambda: service.live_executions == 0) + assert service.run_store.get_run(run_id).status.value == "failed" + assert len(_entries(service.schedule_store, "a", "failed")) == 1 + finally: + await service.stop() + + +async def test_scheduled_interrupt_stays_resumable(tmp_path: Path) -> None: + intended = ts(2026, 9, 8, 12, 0) + service = _service(tmp_path, ScriptedRuntime("interrupt")) + try: + await service.start() + service.schedule_store.create_schedule(_sched_model("a", intended)) + await service.poll_once(intended + timedelta(seconds=1)) + run_id = _only_run_id(service.run_store) + await _wait_for(lambda: service.live_executions == 0) + record = service.run_store.get_run(run_id) + assert record.status.value == "interrupted" + assert record.resume_readiness.value == "ready" + assert len(_entries(service.schedule_store, "a", "interrupted")) == 1 + finally: + await service.stop() + + +async def test_sources_refresh_after_start(tmp_path: Path) -> None: + intended = ts(2026, 9, 8, 12, 0) + service = _service(tmp_path, ScriptedRuntime("complete")) + try: + await service.start() + assert await service.poll_once(intended) == {} + service.schedule_store.create_schedule(_sched_model("a", intended)) + result = await service.poll_once(intended + timedelta(seconds=1)) + assert result["a"].startswith("admit:run-") + await _wait_for(lambda: service.live_executions == 0) + finally: + await service.stop() + + +async def test_hang_then_late_settlement_frees_capacity(tmp_path: Path) -> None: + intended = ts(2026, 9, 8, 12, 0) + runtime = ScriptedRuntime("hang") + service = _service( + tmp_path, + runtime, + SchedulerServiceConfig(poll_interval_s=0.01, capacity=1, auto_tick=False), + ) + try: + await service.start() + service.schedule_store.create_schedule(_sched_model("a", intended)) + service.schedule_store.create_schedule(_sched_model("b", intended)) + first = await service.poll_once(intended + timedelta(seconds=1)) + assert first["a"].startswith("admit:run-") + await _wait_for(lambda: runtime.started.is_set()) + run_a = _only_run_id(service.run_store) + assert service.run_store.get_run(run_a).status.value == "admitted" + assert service.live_executions == 1 + # Capacity is saturated: no second admission while the first hangs. + await service.poll_once(intended + timedelta(seconds=2)) + assert [r.id for r in list(service.run_store.list_runs())] == [run_a] + runtime.release.set() + await _wait_for(lambda: service.live_executions == 0) + assert service.run_store.get_run(run_a).status.value == "completed" + # The freed slot admits the waiting schedule on the next tick. + later = await service.poll_once(intended + timedelta(seconds=3)) + assert later["b"].startswith("admit:run-") + finally: + runtime.release.set() + await service.stop() + + +async def test_long_run_does_not_block_calendar_polling(tmp_path: Path) -> None: + intended = ts(2026, 9, 8, 12, 0) + calls = {"n": 0} + + def outcome(n: int) -> str: + calls["n"] = n + return "hang" if n == 1 else "complete" + + runtime = ScriptedRuntime(outcome) + service = _service( + tmp_path, + runtime, + SchedulerServiceConfig(poll_interval_s=0.01, capacity=2, auto_tick=False), + ) + try: + await service.start() + service.schedule_store.create_schedule(_sched_model("a", intended)) + await service.poll_once(intended + timedelta(seconds=1)) + await _wait_for(lambda: runtime.started.is_set()) + run_a = _only_run_id(service.run_store) + # A second schedule admitted while the first still executes. + service.schedule_store.create_schedule(_sched_model("b", intended)) + result = await service.poll_once(intended + timedelta(seconds=2)) + assert result["b"].startswith("admit:run-") + await _wait_for(lambda: service.live_executions == 1) + run_b = result["b"].split(":", 1)[1] + assert service.run_store.get_run(run_b).status.value == "completed" + assert service.run_store.get_run(run_a).status.value == "admitted" + assert service.live_executions <= 2 + finally: + runtime.release.set() + await service.stop() + + +async def test_graceful_shutdown_cancels_hanging_execution(tmp_path: Path) -> None: + intended = ts(2026, 9, 8, 12, 0) + runtime = ScriptedRuntime("hang") + service = _service( + tmp_path, + runtime, + SchedulerServiceConfig( + poll_interval_s=0.01, drain_grace_s=0.05, auto_tick=False + ), + ) + try: + await service.start() + service.schedule_store.create_schedule(_sched_model("a", intended)) + await service.poll_once(intended + timedelta(seconds=1)) + await _wait_for(lambda: runtime.started.is_set()) + run_id = _only_run_id(service.run_store) + report = await service.stop() + assert report.cancelled == 1 + assert report.settled == 0 + # Truthful handling: the run keeps its executing mark, settled by + # restart recovery instead of a fabricated outcome. + record = FileRunStore(tmp_path).get_run(run_id) + assert record.status.value == "admitted" + assert FileRunStore(tmp_path).is_executing(run_id) is True + assert service.live_executions == 0 + finally: + runtime.release.set() + await service.stop() + + +async def test_restart_after_shutdown_abandons_without_replay( + tmp_path: Path, +) -> None: + intended = ts(2026, 9, 8, 12, 0) + runtime = ScriptedRuntime("hang") + first = _service( + tmp_path, + runtime, + SchedulerServiceConfig( + poll_interval_s=0.01, drain_grace_s=0.05, auto_tick=False + ), + ) + try: + await first.start() + first.schedule_store.create_schedule(_sched_model("a", intended)) + await first.poll_once(intended + timedelta(seconds=1)) + await _wait_for(lambda: runtime.started.is_set()) + run_id = _only_run_id(first.run_store) + await first.stop() + finally: + runtime.release.set() + await first.stop() + assert len(runtime.calls) == 1 + # Fresh instances across the restart, like a real process boundary. + second = _service(tmp_path, ScriptedRuntime("complete")) + try: + await second.start() + record = second.run_store.get_run(run_id) + assert record.status.value == "failed" + assert record.resume_readiness.value == "not_applicable" + assert len(_entries(second.schedule_store, "a", "failed")) == 1 + assert len(second.runtime.calls) == 0 + finally: + await second.stop() + + +async def test_dispatcher_error_abandons_without_replay(tmp_path: Path) -> None: + intended = ts(2026, 9, 8, 12, 0) + service = _service(tmp_path, ScriptedRuntime("raise")) + try: + await service.start() + service.schedule_store.create_schedule(_sched_model("a", intended)) + await service.poll_once(intended + timedelta(seconds=1)) + run_id = _only_run_id(service.run_store) + await _wait_for(lambda: service.live_executions == 0) + record = service.run_store.get_run(run_id) + assert record.status.value == "failed" + assert record.resume_readiness.value == "not_applicable" + assert len(_entries(service.schedule_store, "a", "failed")) == 1 + assert any(run_id in message for message in service.errors) + # The slot is freed: a later schedule still admits. + service.schedule_store.create_schedule(_sched_model("b", intended)) + later = await service.poll_once(intended + timedelta(seconds=2)) + assert later["b"].startswith("admit:run-") + finally: + await service.stop() + + +async def test_failed_startup_releases_lock(tmp_path: Path) -> None: + squatter = SchedulerOwnership(tmp_path, owner="squatter").acquire() + try: + service = _service(tmp_path, ScriptedRuntime("complete")) + with pytest.raises(SchedulerStartupError): + await service.start() + assert service.running is False + finally: + squatter.release() + service = _service(tmp_path, ScriptedRuntime("complete")) + try: + await service.start() + assert service.running is True + finally: + await service.stop() + + +async def test_corrupt_store_startup_releases_lock(tmp_path: Path) -> None: + class BrokenRuns(FileRunStore): + def list_admissions(self) -> list[Any]: + raise OSError("injected store failure") + + service = SchedulerService( + schedule_store=FileScheduleStore(tmp_path), + run_store=BrokenRuns(tmp_path), + runtime=ScriptedRuntime("complete"), + artifact_store=_artifact_store(tmp_path), + ownership=SchedulerOwnership(tmp_path, owner="test"), + config=SchedulerServiceConfig(poll_interval_s=0.01, auto_tick=False), + ) + with pytest.raises(SchedulerStartupError): + await service.start() + # The lock is free for a retry after the operator fixes the store. + retry = SchedulerOwnership(tmp_path, owner="retry").acquire() + retry.release() + + +async def test_double_start_rejected_and_stop_idempotent( + tmp_path: Path, +) -> None: + assert await _service(tmp_path, ScriptedRuntime()).stop() == DrainReport() + service = _service(tmp_path, ScriptedRuntime("complete")) + try: + await service.start() + with pytest.raises(SchedulerStartupError): + await service.start() + finally: + await service.stop() + # Ownership is released only after the drain finishes. + freed = SchedulerOwnership(tmp_path, owner="freed").acquire() + freed.release() + + +async def test_second_owner_rejected_while_running(tmp_path: Path) -> None: + service = _service(tmp_path, ScriptedRuntime("complete")) + try: + await service.start() + with pytest.raises(SecondOwnerError): + SchedulerOwnership(tmp_path, owner="intruder").acquire() + finally: + await service.stop() + + +async def test_tick_error_recorded_and_loop_continues(tmp_path: Path) -> None: + intended = ts(2026, 9, 8, 12, 0) + service = _service(tmp_path, ScriptedRuntime("complete")) + try: + await service.start() + service.schedule_store.create_schedule(_sched_model("a", intended)) + service.schedule_store.create_schedule(_sched_model("broken", intended)) + # Corrupt one schedule file behind the model's back: the listing + # fails loudly instead of ticking a half-blind schedule set. + broken_path = tmp_path / "schedules" / "broken" / "schedule.json" + good_text = broken_path.read_text(encoding="utf-8") + broken_path.write_text('{"id": "broken", "trigger": {"kind": "bogus"}}') + with pytest.raises(Exception): + await service.poll_once(intended + timedelta(seconds=1)) + assert service.last_tick_error is not None + assert service.run_store.list_runs() == [] + # Repairing the file resumes normal ticking on the next call. + broken_path.write_text(good_text, encoding="utf-8") + result = await service.poll_once(intended + timedelta(seconds=2)) + assert result["a"].startswith("admit:run-") + await _wait_for(lambda: service.live_executions == 0) + finally: + await service.stop() + + +async def test_poll_once_before_start_rejected(tmp_path: Path) -> None: + service = _service(tmp_path, ScriptedRuntime("complete")) + with pytest.raises(SchedulerStartupError): + await service.poll_once(ts(2026, 9, 8, 12, 0)) + await service.stop() + + +async def test_background_loop_ticks_and_stops(tmp_path: Path) -> None: + from datetime import timezone + + service = _service( + tmp_path, + ScriptedRuntime("complete"), + SchedulerServiceConfig(poll_interval_s=0.01, auto_tick=True), + ) + try: + await service.start() + # A one-shot due in real time: the background loop admits it + # without any explicit poll_once call. + due = datetime.now(timezone.utc) + timedelta(seconds=0.2) + service.schedule_store.create_schedule(_sched_model("a", due)) + + def _status() -> str | None: + runs = list(service.run_store.list_runs()) + return runs[0].status.value if runs else None + + # Status moves monotonically admitted -> completed: waiting on it + # cannot miss a fast execution the way a live-task edge can. + await _wait_for(lambda: _status() == "completed") + run_id = _only_run_id(service.run_store) + assert service.run_store.get_run(run_id).status.value == "completed" + finally: + report = await service.stop() + assert report.settled == 1 + assert service.tick_count >= 2 + + +async def test_store_deployment_directory_contract(tmp_path: Path) -> None: + store = _artifact_store(tmp_path) + directory = StoreDeploymentDirectory(store) + assert directory.deployment_revision("parent.personal") >= 1 + assert directory.required_inputs("parent.personal") == [] + with pytest.raises(KeyError): + directory.deployment_revision("missing.deployment") + with pytest.raises(KeyError): + directory.required_inputs("missing.deployment") diff --git a/tests/scheduling/test_lock_identity.py b/tests/scheduling/test_lock_identity.py index 12c580b6..80302d2c 100644 --- a/tests/scheduling/test_lock_identity.py +++ b/tests/scheduling/test_lock_identity.py @@ -208,12 +208,12 @@ def test_canonical_windows_path_handling_retained(tmp_path: Path) -> None: def test_non_sibling_store_pairs_have_no_lock_identity(tmp_path: Path) -> None: - """Only distinct siblings under one composition root share a lock file. + """Cross pairs reusing one protected store have no single lock file. - Cross pairs that reuse one protected store (shared schedule store or - shared run store), nested pairs, and same-directory dual use would map - to different lock files while covering the same store files, so they - have no lock identity at all: guards must reject them outright. + A shared schedule store (or run store) with a different partner maps + to no lock identity: no held lock can authorize such a pair, so two + compositions can never gain independent authority over the same + store files through different layouts. """ overlap = tmp_path / "overlap" sched_store, run_store = _stores(overlap) @@ -222,11 +222,40 @@ def test_non_sibling_store_pairs_have_no_lock_identity(tmp_path: Path) -> None: assert canonical_lock_root(sched_store.root, other_runs.root) is None # Shared run store, different schedule store. assert canonical_lock_root(other_sched.root, run_store.root) is None - # Nested pairs: one store inside the other. - assert canonical_lock_root(overlap, sched_store.root) is None - assert canonical_lock_root(sched_store.root, overlap) is None - # Same directory serving as both stores. - assert canonical_lock_root(sched_store.root, sched_store.root) is None + # Split layouts across different parents. + assert canonical_lock_root(sched_store.root, tmp_path / "elsewhere") is None + + +def test_nested_and_shared_roots_unify_on_one_lock(tmp_path: Path) -> None: + """Identical and nested roots map to the same single lock file. + + The server layout points both stores at the composition root itself, + and a nested pair shares its outer root: every composition covering + the same store files through identical, sibling, or nested roots + contends on one lock file instead of holding independent locks. + """ + overlap = tmp_path / "overlap" + sched_store, run_store = _stores(overlap) + assert canonical_lock_root(overlap, overlap) == canonical_store_path(overlap) + assert canonical_lock_root( + sched_store.root, sched_store.root + ) == canonical_store_path(sched_store.root) + assert canonical_lock_root(overlap, sched_store.root) == canonical_store_path( + overlap + ) + assert canonical_lock_root(sched_store.root, overlap) == canonical_store_path( + overlap + ) + # One lock file: a second owner of the unified identity is rejected. + owner = SchedulerOwnership(overlap, owner="owner").acquire() + try: + assert owner.covers(overlap, overlap) is True + with pytest.raises(SecondOwnerError): + SchedulerOwnership(overlap, owner="second").acquire() + nested = SchedulerOwnership(sched_store.root, owner="nested") + assert nested.covers(sched_store.root, overlap) is False + finally: + owner.release() def test_shared_schedule_store_cannot_gain_independent_authority( diff --git a/tests/scheduling/test_preparation.py b/tests/scheduling/test_preparation.py index 2620d656..b5a712ad 100644 --- a/tests/scheduling/test_preparation.py +++ b/tests/scheduling/test_preparation.py @@ -223,3 +223,19 @@ def test_preparer_rejection_type_shape() -> None: rejected = PreparationRejected(reason="deployment-deleted") assert rejected.reason == "deployment-deleted" assert dataclasses.is_dataclass(PreparationRejected) + + +def test_vanished_deployment_in_environment_build_is_preflight() -> None: + def boom(sched: Any) -> Any: + raise KeyError("dep-1") + + preparer = SchedulePreparer( + DictDeployments({"dep-1": {"rev": 3, "required": []}}), + boom, + ) + intended = ts(2026, 9, 8, 12, 0) + result = preparer.prepare( + sched=_sched_model("gone"), intended=intended, now=intended + ) + assert isinstance(result, PreparationRejected) + assert result.reason == "deployment-deleted"