diff --git a/.github/scripts/type-checks.sh b/.github/scripts/type-checks.sh index 5e6ba236f..00c745d0e 100644 --- a/.github/scripts/type-checks.sh +++ b/.github/scripts/type-checks.sh @@ -21,6 +21,10 @@ mypy --install-types --non-interactive \ packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker \ packages/aws-durable-execution-sdk-python-microvm-worker/tests +mypy --install-types --non-interactive \ + packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras \ + packages/aws-durable-execution-sdk-python-extras/tests + # comment out this for now as there are many type check errors in this package #mypy --install-types --non-interactive \ # packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing \ diff --git a/AGENTS.md b/AGENTS.md index 4052f02d2..6cd88099e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -84,7 +84,9 @@ Do NOT add or modify **conformance tests** without coordinating with the team. C (`aws-durable-execution-sdk-python-testing`), the OpenTelemetry plugin (`aws-durable-execution-sdk-python-otel`), the Workflow Insight plugin (`aws-durable-execution-sdk-python-insight`), the Lambda MicroVM worker - (`aws-durable-execution-sdk-python-microvm-worker`), and examples + (`aws-durable-execution-sdk-python-microvm-worker`), the extras package with + custom operations such as `microvm` + (`aws-durable-execution-sdk-python-extras`), and examples (`aws-durable-execution-sdk-python-examples`). - Read [CONTRIBUTING.md](CONTRIBUTING.md) before making changes. Use `hatch` for all tests, type checks, and formatting (for example diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md new file mode 100644 index 000000000..3846c44b4 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -0,0 +1,113 @@ +# AWS Durable Execution SDK for Python: extras + +> This package is experimental. Its API can change in any release. + +Custom operations for the [AWS Durable Execution SDK for Python](../aws-durable-execution-sdk-python). Each operation composes the SDK's built-in operations and calls AWS APIs. The operations live in this package, so the core SDK does not depend on those APIs' clients. + +The package is the Python counterpart of [`@aws/durable-execution-sdk-js-extras`](https://github.com/aws/aws-durable-execution-sdk-js/tree/main/packages/aws-durable-execution-sdk-js-extras). + +## `microvm`: run a job in an AWS Lambda MicroVM + +`microvm` launches a new [Lambda MicroVM](https://docs.aws.amazon.com/lambda/latest/dg/lambda-microvms-guide.html), delivers the job to it, waits for the job's result, and terminates the MicroVM. The function invocation ends while the MicroVM works, so no Lambda compute is billed during the wait. + +```python +import os + +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 MicrovmConfig, microvm + +CONFIG = MicrovmConfig( + image_identifier=os.environ["MICROVM_IMAGE_ARN"], + execution_role_arn=os.environ["MICROVM_ROLE_ARN"], + timeout=Duration.from_hours(1), + heartbeat_timeout=Duration.from_minutes(5), +) + + +@durable_execution +def handler(event: dict, context: DurableContext) -> dict: + return microvm(context, "build", job_input={"repo": event["repo"]}, config=CONFIG) +``` + +The first argument is the context that the call runs in. Inside a child context or a `map` item, pass that context, so the operation's steps belong to it. `job_input` and `config` are keyword-only. A job can be a string, and so is the name. So positional arguments in the wrong order would pass validation and run the wrong job. + +The code in the MicroVM receives the job and reports its result. The worker package [`aws-durable-execution-sdk-python-microvm-worker`](../aws-durable-execution-sdk-python-microvm-worker) does this, and so does the JavaScript worker. The job document is the same in every SDK language. + +### What it does + +The operation runs its durable operations in one child context named after the second argument: + +1. `.callback` creates a durable callback. The service generates the callback ID, and the SDK checkpoints it. So the ID is unique to this call and stays the same on replay. +2. `.launch` calls `RunMicrovm`. The job goes in the `runHookPayload`, and Lambda passes it to the worker's `run` lifecycle hook. The client token is the SHA-256 hex digest of the callback ID. So a retried or replayed launch returns the same MicroVM instead of starting a second one. The step checkpoints plain JSON, `{"microvmId": "...", "delivery": "run-hook"}`, as the JavaScript operation does. So a reader of the execution history finds the MicroVM ID there. +3. The child context waits for the MicroVM to complete the callback. +4. `.terminate` calls `TerminateMicrovm`. It runs after success, failure, and timeout alike. + +A terminate failure does not fail the operation. The job's outcome is already recorded at that point. So the operation logs an error instead of raising. The MicroVM then keeps running, and is billed, until the platform ends it at its maximum duration: the operation timeout plus 5 minutes, at most 8 hours. + +The terminate step runs only when the operation sees the job's outcome. In these cases it never does, and the MicroVM also runs until its maximum duration: + +- A `map` or `parallel` completes early, for example with `min_successful`, while an item still waits for its MicroVM. The SDK never runs that item again. +- The execution fails, times out, or is stopped while the operation waits. +- Every launch attempt timed out or lost its connection, but the service created the MicroVM anyway. The operation never learned the MicroVM's ID. + +The `run` hook payload holds at most 4,096 Unicode code points. It includes the callback ID and the job's fields, so an input of about 3,600 code points fits. The operation fails a larger job with `MicrovmDeliveryError` before the launch. When the job does not fit even without the callback ID, it fails before the callback exists. Delivery over HTTP, for larger jobs, comes in a later release. + +### Errors + +Every MicroVM failure is a `MicrovmError`. Catch the base class for all of them, or a subclass for one kind: + +| Error | When | +|---|---| +| `MicrovmLaunchError` | `RunMicrovm` failed after all retries. When the last attempt timed out or lost its connection, the service may have created the MicroVM anyway, and it then runs the job. So running the job again after this error can run it twice. | +| `MicrovmDeliveryError` | The job is too large for the `run` hook. | +| `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. | + +```python +from aws_durable_execution_sdk_python_extras.microvm import ( + MicrovmError, + MicrovmTimeoutError, +) + +try: + result = microvm(context, "build", job_input=job_input, config=CONFIG) +except MicrovmTimeoutError: + ... # The MicroVM crashed, or took too long. +except MicrovmError: + ... # Any other MicroVM failure. +``` + +The class is the same on the first run and on every replay, so code that branches on it stays deterministic. A `microvm` call inside your own child context, or inside a `map` or `parallel` item, reaches the code around that context as the SDK's `ChildContextError`. Its `__cause__` is the MicroVM error, so check `isinstance(error.__cause__, MicrovmTimeoutError)` there. An invalid name or config raises the SDK's `ValidationError`, and an input that is not JSON-serializable raises `TypeError`, before any durable operation starts. `Duration` raises the same `ValidationError` for a negative value. + +When the error escapes the handler, the execution's error type is the class's fully-qualified name, such as `aws_durable_execution_sdk_python_extras.microvm.errors.MicrovmJobFailedError`. The SDK uses that name to rebuild the class on replay. The JavaScript operation reports the short name, `MicrovmJobFailedError`. So a rule that matches the error type of both SDKs, such as a Step Functions `Catch` or a metric filter, must match both forms. For example, match on the suffix. + +### Retries + +The launch and terminate steps retry with `default_microvm_retry_strategy`, unless the config sets `retry_strategy`. It retries throttling, server errors, connection failures, and timeouts, for at most 5 attempts, with a delay of 2, 4, 8, and 16 seconds. Every other error, such as `ValidationException`, `AccessDeniedException`, or `InsufficientCapacityException`, fails the step at once. The boto3 client also retries throttling and server errors inside each attempt. + +### Defaults and limits + +- `timeout` must be between 1 second and 8 hours, because a MicroVM lives at most 8 hours. The MicroVM's `maximumDurationInSeconds` is `timeout` plus 5 minutes, at most 8 hours. So a timeout above 7 hours 55 minutes leaves the MicroVM less than 5 minutes beyond the callback's timeout. +- `job_input` is sent as JSON. So the MicroVM receives what `json.dumps` writes: a tuple arrives as a list, and a dict key that is not a string arrives as a string. +- Without connectors in the config, the MicroVM gets the `NO_INGRESS` connector, because the job arrives in the `run` hook, and the `INTERNET_EGRESS` connector, because it needs egress to reach the Lambda callback APIs. An empty tuple sends no connector. +- `heartbeat_timeout` starts when the callback is created, before the launch. So it must cover the launch, with its retries, and the boot. A throttled launch can take several minutes. With a shorter heartbeat timeout, the operation fails with `MicrovmTimeoutError` before the MicroVM sends its first heartbeat. +- The Region comes from the durable execution ARN, or from `AWS_REGION`. +- Without a `client`, the operation creates one `lambda-microvms` client per Region, on first use. It has a 5-second connect timeout, a 30-second read timeout, and 3 attempts in botocore's standard retry mode. So one launch or terminate attempt takes at most about 2 minutes. A call that runs past the function's timeout records no outcome for its step attempt, and the step runs again. The launch is idempotent through its client token, so it is repeated, not doubled. + +### Permissions + +The function's execution role needs: + +- `lambda:RunMicrovm` and `lambda:TerminateMicrovm` on the MicroVM image. Without `lambda:TerminateMicrovm`, every operation still returns its result, but every MicroVM runs until its timeout plus 5 minutes, at most 8 hours. The only sign is an error in the function's log. +- `lambda:PassNetworkConnector` on the connectors that the MicroVM uses. With the defaults, that is `arn:aws:lambda::aws:network-connector:aws-network-connector:*`. +- `iam:PassRole` on the MicroVM execution role. + +The MicroVM execution role (`execution_role_arn`) needs `lambda:SendDurableExecutionCallbackSuccess`, `lambda:SendDurableExecutionCallbackFailure`, and `lambda:SendDurableExecutionCallbackHeartbeat` on the function ARN (`arn:aws:lambda:::function::*`). + +### 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. + +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/pyproject.toml b/packages/aws-durable-execution-sdk-python-extras/pyproject.toml new file mode 100644 index 000000000..2b23096a7 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/pyproject.toml @@ -0,0 +1,79 @@ +[build-system] +requires = ["hatchling!=1.32.1,<1.33"] +build-backend = "hatchling.build" + +[project] +name = "aws-durable-execution-sdk-python-extras" +dynamic = ["version"] +description = 'Custom durable operations for the AWS Durable Execution SDK for Python, such as running a job in an AWS Lambda MicroVM' +readme = "README.md" +requires-python = ">=3.11" +license = "Apache-2.0" +keywords = ["lambda", "microvm", "durable-execution"] +authors = [{ name = "AWS durable-execution-dev", email = "durable-execution-dev@amazon.com" }] +classifiers = [ + "Development Status :: 3 - Alpha", + "Programming Language :: Python", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Programming Language :: Python :: Implementation :: CPython", +] +dependencies = [ + # 2.0.0 is the first release with the graded callback errors + # (CallbackTimeoutError and CallbackExternalError) and register_operation_error. + # The upper bound guards one read of the SDK's internal ExecutionState, for + # the error type that a MicroVM reports. See _reported_error_type. + "aws-durable-execution-sdk-python>=2.0.0,<3", + # botocore 1.43.35 is the first release with the lambda-microvms client. + # The core SDK requires only boto3 1.42.1. + "boto3>=1.43.35", +] + +[project.urls] +Documentation = "https://github.com/aws/aws-durable-execution-sdk-python#readme" +Issues = "https://github.com/aws/aws-durable-execution-sdk-python/issues" +Source = "https://github.com/aws/aws-durable-execution-sdk-python" + +[tool.hatch.build.targets.sdist.force-include] +"../../LICENSE" = "LICENSE" +"../../NOTICE" = "NOTICE" + +[tool.hatch.build.targets.wheel] +packages = ["src/aws_durable_execution_sdk_python_extras"] + +[tool.hatch.build.targets.wheel.force-include] +"../../LICENSE" = "aws_durable_execution_sdk_python_extras/LICENSE" +"../../NOTICE" = "aws_durable_execution_sdk_python_extras/NOTICE" + +[tool.hatch.version] +path = "src/aws_durable_execution_sdk_python_extras/__about__.py" + +[tool.hatch.publish.index] +disable = true + +[tool.coverage.run] +source_pkgs = ["aws_durable_execution_sdk_python_extras"] +branch = true +parallel = true +omit = ["src/aws_durable_execution_sdk_python_extras/__about__.py"] + +[tool.coverage.report] +exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"] + +[tool.ruff] +line-length = 88 +target-version = "py311" + +[tool.ruff.lint] +preview = true +select = ["E4", "E7", "E9", "F", "TID252"] + +[tool.ruff.lint.isort] +known-first-party = ["aws_durable_execution_sdk_python_extras"] +force-single-line = false +lines-after-imports = 2 + +[tool.ruff.lint.per-file-ignores] +"tests/**" = ["ARG001", "ARG002", "ARG005", "S101", "PLR2004", "PLR6301", "SIM117", "TRY301"] diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/__about__.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/__about__.py new file mode 100644 index 000000000..c7c5adad4 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/__about__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +__version__ = "0.0.1" diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/__init__.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/__init__.py new file mode 100644 index 000000000..b2e07a0d3 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Custom durable operations for the AWS Durable Execution SDK for Python. + +This package is experimental. Its API can change in any release. + +Each operation composes the SDK's built-in operations and calls AWS APIs. The +operations live in this package, so the core SDK does not depend on those +APIs' clients. Import each operation from its module, for example +``aws_durable_execution_sdk_python_extras.microvm``. +""" + +from aws_durable_execution_sdk_python_extras.__about__ import __version__ + + +__all__ = ["__version__"] 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 new file mode 100644 index 000000000..579941227 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/__init__.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Run durable jobs in AWS Lambda MicroVMs. + +This module is experimental. Its API can change in any release. + +The exports match the parts of ``@aws/durable-execution-sdk-js-extras/microvm`` +that this package has so far. The session operation and delivery over HTTP +come in later changes. +""" + +from aws_durable_execution_sdk_python_extras.microvm.config import ( + MAX_MICROVM_DURATION_SECONDS, + MicrovmConfig, +) +from aws_durable_execution_sdk_python_extras.microvm.errors import ( + MicrovmDeliveryError, + MicrovmError, + MicrovmJobFailedError, + MicrovmLaunchError, + MicrovmTimeoutError, +) +from aws_durable_execution_sdk_python_extras.microvm.job import ( + MAX_RUN_HOOK_PAYLOAD_LENGTH, +) +from aws_durable_execution_sdk_python_extras.microvm.operation import microvm +from aws_durable_execution_sdk_python_extras.microvm.retry import ( + default_microvm_retry_strategy, +) + + +__all__ = [ + "MAX_MICROVM_DURATION_SECONDS", + "MAX_RUN_HOOK_PAYLOAD_LENGTH", + "MicrovmConfig", + "MicrovmDeliveryError", + "MicrovmError", + "MicrovmJobFailedError", + "MicrovmLaunchError", + "MicrovmTimeoutError", + "default_microvm_retry_strategy", + "microvm", +] diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py new file mode 100644 index 000000000..83f825494 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""The configuration of :func:`microvm`.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.retries import RetryDecision + + +if TYPE_CHECKING: + from mypy_boto3_lambda_microvms import LambdaMicroVMsClient as LambdaMicrovmsClient +else: + LambdaMicrovmsClient = Any + + +MAX_MICROVM_DURATION_SECONDS = 28_800 +"""A MicroVM lives at most 8 hours, so ``RunMicrovm`` accepts +``maximumDurationInSeconds`` of at most 28,800.""" + + +@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. + """ + + image_identifier: str + """The ARN or ID of the MicroVM image to run. + + The image must enable the ``run`` lifecycle hook. The service then also + requires the ``ready`` image hook. The worker package + ``aws-durable-execution-sdk-python-microvm-worker`` answers both. + """ + + execution_role_arn: str + """The IAM role that the MicroVM assumes. + + The MicroVM reports its result through the durable callback APIs. So the + role needs ``lambda:SendDurableExecutionCallbackSuccess``, + ``lambda:SendDurableExecutionCallbackFailure``, and + ``lambda:SendDurableExecutionCallbackHeartbeat`` on the durable function's + 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. So it must cover the launch, with its retries, and the boot. A + throttled launch can take several minutes. A heartbeat timeout that is + shorter fails the operation with ``MicrovmTimeoutError`` before the + MicroVM can send its first heartbeat. + """ + + image_version: str | None = None + """The image version to run. Defaults to the latest active version.""" + + ingress_network_connectors: tuple[str, ...] | None = None + """Ingress network connector ARNs. + + ``None`` gets the Lambda-provided ``NO_INGRESS`` connector, because the + job arrives in the ``run`` hook and needs no inbound connection. An empty + tuple sends no ingress connector. + """ + + egress_network_connectors: tuple[str, ...] | None = None + """Egress network connector ARNs. + + ``None`` gets the Lambda-provided ``INTERNET_EGRESS`` connector. The + MicroVM needs egress to reach the Lambda callback APIs. An empty tuple + sends no egress connector. + """ + + logging: Mapping[str, Any] | None = None + """The ``logging`` member of ``RunMicrovm``. Defaults to the service + default.""" + + retry_strategy: Callable[[Exception, int], RetryDecision] | None = None + """The step retry strategy of the launch and terminate steps. + + Defaults to :func:`default_microvm_retry_strategy`. + """ + + client: LambdaMicrovmsClient | None = None + """The boto3 ``lambda-microvms`` client. + + Defaults to one client per Region, created on first use. The Region comes + from the durable execution ARN. The default client has a 5-second connect + timeout, a 30-second read timeout, and 3 attempts in botocore's standard + retry mode. A client that you pass keeps its own settings. Keep each call + shorter than the function's timeout, because a call that runs past it + records no outcome for the step attempt. + """ 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 new file mode 100644 index 000000000..8480ab52c --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/errors.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""The errors that the MicroVM operations raise. + +Every MicroVM failure is a :class:`MicrovmError`. Catch the base class to +handle every failure the same way, or a subclass to handle one kind. + +The operation raises the same class on the first run and on every replay: + +1. Inside its child context, the operation raises a subclass, such as + :class:`MicrovmLaunchError`. +2. The SDK checkpoints the child context's failure with the class's + fully-qualified name as the error type, and the message, data, and stack + trace. +3. The SDK then raises ``ChildContextError``, on the first run and on replay + alike. It rebuilds the escaping error as ``__cause__`` from its registry + of error types. +4. This module registers every MicroVM class with that registry. So + ``__cause__`` is the MicroVM class, and the operation raises it. + +The registry also applies to a child context that the caller wraps around the +operation, and to a ``map`` or ``parallel`` item. The caller then gets +``ChildContextError``, and its ``__cause__`` is the MicroVM error. + +The fully-qualified name is part of the checkpoint format. A checkpoint +written by one release is replayed by the next. So the classes must stay in +this module, with these names. A moved or renamed class is not found in the +registry on replay, and the caller gets a plain ``MicrovmError`` instead. +""" + +from __future__ import annotations + +from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + DurableOperationError, + register_operation_error, +) + + +class MicrovmError(DurableOperationError): + """The base class of every MicroVM failure. + + ``message`` names the operation and keeps the underlying message. + ``data`` and ``stack_trace`` are those of the failure, when it had any. + """ + + @classmethod + def from_child_context_error( + cls, name: str, error: ChildContextError + ) -> MicrovmError: + """Return the MicroVM error that the operation's child context raised. + + 1. The SDK rebuilds a registered MicroVM error type as its own class, + in ``__cause__``. A copy of that error is returned. + 2. Any other error type becomes a plain :class:`MicrovmError`. It keeps + the original error type in ``error_type``. The operation contains no + caller code, so every failure in it is a MicroVM failure. + + The caller raises the result ``from error``. That sets the result's + ``__cause__`` to ``error``. So the result must not be + ``error.__cause__`` itself. Otherwise each error is the other's + ``__cause__``, and a loop that follows ``__cause__`` never ends. + + Args: + name: The operation name, for the message of a plain MicrovmError. + error: The error that ``run_in_child_context`` raised. + """ + cause = error.__cause__ + if isinstance(cause, MicrovmError): + return type(cause)( + message=cause.message, + error_type=cause.error_type, + data=cause.data, + stack_trace=cause.stack_trace, + ) + return MicrovmError( + message=f'MicroVM "{name}": {error.message}', + error_type=error.error_type, + data=error.data, + stack_trace=error.stack_trace, + ) + + +class MicrovmLaunchError(MicrovmError): + """``RunMicrovm`` failed after all retries. + + When the last attempt got an answer, such as a ``ValidationException``, + no MicroVM runs the job. When the last attempt timed out or lost its + connection, the service may have created the MicroVM anyway. That MicroVM + then runs the job, and the operation cannot terminate it, because it never + learned the MicroVM's ID. The MicroVM ends at its maximum duration. So a + caller that runs the job again after this error can run it twice. + """ + + +class MicrovmDeliveryError(MicrovmError): + """The job could not be delivered to the MicroVM.""" + + +class MicrovmJobFailedError(MicrovmError): + """The job handler in the MicroVM failed. + + The MicroVM reported the failure with + ``SendDurableExecutionCallbackFailure``. The message includes the error + type that the MicroVM reported, such as the job handler's exception class. + + The operation also raises it when the MicroVM sent a result that is not + valid JSON. The JavaScript operation differs here: its SDK fails the + invocation with ``SERDES_FAILED`` before the operation sees an error. + """ + + +class MicrovmTimeoutError(MicrovmError): + """No result arrived within the job's timeout, or no heartbeat arrived + within its heartbeat timeout. + + A MicroVM that crashed or never booted ends this way. The message keeps + the service's reason. + """ + + +for _error_class in ( + MicrovmError, + MicrovmLaunchError, + MicrovmDeliveryError, + MicrovmJobFailedError, + MicrovmTimeoutError, +): + # Each class takes the four keywords that the registry passes: message, + # error_type, data, and stack_trace. + register_operation_error(_error_class) 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 new file mode 100644 index 000000000..93edb8d6a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/job.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""The job document that the operation sends to the MicroVM. + +The worker in the MicroVM reads the same document in every SDK language. The +Python worker in ``aws-durable-execution-sdk-python-microvm-worker`` and the +JavaScript worker in ``@aws/durable-execution-sdk-js-microvm-worker`` parse +it. A change to the document must be made in every language. +""" + +from __future__ import annotations + +import json +import re +from dataclasses import dataclass +from typing import Any + + +PAYLOAD_VERSION = 1 +"""The version of the job document.""" + +MAX_RUN_HOOK_PAYLOAD_LENGTH = 4_096 +"""``RunMicrovm`` accepts a ``runHookPayload`` of at most 4,096 code points. + +The service counts Unicode code points, not UTF-8 bytes. A probe against the +service from the JavaScript SDK accepted 4,096 code points of an emoji, and +rejected 4,097. A Python ``str`` has one item per code point. So ``len()`` +counts what the service counts. +""" + + +@dataclass(frozen=True, kw_only=True) +class RunHookPayload: + """The ``runHookPayload`` document that carries one job in ``RunMicrovm``. + + Wire names: ``version``, ``region``, and ``job``, whose fields are + ``callbackId``, ``heartbeatTimeoutSeconds``, and ``input``. + """ + + region: str + """The Region of the durable function, for the callback API calls.""" + callback_id: str + """The callback that the worker completes with the job's outcome.""" + job_input: Any + """The caller's input. It must be JSON-serializable.""" + heartbeat_timeout_seconds: int | None = None + """The heartbeat timeout, when the caller set one.""" + + def to_json(self) -> str: + """Serialize the payload as compact JSON. + + 1. Non-ASCII text stays as it is, as ``JSON.stringify`` writes it. So + the length counts the same code points as the JavaScript SDK's + payload. + 2. A lone surrogate is not valid in the UTF-8 request body. So each + one is written as a ``\\uXXXX`` escape, as ``JSON.stringify`` + writes it. A Python ``str`` holds an astral character as one code + point, so every surrogate in it is a lone one. + 3. ``allow_nan=False`` rejects ``nan`` and infinite floats, because + the worker rejects the ``NaN`` and ``Infinity`` tokens. + + Raises: + TypeError: When the input is not JSON-serializable. + ValueError: When the input holds ``nan`` or an infinite float. + """ + job: dict[str, Any] = {"callbackId": self.callback_id} + if self.heartbeat_timeout_seconds is not None: + job["heartbeatTimeoutSeconds"] = self.heartbeat_timeout_seconds + job["input"] = self.job_input + document = {"version": PAYLOAD_VERSION, "region": self.region, "job": job} + text = json.dumps( + document, allow_nan=False, ensure_ascii=False, separators=(",", ":") + ) + # A surrogate appears only inside a JSON string. So the escape is + # valid JSON, and it decodes to the same surrogate. + return _SURROGATE.sub(lambda match: f"\\u{ord(match.group()):04x}", text) + + +_SURROGATE = re.compile("[\ud800-\udfff]") + + +def fits_run_hook(payload: str) -> bool: + """Whether ``payload`` fits in the ``runHookPayload`` of ``RunMicrovm``.""" + return len(payload) <= MAX_RUN_HOOK_PAYLOAD_LENGTH 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 new file mode 100644 index 000000000..84781beb9 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/operation.py @@ -0,0 +1,550 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""The ``microvm`` operation: run one job in a new AWS Lambda MicroVM.""" + +from __future__ import annotations + +import hashlib +import json +import os +import threading +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import boto3 +from botocore.config import Config +from aws_durable_execution_sdk_python.config import ( + CallbackConfig, + Duration, + StepConfig, +) +from aws_durable_execution_sdk_python.exceptions import ( + CallbackError, + CallbackTimeoutError, + ChildContextError, + StepError, + ValidationError, +) +from aws_durable_execution_sdk_python.retries import RetryDecision +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from botocore.exceptions import ClientError + +from aws_durable_execution_sdk_python_extras.microvm.config import ( + MAX_MICROVM_DURATION_SECONDS, + LambdaMicrovmsClient, + MicrovmConfig, +) +from aws_durable_execution_sdk_python_extras.microvm.errors import ( + MicrovmDeliveryError, + MicrovmError, + MicrovmJobFailedError, + MicrovmLaunchError, + MicrovmTimeoutError, +) +from aws_durable_execution_sdk_python_extras.microvm.job import ( + MAX_RUN_HOOK_PAYLOAD_LENGTH, + RunHookPayload, + fits_run_hook, +) +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 Callback, DurableContext + + +LAUNCH_MARGIN_SECONDS = 300 +"""The time added to the operation timeout to get the MicroVM's maximum +duration. + +The MicroVM should outlive the job's callback, or the platform terminates it +before the job can report. The margin covers the launch. The maximum duration +is capped at 8 hours. So with a timeout near 8 hours, the margin shrinks, and +at 8 hours only the launch delay remains: the callback's timeout starts before +the launch, and the MicroVM's lifetime starts after it. +""" + +_CONNECT_TIMEOUT_SECONDS = 5 +_READ_TIMEOUT_SECONDS = 30 +_CLIENT_ATTEMPTS = 3 +"""The default client's limits. + +botocore defaults to 60-second connect and read timeouts, with its legacy +retries. A launch or terminate call that runs past the Lambda invocation's +timeout records no outcome for its step attempt. These limits keep one step +attempt to about 2 minutes. The launch is idempotent through its client token, +so an attempt that ran out of time is repeated, not doubled. +""" + +RUN_HOOK_DELIVERY = "run-hook" +"""The delivery in which the job goes in the ``runHookPayload`` of the launch.""" + + +@dataclass(frozen=True, kw_only=True) +class LaunchResult: + """What the launch step checkpoints. + + Wire names: ``microvmId`` and ``delivery``. + """ + + microvm_id: str + delivery: str + """How the job reaches the MicroVM. Replay reads it from the checkpoint, + so every replay follows the delivery that the launch prepared for.""" + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> LaunchResult: + # A checkpoint without a delivery gets an empty one. The operation + # then terminates the MicroVM as for any delivery it does not know. + return cls(microvm_id=data["microvmId"], delivery=data.get("delivery", "")) + + def to_dict(self) -> dict[str, Any]: + return {"microvmId": self.microvm_id, "delivery": self.delivery} + + +def microvm( + context: DurableContext, + name: str, + *, + job_input: Any, + config: MicrovmConfig, +) -> Any: + """Run one job in a new AWS Lambda MicroVM, and return the job's result. + + The operation runs its durable operations in one child context named + ``name``: + + 1. ``.callback`` creates a durable callback. The service generates + the callback ID, and the SDK checkpoints it. So the ID is unique to this + call and stays the same on every replay. + 2. ``.launch`` calls ``RunMicrovm``. The job goes in the + ``runHookPayload``, and Lambda passes it to the worker's ``run`` hook. + The client token is the SHA-256 hex digest of the callback ID. So a + retried or replayed launch returns the same MicroVM instead of starting + a second one. + 3. The child context waits on the callback. The invocation ends during the + wait, so no Lambda compute is billed. The worker in the MicroVM + completes the callback with the job's result or failure. + 4. ``.terminate`` calls ``TerminateMicrovm``. It runs after success, + failure, and timeout alike. + + A terminate failure does not fail the operation. The job's outcome is + already recorded when terminate runs, and the MicroVM's maximum duration + bounds how long it keeps running. So the operation logs an error and + returns the job's outcome. + + Pass the context that the call runs in. Inside a ``map`` item or a child + context, pass that context, so the operation's steps belong to it. + + ``job_input`` and ``config`` are keyword-only. A job can be a string, and + the name is a string. So positional arguments in the wrong order would + pass validation and run the wrong job. + + Args: + context: The context to create the durable operations in. + name: The child context name. It also prefixes the inner operations. + job_input: The job's input. It must be JSON-serializable, and the same + on every replay. It is sent as JSON, so the MicroVM receives what + ``json.dumps`` writes: a tuple arrives as a list, and a dict key + that is not a string arrives as a string. + config: The MicroVM and the timeouts. + + Returns: + The value that the MicroVM passed to + ``SendDurableExecutionCallbackSuccess``, parsed as JSON. + + Raises: + ValidationError: When ``name`` or ``config`` is not valid, or the + Region cannot be determined. No durable operation starts. It is + the SDK's class, which ``Duration`` also raises for a negative + value. + TypeError: When ``job_input`` is not JSON-serializable. No durable + operation starts. + MicrovmLaunchError: When ``RunMicrovm`` fails after all retries. + MicrovmDeliveryError: When the job is too large for the ``run`` hook. + 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 = _OperationScope.create(context, name, job_input, config) + + def run(child: DurableContext) -> Any: + return scope.run_job(child, job_input) + + 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 _OperationScope: + """The validated values that every step of one operation uses.""" + + name: str + config: MicrovmConfig + partition: str + region: str + timeout_seconds: int + heartbeat_timeout_seconds: int | None + 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. + + Raises: + ValidationError: When ``name`` or ``config`` is not valid, or the + Region cannot be determined. + TypeError: When ``job_input`` is not JSON-serializable. + """ + 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) + 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, + retry_strategy=config.retry_strategy or default_microvm_retry_strategy, + ) + + def run_job(self, child: DurableContext, job_input: Any) -> Any: + """Run the operation's durable operations in its child context.""" + name = self.name + # The payload without the callback ID is shorter than the full one. + # So a job that does not fit without the ID fails before the callback + # exists, and no callback is left open until its timeout. + self._require_run_hook(self._payload("", job_input), with_callback_id=False) + callback: Callback = child.create_callback( + name=f"{name}.callback", + config=CallbackConfig( + 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(), + ), + ) + payload = self._payload(callback.callback_id, job_input) + # The full payload contains the callback ID, so its size is known only + # after the callback exists. The ID comes from the checkpoint on + # replay, and the input must not change. So every replay reaches the + # same decision. + self._require_run_hook(payload, with_callback_id=True) + + try: + launched = LaunchResult.from_dict( + child.step( + lambda _: self._launch(callback.callback_id, payload).to_dict(), + name=f"{name}.launch", + config=StepConfig( + retry_strategy=self.retry_strategy, + # Plain JSON, as the JavaScript SDK checkpoints it. So + # a reader of the execution history finds the + # MicroVM ID. Changing the serdes later would break + # the replay of executions in flight. + serdes=JsonSerDes(), + ), + ) + ) + except StepError as error: + # An exception without a message checkpoints none. Its type is + # then the only detail. DurableOperationError always sets a type. + detail = error.message or error.error_type + msg = f'MicroVM "{name}": the launch failed: {detail}' + raise MicrovmLaunchError( + message=msg, data=error.data, stack_trace=error.stack_trace + ) from error + microvm_id = launched.microvm_id + if launched.delivery != RUN_HOOK_DELIVERY: + # 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) + + 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 _payload(self, callback_id: str, job_input: Any) -> str: + return RunHookPayload( + region=self.region, + callback_id=callback_id, + job_input=job_input, + heartbeat_timeout_seconds=self.heartbeat_timeout_seconds, + ).to_json() + + def _require_run_hook(self, payload: str, *, with_callback_id: bool) -> None: + """Raise ``MicrovmDeliveryError`` when ``payload`` is too large.""" + if fits_run_hook(payload): + return + size = f"{len(payload)}" if with_callback_id else f"at least {len(payload)}" + msg = ( + f'MicroVM "{self.name}": the job is {size} code points in the run ' + f"hook payload, and the run hook accepts at most " + f"{MAX_RUN_HOOK_PAYLOAD_LENGTH}. Delivery over HTTP for larger jobs " + "is not available yet." + ) + raise MicrovmDeliveryError(msg) + + def _launch(self, callback_id: str, payload: str) -> LaunchResult: + """Call ``RunMicrovm``, and return what the launch step checkpoints.""" + config = self.config + request: dict[str, Any] = { + "imageIdentifier": config.image_identifier, + "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(), + # An explicit empty tuple stays empty. The service accepts 0 to + # 10 connectors, and only None gets the default. + "ingressNetworkConnectors": list( + (self._connector_arn("NO_INGRESS"),) + if config.ingress_network_connectors is None + else config.ingress_network_connectors + ), + "egressNetworkConnectors": list( + (self._connector_arn("INTERNET_EGRESS"),) + if config.egress_network_connectors is None + else config.egress_network_connectors + ), + "maximumDurationInSeconds": min( + self.timeout_seconds + LAUNCH_MARGIN_SECONDS, + MAX_MICROVM_DURATION_SECONDS, + ), + "runHookPayload": payload, + } + if config.image_version is not None: + request["imageVersion"] = config.image_version + if config.logging is not None: + request["logging"] = dict(config.logging) + response = self._client().run_microvm(**request) + microvm_id = response.get("microvmId") + if not microvm_id: + msg = "RunMicrovm returned no microvmId" + raise RuntimeError(msg) + return LaunchResult(microvm_id=microvm_id, delivery=RUN_HOOK_DELIVERY) + + 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) + except ClientError as error: + # TerminateMicrovm succeeds for a terminated MicroVM. The service + # answers ResourceNotFoundException for a MicroVM whose record it + # has removed. Either way, the MicroVM is gone. + if error.response.get("Error", {}).get("Code") != ( + "ResourceNotFoundException" + ): + 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:" + f"aws-network-connector:{connector}" + ) + + def _client(self) -> LambdaMicrovmsClient: + if self.config.client is not None: + return self.config.client + return _DEFAULT_CLIENTS.get(self.region) + + +class _DefaultClients: + """One ``lambda-microvms`` client per Region, created on first use. + + A boto3 client is safe to use from several threads. Creating one from a + shared session is not. So creation runs under a lock, and each client gets + its own session. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._clients: dict[str, LambdaMicrovmsClient] = {} + + def get(self, region: str) -> LambdaMicrovmsClient: + with self._lock: + client = self._clients.get(region) + if client is None: + client = boto3.session.Session().client( + "lambda-microvms", + region_name=region, + config=Config( + connect_timeout=_CONNECT_TIMEOUT_SECONDS, + read_timeout=_READ_TIMEOUT_SECONDS, + retries={ + "mode": "standard", + # total_max_attempts counts the first attempt too. + # max_attempts would count only the retries. + "total_max_attempts": _CLIENT_ATTEMPTS, + }, + ), + ) + self._clients[region] = client + return client + + +_DEFAULT_CLIENTS = _DefaultClients() + + +def _reported_error_type(callback: Callback) -> str | None: + """The error type that the MicroVM reported for a failed callback. + + ``CallbackExternalError`` carries the message and the data, not the type. + The callback's checkpoint holds the type, and it is the same on the first + run and on replay. ``ExecutionState`` is internal to the SDK. So the + package's dependency has an upper bound on the SDK's major version. + + ``Error`` is the type that a worker reports for a value that is not an + exception. It names nothing, so it is left out. + """ + try: + error = callback.state.get_checkpoint_result(callback.operation_id).error + reported = None if error is None else error.type + except (AttributeError, KeyError, TypeError): + # The internal shape changed in an SDK release. The type is only + # detail in the message, so the operation goes on without it. + return None + if not isinstance(reported, str) or reported in {"", "Error"}: + return None + return reported + + +def _partition_and_region(arn: str) -> tuple[str, str]: + """Read the partition and Region from a durable execution ARN. + + The ARN has the form ``arn::lambda:::...``. + Without that form, the partition is ``aws`` and the Region comes from + ``AWS_REGION``. + """ + parts = arn.split(":") + if len(parts) > 3 and parts[0] == "arn" and parts[1] and parts[3]: + return parts[1], parts[3] + return "aws", os.environ.get("AWS_REGION", "") diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/retry.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/retry.py new file mode 100644 index 000000000..41fc18a05 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/retry.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""The default step retry strategy of the MicroVM operations.""" + +from __future__ import annotations + +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.retries import RetryDecision +from botocore.exceptions import ClientError, HTTPClientError +from botocore.exceptions import ConnectionError as BotocoreConnectionError + + +MAX_ATTEMPTS = 5 +"""Attempts of one step, the first one included.""" + +MAX_DELAY_SECONDS = 60 +"""The longest delay between two attempts.""" + +_RETRYABLE_ERROR_CODES = frozenset( + { + "ThrottlingException", + "TooManyRequestsException", + "InternalServerException", + "ServiceException", + "ServiceUnavailableException", + # Another transition of the same MicroVM is in progress. + "ConflictException", + # The transient codes of botocore's standard retry mode. + "RequestTimeout", + "RequestTimeoutException", + "PriorRequestNotComplete", + } +) + +_NON_RETRYABLE_ERROR_CODES = frozenset( + { + # HTTP 500, but the Region has no capacity for the MicroVM. A retry + # seconds later fails the same way, and the JavaScript SDK does not + # retry it either. The check runs before the status check. + "InsufficientCapacityException", + } +) + +_RETRYABLE_ERROR_TYPES: tuple[type[Exception], ...] = ( + TimeoutError, + ConnectionError, + # botocore's connection and HTTP errors do not subclass the builtin + # ConnectionError. These two bases cover EndpointConnectionError, + # ConnectTimeoutError, ProxyConnectionError, SSLError, ReadTimeoutError, + # ConnectionClosedError, and ResponseStreamingError. botocore's standard + # retry mode treats the same two bases as transient. + BotocoreConnectionError, + HTTPClientError, +) + + +def default_microvm_retry_strategy(error: Exception, attempt: int) -> RetryDecision: + """Retry a step that failed for a reason that a later attempt can fix. + + The launch and terminate steps use this strategy unless the config sets + ``retry_strategy``. It is the second retry tier: the boto3 client already + retried throttling and server errors inside the attempt. + + An error is retryable when it is one of these: + + 1. A service error whose code shows a transient failure, such as + ``ThrottlingException`` or ``InternalServerException``. + 2. A service error with HTTP status 429, or 500 and above, except + ``InsufficientCapacityException``. + 3. A connection failure or a timeout, including botocore's + ``HTTPClientError`` family. + + Every other error fails the step at once. A ``ValidationException``, for + example, fails the same way on every attempt. + + Args: + error: The exception that the step raised. + attempt: The number of the attempt that failed, starting at 1. + + Returns: + A retry after 2 to the power ``attempt`` seconds, at most 60 seconds, + for at most 5 attempts in all. Otherwise no retry. + """ + if attempt >= MAX_ATTEMPTS or not _is_retryable(error): + return RetryDecision.no_retry() + return RetryDecision.retry( + Duration.from_seconds(min(2**attempt, MAX_DELAY_SECONDS)) + ) + + +def _is_retryable(error: Exception) -> bool: + if isinstance(error, _RETRYABLE_ERROR_TYPES): + return True + if not isinstance(error, ClientError): + return False + response = error.response + code = response.get("Error", {}).get("Code") + if code in _RETRYABLE_ERROR_CODES: + return True + if code in _NON_RETRYABLE_ERROR_CODES: + return False + status = response.get("ResponseMetadata", {}).get("HTTPStatusCode") + return isinstance(status, int) and (status == 429 or status >= 500) diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/py.typed b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py b/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py new file mode 100644 index 000000000..b8a62e80f --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Test doubles for the Lambda MicroVMs client.""" + +from __future__ import annotations + +import threading +from collections.abc import Callable +from typing import Any + +import pytest +from botocore.exceptions import ClientError + + +def client_error(code: str, status: int = 400) -> ClientError: + """A botocore ClientError with a code and an HTTP status.""" + response: Any = { + "Error": {"Code": code, "Message": f"{code} message"}, + "ResponseMetadata": {"HTTPStatusCode": status}, + } + return ClientError(response, "RunMicrovm") + + +Answer = BaseException | Callable[[dict[str, Any]], Any] | None + + +class FakeMicrovmsClient: + """Records Lambda MicroVMs API calls, and answers from scripted lists. + + ``launches`` and ``terminates`` hold the answers of ``run_microvm`` and + ``terminate_microvm``, in order. Each answer is an exception to raise, a + function that receives the request, or ``None`` for success. After a list + ends, every call succeeds. + """ + + def __init__( + self, + *, + launches: list[Answer] | None = None, + terminates: list[Answer] | None = None, + ) -> None: + self.calls: list[tuple[str, dict[str, Any]]] = [] + self._answers: dict[str, list[Answer]] = { + "run_microvm": list(launches or []), + "terminate_microvm": list(terminates or []), + } + self._lock = threading.Lock() + + def requests(self, name: str) -> list[dict[str, Any]]: + with self._lock: + return [request for call, request in self.calls if call == name] + + def _answer(self, name: str, request: dict[str, Any]) -> Any: + with self._lock: + self.calls.append((name, request)) + answers = self._answers[name] + answer = answers.pop(0) if answers else None + if isinstance(answer, BaseException): + raise answer + if callable(answer): + return answer(request) + return None + + def run_microvm(self, **request: Any) -> dict[str, Any]: + self._answer("run_microvm", request) + return {"microvmId": "mvm-1", "state": "PENDING"} + + def terminate_microvm(self, **request: Any) -> dict[str, Any]: + self._answer("terminate_microvm", request) + return {} + + +# The fakes are fixtures, not imports. Several packages in this repository +# have a top-level ``tests`` package, so ``from tests.conftest import ...`` +# could resolve to another package's tests. + + +@pytest.fixture(autouse=True) +def aws_region(monkeypatch: pytest.MonkeyPatch) -> None: + # The local runner's execution ARN is not an ARN. So the operation takes + # the Region from AWS_REGION, as it does outside Lambda. + monkeypatch.setenv("AWS_REGION", "us-west-2") + + +@pytest.fixture +def fake_client() -> type[FakeMicrovmsClient]: + return FakeMicrovmsClient + + +@pytest.fixture +def make_error() -> Callable[..., ClientError]: + return client_error 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 new file mode 100644 index 000000000..9324e5bcf --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/errors_test.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Tests for rebuilding MicroVM errors from a child context failure.""" + +from __future__ import annotations + +import pytest +from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + DurableOperationError, +) +from aws_durable_execution_sdk_python.lambda_service import ErrorObject + +from aws_durable_execution_sdk_python_extras.microvm import ( + MicrovmDeliveryError, + MicrovmError, + MicrovmJobFailedError, + MicrovmLaunchError, + MicrovmTimeoutError, +) + + +def child_context_error(raised: Exception) -> ChildContextError: + """The error that the SDK raises for a child context that raised + ``raised``, built the way the SDK builds it, from the checkpoint fields.""" + try: + ErrorObject.from_exception(raised).raise_as_operation_error(ChildContextError) + except ChildContextError as error: + return error + msg = "unreachable" + raise AssertionError(msg) + + +@pytest.mark.parametrize( + "error_class", + [ + MicrovmError, + MicrovmLaunchError, + MicrovmDeliveryError, + MicrovmJobFailedError, + MicrovmTimeoutError, + ], +) +def test_rebuilds_each_microvm_error_class(error_class: type[MicrovmError]) -> None: + raised = error_class( + message='MicroVM "build": broken', data='{"k":1}', stack_trace=["at x"] + ) + + rebuilt = MicrovmError.from_child_context_error( + "build", child_context_error(raised) + ) + + assert type(rebuilt) is error_class + assert rebuilt.message == 'MicroVM "build": broken' + assert rebuilt.data == '{"k":1}' + assert rebuilt.stack_trace == ["at x"] + assert rebuilt.error_type == raised.error_type + + +def test_wraps_another_error_type_in_a_plain_microvm_error() -> None: + rebuilt = MicrovmError.from_child_context_error( + "build", child_context_error(RuntimeError("unexpected")) + ) + + assert type(rebuilt) is MicrovmError + assert rebuilt.message == 'MicroVM "build": unexpected' + assert rebuilt.error_type == "RuntimeError" + + +def test_every_microvm_error_is_a_durable_operation_error() -> None: + assert issubclass(MicrovmError, DurableOperationError) + for error_class in ( + MicrovmLaunchError, + MicrovmDeliveryError, + MicrovmJobFailedError, + MicrovmTimeoutError, + ): + assert issubclass(error_class, MicrovmError) + + +@pytest.mark.parametrize( + "error_class", + [ + MicrovmError, + MicrovmLaunchError, + MicrovmDeliveryError, + MicrovmJobFailedError, + MicrovmTimeoutError, + ], +) +def test_registers_each_class_with_the_sdk(error_class: type[MicrovmError]) -> None: + rebuilt = DurableOperationError.from_error_fields( + error_type=f"{error_class.__module__}.{error_class.__qualname__}", + message="m", + data="d", + stack_trace=["s"], + ) + + assert type(rebuilt) is error_class + assert (rebuilt.message, rebuilt.data, rebuilt.stack_trace) == ("m", "d", ["s"]) 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 new file mode 100644 index 000000000..4f841ce00 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/job_test.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the job document.""" + +from __future__ import annotations + +import json +import math + +import pytest +from aws_durable_execution_sdk_python_microvm_worker import RunHookRequest + +from aws_durable_execution_sdk_python_extras.microvm.job import ( + MAX_RUN_HOOK_PAYLOAD_LENGTH, + RunHookPayload, + fits_run_hook, +) + + +def test_writes_compact_json_with_the_job_fields() -> None: + payload = RunHookPayload( + region="us-east-1", + callback_id="cb-1", + job_input={"repo": "org/app"}, + heartbeat_timeout_seconds=60, + ).to_json() + + assert payload == ( + '{"version":1,"region":"us-east-1","job":{"callbackId":"cb-1",' + '"heartbeatTimeoutSeconds":60,"input":{"repo":"org/app"}}}' + ) + + +def test_leaves_out_an_unset_heartbeat_timeout() -> None: + payload = RunHookPayload(region="r", callback_id="cb", job_input=None).to_json() + + assert json.loads(payload)["job"] == {"callbackId": "cb", "input": None} + + +def test_keeps_non_ascii_text_so_the_length_counts_code_points() -> None: + payload = RunHookPayload(region="r", callback_id="cb", job_input="日😀").to_json() + + assert "日😀" in payload + + +def test_escapes_only_a_lone_surrogate() -> None: + # JSON.stringify escapes a lone surrogate and keeps other text. So the + # same job has the same length in both languages. + payload = RunHookPayload( + region="r", callback_id="cb", job_input="é" * 1000 + "\ud800" + ).to_json() + plain = RunHookPayload(region="r", callback_id="cb", job_input="é" * 1000).to_json() + + payload.encode() + assert "\\ud800" in payload + assert "é" * 1000 in payload + assert len(payload) == len(plain) + len("\\ud800") + assert json.loads(payload)["job"]["input"] == "é" * 1000 + "\ud800" + + +@pytest.mark.parametrize("value", [math.nan, math.inf]) +def test_rejects_a_float_that_is_not_json(value: float) -> None: + with pytest.raises(ValueError, match="JSON compliant"): + RunHookPayload(region="r", callback_id="cb", job_input=value).to_json() + + +def test_counts_code_points_against_the_run_hook_limit() -> None: + assert fits_run_hook("😀" * MAX_RUN_HOOK_PAYLOAD_LENGTH) + assert not fits_run_hook("😀" * (MAX_RUN_HOOK_PAYLOAD_LENGTH + 1)) + + +def test_the_worker_reads_the_payload() -> None: + # The worker package parses the same contract. A change on one side must + # be made on the other. + payload = RunHookPayload( + region="us-east-1", + callback_id="cb-1", + job_input={"n": 1}, + heartbeat_timeout_seconds=30, + ).to_json() + + request = RunHookRequest.from_dict( + {"microvmId": "mvm-1", "runHookPayload": payload} + ) + + assert request.payload is not None + assert request.payload.region == "us-east-1" + job = request.payload.job + assert job is not None + assert (job.callback_id, job.input, job.heartbeat_timeout_seconds) == ( + "cb-1", + {"n": 1}, + 30.0, + ) 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 new file mode 100644 index 000000000..ada06b829 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py @@ -0,0 +1,694 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the microvm operation, run on the local test runner.""" + +from __future__ import annotations + +import datetime +import hashlib +import json +from collections.abc import Callable, Iterator +from types import SimpleNamespace +from typing import Any + +import boto3 +import pytest +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.exceptions import ChildContextError +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.retries import RetryDecision +from aws_durable_execution_sdk_python_testing.runner import ( + DurableFunctionTestResult, + DurableFunctionTestRunner, +) +from botocore.stub import ANY, Stubber + +from aws_durable_execution_sdk_python_extras.microvm import ( + MAX_RUN_HOOK_PAYLOAD_LENGTH, + MicrovmConfig, + MicrovmError, + MicrovmJobFailedError, + microvm, + operation, +) +from aws_durable_execution_sdk_python_extras.microvm.job import RunHookPayload + + +WAIT = 30 + + +def config_with(client: Any, **fields: Any) -> MicrovmConfig: + 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_minutes(5), + "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 MicrovmConfig(**values) + + +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[..., Any]]: + """Run a handler on the local test runner, and close the runner after.""" + runners: list[DurableFunctionTestRunner] = [] + + def _run( + handler: Callable[..., Any], + event: Any, + complete: Callable[[DurableFunctionTestRunner, str], None] | 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) + if complete is not None: + callback_id = runner.wait_for_callback( + arn, name="build.callback", timeout=WAIT + ) + complete(runner, callback_id) + return runner.wait_for_result(arn, timeout=WAIT) + + yield _run + for runner in runners: + runner.close() + + +def succeed_with(value: Any) -> Callable[[DurableFunctionTestRunner, str], None]: + 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 | None, message: str +) -> Callable[[DurableFunctionTestRunner, str], None]: + 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 catching(config: MicrovmConfig) -> Callable[..., Any]: + """A handler that returns the class and message of a MicroVM error.""" + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + try: + return {"result": microvm(context, "build", job_input=event, config=config)} + except MicrovmError as error: + return {"error": type(error).__name__, "message": error.message} + + return handler + + +def test_launches_returns_the_result_and_terminates(run: Any, fake_client: Any) -> None: + client = fake_client() + config = config_with(client, heartbeat_timeout=Duration.from_seconds(60)) + + result = run(catching(config), {"repo": "org/app"}, succeed_with({"passed": True})) + + assert outcome(result) == {"result": {"passed": True}} + [launch] = client.requests("run_microvm") + payload = json.loads(launch["runHookPayload"]) + callback_id = payload["job"]["callbackId"] + assert payload == { + "version": 1, + "region": "us-west-2", + "job": { + "callbackId": callback_id, + "heartbeatTimeoutSeconds": 60, + "input": {"repo": "org/app"}, + }, + } + assert launch["clientToken"] == hashlib.sha256(callback_id.encode()).hexdigest() + assert launch["imageIdentifier"] == config.image_identifier + assert launch["executionRoleArn"] == config.execution_role_arn + assert launch["maximumDurationInSeconds"] == 300 + 300 + connector = "arn:aws:lambda:us-west-2:aws:network-connector:aws-network-connector" + assert launch["ingressNetworkConnectors"] == [f"{connector}:NO_INGRESS"] + assert launch["egressNetworkConnectors"] == [f"{connector}:INTERNET_EGRESS"] + assert "imageVersion" not in launch + assert "logging" not in launch + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +def test_records_the_operations_in_one_child_context( + run: Any, fake_client: Any +) -> None: + result = run(catching(config_with(fake_client())), {}, succeed_with(1)) + + context = result.get_context("build") + assert [operation.name for operation in context.child_operations] == [ + "build.callback", + "build.launch", + "build.terminate", + ] + + +def test_checkpoints_the_launch_as_plain_json(run: Any, fake_client: Any) -> None: + result = run(catching(config_with(fake_client())), {}, succeed_with(1)) + + launch = result.get_context("build").get_step("build.launch") + # The JavaScript operation checkpoints the same document, so a reader of + # the execution history finds the MicroVM ID in both. + assert launch.result is not None + assert json.loads(launch.result) == {"microvmId": "mvm-1", "delivery": "run-hook"} + + +def test_keeps_empty_connector_tuples_empty(run: Any, fake_client: Any) -> None: + client = fake_client() + config = config_with( + client, ingress_network_connectors=(), egress_network_connectors=() + ) + + run(catching(config), {}, succeed_with(None)) + + [launch] = client.requests("run_microvm") + assert launch["ingressNetworkConnectors"] == [] + assert launch["egressNetworkConnectors"] == [] + + +def test_passes_the_optional_launch_settings(run: Any, fake_client: Any) -> None: + client = fake_client() + config = config_with( + client, + image_version="7", + logging={"cloudWatch": {"logGroup": "/microvm"}}, + ingress_network_connectors=("arn:ingress",), + egress_network_connectors=("arn:egress",), + timeout=Duration.from_hours(8), + ) + + run(catching(config), {}, succeed_with(None)) + + [launch] = client.requests("run_microvm") + assert launch["imageVersion"] == "7" + assert launch["logging"] == {"cloudWatch": {"logGroup": "/microvm"}} + assert launch["ingressNetworkConnectors"] == ["arn:ingress"] + assert launch["egressNetworkConnectors"] == ["arn:egress"] + # The maximum duration is capped at 8 hours. + assert launch["maximumDurationInSeconds"] == 28_800 + assert "heartbeatTimeoutSeconds" not in json.loads(launch["runHookPayload"])["job"] + + +def test_fails_with_a_job_failed_error_for_a_result_that_is_not_json( + run: Any, fake_client: Any +) -> None: + client = fake_client() + + def complete(runner: DurableFunctionTestRunner, callback_id: str) -> None: + runner.send_callback_success(callback_id, b"not json") + + result = run(catching(config_with(client)), {}, complete) + + error = outcome(result) + assert error["error"] == "MicrovmJobFailedError" + assert error["message"].startswith( + 'MicroVM "build": the job result could not be read (SerDesError): ' + ) + assert len(client.requests("terminate_microvm")) == 1 + + +def test_uses_the_default_client_of_the_region( + run: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + client = fake_client() + regions: list[str] = [] + + class Clients: + def get(self, region: str) -> Any: + regions.append(region) + return client + + monkeypatch.setattr(operation, "_DEFAULT_CLIENTS", Clients()) + + result = run(catching(config_with(None)), {}, succeed_with("done")) + + assert outcome(result) == {"result": "done"} + assert set(regions) == {"us-west-2"} + assert len(client.requests("run_microvm")) == 1 + + +def test_fails_with_the_microvm_error_and_still_terminates( + run: Any, fake_client: Any +) -> None: + client = fake_client() + + result = run( + catching(config_with(client)), {}, fail_with("RuntimeError", "3 tests failed") + ) + + assert outcome(result) == { + "error": "MicrovmJobFailedError", + "message": 'MicroVM "build": the job failed (RuntimeError): 3 tests failed', + } + assert len(client.requests("terminate_microvm")) == 1 + + +@pytest.mark.parametrize("error_type", [None, "Error"]) +def test_leaves_out_an_error_type_that_names_nothing( + run: Any, fake_client: Any, error_type: str | None +) -> None: + result = run( + catching(config_with(fake_client())), {}, fail_with(error_type, "broken") + ) + + assert outcome(result)["message"] == 'MicroVM "build": the job failed: broken' + + +def test_fails_with_a_timeout_error_when_no_result_arrives( + run: Any, fake_client: Any +) -> None: + client = fake_client() + config = config_with(client, timeout=Duration.from_seconds(2)) + + result = run(catching(config), {}) + + error = outcome(result) + assert error["error"] == "MicrovmTimeoutError" + assert error["message"].startswith('MicroVM "build": the job timed out: ') + assert len(client.requests("terminate_microvm")) == 1 + + +def test_returns_the_result_when_terminate_fails( + run: Any, fake_client: Any, make_error: Any +) -> None: + client = fake_client(terminates=[make_error("AccessDeniedException", 403)]) + + result = run(catching(config_with(client)), {}, succeed_with("done")) + + assert outcome(result) == {"result": "done"} + assert len(client.requests("terminate_microvm")) == 1 + terminate = result.get_context("build").get_step("build.terminate") + assert terminate.error is not None + + +def test_treats_resource_not_found_from_terminate_as_terminated( + run: Any, fake_client: Any, make_error: Any +) -> None: + client = fake_client(terminates=[make_error("ResourceNotFoundException", 404)]) + + result = run(catching(config_with(client)), {}, succeed_with("done")) + + assert outcome(result) == {"result": "done"} + terminate = result.get_context("build").get_step("build.terminate") + assert terminate.error is None + + +def test_retries_a_throttled_launch_with_the_same_client_token( + run: Any, fake_client: Any, make_error: Any +) -> None: + client = fake_client(launches=[make_error("ThrottlingException", 400)]) + + result = run(catching(config_with(client)), {}, succeed_with("done")) + + assert outcome(result) == {"result": "done"} + first, second = client.requests("run_microvm") + assert first["clientToken"] == second["clientToken"] + assert first["runHookPayload"] == second["runHookPayload"] + + +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", 400)]) + + result = run(catching(config_with(client)), {}) + + error = outcome(result) + assert error["error"] == "MicrovmLaunchError" + assert error["message"].startswith('MicroVM "build": the launch failed: ') + assert "ValidationException" in error["message"] + assert client.requests("terminate_microvm") == [] + + +def test_fails_a_launch_without_a_microvm_id(run: Any, fake_client: Any) -> None: + class NoId(fake_client): # type: ignore[misc, valid-type] + def run_microvm(self, **request: Any) -> dict[str, Any]: + super().run_microvm(**request) + return {} + + result = run(catching(config_with(NoId())), {}) + + error = outcome(result) + assert error["error"] == "MicrovmLaunchError" + assert "RunMicrovm returned no microvmId" in error["message"] + + +def test_fails_a_job_too_large_for_the_run_hook_before_the_launch( + run: Any, fake_client: Any +) -> None: + client = fake_client() + + result = run(catching(config_with(client)), "x" * MAX_RUN_HOOK_PAYLOAD_LENGTH) + + error = outcome(result) + assert error["error"] == "MicrovmDeliveryError" + assert "the job is at least " in error["message"] + assert "the run hook accepts at most 4096" in error["message"] + assert client.calls == [] + # The job does not fit even without the callback ID. So no callback is + # left open until its timeout. + assert result.get_context("build").child_operations == [] + + +def test_fails_a_job_that_fits_only_without_the_callback_id( + run: Any, fake_client: Any +) -> None: + client = fake_client() + empty = RunHookPayload( + region="us-west-2", + callback_id="", + job_input="", + heartbeat_timeout_seconds=None, + ).to_json() + job_input = "x" * (MAX_RUN_HOOK_PAYLOAD_LENGTH - len(empty)) + + result = run(catching(config_with(client)), job_input) + + error = outcome(result) + assert error["error"] == "MicrovmDeliveryError" + assert "the job is at least " not in error["message"] + assert "the run hook accepts at most 4096" in error["message"] + assert client.calls == [] + [callback] = result.get_context("build").child_operations + assert callback.name == "build.callback" + + +def test_raises_the_same_class_after_the_handler_stops_catching( + run: Any, fake_client: Any +) -> None: + config = config_with(fake_client()) + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + return microvm(context, "build", job_input=event, config=config) + + result = run(handler, {}, fail_with("RuntimeError", "broken")) + + assert result.status is InvocationStatus.FAILED + assert result.error is not None + assert result.error.type == ( + "aws_durable_execution_sdk_python_extras.microvm.errors.MicrovmJobFailedError" + ) + + +def test_runs_in_the_context_that_the_caller_passes(run: Any, fake_client: Any) -> None: + config = config_with(fake_client()) + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + return context.run_in_child_context( + lambda child: microvm(child, "build", job_input=event, config=config), + name="outer", + ) + + result = run(handler, {}, succeed_with(7)) + + assert outcome(result) == 7 + outer = result.get_context("outer") + assert [operation.name for operation in outer.child_operations] == ["build"] + + +def test_runs_one_microvm_per_map_item(run: Any, fake_client: Any) -> None: + client = fake_client() + config = config_with(client) + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + return context.map( + [1, 2, 3], + lambda item_context, item, index, items: microvm( + item_context, f"shard-{item}", job_input=item, config=config + ), + name="shards", + ).get_results() + + runner = DurableFunctionTestRunner(handler=handler, poll_interval=0.05) + try: + arn = runner.run_async(input="{}", execution_timeout=WAIT) + for item in [1, 2, 3]: + callback_id = runner.wait_for_callback( + arn, name=f"shard-{item}.callback", timeout=WAIT + ) + runner.send_callback_success(callback_id, json.dumps(item * 10).encode()) + result = runner.wait_for_result(arn, timeout=WAIT) + finally: + runner.close() + + assert outcome(result) == [10, 20, 30] + tokens = {request["clientToken"] for request in client.requests("run_microvm")} + assert len(tokens) == 3 + assert len(client.requests("terminate_microvm")) == 3 + + +def test_requests_match_the_lambda_microvms_api_model(run: Any) -> None: + """botocore validates each request against the service model.""" + + client = boto3.session.Session().client("lambda-microvms", region_name="us-west-2") + config = config_with( + client, + image_version="7", + heartbeat_timeout=Duration.from_seconds(60), + logging={"cloudWatch": {"logGroup": "/microvm"}}, + ) + with Stubber(client) as stub: + stub.add_response( + "run_microvm", + { + "microvmId": "mvm-1", + "state": "PENDING", + "endpoint": "mvm-1.example", + "imageArn": config.image_identifier, + "imageVersion": "7", + "maximumDurationInSeconds": 600, + "startedAt": datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC), + }, + { + "imageIdentifier": config.image_identifier, + "imageVersion": "7", + "executionRoleArn": config.execution_role_arn, + "clientToken": ANY, + "ingressNetworkConnectors": ANY, + "egressNetworkConnectors": ANY, + "maximumDurationInSeconds": 600, + "logging": {"cloudWatch": {"logGroup": "/microvm"}}, + "runHookPayload": ANY, + }, + ) + stub.add_response("terminate_microvm", {}, {"microvmIdentifier": "mvm-1"}) + + result = run(catching(config), {"n": 1}, succeed_with("done")) + + stub.assert_no_pending_responses() + client.close() + assert outcome(result) == {"result": "done"} + + +def test_keeps_the_class_inside_the_callers_child_context( + run: Any, fake_client: Any +) -> None: + # The MicroVM classes are in the SDK's error registry. So a caller's own + # child context rebuilds the MicroVM error as its __cause__, on the first + # run and on replay alike. + + config = config_with(fake_client()) + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + try: + context.run_in_child_context( + lambda child: microvm(child, "build", job_input=event, config=config), + name="outer", + ) + except ChildContextError as error: + cause = error.__cause__ + seen = { + "cause": type(cause).__name__, + "is_job_failed": isinstance(cause, MicrovmJobFailedError), + } + first = context.step(lambda _: seen, name="first") + # The wait makes the SDK replay the handler, and the outer context + # fails again from its checkpoint. + context.wait(Duration.from_seconds(1), name="replay") + return {"first": first, "replay": seen} + + result = run(handler, {}, fail_with("RuntimeError", "broken")) + + expected = {"cause": "MicrovmJobFailedError", "is_job_failed": True} + assert outcome(result) == {"first": expected, "replay": expected} + + +def test_keeps_the_stack_trace_that_the_microvm_reports( + run: Any, fake_client: Any +) -> None: + config = config_with(fake_client()) + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + try: + microvm(context, "build", job_input=event, config=config) + except MicrovmError as error: + return error.stack_trace + return None + + def complete(runner: DurableFunctionTestRunner, callback_id: str) -> None: + runner.send_callback_failure( + callback_id, + ErrorObject( + message="broken", type="Error", data=None, stack_trace=["at job"] + ), + ) + + result = run(handler, {}, complete) + + assert outcome(result) == ["at job"] + + +def test_terminates_and_fails_for_a_delivery_that_it_does_not_know( + run: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # A launch checkpoint from another release, for example after a rollback. + + original = operation._OperationScope._launch # noqa: SLF001 + + def launch_for_http(self: Any, callback_id: str, payload: str) -> Any: + launched = original(self, callback_id, payload) + return operation.LaunchResult(microvm_id=launched.microvm_id, delivery="http") + + monkeypatch.setattr(operation._OperationScope, "_launch", launch_for_http) # noqa: SLF001 + client = fake_client() + + result = run(catching(config_with(client)), {}) + + assert outcome(result) == { + "error": "MicrovmDeliveryError", + "message": 'MicroVM "build": the launch prepared the delivery "http", ' + "which this release does not support", + } + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +def test_terminates_and_fails_for_a_launch_checkpoint_without_a_delivery( + run: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + original = operation.LaunchResult.to_dict + + def without_delivery(self: Any) -> dict[str, Any]: + data = original(self) + del data["delivery"] + return data + + monkeypatch.setattr(operation.LaunchResult, "to_dict", without_delivery) + client = fake_client() + + result = run(catching(config_with(client)), {}) + + assert outcome(result)["error"] == "MicrovmDeliveryError" + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +def test_leaves_out_the_error_type_when_the_checkpoint_shape_changed() -> None: + # The checkpoint state is internal to the SDK, and can change in a 2.x + # release. + + callback: Any = SimpleNamespace(state=object(), operation_id="op") + + assert operation._reported_error_type(callback) is None # noqa: SLF001 + + +def test_names_the_type_of_a_launch_error_without_a_message( + run: Any, fake_client: Any +) -> None: + class Silent(fake_client): # type: ignore[misc, valid-type] + def run_microvm(self, **request: Any) -> dict[str, Any]: + raise RuntimeError + + result = run(catching(config_with(Silent())), {}) + + assert outcome(result) == { + "error": "MicrovmLaunchError", + "message": 'MicroVM "build": the launch failed: RuntimeError', + } + + +def test_ends_the_cause_chain_of_the_raised_error(run: Any, fake_client: Any) -> None: + config = config_with(fake_client()) + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + try: + microvm(context, "build", job_input=event, config=config) + except MicrovmError as error: + chain: list[str] = [] + seen: set[int] = set() + current: BaseException | None = error + # A loop that follows __cause__ must reach None. The seen set + # stops this test on a cycle instead of hanging. + while current is not None and id(current) not in seen: + seen.add(id(current)) + chain.append(type(current).__name__) + current = current.__cause__ + return {"chain": chain, "ends": current is None} + return None + + result = run(handler, {}, fail_with("RuntimeError", "broken")) + + assert outcome(result) == { + "chain": [ + "MicrovmJobFailedError", + "ChildContextError", + "MicrovmJobFailedError", + ], + "ends": True, + } + + +def test_terminates_only_after_the_wait_for_the_result_ends( + run: Any, fake_client: Any +) -> None: + # The first invocation suspends on the callback. A terminate in a + # finally around callback.result() would run in that invocation, while + # the MicroVM still runs the job. + invocations: list[int] = [] + terminated_in: list[int] = [] + client = fake_client( + terminates=[lambda request: terminated_in.append(len(invocations))] + ) + config = config_with(client, timeout=Duration.from_seconds(2)) + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + invocations.append(1) + try: + return microvm(context, "build", job_input=event, config=config) + except MicrovmError as error: + return type(error).__name__ + + result = run(handler, {}) + + assert outcome(result) == "MicrovmTimeoutError" + assert terminated_in == [2] diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/retry_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/retry_test.py new file mode 100644 index 000000000..2cd569bda --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/retry_test.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the default MicroVM retry strategy.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from botocore.exceptions import ( + ConnectionClosedError, + ConnectTimeoutError, + EndpointConnectionError, + HTTPClientError, + ProxyConnectionError, + ReadTimeoutError, + ResponseStreamingError, + SSLError, +) + +from aws_durable_execution_sdk_python_extras.microvm import ( + default_microvm_retry_strategy, + retry, +) + + +@pytest.mark.parametrize( + ("code", "status"), + [ + ("ThrottlingException", 400), + ("TooManyRequestsException", 400), + ("InternalServerException", 500), + ("ServiceException", 500), + ("ServiceUnavailableException", 503), + ("ConflictException", 409), + ("SomethingNew", 503), + ("SomethingNew", 429), + ("RequestTimeout", 408), + ("RequestTimeoutException", 408), + ("PriorRequestNotComplete", 400), + ], +) +def test_retries_a_transient_service_error( + make_error: Any, code: str, status: int +) -> None: + decision = default_microvm_retry_strategy(make_error(code, status), 1) + + assert decision.should_retry + assert decision.delay_seconds == 2 + + +@pytest.mark.parametrize( + "error", + [ + TimeoutError("slow"), + ConnectionResetError("reset"), + EndpointConnectionError(endpoint_url="https://example"), + ConnectTimeoutError(endpoint_url="https://example"), + ReadTimeoutError(endpoint_url="https://example"), + ConnectionClosedError(endpoint_url="https://example"), + ProxyConnectionError(proxy_url="https://proxy", error="refused"), + SSLError(endpoint_url="https://example", error=OSError("handshake")), + ResponseStreamingError(error="truncated"), + HTTPClientError(error="reset"), + ], +) +def test_retries_a_connection_failure_or_a_timeout(error: Exception) -> None: + assert default_microvm_retry_strategy(error, 1).should_retry + + +@pytest.mark.parametrize( + ("code", "status"), + [ + ("ValidationException", 400), + ("AccessDeniedException", 403), + ("ResourceNotFoundException", 404), + # The service answers it with HTTP 500, so the code check must run + # before the status check. + ("InsufficientCapacityException", 500), + ], +) +def test_does_not_retry_a_permanent_service_error( + make_error: Any, code: str, status: int +) -> None: + assert not default_microvm_retry_strategy(make_error(code, status), 1).should_retry + + +def test_does_not_retry_another_error() -> None: + assert not default_microvm_retry_strategy(RuntimeError("bug"), 1).should_retry + + +def test_stops_after_5_attempts_and_caps_the_delay_at_60_seconds( + make_error: Any, +) -> None: + error = make_error("ThrottlingException", 400) + delays = [ + default_microvm_retry_strategy(error, attempt).delay_seconds + for attempt in range(1, 5) + ] + + assert delays == [2, 4, 8, 16] + assert not default_microvm_retry_strategy(error, 5).should_retry + + +def test_caps_the_delay_at_60_seconds(make_error: Any, monkeypatch: Any) -> None: + monkeypatch.setattr(retry, "MAX_ATTEMPTS", 10) + + decision = default_microvm_retry_strategy(make_error("ThrottlingException"), 7) + + assert decision.delay_seconds == 60 diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/validation_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/validation_test.py new file mode 100644 index 000000000..bd337aae3 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/validation_test.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the checks that run before any durable operation.""" + +from __future__ import annotations + +import math +from typing import Any + +import pytest +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.exceptions import ValidationError + +from aws_durable_execution_sdk_python_extras.microvm import MicrovmConfig, microvm +from aws_durable_execution_sdk_python_extras.microvm import operation + + +class NoOperations: + """A context that fails the test when the operation starts any durable + operation.""" + + def __init__(self, arn: str = "arn:aws:lambda:eu-west-1:1:function:f") -> None: + self.execution_context = type("ExecutionContext", (), {})() + self.execution_context.durable_execution_arn = arn + + def run_in_child_context(self, *args: Any, **kwargs: Any) -> Any: + msg = "no durable operation may start" + raise AssertionError(msg) + + +def config(**fields: Any) -> MicrovmConfig: + values: dict[str, Any] = { + "image_identifier": "img", + "execution_role_arn": "role", + "timeout": Duration.from_minutes(5), + **fields, + } + return MicrovmConfig(**values) + + +@pytest.mark.parametrize( + ("name", "fields", "message"), + [ + ("", {}, "requires a non-empty name"), + ("build", {"image_identifier": ""}, "image_identifier is required"), + ("build", {"execution_role_arn": ""}, "execution_role_arn is required"), + ("build", {"timeout": Duration(0)}, "timeout must be between 1 second"), + ( + "build", + {"timeout": Duration.from_seconds(28_801)}, + "Got 28801 seconds", + ), + ( + "build", + {"heartbeat_timeout": Duration(0)}, + "heartbeat_timeout must be at least 1 second", + ), + ], +) +def test_rejects_an_invalid_call_before_any_durable_operation( + name: str, fields: dict[str, Any], message: str +) -> None: + with pytest.raises(ValidationError, match=message): + microvm(NoOperations(), name, job_input={}, config=config(**fields)) # type: ignore[arg-type] + + +def test_raises_the_class_of_duration_for_every_invalid_timeout() -> None: + # Duration rejects a negative value itself. The operation rejects zero. + # Both raise the SDK's ValidationError. + with pytest.raises(ValidationError): + config(timeout=Duration.from_seconds(-1)) + with pytest.raises(ValidationError, match="timeout must be between"): + microvm( + NoOperations(), # type: ignore[arg-type] + "build", + job_input={}, + config=config(timeout=Duration(0)), + ) + + +@pytest.mark.parametrize("job_input", [{1, 2}, math.nan, {"x": math.inf}]) +def test_rejects_an_input_that_is_not_json(job_input: Any) -> None: + with pytest.raises(TypeError, match="must be JSON-serializable"): + microvm(NoOperations(), "build", job_input=job_input, config=config()) # type: ignore[arg-type] + + +def test_rejects_a_call_without_a_region(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AWS_REGION") + + with pytest.raises(ValidationError, match="cannot determine the AWS Region"): + microvm(NoOperations("not-an-arn"), "build", job_input={}, config=config()) # type: ignore[arg-type] + + +@pytest.mark.parametrize( + ("arn", "expected"), + [ + ("arn:aws:lambda:eu-west-1:1:function:f:$LATEST/x", ("aws", "eu-west-1")), + ("arn:aws-cn:lambda:cn-north-1:1:function:f", ("aws-cn", "cn-north-1")), + ("not-an-arn", ("aws", "us-west-2")), + ("arn:aws:lambda::1:function:f", ("aws", "us-west-2")), + ], +) +def test_reads_the_partition_and_region(arn: str, expected: tuple[str, str]) -> None: + assert operation._partition_and_region(arn) == expected # noqa: SLF001 + + +def test_creates_one_default_client_per_region( + monkeypatch: pytest.MonkeyPatch, +) -> None: + created: list[tuple[str, str]] = [] + configs: list[Any] = [] + + class Session: + def client(self, service: str, region_name: str, config: Any) -> object: + created.append((service, region_name)) + configs.append(config) + return object() + + monkeypatch.setattr(operation.boto3.session, "Session", Session) + clients = operation._DefaultClients() # noqa: SLF001 + + first = clients.get("us-east-1") + + assert clients.get("us-east-1") is first + assert clients.get("eu-west-1") is not first + assert created == [ + ("lambda-microvms", "us-east-1"), + ("lambda-microvms", "eu-west-1"), + ] + # botocore's defaults are 60 seconds each, with legacy retries. + assert (configs[0].connect_timeout, configs[0].read_timeout) == (5, 30) + assert configs[0].retries == {"mode": "standard", "total_max_attempts": 3} + + +def test_takes_the_input_and_the_config_by_keyword_only() -> None: + with pytest.raises(TypeError, match="positional argument"): + microvm(NoOperations(), "build", {}, config()) # type: ignore[arg-type, misc] 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 new file mode 100644 index 000000000..bc78bcfd9 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""The microvm operation and the real MicroVM worker, end to end, in one process. + +1. The durable function runs on the local test runner. +2. The fake ``RunMicrovm`` sends the ``run`` hook to a real worker on a free + port, as Lambda does when the MicroVM starts. +3. The worker runs the job, and reports through a Lambda client double that + forwards each callback call to the runner. + +So the test covers the job document and the callback contract between the two +packages. +""" + +from __future__ import annotations + +import http.client +import json +from collections.abc import Callable, Iterator +from typing import Any + +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.execution import ( + InvocationStatus, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python_microvm_worker import ( + HOOK_PATH_PREFIX, + MicrovmJobContext, + MicrovmWorker, + start_microvm_worker, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner + +from aws_durable_execution_sdk_python_extras.microvm import ( + MicrovmConfig, + MicrovmError, + microvm, +) + + +WAIT = 30 + + +class RunnerCallbacks: + """A Lambda client for the worker that sends each callback call to the + local test runner.""" + + def __init__(self, runner: DurableFunctionTestRunner) -> None: + self._runner = runner + self.heartbeats = 0 + + def send_durable_execution_callback_success( + self, CallbackId: str, Result: bytes | None = None + ) -> dict[str, Any]: + self._runner.send_callback_success(CallbackId, Result) + return {} + + def send_durable_execution_callback_failure( + self, CallbackId: str, Error: dict[str, Any] + ) -> dict[str, Any]: + self._runner.send_callback_failure( + CallbackId, + ErrorObject( + message=Error.get("ErrorMessage"), + type=Error.get("ErrorType"), + data=Error.get("ErrorData"), + stack_trace=None, + ), + ) + return {} + + def send_durable_execution_callback_heartbeat( + self, CallbackId: str + ) -> dict[str, Any]: + self.heartbeats += 1 + self._runner.send_callback_heartbeat(CallbackId) + return {} + + +def run_hook_to(worker: MicrovmWorker) -> Callable[[dict[str, Any]], None]: + """A RunMicrovm answer that sends the run hook to ``worker``.""" + + def deliver(request: dict[str, Any]) -> None: + connection = http.client.HTTPConnection("127.0.0.1", worker.port, timeout=WAIT) + try: + body = {"microvmId": "mvm-1", "runHookPayload": request["runHookPayload"]} + connection.request("POST", f"{HOOK_PATH_PREFIX}run", body=json.dumps(body)) + assert connection.getresponse().status == 200 + finally: + connection.close() + + return deliver + + +@pytest.fixture +def composed(fake_client: Any) -> Iterator[Callable[..., Any]]: + """Run a job through the operation and a real worker.""" + closers: list[Callable[[], None]] = [] + + def _run( + job: Callable[[Any, MicrovmJobContext], Any], + job_input: Any, + **fields: Any, + ) -> tuple[Any, Any, RunnerCallbacks]: + holder: dict[str, MicrovmConfig] = {} + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + try: + return { + "result": microvm( + context, "build", job_input=event, config=holder["config"] + ) + } + except MicrovmError as error: + return {"error": type(error).__name__, "message": error.message} + + runner = DurableFunctionTestRunner(handler=handler, poll_interval=0.05) + closers.append(runner.close) + callbacks = RunnerCallbacks(runner) + worker = start_microvm_worker( + handler=job, + # The double has the three callback methods that the worker calls. + create_client=lambda region: callbacks, # type: ignore[arg-type, return-value] + heartbeat_interval_seconds=fields.pop("heartbeat_interval", None), + port=0, + ) + closers.append(worker.close) + client = fake_client(launches=[run_hook_to(worker)]) + holder["config"] = MicrovmConfig( + image_identifier="img", + execution_role_arn="role", + timeout=Duration.from_minutes(5), + client=client, + **fields, + ) + arn = runner.run_async(input=json.dumps(job_input), execution_timeout=WAIT) + result = runner.wait_for_result(arn, timeout=WAIT) + assert worker.idle(WAIT) + assert result.status is InvocationStatus.SUCCEEDED, result.error + return json.loads(result.result or "null"), client, callbacks + + yield _run + for close in reversed(closers): + close() + + +def test_returns_the_result_of_the_job_in_the_worker(composed: Any) -> None: + def job(job_input: Any, context: MicrovmJobContext) -> Any: + return {"sum": sum(job_input["numbers"]), "microvm": context.microvm_id} + + outcome, client, _ = composed(job, {"numbers": [1, 2, 3]}) + + assert outcome == {"result": {"sum": 6, "microvm": "mvm-1"}} + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +def test_fails_with_the_error_type_that_the_worker_reports(composed: Any) -> None: + def job(job_input: Any, context: MicrovmJobContext) -> Any: + msg = "3 tests failed" + raise RuntimeError(msg) + + outcome, client, _ = composed(job, {}) + + assert outcome == { + "error": "MicrovmJobFailedError", + "message": 'MicroVM "build": the job failed (RuntimeError): 3 tests failed', + } + assert len(client.requests("terminate_microvm")) == 1 + + +def test_sends_heartbeats_when_the_job_has_a_heartbeat_timeout( + composed: Any, +) -> None: + def job(job_input: Any, context: MicrovmJobContext) -> Any: + context.cancelled.wait(0.3) + return "done" + + outcome, _, callbacks = composed( + job, + {}, + heartbeat_timeout=Duration.from_seconds(30), + heartbeat_interval=0.05, + ) + + assert outcome == {"result": "done"} + assert callbacks.heartbeats >= 2 diff --git a/pyproject.toml b/pyproject.toml index 9c2732b29..365a7d25b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ examples-integration = "pytest --runner-mode=cloud -m example packages/aws-durab addopts = "-v --strict-markers --import-mode=importlib" testpaths = [ "packages/aws-durable-execution-sdk-python/tests", + "packages/aws-durable-execution-sdk-python-extras/tests", "packages/aws-durable-execution-sdk-python-insight/tests", "packages/aws-durable-execution-sdk-python-microvm-worker/tests", "packages/aws-durable-execution-sdk-python-otel/tests", @@ -51,7 +52,7 @@ workspace.members = ["packages/*"] extra-dependencies = [ "mypy>=1.0.0", "pytest", - "boto3-stubs[lambda]", + "boto3-stubs[lambda,lambda-microvms]", "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", ] @@ -98,6 +99,26 @@ test = "pytest packages/aws-durable-execution-sdk-python-otel/tests {args}" cov = "pytest --cov-report=term-missing --cov-config=packages/aws-durable-execution-sdk-python-otel/pyproject.toml --cov=packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel packages/aws-durable-execution-sdk-python-otel/tests {args}" typecheck = "mypy --install-types --non-interactive packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel packages/aws-durable-execution-sdk-python-otel/tests" +[tool.hatch.envs.dev-extras] +workspace.members = [ + "packages/aws-durable-execution-sdk-python", + "packages/aws-durable-execution-sdk-python-extras", + "packages/aws-durable-execution-sdk-python-microvm-worker", + "packages/aws-durable-execution-sdk-python-testing", +] +dependencies = [ + "pytest", + "pytest-cov", + "coverage[toml]", + "mypy>=1.0.0", + "boto3-stubs[lambda,lambda-microvms]", +] + +[tool.hatch.envs.dev-extras.scripts] +test = "pytest packages/aws-durable-execution-sdk-python-extras/tests {args}" +cov = "pytest --cov-report=term-missing --cov-config=packages/aws-durable-execution-sdk-python-extras/pyproject.toml --cov=packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras packages/aws-durable-execution-sdk-python-extras/tests {args}" +typecheck = "mypy --install-types --non-interactive packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras packages/aws-durable-execution-sdk-python-extras/tests" + [tool.hatch.envs.dev-microvm-worker] workspace.members = ["packages/aws-durable-execution-sdk-python-microvm-worker"] dependencies = [ @@ -233,6 +254,7 @@ select = ["E4", "E7", "E9", "F", "TID252"] # pycodestyle (E4/E7/E9) + Pyflakes [tool.ruff.lint.isort] known-first-party = [ "aws_durable_execution_sdk_python", + "aws_durable_execution_sdk_python_extras", "aws_durable_execution_sdk_python_insight", "aws_durable_execution_sdk_python_microvm_worker", "aws_durable_execution_sdk_python_otel",