From 0d647621c1eb65a056001bc841debf0bc37aa48d Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 25 Oct 2026 16:24:04 -0700 Subject: [PATCH 01/12] fix(microvm-worker): suspend with one attempt A SuspendMicrovm retry can run after the MicroVM resumed and freeze the next job. The default client now makes one attempt, and the before-send handler starts no attempt after a resume hook. --- .../README.md | 8 +- .../self_suspend.py | 96 +++++++--- .../worker.py | 64 +++++-- .../tests/lifecycle_hooks_test.py | 129 ++++++++++++- .../tests/self_suspend_test.py | 178 ++++++++++++++++-- 5 files changed, 415 insertions(+), 60 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md index 851063f90..8aa5f00e4 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -106,19 +106,19 @@ A session's run hook payload can set `autoSuspendIdleSeconds`. The session sets 1. The idle time starts when the `run` hook arrives without a job, when a job ends, and when the `resume` lifecycle hook arrives, if the image enables it. 2. A job that starts cancels the idle time. So the worker never suspends the MicroVM during a job. -3. When the idle time ends, the worker calls `SuspendMicrovm` with the MicroVM ID from the `run` hook, never one from a job request. The call returned in under 100 milliseconds in testing, before the MicroVM froze. No attempt of the call starts after 10 seconds, because a `before-send` handler on the client stops botocore's retries at that deadline. botocore cannot cancel an attempt in flight, and it opens the connection after `before-send`. So an attempt that started before the deadline can still reach the service after it. The worker waits for that attempt to end, and refuses jobs until then. The client's connect and read timeouts bound the attempt: 5 and 10 seconds for the default client. A client from `create_microvms_client` should use similar timeouts. botocore's own defaults are 60 seconds each, so with them the refusal can last minutes. The durable function resumes the MicroVM before the next job, and the process continues with its memory and files. +3. When the idle time ends, the worker calls `SuspendMicrovm` with the MicroVM ID from the `run` hook, never one from a job request. The call returned in under 100 milliseconds in testing, before the MicroVM froze. The default client makes one attempt. An attempt can suspend the MicroVM and lose its answer in the freeze. A retry would then run after the MicroVM resumed. The `resume` hook has ended the refusal of jobs by then. So the retry could suspend the MicroVM with the next job inside. No attempt of the call starts after 10 seconds, because a `before-send` handler on the client stops botocore's retries at that deadline. The same handler starts no attempt after a `resume` hook. This covers a client from `create_microvms_client` that retries. botocore cannot cancel an attempt in flight, and it opens the connection after `before-send`. So an attempt that started before the deadline can still reach the service after it. The worker waits for that attempt to end, and refuses jobs until then. The client's connect and read timeouts bound the attempt: 5 and 10 seconds for the default client. A client from `create_microvms_client` should use similar timeouts. botocore's own defaults are 60 seconds each, so with them the refusal can last minutes. The durable function resumes the MicroVM before the next job, and the process continues with its memory and files. 4. From the moment the idle time ends, the worker answers every job with HTTP 503 and does not start it. A job accepted then would freeze with the MicroVM, and nothing would resume it. After a 503, the durable function checks the state, resumes the MicroVM, and delivers the job again. 5. The refusal ends when the `resume` lifecycle hook arrives, or 30 seconds after the call returned, if no hook arrives. The idle time then starts again. At most one call runs at a time. A `resume` hook that arrives during the call ends the refusal at once, and the next idle time starts when the call ends. So enable the `resume` hook in the image. Without it, the first job after a resume can wait a little over 30 seconds: the 30 seconds, plus up to 4 seconds until the durable function's next delivery attempt. A test measured 32 seconds. 6. A call that the service rejected on its first attempt, such as `AccessDeniedException` for a missing permission, is logged as a warning, and the MicroVM keeps running. Missing credentials count as a rejection. A later call in the same MicroVM most likely finds no credentials either. The worker accepts jobs again at once and does not retry the call. The next job that ends starts the idle time again. The exception is a `resume` hook that arrived during the call. The idle time then starts when the call ends. So the worker retries the call one idle time after the call ends, if no job runs and no `suspend` hook refusal is in effect when the call ends. A second exception is a `suspend` hook that arrived during the call and whose refusal is still in effect when the call ends. The MicroVM is then suspending anyway. So the rejection does not end the hook's refusal. That refusal ends on the `resume` hook, or 30 seconds after the last `suspend` hook. The idle time then starts again. -7. A call that the service refused for now on its first attempt is logged as a warning too, and the worker accepts jobs again at once. This covers expired credentials and a signature error, with or without a clock offset, whatever the HTTP status. The service answers these before it acts. botocore retries none of these errors, so they reach the worker on the first attempt. The idle time then starts again, so the worker tries the call again. That idle time lasts at least 30 seconds, also when `autoSuspendIdleSeconds` is shorter. The same error can repeat on every call, for example with a wrong secret key or an expired session token from environment variables. So while no job ends and no lifecycle hook arrives, the worker makes at most one call per 30 seconds, and logs one warning per call. A job that ends after the refusal, or a `resume` hook after it, starts the normal idle time instead. A `resume` hook or a job that ends during the call does not shorten the idle time after the refusal, because the refused outcome arrives after it. A `suspend` hook whose refusal is still in effect when the call ends is the exception, as in item 6. Its refusal stays, and the idle time starts when it ends. The default client retries a throttle, so a throttle from it usually ends as an unknown outcome, as in item 8. A throttle counts as refused for now only from a client that does not retry it, such as a client from `create_microvms_client` with one attempt. The JavaScript worker classifies these errors the same way. Its default client retries a clock-skew error. So there, a clock-skew error usually ends as an unknown outcome. -8. A call whose outcome is unknown is handled like a successful call, because the service may have suspended the MicroVM anyway. This covers a timeout, a conflict, a server error, and a dropped connection. An error after botocore retried an earlier attempt also counts as unknown, because the earlier attempt may have suspended the MicroVM. The default client uses botocore's standard retry mode with 3 attempts, and that mode retries a throttle. So a throttle usually ends here. The refusal continues until the `resume` hook or the 30-second limit. The idle time then starts again, so the worker tries the call again. A `resume` hook during the call, or the end of a `suspend` hook's 30 seconds during the call, has already ended the refusal. The worker then accepts jobs at once, and tries the call again one idle time after the call and every job end. +7. A call that the service refused for now on its first attempt is logged as a warning too, and the worker accepts jobs again at once. This covers expired credentials and a signature error, with or without a clock offset, whatever the HTTP status. The service answers these before it acts. botocore retries none of these errors, so they reach the worker on the first attempt. The idle time then starts again, so the worker tries the call again. That idle time lasts at least 30 seconds, also when `autoSuspendIdleSeconds` is shorter. The same error can repeat on every call, for example with a wrong secret key or an expired session token from environment variables. So while no job ends and no lifecycle hook arrives, the worker makes at most one call per 30 seconds, and logs one warning per call. A job that ends after the refusal, or a `resume` hook after it, starts the normal idle time instead. A `resume` hook or a job that ends during the call does not shorten the idle time after the refusal, because the refused outcome arrives after it. A `suspend` hook whose refusal is still in effect when the call ends is the exception, as in item 6. Its refusal stays, and the idle time starts when it ends. The default client makes one attempt. So a throttle from it counts as refused for now. A client from `create_microvms_client` that retries a throttle turns it into an unknown outcome, as in item 8. The JavaScript worker classifies these errors the same way. +8. A call whose outcome is unknown is handled like a successful call, because the service may have suspended the MicroVM anyway. This covers a timeout, a conflict, a server error, and a dropped connection. An error after botocore retried an earlier attempt also counts as unknown, because the earlier attempt may have suspended the MicroVM. Only a client from `create_microvms_client` that retries can raise such an error. A call that the `before-send` handler stopped after a `resume` hook counts as unknown too. That `resume` hook has already ended the refusal. So the worker accepts jobs at once, and starts no new 30-second limit for that call. The refusal continues until the `resume` hook or the 30-second limit. The idle time then starts again, so the worker tries the call again. A `resume` hook during the call, or the end of a `suspend` hook's 30 seconds during the call, has already ended the refusal. The worker then accepts jobs at once, and tries the call again one idle time after the call and every job end. 9. The worker logs one warning for each failed call. The message names the outcome. The fields are, in order, `microvmId`, `error`, `outcome`, `suspendHookArrived`, `terminating`, `closed`, and `refusingJobs`. `microvmId` is the MicroVM ID of the call. `error` holds the name and the message of the error. `outcome` is `rejected`, `refusedForNow`, or `unknown`. `suspendHookArrived` states whether a `suspend` hook arrived since the idle time ended. `terminating` states whether a `terminate` hook came before the worker acted on the outcome. `closed` states whether `worker.close()` came before the worker acted on the outcome. `refusingJobs` states whether the worker answers jobs with HTTP 503 after it has acted on the outcome. `worker.close()` stops the idle time. A closed worker never suspends its MicroVM again. A refusal that is already running still ends on the `resume` hook or after 30 seconds. The worker counts only jobs. It cannot see background threads or processes, or inbound traffic that is not a job. A session with such work leaves `auto_suspend_on_idle` off, and the payload then has no `autoSuspendIdleSeconds`. -The worker creates a `lambda-microvms` client for each suspend, and closes it after the call. Pass `create_microvms_client` to provide the client. The worker never closes a client that it returns. A client that retries must report its retries in `ResponseMetadata.RetryAttempts` of the error that it raises, as botocore does. The worker otherwise counts a retried call as a first attempt. An earlier attempt of that call may have suspended the MicroVM. So a retried throttle would count as refused for now, and the worker would accept a job that then freezes with the MicroVM. +The worker creates a `lambda-microvms` client for each suspend, and closes it after the call. Pass `create_microvms_client` to provide the client. The worker never closes a client that it returns. Give the client one attempt, as the default client has. The worker starts no attempt after a `resume` hook. A retry that starts before the `resume` hook arrives can still suspend the MicroVM after it resumed. A client that retries must report its retries in `ResponseMetadata.RetryAttempts` of the error that it raises, as botocore does. The worker otherwise counts a retried call as a first attempt. An earlier attempt of that call may have suspended the MicroVM. So a retried throttle would count as refused for now, and the worker would accept a job that then freezes with the MicroVM. ## Suspend and terminate hooks diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py index 21ecaf6b0..80cf36090 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py @@ -1,12 +1,22 @@ # SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. # # SPDX-License-Identifier: Apache-2.0 -"""The worker's own SuspendMicrovm call, for a session that suspends when idle.""" +"""The worker's own SuspendMicrovm call, for a session that suspends when idle. + +The call can suspend the MicroVM and lose its answer in the freeze. A retry +of the call then runs after the MicroVM resumed. The ``resume`` hook has +already ended the worker's refusal of jobs at that point. So the retry would +suspend the MicroVM with the next job inside it. So the worker stops the +call's attempts after a ``resume`` hook, and its default client makes one +attempt. After a failed call, the worker can call again after a later idle +time, with a new refusal of jobs. +""" from __future__ import annotations import time import uuid +from collections.abc import Callable from enum import Enum from typing import TYPE_CHECKING, Any @@ -53,10 +63,13 @@ _CONNECT_TIMEOUT_SECONDS = 5 -_CLIENT_ATTEMPTS = 3 +_CLIENT_ATTEMPTS = 1 """The attempts of the default client, the first attempt included. -The JavaScript worker's default client makes 3 attempts too. +An attempt can suspend the MicroVM and lose its answer in the freeze. A +retry then runs after the MicroVM resumed, and can suspend it with a job +inside. So the default client makes one attempt. After a failed call, the +worker can call again after a later idle time, with a new refusal of jobs. """ _BEFORE_SEND = "before-send.lambda-microvms.SuspendMicrovm" @@ -77,12 +90,10 @@ class SuspendOutcome(Enum): """The service did nothing on the call's first and only attempt, and a later call can succeed. - Examples are expired credentials and a signature error, with or without - a clock offset. botocore retries none of these, so they keep this - outcome. The default client retries a throttle. So a throttle usually - ends as :attr:`UNKNOWN`. A throttle has this outcome only when it - reaches the worker without a retry, for example from a caller's client - that does not retry it.""" + Examples are expired credentials, a signature error with or without a + clock offset, and a throttle. The default client makes one attempt, so + these errors keep this outcome. A caller's client that retries a + throttle turns it into :attr:`UNKNOWN`.""" UNKNOWN = "unknown" """The service may have suspended the MicroVM.""" @@ -95,17 +106,30 @@ class SuspendDeadlineError(TimeoutError): """ +class SuspendResumedError(Exception): + """A ``resume`` hook arrived before an attempt of the SuspendMicrovm call + could start. + + The ``resume`` hook ends the worker's refusal of jobs. So a job can run + in the MicroVM from then on. The attempt would suspend the MicroVM with + that job inside. So the call stops before the attempt. botocore does not + retry this error. + """ + + def create_microvms_client(region: str) -> LambdaMicrovmsClient: """Create the default ``lambda-microvms`` client for one call. Each call gets its own session, because boto3's default session is not safe to use from several threads at once. - The client uses botocore's standard retry mode with 3 attempts. That - mode retries a throttle error code, such as ``ThrottlingException``, a - 5xx answer, and a connection error. It retries no credential or - signature error, also with a clock offset. Each ``ClientError`` reports the - retries in ``ResponseMetadata.RetryAttempts``. + The client makes one attempt. An attempt can suspend the MicroVM and + lose its answer in the freeze. A retry would then run after the MicroVM + resumed, and could suspend it with the next job inside. The worker calls + again after a later idle time instead, with a new refusal of jobs. + :func:`suspend_microvm` also stops the attempts of a caller's client + after the MicroVM resumed. One attempt is still the recommended setting + for such a client. """ return boto3.session.Session().client( "lambda-microvms", @@ -123,10 +147,13 @@ def create_microvms_client(region: str) -> LambdaMicrovmsClient: def suspend_microvm( - client: LambdaMicrovmsClient, microvm_id: str, deadline_seconds: float + client: LambdaMicrovmsClient, + microvm_id: str, + deadline_seconds: float, + resumed: Callable[[], bool] | None = None, ) -> None: - """Call SuspendMicrovm in this thread, and start no attempt after - ``deadline_seconds``. + """Call SuspendMicrovm in this thread. Start no attempt after + ``deadline_seconds``, and none after ``resumed()`` returns ``True``. botocore has no deadline for a whole call. Its read timeout limits one socket read, and it covers neither credential resolution nor its own @@ -136,6 +163,18 @@ def suspend_microvm( botocore went on retrying, and an attempt started 22.5 seconds after the first one. + The same handler stops an attempt after the MicroVM resumed: + + 1. An attempt can suspend the MicroVM and lose its answer in the freeze. + 2. botocore then retries the attempt after the MicroVM resumed. + 3. The ``resume`` hook has ended the worker's refusal of jobs by then. + 4. So the retry would suspend the MicroVM with a new job inside. + + So the handler raises :class:`SuspendResumedError` when ``resumed()`` + returns ``True``. botocore does not retry that error either. The handler + runs in this thread. So ``resumed`` must not block, and must not take a + lock that the caller of this function can hold. + botocore emits ``before-send`` before it opens the connection, and it cannot cancel an attempt in flight. So an attempt that started before the deadline can still reach the service after it. This function @@ -149,6 +188,8 @@ def suspend_microvm( Raises: SuspendDeadlineError: When an attempt would start after the deadline. + SuspendResumedError: When an attempt would start after ``resumed()`` + returned ``True``. Exception: The call's own error. """ deadline = time.monotonic() + deadline_seconds @@ -157,7 +198,10 @@ def suspend_microvm( unique_id = f"aws-durable-execution-microvm-worker-deadline-{uuid.uuid4().hex}" events = safe_get(lambda: client.meta.events) - def stop_after_deadline(**_: object) -> None: + def stop_attempt(**_: object) -> None: + if resumed is not None and resumed(): + msg = "the MicroVM resumed before this SuspendMicrovm attempt" + raise SuspendResumedError(msg) if time.monotonic() >= deadline: msg = ( f"the SuspendMicrovm deadline of {deadline_seconds} seconds has passed" @@ -165,7 +209,7 @@ def stop_after_deadline(**_: object) -> None: raise SuspendDeadlineError(msg) if events is not None: - events.register_first(_BEFORE_SEND, stop_after_deadline, unique_id=unique_id) + events.register_first(_BEFORE_SEND, stop_attempt, unique_id=unique_id) try: client.suspend_microvm(microvmIdentifier=microvm_id) finally: @@ -179,7 +223,8 @@ def classify_suspend_error(error: BaseException) -> SuspendOutcome: 1. A botocore retry inside the call means that an earlier attempt may have suspended the MicroVM before its answer was lost. So any error after a retry is :attr:`SuspendOutcome.UNKNOWN`. The default client - retries a throttle. So a throttle from it usually ends here. + makes one attempt. So only a caller's client that retries reaches + this rule. 2. Missing credentials are :attr:`SuspendOutcome.REJECTED`, as in the JavaScript worker. A later call in the same MicroVM most likely finds no credentials either. @@ -191,13 +236,16 @@ def classify_suspend_error(error: BaseException) -> SuspendOutcome: status. The service answers them before it acts, and a later call can succeed. Some services answer a throttle with HTTP 400. botocore retries no credential or signature error, also with a clock offset. - So these errors reach this rule. A throttle reaches it only from a - client that did not retry it. + So these errors reach this rule. A throttle reaches it from the + default client, and from any client that did not retry it. 6. Any other 4xx answer except 408, 409, and 429 is :attr:`SuspendOutcome.REJECTED`. A 429 is :attr:`SuspendOutcome.REFUSED_FOR_NOW`. - 7. Any other error, such as a timeout, a dropped connection, a 408, or a - server error, is :attr:`SuspendOutcome.UNKNOWN`. + 7. Any other error is :attr:`SuspendOutcome.UNKNOWN`. Examples are a + timeout, a dropped connection, a 408, a server error, + :class:`SuspendDeadlineError`, and :class:`SuspendResumedError`. The + last two stop a later attempt, and an earlier attempt may have + suspended the MicroVM. """ if isinstance(error, NoCredentialsError): return SuspendOutcome.REJECTED diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py index 3e891e580..62fcc864b 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py @@ -330,14 +330,15 @@ class MicrovmWorkerListener: arrives after it. A ``suspend`` hook whose refusal is still in effect when the call ends is the exception, as in item 6. Its refusal stays, and the idle time starts when it ends. - The default client retries a throttle. So a throttle from it usually - ends as an unknown outcome, as in item 8. A throttle is refused for - now only from a caller's client that does not retry it. + The default client makes one attempt. So a throttle from it is + refused for now. A caller's client that retries a throttle turns it + into an unknown outcome, as in item 8. 8. A call whose outcome is unknown, such as a timeout, a conflict, a server error, or an error after a retried attempt, is handled like a successful call. The service may have suspended the MicroVM anyway. - No attempt of the call starts after 10 seconds. An attempt already in - flight runs until it ends, and the refusal lasts at least until then. + No attempt of the call starts after 10 seconds, and none starts + after a ``resume`` hook. An attempt already in flight runs until it + ends, and the refusal lasts at least until then. A `resume` hook during the call, or the end of a `suspend` hook's 30 seconds during the call, has already ended the refusal. The worker then accepts jobs at once, and tries the call again one idle time @@ -454,6 +455,13 @@ def __init__( # worker's own that the service rejected must not end a refusal that # a suspend hook started meanwhile. self._suspend_hooks = 0 + # Counts the resume hooks from Lambda. Each one ends the refusal, so + # a job can run from then on. A SuspendMicrovm attempt after it would + # freeze that job. So the call's before-send handler starts no + # attempt once this count has changed. The handler reads the count + # without the lock: an int read is atomic, and the handler then + # waits for no listener lock inside botocore. + self._resume_hooks = 0 # SuspendMicrovm calls in flight. idle() waits for them. At most one # runs at a time. self._suspend_calls = 0 @@ -511,7 +519,11 @@ def create( idle, once per suspend. It never closes a client that this function returns. Defaults to a client of its own per suspend, from the default credential chain. The listener - closes that client after the call. A client that retries + closes that client after the call. The listener stops the + client's attempts once a ``resume`` hook arrived. Give the + client one attempt, as the default client has. A retry + that starts before the ``resume`` hook arrives can still + suspend the MicroVM after it resumed. A client that retries must report its retries in ``ResponseMetadata.RetryAttempts`` of the error that it raises. The listener otherwise counts a retried call as a first attempt. An earlier attempt of that @@ -708,9 +720,11 @@ def _handle_resume_hook(self) -> list[tuple[MicrovmJobDocument, str]]: A resume without a job would otherwise leave the MicroVM running until the next job ends. Returns the deferred run hook jobs, - registered. The caller starts them after the answer. + registered. The caller starts them after the answer. A + SuspendMicrovm call in flight starts no attempt after this hook. """ with self._locked(): + self._resume_hooks += 1 if self._suspending: return self._end_suspending() self._start_idle_time() @@ -1438,10 +1452,13 @@ def _idle_time_ended(self, generation: int, waited_seconds: float) -> None: # thread starts, while the client is created, or during the call. # Each path that ends the refusal compares with this count. hooks_before = self._suspend_hooks + # A resume hook from here on ends the refusal. The call then + # starts no further attempt. + resumes_before = self._resume_hooks try: threading.Thread( target=self._suspend_self, - args=(auto_suspend, hooks_before, waited_seconds), + args=(auto_suspend, hooks_before, resumes_before, waited_seconds), name="microvm-suspend", daemon=True, ).start() @@ -1458,10 +1475,16 @@ def _idle_time_ended(self, generation: int, waited_seconds: float) -> None: self._launch_deferred(deferred) def _suspend_self( - self, auto_suspend: _AutoSuspend, hooks_before: int, waited_seconds: float + self, + auto_suspend: _AutoSuspend, + hooks_before: int, + resumes_before: int, + waited_seconds: float, ) -> None: try: - self._call_suspend(auto_suspend, hooks_before, waited_seconds) + self._call_suspend( + auto_suspend, hooks_before, resumes_before, waited_seconds + ) finally: with self._locked(): self._suspend_call_ended() @@ -1490,7 +1513,11 @@ def _suspend_call_ended(self) -> None: self._changed.notify_all() def _call_suspend( - self, auto_suspend: _AutoSuspend, hooks_before: int, waited_seconds: float + self, + auto_suspend: _AutoSuspend, + hooks_before: int, + resumes_before: int, + waited_seconds: float, ) -> None: target = auto_suspend.microvm_id fields = {"microvmId": target} @@ -1530,7 +1557,15 @@ def _call_suspend( # close(). unknown_described: object | None = None try: - suspend_microvm(client, target, SUSPEND_CALL_TIMEOUT_SECONDS) + # The before-send handler calls resumed() inside botocore, in + # this thread. It reads one int without the lock. So it runs no + # code of the caller, and waits for no listener lock. + suspend_microvm( + client, + target, + SUSPEND_CALL_TIMEOUT_SECONDS, + resumed=lambda: self._resume_hooks != resumes_before, + ) except Exception as error: # noqa: BLE001 outcome = classify_suspend_error(error) # describe() calls str(error). A client from @@ -1544,7 +1579,10 @@ def _call_suspend( # A timeout, a conflict, a server error, or an error after a # retried attempt does not show whether the service suspended # the MicroVM. So the refusal continues, as after a successful - # call. The block after the client is closed acts on it. + # call. The block after the client is closed acts on it. A + # resume hook that stopped the call's next attempt has + # already ended the refusal. So that block starts a grace + # period only for a suspend hook that came after the resume. unknown_described = described else: deferred = [] diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py index 01463abbf..458c64fd4 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py @@ -20,8 +20,12 @@ from dataclasses import dataclass, field from typing import Any +import boto3 import pytest -from botocore.exceptions import ClientError, NoCredentialsError +from botocore.awsrequest import AWSResponse, HTTPHeaders +from botocore.config import Config +from botocore.exceptions import ClientError, ConnectionClosedError, NoCredentialsError +from botocore.retries.standard import ExponentialBackoff from aws_durable_execution_sdk_python_microvm_worker import ( HOOK_PATH_PREFIX, @@ -915,6 +919,129 @@ def test_accepts_jobs_at_once_when_the_resume_hook_arrives_during_the_call( assert eventually(lambda: sorted(h.jobs.started()) == ["cb-1", "cb-2"]) +class LostAnswerService: + """Stands in for the Lambda MicroVMs service in a botocore client. + + A ``before-send`` handler answers each attempt instead of HTTP. The + first attempt "suspends" the MicroVM. Its answer is lost: the attempt + waits until :meth:`thaw`, and then fails with a closed connection, as + after a freeze. Each later attempt also "suspends" the MicroVM, answers + 200, and records the jobs that it would freeze. + """ + + def __init__(self, listener: Callable[[], MicrovmWorkerListener]) -> None: + self._listener = listener + self._thawed = threading.Event() + self._lock = threading.Lock() + self.attempts = 0 + self.frozen_jobs: list[str] = [] + + def install(self, client: Any) -> Any: + client.meta.events.register( + "before-send.lambda-microvms.SuspendMicrovm", self._send + ) + return client + + def thaw(self) -> None: + self._thawed.set() + + def _send(self, **_: Any) -> AWSResponse: + with self._lock: + self.attempts += 1 + first = self.attempts == 1 + if first: + self._thawed.wait(WAIT) + raise ConnectionClosedError(endpoint_url="https://example.com") + with self._lock: + self.frozen_jobs.extend(self._listener()._jobs) # noqa: SLF001 + headers = HTTPHeaders() + headers["Content-Type"] = "application/json" + return AWSResponse("https://example.com", 200, headers, _Body(b"{}")) + + +class _Body: + def __init__(self, data: bytes) -> None: + self._data = data + + def stream(self, **_: Any) -> Any: + yield self._data + + +@pytest.fixture +def credentials(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AWS_PROFILE", raising=False) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "test") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test") + + +@pytest.mark.parametrize("client_kind", ["default", "retrying"]) +def test_sends_no_suspend_attempt_after_the_resume_hook( + start: Any, + logger: Any, + credentials: None, + monkeypatch: pytest.MonkeyPatch, + client_kind: str, +) -> None: + # 1. The first SuspendMicrovm attempt suspends the MicroVM. + # 2. Its answer is lost in the freeze. + # 3. The resume hook ends the refusal, and the next job starts. + # 4. botocore would then retry the closed connection. The retry would + # suspend the MicroVM with that job inside. + # The default client makes one attempt. A caller's client can retry, + # and the before-send handler stops its retry after the resume hook. + monkeypatch.setattr(ExponentialBackoff, "delay_amount", lambda self, context: 0) + harness: list[Harness] = [] + service = LostAnswerService(lambda: harness[0].listener) + if client_kind == "default": + default = worker_module.create_microvms_client + monkeypatch.setattr( + worker_module, + "create_microvms_client", + lambda region: service.install(default(region)), + ) + h = start(create_microvms_client=None) + else: + retrying = service.install( + boto3.session.Session().client( + "lambda-microvms", + region_name="us-east-1", + config=Config(retries={"mode": "standard", "total_max_attempts": 3}), + ) + ) + h = start(create_microvms_client=lambda region: retrying) + harness.append(h) + + h.run_hook(IDLE) + h.timers.advance(IDLE) + assert eventually(lambda: service.attempts == 1) + assert h.job("cb-0") == 503 + assert h.hook("resume") == 200 + assert h.job("cb-1") == 202 + assert eventually(lambda: h.started("cb-1")) + service.thaw() + h.settle() + + assert service.attempts == 1 + assert service.frozen_jobs == [] + # The resume hook has ended the refusal. The failed call starts no new + # refusal and no grace period. + assert_one_warning(logger, "unknown", refusing_jobs=False) + assert h.listener._cancel_grace_timer is None # noqa: SLF001 + assert h.job("cb-2") == 202 + h.finish_job("cb-1") + h.finish_job("cb-2") + assert eventually( + lambda: {c["callbackId"] for c in h.lambda_client.recorded()} + == {"cb-1", "cb-2"} + ) + # The next idle time makes a new call, which the resume hook above does + # not stop. + h.timers.advance(IDLE) + h.settle() + assert service.attempts == 2 + assert service.frozen_jobs == [] + + def test_logs_a_rejected_suspend_accepts_jobs_at_once_and_does_not_retry( start: Any, logger: Any ) -> None: diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py index 671e89137..71dc12b3b 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py @@ -20,6 +20,7 @@ from botocore.config import Config from botocore.exceptions import ( ClientError, + ConnectionClosedError, EndpointConnectionError, NoCredentialsError, ) @@ -29,12 +30,16 @@ SUSPEND_CALL_TIMEOUT_SECONDS, SuspendDeadlineError, SuspendOutcome, + SuspendResumedError, classify_suspend_error, create_microvms_client, suspend_microvm, ) +_BEFORE_SEND = "before-send.lambda-microvms.SuspendMicrovm" + + def suspend_error(code: str, status: int, retries: int = 0) -> ClientError: response: Any = { "Error": {"Code": code, "Message": f"{code} message"}, @@ -93,6 +98,9 @@ def test_classifies_a_first_attempt_that_can_succeed_later_as_refused_for_now( suspend_error("InternalServerException", 500), TimeoutError("timed out"), SuspendDeadlineError("passed"), + # A resume hook stopped a later attempt. An earlier attempt may have + # suspended the MicroVM. + SuspendResumedError("resumed"), EndpointConnectionError(endpoint_url="https://example.com"), RuntimeError("boom"), ], @@ -169,7 +177,9 @@ def test_the_default_client_ends_a_read_at_the_call_timeout() -> None: assert client.meta.region_name == "us-west-2" assert config.read_timeout == SUSPEND_CALL_TIMEOUT_SECONDS assert config.connect_timeout == 5 - assert config.retries == {"mode": "standard", "total_max_attempts": 3} + # One attempt: a retry after a lost answer can run after the MicroVM + # resumed, and suspend it with the next job inside. + assert config.retries == {"mode": "standard", "total_max_attempts": 1} finally: client.close() @@ -212,36 +222,40 @@ def send(**_: Any) -> AWSResponse: @pytest.mark.parametrize( - ("status", "code", "headers", "attempts", "outcome"), + ("status", "code", "headers", "outcome"), [ - # A throttle is retried. So its error reports 2 retries, and an - # earlier attempt may have reached the service. - (429, "ThrottlingException", {}, 3, SuspendOutcome.UNKNOWN), - # A signature error with a clock offset is not retried. botocore does - # not correct the clock from the server's Date header. + # The service did nothing. The worker calls again after the next + # idle time. + (429, "ThrottlingException", {}, SuspendOutcome.REFUSED_FOR_NOW), + # A server error does not show whether the service acted. + (500, "InternalServerException", {}, SuspendOutcome.UNKNOWN), + (503, "ServiceUnavailableException", {}, SuspendOutcome.UNKNOWN), + # botocore does not correct the clock from the server's Date header. ( 403, "InvalidSignatureException", {"Date": formatdate(time.time() + 600, usegmt=True)}, - 1, SuspendOutcome.REFUSED_FOR_NOW, ), - (403, "InvalidSignatureException", {}, 1, SuspendOutcome.REFUSED_FOR_NOW), - (403, "ExpiredTokenException", {}, 1, SuspendOutcome.REFUSED_FOR_NOW), + (403, "InvalidSignatureException", {}, SuspendOutcome.REFUSED_FOR_NOW), + (403, "ExpiredTokenException", {}, SuspendOutcome.REFUSED_FOR_NOW), + ], + ids=[ + "throttle", + "server-error", + "unavailable", + "signature-clock-offset", + "signature-mismatch", + "expired-token", ], - ids=["throttle", "signature-clock-offset", "signature-mismatch", "expired-token"], ) -def test_the_default_client_retries_a_throttle_and_no_signature_or_token_error( +def test_the_default_client_makes_one_attempt_for_an_error_answer( credentials: None, - monkeypatch: pytest.MonkeyPatch, status: int, code: str, headers: dict[str, str], - attempts: int, outcome: SuspendOutcome, ) -> None: - # No backoff, so that the retries do not slow the test. - monkeypatch.setattr(ExponentialBackoff, "delay_amount", lambda self, context: 0) client = create_microvms_client("us-east-1") try: sent = _answer_every_attempt(client, status, code, **headers) @@ -249,14 +263,127 @@ def test_the_default_client_retries_a_throttle_and_no_signature_or_token_error( with pytest.raises(ClientError) as raised: suspend_microvm(client, "mvm-1", 5.0) - assert len(sent) == attempts + assert len(sent) == 1 metadata = raised.value.response["ResponseMetadata"] - assert metadata["RetryAttempts"] == attempts - 1 + assert metadata["RetryAttempts"] == 0 assert classify_suspend_error(raised.value) is outcome finally: client.close() +def test_the_default_client_makes_one_attempt_after_a_connection_error( + credentials: None, +) -> None: + # botocore retries a closed connection in its standard mode. The default + # client makes no second attempt, and the outcome is unknown. + client = create_microvms_client("us-east-1") + try: + sent: list[int] = [] + + def close(**_: Any) -> None: + sent.append(1) + raise ConnectionClosedError(endpoint_url="https://example.com") + + client.meta.events.register(_BEFORE_SEND, close) + + with pytest.raises(ConnectionClosedError) as raised: + suspend_microvm(client, "mvm-1", 5.0) + + assert sent == [1] + assert classify_suspend_error(raised.value) is SuspendOutcome.UNKNOWN + finally: + client.close() + + +def test_a_retrying_client_retries_a_throttle_and_its_error_is_unknown( + credentials: None, monkeypatch: pytest.MonkeyPatch +) -> None: + # A caller's client can retry. Its error then reports the retries, and an + # earlier attempt may have reached the service. + monkeypatch.setattr(ExponentialBackoff, "delay_amount", lambda self, context: 0) + client = _retrying_client() + try: + sent = _answer_every_attempt(client, 429, "ThrottlingException") + + with pytest.raises(ClientError) as raised: + suspend_microvm(client, "mvm-1", 5.0) + + assert len(sent) == 3 + assert raised.value.response["ResponseMetadata"]["RetryAttempts"] == 2 + assert classify_suspend_error(raised.value) is SuspendOutcome.UNKNOWN + finally: + client.close() + + +def _retrying_client() -> Any: + """A caller's client that makes 3 attempts.""" + return boto3.session.Session().client( + "lambda-microvms", + region_name="us-east-1", + config=Config(retries={"mode": "standard", "total_max_attempts": 3}), + ) + + +def test_starts_no_attempt_after_the_microvm_resumed( + credentials: None, monkeypatch: pytest.MonkeyPatch +) -> None: + # 1. The first attempt reaches the service, and the service suspends the + # MicroVM. + # 2. The freeze cuts the connection, so the answer is lost. + # 3. The resume hook arrives, and the worker accepts jobs again. + # 4. botocore then retries the closed connection. That attempt would + # suspend the MicroVM with a new job inside. + # So the handler stops the retry once resumed() is True. + monkeypatch.setattr(ExponentialBackoff, "delay_amount", lambda self, context: 0) + client = _retrying_client() + resumes = 0 + try: + sent: list[int] = [] + + def land_then_lose_the_answer(**_: Any) -> None: + nonlocal resumes + sent.append(1) + resumes += 1 + raise ConnectionClosedError(endpoint_url="https://example.com") + + client.meta.events.register(_BEFORE_SEND, land_then_lose_the_answer) + + with pytest.raises(SuspendResumedError): + suspend_microvm(client, "mvm-1", 5.0, resumed=lambda: resumes > 0) + + assert sent == [1] + assert _deadline_handlers(client) == [] + finally: + client.close() + + +def test_starts_no_first_attempt_when_the_microvm_resumed_before_the_call() -> None: + events = _RecordingEvents() + client: Any = _EmittingClient(events) + + with pytest.raises(SuspendResumedError): + suspend_microvm(client, "mvm-1", 5.0, resumed=lambda: True) + + assert client.sent == 0 + assert events.handlers == {} + + +def test_starts_the_attempt_while_the_microvm_has_not_resumed() -> None: + events = _RecordingEvents() + client: Any = _EmittingClient(events) + checks: list[int] = [] + + def resumed() -> bool: + checks.append(1) + return False + + suspend_microvm(client, "mvm-1", 5.0, resumed=resumed) + + assert client.sent == 1 + assert checks == [1] + assert events.handlers == {} + + class _SilentService: """A local HTTP service that accepts connections and never answers, and records when each attempt arrives.""" @@ -427,6 +554,21 @@ def unregister(self, event: str, unique_id: str) -> None: self.handlers.pop(unique_id, None) +class _EmittingClient: + """A client that runs the registered handlers before its one attempt, + as botocore emits ``before-send``.""" + + def __init__(self, events: _RecordingEvents) -> None: + self.meta = type("Meta", (), {"events": events})() + self.sent = 0 + + def suspend_microvm(self, **kwargs: Any) -> dict[str, Any]: + for handler in list(self.meta.events.handlers.values()): + handler(request=None) + self.sent += 1 + return {} + + class _HeldClient: """A client whose calls wait until two of them run at once.""" From e4b8268fa5c55017e1f31bec39ae24883ae6cb4b Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 25 Oct 2026 16:24:10 -0700 Subject: [PATCH 02/12] fix(microvm-worker): answer duplicates during a suspend A resend of a known job got 503 while the worker refused jobs. The durable function then resumed the MicroVM only to get the duplicate answer. The worker now answers a known callback ID with 202 first. --- .../README.md | 8 +- .../worker.py | 41 +++++---- .../tests/lifecycle_hooks_test.py | 84 +++++++++++++++++++ 3 files changed, 112 insertions(+), 21 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md index 8aa5f00e4..c446af227 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -92,7 +92,7 @@ Each job runs in its own thread. For every job, the worker: 3. Completes the callback with the handler's result, or fails it with the handler's exception. A `SystemExit` from the handler is reported as a failure too. 4. Reports a result that is over 256 KB, or that is not JSON-serializable, as a failure. So the durable function learns why at once, instead of at its timeout. -A second delivery of a job is answered like the first and ignored, while the job runs and after it ended. The durable function can deliver the same job again: it retries a delivery whose answer it did not receive, and its request step runs again on replay when its result was not checkpointed. The worker remembers the callback IDs of the last 1,000 finished jobs. When a job ends and the worker already holds 1,000 IDs, it removes the oldest one. So a job runs once, unless at least 1,000 other jobs ended in between, or the worker process restarted. +A second delivery of a job is answered like the first and ignored, while the job runs and after it ended. The durable function can deliver the same job again: it retries a delivery whose answer it did not receive, and its request step runs again on replay when its result was not checkpointed. The worker remembers the callback IDs of the last 1,000 finished jobs. When a job ends and the worker already holds 1,000 IDs, it removes the oldest one. So a job runs once, unless at least 1,000 other jobs ended in between, or the worker process restarted. A second delivery gets this answer also while the worker refuses new jobs after a suspend or after the `terminate` hook. The worker checks the callback ID before it checks the refusal. A second delivery starts nothing, so nothing can freeze with the MicroVM or end with it. A 503 would make the durable function resume a suspending MicroVM only to get the same answer, and nothing would suspend the MicroVM again. The worker creates a Lambda client for each job, after the job arrives, not at image build. A client created before the snapshot would carry build-time state into every MicroVM. A client kept across jobs would keep connections that a suspend of the MicroVM can leave dead. The worker closes each job's default client when the job ends. It never closes a client from `create_client`. @@ -107,7 +107,7 @@ A session's run hook payload can set `autoSuspendIdleSeconds`. The session sets 1. The idle time starts when the `run` hook arrives without a job, when a job ends, and when the `resume` lifecycle hook arrives, if the image enables it. 2. A job that starts cancels the idle time. So the worker never suspends the MicroVM during a job. 3. When the idle time ends, the worker calls `SuspendMicrovm` with the MicroVM ID from the `run` hook, never one from a job request. The call returned in under 100 milliseconds in testing, before the MicroVM froze. The default client makes one attempt. An attempt can suspend the MicroVM and lose its answer in the freeze. A retry would then run after the MicroVM resumed. The `resume` hook has ended the refusal of jobs by then. So the retry could suspend the MicroVM with the next job inside. No attempt of the call starts after 10 seconds, because a `before-send` handler on the client stops botocore's retries at that deadline. The same handler starts no attempt after a `resume` hook. This covers a client from `create_microvms_client` that retries. botocore cannot cancel an attempt in flight, and it opens the connection after `before-send`. So an attempt that started before the deadline can still reach the service after it. The worker waits for that attempt to end, and refuses jobs until then. The client's connect and read timeouts bound the attempt: 5 and 10 seconds for the default client. A client from `create_microvms_client` should use similar timeouts. botocore's own defaults are 60 seconds each, so with them the refusal can last minutes. The durable function resumes the MicroVM before the next job, and the process continues with its memory and files. -4. From the moment the idle time ends, the worker answers every job with HTTP 503 and does not start it. A job accepted then would freeze with the MicroVM, and nothing would resume it. After a 503, the durable function checks the state, resumes the MicroVM, and delivers the job again. +4. From the moment the idle time ends, the worker answers every new job with HTTP 503 and does not start it. A job accepted then would freeze with the MicroVM, and nothing would resume it. After a 503, the durable function checks the state, resumes the MicroVM, and delivers the job again. 5. The refusal ends when the `resume` lifecycle hook arrives, or 30 seconds after the call returned, if no hook arrives. The idle time then starts again. At most one call runs at a time. A `resume` hook that arrives during the call ends the refusal at once, and the next idle time starts when the call ends. So enable the `resume` hook in the image. Without it, the first job after a resume can wait a little over 30 seconds: the 30 seconds, plus up to 4 seconds until the durable function's next delivery attempt. A test measured 32 seconds. 6. A call that the service rejected on its first attempt, such as `AccessDeniedException` for a missing permission, is logged as a warning, and the MicroVM keeps running. Missing credentials count as a rejection. A later call in the same MicroVM most likely finds no credentials either. The worker accepts jobs again at once and does not retry the call. The next job that ends starts the idle time again. The exception is a `resume` hook that arrived during the call. The idle time then starts when the call ends. So the worker retries the call one idle time after the call ends, if no job runs and no `suspend` hook refusal is in effect when the call ends. A second exception is a `suspend` hook that arrived during the call and whose refusal is still in effect when the call ends. The MicroVM is then suspending anyway. So the rejection does not end the hook's refusal. That refusal ends on the `resume` hook, or 30 seconds after the last `suspend` hook. The idle time then starts again. 7. A call that the service refused for now on its first attempt is logged as a warning too, and the worker accepts jobs again at once. This covers expired credentials and a signature error, with or without a clock offset, whatever the HTTP status. The service answers these before it acts. botocore retries none of these errors, so they reach the worker on the first attempt. The idle time then starts again, so the worker tries the call again. That idle time lasts at least 30 seconds, also when `autoSuspendIdleSeconds` is shorter. The same error can repeat on every call, for example with a wrong secret key or an expired session token from environment variables. So while no job ends and no lifecycle hook arrives, the worker makes at most one call per 30 seconds, and logs one warning per call. A job that ends after the refusal, or a `resume` hook after it, starts the normal idle time instead. A `resume` hook or a job that ends during the call does not shorten the idle time after the refusal, because the refused outcome arrives after it. A `suspend` hook whose refusal is still in effect when the call ends is the exception, as in item 6. Its refusal stays, and the idle time starts when it ends. The default client makes one attempt. So a throttle from it counts as refused for now. A client from `create_microvms_client` that retries a throttle turns it into an unknown outcome, as in item 8. The JavaScript worker classifies these errors the same way. @@ -127,7 +127,7 @@ Lambda waits for the answer to a `suspend` or `terminate` hook, up to the hook's On the `suspend` hook the worker: 1. Answers HTTP 200 at once. A suspend hook that ran past its timeout terminated the MicroVM instead of suspending it. -2. Refuses every job with HTTP 503, as after its own suspend. A job accepted now would freeze with the MicroVM. The refusal ends when the `resume` hook arrives, or after 30 seconds. +2. Refuses every new job with HTTP 503, as after its own suspend. A job accepted now would freeze with the MicroVM. The refusal ends when the `resume` hook arrives, or after 30 seconds. 3. Leaves running jobs alone. They freeze with the MicroVM, and continue after a resume. 4. Defers a job from a `run` hook whose body was still arriving. The durable function cannot deliver a `run` hook job again, so the worker does not refuse it. It starts the job when the refusal ends. @@ -139,7 +139,7 @@ On the `terminate` hook the worker: 2. Fails each running job's callback with `MicrovmTerminatedError`. The worker reports `MicrovmTerminatedError` as the callback's error type. The `microvm` operation then fails with a `MicrovmJobFailedError` whose message contains `(MicrovmTerminatedError)`. The worker does not report the handler's outcome when it returns later. 3. Skips a job whose callback a heartbeat already found gone. A report would fail the same way. 4. Answers HTTP 200 when the reports end, or after 5 seconds. A report took 20 to 1,700 milliseconds in testing. The answer also waits for a job that is reporting its own outcome, for a `run` hook request whose body is still arriving, and for every callback failure report in flight, such as the report for an invalid request. Work that begins during the wait is included. A second `terminate` hook that arrives during the wait waits for the same work. Lambda can end the MicroVM as soon as it gets the answer, so a report sent after it can be lost. -5. Refuses every later job with HTTP 503, and never suspends the MicroVM again. A job from a `run` hook whose body was still arriving is not started. Its callback fails with `MicrovmTerminatedError`, as does a job deferred by an earlier `suspend` hook. +5. Refuses every later new job with HTTP 503, and never suspends the MicroVM again. A second delivery of a known job gets HTTP 202. That job has ended, or this hook fails its callback. So the durable function learns the outcome from the callback, and a 503 would only make it deliver the job again. A job from a `run` hook whose body was still arriving is not started. Its callback fails with `MicrovmTerminatedError`, as does a job deferred by an earlier `suspend` hook. Lambda called the hook after `TerminateMicrovm` on a running MicroVM, and at the end of `maximumDurationInSeconds`. It did not call the hook when it terminated a suspended MicroVM, after `TerminateMicrovm` or at the end of the idle policy's `suspendedDurationSeconds`. The process is frozen then. A job of such a MicroVM fails at its heartbeat timeout or its timeout. diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py index 62fcc864b..9cf9b2956 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py @@ -284,7 +284,8 @@ class MicrovmWorkerListener: A second delivery of a callback ID is answered like the first and ignored, while that job runs and after it ended. The listener remembers - the last 1,000 finished callback IDs. + the last 1,000 finished callback IDs. A second delivery gets this answer + during a suspend or terminate refusal too, because it starts nothing. A document that does not match the contract gets HTTP 400. If the document still names a callback, the listener fails that callback. So @@ -301,9 +302,9 @@ class MicrovmWorkerListener: 3. When the idle time ends, the listener calls SuspendMicrovm with the ``run`` hook's MicroVM ID. The call returns before the MicroVM freezes. - 4. From the moment the idle time ends, the listener answers every job - with HTTP 503 and does not start it. A job accepted then would freeze - with the MicroVM. After a 503, the durable function resumes the + 4. From the moment the idle time ends, the listener answers every new + job with HTTP 503 and does not start it. A job accepted then would + freeze with the MicroVM. After a 503, the durable function resumes the MicroVM and delivers the job again. 5. The refusal ends when the ``resume`` hook arrives, or 30 seconds after the call returned. The idle time then starts again. @@ -361,7 +362,7 @@ class MicrovmWorkerListener: - ``suspend``: the listener answers HTTP 200 at once, because a suspend hook that runs past its timeout terminates the MicroVM. It then - refuses jobs with HTTP 503, as after its own suspend, until the + refuses new jobs with HTTP 503, as after its own suspend, until the ``resume`` hook or for 30 seconds. Running jobs freeze with the MicroVM. A ``run`` hook job that arrives during the refusal starts when the refusal ends, because the durable function cannot deliver it @@ -371,8 +372,8 @@ class MicrovmWorkerListener: :class:`MicrovmTerminatedError`. It answers HTTP 200 when the reports end, or after 5 seconds. The answer also waits for jobs that report their own outcome, for ``run`` hook requests in flight, and for the - failure reports of invalid requests. From then on it refuses jobs with - HTTP 503, and fails the callback of a ``run`` hook job instead of + failure reports of invalid requests. From then on it refuses new jobs + with HTTP 503, and fails the callback of a ``run`` hook job instead of starting it. The listener acts on these two hooks, and on ``resume``, only when the @@ -965,11 +966,24 @@ def _handle_job_request( # arrives. self._microvm_id = request.microvm_id microvm_id = self._microvm_id - terminating = self._terminating - suspending = self._suspending + # A known callback ID gets the duplicate answer before the refusal + # checks. Its job runs now or has ended, and the answer starts + # nothing. So nothing can freeze with the MicroVM or end with it. + # A 503 would make the durable function resume a suspending + # MicroVM only to get this answer, and nothing would suspend the + # MicroVM again. duplicate = self._is_known(job.callback_id) - if not (terminating or suspending or duplicate): + terminating = not duplicate and self._terminating + suspending = not duplicate and self._suspending + if not (duplicate or terminating or suspending): self._register_job(job.callback_id, request.region) + if duplicate: + exchange.respond(202, {}) + self._logger.warning( + "duplicate job ignored", + extra={"microvmId": microvm_id, "callbackId": job.callback_id}, + ) + return if terminating or suspending: # Terminating: the MicroVM is about to end, and the job could not # finish. Suspending: the MicroVM is about to freeze. The durable @@ -981,13 +995,6 @@ def _handle_job_request( ) exchange.respond(503, {"error": f"the MicroVM is {state}"}) return - if duplicate: - exchange.respond(202, {}) - self._logger.warning( - "duplicate job ignored", - extra={"microvmId": microvm_id, "callbackId": job.callback_id}, - ) - return # The thread starts before the answer. A thread that cannot start # raises, and the caller gets HTTP 500, which the durable function # retries. After a 202, the durable function does not deliver the job diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py index 458c64fd4..16e941627 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py @@ -2282,6 +2282,90 @@ def test_refuses_jobs_with_503_after_the_terminate_hook(start: Any) -> None: assert not h.started("cb-1") +def test_answers_a_resent_known_job_with_202_after_the_terminate_hook( + start: Any, logger: Any +) -> None: + # A resend starts nothing. The terminate hook fails the running job's + # callback, and the finished job has reported its outcome. So the durable + # function learns each outcome from the callback. A 503 would only make + # it deliver the job again to a MicroVM that is ending. + h = start() + h.run_hook() + assert h.job("cb-done") == 202 + h.finish_job("cb-done") + assert h.job("cb-running") == 202 + assert eventually(lambda: h.started("cb-running")) + assert h.hook("terminate") == 200 + + assert h.job("cb-running") == 202 + assert h.job("cb-done") == 202 + assert h.job("cb-new") == 503 + h.jobs.release_all() + assert h.worker.idle(WAIT) + + assert h.jobs.started() == ["cb-done", "cb-running"] + assert eventually( + lambda: logger.messages("warning").count("duplicate job ignored") == 2 + ) + assert h.lambda_client.recorded() == [ + {"kind": "success", "callbackId": "cb-done"}, + terminated_failure("cb-running"), + ] + + +@pytest.mark.parametrize("refusal", ["own-suspend", "suspend-hook"]) +def test_answers_a_resent_finished_job_with_202_during_a_suspend_refusal( + start: Any, logger: Any, refusal: str +) -> None: + # The job has ended, so a resend starts nothing and cannot freeze with + # the MicroVM. A 503 would make the durable function resume the MicroVM + # only to get this answer. + h = start() + h.run_hook(IDLE) + assert h.job("cb-done") == 202 + h.finish_job("cb-done") + if refusal == "own-suspend": + h.advance(IDLE) + assert h.microvms.suspended == ["mvm-1"] + else: + assert h.hook("suspend") == 200 + + assert h.job("cb-done") == 202 + assert h.job("cb-new") == 503 + + assert h.jobs.started() == ["cb-done"] + assert eventually( + lambda: logger.messages("warning").count("duplicate job ignored") == 1 + ) + # The duplicate did not end the refusal. + assert h.job("cb-new") == 503 + + +def test_answers_a_resent_running_job_with_202_during_a_suspend_refusal( + start: Any, logger: Any +) -> None: + # The job runs and freezes with the MicroVM anyway. A resend starts + # nothing. + h = start() + h.run_hook(IDLE) + assert h.job("cb-1") == 202 + assert eventually(lambda: h.started("cb-1")) + assert h.hook("suspend") == 200 + + assert h.job("cb-1") == 202 + assert h.job("cb-2") == 503 + h.finish_job("cb-1") + + assert h.jobs.started() == ["cb-1"] + assert eventually( + lambda: h.lambda_client.recorded() + == [{"kind": "success", "callbackId": "cb-1"}] + ) + assert eventually( + lambda: logger.messages("warning").count("duplicate job ignored") == 1 + ) + + def test_never_suspends_after_the_terminate_hook(start: Any) -> None: h = start() h.run_hook(IDLE) From fd74e40459941b8c917ddf92ac8c5fae12e0367f Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 25 Oct 2026 17:49:47 -0700 Subject: [PATCH 03/12] docs(microvm-worker): correct suspend wording The 503 note for a resent job said that nothing would suspend the MicroVM again. The idle time restarts after the refusal, so the worker does suspend it again. The real cost is a resume of the MicroVM for nothing. A resume hook can stop the first SuspendMicrovm attempt. Nothing was sent then. The classification docs now say that an earlier attempt, if there was one, may have suspended the MicroVM. The README now states the cost of one attempt. A first-attempt 5xx or dropped connection counts as unknown, and the worker refuses jobs for up to 30 seconds on a MicroVM that did not suspend. It also states why a resent job whose callback is gone gets 202 after the terminate hook. --- .../README.md | 8 ++++---- .../self_suspend.py | 5 +++-- .../worker.py | 14 ++++++++------ 3 files changed, 15 insertions(+), 12 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md index c446af227..0e8b28150 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -92,7 +92,7 @@ Each job runs in its own thread. For every job, the worker: 3. Completes the callback with the handler's result, or fails it with the handler's exception. A `SystemExit` from the handler is reported as a failure too. 4. Reports a result that is over 256 KB, or that is not JSON-serializable, as a failure. So the durable function learns why at once, instead of at its timeout. -A second delivery of a job is answered like the first and ignored, while the job runs and after it ended. The durable function can deliver the same job again: it retries a delivery whose answer it did not receive, and its request step runs again on replay when its result was not checkpointed. The worker remembers the callback IDs of the last 1,000 finished jobs. When a job ends and the worker already holds 1,000 IDs, it removes the oldest one. So a job runs once, unless at least 1,000 other jobs ended in between, or the worker process restarted. A second delivery gets this answer also while the worker refuses new jobs after a suspend or after the `terminate` hook. The worker checks the callback ID before it checks the refusal. A second delivery starts nothing, so nothing can freeze with the MicroVM or end with it. A 503 would make the durable function resume a suspending MicroVM only to get the same answer, and nothing would suspend the MicroVM again. +A second delivery of a job is answered like the first and ignored, while the job runs and after it ended. The durable function can deliver the same job again: it retries a delivery whose answer it did not receive, and its request step runs again on replay when its result was not checkpointed. The worker remembers the callback IDs of the last 1,000 finished jobs. When a job ends and the worker already holds 1,000 IDs, it removes the oldest one. So a job runs once, unless at least 1,000 other jobs ended in between, or the worker process restarted. A second delivery gets this answer also while the worker refuses new jobs after a suspend or after the `terminate` hook. The worker checks the callback ID before it checks the refusal. A second delivery starts nothing, so nothing can freeze with the MicroVM or end with it. After a 503, the durable function resumes the MicroVM and delivers the job again. That delivery gets the same answer. So the MicroVM ran again for nothing. The worker creates a Lambda client for each job, after the job arrives, not at image build. A client created before the snapshot would carry build-time state into every MicroVM. A client kept across jobs would keep connections that a suspend of the MicroVM can leave dead. The worker closes each job's default client when the job ends. It never closes a client from `create_client`. @@ -106,12 +106,12 @@ A session's run hook payload can set `autoSuspendIdleSeconds`. The session sets 1. The idle time starts when the `run` hook arrives without a job, when a job ends, and when the `resume` lifecycle hook arrives, if the image enables it. 2. A job that starts cancels the idle time. So the worker never suspends the MicroVM during a job. -3. When the idle time ends, the worker calls `SuspendMicrovm` with the MicroVM ID from the `run` hook, never one from a job request. The call returned in under 100 milliseconds in testing, before the MicroVM froze. The default client makes one attempt. An attempt can suspend the MicroVM and lose its answer in the freeze. A retry would then run after the MicroVM resumed. The `resume` hook has ended the refusal of jobs by then. So the retry could suspend the MicroVM with the next job inside. No attempt of the call starts after 10 seconds, because a `before-send` handler on the client stops botocore's retries at that deadline. The same handler starts no attempt after a `resume` hook. This covers a client from `create_microvms_client` that retries. botocore cannot cancel an attempt in flight, and it opens the connection after `before-send`. So an attempt that started before the deadline can still reach the service after it. The worker waits for that attempt to end, and refuses jobs until then. The client's connect and read timeouts bound the attempt: 5 and 10 seconds for the default client. A client from `create_microvms_client` should use similar timeouts. botocore's own defaults are 60 seconds each, so with them the refusal can last minutes. The durable function resumes the MicroVM before the next job, and the process continues with its memory and files. +3. When the idle time ends, the worker calls `SuspendMicrovm` with the MicroVM ID from the `run` hook, never one from a job request. The call returned in under 100 milliseconds in testing, before the MicroVM froze. The default client makes one attempt. An attempt can suspend the MicroVM and lose its answer in the freeze. A retry would then run after the MicroVM resumed. The `resume` hook has ended the refusal of jobs by then. So the retry could suspend the MicroVM with the next job inside. With one attempt, a 5xx answer or a dropped connection on the first attempt of the default client counts as unknown, as in item 8. botocore's retry would often have succeeded. So the worker refuses jobs for up to 30 seconds on a MicroVM that did not suspend. No `resume` hook ends that refusal, because nothing suspended. This delay is the accepted cost of one attempt. No attempt of the call starts after 10 seconds, because a `before-send` handler on the client stops botocore's retries at that deadline. The same handler starts no attempt after a `resume` hook. This covers a client from `create_microvms_client` that retries. botocore cannot cancel an attempt in flight, and it opens the connection after `before-send`. So an attempt that started before the deadline can still reach the service after it. The worker waits for that attempt to end, and refuses jobs until then. The client's connect and read timeouts bound the attempt: 5 and 10 seconds for the default client. A client from `create_microvms_client` should use similar timeouts. botocore's own defaults are 60 seconds each, so with them the refusal can last minutes. The durable function resumes the MicroVM before the next job, and the process continues with its memory and files. 4. From the moment the idle time ends, the worker answers every new job with HTTP 503 and does not start it. A job accepted then would freeze with the MicroVM, and nothing would resume it. After a 503, the durable function checks the state, resumes the MicroVM, and delivers the job again. 5. The refusal ends when the `resume` lifecycle hook arrives, or 30 seconds after the call returned, if no hook arrives. The idle time then starts again. At most one call runs at a time. A `resume` hook that arrives during the call ends the refusal at once, and the next idle time starts when the call ends. So enable the `resume` hook in the image. Without it, the first job after a resume can wait a little over 30 seconds: the 30 seconds, plus up to 4 seconds until the durable function's next delivery attempt. A test measured 32 seconds. 6. A call that the service rejected on its first attempt, such as `AccessDeniedException` for a missing permission, is logged as a warning, and the MicroVM keeps running. Missing credentials count as a rejection. A later call in the same MicroVM most likely finds no credentials either. The worker accepts jobs again at once and does not retry the call. The next job that ends starts the idle time again. The exception is a `resume` hook that arrived during the call. The idle time then starts when the call ends. So the worker retries the call one idle time after the call ends, if no job runs and no `suspend` hook refusal is in effect when the call ends. A second exception is a `suspend` hook that arrived during the call and whose refusal is still in effect when the call ends. The MicroVM is then suspending anyway. So the rejection does not end the hook's refusal. That refusal ends on the `resume` hook, or 30 seconds after the last `suspend` hook. The idle time then starts again. 7. A call that the service refused for now on its first attempt is logged as a warning too, and the worker accepts jobs again at once. This covers expired credentials and a signature error, with or without a clock offset, whatever the HTTP status. The service answers these before it acts. botocore retries none of these errors, so they reach the worker on the first attempt. The idle time then starts again, so the worker tries the call again. That idle time lasts at least 30 seconds, also when `autoSuspendIdleSeconds` is shorter. The same error can repeat on every call, for example with a wrong secret key or an expired session token from environment variables. So while no job ends and no lifecycle hook arrives, the worker makes at most one call per 30 seconds, and logs one warning per call. A job that ends after the refusal, or a `resume` hook after it, starts the normal idle time instead. A `resume` hook or a job that ends during the call does not shorten the idle time after the refusal, because the refused outcome arrives after it. A `suspend` hook whose refusal is still in effect when the call ends is the exception, as in item 6. Its refusal stays, and the idle time starts when it ends. The default client makes one attempt. So a throttle from it counts as refused for now. A client from `create_microvms_client` that retries a throttle turns it into an unknown outcome, as in item 8. The JavaScript worker classifies these errors the same way. -8. A call whose outcome is unknown is handled like a successful call, because the service may have suspended the MicroVM anyway. This covers a timeout, a conflict, a server error, and a dropped connection. An error after botocore retried an earlier attempt also counts as unknown, because the earlier attempt may have suspended the MicroVM. Only a client from `create_microvms_client` that retries can raise such an error. A call that the `before-send` handler stopped after a `resume` hook counts as unknown too. That `resume` hook has already ended the refusal. So the worker accepts jobs at once, and starts no new 30-second limit for that call. The refusal continues until the `resume` hook or the 30-second limit. The idle time then starts again, so the worker tries the call again. A `resume` hook during the call, or the end of a `suspend` hook's 30 seconds during the call, has already ended the refusal. The worker then accepts jobs at once, and tries the call again one idle time after the call and every job end. +8. A call whose outcome is unknown is handled like a successful call, because the service may have suspended the MicroVM anyway. This covers a timeout, a conflict, a server error, and a dropped connection. An error after botocore retried an earlier attempt also counts as unknown, because the earlier attempt may have suspended the MicroVM. Only a client from `create_microvms_client` that retries can raise such an error. A call that the `before-send` handler stopped after a `resume` hook counts as unknown too. The handler can stop the first attempt, and then nothing was sent. An earlier attempt, if there was one, may have suspended the MicroVM. That `resume` hook has already ended the refusal. So the worker accepts jobs at once, and starts no new 30-second limit for that call. The refusal continues until the `resume` hook or the 30-second limit. The idle time then starts again, so the worker tries the call again. A `resume` hook during the call, or the end of a `suspend` hook's 30 seconds during the call, has already ended the refusal. The worker then accepts jobs at once, and tries the call again one idle time after the call and every job end. 9. The worker logs one warning for each failed call. The message names the outcome. The fields are, in order, `microvmId`, `error`, `outcome`, `suspendHookArrived`, `terminating`, `closed`, and `refusingJobs`. `microvmId` is the MicroVM ID of the call. `error` holds the name and the message of the error. `outcome` is `rejected`, `refusedForNow`, or `unknown`. `suspendHookArrived` states whether a `suspend` hook arrived since the idle time ended. `terminating` states whether a `terminate` hook came before the worker acted on the outcome. `closed` states whether `worker.close()` came before the worker acted on the outcome. `refusingJobs` states whether the worker answers jobs with HTTP 503 after it has acted on the outcome. `worker.close()` stops the idle time. A closed worker never suspends its MicroVM again. A refusal that is already running still ends on the `resume` hook or after 30 seconds. @@ -139,7 +139,7 @@ On the `terminate` hook the worker: 2. Fails each running job's callback with `MicrovmTerminatedError`. The worker reports `MicrovmTerminatedError` as the callback's error type. The `microvm` operation then fails with a `MicrovmJobFailedError` whose message contains `(MicrovmTerminatedError)`. The worker does not report the handler's outcome when it returns later. 3. Skips a job whose callback a heartbeat already found gone. A report would fail the same way. 4. Answers HTTP 200 when the reports end, or after 5 seconds. A report took 20 to 1,700 milliseconds in testing. The answer also waits for a job that is reporting its own outcome, for a `run` hook request whose body is still arriving, and for every callback failure report in flight, such as the report for an invalid request. Work that begins during the wait is included. A second `terminate` hook that arrives during the wait waits for the same work. Lambda can end the MicroVM as soon as it gets the answer, so a report sent after it can be lost. -5. Refuses every later new job with HTTP 503, and never suspends the MicroVM again. A second delivery of a known job gets HTTP 202. That job has ended, or this hook fails its callback. So the durable function learns the outcome from the callback, and a 503 would only make it deliver the job again. A job from a `run` hook whose body was still arriving is not started. Its callback fails with `MicrovmTerminatedError`, as does a job deferred by an earlier `suspend` hook. +5. Refuses every later new job with HTTP 503, and never suspends the MicroVM again. A second delivery of a known job gets HTTP 202. That job has ended, or this hook fails its callback. So the durable function learns the outcome from the callback, and a 503 would only make it deliver the job again. A job whose callback a heartbeat already found gone gets HTTP 202 too. That callback is gone, so the durable function no longer waits for it. A job from a `run` hook whose body was still arriving is not started. Its callback fails with `MicrovmTerminatedError`, as does a job deferred by an earlier `suspend` hook. Lambda called the hook after `TerminateMicrovm` on a running MicroVM, and at the end of `maximumDurationInSeconds`. It did not call the hook when it terminated a suspended MicroVM, after `TerminateMicrovm` or at the end of the idle policy's `suspendedDurationSeconds`. The process is frozen then. A job of such a MicroVM fails at its heartbeat timeout or its timeout. diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py index 80cf36090..948bfe422 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py @@ -244,8 +244,9 @@ def classify_suspend_error(error: BaseException) -> SuspendOutcome: 7. Any other error is :attr:`SuspendOutcome.UNKNOWN`. Examples are a timeout, a dropped connection, a 408, a server error, :class:`SuspendDeadlineError`, and :class:`SuspendResumedError`. The - last two stop a later attempt, and an earlier attempt may have - suspended the MicroVM. + last two stop an attempt. When they stop the first attempt, nothing + was sent. An earlier attempt, if there was one, may have suspended + the MicroVM. """ if isinstance(error, NoCredentialsError): return SuspendOutcome.REJECTED diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py index 9cf9b2956..5057810b3 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py @@ -969,9 +969,9 @@ def _handle_job_request( # A known callback ID gets the duplicate answer before the refusal # checks. Its job runs now or has ended, and the answer starts # nothing. So nothing can freeze with the MicroVM or end with it. - # A 503 would make the durable function resume a suspending - # MicroVM only to get this answer, and nothing would suspend the - # MicroVM again. + # After a 503, the durable function resumes the MicroVM and + # delivers the job again. That delivery gets this answer. So the + # MicroVM ran again for nothing. duplicate = self._is_known(job.callback_id) terminating = not duplicate and self._terminating suspending = not duplicate and self._suspending @@ -1587,9 +1587,11 @@ def _call_suspend( # retried attempt does not show whether the service suspended # the MicroVM. So the refusal continues, as after a successful # call. The block after the client is closed acts on it. A - # resume hook that stopped the call's next attempt has - # already ended the refusal. So that block starts a grace - # period only for a suspend hook that came after the resume. + # resume hook can stop any attempt, the first one included. An + # earlier attempt, if there was one, may have suspended the + # MicroVM. That resume hook has already ended the refusal. So + # that block starts a grace period only for a suspend hook + # that came after the resume. unknown_described = described else: deferred = [] From 9a0a6e0d53297693060a748f397d879a8d6f7313 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 25 Oct 2026 17:49:53 -0700 Subject: [PATCH 04/12] fix(microvm-worker): log a stopped suspend call The JavaScript worker logs an info line when a resume hook arrives during its SuspendMicrovm call. This worker logged nothing until the unknown warning at the end of the call. It now logs the same line, with the same text and the call's microvmId. The resume hook decides under the lock and logs after it. Only the first resume hook during a call logs the line. A resume hook after the call returned logs nothing, because no attempt is left to stop. New listener tests cover a resume hook before the first attempt, a suspend hook after that resume and before the call returns, a second resume hook during the call, and a resume hook after the call returned. The README documents the line. It also states that the JavaScript worker cancels the attempt in flight, while this worker can only stop the next attempt. --- .../README.md | 4 +- .../worker.py | 59 ++++-- .../tests/lifecycle_hooks_test.py | 181 ++++++++++++++++++ 3 files changed, 230 insertions(+), 14 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md index 0e8b28150..c5389485c 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -108,11 +108,11 @@ A session's run hook payload can set `autoSuspendIdleSeconds`. The session sets 2. A job that starts cancels the idle time. So the worker never suspends the MicroVM during a job. 3. When the idle time ends, the worker calls `SuspendMicrovm` with the MicroVM ID from the `run` hook, never one from a job request. The call returned in under 100 milliseconds in testing, before the MicroVM froze. The default client makes one attempt. An attempt can suspend the MicroVM and lose its answer in the freeze. A retry would then run after the MicroVM resumed. The `resume` hook has ended the refusal of jobs by then. So the retry could suspend the MicroVM with the next job inside. With one attempt, a 5xx answer or a dropped connection on the first attempt of the default client counts as unknown, as in item 8. botocore's retry would often have succeeded. So the worker refuses jobs for up to 30 seconds on a MicroVM that did not suspend. No `resume` hook ends that refusal, because nothing suspended. This delay is the accepted cost of one attempt. No attempt of the call starts after 10 seconds, because a `before-send` handler on the client stops botocore's retries at that deadline. The same handler starts no attempt after a `resume` hook. This covers a client from `create_microvms_client` that retries. botocore cannot cancel an attempt in flight, and it opens the connection after `before-send`. So an attempt that started before the deadline can still reach the service after it. The worker waits for that attempt to end, and refuses jobs until then. The client's connect and read timeouts bound the attempt: 5 and 10 seconds for the default client. A client from `create_microvms_client` should use similar timeouts. botocore's own defaults are 60 seconds each, so with them the refusal can last minutes. The durable function resumes the MicroVM before the next job, and the process continues with its memory and files. 4. From the moment the idle time ends, the worker answers every new job with HTTP 503 and does not start it. A job accepted then would freeze with the MicroVM, and nothing would resume it. After a 503, the durable function checks the state, resumes the MicroVM, and delivers the job again. -5. The refusal ends when the `resume` lifecycle hook arrives, or 30 seconds after the call returned, if no hook arrives. The idle time then starts again. At most one call runs at a time. A `resume` hook that arrives during the call ends the refusal at once, and the next idle time starts when the call ends. So enable the `resume` hook in the image. Without it, the first job after a resume can wait a little over 30 seconds: the 30 seconds, plus up to 4 seconds until the durable function's next delivery attempt. A test measured 32 seconds. +5. The refusal ends when the `resume` lifecycle hook arrives, or 30 seconds after the call returned, if no hook arrives. The idle time then starts again. At most one call runs at a time. A `resume` hook that arrives during the call ends the refusal at once, and the next idle time starts when the call ends. The call then starts no further attempt, and the worker logs `stopping the suspend call: the MicroVM resumed during it` at info level. A second `resume` hook during the same call logs nothing. The JavaScript worker cancels the attempt in flight. This worker can only stop the next attempt, because botocore cannot cancel an attempt in flight. So enable the `resume` hook in the image. Without it, the first job after a resume can wait a little over 30 seconds: the 30 seconds, plus up to 4 seconds until the durable function's next delivery attempt. A test measured 32 seconds. 6. A call that the service rejected on its first attempt, such as `AccessDeniedException` for a missing permission, is logged as a warning, and the MicroVM keeps running. Missing credentials count as a rejection. A later call in the same MicroVM most likely finds no credentials either. The worker accepts jobs again at once and does not retry the call. The next job that ends starts the idle time again. The exception is a `resume` hook that arrived during the call. The idle time then starts when the call ends. So the worker retries the call one idle time after the call ends, if no job runs and no `suspend` hook refusal is in effect when the call ends. A second exception is a `suspend` hook that arrived during the call and whose refusal is still in effect when the call ends. The MicroVM is then suspending anyway. So the rejection does not end the hook's refusal. That refusal ends on the `resume` hook, or 30 seconds after the last `suspend` hook. The idle time then starts again. 7. A call that the service refused for now on its first attempt is logged as a warning too, and the worker accepts jobs again at once. This covers expired credentials and a signature error, with or without a clock offset, whatever the HTTP status. The service answers these before it acts. botocore retries none of these errors, so they reach the worker on the first attempt. The idle time then starts again, so the worker tries the call again. That idle time lasts at least 30 seconds, also when `autoSuspendIdleSeconds` is shorter. The same error can repeat on every call, for example with a wrong secret key or an expired session token from environment variables. So while no job ends and no lifecycle hook arrives, the worker makes at most one call per 30 seconds, and logs one warning per call. A job that ends after the refusal, or a `resume` hook after it, starts the normal idle time instead. A `resume` hook or a job that ends during the call does not shorten the idle time after the refusal, because the refused outcome arrives after it. A `suspend` hook whose refusal is still in effect when the call ends is the exception, as in item 6. Its refusal stays, and the idle time starts when it ends. The default client makes one attempt. So a throttle from it counts as refused for now. A client from `create_microvms_client` that retries a throttle turns it into an unknown outcome, as in item 8. The JavaScript worker classifies these errors the same way. 8. A call whose outcome is unknown is handled like a successful call, because the service may have suspended the MicroVM anyway. This covers a timeout, a conflict, a server error, and a dropped connection. An error after botocore retried an earlier attempt also counts as unknown, because the earlier attempt may have suspended the MicroVM. Only a client from `create_microvms_client` that retries can raise such an error. A call that the `before-send` handler stopped after a `resume` hook counts as unknown too. The handler can stop the first attempt, and then nothing was sent. An earlier attempt, if there was one, may have suspended the MicroVM. That `resume` hook has already ended the refusal. So the worker accepts jobs at once, and starts no new 30-second limit for that call. The refusal continues until the `resume` hook or the 30-second limit. The idle time then starts again, so the worker tries the call again. A `resume` hook during the call, or the end of a `suspend` hook's 30 seconds during the call, has already ended the refusal. The worker then accepts jobs at once, and tries the call again one idle time after the call and every job end. -9. The worker logs one warning for each failed call. The message names the outcome. The fields are, in order, `microvmId`, `error`, `outcome`, `suspendHookArrived`, `terminating`, `closed`, and `refusingJobs`. `microvmId` is the MicroVM ID of the call. `error` holds the name and the message of the error. `outcome` is `rejected`, `refusedForNow`, or `unknown`. `suspendHookArrived` states whether a `suspend` hook arrived since the idle time ended. `terminating` states whether a `terminate` hook came before the worker acted on the outcome. `closed` states whether `worker.close()` came before the worker acted on the outcome. `refusingJobs` states whether the worker answers jobs with HTTP 503 after it has acted on the outcome. +9. The worker logs one warning for each failed call. The message names the outcome. The fields are, in order, `microvmId`, `error`, `outcome`, `suspendHookArrived`, `terminating`, `closed`, and `refusingJobs`. `microvmId` is the MicroVM ID of the call. `error` holds the name and the message of the error. `outcome` is `rejected`, `refusedForNow`, or `unknown`. `suspendHookArrived` states whether a `suspend` hook arrived since the idle time ended. `terminating` states whether a `terminate` hook came before the worker acted on the outcome. `closed` states whether `worker.close()` came before the worker acted on the outcome. `refusingJobs` states whether the worker answers jobs with HTTP 503 after it has acted on the outcome. The worker also logs two info lines about the call. It logs `suspending the idle MicroVM` before each call, with `microvmId`, `idleSeconds`, and `waitedSeconds`. It logs `stopping the suspend call: the MicroVM resumed during it` for the first `resume` hook during a call, with `microvmId`. The JavaScript worker logs the same stop line. `worker.close()` stops the idle time. A closed worker never suspends its MicroVM again. A refusal that is already running still ends on the `resume` hook or after 30 seconds. diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py index 5057810b3..ff7651a78 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py @@ -463,6 +463,12 @@ def __init__( # without the lock: an int read is atomic, and the handler then # waits for no listener lock inside botocore. self._resume_hooks = 0 + # The MicroVM ID of the SuspendMicrovm call in flight. It is set when + # the idle time ends, and cleared when the call returns or the first + # resume hook during the call arrives. That hook logs that it stops + # the call. So a second resume hook during the same call logs + # nothing. + self._stoppable_suspend: str | None = None # SuspendMicrovm calls in flight. idle() waits for them. At most one # runs at a time. self._suspend_calls = 0 @@ -723,13 +729,26 @@ def _handle_resume_hook(self) -> list[tuple[MicrovmJobDocument, str]]: until the next job ends. Returns the deferred run hook jobs, registered. The caller starts them after the answer. A SuspendMicrovm call in flight starts no attempt after this hook. + The first such hook logs that it stops the call. """ with self._locked(): self._resume_hooks += 1 + stopped = self._stoppable_suspend + self._stoppable_suspend = None if self._suspending: - return self._end_suspending() - self._start_idle_time() - return [] + deferred = self._end_suspending() + else: + self._start_idle_time() + deferred = [] + # The caller of the listener passes the logger. So the line is + # written after the lock is released. botocore cannot cancel an + # attempt in flight. So the call stops at its next attempt. + if stopped is not None: + self._logger.info( + "stopping the suspend call: the MicroVM resumed during it", + extra={"microvmId": stopped}, + ) + return deferred def _handle_terminate_hook(self, exchange: _Exchange) -> None: """Fail the callback of every running job, then answer the hook. @@ -1462,6 +1481,7 @@ def _idle_time_ended(self, generation: int, waited_seconds: float) -> None: # A resume hook from here on ends the refusal. The call then # starts no further attempt. resumes_before = self._resume_hooks + self._stoppable_suspend = auto_suspend.microvm_id try: threading.Thread( target=self._suspend_self, @@ -1514,6 +1534,8 @@ def _suspend_call_ended(self) -> None: a resume hook, a job, or a refusal for now asked for during it. The caller holds the lock through :meth:`_locked`.""" self._suspend_calls -= 1 + # The thread or the client could fail before the call was sent. + self._stoppable_suspend = None if self._idle_after_call: self._idle_after_call = False self._start_idle_time(after_refusal=self._idle_after_call_floor) @@ -1564,15 +1586,7 @@ def _call_suspend( # close(). unknown_described: object | None = None try: - # The before-send handler calls resumed() inside botocore, in - # this thread. It reads one int without the lock. So it runs no - # code of the caller, and waits for no listener lock. - suspend_microvm( - client, - target, - SUSPEND_CALL_TIMEOUT_SECONDS, - resumed=lambda: self._resume_hooks != resumes_before, - ) + self._send_suspend(client, target, resumes_before) except Exception as error: # noqa: BLE001 outcome = classify_suspend_error(error) # describe() calls str(error). A client from @@ -1639,6 +1653,27 @@ def _call_suspend( ) self._logger.warning(message, extra=extra) + def _send_suspend( + self, client: LambdaMicrovmsClient, target: str, resumes_before: int + ) -> None: + """Make the SuspendMicrovm call. Start no attempt after a resume hook + that arrived since the idle time ended.""" + try: + # The before-send handler calls resumed() inside botocore, in + # this thread. It reads one int without the lock. So it runs no + # code of the caller, and waits for no listener lock. + suspend_microvm( + client, + target, + SUSPEND_CALL_TIMEOUT_SECONDS, + resumed=lambda: self._resume_hooks != resumes_before, + ) + finally: + # The call has returned. A resume hook from now on stops nothing, + # so it logs no stop. + with self._lock: + self._stoppable_suspend = None + def _failed_call_log( self, outcome: SuspendOutcome, diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py index 16e941627..f1f0bb2d1 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py @@ -1042,6 +1042,163 @@ def test_sends_no_suspend_attempt_after_the_resume_hook( assert service.frozen_jobs == [] +STOPPED = "stopping the suspend call: the MicroVM resumed during it" + + +def stop_lines(logger: Any) -> list[dict[str, Any]]: + """The fields of each info line that a resume hook logged when it + stopped a SuspendMicrovm call.""" + return [d for level, m, d in logger.lines if level == "info" and m == STOPPED] + + +class AnsweringService: + """Stands in for the Lambda MicroVMs service in a botocore client. + + A ``before-send`` handler counts each attempt and answers 200 instead + of HTTP. The worker's own handler runs first. So an attempt that it + stops is not counted. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self.attempts = 0 + + def install(self, client: Any) -> Any: + client.meta.events.register( + "before-send.lambda-microvms.SuspendMicrovm", self._send + ) + return client + + def _send(self, **_: Any) -> AWSResponse: + with self._lock: + self.attempts += 1 + headers = HTTPHeaders() + headers["Content-Type"] = "application/json" + return AWSResponse("https://example.com", 200, headers, _Body(b"{}")) + + +class SlowClientFactory: + """A ``create_microvms_client`` that blocks until :meth:`release`. + + It then returns the default client, with ``service`` installed. So a + test can send hooks after the idle time ended and before the first + attempt. + """ + + def __init__(self, service: AnsweringService) -> None: + self.service = service + self.creating = threading.Event() + self._release = threading.Event() + + def __call__(self, region: str) -> Any: + self.creating.set() + self._release.wait(WAIT) + return self.service.install(worker_module.create_microvms_client(region)) + + def release(self) -> None: + self._release.set() + + +def test_starts_no_attempt_after_a_resume_hook_while_the_client_is_created( + start: Any, logger: Any, credentials: None +) -> None: + # The idle time has ended, and the worker still creates the client. A + # resume hook now ends the refusal. A job can then start, and an attempt + # would freeze it. So the call sends nothing. + service = AnsweringService() + factory = SlowClientFactory(service) + h = start(create_microvms_client=factory) + h.run_hook(IDLE) + h.timers.advance(IDLE) + assert factory.creating.wait(WAIT) + assert h.job("cb-0") == 503 + assert h.hook("resume") == 200 + factory.release() + h.settle() + + assert service.attempts == 0 + data = assert_one_warning(logger, "unknown", refusing_jobs=False) + assert data["error"]["name"] == "SuspendResumedError" + assert h.listener._cancel_grace_timer is None # noqa: SLF001 + assert stop_lines(logger) == [{"microvmId": "mvm-1"}] + # The resume hook asked for the idle time during the call. So the idle + # time starts when the call ends, and its call is not stopped. + h.advance(IDLE - 0.001) + assert service.attempts == 0 + h.advance(0.001) + assert service.attempts == 1 + assert stop_lines(logger) == [{"microvmId": "mvm-1"}] + + +@pytest.mark.parametrize("end", ["resume", "grace"]) +def test_keeps_a_suspend_hook_refusal_that_follows_a_resume_that_stopped_the_call( + start: Any, logger: Any, credentials: None, end: str +) -> None: + # The resume hook ends the refusal and stops the call. A suspend hook + # then starts a refusal of its own, before the stopped call returns. + # Lambda is freezing the MicroVM, so the end of the call keeps that + # refusal. + service = AnsweringService() + factory = SlowClientFactory(service) + h = start(create_microvms_client=factory) + h.run_hook(IDLE) + h.timers.advance(IDLE) + assert factory.creating.wait(WAIT) + assert h.hook("resume") == 200 + assert h.hook("suspend") == 200 + factory.release() + h.settle() + + assert service.attempts == 0 + assert_one_warning(logger, "unknown", refusing_jobs=True, suspend_hook_arrived=True) + assert h.job("cb-1") == 503 + if end == "resume": + assert h.hook("resume") == 200 + else: + h.advance(GRACE - 0.001) + assert h.job("cb-1") == 503 + h.advance(0.001) + assert h.job("cb-1") == 202 + assert eventually(lambda: h.started("cb-1")) + # The call had returned before the second resume hook. So that hook + # stopped nothing, and logged no stop. + assert stop_lines(logger) == [{"microvmId": "mvm-1"}] + + +def test_a_second_resume_hook_during_the_call_stops_it_too_and_logs_once( + start: Any, logger: Any, credentials: None, monkeypatch: pytest.MonkeyPatch +) -> None: + # The first attempt is in flight when both resume hooks arrive. Its + # answer is lost. botocore would retry it, and the retry would freeze + # the next job. The second resume hook must not undo the stop. + monkeypatch.setattr(ExponentialBackoff, "delay_amount", lambda self, context: 0) + harness: list[Harness] = [] + service = LostAnswerService(lambda: harness[0].listener) + retrying = service.install( + boto3.session.Session().client( + "lambda-microvms", + region_name="us-east-1", + config=Config(retries={"mode": "standard", "total_max_attempts": 3}), + ) + ) + h = start(create_microvms_client=lambda region: retrying) + harness.append(h) + h.run_hook(IDLE) + h.timers.advance(IDLE) + assert eventually(lambda: service.attempts == 1) + assert h.hook("resume") == 200 + assert h.hook("resume") == 200 + service.thaw() + h.settle() + + assert service.attempts == 1 + assert_one_warning(logger, "unknown", refusing_jobs=False) + assert h.listener._cancel_grace_timer is None # noqa: SLF001 + assert stop_lines(logger) == [{"microvmId": "mvm-1"}] + assert h.job("cb-1") == 202 + assert eventually(lambda: h.started("cb-1")) + + def test_logs_a_rejected_suspend_accepts_jobs_at_once_and_does_not_retry( start: Any, logger: Any ) -> None: @@ -1363,6 +1520,30 @@ def test_a_job_that_ends_before_a_refused_call_has_ended_uses_the_normal_idle_ti assert client.suspended == ["mvm-1", "mvm-1"] +def test_logs_no_stop_for_a_resume_hook_after_the_call_returned( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # The call has returned, and the worker still closes its client. A + # resume hook now has no attempt left to stop. So it logs no stop. + client = _ClosingClient() + monkeypatch.setattr(worker_module, "create_microvms_client", lambda r: client) + h = start(create_microvms_client=None) + + h.run_hook(IDLE) + h.timers.advance(IDLE) + assert client.closing.wait(WAIT) + assert h.hook("resume") == 200 + client.release_close.set() + h.settle() + + assert client.suspended == ["mvm-1"] + assert stop_lines(logger) == [] + # The resume hook has ended the refusal. So the end of the call starts no + # grace period. + assert h.listener._cancel_grace_timer is None # noqa: SLF001 + assert h.job("cb-1") == 202 + + def test_keeps_a_hook_refusal_after_a_refusal_for_now_of_the_own_call( start: Any, logger: Any ) -> None: From 66cfb136f0293b8c725ec5d389490a4ed8b04e44 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 25 Oct 2026 18:57:40 -0700 Subject: [PATCH 05/12] feat(microvm-worker): suspend after a job A job request with "suspendAfter": true asks the worker to suspend its MicroVM when the job ends. The worker answers such a request with the header X-Microvm-Suspend-After: accepted. It suspends only the MicroVM that the run hook named, once no other work runs. The call uses the same path as the suspend when idle. --- .../README.md | 31 +- .../payload.py | 9 + .../worker.py | 455 ++++++++++++----- .../tests/lifecycle_hooks_test.py | 463 +++++++++++++++++- .../tests/payload_test.py | 29 ++ 5 files changed, 866 insertions(+), 121 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md index c5389485c..cb46707bd 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -65,7 +65,7 @@ The image needs these hook settings in `CreateMicrovmImage`: ``` - The `run` hook delivers the job, or tells the worker that jobs arrive over HTTP. So it must be enabled. -- The `resume` hook tells a session's worker that its MicroVM runs again after a suspend. It is optional, and recommended for sessions. See [Suspending when idle](#suspending-when-idle). +- The `resume` hook tells a session's worker that its MicroVM runs again after a suspend. It is optional, and recommended for sessions. A session that uses `suspendAfter` needs it. See [Suspending when idle](#suspending-when-idle) and [Suspending after a job](#suspending-after-a-job). - The `suspend` hook tells the worker that Lambda is about to freeze the MicroVM. The worker then refuses new jobs until the `resume` hook. It is optional, and recommended for sessions. - The `terminate` hook tells the worker that Lambda is about to end the MicroVM. The worker then fails the callback of each running job, so the durable function fails the job at once instead of at its heartbeat timeout or its timeout. It is optional, and recommended. Give it at least 10 seconds, because the worker reports before it answers, for at most 5 seconds. A hook timeout left unset ended the hook after about 2 seconds in testing. - The service requires the `ready` hook whenever a lifecycle hook is enabled. The worker answers it with HTTP 200. @@ -83,7 +83,7 @@ On `POST /aws/lambda-microvms/runtime/v1/run` the worker: The worker acts only on the first valid `run` hook. The MicroVM endpoint forwards a request to a hook path from any caller that holds an auth token for the MicroVM. So a later `run` hook gets HTTP 200, and is ignored. -On `POST ` the worker validates the job request, answers HTTP 202 at once, and runs the route's handler. `POST /durable-execution/v1/job` (`MICROVM_JOB_PATH`) runs `handler`. That path is reserved, so a route cannot use it. A job request can arrive before the `run` hook. Until the `run` hook arrives, the worker takes the MicroVM ID for `context.microvm_id` from the first job request that has a `microvmId`. The `run` hook's ID then replaces it for later jobs. Without either, `context.microvm_id` is `"unknown"`. +On `POST ` the worker validates the job request, answers HTTP 202 at once, and runs the route's handler. A job request with `"suspendAfter": true` gets the header `X-Microvm-Suspend-After: accepted` in its 202 answer. See [Suspending after a job](#suspending-after-a-job). `POST /durable-execution/v1/job` (`MICROVM_JOB_PATH`) runs `handler`. That path is reserved, so a route cannot use it. A job request can arrive before the `run` hook. Until the `run` hook arrives, the worker takes the MicroVM ID for `context.microvm_id` from the first job request that has a `microvmId`. The `run` hook's ID then replaces it for later jobs. Without either, `context.microvm_id` is `"unknown"`. Each job runs in its own thread. For every job, the worker: @@ -112,14 +112,33 @@ A session's run hook payload can set `autoSuspendIdleSeconds`. The session sets 6. A call that the service rejected on its first attempt, such as `AccessDeniedException` for a missing permission, is logged as a warning, and the MicroVM keeps running. Missing credentials count as a rejection. A later call in the same MicroVM most likely finds no credentials either. The worker accepts jobs again at once and does not retry the call. The next job that ends starts the idle time again. The exception is a `resume` hook that arrived during the call. The idle time then starts when the call ends. So the worker retries the call one idle time after the call ends, if no job runs and no `suspend` hook refusal is in effect when the call ends. A second exception is a `suspend` hook that arrived during the call and whose refusal is still in effect when the call ends. The MicroVM is then suspending anyway. So the rejection does not end the hook's refusal. That refusal ends on the `resume` hook, or 30 seconds after the last `suspend` hook. The idle time then starts again. 7. A call that the service refused for now on its first attempt is logged as a warning too, and the worker accepts jobs again at once. This covers expired credentials and a signature error, with or without a clock offset, whatever the HTTP status. The service answers these before it acts. botocore retries none of these errors, so they reach the worker on the first attempt. The idle time then starts again, so the worker tries the call again. That idle time lasts at least 30 seconds, also when `autoSuspendIdleSeconds` is shorter. The same error can repeat on every call, for example with a wrong secret key or an expired session token from environment variables. So while no job ends and no lifecycle hook arrives, the worker makes at most one call per 30 seconds, and logs one warning per call. A job that ends after the refusal, or a `resume` hook after it, starts the normal idle time instead. A `resume` hook or a job that ends during the call does not shorten the idle time after the refusal, because the refused outcome arrives after it. A `suspend` hook whose refusal is still in effect when the call ends is the exception, as in item 6. Its refusal stays, and the idle time starts when it ends. The default client makes one attempt. So a throttle from it counts as refused for now. A client from `create_microvms_client` that retries a throttle turns it into an unknown outcome, as in item 8. The JavaScript worker classifies these errors the same way. 8. A call whose outcome is unknown is handled like a successful call, because the service may have suspended the MicroVM anyway. This covers a timeout, a conflict, a server error, and a dropped connection. An error after botocore retried an earlier attempt also counts as unknown, because the earlier attempt may have suspended the MicroVM. Only a client from `create_microvms_client` that retries can raise such an error. A call that the `before-send` handler stopped after a `resume` hook counts as unknown too. The handler can stop the first attempt, and then nothing was sent. An earlier attempt, if there was one, may have suspended the MicroVM. That `resume` hook has already ended the refusal. So the worker accepts jobs at once, and starts no new 30-second limit for that call. The refusal continues until the `resume` hook or the 30-second limit. The idle time then starts again, so the worker tries the call again. A `resume` hook during the call, or the end of a `suspend` hook's 30 seconds during the call, has already ended the refusal. The worker then accepts jobs at once, and tries the call again one idle time after the call and every job end. -9. The worker logs one warning for each failed call. The message names the outcome. The fields are, in order, `microvmId`, `error`, `outcome`, `suspendHookArrived`, `terminating`, `closed`, and `refusingJobs`. `microvmId` is the MicroVM ID of the call. `error` holds the name and the message of the error. `outcome` is `rejected`, `refusedForNow`, or `unknown`. `suspendHookArrived` states whether a `suspend` hook arrived since the idle time ended. `terminating` states whether a `terminate` hook came before the worker acted on the outcome. `closed` states whether `worker.close()` came before the worker acted on the outcome. `refusingJobs` states whether the worker answers jobs with HTTP 503 after it has acted on the outcome. The worker also logs two info lines about the call. It logs `suspending the idle MicroVM` before each call, with `microvmId`, `idleSeconds`, and `waitedSeconds`. It logs `stopping the suspend call: the MicroVM resumed during it` for the first `resume` hook during a call, with `microvmId`. The JavaScript worker logs the same stop line. +9. The worker logs one warning for each failed call. The message names the outcome. The fields are, in order, `microvmId`, `error`, `outcome`, `suspendHookArrived`, `terminating`, `closed`, `refusingJobs`, and `trigger`. `microvmId` is the MicroVM ID of the call. `error` holds the name and the message of the error. `outcome` is `rejected`, `refusedForNow`, or `unknown`. `suspendHookArrived` states whether a `suspend` hook arrived since the worker decided to make the call. `terminating` states whether a `terminate` hook came before the worker acted on the outcome. `closed` states whether `worker.close()` came before the worker acted on the outcome. `refusingJobs` states whether the worker answers jobs with HTTP 503 after it has acted on the outcome. `trigger` is `idle` for a call after the idle time, and `job` for a call after a `suspendAfter` job. The worker also logs two info lines about the call. It logs `suspending the idle MicroVM` before each call, with `microvmId`, `idleSeconds`, and `waitedSeconds`. It logs `stopping the suspend call: the MicroVM resumed during it` for the first `resume` hook during a call, with `microvmId`. The JavaScript worker logs the same stop line. -`worker.close()` stops the idle time. A closed worker never suspends its MicroVM again. A refusal that is already running still ends on the `resume` hook or after 30 seconds. +`worker.close()` stops the idle time, and drops a pending suspend after a job. A closed worker never suspends its MicroVM again. A refusal that is already running still ends on the `resume` hook or after 30 seconds. The worker counts only jobs. It cannot see background threads or processes, or inbound traffic that is not a job. A session with such work leaves `auto_suspend_on_idle` off, and the payload then has no `autoSuspendIdleSeconds`. The worker creates a `lambda-microvms` client for each suspend, and closes it after the call. Pass `create_microvms_client` to provide the client. The worker never closes a client that it returns. Give the client one attempt, as the default client has. The worker starts no attempt after a `resume` hook. A retry that starts before the `resume` hook arrives can still suspend the MicroVM after it resumed. A client that retries must report its retries in `ResponseMetadata.RetryAttempts` of the error that it raises, as botocore does. The worker otherwise counts a retried call as a first attempt. An earlier attempt of that call may have suspended the MicroVM. So a retried throttle would count as refused for now, and the worker would accept a job that then freezes with the MicroVM. +## Suspending after a job + +A session job can carry `"suspendAfter": true`. The durable function sets it for `vm.invoke` with `suspend_after=True`. The worker then suspends its own MicroVM after that job. This works in every session, with or without `autoSuspendIdleSeconds`. + +1. Only a JSON `true` turns it on. The worker ignores any other value, and does not reject the job. The `run` hook payload never carries the field. +2. The worker answers the job with HTTP 202 and the header `X-Microvm-Suspend-After: accepted`. A second delivery of the job gets the header when the first delivery carried the field. An older worker ignores the field and sends no header. The durable function then logs a warning, and the MicroVM keeps running. +3. When the job has ended, and its outcome reports have ended, the worker sets a pending suspend request. This happens whether the handler succeeded or failed. +4. The request fires when no job runs, no callback failure report runs, no `run` hook request is being handled, no `SuspendMicrovm` call runs, and the worker is not terminating, not closed, and not refusing jobs for a suspend. The worker checks these conditions again when a job ends, a report ends, a `run` hook request ends, and a call ends. So a request from a job that ends while another job runs fires when the last job ends. +5. Firing clears the request before the call starts. A `resume` hook never sets it again. A call after the idle time clears it too, because that call suspends the MicroVM as well. +6. Firing stops the idle time, when the session also suspends when idle. +7. The `terminate` hook and `worker.close()` drop the request. +8. The worker calls `SuspendMicrovm` with the MicroVM ID of the first accepted `run` hook with a payload. It never uses the MicroVM ID of a job request. A forged request could otherwise suspend another MicroVM of the same image. Without such a `run` hook, the worker logs the warning `cannot suspend the MicroVM after a job: the run hook named no MicroVM`, with `callbackId`, and does not suspend. +9. The call follows items 3 to 9 of [Suspending when idle](#suspending-when-idle): the refusal of jobs from the moment the request fires, one attempt and none after a `resume` hook, the outcome classification, and the 30-second refusal after the call. +10. A call that the service refused for now is not retried without `autoSuspendIdleSeconds`. With it, the worker tries again after the idle time, at least 30 seconds later, as in item 7 of [Suspending when idle](#suspending-when-idle). A rejected call is not retried. + +Before the call, the worker logs `suspending the MicroVM after a job` at info level, with `microvmId` and `callbackId`. A failed call logs the warning of item 9 of [Suspending when idle](#suspending-when-idle), with `trigger` set to `job`. + +Enable the `resume` hook in the image. Without it, the worker refuses the next job with HTTP 503 until 30 seconds after the call. The next job's heartbeat timeout counts that time. The MicroVM's execution role needs `lambda:SuspendMicrovm` on the image. Without it, each call is rejected and logged, and the MicroVM keeps running. + ## Suspend and terminate hooks Lambda waits for the answer to a `suspend` or `terminate` hook, up to the hook's timeout. The network and the execution role's credentials work while the hook runs. These facts come from probes in us-east-1, which the JavaScript SDK ran. @@ -169,7 +188,7 @@ The `runHookPayload` string holds: } ``` -`job` is absent when the job arrives over HTTP. An HTTP job request body has the job fields at the top level, with `version`, `region`, and an optional `microvmId`. `MicrovmJobRequest` holds those job fields in its `job` attribute. +`job` is absent when the job arrives over HTTP. An HTTP job request body has the job fields at the top level, with `version`, `region`, an optional `microvmId`, and an optional `suspendAfter`. `MicrovmJobRequest` holds those job fields in its `job` attribute, and `suspendAfter` in its `suspend_after` attribute. A document that does not match raises `InvalidRunHookPayloadError`. When the document named a non-empty callback ID, the error carries `callback_id` and `region`, so the worker can fail that callback at once. @@ -203,7 +222,7 @@ The worker logs to the standard library logger `aws_durable_execution_sdk_python The MicroVM's execution role needs `lambda:SendDurableExecutionCallbackSuccess`, `lambda:SendDurableExecutionCallbackFailure`, and `lambda:SendDurableExecutionCallbackHeartbeat` on the durable function's ARN (`arn:aws:lambda:::function::*`). -A session that suspends when idle also needs `lambda:SuspendMicrovm` on the MicroVM image (`arn:aws:lambda:::microvm-image:`). The action authorizes on the image, not on the MicroVM. So a MicroVM with this permission can suspend any MicroVM from the same image, even during a job. That job then waits until its heartbeat timeout or its timeout. Use one image per trust boundary. +A session that suspends when idle, or that sends a job with `suspendAfter`, also needs `lambda:SuspendMicrovm` on the MicroVM image (`arn:aws:lambda:::microvm-image:`). The action authorizes on the image, not on the MicroVM. So a MicroVM with this permission can suspend any MicroVM from the same image, even during a job. That job then waits until its heartbeat timeout or its timeout. Use one image per trust boundary. The worker needs boto3 1.43.35 or later, the first release with the `lambda-microvms` client. diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/payload.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/payload.py index 18e047ced..30a780ada 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/payload.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/payload.py @@ -235,6 +235,14 @@ class MicrovmJobRequest: """The MicroVM identifier. A job request can reach the worker before the ``run`` hook, which also carries the identifier. Until the ``run`` hook arrives, the worker takes it from the first job request that has one.""" + suspend_after: bool = False + """Whether the worker suspends its MicroVM after the job. + + Wire name: ``suspendAfter``. Only a JSON ``true`` turns it on. The + worker ignores any other value, as the JavaScript worker does. So a + document from either SDK has the same meaning in both workers. The + worker suspends the MicroVM named by the ``run`` hook, never the + ``microvmId`` of the request.""" @classmethod def from_dict(cls, body: object) -> MicrovmJobRequest: @@ -267,6 +275,7 @@ def from_dict(cls, body: object) -> MicrovmJobRequest: region=region, job=job, microvm_id=microvm_id, + suspend_after=body.get("suspendAfter") is True, ) diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py index ff7651a78..f8aa1aec3 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py @@ -112,6 +112,16 @@ _UNKNOWN_MICROVM_ID = "unknown" +SUSPEND_AFTER_HEADER = "X-Microvm-Suspend-After" +"""The header of a 202 answer that acknowledges ``"suspendAfter": true``. + +Its value is ``accepted``. The durable function reads it for a +``suspendAfter`` job. An older worker ignores the field and sends no header. +The durable function then logs a warning. The JavaScript worker sends the +same header.""" + +_SUSPEND_AFTER_ACK = {SUSPEND_AFTER_HEADER: "accepted"} + # At most 19 digits, so that int() stays under Python's limit on digits, and # any declared length still exceeds every body limit. _CONTENT_LENGTH = re.compile(r"[0-9]{1,19}") @@ -191,17 +201,57 @@ def _start_timer(delay: float, action: Callable[[], None]) -> Callable[[], None] } +@dataclass(frozen=True, kw_only=True) +class _SuspendTarget: + """The MicroVM that the worker's own SuspendMicrovm call suspends. + + It comes from the first accepted run hook with a payload, never from a + job request. SuspendMicrovm authorizes on the image. So a MicroVM ID + from a job body could suspend another MicroVM of the same image. + """ + + microvm_id: str + region: str + + @dataclass(frozen=True, kw_only=True) class _AutoSuspend: """The run hook's request to suspend the MicroVM when idle.""" idle_seconds: float - region: str - microvm_id: str - """The run hook's MicroVM ID. SuspendMicrovm uses it, never an ID from a + target: _SuspendTarget + """The run hook's MicroVM. SuspendMicrovm uses it, never an ID from a job request.""" +@dataclass(frozen=True, kw_only=True) +class _SuspendCall: + """One SuspendMicrovm call of the worker, as the locked block that + started the refusal decided it.""" + + target: _SuspendTarget + trigger: str + """``idle`` for the idle time, ``job`` for a ``suspendAfter`` job.""" + hooks_before: int + """The suspend hooks counted when the refusal started.""" + resumes_before: int + """The resume hooks counted when the refusal started.""" + message: str + """The info line before the call.""" + fields: Mapping[str, object] + """The fields of that line.""" + + +@dataclass(frozen=True, kw_only=True) +class _JobSuspend: + """A pending suspend after a ``suspendAfter`` job that fires now.""" + + callback_id: str + """The job whose end asked for the suspend.""" + call: _SuspendCall | None + """The call to make. ``None`` when no run hook named a MicroVM.""" + + @dataclass(kw_only=True) class _ActiveJob: """A job whose handler has not settled yet, with what the ``terminate`` @@ -347,15 +397,35 @@ class MicrovmWorkerListener: 9. The listener logs one warning for each failed call. The message names the outcome. The fields are, in order, ``microvmId``, ``error``, ``outcome``, ``suspendHookArrived``, ``terminating``, - ``closed``, and ``refusingJobs``. ``microvmId`` is the MicroVM ID of - the call. ``error`` holds the name and the message of the error. - ``outcome`` is ``rejected``, ``refusedForNow``, or ``unknown``. - ``suspendHookArrived`` states whether a ``suspend`` hook arrived - since the idle time ended. ``terminating`` states whether a - ``terminate`` hook came before the listener acted on the outcome. - ``closed`` states whether ``close()`` came before the listener acted - on the outcome. ``refusingJobs`` states whether the listener answers - jobs with HTTP 503 after it has acted on the outcome. + ``closed``, ``refusingJobs``, and ``trigger``. ``microvmId`` is the + MicroVM ID of the call. ``error`` holds the name and the message of + the error. ``outcome`` is ``rejected``, ``refusedForNow``, or + ``unknown``. ``suspendHookArrived`` states whether a ``suspend`` hook + arrived since the listener decided to make the call. ``terminating`` + states whether a ``terminate`` hook came before the listener acted + on the outcome. ``closed`` states whether ``close()`` came before the + listener acted on the outcome. ``refusingJobs`` states whether the + listener answers jobs with HTTP 503 after it has acted on the + outcome. ``trigger`` is ``idle`` or ``job``. + + A job request with ``"suspendAfter": true`` asks the listener to + suspend its MicroVM after the job, in any session: + + 1. The 202 answer carries ``X-Microvm-Suspend-After: accepted``. A + second delivery gets the header when the first one carried the + field. Only a JSON ``true`` counts. Any other value is ignored. + 2. When the job's outcome reports have ended, the listener sets a + pending request. The request fires when no job, failure report, + ``run`` hook request, or SuspendMicrovm call runs, and the listener + is not terminating, not closed, and not refusing jobs. Each of those + ends checks it again. + 3. Firing clears the request, stops the idle time, and makes the same + call as items 3 to 9, with the MicroVM ID of the first accepted + ``run`` hook with a payload. Without that hook, the listener logs a + warning and does not suspend. A ``resume`` hook never repeats the + request. The ``terminate`` hook and ``close()`` drop it. + 4. A call refused for now is retried only when the session suspends + when idle. When the image enables the ``suspend`` and ``terminate`` hooks, the listener acts on them: @@ -416,8 +486,12 @@ def __init__( self._locked_depth = 0 # The callback IDs of the jobs that run now. self._jobs: set[str] = set() - # The callback IDs of jobs that ended, oldest first. - self._finished: OrderedDict[str, None] = OrderedDict() + # The running jobs whose request carried "suspendAfter": true. + self._suspend_after_jobs: set[str] = set() + # The callback IDs of jobs that ended, oldest first. The value states + # whether the job's request carried "suspendAfter": true. A resent + # copy of the job then gets the same acknowledgement. + self._finished: OrderedDict[str, bool] = OrderedDict() # Callback failure reports in flight, from every place that fails a # callback. The terminate hook waits for them, because Lambda can end # the MicroVM as soon as it gets the answer. @@ -442,6 +516,14 @@ def __init__( self._reporting: set[str] = set() # Set from the run hook payload of a session that suspends when idle. self._auto_suspend: _AutoSuspend | None = None + # Set from the first accepted run hook that has a payload, in every + # session. Both triggers of a suspend call use it. + self._suspend_target: _SuspendTarget | None = None + # The callback ID of the suspendAfter job whose end asked for a + # suspend that has not started yet. The request fires when no work + # runs. A call that starts clears it, so a resume hook never repeats + # it. The terminate hook and close() clear it too. + self._suspend_requested: str | None = None # Each timer action carries the generation it was started with, and # acts only when it is still current. Cancelling a timer increments # the generation. @@ -521,22 +603,23 @@ def create( its own per job, from the default credential chain. The listener closes that client when the job ends. create_microvms_client: Creates the ``lambda-microvms`` client - that suspends this MicroVM when it is idle. The listener - calls it only when a session asks the worker to suspend when - idle, once per suspend. It never closes a client that this - function returns. Defaults to a client of its own per - suspend, from the default credential chain. The listener - closes that client after the call. The listener stops the - client's attempts once a ``resume`` hook arrived. Give the - client one attempt, as the default client has. A retry - that starts before the ``resume`` hook arrives can still - suspend the MicroVM after it resumed. A client that retries - must report its retries in ``ResponseMetadata.RetryAttempts`` - of the error that it raises. The listener otherwise counts a - retried call as a first attempt. An earlier attempt of that - call may have suspended the MicroVM. A retried throttle would - then count as refused for now, and the listener would accept - a job that then freezes with the MicroVM. + that suspends this MicroVM. The listener calls it only when + a session asks the worker to suspend when idle, or after a + ``suspendAfter`` job, once per suspend. It never closes a + client that this function returns. Defaults to a client of + its own per suspend, from the default credential chain. The + listener closes that client after the call. The listener + stops the client's attempts once a ``resume`` hook arrived. + Give the client one attempt, as the default client has. A + retry that starts before the ``resume`` hook arrives can + still suspend the MicroVM after it resumed. A client that + retries must report its retries in + ``ResponseMetadata.RetryAttempts`` of the error that it + raises. The listener otherwise counts a retried call as a + first attempt. An earlier attempt of that call may have + suspended the MicroVM. A retried throttle would then count + as refused for now, and the listener would accept a job that + then freezes with the MicroVM. heartbeat_interval_seconds: The heartbeat interval. Defaults to a third of the job's heartbeat timeout, at most 15 minutes. It is cut to a third of the job's heartbeat timeout. @@ -648,7 +731,8 @@ def idle(self, timeout: float | None = None) -> bool: ) def close(self) -> None: - """Stop the idle time. The listener then never suspends the MicroVM again. + """Stop the idle time, and drop a pending suspend after a job. The + listener then never suspends the MicroVM again. Running jobs continue. A refusal that a suspend started still ends on the ``resume`` hook or after 30 seconds. @@ -656,6 +740,7 @@ def close(self) -> None: with self._lock: self._closed = True self._stop_idle_time() + self._suspend_requested = None def _dispatch(self, exchange: _Exchange) -> None: path = exchange.path @@ -667,9 +752,11 @@ def _dispatch(self, exchange: _Exchange) -> None: try: self._handle_run_hook(exchange) finally: - with self._changed: + with self._locked(): self._run_hooks -= 1 self._changed.notify_all() + decision = self._take_job_suspend() + self._act_on_job_suspend(decision) return if hook in {"suspend", "terminate", "resume"} and not _is_lambda_hook_call( exchange @@ -784,6 +871,7 @@ def _terminate(self) -> None: with self._lock: self._terminating = True self._stop_idle_time() + self._suspend_requested = None stopped: list[tuple[str, _ActiveJob]] = [] for callback_id, job in self._active.items(): if job.cancelled.is_set(): @@ -850,15 +938,18 @@ def _handle_run_hook(self, exchange: _Exchange) -> None: if not later: self._run_hook_accepted = True self._microvm_id = parsed.microvm_id - if ( - payload is not None - and payload.auto_suspend_idle_seconds is not None - ): - self._auto_suspend = _AutoSuspend( - idle_seconds=payload.auto_suspend_idle_seconds, - region=payload.region, - microvm_id=parsed.microvm_id, + if payload is not None: + # Every session gets a target, also without auto-suspend. + # A suspendAfter job needs one. + target = _SuspendTarget( + microvm_id=parsed.microvm_id, region=payload.region ) + self._suspend_target = target + if payload.auto_suspend_idle_seconds is not None: + self._auto_suspend = _AutoSuspend( + idle_seconds=payload.auto_suspend_idle_seconds, + target=target, + ) microvm_id = self._microvm_id duplicate = job is not None and self._is_known(job.callback_id) new_job = None if later or duplicate else job @@ -995,9 +1086,21 @@ def _handle_job_request( terminating = not duplicate and self._terminating suspending = not duplicate and self._suspending if not (duplicate or terminating or suspending): - self._register_job(job.callback_id, request.region) + self._register_job( + job.callback_id, + request.region, + suspend_after=request.suspend_after, + ) + # A resent copy gets the acknowledgement that the first request + # got, because only the first request decides the suspend. + acknowledge = ( + self._carried_suspend_after(job.callback_id) + if duplicate + else request.suspend_after + ) + headers = _SUSPEND_AFTER_ACK if acknowledge else None if duplicate: - exchange.respond(202, {}) + exchange.respond(202, {}, headers) self._logger.warning( "duplicate job ignored", extra={"microvmId": microvm_id, "callbackId": job.callback_id}, @@ -1019,7 +1122,7 @@ def _handle_job_request( # retries. After a 202, the durable function does not deliver the job # again. self._start_job(handler, job, request.region) - exchange.respond(202, {}) + exchange.respond(202, {}, headers) def _reject(self, error: BaseException, exchange: _Exchange) -> None: """Answer 400 to a request that the worker could not accept. @@ -1052,7 +1155,16 @@ def _is_known(self, callback_id: str) -> bool: """ return callback_id in self._jobs or callback_id in self._finished - def _register_job(self, callback_id: str, region: str) -> None: + def _carried_suspend_after(self, callback_id: str) -> bool: + """Whether the request of a known job carried ``"suspendAfter": + true``. The caller holds the lock.""" + return callback_id in self._suspend_after_jobs or self._finished.get( + callback_id, False + ) + + def _register_job( + self, callback_id: str, region: str, *, suspend_after: bool = False + ) -> None: """Register a job that is about to start. The caller holds the lock. The job is registered for the terminate hook before its thread @@ -1061,6 +1173,8 @@ def _register_job(self, callback_id: str, region: str) -> None: handler. """ self._jobs.add(callback_id) + if suspend_after: + self._suspend_after_jobs.add(callback_id) self._active[callback_id] = _ActiveJob( region=region, cancelled=threading.Event() ) @@ -1082,12 +1196,16 @@ def _start_job( # The job never ran. A job request then answers 500, so the # durable function delivers the job again, and the retry may run # it. The run hook fails the callback instead, because Lambda - # sends it only once. + # sends it only once. The job never ran, so it asks for no + # suspend. A suspend that another job asked for can fire now. with self._locked(): self._jobs.discard(job.callback_id) + self._suspend_after_jobs.discard(job.callback_id) self._active.pop(job.callback_id, None) self._start_idle_time() self._changed.notify_all() + decision = self._take_job_suspend() + self._act_on_job_suspend(decision) raise def _run_job( @@ -1101,14 +1219,25 @@ def _run_job( extra={"callbackId": job.callback_id, "error": describe(error)}, ) finally: + # The job's report attempts have ended here, whether the handler + # succeeded or failed. So a suspendAfter job asks for the suspend + # now. The suspend runs after this block, so that no lock is + # held while it logs or starts its thread. with self._locked(): - self._jobs.discard(job.callback_id) - self._active.pop(job.callback_id, None) - self._finished[job.callback_id] = None + callback_id = job.callback_id + suspend_after = callback_id in self._suspend_after_jobs + self._jobs.discard(callback_id) + self._suspend_after_jobs.discard(callback_id) + self._active.pop(callback_id, None) + self._finished[callback_id] = suspend_after if len(self._finished) > MAX_FINISHED_JOBS: self._finished.popitem(last=False) self._start_idle_time() + if suspend_after and not (self._terminating or self._closed): + self._suspend_requested = callback_id self._changed.notify_all() + decision = self._take_job_suspend() + self._act_on_job_suspend(decision) def _run_job_and_report( self, handler: MicrovmJobHandler, job: MicrovmJobDocument, region: str @@ -1335,9 +1464,11 @@ def _send_failure( self._report_ended() def _report_ended(self) -> None: - with self._changed: + with self._locked(): self._reports -= 1 self._changed.notify_all() + decision = self._take_job_suspend() + self._act_on_job_suspend(decision) def _reporter_for(self, callback_id: str, region: str) -> CallbackReporter: # Each job gets its own client. A client kept across jobs would keep @@ -1472,20 +1603,120 @@ def _idle_time_ended(self, generation: int, waited_seconds: float) -> None: ): return self._cancel_idle_timer = None - self._suspending = True - self._suspend_calls += 1 + # waitedSeconds differs from idleSeconds after a refusal for now, + # because that idle time lasts at least SUSPEND_GRACE_SECONDS. + call = self._begin_suspend_call( + auto_suspend.target, + "idle", + "suspending the idle MicroVM", + { + "microvmId": auto_suspend.target.microvm_id, + "idleSeconds": auto_suspend.idle_seconds, + "waitedSeconds": waited_seconds, + }, + ) + self._start_suspend_thread( + call, "could not start the suspend of the idle MicroVM" + ) + + def _begin_suspend_call( + self, + target: _SuspendTarget, + trigger: str, + message: str, + fields: Mapping[str, object], + ) -> _SuspendCall: + """Start the refusal of jobs for a SuspendMicrovm call, and count the + call. The caller holds the lock, and starts the call after it + releases the lock. + + The call suspends the MicroVM either way. So a pending suspend after + a job is done by this call too, and it is cleared. + """ + self._suspending = True + self._suspend_calls += 1 + self._suspend_requested = None + self._stoppable_suspend = target.microvm_id + return _SuspendCall( + target=target, + trigger=trigger, # A suspend hook can arrive at any point from here on: while the - # thread starts, while the client is created, or during the call. - # Each path that ends the refusal compares with this count. - hooks_before = self._suspend_hooks + # thread starts, while the client is created, or during the + # call. Each path that ends the refusal compares with this count. + hooks_before=self._suspend_hooks, # A resume hook from here on ends the refusal. The call then # starts no further attempt. - resumes_before = self._resume_hooks - self._stoppable_suspend = auto_suspend.microvm_id + resumes_before=self._resume_hooks, + message=message, + fields=fields, + ) + + def _take_job_suspend(self) -> _JobSuspend | None: + """Decide whether the pending suspend after a ``suspendAfter`` job + fires now. + + The request fires when no job, failure report, run hook, or + SuspendMicrovm call runs, and the listener is not terminating, not + closed, and not refusing jobs for a suspend. Each event that can + change one of these calls this method: a job end, a report end, a + run hook end, and a call end. Firing clears the request, and stops + an idle timer that auto-suspend armed. + + The caller holds the lock through :meth:`_locked`. It passes the + result to :meth:`_act_on_job_suspend` after it releases the lock. So + no log line, ``describe()``, or caller code runs under the lock, and + no locking method runs inside the caller's locked block. + """ + callback_id = self._suspend_requested + if ( + callback_id is None + or self._jobs + or self._reports + or self._run_hooks + or self._terminating + or self._closed + or self._suspending + or self._suspend_calls + ): + return None + self._suspend_requested = None + target = self._suspend_target + if target is None: + return _JobSuspend(callback_id=callback_id, call=None) + self._stop_idle_time() + call = self._begin_suspend_call( + target, + "job", + "suspending the MicroVM after a job", + {"microvmId": target.microvm_id, "callbackId": callback_id}, + ) + return _JobSuspend(callback_id=callback_id, call=call) + + def _act_on_job_suspend(self, decision: _JobSuspend | None) -> None: + """Start the suspend that :meth:`_take_job_suspend` decided. The + caller does not hold the lock.""" + if decision is None: + return + if decision.call is None: + # A job request carries a MicroVM ID, but a forged request could + # name another MicroVM of the image. So only the run hook names + # the MicroVM to suspend. + self._logger.warning( + "cannot suspend the MicroVM after a job: the run hook named no MicroVM", + extra={"callbackId": decision.callback_id}, + ) + return + self._start_suspend_thread( + decision.call, "could not start the suspend of the MicroVM after a job" + ) + + def _start_suspend_thread(self, call: _SuspendCall, failure: str) -> None: + """Run a SuspendMicrovm call in its own thread. The caller started + the refusal and counted the call, and does not hold the lock.""" try: threading.Thread( target=self._suspend_self, - args=(auto_suspend, hooks_before, resumes_before, waited_seconds), + args=(call,), name="microvm-suspend", daemon=True, ).start() @@ -1493,35 +1724,29 @@ def _idle_time_ended(self, generation: int, waited_seconds: float) -> None: # No call was made. So jobs can run again at once, unless a # suspend hook started its own refusal meanwhile. self._logger.error( - "could not start the suspend of the idle MicroVM", - extra={"microvmId": auto_suspend.microvm_id, "error": describe(error)}, + failure, + extra={"microvmId": call.target.microvm_id, "error": describe(error)}, ) with self._locked(): - deferred = self._end_own_refusal(hooks_before) - self._suspend_call_ended() + deferred = self._end_own_refusal(call.hooks_before) + decision = self._suspend_call_ended() self._launch_deferred(deferred) + self._act_on_job_suspend(decision) - def _suspend_self( - self, - auto_suspend: _AutoSuspend, - hooks_before: int, - resumes_before: int, - waited_seconds: float, - ) -> None: + def _suspend_self(self, call: _SuspendCall) -> None: try: - self._call_suspend( - auto_suspend, hooks_before, resumes_before, waited_seconds - ) + self._call_suspend(call) finally: with self._locked(): - self._suspend_call_ended() + decision = self._suspend_call_ended() + self._act_on_job_suspend(decision) def _end_own_refusal( self, hooks_before: int ) -> list[tuple[MicrovmJobDocument, str]]: """End the refusal that the worker's own suspend started. - A suspend hook that arrived since the idle time ended means that + A suspend hook that arrived since the call was decided means that Lambda is freezing the MicroVM. Its refusal then stays. The caller holds the lock, and starts the returned jobs after it releases it. """ @@ -1529,10 +1754,14 @@ def _end_own_refusal( return [] return self._stop_refusing() - def _suspend_call_ended(self) -> None: + def _suspend_call_ended(self) -> _JobSuspend | None: """Count a SuspendMicrovm call as ended, and start the idle time that - a resume hook, a job, or a refusal for now asked for during it. The - caller holds the lock through :meth:`_locked`.""" + a resume hook, a job, or a refusal for now asked for during it. + + A suspend after a job that ended during the call can fire now. The + caller holds the lock through :meth:`_locked`, and passes the result + to :meth:`_act_on_job_suspend` after it releases the lock. + """ self._suspend_calls -= 1 # The thread or the client could fail before the call was sent. self._stoppable_suspend = None @@ -1540,24 +1769,21 @@ def _suspend_call_ended(self) -> None: self._idle_after_call = False self._start_idle_time(after_refusal=self._idle_after_call_floor) self._changed.notify_all() + return self._take_job_suspend() - def _call_suspend( - self, - auto_suspend: _AutoSuspend, - hooks_before: int, - resumes_before: int, - waited_seconds: float, - ) -> None: - target = auto_suspend.microvm_id + def _call_suspend(self, call: _SuspendCall) -> None: + target = call.target.microvm_id + hooks_before = call.hooks_before fields = {"microvmId": target} # The listener closes only a client that it created itself. A client # from create_microvms_client belongs to the caller. owns_client = self._create_microvms_client is None + region = call.target.region try: client = ( - create_microvms_client(auto_suspend.region) + create_microvms_client(region) if self._create_microvms_client is None - else self._create_microvms_client(auto_suspend.region) + else self._create_microvms_client(region) ) except Exception as error: # noqa: BLE001 # No call was made. So jobs can run again at once, unless a @@ -1570,23 +1796,14 @@ def _call_suspend( deferred = self._end_own_refusal(hooks_before) self._launch_deferred(deferred) return - # waitedSeconds differs from idleSeconds after a refusal for now, - # because that idle time lasts at least SUSPEND_GRACE_SECONDS. - self._logger.info( - "suspending the idle MicroVM", - extra={ - **fields, - "idleSeconds": auto_suspend.idle_seconds, - "waitedSeconds": waited_seconds, - }, - ) + self._logger.info(call.message, extra=dict(call.fields)) # The block after the client is closed decides whether the refusal # continues after an unknown outcome. The warning reads its fields in # that block. So for an unknown outcome, the warning comes after # close(). unknown_described: object | None = None try: - self._send_suspend(client, target, resumes_before) + self._send_suspend(client, target, call.resumes_before) except Exception as error: # noqa: BLE001 outcome = classify_suspend_error(error) # describe() calls str(error). A client from @@ -1611,21 +1828,21 @@ def _call_suspend( deferred = [] with self._locked(): # The service did nothing. So jobs can run again at once. - # A suspend hook that arrived since the idle time ended + # A suspend hook that arrived since the call was decided # means that the MicroVM is suspending anyway, so its # refusal stays. if self._suspend_hooks == hooks_before: deferred = self._stop_refusing() # A rejection, such as a missing permission, is not # retried until a job ends. A refusal for now is - # retried after the next idle time. That idle time - # lasts at least 30 seconds, because the same error - # can repeat on every call. + # retried after the next idle time, only with + # auto-suspend. That idle time lasts at least 30 + # seconds, because the same error can repeat on every + # call. A suspendAfter job asked for one call, so + # without auto-suspend nothing retries it. if outcome is SuspendOutcome.REFUSED_FOR_NOW: self._start_idle_time(after_refusal=True) - message, extra = self._failed_call_log( - outcome, described, target, hooks_before - ) + message, extra = self._failed_call_log(outcome, described, call) self._logger.warning(message, extra=extra) self._launch_deferred(deferred) return @@ -1649,7 +1866,7 @@ def _call_suspend( if unknown_described is None: return message, extra = self._failed_call_log( - SuspendOutcome.UNKNOWN, unknown_described, target, hooks_before + SuspendOutcome.UNKNOWN, unknown_described, call ) self._logger.warning(message, extra=extra) @@ -1657,7 +1874,7 @@ def _send_suspend( self, client: LambdaMicrovmsClient, target: str, resumes_before: int ) -> None: """Make the SuspendMicrovm call. Start no attempt after a resume hook - that arrived since the idle time ended.""" + that arrived since the call was decided.""" try: # The before-send handler calls resumed() inside botocore, in # this thread. It reads one int without the lock. So it runs no @@ -1678,14 +1895,15 @@ def _failed_call_log( self, outcome: SuspendOutcome, described: object, - microvm_id: str, - hooks_before: int, + call: _SuspendCall, ) -> tuple[str, dict[str, object]]: """Return the warning for a failed SuspendMicrovm call, and its fields. The message names only the outcome. The fields state whether a - suspend hook arrived since the idle time ended, whether the listener - is terminating or closed, and whether it refuses jobs. + suspend hook arrived since the call was decided, whether the + listener is terminating or closed, whether it refuses jobs, and what + started the call. ``trigger`` is the last field, so the first seven + keep their order. The caller holds the lock, and has already acted on the outcome in the same locked block. So the fields state what the listener does @@ -1699,14 +1917,15 @@ def _failed_call_log( """ name, message = _FAILED_CALL_WARNINGS[outcome] return message, { - "microvmId": microvm_id, + "microvmId": call.target.microvm_id, "error": described, "outcome": name, - "suspendHookArrived": self._suspend_hooks != hooks_before, + "suspendHookArrived": self._suspend_hooks != call.hooks_before, "terminating": self._terminating, "closed": self._closed, # The job path refuses a job with HTTP 503 in both states. "refusingJobs": self._terminating or self._suspending, + "trigger": call.trigger, } def _start_grace(self) -> None: @@ -2037,8 +2256,14 @@ def read_json(self, max_bytes: int) -> Any: msg = "request body is not valid JSON" raise InvalidRunHookPayloadError(msg) from error - def respond(self, status: int, body: Mapping[str, object]) -> None: - """Send the answer. A client that has gone away is logged.""" + def respond( + self, + status: int, + body: Mapping[str, object], + headers: Mapping[str, str] | None = None, + ) -> None: + """Send the answer, with ``headers`` added. A client that has gone + away is logged.""" self.responded = True data = json.dumps(body, separators=(",", ":")).encode() request = self._request @@ -2046,6 +2271,8 @@ def respond(self, status: int, body: Mapping[str, object]) -> None: request.send_response(status) request.send_header("Content-Type", "application/json") request.send_header("Content-Length", str(len(data))) + for name, value in (headers or {}).items(): + request.send_header(name, value) # The header also makes the base class close the connection. request.send_header("Connection", "close") request.end_headers() diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py index f1f0bb2d1..5310f0d15 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py @@ -64,6 +64,8 @@ "terminating", "closed", "refusingJobs", + # The last field, so the first seven keep their order. + "trigger", ] @@ -76,6 +78,7 @@ def assert_warning( suspend_hook_arrived: bool = False, terminating: bool = False, closed: bool = False, + trigger: str = "idle", ) -> None: """Assert that one warning for a failed SuspendMicrovm call has this outcome and these fields.""" @@ -89,6 +92,7 @@ def assert_warning( "terminating": terminating, "closed": closed, "refusingJobs": refusing_jobs, + "trigger": trigger, } @@ -100,6 +104,7 @@ def assert_one_warning( suspend_hook_arrived: bool = False, terminating: bool = False, closed: bool = False, + trigger: str = "idle", ) -> dict[str, Any]: """Assert that the worker logged exactly one warning, for one failed SuspendMicrovm call, with this outcome and these fields. Returns the @@ -115,6 +120,7 @@ def assert_one_warning( suspend_hook_arrived=suspend_hook_arrived, terminating=terminating, closed=closed, + trigger=trigger, ) return data @@ -308,6 +314,9 @@ def __call__(self, job_input: Any, context: MicrovmJobContext) -> Any: self.order.append(context.callback_id) self.contexts[context.callback_id] = context self._release(context.callback_id).wait(WAIT) + if isinstance(job_input, dict) and job_input.get("fail"): + msg = "the job failed" + raise RuntimeError(msg) return "done" def started(self) -> list[str]: @@ -1794,7 +1803,7 @@ def test_accepts_jobs_after_an_unknown_outcome_when_the_refusal_ended_during_the assert logger.messages("warning") == [] end_failed_call(h) - expected = { + expected: dict[str, Any] = { "suspend_hook_arrived": during != "resume", "refusing_jobs": False, } @@ -3453,3 +3462,455 @@ def create(region: str) -> Any: assert h.job("cb-1") == 503 h.hook("resume") assert h.job("cb-1") == 202 + + +# --------------------------------------------------------------------------- +# Suspend after a suspendAfter job +# --------------------------------------------------------------------------- + +ACK_HEADER = "X-Microvm-Suspend-After" +JOB_SUSPEND_INFO = "suspending the MicroVM after a job" +NO_TARGET_WARNING = ( + "cannot suspend the MicroVM after a job: the run hook named no MicroVM" +) + + +def post_job(h: Harness, callback_id: str, **fields: Any) -> tuple[int, str | None]: + """Send a job request, and return the status and the acknowledgement + header.""" + connection = http.client.HTTPConnection( + host_of(h.worker), h.worker.port, timeout=WAIT + ) + body = { + "version": 1, + "region": "us-east-1", + "callbackId": callback_id, + "input": {}, + **fields, + } + try: + connection.request( + "POST", + "/job", + body=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + ) + response = connection.getresponse() + response.read() + finally: + connection.close() + return response.status, response.getheader(ACK_HEADER) + + +def info_lines(logger: Any, message: str) -> list[dict[str, Any]]: + return [data for level, m, data in logger.lines if level == "info" and m == message] + + +def end_job(h: Harness, callback_id: str) -> None: + """Release a job, wait until it has ended, and let a suspend call that + its end started end too.""" + h.finish_job(callback_id) + h.settle() + + +@pytest.mark.parametrize( + ("value", "acknowledged"), + [(True, "accepted"), (False, None), ("true", None), (1, None), (None, None)], + ids=["true", "false", "string", "one", "null"], +) +def test_acknowledges_only_a_json_true_and_suspends_only_for_it( + start: Any, value: Any, acknowledged: str | None +) -> None: + h = start() + h.run_hook() + + assert post_job(h, "cb-1", suspendAfter=value) == (202, acknowledged) + end_job(h, "cb-1") + + assert h.microvms.suspended == (["mvm-1"] if acknowledged else []) + + +def test_acknowledges_no_job_without_the_field(start: Any) -> None: + h = start() + h.run_hook() + + assert post_job(h, "cb-1") == (202, None) + + +@pytest.mark.parametrize("finished", [False, True], ids=["running", "finished"]) +def test_acknowledges_a_resent_copy_when_the_first_request_carried_the_field( + start: Any, finished: bool +) -> None: + h = start() + h.run_hook() + h.microvms.failure = suspend_error("AccessDeniedException", 403) + assert post_job(h, "cb-1", suspendAfter=True) == (202, "accepted") + assert post_job(h, "cb-2") == (202, None) + if finished: + end_job(h, "cb-1") + end_job(h, "cb-2") + + # Only the first request decides the suspend. So a copy gets the same + # acknowledgement, whatever it carries. + assert post_job(h, "cb-1", suspendAfter=True) == (202, "accepted") + assert post_job(h, "cb-1") == (202, "accepted") + assert post_job(h, "cb-2", suspendAfter=True) == (202, None) + + +def test_sends_no_acknowledgement_with_a_refusal(start: Any) -> None: + h = start() + h.run_hook() + h.hook("suspend") + + assert post_job(h, "cb-1", suspendAfter=True) == (503, None) + + +def test_suspends_after_a_suspend_after_job_in_a_session_without_auto_suspend( + start: Any, logger: Any +) -> None: + h = start() + h.run_hook() + assert post_job(h, "cb-1", suspendAfter=True)[0] == 202 + h.advance(IDLE * 10) + assert h.microvms.suspended == [] + + end_job(h, "cb-1") + + assert h.microvms.suspended == ["mvm-1"] + assert info_lines(logger, JOB_SUSPEND_INFO) == [ + {"microvmId": "mvm-1", "callbackId": "cb-1"} + ] + # The refusal runs from the decision, as after an idle suspend. + assert h.job("cb-2") == 503 + h.hook("resume", LAMBDA_HEADERS) + assert h.job("cb-2") == 202 + + +def test_suspends_after_a_job_whose_handler_failed(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True, input={"fail": True}) + + end_job(h, "cb-1") + + assert h.lambda_client.recorded()[0]["kind"] == "failure" + assert h.microvms.suspended == ["mvm-1"] + + +def test_suspends_after_a_job_whose_report_failed(start: Any) -> None: + h = start() + h.run_hook() + h.lambda_client.completion_error = suspend_error("AccessDeniedException", 403) + post_job(h, "cb-1", suspendAfter=True) + + end_job(h, "cb-1") + + assert h.lambda_client.recorded() == [] + assert h.microvms.suspended == ["mvm-1"] + + +def test_suspends_only_when_the_report_attempts_have_ended(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + h.lambda_client.hold = True + + h.jobs.release("cb-1") + assert eventually(lambda: h.lambda_client.waiting() == 1) + assert h.microvms.suspended == [] + + h.lambda_client.finish() + end_job(h, "cb-1") + assert h.lambda_client.recorded() == [{"kind": "success", "callbackId": "cb-1"}] + assert h.microvms.suspended == ["mvm-1"] + + +def test_suspends_when_the_last_running_job_ends(start: Any, logger: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + post_job(h, "cb-2") + + end_job(h, "cb-1") + assert h.microvms.suspended == [] + + end_job(h, "cb-2") + assert h.microvms.suspended == ["mvm-1"] + assert info_lines(logger, JOB_SUSPEND_INFO) == [ + {"microvmId": "mvm-1", "callbackId": "cb-1"} + ] + + +def test_waits_for_a_background_failure_report(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + h.lambda_client.hold = True + # An invalid request names a callback, so the worker fails it in the + # background. + assert h.job("cb-bad", heartbeatTimeoutSeconds=0) == 400 + assert eventually(lambda: h.lambda_client.waiting() == 1) + h.lambda_client.hold = False + + end_job(h, "cb-1") + assert h.microvms.suspended == [] + + h.lambda_client.finish() + assert eventually(lambda: h.microvms.suspended == ["mvm-1"]) + + +def test_waits_for_a_run_hook_that_is_being_handled(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + later = slow_run_hook(h) + + end_job(h, "cb-1") + assert h.microvms.suspended == [] + + assert later.finish() == 200 + assert eventually(lambda: h.microvms.suspended == ["mvm-1"]) + + +def test_suspends_the_run_hook_microvm_not_the_one_of_the_job_request( + start: Any, +) -> None: + h = start() + # The job request arrives before the run hook, and names another MicroVM. + post_job(h, "cb-1", suspendAfter=True, microvmId="mvm-other") + h.run_hook() + + end_job(h, "cb-1") + + assert h.microvms.suspended == ["mvm-1"] + + +@pytest.mark.parametrize("run_hook", ["none", "without-payload"]) +def test_warns_and_does_not_suspend_when_no_run_hook_named_a_microvm( + start: Any, logger: Any, run_hook: str +) -> None: + h = start() + if run_hook == "without-payload": + assert h.post(f"{HOOK_PATH_PREFIX}run", {"microvmId": "mvm-1"}) == 200 + post_job(h, "cb-1", suspendAfter=True, microvmId="mvm-other") + + end_job(h, "cb-1") + + assert h.microvms.suspended == [] + assert [line for line in logger.lines if line[1] == NO_TARGET_WARNING] == [ + ("warning", NO_TARGET_WARNING, {"callbackId": "cb-1"}) + ] + # The request was consumed. A later job end does not warn again. + post_job(h, "cb-2") + end_job(h, "cb-2") + assert logger.messages("warning").count(NO_TARGET_WARNING) == 1 + assert h.job("cb-3") == 202 + + +def test_a_resume_hook_does_not_repeat_the_suspend(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + end_job(h, "cb-1") + assert h.microvms.suspended == ["mvm-1"] + + h.hook("resume", LAMBDA_HEADERS) + h.advance(GRACE * 10) + + assert h.microvms.suspended == ["mvm-1"] + # A job that ends after the resume asks for nothing either. + post_job(h, "cb-2") + end_job(h, "cb-2") + assert h.microvms.suspended == ["mvm-1"] + + +def test_does_not_retry_a_refusal_for_now_without_auto_suspend( + start: Any, logger: Any +) -> None: + h = start() + h.microvms.failure = suspend_error("ExpiredTokenException", 400) + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + + end_job(h, "cb-1") + h.advance(GRACE * 10) + + assert h.microvms.suspended == ["mvm-1"] + assert_one_warning(logger, "refusedForNow", refusing_jobs=False, trigger="job") + assert h.job("cb-2") == 202 + + +def test_retries_a_refusal_for_now_with_auto_suspend(start: Any) -> None: + h = start() + h.microvms.failure = suspend_error("ExpiredTokenException", 400) + h.run_hook(IDLE) + post_job(h, "cb-1", suspendAfter=True) + + end_job(h, "cb-1") + assert h.microvms.suspended == ["mvm-1"] + h.advance(GRACE - 0.001) + assert h.microvms.suspended == ["mvm-1"] + + h.advance(0.001) + assert h.microvms.suspended == ["mvm-1", "mvm-1"] + + +def test_stops_the_idle_timer_that_the_job_end_armed(start: Any, logger: Any) -> None: + # The job end arms the idle time, and fires the suspend in the same + # block. A rejected call does not retry until a job ends. So the stale + # timer must not fire one idle time later. + h = start() + h.microvms.failure = suspend_error("AccessDeniedException", 403) + h.run_hook(IDLE) + post_job(h, "cb-1", suspendAfter=True) + + end_job(h, "cb-1") + h.advance(IDLE * 10) + + assert h.microvms.suspended == ["mvm-1"] + assert_one_warning(logger, "rejected", refusing_jobs=False, trigger="job") + + +def test_logs_an_unknown_outcome_of_the_job_suspend_with_its_trigger( + start: Any, logger: Any +) -> None: + h = start() + h.microvms.failure = suspend_error("InternalServerException", 500) + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + + end_job(h, "cb-1") + + assert_one_warning(logger, "unknown", refusing_jobs=True, trigger="job") + assert h.job("cb-2") == 503 + + +def test_does_not_suspend_after_the_terminate_hook(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + post_job(h, "cb-2") + end_job(h, "cb-1") + + assert h.hook("terminate", LAMBDA_HEADERS) == 200 + h.jobs.release_all() + assert h.worker.idle(WAIT) + + assert h.microvms.suspended == [] + + +def test_does_not_suspend_after_close(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + post_job(h, "cb-2") + end_job(h, "cb-1") + + h.listener.close() + end_job(h, "cb-2") + + assert h.microvms.suspended == [] + assert h.listener._suspend_requested is None # noqa: SLF001 + + +def test_a_job_that_ends_after_close_asks_for_no_suspend(start: Any) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + + h.listener.close() + end_job(h, "cb-1") + + assert h.microvms.suspended == [] + assert h.listener._suspend_requested is None # noqa: SLF001 + + +def test_fires_a_pending_suspend_when_the_running_call_ends(start: Any) -> None: + # 1. The idle call runs, and a resume hook ends its refusal. + # 2. A suspendAfter job runs and ends while the call still runs. + # 3. The request waits for the call, and fires when it ends. + h = start() + h.microvms.hold = True + h.run_hook(IDLE) + h.advance(IDLE) + assert eventually(lambda: h.microvms.suspended == ["mvm-1"]) + h.hook("resume", LAMBDA_HEADERS) + assert post_job(h, "cb-1", suspendAfter=True)[0] == 202 + h.finish_job("cb-1") + assert h.microvms.suspended == ["mvm-1"] + + h.microvms.finish() + + assert eventually(lambda: h.microvms.suspended == ["mvm-1", "mvm-1"]) + h.settle() + + +def test_an_idle_call_does_the_pending_suspend_too(start: Any) -> None: + # A background report holds the request. The idle time ends first, and + # its call suspends the MicroVM. So the request is done, and it does not + # fire again after the resume. + h = start() + h.run_hook(IDLE) + post_job(h, "cb-1", suspendAfter=True) + h.lambda_client.hold = True + assert h.job("cb-bad", heartbeatTimeoutSeconds=0) == 400 + assert eventually(lambda: h.lambda_client.waiting() == 1) + h.lambda_client.hold = False + end_job(h, "cb-1") + assert h.microvms.suspended == [] + + h.advance(IDLE) + assert h.microvms.suspended == ["mvm-1"] + h.lambda_client.finish() + assert h.worker.idle(WAIT) + h.hook("resume", LAMBDA_HEADERS) + post_job(h, "cb-2") + end_job(h, "cb-2") + + assert h.microvms.suspended == ["mvm-1"] + + +def test_idle_waits_for_the_suspend_after_a_job(start: Any) -> None: + h = start() + h.microvms.hold = True + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + + h.finish_job("cb-1") + + assert not h.worker.idle(0.1) + h.microvms.finish() + assert h.worker.idle(WAIT) + + +def test_accepts_jobs_when_the_thread_of_the_job_suspend_cannot_start( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + failing_thread_start(monkeypatch, "microvm-suspend") + + h.finish_job("cb-1") + + assert "could not start the suspend of the MicroVM after a job" in ( + logger.messages("error") + ) + assert h.microvms.suspended == [] + monkeypatch.undo() + assert h.job("cb-2") == 202 + + +def test_a_suspend_after_job_whose_thread_cannot_start_asks_for_nothing( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + h = start() + h.run_hook() + failing_thread_start(monkeypatch, "microvm-job") + + assert post_job(h, "cb-1", suspendAfter=True) == (500, None) + + monkeypatch.undo() + assert h.listener._suspend_after_jobs == set() # noqa: SLF001 + assert h.listener._suspend_requested is None # noqa: SLF001 + assert h.microvms.suspended == [] diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/payload_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/payload_test.py index 2192fd220..dcd8dcf11 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/payload_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/payload_test.py @@ -284,3 +284,32 @@ def test_an_unusable_callback_id_is_not_attached_to_errors(callback_id): # endregion dataclasses and callback IDs + + +@pytest.mark.parametrize( + ("overrides", "suspend_after"), + [ + ({"suspendAfter": True}, True), + ({"suspendAfter": False}, False), + ({"suspendAfter": "true"}, False), + ({"suspendAfter": 1}, False), + ({"suspendAfter": None}, False), + ({"suspendAfter": {}}, False), + ({}, False), + ], + ids=["true", "false", "string", "one", "null", "object", "absent"], +) +def test_job_request_reads_only_json_true_as_suspend_after(overrides, suspend_after): + # Any other value is ignored, not rejected, as in the JavaScript worker. + parsed = MicrovmJobRequest.from_dict(job_request(**overrides)) + + assert parsed.suspend_after is suspend_after + + +def test_run_hook_payload_ignores_suspend_after(): + parsed = RunHookRequest.from_dict( + run_hook({"version": 1, "region": "us-east-1", "suspendAfter": True}) + ) + + assert parsed.payload is not None + assert not hasattr(parsed.payload, "suspend_after") From 7d0d7e40bd88aa9c108e9d3edc9eb92027242a47 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 25 Oct 2026 18:57:40 -0700 Subject: [PATCH 06/12] feat(extras): add suspend_after to vm.invoke A session job with suspend_after=True asks the worker to suspend the MicroVM after the job. The next job of the session checks the state and resumes the MicroVM. A job whose delivery gets a 502 with MICROVM_CONNECT_FAILED, or no answer, also checks the state. So a replay that lost the session's flag still resumes the MicroVM. --- .../examples-catalog.json | 13 + .../microvm/microvm_session_suspend_after.py | 125 ++++++ .../test/README.md | 1 + .../test/microvm/test_microvm_catalog.py | 9 + .../test/microvm/test_microvm_env.py | 8 +- .../test_microvm_session_suspend_after.py | 172 +++++++ .../README.md | 71 ++- .../microvm/config.py | 39 ++ .../microvm/errors.py | 29 +- .../microvm/job.py | 43 +- .../microvm/lifecycle.py | 161 ++++++- .../microvm/operation.py | 136 +++++- .../microvm/request.py | 106 ++++- .../microvm/session.py | 86 +++- .../tests/conftest.py | 47 +- .../tests/microvm/job_test.py | 50 +++ .../tests/microvm/lifecycle_test.py | 210 +++++++++ .../tests/microvm/request_test.py | 224 ++++++++++ .../tests/microvm/session_test.py | 421 +++++++++++++++++- .../tests/microvm/validation_test.py | 34 ++ .../tests/microvm/worker_composed_test.py | 362 ++++++++++++++- 21 files changed, 2268 insertions(+), 79 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend_after.py create mode 100644 packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend_after.py diff --git a/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json b/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json index 3f7c66350..e6c3ce461 100644 --- a/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json +++ b/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json @@ -834,6 +834,19 @@ "ExecutionTimeout": 900 }, "path": "./src/microvm/microvm_session_suspend.py" + }, + { + "name": "MicroVM Session Suspend After", + "description": "Runs two jobs in one Lambda MicroVM that its worker suspends after the first job", + "handler": "microvm_session_suspend_after.handler", + "integration": true, + "usesMicrovm": true, + "timeout": 120, + "durableConfig": { + "RetentionPeriodInDays": 7, + "ExecutionTimeout": 900 + }, + "path": "./src/microvm/microvm_session_suspend_after.py" } ] } diff --git a/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend_after.py b/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend_after.py new file mode 100644 index 000000000..acf42d44d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/microvm/microvm_session_suspend_after.py @@ -0,0 +1,125 @@ +"""Run two jobs in one Lambda MicroVM that the worker suspends after the first. + +The MicroVM runs the image that scripts/ensure_microvm_image.py builds from +microvm-image/. The image's routes /write and /read write and read a file in +the MicroVM. The image enables the resume hook. So the worker accepts the +next job as soon as the MicroVM resumes. + +The session does not suspend when idle. The first job sets suspend_after, so +the worker suspends the MicroVM after that job. A callback named gate holds +the second job. The caller completes the callback with +SendDurableExecutionCallbackSuccess to start the second job. +""" + +import os +from typing import Any + +from aws_durable_execution_sdk_python.config import CallbackConfig, Duration +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python_extras.microvm import ( + MicrovmInvokeOptions, + MicrovmSession, + MicrovmSessionConfig, + microvm_session, +) + + +# The callback named gate holds the second job until the caller completes it. +# If the caller never completes it, the callback times out after 5 minutes. +# The handler then raises an error, and the session terminates the MicroVM. +# So a failed caller leaves the MicroVM suspended for 5 minutes, not for the +# whole execution timeout. +GATE = CallbackConfig(timeout=Duration.from_minutes(5)) + +# The second job waits for the resume. The worker can refuse it with HTTP 503 +# for up to 30 seconds before that, and the request step retries for up to +# 60 seconds. So its heartbeat timeout is longer than both together. +RESUMED_JOB_HEARTBEAT_TIMEOUT = Duration.from_seconds(120) + + +def _require_microvm_env() -> tuple[str, str]: + """Return the image ARN and the MicroVM execution role ARN. + + The SAM template sets both variables on every example with + ``"usesMicrovm": true``. scripts/build_examples.py copies every file under + src/ into one directory. So a shared helper module would need a different + import in Lambda than in the local tests. So each MicroVM example defines + this function itself. + """ + image_arn = os.environ.get("MICROVM_IMAGE_ARN") + execution_role_arn = os.environ.get("MICROVM_EXECUTION_ROLE_ARN") + if not image_arn or not execution_role_arn: + msg = ( + "MICROVM_IMAGE_ARN and MICROVM_EXECUTION_ROLE_ARN must be set. " + "Build the image with scripts/ensure_microvm_image.py." + ) + raise RuntimeError(msg) + return image_arn, execution_role_arn + + +@durable_execution +def handler(event: dict[str, Any], context: DurableContext) -> Any: + image_arn, execution_role_arn = _require_microvm_env() + key: str = event["key"] + + def write_gate_read(vm: MicrovmSession, session_context: DurableContext) -> Any: + # The worker suspends the MicroVM when this job has ended and its + # result is reported. + written = vm.invoke( + session_context, + "write", + job_input={"key": key, "value": f"written-by-{key}"}, + options=MicrovmInvokeOptions( + path="/write", + timeout=Duration.from_minutes(5), + # A job that the MicroVM stops answering fails after this + # time, instead of after the timeout. + heartbeat_timeout=Duration.from_seconds(30), + suspend_after=True, + ), + ) + + # 1. The first job has ended before the gate starts. + # 2. The second job starts only after the caller completes the gate. + # 3. So no job runs while the gate is pending. + # 4. So the worker suspends the MicroVM after the first job. + # 5. So a caller that sees the MicroVM SUSPENDED while the gate is + # pending sees a suspend between the two jobs. + # The callback result carries no data. So the handler ignores it. + session_context.create_callback(name="gate", config=GATE).result() + + # The first job set suspend_after. So the request step calls + # GetMicrovm, finds the MicroVM SUSPENDED, and calls ResumeMicrovm + # before it sends the job. A suspended MicroVM keeps its files. So + # this job reads the file that the first job wrote. + read = vm.invoke( + session_context, + "read", + job_input={"key": key}, + options=MicrovmInvokeOptions( + path="/read", + timeout=Duration.from_minutes(5), + heartbeat_timeout=RESUMED_JOB_HEARTBEAT_TIMEOUT, + ), + ) + + return { + "sessionMicrovmId": vm.microvm_id, + "writtenBy": written["microvmId"], + "readBy": read["microvmId"], + "value": read["value"], + } + + # The session launches one MicroVM, and terminates it when the handler + # returns. It does not suspend when idle. + return microvm_session( + context, + "suspend-after-session", + config=MicrovmSessionConfig( + image_identifier=image_arn, + execution_role_arn=execution_role_arn, + timeout=Duration.from_minutes(10), + ), + handler=write_gate_read, + ) diff --git a/packages/aws-durable-execution-sdk-python-examples/test/README.md b/packages/aws-durable-execution-sdk-python-examples/test/README.md index 5c7535367..ec4ab5079 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/README.md +++ b/packages/aws-durable-execution-sdk-python-examples/test/README.md @@ -155,6 +155,7 @@ def test_my_example(durable_runner): - Pass a cloud-test input as an object, such as a `dict`. The cloud runner encodes the input with `json.dumps`. So a string that already holds JSON is encoded again, and the handler receives a `str` instead of a `dict`. - Start a MicroVM test with `run_async`, and read its result with `wait_for_result`. The synchronous `run` uses botocore's 60-second read timeout, and botocore sends the Invoke again after a read timeout. Each resend starts a new execution, which launches its own MicroVM. - `test_microvm_session_suspend.py` checks that an idle session MicroVM suspends and resumes. The function's role needs `lambda:GetMicrovm` and `lambda:ResumeMicrovm`, and the MicroVM execution role needs `lambda:SuspendMicrovm` on the image. +- `test_microvm_session_suspend_after.py` checks that the worker suspends a session MicroVM after a job with `suspend_after=True`, in a session without auto-suspend, and that the next job resumes it. It needs the same permissions as `test_microvm_session_suspend.py`. ### Cloud-Only Tests `@pytest.mark.durable_execution(..., cloud_only=True)` marks a test that needs a service the local runner cannot provide, such as a MicroVM. Local mode skips the test with the reason `cloud_only: runs only against the deployed function`. Cloud mode runs it. diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_catalog.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_catalog.py index c7adcd722..58f4d678d 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_catalog.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_catalog.py @@ -35,6 +35,15 @@ def test_the_suspend_session_function_timeout_covers_the_http_retry_window(): assert entry["timeout"] >= HTTP_DELIVERY_MIN_TIMEOUT_SECONDS +def test_the_suspend_after_session_function_timeout_covers_the_http_retry_window(): + # Every session job goes over HTTP. The wait for a resume counts toward + # the retry window. So the same minimum applies. + entry = catalog_entry("MicroVM Session Suspend After") + + assert entry["usesMicrovm"] is True + assert entry["timeout"] >= HTTP_DELIVERY_MIN_TIMEOUT_SECONDS + + def test_the_run_job_function_uses_the_global_timeout(): # The input is small, so the job goes in the run hook. The run hook # delivery adds no Lambda time for the boot. diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_env.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_env.py index 31b3d5e7f..e8fae35e9 100644 --- a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_env.py +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_env.py @@ -11,10 +11,16 @@ microvm_run_job, microvm_session, microvm_session_suspend, + microvm_session_suspend_after, ) -EXAMPLES = [microvm_run_job, microvm_session, microvm_session_suspend] +EXAMPLES = [ + microvm_run_job, + microvm_session, + microvm_session_suspend, + microvm_session_suspend_after, +] IMAGE_ARN = "arn:aws:lambda:us-west-2:123456789012:microvm-image:img" ROLE_ARN = "arn:aws:iam::123456789012:role/microvm-execution" diff --git a/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend_after.py b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend_after.py new file mode 100644 index 000000000..0bb51a23b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/test/microvm/test_microvm_session_suspend_after.py @@ -0,0 +1,172 @@ +"""Cloud test for the MicroVM session example that suspends after a job. + +The local runner cannot launch a MicroVM. So the test is cloud-only. +""" + +import uuid + +import pytest +from aws_durable_execution_sdk_python.execution import InvocationStatus +from aws_durable_execution_sdk_python.lambda_service import ( + OperationStatus, + OperationType, +) + + +# The examples package does not depend on the extras. So an environment +# without the extras skips this module instead of failing to import the +# example. +pytest.importorskip("aws_durable_execution_sdk_python_extras") + +from src.microvm import microvm_session_suspend_after # noqa: E402 +from test.conftest import deserialize_operation_payload # noqa: E402 +from test.microvm.microvm_helpers import ( # noqa: E402 + expect_microvm_terminated, + fail_callback_after_test_failure, + find_launched_microvm_id, + wait_for_microvm_state, +) + + +def tree(operations): + """Return (name, type, status, children) for each operation, in order.""" + return [ + ( + operation.name, + operation.operation_type, + operation.status, + tree(getattr(operation, "child_operations", None) or []), + ) + for operation in operations + ] + + +SUCCEEDED = OperationStatus.SUCCEEDED + + +@pytest.mark.example +@pytest.mark.durable_execution( + handler=microvm_session_suspend_after.handler, + lambda_function_name="MicroVM Session Suspend After", + cloud_only=True, +) +def test_suspends_the_microvm_after_the_first_job_and_resumes_it_for_the_next( + durable_runner, +): + key = str(uuid.uuid4()) + + with durable_runner: + # A synchronous Invoke returns only when the execution ends. The cloud + # runner's boto3 client reads for at most 60 seconds, and it sends the + # Invoke again after a read timeout. Each resend starts a new + # execution, which launches its own MicroVM. So the test starts the + # execution asynchronously and polls for its result. + execution_arn = durable_runner.run_async(input={"key": key}, timeout=300) + + # The test reads the MicroVM state while the gate is pending. So it + # reads the MicroVM ID from the launch step in the execution history. + microvm_id = find_launched_microvm_id( + execution_arn, "suspend-after-session.launch" + ) + + callback_id = None + gate_completed = False + try: + # The cloud runner reads the callback ID from the CallbackStarted + # event named gate. It returns the ID only after the invocation + # that created the callback has completed. + callback_id = durable_runner.wait_for_callback( + execution_arn=execution_arn, name="gate", timeout=120 + ) + + # 1. The gate starts after the first job ends. + # 2. The second job starts only after the test completes the gate. + # 3. The session does not suspend when idle. So only the first + # job's suspend_after can suspend the MicroVM. + # 4. So a SUSPENDED state read now lies between the two jobs, and + # comes from suspend_after. + # The poll lasts at most 150 seconds. The gate times out after 5 + # minutes. So the gate stays pending for the whole poll. The + # state can also pass through SUSPENDING. So the test requires + # only that SUSPENDED is seen once. + wait_for_microvm_state(microvm_id, "SUSPENDED") + + durable_runner.send_callback_success(callback_id=callback_id) + gate_completed = True + finally: + # A failure before the gate completes leaves the execution waiting + # for the gate, with the MicroVM still there. A callback failure + # makes the session terminate the MicroVM now, instead of after + # the gate's timeout. If the test fails before it has the callback + # ID, the gate times out after 5 minutes. The session then + # terminates the MicroVM. + if callback_id is not None and not gate_completed: + fail_callback_after_test_failure(durable_runner, callback_id) + + result = durable_runner.wait_for_result( + execution_arn=execution_arn, timeout=300 + ) + + assert result.status is InvocationStatus.SUCCEEDED, result.error + output = deserialize_operation_payload(result.result) + # 1. GetMicrovm reported the MicroVM SUSPENDED while the gate was pending. + # 2. The second job read the file that the first job wrote. + # 3. A suspended MicroVM keeps its files. + # So both jobs ran in the same MicroVM, before and after the suspend. + assert output["sessionMicrovmId"] == microvm_id + assert output["writtenBy"] == microvm_id + assert output["readBy"] == microvm_id + assert output["value"] == f"written-by-{key}" + + expect_microvm_terminated(microvm_id) + + # The tree is the tree of the suspend example. The state check and the + # resume run inside the second request step. So the tree shows them only + # as a request step that succeeded after the suspend. + assert tree(result.operations) == [ + ( + "suspend-after-session", + OperationType.CONTEXT, + SUCCEEDED, + [ + ("suspend-after-session.session", OperationType.STEP, SUCCEEDED, []), + ("suspend-after-session.launch", OperationType.STEP, SUCCEEDED, []), + ( + "write", + OperationType.CONTEXT, + SUCCEEDED, + [ + ("write.callback", OperationType.CALLBACK, SUCCEEDED, []), + ("write.request", OperationType.STEP, SUCCEEDED, []), + ], + ), + ("gate", OperationType.CALLBACK, SUCCEEDED, []), + ( + "read", + OperationType.CONTEXT, + SUCCEEDED, + [ + ("read.callback", OperationType.CALLBACK, SUCCEEDED, []), + ("read.request", OperationType.STEP, SUCCEEDED, []), + ], + ), + ( + "suspend-after-session.terminate", + OperationType.STEP, + SUCCEEDED, + [], + ), + ], + ) + ] + + # The session sets neither auto_suspend_on_idle nor an idle policy. So the + # launch records no resume mode. suspend_after is an option of one job, + # and the launch checkpoint does not record it. + session = result.get_context("suspend-after-session") + launch = deserialize_operation_payload( + session.get_step("suspend-after-session.launch").result + ) + assert launch["delivery"] == "http" + assert launch["microvmId"] == microvm_id + assert launch["resumeBeforeDelivery"] is False diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index f66c5ac64..7ae933303 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -100,7 +100,7 @@ Every MicroVM failure is a `MicrovmError`. Catch the base class for all of them, |---|---| | `MicrovmLaunchError` | `RunMicrovm` failed after all retries, or the launch request was not valid, for example with an empty `image_identifier` or `execution_role_arn`. A MicroVM can still run the job. The list below the table says when. | | `MicrovmDeliveryError` | HTTP delivery failed after all retries, or the route answered a 4xx status such as 404, or a redirect. | -| `MicrovmNotRunningError` | A `MicrovmDeliveryError`: the MicroVM is terminating, terminated, or removed. It is not retried. A session that checks the MicroVM state before each job reports it. Any job over HTTP also reports it when `CreateMicrovmAuthToken` reports the MicroVM removed. See [Suspending an idle MicroVM](#suspending-an-idle-microvm). | +| `MicrovmNotRunningError` | A `MicrovmDeliveryError`: the MicroVM is terminating, terminated, or removed. It is not retried. A session that checks the MicroVM state before each job reports it, and so does a session job that checks the state after a sign of a suspended MicroVM. Any job over HTTP also reports it when `CreateMicrovmAuthToken` reports the MicroVM removed. See [Suspending an idle MicroVM](#suspending-an-idle-microvm) and [Suspending after a job](#suspending-after-a-job). | | `MicrovmJobFailedError` | The job handler in the MicroVM failed, or the MicroVM sent a result that is not valid JSON. The message includes the error type that the MicroVM reported. | | `MicrovmTimeoutError` | No result arrived within `timeout`, or no heartbeat within `heartbeat_timeout`. A MicroVM that crashed ends this way. | @@ -260,7 +260,7 @@ A MicroVM lives at most its `timeout` plus 5 minutes, and never longer than 8 ho - For a session that launched with `auto_suspend_on_idle=True`, or with an `idle_policy`, the request step calls `GetMicrovm` before each job. A MicroVM that is terminating, terminated, or removed then fails the job at once with `MicrovmNotRunningError`, a subclass of `MicrovmDeliveryError`. A retry finds the same state, so the job is not retried. The message ends with the request ID of the call that reported the state, for example `GetMicrovm request ID: .`. - Every session, and every `microvm` job over HTTP, calls `CreateMicrovmAuthToken` before the first request, and again after a 401 or 403. The API model lists `ResourceNotFoundException` for that call. That answer fails the job at once with `MicrovmNotRunningError`, whose message ends with `CreateMicrovmAuthToken request ID: .`. The live service's answer for a removed MicroVM is not verified. -- A session that launched with neither setting makes no `GetMicrovm` call. When `CreateMicrovmAuthToken` returns a token, the request goes to the endpoint. A terminated MicroVM's endpoint answers 502, and the request step retries a 502 as a MicroVM that is still starting. So the job fails with a plain `MicrovmDeliveryError`, after the retries of both tiers. +- A session that launched with neither setting makes no `GetMicrovm` call before delivery, unless an earlier job of the handle set `suspend_after`. When `CreateMicrovmAuthToken` returns a token, the request goes to the endpoint. A terminated MicroVM's endpoint answers 502. Whether that 502 carries `x-aws-proxy-error: MICROVM_CONNECT_FAILED` is not verified. With that header, the check of [Suspending after a job](#suspending-after-a-job) reports the state, and the job fails with `MicrovmNotRunningError`. Without it, the request step retries a 502 as a MicroVM that is still starting. So the job fails with a plain `MicrovmDeliveryError`, after the retries of both tiers. Suspending is off by default, because it needs more permissions and changes what the MicroVM does between jobs: @@ -270,9 +270,71 @@ Suspending is off by default, because it needs more permissions and changes what The worker must support `auto_suspend_on_idle`. The JavaScript worker [`@aws/durable-execution-sdk-js-microvm-worker`](https://github.com/aws/aws-durable-execution-sdk-js/tree/main/packages/aws-durable-execution-sdk-js-microvm-worker) and the Python worker [`aws-durable-execution-sdk-python-microvm-worker`](../aws-durable-execution-sdk-python-microvm-worker) both do. `idle_policy` works with both workers. Enable the `suspend` and `resume` hooks in the image. The worker then refuses jobs while the MicroVM freezes, and accepts them again as soon as it resumes. Without the `resume` hook, the first job after a resume can wait a little over 30 seconds, until the worker's refusal ends and the request step delivers again. +### Suspending after a job + +`MicrovmInvokeOptions(suspend_after=True)` asks the worker to suspend the MicroVM after this job. It is an option of one job, not of the session. `microvm` raises `ValidationError` for it, because `microvm` terminates its MicroVM after its one job. A value that is not a `bool` raises `ValidationError` too, before the job's durable operations start. + +```python +build = vm.invoke( + ctx, + "build", + job_input=event, + options=MicrovmInvokeOptions( + timeout=Duration.from_minutes(15), suspend_after=True + ), +) +approval = ctx.create_callback(name="approval").result() # a long wait +vm.invoke( + ctx, + "deploy", + job_input=build, + options=MicrovmInvokeOptions( + timeout=Duration.from_minutes(10), + heartbeat_timeout=Duration.from_seconds(120), + ), +) +``` + +What happens: + +1. The job request carries `"suspendAfter": true`. The job document stays at version 1, and the field is omitted for every other job. +2. The worker acknowledges the field with the header `X-Microvm-Suspend-After: accepted` in its 202 answer. +3. The worker suspends the MicroVM when the job has ended and its outcome is reported, and no other job runs. A job that ends while another job runs leaves the request pending. The worker then suspends the MicroVM when the last running job ends. +4. The worker suspends the MicroVM that its `run` hook named. It never uses the MicroVM ID of a job request. +5. Each later job of the session handle checks the MicroVM state before delivery, and resumes a suspended MicroVM, as with `auto_suspend_on_idle`. + +The setting works with or without `auto_suspend_on_idle`. With both, the worker suspends right after the job, and the idle time applies after each resume. + +The record that a job set `suspend_after` is kept in memory only. A replay of a completed job, `map` item, `parallel` branch, or child context does not run its `vm.invoke` again. So the next job can skip the state check of a suspended MicroVM. For that case, every session job checks the state after a sign of a suspended MicroVM: + +1. The sign is an HTTP 502 with `x-aws-proxy-error: MICROVM_CONNECT_FAILED`, or a request without an answer. A booting MicroVM gives neither sign, because its endpoint holds the request until the worker listens. So a session that never suspends makes no extra call. +2. The request step calls `GetMicrovm` once. It resumes a `SUSPENDING` or `SUSPENDED` MicroVM, waits for `RUNNING`, and sends the job again. +3. The check runs again after each failed request, at most once per 2 seconds. +4. A denied or failed `GetMicrovm` logs one warning per job: `GetMicrovm failed, so the operation could not check whether the MicroVM is suspended. The request step keeps sending the job.` The request step then keeps sending the job. After a denied call, the final `MicrovmEndpointUnavailableError` message adds `GetMicrovm was denied, so the operation could not check whether the MicroVM is suspended.` +5. A denied `ResumeMicrovm` for a suspending or suspended MicroVM fails the job at once with `MicrovmDeliveryError`. Its message names `lambda:ResumeMicrovm`. +6. A terminating or terminated MicroVM fails the job with `MicrovmNotRunningError`. + +A session with `auto_suspend_on_idle` or an `idle_policy`, and a handle after a `suspend_after` job, keep the state check before delivery. They also keep the check after any 502, 503, or 504, or a request without an answer, and its errors: a denied `GetMicrovm` fails the job. + +Use `suspend_after` only on the job before a long wait, such as a durable wait or a callback: + +- Not on back-to-back jobs. Each suspend writes a snapshot, and each resume adds about 1 second before the next job. +- Not on `map` items. A request waits until no job runs. So a `map` with `suspend_after` on its items suspends and resumes between items that do not overlap. +- Not on the last job. The session terminates the MicroVM after its handler returns, so the snapshot is wasted. + +It needs: + +1. The `resume` hook in the image. Without it, the worker refuses the next job with HTTP 503 for up to 30 seconds after the suspend. +2. A `heartbeat_timeout` on the next job of at least its `retry_window` plus 30 seconds. The heartbeat timeout starts before delivery, so it includes the resume and the refusal. +3. `lambda:GetMicrovm` and `lambda:ResumeMicrovm` in the function's role. +4. `lambda:SuspendMicrovm` on the image in the MicroVM's role. That permission lets code in any MicroVM of the image suspend any other. So use one image per trust boundary. +5. A worker that supports the field. An older worker ignores it, and the MicroVM keeps running. The session then logs one warning: `the worker in MicroVM does not support suspendAfter. The MicroVM keeps running after the job.` + +A suspend that the service refuses for now, such as a throttle or expired credentials, is not retried without `auto_suspend_on_idle`. The MicroVM then keeps running until a later suspend. With `auto_suspend_on_idle`, the worker tries again after the idle time, as it does for an idle suspend. + ### Session permissions -A session needs the permissions of `microvm` with HTTP delivery: `lambda:RunMicrovm`, `lambda:TerminateMicrovm`, `lambda:CreateMicrovmAuthToken`, `lambda:PassNetworkConnector`, and `iam:PassRole`. Suspending adds the permissions listed above. +A session needs the permissions of `microvm` with HTTP delivery: `lambda:RunMicrovm`, `lambda:TerminateMicrovm`, `lambda:CreateMicrovmAuthToken`, `lambda:PassNetworkConnector`, and `iam:PassRole`. Suspending adds the permissions listed above. Grant `lambda:GetMicrovm` to every session's function role. Without it, the check after a sign of a suspended MicroVM logs a warning, and the job keeps its retries. ### Subtypes @@ -342,7 +404,8 @@ The MicroVM operations handle the network, the service, and the worker lifecycle - The idle time is at most the session `timeout`. A value that breaks any one condition raises `ValidationError` when `microvm_session` is called, before any durable operation starts. See [Suspending an idle MicroVM](#suspending-an-idle-microvm). -4. Change `auto_suspend_on_idle` or `idle_policy` only through a new version, as in [Deploys and versions](#deploys-and-versions). The launch records the resume mode, and a deploy changes it only for sessions that launch after it. Every job of a session that launched with suspending needs `lambda:GetMicrovm` and `lambda:ResumeMicrovm`, also after a deploy turns suspending off. So remove them from the function's role only as [Suspending an idle MicroVM](#suspending-an-idle-microvm) describes. +4. Set `suspend_after` only on the job before a long wait, and give the next job a `heartbeat_timeout` of at least its `retry_window` plus 30 seconds. On back-to-back jobs and `map` items, each suspend and resume adds latency and a snapshot. A shorter heartbeat timeout can fail the next job with `MicrovmTimeoutError` while the MicroVM resumes. Keep `lambda:GetMicrovm` and `lambda:ResumeMicrovm` in the function's role after a deploy that removes `suspend_after`, until every session that ran such a job has ended. See [Suspending after a job](#suspending-after-a-job). +5. Change `auto_suspend_on_idle` or `idle_policy` only through a new version, as in [Deploys and versions](#deploys-and-versions). The launch records the resume mode, and a deploy changes it only for sessions that launch after it. Every job of a session that launched with suspending needs `lambda:GetMicrovm` and `lambda:ResumeMicrovm`, also after a deploy turns suspending off. So remove them from the function's role only as [Suspending an idle MicroVM](#suspending-an-idle-microvm) describes. ### Retry strategies diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py index 18a8bd10e..a9b2d0c9e 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py @@ -352,4 +352,43 @@ class MicrovmInvokeOptions(MicrovmRequestConfig): Set it on every job. A MicroVM that ends while suspended, or whose image disables the ``terminate`` hook, reports nothing. The heartbeat timeout is then the only way that the job fails soon after the MicroVM ends. + + A job after a :attr:`suspend_after` job waits for the resume of the + MicroVM. The worker can refuse it with HTTP 503 for up to 30 seconds + before that. So set the heartbeat timeout of such a job to at least the + ``retry_window`` plus 30 seconds. + """ + + suspend_after: bool = False + """Whether the worker suspends the MicroVM after this job. Defaults to + ``False``. + + The worker suspends the MicroVM when the job has ended and its outcome is + reported, and no other job runs. A job that ends while another job runs + leaves the request pending. The worker then suspends the MicroVM when + the last running job ends. The next job of the session resumes the + MicroVM before it is delivered. The setting works with or without + :attr:`MicrovmSessionConfig.auto_suspend_on_idle`. + + Set it on the job before a long wait, such as a durable wait or a + callback. Do not set it on back-to-back jobs, on ``map`` items, or on + the last job of the session. Each suspend and resume costs a snapshot + and about 1 second of latency for the next job. The session terminates + the MicroVM after its last job anyway. + + It needs more than a session without suspending: + + 1. The image enables the ``resume`` hook. Without it, the next job waits + up to 30 seconds of HTTP 503 answers. + 2. The function's role has ``lambda:GetMicrovm`` and + ``lambda:ResumeMicrovm``. + 3. The MicroVM's role has ``lambda:SuspendMicrovm`` on the image. + 4. The worker supports the setting. An older worker ignores it, and the + MicroVM keeps running. The session then logs one warning. + + A suspend that the service refuses for now is not retried, unless + :attr:`MicrovmSessionConfig.auto_suspend_on_idle` is ``True``. Only + :meth:`MicrovmSession.invoke` accepts it. :func:`microvm` raises + ``ValidationError`` for it. A value that is not a ``bool`` raises + ``ValidationError`` too. """ diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/errors.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/errors.py index 18cd84287..5730c3dc1 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/errors.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/errors.py @@ -122,12 +122,16 @@ class MicrovmNotRunningError(MicrovmDeliveryError): the same state, so the job is not retried. To continue a session, start a new session. - Two checks raise it: + Three checks raise it: 1. A session that checks the MicroVM state before each job. It checks - when ``auto_suspend_on_idle`` is ``True``, or when an ``idle_policy`` - is set. ``GetMicrovm`` or ``ResumeMicrovm`` then reports the state. - 2. Every job over HTTP, of a session or of :func:`microvm`. The request + when ``auto_suspend_on_idle`` is ``True``, when an ``idle_policy`` + is set, and after a job with ``suspend_after``. ``GetMicrovm`` or + ``ResumeMicrovm`` then reports the state. + 2. Every session job, after a request without an answer, or after an + HTTP 502 with ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``. The + request step then calls ``GetMicrovm`` once. + 3. Every job over HTTP, of a session or of :func:`microvm`. The request step calls ``CreateMicrovmAuthToken`` before the first request, and again after a 401 or 403. That call can answer ``ResourceNotFoundException``. @@ -137,14 +141,15 @@ class MicrovmNotRunningError(MicrovmDeliveryError): records only the message of the failure. So a support case can quote the request ID from the message, on the first run and on replay. - Without the state check, a session makes no ``GetMicrovm`` call. The - request step still calls ``CreateMicrovmAuthToken`` first. When that - call answers ``ResourceNotFoundException``, the job fails with this - error. When it returns a token, the request goes to the endpoint. The - endpoint of a terminated MicroVM answers 502, and the request step - retries a 502 as a MicroVM that is still starting. So the job then fails - with a plain :class:`MicrovmDeliveryError`, after the retries of both - tiers. + Without the state check before delivery, the request step still calls + ``CreateMicrovmAuthToken`` first. When that call answers + ``ResourceNotFoundException``, the job fails with this error. When it + returns a token, the request goes to the endpoint. The endpoint of a + terminated MicroVM answers 502. Whether that 502 carries + ``MICROVM_CONNECT_FAILED`` is not verified. Without it, the request + step retries the 502 as a MicroVM that is still starting. So the job + then fails with a plain :class:`MicrovmDeliveryError`, after the + retries of both tiers. The live service's answer of ``CreateMicrovmAuthToken`` for a removed MicroVM is not verified. The API model lists diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/job.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/job.py index 59e7bee0e..cee564937 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/job.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/job.py @@ -36,8 +36,9 @@ 1. The worker accepts a job request body of at most 6 MiB. 2. The body holds the input as the same JSON that :func:`dumps` writes. 3. The body also holds fixed fields: the version, the Region, the MicroVM - ID, the callback ID, and the heartbeat timeout. Their names and the JSON - punctuation take 92 bytes. + ID, the suspend-after flag, the callback ID, and the heartbeat timeout. + Their names and the JSON punctuation take 112 bytes. The suspend-after + flag, ``,"suspendAfter":true``, takes 20 of them. 4. The botocore ``lambda-microvms`` model allows a MicroVM ID of 1 to 256 characters, with no pattern. JSON writes a character in at most 6 bytes, as a ``\\u00XX`` escape. So the MicroVM ID takes at most 1,536 bytes. @@ -61,6 +62,9 @@ class RunHookPayload: A job that goes over HTTP needs a ``runHookPayload`` too, because the worker rejects a payload without ``version`` and ``region``. See :func:`run_hook_payload_without_job`. + + The payload never carries ``suspendAfter``. Only a session job can set + it, and every session job goes over HTTP. """ region: str @@ -99,13 +103,19 @@ def to_json(self) -> str: class JobRequest: """The body of the HTTP request that delivers one job to a route. - Wire names: ``version``, ``region``, ``microvmId``, ``callbackId``, - ``heartbeatTimeoutSeconds``, and ``input``. The job fields sit at the top - level, next to ``version``, ``region``, and ``microvmId``. + Wire names: ``version``, ``region``, ``microvmId``, ``suspendAfter``, + ``callbackId``, ``heartbeatTimeoutSeconds``, and ``input``. The job + fields sit at the top level, next to ``version``, ``region``, and + ``microvmId``. The endpoint can accept a request before Lambda has sent the ``run`` hook. The worker learns the MicroVM ID from the ``run`` hook. So the request carries the MicroVM ID too. + + ``suspendAfter`` is ``true`` only for a session job that asks the worker + to suspend the MicroVM after the job. The field is omitted otherwise. An + older worker ignores a field that it does not know. So the version stays + 1, because a new version would make every older worker fail the job. """ region: str @@ -118,6 +128,8 @@ class JobRequest: """The caller's input. It must be JSON-serializable.""" heartbeat_timeout_seconds: int | None = None """The heartbeat timeout, when the caller set one.""" + suspend_after: bool = False + """Whether the worker suspends the MicroVM after the job.""" def to_json(self) -> str: """Serialize the request as compact JSON, as :meth:`RunHookPayload.to_json` does. @@ -126,16 +138,19 @@ def to_json(self) -> str: TypeError: When the input is not JSON-serializable. ValueError: When the input holds ``nan`` or an infinite float. """ - return dumps( - { - "version": PAYLOAD_VERSION, - "region": self.region, - "microvmId": self.microvm_id, - **_job_fields( - self.callback_id, self.job_input, self.heartbeat_timeout_seconds - ), - } + document: dict[str, Any] = { + "version": PAYLOAD_VERSION, + "region": self.region, + "microvmId": self.microvm_id, + } + if self.suspend_after: + document["suspendAfter"] = True + document.update( + _job_fields( + self.callback_id, self.job_input, self.heartbeat_timeout_seconds + ) ) + return dumps(document) def run_hook_payload_without_job( diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py index 5e31324ca..c1ddfe53f 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py @@ -7,11 +7,11 @@ import time from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Literal from aws_durable_execution_sdk_python.types import LoggerInterface -from botocore.exceptions import ClientError +from botocore.exceptions import BotoCoreError, ClientError from aws_durable_execution_sdk_python_extras.microvm.aws_errors import ( add_request_id, @@ -63,6 +63,31 @@ def __init__(self, message: str, state: str | None = None) -> None: """The MicroVM state, or ``None`` when the MicroVM no longer exists.""" +class MicrovmStateCheckError(Exception): + """``GetMicrovm`` failed in the check of :func:`resume_if_suspended`. + + The request step of a session without the resume mode runs that check + only after a sign of a suspended MicroVM. The session may never have + asked to suspend. So the step logs one warning per job, and keeps + sending the job. ``__cause__`` is the error of the call. + """ + + def __init__(self, message: str, *, denied: bool) -> None: + super().__init__(message) + self.denied = denied + """Whether the service answered ``AccessDeniedException``.""" + + +class MicrovmResumeDeniedError(Exception): + """``ResumeMicrovm`` was denied for a suspending or suspended MicroVM. + + Nothing else resumes the MicroVM, and a later attempt is denied the same + way. So the default retry strategy does not retry it, and the delivery + raises ``MicrovmDeliveryError`` at once. The message names the + permission, and ends with the request ID of the call. + """ + + @dataclass(frozen=True, kw_only=True) class EnsureRunningOptions: """The inputs of :func:`ensure_running`.""" @@ -124,12 +149,135 @@ def ensure_running(options: EnsureRunningOptions) -> str | None: ClientError: When ``GetMicrovm`` or ``ResumeMicrovm`` fails otherwise. """ + return _wait_until_running(options, resume_sent=False) + + +def resume_if_suspended(options: EnsureRunningOptions) -> str | None: + """Check the MicroVM state once, resume a suspended MicroVM, and return + its current endpoint. + + The request step of a session without the resume mode calls it after a + sign that the MicroVM may be suspended: an HTTP 502 from the endpoint + with ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``, or no answer. A + worker can suspend its MicroVM after a ``suspend_after`` job. The + session can lose the record of that job on replay. So this check runs + for every session job. + + 1. ``GetMicrovm`` runs once. A failure other than + ``ResourceNotFoundException`` raises :class:`MicrovmStateCheckError`. + The caller then keeps sending the job. + 2. ``RUNNING`` or ``PENDING``: the function returns at once. + 3. ``SUSPENDING`` or ``SUSPENDED``: it calls ``ResumeMicrovm`` once. + The service honors a resume during ``SUSPENDING`` too. It then waits + for ``RUNNING``, as :func:`ensure_running` does after its own resume. + 4. ``TERMINATING``, ``TERMINATED``, or not found: it raises + :class:`MicrovmStateError`. + + A denied ``ResumeMicrovm`` raises :class:`MicrovmResumeDeniedError`. + ``ConflictException`` is not an error: polling continues. Any other + ``ResumeMicrovm`` error is raised with its class and its ``response``, + as in :func:`ensure_running`. So is any error of a ``GetMicrovm`` call + after the resume. + + Returns: + The endpoint that ``GetMicrovm`` reports, or ``None`` when it reports + none. + + Raises: + MicrovmStateCheckError: When the first ``GetMicrovm`` call fails. + MicrovmStateError: When the MicroVM can no longer run a job. + MicrovmResumeDeniedError: When ``ResumeMicrovm`` is denied. + MicrovmEndpointUnavailableError: When the MicroVM does not reach + ``RUNNING`` in time. + ClientError: When ``ResumeMicrovm`` fails otherwise, or a later + ``GetMicrovm`` call fails. + """ + started = time.monotonic() + reserve = reserve_starts_at(options.remaining_seconds) + microvm_id = options.microvm_id + + def unavailable() -> Exception: + elapsed_ms = round((time.monotonic() - started) * 1000) + return MicrovmEndpointUnavailableError( + f"MicroVM {microvm_id} did not reach RUNNING in {elapsed_ms} ms. " + "Last state: unknown." + ) + + try: + current = call_before_deadline( + lambda: options.client.get_microvm(microvmIdentifier=microvm_id), + reserve, + unavailable, + ) + except ClientError as error: + code = error_code(error) + if code == "ResourceNotFoundException": + raise microvm_not_found(microvm_id, "GetMicrovm", error) from error + raise MicrovmStateCheckError( + str(add_request_id("GetMicrovm", error)), + denied=code == "AccessDeniedException", + ) from error + except BotoCoreError as error: + raise MicrovmStateCheckError(str(error), denied=False) from error + state = current.get("state") + _raise_if_ended(microvm_id, current) + if state not in {"SUSPENDING", "SUSPENDED"}: + return current.get("endpoint") + try: + call_before_deadline( + lambda: options.client.resume_microvm(microvmIdentifier=microvm_id), + reserve, + unavailable, + ) + except ClientError as error: + code = error_code(error) + if code == "ResourceNotFoundException": + raise microvm_not_found(microvm_id, "ResumeMicrovm", error) from error + if code == "AccessDeniedException": + msg = ( + f"MicroVM {microvm_id} is {state}. ResumeMicrovm was denied, " + "so the operation cannot resume it. The function's execution " + "role needs lambda:ResumeMicrovm." + ) + response = ( + error.response + if is_microvms_service_error("ResumeMicrovm", error) + else None + ) + raise MicrovmResumeDeniedError( + with_request_id(msg, "ResumeMicrovm", response) + ) from error + if code != "ConflictException": + raise add_request_id("ResumeMicrovm", error) + options.logger.info("resuming MicroVM", extra={"state": state}) + remaining = options.max_wait_seconds - (time.monotonic() - started) + return _wait_until_running( + replace(options, max_wait_seconds=remaining), resume_sent=True + ) + + +def _raise_if_ended(microvm_id: str, current: Any) -> None: + """Raise :class:`MicrovmStateError` for a terminating or terminated + MicroVM.""" + state = current.get("state") + if state in {"TERMINATING", "TERMINATED"}: + reason = current.get("stateReason") + # The GetMicrovm output carries the request ID. The state fails + # the step, so the message names that request ID. + msg = f"MicroVM {microvm_id} is {state}{f': {reason}' if reason else ''}." + raise MicrovmStateError(with_request_id(msg, "GetMicrovm", current), state) + + +def _wait_until_running( + options: EnsureRunningOptions, *, resume_sent: bool +) -> str | None: + """The loop of :func:`ensure_running`. ``resume_sent`` states that the + caller has already called ``ResumeMicrovm``.""" started = time.monotonic() reserve = reserve_starts_at(options.remaining_seconds) deadline = started + options.max_wait_seconds if reserve is not None: deadline = min(deadline, reserve) - resume_sent = False delay = _INITIAL_POLL_SECONDS last_state: str | None = None microvm_id = options.microvm_id @@ -178,12 +326,7 @@ def resume_microvm() -> Any: extra={"waitedMs": round((time.monotonic() - started) * 1000)}, ) return current.get("endpoint") - if state in {"TERMINATING", "TERMINATED"}: - reason = current.get("stateReason") - # The GetMicrovm output carries the request ID. The state fails - # the step, so the message names that request ID. - msg = f"MicroVM {microvm_id} is {state}{f': {reason}' if reason else ''}." - raise MicrovmStateError(with_request_id(msg, "GetMicrovm", current), state) + _raise_if_ended(microvm_id, current) if state == "SUSPENDED" and not resume_sent: try: resume_microvm() diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/operation.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/operation.py index b7e5759ee..39b20b02a 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/operation.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/operation.py @@ -65,9 +65,11 @@ ) from aws_durable_execution_sdk_python_extras.microvm.lifecycle import ( EnsureRunningOptions, + MicrovmStateCheckError, MicrovmStateError, ensure_running, microvm_not_found, + resume_if_suspended, ) from aws_durable_execution_sdk_python_extras.microvm.request import ( DEFAULT_MICROVM_JOB_PATH, @@ -256,7 +258,9 @@ def microvm( ValidationError: When ``name`` or ``config`` is not valid, or the Region cannot be determined. No durable operation starts. It is the SDK's class, which ``Duration`` also raises for a negative - value. + value. A ``config.request`` of type :class:`MicrovmInvokeOptions` + with ``suspend_after`` set is not valid, because only a session + job can suspend its MicroVM after the job. TypeError: When ``job_input`` is not JSON-serializable. No durable operation starts. ValueError: When the JSON of ``job_input`` is over @@ -290,6 +294,20 @@ def run(child: DurableContext) -> Any: """The error type that a request step checkpoints for a MicroVM that can no longer run a job. The SDK records the fully-qualified class name.""" +_STATE_CHECK_WARNING = ( + "GetMicrovm failed, so the operation could not check whether the MicroVM " + "is suspended. The request step keeps sending the job." +) +"""The warning, once per job, after a failed state check of a session +without the resume mode.""" + +_GET_DENIED_NOTE = ( + "GetMicrovm was denied, so the operation could not check whether the " + "MicroVM is suspended." +) +"""The sentence that the first tier adds to its final error after a denied +state check. Without it, the message names only the last HTTP failure.""" + def step_retry_strategy( config: MicrovmBaseConfig, @@ -479,7 +497,10 @@ def deliver( heartbeat_timeout_seconds: int | None, request: MicrovmRequestConfig, retry_window_seconds: int, - resume: bool = False, + resume: bool | Callable[[], bool] = False, + suspend_check: bool = False, + suspend_after: bool = False, + claim_suspend_after_warning: Callable[[], bool] | None = None, ) -> None: """Run the request step, ``.request``: POST the job to the MicroVM's endpoint. @@ -490,7 +511,16 @@ def deliver( With ``resume``, each attempt first makes sure the MicroVM runs, and resumes a suspended one. It then sends the job to the endpoint that ``GetMicrovm`` reports. A session sets it when its MicroVM can be - suspended. + suspended. A callable is read at the start of each attempt. + + Without ``resume``, ``suspend_check`` turns on the check after a sign + of a suspended MicroVM. See :func:`resume_if_suspended`. A session + sets it for every job. A failed ``GetMicrovm`` call in that check + logs one warning per job, and the attempt keeps sending the job. + + ``suspend_after`` adds ``"suspendAfter": true`` to the job request. + When the answer lacks the worker's acknowledgement, the attempt + logs a warning if ``claim_suspend_after_warning`` returns ``True``. ``GetMicrovm``, ``ResumeMicrovm``, and ``CreateMicrovmAuthToken`` answer ``ResourceNotFoundException`` for a removed MicroVM. An @@ -508,6 +538,7 @@ def deliver( callback_id=callback_id, job_input=job_input, heartbeat_timeout_seconds=heartbeat_timeout_seconds, + suspend_after=suspend_after, ) .to_json() .encode() @@ -515,9 +546,14 @@ def deliver( remaining_seconds = _remaining_seconds(child) path = DEFAULT_MICROVM_JOB_PATH if request.path is None else request.path port = DEFAULT_MICROVM_PORT if request.port is None else request.port + microvm_id = launched.microvm_id + # One warning per job for a failed state check. The attempts of the + # step in this invocation share it. + check_warned = threading.Event() def attempt(step_context: StepContext) -> None: client = self._client() + logger = step_context.logger # The state check before delivery and the request's retries share # one window. So the wait for a resume counts toward it. window_started_at = time.monotonic() @@ -525,23 +561,57 @@ def attempt(step_context: StepContext) -> None: # HTTP delivery. The launch step checkpoints only what # _launch_call returned. So an HTTP launch always has an endpoint. endpoint = cast("str", launched.endpoint) + strict = resume() if callable(resume) else resume + + def options(max_wait_seconds: float) -> EnsureRunningOptions: + return EnsureRunningOptions( + client=client, + microvm_id=microvm_id, + max_wait_seconds=max_wait_seconds, + remaining_seconds=remaining_seconds, + logger=logger, + ) + recheck: Callable[[float], str | None] | None = None - if resume: + check: Callable[[float], str | None] | None = None + get_denied = False + if strict: require_time_for_call(remaining_seconds) def running(max_wait_seconds: float) -> str | None: - return ensure_running( - EnsureRunningOptions( - client=client, - microvm_id=launched.microvm_id, - max_wait_seconds=max_wait_seconds, - remaining_seconds=remaining_seconds, - logger=step_context.logger, - ) - ) + return ensure_running(options(max_wait_seconds)) endpoint = running(retry_window_seconds) or endpoint recheck = running + elif suspend_check: + + def check_state(max_wait_seconds: float) -> str | None: + nonlocal get_denied + try: + return resume_if_suspended(options(max_wait_seconds)) + except MicrovmStateCheckError as error: + get_denied = get_denied or error.denied + if not check_warned.is_set(): + check_warned.set() + logger.warning( + _STATE_CHECK_WARNING, + extra={"microvmId": microvm_id, "error": str(error)}, + ) + return None + + check = check_state + + def ignored() -> None: + if ( + claim_suspend_after_warning is not None + and claim_suspend_after_warning() + ): + logger.warning( + f"the worker in MicroVM {microvm_id} does not support " + "suspendAfter. The MicroVM keeps running after the job.", + extra={"microvmId": microvm_id}, + ) + # Only CreateMicrovmAuthToken can raise ResourceNotFoundException # here: # @@ -549,9 +619,9 @@ def running(max_wait_seconds: float) -> str | None: # 2. send_job calls new_token() outside any try. So the error # leaves send_job unchanged. # 3. _run_with_limit re-raises the same error object. - # 4. ensure_running turns the same answer of GetMicrovm and - # ResumeMicrovm into MicrovmStateError. So neither call raises - # a ClientError with that code. + # 4. ensure_running and resume_if_suspended turn the same answer + # of GetMicrovm and ResumeMicrovm into MicrovmStateError. So + # neither call raises a ClientError with that code. # # So the step raises MicrovmStateError for CreateMicrovmAuthToken. # The default retry strategy does not retry that error. @@ -559,23 +629,28 @@ def running(max_wait_seconds: float) -> str | None: send_job( SendJobOptions( client=client, - microvm_id=launched.microvm_id, + microvm_id=microvm_id, endpoint=endpoint, path=path, port=port, body=body, retry_window_seconds=retry_window_seconds, remaining_seconds=remaining_seconds, - logger=step_context.logger, + logger=logger, opener=self.config.opener, window_started_at=window_started_at, recheck=recheck, + suspend_check=check, + unavailable_note=lambda: ( + _GET_DENIED_NOTE if get_denied else None + ), + on_suspend_after_ignored=ignored if suspend_after else None, ) ) except ClientError as error: if error_code(error) == "ResourceNotFoundException": raise microvm_not_found( - launched.microvm_id, "CreateMicrovmAuthToken", error + microvm_id, "CreateMicrovmAuthToken", error ) from error raise @@ -743,6 +818,16 @@ def create( heartbeat_timeout_seconds = _validate_heartbeat(name, config.heartbeat_timeout) request = config.request or MicrovmRequestConfig() retry_window_seconds = _validate_request(name, request) + # MicrovmInvokeOptions subclasses MicrovmRequestConfig. So a session + # job's options pass the type check of config.request. microvm() + # terminates its MicroVM after its one job, so a suspend after the + # job would only delay that. + if _validate_suspend_after(name, getattr(request, "suspend_after", False)): + msg = ( + f'MicroVM "{name}": suspend_after applies only to a session job. ' + "microvm() terminates its MicroVM after the job." + ) + raise ValidationError(msg) _validate_input(name, job_input) partition, region = _validate_region(context, name) return cls( @@ -1038,6 +1123,19 @@ def _validate_request(name: str, request: MicrovmRequestConfig) -> int: return seconds +def _validate_suspend_after(name: str, suspend_after: object) -> bool: + """Validate ``suspend_after``, and return it. + + Raises: + ValidationError: When it is not a ``bool``. An ``int`` is rejected + too, although ``bool`` subclasses ``int``. + """ + if not isinstance(suspend_after, bool): + msg = f'MicroVM "{name}": suspend_after must be a bool' + raise ValidationError(msg) + return suspend_after + + def _remaining_seconds(context: DurableContext) -> Callable[[], float | None]: """Return a reader of the invocation's remaining time, in seconds. diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py index 9a109b9c9..2d77cc0c9 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py @@ -73,6 +73,22 @@ _AUTH_TOKEN_MINUTES = 5 _PORT_HEADER = "X-aws-proxy-port" +SUSPEND_AFTER_HEADER = "X-Microvm-Suspend-After" +"""The header with which a worker acknowledges ``suspendAfter`` in its 202 +answer. Its value is :data:`SUSPEND_AFTER_ACCEPTED`. The JavaScript SDK +uses the same header.""" + +SUSPEND_AFTER_ACCEPTED = "accepted" + +_PROXY_ERROR_HEADER = "x-aws-proxy-error" +_CONNECT_FAILED = "MICROVM_CONNECT_FAILED" +"""The endpoint answers HTTP 502 with ``x-aws-proxy-error: +MICROVM_CONNECT_FAILED`` while the MicroVM is suspending or suspended. A +probe measured the answer in about 60 milliseconds.""" + +_SUSPEND_CHECK_INTERVAL_SECONDS = 2.0 +"""The shortest time between two calls of ``SendJobOptions.suspend_check``.""" + _T = TypeVar("_T") @@ -239,6 +255,31 @@ class SendJobOptions: seconds, so the wait for a resume never extends the window. ``None`` from it keeps the current endpoint. """ + suspend_check: Callable[[float], str | None] | None = None + """Resume the MicroVM if it is suspended, and return its current + endpoint. + + A session job without :attr:`recheck` sets it. A worker can suspend the + MicroVM after a ``suspend_after`` job, and the session can lose the + record of that job on replay. A suspended or suspending MicroVM gets no + answer to a request, or its endpoint answers HTTP 502 with + ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``. A booting MicroVM gets + neither: its endpoint holds the request until the worker listens. So + the next attempt calls it after either sign, at most once per 2 + seconds. A session that never suspends then makes no extra call. It + receives what remains of the retry window, in seconds. ``None`` from it + keeps the current endpoint. + """ + unavailable_note: Callable[[], str | None] | None = None + """A sentence to add to the message of + :class:`MicrovmEndpointUnavailableError`, or ``None``.""" + on_suspend_after_ignored: Callable[[], None] | None = None + """Called when the answer that accepted the job lacks the header + ``X-Microvm-Suspend-After: accepted``. + + A ``suspend_after`` job sets it. A worker that supports the field sends + the header. An older worker ignores the field, and the MicroVM keeps + running after the job.""" def require_time_for_call(remaining_seconds: Callable[[], float | None]) -> None: @@ -315,6 +356,18 @@ def send_job(options: SendJobOptions) -> int: The endpoint can accept a request before Lambda has sent the ``run`` hook. So the body carries the MicroVM ID too. + A session job checks the MicroVM state in one of two ways. With + ``recheck``, it checks after any 502, 503, 504, or request without an + answer. With ``suspend_check``, it checks only after a 502 with + ``x-aws-proxy-error: MICROVM_CONNECT_FAILED`` or a request without an + answer, at most once per 2 seconds. The function then sends the job + again to the endpoint that the check reports. + + A ``suspend_after`` job also reads the answer that accepted it. A worker + that supports the field answers with ``X-Microvm-Suspend-After: + accepted``. Without that header, the function calls + ``on_suspend_after_ignored``. + The same job can reach the route more than once: 1. A request reaches the route, and its answer is lost. @@ -340,8 +393,10 @@ def send_job(options: SendJobOptions) -> int: ``ResourceNotFoundException``. It is the error that the client raised, unchanged. The request step turns it into ``MicrovmStateError``. - MicrovmStateError: When ``recheck`` reports that the MicroVM can no - longer run a job. + MicrovmStateError: When ``recheck`` or ``suspend_check`` reports + that the MicroVM can no longer run a job. + MicrovmResumeDeniedError: When ``suspend_check`` finds a suspended + MicroVM, and ``ResumeMicrovm`` is denied. Exception: When ``CreateMicrovmAuthToken`` fails otherwise. The message of a botocore error ends with the request ID. """ @@ -357,9 +412,11 @@ def send_job(options: SendJobOptions) -> int: def unavailable() -> Exception: elapsed_ms = round((time.monotonic() - started) * 1000) + note = options.unavailable_note() if options.unavailable_note else None return MicrovmEndpointUnavailableError( f"The MicroVM did not accept the job at {options.path} after " f"{attempt} attempts in {elapsed_ms} ms. Last failure: {last_failure}." + + (f" {note}" if note else "") ) def deadline() -> float: @@ -378,6 +435,7 @@ def new_token() -> Mapping[str, str]: token = new_token() refreshed_token = False delay = _INITIAL_DELAY_SECONDS + last_suspend_check: float | None = None while True: attempt += 1 @@ -410,6 +468,11 @@ def new_token() -> Mapping[str, str]: options.logger.info( "MicroVM job accepted", extra={"attempt": attempt, "status": status} ) + if ( + options.on_suspend_after_ignored is not None + and answer.suspend_after != SUSPEND_AFTER_ACCEPTED + ): + options.on_suspend_after_ignored() return status if status is None: last_failure = answer.failure @@ -466,6 +529,25 @@ def new_token() -> Mapping[str, str]: raise unavailable() endpoint = options.recheck(left) or endpoint url = _endpoint_url(endpoint, options.path) + # A session without the resume mode checks the state only after a + # sign of a suspended MicroVM. A booting MicroVM gives no such sign. + # The check repeats after each such failure, at most once per 2 + # seconds. The worker can suspend after the job's first state check, + # and the last call to the service wins. + if options.suspend_check is not None and ( + status is None or answer.connect_failed + ): + now = time.monotonic() + if ( + last_suspend_check is None + or now - last_suspend_check >= _SUSPEND_CHECK_INTERVAL_SECONDS + ): + last_suspend_check = now + left = deadline() - now + if left <= 0: + raise unavailable() + endpoint = options.suspend_check(left) or endpoint + url = _endpoint_url(endpoint, options.path) def _create_auth_token(options: SendJobOptions) -> Mapping[str, str]: @@ -545,6 +627,12 @@ class _Answer: """Whether the answer has a ``Location`` header.""" followed_redirect: bool = False """Whether the opener followed a redirect.""" + connect_failed: bool = False + """Whether the endpoint answered 502 with ``x-aws-proxy-error: + MICROVM_CONNECT_FAILED``. It answers so for a suspending or suspended + MicroVM.""" + suspend_after: str | None = None + """The value of the ``X-Microvm-Suspend-After`` header of a 2xx answer.""" def _post_within( @@ -593,11 +681,21 @@ def _post( # An opener that follows redirects reports the URL it reached. if response.geturl() != url: return _Answer(response.status, followed_redirect=True) - return _Answer(response.status) + return _Answer( + response.status, + suspend_after=response.headers.get(SUSPEND_AFTER_HEADER), + ) except urllib.error.HTTPError as error: has_location = bool(error.headers and error.headers.get("Location")) + connect_failed = ( + error.code == 502 # noqa: PLR2004 + and bool(error.headers) + and error.headers.get(_PROXY_ERROR_HEADER) == _CONNECT_FAILED + ) error.close() - return _Answer(error.code, has_location=has_location) + return _Answer( + error.code, has_location=has_location, connect_failed=connect_failed + ) except (OSError, http.client.HTTPException, ValueError) as error: # URLError, a refused connection, and a timeout are OSError. An # invalid response from the endpoint raises an HTTPException. diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py index 6cd824c0a..c88c3ec76 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py @@ -6,7 +6,9 @@ from __future__ import annotations +import threading import uuid +from collections import OrderedDict from collections.abc import Callable, Mapping from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, TypeVar @@ -35,6 +37,7 @@ _validate_input, _validate_region, _validate_request, + _validate_suspend_after, _validate_timeout, create_job_callback, job_error, @@ -59,6 +62,16 @@ ``timeout`` is under 10 seconds does not suspend by default. """ +_WARNED_SESSIONS: OrderedDict[str, None] = OrderedDict() +"""The IDs of the sessions that warned about a worker that ignores +``suspend_after``, oldest first.""" + +_MAX_WARNED_SESSIONS = 1_000 +"""How many session IDs :data:`_WARNED_SESSIONS` keeps. A session ID is a +short string, so 1,000 of them use well under 1 MB.""" + +_WARNED_LOCK = threading.Lock() + _MIN_AUTO_SUSPEND_IDLE_SECONDS = 10 """The shortest idle time a session accepts. @@ -108,6 +121,14 @@ def microvm_session( records whether the MicroVM can be suspended. Replay follows that record, not the current config. + A job with ``suspend_after`` asks the worker to suspend the MicroVM after + the job. Each later job of the handle then checks the state before + delivery, as with ``auto_suspend_on_idle``. That record is kept in + memory only. So every job of every session also checks the state after + a sign of a suspended MicroVM: no answer, or an HTTP 502 with + ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``. A booting MicroVM gives + neither sign. So a session that never suspends makes no extra call. + A failed :meth:`MicrovmSession.invoke` raises inside the handler. The handler can catch the error and continue with the same MicroVM. An error that leaves the handler ends the session, and the MicroVM is @@ -203,16 +224,50 @@ def __init__( launched: LaunchResult, *, resume_before_delivery: bool, + session_id: str, ) -> None: self._scope = scope self._launched = launched self._resume_before_delivery = resume_before_delivery + # Set by the first suspend_after job of this handle. Later jobs then + # check the state before delivery. The flag lives in memory only. A + # replay of a completed job, map item, or child context does not run + # its invoke() again, so it does not set the flag. The check after a + # sign of a suspended MicroVM covers that case. + self._may_be_suspended = False + self._session_id = session_id @property def microvm_id(self) -> str: """The ID of the session's MicroVM.""" return self._launched.microvm_id + def _resume_mode(self) -> bool: + """Whether a job checks the MicroVM state before delivery. + + The request step reads it at the start of each attempt. So a step + retry, and a ``map`` item that starts after a sibling set the flag, + see the current value. The flag only changes from ``False`` to + ``True``. + """ + return self._resume_before_delivery or self._may_be_suspended + + def _claim_suspend_after_warning(self) -> bool: + """Return ``True`` the first time, so the session warns once about a + worker that does not support ``suspend_after``. + + Each invocation creates a new handle. So the record is kept per + process, by session ID. A later invocation in the same execution + environment then logs no second warning. + """ + with _WARNED_LOCK: + if self._session_id in _WARNED_SESSIONS: + return False + _WARNED_SESSIONS[self._session_id] = None + if len(_WARNED_SESSIONS) > _MAX_WARNED_SESSIONS: + _WARNED_SESSIONS.popitem(last=False) + return True + def invoke( self, context: DurableContext, @@ -237,7 +292,8 @@ def invoke( name: The job's child context name. job_input: The job's input. It must be JSON-serializable, and the same on every replay. - options: The route and the timeouts of this job. + options: The route, the timeouts, and ``suspend_after`` of this + job. Returns: The value that the MicroVM passed to @@ -252,13 +308,16 @@ def invoke( raises it at ``context.run_in_child_context``. ValueError: When the JSON of ``job_input`` is over ``MAX_JOB_INPUT_BYTES`` UTF-8 bytes. - ValidationError: When ``name`` or ``options`` is not valid. + ValidationError: When ``name`` or ``options`` is not valid, for + example a ``suspend_after`` that is not a ``bool``. MicrovmNotRunningError: When the session's MicroVM has ended. The session reports it when it checks the MicroVM state before each job, and for every job when ``CreateMicrovmAuthToken`` - reports the MicroVM removed. + reports the MicroVM removed. A check after a sign of a + suspended MicroVM reports it too. MicrovmDeliveryError: When the delivery fails after all retries, - or the route rejects the job. + or the route rejects the job. A suspended MicroVM whose + ``ResumeMicrovm`` is denied fails the delivery at once. MicrovmJobFailedError: When the MicroVM reports a failure, or sends a result that is not valid JSON. MicrovmTimeoutError: When no result or heartbeat arrives in time. @@ -285,9 +344,15 @@ def invoke( retry_window_seconds = _validate_request(name, options) timeout_seconds = _validate_timeout(name, "timeout", options.timeout) heartbeat_timeout_seconds = _validate_heartbeat(name, options.heartbeat_timeout) + suspend_after = _validate_suspend_after(name, options.suspend_after) _validate_input(name, job_input) launched = self._launched - resume = self._resume_before_delivery + if suspend_after: + # The flag is set before the child context. A replay of a + # completed job returns its result without running the child + # context's function. So the flag is set on every replay of this + # call too. + self._may_be_suspended = True def run(job_context: DurableContext) -> Any: callback = create_job_callback( @@ -302,7 +367,12 @@ def run(job_context: DurableContext) -> Any: heartbeat_timeout_seconds=heartbeat_timeout_seconds, request=options, retry_window_seconds=retry_window_seconds, - resume=resume, + resume=self._resume_mode, + suspend_check=True, + suspend_after=suspend_after, + claim_suspend_after_warning=( + self._claim_suspend_after_warning if suspend_after else None + ), ) try: return callback.result() @@ -401,7 +471,8 @@ def run_session( # 3. So a deploy that changes auto_suspend_on_idle or idle_policy does # not change what the running MicroVM does. # 4. A job that skips the state check of a suspended MicroVM gets a - # 502. The request step retries a 502 until both retry tiers end. + # 502, or no answer. The request step then checks the state, and + # resumes the MicroVM. That costs a failed request per job. # 5. So each job follows the resume mode that the launch recorded, # not the current config. # Every session launch records a boolean. So only `True` turns the @@ -412,6 +483,7 @@ def run_session( self, launched, resume_before_delivery=launched.resume_before_delivery is True, + session_id=session_id, ), child, ) diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py b/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py index 480fbbbb2..0ce3f5522 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py @@ -135,19 +135,43 @@ def make_error() -> Callable[..., ClientError]: return client_error +@pytest.fixture +def connect_failed() -> Status: + """The answer of the endpoint of a suspended MicroVM.""" + return CONNECT_FAILED + + +@pytest.fixture +def suspend_after_ack() -> dict[str, str]: + """The header of a worker that supports ``suspendAfter``.""" + return dict(SUSPEND_AFTER_ACK) + + +Status = int | tuple[int, dict[str, str]] | None + +CONNECT_FAILED: Status = (502, {"x-aws-proxy-error": "MICROVM_CONNECT_FAILED"}) +"""The answer of the endpoint of a suspending or suspended MicroVM.""" + +SUSPEND_AFTER_ACK = {"X-Microvm-Suspend-After": "accepted"} +"""The header with which a worker acknowledges ``suspendAfter``.""" + + class FakeEndpoint: """A local HTTP server that stands in for a MicroVM endpoint. It records each request, and answers from ``statuses`` in order. After the list ends, it answers 202. A status of ``None`` closes the connection - without an answer. A redirect status gets ``Location: ``, - ``/moved`` by default. A ``location`` of ``None`` sends no ``Location``. + without an answer. A ``(status, headers)`` pair adds the headers to the + answer. The 202 after the list ends gets ``accept_headers``. A redirect + status gets ``Location: ``, ``/moved`` by default. A + ``location`` of ``None`` sends no ``Location``. """ def __init__( self, - statuses: list[int | None] | None = None, + statuses: list[Status] | None = None, location: str | None = "/moved", + accept_headers: dict[str, str] | None = None, ) -> None: self.requests: list[dict[str, Any]] = [] self._statuses = list(statuses or []) @@ -181,12 +205,19 @@ def do_POST(self) -> None: # noqa: N802 - the http.server name "body": json.loads(body), } ) - status = endpoint._statuses.pop(0) if endpoint._statuses else 202 - if status is None: + answer = ( + endpoint._statuses.pop(0) + if endpoint._statuses + else (202, dict(accept_headers or {})) + ) + if answer is None: self.close_connection = True self.connection.shutdown(socket.SHUT_RDWR) return + status, headers = answer if isinstance(answer, tuple) else (answer, {}) self.send_response(status) + for name, value in headers.items(): + self.send_header(name, value) if 300 <= status < 400 and location is not None: # noqa: PLR2004 self.send_header("location", location) self.send_header("content-length", "0") @@ -213,9 +244,11 @@ def fake_endpoint() -> Iterator[Callable[..., FakeEndpoint]]: endpoints: list[FakeEndpoint] = [] def _create( - statuses: list[int | None] | None = None, location: str | None = "/moved" + statuses: list[Status] | None = None, + location: str | None = "/moved", + accept_headers: dict[str, str] | None = None, ) -> FakeEndpoint: - endpoint = FakeEndpoint(statuses, location) + endpoint = FakeEndpoint(statuses, location, accept_headers) endpoints.append(endpoint) return endpoint diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/job_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/job_test.py index 54853af27..f9198ac94 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/job_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/job_test.py @@ -7,6 +7,7 @@ import json import math +from typing import Any import pytest from aws_durable_execution_sdk_python_microvm_worker import ( @@ -15,6 +16,7 @@ ) from aws_durable_execution_sdk_python_extras.microvm.job import ( + MAX_JOB_INPUT_BYTES, MAX_RUN_HOOK_PAYLOAD_LENGTH, JobRequest, RunHookPayload, @@ -150,3 +152,51 @@ def test_the_worker_reads_the_idle_time_of_a_session() -> None: assert payload == '{"version":1,"region":"us-east-1","autoSuspendIdleSeconds":60}' assert request.payload is not None assert request.payload.auto_suspend_idle_seconds == 60 + + +def test_writes_suspend_after_only_when_it_is_set() -> None: + fields: dict[str, Any] = { + "region": "us-east-1", + "microvm_id": "mvm-1", + "callback_id": "cb-1", + "job_input": 1, + } + + assert JobRequest(**fields, suspend_after=True).to_json() == ( + '{"version":1,"region":"us-east-1","microvmId":"mvm-1",' + '"suspendAfter":true,"callbackId":"cb-1","input":1}' + ) + assert "suspendAfter" not in JobRequest(**fields).to_json() + + +def test_the_worker_reads_suspend_after() -> None: + body = JobRequest( + region="us-east-1", + microvm_id="mvm-1", + callback_id="cb-1", + job_input=None, + suspend_after=True, + ).to_json() + + assert MicrovmJobRequest.from_dict(json.loads(body)).suspend_after is True + + +def test_the_job_fields_fit_in_the_margin_of_the_input_limit() -> None: + # 1. The longest MicroVM ID is 256 characters, each escaped in 6 bytes. + # 2. The longest callback ID is 1,024 characters of 1 byte. + # 3. The heartbeat timeout is at most 8 hours, and suspendAfter is set. + # So the body of the largest input is at most the 6 MiB that the worker + # reads. + body = JobRequest( + region="ap-southeast-5", + microvm_id="\x01" * 256, + callback_id="A" * 1024, + job_input="x" * (MAX_JOB_INPUT_BYTES - 2), + heartbeat_timeout_seconds=28_800, + suspend_after=True, + ).to_json() + + fixed = len(body.encode()) - MAX_JOB_INPUT_BYTES + assert fixed == 112 + len("ap-southeast-5") + 256 * 6 + 1024 + len("28800") + assert fixed < 3 * 1024 + assert len(body.encode()) <= 6 * 1024 * 1024 diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py index 3f2eed380..a1d6e4007 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py @@ -11,14 +11,18 @@ from typing import Any import pytest +from botocore.exceptions import ClientError, EndpointConnectionError from aws_durable_execution_sdk_python_extras.microvm import ( MicrovmEndpointUnavailableError, ) from aws_durable_execution_sdk_python_extras.microvm.lifecycle import ( EnsureRunningOptions, + MicrovmResumeDeniedError, + MicrovmStateCheckError, MicrovmStateError, ensure_running, + resume_if_suspended, ) from aws_durable_execution_sdk_python_extras.microvm.request import ( INVOCATION_RESERVE_SECONDS, @@ -222,3 +226,209 @@ def test_raises_the_state_error_when_the_microvm_is_removed_before_the_resume( with pytest.raises(MicrovmStateError, match="no longer exists"): run_with(client) + + +# --------------------------------------------------------------------------- +# resume_if_suspended: the check after a sign of a suspended MicroVM +# --------------------------------------------------------------------------- + + +def check_with( + client: Any, + *, + max_wait_seconds: float = 30, + remaining: Callable[[], float | None] = lambda: None, + sleeps: list[float] | None = None, +) -> str | None: + recorded = [] if sleeps is None else sleeps + return resume_if_suspended( + EnsureRunningOptions( + client=client, + microvm_id="mvm-1", + max_wait_seconds=max_wait_seconds, + remaining_seconds=remaining, + logger=logging.getLogger("test"), + sleep=recorded.append, + ) + ) + + +def service_error(operation: str, code: str, status: int = 400) -> ClientError: + response: Any = { + "Error": {"Code": code, "Message": f"{code} message"}, + "ResponseMetadata": {"HTTPStatusCode": status, "RequestId": "req-1"}, + } + return ClientError(response, operation) + + +@pytest.mark.parametrize("state", ["RUNNING", "PENDING"]) +def test_check_returns_at_once_for_a_microvm_that_runs_or_boots( + fake_client: Any, state: str +) -> None: + client = fake_client(states=[state], endpoint="mvm-1.new.example") + + assert check_with(client) == "mvm-1.new.example" + assert calls(client) == ["get_microvm"] + + +@pytest.mark.parametrize("state", ["SUSPENDING", "SUSPENDED"]) +def test_check_resumes_a_suspending_or_suspended_microvm_once( + fake_client: Any, state: str +) -> None: + # The service honors a resume during SUSPENDING. So the check does not + # wait for SUSPENDED. + client = fake_client(states=[state, "SUSPENDED", "PENDING", "RUNNING"]) + sleeps: list[float] = [] + + assert check_with(client, sleeps=sleeps) == "mvm-1.microvm.example" + assert calls(client) == [ + "get_microvm", + "resume_microvm", + "get_microvm", + "get_microvm", + "get_microvm", + ] + assert sleeps == [0.25, 0.5] + + +def test_check_keeps_polling_after_a_resume_conflict( + fake_client: Any, make_error: Any +) -> None: + client = fake_client( + states=["SUSPENDED", "RUNNING"], + resumes=[make_error("ConflictException", 409)], + ) + + assert check_with(client) == "mvm-1.microvm.example" + assert calls(client) == ["get_microvm", "resume_microvm", "get_microvm"] + + +@pytest.mark.parametrize("state", ["TERMINATING", "TERMINATED"]) +def test_check_raises_the_state_error_for_a_microvm_that_ends( + fake_client: Any, state: str +) -> None: + client = fake_client(states=[state]) + + with pytest.raises(MicrovmStateError, match=f"MicroVM mvm-1 is {state}") as raised: + check_with(client) + + assert raised.value.state == state + + +@pytest.mark.parametrize("operation", ["GetMicrovm", "ResumeMicrovm"]) +def test_check_raises_the_state_error_for_a_removed_microvm( + fake_client: Any, operation: str +) -> None: + error = service_error(operation, "ResourceNotFoundException", 404) + client = ( + fake_client(states=[error]) + if operation == "GetMicrovm" + else fake_client(states=["SUSPENDED"], resumes=[error]) + ) + + with pytest.raises(MicrovmStateError) as raised: + check_with(client) + + assert str(raised.value).endswith(f"{operation} request ID: req-1.") + assert raised.value.__cause__ is error + + +@pytest.mark.parametrize( + ("code", "denied"), + [("AccessDeniedException", True), ("ThrottlingException", False)], +) +def test_check_reports_a_failed_get_microvm_without_failing_the_job( + fake_client: Any, code: str, denied: bool +) -> None: + error = service_error("GetMicrovm", code) + client = fake_client(states=[error]) + + with pytest.raises(MicrovmStateCheckError) as raised: + check_with(client) + + assert raised.value.denied is denied + assert f"({code})" in str(raised.value) + assert str(raised.value).endswith(f"{code} message. GetMicrovm request ID: req-1.") + assert raised.value.__cause__ is error + + +def test_check_reports_a_get_microvm_connection_failure(fake_client: Any) -> None: + error = EndpointConnectionError(endpoint_url="https://lambda.example") + client = fake_client(states=[error]) + + with pytest.raises(MicrovmStateCheckError) as raised: + check_with(client) + + assert raised.value.denied is False + assert raised.value.__cause__ is error + + +@pytest.mark.parametrize("state", ["SUSPENDING", "SUSPENDED"]) +def test_check_fails_at_once_when_resume_microvm_is_denied( + fake_client: Any, state: str +) -> None: + error = service_error("ResumeMicrovm", "AccessDeniedException", 403) + client = fake_client(states=[state], resumes=[error]) + + with pytest.raises(MicrovmResumeDeniedError) as raised: + check_with(client) + + assert str(raised.value) == ( + f"MicroVM mvm-1 is {state}. ResumeMicrovm was denied, so the operation " + "cannot resume it. The function's execution role needs " + "lambda:ResumeMicrovm. ResumeMicrovm request ID: req-1." + ) + assert raised.value.__cause__ is error + assert calls(client) == ["get_microvm", "resume_microvm"] + + +def test_check_names_no_request_id_for_an_error_of_another_call( + fake_client: Any, +) -> None: + # An STS error inside the call carries the request ID of STS. + client = fake_client( + states=["SUSPENDED"], + resumes=[service_error("AssumeRole", "AccessDeniedException", 403)], + ) + + with pytest.raises(MicrovmResumeDeniedError) as raised: + check_with(client) + + assert str(raised.value).endswith("needs lambda:ResumeMicrovm.") + + +def test_check_raises_another_resume_error_with_its_request_id( + fake_client: Any, +) -> None: + client = fake_client( + states=["SUSPENDED"], + resumes=[service_error("ResumeMicrovm", "ThrottlingException", 429)], + ) + + with pytest.raises(ClientError) as raised: + check_with(client) + + assert str(raised.value).endswith("ResumeMicrovm request ID: req-1.") + + +def test_check_gives_up_when_the_resume_does_not_end_in_time( + fake_client: Any, +) -> None: + client = fake_client(states=["SUSPENDED"] * 10) + + with pytest.raises(MicrovmEndpointUnavailableError, match="Last state: SUSPENDED"): + check_with(client, max_wait_seconds=0.1) + + +def test_check_stops_a_stalled_get_microvm_at_the_reserve(fake_client: Any) -> None: + release = threading.Event() + client = fake_client() + client.get_microvm = lambda **_: release.wait(30) + + try: + with pytest.raises( + MicrovmEndpointUnavailableError, match="Last state: unknown" + ): + check_with(client, remaining=lambda: INVOCATION_RESERVE_SECONDS + 1.2) + finally: + release.set() diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py index 6efeb8e62..4560477cb 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py @@ -15,6 +15,7 @@ import pytest from aws_durable_execution_sdk_python_microvm_worker import MICROVM_JOB_PATH +from aws_durable_execution_sdk_python_microvm_worker import worker from aws_durable_execution_sdk_python_extras.microvm import ( DEFAULT_MICROVM_JOB_PATH, @@ -57,6 +58,11 @@ def test_matches_the_job_path_of_the_worker() -> None: assert DEFAULT_MICROVM_JOB_PATH == MICROVM_JOB_PATH +def test_matches_the_suspend_after_header_of_the_worker() -> None: + assert request.SUSPEND_AFTER_HEADER == worker.SUSPEND_AFTER_HEADER + assert request.SUSPEND_AFTER_ACCEPTED == "accepted" + + def test_posts_the_job_with_an_auth_token(fake_client: Any, fake_endpoint: Any) -> None: client = fake_client() endpoint = fake_endpoint() @@ -637,3 +643,221 @@ def test_sends_no_auth_token_on_a_redirect_that_the_callers_opener_follows( assert followed["method"] == "GET" assert "x-test-token" not in followed["headers"] assert "content-type" not in followed["headers"] + + +# --------------------------------------------------------------------------- +# The check for a suspended MicroVM, and the suspendAfter acknowledgement +# --------------------------------------------------------------------------- + + +class FakeClock: + """A monotonic clock that advances only when the tier sleeps.""" + + def __init__(self) -> None: + self.now = 1000.0 + self.sleeps: list[float] = [] + + def monotonic(self) -> float: + return self.now + + def sleep(self, seconds: float) -> None: + self.sleeps.append(seconds) + self.now += seconds + + +@pytest.fixture +def clock(monkeypatch: pytest.MonkeyPatch) -> FakeClock: + fake = FakeClock() + monkeypatch.setattr(request, "time", fake) + return fake + + +def test_checks_for_a_suspend_after_a_connect_failure_or_no_answer( + fake_client: Any, fake_endpoint: Any, connect_failed: Any, clock: FakeClock +) -> None: + # 1. A suspended MicroVM's endpoint answers 502 with MICROVM_CONNECT_FAILED. + # 2. A request that the freeze caught gets no answer. + # 3. The clock passes 2 seconds between the two failures. + endpoint = fake_endpoint([connect_failed, 503, 503, 503, None]) + windows: list[float] = [] + + def check(left: float) -> str | None: + windows.append(left) + return None + + status = send_job( + options( + fake_client(), + endpoint.url, + [], + sleep=clock.sleep, + suspend_check=check, + ) + ) + + assert status == 202 + assert clock.sleeps == [0.25, 0.5, 1.0, 2.0, 4.0] + assert windows == [60 - 0.25, 60 - 7.75] + + +@pytest.mark.parametrize( + "answer", + [502, 503, 504, (502, {"x-aws-proxy-error": "OTHER"}), 429, 500], + ids=["502", "503", "504", "502-other-proxy-error", "429", "500"], +) +def test_makes_no_suspend_check_for_another_failure( + fake_client: Any, fake_endpoint: Any, answer: Any +) -> None: + # A booting MicroVM never gets MICROVM_CONNECT_FAILED. So a session that + # never suspends makes no extra call. + endpoint = fake_endpoint([answer]) + checks: list[float] = [] + + status = send_job( + options(fake_client(), endpoint.url, [], suspend_check=checks.append) + ) + + assert status == 202 + assert checks == [] + + +def test_checks_for_a_suspend_at_most_once_per_2_seconds( + fake_client: Any, fake_endpoint: Any, connect_failed: Any, clock: FakeClock +) -> None: + endpoint = fake_endpoint([connect_failed] * 5) + checked_at: list[float] = [] + + def check(left: float) -> str | None: + checked_at.append(clock.now - 1000) + return None + + send_job( + options(fake_client(), endpoint.url, [], sleep=clock.sleep, suspend_check=check) + ) + + # The failures end at 0.25, 0.75, 1.75, 3.75, and 7.75 seconds. A check + # that would come less than 2 seconds after the last one is skipped. + assert checked_at == [0.25, 3.75, 7.75] + assert len(endpoint.requests) == 6 + + +def test_sends_the_next_request_to_the_endpoint_that_the_suspend_check_returns( + fake_client: Any, fake_endpoint: Any, connect_failed: Any +) -> None: + before = fake_endpoint([connect_failed]) + after = fake_endpoint() + + status = send_job( + options(fake_client(), before.url, [], suspend_check=lambda left: after.url) + ) + + assert status == 202 + assert len(before.requests) == 1 + assert len(after.requests) == 1 + + +def test_ends_the_tier_when_the_suspend_check_has_no_time_left( + fake_client: Any, fake_endpoint: Any, connect_failed: Any +) -> None: + endpoint = fake_endpoint([connect_failed]) + checks: list[float] = [] + + with pytest.raises(MicrovmEndpointUnavailableError, match="Last failure: HTTP 502"): + send_job( + options( + fake_client(), + endpoint.url, + [], + retry_window_seconds=1, + window_started_at=time.monotonic() - 0.2, + sleep=lambda seconds: time.sleep(0.9), + suspend_check=checks.append, + ) + ) + + assert checks == [] + + +def test_adds_the_note_to_the_final_error( + fake_client: Any, fake_endpoint: Any, connect_failed: Any +) -> None: + endpoint = fake_endpoint([connect_failed] * 10) + + with pytest.raises(MicrovmEndpointUnavailableError) as raised: + send_job( + options( + fake_client(), + endpoint.url, + [], + retry_window_seconds=1, + unavailable_note=lambda: "GetMicrovm was denied.", + ) + ) + + assert str(raised.value).endswith("Last failure: HTTP 502. GetMicrovm was denied.") + + +def test_adds_no_note_when_the_note_is_none( + fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([503] * 10) + + with pytest.raises(MicrovmEndpointUnavailableError) as raised: + send_job( + options( + fake_client(), + endpoint.url, + [], + retry_window_seconds=1, + unavailable_note=lambda: None, + ) + ) + + assert str(raised.value).endswith("Last failure: HTTP 503.") + + +@pytest.mark.parametrize( + ("headers", "ignored"), + [ + ({"X-Microvm-Suspend-After": "accepted"}, 0), + ({}, 1), + ({"X-Microvm-Suspend-After": "yes"}, 1), + ], + ids=["acknowledged", "no-header", "other-value"], +) +def test_reports_a_suspend_after_job_that_the_worker_did_not_acknowledge( + fake_client: Any, fake_endpoint: Any, headers: dict[str, str], ignored: int +) -> None: + endpoint = fake_endpoint(accept_headers=headers) + calls: list[None] = [] + + status = send_job( + options( + fake_client(), + endpoint.url, + [], + on_suspend_after_ignored=lambda: calls.append(None), + ) + ) + + assert status == 202 + assert len(calls) == ignored + + +def test_reads_the_acknowledgement_of_the_answer_that_accepted_the_job( + fake_client: Any, fake_endpoint: Any, suspend_after_ack: dict[str, str] +) -> None: + # A refused answer that carries the header does not count. + endpoint = fake_endpoint([(503, suspend_after_ack)]) + calls: list[None] = [] + + send_job( + options( + fake_client(), + endpoint.url, + [], + on_suspend_after_ignored=lambda: calls.append(None), + ) + ) + + assert calls == [None] diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py index d09f1fa84..23b8262e0 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py @@ -8,13 +8,15 @@ import datetime import hashlib import json +import logging +from collections import OrderedDict from collections.abc import Callable, Iterator from types import SimpleNamespace from typing import Any import boto3 import pytest -from aws_durable_execution_sdk_python.config import Duration, StepConfig +from aws_durable_execution_sdk_python.config import Duration, MapConfig, StepConfig from aws_durable_execution_sdk_python.context import Callback, DurableContext from aws_durable_execution_sdk_python.exceptions import ( ChildContextError, @@ -48,6 +50,7 @@ microvm_session, operation, ) +from aws_durable_execution_sdk_python_extras.microvm import session as session_module WAIT = 30 @@ -373,6 +376,20 @@ def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: "TypeError", "must be JSON-serializable", ), + ( + lambda vm, ctx: vm.invoke( + ctx, "build", job_input=1, options=job(suspend_after=1) + ), + "ValidationError", + "suspend_after must be a bool", + ), + ( + lambda vm, ctx: vm.invoke( + ctx, "build", job_input=1, options=job(suspend_after="yes") + ), + "ValidationError", + "suspend_after must be a bool", + ), ], ) def test_rejects_an_invalid_job_inside_the_handler_where_it_can_be_caught( @@ -1461,3 +1478,405 @@ def test_makes_no_state_check_for_a_launch_checkpoint_without_a_bool_resume_mode **fields, } assert client.requests("get_microvm") == [] + + +# --------------------------------------------------------------------------- +# suspend_after, and the check after a sign of a suspended MicroVM +# --------------------------------------------------------------------------- + +IGNORED_WARNING = ( + "the worker in MicroVM mvm-1 does not support suspendAfter. The MicroVM " + "keeps running after the job." +) +STATE_CHECK_WARNING = ( + "GetMicrovm failed, so the operation could not check whether the MicroVM " + "is suspended. The request step keeps sending the job." +) + + +def warnings(caplog: pytest.LogCaptureFixture, message: str) -> int: + return sum( + 1 + for record in caplog.records + if record.levelname == "WARNING" and record.getMessage() == message + ) + + +def suspend_then_wait_then_job( + vm: MicrovmSession, ctx: DurableContext, event: Any +) -> Any: + """A suspend_after job, a durable wait, and one more job. + + The wait ends the invocation. So the second job runs in a later + invocation, in which the first job's child context is complete.""" + build = vm.invoke(ctx, "build", job_input=1, options=job(suspend_after=True)) + ctx.wait(Duration.from_seconds(1), name="pause") + return [build, vm.invoke(ctx, "test", job_input=2, options=job())] + + +def test_asks_the_worker_to_suspend_after_the_job_only_when_it_is_set( + run: Any, fake_client: Any, fake_endpoint: Any, suspend_after_ack: Any +) -> None: + endpoint = fake_endpoint(accept_headers=suspend_after_ack) + client = fake_client(endpoint=endpoint.url) + + result = run( + session_handler(session_config(client), suspend_then_wait_then_job), + None, + [("build", succeed_with(1)), ("test", succeed_with(2))], + ) + + assert outcome(result) == {"result": [1, 2]} + first, second = (request["body"] for request in endpoint.requests) + assert first["suspendAfter"] is True + assert "suspendAfter" not in second + # The run hook payload never carries it. + [launch] = client.requests("run_microvm") + assert json.loads(launch["runHookPayload"]) == {"version": 1, "region": "us-west-2"} + # suspend_after is not a session setting. So the launch records no + # resume mode for it. + assert launch_checkpoint(result)["resumeBeforeDelivery"] is False + + +def test_checks_the_state_before_the_jobs_after_a_suspend_after_job( + run: Any, fake_client: Any, fake_endpoint: Any, suspend_after_ack: Any +) -> None: + endpoint = fake_endpoint(accept_headers=suspend_after_ack) + client = fake_client( + endpoint=endpoint.url, states=["RUNNING", "SUSPENDED", "RUNNING"] + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + build = vm.invoke(ctx, "build", job_input=1, options=job(suspend_after=True)) + return [build, vm.invoke(ctx, "test", job_input=2, options=job())] + + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with(1)), ("test", succeed_with(2))], + ) + + assert outcome(result) == {"result": [1, 2]} + # The flag is set before the first job's child context. So the first job + # checks the state too: RUNNING. The second job finds the MicroVM + # SUSPENDED, resumes it, and sends the job when it runs. + assert len(client.requests("get_microvm")) == 3 + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert len(endpoint.requests) == 2 + + +def test_sets_the_flag_again_when_a_completed_suspend_after_job_replays( + run: Any, fake_client: Any, fake_endpoint: Any, suspend_after_ack: Any +) -> None: + # The second job runs in a later invocation. The first job's child + # context is complete then, so its function does not run again. invoke() + # sets the flag before the child context, so the second job still + # checks the state before delivery. + endpoint = fake_endpoint(accept_headers=suspend_after_ack) + client = fake_client( + endpoint=endpoint.url, states=["RUNNING", "SUSPENDED", "RUNNING"] + ) + handler, invocations = handler_across_deploy( + client, {}, {}, suspend_then_wait_then_job + ) + sent_in = record_invocation_of_each_request(endpoint, invocations) + + result = run(handler, None, [("build", succeed_with(1)), ("test", succeed_with(2))]) + + assert outcome(result) == {"result": [1, 2]} + assert sent_in[1] > sent_in[0] + 1 + assert len(client.requests("get_microvm")) == 3 + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert len(endpoint.requests) == 2 + + +def test_resumes_through_the_fallback_after_a_map_item_lost_the_flag( + run: Any, + fake_client: Any, + fake_endpoint: Any, + suspend_after_ack: Any, + connect_failed: Any, +) -> None: + # 1. Item a asks for a suspend after its job. + # 2. The job after the map runs in a later invocation, after a wait. + # 3. The map is complete then, so its items do not run again, and the + # flag is lost. + # 4. So the job goes to the endpoint of a suspended MicroVM, which + # answers 502 MICROVM_CONNECT_FAILED. The step checks the state, + # resumes the MicroVM, and sends the job again. + endpoint = fake_endpoint([(202, suspend_after_ack), 202, connect_failed]) + client = fake_client( + endpoint=endpoint.url, + states=["RUNNING", "RUNNING", "SUSPENDED", "RUNNING"], + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + shards: list[Any] = ctx.map( + ["a", "b"], + lambda item_context, item, index, items: vm.invoke( + item_context, + f"shard-{item}", + job_input=item, + options=job(suspend_after=item == "a"), + ), + name="shards", + config=MapConfig(max_concurrency=1), + ).get_results() + ctx.wait(Duration.from_seconds(1), name="pause") + return [shards, vm.invoke(ctx, "test", job_input=2, options=job())] + + result = run( + session_handler(session_config(client), body), + None, + [ + ("shard-a", succeed_with("A")), + ("shard-b", succeed_with("B")), + ("test", succeed_with(2)), + ], + ) + + assert outcome(result) == {"result": [["A", "B"], 2]} + # Item a and item b check the state before delivery. The job after the + # map checks it after the 502, and once more after the resume. + assert len(client.requests("get_microvm")) == 4 + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert [request["body"]["input"] for request in endpoint.requests] == [ + "a", + "b", + 2, + 2, + ] + + +def test_a_session_that_never_suspends_checks_the_state_after_no_answer( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + endpoint = fake_endpoint([None]) + client = fake_client(endpoint=endpoint.url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke(ctx, "build", job_input=1, options=job()) + + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with("done"))], + ) + + assert outcome(result) == {"result": "done"} + assert len(client.requests("get_microvm")) == 1 + assert client.requests("resume_microvm") == [] + assert len(endpoint.requests) == 2 + + +def test_a_session_that_never_suspends_makes_no_state_check_for_a_boot_answer( + run: Any, fake_client: Any, fake_endpoint: Any +) -> None: + # A plain 502, 503, or 504 is what the endpoint can answer while the + # MicroVM starts. It is no sign of a suspend. + endpoint = fake_endpoint([502, 503, 504]) + client = fake_client(endpoint=endpoint.url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke(ctx, "build", job_input=1, options=job()) + + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with("done"))], + ) + + assert outcome(result) == {"result": "done"} + assert client.requests("get_microvm") == [] + assert len(endpoint.requests) == 4 + + +def test_keeps_sending_after_a_denied_state_check_and_names_it_at_the_end( + run: Any, + fake_client: Any, + fake_endpoint: Any, + make_error: Any, + connect_failed: Any, + caplog: pytest.LogCaptureFixture, +) -> None: + # The failures end at 0.25, 0.75, 1.75, and 3.75 seconds. The checks run + # at 0.25 and 3.75 seconds, and both are denied. The next delay of 4 + # seconds does not fit in the window. + endpoint = fake_endpoint([connect_failed] * 10) + client = fake_client( + endpoint=endpoint.url, + states=[make_error("AccessDeniedException", 403)] * 10, + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke( + ctx, + "build", + job_input=1, + options=job(retry_window=Duration.from_seconds(5)), + ) + + with caplog.at_level(logging.WARNING): + result = run(session_handler(session_config(client), body)) + + error = outcome(result) + assert error["error"] == "MicrovmDeliveryError" + assert error["message"].endswith( + "Last failure: HTTP 502. GetMicrovm was denied, so the operation could " + "not check whether the MicroVM is suspended." + ) + assert len(client.requests("get_microvm")) == 2 + assert len(endpoint.requests) == 5 + assert warnings(caplog, STATE_CHECK_WARNING) == 1 + + +def test_fails_at_once_when_resume_is_denied_for_a_suspended_microvm( + run: Any, + fake_client: Any, + fake_endpoint: Any, + make_error: Any, + connect_failed: Any, +) -> None: + endpoint = fake_endpoint([connect_failed]) + client = fake_client( + endpoint=endpoint.url, + states=["SUSPENDED"], + resumes=[make_error("AccessDeniedException", 403)], + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke(ctx, "build", job_input=1, options=job()) + + # The default strategy decides on the error of the step. + result = run(session_handler(session_config(client, retry_strategy=None), body)) + + assert outcome(result) == { + "error": "MicrovmDeliveryError", + "message": ( + 'MicroVM "build": the job could not be delivered: MicroVM mvm-1 is ' + "SUSPENDED. ResumeMicrovm was denied, so the operation cannot resume " + "it. The function's execution role needs lambda:ResumeMicrovm." + ), + } + assert len(client.requests("resume_microvm")) == 1 + assert len(endpoint.requests) == 1 + + +def test_fails_with_not_running_when_the_check_finds_the_microvm_terminated( + run: Any, fake_client: Any, fake_endpoint: Any, connect_failed: Any +) -> None: + endpoint = fake_endpoint([connect_failed]) + client = fake_client(endpoint=endpoint.url, states=["TERMINATED"]) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke(ctx, "build", job_input=1, options=job()) + + result = run(session_handler(session_config(client), body)) + + error = outcome(result) + assert error["error"] == "MicrovmNotRunningError" + assert "MicroVM mvm-1 is TERMINATED." in error["message"] + assert len(endpoint.requests) == 1 + + +def test_a_flagged_session_keeps_the_strict_state_check( + run: Any, fake_client: Any, fake_endpoint: Any, make_error: Any +) -> None: + # After a suspend_after job, the state check before delivery fails the + # job for a denied GetMicrovm, as in a session with auto-suspend. + endpoint = fake_endpoint() + client = fake_client( + endpoint=endpoint.url, + states=["RUNNING", make_error("AccessDeniedException", 403)], + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + build = vm.invoke(ctx, "build", job_input=1, options=job(suspend_after=True)) + return [build, vm.invoke(ctx, "test", job_input=2, options=job())] + + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with(1))], + ) + + error = outcome(result) + assert error["error"] == "MicrovmDeliveryError" + assert "AccessDeniedException message" in error["message"] + assert len(endpoint.requests) == 1 + + +@pytest.mark.parametrize( + ("acknowledged", "expected"), [(False, 1), (True, 0)], ids=["old", "new"] +) +def test_warns_once_per_session_when_the_worker_ignores_suspend_after( + run: Any, + fake_client: Any, + fake_endpoint: Any, + suspend_after_ack: Any, + caplog: pytest.LogCaptureFixture, + acknowledged: bool, # noqa: FBT001 - a pytest parameter + expected: int, +) -> None: + endpoint = fake_endpoint(accept_headers=suspend_after_ack if acknowledged else {}) + client = fake_client(endpoint=endpoint.url) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + options = job(suspend_after=True) + return [ + vm.invoke(ctx, "build", job_input=1, options=options), + vm.invoke(ctx, "test", job_input=2, options=options), + ] + + with caplog.at_level(logging.WARNING): + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with(1)), ("test", succeed_with(2))], + ) + + assert outcome(result) == {"result": [1, 2]} + assert len(endpoint.requests) == 2 + assert warnings(caplog, IGNORED_WARNING) == expected + + +def test_reads_no_acknowledgement_for_a_job_without_suspend_after( + run: Any, fake_client: Any, fake_endpoint: Any, caplog: pytest.LogCaptureFixture +) -> None: + client = fake_client(endpoint=fake_endpoint().url) + + with caplog.at_level(logging.WARNING): + result = run( + session_handler(session_config(client), two_jobs), + None, + [("build", succeed_with(1)), ("test", succeed_with(2))], + ) + + assert outcome(result)["result"]["test"] == 2 + assert warnings(caplog, IGNORED_WARNING) == 0 + + +def test_remembers_the_warning_of_the_latest_sessions_only( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(session_module, "_WARNED_SESSIONS", OrderedDict()) + monkeypatch.setattr(session_module, "_MAX_WARNED_SESSIONS", 2) + a, b, c = ( + MicrovmSession( + None, # type: ignore[arg-type] + None, # type: ignore[arg-type] + resume_before_delivery=False, + session_id=session_id, + ) + for session_id in ["a", "b", "c"] + ) + + assert [handle._claim_suspend_after_warning() for handle in (a, b, c)] == [ # noqa: SLF001 + True, + True, + True, + ] + # Each session warns once. + assert c._claim_suspend_after_warning() is False # noqa: SLF001 + # The oldest session was forgotten. + assert a._claim_suspend_after_warning() is True # noqa: SLF001 diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/validation_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/validation_test.py index 8623f2a00..9cd63bc5d 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/validation_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/validation_test.py @@ -15,6 +15,7 @@ from aws_durable_execution_sdk_python_extras.microvm import ( MAX_JOB_INPUT_BYTES, MicrovmConfig, + MicrovmInvokeOptions, MicrovmRequestConfig, microvm, ) @@ -328,3 +329,36 @@ def test_accepts_dots_that_are_not_a_dot_segment(path: str) -> None: operation._validate_request("build", MicrovmRequestConfig(path=path)) # noqa: SLF001 == 60 ) + + +@pytest.mark.parametrize( + ("suspend_after", "message"), + [ + (True, "suspend_after applies only to a session job"), + (1, "suspend_after must be a bool"), + (0, "suspend_after must be a bool"), + ("true", "suspend_after must be a bool"), + (None, "suspend_after must be a bool"), + ], +) +def test_rejects_suspend_after_in_the_request_of_microvm( + suspend_after: Any, message: str +) -> None: + # MicrovmInvokeOptions subclasses MicrovmRequestConfig. So it passes the + # type of config.request, and microvm() checks the field itself. + request = MicrovmInvokeOptions( + timeout=Duration.from_minutes(1), + suspend_after=suspend_after, + ) + + with pytest.raises(ValidationError, match=message): + microvm(NoOperations(), "build", job_input={}, config=config(request=request)) # type: ignore[arg-type] + + +def test_accepts_invoke_options_without_suspend_after_as_the_request_of_microvm() -> ( + None +): + request = MicrovmInvokeOptions(timeout=Duration.from_minutes(1), path="/build") + + with pytest.raises(AssertionError, match="no durable operation may start"): + microvm(NoOperations(), "build", job_input={}, config=config(request=request)) # type: ignore[arg-type] diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py index 06aaf0464..e7bf26bea 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py @@ -20,10 +20,12 @@ import functools import http.client import json +import logging import socket import threading import time from collections.abc import Callable, Iterator +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Any import pytest @@ -406,15 +408,17 @@ def post_hook(worker: MicrovmWorker, hook: str) -> None: class RecordingLogger: - """Records the worker's log messages.""" + """Records the worker's log messages and their fields.""" def __init__(self) -> None: self._lock = threading.Lock() self.messages: list[str] = [] + self._extras: list[tuple[str, dict[str, Any]]] = [] def _add(self, msg: object, *args: object, extra: Any = None) -> None: with self._lock: self.messages.append(str(msg)) + self._extras.append((str(msg), dict(extra or {}))) info = warning = error = _add @@ -422,6 +426,11 @@ def count(self, message: str) -> int: with self._lock: return self.messages.count(message) + def extras(self, message: str) -> list[dict[str, Any]]: + """The fields of each line with this message.""" + with self._lock: + return [data for m, data in self._extras if m == message] + @pytest.mark.parametrize("resume_hook", [True, False], ids=["resume-hook", "none"]) def test_resumes_a_worker_that_suspended_itself_between_two_session_jobs( @@ -557,3 +566,354 @@ def handler(event: Any, context: DurableContext) -> Any: assert refusals == 0 else: assert refusals >= 1 + + +# --------------------------------------------------------------------------- +# suspend_after with the real worker +# --------------------------------------------------------------------------- + +IGNORED_WARNING = ( + "the worker in MicroVM mvm-1 does not support suspendAfter. The MicroVM " + "keeps running after the job." +) + + +class SimulatedMicrovm: + """A stand-in for Lambda that keeps the state of one MicroVM. + + 1. The worker's SuspendMicrovm sends the ``suspend`` hook, and then sets + the state to ``SUSPENDED``. Lambda calls the hook before the freeze. + 2. The session's ResumeMicrovm sends the ``resume`` hook, and then sets + the state to ``RUNNING``. A probe saw the hook before ``RUNNING``. + 3. :attr:`endpoint` stands in for the MicroVM endpoint. It forwards each + job to the worker while the MicroVM runs. While it is suspended, it + answers HTTP 502 with ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``, + as the service does. + + ``old_worker`` removes ``suspendAfter`` from each forwarded job. The + worker then behaves as a release that does not know the field. + """ + + def __init__(self, *, old_worker: bool = False) -> None: + self._lock = threading.Lock() + self._state = "RUNNING" + self.worker: MicrovmWorker | None = None + self.suspended: list[str] = [] + self.answers: list[int] = [] + self._old_worker = old_worker + simulated = self + + class Endpoint(BaseHTTPRequestHandler): + def do_POST(self) -> None: # noqa: N802 - the http.server name + body = self.rfile.read(int(self.headers.get("content-length", "0"))) + if simulated.state() == "SUSPENDED": + simulated.record(502) + self.send_response(502) + self.send_header("x-aws-proxy-error", "MICROVM_CONNECT_FAILED") + self.send_header("content-length", "0") + self.end_headers() + return + status, headers, data = simulated.forward(self.path, body) + simulated.record(status) + self.send_response(status) + for name, value in headers.items(): + self.send_header(name, value) + self.send_header("content-length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def log_message(self, *args: Any) -> None: + pass + + self._server = ThreadingHTTPServer(("127.0.0.1", 0), Endpoint) + threading.Thread(target=self._server.serve_forever, daemon=True).start() + + @property + def endpoint(self) -> str: + return f"http://127.0.0.1:{self._server.server_address[1]}" + + def state(self) -> str: + with self._lock: + return self._state + + def record(self, status: int) -> None: + with self._lock: + self.answers.append(status) + + def set_state(self, state: str) -> None: + with self._lock: + self._state = state + + def forward(self, path: str, body: bytes) -> tuple[int, dict[str, str], bytes]: + worker = self.worker + assert worker is not None + if self._old_worker: + document = json.loads(body) + document.pop("suspendAfter", None) + body = json.dumps(document).encode() + connection = http.client.HTTPConnection( + host_of(worker), worker.port, timeout=WAIT + ) + try: + connection.request( + "POST", path, body=body, headers={"content-type": "application/json"} + ) + response = connection.getresponse() + data = response.read() + ack = response.getheader("X-Microvm-Suspend-After") + return ( + response.status, + ({} if ack is None else {"X-Microvm-Suspend-After": ack}), + data, + ) + finally: + connection.close() + + def wait_for(self, state: str) -> str: + deadline = time.monotonic() + WAIT + while time.monotonic() < deadline: + if self.state() == state: + return state + time.sleep(0.01) + return f"not {state}" + + def close(self) -> None: + self._server.shutdown() + self._server.server_close() + + +Pipeline = Callable[[MicrovmSession, DurableContext], Any] + + +def suspend_after_session( + fake_client: Any, + make_pipeline: Callable[[SimulatedMicrovm], Pipeline], + job: Callable[[Any, MicrovmJobContext], Any], + *, + old_worker: bool = False, +) -> tuple[Any, Any, SimulatedMicrovm, RecordingLogger]: + """Run a pipeline in a session without auto-suspend, against a real + worker behind a :class:`SimulatedMicrovm`. + + ``make_pipeline`` receives the simulated MicroVM, so the pipeline can + wait for its state. + """ + microvm = SimulatedMicrovm(old_worker=old_worker) + pipeline = make_pipeline(microvm) + holder: dict[str, Any] = {} + + class SessionMicrovms(fake_client): # type: ignore[misc, valid-type] + def get_microvm(self, **request: Any) -> dict[str, Any]: + response = super().get_microvm(**request) + return {**response, "state": microvm.state()} + + def resume_microvm(self, **request: Any) -> dict[str, Any]: + super().resume_microvm(**request) + post_hook(holder["worker"], "resume") + microvm.set_state("RUNNING") + return {} + + class WorkerMicrovms: + """The worker's lambda-microvms client.""" + + def suspend_microvm(self, *, microvmIdentifier: str) -> dict[str, Any]: # noqa: N803 + microvm.suspended.append(microvmIdentifier) + post_hook(holder["worker"], "suspend") + microvm.set_state("SUSPENDED") + return {} + + @durable_execution + def handler(event: Any, context: DurableContext) -> Any: + return microvm_session( + context, "pipeline", config=holder["config"], handler=pipeline + ) + + runner = DurableFunctionTestRunner(handler=handler, poll_interval=0.05) + callbacks = RunnerCallbacks(runner) + logger = RecordingLogger() + worker = start_reachable_worker( + lambda: start_microvm_worker( + handler=job, + create_client=lambda region: callbacks, # type: ignore[arg-type, return-value] + create_microvms_client=lambda region: WorkerMicrovms(), # type: ignore[arg-type, return-value] + logger=logger, + port=0, + ) + ) + holder["worker"] = worker + microvm.worker = worker + try: + client = SessionMicrovms( + launches=[run_hook_to(worker)], endpoint=microvm.endpoint + ) + holder["config"] = MicrovmSessionConfig( + image_identifier="img", + execution_role_arn="role", + timeout=Duration.from_minutes(5), + client=client, + ) + arn = runner.run_async(input="null", execution_timeout=WAIT) + result = runner.wait_for_result(arn, timeout=WAIT) + assert worker.idle(WAIT) + finally: + worker.close() + runner.close() + microvm.close() + assert result.status is InvocationStatus.SUCCEEDED, result.error + return json.loads(result.result or "null"), client, microvm, logger + + +def options(*, suspend_after: bool = False) -> MicrovmInvokeOptions: + return MicrovmInvokeOptions( + timeout=Duration.from_minutes(1), suspend_after=suspend_after + ) + + +def echo(job_input: Any, context: MicrovmJobContext) -> Any: + return f"{job_input} done" + + +def test_suspends_after_a_job_and_the_next_job_resumes_the_microvm( + fake_client: Any, +) -> None: + def make_pipeline(microvm: SimulatedMicrovm) -> Pipeline: + def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: + first = vm.invoke( + ctx, "build", job_input="build", options=options(suspend_after=True) + ) + between = ctx.step(lambda _: microvm.wait_for("SUSPENDED"), name="pause") + second = vm.invoke(ctx, "test", job_input="test", options=options()) + return [first, between, second] + + return pipeline + + outcome, client, microvm, logger = suspend_after_session( + fake_client, make_pipeline, echo + ) + + assert outcome == ["build done", "SUSPENDED", "test done"] + assert microvm.suspended == ["mvm-1"] + [line] = logger.extras("suspending the MicroVM after a job") + assert list(line) == ["microvmId", "callbackId"] + assert line["microvmId"] == "mvm-1" + # The flag makes the second job check the state before delivery. So it + # resumes the MicroVM before it sends the job, and gets no 502. + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert 502 not in microvm.answers + assert microvm.state() == "RUNNING" + assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}] + + +def test_suspends_after_a_map_with_one_suspend_after_item(fake_client: Any) -> None: + def slow_echo(job_input: Any, context: MicrovmJobContext) -> Any: + # The items run at once. So the worker suspends only when the last + # of them ends. The sleep keeps each job running until every item + # has delivered its job. + time.sleep(1) + return f"{job_input} done" + + def make_pipeline(microvm: SimulatedMicrovm) -> Pipeline: + def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: + shards: list[Any] = ctx.map( + ["a", "b", "c"], + lambda item_context, item, index, items: vm.invoke( + item_context, + f"shard-{item}", + job_input=item, + options=options(suspend_after=item == "a"), + ), + name="shards", + ).get_results() + between = ctx.step(lambda _: microvm.wait_for("SUSPENDED"), name="pause") + after = vm.invoke(ctx, "after", job_input="after", options=options()) + return [shards, between, after] + + return pipeline + + outcome, client, microvm, logger = suspend_after_session( + fake_client, make_pipeline, slow_echo + ) + + assert outcome == [["a done", "b done", "c done"], "SUSPENDED", "after done"] + # One request from item a, and one suspend. + assert microvm.suspended == ["mvm-1"] + assert len(logger.extras("suspending the MicroVM after a job")) == 1 + # The job after the map resumed the MicroVM, with the flag or after a + # 502, depending on which invocation ran item a. + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert microvm.state() == "RUNNING" + + +def test_a_replay_that_lost_the_flag_resumes_through_the_check_after_a_502( + fake_client: Any, +) -> None: + # 1. The map item asks for a suspend after its job. + # 2. A durable wait ends the invocation. The next invocation replays + # the completed map without running its item. So the flag is lost. + # 3. The next job goes to the suspended MicroVM, and gets a 502 + # MICROVM_CONNECT_FAILED. The request step checks the state, resumes + # the MicroVM, and sends the job again. + def make_pipeline(microvm: SimulatedMicrovm) -> Pipeline: + def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: + shards: list[Any] = ctx.map( + ["a"], + lambda item_context, item, index, items: vm.invoke( + item_context, + f"shard-{item}", + job_input=item, + options=options(suspend_after=True), + ), + name="shards", + ).get_results() + between = ctx.step(lambda _: microvm.wait_for("SUSPENDED"), name="pause") + ctx.wait(Duration.from_seconds(1), name="wait") + after = vm.invoke(ctx, "after", job_input="after", options=options()) + return [shards, between, after] + + return pipeline + + outcome, client, microvm, _ = suspend_after_session( + fake_client, make_pipeline, echo + ) + + assert outcome == [["a done"], "SUSPENDED", "after done"] + assert microvm.suspended == ["mvm-1"] + # The job after the wait made no state check before delivery. Its first + # request got the 502. The check after it found SUSPENDED, resumed the + # MicroVM, and polled until RUNNING. The second request reached the + # worker. + assert microvm.answers[-2:] == [502, 202] + assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}] + assert microvm.state() == "RUNNING" + + +def test_an_old_worker_ignores_suspend_after_and_the_session_warns_once( + fake_client: Any, caplog: pytest.LogCaptureFixture +) -> None: + def make_pipeline(microvm: SimulatedMicrovm) -> Pipeline: + def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: + return [ + vm.invoke( + ctx, name, job_input=name, options=options(suspend_after=True) + ) + for name in ["build", "test"] + ] + + return pipeline + + with caplog.at_level(logging.WARNING): + outcome, client, microvm, logger = suspend_after_session( + fake_client, make_pipeline, echo, old_worker=True + ) + + assert outcome == ["build done", "test done"] + # The worker sent no acknowledgement and never suspended the MicroVM. + assert microvm.suspended == [] + assert logger.extras("suspending the MicroVM after a job") == [] + warnings = [ + record + for record in caplog.records + if record.levelname == "WARNING" and record.getMessage() == IGNORED_WARNING + ] + assert len(warnings) == 1 + assert client.requests("resume_microvm") == [] From d0583352f999753318fdddaba83c536b685b072d Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 25 Oct 2026 22:59:31 -0700 Subject: [PATCH 07/12] fix(microvm-worker): drop the request in a refusal A suspend refusal means that the MicroVM is suspending already. So a suspend request that meets one is done. The suspend hook now clears a pending request, and a suspendAfter job that ends during a refusal sets none. Before, the request fired after the next job, even when that job did not ask for it. Tests cover a job that ends between the suspend and resume hooks, and a request that waits when the suspend hook arrives. The terminate test now asserts that the request is cleared. --- .../README.md | 4 +- .../worker.py | 26 ++++++---- .../tests/lifecycle_hooks_test.py | 49 +++++++++++++++++++ 3 files changed, 68 insertions(+), 11 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md index 35ff83872..a20a861c8 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -130,9 +130,9 @@ A session job can carry `"suspendAfter": true`. The durable function sets it for 1. Only a JSON `true` turns it on. The worker ignores any other value, and does not reject the job. The `run` hook payload never carries the field. 2. The worker answers the job with HTTP 202 and the header `X-Microvm-Suspend-After: accepted`. A second delivery of the job that gets the duplicate HTTP 202 carries the header when the first delivery carried the field. A second delivery refused with HTTP 503 carries no header. An older worker ignores the field and sends no header. The durable function then logs a warning, and the MicroVM keeps running. -3. When the job has ended, and its outcome reports have ended, the worker sets a pending suspend request. This happens whether the handler succeeded or failed. +3. When the job has ended, and its outcome reports have ended, the worker sets a pending suspend request. This happens whether the handler succeeded or failed. A job that ends while the worker refuses jobs for a suspend sets no request. That refusal comes from the worker's own call, a `suspend` hook, or the 30 seconds that follow either. The MicroVM is suspending already. So the request is done. 4. The request fires when no job runs, no callback failure report runs, no `run` hook request is being handled, no `SuspendMicrovm` call runs, and the worker is not terminating, not closed, and not refusing jobs for a suspend. The worker checks these conditions again when a job ends, a report ends, a `run` hook request ends, and a call ends. So a request from a job that ends while another job runs fires when the last job ends. -5. Firing clears the request before the call starts. A `resume` hook never sets it again. A call after the idle time clears it too, because that call suspends the MicroVM as well. +5. Firing clears the request before the call starts. A `resume` hook never sets it again. A call after the idle time clears it too, because that call suspends the MicroVM as well. A `suspend` hook clears it too, because Lambda suspends the MicroVM after that hook. 6. Firing stops the idle time, when the session also suspends when idle. 7. The `terminate` hook and `worker.close()` drop the request. 8. The worker calls `SuspendMicrovm` with the MicroVM ID of the first accepted `run` hook with a payload. It never uses the MicroVM ID of a job request. A forged request could otherwise suspend another MicroVM of the same image. Without such a `run` hook, the worker logs the warning `cannot suspend the MicroVM after a job: the run hook named no MicroVM`, with `callbackId`, and does not suspend. diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py index 785b4589b..8e24dcb25 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py @@ -421,15 +421,17 @@ class MicrovmWorkerListener: the first one carried the field. A 503 never carries it. Only a JSON ``true`` counts. Any other value is ignored. 2. When the job's outcome reports have ended, the listener sets a - pending request. The request fires when no job, failure report, - ``run`` hook request, or SuspendMicrovm call runs, and the listener - is not terminating, not closed, and not refusing jobs. Each of those - ends checks it again. + pending request. A job that ends while the listener refuses jobs for + a suspend sets none, because the MicroVM is suspending already. The + request fires when no job, failure report, ``run`` hook request, or + SuspendMicrovm call runs, and the listener is not terminating, not + closed, and not refusing jobs. Each of those ends checks it again. 3. Firing clears the request, stops the idle time, and makes the same call as items 3 to 9, with the MicroVM ID of the first accepted ``run`` hook with a payload. Without that hook, the listener logs a warning and does not suspend. A ``resume`` hook never repeats the - request. The ``terminate`` hook and ``close()`` drop it. + request. A ``suspend`` hook, the ``terminate`` hook, and ``close()`` + drop it. 4. A call refused for now is retried only when the session suspends when idle. @@ -529,7 +531,7 @@ def __init__( # The callback ID of the suspendAfter job whose end asked for a # suspend that has not started yet. The request fires when no work # runs. A call that starts clears it, so a resume hook never repeats - # it. The terminate hook and close() clear it too. + # it. A suspend hook, the terminate hook, and close() clear it too. self._suspend_requested: str | None = None # Each timer action carries the generation it was started with, and # acts only when it is still current. Cancelling a timer increments @@ -808,12 +810,14 @@ def _handle_suspend_hook(self) -> None: """Refuse jobs until the ``resume`` hook, or for 30 seconds. Lambda freezes the MicroVM after the answer. A job accepted now would - freeze with it. + freeze with it. Lambda suspends the MicroVM after this hook. So a + pending suspend after a job is done, and it is cleared. """ with self._locked(): self._suspend_hooks += 1 self._stop_idle_time() self._suspending = True + self._suspend_requested = None self._start_grace() def _handle_resume_hook(self) -> list[tuple[MicrovmJobDocument, str]]: @@ -1243,7 +1247,9 @@ def _run_job( # The job's report attempts have ended here, whether the handler # succeeded or failed. So a suspendAfter job asks for the suspend # now. The suspend runs after this block, so that no lock is - # held while it logs or starts its thread. + # held while it logs or starts its thread. During a suspend + # refusal, the MicroVM is suspending already. So a job that ends + # then asks for nothing. with self._locked(): callback_id = job.callback_id suspend_after = callback_id in self._suspend_after_jobs @@ -1254,7 +1260,9 @@ def _run_job( if len(self._finished) > MAX_FINISHED_JOBS: self._finished.popitem(last=False) self._start_idle_time() - if suspend_after and not (self._terminating or self._closed): + if suspend_after and not ( + self._terminating or self._closed or self._suspending + ): self._suspend_requested = callback_id self._changed.notify_all() decision = self._take_job_suspend() diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py index ec968f4e5..ad7c85bc2 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py @@ -3621,6 +3621,53 @@ def test_refuses_a_resent_running_suspend_after_job_without_acknowledgement( assert h.microvms.suspended == ["mvm-1"] +@pytest.mark.parametrize("refusal_end", ["resume", "grace"]) +def test_a_job_that_ends_during_a_suspend_hook_refusal_asks_for_nothing( + start: Any, refusal_end: str +) -> None: + # Lambda suspends the MicroVM after the suspend hook. So a suspendAfter + # job that ends during the refusal asks for no call of the worker's own, + # neither then nor after the refusal ends. + h = start() + h.run_hook() + assert post_job(h, "cb-1", suspendAfter=True) == (202, "accepted") + assert eventually(lambda: h.started("cb-1")) + assert h.hook("suspend", LAMBDA_HEADERS) == 200 + + end_job(h, "cb-1") + assert h.microvms.suspended == [] + assert h.listener._suspend_requested is None # noqa: SLF001 + + if refusal_end == "resume": + assert h.hook("resume", LAMBDA_HEADERS) == 200 + else: + h.advance(GRACE) + post_job(h, "cb-2") + end_job(h, "cb-2") + + assert h.jobs.started() == ["cb-1", "cb-2"] + assert h.microvms.suspended == [] + + +def test_a_suspend_hook_clears_a_pending_suspend_after_a_job(start: Any) -> None: + # The request waits for the other job. Lambda then suspends the MicroVM + # after the suspend hook. So the request is done, and the other job's end + # after the resume fires nothing. + h = start() + h.run_hook() + post_job(h, "cb-1", suspendAfter=True) + post_job(h, "cb-2") + end_job(h, "cb-1") + assert h.listener._suspend_requested == "cb-1" # noqa: SLF001 + + assert h.hook("suspend", LAMBDA_HEADERS) == 200 + assert h.listener._suspend_requested is None # noqa: SLF001 + assert h.hook("resume", LAMBDA_HEADERS) == 200 + end_job(h, "cb-2") + + assert h.microvms.suspended == [] + + def test_acknowledges_a_resent_running_suspend_after_job_after_the_terminate_hook( start: Any, ) -> None: @@ -3886,10 +3933,12 @@ def test_does_not_suspend_after_the_terminate_hook(start: Any) -> None: end_job(h, "cb-1") assert h.hook("terminate", LAMBDA_HEADERS) == 200 + assert h.listener._suspend_requested is None # noqa: SLF001 h.jobs.release_all() assert h.worker.idle(WAIT) assert h.microvms.suspended == [] + assert h.listener._suspend_requested is None # noqa: SLF001 def test_does_not_suspend_after_close(start: Any) -> None: From 4e26ee455be2dcb08093c4ac8b58a309a2a87806 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 25 Oct 2026 22:59:31 -0700 Subject: [PATCH 08/12] docs(extras): correct the suspend_after replay note A replay of a completed job calls vm.invoke again, so it sets the suspend_after record again. Only a completed map item, parallel branch, or child context skips its function. The suspend_after job itself checks the state, because the record is set before the job starts. A worker that takes more than 10 seconds to start listening gives the first delivery no answer. The operation then makes one GetMicrovm call, so a session that never suspends can make one extra call. The warning about a worker that ignores the field is logged once per session in each execution environment. The README snippet sets heartbeat_timeout on both jobs, as the example does. --- .../README.md | 14 ++++++++------ .../microvm/config.py | 7 +++++-- .../microvm/request.py | 13 +++++++------ .../microvm/session.py | 6 ++++-- .../tests/microvm/request_test.py | 4 ++-- 5 files changed, 26 insertions(+), 18 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index 7ae933303..751367173 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -280,7 +280,9 @@ build = vm.invoke( "build", job_input=event, options=MicrovmInvokeOptions( - timeout=Duration.from_minutes(15), suspend_after=True + timeout=Duration.from_minutes(15), + heartbeat_timeout=Duration.from_seconds(30), + suspend_after=True, ), ) approval = ctx.create_callback(name="approval").result() # a long wait @@ -297,17 +299,17 @@ vm.invoke( What happens: -1. The job request carries `"suspendAfter": true`. The job document stays at version 1, and the field is omitted for every other job. +1. The job request carries `"suspendAfter": true`. The job document stays at version 1, and the field is omitted for every other job. The run hook payload never carries it. 2. The worker acknowledges the field with the header `X-Microvm-Suspend-After: accepted` in its 202 answer. 3. The worker suspends the MicroVM when the job has ended and its outcome is reported, and no other job runs. A job that ends while another job runs leaves the request pending. The worker then suspends the MicroVM when the last running job ends. 4. The worker suspends the MicroVM that its `run` hook named. It never uses the MicroVM ID of a job request. -5. Each later job of the session handle checks the MicroVM state before delivery, and resumes a suspended MicroVM, as with `auto_suspend_on_idle`. +5. Each later job of the session handle checks the MicroVM state before delivery, and resumes a suspended MicroVM, as with `auto_suspend_on_idle`. The `suspend_after` job itself checks the state too, because the record is set before the job starts. The setting works with or without `auto_suspend_on_idle`. With both, the worker suspends right after the job, and the idle time applies after each resume. -The record that a job set `suspend_after` is kept in memory only. A replay of a completed job, `map` item, `parallel` branch, or child context does not run its `vm.invoke` again. So the next job can skip the state check of a suspended MicroVM. For that case, every session job checks the state after a sign of a suspended MicroVM: +The record that a job set `suspend_after` is kept in memory only. A replay of a completed job calls `vm.invoke` again, so it sets the record again. A replay of a completed `map` item, `parallel` branch, or child context does not run its function again. So the next job can skip the state check of a suspended MicroVM. For that case, every session job checks the state after a sign of a suspended MicroVM: -1. The sign is an HTTP 502 with `x-aws-proxy-error: MICROVM_CONNECT_FAILED`, or a request without an answer. A booting MicroVM gives neither sign, because its endpoint holds the request until the worker listens. So a session that never suspends makes no extra call. +1. The sign is an HTTP 502 with `x-aws-proxy-error: MICROVM_CONNECT_FAILED`, or a request without an answer. A booting MicroVM's endpoint holds the request until the worker listens. A boot gives no such sign, unless the worker takes more than 10 seconds to start listening. Then the first delivery gets no answer, and the operation makes one GetMicrovm call. 2. The request step calls `GetMicrovm` once. It resumes a `SUSPENDING` or `SUSPENDED` MicroVM, waits for `RUNNING`, and sends the job again. 3. The check runs again after each failed request, at most once per 2 seconds. 4. A denied or failed `GetMicrovm` logs one warning per job: `GetMicrovm failed, so the operation could not check whether the MicroVM is suspended. The request step keeps sending the job.` The request step then keeps sending the job. After a denied call, the final `MicrovmEndpointUnavailableError` message adds `GetMicrovm was denied, so the operation could not check whether the MicroVM is suspended.` @@ -328,7 +330,7 @@ It needs: 2. A `heartbeat_timeout` on the next job of at least its `retry_window` plus 30 seconds. The heartbeat timeout starts before delivery, so it includes the resume and the refusal. 3. `lambda:GetMicrovm` and `lambda:ResumeMicrovm` in the function's role. 4. `lambda:SuspendMicrovm` on the image in the MicroVM's role. That permission lets code in any MicroVM of the image suspend any other. So use one image per trust boundary. -5. A worker that supports the field. An older worker ignores it, and the MicroVM keeps running. The session then logs one warning: `the worker in MicroVM does not support suspendAfter. The MicroVM keeps running after the job.` +5. A worker that supports the field. An older worker ignores it, and the MicroVM keeps running. Each execution environment then logs one warning per session: `the worker in MicroVM does not support suspendAfter. The MicroVM keeps running after the job.` A suspend that the service refuses for now, such as a throttle or expired credentials, is not retried without `auto_suspend_on_idle`. The MicroVM then keeps running until a later suspend. With `auto_suspend_on_idle`, the worker tries again after the idle time, as it does for an idle suspend. diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py index a9b2d0c9e..14301c0f3 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py @@ -367,7 +367,9 @@ class MicrovmInvokeOptions(MicrovmRequestConfig): reported, and no other job runs. A job that ends while another job runs leaves the request pending. The worker then suspends the MicroVM when the last running job ends. The next job of the session resumes the - MicroVM before it is delivered. The setting works with or without + MicroVM before it is delivered. The ``suspend_after`` job itself checks + the state too, because the record is set before the job starts. The + setting works with or without :attr:`MicrovmSessionConfig.auto_suspend_on_idle`. Set it on the job before a long wait, such as a durable wait or a @@ -384,7 +386,8 @@ class MicrovmInvokeOptions(MicrovmRequestConfig): ``lambda:ResumeMicrovm``. 3. The MicroVM's role has ``lambda:SuspendMicrovm`` on the image. 4. The worker supports the setting. An older worker ignores it, and the - MicroVM keeps running. The session then logs one warning. + MicroVM keeps running. Each execution environment then logs one + warning per session. A suspend that the service refuses for now is not retried, unless :attr:`MicrovmSessionConfig.auto_suspend_on_idle` is ``True``. Only diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py index 2d77cc0c9..b2c28fe05 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/request.py @@ -263,12 +263,13 @@ class SendJobOptions: MicroVM after a ``suspend_after`` job, and the session can lose the record of that job on replay. A suspended or suspending MicroVM gets no answer to a request, or its endpoint answers HTTP 502 with - ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``. A booting MicroVM gets - neither: its endpoint holds the request until the worker listens. So - the next attempt calls it after either sign, at most once per 2 - seconds. A session that never suspends then makes no extra call. It - receives what remains of the retry window, in seconds. ``None`` from it - keeps the current endpoint. + ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``. So the next attempt calls + it after either sign, at most once per 2 seconds. A booting MicroVM's + endpoint holds the request until the worker listens. A boot gives no + such sign, unless the worker takes more than 10 seconds to start + listening. Then the first delivery gets no answer, and the operation + makes one ``GetMicrovm`` call. It receives what remains of the retry + window, in seconds. ``None`` from it keeps the current endpoint. """ unavailable_note: Callable[[], str | None] | None = None """A sentence to add to the message of diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py index c88c3ec76..5a3f6a376 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py @@ -126,8 +126,10 @@ def microvm_session( delivery, as with ``auto_suspend_on_idle``. That record is kept in memory only. So every job of every session also checks the state after a sign of a suspended MicroVM: no answer, or an HTTP 502 with - ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``. A booting MicroVM gives - neither sign. So a session that never suspends makes no extra call. + ``x-aws-proxy-error: MICROVM_CONNECT_FAILED``. A boot gives no such + sign, unless the worker takes more than 10 seconds to start listening. + Then the first delivery gets no answer, and the operation makes one + ``GetMicrovm`` call. A failed :meth:`MicrovmSession.invoke` raises inside the handler. The handler can catch the error and continue with the same MicroVM. An diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py index 4560477cb..fac5b92df 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/request_test.py @@ -708,8 +708,8 @@ def check(left: float) -> str | None: def test_makes_no_suspend_check_for_another_failure( fake_client: Any, fake_endpoint: Any, answer: Any ) -> None: - # A booting MicroVM never gets MICROVM_CONNECT_FAILED. So a session that - # never suspends makes no extra call. + # A booting MicroVM never gets MICROVM_CONNECT_FAILED. So these answers + # make no extra call. endpoint = fake_endpoint([answer]) checks: list[float] = [] From c2edf55755476f7d1c8513d3e5786970e57fed46 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sun, 25 Oct 2026 22:59:33 -0700 Subject: [PATCH 09/12] test(extras): set the map concurrency explicitly The map test with one suspend_after item relies on items that run at once. MapConfig(max_concurrency=3) states that, instead of relying on the default. --- .../tests/microvm/worker_composed_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py index 66dd5983a..ca8b5b30e 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py @@ -29,7 +29,7 @@ from typing import Any import pytest -from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.config import Duration, MapConfig from aws_durable_execution_sdk_python.context import DurableContext from aws_durable_execution_sdk_python.execution import ( InvocationStatus, @@ -823,6 +823,7 @@ def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: options=options(suspend_after=item == "a"), ), name="shards", + config=MapConfig(max_concurrency=3), ).get_results() between = ctx.step(lambda _: microvm.wait_for("SUSPENDED"), name="pause") after = vm.invoke(ctx, "after", job_input="after", options=options()) From 39816683cf047ce341e88dda26665694bc730a42 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Mon, 26 Oct 2026 11:41:18 -0700 Subject: [PATCH 10/12] fix(extras): resume again after a conflict ResumeMicrovm answers ConflictException while a transition is in progress. The state check marked the resume as sent after that answer. So a MicroVM that was still suspending became SUSPENDED and was never resumed. The wait then ended with MicrovmEndpointUnavailableError and "Last state: SUSPENDED". The check after a sign of a suspended MicroVM had this flaw for a conflict during SUSPENDING. The check before delivery and the check after a gateway error had it for a conflict during SUSPENDED. A conflict now leaves the resume unaccepted. So the next SUSPENDED answer sends ResumeMicrovm again. A PENDING answer after a conflict still waits for RUNNING. The poll interval and the deadline bound the calls. --- .../README.md | 2 +- .../microvm/lifecycle.py | 59 ++++++--- .../tests/conftest.py | 13 +- .../tests/microvm/lifecycle_test.py | 116 +++++++++++++++++- .../tests/microvm/session_test.py | 93 +++++++++++++- 5 files changed, 254 insertions(+), 29 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index 751367173..6b1e287cf 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/packages/aws-durable-execution-sdk-python-extras/README.md @@ -310,7 +310,7 @@ The setting works with or without `auto_suspend_on_idle`. With both, the worker The record that a job set `suspend_after` is kept in memory only. A replay of a completed job calls `vm.invoke` again, so it sets the record again. A replay of a completed `map` item, `parallel` branch, or child context does not run its function again. So the next job can skip the state check of a suspended MicroVM. For that case, every session job checks the state after a sign of a suspended MicroVM: 1. The sign is an HTTP 502 with `x-aws-proxy-error: MICROVM_CONNECT_FAILED`, or a request without an answer. A booting MicroVM's endpoint holds the request until the worker listens. A boot gives no such sign, unless the worker takes more than 10 seconds to start listening. Then the first delivery gets no answer, and the operation makes one GetMicrovm call. -2. The request step calls `GetMicrovm` once. It resumes a `SUSPENDING` or `SUSPENDED` MicroVM, waits for `RUNNING`, and sends the job again. +2. The request step calls `GetMicrovm` once. It resumes a `SUSPENDING` or `SUSPENDED` MicroVM, waits for `RUNNING`, and sends the job again. A resume that the service rejects with a conflict while the MicroVM is still suspending is sent again once the MicroVM is suspended. 3. The check runs again after each failed request, at most once per 2 seconds. 4. A denied or failed `GetMicrovm` logs one warning per job: `GetMicrovm failed, so the operation could not check whether the MicroVM is suspended. The request step keeps sending the job.` The request step then keeps sending the job. After a denied call, the final `MicrovmEndpointUnavailableError` message adds `GetMicrovm was denied, so the operation could not check whether the MicroVM is suspended.` 5. A denied `ResumeMicrovm` for a suspending or suspended MicroVM fails the job at once with `MicrovmDeliveryError`. Its message names `lambda:ResumeMicrovm`. diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py index c1ddfe53f..11ef931b3 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py +++ b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/lifecycle.py @@ -113,10 +113,10 @@ def ensure_running(options: EnsureRunningOptions) -> str | None: 2. ``PENDING``: the function returns at once. A new MicroVM reports ``PENDING`` until it boots. The request's own retries wait for the ``run`` hook. - 3. ``SUSPENDED``: it calls ``ResumeMicrovm`` once, then waits for + 3. ``SUSPENDED``: it calls ``ResumeMicrovm``, then waits for ``RUNNING``. The MicroVM stays ``SUSPENDED`` while its ``resume`` hook - runs. So the function does not call ``ResumeMicrovm`` again for each - ``SUSPENDED`` answer. + runs. So after a successful call, the function does not call + ``ResumeMicrovm`` again for each ``SUSPENDED`` answer. 4. ``SUSPENDING``: it waits. The MicroVM becomes ``SUSPENDED``, and the function then resumes it. 5. ``TERMINATING``, ``TERMINATED``, or not found: it raises @@ -129,7 +129,9 @@ def ensure_running(options: EnsureRunningOptions) -> str | None: ``AccessDeniedException``, is raised with its class and its ``response``. Its message then ends with the request ID of the call. See :func:`add_request_id`. A ``ConflictException`` from ``ResumeMicrovm`` - is not an error: polling continues. + is not an error: polling continues. The call did not resume the + MicroVM. So the next ``SUSPENDED`` answer calls ``ResumeMicrovm`` again, + at most once per poll. It polls after 250 milliseconds, and doubles the wait up to 2 seconds. It stops after ``max_wait_seconds``, or 10 seconds before the invocation @@ -149,7 +151,7 @@ def ensure_running(options: EnsureRunningOptions) -> str | None: ClientError: When ``GetMicrovm`` or ``ResumeMicrovm`` fails otherwise. """ - return _wait_until_running(options, resume_sent=False) + return _wait_until_running(options, resuming=False, resume_accepted=False) def resume_if_suspended(options: EnsureRunningOptions) -> str | None: @@ -174,7 +176,9 @@ def resume_if_suspended(options: EnsureRunningOptions) -> str | None: :class:`MicrovmStateError`. A denied ``ResumeMicrovm`` raises :class:`MicrovmResumeDeniedError`. - ``ConflictException`` is not an error: polling continues. Any other + ``ConflictException`` is not an error: polling continues. A resume that + the service rejects with a conflict while the MicroVM is still + suspending is sent again once the MicroVM is suspended. Any other ``ResumeMicrovm`` error is raised with its class and its ``response``, as in :func:`ensure_running`. So is any error of a ``GetMicrovm`` call after the resume. @@ -229,6 +233,7 @@ def unavailable() -> Exception: reserve, unavailable, ) + accepted = True except ClientError as error: code = error_code(error) if code == "ResourceNotFoundException": @@ -249,10 +254,16 @@ def unavailable() -> Exception: ) from error if code != "ConflictException": raise add_request_id("ResumeMicrovm", error) + # A conflict can mean that the suspend has not ended. The MicroVM + # then becomes SUSPENDED and needs a new resume. So the loop sends + # ResumeMicrovm again at the next SUSPENDED answer. + accepted = False options.logger.info("resuming MicroVM", extra={"state": state}) remaining = options.max_wait_seconds - (time.monotonic() - started) return _wait_until_running( - replace(options, max_wait_seconds=remaining), resume_sent=True + replace(options, max_wait_seconds=remaining), + resuming=True, + resume_accepted=accepted, ) @@ -269,10 +280,16 @@ def _raise_if_ended(microvm_id: str, current: Any) -> None: def _wait_until_running( - options: EnsureRunningOptions, *, resume_sent: bool + options: EnsureRunningOptions, *, resuming: bool, resume_accepted: bool ) -> str | None: - """The loop of :func:`ensure_running`. ``resume_sent`` states that the - caller has already called ``ResumeMicrovm``.""" + """The loop of :func:`ensure_running`. + + ``resuming`` states that the caller has already called + ``ResumeMicrovm``. A ``PENDING`` answer is then a step of the resume, so + the loop waits for ``RUNNING``. ``resume_accepted`` states that a + ``ResumeMicrovm`` call succeeded. A call that answered + ``ConflictException`` leaves it ``False``. So the next ``SUSPENDED`` + answer sends ``ResumeMicrovm`` again.""" started = time.monotonic() reserve = reserve_starts_at(options.remaining_seconds) deadline = started + options.max_wait_seconds @@ -319,22 +336,25 @@ def resume_microvm() -> Any: raise add_request_id("GetMicrovm", error) state = current.get("state") last_state = state - if state == "RUNNING" or (state == "PENDING" and not resume_sent): - if resume_sent: + if state == "RUNNING" or (state == "PENDING" and not resuming): + if resuming: options.logger.info( "MicroVM resumed", extra={"waitedMs": round((time.monotonic() - started) * 1000)}, ) return current.get("endpoint") _raise_if_ended(microvm_id, current) - if state == "SUSPENDED" and not resume_sent: + if state == "SUSPENDED" and not resume_accepted: try: resume_microvm() + resume_accepted = True except ClientError as error: - # ConflictException means that a transition is already in - # progress, for example an auto-resume. Polling continues - # either way. ResourceNotFoundException means that the - # MicroVM was removed after the state check. + # ConflictException means that a transition is in progress, + # for example a suspend that has not ended, or an + # auto-resume. The resume stays unaccepted, so the next + # SUSPENDED answer sends it again. The poll interval and the + # deadline bound those calls. ResourceNotFoundException + # means that the MicroVM was removed after the state check. code = error_code(error) if code == "ResourceNotFoundException": raise microvm_not_found( @@ -342,8 +362,9 @@ def resume_microvm() -> Any: ) from error if code != "ConflictException": raise add_request_id("ResumeMicrovm", error) - resume_sent = True - options.logger.info("resuming MicroVM", extra={"state": state}) + if not resuming: + options.logger.info("resuming MicroVM", extra={"state": state}) + resuming = True if time.monotonic() + delay > deadline: raise unavailable() diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py b/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py index 0ce3f5522..5d9ec3699 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/conftest.py @@ -39,7 +39,9 @@ class FakeMicrovmsClient: ``states`` holds the answers of ``get_microvm``: a state such as ``"SUSPENDED"``, or an exception to raise. After the list ends, the state - is ``RUNNING``. + is ``RUNNING``. With ``suspended_until_resumed``, the state after the + list ends is ``SUSPENDED`` until a ``resume_microvm`` call succeeds, and + ``RUNNING`` after it. ``run_microvm`` and ``get_microvm`` return ``endpoint``. The Nth auth token is ``{"x-test-token": "token-N"}``. @@ -54,9 +56,11 @@ def __init__( resumes: list[Answer] | None = None, states: list[str | BaseException] | None = None, endpoint: str | None = "mvm-1.microvm.example", + suspended_until_resumed: bool = False, ) -> None: self.calls: list[tuple[str, dict[str, Any]]] = [] self.endpoint = endpoint + self._suspended = suspended_until_resumed self._answers: dict[str, list[Answer]] = { "run_microvm": list(launches or []), "create_microvm_auth_token": list(tokens or []), @@ -99,12 +103,17 @@ def terminate_microvm(self, **request: Any) -> dict[str, Any]: def resume_microvm(self, **request: Any) -> dict[str, Any]: self._answer("resume_microvm", request) + with self._lock: + self._suspended = False return {} def get_microvm(self, **request: Any) -> dict[str, Any]: with self._lock: self.calls.append(("get_microvm", request)) - state = self._states.pop(0) if self._states else "RUNNING" + if self._states: + state = self._states.pop(0) + else: + state = "SUSPENDED" if self._suspended else "RUNNING" if isinstance(state, BaseException): raise state response: dict[str, Any] = {"microvmId": "mvm-1", "state": state} diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py index a1d6e4007..e6d979580 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/lifecycle_test.py @@ -82,13 +82,71 @@ def test_waits_for_a_suspend_in_progress_then_resumes_once(fake_client: Any) -> def test_keeps_polling_when_resume_reports_a_transition_in_progress( fake_client: Any, make_error: Any ) -> None: + # An auto-resume is in progress. Each SUSPENDED answer sends + # ResumeMicrovm again, and the service answers ConflictException. A + # PENDING answer after a resume is a step of the resume, not a boot. client = fake_client( - states=["SUSPENDED", "SUSPENDED"], - resumes=[make_error("ConflictException", 409)], + states=["SUSPENDED", "SUSPENDED", "PENDING", "RUNNING"], + resumes=[make_error("ConflictException", 409)] * 2, ) assert run_with(client) == "mvm-1.microvm.example" - assert len(client.requests("resume_microvm")) == 1 + assert calls(client) == [ + "get_microvm", + "resume_microvm", + "get_microvm", + "resume_microvm", + "get_microvm", + "get_microvm", + ] + + +def test_resumes_again_when_a_resume_conflicts_with_the_suspend( + fake_client: Any, make_error: Any +) -> None: + # The service rejects the first ResumeMicrovm with ConflictException. + # The MicroVM stays SUSPENDED until a resume succeeds. So the next + # SUSPENDED answer sends ResumeMicrovm again. + client = fake_client( + states=["SUSPENDING", "SUSPENDED"], + resumes=[make_error("ConflictException", 409)], + suspended_until_resumed=True, + ) + sleeps: list[float] = [] + + assert run_with(client, max_wait_seconds=3, sleeps=sleeps) == ( + "mvm-1.microvm.example" + ) + assert calls(client) == [ + "get_microvm", + "get_microvm", + "resume_microvm", + "get_microvm", + "resume_microvm", + "get_microvm", + ] + assert sleeps == [0.25, 0.5, 1.0] + + +def test_gives_up_after_the_wait_limit_when_each_resume_conflicts( + fake_client: Any, make_error: Any +) -> None: + client = fake_client( + resumes=[make_error("ConflictException", 409)] * 10, + suspended_until_resumed=True, + ) + + with pytest.raises(MicrovmEndpointUnavailableError, match="Last state: SUSPENDED"): + run_with(client, max_wait_seconds=0.5) + + # One ResumeMicrovm call per GetMicrovm call. The wait of 0.5 seconds + # after the second GetMicrovm call does not fit in the limit. + assert calls(client) == [ + "get_microvm", + "resume_microvm", + "get_microvm", + "resume_microvm", + ] def test_raises_a_resume_error_that_is_not_a_conflict( @@ -303,6 +361,58 @@ def test_check_keeps_polling_after_a_resume_conflict( assert calls(client) == ["get_microvm", "resume_microvm", "get_microvm"] +def test_check_resumes_again_after_a_conflict_during_the_suspend( + fake_client: Any, make_error: Any +) -> None: + # The service rejects the resume during SUSPENDING with + # ConflictException. The suspend then ends, and the MicroVM stays + # SUSPENDED until a resume succeeds. So the first SUSPENDED answer sends + # ResumeMicrovm again. + client = fake_client( + states=["SUSPENDING", "SUSPENDING"], + resumes=[make_error("ConflictException", 409)], + suspended_until_resumed=True, + ) + sleeps: list[float] = [] + + assert check_with(client, max_wait_seconds=3, sleeps=sleeps) == ( + "mvm-1.microvm.example" + ) + assert calls(client) == [ + "get_microvm", + "resume_microvm", + "get_microvm", + "get_microvm", + "resume_microvm", + "get_microvm", + ] + assert sleeps == [0.25, 0.5] + + +def test_check_gives_up_after_the_wait_limit_when_each_resume_conflicts( + fake_client: Any, make_error: Any +) -> None: + client = fake_client( + states=["SUSPENDING"], + resumes=[make_error("ConflictException", 409)] * 10, + suspended_until_resumed=True, + ) + + with pytest.raises(MicrovmEndpointUnavailableError, match="Last state: SUSPENDED"): + check_with(client, max_wait_seconds=0.5) + + # One ResumeMicrovm call per GetMicrovm call. The wait of 0.5 seconds + # after the third GetMicrovm call does not fit in the limit. + assert calls(client) == [ + "get_microvm", + "resume_microvm", + "get_microvm", + "resume_microvm", + "get_microvm", + "resume_microvm", + ] + + @pytest.mark.parametrize("state", ["TERMINATING", "TERMINATED"]) def test_check_raises_the_state_error_for_a_microvm_that_ends( fake_client: Any, state: str diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py index 23b8262e0..fa5500eb2 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py @@ -715,6 +715,49 @@ def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: assert len(client.requests("resume_microvm")) == 1 +@pytest.mark.parametrize( + ("statuses", "states"), + [([], ["SUSPENDED"]), ([502], ["RUNNING", "SUSPENDED"])], + ids=["before-delivery", "after-a-gateway-error"], +) +def test_resumes_again_when_a_resume_conflicts( + run: Any, + fake_client: Any, + fake_endpoint: Any, + make_error: Any, + statuses: list[int], + states: list[str], +) -> None: + # The service rejects the first ResumeMicrovm with ConflictException. + # The MicroVM stays SUSPENDED until a resume succeeds. So the next + # SUSPENDED answer sends ResumeMicrovm again. + endpoint = fake_endpoint(statuses) + client = fake_client( + endpoint=endpoint.url, + states=states, + resumes=[make_error("ConflictException", 409)], + suspended_until_resumed=True, + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke( + ctx, + "build", + job_input=1, + options=job(retry_window=Duration.from_seconds(5)), + ) + + result = run( + session_handler(session_config(client, auto_suspend_on_idle=True), body), + None, + [("build", succeed_with("done"))], + ) + + assert outcome(result) == {"result": "done"} + assert len(client.requests("resume_microvm")) == 2 + assert len(endpoint.requests) == len(statuses) + 1 + + def test_does_not_check_the_state_again_after_a_rejection( run: Any, fake_client: Any, fake_endpoint: Any ) -> None: @@ -1041,16 +1084,17 @@ def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: assert len(client.requests("terminate_microvm")) == 1 -def test_resumes_once_when_two_items_find_the_microvm_suspended( +def test_runs_both_items_when_a_resume_finds_another_in_progress( run: Any, fake_client: Any, fake_endpoint: Any, make_error: Any ) -> None: # Both items find the MicroVM SUSPENDED and call ResumeMicrovm. The - # second call finds the resume in progress. + # service accepts the first call. It answers ConflictException to each + # later call while the resume runs. endpoint = fake_endpoint() client = fake_client( endpoint=endpoint.url, states=["SUSPENDED"] * 4, - resumes=[None, make_error("ConflictException", 409)], + resumes=[None] + [make_error("ConflictException", 409)] * 4, ) def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: @@ -1069,7 +1113,10 @@ def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: ) assert outcome(result) == {"result": ["A", "B"]} - assert len(client.requests("resume_microvm")) == 2 + # The item whose call conflicts calls ResumeMicrovm again at each later + # SUSPENDED answer. The order of the threads sets how many answers it + # gets. The other item calls once. + assert 2 <= len(client.requests("resume_microvm")) <= 4 assert len(endpoint.requests) == 2 @@ -1648,6 +1695,44 @@ def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: ] +def test_the_fallback_resumes_again_after_a_conflict_during_the_suspend( + run: Any, + fake_client: Any, + fake_endpoint: Any, + make_error: Any, + connect_failed: Any, +) -> None: + # The check finds the MicroVM SUSPENDING. The service rejects the resume + # with ConflictException. The suspend then ends, and the MicroVM stays + # SUSPENDED until a resume succeeds. So the first SUSPENDED answer sends + # ResumeMicrovm again. + endpoint = fake_endpoint([connect_failed]) + client = fake_client( + endpoint=endpoint.url, + states=["SUSPENDING"], + resumes=[make_error("ConflictException", 409)], + suspended_until_resumed=True, + ) + + def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any: + return vm.invoke( + ctx, + "build", + job_input=1, + options=job(retry_window=Duration.from_seconds(5)), + ) + + result = run( + session_handler(session_config(client), body), + None, + [("build", succeed_with("done"))], + ) + + assert outcome(result) == {"result": "done"} + assert len(client.requests("resume_microvm")) == 2 + assert len(endpoint.requests) == 2 + + def test_a_session_that_never_suspends_checks_the_state_after_no_answer( run: Any, fake_client: Any, fake_endpoint: Any ) -> None: From 73588b0620c7f8a7b1447ecb5dc22df7043dc062 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Mon, 26 Oct 2026 11:41:18 -0700 Subject: [PATCH 11/12] test(extras): forward only known proxy headers The simulated endpoint copied the X-Microvm-Suspend-After value from the worker's answer into its own answer. A value from an answer could then reach send_header. The endpoint now sends the header with the fixed value "accepted", and only when the worker answered that value. The session accepts only that value. So the tests still prove that an old worker gets a warning and a current worker does not. --- .../tests/microvm/worker_composed_test.py | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py index ca8b5b30e..dbca22516 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py @@ -613,11 +613,11 @@ def do_POST(self) -> None: # noqa: N802 - the http.server name self.send_header("content-length", "0") self.end_headers() return - status, headers, data = simulated.forward(self.path, body) + status, accepted, data = simulated.forward(self.path, body) simulated.record(status) self.send_response(status) - for name, value in headers.items(): - self.send_header(name, value) + if accepted: + self.send_header("X-Microvm-Suspend-After", "accepted") self.send_header("content-length", str(len(data))) self.end_headers() self.wfile.write(data) @@ -644,7 +644,9 @@ def set_state(self, state: str) -> None: with self._lock: self._state = state - def forward(self, path: str, body: bytes) -> tuple[int, dict[str, str], bytes]: + def forward(self, path: str, body: bytes) -> tuple[int, bool, bytes]: + """Send the job to the worker. Return the status, whether the worker + accepted ``suspendAfter``, and the body of the answer.""" worker = self.worker assert worker is not None if self._old_worker: @@ -660,12 +662,13 @@ def forward(self, path: str, body: bytes) -> tuple[int, dict[str, str], bytes]: ) response = connection.getresponse() data = response.read() - ack = response.getheader("X-Microvm-Suspend-After") - return ( - response.status, - ({} if ack is None else {"X-Microvm-Suspend-After": ack}), - data, - ) + # The endpoint sends one header with a fixed name and a fixed + # value. A value from the worker's answer never reaches + # send_header. The session accepts only this value. So an answer + # with another value reaches the session as an answer without + # the header, which the session treats the same way. + accepted = response.getheader("X-Microvm-Suspend-After") == "accepted" + return response.status, accepted, data finally: connection.close() From 57c2b3b79ccd3f57436ecd5e0b06b4abf2df5d42 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Mon, 26 Oct 2026 11:42:32 -0700 Subject: [PATCH 12/12] test(extras): check the ack of a current worker The simulated endpoint sends X-Microvm-Suspend-After only when the worker accepted suspendAfter. No test checked that a current worker's acknowledgement reaches the session. So a lost header would have passed. The suspend-after test now checks that the session logs no warning about an ignored suspendAfter. Without the header, the test fails. --- .../tests/microvm/worker_composed_test.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py index dbca22516..3fa7a0896 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py @@ -777,7 +777,7 @@ def echo(job_input: Any, context: MicrovmJobContext) -> Any: def test_suspends_after_a_job_and_the_next_job_resumes_the_microvm( - fake_client: Any, + fake_client: Any, caplog: pytest.LogCaptureFixture ) -> None: def make_pipeline(microvm: SimulatedMicrovm) -> Pipeline: def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: @@ -790,11 +790,14 @@ def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any: return pipeline - outcome, client, microvm, logger = suspend_after_session( - fake_client, make_pipeline, echo - ) + with caplog.at_level(logging.WARNING): + outcome, client, microvm, logger = suspend_after_session( + fake_client, make_pipeline, echo + ) assert outcome == ["build done", "SUSPENDED", "test done"] + # The worker acknowledged suspendAfter. So the session logs no warning. + assert IGNORED_WARNING not in [record.getMessage() for record in caplog.records] assert microvm.suspended == ["mvm-1"] [line] = logger.extras("suspending the MicroVM after a job") assert list(line) == ["microvmId", "callbackId"]