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 3f7c66350..e6c3ce461 100644 --- a/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json +++ b/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" } ] } diff --git a/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend_after.py b/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend_after.py new file mode 100644 index 000000000..acf42d44d --- /dev/null +++ b/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, + ) 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 5c7535367..ec4ab5079 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/README.md +++ b/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. 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 c7adcd722..58f4d678d 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 @@ -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. 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 31b3d5e7f..e8fae35e9 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 @@ -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" diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend_after.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend_after.py new file mode 100644 index 000000000..0bb51a23b --- /dev/null +++ b/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 diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index f66c5ac64..6b1e287cf 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -100,7 +100,7 @@ Every MicroVM failure is a `MicrovmError`. Catch the base class for all of them, |---|---| | `MicrovmLaunchError` | `RunMicrovm` failed after all retries, or the launch request was not valid, for example with an empty `image_identifier` or `execution_role_arn`. A MicroVM can still run the job. The list below the table says when. | | `MicrovmDeliveryError` | HTTP delivery failed after all retries, or the route answered a 4xx status such as 404, or a redirect. | -| `MicrovmNotRunningError` | A `MicrovmDeliveryError`: the MicroVM is terminating, terminated, or removed. It is not retried. A session that checks the MicroVM state before each job reports it. Any job over HTTP also reports it when `CreateMicrovmAuthToken` reports the MicroVM removed. See [Suspending an idle MicroVM](#suspending-an-idle-microvm). | +| `MicrovmNotRunningError` | A `MicrovmDeliveryError`: the MicroVM is terminating, terminated, or removed. It is not retried. A session that checks the MicroVM state before each job reports it, and so does a session job that checks the state after a sign of a suspended MicroVM. Any job over HTTP also reports it when `CreateMicrovmAuthToken` reports the MicroVM removed. See [Suspending an idle MicroVM](#suspending-an-idle-microvm) and [Suspending after a job](#suspending-after-a-job). | | `MicrovmJobFailedError` | The job handler in the MicroVM failed, or the MicroVM sent a result that is not valid JSON. The message includes the error type that the MicroVM reported. | | `MicrovmTimeoutError` | No result arrived within `timeout`, or no heartbeat within `heartbeat_timeout`. A MicroVM that crashed ends this way. | @@ -260,7 +260,7 @@ A MicroVM lives at most its `timeout` plus 5 minutes, and never longer than 8 ho - For a session that launched with `auto_suspend_on_idle=True`, or with an `idle_policy`, the request step calls `GetMicrovm` before each job. A MicroVM that is terminating, terminated, or removed then fails the job at once with `MicrovmNotRunningError`, a subclass of `MicrovmDeliveryError`. A retry finds the same state, so the job is not retried. The message ends with the request ID of the call that reported the state, for example `GetMicrovm request ID: .`. - Every session, and every `microvm` job over HTTP, calls `CreateMicrovmAuthToken` before the first request, and again after a 401 or 403. The API model lists `ResourceNotFoundException` for that call. That answer fails the job at once with `MicrovmNotRunningError`, whose message ends with `CreateMicrovmAuthToken request ID: .`. The live service's answer for a removed MicroVM is not verified. -- A session that launched with neither setting makes no `GetMicrovm` call. When `CreateMicrovmAuthToken` returns a token, the request goes to the endpoint. A terminated MicroVM's endpoint answers 502, and the request step retries a 502 as a MicroVM that is still starting. So the job fails with a plain `MicrovmDeliveryError`, after the retries of both tiers. +- A session that launched with neither setting makes no `GetMicrovm` call before delivery, unless an earlier job of the handle set `suspend_after`. When `CreateMicrovmAuthToken` returns a token, the request goes to the endpoint. A terminated MicroVM's endpoint answers 502. Whether that 502 carries `x-aws-proxy-error: MICROVM_CONNECT_FAILED` is not verified. With that header, the check of [Suspending after a job](#suspending-after-a-job) reports the state, and the job fails with `MicrovmNotRunningError`. Without it, the request step retries a 502 as a MicroVM that is still starting. So the job fails with a plain `MicrovmDeliveryError`, after the retries of both tiers. Suspending is off by default, because it needs more permissions and changes what the MicroVM does between jobs: @@ -270,9 +270,73 @@ Suspending is off by default, because it needs more permissions and changes what The worker must support `auto_suspend_on_idle`. The JavaScript worker [`@aws/durable-execution-sdk-js-microvm-worker`](https://github.com/aws/aws-durable-execution-sdk-js/tree/main/packages/aws-durable-execution-sdk-js-microvm-worker) and the Python worker [`aws-durable-execution-sdk-python-microvm-worker`](../aws-durable-execution-sdk-python-microvm-worker) both do. `idle_policy` works with both workers. Enable the `suspend` and `resume` hooks in the image. The worker then refuses jobs while the MicroVM freezes, and accepts them again as soon as it resumes. Without the `resume` hook, the first job after a resume can wait a little over 30 seconds, until the worker's refusal ends and the request step delivers again. +### Suspending after a job + +`MicrovmInvokeOptions(suspend_after=True)` asks the worker to suspend the MicroVM after this job. It is an option of one job, not of the session. `microvm` raises `ValidationError` for it, because `microvm` terminates its MicroVM after its one job. A value that is not a `bool` raises `ValidationError` too, before the job's durable operations start. + +```python +build = vm.invoke( + ctx, + "build", + job_input=event, + options=MicrovmInvokeOptions( + timeout=Duration.from_minutes(15), + heartbeat_timeout=Duration.from_seconds(30), + suspend_after=True, + ), +) +approval = ctx.create_callback(name="approval").result() # a long wait +vm.invoke( + ctx, + "deploy", + job_input=build, + options=MicrovmInvokeOptions( + timeout=Duration.from_minutes(10), + heartbeat_timeout=Duration.from_seconds(120), + ), +) +``` + +What happens: + +1. The job request carries `"suspendAfter": true`. The job document stays at version 1, and the field is omitted for every other job. The run hook payload never carries it. +2. The worker acknowledges the field with the header `X-Microvm-Suspend-After: accepted` in its 202 answer. +3. The worker suspends the MicroVM when the job has ended and its outcome is reported, and no other job runs. A job that ends while another job runs leaves the request pending. The worker then suspends the MicroVM when the last running job ends. +4. The worker suspends the MicroVM that its `run` hook named. It never uses the MicroVM ID of a job request. +5. Each later job of the session handle checks the MicroVM state before delivery, and resumes a suspended MicroVM, as with `auto_suspend_on_idle`. The `suspend_after` job itself checks the state too, because the record is set before the job starts. + +The setting works with or without `auto_suspend_on_idle`. With both, the worker suspends right after the job, and the idle time applies after each resume. + +The record that a job set `suspend_after` is kept in memory only. A replay of a completed job calls `vm.invoke` again, so it sets the record again. A replay of a completed `map` item, `parallel` branch, or child context does not run its function again. So the next job can skip the state check of a suspended MicroVM. For that case, every session job checks the state after a sign of a suspended MicroVM: + +1. The sign is an HTTP 502 with `x-aws-proxy-error: MICROVM_CONNECT_FAILED`, or a request without an answer. A booting MicroVM's endpoint holds the request until the worker listens. A boot gives no such sign, unless the worker takes more than 10 seconds to start listening. Then the first delivery gets no answer, and the operation makes one GetMicrovm call. +2. The request step calls `GetMicrovm` once. It resumes a `SUSPENDING` or `SUSPENDED` MicroVM, waits for `RUNNING`, and sends the job again. A resume that the service rejects with a conflict while the MicroVM is still suspending is sent again once the MicroVM is suspended. +3. The check runs again after each failed request, at most once per 2 seconds. +4. A denied or failed `GetMicrovm` logs one warning per job: `GetMicrovm failed, so the operation could not check whether the MicroVM is suspended. The request step keeps sending the job.` The request step then keeps sending the job. After a denied call, the final `MicrovmEndpointUnavailableError` message adds `GetMicrovm was denied, so the operation could not check whether the MicroVM is suspended.` +5. A denied `ResumeMicrovm` for a suspending or suspended MicroVM fails the job at once with `MicrovmDeliveryError`. Its message names `lambda:ResumeMicrovm`. +6. A terminating or terminated MicroVM fails the job with `MicrovmNotRunningError`. + +A session with `auto_suspend_on_idle` or an `idle_policy`, and a handle after a `suspend_after` job, keep the state check before delivery. They also keep the check after any 502, 503, or 504, or a request without an answer, and its errors: a denied `GetMicrovm` fails the job. + +Use `suspend_after` only on the job before a long wait, such as a durable wait or a callback: + +- Not on back-to-back jobs. Each suspend writes a snapshot, and each resume adds about 1 second before the next job. +- Not on `map` items. A request waits until no job runs. So a `map` with `suspend_after` on its items suspends and resumes between items that do not overlap. +- Not on the last job. The session terminates the MicroVM after its handler returns, so the snapshot is wasted. + +It needs: + +1. The `resume` hook in the image. Without it, the worker refuses the next job with HTTP 503 for up to 30 seconds after the suspend. +2. A `heartbeat_timeout` on the next job of at least its `retry_window` plus 30 seconds. The heartbeat timeout starts before delivery, so it includes the resume and the refusal. +3. `lambda:GetMicrovm` and `lambda:ResumeMicrovm` in the function's role. +4. `lambda:SuspendMicrovm` on the image in the MicroVM's role. That permission lets code in any MicroVM of the image suspend any other. So use one image per trust boundary. +5. A worker that supports the field. An older worker ignores it, and the MicroVM keeps running. Each execution environment then logs one warning per session: `the worker in MicroVM does not support suspendAfter. The MicroVM keeps running after the job.` + +A suspend that the service refuses for now, such as a throttle or expired credentials, is not retried without `auto_suspend_on_idle`. The MicroVM then keeps running until a later suspend. With `auto_suspend_on_idle`, the worker tries again after the idle time, as it does for an idle suspend. + ### Session permissions -A session needs the permissions of `microvm` with HTTP delivery: `lambda:RunMicrovm`, `lambda:TerminateMicrovm`, `lambda:CreateMicrovmAuthToken`, `lambda:PassNetworkConnector`, and `iam:PassRole`. Suspending adds the permissions listed above. +A session needs the permissions of `microvm` with HTTP delivery: `lambda:RunMicrovm`, `lambda:TerminateMicrovm`, `lambda:CreateMicrovmAuthToken`, `lambda:PassNetworkConnector`, and `iam:PassRole`. Suspending adds the permissions listed above. Grant `lambda:GetMicrovm` to every session's function role. Without it, the check after a sign of a suspended MicroVM logs a warning, and the job keeps its retries. ### Subtypes @@ -342,7 +406,8 @@ The MicroVM operations handle the network, the service, and the worker lifecycle - The idle time is at most the session `timeout`. A value that breaks any one condition raises `ValidationError` when `microvm_session` is called, before any durable operation starts. See [Suspending an idle MicroVM](#suspending-an-idle-microvm). -4. Change `auto_suspend_on_idle` or `idle_policy` only through a new version, as in [Deploys and versions](#deploys-and-versions). The launch records the resume mode, and a deploy changes it only for sessions that launch after it. Every job of a session that launched with suspending needs `lambda:GetMicrovm` and `lambda:ResumeMicrovm`, also after a deploy turns suspending off. So remove them from the function's role only as [Suspending an idle MicroVM](#suspending-an-idle-microvm) describes. +4. Set `suspend_after` only on the job before a long wait, and give the next job a `heartbeat_timeout` of at least its `retry_window` plus 30 seconds. On back-to-back jobs and `map` items, each suspend and resume adds latency and a snapshot. A shorter heartbeat timeout can fail the next job with `MicrovmTimeoutError` while the MicroVM resumes. Keep `lambda:GetMicrovm` and `lambda:ResumeMicrovm` in the function's role after a deploy that removes `suspend_after`, until every session that ran such a job has ended. See [Suspending after a job](#suspending-after-a-job). +5. Change `auto_suspend_on_idle` or `idle_policy` only through a new version, as in [Deploys and versions](#deploys-and-versions). The launch records the resume mode, and a deploy changes it only for sessions that launch after it. Every job of a session that launched with suspending needs `lambda:GetMicrovm` and `lambda:ResumeMicrovm`, also after a deploy turns suspending off. So remove them from the function's role only as [Suspending an idle MicroVM](#suspending-an-idle-microvm) describes. ### Retry strategies diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py index 18a8bd10e..14301c0f3 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py @@ -352,4 +352,46 @@ class MicrovmInvokeOptions(MicrovmRequestConfig): Set it on every job. A MicroVM that ends while suspended, or whose image disables the ``terminate`` hook, reports nothing. The heartbeat timeout is then the only way that the job fails soon after the MicroVM ends. + + A job after a :attr:`suspend_after` job waits for the resume of the + MicroVM. The worker can refuse it with HTTP 503 for up to 30 seconds + before that. So set the heartbeat timeout of such a job to at least the + ``retry_window`` plus 30 seconds. + """ + + suspend_after: bool = False + """Whether the worker suspends the MicroVM after this job. Defaults to + ``False``. + + The worker suspends the MicroVM when the job has ended and its outcome is + reported, and no other job runs. A job that ends while another job runs + leaves the request pending. The worker then suspends the MicroVM when + the last running job ends. The next job of the session resumes the + MicroVM before it is delivered. The ``suspend_after`` job itself checks + the state too, because the record is set before the job starts. The + setting works with or without + :attr:`MicrovmSessionConfig.auto_suspend_on_idle`. + + Set it on the job before a long wait, such as a durable wait or a + callback. Do not set it on back-to-back jobs, on ``map`` items, or on + the last job of the session. Each suspend and resume costs a snapshot + and about 1 second of latency for the next job. The session terminates + the MicroVM after its last job anyway. + + It needs more than a session without suspending: + + 1. The image enables the ``resume`` hook. Without it, the next job waits + up to 30 seconds of HTTP 503 answers. + 2. The function's role has ``lambda:GetMicrovm`` and + ``lambda:ResumeMicrovm``. + 3. The MicroVM's role has ``lambda:SuspendMicrovm`` on the image. + 4. The worker supports the setting. An older worker ignores it, and the + MicroVM keeps running. Each execution environment then logs one + warning per session. + + A suspend that the service refuses for now is not retried, unless + :attr:`MicrovmSessionConfig.auto_suspend_on_idle` is ``True``. Only + :meth:`MicrovmSession.invoke` accepts it. :func:`microvm` raises + ``ValidationError`` for it. A value that is not a ``bool`` raises + ``ValidationError`` too. """ diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/errors.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/errors.py index 18cd84287..5730c3dc1 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/errors.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/errors.py @@ -122,12 +122,16 @@ class MicrovmNotRunningError(MicrovmDeliveryError): the same state, so the job is not retried. To continue a session, start a new session. - Two checks raise it: + Three checks raise it: 1. A session that checks the MicroVM state before each job. It checks - when ``auto_suspend_on_idle`` is ``True``, or when an ``idle_policy`` - is set. ``GetMicrovm`` or ``ResumeMicrovm`` then reports the state. - 2. Every job over HTTP, of a session or of :func:`microvm`. The request + when ``auto_suspend_on_idle`` is ``True``, when an ``idle_policy`` + is set, and after a job with ``suspend_after``. ``GetMicrovm`` or + ``ResumeMicrovm`` then reports the state. + 2. Every session job, after a request without an answer, or after an + HTTP 502 with ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``. The + request step then calls ``GetMicrovm`` once. + 3. Every job over HTTP, of a session or of :func:`microvm`. The request step calls ``CreateMicrovmAuthToken`` before the first request, and again after a 401 or 403. That call can answer ``ResourceNotFoundException``. @@ -137,14 +141,15 @@ class MicrovmNotRunningError(MicrovmDeliveryError): records only the message of the failure. So a support case can quote the request ID from the message, on the first run and on replay. - Without the state check, a session makes no ``GetMicrovm`` call. The - request step still calls ``CreateMicrovmAuthToken`` first. When that - call answers ``ResourceNotFoundException``, the job fails with this - error. When it returns a token, the request goes to the endpoint. The - endpoint of a terminated MicroVM answers 502, and the request step - retries a 502 as a MicroVM that is still starting. So the job then fails - with a plain :class:`MicrovmDeliveryError`, after the retries of both - tiers. + Without the state check before delivery, the request step still calls + ``CreateMicrovmAuthToken`` first. When that call answers + ``ResourceNotFoundException``, the job fails with this error. When it + returns a token, the request goes to the endpoint. The endpoint of a + terminated MicroVM answers 502. Whether that 502 carries + ``MICROVM_CONNECT_FAILED`` is not verified. Without it, the request + step retries the 502 as a MicroVM that is still starting. So the job + then fails with a plain :class:`MicrovmDeliveryError`, after the + retries of both tiers. The live service's answer of ``CreateMicrovmAuthToken`` for a removed MicroVM is not verified. The API model lists diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/job.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/job.py index 59e7bee0e..cee564937 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/job.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/job.py @@ -36,8 +36,9 @@ 1. The worker accepts a job request body of at most 6 MiB. 2. The body holds the input as the same JSON that :func:`dumps` writes. 3. The body also holds fixed fields: the version, the Region, the MicroVM - ID, the callback ID, and the heartbeat timeout. Their names and the JSON - punctuation take 92 bytes. + ID, the suspend-after flag, the callback ID, and the heartbeat timeout. + Their names and the JSON punctuation take 112 bytes. The suspend-after + flag, ``,"suspendAfter":true``, takes 20 of them. 4. The botocore ``lambda-microvms`` model allows a MicroVM ID of 1 to 256 characters, with no pattern. JSON writes a character in at most 6 bytes, as a ``\\u00XX`` escape. So the MicroVM ID takes at most 1,536 bytes. @@ -61,6 +62,9 @@ class RunHookPayload: A job that goes over HTTP needs a ``runHookPayload`` too, because the worker rejects a payload without ``version`` and ``region``. See :func:`run_hook_payload_without_job`. + + The payload never carries ``suspendAfter``. Only a session job can set + it, and every session job goes over HTTP. """ region: str @@ -99,13 +103,19 @@ def to_json(self) -> str: class JobRequest: """The body of the HTTP request that delivers one job to a route. - Wire names: ``version``, ``region``, ``microvmId``, ``callbackId``, - ``heartbeatTimeoutSeconds``, and ``input``. The job fields sit at the top - level, next to ``version``, ``region``, and ``microvmId``. + Wire names: ``version``, ``region``, ``microvmId``, ``suspendAfter``, + ``callbackId``, ``heartbeatTimeoutSeconds``, and ``input``. The job + fields sit at the top level, next to ``version``, ``region``, and + ``microvmId``. The endpoint can accept a request before Lambda has sent the ``run`` hook. The worker learns the MicroVM ID from the ``run`` hook. So the request carries the MicroVM ID too. + + ``suspendAfter`` is ``true`` only for a session job that asks the worker + to suspend the MicroVM after the job. The field is omitted otherwise. An + older worker ignores a field that it does not know. So the version stays + 1, because a new version would make every older worker fail the job. """ region: str @@ -118,6 +128,8 @@ class JobRequest: """The caller's input. It must be JSON-serializable.""" heartbeat_timeout_seconds: int | None = None """The heartbeat timeout, when the caller set one.""" + suspend_after: bool = False + """Whether the worker suspends the MicroVM after the job.""" def to_json(self) -> str: """Serialize the request as compact JSON, as :meth:`RunHookPayload.to_json` does. @@ -126,16 +138,19 @@ def to_json(self) -> str: TypeError: When the input is not JSON-serializable. ValueError: When the input holds ``nan`` or an infinite float. """ - return dumps( - { - "version": PAYLOAD_VERSION, - "region": self.region, - "microvmId": self.microvm_id, - **_job_fields( - self.callback_id, self.job_input, self.heartbeat_timeout_seconds - ), - } + document: dict[str, Any] = { + "version": PAYLOAD_VERSION, + "region": self.region, + "microvmId": self.microvm_id, + } + if self.suspend_after: + document["suspendAfter"] = True + document.update( + _job_fields( + self.callback_id, self.job_input, self.heartbeat_timeout_seconds + ) ) + return dumps(document) def run_hook_payload_without_job( diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py index 5e31324ca..11ef931b3 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py @@ -7,11 +7,11 @@ import time from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Literal from aws_durable_execution_sdk_python.types import LoggerInterface -from botocore.exceptions import ClientError +from botocore.exceptions import BotoCoreError, ClientError from aws_durable_execution_sdk_python_extras.microvm.aws_errors import ( add_request_id, @@ -63,6 +63,31 @@ def __init__(self, message: str, state: str | None = None) -> None: """The MicroVM state, or ``None`` when the MicroVM no longer exists.""" +class MicrovmStateCheckError(Exception): + """``GetMicrovm`` failed in the check of :func:`resume_if_suspended`. + + The request step of a session without the resume mode runs that check + only after a sign of a suspended MicroVM. The session may never have + asked to suspend. So the step logs one warning per job, and keeps + sending the job. ``__cause__`` is the error of the call. + """ + + def __init__(self, message: str, *, denied: bool) -> None: + super().__init__(message) + self.denied = denied + """Whether the service answered ``AccessDeniedException``.""" + + +class MicrovmResumeDeniedError(Exception): + """``ResumeMicrovm`` was denied for a suspending or suspended MicroVM. + + Nothing else resumes the MicroVM, and a later attempt is denied the same + way. So the default retry strategy does not retry it, and the delivery + raises ``MicrovmDeliveryError`` at once. The message names the + permission, and ends with the request ID of the call. + """ + + @dataclass(frozen=True, kw_only=True) class EnsureRunningOptions: """The inputs of :func:`ensure_running`.""" @@ -88,10 +113,10 @@ def ensure_running(options: EnsureRunningOptions) -> str | None: 2. ``PENDING``: the function returns at once. A new MicroVM reports ``PENDING`` until it boots. The request's own retries wait for the ``run`` hook. - 3. ``SUSPENDED``: it calls ``ResumeMicrovm`` once, then waits for + 3. ``SUSPENDED``: it calls ``ResumeMicrovm``, then waits for ``RUNNING``. The MicroVM stays ``SUSPENDED`` while its ``resume`` hook - runs. So the function does not call ``ResumeMicrovm`` again for each - ``SUSPENDED`` answer. + runs. So after a successful call, the function does not call + ``ResumeMicrovm`` again for each ``SUSPENDED`` answer. 4. ``SUSPENDING``: it waits. The MicroVM becomes ``SUSPENDED``, and the function then resumes it. 5. ``TERMINATING``, ``TERMINATED``, or not found: it raises @@ -104,7 +129,9 @@ def ensure_running(options: EnsureRunningOptions) -> str | None: ``AccessDeniedException``, is raised with its class and its ``response``. Its message then ends with the request ID of the call. See :func:`add_request_id`. A ``ConflictException`` from ``ResumeMicrovm`` - is not an error: polling continues. + is not an error: polling continues. The call did not resume the + MicroVM. So the next ``SUSPENDED`` answer calls ``ResumeMicrovm`` again, + at most once per poll. It polls after 250 milliseconds, and doubles the wait up to 2 seconds. It stops after ``max_wait_seconds``, or 10 seconds before the invocation @@ -124,12 +151,150 @@ def ensure_running(options: EnsureRunningOptions) -> str | None: ClientError: When ``GetMicrovm`` or ``ResumeMicrovm`` fails otherwise. """ + return _wait_until_running(options, resuming=False, resume_accepted=False) + + +def resume_if_suspended(options: EnsureRunningOptions) -> str | None: + """Check the MicroVM state once, resume a suspended MicroVM, and return + its current endpoint. + + The request step of a session without the resume mode calls it after a + sign that the MicroVM may be suspended: an HTTP 502 from the endpoint + with ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``, or no answer. A + worker can suspend its MicroVM after a ``suspend_after`` job. The + session can lose the record of that job on replay. So this check runs + for every session job. + + 1. ``GetMicrovm`` runs once. A failure other than + ``ResourceNotFoundException`` raises :class:`MicrovmStateCheckError`. + The caller then keeps sending the job. + 2. ``RUNNING`` or ``PENDING``: the function returns at once. + 3. ``SUSPENDING`` or ``SUSPENDED``: it calls ``ResumeMicrovm`` once. + The service honors a resume during ``SUSPENDING`` too. It then waits + for ``RUNNING``, as :func:`ensure_running` does after its own resume. + 4. ``TERMINATING``, ``TERMINATED``, or not found: it raises + :class:`MicrovmStateError`. + + A denied ``ResumeMicrovm`` raises :class:`MicrovmResumeDeniedError`. + ``ConflictException`` is not an error: polling continues. A resume that + the service rejects with a conflict while the MicroVM is still + suspending is sent again once the MicroVM is suspended. Any other + ``ResumeMicrovm`` error is raised with its class and its ``response``, + as in :func:`ensure_running`. So is any error of a ``GetMicrovm`` call + after the resume. + + Returns: + The endpoint that ``GetMicrovm`` reports, or ``None`` when it reports + none. + + Raises: + MicrovmStateCheckError: When the first ``GetMicrovm`` call fails. + MicrovmStateError: When the MicroVM can no longer run a job. + MicrovmResumeDeniedError: When ``ResumeMicrovm`` is denied. + MicrovmEndpointUnavailableError: When the MicroVM does not reach + ``RUNNING`` in time. + ClientError: When ``ResumeMicrovm`` fails otherwise, or a later + ``GetMicrovm`` call fails. + """ + started = time.monotonic() + reserve = reserve_starts_at(options.remaining_seconds) + microvm_id = options.microvm_id + + def unavailable() -> Exception: + elapsed_ms = round((time.monotonic() - started) * 1000) + return MicrovmEndpointUnavailableError( + f"MicroVM {microvm_id} did not reach RUNNING in {elapsed_ms} ms. " + "Last state: unknown." + ) + + try: + current = call_before_deadline( + lambda: options.client.get_microvm(microvmIdentifier=microvm_id), + reserve, + unavailable, + ) + except ClientError as error: + code = error_code(error) + if code == "ResourceNotFoundException": + raise microvm_not_found(microvm_id, "GetMicrovm", error) from error + raise MicrovmStateCheckError( + str(add_request_id("GetMicrovm", error)), + denied=code == "AccessDeniedException", + ) from error + except BotoCoreError as error: + raise MicrovmStateCheckError(str(error), denied=False) from error + state = current.get("state") + _raise_if_ended(microvm_id, current) + if state not in {"SUSPENDING", "SUSPENDED"}: + return current.get("endpoint") + try: + call_before_deadline( + lambda: options.client.resume_microvm(microvmIdentifier=microvm_id), + reserve, + unavailable, + ) + accepted = True + except ClientError as error: + code = error_code(error) + if code == "ResourceNotFoundException": + raise microvm_not_found(microvm_id, "ResumeMicrovm", error) from error + if code == "AccessDeniedException": + msg = ( + f"MicroVM {microvm_id} is {state}. ResumeMicrovm was denied, " + "so the operation cannot resume it. The function's execution " + "role needs lambda:ResumeMicrovm." + ) + response = ( + error.response + if is_microvms_service_error("ResumeMicrovm", error) + else None + ) + raise MicrovmResumeDeniedError( + with_request_id(msg, "ResumeMicrovm", response) + ) from error + if code != "ConflictException": + raise add_request_id("ResumeMicrovm", error) + # A conflict can mean that the suspend has not ended. The MicroVM + # then becomes SUSPENDED and needs a new resume. So the loop sends + # ResumeMicrovm again at the next SUSPENDED answer. + accepted = False + options.logger.info("resuming MicroVM", extra={"state": state}) + remaining = options.max_wait_seconds - (time.monotonic() - started) + return _wait_until_running( + replace(options, max_wait_seconds=remaining), + resuming=True, + resume_accepted=accepted, + ) + + +def _raise_if_ended(microvm_id: str, current: Any) -> None: + """Raise :class:`MicrovmStateError` for a terminating or terminated + MicroVM.""" + state = current.get("state") + if state in {"TERMINATING", "TERMINATED"}: + reason = current.get("stateReason") + # The GetMicrovm output carries the request ID. The state fails + # the step, so the message names that request ID. + msg = f"MicroVM {microvm_id} is {state}{f': {reason}' if reason else ''}." + raise MicrovmStateError(with_request_id(msg, "GetMicrovm", current), state) + + +def _wait_until_running( + options: EnsureRunningOptions, *, resuming: bool, resume_accepted: bool +) -> str | None: + """The loop of :func:`ensure_running`. + + ``resuming`` states that the caller has already called + ``ResumeMicrovm``. A ``PENDING`` answer is then a step of the resume, so + the loop waits for ``RUNNING``. ``resume_accepted`` states that a + ``ResumeMicrovm`` call succeeded. A call that answered + ``ConflictException`` leaves it ``False``. So the next ``SUSPENDED`` + answer sends ``ResumeMicrovm`` again.""" started = time.monotonic() reserve = reserve_starts_at(options.remaining_seconds) deadline = started + options.max_wait_seconds if reserve is not None: deadline = min(deadline, reserve) - resume_sent = False delay = _INITIAL_POLL_SECONDS last_state: str | None = None microvm_id = options.microvm_id @@ -171,27 +336,25 @@ def resume_microvm() -> Any: raise add_request_id("GetMicrovm", error) state = current.get("state") last_state = state - if state == "RUNNING" or (state == "PENDING" and not resume_sent): - if resume_sent: + if state == "RUNNING" or (state == "PENDING" and not resuming): + if resuming: options.logger.info( "MicroVM resumed", extra={"waitedMs": round((time.monotonic() - started) * 1000)}, ) return current.get("endpoint") - if state in {"TERMINATING", "TERMINATED"}: - reason = current.get("stateReason") - # The GetMicrovm output carries the request ID. The state fails - # the step, so the message names that request ID. - msg = f"MicroVM {microvm_id} is {state}{f': {reason}' if reason else ''}." - raise MicrovmStateError(with_request_id(msg, "GetMicrovm", current), state) - if state == "SUSPENDED" and not resume_sent: + _raise_if_ended(microvm_id, current) + if state == "SUSPENDED" and not resume_accepted: try: resume_microvm() + resume_accepted = True except ClientError as error: - # ConflictException means that a transition is already in - # progress, for example an auto-resume. Polling continues - # either way. ResourceNotFoundException means that the - # MicroVM was removed after the state check. + # ConflictException means that a transition is in progress, + # for example a suspend that has not ended, or an + # auto-resume. The resume stays unaccepted, so the next + # SUSPENDED answer sends it again. The poll interval and the + # deadline bound those calls. ResourceNotFoundException + # means that the MicroVM was removed after the state check. code = error_code(error) if code == "ResourceNotFoundException": raise microvm_not_found( @@ -199,8 +362,9 @@ def resume_microvm() -> Any: ) from error if code != "ConflictException": raise add_request_id("ResumeMicrovm", error) - resume_sent = True - options.logger.info("resuming MicroVM", extra={"state": state}) + if not resuming: + options.logger.info("resuming MicroVM", extra={"state": state}) + resuming = True if time.monotonic() + delay > deadline: raise unavailable() diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/operation.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/operation.py index b7e5759ee..39b20b02a 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/operation.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/operation.py @@ -65,9 +65,11 @@ ) from aws_durable_execution_sdk_python_extras.microvm.lifecycle import ( EnsureRunningOptions, + MicrovmStateCheckError, MicrovmStateError, ensure_running, microvm_not_found, + resume_if_suspended, ) from aws_durable_execution_sdk_python_extras.microvm.request import ( DEFAULT_MICROVM_JOB_PATH, @@ -256,7 +258,9 @@ def microvm( ValidationError: When ``name`` or ``config`` is not valid, or the Region cannot be determined. No durable operation starts. It is the SDK's class, which ``Duration`` also raises for a negative - value. + value. A ``config.request`` of type :class:`MicrovmInvokeOptions` + with ``suspend_after`` set is not valid, because only a session + job can suspend its MicroVM after the job. TypeError: When ``job_input`` is not JSON-serializable. No durable operation starts. ValueError: When the JSON of ``job_input`` is over @@ -290,6 +294,20 @@ def run(child: DurableContext) -> Any: """The error type that a request step checkpoints for a MicroVM that can no longer run a job. The SDK records the fully-qualified class name.""" +_STATE_CHECK_WARNING = ( + "GetMicrovm failed, so the operation could not check whether the MicroVM " + "is suspended. The request step keeps sending the job." +) +"""The warning, once per job, after a failed state check of a session +without the resume mode.""" + +_GET_DENIED_NOTE = ( + "GetMicrovm was denied, so the operation could not check whether the " + "MicroVM is suspended." +) +"""The sentence that the first tier adds to its final error after a denied +state check. Without it, the message names only the last HTTP failure.""" + def step_retry_strategy( config: MicrovmBaseConfig, @@ -479,7 +497,10 @@ def deliver( heartbeat_timeout_seconds: int | None, request: MicrovmRequestConfig, retry_window_seconds: int, - resume: bool = False, + resume: bool | Callable[[], bool] = False, + suspend_check: bool = False, + suspend_after: bool = False, + claim_suspend_after_warning: Callable[[], bool] | None = None, ) -> None: """Run the request step, ``.request``: POST the job to the MicroVM's endpoint. @@ -490,7 +511,16 @@ def deliver( With ``resume``, each attempt first makes sure the MicroVM runs, and resumes a suspended one. It then sends the job to the endpoint that ``GetMicrovm`` reports. A session sets it when its MicroVM can be - suspended. + suspended. A callable is read at the start of each attempt. + + Without ``resume``, ``suspend_check`` turns on the check after a sign + of a suspended MicroVM. See :func:`resume_if_suspended`. A session + sets it for every job. A failed ``GetMicrovm`` call in that check + logs one warning per job, and the attempt keeps sending the job. + + ``suspend_after`` adds ``"suspendAfter": true`` to the job request. + When the answer lacks the worker's acknowledgement, the attempt + logs a warning if ``claim_suspend_after_warning`` returns ``True``. ``GetMicrovm``, ``ResumeMicrovm``, and ``CreateMicrovmAuthToken`` answer ``ResourceNotFoundException`` for a removed MicroVM. An @@ -508,6 +538,7 @@ def deliver( callback_id=callback_id, job_input=job_input, heartbeat_timeout_seconds=heartbeat_timeout_seconds, + suspend_after=suspend_after, ) .to_json() .encode() @@ -515,9 +546,14 @@ def deliver( remaining_seconds = _remaining_seconds(child) path = DEFAULT_MICROVM_JOB_PATH if request.path is None else request.path port = DEFAULT_MICROVM_PORT if request.port is None else request.port + microvm_id = launched.microvm_id + # One warning per job for a failed state check. The attempts of the + # step in this invocation share it. + check_warned = threading.Event() def attempt(step_context: StepContext) -> None: client = self._client() + logger = step_context.logger # The state check before delivery and the request's retries share # one window. So the wait for a resume counts toward it. window_started_at = time.monotonic() @@ -525,23 +561,57 @@ def attempt(step_context: StepContext) -> None: # HTTP delivery. The launch step checkpoints only what # _launch_call returned. So an HTTP launch always has an endpoint. endpoint = cast("str", launched.endpoint) + strict = resume() if callable(resume) else resume + + def options(max_wait_seconds: float) -> EnsureRunningOptions: + return EnsureRunningOptions( + client=client, + microvm_id=microvm_id, + max_wait_seconds=max_wait_seconds, + remaining_seconds=remaining_seconds, + logger=logger, + ) + recheck: Callable[[float], str | None] | None = None - if resume: + check: Callable[[float], str | None] | None = None + get_denied = False + if strict: require_time_for_call(remaining_seconds) def running(max_wait_seconds: float) -> str | None: - return ensure_running( - EnsureRunningOptions( - client=client, - microvm_id=launched.microvm_id, - max_wait_seconds=max_wait_seconds, - remaining_seconds=remaining_seconds, - logger=step_context.logger, - ) - ) + return ensure_running(options(max_wait_seconds)) endpoint = running(retry_window_seconds) or endpoint recheck = running + elif suspend_check: + + def check_state(max_wait_seconds: float) -> str | None: + nonlocal get_denied + try: + return resume_if_suspended(options(max_wait_seconds)) + except MicrovmStateCheckError as error: + get_denied = get_denied or error.denied + if not check_warned.is_set(): + check_warned.set() + logger.warning( + _STATE_CHECK_WARNING, + extra={"microvmId": microvm_id, "error": str(error)}, + ) + return None + + check = check_state + + def ignored() -> None: + if ( + claim_suspend_after_warning is not None + and claim_suspend_after_warning() + ): + logger.warning( + f"the worker in MicroVM {microvm_id} does not support " + "suspendAfter. The MicroVM keeps running after the job.", + extra={"microvmId": microvm_id}, + ) + # Only CreateMicrovmAuthToken can raise ResourceNotFoundException # here: # @@ -549,9 +619,9 @@ def running(max_wait_seconds: float) -> str | None: # 2. send_job calls new_token() outside any try. So the error # leaves send_job unchanged. # 3. _run_with_limit re-raises the same error object. - # 4. ensure_running turns the same answer of GetMicrovm and - # ResumeMicrovm into MicrovmStateError. So neither call raises - # a ClientError with that code. + # 4. ensure_running and resume_if_suspended turn the same answer + # of GetMicrovm and ResumeMicrovm into MicrovmStateError. So + # neither call raises a ClientError with that code. # # So the step raises MicrovmStateError for CreateMicrovmAuthToken. # The default retry strategy does not retry that error. @@ -559,23 +629,28 @@ def running(max_wait_seconds: float) -> str | None: send_job( SendJobOptions( client=client, - microvm_id=launched.microvm_id, + microvm_id=microvm_id, endpoint=endpoint, path=path, port=port, body=body, retry_window_seconds=retry_window_seconds, remaining_seconds=remaining_seconds, - logger=step_context.logger, + logger=logger, opener=self.config.opener, window_started_at=window_started_at, recheck=recheck, + suspend_check=check, + unavailable_note=lambda: ( + _GET_DENIED_NOTE if get_denied else None + ), + on_suspend_after_ignored=ignored if suspend_after else None, ) ) except ClientError as error: if error_code(error) == "ResourceNotFoundException": raise microvm_not_found( - launched.microvm_id, "CreateMicrovmAuthToken", error + microvm_id, "CreateMicrovmAuthToken", error ) from error raise @@ -743,6 +818,16 @@ def create( heartbeat_timeout_seconds = _validate_heartbeat(name, config.heartbeat_timeout) request = config.request or MicrovmRequestConfig() retry_window_seconds = _validate_request(name, request) + # MicrovmInvokeOptions subclasses MicrovmRequestConfig. So a session + # job's options pass the type check of config.request. microvm() + # terminates its MicroVM after its one job, so a suspend after the + # job would only delay that. + if _validate_suspend_after(name, getattr(request, "suspend_after", False)): + msg = ( + f'MicroVM "{name}": suspend_after applies only to a session job. ' + "microvm() terminates its MicroVM after the job." + ) + raise ValidationError(msg) _validate_input(name, job_input) partition, region = _validate_region(context, name) return cls( @@ -1038,6 +1123,19 @@ def _validate_request(name: str, request: MicrovmRequestConfig) -> int: return seconds +def _validate_suspend_after(name: str, suspend_after: object) -> bool: + """Validate ``suspend_after``, and return it. + + Raises: + ValidationError: When it is not a ``bool``. An ``int`` is rejected + too, although ``bool`` subclasses ``int``. + """ + if not isinstance(suspend_after, bool): + msg = f'MicroVM "{name}": suspend_after must be a bool' + raise ValidationError(msg) + return suspend_after + + def _remaining_seconds(context: DurableContext) -> Callable[[], float | None]: """Return a reader of the invocation's remaining time, in seconds. diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py index 9a109b9c9..b2c28fe05 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py @@ -73,6 +73,22 @@ _AUTH_TOKEN_MINUTES = 5 _PORT_HEADER = "X-aws-proxy-port" +SUSPEND_AFTER_HEADER = "X-Microvm-Suspend-After" +"""The header with which a worker acknowledges ``suspendAfter`` in its 202 +answer. Its value is :data:`SUSPEND_AFTER_ACCEPTED`. The JavaScript SDK +uses the same header.""" + +SUSPEND_AFTER_ACCEPTED = "accepted" + +_PROXY_ERROR_HEADER = "x-aws-proxy-error" +_CONNECT_FAILED = "MICROVM_CONNECT_FAILED" +"""The endpoint answers HTTP 502 with ``x-aws-proxy-error: +MICROVM_CONNECT_FAILED`` while the MicroVM is suspending or suspended. A +probe measured the answer in about 60 milliseconds.""" + +_SUSPEND_CHECK_INTERVAL_SECONDS = 2.0 +"""The shortest time between two calls of ``SendJobOptions.suspend_check``.""" + _T = TypeVar("_T") @@ -239,6 +255,32 @@ class SendJobOptions: seconds, so the wait for a resume never extends the window. ``None`` from it keeps the current endpoint. """ + suspend_check: Callable[[float], str | None] | None = None + """Resume the MicroVM if it is suspended, and return its current + endpoint. + + A session job without :attr:`recheck` sets it. A worker can suspend the + MicroVM after a ``suspend_after`` job, and the session can lose the + record of that job on replay. A suspended or suspending MicroVM gets no + answer to a request, or its endpoint answers HTTP 502 with + ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``. So the next attempt calls + it after either sign, at most once per 2 seconds. A booting MicroVM's + endpoint holds the request until the worker listens. A boot gives no + such sign, unless the worker takes more than 10 seconds to start + listening. Then the first delivery gets no answer, and the operation + makes one ``GetMicrovm`` call. It receives what remains of the retry + window, in seconds. ``None`` from it keeps the current endpoint. + """ + unavailable_note: Callable[[], str | None] | None = None + """A sentence to add to the message of + :class:`MicrovmEndpointUnavailableError`, or ``None``.""" + on_suspend_after_ignored: Callable[[], None] | None = None + """Called when the answer that accepted the job lacks the header + ``X-Microvm-Suspend-After: accepted``. + + A ``suspend_after`` job sets it. A worker that supports the field sends + the header. An older worker ignores the field, and the MicroVM keeps + running after the job.""" def require_time_for_call(remaining_seconds: Callable[[], float | None]) -> None: @@ -315,6 +357,18 @@ def send_job(options: SendJobOptions) -> int: The endpoint can accept a request before Lambda has sent the ``run`` hook. So the body carries the MicroVM ID too. + A session job checks the MicroVM state in one of two ways. With + ``recheck``, it checks after any 502, 503, 504, or request without an + answer. With ``suspend_check``, it checks only after a 502 with + ``x-aws-proxy-error: MICROVM_CONNECT_FAILED`` or a request without an + answer, at most once per 2 seconds. The function then sends the job + again to the endpoint that the check reports. + + A ``suspend_after`` job also reads the answer that accepted it. A worker + that supports the field answers with ``X-Microvm-Suspend-After: + accepted``. Without that header, the function calls + ``on_suspend_after_ignored``. + The same job can reach the route more than once: 1. A request reaches the route, and its answer is lost. @@ -340,8 +394,10 @@ def send_job(options: SendJobOptions) -> int: ``ResourceNotFoundException``. It is the error that the client raised, unchanged. The request step turns it into ``MicrovmStateError``. - MicrovmStateError: When ``recheck`` reports that the MicroVM can no - longer run a job. + MicrovmStateError: When ``recheck`` or ``suspend_check`` reports + that the MicroVM can no longer run a job. + MicrovmResumeDeniedError: When ``suspend_check`` finds a suspended + MicroVM, and ``ResumeMicrovm`` is denied. Exception: When ``CreateMicrovmAuthToken`` fails otherwise. The message of a botocore error ends with the request ID. """ @@ -357,9 +413,11 @@ def send_job(options: SendJobOptions) -> int: def unavailable() -> Exception: elapsed_ms = round((time.monotonic() - started) * 1000) + note = options.unavailable_note() if options.unavailable_note else None return MicrovmEndpointUnavailableError( f"The MicroVM did not accept the job at {options.path} after " f"{attempt} attempts in {elapsed_ms} ms. Last failure: {last_failure}." + + (f" {note}" if note else "") ) def deadline() -> float: @@ -378,6 +436,7 @@ def new_token() -> Mapping[str, str]: token = new_token() refreshed_token = False delay = _INITIAL_DELAY_SECONDS + last_suspend_check: float | None = None while True: attempt += 1 @@ -410,6 +469,11 @@ def new_token() -> Mapping[str, str]: options.logger.info( "MicroVM job accepted", extra={"attempt": attempt, "status": status} ) + if ( + options.on_suspend_after_ignored is not None + and answer.suspend_after != SUSPEND_AFTER_ACCEPTED + ): + options.on_suspend_after_ignored() return status if status is None: last_failure = answer.failure @@ -466,6 +530,25 @@ def new_token() -> Mapping[str, str]: raise unavailable() endpoint = options.recheck(left) or endpoint url = _endpoint_url(endpoint, options.path) + # A session without the resume mode checks the state only after a + # sign of a suspended MicroVM. A booting MicroVM gives no such sign. + # The check repeats after each such failure, at most once per 2 + # seconds. The worker can suspend after the job's first state check, + # and the last call to the service wins. + if options.suspend_check is not None and ( + status is None or answer.connect_failed + ): + now = time.monotonic() + if ( + last_suspend_check is None + or now - last_suspend_check >= _SUSPEND_CHECK_INTERVAL_SECONDS + ): + last_suspend_check = now + left = deadline() - now + if left <= 0: + raise unavailable() + endpoint = options.suspend_check(left) or endpoint + url = _endpoint_url(endpoint, options.path) def _create_auth_token(options: SendJobOptions) -> Mapping[str, str]: @@ -545,6 +628,12 @@ class _Answer: """Whether the answer has a ``Location`` header.""" followed_redirect: bool = False """Whether the opener followed a redirect.""" + connect_failed: bool = False + """Whether the endpoint answered 502 with ``x-aws-proxy-error: + MICROVM_CONNECT_FAILED``. It answers so for a suspending or suspended + MicroVM.""" + suspend_after: str | None = None + """The value of the ``X-Microvm-Suspend-After`` header of a 2xx answer.""" def _post_within( @@ -593,11 +682,21 @@ def _post( # An opener that follows redirects reports the URL it reached. if response.geturl() != url: return _Answer(response.status, followed_redirect=True) - return _Answer(response.status) + return _Answer( + response.status, + suspend_after=response.headers.get(SUSPEND_AFTER_HEADER), + ) except urllib.error.HTTPError as error: has_location = bool(error.headers and error.headers.get("Location")) + connect_failed = ( + error.code == 502 # noqa: PLR2004 + and bool(error.headers) + and error.headers.get(_PROXY_ERROR_HEADER) == _CONNECT_FAILED + ) error.close() - return _Answer(error.code, has_location=has_location) + return _Answer( + error.code, has_location=has_location, connect_failed=connect_failed + ) except (OSError, http.client.HTTPException, ValueError) as error: # URLError, a refused connection, and a timeout are OSError. An # invalid response from the endpoint raises an HTTPException. diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py index 6cd824c0a..5a3f6a376 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py @@ -6,7 +6,9 @@ from __future__ import annotations +import threading import uuid +from collections import OrderedDict from collections.abc import Callable, Mapping from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, TypeVar @@ -35,6 +37,7 @@ _validate_input, _validate_region, _validate_request, + _validate_suspend_after, _validate_timeout, create_job_callback, job_error, @@ -59,6 +62,16 @@ ``timeout`` is under 10 seconds does not suspend by default. """ +_WARNED_SESSIONS: OrderedDict[str, None] = OrderedDict() +"""The IDs of the sessions that warned about a worker that ignores +``suspend_after``, oldest first.""" + +_MAX_WARNED_SESSIONS = 1_000 +"""How many session IDs :data:`_WARNED_SESSIONS` keeps. A session ID is a +short string, so 1,000 of them use well under 1 MB.""" + +_WARNED_LOCK = threading.Lock() + _MIN_AUTO_SUSPEND_IDLE_SECONDS = 10 """The shortest idle time a session accepts. @@ -108,6 +121,16 @@ def microvm_session( records whether the MicroVM can be suspended. Replay follows that record, not the current config. + A job with ``suspend_after`` asks the worker to suspend the MicroVM after + the job. Each later job of the handle then checks the state before + delivery, as with ``auto_suspend_on_idle``. That record is kept in + memory only. So every job of every session also checks the state after + a sign of a suspended MicroVM: no answer, or an HTTP 502 with + ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``. A boot gives no such + sign, unless the worker takes more than 10 seconds to start listening. + Then the first delivery gets no answer, and the operation makes one + ``GetMicrovm`` call. + A failed :meth:`MicrovmSession.invoke` raises inside the handler. The handler can catch the error and continue with the same MicroVM. An error that leaves the handler ends the session, and the MicroVM is @@ -203,16 +226,50 @@ def __init__( launched: LaunchResult, *, resume_before_delivery: bool, + session_id: str, ) -> None: self._scope = scope self._launched = launched self._resume_before_delivery = resume_before_delivery + # Set by the first suspend_after job of this handle. Later jobs then + # check the state before delivery. The flag lives in memory only. A + # replay of a completed job, map item, or child context does not run + # its invoke() again, so it does not set the flag. The check after a + # sign of a suspended MicroVM covers that case. + self._may_be_suspended = False + self._session_id = session_id @property def microvm_id(self) -> str: """The ID of the session's MicroVM.""" return self._launched.microvm_id + def _resume_mode(self) -> bool: + """Whether a job checks the MicroVM state before delivery. + + The request step reads it at the start of each attempt. So a step + retry, and a ``map`` item that starts after a sibling set the flag, + see the current value. The flag only changes from ``False`` to + ``True``. + """ + return self._resume_before_delivery or self._may_be_suspended + + def _claim_suspend_after_warning(self) -> bool: + """Return ``True`` the first time, so the session warns once about a + worker that does not support ``suspend_after``. + + Each invocation creates a new handle. So the record is kept per + process, by session ID. A later invocation in the same execution + environment then logs no second warning. + """ + with _WARNED_LOCK: + if self._session_id in _WARNED_SESSIONS: + return False + _WARNED_SESSIONS[self._session_id] = None + if len(_WARNED_SESSIONS) > _MAX_WARNED_SESSIONS: + _WARNED_SESSIONS.popitem(last=False) + return True + def invoke( self, context: DurableContext, @@ -237,7 +294,8 @@ def invoke( name: The job's child context name. job_input: The job's input. It must be JSON-serializable, and the same on every replay. - options: The route and the timeouts of this job. + options: The route, the timeouts, and ``suspend_after`` of this + job. Returns: The value that the MicroVM passed to @@ -252,13 +310,16 @@ def invoke( raises it at ``context.run_in_child_context``. ValueError: When the JSON of ``job_input`` is over ``MAX_JOB_INPUT_BYTES`` UTF-8 bytes. - ValidationError: When ``name`` or ``options`` is not valid. + ValidationError: When ``name`` or ``options`` is not valid, for + example a ``suspend_after`` that is not a ``bool``. MicrovmNotRunningError: When the session's MicroVM has ended. The session reports it when it checks the MicroVM state before each job, and for every job when ``CreateMicrovmAuthToken`` - reports the MicroVM removed. + reports the MicroVM removed. A check after a sign of a + suspended MicroVM reports it too. MicrovmDeliveryError: When the delivery fails after all retries, - or the route rejects the job. + or the route rejects the job. A suspended MicroVM whose + ``ResumeMicrovm`` is denied fails the delivery at once. MicrovmJobFailedError: When the MicroVM reports a failure, or sends a result that is not valid JSON. MicrovmTimeoutError: When no result or heartbeat arrives in time. @@ -285,9 +346,15 @@ def invoke( retry_window_seconds = _validate_request(name, options) timeout_seconds = _validate_timeout(name, "timeout", options.timeout) heartbeat_timeout_seconds = _validate_heartbeat(name, options.heartbeat_timeout) + suspend_after = _validate_suspend_after(name, options.suspend_after) _validate_input(name, job_input) launched = self._launched - resume = self._resume_before_delivery + if suspend_after: + # The flag is set before the child context. A replay of a + # completed job returns its result without running the child + # context's function. So the flag is set on every replay of this + # call too. + self._may_be_suspended = True def run(job_context: DurableContext) -> Any: callback = create_job_callback( @@ -302,7 +369,12 @@ def run(job_context: DurableContext) -> Any: heartbeat_timeout_seconds=heartbeat_timeout_seconds, request=options, retry_window_seconds=retry_window_seconds, - resume=resume, + resume=self._resume_mode, + suspend_check=True, + suspend_after=suspend_after, + claim_suspend_after_warning=( + self._claim_suspend_after_warning if suspend_after else None + ), ) try: return callback.result() @@ -401,7 +473,8 @@ def run_session( # 3. So a deploy that changes auto_suspend_on_idle or idle_policy does # not change what the running MicroVM does. # 4. A job that skips the state check of a suspended MicroVM gets a - # 502. The request step retries a 502 until both retry tiers end. + # 502, or no answer. The request step then checks the state, and + # resumes the MicroVM. That costs a failed request per job. # 5. So each job follows the resume mode that the launch recorded, # not the current config. # Every session launch records a boolean. So only `True` turns the @@ -412,6 +485,7 @@ def run_session( self, launched, resume_before_delivery=launched.resume_before_delivery is True, + session_id=session_id, ), child, ) diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py b/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py index 480fbbbb2..5d9ec3699 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py @@ -39,7 +39,9 @@ class FakeMicrovmsClient: ``states`` holds the answers of ``get_microvm``: a state such as ``"SUSPENDED"``, or an exception to raise. After the list ends, the state - is ``RUNNING``. + is ``RUNNING``. With ``suspended_until_resumed``, the state after the + list ends is ``SUSPENDED`` until a ``resume_microvm`` call succeeds, and + ``RUNNING`` after it. ``run_microvm`` and ``get_microvm`` return ``endpoint``. The Nth auth token is ``{"x-test-token": "token-N"}``. @@ -54,9 +56,11 @@ def __init__( resumes: list[Answer] | None = None, states: list[str | BaseException] | None = None, endpoint: str | None = "mvm-1.microvm.example", + suspended_until_resumed: bool = False, ) -> None: self.calls: list[tuple[str, dict[str, Any]]] = [] self.endpoint = endpoint + self._suspended = suspended_until_resumed self._answers: dict[str, list[Answer]] = { "run_microvm": list(launches or []), "create_microvm_auth_token": list(tokens or []), @@ -99,12 +103,17 @@ def terminate_microvm(self, **request: Any) -> dict[str, Any]: def resume_microvm(self, **request: Any) -> dict[str, Any]: self._answer("resume_microvm", request) + with self._lock: + self._suspended = False return {} def get_microvm(self, **request: Any) -> dict[str, Any]: with self._lock: self.calls.append(("get_microvm", request)) - state = self._states.pop(0) if self._states else "RUNNING" + if self._states: + state = self._states.pop(0) + else: + state = "SUSPENDED" if self._suspended else "RUNNING" if isinstance(state, BaseException): raise state response: dict[str, Any] = {"microvmId": "mvm-1", "state": state} @@ -135,19 +144,43 @@ def make_error() -> Callable[..., ClientError]: return client_error +@pytest.fixture +def connect_failed() -> Status: + """The answer of the endpoint of a suspended MicroVM.""" + return CONNECT_FAILED + + +@pytest.fixture +def suspend_after_ack() -> dict[str, str]: + """The header of a worker that supports ``suspendAfter``.""" + return dict(SUSPEND_AFTER_ACK) + + +Status = int | tuple[int, dict[str, str]] | None + +CONNECT_FAILED: Status = (502, {"x-aws-proxy-error": "MICROVM_CONNECT_FAILED"}) +"""The answer of the endpoint of a suspending or suspended MicroVM.""" + +SUSPEND_AFTER_ACK = {"X-Microvm-Suspend-After": "accepted"} +"""The header with which a worker acknowledges ``suspendAfter``.""" + + class FakeEndpoint: """A local HTTP server that stands in for a MicroVM endpoint. It records each request, and answers from ``statuses`` in order. After the list ends, it answers 202. A status of ``None`` closes the connection - without an answer. A redirect status gets ``Location: ``, - ``/moved`` by default. A ``location`` of ``None`` sends no ``Location``. + without an answer. A ``(status, headers)`` pair adds the headers to the + answer. The 202 after the list ends gets ``accept_headers``. A redirect + status gets ``Location: ``, ``/moved`` by default. A + ``location`` of ``None`` sends no ``Location``. """ def __init__( self, - statuses: list[int | None] | None = None, + statuses: list[Status] | None = None, location: str | None = "/moved", + accept_headers: dict[str, str] | None = None, ) -> None: self.requests: list[dict[str, Any]] = [] self._statuses = list(statuses or []) @@ -181,12 +214,19 @@ def do_POST(self) -> None: # noqa: N802 - the http.server name "body": json.loads(body), } ) - status = endpoint._statuses.pop(0) if endpoint._statuses else 202 - if status is None: + answer = ( + endpoint._statuses.pop(0) + if endpoint._statuses + else (202, dict(accept_headers or {})) + ) + if answer is None: self.close_connection = True self.connection.shutdown(socket.SHUT_RDWR) return + status, headers = answer if isinstance(answer, tuple) else (answer, {}) self.send_response(status) + for name, value in headers.items(): + self.send_header(name, value) if 300 <= status < 400 and location is not None: # noqa: PLR2004 self.send_header("location", location) self.send_header("content-length", "0") @@ -213,9 +253,11 @@ def fake_endpoint() -> Iterator[Callable[..., FakeEndpoint]]: endpoints: list[FakeEndpoint] = [] def _create( - statuses: list[int | None] | None = None, location: str | None = "/moved" + statuses: list[Status] | None = None, + location: str | None = "/moved", + accept_headers: dict[str, str] | None = None, ) -> FakeEndpoint: - endpoint = FakeEndpoint(statuses, location) + endpoint = FakeEndpoint(statuses, location, accept_headers) endpoints.append(endpoint) return endpoint diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/job_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/job_test.py index 54853af27..f9198ac94 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/job_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/job_test.py @@ -7,6 +7,7 @@ import json import math +from typing import Any import pytest from aws_durable_execution_sdk_python_microvm_worker import ( @@ -15,6 +16,7 @@ ) from aws_durable_execution_sdk_python_extras.microvm.job import ( + MAX_JOB_INPUT_BYTES, MAX_RUN_HOOK_PAYLOAD_LENGTH, JobRequest, RunHookPayload, @@ -150,3 +152,51 @@ def test_the_worker_reads_the_idle_time_of_a_session() -> None: assert payload == '{"version":1,"region":"us-east-1","autoSuspendIdleSeconds":60}' assert request.payload is not None assert request.payload.auto_suspend_idle_seconds == 60 + + +def test_writes_suspend_after_only_when_it_is_set() -> None: + fields: dict[str, Any] = { + "region": "us-east-1", + "microvm_id": "mvm-1", + "callback_id": "cb-1", + "job_input": 1, + } + + assert JobRequest(**fields, suspend_after=True).to_json() == ( + '{"version":1,"region":"us-east-1","microvmId":"mvm-1",' + '"suspendAfter":true,"callbackId":"cb-1","input":1}' + ) + assert "suspendAfter" not in JobRequest(**fields).to_json() + + +def test_the_worker_reads_suspend_after() -> None: + body = JobRequest( + region="us-east-1", + microvm_id="mvm-1", + callback_id="cb-1", + job_input=None, + suspend_after=True, + ).to_json() + + assert MicrovmJobRequest.from_dict(json.loads(body)).suspend_after is True + + +def test_the_job_fields_fit_in_the_margin_of_the_input_limit() -> None: + # 1. The longest MicroVM ID is 256 characters, each escaped in 6 bytes. + # 2. The longest callback ID is 1,024 characters of 1 byte. + # 3. The heartbeat timeout is at most 8 hours, and suspendAfter is set. + # So the body of the largest input is at most the 6 MiB that the worker + # reads. + body = JobRequest( + region="ap-southeast-5", + microvm_id="\x01" * 256, + callback_id="A" * 1024, + job_input="x" * (MAX_JOB_INPUT_BYTES - 2), + heartbeat_timeout_seconds=28_800, + suspend_after=True, + ).to_json() + + fixed = len(body.encode()) - MAX_JOB_INPUT_BYTES + assert fixed == 112 + len("ap-southeast-5") + 256 * 6 + 1024 + len("28800") + assert fixed < 3 * 1024 + assert len(body.encode()) <= 6 * 1024 * 1024 diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py index 3f2eed380..e6d979580 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py @@ -11,14 +11,18 @@ from typing import Any import pytest +from botocore.exceptions import ClientError, EndpointConnectionError from aws_durable_execution_sdk_python_extras.microvm import ( MicrovmEndpointUnavailableError, ) from aws_durable_execution_sdk_python_extras.microvm.lifecycle import ( EnsureRunningOptions, + MicrovmResumeDeniedError, + MicrovmStateCheckError, MicrovmStateError, ensure_running, + resume_if_suspended, ) from aws_durable_execution_sdk_python_extras.microvm.request import ( INVOCATION_RESERVE_SECONDS, @@ -78,13 +82,71 @@ def test_waits_for_a_suspend_in_progress_then_resumes_once(fake_client: Any) -> def test_keeps_polling_when_resume_reports_a_transition_in_progress( fake_client: Any, make_error: Any ) -> None: + # An auto-resume is in progress. Each SUSPENDED answer sends + # ResumeMicrovm again, and the service answers ConflictException. A + # PENDING answer after a resume is a step of the resume, not a boot. client = fake_client( - states=["SUSPENDED", "SUSPENDED"], - resumes=[make_error("ConflictException", 409)], + states=["SUSPENDED", "SUSPENDED", "PENDING", "RUNNING"], + resumes=[make_error("ConflictException", 409)] * 2, ) assert run_with(client) == "mvm-1.microvm.example" - assert len(client.requests("resume_microvm")) == 1 + assert calls(client) == [ + "get_microvm", + "resume_microvm", + "get_microvm", + "resume_microvm", + "get_microvm", + "get_microvm", + ] + + +def test_resumes_again_when_a_resume_conflicts_with_the_suspend( + fake_client: Any, make_error: Any +) -> None: + # The service rejects the first ResumeMicrovm with ConflictException. + # The MicroVM stays SUSPENDED until a resume succeeds. So the next + # SUSPENDED answer sends ResumeMicrovm again. + client = fake_client( + states=["SUSPENDING", "SUSPENDED"], + resumes=[make_error("ConflictException", 409)], + suspended_until_resumed=True, + ) + sleeps: list[float] = [] + + assert run_with(client, max_wait_seconds=3, sleeps=sleeps) == ( + "mvm-1.microvm.example" + ) + assert calls(client) == [ + "get_microvm", + "get_microvm", + "resume_microvm", + "get_microvm", + "resume_microvm", + "get_microvm", + ] + assert sleeps == [0.25, 0.5, 1.0] + + +def test_gives_up_after_the_wait_limit_when_each_resume_conflicts( + fake_client: Any, make_error: Any +) -> None: + client = fake_client( + resumes=[make_error("ConflictException", 409)] * 10, + suspended_until_resumed=True, + ) + + with pytest.raises(MicrovmEndpointUnavailableError, match="Last state: SUSPENDED"): + run_with(client, max_wait_seconds=0.5) + + # One ResumeMicrovm call per GetMicrovm call. The wait of 0.5 seconds + # after the second GetMicrovm call does not fit in the limit. + assert calls(client) == [ + "get_microvm", + "resume_microvm", + "get_microvm", + "resume_microvm", + ] def test_raises_a_resume_error_that_is_not_a_conflict( @@ -222,3 +284,261 @@ def test_raises_the_state_error_when_the_microvm_is_removed_before_the_resume( with pytest.raises(MicrovmStateError, match="no longer exists"): run_with(client) + + +# --------------------------------------------------------------------------- +# resume_if_suspended: the check after a sign of a suspended MicroVM +# --------------------------------------------------------------------------- + + +def check_with( + client: Any, + *, + max_wait_seconds: float = 30, + remaining: Callable[[], float | None] = lambda: None, + sleeps: list[float] | None = None, +) -> str | None: + recorded = [] if sleeps is None else sleeps + return resume_if_suspended( + EnsureRunningOptions( + client=client, + microvm_id="mvm-1", + max_wait_seconds=max_wait_seconds, + remaining_seconds=remaining, + logger=logging.getLogger("test"), + sleep=recorded.append, + ) + ) + + +def service_error(operation: str, code: str, status: int = 400) -> ClientError: + response: Any = { + "Error": {"Code": code, "Message": f"{code} message"}, + "ResponseMetadata": {"HTTPStatusCode": status, "RequestId": "req-1"}, + } + return ClientError(response, operation) + + +@pytest.mark.parametrize("state", ["RUNNING", "PENDING"]) +def test_check_returns_at_once_for_a_microvm_that_runs_or_boots( + fake_client: Any, state: str +) -> None: + client = fake_client(states=[state], endpoint="mvm-1.new.example") + + assert check_with(client) == "mvm-1.new.example" + assert calls(client) == ["get_microvm"] + + +@pytest.mark.parametrize("state", ["SUSPENDING", "SUSPENDED"]) +def test_check_resumes_a_suspending_or_suspended_microvm_once( + fake_client: Any, state: str +) -> None: + # The service honors a resume during SUSPENDING. So the check does not + # wait for SUSPENDED. + client = fake_client(states=[state, "SUSPENDED", "PENDING", "RUNNING"]) + sleeps: list[float] = [] + + assert check_with(client, sleeps=sleeps) == "mvm-1.microvm.example" + assert calls(client) == [ + "get_microvm", + "resume_microvm", + "get_microvm", + "get_microvm", + "get_microvm", + ] + assert sleeps == [0.25, 0.5] + + +def test_check_keeps_polling_after_a_resume_conflict( + fake_client: Any, make_error: Any +) -> None: + client = fake_client( + states=["SUSPENDED", "RUNNING"], + resumes=[make_error("ConflictException", 409)], + ) + + assert check_with(client) == "mvm-1.microvm.example" + assert calls(client) == ["get_microvm", "resume_microvm", "get_microvm"] + + +def test_check_resumes_again_after_a_conflict_during_the_suspend( + fake_client: Any, make_error: Any +) -> None: + # The service rejects the resume during SUSPENDING with + # ConflictException. The suspend then ends, and the MicroVM stays + # SUSPENDED until a resume succeeds. So the first SUSPENDED answer sends + # ResumeMicrovm again. + client = fake_client( + states=["SUSPENDING", "SUSPENDING"], + resumes=[make_error("ConflictException", 409)], + suspended_until_resumed=True, + ) + sleeps: list[float] = [] + + assert check_with(client, max_wait_seconds=3, sleeps=sleeps) == ( + "mvm-1.microvm.example" + ) + assert calls(client) == [ + "get_microvm", + "resume_microvm", + "get_microvm", + "get_microvm", + "resume_microvm", + "get_microvm", + ] + assert sleeps == [0.25, 0.5] + + +def test_check_gives_up_after_the_wait_limit_when_each_resume_conflicts( + fake_client: Any, make_error: Any +) -> None: + client = fake_client( + states=["SUSPENDING"], + resumes=[make_error("ConflictException", 409)] * 10, + suspended_until_resumed=True, + ) + + with pytest.raises(MicrovmEndpointUnavailableError, match="Last state: SUSPENDED"): + check_with(client, max_wait_seconds=0.5) + + # One ResumeMicrovm call per GetMicrovm call. The wait of 0.5 seconds + # after the third GetMicrovm call does not fit in the limit. + assert calls(client) == [ + "get_microvm", + "resume_microvm", + "get_microvm", + "resume_microvm", + "get_microvm", + "resume_microvm", + ] + + +@pytest.mark.parametrize("state", ["TERMINATING", "TERMINATED"]) +def test_check_raises_the_state_error_for_a_microvm_that_ends( + fake_client: Any, state: str +) -> None: + client = fake_client(states=[state]) + + with pytest.raises(MicrovmStateError, match=f"MicroVM mvm-1 is {state}") as raised: + check_with(client) + + assert raised.value.state == state + + +@pytest.mark.parametrize("operation", ["GetMicrovm", "ResumeMicrovm"]) +def test_check_raises_the_state_error_for_a_removed_microvm( + fake_client: Any, operation: str +) -> None: + error = service_error(operation, "ResourceNotFoundException", 404) + client = ( + fake_client(states=[error]) + if operation == "GetMicrovm" + else fake_client(states=["SUSPENDED"], resumes=[error]) + ) + + with pytest.raises(MicrovmStateError) as raised: + check_with(client) + + assert str(raised.value).endswith(f"{operation} request ID: req-1.") + assert raised.value.__cause__ is error + + +@pytest.mark.parametrize( + ("code", "denied"), + [("AccessDeniedException", True), ("ThrottlingException", False)], +) +def test_check_reports_a_failed_get_microvm_without_failing_the_job( + fake_client: Any, code: str, denied: bool +) -> None: + error = service_error("GetMicrovm", code) + client = fake_client(states=[error]) + + with pytest.raises(MicrovmStateCheckError) as raised: + check_with(client) + + assert raised.value.denied is denied + assert f"({code})" in str(raised.value) + assert str(raised.value).endswith(f"{code} message. GetMicrovm request ID: req-1.") + assert raised.value.__cause__ is error + + +def test_check_reports_a_get_microvm_connection_failure(fake_client: Any) -> None: + error = EndpointConnectionError(endpoint_url="https://lambda.example") + client = fake_client(states=[error]) + + with pytest.raises(MicrovmStateCheckError) as raised: + check_with(client) + + assert raised.value.denied is False + assert raised.value.__cause__ is error + + +@pytest.mark.parametrize("state", ["SUSPENDING", "SUSPENDED"]) +def test_check_fails_at_once_when_resume_microvm_is_denied( + fake_client: Any, state: str +) -> None: + error = service_error("ResumeMicrovm", "AccessDeniedException", 403) + client = fake_client(states=[state], resumes=[error]) + + with pytest.raises(MicrovmResumeDeniedError) as raised: + check_with(client) + + assert str(raised.value) == ( + f"MicroVM mvm-1 is {state}. ResumeMicrovm was denied, so the operation " + "cannot resume it. The function's execution role needs " + "lambda:ResumeMicrovm. ResumeMicrovm request ID: req-1." + ) + assert raised.value.__cause__ is error + assert calls(client) == ["get_microvm", "resume_microvm"] + + +def test_check_names_no_request_id_for_an_error_of_another_call( + fake_client: Any, +) -> None: + # An STS error inside the call carries the request ID of STS. + client = fake_client( + states=["SUSPENDED"], + resumes=[service_error("AssumeRole", "AccessDeniedException", 403)], + ) + + with pytest.raises(MicrovmResumeDeniedError) as raised: + check_with(client) + + assert str(raised.value).endswith("needs lambda:ResumeMicrovm.") + + +def test_check_raises_another_resume_error_with_its_request_id( + fake_client: Any, +) -> None: + client = fake_client( + states=["SUSPENDED"], + resumes=[service_error("ResumeMicrovm", "ThrottlingException", 429)], + ) + + with pytest.raises(ClientError) as raised: + check_with(client) + + assert str(raised.value).endswith("ResumeMicrovm request ID: req-1.") + + +def test_check_gives_up_when_the_resume_does_not_end_in_time( + fake_client: Any, +) -> None: + client = fake_client(states=["SUSPENDED"] * 10) + + with pytest.raises(MicrovmEndpointUnavailableError, match="Last state: SUSPENDED"): + check_with(client, max_wait_seconds=0.1) + + +def test_check_stops_a_stalled_get_microvm_at_the_reserve(fake_client: Any) -> None: + release = threading.Event() + client = fake_client() + client.get_microvm = lambda **_: release.wait(30) + + try: + with pytest.raises( + MicrovmEndpointUnavailableError, match="Last state: unknown" + ): + check_with(client, remaining=lambda: INVOCATION_RESERVE_SECONDS + 1.2) + finally: + release.set() diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py index 6efeb8e62..fac5b92df 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py @@ -15,6 +15,7 @@ import pytest from aws_durable_execution_sdk_python_microvm_worker import MICROVM_JOB_PATH +from aws_durable_execution_sdk_python_microvm_worker import worker from aws_durable_execution_sdk_python_extras.microvm import ( DEFAULT_MICROVM_JOB_PATH, @@ -57,6 +58,11 @@ def test_matches_the_job_path_of_the_worker() -> None: assert DEFAULT_MICROVM_JOB_PATH == MICROVM_JOB_PATH +def test_matches_the_suspend_after_header_of_the_worker() -> None: + assert request.SUSPEND_AFTER_HEADER == worker.SUSPEND_AFTER_HEADER + assert request.SUSPEND_AFTER_ACCEPTED == "accepted" + + def test_posts_the_job_with_an_auth_token(fake_client: Any, fake_endpoint: Any) -> None: client = fake_client() endpoint = fake_endpoint() @@ -637,3 +643,221 @@ def test_sends_no_auth_token_on_a_redirect_that_the_callers_opener_follows( assert followed["method"] == "GET" assert "x-test-token" not in followed["headers"] assert "content-type" not in followed["headers"] + + +# --------------------------------------------------------------------------- +# The check for a suspended MicroVM, and the suspendAfter acknowledgement +# --------------------------------------------------------------------------- + + +class FakeClock: + """A monotonic clock that advances only when the tier sleeps.""" + + def __init__(self) -> None: + self.now = 1000.0 + self.sleeps: list[float] = [] + + def monotonic(self) -> float: + return self.now + + def sleep(self, seconds: float) -> None: + self.sleeps.append(seconds) + self.now += seconds + + +@pytest.fixture +def clock(monkeypatch: pytest.MonkeyPatch) -> FakeClock: + fake = FakeClock() + monkeypatch.setattr(request, "time", fake) + return fake + + +def test_checks_for_a_suspend_after_a_connect_failure_or_no_answer( + fake_client: Any, fake_endpoint: Any, connect_failed: Any, clock: FakeClock +) -> None: + # 1. A suspended MicroVM's endpoint answers 502 with MICROVM_CONNECT_FAILED. + # 2. A request that the freeze caught gets no answer. + # 3. The clock passes 2 seconds between the two failures. + endpoint = fake_endpoint([connect_failed, 503, 503, 503, None]) + windows: list[float] = [] + + def check(left: float) -> str | None: + windows.append(left) + return None + + status = send_job( + options( + fake_client(), + endpoint.url, + [], + sleep=clock.sleep, + suspend_check=check, + ) + ) + + assert status == 202 + assert clock.sleeps == [0.25, 0.5, 1.0, 2.0, 4.0] + assert windows == [60 - 0.25, 60 - 7.75] + + +@pytest.mark.parametrize( + "answer", + [502, 503, 504, (502, {"x-aws-proxy-error": "OTHER"}), 429, 500], + ids=["502", "503", "504", "502-other-proxy-error", "429", "500"], +) +def test_makes_no_suspend_check_for_another_failure( + fake_client: Any, fake_endpoint: Any, answer: Any +) -> None: + # A booting MicroVM never gets MICROVM_CONNECT_FAILED. So these answers + # make no extra call. + endpoint = fake_endpoint([answer]) + checks: list[float] = [] + + status = send_job( + options(fake_client(), endpoint.url, [], suspend_check=checks.append) + ) + + assert status == 202 + assert checks == [] + + +def test_checks_for_a_suspend_at_most_once_per_2_seconds( + fake_client: Any, fake_endpoint: Any, connect_failed: Any, clock: FakeClock +) -> None: + endpoint = fake_endpoint([connect_failed] * 5) + checked_at: list[float] = [] + + def check(left: float) -> str | None: + checked_at.append(clock.now - 1000) + return None + + send_job( + options(fake_client(), endpoint.url, [], sleep=clock.sleep, suspend_check=check) + ) + + # The failures end at 0.25, 0.75, 1.75, 3.75, and 7.75 seconds. A check + # that would come less than 2 seconds after the last one is skipped. + assert checked_at == [0.25, 3.75, 7.75] + assert len(endpoint.requests) == 6 + + +def test_sends_the_next_request_to_the_endpoint_that_the_suspend_check_returns( + fake_client: Any, fake_endpoint: Any, connect_failed: Any +) -> None: + before = fake_endpoint([connect_failed]) + after = fake_endpoint() + + status = send_job( + options(fake_client(), before.url, [], suspend_check=lambda left: after.url) + ) + + assert status == 202 + assert len(before.requests) == 1 + assert len(after.requests) == 1 + + +def test_ends_the_tier_when_the_suspend_check_has_no_time_left( + fake_client: Any, fake_endpoint: Any, connect_failed: Any +) -> None: + endpoint = fake_endpoint([connect_failed]) + checks: list[float] = [] + + with pytest.raises(MicrovmEndpointUnavailableError, match="Last failure: HTTP 502"): + send_job( + options( + fake_client(), + endpoint.url, + [], + retry_window_seconds=1, + window_started_at=time.monotonic() - 0.2, + sleep=lambda seconds: time.sleep(0.9), + suspend_check=checks.append, + ) + ) + + assert checks == [] + + +def test_adds_the_note_to_the_final_error( + fake_client: Any, fake_endpoint: Any, connect_failed: Any +) -> None: + endpoint = fake_endpoint([connect_failed] * 10) + + with pytest.raises(MicrovmEndpointUnavailableError) as raised: + send_job( + options( + fake_client(), + endpoint.url, + [], + retry_window_seconds=1, + unavailable_note=lambda: "GetMicrovm was denied.", + ) + ) + + assert str(raised.value).endswith("Last failure: HTTP 502. GetMicrovm was denied.") + + +def test_adds_no_note_when_the_note_is_none( + fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([503] * 10) + + with pytest.raises(MicrovmEndpointUnavailableError) as raised: + send_job( + options( + fake_client(), + endpoint.url, + [], + retry_window_seconds=1, + unavailable_note=lambda: None, + ) + ) + + assert str(raised.value).endswith("Last failure: HTTP 503.") + + +@pytest.mark.parametrize( + ("headers", "ignored"), + [ + ({"X-Microvm-Suspend-After": "accepted"}, 0), + ({}, 1), + ({"X-Microvm-Suspend-After": "yes"}, 1), + ], + ids=["acknowledged", "no-header", "other-value"], +) +def test_reports_a_suspend_after_job_that_the_worker_did_not_acknowledge( + fake_client: Any, fake_endpoint: Any, headers: dict[str, str], ignored: int +) -> None: + endpoint = fake_endpoint(accept_headers=headers) + calls: list[None] = [] + + status = send_job( + options( + fake_client(), + endpoint.url, + [], + on_suspend_after_ignored=lambda: calls.append(None), + ) + ) + + assert status == 202 + assert len(calls) == ignored + + +def test_reads_the_acknowledgement_of_the_answer_that_accepted_the_job( + fake_client: Any, fake_endpoint: Any, suspend_after_ack: dict[str, str] +) -> None: + # A refused answer that carries the header does not count. + endpoint = fake_endpoint([(503, suspend_after_ack)]) + calls: list[None] = [] + + send_job( + options( + fake_client(), + endpoint.url, + [], + on_suspend_after_ignored=lambda: calls.append(None), + ) + ) + + assert calls == [None] diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py index d09f1fa84..fa5500eb2 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py @@ -8,13 +8,15 @@ import datetime import hashlib import json +import logging +from collections import OrderedDict from collections.abc import Callable, Iterator from types import SimpleNamespace from typing import Any import boto3 import pytest -from aws_durable_execution_sdk_python.config import Duration, StepConfig +from aws_durable_execution_sdk_python.config import Duration, MapConfig, StepConfig from aws_durable_execution_sdk_python.context import Callback, DurableContext from aws_durable_execution_sdk_python.exceptions import ( ChildContextError, @@ -48,6 +50,7 @@ microvm_session, operation, ) +from aws_durable_execution_sdk_python_extras.microvm import session as session_module WAIT = 30 @@ -373,6 +376,20 @@ def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: "TypeError", "must be JSON-serializable", ), + ( + lambda vm, ctx: vm.invoke( + ctx, "build", job_input=1, options=job(suspend_after=1) + ), + "ValidationError", + "suspend_after must be a bool", + ), + ( + lambda vm, ctx: vm.invoke( + ctx, "build", job_input=1, options=job(suspend_after="yes") + ), + "ValidationError", + "suspend_after must be a bool", + ), ], ) def test_rejects_an_invalid_job_inside_the_handler_where_it_can_be_caught( @@ -698,6 +715,49 @@ def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: assert len(client.requests("resume_microvm")) == 1 +@pytest.mark.parametrize( + ("statuses", "states"), + [([], ["SUSPENDED"]), ([502], ["RUNNING", "SUSPENDED"])], + ids=["before-delivery", "after-a-gateway-error"], +) +def test_resumes_again_when_a_resume_conflicts( + run: Any, + fake_client: Any, + fake_endpoint: Any, + make_error: Any, + statuses: list[int], + states: list[str], +) -> None: + # The service rejects the first ResumeMicrovm with ConflictException. + # The MicroVM stays SUSPENDED until a resume succeeds. So the next + # SUSPENDED answer sends ResumeMicrovm again. + endpoint = fake_endpoint(statuses) + client = fake_client( + endpoint=endpoint.url, + states=states, + resumes=[make_error("ConflictException", 409)], + suspended_until_resumed=True, + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke( + ctx, + "build", + job_input=1, + options=job(retry_window=Duration.from_seconds(5)), + ) + + result = run( + session_handler(session_config(client, auto_suspend_on_idle=True), body), + None, + [("build", succeed_with("done"))], + ) + + assert outcome(result) == {"result": "done"} + assert len(client.requests("resume_microvm")) == 2 + assert len(endpoint.requests) == len(statuses) + 1 + + def test_does_not_check_the_state_again_after_a_rejection( run: Any, fake_client: Any, fake_endpoint: Any ) -> None: @@ -1024,16 +1084,17 @@ def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: assert len(client.requests("terminate_microvm")) == 1 -def test_resumes_once_when_two_items_find_the_microvm_suspended( +def test_runs_both_items_when_a_resume_finds_another_in_progress( run: Any, fake_client: Any, fake_endpoint: Any, make_error: Any ) -> None: # Both items find the MicroVM SUSPENDED and call ResumeMicrovm. The - # second call finds the resume in progress. + # service accepts the first call. It answers ConflictException to each + # later call while the resume runs. endpoint = fake_endpoint() client = fake_client( endpoint=endpoint.url, states=["SUSPENDED"] * 4, - resumes=[None, make_error("ConflictException", 409)], + resumes=[None] + [make_error("ConflictException", 409)] * 4, ) def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: @@ -1052,7 +1113,10 @@ def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: ) assert outcome(result) == {"result": ["A", "B"]} - assert len(client.requests("resume_microvm")) == 2 + # The item whose call conflicts calls ResumeMicrovm again at each later + # SUSPENDED answer. The order of the threads sets how many answers it + # gets. The other item calls once. + assert 2 <= len(client.requests("resume_microvm")) <= 4 assert len(endpoint.requests) == 2 @@ -1461,3 +1525,443 @@ def test_makes_no_state_check_for_a_launch_checkpoint_without_a_bool_resume_mode **fields, } assert client.requests("get_microvm") == [] + + +# --------------------------------------------------------------------------- +# suspend_after, and the check after a sign of a suspended MicroVM +# --------------------------------------------------------------------------- + +IGNORED_WARNING = ( + "the worker in MicroVM mvm-1 does not support suspendAfter. The MicroVM " + "keeps running after the job." +) +STATE_CHECK_WARNING = ( + "GetMicrovm failed, so the operation could not check whether the MicroVM " + "is suspended. The request step keeps sending the job." +) + + +def warnings(caplog: pytest.LogCaptureFixture, message: str) -> int: + return sum( + 1 + for record in caplog.records + if record.levelname == "WARNING" and record.getMessage() == message + ) + + +def suspend_then_wait_then_job( + vm: MicrovmSession, ctx: DurableContext, event: Any +) -> Any: + """A suspend_after job, a durable wait, and one more job. + + The wait ends the invocation. So the second job runs in a later + invocation, in which the first job's child context is complete.""" + build = vm.invoke(ctx, "build", job_input=1, options=job(suspend_after=True)) + ctx.wait(Duration.from_seconds(1), name="pause") + return [build, vm.invoke(ctx, "test", job_input=2, options=job())] + + +def test_asks_the_worker_to_suspend_after_the_job_only_when_it_is_set( + run: Any, fake_client: Any, fake_endpoint: Any, suspend_after_ack: Any +) -> None: + endpoint = fake_endpoint(accept_headers=suspend_after_ack) + client = fake_client(endpoint=endpoint.url) + + result = run( + session_handler(session_config(client), suspend_then_wait_then_job), + None, + [("build", succeed_with(1)), ("test", succeed_with(2))], + ) + + assert outcome(result) == {"result": [1, 2]} + first, second = (request["body"] for request in endpoint.requests) + assert first["suspendAfter"] is True + assert "suspendAfter" not in second + # The run hook payload never carries it. + [launch] = client.requests("run_microvm") + assert json.loads(launch["runHookPayload"]) == {"version": 1, "region": "us-west-2"} + # suspend_after is not a session setting. So the launch records no + # resume mode for it. + assert launch_checkpoint(result)["resumeBeforeDelivery"] is False + + +def test_checks_the_state_before_the_jobs_after_a_suspend_after_job( + run: Any, fake_client: Any, fake_endpoint: Any, suspend_after_ack: Any +) -> None: + endpoint = fake_endpoint(accept_headers=suspend_after_ack) + client = fake_client( + endpoint=endpoint.url, states=["RUNNING", "SUSPENDED", "RUNNING"] + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + build = vm.invoke(ctx, "build", job_input=1, options=job(suspend_after=True)) + return [build, vm.invoke(ctx, "test", job_input=2, options=job())] + + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with(1)), ("test", succeed_with(2))], + ) + + assert outcome(result) == {"result": [1, 2]} + # The flag is set before the first job's child context. So the first job + # checks the state too: RUNNING. The second job finds the MicroVM + # SUSPENDED, resumes it, and sends the job when it runs. + assert len(client.requests("get_microvm")) == 3 + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert len(endpoint.requests) == 2 + + +def test_sets_the_flag_again_when_a_completed_suspend_after_job_replays( + run: Any, fake_client: Any, fake_endpoint: Any, suspend_after_ack: Any +) -> None: + # The second job runs in a later invocation. The first job's child + # context is complete then, so its function does not run again. invoke() + # sets the flag before the child context, so the second job still + # checks the state before delivery. + endpoint = fake_endpoint(accept_headers=suspend_after_ack) + client = fake_client( + endpoint=endpoint.url, states=["RUNNING", "SUSPENDED", "RUNNING"] + ) + handler, invocations = handler_across_deploy( + client, {}, {}, suspend_then_wait_then_job + ) + sent_in = record_invocation_of_each_request(endpoint, invocations) + + result = run(handler, None, [("build", succeed_with(1)), ("test", succeed_with(2))]) + + assert outcome(result) == {"result": [1, 2]} + assert sent_in[1] > sent_in[0] + 1 + assert len(client.requests("get_microvm")) == 3 + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert len(endpoint.requests) == 2 + + +def test_resumes_through_the_fallback_after_a_map_item_lost_the_flag( + run: Any, + fake_client: Any, + fake_endpoint: Any, + suspend_after_ack: Any, + connect_failed: Any, +) -> None: + # 1. Item a asks for a suspend after its job. + # 2. The job after the map runs in a later invocation, after a wait. + # 3. The map is complete then, so its items do not run again, and the + # flag is lost. + # 4. So the job goes to the endpoint of a suspended MicroVM, which + # answers 502 MICROVM_CONNECT_FAILED. The step checks the state, + # resumes the MicroVM, and sends the job again. + endpoint = fake_endpoint([(202, suspend_after_ack), 202, connect_failed]) + client = fake_client( + endpoint=endpoint.url, + states=["RUNNING", "RUNNING", "SUSPENDED", "RUNNING"], + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + shards: list[Any] = ctx.map( + ["a", "b"], + lambda item_context, item, index, items: vm.invoke( + item_context, + f"shard-{item}", + job_input=item, + options=job(suspend_after=item == "a"), + ), + name="shards", + config=MapConfig(max_concurrency=1), + ).get_results() + ctx.wait(Duration.from_seconds(1), name="pause") + return [shards, vm.invoke(ctx, "test", job_input=2, options=job())] + + result = run( + session_handler(session_config(client), body), + None, + [ + ("shard-a", succeed_with("A")), + ("shard-b", succeed_with("B")), + ("test", succeed_with(2)), + ], + ) + + assert outcome(result) == {"result": [["A", "B"], 2]} + # Item a and item b check the state before delivery. The job after the + # map checks it after the 502, and once more after the resume. + assert len(client.requests("get_microvm")) == 4 + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert [request["body"]["input"] for request in endpoint.requests] == [ + "a", + "b", + 2, + 2, + ] + + +def test_the_fallback_resumes_again_after_a_conflict_during_the_suspend( + run: Any, + fake_client: Any, + fake_endpoint: Any, + make_error: Any, + connect_failed: Any, +) -> None: + # The check finds the MicroVM SUSPENDING. The service rejects the resume + # with ConflictException. The suspend then ends, and the MicroVM stays + # SUSPENDED until a resume succeeds. So the first SUSPENDED answer sends + # ResumeMicrovm again. + endpoint = fake_endpoint([connect_failed]) + client = fake_client( + endpoint=endpoint.url, + states=["SUSPENDING"], + resumes=[make_error("ConflictException", 409)], + suspended_until_resumed=True, + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke( + ctx, + "build", + job_input=1, + options=job(retry_window=Duration.from_seconds(5)), + ) + + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with("done"))], + ) + + assert outcome(result) == {"result": "done"} + assert len(client.requests("resume_microvm")) == 2 + assert len(endpoint.requests) == 2 + + +def test_a_session_that_never_suspends_checks_the_state_after_no_answer( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([None]) + client = fake_client(endpoint=endpoint.url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke(ctx, "build", job_input=1, options=job()) + + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with("done"))], + ) + + assert outcome(result) == {"result": "done"} + assert len(client.requests("get_microvm")) == 1 + assert client.requests("resume_microvm") == [] + assert len(endpoint.requests) == 2 + + +def test_a_session_that_never_suspends_makes_no_state_check_for_a_boot_answer( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + # A plain 502, 503, or 504 is what the endpoint can answer while the + # MicroVM starts. It is no sign of a suspend. + endpoint = fake_endpoint([502, 503, 504]) + client = fake_client(endpoint=endpoint.url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke(ctx, "build", job_input=1, options=job()) + + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with("done"))], + ) + + assert outcome(result) == {"result": "done"} + assert client.requests("get_microvm") == [] + assert len(endpoint.requests) == 4 + + +def test_keeps_sending_after_a_denied_state_check_and_names_it_at_the_end( + run: Any, + fake_client: Any, + fake_endpoint: Any, + make_error: Any, + connect_failed: Any, + caplog: pytest.LogCaptureFixture, +) -> None: + # The failures end at 0.25, 0.75, 1.75, and 3.75 seconds. The checks run + # at 0.25 and 3.75 seconds, and both are denied. The next delay of 4 + # seconds does not fit in the window. + endpoint = fake_endpoint([connect_failed] * 10) + client = fake_client( + endpoint=endpoint.url, + states=[make_error("AccessDeniedException", 403)] * 10, + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke( + ctx, + "build", + job_input=1, + options=job(retry_window=Duration.from_seconds(5)), + ) + + with caplog.at_level(logging.WARNING): + result = run(session_handler(session_config(client), body)) + + error = outcome(result) + assert error["error"] == "MicrovmDeliveryError" + assert error["message"].endswith( + "Last failure: HTTP 502. GetMicrovm was denied, so the operation could " + "not check whether the MicroVM is suspended." + ) + assert len(client.requests("get_microvm")) == 2 + assert len(endpoint.requests) == 5 + assert warnings(caplog, STATE_CHECK_WARNING) == 1 + + +def test_fails_at_once_when_resume_is_denied_for_a_suspended_microvm( + run: Any, + fake_client: Any, + fake_endpoint: Any, + make_error: Any, + connect_failed: Any, +) -> None: + endpoint = fake_endpoint([connect_failed]) + client = fake_client( + endpoint=endpoint.url, + states=["SUSPENDED"], + resumes=[make_error("AccessDeniedException", 403)], + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke(ctx, "build", job_input=1, options=job()) + + # The default strategy decides on the error of the step. + result = run(session_handler(session_config(client, retry_strategy=None), body)) + + assert outcome(result) == { + "error": "MicrovmDeliveryError", + "message": ( + 'MicroVM "build": the job could not be delivered: MicroVM mvm-1 is ' + "SUSPENDED. ResumeMicrovm was denied, so the operation cannot resume " + "it. The function's execution role needs lambda:ResumeMicrovm." + ), + } + assert len(client.requests("resume_microvm")) == 1 + assert len(endpoint.requests) == 1 + + +def test_fails_with_not_running_when_the_check_finds_the_microvm_terminated( + run: Any, fake_client: Any, fake_endpoint: Any, connect_failed: Any +) -> None: + endpoint = fake_endpoint([connect_failed]) + client = fake_client(endpoint=endpoint.url, states=["TERMINATED"]) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke(ctx, "build", job_input=1, options=job()) + + result = run(session_handler(session_config(client), body)) + + error = outcome(result) + assert error["error"] == "MicrovmNotRunningError" + assert "MicroVM mvm-1 is TERMINATED." in error["message"] + assert len(endpoint.requests) == 1 + + +def test_a_flagged_session_keeps_the_strict_state_check( + run: Any, fake_client: Any, fake_endpoint: Any, make_error: Any +) -> None: + # After a suspend_after job, the state check before delivery fails the + # job for a denied GetMicrovm, as in a session with auto-suspend. + endpoint = fake_endpoint() + client = fake_client( + endpoint=endpoint.url, + states=["RUNNING", make_error("AccessDeniedException", 403)], + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + build = vm.invoke(ctx, "build", job_input=1, options=job(suspend_after=True)) + return [build, vm.invoke(ctx, "test", job_input=2, options=job())] + + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with(1))], + ) + + error = outcome(result) + assert error["error"] == "MicrovmDeliveryError" + assert "AccessDeniedException message" in error["message"] + assert len(endpoint.requests) == 1 + + +@pytest.mark.parametrize( + ("acknowledged", "expected"), [(False, 1), (True, 0)], ids=["old", "new"] +) +def test_warns_once_per_session_when_the_worker_ignores_suspend_after( + run: Any, + fake_client: Any, + fake_endpoint: Any, + suspend_after_ack: Any, + caplog: pytest.LogCaptureFixture, + acknowledged: bool, # noqa: FBT001 - a pytest parameter + expected: int, +) -> None: + endpoint = fake_endpoint(accept_headers=suspend_after_ack if acknowledged else {}) + client = fake_client(endpoint=endpoint.url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + options = job(suspend_after=True) + return [ + vm.invoke(ctx, "build", job_input=1, options=options), + vm.invoke(ctx, "test", job_input=2, options=options), + ] + + with caplog.at_level(logging.WARNING): + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with(1)), ("test", succeed_with(2))], + ) + + assert outcome(result) == {"result": [1, 2]} + assert len(endpoint.requests) == 2 + assert warnings(caplog, IGNORED_WARNING) == expected + + +def test_reads_no_acknowledgement_for_a_job_without_suspend_after( + run: Any, fake_client: Any, fake_endpoint: Any, caplog: pytest.LogCaptureFixture +) -> None: + client = fake_client(endpoint=fake_endpoint().url) + + with caplog.at_level(logging.WARNING): + result = run( + session_handler(session_config(client), two_jobs), + None, + [("build", succeed_with(1)), ("test", succeed_with(2))], + ) + + assert outcome(result)["result"]["test"] == 2 + assert warnings(caplog, IGNORED_WARNING) == 0 + + +def test_remembers_the_warning_of_the_latest_sessions_only( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(session_module, "_WARNED_SESSIONS", OrderedDict()) + monkeypatch.setattr(session_module, "_MAX_WARNED_SESSIONS", 2) + a, b, c = ( + MicrovmSession( + None, # type: ignore[arg-type] + None, # type: ignore[arg-type] + resume_before_delivery=False, + session_id=session_id, + ) + for session_id in ["a", "b", "c"] + ) + + assert [handle._claim_suspend_after_warning() for handle in (a, b, c)] == [ # noqa: SLF001 + True, + True, + True, + ] + # Each session warns once. + assert c._claim_suspend_after_warning() is False # noqa: SLF001 + # The oldest session was forgotten. + assert a._claim_suspend_after_warning() is True # noqa: SLF001 diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/validation_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/validation_test.py index 8623f2a00..9cd63bc5d 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/validation_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/validation_test.py @@ -15,6 +15,7 @@ from aws_durable_execution_sdk_python_extras.microvm import ( MAX_JOB_INPUT_BYTES, MicrovmConfig, + MicrovmInvokeOptions, MicrovmRequestConfig, microvm, ) @@ -328,3 +329,36 @@ def test_accepts_dots_that_are_not_a_dot_segment(path: str) -> None: operation._validate_request("build", MicrovmRequestConfig(path=path)) # noqa: SLF001 == 60 ) + + +@pytest.mark.parametrize( + ("suspend_after", "message"), + [ + (True, "suspend_after applies only to a session job"), + (1, "suspend_after must be a bool"), + (0, "suspend_after must be a bool"), + ("true", "suspend_after must be a bool"), + (None, "suspend_after must be a bool"), + ], +) +def test_rejects_suspend_after_in_the_request_of_microvm( + suspend_after: Any, message: str +) -> None: + # MicrovmInvokeOptions subclasses MicrovmRequestConfig. So it passes the + # type of config.request, and microvm() checks the field itself. + request = MicrovmInvokeOptions( + timeout=Duration.from_minutes(1), + suspend_after=suspend_after, + ) + + with pytest.raises(ValidationError, match=message): + microvm(NoOperations(), "build", job_input={}, config=config(request=request)) # type: ignore[arg-type] + + +def test_accepts_invoke_options_without_suspend_after_as_the_request_of_microvm() -> ( + None +): + request = MicrovmInvokeOptions(timeout=Duration.from_minutes(1), path="/build") + + with pytest.raises(AssertionError, match="no durable operation may start"): + microvm(NoOperations(), "build", job_input={}, config=config(request=request)) # type: ignore[arg-type] diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py index ebc397a1d..3fa7a0896 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py @@ -19,16 +19,17 @@ import functools import http.client -import http.server import json +import logging import socket import threading import time from collections.abc import Callable, Iterator +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Any import pytest -from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.config import Duration, MapConfig from aws_durable_execution_sdk_python.context import DurableContext from aws_durable_execution_sdk_python.execution import ( InvocationStatus, @@ -407,15 +408,17 @@ def post_hook(worker: MicrovmWorker, hook: str) -> None: class RecordingLogger: - """Records the worker's log messages.""" + """Records the worker's log messages and their fields.""" def __init__(self) -> None: self._lock = threading.Lock() self.messages: list[str] = [] + self._extras: list[tuple[str, dict[str, Any]]] = [] def _add(self, msg: object, *args: object, extra: Any = None) -> None: with self._lock: self.messages.append(str(msg)) + self._extras.append((str(msg), dict(extra or {}))) info = warning = error = _add @@ -423,6 +426,11 @@ def count(self, message: str) -> int: with self._lock: return self.messages.count(message) + def extras(self, message: str) -> list[dict[str, Any]]: + """The fields of each line with this message.""" + with self._lock: + return [data for m, data in self._extras if m == message] + @pytest.mark.parametrize("resume_hook", [True, False], ids=["resume-hook", "none"]) def test_resumes_a_worker_that_suspended_itself_between_two_session_jobs( @@ -560,6 +568,364 @@ def handler(event: Any, context: DurableContext) -> Any: assert refusals >= 1 +# --------------------------------------------------------------------------- +# suspend_after with the real worker +# --------------------------------------------------------------------------- + +IGNORED_WARNING = ( + "the worker in MicroVM mvm-1 does not support suspendAfter. The MicroVM " + "keeps running after the job." +) + + +class SimulatedMicrovm: + """A stand-in for Lambda that keeps the state of one MicroVM. + + 1. The worker's SuspendMicrovm sends the ``suspend`` hook, and then sets + the state to ``SUSPENDED``. Lambda calls the hook before the freeze. + 2. The session's ResumeMicrovm sends the ``resume`` hook, and then sets + the state to ``RUNNING``. A probe saw the hook before ``RUNNING``. + 3. :attr:`endpoint` stands in for the MicroVM endpoint. It forwards each + job to the worker while the MicroVM runs. While it is suspended, it + answers HTTP 502 with ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``, + as the service does. + + ``old_worker`` removes ``suspendAfter`` from each forwarded job. The + worker then behaves as a release that does not know the field. + """ + + def __init__(self, *, old_worker: bool = False) -> None: + self._lock = threading.Lock() + self._state = "RUNNING" + self.worker: MicrovmWorker | None = None + self.suspended: list[str] = [] + self.answers: list[int] = [] + self._old_worker = old_worker + simulated = self + + class Endpoint(BaseHTTPRequestHandler): + def do_POST(self) -> None: # noqa: N802 - the http.server name + body = self.rfile.read(int(self.headers.get("content-length", "0"))) + if simulated.state() == "SUSPENDED": + simulated.record(502) + self.send_response(502) + self.send_header("x-aws-proxy-error", "MICROVM_CONNECT_FAILED") + self.send_header("content-length", "0") + self.end_headers() + return + status, accepted, data = simulated.forward(self.path, body) + simulated.record(status) + self.send_response(status) + if accepted: + self.send_header("X-Microvm-Suspend-After", "accepted") + self.send_header("content-length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def log_message(self, *args: Any) -> None: + pass + + self._server = ThreadingHTTPServer(("127.0.0.1", 0), Endpoint) + threading.Thread(target=self._server.serve_forever, daemon=True).start() + + @property + def endpoint(self) -> str: + return f"http://127.0.0.1:{self._server.server_address[1]}" + + def state(self) -> str: + with self._lock: + return self._state + + def record(self, status: int) -> None: + with self._lock: + self.answers.append(status) + + def set_state(self, state: str) -> None: + with self._lock: + self._state = state + + def forward(self, path: str, body: bytes) -> tuple[int, bool, bytes]: + """Send the job to the worker. Return the status, whether the worker + accepted ``suspendAfter``, and the body of the answer.""" + worker = self.worker + assert worker is not None + if self._old_worker: + document = json.loads(body) + document.pop("suspendAfter", None) + body = json.dumps(document).encode() + connection = http.client.HTTPConnection( + host_of(worker), worker.port, timeout=WAIT + ) + try: + connection.request( + "POST", path, body=body, headers={"content-type": "application/json"} + ) + response = connection.getresponse() + data = response.read() + # The endpoint sends one header with a fixed name and a fixed + # value. A value from the worker's answer never reaches + # send_header. The session accepts only this value. So an answer + # with another value reaches the session as an answer without + # the header, which the session treats the same way. + accepted = response.getheader("X-Microvm-Suspend-After") == "accepted" + return response.status, accepted, data + finally: + connection.close() + + def wait_for(self, state: str) -> str: + deadline = time.monotonic() + WAIT + while time.monotonic() < deadline: + if self.state() == state: + return state + time.sleep(0.01) + return f"not {state}" + + def close(self) -> None: + self._server.shutdown() + self._server.server_close() + + +Pipeline = Callable[[MicrovmSession, DurableContext], Any] + + +def suspend_after_session( + fake_client: Any, + make_pipeline: Callable[[SimulatedMicrovm], Pipeline], + job: Callable[[Any, MicrovmJobContext], Any], + *, + old_worker: bool = False, +) -> tuple[Any, Any, SimulatedMicrovm, RecordingLogger]: + """Run a pipeline in a session without auto-suspend, against a real + worker behind a :class:`SimulatedMicrovm`. + + ``make_pipeline`` receives the simulated MicroVM, so the pipeline can + wait for its state. + """ + microvm = SimulatedMicrovm(old_worker=old_worker) + pipeline = make_pipeline(microvm) + holder: dict[str, Any] = {} + + class SessionMicrovms(fake_client): # type: ignore[misc, valid-type] + def get_microvm(self, **request: Any) -> dict[str, Any]: + response = super().get_microvm(**request) + return {**response, "state": microvm.state()} + + def resume_microvm(self, **request: Any) -> dict[str, Any]: + super().resume_microvm(**request) + post_hook(holder["worker"], "resume") + microvm.set_state("RUNNING") + return {} + + class WorkerMicrovms: + """The worker's lambda-microvms client.""" + + def suspend_microvm(self, *, microvmIdentifier: str) -> dict[str, Any]: # noqa: N803 + microvm.suspended.append(microvmIdentifier) + post_hook(holder["worker"], "suspend") + microvm.set_state("SUSPENDED") + return {} + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + return microvm_session( + context, "pipeline", config=holder["config"], handler=pipeline + ) + + runner = DurableFunctionTestRunner(handler=handler, poll_interval=0.05) + callbacks = RunnerCallbacks(runner) + logger = RecordingLogger() + worker = start_reachable_worker( + lambda: start_microvm_worker( + handler=job, + create_client=lambda region: callbacks, # type: ignore[arg-type, return-value] + create_microvms_client=lambda region: WorkerMicrovms(), # type: ignore[arg-type, return-value] + logger=logger, + port=0, + ) + ) + holder["worker"] = worker + microvm.worker = worker + try: + client = SessionMicrovms( + launches=[run_hook_to(worker)], endpoint=microvm.endpoint + ) + holder["config"] = MicrovmSessionConfig( + image_identifier="img", + execution_role_arn="role", + timeout=Duration.from_minutes(5), + client=client, + ) + arn = runner.run_async(input="null", execution_timeout=WAIT) + result = runner.wait_for_result(arn, timeout=WAIT) + assert worker.idle(WAIT) + finally: + worker.close() + runner.close() + microvm.close() + assert result.status is InvocationStatus.SUCCEEDED, result.error + return json.loads(result.result or "null"), client, microvm, logger + + +def options(*, suspend_after: bool = False) -> MicrovmInvokeOptions: + return MicrovmInvokeOptions( + timeout=Duration.from_minutes(1), suspend_after=suspend_after + ) + + +def echo(job_input: Any, context: MicrovmJobContext) -> Any: + return f"{job_input} done" + + +def test_suspends_after_a_job_and_the_next_job_resumes_the_microvm( + fake_client: Any, caplog: pytest.LogCaptureFixture +) -> None: + def make_pipeline(microvm: SimulatedMicrovm) -> Pipeline: + def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: + first = vm.invoke( + ctx, "build", job_input="build", options=options(suspend_after=True) + ) + between = ctx.step(lambda _: microvm.wait_for("SUSPENDED"), name="pause") + second = vm.invoke(ctx, "test", job_input="test", options=options()) + return [first, between, second] + + return pipeline + + with caplog.at_level(logging.WARNING): + outcome, client, microvm, logger = suspend_after_session( + fake_client, make_pipeline, echo + ) + + assert outcome == ["build done", "SUSPENDED", "test done"] + # The worker acknowledged suspendAfter. So the session logs no warning. + assert IGNORED_WARNING not in [record.getMessage() for record in caplog.records] + assert microvm.suspended == ["mvm-1"] + [line] = logger.extras("suspending the MicroVM after a job") + assert list(line) == ["microvmId", "callbackId"] + assert line["microvmId"] == "mvm-1" + # The flag makes the second job check the state before delivery. So it + # resumes the MicroVM before it sends the job, and gets no 502. + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert 502 not in microvm.answers + assert microvm.state() == "RUNNING" + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +def test_suspends_after_a_map_with_one_suspend_after_item(fake_client: Any) -> None: + def slow_echo(job_input: Any, context: MicrovmJobContext) -> Any: + # The items run at once. So the worker suspends only when the last + # of them ends. The sleep keeps each job running until every item + # has delivered its job. + time.sleep(1) + return f"{job_input} done" + + def make_pipeline(microvm: SimulatedMicrovm) -> Pipeline: + def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: + shards: list[Any] = ctx.map( + ["a", "b", "c"], + lambda item_context, item, index, items: vm.invoke( + item_context, + f"shard-{item}", + job_input=item, + options=options(suspend_after=item == "a"), + ), + name="shards", + config=MapConfig(max_concurrency=3), + ).get_results() + between = ctx.step(lambda _: microvm.wait_for("SUSPENDED"), name="pause") + after = vm.invoke(ctx, "after", job_input="after", options=options()) + return [shards, between, after] + + return pipeline + + outcome, client, microvm, logger = suspend_after_session( + fake_client, make_pipeline, slow_echo + ) + + assert outcome == [["a done", "b done", "c done"], "SUSPENDED", "after done"] + # One request from item a, and one suspend. + assert microvm.suspended == ["mvm-1"] + assert len(logger.extras("suspending the MicroVM after a job")) == 1 + # The job after the map resumed the MicroVM, with the flag or after a + # 502, depending on which invocation ran item a. + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert microvm.state() == "RUNNING" + + +def test_a_replay_that_lost_the_flag_resumes_through_the_check_after_a_502( + fake_client: Any, +) -> None: + # 1. The map item asks for a suspend after its job. + # 2. A durable wait ends the invocation. The next invocation replays + # the completed map without running its item. So the flag is lost. + # 3. The next job goes to the suspended MicroVM, and gets a 502 + # MICROVM_CONNECT_FAILED. The request step checks the state, resumes + # the MicroVM, and sends the job again. + def make_pipeline(microvm: SimulatedMicrovm) -> Pipeline: + def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: + shards: list[Any] = ctx.map( + ["a"], + lambda item_context, item, index, items: vm.invoke( + item_context, + f"shard-{item}", + job_input=item, + options=options(suspend_after=True), + ), + name="shards", + ).get_results() + between = ctx.step(lambda _: microvm.wait_for("SUSPENDED"), name="pause") + ctx.wait(Duration.from_seconds(1), name="wait") + after = vm.invoke(ctx, "after", job_input="after", options=options()) + return [shards, between, after] + + return pipeline + + outcome, client, microvm, _ = suspend_after_session( + fake_client, make_pipeline, echo + ) + + assert outcome == [["a done"], "SUSPENDED", "after done"] + assert microvm.suspended == ["mvm-1"] + # The job after the wait made no state check before delivery. Its first + # request got the 502. The check after it found SUSPENDED, resumed the + # MicroVM, and polled until RUNNING. The second request reached the + # worker. + assert microvm.answers[-2:] == [502, 202] + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert microvm.state() == "RUNNING" + + +def test_an_old_worker_ignores_suspend_after_and_the_session_warns_once( + fake_client: Any, caplog: pytest.LogCaptureFixture +) -> None: + def make_pipeline(microvm: SimulatedMicrovm) -> Pipeline: + def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: + return [ + vm.invoke( + ctx, name, job_input=name, options=options(suspend_after=True) + ) + for name in ["build", "test"] + ] + + return pipeline + + with caplog.at_level(logging.WARNING): + outcome, client, microvm, logger = suspend_after_session( + fake_client, make_pipeline, echo, old_worker=True + ) + + assert outcome == ["build done", "test done"] + # The worker sent no acknowledgement and never suspended the MicroVM. + assert microvm.suspended == [] + assert logger.extras("suspending the MicroVM after a job") == [] + warnings = [ + record + for record in caplog.records + if record.levelname == "WARNING" and record.getMessage() == IGNORED_WARNING + ] + assert len(warnings) == 1 + assert client.requests("resume_microvm") == [] + + class LosingProxy: """A stand-in for the MicroVM endpoint in front of a real worker. @@ -580,7 +946,7 @@ def __init__( self._lock = threading.Lock() proxy = self - class Handler(http.server.BaseHTTPRequestHandler): + class Handler(BaseHTTPRequestHandler): def do_POST(self) -> None: # noqa: N802 - the http.server name body = self.rfile.read(int(self.headers.get("content-length", "0"))) with proxy._lock: @@ -617,7 +983,7 @@ def do_POST(self) -> None: # noqa: N802 - the http.server name def log_message(self, *args: Any) -> None: pass - self._server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self._server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) self.url = f"http://127.0.0.1:{self._server.server_address[1]}" threading.Thread(target=self._server.serve_forever, daemon=True).start() diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md index 02ffa19fe..a20a861c8 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -65,7 +65,7 @@ The image needs these hook settings in `CreateMicrovmImage`: ``` - The `run` hook delivers the job, or tells the worker that jobs arrive over HTTP. So it must be enabled. -- The `resume` hook tells a session's worker that its MicroVM runs again after a suspend. It is optional, and recommended for sessions. See [Suspending when idle](#suspending-when-idle). +- The `resume` hook tells a session's worker that its MicroVM runs again after a suspend. It is optional, and recommended for sessions. A session that uses `suspendAfter` needs it. See [Suspending when idle](#suspending-when-idle) and [Suspending after a job](#suspending-after-a-job). - The `suspend` hook tells the worker that Lambda is about to freeze the MicroVM. The worker then refuses new jobs until the `resume` hook. It is optional, and recommended for sessions. - The `terminate` hook tells the worker that Lambda is about to end the MicroVM. The worker then fails the callback of each running job, so the durable function fails the job at once instead of at its heartbeat timeout or its timeout. It is optional, and recommended. Give it at least 10 seconds, because the worker reports before it answers, for at most 5 seconds. A hook timeout left unset ended the hook after about 2 seconds in testing. - The service requires the `ready` hook whenever a lifecycle hook is enabled. The worker answers it with HTTP 200. @@ -83,7 +83,7 @@ On `POST /aws/lambda-microvms/runtime/v1/run` the worker: The worker acts only on the first valid `run` hook. The MicroVM endpoint forwards a request to a hook path from any caller that holds an auth token for the MicroVM. So a later `run` hook gets HTTP 200, and is ignored. -On `POST ` the worker validates the job request, answers HTTP 202 at once, and runs the route's handler. `POST /durable-execution/v1/job` (`MICROVM_JOB_PATH`) runs `handler`. That path is reserved, so a route cannot use it. A job request can arrive before the `run` hook. Until the `run` hook arrives, the worker takes the MicroVM ID for `context.microvm_id` from the first job request that has a `microvmId`. The `run` hook's ID then replaces it for later jobs. Without either, `context.microvm_id` is `"unknown"`. +On `POST ` the worker validates the job request, answers HTTP 202 at once, and runs the route's handler. A job request with `"suspendAfter": true` gets the header `X-Microvm-Suspend-After: accepted` in its 202 answer. See [Suspending after a job](#suspending-after-a-job). `POST /durable-execution/v1/job` (`MICROVM_JOB_PATH`) runs `handler`. That path is reserved, so a route cannot use it. A job request can arrive before the `run` hook. Until the `run` hook arrives, the worker takes the MicroVM ID for `context.microvm_id` from the first job request that has a `microvmId`. The `run` hook's ID then replaces it for later jobs. Without either, `context.microvm_id` is `"unknown"`. Each job runs in its own thread. For every job, the worker: @@ -116,14 +116,33 @@ A session's run hook payload can set `autoSuspendIdleSeconds`. The session sets 6. A call that the service rejected on its first attempt, such as `AccessDeniedException` for a missing permission, is logged as a warning, and the MicroVM keeps running. Missing credentials count as a rejection. A later call in the same MicroVM most likely finds no credentials either. The worker accepts jobs again at once and does not retry the call. The next job that ends starts the idle time again. The exception is a `resume` hook that arrived during the call. The idle time then starts when the call ends. So the worker retries the call one idle time after the call ends, if no job runs and no `suspend` hook refusal is in effect when the call ends. A second exception is a `suspend` hook that arrived during the call and whose refusal is still in effect when the call ends. The MicroVM is then suspending anyway. So the rejection does not end the hook's refusal. That refusal ends on the `resume` hook, or 30 seconds after the last `suspend` hook. The idle time then starts again. 7. A call that the service refused for now on its first attempt is logged as a warning too, and the worker accepts jobs again at once. This covers expired credentials and a signature error, with or without a clock offset, whatever the HTTP status. The service answers these before it acts. botocore retries none of these errors, so they reach the worker on the first attempt. The idle time then starts again, so the worker tries the call again. That idle time lasts at least 30 seconds, also when `autoSuspendIdleSeconds` is shorter. The same error can repeat on every call, for example with a wrong secret key or an expired session token from environment variables. So while no job ends and no lifecycle hook arrives, the worker makes at most one call per 30 seconds, and logs one warning per call. A job that ends after the refusal, or a `resume` hook after it, starts the normal idle time instead. A `resume` hook or a job that ends during the call does not shorten the idle time after the refusal, because the refused outcome arrives after it. A `suspend` hook whose refusal is still in effect when the call ends is the exception, as in item 6. Its refusal stays, and the idle time starts when it ends. The default client makes one attempt. So a throttle from it counts as refused for now. A client from `create_microvms_client` that retries a throttle turns it into an unknown outcome, as in item 8. The JavaScript worker classifies these errors the same way. 8. A call whose outcome is unknown is handled like a successful call, because the service may have suspended the MicroVM anyway. This covers a timeout, a conflict, a server error, and a dropped connection. An error after botocore retried an earlier attempt also counts as unknown, because the earlier attempt may have suspended the MicroVM. Only a client from `create_microvms_client` that retries can raise such an error. A call that the `before-send` handler stopped after a `resume` hook counts as unknown too. The handler can stop the first attempt, and then nothing was sent. An earlier attempt, if there was one, may have suspended the MicroVM. That `resume` hook has already ended the refusal. So the worker accepts jobs at once, and starts no new 30-second limit for that call. The refusal continues until the `resume` hook or the 30-second limit. The idle time then starts again, so the worker tries the call again. A `resume` hook during the call, or the end of a `suspend` hook's 30 seconds during the call, has already ended the refusal. The worker then accepts jobs at once, and tries the call again one idle time after the call and every job end. -9. The worker logs one warning for each failed call. The message names the outcome. The fields are, in order, `microvmId`, `error`, `outcome`, `suspendHookArrived`, `terminating`, `closed`, and `refusingJobs`. `microvmId` is the MicroVM ID of the call. `error` holds the name and the message of the error. `outcome` is `rejected`, `refusedForNow`, or `unknown`. `suspendHookArrived` states whether a `suspend` hook arrived since the idle time ended. `terminating` states whether a `terminate` hook came before the worker acted on the outcome. `closed` states whether `worker.close()` came before the worker acted on the outcome. `refusingJobs` states whether the worker answers jobs with HTTP 503 after it has acted on the outcome. The worker also logs two info lines about the call. It logs `suspending the idle MicroVM` before each call, with `microvmId`, `idleSeconds`, and `waitedSeconds`. It logs `stopping the suspend call: the MicroVM resumed during it` for the first `resume` hook during a call, with `microvmId`. The JavaScript worker logs the same stop line. +9. The worker logs one warning for each failed call. The message names the outcome. The fields are, in order, `microvmId`, `error`, `outcome`, `suspendHookArrived`, `terminating`, `closed`, `refusingJobs`, and `trigger`. `microvmId` is the MicroVM ID of the call. `error` holds the name and the message of the error. `outcome` is `rejected`, `refusedForNow`, or `unknown`. `suspendHookArrived` states whether a `suspend` hook arrived since the worker decided to make the call. `terminating` states whether a `terminate` hook came before the worker acted on the outcome. `closed` states whether `worker.close()` came before the worker acted on the outcome. `refusingJobs` states whether the worker answers jobs with HTTP 503 after it has acted on the outcome. `trigger` is `idle` for a call after the idle time, and `job` for a call after a `suspendAfter` job. The worker also logs two info lines about the call. It logs `suspending the idle MicroVM` before each call, with `microvmId`, `idleSeconds`, and `waitedSeconds`. It logs `stopping the suspend call: the MicroVM resumed during it` for the first `resume` hook during a call, with `microvmId`. The JavaScript worker logs the same stop line. -`worker.close()` stops the idle time. A closed worker never suspends its MicroVM again. A refusal that is already running still ends on the `resume` hook or after 30 seconds. +`worker.close()` stops the idle time, and drops a pending suspend after a job. A closed worker never suspends its MicroVM again. A refusal that is already running still ends on the `resume` hook or after 30 seconds. The worker counts only jobs. It cannot see background threads or processes, or inbound traffic that is not a job. A session with such work leaves `auto_suspend_on_idle` off, and the payload then has no `autoSuspendIdleSeconds`. The worker creates a `lambda-microvms` client for each suspend, and closes it after the call. Pass `create_microvms_client` to provide the client. The worker never closes a client that it returns. Give the client one attempt, as the default client has. The worker starts no attempt after a `resume` hook. A retry that starts before the `resume` hook arrives can still suspend the MicroVM after it resumed. A client that retries must report its retries in `ResponseMetadata.RetryAttempts` of the error that it raises, as botocore does. The worker otherwise counts a retried call as a first attempt. An earlier attempt of that call may have suspended the MicroVM. So a retried throttle would count as refused for now, and the worker would accept a job that then freezes with the MicroVM. +## Suspending after a job + +A session job can carry `"suspendAfter": true`. The durable function sets it for `vm.invoke` with `suspend_after=True`. The worker then suspends its own MicroVM after that job. This works in every session, with or without `autoSuspendIdleSeconds`. + +1. Only a JSON `true` turns it on. The worker ignores any other value, and does not reject the job. The `run` hook payload never carries the field. +2. The worker answers the job with HTTP 202 and the header `X-Microvm-Suspend-After: accepted`. A second delivery of the job that gets the duplicate HTTP 202 carries the header when the first delivery carried the field. A second delivery refused with HTTP 503 carries no header. An older worker ignores the field and sends no header. The durable function then logs a warning, and the MicroVM keeps running. +3. When the job has ended, and its outcome reports have ended, the worker sets a pending suspend request. This happens whether the handler succeeded or failed. A job that ends while the worker refuses jobs for a suspend sets no request. That refusal comes from the worker's own call, a `suspend` hook, or the 30 seconds that follow either. The MicroVM is suspending already. So the request is done. +4. The request fires when no job runs, no callback failure report runs, no `run` hook request is being handled, no `SuspendMicrovm` call runs, and the worker is not terminating, not closed, and not refusing jobs for a suspend. The worker checks these conditions again when a job ends, a report ends, a `run` hook request ends, and a call ends. So a request from a job that ends while another job runs fires when the last job ends. +5. Firing clears the request before the call starts. A `resume` hook never sets it again. A call after the idle time clears it too, because that call suspends the MicroVM as well. A `suspend` hook clears it too, because Lambda suspends the MicroVM after that hook. +6. Firing stops the idle time, when the session also suspends when idle. +7. The `terminate` hook and `worker.close()` drop the request. +8. The worker calls `SuspendMicrovm` with the MicroVM ID of the first accepted `run` hook with a payload. It never uses the MicroVM ID of a job request. A forged request could otherwise suspend another MicroVM of the same image. Without such a `run` hook, the worker logs the warning `cannot suspend the MicroVM after a job: the run hook named no MicroVM`, with `callbackId`, and does not suspend. +9. The call follows items 3 to 9 of [Suspending when idle](#suspending-when-idle): the refusal of jobs from the moment the request fires, one attempt and none after a `resume` hook, the outcome classification, and the 30-second refusal after the call. +10. A call that the service refused for now is not retried without `autoSuspendIdleSeconds`. With it, the worker tries again after the idle time, at least 30 seconds later, as in item 7 of [Suspending when idle](#suspending-when-idle). A rejected call is not retried. + +Before the call, the worker logs `suspending the MicroVM after a job` at info level, with `microvmId` and `callbackId`. A failed call logs the warning of item 9 of [Suspending when idle](#suspending-when-idle), with `trigger` set to `job`. + +Enable the `resume` hook in the image. Without it, the worker refuses the next job with HTTP 503 until 30 seconds after the call. The next job's heartbeat timeout counts that time. The MicroVM's execution role needs `lambda:SuspendMicrovm` on the image. Without it, each call is rejected and logged, and the MicroVM keeps running. + ## Suspend and terminate hooks Lambda waits for the answer to a `suspend` or `terminate` hook, up to the hook's timeout. The network and the execution role's credentials work while the hook runs. These facts come from probes in us-east-1, which the JavaScript SDK ran. @@ -173,7 +192,7 @@ The `runHookPayload` string holds: } ``` -`job` is absent when the job arrives over HTTP. An HTTP job request body has the job fields at the top level, with `version`, `region`, and an optional `microvmId`. `MicrovmJobRequest` holds those job fields in its `job` attribute. +`job` is absent when the job arrives over HTTP. An HTTP job request body has the job fields at the top level, with `version`, `region`, an optional `microvmId`, and an optional `suspendAfter`. `MicrovmJobRequest` holds those job fields in its `job` attribute, and `suspendAfter` in its `suspend_after` attribute. A document that does not match raises `InvalidRunHookPayloadError`. When the document named a non-empty callback ID, the error carries `callback_id` and `region`, so the worker can fail that callback at once. @@ -207,7 +226,7 @@ The worker logs to the standard library logger `aws_durable_execution_sdk_python The MicroVM's execution role needs `lambda:SendDurableExecutionCallbackSuccess`, `lambda:SendDurableExecutionCallbackFailure`, and `lambda:SendDurableExecutionCallbackHeartbeat` on the durable function's ARN (`arn:aws:lambda:::function::*`). -A session that suspends when idle also needs `lambda:SuspendMicrovm` on the MicroVM image (`arn:aws:lambda:::microvm-image:`). The action authorizes on the image, not on the MicroVM. So a MicroVM with this permission can suspend any MicroVM from the same image, even during a job. That job then waits until its heartbeat timeout or its timeout. Use one image per trust boundary. +A session that suspends when idle, or that sends a job with `suspendAfter`, also needs `lambda:SuspendMicrovm` on the MicroVM image (`arn:aws:lambda:::microvm-image:`). The action authorizes on the image, not on the MicroVM. So a MicroVM with this permission can suspend any MicroVM from the same image, even during a job. That job then waits until its heartbeat timeout or its timeout. Use one image per trust boundary. The worker needs boto3 1.43.35 or later, the first release with the `lambda-microvms` client. diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/payload.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/payload.py index 18e047ced..30a780ada 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/payload.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/payload.py @@ -235,6 +235,14 @@ class MicrovmJobRequest: """The MicroVM identifier. A job request can reach the worker before the ``run`` hook, which also carries the identifier. Until the ``run`` hook arrives, the worker takes it from the first job request that has one.""" + suspend_after: bool = False + """Whether the worker suspends its MicroVM after the job. + + Wire name: ``suspendAfter``. Only a JSON ``true`` turns it on. The + worker ignores any other value, as the JavaScript worker does. So a + document from either SDK has the same meaning in both workers. The + worker suspends the MicroVM named by the ``run`` hook, never the + ``microvmId`` of the request.""" @classmethod def from_dict(cls, body: object) -> MicrovmJobRequest: @@ -267,6 +275,7 @@ def from_dict(cls, body: object) -> MicrovmJobRequest: region=region, job=job, microvm_id=microvm_id, + suspend_after=body.get("suspendAfter") is True, ) diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py index f13a6c485..8e24dcb25 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py @@ -112,6 +112,16 @@ _UNKNOWN_MICROVM_ID = "unknown" +SUSPEND_AFTER_HEADER = "X-Microvm-Suspend-After" +"""The header of a 202 answer that acknowledges ``"suspendAfter": true``. + +Its value is ``accepted``. The durable function reads it for a +``suspendAfter`` job. An older worker ignores the field and sends no header. +The durable function then logs a warning. The JavaScript worker sends the +same header.""" + +_SUSPEND_AFTER_ACK = {SUSPEND_AFTER_HEADER: "accepted"} + # At most 19 digits, so that int() stays under Python's limit on digits, and # any declared length still exceeds every body limit. _CONTENT_LENGTH = re.compile(r"[0-9]{1,19}") @@ -191,17 +201,57 @@ def _start_timer(delay: float, action: Callable[[], None]) -> Callable[[], None] } +@dataclass(frozen=True, kw_only=True) +class _SuspendTarget: + """The MicroVM that the worker's own SuspendMicrovm call suspends. + + It comes from the first accepted run hook with a payload, never from a + job request. SuspendMicrovm authorizes on the image. So a MicroVM ID + from a job body could suspend another MicroVM of the same image. + """ + + microvm_id: str + region: str + + @dataclass(frozen=True, kw_only=True) class _AutoSuspend: """The run hook's request to suspend the MicroVM when idle.""" idle_seconds: float - region: str - microvm_id: str - """The run hook's MicroVM ID. SuspendMicrovm uses it, never an ID from a + target: _SuspendTarget + """The run hook's MicroVM. SuspendMicrovm uses it, never an ID from a job request.""" +@dataclass(frozen=True, kw_only=True) +class _SuspendCall: + """One SuspendMicrovm call of the worker, as the locked block that + started the refusal decided it.""" + + target: _SuspendTarget + trigger: str + """``idle`` for the idle time, ``job`` for a ``suspendAfter`` job.""" + hooks_before: int + """The suspend hooks counted when the refusal started.""" + resumes_before: int + """The resume hooks counted when the refusal started.""" + message: str + """The info line before the call.""" + fields: Mapping[str, object] + """The fields of that line.""" + + +@dataclass(frozen=True, kw_only=True) +class _JobSuspend: + """A pending suspend after a ``suspendAfter`` job that fires now.""" + + callback_id: str + """The job whose end asked for the suspend.""" + call: _SuspendCall | None + """The call to make. ``None`` when no run hook named a MicroVM.""" + + @dataclass(kw_only=True) class _ActiveJob: """A job whose handler has not settled yet, with what the ``terminate`` @@ -352,15 +402,38 @@ class MicrovmWorkerListener: 9. The listener logs one warning for each failed call. The message names the outcome. The fields are, in order, ``microvmId``, ``error``, ``outcome``, ``suspendHookArrived``, ``terminating``, - ``closed``, and ``refusingJobs``. ``microvmId`` is the MicroVM ID of - the call. ``error`` holds the name and the message of the error. - ``outcome`` is ``rejected``, ``refusedForNow``, or ``unknown``. - ``suspendHookArrived`` states whether a ``suspend`` hook arrived - since the idle time ended. ``terminating`` states whether a - ``terminate`` hook came before the listener acted on the outcome. - ``closed`` states whether ``close()`` came before the listener acted - on the outcome. ``refusingJobs`` states whether the listener answers - jobs with HTTP 503 after it has acted on the outcome. + ``closed``, ``refusingJobs``, and ``trigger``. ``microvmId`` is the + MicroVM ID of the call. ``error`` holds the name and the message of + the error. ``outcome`` is ``rejected``, ``refusedForNow``, or + ``unknown``. ``suspendHookArrived`` states whether a ``suspend`` hook + arrived since the listener decided to make the call. ``terminating`` + states whether a ``terminate`` hook came before the listener acted + on the outcome. ``closed`` states whether ``close()`` came before the + listener acted on the outcome. ``refusingJobs`` states whether the + listener answers jobs with HTTP 503 after it has acted on the + outcome. ``trigger`` is ``idle`` or ``job``. + + A job request with ``"suspendAfter": true`` asks the listener to + suspend its MicroVM after the job, in any session: + + 1. The 202 answer carries ``X-Microvm-Suspend-After: accepted``. A + second delivery that gets the duplicate 202 carries the header when + the first one carried the field. A 503 never carries it. Only a + JSON ``true`` counts. Any other value is ignored. + 2. When the job's outcome reports have ended, the listener sets a + pending request. A job that ends while the listener refuses jobs for + a suspend sets none, because the MicroVM is suspending already. The + request fires when no job, failure report, ``run`` hook request, or + SuspendMicrovm call runs, and the listener is not terminating, not + closed, and not refusing jobs. Each of those ends checks it again. + 3. Firing clears the request, stops the idle time, and makes the same + call as items 3 to 9, with the MicroVM ID of the first accepted + ``run`` hook with a payload. Without that hook, the listener logs a + warning and does not suspend. A ``resume`` hook never repeats the + request. A ``suspend`` hook, the ``terminate`` hook, and ``close()`` + drop it. + 4. A call refused for now is retried only when the session suspends + when idle. When the image enables the ``suspend`` and ``terminate`` hooks, the listener acts on them: @@ -422,8 +495,12 @@ def __init__( self._locked_depth = 0 # The callback IDs of the jobs that run now. self._jobs: set[str] = set() - # The callback IDs of jobs that ended, oldest first. - self._finished: OrderedDict[str, None] = OrderedDict() + # The running jobs whose request carried "suspendAfter": true. + self._suspend_after_jobs: set[str] = set() + # The callback IDs of jobs that ended, oldest first. The value states + # whether the job's request carried "suspendAfter": true. A resent + # copy of the job then gets the same acknowledgement. + self._finished: OrderedDict[str, bool] = OrderedDict() # Callback failure reports in flight, from every place that fails a # callback. The terminate hook waits for them, because Lambda can end # the MicroVM as soon as it gets the answer. @@ -448,6 +525,14 @@ def __init__( self._reporting: set[str] = set() # Set from the run hook payload of a session that suspends when idle. self._auto_suspend: _AutoSuspend | None = None + # Set from the first accepted run hook that has a payload, in every + # session. Both triggers of a suspend call use it. + self._suspend_target: _SuspendTarget | None = None + # The callback ID of the suspendAfter job whose end asked for a + # suspend that has not started yet. The request fires when no work + # runs. A call that starts clears it, so a resume hook never repeats + # it. A suspend hook, the terminate hook, and close() clear it too. + self._suspend_requested: str | None = None # Each timer action carries the generation it was started with, and # acts only when it is still current. Cancelling a timer increments # the generation. @@ -527,22 +612,23 @@ def create( its own per job, from the default credential chain. The listener closes that client when the job ends. create_microvms_client: Creates the ``lambda-microvms`` client - that suspends this MicroVM when it is idle. The listener - calls it only when a session asks the worker to suspend when - idle, once per suspend. It never closes a client that this - function returns. Defaults to a client of its own per - suspend, from the default credential chain. The listener - closes that client after the call. The listener stops the - client's attempts once a ``resume`` hook arrived. Give the - client one attempt, as the default client has. A retry - that starts before the ``resume`` hook arrives can still - suspend the MicroVM after it resumed. A client that retries - must report its retries in ``ResponseMetadata.RetryAttempts`` - of the error that it raises. The listener otherwise counts a - retried call as a first attempt. An earlier attempt of that - call may have suspended the MicroVM. A retried throttle would - then count as refused for now, and the listener would accept - a job that then freezes with the MicroVM. + that suspends this MicroVM. The listener calls it only when + a session asks the worker to suspend when idle, or after a + ``suspendAfter`` job, once per suspend. It never closes a + client that this function returns. Defaults to a client of + its own per suspend, from the default credential chain. The + listener closes that client after the call. The listener + stops the client's attempts once a ``resume`` hook arrived. + Give the client one attempt, as the default client has. A + retry that starts before the ``resume`` hook arrives can + still suspend the MicroVM after it resumed. A client that + retries must report its retries in + ``ResponseMetadata.RetryAttempts`` of the error that it + raises. The listener otherwise counts a retried call as a + first attempt. An earlier attempt of that call may have + suspended the MicroVM. A retried throttle would then count + as refused for now, and the listener would accept a job that + then freezes with the MicroVM. heartbeat_interval_seconds: The heartbeat interval. Defaults to a third of the job's heartbeat timeout, at most 15 minutes. It is cut to a third of the job's heartbeat timeout. @@ -654,7 +740,8 @@ def idle(self, timeout: float | None = None) -> bool: ) def close(self) -> None: - """Stop the idle time. The listener then never suspends the MicroVM again. + """Stop the idle time, and drop a pending suspend after a job. The + listener then never suspends the MicroVM again. Running jobs continue. A refusal that a suspend started still ends on the ``resume`` hook or after 30 seconds. @@ -662,6 +749,7 @@ def close(self) -> None: with self._lock: self._closed = True self._stop_idle_time() + self._suspend_requested = None def _dispatch(self, exchange: _Exchange) -> None: path = exchange.path @@ -673,9 +761,11 @@ def _dispatch(self, exchange: _Exchange) -> None: try: self._handle_run_hook(exchange) finally: - with self._changed: + with self._locked(): self._run_hooks -= 1 self._changed.notify_all() + decision = self._take_job_suspend() + self._act_on_job_suspend(decision) return if hook in {"suspend", "terminate", "resume"} and not _is_lambda_hook_call( exchange @@ -720,12 +810,14 @@ def _handle_suspend_hook(self) -> None: """Refuse jobs until the ``resume`` hook, or for 30 seconds. Lambda freezes the MicroVM after the answer. A job accepted now would - freeze with it. + freeze with it. Lambda suspends the MicroVM after this hook. So a + pending suspend after a job is done, and it is cleared. """ with self._locked(): self._suspend_hooks += 1 self._stop_idle_time() self._suspending = True + self._suspend_requested = None self._start_grace() def _handle_resume_hook(self) -> list[tuple[MicrovmJobDocument, str]]: @@ -790,6 +882,7 @@ def _terminate(self) -> None: with self._lock: self._terminating = True self._stop_idle_time() + self._suspend_requested = None stopped: list[tuple[str, _ActiveJob]] = [] for callback_id, job in self._active.items(): if job.cancelled.is_set(): @@ -856,15 +949,18 @@ def _handle_run_hook(self, exchange: _Exchange) -> None: if not later: self._run_hook_accepted = True self._microvm_id = parsed.microvm_id - if ( - payload is not None - and payload.auto_suspend_idle_seconds is not None - ): - self._auto_suspend = _AutoSuspend( - idle_seconds=payload.auto_suspend_idle_seconds, - region=payload.region, - microvm_id=parsed.microvm_id, + if payload is not None: + # Every session gets a target, also without auto-suspend. + # A suspendAfter job needs one. + target = _SuspendTarget( + microvm_id=parsed.microvm_id, region=payload.region ) + self._suspend_target = target + if payload.auto_suspend_idle_seconds is not None: + self._auto_suspend = _AutoSuspend( + idle_seconds=payload.auto_suspend_idle_seconds, + target=target, + ) microvm_id = self._microvm_id duplicate = job is not None and self._is_known(job.callback_id) new_job = None if later or duplicate else job @@ -1014,9 +1110,22 @@ def _handle_job_request( terminating = not duplicate and self._terminating suspending = not duplicate and self._suspending if not (duplicate or terminating or suspending): - self._register_job(job.callback_id, request.region) + self._register_job( + job.callback_id, + request.region, + suspend_after=request.suspend_after, + ) + # A resent copy that gets the duplicate 202 carries the + # acknowledgement that the first request got, because only the + # first request decides the suspend. A 503 carries none. + acknowledge = ( + self._carried_suspend_after(job.callback_id) + if duplicate + else request.suspend_after + ) + headers = _SUSPEND_AFTER_ACK if acknowledge else None if duplicate: - exchange.respond(202, {}) + exchange.respond(202, {}, headers) self._logger.warning( "duplicate job ignored", extra={"microvmId": microvm_id, "callbackId": job.callback_id}, @@ -1038,7 +1147,7 @@ def _handle_job_request( # retries. After a 202, the durable function does not deliver the job # again. self._start_job(handler, job, request.region) - exchange.respond(202, {}) + exchange.respond(202, {}, headers) def _reject(self, error: BaseException, exchange: _Exchange) -> None: """Answer 400 to a request that the worker could not accept. @@ -1071,7 +1180,16 @@ def _is_known(self, callback_id: str) -> bool: """ return callback_id in self._jobs or callback_id in self._finished - def _register_job(self, callback_id: str, region: str) -> None: + def _carried_suspend_after(self, callback_id: str) -> bool: + """Whether the request of a known job carried ``"suspendAfter": + true``. The caller holds the lock.""" + return callback_id in self._suspend_after_jobs or self._finished.get( + callback_id, False + ) + + def _register_job( + self, callback_id: str, region: str, *, suspend_after: bool = False + ) -> None: """Register a job that is about to start. The caller holds the lock. The job is registered for the terminate hook before its thread @@ -1080,6 +1198,8 @@ def _register_job(self, callback_id: str, region: str) -> None: handler. """ self._jobs.add(callback_id) + if suspend_after: + self._suspend_after_jobs.add(callback_id) self._active[callback_id] = _ActiveJob( region=region, cancelled=threading.Event() ) @@ -1101,12 +1221,16 @@ def _start_job( # The job never ran. A job request then answers 500, so the # durable function delivers the job again, and the retry may run # it. The run hook fails the callback instead, because Lambda - # sends it only once. + # sends it only once. The job never ran, so it asks for no + # suspend. A suspend that another job asked for can fire now. with self._locked(): self._jobs.discard(job.callback_id) + self._suspend_after_jobs.discard(job.callback_id) self._active.pop(job.callback_id, None) self._start_idle_time() self._changed.notify_all() + decision = self._take_job_suspend() + self._act_on_job_suspend(decision) raise def _run_job( @@ -1120,14 +1244,29 @@ def _run_job( extra={"callbackId": job.callback_id, "error": describe(error)}, ) finally: + # The job's report attempts have ended here, whether the handler + # succeeded or failed. So a suspendAfter job asks for the suspend + # now. The suspend runs after this block, so that no lock is + # held while it logs or starts its thread. During a suspend + # refusal, the MicroVM is suspending already. So a job that ends + # then asks for nothing. with self._locked(): - self._jobs.discard(job.callback_id) - self._active.pop(job.callback_id, None) - self._finished[job.callback_id] = None + callback_id = job.callback_id + suspend_after = callback_id in self._suspend_after_jobs + self._jobs.discard(callback_id) + self._suspend_after_jobs.discard(callback_id) + self._active.pop(callback_id, None) + self._finished[callback_id] = suspend_after if len(self._finished) > MAX_FINISHED_JOBS: self._finished.popitem(last=False) self._start_idle_time() + if suspend_after and not ( + self._terminating or self._closed or self._suspending + ): + self._suspend_requested = callback_id self._changed.notify_all() + decision = self._take_job_suspend() + self._act_on_job_suspend(decision) def _run_job_and_report( self, handler: MicrovmJobHandler, job: MicrovmJobDocument, region: str @@ -1354,9 +1493,11 @@ def _send_failure( self._report_ended() def _report_ended(self) -> None: - with self._changed: + with self._locked(): self._reports -= 1 self._changed.notify_all() + decision = self._take_job_suspend() + self._act_on_job_suspend(decision) def _reporter_for(self, callback_id: str, region: str) -> CallbackReporter: # Each job gets its own client. A client kept across jobs would keep @@ -1491,20 +1632,120 @@ def _idle_time_ended(self, generation: int, waited_seconds: float) -> None: ): return self._cancel_idle_timer = None - self._suspending = True - self._suspend_calls += 1 + # waitedSeconds differs from idleSeconds after a refusal for now, + # because that idle time lasts at least SUSPEND_GRACE_SECONDS. + call = self._begin_suspend_call( + auto_suspend.target, + "idle", + "suspending the idle MicroVM", + { + "microvmId": auto_suspend.target.microvm_id, + "idleSeconds": auto_suspend.idle_seconds, + "waitedSeconds": waited_seconds, + }, + ) + self._start_suspend_thread( + call, "could not start the suspend of the idle MicroVM" + ) + + def _begin_suspend_call( + self, + target: _SuspendTarget, + trigger: str, + message: str, + fields: Mapping[str, object], + ) -> _SuspendCall: + """Start the refusal of jobs for a SuspendMicrovm call, and count the + call. The caller holds the lock, and starts the call after it + releases the lock. + + The call suspends the MicroVM either way. So a pending suspend after + a job is done by this call too, and it is cleared. + """ + self._suspending = True + self._suspend_calls += 1 + self._suspend_requested = None + self._stoppable_suspend = target.microvm_id + return _SuspendCall( + target=target, + trigger=trigger, # A suspend hook can arrive at any point from here on: while the - # thread starts, while the client is created, or during the call. - # Each path that ends the refusal compares with this count. - hooks_before = self._suspend_hooks + # thread starts, while the client is created, or during the + # call. Each path that ends the refusal compares with this count. + hooks_before=self._suspend_hooks, # A resume hook from here on ends the refusal. The call then # starts no further attempt. - resumes_before = self._resume_hooks - self._stoppable_suspend = auto_suspend.microvm_id + resumes_before=self._resume_hooks, + message=message, + fields=fields, + ) + + def _take_job_suspend(self) -> _JobSuspend | None: + """Decide whether the pending suspend after a ``suspendAfter`` job + fires now. + + The request fires when no job, failure report, run hook, or + SuspendMicrovm call runs, and the listener is not terminating, not + closed, and not refusing jobs for a suspend. Each event that can + change one of these calls this method: a job end, a report end, a + run hook end, and a call end. Firing clears the request, and stops + an idle timer that auto-suspend armed. + + The caller holds the lock through :meth:`_locked`. It passes the + result to :meth:`_act_on_job_suspend` after it releases the lock. So + no log line, ``describe()``, or caller code runs under the lock, and + no locking method runs inside the caller's locked block. + """ + callback_id = self._suspend_requested + if ( + callback_id is None + or self._jobs + or self._reports + or self._run_hooks + or self._terminating + or self._closed + or self._suspending + or self._suspend_calls + ): + return None + self._suspend_requested = None + target = self._suspend_target + if target is None: + return _JobSuspend(callback_id=callback_id, call=None) + self._stop_idle_time() + call = self._begin_suspend_call( + target, + "job", + "suspending the MicroVM after a job", + {"microvmId": target.microvm_id, "callbackId": callback_id}, + ) + return _JobSuspend(callback_id=callback_id, call=call) + + def _act_on_job_suspend(self, decision: _JobSuspend | None) -> None: + """Start the suspend that :meth:`_take_job_suspend` decided. The + caller does not hold the lock.""" + if decision is None: + return + if decision.call is None: + # A job request carries a MicroVM ID, but a forged request could + # name another MicroVM of the image. So only the run hook names + # the MicroVM to suspend. + self._logger.warning( + "cannot suspend the MicroVM after a job: the run hook named no MicroVM", + extra={"callbackId": decision.callback_id}, + ) + return + self._start_suspend_thread( + decision.call, "could not start the suspend of the MicroVM after a job" + ) + + def _start_suspend_thread(self, call: _SuspendCall, failure: str) -> None: + """Run a SuspendMicrovm call in its own thread. The caller started + the refusal and counted the call, and does not hold the lock.""" try: threading.Thread( target=self._suspend_self, - args=(auto_suspend, hooks_before, resumes_before, waited_seconds), + args=(call,), name="microvm-suspend", daemon=True, ).start() @@ -1512,35 +1753,29 @@ def _idle_time_ended(self, generation: int, waited_seconds: float) -> None: # No call was made. So jobs can run again at once, unless a # suspend hook started its own refusal meanwhile. self._logger.error( - "could not start the suspend of the idle MicroVM", - extra={"microvmId": auto_suspend.microvm_id, "error": describe(error)}, + failure, + extra={"microvmId": call.target.microvm_id, "error": describe(error)}, ) with self._locked(): - deferred = self._end_own_refusal(hooks_before) - self._suspend_call_ended() + deferred = self._end_own_refusal(call.hooks_before) + decision = self._suspend_call_ended() self._launch_deferred(deferred) + self._act_on_job_suspend(decision) - def _suspend_self( - self, - auto_suspend: _AutoSuspend, - hooks_before: int, - resumes_before: int, - waited_seconds: float, - ) -> None: + def _suspend_self(self, call: _SuspendCall) -> None: try: - self._call_suspend( - auto_suspend, hooks_before, resumes_before, waited_seconds - ) + self._call_suspend(call) finally: with self._locked(): - self._suspend_call_ended() + decision = self._suspend_call_ended() + self._act_on_job_suspend(decision) def _end_own_refusal( self, hooks_before: int ) -> list[tuple[MicrovmJobDocument, str]]: """End the refusal that the worker's own suspend started. - A suspend hook that arrived since the idle time ended means that + A suspend hook that arrived since the call was decided means that Lambda is freezing the MicroVM. Its refusal then stays. The caller holds the lock, and starts the returned jobs after it releases it. """ @@ -1548,10 +1783,14 @@ def _end_own_refusal( return [] return self._stop_refusing() - def _suspend_call_ended(self) -> None: + def _suspend_call_ended(self) -> _JobSuspend | None: """Count a SuspendMicrovm call as ended, and start the idle time that - a resume hook, a job, or a refusal for now asked for during it. The - caller holds the lock through :meth:`_locked`.""" + a resume hook, a job, or a refusal for now asked for during it. + + A suspend after a job that ended during the call can fire now. The + caller holds the lock through :meth:`_locked`, and passes the result + to :meth:`_act_on_job_suspend` after it releases the lock. + """ self._suspend_calls -= 1 # The thread or the client could fail before the call was sent. self._stoppable_suspend = None @@ -1559,24 +1798,21 @@ def _suspend_call_ended(self) -> None: self._idle_after_call = False self._start_idle_time(after_refusal=self._idle_after_call_floor) self._changed.notify_all() + return self._take_job_suspend() - def _call_suspend( - self, - auto_suspend: _AutoSuspend, - hooks_before: int, - resumes_before: int, - waited_seconds: float, - ) -> None: - target = auto_suspend.microvm_id + def _call_suspend(self, call: _SuspendCall) -> None: + target = call.target.microvm_id + hooks_before = call.hooks_before fields = {"microvmId": target} # The listener closes only a client that it created itself. A client # from create_microvms_client belongs to the caller. owns_client = self._create_microvms_client is None + region = call.target.region try: client = ( - create_microvms_client(auto_suspend.region) + create_microvms_client(region) if self._create_microvms_client is None - else self._create_microvms_client(auto_suspend.region) + else self._create_microvms_client(region) ) except Exception as error: # noqa: BLE001 # No call was made. So jobs can run again at once, unless a @@ -1589,23 +1825,14 @@ def _call_suspend( deferred = self._end_own_refusal(hooks_before) self._launch_deferred(deferred) return - # waitedSeconds differs from idleSeconds after a refusal for now, - # because that idle time lasts at least SUSPEND_GRACE_SECONDS. - self._logger.info( - "suspending the idle MicroVM", - extra={ - **fields, - "idleSeconds": auto_suspend.idle_seconds, - "waitedSeconds": waited_seconds, - }, - ) + self._logger.info(call.message, extra=dict(call.fields)) # The block after the client is closed decides whether the refusal # continues after an unknown outcome. The warning reads its fields in # that block. So for an unknown outcome, the warning comes after # close(). unknown_described: object | None = None try: - self._send_suspend(client, target, resumes_before) + self._send_suspend(client, target, call.resumes_before) except Exception as error: # noqa: BLE001 outcome = classify_suspend_error(error) # describe() calls str(error). A client from @@ -1630,21 +1857,21 @@ def _call_suspend( deferred = [] with self._locked(): # The service did nothing. So jobs can run again at once. - # A suspend hook that arrived since the idle time ended + # A suspend hook that arrived since the call was decided # means that the MicroVM is suspending anyway, so its # refusal stays. if self._suspend_hooks == hooks_before: deferred = self._stop_refusing() # A rejection, such as a missing permission, is not # retried until a job ends. A refusal for now is - # retried after the next idle time. That idle time - # lasts at least 30 seconds, because the same error - # can repeat on every call. + # retried after the next idle time, only with + # auto-suspend. That idle time lasts at least 30 + # seconds, because the same error can repeat on every + # call. A suspendAfter job asked for one call, so + # without auto-suspend nothing retries it. if outcome is SuspendOutcome.REFUSED_FOR_NOW: self._start_idle_time(after_refusal=True) - message, extra = self._failed_call_log( - outcome, described, target, hooks_before - ) + message, extra = self._failed_call_log(outcome, described, call) self._logger.warning(message, extra=extra) self._launch_deferred(deferred) return @@ -1668,7 +1895,7 @@ def _call_suspend( if unknown_described is None: return message, extra = self._failed_call_log( - SuspendOutcome.UNKNOWN, unknown_described, target, hooks_before + SuspendOutcome.UNKNOWN, unknown_described, call ) self._logger.warning(message, extra=extra) @@ -1676,7 +1903,7 @@ def _send_suspend( self, client: LambdaMicrovmsClient, target: str, resumes_before: int ) -> None: """Make the SuspendMicrovm call. Start no attempt after a resume hook - that arrived since the idle time ended.""" + that arrived since the call was decided.""" try: # The before-send handler calls resumed() inside botocore, in # this thread. It reads one int without the lock. So it runs no @@ -1697,14 +1924,15 @@ def _failed_call_log( self, outcome: SuspendOutcome, described: object, - microvm_id: str, - hooks_before: int, + call: _SuspendCall, ) -> tuple[str, dict[str, object]]: """Return the warning for a failed SuspendMicrovm call, and its fields. The message names only the outcome. The fields state whether a - suspend hook arrived since the idle time ended, whether the listener - is terminating or closed, and whether it refuses jobs. + suspend hook arrived since the call was decided, whether the + listener is terminating or closed, whether it refuses jobs, and what + started the call. ``trigger`` is the last field, so the first seven + keep their order. The caller holds the lock, and has already acted on the outcome in the same locked block. So the fields state what the listener does @@ -1718,14 +1946,15 @@ def _failed_call_log( """ name, message = _FAILED_CALL_WARNINGS[outcome] return message, { - "microvmId": microvm_id, + "microvmId": call.target.microvm_id, "error": described, "outcome": name, - "suspendHookArrived": self._suspend_hooks != hooks_before, + "suspendHookArrived": self._suspend_hooks != call.hooks_before, "terminating": self._terminating, "closed": self._closed, # The job path refuses a job with HTTP 503 in both states. "refusingJobs": self._terminating or self._suspending, + "trigger": call.trigger, } def _start_grace(self) -> None: @@ -2056,8 +2285,14 @@ def read_json(self, max_bytes: int) -> Any: msg = "request body is not valid JSON" raise InvalidRunHookPayloadError(msg) from error - def respond(self, status: int, body: Mapping[str, object]) -> None: - """Send the answer. A client that has gone away is logged.""" + def respond( + self, + status: int, + body: Mapping[str, object], + headers: Mapping[str, str] | None = None, + ) -> None: + """Send the answer, with ``headers`` added. A client that has gone + away is logged.""" self.responded = True data = json.dumps(body, separators=(",", ":")).encode() request = self._request @@ -2065,6 +2300,8 @@ def respond(self, status: int, body: Mapping[str, object]) -> None: request.send_response(status) request.send_header("Content-Type", "application/json") request.send_header("Content-Length", str(len(data))) + for name, value in (headers or {}).items(): + request.send_header(name, value) # The header also makes the base class close the connection. request.send_header("Connection", "close") request.end_headers() diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py index 94d7c3c7b..ad7c85bc2 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py @@ -64,6 +64,8 @@ "terminating", "closed", "refusingJobs", + # The last field, so the first seven keep their order. + "trigger", ] @@ -76,6 +78,7 @@ def assert_warning( suspend_hook_arrived: bool = False, terminating: bool = False, closed: bool = False, + trigger: str = "idle", ) -> None: """Assert that one warning for a failed SuspendMicrovm call has this outcome and these fields.""" @@ -89,6 +92,7 @@ def assert_warning( "terminating": terminating, "closed": closed, "refusingJobs": refusing_jobs, + "trigger": trigger, } @@ -100,6 +104,7 @@ def assert_one_warning( suspend_hook_arrived: bool = False, terminating: bool = False, closed: bool = False, + trigger: str = "idle", ) -> dict[str, Any]: """Assert that the worker logged exactly one warning, for one failed SuspendMicrovm call, with this outcome and these fields. Returns the @@ -115,6 +120,7 @@ def assert_one_warning( suspend_hook_arrived=suspend_hook_arrived, terminating=terminating, closed=closed, + trigger=trigger, ) return data @@ -308,6 +314,9 @@ def __call__(self, job_input: Any, context: MicrovmJobContext) -> Any: self.order.append(context.callback_id) self.contexts[context.callback_id] = context self._release(context.callback_id).wait(WAIT) + if isinstance(job_input, dict) and job_input.get("fail"): + msg = "the job failed" + raise RuntimeError(msg) return "done" def started(self) -> list[str]: @@ -1794,7 +1803,7 @@ def test_accepts_jobs_after_an_unknown_outcome_when_the_refusal_ended_during_the assert logger.messages("warning") == [] end_failed_call(h) - expected = { + expected: dict[str, Any] = { "suspend_hook_arrived": during != "resume", "refusing_jobs": False, } @@ -3486,3 +3495,564 @@ def create(region: str) -> Any: assert h.job("cb-1") == 503 h.hook("resume") assert h.job("cb-1") == 202 + + +# --------------------------------------------------------------------------- +# Suspend after a suspendAfter job +# --------------------------------------------------------------------------- + +ACK_HEADER = "X-Microvm-Suspend-After" +JOB_SUSPEND_INFO = "suspending the MicroVM after a job" +NO_TARGET_WARNING = ( + "cannot suspend the MicroVM after a job: the run hook named no MicroVM" +) + + +def post_job(h: Harness, callback_id: str, **fields: Any) -> tuple[int, str | None]: + """Send a job request, and return the status and the acknowledgement + header.""" + connection = http.client.HTTPConnection( + host_of(h.worker), h.worker.port, timeout=WAIT + ) + body = { + "version": 1, + "region": "us-east-1", + "callbackId": callback_id, + "input": {}, + **fields, + } + try: + connection.request( + "POST", + "/job", + body=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + ) + response = connection.getresponse() + response.read() + finally: + connection.close() + return response.status, response.getheader(ACK_HEADER) + + +def info_lines(logger: Any, message: str) -> list[dict[str, Any]]: + return [data for level, m, data in logger.lines if level == "info" and m == message] + + +def end_job(h: Harness, callback_id: str) -> None: + """Release a job, wait until it has ended, and let a suspend call that + its end started end too.""" + h.finish_job(callback_id) + h.settle() + + +@pytest.mark.parametrize( + ("value", "acknowledged"), + [(True, "accepted"), (False, None), ("true", None), (1, None), (None, None)], + ids=["true", "false", "string", "one", "null"], +) +def test_acknowledges_only_a_json_true_and_suspends_only_for_it( + start: Any, value: Any, acknowledged: str | None +) -> None: + h = start() + h.run_hook() + + assert post_job(h, "cb-1", suspendAfter=value) == (202, acknowledged) + end_job(h, "cb-1") + + assert h.microvms.suspended == (["mvm-1"] if acknowledged else []) + + +def test_acknowledges_no_job_without_the_field(start: Any) -> None: + h = start() + h.run_hook() + + assert post_job(h, "cb-1") == (202, None) + + +@pytest.mark.parametrize("finished", [False, True], ids=["running", "finished"]) +def test_acknowledges_a_resent_copy_when_the_first_request_carried_the_field( + start: Any, finished: bool +) -> None: + h = start() + h.run_hook() + h.microvms.failure = suspend_error("AccessDeniedException", 403) + assert post_job(h, "cb-1", suspendAfter=True) == (202, "accepted") + assert post_job(h, "cb-2") == (202, None) + if finished: + end_job(h, "cb-1") + end_job(h, "cb-2") + + # Only the first request decides the suspend. So a copy gets the same + # acknowledgement, whatever it carries. + assert post_job(h, "cb-1", suspendAfter=True) == (202, "accepted") + assert post_job(h, "cb-1") == (202, "accepted") + assert post_job(h, "cb-2", suspendAfter=True) == (202, None) + + +def test_sends_no_acknowledgement_with_a_refusal(start: Any) -> None: + h = start() + h.run_hook() + h.hook("suspend") + + assert post_job(h, "cb-1", suspendAfter=True) == (503, None) + + +def test_refuses_a_resent_running_suspend_after_job_without_acknowledgement( + start: Any, +) -> None: + # The suspend refusal answers the resend of a running job with 503, so + # the copy carries no acknowledgement. The delivery after the resume gets + # the duplicate 202 with the acknowledgement of the first request. + h = start() + h.run_hook() + assert post_job(h, "cb-1", suspendAfter=True) == (202, "accepted") + assert eventually(lambda: h.started("cb-1")) + assert h.hook("suspend") == 200 + + assert post_job(h, "cb-1", suspendAfter=True) == (503, None) + assert post_job(h, "cb-1") == (503, None) + + assert h.hook("resume") == 200 + assert post_job(h, "cb-1") == (202, "accepted") + end_job(h, "cb-1") + + assert h.jobs.started() == ["cb-1"] + assert h.microvms.suspended == ["mvm-1"] + + +@pytest.mark.parametrize("refusal_end", ["resume", "grace"]) +def test_a_job_that_ends_during_a_suspend_hook_refusal_asks_for_nothing( + start: Any, refusal_end: str +) -> None: + # Lambda suspends the MicroVM after the suspend hook. So a suspendAfter + # job that ends during the refusal asks for no call of the worker's own, + # neither then nor after the refusal ends. + h = start() + h.run_hook() + assert post_job(h, "cb-1", suspendAfter=True) == (202, "accepted") + assert eventually(lambda: h.started("cb-1")) + assert h.hook("suspend", LAMBDA_HEADERS) == 200 + + end_job(h, "cb-1") + assert h.microvms.suspended == [] + assert h.listener._suspend_requested is None # noqa: SLF001 + + if refusal_end == "resume": + assert h.hook("resume", LAMBDA_HEADERS) == 200 + else: + h.advance(GRACE) + post_job(h, "cb-2") + end_job(h, "cb-2") + + assert h.jobs.started() == ["cb-1", "cb-2"] + assert h.microvms.suspended == [] + + +def test_a_suspend_hook_clears_a_pending_suspend_after_a_job(start: Any) -> None: + # The request waits for the other job. Lambda then suspends the MicroVM + # after the suspend hook. So the request is done, and the other job's end + # after the resume fires nothing. + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + post_job(h, "cb-2") + end_job(h, "cb-1") + assert h.listener._suspend_requested == "cb-1" # noqa: SLF001 + + assert h.hook("suspend", LAMBDA_HEADERS) == 200 + assert h.listener._suspend_requested is None # noqa: SLF001 + assert h.hook("resume", LAMBDA_HEADERS) == 200 + end_job(h, "cb-2") + + assert h.microvms.suspended == [] + + +def test_acknowledges_a_resent_running_suspend_after_job_after_the_terminate_hook( + start: Any, +) -> None: + # A known job gets the duplicate 202 while terminating, also during a + # suspend refusal. That 202 carries the acknowledgement of the first + # request. + h = start() + h.run_hook() + assert post_job(h, "cb-1", suspendAfter=True) == (202, "accepted") + assert eventually(lambda: h.started("cb-1")) + assert h.hook("suspend") == 200 + assert h.hook("terminate") == 200 + + assert post_job(h, "cb-1") == (202, "accepted") + h.jobs.release_all() + assert h.worker.idle(WAIT) + + assert h.jobs.started() == ["cb-1"] + assert h.microvms.suspended == [] + + +def test_acknowledges_a_resent_finished_suspend_after_job_during_its_own_suspend( + start: Any, +) -> None: + # The job has ended, and the worker suspends after it. A resend during + # that refusal gets the duplicate 202 with the acknowledgement. + h = start() + h.run_hook() + assert post_job(h, "cb-1", suspendAfter=True) == (202, "accepted") + end_job(h, "cb-1") + assert h.microvms.suspended == ["mvm-1"] + + assert post_job(h, "cb-1") == (202, "accepted") + assert post_job(h, "cb-2", suspendAfter=True) == (503, None) + assert h.jobs.started() == ["cb-1"] + + +def test_suspends_after_a_suspend_after_job_in_a_session_without_auto_suspend( + start: Any, logger: Any +) -> None: + h = start() + h.run_hook() + assert post_job(h, "cb-1", suspendAfter=True)[0] == 202 + h.advance(IDLE * 10) + assert h.microvms.suspended == [] + + end_job(h, "cb-1") + + assert h.microvms.suspended == ["mvm-1"] + assert info_lines(logger, JOB_SUSPEND_INFO) == [ + {"microvmId": "mvm-1", "callbackId": "cb-1"} + ] + # The refusal runs from the decision, as after an idle suspend. + assert h.job("cb-2") == 503 + h.hook("resume", LAMBDA_HEADERS) + assert h.job("cb-2") == 202 + + +def test_suspends_after_a_job_whose_handler_failed(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True, input={"fail": True}) + + end_job(h, "cb-1") + + assert h.lambda_client.recorded()[0]["kind"] == "failure" + assert h.microvms.suspended == ["mvm-1"] + + +def test_suspends_after_a_job_whose_report_failed(start: Any) -> None: + h = start() + h.run_hook() + h.lambda_client.completion_error = suspend_error("AccessDeniedException", 403) + post_job(h, "cb-1", suspendAfter=True) + + end_job(h, "cb-1") + + assert h.lambda_client.recorded() == [] + assert h.microvms.suspended == ["mvm-1"] + + +def test_suspends_only_when_the_report_attempts_have_ended(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + h.lambda_client.hold = True + + h.jobs.release("cb-1") + assert eventually(lambda: h.lambda_client.waiting() == 1) + assert h.microvms.suspended == [] + + h.lambda_client.finish() + end_job(h, "cb-1") + assert h.lambda_client.recorded() == [{"kind": "success", "callbackId": "cb-1"}] + assert h.microvms.suspended == ["mvm-1"] + + +def test_suspends_when_the_last_running_job_ends(start: Any, logger: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + post_job(h, "cb-2") + + end_job(h, "cb-1") + assert h.microvms.suspended == [] + + end_job(h, "cb-2") + assert h.microvms.suspended == ["mvm-1"] + assert info_lines(logger, JOB_SUSPEND_INFO) == [ + {"microvmId": "mvm-1", "callbackId": "cb-1"} + ] + + +def test_waits_for_a_background_failure_report(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + h.lambda_client.hold = True + # An invalid request names a callback, so the worker fails it in the + # background. + assert h.job("cb-bad", heartbeatTimeoutSeconds=0) == 400 + assert eventually(lambda: h.lambda_client.waiting() == 1) + h.lambda_client.hold = False + + end_job(h, "cb-1") + assert h.microvms.suspended == [] + + h.lambda_client.finish() + assert eventually(lambda: h.microvms.suspended == ["mvm-1"]) + + +def test_waits_for_a_run_hook_that_is_being_handled(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + later = slow_run_hook(h) + + end_job(h, "cb-1") + assert h.microvms.suspended == [] + + assert later.finish() == 200 + assert eventually(lambda: h.microvms.suspended == ["mvm-1"]) + + +def test_suspends_the_run_hook_microvm_not_the_one_of_the_job_request( + start: Any, +) -> None: + h = start() + # The job request arrives before the run hook, and names another MicroVM. + post_job(h, "cb-1", suspendAfter=True, microvmId="mvm-other") + h.run_hook() + + end_job(h, "cb-1") + + assert h.microvms.suspended == ["mvm-1"] + + +@pytest.mark.parametrize("run_hook", ["none", "without-payload"]) +def test_warns_and_does_not_suspend_when_no_run_hook_named_a_microvm( + start: Any, logger: Any, run_hook: str +) -> None: + h = start() + if run_hook == "without-payload": + assert h.post(f"{HOOK_PATH_PREFIX}run", {"microvmId": "mvm-1"}) == 200 + post_job(h, "cb-1", suspendAfter=True, microvmId="mvm-other") + + end_job(h, "cb-1") + + assert h.microvms.suspended == [] + assert [line for line in logger.lines if line[1] == NO_TARGET_WARNING] == [ + ("warning", NO_TARGET_WARNING, {"callbackId": "cb-1"}) + ] + # The request was consumed. A later job end does not warn again. + post_job(h, "cb-2") + end_job(h, "cb-2") + assert logger.messages("warning").count(NO_TARGET_WARNING) == 1 + assert h.job("cb-3") == 202 + + +def test_a_resume_hook_does_not_repeat_the_suspend(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + end_job(h, "cb-1") + assert h.microvms.suspended == ["mvm-1"] + + h.hook("resume", LAMBDA_HEADERS) + h.advance(GRACE * 10) + + assert h.microvms.suspended == ["mvm-1"] + # A job that ends after the resume asks for nothing either. + post_job(h, "cb-2") + end_job(h, "cb-2") + assert h.microvms.suspended == ["mvm-1"] + + +def test_does_not_retry_a_refusal_for_now_without_auto_suspend( + start: Any, logger: Any +) -> None: + h = start() + h.microvms.failure = suspend_error("ExpiredTokenException", 400) + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + + end_job(h, "cb-1") + h.advance(GRACE * 10) + + assert h.microvms.suspended == ["mvm-1"] + assert_one_warning(logger, "refusedForNow", refusing_jobs=False, trigger="job") + assert h.job("cb-2") == 202 + + +def test_retries_a_refusal_for_now_with_auto_suspend(start: Any) -> None: + h = start() + h.microvms.failure = suspend_error("ExpiredTokenException", 400) + h.run_hook(IDLE) + post_job(h, "cb-1", suspendAfter=True) + + end_job(h, "cb-1") + assert h.microvms.suspended == ["mvm-1"] + h.advance(GRACE - 0.001) + assert h.microvms.suspended == ["mvm-1"] + + h.advance(0.001) + assert h.microvms.suspended == ["mvm-1", "mvm-1"] + + +def test_stops_the_idle_timer_that_the_job_end_armed(start: Any, logger: Any) -> None: + # The job end arms the idle time, and fires the suspend in the same + # block. A rejected call does not retry until a job ends. So the stale + # timer must not fire one idle time later. + h = start() + h.microvms.failure = suspend_error("AccessDeniedException", 403) + h.run_hook(IDLE) + post_job(h, "cb-1", suspendAfter=True) + + end_job(h, "cb-1") + h.advance(IDLE * 10) + + assert h.microvms.suspended == ["mvm-1"] + assert_one_warning(logger, "rejected", refusing_jobs=False, trigger="job") + + +def test_logs_an_unknown_outcome_of_the_job_suspend_with_its_trigger( + start: Any, logger: Any +) -> None: + h = start() + h.microvms.failure = suspend_error("InternalServerException", 500) + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + + end_job(h, "cb-1") + + assert_one_warning(logger, "unknown", refusing_jobs=True, trigger="job") + assert h.job("cb-2") == 503 + + +def test_does_not_suspend_after_the_terminate_hook(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + post_job(h, "cb-2") + end_job(h, "cb-1") + + assert h.hook("terminate", LAMBDA_HEADERS) == 200 + assert h.listener._suspend_requested is None # noqa: SLF001 + h.jobs.release_all() + assert h.worker.idle(WAIT) + + assert h.microvms.suspended == [] + assert h.listener._suspend_requested is None # noqa: SLF001 + + +def test_does_not_suspend_after_close(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + post_job(h, "cb-2") + end_job(h, "cb-1") + + h.listener.close() + end_job(h, "cb-2") + + assert h.microvms.suspended == [] + assert h.listener._suspend_requested is None # noqa: SLF001 + + +def test_a_job_that_ends_after_close_asks_for_no_suspend(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + + h.listener.close() + end_job(h, "cb-1") + + assert h.microvms.suspended == [] + assert h.listener._suspend_requested is None # noqa: SLF001 + + +def test_fires_a_pending_suspend_when_the_running_call_ends(start: Any) -> None: + # 1. The idle call runs, and a resume hook ends its refusal. + # 2. A suspendAfter job runs and ends while the call still runs. + # 3. The request waits for the call, and fires when it ends. + h = start() + h.microvms.hold = True + h.run_hook(IDLE) + h.advance(IDLE) + assert eventually(lambda: h.microvms.suspended == ["mvm-1"]) + h.hook("resume", LAMBDA_HEADERS) + assert post_job(h, "cb-1", suspendAfter=True)[0] == 202 + h.finish_job("cb-1") + assert h.microvms.suspended == ["mvm-1"] + + h.microvms.finish() + + assert eventually(lambda: h.microvms.suspended == ["mvm-1", "mvm-1"]) + h.settle() + + +def test_an_idle_call_does_the_pending_suspend_too(start: Any) -> None: + # A background report holds the request. The idle time ends first, and + # its call suspends the MicroVM. So the request is done, and it does not + # fire again after the resume. + h = start() + h.run_hook(IDLE) + post_job(h, "cb-1", suspendAfter=True) + h.lambda_client.hold = True + assert h.job("cb-bad", heartbeatTimeoutSeconds=0) == 400 + assert eventually(lambda: h.lambda_client.waiting() == 1) + h.lambda_client.hold = False + end_job(h, "cb-1") + assert h.microvms.suspended == [] + + h.advance(IDLE) + assert h.microvms.suspended == ["mvm-1"] + h.lambda_client.finish() + assert h.worker.idle(WAIT) + h.hook("resume", LAMBDA_HEADERS) + post_job(h, "cb-2") + end_job(h, "cb-2") + + assert h.microvms.suspended == ["mvm-1"] + + +def test_idle_waits_for_the_suspend_after_a_job(start: Any) -> None: + h = start() + h.microvms.hold = True + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + + h.finish_job("cb-1") + + assert not h.worker.idle(0.1) + h.microvms.finish() + assert h.worker.idle(WAIT) + + +def test_accepts_jobs_when_the_thread_of_the_job_suspend_cannot_start( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + failing_thread_start(monkeypatch, "microvm-suspend") + + h.finish_job("cb-1") + + assert "could not start the suspend of the MicroVM after a job" in ( + logger.messages("error") + ) + assert h.microvms.suspended == [] + monkeypatch.undo() + assert h.job("cb-2") == 202 + + +def test_a_suspend_after_job_whose_thread_cannot_start_asks_for_nothing( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + h = start() + h.run_hook() + failing_thread_start(monkeypatch, "microvm-job") + + assert post_job(h, "cb-1", suspendAfter=True) == (500, None) + + monkeypatch.undo() + assert h.listener._suspend_after_jobs == set() # noqa: SLF001 + assert h.listener._suspend_requested is None # noqa: SLF001 + assert h.microvms.suspended == [] diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/payload_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/payload_test.py index 2192fd220..dcd8dcf11 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/payload_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/payload_test.py @@ -284,3 +284,32 @@ def test_an_unusable_callback_id_is_not_attached_to_errors(callback_id): # endregion dataclasses and callback IDs + + +@pytest.mark.parametrize( + ("overrides", "suspend_after"), + [ + ({"suspendAfter": True}, True), + ({"suspendAfter": False}, False), + ({"suspendAfter": "true"}, False), + ({"suspendAfter": 1}, False), + ({"suspendAfter": None}, False), + ({"suspendAfter": {}}, False), + ({}, False), + ], + ids=["true", "false", "string", "one", "null", "object", "absent"], +) +def test_job_request_reads_only_json_true_as_suspend_after(overrides, suspend_after): + # Any other value is ignored, not rejected, as in the JavaScript worker. + parsed = MicrovmJobRequest.from_dict(job_request(**overrides)) + + assert parsed.suspend_after is suspend_after + + +def test_run_hook_payload_ignores_suspend_after(): + parsed = RunHookRequest.from_dict( + run_hook({"version": 1, "region": "us-east-1", "suspendAfter": True}) + ) + + assert parsed.payload is not None + assert not hasattr(parsed.payload, "suspend_after")