support instantiating dispatch without FastAPI

Signed-off-by: Achille Roussel <achille.roussel@gmail.com>

Author: achille-roussel

src/dispatch/__init__.py

@@ -1,13 +1,22 @@
 """The Dispatch SDK for Python."""
-
 from __future__ import annotations
 
+import os
 import dispatch.integrations
+
+from concurrent import futures
+from http.server import HTTPServer
+from typing import Any, Callable, Coroutine, Optional, TypeVar, overload
+from typing_extensions import ParamSpec, TypeAlias
+from urllib.parse import urlsplit
+
 from dispatch.coroutine import all, any, call, gather, race
-from dispatch.function import DEFAULT_API_URL, Client, Registry, Reset
+from dispatch.function import DEFAULT_API_URL, Client, Function, Registry, Reset
+from dispatch.http import Dispatch
 from dispatch.id import DispatchID
 from dispatch.proto import Call, Error, Input, Output
 from dispatch.status import Status
+from dispatch.sdk.v1 import function_pb2_grpc as function_grpc
 
 __all__ = [
     "Call",
@@ -29,8 +38,43 @@
     "run",
 ]
 
-function = None
-primitive_function = None
 
+P = ParamSpec("P")
+T = TypeVar("T")
+
+_registry: Optional[Registry] = None
+
+def _default_registry():
+    global _registry
+    if not _registry:
+        _registry = Registry()
+    return _registry
+
+@overload
+def function(func: Callable[P, Coroutine[Any, Any, T]]) -> Function[P, T]: ...
+
+@overload
+def function(func: Callable[P, T]) -> Function[P, T]: ...
+
+def function(func):
+    return _default_registry().function(func)
+
+def run(port: str = os.environ.get("DISPATCH_ENDPOINT_ADDR", "[::]:8000")):
+    """Run the default dispatch server on the given port. The default server
+    uses a function registry where functions tagged by the `@dispatch.function`
+    decorator are registered.
+
+    This function is intended to be used with the `dispatch` CLI tool, which
+    automatically configures environment variables to connect the local server
+    to the Dispatch bridge API.
+
+    Args:
+        port: The address to bind the server to. Defaults to the value of the
+          DISPATCH_ENDPOINT_ADDR environment variable, or '[::]:8000' if it
+          wasn't set.
+    """
+    parsed_url = urlsplit('//' + port)
+    server_address = (parsed_url.hostname, parsed_url.port)
+    server = HTTPServer(server_address, Dispatch(_default_registry()))
+    server.serve_forever()
 
-def run(): ...

src/dispatch/fastapi.py

