microvm-ctl 0.3.1__tar.gz → 0.5.0__tar.gz

{microvm_ctl-0.3.1 → microvm_ctl-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: microvm-ctl
-Version: 0.3.1
+Version: 0.5.0
 Summary: Control and execution plane for AWS Lambda MicroVMs: build images, launch Firecracker microVMs, scale fleets inside your quotas, call into them securely, and watch it all live.
 Author: Vivek Raja P S
 License: Apache-2.0
@@ -196,10 +196,12 @@ The four-part article series **Building on AWS Lambda MicroVMs** on the AWS Buil
 1. [Control and scale AWS Lambda MicroVMs with microvm-ctl](https://builder.aws.com/content/3JIDTpz0ZgatSBv24drra3gEod9/control-and-scale-aws-lambda-microvms-with-microvm-ctl): this plane and its measurements.
 2. [Seven workloads Lambda could never run, until MicroVMs](https://builder.aws.com/content/3JJ2oNWY9EsZzivMMx044cSlrFQ/seven-workloads-lambda-could-never-run-until-microvms): the first seven examples through build, run, cost, and gotchas.
 3. [A kernel for every customer: scaling AI agents to 1,000 tenants on AWS Lambda MicroVMs with microvm-ctl](https://builder.aws.com/content/3JJ7tASPSSUUTrcpWnWtxOuu8g3/a-kernel-for-every-customer-scaling-ai-agents-to-1000-tenants-on-aws-lambda-microvms-with-microvm-ctl): the multi-tenant finale with the decision guide.
-4. Hand a task to a MicroVM from anywhere (part 4, [source in awesome-microvm](https://github.com/Vivek0712/awesome-microvm/blob/main/blog/03-handoff.md), Builder Center link to follow): the lease contract, Step Functions, durable functions, your own controller, and leases at scale.
+4. [Hand a task to a MicroVM from anywhere: one lease, Step Functions, durable functions, or your own controller](https://builder.aws.com/content/3Jcx0VcEOq9QQBtqr59y4S8dVFC/hand-a-task-to-a-microvm-from-anywhere-one-lease-step-functions-durable-functions-or-your-own-controller): the lease contract, Step Functions, durable functions, your own controller, and leases at scale.
 
 The article sources and a long-form deep dive per example live under [awesome-microvm/blog](https://github.com/Vivek0712/awesome-microvm/tree/main/blog).
 
+Every lease scenario also runs live, on Step Functions and on a Lambda durable function, in [microvm-handoff-demo](https://github.com/Vivek0712/microvm-handoff-demo): nine scenarios on both orchestrators plus one lease from the CLI, with the execution histories, VM and orchestrator logs, benchmarks, and console screenshots checked in. Its first pass ran on 0.3.0 and found the two things that became 0.3.1: a typed failure whose VM Step Functions never terminated, and a 30 s heartbeat interval against a 30 s heartbeat timeout that lost a durable callback before the first heartbeat landed.
+
 ## Requirements and regions
 
 Python 3.9 or newer on the machine running the plane. The VMs themselves are ARM64 (Graviton) only, so audit binary wheels before you build an image. The service is available in `us-east-1`, `us-east-2`, `us-west-2`, `eu-west-1`, and `ap-northeast-1`.

{microvm_ctl-0.3.1 → microvm_ctl-0.5.0}/README.md

@@ -160,10 +160,12 @@ The four-part article series **Building on AWS Lambda MicroVMs** on the AWS Buil
 1. [Control and scale AWS Lambda MicroVMs with microvm-ctl](https://builder.aws.com/content/3JIDTpz0ZgatSBv24drra3gEod9/control-and-scale-aws-lambda-microvms-with-microvm-ctl): this plane and its measurements.
 2. [Seven workloads Lambda could never run, until MicroVMs](https://builder.aws.com/content/3JJ2oNWY9EsZzivMMx044cSlrFQ/seven-workloads-lambda-could-never-run-until-microvms): the first seven examples through build, run, cost, and gotchas.
 3. [A kernel for every customer: scaling AI agents to 1,000 tenants on AWS Lambda MicroVMs with microvm-ctl](https://builder.aws.com/content/3JJ7tASPSSUUTrcpWnWtxOuu8g3/a-kernel-for-every-customer-scaling-ai-agents-to-1000-tenants-on-aws-lambda-microvms-with-microvm-ctl): the multi-tenant finale with the decision guide.
-4. Hand a task to a MicroVM from anywhere (part 4, [source in awesome-microvm](https://github.com/Vivek0712/awesome-microvm/blob/main/blog/03-handoff.md), Builder Center link to follow): the lease contract, Step Functions, durable functions, your own controller, and leases at scale.
+4. [Hand a task to a MicroVM from anywhere: one lease, Step Functions, durable functions, or your own controller](https://builder.aws.com/content/3Jcx0VcEOq9QQBtqr59y4S8dVFC/hand-a-task-to-a-microvm-from-anywhere-one-lease-step-functions-durable-functions-or-your-own-controller): the lease contract, Step Functions, durable functions, your own controller, and leases at scale.
 
 The article sources and a long-form deep dive per example live under [awesome-microvm/blog](https://github.com/Vivek0712/awesome-microvm/tree/main/blog).
 
+Every lease scenario also runs live, on Step Functions and on a Lambda durable function, in [microvm-handoff-demo](https://github.com/Vivek0712/microvm-handoff-demo): nine scenarios on both orchestrators plus one lease from the CLI, with the execution histories, VM and orchestrator logs, benchmarks, and console screenshots checked in. Its first pass ran on 0.3.0 and found the two things that became 0.3.1: a typed failure whose VM Step Functions never terminated, and a 30 s heartbeat interval against a 30 s heartbeat timeout that lost a durable callback before the first heartbeat landed.
+
 ## Requirements and regions
 
 Python 3.9 or newer on the machine running the plane. The VMs themselves are ARM64 (Graviton) only, so audit binary wheels before you build an image. The service is available in `us-east-1`, `us-east-2`, `us-west-2`, `eu-west-1`, and `ap-northeast-1`.

{microvm_ctl-0.3.1 → microvm_ctl-0.5.0}/microvm/__init__.py

@@ -12,7 +12,7 @@ from microvm.images import ImageBuilder, ImageBuildError
 from microvm.lease import FanoutLimit, Lease, LeasePlan, LeasePlanRejected, LeasePolicy, plan_fanout
 from microvm.monitor import CostModel, FleetMonitor
 
-__version__ = "0.3.1"
+__version__ = "0.5.0"
 
 __all__ = [
     "microvm_client",

{microvm_ctl-0.3.1 → microvm_ctl-0.5.0}/microvm/cli.py

@@ -9,6 +9,7 @@
     mvm scale IMAGE N                      converge fleet to N
     mvm suspend|resume|terminate ID...     lifecycle control
     mvm drain IMAGE                        terminate the whole fleet
+    mvm dispatch IMAGE /path -d '{}' -n 64 spread requests over the warm fleet, --per-vm in flight each
     mvm call ID /path [-X POST -d '{}']    authenticated request into the VM
     mvm status ID | watch ID               job telemetry from the hook runtime (/status, /events)
     mvm watch --image NAME                 one live table for every leased VM of an image
@@ -226,6 +227,39 @@ def cmd_drain(args):
     console.print(f"[bold green]✓[/] terminated {n} microVMs of [cyan]{args.image}[/]")
 
 
+def cmd_dispatch(args):
+    if args.bodies:
+        with open(args.bodies) as f:
+            bodies = [json.loads(line) for line in f if line.strip()]
+    else:
+        bodies = [json.loads(args.data or "{}")] * args.n
+    fleet = Fleet(FleetManager(_cfg(args)), args.image)
+    started = time.time()
+    with console.status(f"[bold]dispatching {len(bodies)} requests over [cyan]{args.image}[/]…"):
+        results = fleet.dispatch(args.path, bodies, per_vm=args.per_vm, method=args.method, port=args.port,
+                                 timeout=args.timeout)
+    wall = time.time() - started
+    if args.out:
+        with open(args.out, "w") as f:
+            for r in results:
+                f.write(json.dumps(r, default=str) + "\n")
+    by_vm: dict = {}
+    for r in results:
+        by_vm.setdefault(r["microvm_id"], []).append(r)
+    rate = len(results) / wall
+    t = Table(title=f"{len(results)} requests in {wall:.2f} s ({rate:.1f}/s), {args.per_vm} in flight per VM")
+    for col in ("microVM", "requests", "ok", "failed", "p50 ms", "max ms"):
+        t.add_column(col)
+    for vid, rs in sorted(by_vm.items()):
+        lat = sorted(r["ms"] for r in rs)
+        ok = sum(1 for r in rs if r["status"] and r["status"] < 400)
+        p50 = lat[len(lat) // 2]
+        t.add_row(vid, str(len(rs)), str(ok), str(len(rs) - ok), f"{p50:.0f}", f"{lat[-1]:.0f}")
+    console.print(t)
+    if any(not r["status"] or r["status"] >= 400 for r in results):
+        sys.exit(1)
+
+
 # ---------------------------------------------------------------- execution plane
 def cmd_call(args):
     cfg = _cfg(args)
@@ -724,6 +758,19 @@ def main(argv: list[str] | None = None):
     d.add_argument("image")
     d.set_defaults(fn=cmd_drain)
 
+    dp = sub.add_parser("dispatch", help="spread requests over the RUNNING members of an image")
+    dp.add_argument("image")
+    dp.add_argument("path")
+    dp.add_argument("-X", "--method", default="POST")
+    dp.add_argument("-d", "--data", help="one JSON body, sent -n times")
+    dp.add_argument("-n", type=int, default=1, help="how many times to send --data")
+    dp.add_argument("--bodies", help="a JSON Lines file, one request body per line")
+    dp.add_argument("--per-vm", type=int, default=4, help="requests in flight per VM")
+    dp.add_argument("--port", type=int, default=8080)
+    dp.add_argument("--timeout", type=float, default=300)
+    dp.add_argument("--out", help="write every result as JSON Lines")
+    dp.set_defaults(fn=cmd_dispatch)
+
     c = sub.add_parser("call", help="authenticated request into a microVM")
     c.add_argument("id")
     c.add_argument("path")

{microvm_ctl-0.3.1 → microvm_ctl-0.5.0}/microvm/fleet.py

@@ -12,9 +12,11 @@ retried with jittered backoff, so `scale_to(50)` is safe to call in one shot.
 from __future__ import annotations
 
 import concurrent.futures as futures
+import queue
+import threading
 import time
 from dataclasses import dataclass, field
-from typing import Callable
+from typing import Any, Callable
 
 from microvm.client import image_arn, microvm_client
 from microvm.config import TPS, PlaneConfig
@@ -311,15 +313,23 @@ class Fleet:
     _pool: futures.ThreadPoolExecutor = field(
         default_factory=lambda: futures.ThreadPoolExecutor(max_workers=Fleet.MAX_WORKERS), repr=False
     )
+    #: VMs this fleet has asked to terminate. ListMicrovms keeps reporting a terminated VM in its old
+    #: state for about a second, so without this a size() right after a scale-down counts the victim
+    #: and a drain terminates it twice.
+    _terminated: set = field(default_factory=set, repr=False)
 
     # -- observation -------------------------------------------------------------
     def members(self) -> list[Microvm]:
         return [
             vm
             for vm in self.manager.list(self.image, self.version)
-            if vm.state in ACTIVE_STATES
+            if vm.state in ACTIVE_STATES and vm.microvm_id not in self._terminated
         ]
 
+    def _terminate(self, microvm_id: str) -> None:
+        self.manager.terminate(microvm_id)
+        self._terminated.add(microvm_id)
+
     def size(self) -> int:
         return len(self.members())
 
@@ -345,7 +355,7 @@ class Fleet:
             return launched
         if delta < 0:
             for vm in self.scale_down_victims(current, -delta):
-                self.manager.terminate(vm.microvm_id)
+                self._terminate(vm.microvm_id)
         return []
 
     @staticmethod
@@ -384,7 +394,7 @@ class Fleet:
     def drain(self) -> int:
         """Terminate every member of the fleet."""
         vms = self.members()
-        list(self._pool.map(lambda v: self.manager.terminate(v.microvm_id), vms))
+        list(self._pool.map(lambda v: self._terminate(v.microvm_id), vms))
         return len(vms)
 
     # -- reaper ------------------------------------------------------------------
@@ -396,6 +406,82 @@ class Fleet:
         for vm in self.members():
             started = vm.started_epoch
             if started and now - started > max_age_seconds:
-                self.manager.terminate(vm.microvm_id)
+                self._terminate(vm.microvm_id)
                 reaped.append(vm.microvm_id)
         return reaped
+
+    # -- work distribution -------------------------------------------------------
+    def dispatch(
+        self,
+        path: str,
+        bodies: list,
+        *,
+        per_vm: int = 4,
+        method: str = "POST",
+        port: int = 8080,
+        timeout: float = 300,
+        client_factory: Callable[[str], Any] | None = None,
+        on_result: Callable[[dict], None] | None = None,
+    ) -> list[dict]:
+        """Send one request per body to the fleet's RUNNING members and return the results in order.
+
+        Each RUNNING member gets `per_vm` workers that pull from one shared queue, so a slow VM takes
+        fewer items and no VM ever holds more than `per_vm` requests in flight. This is the warm-fleet
+        path: hundreds of short tasks over VMs that are already up, where a lease per task would pay a
+        launch per task. A request that raises is recorded with its error and does not stop the others.
+
+        Each result is {"index", "microvm_id", "status", "ms", "body"} or, on an exception,
+        {"index", "microvm_id", "status": None, "ms", "error"}. `on_result` is called with each result
+        as it lands, from a worker thread.
+        """
+        if per_vm < 1:
+            raise ValueError("per_vm must be at least 1")
+        running = [vm for vm in self.members() if vm.state == "RUNNING"]
+        if not running:
+            raise RuntimeError(f"fleet {self.image} has no RUNNING members to dispatch to")
+        if client_factory is None:
+            from microvm.endpoint import EndpointClient
+
+            def client_factory(microvm_id: str):
+                return EndpointClient(self.manager.cfg, microvm_id, ports=[port])
+
+        work: queue.Queue = queue.Queue()
+        for i, body in enumerate(bodies):
+            work.put((i, body))
+        results: list = [None] * len(bodies)
+        lock = threading.Lock()
+
+        def worker(microvm_id: str, client) -> None:
+            while True:
+                try:
+                    i, body = work.get_nowait()
+                except queue.Empty:
+                    return
+                started = time.time()
+                try:
+                    resp = client.request(method, path, json=body, timeout=timeout, port=port)
+                    try:
+                        payload = resp.json()
+                    except ValueError:
+                        payload = resp.text
+                    out = {"index": i, "microvm_id": microvm_id, "status": resp.status_code,
+                           "ms": round((time.time() - started) * 1000, 1), "body": payload}
+                except Exception as e:  # one failed request must not stop the batch
+                    out = {"index": i, "microvm_id": microvm_id, "status": None,
+                           "ms": round((time.time() - started) * 1000, 1),
+                           "error": f"{type(e).__name__}: {e}"}
+                with lock:
+                    results[i] = out
+                if on_result:
+                    on_result(out)
+
+        threads = []
+        for vm in running:
+            client = client_factory(vm.microvm_id)
+            for _ in range(per_vm):
+                t = threading.Thread(target=worker, args=(vm.microvm_id, client), daemon=True)
+                t.start()
+                threads.append(t)
+        for t in threads:
+            t.join()
+        return results

{microvm_ctl-0.3.1 → microvm_ctl-0.5.0}/microvm/hooks/server.py

@@ -52,6 +52,12 @@ import traceback
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 
 HOOK_BASE = "/aws/lambda-microvms/runtime/v1"
+#: The job path of the AWS durable execution SDK's microvm operation (@aws/durable-execution-sdk-js-extras).
+AWS_JOB_PATH = "/durable-execution/v1/job"
+#: Error names from the durable callback APIs that mean no later attempt can succeed.
+AWS_TERMINAL_CALLBACK_ERRORS = ("CallbackTimeoutException", "InvalidParameterValueException",
+                                "ResourceNotFoundException")
+AWS_MAX_RESULT_BYTES = 256 * 1024
 LEASE_KINDS = ("sfn", "durable", "http", "sqs", "eventbridge", "none")
 RESULT_LIMIT = 240 * 1024
 SUMMARY_LIMIT = 4096
@@ -276,7 +282,8 @@ def _closed_if(call, names: tuple, **kw):
     try:
         return call(**kw)
     except Exception as e:
-        if type(e).__name__ in names:
+        # botocore raises a modeled class named after the code, or a plain ClientError carrying it
+        if type(e).__name__ in names or _error_code(e) in names:
             raise _TokenClosed(str(e)) from e
         raise
 
@@ -301,7 +308,10 @@ class _SfnCompleter(_Completer):
 
 
 class _DurableCompleter(_Completer):
-    CLOSED = ("CallbackTimeoutException", "InvalidToken")
+    # measured by the durable SDK team: a closed callback answers CallbackTimeoutException; an unknown or
+    # malformed id answers InvalidParameterValueException or ResourceNotFoundException
+    CLOSED = ("CallbackTimeoutException", "InvalidParameterValueException", "ResourceNotFoundException",
+              "InvalidToken")
 
     def __init__(self, lease: dict, factory):
         super().__init__(lease, factory)
@@ -374,9 +384,29 @@ COMPLETERS = {"none": _NoneCompleter, "sfn": _SfnCompleter, "durable": _DurableC
               "http": _HttpCompleter, "sqs": _SqsCompleter, "eventbridge": _EventBridgeCompleter}
 
 
+def _error_code(e) -> str | None:
+    resp = getattr(e, "response", None)
+    if isinstance(resp, dict):
+        return (resp.get("Error") or {}).get("Code")
+    return getattr(e, "code", None) or type(e).__name__
+
+
+class AwsJobContext:
+    """One job of the AWS durable execution SDK contract, as a durable_job handler sees it."""
+
+    def __init__(self, callback_id: str, input, microvm_id: str | None, job: Job, path: str):
+        self.callback_id, self.input, self.path = callback_id, input, path
+        self.microvm_id, self.job = microvm_id, job
+        self.cancelled = threading.Event()  # the durable function stopped waiting, or the microVM terminates
+        self.closed = False
+        self.heartbeats = 0
+
+
 class HookApp:
     #: (service, region) -> boto3-like client; tests replace this with a fake factory.
     lease_client_factory = staticmethod(_default_client_factory)
+    #: multiplier on the completion retry backoff (1, 2, 4, 8, 16 s); tests set it to 0
+    completion_backoff = 1.0
 
     def __init__(self):
         self._hooks: dict[str, object] = {}
@@ -386,6 +416,16 @@ class HookApp:
         self.run_payload: str | None = None
         self.job = Job(self.microvm_id)
         self.lease: LeaseContext | None = None
+        # AWS durable execution SDK job contract (see durable_job)
+        self._aws_handlers: dict[str, object] = {}
+        self._aws_seen: collections.OrderedDict = collections.OrderedDict()  # callback id -> "running"/"done"
+        self._aws_running: dict[str, AwsJobContext] = {}
+        self._aws_lock = threading.Lock()
+        self._aws_region: str | None = None
+        self._aws_refuse_until = 0.0
+        self._aws_terminated = False
+        self._aws_idle_s: float | None = None
+        self._aws_idle_timer: threading.Timer | None = None
 
         @self.route("GET", "/healthz")
         def _healthz(body, headers):
@@ -432,6 +472,24 @@ class HookApp:
         self._hooks["lease"] = fn
         return fn
 
+    def durable_job(self, fn=None, *, path: str | None = None):
+        """Serve jobs from the AWS durable execution SDK's ``microvm`` and ``microvmSession`` operations.
+
+        ``@app.durable_job`` handles a job that names no route: one in the run hook payload, or one
+        POSTed to /durable-execution/v1/job. ``@app.durable_job(path="/build")`` handles jobs sent with
+        ``request: {path: "/build"}``. ``fn(input, ctx)`` returns the JSON-serializable result, which
+        completes the durable callback; an exception fails it with the exception's ``error_type`` (a
+        LeaseError) or class name. ``ctx`` is an AwsJobContext: ``.cancelled`` is set when the durable
+        function stops waiting or the microVM terminates, and ``.job`` is the job telemetry that
+        GET /status and /events serve.
+        """
+        def deco(f):
+            key = path or AWS_JOB_PATH
+            self._aws_handlers[key] = f
+            self._routes[("POST", key)] = lambda body, headers: self._aws_http_job(key, body)
+            return f
+        return deco(fn) if fn is not None else deco
+
     def route(self, method: str, path: str):
         def deco(fn):
             self._routes[(method.upper(), path)] = fn
@@ -439,7 +497,27 @@ class HookApp:
         return deco
 
     # -- hook dispatch -----------------------------------------------------------
-    def _dispatch_hook(self, name: str, ctx: dict) -> int:
+    @staticmethod
+    def _from_lambda(headers: dict | None) -> bool:
+        """Lambda's own lifecycle calls carry a localhost Host and no x-amzn-requestid. The endpoint
+        forwards a hook path from anyone holding an auth token for the microVM, with the endpoint's
+        host name and an x-amzn-requestid header, so suspend, resume and terminate act only on the first."""
+        if headers is None:
+            return True
+        h = {k.lower(): v for k, v in headers.items()}
+        host = (h.get("host") or "").rsplit(":", 1)[0] if not (h.get("host") or "").startswith("[") \
+            else (h.get("host") or "").split("]")[0] + "]"
+        return host in ("localhost", "127.0.0.1", "[::1]") and "x-amzn-requestid" not in h
+
+    def _dispatch_hook(self, name: str, ctx: dict, headers: dict | None = None) -> int:
+        if name in ("suspend", "resume", "terminate") and not self._from_lambda(headers):
+            self.job.log(f"ignored a forwarded {name} hook call", level="warning")
+            return 200
+        if name == "suspend":
+            self._aws_refuse(30)
+        if name == "resume":
+            self._aws_refuse_until = 0.0
+            self._aws_arm_idle()
         if name == "ready":
             fn = self._hooks.get("ready")
             ok = fn(ctx) if fn else True
@@ -462,11 +540,13 @@ class HookApp:
             random.seed()
         if name == "terminate":
             self._terminate_lease()
+            self._aws_terminate()
         fn = self._hooks.get(name)
         if fn:
             fn(ctx)
         if name == "run":
-            self._maybe_start_lease(self.run_payload)
+            if not self._aws_run_hook(self.run_payload):
+                self._maybe_start_lease(self.run_payload)
         return 200
 
     # -- lease worker ------------------------------------------------------------
@@ -541,14 +621,19 @@ class HookApp:
         if ctx.lost:
             self.job.log(f"lease {outcome} not delivered: orchestrator stopped waiting", level="warning")
             return
-        try:
-            completer.send(outcome, payload)
-            self.job.log(f"lease {outcome} delivered", kind=ctx.lease.get("kind"))
-        except _TokenClosed as e:
-            ctx.lost = True
-            self.job.log(f"lease {outcome} not delivered: {e}", level="warning")
-        except Exception as e:
-            self.job.log(f"lease {outcome} delivery failed: {type(e).__name__}: {e}", level="error")
+        for attempt in range(5):  # a lost completion leaves the orchestrator waiting until its timeout
+            try:
+                completer.send(outcome, payload)
+                self.job.log(f"lease {outcome} delivered", kind=ctx.lease.get("kind"), attempt=attempt + 1)
+                return
+            except _TokenClosed as e:
+                ctx.lost = True
+                self.job.log(f"lease {outcome} not delivered: {e}", level="warning")
+                return
+            except Exception as e:
+                self.job.log(f"lease {outcome} delivery failed (attempt {attempt + 1}): "
+                             f"{type(e).__name__}: {e}", level="error")
+                time.sleep(min(2 ** attempt, 16) * type(self).completion_backoff)
 
     def _terminate_lease(self) -> None:
         ctx = self.lease
@@ -560,6 +645,185 @@ class HookApp:
         except Exception:
             pass
 
+    # -- AWS durable execution SDK job contract --------------------------------------
+    def _aws_run_hook(self, raw) -> bool:
+        """True when the run hook payload is the SDK's (version 1); starts its job when it carries one."""
+        try:
+            obj = json.loads(raw) if isinstance(raw, str) else raw
+        except (TypeError, ValueError):
+            return False
+        if not isinstance(obj, dict) or obj.get("version") != 1 or "lease" in obj:
+            return False
+        self._aws_region = obj.get("region") or self._aws_region
+        idle = obj.get("autoSuspendIdleSeconds")
+        self._aws_idle_s = float(idle) if isinstance(idle, (int, float)) and idle > 0 else None
+        job = obj.get("job")
+        if isinstance(job, dict) and self._aws_handlers:
+            self._aws_start(AWS_JOB_PATH, job)
+        else:
+            self._aws_arm_idle()
+        return True
+
+    def _aws_http_job(self, path: str, body: dict):
+        cb = body.get("callbackId") if isinstance(body, dict) else None
+        valid = isinstance(cb, str) and body.get("version") == 1
+        if not valid:
+            if isinstance(body, dict) and isinstance(body.get("callbackId"), str):
+                self._aws_report_async(body.get("region"), body["callbackId"], "failure",
+                                       {"ErrorType": "InvalidJobRequest",
+                                        "ErrorMessage": "malformed job request"})
+            return 400, {"error": "not a durable execution job request"}
+        if self._aws_terminated or time.time() < self._aws_refuse_until:
+            return 503, {"error": "suspending or terminating; send the job again after the resume"}
+        if body.get("microvmId") and not self.microvm_id:
+            self.microvm_id = self.job.microvm_id = body["microvmId"]
+        self._aws_region = self._aws_region or body.get("region")
+        self._aws_start(path, body)
+        return 202, {"accepted": body["callbackId"]}
+
+    def _aws_start(self, path: str, job: dict) -> None:
+        cb = job.get("callbackId")
+        if not isinstance(cb, str):
+            return
+        with self._aws_lock:
+            if cb in self._aws_seen:  # a second delivery of the same job
+                return
+            self._aws_seen[cb] = "running"
+            while len(self._aws_seen) > 1000:
+                self._aws_seen.popitem(last=False)
+            ctx = AwsJobContext(cb, job.get("input"), self.microvm_id, self.job, path)
+            self._aws_running[cb] = ctx
+            if self._aws_idle_timer:
+                self._aws_idle_timer.cancel()
+        handler = self._aws_handlers.get(path) or self._aws_handlers.get(AWS_JOB_PATH)
+        region = job.get("region") or self._aws_region or os.environ.get("AWS_REGION", "us-east-1")
+        hb = job.get("heartbeatTimeoutSeconds")
+        threading.Thread(target=self._aws_run, args=(handler, ctx, region, hb), daemon=True).start()
+
+    def _aws_run(self, handler, ctx: AwsJobContext, region: str, hb_timeout) -> None:
+        client = type(self).lease_client_factory("lambda", region)  # per job: none frozen in the snapshot
+        if isinstance(hb_timeout, (int, float)) and hb_timeout > 0:
+            threading.Thread(target=self._aws_heartbeats, args=(client, ctx, float(hb_timeout)),
+                             daemon=True).start()
+        self.job.phase(f"job {ctx.path}")
+        self.job.log("durable job started", callback=ctx.callback_id[-12:], path=ctx.path)
+        try:
+            if handler is None:
+                raise LeaseError("NoHandler", f"no durable_job handler for {ctx.path}")
+            result = handler(ctx.input, ctx)
+            body = json.dumps(result, default=str)
+            if len(body.encode()) > AWS_MAX_RESULT_BYTES:
+                raise LeaseError("ResultTooLargeError",
+                                 f"the result is {len(body.encode())} bytes; the limit is 256 KB")
+            outcome, kw = "success", {"Result": body.encode()}
+        except Exception as e:
+            data = getattr(e, "data", None)
+            outcome, kw = "failure", {"Error": {
+                "ErrorType": str(getattr(e, "error_type", type(e).__name__))[:256],
+                "ErrorMessage": str(getattr(e, "message", None) or e)[:8192],
+                **({"ErrorData": json.dumps(data, default=str)[:8192]} if data else {})}}
+        ctx.cancelled.set()
+        if not ctx.closed:
+            self._aws_report(client, ctx.callback_id, outcome, kw)
+        self._aws_finish(ctx, outcome)
+
+    def _aws_finish(self, ctx: AwsJobContext, outcome: str) -> None:
+        with self._aws_lock:
+            self._aws_running.pop(ctx.callback_id, None)
+            self._aws_seen[ctx.callback_id] = "done"
+            idle = not self._aws_running
+        self.job.log(f"durable job {outcome}", callback=ctx.callback_id[-12:])
+        self.job.counter("durable_jobs_" + outcome, self.job.counters.get("durable_jobs_" + outcome, 0) + 1)
+        if idle:
+            self.job.phase("idle")
+            self._aws_arm_idle()
+
+    def _aws_heartbeats(self, client, ctx: AwsJobContext, timeout_s: float) -> None:
+        # a third of the timeout, shortened by 1 to 2 s from a hash of the callback id: clones of one
+        # snapshot share Math.random-style state, so jitter must come from the job itself
+        jitter = 1 + (int.from_bytes(ctx.callback_id.encode()[-4:], "big") % 1000) / 1000
+        interval = max(1.0, timeout_s / 3 - jitter)
+        first = True
+        while first or not ctx.cancelled.wait(interval):
+            first = False
+            try:
+                client.send_durable_execution_callback_heartbeat(CallbackId=ctx.callback_id)
+                ctx.heartbeats += 1
+            except Exception as e:
+                if _error_code(e) in AWS_TERMINAL_CALLBACK_ERRORS:
+                    ctx.closed = True
+                    ctx.cancelled.set()
+                    self.job.log("durable function stopped waiting; cancelling the job", level="warning")
+                    return
+                self.job.log(f"heartbeat failed: {e}", level="warning")
+
+    def _aws_report(self, client, callback_id: str, outcome: str, kw: dict) -> None:
+        call = (client.send_durable_execution_callback_success if outcome == "success"
+                else client.send_durable_execution_callback_failure)
+        for attempt in range(5):
+            try:
+                call(CallbackId=callback_id, **kw)
+                return
+            except Exception as e:
+                if _error_code(e) in AWS_TERMINAL_CALLBACK_ERRORS:
+                    self.job.log(f"callback already closed: {_error_code(e)}", level="warning")
+                    return
+                self.job.log(f"{outcome} report failed (attempt {attempt + 1}): {e}", level="warning")
+                time.sleep(min(2 ** attempt, 16))
+
+    def _aws_report_async(self, region, callback_id: str, outcome: str, error: dict) -> None:
+        region = region or self._aws_region or os.environ.get("AWS_REGION", "us-east-1")
+        client = type(self).lease_client_factory("lambda", region)
+        threading.Thread(target=self._aws_report, args=(client, callback_id, outcome, {"Error": error}),
+                         daemon=True).start()
+
+    def _aws_terminate(self) -> None:
+        self._aws_terminated = True
+        with self._aws_lock:
+            running = list(self._aws_running.values())
+        for ctx in running:
+            ctx.closed = True
+            ctx.cancelled.set()
+            region = self._aws_region or os.environ.get("AWS_REGION", "us-east-1")
+            try:
+                client = type(self).lease_client_factory("lambda", region)
+                self._aws_report(client, ctx.callback_id, "failure",
+                                 {"Error": {"ErrorType": "MicrovmTerminatedError",
+                                            "ErrorMessage": "the microVM terminated while the job ran"}})
+            except Exception:
+                pass
+
+    def _aws_refuse(self, seconds: float) -> None:
+        self._aws_refuse_until = time.time() + seconds
+        if self._aws_idle_timer:
+            self._aws_idle_timer.cancel()
+
+    def _aws_arm_idle(self) -> None:
+        """With autoSuspendIdleSeconds, suspend this microVM after that long without a running job."""
+        if not self._aws_idle_s or self._aws_terminated:
+            return
+        if self._aws_idle_timer:
+            self._aws_idle_timer.cancel()
+        self._aws_idle_timer = threading.Timer(self._aws_idle_s, self._aws_self_suspend)
+        self._aws_idle_timer.daemon = True
+        self._aws_idle_timer.start()
+
+    def _aws_self_suspend(self) -> None:
+        with self._aws_lock:
+            if self._aws_running or not self.microvm_id:
+                return
+        self._aws_refuse(30)
+        try:
+            region = self._aws_region or os.environ.get("AWS_REGION", "us-east-1")
+            type(self).lease_client_factory("lambda-microvms", region).suspend_microvm(
+                microvmIdentifier=self.microvm_id)
+            self.job.log("suspended itself after the idle time", idle_s=self._aws_idle_s)
+        except Exception as e:
+            code = _error_code(e)
+            if code in ("AccessDeniedException", "ValidationException", "ResourceNotFoundException"):
+                self._aws_refuse_until = 0.0  # rejected outright: keep running and accept jobs
+            self.job.log(f"self-suspend failed: {code or e}", level="warning")
+
     # -- server ------------------------------------------------------------------
     def serve(self, port: int = 8080, background: bool = False):
         app = self
@@ -648,7 +912,7 @@ class HookApp:
                 try:
                     if path.startswith(HOOK_BASE + "/"):
                         hook = path[len(HOOK_BASE) + 1:]
-                        status = app._dispatch_hook(hook, self._read_json())
+                        status = app._dispatch_hook(hook, self._read_json(), dict(self.headers))
                         return self._send(status, {"hook": hook, "status": status})
                     fn = app._routes.get((method, path))
                     if fn is None:

{microvm_ctl-0.3.1 → microvm_ctl-0.5.0}/microvm_ctl.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: microvm-ctl
-Version: 0.3.1
+Version: 0.5.0
 Summary: Control and execution plane for AWS Lambda MicroVMs: build images, launch Firecracker microVMs, scale fleets inside your quotas, call into them securely, and watch it all live.
 Author: Vivek Raja P S
 License: Apache-2.0
@@ -196,10 +196,12 @@ The four-part article series **Building on AWS Lambda MicroVMs** on the AWS Buil
 1. [Control and scale AWS Lambda MicroVMs with microvm-ctl](https://builder.aws.com/content/3JIDTpz0ZgatSBv24drra3gEod9/control-and-scale-aws-lambda-microvms-with-microvm-ctl): this plane and its measurements.
 2. [Seven workloads Lambda could never run, until MicroVMs](https://builder.aws.com/content/3JJ2oNWY9EsZzivMMx044cSlrFQ/seven-workloads-lambda-could-never-run-until-microvms): the first seven examples through build, run, cost, and gotchas.
 3. [A kernel for every customer: scaling AI agents to 1,000 tenants on AWS Lambda MicroVMs with microvm-ctl](https://builder.aws.com/content/3JJ7tASPSSUUTrcpWnWtxOuu8g3/a-kernel-for-every-customer-scaling-ai-agents-to-1000-tenants-on-aws-lambda-microvms-with-microvm-ctl): the multi-tenant finale with the decision guide.
-4. Hand a task to a MicroVM from anywhere (part 4, [source in awesome-microvm](https://github.com/Vivek0712/awesome-microvm/blob/main/blog/03-handoff.md), Builder Center link to follow): the lease contract, Step Functions, durable functions, your own controller, and leases at scale.
+4. [Hand a task to a MicroVM from anywhere: one lease, Step Functions, durable functions, or your own controller](https://builder.aws.com/content/3Jcx0VcEOq9QQBtqr59y4S8dVFC/hand-a-task-to-a-microvm-from-anywhere-one-lease-step-functions-durable-functions-or-your-own-controller): the lease contract, Step Functions, durable functions, your own controller, and leases at scale.
 
 The article sources and a long-form deep dive per example live under [awesome-microvm/blog](https://github.com/Vivek0712/awesome-microvm/tree/main/blog).
 
+Every lease scenario also runs live, on Step Functions and on a Lambda durable function, in [microvm-handoff-demo](https://github.com/Vivek0712/microvm-handoff-demo): nine scenarios on both orchestrators plus one lease from the CLI, with the execution histories, VM and orchestrator logs, benchmarks, and console screenshots checked in. Its first pass ran on 0.3.0 and found the two things that became 0.3.1: a typed failure whose VM Step Functions never terminated, and a 30 s heartbeat interval against a 30 s heartbeat timeout that lost a durable callback before the first heartbeat landed.
+
 ## Requirements and regions
 
 Python 3.9 or newer on the machine running the plane. The VMs themselves are ARM64 (Graviton) only, so audit binary wheels before you build an image. The service is available in `us-east-1`, `us-east-2`, `us-west-2`, `eu-west-1`, and `ap-northeast-1`.

{microvm_ctl-0.3.1 → microvm_ctl-0.5.0}/microvm_ctl.egg-info/SOURCES.txt

@@ -35,6 +35,7 @@ tests/test_cost_model.py
 tests/test_endpoint.py
 tests/test_fleet.py
 tests/test_hooks.py
+tests/test_hooks_aws_contract.py
 tests/test_hooks_lease.py
 tests/test_hooks_validate.py
 tests/test_integrations.py

{microvm_ctl-0.3.1 → microvm_ctl-0.5.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "microvm-ctl"
-version = "0.3.1"
+version = "0.5.0"
 description = "Control and execution plane for AWS Lambda MicroVMs: build images, launch Firecracker microVMs, scale fleets inside your quotas, call into them securely, and watch it all live."
 readme = "README.md"
 requires-python = ">=3.9"

{microvm_ctl-0.3.1 → microvm_ctl-0.5.0}/tests/test_fleet.py

@@ -84,3 +84,67 @@ def test_run_params_carries_client_token(monkeypatch):
     assert params["runHookPayload"] == '{"callback_id": "abc"}'
     assert params["maximumDurationInSeconds"] == 900
     assert "clientToken" not in fm.run_params("img")
+
+
+def test_scale_down_victim_is_not_counted_while_list_still_shows_it():
+    # ListMicrovms reports a terminated VM in its old state for a moment; the fleet must not count it.
+    members = [_vm(1, "RUNNING", 600), _vm(2, "SUSPENDED", 300), _vm(3, "RUNNING", 10)]
+    fm = FakeManager(members)
+    fleet = Fleet(fm, "img")
+    fleet.scale_to(2)
+    assert fm.terminated == ["vm-2"]
+    assert fleet.size() == 2
+    assert fleet.drain() == 2
+    assert sorted(fm.terminated) == ["vm-1", "vm-2", "vm-3"]
+
+
+class FakeResponse:
+    def __init__(self, status, payload):
+        self.status_code, self._payload = status, payload
+
+    def json(self):
+        return self._payload
+
+
+class FakeClient:
+    def __init__(self, microvm_id, log, fail_on=None, delay=0.0):
+        self.microvm_id, self.log, self.fail_on, self.delay = microvm_id, log, fail_on, delay
+
+    def request(self, method, path, json=None, timeout=None, port=None):
+        import time as _t
+        _t.sleep(self.delay)
+        self.log.append((self.microvm_id, method, path, json))
+        if self.fail_on is not None and json == self.fail_on:
+            raise ConnectionError("boom")
+        return FakeResponse(200, {"echo": json, "vm": self.microvm_id})
+
+
+def test_dispatch_spreads_bodies_over_running_members_and_keeps_order():
+    fm = FakeManager([_vm(1, "RUNNING", 60), _vm(2, "RUNNING", 60), _vm(3, "SUSPENDED", 60)])
+    log = []
+    fleet = Fleet(fm, "img")
+    seen = []
+    results = fleet.dispatch("/task", [{"n": i} for i in range(10)], per_vm=2,
+                             client_factory=lambda vm_id: FakeClient(vm_id, log, delay=0.01),
+                             on_result=seen.append)
+    assert [r["index"] for r in results] == list(range(10))
+    assert [r["body"]["echo"] for r in results] == [{"n": i} for i in range(10)]
+    assert {r["microvm_id"] for r in results} == {"vm-1", "vm-2"}  # the SUSPENDED member gets nothing
+    assert all(m == "POST" and p == "/task" for _, m, p, _ in log)
+    assert len(seen) == 10
+
+
+def test_dispatch_records_a_failed_request_and_finishes_the_rest():
+    fm = FakeManager([_vm(1, "RUNNING", 60)])
+    fleet = Fleet(fm, "img")
+    results = fleet.dispatch("/task", [{"n": 0}, {"n": 1}, {"n": 2}], per_vm=1,
+                             client_factory=lambda vm_id: FakeClient(vm_id, [], fail_on={"n": 1}))
+    assert results[0]["status"] == 200 and results[2]["status"] == 200
+    assert results[1]["status"] is None and "ConnectionError" in results[1]["error"]
+
+
+def test_dispatch_needs_a_running_member():
+    import pytest
+    fleet = Fleet(FakeManager([_vm(1, "SUSPENDED", 60)]), "img")
+    with pytest.raises(RuntimeError, match="no RUNNING members"):
+        fleet.dispatch("/task", [{}], client_factory=lambda vm_id: None)

microvm_ctl-0.5.0/tests/test_hooks_aws_contract.py

@@ -0,0 +1,234 @@
+"""The AWS durable execution SDK job contract in the hook server (loopback only, no AWS).
+
+The contract is the one @aws/durable-execution-sdk-js-extras's `microvm` and `microvmSession` operations
+speak: a version 1 run hook payload with an optional job, jobs POSTed over HTTP, durable callback
+heartbeats and completion, and the suspend, resume and terminate lifecycle hooks.
+"""
+
+import json
+import threading
+import time
+import urllib.error
+import urllib.request
+
+import pytest
+
+from microvm.hooks.server import AWS_JOB_PATH, HookApp, LeaseError
+
+HOOK = "/aws/lambda-microvms/runtime/v1/"
+
+
+class ClientError(Exception):
+    def __init__(self, code):
+        super().__init__(code)
+        self.response = {"Error": {"Code": code}}
+
+
+class FakeLambda:
+    def __init__(self):
+        self.calls = []
+        self.fail = {}  # method name -> error code to raise
+        self.cv = threading.Condition()
+
+    def _call(self, name, **kw):
+        with self.cv:
+            self.calls.append((name, kw))
+            self.cv.notify_all()
+        if name in self.fail:
+            raise ClientError(self.fail[name])
+        return {}
+
+    def __getattr__(self, name):
+        return lambda **kw: self._call(name, **kw)
+
+    def wait(self, name, n=1, timeout=5.0):
+        deadline = time.time() + timeout
+        with self.cv:
+            while len([c for c in self.calls if c[0] == name]) < n:
+                if time.time() > deadline:
+                    raise AssertionError(f"no {n}x {name}: {self.calls}")
+                self.cv.wait(0.05)
+        return [kw for c, kw in self.calls if c == name]
+
+
+@pytest.fixture
+def env():
+    fake = FakeLambda()
+
+    class App(HookApp):
+        lease_client_factory = staticmethod(lambda service, region: fake)
+
+    app = App()
+    server = app.serve(port=0, background=True)
+    port = server.server_address[1]
+
+    def post(path, body, host=f"127.0.0.1:{port}", extra=None):
+        req = urllib.request.Request(f"http://127.0.0.1:{port}{path}", data=json.dumps(body).encode(),
+                                     method="POST", headers={"Content-Type": "application/json", "Host": host,
+                                                             **(extra or {})})
+        try:
+            with urllib.request.urlopen(req, timeout=5) as r:
+                return r.status, json.loads(r.read() or b"{}")
+        except urllib.error.HTTPError as e:
+            return e.code, json.loads(e.read() or b"{}")
+
+    yield app, fake, post
+    server.shutdown()
+
+
+def run_hook(post, payload):
+    return post(HOOK + "run", {"microvmId": "mvm-1", "runHookPayload": json.dumps(payload)})
+
+
+def test_run_hook_job_succeeds_with_the_raw_result_and_heartbeats(env):
+    app, fake, post = env
+
+    @app.durable_job
+    def handler(inp, ctx):
+        time.sleep(0.2)
+        return {"echo": inp["text"], "vm": ctx.microvm_id}
+
+    job = {"callbackId": "cb-1", "heartbeatTimeoutSeconds": 6, "input": {"text": "hi"}}
+    assert run_hook(post, {"version": 1, "region": "us-east-1", "job": job})[0] == 200
+    ok = fake.wait("send_durable_execution_callback_success")[0]
+    assert ok["CallbackId"] == "cb-1"
+    assert json.loads(ok["Result"]) == {"echo": "hi", "vm": "mvm-1"}  # the handler's value, no envelope
+    assert fake.wait("send_durable_execution_callback_heartbeat")[0]["CallbackId"] == "cb-1"
+
+
+def test_failure_carries_the_error_type_and_data(env):
+    app, fake, post = env
+
+    @app.durable_job
+    def handler(inp, ctx):
+        raise LeaseError("BuildFailed", "exit 2", data={"step": 3})
+
+    run_hook(post, {"version": 1, "job": {"callbackId": "cb-2", "input": {}}})
+    err = fake.wait("send_durable_execution_callback_failure")[0]["Error"]
+    assert err["ErrorType"] == "BuildFailed" and err["ErrorMessage"] == "exit 2"
+    assert json.loads(err["ErrorData"]) == {"step": 3}
+
+
+def test_http_job_on_a_named_route_runs_once_per_callback(env):
+    app, fake, post = env
+    seen = []
+
+    @app.durable_job(path="/build")
+    def build(inp, ctx):
+        seen.append(inp)
+        return {"n": len(seen)}
+
+    run_hook(post, {"version": 1, "region": "us-east-1"})
+    job = {"version": 1, "region": "us-east-1", "microvmId": "mvm-1", "callbackId": "cb-3", "input": {"x": 1}}
+    assert post("/build", job)[0] == 202
+    assert post("/build", job)[0] == 202  # a second delivery is accepted and ignored
+    fake.wait("send_durable_execution_callback_success")
+    time.sleep(0.2)
+    assert seen == [{"x": 1}]
+
+
+def test_default_job_path_and_malformed_job_fails_its_callback(env):
+    app, fake, post = env
+    app.durable_job(lambda inp, ctx: "ok")
+    assert post(AWS_JOB_PATH, {"version": 1, "callbackId": "cb-4", "input": 1})[0] == 202
+    fake.wait("send_durable_execution_callback_success")
+    assert post(AWS_JOB_PATH, {"version": 9, "callbackId": "cb-5"})[0] == 400
+    assert fake.wait("send_durable_execution_callback_failure")[0]["CallbackId"] == "cb-5"
+
+
+def test_closed_callback_cancels_the_job_and_reports_nothing(env):
+    app, fake, post = env
+    fake.fail["send_durable_execution_callback_heartbeat"] = "CallbackTimeoutException"
+    cancelled = threading.Event()
+
+    @app.durable_job
+    def handler(inp, ctx):
+        ctx.cancelled.wait(5)
+        cancelled.set()
+        return "late"
+
+    run_hook(post, {"version": 1, "job": {"callbackId": "cb-6", "heartbeatTimeoutSeconds": 30, "input": {}}})
+    assert cancelled.wait(5)
+    time.sleep(0.2)
+    assert not [c for c in fake.calls if c[0] == "send_durable_execution_callback_success"]
+
+
+def test_terminate_from_lambda_fails_running_jobs_but_a_forwarded_one_is_ignored(env):
+    app, fake, post = env
+    release = threading.Event()
+    app.durable_job(lambda inp, ctx: release.wait(5) or "done")
+    run_hook(post, {"version": 1, "job": {"callbackId": "cb-7", "input": {}}})
+    time.sleep(0.2)
+    # forwarded through the endpoint: endpoint host name and a request id
+    assert post(HOOK + "terminate", {}, host="abc.lambda-microvm.us-east-1.on.aws",
+                extra={"x-amzn-requestid": "r-1"})[0] == 200
+    time.sleep(0.2)
+    assert not [c for c in fake.calls if c[0] == "send_durable_execution_callback_failure"]
+    assert post(HOOK + "terminate", {}, host="localhost:8080")[0] == 200
+    err = fake.wait("send_durable_execution_callback_failure")[0]["Error"]
+    assert err["ErrorType"] == "MicrovmTerminatedError"
+    assert post(AWS_JOB_PATH, {"version": 1, "callbackId": "cb-8", "input": {}})[0] == 503
+    release.set()
+
+
+def test_suspend_hook_refuses_jobs_until_resume(env):
+    app, fake, post = env
+    app.durable_job(lambda inp, ctx: "ok")
+    run_hook(post, {"version": 1})
+    post(HOOK + "suspend", {}, host="localhost:8080")
+    assert post(AWS_JOB_PATH, {"version": 1, "callbackId": "cb-9", "input": {}})[0] == 503
+    post(HOOK + "resume", {}, host="localhost:8080")
+    assert post(AWS_JOB_PATH, {"version": 1, "callbackId": "cb-9", "input": {}})[0] == 202
+
+
+def test_idle_microvm_suspends_itself(env):
+    app, fake, post = env
+    app.durable_job(lambda inp, ctx: "ok")
+    run_hook(post, {"version": 1, "autoSuspendIdleSeconds": 0.3})
+    call = fake.wait("suspend_microvm", timeout=3)[0]
+    assert call == {"microvmIdentifier": "mvm-1"}
+    assert post(AWS_JOB_PATH, {"version": 1, "callbackId": "cb-10", "input": {}})[0] == 503
+
+
+def test_a_lease_payload_still_goes_to_on_lease(env):
+    app, fake, post = env
+    got = threading.Event()
+    app.on_lease(lambda task, lease: got.set() or {"ok": True})
+    run_hook(post, {"lease": {"kind": "none", "id": "l-1"}, "task": {"steps": []}})
+    assert got.wait(5)
+
+
+def test_durable_lease_retries_a_transient_completion_and_reads_codes_from_client_errors(env):
+    app, fake, post = env
+    type(app).completion_backoff = 0
+    attempts = {"n": 0}
+    real = fake._call
+
+    def flaky(name, **kw):
+        if name == "send_durable_execution_callback_success":
+            attempts["n"] += 1
+            if attempts["n"] == 1:
+                fake.calls.append((name, kw))
+                raise ClientError("ServiceUnavailableException")
+        return real(name, **kw)
+
+    fake._call = flaky
+    app.on_lease(lambda task, lease: {"ok": True})
+    run_hook(post, {"lease": {"kind": "durable", "token": "cb-11", "heartbeat_s": 30}, "task": {}})
+    fake.wait("send_durable_execution_callback_success", n=2)
+    assert attempts["n"] == 2
+
+
+def test_durable_lease_treats_an_unknown_callback_as_closed(env):
+    app, fake, post = env
+    type(app).completion_backoff = 0
+    fake.fail["send_durable_execution_callback_success"] = "ResourceNotFoundException"
+    fake.fail["send_durable_execution_callback_heartbeat"] = "ResourceNotFoundException"
+    app.on_lease(lambda task, lease: {"ok": True})
+    run_hook(post, {"lease": {"kind": "durable", "token": "cb-12", "heartbeat_s": 30}, "task": {}})
+    deadline = time.time() + 5
+    while time.time() < deadline and not (app.lease and app.lease.lost):
+        time.sleep(0.05)
+    assert app.lease.lost  # the first heartbeat found the callback gone
+    time.sleep(0.2)
+    assert not [c for c in fake.calls if c[0] == "send_durable_execution_callback_success"]

{microvm_ctl-0.3.1 → microvm_ctl-0.5.0}/tests/test_integrations_durable_scale.py

@@ -118,7 +118,8 @@ def test_two_shards_happy_path(fm):
         cb1 = runner.wait_for_callback(arn, name="job-1-0-callback", timeout=30)
         assert cb0 != cb1
         assert fm.plans == [(2, 2048, POLICY)]  # the plan step ran once, with the reconstructed policy
-        assert [launch["task"] for launch in fm.leases] == SHARDS
+        # map items launch concurrently, so compare as a set of shards, not in launch order
+        assert sorted((launch["task"] for launch in fm.leases), key=lambda t: t["i"]) == SHARDS
         assert all(launch["version"] == "3" and launch["policy"] == POLICY for launch in fm.leases)
         _succeed(runner, cb0, 1)
         _succeed(runner, cb1, 2)