Skip to content

Commit 9d704ff

Browse files
authored
feat(extras): suspend a session MicroVM after a job (#845)
A session job with suspend_after=True asks the worker to suspend its MicroVM when the job ends. The worker suspends only the MicroVM that the run hook named, once no other work runs, and acknowledges the request with X-Microvm-Suspend-After: accepted. The next job of the session checks the state and resumes the MicroVM. Every session job also has a fallback: a delivery that gets HTTP 502 with MICROVM_CONNECT_FAILED, or no answer, checks the state, resumes a suspended MicroVM, and sends the job again. So a replay that lost the session's in-memory flag still resumes the MicroVM. A resume rejected with a conflict is sent again once the MicroVM is suspended. Adds the example microvm_session_suspend_after.py and its cloud test.
1 parent 7c16963 commit 9d704ff

26 files changed

Lines changed: 3519 additions & 227 deletions

File tree

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -834,6 +834,19 @@
834834
"ExecutionTimeout": 900
835835
},
836836
"path": "./src/microvm/microvm_session_suspend.py"
837+
},
838+
{
839+
"name": "MicroVM Session Suspend After",
840+
"description": "Runs two jobs in one Lambda MicroVM that its worker suspends after the first job",
841+
"handler": "microvm_session_suspend_after.handler",
842+
"integration": true,
843+
"usesMicrovm": true,
844+
"timeout": 120,
845+
"durableConfig": {
846+
"RetentionPeriodInDays": 7,
847+
"ExecutionTimeout": 900
848+
},
849+
"path": "./src/microvm/microvm_session_suspend_after.py"
837850
}
838851
]
839852
}
Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
"""Run two jobs in one Lambda MicroVM that the worker suspends after the first.
2+
3+
The MicroVM runs the image that scripts/ensure_microvm_image.py builds from
4+
microvm-image/. The image's routes /write and /read write and read a file in
5+
the MicroVM. The image enables the resume hook. So the worker accepts the
6+
next job as soon as the MicroVM resumes.
7+
8+
The session does not suspend when idle. The first job sets suspend_after, so
9+
the worker suspends the MicroVM after that job. A callback named gate holds
10+
the second job. The caller completes the callback with
11+
SendDurableExecutionCallbackSuccess to start the second job.
12+
"""
13+
14+
import os
15+
from typing import Any
16+
17+
from aws_durable_execution_sdk_python.config import CallbackConfig, Duration
18+
from aws_durable_execution_sdk_python.context import DurableContext
19+
from aws_durable_execution_sdk_python.execution import durable_execution
20+
from aws_durable_execution_sdk_python_extras.microvm import (
21+
MicrovmInvokeOptions,
22+
MicrovmSession,
23+
MicrovmSessionConfig,
24+
microvm_session,
25+
)
26+
27+
28+
# The callback named gate holds the second job until the caller completes it.
29+
# If the caller never completes it, the callback times out after 5 minutes.
30+
# The handler then raises an error, and the session terminates the MicroVM.
31+
# So a failed caller leaves the MicroVM suspended for 5 minutes, not for the
32+
# whole execution timeout.
33+
GATE = CallbackConfig(timeout=Duration.from_minutes(5))
34+
35+
# The second job waits for the resume. The worker can refuse it with HTTP 503
36+
# for up to 30 seconds before that, and the request step retries for up to
37+
# 60 seconds. So its heartbeat timeout is longer than both together.
38+
RESUMED_JOB_HEARTBEAT_TIMEOUT = Duration.from_seconds(120)
39+
40+
41+
def _require_microvm_env() -> tuple[str, str]:
42+
"""Return the image ARN and the MicroVM execution role ARN.
43+
44+
The SAM template sets both variables on every example with
45+
``"usesMicrovm": true``. scripts/build_examples.py copies every file under
46+
src/ into one directory. So a shared helper module would need a different
47+
import in Lambda than in the local tests. So each MicroVM example defines
48+
this function itself.
49+
"""
50+
image_arn = os.environ.get("MICROVM_IMAGE_ARN")
51+
execution_role_arn = os.environ.get("MICROVM_EXECUTION_ROLE_ARN")
52+
if not image_arn or not execution_role_arn:
53+
msg = (
54+
"MICROVM_IMAGE_ARN and MICROVM_EXECUTION_ROLE_ARN must be set. "
55+
"Build the image with scripts/ensure_microvm_image.py."
56+
)
57+
raise RuntimeError(msg)
58+
return image_arn, execution_role_arn
59+
60+
61+
@durable_execution
62+
def handler(event: dict[str, Any], context: DurableContext) -> Any:
63+
image_arn, execution_role_arn = _require_microvm_env()
64+
key: str = event["key"]
65+
66+
def write_gate_read(vm: MicrovmSession, session_context: DurableContext) -> Any:
67+
# The worker suspends the MicroVM when this job has ended and its
68+
# result is reported.
69+
written = vm.invoke(
70+
session_context,
71+
"write",
72+
job_input={"key": key, "value": f"written-by-{key}"},
73+
options=MicrovmInvokeOptions(
74+
path="/write",
75+
timeout=Duration.from_minutes(5),
76+
# A job that the MicroVM stops answering fails after this
77+
# time, instead of after the timeout.
78+
heartbeat_timeout=Duration.from_seconds(30),
79+
suspend_after=True,
80+
),
81+
)
82+
83+
# 1. The first job has ended before the gate starts.
84+
# 2. The second job starts only after the caller completes the gate.
85+
# 3. So no job runs while the gate is pending.
86+
# 4. So the worker suspends the MicroVM after the first job.
87+
# 5. So a caller that sees the MicroVM SUSPENDED while the gate is
88+
# pending sees a suspend between the two jobs.
89+
# The callback result carries no data. So the handler ignores it.
90+
session_context.create_callback(name="gate", config=GATE).result()
91+
92+
# The first job set suspend_after. So the request step calls
93+
# GetMicrovm, finds the MicroVM SUSPENDED, and calls ResumeMicrovm
94+
# before it sends the job. A suspended MicroVM keeps its files. So
95+
# this job reads the file that the first job wrote.
96+
read = vm.invoke(
97+
session_context,
98+
"read",
99+
job_input={"key": key},
100+
options=MicrovmInvokeOptions(
101+
path="/read",
102+
timeout=Duration.from_minutes(5),
103+
heartbeat_timeout=RESUMED_JOB_HEARTBEAT_TIMEOUT,
104+
),
105+
)
106+
107+
return {
108+
"sessionMicrovmId": vm.microvm_id,
109+
"writtenBy": written["microvmId"],
110+
"readBy": read["microvmId"],
111+
"value": read["value"],
112+
}
113+
114+
# The session launches one MicroVM, and terminates it when the handler
115+
# returns. It does not suspend when idle.
116+
return microvm_session(
117+
context,
118+
"suspend-after-session",
119+
config=MicrovmSessionConfig(
120+
image_identifier=image_arn,
121+
execution_role_arn=execution_role_arn,
122+
timeout=Duration.from_minutes(10),
123+
),
124+
handler=write_gate_read,
125+
)

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,7 @@ def test_my_example(durable_runner):
155155
- 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`.
156156
- 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.
157157
- `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.
158+
- `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`.
158159

