Review feedback, remove the entrypoint decorator, fix tests and linter

Author: Pryz

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Test
 
 on:
-- push
+  - push
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.number || github.ref }}
@@ -15,43 +15,43 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python: ['3.8', '3.9', '3.10', '3.11', '3.12']
+        python: ["3.8", "3.9", "3.10", "3.11", "3.12"]
     steps:
-    - uses: actions/checkout@v4
-    - name: Set up Python ${{ matrix.python }}
-      uses: actions/setup-python@v5
-      with:
-        python-version: ${{ matrix.python }}
-        cache: pip
-    - run: make dev
-    - run: make test
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+          cache: pip
+      - run: make dev lambda
+      - run: make test
 
   examples:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python: ['3.12']
+        python: ["3.12"]
     steps:
-    - uses: actions/checkout@v4
-    - name: Set up Python ${{ matrix.python }}
-      uses: actions/setup-python@v5
-      with:
-        python-version: ${{ matrix.python }}
-        cache: pip
-    - run: make dev
-    - run: make exampletest
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+          cache: pip
+      - run: make dev
+      - run: make exampletest
 
   format:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python: ['3.12']
+        python: ["3.12"]
     steps:
-    - uses: actions/checkout@v4
-    - name: Set up Python ${{ matrix.python }}
-      uses: actions/setup-python@v5
-      with:
-        python-version: ${{ matrix.python }}
-        cache: pip
-    - run: make dev
-    - run: make fmt-check
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+          cache: pip
+      - run: make dev lambda
+      - run: make fmt-check

Makefile

@@ -17,6 +17,9 @@ install:
 dev:
 	$(PYTHON) -m pip install -e .[dev]
 
+lambda:
+	$(PYTHON) -m pip install -e .[lambda]
+
 fmt:
 	$(PYTHON) -m isort .
 	$(PYTHON) -m black .

examples/auto_retry/test_app.py

@@ -28,7 +28,6 @@ def test_app(self):
         endpoint_client = EndpointClient.from_app(app)
         dispatch_service = DispatchService(endpoint_client, collect_roundtrips=True)
         with DispatchServer(dispatch_service) as dispatch_server:
-
             # Use it when dispatching function calls.
             dispatch.set_client(Client(api_url=dispatch_server.url))
 

examples/getting_started/test_app.py

@@ -27,7 +27,6 @@ def test_app(self):
         endpoint_client = EndpointClient.from_app(app)
         dispatch_service = DispatchService(endpoint_client, collect_roundtrips=True)
         with DispatchServer(dispatch_service) as dispatch_server:
-
             # Use it when dispatching function calls.
             dispatch.set_client(Client(api_url=dispatch_server.url))
 

examples/github_stats/test_app.py

@@ -27,7 +27,6 @@ def test_app(self):
         endpoint_client = EndpointClient.from_app(app)
         dispatch_service = DispatchService(endpoint_client, collect_roundtrips=True)
         with DispatchServer(dispatch_service) as dispatch_server:
-
             # Use it when dispatching function calls.
             dispatch.set_client(Client(api_url=dispatch_server.url))
 

pyproject.toml

