feat: enforce step budget during sync dispatch
This commit is contained in:
@@ -42,7 +42,9 @@ def _inject_v1_budget_defaults(state: dict[str, Any]) -> dict[str, Any]:
|
||||
|
||||
Version-1 envelopes predate step budgets, so they receive the default
|
||||
limit, a zeroed counter, and an unassigned number per frame exactly once
|
||||
at load time. Attempts made before the upgrade are outside the new budget.
|
||||
at load time. Pre-budget trace entries and any outstanding interrupt keep
|
||||
an unassigned (``None``) number: attempts made before the upgrade are
|
||||
outside the new budget.
|
||||
"""
|
||||
upgraded = deepcopy(state)
|
||||
upgraded.setdefault("limits", {"max_steps": RunLimits().max_steps})
|
||||
@@ -52,6 +54,14 @@ def _inject_v1_budget_defaults(state: dict[str, Any]) -> dict[str, Any]:
|
||||
for frame in frames.values():
|
||||
if isinstance(frame, dict):
|
||||
frame.setdefault("step_number", None)
|
||||
trace = upgraded.get("trace")
|
||||
if isinstance(trace, list):
|
||||
for entry in trace:
|
||||
if isinstance(entry, dict):
|
||||
entry.setdefault("step_number", None)
|
||||
interrupt = upgraded.get("interrupt")
|
||||
if isinstance(interrupt, dict):
|
||||
interrupt.setdefault("step_number", None)
|
||||
return upgraded
|
||||
|
||||
|
||||
@@ -59,7 +69,10 @@ def _require_v2_budget_fields(state: dict[str, Any]) -> None:
|
||||
"""Reject v2 payloads missing budget fields as corrupt state.
|
||||
|
||||
Unlike v1, a v2 envelope promises budget fields; a missing counter is
|
||||
corruption, not another request for defaults.
|
||||
corruption, not another request for defaults. Trace and interrupt entries
|
||||
always carry the key (``None`` only for upgraded pre-budget history), so a
|
||||
missing key is likewise corrupt even though the dataclass default would
|
||||
otherwise mask it.
|
||||
"""
|
||||
if "limits" not in state or "steps_executed" not in state:
|
||||
raise ValueError("invalid persisted workflow run state: missing step budget")
|
||||
@@ -72,6 +85,21 @@ def _require_v2_budget_fields(state: dict[str, Any]) -> None:
|
||||
"invalid persisted workflow run state: "
|
||||
f"frame {frame_id!r} is missing its step number"
|
||||
)
|
||||
trace = state.get("trace")
|
||||
if isinstance(trace, list):
|
||||
for position, entry in enumerate(trace):
|
||||
if not isinstance(entry, dict) or "step_number" not in entry:
|
||||
raise ValueError(
|
||||
"invalid persisted workflow run state: "
|
||||
f"trace entry {position!r} is missing its step number"
|
||||
)
|
||||
interrupt = state.get("interrupt")
|
||||
if interrupt is not None:
|
||||
if not isinstance(interrupt, dict) or "step_number" not in interrupt:
|
||||
raise ValueError(
|
||||
"invalid persisted workflow run state: "
|
||||
"interrupt is missing its step number"
|
||||
)
|
||||
|
||||
|
||||
def _restore_root_alias(run: RunState) -> RunState:
|
||||
|
||||
@@ -134,6 +134,12 @@ class TraceEntry:
|
||||
next_node_id: str
|
||||
output: dict[str, Any] = field(default_factory=dict)
|
||||
state_changes: dict[str, Any] = field(default_factory=dict)
|
||||
# One-based step number assigned at admission. Every trace emitted for an
|
||||
# admitted step carries its frame's number; ``None`` only survives on
|
||||
# pre-budget (v1) entries. Gaps are valid when an attempt fails or
|
||||
# interrupts before emitting a trace, so ``RunState.steps_executed`` stays
|
||||
# authoritative for enforcement and resume.
|
||||
step_number: int | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -179,6 +185,12 @@ class InterruptRequest:
|
||||
request_schema: dict[str, object] = field(default_factory=_object_schema)
|
||||
resume_schema: dict[str, object] = field(default_factory=_object_schema)
|
||||
typed: bool = False
|
||||
# Step number of the admitted interrupt activation. The initial interrupt
|
||||
# trace and the later resume-completion trace both reuse this one number
|
||||
# because they describe a single activation; resume never admits again.
|
||||
# ``None`` only survives on pre-budget (v1) checkpoints whose activation
|
||||
# predates the counter.
|
||||
step_number: int | None = None
|
||||
|
||||
|
||||
def _default_run_limits() -> RunLimits:
|
||||
|
||||
@@ -22,6 +22,12 @@ from wf_core.runtime.scheduler import (
|
||||
)
|
||||
from wf_core.tokens import END
|
||||
|
||||
# Sentinel for ``append_trace()``: copy the named frame's admitted step number
|
||||
# (failing closed when unassigned). Interrupt resume passes its stored
|
||||
# activation number explicitly instead, so one activation keeps one number
|
||||
# across its interrupt and resume-completion entries without a second admission.
|
||||
_FROM_FRAME: Any = object()
|
||||
|
||||
|
||||
def append_trace(
|
||||
run: RunState,
|
||||
@@ -34,7 +40,28 @@ def append_trace(
|
||||
next_node_id: str,
|
||||
output: dict[str, Any],
|
||||
state_changes: dict[str, Any],
|
||||
step_number: int | None | Any = _FROM_FRAME,
|
||||
) -> None:
|
||||
"""Append one trace entry carrying its admitted step number.
|
||||
|
||||
By default the number is copied from the named frame, which must have been
|
||||
assigned by ``admit_step_attempt()`` during this dispatch; otherwise the
|
||||
trace would silently describe an uncounted step, so fail closed with
|
||||
``WorkflowExecutionError``. Pass ``step_number`` explicitly only to reuse a
|
||||
persisted activation number (interrupt resume-completion).
|
||||
"""
|
||||
if step_number is _FROM_FRAME:
|
||||
frame = run.frames.get(frame_id)
|
||||
if frame is None:
|
||||
raise WorkflowExecutionError(
|
||||
f"cannot trace step for unknown frame {frame_id!r}"
|
||||
)
|
||||
if frame.step_number is None:
|
||||
raise WorkflowExecutionError(
|
||||
f"cannot trace unadmitted step for frame {frame_id!r} "
|
||||
f"at node {node_id!r}; admit the step before dispatch"
|
||||
)
|
||||
step_number = frame.step_number
|
||||
run.trace.append(
|
||||
TraceEntry(
|
||||
frame_id=frame_id,
|
||||
@@ -45,6 +72,7 @@ def append_trace(
|
||||
next_node_id=next_node_id,
|
||||
output=output,
|
||||
state_changes=state_changes,
|
||||
step_number=step_number,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -57,6 +85,7 @@ def append_step_result_trace(
|
||||
step_type: str,
|
||||
next_node_id: str,
|
||||
result: StepExecutionResult,
|
||||
step_number: int | None | Any = _FROM_FRAME,
|
||||
) -> None:
|
||||
append_trace(
|
||||
run,
|
||||
@@ -68,6 +97,7 @@ def append_step_result_trace(
|
||||
next_node_id=next_node_id,
|
||||
output=result.output,
|
||||
state_changes=result.state_changes,
|
||||
step_number=step_number,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -73,6 +73,11 @@ def handle_interrupt_step(
|
||||
public_node_id=public_frame.node_id,
|
||||
route=route,
|
||||
)
|
||||
# The current dispatch was already admitted by step_workflow(), so the
|
||||
# frame carries this activation's number. Persist it on the request: resume
|
||||
# completes the same activation without admitting again, and its
|
||||
# completion trace reuses this stored number.
|
||||
interrupt_request.step_number = frame.step_number
|
||||
run.interrupt = interrupt_request
|
||||
run.status = RunStatus.INTERRUPTED
|
||||
frame.status = FrameStatus.INTERRUPTED
|
||||
|
||||
@@ -69,6 +69,15 @@ def resume_interrupt(
|
||||
resume_outcome: str,
|
||||
reducers: Mapping[str, ReducerDefinition] | None = None,
|
||||
) -> None:
|
||||
"""Complete a previously admitted interrupt activation.
|
||||
|
||||
Resume never admits a new step attempt: supplying the external payload
|
||||
finishes the activation counted at interrupt time. The completion trace
|
||||
therefore reuses the stored ``InterruptRequest.step_number`` (which may be
|
||||
``None`` only for pre-budget legacy activations) instead of the current
|
||||
frame number, and execution after resume continues from the persisted
|
||||
cumulative counter.
|
||||
"""
|
||||
if run.interrupt is None:
|
||||
raise WorkflowExecutionError("run is interrupted but has no interrupt request")
|
||||
|
||||
@@ -119,6 +128,9 @@ def resume_interrupt(
|
||||
# scope, a concurrent one buffers in the item lineage for barrier merge.
|
||||
state_changes = commit_foreach_aware_patch(run, frame, patch)
|
||||
next_node_id = index.next_node_id(frame.node_id, resume_outcome)
|
||||
# Reuse the activation's stored number (not the frame's current number, and
|
||||
# without admitting): both entries describe one admitted activation.
|
||||
activation_number = run.interrupt.step_number
|
||||
append_step_result_trace(
|
||||
run,
|
||||
frame_id=frame.id,
|
||||
@@ -131,6 +143,7 @@ def resume_interrupt(
|
||||
output=resume_payload,
|
||||
state_changes=state_changes,
|
||||
),
|
||||
step_number=activation_number,
|
||||
)
|
||||
run.interrupt = None
|
||||
advance_frame(run, frame, outcome=resume_outcome, next_node_id=next_node_id)
|
||||
|
||||
@@ -16,6 +16,7 @@ from wf_core.models.steps import (
|
||||
from wf_core.models.workflow import Workflow
|
||||
from wf_core.run_state import ExecutionFrame, FrameStatus, RunState, StepExecutionResult
|
||||
from wf_core.runtime.foreach_state import item_frame_owner, load_foreach_activation
|
||||
from wf_core.runtime.limits import admit_step_attempt
|
||||
from wf_core.runtime.ops.flow import advance_frame, append_step_result_trace
|
||||
from wf_core.runtime.ops.foreach import step_foreach
|
||||
from wf_core.runtime.ops.handlers import (
|
||||
@@ -129,6 +130,11 @@ def step_workflow(
|
||||
return run
|
||||
index, step = prepared
|
||||
frame = run.current_frame()
|
||||
# One admission per selected step, immediately before dispatch: the counter
|
||||
# increments before any handler runs, so failures and interrupts consume
|
||||
# their attempt, while a denied dispatch raises before its handler runs.
|
||||
# prepare_step() returning None (legacy END, interrupted) admits nothing.
|
||||
admit_step_attempt(run, frame, frame.node_id)
|
||||
|
||||
if isinstance(step, NodeUse):
|
||||
node_def = index.node_defs[step.node]
|
||||
@@ -241,6 +247,9 @@ async def step_workflow_async(
|
||||
return run
|
||||
index, step = prepared
|
||||
frame = run.current_frame()
|
||||
# Same single-admission rule as the sync path; only the concurrent foreach
|
||||
# batch below (Task 3) reserves differently.
|
||||
admit_step_attempt(run, frame, frame.node_id)
|
||||
|
||||
if isinstance(step, NodeUse):
|
||||
node_def = index.node_defs[step.node]
|
||||
@@ -312,8 +321,15 @@ async def _step_async_foreach_item_batch(
|
||||
Only handler awaits run concurrently. Finalization, tracing, and frame
|
||||
advancement happen afterward in ready-queue order so `RunState` is mutated
|
||||
deterministically.
|
||||
|
||||
Task 3 will bound the claimed siblings by the remaining budget and pin the
|
||||
reservation/failure semantics. Until then every frame in the batch is
|
||||
admitted in ready-queue order before any handler starts, so each trace has
|
||||
a number; a denied frame raises before any handler in the batch runs.
|
||||
"""
|
||||
frames = [first_frame, *_claim_matching_async_item_frames(run, index, first_frame)]
|
||||
for frame in frames:
|
||||
admit_step_attempt(run, frame, frame.node_id)
|
||||
tasks = []
|
||||
for frame in frames:
|
||||
node = _node_use_for_frame(index, frame)
|
||||
|
||||
Reference in New Issue
Block a user