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..03a1530ec --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend.py @@ -0,0 +1,119 @@ +"""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 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 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 ( + MicrovmInvokeOptions, + MicrovmSession, + MicrovmSessionConfig, + microvm_session, +) + + +# 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]: + """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_gate_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"), + ) + + # 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 + # 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_gate_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..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,16 +1,51 @@ """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 +HISTORY_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 +72,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 +92,184 @@ def expect_microvm_terminated( ) raise AssertionError(msg) sleep(POLL_INTERVAL_SECONDS) + + +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") == event_type and event.get("Name") == name: + return event + return None + + +def wait_for_history_event( + execution_arn: str, + event_type: str, + name: str, + *, + client: Any = None, + sleep: Callable[[float], None] = time.sleep, + clock: Callable[[], float] = time.monotonic, +) -> dict[str, Any]: + """Return the first history event with ``event_type`` and ``name``. + + 1. GetDurableExecutionHistory returns the events while the execution + still runs. + 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 + answer. The history is paginated. The check reads every page. + """ + if client is None: + client = _default_lambda_client() + + deadline = clock() + HISTORY_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 + event = _find_event(pages, event_type, name) + if event is not None: + return event + if clock() >= deadline: + msg = ( + 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, + *, + 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) + + +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_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..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 @@ -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,16 @@ from botocore.exceptions import ClientError from test.microvm.microvm_helpers import ( + 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, ) @@ -102,3 +109,257 @@ 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="no StepSucceeded event named s.launch within 120" + ): + find(client, clock) + + assert sum(clock.sleeps) == HISTORY_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) + + +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 + + +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 new file mode 100644 index 000000000..1167975e1 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend.py @@ -0,0 +1,169 @@ +"""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, + fail_callback_after_test_failure, + 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 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") + + 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 + ) + + assert result.status is InvocationStatus.SUCCEEDED, result.error + output = deserialize_operation_payload(result.result) + # 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. + 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 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", + 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, []), + ], + ), + ("gate", OperationType.CALLBACK, 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. 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 + ) + assert launch["delivery"] == "http" + assert launch["microvmId"] == microvm_id + assert launch["resumeBeforeDelivery"] is True