Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/tests/otel_lifecycle_compatibility_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,7 @@ def handler(event, durable_context):
},
}
lambda_context = SimpleNamespace(
get_remaining_time_in_millis=lambda: 900000,
aws_request_id="installed",
client_context=None,
identity=None,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@

def _lambda_context() -> Mock:
context = Mock()
context.get_remaining_time_in_millis.return_value = 900000
context.aws_request_id = "test-request-id"
context.client_context = None
context.identity = None
Expand Down
37 changes: 37 additions & 0 deletions packages/aws-durable-execution-sdk-python/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,43 @@ core alone with OTel 1.0 isolates worker bindings, but does not add the newer
Invocation-view fallback to that older plugin; upgrade both packages for it.
Provider API version 1 and the independent registration API are unchanged.

### Invocation deadlines and Lambda Managed Instances

The SDK captures `LambdaContext.get_remaining_time_in_millis()` at invocation
entry and uses one monotonic budget, reserving 100 ms before the platform
deadline for cooperative cleanup. An already-expired invocation starts no user
handler. Once the budget is exhausted, new user operations and checkpoint/state
requests are rejected with `InvocationDeadlineExceeded`, a retryable
`InvocationError`. Step retry strategies do not turn this interruption into a
persisted step failure. Completed checkpoints retain their stored results and
errors for replay.

Synchronous checkpoint waits, scheduled refresh waits, checkpoint queue reads,
and map/parallel coordinator waits share that budget. In particular, a losing
branch parked in an SDK refresh wait after early parallel completion unwinds at
the cutoff, allowing invocation cleanup to join it. A newly created descendant
of a completed context is rejected before its body, START hook, or checkpoint
runs, even if it did not exist when its ancestor completed.

This is cooperative interruption, **not a guarantee of bounded invocation
shutdown or LMI worker reclamation**. Python cannot safely terminate arbitrary
running threads. The SDK still joins invocation-owned workers before returning:
user code or I/O that does not cooperate, an already-running synchronous service
client call (including its retries), or a blocked plugin hook can delay those
joins indefinitely. The shared service client is not reconfigured or closed to
interrupt one invocation. Give application I/O explicit finite timeouts and keep
plugin hooks bounded. A runtime-supported worker-retirement strategy is still
required for non-cooperative work; [issue #741](https://github.com/aws/aws-durable-execution-sdk-python/issues/741)
tracks that remaining work, with real-LMI coverage in [test PR #743](https://github.com/aws/aws-durable-execution-sdk-python/pull/743).

An expired invocation does not report an unacknowledged result as `SUCCEEDED` or
`PENDING`. If the service has already acknowledged this invocation's terminal
execution success/failure checkpoint, that outcome remains authoritative while
local delivery hooks finish. At the cutoff, further response metadata pages are
skipped for that accepted terminal checkpoint; inline and already-read operations
are retained. Already-running side effects cannot be undone;
interrupted work continues to require the documented retry/idempotency strategy.

## 🚀 Quick Start

Install the execution SDK:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -292,6 +292,7 @@ def execute(
execution_state.register_branch_pool(pool)

def submit(branch: Branch[CallableType, ResultType]) -> None:
execution_state.invocation_deadline.check()
branch.start()
if execution_state._plugin_executor._plugins: # noqa: SLF001
# Every admission/resume gets its own Context. Branch bindings
Expand Down Expand Up @@ -323,6 +324,7 @@ def submit(branch: Branch[CallableType, ResultType]) -> None:
# and the reason, so user code runs at most once per change.
decision: CompletionDecision | None = None
while True:
execution_state.invocation_deadline.check()
if needs_snapshot_rebuild:
items_snapshot = (
self._build_items_snapshot()
Expand Down Expand Up @@ -386,7 +388,9 @@ def submit(branch: Branch[CallableType, ResultType]) -> None:
if resume_times:
timeout = max(resume_times[0] - time.time(), 0)
try:
event: BranchEvent[ResultType] = events.get(timeout=timeout)
event: BranchEvent[ResultType] = events.get(
timeout=execution_state.invocation_deadline.remaining(timeout)
)
except queue.Empty:
# A timed resume came due while branches were running.
continue
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
"""One monotonic time budget for cooperative invocation interruption."""

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import TYPE_CHECKING

from aws_durable_execution_sdk_python.exceptions import InvocationDeadlineExceeded

if TYPE_CHECKING:
from aws_durable_execution_sdk_python.threading import CompletionEvent
from aws_durable_execution_sdk_python.types import LambdaContext

# Stop admitting work before the platform deadline to leave time for cooperative
# cleanup. This is not a bound on arbitrary user code, clients, or plugin hooks.
_CLEANUP_MARGIN_SECONDS = 0.1


@dataclass(frozen=True)
class InvocationDeadline:
"""Internal invocation-local deadline; None keeps standalone state unbounded."""

expires_at: float | None = None

@classmethod
def from_lambda_context(cls, context: LambdaContext) -> InvocationDeadline:
# Capture once: wall-clock adjustments cannot extend this invocation.
started_at = time.monotonic()
remaining_seconds = context.get_remaining_time_in_millis() / 1000
return cls(started_at + max(0.0, remaining_seconds - _CLEANUP_MARGIN_SECONDS))

def remaining(self, timeout: float | None = None) -> float | None:
"""Clamp an existing wait to the same budget, without resetting it."""
if self.expires_at is None:
return timeout
remaining = max(0.0, self.expires_at - time.monotonic())
return remaining if timeout is None else min(timeout, remaining)

def check(self) -> None:
if self.remaining() == 0:
raise InvocationDeadlineExceeded("Lambda invocation time budget exhausted")

def wait(self, event: CompletionEvent, timeout: float | None = None) -> bool:
"""Do not turn an acknowledged checkpoint into an error after its wait."""
completed = event.wait(self.remaining(timeout))
if not completed:
self.check()
return completed
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,10 @@ def is_retryable(self) -> bool:
return True


class InvocationDeadlineExceeded(InvocationError):
"""The invocation must retry after its cooperative time budget expires."""


class DurableApiErrorCategory(Enum):
INVOCATION = "INVOCATION"
EXECUTION = "EXECUTION"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from typing import TYPE_CHECKING, Any

from aws_durable_execution_sdk_python.context import DurableContext
from aws_durable_execution_sdk_python.deadline import InvocationDeadline
from aws_durable_execution_sdk_python.exceptions import (
BackgroundThreadError,
BotoClientError,
Expand Down Expand Up @@ -193,6 +194,8 @@ def durable_execution(
plugin_executor = PluginExecutor(configured_plugins)

def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
deadline = InvocationDeadline.from_lambda_context(context)
deadline.check()
invocation_input: DurableExecutionInvocationInput
service_client: DurableServiceClient

Expand Down Expand Up @@ -228,6 +231,7 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
service_client=service_client,
plugin_executor=plugin_executor,
updated_operation_ids=invocation_input.updated_operation_ids,
invocation_deadline=deadline,
)

try:
Expand Down Expand Up @@ -286,6 +290,7 @@ def invoke_handler() -> MutableMapping[str, Any]:
)
try:
# Background checkpointing errors will propagate through CompletionEvent.wait() as BackgroundThreadError
deadline.check()
result = func(input_event, durable_context)

# done with userland
Expand Down Expand Up @@ -445,6 +450,7 @@ def wait_for_checkpoint() -> None:
)
execution_operation = execution_state.get_execution_operation()

deadline.check()
# execute the plugins
plugin_executor.on_invocation_start(
execution_arn=invocation_input.durable_execution_arn,
Expand All @@ -470,6 +476,7 @@ def wait_for_checkpoint() -> None:
output = plugin_executor._run_in_invocation_context(
invoke_handler
)
execution_state.check_invocation_result()
plugin_executor.on_invocation_end(
DurableExecutionInvocationOutput.from_dict(output)
)
Expand All @@ -492,7 +499,11 @@ def wait_for_checkpoint() -> None:
else:
# Preserve the original fresh-worker context without plugins.
invocation_future = executor.submit(invoke_worker)
return invocation_future.result()
result = invocation_future.result()
# Cleanup hooks and all owned workers have returned. A hook can consume
# the remaining budget; do not call End again when this check fails.
execution_state.check_invocation_result()
return result

return wrapper

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,7 @@ def process(self) -> T:
InvalidStateError: If the check result is in an invalid state
May raise operation-specific errors from check_result_status() or execute()
"""
self.state.check_operation_active(self.operation_identifier)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Codex AI review · Finding arf_v1_ygsoiugckiy5x5urrhaxumbgjl

[P1] Preserve completed-child replay during in-process resumes

A parallel/map branch can complete a child context with a large ReplayChildren result, then suspend on a wait while another branch remains running. When the coordinator resumes it within the same invocation, that child's terminal checkpoint has already marked its inner operations in _parent_done. This new check rejects those operations before their saved results can be read. _branch_worker reports ORPHANED, causing the coordinator to exit and checkpoint a successful parent with missing branch results. Allow saved checkpoint reads and legitimate ReplayChildren reconstruction through the orphan guard while retaining checks before new work and checkpoints. Add an integration test covering this same-invocation resume.

To suppress repeats, reply /ai verdict deferred, /ai verdict false-positive, or /ai verdict out-of-scope on the first nonblank line. Add rationale on later lines. A new /ai verdict accepted allows future repeats again.

# Check 1: Entry (handles replay and existing checkpoints)
result = self.check_result_status()

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from aws_durable_execution_sdk_python.exceptions import (
ExecutionError,
InvalidStateError,
InvocationDeadlineExceeded,
RetryableSerDesError,
SerDesError,
StepError,
Expand Down Expand Up @@ -266,6 +267,10 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T:
self.operation_identifier.name,
)
return return_value
except InvocationDeadlineExceeded:
# Deadline interruption belongs to the invocation, not the step
# retry policy or persisted operation failure.
raise
except RetryableSerDesError:
# Transient serdes failure: fail the invocation for backend retry,
# bypassing the step retry strategy. For AT_MOST_ONCE steps the
Expand Down
Loading
Loading