feat(extras): reject an oversized MicroVM input

packages/aws-durable-execution-sdk-python-extras/README.md

@@ -123,7 +123,7 @@ except MicrovmError:
     ...  # Any other MicroVM failure.
 ```
 
-The class is the same on the first run and on every replay, so code that branches on it stays deterministic. A `microvm` call inside your own child context, or inside a `map` or `parallel` item, reaches the code around that context as the SDK's `ChildContextError`. Its `__cause__` is the MicroVM error, so check `isinstance(error.__cause__, MicrovmTimeoutError)` there. An invalid name or config raises the SDK's `ValidationError`, and an input that is not JSON-serializable raises `TypeError`, before any durable operation starts. `Duration` raises the same `ValidationError` for a negative value. The operation does not check `image_identifier` and `execution_role_arn` before it starts. boto3 checks them against the API model before it sends `RunMicrovm`, and the service checks them again. So an invalid value fails the launch step at once with `MicrovmLaunchError`, after the callback step.
+The class is the same on the first run and on every replay, so code that branches on it stays deterministic. A `microvm` call inside your own child context, or inside a `map` or `parallel` item, reaches the code around that context as the SDK's `ChildContextError`. Its `__cause__` is the MicroVM error, so check `isinstance(error.__cause__, MicrovmTimeoutError)` there. The operation raises three errors before any durable operation starts. An invalid name or config raises the SDK's `ValidationError`. An input that is not JSON-serializable raises `TypeError`. An input over `MAX_JOB_INPUT_BYTES` raises `ValueError`. `Duration` raises the same `ValidationError` for a negative value. The operation does not check `image_identifier` and `execution_role_arn` before it starts. boto3 checks them against the API model before it sends `RunMicrovm`, and the service checks them again. So an invalid value fails the launch step at once with `MicrovmLaunchError`, after the callback step.
 
 A failed AWS call keeps its request ID in the error message, in the form `<Operation> request ID: <id>.`, for example `RunMicrovm request ID: <id>.`. A support case can quote it. This applies to `RunMicrovm`, `GetMicrovm`, `ResumeMicrovm`, `CreateMicrovmAuthToken`, and `TerminateMicrovm`. The SDK checkpoints the message of a failed step, and not its `__cause__`. So the request ID is in the message, the handler sees it on the first run and on replay, and the execution history keeps it. A failed `TerminateMicrovm` call does not fail the operation. So a `TerminateMicrovm` request ID appears only in the terminate log entry and in the terminate step record. A message has no request ID when the answer carries none, for example after a connection failure. A message also has no request ID when the error is not the answer of the Lambda MicroVMs call. For example, a `client` with role-assuming credentials calls STS `AssumeRole` inside each call, and an STS error carries the request ID of STS. The operation recognizes such an error by its `operation_name`. The step raises a copy of the botocore error, of the same class, with the same `response`, and with the request ID in its message. So `default_microvm_retry_strategy` decides as before. A custom `retry_strategy` receives the error that botocore raised, without the request ID. So a strategy that matches the message, such as `create_retry_strategy(RetryStrategyConfig(retryable_errors=[...]))`, decides on the message of the service. A `ResourceNotFoundException` from `GetMicrovm`, `ResumeMicrovm`, or `CreateMicrovmAuthToken` is the exception. It means that the MicroVM no longer exists. So the operation raises `MicrovmStateError` instead, and the strategy receives that error. Its `__cause__` is the botocore error. An error that the operation creates itself, such as `MicrovmStateError` or "RunMicrovm returned no microvmId", has the request ID in its message when the strategy receives it. A custom client can raise one error object for several calls. Each step message then names only its own call, and the object keeps its message.
 
@@ -146,6 +146,7 @@ The 10-second reserve keeps time for the step to record its failure. A Lambda ti
 
 - `timeout` must be between 1 second and 8 hours, because a MicroVM lives at most 8 hours. The MicroVM's `maximumDurationInSeconds` is `timeout` plus 5 minutes, at most 8 hours. So a timeout above 7 hours 55 minutes leaves the MicroVM less than 5 minutes beyond the callback's timeout.
 - `job_input` is sent as JSON. So the MicroVM receives what `json.dumps` writes: a tuple arrives as a list, and a dict key that is not a string arrives as a string.
+- The worker accepts a job request body of at most 6 MiB. The body holds the JSON of `job_input` and a few job fields. The job fields take less than 3 KiB. So `job_input` can be at most `MAX_JOB_INPUT_BYTES`, 6 MiB minus 8 KiB, as UTF-8 JSON. A larger input raises `ValueError` from `microvm` or `vm.invoke` before any durable operation starts.
 - The MicroVM gets the `INTERNET_EGRESS` connector, because it needs egress to reach the Lambda callback APIs. For ingress, it gets `NO_INGRESS` with `run` hook delivery and `ALL_INGRESS` with HTTP delivery. If you set `ingress_network_connectors`, allow ingress for every job that goes over HTTP. A `microvm` job goes over HTTP when its `run` hook payload is over 4,096 code points, or when its config sets `request.path`. Every session job goes over HTTP. An empty tuple sends no connector.
 - `heartbeat_timeout` starts when the callback is created, before the launch. The worker sends its first heartbeat when the job arrives. So the first period must cover the launch, the boot, and, for HTTP delivery, the request step, each with its retries. A throttled launch can take several minutes, and the request step retries for up to `retry_window` in each attempt. With a shorter heartbeat timeout, the operation fails with `MicrovmTimeoutError` before the MicroVM sends its first heartbeat.
 - The request step needs at least 11 seconds left in the invocation when it starts. It keeps a 10-second reserve before the invocation times out, and it starts a call only with at least 1 second left before the reserve. The launch step can run earlier in the same invocation, so the function timeout must cover it too. A function timeout of `retry_window` plus 11 seconds, 71 seconds by default, leaves the first retry tier about its whole window. With less than 11 seconds left, every attempt fails with `MicrovmEndpointUnavailableError`, and the message of the final `MicrovmDeliveryError` names the function timeout. The Lambda default timeout is 3 seconds. With it, `microvm` runs small jobs in the `run` hook and fails every job that goes over HTTP: a job too large for the hook, and a job with `request.path`.
@@ -240,7 +241,7 @@ A failed `vm.invoke` raises its MicroVM error inside the handler. The handler ca
 1. A MicroVM error that the handler raises itself reaches the caller with its own class, such as `MicrovmJobFailedError`.
 2. Any other error reaches the caller as the SDK's `ChildContextError`, as from `run_in_child_context`. This includes a MicroVM error that leaves a child context that the handler created: that child context records `ChildContextError` as its error type. The caller then gets the MicroVM error's message, but not its class. Inside the handler, the `ChildContextError` of that child context has the MicroVM error as its `__cause__`. An SDK error that leaves the handler, such as the `StepError` of a failed step, also reaches the caller as `ChildContextError`, with the SDK error as its `__cause__`. This rule matches `run_in_child_context`. The JavaScript operation differs here: it passes an SDK error to the caller with its own class.
 
-An invalid job raises `ValidationError` or `TypeError` from `vm.invoke`, before the job's durable operations start. The handler can catch it too. A call without the context, `vm.invoke("build", job_input=..., options=...)`, raises Python's `TypeError`. A first argument that is not a durable context, as in `vm.invoke("build", "build", ...)`, raises `AttributeError`.
+An invalid job raises `ValidationError`, `TypeError`, or `ValueError` from `vm.invoke`, before the job's durable operations start. The handler can catch it too. A call without the context, `vm.invoke("build", job_input=..., options=...)`, raises Python's `TypeError`. A first argument that is not a durable context, as in `vm.invoke("build", "build", ...)`, raises `AttributeError`.
 
 The handler's return value is checkpointed as the session's result. So it must be serializable by the SDK, and it counts toward the checkpoint size limit.
 
@@ -299,9 +300,9 @@ The MicroVM operations handle the network, the service, and the worker lifecycle
 ### Job inputs
 
 1. Pass a JSON-serializable input with finite numbers. Any other input raises `TypeError` before any durable operation starts. See [Errors](#errors).
-2. Keep the serialized job request at or under 6 MiB. The job request holds the input and a few job fields. Pass larger data by reference, for example as an S3 key.
-   - The Python worker reads a job request body of at most 6 MiB. It answers a larger body with HTTP 400.
-   - The operation does not check the size before it sends the job. So a larger body reaches the worker, and the job fails with `MicrovmDeliveryError`. It fails at once when the 400 reaches the operation. It fails only after the retries of both tiers when the connection is reset instead, and the MicroVM is billed meanwhile. [What the worker does](../aws-durable-execution-sdk-python-microvm-worker/README.md#what-the-worker-does) says when each case happens.
+2. Keep the job input's JSON at or under `MAX_JOB_INPUT_BYTES`, 6 MiB minus 8 KiB, in UTF-8. Pass larger data by reference, for example as an S3 key.
+   - The Python worker reads a job request body of at most 6 MiB. The body holds the input's JSON and a few job fields.
+   - `microvm` and `vm.invoke` measure the input's JSON before any durable operation starts. A larger input raises `ValueError`. So `microvm` launches no MicroVM for it.
 3. Expect a job to go over HTTP in three cases:
    - The `run` hook payload of a `microvm` job is over 4,096 code points. The payload holds the input, the callback ID, and a few other fields.
    - The `microvm` config sets `request.path`. The job then goes over HTTP for any input size.

packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/__init__.py

@@ -25,6 +25,7 @@
     MicrovmTimeoutError,
 )
 from aws_durable_execution_sdk_python_extras.microvm.job import (
+    MAX_JOB_INPUT_BYTES,
     MAX_RUN_HOOK_PAYLOAD_LENGTH,
 )
 from aws_durable_execution_sdk_python_extras.microvm.operation import microvm
@@ -49,6 +50,7 @@
     "DEFAULT_AUTO_SUSPEND_IDLE_SECONDS",
     "DEFAULT_MICROVM_JOB_PATH",
     "DEFAULT_MICROVM_PORT",
+    "MAX_JOB_INPUT_BYTES",
     "MAX_MICROVM_DURATION_SECONDS",
     "MAX_RUN_HOOK_PAYLOAD_LENGTH",
     "MicrovmBaseConfig",

packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/job.py

@@ -29,6 +29,27 @@
 counts what the service counts.
 """
 
+MAX_JOB_INPUT_BYTES = 6 * 1024 * 1024 - 8 * 1024
+"""The largest job input, in UTF-8 bytes of its JSON, that ``microvm`` and
+``vm.invoke`` accept.
+
+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.
+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.
+5. The Lambda model allows a callback ID of at most 1,024 characters of
+   ``[A-Za-z0-9+/]+={0,2}``. JSON writes each of these characters in 1 byte.
+   So the callback ID takes at most 1,024 bytes.
+6. The Region is short ASCII. The heartbeat timeout is a few digits.
+7. So the fixed fields take less than 3 KiB. So 8 KiB covers them.
+
+So an input whose JSON has at most 6 MiB minus 8 KiB fits in the body.
+"""
+
 
 @dataclass(frozen=True, kw_only=True)
 class RunHookPayload:

packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/operation.py

@@ -6,7 +6,6 @@
 from __future__ import annotations
 
 import hashlib
-import json
 import os
 import re
 import threading
@@ -57,8 +56,10 @@
     MicrovmTimeoutError,
 )
 from aws_durable_execution_sdk_python_extras.microvm.job import (
+    MAX_JOB_INPUT_BYTES,
     JobRequest,
     RunHookPayload,
+    dumps,
     fits_run_hook,
     run_hook_payload_without_job,
 )
@@ -258,6 +259,9 @@ def microvm(
             value.
         TypeError: When ``job_input`` is not JSON-serializable. No durable
             operation starts.
+        ValueError: When the JSON of ``job_input`` is over
+            :data:`MAX_JOB_INPUT_BYTES` UTF-8 bytes. No durable operation
+            starts.
         MicrovmLaunchError: When ``RunMicrovm`` fails after all retries, or
             rejects the config, for example an empty ``image_identifier``.
         MicrovmNotRunningError: When ``CreateMicrovmAuthToken`` reports that
@@ -356,16 +360,29 @@ def _validate_heartbeat(name: str, heartbeat_timeout: Duration | None) -> int |
 
 
 def _validate_input(name: str, job_input: Any) -> None:
-    """Check that the job input is JSON-serializable.
+    """Check that the job input is JSON-serializable, and fits in a job body.
+
+    The size is the UTF-8 length of the JSON that the job body holds. So it
+    is measured on the output of :func:`dumps`.
 
     Raises:
-        TypeError: When it is not.
+        TypeError: When it is not JSON-serializable.
+        ValueError: When its JSON is over :data:`MAX_JOB_INPUT_BYTES` bytes.
     """
     try:
-        json.dumps(job_input, allow_nan=False)
+        text = dumps(job_input)
     except (TypeError, ValueError) as error:
         msg = f'MicroVM "{name}": the job input must be JSON-serializable: {error}'
         raise TypeError(msg) from error
+    size = len(text.encode())
+    if size > MAX_JOB_INPUT_BYTES:
+        msg = (
+            f'MicroVM "{name}": the job input is {size} bytes as JSON. The limit '
+            f"is {MAX_JOB_INPUT_BYTES} bytes, because the worker accepts a job "
+            "body of at most 6 MiB. Pass larger data by reference, for example "
+            "as an S3 key."
+        )
+        raise ValueError(msg)
 
 
 def _validate_region(context: DurableContext, name: str) -> tuple[str, str]:
@@ -720,6 +737,7 @@ def create(
             ValidationError: When ``name`` or ``config`` is not valid, or the
                 Region cannot be determined.
             TypeError: When ``job_input`` is not JSON-serializable.
+            ValueError: When the JSON of ``job_input`` is too large.
         """
         timeout_seconds = _validate_base(name, config.timeout)
         heartbeat_timeout_seconds = _validate_heartbeat(name, config.heartbeat_timeout)

packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py

@@ -250,6 +250,8 @@ def invoke(
             AttributeError: When ``context`` is not a durable context, for
                 example a job name passed as the first argument. Python
                 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.
             MicrovmNotRunningError: When the session's MicroVM has ended. The
                 session reports it when it checks the MicroVM state before

packages/aws-durable-execution-sdk-python-extras/tests/microvm/operation_test.py

@@ -36,6 +36,7 @@
 from botocore.stub import ANY, Stubber
 
 from aws_durable_execution_sdk_python_extras.microvm import (
+    MAX_JOB_INPUT_BYTES,
     MAX_RUN_HOOK_PAYLOAD_LENGTH,
     MicrovmConfig,
     MicrovmError,
@@ -366,6 +367,33 @@ def run_microvm(self, **request: Any) -> dict[str, Any]:
     assert "RunMicrovm returned no microvmId" in error["message"]
 
 
+def test_rejects_a_job_input_over_the_size_limit_before_the_launch(
+    run: Any, fake_client: Any
+) -> None:
+    client = fake_client()
+
+    @durable_execution
+    def handler(event: Any, context: DurableContext) -> Any:
+        # The JSON of the input is the limit plus 1 byte, with the 2 quotes.
+        job_input = "x" * (MAX_JOB_INPUT_BYTES - 1)
+        try:
+            microvm(context, "build", job_input=job_input, config=config_with(client))
+        except ValueError as error:
+            return str(error)
+        return None
+
+    result = run(handler, {})
+
+    assert outcome(result) == (
+        f'MicroVM "build": the job input is {MAX_JOB_INPUT_BYTES + 1} bytes as '
+        f"JSON. The limit is {MAX_JOB_INPUT_BYTES} bytes, because the worker "
+        "accepts a job body of at most 6 MiB. Pass larger data by reference, "
+        "for example as an S3 key."
+    )
+    assert client.calls == []
+    assert result.operations == []
+
+
 def test_sends_a_job_too_large_for_the_run_hook_over_http(
     run: Any, fake_client: Any, fake_endpoint: Any
 ) -> None:

packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py

@@ -38,6 +38,7 @@
 
 from aws_durable_execution_sdk_python_extras.microvm import (
     DEFAULT_MICROVM_JOB_PATH,
+    MAX_JOB_INPUT_BYTES,
     MicrovmError,
     MicrovmInvokeOptions,
     MicrovmJobFailedError,
@@ -405,6 +406,68 @@ def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any:
     ]
 
 
+def test_rejects_a_job_input_over_the_size_limit_before_the_job_starts(
+    run: Any, fake_client: Any, fake_endpoint: Any
+) -> None:
+    endpoint = fake_endpoint()
+    client = fake_client(endpoint=endpoint.url)
+
+    def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any:
+        # The JSON of the input is the limit plus 1 byte, with the 2 quotes.
+        job_input = "x" * (MAX_JOB_INPUT_BYTES - 1)
+        try:
+            vm.invoke(ctx, "build", job_input=job_input, options=job())
+        except ValueError as error:
+            return str(error)
+        return None
+
+    result = run(session_handler(session_config(client), body))
+
+    assert outcome(result) == {
+        "result": (
+            f'MicroVM "build": the job input is {MAX_JOB_INPUT_BYTES + 1} bytes '
+            f"as JSON. The limit is {MAX_JOB_INPUT_BYTES} bytes, because the "
+            "worker accepts a job body of at most 6 MiB. Pass larger data by "
+            "reference, for example as an S3 key."
+        )
+    }
+    # Only the session's own launch ran. No job operation started, and no
+    # request reached the MicroVM.
+    assert len(client.requests("run_microvm")) == 1
+    assert client.requests("create_microvm_auth_token") == []
+    assert endpoint.requests == []
+    session = result.get_context("pipeline")
+    assert [operation.name for operation in session.child_operations] == [
+        "pipeline.session",
+        "pipeline.launch",
+        "pipeline.terminate",
+    ]
+
+
+def test_sends_a_job_input_at_the_size_limit(
+    run: Any, fake_client: Any, fake_endpoint: Any
+) -> None:
+    endpoint = fake_endpoint()
+    client = fake_client(endpoint=endpoint.url)
+    # The JSON of the input is the limit, with the 2 quotes. Each "é" takes
+    # 2 bytes.
+    job_input = "é" * ((MAX_JOB_INPUT_BYTES - 2) // 2)
+
+    def body(vm: MicrovmSession, ctx: DurableContext, event: Any) -> Any:
+        return vm.invoke(ctx, "build", job_input=job_input, options=job())
+
+    result = run(
+        session_handler(session_config(client), body),
+        None,
+        [("build", succeed_with("done"))],
+    )
+
+    assert outcome(result) == {"result": "done"}
+    [posted] = endpoint.requests
+    assert posted["body"]["input"] == job_input
+    assert int(posted["headers"]["content-length"]) <= 6 * 1024 * 1024
+
+
 def test_retries_a_throttled_launch_with_the_same_client_token(
     run: Any, fake_client: Any, fake_endpoint: Any, make_error: Any
 ) -> None: