feat: add persisted run step budget state
This commit is contained in:
@@ -0,0 +1,53 @@
|
||||
"""Run-wide step budget policy and admission.
|
||||
|
||||
One finite, persisted counter (``RunState.steps_executed``) covers every frame
|
||||
and subgraph scope in a run. Admission happens immediately before step dispatch:
|
||||
an admitted attempt increments the counter and stamps the selected frame with
|
||||
its one-based step number; a denied attempt raises without incrementing and
|
||||
without invoking any handler. This slice deliberately adds no per-step durable
|
||||
checkpoints; the counter is persisted inside the existing stopped-run envelope.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from wf_core.errors import WorkflowStepLimitExceeded
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from wf_core.run_state import ExecutionFrame, RunState
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RunLimits:
|
||||
"""Immutable step budget captured when a run is created."""
|
||||
|
||||
max_steps: int = 10_000
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if isinstance(self.max_steps, bool) or not isinstance(self.max_steps, int):
|
||||
raise TypeError("max_steps must be an integer")
|
||||
if self.max_steps < 1:
|
||||
raise ValueError("max_steps must be positive")
|
||||
|
||||
|
||||
def admit_step_attempt(run: RunState, frame: ExecutionFrame, node_id: str) -> int:
|
||||
"""Admit one step attempt for ``frame`` about to dispatch ``node_id``.
|
||||
|
||||
On success the run-wide counter is incremented, the frame remembers the
|
||||
assigned step number, and that number is returned. When the budget is
|
||||
already exhausted the counter is left untouched, the frame keeps its
|
||||
previous number, and ``WorkflowStepLimitExceeded`` is raised before any
|
||||
handler runs.
|
||||
"""
|
||||
if run.steps_executed >= run.limits.max_steps:
|
||||
raise WorkflowStepLimitExceeded.from_run(run, frame, node_id)
|
||||
run.steps_executed += 1
|
||||
frame.step_number = run.steps_executed
|
||||
return frame.step_number
|
||||
|
||||
|
||||
def remaining_step_attempts(run: RunState) -> int:
|
||||
"""Return the unspent budget, floored at zero (never negative)."""
|
||||
return max(run.limits.max_steps - run.steps_executed, 0)
|
||||
Reference in New Issue
Block a user