mirror of
https://github.com/wassname/ray.git
synced 2026-06-28 02:01:24 +08:00
architkulkarni
78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
import json
|
|
|
|
import starlette.requests
|
|
|
|
|
|
def build_starlette_request(scope, serialized_body: bytes):
|
|
"""Build and return a Starlette Request from ASGI payload.
|
|
|
|
This function is intended to be used immediately before task invocation
|
|
happens.
|
|
"""
|
|
|
|
# Simulates receiving HTTP body from TCP socket. In reality, the body has
|
|
# already been streamed in chunks and stored in serialized_body.
|
|
async def mock_receive():
|
|
return {
|
|
"body": serialized_body,
|
|
"type": "http.request",
|
|
"more_body": False
|
|
}
|
|
|
|
return starlette.requests.Request(scope, mock_receive)
|
|
|
|
|
|
class Response:
|
|
"""ASGI compliant response class.
|
|
|
|
It is expected to be called in async context and pass along
|
|
`scope, receive, send` as in ASGI spec.
|
|
|
|
>>> await Response({"k": "v"}).send(scope, receive, send)
|
|
"""
|
|
|
|
def __init__(self, content=None, status_code=200):
|
|
"""Construct a HTTP Response based on input type.
|
|
|
|
Args:
|
|
content (optional): Any JSON serializable object.
|
|
status_code (int, optional): Default status code is 200.
|
|
"""
|
|
self.status_code = status_code
|
|
self.raw_headers = []
|
|
|
|
if content is None:
|
|
self.body = b""
|
|
self.set_content_type("text")
|
|
elif isinstance(content, bytes):
|
|
self.body = content
|
|
self.set_content_type("text")
|
|
elif isinstance(content, str):
|
|
self.body = content.encode("utf-8")
|
|
self.set_content_type("text-utf8")
|
|
else:
|
|
# Delayed import since utils depends on http_util
|
|
from ray.serve.utils import ServeEncoder
|
|
self.body = json.dumps(
|
|
content, cls=ServeEncoder, indent=2).encode()
|
|
self.set_content_type("json")
|
|
|
|
def set_content_type(self, content_type):
|
|
if content_type == "text":
|
|
self.raw_headers.append([b"content-type", b"text/plain"])
|
|
elif content_type == "text-utf8":
|
|
self.raw_headers.append(
|
|
[b"content-type", b"text/plain; charset=utf-8"])
|
|
elif content_type == "json":
|
|
self.raw_headers.append([b"content-type", b"application/json"])
|
|
else:
|
|
raise ValueError("Invalid content type {}".format(content_type))
|
|
|
|
async def send(self, scope, receive, send):
|
|
await send({
|
|
"type": "http.response.start",
|
|
"status": self.status_code,
|
|
"headers": self.raw_headers,
|
|
})
|
|
await send({"type": "http.response.body", "body": self.body})
|