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)
|
||||
@@ -15,6 +15,7 @@ from wf_core.run_state import (
|
||||
RunStatus,
|
||||
RuntimeScope,
|
||||
)
|
||||
from wf_core.runtime.limits import RunLimits
|
||||
from wf_core.runtime.scheduler import add_frame
|
||||
|
||||
|
||||
@@ -30,13 +31,19 @@ def initial_state(
|
||||
return state
|
||||
|
||||
|
||||
def create_run_state(workflow: Workflow, workflow_input: dict[str, object]) -> RunState:
|
||||
def create_run_state(
|
||||
workflow: Workflow,
|
||||
workflow_input: dict[str, object],
|
||||
*,
|
||||
limits: RunLimits | None = None,
|
||||
) -> RunState:
|
||||
state = initial_state(workflow, workflow_input)
|
||||
run = RunState(
|
||||
workflow_name=workflow.name,
|
||||
status=RunStatus.PENDING,
|
||||
workflow_input=dict(workflow_input),
|
||||
state=state,
|
||||
limits=limits if limits is not None else RunLimits(),
|
||||
scopes={
|
||||
ROOT_SCOPE_ID: RuntimeScope(
|
||||
id=ROOT_SCOPE_ID,
|
||||
|
||||
Reference in New Issue
Block a user