mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-08-18 16:02:10 +00:00
feat: workflow-as-code (WAC) v2 (#8172)
* feat: workflow-as-code v2 with @task decorator API
Replace ctx.step("name", "script") API with @task decorators where
functions are called directly. Users no longer need to pass WorkflowCtx
or use string-based step names/script paths.
Python: @task decorator with contextvars-based implicit context
TypeScript: task() wrapper with module-level context variable
Parsers: detect @task function calls instead of ctx.step() calls
Worker: updated wrappers to set implicit context
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat: WAC v2 checkpoint/replay with _executing_key child dispatch
- Rust-side orchestration: parent dispatches child jobs, suspends, resumes on completion
- _executing_key in checkpoint tells child which step to execute directly
- task() throws StepSuspend(mode="step_complete") after executing target step
- result_processor handles child completion and updates parent checkpoint
- WacGraph.svelte for runtime execution visualization
- Sequential and parallel workflows tested end-to-end
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: WAC v2 bundle cache, globalThis ctx sharing, description optional
- Disable bun bundle caching for WAC v2 scripts (wrapper needs
windmill-client from node_modules, not available in bundle mode)
- Use Reflect.set/get(globalThis, "__wmill_wf_ctx") to share workflow
context across dual module instances (wrapper vs user script)
- Never-resolving thenable for non-matching steps in child job mode
prevents Promise.all race conditions
- Make description field optional in NewScript API (defaults to "")
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat: add step() primitive for inline checkpointed steps
step() executes a function inline (no child job) and persists the result
to the checkpoint. On replay, the cached value is returned — ensuring
deterministic behavior for non-deterministic operations like Date.now()
or Math.random().
- TypeScript: step(name, fn) — executes inline, throws StepSuspend with
mode "inline_checkpoint" to persist before continuing
- Rust: InlineCheckpoint variant in WacOutput, saves to checkpoint and
resets running=false for immediate re-pickup (no zombie wait)
- Shared step counter between task() and step() via _allocKey()
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat: add Python WAC v2 support with task(), step(), workflow()
- Python SDK: WorkflowCtx with _executing_key child mode, _alloc_key
shared counter, _run_inline_step for step(), _execute_directly and
_never_resolve for child mode, step() async function
- Python executor: WAC v2 detection, checkpoint.json writing, WAC
wrapper.py generation calling _run_workflow(), post-execution hook
into shared handle_wac_v2_output()
- Make handle_wac_v2_output pub so both bun and python executors share
the same dispatch/suspend/inline-checkpoint logic
- 17 Python tests covering dispatch, replay, parallel, conditional,
inline checkpoint, and child mode
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* chore: update sqlx prepared queries
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: WacGraph Tooltip→Popover, simplify wacToFlow parsers
- Fix type error: Tooltip doesn't accept text snippet, use Popover
- Extract shared helpers for task matching and block collection
- Replace linear tasks.find() with Map lookups
- Remove mutable module-level counter
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: Box::pin WAC v2 output handler to prevent stack overflow
handle_python_job's async state machine was too large when combined
with handle_wac_v2_output. Box::pin heap-allocates the future.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: merge WAC v1 and v2 task decorators to preserve backward compat
The v2 @task decorator was shadowing the v1 one, breaking WAC v1
scripts that rely on HTTP-based dispatch via /workflow_as_code/ API.
The merged decorator handles three modes:
- v2: inside @workflow context → checkpoint/replay dispatch
- v1: WM_JOB_ID set, no @workflow → HTTP API dispatch + wait_job
- standalone: no Windmill env → execute function body directly
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: skip no_main_func detection for WAC v2 scripts in TS and Python parsers
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: prevent empty/noop dispatch causing infinite requeue loop
- Validate steps.len() > 0 in WAC dispatch handler (issue 3)
- Replace noop StepSuspend throw with never-resolving promise so it
can't reach the backend as an empty dispatch (issue 4)
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: Python task wrapper now converts positional args to kwargs in v2 mode
Previously only **kwargs were passed to _next_step(), silently dropping
positional arguments. Extract shared _merge_args() helper used by both
v1 and v2 paths.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: replace unwrap() with proper error propagation in WAC arg serialization
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: add workspace_id filter to v2_job queries in WAC dispatch
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: prevent race condition in WAC child dispatch
Restructure dispatch to save checkpoint + suspend parent + seed child
checkpoints in a single transaction BEFORE pushing child jobs. This
ensures a fast child can't complete before the parent is suspended.
Also wrap InlineCheckpoint save + running reset in a transaction to
prevent corrupted state on crash.
Use ULID for pre-generated child job IDs (consistent with rest of API).
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: include step key and child job ID in WAC error propagation
Move step_key lookup before the success check so failed child errors
include which task failed, the child job ID, and the original error.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* docs: document WAC determinism contract and step dispatch semantics
- Document that workflow functions must be deterministic across replays
- Document that WacStepDispatch.script/args are metadata, not dispatch targets
- Add comments on counter-based key allocation
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: tighten WAC v2 detection to reduce false positives
Replace naive substring matching with line-aware checks that skip
comments and look for specific patterns:
- TS: import from "windmill-client" containing workflow/task
- Python: @workflow and @task decorators with wmill import
Extracted shared helpers in wac_executor.rs used by both executors.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: show failed steps in WacGraph when workflow completes with errors
When flowDone is true and a pending step isn't in completedSteps,
mark it as 'failed' instead of 'running'. The failed state CSS and
XCircle icon were already defined but never triggered.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: unsuspend and fail parent when WAC child push fails
Previously if a child push failed mid-batch, the parent remained
suspended with suspend = num_steps but fewer children, hanging until
the 14-day timeout. Now the push loop catches errors and unsuspends
the parent before returning the error.
Also adds source hash validation: if the script content changes between
replays, the job fails with a clear error instead of silently feeding
stale checkpoint data into wrong steps.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: clear suspend_until when unsuspending WAC parent
Set suspend_until = NULL alongside suspend = 0 in both the child
failure and all-children-complete paths, so the parent doesn't rely
on subtle pull query invariants to be re-picked-up.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* test: add exhaustive edge case tests for WAC v2 SDK
fix: make TS task wrapper non-async to fix unawaited task flush
The async wrapper caused microtask-based thenable auto-resolution that
fired .then() and threw StepSuspend before _flushPending() could capture
unawaited steps — making the flush mechanism completely broken. Now the
thenable is returned directly without async wrapping. Backward compatible
with v1 (all code paths still return awaitables).
Tests added (59 TS + 66 Python) covering: full sequential lifecycle,
step after parallel, parallel after parallel, conditional on step result,
empty/single-task workflows, 10+ steps, falsy value preservation, inline
steps, mixed step/task, unawaited flush, child mode with parallel,
key determinism, large parallel groups, and complex mixed patterns.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: atomic checkpoint updates to prevent parallel child race condition
Replace read-modify-write pattern in handle_wac_child_completion with
atomic SQL operations:
- completed_steps merged via jsonb_set(... || jsonb_build_object(...))
so concurrent children on different workers don't overwrite each other
- suspend counter decremented atomically with RETURNING to determine
"all done" condition (instead of checking completed_steps in memory)
- suspend_until cleared in the same atomic decrement statement
Before this fix, two parallel children completing simultaneously could
both load the same checkpoint, each add their step, and save — the
second write would overwrite the first, silently losing a child result
and leaving the parent suspended forever.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: cancel already-pushed children on partial WAC dispatch failure
When pushing child jobs sequentially, if pushing child N fails, children
1..N-1 are already running. Previously the error handler only unsuspended
the parent, leaving orphaned children that would complete and corrupt the
checkpoint state (decrementing suspend on an already-unsuspended parent,
potentially causing duplicate step execution on re-run).
Now on partial failure:
1. Cancel all already-pushed children (prevents them from completing
and corrupting checkpoint state)
2. Clear pending_steps from checkpoint (so parent doesn't think
children are outstanding on re-run)
3. Then unsuspend parent (so the error propagates)
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: skip WAC duration write and child check for non-WAC parents
The duration write to workflow_as_code_status was running for every
non-flow child with a parent (error handlers, success handlers,
run_script children), even though it was only intended for WAC jobs.
Add WHERE workflow_as_code_status IS NOT NULL to skip non-WAC parents
entirely. Piggyback RETURNING pending_steps.job_ids on the same query
so WAC v2 child completion needs zero extra DB round-trips on the
success path.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: seed child checkpoint in same transaction as push
The child checkpoint insert was happening before the child job was
pushed, violating the FK constraint on v2_job_status. Move it into
the push transaction so the job row exists and the child can't be
picked up before its checkpoint is ready.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: set running=false when WAC parent suspends for child dispatch
The parent job kept running=true after suspending, so workers wouldn't
pick it up when children completed and suspend reached 0. The parent
only advanced when the zombie job detector reset it (~90s). Now the
dispatch suspend sets running=false so the parent is immediately
eligible for pickup.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: WAC parent suspend/unsuspend lifecycle
Keep running=true when suspending the parent so the normal pull query
(WHERE running=false) never picks it up. Keep suspend_until non-null
when decrementing suspend to 0 so the suspended pull query
(WHERE suspend_until IS NOT NULL AND suspend<=0) picks it up.
Previously: setting running=false caused infinite restart loops because
the normal pull query has no suspend check and would immediately re-pick
the parent. Clearing suspend_until on the last child prevented the
suspended pull from ever seeing it, requiring the 90s zombie detector.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat: add approval primitive, flow child completion, timeline fixes for WAC v2
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat: add error propagation, task options, sleep, and parallel for WAC v2
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* test: fix python SDK tests to use name-based keys and add new test coverage
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: address WAC v2 review findings (sleep timing, error marker, atomicity)
- Fix sleep using suspend=1 instead of 0 to enforce actual delay
- Add approval/sleep resume injection to Python executor
- Fix TS SDK concurrency_limit mapping (was reading wrong property)
- Namespace error marker as __wmill_error to avoid user data collision
- Wrap child completion SQL in transaction for atomicity
- Decrement suspend even when step key is missing (prevents hang)
- Expand TASK_RE to handle export const, let, var, generics
- Validate step key uniqueness before dispatch
- Log warning on checkpoint deserialization failure
- Remove unimplemented delete_after_use from SDKs
- Add TaskError exception class to Python SDK with diagnostic context
- Fix extra positional args handling and add functools.wraps
- Improve getParamNames to handle typed/destructured params
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* sqlx
* sqlx
* test: add WAC v1 e2e integration tests for TS and Python
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: revert fake test versions in typescript-client
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* refactor: remove unused WacGraph component and strip wacToFlow to isWorkflowAsCode
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* refactor: extract shared approval/sleep resume logic into wac_executor
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Ruben Fiszel
GitHub
@@ -22,6 +22,12 @@ httpx = ">=0.24"
|
||||
requires = ["poetry>=1.0.2", "poetry-dynamic-versioning"]
|
||||
build-backend = "poetry.masonry.api"
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"httpx>=0.28.1",
|
||||
"pytest>=9.0.2",
|
||||
]
|
||||
|
||||
[tool.poetry-dynamic-versioning]
|
||||
enable = true
|
||||
vcs = "git"
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -2151,69 +2151,6 @@ def ducklake(name: str = "main") -> DucklakeClient:
|
||||
"""
|
||||
return _client.ducklake(name)
|
||||
|
||||
def task(*args, **kwargs):
|
||||
"""Decorator to mark a function as a workflow task.
|
||||
|
||||
When executed inside a Windmill job, the decorated function runs as a
|
||||
separate workflow step. Outside Windmill, it executes normally.
|
||||
|
||||
Args:
|
||||
tag: Optional worker tag for execution
|
||||
|
||||
Returns:
|
||||
Decorated function
|
||||
"""
|
||||
from inspect import signature
|
||||
|
||||
def f(func, tag: str | None = None):
|
||||
if (
|
||||
os.environ.get("WM_JOB_ID") is None
|
||||
or os.environ.get("MAIN_OVERRIDE") == func.__name__
|
||||
):
|
||||
|
||||
def inner(*args, **kwargs):
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return inner
|
||||
else:
|
||||
|
||||
def inner(*args, **kwargs):
|
||||
global _client
|
||||
if _client is None:
|
||||
_client = Windmill()
|
||||
w_id = os.environ.get("WM_WORKSPACE")
|
||||
job_id = os.environ.get("WM_JOB_ID")
|
||||
f_name = func.__name__
|
||||
json = kwargs
|
||||
params = list(signature(func).parameters)
|
||||
for i, arg in enumerate(args):
|
||||
if i < len(params):
|
||||
p = params[i]
|
||||
key = p
|
||||
if key not in kwargs:
|
||||
json[key] = arg
|
||||
|
||||
params = {}
|
||||
if tag is not None:
|
||||
params["tag"] = tag
|
||||
w_as_code_response = _client.post(
|
||||
f"/w/{w_id}/jobs/run/workflow_as_code/{job_id}/{f_name}",
|
||||
json={"args": json},
|
||||
params=params,
|
||||
)
|
||||
job_id = w_as_code_response.text
|
||||
print(f"Executing task {func.__name__} on job {job_id}")
|
||||
job_result = _client.wait_job(job_id)
|
||||
print(f"Task {func.__name__} ({job_id}) completed")
|
||||
return job_result
|
||||
|
||||
return inner
|
||||
|
||||
if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
|
||||
return f(args[0], None)
|
||||
else:
|
||||
return lambda x: f(x, kwargs.get("tag"))
|
||||
|
||||
def parse_resource_syntax(s: str) -> Optional[str]:
|
||||
"""Parse resource syntax from string."""
|
||||
if s is None:
|
||||
@@ -2413,7 +2350,495 @@ def parse_sql_client_name(name: str) -> tuple[str, Optional[str]]:
|
||||
name = name
|
||||
schema = None
|
||||
if ":" in name:
|
||||
name, schema = name.split(":", 1)
|
||||
name, schema = name.split(":", 1)
|
||||
if not name:
|
||||
name = "main"
|
||||
return name, schema
|
||||
|
||||
|
||||
# ── Workflow-as-Code SDK ──────────────────────────────────────────────
|
||||
|
||||
import asyncio as _asyncio
|
||||
import contextvars as _contextvars
|
||||
|
||||
|
||||
class _StepSuspend(BaseException):
|
||||
"""Raised to suspend workflow execution. Inherits from BaseException
|
||||
so it is not caught by bare `except Exception:` blocks."""
|
||||
|
||||
def __init__(self, dispatch_info: dict):
|
||||
self.dispatch_info = dispatch_info
|
||||
|
||||
|
||||
class TaskError(Exception):
|
||||
"""Raised when a WAC task step failed.
|
||||
|
||||
Attributes:
|
||||
step_key: The checkpoint key of the failed step.
|
||||
child_job_id: The UUID of the failed child job.
|
||||
result: The error result from the child job.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str, *, step_key: str = "", child_job_id: str = "", result=None):
|
||||
super().__init__(message)
|
||||
self.step_key = step_key
|
||||
self.child_job_id = child_job_id
|
||||
self.result = result
|
||||
|
||||
|
||||
_workflow_ctx: _contextvars.ContextVar["WorkflowCtx"] = _contextvars.ContextVar(
|
||||
"_workflow_ctx"
|
||||
)
|
||||
|
||||
|
||||
class WorkflowCtx:
|
||||
"""Internal context for workflow replay/suspension.
|
||||
|
||||
Not user-facing — set implicitly by ``@workflow`` via contextvars.
|
||||
"""
|
||||
|
||||
def __init__(self, checkpoint: dict | None = None):
|
||||
checkpoint = checkpoint or {}
|
||||
self._completed: dict = checkpoint.get("completed_steps", {})
|
||||
self._counters: dict[str, int] = {}
|
||||
self._pending: list = []
|
||||
self._executing_key: str | None = checkpoint.get("_executing_key")
|
||||
|
||||
def _alloc_key(self, name: str = "step") -> str:
|
||||
"""Name-based key: ``double`` for first call, ``double_2``, ``double_3`` for subsequent."""
|
||||
n = self._counters.get(name, 0) + 1
|
||||
self._counters[name] = n
|
||||
return name if n == 1 else f"{name}_{n}"
|
||||
|
||||
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")
|
||||
|
||||
if key in self._completed:
|
||||
val = self._completed[key]
|
||||
if isinstance(val, dict) and val.get("__wmill_error"):
|
||||
raise TaskError(
|
||||
val.get("message", f"Task '{name}' failed"),
|
||||
step_key=val.get("step_key", ""),
|
||||
child_job_id=val.get("child_job_id", ""),
|
||||
result=val.get("result"),
|
||||
)
|
||||
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()
|
||||
|
||||
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()
|
||||
|
||||
async def _resolved(self, value):
|
||||
return value
|
||||
|
||||
async def _execute_directly(self, func, **kwargs):
|
||||
result = func(**kwargs)
|
||||
if _asyncio.iscoroutine(result):
|
||||
result = await result
|
||||
raise _StepSuspend({"mode": "step_complete", "steps": [], "result": result})
|
||||
|
||||
async def _never_resolve(self):
|
||||
await _asyncio.Future()
|
||||
|
||||
async def _suspend(self):
|
||||
steps = list(self._pending)
|
||||
self._pending.clear()
|
||||
raise _StepSuspend(
|
||||
{
|
||||
"mode": "parallel" if len(steps) > 1 else "sequential",
|
||||
"steps": steps,
|
||||
}
|
||||
)
|
||||
|
||||
async def _wait_for_approval(
|
||||
self, timeout: int = 1800, form: dict | None = None
|
||||
):
|
||||
key = self._alloc_key("approval")
|
||||
|
||||
if key in self._completed:
|
||||
return self._completed[key]
|
||||
|
||||
if self._executing_key is not None:
|
||||
await _asyncio.Future()
|
||||
|
||||
raise _StepSuspend({
|
||||
"mode": "approval",
|
||||
"key": key,
|
||||
"timeout": timeout,
|
||||
"form": form,
|
||||
"steps": [],
|
||||
})
|
||||
|
||||
async def _sleep(self, seconds: int):
|
||||
key = self._alloc_key("sleep")
|
||||
|
||||
if key in self._completed:
|
||||
return
|
||||
|
||||
if self._executing_key is not None:
|
||||
await _asyncio.Future()
|
||||
|
||||
raise _StepSuspend({
|
||||
"mode": "sleep",
|
||||
"key": key,
|
||||
"seconds": max(1, int(seconds)),
|
||||
"steps": [],
|
||||
})
|
||||
|
||||
async def _run_inline_step(self, name: str, fn):
|
||||
key = self._alloc_key(name or "step")
|
||||
|
||||
if key in self._completed:
|
||||
val = self._completed[key]
|
||||
if isinstance(val, dict) and val.get("__wmill_error"):
|
||||
raise TaskError(
|
||||
val.get("message", f"Step '{name}' failed"),
|
||||
step_key=val.get("step_key", ""),
|
||||
child_job_id=val.get("child_job_id", ""),
|
||||
result=val.get("result"),
|
||||
)
|
||||
return val
|
||||
|
||||
if self._executing_key is not None:
|
||||
await _asyncio.Future()
|
||||
|
||||
result = fn()
|
||||
if _asyncio.iscoroutine(result):
|
||||
result = await result
|
||||
|
||||
raise _StepSuspend({
|
||||
"mode": "inline_checkpoint",
|
||||
"steps": [],
|
||||
"key": key,
|
||||
"result": result,
|
||||
})
|
||||
|
||||
|
||||
def task(
|
||||
_func=None,
|
||||
*,
|
||||
path: Optional[str] = None,
|
||||
tag: Optional[str] = None,
|
||||
timeout: Optional[int] = None,
|
||||
cache_ttl: Optional[int] = None,
|
||||
priority: Optional[int] = None,
|
||||
concurrency_limit: Optional[int] = None,
|
||||
concurrency_key: Optional[str] = None,
|
||||
concurrency_time_window_s: Optional[int] = None,
|
||||
):
|
||||
"""Decorator that marks a function as a workflow task.
|
||||
|
||||
Works in both WAC v1 (sync, HTTP-based dispatch) and WAC v2
|
||||
(async, checkpoint/replay) modes:
|
||||
|
||||
- **v2 (inside @workflow)**: dispatches as a checkpoint step.
|
||||
- **v1 (WM_JOB_ID set, no @workflow)**: dispatches via HTTP API.
|
||||
- **Standalone**: executes the function body directly.
|
||||
|
||||
Usage::
|
||||
|
||||
@task
|
||||
async def extract_data(url: str): ...
|
||||
|
||||
@task(path="f/external_script", timeout=600, tag="gpu")
|
||||
async def run_external(x: int): ...
|
||||
"""
|
||||
from inspect import signature as _sig
|
||||
|
||||
_task_opts = {
|
||||
"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,
|
||||
}
|
||||
# Remove None values
|
||||
_task_opts = {k: v for k, v in _task_opts.items() if v is not None} or None
|
||||
|
||||
def decorator(func):
|
||||
task_path = path
|
||||
task_name = func.__name__
|
||||
|
||||
_params_list = list(_sig(func).parameters)
|
||||
|
||||
def _merge_args(args, kwargs):
|
||||
merged = dict(kwargs)
|
||||
for i, arg in enumerate(args):
|
||||
if i < len(_params_list):
|
||||
key = _params_list[i]
|
||||
if key not in merged:
|
||||
merged[key] = arg
|
||||
else:
|
||||
merged[f"arg{i}"] = arg
|
||||
return merged
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
# WAC v2: inside a @workflow context
|
||||
ctx = _workflow_ctx.get(None)
|
||||
if ctx is not None:
|
||||
script = task_path if task_path else task_name
|
||||
merged = _merge_args(args, kwargs)
|
||||
return ctx._next_step(task_name, script, func, _task_options=_task_opts, **merged)
|
||||
|
||||
# WAC v1: running inside a Windmill job but not in a @workflow
|
||||
if (
|
||||
os.environ.get("WM_JOB_ID") is not None
|
||||
and os.environ.get("MAIN_OVERRIDE") != func.__name__
|
||||
):
|
||||
global _client
|
||||
if _client is None:
|
||||
_client = Windmill()
|
||||
w_id = os.environ.get("WM_WORKSPACE")
|
||||
job_id = os.environ.get("WM_JOB_ID")
|
||||
json_args = _merge_args(args, kwargs)
|
||||
api_params = {}
|
||||
if tag is not None:
|
||||
api_params["tag"] = tag
|
||||
resp = _client.post(
|
||||
f"/w/{w_id}/jobs/run/workflow_as_code/{job_id}/{func.__name__}",
|
||||
json={"args": json_args},
|
||||
params=api_params,
|
||||
)
|
||||
child_job_id = resp.text
|
||||
print(f"Executing task {func.__name__} on job {child_job_id}")
|
||||
job_result = _client.wait_job(child_job_id)
|
||||
print(f"Task {func.__name__} ({child_job_id}) completed")
|
||||
return job_result
|
||||
|
||||
# Standalone — execute directly
|
||||
return func(*args, **kwargs)
|
||||
|
||||
wrapper._is_task = True
|
||||
wrapper._task_path = task_path
|
||||
return wrapper
|
||||
|
||||
if _func is not None:
|
||||
# @task without parentheses
|
||||
return decorator(_func)
|
||||
# @task() or @task(path="...", tag="...")
|
||||
return decorator
|
||||
|
||||
|
||||
def task_script(
|
||||
path: str,
|
||||
*,
|
||||
timeout: Optional[int] = None,
|
||||
tag: Optional[str] = None,
|
||||
cache_ttl: Optional[int] = None,
|
||||
priority: Optional[int] = None,
|
||||
concurrency_limit: Optional[int] = None,
|
||||
concurrency_key: Optional[str] = None,
|
||||
concurrency_time_window_s: Optional[int] = None,
|
||||
):
|
||||
"""Create a task that dispatches to a separate Windmill script.
|
||||
|
||||
Usage::
|
||||
|
||||
extract = task_script("f/data/extract", timeout=600)
|
||||
|
||||
@workflow
|
||||
async def main():
|
||||
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
|
||||
|
||||
def wrapper(**kwargs):
|
||||
ctx = _workflow_ctx.get(None)
|
||||
if ctx is not None:
|
||||
return ctx._next_step(name, path, dispatch_type="script", _task_options=_opts, **kwargs)
|
||||
raise RuntimeError(f'task_script("{path}") can only be called inside a @workflow')
|
||||
|
||||
wrapper.__name__ = name
|
||||
wrapper._is_task = True
|
||||
wrapper._task_path = path
|
||||
return wrapper
|
||||
|
||||
|
||||
def task_flow(
|
||||
path: str,
|
||||
*,
|
||||
timeout: Optional[int] = None,
|
||||
tag: Optional[str] = None,
|
||||
cache_ttl: Optional[int] = None,
|
||||
priority: Optional[int] = None,
|
||||
concurrency_limit: Optional[int] = None,
|
||||
concurrency_key: Optional[str] = None,
|
||||
concurrency_time_window_s: Optional[int] = None,
|
||||
):
|
||||
"""Create a task that dispatches to a separate Windmill flow.
|
||||
|
||||
Usage::
|
||||
|
||||
pipeline = task_flow("f/etl/pipeline", priority=10)
|
||||
|
||||
@workflow
|
||||
async def main():
|
||||
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
|
||||
|
||||
def wrapper(**kwargs):
|
||||
ctx = _workflow_ctx.get(None)
|
||||
if ctx is not None:
|
||||
return ctx._next_step(name, path, dispatch_type="flow", _task_options=_opts, **kwargs)
|
||||
raise RuntimeError(f'task_flow("{path}") can only be called inside a @workflow')
|
||||
|
||||
wrapper.__name__ = name
|
||||
wrapper._is_task = True
|
||||
wrapper._task_path = path
|
||||
return wrapper
|
||||
|
||||
|
||||
def workflow(func):
|
||||
"""Decorator marking an async function as a workflow-as-code entry point.
|
||||
|
||||
The function must be **deterministic**: given the same inputs it must call
|
||||
tasks in the same order on every replay. Branching on task results is fine
|
||||
(results are replayed from checkpoint), but branching on external state
|
||||
(current time, random values, external API calls) must use ``step()`` to
|
||||
checkpoint the value so replays see the same result.
|
||||
"""
|
||||
func._is_workflow = True
|
||||
return func
|
||||
|
||||
|
||||
async def step(name: str, fn):
|
||||
"""Execute ``fn`` inline and checkpoint the result.
|
||||
|
||||
On replay the cached value is returned without re-executing ``fn``.
|
||||
Use for lightweight deterministic operations (timestamps, random IDs,
|
||||
config reads) that should not incur the overhead of a child job.
|
||||
"""
|
||||
ctx: WorkflowCtx | None = _workflow_ctx.get(None)
|
||||
if ctx is not None:
|
||||
return await ctx._run_inline_step(name, fn)
|
||||
result = fn()
|
||||
if _asyncio.iscoroutine(result):
|
||||
result = await result
|
||||
return result
|
||||
|
||||
|
||||
async def sleep(seconds: int):
|
||||
"""Server-side sleep — suspend the workflow for the given duration without holding a worker.
|
||||
|
||||
Inside a @workflow, the parent job suspends and auto-resumes after ``seconds``.
|
||||
Outside a workflow, falls back to ``asyncio.sleep``.
|
||||
"""
|
||||
ctx: WorkflowCtx | None = _workflow_ctx.get(None)
|
||||
if ctx is not None:
|
||||
return await ctx._sleep(seconds)
|
||||
await _asyncio.sleep(seconds)
|
||||
|
||||
|
||||
async def wait_for_approval(
|
||||
timeout: int = 1800,
|
||||
form: dict | None = None,
|
||||
) -> dict:
|
||||
"""Suspend the workflow and wait for an external approval.
|
||||
|
||||
Use ``get_resume_urls()`` (wrapped in ``step()``) to obtain
|
||||
resume/cancel/approval URLs before calling this function.
|
||||
|
||||
Returns a dict with ``value`` (form data), ``approver``, and ``approved``.
|
||||
|
||||
Example::
|
||||
|
||||
urls = await step("urls", lambda: get_resume_urls())
|
||||
await step("notify", lambda: send_email(urls["approvalPage"]))
|
||||
result = await wait_for_approval(timeout=3600)
|
||||
"""
|
||||
ctx: WorkflowCtx | None = _workflow_ctx.get(None)
|
||||
if ctx is not None:
|
||||
return await ctx._wait_for_approval(timeout=timeout, form=form)
|
||||
raise RuntimeError("wait_for_approval can only be called inside a @workflow")
|
||||
|
||||
|
||||
async def parallel(items, fn, *, concurrency: Optional[int] = None):
|
||||
"""Process items in parallel with optional concurrency control.
|
||||
|
||||
Each item is processed by calling ``fn(item)``, which should be a @task.
|
||||
Items are dispatched in batches of ``concurrency`` (default: all at once).
|
||||
|
||||
Example::
|
||||
|
||||
@task
|
||||
async def process(item: str):
|
||||
...
|
||||
|
||||
results = await parallel(items, process, concurrency=5)
|
||||
"""
|
||||
if not items:
|
||||
return []
|
||||
batch_size = concurrency if concurrency and concurrency > 0 else len(items)
|
||||
results = []
|
||||
for i in range(0, len(items), batch_size):
|
||||
batch = items[i : i + batch_size]
|
||||
batch_results = await _asyncio.gather(*(fn(item) for item in batch))
|
||||
results.extend(batch_results)
|
||||
return results
|
||||
|
||||
|
||||
async def _run_workflow_async(func, checkpoint: dict, input_args: dict):
|
||||
ctx = WorkflowCtx(checkpoint)
|
||||
token = _workflow_ctx.set(ctx)
|
||||
try:
|
||||
result = await func(**input_args)
|
||||
# Flush any unawaited tasks (e.g. forgotten await on last statement)
|
||||
if ctx._pending:
|
||||
steps = list(ctx._pending)
|
||||
ctx._pending.clear()
|
||||
return {
|
||||
"type": "dispatch",
|
||||
"mode": "parallel" if len(steps) > 1 else "sequential",
|
||||
"steps": steps,
|
||||
}
|
||||
return {"type": "complete", "result": result}
|
||||
except _StepSuspend as e:
|
||||
info = e.dispatch_info
|
||||
mode = info.get("mode")
|
||||
if mode == "step_complete":
|
||||
return {"type": "complete", "result": info.get("result")}
|
||||
if mode == "inline_checkpoint":
|
||||
return {
|
||||
"type": "inline_checkpoint",
|
||||
"key": info["key"],
|
||||
"result": info.get("result"),
|
||||
}
|
||||
if mode == "approval":
|
||||
return {
|
||||
"type": "approval",
|
||||
"key": info["key"],
|
||||
"timeout": info.get("timeout"),
|
||||
"form": info.get("form"),
|
||||
}
|
||||
if mode == "sleep":
|
||||
return {
|
||||
"type": "sleep",
|
||||
"key": info["key"],
|
||||
"seconds": info.get("seconds"),
|
||||
}
|
||||
return {"type": "dispatch", **info}
|
||||
finally:
|
||||
_workflow_ctx.reset(token)
|
||||
|
||||
|
||||
def _run_workflow(func, checkpoint: dict, input_args: dict):
|
||||
"""Synchronous wrapper that runs the workflow coroutine to completion
|
||||
or until it suspends."""
|
||||
return _asyncio.run(_run_workflow_async(func, checkpoint, input_args))
|
||||
|
||||
Reference in New Issue
Block a user