159160
### Cloud-Only Tests
160161
`@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‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,15 @@ def test_the_suspend_session_function_timeout_covers_the_http_retry_window():
3535
assert entry["timeout"] >= HTTP_DELIVERY_MIN_TIMEOUT_SECONDS
3636

3737

38+
def test_the_suspend_after_session_function_timeout_covers_the_http_retry_window():
39+
# Every session job goes over HTTP. The wait for a resume counts toward
40+
# the retry window. So the same minimum applies.
41+
entry = catalog_entry("MicroVM Session Suspend After")
42+
43+
assert entry["usesMicrovm"] is True
44+
assert entry["timeout"] >= HTTP_DELIVERY_MIN_TIMEOUT_SECONDS
45+
46+
3847
def test_the_run_job_function_uses_the_global_timeout():
3948
# The input is small, so the job goes in the run hook. The run hook
4049
# delivery adds no Lambda time for the boot.

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,16 @@
1111
microvm_run_job,
1212
microvm_session,
1313
microvm_session_suspend,
14+
microvm_session_suspend_after,
1415
)
1516

1617

17-
EXAMPLES = [microvm_run_job, microvm_session, microvm_session_suspend]
18+
EXAMPLES = [
19+
microvm_run_job,
20+
microvm_session,
21+
microvm_session_suspend,
22+
microvm_session_suspend_after,
23+
]
1824
IMAGE_ARN = "arn:aws:lambda:us-west-2:123456789012:microvm-image:img"
1925
ROLE_ARN = "arn:aws:iam::123456789012:role/microvm-execution"
2026

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
"""Cloud test for the MicroVM session example that suspends after a job.
2+
3+
The local runner cannot launch a MicroVM. So the test is cloud-only.
4+
"""
5+
6+
import uuid
7+
8+
import pytest
9+
from aws_durable_execution_sdk_python.execution import InvocationStatus
10+
from aws_durable_execution_sdk_python.lambda_service import (
11+
OperationStatus,
12+
OperationType,
13+
)
14+
15+
16+
# The examples package does not depend on the extras. So an environment
17+
# without the extras skips this module instead of failing to import the
18+
# example.
19+
pytest.importorskip("aws_durable_execution_sdk_python_extras")
20+
21+
from src.microvm import microvm_session_suspend_after # noqa: E402
22+
from test.conftest import deserialize_operation_payload # noqa: E402
23+
from test.microvm.microvm_helpers import ( # noqa: E402
24+
expect_microvm_terminated,
25+
fail_callback_after_test_failure,
26+
find_launched_microvm_id,
27+
wait_for_microvm_state,
28+
)
29+
30+
31+
def tree(operations):
32+
"""Return (name, type, status, children) for each operation, in order."""
33+
return [
34+
(
35+
operation.name,
36+
operation.operation_type,
37+
operation.status,
38+
tree(getattr(operation, "child_operations", None) or []),
39+
)
40+
for operation in operations
41+
]
42+
43+
44+
SUCCEEDED = OperationStatus.SUCCEEDED
45+
46+
47+
@pytest.mark.example
48+
@pytest.mark.durable_execution(
49+
handler=microvm_session_suspend_after.handler,
50+
lambda_function_name="MicroVM Session Suspend After",
51+
cloud_only=True,
52+
)
53+
def test_suspends_the_microvm_after_the_first_job_and_resumes_it_for_the_next(
54+
durable_runner,
55+
):
56+
key = str(uuid.uuid4())
57+
58+
with durable_runner:
59+
# A synchronous Invoke returns only when the execution ends. The cloud
60+
# runner's boto3 client reads for at most 60 seconds, and it sends the
61+
# Invoke again after a read timeout. Each resend starts a new
62+
# execution, which launches its own MicroVM. So the test starts the
63+
# execution asynchronously and polls for its result.
64+
execution_arn = durable_runner.run_async(input={"key": key}, timeout=300)
65+
66+
# The test reads the MicroVM state while the gate is pending. So it
67+
# reads the MicroVM ID from the launch step in the execution history.
68+
microvm_id = find_launched_microvm_id(
69+
execution_arn, "suspend-after-session.launch"
70+
)
71+
72+
callback_id = None
73+
gate_completed = False
74+
try:
75+
# The cloud runner reads the callback ID from the CallbackStarted
76+
# event named gate. It returns the ID only after the invocation
77+
# that created the callback has completed.
78+
callback_id = durable_runner.wait_for_callback(
79+
execution_arn=execution_arn, name="gate", timeout=120
80+
)
81+
82+
# 1. The gate starts after the first job ends.
83+
# 2. The second job starts only after the test completes the gate.
84+
# 3. The session does not suspend when idle. So only the first
85+
# job's suspend_after can suspend the MicroVM.
86+
# 4. So a SUSPENDED state read now lies between the two jobs, and
87+
# comes from suspend_after.
88+
# The poll lasts at most 150 seconds. The gate times out after 5
89+
# minutes. So the gate stays pending for the whole poll. The
90+
# state can also pass through SUSPENDING. So the test requires
91+
# only that SUSPENDED is seen once.
92+
wait_for_microvm_state(microvm_id, "SUSPENDED")
93+
94+
durable_runner.send_callback_success(callback_id=callback_id)
95+
gate_completed = True
96+
finally:
97+
# A failure before the gate completes leaves the execution waiting
98+
# for the gate, with the MicroVM still there. A callback failure
99+
# makes the session terminate the MicroVM now, instead of after
100+
# the gate's timeout. If the test fails before it has the callback
101+
# ID, the gate times out after 5 minutes. The session then
102+
# terminates the MicroVM.
103+
if callback_id is not None and not gate_completed:
104+
fail_callback_after_test_failure(durable_runner, callback_id)
105+
106+
result = durable_runner.wait_for_result(
107+
execution_arn=execution_arn, timeout=300
108+
)
109+
110+
assert result.status is InvocationStatus.SUCCEEDED, result.error
111+
output = deserialize_operation_payload(result.result)
112+
# 1. GetMicrovm reported the MicroVM SUSPENDED while the gate was pending.
113+
# 2. The second job read the file that the first job wrote.
114+
# 3. A suspended MicroVM keeps its files.
115+
# So both jobs ran in the same MicroVM, before and after the suspend.
116+
assert output["sessionMicrovmId"] == microvm_id
117+
assert output["writtenBy"] == microvm_id
118+
assert output["readBy"] == microvm_id
119+
assert output["value"] == f"written-by-{key}"
120+
121+
expect_microvm_terminated(microvm_id)
122+
123+
# The tree is the tree of the suspend example. The state check and the
124+
# resume run inside the second request step. So the tree shows them only
125+
# as a request step that succeeded after the suspend.
126+
assert tree(result.operations) == [
127+
(
128+
"suspend-after-session",
129+
OperationType.CONTEXT,
130+
SUCCEEDED,
131+
[
132+
("suspend-after-session.session", OperationType.STEP, SUCCEEDED, []),
133+
("suspend-after-session.launch", OperationType.STEP, SUCCEEDED, []),
134+
(
135+
"write",
136+
OperationType.CONTEXT,
137+
SUCCEEDED,
138+
[
139+
("write.callback", OperationType.CALLBACK, SUCCEEDED, []),
140+
("write.request", OperationType.STEP, SUCCEEDED, []),
141+
],
142+
),
143+
("gate", OperationType.CALLBACK, SUCCEEDED, []),
144+
(
145+
"read",
146+
OperationType.CONTEXT,
147+
SUCCEEDED,
148+
[
149+
("read.callback", OperationType.CALLBACK, SUCCEEDED, []),
150+
("read.request", OperationType.STEP, SUCCEEDED, []),
151+
],
152+
),
153+
(
154+
"suspend-after-session.terminate",
155+
OperationType.STEP,
156+
SUCCEEDED,
157+
[],
158+
),
159+
],
160+
)
161+
]
162+
163+
# The session sets neither auto_suspend_on_idle nor an idle policy. So the
164+
# launch records no resume mode. suspend_after is an option of one job,
165+
# and the launch checkpoint does not record it.
166+
session = result.get_context("suspend-after-session")
167+
launch = deserialize_operation_payload(
168+
session.get_step("suspend-after-session.launch").result
169+
)
170+
assert launch["delivery"] == "http"
171+
assert launch["microvmId"] == microvm_id
172+
assert launch["resumeBeforeDelivery"] is False

0 commit comments

Comments
 (0)