diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/__init__.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_callback.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_callback.py new file mode 100644 index 00000000..6fe72b2a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_callback.py @@ -0,0 +1,69 @@ +"""DAG conformance 10-11: task that is a wait-for-callback (flat). + +pre(step->"ready") -> cb(callback[dep pre]) -> post(step[dep cb] -> +cb + "_done"). The callback's native WaitForCallback op is checkpointed directly +under the Dag container (flat, name-based). The submitter receives the generated +callback id and does nothing durable (same as the 7-1 wait_for_callback +handler); the conformance runner completes the callback externally with a +success payload, which ``cb`` resolves to (a string). max_concurrency=1 for a +deterministic topological order. + +This scenario suspends until the external callback arrives. + +Returns the canonical summary defined by test-requirements/dag/10-11.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import ( + DurableContext, + WaitForCallbackContext, +) +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +def _normalize(value: str) -> str: + """Strip a single pair of surrounding double-quote characters if present. + + The default callback deserializer returns the raw payload text, which in + some SDKs includes the surrounding quote characters. The runner's payload is + alphanumeric, so stripping one surrounding pair is unambiguous. + """ + if len(value) >= 2 and value[0] == '"' and value[-1] == '"': + return value[1:-1] + return value + + +def _submitter(_deps, _callback_id: str, _ctx: WaitForCallbackContext) -> None: + """Receives the generated callback id; does nothing durable.""" + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + pre = d.step(lambda deps, sc: "ready", name="pre") + cb = d.wait_for_callback(_submitter, deps=[pre], name="cb") + d.step(lambda deps, sc: _normalize(deps[cb]) + "_done", deps=[cb], name="post") + + result = context.dag( + register, name="callbackdag", config=DagConfig(max_concurrency=1) + ) + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "cb": _normalize(result.results["cb"].result), + "post": result.results["post"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_child.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_child.py new file mode 100644 index 00000000..a21a2670 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_child.py @@ -0,0 +1,49 @@ +"""DAG conformance 10-5: task that is a runInChildContext (flat child container). + +seed(step->1) -> group(runInChildContext[dep seed]: inner-a->2, inner-b->3, +returns 5) -> done(step[dep group]->group*2=10). The child's native +RunInChildContext op is checkpointed directly under the Dag container (flat, +name-based). max_concurrency=1 for a deterministic topological order. Returns +the canonical summary defined by test-requirements/dag/10-5.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + seed = d.step(lambda deps, sc: 1, name="seed") + + def group_body(deps, child: DurableContext) -> int: + seed_val = deps[seed] + a = child.step(lambda _sc: seed_val + 1, name="inner-a") + b = child.step(lambda _sc: seed_val + 2, name="inner-b") + return a + b + + group = d.run_in_child_context(group_body, deps=[seed], name="group") + d.step(lambda deps, sc: deps[group] * 2, deps=[group], name="done") + + result = context.dag(register, name="childdag", config=DagConfig(max_concurrency=1)) + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "group": result.results["group"].result, + "done": result.results["done"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_compensate.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_compensate.py new file mode 100644 index 00000000..13e0e1b5 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_compensate.py @@ -0,0 +1,91 @@ +"""DAG conformance 10-18: compensation dependency read on a FAILED upstream is +ABSENT, not present (the deps-nullability contract). + +A DAG "compensate" with two step tasks: charge -> audit. + +- ``charge`` is a root step that ALWAYS fails. Its retry strategy is disabled + (max_attempts=1) so it ends FAILED deterministically after a single attempt + (exactly one StepFailed). +- ``audit`` depends on ``charge`` via an INLINE (typed) dependency and uses the + ALL_DONE trigger rule, so it runs even though ``charge`` FAILED and receives + ``charge`` in its resolved deps map. Its body reads its dependency's result + for ``charge``: a dependency that did not SUCCEED resolves to ``None`` + (absent), never a stale/fabricated value. ``audit`` returns ``"absent"`` when + it observes ``None`` and ``"present"`` otherwise. + +The DAG drains to COMPLETED_WITH_FAILURES without throwing: ``charge`` FAILED, +``audit`` SUCCEEDED with result ``"absent"``. Returns the canonical summary +defined by test-requirements/dag/10-18.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.config import ( + Duration, + JitterStrategy, + StepConfig, +) +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import ( + DagConfig, + DagContext, + DagResult, + TriggerRule, +) +from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python.retries import ( + RetryStrategyConfig, + create_retry_strategy, +) + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +def _charge_fails(_deps: Any, _sc: Any) -> Any: + raise RuntimeError("charge failed") + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + # Single attempt (no retry) so charge ends FAILED deterministically. + no_retry = create_retry_strategy( + RetryStrategyConfig( + max_attempts=1, + initial_delay=Duration.from_seconds(1), + backoff_rate=1, + jitter_strategy=JitterStrategy.NONE, + ) + ) + + def register(d: DagContext) -> None: + charge = d.step( + _charge_fails, + name="charge", + config=StepConfig(retry_strategy=no_retry), + ) + # Inline dep on charge + ALL_DONE: audit runs and reads charge. A failed + # dependency's value is absent (None), so audit returns "absent". + d.step( + lambda deps, sc: "absent" if deps.get(charge) is None else "present", + deps=[charge], + name="audit", + ).trigger_rule(TriggerRule.ALL_DONE) + + result = context.dag( + register, name="compensate", config=DagConfig(max_concurrency=1) + ) + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "audit": result.results["audit"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_compensation.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_compensation.py new file mode 100644 index 00000000..a4e60978 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_compensation.py @@ -0,0 +1,56 @@ +"""DAG conformance 10-2: trigger-rule compensation (COMPLETED_WITH_FAILURES). + +charge (root) always fails. fulfill uses the default (ALL_SUCCESS) trigger, so it +is SKIPPED. refund uses ALL_FAILED and runs ("refunded"). audit uses ALL_DONE and +runs ("logged"). charge exhausts the DAG default retry policy before failing +terminally, so the DAG drains to COMPLETED_WITH_FAILURES without throwing. +Returns the canonical summary defined by test-requirements/dag/10-2.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import ( + DagConfig, + DagContext, + DagResult, + TriggerRule, +) +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +def _charge_declined(_deps: Any, _sc: Any) -> Any: + raise RuntimeError("payment declined") + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + charge = d.step(_charge_declined, name="charge") + d.step(lambda deps, sc: "fulfilled", name="fulfill").after(charge) + d.step(lambda deps, sc: "refunded", name="refund").after(charge).trigger_rule( + TriggerRule.ALL_FAILED + ) + d.step(lambda deps, sc: "logged", name="audit").after(charge).trigger_rule( + TriggerRule.ALL_DONE + ) + + result = context.dag( + register, name="compensation", config=DagConfig(max_concurrency=1) + ) + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_concurrent_overlap.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_concurrent_overlap.py new file mode 100644 index 00000000..7e9e6fe5 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_concurrent_overlap.py @@ -0,0 +1,100 @@ +"""DAG conformance 10-13: real overlap of two tasks inside one invocation. + +max_concurrency is UNSET, so ready siblings run concurrently on real threads. + +root returns 1. slow and fast both depend on root; slow (registered FIRST) +sleeps ~2s in-body, fast sleeps ~200ms — so fast finishes first even though it +was registered second. afterSlow (registered FIRST) depends on slow and returns +deps.slow + "s" ("Ss"); afterFast depends on fast and returns deps.fast + "f" +("Ff"). Because fast is ready and completes before slow, afterFast starts before +afterSlow — the inversion of registration order versus start order that a +counter-based id scheme cannot survive (it would hand out different ids on +replay and terminate the execution with a replay-consistency error). merge +depends on [afterSlow, afterFast] and returns "SsFf". + +Peak-concurrency instrumentation: a shared counter is incremented on entry to +slow/fast and decremented on exit, tracking the maximum observed. Python runs +tasks on real OS threads, so the counter is guarded by a lock. It is returned as +peakConcurrency (expected 2); without it the scenario would silently go vacuous +if a future change serialized the scheduler. + +Returns the canonical summary defined by test-requirements/dag/10-13.yaml. +""" + +from __future__ import annotations + +import threading +import time +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + # Shared peak-concurrency tracker. Bodies run on real threads, so mutation + # is serialized by a lock; ``current`` is the live in-flight count and + # ``peak`` the maximum ever observed simultaneously. + tracker = {"current": 0, "peak": 0} + lock = threading.Lock() + + def _enter() -> None: + with lock: + tracker["current"] += 1 + tracker["peak"] = max(tracker["peak"], tracker["current"]) + + def _leave() -> None: + with lock: + tracker["current"] -= 1 + + def slow(_deps: Any, _sc: Any) -> str: + _enter() + try: + time.sleep(2) + return "S" + finally: + _leave() + + def fast(_deps: Any, _sc: Any) -> str: + _enter() + try: + time.sleep(0.2) + return "F" + finally: + _leave() + + def register(d: DagContext) -> None: + root = d.step(lambda deps, sc: 1, name="root") + slow_h = d.step(slow, deps=[root], name="slow") # registered FIRST + fast_h = d.step(fast, deps=[root], name="fast") + after_slow = d.step( # registered FIRST + lambda deps, sc: deps[slow_h] + "s", deps=[slow_h], name="afterSlow" + ) + after_fast = d.step( + lambda deps, sc: deps[fast_h] + "f", deps=[fast_h], name="afterFast" + ) + d.step( + lambda deps, sc: deps[after_slow] + deps[after_fast], + deps=[after_slow, after_fast], + name="merge", + ) + + result = context.dag(register, name="overlapdag") + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "merge": result.results["merge"].result, + "peakConcurrency": tracker["peak"], + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_concurrent_suspend.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_concurrent_suspend.py new file mode 100644 index 00000000..6051c74c --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_concurrent_suspend.py @@ -0,0 +1,65 @@ +"""DAG conformance 10-14: inverted readiness across a suspend. + +max_concurrency is UNSET. This is the replay-flip case: two Wait tasks are in +flight when the invocation suspends, and the downstream pair becomes ready in +the reverse of registration order across different invocations. + +root returns 1. slow is an 8s Wait and fast is a 2s Wait, both depending on root; +slow is registered FIRST. Both waits start in the first invocation, so the +invocation suspends with TWO tasks in flight and resumes twice. afterSlow +(registered FIRST) has an ordering-only edge .after(slow) and returns "S"; +afterFast has .after(fast) and returns "F". Because fast's timer fires first, +afterFast becomes ready — and starts — one invocation before afterSlow, the +inversion of registration order that a counter-based id scheme cannot survive. +merge depends on [afterSlow, afterFast] and returns "SF". + +Timers, not races, decide the order, so the outcome is deterministic; the gap +between the two waits is 6s (>> the ~4s floor). No peak-concurrency assertion is +possible or needed here — the waits are not user code and the suspend boundary +is the point. + +Returns the canonical summary defined by test-requirements/dag/10-14.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.dag import DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + root = d.step(lambda deps, sc: 1, name="root") + slow = d.wait(Duration.from_seconds(8), deps=[root], name="slow") # registered FIRST + fast = d.wait(Duration.from_seconds(2), deps=[root], name="fast") + after_slow = d.step(lambda deps, sc: "S", name="afterSlow").after( + slow + ) # registered FIRST + after_fast = d.step(lambda deps, sc: "F", name="afterFast").after(fast) + d.step( + lambda deps, sc: deps[after_slow] + deps[after_fast], + deps=[after_slow, after_fast], + name="merge", + ) + + result = context.dag(register, name="suspenddag") + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "merge": result.results["merge"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_diamond.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_diamond.py new file mode 100644 index 00000000..ba169f0b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_diamond.py @@ -0,0 +1,40 @@ +"""DAG conformance 10-1: diamond fan-out/fan-in (all tasks complete). + +fetch(10) -> {ta(=fetch+1=11), tb(=fetch*2=20)} -> merge(=ta+tb=31). +max_concurrency=1 for a deterministic topological order. Returns the canonical +cross-language summary defined by test-requirements/dag/10-1.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + fetch = d.step(lambda deps, sc: 10, name="fetch") + ta = d.step(lambda deps, sc: deps[fetch] + 1, deps=[fetch], name="ta") + tb = d.step(lambda deps, sc: deps[fetch] * 2, deps=[fetch], name="tb") + d.step(lambda deps, sc: deps[ta] + deps[tb], deps=[ta, tb], name="merge") + + result = context.dag(register, name="diamond", config=DagConfig(max_concurrency=1)) + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "merge": result.results["merge"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_id_stability.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_id_stability.py new file mode 100644 index 00000000..8303d6f2 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_id_stability.py @@ -0,0 +1,70 @@ +"""DAG conformance 10-20: task-id stability across independently forced completion orders. + +Identical shape to 10-13's overlap (DagConcurrentOverlap) -- root -> {a, b} -> +{afterA, afterB} -> merge, max_concurrency unset -- except which sibling +sleeps longer is driven by event["swap"]: swap=False makes a finish first; +swap=True makes b finish first. Both invocations register the SAME task names +in the SAME order every time -- only the RUNTIME completion order changes. + +This is the harness-level counterpart to 10-13: 10-13 proves out-of-order +completion doesn't fail the execution (an INDIRECT proof of name-based ids, +since a counter-based scheme would trip the SDK's own replay-consistency +check). This scenario is invoked TWICE by a dedicated script +(id_stability.py, not the normal single-invocation validator) with swap +flipped between runs, and asserts each task's Id field in the captured +execution history is IDENTICAL across both runs -- the direct proof that ids +are derived from the task name, not from completion order or a counter. + +Returns the canonical summary defined by test-requirements/dag/10-20.yaml. +""" + +from __future__ import annotations + +import time +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +@durable_execution +def handler(event: Any, context: DurableContext) -> dict[str, Any]: + swap = bool((event or {}).get("swap", False)) + + def a_body(_deps: Any, _sc: Any) -> str: + time.sleep(2 if swap else 0.2) + return "A" + + def b_body(_deps: Any, _sc: Any) -> str: + time.sleep(0.2 if swap else 2) + return "B" + + def register(d: DagContext) -> None: + root = d.step(lambda deps, sc: 1, name="root") + a_h = d.step(a_body, deps=[root], name="a") + b_h = d.step(b_body, deps=[root], name="b") + after_a = d.step(lambda deps, sc: deps[a_h] + "a", deps=[a_h], name="afterA") + after_b = d.step(lambda deps, sc: deps[b_h] + "b", deps=[b_h], name="afterB") + d.step( + lambda deps, sc: deps[after_a] + deps[after_b], + deps=[after_a, after_b], + name="merge", + ) + + result = context.dag(register, name="idstabilitydag") + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "merge": result.results["merge"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_invoke.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_invoke.py new file mode 100644 index 00000000..349f2350 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_invoke.py @@ -0,0 +1,58 @@ +"""DAG conformance 10-10: task that is an invoke of another Lambda (flat). + +prep(step->21) -> call(invoke[dep prep]: target echoes the payload -> 21) -> +done(step[dep call] -> call*2 = 42). The invoke's native Invoke op is +checkpointed directly under the Dag container (flat, name-based). The target is +the shared echo function (``invoke.target_echo``) reached via the +``TARGET_FUNCTION_NAME`` environment variable and returns whatever it receives, +so ``call`` resolves to the integer ``prep`` produced. max_concurrency=1 for a +deterministic topological order. + +This scenario suspends and resumes: the invoke completes in a later invocation. + +Returns the canonical summary defined by test-requirements/dag/10-10.yaml. +""" + +from __future__ import annotations + +import os +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + function_name = os.environ["TARGET_FUNCTION_NAME"] + + def register(d: DagContext) -> None: + prep = d.step(lambda deps, sc: 21, name="prep") + call = d.invoke( + function_name, + lambda deps: deps[prep], + deps=[prep], + name="call", + ) + d.step(lambda deps, sc: deps[call] * 2, deps=[call], name="done") + + result = context.dag( + register, name="invokedag", config=DagConfig(max_concurrency=1) + ) + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "call": result.results["call"].result, + "done": result.results["done"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_large_payload.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_large_payload.py new file mode 100644 index 00000000..ba4ad8f5 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_large_payload.py @@ -0,0 +1,110 @@ +"""DAG conformance 10-15: large-payload aggregate offload survives a replay. + +max_concurrency=1 for determinism. DAG name ``bigdag``. Eight root step tasks +``p1``..``p8`` with no deps; task ``pN`` returns its own letter repeated 51200 +times (``p1`` -> "a" x 51200, ``p2`` -> "b" x 51200, ... ``p8`` -> "h" x 51200). +The aggregate is ~410KB (8 * 51200 = 409600 chars), comfortably over the 256KB +checkpoint threshold, so the container result is OFFLOADED; every individual +task result stays far under it, so only the aggregate is offloaded. + +The reconstruct-vs-re-execute divergence (JS writes a DagSummary envelope and +reconstructs from it; Python/Java/Go re-execute the DAG child body via +ReplayChildren with no envelope) only fires when a SUCCEEDED CONTAINER IS +REPLAYED. A DAG that completes and returns in one invocation never exercises it. +So this handler deliberately SUSPENDS after the DAG resolves, via a 2s wait, so +the next invocation replays the completed container: + +1. ``dag(...)`` resolves the ~410KB aggregate. +2. A step (outside the DAG) computes a digest from the DagResult: + ``"::"`` -> exactly + ``"8:409600:abcdefgh"``. Because it is a step it is checkpointed and survives + the suspend as ``digestBefore``. +3. A ``wait`` of 2 seconds ends the invocation. +4. After the resume, the same digest is recomputed from the REPLAYED DagResult + -> ``digestAfter``. + +The language-neutral assertion is ``digestBefore == digestAfter == +"8:409600:abcdefgh"``: the aggregate survived the offload AND came back +identical through whichever replay strategy the SDK uses (child-body +re-execution here). Assert outcome only -- the container's succeeded payload +legitimately differs across SDKs, so 10-15.yaml pins no ExpectedExecutionHistory. + +This scenario deliberately uses NO completionConfig: Python has a documented +exception where a faithful STARTED-set is not reproduced under large-payload +early completion (inherited from map/parallel). All eight tasks complete, so it +does not walk into that. The returned summary is kept small -- the digest, never +the payload. + +Returns the canonical summary defined by test-requirements/dag/10-15.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + +_TASK_COUNT = 8 +_REPEAT = 51200 +_TASK_NAMES = [f"p{i}" for i in range(1, _TASK_COUNT + 1)] + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +def _digest(result: DagResult) -> str: + """``"::"``. + + Computed over p1..p8 in order, so the first-char run is deterministic + regardless of task completion order. For the 10-15 graph this is exactly + ``"8:409600:abcdefgh"``. + """ + total_length = 0 + first_chars = [] + for name in _TASK_NAMES: + value = result.get_result(name) + total_length += len(value) + first_chars.append(value[0]) + return f"{len(_TASK_NAMES)}:{total_length}:{''.join(first_chars)}" + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + for i, name in enumerate(_TASK_NAMES): + letter = chr(ord("a") + i) + # Default-arg binding captures this task's letter (avoids the + # late-binding closure trap over the loop variable). + d.step(lambda deps, sc, _letter=letter: _letter * _REPEAT, name=name) + + result = context.dag(register, name="bigdag", config=DagConfig(max_concurrency=1)) + + # Checkpointed step: computed once from the live DagResult, then fast-pathed + # from its own checkpoint on the post-suspend replay, so it carries the + # pre-suspend digest across the boundary. + digest_before: str = context.step( + lambda _sc: _digest(result), name="digestBefore" + ) + + # Forces the invocation to end; the next one replays the completed container. + context.wait(Duration.from_seconds(2), name="pauseForReplay") + + # Recomputed from the REPLAYED DagResult after resume. + digest_after = _digest(result) + + return { + "reason": result.completion_reason.value, + "counts": _counts(result), + "digestBefore": digest_before, + "digestAfter": digest_after, + "match": digest_before == digest_after, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_map.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_map.py new file mode 100644 index 00000000..1bdc8e81 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_map.py @@ -0,0 +1,50 @@ +"""DAG conformance 10-6: task that is a map over a fixed item list (flat map). + +squares(map over [1, 2]; each item one step -> item*item => [1, 4]) -> +sum(step[dep squares] -> sum of successful results = 5). The map's native Map op +is checkpointed directly under the Dag container (flat, name-based). +max_concurrency=1 (both DAG and map) for a deterministic history. Returns the +canonical summary defined by test-requirements/dag/10-6.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.config import MapConfig +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + def square(ctx: DurableContext, item: int, _index: int, _items: Any) -> int: + return ctx.step(lambda _sc: item * item, name="square") + + squares = d.map( + [1, 2], square, name="squares", config=MapConfig(max_concurrency=1) + ) + d.step( + lambda deps, sc: sum(deps[squares].get_results()), + deps=[squares], + name="sum", + ) + + result = context.dag(register, name="mapdag", config=DagConfig(max_concurrency=1)) + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "sum": result.results["sum"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_nested.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_nested.py new file mode 100644 index 00000000..ff1394d1 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_nested.py @@ -0,0 +1,55 @@ +"""DAG conformance 10-9: task that is itself a nested DAG / sub-dag (flat). + +pre(step->1) -> sub(nested dag[dep pre]: n1->2, n2[dep n1]->n1+3=5) -> +post(step[dep sub] -> nested n2 result * 10 = 50). The nested DAG's native Dag +op is checkpointed directly under the outer Dag container (flat, name-based). +max_concurrency=1 at both DAG levels for a deterministic topological order. +Returns the canonical summary defined by test-requirements/dag/10-9.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + pre = d.step(lambda deps, sc: 1, name="pre") + + def sub_register(sd: DagContext) -> None: + n1 = sd.step(lambda deps, sc: 2, name="n1") + sd.step(lambda deps, sc: deps[n1] + 3, deps=[n1], name="n2") + + sub = d.dag( + sub_register, + deps=[pre], + name="sub", + config=DagConfig(max_concurrency=1), + ) + d.step( + lambda deps, sc: deps[sub].get_result("n2") * 10, + deps=[sub], + name="post", + ) + + result = context.dag(register, name="outerdag", config=DagConfig(max_concurrency=1)) + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "post": result.results["post"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_nested_large_payload.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_nested_large_payload.py new file mode 100644 index 00000000..121699fb --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_nested_large_payload.py @@ -0,0 +1,124 @@ +"""DAG conformance 10-17: nested DAG whose inner aggregate offloads, across a +container replay -- the untested intersection of nesting and large payloads. + +Modeled on ``10-15`` (flat large payload) but with the large aggregate produced +by a NESTED DAG. ``max_concurrency=1`` everywhere for determinism. + +Outer DAG ``outernested`` has a single task ``inner`` that is itself a nested DAG +with six root step tasks ``p1``..``p6``; task ``pN`` returns its letter repeated +51200 times (``p1`` -> "a" x 51200, ... ``p6`` -> "f" x 51200). The inner +aggregate is ~307KB (6 * 51200 = 307200 chars), comfortably over the 256KB +checkpoint threshold, so the INNER container is OFFLOADED. Because the outer +embeds the inner result in full, the OUTER aggregate is over the limit too, so +the outer container is offloaded as well. Every individual inner task result +stays far under the limit, so only the two aggregates offload. + +The reconstruct-vs-inline divergence only fires when a SUCCEEDED CONTAINER IS +REPLAYED, so -- exactly as in 10-15 -- the handler resolves the DAG, records a +checkpointed digest, then SUSPENDS on an outer 2s wait so the next invocation +replays BOTH completed containers: + +1. ``dag(...)`` resolves ``outernested``; its one task ``inner`` resolves the + ~307KB inner aggregate (inner container offloaded, outer container offloaded). +2. An outer (handler-level) step reads the inner ``DagResult`` and computes a + compact digest ``"::"`` + -> exactly ``"6:307200:abcdef"``. Being a checkpointed step it survives the + suspend as ``digestBefore``. +3. An outer ``wait`` of 2 seconds ends the invocation; the next one replays both + completed containers. +4. After the resume the identical digest is recomputed from the REPLAYED inner + result -> ``digestAfter``. + +On replay the outer container reconstructs from its retained child checkpoints; +its ``inner`` task re-runs through the DAG container executor, which detects the +offloaded inner and reconstructs it RECURSIVELY from the inner's own child +checkpoints, restoring full per-task detail. The decisive, language-neutral +assertion is ``digestBefore == digestAfter == "6:307200:abcdef"`` with +``match: true``: it proves the inner per-task detail survived the offload of BOTH +containers. Under the bug the inner comes back empty, so ``digestAfter`` differs +while ``innerReason`` would still read ``ALL_COMPLETED`` from a fabricated +result -- which is exactly why the digest, not the reason, is the decisive check. + +Returns the canonical summary defined by test-requirements/dag/10-17.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + +_INNER_COUNT = 6 +_REPEAT = 51200 +_INNER_NAMES = [f"p{i}" for i in range(1, _INNER_COUNT + 1)] + + +def _inner_counts(inner: DagResult) -> list[int]: + """``[total, failed, skipped, succeeded]`` for the inner DagResult.""" + return [ + inner.total_count, + inner.failure_count, + inner.skipped_count, + inner.success_count, + ] + + +def _digest(inner: DagResult) -> str: + """``"::"``. + + Computed over p1..p6 in order, so the first-char run is deterministic + regardless of completion order. For the 10-17 inner graph this is exactly + ``"6:307200:abcdef"``. + """ + total_length = 0 + first_chars = [] + for name in _INNER_NAMES: + value = inner.get_result(name) + total_length += len(value) + first_chars.append(value[0]) + return f"{len(_INNER_NAMES)}:{total_length}:{''.join(first_chars)}" + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + def inner_register(sd: DagContext) -> None: + for i, name in enumerate(_INNER_NAMES): + letter = chr(ord("a") + i) + # Default-arg binding captures this task's letter (avoids the + # late-binding closure trap over the loop variable). + sd.step( + lambda deps, sc, _letter=letter: _letter * _REPEAT, name=name + ) + + d.dag(inner_register, name="inner", config=DagConfig(max_concurrency=1)) + + result = context.dag( + register, name="outernested", config=DagConfig(max_concurrency=1) + ) + + # Checkpointed step: computed once from the live inner DagResult, then + # fast-pathed from its own checkpoint on the post-suspend replay, so it + # carries the pre-suspend digest across the boundary. + digest_before: str = context.step( + lambda _sc: _digest(result.get_result("inner")), name="digestBefore" + ) + + # Forces the invocation to end; the next one replays both completed containers. + context.wait(Duration.from_seconds(2), name="settle") + + # Recomputed from the REPLAYED (recursively reconstructed) inner result. + inner = result.get_result("inner") + digest_after = _digest(inner) + + return { + "reason": result.completion_reason.value, + "innerReason": inner.completion_reason.value, + "innerCounts": _inner_counts(inner), + "digestBefore": digest_before, + "digestAfter": digest_after, + "match": digest_before == digest_after, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_parallel.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_parallel.py new file mode 100644 index 00000000..334c5d2f --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_parallel.py @@ -0,0 +1,69 @@ +"""DAG conformance 10-7: task that is a parallel of two named branches (flat). + +fork(parallel of branches left->"L", right->"R") -> join(step[dep fork]). The +parallel's native Parallel op is checkpointed directly under the Dag container +(flat, name-based). max_concurrency=1 (both DAG and parallel) for a +deterministic history. + +Aggregate-only join: ``join`` reads ONLY the aggregate ParallelResult / +BatchResult handed to it as the dep value and returns ``"/"`` +(``"2/2"``). It does not read individual branch values, keeping the scenario +expressible in every SDK (Java's ``ParallelResult`` is aggregate-only, and a +durable-op read from inside a step body is illegal there). Reading child branch +values is still covered by 10-6 (map). + +Returns the canonical summary defined by test-requirements/dag/10-7.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.config import ParallelBranch, ParallelConfig +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + def left(ctx: DurableContext) -> str: + return ctx.step(lambda _sc: "L", name="left-step") + + def right(ctx: DurableContext) -> str: + return ctx.step(lambda _sc: "R", name="right-step") + + fork = d.parallel( + [ + ParallelBranch(func=left, name="left"), + ParallelBranch(func=right, name="right"), + ], + name="fork", + config=ParallelConfig(max_concurrency=1), + ) + + def join(deps, _sc) -> str: + aggregate = deps[fork] + return f"{aggregate.success_count}/{aggregate.total_count}" + + d.step(join, deps=[fork], name="join") + + result = context.dag( + register, name="paralleldag", config=DagConfig(max_concurrency=1) + ) + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "join": result.results["join"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_retry.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_retry.py new file mode 100644 index 00000000..0cb77abf --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_retry.py @@ -0,0 +1,76 @@ +"""DAG conformance 10-16: per-task retry inside a DAG (DagRetry). + +A DAG where a task retries and eventually succeeds, proving that a retried +task's result flows downstream normally: + +- ``flaky`` is a step with a per-task retry strategy allowing at least three + attempts and no meaningful backoff. Its body reads the 1-based attempt number + from the step context (``sc.attempt``). It throws while the attempt is not yet + the third, and returns the attempt number (``3``) on the third attempt. +- ``after`` is a step depending on ``flaky`` that returns ``flaky``'s result + doubled (``6``). + +``flaky`` must end SUCCEEDED (not FAILED) and ``after`` must run (not skip), +which is exactly what a broken retry inside a DAG would break. ``max_concurrency`` +is 1 for a deterministic single-lane order. + +Handler returns ``{"flaky": , "after": }`` — i.e. +``{"flaky": 3, "after": 6}`` — the canonical values pinned by +test-requirements/dag/10-16.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.config import ( + Duration, + JitterStrategy, + StepConfig, +) +from aws_durable_execution_sdk_python.context import DurableContext, StepContext +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python.retries import ( + RetryStrategyConfig, + create_retry_strategy, +) + + +def _flaky(_deps: Any, sc: StepContext) -> int: + # Read the SDK's built-in durable attempt counter (1-based) from the step + # context. Fail until the third attempt, then return the attempt number. + if sc.attempt < 3: + msg = f"attempt {sc.attempt} not yet the third" + raise RuntimeError(msg) + return sc.attempt + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + # At least three attempts, no backoff delay worth waiting on (the step + # executor floors any retry delay at 1 second regardless). + retry_strategy = create_retry_strategy( + RetryStrategyConfig( + max_attempts=5, + initial_delay=Duration.from_seconds(1), + backoff_rate=1, + jitter_strategy=JitterStrategy.NONE, + ) + ) + + def register(d: DagContext) -> None: + flaky = d.step( + _flaky, + name="flaky", + config=StepConfig(retry_strategy=retry_strategy), + ) + d.step(lambda deps, sc: deps[flaky] * 2, deps=[flaky], name="after") + + result: DagResult = context.dag( + register, name="retrydag", config=DagConfig(max_concurrency=1) + ) + return { + "flaky": result.results["flaky"].result, + "after": result.results["after"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_rules_engine.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_rules_engine.py new file mode 100644 index 00000000..87da93a9 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_rules_engine.py @@ -0,0 +1,85 @@ +"""DAG conformance 10-19: custom result-based completion. A rules-engine +predicate short-circuits the moment any task's SUCCEEDED result carries a +REJECT verdict -- expressible only because the custom-completion predicate can +inspect task RESULTS, not just aggregate counts. + +A DAG "rulesengine" with max-concurrency 1 and a linear chain of three step +tasks: r1 -> r2 -> r3, each returning a verdict dict. r1 -> ACCEPT, r2 -> +REJECT, r3 (never runs) -> ACCEPT. + +The DAG's completion_config is a custom predicate (DagCustomCompletionConfig), +not a threshold: after every settlement it receives a live DagCompletionStatus +snapshot and inspects every SUCCEEDED item's result for a REJECT verdict. The +moment it sees one, it returns complete_dag(FAILED). r3 is never started and +is absent from the results map. The DAG completes with +CUSTOM_COMPLETION_FAILED -- a dedicated reason distinct from +COMPLETED_WITH_FAILURES, since no individual task FAILED. throw_if_error() +MUST still raise in this case (the contract keys off completion_reason too, +not failure_count alone). + +Returns the canonical summary defined by test-requirements/dag/10-19.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import ( + DagCompletionOutcome, + DagCompletionStatus, + DagConfig, + DagContext, + DagCustomCompletionConfig, + DagResult, + TaskStatus, + complete_dag, + continue_dag, +) +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +def _should_complete(status: DagCompletionStatus) -> Any: + rejected = any( + item.status is TaskStatus.SUCCEEDED + and isinstance(item.result, dict) + and item.result.get("verdict") == "REJECT" + for item in status.items + ) + if rejected: + return complete_dag(DagCompletionOutcome.FAILED) + return continue_dag() + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + r1 = d.step(lambda deps, sc: {"verdict": "ACCEPT"}, name="r1") + r2 = d.step( + lambda deps, sc: {"verdict": "REJECT"}, deps=[r1], name="r2" + ) + d.step(lambda deps, sc: {"verdict": "ACCEPT"}, deps=[r2], name="r3") + + result = context.dag( + register, + name="rulesengine", + config=DagConfig( + max_concurrency=1, + completion_config=DagCustomCompletionConfig(_should_complete), + ), + ) + return { + "reason": result.completion_reason.value, + "counts": _counts(result), + "r1": result.get_result("r1"), + "r2": result.get_result("r2"), + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_run_if.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_run_if.py new file mode 100644 index 00000000..0607be9a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_run_if.py @@ -0,0 +1,52 @@ +"""DAG conformance 10-3: per-task conditional execution (run_if). + +classify returns "review". publish/review/block each depend on classify and are +guarded by a run_if predicate that runs the branch only when classify's result +equals the branch's own name. Only review runs; publish and block are SKIPPED and +emit no events. Returns the canonical summary defined by +test-requirements/dag/10-3.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + +_BRANCHES = ("publish", "review", "block") + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + classify = d.step(lambda deps, sc: "review", name="classify") + for branch in _BRANCHES: + d.step( + (lambda name: (lambda deps, sc: name))(branch), + deps=[classify], + name=branch, + run_if=(lambda name: (lambda deps: deps[classify] == name))(branch), + ) + + result = context.dag(register, name="runif", config=DagConfig(max_concurrency=1)) + statuses = {name: te.status.value for name, te in result.results.items()} + branch = next( + name for name in _BRANCHES if statuses.get(name) == "SUCCEEDED" + ) + return { + "reason": result.completion_reason.value, + "statuses": statuses, + "counts": _counts(result), + "branch": branch, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_run_if_abort.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_run_if_abort.py new file mode 100644 index 00000000..da191e09 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_run_if_abort.py @@ -0,0 +1,65 @@ +"""DAG conformance 10-12: a throwing run_if aborts the DAG with a typed error. + +Serial (max_concurrency=1), so this scenario keeps full history assertions. + +gate (root) returns 1. guarded depends on gate and its run_if THROWS ("predicate +boom") — a defect in deterministic code, not a business outcome. Per the runIf +abort contract the scheduler neither records guarded FAILED nor SKIPPED: it +aborts, starts no further tasks, and context.dag() fails with DagPredicateError. +guarded's body ("ran") is never invoked. refund has an ordering-only edge +.after(guarded) with an ALL_FAILED trigger and returns "refunded"; it MUST NOT +run — the whole point of the abort contract is that a predicate defect does not +drive compensation. + +The handler does NOT catch the error: the DagPredicateError propagates so the +execution FAILS. The error-type token differs per language, so the canonical +YAML wildcards it. Returns nothing on the success path (never reached). +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import ( + DagConfig, + DagContext, + DagResult, + TriggerRule, +) +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +def _boom(_deps: Any) -> bool: + raise RuntimeError("predicate boom") + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + gate = d.step(lambda deps, sc: 1, name="gate") + guarded = d.step( + lambda deps, sc: "ran", deps=[gate], name="guarded", run_if=_boom + ) + d.step(lambda deps, sc: "refunded", name="refund").after(guarded).trigger_rule( + TriggerRule.ALL_FAILED + ) + + # A throwing run_if aborts the DAG: this raises DagPredicateError, failing + # the execution. The code below is unreachable and exists only to mirror the + # other handlers' summary shape. + result = context.dag(register, name="abortdag", config=DagConfig(max_concurrency=1)) + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_wait_for_condition.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_wait_for_condition.py new file mode 100644 index 00000000..c462ce86 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_wait_for_condition.py @@ -0,0 +1,60 @@ +"""DAG conformance 10-8: task that is a waitForCondition (flat WaitForCondition). + +poll(waitForCondition: initial state 0, +1 per poll, stop at 2 => returns 2) -> +done(step[dep poll] -> poll*5 = 10). The waitForCondition's native +WaitForCondition op is checkpointed directly under the Dag container (flat, +name-based); the first poll suspends and the DAG resumes across the +suspend/resume boundary. max_concurrency=1. Returns the canonical summary +defined by test-requirements/dag/10-8.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python.types import WaitForConditionCheckContext +from aws_durable_execution_sdk_python.waits import ( + WaitForConditionConfig, + WaitForConditionDecision, +) + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + def check(_deps, state: int, _ctx: WaitForConditionCheckContext) -> int: + return state + 1 + + def wait_strategy(state: int, _attempt: int) -> WaitForConditionDecision: + if state >= 2: + return WaitForConditionDecision.stop_polling() + return WaitForConditionDecision.continue_waiting(Duration.from_seconds(1)) + + poll = d.wait_for_condition( + check, + WaitForConditionConfig(wait_strategy=wait_strategy, initial_state=0), + name="poll", + ) + d.step(lambda deps, sc: deps[poll] * 5, deps=[poll], name="done") + + result = context.dag(register, name="wfcdag", config=DagConfig(max_concurrency=1)) + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "poll": result.results["poll"].result, + "done": result.results["done"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_wait_resume.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_wait_resume.py new file mode 100644 index 00000000..ae6e2660 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_wait_resume.py @@ -0,0 +1,43 @@ +"""DAG conformance 10-4: in-graph Wait task (suspend and resume). + +start -> pause(Wait 5s) -> finish. pause suspends the whole invocation without +compute charges until the wait elapses, then resumes in a fresh invocation. +finish returns "resumed", proving the DAG ran across the suspend/resume boundary. +Returns the canonical summary defined by test-requirements/dag/10-4.yaml. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.dag import DagConfig, DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _counts(result: DagResult) -> list[int]: + return [ + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ] + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + def register(d: DagContext) -> None: + start = d.step(lambda deps, sc: "started", name="start") + pause = d.wait(Duration.from_seconds(5), deps=[start], name="pause") + d.step(lambda deps, sc: "resumed", deps=[pause], name="finish") + + result = context.dag( + register, name="waitresume", config=DagConfig(max_concurrency=1) + ) + return { + "reason": result.completion_reason.value, + "statuses": {name: te.status.value for name, te in result.results.items()}, + "counts": _counts(result), + "marker": result.results["finish"].result, + } diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml new file mode 100644 index 00000000..7258a92d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml @@ -0,0 +1,359 @@ +AWSTemplateFormatVersion: '2010-09-09' +Transform: AWS::Serverless-2016-10-31 +Globals: + Function: + Runtime: python3.13 + Timeout: 60 + MemorySize: 128 +Resources: + DurableFunctionRole: + Type: AWS::IAM::Role + Properties: + AssumeRolePolicyDocument: + Version: '2012-10-17' + Statement: + - Effect: Allow + Principal: + Service: lambda.amazonaws.com + Action: sts:AssumeRole + ManagedPolicyArns: + - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole + Policies: + - PolicyName: DurableExecutionPolicy + PolicyDocument: + Version: '2012-10-17' + Statement: + - Effect: Allow + Action: + - lambda:CheckpointDurableExecution + - lambda:GetDurableExecutionState + - lambda:InvokeFunction + Resource: '*' + DagDiamond: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-1"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_diamond.handler + Description: DAG diamond fan-out/fan-in (all tasks complete) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + DagCompensation: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-2"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_compensation.handler + Description: DAG trigger-rule compensation (COMPLETED_WITH_FAILURES) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + DagRunIf: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-3"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_run_if.handler + Description: DAG per-task conditional execution (run_if) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + DagWaitResume: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-4"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_wait_resume.handler + Description: DAG in-graph Wait task (suspend and resume) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + DagChild: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-5"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_child.handler + Description: DAG task that is a runInChildContext (flat child container) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + DagMap: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-6"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_map.handler + Description: DAG task that is a map over a fixed item list (flat map container) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + DagParallel: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-7"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_parallel.handler + Description: DAG task that is a parallel of two named branches (flat parallel container) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + DagWaitForCondition: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-8"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_wait_for_condition.handler + Description: DAG task that is a waitForCondition (flat WaitForCondition op) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + DagNested: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-9"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_nested.handler + Description: DAG task that is itself a nested sub-DAG (flat nested Dag container) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + + # Target function for DagInvoke (10-10); no TestingMetadata (not a test itself) + TargetEcho: + Type: AWS::Serverless::Function + Properties: + CodeUri: lambda-build/ + Handler: invoke.target_echo.handler + Description: Echo target function — returns whatever it receives + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + + DagInvoke: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-10"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_invoke.handler + Description: DAG task that is an invoke of another Lambda (flat invoke container) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + Environment: + Variables: + TARGET_FUNCTION_NAME: !Sub "${TargetEcho.Arn}:$LATEST" + + DagCallback: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-11"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_callback.handler + Description: DAG task that is a wait-for-callback (flat callback container) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + + DagRunIfAbort: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-12"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_run_if_abort.handler + Description: DAG throwing run_if aborts the DAG with a typed error + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + + DagConcurrentOverlap: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-13"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_concurrent_overlap.handler + Description: DAG real overlap of two tasks inside one invocation + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + + DagIdStability: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-20"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_id_stability.handler + Description: DAG task ids are name-based, verified by invoking twice with completion order swapped + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + + DagConcurrentSuspend: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-14"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_concurrent_suspend.handler + Description: DAG inverted readiness across a suspend (two waits in flight) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + + DagLargePayload: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-15"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_large_payload.handler + Description: DAG large-payload aggregate offload survives a container replay + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + + DagRetry: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-16"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_retry.handler + Description: DAG per-task retry inside a DAG succeeds and flows downstream + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + + DagNestedLargePayload: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-17"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_nested_large_payload.handler + Description: Nested DAG whose inner aggregate offloads both containers, survives a container replay + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + + DagCompensate: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-18"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_compensate.handler + Description: DAG compensation dependency read on a failed upstream is absent, not present (deps-nullability) + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 + DagRulesEngine: + Type: AWS::Serverless::Function + TestingMetadata: + TestDescription: ["10-19"] + Properties: + CodeUri: lambda-build/ + Handler: dag.dag_rules_engine.handler + Description: DAG custom result-based completion short-circuits on a rejected verdict + Role: + Fn::GetAtt: + - DurableFunctionRole + - Arn + DurableConfig: + RetentionPeriodInDays: 7 + ExecutionTimeout: 300 \ No newline at end of file diff --git a/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json b/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json index 030fce1a..6993f38c 100644 --- a/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json +++ b/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json @@ -774,6 +774,50 @@ "ExecutionTimeout": 300 }, "path": "./src/error_handling/catch_typed_step_error.py" + }, + { + "name": "DAG Diamond", + "description": "Diamond DAG (fetch -> ta,tb -> merge) using context.dag() with typed dependency access", + "handler": "dag_diamond.handler", + "integration": true, + "durableConfig": { + "RetentionPeriodInDays": 7, + "ExecutionTimeout": 300 + }, + "path": "./src/dag/dag_diamond.py" + }, + { + "name": "DAG Compensation", + "description": "Compensation DAG using trigger rules (ALL_FAILED refund, ALL_SUCCESS fulfill, ALL_DONE audit) draining to COMPLETED_WITH_FAILURES", + "handler": "dag_compensation.handler", + "integration": true, + "durableConfig": { + "RetentionPeriodInDays": 7, + "ExecutionTimeout": 300 + }, + "path": "./src/dag/dag_compensation.py" + }, + { + "name": "DAG Run If", + "description": "Conditional-branching DAG using run_if predicates so only the matching branch runs", + "handler": "dag_run_if.handler", + "integration": true, + "durableConfig": { + "RetentionPeriodInDays": 7, + "ExecutionTimeout": 300 + }, + "path": "./src/dag/dag_run_if.py" + }, + { + "name": "DAG Wait Resume", + "description": "DAG with a wait task (start -> pause -> finish) exercising suspend/resume", + "handler": "dag_wait_resume.handler", + "integration": true, + "durableConfig": { + "RetentionPeriodInDays": 7, + "ExecutionTimeout": 300 + }, + "path": "./src/dag/dag_wait_resume.py" } ] } diff --git a/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_compensation.py b/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_compensation.py new file mode 100644 index 00000000..275625dd --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_compensation.py @@ -0,0 +1,75 @@ +"""DAG example: compensation via trigger rules. + +Demonstrates the saga/compensation pattern with ``context.dag()`` trigger rules: +- ``fulfill`` runs only if ``charge`` SUCCEEDED (ALL_SUCCESS, the default) +- ``refund`` runs only if ``charge`` FAILED (ALL_FAILED) +- ``audit`` always runs once ``charge`` is terminal (ALL_DONE) + +By default ``charge`` fails, so the DAG drains to COMPLETED_WITH_FAILURES with +``refund`` + ``audit`` succeeding and ``fulfill`` skipped. Pass +``{"charge_ok": true}`` to exercise the success path. + +.. warning:: + Uses the EXPERIMENTAL ``context.dag()`` API. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.config import StepConfig +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagContext, DagResult, TriggerRule +from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python.retries import RetryPresets + + +def _summarize(result: DagResult) -> dict[str, Any]: + return { + "completion_reason": result.completion_reason.value, + "counts": { + "success": result.success_count, + "failure": result.failure_count, + "skipped": result.skipped_count, + "total": result.total_count, + }, + "tasks": { + name: { + "status": te.status.value, + "skip_reason": te.skip_reason.value if te.skip_reason else None, + "result": te.result, + } + for name, te in result.results.items() + }, + } + + +def _charge_declined(_deps: Any, _sc: Any) -> Any: + raise RuntimeError("charge declined") + + +@durable_execution +def handler(event: Any, context: DurableContext) -> dict[str, Any]: + """Run a compensation DAG driven by the terminal state of ``charge``.""" + charge_ok = bool(event.get("charge_ok", False)) if isinstance(event, dict) else False + + def register(d: DagContext) -> None: + # Disable retries on the charge so an intentional failure terminates promptly. + no_retry = StepConfig(retry_strategy=RetryPresets.none()) + if charge_ok: + charge = d.step(lambda deps, sc: "charged", name="charge", config=no_retry) + else: + charge = d.step(_charge_declined, name="charge", config=no_retry) + d.step(lambda deps, sc: "fulfilled", name="fulfill").after(charge) + d.step(lambda deps, sc: "refunded", name="refund").after(charge).trigger_rule( + TriggerRule.ALL_FAILED + ) + d.step(lambda deps, sc: "audited", name="audit").after(charge).trigger_rule( + TriggerRule.ALL_DONE + ) + + result = context.dag( + register, + name="compensation", + ) + return _summarize(result) diff --git a/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_diamond.py b/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_diamond.py new file mode 100644 index 00000000..f136e121 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_diamond.py @@ -0,0 +1,51 @@ +"""DAG example: diamond topology with typed dependency access. + +Demonstrates ``context.dag()`` with a fan-out/fan-in ("diamond") graph where +downstream tasks read upstream results via the typed ``deps[handle]`` accessor. + +.. warning:: + Uses the EXPERIMENTAL ``context.dag()`` API. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _summarize(result: DagResult) -> dict[str, Any]: + """Build a JSON-serializable summary of a DagResult for assertion.""" + return { + "completion_reason": result.completion_reason.value, + "counts": { + "success": result.success_count, + "failure": result.failure_count, + "skipped": result.skipped_count, + "total": result.total_count, + }, + "tasks": { + name: { + "status": te.status.value, + "skip_reason": te.skip_reason.value if te.skip_reason else None, + "result": te.result, + } + for name, te in result.results.items() + }, + } + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + """Run a diamond DAG: fetch -> (ta, tb) -> merge.""" + + def register(d: DagContext) -> None: + fetch = d.step(lambda deps, sc: 10, name="fetch") + ta = d.step(lambda deps, sc: deps[fetch] + 1, deps=[fetch], name="ta") + tb = d.step(lambda deps, sc: deps[fetch] * 2, deps=[fetch], name="tb") + d.step(lambda deps, sc: deps[ta] + deps[tb], deps=[ta, tb], name="merge") + + result = context.dag(register, name="diamond") + return _summarize(result) diff --git a/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_run_if.py b/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_run_if.py new file mode 100644 index 00000000..d86d86a5 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_run_if.py @@ -0,0 +1,68 @@ +"""DAG example: conditional branching with ``run_if``. + +A ``classify`` task produces a category; three downstream branches each declare a +``run_if`` predicate so only the matching branch runs and the others are SKIPPED +with reason RUN_IF_PREDICATE. The category defaults to ``"review"`` and can be +overridden via ``{"category": "publish" | "review" | "block"}``. + +.. warning:: + Uses the EXPERIMENTAL ``context.dag()`` API. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.dag import DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _summarize(result: DagResult) -> dict[str, Any]: + return { + "completion_reason": result.completion_reason.value, + "counts": { + "success": result.success_count, + "failure": result.failure_count, + "skipped": result.skipped_count, + "total": result.total_count, + }, + "tasks": { + name: { + "status": te.status.value, + "skip_reason": te.skip_reason.value if te.skip_reason else None, + "result": te.result, + } + for name, te in result.results.items() + }, + } + + +@durable_execution +def handler(event: Any, context: DurableContext) -> dict[str, Any]: + """Run a run_if-branching DAG selecting exactly one downstream branch.""" + category = event.get("category", "review") if isinstance(event, dict) else "review" + + def register(d: DagContext) -> None: + classify = d.step(lambda deps, sc: category, name="classify") + d.step( + lambda deps, sc: "published", + deps=[classify], + name="publish", + run_if=lambda deps: deps[classify] == "publish", + ) + d.step( + lambda deps, sc: "reviewed", + deps=[classify], + name="review", + run_if=lambda deps: deps[classify] == "review", + ) + d.step( + lambda deps, sc: "blocked", + deps=[classify], + name="block", + run_if=lambda deps: deps[classify] == "block", + ) + + result = context.dag(register, name="run_if_branching") + return _summarize(result) diff --git a/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_wait_resume.py b/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_wait_resume.py new file mode 100644 index 00000000..5933674f --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_wait_resume.py @@ -0,0 +1,51 @@ +"""DAG example: wait task forcing suspend/resume. + +A ``wait`` task sits between two steps so the execution suspends during the wait +and resumes to run the downstream task. Verifies DAG scheduling survives a real +suspend/resume cycle in the cloud. + +.. warning:: + Uses the EXPERIMENTAL ``context.dag()`` API. +""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.dag import DagContext, DagResult +from aws_durable_execution_sdk_python.execution import durable_execution + + +def _summarize(result: DagResult) -> dict[str, Any]: + return { + "completion_reason": result.completion_reason.value, + "counts": { + "success": result.success_count, + "failure": result.failure_count, + "skipped": result.skipped_count, + "total": result.total_count, + }, + "tasks": { + name: { + "status": te.status.value, + "skip_reason": te.skip_reason.value if te.skip_reason else None, + "result": te.result, + } + for name, te in result.results.items() + }, + } + + +@durable_execution +def handler(_event: Any, context: DurableContext) -> dict[str, Any]: + """Run a DAG with a wait task: start -> pause(wait) -> finish.""" + + def register(d: DagContext) -> None: + start = d.step(lambda deps, sc: "started", name="start") + pause = d.wait(Duration.from_seconds(3), deps=[start], name="pause") + d.step(lambda deps, sc: "resumed", deps=[pause], name="finish") + + result = context.dag(register, name="wait_resume") + return _summarize(result) diff --git a/packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_compensation.py b/packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_compensation.py new file mode 100644 index 00000000..d8fbce70 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_compensation.py @@ -0,0 +1,37 @@ +"""Cloud e2e tests for the DAG compensation example.""" + +import pytest +from aws_durable_execution_sdk_python.execution import InvocationStatus + +from src.dag import dag_compensation +from test.conftest import deserialize_operation_payload + + +@pytest.mark.example +@pytest.mark.durable_execution( + handler=dag_compensation.handler, + lambda_function_name="DAG Compensation", +) +def test_dag_compensation_charge_fails(durable_runner): + """Failed charge triggers refund + audit, skips fulfill, drains with failures.""" + with durable_runner: + result = durable_runner.run(input={"charge_ok": False}, timeout=60) + + assert result.status is InvocationStatus.SUCCEEDED + summary = deserialize_operation_payload(result.result) + + assert summary["completion_reason"] == "COMPLETED_WITH_FAILURES" + assert summary["counts"] == { + "success": 2, + "failure": 1, + "skipped": 1, + "total": 4, + } + tasks = summary["tasks"] + assert tasks["charge"]["status"] == "FAILED" + assert tasks["fulfill"]["status"] == "SKIPPED" + assert tasks["fulfill"]["skip_reason"] == "TRIGGER_RULE" + assert tasks["refund"]["status"] == "SUCCEEDED" + assert tasks["refund"]["result"] == "refunded" + assert tasks["audit"]["status"] == "SUCCEEDED" + assert tasks["audit"]["result"] == "audited" diff --git a/packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_diamond.py b/packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_diamond.py new file mode 100644 index 00000000..2c26fc7f --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_diamond.py @@ -0,0 +1,35 @@ +"""Cloud e2e tests for the DAG diamond example.""" + +import pytest +from aws_durable_execution_sdk_python.execution import InvocationStatus + +from src.dag import dag_diamond +from test.conftest import deserialize_operation_payload + + +@pytest.mark.example +@pytest.mark.durable_execution( + handler=dag_diamond.handler, + lambda_function_name="DAG Diamond", +) +def test_dag_diamond(durable_runner): + """Diamond DAG resolves typed deps and fans in to a merge task.""" + with durable_runner: + result = durable_runner.run(input={}, timeout=60) + + assert result.status is InvocationStatus.SUCCEEDED + summary = deserialize_operation_payload(result.result) + + assert summary["completion_reason"] == "ALL_COMPLETED" + assert summary["counts"] == { + "success": 4, + "failure": 0, + "skipped": 0, + "total": 4, + } + tasks = summary["tasks"] + assert tasks["fetch"]["result"] == 10 + assert tasks["ta"]["result"] == 11 + assert tasks["tb"]["result"] == 20 + assert tasks["merge"]["result"] == 31 + assert all(t["status"] == "SUCCEEDED" for t in tasks.values()) diff --git a/packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_run_if.py b/packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_run_if.py new file mode 100644 index 00000000..9f0c32c5 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_run_if.py @@ -0,0 +1,37 @@ +"""Cloud e2e tests for the DAG run_if branching example.""" + +import pytest +from aws_durable_execution_sdk_python.execution import InvocationStatus + +from src.dag import dag_run_if +from test.conftest import deserialize_operation_payload + + +@pytest.mark.example +@pytest.mark.durable_execution( + handler=dag_run_if.handler, + lambda_function_name="DAG Run If", +) +def test_dag_run_if_default_review(durable_runner): + """Only the branch whose run_if predicate matches runs; others are skipped.""" + with durable_runner: + result = durable_runner.run(input={"category": "review"}, timeout=60) + + assert result.status is InvocationStatus.SUCCEEDED + summary = deserialize_operation_payload(result.result) + + assert summary["completion_reason"] == "ALL_COMPLETED" + assert summary["counts"] == { + "success": 2, + "failure": 0, + "skipped": 2, + "total": 4, + } + tasks = summary["tasks"] + assert tasks["classify"]["result"] == "review" + assert tasks["review"]["status"] == "SUCCEEDED" + assert tasks["review"]["result"] == "reviewed" + assert tasks["publish"]["status"] == "SKIPPED" + assert tasks["publish"]["skip_reason"] == "RUN_IF_PREDICATE" + assert tasks["block"]["status"] == "SKIPPED" + assert tasks["block"]["skip_reason"] == "RUN_IF_PREDICATE" diff --git a/packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_wait_resume.py b/packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_wait_resume.py new file mode 100644 index 00000000..a28f6745 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_wait_resume.py @@ -0,0 +1,35 @@ +"""Cloud e2e tests for the DAG wait/suspend-resume example.""" + +import pytest +from aws_durable_execution_sdk_python.execution import InvocationStatus + +from src.dag import dag_wait_resume +from test.conftest import deserialize_operation_payload + + +@pytest.mark.example +@pytest.mark.durable_execution( + handler=dag_wait_resume.handler, + lambda_function_name="DAG Wait Resume", +) +def test_dag_wait_resume(durable_runner): + """DAG scheduling survives a real suspend/resume across a wait task.""" + with durable_runner: + result = durable_runner.run(input={}, timeout=90) + + assert result.status is InvocationStatus.SUCCEEDED + summary = deserialize_operation_payload(result.result) + + assert summary["completion_reason"] == "ALL_COMPLETED" + assert summary["counts"] == { + "success": 3, + "failure": 0, + "skipped": 0, + "total": 3, + } + tasks = summary["tasks"] + assert tasks["start"]["status"] == "SUCCEEDED" + assert tasks["start"]["result"] == "started" + assert tasks["pause"]["status"] == "SUCCEEDED" + assert tasks["finish"]["status"] == "SUCCEEDED" + assert tasks["finish"]["result"] == "resumed" diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py index 6871c5d3..4952f121 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py @@ -15,6 +15,18 @@ durable_wait_for_callback, durable_with_child_context, ) +from aws_durable_execution_sdk_python.dag import ( + DagCompletionReason, + DagConfig, + DagContext, + DagResult, + DepsMap, + SkipReason, + TaskExecution, + TaskHandle, + TaskStatus, + TriggerRule, +) # Most common exceptions - users need to handle these exceptions from aws_durable_execution_sdk_python.exceptions import ( @@ -23,6 +35,12 @@ CallbackSubmitterError, CallbackTimeoutError, ChildContextError, + DagCyclicDependencyError, + DagDuplicateTaskError, + DagExecutionError, + DagInvalidDependencyError, + DagInvalidTaskNameError, + DagPredicateError, DurableExecutionsError, DurableOperationError, ExecutionError, @@ -48,6 +66,17 @@ "CallbackSubmitterError", "CallbackTimeoutError", "ChildContextError", + "DagCompletionReason", + "DagConfig", + "DagContext", + "DagCyclicDependencyError", + "DagDuplicateTaskError", + "DagExecutionError", + "DagInvalidDependencyError", + "DagInvalidTaskNameError", + "DagPredicateError", + "DagResult", + "DepsMap", "DurableContext", "DurableExecutionsError", "DurableOperationError", @@ -55,8 +84,13 @@ "InvocationError", "InvokeError", "ParallelBranch", + "SkipReason", "StepContext", "StepError", + "TaskExecution", + "TaskHandle", + "TaskStatus", + "TriggerRule", "ValidationError", "WaitForConditionError", "WithRetryConfig", diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py index ff1ab6bc..9984d171 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py @@ -1,5 +1,6 @@ from __future__ import annotations +import hashlib import logging from contextlib import contextmanager from dataclasses import dataclass @@ -488,6 +489,74 @@ def _create_step_id(self) -> str: new_counter: int = self._step_counter.increment() return self._create_step_id_for_logical_step(new_counter) + # region DAG (experimental) + def _create_task_id(self, name: str) -> str: + """Generate a deterministic, name-based operation id for a DAG task. + + Unlike :meth:`_create_step_id`, this does NOT touch the per-context + counter, so the id is independent of run-time task-completion ordering + (which can vary across replays). The reserved ``DAG_NODE_T_`` token keeps + the *pre-image* of a DAG task id disjoint from counter-based sibling + pre-images: counter pre-images are ``{prefix}-{int}`` while a task + pre-image is ``{prefix}-DAG_NODE_T_{name}`` (uppercase + underscore, so + it can never equal a ``{prefix}-{int}`` counter pre-image). The container + prefix is this context's ``parent_id`` (the DAG container's own operation + id when the scheduler runs inside the DAG child context). + + The composed name-based pre-image is then fed through the SAME + ``blake2b(...).hexdigest()[:64]`` bounding the core applies to ordinary + counter operation ids (see :meth:`_create_step_id_for_logical_step`), so + the final backend operation id is a fixed 64-hex-char digest. This keeps + the id within the backend's 64-char ``updates[].id`` limit while + preserving determinism, injectivity (distinct pre-image => distinct + digest w.h.p.), and name-based replay stability (same ``(scope, name)`` + => same digest on every replay). Nested DAGs re-hash per level: a nested + container's digest becomes the child scope's ``parent_id`` prefix. + + .. warning:: + **Experimental.** Internal implementation detail; the id format is + subject to change without notice. + """ + prefix = self._parent_id + logical_id = ( + f"{prefix}-DAG_NODE_T_{name}" if prefix else f"DAG_NODE_T_{name}" + ) + return hashlib.blake2b(logical_id.encode()).hexdigest()[:64] + + def _run_step_with_task_id( + self, + name: str, + func: Callable[[StepContext], T], + config: StepConfig | None = None, + ) -> T: + """Run a step under a name-based (DAG) operation id. + + This is the explicit-id seam used by the DAG scheduler: it builds an + :class:`OperationIdentifier` from :meth:`_create_task_id` and drives the + step executor directly. It relies on the executor's checkpoint fast path + (keyed on the explicit operation id) for replay correctness rather than + on the per-context counter. Mirrors the order-independent child-id + pattern that ``concurrency`` already uses for map/parallel branches. + + .. warning:: + **Experimental.** Internal implementation detail. + """ + executor: StepOperationExecutor[T] = StepOperationExecutor( + func=func, + config=config or StepConfig(), + state=self.state, + operation_identifier=OperationIdentifier( + operation_id=self._create_task_id(name), + sub_type=OperationSubType.STEP, + parent_id=self._parent_id, + name=name, + ), + context_logger=self.logger, + ) + return executor.process() + + # endregion DAG (experimental) + # region replay status def is_replaying(self) -> bool: @@ -949,5 +1018,35 @@ def wait_for_condition( ) return executor.process() + def dag( + self, + register: Callable[[Any], None], + name: str | None = None, + config: Any = None, + ) -> Any: + """Declare and run a DAG of tasks, returning a ``DagResult`` synchronously. + + ``register`` is a synchronous, deterministic callback that declares tasks + on the provided ``DagContext``. The runtime then validates the graph and + schedules tasks topologically, running independent chains concurrently. + + .. warning:: + **Experimental.** This API is experimental and may be changed or + removed in future releases. First use emits a ``FutureWarning``. + """ + from aws_durable_execution_sdk_python.operation.dag import ( + dag_handler, + emit_experimental_warning_once, + ) + + emit_experimental_warning_once() + resolved_name: str | None = self._resolve_step_name(name, register) + return dag_handler( + ctx=self, + name=resolved_name, + register=register, + config=config, + ) + # endregion Operations diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dag.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dag.py new file mode 100644 index 00000000..87eb57c1 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dag.py @@ -0,0 +1,616 @@ +"""Public API for the experimental DAG operation (``context.dag()``). + +.. warning:: + **Experimental.** This API is experimental and may be changed or removed in + future releases without a major-version bump. Do not depend on it in + production until it is promoted to stable. + +This module defines the public, user-facing surface of the DAG primitive: +enums (:class:`TriggerRule`, :class:`TaskStatus`, :class:`SkipReason`, +:class:`DagCompletionReason`), the per-task result record +(:class:`TaskExecution`), configuration (:class:`DagConfig`), the registration +handle (:class:`TaskHandle`), the resolved-dependency mapping +(:class:`DepsMap`), the registration protocol (:class:`DagContext`), and the +aggregate result type (:class:`DagResult`). + +Implementation lives under ``operation/dag*.py``; this module holds only the +declarative public types so it can sit low in the dependency graph. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Mapping +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + from datetime import datetime + + from aws_durable_execution_sdk_python.config import ( + ChildConfig, + CompletionConfig, + Duration, + InvokeConfig, + MapConfig, + ParallelConfig, + StepConfig, + WaitForCallbackConfig, + ) + from aws_durable_execution_sdk_python.lambda_service import ErrorObject + from aws_durable_execution_sdk_python.serdes import SerDes + from aws_durable_execution_sdk_python.types import ( + DurableContext, + ) + from aws_durable_execution_sdk_python.waits import WaitForConditionConfig + +T = TypeVar("T") +P = TypeVar("P") +R = TypeVar("R") +U = TypeVar("U") + + +# region enums +class TriggerRule(Enum): + """Controls when a task runs based on the terminal state of its upstream deps. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + ALL_SUCCESS = "ALL_SUCCESS" # default: run only if every dep SUCCEEDED + ALL_FAILED = "ALL_FAILED" # run only if every dep FAILED (and there is >=1) + ALL_DONE = "ALL_DONE" # run once every dep is terminal, regardless of outcome + ANY_SUCCESS = "ANY_SUCCESS" # run if at least one dep SUCCEEDED + ANY_FAILED = "ANY_FAILED" # run if at least one dep FAILED + NONE_FAILED = "NONE_FAILED" # run if no dep FAILED (SUCCEEDED/SKIPPED allowed) + + +class TaskStatus(Enum): + """Terminal (or in-flight) status of a single DAG task. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + SUCCEEDED = "SUCCEEDED" + FAILED = "FAILED" + SKIPPED = "SKIPPED" + STARTED = "STARTED" + + +class SkipReason(Enum): + """Why a task was SKIPPED. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + TRIGGER_RULE = "TRIGGER_RULE" + RUN_IF_PREDICATE = "RUN_IF_PREDICATE" + + +class DagCompletionReason(Enum): + """Why the DAG stopped scheduling. + + The first three members are value-compatible with + ``concurrency.CompletionReason``; ``COMPLETED_WITH_FAILURES`` is DAG-only and + signals the default drain-on-failure path finished with >=1 failed task. + ``CUSTOM_COMPLETION_SUCCEEDED``/``CUSTOM_COMPLETION_FAILED`` are early + completion via a custom :attr:`DagCustomCompletionConfig.should_complete` + predicate rather than a threshold. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + ALL_COMPLETED = "ALL_COMPLETED" + MIN_SUCCESSFUL_REACHED = "MIN_SUCCESSFUL_REACHED" + FAILURE_TOLERANCE_EXCEEDED = "FAILURE_TOLERANCE_EXCEEDED" + COMPLETED_WITH_FAILURES = "COMPLETED_WITH_FAILURES" + CUSTOM_COMPLETION_SUCCEEDED = "CUSTOM_COMPLETION_SUCCEEDED" + CUSTOM_COMPLETION_FAILED = "CUSTOM_COMPLETION_FAILED" + + +class DagCompletionOutcome(Enum): + """The terminal disposition a custom DAG completion predicate assigns to an + early completion. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + SUCCEEDED = "SUCCEEDED" + FAILED = "FAILED" + + +@dataclass(frozen=True) +class DagCompletionDecision: + """The value a DAG custom completion predicate returns. + + Use :meth:`continue_dag` to keep scheduling, or :meth:`complete_dag` to stop + the DAG now with a given outcome (default + :attr:`DagCompletionOutcome.SUCCEEDED`). + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + complete: bool + outcome: DagCompletionOutcome | None = None + + def __post_init__(self) -> None: + if self.complete and self.outcome is None: + object.__setattr__(self, "outcome", DagCompletionOutcome.SUCCEEDED) + if not self.complete and self.outcome is not None: + msg = "outcome must be None when complete is False" + raise ValueError(msg) + + +def continue_dag() -> DagCompletionDecision: + """Returns a decision meaning "keep scheduling ready tasks".""" + return DagCompletionDecision(complete=False) + + +def complete_dag( + outcome: DagCompletionOutcome = DagCompletionOutcome.SUCCEEDED, +) -> DagCompletionDecision: + """Returns a decision meaning "complete the DAG now" with the given outcome.""" + return DagCompletionDecision(complete=True, outcome=outcome) + + +# endregion enums + + +@dataclass(frozen=True) +class TaskExecution(Generic[T]): + """Immutable record of a single task's outcome within a DAG. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + name: str + status: TaskStatus + skip_reason: SkipReason | None = None + result: T | None = None + error: ErrorObject | None = None + started_at: datetime | None = None + completed_at: datetime | None = None + + +@dataclass(frozen=True) +class DagCompletionItemStatus: + """Per-task snapshot passed to a DAG custom completion predicate. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + name: str + status: TaskStatus | None = None + """The task's full status, including ``SKIPPED``; ``None`` if not started.""" + result: Any | None = None + """Present only when ``status`` is ``TaskStatus.SUCCEEDED``.""" + skip_reason: SkipReason | None = None + """Present only when ``status`` is ``TaskStatus.SKIPPED``.""" + + +@dataclass(frozen=True) +class DagCompletionStatus: + """Progress snapshot passed to a DAG custom completion predicate. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + success_count: int + failure_count: int + skipped_count: int + completed_count: int + """``success_count + failure_count + skipped_count`` (all terminal states).""" + total_count: int + items: Sequence[DagCompletionItemStatus] + """Per-task snapshot, ordered by registration order.""" + results: Mapping[str, DagCompletionItemStatus] + """Live view of terminal task snapshots by name.""" + + +@dataclass(frozen=True) +class DagCustomCompletionConfig: + """Custom-predicate DAG completion: a deterministic predicate evaluated over + the DAG's live progress and task results after every task settlement. + + Unlike threshold-based :class:`~...config.CompletionConfig`, this predicate + can inspect individual tasks' results (via :attr:`DagCompletionStatus.items` + / :attr:`DagCompletionStatus.results`), not just aggregate counts. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + should_complete: Callable[[DagCompletionStatus], DagCompletionDecision] + + +@dataclass(frozen=True) +class DagConfig: + """Configuration for a DAG. + + ``completion_config`` accepts either the base SDK's threshold-only + :class:`~...config.CompletionConfig`, reused verbatim, or a + :class:`DagCustomCompletionConfig` for a results-aware custom predicate. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + max_concurrency: int | None = None + """Maximum number of top-level tasks this DAG's scheduler runs concurrently. + + When ``None`` (unset), the scheduler applies a default cap of + ``DEFAULT_DAG_MAX_CONCURRENCY`` (40) rather than running unbounded. An + explicit value always wins, including a value above 40; the only validation + is that it must be ``>= 1``. + + Scope: this bounds the DAG SCHEDULER ONLY -- the top-level tasks of *this* + DAG. It is deliberately NOT inherited by a task's own internal fan-out: a + ``map`` or ``parallel`` task keeps its own default (unlimited) unless + configured on that task, and a nested ``dag`` task gets its own independent + default of 40. + """ + completion_config: CompletionConfig | DagCustomCompletionConfig | None = None + default_trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS + serdes: SerDes | None = None + + +# region handle & deps-map +@dataclass(eq=False) +class TaskHandle(Generic[T]): + """Registration-time reference to a task, plus a small chaining builder. + + Never serialized. Identity is the task ``name`` (unique within a DAG scope). + ``__hash__`` is keyed on the name so handles can live in ``set``\\ s (e.g. the + dependency edge set); note this does NOT make ``deps[handle]`` resolve against + a name-keyed dict — :meth:`DepsMap.__getitem__` dispatches on the handle type + explicitly. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + _name: str + _dag: Any # back-ref to DagContextImpl; typed Any to avoid a circular import + + def __hash__(self) -> int: + return hash(self._name) + + @property + def name(self) -> str: + """The resolved task name.""" + return self._name + + def after(self, *deps: TaskHandle[Any]) -> TaskHandle[T]: + """Add ordering-only dependencies (wait for them, but do not receive + their results in the ``DepsMap``). Returns ``self`` for chaining. + + .. warning:: + **Experimental.** + """ + self._dag._register_after(self, deps) + return self + + def trigger_rule(self, rule: TriggerRule) -> TaskHandle[T]: + """Override this task's trigger rule. Returns ``self`` for chaining. + + .. warning:: + **Experimental.** + """ + self._dag._register_trigger_rule(self, rule) + return self + + +class DepsMap(Mapping[str, Any]): + """Resolved upstream results, keyed by dependency task name. + + Access by string name (``deps["fetch"] -> Any``) or, for static typing, by + the originating :class:`TaskHandle` (``deps[handle] -> T | None``). The handle + path dispatches at runtime on ``isinstance(key, TaskHandle)`` and extracts the + name; it does not rely on hashing. + + The handle overload returns ``T | None`` (not bare ``T``): a dependency's + result is only present when that upstream task SUCCEEDED. Under a + non-``ALL_SUCCESS`` trigger rule (e.g. ``ALL_DONE``/``ANY_FAILED``) a task + body can legitimately run while an upstream dep FAILED or was SKIPPED, in + which case its value here is ``None`` -- the type reflects that + long-standing runtime behavior rather than pretending the value is always + present. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + def __init__(self, by_name: dict[str, Any]) -> None: + self._by_name = by_name + + @overload + def __getitem__(self, key: TaskHandle[T]) -> T | None: ... + @overload + def __getitem__(self, key: str) -> Any: ... + def __getitem__(self, key: str | TaskHandle[Any]) -> Any: + name = key._name if isinstance(key, TaskHandle) else key + return self._by_name[name] + + def __iter__(self): + return iter(self._by_name) + + def __len__(self) -> int: + return len(self._by_name) + + def __contains__(self, key: object) -> bool: + name = key._name if isinstance(key, TaskHandle) else key + return name in self._by_name + + +# endregion handle & deps-map + + +# Type alias: dependencies are declared as a list of TaskHandles (or None). +DepsArg = "Sequence[TaskHandle[Any]] | None" + + +# region DagContext protocol +class DagContext(ABC): + """Declarative task-registration surface passed to a DAG ``register`` callback. + + Each method registers exactly one task and returns a :class:`TaskHandle`. + Task bodies always receive the resolved :class:`DepsMap` as their first + positional argument (empty for root tasks). + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + @abstractmethod + def step( + self, + func: Callable[..., T], + deps: Sequence[TaskHandle[Any]] | None = None, + name: str | None = None, + config: StepConfig | None = None, + *, + trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS, + run_if: Callable[[DepsMap], bool] | None = None, + ) -> TaskHandle[T]: + """Register a step task. Body signature: ``(deps, step_ctx)``. + + .. warning:: + **Experimental.** + """ + ... # pragma: no cover + + @abstractmethod + def invoke( + self, + function_name: str, + payload_fn: Callable[[DepsMap], Any] | Any, + deps: Sequence[TaskHandle[Any]] | None = None, + name: str | None = None, + config: InvokeConfig[Any, Any] | None = None, + *, + trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS, + run_if: Callable[[DepsMap], bool] | None = None, + ) -> TaskHandle[Any]: + """Register an invoke task. ``payload_fn`` is a plain value or a + ``(deps) -> payload`` callable materialized at scheduling time. + + .. warning:: + **Experimental.** + """ + ... # pragma: no cover + + @abstractmethod + def wait_for_callback( + self, + submitter: Callable[..., None], + deps: Sequence[TaskHandle[Any]] | None = None, + name: str | None = None, + config: WaitForCallbackConfig | None = None, + *, + trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS, + run_if: Callable[[DepsMap], bool] | None = None, + ) -> TaskHandle[Any]: + """Register a wait-for-callback task. Body signature: ``(deps, callback_id, ctx)``. + + .. warning:: + **Experimental.** + """ + ... # pragma: no cover + + @abstractmethod + def wait( + self, + duration: Duration, + deps: Sequence[TaskHandle[Any]] | None = None, + name: str | None = None, + *, + trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS, + run_if: Callable[[DepsMap], bool] | None = None, + ) -> TaskHandle[None]: + """Register a wait task. + + .. warning:: + **Experimental.** + """ + ... # pragma: no cover + + @abstractmethod + def wait_for_condition( + self, + check: Callable[..., T], + config: WaitForConditionConfig[T], + deps: Sequence[TaskHandle[Any]] | None = None, + name: str | None = None, + *, + trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS, + run_if: Callable[[DepsMap], bool] | None = None, + ) -> TaskHandle[T]: + """Register a wait-for-condition task. Body signature: ``(deps, state, ctx)``. + + .. warning:: + **Experimental.** + """ + ... # pragma: no cover + + @abstractmethod + def run_in_child_context( + self, + func: Callable[..., T], + deps: Sequence[TaskHandle[Any]] | None = None, + name: str | None = None, + config: ChildConfig | None = None, + *, + trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS, + run_if: Callable[[DepsMap], bool] | None = None, + ) -> TaskHandle[T]: + """Register a child-context task. Body signature: ``(deps, child_ctx)``. + + .. warning:: + **Experimental.** + """ + ... # pragma: no cover + + @abstractmethod + def map( + self, + inputs: Sequence[U] | Callable[[DepsMap], Sequence[U]], + func: Callable[..., T], + deps: Sequence[TaskHandle[Any]] | None = None, + name: str | None = None, + config: MapConfig | None = None, + *, + trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS, + run_if: Callable[[DepsMap], bool] | None = None, + ) -> TaskHandle[Any]: + """Register a map task. ``inputs`` may be a sequence or ``(deps) -> sequence``. + + .. warning:: + **Experimental.** + """ + ... # pragma: no cover + + @abstractmethod + def parallel( + self, + functions: Sequence[Callable[[DurableContext], Any]], + deps: Sequence[TaskHandle[Any]] | None = None, + name: str | None = None, + config: ParallelConfig | None = None, + *, + trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS, + run_if: Callable[[DepsMap], bool] | None = None, + ) -> TaskHandle[Any]: + """Register a parallel task. + + .. warning:: + **Experimental.** + """ + ... # pragma: no cover + + @abstractmethod + def dag( + self, + register: Callable[[DagContext], None], + deps: Sequence[TaskHandle[Any]] | None = None, + name: str | None = None, + config: DagConfig | None = None, + *, + trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS, + run_if: Callable[[DepsMap], bool] | None = None, + ) -> TaskHandle[Any]: + """Register a nested DAG task. + + .. warning:: + **Experimental.** + """ + ... # pragma: no cover + + +# endregion DagContext protocol + + +# region DagResult surface +class DagResult(ABC): + """Aggregate result of a DAG run. Concrete impl lives in + ``operation/dag_result.py``. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + @overload + def get_result(self, task: TaskHandle[T]) -> T | None: ... # pragma: no cover + @overload + def get_result(self, task: str) -> Any: ... # pragma: no cover + @abstractmethod + def get_result(self, task: str | TaskHandle[Any]) -> Any: ... # pragma: no cover + + @abstractmethod + def get_status( + self, task: str | TaskHandle[Any] + ) -> TaskStatus | None: ... # pragma: no cover + + @abstractmethod + def succeeded(self) -> list[TaskExecution]: ... # pragma: no cover + + @abstractmethod + def failed(self) -> list[TaskExecution]: ... # pragma: no cover + + @abstractmethod + def skipped(self) -> list[TaskExecution]: ... # pragma: no cover + + @property + @abstractmethod + def results(self) -> Mapping[str, TaskExecution]: ... # pragma: no cover + + @property + @abstractmethod + def success_count(self) -> int: ... # pragma: no cover + + @property + @abstractmethod + def failure_count(self) -> int: ... # pragma: no cover + + @property + @abstractmethod + def skipped_count(self) -> int: ... # pragma: no cover + + @property + @abstractmethod + def total_count(self) -> int: ... # pragma: no cover + + @property + @abstractmethod + def completion_reason(self) -> DagCompletionReason: ... # pragma: no cover + + @abstractmethod + def throw_if_error(self) -> None: ... # pragma: no cover + + +# endregion DagResult surface diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py index 46c23356..f9b28a11 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py @@ -256,6 +256,85 @@ class InvalidStateError(DurableExecutionsError): """Raised when an operation is attempted on an object in an invalid state.""" +class DagExecutionError(DurableExecutionsError): + """Raised by ``DagResult.throw_if_error()`` when one or more DAG tasks FAILED. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + +class DagCyclicDependencyError(ValidationError): + """Raised when the DAG contains a dependency cycle. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + +class DagInvalidTaskNameError(ValidationError): + """Raised when a task name is empty, unresolvable, too long, or uses a + reserved / disallowed character. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + +class DagDuplicateTaskError(ValidationError): + """Raised when two tasks in the same DAG scope share a name. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + +class DagInvalidDependencyError(ValidationError): + """Raised when a task depends on a handle not registered in this DAG scope. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + +class DagPredicateError(DurableExecutionsError): + """Raised when a task's ``run_if`` predicate raises, aborting the DAG. + + ``run_if`` is specified as a synchronous, deterministic, pure predicate over + resolved upstream results; it is re-evaluated on every replay and is not a + checkpointed operation. A predicate that raises is therefore a *defect in + deterministic code*, not a business outcome. Rather than record the task as + ``FAILED`` (which would silently drive every downstream ``ALL_FAILED`` / + ``ANY_FAILED`` / ``ALL_DONE`` compensation path) or ``SKIPPED``, the + scheduler aborts: the offending task gets no terminal state, no further + tasks start, and the ``dag()`` operation fails with this error. + + The offending task name is available as :attr:`task_name` and is also named + in the message. Where it is raised (the scheduler), the original exception + is preserved as ``__cause__`` (``raise ... from e``) and ``task_name`` is + set. Across the durable ``dag()`` child-context boundary the error is rebuilt + from serialized fields (type name + message) so the first run and replay are + identical; through that boundary ``__cause__`` and ``task_name`` are not + reconstructed, but the offending task name remains in the message. This is + the same boundary behaviour as the rest of the ``Dag*`` error family. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + def __init__( + self, message: str | None = None, task_name: str | None = None + ) -> None: + super().__init__(message) + self.task_name: str | None = task_name + + class DurableOperationError(DurableExecutionsError): """Base class for typed, per-operation failures. diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py index 1db42ec3..9da2ad70 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py @@ -92,6 +92,7 @@ def from_sub_type(cls, sub_type: OperationSubType) -> OperationType: | OperationSubType.MAP_ITERATION | OperationSubType.PARALLEL | OperationSubType.PARALLEL_BRANCH + | OperationSubType.DAG ): return OperationType.CONTEXT case _: @@ -115,6 +116,7 @@ class OperationSubType(Enum): WAIT_FOR_CALLBACK = "WaitForCallback" WAIT_FOR_CONDITION = "WaitForCondition" CHAINED_INVOKE = "ChainedInvoke" + DAG = "Dag" class InvocationStatus(Enum): diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag.py new file mode 100644 index 00000000..13a16fd2 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag.py @@ -0,0 +1,376 @@ +"""dag_handler: wraps register + validate + schedule inside a DAG container. + +The DAG container is a child-context operation whose result is the converged +cross-language envelope (see ``dag_result.DagResultImpl.to_dict``). This module +owns the two behaviours the generic child-context executor cannot express: + +* the **degradation ladder** that serializes the envelope and, if it exceeds the + checkpoint size limit, drops ``tasks`` (setting ``ReplayChildren``) and then + ``failedTaskNames`` -- never the counts, ``completionReason`` or + ``startedTaskNames``; and +* the **reconstruct** replay strategy for the offloaded case, which re-runs the + deterministic register graph (each task fast-paths from its retained child + checkpoint) and seeds the STARTED set from the envelope so an in-flight task + is never restarted. + +.. warning:: + **Experimental.** Internal wiring for ``context.dag()``. +""" + +from __future__ import annotations + +import json +import warnings +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from aws_durable_execution_sdk_python.constants import CHECKPOINT_SIZE_LIMIT_BYTES +from aws_durable_execution_sdk_python.dag import DagCompletionReason, DagConfig, DagResult +from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + DagCyclicDependencyError, + DagDuplicateTaskError, + DagExecutionError, + DagInvalidDependencyError, + DagInvalidTaskNameError, + DagPredicateError, + InvocationError, + SuspendExecution, + ValidationError, +) +from aws_durable_execution_sdk_python.identifier import OperationIdentifier +from aws_durable_execution_sdk_python.lambda_service import ( + ContextOptions, + ErrorObject, + OperationSubType, + OperationUpdate, +) +from aws_durable_execution_sdk_python.operation.base import CheckResult, OperationExecutor +from aws_durable_execution_sdk_python.operation.dag_context import DagContextImpl +from aws_durable_execution_sdk_python.operation.dag_executor import DagExecutor +from aws_durable_execution_sdk_python.operation.dag_result import DagResultImpl +from aws_durable_execution_sdk_python.operation.dag_validator import validate_dag + +if TYPE_CHECKING: + from collections.abc import Callable + + from aws_durable_execution_sdk_python.context import DurableContext + from aws_durable_execution_sdk_python.dag import DagContext + from aws_durable_execution_sdk_python.state import CheckpointedResult, ExecutionState + +# Typed Dag* errors that ``unwrap_dag_error`` surfaces cleanly through the child +# context boundary (on the first run via ``__cause__`` and on replay via +# ``error_type``). Not all are validation errors: ``DagExecutionError`` and +# ``DagPredicateError`` are execution-time. +_DAG_VALIDATION_ERRORS = ( + DagCyclicDependencyError, + DagInvalidTaskNameError, + DagDuplicateTaskError, + DagInvalidDependencyError, + DagExecutionError, + DagPredicateError, +) + +_DAG_ERROR_BY_NAME = {cls.__name__: cls for cls in _DAG_VALIDATION_ERRORS} + +_warned = False + + +def emit_experimental_warning_once() -> None: + """Emit a one-time ``FutureWarning`` on first use of ``context.dag()``.""" + global _warned + if not _warned: + _warned = True + warnings.warn( + "context.dag() is an EXPERIMENTAL API and may change or be removed " + "in a future release without a major-version bump.", + FutureWarning, + stacklevel=3, + ) + + +def _check_max_concurrency(config: DagConfig) -> None: + if config.max_concurrency is not None and config.max_concurrency <= 0: + msg = f"Invalid max_concurrency: {config.max_concurrency}" + raise ValidationError(msg) + + +@dataclass(frozen=True) +class _ReconstructInfo: + """The subset of the offloaded envelope the reconstruct path consumes. + + ``started_task_names`` seeds the STARTED set (never restart an in-flight + task); ``completion_reason``/``total_count`` are taken as authoritative + rather than re-derived from the fast-pathed results. + """ + + started_task_names: set[str] + completion_reason: DagCompletionReason | None + total_count: int | None + + +def _run_dag_body( + dag_child_ctx: DurableContext, + register: Callable[[DagContext], None], + config: DagConfig, + reconstruct: _ReconstructInfo | None = None, +) -> DagResult: + dag_ctx = DagContextImpl(dag_child_ctx, config) + register(dag_ctx) + validate_dag(dag_ctx) + executor = DagExecutor(dag_child_ctx, dag_ctx.get_tasks(), config) + if reconstruct is None: + return executor.run() + return executor.run( + reconstruct_started=reconstruct.started_task_names, + reconstruct_reason=reconstruct.completion_reason, + reconstruct_total=reconstruct.total_count, + ) + + +class DagContainerExecutor(OperationExecutor[DagResult]): + """Checkpoint orchestration for the DAG container operation. + + Mirrors ``ChildOperationExecutor``'s START/SUCCEED/FAIL contract (so a nested + DAG's error unwraps identically and replay reconstruction of a FAILED + container is unchanged), but replaces the generic size branch with the + DAG degradation ladder and the generic ``ReplayChildren`` re-execute with the + DAG reconstruct strategy. + """ + + def __init__( + self, + *, + run_body: Callable[[_ReconstructInfo | None], DagResult], + state: ExecutionState, + operation_identifier: OperationIdentifier, + ) -> None: + self._run_body = run_body + self.state = state + self.operation_identifier = operation_identifier + self.sub_type = OperationSubType.DAG + + def check_result_status(self) -> CheckResult[DagResult]: + cr: CheckpointedResult = self.state.get_checkpoint_result( + self.operation_identifier.operation_id + ) + + # Terminal success, tasks present (not offloaded): deserialize the full + # envelope and return. Do not read children, do not re-run the body. + if cr.is_succeeded() and not cr.is_replay_children(): + return CheckResult.create_completed(_deserialize_inline(cr.result)) + + # Terminal success, offloaded (ReplayChildren): reconstruct from the + # retained child checkpoints plus the envelope. + if cr.is_succeeded() and cr.is_replay_children(): + return CheckResult.create_is_ready_to_execute(cr) + + # Terminal failure: surface as ChildContextError (unwrap_dag_error maps it + # back to the typed Dag* error), identical on first run and replay. + if cr.is_failed(): + cr.raise_operation_error(ChildContextError) + + # Create the START checkpoint if the container has not started. Fire and + # forget (is_sync=False), matching the child-context executor. + if not cr.is_existent(): + start = OperationUpdate.create_context_start( + identifier=self.operation_identifier, sub_type=self.sub_type + ) + self.state.create_checkpoint(operation_update=start, is_sync=False) + + return CheckResult.create_is_ready_to_execute(cr) + + def execute(self, checkpointed_result: CheckpointedResult) -> DagResult: + reconstruct: _ReconstructInfo | None = None + if checkpointed_result.is_succeeded() and ( + checkpointed_result.is_replay_children() + ): + reconstruct = _reconstruct_info(checkpointed_result.result) + + try: + result = self._run_body(reconstruct) + except SuspendExecution: + # The DAG suspended (a task is waiting): bubble without checkpointing. + raise + except Exception as e: # noqa: BLE001 + # Retryable InvocationError: re-raise with no FAIL checkpoint so the + # backend retry re-runs. Everything else is terminal. + if isinstance(e, InvocationError) and e.is_retryable(): + raise + error_object = ErrorObject.from_exception(e) + fail = OperationUpdate.create_context_fail( + identifier=self.operation_identifier, + error=error_object, + sub_type=self.sub_type, + ) + self.state.create_checkpoint(operation_update=fail) + error_object.raise_as_operation_error(ChildContextError) + + if reconstruct is not None: + # Offloaded reconstruct: the container is already SUCCEEDED with the + # offloaded envelope; do not re-checkpoint. + return result + + self._checkpoint_with_ladder(result) # type: ignore[arg-type] + return result + + def _checkpoint_with_ladder(self, result: DagResultImpl) -> None: + """Serialize the envelope and degrade until it fits, in the exact + contract order. + + 1. Full envelope with ``tasks`` -- checkpoint, no ReplayChildren. + 2. Too large: drop ``tasks``, set ReplayChildren so the backend retains + the child operations that hold the per-task results. + 3. Still too large: drop ``failedTaskNames``. + + Counts, ``completionReason`` and ``startedTaskNames`` are never dropped: + a DAG must never fail to checkpoint because its own summary did not fit, + and ``startedTaskNames`` (bounded by ``max_concurrency``) is what replay + needs to avoid restarting an in-flight task. + """ + envelope = result.to_dict() + payload = json.dumps(envelope) + replay_children = False + if _too_large(payload): + envelope.pop("tasks", None) + payload = json.dumps(envelope) + replay_children = True + if _too_large(payload): + envelope.pop("failedTaskNames", None) + payload = json.dumps(envelope) + + succeed = OperationUpdate.create_context_succeed( + identifier=self.operation_identifier, + payload=payload, + sub_type=self.sub_type, + context_options=ContextOptions(replay_children=replay_children), + ) + self.state.create_checkpoint(operation_update=succeed) + + +def _too_large(payload: str) -> bool: + return len(payload.encode("utf-8")) > CHECKPOINT_SIZE_LIMIT_BYTES + + +def _deserialize_inline(payload: str | None) -> DagResult: + if not payload: + # Defensive: a succeeded, non-offloaded container always carries the + # envelope; an empty payload can only mean an empty DAG round-trip. + return DagResultImpl({}, DagCompletionReason.ALL_COMPLETED) + return DagResultImpl.from_dict(json.loads(payload)) + + +def _reconstruct_info(payload: str | None) -> _ReconstructInfo: + data = json.loads(payload) if payload else {} + reason_value = data.get("completionReason") + return _ReconstructInfo( + started_task_names=set(data.get("startedTaskNames") or []), + completion_reason=( + DagCompletionReason(reason_value) if reason_value else None + ), + total_count=data.get("totalCount"), + ) + + +def unwrap_dag_error(exc: ChildContextError) -> None: + """Re-raise the typed Dag* cause of a wrapped ``ChildContextError``. + + ``DagContainerExecutor`` surfaces body exceptions as ``ChildContextError`` + with the original on ``__cause__`` (first run) and ``error_type`` set to the + original class name (both first run and replay, reconstructed from the + checkpoint). This restores the clean typed throw for DAG validation / + execution errors, mirroring the ``wait_for_callback`` precedent. + + On **replay** the failure is rebuilt from a checkpoint, which sets + ``error_type`` (the original class name) but leaves ``__cause__`` as + ``None``. In that case we reconstruct the typed Dag* error from + ``error_type`` so a nested DAG's error surfaces identically on the first run + and on replay. If neither path identifies a Dag* error, re-raise the + original wrapper unchanged. + """ + cause = exc.__cause__ + if isinstance(cause, _DAG_VALIDATION_ERRORS): + # Re-raise the typed error, preserving ITS OWN original cause so a + # DagPredicateError still exposes the raising predicate's exception as + # __cause__ (contract: the original error must remain the retrievable + # cause). ``from inner`` also suppresses the ChildContextError wrapper + # from the traceback; when inner is None (the validation errors) this is + # exactly the previous ``raise cause from None`` behaviour. + inner = cause.__cause__ + raise cause from inner + dag_cls = _DAG_ERROR_BY_NAME.get(exc.error_type or "") + if dag_cls is not None: + raise dag_cls(exc.message) from None + raise exc + + +def dag_handler( + ctx: DurableContext, + name: str | None, + register: Callable[[DagContext], None], + config: DagConfig | None, +) -> DagResult: + """Run a top-level DAG as a container child context and return its DagResult. + + The container takes a counter-based operation id (like + ``run_in_child_context``), so ``_replay_aware`` bookkeeping is identical. + """ + config = config or DagConfig() + _check_max_concurrency(config) + + with ctx._replay_aware(): + operation_id = ctx._create_step_id() + identifier = OperationIdentifier( + operation_id=operation_id, + sub_type=OperationSubType.DAG, + parent_id=ctx._parent_id, + name=name, + ) + + def run_body(reconstruct: _ReconstructInfo | None) -> DagResult: + child = ctx.create_child_context(operation_id=operation_id) + return _run_dag_body(child, register, config, reconstruct) + + executor = DagContainerExecutor( + run_body=run_body, + state=ctx.state, + operation_identifier=identifier, + ) + try: + return executor.process() + except ChildContextError as e: + unwrap_dag_error(e) + raise # pragma: no cover - unwrap_dag_error always raises + + +def run_nested_dag( + ctx: DurableContext, + name: str, + register: Callable[[DagContext], None], + config: DagConfig | None, +) -> DagResult: + """Run a nested DAG task under a name-based (``DAG_NODE_T_``) container id.""" + config = config or DagConfig() + _check_max_concurrency(config) + task_id = ctx._create_task_id(name) + identifier = OperationIdentifier( + operation_id=task_id, + sub_type=OperationSubType.DAG, + parent_id=ctx._parent_id, + name=name, + ) + + def run_body(reconstruct: _ReconstructInfo | None) -> DagResult: + child = ctx.create_child_context(operation_id=task_id) + return _run_dag_body(child, register, config, reconstruct) + + executor = DagContainerExecutor( + run_body=run_body, + state=ctx.state, + operation_identifier=identifier, + ) + try: + return executor.process() + except ChildContextError as e: + unwrap_dag_error(e) + raise # pragma: no cover - unwrap_dag_error always raises diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_context.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_context.py new file mode 100644 index 00000000..67057834 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_context.py @@ -0,0 +1,426 @@ +"""DagContextImpl: registers TaskDefs and hands back TaskHandles. + +.. warning:: + **Experimental.** Internal implementation of the DAG registration phase. +""" + +from __future__ import annotations + +import json +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from aws_durable_execution_sdk_python.concurrency.models import BatchResult +from aws_durable_execution_sdk_python.config import ChildConfig +from aws_durable_execution_sdk_python.dag import ( + DagContext, + DepsMap, + TaskHandle, + TriggerRule, +) +from aws_durable_execution_sdk_python.exceptions import ( + DagInvalidTaskNameError, + ValidationError, +) +from aws_durable_execution_sdk_python.identifier import ( + OperationIdentifier, + OperationIdNamespace, +) +from aws_durable_execution_sdk_python.lambda_service import OperationSubType +from aws_durable_execution_sdk_python.operation.child import child_handler +from aws_durable_execution_sdk_python.operation.invoke import InvokeOperationExecutor +from aws_durable_execution_sdk_python.operation.map import map_handler +from aws_durable_execution_sdk_python.operation.parallel import parallel_handler +from aws_durable_execution_sdk_python.operation.wait import WaitOperationExecutor +from aws_durable_execution_sdk_python.operation.wait_for_condition import ( + WaitForConditionOperationExecutor, +) +from aws_durable_execution_sdk_python.serdes import SerDes + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from aws_durable_execution_sdk_python.context import DurableContext + from aws_durable_execution_sdk_python.dag import DagConfig + from aws_durable_execution_sdk_python.serdes import SerDesContext + +logger = logging.getLogger(__name__) + + +class _BatchResultSerDes(SerDes): + """Serialize a map/parallel task's ``BatchResult`` container payload.""" + + def serialize(self, value: BatchResult, serdes_context: SerDesContext) -> str: + return json.dumps(value.to_dict()) + + def deserialize(self, data: str, serdes_context: SerDesContext) -> BatchResult: + return BatchResult.from_dict(json.loads(data)) + + +@dataclass +class TaskDef: + """Internal record for one registered DAG task. + + ``inline_deps`` drive the :class:`DepsMap`; ``all_deps`` (inline ∪ + ``.after`` edges) drive readiness / trigger-rule / cycle checks. ``executor`` + binds the name-based (explicit-id) runner for this task's operation kind. + """ + + name: str + kind: str + inline_deps: list[TaskHandle[Any]] + all_deps: list[TaskHandle[Any]] + trigger_rule: TriggerRule + run_if: Callable[[DepsMap], bool] | None + config: Any + executor: Callable[[DurableContext, DepsMap], Any] + + +def _resolve_name(name: str | None, func: Any) -> str: + resolved = name if name else getattr(func, "_original_name", None) + if not resolved: + msg = ( + "Could not resolve a task name. Pass an explicit `name=` (bare " + "lambdas have no resolvable name)." + ) + raise DagInvalidTaskNameError(msg) + return resolved + + +class DagContextImpl(DagContext): + """Concrete DagContext used during the registration phase.""" + + def __init__(self, ctx: DurableContext, config: DagConfig) -> None: + self._ctx = ctx + self._config = config + self._tasks: dict[str, TaskDef] = {} + # ordered record of every registration (incl. duplicate names) so the + # validator can detect duplicates; the dict is for name lookup. + self._registration_order: list[TaskDef] = [] + + def get_tasks(self) -> dict[str, TaskDef]: + """Return the registered tasks by name (last registration wins).""" + return self._tasks + + def get_registration_order(self) -> list[TaskDef]: + """Return every registered task in order, including duplicate names.""" + return self._registration_order + + # region builder mutations (called by TaskHandle) + def _register_after( + self, handle: TaskHandle[Any], deps: Sequence[TaskHandle[Any]] + ) -> None: + task = self._tasks[handle.name] + for d in deps: + if d not in task.all_deps: + task.all_deps.append(d) + + def _register_trigger_rule( + self, handle: TaskHandle[Any], rule: TriggerRule + ) -> None: + self._tasks[handle.name].trigger_rule = rule + + # endregion builder mutations + + def _add( + self, + name: str, + kind: str, + inline_deps: Sequence[TaskHandle[Any]] | None, + trigger_rule: TriggerRule, + run_if: Callable[[DepsMap], bool] | None, + config: Any, + executor: Callable[[DurableContext, DepsMap], Any], + ) -> TaskHandle[Any]: + inline = list(inline_deps) if inline_deps else [] + # A per-method `trigger_rule=None` sentinel means "not explicitly set", + # so it falls back to the DAG-wide default from DagConfig. + resolved_trigger = ( + trigger_rule + if trigger_rule is not None + else self._config.default_trigger_rule + ) + task = TaskDef( + name=name, + kind=kind, + inline_deps=inline, + all_deps=list(inline), + trigger_rule=resolved_trigger, + run_if=run_if, + config=config, + executor=executor, + ) + self._registration_order.append(task) + self._tasks[name] = task + return TaskHandle(_name=name, _dag=self) + + # region task kinds + def step( + self, func, deps=None, name=None, config=None, *, + trigger_rule=None, run_if=None, + ): + task_name = _resolve_name(name, func) + cfg = config + + def executor(ctx: DurableContext, deps_map: DepsMap): + return ctx._run_step_with_task_id( + task_name, lambda step_ctx: func(deps_map, step_ctx), cfg + ) + + return self._add(task_name, "step", deps, trigger_rule, run_if, cfg, executor) + + def invoke( + self, function_name, payload_fn, deps=None, name=None, config=None, *, + trigger_rule=None, run_if=None, + ): + task_name = _resolve_name(name, payload_fn) + + def executor(ctx: DurableContext, deps_map: DepsMap): + from aws_durable_execution_sdk_python.config import InvokeConfig + + payload = payload_fn(deps_map) if callable(payload_fn) else payload_fn + return InvokeOperationExecutor( + function_name=function_name, + payload=payload, + state=ctx.state, + operation_identifier=OperationIdentifier( + operation_id=ctx._create_task_id(task_name), + sub_type=OperationSubType.CHAINED_INVOKE, + parent_id=ctx._parent_id, + name=task_name, + ), + config=config or InvokeConfig(), + ).process() + + return self._add( + task_name, "invoke", deps, trigger_rule, run_if, config, executor + ) + + def wait_for_callback( + self, submitter, deps=None, name=None, config=None, *, + trigger_rule=None, run_if=None, + ): + task_name = _resolve_name(name, submitter) + + def executor(ctx: DurableContext, deps_map: DepsMap): + # A DAG callback task materializes as a Callback *container* context: + # a callback operation cannot take an explicit (name-based) operation + # id directly, so the task wraps the SDK's native wait_for_callback in + # a container that carries the task's name-based id. The native op + # then creates the inner WaitForCallback child context, the callback, + # and the submitter step. The resulting on-the-wire shape is + # Callback(container) -> WaitForCallback -> {CallbackStarted, submitter} + # matching the frozen DAG contract (and the JS reference). + def body(child: DurableContext): + return child.wait_for_callback( + lambda cb_id, cb_ctx: submitter(deps_map, cb_id, cb_ctx), + name=task_name, + config=config, + ) + + return self._run_child( + ctx, + task_name, + body, + ChildConfig(sub_type=OperationSubType.CALLBACK), + OperationSubType.CALLBACK, + ) + + return self._add( + task_name, "wait_for_callback", deps, trigger_rule, run_if, config, executor + ) + + def wait( + self, duration, deps=None, name=None, *, + trigger_rule=None, run_if=None, + ): + if not name: + msg = "wait tasks require an explicit `name=`." + raise DagInvalidTaskNameError(msg) + task_name = name + + seconds = duration.to_seconds() + if seconds < 1: + msg = "duration must be at least 1 second" + raise ValidationError(msg) + + def executor(ctx: DurableContext, deps_map: DepsMap): + return WaitOperationExecutor( + seconds=seconds, + state=ctx.state, + operation_identifier=OperationIdentifier( + operation_id=ctx._create_task_id(task_name), + sub_type=OperationSubType.WAIT, + parent_id=ctx._parent_id, + name=task_name, + ), + ).process() + + return self._add(task_name, "wait", deps, trigger_rule, run_if, None, executor) + + def wait_for_condition( + self, check, config, deps=None, name=None, *, + trigger_rule=None, run_if=None, + ): + task_name = _resolve_name(name, check) + + def executor(ctx: DurableContext, deps_map: DepsMap): + return WaitForConditionOperationExecutor( + check=lambda state, cctx: check(deps_map, state, cctx), + config=config, + state=ctx.state, + operation_identifier=OperationIdentifier( + operation_id=ctx._create_task_id(task_name), + sub_type=OperationSubType.WAIT_FOR_CONDITION, + parent_id=ctx._parent_id, + name=task_name, + ), + context_logger=ctx.logger, + ).process() + + return self._add( + task_name, "wait_for_condition", deps, trigger_rule, run_if, config, executor + ) + + def run_in_child_context( + self, func, deps=None, name=None, config=None, *, + trigger_rule=None, run_if=None, + ): + task_name = _resolve_name(name, func) + cfg = config or ChildConfig() + + def executor(ctx: DurableContext, deps_map: DepsMap): + return self._run_child( + ctx, + task_name, + lambda child: func(deps_map, child), + cfg, + OperationSubType.RUN_IN_CHILD_CONTEXT, + ) + + return self._add( + task_name, "child", deps, trigger_rule, run_if, cfg, executor + ) + + def map( + self, inputs, func, deps=None, name=None, config=None, *, + trigger_rule=None, run_if=None, + ): + task_name = _resolve_name(name, func) + + def executor(ctx: DurableContext, deps_map: DepsMap): + resolved = inputs(deps_map) if callable(inputs) else inputs + serdes = config.serdes if (config and config.serdes) else _BatchResultSerDes() + task_id = ctx._create_task_id(task_name) + operation_identifier = OperationIdentifier( + operation_id=task_id, + sub_type=OperationSubType.MAP, + parent_id=ctx._parent_id, + name=task_name, + ) + map_context = ctx.create_child_context(operation_id=task_id) + + def body(): + return map_handler( + items=resolved, + func=func, + config=config, + execution_state=ctx.state, + map_context=map_context, + operation_identifier=operation_identifier, + operation_id_namespace=OperationIdNamespace(task_id), + ) + + return child_handler( + func=body, + state=ctx.state, + operation_identifier=operation_identifier, + config=ChildConfig(sub_type=OperationSubType.MAP, serdes=serdes), + ) + + return self._add(task_name, "map", deps, trigger_rule, run_if, config, executor) + + def parallel( + self, functions, deps=None, name=None, config=None, *, + trigger_rule=None, run_if=None, + ): + if not name: + msg = "parallel tasks require an explicit `name=`." + raise DagInvalidTaskNameError(msg) + task_name = name + + def executor(ctx: DurableContext, deps_map: DepsMap): + serdes = config.serdes if (config and config.serdes) else _BatchResultSerDes() + task_id = ctx._create_task_id(task_name) + operation_identifier = OperationIdentifier( + operation_id=task_id, + sub_type=OperationSubType.PARALLEL, + parent_id=ctx._parent_id, + name=task_name, + ) + parallel_context = ctx.create_child_context(operation_id=task_id) + + def body(): + return parallel_handler( + callables=functions, + config=config, + execution_state=ctx.state, + parallel_context=parallel_context, + operation_identifier=operation_identifier, + operation_id_namespace=OperationIdNamespace(task_id), + ) + + return child_handler( + func=body, + state=ctx.state, + operation_identifier=operation_identifier, + config=ChildConfig(sub_type=OperationSubType.PARALLEL, serdes=serdes), + ) + + return self._add( + task_name, "parallel", deps, trigger_rule, run_if, config, executor + ) + + def dag( + self, register, deps=None, name=None, config=None, *, + trigger_rule=None, run_if=None, + ): + task_name = _resolve_name(name, register) + + def executor(ctx: DurableContext, deps_map: DepsMap): + # Deferred import to avoid a circular import (operation.dag imports us). + from aws_durable_execution_sdk_python.operation.dag import ( + run_nested_dag, + ) + + return run_nested_dag(ctx, task_name, register, config) + + return self._add(task_name, "dag", deps, trigger_rule, run_if, config, executor) + + # endregion task kinds + + def _run_child( + self, + ctx: DurableContext, + name: str, + body_takes_child: Callable[[DurableContext], Any], + config: ChildConfig, + sub_type: OperationSubType = OperationSubType.RUN_IN_CHILD_CONTEXT, + ) -> Any: + task_id = ctx._create_task_id(name) + + def wrapped(): + return body_takes_child(ctx.create_child_context(operation_id=task_id)) + + return child_handler( + func=wrapped, + state=ctx.state, + operation_identifier=OperationIdentifier( + operation_id=task_id, + sub_type=sub_type, + parent_id=ctx._parent_id, + name=name, + ), + config=config, + ) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_executor.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_executor.py new file mode 100644 index 00000000..63b8c8bd --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_executor.py @@ -0,0 +1,749 @@ +"""DagExecutor: a dedicated topological scheduler for the DAG operation. + +Reuses the SDK's worker-thread primitives (``ThreadPoolExecutor``, the +``SuspendExecution`` protocol) but is a *separate* component from +``ConcurrentExecutor`` (which is hard-wired for the flat map/parallel shape). +It gates task submission on dependency readiness, evaluates trigger rules and +``run_if`` predicates, drains on task *failure* by default (a task body that +raises is a terminal FAILED state, not an abort — spec §5.5), and computes +DAG-global success/failure/skip counts, feeding only success+failure into the +reused threshold ``CompletionConfig``. A ``run_if`` predicate that *raises* is +different: it is a defect in deterministic code, so it aborts the DAG with +``DagPredicateError`` rather than being recorded as a task failure. + +.. warning:: + **Experimental.** Internal implementation of the DAG scheduler. +""" + +from __future__ import annotations + +import heapq +import datetime +import itertools +import logging +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING, Any + +from aws_durable_execution_sdk_python.dag import ( + DagCompletionItemStatus, + DagCompletionOutcome, + DagCompletionReason, + DagCompletionStatus, + DagCustomCompletionConfig, + DepsMap, + SkipReason, + TaskExecution, + TaskStatus, +) +from aws_durable_execution_sdk_python.exceptions import ( + DagPredicateError, + SuspendExecution, + TimedSuspendExecution, + ValidationError, +) +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.operation.dag_result import DagResultImpl + +if TYPE_CHECKING: + from collections.abc import Callable + from concurrent.futures import Future + from typing import Self + + from aws_durable_execution_sdk_python.context import DurableContext + from aws_durable_execution_sdk_python.dag import DagConfig + from aws_durable_execution_sdk_python.operation.dag_context import TaskDef + +logger = logging.getLogger(__name__) + +_TERMINAL = (TaskStatus.SUCCEEDED, TaskStatus.FAILED, TaskStatus.SKIPPED) + +# Default cap on how many top-level DAG tasks the scheduler runs concurrently +# when the config leaves ``max_concurrency`` unset. Previously the DAG was +# unbounded: ``max_workers`` fell back to the task count, so an N-task DAG +# spawned N OS threads inside the Lambda sandbox (a 500-task DAG -> 500 threads). +# 40 is a pragmatic bound -- high enough that realistic graphs are unaffected, +# low enough to keep thread/socket usage sane in the smallest Lambda configs. It +# governs the DAG SCHEDULER ONLY (top-level tasks of THIS DAG); it is not +# inherited by a task's own map/parallel fan-out, and a nested dag task resolves +# its own independent default of 40. An explicit ``max_concurrency`` always wins, +# including a value above 40. See dag-review/DEFAULT_CONCURRENCY_CONTRACT.md. +DEFAULT_DAG_MAX_CONCURRENCY = 40 + +# task scheduling decisions +_RUN = "RUN" +_SKIP = "SKIP" + + +def _trigger_passes(rule, statuses: list[TaskStatus]) -> bool: + """Trigger-rule truth table over upstream terminal statuses. + + Ports the JS truth table verbatim, incl. the empty-upstream rows and the + ``ALL_FAILED`` ``len > 0`` guard. + """ + from aws_durable_execution_sdk_python.dag import TriggerRule + + has_failed = any(s is TaskStatus.FAILED for s in statuses) + has_succeeded = any(s is TaskStatus.SUCCEEDED for s in statuses) + if rule is TriggerRule.ALL_SUCCESS: + return all(s is TaskStatus.SUCCEEDED for s in statuses) + if rule is TriggerRule.ALL_FAILED: + return len(statuses) > 0 and all(s is TaskStatus.FAILED for s in statuses) + if rule is TriggerRule.ALL_DONE: + return True + if rule is TriggerRule.ANY_SUCCESS: + return has_succeeded + if rule is TriggerRule.ANY_FAILED: + return has_failed + if rule is TriggerRule.NONE_FAILED: + return not has_failed + msg = f"Unknown trigger rule: {rule}" # pragma: no cover + raise ValidationError(msg) # pragma: no cover + + +_resume_seq = itertools.count() + + +class _TimedResume: + """A name-keyed timed-resume record driven by the DAG's ``TimerScheduler``. + + The DAG-owned :class:`TimerScheduler` fires resume records on a background + timer thread: on fire it checks ``can_resume``, calls ``reset_to_pending()`` + then hands the record to its resubmit callback. The DAG tracks task state by + *name* rather than by an executable instance, so this record carries only + the task name. ``__lt__`` (via a monotonic sequence) keeps heap ties in the + scheduler total-orderable when two resumes share a timestamp. + """ + + __slots__ = ("_seq", "name") + + def __init__(self, name: str) -> None: + self.name = name + self._seq = next(_resume_seq) + + @property + def can_resume(self) -> bool: + return True + + def reset_to_pending(self) -> None: + """No-op: the DAG resets its own task bookkeeping in ``_resubmit``.""" + + def __lt__(self, other: _TimedResume) -> bool: + return self._seq < other._seq + + +class TimerScheduler: + """DAG-owned timer for in-process timed resumes. + + Manages timed suspend records with a background timer thread. This is a + self-contained copy of the mechanism the core map/parallel executor used + before that logic was inlined into ``ConcurrentExecutor``; the DAG owns it + so timed waits resume in-process (earliest-timed wins) while indefinite and + callback suspends still bubble to the platform. It drives :class:`_TimedResume` + records: on fire it transitions each resumable record to pending under the + lock, then hands the ready wave to ``resubmit_callback`` off the lock. + """ + + def __init__( + self, resubmit_callback: Callable[[list[_TimedResume]], None] + ) -> None: + self.resubmit_callback = resubmit_callback + self._pending_resumes: list[tuple[float, int, _TimedResume]] = [] + self._lock = threading.Lock() + self._schedule_counter = 0 + self._shutdown = threading.Event() + self._timer_thread = threading.Thread(target=self._timer_loop, daemon=True) + self._timer_thread.start() + + def __enter__(self) -> Self: + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + self.shutdown() + + def schedule_resume(self, exe_state: _TimedResume, resume_time: float) -> None: + """Schedule a record to resume at the specified time. + + Uses a counter as a tie-breaker to ensure FIFO ordering when multiple + records share a ``resume_time``. + """ + with self._lock: + heapq.heappush( + self._pending_resumes, + (resume_time, self._schedule_counter, exe_state), + ) + self._schedule_counter += 1 + + def shutdown(self) -> None: + """Shutdown the timer thread and cancel all pending resumes.""" + self._shutdown.set() + self._timer_thread.join(timeout=1.0) + with self._lock: + self._pending_resumes.clear() + + def _timer_loop(self) -> None: + """Background thread that processes timed resumes.""" + while not self._shutdown.is_set(): + next_resume_time = None + + with self._lock: + if self._pending_resumes: + next_resume_time = self._pending_resumes[0][0] + + if next_resume_time is None: + # No pending resumes, wait a bit and check again + self._shutdown.wait(timeout=0.1) + continue + + current_time = time.time() + if current_time >= next_resume_time: + # Drain every due resume under the lock, transitioning each to + # PENDING atomically with the pop, then resubmit off the lock so + # timed resumes don't serialize behind the resubmit callback and + # the timer thread can't re-enter this non-reentrant lock when a + # resubmit schedules another resume inline. + ready: list[_TimedResume] = [] + with self._lock: + while ( + self._pending_resumes + and self._pending_resumes[0][0] <= current_time + ): + _, _, exe_state = heapq.heappop(self._pending_resumes) + if exe_state.can_resume: + exe_state.reset_to_pending() + ready.append(exe_state) + if ready: + self.resubmit_callback(ready) + else: + # Wait until next resume time + wait_time = min(next_resume_time - current_time, 0.1) + self._shutdown.wait(timeout=wait_time) + + +class DagExecutor: + """Topological scheduler for a validated DAG.""" + + def __init__( + self, + ctx: DurableContext, + tasks: dict[str, TaskDef], + config: DagConfig, + ) -> None: + if config.max_concurrency is not None and config.max_concurrency <= 0: + msg = f"Invalid max_concurrency: {config.max_concurrency}" + raise ValidationError(msg) + self._ctx = ctx + self._tasks = tasks + self._config = config + self._lock = threading.Lock() + self._completion_event = threading.Event() + self._results: dict[str, TaskExecution] = {} + self._scheduled: set[str] = set() + self._in_flight: set[str] = set() + # First-observed wall-clock start per task name, stamped when the body is + # about to run. Terminal records copy it into TaskExecution.started_at so + # the envelope's startedAt/completedAt are populated (Python omitted them + # entirely before). Set once per name so a timed re-run keeps the first + # start. + self._started_at: dict[str, datetime.datetime] = {} + self._success = 0 + self._failure = 0 + self._skip = 0 + # All suspends raised by tasks this run. We do NOT re-raise the first + # one captured; when stopping we resolve which suspend to surface with + # the same precedence as ConcurrentExecutor.should_execution_suspend + # (earliest timed wins over indefinite) so a concurrent short timer is + # never dropped behind an indefinite wait_for_callback. + self._pending_suspends: list[SuspendExecution] = [] + # In-process timed-resume bookkeeping: a timed suspend does NOT stop the + # DAG. While other tasks make progress the DAG-owned TimerScheduler + # re-runs the timed task in this same invocation at its scheduled + # timestamp. Only an *indefinite* (callback) suspend forces leaving the + # invocation for platform replay. + self._scheduler: TimerScheduler | None = None + self._pending_timers: set[str] = set() + self._timed_suspend_by_name: dict[str, TimedSuspendExecution] = {} + self._scheduler_exception: Exception | None = None + self._early_reason: DagCompletionReason | None = None + self._pool: ThreadPoolExecutor | None = None + + # region public + def run( + self, + *, + reconstruct_started: set[str] | None = None, + reconstruct_reason: DagCompletionReason | None = None, + reconstruct_total: int | None = None, + ) -> DagResultImpl: + """Schedule and run the DAG; return a DagResult (may raise to suspend). + + Offloaded-replay reconstruct (contract "replay rule"): when the container + payload had ``tasks`` dropped, the caller passes the envelope's + ``startedTaskNames`` as ``reconstruct_started`` and its + ``completionReason``/``totalCount`` as ``reconstruct_reason``/ + ``reconstruct_total``. Reconstruct then re-runs this deterministic + register graph exactly as a first run would -- each task fast-paths from + its own retained child checkpoint, so bodies never re-execute -- except + that a task named in ``reconstruct_started`` is seeded STARTED and never + scheduled. That is the fix for the documented STARTED-set loss: an + in-flight task recorded STARTED in the offloaded envelope is reproduced + as STARTED instead of being restarted. The completion reason and total + are taken from the envelope rather than re-derived. + """ + total = len(self._tasks) + if total == 0: + reason = reconstruct_reason or DagCompletionReason.ALL_COMPLETED + return DagResultImpl( + {}, reason, total_count=reconstruct_total + ) + + if reconstruct_started: + # Seed the started set before pumping: mark each as STARTED and + # already-scheduled so _pump never submits it (no body run) and + # downstream deps see it as non-terminal (stay unscheduled), exactly + # reproducing the live in-flight snapshot. + with self._lock: + for name in reconstruct_started: + if name in self._tasks and name not in self._results: + self._results[name] = TaskExecution( + name=name, status=TaskStatus.STARTED + ) + self._scheduled.add(name) + + # Resolve the single effective concurrency bound: an explicit + # max_concurrency always wins (including a value above the default); + # otherwise cap at DEFAULT_DAG_MAX_CONCURRENCY. This is BOTH the + # scheduler's in-flight bound and the pool's max_workers -- the pool is + # the resource that made an unbounded DAG spawn one OS thread per task. + max_workers = self._config.max_concurrency or min( + total, DEFAULT_DAG_MAX_CONCURRENCY + ) + # Mirror ConcurrentExecutor.execute: scheduler OUTER, pool INNER, so the + # pool drains (joins in-flight tasks) before the timer thread is torn + # down. Any suspend is raised inside the pool ``with`` (as before). + with ( + TimerScheduler(self._resubmit) as scheduler, + ThreadPoolExecutor(max_workers=max_workers) as pool, + ): + self._scheduler = scheduler + self._pool = pool + self._pump() + self._completion_event.wait() + if self._scheduler_exception is not None: + raise self._scheduler_exception + suspend = self._resolve_suspend() + if suspend is not None: + raise suspend + return self._build_result( + reconstruct_reason=reconstruct_reason, + reconstruct_total=reconstruct_total, + ) + + # endregion public + + # region scheduling + def _pump(self) -> None: + """Resolve newly-ready tasks (skip or submit); set completion if done.""" + to_submit: list[tuple[str, TaskDef]] = [] + with self._lock: + progressed = True + while progressed: + progressed = False + if self._stopping_locked(): + break + for name, task in self._tasks.items(): + if name in self._scheduled or not self._deps_terminal_locked(name): + continue + decision, payload = self._evaluate_locked(task) + self._scheduled.add(name) + if decision == _RUN: + self._in_flight.add(name) + to_submit.append((name, task)) + else: # _SKIP: trigger rule or run_if predicate returned False + self._results[name] = TaskExecution( + name=name, status=TaskStatus.SKIPPED, skip_reason=payload + ) + self._skip += 1 + progressed = True + done = self._is_done_locked() + + for name, task in to_submit: + future = self._pool.submit(self._run_task, name, task) # type: ignore[union-attr] + + def _done(f: Future, n: str = name) -> None: + self._on_done(n, f) + + future.add_done_callback(_done) + + if done: + self._completion_event.set() + + def _run_task(self, name: str, task: TaskDef) -> Any: + # Snapshot deps under the lock: this runs on a worker thread and + # _build_deps_map reads self._results, which the scheduler mutates + # concurrently (the run_if path already builds deps under the lock). + with self._lock: + deps_map = self._build_deps_map(task) + # Stamp the first-observed start for this task name (kept across a + # timed re-run). Copied into the terminal/STARTED record so the + # envelope carries startedAt. + self._started_at.setdefault(name, datetime.datetime.now(datetime.UTC)) + logger.debug("DAG task %s starting", name) + return task.executor(self._ctx, deps_map) + + def _on_done(self, name: str, future: Future) -> None: + completed_at = datetime.datetime.now(datetime.UTC) + try: + result = future.result() + with self._lock: + self._results[name] = TaskExecution( + name=name, + status=TaskStatus.SUCCEEDED, + result=result, + started_at=self._started_at.get(name), + completed_at=completed_at, + ) + self._success += 1 + self._in_flight.discard(name) + except SuspendExecution as se: # includes TimedSuspendExecution + schedule_ts: float | None = None + with self._lock: + # Record every suspend (timed + indefinite); precedence is + # resolved in _resolve_suspend when the DAG stops. A STARTED task + # has begun but not completed, so it carries startedAt but no + # completedAt. + self._pending_suspends.append(se) + self._results[name] = TaskExecution( + name=name, + status=TaskStatus.STARTED, + started_at=self._started_at.get(name), + ) + self._in_flight.discard(name) + # Timed suspend: register an in-process resume so the base timer + # thread re-runs this task at its timestamp WITHOUT leaving the + # invocation. Indefinite (callback) suspends get no timer and + # fall through to platform replay via _resolve_suspend. + if isinstance(se, TimedSuspendExecution): + self._timed_suspend_by_name[name] = se + self._pending_timers.add(name) + schedule_ts = se.scheduled_timestamp + if schedule_ts is not None and self._scheduler is not None: + self._scheduler.schedule_resume(_TimedResume(name), schedule_ts) + except Exception as e: # noqa: BLE001 + with self._lock: + self._results[name] = TaskExecution( + name=name, + status=TaskStatus.FAILED, + error=ErrorObject.from_exception(e), + started_at=self._started_at.get(name), + completed_at=completed_at, + ) + self._failure += 1 + self._in_flight.discard(name) + self._safe_pump() + + def _safe_pump(self) -> None: + """Run ``_pump`` from a worker-thread completion callback. + + ``concurrent.futures`` swallows exceptions raised inside + ``add_done_callback``. If ``_pump`` ever raised there (e.g. an + unexpected scheduler bug) the completion event would never be set and + ``run()`` would block forever. Capture any escaping exception and set + the event so ``run()`` re-raises it instead of hanging. + """ + try: + self._pump() + except Exception as e: # noqa: BLE001 + with self._lock: + if self._scheduler_exception is None: + self._scheduler_exception = e + self._completion_event.set() + # endregion scheduling + + def _resubmit(self, resumes: list[_TimedResume]) -> None: + """DAG-owned TimerScheduler callback: re-run a wave of timed tasks in-process. + + Fires on the scheduler's timer thread once tasks' scheduled timestamps + elapse. The DAG-owned ``TimerScheduler`` batches all due resumes into one + callback invocation (one checkpoint refresh serves the whole wave), so + this accepts a list. It clears each task's timed-suspend bookkeeping and its STARTED placeholder so ``_pump`` sees them as fresh, + ready tasks and re-runs them within the same invocation. Tasks that + already left the timer set (e.g. the DAG bubbled to the platform and the + scheduler was torn down) are skipped. + """ + with self._lock: + # Abort guard: once the DAG has decided to abort (a run_if predicate + # raised, so _scheduler_exception is set), no further checkpoint may + # be written and no task may re-run. The scheduler is the OUTER + # context manager and the pool the INNER one, so the timer thread is + # still alive during the pool-drain window; a resume whose timestamp + # elapsed in that window would otherwise fire create_checkpoint() + # AFTER the abort decision — a stray flush that contradicts the abort + # contract (run() re-raises _scheduler_exception once the pool + # drains). Bail before mutating state or checkpointing. This is the + # only scheduler entry point that checkpoints; _on_done/_pump are + # already gated by _stopping_locked (which checks _scheduler_exception + # first), so no other teardown-window checkpoint exists. + # + # create_checkpoint() itself is called INSIDE this same lock + # acquisition (below), not after releasing it: _scheduler_exception + # is set by _safe_pump on a worker-thread completion callback (a + # different thread than this timer thread), so if the checkpoint + # call were outside the lock, that worker thread could set the + # exception in the gap between releasing the lock and reaching + # create_checkpoint() -- the exact stray-flush-after-abort this + # guard exists to prevent, just via a narrower window than a + # lock-scoped check alone would close. create_checkpoint() blocks + # on the background checkpoint-batching thread's queue, not on + # this lock, so holding self._lock across the call is deadlock-safe. + if self._scheduler_exception is not None: + return + for resume in resumes: + name = resume.name + if name not in self._pending_timers: + continue + self._pending_timers.discard(name) + se = self._timed_suspend_by_name.pop(name, None) + if se is not None: + try: + self._pending_suspends.remove(se) + except ValueError: # pragma: no cover - defensive + pass + # Make the task schedulable again: drop its STARTED placeholder + # and its scheduled mark so _pump re-evaluates and re-runs it. + self._scheduled.discard(name) + self._results.pop(name, None) + # Checkpoint before re-running, matching ConcurrentExecutor.resubmitter. + # Deliberately still under self._lock -- see the guard comment above. + self._ctx.state.create_checkpoint() + self._safe_pump() + + def _resolve_suspend(self) -> SuspendExecution | None: + """Pick which suspend to surface, matching the base executor's contract. + + Ports ``ConcurrentExecutor.should_execution_suspend`` precedence: if any + timed suspend is pending, raise a ``TimedSuspendExecution`` with the + EARLIEST ``scheduled_timestamp`` (timed wins over indefinite so the + platform resumes at the soonest timer); otherwise raise the indefinite + ``SuspendExecution``. Returns ``None`` when nothing suspended. Called + after the completion event fires, so no lock is needed. + """ + earliest_timestamp = float("inf") + indefinite: SuspendExecution | None = None + for se in self._pending_suspends: + if isinstance(se, TimedSuspendExecution): + if se.scheduled_timestamp < earliest_timestamp: + earliest_timestamp = se.scheduled_timestamp + else: + indefinite = se + if earliest_timestamp != float("inf"): + return TimedSuspendExecution( + "DAG suspended; resuming at the earliest pending timer.", + earliest_timestamp, + ) + return indefinite + + # region helpers (lock held) + def _deps_terminal_locked(self, name: str) -> bool: + task = self._tasks[name] + for dep in task.all_deps: + te = self._results.get(dep.name) + if te is None or te.status not in _TERMINAL: + return False + return True + + def _evaluate_locked(self, task: TaskDef) -> tuple[str, Any]: + """Decide a ready task's fate: ``(_RUN, None)`` or ``(_SKIP, SkipReason)``. + + The trigger rule is a pure function of upstream enum statuses. ``run_if`` + is user-supplied code, but it is specified as a synchronous, + deterministic, pure predicate over resolved upstream results, not a + checkpointed operation. A ``run_if`` that raises is therefore a defect in + deterministic code, not a business outcome: we neither record the task + ``FAILED`` (which would silently drive every downstream ``ALL_FAILED`` / + ``ANY_FAILED`` / ``ALL_DONE`` compensation path off a scheduler defect) + nor ``SKIPPED``. Instead we raise :class:`DagPredicateError`, chaining the + original exception, so the whole DAG aborts and ``dag()`` fails loudly. + The offending task is left with no terminal state (it is never added to + ``self._results``). This is identical for root and non-root tasks. The + raise propagates out of ``_pump``: on the caller thread (first pump) it + leaves ``run()`` directly; on a worker/timer thread it is captured by + ``_safe_pump`` into ``self._scheduler_exception`` and re-raised by + ``run()`` after the pool drains. + """ + statuses = [self._results[dep.name].status for dep in task.all_deps] + if not _trigger_passes(task.trigger_rule, statuses): + return (_SKIP, SkipReason.TRIGGER_RULE) + if task.run_if is not None: + deps_map = self._build_deps_map(task) + try: + should_run = task.run_if(deps_map) + except Exception as e: + msg = ( + f"run_if predicate for DAG task {task.name!r} raised " + f"{type(e).__name__}: {e}" + ) + raise DagPredicateError(msg, task_name=task.name) from e + if not should_run: + return (_SKIP, SkipReason.RUN_IF_PREDICATE) + return (_RUN, None) + + def _build_deps_map(self, task: TaskDef) -> DepsMap: + by_name: dict[str, Any] = {} + for dep in task.inline_deps: + te = self._results.get(dep.name) + by_name[dep.name] = te.result if te else None + return DepsMap(by_name) + + def _threshold_reason_locked(self) -> DagCompletionReason | None: + """Early-completion reason: either a custom ``should_complete`` + predicate's decision, or the threshold logic mirroring + ``ExecutionCounters.should_complete``. + + Threshold order matches the reused batch logic: success threshold + first, then the failure-tolerance conditions, then the + impossible-to-succeed early stop (which batch reports as + ``FAILURE_TOLERANCE_EXCEEDED`` — see + ``ConcurrentExecutor._create_result``). The failure-percentage + denominator excludes SKIPPED tasks (they neither succeed nor fail) so + skips do not dilute the ratio. + """ + cc = self._config.completion_config + if cc is None: + return None + if isinstance(cc, DagCustomCompletionConfig): + status = self._build_completion_status_locked() + decision = cc.should_complete(status) + if decision.complete: + if decision.outcome == DagCompletionOutcome.FAILED: + return DagCompletionReason.CUSTOM_COMPLETION_FAILED + return DagCompletionReason.CUSTOM_COMPLETION_SUCCEEDED + return None + min_successful = cc.min_successful + # Success condition (checked before failure, matching batch semantics). + if min_successful is not None and self._success >= min_successful: + return DagCompletionReason.MIN_SUCCESSFUL_REACHED + # Failure-tolerance conditions (count, then percentage). + if ( + cc.tolerated_failure_count is not None + and self._failure > cc.tolerated_failure_count + ): + return DagCompletionReason.FAILURE_TOLERANCE_EXCEEDED + if cc.tolerated_failure_percentage is not None: + denom = len(self._tasks) - self._skip + if denom > 0: + pct = (self._failure / denom) * 100 + if pct > cc.tolerated_failure_percentage: + return DagCompletionReason.FAILURE_TOLERANCE_EXCEEDED + # Impossible-to-succeed early stop: max reachable successes is every task + # that has not already failed or been skipped. + if min_successful is not None: + reachable = len(self._tasks) - self._failure - self._skip + if reachable < min_successful: + return DagCompletionReason.FAILURE_TOLERANCE_EXCEEDED + return None + + def _build_completion_status_locked(self) -> DagCompletionStatus: + """Builds the live progress snapshot passed to a custom + ``should_complete`` predicate: every task in registration order, keyed + by name, reflecting exactly what has settled so far (tasks with no + entry in ``self._results`` yet are reported with a ``None`` status, + i.e. not yet started). + """ + items: list[DagCompletionItemStatus] = [] + by_name: dict[str, DagCompletionItemStatus] = {} + for name in self._tasks: + te = self._results.get(name) + item = ( + DagCompletionItemStatus(name=name) + if te is None + else DagCompletionItemStatus( + name=name, + status=te.status, + result=te.result, + skip_reason=te.skip_reason, + ) + ) + items.append(item) + by_name[name] = item + completed = self._success + self._failure + self._skip + return DagCompletionStatus( + success_count=self._success, + failure_count=self._failure, + skipped_count=self._skip, + completed_count=completed, + total_count=len(self._tasks), + items=items, + results=by_name, + ) + + def _has_indefinite_locked(self) -> bool: + """True if any *indefinite* (non-timed) suspend is outstanding. + + Only indefinite suspends (e.g. ``wait_for_callback``) force leaving the + invocation for platform replay; timed suspends are resumed in-process by + the DAG-owned ``TimerScheduler``. + """ + return any( + not isinstance(se, TimedSuspendExecution) for se in self._pending_suspends + ) + + def _stopping_locked(self) -> bool: + # A captured scheduler exception (e.g. a run_if predicate raised and the + # DAG is aborting with DagPredicateError) stops all further scheduling: + # no new tasks start while the pool drains any in-flight work. Checked + # first so an abort is never downgraded by a threshold reason. + if self._scheduler_exception is not None: + return True + # An indefinite suspend can only be resolved by the platform, so we stop + # scheduling new work and drain (unchanged pre-timer behaviour). Timed + # suspends do NOT stop the DAG: they are resumed in-process while other + # tasks keep making progress (parity with ConcurrentExecutor). + if self._has_indefinite_locked(): + return True + reason = self._threshold_reason_locked() + if reason is not None: + self._early_reason = reason + return True + return False + + def _has_schedulable_locked(self) -> bool: + for name in self._tasks: + if name not in self._scheduled and self._deps_terminal_locked(name): + return True + return False + + def _is_done_locked(self) -> bool: + if self._stopping_locked(): + return len(self._in_flight) == 0 + if self._in_flight: + return False + return not self._has_schedulable_locked() + + # endregion helpers + + def _build_result( + self, + *, + reconstruct_reason: DagCompletionReason | None = None, + reconstruct_total: int | None = None, + ) -> DagResultImpl: + if reconstruct_reason is not None: + # Offloaded reconstruct: the completion reason is authoritative from + # the envelope, not re-derived (a re-derivation over fast-pathed + # results could disagree at an early-completion boundary). + reason = reconstruct_reason + elif self._early_reason is not None: + reason = self._early_reason + elif self._failure == 0: + reason = DagCompletionReason.ALL_COMPLETED + else: + reason = DagCompletionReason.COMPLETED_WITH_FAILURES + task_kinds = {name: task.kind for name, task in self._tasks.items()} + total = ( + reconstruct_total if reconstruct_total is not None else len(self._tasks) + ) + return DagResultImpl( + dict(self._results), reason, task_kinds, total_count=total + ) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_result.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_result.py new file mode 100644 index 00000000..d0bf7e9f --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_result.py @@ -0,0 +1,368 @@ +"""DagResultImpl + serialization for the DAG operation. + +.. warning:: + **Experimental.** Internal implementation of :class:`~...dag.DagResult`. +""" + +from __future__ import annotations + +import datetime +from typing import TYPE_CHECKING, Any, TypeVar, overload + +from aws_durable_execution_sdk_python.concurrency.models import ( + BatchResult, + CompletionReason, +) +from aws_durable_execution_sdk_python.dag import ( + DagCompletionReason, + DagResult, + SkipReason, + TaskExecution, + TaskHandle, + TaskStatus, +) +from aws_durable_execution_sdk_python.exceptions import DagExecutionError +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.serdes import SerDes + +if TYPE_CHECKING: + from collections.abc import Mapping + + from aws_durable_execution_sdk_python.serdes import SerDesContext + +T = TypeVar("T") + +# The single converged envelope discriminator (contract: ``type: "DagResult"``). +_ENVELOPE_TYPE = "DagResult" + +# result_kind discriminators +_KIND_PLAIN = "plain" +_KIND_BATCH = "batch" +_KIND_DAG = "dag" + +# TaskDef.kind values whose result is a BatchResult / DagResult +_BATCH_KINDS = frozenset({"map", "parallel"}) +_DAG_KINDS = frozenset({"dag"}) + + +def _iso_millis(dt: datetime.datetime | None) -> str | None: + """Format a datetime as the contract timestamp: UTC, millisecond precision, + ``Z`` suffix (e.g. ``2026-07-26T03:19:01.884Z``). ``None`` stays ``None`` + (the value is genuinely unknown).""" + if dt is None: + return None + utc = dt.astimezone(datetime.UTC) + return f"{utc.strftime('%Y-%m-%dT%H:%M:%S')}.{utc.microsecond // 1000:03d}Z" + + +def _parse_iso(value: str | None) -> datetime.datetime | None: + """Parse a contract timestamp back to an aware UTC datetime. Tolerant of a + trailing ``Z`` (which older Pythons' ``fromisoformat`` rejected). ``None`` + and unparseable values yield ``None`` rather than raising -- a timestamp is + informational and must never break deserialization (contract rule 4).""" + if not value: + return None + try: + return datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + except (ValueError, TypeError): # pragma: no cover - defensive + return None + + +def _error_to_dict(err: ErrorObject | None) -> dict[str, Any] | None: + """Serialize an error to the canonical PascalCase object with explicit nulls. + + The three canonical keys (``ErrorType``, ``ErrorMessage``, ``StackTrace``) + are always present (``null`` when absent) so the payload is diffable across + languages; any extra platform field (e.g. ``ErrorData``) is preserved.""" + if err is None: + return None + d = dict(err.to_dict()) + d.setdefault("ErrorType", None) + d.setdefault("ErrorMessage", None) + d.setdefault("StackTrace", None) + return d + + +def dag_reason_from_core(core: CompletionReason) -> DagCompletionReason: + """Bridge a batch ``CompletionReason`` into the DAG's superset enum.""" + return DagCompletionReason(core.value) + + +def _result_kind(task_kind: str | None) -> str: + if task_kind in _BATCH_KINDS: + return _KIND_BATCH + if task_kind in _DAG_KINDS: + return _KIND_DAG + return _KIND_PLAIN + + + +def _name_of(task: str | TaskHandle[Any]) -> str: + return task.name if isinstance(task, TaskHandle) else task + + +class DagResultImpl(DagResult): + """Concrete DAG result. Mirrors the ``BatchResult`` accessor surface. + + .. warning:: + **Experimental.** + """ + + def __init__( + self, + results: dict[str, TaskExecution], + completion_reason: DagCompletionReason, + task_kinds: dict[str, str] | None = None, + total_count: int | None = None, + success_count: int | None = None, + failure_count: int | None = None, + skipped_count: int | None = None, + ) -> None: + self._results = results + self._completion_reason = completion_reason + self._task_kinds = task_kinds or {} + # total_count is the number of REGISTERED tasks in the DAG (spec §2.8), + # a fixed value independent of early completion / never-started tasks. + # Defaults to len(results) when omitted (fully-recorded DAGs). + self._total_count = total_count if total_count is not None else len(results) + # The three aggregate counts are normally DERIVED from the per-task map. + # But an offloaded envelope (no ``tasks``) still carries them, and the + # map is legitimately empty when restored from that envelope: honouring + # the stored value is contract rule 1 -- a tasks-less restore MUST + # preserve the counts rather than fabricate zeros. ``None`` means "derive + # from the map" (the normal, fully-recorded path); an explicit value + # (from ``from_dict``) always wins, mirroring ``total_count``. + self._success_count = success_count + self._failure_count = failure_count + self._skipped_count = skipped_count + + # region accessors + @overload + def get_result(self, task: TaskHandle[T]) -> T | None: ... + @overload + def get_result(self, task: str) -> Any: ... + def get_result(self, task: str | TaskHandle[Any]) -> Any: + """Return a task's result (or ``None`` if absent / not succeeded). + + Passing the originating :class:`TaskHandle` preserves the task's result + type for static typing (``get_result(handle) -> T | None``); passing a + name string returns ``Any``. Both resolve by task name at runtime. The + handle overload is ``T | None`` (not bare ``T``) because a missing, + FAILED, or SKIPPED task has no result and yields ``None`` -- the type + reflects that runtime behavior. + """ + te = self._results.get(_name_of(task)) + return te.result if te else None + + def get_status(self, task: str | TaskHandle[Any]) -> TaskStatus | None: + """Return a task's status, or ``None`` if the task never ran.""" + te = self._results.get(_name_of(task)) + return te.status if te else None + + def succeeded(self) -> list[TaskExecution]: + """Tasks that SUCCEEDED.""" + return [t for t in self._results.values() if t.status is TaskStatus.SUCCEEDED] + + def failed(self) -> list[TaskExecution]: + """Tasks that FAILED.""" + return [t for t in self._results.values() if t.status is TaskStatus.FAILED] + + def skipped(self) -> list[TaskExecution]: + """Tasks that were SKIPPED.""" + return [t for t in self._results.values() if t.status is TaskStatus.SKIPPED] + + @property + def results(self) -> Mapping[str, TaskExecution]: + """All recorded task executions, keyed by name.""" + return self._results + + @property + def success_count(self) -> int: + if self._success_count is not None: + return self._success_count + return len(self.succeeded()) + + @property + def failure_count(self) -> int: + if self._failure_count is not None: + return self._failure_count + return len(self.failed()) + + @property + def skipped_count(self) -> int: + if self._skipped_count is not None: + return self._skipped_count + return len(self.skipped()) + + @property + def total_count(self) -> int: + return self._total_count + + @property + def completion_reason(self) -> DagCompletionReason: + return self._completion_reason + + def throw_if_error(self) -> None: + """Raise :class:`DagExecutionError` if any task FAILED, or if a custom + completion predicate completed the DAG as ``CUSTOM_COMPLETION_FAILED`` + (which can happen with zero individually-failed tasks -- the predicate's + verdict, not a task exception, is what failed the DAG). + """ + failures = self.failed() + if failures: + first = failures[0] + detail = first.error.message if first.error else "unknown error" + msg = ( + f"DAG completed with {len(failures)} failed task(s); " + f"first failure '{first.name}': {detail}" + ) + raise DagExecutionError(msg) + if self.completion_reason is DagCompletionReason.CUSTOM_COMPLETION_FAILED: + msg = ( + "DAG completed with reason CUSTOM_COMPLETION_FAILED (a custom " + "completion predicate completed the DAG as a failure)" + ) + raise DagExecutionError(msg) + + # endregion accessors + + # region serialization + def to_dict(self) -> dict[str, Any]: + """Serialize to the converged cross-language DAG envelope. + + Single shape for both the inline and offloaded cases (the offloaded case + drops only ``tasks``; see the degradation ladder in ``operation/dag.py``). + Every canonical field is always present; absent values are ``null``, + never omitted. Aggregate fields are always present even though they are + derivable from ``tasks`` -- that redundancy is what lets the offloaded + payload keep the same shape after ``tasks`` is dropped. Field order + follows the contract listing for console readability (structural + comparison ignores order). + """ + return { + "type": _ENVELOPE_TYPE, + "totalCount": self._total_count, + "successCount": self.success_count, + "failureCount": self.failure_count, + "skippedCount": self.skipped_count, + "completionReason": self._completion_reason.value, + "startedTaskNames": [ + te.name + for te in self._results.values() + if te.status is TaskStatus.STARTED + ], + "failedTaskNames": [ + te.name + for te in self._results.values() + if te.status is TaskStatus.FAILED + ], + "tasks": [self._task_to_dict(te) for te in self._results.values()], + } + + def _task_to_dict(self, te: TaskExecution) -> dict[str, Any]: + kind = _result_kind(self._task_kinds.get(te.name)) + result_value: Any = None + if te.result is not None: + if ( + kind == _KIND_BATCH + and isinstance(te.result, BatchResult) + or kind == _KIND_DAG + and isinstance(te.result, DagResultImpl) + ): + result_value = te.result.to_dict() + else: + result_value = te.result + # resultKind describes how to interpret ``result``, so it is null when there + # is no result to interpret: a FAILED or SKIPPED task carries null for both. + # All four SDKs agree on this (envelope contract rule 1, explicit nulls). + serialized_kind = kind if te.status is TaskStatus.SUCCEEDED else None + return { + "name": te.name, + "status": te.status.value, + "skipReason": te.skip_reason.value if te.skip_reason else None, + "resultKind": serialized_kind, + "result": result_value, + "error": _error_to_dict(te.error), + "startedAt": _iso_millis(te.started_at), + "completedAt": _iso_millis(te.completed_at), + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> DagResultImpl: + """Deserialize the converged envelope. + + Reads the ``tasks`` array; unknown fields are ignored and a missing + field is treated as absent rather than an error (contract rule 4, + additive-only evolution). An envelope with no ``tasks`` (the offloaded + case) yields an empty results map, but the aggregate summary -- + ``totalCount``, the three counts, and ``completionReason`` -- is read + straight from the envelope and preserved (contract rule 1: a tasks-less + restore MUST NOT fabricate zeroed counts or ``ALL_COMPLETED``). The + offloaded reconstruct path repopulates the per-task detail from the child + checkpoints instead (see ``operation/dag.py``); it does not call this + method to rebuild tasks. + """ + results: dict[str, TaskExecution] = {} + task_kinds: dict[str, str] = {} + for td in data.get("tasks") or []: + name = td["name"] + kind = td.get("resultKind", _KIND_PLAIN) + result_value = td.get("result") + if result_value is not None: + if kind == _KIND_BATCH: + result_value = BatchResult.from_dict(result_value) + elif kind == _KIND_DAG: + result_value = cls.from_dict(result_value) + error_raw = td.get("error") + results[name] = TaskExecution( + name=name, + status=TaskStatus(td["status"]), + skip_reason=( + SkipReason(td["skipReason"]) if td.get("skipReason") else None + ), + result=result_value, + error=ErrorObject.from_dict(error_raw) if error_raw else None, + started_at=_parse_iso(td.get("startedAt")), + completed_at=_parse_iso(td.get("completedAt")), + ) + task_kinds[name] = ( + "dag" + if kind == _KIND_DAG + else ("map" if kind == _KIND_BATCH else "step") + ) + return cls( + results=results, + completion_reason=DagCompletionReason(data["completionReason"]), + task_kinds=task_kinds, + total_count=data.get("totalCount"), + success_count=data.get("successCount"), + failure_count=data.get("failureCount"), + skipped_count=data.get("skippedCount"), + ) + + # endregion serialization + + +class DagResultSerDes(SerDes): + """SerDes for the inline DagResult container payload. + + Serializes/deserializes the full converged envelope (with ``tasks``). The + offloaded degradation ladder and the reconstruct path live in + ``operation/dag.py`` because they need to manipulate the envelope structure + and read the retained child checkpoints, which a plain SerDes cannot do. + """ + + def serialize(self, value: DagResultImpl, serdes_context: SerDesContext) -> str: + import json + + return json.dumps(value.to_dict()) + + def deserialize(self, data: str, serdes_context: SerDesContext) -> DagResultImpl: + import json + + return DagResultImpl.from_dict(json.loads(data)) + + +def create_dag_result_serdes() -> SerDes: + """Return a SerDes that round-trips a :class:`DagResultImpl`.""" + return DagResultSerDes() diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_validator.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_validator.py new file mode 100644 index 00000000..85af0476 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_validator.py @@ -0,0 +1,106 @@ +"""DAG validation: names, duplicates, foreign deps, and cycle detection. + +.. warning:: + **Experimental.** Internal validation for the DAG operation. +""" + +from __future__ import annotations + +import re +from collections import deque +from typing import TYPE_CHECKING + +from aws_durable_execution_sdk_python.exceptions import ( + DagCyclicDependencyError, + DagDuplicateTaskError, + DagInvalidDependencyError, + DagInvalidTaskNameError, +) + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python.operation.dag_context import ( + DagContextImpl, + TaskDef, + ) + +_MAX_NAME_LEN = 100 +_NAME_RE = re.compile(r"[a-zA-Z0-9_]+") +_RESERVED_TOKEN = "DAG_NODE_T_" + + +def _validate_name(name: str) -> None: + if not name: + msg = "Task name must be non-empty." + raise DagInvalidTaskNameError(msg) + if len(name) > _MAX_NAME_LEN: + msg = f"Task name '{name}' exceeds {_MAX_NAME_LEN} characters." + raise DagInvalidTaskNameError(msg) + if not _NAME_RE.fullmatch(name): + msg = ( + f"Task name '{name}' is invalid: must match ^[a-zA-Z0-9_]+$ " + "(no '-' or other special characters)." + ) + raise DagInvalidTaskNameError(msg) + if _RESERVED_TOKEN in name: + msg = f"Task name '{name}' must not contain the reserved token '{_RESERVED_TOKEN}'." + raise DagInvalidTaskNameError(msg) + + +def validate_dag(dag_ctx: DagContextImpl) -> None: + """Validate a registered DAG. Deterministic; identical result on replay. + + Raises: + DagInvalidTaskNameError: invalid or reserved task name. + DagDuplicateTaskError: two tasks share a name. + DagInvalidDependencyError: a dep is not registered in this scope. + DagCyclicDependencyError: the dependency graph contains a cycle. + """ + registration_order: list[TaskDef] = dag_ctx.get_registration_order() + tasks: dict[str, TaskDef] = dag_ctx.get_tasks() + + # 1. names + 2. duplicates (in registration order for determinism) + seen: set[str] = set() + for task in registration_order: + _validate_name(task.name) + if task.name in seen: + msg = f"Duplicate task name '{task.name}' in DAG scope." + raise DagDuplicateTaskError(msg) + seen.add(task.name) + + # 3. deps must be registered in this scope + for task in tasks.values(): + for dep in task.all_deps: + if dep.name not in tasks: + msg = ( + f"Task '{task.name}' depends on '{dep.name}', which is not " + "registered in this DAG scope." + ) + raise DagInvalidDependencyError(msg) + + # 4. cycle detection via Kahn's algorithm over all_deps + _detect_cycle(tasks) + + +def _detect_cycle(tasks: dict[str, TaskDef]) -> None: + indegree: dict[str, int] = {name: 0 for name in tasks} + dependents: dict[str, list[str]] = {name: [] for name in tasks} + for name, task in tasks.items(): + dep_names = {dep.name for dep in task.all_deps} + indegree[name] = len(dep_names) + for dep_name in dep_names: + dependents[dep_name].append(name) + + queue: deque[str] = deque(n for n, d in indegree.items() if d == 0) + processed = 0 + while queue: + current = queue.popleft() + processed += 1 + for child in dependents[current]: + indegree[child] -= 1 + if indegree[child] == 0: + queue.append(child) + + if processed != len(tasks): + cyclic = sorted(n for n, d in indegree.items() if d > 0) + msg = f"DAG contains a dependency cycle among tasks: {cyclic}" + raise DagCyclicDependencyError(msg) diff --git a/packages/aws-durable-execution-sdk-python/tests/conformance/__init__.py b/packages/aws-durable-execution-sdk-python/tests/conformance/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/packages/aws-durable-execution-sdk-python/tests/conformance/dag_conformance_test.py b/packages/aws-durable-execution-sdk-python/tests/conformance/dag_conformance_test.py new file mode 100644 index 00000000..ba8e6f89 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/conformance/dag_conformance_test.py @@ -0,0 +1,661 @@ +"""Cross-language DAG conformance suite (Python side). + +Implements the applicable scenarios from the canonical catalog +``aws-durable-execution-sdk-js/docs/DAG_CONFORMANCE.md`` against the shipped +``context.dag()`` API, asserts each actual outcome against the catalog's +expected *semantic* outcome, and emits one key-sorted normalized JSON record +file to ``dag-conformance-out/python.json`` (schema per catalog Part B). + +Applicability (catalog Part C): all 19 scenarios (DAG-1..19) apply to Python. +DAG-18 (custom result-based completion) originally shipped TS+Go only; Python +added `DagCustomCompletionConfig`/`should_complete` after v1 (the DAG module +needed no base-SDK predicate hook to build it on, unlike what the catalog +originally assumed). + +All scenarios (including DAG-16/17 early completion) now conform to the +canonical catalog outcomes; ``total_count`` equals the registered task count +per spec §2.8. +""" + +from __future__ import annotations + +import hashlib +import json +import re +from pathlib import Path +from typing import Any + +import pytest + +from aws_durable_execution_sdk_python import exceptions +from aws_durable_execution_sdk_python.config import CompletionConfig, StepConfig +from aws_durable_execution_sdk_python.dag import ( + DagCompletionOutcome, + DagConfig, + DagCustomCompletionConfig, + DagResult, + TaskStatus, + TriggerRule, + complete_dag, + continue_dag, +) +from aws_durable_execution_sdk_python.retries import RetryPresets +from tests.dag_support import InMemoryServiceClient, make_context, make_state + +NO_RETRY = RetryPresets.none() +# Per-task step config that disables retries so an intentionally failing step +# fails promptly instead of falling back to RetryPresets.default(). +NO_RETRY_CFG = StepConfig(retry_strategy=NO_RETRY) + +OUT_PATH = Path("/Users/parpooya/workplace/dag-conformance-out/python.json") +_NAME_CHARSET = re.compile(r"[A-Za-z0-9_]+") + +# Records accumulate here across scenario helpers, then get written once. +RECORDS: dict[str, dict[str, Any]] = {} + + +# region helpers +def _run(register: Any, name: str, config: DagConfig | None = None): + state, client = make_state() + result = make_context(state).dag(register, name=name, config=config) + return result, client + + +def _fail(_deps: Any, _sc: Any) -> Any: + raise RuntimeError("boom") + + +def _structural_checks(client: InMemoryServiceClient) -> dict[str, bool]: + """Compute the four per-language structural entity-ID invariants. + + Per the conformance catalog (Part A, note 2), hashes are NOT expected to + match across languages and the invariants MUST be proven on each SDK's own + *pre-image*. Python re-hashes per level: a task's backend id is + ``blake2b(f"{parent}-DAG_NODE_T_{name}")[:64]``. A task op is therefore any + operation whose id equals that recomputation from its ``(parent_id, name)``; + counter ops (e.g. the top-level DAG container) never match. For the empty + DAG there are no task ops, so all four are vacuously ``true``. + """ + + def task_preimage(op: Any) -> str: + return ( + f"{op.parent_id}-DAG_NODE_T_{op.name}" + if op.parent_id + else f"DAG_NODE_T_{op.name}" + ) + + def is_task(op: Any) -> bool: + if op.name is None: + return False + digest = hashlib.blake2b(task_preimage(op).encode()).hexdigest()[:64] + return digest == op.operation_id + + ops = list(client.operations.values()) + task_ops = [op for op in ops if is_task(op)] + counter_ids = [op.operation_id for op in ops if not is_task(op)] + if not task_ops: + return { + "name_based": True, + "has_delimiter": True, + "dash_free": True, + "disjoint_from_counter": True, + } + # name_based: id is exactly the blake2b bound of the name-based pre-image. + name_based = all( + op.operation_id + == hashlib.blake2b(task_preimage(op).encode()).hexdigest()[:64] + for op in task_ops + ) + # has_delimiter: the pre-image carries the reserved token once per level. + has_delimiter = all(task_preimage(op).count("DAG_NODE_T_") >= 1 for op in task_ops) + # dash_free: task names contain no dash and not the reserved token. + dash_free = all( + _NAME_CHARSET.fullmatch(op.name) is not None + and "DAG_NODE_T_" not in op.name + for op in task_ops + ) + # disjoint_from_counter: task ids never collide with counter ids. + disjoint = {op.operation_id for op in task_ops}.isdisjoint(counter_ids) + return { + "name_based": name_based, + "has_delimiter": has_delimiter, + "dash_free": dash_free, + "disjoint_from_counter": disjoint, + } + + +def _normalize_result(value: Any) -> Any: + """Serialize a task result to a language-neutral JSON value. + + Nested-``dag`` results collapse to ``{completion_reason, counts}`` (never + raw sub-task values), per catalog Part B.2. + """ + if isinstance(value, DagResult): + return { + "completion_reason": value.completion_reason.value, + "counts": { + "success": value.success_count, + "failure": value.failure_count, + "skipped": value.skipped_count, + "total": value.total_count, + }, + } + return value + + +def _task_record(te: Any) -> dict[str, Any]: + status = te.status.value + result_val: Any = None + error_type: str | None = None + skip_reason: str | None = te.skip_reason.value if te.skip_reason else None + if te.status is TaskStatus.SUCCEEDED: + result_val = _normalize_result(te.result) + elif te.status is TaskStatus.FAILED: + # Native step failure (CallableRuntimeError/RuntimeError) -> normalized. + error_type = "StepError" + return { + "status": status, + "result": result_val, + "error_type": error_type, + "skip_reason": skip_reason, + } + + +def _record(scenario: str, result: Any, client: InMemoryServiceClient) -> dict[str, Any]: + tasks = {name: _task_record(te) for name, te in result.results.items()} + rec = { + "scenario": scenario, + "tasks": tasks, + "completion_reason": result.completion_reason.value, + "counts": { + "success": result.success_count, + "failure": result.failure_count, + "skipped": result.skipped_count, + "total": result.total_count, + }, + "structural_id_checks": _structural_checks(client), + "validation_error": None, + } + RECORDS[scenario] = rec + return rec + + +def _validation_record(scenario: str, token: str) -> dict[str, Any]: + rec = { + "scenario": scenario, + "tasks": {}, + "completion_reason": None, + "counts": {"success": 0, "failure": 0, "skipped": 0, "total": 0}, + "structural_id_checks": { + "name_based": False, + "has_delimiter": False, + "dash_free": False, + "disjoint_from_counter": False, + }, + "validation_error": token, + } + RECORDS[scenario] = rec + return rec + + +def _status(rec: dict[str, Any], name: str) -> str | None: + t = rec["tasks"].get(name) + return t["status"] if t else None + + +# endregion helpers + + +# region scenarios +def test_dag_1_diamond() -> None: + def reg(d: Any) -> None: + fetch = d.step(lambda deps, sc: 10, name="fetch") + ta = d.step(lambda deps, sc: deps["fetch"] + 1, deps=[fetch], name="ta") + tb = d.step(lambda deps, sc: deps["fetch"] * 2, deps=[fetch], name="tb") + d.step(lambda deps, sc: deps["ta"] + deps["tb"], deps=[ta, tb], name="merge") + + result, client = _run(reg, "DAG-1") + rec = _record("DAG-1", result, client) + assert rec["tasks"]["fetch"]["result"] == 10 + assert rec["tasks"]["ta"]["result"] == 11 + assert rec["tasks"]["tb"]["result"] == 20 + assert rec["tasks"]["merge"]["result"] == 31 + assert rec["completion_reason"] == "ALL_COMPLETED" + assert rec["counts"] == {"success": 4, "failure": 0, "skipped": 0, "total": 4} + assert all(rec["structural_id_checks"].values()) + + +def _compensation(scenario: str, charge_ok: bool) -> dict[str, Any]: + def reg(d: Any) -> None: + if charge_ok: + c = d.step(lambda deps, sc: "charged", name="charge") + else: + c = d.step(_fail, name="charge", config=NO_RETRY_CFG) + d.step(lambda deps, sc: "fulfilled", name="fulfill").after(c) + d.step(lambda deps, sc: "refunded", name="refund").after(c).trigger_rule( + TriggerRule.ALL_FAILED + ) + d.step(lambda deps, sc: "audited", name="audit").after(c).trigger_rule( + TriggerRule.ALL_DONE + ) + + result, client = _run(reg, scenario) + return _record(scenario, result, client) + + +def test_dag_2_compensation_charge_fails() -> None: + rec = _compensation("DAG-2", charge_ok=False) + assert _status(rec, "charge") == "FAILED" + assert rec["tasks"]["charge"]["error_type"] == "StepError" + assert _status(rec, "fulfill") == "SKIPPED" + assert rec["tasks"]["fulfill"]["skip_reason"] == "TRIGGER_RULE" + assert rec["tasks"]["refund"]["result"] == "refunded" + assert rec["tasks"]["audit"]["result"] == "audited" + assert rec["completion_reason"] == "COMPLETED_WITH_FAILURES" + assert rec["counts"] == {"success": 2, "failure": 1, "skipped": 1, "total": 4} + + +def test_dag_3_compensation_charge_succeeds() -> None: + rec = _compensation("DAG-3", charge_ok=True) + assert rec["tasks"]["charge"]["result"] == "charged" + assert rec["tasks"]["fulfill"]["result"] == "fulfilled" + assert _status(rec, "refund") == "SKIPPED" + assert rec["tasks"]["refund"]["skip_reason"] == "TRIGGER_RULE" + assert rec["tasks"]["audit"]["result"] == "audited" + assert rec["completion_reason"] == "ALL_COMPLETED" + assert rec["counts"] == {"success": 3, "failure": 0, "skipped": 1, "total": 4} + + +def test_dag_4_run_if_branching() -> None: + def reg(d: Any) -> None: + classify = d.step(lambda deps, sc: "review", name="classify") + d.step( + lambda deps, sc: "published", + deps=[classify], + name="publish", + run_if=lambda deps: deps["classify"] == "publish", + ) + d.step( + lambda deps, sc: "reviewed", + deps=[classify], + name="review", + run_if=lambda deps: deps["classify"] == "review", + ) + d.step( + lambda deps, sc: "blocked", + deps=[classify], + name="block", + run_if=lambda deps: deps["classify"] == "block", + ) + + result, client = _run(reg, "DAG-4") + rec = _record("DAG-4", result, client) + assert rec["tasks"]["classify"]["result"] == "review" + assert rec["tasks"]["review"]["result"] == "reviewed" + assert _status(rec, "publish") == "SKIPPED" + assert rec["tasks"]["publish"]["skip_reason"] == "RUN_IF_PREDICATE" + assert _status(rec, "block") == "SKIPPED" + assert rec["tasks"]["block"]["skip_reason"] == "RUN_IF_PREDICATE" + assert rec["counts"] == {"success": 2, "failure": 0, "skipped": 2, "total": 4} + assert rec["completion_reason"] == "ALL_COMPLETED" + + +def test_dag_5_trigger_matrix_empty_upstream() -> None: + rules = [ + ("r_all_success", TriggerRule.ALL_SUCCESS, "SUCCEEDED"), + ("r_all_failed", TriggerRule.ALL_FAILED, "SKIPPED"), + ("r_all_done", TriggerRule.ALL_DONE, "SUCCEEDED"), + ("r_one_success", TriggerRule.ANY_SUCCESS, "SKIPPED"), + ("r_one_failed", TriggerRule.ANY_FAILED, "SKIPPED"), + ("r_none_failed", TriggerRule.NONE_FAILED, "SUCCEEDED"), + ] + + def reg(d: Any) -> None: + for name, rule, _ in rules: + d.step(lambda deps, sc: "ok", name=name).trigger_rule(rule) + + result, client = _run(reg, "DAG-5") + rec = _record("DAG-5", result, client) + for name, _, expected in rules: + assert _status(rec, name) == expected, name + if expected == "SUCCEEDED": + assert rec["tasks"][name]["result"] == "ok" + else: + assert rec["tasks"][name]["skip_reason"] == "TRIGGER_RULE" + assert rec["counts"] == {"success": 3, "failure": 0, "skipped": 3, "total": 6} + assert rec["completion_reason"] == "ALL_COMPLETED" + + +def test_dag_6_trigger_matrix_mixed() -> None: + consumers = [ + ("c_all_success", TriggerRule.ALL_SUCCESS, "SKIPPED"), + ("c_all_failed", TriggerRule.ALL_FAILED, "SKIPPED"), + ("c_all_done", TriggerRule.ALL_DONE, "SUCCEEDED"), + ("c_one_success", TriggerRule.ANY_SUCCESS, "SUCCEEDED"), + ("c_one_failed", TriggerRule.ANY_FAILED, "SUCCEEDED"), + ("c_none_failed", TriggerRule.NONE_FAILED, "SKIPPED"), + ] + + def reg(d: Any) -> None: + up_ok = d.step(lambda deps, sc: "ok", name="up_ok") + up_fail = d.step(_fail, name="up_fail", config=NO_RETRY_CFG) + for name, rule, _ in consumers: + d.step(lambda deps, sc: "c", name=name).after(up_ok, up_fail).trigger_rule( + rule + ) + + result, client = _run(reg, "DAG-6") + rec = _record("DAG-6", result, client) + assert rec["tasks"]["up_ok"]["result"] == "ok" + assert _status(rec, "up_fail") == "FAILED" + for name, _, expected in consumers: + assert _status(rec, name) == expected, name + assert rec["counts"] == {"success": 4, "failure": 1, "skipped": 3, "total": 8} + assert rec["completion_reason"] == "COMPLETED_WITH_FAILURES" + + +def test_dag_7_trigger_matrix_all_failed() -> None: + consumers = [ + ("k_all_success", TriggerRule.ALL_SUCCESS, "SKIPPED"), + ("k_all_failed", TriggerRule.ALL_FAILED, "SUCCEEDED"), + ("k_all_done", TriggerRule.ALL_DONE, "SUCCEEDED"), + ("k_one_success", TriggerRule.ANY_SUCCESS, "SKIPPED"), + ("k_one_failed", TriggerRule.ANY_FAILED, "SUCCEEDED"), + ("k_none_failed", TriggerRule.NONE_FAILED, "SKIPPED"), + ] + + def reg(d: Any) -> None: + u1 = d.step(_fail, name="u1", config=NO_RETRY_CFG) + u2 = d.step(_fail, name="u2", config=NO_RETRY_CFG) + for name, rule, _ in consumers: + d.step(lambda deps, sc: "k", name=name).after(u1, u2).trigger_rule(rule) + + result, client = _run(reg, "DAG-7") + rec = _record("DAG-7", result, client) + assert _status(rec, "u1") == "FAILED" + assert _status(rec, "u2") == "FAILED" + for name, _, expected in consumers: + assert _status(rec, name) == expected, name + assert rec["counts"] == {"success": 3, "failure": 2, "skipped": 3, "total": 8} + assert rec["completion_reason"] == "COMPLETED_WITH_FAILURES" + + +def test_dag_8_skip_cascade() -> None: + def reg(d: Any) -> None: + seed = d.step(lambda deps, sc: 1, name="seed") + gate = d.step( + lambda deps, sc: "gate", + deps=[seed], + name="gate", + run_if=lambda deps: deps["seed"] > 100, + ) + d1 = d.step(lambda deps, sc: "d1", deps=[gate], name="d1") + d.step(lambda deps, sc: "d2", deps=[d1], name="d2") + d.step(lambda deps, sc: "sink", name="sink").after(gate).trigger_rule( + TriggerRule.ALL_DONE + ) + + result, client = _run(reg, "DAG-8") + rec = _record("DAG-8", result, client) + assert rec["tasks"]["seed"]["result"] == 1 + assert _status(rec, "gate") == "SKIPPED" + assert rec["tasks"]["gate"]["skip_reason"] == "RUN_IF_PREDICATE" + assert rec["tasks"]["d1"]["skip_reason"] == "TRIGGER_RULE" + assert rec["tasks"]["d2"]["skip_reason"] == "TRIGGER_RULE" + assert rec["tasks"]["sink"]["result"] == "sink" + assert rec["counts"] == {"success": 2, "failure": 0, "skipped": 3, "total": 5} + assert rec["completion_reason"] == "ALL_COMPLETED" + + +def test_dag_9_nested_dag() -> None: + def inner(d: Any) -> None: + x = d.step(lambda deps, sc: 3, name="x") + d.step(lambda deps, sc: deps["x"] * 10, deps=[x], name="y") + + def outer(d: Any) -> None: + a = d.step(lambda deps, sc: 2, name="a") + inn = d.dag(inner, deps=[a], name="inner") + d.step( + lambda deps, sc: deps["inner"].get_result("y") + 5, + deps=[inn], + name="consume", + ) + + result, client = _run(outer, "DAG-9") + # scope isolation: inner task names invisible in outer scope + assert result.get_status("x") is None + assert result.get_status("y") is None + rec = _record("DAG-9", result, client) + assert rec["tasks"]["a"]["result"] == 2 + assert rec["tasks"]["consume"]["result"] == 35 + assert rec["tasks"]["inner"]["result"] == { + "completion_reason": "ALL_COMPLETED", + "counts": {"success": 2, "failure": 0, "skipped": 0, "total": 2}, + } + assert rec["counts"] == {"success": 3, "failure": 0, "skipped": 0, "total": 3} + assert rec["completion_reason"] == "ALL_COMPLETED" + + +def test_dag_10_empty() -> None: + result, client = _run(lambda d: None, "DAG-10") + rec = _record("DAG-10", result, client) + assert rec["tasks"] == {} + assert rec["counts"] == {"success": 0, "failure": 0, "skipped": 0, "total": 0} + assert rec["completion_reason"] == "ALL_COMPLETED" + assert all(rec["structural_id_checks"].values()) + + +def test_dag_11_cycle() -> None: + def reg(d: Any) -> None: + p = d.step(lambda deps, sc: 1, name="p") + q = d.step(lambda deps, sc: 1, deps=[p], name="q") + p.after(q) + + with pytest.raises(exceptions.DagCyclicDependencyError): + _run(reg, "DAG-11") + _validation_record("DAG-11", "DagCyclicDependencyError") + + +def test_dag_12_duplicate() -> None: + def reg(d: Any) -> None: + d.step(lambda deps, sc: 1, name="dup") + d.step(lambda deps, sc: 2, name="dup") + + with pytest.raises(exceptions.DagDuplicateTaskError): + _run(reg, "DAG-12") + _validation_record("DAG-12", "DagDuplicateTaskError") + + +def test_dag_13_invalid_name_dash() -> None: + def reg(d: Any) -> None: + d.step(lambda deps, sc: 1, name="fetch-data") + + with pytest.raises(exceptions.DagInvalidTaskNameError): + _run(reg, "DAG-13") + _validation_record("DAG-13", "DagInvalidTaskNameError") + + +def test_dag_14_invalid_name_reserved() -> None: + def reg(d: Any) -> None: + d.step(lambda deps, sc: 1, name="DAG_NODE_T_root") + + with pytest.raises(exceptions.DagInvalidTaskNameError): + _run(reg, "DAG-14") + _validation_record("DAG-14", "DagInvalidTaskNameError") + + +def test_dag_15_foreign_dep() -> None: + captured: dict[str, Any] = {} + + def sibling(d: Any) -> None: + captured["h"] = d.step(lambda deps, sc: 1, name="foreign") + + _run(sibling, "DAG-15-sibling") + + def reg(d: Any) -> None: + d.step(lambda deps, sc: 1, deps=[captured["h"]], name="t") + + with pytest.raises(exceptions.DagInvalidDependencyError): + _run(reg, "DAG-15") + _validation_record("DAG-15", "DagInvalidDependencyError") + + +def test_dag_16_min_successful() -> None: + def reg(d: Any) -> None: + prev = None + for i in range(1, 6): + deps = [prev] if prev is not None else None + prev = d.step( + (lambda i: lambda deps, sc: i)(i), deps=deps, name=f"s{i}" + ) + + result, client = _run( + reg, + "DAG-16", + DagConfig(max_concurrency=1, completion_config=CompletionConfig(min_successful=3)), + ) + rec = _record("DAG-16", result, client) + assert rec["tasks"]["s1"]["result"] == 1 + assert rec["tasks"]["s2"]["result"] == 2 + assert rec["tasks"]["s3"]["result"] == 3 + assert "s4" not in rec["tasks"] and "s5" not in rec["tasks"] # absent + assert rec["completion_reason"] == "MIN_SUCCESSFUL_REACHED" + assert rec["counts"]["success"] == 3 + assert rec["counts"]["failure"] == 0 + assert rec["counts"]["skipped"] == 0 + # total_count = registered task count (spec §2.8): 5 tasks registered even + # though s4/s5 never started (absent from results) under min_successful. + assert rec["counts"]["total"] == 5 + + +def test_dag_17_tolerated_failures() -> None: + def reg(d: Any) -> None: + prev = None + for i in range(1, 5): + deps = [prev] if prev is not None else None + h = d.step(_fail, deps=deps, name=f"t{i}", config=NO_RETRY_CFG) + if i > 1: + h.trigger_rule(TriggerRule.ALL_DONE) + prev = h + + result, client = _run( + reg, + "DAG-17", + DagConfig( + max_concurrency=1, + completion_config=CompletionConfig(tolerated_failure_count=1), + ), + ) + rec = _record("DAG-17", result, client) + assert _status(rec, "t1") == "FAILED" + assert _status(rec, "t2") == "FAILED" + assert "t3" not in rec["tasks"] and "t4" not in rec["tasks"] # absent + assert rec["completion_reason"] == "FAILURE_TOLERANCE_EXCEEDED" + assert rec["counts"]["success"] == 0 + assert rec["counts"]["failure"] == 2 + assert rec["counts"]["skipped"] == 0 + # total_count = registered task count (spec §2.8): 4 tasks registered even + # though t3/t4 never started (absent) under tolerated_failure_count. + assert rec["counts"]["total"] == 4 + + +def test_dag_18_custom_completion() -> None: + def should_complete(status: Any) -> Any: + rejected = any( + item.status is TaskStatus.SUCCEEDED + and isinstance(item.result, dict) + and item.result.get("verdict") == "REJECT" + for item in status.items + ) + return ( + complete_dag(DagCompletionOutcome.FAILED) if rejected else continue_dag() + ) + + def reg(d: Any) -> None: + r1 = d.step(lambda deps, sc: {"verdict": "ACCEPT"}, name="r1") + r2 = d.step(lambda deps, sc: {"verdict": "REJECT"}, deps=[r1], name="r2") + d.step(lambda deps, sc: {"verdict": "ACCEPT"}, deps=[r2], name="r3") + + result, client = _run( + reg, + "DAG-18", + DagConfig( + max_concurrency=1, + completion_config=DagCustomCompletionConfig(should_complete), + ), + ) + rec = _record("DAG-18", result, client) + assert rec["tasks"]["r1"]["result"] == {"verdict": "ACCEPT"} + assert rec["tasks"]["r2"]["result"] == {"verdict": "REJECT"} + assert "r3" not in rec["tasks"] # absent: never started + assert rec["completion_reason"] == "CUSTOM_COMPLETION_FAILED" + assert rec["counts"] == {"success": 2, "failure": 0, "skipped": 0, "total": 3} + with pytest.raises(exceptions.DagExecutionError): + result.throw_if_error() + + +def test_dag_19_order_independence() -> None: + def make_reg(swap: bool) -> Any: + def reg(d: Any) -> None: + root = d.step(lambda deps, sc: 100, name="root") + if swap: + c = d.step(lambda deps, sc: deps["root"] + 2, deps=[root], name="c") + b = d.step(lambda deps, sc: deps["root"] + 1, deps=[root], name="b") + else: + b = d.step(lambda deps, sc: deps["root"] + 1, deps=[root], name="b") + c = d.step(lambda deps, sc: deps["root"] + 2, deps=[root], name="c") + d.step(lambda deps, sc: deps["b"] + deps["c"], deps=[b, c], name="merge") + + return reg + + r1, c1 = _run(make_reg(False), "DAG-19") + rec1 = _record("DAG-19", r1, c1) + r2, c2 = _run(make_reg(True), "DAG-19-swapped") + rec2 = dict(rec1) # keep DAG-19 as the emitted record + rec2_actual = { + "scenario": "DAG-19", + "tasks": {name: _task_record(te) for name, te in r2.results.items()}, + "completion_reason": r2.completion_reason.value, + "counts": { + "success": r2.success_count, + "failure": r2.failure_count, + "skipped": r2.skipped_count, + "total": r2.total_count, + }, + "structural_id_checks": _structural_checks(c2), + "validation_error": None, + } + RECORDS.pop("DAG-19-swapped", None) + # order-independence: both completion orders yield identical records + assert rec1 == rec2_actual + assert rec2["tasks"]["merge"]["result"] == 203 + assert rec1["tasks"]["root"]["result"] == 100 + assert rec1["tasks"]["b"]["result"] == 101 + assert rec1["tasks"]["c"]["result"] == 102 + assert rec1["counts"] == {"success": 4, "failure": 0, "skipped": 0, "total": 4} + assert rec1["completion_reason"] == "ALL_COMPLETED" + + +def test_zz_emit_python_json() -> None: + """Runs last (name-sorted): assemble, validate, and write python.json. + + Emits exactly 19 records (DAG-1..19) as key-sorted UTF-8 JSON with 2-space + indent + trailing newline. + """ + expected_scenarios = {f"DAG-{i}" for i in range(1, 20)} + assert set(RECORDS) == expected_scenarios, ( + f"missing/extra: {expected_scenarios ^ set(RECORDS)}" + ) + assert len(RECORDS) == 19 + + OUT_PATH.parent.mkdir(parents=True, exist_ok=True) + text = json.dumps(RECORDS, indent=2, sort_keys=True, ensure_ascii=False) + "\n" + OUT_PATH.write_text(text, encoding="utf-8") + + # round-trip + byte-shape sanity + reloaded = json.loads(OUT_PATH.read_text(encoding="utf-8")) + assert reloaded == RECORDS + assert text.endswith("\n") + + +# endregion scenarios diff --git a/packages/aws-durable-execution-sdk-python/tests/dag_support.py b/packages/aws-durable-execution-sdk-python/tests/dag_support.py new file mode 100644 index 00000000..ea4116d0 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/dag_support.py @@ -0,0 +1,169 @@ +"""Shared test support: an in-memory DurableServiceClient that actually stores +operations, so checkpoint fast paths / replay behave realistically for DAG tests. +""" + +from __future__ import annotations + +import datetime +import threading +from typing import TYPE_CHECKING + +from aws_durable_execution_sdk_python.context import DurableContext, ExecutionContext +from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, + ChainedInvokeDetails, + CheckpointOutput, + CheckpointUpdatedExecutionState, + ContextDetails, + Operation, + OperationAction, + OperationStatus, + OperationType, + StepDetails, + WaitDetails, +) +from aws_durable_execution_sdk_python.plugin import PluginExecutor +from aws_durable_execution_sdk_python.state import ExecutionState + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python.lambda_service import OperationUpdate + +_ACTION_STATUS = { + OperationAction.START: OperationStatus.STARTED, + OperationAction.SUCCEED: OperationStatus.SUCCEEDED, + OperationAction.FAIL: OperationStatus.FAILED, + OperationAction.RETRY: OperationStatus.STARTED, + OperationAction.CANCEL: OperationStatus.CANCELLED, +} + + +class InMemoryServiceClient: + """A minimal in-memory backend that persists operations by id. + + Enough fidelity to exercise checkpoint fast paths and replay: converts each + ``OperationUpdate`` into a stored ``Operation`` and returns the full set on + every checkpoint so ``ExecutionState`` reflects completion. + """ + + def __init__(self) -> None: + self.operations: dict[str, Operation] = {} + self._lock = threading.Lock() + self.checkpoint_count = 0 + + def _to_operation(self, update: OperationUpdate) -> Operation: + status = _ACTION_STATUS[update.action] + prev = self.operations.get(update.operation_id) + step_details = None + context_details = None + chained_invoke_details = None + wait_details = None + callback_details = None + + if update.operation_type is OperationType.STEP: + attempt = 0 + if prev and prev.step_details: + attempt = prev.step_details.attempt + if update.action is OperationAction.RETRY: + attempt += 1 + step_details = StepDetails( + attempt=attempt, result=update.payload, error=update.error + ) + elif update.operation_type is OperationType.CONTEXT: + replay = ( + update.context_options.replay_children + if update.context_options + else False + ) + context_details = ContextDetails( + replay_children=replay, result=update.payload, error=update.error + ) + elif update.operation_type is OperationType.CHAINED_INVOKE: + chained_invoke_details = ChainedInvokeDetails( + result=update.payload, + error=update.error, + ) + elif update.operation_type is OperationType.WAIT: + wait_details = WaitDetails() + elif update.operation_type is OperationType.CALLBACK: + callback_details = CallbackDetails( + callback_id=update.operation_id, + result=update.payload, + error=update.error, + ) + + return Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=status, + parent_id=update.parent_id, + name=update.name, + sub_type=update.sub_type, + step_details=step_details, + context_details=context_details, + chained_invoke_details=chained_invoke_details, + wait_details=wait_details, + callback_details=callback_details, + ) + + def checkpoint( + self, durable_execution_arn, checkpoint_token, updates, client_token + ) -> CheckpointOutput: + with self._lock: + self.checkpoint_count += 1 + for update in updates: + self.operations[update.operation_id] = self._to_operation(update) + all_ops = list(self.operations.values()) + return CheckpointOutput( + checkpoint_token="token", + new_execution_state=CheckpointUpdatedExecutionState( + operations=all_ops, next_marker=None + ), + ) + + def get_execution_state( + self, durable_execution_arn, checkpoint_token, next_marker, max_items=1000 + ): # pragma: no cover - not used by these tests + raise NotImplementedError + + def stop(self, execution_arn, payload) -> datetime.datetime: # pragma: no cover + return datetime.datetime.now(tz=datetime.UTC) + + +def make_state( + client: InMemoryServiceClient | None = None, +) -> tuple[ExecutionState, InMemoryServiceClient]: + """Build a fresh ExecutionState wired to an in-memory client. + + Starts the background checkpoint-processing thread exactly like the real + runtime (``execution.py`` submits ``checkpoint_batches_forever`` to its + executor). Without this thread every synchronous checkpoint blocks forever + on its completion event — which is the deadlock DAG tasks hit when they run + their operations on the executor's worker threads. The thread is a daemon so + it never blocks interpreter exit; it idles on the queue between tests. + """ + client = client or InMemoryServiceClient() + state = ExecutionState( + durable_execution_arn="test-arn", + initial_checkpoint_token="token", # noqa: S106 + operations={}, + service_client=client, + plugin_executor=PluginExecutor(plugins=None), + ) + checkpoint_thread = threading.Thread( + target=state.checkpoint_batches_forever, + name="dag-test-checkpointer", + daemon=True, + ) + checkpoint_thread.start() + return state, client + + +def make_context(state: ExecutionState, parent_id: str | None = None) -> DurableContext: + """Build a DurableContext over the given state.""" + return DurableContext( + state=state, + execution_context=ExecutionContext( + durable_execution_arn=state.durable_execution_arn + ), + parent_id=parent_id, + ) diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_concurrency_coverage_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_concurrency_coverage_test.py new file mode 100644 index 00000000..f40c5bef --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_concurrency_coverage_test.py @@ -0,0 +1,343 @@ +"""Integration coverage for the concurrency + abort gaps (conformance 10-12..10-14). + +These run the *exact* graphs of the three new conformance scenarios in-process so +the regressions the cloud suite cannot catch are caught here without a deploy: + +* **10-13 concurrent overlap** — asserts the result AND that peak observed + concurrency was >= 2 AND that every task's recorded operation id is name-based + (derived from the ``DAG_NODE_T_`` pre-image). The id assertion is the one + the cloud suite deliberately cannot make: at unset ``max_concurrency`` with + out-of-order completion a counter-based regression would hand out different ids + on replay and terminate the execution, but only a per-SDK test can prove the + ids are *positively* name-based rather than merely internally consistent. Runs + on the lightweight in-memory harness (``tests/dag_support``). +* **10-14 inverted readiness** — drives the graph through a multi-invocation + replay loop on the in-memory harness: the two waits suspend on the first + invocation, then a driver marks each wait's checkpoint SUCCEEDED in resume + order (fast, then slow) and re-invokes on a state seeded from the accumulated + operations, simulating the platform resuming each timer. It asserts the result + AND that the downstream pair ran in the reverse of registration order across + the suspend (afterFast before afterSlow) with no replay-consistency error. +* **10-12 abort** — asserts the typed ``DagPredicateError`` surfaced to the caller + and that the ``ALL_FAILED`` compensation body was never invoked (external + counter). Runs on the in-memory harness. + +Python bounds every backend operation id through ``blake2b(...)[:64]``, so the +``DAG_NODE_T_`` token lives in the id's *pre-image*, not the digest. The id +helper recomputes the digest from ``(parent_id, name)`` and asserts the recorded +id equals it — proving, per task, that the id is the name-based one and could +never be produced by a counter scheme. This mirrors the ``_structural_checks`` +technique in ``tests/conformance/dag_conformance_test.py``. +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import signal +import threading +import time +from contextlib import contextmanager +from typing import Any + +import pytest + +from aws_durable_execution_sdk_python.context import DurableContext, ExecutionContext +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.dag import ( + DagCompletionReason, + DagConfig, + TriggerRule, +) +from aws_durable_execution_sdk_python.exceptions import ( + DagPredicateError, + SuspendExecution, +) +from aws_durable_execution_sdk_python.lambda_service import ( + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import PluginExecutor +from aws_durable_execution_sdk_python.state import ExecutionState +from tests.dag_support import InMemoryServiceClient, make_context, make_state + + +@contextmanager +def _fail_on_hang(seconds: int = 30): + """Turn a scheduler hang into an assertion failure rather than blocking the + whole test session. SIGALRM fires on the main thread (where pytest runs).""" + + def _handler(_signum, _frame): + raise AssertionError("context.dag() hung (concurrency/abort regression)") + + old = signal.signal(signal.SIGALRM, _handler) + signal.alarm(seconds) + try: + yield + finally: + signal.alarm(0) + signal.signal(signal.SIGALRM, old) + + +def _name_based_task_ops(client: InMemoryServiceClient) -> dict[str, Any]: + """Return recorded ops whose id is exactly the name-based DAG-task digest. + + A DAG task's backend id is ``blake2b(f"{parent}-DAG_NODE_T_{name}")[:64]``. + Any op whose id equals that recomputation from its own ``(parent_id, name)`` + is a name-based task op; counter ops (e.g. the top-level DAG container) never + match, so this positively isolates the task ops. + """ + by_name: dict[str, Any] = {} + for op in client.operations.values(): + if op.name is None: + continue + preimage = ( + f"{op.parent_id}-DAG_NODE_T_{op.name}" + if op.parent_id + else f"DAG_NODE_T_{op.name}" + ) + expected = hashlib.blake2b(preimage.encode()).hexdigest()[:64] + if op.operation_id == expected: + by_name[op.name] = op + return by_name + + +def _assert_ids_name_based(client: InMemoryServiceClient, names: set[str]) -> None: + task_ops = _name_based_task_ops(client) + for name in names: + assert name in task_ops, f"no name-based op recorded for task {name!r}" + op = task_ops[name] + preimage = ( + f"{op.parent_id}-DAG_NODE_T_{op.name}" + if op.parent_id + else f"DAG_NODE_T_{op.name}" + ) + # The DAG_NODE_T_ segment lives in the id's PRE-IMAGE; asserting + # the recorded digest equals the hash of that pre-image proves the id is + # name-based for THIS task and could not come from a counter scheme. + assert f"DAG_NODE_T_{name}" in preimage + assert op.operation_id == hashlib.blake2b(preimage.encode()).hexdigest()[:64] + + +_OVERLAP_TASKS = {"root", "slow", "fast", "afterSlow", "afterFast", "merge"} + + +def test_10_13_concurrent_overlap() -> None: + """The 10-13 graph: slow + fast overlap inside one invocation, fast finishes + first (inverting registration order), and every task id is name-based.""" + tracker = {"current": 0, "peak": 0} + lock = threading.Lock() + # Both slow and fast must be simultaneously in-flight to clear this barrier; + # if the scheduler serialized them it would time out -> BrokenBarrierError, + # failing the task and this test. This makes the overlap *deterministic* + # rather than relying on sleep timing. + barrier = threading.Barrier(2, timeout=10) + + def _enter() -> None: + with lock: + tracker["current"] += 1 + tracker["peak"] = max(tracker["peak"], tracker["current"]) + + def _leave() -> None: + with lock: + tracker["current"] -= 1 + + def slow(_deps: Any, _sc: Any) -> str: + _enter() + try: + barrier.wait() + time.sleep(0.2) # ensure fast completes first -> out-of-order + return "S" + finally: + _leave() + + def fast(_deps: Any, _sc: Any) -> str: + _enter() + try: + barrier.wait() + return "F" + finally: + _leave() + + def register(d: Any) -> None: + root = d.step(lambda deps, sc: 1, name="root") + slow_h = d.step(slow, deps=[root], name="slow") # registered FIRST + fast_h = d.step(fast, deps=[root], name="fast") + after_slow = d.step( # registered FIRST + lambda deps, sc: deps[slow_h] + "s", deps=[slow_h], name="afterSlow" + ) + after_fast = d.step( + lambda deps, sc: deps[fast_h] + "f", deps=[fast_h], name="afterFast" + ) + d.step( + lambda deps, sc: deps[after_slow] + deps[after_fast], + deps=[after_slow, after_fast], + name="merge", + ) + + state, client = make_state() + with _fail_on_hang(): + result = make_context(state).dag(register, name="overlapdag") + + assert result.get_result("merge") == "SsFf" + assert result.completion_reason is DagCompletionReason.ALL_COMPLETED + assert ( + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ) == (6, 0, 0, 6) + # Genuine overlap: both bodies were simultaneously in-flight. + assert tracker["peak"] >= 2 + _assert_ids_name_based(client, _OVERLAP_TASKS) + + +def _seeded_state( + client: InMemoryServiceClient, operations: dict[str, Any] +) -> ExecutionState: + """A fresh ExecutionState seeded with prior operations, wired to ``client``. + + Mirrors ``dag_support.make_state`` but seeds the operations a real + re-invocation would have fetched from the backend, and starts the background + checkpoint thread so synchronous checkpoints don't deadlock. + """ + state = ExecutionState( + durable_execution_arn="test-arn", + initial_checkpoint_token="token", # noqa: S106 + operations=operations, + service_client=client, + plugin_executor=PluginExecutor(plugins=None), + ) + thread = threading.Thread( + target=state.checkpoint_batches_forever, + name="dag-suspend-checkpointer", + daemon=True, + ) + thread.start() + return state + + +def _complete_wait(client: InMemoryServiceClient, name: str) -> bool: + """Mark the wait op with ``name`` SUCCEEDED, simulating its timer firing.""" + for op_id, op in list(client.operations.items()): + if op.name == name and op.operation_type is OperationType.WAIT: + client.operations[op_id] = dataclasses.replace( + op, status=OperationStatus.SUCCEEDED + ) + return True + return False + + +def test_10_14_inverted_readiness_across_suspend() -> None: + """The 10-14 graph across a real suspend: two waits are in flight when the + invocation suspends; resuming fast (2s) before slow (8s) makes afterFast ready + an invocation before afterSlow. The downstream pair therefore starts in the + reverse of registration order across invocations, and the run completes with + merge == "SF" and no replay-consistency error.""" + run_order: list[str] = [] + order_lock = threading.Lock() + + def _record(name: str) -> str: + with order_lock: + run_order.append(name) + return "S" if name == "afterSlow" else "F" + + def register(d: Any) -> None: + root = d.step(lambda deps, sc: 1, name="root") + slow = d.wait(Duration.from_seconds(8), deps=[root], name="slow") # registered FIRST + fast = d.wait(Duration.from_seconds(2), deps=[root], name="fast") + after_slow = d.step(lambda deps, sc: _record("afterSlow"), name="afterSlow") + after_slow.after(slow) # registered FIRST + after_fast = d.step(lambda deps, sc: _record("afterFast"), name="afterFast") + after_fast.after(fast) + d.step( + lambda deps, sc: deps[after_slow] + deps[after_fast], + deps=[after_slow, after_fast], + name="merge", + ) + + client = InMemoryServiceClient() + # fast's timer (2s) fires before slow's (8s), so the platform resumes fast + # first; the driver mirrors that resume order. + resume_order = ["fast", "slow"] + result = None + with _fail_on_hang(): + for _ in range(len(resume_order) + 1): + state = _seeded_state(client, dict(client.operations)) + ctx = DurableContext( + state=state, + execution_context=ExecutionContext( + durable_execution_arn=state.durable_execution_arn + ), + parent_id=None, + ) + try: + result = ctx.dag(register, name="suspenddag") + break + except SuspendExecution: + assert resume_order, "suspended more times than there are waits" + assert _complete_wait(client, resume_order.pop(0)) + + assert result is not None, "DAG never completed across resumes" + assert result.get_result("merge") == "SF" + assert result.completion_reason is DagCompletionReason.ALL_COMPLETED + assert ( + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ) == (6, 0, 0, 6) + # Inverted readiness across the suspend: afterFast ran an invocation before + # afterSlow, i.e. the reverse of their registration order. A counter-based id + # scheme could not survive this out-of-order resume without a replay error. + assert run_order == ["afterFast", "afterSlow"] + + +def test_10_12_run_if_abort() -> None: + """The 10-12 graph: a throwing run_if aborts the DAG with a typed error and + the ALL_FAILED compensation body is never invoked. + + Through the ``context.dag()`` child-context boundary the caller observes a + ``DagPredicateError`` whose *message* names the offending task and embeds the + original error; ``task_name`` / ``__cause__`` are intentionally not + reconstructed across that boundary (so the first run matches replay — the + executor-level richness is covered by ``dag_executor_test``). We assert + exactly what the caller observes. + """ + calls = {"guarded_body": 0, "refund_body": 0} + calls_lock = threading.Lock() + + def _bump(key: str) -> None: + with calls_lock: + calls[key] += 1 + + def guarded_body(_deps: Any, _sc: Any) -> str: + _bump("guarded_body") # MUST NOT happen + return "ran" + + def refund_body(_deps: Any, _sc: Any) -> str: + _bump("refund_body") # MUST NOT happen + return "refunded" + + def boom(_deps: Any) -> bool: + raise RuntimeError("predicate boom") + + def register(d: Any) -> None: + gate = d.step(lambda deps, sc: 1, name="gate") + guarded = d.step(guarded_body, deps=[gate], name="guarded", run_if=boom) + d.step(refund_body, name="refund").after(guarded).trigger_rule( + TriggerRule.ALL_FAILED + ) + + state, _ = make_state() + with _fail_on_hang(), pytest.raises(DagPredicateError) as ei: + make_context(state).dag( + register, name="abortdag", config=DagConfig(max_concurrency=1) + ) + + # The typed error names the offending task and embeds the original cause. + assert "guarded" in str(ei.value) + assert "predicate boom" in str(ei.value) + # Neither the guarded body nor the ALL_FAILED compensation ever ran. + assert calls["guarded_body"] == 0 + assert calls["refund_body"] == 0 diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_context_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_context_test.py new file mode 100644 index 00000000..d6d0cf1a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_context_test.py @@ -0,0 +1,113 @@ +"""T3: DagContext registration, TaskHandle chaining, DepsMap access.""" + +from __future__ import annotations + +import pytest + +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.dag import ( + DagConfig, + DepsMap, + TriggerRule, +) +from aws_durable_execution_sdk_python.exceptions import ( + DagInvalidTaskNameError, + ValidationError, +) +from aws_durable_execution_sdk_python.operation.dag_context import DagContextImpl +from tests.dag_support import make_context, make_state + + +def _impl() -> DagContextImpl: + state, _ = make_state() + ctx = make_context(state, parent_id="c") + return DagContextImpl(ctx, DagConfig()) + + +def test_depsmap_string_and_handle_access(): + d = _impl() + fetch = d.step(lambda deps, sc: "x", name="fetch") + dm = DepsMap({"fetch": "result-value"}) + assert dm["fetch"] == "result-value" + assert dm[fetch] == "result-value" + assert "fetch" in dm + assert fetch in dm + assert len(dm) == 1 + assert list(dm) == ["fetch"] + + +def test_taskhandle_hash_by_name_and_identity_eq(): + d = _impl() + a = d.step(lambda deps, sc: 1, name="a") + b = d.step(lambda deps, sc: 1, name="b") + # identity equality (eq=False) -> distinct handles are not equal + assert a is a + assert a != b + assert hash(a) == hash("a") + + +def test_after_adds_ordering_only_dep(): + d = _impl() + a = d.step(lambda deps, sc: 1, name="a") + b = d.step(lambda deps, sc: 2, name="b") + d.step(lambda deps, sc: 3, deps=[a], name="c").after(b) + tasks = d.get_tasks() + cdef = tasks["c"] + assert [h.name for h in cdef.inline_deps] == ["a"] + assert {h.name for h in cdef.all_deps} == {"a", "b"} + + +def test_trigger_rule_chaining_mutates_taskdef(): + d = _impl() + a = d.step(lambda deps, sc: 1, name="a") + d.step(lambda deps, sc: 2, deps=[a], name="b").trigger_rule(TriggerRule.ALL_DONE) + assert d.get_tasks()["b"].trigger_rule is TriggerRule.ALL_DONE + + +def test_name_resolution_from_original_name(): + from aws_durable_execution_sdk_python.context import durable_step + + d = _impl() + + @durable_step + def my_step(step_ctx): + return 1 + + h = d.step(my_step()) + assert h.name == "my_step" + + +def test_unresolvable_name_raises(): + d = _impl() + with pytest.raises(DagInvalidTaskNameError): + d.step(lambda deps, sc: 1) # bare lambda, no name + + +def test_wait_requires_name(): + d = _impl() + with pytest.raises(DagInvalidTaskNameError): + d.wait(Duration.from_seconds(5)) + + +def test_wait_requires_duration_of_at_least_one_second(): + # Matches context.wait(duration: Duration)'s own validation -- a DAG wait + # task takes the same Duration type and enforces the same floor. + d = _impl() + with pytest.raises(ValidationError, match="at least 1 second"): + d.wait(Duration.from_seconds(0), name="too-short") + + +def test_duplicate_registration_recorded_in_order(): + d = _impl() + d.step(lambda deps, sc: 1, name="dup") + d.step(lambda deps, sc: 2, name="dup") + # dict keeps last; registration order keeps both for the validator + assert len(d.get_tasks()) == 1 + assert [t.name for t in d.get_registration_order()] == ["dup", "dup"] + + +def test_invoke_registers_with_deferred_payload(): + d = _impl() + h = d.invoke("fn:prod", lambda deps: {"a": 1}, name="charge") + assert h.name == "charge" + assert d.get_tasks()["charge"].kind == "invoke" diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_executor_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_executor_test.py new file mode 100644 index 00000000..f98c676b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_executor_test.py @@ -0,0 +1,1159 @@ +"""T5: DagExecutor scheduler tests.""" + +from __future__ import annotations + +import threading +import time + +import pytest + +from aws_durable_execution_sdk_python.config import CompletionConfig, StepConfig +from aws_durable_execution_sdk_python.dag import ( + DagCompletionOutcome, + DagCompletionReason, + DagConfig, + DagCustomCompletionConfig, + SkipReason, + TaskStatus, + TriggerRule, + complete_dag, + continue_dag, +) +from aws_durable_execution_sdk_python.exceptions import ( + DagExecutionError, + ValidationError, +) +from aws_durable_execution_sdk_python.operation.dag_context import DagContextImpl +from aws_durable_execution_sdk_python.operation.dag_executor import DagExecutor +from aws_durable_execution_sdk_python.operation.dag_validator import validate_dag +from aws_durable_execution_sdk_python.retries import RetryPresets +from tests.dag_support import make_context, make_state + +NO_RETRY = RetryPresets.none() +# Per-task step config that disables retries, so an intentionally failing step +# fails promptly (attempt 1) instead of falling back to RetryPresets.default(). +NO_RETRY_CFG = StepConfig(retry_strategy=NO_RETRY) + + +def run_dag(register, config=None, parent_id="dag"): + config = config or DagConfig() + state, client = make_state() + ctx = make_context(state, parent_id=parent_id) + d = DagContextImpl(ctx, config) + register(d) + validate_dag(d) + result = DagExecutor(ctx, d.get_tasks(), config).run() + return result, client + + +def test_diamond_topological_order_and_results(): + order = [] + order_lock = threading.Lock() + + def rec(name): + with order_lock: + order.append(name) + + def register(d): + a = d.step(lambda deps, sc: (rec("a"), "A")[1], name="a") + b = d.step(lambda deps, sc: (rec("b"), deps["a"] + "B")[1], deps=[a], name="b") + c = d.step(lambda deps, sc: (rec("c"), deps["a"] + "C")[1], deps=[a], name="c") + d.step( + lambda deps, sc: (rec("d"), deps["b"] + deps["c"])[1], + deps=[b, c], + name="d", + ) + + result, _ = run_dag(register) + assert result.get_status("d") is TaskStatus.SUCCEEDED + assert result.get_result("d") == "ABAC" + assert result.success_count == 4 + assert result.completion_reason is DagCompletionReason.ALL_COMPLETED + # a before b,c before d + assert order[0] == "a" + assert order[-1] == "d" + + +def test_branches_run_concurrently(): + barrier = threading.Barrier(2, timeout=3) + both = {"ok": False} + + def branch(_deps, _sc): + try: + barrier.wait() + both["ok"] = True + except threading.BrokenBarrierError: # pragma: no cover + pass + return 1 + + def register(d): + a = d.step(lambda deps, sc: 0, name="a") + d.step(branch, deps=[a], name="b") + d.step(branch, deps=[a], name="c") + + result, _ = run_dag(register) + assert both["ok"] is True # b and c reached the barrier simultaneously + assert result.success_count == 3 + + +def test_max_concurrency_throttles(): + current = {"n": 0, "max": 0} + lock = threading.Lock() + + def slow(_deps, _sc): + with lock: + current["n"] += 1 + current["max"] = max(current["max"], current["n"]) + time.sleep(0.05) + with lock: + current["n"] -= 1 + return 1 + + def register(d): + for i in range(5): + d.step(slow, name=f"t{i}") + + result, _ = run_dag(register, DagConfig(max_concurrency=2)) + assert result.success_count == 5 + assert current["max"] <= 2 + + +def test_trigger_rule_skip_propagation(): + def boom(_deps, _sc): + raise ValueError("boom") + + def register(d): + a = d.step(boom, name="a", config=NO_RETRY_CFG) + # default ALL_SUCCESS -> skipped because a FAILED + d.step(lambda deps, sc: 1, deps=[a], name="b") + + result, _ = run_dag(register) + assert result.get_status("a") is TaskStatus.FAILED + assert result.get_status("b") is TaskStatus.SKIPPED + assert result.results["b"].skip_reason is SkipReason.TRIGGER_RULE + assert result.completion_reason is DagCompletionReason.COMPLETED_WITH_FAILURES + + +def test_compensation_all_failed_runs_on_failure(): + def charge(_deps, _sc): + raise RuntimeError("charge failed") + + def register(d): + c = d.step(charge, name="charge", config=NO_RETRY_CFG) + # refund runs when charge FAILED + d.step(lambda deps, sc: "refunded", deps=[c], name="refund").trigger_rule( + TriggerRule.ALL_FAILED + ) + # fulfill only on success -> skipped + d.step(lambda deps, sc: "fulfilled", deps=[c], name="fulfill") + # audit always runs + d.step(lambda deps, sc: "audited", deps=[c], name="audit").trigger_rule( + TriggerRule.ALL_DONE + ) + + result, _ = run_dag(register) + assert result.get_status("charge") is TaskStatus.FAILED + assert result.get_result("refund") == "refunded" + assert result.get_status("fulfill") is TaskStatus.SKIPPED + assert result.get_result("audit") == "audited" + + +def test_deps_value_is_none_for_failed_upstream_under_all_done(): + """A non-ALL_SUCCESS task (ALL_DONE) may run while an upstream FAILED. Reading + that dependency's result inside the body yields ``None`` at runtime — the + long-standing behavior that the ``DepsMap[handle] -> T | None`` type reflects. + + Exercises the ``TaskHandle`` (typed) access path specifically, since that is + the overload whose return type was corrected from bare ``T`` to ``T | None``. + Also asserts ``DagResult.get_result(handle)`` returns ``None`` for the same + failed task (its handle overload has the identical fix). + """ + seen = {} + + def boom(_deps, _sc): + raise ValueError("boom") + + def register(d): + charge = d.step(boom, name="charge", config=NO_RETRY_CFG) + + def audit(deps, _sc): + # Handle-typed access: value is None because `charge` FAILED, even + # though this ALL_DONE task legitimately runs. Also confirm the + # string-keyed access agrees. + seen["by_handle"] = deps[charge] + seen["by_name"] = deps["charge"] + return "audited" + + d.step(audit, deps=[charge], name="audit").trigger_rule(TriggerRule.ALL_DONE) + # Expose the handle to the assertions below. + seen["charge_handle"] = charge + + result, _ = run_dag(register) + + assert result.get_status("charge") is TaskStatus.FAILED + assert result.get_status("audit") is TaskStatus.SUCCEEDED + assert result.get_result("audit") == "audited" + # The dependency's value inside the body was None (not the bare result type). + assert seen["by_handle"] is None + assert seen["by_name"] is None + # DagResult.get_result for the failed task is likewise None (its handle + # overload was corrected to T | None as well). + assert result.get_result(seen["charge_handle"]) is None + assert result.get_result("charge") is None + + +def test_run_if_skip(): + def register(d): + a = d.step(lambda deps, sc: 10, name="a") + d.step( + lambda deps, sc: "ran", + deps=[a], + name="b", + run_if=lambda deps: deps["a"] > 100, + ) + + result, _ = run_dag(register) + assert result.get_status("b") is TaskStatus.SKIPPED + assert result.results["b"].skip_reason is SkipReason.RUN_IF_PREDICATE + + +def test_min_successful_early_completion(): + def register(d): + for i in range(4): + d.step(lambda deps, sc: 1, name=f"t{i}") + + result, _ = run_dag( + register, DagConfig(completion_config=CompletionConfig(min_successful=2)) + ) + assert result.completion_reason is DagCompletionReason.MIN_SUCCESSFUL_REACHED + assert result.success_count >= 2 + + +def test_custom_completion_short_circuits_on_rejected_verdict(): + """DAG-18-style rules engine: a linear chain r1 -> r2 -> r3, max_concurrency + 1, where each task returns a verdict. The custom predicate inspects + SUCCEEDED items' RESULTS (not just counts) and stops the moment any task's + verdict is REJECT -- something no threshold config can express, since + thresholds only ever see aggregate counts. r2 rejects, so r3 must never + run. + """ + ran: list[str] = [] + + def should_complete(status): + any_rejected = any( + item.status is TaskStatus.SUCCEEDED and item.result == "REJECT" + for item in status.items + ) + if any_rejected: + return complete_dag(DagCompletionOutcome.FAILED) + return continue_dag() + + def register(d): + r1 = d.step(lambda deps, sc: (ran.append("r1"), "ACCEPT")[1], name="r1") + r2 = d.step( + lambda deps, sc: (ran.append("r2"), "REJECT")[1], deps=[r1], name="r2" + ) + d.step(lambda deps, sc: (ran.append("r3"), "ACCEPT")[1], deps=[r2], name="r3") + + result, _ = run_dag( + register, + DagConfig( + max_concurrency=1, + completion_config=DagCustomCompletionConfig(should_complete), + ), + ) + assert result.completion_reason is DagCompletionReason.CUSTOM_COMPLETION_FAILED + assert result.success_count == 2 + assert "r1" in ran + assert "r2" in ran + assert "r3" not in ran + + +def test_custom_completion_succeeds_when_predicate_never_rejects(): + def should_complete(status): + if status.completed_count >= status.total_count: + return complete_dag() + return continue_dag() + + def register(d): + d.step(lambda deps, sc: "ACCEPT", name="a") + d.step(lambda deps, sc: "ACCEPT", name="b") + + result, _ = run_dag( + register, + DagConfig(completion_config=DagCustomCompletionConfig(should_complete)), + ) + assert result.completion_reason is DagCompletionReason.CUSTOM_COMPLETION_SUCCEEDED + assert result.success_count == 2 + + +def test_custom_completion_predicate_sees_accurate_live_snapshot(): + """The predicate must see exactly what has settled so far: unsettled tasks + report a None status, settled tasks report their real result/skip reason, + and the aggregate counts always match the per-item list. + """ + snapshots = [] + + def should_complete(status): + snapshots.append(status) + if status.completed_count >= 4: + return complete_dag() + return continue_dag() + + def register(d): + a = d.step(lambda deps, sc: "A", name="a") + d.step(lambda deps, sc: "B", name="b") + d.step(lambda deps, sc: "C", deps=[a], name="c", trigger_rule=TriggerRule.ALL_FAILED) + d.step(lambda deps, sc: "D", name="d") + + result, _ = run_dag( + register, + DagConfig(completion_config=DagCustomCompletionConfig(should_complete)), + ) + assert result.completion_reason is DagCompletionReason.CUSTOM_COMPLETION_SUCCEEDED + assert snapshots, "the predicate must have been invoked at least once" + for snap in snapshots: + derived_succeeded = sum( + 1 for item in snap.items if item.status is TaskStatus.SUCCEEDED + ) + derived_skipped = sum( + 1 for item in snap.items if item.status is TaskStatus.SKIPPED + ) + assert derived_succeeded == snap.success_count + assert derived_skipped == snap.skipped_count + assert len(snap.items) == len(snap.results) + assert snap.total_count == 4 + last = snapshots[-1] + assert last.success_count == 3, "a, b, d succeed" + assert last.skipped_count == 1, "c is skipped: ALL_FAILED with no failed upstream" + assert last.results["c"].skip_reason is SkipReason.TRIGGER_RULE + + +def test_failure_tolerance_exceeded(): + def boom(_deps, _sc): + raise ValueError("x") + + def register(d): + for i in range(3): + d.step(boom, name=f"t{i}", config=NO_RETRY_CFG) + + result, _ = run_dag( + register, + DagConfig( + completion_config=CompletionConfig(tolerated_failure_count=0), + ), + ) + assert result.completion_reason is DagCompletionReason.FAILURE_TOLERANCE_EXCEEDED + + +def test_default_drains_on_failure_no_fail_fast(): + """A failure does not abort; independent tasks still run (drain).""" + ran = {"b": False} + + def boom(_deps, _sc): + raise ValueError("x") + + def register(d): + d.step(boom, name="a", config=NO_RETRY_CFG) + d.step(lambda deps, sc: ran.__setitem__("b", True), name="b") + + result, _ = run_dag(register) + assert ran["b"] is True + assert result.failure_count == 1 + assert result.success_count == 1 + + +def test_throw_if_error(): + def boom(_deps, _sc): + raise ValueError("bad") + + def register(d): + d.step(boom, name="a", config=NO_RETRY_CFG) + + result, _ = run_dag(register) + with pytest.raises(DagExecutionError): + result.throw_if_error() + + +def test_empty_dag(): + result, _ = run_dag(lambda d: None) + assert result.total_count == 0 + assert result.completion_reason is DagCompletionReason.ALL_COMPLETED + + +def test_failure_tolerance_percentage_exceeded(): + def boom(_deps, _sc): + raise ValueError("x") + + def register(d): + d.step(boom, name="a", config=NO_RETRY_CFG) + d.step(lambda deps, sc: 1, name="b") + + result, _ = run_dag( + register, + DagConfig( + completion_config=CompletionConfig(tolerated_failure_percentage=10), + ), + ) + assert result.completion_reason is DagCompletionReason.FAILURE_TOLERANCE_EXCEEDED + + +def test_invalid_max_concurrency(): + state, _ = make_state() + ctx = make_context(state, parent_id="dag") + d = DagContextImpl(ctx, DagConfig()) + d.step(lambda deps, sc: 1, name="a") + with pytest.raises(ValidationError): + DagExecutor(ctx, d.get_tasks(), DagConfig(max_concurrency=0)) + + +def test_default_trigger_rule_from_config_applies(): + """DagConfig.default_trigger_rule is used when a task sets no explicit rule.""" + + def boom(_deps, _sc): + raise ValueError("x") + + def register(d): + a = d.step(boom, name="a", config=NO_RETRY_CFG) + # no explicit trigger_rule -> inherits config default ALL_DONE, so it + # runs even though its upstream FAILED. + d.step(lambda deps, sc: "ran", deps=[a], name="b") + + result, _ = run_dag( + register, + DagConfig( + default_trigger_rule=TriggerRule.ALL_DONE, + ), + ) + assert result.get_status("a") is TaskStatus.FAILED + assert result.get_status("b") is TaskStatus.SUCCEEDED + assert result.get_result("b") == "ran" + + +def test_explicit_trigger_rule_overrides_config_default(): + """An explicit per-task trigger_rule wins over DagConfig.default_trigger_rule.""" + + def register(d): + a = d.step(lambda deps, sc: 1, name="a") + # config default is ALL_DONE, but explicit ALL_FAILED + a SUCCEEDED => skip + d.step( + lambda deps, sc: "ran", + deps=[a], + name="b", + trigger_rule=TriggerRule.ALL_FAILED, + ) + + result, _ = run_dag(register, DagConfig(default_trigger_rule=TriggerRule.ALL_DONE)) + assert result.get_status("b") is TaskStatus.SKIPPED + + +# region run_if-raises abort (a raising predicate ABORTS the DAG) +import signal # noqa: E402 +from contextlib import contextmanager # noqa: E402 + +from aws_durable_execution_sdk_python.exceptions import DagPredicateError # noqa: E402 + + +@contextmanager +def _fail_on_hang(seconds: int = 10): + """Turn a scheduler hang into an assertion failure instead of blocking the + whole test session. SIGALRM fires on the main thread (where pytest runs).""" + + def _handler(_signum, _frame): + raise AssertionError("DagExecutor.run() hung (run_if abort regression)") + + old = signal.signal(signal.SIGALRM, _handler) + signal.alarm(seconds) + try: + yield + finally: + signal.alarm(0) + signal.signal(signal.SIGALRM, old) + + +def _make_executor(register, config=None, parent_id="dag"): + """Build a DagExecutor so a test can inspect ``_results`` after ``run()`` + raises (``run_dag`` can't, because the abort means there is no DagResult).""" + config = config or DagConfig() + state, _ = make_state() + ctx = make_context(state, parent_id=parent_id) + d = DagContextImpl(ctx, config) + register(d) + validate_dag(d) + return DagExecutor(ctx, d.get_tasks(), config) + + +def test_run_if_raises_on_non_root_aborts_dag(): + """A run_if that raises on a downstream task (evaluated inside a worker-thread + completion callback) ABORTS the DAG with DagPredicateError: the offending + task gets no terminal state and a downstream ALL_FAILED compensation task + never runs.""" + + def register(d): + a = d.step(lambda deps, sc: "A", name="a") + # run_if dereferences a missing dep -> KeyError, evaluated after `a` done + b = d.step( + lambda deps, sc: "ran", + deps=[a], + name="b", + run_if=lambda deps: deps["missing"] > 0, + ) + # ALL_FAILED compensation on the offending task: MUST NOT run, because a + # predicate defect must never drive a compensation path. + d.step( + lambda deps, sc: "refunded", + deps=[b], + name="refund", + trigger_rule=TriggerRule.ALL_FAILED, + ) + + ex = _make_executor(register) + with _fail_on_hang(), pytest.raises(DagPredicateError) as ei: + ex.run() + + assert isinstance(ei.value.__cause__, KeyError) + assert ei.value.task_name == "b" + assert "b" in str(ei.value) + # `a` completed normally; `b` (offending) has NO terminal state; `refund` + # (downstream ALL_FAILED) never ran. + assert ex._results["a"].status is TaskStatus.SUCCEEDED # noqa: SLF001 + assert "b" not in ex._results # noqa: SLF001 + assert "refund" not in ex._results # noqa: SLF001 + + +def test_run_if_raises_on_root_aborts_dag(): + """A raising run_if on a root task ABORTS the DAG (propagates out of the very + first pump on the caller thread) rather than failing that task.""" + + def register(d): + a = d.step( + lambda deps, sc: "ran", + name="a", + run_if=lambda deps: 1 // 0 == 0, + ) + # ALL_FAILED compensation on the offending root: MUST NOT run. + d.step( + lambda deps, sc: "refunded", + deps=[a], + name="refund", + trigger_rule=TriggerRule.ALL_FAILED, + ) + + ex = _make_executor(register) + with _fail_on_hang(), pytest.raises(DagPredicateError) as ei: + ex.run() + + assert isinstance(ei.value.__cause__, ZeroDivisionError) + assert ei.value.task_name == "a" + assert "a" not in ex._results # noqa: SLF001 - no terminal state + assert "refund" not in ex._results # noqa: SLF001 - compensation did not run + + +# endregion run_if-raises abort + + +# region threshold-completion fidelity (mirrors ExecutionCounters.should_complete) +def _threshold_executor(task_count, config): + state, _ = make_state() + ctx = make_context(state, parent_id="dag") + d = DagContextImpl(ctx, config) + for i in range(task_count): + d.step(lambda deps, sc: 1, name=f"t{i}") + return DagExecutor(ctx, d.get_tasks(), config) + + +def test_threshold_success_checked_before_failure(): + """When both min_successful and failure-tolerance fire, success wins (matches + batch ExecutionCounters ordering).""" + ex = _threshold_executor( + 3, + DagConfig( + completion_config=CompletionConfig( + min_successful=2, tolerated_failure_count=0 + ) + ), + ) + ex._success = 2 + ex._failure = 1 + assert ex._threshold_reason_locked() is DagCompletionReason.MIN_SUCCESSFUL_REACHED + + +def test_threshold_impossible_to_succeed_stops_early(): + """Once min_successful can no longer be reached, stop (reported as + FAILURE_TOLERANCE_EXCEEDED, matching batch _create_result).""" + ex = _threshold_executor( + 3, DagConfig(completion_config=CompletionConfig(min_successful=3)) + ) + ex._failure = 1 # max reachable successes = 3 - 1 = 2 < 3 + assert ( + ex._threshold_reason_locked() is DagCompletionReason.FAILURE_TOLERANCE_EXCEEDED + ) + + +def test_threshold_percentage_denominator_excludes_skipped(): + """Skipped tasks are excluded from the failure-percentage denominator so + they do not dilute the ratio.""" + ex = _threshold_executor( + 4, DagConfig(completion_config=CompletionConfig(tolerated_failure_percentage=40)) + ) + ex._skip = 2 + ex._failure = 1 + ex._success = 1 + # denom = 4 - 2 = 2 -> 50% > 40% -> exceeded. (Old denom=4 -> 25%, would NOT.) + assert ( + ex._threshold_reason_locked() is DagCompletionReason.FAILURE_TOLERANCE_EXCEEDED + ) + + +# endregion threshold-completion fidelity + + +# region multi-suspend precedence (earliest timed wins over indefinite) +from aws_durable_execution_sdk_python.exceptions import ( # noqa: E402 + SuspendExecution, + TimedSuspendExecution, +) +from aws_durable_execution_sdk_python.operation.dag_context import TaskDef # noqa: E402 + + +def _suspend_executor(specs): + """Build a DagExecutor of independent root tasks that each raise `exc`. + + `specs` is a list of (name, exception) pairs; every task is a root + (empty deps, ALL_SUCCESS) so all are submitted concurrently. + """ + state, _ = make_state() + ctx = make_context(state, parent_id="dag") + + def make_executor(exc): + def executor(_ctx, _deps_map): + raise exc + + return executor + + tasks = { + name: TaskDef( + name=name, + kind="step", + inline_deps=[], + all_deps=[], + trigger_rule=TriggerRule.ALL_SUCCESS, + run_if=None, + config=None, + executor=make_executor(exc), + ) + for name, exc in specs + } + return DagExecutor(ctx, tasks, DagConfig()) + + +def test_two_concurrent_timed_waits_raise_earliest_timestamp(): + """(a) Two concurrent timed suspends -> the EARLIEST timestamp is raised.""" + now = time.time() + ex = _suspend_executor( + [ + ("slow", TimedSuspendExecution("slow", now + 100)), + ("fast", TimedSuspendExecution("fast", now + 5)), + ] + ) + with pytest.raises(TimedSuspendExecution) as ei: + ex.run() + assert ei.value.scheduled_timestamp == pytest.approx(now + 5) + + +def test_timed_suspend_wins_over_indefinite(): + """(b) Timed + indefinite concurrent -> timed wins (timer not dropped).""" + now = time.time() + ex = _suspend_executor( + [ + ("callback", SuspendExecution("waiting for external callback")), + ("timer", TimedSuspendExecution("timer", now + 7)), + ] + ) + with pytest.raises(TimedSuspendExecution) as ei: + ex.run() + assert ei.value.scheduled_timestamp == pytest.approx(now + 7) + + +def test_indefinite_only_raises_indefinite_suspend(): + """No timed suspend pending -> the indefinite SuspendExecution is raised.""" + ex = _suspend_executor([("callback", SuspendExecution("external callback"))]) + with pytest.raises(SuspendExecution) as ei: + ex.run() + assert not isinstance(ei.value, TimedSuspendExecution) + + +# endregion multi-suspend precedence + + +# region in-process timed resume (map/parallel parity) +def _root_executor(specs, config=None): + """Build a DagExecutor whose independent root tasks run `specs` concurrently. + + `specs` is a list of (name, func) where func(ctx, deps_map) returns a result + or raises. Every task is a root (empty deps, ALL_SUCCESS) so all are + submitted at once. Returns (executor, in_memory_client). + """ + state, client = make_state() + ctx = make_context(state, parent_id="dag") + tasks = { + name: TaskDef( + name=name, + kind="step", + inline_deps=[], + all_deps=[], + trigger_rule=TriggerRule.ALL_SUCCESS, + run_if=None, + config=None, + executor=func, + ) + for name, func in specs + } + return DagExecutor(ctx, tasks, config or DagConfig()), client + + +def test_timed_wait_resumes_in_process_within_single_invocation(): + """(a) A timed suspend is resumed IN-PROCESS by the DAG-owned TimerScheduler while + a concurrent task keeps the invocation alive: run() returns a success result + (no SuspendExecution bubbles to the platform) and the timed task re-runs.""" + calls = {"x": 0, "y": 0} + x_done = threading.Event() + + def x(_ctx, _deps): + calls["x"] += 1 + if calls["x"] == 1: + # First pass suspends with a short timer; the scheduler must re-run + # this task in-process rather than surfacing a platform suspend. + raise TimedSuspendExecution("wait", time.time() + 0.05) + x_done.set() + return "x-done" + + def y(_ctx, _deps): + calls["y"] += 1 + # Stay RUNNING until X has resumed + completed, so the DAG never settles + # into a platform suspend for the pure-timed case. + assert x_done.wait(timeout=3) + return "y-done" + + ex, client = _root_executor([("x", x), ("y", y)]) + result = ex.run() # must NOT raise -> resumed within a single invocation + + assert result.get_status("x") is TaskStatus.SUCCEEDED + assert result.get_status("y") is TaskStatus.SUCCEEDED + assert result.get_result("x") == "x-done" + assert result.completion_reason is DagCompletionReason.ALL_COMPLETED + assert calls["x"] == 2 # initial + in-process timed resume + assert calls["y"] == 1 + # The resume checkpoints before re-running (mirrors ConcurrentExecutor). + assert client.checkpoint_count >= 1 + + +def test_indefinite_callback_suspends_the_invocation(): + """(b) An indefinite (callback) suspend can only be resolved by the platform: + run() raises a plain SuspendExecution and the task is never re-run.""" + calls = {"n": 0} + + def approval(_ctx, _deps): + calls["n"] += 1 + raise SuspendExecution("waiting for external callback") + + ex, _ = _root_executor([("approval", approval)]) + with pytest.raises(SuspendExecution) as ei: + ex.run() + assert not isinstance(ei.value, TimedSuspendExecution) + assert calls["n"] == 1 # no in-process resume for indefinite suspends + + +def test_mixed_timed_and_indefinite_forces_platform_suspend(): + """(c) Timed + indefinite concurrently: the indefinite one forces a platform + suspend, and timed-wins precedence still surfaces the EARLIEST timer so the + platform resumes as soon as possible. Neither task resumes in-process.""" + now = time.time() + calls = {"cb": 0, "timer": 0} + + def cb(_ctx, _deps): + calls["cb"] += 1 + raise SuspendExecution("external callback") + + def timer(_ctx, _deps): + calls["timer"] += 1 + raise TimedSuspendExecution("timer", now + 5) + + ex, _ = _root_executor([("cb", cb), ("timer", timer)]) + with pytest.raises(TimedSuspendExecution) as ei: + ex.run() + assert ei.value.scheduled_timestamp == pytest.approx(now + 5) + # Indefinite forces platform suspend -> no in-process re-run of either task. + assert calls["cb"] == 1 + assert calls["timer"] == 1 + + +# endregion in-process timed resume + + +# region no checkpoint after abort (teardown-window regression) +from aws_durable_execution_sdk_python.dag import TaskHandle # noqa: E402 +from aws_durable_execution_sdk_python.operation.dag_executor import ( # noqa: E402 + _TimedResume, +) + + +def _bare_executor(): + """A DagExecutor with no tasks, wired to an in-memory client so a direct + ``_resubmit`` call exercises the real ``create_checkpoint`` path (it lands + on ``InMemoryServiceClient.checkpoint_count``).""" + state, client = make_state() + ctx = make_context(state, parent_id="dag") + return DagExecutor(ctx, {}, DagConfig()), client + + +def test_resubmit_checkpoints_when_not_aborting_control(): + """Sensitivity anchor. With the abort flag UNSET, a timed resume writes + exactly one checkpoint before re-running (mirrors + ``ConcurrentExecutor.resubmitter``) and clears the task's timer bookkeeping. + + This is the checkpoint the abort guard must suppress. Pinning it with the + *same* setup as the guard test below — the only difference being + ``_scheduler_exception`` — proves the guard test is not vacuous: the count + flips from 1 to 0 solely because of the abort flag. + """ + ex, client = _bare_executor() + ex._pending_timers.add("t") # noqa: SLF001 + assert client.checkpoint_count == 0 + ex._resubmit([_TimedResume("t")]) + assert client.checkpoint_count == 1 # resume checkpointed + assert "t" not in ex._pending_timers # noqa: SLF001 - would re-run + + +def test_resubmit_writes_no_checkpoint_after_abort_decision(): + """Regression (direct): once the DAG has decided to abort + (``_scheduler_exception`` set — e.g. a ``run_if`` predicate raised), a timed + resume that fires during the teardown/drain window must write NO checkpoint + and must not touch task state. Identical setup to the control above; only + the abort flag differs, and it takes the checkpoint count from 1 to 0. + """ + ex, client = _bare_executor() + ex._pending_timers.add("t") # noqa: SLF001 + ex._scheduler_exception = DagPredicateError("aborted", task_name="x") # noqa: SLF001 + ex._resubmit([_TimedResume("t")]) + assert client.checkpoint_count == 0 # zero checkpoints after the abort decision + assert "t" in ex._pending_timers # noqa: SLF001 - guarded path left state intact + + +def test_no_late_checkpoint_in_abort_drain_window(): + """Regression (end-to-end race). Reproduces the reviewer's observation: the + DAG-owned ``TimerScheduler`` is the OUTER context manager and the pool the + INNER one, so the timer thread is still alive while the pool drains. A task + that timed-suspended has a resume pending; when a *different* task's + ``run_if`` aborts the DAG, the pending resume can fire ``_resubmit`` during + the drain window and — before the fix — write a checkpoint AFTER the abort + decision. + + Topology (raw TaskDefs so the step machinery does not emit its own + checkpoints; the ONLY checkpoint source is ``_resubmit``): + + * ``seed`` (root) completes immediately -> makes ``gate`` ready on a + worker thread, so the abort is captured into + ``_scheduler_exception`` (the non-root abort path) rather + than raising out of the first pump. + * ``timer`` (root) timed-suspends with a resume due *now* -> a resume is + queued on the scheduler heap and ``timer`` stays in + ``_pending_timers`` across the abort. + * ``gate`` (deps=[seed]) ``run_if`` raises -> DAG aborts. It records + ``checkpoint_count`` at that instant. + * ``blocker`` (root) stays in-flight ~0.5s to hold the pool-drain window + open, giving the ~0.1s timer loop several chances to fire + the due resume before the scheduler is torn down. + + Sensitivity: with the guard removed I observed exactly one late checkpoint + (final == at-abort + 1) and this assertion fails; with the guard, the timer + loop fires ``_resubmit`` in the same window but it early-returns, so the + count is unchanged. + """ + from aws_durable_execution_sdk_python.dag import TriggerRule # noqa: PLC0415 + from aws_durable_execution_sdk_python.operation.dag_context import ( # noqa: PLC0415 + TaskDef, + ) + + state, client = make_state() + ctx = make_context(state, parent_id="dag") + at_abort = {"count": None} + + def seed_exec(_ctx, _deps): + return "seed" + + def timer_exec(_ctx, _deps): + # Due immediately: the scheduler queues a resume the timer thread will + # fire on its next (<=0.1s) loop, i.e. squarely inside the drain window. + raise TimedSuspendExecution("timer", time.time()) + + def blocker_exec(_ctx, _deps): + # Keep the pool draining so the scheduler (outer CM) is not yet torn + # down while the due resume fires. + time.sleep(0.5) + return "blocker" + + def gate_run_if(_deps): + # The abort decision. Snapshot the checkpoint count at this instant; + # nothing legitimate may checkpoint afterwards. + at_abort["count"] = client.checkpoint_count + raise KeyError("predicate defect") + + seed_ref = TaskHandle(_name="seed", _dag=None) + + def _root(name, executor): + return TaskDef( + name=name, + kind="step", + inline_deps=[], + all_deps=[], + trigger_rule=TriggerRule.ALL_SUCCESS, + run_if=None, + config=None, + executor=executor, + ) + + tasks = { + "seed": _root("seed", seed_exec), + "timer": _root("timer", timer_exec), + "blocker": _root("blocker", blocker_exec), + "gate": TaskDef( + name="gate", + kind="step", + inline_deps=[], + all_deps=[seed_ref], + trigger_rule=TriggerRule.ALL_SUCCESS, + run_if=gate_run_if, + config=None, + executor=seed_exec, + ), + } + ex = DagExecutor(ctx, tasks, DagConfig()) + + with _fail_on_hang(), pytest.raises(DagPredicateError): + ex.run() + + # The abort actually happened via the predicate. + assert at_abort["count"] is not None + # Nothing checkpointed before the abort in this DAG, and — the regression — + # nothing checkpointed after it either, despite the resume firing in-window. + assert at_abort["count"] == 0 + assert client.checkpoint_count == 0, ( + f"late checkpoint after abort decision: {client.checkpoint_count}" + ) + # And the aborting predicate's task never got a terminal state. + assert "gate" not in ex._results # noqa: SLF001 + + +def test_resubmit_checkpoint_is_inside_the_abort_check_lock(): + """Regression (deterministic, no sleep-based timing): closes the narrower + race the end-to-end drain-window test above cannot reliably force. + + ``_scheduler_exception`` is set by ``_safe_pump`` on a WORKER thread (a + completion callback) UNDER ``self._lock`` (see its handler, guarded the + same way). ``_resubmit`` runs on the TIMER thread. The abort guard's first + check only proves no abort had been decided at the instant the timer + thread entered its critical section; if ``create_checkpoint`` were called + AFTER releasing ``self._lock``, a worker thread's ``_safe_pump`` could + acquire the lock and set the exception in that gap, and the checkpoint + would still fire moments after an abort was decided elsewhere -- + reachable only via an exact interleaving that 200 iterations of the + timing-based test above will not reliably hit. + + Proven directly here with two real threads and a non-blocking lock probe + (no sleeps anywhere): ``_resubmit`` itself runs on a background thread + (playing the timer thread), with ``create_checkpoint`` monkeypatched to + signal it has started and then block on a gate this test controls -- + holding the checkpoint call "in flight" indefinitely, independent of wall + time. Once that signal fires, the main thread (playing the racing worker + thread) performs a non-blocking ``self._lock.acquire(blocking=False)`` + probe -- exactly what ``_safe_pump`` would need to do to set + ``_scheduler_exception``. With the fix (checkpoint call genuinely inside + the lock), the probe MUST fail: ``_resubmit``'s thread still holds + ``self._lock`` for as long as its call to ``create_checkpoint`` is + outstanding. If the checkpoint call were outside the lock (the bug), + ``_resubmit`` would have already released the lock before ever calling + ``create_checkpoint``, and the probe would succeed. + """ + ex, client = _bare_executor() + ex._pending_timers.add("t") # noqa: SLF001 + + checkpoint_started = threading.Event() + release_checkpoint = threading.Event() + real_create_checkpoint = ex._ctx.state.create_checkpoint + + def _held_open_create_checkpoint(*args, **kwargs): + checkpoint_started.set() + # Held open until THIS test explicitly releases it -- independent of + # wall-clock time, so the probe below has an unbounded window to run. + release_checkpoint.wait(timeout=5) + return real_create_checkpoint(*args, **kwargs) + + ex._ctx.state.create_checkpoint = _held_open_create_checkpoint # noqa: SLF001 + + resubmit_thread = threading.Thread( + target=lambda: ex._resubmit([_TimedResume("t")]), # noqa: SLF001 + daemon=True, + ) + resubmit_thread.start() + + # Wait for _resubmit (on its own thread) to be blocked inside + # create_checkpoint, i.e. genuinely "in flight" with self._lock held if + # and only if the fix holds. + assert checkpoint_started.wait(timeout=5), "create_checkpoint was never called" + + # Mirrors exactly what _safe_pump needs to do to set the abort flag: + # acquire self._lock. Non-blocking, so this cannot itself hang the test + # regardless of which way the bug/fix goes. + acquired = ex._lock.acquire(blocking=False) # noqa: SLF001 + if acquired: + ex._lock.release() # noqa: SLF001 + + release_checkpoint.set() + resubmit_thread.join(timeout=5) + + assert not resubmit_thread.is_alive() + assert not acquired, ( + "a second thread was able to acquire self._lock while create_checkpoint " + "was in flight inside _resubmit -- the checkpoint call is reachable " + "outside the lock-protected abort check (race window regression)" + ) + assert client.checkpoint_count == 1 # the (unraced) checkpoint still completed + + +# endregion no checkpoint after abort (teardown-window regression) + + +# region default max_concurrency cap (contract: unset -> 40, previously unbounded) +from aws_durable_execution_sdk_python.operation import dag_executor as _dag_executor # noqa: E402 +from aws_durable_execution_sdk_python.operation.dag_executor import ( # noqa: E402 + DEFAULT_DAG_MAX_CONCURRENCY, +) + + +def test_default_dag_max_concurrency_constant_is_40(): + """Pin the shared cross-language default. The behavioural tests below size + themselves off this constant, so this guards against a silent retune.""" + assert DEFAULT_DAG_MAX_CONCURRENCY == 40 + + +def _spy_pool(monkeypatch): + """Record the ``max_workers`` every ThreadPoolExecutor is built with. + + Returns the list the executor's real constructor is still invoked, so the + DAG runs for real; we only observe the pool size.""" + captured: list[int] = [] + real = _dag_executor.ThreadPoolExecutor + + def spy(*args, **kwargs): + captured.append(kwargs.get("max_workers", args[0] if args else None)) + return real(*args, **kwargs) + + monkeypatch.setattr(_dag_executor, "ThreadPoolExecutor", spy) + return captured + + +def test_default_caps_pool_max_workers_when_unset(monkeypatch): + """A DAG wider than the default and with NO ``max_concurrency`` must build + its pool with exactly 40 workers, not one-per-task (the previously unbounded + behaviour that spawned N OS threads). The pool is the actual resource being + protected, so assert on the size it was constructed with.""" + captured = _spy_pool(monkeypatch) + width = DEFAULT_DAG_MAX_CONCURRENCY + 20 # 60: comfortably wider than the cap + + def register(d): + for i in range(width): + d.step(lambda deps, sc: 1, name=f"t{i}") + + result, _ = run_dag(register) + assert result.success_count == width + assert captured == [DEFAULT_DAG_MAX_CONCURRENCY] + assert captured[0] <= DEFAULT_DAG_MAX_CONCURRENCY + + +def test_default_never_exceeds_40_in_flight_when_unset(): + """The sensitive one. A graph wider than 40 with no ``max_concurrency`` must + never run more than 40 task bodies concurrently. This asserts an OBSERVED + peak (a lock-guarded counter), not a config value. + + A ``Barrier`` sized to the cap makes the assertion two-sided: every worker + increments the live counter, then blocks on the barrier, so a full wave of + exactly 40 bodies is simultaneously in flight before any releases — proving + the pool genuinely reaches 40 (not merely stays under it). Because only the + pool's threads ever run a body and each runs one at a time, the counter can + exceed 40 only if the pool was built with >40 workers. Width is a whole + multiple of the cap so the barrier drains in exact waves and never deadlocks; + the same graph under the old unbounded behaviour would put all `width` + bodies in flight at once, pushing the peak to `width`.""" + width = DEFAULT_DAG_MAX_CONCURRENCY * 2 # 80: two exact waves of 40 + tracker = {"current": 0, "peak": 0} + lock = threading.Lock() + barrier = threading.Barrier(DEFAULT_DAG_MAX_CONCURRENCY, timeout=10) + + def body(_deps, _sc): + with lock: + tracker["current"] += 1 + tracker["peak"] = max(tracker["peak"], tracker["current"]) + try: + barrier.wait() + except threading.BrokenBarrierError: # pragma: no cover - only on regression + pass + finally: + with lock: + tracker["current"] -= 1 + return 1 + + def register(d): + for i in range(width): + d.step(body, name=f"t{i}") + + result, _ = run_dag(register) + assert result.success_count == width + # Two-sided: exactly the cap was reached, and it was never exceeded. + assert tracker["peak"] == DEFAULT_DAG_MAX_CONCURRENCY + + +def test_explicit_max_concurrency_below_40_wins(monkeypatch): + """An explicit value below the default still wins: the pool is built with + that value, not the 40 cap.""" + captured = _spy_pool(monkeypatch) + + def register(d): + for i in range(60): + d.step(lambda deps, sc: 1, name=f"t{i}") + + result, _ = run_dag(register, DagConfig(max_concurrency=5)) + assert result.success_count == 60 + assert captured == [5] + + +def test_explicit_max_concurrency_above_40_wins(monkeypatch): + """An explicit value ABOVE the default still wins (the cap is only a default, + never a ceiling): the pool is built with 50 workers for a 60-task graph.""" + captured = _spy_pool(monkeypatch) + + def register(d): + for i in range(60): + d.step(lambda deps, sc: 1, name=f"t{i}") + + result, _ = run_dag(register, DagConfig(max_concurrency=50)) + assert result.success_count == 60 + assert captured == [50] + + +def test_default_cap_does_not_over_allocate_for_small_graphs(monkeypatch): + """A DAG narrower than the cap and unset ``max_concurrency`` still builds a + pool sized to the task count (min(total, 40)), preserving the pre-change + small-graph behaviour rather than always allocating 40.""" + captured = _spy_pool(monkeypatch) + + def register(d): + for i in range(3): + d.step(lambda deps, sc: 1, name=f"t{i}") + + result, _ = run_dag(register) + assert result.success_count == 3 + assert captured == [3] + + +# endregion default max_concurrency cap diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_handler_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_handler_test.py new file mode 100644 index 00000000..39f08de7 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_handler_test.py @@ -0,0 +1,391 @@ +"""T7: context.dag() wiring, FutureWarning, error unwrapping, nested DAG, exports.""" + +from __future__ import annotations + +import warnings + +import pytest + +from aws_durable_execution_sdk_python.dag import ( + DagCompletionOutcome, + DagCompletionReason, + DagConfig, + DagCustomCompletionConfig, + TaskStatus, + complete_dag, + continue_dag, +) +from aws_durable_execution_sdk_python.exceptions import DagCyclicDependencyError +from tests.dag_support import make_context, make_state + + +def _diamond(d): + a = d.step(lambda deps, sc: "A", name="a") + b = d.step(lambda deps, sc: deps["a"] + "B", deps=[a], name="b") + c = d.step(lambda deps, sc: deps["a"] + "C", deps=[a], name="c") + d.step(lambda deps, sc: deps["b"] + deps["c"], deps=[b, c], name="d") + + +def test_context_dag_end_to_end(): + state, _ = make_state() + ctx = make_context(state) + result = ctx.dag(_diamond, name="pipeline") + assert result.get_result("d") == "ABAC" + assert result.completion_reason is DagCompletionReason.ALL_COMPLETED + assert result.success_count == 4 + + +def test_context_dag_custom_completion_end_to_end(): + """DAG-18-style rules engine through the real public context.dag() entry + point: a linear chain r1 -> r2 -> r3, max_concurrency 1, where the custom + predicate stops the moment any task's result is REJECT. r2 rejects, so r3 + must never run. + """ + ran: list[str] = [] + + def should_complete(status): + any_rejected = any( + item.status is TaskStatus.SUCCEEDED and item.result == "REJECT" + for item in status.items + ) + return complete_dag(DagCompletionOutcome.FAILED) if any_rejected else continue_dag() + + def register(d): + r1 = d.step(lambda deps, sc: (ran.append("r1"), "ACCEPT")[1], name="r1") + r2 = d.step( + lambda deps, sc: (ran.append("r2"), "REJECT")[1], deps=[r1], name="r2" + ) + d.step(lambda deps, sc: (ran.append("r3"), "ACCEPT")[1], deps=[r2], name="r3") + + state, _ = make_state() + ctx = make_context(state) + config = DagConfig( + max_concurrency=1, + completion_config=DagCustomCompletionConfig(should_complete), + ) + result = ctx.dag(register, name="rules-engine", config=config) + assert result.completion_reason is DagCompletionReason.CUSTOM_COMPLETION_FAILED + assert result.success_count == 2 + assert ran == ["r1", "r2"] + + +def test_future_warning_emitted_once(): + import aws_durable_execution_sdk_python.operation.dag as dag_mod + + dag_mod._warned = False # reset for the test + state, _ = make_state() + ctx = make_context(state) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + ctx.dag(lambda d: d.step(lambda deps, sc: 1, name="a"), name="p1") + ctx.dag(lambda d: d.step(lambda deps, sc: 1, name="a"), name="p2") + future_warnings = [w for w in caught if issubclass(w.category, FutureWarning)] + assert len(future_warnings) == 1 + + +def test_cycle_surfaces_typed_error(): + state, _ = make_state() + ctx = make_context(state) + + def register(d): + a = d.step(lambda deps, sc: 1, name="a") + b = d.step(lambda deps, sc: 2, deps=[a], name="b") + a.after(b) + + with pytest.raises(DagCyclicDependencyError): + ctx.dag(register, name="cyclic") + + +def test_nested_dag_scope_isolation(): + state, _ = make_state() + ctx = make_context(state) + + def inner(d): + d.step(lambda deps, sc: "inner-x", name="x") + + def outer(d): + d.step(lambda deps, sc: "outer-a", name="a") + d.dag(inner, name="inner") + + result = ctx.dag(outer, name="outer") + assert result.get_status("a") is TaskStatus.SUCCEEDED + nested = result.get_result("inner") + assert nested.get_result("x") == "inner-x" + + +def test_invalid_max_concurrency_raises_at_handler(): + from aws_durable_execution_sdk_python.exceptions import ValidationError + + state, _ = make_state() + ctx = make_context(state) + with pytest.raises(ValidationError): + ctx.dag( + lambda d: d.step(lambda deps, sc: 1, name="a"), + name="p", + config=DagConfig(max_concurrency=-1), + ) + + +def test_public_exports(): + import aws_durable_execution_sdk_python as sdk + + for symbol in [ + "DagContext", + "TaskHandle", + "DagResult", + "DagConfig", + "TriggerRule", + "TaskStatus", + "SkipReason", + "DagCompletionReason", + "DagExecutionError", + "DagCyclicDependencyError", + "DagInvalidTaskNameError", + "DagDuplicateTaskError", + "DagInvalidDependencyError", + "DagPredicateError", + ]: + assert hasattr(sdk, symbol), symbol + + +def test_degradation_ladder_drops_failed_task_names_at_rung_three(): + """When the no-``tasks`` envelope still exceeds the checkpoint limit (many + failed task names), the ladder drops ``failedTaskNames`` too -- but never the + counts, ``completionReason`` or ``startedTaskNames``.""" + import json + + from aws_durable_execution_sdk_python.dag import ( + DagCompletionReason, + TaskExecution, + TaskStatus, + ) + from aws_durable_execution_sdk_python.identifier import OperationIdentifier + from aws_durable_execution_sdk_python.lambda_service import ( + ErrorObject, + OperationSubType, + ) + from aws_durable_execution_sdk_python.operation.dag import DagContainerExecutor + from aws_durable_execution_sdk_python.operation.dag_result import DagResultImpl + + # ~5000 failed tasks with long names: the failedTaskNames list alone exceeds + # the 256KB limit, so rung 2 (drop tasks) is not enough and rung 3 fires. + results = {} + kinds = {} + for i in range(5000): + name = f"task_{i:06d}_" + "x" * 60 + results[name] = TaskExecution( + name, TaskStatus.FAILED, error=ErrorObject.from_message("e") + ) + kinds[name] = "step" + result = DagResultImpl( + results, DagCompletionReason.COMPLETED_WITH_FAILURES, kinds, total_count=5000 + ) + + captured: dict = {} + + class _FakeState: + def create_checkpoint(self, operation_update, is_sync=True): + captured["update"] = operation_update + + executor = DagContainerExecutor( + run_body=lambda _r: result, + state=_FakeState(), # type: ignore[arg-type] + operation_identifier=OperationIdentifier( + operation_id="container", + sub_type=OperationSubType.DAG, + parent_id=None, + name="p", + ), + ) + executor._checkpoint_with_ladder(result) + + upd = captured["update"] + assert upd.context_options.replay_children is True + env = json.loads(upd.payload) + assert "tasks" not in env + assert "failedTaskNames" not in env # dropped at rung 3 + # Never dropped: counts, completionReason, startedTaskNames. + assert env["completionReason"] == "COMPLETED_WITH_FAILURES" + assert env["startedTaskNames"] == [] + assert env["successCount"] == 0 + assert env["failureCount"] == 5000 + assert env["totalCount"] == 5000 + + +def test_unwrap_dag_error_reconstructs_typed_error_on_replay(): + """On replay the checkpointed failure rebuilds a CallableRuntimeError with + error_type set but __cause__ absent; unwrap must still surface the typed + Dag* error so replay matches the first run.""" + from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + DagExecutionError, + ) + from aws_durable_execution_sdk_python.operation.dag import unwrap_dag_error + + exc = ChildContextError( + message="2 task(s) FAILED", + error_type="DagExecutionError", + data=None, + stack_trace=None, + ) + assert exc.__cause__ is None + with pytest.raises(DagExecutionError, match="FAILED"): + unwrap_dag_error(exc) + + +def test_unwrap_dag_error_passthrough_for_non_dag_error(): + from aws_durable_execution_sdk_python.exceptions import ChildContextError + from aws_durable_execution_sdk_python.operation.dag import unwrap_dag_error + + exc = ChildContextError( + message="boom", error_type="ValueError", data=None, stack_trace=None + ) + with pytest.raises(ChildContextError): + unwrap_dag_error(exc) + + +def test_unwrap_dag_error_preserves_live_cause_when_present(): + """When a live DagPredicateError is the ChildContextError cause (before the + durable boundary rebuilds it), unwrap surfaces it with its OWN original + cause and task_name intact, suppressing the ChildContextError wrapper.""" + from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + DagPredicateError, + ) + from aws_durable_execution_sdk_python.operation.dag import unwrap_dag_error + + original = KeyError("missing") + predicate_error = DagPredicateError( + "run_if predicate for DAG task 'b' raised KeyError: 'missing'", + task_name="b", + ) + predicate_error.__cause__ = original + wrapper = ChildContextError( + message="run_if predicate for DAG task 'b' raised KeyError: 'missing'", + error_type="DagPredicateError", + data=None, + stack_trace=None, + ) + wrapper.__cause__ = predicate_error + + with pytest.raises(DagPredicateError) as ei: + unwrap_dag_error(wrapper) + + assert ei.value is predicate_error + assert ei.value.task_name == "b" + assert ei.value.__cause__ is original + + +def test_run_if_raise_aborts_dag_through_context(): + """A raising run_if surfaces DagPredicateError (not a DagResult with a FAILED + task) out of ctx.dag(). Across the durable child-context boundary the error + is rebuilt from serialized fields (type name + message) to keep first run and + replay identical, so ``task_name`` and the live ``__cause__`` are erased here; + the offending task survives IN the message. The wrapped chain is verified at + the scheduler level in dag_executor_test.py instead.""" + from aws_durable_execution_sdk_python.dag import TriggerRule + from aws_durable_execution_sdk_python.exceptions import DagPredicateError + + state, _ = make_state() + ctx = make_context(state) + + def register(d): + a = d.step(lambda deps, sc: "A", name="a") + d.step( + lambda deps, sc: "b", + deps=[a], + name="b", + run_if=lambda deps: deps["missing"] > 0, # KeyError + ) + # ALL_FAILED compensation MUST NOT run. + d.step( + lambda deps, sc: pytest.fail("compensation ran on a predicate defect"), + deps=[a], + name="refund", + trigger_rule=TriggerRule.ALL_FAILED, + ) + + with pytest.raises(DagPredicateError) as ei: + ctx.dag( + register, + name="pipeline", + config=DagConfig(), + ) + + # Type surfaces cleanly; the offending task is named in the durable message. + assert "b" in str(ei.value) + # The durable boundary rebuilds from serialized fields, so task_name and the + # live cause are not retrievable here (parity with the rest of the Dag* + # family and identical on replay). + assert ei.value.task_name is None + assert ei.value.__cause__ is None + + +def test_unwrap_dag_error_reconstructs_predicate_error_on_replay(): + """On replay the checkpointed predicate abort rebuilds with error_type set + but __cause__ absent; unwrap must still surface DagPredicateError (task name + survives in the message).""" + from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + DagPredicateError, + ) + from aws_durable_execution_sdk_python.operation.dag import unwrap_dag_error + + exc = ChildContextError( + message="run_if predicate for DAG task 'b' raised KeyError: 'missing'", + error_type="DagPredicateError", + data=None, + stack_trace=None, + ) + assert exc.__cause__ is None + with pytest.raises(DagPredicateError, match="'b'"): + unwrap_dag_error(exc) + + +def test_deps_none_for_failed_upstream_end_to_end_through_context_dag(): + """Integration (public surface): a customer-shaped handler that runs a DAG via + ``ctx.dag(...)`` -- the real public entry point -- with a ``charge`` task that + raises and a downstream ``audit`` task (``TriggerRule.ALL_DONE``) depending on + it. ``audit`` legitimately runs even though ``charge`` FAILED and, reading the + failed dependency's value inside its body, observes ``None`` (the runtime + behaviour the ``DepsMap[handle] -> T | None`` contract encodes). + + Unlike ``dag_executor_test.test_deps_value_is_none_for_failed_upstream_under_all_done`` + -- which drives ``DagExecutor`` directly -- this goes end-to-end through + ``dag_handler`` (validation, the durable checkpoint boundary and result + reconstruction), so the ``None`` observation is exercised on the surface a + customer actually calls and survives round-tripping into the returned + ``DagResult``. + """ + from aws_durable_execution_sdk_python.config import StepConfig + from aws_durable_execution_sdk_python.dag import TriggerRule + from aws_durable_execution_sdk_python.retries import RetryPresets + + # Disable retries so the intentionally-failing task fails promptly (attempt 1). + no_retry = StepConfig(retry_strategy=RetryPresets.none()) + + state, _ = make_state() + ctx = make_context(state) + + def register(d): + def charge(deps, sc): + raise RuntimeError("charge failed") + + charge_task = d.step(charge, name="charge", config=no_retry) + + def audit(deps, sc): + # Handle-typed access: ``charge`` FAILED, so its value is None even + # though this ALL_DONE task legitimately runs. Return the observation + # so it round-trips through the DagResult the caller receives. + return deps[charge_task] is None + + d.step(audit, deps=[charge_task], name="audit").trigger_rule( + TriggerRule.ALL_DONE + ) + + result = ctx.dag(register, name="pipeline") + + assert result.get_status("charge") is TaskStatus.FAILED + assert result.get_status("audit") is TaskStatus.SUCCEEDED + # audit succeeded having observed the failed upstream's value as None. + assert result.get_result("audit") is True diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_large_payload_coverage_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_large_payload_coverage_test.py new file mode 100644 index 00000000..cd27c7a6 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_large_payload_coverage_test.py @@ -0,0 +1,212 @@ +"""Integration coverage for the large-payload gap (conformance 10-15). + +Runs the *exact* ``10-15`` graph in-process across a real container replay, so +the regressions the cloud suite cannot catch are caught here without a deploy. + +The DAG ``bigdag`` has eight root step tasks ``p1``..``p8``; task ``pN`` returns +its letter repeated 51200 times, so the aggregate is ~410KB (8 * 51200 = 409600 +chars) -- comfortably over the 256KB checkpoint limit -- while every individual +task result stays far under it. When the container result is checkpointed the +aggregate is OFFLOADED: Python drops ``tasks`` from the envelope and marks the +container ``ReplayChildren=true`` while still writing the aggregate summary +(counts, completionReason, startedTaskNames), then RECONSTRUCTS on replay from +that envelope plus the retained child checkpoints. + +The reconstruct-vs-inline divergence only fires when a SUCCEEDED container is +REPLAYED, so the driver mirrors the handler: it runs ``dag()``, a checkpointed +``digestBefore`` step, then a ``wait`` that SUSPENDS the invocation; the next +invocation replays the completed container. Three things are asserted: + +* **Aggregate fidelity** across the replay -- every task result is individually + retrievable and byte-identical afterwards, including one full 51200-char value, + and the language-neutral digest ``"8:409600:abcdefgh"`` matches before and + after the suspend. +* **Task bodies are not re-invoked** -- external per-task counters prove each body + ran exactly once across the offload and the container replay. Under + reconstruct the DAG register graph re-runs, but each task's step operation must + fast-path from its own (small, normally-checkpointed) result; a body running + twice is a duplicated customer side effect, the bug this test exists to catch. +* **The offload + reconstruct path was actually taken** -- the container + operation carries ``replay_children is True`` and an envelope with ``tasks`` + dropped but the aggregate summary present. This is the observable hook that + distinguishes Python's reconstruct strategy. +""" + +from __future__ import annotations + +import threading +from typing import Any + +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.context import DurableContext, ExecutionContext +from aws_durable_execution_sdk_python.dag import DagCompletionReason, DagConfig +from aws_durable_execution_sdk_python.exceptions import SuspendExecution +from aws_durable_execution_sdk_python.lambda_service import ( + OperationSubType, + OperationType, +) +from tests.dag_support import InMemoryServiceClient +from tests.operation.dag_concurrency_coverage_test import ( # reuse the proven harness + _complete_wait, + _fail_on_hang, + _seeded_state, +) + +_TASK_COUNT = 8 +_REPEAT = 51200 +_TASK_NAMES = [f"p{i}" for i in range(1, _TASK_COUNT + 1)] +_EXPECTED_DIGEST = "8:409600:abcdefgh" + + +def _digest(result: Any) -> str: + """``"::"`` over p1..p8.""" + total_length = 0 + first_chars = [] + for name in _TASK_NAMES: + value = result.get_result(name) + total_length += len(value) + first_chars.append(value[0]) + return f"{len(_TASK_NAMES)}:{total_length}:{''.join(first_chars)}" + + +def _dag_container_op(client: InMemoryServiceClient) -> Any: + """The single DAG-container operation (CONTEXT op with SubType=Dag).""" + containers = [ + op + for op in client.operations.values() + if op.operation_type is OperationType.CONTEXT + and op.sub_type is OperationSubType.DAG + ] + assert len(containers) == 1, f"expected one DAG container, found {len(containers)}" + return containers[0] + + +def test_10_15_large_payload_survives_container_replay() -> None: + """The 10-15 graph across a real suspend: the ~410KB aggregate is offloaded, + the completed container is replayed on the next invocation, and the aggregate + comes back byte-identical via child-body re-execution. + + A single driver run exercises all three assertions (fidelity, single + invocation per body, and the re-execution path) so the counters and the + stored container operation observe the same replay. + """ + # External per-task counters: a body that runs twice increments twice. + calls = {name: 0 for name in _TASK_NAMES} + calls_lock = threading.Lock() + + def _bump(name: str) -> None: + with calls_lock: + calls[name] += 1 + + def _make_body(name: str, letter: str): + def _body(_deps: Any, _sc: Any) -> str: + _bump(name) + return letter * _REPEAT + + return _body + + # register runs once per DAG-body execution. It runs a second time on the + # replay iff the container is re-executed (ReplayChildren) rather than + # reconstructed from an envelope -- direct evidence of the re-execution path. + register_calls = {"n": 0} + + def register(d: Any) -> None: + register_calls["n"] += 1 + for i, name in enumerate(_TASK_NAMES): + d.step(_make_body(name, chr(ord("a") + i)), name=name) + + invocations = {"n": 0} + + def run(ctx: DurableContext): + invocations["n"] += 1 + result = ctx.dag( + register, name="bigdag", config=DagConfig(max_concurrency=1) + ) + # Checkpointed step: computed once from the live DagResult, fast-pathed + # from its own checkpoint after the suspend, so it carries the + # pre-suspend digest across the boundary. + digest_before: str = ctx.step( + lambda _sc: _digest(result), name="digestBefore" + ) + # Ends the invocation; the next one replays the completed container. + ctx.wait(Duration.from_seconds(2), name="pauseForReplay") + # Recomputed from the REPLAYED DagResult after resume. + digest_after = _digest(result) + return result, digest_before, digest_after + + client = InMemoryServiceClient() + final: tuple[Any, str, str] | None = None + with _fail_on_hang(): + # First invocation resolves the DAG then suspends on the wait; the + # second replays the completed container and finishes. + for _ in range(2): + state = _seeded_state(client, dict(client.operations)) + ctx = DurableContext( + state=state, + execution_context=ExecutionContext( + durable_execution_arn=state.durable_execution_arn + ), + parent_id=None, + ) + try: + final = run(ctx) + break + except SuspendExecution: + assert _complete_wait(client, "pauseForReplay") + + assert final is not None, "handler never completed across the suspend" + result, digest_before, digest_after = final + + # --- Aggregate fidelity across the replay -------------------------------- + assert result.completion_reason is DagCompletionReason.ALL_COMPLETED + assert ( + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ) == (8, 0, 0, 8) + # Every task result is individually retrievable and byte-identical, and at + # least one full 51200-char value is checked in its entirety. + for i, name in enumerate(_TASK_NAMES): + expected = chr(ord("a") + i) * _REPEAT + assert result.get_result(name) == expected + assert result.get_result("p1") == "a" * _REPEAT # full-value check + assert len(result.get_result("p1")) == _REPEAT + # The language-neutral assertion: the digest survived the offload and came + # back identical through the replay strategy. + assert digest_before == _EXPECTED_DIGEST + assert digest_after == _EXPECTED_DIGEST + assert digest_before == digest_after + + # --- Task bodies were not re-invoked ------------------------------------- + # The DAG child body re-runs under ReplayChildren, but each task step must + # fast-path from its own checkpoint. Exactly one invocation per body. + for name in _TASK_NAMES: + assert calls[name] == 1, f"task {name} body ran {calls[name]} times, expected 1" + + # --- The offload + reconstruct path was actually taken ------------------- + # Python offloads via the degradation ladder: it drops ``tasks`` and sets + # ReplayChildren, but -- unlike the pre-convergence behaviour -- it still + # writes the aggregate envelope (counts, completionReason, startedTaskNames). + # On replay it RECONSTRUCTS from that envelope plus the retained child + # checkpoints rather than blindly re-executing. + import json + + container = _dag_container_op(client) + assert container.context_details is not None + assert container.context_details.replay_children is True + envelope = json.loads(container.context_details.result) + assert envelope["type"] == "DagResult" + assert "tasks" not in envelope # per-task detail offloaded to the children + assert envelope["totalCount"] == _TASK_COUNT + assert envelope["successCount"] == _TASK_COUNT + assert envelope["completionReason"] == "ALL_COMPLETED" + # No task was in-flight at completion, so the started set is empty. + assert envelope["startedTaskNames"] == [] + + # Corroborating evidence that the interesting path was genuinely exercised: + # the invocation actually suspended and resumed (two invocations), and the + # DAG register graph was re-run on the reconstruct (register ran twice) so + # each task could fast-path from its own retained child checkpoint. + assert invocations["n"] == 2, "container was not replayed across a suspend" + assert register_calls["n"] == 2, "DAG graph was not re-run on reconstruct" diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_nested_large_payload_coverage_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_nested_large_payload_coverage_test.py new file mode 100644 index 00000000..20939a70 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_nested_large_payload_coverage_test.py @@ -0,0 +1,228 @@ +"""Integration coverage for the nested + large-payload intersection (10-17). + +Runs the *exact* ``10-17`` graph in-process across a real container replay, so +the regression the cloud suite cannot catch is caught here without a deploy. + +The outer DAG ``outernested`` (``max_concurrency=1``) has a single task +``inner`` that is itself a nested DAG (``max_concurrency=1``) with six root step +tasks ``p1``..``p6``; task ``pN`` returns its letter repeated 51200 times, so the +inner aggregate is ~307KB (6 * 51200 = 307200 chars) -- comfortably over the +256KB checkpoint limit. The inner container therefore OFFLOADS (drops ``tasks``, +sets ``ReplayChildren``), and because the outer embeds the inner result in full, +the outer aggregate is over the limit too, so the OUTER container offloads as +well. This is the untested intersection: an offloaded outer whose one task is an +offloaded inner. + +The bug this guards against (confirmed in TypeScript): on the outer's +reconstruct path the inner DagResult is rebuilt from the inner container's +tasks-less envelope alone, coming back EMPTY while still claiming +``ALL_COMPLETED`` -- so the per-task detail is silently lost. The fix is that the +outer's reconstruct re-runs its register graph, and the ``inner`` nested-dag task +re-runs through the container executor, which detects the offloaded inner and +RECONSTRUCTS it from the inner's own retained child checkpoints, recursively +(contract rule 2). Three things are asserted: + +* **Recursive fidelity** -- after the outer container is replayed, the inner + DagResult reports ``ALL_COMPLETED`` with counts ``[6,0,0,6]`` AND every inner + per-task result is individually retrievable and byte-identical, so the + language-neutral digest ``"6:307200:abcdef"`` matches before and after the + suspend. The digest, not the reason, is the decisive check: under the bug the + reason still reads ``ALL_COMPLETED`` from the honest aggregate while the digest + after replay would differ (empty inner). +* **Task bodies are not re-invoked** -- external per-task counters prove each + inner body ran exactly once across the offload of BOTH containers and the + double replay. Nesting doubles the number of containers that replay, so a body + running twice is the duplicated-side-effect bug this test exists to catch. +* **Both containers offloaded + reconstructed** -- the outer AND the inner DAG + container operations each carry ``replay_children is True`` with ``tasks`` + dropped but the aggregate summary present, and both register graphs re-ran on + the reconstruct. +""" + +from __future__ import annotations + +import json +import threading +from typing import Any + +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.context import DurableContext, ExecutionContext +from aws_durable_execution_sdk_python.dag import DagCompletionReason, DagConfig +from aws_durable_execution_sdk_python.exceptions import SuspendExecution +from aws_durable_execution_sdk_python.lambda_service import ( + OperationSubType, + OperationType, +) +from tests.dag_support import InMemoryServiceClient +from tests.operation.dag_concurrency_coverage_test import ( # reuse the proven harness + _complete_wait, + _fail_on_hang, + _seeded_state, +) + +_INNER_COUNT = 6 +_REPEAT = 51200 +_INNER_NAMES = [f"p{i}" for i in range(1, _INNER_COUNT + 1)] +_EXPECTED_DIGEST = "6:307200:abcdef" + + +def _digest(inner: Any) -> str: + """``"::"`` over + p1..p6, so the first-char run is deterministic regardless of completion + order. For the 10-17 inner graph this is exactly ``"6:307200:abcdef"``.""" + total_length = 0 + first_chars = [] + for name in _INNER_NAMES: + value = inner.get_result(name) + total_length += len(value) + first_chars.append(value[0]) + return f"{len(_INNER_NAMES)}:{total_length}:{''.join(first_chars)}" + + +def _dag_container_ops(client: InMemoryServiceClient) -> list[Any]: + """Every DAG-container operation (CONTEXT op with SubType=Dag).""" + return [ + op + for op in client.operations.values() + if op.operation_type is OperationType.CONTEXT + and op.sub_type is OperationSubType.DAG + ] + + +def test_10_17_nested_large_payload_survives_container_replay() -> None: + """The 10-17 graph across a real suspend: the ~307KB inner aggregate offloads + both the inner AND the outer container; the completed outer is replayed on + the next invocation and the inner per-task detail comes back byte-identical + via recursive reconstruction from the inner's own child checkpoints.""" + # External per-inner-task counters: a body that runs twice increments twice. + calls = {name: 0 for name in _INNER_NAMES} + calls_lock = threading.Lock() + + def _bump(name: str) -> None: + with calls_lock: + calls[name] += 1 + + def _make_body(name: str, letter: str): + def _body(_deps: Any, _sc: Any) -> str: + _bump(name) + return letter * _REPEAT + + return _body + + # Each register runs once per DAG-body execution; it runs a second time on + # the replay iff its container is reconstructed (ReplayChildren) -- direct + # evidence that both the outer and inner reconstruct paths were taken. + outer_register_calls = {"n": 0} + inner_register_calls = {"n": 0} + + def register(d: Any) -> None: + outer_register_calls["n"] += 1 + + def inner_register(sd: Any) -> None: + inner_register_calls["n"] += 1 + for i, name in enumerate(_INNER_NAMES): + sd.step(_make_body(name, chr(ord("a") + i)), name=name) + + d.dag(inner_register, name="inner", config=DagConfig(max_concurrency=1)) + + invocations = {"n": 0} + + def run(ctx: DurableContext): + invocations["n"] += 1 + outer = ctx.dag( + register, name="outernested", config=DagConfig(max_concurrency=1) + ) + # Checkpointed step: the inner digest computed once from the live inner + # DagResult, fast-pathed from its own checkpoint after the suspend, so it + # carries the pre-suspend digest across the boundary. + digest_before: str = ctx.step( + lambda _sc: _digest(outer.get_result("inner")), name="digestBefore" + ) + # Ends the invocation; the next one replays the completed outer container. + ctx.wait(Duration.from_seconds(2), name="pauseForReplay") + # Recomputed from the REPLAYED (recursively reconstructed) inner result. + digest_after = _digest(outer.get_result("inner")) + return outer, digest_before, digest_after + + client = InMemoryServiceClient() + final: tuple[Any, str, str] | None = None + with _fail_on_hang(): + # First invocation resolves the nested DAG then suspends on the wait; the + # second replays the completed outer container and finishes. + for _ in range(2): + state = _seeded_state(client, dict(client.operations)) + ctx = DurableContext( + state=state, + execution_context=ExecutionContext( + durable_execution_arn=state.durable_execution_arn + ), + parent_id=None, + ) + try: + final = run(ctx) + break + except SuspendExecution: + assert _complete_wait(client, "pauseForReplay") + + assert final is not None, "handler never completed across the suspend" + outer, digest_before, digest_after = final + + # --- Recursive fidelity: the inner survived the offload of BOTH containers - + inner = outer.get_result("inner") + assert inner is not None, "inner DagResult was lost across the replay" + assert inner.completion_reason is DagCompletionReason.ALL_COMPLETED + assert ( + inner.success_count, + inner.failure_count, + inner.skipped_count, + inner.total_count, + ) == (6, 0, 0, 6) + # Every inner task result is individually retrievable and byte-identical, + # and at least one full 51200-char value is checked in its entirety. This is + # the rule-2 assertion: the per-task detail is PRESENT, not just the honest + # aggregate. + for i, name in enumerate(_INNER_NAMES): + expected = chr(ord("a") + i) * _REPEAT + assert inner.get_result(name) == expected + assert inner.get_result("p1") == "a" * _REPEAT # full-value check + assert len(inner.get_result("p1")) == _REPEAT + # The decisive language-neutral assertion. + assert digest_before == _EXPECTED_DIGEST + assert digest_after == _EXPECTED_DIGEST + assert digest_before == digest_after + + # --- Inner task bodies were not re-invoked -------------------------------- + for name in _INNER_NAMES: + assert calls[name] == 1, f"inner body {name} ran {calls[name]} times, expected 1" + + # --- Both containers offloaded + reconstructed ---------------------------- + containers = _dag_container_ops(client) + assert len(containers) == 2, ( + f"expected outer + inner DAG containers, found {len(containers)}" + ) + for container in containers: + assert container.context_details is not None + assert container.context_details.replay_children is True, ( + "a DAG container did not offload (ReplayChildren) as expected" + ) + envelope = json.loads(container.context_details.result) + assert envelope["type"] == "DagResult" + assert "tasks" not in envelope # per-task detail offloaded to children + assert envelope["completionReason"] == "ALL_COMPLETED" + + # The inner container's aggregate summary is canonical and correct. + inner_container = next(c for c in containers if c.name == "inner") + inner_env = json.loads(inner_container.context_details.result) + assert inner_env["totalCount"] == _INNER_COUNT + assert inner_env["successCount"] == _INNER_COUNT + assert inner_env["failureCount"] == 0 + assert inner_env["skippedCount"] == 0 + assert inner_env["startedTaskNames"] == [] + + # Corroborating evidence the interesting path was genuinely exercised: the + # invocation suspended and resumed (two invocations), and BOTH register + # graphs re-ran on the reconstruct so each inner task could fast-path from + # its own retained child checkpoint. + assert invocations["n"] == 2, "outer container was not replayed across a suspend" + assert outer_register_calls["n"] == 2, "outer DAG graph was not re-run on reconstruct" + assert inner_register_calls["n"] == 2, "inner DAG graph was not re-run on reconstruct" diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_order_independence_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_order_independence_test.py new file mode 100644 index 00000000..d1603ada --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_order_independence_test.py @@ -0,0 +1,357 @@ +"""DAG order-independence guards (Python analogue of JS ``DAG-19``, Go +``TestDagE2E_OrderIndependenceReplay`` and Java ``diamondWithWaitReplays…`` / +``wideFanOutTasksAlwaysObserveUpstreamValue``). + +The whole DAG design rests on one property: a task's backend entity/operation id +is a **pure function of ``(scope-prefix, name)``** and carries *no* completion- +or registration-order input (``context.py::_create_task_id`` -> +``blake2b(f"{parent}-DAG_NODE_T_{name}")[:64]``, deliberately NOT touching the +per-context step counter). If that ever regressed to a counter-based id the DAG +would trip ``NonDeterministicExecutionError`` (or, worse, alias two tasks onto +one checkpoint) the first time tasks completed in a different order across a +replay boundary — the exact failure name-based ids exist to prevent. + +Before this file Python had concurrency/throttle tests but *no* order- +independence guard, unlike the other three SDKs (see +``dag-review/GAP_concurrent_completion.md`` §2). These tests close that hole from +three complementary angles: + +(a) ``test_registration_order_independence_identical_record`` — the DAG-19 + equivalent, the strongest signal: build the same diamond twice with the task + *registration order permuted* and assert the derived per-task ids AND the + normalized checkpoint record are identical. Asserts on the actual derived + ids (not just results), so a counter-based-id regression cannot hide. +(b) ``test_completion_order_independence_under_concurrency`` — with + ``max_concurrency=2`` and an event seam that forces a later-registered branch + to complete FIRST (deterministic, not racy), assert the per-task ids and the + DagResult equal a serial run. +(c) ``test_wide_fan_out_readers_observe_correct_upstream_value`` — the Java B1 + analogue: a barrier makes many readers complete simultaneously (concurrent + writes to the shared results map) and a collector then verifies every task + observed the correct upstream value, across many iterations. +""" + +from __future__ import annotations + +import hashlib +import threading +from typing import Any + +from aws_durable_execution_sdk_python.dag import ( + DagCompletionReason, + DagConfig, + TaskStatus, +) +from aws_durable_execution_sdk_python.lambda_service import OperationSubType +from tests.dag_support import InMemoryServiceClient, make_context, make_state + +# ───────────────────────────────────────────────────────────────────────── +# Helpers: extract the derived ids / normalized record from the checkpoint +# stream that ``InMemoryServiceClient`` persisted during a ``.dag()`` run. +# ───────────────────────────────────────────────────────────────────────── + + +def _ids_by_name(client: InMemoryServiceClient) -> dict[str | None, str]: + """Map every persisted operation's ``name`` -> its backend ``operation_id``. + + Includes the DAG container itself. Each DAG task mints exactly one operation + id (its Start/Succeed updates reuse it), so this is 1:1. + """ + return {op.name: op.operation_id for op in client.operations.values()} + + +def _task_ops(client: InMemoryServiceClient) -> list[Any]: + """Every non-container (i.e. real task) operation.""" + return [ + op + for op in client.operations.values() + if op.sub_type is not OperationSubType.DAG and op.name is not None + ] + + +def _assert_name_based(client: InMemoryServiceClient) -> None: + """Assert every task id is *exactly* the name-based blake2b derivation. + + This is the direct structural guard: it recomputes + ``blake2b(f"{parent}-DAG_NODE_T_{name}")[:64]`` from each op's own + ``(parent_id, name)`` and asserts equality. A counter-based (or otherwise + order-dependent) id would not match this recomputation. + """ + task_ops = _task_ops(client) + assert task_ops, "expected at least one task operation in the checkpoint stream" + for op in task_ops: + preimage = ( + f"{op.parent_id}-DAG_NODE_T_{op.name}" + if op.parent_id + else f"DAG_NODE_T_{op.name}" + ) + expected = hashlib.blake2b(preimage.encode()).hexdigest()[:64] + assert op.operation_id == expected, ( + f"task {op.name!r} id is not the name-based derivation: " + f"{op.operation_id} != {expected}" + ) + + +def _dag_result_view(result: Any, names: list[str]) -> dict[str, Any]: + """A name-keyed, order-independent semantic view of a ``DagResult``.""" + return { + "results": {n: result.get_result(n) for n in names}, + "statuses": {n: result.get_status(n).name for n in names}, + "counts": ( + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ), + "reason": result.completion_reason.name, + } + + +def _normalized_record( + client: InMemoryServiceClient, result: Any, names: list[str] +) -> dict[str, Any]: + """A normalized (sorted, order-independent) checkpoint record. + + Mirrors JS ``DAG-19``'s ``sortDeep`` record: the per-operation projection is + sorted by ``(name, operation_id)`` so that a mere change in + completion/registration *ordering* does not perturb it — but a change in any + ``operation_id`` (the counter-based-id regression) does. + """ + ops = sorted( + ( + { + "name": op.name, + "operation_id": op.operation_id, + "parent_id": op.parent_id, + "sub_type": op.sub_type.name if op.sub_type else None, + "status": op.status.name, + # Step results are simple JSON-able values here; the container + # op carries no step_details (its serialized DagResult payload is + # completion-order sensitive and is asserted semantically via + # ``_dag_result_view`` instead). + "result": op.step_details.result if op.step_details else None, + } + for op in client.operations.values() + ), + key=lambda r: (r["name"] or "", r["operation_id"]), + ) + return {"ops": ops, "dag_result": _dag_result_view(result, names)} + + +# ───────────────────────────────────────────────────────────────────────── +# Diamond builder used by (a) and (b): root -> {b, c} -> merge. +# ───────────────────────────────────────────────────────────────────────── + +_DIAMOND_NAMES = ["root", "b", "c", "merge"] + + +def _diamond_register(order: str) -> Any: + """Return a ``register`` callback building root -> {b, c} -> merge, with the + two middle branches declared in ``order`` ('bc' or 'cb'). The *logical* graph + is identical for both; only the registration order changes. + """ + + def register(d: Any) -> None: + root = d.step(lambda deps, sc: 100, name="root") + if order == "bc": + b = d.step(lambda deps, sc: deps["root"] + 1, deps=[root], name="b") + c = d.step(lambda deps, sc: deps["root"] + 2, deps=[root], name="c") + else: + c = d.step(lambda deps, sc: deps["root"] + 2, deps=[root], name="c") + b = d.step(lambda deps, sc: deps["root"] + 1, deps=[root], name="b") + d.step(lambda deps, sc: deps["b"] + deps["c"], deps=[b, c], name="merge") + + return register + + +def _assert_diamond_results(result: Any) -> None: + assert result.get_result("root") == 100 + assert result.get_result("b") == 101 + assert result.get_result("c") == 102 + assert result.get_result("merge") == 203 + assert result.completion_reason is DagCompletionReason.ALL_COMPLETED + assert ( + result.success_count, + result.failure_count, + result.skipped_count, + result.total_count, + ) == (4, 0, 0, 4) + + +# ───────────────────────────────────────────────────────────────────────── +# (a) Registration-order independence — the DAG-19 equivalent. +# ───────────────────────────────────────────────────────────────────────── + + +def test_registration_order_independence_identical_record(): + """Build the same diamond twice with b/c registered in swapped order and + assert the derived per-task ids and the whole normalized checkpoint record + are IDENTICAL. + + This directly proves ids are a pure function of ``(scope, name)`` and not of + order. ``max_concurrency=1`` makes the scheduler fully deterministic so that, + were ids counter-based, the permuted registration would deterministically + assign b/c *different* ids between the two runs — i.e. this test has teeth + (verified by the sabotage check in the implementation note). + """ + cfg = DagConfig(max_concurrency=1) + + state1, client1 = make_state() + r1 = make_context(state1).dag(_diamond_register("bc"), name="dag_oi", config=cfg) + state2, client2 = make_state() + r2 = make_context(state2).dag(_diamond_register("cb"), name="dag_oi", config=cfg) + + _assert_diamond_results(r1) + _assert_diamond_results(r2) + + # Every task id is exactly the name-based blake2b derivation ... + _assert_name_based(client1) + _assert_name_based(client2) + # ... and is byte-for-byte identical across the two registration orders. + assert _ids_by_name(client1) == _ids_by_name(client2) + # The full normalized checkpoint record is identical (the DAG-19 property). + assert _normalized_record(client1, r1, _DIAMOND_NAMES) == _normalized_record( + client2, r2, _DIAMOND_NAMES + ) + + +# ───────────────────────────────────────────────────────────────────────── +# (b) Completion-order independence under real concurrency. +# ───────────────────────────────────────────────────────────────────────── + + +def test_completion_order_independence_under_concurrency(): + """With ``max_concurrency=2`` force a later-registered branch to complete + FIRST (deterministically, via an event seam — NOT sleeps/races) and assert + the derived per-task ids and the DagResult equal a serial baseline. + + Determinism: ``b`` is registered first but blocks on an event that ``c`` + (registered second) sets right before it returns, so the observed completion + order is always ``[c, b]`` — the reverse of registration. Because both must + be in flight at once for ``c`` to unblock ``b``, this also proves the two + branches genuinely ran concurrently (it would deadlock at concurrency 1). + The core assertion (ids/result equal the serial run) is itself timing- + independent since ids are name-based, so the test cannot flake. + """ + # Serial baseline. + state_s, client_s = make_state() + r_s = make_context(state_s).dag( + _diamond_register("bc"), name="dag_co", config=DagConfig(max_concurrency=1) + ) + + completion: list[str] = [] + completion_lock = threading.Lock() + c_done = threading.Event() + + def register(d: Any) -> None: + root = d.step(lambda deps, sc: 100, name="root") + + def b_body(deps: Any, sc: Any) -> int: + # Registered FIRST, but wait until c has completed -> finishes SECOND. + assert c_done.wait(timeout=5), "concurrency seam broke: c never completed" + with completion_lock: + completion.append("b") + return deps["root"] + 1 + + b = d.step(b_body, deps=[root], name="b") + + def c_body(deps: Any, sc: Any) -> int: + # Registered SECOND, returns immediately -> finishes FIRST, then + # unblocks b. Record completion before signalling so the order is + # deterministic. + with completion_lock: + completion.append("c") + value = deps["root"] + 2 + c_done.set() + return value + + c = d.step(c_body, deps=[root], name="c") + d.step(lambda deps, sc: deps["b"] + deps["c"], deps=[b, c], name="merge") + + state_c, client_c = make_state() + r_c = make_context(state_c).dag( + register, name="dag_co", config=DagConfig(max_concurrency=2) + ) + + _assert_diamond_results(r_s) + _assert_diamond_results(r_c) + + # Completion order was the REVERSE of registration, deterministically. + assert completion == ["c", "b"] + # Ids are name-based -> identical to the serial run despite inverted timing. + _assert_name_based(client_c) + assert _ids_by_name(client_c) == _ids_by_name(client_s) + # And the DagResult is identical. + assert _dag_result_view(r_c, _DIAMOND_NAMES) == _dag_result_view( + r_s, _DIAMOND_NAMES + ) + + +# ───────────────────────────────────────────────────────────────────────── +# (c) Wide fan-out — Java B1 analogue (shared results-map regression guard). +# ───────────────────────────────────────────────────────────────────────── + + +def test_wide_fan_out_readers_observe_correct_upstream_value(): + """Many readers of a common upstream, completing simultaneously under high + concurrency, must each observe the CORRECT upstream value; a downstream + collector must then observe every reader's correct value. + + A ``Barrier`` sized to the reader count forces all readers to be in flight + and to complete at the same instant, maximizing concurrent writes to the + shared results map (the exact condition of the Java B1 race, + ``wideFanOutTasksAlwaysObserveUpstreamValue``). Python's results map is + lock-guarded, so this is a regression guard that must pass reliably; it is + repeated over many iterations to be meaningful while staying fast. + """ + fan_out = 16 + iterations = 20 + sentinel = {"marker": "root-value", "n": 987654321} + mismatches: list[Any] = [] + mismatch_lock = threading.Lock() + + def make_reader(idx: int) -> Any: + def reader(deps: Any, sc: Any) -> int: + if deps["root"] != sentinel: + with mismatch_lock: + mismatches.append(("reader", idx, deps["root"])) + # Complete simultaneously with the other readers -> concurrent writes + # into the shared results map. + try: + barrier.wait() + except threading.BrokenBarrierError: # pragma: no cover - only on hang + with mismatch_lock: + mismatches.append(("barrier-broken", idx)) + return idx + + return reader + + def collector(deps: Any, sc: Any) -> str: + for i in range(fan_out): + if deps[f"r{i}"] != i: + with mismatch_lock: + mismatches.append(("collector", i, deps[f"r{i}"])) + return "ok" + + for _ in range(iterations): + barrier = threading.Barrier(fan_out, timeout=10) + + def register(d: Any) -> None: + root = d.step(lambda deps, sc: sentinel, name="root") + readers = [ + d.step(make_reader(i), deps=[root], name=f"r{i}") + for i in range(fan_out) + ] + d.step(collector, deps=readers, name="collector") + + state, _ = make_state() + result = make_context(state).dag( + register, name="dag_fanout", config=DagConfig(max_concurrency=fan_out) + ) + + assert result.get_status("collector") is TaskStatus.SUCCEEDED + assert result.get_result("collector") == "ok" + assert result.success_count == fan_out + 2 # root + readers + collector + assert result.completion_reason is DagCompletionReason.ALL_COMPLETED + + assert mismatches == [], f"tasks observed wrong dep values under concurrency: {mismatches}" diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_result_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_result_test.py new file mode 100644 index 00000000..de79ddc0 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_result_test.py @@ -0,0 +1,259 @@ +"""T6: DagResult accessors + serialization round-trip.""" + +from __future__ import annotations + +from aws_durable_execution_sdk_python.concurrency.models import ( + BatchItem, + BatchItemStatus, + BatchResult, + CompletionReason, +) +from aws_durable_execution_sdk_python.dag import ( + DagCompletionReason, + SkipReason, + TaskExecution, + TaskHandle, + TaskStatus, +) +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.operation.dag_result import ( + DagResultImpl, + create_dag_result_serdes, + dag_reason_from_core, +) + + +def _sample_results(): + return { + "a": TaskExecution("a", TaskStatus.SUCCEEDED, result={"v": 1}), + "b": TaskExecution( + "b", TaskStatus.FAILED, error=ErrorObject.from_message("boom") + ), + "c": TaskExecution("c", TaskStatus.SKIPPED, skip_reason=SkipReason.TRIGGER_RULE), + } + + +def test_accessors(): + r = DagResultImpl( + _sample_results(), + DagCompletionReason.COMPLETED_WITH_FAILURES, + {"a": "step", "b": "step", "c": "step"}, + ) + assert r.success_count == 1 + assert r.failure_count == 1 + assert r.skipped_count == 1 + assert r.total_count == 3 + assert r.get_status("a") is TaskStatus.SUCCEEDED + assert r.get_result("a") == {"v": 1} + # C3: a TaskHandle arg resolves by its name (typed path), same as string. + handle: TaskHandle = TaskHandle(_name="a", _dag=None) + assert r.get_result(handle) == {"v": 1} + assert r.get_status(handle) is TaskStatus.SUCCEEDED + assert r.get_status("missing") is None + assert [t.name for t in r.succeeded()] == ["a"] + assert [t.name for t in r.failed()] == ["b"] + assert [t.name for t in r.skipped()] == ["c"] + + +def test_dag_reason_from_core(): + assert ( + dag_reason_from_core(CompletionReason.ALL_COMPLETED) + is DagCompletionReason.ALL_COMPLETED + ) + assert ( + dag_reason_from_core(CompletionReason.MIN_SUCCESSFUL_REACHED) + is DagCompletionReason.MIN_SUCCESSFUL_REACHED + ) + + +def test_roundtrip_plain_and_error(): + r = DagResultImpl( + _sample_results(), + DagCompletionReason.COMPLETED_WITH_FAILURES, + {"a": "step", "b": "step", "c": "step"}, + ) + serdes = create_dag_result_serdes() + data = serdes.serialize(r, None) + restored = serdes.deserialize(data, None) + assert restored.completion_reason is DagCompletionReason.COMPLETED_WITH_FAILURES + assert restored.get_result("a") == {"v": 1} + assert restored.get_status("c") is TaskStatus.SKIPPED + assert restored.results["c"].skip_reason is SkipReason.TRIGGER_RULE + assert restored.results["b"].error.message == "boom" + + +def test_roundtrip_batch_result_kind(): + batch = BatchResult( + [BatchItem(0, BatchItemStatus.SUCCEEDED, "x")], CompletionReason.ALL_COMPLETED + ) + results = {"m": TaskExecution("m", TaskStatus.SUCCEEDED, result=batch)} + r = DagResultImpl(results, DagCompletionReason.ALL_COMPLETED, {"m": "map"}) + restored = DagResultImpl.from_dict(r.to_dict()) + inner = restored.get_result("m") + assert isinstance(inner, BatchResult) + assert inner.get_results() == ["x"] + + +def test_roundtrip_nested_dag_result_kind(): + inner = DagResultImpl( + {"x": TaskExecution("x", TaskStatus.SUCCEEDED, result=42)}, + DagCompletionReason.ALL_COMPLETED, + {"x": "step"}, + ) + outer = DagResultImpl( + {"nested": TaskExecution("nested", TaskStatus.SUCCEEDED, result=inner)}, + DagCompletionReason.ALL_COMPLETED, + {"nested": "dag"}, + ) + restored = DagResultImpl.from_dict(outer.to_dict()) + nested = restored.get_result("nested") + assert isinstance(nested, DagResultImpl) + assert nested.get_result("x") == 42 + + +def test_envelope_shape_with_tasks(): + """The converged envelope: type, always-present aggregates + null-explicit + per-task fields (contract rules 1 and 3).""" + import datetime + + started = datetime.datetime(2026, 7, 26, 3, 19, 1, 884000, tzinfo=datetime.UTC) + completed = datetime.datetime(2026, 7, 26, 3, 19, 1, 885000, tzinfo=datetime.UTC) + results = { + "load": TaskExecution( + "load", + TaskStatus.SUCCEEDED, + result="ok", + started_at=started, + completed_at=completed, + ), + "charge": TaskExecution( + "charge", TaskStatus.FAILED, error=ErrorObject.from_message("boom") + ), + "ship": TaskExecution( + "ship", TaskStatus.SKIPPED, skip_reason=SkipReason.TRIGGER_RULE + ), + } + env = DagResultImpl( + results, + DagCompletionReason.COMPLETED_WITH_FAILURES, + {"load": "step", "charge": "step", "ship": "step"}, + ).to_dict() + + assert env["type"] == "DagResult" + assert env["totalCount"] == 3 + assert env["successCount"] == 1 + assert env["failureCount"] == 1 + assert env["skippedCount"] == 1 + assert env["completionReason"] == "COMPLETED_WITH_FAILURES" + assert env["startedTaskNames"] == [] + assert env["failedTaskNames"] == ["charge"] + + by_name = {t["name"]: t for t in env["tasks"]} + # Every canonical per-task field is present, null when unset. + for t in env["tasks"]: + assert set(t) == { + "name", + "status", + "skipReason", + "resultKind", + "result", + "error", + "startedAt", + "completedAt", + } + assert by_name["load"]["startedAt"] == "2026-07-26T03:19:01.884Z" + assert by_name["load"]["completedAt"] == "2026-07-26T03:19:01.885Z" + assert by_name["load"]["skipReason"] is None + assert by_name["load"]["resultKind"] == "plain" + assert by_name["ship"]["skipReason"] == "TRIGGER_RULE" + assert by_name["ship"]["startedAt"] is None + # Canonical PascalCase error object with explicit nulls. + err = by_name["charge"]["error"] + assert err == {"ErrorMessage": "boom", "ErrorType": None, "StackTrace": None} + + +def test_envelope_tasks_dropped_is_valid(): + """The offloaded case is the same envelope minus ``tasks``; from_dict yields + an empty results map and preserves the aggregates.""" + data = { + "type": "DagResult", + "totalCount": 8, + "successCount": 6, + "failureCount": 1, + "skippedCount": 1, + "completionReason": "COMPLETED_WITH_FAILURES", + "startedTaskNames": ["reserve"], + "failedTaskNames": ["charge"], + # no "tasks" + } + restored = DagResultImpl.from_dict(data) + assert restored.completion_reason is DagCompletionReason.COMPLETED_WITH_FAILURES + assert restored.total_count == 8 + assert dict(restored.results) == {} + + +def test_tasks_less_envelope_preserves_counts_and_reason(): + """Contract rule 1: restoring a tasks-less envelope that reports failures + MUST preserve ``totalCount``, the three counts and ``completionReason`` from + the envelope, and MUST NOT fabricate ``ALL_COMPLETED`` or zeroed counts. + + The per-task map is legitimately empty (the detail was offloaded to the + child checkpoints), but the aggregate summary is authoritative and survives. + """ + data = { + "type": "DagResult", + "totalCount": 8, + "successCount": 6, + "failureCount": 1, + "skippedCount": 1, + "completionReason": "COMPLETED_WITH_FAILURES", + "startedTaskNames": [], + "failedTaskNames": ["charge"], + # no "tasks" -- offloaded + } + restored = DagResultImpl.from_dict(data) + + # completionReason is preserved and is NOT the fabricated ALL_COMPLETED. + assert restored.completion_reason is DagCompletionReason.COMPLETED_WITH_FAILURES + assert restored.completion_reason is not DagCompletionReason.ALL_COMPLETED + # All four counts come straight from the envelope, not the (empty) map. + assert restored.total_count == 8 + assert restored.success_count == 6 + assert restored.failure_count == 1 + assert restored.skipped_count == 1 + # The per-task map is empty, as allowed; the accessor lists reflect that. + assert dict(restored.results) == {} + assert restored.succeeded() == [] + assert restored.failed() == [] + assert restored.skipped() == [] + + +def test_tasks_present_counts_still_derive_consistently(): + """The count override never diverges from the map on a fully-recorded + round-trip: to_dict writes the map-derived counts, from_dict reads them back, + and they equal the map counts (a present ``tasks`` array stays canonical).""" + r = DagResultImpl( + _sample_results(), + DagCompletionReason.COMPLETED_WITH_FAILURES, + {"a": "step", "b": "step", "c": "step"}, + ) + restored = DagResultImpl.from_dict(r.to_dict()) + assert restored.success_count == 1 == len(restored.succeeded()) + assert restored.failure_count == 1 == len(restored.failed()) + assert restored.skipped_count == 1 == len(restored.skipped()) + assert restored.total_count == 3 + + +def test_from_dict_ignores_unknown_fields(): + """Contract rule 4: readers MUST ignore unknown fields and treat a missing + field as absent rather than failing (additive-only evolution).""" + env = DagResultImpl( + {"a": TaskExecution("a", TaskStatus.SUCCEEDED, result=1)}, + DagCompletionReason.ALL_COMPLETED, + {"a": "step"}, + ).to_dict() + env["schemaVersion"] = "v99" # unknown top-level field + env["tasks"][0]["futureField"] = {"any": "thing"} # unknown per-task field + restored = DagResultImpl.from_dict(env) + assert restored.get_result("a") == 1 + assert restored.completion_reason is DagCompletionReason.ALL_COMPLETED diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_seam_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_seam_test.py new file mode 100644 index 00000000..8217e5fe --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_seam_test.py @@ -0,0 +1,99 @@ +"""T2: name-based entity-id seam tests.""" + +from __future__ import annotations + +import hashlib + +from tests.dag_support import make_context, make_state + + +def _task_digest(prefix: str | None, name: str) -> str: + """Expected DAG task id: the name-based pre-image, blake2b-bounded to 64 hex. + + Mirrors ``DurableContext._create_task_id`` / the core's + ``_create_step_id_for_logical_step`` bounding so the backend operation id + stays within the 64-char ``updates[].id`` limit. + """ + logical_id = f"{prefix}-DAG_NODE_T_{name}" if prefix else f"DAG_NODE_T_{name}" + return hashlib.blake2b(logical_id.encode()).hexdigest()[:64] + + +def test_create_task_id_unprefixed(): + state, _ = make_state() + ctx = make_context(state, parent_id=None) + task_id = ctx._create_task_id("fetch") + assert task_id == _task_digest(None, "fetch") + assert len(task_id) <= 64 + + +def test_create_task_id_prefixed(): + state, _ = make_state() + ctx = make_context(state, parent_id="container") + task_id = ctx._create_task_id("fetch") + assert task_id == _task_digest("container", "fetch") + assert len(task_id) <= 64 + + +def test_create_task_id_does_not_touch_counter(): + state, _ = make_state() + ctx = make_context(state, parent_id="c") + before = ctx._step_counter.get_current() + ctx._create_task_id("a") + ctx._create_task_id("b") + assert ctx._step_counter.get_current() == before + + +def test_no_collision_with_counter_ids(): + """Counter pre-images are {prefix}-{int}; task pre-images are + {prefix}-DAG_NODE_T_{name}. Both are blake2b-bounded, so the digests differ + because the pre-images differ (the reserved token guarantees disjointness). + """ + state, _ = make_state() + ctx = make_context(state, parent_id="c") + counter_id = ctx._create_step_id() + task_id = ctx._create_task_id("1") + assert counter_id != task_id + assert task_id == _task_digest("c", "1") + + +def test_seam_checkpoints_under_task_id_and_fast_path_on_replay(): + """Drive one explicit-id step through the seam; confirm checkpoint id and fast path.""" + state, client = make_state() + ctx = make_context(state, parent_id="dagc") + + calls = {"n": 0} + + def body(_step_ctx): + calls["n"] += 1 + return "value" + + # first call runs and checkpoints under the name-based id + result = ctx._run_step_with_task_id("mytask", body) + assert result == "value" + assert calls["n"] == 1 + assert _task_digest("dagc", "mytask") in client.operations + + # second call (simulated replay) hits the checkpoint fast path, no re-exec + result2 = ctx._run_step_with_task_id("mytask", body) + assert result2 == "value" + assert calls["n"] == 1 # not re-executed + + +def test_per_level_hashing_no_multi_level_preimage(): + """A nested DAG container id becomes the child's parent_id, so sub-task ids + are blake2b({container}-DAG_NODE_T_{name})[:64] per level — the container's + already-hashed digest is the prefix, never a single raw multi-level string + like ...DAG_NODE_T_validation-DAG_NODE_T_rule_a hashed once at one level.""" + state, _ = make_state() + outer = make_context(state, parent_id="root") + container_id = outer._create_task_id("validation") + assert container_id == _task_digest("root", "validation") + assert len(container_id) <= 64 + # nested scope: the container digest is the child's parent prefix + nested = make_context(state, parent_id=container_id) + sub_id = nested._create_task_id("rule_a") + # sub-task pre-image uses the container DIGEST as prefix (re-hashed per level) + assert sub_id == _task_digest(container_id, "rule_a") + assert len(sub_id) <= 64 + # it is NOT the single raw multi-level pre-image hashed at one level + assert sub_id != _task_digest("root", "validation-DAG_NODE_T_rule_a") diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_started_set_reconstruct_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_started_set_reconstruct_test.py new file mode 100644 index 00000000..6c204c5c --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_started_set_reconstruct_test.py @@ -0,0 +1,132 @@ +"""The STARTED-set fix: offloaded reconstruct must not restart an in-flight task. + +Python previously documented losing the STARTED set under large-payload early +completion -- an in-flight task could restart on replay (a duplicated customer +side effect). The converged envelope carries ``startedTaskNames`` precisely so +the offloaded reconstruct path can seed those tasks as STARTED and never +reschedule them. + +Because Python's drain-based scheduler bubbles the whole DAG on any suspend, a +*returned* DagResult never itself contains a STARTED task, so a natural offloaded +envelope has an empty started set. To exercise the fix deterministically this +test crafts the offloaded container checkpoint directly: a two-root DAG is run +once (so ``winner`` has a real, name-based child checkpoint), then the container +is rewritten to the offloaded shape (``ReplayChildren``, ``tasks`` dropped) and +``laggard``'s child checkpoint is removed so that -- absent the started set -- +reconstruct WOULD reschedule and re-run its body. The two cases are contrasted: + +* with ``startedTaskNames == ["laggard"]`` the body never runs and the task is + reproduced STARTED (the fix); +* with ``startedTaskNames == []`` the body runs again (the pre-fix restart), + proving the started set is what prevents it. +""" + +from __future__ import annotations + +import dataclasses +import json +from typing import Any + +from aws_durable_execution_sdk_python.dag import DagCompletionReason, TaskStatus +from aws_durable_execution_sdk_python.lambda_service import ( + ContextDetails, + OperationStatus, + OperationSubType, + OperationType, +) +from tests.dag_support import InMemoryServiceClient, make_context, make_state +from tests.operation.dag_concurrency_coverage_test import _seeded_state + +_DAG_NAME = "bigdag" + + +def _register_factory(calls: dict[str, int]): + def register(d: Any) -> None: + d.step(lambda deps, sc: "W", name="winner") + + def laggard_body(_deps: Any, _sc: Any) -> str: + calls["laggard"] += 1 + return "L" + + d.step(laggard_body, name="laggard") + + return register + + +def _craft_offloaded_container( + client: InMemoryServiceClient, started_task_names: list[str] +) -> None: + """Rewrite the DAG container to the offloaded shape and drop laggard's + child checkpoint so reconstruct must rely on the started set.""" + # Rewrite the container: SUCCEEDED + ReplayChildren + envelope w/o tasks. + for op_id, op in list(client.operations.items()): + if ( + op.operation_type is OperationType.CONTEXT + and op.sub_type is OperationSubType.DAG + ): + envelope = { + "type": "DagResult", + "totalCount": 2, + "successCount": 1, + "failureCount": 0, + "skippedCount": 0, + "completionReason": "MIN_SUCCESSFUL_REACHED", + "startedTaskNames": started_task_names, + "failedTaskNames": [], + } + client.operations[op_id] = dataclasses.replace( + op, + status=OperationStatus.SUCCEEDED, + context_details=ContextDetails( + replay_children=True, result=json.dumps(envelope), error=None + ), + ) + # Remove laggard's own child checkpoint: without the started set, reconstruct + # would find it un-checkpointed and re-run the body. + for op_id, op in list(client.operations.items()): + if op.name == "laggard" and op.sub_type is not OperationSubType.DAG: + del client.operations[op_id] + + +def _reconstruct_once(started_task_names: list[str]) -> tuple[Any, dict[str, int]]: + calls = {"laggard": 0} + register = _register_factory(calls) + + # First run: both roots complete and checkpoint (small, inline). + state, client = make_state() + make_context(state).dag(register, name=_DAG_NAME) + assert calls["laggard"] == 1 # body ran once on the first invocation + + # Craft the offloaded container and reset the counter to observe reconstruct. + _craft_offloaded_container(client, started_task_names) + calls["laggard"] = 0 + + # Re-invoke: the container is SUCCEEDED + ReplayChildren -> reconstruct. + recon_state = _seeded_state(client, dict(client.operations)) + result = make_context(recon_state).dag(register, name=_DAG_NAME) + return result, calls + + +def test_started_task_is_not_restarted_on_reconstruct() -> None: + """With ``laggard`` in the started set, reconstruct seeds it STARTED and its + body never runs; ``winner`` fast-paths from its checkpoint.""" + result, calls = _reconstruct_once(started_task_names=["laggard"]) + + assert calls["laggard"] == 0, "in-flight task was restarted on reconstruct" + assert result.get_status("laggard") is TaskStatus.STARTED + assert result.get_status("winner") is TaskStatus.SUCCEEDED + assert result.get_result("winner") == "W" + # Aggregates come from the envelope, not re-derivation. + assert result.completion_reason is DagCompletionReason.MIN_SUCCESSFUL_REACHED + assert result.total_count == 2 + assert result.success_count == 1 + + +def test_without_started_set_the_task_restarts() -> None: + """Control: with an empty started set the same reconstruct re-runs the + laggard body -- the exact restart the started set exists to prevent.""" + result, calls = _reconstruct_once(started_task_names=[]) + + assert calls["laggard"] == 1, "expected the pre-fix restart without the set" + # Re-run, so it lands SUCCEEDED rather than STARTED. + assert result.get_status("laggard") is TaskStatus.SUCCEEDED diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_task_kinds_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_task_kinds_test.py new file mode 100644 index 00000000..7a25fd7b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_task_kinds_test.py @@ -0,0 +1,120 @@ +"""T8/T9: exercise every DagContext task-kind executor closure end-to-end.""" + +from __future__ import annotations + +import pytest + +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.waits import ( + WaitForConditionConfig, + WaitForConditionDecision, +) +from aws_durable_execution_sdk_python.dag import DagConfig, TaskStatus +from aws_durable_execution_sdk_python.exceptions import SuspendExecution +from aws_durable_execution_sdk_python.retries import RetryPresets +from tests.dag_support import make_context, make_state + +NO_RETRY = RetryPresets.none() + + +def test_map_task(): + def register(d): + d.map( + [1, 2, 3], + lambda ctx, item, idx, items: ctx.step(lambda sc: item * 2, name=f"m{idx}"), + name="mymap", + ) + + state, _ = make_state() + result = make_context(state).dag(register, name="p") + assert result.get_status("mymap") is TaskStatus.SUCCEEDED + batch = result.get_result("mymap") + assert sorted(batch.get_results()) == [2, 4, 6] + + +def test_parallel_task(): + def register(d): + d.parallel( + [ + lambda ctx: ctx.step(lambda sc: "one", name="b1"), + lambda ctx: ctx.step(lambda sc: "two", name="b2"), + ], + name="par", + ) + + state, _ = make_state() + result = make_context(state).dag(register, name="p") + assert result.get_status("par") is TaskStatus.SUCCEEDED + assert sorted(result.get_result("par").get_results()) == ["one", "two"] + + +def test_wait_for_condition_task_completes(): + def check(deps, state_val, cctx): + return {"done": True} + + cfg = WaitForConditionConfig( + wait_strategy=lambda s, attempt: WaitForConditionDecision.stop_polling(), + initial_state={"done": False}, + ) + + def register(d): + d.wait_for_condition(check, cfg, name="poll") + + state, _ = make_state() + result = make_context(state).dag(register, name="p") + assert result.get_status("poll") is TaskStatus.SUCCEEDED + assert result.get_result("poll") == {"done": True} + + +def test_invoke_task_suspends(): + def register(d): + d.invoke("fn:prod", lambda deps: {"x": 1}, name="charge") + + state, _ = make_state() + with pytest.raises(SuspendExecution): + make_context(state).dag(register, name="p") + + +def test_wait_task_suspends(): + def register(d): + d.wait(Duration.from_seconds(30), name="cooldown") + + state, _ = make_state() + with pytest.raises(SuspendExecution): + make_context(state).dag(register, name="p") + + +def test_wait_for_callback_task_suspends(): + def register(d): + d.wait_for_callback( + lambda deps, cb_id, ctx: None, name="approval" + ) + + state, _ = make_state() + with pytest.raises(SuspendExecution): + make_context(state).dag(register, name="p") + + +def test_invoke_eager_payload(): + """Non-callable payload is used as-is (covers the eager branch).""" + def register(d): + d.invoke("fn:prod", {"eager": True}, name="charge") + + state, _ = make_state() + with pytest.raises(SuspendExecution): + make_context(state).dag(register, name="p") + + +def test_map_inputs_from_deps(): + def register(d): + src = d.step(lambda deps, sc: [1, 2], name="src") + d.map( + lambda deps: deps["src"], + lambda ctx, item, idx, items: ctx.step(lambda sc: item + 10, name=f"m{idx}"), + deps=[src], + name="mymap", + ) + + state, _ = make_state() + result = make_context(state).dag(register, name="p") + assert sorted(result.get_result("mymap").get_results()) == [11, 12] diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py new file mode 100644 index 00000000..9eee5922 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py @@ -0,0 +1,182 @@ +"""T9: DAG replay / interruption / large-payload re-execute tests.""" + +from __future__ import annotations + +import pytest + +from aws_durable_execution_sdk_python.dag import ( + DagCompletionReason, + TaskStatus, +) +from aws_durable_execution_sdk_python.exceptions import SuspendExecution +from tests.dag_support import make_context, make_state + + +def test_container_checkpointed_with_dag_subtype_and_serialized_result(): + """The DAG container is checkpointed under sub_type=DAG and its DagResult + serializes into the container payload (round-trips via the DAG serdes). + + Note: small-payload child-context *replay* short-circuit returns None in this + base-SDK snapshot because ``CheckpointedResult`` does not read back CONTEXT + results (a pre-existing limitation, unrelated to the DAG). The DAG's + re-execute model (spec §8.2) is exercised by the interrupt-resume and + large-payload tests below. + """ + from aws_durable_execution_sdk_python.lambda_service import OperationSubType + from aws_durable_execution_sdk_python.operation.dag_result import ( + create_dag_result_serdes, + ) + + side = {"a": 0} + + def register(d): + d.step(lambda deps, sc: side.__setitem__("a", side["a"] + 1) or "A", name="a") + + state, client = make_state() + r1 = make_context(state).dag(register, name="p") + assert side["a"] == 1 + assert r1.get_result("a") == "A" + + container = next( + o for o in client.operations.values() if o.sub_type is OperationSubType.DAG + ) + assert container.sub_type is OperationSubType.DAG + # step task checkpointed under a name-based id: the {parent}-DAG_NODE_T_{name} + # pre-image, blake2b-bounded to <=64 hex (backend id-length limit). + import hashlib + + task_a = next( + o + for o in client.operations.values() + if o.name == "a" and o.sub_type is not OperationSubType.DAG + ) + preimage = f"{task_a.parent_id}-DAG_NODE_T_a" + assert task_a.operation_id == hashlib.blake2b(preimage.encode()).hexdigest()[:64] + assert len(task_a.operation_id) <= 64 + # the serialized container payload round-trips to an equal DagResult + payload = container.context_details.result + restored = create_dag_result_serdes().deserialize(payload, None) + assert restored.get_result("a") == "A" + assert restored.completion_reason is DagCompletionReason.ALL_COMPLETED + + +def test_inline_container_replay_deserializes_without_rerun(): + """Replaying a completed, non-offloaded container deserializes the envelope + (tasks present) and returns it without re-running any task body.""" + from tests.operation.dag_concurrency_coverage_test import _seeded_state + + side = {"a": 0} + + def register(d): + d.step( + lambda deps, sc: side.__setitem__("a", side["a"] + 1) or "A", name="a" + ) + + state, client = make_state() + r1 = make_context(state).dag(register, name="p") + assert side["a"] == 1 + assert r1.get_result("a") == "A" + + # Second invocation on the seeded operations: the container is SUCCEEDED with + # tasks present, so it deserializes rather than re-running the body. + state2 = _seeded_state(client, dict(client.operations)) + r2 = make_context(state2).dag(register, name="p") + assert side["a"] == 1 # body NOT re-run + assert r2.get_result("a") == "A" + assert r2.completion_reason is DagCompletionReason.ALL_COMPLETED + assert r2.total_count == 1 + + +def test_interrupt_and_resume(): + """Interrupt via a gate that suspends on run 1; resume completes remaining once.""" + side = {"a": 0, "work": 0} + control = {"suspend": True} + + def register(d): + a = d.step( + lambda deps, sc: side.__setitem__("a", side["a"] + 1) or "A", name="a" + ) + + def gate(deps, sc): + if control["suspend"]: + raise SuspendExecution("gate not ready") + return "open" + + g = d.step(gate, deps=[a], name="gate") + d.step( + lambda deps, sc: side.__setitem__("work", side["work"] + 1) or "done", + deps=[g], + name="work", + ) + # a run_if-skipped task that must stay skipped across replay + d.step( + lambda deps, sc: "never", + deps=[a], + name="skipme", + run_if=lambda deps: False, + ) + + state, client = make_state() + + # run 1: gate suspends -> whole DAG suspends + with pytest.raises(SuspendExecution): + make_context(state).dag(register, name="p") + assert side["a"] == 1 # a completed and checkpointed + assert side["work"] == 0 # work never scheduled (gate not terminal) + # skipped task minted no checkpoint (no operation carries its name) + assert not any(o.name == "skipme" for o in client.operations.values()) + + # run 2: gate now succeeds + control["suspend"] = False + result = make_context(state).dag(register, name="p") + assert side["a"] == 1 # a hit the checkpoint fast path, not re-executed + assert side["work"] == 1 # remaining task ran exactly once + assert result.get_result("work") == "done" + assert result.get_status("skipme") is TaskStatus.SKIPPED + assert not any(o.name == "skipme" for o in client.operations.values()) + + +def test_large_payload_reexecutes_to_equal_result(): + """A >256KB DagResult forces the offload ladder (drop ``tasks`` + + ReplayChildren); the DAG reconstructs to an equal result on replay and each + task body runs exactly once (fast-pathed from its own checkpoint).""" + big = "x" * (300 * 1024) + side = {"big": 0, "small": 0} + + def register(d): + d.step( + lambda deps, sc: side.__setitem__("big", side["big"] + 1) or big, + name="big", + ) + d.step( + lambda deps, sc: side.__setitem__("small", side["small"] + 1) or "s", + name="small", + ) + + state, client = make_state() + + r1 = make_context(state).dag(register, name="p") + assert side["big"] == 1 + # container stored with replay_children due to large payload + import json + + from aws_durable_execution_sdk_python.lambda_service import OperationSubType + + container = next( + o for o in client.operations.values() if o.sub_type is OperationSubType.DAG + ) + assert container.context_details.replay_children is True + # Offloaded envelope: tasks dropped, but the aggregate summary survives. + envelope = json.loads(container.context_details.result) + assert "tasks" not in envelope + assert envelope["type"] == "DagResult" + assert envelope["completionReason"] == "ALL_COMPLETED" + assert envelope["successCount"] == 2 + assert envelope["startedTaskNames"] == [] + + # replay: container reconstructs (ReplayChildren), tasks fast-path + r2 = make_context(state).dag(register, name="p") + assert side["big"] == 1 # big task not re-executed (own checkpoint fast path) + assert side["small"] == 1 + assert r2.get_result("big") == r1.get_result("big") == big + assert r2.completion_reason is DagCompletionReason.ALL_COMPLETED diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_trigger_rules_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_trigger_rules_test.py new file mode 100644 index 00000000..d3a1ea57 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_trigger_rules_test.py @@ -0,0 +1,57 @@ +"""T8: trigger-rule truth table over upstream statuses.""" + +from __future__ import annotations + +import pytest + +from aws_durable_execution_sdk_python.dag import TaskStatus, TriggerRule +from aws_durable_execution_sdk_python.operation.dag_executor import _trigger_passes + +S = TaskStatus.SUCCEEDED +F = TaskStatus.FAILED +K = TaskStatus.SKIPPED + +ALL_SUCC = [S, S] +ALL_FAIL = [F, F] +MIXED = [S, F] +INCLUDES_SKIP = [S, K] +EMPTY: list[TaskStatus] = [] + + +@pytest.mark.parametrize( + ("rule", "statuses", "expected"), + [ + # ALL_SUCCESS + (TriggerRule.ALL_SUCCESS, ALL_SUCC, True), + (TriggerRule.ALL_SUCCESS, ALL_FAIL, False), + (TriggerRule.ALL_SUCCESS, MIXED, False), + (TriggerRule.ALL_SUCCESS, INCLUDES_SKIP, False), + (TriggerRule.ALL_SUCCESS, EMPTY, True), # root + # ALL_FAILED (len>0 guard) + (TriggerRule.ALL_FAILED, ALL_FAIL, True), + (TriggerRule.ALL_FAILED, ALL_SUCC, False), + (TriggerRule.ALL_FAILED, MIXED, False), + (TriggerRule.ALL_FAILED, EMPTY, False), + # ALL_DONE + (TriggerRule.ALL_DONE, ALL_SUCC, True), + (TriggerRule.ALL_DONE, ALL_FAIL, True), + (TriggerRule.ALL_DONE, MIXED, True), + (TriggerRule.ALL_DONE, INCLUDES_SKIP, True), + (TriggerRule.ALL_DONE, EMPTY, True), + # ANY_SUCCESS + (TriggerRule.ANY_SUCCESS, MIXED, True), + (TriggerRule.ANY_SUCCESS, ALL_FAIL, False), + (TriggerRule.ANY_SUCCESS, EMPTY, False), + # ANY_FAILED + (TriggerRule.ANY_FAILED, MIXED, True), + (TriggerRule.ANY_FAILED, ALL_SUCC, False), + (TriggerRule.ANY_FAILED, EMPTY, False), + # NONE_FAILED + (TriggerRule.NONE_FAILED, ALL_SUCC, True), + (TriggerRule.NONE_FAILED, INCLUDES_SKIP, True), + (TriggerRule.NONE_FAILED, MIXED, False), + (TriggerRule.NONE_FAILED, EMPTY, True), + ], +) +def test_trigger_truth_table(rule, statuses, expected): + assert _trigger_passes(rule, statuses) is expected diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/dag_validator_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_validator_test.py new file mode 100644 index 00000000..49af23a8 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_validator_test.py @@ -0,0 +1,98 @@ +"""T4: DAG validator tests.""" + +from __future__ import annotations + +import pytest + +from aws_durable_execution_sdk_python.dag import DagConfig +from aws_durable_execution_sdk_python.exceptions import ( + DagCyclicDependencyError, + DagDuplicateTaskError, + DagInvalidDependencyError, + DagInvalidTaskNameError, +) +from aws_durable_execution_sdk_python.operation.dag_context import DagContextImpl +from aws_durable_execution_sdk_python.operation.dag_validator import validate_dag +from tests.dag_support import make_context, make_state + + +def _impl() -> DagContextImpl: + state, _ = make_state() + return DagContextImpl(make_context(state, parent_id="c"), DagConfig()) + + +def _step(d, name, deps=None): + return d.step(lambda deps_map, sc: name, deps=deps, name=name) + + +def test_valid_diamond_passes(): + d = _impl() + a = _step(d, "a") + b = _step(d, "b", [a]) + c = _step(d, "c", [a]) + _step(d, "d", [b, c]) + validate_dag(d) # no raise + + +def test_self_loop_detected(): + d = _impl() + a = _step(d, "a") + a.after(a) # self dependency + with pytest.raises(DagCyclicDependencyError): + validate_dag(d) + + +def test_two_cycle_detected(): + d = _impl() + a = _step(d, "a") + b = _step(d, "b", [a]) + a.after(b) # a<->b + with pytest.raises(DagCyclicDependencyError) as exc: + validate_dag(d) + assert "a" in str(exc.value) and "b" in str(exc.value) + + +def test_deep_cycle_detected(): + d = _impl() + a = _step(d, "a") + b = _step(d, "b", [a]) + c = _step(d, "c", [b]) + a.after(c) # a->b->c->a + with pytest.raises(DagCyclicDependencyError): + validate_dag(d) + + +@pytest.mark.parametrize("bad", ["", "a-b", "has space", "x" * 101, "DAG_NODE_T_x"]) +def test_invalid_names(bad): + d = _impl() + # register with a valid name, then tamper the name to bypass registration guards + _step(d, "valid") + d.get_tasks()["valid"].name = bad + with pytest.raises(DagInvalidTaskNameError): + validate_dag(d) + + +def test_duplicate_names(): + d = _impl() + _step(d, "dup") + _step(d, "dup") + with pytest.raises(DagDuplicateTaskError): + validate_dag(d) + + +def test_foreign_scope_dep(): + d1 = _impl() + foreign = _step(d1, "foreign") + + d2 = _impl() + _step(d2, "a") + d2.get_tasks()["a"].all_deps.append(foreign) + with pytest.raises(DagInvalidDependencyError): + validate_dag(d2) + + +def test_valid_names_pass(): + d = _impl() + _step(d, "abc_123") + _step(d, "A9") + validate_dag(d)