Merge branch 'aws-main'

Author: gorloffslava

.pre-commit-config.yaml

@@ -3,5 +3,5 @@ repos:
     rev: 19.3b0
     hooks:
       - id: black
-        language_version: python3.6
+        language_version: python3.9
         exclude_types: ['markdown', 'ini', 'toml', 'rst']

README.md

@@ -5,12 +5,12 @@ We have open-sourced a set of software packages, Runtime Interface Clients (RIC)
   base images to be Lambda compatible.
 The Lambda Runtime Interface Client is a lightweight interface that allows your runtime to receive requests from and send requests to the Lambda service.
 
-The Lambda Python Runtime Interface Client is vended through [pip](https://pypi.org/project/awslambdaric). 
+The Lambda Python Runtime Interface Client is vended through [pip](https://pypi.org/project/awslambdaric).
 You can include this package in your preferred base image to make that base image Lambda compatible.
 
 ## Requirements
 The Python Runtime Interface Client package currently supports Python versions:
- - 3.7.x up to and including 3.12.x
+ - 3.9.x up to and including 3.13.x
 
 ## Usage
 
@@ -103,18 +103,18 @@ def handler(event, context):
 
 ### Local Testing
 
-To make it easy to locally test Lambda functions packaged as container images we open-sourced a lightweight web-server, Lambda Runtime Interface Emulator (RIE), which allows your function packaged as a container image to accept HTTP requests. You can install the [AWS Lambda Runtime Interface Emulator](https://github.com/aws/aws-lambda-runtime-interface-emulator) on your local machine to test your function. Then when you run the image function, you set the entrypoint to be the emulator. 
+To make it easy to locally test Lambda functions packaged as container images we open-sourced a lightweight web-server, Lambda Runtime Interface Emulator (RIE), which allows your function packaged as a container image to accept HTTP requests. You can install the [AWS Lambda Runtime Interface Emulator](https://github.com/aws/aws-lambda-runtime-interface-emulator) on your local machine to test your function. Then when you run the image function, you set the entrypoint to be the emulator.
 
 *To install the emulator and test your Lambda function*
 
-1) From your project directory, run the following command to download the RIE from GitHub and install it on your local machine. 
+1) From your project directory, run the following command to download the RIE from GitHub and install it on your local machine.
 
 ```shell script
 mkdir -p ~/.aws-lambda-rie && \
     curl -Lo ~/.aws-lambda-rie/aws-lambda-rie https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/latest/download/aws-lambda-rie && \
     chmod +x ~/.aws-lambda-rie/aws-lambda-rie
 ```
-2) Run your Lambda image function using the docker run command. 
+2) Run your Lambda image function using the docker run command.
 
 ```shell script
 docker run -d -v ~/.aws-lambda-rie:/aws-lambda -p 9000:8080 \
@@ -123,9 +123,9 @@ docker run -d -v ~/.aws-lambda-rie:/aws-lambda -p 9000:8080 \
         /usr/local/bin/python -m awslambdaric app.handler
 ```
 
-This runs the image as a container and starts up an endpoint locally at `http://localhost:9000/2015-03-31/functions/function/invocations`. 
+This runs the image as a container and starts up an endpoint locally at `http://localhost:9000/2015-03-31/functions/function/invocations`.
 
