make HTTP routes async

python/cog/server/http.py

@@ -1,4 +1,6 @@
 import argparse
+import asyncio
+import functools
 import logging
 import os
 import signal
@@ -11,8 +13,6 @@
 
 import structlog
 import uvicorn
-from anyio import CapacityLimiter
-from anyio.lowlevel import RunVar
 from fastapi import Body, FastAPI, Header, HTTPException, Path, Response
 from fastapi.encoders import jsonable_encoder
 from fastapi.exceptions import RequestValidationError
@@ -79,27 +79,34 @@ def create_app(
         input_type=InputType, output_type=OutputType
     )
 
+    http_semaphore = asyncio.Semaphore(threads)
+
+    def limited(f: Callable) -> Callable:
+        @functools.wraps(f)
+        async def wrapped(*args: Any, **kwargs: Any) -> Any:
+            async with http_semaphore:
+                return await f(*args, **kwargs)
+
+        return wrapped
+
     @app.on_event("startup")
     def startup() -> None:
-        # https://github.com/tiangolo/fastapi/issues/4221
-        RunVar("_default_thread_limiter").set(CapacityLimiter(threads))  # type: ignore
-
         app.state.setup_result = runner.setup()
 
     @app.on_event("shutdown")
     def shutdown() -> None:
         runner.shutdown()
 
     @app.get("/")
-    def root() -> Any:
+    async def root() -> Any:
         return {
             # "cog_version": "", # TODO
             "docs_url": "/docs",
             "openapi_url": "/openapi.json",
         }
 
     @app.get("/health-check")
-    def healthcheck() -> Any:
+    async def healthcheck() -> Any:
         _check_setup_result()
         if app.state.health == Health.READY:
             health = Health.BUSY if runner.is_busy() else Health.READY
@@ -112,12 +119,13 @@ def healthcheck() -> Any:
             }
         )
 
+    @limited
     @app.post(
         "/predictions",
         response_model=PredictionResponse,
         response_model_exclude_unset=True,
     )
-    def predict(request: PredictionRequest = Body(default=None), prefer: Union[str, None] = Header(default=None)) -> Any:  # type: ignore
+    async def predict(request: PredictionRequest = Body(default=None), prefer: Union[str, None] = Header(default=None)) -> Any:  # type: ignore
         """
         Run a single prediction on the model
         """
@@ -131,12 +139,13 @@ def predict(request: PredictionRequest = Body(default=None), prefer: Union[str,
 
         return _predict(request=request, respond_async=respond_async)
 
+    @limited
     @app.put(
         "/predictions/{prediction_id}",
         response_model=PredictionResponse,
         response_model_exclude_unset=True,
     )
-    def predict_idempotent(
+    async def predict_idempotent(
         prediction_id: str = Path(..., title="Prediction ID"),
         request: PredictionRequest = Body(..., title="Prediction Request"),
         prefer: Union[str, None] = Header(default=None),
@@ -210,7 +219,7 @@ def _predict(
         return JSONResponse(content=encoded_response)
 
     @app.post("/predictions/{prediction_id}/cancel")
-    def cancel(prediction_id: str = Path(..., title="Prediction ID")) -> Any:
+    async def cancel(prediction_id: str = Path(..., title="Prediction ID")) -> Any:
         """
         Cancel a running prediction
         """
@@ -224,7 +233,7 @@ def cancel(prediction_id: str = Path(..., title="Prediction ID")) -> Any:
             return JSONResponse({}, status_code=200)
 
     @app.post("/shutdown")
-    def start_shutdown() -> Any:
+    async def start_shutdown() -> Any:
         log.info("shutdown requested via http")
         if shutdown_event is not None:
             shutdown_event.set()