@@ -86,26 +86,7 @@ def __init__(
             raise ValueError(
                 "missing FastAPI app as first argument of the Dispatch constructor"
             )
-
-        endpoint_from = "endpoint argument"
-        if not endpoint:
-            endpoint = os.getenv("DISPATCH_ENDPOINT_URL")
-            endpoint_from = "DISPATCH_ENDPOINT_URL"
-        if not endpoint:
-            raise ValueError(
-                "missing application endpoint: set it with the DISPATCH_ENDPOINT_URL environment variable"
-            )
-
-        logger.info("configuring Dispatch endpoint %s", endpoint)
-
-        parsed_url = urlparse(endpoint)
-        if not parsed_url.netloc or not parsed_url.scheme:
-            raise ValueError(
-                f"{endpoint_from} must be a full URL with protocol and domain (e.g., https://example.com)"
-            )
-
         super().__init__(endpoint, api_key=api_key, api_url=api_url)
-
         verification_key = parse_verification_key(verification_key, url_scheme=parsed_url.scheme)
         function_service = _new_app(self, verification_key)
         app.mount("/dispatch.sdk.v1.FunctionService", function_service)
@@ -200,39 +181,39 @@ async def execute(request: fastapi.Request):
             raise _ConnectError(
                 500, "internal", f"function '{req.function}' fatal error"
             )
-        else:
-            response = output._message
-            status = Status(response.status)
 
-            if response.HasField("poll"):
-                logger.debug(
-                    "function '%s' polling with %d call(s)",
-                    req.function,
-                    len(response.poll.calls),
-                )
-            elif response.HasField("exit"):
-                exit = response.exit
-                if not exit.HasField("result"):
-                    logger.debug("function '%s' exiting with no result", req.function)
-                else:
-                    result = exit.result
-                    if result.HasField("output"):
-                        logger.debug(
-                            "function '%s' exiting with output value", req.function
-                        )
-                    elif result.HasField("error"):
-                        err = result.error
-                        logger.debug(
-                            "function '%s' exiting with error: %s (%s)",
-                            req.function,
-                            err.message,
-                            err.type,
-                        )
-                if exit.HasField("tail_call"):
+        response = output._message
+        status = Status(response.status)
+
+        if response.HasField("poll"):
+            logger.debug(
+                "function '%s' polling with %d call(s)",
+                req.function,
+                len(response.poll.calls),
+            )
+        elif response.HasField("exit"):
+            exit = response.exit
+            if not exit.HasField("result"):
+                logger.debug("function '%s' exiting with no result", req.function)
+            else:
+                result = exit.result
+                if result.HasField("output"):
                     logger.debug(
-                        "function '%s' tail calling function '%s'",
-                        exit.tail_call.function,
+                        "function '%s' exiting with output value", req.function
                     )
+                elif result.HasField("error"):
+                    err = result.error
+                    logger.debug(
+                        "function '%s' exiting with error: %s (%s)",
+                        req.function,
+                        err.message,
+                        err.type,
+                    )
+            if exit.HasField("tail_call"):
+                logger.debug(
+                    "function '%s' tail calling function '%s'",
+                    exit.tail_call.function,
+                )
 
         logger.debug("finished handling run request with status %s", status.name)
         return fastapi.Response(

src/dispatch/function.py

@@ -176,14 +176,16 @@ class Registry:
 
     def __init__(
         self,
-        endpoint: str,
+        endpoint: Optional[str] = None,
         api_key: Optional[str] = None,
         api_url: Optional[str] = None,
     ):
         """Initialize a function registry.
 
         Args:
             endpoint: URL of the endpoint that the function is accessible from.
+                Uses the value of the DISPATCH_ENDPOINT_URL environment variable
+                by default.
 
             api_key: Dispatch API key to use for authentication when
                 dispatching calls to functions. Uses the value of the
@@ -193,7 +195,24 @@ def __init__(
                 to functions. Uses the value of the DISPATCH_API_URL environment
                 variable if set, otherwise defaults to the public Dispatch API
                 (DEFAULT_API_URL).
+
+        Raises:
+            ValueError: If any of the required arguments are missing.
         """
+        endpoint_from = "endpoint argument"
+        if not endpoint:
+            endpoint = os.getenv("DISPATCH_ENDPOINT_URL")
+            endpoint_from = "DISPATCH_ENDPOINT_URL"
+        if not endpoint:
+            raise ValueError(
+                "missing application endpoint: set it with the DISPATCH_ENDPOINT_URL environment variable"
+            )
+        parsed_url = urlparse(endpoint)
+        if not parsed_url.netloc or not parsed_url.scheme:
+            raise ValueError(
+                f"{endpoint_from} must be a full URL with protocol and domain (e.g., https://example.com)"
+            )
+        logger.info("configuring Dispatch endpoint %s", endpoint)
         self.functions: Dict[str, PrimitiveFunction] = {}
         self.endpoint = endpoint
         self.client = Client(api_key=api_key, api_url=api_url)

src/dispatch/grpc.py

@@ -1,7 +1,50 @@
 """Integration of Dispatch functions with gRPC."""
 
-from dispatch.function import Batch, Registry
+import logging
+import os
+from typing import Optional, Union
 
+from dispatch.function import Registry
+from dispatch.proto import Input
+from dispatch.sdk.v1.function_pb2_grpc import FunctionServiceServicer
 
-class Dispatch(Registry):
+logger = logging.getLogger(__name__)
+
+
+class FunctionService(FunctionServiceServicer):
     """A Dispatch instance to be serviced by a gRPC server."""
+
+    def __init__(        self,            registry: Registry    ):
+        """Initialize a Dispatch gRPC service.
+
+        Args:
+            registry: The registry of functions to be serviced.
+        """
+        self.registry = registry
+
+    def Run(self, request, context):
+        if not request.function:
+            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
+            context.set_details("function is required")
+            return
+
+        try:
+            func = self.registry.functions[request.function]
+        except KeyError:
+            logger.debug("function '%s' not found", request.function)
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details(f"function '{request.function}' does not exist")
+            return
+
+        logger.info("running function '%s'", request.function)
+        try:
+            output = func._primitive_call(Input(request))
+        except Exception as e:
+            logger.error("function '%s' fatal error", request.function, exc_info=True)
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(f"function '{request.function}' fatal error")
+            return
+
+        return output._message
+
+# TODO: interceptor for verification key