Skip to content

Commit 6661a09

Browse files
achille-rousselachille-roussel
authored
Merge pull request #175 from dispatchrun/aiohttp
interoperability with asyncio (part 2): integration with aiohttp
2 parents ffc358b + 58f5b59 commit 6661a09

File tree

5 files changed

+287
-18
lines changed

5 files changed

+287
-18
lines changed

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ readme = "README.md"
99
dynamic = ["version"]
1010
requires-python = ">= 3.8"
1111
dependencies = [
12+
"aiohttp >= 3.9.4",
1213
"grpcio >= 1.60.0",
1314
"protobuf >= 4.24.0",
1415
"types-protobuf >= 4.24.0.20240129",

‎src/dispatch/aiohttp.py‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
from typing import Optional, Union
2+
3+
from aiohttp import web
4+
5+
from dispatch.function import Registry
6+
from dispatch.http import (
7+
FunctionServiceError,
8+
function_service_run,
9+
make_error_response_body,
10+
)
11+
from dispatch.signature import Ed25519PublicKey, parse_verification_key
12+
13+
14+
class Dispatch(web.Application):
15+
"""A Dispatch instance servicing as a http server."""
16+
17+
registry: Registry
18+
verification_key: Optional[Ed25519PublicKey]
19+
20+
def __init__(
21+
self,
22+
registry: Registry,
23+
verification_key: Optional[Union[Ed25519PublicKey, str, bytes]] = None,
24+
):
25+
"""Initialize a Dispatch application.
26+
27+
Args:
28+
registry: The registry of functions to be serviced.
29+
30+
verification_key: The verification key to use for requests.
31+
"""
32+
super().__init__()
33+
self.registry = registry
34+
self.verification_key = parse_verification_key(verification_key)
35+
self.add_routes(
36+
[
37+
web.post(
38+
"/dispatch.sdk.v1.FunctionService/Run", self.handle_run_request
39+
),
40+
]
41+
)
42+
43+
async def handle_run_request(self, request: web.Request) -> web.Response:
44+
return await function_service_run_handler(
45+
request, self.registry, self.verification_key
46+
)
47+
48+
49+
class Server:
50+
host: str
51+
port: int
52+
app: Dispatch
53+
54+
_runner: web.AppRunner
55+
_site: web.TCPSite
56+
57+
def __init__(self, host: str, port: int, app: Dispatch):
58+
self.host = host
59+
self.port = port
60+
self.app = app
61+
62+
async def __aenter__(self):
63+
await self.start()
64+
return self
65+
66+
async def __aexit__(self, exc_type, exc_value, traceback):
67+
await self.stop()
68+
69+
async def start(self):
70+
self._runner = web.AppRunner(self.app)
71+
await self._runner.setup()
72+
73+
self._site = web.TCPSite(self._runner, self.host, self.port)
74+
await self._site.start()
75+
76+
async def stop(self):
77+
await self._site.stop()
78+
await self._runner.cleanup()
79+
80+
81+
def make_error_response(status: int, code: str, message: str) -> web.Response:
82+
body = make_error_response_body(code, message)
83+
return web.Response(status=status, content_type="application/json", body=body)
84+
85+
86+
def make_error_response_invalid_argument(message: str) -> web.Response:
87+
return make_error_response(400, "invalid_argument", message)
88+
89+
90+
def make_error_response_not_found(message: str) -> web.Response:
91+
return make_error_response(404, "not_found", message)
92+
93+
94+
def make_error_response_unauthenticated(message: str) -> web.Response:
95+
return make_error_response(401, "unauthenticated", message)
96+
97+
98+
def make_error_response_permission_denied(message: str) -> web.Response:
99+
return make_error_response(403, "permission_denied", message)
100+
101+
102+
def make_error_response_internal(message: str) -> web.Response:
103+
return make_error_response(500, "internal", message)
104+
105+
106+
async def function_service_run_handler(
107+
request: web.Request,
108+
function_registry: Registry,
109+
verification_key: Optional[Ed25519PublicKey],
110+
) -> web.Response:
111+
content_length = request.content_length
112+
if content_length is None or content_length == 0:
113+
return make_error_response_invalid_argument("content length is required")
114+
if content_length < 0:
115+
return make_error_response_invalid_argument("content length is negative")
116+
if content_length > 16_000_000:
117+
return make_error_response_invalid_argument("content length is too large")
118+
119+
data: bytes = await request.read()
120+
try:
121+
content = await function_service_run(
122+
str(request.url),
123+
request.method,
124+
dict(request.headers),
125+
data,
126+
function_registry,
127+
verification_key,
128+
)
129+
except FunctionServiceError as e:
130+
return make_error_response(e.status, e.code, e.message)
131+
return web.Response(status=200, content_type="application/proto", body=content)

‎src/dispatch/http.py‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,10 @@
2525

2626

2727
class Dispatch:
28-
"""A Dispatch instance to be serviced by a http server. The Dispatch class
29-
acts as a factory for DispatchHandler objects, by capturing the variables
30-
that would be shared between all DispatchHandler instances it created."""
28+
"""A Dispatch instance servicing as a http server."""
29+
30+
registry: Registry
31+
verification_key: Optional[Ed25519PublicKey]
3132

3233
def __init__(
3334
self,
@@ -38,6 +39,8 @@ def __init__(
3839
3940
Args:
4041
registry: The registry of functions to be serviced.
42+
43+
verification_key: The verification key to use for requests.
4144
"""
4245
self.registry = registry
4346
self.verification_key = parse_verification_key(verification_key)
@@ -92,7 +95,7 @@ def send_error_response_internal(self, message: str):
9295
self.send_error_response(500, "internal", message)
9396

9497
def send_error_response(self, status: int, code: str, message: str):
95-
body = f'{{"code":"{code}","message":"{message}"}}'.encode()
98+
body = make_error_response_body(code, message)
9699
self.send_response(status)
97100
self.send_header("Content-Type", self.error_content_type)
98101
self.send_header("Content-Length", str(len(body)))
@@ -234,3 +237,7 @@ async def function_service_run(
234237

235238
logger.debug("finished handling run request with status %s", status.name)
236239
return response.SerializeToString()
240+
241+
242+
def make_error_response_body(code: str, message: str) -> bytes:
243+
return f'{{"code":"{code}","message":"{message}"}}'.encode()

‎tests/test_aiohttp.py‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
import asyncio
2+
import base64
3+
import os
4+
import pickle
5+
import struct
6+
import threading
7+
import unittest
8+
from typing import Any, Tuple
9+
from unittest import mock
10+
11+
import fastapi
12+
import google.protobuf.any_pb2
13+
import google.protobuf.wrappers_pb2
14+
import httpx
15+
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
16+
17+
import dispatch.test.httpx
18+
from dispatch.aiohttp import Dispatch, Server
19+
from dispatch.asyncio import Runner
20+
from dispatch.experimental.durable.registry import clear_functions
21+
from dispatch.function import Arguments, Error, Function, Input, Output, Registry
22+
from dispatch.proto import _any_unpickle as any_unpickle
23+
from dispatch.sdk.v1 import call_pb2 as call_pb
24+
from dispatch.sdk.v1 import function_pb2 as function_pb
25+
from dispatch.signature import parse_verification_key, public_key_from_pem
26+
from dispatch.status import Status
27+
from dispatch.test import EndpointClient
28+
29+
public_key_pem = "-----BEGIN PUBLIC KEY-----\nMCowBQYDK2VwAyEAJrQLj5P/89iXES9+vFgrIy29clF9CC/oPPsw3c5D0bs=\n-----END PUBLIC KEY-----"
30+
public_key_pem2 = "-----BEGIN PUBLIC KEY-----\\nMCowBQYDK2VwAyEAJrQLj5P/89iXES9+vFgrIy29clF9CC/oPPsw3c5D0bs=\\n-----END PUBLIC KEY-----"
31+
public_key = public_key_from_pem(public_key_pem)
32+
public_key_bytes = public_key.public_bytes_raw()
33+
public_key_b64 = base64.b64encode(public_key_bytes)
34+
35+
from datetime import datetime
36+
37+
38+
def run(runner: Runner, server: Server, ready: threading.Event):
39+
try:
40+
with runner:
41+
runner.run(serve(server, ready))
42+
except RuntimeError as e:
43+
pass # silence errors triggered by stopping the loop after tests are done
44+
45+
46+
async def serve(server: Server, ready: threading.Event):
47+
async with server:
48+
ready.set() # allow the test to continue after the server started
49+
await asyncio.Event().wait()
50+
51+
52+
class TestAIOHTTP(unittest.TestCase):
53+
def setUp(self):
54+
ready = threading.Event()
55+
self.runner = Runner()
56+
57+
host = "127.0.0.1"
58+
port = 9997
59+
60+
self.endpoint = f"http://{host}:{port}"
61+
self.dispatch = Dispatch(
62+
Registry(
63+
endpoint=self.endpoint,
64+
api_key="0000000000000000",
65+
api_url="http://127.0.0.1:10000",
66+
),
67+
)
68+
69+
self.client = httpx.Client(timeout=1.0)
70+
self.server = Server(host, port, self.dispatch)
71+
self.thread = threading.Thread(
72+
target=lambda: run(self.runner, self.server, ready)
73+
)
74+
self.thread.start()
75+
ready.wait()
76+
77+
def tearDown(self):
78+
loop = self.runner.get_loop()
79+
loop.call_soon_threadsafe(loop.stop)
80+
self.thread.join(timeout=1.0)
81+
self.client.close()
82+
83+
def test_content_length_missing(self):
84+
resp = self.client.post(f"{self.endpoint}/dispatch.sdk.v1.FunctionService/Run")
85+
body = resp.read()
86+
self.assertEqual(resp.status_code, 400)
87+
self.assertEqual(
88+
body, b'{"code":"invalid_argument","message":"content length is required"}'
89+
)
90+
91+
def test_content_length_too_large(self):
92+
resp = self.client.post(
93+
f"{self.endpoint}/dispatch.sdk.v1.FunctionService/Run",
94+
data={"msg": "a" * 16_000_001},
95+
)
96+
body = resp.read()
97+
self.assertEqual(resp.status_code, 400)
98+
self.assertEqual(
99+
body, b'{"code":"invalid_argument","message":"content length is too large"}'
100+
)
101+
102+
def test_simple_request(self):
103+
@self.dispatch.registry.primitive_function
104+
async def my_function(input: Input) -> Output:
105+
return Output.value(
106+
f"You told me: '{input.input}' ({len(input.input)} characters)"
107+
)
108+
109+
http_client = dispatch.test.httpx.Client(httpx.Client(base_url=self.endpoint))
110+
client = EndpointClient(http_client)
111+
112+
pickled = pickle.dumps("Hello World!")
113+
input_any = google.protobuf.any_pb2.Any()
114+
input_any.Pack(google.protobuf.wrappers_pb2.BytesValue(value=pickled))
115+
116+
req = function_pb.RunRequest(
117+
function=my_function.name,
118+
input=input_any,
119+
)
120+
121+
resp = client.run(req)
122+
123+
self.assertIsInstance(resp, function_pb.RunResponse)
124+
125+
resp.exit.result.output.Unpack(
126+
output_bytes := google.protobuf.wrappers_pb2.BytesValue()
127+
)
128+
output = pickle.loads(output_bytes.value)
129+
130+
self.assertEqual(output, "You told me: 'Hello World!' (12 characters)")

‎tests/test_http.py‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import threading
66
import unittest
77
from http.server import HTTPServer
8-
from typing import Any
8+
from typing import Any, Tuple
99
from unittest import mock
1010

1111
import fastapi
@@ -34,21 +34,21 @@
3434
from datetime import datetime
3535

3636

37-
def create_dispatch_instance(endpoint: str):
38-
return Dispatch(
39-
Registry(
40-
endpoint=endpoint,
41-
api_key="0000000000000000",
42-
api_url="http://127.0.0.1:10000",
43-
),
44-
)
45-
46-
4737
class TestHTTP(unittest.TestCase):
4838
def setUp(self):
49-
self.server_address = ("127.0.0.1", 9999)
50-
self.endpoint = f"http://{self.server_address[0]}:{self.server_address[1]}"
51-
self.dispatch = create_dispatch_instance(self.endpoint)
39+
host = "127.0.0.1"
40+
port = 9999
41+
42+
self.server_address = (host, port)
43+
self.endpoint = f"http://{host}:{port}"
44+
self.dispatch = Dispatch(
45+
Registry(
46+
endpoint=self.endpoint,
47+
api_key="0000000000000000",
48+
api_url="http://127.0.0.1:10000",
49+
),
50+
)
51+
5252
self.client = httpx.Client(timeout=1.0)
5353
self.server = HTTPServer(self.server_address, self.dispatch)
5454
self.thread = threading.Thread(

0 commit comments

Comments
 (0)