Fix various issues of async servers (#135)

cacheflow/entrypoints/simple_fastapi_frontend.py

@@ -2,15 +2,16 @@ import argparse
 import json
 from typing import AsyncGenerator
 
-from fastapi import FastAPI, Request
+from fastapi import BackgroundTasks, FastAPI, Request
 from fastapi.responses import StreamingResponse
 import uvicorn
 
 from cacheflow.sampling_params import SamplingParams
-from cacheflow.server.arg_utils import ServerArgs
+from cacheflow.server.arg_utils import AsyncServerArgs
 from cacheflow.server.async_llm_server import AsyncLLMServer
-from cacheflow.server.ray_utils import initialize_cluster
+from cacheflow.utils import random_uuid
 
+TIMEOUT_KEEP_ALIVE = 5 # seconds.
 TIMEOUT_TO_PREVENT_DEADLOCK = 1 # seconds
 app = FastAPI()
 
@@ -20,7 +21,8 @@ async def generate_stream(request: Request) -> StreamingResponse:
     request_dict = await request.json()
     prompt = request_dict.pop("prompt")
     sampling_params = SamplingParams(**request_dict)
-    results_generator = server.generate(prompt, sampling_params)
+    request_id = random_uuid()
+    results_generator = server.generate(prompt, sampling_params, request_id)
 
     async def stream_results() -> AsyncGenerator[bytes, None]:
         async for request_output in results_generator:
@@ -35,17 +37,24 @@ async def generate_stream(request: Request) -> StreamingResponse:
             }
             yield (json.dumps(ret) + "\0").encode("utf-8")
 
-    return StreamingResponse(stream_results())
+    async def abort_request() -> None:
+        await server.abort(request_id)
+
+    background_tasks = BackgroundTasks()
+    # Abort the request if the client disconnects.
+    background_tasks.add_task(abort_request)
+    return StreamingResponse(stream_results(), background=background_tasks)
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--host", type=str, default="localhost")
     parser.add_argument("--port", type=int, default=8001)
-    parser = ServerArgs.add_cli_args(parser)
+    parser = AsyncServerArgs.add_cli_args(parser)
     args = parser.parse_args()
 
-    server_args = ServerArgs.from_cli_args(args)
+    server_args = AsyncServerArgs.from_cli_args(args)
     server = AsyncLLMServer.from_server_args(server_args)
 
-    uvicorn.run(app, host=args.host, port=args.port, log_level="info")
+    uvicorn.run(app, host=args.host, port=args.port, log_level="debug",
+                timeout_keep_alive=TIMEOUT_KEEP_ALIVE)