@@ -23,6 +23,7 @@ dependencies = [
 
 [project.optional-dependencies]
 fastapi = ["fastapi", "httpx"]
+lambda = ["awslambdaric"]
 
 dev = [
     "black >= 24.1.0",
@@ -33,7 +34,8 @@ dev = [
     "coverage >= 7.4.1",
     "requests >= 2.31.0",
     "types-requests >= 2.31.0.20240125",
-    "uvicorn >= 0.28.0"
+    "uvicorn >= 0.28.0",
+    "awslambdaric-stubs"
 ]
 
 docs = [

src/dispatch/experimental/lambda_handler.py

@@ -10,17 +10,20 @@
     def my_function():
         return "Hello World!"
 
-    @dispatch.entrypoint
+    @dispatch.function
     def entrypoint():
         my_function()
 
     def handler(event, context):
-        dispatch.handle(event, context)
+        dispatch.handle(event, context, entrypoint="entrypoint")
     """
 
 import base64
-import logging
 import json
+import logging
+from typing import Optional
+
+from awslambdaric.lambda_context import LambdaContext
 
 from dispatch.function import Registry
 from dispatch.proto import Input
@@ -33,43 +36,52 @@ def handler(event, context):
 class Dispatch(Registry):
     def __init__(
         self,
-        api_key: str | None = None,
-        api_url: str | None = None,
+        api_key: Optional[str] = None,
+        api_url: Optional[str] = None,
     ):
         """Initializes a Dispatch Lambda handler.
 
         Args:
-            api_key: Dispatch API key to use for authentication. Uses the value of
-                the DISPATCH_API_KEY environment variable by default.
+            api_key: Dispatch API key to use for authentication. Uses the value
+                of the DISPATCH_API_KEY environment variable by default.
 
             api_url: The URL of the Dispatch API to use. Uses the value of the
                 DISPATCH_API_URL environment variable if set, otherwise
                 defaults to the public Dispatch API (DEFAULT_API_URL).
 
         """
 
-        # The endpoint (AWS Lambda ARN) is set when handling the first request.
-        super().__init__(endpoint=None, api_key=api_key, api_url=api_url)
+        super().__init__(endpoint="not configured", api_key=api_key, api_url=api_url)
 
-    def handle(self, event, context):
-        # Use the context to determine the ARN of the Lambda function.
-        self.endpoint = context.invoked_function_arn
-
-        logger.debug("Dispatch handler invoked for %s with event: %s", self.endpoint, event)
+    def handle(
+        self, event: str, context: LambdaContext, entrypoint: Optional[str] = None
+    ):
+        # The ARN is not none until the first invocation of the Lambda function.
+        # We override the endpoint of all registered functions before any execution.
+        if context.invoked_function_arn:
+            self.endpoint = context.invoked_function_arn
+        self.override_endpoint(self.endpoint)
 
         if not event:
             raise ValueError("event is required")
 
         try:
             raw = base64.b64decode(event)
         except Exception as e:
-            raise ValueError(f"event is not base64 encoded: {e}")
+            raise ValueError("event is not base64 encoded") from e
 
         req = function_pb.RunRequest.FromString(raw)
-        print(req)
-        if not req.function:
-            req.function = "entrypoint"
-            # FIXME raise ValueError("function is required")
+
+        function: Optional[str] = req.function if req.function else entrypoint
+        if not function:
+            raise ValueError("function is required")
+
+        logger.debug(
+            "Dispatch handler invoked for %s function %s with runRequest: %s",
+            self.endpoint,
+            function,
+            req,
+        )
 
         try:
             func = self.functions[req.function]
@@ -117,6 +129,6 @@ def handle(self, event, context):
                     )
 
             logger.debug("finished handling run request with status %s", status.name)
-            resp = response.SerializeToString()
-            resp = base64.b64encode(resp).decode("utf-8")
-            return bytes(json.dumps(resp), "utf-8")
+            respBytes = response.SerializeToString()
+            respStr = base64.b64encode(respBytes).decode("utf-8")
+            return bytes(json.dumps(respStr), "utf-8")

src/dispatch/function.py

@@ -62,6 +62,10 @@ def __init__(
     def endpoint(self) -> str:
         return self._endpoint
 
+    @endpoint.setter
+    def endpoint(self, value: str):
+        self._endpoint = value
+
     @property
     def name(self) -> str:
         return self._name
@@ -210,17 +214,6 @@ def function(self, func):
         logger.info("registering coroutine: %s", name)
         return self._register_coroutine(name, func)
 
-    # TODO: ensure we have only 1 entrypoint per app.
-    def entrypoint(self, func):
-        """Decorated that registers the program entrypoint."""
-        name = "entrypoint"
-        if not inspect.iscoroutinefunction(func):
-            logger.info("registering entrypoint function: %s", name)
-            return self._register_function(name, func)
-
-        logger.info("registering entrypoint coroutine: %s", name)
-        return self._register_coroutine(name, func)
-
     def _register_function(self, name: str, func: Callable[P, T]) -> Function[P, T]:
         func = durable(func)
 
@@ -278,6 +271,10 @@ def batch(self) -> Batch:
         a set of calls to dispatch."""
         return self.client.batch()
 
+    def override_endpoint(self, endpoint: str):
+        for fn in self.functions.values():
+            fn.endpoint = endpoint
+
 
 class Client:
     """Client for the Dispatch API."""

src/dispatch/proto.py

@@ -10,7 +10,7 @@
 import google.protobuf.message
 import google.protobuf.wrappers_pb2
 import tblib  # type: ignore[import-untyped]
-from google.protobuf import duration_pb2
+from google.protobuf import descriptor_pool, duration_pb2, message_factory
 
 from dispatch.sdk.v1 import call_pb2 as call_pb
 from dispatch.sdk.v1 import error_pb2 as error_pb
@@ -61,13 +61,16 @@ class Input:
     def __init__(self, req: function_pb.RunRequest):
         self._has_input = req.HasField("input")
         if self._has_input:
-            input_pb = google.protobuf.wrappers_pb2.BytesValue()
-            req.input.Unpack(input_pb)
-            input_bytes = input_pb.value
-            try:
-                self._input = pickle.loads(input_bytes)
-            except EOFError:
-                self._input = Arguments(args=(req.input,), kwargs={})
+            if req.input.Is(google.protobuf.wrappers_pb2.BytesValue.DESCRIPTOR):
+                input_pb = google.protobuf.wrappers_pb2.BytesValue()
+                req.input.Unpack(input_pb)
+                input_bytes = input_pb.value
+                try:
+                    self._input = pickle.loads(input_bytes)
+                except Exception as e:
+                    self._input = input_bytes
+            else:
+                self._input = _pb_any_unpack(req.input)
         else:
             state_bytes = req.poll_result.coroutine_state
             if len(state_bytes) > 0:
@@ -436,3 +439,11 @@ def _pb_any_pickle(x: Any) -> google.protobuf.any_pb2.Any:
     pb_any = google.protobuf.any_pb2.Any()
     pb_any.Pack(pb_bytes)
     return pb_any
+
+
+def _pb_any_unpack(x: google.protobuf.any_pb2.Any) -> Any:
+    pool = descriptor_pool.Default()
+    msg_descriptor = pool.FindMessageTypeByName(x.TypeName())
+    proto = message_factory.GetMessageClass(msg_descriptor)()
+    x.Unpack(proto)
+    return proto

src/dispatch/scheduler.py

@@ -374,7 +374,6 @@ def _rebuild_state(self, input: Input):
             raise IncompatibleStateError from e
 
     def _run(self, input: Input) -> Output:
-
         if input.is_first_call:
             state = self._init_state(input)
         else: