starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-09-09 00:04:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: retry a workflow-as-code task from its task options (#11013)

Browse Source 
* feat: retry a workflow-as-code task from its task options

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WvjgKMRNtRNPnAkg6MkiTA

* fix: claim every retry attempt key up front, so a step cannot alias one

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WvjgKMRNtRNPnAkg6MkiTA

* fix: bound retry attempts, which now claim their keys up front

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WvjgKMRNtRNPnAkg6MkiTA

* fix: honour an explicit zero retry multiplier in the python client

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WvjgKMRNtRNPnAkg6MkiTA

* docs: state the retry validation rules once in the task docstring

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WvjgKMRNtRNPnAkg6MkiTA

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Ruben Fiszel
2026-09-08 11:37:56 +02:00
committed by GitHub
co-authored by Claude Opus 5
parent 33f9828c3e
commit d3f305db98
15 changed files with 1126 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python-client/wmill/tests/test_workflow.py
+211
View File
@@ -1436,6 +1436,217 @@ class TestTaskOptions:
assert step_info["concurrency_time_window_s"] == 60
# =====================================================================
# TASK RETRY TESTS
# =====================================================================
# What the failed child job leaves in ``completed_steps``.
_FAILED = {"__wmill_error": True, "message": "boom", "result": {}}
class TestTaskRetry:
"""Nothing carries a retry across rounds: every round re-derives which
attempt comes next from the checkpoint alone."""
def test_each_failure_buys_a_backoff_sleep_and_one_more_attempt(self):
@task(retry={"attempts": 2, "delay": 30, "multiplier": 2})
async def call_api(x: int):
return x
@workflow
async def wf(x: int):
return await call_api(x=x)
completed = {"call_api": _FAILED}
r = _run_workflow(wf, {"completed_steps": dict(completed)}, {"x": 1})
assert r == {"type": "sleep", "key": "call_api#retry2", "seconds": 30}
completed["call_api#retry2"] = None
r = _run_workflow(wf, {"completed_steps": dict(completed)}, {"x": 1})
assert [s["key"] for s in r["steps"]] == ["call_api#2"]
# the delay grows by the multiplier for the second retry
completed["call_api#2"] = _FAILED
r = _run_workflow(wf, {"completed_steps": dict(completed)}, {"x": 1})
assert r == {"type": "sleep", "key": "call_api#retry3", "seconds": 60}
completed["call_api#retry3"] = None
r = _run_workflow(wf, {"completed_steps": dict(completed)}, {"x": 1})
assert [s["key"] for s in r["steps"]] == ["call_api#3"]
# attempts spent: the failure reaches the body
completed["call_api#3"] = _FAILED
with pytest.raises(TaskError):
_run_workflow(wf, {"completed_steps": dict(completed)}, {"x": 1})
def test_a_task_the_body_never_awaits_still_sleeps_and_retries(self):
# The runner dispatches unawaited task calls by flushing ``_pending``, so
# a backoff that only fired when awaited would drop the retry silently
# and report the workflow complete.
@task(retry={"attempts": 1, "delay": 30})
async def fire(x: int):
return x
@workflow
async def wf():
fire(x=1)
return "done"
r = _run_workflow(wf, {"completed_steps": {"fire": _FAILED}}, {})
assert r == {"type": "sleep", "key": "fire#retry2", "seconds": 30}
# ``step()`` names are arbitrary strings, so a step really can be called
# ``t#2``. Whichever of the two allocates second is the one renamed, and it
# has to be the same one in every round — hence claiming the attempt keys up
# front.
def test_an_inline_step_named_like_an_attempt_key_before_the_task_keeps_it(self):
@task(retry={"attempts": 1})
async def t(x: int):
return x
@workflow
async def wf():
decoy = await step("t#2", lambda: "not an attempt")
return [decoy, await t(x=1)]
r = _run_workflow(
wf, {"completed_steps": {"t#2": "not an attempt", "t": _FAILED}}, {}
)
assert [s["key"] for s in r["steps"]] == ["t#2_2"]
def test_an_inline_step_named_like_an_attempt_key_after_the_task_is_not_it(self):
@task(retry={"attempts": 1})
async def t(x: int):
return x
@workflow
async def wf():
pending = t(x=1)
decoy = await step("t#2", lambda: "not an attempt")
return [decoy, await pending]
# The first round records the step under the key left over after the
# task claimed ``t#2``, so the retry re-dispatches instead of reading
# the step's value.
r = _run_workflow(wf, {}, {})
assert r["key"] == "t#2_2"
r = _run_workflow(
wf, {"completed_steps": {"t": _FAILED, "t#2_2": "not an attempt"}}, {}
)
assert [s["key"] for s in r["steps"]] == ["t#2"]
def test_the_child_dispatched_for_an_attempt_walks_past_failure_and_backoff(self):
# The non-matching branch awaits a future that never resolves, so a child
# that walks the loop wrong parks the run until its timeout.
@task(retry={"attempts": 1, "delay": 30})
async def t(x: int):
return x * 10
@workflow
async def wf(x: int):
return await t(x=x)
r = _run_workflow(
wf,
{
"completed_steps": {"t": _FAILED, "t#retry2": None},
"_executing_key": "t#2",
},
{"x": 4},
)
assert r == {"type": "complete", "result": 40}
def test_a_misspelled_retry_option_raises_instead_of_being_ignored(self):
# The policy is a plain dict here, unlike the TS `TaskRetry` type, so
# nothing else would tell the author the option never took effect.
with pytest.raises(ValueError, match="max_delay_s"):
@task(retry={"attempts": 2, "max_delay_s": 300})
async def t(x: int):
return x
def test_an_out_of_range_attempt_count_is_rejected_where_it_is_written(self):
# Each attempt claims its keys before the first one is dispatched, so an
# unbounded count would hang the workflow allocating them.
for attempts in (10_000, -1, 2.5, float("inf")):
with pytest.raises(ValueError, match="whole number"):
@task(retry={"attempts": attempts})
async def t(x: int):
return x
def test_a_zero_multiplier_is_honoured_not_read_as_the_default_of_one(self):
@task(retry={"attempts": 2, "delay": 30, "multiplier": 0})
async def t(x: int):
return x
@workflow
async def wf(x: int):
return await t(x=x)
r = _run_workflow(wf, {"completed_steps": {"t": _FAILED}}, {"x": 1})
assert r == {"type": "sleep", "key": "t#retry2", "seconds": 30}
# 30 * 0 — the second retry goes out with no wait at all
r = _run_workflow(
wf,
{"completed_steps": {"t": _FAILED, "t#retry2": None, "t#2": _FAILED}},
{"x": 1},
)
assert [s["key"] for s in r["steps"]] == ["t#3"]
def test_max_delay_caps_the_backoff_a_multiplier_grows(self):
@task(retry={"attempts": 2, "delay": 60, "multiplier": 100, "max_delay": 300})
async def t(x: int):
return x
@workflow
async def wf(x: int):
return await t(x=x)
r = _run_workflow(
wf,
{"completed_steps": {"t": _FAILED, "t#retry2": None, "t#2": _FAILED}},
{"x": 1},
)
assert r == {"type": "sleep", "key": "t#retry3", "seconds": 300}
def test_retry_that_succeeds_resolves_and_later_steps_keep_their_keys(self):
# No delay, so the retry is dispatched without a sleep round in between.
@task(retry={"attempts": 1})
async def flaky(x: int):
return x
@workflow
async def wf(x: int):
return await double(x=await flaky(x=x))
r = _run_workflow(
wf, {"completed_steps": {"flaky": _FAILED, "flaky#2": 7}}, {"x": 1}
)
assert r["steps"][0]["key"] == "double"
assert r["steps"][0]["args"] == {"x": 7}
def test_retrying_one_call_does_not_move_the_keys_of_the_calls_beside_it(self):
@task(retry={"attempts": 1})
async def t(x: int):
return x
@workflow
async def wf():
return await asyncio.gather(t(x=1), t(x=2))
r = _run_workflow(wf, {}, {})
assert [s["key"] for s in r["steps"]] == ["t", "t_2"]
# the first call retries as ``t#2``; the second keeps the ``t_2`` it was
# dispatched under, rather than being read as the first call's retry
r = _run_workflow(wf, {"completed_steps": {"t": _FAILED, "t_2": 20}}, {})
assert [s["key"] for s in r["steps"]] == ["t#2"]
# =====================================================================
# SLEEP TESTS
# =====================================================================
python-client/wmill/wmill/client.py
+150 -21
View File
@@ -2874,6 +2874,59 @@ def _task_error_from_marker(marker: dict, fallback_message: str) -> TaskError:
)
# The worker deserializes a sleep into a ``u32`` of seconds and fails the whole
# job on anything wider, so a delay a multiplier has run away with has to be
# capped here rather than sent.
_MAX_SLEEP_SECONDS = 2**32 - 1
_RETRY_KEYS = ("attempts", "delay", "multiplier", "max_delay")
# Every attempt claims its keys before the first one is dispatched, so an
# unbounded ``attempts`` is a workflow that hangs allocating rather than a very
# patient one.
_MAX_RETRY_ATTEMPTS = 100
def _checked_retry(retry: Optional[dict]) -> Optional[dict]:
"""Reject a policy where it is written, rather than mid-run on a replay: the
policy is a plain dict, so a misspelled key would otherwise be dropped in
silence and the task would retry on a policy nobody wrote."""
if retry is None:
return None
unknown = sorted(k for k in retry if k not in _RETRY_KEYS)
if unknown:
raise ValueError(
f"unknown retry option(s): {', '.join(unknown)}. Expected any of: {', '.join(_RETRY_KEYS)}"
)
attempts = retry.get("attempts")
if isinstance(attempts, bool) or not isinstance(attempts, int) or not 0 <= attempts <= _MAX_RETRY_ATTEMPTS:
raise ValueError(
f"retry attempts must be a whole number between 0 and {_MAX_RETRY_ATTEMPTS}, got {attempts!r}"
)
return retry
def _retry_delay_seconds(retry: dict, attempt: int) -> int:
"""Seconds to wait before retry number ``attempt`` (0 is the first retry)."""
base = retry.get("delay") or 0
if base <= 0:
return 0
# `or 1` would read an explicit `multiplier: 0` — every retry after the
# first going out with no wait — as the default of 1.
multiplier = retry.get("multiplier")
if multiplier is None:
multiplier = 1
try:
grown = base * multiplier**attempt
except OverflowError:
# A float delay times an integer multiplier raised past ~1e308.
grown = _MAX_SLEEP_SECONDS
max_delay = retry.get("max_delay")
if max_delay is not None:
grown = min(grown, max_delay)
return max(0, int(min(grown, _MAX_SLEEP_SECONDS)))
_workflow_ctx: _contextvars.ContextVar["WorkflowCtx"] = _contextvars.ContextVar(
"_workflow_ctx"
)
@@ -2929,28 +2982,77 @@ class WorkflowCtx:
def _next_step(self, name: str, script: str, func=None, dispatch_type: str = "inline", _task_options: Optional[dict] = None, **kwargs):
"""Return an awaitable that either resolves from cache or suspends."""
key = self._alloc_key(name or script or "step")
step_name = name or script or "step"
retry = (_task_options or {}).get("retry") or {}
# Clamped as well as validated at decoration: a policy that reached here
# another way must not spin the key loop below.
max_retries = min(max(0, int(retry.get("attempts") or 0)), _MAX_RETRY_ATTEMPTS)
# Claimed up front, all of them, and named off the first attempt's key:
# one allocated later would shift the keys of the steps beside it, and a
# ``step()`` named ``t#2`` — names are arbitrary — could alias one.
# Whichever is allocated second is the one renamed, in every round alike.
base_key = self._alloc_key(step_name)
attempt_keys = [base_key]
backoff_keys = []
for i in range(max_retries):
backoff_keys.append(self._alloc_key(f"{base_key}#retry{i + 2}"))
attempt_keys.append(self._alloc_key(f"{base_key}#{i + 2}"))
# One pass per attempt. Every attempt the checkpoint already holds is
# decided here — a failed one either retries (moving to the next key) or
# is handed back to the body — so the loop always ends at the first
# attempt that has yet to run.
attempt = 0
while True:
key = attempt_keys[attempt]
if key in self._completed:
val = self._completed[key]
if isinstance(val, dict) and val.get("__wmill_error"):
if attempt < max_retries:
self._retry_backoff(backoff_keys[attempt], base_key, retry, attempt)
attempt += 1
continue
raise _task_error_from_marker(val, f"Task '{name}' failed")
return self._resolved(val)
if self._executing_key is not None:
if key == self._executing_key:
return self._execute_directly(func, **kwargs)
else:
return self._never_resolve()
print(f"\n--- WAC: {key} ---")
info = {"name": name or key, "script": script or key, "args": kwargs, "key": key, "dispatch_type": dispatch_type}
if _task_options:
for opt_key in ("timeout", "tag", "cache_ttl", "priority", "concurrent_limit", "concurrency_key", "concurrency_time_window_s"):
if opt_key in _task_options and _task_options[opt_key] is not None:
info[opt_key] = _task_options[opt_key]
self._pending.append(info)
return self._suspend()
def _retry_backoff(self, key: str, base_key: str, retry: dict, attempt: int) -> None:
"""Wait out the backoff between two attempts of a retried task, as a
durable sleep, and return once there is nothing to wait for — no delay
configured, or the sleep already in the checkpoint.
Raises where it stands rather than from a coroutine the caller has to
await: a task call the body never awaits is still dispatched (the runner
flushes ``_pending``), so a backoff that only fired when awaited would
drop the retry and let the round report the workflow complete."""
seconds = _retry_delay_seconds(retry, attempt)
if seconds < 1:
return
if key in self._completed:
val = self._completed[key]
if isinstance(val, dict) and val.get("__wmill_error"):
raise _task_error_from_marker(val, f"Task '{name}' failed")
return self._resolved(val)
return
# Child mode never raises: the parent dispatched this child only after
# its own round had slept, so the loop moves on to the attempt being
# executed.
if self._executing_key is not None:
if key == self._executing_key:
return self._execute_directly(func, **kwargs)
else:
return self._never_resolve()
print(f"\n--- WAC: {key} ---")
info = {"name": name or key, "script": script or key, "args": kwargs, "key": key, "dispatch_type": dispatch_type}
if _task_options:
for opt_key in ("timeout", "tag", "cache_ttl", "priority", "concurrent_limit", "concurrency_key", "concurrency_time_window_s"):
if opt_key in _task_options and _task_options[opt_key] is not None:
info[opt_key] = _task_options[opt_key]
self._pending.append(info)
return self._suspend()
return
print(f"\n--- WAC: sleep({key}, {seconds}s) before retrying {base_key} ---")
raise _StepSuspend({"mode": "sleep", "key": key, "seconds": seconds, "steps": []})
async def _resolved(self, value):
return value
@@ -3190,6 +3292,7 @@ def task(
concurrency_limit: Optional[int] = None,
concurrency_key: Optional[str] = None,
concurrency_time_window_s: Optional[int] = None,
retry: Optional[dict] = None,
):
"""Decorator that marks a function as a workflow task.
@@ -3204,6 +3307,22 @@ def task(
decoded back before the caller sees it: a ``datetime`` comes back as a
string, a tuple as a list.
``retry`` re-dispatches the task after a failure, inside ``@workflow`` only.
Every attempt is a step of its own (``call_api``, ``call_api#2``, ...) and
the wait between two of them is a durable sleep, so a retrying task holds no
worker while it backs off. Keys: ``attempts`` (retries after the first
failure, a whole number from 0 to 100), ``delay`` (seconds before the first
retry, sub-second delays dropped), ``multiplier`` (applied to the delay
after each attempt, 1 keeps it constant), ``max_delay`` (ceiling in
seconds). ``attempts`` is required, and an out-of-range or unknown key is
rejected where the policy is written.
A workflow sleeps once per round, so tasks backing off in the same fan-out
wait one after another rather than together: the delay before a fan-out
retries is the sum of every backoff pending in it, not the longest one, and
it grows with both the width of the fan-out and ``attempts``. Retries with
no ``delay`` all go out in a single round.
Usage::
@task
@@ -3211,6 +3330,9 @@ def task(
@task(path="f/external_script", timeout=600, tag="gpu")
async def run_external(x: int): ...
@task(retry={"attempts": 3, "delay": 30, "multiplier": 2})
async def call_api(payload: dict): ...
"""
from inspect import signature as _sig
@@ -3222,6 +3344,7 @@ def task(
"concurrent_limit": concurrency_limit,
"concurrency_key": concurrency_key,
"concurrency_time_window_s": concurrency_time_window_s,
"retry": _checked_retry(retry),
}
# Remove None values
_task_opts = {k: v for k, v in _task_opts.items() if v is not None} or None
@@ -3316,9 +3439,12 @@ def task_script(
concurrency_limit: Optional[int] = None,
concurrency_key: Optional[str] = None,
concurrency_time_window_s: Optional[int] = None,
retry: Optional[dict] = None,
):
"""Create a task that dispatches to a separate Windmill script.
``retry`` takes the same policy as :func:`task`.
Usage::
extract = task_script("f/data/extract", timeout=600)
@@ -3328,7 +3454,7 @@ def task_script(
data = await extract(url="https://...")
"""
name = path.rsplit("/", 1)[-1]
_opts = {k: v for k, v in {"timeout": timeout, "tag": tag, "cache_ttl": cache_ttl, "priority": priority, "concurrent_limit": concurrency_limit, "concurrency_key": concurrency_key, "concurrency_time_window_s": concurrency_time_window_s}.items() if v is not None} or None
_opts = {k: v for k, v in {"timeout": timeout, "tag": tag, "cache_ttl": cache_ttl, "priority": priority, "concurrent_limit": concurrency_limit, "concurrency_key": concurrency_key, "concurrency_time_window_s": concurrency_time_window_s, "retry": _checked_retry(retry)}.items() if v is not None} or None
def wrapper(**kwargs):
ctx = _workflow_ctx.get(None)
@@ -3352,9 +3478,12 @@ def task_flow(
concurrency_limit: Optional[int] = None,
concurrency_key: Optional[str] = None,
concurrency_time_window_s: Optional[int] = None,
retry: Optional[dict] = None,
):
"""Create a task that dispatches to a separate Windmill flow.
``retry`` takes the same policy as :func:`task`.
Usage::
pipeline = task_flow("f/etl/pipeline", priority=10)
@@ -3364,7 +3493,7 @@ def task_flow(
result = await pipeline(input=data)
"""
name = path.rsplit("/", 1)[-1]
_opts = {k: v for k, v in {"timeout": timeout, "tag": tag, "cache_ttl": cache_ttl, "priority": priority, "concurrent_limit": concurrency_limit, "concurrency_key": concurrency_key, "concurrency_time_window_s": concurrency_time_window_s}.items() if v is not None} or None
_opts = {k: v for k, v in {"timeout": timeout, "tag": tag, "cache_ttl": cache_ttl, "priority": priority, "concurrent_limit": concurrency_limit, "concurrency_key": concurrency_key, "concurrency_time_window_s": concurrency_time_window_s, "retry": _checked_retry(retry)}.items() if v is not None} or None
def wrapper(**kwargs):
ctx = _workflow_ctx.get(None)