-3) Post an event to the following endpoint using a curl command: 
+3) Post an event to the following endpoint using a curl command:
 
 ```shell script
 curl -XPOST "http://localhost:9000/2015-03-31/functions/function/invocations" -d '{}'
@@ -174,4 +174,4 @@ If you discover a potential security issue in this project we ask that you notif
 
 ## License
 
-This project is licensed under the Apache-2.0 License.
+This project is licensed under the Apache-2.0 License.

RELEASE.CHANGELOG.md

@@ -1,3 +1,8 @@
+### November 19, 2024
+`3.0.0`
+- Drop support for deprecated python versions ([#179](https://github.com/aws/aws-lambda-python-runtime-interface-client/pull/179))
+- Add support for snapstart runtime hooks ([#176](https://github.com/aws/aws-lambda-python-runtime-interface-client/pull/176))
+
 ### August 23, 2024
 `2.2.1`:
 - Patch libcurl configure.ac to work with later versions of autoconf ([#166](https://github.com/aws/aws-lambda-python-runtime-interface-client/pull/168))

awslambdaric/__init__.py

@@ -2,4 +2,4 @@
 Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 """
 
-__version__ = "2.2.1"
+__version__ = "3.0.0"

awslambdaric/bootstrap.py

@@ -31,6 +31,8 @@
 _AWS_LAMBDA_LOG_LEVEL = _get_log_level_from_env_var(
     os.environ.get("AWS_LAMBDA_LOG_LEVEL")
 )
+AWS_LAMBDA_INITIALIZATION_TYPE = "AWS_LAMBDA_INITIALIZATION_TYPE"
+INIT_TYPE_SNAP_START = "snap-start"
 
 
 def _get_handler(handler):
@@ -286,6 +288,29 @@ def extract_traceback(tb):
     ]
 
 
+def on_init_complete(lambda_runtime_client, log_sink):
+    from . import lambda_runtime_hooks_runner
+
+    try:
+        lambda_runtime_hooks_runner.run_before_snapshot()
+        lambda_runtime_client.restore_next()
+    except:
+        error_result = build_fault_result(sys.exc_info(), None)
+        log_error(error_result, log_sink)
+        lambda_runtime_client.post_init_error(
+            error_result, FaultException.BEFORE_SNAPSHOT_ERROR
+        )
+        sys.exit(64)
+
+    try:
+        lambda_runtime_hooks_runner.run_after_restore()
+    except:
+        error_result = build_fault_result(sys.exc_info(), None)
+        log_error(error_result, log_sink)
+        lambda_runtime_client.report_restore_error(error_result)
+        sys.exit(65)
+
+
 class LambdaLoggerHandler(logging.Handler):
     def __init__(self, log_sink):
         logging.Handler.__init__(self)
@@ -454,9 +479,10 @@ def run(app_root, handler, lambda_runtime_api_addr):
     sys.stdout = Unbuffered(sys.stdout)
     sys.stderr = Unbuffered(sys.stderr)
 
-    use_thread_for_polling_next = (
-        os.environ.get("AWS_EXECUTION_ENV") == "AWS_Lambda_python3.12"
-    )
+    use_thread_for_polling_next = os.environ.get("AWS_EXECUTION_ENV") in {
+        "AWS_Lambda_python3.12",
+        "AWS_Lambda_python3.13",
+    }
 
     with create_log_sink() as log_sink:
         lambda_runtime_client = LambdaRuntimeClient(
@@ -484,6 +510,9 @@ def run(app_root, handler, lambda_runtime_api_addr):
 
             sys.exit(1)
 
+        if os.environ.get(AWS_LAMBDA_INITIALIZATION_TYPE) == INIT_TYPE_SNAP_START:
+            on_init_complete(lambda_runtime_client, log_sink)
+
         while True:
             event_request = lambda_runtime_client.wait_next_invocation()
 

awslambdaric/lambda_runtime_client.py

@@ -62,25 +62,57 @@ def __init__(self, lambda_runtime_address, use_thread_for_polling_next=False):
             # Not defining symbol as global to avoid relying on TPE being imported unconditionally.
             self.ThreadPoolExecutor = ThreadPoolExecutor
 
-    def post_init_error(self, error_response_data):
+    def call_rapid(
+        self, http_method, endpoint, expected_http_code, payload=None, headers=None
+    ):
         # These imports are heavy-weight. They implicitly trigger `import ssl, hashlib`.
         # Importing them lazily to speed up critical path of a common case.
-        import http
         import http.client
 
         runtime_connection = http.client.HTTPConnection(self.lambda_runtime_address)
         runtime_connection.connect()
-        endpoint = "/2018-06-01/runtime/init/error"
-        headers = {ERROR_TYPE_HEADER: error_response_data["errorType"]}
-        runtime_connection.request(
-            "POST", endpoint, to_json(error_response_data), headers=headers
-        )
+        if http_method == "GET":
+            runtime_connection.request(http_method, endpoint)
+        else:
+            runtime_connection.request(
+                http_method, endpoint, to_json(payload), headers=headers
+            )
+
         response = runtime_connection.getresponse()
         response_body = response.read()
-
-        if response.code != http.HTTPStatus.ACCEPTED:
+        if response.code != expected_http_code:
             raise LambdaRuntimeClientError(endpoint, response.code, response_body)
 
+    def post_init_error(self, error_response_data, error_type_override=None):
+        import http
+
+        endpoint = "/2018-06-01/runtime/init/error"
+        headers = {
+            ERROR_TYPE_HEADER: (
+                error_type_override
+                if error_type_override
+                else error_response_data["errorType"]
+            )
+        }
+        self.call_rapid(
+            "POST", endpoint, http.HTTPStatus.ACCEPTED, error_response_data, headers
+        )
+
+    def restore_next(self):
+        import http
+
+        endpoint = "/2018-06-01/runtime/restore/next"
+        self.call_rapid("GET", endpoint, http.HTTPStatus.OK)
+
+    def report_restore_error(self, restore_error_data):
+        import http
+
+        endpoint = "/2018-06-01/runtime/restore/error"
+        headers = {ERROR_TYPE_HEADER: FaultException.AFTER_RESTORE_ERROR}
+        self.call_rapid(
+            "POST", endpoint, http.HTTPStatus.ACCEPTED, restore_error_data, headers
+        )
+
     def wait_next_invocation(self):
         # Calling runtime_client.next() from a separate thread unblocks the main thread,
         # which can then process signals.

awslambdaric/lambda_runtime_exception.py

@@ -11,6 +11,8 @@ class FaultException(Exception):
     IMPORT_MODULE_ERROR = "Runtime.ImportModuleError"
     BUILT_IN_MODULE_CONFLICT = "Runtime.BuiltInModuleConflict"
     MALFORMED_HANDLER_NAME = "Runtime.MalformedHandlerName"
+    BEFORE_SNAPSHOT_ERROR = "Runtime.BeforeSnapshotError"
+    AFTER_RESTORE_ERROR = "Runtime.AfterRestoreError"
     LAMBDA_CONTEXT_UNMARSHAL_ERROR = "Runtime.LambdaContextUnmarshalError"
     LAMBDA_RUNTIME_CLIENT_ERROR = "Runtime.LambdaRuntimeClientError"
 

awslambdaric/lambda_runtime_hooks_runner.py

@@ -0,0 +1,18 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from snapshot_restore_py import get_before_snapshot, get_after_restore
+
+
+def run_before_snapshot():
+    before_snapshot_callables = get_before_snapshot()
+    while before_snapshot_callables:
+        # Using pop as before checkpoint callables are executed in the reverse order of their registration
+        func, args, kwargs = before_snapshot_callables.pop()
+        func(*args, **kwargs)
+
+
+def run_after_restore():
+    after_restore_callables = get_after_restore()
+    for func, args, kwargs in after_restore_callables:
+        func(*args, **kwargs)

awslambdaric/lambda_runtime_marshaller.py

@@ -15,7 +15,10 @@
 # We also set 'ensure_ascii=False' so that the encoded json contains unicode characters instead of unicode escape sequences
 class Encoder(json.JSONEncoder):
     def __init__(self):
-        if os.environ.get("AWS_EXECUTION_ENV") == "AWS_Lambda_python3.12":
+        if os.environ.get("AWS_EXECUTION_ENV") in {
+            "AWS_Lambda_python3.12",
+            "AWS_Lambda_python3.13",
+        }:
             super().__init__(use_decimal=False, ensure_ascii=False, allow_nan=True)
         else:
             super().__init__(use_decimal=False, allow_nan=True)

requirements/base.txt

@@ -1 +1,2 @@
 simplejson>=3.18.4
+snapshot-restore-py>=1.0.0