Gracefully handle client disconnects (#625)

* More gracefully handle client disconnect.

* Handle prod servers, bump pyproject.

* Adjust error message

* add log statement to see if it shows up.

* Bump version.

* remove log.

* Copy the control server logic for inference server.

* Bump project version.

* Remove log line.

Author: squidarth

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.6.5rc2"
+version = "0.6.5rc5"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

truss/templates/control/control/endpoints.py

@@ -6,7 +6,7 @@
 from fastapi.responses import JSONResponse, StreamingResponse
 from helpers.errors import ModelLoadFailed, ModelNotReady
 from httpx import URL, ConnectError, RemoteProtocolError
-from starlette.requests import Request
+from starlette.requests import ClientDisconnect, Request
 from starlette.responses import Response
 from tenacity import Retrying, retry_if_exception_type, stop_after_attempt, wait_fixed
 
@@ -28,15 +28,20 @@ async def proxy(request: Request):
     client: httpx.AsyncClient = request.app.state.proxy_client
     url = URL(path=request.url.path, query=request.url.query.encode("utf-8"))
 
-    # 5 mins request and 2 min connect timeouts
-    # Large values; we don't want requests to fail due to timeout on the proxy
-    timeout = httpx.Timeout(5 * 60.0, connect=2 * 60.0)
+    # 2 min connect timeouts, no timeout for requests.
+    # We don't want requests to fail due to timeout on the proxy
+    timeout = httpx.Timeout(None, connect=2 * 60.0)
+    try:
+        request_body = await request.body()
+    except ClientDisconnect:
+        # If the client disconnects, we don't need to proxy the request
+        return Response(status_code=499)
 
     inf_serv_req = client.build_request(
         request.method,
         url,
         headers=request.headers.raw,
-        content=await request.body(),
+        content=request_body,
         timeout=timeout,
     )
 

truss/templates/server/common/truss_server.py

@@ -15,7 +15,7 @@
 import shared.util as utils
 import uvicorn
 from common.termination_handler_middleware import TerminationHandlerMiddleware
-from fastapi import Depends, FastAPI, Request
+from fastapi import Depends, FastAPI, HTTPException, Request
 from fastapi.responses import ORJSONResponse, StreamingResponse
 from fastapi.routing import APIRoute as FastAPIRoute
 from model_wrapper import ModelWrapper
@@ -26,6 +26,7 @@
     truss_msgpack_serialize,
 )
 from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import ClientDisconnect
 from starlette.responses import Response
 
 # [IMPORTANT] A lot of things depend on this currently.
@@ -41,7 +42,10 @@ async def parse_body(request: Request) -> bytes:
     """
     Used by FastAPI to read body in an asynchronous manner
     """
-    return await request.body()
+    try:
+        return await request.body()
+    except ClientDisconnect as exc:
+        raise HTTPException(status_code=499, detail="Client disconnected") from exc
 
 
 FORMAT = "%(asctime)s.%(msecs)03d %(name)s %(levelname)s [%(funcName)s():%(lineno)s] %(message)s"