From 9acabbf4372770202f311923770b1700680c43ce Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Wed, 14 Oct 2026 12:50:35 -0700 Subject: [PATCH 01/10] feat(extras): deliver large MicroVM jobs over HTTP A job whose run hook payload exceeds 4,096 code points, or any job when the config sets a request path, now goes over HTTP instead of failing with MicrovmDeliveryError. This matches the JavaScript operation. - The launch sends a run hook payload without the job, gives the MicroVM the ALL_INGRESS connector, and checkpoints the endpoint. - A new .request step creates an auth token and POSTs the job to the endpoint. Inside one step attempt it retries connection errors, timeouts, 429, and 5xx from 250 ms to 4 s, and one 401 or 403 with a new token. It stops at the retry window, 60 seconds by default, or 10 seconds before the invocation times out. - boto3 cannot cancel a call. So a CreateMicrovmAuthToken call with a deadline runs in a daemon thread, and the step stops waiting at the reserve. - The default retry strategy retries the end of the first tier, and does not retry a route that rejects the job with another 4xx. - MicrovmRequestConfig sets the path, the port, and the retry window. - Replay follows the delivery that the launch checkpointed. --- .../README.md | 53 ++- .../microvm/__init__.py | 17 +- .../microvm/config.py | 65 ++- .../microvm/job.py | 98 ++++- .../microvm/operation.py | 302 +++++++++++--- .../microvm/request.py | 377 ++++++++++++++++++ .../microvm/retry.py | 22 +- .../tests/conftest.py | 96 ++++- .../tests/microvm/job_test.py | 47 ++- .../tests/microvm/operation_test.py | 271 +++++++++++-- .../tests/microvm/request_test.py | 292 ++++++++++++++ .../tests/microvm/retry_test.py | 17 + .../tests/microvm/validation_test.py | 84 +++- .../tests/microvm/worker_composed_test.py | 37 +- 14 files changed, 1662 insertions(+), 116 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py create mode 100644 packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index 3846c44b4..be3e68eb9 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -35,14 +35,41 @@ The first argument is the context that the call runs in. Inside a child context The code in the MicroVM receives the job and reports its result. The worker package [`aws-durable-execution-sdk-python-microvm-worker`](../aws-durable-execution-sdk-python-microvm-worker) does this, and so does the JavaScript worker. The job document is the same in every SDK language. +### How the job reaches the MicroVM + +The operation picks the delivery from the input size and the config: + +1. A job whose `run` hook payload fits in 4,096 Unicode code points goes in the `RunMicrovm` request. Lambda passes it to the worker's `run` lifecycle hook. The payload also holds the callback ID, so an input of about 3,600 code points fits. The MicroVM gets the `NO_INGRESS` connector, so nothing can reach it. +2. A larger job goes over HTTP. The launch request carries no job, and a request step POSTs the job to the MicroVM's endpoint after it starts. The Python worker accepts a request body of at most 6 MiB. The MicroVM gets the `ALL_INGRESS` connector, and every request carries a short-lived auth token that the operation creates. + +The worker package runs its `handler` for both deliveries. So the code in the MicroVM does not change with the input size. + +The HTTP delivery waits inside the Lambda invocation until the MicroVM's endpoint accepts the job, usually one to two seconds. The `run` hook delivery adds no Lambda time for the boot. + +An application that already serves HTTP routes can receive every job on its own route: + +```python +import dataclasses + +from aws_durable_execution_sdk_python_extras.microvm import MicrovmRequestConfig + +config = dataclasses.replace(CONFIG, request=MicrovmRequestConfig(path="/clone-build")) +microvm(context, "build", job_input=job_input, config=config) +``` + +A `path` sends every job over HTTP, whatever its size. The route must answer with a 2xx status at once and run the job in the background. `port` names a port other than 8080. `retry_window` sets how long the request step retries inside one attempt, 60 seconds by default. + ### What it does The operation runs its durable operations in one child context named after the second argument: 1. `.callback` creates a durable callback. The service generates the callback ID, and the SDK checkpoints it. So the ID is unique to this call and stays the same on replay. -2. `.launch` calls `RunMicrovm`. The job goes in the `runHookPayload`, and Lambda passes it to the worker's `run` lifecycle hook. The client token is the SHA-256 hex digest of the callback ID. So a retried or replayed launch returns the same MicroVM instead of starting a second one. The step checkpoints plain JSON, `{"microvmId": "...", "delivery": "run-hook"}`, as the JavaScript operation does. So a reader of the execution history finds the MicroVM ID there. -3. The child context waits for the MicroVM to complete the callback. -4. `.terminate` calls `TerminateMicrovm`. It runs after success, failure, and timeout alike. +2. `.launch` calls `RunMicrovm`. The client token is the SHA-256 hex digest of the callback ID. So a retried or replayed launch returns the same MicroVM instead of starting a second one. With `run` hook delivery, the launch request contains the job. The step checkpoints plain JSON, `{"microvmId": "...", "endpoint": "...", "delivery": "http"}`, as the JavaScript operation does. The `endpoint` is present only for HTTP delivery. So a reader of the execution history finds the MicroVM ID there. +3. With HTTP delivery, `.request` creates an auth token and POSTs the job. +4. The child context waits for the MicroVM to complete the callback. +5. `.terminate` calls `TerminateMicrovm`. It runs after success, failure, and timeout alike, and after a failed delivery. + +Replay follows the delivery that the launch checkpointed, not a new choice. A deploy between the launch and a replay can change the input or the config. The launch result comes from the checkpoint, so the job still goes the way the MicroVM was launched for. A terminate failure does not fail the operation. The job's outcome is already recorded at that point. So the operation logs an error instead of raising. The MicroVM then keeps running, and is billed, until the platform ends it at its maximum duration: the operation timeout plus 5 minutes, at most 8 hours. @@ -52,8 +79,6 @@ The terminate step runs only when the operation sees the job's outcome. In these - The execution fails, times out, or is stopped while the operation waits. - Every launch attempt timed out or lost its connection, but the service created the MicroVM anyway. The operation never learned the MicroVM's ID. -The `run` hook payload holds at most 4,096 Unicode code points. It includes the callback ID and the job's fields, so an input of about 3,600 code points fits. The operation fails a larger job with `MicrovmDeliveryError` before the launch. When the job does not fit even without the callback ID, it fails before the callback exists. Delivery over HTTP, for larger jobs, comes in a later release. - ### Errors Every MicroVM failure is a `MicrovmError`. Catch the base class for all of them, or a subclass for one kind: @@ -61,7 +86,7 @@ Every MicroVM failure is a `MicrovmError`. Catch the base class for all of them, | Error | When | |---|---| | `MicrovmLaunchError` | `RunMicrovm` failed after all retries. When the last attempt timed out or lost its connection, the service may have created the MicroVM anyway, and it then runs the job. So running the job again after this error can run it twice. | -| `MicrovmDeliveryError` | The job is too large for the `run` hook. | +| `MicrovmDeliveryError` | HTTP delivery failed after all retries, or the route answered a 4xx status such as 404. | | `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. | @@ -83,15 +108,22 @@ The class is the same on the first run and on every replay, so code that branche When the error escapes the handler, the execution's error type is the class's fully-qualified name, such as `aws_durable_execution_sdk_python_extras.microvm.errors.MicrovmJobFailedError`. The SDK uses that name to rebuild the class on replay. The JavaScript operation reports the short name, `MicrovmJobFailedError`. So a rule that matches the error type of both SDKs, such as a Step Functions `Catch` or a metric filter, must match both forms. For example, match on the suffix. -### Retries +### Retries in two tiers + +A step retry costs one durable operation and one more Lambda invocation. Most failures here are short, so the steps retry in two tiers: + +1. **Inside one step attempt.** The boto3 client retries throttling and server errors of `RunMicrovm`, `CreateMicrovmAuthToken`, and `TerminateMicrovm`. The request step retries the POST itself: connection errors, timeouts, 429, and 5xx, with a delay from 250 milliseconds to 4 seconds. It retries one 401 or 403 with a new auth token. The endpoint answers this way until the MicroVM's `run` hook returns. This tier ends after `retry_window`, or 10 seconds before the invocation times out, whichever comes first. +2. **The step retry strategy.** When the first tier gives up, the step retries after a longer delay, outside the invocation. `default_microvm_retry_strategy` is the default, unless the config sets `retry_strategy`. It retries throttling, server errors, connection failures, timeouts, and the end of the first tier, for at most 5 attempts, with a delay of 2, 4, 8, and 16 seconds. Every other error, such as `ValidationException`, `AccessDeniedException`, or `InsufficientCapacityException`, fails the step at once. + +A route that answers another 4xx status fails the operation without a retry, with `MicrovmDeliveryError`. -The launch and terminate steps retry with `default_microvm_retry_strategy`, unless the config sets `retry_strategy`. It retries throttling, server errors, connection failures, and timeouts, for at most 5 attempts, with a delay of 2, 4, 8, and 16 seconds. Every other error, such as `ValidationException`, `AccessDeniedException`, or `InsufficientCapacityException`, fails the step at once. The boto3 client also retries throttling and server errors inside each attempt. +The 10-second reserve keeps time for the step to record its failure. A Lambda timeout during a step attempt records no outcome for the attempt. So every call of the first tier must end before the reserve. boto3 cannot cancel a call, so a `CreateMicrovmAuthToken` call runs in a separate thread, and the step stops waiting for it at the reserve. The POST's socket timeout is at most the time left before the reserve. ### Defaults and limits - `timeout` must be between 1 second and 8 hours, because a MicroVM lives at most 8 hours. The MicroVM's `maximumDurationInSeconds` is `timeout` plus 5 minutes, at most 8 hours. So a timeout above 7 hours 55 minutes leaves the MicroVM less than 5 minutes beyond the callback's timeout. - `job_input` is sent as JSON. So the MicroVM receives what `json.dumps` writes: a tuple arrives as a list, and a dict key that is not a string arrives as a string. -- Without connectors in the config, the MicroVM gets the `NO_INGRESS` connector, because the job arrives in the `run` hook, and the `INTERNET_EGRESS` connector, because it needs egress to reach the Lambda callback APIs. An empty tuple sends no connector. +- The MicroVM gets the `INTERNET_EGRESS` connector, because it needs egress to reach the Lambda callback APIs. For ingress, it gets `NO_INGRESS` with `run` hook delivery and `ALL_INGRESS` with HTTP delivery. If you set `ingress_network_connectors`, allow ingress whenever an input can exceed the `run` hook limit. An empty tuple sends no connector. - `heartbeat_timeout` starts when the callback is created, before the launch. So it must cover the launch, with its retries, and the boot. A throttled launch can take several minutes. With a shorter heartbeat timeout, the operation fails with `MicrovmTimeoutError` before the MicroVM sends its first heartbeat. - The Region comes from the durable execution ARN, or from `AWS_REGION`. - Without a `client`, the operation creates one `lambda-microvms` client per Region, on first use. It has a 5-second connect timeout, a 30-second read timeout, and 3 attempts in botocore's standard retry mode. So one launch or terminate attempt takes at most about 2 minutes. A call that runs past the function's timeout records no outcome for its step attempt, and the step runs again. The launch is idempotent through its client token, so it is repeated, not doubled. @@ -101,7 +133,8 @@ The launch and terminate steps retry with `default_microvm_retry_strategy`, unle The function's execution role needs: - `lambda:RunMicrovm` and `lambda:TerminateMicrovm` on the MicroVM image. Without `lambda:TerminateMicrovm`, every operation still returns its result, but every MicroVM runs until its timeout plus 5 minutes, at most 8 hours. The only sign is an error in the function's log. -- `lambda:PassNetworkConnector` on the connectors that the MicroVM uses. With the defaults, that is `arn:aws:lambda::aws:network-connector:aws-network-connector:*`. +- `lambda:CreateMicrovmAuthToken`, for HTTP delivery. Grant it unless every input is small and no config sets a `path`. +- `lambda:PassNetworkConnector` on the connectors that the MicroVM uses. `RunMicrovm` authorizes each connector separately. With the defaults, that is `arn:aws:lambda::aws:network-connector:aws-network-connector:*`. - `iam:PassRole` on the MicroVM execution role. The MicroVM execution role (`execution_role_arn`) needs `lambda:SendDurableExecutionCallbackSuccess`, `lambda:SendDurableExecutionCallbackFailure`, and `lambda:SendDurableExecutionCallbackHeartbeat` on the function ARN (`arn:aws:lambda:::function::*`). diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/__init__.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/__init__.py index 579941227..32572cbde 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/__init__.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/__init__.py @@ -6,13 +6,13 @@ This module is experimental. Its API can change in any release. The exports match the parts of ``@aws/durable-execution-sdk-js-extras/microvm`` -that this package has so far. The session operation and delivery over HTTP -come in later changes. +that this package has so far. The session operation comes in a later change. """ from aws_durable_execution_sdk_python_extras.microvm.config import ( MAX_MICROVM_DURATION_SECONDS, MicrovmConfig, + MicrovmRequestConfig, ) from aws_durable_execution_sdk_python_extras.microvm.errors import ( MicrovmDeliveryError, @@ -25,19 +25,32 @@ MAX_RUN_HOOK_PAYLOAD_LENGTH, ) from aws_durable_execution_sdk_python_extras.microvm.operation import microvm +from aws_durable_execution_sdk_python_extras.microvm.request import ( + DEFAULT_MICROVM_JOB_PATH, + DEFAULT_MICROVM_PORT, + MicrovmEndpointUnavailableError, + MicrovmRequestRejectedError, + MicrovmRequestUnauthorizedError, +) from aws_durable_execution_sdk_python_extras.microvm.retry import ( default_microvm_retry_strategy, ) __all__ = [ + "DEFAULT_MICROVM_JOB_PATH", + "DEFAULT_MICROVM_PORT", "MAX_MICROVM_DURATION_SECONDS", "MAX_RUN_HOOK_PAYLOAD_LENGTH", "MicrovmConfig", "MicrovmDeliveryError", + "MicrovmEndpointUnavailableError", "MicrovmError", "MicrovmJobFailedError", "MicrovmLaunchError", + "MicrovmRequestConfig", + "MicrovmRequestRejectedError", + "MicrovmRequestUnauthorizedError", "MicrovmTimeoutError", "default_microvm_retry_strategy", "microvm", 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 83f825494..ead8d51f9 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 @@ -24,6 +24,45 @@ ``maximumDurationInSeconds`` of at most 28,800.""" +@dataclass(frozen=True, kw_only=True) +class MicrovmRequestConfig: + """How :func:`microvm` delivers a job over HTTP. + + An HTTP delivery POSTs the job document to a route of the application in + the MicroVM. The route must answer with a 2xx status at once, and run the + job in the background. + """ + + path: str | None = None + """The route path, for example ``/clone-build``. It must start with ``/``. + + ``None`` sends the job to :data:`DEFAULT_MICROVM_JOB_PATH`. The worker + package serves that path with its ``handler``. A path sends the job to one + of your own routes instead, and it selects HTTP delivery for every input + size. The worker runs a job that arrives in the ``run`` hook with its + ``handler``, not with a route. So a job for a route must go over HTTP. + """ + + port: int | None = None + """The port of the route in the MicroVM, between 1 and 65535. + + ``None`` uses port 8080, the port that the MicroVM endpoint forwards to by + default. The port applies only when the job goes over HTTP. Lambda sends + the ``run`` hook to the port in the image's hook configuration, not to + this port. So without a ``path``, the port is used only for an input too + large for the ``run`` hook. + """ + + retry_window: Duration | None = None + """How long the request step retries inside one step attempt. + + ``None`` means 60 seconds. It must be between 1 second and 10 minutes. + After the window, the step fails the attempt, and the step retry strategy + takes over. The step also stops 10 seconds before the Lambda invocation + times out. + """ + + @dataclass(frozen=True, kw_only=True) class MicrovmConfig: """The MicroVM, and how long to wait for its result. @@ -79,9 +118,16 @@ class MicrovmConfig: ingress_network_connectors: tuple[str, ...] | None = None """Ingress network connector ARNs. - ``None`` gets the Lambda-provided ``NO_INGRESS`` connector, because the - job arrives in the ``run`` hook and needs no inbound connection. An empty - tuple sends no ingress connector. + ``None`` gets a Lambda-provided connector that depends on the delivery: + + 1. ``NO_INGRESS`` when the job goes in the ``run`` hook. The job then + needs no inbound connection. + 2. ``ALL_INGRESS`` when the job goes over HTTP. The request must reach + the MicroVM's endpoint. + + An empty tuple sends no ingress connector. Connectors that you set must + allow ingress whenever an input can exceed the ``run`` hook limit, because + such an input goes over HTTP. """ egress_network_connectors: tuple[str, ...] | None = None @@ -96,10 +142,19 @@ class MicrovmConfig: """The ``logging`` member of ``RunMicrovm``. Defaults to the service default.""" + request: MicrovmRequestConfig | None = None + """How to deliver the job over HTTP. + + ``None`` delivers a job in the ``run`` hook when it fits, and over HTTP to + :data:`DEFAULT_MICROVM_JOB_PATH` on port 8080 when it does not. + """ + retry_strategy: Callable[[Exception, int], RetryDecision] | None = None - """The step retry strategy of the launch and terminate steps. + """The step retry strategy of the launch, request, and terminate steps. - Defaults to :func:`default_microvm_retry_strategy`. + It is the second retry tier. The first tier retries inside one step + attempt: the boto3 client retries each call, and the request step retries + the POST. Defaults to :func:`default_microvm_retry_strategy`. """ client: LambdaMicrovmsClient | None = None 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 93edb8d6a..454711ccd 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,6 +36,10 @@ class RunHookPayload: Wire names: ``version``, ``region``, and ``job``, whose fields are ``callbackId``, ``heartbeatTimeoutSeconds``, and ``input``. + + 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`. """ region: str @@ -64,17 +68,91 @@ def to_json(self) -> str: TypeError: When the input is not JSON-serializable. ValueError: When the input holds ``nan`` or an infinite float. """ - job: dict[str, Any] = {"callbackId": self.callback_id} - if self.heartbeat_timeout_seconds is not None: - job["heartbeatTimeoutSeconds"] = self.heartbeat_timeout_seconds - job["input"] = self.job_input - document = {"version": PAYLOAD_VERSION, "region": self.region, "job": job} - text = json.dumps( - document, allow_nan=False, ensure_ascii=False, separators=(",", ":") + job = _job_fields( + self.callback_id, self.job_input, self.heartbeat_timeout_seconds ) - # A surrogate appears only inside a JSON string. So the escape is - # valid JSON, and it decodes to the same surrogate. - return _SURROGATE.sub(lambda match: f"\\u{ord(match.group()):04x}", text) + return dumps({"version": PAYLOAD_VERSION, "region": self.region, "job": job}) + + +@dataclass(frozen=True, kw_only=True) +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``. + + 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. + """ + + region: str + """The Region of the durable function, for the callback API calls.""" + microvm_id: str + """The MicroVM that runs the job.""" + callback_id: str + """The callback that the worker completes with the job's outcome.""" + job_input: Any + """The caller's input. It must be JSON-serializable.""" + heartbeat_timeout_seconds: int | None = None + """The heartbeat timeout, when the caller set one.""" + + def to_json(self) -> str: + """Serialize the request as compact JSON, as :meth:`RunHookPayload.to_json` does. + + Raises: + 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 + ), + } + ) + + +def run_hook_payload_without_job(region: str) -> str: + """The ``runHookPayload`` of a launch whose job goes over HTTP. + + The worker rejects a ``run`` hook payload without ``version`` and + ``region``. So the payload keeps those two fields, and omits the job. + """ + return dumps({"version": PAYLOAD_VERSION, "region": region}) + + +def dumps(document: Any) -> str: + """Serialize ``document`` as compact JSON, as ``JSON.stringify`` writes it. + + 1. Non-ASCII text stays as it is. + 2. A lone surrogate is written as a ``\\uXXXX`` escape. + 3. ``nan`` and infinite floats are rejected. + + Raises: + TypeError: When the document is not JSON-serializable. + ValueError: When the document holds ``nan`` or an infinite float. + """ + text = json.dumps( + document, allow_nan=False, ensure_ascii=False, separators=(",", ":") + ) + # A surrogate appears only inside a JSON string. So the escape is valid + # JSON, and it decodes to the same surrogate. + return _SURROGATE.sub(lambda match: f"\\u{ord(match.group()):04x}", text) + + +def _job_fields( + callback_id: str, job_input: Any, heartbeat_timeout_seconds: int | None +) -> dict[str, Any]: + job: dict[str, Any] = {"callbackId": callback_id} + if heartbeat_timeout_seconds is not None: + job["heartbeatTimeoutSeconds"] = heartbeat_timeout_seconds + job["input"] = job_input + return job _SURROGATE = re.compile("[\ud800-\udfff]") 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 84781beb9..c05e32e76 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 @@ -7,6 +7,7 @@ import hashlib import json +import math import os import threading from collections.abc import Callable @@ -35,6 +36,7 @@ MAX_MICROVM_DURATION_SECONDS, LambdaMicrovmsClient, MicrovmConfig, + MicrovmRequestConfig, ) from aws_durable_execution_sdk_python_extras.microvm.errors import ( MicrovmDeliveryError, @@ -44,9 +46,18 @@ MicrovmTimeoutError, ) from aws_durable_execution_sdk_python_extras.microvm.job import ( - MAX_RUN_HOOK_PAYLOAD_LENGTH, + JobRequest, RunHookPayload, fits_run_hook, + run_hook_payload_without_job, +) +from aws_durable_execution_sdk_python_extras.microvm.request import ( + DEFAULT_MICROVM_JOB_PATH, + DEFAULT_MICROVM_PORT, + DEFAULT_REQUEST_RETRY_WINDOW_SECONDS, + MAX_REQUEST_RETRY_WINDOW_SECONDS, + SendJobOptions, + send_job, ) from aws_durable_execution_sdk_python_extras.microvm.retry import ( default_microvm_retry_strategy, @@ -55,6 +66,7 @@ if TYPE_CHECKING: from aws_durable_execution_sdk_python.context import Callback, DurableContext + from aws_durable_execution_sdk_python.types import StepContext LAUNCH_MARGIN_SECONDS = 300 @@ -83,27 +95,47 @@ RUN_HOOK_DELIVERY = "run-hook" """The delivery in which the job goes in the ``runHookPayload`` of the launch.""" +HTTP_DELIVERY = "http" +"""The delivery in which a request step POSTs the job to the MicroVM's +endpoint after the launch.""" + +_HOOK_PATH_PREFIX = "/aws/lambda-microvms/" +"""Lambda sends the lifecycle hooks to paths under this prefix.""" + +_MAX_PORT = 65_535 + @dataclass(frozen=True, kw_only=True) class LaunchResult: """What the launch step checkpoints. - Wire names: ``microvmId`` and ``delivery``. + Wire names: ``microvmId``, ``endpoint``, and ``delivery``. The JavaScript + operation checkpoints the same document. """ microvm_id: str delivery: str """How the job reaches the MicroVM. Replay reads it from the checkpoint, so every replay follows the delivery that the launch prepared for.""" + endpoint: str | None = None + """The MicroVM's endpoint. Present only for HTTP delivery.""" @classmethod def from_dict(cls, data: dict[str, Any]) -> LaunchResult: # A checkpoint without a delivery gets an empty one. The operation # then terminates the MicroVM as for any delivery it does not know. - return cls(microvm_id=data["microvmId"], delivery=data.get("delivery", "")) + return cls( + microvm_id=data["microvmId"], + delivery=data.get("delivery", ""), + endpoint=data.get("endpoint"), + ) def to_dict(self) -> dict[str, Any]: - return {"microvmId": self.microvm_id, "delivery": self.delivery} + data: dict[str, Any] = {"microvmId": self.microvm_id} + if self.endpoint is not None: + data["endpoint"] = self.endpoint + data["delivery"] = self.delivery + return data def microvm( @@ -121,16 +153,26 @@ def microvm( 1. ``.callback`` creates a durable callback. The service generates the callback ID, and the SDK checkpoints it. So the ID is unique to this call and stays the same on every replay. - 2. ``.launch`` calls ``RunMicrovm``. The job goes in the - ``runHookPayload``, and Lambda passes it to the worker's ``run`` hook. - The client token is the SHA-256 hex digest of the callback ID. So a - retried or replayed launch returns the same MicroVM instead of starting - a second one. - 3. The child context waits on the callback. The invocation ends during the + 2. ``.launch`` calls ``RunMicrovm``. The client token is the + SHA-256 hex digest of the callback ID. So a retried or replayed launch + returns the same MicroVM instead of starting a second one. With + ``run`` hook delivery, the job goes in the ``runHookPayload``, and + Lambda passes it to the worker's ``run`` hook. + 3. With HTTP delivery, ``.request`` creates an auth token and POSTs + the job to the MicroVM's endpoint. + 4. The child context waits on the callback. The invocation ends during the wait, so no Lambda compute is billed. The worker in the MicroVM completes the callback with the job's result or failure. - 4. ``.terminate`` calls ``TerminateMicrovm``. It runs after success, - failure, and timeout alike. + 5. ``.terminate`` calls ``TerminateMicrovm``. It runs after success, + failure, and timeout alike, and after a failed delivery. + + The operation picks the delivery: + + 1. A job whose ``run`` hook payload fits in + :data:`MAX_RUN_HOOK_PAYLOAD_LENGTH` code points goes in the ``run`` + hook. The MicroVM then gets no ingress. + 2. A larger job, or any job when ``config.request.path`` is set, goes + over HTTP. The MicroVM then gets the ``ALL_INGRESS`` connector. A terminate failure does not fail the operation. The job's outcome is already recorded when terminate runs, and the MicroVM's maximum duration @@ -165,7 +207,8 @@ def microvm( TypeError: When ``job_input`` is not JSON-serializable. No durable operation starts. MicrovmLaunchError: When ``RunMicrovm`` fails after all retries. - MicrovmDeliveryError: When the job is too large for the ``run`` hook. + MicrovmDeliveryError: When the HTTP delivery fails after all retries, + or the route answers a 4xx status such as 404. 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. @@ -191,6 +234,8 @@ class _OperationScope: region: str timeout_seconds: int heartbeat_timeout_seconds: int | None + request: MicrovmRequestConfig + retry_window_seconds: int retry_strategy: Callable[[Exception, int], RetryDecision] @classmethod @@ -233,6 +278,8 @@ def create( if heartbeat_timeout_seconds is not None and heartbeat_timeout_seconds < 1: msg = f'MicroVM "{name}": heartbeat_timeout must be at least 1 second' raise ValidationError(msg) + request = config.request or MicrovmRequestConfig() + retry_window_seconds = _validate_request(name, request) try: json.dumps(job_input, allow_nan=False) except (TypeError, ValueError) as error: @@ -256,16 +303,14 @@ def create( region=region, timeout_seconds=timeout_seconds, heartbeat_timeout_seconds=heartbeat_timeout_seconds, + request=request, + retry_window_seconds=retry_window_seconds, retry_strategy=config.retry_strategy or default_microvm_retry_strategy, ) def run_job(self, child: DurableContext, job_input: Any) -> Any: """Run the operation's durable operations in its child context.""" name = self.name - # The payload without the callback ID is shorter than the full one. - # So a job that does not fit without the ID fails before the callback - # exists, and no callback is left open until its timeout. - self._require_run_hook(self._payload("", job_input), with_callback_id=False) callback: Callback = child.create_callback( name=f"{name}.callback", config=CallbackConfig( @@ -278,17 +323,36 @@ def run_job(self, child: DurableContext, job_input: Any) -> Any: serdes=JsonSerDes(), ), ) - payload = self._payload(callback.callback_id, job_input) - # The full payload contains the callback ID, so its size is known only + # The operation picks the delivery: + # + # 1. The run hook is preferred. It opens no ingress, and it needs no + # extra durable step. + # 2. A config with a request path goes over HTTP. The worker runs a + # run hook job with its handler, not with a route. + # 3. A payload larger than the run hook limit goes over HTTP. + # RunMicrovm rejects it with ValidationException, which the + # default retry strategy does not retry. + # + # The payload contains the callback ID, so the size is known only # after the callback exists. The ID comes from the checkpoint on - # replay, and the input must not change. So every replay reaches the - # same decision. - self._require_run_hook(payload, with_callback_id=True) + # replay, and the input must not change. So every replay picks the + # same delivery as the first attempt. + with_job = RunHookPayload( + region=self.region, + callback_id=callback.callback_id, + job_input=job_input, + heartbeat_timeout_seconds=self.heartbeat_timeout_seconds, + ).to_json() + over_http = self.request.path is not None or not fits_run_hook(with_job) + delivery = HTTP_DELIVERY if over_http else RUN_HOOK_DELIVERY + payload = run_hook_payload_without_job(self.region) if over_http else with_job try: launched = LaunchResult.from_dict( child.step( - lambda _: self._launch(callback.callback_id, payload).to_dict(), + lambda _: self._launch( + callback.callback_id, payload, delivery + ).to_dict(), name=f"{name}.launch", config=StepConfig( retry_strategy=self.retry_strategy, @@ -301,15 +365,17 @@ def run_job(self, child: DurableContext, job_input: Any) -> Any: ) ) except StepError as error: - # An exception without a message checkpoints none. Its type is - # then the only detail. DurableOperationError always sets a type. - detail = error.message or error.error_type - msg = f'MicroVM "{name}": the launch failed: {detail}' + msg = f'MicroVM "{name}": the launch failed: {_detail(error)}' raise MicrovmLaunchError( message=msg, data=error.data, stack_trace=error.stack_trace ) from error microvm_id = launched.microvm_id - if launched.delivery != RUN_HOOK_DELIVERY: + # Replay computes the delivery again from the current input and + # config. A deploy between the launch and a replay can change both. + # The launch result comes from the checkpoint on replay, and it + # records the delivery that the launch prepared for. So the operation + # follows the launch's delivery, not the new computation. + if launched.delivery not in {RUN_HOOK_DELIVERY, HTTP_DELIVERY}: # The launch checkpoint comes from a release with a delivery that # this release does not know, for example after a rollback. This # release cannot deliver the job. So it terminates the MicroVM and @@ -321,6 +387,16 @@ def run_job(self, child: DurableContext, job_input: Any) -> Any: ) raise MicrovmDeliveryError(msg) + if launched.delivery == HTTP_DELIVERY: + try: + self._deliver(child, launched, callback.callback_id, job_input) + except StepError as error: + self._terminate(child, microvm_id) + msg = f'MicroVM "{name}": the job could not be delivered: {_detail(error)}' + raise MicrovmDeliveryError( + message=msg, data=error.data, stack_trace=error.stack_trace + ) from error + try: result = callback.result() except Exception as error: @@ -331,29 +407,69 @@ def run_job(self, child: DurableContext, job_input: Any) -> Any: self._terminate(child, microvm_id) return result - def _payload(self, callback_id: str, job_input: Any) -> str: - return RunHookPayload( - region=self.region, - callback_id=callback_id, - job_input=job_input, - heartbeat_timeout_seconds=self.heartbeat_timeout_seconds, - ).to_json() + def _deliver( + self, + child: DurableContext, + launched: LaunchResult, + callback_id: str, + job_input: Any, + ) -> None: + """Run the request step: POST the job to the MicroVM's endpoint. - def _require_run_hook(self, payload: str, *, with_callback_id: bool) -> None: - """Raise ``MicrovmDeliveryError`` when ``payload`` is too large.""" - if fits_run_hook(payload): - return - size = f"{len(payload)}" if with_callback_id else f"at least {len(payload)}" - msg = ( - f'MicroVM "{self.name}": the job is {size} code points in the run ' - f"hook payload, and the run hook accepts at most " - f"{MAX_RUN_HOOK_PAYLOAD_LENGTH}. Delivery over HTTP for larger jobs " - "is not available yet." + The step retries inside each attempt with :func:`send_job`, the first + tier. The step retry strategy is the second tier. + """ + body = ( + JobRequest( + region=self.region, + microvm_id=launched.microvm_id, + callback_id=callback_id, + job_input=job_input, + heartbeat_timeout_seconds=self.heartbeat_timeout_seconds, + ) + .to_json() + .encode() + ) + request = self.request + remaining_seconds = _remaining_seconds(child) + + def attempt(step_context: StepContext) -> None: + # The launch records an endpoint for every HTTP delivery. So a + # missing endpoint is a defect, not a transient failure. The + # default retry strategy does not retry RuntimeError. + if launched.endpoint is None: + msg = ( + f"the launch recorded no endpoint for MicroVM " + f"{launched.microvm_id}, so the job cannot be sent over HTTP" + ) + raise RuntimeError(msg) + send_job( + SendJobOptions( + client=self._client(), + microvm_id=launched.microvm_id, + endpoint=launched.endpoint, + 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, + body=body, + retry_window_seconds=self.retry_window_seconds, + remaining_seconds=remaining_seconds, + logger=step_context.logger, + ) + ) + + child.step( + attempt, + name=f"{self.name}.request", + config=StepConfig(retry_strategy=self.retry_strategy), ) - raise MicrovmDeliveryError(msg) - def _launch(self, callback_id: str, payload: str) -> LaunchResult: + def _launch(self, callback_id: str, payload: str, delivery: str) -> LaunchResult: """Call ``RunMicrovm``, and return what the launch step checkpoints.""" + over_http = delivery == HTTP_DELIVERY config = self.config request: dict[str, Any] = { "imageIdentifier": config.image_identifier, @@ -363,8 +479,10 @@ def _launch(self, callback_id: str, payload: str) -> LaunchResult: "clientToken": hashlib.sha256(callback_id.encode()).hexdigest(), # An explicit empty tuple stays empty. The service accepts 0 to # 10 connectors, and only None gets the default. + # An HTTP delivery must reach the MicroVM's endpoint, and a run + # hook delivery needs no inbound connection. "ingressNetworkConnectors": list( - (self._connector_arn("NO_INGRESS"),) + (self._connector_arn("ALL_INGRESS" if over_http else "NO_INGRESS"),) if config.ingress_network_connectors is None else config.ingress_network_connectors ), @@ -388,7 +506,15 @@ def _launch(self, callback_id: str, payload: str) -> LaunchResult: if not microvm_id: msg = "RunMicrovm returned no microvmId" raise RuntimeError(msg) - return LaunchResult(microvm_id=microvm_id, delivery=RUN_HOOK_DELIVERY) + endpoint = response.get("endpoint") + if over_http and not endpoint: + msg = "RunMicrovm returned no endpoint" + raise RuntimeError(msg) + return LaunchResult( + microvm_id=microvm_id, + delivery=delivery, + endpoint=endpoint if over_http else None, + ) def _terminate(self, child: DurableContext, microvm_id: str) -> None: """Run the terminate step, and log a final failure instead of raising. @@ -548,3 +674,81 @@ def _partition_and_region(arn: str) -> tuple[str, str]: if len(parts) > 3 and parts[0] == "arn" and parts[1] and parts[3]: return parts[1], parts[3] return "aws", os.environ.get("AWS_REGION", "") + + +def _validate_request(name: str, request: MicrovmRequestConfig) -> int: + """Validate the HTTP delivery settings. + + Returns: + The retry window in seconds. + + Raises: + ValidationError: When a setting is not valid. + """ + path = request.path + if path is not None and (not isinstance(path, str) or not path.startswith("/")): + msg = ( + f'MicroVM "{name}": the request path must be a string that starts with "/"' + ) + raise ValidationError(msg) + if path is not None and path.startswith(_HOOK_PATH_PREFIX): + msg = ( + f'MicroVM "{name}": the request path must not use the lifecycle hook ' + f"prefix {_HOOK_PATH_PREFIX}" + ) + raise ValidationError(msg) + port = request.port + # bool is an int subclass, and True would pass as port 1. + if port is not None and ( + isinstance(port, bool) + or not isinstance(port, int) + or not 1 <= port <= _MAX_PORT + ): + msg = f'MicroVM "{name}": the request port must be an integer between 1 and {_MAX_PORT}' + raise ValidationError(msg) + if request.retry_window is None: + return DEFAULT_REQUEST_RETRY_WINDOW_SECONDS + seconds = request.retry_window.to_seconds() + # A window over 10 minutes would outlive the auth token and its one + # refresh. + if not 1 <= seconds <= MAX_REQUEST_RETRY_WINDOW_SECONDS: + msg = ( + f'MicroVM "{name}": the request retry_window must be between 1 second ' + f"and {MAX_REQUEST_RETRY_WINDOW_SECONDS // 60} minutes" + ) + raise ValidationError(msg) + return seconds + + +def _remaining_seconds(context: DurableContext) -> Callable[[], float | None]: + """Return a reader of the invocation's remaining time, in seconds. + + The first retry tier uses it to stop before the invocation times out. The + Lambda context has ``get_remaining_time_in_millis()``. Another compute's + context may lack it, or the method may raise or return no finite number. + The reader returns ``None`` in each of these cases. So the request step + falls back to its retry window, and nothing fails. + """ + lambda_context = context.lambda_context + read = getattr(lambda_context, "get_remaining_time_in_millis", None) + if not callable(read): + return lambda: None + + def remaining() -> float | None: + try: + value = read() + except Exception: # noqa: BLE001 - any failure means no deadline + return None + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + if not math.isfinite(value): + return None + return value / 1000 + + return remaining + + +def _detail(error: StepError) -> str: + # An exception without a message checkpoints none. Its type is then the + # only detail. DurableOperationError always sets a type. + return error.message or error.error_type 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 new file mode 100644 index 000000000..1cc074ff8 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py @@ -0,0 +1,377 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Delivery of one job to a MicroVM route over HTTP. + +This module is the first retry tier of the request step. It runs inside one +step attempt. The step retry strategy is the second tier. +""" + +from __future__ import annotations + +import http.client +import threading +import time +import urllib.error +import urllib.request +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, TypeVar + +from aws_durable_execution_sdk_python.types import LoggerInterface + + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python_extras.microvm.config import ( + LambdaMicrovmsClient, + ) + + +DEFAULT_MICROVM_PORT = 8080 +"""The MicroVM endpoint forwards to this port when no port header is sent.""" + +DEFAULT_MICROVM_JOB_PATH = "/durable-execution/v1/job" +"""The path that receives a job over HTTP when the config names no route. + +The worker package serves this path with its ``handler``. So the same handler +receives a job whether it arrives in the ``run`` hook or over HTTP. The value +must equal ``MICROVM_JOB_PATH`` in ``aws-durable-execution-sdk-python-microvm-worker`` +and ``DEFAULT_MICROVM_JOB_PATH`` in the JavaScript SDK. +""" + +DEFAULT_REQUEST_RETRY_WINDOW_SECONDS = 60 +"""The default length of the first retry tier.""" + +MAX_REQUEST_RETRY_WINDOW_SECONDS = 600 +"""The longest retry window. A longer window would outlive the auth token +and its one refresh.""" + +INVOCATION_RESERVE_SECONDS = 10.0 +"""The first retry tier stops this long before the invocation times out. + +A Lambda timeout during a step attempt records no outcome for the attempt. +So the step keeps time to raise and to checkpoint the second-tier retry. +""" + +MIN_CALL_SECONDS = 1.0 +"""A call that would start with less than this much time before the reserve +does not start. A shorter call is unlikely to finish, and the step retries +the job in a later invocation anyway.""" + +_INITIAL_DELAY_SECONDS = 0.25 +_MAX_DELAY_SECONDS = 4.0 +_REQUEST_TIMEOUT_SECONDS = 10.0 +_AUTH_TOKEN_MINUTES = 5 +_PORT_HEADER = "X-aws-proxy-port" + +_T = TypeVar("_T") + + +class MicrovmRequestRejectedError(Exception): + """The endpoint rejected the job with a status that a retry cannot change. + + Examples are 400, 404, and 413. The default step retry strategy does not + retry it. + """ + + def __init__(self, message: str, status: int) -> None: + super().__init__(message) + self.status = status + + +class MicrovmRequestUnauthorizedError(Exception): + """The endpoint refused the job with 401 or 403 after a new auth token. + + The default step retry strategy retries it, outside the invocation. + """ + + def __init__(self, message: str, status: int) -> None: + super().__init__(message) + self.status = status + + +class MicrovmEndpointUnavailableError(Exception): + """The first retry tier ended without an accepted request. + + The default step retry strategy retries it, outside the invocation. + """ + + +def reserve_starts_at( + remaining_seconds: Callable[[], float | None], +) -> float | None: + """Return the ``time.monotonic()`` time at which the reserve starts. + + Every call of the first tier must end by then. ``None`` means that the + compute reports no deadline. + """ + remaining = remaining_seconds() + if remaining is None: + return None + return time.monotonic() + remaining - INVOCATION_RESERVE_SECONDS + + +def _call_budget( + hard_deadline: float | None, + cap: float, + on_expired: Callable[[], Exception], +) -> float: + """Return how long the next call may run. + + It is at most ``cap``, and it never runs past ``hard_deadline``. + + Raises: + Exception: The result of ``on_expired``, when less than + :data:`MIN_CALL_SECONDS` remains before ``hard_deadline``. + """ + if hard_deadline is None: + return cap + left = hard_deadline - time.monotonic() + if left < MIN_CALL_SECONDS: + raise on_expired() + return min(cap, left) + + +def call_before_deadline( + call: Callable[[], _T], + hard_deadline: float | None, + on_expired: Callable[[], Exception], +) -> _T: + """Run one boto3 call that must end before ``hard_deadline``. + + 1. Without a deadline, the call runs with the client's own settings. + 2. With a deadline, the call runs in a daemon thread, and this function + waits for it until the deadline. boto3 cannot cancel a call, and its + own retries can run longer than the time that remains. + 3. A call that has not ended at the deadline raises the result of + ``on_expired``. That error is retryable, so the step retries the job + in a later invocation. The thread keeps running until the call + returns. Its result is discarded. + + Raises: + Exception: The exception that the call raised, or the result of + ``on_expired``. + """ + if hard_deadline is None: + return call() + budget = _call_budget(hard_deadline, float("inf"), on_expired) + outcome: dict[str, Any] = {} + done = threading.Event() + + def run() -> None: + try: + outcome["value"] = call() + except Exception as error: # noqa: BLE001 - re-raised by the caller below + outcome["error"] = error + finally: + done.set() + + threading.Thread(target=run, name="microvm-call", daemon=True).start() + if not done.wait(budget): + raise on_expired() + if "error" in outcome: + raise outcome["error"] + return outcome["value"] + + +@dataclass(frozen=True, kw_only=True) +class SendJobOptions: + """The inputs of :func:`send_job`.""" + + client: LambdaMicrovmsClient + microvm_id: str + endpoint: str + path: str + port: int + body: bytes + retry_window_seconds: float + """The first-tier retry window.""" + remaining_seconds: Callable[[], float | None] + """The remaining invocation time. ``None`` means that the compute reports + no deadline, and only the retry window applies.""" + logger: LoggerInterface + sleep: Callable[[float], None] = time.sleep + + +def send_job(options: SendJobOptions) -> int: + """Deliver one job to a MicroVM route, with the first retry tier. + + The function handles each failure this way: + + 1. A connection error, a timeout, 429, or any 5xx: wait and retry. The + endpoint answers this way while the MicroVM starts. + 2. 401 or 403: create a new auth token and retry once. A second 401 or + 403 raises :class:`MicrovmRequestUnauthorizedError`, and the step + retry strategy takes over. + 3. Any other 4xx: raise :class:`MicrovmRequestRejectedError` at once. + + The delay starts at 250 milliseconds and doubles up to 4 seconds. The + tier ends at the retry window, or 10 seconds before the invocation times + out, whichever comes first. It then raises + :class:`MicrovmEndpointUnavailableError`. + + The two limits differ: + + 1. The retry window limits when an attempt may start. An attempt that + starts inside the window may end after it. + 2. The invocation reserve limits when every call must end. Each request + and each ``CreateMicrovmAuthToken`` call stops at the reserve. A call + that would start with less than 1 second before the reserve does not + start. A Lambda timeout during the attempt would record no outcome, + and the reserve prevents that. + + The endpoint can accept a request before Lambda has sent the ``run`` + hook. So the body carries the MicroVM ID too. + + A request can reach the MicroVM while its response is lost. The next + request then delivers the same job again. The worker ignores a second + request for a callback ID that runs or recently ended, so the job still + runs once. + + Returns: + The status that accepted the job. + + Raises: + MicrovmRequestRejectedError: When the endpoint rejects the job. + MicrovmRequestUnauthorizedError: When a new auth token is refused. + MicrovmEndpointUnavailableError: When the tier ends. + Exception: When ``CreateMicrovmAuthToken`` fails. + """ + started = time.monotonic() + window_ends = started + options.retry_window_seconds + attempt = 0 + last_failure = "no attempt" + + def unavailable() -> Exception: + elapsed_ms = round((time.monotonic() - started) * 1000) + 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}." + ) + + def deadline() -> float: + reserve = reserve_starts_at(options.remaining_seconds) + return window_ends if reserve is None else min(window_ends, reserve) + + def new_token() -> Mapping[str, str]: + return call_before_deadline( + lambda: _create_auth_token(options), + reserve_starts_at(options.remaining_seconds), + unavailable, + ) + + url = _endpoint_url(options.endpoint, options.path) + token = new_token() + refreshed_token = False + delay = _INITIAL_DELAY_SECONDS + + while True: + attempt += 1 + # The request must end before the reserve, like every call in the + # tier. + timeout = _call_budget( + reserve_starts_at(options.remaining_seconds), + _REQUEST_TIMEOUT_SECONDS, + unavailable, + ) + headers = {"content-type": "application/json", **token} + if options.port != DEFAULT_MICROVM_PORT: + headers[_PORT_HEADER] = str(options.port) + status, failure = _post(url, options.body, headers, timeout) + + if status is not None and 200 <= status < 300: + options.logger.info( + "MicroVM job accepted", extra={"attempt": attempt, "status": status} + ) + return status + if status is None: + last_failure = failure + else: + last_failure = f"HTTP {status}" + if status in {401, 403}: + if refreshed_token: + msg = ( + f"The MicroVM refused a new auth token at {options.path} " + f"with HTTP {status}" + ) + raise MicrovmRequestUnauthorizedError(msg, status) + refreshed_token = True + # The refresh checks the reserve like every other call. The + # next request then checks it again. + token = new_token() + options.logger.info( + "MicroVM auth token refreshed", + extra={"attempt": attempt, "status": status}, + ) + continue + if 400 <= status < 500 and status != 429: # noqa: PLR2004 + msg = ( + f"The MicroVM rejected the job at {options.path} with HTTP {status}" + ) + raise MicrovmRequestRejectedError(msg, status) + + if time.monotonic() + delay > deadline(): + raise unavailable() + options.logger.info( + "MicroVM job not accepted yet", + extra={"attempt": attempt, "failure": last_failure}, + ) + options.sleep(delay) + delay = min(delay * 2, _MAX_DELAY_SECONDS) + + +def _create_auth_token(options: SendJobOptions) -> Mapping[str, str]: + response = options.client.create_microvm_auth_token( + microvmIdentifier=options.microvm_id, + expirationInMinutes=_AUTH_TOKEN_MINUTES, + allowedPorts=[{"port": options.port}], + ) + token = response.get("authToken") + if not token: + msg = "CreateMicrovmAuthToken returned no authToken" + raise RuntimeError(msg) + return token + + +# Proxy settings from the environment do not apply. The request goes to the +# MicroVM endpoint directly, as the JavaScript SDK's fetch sends it. +_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({})) + + +def _post( + url: str, body: bytes, headers: Mapping[str, str], timeout: float +) -> tuple[int | None, str]: + """POST ``body``, and return the status, or ``None`` and the failure. + + The body of the response is not used. So the function reads no body, and + closes the response. + + ``timeout`` applies to each socket operation: the connect, and each send + and receive. The endpoint answers at once, so it bounds the request. + """ + request = urllib.request.Request( # noqa: S310 - the URL is https or the endpoint's own + url, data=body, headers=dict(headers), method="POST" + ) + try: + with _OPENER.open(request, timeout=timeout) as response: + return response.status, "" + except urllib.error.HTTPError as error: + error.close() + return error.code, "" + 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. + return None, f"{type(error).__name__}: {error}" + + +def _endpoint_url(endpoint: str, path: str) -> str: + """Join the endpoint and the path. + + ``RunMicrovm`` returns the endpoint as a host name without a scheme. + """ + base = ( + endpoint + if endpoint.startswith(("https://", "http://")) + else (f"https://{endpoint}") + ) + return f"{base.rstrip('/')}{path}" diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/retry.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/retry.py index 41fc18a05..08da15ccb 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/retry.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/retry.py @@ -10,6 +10,11 @@ from botocore.exceptions import ClientError, HTTPClientError from botocore.exceptions import ConnectionError as BotocoreConnectionError +from aws_durable_execution_sdk_python_extras.microvm.request import ( + MicrovmEndpointUnavailableError, + MicrovmRequestUnauthorizedError, +) + MAX_ATTEMPTS = 5 """Attempts of one step, the first one included.""" @@ -52,15 +57,20 @@ # retry mode treats the same two bases as transient. BotocoreConnectionError, HTTPClientError, + # The request step's first retry tier ended. A later attempt runs after + # a longer delay, outside the invocation. + MicrovmEndpointUnavailableError, + MicrovmRequestUnauthorizedError, ) def default_microvm_retry_strategy(error: Exception, attempt: int) -> RetryDecision: """Retry a step that failed for a reason that a later attempt can fix. - The launch and terminate steps use this strategy unless the config sets - ``retry_strategy``. It is the second retry tier: the boto3 client already - retried throttling and server errors inside the attempt. + The launch, request, and terminate steps use this strategy unless the + config sets ``retry_strategy``. It is the second retry tier. Inside each + attempt, the boto3 client already retried throttling and server errors, + and the request step already retried the POST. An error is retryable when it is one of these: @@ -70,6 +80,12 @@ def default_microvm_retry_strategy(error: Exception, attempt: int) -> RetryDecis ``InsufficientCapacityException``. 3. A connection failure or a timeout, including botocore's ``HTTPClientError`` family. + 4. The end of the request step's first retry tier: + ``MicrovmEndpointUnavailableError``, or + ``MicrovmRequestUnauthorizedError`` after a refused new token. + + ``MicrovmRequestRejectedError`` is not retried. The route rejected the + job with a 4xx status, and the same job gets the same answer. Every other error fails the step at once. A ``ValidationException``, for example, fails the same way on every attempt. 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 b8a62e80f..94ae9cab7 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py @@ -5,8 +5,11 @@ from __future__ import annotations +import json +import socket import threading -from collections.abc import Callable +from collections.abc import Callable, Iterator +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Any import pytest @@ -28,21 +31,29 @@ def client_error(code: str, status: int = 400) -> ClientError: class FakeMicrovmsClient: """Records Lambda MicroVMs API calls, and answers from scripted lists. - ``launches`` and ``terminates`` hold the answers of ``run_microvm`` and - ``terminate_microvm``, in order. Each answer is an exception to raise, a - function that receives the request, or ``None`` for success. After a list - ends, every call succeeds. + ``launches``, ``tokens``, and ``terminates`` hold the answers of + ``run_microvm``, ``create_microvm_auth_token``, and ``terminate_microvm``, + in order. Each answer is an exception to raise, a function that receives + the request, or ``None`` for success. After a list ends, every call + succeeds. + + ``run_microvm`` returns ``endpoint``. The Nth auth token is + ``{"x-test-token": "token-N"}``. """ def __init__( self, *, launches: list[Answer] | None = None, + tokens: list[Answer] | None = None, terminates: list[Answer] | None = None, + endpoint: str | None = "mvm-1.microvm.example", ) -> None: self.calls: list[tuple[str, dict[str, Any]]] = [] + self.endpoint = endpoint self._answers: dict[str, list[Answer]] = { "run_microvm": list(launches or []), + "create_microvm_auth_token": list(tokens or []), "terminate_microvm": list(terminates or []), } self._lock = threading.Lock() @@ -64,7 +75,15 @@ def _answer(self, name: str, request: dict[str, Any]) -> Any: def run_microvm(self, **request: Any) -> dict[str, Any]: self._answer("run_microvm", request) - return {"microvmId": "mvm-1", "state": "PENDING"} + response: dict[str, Any] = {"microvmId": "mvm-1", "state": "PENDING"} + if self.endpoint is not None: + response["endpoint"] = self.endpoint + return response + + def create_microvm_auth_token(self, **request: Any) -> dict[str, Any]: + self._answer("create_microvm_auth_token", request) + count = len(self.requests("create_microvm_auth_token")) + return {"authToken": {"x-test-token": f"token-{count}"}} def terminate_microvm(self, **request: Any) -> dict[str, Any]: self._answer("terminate_microvm", request) @@ -91,3 +110,68 @@ def fake_client() -> type[FakeMicrovmsClient]: @pytest.fixture def make_error() -> Callable[..., ClientError]: return client_error + + +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. + """ + + def __init__(self, statuses: list[int | None] | None = None) -> None: + self.requests: list[dict[str, Any]] = [] + self._statuses = list(statuses or []) + self._lock = threading.Lock() + endpoint = self + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: # noqa: N802 - the http.server name + length = int(self.headers.get("content-length", "0")) + body = self.rfile.read(length) + with endpoint._lock: + endpoint.requests.append( + { + "path": self.path, + "headers": {k.lower(): v for k, v in self.headers.items()}, + "body": json.loads(body), + } + ) + status = endpoint._statuses.pop(0) if endpoint._statuses else 202 + if status is None: + self.close_connection = True + self.connection.shutdown(socket.SHUT_RDWR) + return + self.send_response(status) + self.send_header("content-length", "0") + self.end_headers() + + def log_message(self, *args: Any) -> None: + pass + + self._server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + self._thread.start() + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self._server.server_address[1]}" + + def close(self) -> None: + self._server.shutdown() + self._server.server_close() + + +@pytest.fixture +def fake_endpoint() -> Iterator[Callable[..., FakeEndpoint]]: + endpoints: list[FakeEndpoint] = [] + + def _create(statuses: list[int | None] | None = None) -> FakeEndpoint: + endpoint = FakeEndpoint(statuses) + endpoints.append(endpoint) + return endpoint + + yield _create + for endpoint in endpoints: + endpoint.close() 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 4f841ce00..1bc2df308 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 @@ -9,12 +9,17 @@ import math import pytest -from aws_durable_execution_sdk_python_microvm_worker import RunHookRequest +from aws_durable_execution_sdk_python_microvm_worker import ( + MicrovmJobRequest, + RunHookRequest, +) from aws_durable_execution_sdk_python_extras.microvm.job import ( MAX_RUN_HOOK_PAYLOAD_LENGTH, + JobRequest, RunHookPayload, fits_run_hook, + run_hook_payload_without_job, ) @@ -93,3 +98,43 @@ def test_the_worker_reads_the_payload() -> None: {"n": 1}, 30.0, ) + + +def test_writes_the_job_request_with_the_job_fields_at_the_top_level() -> None: + body = JobRequest( + region="us-east-1", + microvm_id="mvm-1", + callback_id="cb-1", + job_input={"n": "日"}, + heartbeat_timeout_seconds=30, + ).to_json() + + assert body == ( + '{"version":1,"region":"us-east-1","microvmId":"mvm-1",' + '"callbackId":"cb-1","heartbeatTimeoutSeconds":30,"input":{"n":"日"}}' + ) + + +def test_the_worker_reads_the_job_request() -> None: + body = JobRequest( + region="us-east-1", microvm_id="mvm-1", callback_id="cb-1", job_input=[1] + ).to_json() + + request = MicrovmJobRequest.from_dict(json.loads(body)) + + assert request.microvm_id == "mvm-1" + assert request.region == "us-east-1" + assert (request.job.callback_id, request.job.input) == ("cb-1", [1]) + assert request.job.heartbeat_timeout_seconds is None + + +def test_the_worker_reads_a_run_hook_payload_without_a_job() -> None: + payload = run_hook_payload_without_job("us-east-1") + + request = RunHookRequest.from_dict( + {"microvmId": "mvm-1", "runHookPayload": payload} + ) + + assert payload == '{"version":1,"region":"us-east-1"}' + assert request.payload is not None + assert request.payload.job is None diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py index ada06b829..04cf6839f 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py @@ -34,6 +34,8 @@ MicrovmConfig, MicrovmError, MicrovmJobFailedError, + MicrovmRequestConfig, + default_microvm_retry_strategy, microvm, operation, ) @@ -358,44 +360,208 @@ def run_microvm(self, **request: Any) -> dict[str, Any]: assert "RunMicrovm returned no microvmId" in error["message"] -def test_fails_a_job_too_large_for_the_run_hook_before_the_launch( - run: Any, fake_client: Any +def test_sends_a_job_too_large_for_the_run_hook_over_http( + run: Any, fake_client: Any, fake_endpoint: Any ) -> None: - client = fake_client() + endpoint = fake_endpoint() + client = fake_client(endpoint=endpoint.url) + job_input = "x" * MAX_RUN_HOOK_PAYLOAD_LENGTH + config = config_with(client, heartbeat_timeout=Duration.from_seconds(60)) - result = run(catching(config_with(client)), "x" * MAX_RUN_HOOK_PAYLOAD_LENGTH) + result = run(catching(config), job_input, succeed_with("done")) - error = outcome(result) - assert error["error"] == "MicrovmDeliveryError" - assert "the job is at least " in error["message"] - assert "the run hook accepts at most 4096" in error["message"] - assert client.calls == [] - # The job does not fit even without the callback ID. So no callback is - # left open until its timeout. - assert result.get_context("build").child_operations == [] + assert outcome(result) == {"result": "done"} + [launch] = client.requests("run_microvm") + # The worker rejects a run hook payload without a version and a Region. + assert json.loads(launch["runHookPayload"]) == { + "version": 1, + "region": "us-west-2", + } + connector = "arn:aws:lambda:us-west-2:aws:network-connector:aws-network-connector" + assert launch["ingressNetworkConnectors"] == [f"{connector}:ALL_INGRESS"] + assert launch["egressNetworkConnectors"] == [f"{connector}:INTERNET_EGRESS"] + [posted] = endpoint.requests + callback_id = posted["body"]["callbackId"] + assert posted["path"] == "/durable-execution/v1/job" + assert posted["headers"]["x-test-token"] == "token-1" + assert posted["body"] == { + "version": 1, + "region": "us-west-2", + "microvmId": "mvm-1", + "callbackId": callback_id, + "heartbeatTimeoutSeconds": 60, + "input": job_input, + } + assert launch["clientToken"] == hashlib.sha256(callback_id.encode()).hexdigest() + assert client.requests("create_microvm_auth_token") == [ + { + "microvmIdentifier": "mvm-1", + "expirationInMinutes": 5, + "allowedPorts": [{"port": 8080}], + } + ] + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + context = result.get_context("build") + assert [operation.name for operation in context.child_operations] == [ + "build.callback", + "build.launch", + "build.request", + "build.terminate", + ] + launch_step = context.get_step("build.launch") + assert launch_step.result is not None + assert json.loads(launch_step.result) == { + "microvmId": "mvm-1", + "endpoint": endpoint.url, + "delivery": "http", + } -def test_fails_a_job_that_fits_only_without_the_callback_id( - run: Any, fake_client: Any +def test_picks_the_delivery_by_the_payload_with_the_callback_id( + run: Any, fake_client: Any, fake_endpoint: Any ) -> None: - client = fake_client() + endpoint = fake_endpoint() + client = fake_client(endpoint=endpoint.url) empty = RunHookPayload( region="us-west-2", callback_id="", job_input="", heartbeat_timeout_seconds=None, ).to_json() + # The payload fits only without the callback ID. job_input = "x" * (MAX_RUN_HOOK_PAYLOAD_LENGTH - len(empty)) - result = run(catching(config_with(client)), job_input) + result = run(catching(config_with(client)), job_input, succeed_with("done")) + + assert outcome(result) == {"result": "done"} + assert len(endpoint.requests) == 1 + assert "job" not in json.loads(client.requests("run_microvm")[0]["runHookPayload"]) + + +def test_sends_every_job_over_http_to_a_configured_route( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint() + client = fake_client(endpoint=endpoint.url) + config = config_with( + client, request=MicrovmRequestConfig(path="/clone-build", port=3000) + ) + + result = run(catching(config), {"repo": "org/app"}, succeed_with("done")) + + assert outcome(result) == {"result": "done"} + [posted] = endpoint.requests + assert posted["path"] == "/clone-build" + assert posted["headers"]["x-aws-proxy-port"] == "3000" + assert posted["body"]["input"] == {"repo": "org/app"} + assert client.requests("create_microvm_auth_token")[0]["allowedPorts"] == [ + {"port": 3000} + ] + + +def test_keeps_the_configured_ingress_for_http( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint() + client = fake_client(endpoint=endpoint.url) + config = config_with( + client, + ingress_network_connectors=("arn:ingress",), + request=MicrovmRequestConfig(path="/job"), + ) + + run(catching(config), {}, succeed_with(None)) + + assert client.requests("run_microvm")[0]["ingressNetworkConnectors"] == [ + "arn:ingress" + ] + + +def test_fails_with_a_delivery_error_when_the_route_rejects_the_job( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([404]) + client = fake_client(endpoint=endpoint.url) + config = config_with(client, request=MicrovmRequestConfig(path="/missing")) + + result = run(catching(config), {}) + + assert outcome(result) == { + "error": "MicrovmDeliveryError", + "message": 'MicroVM "build": the job could not be delivered: The MicroVM ' + "rejected the job at /missing with HTTP 404", + } + # The default strategy does not retry a rejected job. + assert len(endpoint.requests) == 1 + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +def test_retries_the_request_step_after_the_first_tier_ends( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + # In a 1-second window, the requests start at 0, 0.25, and 0.75 seconds. + # The next delay of 1 second ends past the window. So the first step + # attempt sends 3 requests and fails. + endpoint = fake_endpoint([503, 503, 503]) + client = fake_client(endpoint=endpoint.url) + failures: list[str] = [] + + def once(error: Exception, attempt: int) -> RetryDecision: + failures.append(type(error).__name__) + if attempt < 2 and default_microvm_retry_strategy(error, attempt).should_retry: + return RetryDecision.retry(Duration.from_seconds(1)) + return RetryDecision.no_retry() + + config = config_with( + client, + retry_strategy=once, + request=MicrovmRequestConfig( + path="/job", retry_window=Duration.from_seconds(1) + ), + ) + + result = run(catching(config), {}, succeed_with("done")) + + assert outcome(result) == {"result": "done"} + assert failures == ["MicrovmEndpointUnavailableError"] + assert len(endpoint.requests) == 4 + assert len(client.requests("run_microvm")) == 1 + + +def test_fails_a_launch_without_an_endpoint_for_http( + run: Any, fake_client: Any +) -> None: + client = fake_client(endpoint=None) + config = config_with(client, request=MicrovmRequestConfig(path="/job")) + + result = run(catching(config), {}) + + error = outcome(result) + assert error["error"] == "MicrovmLaunchError" + assert "RunMicrovm returned no endpoint" in error["message"] + + +def test_fails_the_request_for_a_launch_checkpoint_without_an_endpoint( + run: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + original = operation._OperationScope._launch # noqa: SLF001 + + def without_endpoint(self: Any, *args: Any) -> Any: + launched = original(self, *args) + return operation.LaunchResult( + microvm_id=launched.microvm_id, delivery=launched.delivery + ) + + monkeypatch.setattr(operation._OperationScope, "_launch", without_endpoint) # noqa: SLF001 + client = fake_client() + config = config_with(client, request=MicrovmRequestConfig(path="/job")) + + result = run(catching(config), {}) error = outcome(result) assert error["error"] == "MicrovmDeliveryError" - assert "the job is at least " not in error["message"] - assert "the run hook accepts at most 4096" in error["message"] - assert client.calls == [] - [callback] = result.get_context("build").child_operations - assert callback.name == "build.callback" + assert "the launch recorded no endpoint for MicroVM mvm-1" in error["message"] + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] def test_raises_the_same_class_after_the_handler_stops_catching( @@ -508,6 +674,55 @@ def test_requests_match_the_lambda_microvms_api_model(run: Any) -> None: assert outcome(result) == {"result": "done"} +def test_http_requests_match_the_lambda_microvms_api_model( + run: Any, fake_endpoint: Any +) -> None: + """botocore validates the RunMicrovm and CreateMicrovmAuthToken requests.""" + + endpoint = fake_endpoint() + client = boto3.session.Session().client("lambda-microvms", region_name="us-west-2") + config = config_with(client, request=MicrovmRequestConfig(path="/job")) + with Stubber(client) as stub: + stub.add_response( + "run_microvm", + { + "microvmId": "mvm-1", + "state": "PENDING", + "endpoint": endpoint.url, + "imageArn": config.image_identifier, + "imageVersion": "1", + "maximumDurationInSeconds": 600, + "startedAt": datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC), + }, + { + "imageIdentifier": config.image_identifier, + "executionRoleArn": config.execution_role_arn, + "clientToken": ANY, + "ingressNetworkConnectors": ANY, + "egressNetworkConnectors": ANY, + "maximumDurationInSeconds": 600, + "runHookPayload": ANY, + }, + ) + stub.add_response( + "create_microvm_auth_token", + {"authToken": {"x-amz-microvm-token": "t"}}, + { + "microvmIdentifier": "mvm-1", + "expirationInMinutes": 5, + "allowedPorts": [{"port": 8080}], + }, + ) + stub.add_response("terminate_microvm", {}, {"microvmIdentifier": "mvm-1"}) + + result = run(catching(config), {"n": 1}, succeed_with("done")) + + stub.assert_no_pending_responses() + client.close() + assert outcome(result) == {"result": "done"} + assert endpoint.requests[0]["headers"]["x-amz-microvm-token"] == "t" + + def test_keeps_the_class_inside_the_callers_child_context( run: Any, fake_client: Any ) -> None: @@ -575,19 +790,21 @@ def test_terminates_and_fails_for_a_delivery_that_it_does_not_know( original = operation._OperationScope._launch # noqa: SLF001 - def launch_for_http(self: Any, callback_id: str, payload: str) -> Any: - launched = original(self, callback_id, payload) - return operation.LaunchResult(microvm_id=launched.microvm_id, delivery="http") + def launch_for_another(self: Any, *args: Any) -> Any: + launched = original(self, *args) + return operation.LaunchResult( + microvm_id=launched.microvm_id, delivery="smoke-signal" + ) - monkeypatch.setattr(operation._OperationScope, "_launch", launch_for_http) # noqa: SLF001 + monkeypatch.setattr(operation._OperationScope, "_launch", launch_for_another) # noqa: SLF001 client = fake_client() result = run(catching(config_with(client)), {}) assert outcome(result) == { "error": "MicrovmDeliveryError", - "message": 'MicroVM "build": the launch prepared the delivery "http", ' - "which this release does not support", + "message": 'MicroVM "build": the launch prepared the delivery ' + '"smoke-signal", which this release does not support', } assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] 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 new file mode 100644 index 000000000..f59af8550 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the first retry tier of the HTTP delivery.""" + +from __future__ import annotations + +import logging +import threading +import time +from typing import Any + +import pytest +from aws_durable_execution_sdk_python_microvm_worker import MICROVM_JOB_PATH + +from aws_durable_execution_sdk_python_extras.microvm import ( + DEFAULT_MICROVM_JOB_PATH, + MicrovmEndpointUnavailableError, + MicrovmRequestRejectedError, + MicrovmRequestUnauthorizedError, + request, +) +from aws_durable_execution_sdk_python_extras.microvm.request import ( + SendJobOptions, + call_before_deadline, + send_job, +) + + +BODY = b'{"version":1}' + + +def options( + client: Any, endpoint: str, sleeps: list[float], **fields: Any +) -> SendJobOptions: + values: dict[str, Any] = { + "client": client, + "microvm_id": "mvm-1", + "endpoint": endpoint, + "path": "/durable-execution/v1/job", + "port": 8080, + "body": BODY, + "retry_window_seconds": 60, + "remaining_seconds": lambda: None, + "logger": logging.getLogger("test"), + # The tests record the delays instead of waiting. + "sleep": sleeps.append, + **fields, + } + return SendJobOptions(**values) + + +def test_matches_the_job_path_of_the_worker() -> None: + assert DEFAULT_MICROVM_JOB_PATH == MICROVM_JOB_PATH + + +def test_posts_the_job_with_an_auth_token(fake_client: Any, fake_endpoint: Any) -> None: + client = fake_client() + endpoint = fake_endpoint() + sleeps: list[float] = [] + + status = send_job(options(client, endpoint.url, sleeps)) + + assert status == 202 + [token_request] = client.requests("create_microvm_auth_token") + assert token_request == { + "microvmIdentifier": "mvm-1", + "expirationInMinutes": 5, + "allowedPorts": [{"port": 8080}], + } + [posted] = endpoint.requests + assert posted["path"] == "/durable-execution/v1/job" + assert posted["headers"]["content-type"] == "application/json" + assert posted["headers"]["x-test-token"] == "token-1" + # The endpoint forwards to 8080 when no port header is sent. + assert "x-aws-proxy-port" not in posted["headers"] + assert posted["body"] == {"version": 1} + assert sleeps == [] + + +def test_sends_a_port_header_for_another_port( + fake_client: Any, fake_endpoint: Any +) -> None: + client = fake_client() + endpoint = fake_endpoint() + + send_job(options(client, endpoint.url, [], port=3000)) + + assert client.requests("create_microvm_auth_token")[0]["allowedPorts"] == [ + {"port": 3000} + ] + assert endpoint.requests[0]["headers"]["x-aws-proxy-port"] == "3000" + + +def test_retries_while_the_microvm_starts(fake_client: Any, fake_endpoint: Any) -> None: + # None closes the connection without an answer. + endpoint = fake_endpoint([503, 429, 500, None]) + sleeps: list[float] = [] + + status = send_job(options(fake_client(), endpoint.url, sleeps)) + + assert status == 202 + assert len(endpoint.requests) == 5 + assert sleeps == [0.25, 0.5, 1.0, 2.0] + + +def test_caps_the_delay_at_4_seconds(fake_client: Any, fake_endpoint: Any) -> None: + endpoint = fake_endpoint([503] * 7) + sleeps: list[float] = [] + + send_job(options(fake_client(), endpoint.url, sleeps)) + + assert sleeps == [0.25, 0.5, 1.0, 2.0, 4.0, 4.0, 4.0] + + +@pytest.mark.parametrize("status", [400, 404, 413]) +def test_rejects_the_job_at_once_for_another_4xx( + fake_client: Any, fake_endpoint: Any, status: int +) -> None: + endpoint = fake_endpoint([status]) + + with pytest.raises(MicrovmRequestRejectedError) as raised: + send_job(options(fake_client(), endpoint.url, [])) + + assert raised.value.status == status + assert str(raised.value) == ( + f"The MicroVM rejected the job at /durable-execution/v1/job with HTTP {status}" + ) + assert len(endpoint.requests) == 1 + + +@pytest.mark.parametrize("status", [401, 403]) +def test_retries_once_with_a_new_auth_token( + fake_client: Any, fake_endpoint: Any, status: int +) -> None: + client = fake_client() + endpoint = fake_endpoint([status]) + sleeps: list[float] = [] + + assert send_job(options(client, endpoint.url, sleeps)) == 202 + + assert len(client.requests("create_microvm_auth_token")) == 2 + assert [r["headers"]["x-test-token"] for r in endpoint.requests] == [ + "token-1", + "token-2", + ] + # The refresh retries at once. + assert sleeps == [] + + +def test_raises_unauthorized_when_the_new_token_is_refused( + fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([401, 403]) + + with pytest.raises(MicrovmRequestUnauthorizedError) as raised: + send_job(options(fake_client(), endpoint.url, [])) + + assert raised.value.status == 403 + assert len(endpoint.requests) == 2 + + +def test_ends_the_tier_at_the_retry_window( + fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([503] * 10) + sleeps: list[float] = [] + + # The first delay fits in the window, and the second does not. + with pytest.raises(MicrovmEndpointUnavailableError) as raised: + send_job(options(fake_client(), endpoint.url, sleeps, retry_window_seconds=0.4)) + + assert sleeps == [0.25] + message = str(raised.value) + assert message.startswith( + "The MicroVM did not accept the job at /durable-execution/v1/job after " + "2 attempts in " + ) + assert message.endswith("Last failure: HTTP 503.") + + +def test_reports_a_connection_failure_as_the_last_failure( + fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([None, None]) + + with pytest.raises(MicrovmEndpointUnavailableError) as raised: + send_job(options(fake_client(), endpoint.url, [], retry_window_seconds=0.4)) + + assert "Last failure: HTTP 503" not in str(raised.value) + assert "Last failure: " in str(raised.value) + + +def test_ends_the_tier_before_the_invocation_reserve( + fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([503] * 10) + sleeps: list[float] = [] + + # The reserve starts 3 seconds from now. A 4-second delay crosses it. + with pytest.raises(MicrovmEndpointUnavailableError): + send_job( + options(fake_client(), endpoint.url, sleeps, remaining_seconds=lambda: 13.0) + ) + + assert sleeps == [0.25, 0.5, 1.0, 2.0] + + +def test_starts_no_call_within_1_second_of_the_reserve( + fake_client: Any, fake_endpoint: Any +) -> None: + client = fake_client() + endpoint = fake_endpoint() + + with pytest.raises(MicrovmEndpointUnavailableError, match="after 0 attempts"): + send_job(options(client, endpoint.url, [], remaining_seconds=lambda: 10.5)) + + assert client.calls == [] + assert endpoint.requests == [] + + +def test_bounds_the_request_timeout_by_the_reserve( + fake_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + timeouts: list[float] = [] + + def post(url: str, body: bytes, headers: Any, timeout: float) -> Any: + timeouts.append(timeout) + return 202, "" + + monkeypatch.setattr(request, "_post", post) + + send_job( + options(fake_client(), "mvm-1.example", [], remaining_seconds=lambda: 15.0) + ) + send_job(options(fake_client(), "mvm-1.example", [])) + + # The reserve starts 5 seconds from now. Without a deadline, the cap is + # 10 seconds. + assert 4.5 < timeouts[0] <= 5.0 + assert timeouts[1] == 10.0 + + +@pytest.mark.parametrize( + ("endpoint", "url"), + [ + ("mvm-1.example", "https://mvm-1.example/job"), + ("mvm-1.example/", "https://mvm-1.example/job"), + ("http://127.0.0.1:1", "http://127.0.0.1:1/job"), + ("https://mvm-1.example", "https://mvm-1.example/job"), + ], +) +def test_adds_https_to_an_endpoint_without_a_scheme(endpoint: str, url: str) -> None: + assert request._endpoint_url(endpoint, "/job") == url # noqa: SLF001 + + +def test_fails_when_the_service_returns_no_auth_token(fake_client: Any) -> None: + class NoToken(fake_client): # type: ignore[misc, valid-type] + def create_microvm_auth_token(self, **request: Any) -> dict[str, Any]: + return {} + + with pytest.raises(RuntimeError, match="returned no authToken"): + send_job(options(NoToken(), "mvm-1.example", [])) + + +def test_passes_on_a_failure_of_the_auth_token_call( + fake_client: Any, make_error: Any +) -> None: + client = fake_client(tokens=[make_error("AccessDeniedException", 403)]) + + with pytest.raises(Exception, match="AccessDeniedException"): + send_job(options(client, "mvm-1.example", [], remaining_seconds=lambda: 60.0)) + + +def test_stops_a_stalled_call_at_the_deadline() -> None: + release = threading.Event() + started = time.monotonic() + + with pytest.raises(TimeoutError, match="expired"): + call_before_deadline( + lambda: release.wait(30), + time.monotonic() + 1.2, + lambda: TimeoutError("expired"), + ) + + assert time.monotonic() - started < 5 + release.set() + + +def test_returns_the_result_of_a_call_before_the_deadline() -> None: + assert call_before_deadline(lambda: 7, time.monotonic() + 30, RuntimeError) == 7 + assert call_before_deadline(lambda: 8, None, RuntimeError) == 8 diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/retry_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/retry_test.py index 2cd569bda..596a53f2c 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/retry_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/retry_test.py @@ -20,6 +20,9 @@ ) from aws_durable_execution_sdk_python_extras.microvm import ( + MicrovmEndpointUnavailableError, + MicrovmRequestRejectedError, + MicrovmRequestUnauthorizedError, default_microvm_retry_strategy, retry, ) @@ -86,6 +89,20 @@ def test_does_not_retry_a_permanent_service_error( assert not default_microvm_retry_strategy(make_error(code, status), 1).should_retry +@pytest.mark.parametrize( + ("error", "retried"), + [ + (MicrovmEndpointUnavailableError("tier ended"), True), + (MicrovmRequestUnauthorizedError("refused", 401), True), + (MicrovmRequestRejectedError("not found", 404), False), + ], +) +def test_decides_on_the_end_of_the_request_tier( + error: Exception, retried: bool +) -> None: + assert default_microvm_retry_strategy(error, 1).should_retry is retried + + def test_does_not_retry_another_error() -> None: assert not default_microvm_retry_strategy(RuntimeError("bug"), 1).should_retry 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 bd337aae3..dcbb26604 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 @@ -12,7 +12,11 @@ from aws_durable_execution_sdk_python.config import Duration from aws_durable_execution_sdk_python.exceptions import ValidationError -from aws_durable_execution_sdk_python_extras.microvm import MicrovmConfig, microvm +from aws_durable_execution_sdk_python_extras.microvm import ( + MicrovmConfig, + MicrovmRequestConfig, + microvm, +) from aws_durable_execution_sdk_python_extras.microvm import operation @@ -56,6 +60,51 @@ def config(**fields: Any) -> MicrovmConfig: {"heartbeat_timeout": Duration(0)}, "heartbeat_timeout must be at least 1 second", ), + ( + "build", + {"request": MicrovmRequestConfig(path="job")}, + 'the request path must be a string that starts with "/"', + ), + ( + "build", + {"request": MicrovmRequestConfig(path=7)}, # type: ignore[arg-type] + 'the request path must be a string that starts with "/"', + ), + ( + "build", + {"request": MicrovmRequestConfig(path="/aws/lambda-microvms/run")}, + "must not use the lifecycle hook prefix /aws/lambda-microvms/", + ), + ( + "build", + {"request": MicrovmRequestConfig(port=0)}, + "the request port must be an integer between 1 and 65535", + ), + ( + "build", + {"request": MicrovmRequestConfig(port=65_536)}, + "the request port must be an integer between 1 and 65535", + ), + ( + "build", + {"request": MicrovmRequestConfig(port=True)}, + "the request port must be an integer between 1 and 65535", + ), + ( + "build", + {"request": MicrovmRequestConfig(port=8080.0)}, # type: ignore[arg-type] + "the request port must be an integer between 1 and 65535", + ), + ( + "build", + {"request": MicrovmRequestConfig(retry_window=Duration(0))}, + "the request retry_window must be between 1 second and 10 minutes", + ), + ( + "build", + {"request": MicrovmRequestConfig(retry_window=Duration.from_seconds(601))}, + "the request retry_window must be between 1 second and 10 minutes", + ), ], ) def test_rejects_an_invalid_call_before_any_durable_operation( @@ -136,3 +185,36 @@ def client(self, service: str, region_name: str, config: Any) -> object: def test_takes_the_input_and_the_config_by_keyword_only() -> None: with pytest.raises(TypeError, match="positional argument"): microvm(NoOperations(), "build", {}, config()) # type: ignore[arg-type, misc] + + +class LambdaContext: + def __init__(self, read: Any) -> None: + self.get_remaining_time_in_millis = read + + +def boom() -> int: + msg = "no deadline here" + raise RuntimeError(msg) + + +@pytest.mark.parametrize( + ("lambda_context", "expected"), + [ + (LambdaContext(lambda: 12_500), 12.5), + (LambdaContext(lambda: 2.5), 0.0025), + (None, None), + (object(), None), + (LambdaContext("not callable"), None), + (LambdaContext(boom), None), + (LambdaContext(lambda: math.inf), None), + (LambdaContext(lambda: math.nan), None), + (LambdaContext(lambda: "12"), None), + (LambdaContext(lambda: True), None), + ], +) +def test_reads_the_remaining_time_or_reports_no_deadline( + lambda_context: Any, expected: float | None +) -> None: + context: Any = type("Context", (), {"lambda_context": lambda_context})() + + assert operation._remaining_seconds(context)() == expected # noqa: SLF001 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 bc78bcfd9..f3e9f1da6 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 @@ -5,7 +5,9 @@ 1. The durable function runs on the local test runner. 2. The fake ``RunMicrovm`` sends the ``run`` hook to a real worker on a free - port, as Lambda does when the MicroVM starts. + port, as Lambda does when the MicroVM starts. It returns the worker's + address as the MicroVM's endpoint. So an HTTP delivery reaches the worker + too. 3. The worker runs the job, and reports through a Lambda client double that forwards each callback call to the runner. @@ -37,8 +39,10 @@ from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner from aws_durable_execution_sdk_python_extras.microvm import ( + MAX_RUN_HOOK_PAYLOAD_LENGTH, MicrovmConfig, MicrovmError, + MicrovmRequestConfig, microvm, ) @@ -125,13 +129,17 @@ def handler(event: Any, context: DurableContext) -> Any: callbacks = RunnerCallbacks(runner) worker = start_microvm_worker( handler=job, + routes=fields.pop("routes", None), # The double has the three callback methods that the worker calls. create_client=lambda region: callbacks, # type: ignore[arg-type, return-value] heartbeat_interval_seconds=fields.pop("heartbeat_interval", None), port=0, ) closers.append(worker.close) - client = fake_client(launches=[run_hook_to(worker)]) + client = fake_client( + launches=[run_hook_to(worker)], + endpoint=f"http://127.0.0.1:{worker.port}", + ) holder["config"] = MicrovmConfig( image_identifier="img", execution_role_arn="role", @@ -190,3 +198,28 @@ def job(job_input: Any, context: MicrovmJobContext) -> Any: assert outcome == {"result": "done"} assert callbacks.heartbeats >= 2 + + +def test_runs_a_job_too_large_for_the_run_hook_over_http(composed: Any) -> None: + def job(job_input: Any, context: MicrovmJobContext) -> Any: + return {"length": len(job_input), "microvm": context.microvm_id} + + outcome, client, _ = composed(job, "日" * MAX_RUN_HOOK_PAYLOAD_LENGTH) + + assert outcome == {"result": {"length": 4096, "microvm": "mvm-1"}} + assert len(client.requests("create_microvm_auth_token")) == 1 + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +def test_runs_a_job_on_a_route_of_the_worker(composed: Any) -> None: + def job(job_input: Any, context: MicrovmJobContext) -> Any: + return "handler" + + outcome, _, _ = composed( + job, + {"n": 1}, + request=MicrovmRequestConfig(path="/clone-build"), + routes={"/clone-build": lambda job_input, context: {"route": job_input}}, + ) + + assert outcome == {"result": {"route": {"n": 1}}} From 3dd6cb68929472777042c9c943759f8960af1bfa Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Wed, 14 Oct 2026 13:09:57 -0700 Subject: [PATCH 02/10] fix(extras): apply earlier MicroVM review lessons The heartbeat timeout's first period now covers the request step too, in the config docs and the README. A new test checks that a replay sends no second POST. It fails when the POST runs outside a step. --- .../README.md | 2 +- .../microvm/config.py | 11 ++++--- .../tests/microvm/operation_test.py | 32 +++++++++++++++++++ 3 files changed, 40 insertions(+), 5 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index be3e68eb9..974523aa9 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -124,7 +124,7 @@ The 10-second reserve keeps time for the step to record its failure. A Lambda ti - `timeout` must be between 1 second and 8 hours, because a MicroVM lives at most 8 hours. The MicroVM's `maximumDurationInSeconds` is `timeout` plus 5 minutes, at most 8 hours. So a timeout above 7 hours 55 minutes leaves the MicroVM less than 5 minutes beyond the callback's timeout. - `job_input` is sent as JSON. So the MicroVM receives what `json.dumps` writes: a tuple arrives as a list, and a dict key that is not a string arrives as a string. - The MicroVM gets the `INTERNET_EGRESS` connector, because it needs egress to reach the Lambda callback APIs. For ingress, it gets `NO_INGRESS` with `run` hook delivery and `ALL_INGRESS` with HTTP delivery. If you set `ingress_network_connectors`, allow ingress whenever an input can exceed the `run` hook limit. An empty tuple sends no connector. -- `heartbeat_timeout` starts when the callback is created, before the launch. So it must cover the launch, with its retries, and the boot. A throttled launch can take several minutes. With a shorter heartbeat timeout, the operation fails with `MicrovmTimeoutError` before the MicroVM sends its first heartbeat. +- `heartbeat_timeout` starts when the callback is created, before the launch. The worker sends its first heartbeat when the job arrives. So the first period must cover the launch, the boot, and, for HTTP delivery, the request step, each with its retries. A throttled launch can take several minutes, and the request step retries for up to `retry_window` in each attempt. With a shorter heartbeat timeout, the operation fails with `MicrovmTimeoutError` before the MicroVM sends its first heartbeat. - The Region comes from the durable execution ARN, or from `AWS_REGION`. - Without a `client`, the operation creates one `lambda-microvms` client per Region, on first use. It has a 5-second connect timeout, a 30-second read timeout, and 3 attempts in botocore's standard retry mode. So one launch or terminate attempt takes at most about 2 minutes. A call that runs past the function's timeout records no outcome for its step attempt, and the step runs again. The launch is idempotent through its client token, so it is repeated, not doubled. 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 ead8d51f9..06e468440 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 @@ -106,10 +106,13 @@ class MicrovmConfig: receives the value in the job document. The heartbeat timeout starts when the callback is created, before the - launch. So it must cover the launch, with its retries, and the boot. A - throttled launch can take several minutes. A heartbeat timeout that is - shorter fails the operation with ``MicrovmTimeoutError`` before the - MicroVM can send its first heartbeat. + launch. The worker sends its first heartbeat when the job arrives. So the + first period must cover the launch, the boot, and, for HTTP delivery, the + request step, each with its retries. A throttled launch can take several + minutes, and the request step retries for up to ``request.retry_window`` + in each attempt. A heartbeat timeout that is shorter fails the operation + with ``MicrovmTimeoutError`` before the MicroVM can send its first + heartbeat. """ image_version: str | None = None diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py index 04cf6839f..ccf1ce456 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py @@ -909,3 +909,35 @@ def handler(event: Any, context: DurableContext) -> Any: assert outcome(result) == "MicrovmTimeoutError" assert terminated_in == [2] + + +def test_posts_the_job_once_across_invocations( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + # The first invocation suspends on the callback, and the replay times it + # out. The request step is checkpointed, so the replay sends no second + # POST. + endpoint = fake_endpoint() + client = fake_client(endpoint=endpoint.url) + invocations: list[int] = [] + config = config_with( + client, + timeout=Duration.from_seconds(2), + request=MicrovmRequestConfig(path="/job"), + ) + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + invocations.append(1) + try: + return microvm(context, "build", job_input=event, config=config) + except MicrovmError as error: + return type(error).__name__ + + result = run(handler, {}) + + assert outcome(result) == "MicrovmTimeoutError" + assert len(invocations) >= 2 + assert len(endpoint.requests) == 1 + assert len(client.requests("create_microvm_auth_token")) == 1 + assert len(client.requests("run_microvm")) == 1 From c2d7db62d3cb2207c6dda70abbe016047a0d48ef Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Wed, 14 Oct 2026 13:58:41 -0700 Subject: [PATCH 03/10] fix(extras): bound each MicroVM POST and stop redirects - Run each POST in a daemon thread, and wait at most 10 seconds or until the invocation reserve. A socket timeout applies to each socket operation, and each byte received starts it again. So a slow endpoint kept the request step running past the end of the invocation. - Do not follow redirects. A 301, 302, or 303 turned the POST into a GET, and its 2xx counted as accepted although the job never started. A redirect to another host received the auth token. A 3xx now raises MicrovmRequestRejectedError with the Location in the message. - Percent-encode the path as fetch does. A space or a non-ASCII character failed every attempt with a retryable error. - Re-raise a BaseException from the deadline thread with its own type. - Add tests: replay keeps the launch's delivery in both directions, the whole POST is bounded, redirects are rejected, the path is encoded, and proxy settings are ignored. Each fails under the mutation it guards against. --- .../README.md | 8 +- .../microvm/operation.py | 2 +- .../microvm/request.py | 135 ++++++++++++++---- .../microvm/retry.py | 3 +- .../tests/conftest.py | 19 ++- .../tests/microvm/operation_test.py | 87 +++++++++++ .../tests/microvm/request_test.py | 109 +++++++++++++- 7 files changed, 332 insertions(+), 31 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index 974523aa9..77549b3ac 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -86,7 +86,7 @@ Every MicroVM failure is a `MicrovmError`. Catch the base class for all of them, | Error | When | |---|---| | `MicrovmLaunchError` | `RunMicrovm` failed after all retries. When the last attempt timed out or lost its connection, the service may have created the MicroVM anyway, and it then runs the job. So running the job again after this error can run it twice. | -| `MicrovmDeliveryError` | HTTP delivery failed after all retries, or the route answered a 4xx status such as 404. | +| `MicrovmDeliveryError` | HTTP delivery failed after all retries, or the route answered a 4xx status such as 404, or a redirect. | | `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. | @@ -115,9 +115,11 @@ A step retry costs one durable operation and one more Lambda invocation. Most fa 1. **Inside one step attempt.** The boto3 client retries throttling and server errors of `RunMicrovm`, `CreateMicrovmAuthToken`, and `TerminateMicrovm`. The request step retries the POST itself: connection errors, timeouts, 429, and 5xx, with a delay from 250 milliseconds to 4 seconds. It retries one 401 or 403 with a new auth token. The endpoint answers this way until the MicroVM's `run` hook returns. This tier ends after `retry_window`, or 10 seconds before the invocation times out, whichever comes first. 2. **The step retry strategy.** When the first tier gives up, the step retries after a longer delay, outside the invocation. `default_microvm_retry_strategy` is the default, unless the config sets `retry_strategy`. It retries throttling, server errors, connection failures, timeouts, and the end of the first tier, for at most 5 attempts, with a delay of 2, 4, 8, and 16 seconds. Every other error, such as `ValidationException`, `AccessDeniedException`, or `InsufficientCapacityException`, fails the step at once. -A route that answers another 4xx status fails the operation without a retry, with `MicrovmDeliveryError`. +A route that answers another 4xx status fails the operation without a retry, with `MicrovmDeliveryError`. So does a redirect, 3xx. The operation does not follow redirects, and the message names the `Location`. Set the path that the route serves, for example with its trailing slash. The JavaScript operation follows redirects. -The 10-second reserve keeps time for the step to record its failure. A Lambda timeout during a step attempt records no outcome for the attempt. So every call of the first tier must end before the reserve. boto3 cannot cancel a call, so a `CreateMicrovmAuthToken` call runs in a separate thread, and the step stops waiting for it at the reserve. The POST's socket timeout is at most the time left before the reserve. +The path is percent-encoded as the JavaScript SDK's `fetch` encodes it. A space becomes `%20`, and a non-ASCII character becomes its UTF-8 bytes. A `%` stays, so an encoded path is sent as it is. + +The 10-second reserve keeps time for the step to record its failure. A Lambda timeout during a step attempt records no outcome for the attempt. So every call of the first tier must end before the reserve. boto3 cannot cancel a call, so a `CreateMicrovmAuthToken` call runs in a separate thread, and the step stops waiting for it at the reserve. Each POST runs in a separate thread too, and the step waits for it at most 10 seconds, or until the reserve. A socket timeout alone would not bound it, because each byte received starts the socket timeout again. ### Defaults and limits 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 c05e32e76..82ffa3cfe 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 @@ -208,7 +208,7 @@ def microvm( operation starts. MicrovmLaunchError: When ``RunMicrovm`` fails after all retries. MicrovmDeliveryError: When the HTTP delivery fails after all retries, - or the route answers a 4xx status such as 404. + or the route answers a 4xx status such as 404, or a redirect. 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. 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 1cc074ff8..fb0b41dd1 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 @@ -13,10 +13,11 @@ import threading import time import urllib.error +import urllib.parse import urllib.request from collections.abc import Callable, Mapping from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, TypeVar +from typing import TYPE_CHECKING, Any, TypeVar, cast from aws_durable_execution_sdk_python.types import LoggerInterface @@ -70,8 +71,9 @@ class MicrovmRequestRejectedError(Exception): """The endpoint rejected the job with a status that a retry cannot change. - Examples are 400, 404, and 413. The default step retry strategy does not - retry it. + Examples are 400, 404, and 413. A redirect, 3xx, is rejected too, because + the operation does not follow redirects. The default step retry strategy + does not retry it. """ def __init__(self, message: str, status: int) -> None: @@ -149,29 +151,47 @@ def call_before_deadline( returns. Its result is discarded. Raises: - Exception: The exception that the call raised, or the result of + BaseException: The exception that the call raised, or the result of ``on_expired``. """ if hard_deadline is None: return call() budget = _call_budget(hard_deadline, float("inf"), on_expired) + finished, value = _run_with_limit(call, budget) + if not finished: + raise on_expired() + return cast("_T", value) + + +def _run_with_limit(call: Callable[[], _T], seconds: float) -> tuple[bool, _T | None]: + """Run ``call`` in a daemon thread, and wait for it at most ``seconds``. + + Returns: + ``(True, value)`` when the call returned in time. ``(False, None)`` + when it did not. The thread then keeps running until the call ends, + and its result is discarded. + + Raises: + BaseException: The exception that the call raised in time. It is + re-raised with its own type, a ``BaseException`` included. + """ outcome: dict[str, Any] = {} done = threading.Event() def run() -> None: try: outcome["value"] = call() - except Exception as error: # noqa: BLE001 - re-raised by the caller below + except BaseException as error: # noqa: BLE001 - re-raised by the caller below outcome["error"] = error finally: done.set() threading.Thread(target=run, name="microvm-call", daemon=True).start() - if not done.wait(budget): - raise on_expired() + if not done.wait(seconds): + return False, None if "error" in outcome: raise outcome["error"] - return outcome["value"] + return True, outcome["value"] @dataclass(frozen=True, kw_only=True) @@ -203,7 +223,15 @@ def send_job(options: SendJobOptions) -> int: 2. 401 or 403: create a new auth token and retry once. A second 401 or 403 raises :class:`MicrovmRequestUnauthorizedError`, and the step retry strategy takes over. - 3. Any other 4xx: raise :class:`MicrovmRequestRejectedError` at once. + 3. Any other 4xx, or a 3xx: raise :class:`MicrovmRequestRejectedError` + at once. + + The function does not follow a redirect. The JavaScript SDK's ``fetch`` + follows one, and it differs in three ways. A 301, 302, or 303 turns the + POST into a GET without the job, and the GET's 2xx counts as accepted. + So the job never starts, and the operation waits until its timeout. A + 307 or 308 to another host receives the auth token. So a redirect fails + the delivery at once, with the ``Location`` in the message. The delay starts at 250 milliseconds and doubles up to 4 seconds. The tier ends at the retry window, or 10 seconds before the invocation times @@ -220,6 +248,11 @@ def send_job(options: SendJobOptions) -> int: start. A Lambda timeout during the attempt would record no outcome, and the reserve prevents that. + Each request runs in a daemon thread, and the function waits for it at + most 10 seconds, or until the reserve. A socket timeout alone does not + bound a request: it applies to each socket operation, and each byte + received starts it again. + The endpoint can accept a request before Lambda has sent the ``run`` hook. So the body carries the MicroVM ID too. @@ -277,7 +310,8 @@ def new_token() -> Mapping[str, str]: headers = {"content-type": "application/json", **token} if options.port != DEFAULT_MICROVM_PORT: headers[_PORT_HEADER] = str(options.port) - status, failure = _post(url, options.body, headers, timeout) + answer = _post_within(url, options.body, headers, timeout) + status = answer.status if status is not None and 200 <= status < 300: options.logger.info( @@ -285,7 +319,7 @@ def new_token() -> Mapping[str, str]: ) return status if status is None: - last_failure = failure + last_failure = answer.failure else: last_failure = f"HTTP {status}" if status in {401, 403}: @@ -304,6 +338,14 @@ def new_token() -> Mapping[str, str]: extra={"attempt": attempt, "status": status}, ) continue + if 300 <= status < 400: # noqa: PLR2004 + msg = ( + f"The MicroVM answered the job at {options.path} with HTTP " + f"{status}, a redirect to {answer.location or 'no location'}. " + "The operation does not follow redirects, so set the path " + "that the route serves." + ) + raise MicrovmRequestRejectedError(msg, status) if 400 <= status < 500 and status != 429: # noqa: PLR2004 msg = ( f"The MicroVM rejected the job at {options.path} with HTTP {status}" @@ -333,39 +375,84 @@ def _create_auth_token(options: SendJobOptions) -> Mapping[str, str]: return token +class _NoRedirect(urllib.request.HTTPRedirectHandler): + """Do not follow a redirect. urllib then raises ``HTTPError`` with the + 3xx status.""" + + def redirect_request(self, *args: Any, **kwargs: Any) -> None: + return None + + # Proxy settings from the environment do not apply. The request goes to the # MicroVM endpoint directly, as the JavaScript SDK's fetch sends it. -_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({})) +_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({}), _NoRedirect()) + + +@dataclass(frozen=True) +class _Answer: + """The outcome of one POST.""" + status: int | None + """The HTTP status. ``None`` when no answer arrived.""" + failure: str = "" + """Why no answer arrived.""" + location: str | None = None + """The ``Location`` header of a redirect.""" -def _post( + +def _post_within( url: str, body: bytes, headers: Mapping[str, str], timeout: float -) -> tuple[int | None, str]: - """POST ``body``, and return the status, or ``None`` and the failure. +) -> _Answer: + """POST ``body``, and wait for the answer at most ``timeout`` seconds. - The body of the response is not used. So the function reads no body, and - closes the response. + The POST runs in a daemon thread, because a socket timeout alone does not + bound the whole request. A POST without an answer in time counts as one + failed attempt, like a refused connection. + """ + finished, answer = _run_with_limit( + lambda: _post(url, body, headers, timeout), timeout + ) + if not finished: + return _Answer(None, f"TimeoutError: no answer in {timeout:.1f} s") + return cast("_Answer", answer) - ``timeout`` applies to each socket operation: the connect, and each send - and receive. The endpoint answers at once, so it bounds the request. + +def _post(url: str, body: bytes, headers: Mapping[str, str], timeout: float) -> _Answer: + """POST ``body``, and return the status, or the failure. + + The body of the response is not used. So the function reads no body, and + closes the response. ``timeout`` is the socket timeout. It applies to + each socket operation, so :func:`_post_within` bounds the whole request. """ request = urllib.request.Request( # noqa: S310 - the URL is https or the endpoint's own url, data=body, headers=dict(headers), method="POST" ) try: with _OPENER.open(request, timeout=timeout) as response: - return response.status, "" + return _Answer(response.status) except urllib.error.HTTPError as error: + location = error.headers.get("Location") if error.headers else None error.close() - return error.code, "" + return _Answer(error.code, location=location) 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. - return None, f"{type(error).__name__}: {error}" + return _Answer(None, f"{type(error).__name__}: {error}") + + +_PATH_SAFE = "/%!$&'()*+,;=:@-._~?#[]|^" +"""The characters that stay as they are in the path. + +The WHATWG URL parser, which the JavaScript SDK's ``fetch`` uses, +percent-encodes a space, ``"``, ``<``, ``>``, a backtick, ``{``, ``}``, the +control characters, and every non-ASCII character in a path, as UTF-8. +``urllib.parse.quote`` encodes those too, and leaves this set as it is. A +``%`` stays, so a path that is already encoded is not encoded twice. +""" def _endpoint_url(endpoint: str, path: str) -> str: - """Join the endpoint and the path. + """Join the endpoint and the percent-encoded path. ``RunMicrovm`` returns the endpoint as a host name without a scheme. """ @@ -374,4 +461,4 @@ def _endpoint_url(endpoint: str, path: str) -> str: if endpoint.startswith(("https://", "http://")) else (f"https://{endpoint}") ) - return f"{base.rstrip('/')}{path}" + return f"{base.rstrip('/')}{urllib.parse.quote(path, safe=_PATH_SAFE)}" diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/retry.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/retry.py index 08da15ccb..46387694a 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/retry.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/retry.py @@ -85,7 +85,8 @@ def default_microvm_retry_strategy(error: Exception, attempt: int) -> RetryDecis ``MicrovmRequestUnauthorizedError`` after a refused new token. ``MicrovmRequestRejectedError`` is not retried. The route rejected the - job with a 4xx status, and the same job gets the same answer. + job with a 4xx status or a redirect, and the same job gets the same + answer. Every other error fails the step at once. A ``ValidationException``, for example, fails the same way on every attempt. 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 94ae9cab7..253f2043a 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py @@ -117,7 +117,7 @@ class FakeEndpoint: 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. + without an answer. A redirect status gets ``Location: /moved``. """ def __init__(self, statuses: list[int | None] | None = None) -> None: @@ -127,6 +127,21 @@ def __init__(self, statuses: list[int | None] | None = None) -> None: endpoint = self class Handler(BaseHTTPRequestHandler): + def do_GET(self) -> None: # noqa: N802 - the http.server name + # A followed 301, 302, or 303 arrives as a GET. + with endpoint._lock: + endpoint.requests.append( + { + "method": "GET", + "path": self.path, + "headers": {}, + "body": None, + } + ) + self.send_response(200) + self.send_header("content-length", "0") + self.end_headers() + def do_POST(self) -> None: # noqa: N802 - the http.server name length = int(self.headers.get("content-length", "0")) body = self.rfile.read(length) @@ -144,6 +159,8 @@ def do_POST(self) -> None: # noqa: N802 - the http.server name self.connection.shutdown(socket.SHUT_RDWR) return self.send_response(status) + if 300 <= status < 400: # noqa: PLR2004 + self.send_header("location", "/moved") self.send_header("content-length", "0") self.end_headers() diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py index ccf1ce456..61e353e80 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py @@ -8,6 +8,7 @@ import datetime import hashlib import json +import threading from collections.abc import Callable, Iterator from types import SimpleNamespace from typing import Any @@ -941,3 +942,89 @@ def handler(event: Any, context: DurableContext) -> Any: assert len(endpoint.requests) == 1 assert len(client.requests("create_microvm_auth_token")) == 1 assert len(client.requests("run_microvm")) == 1 + + +def test_replay_keeps_the_run_hook_delivery_of_the_launch( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + # The launch uses the run hook. A deploy before the replay sets a path, + # which would select HTTP. The replay must follow the checkpoint. + endpoint = fake_endpoint() + client = fake_client(endpoint=endpoint.url) + first = config_with(client) + second = config_with(client, request=MicrovmRequestConfig(path="/job")) + invocations: list[int] = [] + first_ended = threading.Event() + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + invocations.append(1) + config = first if len(invocations) == 1 else second + try: + return {"result": microvm(context, "build", job_input=event, config=config)} + except MicrovmError as error: + return {"error": type(error).__name__, "message": error.message} + finally: + first_ended.set() + + def complete(runner: DurableFunctionTestRunner, callback_id: str) -> None: + # The first invocation suspends on the callback before it completes. + # So the result arrives in a replay. + assert first_ended.wait(WAIT) + runner.send_callback_success(callback_id, json.dumps("done").encode()) + + result = run(handler, {}, complete) + + assert outcome(result) == {"result": "done"} + assert len(invocations) == 2 + assert endpoint.requests == [] + assert client.requests("create_microvm_auth_token") == [] + + +def test_replay_keeps_the_http_delivery_of_the_launch( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + # The launch uses HTTP, and the first request step attempt fails. A + # deploy before the replay removes the path, which would select the run + # hook for this small input. The replay must still deliver over HTTP. + endpoint = fake_endpoint([503, 503, 503]) + client = fake_client(endpoint=endpoint.url) + + def once(error: Exception, attempt: int) -> RetryDecision: + if attempt < 2 and default_microvm_retry_strategy(error, attempt).should_retry: + return RetryDecision.retry(Duration.from_seconds(1)) + return RetryDecision.no_retry() + + window = Duration.from_seconds(1) + first = config_with( + client, + retry_strategy=once, + request=MicrovmRequestConfig(path="/job", retry_window=window), + ) + second = config_with( + client, retry_strategy=once, request=MicrovmRequestConfig(retry_window=window) + ) + invocations: list[int] = [] + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + invocations.append(1) + config = first if len(invocations) == 1 else second + try: + return {"result": microvm(context, "build", job_input=event, config=config)} + except MicrovmError as error: + return {"error": type(error).__name__, "message": error.message} + + result = run(handler, {}, succeed_with("done")) + + assert outcome(result) == {"result": "done"} + assert len(invocations) >= 2 + # 3 requests in the first attempt, and 1 in the replay. The delivery + # comes from the checkpoint. The route comes from the current config, as + # in the JavaScript operation, which checkpoints no route either. + assert [r["path"] for r in endpoint.requests] == [ + "/job", + "/job", + "/job", + "/durable-execution/v1/job", + ] 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 f59af8550..ca6236674 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 @@ -6,8 +6,10 @@ from __future__ import annotations import logging +import socket import threading import time +from collections.abc import Iterator from typing import Any import pytest @@ -226,7 +228,7 @@ def test_bounds_the_request_timeout_by_the_reserve( def post(url: str, body: bytes, headers: Any, timeout: float) -> Any: timeouts.append(timeout) - return 202, "" + return request._Answer(202) # noqa: SLF001 monkeypatch.setattr(request, "_post", post) @@ -290,3 +292,108 @@ def test_stops_a_stalled_call_at_the_deadline() -> None: def test_returns_the_result_of_a_call_before_the_deadline() -> None: assert call_before_deadline(lambda: 7, time.monotonic() + 30, RuntimeError) == 7 assert call_before_deadline(lambda: 8, None, RuntimeError) == 8 + + +@pytest.fixture +def trickle_server() -> Iterator[str]: + """A server that answers one byte every 0.3 seconds. + + Each byte starts the socket timeout again. So only a bound on the whole + request stops the POST in time. + """ + listener = socket.create_server(("127.0.0.1", 0)) + stop = threading.Event() + + def serve() -> None: + connection, _ = listener.accept() + with connection: + connection.recv(65536) + for byte in b"HTTP/1.1 202 Accepted\r\ncontent-length: 0\r\n\r\n": + if stop.wait(0.3): + return + try: + connection.sendall(bytes([byte])) + except OSError: + return + + thread = threading.Thread(target=serve, daemon=True) + thread.start() + yield f"http://127.0.0.1:{listener.getsockname()[1]}" + stop.set() + listener.close() + + +def test_bounds_the_whole_request_and_not_each_socket_operation( + trickle_server: str, +) -> None: + started = time.monotonic() + + answer = request._post_within( # noqa: SLF001 + f"{trickle_server}/job", BODY, {}, 1.0 + ) + + assert time.monotonic() - started < 2 + assert answer.status is None + assert answer.failure == "TimeoutError: no answer in 1.0 s" + + +@pytest.mark.parametrize("status", [301, 302, 303, 307, 308]) +def test_rejects_a_redirect_without_following_it( + fake_client: Any, fake_endpoint: Any, status: int +) -> None: + endpoint = fake_endpoint([status]) + + with pytest.raises(MicrovmRequestRejectedError) as raised: + send_job(options(fake_client(), endpoint.url, [])) + + assert raised.value.status == status + assert str(raised.value) == ( + f"The MicroVM answered the job at /durable-execution/v1/job with HTTP " + f"{status}, a redirect to /moved. The operation does not follow " + "redirects, so set the path that the route serves." + ) + # A followed 301, 302, or 303 would arrive as a GET without the job. + assert [r["path"] for r in endpoint.requests] == ["/durable-execution/v1/job"] + + +@pytest.mark.parametrize( + ("path", "sent"), + [ + ("/clone build", "/clone%20build"), + ("/café", "/caf%C3%A9"), + ("/already%20encoded", "/already%20encoded"), + ("/a/b-c_d.e~f!$&'()*+,;=:@[]|^", "/a/b-c_d.e~f!$&'()*+,;=:@[]|^"), + ('/q"<>`{}', "/q%22%3C%3E%60%7B%7D"), + ], +) +def test_percent_encodes_the_path_as_fetch_does( + fake_client: Any, fake_endpoint: Any, path: str, sent: str +) -> None: + endpoint = fake_endpoint() + + send_job(options(fake_client(), endpoint.url, [], path=path)) + + assert endpoint.requests[0]["path"] == sent + + +def test_ignores_proxy_settings_in_the_environment( + fake_client: Any, fake_endpoint: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # Port 1 refuses every connection. A request that used the proxy would + # fail every attempt. + for name in ("HTTP_PROXY", "http_proxy", "HTTPS_PROXY", "https_proxy"): + monkeypatch.setenv(name, "http://127.0.0.1:1") + for name in ("NO_PROXY", "no_proxy"): + monkeypatch.setenv(name, "") + endpoint = fake_endpoint() + + assert send_job(options(fake_client(), endpoint.url, [])) == 202 + assert len(endpoint.requests) == 1 + + +def test_passes_on_a_base_exception_of_the_call() -> None: + def interrupted() -> None: + raise KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + call_before_deadline(interrupted, time.monotonic() + 30, RuntimeError) From 029efc7ac5eeee0368f352e8fffb82348be2b831 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Thu, 15 Oct 2026 08:01:10 -0700 Subject: [PATCH 04/10] fix(extras): address the first review of HTTP delivery - Fail the request step at once, with a message that names the function timeout, when less than 11 seconds remain in the invocation. The error stays retryable. Before, the message said that the MicroVM did not accept the job after 0 attempts. - Send the request path as written, and reject a path that is not an RFC 3986 path with ValidationError. The workers match the raw request path, so an encoded path missed a route registered with the same string. - Add an optional opener to the config, for a network that reaches the endpoint only through a proxy. Reject a redirect that such an opener followed. Build the default opener on first use. - Read the remaining time from the Lambda context directly. A failure of the call now raises, instead of turning off the reserve. - Document that a route can receive the same job again, and must ignore a repeated callback ID. --- .../README.md | 11 +- .../microvm/config.py | 33 +++++- .../microvm/operation.py | 45 ++++---- .../microvm/request.py | 109 +++++++++++++----- .../tests/microvm/request_test.py | 102 +++++++++++++--- .../tests/microvm/validation_test.py | 48 ++++---- 6 files changed, 260 insertions(+), 88 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index 77549b3ac..579cea86c 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -59,6 +59,14 @@ microvm(context, "build", job_input=job_input, config=config) A `path` sends every job over HTTP, whatever its size. The route must answer with a 2xx status at once and run the job in the background. `port` names a port other than 8080. `retry_window` sets how long the request step retries inside one attempt, 60 seconds by default. +The same job can arrive again, also after it ended, so ignore a second request for a callback ID that you already ran. The worker package does this. A route served without it must do it itself. A job arrives again in these cases: + +1. A request reaches the route, and its answer is lost. +2. The request step stops waiting for a request after at most 10 seconds, but the request keeps running in its thread and reaches the route. The next attempt sends the job again. +3. The invocation ends after the route accepted the job, and before the request step checkpointed its success. The step then runs again in the next invocation. + +The path must be an RFC 3986 path, and the operation sends it exactly as written. Percent-encode any other character, for example a space as `%20`. The workers match the raw request path, so register the same string as the worker's route. An invalid path raises `ValidationError` before any durable operation starts. + ### What it does The operation runs its durable operations in one child context named after the second argument: @@ -117,7 +125,7 @@ A step retry costs one durable operation and one more Lambda invocation. Most fa A route that answers another 4xx status fails the operation without a retry, with `MicrovmDeliveryError`. So does a redirect, 3xx. The operation does not follow redirects, and the message names the `Location`. Set the path that the route serves, for example with its trailing slash. The JavaScript operation follows redirects. -The path is percent-encoded as the JavaScript SDK's `fetch` encodes it. A space becomes `%20`, and a non-ASCII character becomes its UTF-8 bytes. A `%` stays, so an encoded path is sent as it is. +The POST ignores the proxy settings in the environment, as the JavaScript SDK's `fetch` does. boto3 reads `HTTPS_PROXY`. So in a network that reaches the MicroVM endpoint only through a proxy, the launch succeeds and every POST fails. Pass an `opener` in the config for that network, for example `urllib.request.build_opener(urllib.request.ProxyHandler({"https": "http://proxy.internal:3128"}))`. When an opener follows a redirect, the delivery still fails at once with `MicrovmDeliveryError`. The 10-second reserve keeps time for the step to record its failure. A Lambda timeout during a step attempt records no outcome for the attempt. So every call of the first tier must end before the reserve. boto3 cannot cancel a call, so a `CreateMicrovmAuthToken` call runs in a separate thread, and the step stops waiting for it at the reserve. Each POST runs in a separate thread too, and the step waits for it at most 10 seconds, or until the reserve. A socket timeout alone would not bound it, because each byte received starts the socket timeout again. @@ -127,6 +135,7 @@ The 10-second reserve keeps time for the step to record its failure. A Lambda ti - `job_input` is sent as JSON. So the MicroVM receives what `json.dumps` writes: a tuple arrives as a list, and a dict key that is not a string arrives as a string. - The MicroVM gets the `INTERNET_EGRESS` connector, because it needs egress to reach the Lambda callback APIs. For ingress, it gets `NO_INGRESS` with `run` hook delivery and `ALL_INGRESS` with HTTP delivery. If you set `ingress_network_connectors`, allow ingress whenever an input can exceed the `run` hook limit. An empty tuple sends no connector. - `heartbeat_timeout` starts when the callback is created, before the launch. The worker sends its first heartbeat when the job arrives. So the first period must cover the launch, the boot, and, for HTTP delivery, the request step, each with its retries. A throttled launch can take several minutes, and the request step retries for up to `retry_window` in each attempt. With a shorter heartbeat timeout, the operation fails with `MicrovmTimeoutError` before the MicroVM sends its first heartbeat. +- HTTP delivery needs a function timeout of at least 11 seconds. The request step keeps a 10-second reserve before the invocation times out, and it starts a call only with at least 1 second left before the reserve. With less time, every attempt fails with `MicrovmEndpointUnavailableError`, and the message names the function timeout. The Lambda default timeout is 3 seconds. A function with that timeout runs small jobs in the `run` hook, and fails every job too large for it. - The Region comes from the durable execution ARN, or from `AWS_REGION`. - Without a `client`, the operation creates one `lambda-microvms` client per Region, on first use. It has a 5-second connect timeout, a 30-second read timeout, and 3 attempts in botocore's standard retry mode. So one launch or terminate attempt takes at most about 2 minutes. A call that runs past the function's timeout records no outcome for its step attempt, and the step runs again. The launch is idempotent through its client token, so it is repeated, not doubled. 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 06e468440..35290ae35 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 @@ -5,6 +5,7 @@ from __future__ import annotations +import urllib.request from collections.abc import Callable, Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -31,10 +32,20 @@ class MicrovmRequestConfig: An HTTP delivery POSTs the job document to a route of the application in the MicroVM. The route must answer with a 2xx status at once, and run the job in the background. + + The same job can arrive again, also after it ended. So a route must + ignore a second request for a callback ID that it already ran. The + worker package does this. A route served without the worker package must + do it itself. """ path: str | None = None - """The route path, for example ``/clone-build``. It must start with ``/``. + """The route path, for example ``/clone-build``. + + It must be an RFC 3986 path that starts with ``/``. The operation sends + it exactly as written. Percent-encode any other character, for example + a space as ``%20``. Register the same string as the worker's route, because + the worker matches the raw request path. ``None`` sends the job to :data:`DEFAULT_MICROVM_JOB_PATH`. The worker package serves that path with its ``handler``. A path sends the job to one @@ -160,6 +171,26 @@ class MicrovmConfig: the POST. Defaults to :func:`default_microvm_retry_strategy`. """ + opener: urllib.request.OpenerDirector | None = None + """The urllib opener that POSTs the job for HTTP delivery. + + ``None`` uses an opener that ignores the proxy settings in the + environment and follows no redirect. boto3 reads ``HTTPS_PROXY``, so in + a network that reaches the MicroVM endpoint only through a proxy, pass + an opener with that proxy: + + .. code-block:: python + + urllib.request.build_opener( + urllib.request.ProxyHandler({"https": "http://proxy.internal:3128"}) + ) + + The operation does not follow redirects. When an opener follows one, the + delivery fails at once with ``MicrovmRequestRejectedError``. A followed + redirect to another host also receives the auth token, so prefer an + opener that follows no redirect. + """ + client: LambdaMicrovmsClient | None = None """The boto3 ``lambda-microvms`` client. 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 82ffa3cfe..53a50b3c6 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 @@ -7,8 +7,8 @@ import hashlib import json -import math import os +import re import threading from collections.abc import Callable from dataclasses import dataclass @@ -104,6 +104,13 @@ _MAX_PORT = 65_535 +_ROUTE_PATH = re.compile(r"(?:/(?:[A-Za-z0-9\-._~!$&'()*+,;=:@]|%[0-9A-Fa-f]{2})*)+") +"""An RFC 3986 path, section 3.3: segments of pchar characters and %XX escapes. + +The operation sends the path as written. Both workers match the raw request +path, so the config and the route use the same string. +""" + @dataclass(frozen=True, kw_only=True) class LaunchResult: @@ -458,6 +465,7 @@ def attempt(step_context: StepContext) -> None: retry_window_seconds=self.retry_window_seconds, remaining_seconds=remaining_seconds, logger=step_context.logger, + opener=self.config.opener, ) ) @@ -686,9 +694,13 @@ def _validate_request(name: str, request: MicrovmRequestConfig) -> int: ValidationError: When a setting is not valid. """ path = request.path - if path is not None and (not isinstance(path, str) or not path.startswith("/")): + if path is not None and ( + not isinstance(path, str) or not _ROUTE_PATH.fullmatch(path) + ): msg = ( - f'MicroVM "{name}": the request path must be a string that starts with "/"' + f'MicroVM "{name}": the request path must be an RFC 3986 path that ' + 'starts with "/". Percent-encode any other character, for example a ' + "space as %20." ) raise ValidationError(msg) if path is not None and path.startswith(_HOOK_PATH_PREFIX): @@ -723,29 +735,16 @@ def _validate_request(name: str, request: MicrovmRequestConfig) -> int: def _remaining_seconds(context: DurableContext) -> Callable[[], float | None]: """Return a reader of the invocation's remaining time, in seconds. - The first retry tier uses it to stop before the invocation times out. The - Lambda context has ``get_remaining_time_in_millis()``. Another compute's - context may lack it, or the method may raise or return no finite number. - The reader returns ``None`` in each of these cases. So the request step - falls back to its retry window, and nothing fails. + The first retry tier uses it to stop before the invocation times out. + ``None`` means that the compute reports no deadline. The core passes the + Lambda handler's context, whose ``get_remaining_time_in_millis()`` + returns an ``int``. A failure of that call raises, so the reserve is + never turned off without a sign. """ lambda_context = context.lambda_context - read = getattr(lambda_context, "get_remaining_time_in_millis", None) - if not callable(read): + if lambda_context is None: return lambda: None - - def remaining() -> float | None: - try: - value = read() - except Exception: # noqa: BLE001 - any failure means no deadline - return None - if isinstance(value, bool) or not isinstance(value, (int, float)): - return None - if not math.isfinite(value): - return None - return value / 1000 - - return remaining + return lambda: lambda_context.get_remaining_time_in_millis() / 1000 def _detail(error: StepError) -> str: 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 fb0b41dd1..c4813f228 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 @@ -9,11 +9,11 @@ from __future__ import annotations +import functools import http.client import threading import time import urllib.error -import urllib.parse import urllib.request from collections.abc import Callable, Mapping from dataclasses import dataclass @@ -210,6 +210,8 @@ class SendJobOptions: """The remaining invocation time. ``None`` means that the compute reports no deadline, and only the retry window applies.""" logger: LoggerInterface + opener: urllib.request.OpenerDirector | None = None + """The urllib opener for the POST. ``None`` uses :func:`default_opener`.""" sleep: Callable[[float], None] = time.sleep @@ -256,10 +258,19 @@ 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 request can reach the MicroVM while its response is lost. The next - request then delivers the same job again. The worker ignores a second - request for a callback ID that runs or recently ended, so the job still - runs once. + The same job can reach the route more than once: + + 1. A request reaches the route, and its answer is lost. + 2. The function stops waiting for a request, but the request keeps + running in its thread and reaches the route. The next attempt sends + the job again. + 3. The invocation ends after the route accepted the job, and before the + request step checkpointed its success. The step then runs again in + the next invocation. + + The worker package ignores a second request for a callback ID that runs + or recently ended. A route served without the worker package must ignore + it too. Returns: The status that accepted the job. @@ -270,6 +281,21 @@ def send_job(options: SendJobOptions) -> int: MicrovmEndpointUnavailableError: When the tier ends. Exception: When ``CreateMicrovmAuthToken`` fails. """ + # A step that starts too close to the end of the invocation cannot make + # one call before the reserve. Say so, instead of reporting that the + # MicroVM did not accept the job. The error stays retryable: a step + # retry runs early in a new invocation, with more time left. A function + # timeout under the needed time fails every attempt, and the final error + # names the fix. + remaining = options.remaining_seconds() + needed = INVOCATION_RESERVE_SECONDS + MIN_CALL_SECONDS + if remaining is not None and remaining < needed: + msg = ( + f"The request step had {remaining:.1f} seconds left in the Lambda " + f"invocation, and it needs at least {needed:.0f} seconds to send " + "the job. Raise the function timeout." + ) + raise MicrovmEndpointUnavailableError(msg) started = time.monotonic() window_ends = started + options.retry_window_seconds attempt = 0 @@ -310,8 +336,20 @@ def new_token() -> Mapping[str, str]: headers = {"content-type": "application/json", **token} if options.port != DEFAULT_MICROVM_PORT: headers[_PORT_HEADER] = str(options.port) - answer = _post_within(url, options.body, headers, timeout) + answer = _post_within( + url, options.body, headers, timeout, options.opener or default_opener() + ) status = answer.status + if answer.followed_redirect: + # A caller's opener followed a redirect. The route never received + # the job: a 301, 302, or 303 turns the POST into a GET. + msg = ( + f"The MicroVM answered the job at {options.path} with a " + f"redirect to {answer.location}, and the opener followed it. " + "The operation does not follow redirects, so set the path that " + "the route serves." + ) + raise MicrovmRequestRejectedError(msg, status or 0) if status is not None and 200 <= status < 300: options.logger.info( @@ -383,9 +421,19 @@ def redirect_request(self, *args: Any, **kwargs: Any) -> None: return None -# Proxy settings from the environment do not apply. The request goes to the -# MicroVM endpoint directly, as the JavaScript SDK's fetch sends it. -_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({}), _NoRedirect()) +@functools.cache +def default_opener() -> urllib.request.OpenerDirector: + """The urllib opener that the POST uses when the config sets none. + + 1. It ignores the proxy settings in the environment. The request goes to + the MicroVM endpoint directly, as the JavaScript SDK's ``fetch`` sends + it. + 2. It follows no redirect. + + It is built on first use. Building it loads the CA certificates, so a + function whose jobs all go in the ``run`` hook does not pay for it. + """ + return urllib.request.build_opener(urllib.request.ProxyHandler({}), _NoRedirect()) @dataclass(frozen=True) @@ -397,11 +445,18 @@ class _Answer: failure: str = "" """Why no answer arrived.""" location: str | None = None - """The ``Location`` header of a redirect.""" + """The ``Location`` header of a redirect, or the URL that a followed + redirect reached.""" + followed_redirect: bool = False + """Whether the opener followed a redirect.""" def _post_within( - url: str, body: bytes, headers: Mapping[str, str], timeout: float + url: str, + body: bytes, + headers: Mapping[str, str], + timeout: float, + opener: urllib.request.OpenerDirector, ) -> _Answer: """POST ``body``, and wait for the answer at most ``timeout`` seconds. @@ -410,14 +465,20 @@ def _post_within( failed attempt, like a refused connection. """ finished, answer = _run_with_limit( - lambda: _post(url, body, headers, timeout), timeout + lambda: _post(url, body, headers, timeout, opener), timeout ) if not finished: return _Answer(None, f"TimeoutError: no answer in {timeout:.1f} s") return cast("_Answer", answer) -def _post(url: str, body: bytes, headers: Mapping[str, str], timeout: float) -> _Answer: +def _post( + url: str, + body: bytes, + headers: Mapping[str, str], + timeout: float, + opener: urllib.request.OpenerDirector, +) -> _Answer: """POST ``body``, and return the status, or the failure. The body of the response is not used. So the function reads no body, and @@ -428,7 +489,12 @@ def _post(url: str, body: bytes, headers: Mapping[str, str], timeout: float) -> url, data=body, headers=dict(headers), method="POST" ) try: - with _OPENER.open(request, timeout=timeout) as response: + with opener.open(request, timeout=timeout) as response: + # An opener that follows redirects reports the URL it reached. + if response.geturl() != url: + return _Answer( + response.status, location=response.geturl(), followed_redirect=True + ) return _Answer(response.status) except urllib.error.HTTPError as error: location = error.headers.get("Location") if error.headers else None @@ -440,19 +506,8 @@ def _post(url: str, body: bytes, headers: Mapping[str, str], timeout: float) -> return _Answer(None, f"{type(error).__name__}: {error}") -_PATH_SAFE = "/%!$&'()*+,;=:@-._~?#[]|^" -"""The characters that stay as they are in the path. - -The WHATWG URL parser, which the JavaScript SDK's ``fetch`` uses, -percent-encodes a space, ``"``, ``<``, ``>``, a backtick, ``{``, ``}``, the -control characters, and every non-ASCII character in a path, as UTF-8. -``urllib.parse.quote`` encodes those too, and leaves this set as it is. A -``%`` stays, so a path that is already encoded is not encoded twice. -""" - - def _endpoint_url(endpoint: str, path: str) -> str: - """Join the endpoint and the percent-encoded path. + """Join the endpoint and the path, as the caller wrote the path. ``RunMicrovm`` returns the endpoint as a host name without a scheme. """ @@ -461,4 +516,4 @@ def _endpoint_url(endpoint: str, path: str) -> str: if endpoint.startswith(("https://", "http://")) else (f"https://{endpoint}") ) - return f"{base.rstrip('/')}{urllib.parse.quote(path, safe=_PATH_SAFE)}" + return f"{base.rstrip('/')}{path}" 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 ca6236674..b40104e6a 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 @@ -9,6 +9,7 @@ import socket import threading import time +import urllib.request from collections.abc import Iterator from typing import Any @@ -208,14 +209,44 @@ def test_ends_the_tier_before_the_invocation_reserve( assert sleeps == [0.25, 0.5, 1.0, 2.0] +@pytest.mark.parametrize("remaining", [0.5, 2.9, 10.9]) +def test_names_the_function_timeout_when_too_little_time_is_left( + fake_client: Any, fake_endpoint: Any, remaining: float +) -> None: + # The Lambda default timeout is 3 seconds. Such a function can never + # send a job over HTTP. + client = fake_client() + endpoint = fake_endpoint() + + with pytest.raises(MicrovmEndpointUnavailableError) as raised: + send_job(options(client, endpoint.url, [], remaining_seconds=lambda: remaining)) + + assert str(raised.value) == ( + f"The request step had {remaining:.1f} seconds left in the Lambda " + "invocation, and it needs at least 11 seconds to send the job. Raise " + "the function timeout." + ) + assert client.calls == [] + assert endpoint.requests == [] + + def test_starts_no_call_within_1_second_of_the_reserve( fake_client: Any, fake_endpoint: Any ) -> None: + # The time left falls below the minimum after the first check. client = fake_client() endpoint = fake_endpoint() + readings = iter([12.0]) with pytest.raises(MicrovmEndpointUnavailableError, match="after 0 attempts"): - send_job(options(client, endpoint.url, [], remaining_seconds=lambda: 10.5)) + send_job( + options( + client, + endpoint.url, + [], + remaining_seconds=lambda: next(readings, 10.5), + ) + ) assert client.calls == [] assert endpoint.requests == [] @@ -226,7 +257,7 @@ def test_bounds_the_request_timeout_by_the_reserve( ) -> None: timeouts: list[float] = [] - def post(url: str, body: bytes, headers: Any, timeout: float) -> Any: + def post(url: str, body: bytes, headers: Any, timeout: float, opener: Any) -> Any: timeouts.append(timeout) return request._Answer(202) # noqa: SLF001 @@ -329,7 +360,7 @@ def test_bounds_the_whole_request_and_not_each_socket_operation( started = time.monotonic() answer = request._post_within( # noqa: SLF001 - f"{trickle_server}/job", BODY, {}, 1.0 + f"{trickle_server}/job", BODY, {}, 1.0, request.default_opener() ) assert time.monotonic() - started < 2 @@ -357,23 +388,66 @@ def test_rejects_a_redirect_without_following_it( @pytest.mark.parametrize( - ("path", "sent"), - [ - ("/clone build", "/clone%20build"), - ("/café", "/caf%C3%A9"), - ("/already%20encoded", "/already%20encoded"), - ("/a/b-c_d.e~f!$&'()*+,;=:@[]|^", "/a/b-c_d.e~f!$&'()*+,;=:@[]|^"), - ('/q"<>`{}', "/q%22%3C%3E%60%7B%7D"), - ], + "path", ["/clone%20build", "/caf%C3%A9", "/a/b-c_d.e~f!$&'()*+,;=:@", "/a/"] ) -def test_percent_encodes_the_path_as_fetch_does( - fake_client: Any, fake_endpoint: Any, path: str, sent: str +def test_sends_the_path_as_written( + fake_client: Any, fake_endpoint: Any, path: str ) -> None: + # The workers match the raw request path. So the config and the route + # use the same string only when the path is sent unchanged. endpoint = fake_endpoint() send_job(options(fake_client(), endpoint.url, [], path=path)) - assert endpoint.requests[0]["path"] == sent + assert endpoint.requests[0]["path"] == path + + +def test_uses_the_opener_that_the_caller_passes( + fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint() + opened: list[str] = [] + + class Recording(urllib.request.BaseHandler): + def http_request(self, req: urllib.request.Request) -> urllib.request.Request: + opened.append(req.full_url) + return req + + opener = urllib.request.build_opener(Recording()) + + assert send_job(options(fake_client(), endpoint.url, [], opener=opener)) == 202 + assert opened == [f"{endpoint.url}/durable-execution/v1/job"] + + +@pytest.mark.parametrize("status", [301, 302, 303]) +def test_rejects_a_redirect_that_the_callers_opener_followed( + fake_client: Any, fake_endpoint: Any, status: int +) -> None: + # The default urllib opener follows the redirect with a GET, and the GET + # gets 200. The route never received the job. + endpoint = fake_endpoint([status]) + + with pytest.raises(MicrovmRequestRejectedError) as raised: + send_job( + options( + fake_client(), endpoint.url, [], opener=urllib.request.build_opener() + ) + ) + + assert str(raised.value) == ( + "The MicroVM answered the job at /durable-execution/v1/job with a " + f"redirect to {endpoint.url}/moved, and the opener followed it. The " + "operation does not follow redirects, so set the path that the route " + "serves." + ) + assert [r["path"] for r in endpoint.requests] == [ + "/durable-execution/v1/job", + "/moved", + ] + + +def test_builds_the_default_opener_once() -> None: + assert request.default_opener() is request.default_opener() def test_ignores_proxy_settings_in_the_environment( 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 dcbb26604..c8aed4c16 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 @@ -63,12 +63,33 @@ def config(**fields: Any) -> MicrovmConfig: ( "build", {"request": MicrovmRequestConfig(path="job")}, - 'the request path must be a string that starts with "/"', + "the request path must be an RFC 3986 path", ), ( "build", {"request": MicrovmRequestConfig(path=7)}, # type: ignore[arg-type] - 'the request path must be a string that starts with "/"', + "the request path must be an RFC 3986 path", + ), + ( + "build", + {"request": MicrovmRequestConfig(path="/clone build")}, + 'starts with "/". Percent-encode any other character, for example a ' + "space as %20.", + ), + ( + "build", + {"request": MicrovmRequestConfig(path="/café")}, + "the request path must be an RFC 3986 path", + ), + ( + "build", + {"request": MicrovmRequestConfig(path="/job?x=1")}, + "the request path must be an RFC 3986 path", + ), + ( + "build", + {"request": MicrovmRequestConfig(path="/bad%zz")}, + "the request path must be an RFC 3986 path", ), ( "build", @@ -188,29 +209,12 @@ def test_takes_the_input_and_the_config_by_keyword_only() -> None: class LambdaContext: - def __init__(self, read: Any) -> None: - self.get_remaining_time_in_millis = read - - -def boom() -> int: - msg = "no deadline here" - raise RuntimeError(msg) + def get_remaining_time_in_millis(self) -> int: + return 12_500 @pytest.mark.parametrize( - ("lambda_context", "expected"), - [ - (LambdaContext(lambda: 12_500), 12.5), - (LambdaContext(lambda: 2.5), 0.0025), - (None, None), - (object(), None), - (LambdaContext("not callable"), None), - (LambdaContext(boom), None), - (LambdaContext(lambda: math.inf), None), - (LambdaContext(lambda: math.nan), None), - (LambdaContext(lambda: "12"), None), - (LambdaContext(lambda: True), None), - ], + ("lambda_context", "expected"), [(LambdaContext(), 12.5), (None, None)] ) def test_reads_the_remaining_time_or_reports_no_deadline( lambda_context: Any, expected: float | None From 5eb82c4e819fcb1bfaa0b03ef6b645efca75e2a0 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Thu, 15 Oct 2026 09:25:20 -0700 Subject: [PATCH 05/10] feat(extras): add microvm_session microvm_session launches one MicroVM and runs a handler that sends several jobs to it with vm.invoke. The jobs share the MicroVM's files and processes. The MicroVM is terminated when the handler returns or raises. It ports microvmSession from the JavaScript SDK. - Each job is a child context with a callback and a request step, and always goes over HTTP. - With auto_suspend_on_idle or an idle_policy, each request step calls GetMicrovm first and resumes a suspended MicroVM. After a 502, 503, 504, or no answer, it checks the state again. A MicroVM that ended fails the job at once with MicrovmNotRunningError. - An uncaught MicroVM error reaches the caller with its own class. Any other handler error stays the SDK's ChildContextError. - microvm and the session share the launch, request, and terminate steps. MicrovmBaseConfig holds the settings that both configs share. The Python worker does not suspend its MicroVM yet. The README says so. --- .../README.md | 103 ++- .../microvm/__init__.py | 19 +- .../microvm/config.py | 217 ++++- .../microvm/errors.py | 22 + .../microvm/job.py | 14 +- .../microvm/lifecycle.py | 177 ++++ .../microvm/operation.py | 648 ++++++++------ .../microvm/request.py | 79 +- .../microvm/session.py | 417 +++++++++ .../tests/conftest.py | 39 +- .../tests/microvm/errors_test.py | 4 + .../tests/microvm/job_test.py | 12 + .../tests/microvm/lifecycle_test.py | 213 +++++ .../tests/microvm/operation_test.py | 8 +- .../tests/microvm/request_test.py | 78 ++ .../tests/microvm/session_test.py | 805 ++++++++++++++++++ .../tests/microvm/worker_composed_test.py | 61 ++ 17 files changed, 2600 insertions(+), 316 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py create mode 100644 packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py create mode 100644 packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py create mode 100644 packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index 579cea86c..5c488ddbc 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -95,6 +95,7 @@ Every MicroVM failure is a `MicrovmError`. Catch the base class for all of them, |---|---| | `MicrovmLaunchError` | `RunMicrovm` failed after all retries. When the last attempt timed out or lost its connection, the service may have created the MicroVM anyway, and it then runs the job. So running the job again after this error can run it twice. | | `MicrovmDeliveryError` | HTTP delivery failed after all retries, or the route answered a 4xx status such as 404, or a redirect. | +| `MicrovmNotRunningError` | A session job found its MicroVM terminating, terminated, or removed. It is a `MicrovmDeliveryError`. See [Suspending an idle MicroVM](#suspending-an-idle-microvm). | | `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. | @@ -150,8 +151,108 @@ The function's execution role needs: The MicroVM execution role (`execution_role_arn`) needs `lambda:SendDurableExecutionCallbackSuccess`, `lambda:SendDurableExecutionCallbackFailure`, and `lambda:SendDurableExecutionCallbackHeartbeat` on the function ARN (`arn:aws:lambda:::function::*`). +## `microvm_session`: run several jobs in one MicroVM + +`microvm_session` launches one MicroVM and runs a handler that sends jobs to it. The jobs share the MicroVM's files and processes. So a later job can use what an earlier job wrote. The MicroVM is terminated when the handler returns or raises. + +```python +from typing import Any + +from aws_durable_execution_sdk_python_extras.microvm import ( + MicrovmInvokeOptions, + MicrovmSession, + MicrovmSessionConfig, + microvm_session, +) + +SESSION = MicrovmSessionConfig( + image_identifier=os.environ["MICROVM_IMAGE_ARN"], + execution_role_arn=os.environ["MICROVM_ROLE_ARN"], + timeout=Duration.from_hours(1), +) + + +def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: + build = vm.invoke( + ctx, + "clone-build", + job_input={"repo": "org/app"}, + options=MicrovmInvokeOptions( + path="/clone-build", + timeout=Duration.from_minutes(15), + heartbeat_timeout=Duration.from_minutes(2), + ), + ) + ctx.wait(Duration.from_minutes(5), name="cool-down") + return vm.invoke( + ctx, + "test", + job_input=build, + options=MicrovmInvokeOptions(timeout=Duration.from_minutes(10)), + ) + + +@durable_execution +def handler(event: dict, context: DurableContext) -> Any: + return microvm_session(context, "pipeline", config=SESSION, handler=pipeline) +``` + +The operation runs its durable operations in one child context named `name`: + +``` +pipeline child context +├── pipeline.session step: generates the session ID +├── pipeline.launch step: RunMicrovm +├── clone-build child context, one per job +│ ├── clone-build.callback callback +│ └── clone-build.request step: POSTs the job +├── cool-down wait +├── test +│ ├── test.callback +│ └── test.request +└── pipeline.terminate step: TerminateMicrovm +``` + +- The client token of the launch is the SHA-256 hex digest of the session ID. The session ID is random, and it is checkpointed. So a retried or replayed launch returns the same MicroVM. +- Every job goes over HTTP. The MicroVM gets the `ALL_INGRESS` connector. A job without a `path` goes to the worker's `handler`. +- The first job is sent right after the launch. The MicroVM endpoint refuses requests until the `run` hook returns, and the request step retries a refused request inside its attempt. So the session needs no separate readiness wait. +- `vm.invoke` takes the context that the call runs in, as `microvm` does. Inside a `map` or `parallel` item, pass the item's context. +- `config.timeout` bounds the MicroVM's lifetime: it ends at `timeout` plus 5 minutes, at most 8 hours. Each job waits for its own `timeout`. So set a `heartbeat_timeout` on every job. A job still running when the MicroVM ends then fails at its heartbeat timeout. + +A failed `vm.invoke` raises its MicroVM error inside the handler. The handler can catch it and send another job to the same MicroVM. An error that leaves the handler ends the session: + +1. A MicroVM error reaches the caller with its own class, such as `MicrovmJobFailedError`. +2. Any other error reaches the caller as the SDK's `ChildContextError`, as from `run_in_child_context`. + +An invalid job raises `ValidationError` or `TypeError` from `vm.invoke`, before the job's durable operations start. The handler can catch it too. + +The handler's return value is checkpointed as the session's result. So it must be serializable by the SDK, and it counts toward the checkpoint size limit. + +### Suspending an idle MicroVM + +A running MicroVM pays compute charges. A suspended MicroVM pays only for snapshot storage, and it keeps its memory and files. A session can suspend its MicroVM between jobs in two ways: + +1. `auto_suspend_on_idle=True` asks the worker in the MicroVM to suspend its own MicroVM when no job has run for `auto_suspend_idle_time`. The default idle time is 60 seconds, or the session `timeout` if that is shorter. The idle time must be at least 10 seconds. +2. `idle_policy` sets the `RunMicrovm` idle policy. The service then suspends the MicroVM when its endpoint receives no traffic. A running job receives no traffic after its delivery. So `maxIdleDurationSeconds` must be longer than the longest job. + +With either setting, each request step calls `GetMicrovm` before it sends the job. It resumes a suspended MicroVM, and it sends the job to the endpoint that `GetMicrovm` reports. The MicroVM can suspend itself between that check and the request. The endpoint then answers 502, or the worker answers 503. So after a 502, 503, or 504, or a request without an answer, the step checks the state again before it retries. The wait for a resume counts toward the request's `retry_window`. + +A MicroVM that is terminating, terminated, or removed fails the job at once with `MicrovmNotRunningError`, a subclass of `MicrovmDeliveryError`. A retry finds the same state, so the job is not retried. Without either setting, the session makes no `GetMicrovm` call. 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. + +Suspending is off by default, because it needs more permissions and changes what the MicroVM does between jobs: + +- The function's role needs `lambda:GetMicrovm` and `lambda:ResumeMicrovm`. Without them, every job fails. +- With `auto_suspend_on_idle`, the MicroVM's role needs `lambda:SuspendMicrovm` on the image. That permission lets code in any MicroVM of the image suspend any other. So use one image per trust boundary. +- Background processes that an earlier job started stop while the MicroVM is suspended. The worker counts only its running jobs. + +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) does. The Python worker in this repository does not yet. It accepts the setting, and its MicroVM keeps running. `idle_policy` works with both workers. + +### 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. + ### Subtypes -The JavaScript operation records MicroVM subtypes, such as `Microvm` on the child context and `MicrovmLaunch` on the launch step. This package records the SDK's default subtypes for now, because the Python SDK does not accept custom subtypes yet. +The JavaScript operations record MicroVM subtypes, such as `Microvm` on the child context, `MicrovmSession` on a session, and `MicrovmLaunch` on the launch step. This package records the SDK's default subtypes for now, because the Python SDK does not accept custom subtypes yet. On replay, the SDK compares the subtype of each operation with its checkpoint. A mismatch fails the execution with `NonDeterministicExecutionError`. So a release that changed the subtypes would fail every execution that waits on a MicroVM during the deploy. So this package is published only after the SDK accepts custom subtypes, and it records the MicroVM subtypes from its first release. diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/__init__.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/__init__.py index 32572cbde..1c0c04634 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/__init__.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/__init__.py @@ -5,20 +5,23 @@ This module is experimental. Its API can change in any release. -The exports match the parts of ``@aws/durable-execution-sdk-js-extras/microvm`` -that this package has so far. The session operation comes in a later change. +The exports match ``@aws/durable-execution-sdk-js-extras/microvm``. """ from aws_durable_execution_sdk_python_extras.microvm.config import ( MAX_MICROVM_DURATION_SECONDS, + MicrovmBaseConfig, MicrovmConfig, + MicrovmInvokeOptions, MicrovmRequestConfig, + MicrovmSessionConfig, ) from aws_durable_execution_sdk_python_extras.microvm.errors import ( MicrovmDeliveryError, MicrovmError, MicrovmJobFailedError, MicrovmLaunchError, + MicrovmNotRunningError, MicrovmTimeoutError, ) from aws_durable_execution_sdk_python_extras.microvm.job import ( @@ -35,23 +38,35 @@ from aws_durable_execution_sdk_python_extras.microvm.retry import ( default_microvm_retry_strategy, ) +from aws_durable_execution_sdk_python_extras.microvm.session import ( + DEFAULT_AUTO_SUSPEND_IDLE_SECONDS, + MicrovmSession, + microvm_session, +) __all__ = [ + "DEFAULT_AUTO_SUSPEND_IDLE_SECONDS", "DEFAULT_MICROVM_JOB_PATH", "DEFAULT_MICROVM_PORT", "MAX_MICROVM_DURATION_SECONDS", "MAX_RUN_HOOK_PAYLOAD_LENGTH", + "MicrovmBaseConfig", "MicrovmConfig", "MicrovmDeliveryError", "MicrovmEndpointUnavailableError", "MicrovmError", + "MicrovmInvokeOptions", "MicrovmJobFailedError", "MicrovmLaunchError", + "MicrovmNotRunningError", "MicrovmRequestConfig", "MicrovmRequestRejectedError", "MicrovmRequestUnauthorizedError", + "MicrovmSession", + "MicrovmSessionConfig", "MicrovmTimeoutError", "default_microvm_retry_strategy", "microvm", + "microvm_session", ] 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 35290ae35..d9ac89c74 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 @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. # # SPDX-License-Identifier: Apache-2.0 -"""The configuration of :func:`microvm`.""" +"""The configuration of :func:`microvm` and :func:`microvm_session`.""" from __future__ import annotations @@ -75,12 +75,9 @@ class MicrovmRequestConfig: @dataclass(frozen=True, kw_only=True) -class MicrovmConfig: - """The MicroVM, and how long to wait for its result. - - :func:`microvm` validates the config before it starts any durable - operation. - """ +class MicrovmBaseConfig: + """The settings that :class:`MicrovmConfig` and + :class:`MicrovmSessionConfig` share.""" image_identifier: str """The ARN or ID of the MicroVM image to run. @@ -100,32 +97,6 @@ class MicrovmConfig: ARN. """ - timeout: Duration - """The longest time to wait for the MicroVM to report a result. - - It must be between 1 second and 8 hours, because a MicroVM lives at most - 8 hours. The MicroVM's ``maximumDurationInSeconds`` is this timeout plus 5 - minutes, at most 8 hours. So a timeout above 7 hours 55 minutes leaves the - MicroVM less than 5 minutes beyond the callback's timeout. - """ - - heartbeat_timeout: Duration | None = None - """The longest time between two heartbeats from the MicroVM. - - When it is set, a MicroVM that fails to boot or stops responding fails the - operation after this time, instead of after ``timeout``. The MicroVM - receives the value in the job document. - - The heartbeat timeout starts when the callback is created, before the - launch. The worker sends its first heartbeat when the job arrives. So the - first period must cover the launch, the boot, and, for HTTP delivery, the - request step, each with its retries. A throttled launch can take several - minutes, and the request step retries for up to ``request.retry_window`` - in each attempt. A heartbeat timeout that is shorter fails the operation - with ``MicrovmTimeoutError`` before the MicroVM can send its first - heartbeat. - """ - image_version: str | None = None """The image version to run. Defaults to the latest active version.""" @@ -137,11 +108,12 @@ class MicrovmConfig: 1. ``NO_INGRESS`` when the job goes in the ``run`` hook. The job then needs no inbound connection. 2. ``ALL_INGRESS`` when the job goes over HTTP. The request must reach - the MicroVM's endpoint. + the MicroVM's endpoint. Every job of a session goes over HTTP. An empty tuple sends no ingress connector. Connectors that you set must - allow ingress whenever an input can exceed the ``run`` hook limit, because - such an input goes over HTTP. + allow ingress whenever a job goes over HTTP: for :func:`microvm`, + whenever an input can exceed the ``run`` hook limit, and for a session, + always. """ egress_network_connectors: tuple[str, ...] | None = None @@ -156,13 +128,6 @@ class MicrovmConfig: """The ``logging`` member of ``RunMicrovm``. Defaults to the service default.""" - request: MicrovmRequestConfig | None = None - """How to deliver the job over HTTP. - - ``None`` delivers a job in the ``run`` hook when it fits, and over HTTP to - :data:`DEFAULT_MICROVM_JOB_PATH` on port 8080 when it does not. - """ - retry_strategy: Callable[[Exception, int], RetryDecision] | None = None """The step retry strategy of the launch, request, and terminate steps. @@ -201,3 +166,169 @@ class MicrovmConfig: shorter than the function's timeout, because a call that runs past it records no outcome for the step attempt. """ + + +@dataclass(frozen=True, kw_only=True) +class MicrovmConfig(MicrovmBaseConfig): + """The MicroVM of :func:`microvm`, and how long to wait for its result. + + :func:`microvm` validates the config before it starts any durable + operation. + """ + + timeout: Duration + """The longest time to wait for the MicroVM to report a result. + + It must be between 1 second and 8 hours, because a MicroVM lives at most + 8 hours. The MicroVM's ``maximumDurationInSeconds`` is this timeout plus 5 + minutes, at most 8 hours. So a timeout above 7 hours 55 minutes leaves the + MicroVM less than 5 minutes beyond the callback's timeout. + """ + + heartbeat_timeout: Duration | None = None + """The longest time between two heartbeats from the MicroVM. + + When it is set, a MicroVM that fails to boot or stops responding fails the + operation after this time, instead of after ``timeout``. The MicroVM + receives the value in the job document. + + The heartbeat timeout starts when the callback is created, before the + launch. The worker sends its first heartbeat when the job arrives. So the + first period must cover the launch, the boot, and, for HTTP delivery, the + request step, each with its retries. A throttled launch can take several + minutes, and the request step retries for up to ``request.retry_window`` + in each attempt. A heartbeat timeout that is shorter fails the operation + with ``MicrovmTimeoutError`` before the MicroVM can send its first + heartbeat. + """ + + request: MicrovmRequestConfig | None = None + """How to deliver the job over HTTP. + + ``None`` delivers a job in the ``run`` hook when it fits, and over HTTP to + :data:`DEFAULT_MICROVM_JOB_PATH` on port 8080 when it does not. + """ + + +@dataclass(frozen=True, kw_only=True) +class MicrovmSessionConfig(MicrovmBaseConfig): + """The MicroVM of :func:`microvm_session`. + + :func:`microvm_session` validates the config before it starts any durable + operation. + """ + + timeout: Duration + """The lifetime of the session's MicroVM, including durable waits between + jobs. + + The MicroVM's ``maximumDurationInSeconds`` is this timeout plus 5 + minutes, at most 8 hours. The service then terminates the MicroVM. A + MicroVM lives at most 8 hours, running or suspended. So the timeout must + be between 1 second and 8 hours. The session also uses it to cap the + default :attr:`auto_suspend_idle_time`. + + The timeout does not limit the handler or its jobs. Each job waits for + its own ``timeout``. For example, a session with a 1-hour timeout starts + a job with an 8-hour timeout at 0:58. The service terminates the MicroVM + at about 1:05. + + 1. When the image enables the ``terminate`` hook, a worker that acts on + it fails the job at about 1:05. + 2. Lambda does not call the hook for a suspended MicroVM, or when the + image disables it. The job then fails at its ``heartbeat_timeout``, or + at 8:58 without one. + + So set a ``heartbeat_timeout`` on every :meth:`MicrovmSession.invoke`. + """ + + idle_policy: Mapping[str, Any] | None = None + """The ``idlePolicy`` member of ``RunMicrovm``. Not set by default. + + It needs all three members, for example ``{"maxIdleDurationSeconds": 300, + "suspendedDurationSeconds": 3600, "autoResumeEnabled": True}``. + + An idle policy counts only inbound traffic to the endpoint, and a running + job receives none after its delivery. So ``maxIdleDurationSeconds`` must + be longer than the longest job, or the service suspends the MicroVM + during that job. Set it only for a MicroVM that serves inbound traffic of + its own. With an idle policy, the request step checks the MicroVM state + before each job, and resumes a suspended MicroVM. That needs + ``lambda:GetMicrovm`` and ``lambda:ResumeMicrovm`` in the function's + role. + """ + + auto_suspend_on_idle: bool = False + """Whether the worker suspends the MicroVM when no job has run for + :attr:`auto_suspend_idle_time`. Defaults to ``False``. + + A running MicroVM pays compute charges. A suspended MicroVM pays only for + snapshot storage, and keeps its memory and files. With this setting, the + session asks the worker in the ``run`` hook payload to suspend its own + MicroVM when idle. Before each job, the request step checks the MicroVM + state, and resumes a suspended MicroVM. + + It is off by default, because it needs more permissions and changes what + the MicroVM does between jobs: + + 1. The function's role needs ``lambda:GetMicrovm`` and + ``lambda:ResumeMicrovm``. Without them, every job fails. + 2. The MicroVM's role needs ``lambda:SuspendMicrovm`` on the image. That + permission lets code in any MicroVM of the image suspend any other. + Use one image per trust boundary. + 3. Background processes that an earlier job started stop while the + MicroVM is suspended. "Idle" means only that no job runs in the + worker. + + A resume adds latency to the next job, about 1 second in us-east-1, and + the job's ``timeout`` and ``heartbeat_timeout`` count it. Suspending does + not extend the MicroVM's lifetime. + + The worker must be a version that supports the setting. The JavaScript + worker ``@aws/durable-execution-sdk-js-microvm-worker`` does. The Python + worker in this repository does not yet: it accepts the setting, and its + MicroVM keeps running. + """ + + auto_suspend_idle_time: Duration | None = None + """How long the worker waits with no running job before it suspends the + MicroVM. + + ``None`` means 60 seconds, or the session ``timeout`` if that is shorter. + A session whose ``timeout`` is under 10 seconds then does not suspend. + It applies only when :attr:`auto_suspend_on_idle` is ``True``, and + setting it otherwise fails validation. It must be at least 10 seconds, + and not longer than the session ``timeout``. A resumed MicroVM needs a + few seconds before its job arrives, so a shorter time could suspend it + again first. + """ + + +@dataclass(frozen=True, kw_only=True) +class MicrovmInvokeOptions(MicrovmRequestConfig): + """One job of :meth:`MicrovmSession.invoke`. + + The job always goes over HTTP: to ``path`` when it is set, and otherwise + to :data:`DEFAULT_MICROVM_JOB_PATH`, which the worker serves with its + ``handler``. + """ + + timeout: Duration + """The longest time to wait for this job's result. At most 8 hours. + + The time starts before the job is delivered. So it includes a resume of + a suspended MicroVM. The session's ``timeout`` does not shorten it. + """ + + heartbeat_timeout: Duration | None = None + """The longest time between two heartbeats from the MicroVM for this job. + + The first period starts before the job is delivered, and the worker sends + its first heartbeat when the job arrives. So the first period includes a + resume of a suspended MicroVM and the request step with its retries. + Keep it longer than both together. + + 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. + """ 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 8480ab52c..2773cbf90 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 @@ -98,6 +98,27 @@ class MicrovmDeliveryError(MicrovmError): """The job could not be delivered to the MicroVM.""" +class MicrovmNotRunningError(MicrovmDeliveryError): + """The session's MicroVM can no longer run a job. + + The MicroVM is terminating, terminated, or removed. A session MicroVM + lives at most its session ``timeout`` plus 5 minutes, and never longer + than 8 hours, running or suspended. A ``TerminateMicrovm`` call from + outside the session ends it too. It is a :class:`MicrovmDeliveryError`, + so code that catches delivery failures also catches it. + + Only a session that checks the MicroVM state before each job raises it. + It checks when ``auto_suspend_on_idle`` is ``True``, or when an + ``idle_policy`` is set. A later attempt finds the same state, so the job + is not retried. To continue, start a new session. + + Without the check, the session makes no ``GetMicrovm`` call. 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 fails with a plain + :class:`MicrovmDeliveryError`, after the retries of both tiers. + """ + + class MicrovmJobFailedError(MicrovmError): """The job handler in the MicroVM failed. @@ -124,6 +145,7 @@ class MicrovmTimeoutError(MicrovmError): MicrovmError, MicrovmLaunchError, MicrovmDeliveryError, + MicrovmNotRunningError, MicrovmJobFailedError, MicrovmTimeoutError, ): 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 454711ccd..9cc7979d6 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 @@ -117,13 +117,21 @@ def to_json(self) -> str: ) -def run_hook_payload_without_job(region: str) -> str: - """The ``runHookPayload`` of a launch whose job goes over HTTP. +def run_hook_payload_without_job( + region: str, auto_suspend_idle_seconds: int | None = None +) -> str: + """The ``runHookPayload`` of a launch whose jobs go over HTTP. The worker rejects a ``run`` hook payload without ``version`` and ``region``. So the payload keeps those two fields, and omits the job. + + A session that suspends when idle adds ``autoSuspendIdleSeconds``. The + worker then suspends its own MicroVM when no job has run for that long. """ - return dumps({"version": PAYLOAD_VERSION, "region": region}) + document: dict[str, Any] = {"version": PAYLOAD_VERSION, "region": region} + if auto_suspend_idle_seconds is not None: + document["autoSuspendIdleSeconds"] = auto_suspend_idle_seconds + return dumps(document) def dumps(document: Any) -> str: 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 new file mode 100644 index 000000000..b4754c6b4 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Resume a suspended session MicroVM before its next job.""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from aws_durable_execution_sdk_python.types import LoggerInterface +from botocore.exceptions import ClientError + +from aws_durable_execution_sdk_python_extras.microvm.request import ( + MicrovmEndpointUnavailableError, + call_before_deadline, + reserve_starts_at, +) + + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python_extras.microvm.config import ( + LambdaMicrovmsClient, + ) + + +_INITIAL_POLL_SECONDS = 0.25 +_MAX_POLL_SECONDS = 2.0 +"""The first poll comes soon, because a resume measured in us-east-1 took +about 1 second. Later polls back off, so a slow transition costs few +``GetMicrovm`` calls.""" + + +class MicrovmStateError(Exception): + """The MicroVM is terminating, terminated, or removed. + + The request step raises it. A later attempt finds the same state, so the + default retry strategy does not retry it. The session job then raises + :class:`MicrovmNotRunningError`. + """ + + def __init__(self, message: str, state: str | None = None) -> None: + super().__init__(message) + self.state = state + """The MicroVM state, or ``None`` when the MicroVM no longer exists.""" + + +@dataclass(frozen=True, kw_only=True) +class EnsureRunningOptions: + """The inputs of :func:`ensure_running`.""" + + client: LambdaMicrovmsClient + microvm_id: str + max_wait_seconds: float + """The longest time to wait for a transition.""" + remaining_seconds: Callable[[], float | None] + """The remaining invocation time. ``None`` means that the compute reports + no deadline, and only ``max_wait_seconds`` applies.""" + logger: LoggerInterface + sleep: Callable[[float], None] = time.sleep + + +def ensure_running(options: EnsureRunningOptions) -> str | None: + """Make sure the MicroVM can take a job, and return its current endpoint. + + A session MicroVM can be suspended between jobs, by its worker or by an + idle policy. So the next job can find the MicroVM in any of these states: + + 1. ``RUNNING``: the function returns at once. + 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 + ``RUNNING``. The MicroVM stays ``SUSPENDED`` while its ``resume`` hook + runs. So 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 + :class:`MicrovmStateError`. + + 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 + times out. It then raises ``MicrovmEndpointUnavailableError``, and the + step retry strategy takes over. Each ``GetMicrovm`` and ``ResumeMicrovm`` + call also stops 10 seconds before the invocation times out, with the same + error. + + Returns: + The endpoint that ``GetMicrovm`` reports, or ``None`` when it reports + none. + + Raises: + MicrovmStateError: When the MicroVM can no longer run a job. + MicrovmEndpointUnavailableError: When the MicroVM does not reach + ``RUNNING`` in time. + """ + 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 + + def unavailable() -> Exception: + elapsed_ms = round((time.monotonic() - started) * 1000) + return MicrovmEndpointUnavailableError( + f"MicroVM {microvm_id} did not reach RUNNING in {elapsed_ms} ms. " + f"Last state: {last_state or 'unknown'}." + ) + + # Each call must end before the invocation reserve. A Lambda timeout + # during the step attempt would record no outcome. + def get_microvm() -> Any: + return call_before_deadline( + lambda: options.client.get_microvm(microvmIdentifier=microvm_id), + reserve, + unavailable, + ) + + def resume_microvm() -> Any: + return call_before_deadline( + lambda: options.client.resume_microvm(microvmIdentifier=microvm_id), + reserve, + unavailable, + ) + + while True: + try: + current = get_microvm() + except ClientError as error: + if _code(error) == "ResourceNotFoundException": + msg = ( + f"MicroVM {microvm_id} no longer exists. A session MicroVM " + "lives at most its session timeout plus 5 minutes, and never " + "longer than 8 hours." + ) + raise MicrovmStateError(msg) from error + raise + state = current.get("state") + last_state = state + if state == "RUNNING" or (state == "PENDING" and not resume_sent): + if resume_sent: + 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") + msg = f"MicroVM {microvm_id} is {state}{f': {reason}' if reason else ''}." + raise MicrovmStateError(msg, state) + if state == "SUSPENDED" and not resume_sent: + try: + resume_microvm() + except ClientError as error: + # ConflictException means that a transition is already in + # progress, for example an auto-resume. Polling continues + # either way. + if _code(error) != "ConflictException": + raise + resume_sent = True + options.logger.info("resuming MicroVM", extra={"state": state}) + + if time.monotonic() + delay > deadline: + raise unavailable() + options.sleep(delay) + delay = min(delay * 2, _MAX_POLL_SECONDS) + + +def _code(error: ClientError) -> str | None: + return error.response.get("Error", {}).get("Code") 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 53a50b3c6..1f2b79051 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 @@ -10,7 +10,8 @@ import os import re import threading -from collections.abc import Callable +import time +from collections.abc import Callable, Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -35,6 +36,7 @@ from aws_durable_execution_sdk_python_extras.microvm.config import ( MAX_MICROVM_DURATION_SECONDS, LambdaMicrovmsClient, + MicrovmBaseConfig, MicrovmConfig, MicrovmRequestConfig, ) @@ -43,6 +45,7 @@ MicrovmError, MicrovmJobFailedError, MicrovmLaunchError, + MicrovmNotRunningError, MicrovmTimeoutError, ) from aws_durable_execution_sdk_python_extras.microvm.job import ( @@ -51,12 +54,18 @@ fits_run_hook, run_hook_payload_without_job, ) +from aws_durable_execution_sdk_python_extras.microvm.lifecycle import ( + EnsureRunningOptions, + MicrovmStateError, + ensure_running, +) from aws_durable_execution_sdk_python_extras.microvm.request import ( DEFAULT_MICROVM_JOB_PATH, DEFAULT_MICROVM_PORT, DEFAULT_REQUEST_RETRY_WINDOW_SECONDS, MAX_REQUEST_RETRY_WINDOW_SECONDS, SendJobOptions, + require_time_for_call, send_job, ) from aws_durable_execution_sdk_python_extras.microvm.retry import ( @@ -231,136 +240,136 @@ def run(child: DurableContext) -> Any: raise MicrovmError.from_child_context_error(name, error) from error +_STATE_ERROR_TYPE = f"{MicrovmStateError.__module__}.{MicrovmStateError.__qualname__}" +"""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.""" + + +def _validate_base(name: str, config: MicrovmBaseConfig, timeout: Duration) -> int: + """Validate the settings that every MicroVM operation shares. + + Returns: + The timeout in seconds. + + Raises: + ValidationError: When ``name`` or a setting is not valid. + """ + if not isinstance(name, str) or not name: + msg = "a MicroVM operation requires a non-empty name" + raise ValidationError(msg) + if not config.image_identifier: + msg = f'MicroVM "{name}": image_identifier is required' + raise ValidationError(msg) + if not config.execution_role_arn: + msg = f'MicroVM "{name}": execution_role_arn is required' + raise ValidationError(msg) + return _validate_timeout(name, "timeout", timeout) + + +def _validate_timeout(name: str, field: str, timeout: Duration) -> int: + """Validate a job or MicroVM timeout, and return it in seconds. + + Raises: + ValidationError: When it is not between 1 second and 8 hours. + """ + seconds = timeout.to_seconds() + if not 1 <= seconds <= MAX_MICROVM_DURATION_SECONDS: + msg = ( + f'MicroVM "{name}": {field} must be between 1 second and ' + f"{MAX_MICROVM_DURATION_SECONDS} seconds (8 hours), because a " + f"MicroVM lives at most 8 hours. Got {seconds} seconds." + ) + raise ValidationError(msg) + return seconds + + +def _validate_heartbeat(name: str, heartbeat_timeout: Duration | None) -> int | None: + """Validate a heartbeat timeout, and return it in seconds. + + Raises: + ValidationError: When it is shorter than 1 second. + """ + if heartbeat_timeout is None: + return None + seconds = heartbeat_timeout.to_seconds() + if seconds < 1: + msg = f'MicroVM "{name}": heartbeat_timeout must be at least 1 second' + raise ValidationError(msg) + return seconds + + +def _validate_input(name: str, job_input: Any) -> None: + """Check that the job input is JSON-serializable. + + Raises: + TypeError: When it is not. + """ + try: + json.dumps(job_input, allow_nan=False) + except (TypeError, ValueError) as error: + msg = f'MicroVM "{name}": the job input must be JSON-serializable: {error}' + raise TypeError(msg) from error + + +def _validate_region(context: DurableContext, name: str) -> tuple[str, str]: + """Return the partition and the Region of the durable execution. + + Raises: + ValidationError: When the Region cannot be determined. + """ + partition, region = _partition_and_region( + context.execution_context.durable_execution_arn + ) + if not region: + # The worker rejects a payload without a Region. So the operation + # fails before any durable operation instead. + msg = ( + f'MicroVM "{name}": cannot determine the AWS Region. The durable ' + "execution ARN has none, and AWS_REGION is not set." + ) + raise ValidationError(msg) + return partition, region + + @dataclass(frozen=True, kw_only=True) -class _OperationScope: - """The validated values that every step of one operation uses.""" +class _MicrovmScope: + """The validated values and the steps that every MicroVM operation uses. + + :func:`microvm` and :func:`microvm_session` share the launch, request, + and terminate steps. + """ name: str - config: MicrovmConfig + config: MicrovmBaseConfig partition: str region: str timeout_seconds: int - heartbeat_timeout_seconds: int | None - request: MicrovmRequestConfig - retry_window_seconds: int retry_strategy: Callable[[Exception, int], RetryDecision] - @classmethod - def create( - cls, - context: DurableContext, - name: str, - job_input: Any, - config: MicrovmConfig, - ) -> _OperationScope: - """Validate the call, before any durable operation starts. + def launch( + self, + child: DurableContext, + *, + token_source: str, + payload: str, + delivery: str, + idle_policy: Mapping[str, Any] | None = None, + ) -> LaunchResult: + """Run the launch step, ``.launch``. + + The client token is the SHA-256 hex digest of ``token_source``. It + must be unique per launch and stable across retries and replays. Raises: - ValidationError: When ``name`` or ``config`` is not valid, or the - Region cannot be determined. - TypeError: When ``job_input`` is not JSON-serializable. + MicrovmLaunchError: When ``RunMicrovm`` fails after all retries. """ - if not isinstance(name, str) or not name: - msg = "a MicroVM operation requires a non-empty name" - raise ValidationError(msg) - if not config.image_identifier: - msg = f'MicroVM "{name}": image_identifier is required' - raise ValidationError(msg) - if not config.execution_role_arn: - msg = f'MicroVM "{name}": execution_role_arn is required' - raise ValidationError(msg) - timeout_seconds = config.timeout.to_seconds() - if not 1 <= timeout_seconds <= MAX_MICROVM_DURATION_SECONDS: - msg = ( - f'MicroVM "{name}": timeout must be between 1 second and ' - f"{MAX_MICROVM_DURATION_SECONDS} seconds (8 hours), because a " - f"MicroVM lives at most 8 hours. Got {timeout_seconds} seconds." - ) - raise ValidationError(msg) - heartbeat_timeout_seconds = ( - None - if config.heartbeat_timeout is None - else config.heartbeat_timeout.to_seconds() - ) - if heartbeat_timeout_seconds is not None and heartbeat_timeout_seconds < 1: - msg = f'MicroVM "{name}": heartbeat_timeout must be at least 1 second' - raise ValidationError(msg) - request = config.request or MicrovmRequestConfig() - retry_window_seconds = _validate_request(name, request) try: - json.dumps(job_input, allow_nan=False) - except (TypeError, ValueError) as error: - msg = f'MicroVM "{name}": the job input must be JSON-serializable: {error}' - raise TypeError(msg) from error - partition, region = _partition_and_region( - context.execution_context.durable_execution_arn - ) - if not region: - # The worker rejects a payload without a Region. So the operation - # fails before any durable operation instead. - msg = ( - f'MicroVM "{name}": cannot determine the AWS Region. The durable ' - "execution ARN has none, and AWS_REGION is not set." - ) - raise ValidationError(msg) - return cls( - name=name, - config=config, - partition=partition, - region=region, - timeout_seconds=timeout_seconds, - heartbeat_timeout_seconds=heartbeat_timeout_seconds, - request=request, - retry_window_seconds=retry_window_seconds, - retry_strategy=config.retry_strategy or default_microvm_retry_strategy, - ) - - def run_job(self, child: DurableContext, job_input: Any) -> Any: - """Run the operation's durable operations in its child context.""" - name = self.name - callback: Callback = child.create_callback( - name=f"{name}.callback", - config=CallbackConfig( - timeout=Duration.from_seconds(self.timeout_seconds), - heartbeat_timeout=Duration.from_seconds( - self.heartbeat_timeout_seconds or 0 - ), - # The worker sends the result as a JSON string. The default - # callback serdes would return that string unparsed. - serdes=JsonSerDes(), - ), - ) - # The operation picks the delivery: - # - # 1. The run hook is preferred. It opens no ingress, and it needs no - # extra durable step. - # 2. A config with a request path goes over HTTP. The worker runs a - # run hook job with its handler, not with a route. - # 3. A payload larger than the run hook limit goes over HTTP. - # RunMicrovm rejects it with ValidationException, which the - # default retry strategy does not retry. - # - # The payload contains the callback ID, so the size is known only - # after the callback exists. The ID comes from the checkpoint on - # replay, and the input must not change. So every replay picks the - # same delivery as the first attempt. - with_job = RunHookPayload( - region=self.region, - callback_id=callback.callback_id, - job_input=job_input, - heartbeat_timeout_seconds=self.heartbeat_timeout_seconds, - ).to_json() - over_http = self.request.path is not None or not fits_run_hook(with_job) - delivery = HTTP_DELIVERY if over_http else RUN_HOOK_DELIVERY - payload = run_hook_payload_without_job(self.region) if over_http else with_job - - try: - launched = LaunchResult.from_dict( + return LaunchResult.from_dict( child.step( - lambda _: self._launch( - callback.callback_id, payload, delivery + lambda _: self._launch_call( + token_source, payload, delivery, idle_policy ).to_dict(), - name=f"{name}.launch", + name=f"{self.name}.launch", config=StepConfig( retry_strategy=self.retry_strategy, # Plain JSON, as the JavaScript SDK checkpoints it. So @@ -372,59 +381,38 @@ def run_job(self, child: DurableContext, job_input: Any) -> Any: ) ) except StepError as error: - msg = f'MicroVM "{name}": the launch failed: {_detail(error)}' + msg = f'MicroVM "{self.name}": the launch failed: {_detail(error)}' raise MicrovmLaunchError( message=msg, data=error.data, stack_trace=error.stack_trace ) from error - microvm_id = launched.microvm_id - # Replay computes the delivery again from the current input and - # config. A deploy between the launch and a replay can change both. - # The launch result comes from the checkpoint on replay, and it - # records the delivery that the launch prepared for. So the operation - # follows the launch's delivery, not the new computation. - if launched.delivery not in {RUN_HOOK_DELIVERY, HTTP_DELIVERY}: - # The launch checkpoint comes from a release with a delivery that - # this release does not know, for example after a rollback. This - # release cannot deliver the job. So it terminates the MicroVM and - # fails, the same way on every replay. - self._terminate(child, microvm_id) - msg = ( - f'MicroVM "{name}": the launch prepared the delivery ' - f'"{launched.delivery}", which this release does not support' - ) - raise MicrovmDeliveryError(msg) - - if launched.delivery == HTTP_DELIVERY: - try: - self._deliver(child, launched, callback.callback_id, job_input) - except StepError as error: - self._terminate(child, microvm_id) - msg = f'MicroVM "{name}": the job could not be delivered: {_detail(error)}' - raise MicrovmDeliveryError( - message=msg, data=error.data, stack_trace=error.stack_trace - ) from error - try: - result = callback.result() - except Exception as error: - # SuspendExecution is a BaseException, so a wait for the callback - # passes through here without terminating the MicroVM. - self._terminate(child, microvm_id) - raise self._job_error(callback, error) from error - self._terminate(child, microvm_id) - return result - - def _deliver( + def deliver( self, child: DurableContext, + *, + job_name: str, launched: LaunchResult, callback_id: str, job_input: Any, + heartbeat_timeout_seconds: int | None, + request: MicrovmRequestConfig, + retry_window_seconds: int, + resume: bool = False, ) -> None: - """Run the request step: POST the job to the MicroVM's endpoint. + """Run the request step, ``.request``: POST the job to the + MicroVM's endpoint. The step retries inside each attempt with :func:`send_job`, the first tier. The step retry strategy is the second tier. + + 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. + + Raises: + MicrovmNotRunningError: When the MicroVM can no longer run a job. + MicrovmDeliveryError: When the delivery fails after all retries. """ body = ( JobRequest( @@ -432,13 +420,14 @@ def _deliver( microvm_id=launched.microvm_id, callback_id=callback_id, job_input=job_input, - heartbeat_timeout_seconds=self.heartbeat_timeout_seconds, + heartbeat_timeout_seconds=heartbeat_timeout_seconds, ) .to_json() .encode() ) - request = self.request 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 def attempt(step_context: StepContext) -> None: # The launch records an endpoint for every HTTP delivery. So a @@ -450,32 +439,99 @@ def attempt(step_context: StepContext) -> None: f"{launched.microvm_id}, so the job cannot be sent over HTTP" ) raise RuntimeError(msg) + client = self._client() + # 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() + endpoint = launched.endpoint + recheck: Callable[[float], str | None] | None = None + if resume: + 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, + ) + ) + + endpoint = running(retry_window_seconds) or endpoint + recheck = running send_job( SendJobOptions( - client=self._client(), + client=client, microvm_id=launched.microvm_id, - endpoint=launched.endpoint, - 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, + endpoint=endpoint, + path=path, + port=port, body=body, - retry_window_seconds=self.retry_window_seconds, + retry_window_seconds=retry_window_seconds, remaining_seconds=remaining_seconds, logger=step_context.logger, opener=self.config.opener, + window_started_at=window_started_at, + recheck=recheck, ) ) - child.step( - attempt, - name=f"{self.name}.request", - config=StepConfig(retry_strategy=self.retry_strategy), - ) + try: + child.step( + attempt, + name=f"{job_name}.request", + config=StepConfig(retry_strategy=self.retry_strategy), + ) + except StepError as error: + error_class = ( + MicrovmNotRunningError + if error.error_type == _STATE_ERROR_TYPE + else MicrovmDeliveryError + ) + msg = f'MicroVM "{job_name}": the job could not be delivered: {_detail(error)}' + raise error_class( + message=msg, data=error.data, stack_trace=error.stack_trace + ) from error + + def terminate(self, child: DurableContext, microvm_id: str) -> None: + """Run the terminate step, and log a final failure instead of raising. + + The job's outcome is already recorded when terminate runs, and the + MicroVM's maximum duration bounds how long it can keep running. So a + terminate failure must not discard the outcome. The step's failure is + checkpointed. So on replay the step fails again at once, with the same + outcome. + """ + try: + child.step( + lambda _: self._terminate_call(microvm_id), + name=f"{self.name}.terminate", + config=StepConfig(retry_strategy=self.retry_strategy), + ) + except StepError as error: + # The MicroVM keeps running, and is billed, until the platform + # ends it. A missing permission fails every terminate this way, + # and the default retry strategy does not retry + # AccessDeniedException. This log is the only sign of it. So it + # is an error, not a warning. + child.logger.error( + f'MicroVM "{self.name}": terminate failed for {microvm_id}. The ' + "MicroVM keeps running, and is billed, until the platform ends " + "it at its maximum duration: the operation timeout plus 5 " + "minutes, at most 8 hours. A missing permission needs " + "lambda:TerminateMicrovm in the function's execution role.", + extra={"microvmId": microvm_id, "error": error.message}, + ) - def _launch(self, callback_id: str, payload: str, delivery: str) -> LaunchResult: + def _launch_call( + self, + token_source: str, + payload: str, + delivery: str, + idle_policy: Mapping[str, Any] | None, + ) -> LaunchResult: """Call ``RunMicrovm``, and return what the launch step checkpoints.""" over_http = delivery == HTTP_DELIVERY config = self.config @@ -484,7 +540,7 @@ def _launch(self, callback_id: str, payload: str, delivery: str) -> LaunchResult "executionRoleArn": config.execution_role_arn, # A callback ID can be 1,024 characters, and the client token # allows 128. The SHA-256 hex digest is 64 characters. - "clientToken": hashlib.sha256(callback_id.encode()).hexdigest(), + "clientToken": hashlib.sha256(token_source.encode()).hexdigest(), # An explicit empty tuple stays empty. The service accepts 0 to # 10 connectors, and only None gets the default. # An HTTP delivery must reach the MicroVM's endpoint, and a run @@ -509,6 +565,8 @@ def _launch(self, callback_id: str, payload: str, delivery: str) -> LaunchResult request["imageVersion"] = config.image_version if config.logging is not None: request["logging"] = dict(config.logging) + if idle_policy is not None: + request["idlePolicy"] = dict(idle_policy) response = self._client().run_microvm(**request) microvm_id = response.get("microvmId") if not microvm_id: @@ -524,36 +582,6 @@ def _launch(self, callback_id: str, payload: str, delivery: str) -> LaunchResult endpoint=endpoint if over_http else None, ) - def _terminate(self, child: DurableContext, microvm_id: str) -> None: - """Run the terminate step, and log a final failure instead of raising. - - The job's outcome is already recorded when terminate runs, and the - MicroVM's maximum duration bounds how long it can keep running. So a - terminate failure must not discard the outcome. The step's failure is - checkpointed. So on replay the step fails again at once, with the same - outcome. - """ - try: - child.step( - lambda _: self._terminate_call(microvm_id), - name=f"{self.name}.terminate", - config=StepConfig(retry_strategy=self.retry_strategy), - ) - except StepError as error: - # The MicroVM keeps running, and is billed, until the platform - # ends it. A missing permission fails every terminate this way, - # and the default retry strategy does not retry - # AccessDeniedException. This log is the only sign of it. So it - # is an error, not a warning. - child.logger.error( - f'MicroVM "{self.name}": terminate failed for {microvm_id}. The ' - "MicroVM keeps running, and is billed, until the platform ends " - "it at its maximum duration: the operation timeout plus 5 " - "minutes, at most 8 hours. A missing permission needs " - "lambda:TerminateMicrovm in the function's execution role.", - extra={"microvmId": microvm_id, "error": error.message}, - ) - def _terminate_call(self, microvm_id: str) -> None: try: self._client().terminate_microvm(microvmIdentifier=microvm_id) @@ -566,39 +594,6 @@ def _terminate_call(self, microvm_id: str) -> None: ): raise - def _job_error(self, callback: Callback, error: Exception) -> Exception: - """Convert a failure of the callback wait into its MicroVM error.""" - name = self.name - if isinstance(error, CallbackTimeoutError): - return MicrovmTimeoutError( - message=f'MicroVM "{name}": the job timed out: {error.message}', - data=error.data, - stack_trace=error.stack_trace, - ) - if isinstance(error, CallbackError): - # The SDK raises CallbackExternalError with the message and the - # data. The error type that the MicroVM reported stays on the - # callback's checkpoint, so the message carries it from there. - reported = _reported_error_type(callback) - detail = f" ({reported})" if reported else "" - return MicrovmJobFailedError( - message=f'MicroVM "{name}": the job failed{detail}: {error.message}', - data=error.data, - stack_trace=error.stack_trace, - ) - # Any other failure, such as a result that is not valid JSON. The - # MicroVM sent the result, so the job failed. The JavaScript operation - # differs here: its SDK fails the invocation with SERDES_FAILED before - # the operation sees an error. The SDK checkpoints the class of a - # MicroVM error as its type. So the message names the original type. - detail = getattr(error, "message", None) or str(error) - return MicrovmJobFailedError( - message=( - f'MicroVM "{name}": the job result could not be read ' - f"({type(error).__name__}): {detail}" - ) - ) - def _connector_arn(self, connector: str) -> str: return ( f"arn:{self.partition}:lambda:{self.region}:aws:network-connector:" @@ -611,6 +606,183 @@ def _client(self) -> LambdaMicrovmsClient: return _DEFAULT_CLIENTS.get(self.region) +@dataclass(frozen=True, kw_only=True) +class _OperationScope(_MicrovmScope): + """The validated values of one :func:`microvm` call.""" + + heartbeat_timeout_seconds: int | None + request: MicrovmRequestConfig + retry_window_seconds: int + + @classmethod + def create( + cls, + context: DurableContext, + name: str, + job_input: Any, + config: MicrovmConfig, + ) -> _OperationScope: + """Validate the call, before any durable operation starts. + + Raises: + ValidationError: When ``name`` or ``config`` is not valid, or the + Region cannot be determined. + TypeError: When ``job_input`` is not JSON-serializable. + """ + timeout_seconds = _validate_base(name, config, config.timeout) + heartbeat_timeout_seconds = _validate_heartbeat(name, config.heartbeat_timeout) + request = config.request or MicrovmRequestConfig() + retry_window_seconds = _validate_request(name, request) + _validate_input(name, job_input) + partition, region = _validate_region(context, name) + return cls( + name=name, + config=config, + partition=partition, + region=region, + timeout_seconds=timeout_seconds, + heartbeat_timeout_seconds=heartbeat_timeout_seconds, + request=request, + retry_window_seconds=retry_window_seconds, + retry_strategy=config.retry_strategy or default_microvm_retry_strategy, + ) + + def run_job(self, child: DurableContext, job_input: Any) -> Any: + """Run the operation's durable operations in its child context.""" + name = self.name + callback = create_job_callback( + child, name, self.timeout_seconds, self.heartbeat_timeout_seconds + ) + # The operation picks the delivery: + # + # 1. The run hook is preferred. It opens no ingress, and it needs no + # extra durable step. + # 2. A config with a request path goes over HTTP. The worker runs a + # run hook job with its handler, not with a route. + # 3. A payload larger than the run hook limit goes over HTTP. + # RunMicrovm rejects it with ValidationException, which the + # default retry strategy does not retry. + # + # The payload contains the callback ID, so the size is known only + # after the callback exists. The ID comes from the checkpoint on + # replay, and the input must not change. So every replay picks the + # same delivery as the first attempt. + with_job = RunHookPayload( + region=self.region, + callback_id=callback.callback_id, + job_input=job_input, + heartbeat_timeout_seconds=self.heartbeat_timeout_seconds, + ).to_json() + over_http = self.request.path is not None or not fits_run_hook(with_job) + delivery = HTTP_DELIVERY if over_http else RUN_HOOK_DELIVERY + payload = run_hook_payload_without_job(self.region) if over_http else with_job + + launched = self.launch( + child, token_source=callback.callback_id, payload=payload, delivery=delivery + ) + microvm_id = launched.microvm_id + # Replay computes the delivery again from the current input and + # config. A deploy between the launch and a replay can change both. + # The launch result comes from the checkpoint on replay, and it + # records the delivery that the launch prepared for. So the operation + # follows the launch's delivery, not the new computation. + if launched.delivery not in {RUN_HOOK_DELIVERY, HTTP_DELIVERY}: + # The launch checkpoint comes from a release with a delivery that + # this release does not know, for example after a rollback. This + # release cannot deliver the job. So it terminates the MicroVM and + # fails, the same way on every replay. + self.terminate(child, microvm_id) + raise unknown_delivery_error(name, launched) + + if launched.delivery == HTTP_DELIVERY: + try: + self.deliver( + child, + job_name=name, + launched=launched, + callback_id=callback.callback_id, + job_input=job_input, + heartbeat_timeout_seconds=self.heartbeat_timeout_seconds, + request=self.request, + retry_window_seconds=self.retry_window_seconds, + ) + except MicrovmDeliveryError: + self.terminate(child, microvm_id) + raise + + try: + result = callback.result() + except Exception as error: + # SuspendExecution is a BaseException, so a wait for the callback + # passes through here without terminating the MicroVM. + self.terminate(child, microvm_id) + raise job_error(name, callback, error) from error + self.terminate(child, microvm_id) + return result + + +def create_job_callback( + child: DurableContext, + job_name: str, + timeout_seconds: int, + heartbeat_timeout_seconds: int | None, +) -> Callback: + """Create the callback ``.callback`` that the MicroVM completes.""" + return child.create_callback( + name=f"{job_name}.callback", + config=CallbackConfig( + timeout=Duration.from_seconds(timeout_seconds), + heartbeat_timeout=Duration.from_seconds(heartbeat_timeout_seconds or 0), + # The worker sends the result as a JSON string. The default + # callback serdes would return that string unparsed. + serdes=JsonSerDes(), + ), + ) + + +def unknown_delivery_error(name: str, launched: LaunchResult) -> MicrovmDeliveryError: + """The error for a launch checkpoint whose delivery this release does not + know.""" + msg = ( + f'MicroVM "{name}": the launch prepared the delivery ' + f'"{launched.delivery}", which this release does not support' + ) + return MicrovmDeliveryError(msg) + + +def job_error(name: str, callback: Callback, error: Exception) -> Exception: + """Convert a failure of the callback wait into its MicroVM error.""" + if isinstance(error, CallbackTimeoutError): + return MicrovmTimeoutError( + message=f'MicroVM "{name}": the job timed out: {error.message}', + data=error.data, + stack_trace=error.stack_trace, + ) + if isinstance(error, CallbackError): + # The SDK raises CallbackExternalError with the message and the + # data. The error type that the MicroVM reported stays on the + # callback's checkpoint, so the message carries it from there. + reported = _reported_error_type(callback) + detail = f" ({reported})" if reported else "" + return MicrovmJobFailedError( + message=f'MicroVM "{name}": the job failed{detail}: {error.message}', + data=error.data, + stack_trace=error.stack_trace, + ) + # Any other failure, such as a result that is not valid JSON. The + # MicroVM sent the result, so the job failed. The JavaScript operation + # differs here: its SDK fails the invocation with SERDES_FAILED before + # the operation sees an error. The SDK checkpoints the class of a + # MicroVM error as its type. So the message names the original type. + detail = getattr(error, "message", None) or str(error) + return MicrovmJobFailedError( + message=( + f'MicroVM "{name}": the job result could not be read ' + f"({type(error).__name__}): {detail}" + ) + ) + + class _DefaultClients: """One ``lambda-microvms`` client per Region, created on first use. 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 c4813f228..ac9676b4b 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 @@ -213,6 +213,49 @@ class SendJobOptions: opener: urllib.request.OpenerDirector | None = None """The urllib opener for the POST. ``None`` uses :func:`default_opener`.""" sleep: Callable[[float], None] = time.sleep + window_started_at: float | None = None + """When the retry window started, as a ``time.monotonic()`` time. + + ``None`` means the moment :func:`send_job` is called. A session job sets + it to the start of the state check that runs before delivery. So the + wait for a resume and the request's own retries share one window. + """ + recheck: Callable[[float], str | None] | None = None + """Make sure the MicroVM runs, and return its current endpoint. + + A session that can suspend its MicroVM sets it. The suspend can land + between the state check before delivery and the request. The endpoint + then answers 502, or the worker answers 503 while it suspends. So after + a 502, 503, or 504, or a request without an answer, the next attempt + calls it first. It receives what remains of the retry window, in + seconds, so the wait for a resume never extends the window. ``None`` + from it keeps the current endpoint. + """ + + +def require_time_for_call(remaining_seconds: Callable[[], float | None]) -> None: + """Fail at once when the invocation has too little time for one call. + + A step that starts too close to the end of the invocation cannot make + one call before the reserve. Without this check, the first tier would + report that the MicroVM did not accept the job, or did not reach + ``RUNNING``. The cause is the function timeout. The error stays + retryable: a step retry runs early in a new invocation, with more time + left. A function timeout under the needed time fails every attempt, and + the final error names the fix. + + Raises: + MicrovmEndpointUnavailableError: When less than 11 seconds remain. + """ + remaining = remaining_seconds() + needed = INVOCATION_RESERVE_SECONDS + MIN_CALL_SECONDS + if remaining is not None and remaining < needed: + msg = ( + f"The request step had {remaining:.1f} seconds left in the Lambda " + f"invocation, and it needs at least {needed:.0f} seconds to send " + "the job. Raise the function timeout." + ) + raise MicrovmEndpointUnavailableError(msg) def send_job(options: SendJobOptions) -> int: @@ -281,22 +324,12 @@ def send_job(options: SendJobOptions) -> int: MicrovmEndpointUnavailableError: When the tier ends. Exception: When ``CreateMicrovmAuthToken`` fails. """ - # A step that starts too close to the end of the invocation cannot make - # one call before the reserve. Say so, instead of reporting that the - # MicroVM did not accept the job. The error stays retryable: a step - # retry runs early in a new invocation, with more time left. A function - # timeout under the needed time fails every attempt, and the final error - # names the fix. - remaining = options.remaining_seconds() - needed = INVOCATION_RESERVE_SECONDS + MIN_CALL_SECONDS - if remaining is not None and remaining < needed: - msg = ( - f"The request step had {remaining:.1f} seconds left in the Lambda " - f"invocation, and it needs at least {needed:.0f} seconds to send " - "the job. Raise the function timeout." - ) - raise MicrovmEndpointUnavailableError(msg) - started = time.monotonic() + require_time_for_call(options.remaining_seconds) + started = ( + time.monotonic() + if options.window_started_at is None + else options.window_started_at + ) window_ends = started + options.retry_window_seconds attempt = 0 last_failure = "no attempt" @@ -319,7 +352,8 @@ def new_token() -> Mapping[str, str]: unavailable, ) - url = _endpoint_url(options.endpoint, options.path) + endpoint = options.endpoint + url = _endpoint_url(endpoint, options.path) token = new_token() refreshed_token = False delay = _INITIAL_DELAY_SECONDS @@ -398,6 +432,17 @@ def new_token() -> Mapping[str, str]: ) options.sleep(delay) delay = min(delay * 2, _MAX_DELAY_SECONDS) + # A request without an answer, or a gateway error, can mean that the + # MicroVM suspended itself. So a session checks the state, and + # resumes the MicroVM, before it sends the job again. + if options.recheck is not None and ( + status is None or 502 <= status <= 504 # noqa: PLR2004 + ): + left = deadline() - time.monotonic() + if left <= 0: + raise unavailable() + endpoint = options.recheck(left) or endpoint + url = _endpoint_url(endpoint, options.path) def _create_auth_token(options: SendJobOptions) -> Mapping[str, str]: 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 new file mode 100644 index 000000000..4536fc9a4 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py @@ -0,0 +1,417 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""The ``microvm_session`` operation: run several jobs in one AWS Lambda +MicroVM.""" + +from __future__ import annotations + +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, TypeVar + +from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + ValidationError, +) + +from aws_durable_execution_sdk_python_extras.microvm.config import ( + MicrovmInvokeOptions, + MicrovmSessionConfig, +) +from aws_durable_execution_sdk_python_extras.microvm.errors import MicrovmError +from aws_durable_execution_sdk_python_extras.microvm.job import ( + run_hook_payload_without_job, +) +from aws_durable_execution_sdk_python_extras.microvm.operation import ( + HTTP_DELIVERY, + LaunchResult, + _MicrovmScope, + _validate_base, + _validate_heartbeat, + _validate_input, + _validate_region, + _validate_request, + _validate_timeout, + create_job_callback, + job_error, + unknown_delivery_error, +) +from aws_durable_execution_sdk_python_extras.microvm.retry import ( + default_microvm_retry_strategy, +) + + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python.context import DurableContext + + +_T = TypeVar("_T") + +DEFAULT_AUTO_SUSPEND_IDLE_SECONDS = 60 +"""The default :attr:`MicrovmSessionConfig.auto_suspend_idle_time`, in seconds. + +A resume took about 1 second in us-east-1. So 60 seconds keeps back-to-back +jobs from paying for a suspend and a resume that they do not need. A long +wait costs at most 60 seconds of compute before the MicroVM suspends. A +session whose ``timeout`` is shorter uses its ``timeout``. A session whose +``timeout`` is under 10 seconds does not suspend by default. +""" + +_MIN_AUTO_SUSPEND_IDLE_SECONDS = 10 +"""The shortest idle time a session accepts. + +A resumed MicroVM runs its ``resume`` hook, reports ``RUNNING`` at the next +state poll, and then receives the job after a new auth token. That path can +take a few seconds. A shorter idle time could suspend the MicroVM again +before its job arrives. +""" + + +def microvm_session( + context: DurableContext, + name: str, + *, + config: MicrovmSessionConfig, + handler: Callable[[MicrovmSession, DurableContext], _T], +) -> _T: + """Run a handler that sends several jobs to one AWS Lambda MicroVM, and + return the handler's value. + + The jobs share the MicroVM's files and processes. So a later job can use + what an earlier job wrote. The operation runs its durable operations in + one child context named ``name``: + + 1. ``.session`` generates a session ID and checkpoints it. + 2. ``.launch`` calls ``RunMicrovm``. The client token is the + SHA-256 hex digest of the session ID. The session ID is random, so the + token is unique per session. It is checkpointed, so every retry and + replay of the launch returns the same MicroVM. + 3. The handler runs with a :class:`MicrovmSession` and the child context. + Each :meth:`MicrovmSession.invoke` delivers one job over HTTP and + waits for its result. The handler can use the child context for + durable operations between jobs, such as waits and callbacks. + 4. ``.terminate`` calls ``TerminateMicrovm`` after the handler + returns or raises. + + The session sends its first job right after the launch. The MicroVM + endpoint holds or refuses requests until the ``run`` hook returns, and + the request step retries a refused request inside its own attempt. So the + session needs no separate readiness wait. + + With ``config.auto_suspend_on_idle``, the worker in the MicroVM suspends + its own MicroVM when no job has run for ``auto_suspend_idle_time``. With + that setting or an ``idle_policy``, the request step of each job calls + ``GetMicrovm`` first. It resumes a suspended MicroVM, and it returns at + once for a new MicroVM that is still booting. + + 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 + terminated: + + 1. A MicroVM error, such as :class:`MicrovmJobFailedError`, reaches the + caller with its own class. + 2. Any other error reaches the caller as ``ChildContextError``, as from + ``run_in_child_context``. + + The session's value is the handler's return value, checkpointed as the + child context result. So it must be serializable by the SDK's default + serdes, and it counts toward the checkpoint size limit. + + Pass the context that the call runs in. ``invoke`` takes its context the + same way: + + .. code-block:: python + + def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: + build = vm.invoke( + ctx, + "build", + job_input=event, + options=MicrovmInvokeOptions( + path="/build", timeout=Duration.from_minutes(15) + ), + ) + return vm.invoke( + ctx, + "test", + job_input=build, + options=MicrovmInvokeOptions( + path="/test", timeout=Duration.from_minutes(10) + ), + ) + + + result = microvm_session(context, "pipeline", config=CONFIG, handler=pipeline) + + Args: + context: The context to create the durable operations in. + name: The child context name. It also prefixes the session's own + operation names. + config: The MicroVM. ``timeout`` bounds the MicroVM's lifetime. + handler: Sends the jobs. It receives the session handle and the + session's child context. + + Returns: + The handler's return value. + + Raises: + ValidationError: When ``name`` or ``config`` is not valid, or the + Region cannot be determined. No durable operation starts. + TypeError: When ``handler`` is not callable. No durable operation + starts. + MicrovmLaunchError: When ``RunMicrovm`` fails after all retries. + MicrovmError: A MicroVM error that the handler did not catch. + ChildContextError: Any other error that the handler raised. + """ + scope = _SessionScope.create(context, name, config) + if not callable(handler): + msg = f'MicroVM session "{name}": handler must be callable' + raise TypeError(msg) + + def run(child: DurableContext) -> _T: + return scope.run_session(child, handler) + + try: + return context.run_in_child_context(run, name=name) + except ChildContextError as error: + # The handler is caller code. So only a MicroVM error is restored to + # its class. Every other error keeps the SDK's ChildContextError. + if isinstance(error.__cause__, MicrovmError): + raise MicrovmError.from_child_context_error(name, error) from error + raise + + +class MicrovmSession: + """The handle that a :func:`microvm_session` handler receives. + + The handle holds no context of its own. Each :meth:`invoke` creates its + durable operations in the context that the caller passes. So one handle + works in the session's child context, in a nested child context, and in + ``map`` and ``parallel`` items. + """ + + def __init__(self, scope: _SessionScope, launched: LaunchResult) -> None: + self._scope = scope + self._launched = launched + + @property + def microvm_id(self) -> str: + """The ID of the session's MicroVM.""" + return self._launched.microvm_id + + def invoke( + self, + context: DurableContext, + name: str, + *, + job_input: Any, + options: MicrovmInvokeOptions, + ) -> Any: + """Send one job to the session's MicroVM, and return the job's result. + + The job goes over HTTP to ``options.path``. Without a path, it goes to + :data:`DEFAULT_MICROVM_JOB_PATH`, which the worker serves with its + ``handler``. + + Each call creates a child context named ``name`` in ``context``, with + a callback ``.callback`` and a request step ``.request``. + Pass the context that the call runs in: the session's child context, + or a nested one, such as a ``map`` item's context. + + Args: + context: The context to create the job's durable operations in. + 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. + + Returns: + The value that the MicroVM passed to + ``SendDurableExecutionCallbackSuccess``, parsed as JSON. + + Raises: + TypeError: When ``context`` is not a durable context, or + ``job_input`` is not JSON-serializable. + ValidationError: When ``name`` or ``options`` is not valid. + MicrovmNotRunningError: When the session's MicroVM has ended, and + the session checks the MicroVM state before each job. + MicrovmDeliveryError: When the delivery fails after all retries, + or the route rejects the job. + 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. + """ + scope = self._scope + # A call in the shape vm.invoke("build", ...) passes a string here. + # This check names the mistake, instead of an AttributeError from + # inside. + if not callable(getattr(context, "run_in_child_context", None)): + msg = ( + f'MicroVM session "{scope.name}": invoke requires the durable ' + "context to create the job in, as its first argument" + ) + raise TypeError(msg) + if not isinstance(name, str) or not name: + msg = ( + f'MicroVM session "{scope.name}": invoke requires a non-empty job name' + ) + raise ValidationError(msg) + retry_window_seconds = _validate_request(name, options) + timeout_seconds = _validate_timeout(name, "timeout", options.timeout) + heartbeat_timeout_seconds = _validate_heartbeat(name, options.heartbeat_timeout) + _validate_input(name, job_input) + launched = self._launched + + def run(job_context: DurableContext) -> Any: + callback = create_job_callback( + job_context, name, timeout_seconds, heartbeat_timeout_seconds + ) + scope.deliver( + job_context, + job_name=name, + launched=launched, + callback_id=callback.callback_id, + job_input=job_input, + heartbeat_timeout_seconds=heartbeat_timeout_seconds, + request=options, + retry_window_seconds=retry_window_seconds, + resume=scope.resume_before_delivery, + ) + try: + return callback.result() + except Exception as error: + # SuspendExecution is a BaseException, so a wait for the + # callback passes through here. + raise job_error(name, callback, error) from error + + try: + return context.run_in_child_context(run, name=name) + except ChildContextError as error: + raise MicrovmError.from_child_context_error(name, error) from error + + +@dataclass(frozen=True, kw_only=True) +class _SessionScope(_MicrovmScope): + """The validated values of one :func:`microvm_session` call.""" + + auto_suspend_idle_seconds: int | None + idle_policy: Mapping[str, Any] | None + + @property + def resume_before_delivery(self) -> bool: + """Whether each job must check the MicroVM state, and resume it. + + The worker suspends an idle MicroVM, and an idle policy can suspend + it too. Either way, each job must resume it first. + """ + return ( + self.auto_suspend_idle_seconds is not None or self.idle_policy is not None + ) + + @classmethod + def create( + cls, context: DurableContext, name: str, config: MicrovmSessionConfig + ) -> _SessionScope: + """Validate the call, before any durable operation starts. + + Raises: + ValidationError: When ``name`` or ``config`` is not valid, or the + Region cannot be determined. + """ + timeout_seconds = _validate_base(name, config, config.timeout) + auto_suspend_idle_seconds = _validate_auto_suspend( + name, config, timeout_seconds + ) + partition, region = _validate_region(context, name) + return cls( + name=name, + config=config, + partition=partition, + region=region, + timeout_seconds=timeout_seconds, + retry_strategy=config.retry_strategy or default_microvm_retry_strategy, + auto_suspend_idle_seconds=auto_suspend_idle_seconds, + idle_policy=config.idle_policy, + ) + + def run_session( + self, + child: DurableContext, + handler: Callable[[MicrovmSession, DurableContext], _T], + ) -> _T: + """Run the session's durable operations in its child context.""" + session_id: str = child.step( + lambda _: str(uuid.uuid4()), name=f"{self.name}.session" + ) + launched = self.launch( + child, + token_source=session_id, + payload=run_hook_payload_without_job( + self.region, self.auto_suspend_idle_seconds + ), + delivery=HTTP_DELIVERY, + idle_policy=self.idle_policy, + ) + if launched.delivery != HTTP_DELIVERY: + # A launch checkpoint from a release with another delivery. Every + # session job goes over HTTP, so this release cannot run one. + self.terminate(child, launched.microvm_id) + raise unknown_delivery_error(self.name, launched) + try: + result = handler(MicrovmSession(self, launched), child) + except Exception: + # SuspendExecution is a BaseException. A handler that waits, for + # a job or a durable wait, passes through here without + # terminating the MicroVM. + self.terminate(child, launched.microvm_id) + raise + self.terminate(child, launched.microvm_id) + return result + + +def _validate_auto_suspend( + name: str, config: MicrovmSessionConfig, timeout_seconds: int +) -> int | None: + """Validate the auto-suspend settings. + + Returns: + The idle time in seconds, or ``None`` when the MicroVM must not + suspend itself. + + Raises: + ValidationError: When a setting is not valid. + """ + # Suspending is opt-in. It needs lambda:GetMicrovm and + # lambda:ResumeMicrovm in the function's role, and lambda:SuspendMicrovm + # in the MicroVM's role. It also stops background processes between jobs. + if config.auto_suspend_on_idle is not True: + if config.auto_suspend_idle_time is not None: + msg = ( + f'MicroVM session "{name}": auto_suspend_idle_time is set, but ' + "auto_suspend_on_idle is not True" + ) + raise ValidationError(msg) + return None + if config.auto_suspend_idle_time is None: + # A session shorter than the minimum idle time does not suspend. Its + # MicroVM would suspend again before a resumed job arrives. + if timeout_seconds < _MIN_AUTO_SUSPEND_IDLE_SECONDS: + return None + return min(DEFAULT_AUTO_SUSPEND_IDLE_SECONDS, timeout_seconds) + idle_seconds = config.auto_suspend_idle_time.to_seconds() + if idle_seconds < _MIN_AUTO_SUSPEND_IDLE_SECONDS: + msg = ( + f'MicroVM session "{name}": auto_suspend_idle_time must be at least ' + f"{_MIN_AUTO_SUSPEND_IDLE_SECONDS} seconds" + ) + raise ValidationError(msg) + if idle_seconds > timeout_seconds: + msg = ( + f'MicroVM session "{name}": auto_suspend_idle_time must not be longer ' + "than timeout" + ) + raise ValidationError(msg) + return idle_seconds 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 253f2043a..2bcd000b0 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py @@ -31,14 +31,18 @@ def client_error(code: str, status: int = 400) -> ClientError: class FakeMicrovmsClient: """Records Lambda MicroVMs API calls, and answers from scripted lists. - ``launches``, ``tokens``, and ``terminates`` hold the answers of - ``run_microvm``, ``create_microvm_auth_token``, and ``terminate_microvm``, - in order. Each answer is an exception to raise, a function that receives - the request, or ``None`` for success. After a list ends, every call - succeeds. - - ``run_microvm`` returns ``endpoint``. The Nth auth token is - ``{"x-test-token": "token-N"}``. + ``launches``, ``tokens``, ``terminates``, and ``resumes`` hold the + answers of ``run_microvm``, ``create_microvm_auth_token``, + ``terminate_microvm``, and ``resume_microvm``, in order. Each answer is an + exception to raise, a function that receives the request, or ``None`` for + success. After a list ends, every call succeeds. + + ``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``. + + ``run_microvm`` and ``get_microvm`` return ``endpoint``. The Nth auth + token is ``{"x-test-token": "token-N"}``. """ def __init__( @@ -47,6 +51,8 @@ def __init__( launches: list[Answer] | None = None, tokens: list[Answer] | None = None, terminates: list[Answer] | None = None, + resumes: list[Answer] | None = None, + states: list[str | BaseException] | None = None, endpoint: str | None = "mvm-1.microvm.example", ) -> None: self.calls: list[tuple[str, dict[str, Any]]] = [] @@ -55,7 +61,9 @@ def __init__( "run_microvm": list(launches or []), "create_microvm_auth_token": list(tokens or []), "terminate_microvm": list(terminates or []), + "resume_microvm": list(resumes or []), } + self._states = list(states or []) self._lock = threading.Lock() def requests(self, name: str) -> list[dict[str, Any]]: @@ -89,6 +97,21 @@ def terminate_microvm(self, **request: Any) -> dict[str, Any]: self._answer("terminate_microvm", request) return {} + def resume_microvm(self, **request: Any) -> dict[str, Any]: + self._answer("resume_microvm", request) + 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 isinstance(state, BaseException): + raise state + response: dict[str, Any] = {"microvmId": "mvm-1", "state": state} + if self.endpoint is not None: + response["endpoint"] = self.endpoint + return response + # The fakes are fixtures, not imports. Several packages in this repository # have a top-level ``tests`` package, so ``from tests.conftest import ...`` diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/errors_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/errors_test.py index 9324e5bcf..1ad3f5750 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/errors_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/errors_test.py @@ -17,6 +17,7 @@ MicrovmError, MicrovmJobFailedError, MicrovmLaunchError, + MicrovmNotRunningError, MicrovmTimeoutError, ) @@ -39,6 +40,7 @@ def child_context_error(raised: Exception) -> ChildContextError: MicrovmLaunchError, MicrovmDeliveryError, MicrovmJobFailedError, + MicrovmNotRunningError, MicrovmTimeoutError, ], ) @@ -74,6 +76,7 @@ def test_every_microvm_error_is_a_durable_operation_error() -> None: MicrovmLaunchError, MicrovmDeliveryError, MicrovmJobFailedError, + MicrovmNotRunningError, MicrovmTimeoutError, ): assert issubclass(error_class, MicrovmError) @@ -86,6 +89,7 @@ def test_every_microvm_error_is_a_durable_operation_error() -> None: MicrovmLaunchError, MicrovmDeliveryError, MicrovmJobFailedError, + MicrovmNotRunningError, MicrovmTimeoutError, ], ) 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 1bc2df308..54853af27 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 @@ -138,3 +138,15 @@ def test_the_worker_reads_a_run_hook_payload_without_a_job() -> None: assert payload == '{"version":1,"region":"us-east-1"}' assert request.payload is not None assert request.payload.job is None + + +def test_the_worker_reads_the_idle_time_of_a_session() -> None: + payload = run_hook_payload_without_job("us-east-1", 60) + + request = RunHookRequest.from_dict( + {"microvmId": "mvm-1", "runHookPayload": payload} + ) + + assert payload == '{"version":1,"region":"us-east-1","autoSuspendIdleSeconds":60}' + assert request.payload is not None + assert request.payload.auto_suspend_idle_seconds == 60 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 new file mode 100644 index 000000000..be0f99c4b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py @@ -0,0 +1,213 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the state check that resumes a suspended session MicroVM.""" + +from __future__ import annotations + +import logging +import threading +from collections.abc import Callable +from typing import Any + +import pytest + +from aws_durable_execution_sdk_python_extras.microvm import ( + MicrovmEndpointUnavailableError, +) +from aws_durable_execution_sdk_python_extras.microvm.lifecycle import ( + EnsureRunningOptions, + MicrovmStateError, + ensure_running, +) +from aws_durable_execution_sdk_python_extras.microvm.request import ( + INVOCATION_RESERVE_SECONDS, +) + + +def run_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 ensure_running( + EnsureRunningOptions( + client=client, + microvm_id="mvm-1", + max_wait_seconds=max_wait_seconds, + remaining_seconds=remaining, + logger=logging.getLogger("test"), + sleep=recorded.append, + ) + ) + + +def calls(client: Any) -> list[str]: + return [name for name, _ in client.calls] + + +def test_returns_the_current_endpoint_of_a_running_microvm(fake_client: Any) -> None: + client = fake_client(endpoint="mvm-1.new.example") + + assert run_with(client) == "mvm-1.new.example" + assert calls(client) == ["get_microvm"] + + +def test_waits_for_a_suspend_in_progress_then_resumes_once(fake_client: Any) -> None: + client = fake_client(states=["SUSPENDING", "SUSPENDED", "SUSPENDED", "RUNNING"]) + sleeps: list[float] = [] + + assert run_with(client, sleeps=sleeps) == "mvm-1.microvm.example" + + # The MicroVM stays SUSPENDED while its resume hook runs. So one + # ResumeMicrovm call is enough. + assert calls(client) == [ + "get_microvm", + "get_microvm", + "resume_microvm", + "get_microvm", + "get_microvm", + ] + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert sleeps == [0.25, 0.5, 1.0] + + +def test_keeps_polling_when_resume_reports_a_transition_in_progress( + fake_client: Any, make_error: Any +) -> None: + client = fake_client( + states=["SUSPENDED", "SUSPENDED"], + resumes=[make_error("ConflictException", 409)], + ) + + assert run_with(client) == "mvm-1.microvm.example" + assert len(client.requests("resume_microvm")) == 1 + + +def test_raises_a_resume_error_that_is_not_a_conflict( + fake_client: Any, make_error: Any +) -> None: + denied = make_error("AccessDeniedException", 403) + client = fake_client(states=["SUSPENDED"], resumes=[denied]) + + with pytest.raises(type(denied)) as raised: + run_with(client) + + assert raised.value is denied + + +def test_returns_at_once_for_a_new_microvm_that_is_still_booting( + fake_client: Any, +) -> None: + client = fake_client(states=["PENDING"]) + + assert run_with(client) == "mvm-1.microvm.example" + assert calls(client) == ["get_microvm"] + + +def test_waits_for_running_after_a_resume_that_reports_pending( + fake_client: Any, +) -> None: + client = fake_client(states=["SUSPENDED", "PENDING", "RUNNING"]) + + assert run_with(client) == "mvm-1.microvm.example" + assert calls(client).count("get_microvm") == 3 + + +@pytest.mark.parametrize("state", ["TERMINATING", "TERMINATED"]) +def test_raises_the_state_error_for_a_microvm_that_ends( + fake_client: Any, state: str +) -> None: + client = fake_client(states=[state]) + + with pytest.raises(MicrovmStateError) as raised: + run_with(client) + + assert raised.value.state == state + assert str(raised.value) == f"MicroVM mvm-1 is {state}." + + +def test_includes_the_state_reason(fake_client: Any) -> None: + client = fake_client() + client.get_microvm = lambda **_: { + "state": "TERMINATED", + "stateReason": "Maximum duration reached", + } + + with pytest.raises(MicrovmStateError, match="is TERMINATED: Maximum duration"): + run_with(client) + + +def test_raises_the_state_error_when_the_microvm_no_longer_exists( + fake_client: Any, make_error: Any +) -> None: + client = fake_client(states=[make_error("ResourceNotFoundException", 404)]) + + with pytest.raises(MicrovmStateError) as raised: + run_with(client) + + assert raised.value.state is None + assert "no longer exists" in str(raised.value) + + +def test_raises_another_get_microvm_error(fake_client: Any, make_error: Any) -> None: + throttled = make_error("ThrottlingException", 429) + client = fake_client(states=[throttled]) + + with pytest.raises(type(throttled)) as raised: + run_with(client) + + assert raised.value is throttled + + +def test_gives_up_after_the_wait_limit(fake_client: Any) -> None: + client = fake_client(states=["SUSPENDING"] * 10) + + with pytest.raises(MicrovmEndpointUnavailableError, match="Last state: SUSPENDING"): + run_with(client, max_wait_seconds=0.5) + + # 0.25 fits in the limit, and the next 0.5 does not. + assert calls(client) == ["get_microvm", "get_microvm"] + + +def test_gives_up_before_the_invocation_reserve(fake_client: Any) -> None: + client = fake_client(states=["SUSPENDING"] * 10) + + with pytest.raises(MicrovmEndpointUnavailableError): + run_with( + client, + # The reserve starts in 1.3 seconds. The fake sleep takes no time, + # so the waits of 0.25, 0.5, and 1 second fit, and 2 seconds + # does not. + remaining=lambda: INVOCATION_RESERVE_SECONDS + 1.3, + ) + + assert calls(client) == ["get_microvm"] * 4 + + +def test_starts_no_call_when_less_than_1_second_remains_before_the_reserve( + fake_client: Any, +) -> None: + client = fake_client() + + with pytest.raises(MicrovmEndpointUnavailableError): + run_with(client, remaining=lambda: INVOCATION_RESERVE_SECONDS + 0.5) + + assert calls(client) == [] + + +def test_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" + ): + run_with(client, remaining=lambda: INVOCATION_RESERVE_SECONDS + 1.2) + finally: + release.set() diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py index 61e353e80..c03677e9d 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py @@ -545,7 +545,7 @@ def test_fails_a_launch_without_an_endpoint_for_http( def test_fails_the_request_for_a_launch_checkpoint_without_an_endpoint( run: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch ) -> None: - original = operation._OperationScope._launch # noqa: SLF001 + original = operation._MicrovmScope._launch_call # noqa: SLF001 def without_endpoint(self: Any, *args: Any) -> Any: launched = original(self, *args) @@ -553,7 +553,7 @@ def without_endpoint(self: Any, *args: Any) -> Any: microvm_id=launched.microvm_id, delivery=launched.delivery ) - monkeypatch.setattr(operation._OperationScope, "_launch", without_endpoint) # noqa: SLF001 + monkeypatch.setattr(operation._MicrovmScope, "_launch_call", without_endpoint) # noqa: SLF001 client = fake_client() config = config_with(client, request=MicrovmRequestConfig(path="/job")) @@ -789,7 +789,7 @@ def test_terminates_and_fails_for_a_delivery_that_it_does_not_know( ) -> None: # A launch checkpoint from another release, for example after a rollback. - original = operation._OperationScope._launch # noqa: SLF001 + original = operation._MicrovmScope._launch_call # noqa: SLF001 def launch_for_another(self: Any, *args: Any) -> Any: launched = original(self, *args) @@ -797,7 +797,7 @@ def launch_for_another(self: Any, *args: Any) -> Any: microvm_id=launched.microvm_id, delivery="smoke-signal" ) - monkeypatch.setattr(operation._OperationScope, "_launch", launch_for_another) # noqa: SLF001 + monkeypatch.setattr(operation._MicrovmScope, "_launch_call", launch_for_another) # noqa: SLF001 client = fake_client() result = run(catching(config_with(client)), {}) 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 b40104e6a..0a42c678a 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 @@ -471,3 +471,81 @@ def interrupted() -> None: with pytest.raises(KeyboardInterrupt): call_before_deadline(interrupted, time.monotonic() + 30, RuntimeError) + + +def test_checks_the_state_again_after_a_gateway_error_or_no_answer( + fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([502, None, 503, 504, 429, 500]) + sleeps: list[float] = [] + windows: list[float] = [] + + def recheck(left: float) -> str | None: + windows.append(left) + return None + + status = send_job(options(fake_client(), endpoint.url, sleeps, recheck=recheck)) + + assert status == 202 + # 502, no answer, 503, and 504 can mean a suspend. 429 and 500 cannot. + assert len(windows) == 4 + assert all(0 < left <= 60 for left in windows) + + +def test_sends_the_next_request_to_the_endpoint_that_the_recheck_returns( + fake_client: Any, fake_endpoint: Any +) -> None: + before = fake_endpoint([502]) + after = fake_endpoint() + + status = send_job( + options(fake_client(), before.url, [], recheck=lambda left: after.url) + ) + + assert status == 202 + assert len(before.requests) == 1 + assert len(after.requests) == 1 + + +def test_ends_the_tier_when_the_recheck_has_no_time_left( + fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([502]) + rechecks: list[float] = [] + + def sleep(seconds: float) -> None: + # The delay uses the rest of the window. + time.sleep(0.9) + + 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=sleep, + recheck=rechecks.append, + ) + ) + + assert rechecks == [] + + +def test_counts_the_window_from_the_given_start( + fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([503]) + + # The state check before delivery used 59.9 of the 60 seconds. So the + # first delay of 0.25 seconds does not fit. + with pytest.raises(MicrovmEndpointUnavailableError, match="after 1 attempts"): + send_job( + options( + fake_client(), + endpoint.url, + [], + window_started_at=time.monotonic() - 59.9, + ) + ) 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 new file mode 100644 index 000000000..20c53faf5 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py @@ -0,0 +1,805 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the microvm_session operation, run on the local test runner.""" + +from __future__ import annotations + +import datetime +import hashlib +import json +from collections.abc import Callable, Iterator +from typing import Any + +import boto3 +import pytest +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + ValidationError, +) +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.retries import RetryDecision +from aws_durable_execution_sdk_python_testing.runner import ( + DurableFunctionTestResult, + DurableFunctionTestRunner, +) +from botocore.stub import ANY, Stubber + +from aws_durable_execution_sdk_python_extras.microvm import ( + DEFAULT_MICROVM_JOB_PATH, + MicrovmError, + MicrovmInvokeOptions, + MicrovmJobFailedError, + MicrovmNotRunningError, + MicrovmSession, + MicrovmSessionConfig, + microvm_session, + operation, +) + + +WAIT = 30 +INGRESS = ( + "arn:aws:lambda:us-west-2:aws:network-connector:aws-network-connector:ALL_INGRESS" +) + +Complete = Callable[[DurableFunctionTestRunner, str], None] + + +def session_config(client: Any, **fields: Any) -> MicrovmSessionConfig: + values: dict[str, Any] = { + "image_identifier": "arn:aws:lambda:us-west-2:123456789012:microvm-image:img", + "execution_role_arn": "arn:aws:iam::123456789012:role/microvm", + "timeout": Duration.from_hours(1), + "client": client, + # One quick retry keeps the retry tests fast. + "retry_strategy": lambda error, attempt: ( + RetryDecision.retry(Duration.from_seconds(1)) + if attempt < 2 + and getattr(error, "response", {}).get("Error", {}).get("Code") + == "ThrottlingException" + else RetryDecision.no_retry() + ), + **fields, + } + return MicrovmSessionConfig(**values) + + +def job(path: str | None = None, **fields: Any) -> MicrovmInvokeOptions: + return MicrovmInvokeOptions(path=path, timeout=Duration.from_minutes(5), **fields) + + +def succeed_with(value: Any) -> Complete: + def complete(runner: DurableFunctionTestRunner, callback_id: str) -> None: + runner.send_callback_success(callback_id, json.dumps(value).encode()) + + return complete + + +def fail_with(error_type: str, message: str) -> Complete: + def complete(runner: DurableFunctionTestRunner, callback_id: str) -> None: + runner.send_callback_failure( + callback_id, + ErrorObject(message=message, type=error_type, data=None, stack_trace=None), + ) + + return complete + + +def outcome(result: DurableFunctionTestResult) -> Any: + assert result.status is InvocationStatus.SUCCEEDED, result.error + return json.loads(result.result) if result.result is not None else None + + +@pytest.fixture +def run() -> Iterator[Callable[..., DurableFunctionTestResult]]: + """Run a handler on the local test runner, and complete the named callbacks + in order.""" + runners: list[DurableFunctionTestRunner] = [] + + def _run( + handler: Callable[..., Any], + event: Any = None, + completions: list[tuple[str, Complete]] | None = None, + ) -> DurableFunctionTestResult: + runner = DurableFunctionTestRunner(handler=handler, poll_interval=0.05) + runners.append(runner) + arn = runner.run_async(input=json.dumps(event), execution_timeout=WAIT) + for name, complete in completions or []: + callback_id = runner.wait_for_callback( + arn, name=f"{name}.callback", timeout=WAIT + ) + complete(runner, callback_id) + return runner.wait_for_result(arn, timeout=WAIT) + + yield _run + for runner in runners: + runner.close() + + +def session_handler( + config: MicrovmSessionConfig, + body: Callable[[MicrovmSession, DurableContext, Any], Any], +) -> Callable[..., Any]: + """A durable handler that runs ``body`` in a session named ``pipeline``, + and returns the class and message of an error that leaves it.""" + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + try: + return { + "result": microvm_session( + context, + "pipeline", + config=config, + handler=lambda vm, ctx: body(vm, ctx, event), + ) + } + except MicrovmError as error: + return {"error": type(error).__name__, "message": error.message} + except ChildContextError as error: + return {"error": "ChildContextError", "message": error.message} + + return handler + + +def two_jobs(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + build = vm.invoke(ctx, "build", job_input=event, options=job("/build")) + test = vm.invoke(ctx, "test", job_input=build, options=job()) + return {"build": build, "test": test, "microvm": vm.microvm_id} + + +def test_sends_two_jobs_to_one_microvm_and_returns_the_handlers_value( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint() + client = fake_client(endpoint=endpoint.url) + config = session_config(client) + + result = run( + session_handler(config, two_jobs), + {"repo": "org/app"}, + [("build", succeed_with({"artifact": "a.zip"})), ("test", succeed_with("ok"))], + ) + + assert outcome(result) == { + "result": {"build": {"artifact": "a.zip"}, "test": "ok", "microvm": "mvm-1"} + } + session = result.get_context("pipeline") + assert [operation.name for operation in session.child_operations] == [ + "pipeline.session", + "pipeline.launch", + "build", + "test", + "pipeline.terminate", + ] + assert [ + operation.name for operation in session.get_context("build").child_operations + ] == [ + "build.callback", + "build.request", + ] + session_id = json.loads(session.get_step("pipeline.session").result or "null") + [launch] = client.requests("run_microvm") + assert launch["clientToken"] == hashlib.sha256(session_id.encode()).hexdigest() + assert launch["ingressNetworkConnectors"] == [INGRESS] + assert launch["maximumDurationInSeconds"] == 3600 + 300 + assert json.loads(launch["runHookPayload"]) == {"version": 1, "region": "us-west-2"} + assert "idlePolicy" not in launch + # The job without a path goes to the worker's handler. + assert [request["path"] for request in endpoint.requests] == [ + "/build", + DEFAULT_MICROVM_JOB_PATH, + ] + first, second = (request["body"] for request in endpoint.requests) + assert first["microvmId"] == "mvm-1" + assert first["input"] == {"repo": "org/app"} + assert second["input"] == {"artifact": "a.zip"} + assert first["callbackId"] != second["callbackId"] + # Without suspending, the session checks no state. + assert client.requests("get_microvm") == [] + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +def test_allows_durable_waits_between_jobs( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + client = fake_client(endpoint=fake_endpoint().url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + first = vm.invoke(ctx, "build", job_input=1, options=job()) + ctx.wait(Duration.from_seconds(1), name="cool-down") + return [first, vm.invoke(ctx, "test", job_input=2, options=job())] + + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with(10)), ("test", succeed_with(20))], + ) + + assert outcome(result) == {"result": [10, 20]} + assert len(client.requests("run_microvm")) == 1 + assert len(client.requests("terminate_microvm")) == 1 + + +def test_runs_jobs_from_map_items_with_each_items_context( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint() + client = fake_client(endpoint=endpoint.url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return ctx.map( + [1, 2, 3], + lambda item_context, item, index, items: vm.invoke( + item_context, f"shard-{item}", job_input=item, options=job() + ), + name="shards", + ).get_results() + + result = run( + session_handler(session_config(client), body), + None, + [(f"shard-{item}", succeed_with(item * 10)) for item in [1, 2, 3]], + ) + + assert outcome(result) == {"result": [10, 20, 30]} + assert len(client.requests("run_microvm")) == 1 + assert sorted(request["body"]["input"] for request in endpoint.requests) == [ + 1, + 2, + 3, + ] + + +def test_lets_the_handler_catch_a_failed_job_and_continue( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + client = fake_client(endpoint=fake_endpoint().url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + try: + vm.invoke(ctx, "build", job_input=1, options=job()) + except MicrovmJobFailedError as error: + caught = error.message + return [caught, vm.invoke(ctx, "retry", job_input=2, options=job())] + + result = run( + session_handler(session_config(client), body), + None, + [ + ("build", fail_with("BuildError", "compile failed")), + ("retry", succeed_with("ok")), + ], + ) + + assert outcome(result) == { + "result": ['MicroVM "build": the job failed (BuildError): compile failed', "ok"] + } + assert len(client.requests("terminate_microvm")) == 1 + + +def test_fails_with_an_uncaught_job_error_and_still_terminates( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + client = fake_client(endpoint=fake_endpoint().url) + + result = run( + session_handler(session_config(client), two_jobs), + None, + [("build", fail_with("BuildError", "compile failed"))], + ) + + # The caller gets the MicroVM class, not ChildContextError. + assert outcome(result) == { + "error": "MicrovmJobFailedError", + "message": 'MicroVM "build": the job failed (BuildError): compile failed', + } + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +def test_fails_with_the_handlers_own_error_and_still_terminates( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + client = fake_client(endpoint=fake_endpoint().url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + msg = "no artifact" + raise LookupError(msg) + + result = run(session_handler(session_config(client), body)) + + # Caller code keeps the SDK's error. + assert outcome(result) == {"error": "ChildContextError", "message": "no artifact"} + assert len(client.requests("terminate_microvm")) == 1 + + +@pytest.mark.parametrize( + ("call", "error", "message"), + [ + ( + lambda vm, ctx: vm.invoke("build", "build", job_input=1, options=job()), # type: ignore[arg-type] + "TypeError", + "invoke requires the durable context", + ), + ( + lambda vm, ctx: vm.invoke(ctx, "", job_input=1, options=job()), + "ValidationError", + "non-empty job name", + ), + ( + lambda vm, ctx: vm.invoke(ctx, "build", job_input=1, options=job("job")), + "ValidationError", + "RFC 3986 path", + ), + ( + lambda vm, ctx: vm.invoke( + ctx, + "build", + job_input=1, + options=MicrovmInvokeOptions(timeout=Duration.from_hours(9)), + ), + "ValidationError", + "timeout must be between 1 second", + ), + ( + lambda vm, ctx: vm.invoke( + ctx, "build", job_input=1, options=job(heartbeat_timeout=Duration()) + ), + "ValidationError", + "heartbeat_timeout must be at least 1 second", + ), + ( + lambda vm, ctx: vm.invoke(ctx, "build", job_input={1, 2}, options=job()), + "TypeError", + "must be JSON-serializable", + ), + ], +) +def test_rejects_an_invalid_job_inside_the_handler_where_it_can_be_caught( + run: Any, + fake_client: Any, + fake_endpoint: Any, + call: Callable[[MicrovmSession, DurableContext], Any], + error: str, + message: str, +) -> None: + client = fake_client(endpoint=fake_endpoint().url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + try: + call(vm, ctx) + except (TypeError, ValidationError) as raised: + return [type(raised).__name__, str(raised)] + return None + + result = run(session_handler(session_config(client), body)) + + [name, text] = outcome(result)["result"] + assert name == error + assert message in text + # No job operation started. + session = result.get_context("pipeline") + assert [operation.name for operation in session.child_operations] == [ + "pipeline.session", + "pipeline.launch", + "pipeline.terminate", + ] + + +def test_retries_a_throttled_launch_with_the_same_client_token( + run: Any, fake_client: Any, fake_endpoint: Any, make_error: Any +) -> None: + client = fake_client( + launches=[make_error("ThrottlingException", 429)], + endpoint=fake_endpoint().url, + ) + + result = run( + session_handler(session_config(client), lambda vm, ctx, event: vm.microvm_id) + ) + + assert outcome(result) == {"result": "mvm-1"} + first, second = client.requests("run_microvm") + assert first["clientToken"] == second["clientToken"] + + +def test_fails_with_the_launch_error_and_skips_terminate( + run: Any, fake_client: Any, make_error: Any +) -> None: + client = fake_client(launches=[make_error("ValidationException")]) + + result = run(session_handler(session_config(client), lambda vm, ctx, event: None)) + + error = outcome(result) + assert error["error"] == "MicrovmLaunchError" + assert error["message"].startswith('MicroVM "pipeline": the launch failed:') + assert client.requests("terminate_microvm") == [] + + +@pytest.mark.parametrize( + ("fields", "message"), + [ + ({"image_identifier": ""}, "image_identifier is required"), + ({"timeout": Duration.from_hours(9)}, "timeout must be between 1 second"), + ( + {"auto_suspend_idle_time": Duration.from_seconds(30)}, + "auto_suspend_idle_time is set, but auto_suspend_on_idle is not True", + ), + ( + { + "auto_suspend_on_idle": True, + "auto_suspend_idle_time": Duration.from_seconds(9), + }, + "auto_suspend_idle_time must be at least 10 seconds", + ), + ( + { + "auto_suspend_on_idle": True, + "timeout": Duration.from_minutes(1), + "auto_suspend_idle_time": Duration.from_minutes(2), + }, + "auto_suspend_idle_time must not be longer than timeout", + ), + ], +) +def test_rejects_an_invalid_config_before_any_durable_operation( + run: Any, fake_client: Any, fields: dict[str, Any], message: str +) -> None: + client = fake_client() + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + try: + microvm_session( + context, + "pipeline", + config=session_config(client, **fields), + handler=lambda vm, ctx: None, + ) + except ValidationError as error: + return str(error) + return None + + result = run(handler) + + assert message in outcome(result) + assert result.get_all_operations() == [] + assert client.calls == [] + + +def test_rejects_a_handler_that_is_not_callable(run: Any, fake_client: Any) -> None: + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + try: + microvm_session( + context, + "pipeline", + config=session_config(fake_client()), + handler="build", # type: ignore[arg-type] + ) + except TypeError as error: + return str(error) + return None + + assert ( + outcome(run(handler)) == 'MicroVM session "pipeline": handler must be callable' + ) + + +def test_passes_an_idle_policy_through_and_resumes_before_each_job( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + client = fake_client( + endpoint=fake_endpoint().url, states=["RUNNING", "SUSPENDED", "RUNNING"] + ) + policy = { + "maxIdleDurationSeconds": 300, + "suspendedDurationSeconds": 3600, + "autoResumeEnabled": True, + } + + result = run( + session_handler(session_config(client, idle_policy=policy), two_jobs), + None, + [("build", succeed_with(1)), ("test", succeed_with(2))], + ) + + assert outcome(result)["result"]["test"] == 2 + [launch] = client.requests("run_microvm") + assert launch["idlePolicy"] == policy + # An idle policy does not ask the worker to suspend. + assert "autoSuspendIdleSeconds" not in json.loads(launch["runHookPayload"]) + assert len(client.requests("get_microvm")) == 3 + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +def test_asks_the_worker_to_suspend_after_60_seconds_by_default( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + client = fake_client(endpoint=fake_endpoint().url) + + result = run( + session_handler(session_config(client, auto_suspend_on_idle=True), two_jobs), + None, + [("build", succeed_with(1)), ("test", succeed_with(2))], + ) + + assert outcome(result)["result"]["test"] == 2 + [launch] = client.requests("run_microvm") + assert json.loads(launch["runHookPayload"]) == { + "version": 1, + "region": "us-west-2", + "autoSuspendIdleSeconds": 60, + } + assert len(client.requests("get_microvm")) == 2 + + +@pytest.mark.parametrize( + ("fields", "idle_seconds"), + [ + ({"auto_suspend_idle_time": Duration.from_seconds(90)}, 90), + ({"timeout": Duration.from_seconds(30)}, 30), + ({"timeout": Duration.from_seconds(9)}, None), + ], +) +def test_sends_the_idle_time_to_the_worker( + run: Any, + fake_client: Any, + fake_endpoint: Any, + fields: dict[str, Any], + idle_seconds: int | None, +) -> None: + client = fake_client(endpoint=fake_endpoint().url) + config = session_config(client, auto_suspend_on_idle=True, **fields) + + result = run(session_handler(config, lambda vm, ctx, event: None)) + + assert outcome(result) == {"result": None} + payload = json.loads(client.requests("run_microvm")[0]["runHookPayload"]) + assert payload.get("autoSuspendIdleSeconds") == idle_seconds + + +def test_resumes_a_microvm_that_the_worker_suspended_and_uses_its_endpoint( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + launched_at = fake_endpoint() + resumed_at = fake_endpoint() + client = fake_client( + endpoint=launched_at.url, states=["SUSPENDED", "SUSPENDED", "RUNNING"] + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + # The resumed MicroVM reports a new endpoint. + client.endpoint = resumed_at.url + return vm.invoke(ctx, "build", job_input=1, options=job()) + + result = run( + session_handler(session_config(client, auto_suspend_on_idle=True), body), + None, + [("build", succeed_with("done"))], + ) + + assert outcome(result) == {"result": "done"} + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert launched_at.requests == [] + assert len(resumed_at.requests) == 1 + + +def test_checks_the_state_again_after_a_gateway_error( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + # The MicroVM suspends itself between the state check and the request. + # The endpoint answers 502, and the worker answers 503 while it suspends. + endpoint = fake_endpoint([502, 503]) + client = fake_client( + endpoint=endpoint.url, states=["RUNNING", "SUSPENDED", "RUNNING", "RUNNING"] + ) + + 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, auto_suspend_on_idle=True), body), + None, + [("build", succeed_with("done"))], + ) + + assert outcome(result) == {"result": "done"} + assert len(endpoint.requests) == 3 + assert len(client.requests("get_microvm")) == 4 + assert len(client.requests("resume_microvm")) == 1 + + +def test_does_not_check_the_state_again_after_a_rejection( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([429]) + 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, auto_suspend_on_idle=True), body), + None, + [("build", succeed_with("done"))], + ) + + assert outcome(result) == {"result": "done"} + assert len(endpoint.requests) == 2 + assert len(client.requests("get_microvm")) == 1 + + +def test_fails_the_job_at_once_when_the_microvm_no_longer_exists( + run: Any, fake_client: Any, fake_endpoint: Any, make_error: Any +) -> None: + endpoint = fake_endpoint() + client = fake_client( + endpoint=endpoint.url, + states=[make_error("ResourceNotFoundException", 404)], + ) + + 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, auto_suspend_on_idle=True), body) + ) + + error = outcome(result) + assert error["error"] == "MicrovmNotRunningError" + assert error["message"].startswith( + 'MicroVM "build": the job could not be delivered: MicroVM mvm-1 no longer exists.' + ) + # The state error is not retried, and no request is sent. + assert len(client.requests("get_microvm")) == 1 + assert endpoint.requests == [] + assert len(client.requests("terminate_microvm")) == 1 + + +def test_lets_the_handler_catch_the_not_running_error_on_replay_too( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + client = fake_client(endpoint=fake_endpoint().url, states=["TERMINATED"]) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + try: + vm.invoke(ctx, "build", job_input=1, options=job()) + except MicrovmNotRunningError as error: + caught = type(error).__name__ + # The wait ends the invocation. The next invocation replays the + # failed job, and the handler catches the same class. + ctx.wait(Duration.from_seconds(1), name="pause") + return caught + + result = run( + session_handler(session_config(client, auto_suspend_on_idle=True), body) + ) + + assert outcome(result) == {"result": "MicrovmNotRunningError"} + assert len(client.requests("get_microvm")) == 1 + + +def test_terminates_only_after_the_handler_ends( + fake_client: Any, fake_endpoint: Any +) -> None: + client = fake_client(endpoint=fake_endpoint().url) + handler = session_handler(session_config(client), two_jobs) + runner = DurableFunctionTestRunner(handler=handler, poll_interval=0.05) + try: + arn = runner.run_async(input="null", execution_timeout=WAIT) + build = runner.wait_for_callback(arn, name="build.callback", timeout=WAIT) + # The invocation ended while the handler waited for the first job. + assert client.requests("terminate_microvm") == [] + runner.send_callback_success(build, b"1") + test = runner.wait_for_callback(arn, name="test.callback", timeout=WAIT) + assert client.requests("terminate_microvm") == [] + runner.send_callback_success(test, b"2") + result = runner.wait_for_result(arn, timeout=WAIT) + finally: + runner.close() + + assert outcome(result)["result"]["test"] == 2 + assert len(client.requests("terminate_microvm")) == 1 + + +def test_session_requests_match_the_lambda_microvms_api_model( + run: Any, fake_endpoint: Any +) -> None: + """botocore validates RunMicrovm with an idle policy, GetMicrovm, and + ResumeMicrovm against the service model.""" + + endpoint = fake_endpoint() + client = boto3.session.Session().client("lambda-microvms", region_name="us-west-2") + policy = { + "maxIdleDurationSeconds": 300, + "suspendedDurationSeconds": 3600, + "autoResumeEnabled": True, + } + config = session_config(client, idle_policy=policy, auto_suspend_on_idle=True) + started = datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC) + microvm = { + "microvmId": "mvm-1", + "endpoint": endpoint.url, + "imageArn": config.image_identifier, + "imageVersion": "1", + "maximumDurationInSeconds": 3900, + "startedAt": started, + } + with Stubber(client) as stub: + stub.add_response( + "run_microvm", + {**microvm, "state": "PENDING"}, + { + "imageIdentifier": config.image_identifier, + "executionRoleArn": config.execution_role_arn, + "clientToken": ANY, + "ingressNetworkConnectors": [INGRESS], + "egressNetworkConnectors": ANY, + "maximumDurationInSeconds": 3900, + "runHookPayload": ANY, + "idlePolicy": policy, + }, + ) + stub.add_response( + "get_microvm", + {**microvm, "state": "SUSPENDED"}, + {"microvmIdentifier": "mvm-1"}, + ) + stub.add_response("resume_microvm", {}, {"microvmIdentifier": "mvm-1"}) + stub.add_response( + "get_microvm", + {**microvm, "state": "RUNNING"}, + {"microvmIdentifier": "mvm-1"}, + ) + stub.add_response( + "create_microvm_auth_token", + {"authToken": {"x-amz-microvm-token": "t"}}, + { + "microvmIdentifier": "mvm-1", + "expirationInMinutes": 5, + "allowedPorts": [{"port": 8080}], + }, + ) + stub.add_response("terminate_microvm", {}, {"microvmIdentifier": "mvm-1"}) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke(ctx, "build", job_input=1, options=job()) + + result = run(session_handler(config, body), None, [("build", succeed_with(3))]) + + stub.assert_no_pending_responses() + client.close() + assert outcome(result) == {"result": 3} + + +def test_terminates_and_fails_for_a_launch_checkpoint_of_another_delivery( + run: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # A launch checkpoint from another release, for example after a rollback. + original = operation._MicrovmScope._launch_call # noqa: SLF001 + + def launch_for_another(self: Any, *args: Any) -> Any: + launched = original(self, *args) + return operation.LaunchResult( + microvm_id=launched.microvm_id, delivery="smoke-signal" + ) + + monkeypatch.setattr(operation._MicrovmScope, "_launch_call", launch_for_another) # noqa: SLF001 + client = fake_client() + + result = run(session_handler(session_config(client), two_jobs)) + + assert outcome(result) == { + "error": "MicrovmDeliveryError", + "message": 'MicroVM "pipeline": the launch prepared the delivery ' + '"smoke-signal", which this release does not support', + } + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] 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 f3e9f1da6..d86415077 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 @@ -42,8 +42,12 @@ MAX_RUN_HOOK_PAYLOAD_LENGTH, MicrovmConfig, MicrovmError, + MicrovmInvokeOptions, MicrovmRequestConfig, + MicrovmSession, + MicrovmSessionConfig, microvm, + microvm_session, ) @@ -223,3 +227,60 @@ def job(job_input: Any, context: MicrovmJobContext) -> Any: ) assert outcome == {"result": {"route": {"n": 1}}} + + +def test_runs_two_session_jobs_in_one_worker(fake_client: Any, tmp_path: Any) -> None: + """The second job reads what the first job wrote in the same MicroVM.""" + + def job(job_input: Any, context: MicrovmJobContext) -> Any: + state = tmp_path / "artifact.txt" + if job_input == "build": + state.write_text(f"built in {context.microvm_id}") + return "built" + return state.read_text() + + holder: dict[str, MicrovmSessionConfig] = {} + + def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: + options = MicrovmInvokeOptions(timeout=Duration.from_minutes(1)) + return [ + vm.invoke(ctx, "build", job_input="build", options=options), + vm.invoke(ctx, "test", job_input="test", options=options), + ] + + @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) + worker = start_microvm_worker( + handler=job, + create_client=lambda region: callbacks, # type: ignore[arg-type, return-value] + port=0, + ) + try: + client = fake_client( + launches=[run_hook_to(worker)], + endpoint=f"http://127.0.0.1:{worker.port}", + ) + 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() + + assert result.status is InvocationStatus.SUCCEEDED, result.error + assert json.loads(result.result or "null") == ["built", "built in mvm-1"] + assert len(client.requests("run_microvm")) == 1 + assert len(client.requests("create_microvm_auth_token")) == 2 + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] From cfe6737087020f6298d8fec89f8e5db952cebb1b Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Thu, 15 Oct 2026 10:22:15 -0700 Subject: [PATCH 06/10] fix(extras): keep the MicroVM for invocation retries A retryable InvocationError, such as RetryableSerDesError, makes Lambda run the invocation again. The session terminated the MicroVM for it. The terminate step then took the position of the handler's next operation, and the replay failed with NonDeterministicExecutionError. The session now terminates only for an error that ends it. - The job wait in microvm and in vm.invoke passes an InvocationError through, instead of recording it as a job failure. microvm keeps the MicroVM for a retryable one. - A ResumeMicrovm answer of ResourceNotFoundException now fails the job with MicrovmNotRunningError, as a GetMicrovm answer does. - The README session example has its imports. - New tests: parallel jobs, two items that resume one MicroVM, and the invocation error paths. --- .../README.md | 4 + .../microvm/lifecycle.py | 23 ++- .../microvm/operation.py | 8 + .../microvm/session.py | 16 ++ .../tests/microvm/lifecycle_test.py | 11 ++ .../tests/microvm/operation_test.py | 62 +++++- .../tests/microvm/session_test.py | 177 +++++++++++++++++- 7 files changed, 289 insertions(+), 12 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index 5c488ddbc..3dd7f38d2 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -156,8 +156,12 @@ The MicroVM execution role (`execution_role_arn`) needs `lambda:SendDurableExecu `microvm_session` launches one MicroVM and runs a handler that sends jobs to it. The jobs share the MicroVM's files and processes. So a later job can use what an earlier job wrote. The MicroVM is terminated when the handler returns or raises. ```python +import os from typing import Any +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.execution import durable_execution from aws_durable_execution_sdk_python_extras.microvm import ( MicrovmInvokeOptions, MicrovmSession, 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 b4754c6b4..539b5175a 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 @@ -130,17 +130,20 @@ def resume_microvm() -> Any: unavailable, ) + def not_found() -> MicrovmStateError: + msg = ( + f"MicroVM {microvm_id} no longer exists. A session MicroVM lives at " + "most its session timeout plus 5 minutes, and never longer than 8 " + "hours." + ) + return MicrovmStateError(msg) + while True: try: current = get_microvm() except ClientError as error: if _code(error) == "ResourceNotFoundException": - msg = ( - f"MicroVM {microvm_id} no longer exists. A session MicroVM " - "lives at most its session timeout plus 5 minutes, and never " - "longer than 8 hours." - ) - raise MicrovmStateError(msg) from error + raise not_found() from error raise state = current.get("state") last_state = state @@ -161,8 +164,12 @@ def resume_microvm() -> Any: except ClientError as error: # ConflictException means that a transition is already in # progress, for example an auto-resume. Polling continues - # either way. - if _code(error) != "ConflictException": + # either way. ResourceNotFoundException means that the + # MicroVM was removed after the state check. + code = _code(error) + if code == "ResourceNotFoundException": + raise not_found() from error + if code != "ConflictException": raise resume_sent = True options.logger.info("resuming MicroVM", extra={"state": state}) 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 1f2b79051..971275f30 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 @@ -26,6 +26,7 @@ CallbackError, CallbackTimeoutError, ChildContextError, + InvocationError, StepError, ValidationError, ) @@ -712,6 +713,13 @@ def run_job(self, child: DurableContext, job_input: Any) -> Any: try: result = callback.result() + except InvocationError as error: + # An SDK invocation error is not a job outcome. A retryable one + # makes Lambda run the invocation again, and the replay waits for + # the same MicroVM. So the MicroVM stays up for it. + if not error.is_retryable(): + self.terminate(child, microvm_id) + raise except Exception as error: # SuspendExecution is a BaseException, so a wait for the callback # passes through here without terminating the MicroVM. 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 4536fc9a4..5000a3def 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 @@ -13,6 +13,7 @@ from aws_durable_execution_sdk_python.exceptions import ( ChildContextError, + InvocationError, ValidationError, ) @@ -282,6 +283,10 @@ def run(job_context: DurableContext) -> Any: ) try: return callback.result() + except InvocationError: + # An SDK invocation error is not a job outcome. A retryable + # one makes Lambda run the invocation again. + raise except Exception as error: # SuspendExecution is a BaseException, so a wait for the # callback passes through here. @@ -362,6 +367,17 @@ def run_session( raise unknown_delivery_error(self.name, launched) try: result = handler(MicrovmSession(self, launched), child) + except InvocationError as error: + # A retryable InvocationError, such as RetryableSerDesError, makes + # Lambda run the invocation again. The child context records no + # outcome for it, and the handler replays on the same MicroVM. + # A terminate step recorded now would take the position of the + # handler's next operation, and the replay would fail with + # NonDeterministicExecutionError. So the session terminates only + # for an error that ends the session. + if not error.is_retryable(): + self.terminate(child, launched.microvm_id) + raise except Exception: # SuspendExecution is a BaseException. A handler that waits, for # a job or a durable wait, passes through here without 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 be0f99c4b..3f2eed380 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 @@ -211,3 +211,14 @@ def test_stops_a_stalled_get_microvm_at_the_reserve(fake_client: Any) -> None: run_with(client, remaining=lambda: INVOCATION_RESERVE_SECONDS + 1.2) finally: release.set() + + +def test_raises_the_state_error_when_the_microvm_is_removed_before_the_resume( + fake_client: Any, make_error: Any +) -> None: + client = fake_client( + states=["SUSPENDED"], resumes=[make_error("ResourceNotFoundException", 404)] + ) + + with pytest.raises(MicrovmStateError, match="no longer exists"): + run_with(client) diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py index c03677e9d..0ef54173d 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py @@ -9,6 +9,7 @@ import hashlib import json import threading +import time from collections.abc import Callable, Iterator from types import SimpleNamespace from typing import Any @@ -16,8 +17,11 @@ import boto3 import pytest from aws_durable_execution_sdk_python.config import Duration -from aws_durable_execution_sdk_python.context import DurableContext -from aws_durable_execution_sdk_python.exceptions import ChildContextError +from aws_durable_execution_sdk_python.context import Callback, DurableContext +from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + InvocationError, +) from aws_durable_execution_sdk_python.execution import ( InvocationStatus, durable_execution, @@ -1028,3 +1032,57 @@ def handler(event: Any, context: DurableContext) -> Any: "/job", "/durable-execution/v1/job", ] + + +@pytest.mark.parametrize("retryable", [True, False]) +def test_passes_an_invocation_error_of_the_wait_through( + fake_client: Any, monkeypatch: pytest.MonkeyPatch, retryable: bool +) -> None: + # An SDK invocation error is not a job outcome. A retryable one makes + # Lambda run the invocation again, and the replay waits for the same + # MicroVM. So only an error that is not retryable terminates it. + class Interrupted(InvocationError): + def is_retryable(self) -> bool: + return retryable + + original = Callback.result + raised = threading.Event() + + def result_once(self: Any) -> Any: + if not raised.is_set(): + raised.set() + msg = "interrupted" + raise Interrupted(msg) + return original(self) + + monkeypatch.setattr(Callback, "result", result_once) + order: list[str] = [] + client = fake_client(terminates=[lambda request: order.append("terminate")]) + runner = DurableFunctionTestRunner( + handler=catching(config_with(client)), poll_interval=0.05 + ) + try: + arn = runner.run_async(input="{}", execution_timeout=WAIT) + if retryable: + callback_id = runner.wait_for_callback( + arn, name="build.callback", timeout=WAIT + ) + # A terminate for the error would run at once after the raise, + # in the same thread. So 0.5 seconds after it, the order shows + # whether one ran. + assert raised.wait(WAIT) + time.sleep(0.5) + order.append("complete") + runner.send_callback_success(callback_id, b'"done"') + result = runner.wait_for_result(arn, timeout=WAIT) + finally: + runner.close() + + # An error that is not retryable is recorded as the operation's failure. + # It is not a MicroVM error type, so it becomes a plain MicrovmError. + assert outcome(result) == ( + {"result": "done"} + if retryable + else {"error": "MicrovmError", "message": 'MicroVM "build": interrupted'} + ) + assert order == (["complete", "terminate"] if retryable else ["terminate"]) 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 20c53faf5..30273a8d8 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 @@ -13,10 +13,12 @@ import boto3 import pytest -from aws_durable_execution_sdk_python.config import Duration -from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.config import Duration, StepConfig +from aws_durable_execution_sdk_python.context import Callback, DurableContext from aws_durable_execution_sdk_python.exceptions import ( ChildContextError, + InvocationError, + RetryableSerDesError, ValidationError, ) from aws_durable_execution_sdk_python.execution import ( @@ -25,6 +27,7 @@ ) from aws_durable_execution_sdk_python.lambda_service import ErrorObject from aws_durable_execution_sdk_python.retries import RetryDecision +from aws_durable_execution_sdk_python.serdes import SerDes from aws_durable_execution_sdk_python_testing.runner import ( DurableFunctionTestResult, DurableFunctionTestRunner, @@ -803,3 +806,173 @@ def launch_for_another(self: Any, *args: Any) -> Any: '"smoke-signal", which this release does not support', } assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +class FlakySerDes(SerDes[Any]): + """A serdes whose store fails once, for a reason that a retry can fix.""" + + def __init__(self) -> None: + self.failed = False + + def serialize(self, value: Any, serdes_context: Any) -> str: + if not self.failed: + self.failed = True + msg = "store throttled" + raise RetryableSerDesError(msg) + return json.dumps(value) + + def deserialize(self, data: str, serdes_context: Any) -> Any: + return json.loads(data) + + +def test_keeps_the_microvm_when_the_handler_raises_a_retryable_invocation_error( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + # A retryable InvocationError makes Lambda run the invocation again. A + # terminate step recorded for it would take the position of the next + # job, and the replay would fail with NonDeterministicExecutionError. + client = fake_client(endpoint=fake_endpoint().url) + serdes = FlakySerDes() + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + fetched = ctx.step( + lambda _: {"n": 1}, name="fetch", config=StepConfig(serdes=serdes) + ) + return vm.invoke(ctx, "build", job_input=fetched, options=job()) + + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with("built"))], + ) + + assert outcome(result) == {"result": "built"} + assert serdes.failed + assert [name for name, _ in client.calls] == [ + "run_microvm", + "create_microvm_auth_token", + "terminate_microvm", + ] + + +def test_passes_a_retryable_invocation_error_of_a_job_wait_through( + run: Any, + fake_client: Any, + fake_endpoint: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + # The SDK error is not a job outcome. So the job does not fail with it, + # and the next invocation waits for the same job. + original = Callback.result + raised: list[bool] = [] + + def result_once(self: Any) -> Any: + if not raised: + raised.append(True) + msg = "store throttled" + raise RetryableSerDesError(msg) + return original(self) + + monkeypatch.setattr(Callback, "result", result_once) + endpoint = fake_endpoint() + client = fake_client(endpoint=endpoint.url) + + result = run( + session_handler( + session_config(client), + lambda vm, ctx, event: vm.invoke(ctx, "build", job_input=1, options=job()), + ), + None, + [("build", succeed_with("built"))], + ) + + assert outcome(result) == {"result": "built"} + assert raised == [True] + assert len(endpoint.requests) == 1 + assert len(client.requests("terminate_microvm")) == 1 + + +def test_terminates_for_an_invocation_error_that_is_not_retryable( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + client = fake_client(endpoint=fake_endpoint().url) + + class Fatal(InvocationError): + def is_retryable(self) -> bool: + return False + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + msg = "checkpoint rejected" + raise Fatal(msg) + + result = run(session_handler(session_config(client), body)) + + # The SDK records the error as the session's failure. + assert outcome(result) == { + "error": "ChildContextError", + "message": "checkpoint rejected", + } + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +def test_runs_jobs_in_parallel_branches( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint() + client = fake_client(endpoint=endpoint.url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return ctx.parallel( + [ + lambda branch: vm.invoke(branch, "lint", job_input="a", options=job()), + lambda branch: vm.invoke(branch, "unit", job_input="b", options=job()), + ], + name="checks", + ).get_results() + + result = run( + session_handler(session_config(client, auto_suspend_on_idle=True), body), + None, + [("lint", succeed_with("A")), ("unit", succeed_with("B"))], + ) + + assert outcome(result) == {"result": ["A", "B"]} + assert len(client.requests("run_microvm")) == 1 + assert sorted(request["body"]["input"] for request in endpoint.requests) == [ + "a", + "b", + ] + assert len(client.requests("get_microvm")) == 2 + assert len(client.requests("terminate_microvm")) == 1 + + +def test_resumes_once_when_two_items_find_the_microvm_suspended( + 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. + endpoint = fake_endpoint() + client = fake_client( + endpoint=endpoint.url, + states=["SUSPENDED"] * 4, + resumes=[None, make_error("ConflictException", 409)], + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return ctx.map( + ["a", "b"], + lambda item_context, item, index, items: vm.invoke( + item_context, f"shard-{item}", job_input=item, options=job() + ), + name="shards", + ).get_results() + + result = run( + session_handler(session_config(client, auto_suspend_on_idle=True), body), + None, + [("shard-a", succeed_with("A")), ("shard-b", succeed_with("B"))], + ) + + assert outcome(result) == {"result": ["A", "B"]} + assert len(client.requests("resume_microvm")) == 2 + assert len(endpoint.requests) == 2 From f2fe5fa504db240d884979c078b7f6f62a4143cd Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Thu, 15 Oct 2026 10:32:56 -0700 Subject: [PATCH 07/10] fix(extras): keep the token off followed redirects A caller's opener that follows redirects sent the auth token to the redirect target, also on another host. The POST now adds its headers as unredirected headers, which urllib does not copy to a followed request. Two fixes from the review of the JavaScript port apply here too: - The request path must not have a "." or ".." segment, also with %2e. The JavaScript SDK's fetch resolves such a segment, so both SDKs now accept the same paths. - The short-timeout message rounds the time down. So it never reads 11.0 seconds next to "at least 11 seconds". The README states that the 11 seconds must be left when the request step starts, and which jobs go over HTTP. --- .../README.md | 6 +-- .../microvm/config.py | 10 ++--- .../microvm/operation.py | 19 ++++++++ .../microvm/request.py | 16 +++++-- .../tests/conftest.py | 17 ++++--- .../tests/microvm/request_test.py | 31 +++++++++++-- .../tests/microvm/validation_test.py | 44 +++++++++++++++++++ 7 files changed, 122 insertions(+), 21 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index 579cea86c..2622f84a7 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -65,7 +65,7 @@ The same job can arrive again, also after it ended, so ignore a second request f 2. The request step stops waiting for a request after at most 10 seconds, but the request keeps running in its thread and reaches the route. The next attempt sends the job again. 3. The invocation ends after the route accepted the job, and before the request step checkpointed its success. The step then runs again in the next invocation. -The path must be an RFC 3986 path, and the operation sends it exactly as written. Percent-encode any other character, for example a space as `%20`. The workers match the raw request path, so register the same string as the worker's route. An invalid path raises `ValidationError` before any durable operation starts. +The path must be an RFC 3986 path without a `.` or `..` segment, and the operation sends it exactly as written. A dot segment is rejected also when the dot is written as `%2e`, because the JavaScript SDK's `fetch` resolves such a segment, and a proxy or server can too. So both SDKs accept the same paths. Percent-encode any other character, for example a space as `%20`. The workers match the raw request path, so register the same string as the worker's route. An invalid path raises `ValidationError` before any durable operation starts. ### What it does @@ -125,7 +125,7 @@ A step retry costs one durable operation and one more Lambda invocation. Most fa A route that answers another 4xx status fails the operation without a retry, with `MicrovmDeliveryError`. So does a redirect, 3xx. The operation does not follow redirects, and the message names the `Location`. Set the path that the route serves, for example with its trailing slash. The JavaScript operation follows redirects. -The POST ignores the proxy settings in the environment, as the JavaScript SDK's `fetch` does. boto3 reads `HTTPS_PROXY`. So in a network that reaches the MicroVM endpoint only through a proxy, the launch succeeds and every POST fails. Pass an `opener` in the config for that network, for example `urllib.request.build_opener(urllib.request.ProxyHandler({"https": "http://proxy.internal:3128"}))`. When an opener follows a redirect, the delivery still fails at once with `MicrovmDeliveryError`. +The POST ignores the proxy settings in the environment, as the JavaScript SDK's `fetch` does. boto3 reads `HTTPS_PROXY`. So in a network that reaches the MicroVM endpoint only through a proxy, the launch succeeds and every POST fails. Pass an `opener` in the config for that network, for example `urllib.request.build_opener(urllib.request.ProxyHandler({"https": "http://proxy.internal:3128"}))`. When an opener follows a redirect, the delivery still fails at once with `MicrovmDeliveryError`. The followed request carries no auth token, because the operation adds its headers as unredirected headers. The 10-second reserve keeps time for the step to record its failure. A Lambda timeout during a step attempt records no outcome for the attempt. So every call of the first tier must end before the reserve. boto3 cannot cancel a call, so a `CreateMicrovmAuthToken` call runs in a separate thread, and the step stops waiting for it at the reserve. Each POST runs in a separate thread too, and the step waits for it at most 10 seconds, or until the reserve. A socket timeout alone would not bound it, because each byte received starts the socket timeout again. @@ -135,7 +135,7 @@ The 10-second reserve keeps time for the step to record its failure. A Lambda ti - `job_input` is sent as JSON. So the MicroVM receives what `json.dumps` writes: a tuple arrives as a list, and a dict key that is not a string arrives as a string. - The MicroVM gets the `INTERNET_EGRESS` connector, because it needs egress to reach the Lambda callback APIs. For ingress, it gets `NO_INGRESS` with `run` hook delivery and `ALL_INGRESS` with HTTP delivery. If you set `ingress_network_connectors`, allow ingress whenever an input can exceed the `run` hook limit. An empty tuple sends no connector. - `heartbeat_timeout` starts when the callback is created, before the launch. The worker sends its first heartbeat when the job arrives. So the first period must cover the launch, the boot, and, for HTTP delivery, the request step, each with its retries. A throttled launch can take several minutes, and the request step retries for up to `retry_window` in each attempt. With a shorter heartbeat timeout, the operation fails with `MicrovmTimeoutError` before the MicroVM sends its first heartbeat. -- HTTP delivery needs a function timeout of at least 11 seconds. The request step keeps a 10-second reserve before the invocation times out, and it starts a call only with at least 1 second left before the reserve. With less time, every attempt fails with `MicrovmEndpointUnavailableError`, and the message names the function timeout. The Lambda default timeout is 3 seconds. A function with that timeout runs small jobs in the `run` hook, and fails every job too large for it. +- The request step needs at least 11 seconds left in the invocation when it starts. It keeps a 10-second reserve before the invocation times out, and it starts a call only with at least 1 second left before the reserve. The launch step can run earlier in the same invocation, so the function timeout must cover it too. A function timeout of `retry_window` plus 11 seconds, 71 seconds by default, leaves the first retry tier about its whole window. With less than 11 seconds left, every attempt fails with `MicrovmEndpointUnavailableError`, and the message of the final `MicrovmDeliveryError` names the function timeout. The Lambda default timeout is 3 seconds. With it, `microvm` runs small jobs in the `run` hook and fails every job that goes over HTTP: a job too large for the hook, and a job with `request.path`. - The Region comes from the durable execution ARN, or from `AWS_REGION`. - Without a `client`, the operation creates one `lambda-microvms` client per Region, on first use. It has a 5-second connect timeout, a 30-second read timeout, and 3 attempts in botocore's standard retry mode. So one launch or terminate attempt takes at most about 2 minutes. A call that runs past the function's timeout records no outcome for its step attempt, and the step runs again. The launch is idempotent through its client token, so it is repeated, not doubled. 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 35290ae35..989c1ee92 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 @@ -42,8 +42,8 @@ class MicrovmRequestConfig: path: str | None = None """The route path, for example ``/clone-build``. - It must be an RFC 3986 path that starts with ``/``. The operation sends - it exactly as written. Percent-encode any other character, for example + It must be an RFC 3986 path that starts with ``/``, without a ``.`` or + ``..`` segment. The operation sends it exactly as written. Percent-encode any other character, for example a space as ``%20``. Register the same string as the worker's route, because the worker matches the raw request path. @@ -186,9 +186,9 @@ class MicrovmConfig: ) The operation does not follow redirects. When an opener follows one, the - delivery fails at once with ``MicrovmRequestRejectedError``. A followed - redirect to another host also receives the auth token, so prefer an - opener that follows no redirect. + delivery fails at once with ``MicrovmRequestRejectedError``. The followed + request carries no auth token, because the operation adds its headers as + unredirected headers. """ client: LambdaMicrovmsClient | None = None 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 53a50b3c6..ceb6334aa 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 @@ -111,6 +111,13 @@ path, so the config and the route use the same string. """ +_DOT_SEGMENT = re.compile(r"/(?:\.|%2e){1,2}(?=/|$)", re.IGNORECASE) +"""A ``.`` or ``..`` segment, also with ``%2e`` for a dot. + +A segment with other characters next to its dots, such as ``/.well-known``, +and a segment of three dots are not dot segments. +""" + @dataclass(frozen=True, kw_only=True) class LaunchResult: @@ -703,6 +710,18 @@ def _validate_request(name: str, request: MicrovmRequestConfig) -> int: "space as %20." ) raise ValidationError(msg) + # The JavaScript SDK's fetch resolves a dot segment before it sends the + # request, so the JavaScript operation rejects such a path. urllib sends + # the path unchanged, but a proxy or server in front of the worker can + # resolve it too. So both SDKs accept the same paths, and each path + # reaches the worker as written. + if path is not None and _DOT_SEGMENT.search(path): + msg = ( + f'MicroVM "{name}": the request path must not have a "." or ".." ' + "segment, because an HTTP client or server can resolve it before " + "the route sees it." + ) + raise ValidationError(msg) if path is not None and path.startswith(_HOOK_PATH_PREFIX): msg = ( f'MicroVM "{name}": the request path must not use the lifecycle hook ' 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 c4813f228..d2345fa75 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 @@ -10,6 +10,7 @@ from __future__ import annotations import functools +import math import http.client import threading import time @@ -229,11 +230,13 @@ def send_job(options: SendJobOptions) -> int: at once. The function does not follow a redirect. The JavaScript SDK's ``fetch`` - follows one, and it differs in three ways. A 301, 302, or 303 turns the + follows one, and it differs in two ways. A 301, 302, or 303 turns the POST into a GET without the job, and the GET's 2xx counts as accepted. So the job never starts, and the operation waits until its timeout. A 307 or 308 to another host receives the auth token. So a redirect fails - the delivery at once, with the ``Location`` in the message. + the delivery at once, with the ``Location`` in the message. The headers + are added as unredirected headers, so an opener that follows a redirect + does not send the auth token on either. The delay starts at 250 milliseconds and doubles up to 4 seconds. The tier ends at the retry window, or 10 seconds before the invocation times @@ -291,7 +294,8 @@ def send_job(options: SendJobOptions) -> int: needed = INVOCATION_RESERVE_SECONDS + MIN_CALL_SECONDS if remaining is not None and remaining < needed: msg = ( - f"The request step had {remaining:.1f} seconds left in the Lambda " + f"The request step had {math.floor(remaining * 10) / 10:.1f} seconds " + "left in the Lambda " f"invocation, and it needs at least {needed:.0f} seconds to send " "the job. Raise the function timeout." ) @@ -486,8 +490,12 @@ def _post( each socket operation, so :func:`_post_within` bounds the whole request. """ request = urllib.request.Request( # noqa: S310 - the URL is https or the endpoint's own - url, data=body, headers=dict(headers), method="POST" + url, data=body, method="POST" ) + # urllib does not copy these headers to a redirected request. So an + # opener that follows a redirect does not send the auth token on. + for name, value in headers.items(): + request.add_unredirected_header(name, value) try: with opener.open(request, timeout=timeout) as response: # An opener that follows redirects reports the URL it reached. 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 253f2043a..50536a747 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py @@ -117,10 +117,13 @@ class FakeEndpoint: 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``. + without an answer. A redirect status gets ``Location: ``, + ``/moved`` by default. """ - def __init__(self, statuses: list[int | None] | None = None) -> None: + def __init__( + self, statuses: list[int | None] | None = None, location: str = "/moved" + ) -> None: self.requests: list[dict[str, Any]] = [] self._statuses = list(statuses or []) self._lock = threading.Lock() @@ -134,7 +137,7 @@ def do_GET(self) -> None: # noqa: N802 - the http.server name { "method": "GET", "path": self.path, - "headers": {}, + "headers": {k.lower(): v for k, v in self.headers.items()}, "body": None, } ) @@ -160,7 +163,7 @@ def do_POST(self) -> None: # noqa: N802 - the http.server name return self.send_response(status) if 300 <= status < 400: # noqa: PLR2004 - self.send_header("location", "/moved") + self.send_header("location", location) self.send_header("content-length", "0") self.end_headers() @@ -184,8 +187,10 @@ def close(self) -> None: def fake_endpoint() -> Iterator[Callable[..., FakeEndpoint]]: endpoints: list[FakeEndpoint] = [] - def _create(statuses: list[int | None] | None = None) -> FakeEndpoint: - endpoint = FakeEndpoint(statuses) + def _create( + statuses: list[int | None] | None = None, location: str = "/moved" + ) -> FakeEndpoint: + endpoint = FakeEndpoint(statuses, location) endpoints.append(endpoint) return endpoint 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 b40104e6a..4af60f6d8 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 @@ -209,9 +209,12 @@ def test_ends_the_tier_before_the_invocation_reserve( assert sleeps == [0.25, 0.5, 1.0, 2.0] -@pytest.mark.parametrize("remaining", [0.5, 2.9, 10.9]) +@pytest.mark.parametrize( + ("remaining", "shown"), + [(0.5, "0.5"), (2.9, "2.9"), (10.9, "10.9"), (10.96, "10.9")], +) def test_names_the_function_timeout_when_too_little_time_is_left( - fake_client: Any, fake_endpoint: Any, remaining: float + fake_client: Any, fake_endpoint: Any, remaining: float, shown: str ) -> None: # The Lambda default timeout is 3 seconds. Such a function can never # send a job over HTTP. @@ -222,7 +225,7 @@ def test_names_the_function_timeout_when_too_little_time_is_left( send_job(options(client, endpoint.url, [], remaining_seconds=lambda: remaining)) assert str(raised.value) == ( - f"The request step had {remaining:.1f} seconds left in the Lambda " + f"The request step had {shown} seconds left in the Lambda " "invocation, and it needs at least 11 seconds to send the job. Raise " "the function timeout." ) @@ -471,3 +474,25 @@ def interrupted() -> None: with pytest.raises(KeyboardInterrupt): call_before_deadline(interrupted, time.monotonic() + 30, RuntimeError) + + +def test_sends_no_auth_token_on_a_redirect_that_the_callers_opener_follows( + fake_client: Any, fake_endpoint: Any +) -> None: + # The opener follows a 302 to another host. That host must not receive + # the token, because urllib does not copy an unredirected header. + other_host = fake_endpoint() + endpoint = fake_endpoint([302], location=f"{other_host.url}/moved") + + with pytest.raises(MicrovmRequestRejectedError): + send_job( + options( + fake_client(), endpoint.url, [], opener=urllib.request.build_opener() + ) + ) + + assert endpoint.requests[0]["headers"]["x-test-token"] == "token-1" + [followed] = other_host.requests + assert followed["method"] == "GET" + assert "x-test-token" not in followed["headers"] + assert "content-type" not in followed["headers"] 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 c8aed4c16..2ce94ede1 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 @@ -96,6 +96,40 @@ def config(**fields: Any) -> MicrovmConfig: {"request": MicrovmRequestConfig(path="/aws/lambda-microvms/run")}, "must not use the lifecycle hook prefix /aws/lambda-microvms/", ), + ( + "build", + {"request": MicrovmRequestConfig(path="/jobs/../build")}, + 'must not have a "." or ".." segment', + ), + ( + "build", + {"request": MicrovmRequestConfig(path="/jobs/./build")}, + 'must not have a "." or ".." segment', + ), + ( + "build", + {"request": MicrovmRequestConfig(path="/jobs/..")}, + 'must not have a "." or ".." segment', + ), + ( + "build", + {"request": MicrovmRequestConfig(path="/jobs/%2e/build")}, + 'must not have a "." or ".." segment', + ), + ( + "build", + {"request": MicrovmRequestConfig(path="/jobs/%2E%2e/build")}, + 'must not have a "." or ".." segment', + ), + ( + "build", + { + "request": MicrovmRequestConfig( + path="/jobs/../aws/lambda-microvms/runtime/v1/run" + ) + }, + 'must not have a "." or ".." segment', + ), ( "build", {"request": MicrovmRequestConfig(port=0)}, @@ -222,3 +256,13 @@ def test_reads_the_remaining_time_or_reports_no_deadline( context: Any = type("Context", (), {"lambda_context": lambda_context})() assert operation._remaining_seconds(context)() == expected # noqa: SLF001 + + +@pytest.mark.parametrize("path", ["/.well-known/job", "/.../job", "/a.b/c..d", "/"]) +def test_accepts_dots_that_are_not_a_dot_segment(path: str) -> None: + # Dots next to other characters, and three dots, are not dot segments. + # fetch and urllib leave such a path unchanged. + assert ( + operation._validate_request("build", MicrovmRequestConfig(path=path)) # noqa: SLF001 + == 60 + ) From a4a3494be8d881eabd4da64215e7120c31ca7896 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Thu, 15 Oct 2026 10:35:36 -0700 Subject: [PATCH 08/10] docs(extras): rewrap the request path docstring --- .../microvm/config.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) 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 989c1ee92..38e243778 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 @@ -43,9 +43,10 @@ class MicrovmRequestConfig: """The route path, for example ``/clone-build``. It must be an RFC 3986 path that starts with ``/``, without a ``.`` or - ``..`` segment. The operation sends it exactly as written. Percent-encode any other character, for example - a space as ``%20``. Register the same string as the worker's route, because - the worker matches the raw request path. + ``..`` segment. The operation sends it exactly as written. Percent-encode + any other character, for example a space as ``%20``. Register the same + string as the worker's route, because the worker matches the raw request + path. ``None`` sends the job to :data:`DEFAULT_MICROVM_JOB_PATH`. The worker package serves that path with its ``handler``. A path sends the job to one From cac75430beb0fae703d113af79622ba51d4d0dcf Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Thu, 15 Oct 2026 11:45:28 -0700 Subject: [PATCH 09/10] fix(extras): refuse the outer context in invoke A job created in the context that microvm_session was called with lands outside the session's child context. The core SDK does not refuse that. So the job ran, and the next replay failed the execution with NonDeterministicExecutionError. vm.invoke now raises TypeError for that context, before the job's durable operations start. The docstring and the README said that an uncaught MicroVM error keeps its class. That holds only for an error that the handler raises itself. A child context that the handler created records ChildContextError as its error type, so the caller gets the message but not the class. --- .../README.md | 6 +- .../microvm/session.py | 31 ++++++-- .../tests/microvm/session_test.py | 77 +++++++++++++++++++ 3 files changed, 106 insertions(+), 8 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index c8281577a..571ec8317 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -220,13 +220,13 @@ pipeline child context - The client token of the launch is the SHA-256 hex digest of the session ID. The session ID is random, and it is checkpointed. So a retried or replayed launch returns the same MicroVM. - Every job goes over HTTP. The MicroVM gets the `ALL_INGRESS` connector. A job without a `path` goes to the worker's `handler`. - The first job is sent right after the launch. The MicroVM endpoint refuses requests until the `run` hook returns, and the request step retries a refused request inside its attempt. So the session needs no separate readiness wait. -- `vm.invoke` takes the context that the call runs in, as `microvm` does. Inside a `map` or `parallel` item, pass the item's context. +- `vm.invoke` takes the context that the call runs in, as `microvm` does. Inside a `map` or `parallel` item, pass the item's context. A job must be created inside the session's child context. A job created in the context that `microvm_session` was called with runs, and the next replay then fails the execution with `NonDeterministicExecutionError`. So `vm.invoke` raises `TypeError` for that context. The core SDK does not check other contexts outside the session. - `config.timeout` bounds the MicroVM's lifetime: it ends at `timeout` plus 5 minutes, at most 8 hours. Each job waits for its own `timeout`. So set a `heartbeat_timeout` on every job. A job still running when the MicroVM ends then fails at its heartbeat timeout. A failed `vm.invoke` raises its MicroVM error inside the handler. The handler can catch it and send another job to the same MicroVM. An error that leaves the handler ends the session: -1. A MicroVM error reaches the caller with its own class, such as `MicrovmJobFailedError`. -2. Any other error reaches the caller as the SDK's `ChildContextError`, as from `run_in_child_context`. +1. A MicroVM error that the handler raises itself reaches the caller with its own class, such as `MicrovmJobFailedError`. +2. Any other error reaches the caller as the SDK's `ChildContextError`, as from `run_in_child_context`. This includes a MicroVM error that leaves a child context that the handler created: that child context records `ChildContextError` as its error type. The caller then gets the MicroVM error's message, but not its class. Inside the handler, the `ChildContextError` of that child context has the MicroVM error as its `__cause__`. The JavaScript operation behaves the same way. An invalid job raises `ValidationError` or `TypeError` from `vm.invoke`, before the job's durable operations start. The handler can catch it too. 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 5000a3def..3c3fb6137 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 @@ -8,7 +8,7 @@ import uuid from collections.abc import Callable, Mapping -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, TypeVar from aws_durable_execution_sdk_python.exceptions import ( @@ -112,10 +112,13 @@ def microvm_session( error that leaves the handler ends the session, and the MicroVM is terminated: - 1. A MicroVM error, such as :class:`MicrovmJobFailedError`, reaches the - caller with its own class. + 1. A MicroVM error, such as :class:`MicrovmJobFailedError`, that the + handler raises itself reaches the caller with its own class. 2. Any other error reaches the caller as ``ChildContextError``, as from - ``run_in_child_context``. + ``run_in_child_context``. This includes a MicroVM error that leaves + a child context that the handler created: that child context records + ``ChildContextError`` as its error type. The caller then gets the + MicroVM error's message, but not its class. The session's value is the handler's return value, checkpointed as the child context result. So it must be serializable by the SDK's default @@ -234,7 +237,8 @@ def invoke( ``SendDurableExecutionCallbackSuccess``, parsed as JSON. Raises: - TypeError: When ``context`` is not a durable context, or + TypeError: When ``context`` is not a durable context, or is the + context that ``microvm_session`` was called with, or ``job_input`` is not JSON-serializable. ValidationError: When ``name`` or ``options`` is not valid. MicrovmNotRunningError: When the session's MicroVM has ended, and @@ -255,6 +259,19 @@ def invoke( "context to create the job in, as its first argument" ) raise TypeError(msg) + # The caller's own context is often in scope in the handler's + # closure. A job created in it lands outside the session's child + # context. The core SDK does not refuse that, and the next replay + # fails with NonDeterministicExecutionError after the job has run. + # So the most likely wrong context is refused here. A full ancestry + # check would need support in the core SDK. + if context is scope.outer_context: + msg = ( + f'MicroVM session "{scope.name}": invoke received the context ' + "that microvm_session was called with. Pass the context that " + "the handler receives, or a context nested in it." + ) + raise TypeError(msg) if not isinstance(name, str) or not name: msg = ( f'MicroVM session "{scope.name}": invoke requires a non-empty job name' @@ -304,6 +321,9 @@ class _SessionScope(_MicrovmScope): auto_suspend_idle_seconds: int | None idle_policy: Mapping[str, Any] | None + outer_context: DurableContext = field(compare=False, repr=False) + """The context that ``microvm_session`` was called with. A job created + in it would be outside the session's child context.""" @property def resume_before_delivery(self) -> bool: @@ -340,6 +360,7 @@ def create( retry_strategy=config.retry_strategy or default_microvm_retry_strategy, auto_suspend_idle_seconds=auto_suspend_idle_seconds, idle_policy=config.idle_policy, + outer_context=context, ) def run_session( 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 30273a8d8..6b689952f 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 @@ -976,3 +976,80 @@ def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: assert outcome(result) == {"result": ["A", "B"]} assert len(client.requests("resume_microvm")) == 2 assert len(endpoint.requests) == 2 + + +def test_rejects_the_outer_context_before_the_job_starts( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + # A job created in the outer context lands outside the session's child + # context. The core SDK does not refuse it, and the next replay fails + # with NonDeterministicExecutionError after the job has run. + endpoint = fake_endpoint() + client = fake_client(endpoint=endpoint.url) + config = session_config(client) + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + def body(vm: MicrovmSession, ctx: DurableContext) -> Any: + try: + vm.invoke(context, "build", job_input=1, options=job()) + except TypeError as error: + return str(error) + return None + + return microvm_session(context, "pipeline", config=config, handler=body) + + result = run(handler) + + assert outcome(result) == ( + 'MicroVM session "pipeline": invoke received the context that ' + "microvm_session was called with. Pass the context that the handler " + "receives, or a context nested in it." + ) + assert endpoint.requests == [] + session = result.get_context("pipeline") + assert [operation.name for operation in session.child_operations] == [ + "pipeline.session", + "pipeline.launch", + "pipeline.terminate", + ] + + +def test_wraps_a_job_error_that_leaves_a_child_context_of_the_handler( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + # Only an error that the handler raises itself keeps its MicroVM class. + # The SDK records one error type per child context. So an error from the + # handler's own child context reaches the caller as ChildContextError, + # with the MicroVM error's message. Inside the handler, the + # ChildContextError's __cause__ is still the MicroVM class. + client = fake_client(endpoint=fake_endpoint().url) + config = session_config(client) + inside: list[str] = [] + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + def body(vm: MicrovmSession, ctx: DurableContext) -> Any: + try: + return ctx.run_in_child_context( + lambda stage: vm.invoke(stage, "build", job_input=1, options=job()), + name="stage", + ) + except ChildContextError as error: + inside.append(type(error.__cause__).__name__) + raise + + try: + microvm_session(context, "pipeline", config=config, handler=body) + except ChildContextError as error: + return [type(error.__cause__).__name__, error.message] + return None + + result = run(handler, None, [("build", fail_with("BuildError", "failed"))]) + + assert outcome(result) == [ + "ChildContextError", + 'MicroVM "build": the job failed (BuildError): failed', + ] + assert "MicrovmJobFailedError" in inside + assert len(client.requests("terminate_microvm")) == 1 From 8b762da7f418f659a91a8453d87caf483312a6ce Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Thu, 15 Oct 2026 18:57:12 -0700 Subject: [PATCH 10/10] fix(extras): pass ExecutionError through MicroVM ops NonDeterministicExecutionError is an ExecutionError, and the session's except Exception ran the terminate step for it. The core SDK lets an ExecutionError leave a child context without a checkpoint. So the terminate step took the ID of an operation already in the history, and the execution failed with a second error, or it added an operation to the history of a failed execution. The session and both job waits now pass an ExecutionError through. The MicroVM then runs until its maximum duration, as in JS. - The README states that an SDK error that leaves the session handler reaches the caller as ChildContextError, and that JS differs here. - The path and port docstrings limit the run hook sentences to microvm(), because a session job always goes over HTTP. --- .../README.md | 2 +- .../microvm/config.py | 15 +-- .../microvm/operation.py | 6 ++ .../microvm/session.py | 16 +++- .../tests/microvm/operation_test.py | 28 ++++++ .../tests/microvm/session_test.py | 95 +++++++++++++++++++ 6 files changed, 152 insertions(+), 10 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index 571ec8317..8886d99ec 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -226,7 +226,7 @@ pipeline child context A failed `vm.invoke` raises its MicroVM error inside the handler. The handler can catch it and send another job to the same MicroVM. An error that leaves the handler ends the session: 1. A MicroVM error that the handler raises itself reaches the caller with its own class, such as `MicrovmJobFailedError`. -2. Any other error reaches the caller as the SDK's `ChildContextError`, as from `run_in_child_context`. This includes a MicroVM error that leaves a child context that the handler created: that child context records `ChildContextError` as its error type. The caller then gets the MicroVM error's message, but not its class. Inside the handler, the `ChildContextError` of that child context has the MicroVM error as its `__cause__`. The JavaScript operation behaves the same way. +2. Any other error reaches the caller as the SDK's `ChildContextError`, as from `run_in_child_context`. This includes a MicroVM error that leaves a child context that the handler created: that child context records `ChildContextError` as its error type. The caller then gets the MicroVM error's message, but not its class. Inside the handler, the `ChildContextError` of that child context has the MicroVM error as its `__cause__`. An SDK error that leaves the handler, such as the `StepError` of a failed step, also reaches the caller as `ChildContextError`, with the SDK error as its `__cause__`. This rule matches `run_in_child_context`. The JavaScript operation differs here: it passes an SDK error to the caller with its own class. An invalid job raises `ValidationError` or `TypeError` from `vm.invoke`, before the job's durable operations start. The handler can catch it too. 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 33d873bd2..47cbe0db4 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 @@ -27,7 +27,8 @@ @dataclass(frozen=True, kw_only=True) class MicrovmRequestConfig: - """How :func:`microvm` delivers a job over HTTP. + """How :func:`microvm` and :meth:`MicrovmSession.invoke` deliver a job over + HTTP. An HTTP delivery POSTs the job document to a route of the application in the MicroVM. The route must answer with a 2xx status at once, and run the @@ -50,9 +51,10 @@ class MicrovmRequestConfig: ``None`` sends the job to :data:`DEFAULT_MICROVM_JOB_PATH`. The worker package serves that path with its ``handler``. A path sends the job to one - of your own routes instead, and it selects HTTP delivery for every input - size. The worker runs a job that arrives in the ``run`` hook with its - ``handler``, not with a route. So a job for a route must go over HTTP. + of your own routes instead. For :func:`microvm`, a path also selects HTTP + delivery for every input size. The worker runs a job that arrives in the + ``run`` hook with its ``handler``, not with a route. So a job for a route + must go over HTTP. A session job always goes over HTTP. """ port: int | None = None @@ -61,8 +63,9 @@ class MicrovmRequestConfig: ``None`` uses port 8080, the port that the MicroVM endpoint forwards to by default. The port applies only when the job goes over HTTP. Lambda sends the ``run`` hook to the port in the image's hook configuration, not to - this port. So without a ``path``, the port is used only for an input too - large for the ``run`` hook. + this port. So for :func:`microvm` without a ``path``, the port is used + only for an input too large for the ``run`` hook. A session job always + uses the port. """ retry_window: Duration | None = None 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 da460be7c..116a3dbcc 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 @@ -26,6 +26,7 @@ CallbackError, CallbackTimeoutError, ChildContextError, + ExecutionError, InvocationError, StepError, ValidationError, @@ -720,6 +721,11 @@ def run_job(self, child: DurableContext, job_input: Any) -> Any: try: result = callback.result() + except ExecutionError: + # The execution fails without a retry, and the core SDK records + # no outcome for the child context. A terminate step recorded now + # would change that history. So the error passes through. + raise except InvocationError as error: # An SDK invocation error is not a job outcome. A retryable one # makes Lambda run the invocation again, and the replay waits for 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 3c3fb6137..591048b83 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 @@ -13,6 +13,7 @@ from aws_durable_execution_sdk_python.exceptions import ( ChildContextError, + ExecutionError, InvocationError, ValidationError, ) @@ -300,9 +301,9 @@ def run(job_context: DurableContext) -> Any: ) try: return callback.result() - except InvocationError: - # An SDK invocation error is not a job outcome. A retryable - # one makes Lambda run the invocation again. + except (ExecutionError, InvocationError): + # An SDK execution or invocation error is not a job outcome. + # The core SDK records neither as the child context's failure. raise except Exception as error: # SuspendExecution is a BaseException, so a wait for the @@ -388,6 +389,15 @@ def run_session( raise unknown_delivery_error(self.name, launched) try: result = handler(MicrovmSession(self, launched), child) + except ExecutionError: + # A replay that does not match the history raises + # NonDeterministicExecutionError, and the execution fails without + # a retry. The core SDK lets an ExecutionError leave a child + # context without a checkpoint. A terminate step recorded now + # would add an operation to that history, or take the ID of an + # operation already in it. So the error passes through, and the + # MicroVM runs until its maximum duration, as in JS. + raise except InvocationError as error: # A retryable InvocationError, such as RetryableSerDesError, makes # Lambda run the invocation again. The child context records no diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py index 0ef54173d..6b45d9866 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py @@ -21,6 +21,7 @@ from aws_durable_execution_sdk_python.exceptions import ( ChildContextError, InvocationError, + NonDeterministicExecutionError, ) from aws_durable_execution_sdk_python.execution import ( InvocationStatus, @@ -1086,3 +1087,30 @@ def result_once(self: Any) -> Any: else {"error": "MicrovmError", "message": 'MicroVM "build": interrupted'} ) assert order == (["complete", "terminate"] if retryable else ["terminate"]) + + +def test_passes_an_execution_error_of_the_wait_through( + fake_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # The core SDK records no outcome for an ExecutionError, and the + # execution fails without a retry. A terminate step recorded for it would + # change that history. So the MicroVM runs until its maximum duration. + def result_fails(self: Any) -> Any: + msg = "history does not match" + raise NonDeterministicExecutionError(msg) + + monkeypatch.setattr(Callback, "result", result_fails) + client = fake_client() + runner = DurableFunctionTestRunner( + handler=catching(config_with(client)), poll_interval=0.05 + ) + try: + arn = runner.run_async(input="{}", execution_timeout=WAIT) + result = runner.wait_for_result(arn, timeout=WAIT) + finally: + runner.close() + + assert result.status is InvocationStatus.FAILED + assert result.error is not None + assert "history does not match" in (result.error.message or "") + assert client.requests("terminate_microvm") == [] 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 6b689952f..2e914e265 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 @@ -18,6 +18,7 @@ from aws_durable_execution_sdk_python.exceptions import ( ChildContextError, InvocationError, + NonDeterministicExecutionError, RetryableSerDesError, ValidationError, ) @@ -1053,3 +1054,97 @@ def body(vm: MicrovmSession, ctx: DurableContext) -> Any: ] assert "MicrovmJobFailedError" in inside assert len(client.requests("terminate_microvm")) == 1 + + +def test_records_the_handlers_non_determinism_and_does_not_terminate( + fake_client: Any, fake_endpoint: Any +) -> None: + client = fake_client(endpoint=fake_endpoint().url) + names = ["build"] + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + first = vm.invoke(ctx, names[0], job_input=1, options=job()) + return [first, vm.invoke(ctx, "test", job_input=2, options=job())] + + runner = DurableFunctionTestRunner( + handler=session_handler(session_config(client), body), poll_interval=0.05 + ) + try: + arn = runner.run_async(input="null", execution_timeout=WAIT) + build = runner.wait_for_callback(arn, name="build.callback", timeout=WAIT) + runner.send_callback_success(build, b"1") + test = runner.wait_for_callback(arn, name="test.callback", timeout=WAIT) + # A deploy renames the first job while the execution waits. + names[0] = "compile" + runner.send_callback_success(test, b"2") + result = runner.wait_for_result(arn, timeout=WAIT) + finally: + runner.close() + + assert result.status is InvocationStatus.FAILED + assert result.error is not None + assert "current='compile'" in (result.error.message or "") + assert client.requests("terminate_microvm") == [] + + +def test_wraps_an_sdk_error_that_leaves_the_handler( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + # The JavaScript operation passes an SDK error to the caller with its own + # class. This port follows run_in_child_context instead. + client = fake_client(endpoint=fake_endpoint().url) + config = session_config(client) + + def fail(_: Any) -> Any: + msg = "no artifact" + raise LookupError(msg) + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + def body(vm: MicrovmSession, ctx: DurableContext) -> Any: + return ctx.step( + fail, + name="fetch", + config=StepConfig(retry_strategy=lambda e, a: RetryDecision.no_retry()), + ) + + try: + microvm_session(context, "pipeline", config=config, handler=body) + except ChildContextError as error: + return [type(error.__cause__).__name__, error.message] + return None + + result = run(handler) + + assert outcome(result) == ["StepError", "no artifact"] + assert len(client.requests("terminate_microvm")) == 1 + + +def test_passes_an_execution_error_of_a_job_wait_through( + fake_client: Any, fake_endpoint: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # The core SDK records no outcome for an ExecutionError. So the job does + # not fail with it, and the session does not terminate for it. + def result_fails(self: Any) -> Any: + msg = "history does not match" + raise NonDeterministicExecutionError(msg) + + monkeypatch.setattr(Callback, "result", result_fails) + client = fake_client(endpoint=fake_endpoint().url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke(ctx, "build", job_input=1, options=job()) + + runner = DurableFunctionTestRunner( + handler=session_handler(session_config(client), body), poll_interval=0.05 + ) + try: + arn = runner.run_async(input="null", execution_timeout=WAIT) + result = runner.wait_for_result(arn, timeout=WAIT) + finally: + runner.close() + + assert result.status is InvocationStatus.FAILED + assert result.error is not None + assert "history does not match" in (result.error.message or "") + assert client.requests("terminate_microvm") == []