diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index 2622f84a..8886d99e 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,112 @@ 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 +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, + 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. 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 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__`. 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. + +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 32572cbd..1c0c0463 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 38e24377..47cbe0db 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 @@ -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 @@ -76,12 +79,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. @@ -101,32 +101,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.""" @@ -138,11 +112,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 @@ -157,13 +132,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. @@ -202,3 +170,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 8480ab52..2773cbf9 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 454711cc..9cc7979d 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 00000000..539b5175 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py @@ -0,0 +1,184 @@ +# 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, + ) + + 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": + raise not_found() 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. 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}) + + 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 ceb6334a..116a3dbc 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 @@ -25,6 +26,8 @@ CallbackError, CallbackTimeoutError, ChildContextError, + ExecutionError, + InvocationError, StepError, ValidationError, ) @@ -35,6 +38,7 @@ from aws_durable_execution_sdk_python_extras.microvm.config import ( MAX_MICROVM_DURATION_SECONDS, LambdaMicrovmsClient, + MicrovmBaseConfig, MicrovmConfig, MicrovmRequestConfig, ) @@ -43,6 +47,7 @@ MicrovmError, MicrovmJobFailedError, MicrovmLaunchError, + MicrovmNotRunningError, MicrovmTimeoutError, ) from aws_durable_execution_sdk_python_extras.microvm.job import ( @@ -51,12 +56,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 ( @@ -238,136 +249,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 @@ -379,59 +390,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( @@ -439,13 +429,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 @@ -457,32 +448,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 @@ -491,7 +549,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 @@ -516,6 +574,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: @@ -531,36 +591,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) @@ -573,39 +603,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:" @@ -618,6 +615,195 @@ 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 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 + # 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. + 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 d2345fa7..c2b98a19 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 @@ -214,6 +214,50 @@ 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 {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." + ) + raise MicrovmEndpointUnavailableError(msg) def send_job(options: SendJobOptions) -> int: @@ -284,23 +328,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 {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." - ) - 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" @@ -323,7 +356,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 @@ -402,6 +436,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 00000000..591048b8 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py @@ -0,0 +1,464 @@ +# 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, field +from typing import TYPE_CHECKING, Any, TypeVar + +from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + ExecutionError, + InvocationError, + 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`, 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``. 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 + 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 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 + 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) + # 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' + ) + 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 (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 + # 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 + 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: + """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, + outer_context=context, + ) + + 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 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 + # 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 + # 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 50536a74..92f1fb02 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 9324e5bc..1ad3f575 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 1bc2df30..54853af2 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 00000000..3f2eed38 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py @@ -0,0 +1,224 @@ +# 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() + + +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 61e353e8..6b45d986 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,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.exceptions import ChildContextError +from aws_durable_execution_sdk_python.context import Callback, DurableContext +from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + InvocationError, + NonDeterministicExecutionError, +) from aws_durable_execution_sdk_python.execution import ( InvocationStatus, durable_execution, @@ -545,7 +550,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 +558,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 +794,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 +802,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)), {}) @@ -1028,3 +1033,84 @@ 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"]) + + +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/request_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py index 4af60f6d..638d651d 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 @@ -476,6 +476,84 @@ def interrupted() -> None: 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, + ) + ) + + def test_sends_no_auth_token_on_a_redirect_that_the_callers_opener_follows( fake_client: Any, fake_endpoint: Any ) -> None: diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py new file mode 100644 index 00000000..2e914e26 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py @@ -0,0 +1,1150 @@ +# 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, StepConfig +from aws_durable_execution_sdk_python.context import Callback, DurableContext +from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + InvocationError, + NonDeterministicExecutionError, + RetryableSerDesError, + 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.serdes import SerDes +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"}] + + +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 + + +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 + + +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") == [] 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 f3e9f1da..d8641507 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"}]