From 787a3c433f98708af17873559c802fb2e5a60df3 Mon Sep 17 00:00:00 2001 From: lda Date: Wed, 9 Sep 2026 19:13:17 +0700 Subject: [PATCH] fix: close scheduler persistence and capacity gaps --- src/wf_api/schedules.py | 17 ++ src/wf_scheduling/poll.py | 62 ++++-- src/wf_scheduling/recovery.py | 4 + src/wf_scheduling/store.py | 26 ++- src/wf_server/scheduling.py | 11 +- tests/scheduling/test_poll.py | 179 ++++++++++++++++++ tests/wf_server/test_scheduler_composition.py | 10 + 7 files changed, 296 insertions(+), 13 deletions(-) diff --git a/src/wf_api/schedules.py b/src/wf_api/schedules.py index 3be3a10e..3048a5ce 100644 --- a/src/wf_api/schedules.py +++ b/src/wf_api/schedules.py @@ -110,6 +110,23 @@ class WorkflowScheduleApi: raise KeyError("workflow schedule store is not configured") return store + def bind_schedule_store(self, store: Any) -> None: + """Bind the composition-owned schedule store exactly once. + + Server composition uses this seam when the scheduler is enabled on a + server that was built without the optional schedule surface. Refuse + to replace an already configured store: API and scheduler writes must + remain on one same-process lock and one durable root. + """ + if self._explicit_schedule_store is not None: + if self._explicit_schedule_store is not store: + raise ValueError("workflow schedule store is already configured") + return + context_store = getattr(self.context, "schedule_store", None) + if context_store is not None and context_store is not store: + raise ValueError("workflow schedule store is already configured") + self._explicit_schedule_store = store + def _artifact_store(self) -> Any: if self.context.artifact_store is None: raise KeyError("workflow artifact store is not configured") diff --git a/src/wf_scheduling/poll.py b/src/wf_scheduling/poll.py index 990f6772..e56d0b40 100644 --- a/src/wf_scheduling/poll.py +++ b/src/wf_scheduling/poll.py @@ -189,11 +189,44 @@ class Scheduler: return None return admission.schedule_id + def _fail_unattributed_active_runs(self, now: datetime) -> None: + """Fail corrupt active views without assigning blame to a sibling. + + A run view without an admission has no schedule identity and cannot + safely participate in overlap or capacity accounting. It is failed + closed in place, while healthy schedules continue their own poll. + """ + from wf_scheduling.recovery import ( + CORRUPT_VIEW_REASON, + _fail_run, + clear_executing, + clear_pending, + ) + + for run in self.run_store.list_runs(): + if self._status_value(run) not in ("admitted", "interrupted"): + continue + try: + self.run_store.get_admission(run.id) + except KeyError: + _fail_run(self.run_store, run, CORRUPT_VIEW_REASON, now, self.history) + clear_pending(self.run_store, run.id) + clear_executing(self.run_store, run.id) + def _task_load(self) -> int: from wf_scheduling.recovery import _is_pending, is_executing count = 0 for run in self.run_store.list_runs(): + # Manual API runs share the run store but do not consume the + # scheduler's bounded dispatch capacity. An admission with no + # schedule owner is the durable discriminator for that path. + try: + admission = self.run_store.get_admission(run.id) + except KeyError: + continue + if admission.schedule_id is None: + continue if self._status_value(run) != "admitted": # A mid-resume scheduled run is interrupted but holds the # durable executing mark: it occupies a live execution @@ -255,6 +288,17 @@ class Scheduler: """ for admission in self.run_store.list_admissions(): if admission.schedule_id == sched_id and admission.scheduled_at == intended: + # Admission is the occurrence authority. If a later write + # failed before the view/pending marker, rebuild only this + # run so the next pending sweep can dispatch it; never admit + # the same occurrence again. + try: + self.run_store.get_run(admission.id) + except KeyError: + from wf_api.run_lifecycle import materialize_admitted_view + + materialize_admitted_view(store=self.run_store, admission=admission) + self.run_store.mark_pending_dispatch(admission.id) return admission.id return None @@ -555,6 +599,7 @@ class Scheduler: # -- polling -------------------------------------------------------- def poll(self, now: datetime) -> dict[str, str]: self._require_ownership() + self._fail_unattributed_active_runs(now) self._dispatch_pending(now) schedules = self.schedule_store.list_schedules(include_deleted=True) ids = sorted(item.id for item in schedules) @@ -666,6 +711,11 @@ class Scheduler: self.schedule_store.save_consumed(sched.id, max(consumed, now)) return "disabled" if sched.exhausted: + # Exhaustion is durable terminal state. A torn one-shot + # transition may have written the history/flag before its stale + # candidate was cleared; retrying must finish that cleanup. + if self.schedule_store.get_candidate(sched.id) is not None: + self.schedule_store.save_candidate(None, schedule_id=sched.id) return "exhausted" if sched.paused: if self.schedule_store.get_candidate(sched.id) is not None: @@ -673,17 +723,6 @@ class Scheduler: consumed = self.schedule_store.get_consumed(sched.id) or EPOCH self.schedule_store.save_consumed(sched.id, max(consumed, now)) return "paused" - # Fail closed on corrupt active views before calendar work (F5). - for run in self.run_store.list_runs(): - if self._status_value(run) not in ("admitted", "interrupted"): - continue - try: - self.run_store.get_admission(run.id) - except KeyError: - reason = f"corrupt run view without admission: {run.id}" - sched.blocked_reason = reason - self.schedule_store.save_schedule(sched) - raise BlockedSchedule(reason) from None try: src = self.sources[sched.id] except KeyError as exc: @@ -838,6 +877,7 @@ class Scheduler: sched.exhausted = True self.schedule_store.save_schedule(sched) self.schedule_store.save_consumed(sched.id, instant) + self.schedule_store.save_candidate(None, schedule_id=sched.id) last_result = "exhausted" continue self._record( diff --git a/src/wf_scheduling/recovery.py b/src/wf_scheduling/recovery.py index 8a893e93..ef901faa 100644 --- a/src/wf_scheduling/recovery.py +++ b/src/wf_scheduling/recovery.py @@ -82,6 +82,10 @@ CORRUPT_PENDING_REASON = ( "corrupt pending marker: no admission owns this run; the marker cannot " "be trusted for dispatch" ) +CORRUPT_VIEW_REASON = ( + "corrupt run view: no admission owns this run; the view cannot be " + "trusted for overlap or dispatch" +) def recover( diff --git a/src/wf_scheduling/store.py b/src/wf_scheduling/store.py index d50ae267..4cf1ba57 100644 --- a/src/wf_scheduling/store.py +++ b/src/wf_scheduling/store.py @@ -28,6 +28,20 @@ from .models import OccurrenceRecord, PendingCandidate, Schedule, ensure_schedul UTC = timezone.utc +def _history_identity(record: OccurrenceRecord) -> tuple[object, ...]: + """Return fields that identify one durable occurrence transition.""" + return ( + record.schedule_id, + record.occurrence_id, + record.kind, + record.resolved_at, + record.run_id, + record.checkpoint_id, + record.interval_start, + record.interval_end, + ) + + class ScheduleExistsError(ValueError): """A schedule id is already taken (deleted ids are never reused).""" @@ -167,11 +181,21 @@ class FileScheduleStore: Stamps the entry ordinal (one past the highest effective ordinal so far) so pagination distinguishes multiple entries sharing one occurrence's ``(resolved_at, occurrence_id)`` — admission and each - stopped result each keep their own position. + stopped result each keep their own position. Repeating the same + occurrence transition is a no-op, which makes terminal bookkeeping + safe across a history-write/flag-write tear. """ with self._lock: path = self._history_path(record.schedule_id) entries = self._read_history_locked(record.schedule_id) + identity = _history_identity(record) + for item in entries: + try: + existing = OccurrenceRecord.model_validate(item) + except TypeError, ValueError: + continue + if _history_identity(existing) == identity: + return highest = -1 for index, item in enumerate(entries): ordinal = item.get("seq") diff --git a/src/wf_server/scheduling.py b/src/wf_server/scheduling.py index feb5703d..162a084d 100644 --- a/src/wf_server/scheduling.py +++ b/src/wf_server/scheduling.py @@ -39,8 +39,17 @@ def build_scheduler_service( ever applies to schedule-owned runs); manual runs keep the legacy path either way. """ + # The API and scheduler must share the same store instance. A + # process-local lock on two same-root FileScheduleStore instances would + # not serialize their read-modify-write history updates, so attach the + # scheduler store to a server built without ``schedules=True`` as well. + try: + schedule_store = server.api.schedules._schedule_store() + except KeyError: + schedule_store = FileScheduleStore(server.config.store_root) + server.api.schedules.bind_schedule_store(schedule_store) service = SchedulerService( - schedule_store=FileScheduleStore(server.config.store_root), + schedule_store=schedule_store, run_store=server.stores.run_store, runtime=server.context.runtime, artifact_store=server.stores.artifact_store, diff --git a/tests/scheduling/test_poll.py b/tests/scheduling/test_poll.py index 1dc72087..b15f8ce1 100644 --- a/tests/scheduling/test_poll.py +++ b/tests/scheduling/test_poll.py @@ -392,6 +392,185 @@ def test_manual_and_other_schedule_runs_are_overlap_independent( sched.ownership.release() +def test_manual_admission_does_not_consume_scheduler_capacity(tmp_path: Path) -> None: + """Manual runs bypass scheduler capacity; scheduled work still dispatches.""" + from wf_api.run_lifecycle import materialize_admitted_view, persist_admission + + sched, store, runs, sources = _harness(tmp_path, capacity=1, script={"*": "hang"}) + t0 = ts(2026, 9, 8, 12, 0) + manual = persist_admission( + store=runs, + run_id=runs.allocate_run_id(), + environment=fixture_environment(object()), + resolved_input={}, + max_steps=None, + ) + materialize_admitted_view(store=runs, admission=manual) + assert sched._task_load() == 0 + + _add(sched, store, sources, "a", OneShotSource(t0), t0 - timedelta(hours=1)) + assert sched.poll(t0)["a"].startswith("admit:run-") + sched.ownership.release() + + +def test_existing_admission_rebuilds_missing_view_and_pending_marker( + tmp_path: Path, +) -> None: + """A history tear after admission remains dispatchable on the next poll.""" + + class FailHistoryOnce(FileScheduleStore): + def __init__(self, root: Path) -> None: + super().__init__(root) + self.armed = True + + def append_history(self, record: Any) -> None: + if self.armed: + self.armed = False + raise OSError("injected history failure") + super().append_history(record) + + sched_store = FailHistoryOnce(tmp_path / "sched") + run_store = FileRunStore(tmp_path / "runs") + sources: dict[str, Any] = {} + t0 = ts(2026, 9, 8, 12, 0) + sched = Scheduler( + schedule_store=sched_store, + run_store=run_store, + sources=sources, + capacity=1, + preparer=SchedulePreparer( + DictDeployments({"dep-1": {"rev": 1, "required": []}}), + fixture_environment, + ), + dispatcher=ScriptedDispatcher({"*": "complete"}), + ownership=SchedulerOwnership(tmp_path, owner="test").acquire(), + ) + try: + _add( + sched, sched_store, sources, "a", OneShotSource(t0), t0 - timedelta(hours=1) + ) + try: + sched.poll(t0) + raise AssertionError("history failure must propagate") + except OSError as exc: + assert "injected history failure" in str(exc) + [admission] = run_store.list_admissions() + assert run_store.list_runs() == [] + assert not run_store.is_pending_dispatch(admission.id) + + result = sched.poll(t0 + timedelta(seconds=1)) + assert result["a"] == f"admit:{admission.id}" + assert run_store.get_run(admission.id).status.value == "admitted" + assert run_store.is_pending_dispatch(admission.id) + + sched.poll(t0 + timedelta(seconds=2)) + assert run_store.get_run(admission.id).status.value == "completed" + assert not run_store.is_pending_dispatch(admission.id) + finally: + sched.ownership.release() + + +def test_oneshot_terminal_bookkeeping_is_idempotent_after_torn_write( + tmp_path: Path, +) -> None: + """A torn exhausted transition does not duplicate history or retain a candidate.""" + + class FailScheduleOnce(FileScheduleStore): + def __init__(self, root: Path) -> None: + super().__init__(root) + self.armed = False + + def save_schedule(self, schedule: Any) -> None: + if self.armed: + self.armed = False + raise OSError("injected schedule failure") + super().save_schedule(schedule) + + sched_store = FailScheduleOnce(tmp_path / "sched") + run_store = FileRunStore(tmp_path / "runs") + sources: dict[str, Any] = {} + t0 = ts(2026, 9, 8, 12, 0) + sched = Scheduler( + schedule_store=sched_store, + run_store=run_store, + sources=sources, + capacity=1, + preparer=SchedulePreparer( + DictDeployments({"dep-1": {"rev": 1, "required": []}}), + fixture_environment, + ), + dispatcher=ScriptedDispatcher({"*": "complete"}), + ownership=SchedulerOwnership(tmp_path, owner="test").acquire(), + ) + try: + _add( + sched, sched_store, sources, "a", OneShotSource(t0), t0 - timedelta(hours=1) + ) + from wf_scheduling.models import PendingCandidate + + sched_store.save_candidate( + PendingCandidate(schedule_id="a", intended_at=t0, revision=1), + schedule_id="a", + ) + sched_store.armed = True + try: + sched.poll(t0 + timedelta(minutes=5)) + raise AssertionError("schedule failure must propagate") + except OSError as exc: + assert "injected schedule failure" in str(exc) + assert ( + len([r for r in _history(sched_store, "a") if r["kind"] == "exhausted"]) + == 1 + ) + assert sched_store.get_candidate("a") is not None + + assert sched.poll(t0 + timedelta(minutes=6))["a"] == "exhausted" + assert ( + len([r for r in _history(sched_store, "a") if r["kind"] == "exhausted"]) + == 1 + ) + assert sched_store.get_candidate("a") is None + finally: + sched.ownership.release() + + +def test_corrupt_unattributed_view_isolated_from_healthy_schedules( + tmp_path: Path, +) -> None: + """A view without an admission fails closed without blocking siblings.""" + from wf_artifacts import ResumeReadiness, WorkflowRunRecord + from wf_artifacts.runs.models import StoredRunStatus + + sched, store, runs, sources = _harness(tmp_path, script={"*": "complete"}) + t0 = ts(2026, 9, 8, 12, 0) + _add(sched, store, sources, "a", OneShotSource(t0), t0 - timedelta(hours=1)) + _add(sched, store, sources, "b", OneShotSource(t0), t0 - timedelta(hours=1)) + corrupt_id = runs.allocate_run_id() + runs.save_run( + WorkflowRunRecord( + id=corrupt_id, + status=StoredRunStatus.ADMITTED, + resume_readiness=ResumeReadiness.NOT_APPLICABLE, + environment=fixture_environment(object()), + latest_checkpoint_id=None, + created_at=t0, + updated_at=t0, + ) + ) + + results = sched.poll(t0) + + assert results["a"].startswith("admit:run-") + assert results["b"].startswith("admit:run-") + assert runs.get_run(corrupt_id).status.value == "failed" + assert all( + runs.get_run(run.id).status.value == "completed" + for run in runs.list_runs() + if run.id != corrupt_id + ) + sched.ownership.release() + + def test_record_resumed_stopped_result_is_idempotent_and_attributed( tmp_path: Path, ) -> None: diff --git a/tests/wf_server/test_scheduler_composition.py b/tests/wf_server/test_scheduler_composition.py index bc30ee01..cba157d7 100644 --- a/tests/wf_server/test_scheduler_composition.py +++ b/tests/wf_server/test_scheduler_composition.py @@ -97,6 +97,16 @@ def test_build_scheduler_service_wires_server_stores(tmp_path: Path) -> None: assert service.ownership.lock_path == server.config.store_root / "scheduler.lock" +def test_build_scheduler_service_shares_schedule_store_with_server_api( + tmp_path: Path, +) -> None: + """Enabled scheduler composition exposes the same store to API and poller.""" + server = build_local_static_workflow_server(tmp_path) + service = build_scheduler_service(server, SchedulerServiceConfig(auto_tick=False)) + + assert server.api.schedules._schedule_store() is service.schedule_store + + async def test_scheduler_service_start_stop_on_server_stores( tmp_path: Path, ) -> None: