From 17445ba0674edd0f6b679eff247f89a54538d15d Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 24 Jul 2026 10:05:42 -0700 Subject: [PATCH 01/25] feat(dag): add core seam for experimental context.dag() - context.py: add dag(), name-based _create_task_id, and _run_step_with_task_id explicit-id step seam (bypasses _replay_aware, mirrors map/parallel child-id pattern) - exceptions.py: Dag{Execution,CyclicDependency,InvalidTaskName,DuplicateTask,InvalidDependency}Error - lambda_service.py: OperationSubType.DAG (CONTEXT-typed container) - concurrency/executor.py: make TimerScheduler generic (SupportsResume protocol) so the DAG scheduler can reuse it with a name-keyed timed-resume adapter (backward compatible; ExecutableWithState already satisfies the protocol) - __init__.py: export public DAG surface --- .../__init__.py | 32 + .../context.py | 87 +++ .../aws_durable_execution_sdk_python/dag.py | 486 ++++++++++++++ .../exceptions.py | 46 ++ .../lambda_service.py | 2 + .../operation/dag.py | 168 +++++ .../operation/dag_context.py | 421 ++++++++++++ .../operation/dag_executor.py | 563 ++++++++++++++++ .../operation/dag_result.py | 232 +++++++ .../operation/dag_validator.py | 106 +++ .../tests/conformance/__init__.py | 0 .../tests/conformance/dag_conformance_test.py | 594 +++++++++++++++++ .../tests/dag_support.py | 155 +++++ .../tests/operation/dag_context_test.py | 101 +++ .../tests/operation/dag_executor_test.py | 612 ++++++++++++++++++ .../tests/operation/dag_handler_test.py | 197 ++++++ .../tests/operation/dag_result_test.py | 111 ++++ .../tests/operation/dag_seam_test.py | 77 +++ .../tests/operation/dag_task_kinds_test.py | 119 ++++ .../tests/operation/dag_test.py | 131 ++++ .../tests/operation/dag_trigger_rules_test.py | 57 ++ .../tests/operation/dag_validator_test.py | 98 +++ 22 files changed, 4395 insertions(+) create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dag.py create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag.py create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_context.py create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_executor.py create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_result.py create mode 100644 packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_validator.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/conformance/__init__.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/conformance/dag_conformance_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/dag_support.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_context_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_executor_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_handler_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_result_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_seam_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_task_kinds_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_trigger_rules_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_validator_test.py 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..b1bd30bf 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,11 @@ CallbackSubmitterError, CallbackTimeoutError, ChildContextError, + DagCyclicDependencyError, + DagDuplicateTaskError, + DagExecutionError, + DagInvalidDependencyError, + DagInvalidTaskNameError, DurableExecutionsError, DurableOperationError, ExecutionError, @@ -48,6 +65,16 @@ "CallbackSubmitterError", "CallbackTimeoutError", "ChildContextError", + "DagCompletionReason", + "DagConfig", + "DagContext", + "DagCyclicDependencyError", + "DagDuplicateTaskError", + "DagExecutionError", + "DagInvalidDependencyError", + "DagInvalidTaskNameError", + "DagResult", + "DepsMap", "DurableContext", "DurableExecutionsError", "DurableOperationError", @@ -55,8 +82,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..bbdeba41 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,61 @@ 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 + DAG task ids disjoint from counter-based sibling ids: counter ids are + opaque blake2b hex digests (``[0-9a-f]{64}``) while a task id is the + plain string ``{prefix}-DAG_NODE_T_{name}`` (uppercase + underscore, so + it can never equal a lowercase-hex digest). 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). + + .. warning:: + **Experimental.** Internal implementation detail; the id format is + subject to change without notice. + """ + prefix = self._parent_id + return f"{prefix}-DAG_NODE_T_{name}" if prefix else f"DAG_NODE_T_{name}" + + 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 +1005,36 @@ 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( + run_in_child_context=self.run_in_child_context, + state=self.state, + 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..ba37e818 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/dag.py @@ -0,0 +1,486 @@ +"""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, + InvokeConfig, + MapConfig, + ParallelConfig, + StepConfig, + WaitForCallbackConfig, + ) + from aws_durable_execution_sdk_python.lambda_service import ErrorObject + from aws_durable_execution_sdk_python.retries import RetryDecision + 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. + + .. 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" + + +# 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 DagConfig: + """Configuration for a DAG. + + Reuses the existing threshold-only :class:`~...config.CompletionConfig` + verbatim; result-based custom completion is v2-deferred. + + .. warning:: + **Experimental.** This API is experimental and may be changed or removed + in future releases. + """ + + max_concurrency: int | None = None + completion_config: CompletionConfig | None = None + default_retry_strategy: Callable[[Exception, int], RetryDecision] | None = None + default_trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS + serdes: SerDes | None = None + summary_generator: Callable[[Any], str] | 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``). The handle path + dispatches at runtime on ``isinstance(key, TaskHandle)`` and extracts the + name; it does not rely on hashing. + + .. 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: ... + @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, + seconds: int, + 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: ... # 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..da637b11 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,52 @@ 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 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..b74f869b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag.py @@ -0,0 +1,168 @@ +"""dag_handler: wraps register + validate + schedule inside a child context. + +.. warning:: + **Experimental.** Internal wiring for ``context.dag()``. +""" + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +from aws_durable_execution_sdk_python.config import ChildConfig +from aws_durable_execution_sdk_python.dag import DagConfig +from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + DagCyclicDependencyError, + DagDuplicateTaskError, + DagExecutionError, + DagInvalidDependencyError, + DagInvalidTaskNameError, + ValidationError, +) +from aws_durable_execution_sdk_python.identifier import OperationIdentifier +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.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 ( + create_dag_result_serdes, +) +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, DagResult + from aws_durable_execution_sdk_python.state import ExecutionState + +_DAG_VALIDATION_ERRORS = ( + DagCyclicDependencyError, + DagInvalidTaskNameError, + DagDuplicateTaskError, + DagInvalidDependencyError, + DagExecutionError, +) + +_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) + + +def _run_dag_body( + dag_child_ctx: DurableContext, + register: Callable[[DagContext], None], + config: DagConfig, +) -> DagResult: + dag_ctx = DagContextImpl(dag_child_ctx, config) + register(dag_ctx) + validate_dag(dag_ctx) + return DagExecutor(dag_child_ctx, dag_ctx.get_tasks(), config).run() + + +def unwrap_dag_error(exc: ChildContextError) -> None: + """Re-raise the typed Dag* cause of a wrapped ``ChildContextError``. + + ``child_handler`` 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): + raise cause from None + 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( + run_in_child_context: Callable[..., DagResult], + state: ExecutionState, + name: str | None, + register: Callable[[DagContext], None], + config: DagConfig | None, +) -> DagResult: + """Run a DAG as a child context and return its DagResult synchronously.""" + config = config or DagConfig() + _check_max_concurrency(config) + + def body(dag_child_ctx: DurableContext) -> DagResult: + return _run_dag_body(dag_child_ctx, register, config) + + child_config: ChildConfig = ChildConfig( + sub_type=OperationSubType.DAG, + serdes=config.serdes or create_dag_result_serdes(), + summary_generator=config.summary_generator, + ) + try: + return run_in_child_context(body, name, child_config) + 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) + + def body() -> DagResult: + child = ctx.create_child_context(operation_id=task_id) + return _run_dag_body(child, register, config) + + child_config: ChildConfig = ChildConfig( + sub_type=OperationSubType.DAG, + serdes=config.serdes or create_dag_result_serdes(), + summary_generator=config.summary_generator, + ) + try: + return child_handler( + func=body, + state=ctx.state, + operation_identifier=OperationIdentifier( + operation_id=task_id, + sub_type=OperationSubType.DAG, + parent_id=ctx._parent_id, + name=name, + ), + config=child_config, + ) + except ChildContextError as e: + unwrap_dag_error(e) + raise # pragma: no cover 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..62fecfab --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_context.py @@ -0,0 +1,421 @@ +"""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 +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.config import StepConfig + 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 or self._default_step_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): + 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, + ).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): + from aws_durable_execution_sdk_python.operation.callback import ( + wait_for_callback_handler, + ) + + def body(child: DurableContext): + return wait_for_callback_handler( + child, + lambda cb_id: submitter(deps_map, cb_id, child), + task_name, + config, + ) + + return self._run_child( + ctx, + task_name, + body, + ChildConfig(sub_type=OperationSubType.WAIT_FOR_CALLBACK), + OperationSubType.WAIT_FOR_CALLBACK, + ) + + return self._add( + task_name, "wait_for_callback", deps, trigger_rule, run_if, config, executor + ) + + def wait( + self, seconds, 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 + + 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, + ) + + def _default_step_config(self) -> StepConfig | None: + from aws_durable_execution_sdk_python.config import StepConfig + + if self._config.default_retry_strategy is not None: + return StepConfig(retry_strategy=self._config.default_retry_strategy) + return None 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..ec74aef4 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_executor.py @@ -0,0 +1,563 @@ +"""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 failure by default (failure is a terminal +state, not an abort), and computes DAG-global success/failure/skip counts, +feeding only success+failure into the reused threshold ``CompletionConfig``. + +.. warning:: + **Experimental.** Internal implementation of the DAG scheduler. +""" + +from __future__ import annotations + +import heapq +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 ( + DagCompletionReason, + DepsMap, + SkipReason, + TaskExecution, + TaskStatus, +) +from aws_durable_execution_sdk_python.exceptions import ( + 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) + +# task scheduling decisions +_RUN = "RUN" +_SKIP = "SKIP" +_FAIL = "FAIL" + + +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() + 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) -> DagResultImpl: + """Schedule and run the DAG; return a DagResult (may raise to suspend).""" + total = len(self._tasks) + if total == 0: + return DagResultImpl({}, DagCompletionReason.ALL_COMPLETED) + + max_workers = self._config.max_concurrency or total + # 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() + + # 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)) + elif decision == _SKIP: + self._results[name] = TaskExecution( + name=name, status=TaskStatus.SKIPPED, skip_reason=payload + ) + self._skip += 1 + else: # _FAIL: run_if raised — treat as task failure (drain) + self._results[name] = TaskExecution( + name=name, status=TaskStatus.FAILED, error=payload + ) + self._failure += 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) + logger.debug("▶️ DAG task %s starting", name) + return task.executor(self._ctx, deps_map) + + def _on_done(self, name: str, future: Future) -> None: + try: + result = future.result() + with self._lock: + self._results[name] = TaskExecution( + name=name, status=TaskStatus.SUCCEEDED, result=result + ) + 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. + self._pending_suspends.append(se) + self._results[name] = TaskExecution( + name=name, status=TaskStatus.STARTED + ) + 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), + ) + 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: + 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. + 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)``, ``(_SKIP, SkipReason)`` + or ``(_FAIL, ErrorObject)``. + + The trigger rule is a pure function of upstream enum statuses. ``run_if`` + is user code: if it raises, we treat the task as FAILED (spec §5.5's + "raise ⇒ FAILED", drain model) rather than letting the exception escape + the scheduler. This keeps root and non-root ``run_if`` behaviour + identical and never hangs the scheduler. + """ + 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: # noqa: BLE001 + return (_FAIL, ErrorObject.from_exception(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, mirroring ``ExecutionCounters.should_complete``. + + 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 + 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 _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: + # 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) -> DagResultImpl: + if 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()} + return DagResultImpl( + dict(self._results), reason, task_kinds, total_count=len(self._tasks) + ) 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..106005c9 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_result.py @@ -0,0 +1,232 @@ +"""DagResultImpl + serialization for the DAG operation. + +.. warning:: + **Experimental.** Internal implementation of :class:`~...dag.DagResult`. +""" + +from __future__ import annotations + +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") + +# 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 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, + ) -> 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) + + # region accessors + @overload + def get_result(self, task: TaskHandle[T]) -> T: ... + @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``); passing a name + string returns ``Any``. Both resolve by task name at runtime. + """ + 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: + return len(self.succeeded()) + + @property + def failure_count(self) -> int: + return len(self.failed()) + + @property + def skipped_count(self) -> int: + 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.""" + 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) + + # endregion accessors + + # region serialization + def to_dict(self) -> dict[str, Any]: + return { + "results": { + name: self._task_to_dict(te) for name, te in self._results.items() + }, + "completionReason": self._completion_reason.value, + "totalCount": self._total_count, + } + + 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 + return { + "name": te.name, + "status": te.status.value, + "resultKind": kind, + "skipReason": te.skip_reason.value if te.skip_reason else None, + "result": result_value, + "error": te.error.to_dict() if te.error else None, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> DagResultImpl: + results: dict[str, TaskExecution] = {} + task_kinds: dict[str, str] = {} + for name, td in data["results"].items(): + 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=td["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, + ) + 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"), + ) + + # endregion serialization + + +class DagResultSerDes(SerDes): + """SerDes for the DagResult container payload.""" + + 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..9a09cbc9 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/conformance/dag_conformance_test.py @@ -0,0 +1,594 @@ +"""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): DAG-1..17 and DAG-19 apply to Python (18 +records). DAG-18 (custom result-based completion) is TS+Go only and is OMITTED +here per the catalog's "emit only applicable scenarios" convention. + +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 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 +from aws_durable_execution_sdk_python.dag import ( + DagConfig, + DagResult, + TaskStatus, + TriggerRule, +) +from aws_durable_execution_sdk_python.retries import RetryPresets +from tests.dag_support import InMemoryServiceClient, make_context, make_state + +NO_RETRY = RetryPresets.none() + +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. + + Task ops are exactly the operations whose id contains the reserved + ``DAG_NODE_T_`` delimiter; counter/other ops never do. For the empty DAG + there are no task ops, so all four are vacuously ``true``. + """ + task_ids = [oid for oid in client.operations if "DAG_NODE_T_" in oid] + counter_ids = [oid for oid in client.operations if "DAG_NODE_T_" not in oid] + if not task_ids: + return { + "name_based": True, + "has_delimiter": True, + "dash_free": True, + "disjoint_from_counter": True, + } + short_names = [oid.split("DAG_NODE_T_")[-1] for oid in task_ids] + name_based = all( + client.operations[oid].name == oid.split("DAG_NODE_T_")[-1] + for oid in task_ids + ) + has_delimiter = all("DAG_NODE_T_" in oid for oid in task_ids) + dash_free = all(_NAME_CHARSET.fullmatch(s) is not None for s in short_names) + disjoint = set(task_ids).isdisjoint(counter_ids) and all( + "DAG_NODE_T_" not in cid for cid in 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") + 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, DagConfig(default_retry_strategy=NO_RETRY)) + 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") + 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", DagConfig(default_retry_strategy=NO_RETRY)) + 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") + u2 = d.step(_fail, name="u2") + 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", DagConfig(default_retry_strategy=NO_RETRY)) + 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}") + 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), + default_retry_strategy=NO_RETRY, + ), + ) + 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_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. + + DAG-18 is intentionally OMITTED (TS+Go only, catalog Part C). Emits exactly + 18 records as key-sorted UTF-8 JSON with 2-space indent + trailing newline. + """ + expected_scenarios = {f"DAG-{i}" for i in range(1, 18)} | {"DAG-19"} + assert set(RECORDS) == expected_scenarios, ( + f"missing/extra: {expected_scenarios ^ set(RECORDS)}" + ) + assert "DAG-18" not in RECORDS # NOT-APPLICABLE for Python (custom completion) + assert len(RECORDS) == 18 + + 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") + assert "DAG-18" not in reloaded + + +# 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..acb9940f --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/dag_support.py @@ -0,0 +1,155 @@ +"""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.""" + 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), + ) + 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_context_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_context_test.py new file mode 100644 index 00000000..4a3f46b6 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_context_test.py @@ -0,0 +1,101 @@ +"""T3: DagContext registration, TaskHandle chaining, DepsMap access.""" + +from __future__ import annotations + +import pytest + +from aws_durable_execution_sdk_python.dag import ( + DagConfig, + DepsMap, + TriggerRule, +) +from aws_durable_execution_sdk_python.exceptions import DagInvalidTaskNameError +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(5) + + +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..819490bd --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_executor_test.py @@ -0,0 +1,612 @@ +"""T5: DagExecutor scheduler tests.""" + +from __future__ import annotations + +import threading +import time + +import pytest + +from aws_durable_execution_sdk_python.config import CompletionConfig +from aws_durable_execution_sdk_python.dag import ( + DagCompletionReason, + DagConfig, + SkipReason, + TaskStatus, + TriggerRule, +) +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() + + +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") + # default ALL_SUCCESS -> skipped because a FAILED + d.step(lambda deps, sc: 1, deps=[a], name="b") + + result, _ = run_dag(register, DagConfig(default_retry_strategy=NO_RETRY)) + 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") + # 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, DagConfig(default_retry_strategy=NO_RETRY)) + 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_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_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}") + + result, _ = run_dag( + register, + DagConfig( + default_retry_strategy=NO_RETRY, + 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") + d.step(lambda deps, sc: ran.__setitem__("b", True), name="b") + + result, _ = run_dag(register, DagConfig(default_retry_strategy=NO_RETRY)) + 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") + + result, _ = run_dag(register, DagConfig(default_retry_strategy=NO_RETRY)) + 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") + d.step(lambda deps, sc: 1, name="b") + + result, _ = run_dag( + register, + DagConfig( + default_retry_strategy=NO_RETRY, + 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") + # 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_retry_strategy=NO_RETRY, + 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 regression (worker-thread callback swallowed exceptions) +import signal # noqa: E402 +from contextlib import contextmanager # 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 regression)") + + old = signal.signal(signal.SIGALRM, _handler) + signal.alarm(seconds) + try: + yield + finally: + signal.alarm(0) + signal.signal(signal.SIGALRM, old) + + +def test_run_if_raises_on_non_root_fails_task_and_drains(): + """A run_if that raises on a downstream task (evaluated inside a worker-thread + completion callback) must not hang: the task FAILS and the DAG drains.""" + + def register(d): + a = d.step(lambda deps, sc: "A", name="a") + # run_if dereferences a missing dep -> KeyError, evaluated after `a` done + d.step( + lambda deps, sc: "ran", + deps=[a], + name="b", + run_if=lambda deps: deps["missing"] > 0, + ) + # independent root task must still run (drain, not fail-fast) + d.step(lambda deps, sc: "c-ran", name="c") + + with _fail_on_hang(): + result, _ = run_dag(register, DagConfig(default_retry_strategy=NO_RETRY)) + + assert result.get_status("a") is TaskStatus.SUCCEEDED + assert result.get_status("b") is TaskStatus.FAILED + assert result.results["b"].error is not None + assert result.get_result("c") == "c-ran" + assert result.failure_count == 1 + assert result.completion_reason is DagCompletionReason.COMPLETED_WITH_FAILURES + with pytest.raises(DagExecutionError): + result.throw_if_error() + + +def test_run_if_raises_on_root_fails_task(): + """A raising run_if on a root task fails that task (consistent with the + non-root path) rather than aborting the whole DAG.""" + + def register(d): + d.step( + lambda deps, sc: "ran", + name="a", + run_if=lambda deps: 1 // 0 == 0, + ) + d.step(lambda deps, sc: "b-ran", name="b") + + with _fail_on_hang(): + result, _ = run_dag(register, DagConfig(default_retry_strategy=NO_RETRY)) + + assert result.get_status("a") is TaskStatus.FAILED + assert result.get_result("b") == "b-ran" + assert result.failure_count == 1 + + +# endregion run_if-raises regression + + +# 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 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..05e444a5 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_handler_test.py @@ -0,0 +1,197 @@ +"""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 ( + DagCompletionReason, + DagConfig, + TaskStatus, +) +from aws_durable_execution_sdk_python.exceptions import DagCyclicDependencyError +from aws_durable_execution_sdk_python.retries import RetryPresets +from tests.dag_support import make_context, make_state + +NO_RETRY = RetryPresets.none() + + +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_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", + ]: + assert hasattr(sdk, symbol), symbol + + +def test_summary_generator_wired_into_child_config(): + """DagConfig.summary_generator is passed through to the container ChildConfig.""" + from aws_durable_execution_sdk_python.operation.dag import dag_handler + + captured = {} + + def fake_run_in_child_context(body, name, child_config): + captured["config"] = child_config + + def gen(_result): # pragma: no cover - not invoked (small payload) + return "summary" + + state, _ = make_state() + dag_handler( + run_in_child_context=fake_run_in_child_context, + state=state, + name="p", + register=lambda d: None, + config=DagConfig(summary_generator=gen), + ) + assert captured["config"].summary_generator is gen + + +def test_nested_dag_summary_generator_wired(): + """run_nested_dag builds a container ChildConfig carrying summary_generator.""" + import aws_durable_execution_sdk_python.operation.dag as dag_mod + + captured = {} + + def fake_child_handler(func, state, operation_identifier, config): + captured["config"] = config + return func() + + def gen(_result): # pragma: no cover - not invoked (small payload) + return "nested-summary" + + original = dag_mod.child_handler + dag_mod.child_handler = fake_child_handler # type: ignore[assignment] + try: + state, _ = make_state() + ctx = make_context(state) + dag_mod.run_nested_dag( + ctx, + "inner", + lambda d: d.step(lambda deps, sc: 1, name="x"), + DagConfig(summary_generator=gen), + ) + finally: + dag_mod.child_handler = original # type: ignore[assignment] + assert captured["config"].summary_generator is gen + + + +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 ( + CallableRuntimeError, + DagExecutionError, + ) + from aws_durable_execution_sdk_python.operation.dag import unwrap_dag_error + + exc = CallableRuntimeError( + 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 CallableRuntimeError + from aws_durable_execution_sdk_python.operation.dag import unwrap_dag_error + + exc = CallableRuntimeError( + message="boom", error_type="ValueError", data=None, stack_trace=None + ) + with pytest.raises(CallableRuntimeError): + unwrap_dag_error(exc) 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..1d6cf67f --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_result_test.py @@ -0,0 +1,111 @@ +"""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 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..db5d5906 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_seam_test.py @@ -0,0 +1,77 @@ +"""T2: name-based entity-id seam tests.""" + +from __future__ import annotations + +from tests.dag_support import make_context, make_state + + +def test_create_task_id_unprefixed(): + state, _ = make_state() + ctx = make_context(state, parent_id=None) + assert ctx._create_task_id("fetch") == "DAG_NODE_T_fetch" + + +def test_create_task_id_prefixed(): + state, _ = make_state() + ctx = make_context(state, parent_id="container") + assert ( + ctx._create_task_id("fetch") == "container-DAG_NODE_T_fetch" + ) + + +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 ids are {prefix}-{int}; task ids are {prefix}-DAG_NODE_T_{name}.""" + 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 == "c-DAG_NODE_T_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 "dagc-DAG_NODE_T_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 + {container}-DAG_NODE_T_{name} per level, never a single raw multi-level string + like ...DAG_NODE_T_validation-DAG_NODE_T_rule_a composed at one level.""" + state, _ = make_state() + outer = make_context(state, parent_id="root") + container_id = outer._create_task_id("validation") + assert container_id == "root-DAG_NODE_T_validation" + # nested scope: the container id is the child's parent prefix + nested = make_context(state, parent_id=container_id) + sub_id = nested._create_task_id("rule_a") + assert sub_id == "root-DAG_NODE_T_validation-DAG_NODE_T_rule_a" + # the sub-task prefix is exactly the container id (composed per level) + assert sub_id.startswith(container_id + "-DAG_NODE_T_") 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..0d6ae168 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_task_kinds_test.py @@ -0,0 +1,119 @@ +"""T8/T9: exercise every DagContext task-kind executor closure end-to-end.""" + +from __future__ import annotations + +import pytest + +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(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..0e39fd56 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py @@ -0,0 +1,131 @@ +"""T9: DAG replay / interruption / large-payload re-execute tests.""" + +from __future__ import annotations + +import pytest + +from aws_durable_execution_sdk_python.dag import ( + DagCompletionReason, + DagConfig, + 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 = client.operations["1"] + assert container.sub_type is OperationSubType.DAG + # step task checkpointed under a DAG_NODE_T_ id + assert "1-DAG_NODE_T_a" in client.operations + # 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_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 + assert "1-DAG_NODE_T_skipme" not in client.operations + + # 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 "1-DAG_NODE_T_skipme" not in client.operations + + +def test_large_payload_reexecutes_to_equal_result(): + """A >256KB DagResult forces ReplayChildren; the DAG re-executes to an equal + result and a custom summary_generator neither changes it nor hangs replay.""" + 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", + ) + + config = DagConfig(summary_generator=lambda r: "custom-summary-string") + state, client = make_state() + + r1 = make_context(state).dag(register, name="p", config=config) + assert side["big"] == 1 + # container stored with replay_children due to large payload + container = client.operations["1"] + assert container.context_details.replay_children is True + + # replay: container re-executes body (ReplayChildren), tasks fast-path + r2 = make_context(state).dag(register, name="p", config=config) + 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) From f50a38dbc9d73c85707560e60ddba5be1ccbeef4 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 24 Jul 2026 10:05:53 -0700 Subject: [PATCH 02/25] feat(dag): implement DAG primitive wired to current core - dag.py: public API (DagContext, TaskHandle, DepsMap, DagResult, DagConfig, TriggerRule/TaskStatus/SkipReason/DagCompletionReason, TaskExecution) - operation/dag_executor.py: dedicated topological scheduler mirroring ConcurrentExecutor's worker-thread model (ThreadPoolExecutor + TimerScheduler, add_done_callback, drain-not-fail-fast, timed-earliest suspend precedence) - operation/dag_context.py: name-based explicit-id task runners for every kind - operation/dag_result.py: DagResult impl + result_kind serdes (batch/dag nesting) - operation/dag_validator.py: name/duplicate/foreign-dep/cycle validation - operation/dag.py: dag_handler (child-context wrap + FutureWarning + error unwrap) --- .../src/aws_durable_execution_sdk_python/dag.py | 5 ++++- .../operation/dag_context.py | 6 ++++-- 2 files changed, 8 insertions(+), 3 deletions(-) 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 index ba37e818..82ed031c 100644 --- 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 @@ -44,6 +44,9 @@ from aws_durable_execution_sdk_python.types import ( DurableContext, ) + from aws_durable_execution_sdk_python.types import ( + SummaryGenerator, + ) from aws_durable_execution_sdk_python.waits import WaitForConditionConfig T = TypeVar("T") @@ -151,7 +154,7 @@ class DagConfig: default_retry_strategy: Callable[[Exception, int], RetryDecision] | None = None default_trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS serdes: SerDes | None = None - summary_generator: Callable[[Any], str] | None = None + summary_generator: SummaryGenerator[Any] | None = None # region handle & deps-map 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 index 62fecfab..fd6cb0d1 100644 --- 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 @@ -175,6 +175,8 @@ def invoke( 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, @@ -186,7 +188,7 @@ def executor(ctx: DurableContext, deps_map: DepsMap): parent_id=ctx._parent_id, name=task_name, ), - config=config, + config=config or InvokeConfig(), ).process() return self._add( @@ -207,7 +209,7 @@ def executor(ctx: DurableContext, deps_map: DepsMap): def body(child: DurableContext): return wait_for_callback_handler( child, - lambda cb_id: submitter(deps_map, cb_id, child), + lambda cb_id, cb_ctx: submitter(deps_map, cb_id, cb_ctx), task_name, config, ) From d2c619cf10de70b3130ead20bbf4ee0675e1a8c5 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 24 Jul 2026 10:05:53 -0700 Subject: [PATCH 03/25] test(dag): DAG unit + conformance suite adapted to current core - dag_support.py: in-memory harness starts checkpoint_batches_forever daemon thread (as the real runtime does in execution.py) so tasks' synchronous checkpoints complete on worker threads instead of deadlocking - 122 bounded tests: validator, result, trigger rules, context, executor, handler, task kinds, replay/large-payload, seam, and conformance scenarios --- .../tests/dag_support.py | 16 +++++++++++++++- .../tests/operation/dag_handler_test.py | 10 +++++----- .../tests/operation/dag_test.py | 15 +++++++++++---- 3 files changed, 31 insertions(+), 10 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python/tests/dag_support.py b/packages/aws-durable-execution-sdk-python/tests/dag_support.py index acb9940f..ea4116d0 100644 --- a/packages/aws-durable-execution-sdk-python/tests/dag_support.py +++ b/packages/aws-durable-execution-sdk-python/tests/dag_support.py @@ -132,7 +132,15 @@ def stop(self, execution_arn, payload) -> datetime.datetime: # pragma: no cover def make_state( client: InMemoryServiceClient | None = None, ) -> tuple[ExecutionState, InMemoryServiceClient]: - """Build a fresh ExecutionState wired to an in-memory client.""" + """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", @@ -141,6 +149,12 @@ def make_state( 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 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 index 05e444a5..e03c1d3d 100644 --- 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 @@ -170,12 +170,12 @@ def test_unwrap_dag_error_reconstructs_typed_error_on_replay(): 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 ( - CallableRuntimeError, + ChildContextError, DagExecutionError, ) from aws_durable_execution_sdk_python.operation.dag import unwrap_dag_error - exc = CallableRuntimeError( + exc = ChildContextError( message="2 task(s) FAILED", error_type="DagExecutionError", data=None, @@ -187,11 +187,11 @@ def test_unwrap_dag_error_reconstructs_typed_error_on_replay(): def test_unwrap_dag_error_passthrough_for_non_dag_error(): - from aws_durable_execution_sdk_python.exceptions import CallableRuntimeError + from aws_durable_execution_sdk_python.exceptions import ChildContextError from aws_durable_execution_sdk_python.operation.dag import unwrap_dag_error - exc = CallableRuntimeError( + exc = ChildContextError( message="boom", error_type="ValueError", data=None, stack_trace=None ) - with pytest.raises(CallableRuntimeError): + with pytest.raises(ChildContextError): unwrap_dag_error(exc) 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 index 0e39fd56..95d2984e 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py @@ -38,10 +38,13 @@ def register(d): assert side["a"] == 1 assert r1.get_result("a") == "A" - container = client.operations["1"] + 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 DAG_NODE_T_ id - assert "1-DAG_NODE_T_a" in client.operations + # step task checkpointed under a DAG_NODE_T_ id (prefixed by the hashed + # container id in the current core, so match by suffix) + assert any(k.endswith("-DAG_NODE_T_a") for k in client.operations) # the serialized container payload round-trips to an equal DagResult payload = container.context_details.result restored = create_dag_result_serdes().deserialize(payload, None) @@ -120,7 +123,11 @@ def register(d): r1 = make_context(state).dag(register, name="p", config=config) assert side["big"] == 1 # container stored with replay_children due to large payload - container = client.operations["1"] + 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 # replay: container re-executes body (ReplayChildren), tasks fast-path From bbd5875661880c656220a506c2e5bea5ccf36f84 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 24 Jul 2026 11:19:36 -0700 Subject: [PATCH 04/25] test(dag): cloud integration tests --- .../examples-catalog.json | 44 +++++++++++ .../src/dag/dag_compensation.py | 74 +++++++++++++++++++ .../src/dag/dag_diamond.py | 51 +++++++++++++ .../src/dag/dag_run_if.py | 68 +++++++++++++++++ .../src/dag/dag_wait_resume.py | 50 +++++++++++++ .../test/dag/test_dag_compensation.py | 37 ++++++++++ .../test/dag/test_dag_diamond.py | 35 +++++++++ .../test/dag/test_dag_run_if.py | 37 ++++++++++ .../test/dag/test_dag_wait_resume.py | 35 +++++++++ 9 files changed, 431 insertions(+) create mode 100644 packages/aws-durable-execution-sdk-python-examples/src/dag/dag_compensation.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/src/dag/dag_diamond.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/src/dag/dag_run_if.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/src/dag/dag_wait_resume.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_compensation.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_diamond.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_run_if.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/test/dag/test_dag_wait_resume.py 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..9f28a5ea --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_compensation.py @@ -0,0 +1,74 @@ +"""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.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 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: + if charge_ok: + charge = d.step(lambda deps, sc: "charged", name="charge") + else: + 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: "audited", name="audit").after(charge).trigger_rule( + TriggerRule.ALL_DONE + ) + + # Disable retries so an intentionally failing charge terminates promptly. + result = context.dag( + register, + name="compensation", + config=DagConfig(default_retry_strategy=RetryPresets.none()), + ) + 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..06a744b5 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/dag/dag_wait_resume.py @@ -0,0 +1,50 @@ +"""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.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(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" From 5108feb85eb83dcb719f12b743e079716b99a3b8 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 24 Jul 2026 11:27:38 -0700 Subject: [PATCH 05/25] fix(dag): bound task operation id length for backend --- .../context.py | 27 +++++++--- .../tests/conformance/dag_conformance_test.py | 52 ++++++++++++++----- .../tests/operation/dag_seam_test.py | 52 +++++++++++++------ .../tests/operation/dag_test.py | 21 +++++--- 4 files changed, 111 insertions(+), 41 deletions(-) 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 bbdeba41..2fd083dd 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 @@ -496,19 +496,32 @@ def _create_task_id(self, name: str) -> str: 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 - DAG task ids disjoint from counter-based sibling ids: counter ids are - opaque blake2b hex digests (``[0-9a-f]{64}``) while a task id is the - plain string ``{prefix}-DAG_NODE_T_{name}`` (uppercase + underscore, so - it can never equal a lowercase-hex digest). 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 *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 - return f"{prefix}-DAG_NODE_T_{name}" if prefix else f"DAG_NODE_T_{name}" + 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, 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 index 9a09cbc9..fb1c158e 100644 --- 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 @@ -17,6 +17,7 @@ from __future__ import annotations +import hashlib import json import re from pathlib import Path @@ -58,29 +59,54 @@ def _fail(_deps: Any, _sc: Any) -> Any: def _structural_checks(client: InMemoryServiceClient) -> dict[str, bool]: """Compute the four per-language structural entity-ID invariants. - Task ops are exactly the operations whose id contains the reserved - ``DAG_NODE_T_`` delimiter; counter/other ops never do. For the empty DAG - there are no task ops, so all four are vacuously ``true``. + 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``. """ - task_ids = [oid for oid in client.operations if "DAG_NODE_T_" in oid] - counter_ids = [oid for oid in client.operations if "DAG_NODE_T_" not in oid] - if not task_ids: + + 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, } - short_names = [oid.split("DAG_NODE_T_")[-1] for oid in task_ids] + # name_based: id is exactly the blake2b bound of the name-based pre-image. name_based = all( - client.operations[oid].name == oid.split("DAG_NODE_T_")[-1] - for oid in task_ids + op.operation_id + == hashlib.blake2b(task_preimage(op).encode()).hexdigest()[:64] + for op in task_ops ) - has_delimiter = all("DAG_NODE_T_" in oid for oid in task_ids) - dash_free = all(_NAME_CHARSET.fullmatch(s) is not None for s in short_names) - disjoint = set(task_ids).isdisjoint(counter_ids) and all( - "DAG_NODE_T_" not in cid for cid in counter_ids + # 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, 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 index db5d5906..8217e5fe 100644 --- 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 @@ -2,21 +2,36 @@ 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) - assert ctx._create_task_id("fetch") == "DAG_NODE_T_fetch" + 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") - assert ( - ctx._create_task_id("fetch") == "container-DAG_NODE_T_fetch" - ) + 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(): @@ -29,13 +44,16 @@ def test_create_task_id_does_not_touch_counter(): def test_no_collision_with_counter_ids(): - """Counter ids are {prefix}-{int}; task ids are {prefix}-DAG_NODE_T_{name}.""" + """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 == "c-DAG_NODE_T_1" + assert task_id == _task_digest("c", "1") def test_seam_checkpoints_under_task_id_and_fast_path_on_replay(): @@ -53,7 +71,7 @@ def body(_step_ctx): result = ctx._run_step_with_task_id("mytask", body) assert result == "value" assert calls["n"] == 1 - assert "dagc-DAG_NODE_T_mytask" in client.operations + 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) @@ -62,16 +80,20 @@ def body(_step_ctx): 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 - {container}-DAG_NODE_T_{name} per level, never a single raw multi-level string - like ...DAG_NODE_T_validation-DAG_NODE_T_rule_a composed at one level.""" + """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 == "root-DAG_NODE_T_validation" - # nested scope: the container id is the child's parent prefix + 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") - assert sub_id == "root-DAG_NODE_T_validation-DAG_NODE_T_rule_a" - # the sub-task prefix is exactly the container id (composed per level) - assert sub_id.startswith(container_id + "-DAG_NODE_T_") + # 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_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py index 95d2984e..394a5c86 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py @@ -42,9 +42,18 @@ def register(d): 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 DAG_NODE_T_ id (prefixed by the hashed - # container id in the current core, so match by suffix) - assert any(k.endswith("-DAG_NODE_T_a") for k in client.operations) + # 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) @@ -88,8 +97,8 @@ def gate(deps, sc): 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 - assert "1-DAG_NODE_T_skipme" not in client.operations + # 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 @@ -98,7 +107,7 @@ def gate(deps, sc): assert side["work"] == 1 # remaining task ran exactly once assert result.get_result("work") == "done" assert result.get_status("skipme") is TaskStatus.SKIPPED - assert "1-DAG_NODE_T_skipme" not in client.operations + assert not any(o.name == "skipme" for o in client.operations.values()) def test_large_payload_reexecutes_to_equal_result(): From e460101436db369cb3bb3e108f2b64a6728d1b4a Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 24 Jul 2026 19:11:30 -0700 Subject: [PATCH 06/25] test(dag): conformance handlers + SAM template for dag suite (10-1..10-4) Add 4 DAG conformance handlers (diamond, compensation, run_if, wait_resume) returning the canonical cross-language result shapes, and template_dag.yaml mapping each to requirement ids 10-1..10-4 via TestingMetadata.TestDescription with the durable role/policies and DurableConfig. Runner (account 730758745077, us-west-2): all 4 result shapes PASS (0 result errors); history FAILs on the missing DagTask CONTEXT-op layer (see report). No SDK change: aligning to the DagTask contract is a large redesign, reported rather than forced. Existing DAG unit tests remain green (44 passed). --- .../handlers/dag/__init__.py | 0 .../handlers/dag/dag_compensation.py | 56 ++++++++++++ .../handlers/dag/dag_diamond.py | 40 +++++++++ .../handlers/dag/dag_run_if.py | 52 +++++++++++ .../handlers/dag/dag_wait_resume.py | 42 +++++++++ .../template_dag.yaml | 90 +++++++++++++++++++ 6 files changed, 280 insertions(+) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/__init__.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_compensation.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_diamond.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_run_if.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_wait_resume.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml 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_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_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_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_wait_resume.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_wait_resume.py new file mode 100644 index 00000000..664dd242 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_wait_resume.py @@ -0,0 +1,42 @@ +"""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.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(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..9a614fac --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml @@ -0,0 +1,90 @@ +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 + 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 From 5f0064bddc83cd21c49a9b49b7f6a45e2ee2d818 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 24 Jul 2026 19:30:17 -0700 Subject: [PATCH 07/25] fix(dag): materialize tasks as DagTask CONTEXT ops for cross-language conformance --- .../context.py | 32 ------------ .../lambda_service.py | 2 + .../operation/dag_context.py | 41 ++++++++++----- .../tests/operation/dag_seam_test.py | 51 ++++++++++++++----- .../tests/operation/dag_test.py | 24 ++++++--- 5 files changed, 85 insertions(+), 65 deletions(-) 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 2fd083dd..8f490d31 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 @@ -523,38 +523,6 @@ def _create_task_id(self, name: str) -> str: ) 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 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 9da2ad70..fa30750f 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 @@ -93,6 +93,7 @@ def from_sub_type(cls, sub_type: OperationSubType) -> OperationType: | OperationSubType.PARALLEL | OperationSubType.PARALLEL_BRANCH | OperationSubType.DAG + | OperationSubType.DAG_TASK ): return OperationType.CONTEXT case _: @@ -117,6 +118,7 @@ class OperationSubType(Enum): WAIT_FOR_CONDITION = "WaitForCondition" CHAINED_INVOKE = "ChainedInvoke" DAG = "Dag" + DAG_TASK = "DagTask" class InvocationStatus(Enum): 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 index fd6cb0d1..59e83e47 100644 --- 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 @@ -29,7 +29,6 @@ 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, ) @@ -162,8 +161,22 @@ def step( cfg = config or self._default_step_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 + # Materialize the task as a DagTask CONTEXT op (name-based id, + # parented to the Dag container) and run the inner step nested one + # level beneath it (child, counter-based id). + return self._run_child( + ctx, + task_name, + lambda child: child.step( + lambda step_ctx: func(deps_map, step_ctx), + name=task_name, + config=cfg, + ), + ChildConfig( + sub_type=OperationSubType.DAG_TASK, + serdes=cfg.serdes if cfg else None, + ), + OperationSubType.DAG_TASK, ) return self._add(task_name, "step", deps, trigger_rule, run_if, cfg, executor) @@ -236,16 +249,20 @@ def wait( task_name = name 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, + from aws_durable_execution_sdk_python.config import Duration + + # Materialize the task as a DagTask CONTEXT op with the wait nested + # one level beneath it. The name-based DagTask id is stable across + # the suspend/resume boundary so the nested wait replays deterministically. + return self._run_child( + ctx, + task_name, + lambda child: child.wait( + Duration.from_seconds(seconds), name=task_name ), - ).process() + ChildConfig(sub_type=OperationSubType.DAG_TASK), + OperationSubType.DAG_TASK, + ) return self._add(task_name, "wait", deps, trigger_rule, run_if, None, executor) 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 index 8217e5fe..740369be 100644 --- 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 @@ -56,26 +56,49 @@ def test_no_collision_with_counter_ids(): 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") +def test_step_task_materializes_dagtask_and_fast_path_on_replay(): + """A DAG step task materializes as a DagTask CONTEXT op (name-based id) with + its inner Step nested beneath, and on replay hits the checkpoint fast path. + + Replaces the pre-conformance flat seam assertion (step checkpointed directly + under the container with no wrapper) with the canonical + Dag -> DagTask -> Step structure that matches the cross-language contract. + """ + from aws_durable_execution_sdk_python.lambda_service import OperationSubType + state, client = make_state() calls = {"n": 0} - def body(_step_ctx): - calls["n"] += 1 - return "value" + def register(d): + d.step( + lambda deps, sc: calls.__setitem__("n", calls["n"] + 1) or "value", + name="mytask", + ) - # first call runs and checkpoints under the name-based id - result = ctx._run_step_with_task_id("mytask", body) - assert result == "value" + # first run: task runs once and materializes the DagTask + nested step + result = make_context(state).dag(register, name="p") + assert result.get_result("mytask") == "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" + dag_task = next( + o + for o in client.operations.values() + if o.name == "mytask" and o.sub_type is OperationSubType.DAG_TASK + ) + # DagTask id is the name-based digest under its parent (the Dag container) + assert dag_task.operation_id == _task_digest(dag_task.parent_id, "mytask") + inner_step = next( + o + for o in client.operations.values() + if o.name == "mytask" and o.sub_type is OperationSubType.STEP + ) + # the inner step nests one level beneath the DagTask (distinct child id) + assert inner_step.parent_id == dag_task.operation_id + assert inner_step.operation_id != dag_task.operation_id + + # second run (simulated replay): DagTask fast-path, inner step not re-executed + result2 = make_context(state).dag(register, name="p") + assert result2.get_result("mytask") == "value" assert calls["n"] == 1 # not re-executed 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 index 394a5c86..4b2b9aae 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py @@ -42,18 +42,28 @@ def register(d): 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). + # Canonical structure: the task materializes as a DagTask CONTEXT op + # (name-based id, parented to the Dag container) with the inner Step nested + # one level beneath it (ParentId = the DagTask op, counter-based child id). import hashlib - task_a = next( + dag_task = next( o for o in client.operations.values() - if o.name == "a" and o.sub_type is not OperationSubType.DAG + if o.name == "a" and o.sub_type is OperationSubType.DAG_TASK ) - 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 + preimage = f"{dag_task.parent_id}-DAG_NODE_T_a" + assert dag_task.operation_id == hashlib.blake2b(preimage.encode()).hexdigest()[:64] + assert len(dag_task.operation_id) <= 64 + assert dag_task.parent_id == container.operation_id + # the inner step is a STEP op nested under the DagTask op + inner_step = next( + o + for o in client.operations.values() + if o.name == "a" and o.sub_type is OperationSubType.STEP + ) + assert inner_step.parent_id == dag_task.operation_id + assert inner_step.operation_id != dag_task.operation_id # the serialized container payload round-trips to an equal DagResult payload = container.context_details.result restored = create_dag_result_serdes().deserialize(payload, None) From 02e636e9cf36771cb6be10e2ec2750a3b42612c0 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 24 Jul 2026 21:44:00 -0700 Subject: [PATCH 08/25] revert(dag): restore flat name-based task materialization --- .../context.py | 32 ++++++++++++ .../lambda_service.py | 2 - .../operation/dag_context.py | 41 +++++---------- .../tests/operation/dag_seam_test.py | 51 +++++-------------- .../tests/operation/dag_test.py | 24 +++------ 5 files changed, 65 insertions(+), 85 deletions(-) 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 8f490d31..2fd083dd 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 @@ -523,6 +523,38 @@ def _create_task_id(self, name: str) -> str: ) 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 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 fa30750f..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 @@ -93,7 +93,6 @@ def from_sub_type(cls, sub_type: OperationSubType) -> OperationType: | OperationSubType.PARALLEL | OperationSubType.PARALLEL_BRANCH | OperationSubType.DAG - | OperationSubType.DAG_TASK ): return OperationType.CONTEXT case _: @@ -118,7 +117,6 @@ class OperationSubType(Enum): WAIT_FOR_CONDITION = "WaitForCondition" CHAINED_INVOKE = "ChainedInvoke" DAG = "Dag" - DAG_TASK = "DagTask" class InvocationStatus(Enum): 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 index 59e83e47..fd6cb0d1 100644 --- 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 @@ -29,6 +29,7 @@ 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, ) @@ -161,22 +162,8 @@ def step( cfg = config or self._default_step_config() def executor(ctx: DurableContext, deps_map: DepsMap): - # Materialize the task as a DagTask CONTEXT op (name-based id, - # parented to the Dag container) and run the inner step nested one - # level beneath it (child, counter-based id). - return self._run_child( - ctx, - task_name, - lambda child: child.step( - lambda step_ctx: func(deps_map, step_ctx), - name=task_name, - config=cfg, - ), - ChildConfig( - sub_type=OperationSubType.DAG_TASK, - serdes=cfg.serdes if cfg else None, - ), - OperationSubType.DAG_TASK, + 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) @@ -249,20 +236,16 @@ def wait( task_name = name def executor(ctx: DurableContext, deps_map: DepsMap): - from aws_durable_execution_sdk_python.config import Duration - - # Materialize the task as a DagTask CONTEXT op with the wait nested - # one level beneath it. The name-based DagTask id is stable across - # the suspend/resume boundary so the nested wait replays deterministically. - return self._run_child( - ctx, - task_name, - lambda child: child.wait( - Duration.from_seconds(seconds), name=task_name + 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, ), - ChildConfig(sub_type=OperationSubType.DAG_TASK), - OperationSubType.DAG_TASK, - ) + ).process() return self._add(task_name, "wait", deps, trigger_rule, run_if, None, executor) 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 index 740369be..8217e5fe 100644 --- 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 @@ -56,49 +56,26 @@ def test_no_collision_with_counter_ids(): assert task_id == _task_digest("c", "1") -def test_step_task_materializes_dagtask_and_fast_path_on_replay(): - """A DAG step task materializes as a DagTask CONTEXT op (name-based id) with - its inner Step nested beneath, and on replay hits the checkpoint fast path. - - Replaces the pre-conformance flat seam assertion (step checkpointed directly - under the container with no wrapper) with the canonical - Dag -> DagTask -> Step structure that matches the cross-language contract. - """ - from aws_durable_execution_sdk_python.lambda_service import OperationSubType - +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 register(d): - d.step( - lambda deps, sc: calls.__setitem__("n", calls["n"] + 1) or "value", - name="mytask", - ) + def body(_step_ctx): + calls["n"] += 1 + return "value" - # first run: task runs once and materializes the DagTask + nested step - result = make_context(state).dag(register, name="p") - assert result.get_result("mytask") == "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 - dag_task = next( - o - for o in client.operations.values() - if o.name == "mytask" and o.sub_type is OperationSubType.DAG_TASK - ) - # DagTask id is the name-based digest under its parent (the Dag container) - assert dag_task.operation_id == _task_digest(dag_task.parent_id, "mytask") - inner_step = next( - o - for o in client.operations.values() - if o.name == "mytask" and o.sub_type is OperationSubType.STEP - ) - # the inner step nests one level beneath the DagTask (distinct child id) - assert inner_step.parent_id == dag_task.operation_id - assert inner_step.operation_id != dag_task.operation_id - - # second run (simulated replay): DagTask fast-path, inner step not re-executed - result2 = make_context(state).dag(register, name="p") - assert result2.get_result("mytask") == "value" + # 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 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 index 4b2b9aae..394a5c86 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py @@ -42,28 +42,18 @@ def register(d): o for o in client.operations.values() if o.sub_type is OperationSubType.DAG ) assert container.sub_type is OperationSubType.DAG - # Canonical structure: the task materializes as a DagTask CONTEXT op - # (name-based id, parented to the Dag container) with the inner Step nested - # one level beneath it (ParentId = the DagTask op, counter-based child id). + # 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 - dag_task = next( + task_a = next( o for o in client.operations.values() - if o.name == "a" and o.sub_type is OperationSubType.DAG_TASK + if o.name == "a" and o.sub_type is not OperationSubType.DAG ) - preimage = f"{dag_task.parent_id}-DAG_NODE_T_a" - assert dag_task.operation_id == hashlib.blake2b(preimage.encode()).hexdigest()[:64] - assert len(dag_task.operation_id) <= 64 - assert dag_task.parent_id == container.operation_id - # the inner step is a STEP op nested under the DagTask op - inner_step = next( - o - for o in client.operations.values() - if o.name == "a" and o.sub_type is OperationSubType.STEP - ) - assert inner_step.parent_id == dag_task.operation_id - assert inner_step.operation_id != dag_task.operation_id + 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) From 21fd3640041c08e0cce601b3f8ef7b889e9b53b0 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 24 Jul 2026 22:11:01 -0700 Subject: [PATCH 09/25] test(dag): conformance handlers for tier-1 task kinds --- .../handlers/dag/dag_child.py | 49 ++++++++++++ .../handlers/dag/dag_map.py | 50 +++++++++++++ .../handlers/dag/dag_nested.py | 55 ++++++++++++++ .../handlers/dag/dag_parallel.py | 60 +++++++++++++++ .../handlers/dag/dag_wait_for_condition.py | 60 +++++++++++++++ .../template_dag.yaml | 75 +++++++++++++++++++ 6 files changed, 349 insertions(+) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_child.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_map.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_nested.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_parallel.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_wait_for_condition.py 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_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_parallel.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_parallel.py new file mode 100644 index 00000000..0c6578f0 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_parallel.py @@ -0,0 +1,60 @@ +"""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] -> +successful results joined by "-" = "L-R"). 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. 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), + ) + d.step( + lambda deps, sc: "-".join(deps[fork].get_results()), + 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_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/template_dag.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml index 9a614fac..a6f88126 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml @@ -88,3 +88,78 @@ Resources: 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 From c2dae48abe193dc3f6bc220e0dc251a3e409ddd7 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sat, 25 Jul 2026 00:33:26 -0700 Subject: [PATCH 10/25] fix(dag): callback task materializes as a Callback container A DAG callback task emitted a bare WaitForCallback context directly under the DAG container, one nesting level shallower than the cross-language contract. It now wraps the native context.wait_for_callback in a Callback-subtype container carrying the task's name-based id, producing Callback -> WaitForCallback -> {CallbackStarted, submitter step}. Caught by conformance scenario dag/10-11; see DAG_SPEC_CROSS_LANGUAGE.md 2.A.5. --- .../operation/dag_context.py | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) 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 index fd6cb0d1..d2d7bf0a 100644 --- 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 @@ -202,24 +202,27 @@ def wait_for_callback( task_name = _resolve_name(name, submitter) def executor(ctx: DurableContext, deps_map: DepsMap): - from aws_durable_execution_sdk_python.operation.callback import ( - wait_for_callback_handler, - ) - + # 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 wait_for_callback_handler( - child, + return child.wait_for_callback( lambda cb_id, cb_ctx: submitter(deps_map, cb_id, cb_ctx), - task_name, - config, + name=task_name, + config=config, ) return self._run_child( ctx, task_name, body, - ChildConfig(sub_type=OperationSubType.WAIT_FOR_CALLBACK), - OperationSubType.WAIT_FOR_CALLBACK, + ChildConfig(sub_type=OperationSubType.CALLBACK), + OperationSubType.CALLBACK, ) return self._add( From da375eb73fdcf82a3ed7c01047e2ca38e39c0ceb Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sat, 25 Jul 2026 00:33:26 -0700 Subject: [PATCH 11/25] test(dag): conformance handlers for invoke and callback task kinds Adds dag_invoke (10-10) and dag_callback (10-11) handlers plus the echo target and function resources they need, and reworks dag_parallel (10-7) to read only the parallel task's aggregate result per the updated cross-language contract. --- .../handlers/dag/dag_callback.py | 69 +++++++++++++++++++ .../handlers/dag/dag_invoke.py | 58 ++++++++++++++++ .../handlers/dag/dag_parallel.py | 29 +++++--- .../template_dag.yaml | 51 ++++++++++++++ 4 files changed, 197 insertions(+), 10 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_callback.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_invoke.py 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_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_parallel.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_parallel.py index 0c6578f0..334c5d2f 100644 --- 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 @@ -1,10 +1,18 @@ """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] -> -successful results joined by "-" = "L-R"). 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. Returns -the canonical summary defined by test-requirements/dag/10-7.yaml. +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 @@ -43,11 +51,12 @@ def right(ctx: DurableContext) -> str: name="fork", config=ParallelConfig(max_concurrency=1), ) - d.step( - lambda deps, sc: "-".join(deps[fork].get_results()), - deps=[fork], - name="join", - ) + + 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) 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 index a6f88126..f1c6b47a 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml @@ -27,6 +27,7 @@ Resources: Action: - lambda:CheckpointDurableExecution - lambda:GetDurableExecutionState + - lambda:InvokeFunction Resource: '*' DagDiamond: Type: AWS::Serverless::Function @@ -163,3 +164,53 @@ Resources: 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 From 1170255d2aba6670f9f5a337df92872420107c92 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sat, 25 Jul 2026 15:17:47 -0700 Subject: [PATCH 12/25] feat(dag): a throwing run_if aborts the DAG with DagPredicateError Reverses this SDK's previous behaviour, which caught a raising run_if and recorded the task FAILED, citing "spec 5.5 raise => FAILED" - but 5.5 governs the task BODY, not the predicate, so that justification was wrong. The four SDKs had diverged 2-2 on this (review finding H5) and the contract is now uniform: run_if is a pure deterministic predicate, so a raise is a DEFECT. The offending task gets no terminal state, no further tasks start, and dag() raises DagPredicateError naming the task. Recording it as a task failure would fire every downstream ALL_FAILED / ANY_FAILED / ALL_DONE compensation task, letting a predicate bug issue a refund. In-flight tasks are drained rather than cancelled - there is no cancellation in the model. A raise on a worker or timer thread is captured into the scheduler exception, the completion event is set, and run() re-raises inside the timer and pool context managers so already-running bodies are joined and the timer thread is torn down. A stopping guard prevents new tasks starting during the drain, and a sibling's outcome cannot downgrade the abort. Known boundary erasure, documented in the docstring: at the scheduler level the error carries .task_name and .__cause__, but through dag() the durable child-context boundary rebuilds it from serialized type and message only, so both are None there and the task name survives in the message. This is the same erasure the whole Dag* family already has, kept deliberately consistent rather than special-cased, and it is identical on first run and replay. Behavior change: a raising run_if now raises DagPredicateError out of dag() instead of recording the task as FAILED. The feature is unreleased. --- .../__init__.py | 2 + .../exceptions.py | 33 +++++++ .../operation/dag.py | 15 ++- .../operation/dag_executor.py | 53 ++++++---- .../tests/operation/dag_executor_test.py | 87 ++++++++++------ .../tests/operation/dag_handler_test.py | 99 +++++++++++++++++++ 6 files changed, 242 insertions(+), 47 deletions(-) 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 b1bd30bf..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 @@ -40,6 +40,7 @@ DagExecutionError, DagInvalidDependencyError, DagInvalidTaskNameError, + DagPredicateError, DurableExecutionsError, DurableOperationError, ExecutionError, @@ -73,6 +74,7 @@ "DagExecutionError", "DagInvalidDependencyError", "DagInvalidTaskNameError", + "DagPredicateError", "DagResult", "DepsMap", "DurableContext", 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 da637b11..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 @@ -302,6 +302,39 @@ class DagInvalidDependencyError(ValidationError): """ +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/operation/dag.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag.py index b74f869b..ec720748 100644 --- 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 @@ -18,6 +18,7 @@ DagExecutionError, DagInvalidDependencyError, DagInvalidTaskNameError, + DagPredicateError, ValidationError, ) from aws_durable_execution_sdk_python.identifier import OperationIdentifier @@ -37,12 +38,17 @@ from aws_durable_execution_sdk_python.dag import DagContext, DagResult from aws_durable_execution_sdk_python.state import 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} @@ -98,7 +104,14 @@ class name (both first run and replay, reconstructed from the checkpoint). """ cause = exc.__cause__ if isinstance(cause, _DAG_VALIDATION_ERRORS): - raise cause from None + # 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 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 index ec74aef4..c7e99712 100644 --- 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 @@ -4,9 +4,12 @@ ``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 failure by default (failure is a terminal -state, not an abort), and computes DAG-global success/failure/skip counts, -feeding only success+failure into the reused threshold ``CompletionConfig``. +``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. @@ -30,6 +33,7 @@ TaskStatus, ) from aws_durable_execution_sdk_python.exceptions import ( + DagPredicateError, SuspendExecution, TimedSuspendExecution, ValidationError, @@ -53,7 +57,6 @@ # task scheduling decisions _RUN = "RUN" _SKIP = "SKIP" -_FAIL = "FAIL" def _trigger_passes(rule, statuses: list[TaskStatus]) -> bool: @@ -288,16 +291,11 @@ def _pump(self) -> None: if decision == _RUN: self._in_flight.add(name) to_submit.append((name, task)) - elif decision == _SKIP: + 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 - else: # _FAIL: run_if raised — treat as task failure (drain) - self._results[name] = TaskExecution( - name=name, status=TaskStatus.FAILED, error=payload - ) - self._failure += 1 progressed = True done = self._is_done_locked() @@ -445,14 +443,23 @@ def _deps_terminal_locked(self, name: str) -> bool: return True def _evaluate_locked(self, task: TaskDef) -> tuple[str, Any]: - """Decide a ready task's fate: ``(_RUN, None)``, ``(_SKIP, SkipReason)`` - or ``(_FAIL, ErrorObject)``. + """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 code: if it raises, we treat the task as FAILED (spec §5.5's - "raise ⇒ FAILED", drain model) rather than letting the exception escape - the scheduler. This keeps root and non-root ``run_if`` behaviour - identical and never hangs the scheduler. + 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): @@ -461,8 +468,12 @@ def _evaluate_locked(self, task: TaskDef) -> tuple[str, Any]: deps_map = self._build_deps_map(task) try: should_run = task.run_if(deps_map) - except Exception as e: # noqa: BLE001 - return (_FAIL, ErrorObject.from_exception(e)) + 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) @@ -523,6 +534,12 @@ def _has_indefinite_locked(self) -> bool: ) 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 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 index 819490bd..ca724b98 100644 --- 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 @@ -299,10 +299,12 @@ def register(d): assert result.get_status("b") is TaskStatus.SKIPPED -# region run_if-raises regression (worker-thread callback swallowed exceptions) +# 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): @@ -310,7 +312,7 @@ def _fail_on_hang(seconds: int = 10): whole test session. SIGALRM fires on the main thread (where pytest runs).""" def _handler(_signum, _frame): - raise AssertionError("DagExecutor.run() hung (run_if regression)") + raise AssertionError("DagExecutor.run() hung (run_if abort regression)") old = signal.signal(signal.SIGALRM, _handler) signal.alarm(seconds) @@ -321,56 +323,85 @@ def _handler(_signum, _frame): signal.signal(signal.SIGALRM, old) -def test_run_if_raises_on_non_root_fails_task_and_drains(): +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) must not hang: the task FAILS and the DAG drains.""" + 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 - d.step( + b = d.step( lambda deps, sc: "ran", deps=[a], name="b", run_if=lambda deps: deps["missing"] > 0, ) - # independent root task must still run (drain, not fail-fast) - d.step(lambda deps, sc: "c-ran", name="c") + # 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, + ) - with _fail_on_hang(): - result, _ = run_dag(register, DagConfig(default_retry_strategy=NO_RETRY)) + ex = _make_executor(register, DagConfig(default_retry_strategy=NO_RETRY)) + with _fail_on_hang(), pytest.raises(DagPredicateError) as ei: + ex.run() - assert result.get_status("a") is TaskStatus.SUCCEEDED - assert result.get_status("b") is TaskStatus.FAILED - assert result.results["b"].error is not None - assert result.get_result("c") == "c-ran" - assert result.failure_count == 1 - assert result.completion_reason is DagCompletionReason.COMPLETED_WITH_FAILURES - with pytest.raises(DagExecutionError): - result.throw_if_error() + 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_fails_task(): - """A raising run_if on a root task fails that task (consistent with the - non-root path) rather than aborting the whole DAG.""" +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): - d.step( + a = d.step( lambda deps, sc: "ran", name="a", run_if=lambda deps: 1 // 0 == 0, ) - d.step(lambda deps, sc: "b-ran", name="b") + # 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, + ) - with _fail_on_hang(): - result, _ = run_dag(register, DagConfig(default_retry_strategy=NO_RETRY)) + ex = _make_executor(register, DagConfig(default_retry_strategy=NO_RETRY)) + with _fail_on_hang(), pytest.raises(DagPredicateError) as ei: + ex.run() - assert result.get_status("a") is TaskStatus.FAILED - assert result.get_result("b") == "b-ran" - assert result.failure_count == 1 + 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 regression +# endregion run_if-raises abort # region threshold-completion fidelity (mirrors ExecutionCounters.should_complete) 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 index e03c1d3d..c775e171 100644 --- 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 @@ -108,6 +108,7 @@ def test_public_exports(): "DagInvalidTaskNameError", "DagDuplicateTaskError", "DagInvalidDependencyError", + "DagPredicateError", ]: assert hasattr(sdk, symbol), symbol @@ -195,3 +196,101 @@ def test_unwrap_dag_error_passthrough_for_non_dag_error(): ) 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(default_retry_strategy=NO_RETRY), + ) + + # 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) From bf6dcbe927e84506adc080d08694cbdfbeb09447 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sat, 25 Jul 2026 16:35:04 -0700 Subject: [PATCH 13/25] fix(dag): no checkpoint after the abort decision Round-2 review, empirically reproduced: the DAG-owned TimerScheduler daemon thread is torn down only after the pool drains - the scheduler is the outer context manager, the pool the inner - so during the drain window a pending timed resume fired _resubmit() and wrote a checkpoint without consulting _scheduler_exception. Exactly one late checkpoint was observed. No task double-execution, so the blast radius was a stray flush, but it contradicted the abort contract's own docstring. _resubmit now early-returns when the abort has been decided. This is the only scheduler entry point that checkpoints; _on_done and _pump are already gated by _stopping_locked, which consults _scheduler_exception first. Three tests, including a control that asserts a checkpoint IS written when not aborting. Verified sensitive: removing the guard fails both abort tests. --- .../operation/dag_executor.py | 14 ++ .../tests/operation/dag_executor_test.py | 157 ++++++++++++++++++ 2 files changed, 171 insertions(+) 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 index c7e99712..61932be4 100644 --- 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 @@ -389,6 +389,20 @@ def _resubmit(self, resumes: list[_TimedResume]) -> None: 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. + if self._scheduler_exception is not None: + return for resume in resumes: name = resume.name if name not in self._pending_timers: 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 index ca724b98..946ba015 100644 --- 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 @@ -641,3 +641,160 @@ def timer(_ctx, _deps): # 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 + + +# endregion no checkpoint after abort (teardown-window regression) From 9cc269b4b25b42eb3fc29624e47cf904a2146e10 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sat, 25 Jul 2026 19:22:40 -0700 Subject: [PATCH 14/25] test(dag): conformance handlers and integration tests for concurrency and abort Handlers for DagRunIfAbort (10-12), DagConcurrentOverlap (10-13) and DagConcurrentSuspend (10-14) with matching template resources. Tasks run on real threads here, so the 10-13 peak-concurrency counter is lock-guarded. The overlap integration test proves the overlap rather than assuming it: both task bodies must clear a threading.Barrier(2) before either returns, which is impossible unless they were simultaneously in flight on separate threads - serialisation would time out instead of passing. It also recomputes each task's blake2b name-based id and asserts equality, the check the cloud suite cannot make. The suspend test drives the resumes explicitly, marking each wait SUCCEEDED in resume order, and asserts afterFast ran before afterSlow - the reverse of registration order across the suspend. --- .../handlers/dag/dag_concurrent_overlap.py | 100 +++++ .../handlers/dag/dag_concurrent_suspend.py | 64 ++++ .../handlers/dag/dag_run_if_abort.py | 65 ++++ .../template_dag.yaml | 48 +++ .../dag_concurrency_coverage_test.py | 342 ++++++++++++++++++ 5 files changed, 619 insertions(+) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_concurrent_overlap.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_concurrent_suspend.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_run_if_abort.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_concurrency_coverage_test.py 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..1f29d2ee --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_concurrent_suspend.py @@ -0,0 +1,64 @@ +"""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.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(8, deps=[root], name="slow") # registered FIRST + fast = d.wait(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_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/template_dag.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml index f1c6b47a..6419fb51 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml @@ -214,3 +214,51 @@ Resources: 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 + + 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 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..608ec17b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_concurrency_coverage_test.py @@ -0,0 +1,342 @@ +"""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.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(8, deps=[root], name="slow") # registered FIRST + fast = d.wait(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 From 591562666d95f637ae0f6c464600ef4d227f3225 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sat, 25 Jul 2026 20:43:53 -0700 Subject: [PATCH 15/25] feat(dag): default maxConcurrency to 40 Closes the surviving half of review finding H2. Every SDK previously ran a DAG's top-level tasks unbounded when no concurrency was configured. The change exists for Python, where it is a resource bound rather than a scheduling preference: the executor sized its ThreadPoolExecutor from the same number, so an unbounded 500-task DAG spawned 500 OS threads inside a Lambda sandbox. JS promises on one event loop and Go goroutines are far cheaper, but a single shared value is chosen over per-language tuning so that a graph behaves identically in every language, which is the point of the cross-language contract. Scope is the DAG scheduler only, one level. It is deliberately NOT inherited by a task's internal fan-out - a map or parallel task keeps its own default, which stays unlimited - and a nested dag task resolves its own independent 40. An explicit value always wins, including above 40; the <= 0 registration error is unchanged. Two consequences are documented rather than left to be discovered: the divergence from map and parallel, and that a graph wider than 40 under completionConfig shifts tasks from a terminal state to absent (never started), which count only toward totalCount. Tests assert the OBSERVED peak rather than a config value, since asserting the config would pass against an unbounded scheduler: a graph wider than 40 must both stay at or under 40 and actually reach it, so the test fails if the default reverts to unbounded or if the scheduler serialises. Wire-invisible for the conformance suite: 10-1..10-12 pin maxConcurrency to 1 and 10-13/10-14 leave it unset with only 6 tasks, so 40 never binds. Behavior change: a DAG with no configured concurrency now runs at most 40 tasks concurrently instead of unbounded. The feature is unreleased. --- .../aws_durable_execution_sdk_python/dag.py | 13 ++ .../operation/dag_executor.py | 21 ++- .../tests/operation/dag_executor_test.py | 135 ++++++++++++++++++ 3 files changed, 168 insertions(+), 1 deletion(-) 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 index 82ed031c..27497bba 100644 --- 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 @@ -150,6 +150,19 @@ class DagConfig: """ 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 | None = None default_retry_strategy: Callable[[Exception, int], RetryDecision] | None = None default_trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS 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 index 61932be4..86733f70 100644 --- 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 @@ -54,6 +54,18 @@ _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" @@ -252,7 +264,14 @@ def run(self) -> DagResultImpl: if total == 0: return DagResultImpl({}, DagCompletionReason.ALL_COMPLETED) - max_workers = self._config.max_concurrency or total + # 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). 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 index 946ba015..e053b256 100644 --- 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 @@ -798,3 +798,138 @@ def _root(name, executor): # 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 From ff05eac5487d325c89b98a38ade746f07026805c Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sat, 25 Jul 2026 21:29:15 -0700 Subject: [PATCH 16/25] test(sdk): large-payload conformance handler and integration tests DagLargePayload (scenario 10-15): eight tasks each returning 51200 characters, an aggregate around 410KB that exceeds the 256KB checkpoint limit, then a 2s wait so the next invocation replays the completed container - which is the only way to exercise the offload replay path. Integration tests assert aggregate fidelity across the container replay with a full 51200-character value checked byte-for-byte, and that every task body runs EXACTLY ONCE across the offload and replay. That second one is the assertion that matters: under re-execution the scheduler re-runs, so if a task body were invoked again a customer's side effect would happen twice. --- .../handlers/dag/dag_large_payload.py | 110 ++++++++++ .../template_dag.yaml | 16 ++ .../dag_large_payload_coverage_test.py | 200 ++++++++++++++++++ 3 files changed, 326 insertions(+) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_large_payload.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_large_payload_coverage_test.py 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/template_dag.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml index 6419fb51..86427de2 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml @@ -262,3 +262,19 @@ Resources: 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 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..4ed8423f --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_large_payload_coverage_test.py @@ -0,0 +1,200 @@ +"""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 marks the container ``ReplayChildren=true`` and +writes no envelope (unlike JS, which writes an SDK-owned ``DagSummary``). + +The reconstruct-vs-re-execute 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 + re-execution the DAG child body 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 re-execution path was actually taken** -- Python uses ``ReplayChildren``, + not an envelope, so the container operation is asserted to carry + ``replay_children is True`` with an empty checkpointed payload (no envelope). + This is the observable hook that distinguishes Python's strategy from JS's. +""" + +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 re-execution path was actually taken ---------------------------- + # Python offloads via ReplayChildren with NO envelope. The container op + # carries replay_children=True (proving the aggregate exceeded the 256KB + # limit and was offloaded) and an empty checkpointed payload (proving no + # DagSummary envelope was written -- the JS-only strategy). + container = _dag_container_op(client) + assert container.context_details is not None + assert container.context_details.replay_children is True + assert (container.context_details.result or "") == "" + + # Corroborating evidence that the interesting path was genuinely exercised: + # the invocation actually suspended and resumed (two invocations), and the + # DAG body was re-executed on the replay (register ran twice) rather than + # reconstructed from an envelope. + assert invocations["n"] == 2, "container was not replayed across a suspend" + assert register_calls["n"] == 2, "DAG body was not re-executed on replay" From 92fdafde0559e1cabeb3bbe0dc24dae068cd3d3e Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 26 Jul 2026 20:57:33 -0700 Subject: [PATCH 17/25] feat(sdk): converge the DAG container checkpoint envelope One envelope for the DAG container checkpoint, written identically by all four SDKs, in both the inline and the offloaded case. The payload is returned by GetExecutionHistory and rendered in the AWS console, so it is a customer-facing contract. It was not being treated as one: the same three-task DAG produced four different shapes. JS and Java used a tasks array while Python used a results map, Java spelled resultKind in uppercase and added a resultType field the others lacked, Java used camelCase error keys against PascalCase elsewhere, only JS populated timestamps, and Go wrote no payload at all because it set ReplayChildren unconditionally with no size test and no serializer. Conformance asserted ContextSucceededDetails as a wildcard, which is exactly why four shapes lived undetected. The envelope carries type, the four counts, completionReason, startedTaskNames and failedTaskNames, plus an optional tasks array. Only tasks is optional: its absence is the signal that per-task detail exceeded the checkpoint limit and now lives in the retained child operations. Every other field is always present, with explicit nulls rather than omissions, so the two cases share one shape and a console reader sees the same fields either way. Degradation is ordered: full envelope, then drop tasks and set ReplayChildren, then drop failedTaskNames, and counts plus completionReason plus startedTaskNames never drop. startedTaskNames is bounded by maxConcurrency, so a DAG can never fail to checkpoint because its own summary did not fit. There is no schemaVersion. Evolution is additive only, so every reader ignores unknown fields and treats a missing field as absent, which each SDK now proves with a test. Replay is one strategy everywhere. With tasks present, deserialize and return. Without, reconstruct: re-run the deterministic registration graph, recompute skip and trigger decisions, read each result from its child checkpoint, and take counts, reason and the in-flight set from the envelope. Task bodies never re-execute in either case. summaryGenerator and summary_generator are deleted. The envelope is readable on its own, so a customer-supplied summary string is no longer written into a payload the SDK parses back, which also removes the corruption vector that regression guard existed for. --- .../context.py | 3 +- .../aws_durable_execution_sdk_python/dag.py | 4 - .../operation/dag.py | 293 +++++++++++++++--- .../operation/dag_executor.py | 96 +++++- .../operation/dag_result.py | 121 +++++++- .../tests/operation/dag_handler_test.py | 105 ++++--- .../dag_large_payload_coverage_test.py | 46 ++- .../tests/operation/dag_result_test.py | 96 ++++++ .../dag_started_set_reconstruct_test.py | 132 ++++++++ .../tests/operation/dag_test.py | 51 ++- 10 files changed, 797 insertions(+), 150 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_started_set_reconstruct_test.py 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 2fd083dd..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 @@ -1042,8 +1042,7 @@ def dag( emit_experimental_warning_once() resolved_name: str | None = self._resolve_step_name(name, register) return dag_handler( - run_in_child_context=self.run_in_child_context, - state=self.state, + ctx=self, name=resolved_name, register=register, config=config, 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 index 27497bba..9ac6950f 100644 --- 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 @@ -44,9 +44,6 @@ from aws_durable_execution_sdk_python.types import ( DurableContext, ) - from aws_durable_execution_sdk_python.types import ( - SummaryGenerator, - ) from aws_durable_execution_sdk_python.waits import WaitForConditionConfig T = TypeVar("T") @@ -167,7 +164,6 @@ class DagConfig: default_retry_strategy: Callable[[Exception, int], RetryDecision] | None = None default_trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS serdes: SerDes | None = None - summary_generator: SummaryGenerator[Any] | None = None # region handle & deps-map 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 index ec720748..13a16fd2 100644 --- 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 @@ -1,4 +1,17 @@ -"""dag_handler: wraps register + validate + schedule inside a child context. +"""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()``. @@ -6,11 +19,13 @@ from __future__ import annotations +import json import warnings +from dataclasses import dataclass from typing import TYPE_CHECKING -from aws_durable_execution_sdk_python.config import ChildConfig -from aws_durable_execution_sdk_python.dag import DagConfig +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, @@ -19,24 +34,29 @@ DagInvalidDependencyError, DagInvalidTaskNameError, DagPredicateError, + InvocationError, + SuspendExecution, ValidationError, ) from aws_durable_execution_sdk_python.identifier import OperationIdentifier -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.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 ( - create_dag_result_serdes, -) +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, DagResult - from aws_durable_execution_sdk_python.state import ExecutionState + 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 @@ -75,25 +95,191 @@ def _check_max_concurrency(config: DagConfig) -> None: 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) - return DagExecutor(dag_child_ctx, dag_ctx.get_tasks(), config).run() + 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``. - ``child_handler`` 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. + ``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 @@ -119,29 +305,42 @@ class name (both first run and replay, reconstructed from the checkpoint). def dag_handler( - run_in_child_context: Callable[..., DagResult], - state: ExecutionState, + ctx: DurableContext, name: str | None, register: Callable[[DagContext], None], config: DagConfig | None, ) -> DagResult: - """Run a DAG as a child context and return its DagResult synchronously.""" + """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) - def body(dag_child_ctx: DurableContext) -> DagResult: - return _run_dag_body(dag_child_ctx, register, 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, + ) - child_config: ChildConfig = ChildConfig( - sub_type=OperationSubType.DAG, - serdes=config.serdes or create_dag_result_serdes(), - summary_generator=config.summary_generator, - ) - try: - return run_in_child_context(body, name, child_config) - except ChildContextError as e: - unwrap_dag_error(e) - raise # pragma: no cover - unwrap_dag_error always raises + 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( @@ -154,28 +353,24 @@ def run_nested_dag( 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 body() -> DagResult: + def run_body(reconstruct: _ReconstructInfo | None) -> DagResult: child = ctx.create_child_context(operation_id=task_id) - return _run_dag_body(child, register, config) + return _run_dag_body(child, register, config, reconstruct) - child_config: ChildConfig = ChildConfig( - sub_type=OperationSubType.DAG, - serdes=config.serdes or create_dag_result_serdes(), - summary_generator=config.summary_generator, + executor = DagContainerExecutor( + run_body=run_body, + state=ctx.state, + operation_identifier=identifier, ) try: - return child_handler( - func=body, - state=ctx.state, - operation_identifier=OperationIdentifier( - operation_id=task_id, - sub_type=OperationSubType.DAG, - parent_id=ctx._parent_id, - name=name, - ), - config=child_config, - ) + return executor.process() except ChildContextError as e: unwrap_dag_error(e) - raise # pragma: no cover + 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_executor.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/dag_executor.py index 86733f70..ee2d6a8e 100644 --- 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 @@ -18,6 +18,7 @@ from __future__ import annotations import heapq +import datetime import itertools import logging import threading @@ -236,6 +237,12 @@ def __init__( 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 @@ -258,11 +265,47 @@ def __init__( self._pool: ThreadPoolExecutor | None = None # region public - def run(self) -> DagResultImpl: - """Schedule and run the DAG; return a DagResult (may raise to suspend).""" + 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: - return DagResultImpl({}, DagCompletionReason.ALL_COMPLETED) + 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); @@ -288,7 +331,10 @@ def run(self) -> DagResultImpl: suspend = self._resolve_suspend() if suspend is not None: raise suspend - return self._build_result() + return self._build_result( + reconstruct_reason=reconstruct_reason, + reconstruct_total=reconstruct_total, + ) # endregion public @@ -335,15 +381,24 @@ def _run_task(self, name: str, task: TaskDef) -> Any: # concurrently (the run_if path already builds deps under the lock). with self._lock: deps_map = self._build_deps_map(task) - logger.debug("▶️ DAG task %s starting", name) + # 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 + 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) @@ -351,10 +406,14 @@ def _on_done(self, name: str, future: Future) -> None: schedule_ts: float | None = None with self._lock: # Record every suspend (timed + indefinite); precedence is - # resolved in _resolve_suspend when the DAG stops. + # 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 + 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 @@ -373,6 +432,8 @@ def _on_done(self, name: str, future: Future) -> None: 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) @@ -600,14 +661,27 @@ def _is_done_locked(self) -> bool: # endregion helpers - def _build_result(self) -> DagResultImpl: - if self._early_reason is not None: + 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=len(self._tasks) + 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 index 106005c9..51a36969 100644 --- 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 @@ -6,6 +6,7 @@ from __future__ import annotations +import datetime from typing import TYPE_CHECKING, Any, TypeVar, overload from aws_durable_execution_sdk_python.concurrency.models import ( @@ -31,6 +32,9 @@ T = TypeVar("T") +# The single converged envelope discriminator (contract: ``type: "DagResult"``). +_ENVELOPE_TYPE = "DagResult" + # result_kind discriminators _KIND_PLAIN = "plain" _KIND_BATCH = "batch" @@ -41,6 +45,44 @@ _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) @@ -54,6 +96,7 @@ def _result_kind(task_kind: str | None) -> str: return _KIND_PLAIN + def _name_of(task: str | TaskHandle[Any]) -> str: return task.name if isinstance(task, TaskHandle) else task @@ -153,36 +196,80 @@ def throw_if_error(self) -> None: # 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 { - "results": { - name: self._task_to_dict(te) for name, te in self._results.items() - }, - "completionReason": self._completion_reason.value, + "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): + 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, - "resultKind": kind, "skipReason": te.skip_reason.value if te.skip_reason else None, + "resultKind": serialized_kind, "result": result_value, - "error": te.error.to_dict() if te.error else None, + "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 -- the offloaded path reconstructs 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 name, td in data["results"].items(): + 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: @@ -192,16 +279,20 @@ def from_dict(cls, data: dict[str, Any]) -> DagResultImpl: result_value = cls.from_dict(result_value) error_raw = td.get("error") results[name] = TaskExecution( - name=td["name"], + 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" + task_kinds[name] = ( + "dag" + if kind == _KIND_DAG + else ("map" if kind == _KIND_BATCH else "step") ) return cls( results=results, @@ -214,7 +305,13 @@ def from_dict(cls, data: dict[str, Any]) -> DagResultImpl: class DagResultSerDes(SerDes): - """SerDes for the DagResult container payload.""" + """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 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 index c775e171..e7b3d297 100644 --- 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 @@ -113,57 +113,68 @@ def test_public_exports(): assert hasattr(sdk, symbol), symbol -def test_summary_generator_wired_into_child_config(): - """DagConfig.summary_generator is passed through to the container ChildConfig.""" - from aws_durable_execution_sdk_python.operation.dag import dag_handler - - captured = {} - - def fake_run_in_child_context(body, name, child_config): - captured["config"] = child_config - - def gen(_result): # pragma: no cover - not invoked (small payload) - return "summary" - - state, _ = make_state() - dag_handler( - run_in_child_context=fake_run_in_child_context, - state=state, - name="p", - register=lambda d: None, - config=DagConfig(summary_generator=gen), +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 ) - assert captured["config"].summary_generator is gen - -def test_nested_dag_summary_generator_wired(): - """run_nested_dag builds a container ChildConfig carrying summary_generator.""" - import aws_durable_execution_sdk_python.operation.dag as dag_mod + captured: dict = {} - captured = {} - - def fake_child_handler(func, state, operation_identifier, config): - captured["config"] = config - return func() - - def gen(_result): # pragma: no cover - not invoked (small payload) - return "nested-summary" - - original = dag_mod.child_handler - dag_mod.child_handler = fake_child_handler # type: ignore[assignment] - try: - state, _ = make_state() - ctx = make_context(state) - dag_mod.run_nested_dag( - ctx, - "inner", - lambda d: d.step(lambda deps, sc: 1, name="x"), - DagConfig(summary_generator=gen), - ) - finally: - dag_mod.child_handler = original # type: ignore[assignment] - assert captured["config"].summary_generator is gen + 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(): 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 index 4ed8423f..cd27c7a6 100644 --- 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 @@ -7,10 +7,12 @@ 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 marks the container ``ReplayChildren=true`` and -writes no envelope (unlike JS, which writes an SDK-owned ``DagSummary``). +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-re-execute divergence only fires when a SUCCEEDED container is +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: @@ -21,13 +23,13 @@ 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 - re-execution the DAG child body re-runs, but each task's step operation must + 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 re-execution path was actually taken** -- Python uses ``ReplayChildren``, - not an envelope, so the container operation is asserted to carry - ``replay_children is True`` with an empty checkpointed payload (no envelope). - This is the observable hook that distinguishes Python's strategy from JS's. +* **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 @@ -182,19 +184,29 @@ def run(ctx: DurableContext): for name in _TASK_NAMES: assert calls[name] == 1, f"task {name} body ran {calls[name]} times, expected 1" - # --- The re-execution path was actually taken ---------------------------- - # Python offloads via ReplayChildren with NO envelope. The container op - # carries replay_children=True (proving the aggregate exceeded the 256KB - # limit and was offloaded) and an empty checkpointed payload (proving no - # DagSummary envelope was written -- the JS-only strategy). + # --- 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 - assert (container.context_details.result or "") == "" + 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 body was re-executed on the replay (register ran twice) rather than - # reconstructed from an envelope. + # 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 body was not re-executed on replay" + 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_result_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_result_test.py index 1d6cf67f..6e1bc997 100644 --- 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 @@ -109,3 +109,99 @@ def test_roundtrip_nested_dag_result_kind(): 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_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_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_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py index 394a5c86..9eee5922 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/dag_test.py @@ -6,7 +6,6 @@ from aws_durable_execution_sdk_python.dag import ( DagCompletionReason, - DagConfig, TaskStatus, ) from aws_durable_execution_sdk_python.exceptions import SuspendExecution @@ -61,6 +60,33 @@ def register(d): 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} @@ -111,8 +137,9 @@ def gate(deps, sc): def test_large_payload_reexecutes_to_equal_result(): - """A >256KB DagResult forces ReplayChildren; the DAG re-executes to an equal - result and a custom summary_generator neither changes it nor hangs replay.""" + """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} @@ -126,21 +153,29 @@ def register(d): name="small", ) - config = DagConfig(summary_generator=lambda r: "custom-summary-string") state, client = make_state() - r1 = make_context(state).dag(register, name="p", config=config) + 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 - - # replay: container re-executes body (ReplayChildren), tasks fast-path - r2 = make_context(state).dag(register, name="p", config=config) + # 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 From 3b3acd3255128ece57c30871e6ba318f1bf65b6f Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 26 Jul 2026 22:35:46 -0700 Subject: [PATCH 18/25] feat(sdk): remove the DAG default retry strategy A DAG-level default that four SDKs interpreted four different ways. JS merged it into step and callback configs. Go merged it at the single task-registration point, so every kind picked it up. Python applied it to step only, and only when the task passed no config at all, so a task supplying a StepConfig without a retry silently lost the default. Java never read the field at all, while its javadoc claimed it was applied to tasks that do not specify one, and the published docs advertised it in all three language tabs. Per-task retryStrategy already expresses this, the withRetry helper covers the wrapper case, and a DAG container cannot be retried at all because a child context has no retry, so the name promised something the container could never do. Sharing one policy across tasks is now a local variable passed to each task. Removing it while the API is experimental is free. Keeping it would have meant implementing it in Java, fixing Python to merge rather than overwrite, and defining which of the nine task kinds a retry default reaches, since seven of them ignore it. --- .../handlers/dag/dag_retry.py | 76 +++++++++++++++++++ .../template_dag.yaml | 16 ++++ .../src/dag/dag_compensation.py | 11 +-- .../aws_durable_execution_sdk_python/dag.py | 2 - .../operation/dag_context.py | 10 +-- .../tests/conformance/dag_conformance_test.py | 22 +++--- .../tests/operation/dag_executor_test.py | 34 ++++----- .../tests/operation/dag_handler_test.py | 5 +- 8 files changed, 129 insertions(+), 47 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_retry.py 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/template_dag.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml index 86427de2..e527c82d 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml @@ -278,3 +278,19 @@ Resources: 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 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 index 9f28a5ea..275625dd 100644 --- 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 @@ -17,8 +17,9 @@ 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 DagConfig, DagContext, DagResult, TriggerRule +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 @@ -53,10 +54,12 @@ def handler(event: Any, context: DurableContext) -> dict[str, Any]: 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") + charge = d.step(lambda deps, sc: "charged", name="charge", config=no_retry) else: - charge = d.step(_charge_declined, name="charge") + 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 @@ -65,10 +68,8 @@ def register(d: DagContext) -> None: TriggerRule.ALL_DONE ) - # Disable retries so an intentionally failing charge terminates promptly. result = context.dag( register, name="compensation", - config=DagConfig(default_retry_strategy=RetryPresets.none()), ) return _summarize(result) 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 index 9ac6950f..bd03ac7e 100644 --- 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 @@ -39,7 +39,6 @@ WaitForCallbackConfig, ) from aws_durable_execution_sdk_python.lambda_service import ErrorObject - from aws_durable_execution_sdk_python.retries import RetryDecision from aws_durable_execution_sdk_python.serdes import SerDes from aws_durable_execution_sdk_python.types import ( DurableContext, @@ -161,7 +160,6 @@ class DagConfig: default of 40. """ completion_config: CompletionConfig | None = None - default_retry_strategy: Callable[[Exception, int], RetryDecision] | None = None default_trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS serdes: SerDes | None = None 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 index d2d7bf0a..72e2201b 100644 --- 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 @@ -38,7 +38,6 @@ if TYPE_CHECKING: from collections.abc import Callable, Sequence - 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 DagConfig from aws_durable_execution_sdk_python.serdes import SerDesContext @@ -159,7 +158,7 @@ def step( trigger_rule=None, run_if=None, ): task_name = _resolve_name(name, func) - cfg = config or self._default_step_config() + cfg = config def executor(ctx: DurableContext, deps_map: DepsMap): return ctx._run_step_with_task_id( @@ -417,10 +416,3 @@ def wrapped(): ), config=config, ) - - def _default_step_config(self) -> StepConfig | None: - from aws_durable_execution_sdk_python.config import StepConfig - - if self._config.default_retry_strategy is not None: - return StepConfig(retry_strategy=self._config.default_retry_strategy) - return None 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 index fb1c158e..d2abf16d 100644 --- 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 @@ -26,7 +26,7 @@ import pytest from aws_durable_execution_sdk_python import exceptions -from aws_durable_execution_sdk_python.config import CompletionConfig +from aws_durable_execution_sdk_python.config import CompletionConfig, StepConfig from aws_durable_execution_sdk_python.dag import ( DagConfig, DagResult, @@ -37,6 +37,9 @@ 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_]+") @@ -221,7 +224,7 @@ def reg(d: Any) -> None: if charge_ok: c = d.step(lambda deps, sc: "charged", name="charge") else: - c = d.step(_fail, name="charge") + 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 @@ -230,7 +233,7 @@ def reg(d: Any) -> None: TriggerRule.ALL_DONE ) - result, client = _run(reg, scenario, DagConfig(default_retry_strategy=NO_RETRY)) + result, client = _run(reg, scenario) return _record(scenario, result, client) @@ -329,13 +332,13 @@ def test_dag_6_trigger_matrix_mixed() -> None: def reg(d: Any) -> None: up_ok = d.step(lambda deps, sc: "ok", name="up_ok") - up_fail = d.step(_fail, name="up_fail") + 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", DagConfig(default_retry_strategy=NO_RETRY)) + 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" @@ -356,12 +359,12 @@ def test_dag_7_trigger_matrix_all_failed() -> None: ] def reg(d: Any) -> None: - u1 = d.step(_fail, name="u1") - u2 = d.step(_fail, name="u2") + 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", DagConfig(default_retry_strategy=NO_RETRY)) + result, client = _run(reg, "DAG-7") rec = _record("DAG-7", result, client) assert _status(rec, "u1") == "FAILED" assert _status(rec, "u2") == "FAILED" @@ -524,7 +527,7 @@ 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}") + 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 @@ -535,7 +538,6 @@ def reg(d: Any) -> None: DagConfig( max_concurrency=1, completion_config=CompletionConfig(tolerated_failure_count=1), - default_retry_strategy=NO_RETRY, ), ) rec = _record("DAG-17", result, client) 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 index e053b256..abcd76e4 100644 --- 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 @@ -7,7 +7,7 @@ import pytest -from aws_durable_execution_sdk_python.config import CompletionConfig +from aws_durable_execution_sdk_python.config import CompletionConfig, StepConfig from aws_durable_execution_sdk_python.dag import ( DagCompletionReason, DagConfig, @@ -26,6 +26,9 @@ 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"): @@ -116,11 +119,11 @@ def boom(_deps, _sc): raise ValueError("boom") def register(d): - a = d.step(boom, name="a") + 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, DagConfig(default_retry_strategy=NO_RETRY)) + 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 @@ -132,7 +135,7 @@ def charge(_deps, _sc): raise RuntimeError("charge failed") def register(d): - c = d.step(charge, name="charge") + 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 @@ -144,7 +147,7 @@ def register(d): TriggerRule.ALL_DONE ) - result, _ = run_dag(register, DagConfig(default_retry_strategy=NO_RETRY)) + 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 @@ -184,12 +187,11 @@ def boom(_deps, _sc): def register(d): for i in range(3): - d.step(boom, name=f"t{i}") + d.step(boom, name=f"t{i}", config=NO_RETRY_CFG) result, _ = run_dag( register, DagConfig( - default_retry_strategy=NO_RETRY, completion_config=CompletionConfig(tolerated_failure_count=0), ), ) @@ -204,10 +206,10 @@ def boom(_deps, _sc): raise ValueError("x") def register(d): - d.step(boom, name="a") + d.step(boom, name="a", config=NO_RETRY_CFG) d.step(lambda deps, sc: ran.__setitem__("b", True), name="b") - result, _ = run_dag(register, DagConfig(default_retry_strategy=NO_RETRY)) + result, _ = run_dag(register) assert ran["b"] is True assert result.failure_count == 1 assert result.success_count == 1 @@ -218,9 +220,9 @@ def boom(_deps, _sc): raise ValueError("bad") def register(d): - d.step(boom, name="a") + d.step(boom, name="a", config=NO_RETRY_CFG) - result, _ = run_dag(register, DagConfig(default_retry_strategy=NO_RETRY)) + result, _ = run_dag(register) with pytest.raises(DagExecutionError): result.throw_if_error() @@ -236,13 +238,12 @@ def boom(_deps, _sc): raise ValueError("x") def register(d): - d.step(boom, name="a") + d.step(boom, name="a", config=NO_RETRY_CFG) d.step(lambda deps, sc: 1, name="b") result, _ = run_dag( register, DagConfig( - default_retry_strategy=NO_RETRY, completion_config=CompletionConfig(tolerated_failure_percentage=10), ), ) @@ -265,7 +266,7 @@ def boom(_deps, _sc): raise ValueError("x") def register(d): - a = d.step(boom, name="a") + 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") @@ -273,7 +274,6 @@ def register(d): result, _ = run_dag( register, DagConfig( - default_retry_strategy=NO_RETRY, default_trigger_rule=TriggerRule.ALL_DONE, ), ) @@ -359,7 +359,7 @@ def register(d): trigger_rule=TriggerRule.ALL_FAILED, ) - ex = _make_executor(register, DagConfig(default_retry_strategy=NO_RETRY)) + ex = _make_executor(register) with _fail_on_hang(), pytest.raises(DagPredicateError) as ei: ex.run() @@ -391,7 +391,7 @@ def register(d): trigger_rule=TriggerRule.ALL_FAILED, ) - ex = _make_executor(register, DagConfig(default_retry_strategy=NO_RETRY)) + ex = _make_executor(register) with _fail_on_hang(), pytest.raises(DagPredicateError) as ei: ex.run() 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 index e7b3d297..4e526465 100644 --- 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 @@ -12,11 +12,8 @@ TaskStatus, ) from aws_durable_execution_sdk_python.exceptions import DagCyclicDependencyError -from aws_durable_execution_sdk_python.retries import RetryPresets from tests.dag_support import make_context, make_state -NO_RETRY = RetryPresets.none() - def _diamond(d): a = d.step(lambda deps, sc: "A", name="a") @@ -274,7 +271,7 @@ def register(d): ctx.dag( register, name="pipeline", - config=DagConfig(default_retry_strategy=NO_RETRY), + config=DagConfig(), ) # Type surfaces cleanly; the offending task is named in the durable message. From 5258192c1cdae5267d3f1a58d88e26013c301e3a Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Mon, 27 Jul 2026 00:15:54 -0700 Subject: [PATCH 19/25] fix(sdk): preserve a nested DAG result when the inner container offloads A nested DAG whose own aggregate exceeded the checkpoint limit could come back from replay as an EMPTY result that claimed success. A nested DAG is both a task of the outer DAG and a child context of the outer container, so the inner container checkpoints its own envelope under the same branch rule, and the outer embeds the inner result recursively. Restoring only rebuilt when the envelope carried a tasks array, which the OFFLOADED shape by definition does not, so it fell through to a fabricated empty result: ALL_COMPLETED with totalCount 0, discarding the counts and completion reason the envelope actually carries. A nested DAG with failed tasks reported success. Reachability was not exotic. The outer embeds the inner tasks in full, so an inner aggregate over the limit guarantees the outer is over it too, meaning a nested DAG with a large aggregate offloads BOTH containers, which is exactly the combination that hit the fallback. Coverage missed it because the nested scenario is small enough never to offload and the large-payload scenario is flat with no nesting. Restoring a tasks-less envelope now preserves its counts and completion reason, so a caller is never told a DAG succeeded when the checkpoint says otherwise, and the reconstruct path recovers the inner per-task detail from the inner container own child checkpoints rather than parsing the inner envelope alone. --- .../handlers/dag/dag_nested_large_payload.py | 124 ++++++ .../template_dag.yaml | 16 + .../operation/dag_result.py | 32 +- .../dag_nested_large_payload_coverage_test.py | 228 +++++++++++ .../operation/dag_order_independence_test.py | 357 ++++++++++++++++++ .../tests/operation/dag_result_test.py | 52 +++ 6 files changed, 806 insertions(+), 3 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_nested_large_payload.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_nested_large_payload_coverage_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/operation/dag_order_independence_test.py 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..7caaf52e --- /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="pauseForReplay") + + # 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/template_dag.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml index e527c82d..b4537f88 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml @@ -294,3 +294,19 @@ Resources: 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 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 index 51a36969..958bfbaa 100644 --- 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 @@ -114,6 +114,9 @@ def __init__( 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 @@ -122,6 +125,16 @@ def __init__( # 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 @@ -162,14 +175,20 @@ def results(self) -> Mapping[str, TaskExecution]: @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 @@ -262,9 +281,13 @@ def from_dict(cls, data: dict[str, Any]) -> DagResultImpl: 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 -- the offloaded path reconstructs the - per-task detail from the child checkpoints instead (see - ``operation/dag.py``); it does not call this method to rebuild tasks. + 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] = {} @@ -299,6 +322,9 @@ def from_dict(cls, data: dict[str, Any]) -> DagResultImpl: 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 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 index 6e1bc997..de79ddc0 100644 --- 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 @@ -192,6 +192,58 @@ def test_envelope_tasks_dropped_is_valid(): 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).""" From a65bbd2ac4ec2fab15d1af0ad4dcbb47c7297f50 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Mon, 27 Jul 2026 12:13:05 -0700 Subject: [PATCH 20/25] fix(sdk): type DAG task deps as Optional[T] for non-SUCCEEDED upstreams DepsMap.__getitem__ and DagResult.get_result typed a handle-keyed dependency lookup as bare T, but the runtime has always returned None for a dependency that did not SUCCEED (FAILED/SKIPPED) under a non-ALL_SUCCESS trigger rule. A compensation task reading such a dependency got no static signal and could crash or misbehave silently. - dag.py / operation/dag_result.py: handle-keyed lookups now type as T | None, matching what the executor has always produced. - Add a component-level test in dag_executor_test.py plus a true integration test in dag_handler_test.py that goes end-to-end through the public ctx.dag() entry point, proving a downstream ALL_DONE task observes a failed upstream's value as None. - Add conformance scenario 10-18 (DagCompensate) handler, verified against a real cloud run (account 730758745077/us-west-2), 18/18 passed with zero regressions. --- .../handlers/dag/dag_compensate.py | 91 +++++++++++++++++++ .../template_dag.yaml | 16 ++++ .../aws_durable_execution_sdk_python/dag.py | 16 +++- .../operation/dag_result.py | 9 +- .../tests/operation/dag_executor_test.py | 44 +++++++++ .../tests/operation/dag_handler_test.py | 49 ++++++++++ 6 files changed, 218 insertions(+), 7 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_compensate.py 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/template_dag.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml index b4537f88..759a06f8 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml @@ -310,3 +310,19 @@ Resources: 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 \ No newline at end of file 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 index bd03ac7e..47b3c061 100644 --- 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 @@ -215,10 +215,18 @@ 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``). The handle path - dispatches at runtime on ``isinstance(key, TaskHandle)`` and extracts the + 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. @@ -228,7 +236,7 @@ def __init__(self, by_name: dict[str, Any]) -> None: self._by_name = by_name @overload - def __getitem__(self, key: TaskHandle[T]) -> T: ... + def __getitem__(self, key: TaskHandle[T]) -> T | None: ... @overload def __getitem__(self, key: str) -> Any: ... def __getitem__(self, key: str | TaskHandle[Any]) -> Any: @@ -445,7 +453,7 @@ class DagResult(ABC): """ @overload - def get_result(self, task: TaskHandle[T]) -> T: ... # pragma: no cover + def get_result(self, task: TaskHandle[T]) -> T | None: ... # pragma: no cover @overload def get_result(self, task: str) -> Any: ... # pragma: no cover @abstractmethod 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 index 958bfbaa..bae39d4a 100644 --- 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 @@ -138,15 +138,18 @@ def __init__( # region accessors @overload - def get_result(self, task: TaskHandle[T]) -> T: ... + 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``); passing a name - string returns ``Any``. Both resolve by task name at runtime. + 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 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 index abcd76e4..cce9ad13 100644 --- 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 @@ -154,6 +154,50 @@ def register(d): 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") 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 index 4e526465..12ae832f 100644 --- 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 @@ -302,3 +302,52 @@ def test_unwrap_dag_error_reconstructs_predicate_error_on_replay(): 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 From 17497063b710476f82245a412901435a24bdf1e7 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Mon, 27 Jul 2026 19:56:55 -0700 Subject: [PATCH 21/25] fix(dag): unify DagContext.wait's signature with context.wait (Duration) Item #5 from the cross-language DAG review: DagContext.wait took seconds: int, while this same SDK's own context.wait(duration: Duration) takes a Duration object -- the same concept, two different shapes, in one SDK. Changed the DAG variant to take Duration, matching the base primitive exactly (same duration.to_seconds() conversion, same 'duration must be at least 1 second' validation via ValidationError, which the DAG variant previously lacked entirely). - dag.py: abstract DagContext.wait signature now takes Duration. - operation/dag_context.py: DagContextImpl.wait converts via duration.to_seconds() and validates the same >= 1 second floor as context.wait, before this DAG's previous behavior on TooShort was to silently pass an invalid seconds value straight to WaitOperationExecutor. - Updated all in-repo call sites: two conformance handlers (dag_wait_resume, dag_concurrent_suspend), one customer example (dag_wait_resume), and three unit tests, all now passing Duration.from_seconds(n) instead of a bare int. - Added test_wait_requires_duration_of_at_least_one_second, matching context.wait's own validation coverage (which the DAG side lacked). This is a breaking signature change on an @Experimental, pre-launch, no-customers surface -- acceptable now, expensive later. Verified: pytest 1609/1609 (151 DAG-specific), mypy clean (2 pre-existing, unrelated errors in dag_conformance_test.py unchanged). --- .../handlers/dag/dag_concurrent_suspend.py | 5 +++-- .../handlers/dag/dag_wait_resume.py | 3 ++- .../src/dag/dag_wait_resume.py | 3 ++- .../src/aws_durable_execution_sdk_python/dag.py | 3 ++- .../operation/dag_context.py | 12 ++++++++++-- .../operation/dag_concurrency_coverage_test.py | 5 +++-- .../tests/operation/dag_context_test.py | 16 ++++++++++++++-- .../tests/operation/dag_task_kinds_test.py | 3 ++- 8 files changed, 38 insertions(+), 12 deletions(-) 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 index 1f29d2ee..6051c74c 100644 --- 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 @@ -26,6 +26,7 @@ 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 @@ -43,8 +44,8 @@ def _counts(result: DagResult) -> list[int]: 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(8, deps=[root], name="slow") # registered FIRST - fast = d.wait(2, deps=[root], name="fast") + 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 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 index 664dd242..ae6e2660 100644 --- 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 @@ -11,6 +11,7 @@ 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 @@ -28,7 +29,7 @@ def _counts(result: DagResult) -> list[int]: 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(5, deps=[start], name="pause") + pause = d.wait(Duration.from_seconds(5), deps=[start], name="pause") d.step(lambda deps, sc: "resumed", deps=[pause], name="finish") result = context.dag( 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 index 06a744b5..5933674f 100644 --- 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 @@ -13,6 +13,7 @@ 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 @@ -43,7 +44,7 @@ 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(3, deps=[start], name="pause") + 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") 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 index 47b3c061..d7d878df 100644 --- 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 @@ -32,6 +32,7 @@ from aws_durable_execution_sdk_python.config import ( ChildConfig, CompletionConfig, + Duration, InvokeConfig, MapConfig, ParallelConfig, @@ -333,7 +334,7 @@ def wait_for_callback( @abstractmethod def wait( self, - seconds: int, + duration: Duration, deps: Sequence[TaskHandle[Any]] | None = None, name: str | None = None, *, 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 index 72e2201b..67057834 100644 --- 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 @@ -19,7 +19,10 @@ TaskHandle, TriggerRule, ) -from aws_durable_execution_sdk_python.exceptions import DagInvalidTaskNameError +from aws_durable_execution_sdk_python.exceptions import ( + DagInvalidTaskNameError, + ValidationError, +) from aws_durable_execution_sdk_python.identifier import ( OperationIdentifier, OperationIdNamespace, @@ -229,7 +232,7 @@ def body(child: DurableContext): ) def wait( - self, seconds, deps=None, name=None, *, + self, duration, deps=None, name=None, *, trigger_rule=None, run_if=None, ): if not name: @@ -237,6 +240,11 @@ def wait( 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, 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 index 608ec17b..f40c5bef 100644 --- 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 @@ -43,6 +43,7 @@ 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, @@ -243,8 +244,8 @@ def _record(name: str) -> str: def register(d: Any) -> None: root = d.step(lambda deps, sc: 1, name="root") - slow = d.wait(8, deps=[root], name="slow") # registered FIRST - fast = d.wait(2, deps=[root], name="fast") + 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") 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 index 4a3f46b6..d6d0cf1a 100644 --- 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 @@ -4,12 +4,16 @@ 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 +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 @@ -82,7 +86,15 @@ def test_unresolvable_name_raises(): def test_wait_requires_name(): d = _impl() with pytest.raises(DagInvalidTaskNameError): - d.wait(5) + 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(): 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 index 0d6ae168..7a25fd7b 100644 --- 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 @@ -4,6 +4,7 @@ import pytest +from aws_durable_execution_sdk_python.config import Duration from aws_durable_execution_sdk_python.waits import ( WaitForConditionConfig, WaitForConditionDecision, @@ -76,7 +77,7 @@ def register(d): def test_wait_task_suspends(): def register(d): - d.wait(30, name="cooldown") + d.wait(Duration.from_seconds(30), name="cooldown") state, _ = make_state() with pytest.raises(SuspendExecution): From 0051d69592cbb3592156742df2f0d8cbfd24d134 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Tue, 28 Jul 2026 14:55:30 -0700 Subject: [PATCH 22/25] feat(dag): implement custom completion predicate (shouldComplete) DagConfig.completion_config now accepts either the existing threshold CompletionConfig or a new DagCustomCompletionConfig -- a results-aware predicate evaluated after every task settlement via a live DagCompletionStatus snapshot (per-task status/result, registration order). complete_dag()/continue_dag() build the decision; two new DagCompletionReason values (CUSTOM_COMPLETION_SUCCEEDED, CUSTOM_COMPLETION_FAILED) record the outcome. The base SDK's map/parallel CompletionConfig has no predicate hook to build on, so this is a DAG-owned mechanism built fresh, matching what DAG_SPEC_CROSS_LANGUAGE.md calls a net-new feature. The 'v2-deferred' note this replaced referred to the base SDK's own feature, not DAG. Also fixes a real bug: DagResult.throw_if_error() only checked failed(), so it could never raise for a CUSTOM_COMPLETION_FAILED DAG with zero individually-failed tasks (a predicate's verdict, not a task exception, failed the DAG). Now also checks completion_reason. Adds the DAG-18 custom-completion conformance scenario (previously TS+Go only) plus the new DAG-19 cloud scenario (dag_rules_engine.py). Tests: 1614/1614 passing, mypy clean (2 pre-existing unrelated errors). --- .../handlers/dag/dag_rules_engine.py | 85 +++++++++++++ .../template_dag.yaml | 15 +++ .../aws_durable_execution_sdk_python/dag.py | 117 +++++++++++++++++- .../operation/dag_executor.py | 59 ++++++++- .../operation/dag_result.py | 12 +- .../tests/conformance/dag_conformance_test.py | 57 +++++++-- .../tests/operation/dag_executor_test.py | 104 ++++++++++++++++ .../tests/operation/dag_handler_test.py | 38 ++++++ 8 files changed, 469 insertions(+), 18 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_rules_engine.py 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/template_dag.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml index 759a06f8..ffe1081c 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml @@ -319,6 +319,21 @@ Resources: 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 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 index d7d878df..87eb57c1 100644 --- 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 @@ -101,6 +101,9 @@ class DagCompletionReason(Enum): 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 @@ -111,6 +114,57 @@ class DagCompletionReason(Enum): 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 @@ -134,12 +188,69 @@ class TaskExecution(Generic[T]): 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. - Reuses the existing threshold-only :class:`~...config.CompletionConfig` - verbatim; result-based custom completion is v2-deferred. + ``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 @@ -160,7 +271,7 @@ class DagConfig: configured on that task, and a nested ``dag`` task gets its own independent default of 40. """ - completion_config: CompletionConfig | None = None + completion_config: CompletionConfig | DagCustomCompletionConfig | None = None default_trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS serdes: SerDes | None = None 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 index ee2d6a8e..ddfbdea3 100644 --- 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 @@ -27,7 +27,11 @@ from typing import TYPE_CHECKING, Any from aws_durable_execution_sdk_python.dag import ( + DagCompletionItemStatus, + DagCompletionOutcome, DagCompletionReason, + DagCompletionStatus, + DagCustomCompletionConfig, DepsMap, SkipReason, TaskExecution, @@ -580,11 +584,14 @@ def _build_deps_map(self, task: TaskDef) -> DepsMap: return DepsMap(by_name) def _threshold_reason_locked(self) -> DagCompletionReason | None: - """Early-completion reason, mirroring ``ExecutionCounters.should_complete``. - - 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 + """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. @@ -592,6 +599,14 @@ def _threshold_reason_locked(self) -> DagCompletionReason | None: 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: @@ -616,6 +631,40 @@ def _threshold_reason_locked(self) -> DagCompletionReason | None: 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. 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 index bae39d4a..d0bf7e9f 100644 --- 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 @@ -203,7 +203,11 @@ def completion_reason(self) -> DagCompletionReason: return self._completion_reason def throw_if_error(self) -> None: - """Raise :class:`DagExecutionError` if any task FAILED.""" + """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] @@ -213,6 +217,12 @@ def throw_if_error(self) -> None: 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 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 index d2abf16d..ba8e6f89 100644 --- 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 @@ -6,9 +6,11 @@ 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): DAG-1..17 and DAG-19 apply to Python (18 -records). DAG-18 (custom result-based completion) is TS+Go only and is OMITTED -here per the catalog's "emit only applicable scenarios" convention. +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 @@ -28,10 +30,14 @@ 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 @@ -553,6 +559,41 @@ def reg(d: Any) -> None: 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: @@ -598,15 +639,14 @@ def reg(d: Any) -> None: def test_zz_emit_python_json() -> None: """Runs last (name-sorted): assemble, validate, and write python.json. - DAG-18 is intentionally OMITTED (TS+Go only, catalog Part C). Emits exactly - 18 records as key-sorted UTF-8 JSON with 2-space indent + trailing newline. + 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, 18)} | {"DAG-19"} + expected_scenarios = {f"DAG-{i}" for i in range(1, 20)} assert set(RECORDS) == expected_scenarios, ( f"missing/extra: {expected_scenarios ^ set(RECORDS)}" ) - assert "DAG-18" not in RECORDS # NOT-APPLICABLE for Python (custom completion) - assert len(RECORDS) == 18 + 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" @@ -616,7 +656,6 @@ def test_zz_emit_python_json() -> None: reloaded = json.loads(OUT_PATH.read_text(encoding="utf-8")) assert reloaded == RECORDS assert text.endswith("\n") - assert "DAG-18" not in reloaded # endregion scenarios 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 index cce9ad13..e1d0d938 100644 --- 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 @@ -9,11 +9,15 @@ 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, @@ -225,6 +229,106 @@ def register(d): 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") 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 index 12ae832f..39f08de7 100644 --- 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 @@ -7,9 +7,13 @@ 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 @@ -31,6 +35,40 @@ def test_context_dag_end_to_end(): 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 From 399a47c327aebede6584973002743d39f0adc71d Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Tue, 28 Jul 2026 17:55:37 -0700 Subject: [PATCH 23/25] test(dag): name the 10-17 settle wait consistently with other SDKs The forced-replay wait in dag_nested_large_payload was 'pauseForReplay' in Python, unnamed in TS, 'wait' in Java, and 'settle' in Go -- a four-way naming inconsistency across otherwise-identical conformance handlers. 10-17.yaml never asserts the specific name (only event types/timing), so this was cosmetic, not a conformance violation. Standardized on Go's 'settle', which best describes why the wait exists. Verified: 19/19 dag scenarios pass on real cloud, 10-17 included. --- .../handlers/dag/dag_nested_large_payload.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 7caaf52e..121699fb 100644 --- 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 @@ -108,7 +108,7 @@ def inner_register(sd: DagContext) -> None: ) # Forces the invocation to end; the next one replays both completed containers. - context.wait(Duration.from_seconds(2), name="pauseForReplay") + context.wait(Duration.from_seconds(2), name="settle") # Recomputed from the REPLAYED (recursively reconstructed) inner result. inner = result.get_result("inner") From 688cba219bc458e9fb1e58d71ac5a99c418ca7e0 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Tue, 28 Jul 2026 19:46:31 -0700 Subject: [PATCH 24/25] fix(dag): close residual off-lock race in the abort/checkpoint guard _resubmit's abort guard checked _scheduler_exception before mutating state, but called create_checkpoint() AFTER releasing self._lock. _scheduler_exception is set by _safe_pump on a worker thread (a completion callback), a different thread than the timer thread that runs _resubmit -- so a worker thread could set the abort flag in the gap between _resubmit releasing the lock and reaching create_checkpoint(), letting a checkpoint fire moments after an abort was decided elsewhere. Not reachable in ~200 iterations of the existing timing-based drain-window test; only surfaces via an exact thread interleaving. Fix: call create_checkpoint() while still holding self._lock. create_checkpoint() blocks on the background checkpoint-batching thread's queue, not on this lock, so this is deadlock-safe. Added a deterministic (no sleep-based timing) regression test that runs _resubmit on its own thread, holds its create_checkpoint call open via a monkeypatch, and proves a second thread's non-blocking self._lock.acquire(blocking=False) probe fails while the checkpoint is in flight. Verified the test fails against the pre-fix code (checkpoint call outside the lock) and passes against the fix. Full suite: 3193 passed, 2 skipped. --- .../operation/dag_executor.py | 17 ++++- .../tests/operation/dag_executor_test.py | 76 +++++++++++++++++++ 2 files changed, 91 insertions(+), 2 deletions(-) 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 index ddfbdea3..63b8c8bd 100644 --- 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 @@ -485,6 +485,18 @@ def _resubmit(self, resumes: list[_TimedResume]) -> None: # 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: @@ -502,8 +514,9 @@ def _resubmit(self, resumes: list[_TimedResume]) -> None: # 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. - self._ctx.state.create_checkpoint() + # 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: 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 index e1d0d938..f98c676b 100644 --- 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 @@ -945,6 +945,82 @@ def _root(name, executor): 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) From b1a37095de76b8a1c5b66f42fe1a85b21e3002b6 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Tue, 28 Jul 2026 22:22:53 -0700 Subject: [PATCH 25/25] test(dag): add DagIdStability handler for 10-20 (task-id stability) Same shape as DagConcurrentOverlap (10-13) -- root -> {a,b} -> {afterA, afterB} -> merge, unbounded concurrency -- except which sibling sleeps longer is driven by event['swap'], forcing genuinely different completion order across two invocations of the SAME deployed function while registering identical task names in identical order every time. Paired with the conformance-tests repo's id_stability.py script, which invokes this function twice (swap=false, then swap=true) and asserts every task's Id field is identical across both runs. Verified against real cloud: both runs SUCCEEDED, all 7 task ids byte-identical across runs. --- .../handlers/dag/dag_id_stability.py | 70 +++++++++++++++++++ .../template_dag.yaml | 16 +++++ 2 files changed, 86 insertions(+) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/handlers/dag/dag_id_stability.py 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/template_dag.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml index ffe1081c..7258a92d 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/template_dag.yaml @@ -247,6 +247,22 @@ Resources: 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: