From 9506c7f07f2af1986a9a96e40f758a356c2a760b Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 25 Oct 2026 12:27:39 -0700 Subject: [PATCH 1/3] test(examples): suspend and resume a MicroVM The new session example lets its idle MicroVM suspend during a 60-second wait. Its cloud test reads the MicroVM ID from the launch step, sees the state SUSPENDED, and checks that the next job resumed the same MicroVM. --- .../examples-catalog.json | 13 ++ .../src/microvm/microvm_session_suspend.py | 118 +++++++++++++ .../test/README.md | 1 + .../test/microvm/microvm_helpers.py | 167 +++++++++++++++++- .../test/microvm/test_microvm_catalog.py | 9 + .../test/microvm/test_microvm_env.py | 8 +- .../test/microvm/test_microvm_helpers.py | 160 ++++++++++++++++- .../microvm/test_microvm_session_suspend.py | 138 +++++++++++++++ 8 files changed, 602 insertions(+), 12 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.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 454522a11..3f7c66350 100644 --- a/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json +++ b/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json @@ -821,6 +821,19 @@ "ExecutionTimeout": 900 }, "path": "./src/microvm/microvm_session.py" + }, + { + "name": "MicroVM Session Suspend", + "description": "Runs two jobs in one Lambda MicroVM that suspends itself while idle between them", + "handler": "microvm_session_suspend.handler", + "integration": true, + "usesMicrovm": true, + "timeout": 120, + "durableConfig": { + "RetentionPeriodInDays": 7, + "ExecutionTimeout": 900 + }, + "path": "./src/microvm/microvm_session_suspend.py" } ] } diff --git a/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py b/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py new file mode 100644 index 000000000..815921f74 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py @@ -0,0 +1,118 @@ +"""Run two jobs in one Lambda MicroVM that suspends itself between them. + +The MicroVM runs the image that scripts/ensure_microvm_image.py builds from +microvm-image/. The image's routes /write and /read write and read a file in +the MicroVM. The image enables the suspend and resume hooks, which the worker +needs to suspend its own MicroVM. +""" + +import os +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.execution import durable_execution +from aws_durable_execution_sdk_python_extras.microvm import ( + MicrovmInvokeOptions, + MicrovmSession, + MicrovmSessionConfig, + microvm_session, +) + + +# 1. The worker starts its idle time of 10 seconds when the first job ends. +# 2. When the idle time ends, the worker calls SuspendMicrovm. The call +# returns before the MicroVM freezes. The image gives the suspend hook up +# to 10 seconds. +# 3. The service does not document how long the freeze takes. So this +# example assumes that the MicroVM is SUSPENDED within about 25 seconds +# of the end of the first job. +# 4. The wait starts after the first job ends. So the wait ends more than 60 +# seconds after the first job ends. +# 5. So the MicroVM stays SUSPENDED for about 35 seconds or more before the +# second job resumes it. +# 6. The cloud test reads the MicroVM state every 2 seconds. So it observes +# the SUSPENDED state. +# A longer wait would only make the test slower. +PAUSE = Duration.from_seconds(60) +IDLE_TIME = Duration.from_seconds(10) + + +def _require_microvm_env() -> tuple[str, str]: + """Return the image ARN and the MicroVM execution role ARN. + + The SAM template sets both variables on every example with + ``"usesMicrovm": true``. scripts/build_examples.py copies every file under + src/ into one directory. So a shared helper module would need a different + import in Lambda than in the local tests. So each MicroVM example defines + this function itself. + """ + image_arn = os.environ.get("MICROVM_IMAGE_ARN") + execution_role_arn = os.environ.get("MICROVM_EXECUTION_ROLE_ARN") + if not image_arn or not execution_role_arn: + msg = ( + "MICROVM_IMAGE_ARN and MICROVM_EXECUTION_ROLE_ARN must be set. " + "Build the image with scripts/ensure_microvm_image.py." + ) + raise RuntimeError(msg) + return image_arn, execution_role_arn + + +def _job_options(path: str) -> MicrovmInvokeOptions: + return MicrovmInvokeOptions( + path=path, + timeout=Duration.from_minutes(5), + # A job that the MicroVM stops answering fails after this time, + # instead of after the timeout. + heartbeat_timeout=Duration.from_seconds(30), + ) + + +@durable_execution +def handler(event: dict[str, Any], context: DurableContext) -> Any: + image_arn, execution_role_arn = _require_microvm_env() + key: str = event["key"] + + def write_pause_read(vm: MicrovmSession, session_context: DurableContext) -> Any: + written = vm.invoke( + session_context, + "write", + job_input={"key": key, "value": f"written-by-{key}"}, + options=_job_options("/write"), + ) + + # No job runs during the wait. So the worker suspends the MicroVM. + session_context.wait(PAUSE, name="pause") + + # The request step calls GetMicrovm, finds the MicroVM SUSPENDED, and + # calls ResumeMicrovm before it sends the job. A suspended MicroVM + # keeps its files. So this job reads the file that the first job + # wrote. + read = vm.invoke( + session_context, + "read", + job_input={"key": key}, + options=_job_options("/read"), + ) + + return { + "sessionMicrovmId": vm.microvm_id, + "writtenBy": written["microvmId"], + "readBy": read["microvmId"], + "value": read["value"], + } + + # The session launches one MicroVM, and terminates it when the handler + # returns. + return microvm_session( + context, + "suspend-session", + config=MicrovmSessionConfig( + image_identifier=image_arn, + execution_role_arn=execution_role_arn, + timeout=Duration.from_minutes(10), + auto_suspend_on_idle=True, + auto_suspend_idle_time=IDLE_TIME, + ), + handler=write_pause_read, + ) diff --git a/packages/aws-durable-execution-sdk-python-examples/test/README.md b/packages/aws-durable-execution-sdk-python-examples/test/README.md index 323c8eecb..5c7535367 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/README.md +++ b/packages/aws-durable-execution-sdk-python-examples/test/README.md @@ -154,6 +154,7 @@ def test_my_example(durable_runner): - Deploy the MicroVM examples with the SAM template, with `MicrovmImageArn` and `MicrovmExecutionRoleArn` set. `cli.py deploy` sets neither `MICROVM_IMAGE_ARN` nor `MICROVM_EXECUTION_ROLE_ARN`. So the example raises `RuntimeError` when it runs. When `build/` is missing, `cli.py` also builds it without the extras. So the example then fails at import. - Pass a cloud-test input as an object, such as a `dict`. The cloud runner encodes the input with `json.dumps`. So a string that already holds JSON is encoded again, and the handler receives a `str` instead of a `dict`. - Start a MicroVM test with `run_async`, and read its result with `wait_for_result`. The synchronous `run` uses botocore's 60-second read timeout, and botocore sends the Invoke again after a read timeout. Each resend starts a new execution, which launches its own MicroVM. +- `test_microvm_session_suspend.py` checks that an idle session MicroVM suspends and resumes. The function's role needs `lambda:GetMicrovm` and `lambda:ResumeMicrovm`, and the MicroVM execution role needs `lambda:SuspendMicrovm` on the image. ### Cloud-Only Tests `@pytest.mark.durable_execution(..., cloud_only=True)` marks a test that needs a service the local runner cannot provide, such as a MicroVM. Local mode skips the test with the reason `cloud_only: runs only against the deployed function`. Cloud mode runs it. diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/microvm_helpers.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/microvm_helpers.py index e5f6914ff..6e1b20b4b 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/microvm/microvm_helpers.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/microvm_helpers.py @@ -1,5 +1,6 @@ """Checks for the cloud tests of the MicroVM examples.""" +import json import os import time from collections.abc import Callable @@ -11,6 +12,36 @@ ENDED_STATES = frozenset({"TERMINATING", "TERMINATED"}) POLL_INTERVAL_SECONDS = 2.0 TIMEOUT_SECONDS = 30.0 +LAUNCH_TIMEOUT_SECONDS = 120.0 +STATE_TIMEOUT_SECONDS = 150.0 + + +def _region() -> str: + return os.environ.get("AWS_REGION", "us-west-2") + + +def _default_microvms_client() -> Any: + # The client is created on first use, not at import. So the test modules + # import without AWS credentials, and the local run can skip them. + import boto3 + + return boto3.client("lambda-microvms", region_name=_region()) + + +def _default_lambda_client() -> Any: + # The cloud runner calls the Lambda endpoint in LAMBDA_ENDPOINT. The + # execution history is on the same endpoint. So this client uses it too. + import boto3 + + return boto3.client( + "lambda", + region_name=_region(), + endpoint_url=os.environ.get("LAMBDA_ENDPOINT") or None, + ) + + +def _is_not_found(error: ClientError) -> bool: + return error.response.get("Error", {}).get("Code") == "ResourceNotFoundException" def expect_microvm_terminated( @@ -37,21 +68,14 @@ def expect_microvm_terminated( import without AWS credentials, and the local run can skip them. """ if client is None: - import boto3 - - client = boto3.client( - "lambda-microvms", region_name=os.environ.get("AWS_REGION", "us-west-2") - ) + client = _default_microvms_client() deadline = clock() + TIMEOUT_SECONDS while True: try: state = client.get_microvm(microvmIdentifier=microvm_id).get("state") except ClientError as error: - if ( - error.response.get("Error", {}).get("Code") - == "ResourceNotFoundException" - ): + if _is_not_found(error): return raise if state in ENDED_STATES: @@ -64,3 +88,128 @@ def expect_microvm_terminated( ) raise AssertionError(msg) sleep(POLL_INTERVAL_SECONDS) + + +def _launch_result(history_pages: list[dict[str, Any]], step_name: str) -> Any: + for page in history_pages: + for event in page.get("Events", []): + if event.get("EventType") == "StepSucceeded" and ( + event.get("Name") == step_name + ): + payload = ( + event.get("StepSucceededDetails", {}) + .get("Result", {}) + .get("Payload") + ) + # The launch step checkpoints its result with JsonSerDes. + return json.loads(payload) + return None + + +def find_launched_microvm_id( + execution_arn: str, + launch_step_name: str, + *, + client: Any = None, + sleep: Callable[[float], None] = time.sleep, + clock: Callable[[], float] = time.monotonic, +) -> str: + """Return the MicroVM ID from the launch step of a running execution. + + 1. The launch step checkpoints the RunMicrovm result, with the MicroVM ID. + 2. GetDurableExecutionHistory returns the checkpoint while the execution + still runs. + 3. So the check reads the history every 2 seconds, until the launch step + has succeeded, for up to 120 seconds. + + Other tests launch MicroVMs from the same image at the same time. So a + list of the image's MicroVMs could name another test's MicroVM. The + history names only this execution's MicroVM. + + GetDurableExecutionHistory answers ResourceNotFoundException for a short + time after an asynchronous Invoke. The check reads again after that + answer. The history is paginated. The check reads every page. + """ + if client is None: + client = _default_lambda_client() + + deadline = clock() + LAUNCH_TIMEOUT_SECONDS + while True: + pages: list[dict[str, Any]] = [] + request: dict[str, Any] = { + "DurableExecutionArn": execution_arn, + "IncludeExecutionData": True, + } + try: + while True: + page = client.get_durable_execution_history(**request) + pages.append(page) + if not page.get("NextMarker"): + break + request["Marker"] = page["NextMarker"] + except ClientError as error: + if not _is_not_found(error): + raise + launch = _launch_result(pages, launch_step_name) + if launch is not None: + return launch["microvmId"] + if clock() >= deadline: + msg = ( + f"The step {launch_step_name} did not succeed within " + f"{LAUNCH_TIMEOUT_SECONDS:.0f} seconds of the execution start." + ) + raise AssertionError(msg) + sleep(POLL_INTERVAL_SECONDS) + + +def wait_for_microvm_state( + microvm_id: str, + expected_state: str, + *, + client: Any = None, + sleep: Callable[[float], None] = time.sleep, + clock: Callable[[], float] = time.monotonic, +) -> list[str]: + """Wait until GetMicrovm reports ``expected_state``. + + Returns the states that GetMicrovm reported, in order, with each run of + one state reported once. + + The check reads the state every 2 seconds, for up to 150 seconds. A + MicroVM that ends, or that GetMicrovm does not find, cannot reach the + state any more. So the check fails at once in that case. A MicroVM in + another state after 150 seconds fails the check too. + """ + if client is None: + client = _default_microvms_client() + + observed: list[str] = [] + deadline = clock() + STATE_TIMEOUT_SECONDS + while True: + try: + state = client.get_microvm(microvmIdentifier=microvm_id).get("state") + except ClientError as error: + if _is_not_found(error): + msg = ( + f"GetMicrovm did not find MicroVM {microvm_id} before it " + f"was {expected_state}. States seen: {observed}" + ) + raise AssertionError(msg) from error + raise + if not observed or observed[-1] != state: + observed.append(state) + if state == expected_state: + return observed + if state in ENDED_STATES: + msg = ( + f"MicroVM {microvm_id} is {state} before it was " + f"{expected_state}. States seen: {observed}" + ) + raise AssertionError(msg) + if clock() >= deadline: + msg = ( + f"MicroVM {microvm_id} was not {expected_state} within " + f"{STATE_TIMEOUT_SECONDS:.0f} seconds. States seen: {observed}" + ) + raise AssertionError(msg) + sleep(POLL_INTERVAL_SECONDS) diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_catalog.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_catalog.py index eac08869c..c7adcd722 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_catalog.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_catalog.py @@ -26,6 +26,15 @@ def test_the_session_function_timeout_covers_the_http_retry_window(): assert entry["timeout"] >= HTTP_DELIVERY_MIN_TIMEOUT_SECONDS +def test_the_suspend_session_function_timeout_covers_the_http_retry_window(): + # Every session job goes over HTTP. The wait for a resume counts toward + # the retry window. So the same minimum applies. + entry = catalog_entry("MicroVM Session Suspend") + + assert entry["usesMicrovm"] is True + assert entry["timeout"] >= HTTP_DELIVERY_MIN_TIMEOUT_SECONDS + + def test_the_run_job_function_uses_the_global_timeout(): # The input is small, so the job goes in the run hook. The run hook # delivery adds no Lambda time for the boot. diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_env.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_env.py index 0d17deb3b..31b3d5e7f 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_env.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_env.py @@ -7,10 +7,14 @@ # this module. pytest.importorskip("aws_durable_execution_sdk_python_extras") -from src.microvm import microvm_run_job, microvm_session # noqa: E402 +from src.microvm import ( # noqa: E402 + microvm_run_job, + microvm_session, + microvm_session_suspend, +) -EXAMPLES = [microvm_run_job, microvm_session] +EXAMPLES = [microvm_run_job, microvm_session, microvm_session_suspend] IMAGE_ARN = "arn:aws:lambda:us-west-2:123456789012:microvm-image:img" ROLE_ARN = "arn:aws:iam::123456789012:role/microvm-execution" diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_helpers.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_helpers.py index 3d8446307..f5300ff82 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_helpers.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_helpers.py @@ -1,4 +1,4 @@ -"""Tests for expect_microvm_terminated with a fake client and a fake clock.""" +"""Tests for the MicroVM cloud-test helpers, with fake clients and a fake clock.""" from typing import Any @@ -6,9 +6,13 @@ from botocore.exceptions import ClientError from test.microvm.microvm_helpers import ( + LAUNCH_TIMEOUT_SECONDS, POLL_INTERVAL_SECONDS, + STATE_TIMEOUT_SECONDS, TIMEOUT_SECONDS, expect_microvm_terminated, + find_launched_microvm_id, + wait_for_microvm_state, ) @@ -102,3 +106,157 @@ def test_raises_another_get_microvm_error(): with pytest.raises(ClientError, match="AccessDeniedException"): check(client, clock) + + +def wait(client: FakeMicrovmsClient, clock: FakeClock) -> list[str]: + return wait_for_microvm_state( + "mvm-1", "SUSPENDED", client=client, sleep=clock.sleep, clock=clock + ) + + +def test_wait_returns_the_states_seen_until_the_expected_state(): + client = FakeMicrovmsClient( + ["RUNNING", "RUNNING", "SUSPENDING", "SUSPENDED", "RUNNING"] + ) + clock = FakeClock() + + assert wait(client, clock) == ["RUNNING", "SUSPENDING", "SUSPENDED"] + assert len(client.requests) == 4 + assert clock.sleeps == [POLL_INTERVAL_SECONDS] * 3 + + +@pytest.mark.parametrize("state", ["TERMINATING", "TERMINATED"]) +def test_wait_fails_at_once_for_an_ended_microvm(state): + client = FakeMicrovmsClient(["RUNNING", state]) + clock = FakeClock() + + with pytest.raises(AssertionError, match=f"is {state} before it was SUSPENDED"): + wait(client, clock) + + assert len(client.requests) == 2 + + +def test_wait_fails_at_once_for_a_microvm_that_get_microvm_does_not_find(): + client = FakeMicrovmsClient([client_error("ResourceNotFoundException")]) + clock = FakeClock() + + with pytest.raises(AssertionError, match="did not find MicroVM mvm-1"): + wait(client, clock) + + assert clock.sleeps == [] + + +def test_wait_fails_after_the_timeout(): + client = FakeMicrovmsClient(["RUNNING"]) + clock = FakeClock() + + with pytest.raises(AssertionError, match="was not SUSPENDED within 150 seconds"): + wait(client, clock) + + assert sum(clock.sleeps) == STATE_TIMEOUT_SECONDS + + +def test_wait_raises_another_get_microvm_error(): + client = FakeMicrovmsClient([client_error("AccessDeniedException")]) + clock = FakeClock() + + with pytest.raises(ClientError, match="AccessDeniedException"): + wait(client, clock) + + +def launch_event(name: str, microvm_id: str) -> dict[str, Any]: + return { + "EventType": "StepSucceeded", + "Name": name, + "StepSucceededDetails": { + "Result": { + "Payload": ( + f'{{"microvmId": "{microvm_id}", "delivery": "http", ' + '"resumeBeforeDelivery": true}' + ) + } + }, + } + + +class FakeLambdaClient: + """Answers GetDurableExecutionHistory from a list of answers, in order. + + Each answer is a list of pages, or an exception. After the list ends, it + repeats the last answer. Each page is a list of events. Every page but the + last carries a NextMarker. + """ + + def __init__(self, answers: list[Any]) -> None: + self._answers = answers + self._polls = 0 + self.requests: list[dict[str, Any]] = [] + + def get_durable_execution_history(self, **request: Any) -> dict[str, Any]: + self.requests.append(request) + if "Marker" not in request: + self._polls += 1 + answer = self._answers[min(self._polls, len(self._answers)) - 1] + if isinstance(answer, BaseException): + raise answer + index = int(request.get("Marker", "0")) + page: dict[str, Any] = {"Events": answer[index]} + if index + 1 < len(answer): + page["NextMarker"] = str(index + 1) + return page + + +def find(client: FakeLambdaClient, clock: FakeClock) -> str: + return find_launched_microvm_id( + "arn:execution", "s.launch", client=client, sleep=clock.sleep, clock=clock + ) + + +def test_find_reads_the_microvm_id_from_a_later_history_page(): + client = FakeLambdaClient( + [[[{"EventType": "ExecutionStarted"}], [launch_event("s.launch", "mvm-1")]]] + ) + clock = FakeClock() + + assert find(client, clock) == "mvm-1" + assert client.requests == [ + {"DurableExecutionArn": "arn:execution", "IncludeExecutionData": True}, + { + "DurableExecutionArn": "arn:execution", + "IncludeExecutionData": True, + "Marker": "1", + }, + ] + assert clock.sleeps == [] + + +def test_find_reads_again_until_the_launch_step_succeeds(): + client = FakeLambdaClient( + [ + client_error("ResourceNotFoundException"), + [[launch_event("s.session", "ignored")]], + [[launch_event("s.session", "ignored"), launch_event("s.launch", "mvm-1")]], + ] + ) + clock = FakeClock() + + assert find(client, clock) == "mvm-1" + assert clock.sleeps == [POLL_INTERVAL_SECONDS] * 2 + + +def test_find_fails_after_the_timeout(): + client = FakeLambdaClient([[[{"EventType": "ExecutionStarted"}]]]) + clock = FakeClock() + + with pytest.raises(AssertionError, match="s.launch did not succeed within 120"): + find(client, clock) + + assert sum(clock.sleeps) == LAUNCH_TIMEOUT_SECONDS + + +def test_find_raises_another_history_error(): + client = FakeLambdaClient([client_error("AccessDeniedException")]) + clock = FakeClock() + + with pytest.raises(ClientError, match="AccessDeniedException"): + find(client, clock) diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.py new file mode 100644 index 000000000..06c8c359a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.py @@ -0,0 +1,138 @@ +"""Cloud test for the MicroVM session example that suspends its MicroVM. + +The local runner cannot launch a MicroVM. So the test is cloud-only. +""" + +import uuid + +import pytest +from aws_durable_execution_sdk_python.execution import InvocationStatus +from aws_durable_execution_sdk_python.lambda_service import ( + OperationStatus, + OperationType, +) + + +# The examples package does not depend on the extras. So an environment +# without the extras skips this module instead of failing to import the +# example. +pytest.importorskip("aws_durable_execution_sdk_python_extras") + +from src.microvm import microvm_session_suspend # noqa: E402 +from test.conftest import deserialize_operation_payload # noqa: E402 +from test.microvm.microvm_helpers import ( # noqa: E402 + expect_microvm_terminated, + find_launched_microvm_id, + wait_for_microvm_state, +) + + +def tree(operations): + """Return (name, type, status, children) for each operation, in order.""" + return [ + ( + operation.name, + operation.operation_type, + operation.status, + tree(getattr(operation, "child_operations", None) or []), + ) + for operation in operations + ] + + +SUCCEEDED = OperationStatus.SUCCEEDED + + +@pytest.mark.example +@pytest.mark.durable_execution( + handler=microvm_session_suspend.handler, + lambda_function_name="MicroVM Session Suspend", + cloud_only=True, +) +def test_suspends_the_idle_microvm_and_resumes_it_for_the_next_job(durable_runner): + key = str(uuid.uuid4()) + + with durable_runner: + # A synchronous Invoke returns only when the execution ends. The cloud + # runner's boto3 client reads for at most 60 seconds, and it sends the + # Invoke again after a read timeout. Each resend starts a new + # execution, which launches its own MicroVM. So the test starts the + # execution asynchronously and polls for its result. + execution_arn = durable_runner.run_async(input={"key": key}, timeout=300) + + # The test reads the MicroVM state while the execution waits. So it + # reads the MicroVM ID from the launch step in the execution history. + microvm_id = find_launched_microvm_id(execution_arn, "suspend-session.launch") + + # 1. The worker suspends the MicroVM 10 seconds after the first job + # ends. + # 2. The second job resumes it about 60 seconds after the first job + # ends. + # 3. So GetMicrovm reports SUSPENDED for about 35 seconds or more. + # The state can also pass through SUSPENDING. The service controls + # that timing. So the test requires only that SUSPENDED is seen once. + states = wait_for_microvm_state(microvm_id, "SUSPENDED") + assert states[-1] == "SUSPENDED" + + result = durable_runner.wait_for_result( + execution_arn=execution_arn, timeout=300 + ) + + assert result.status is InvocationStatus.SUCCEEDED, result.error + output = deserialize_operation_payload(result.result) + # 1. GetMicrovm reported the MicroVM SUSPENDED before the execution ended. + # 2. The second job read the file that the first job wrote. + # 3. A suspended MicroVM keeps its files. + # So both jobs ran in the same MicroVM, before and after the suspend. + assert output["sessionMicrovmId"] == microvm_id + assert output["writtenBy"] == microvm_id + assert output["readBy"] == microvm_id + assert output["value"] == f"written-by-{key}" + + expect_microvm_terminated(microvm_id) + + # The tree is the same as in the session example without suspending. The + # state check and the resume run inside each request step. So the tree + # shows them only as a request step that succeeded after the suspend. + assert tree(result.operations) == [ + ( + "suspend-session", + OperationType.CONTEXT, + SUCCEEDED, + [ + ("suspend-session.session", OperationType.STEP, SUCCEEDED, []), + ("suspend-session.launch", OperationType.STEP, SUCCEEDED, []), + ( + "write", + OperationType.CONTEXT, + SUCCEEDED, + [ + ("write.callback", OperationType.CALLBACK, SUCCEEDED, []), + ("write.request", OperationType.STEP, SUCCEEDED, []), + ], + ), + ("pause", OperationType.WAIT, SUCCEEDED, []), + ( + "read", + OperationType.CONTEXT, + SUCCEEDED, + [ + ("read.callback", OperationType.CALLBACK, SUCCEEDED, []), + ("read.request", OperationType.STEP, SUCCEEDED, []), + ], + ), + ("suspend-session.terminate", OperationType.STEP, SUCCEEDED, []), + ], + ) + ] + + # With auto_suspend_on_idle, the launch records the resume mode. Each + # request step then calls GetMicrovm before it sends the job, and calls + # ResumeMicrovm for a suspended MicroVM. + session = result.get_context("suspend-session") + launch = deserialize_operation_payload( + session.get_step("suspend-session.launch").result + ) + assert launch["delivery"] == "http" + assert launch["microvmId"] == microvm_id + assert launch["resumeBeforeDelivery"] is True From b8882db59cf3de27075e959beb5a504587c11e73 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 25 Oct 2026 13:09:11 -0700 Subject: [PATCH 2/3] test(examples): check the suspend between jobs This addresses the review of the suspend test. The test now reads the MicroVM state only after the pause wait has started in the history. --- .../src/microvm/microvm_session_suspend.py | 7 +- .../test/microvm/microvm_helpers.py | 85 ++++++++++++------- .../test/microvm/test_microvm_helpers.py | 76 ++++++++++++++++- .../microvm/test_microvm_session_suspend.py | 34 ++++++-- 4 files changed, 160 insertions(+), 42 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py b/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py index 815921f74..d2a9344c4 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py @@ -2,8 +2,8 @@ The MicroVM runs the image that scripts/ensure_microvm_image.py builds from microvm-image/. The image's routes /write and /read write and read a file in -the MicroVM. The image enables the suspend and resume hooks, which the worker -needs to suspend its own MicroVM. +the MicroVM. The image enables the resume hook. So the worker accepts the +next job as soon as the MicroVM resumes. """ import os @@ -20,7 +20,8 @@ ) -# 1. The worker starts its idle time of 10 seconds when the first job ends. +# 1. The worker starts its idle time of 10 seconds at the run hook, and again +# after each job. So the idle time starts again when the first job ends. # 2. When the idle time ends, the worker calls SuspendMicrovm. The call # returns before the MicroVM freezes. The image gives the suspend hook up # to 10 seconds. diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/microvm_helpers.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/microvm_helpers.py index 6e1b20b4b..0fd8113d2 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/microvm/microvm_helpers.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/microvm_helpers.py @@ -12,7 +12,7 @@ ENDED_STATES = frozenset({"TERMINATING", "TERMINATED"}) POLL_INTERVAL_SECONDS = 2.0 TIMEOUT_SECONDS = 30.0 -LAUNCH_TIMEOUT_SECONDS = 120.0 +HISTORY_TIMEOUT_SECONDS = 120.0 STATE_TIMEOUT_SECONDS = 150.0 @@ -90,41 +90,31 @@ def expect_microvm_terminated( sleep(POLL_INTERVAL_SECONDS) -def _launch_result(history_pages: list[dict[str, Any]], step_name: str) -> Any: +def _find_event( + history_pages: list[dict[str, Any]], event_type: str, name: str +) -> dict[str, Any] | None: for page in history_pages: for event in page.get("Events", []): - if event.get("EventType") == "StepSucceeded" and ( - event.get("Name") == step_name - ): - payload = ( - event.get("StepSucceededDetails", {}) - .get("Result", {}) - .get("Payload") - ) - # The launch step checkpoints its result with JsonSerDes. - return json.loads(payload) + if event.get("EventType") == event_type and event.get("Name") == name: + return event return None -def find_launched_microvm_id( +def wait_for_history_event( execution_arn: str, - launch_step_name: str, + event_type: str, + name: str, *, client: Any = None, sleep: Callable[[float], None] = time.sleep, clock: Callable[[], float] = time.monotonic, -) -> str: - """Return the MicroVM ID from the launch step of a running execution. +) -> dict[str, Any]: + """Return the first history event with ``event_type`` and ``name``. - 1. The launch step checkpoints the RunMicrovm result, with the MicroVM ID. - 2. GetDurableExecutionHistory returns the checkpoint while the execution + 1. GetDurableExecutionHistory returns the events while the execution still runs. - 3. So the check reads the history every 2 seconds, until the launch step - has succeeded, for up to 120 seconds. - - Other tests launch MicroVMs from the same image at the same time. So a - list of the image's MicroVMs could name another test's MicroVM. The - history names only this execution's MicroVM. + 2. So the check reads the history every 2 seconds, until it contains the + event, for up to 120 seconds. GetDurableExecutionHistory answers ResourceNotFoundException for a short time after an asynchronous Invoke. The check reads again after that @@ -133,7 +123,7 @@ def find_launched_microvm_id( if client is None: client = _default_lambda_client() - deadline = clock() + LAUNCH_TIMEOUT_SECONDS + deadline = clock() + HISTORY_TIMEOUT_SECONDS while True: pages: list[dict[str, Any]] = [] request: dict[str, Any] = { @@ -150,18 +140,53 @@ def find_launched_microvm_id( except ClientError as error: if not _is_not_found(error): raise - launch = _launch_result(pages, launch_step_name) - if launch is not None: - return launch["microvmId"] + event = _find_event(pages, event_type, name) + if event is not None: + return event if clock() >= deadline: msg = ( - f"The step {launch_step_name} did not succeed within " - f"{LAUNCH_TIMEOUT_SECONDS:.0f} seconds of the execution start." + f"The execution history had no {event_type} event named " + f"{name} within {HISTORY_TIMEOUT_SECONDS:.0f} seconds." ) raise AssertionError(msg) sleep(POLL_INTERVAL_SECONDS) +def find_launched_microvm_id( + execution_arn: str, + launch_step_name: str, + *, + client: Any = None, + sleep: Callable[[float], None] = time.sleep, + clock: Callable[[], float] = time.monotonic, +) -> str: + """Return the MicroVM ID from the launch step of a running execution. + + 1. The launch step checkpoints the RunMicrovm result, with the MicroVM ID. + 2. The history records the checkpoint as a StepSucceeded event with the + step's name. + 3. So the check waits for that event with wait_for_history_event. + + Other tests launch MicroVMs from the same image at the same time. So a + list of the image's MicroVMs could name another test's MicroVM. The + history names only this execution's MicroVM. + + The result payload is execution data. So wait_for_history_event requests + the history with IncludeExecutionData. + """ + event = wait_for_history_event( + execution_arn, + "StepSucceeded", + launch_step_name, + client=client, + sleep=sleep, + clock=clock, + ) + payload = event.get("StepSucceededDetails", {}).get("Result", {}).get("Payload") + # The launch step checkpoints its result with JsonSerDes. + return json.loads(payload)["microvmId"] + + def wait_for_microvm_state( microvm_id: str, expected_state: str, diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_helpers.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_helpers.py index f5300ff82..8b45fd645 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_helpers.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_helpers.py @@ -6,12 +6,13 @@ from botocore.exceptions import ClientError from test.microvm.microvm_helpers import ( - LAUNCH_TIMEOUT_SECONDS, + HISTORY_TIMEOUT_SECONDS, POLL_INTERVAL_SECONDS, STATE_TIMEOUT_SECONDS, TIMEOUT_SECONDS, expect_microvm_terminated, find_launched_microvm_id, + wait_for_history_event, wait_for_microvm_state, ) @@ -248,10 +249,12 @@ def test_find_fails_after_the_timeout(): client = FakeLambdaClient([[[{"EventType": "ExecutionStarted"}]]]) clock = FakeClock() - with pytest.raises(AssertionError, match="s.launch did not succeed within 120"): + with pytest.raises( + AssertionError, match="no StepSucceeded event named s.launch within 120" + ): find(client, clock) - assert sum(clock.sleeps) == LAUNCH_TIMEOUT_SECONDS + assert sum(clock.sleeps) == HISTORY_TIMEOUT_SECONDS def test_find_raises_another_history_error(): @@ -260,3 +263,70 @@ def test_find_raises_another_history_error(): with pytest.raises(ClientError, match="AccessDeniedException"): find(client, clock) + + +WAIT_STARTED = {"EventType": "WaitStarted", "Name": "pause", "Id": "w-1"} + + +def wait_for_pause(client: FakeLambdaClient, clock: FakeClock) -> dict[str, Any]: + return wait_for_history_event( + "arn:execution", + "WaitStarted", + "pause", + client=client, + sleep=clock.sleep, + clock=clock, + ) + + +def test_wait_for_event_returns_the_event_from_a_later_history_page(): + client = FakeLambdaClient( + [ + [ + [{"EventType": "ExecutionStarted"}], + # Same type with another name, and same name with another type. + [ + {"EventType": "WaitStarted", "Name": "other"}, + {"EventType": "StepStarted", "Name": "pause"}, + ], + [WAIT_STARTED], + ] + ] + ) + clock = FakeClock() + + assert wait_for_pause(client, clock) == WAIT_STARTED + assert [request.get("Marker") for request in client.requests] == [ + None, + "1", + "2", + ] + assert clock.sleeps == [] + + +def test_wait_for_event_reads_again_after_resource_not_found(): + client = FakeLambdaClient( + [ + client_error("ResourceNotFoundException"), + client_error("ResourceNotFoundException"), + [[{"EventType": "ExecutionStarted"}, WAIT_STARTED]], + ] + ) + clock = FakeClock() + + assert wait_for_pause(client, clock) == WAIT_STARTED + assert clock.sleeps == [POLL_INTERVAL_SECONDS] * 2 + + +def test_wait_for_event_fails_after_the_timeout(): + client = FakeLambdaClient([[[{"EventType": "ExecutionStarted"}]]]) + clock = FakeClock() + + with pytest.raises( + AssertionError, match="no WaitStarted event named pause within 120 seconds" + ): + wait_for_pause(client, clock) + + # The check reads the history at 0, 2, ..., 120 seconds. + assert len(client.requests) == HISTORY_TIMEOUT_SECONDS / POLL_INTERVAL_SECONDS + 1 + assert sum(clock.sleeps) == HISTORY_TIMEOUT_SECONDS diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.py index 06c8c359a..6911572da 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.py @@ -23,6 +23,7 @@ from test.microvm.microvm_helpers import ( # noqa: E402 expect_microvm_terminated, find_launched_microvm_id, + wait_for_history_event, wait_for_microvm_state, ) @@ -64,15 +65,35 @@ def test_suspends_the_idle_microvm_and_resumes_it_for_the_next_job(durable_runne # reads the MicroVM ID from the launch step in the execution history. microvm_id = find_launched_microvm_id(execution_arn, "suspend-session.launch") + # 1. The worker starts its idle time at the run hook, and again after + # each job. + # 2. If the first job arrives more than 10 seconds after the run hook, + # the worker suspends the MicroVM before the first job. + # 3. The wait named pause starts only after the first job ends. + # 4. So a SUSPENDED state read after the wait has started lies between + # the two jobs. So the test reads the state only after the history + # shows that the wait has started. + wait_for_history_event(execution_arn, "WaitStarted", "pause") + # 1. The worker suspends the MicroVM 10 seconds after the first job # ends. - # 2. The second job resumes it about 60 seconds after the first job + # 2. The freeze takes up to about 15 seconds more. + # 3. The second job resumes it about 60 seconds after the first job # ends. - # 3. So GetMicrovm reports SUSPENDED for about 35 seconds or more. + # 4. So GetMicrovm reports SUSPENDED for about 35 seconds or more. # The state can also pass through SUSPENDING. The service controls # that timing. So the test requires only that SUSPENDED is seen once. - states = wait_for_microvm_state(microvm_id, "SUSPENDED") - assert states[-1] == "SUSPENDED" + # + # The helper returns only after it reads SUSPENDED. The test does not + # also require a RUNNING state before SUSPENDED: + # 1. The idle time starts when the first job ends. + # 2. The wait starts later. The service first completes the callback, + # and then invokes the function again. That time has no upper + # bound. + # 3. So the wait can start more than 10 seconds after the idle time + # starts. The first read can then already be SUSPENDING or + # SUSPENDED. + wait_for_microvm_state(microvm_id, "SUSPENDED") result = durable_runner.wait_for_result( execution_arn=execution_arn, timeout=300 @@ -80,7 +101,7 @@ def test_suspends_the_idle_microvm_and_resumes_it_for_the_next_job(durable_runne assert result.status is InvocationStatus.SUCCEEDED, result.error output = deserialize_operation_payload(result.result) - # 1. GetMicrovm reported the MicroVM SUSPENDED before the execution ended. + # 1. GetMicrovm reported the MicroVM SUSPENDED after the wait started. # 2. The second job read the file that the first job wrote. # 3. A suspended MicroVM keeps its files. # So both jobs ran in the same MicroVM, before and after the suspend. @@ -128,7 +149,8 @@ def test_suspends_the_idle_microvm_and_resumes_it_for_the_next_job(durable_runne # With auto_suspend_on_idle, the launch records the resume mode. Each # request step then calls GetMicrovm before it sends the job, and calls - # ResumeMicrovm for a suspended MicroVM. + # ResumeMicrovm for a suspended MicroVM. No idle policy is set. So only + # the request step can have called ResumeMicrovm. session = result.get_context("suspend-session") launch = deserialize_operation_payload( session.get_step("suspend-session.launch").result From be283c54c77249901e815f551782d88afddd991a Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 25 Oct 2026 14:28:26 -0700 Subject: [PATCH 3/3] test(examples): gate the read job on a callback Address review: a callback named gate replaces the 60-second wait. The test completes the gate only after it sees SUSPENDED, so the suspend lies between the jobs with no timing assumption. A failed test fails the gate, so the session terminates the MicroVM. --- .../src/microvm/microvm_session_suspend.py | 42 +++++----- .../test/microvm/microvm_helpers.py | 35 ++++++++ .../test/microvm/test_microvm_helpers.py | 33 ++++++++ .../microvm/test_microvm_session_suspend.py | 81 ++++++++++--------- 4 files changed, 134 insertions(+), 57 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py b/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py index d2a9344c4..03a1530ec 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py @@ -4,12 +4,15 @@ microvm-image/. The image's routes /write and /read write and read a file in the MicroVM. The image enables the resume hook. So the worker accepts the next job as soon as the MicroVM resumes. + +A callback named gate holds the second job. The caller completes the callback +with SendDurableExecutionCallbackSuccess to start the second job. """ import os from typing import Any -from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.config import CallbackConfig, Duration from aws_durable_execution_sdk_python.context import DurableContext from aws_durable_execution_sdk_python.execution import durable_execution from aws_durable_execution_sdk_python_extras.microvm import ( @@ -20,23 +23,14 @@ ) -# 1. The worker starts its idle time of 10 seconds at the run hook, and again -# after each job. So the idle time starts again when the first job ends. -# 2. When the idle time ends, the worker calls SuspendMicrovm. The call -# returns before the MicroVM freezes. The image gives the suspend hook up -# to 10 seconds. -# 3. The service does not document how long the freeze takes. So this -# example assumes that the MicroVM is SUSPENDED within about 25 seconds -# of the end of the first job. -# 4. The wait starts after the first job ends. So the wait ends more than 60 -# seconds after the first job ends. -# 5. So the MicroVM stays SUSPENDED for about 35 seconds or more before the -# second job resumes it. -# 6. The cloud test reads the MicroVM state every 2 seconds. So it observes -# the SUSPENDED state. -# A longer wait would only make the test slower. -PAUSE = Duration.from_seconds(60) +# The worker suspends the MicroVM after it has run no job for this time. IDLE_TIME = Duration.from_seconds(10) +# The callback named gate holds the second job until the caller completes it. +# If the caller never completes it, the callback times out after 5 minutes. +# The handler then raises an error, and the session terminates the MicroVM. +# So a failed caller leaves the MicroVM running for 5 minutes, not for the +# whole execution timeout. +GATE = CallbackConfig(timeout=Duration.from_minutes(5)) def _require_microvm_env() -> tuple[str, str]: @@ -74,7 +68,7 @@ def handler(event: dict[str, Any], context: DurableContext) -> Any: image_arn, execution_role_arn = _require_microvm_env() key: str = event["key"] - def write_pause_read(vm: MicrovmSession, session_context: DurableContext) -> Any: + def write_gate_read(vm: MicrovmSession, session_context: DurableContext) -> Any: written = vm.invoke( session_context, "write", @@ -82,8 +76,14 @@ def write_pause_read(vm: MicrovmSession, session_context: DurableContext) -> Any options=_job_options("/write"), ) - # No job runs during the wait. So the worker suspends the MicroVM. - session_context.wait(PAUSE, name="pause") + # 1. The first job has ended before the gate starts. + # 2. The second job starts only after the caller completes the gate. + # 3. So no job runs while the gate is pending. + # 4. So the worker suspends the MicroVM after the idle time. + # 5. So a caller that sees the MicroVM SUSPENDED while the gate is + # pending sees a suspend between the two jobs. + # The callback result carries no data. So the handler ignores it. + session_context.create_callback(name="gate", config=GATE).result() # The request step calls GetMicrovm, finds the MicroVM SUSPENDED, and # calls ResumeMicrovm before it sends the job. A suspended MicroVM @@ -115,5 +115,5 @@ def write_pause_read(vm: MicrovmSession, session_context: DurableContext) -> Any auto_suspend_on_idle=True, auto_suspend_idle_time=IDLE_TIME, ), - handler=write_pause_read, + handler=write_gate_read, ) diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/microvm_helpers.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/microvm_helpers.py index 0fd8113d2..a858a39be 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/microvm/microvm_helpers.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/microvm_helpers.py @@ -1,14 +1,18 @@ """Checks for the cloud tests of the MicroVM examples.""" import json +import logging import os import time from collections.abc import Callable from typing import Any +from aws_durable_execution_sdk_python.lambda_service import ErrorObject from botocore.exceptions import ClientError +logger = logging.getLogger(__name__) + ENDED_STATES = frozenset({"TERMINATING", "TERMINATED"}) POLL_INTERVAL_SECONDS = 2.0 TIMEOUT_SECONDS = 30.0 @@ -238,3 +242,34 @@ def wait_for_microvm_state( ) raise AssertionError(msg) sleep(POLL_INTERVAL_SECONDS) + + +TEST_FAILED_ERROR = ErrorObject( + message="The test failed before it completed the callback.", + type="TestFailed", + data=None, + stack_trace=None, +) + + +def fail_callback_after_test_failure(runner: Any, callback_id: str) -> None: + """Fail a callback that a failed test left pending. + + 1. A pending callback holds the execution until the callback times out. + 2. A callback failure ends the wait at once. The handler then raises + CallbackError, and the session terminates its MicroVM. + 3. So the test calls this function when it fails before it completes the + callback. + + The call runs while the test's own error propagates. An error from the + call would replace that error. So this function logs the error and + returns. + """ + try: + runner.send_callback_failure(callback_id=callback_id, error=TEST_FAILED_ERROR) + except Exception: + logger.warning( + "Could not fail callback %s after the test failed.", + callback_id, + exc_info=True, + ) diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_helpers.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_helpers.py index 8b45fd645..f4f256d74 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_helpers.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_helpers.py @@ -9,8 +9,10 @@ HISTORY_TIMEOUT_SECONDS, POLL_INTERVAL_SECONDS, STATE_TIMEOUT_SECONDS, + TEST_FAILED_ERROR, TIMEOUT_SECONDS, expect_microvm_terminated, + fail_callback_after_test_failure, find_launched_microvm_id, wait_for_history_event, wait_for_microvm_state, @@ -330,3 +332,34 @@ def test_wait_for_event_fails_after_the_timeout(): # The check reads the history at 0, 2, ..., 120 seconds. assert len(client.requests) == HISTORY_TIMEOUT_SECONDS / POLL_INTERVAL_SECONDS + 1 assert sum(clock.sleeps) == HISTORY_TIMEOUT_SECONDS + + +class FakeRunner: + """Records send_callback_failure calls, and raises ``error`` if set.""" + + def __init__(self, error: BaseException | None = None) -> None: + self._error = error + self.failures: list[dict[str, Any]] = [] + + def send_callback_failure(self, **request: Any) -> None: + self.failures.append(request) + if self._error is not None: + raise self._error + + +def test_fail_callback_sends_a_callback_failure(): + runner = FakeRunner() + + fail_callback_after_test_failure(runner, "cb-1") + + assert runner.failures == [{"callback_id": "cb-1", "error": TEST_FAILED_ERROR}] + + +def test_fail_callback_logs_an_error_instead_of_raising_it(caplog): + runner = FakeRunner(RuntimeError("callback already completed")) + + fail_callback_after_test_failure(runner, "cb-1") + + assert len(runner.failures) == 1 + assert "Could not fail callback cb-1" in caplog.text + assert "callback already completed" in caplog.text diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.py index 6911572da..1167975e1 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.py @@ -22,8 +22,8 @@ from test.conftest import deserialize_operation_payload # noqa: E402 from test.microvm.microvm_helpers import ( # noqa: E402 expect_microvm_terminated, + fail_callback_after_test_failure, find_launched_microvm_id, - wait_for_history_event, wait_for_microvm_state, ) @@ -61,39 +61,47 @@ def test_suspends_the_idle_microvm_and_resumes_it_for_the_next_job(durable_runne # execution asynchronously and polls for its result. execution_arn = durable_runner.run_async(input={"key": key}, timeout=300) - # The test reads the MicroVM state while the execution waits. So it + # The test reads the MicroVM state while the gate is pending. So it # reads the MicroVM ID from the launch step in the execution history. microvm_id = find_launched_microvm_id(execution_arn, "suspend-session.launch") - # 1. The worker starts its idle time at the run hook, and again after - # each job. - # 2. If the first job arrives more than 10 seconds after the run hook, - # the worker suspends the MicroVM before the first job. - # 3. The wait named pause starts only after the first job ends. - # 4. So a SUSPENDED state read after the wait has started lies between - # the two jobs. So the test reads the state only after the history - # shows that the wait has started. - wait_for_history_event(execution_arn, "WaitStarted", "pause") - - # 1. The worker suspends the MicroVM 10 seconds after the first job - # ends. - # 2. The freeze takes up to about 15 seconds more. - # 3. The second job resumes it about 60 seconds after the first job - # ends. - # 4. So GetMicrovm reports SUSPENDED for about 35 seconds or more. - # The state can also pass through SUSPENDING. The service controls - # that timing. So the test requires only that SUSPENDED is seen once. - # - # The helper returns only after it reads SUSPENDED. The test does not - # also require a RUNNING state before SUSPENDED: - # 1. The idle time starts when the first job ends. - # 2. The wait starts later. The service first completes the callback, - # and then invokes the function again. That time has no upper - # bound. - # 3. So the wait can start more than 10 seconds after the idle time - # starts. The first read can then already be SUSPENDING or - # SUSPENDED. - wait_for_microvm_state(microvm_id, "SUSPENDED") + callback_id = None + gate_completed = False + try: + # The cloud runner reads the callback ID from the CallbackStarted + # event named gate. It returns the ID only after the invocation + # that created the callback has completed. + callback_id = durable_runner.wait_for_callback( + execution_arn=execution_arn, name="gate", timeout=120 + ) + + # 1. The gate starts after the first job ends. + # 2. The second job starts only after the test completes the gate. + # 3. So a SUSPENDED state read now lies between the two jobs. + # 4. So the result does not depend on how long the suspend takes. + # The poll lasts at most 150 seconds. The gate times out after + # 5 minutes. So the gate stays pending for the whole poll. + # The state can also pass through SUSPENDING. So the test requires + # only that SUSPENDED is seen once. + # + # The test does not also require a RUNNING state before SUSPENDED: + # 1. The idle time starts when the first job ends. + # 2. The runner returns the callback ID after the invocation ends. + # That time has no upper bound. + # 3. So the first read can already be SUSPENDING or SUSPENDED. + wait_for_microvm_state(microvm_id, "SUSPENDED") + + durable_runner.send_callback_success(callback_id=callback_id) + gate_completed = True + finally: + # A failure before the gate completes leaves the execution waiting + # for the gate, with the MicroVM still running. A callback failure + # makes the session terminate the MicroVM now, instead of after + # the gate's timeout. If the test fails before it has the callback + # ID, the gate times out after 5 minutes. The session then + # terminates the MicroVM. + if callback_id is not None and not gate_completed: + fail_callback_after_test_failure(durable_runner, callback_id) result = durable_runner.wait_for_result( execution_arn=execution_arn, timeout=300 @@ -101,7 +109,7 @@ def test_suspends_the_idle_microvm_and_resumes_it_for_the_next_job(durable_runne assert result.status is InvocationStatus.SUCCEEDED, result.error output = deserialize_operation_payload(result.result) - # 1. GetMicrovm reported the MicroVM SUSPENDED after the wait started. + # 1. GetMicrovm reported the MicroVM SUSPENDED while the gate was pending. # 2. The second job read the file that the first job wrote. # 3. A suspended MicroVM keeps its files. # So both jobs ran in the same MicroVM, before and after the suspend. @@ -112,9 +120,10 @@ def test_suspends_the_idle_microvm_and_resumes_it_for_the_next_job(durable_runne expect_microvm_terminated(microvm_id) - # The tree is the same as in the session example without suspending. The - # state check and the resume run inside each request step. So the tree - # shows them only as a request step that succeeded after the suspend. + # The tree is the tree of the session example without suspending, with + # the callback named gate in place of that example's wait. The state + # check and the resume run inside each request step. So the tree shows + # them only as a request step that succeeded after the suspend. assert tree(result.operations) == [ ( "suspend-session", @@ -132,7 +141,7 @@ def test_suspends_the_idle_microvm_and_resumes_it_for_the_next_job(durable_runne ("write.request", OperationType.STEP, SUCCEEDED, []), ], ), - ("pause", OperationType.WAIT, SUCCEEDED, []), + ("gate", OperationType.CALLBACK, SUCCEEDED, []), ( "read", OperationType.CONTEXT,