feat(extras): suspend a session MicroVM after a job

packages/aws-durable-execution-sdk-python-examples/examples-catalog.json

@@ -834,6 +834,19 @@
         "ExecutionTimeout": 900
       },
       "path": "./src/microvm/microvm_session_suspend.py"
+    },
+    {
+      "name": "MicroVM Session Suspend After",
+      "description": "Runs two jobs in one Lambda MicroVM that its worker suspends after the first job",
+      "handler": "microvm_session_suspend_after.handler",
+      "integration": true,
+      "usesMicrovm": true,
+      "timeout": 120,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 900
+      },
+      "path": "./src/microvm/microvm_session_suspend_after.py"
     }
   ]
 }

packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend_after.py

@@ -0,0 +1,125 @@
+"""Run two jobs in one Lambda MicroVM that the worker suspends after the first.
+
+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.
+
+The session does not suspend when idle. The first job sets suspend_after, so
+the worker suspends the MicroVM after that job. 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 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 suspended for 5 minutes, not for the
+# whole execution timeout.
+GATE = CallbackConfig(timeout=Duration.from_minutes(5))
+
+# The second job waits for the resume. The worker can refuse it with HTTP 503
+# for up to 30 seconds before that, and the request step retries for up to
+# 60 seconds. So its heartbeat timeout is longer than both together.
+RESUMED_JOB_HEARTBEAT_TIMEOUT = Duration.from_seconds(120)
+
+
+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
+
+
+@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:
+        # The worker suspends the MicroVM when this job has ended and its
+        # result is reported.
+        written = vm.invoke(
+            session_context,
+            "write",
+            job_input={"key": key, "value": f"written-by-{key}"},
+            options=MicrovmInvokeOptions(
+                path="/write",
+                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),
+                suspend_after=True,
+            ),
+        )
+
+        # 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 first job.
+        # 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 first job set suspend_after. So 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=MicrovmInvokeOptions(
+                path="/read",
+                timeout=Duration.from_minutes(5),
+                heartbeat_timeout=RESUMED_JOB_HEARTBEAT_TIMEOUT,
+            ),
+        )
+
+        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. It does not suspend when idle.
+    return microvm_session(
+        context,
+        "suspend-after-session",
+        config=MicrovmSessionConfig(
+            image_identifier=image_arn,
+            execution_role_arn=execution_role_arn,
+            timeout=Duration.from_minutes(10),
+        ),
+        handler=write_gate_read,
+    )

packages/aws-durable-execution-sdk-python-examples/test/README.md

@@ -155,6 +155,7 @@ def test_my_example(durable_runner):
 - 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.
+- `test_microvm_session_suspend_after.py` checks that the worker suspends a session MicroVM after a job with `suspend_after=True`, in a session without auto-suspend, and that the next job resumes it. It needs the same permissions as `test_microvm_session_suspend.py`.
 
 ### 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.

packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_catalog.py

@@ -35,6 +35,15 @@ def test_the_suspend_session_function_timeout_covers_the_http_retry_window():
     assert entry["timeout"] >= HTTP_DELIVERY_MIN_TIMEOUT_SECONDS
 
 
+def test_the_suspend_after_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 After")
+
+    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.

packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_env.py

@@ -11,10 +11,16 @@
     microvm_run_job,
     microvm_session,
     microvm_session_suspend,
+    microvm_session_suspend_after,
 )
 
 
-EXAMPLES = [microvm_run_job, microvm_session, microvm_session_suspend]
+EXAMPLES = [
+    microvm_run_job,
+    microvm_session,
+    microvm_session_suspend,
+    microvm_session_suspend_after,
+]
 IMAGE_ARN = "arn:aws:lambda:us-west-2:123456789012:microvm-image:img"
 ROLE_ARN = "arn:aws:iam::123456789012:role/microvm-execution"
 

packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend_after.py

@@ -0,0 +1,172 @@
+"""Cloud test for the MicroVM session example that suspends after a job.
+
+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_after  # 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_after.handler,
+    lambda_function_name="MicroVM Session Suspend After",
+    cloud_only=True,
+)
+def test_suspends_the_microvm_after_the_first_job_and_resumes_it_for_the_next(
+    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-after-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. The session does not suspend when idle. So only the first
+            #    job's suspend_after can suspend the MicroVM.
+            # 4. So a SUSPENDED state read now lies between the two jobs, and
+            #    comes from suspend_after.
+            # 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.
+            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 there. 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 suspend example. The state check and the
+    # resume run inside the second request step. So the tree shows them only
+    # as a request step that succeeded after the suspend.
+    assert tree(result.operations) == [
+        (
+            "suspend-after-session",
+            OperationType.CONTEXT,
+            SUCCEEDED,
+            [
+                ("suspend-after-session.session", OperationType.STEP, SUCCEEDED, []),
+                ("suspend-after-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-after-session.terminate",
+                    OperationType.STEP,
+                    SUCCEEDED,
+                    [],
+                ),
+            ],
+        )
+    ]
+
+    # The session sets neither auto_suspend_on_idle nor an idle policy. So the
+    # launch records no resume mode. suspend_after is an option of one job,
+    # and the launch checkpoint does not record it.
+    session = result.get_context("suspend-after-session")
+    launch = deserialize_operation_payload(
+        session.get_step("suspend-after-session.launch").result
+    )
+    assert launch["delivery"] == "http"
+    assert launch["microvmId"] == microvm_id
+    assert launch["resumeBeforeDelivery"] is False