review changes to tests and server

Signed-off-by: technillogue <technillogue@gmail.com>

python/cog/server/http.py

@@ -203,7 +203,8 @@ async def _predict(
             return JSONResponse(jsonable_encoder(initial_response), status_code=202)
 
         try:
-            response = PredictionResponse(**(await async_result).dict())
+            prediction = await async_result
+            response = PredictionResponse(**prediction.dict())
         except ValidationError as e:
             _log_invalid_output(e)
             raise HTTPException(status_code=500, detail=str(e)) from e
@@ -246,7 +247,8 @@ async def _check_setup_result() -> Any:
         if not app.state.setup_result.done():
             return
 
-        result = await app.state.setup_result
+        # this can raise CancelledError
+        result = app.state.setup_result.result()
 
         if result["status"] == schema.Status.SUCCEEDED:
             app.state.health = Health.READY

python/cog/server/runner.py

@@ -129,8 +129,9 @@ def is_busy(self) -> bool:
         return False
 
     def shutdown(self) -> None:
+        if self._result:
+            self._result.cancel()
         self._worker.terminate()
-        # TODO: cancel setup or predict task
 
     def cancel(self, prediction_id: Optional[str] = None) -> None:
         if not self.is_busy():
@@ -280,6 +281,7 @@ async def setup(*, worker: Worker) -> Dict[str, Any]:
     try:
         # will be async
         for event in worker.setup():
+            await asyncio.sleep(0)
             if isinstance(event, Log):
                 logs.append(event.message)
             elif isinstance(event, Done):

python/tests/server/test_http.py

@@ -425,7 +425,6 @@ def test_prediction_idempotent_endpoint_conflict(client, match):
         json={"input": {"sleep": 1}},
         headers={"Prefer": "respond-async"},
     )
-    time.sleep(0.001)
     resp2 = client.put(
         "/predictions/5678efgh",
         json={"input": {"sleep": 1}},
@@ -493,12 +492,12 @@ def test_prediction_cancel(client):
     )
     assert resp.status_code == 202
 
-    resp = client.post("/predictions/123/cancel")
-    assert resp.status_code == 200
-
     resp = client.post("/predictions/456/cancel")
     assert resp.status_code == 404
 
+    resp = client.post("/predictions/123/cancel")
+    assert resp.status_code == 200
+
 
 @uses_predictor_with_client_options(
     "setup_weights",

python/tests/server/test_runner.py

@@ -1,3 +1,4 @@
+import asyncio
 import os
 import threading
 from datetime import datetime
@@ -75,7 +76,8 @@ async def test_prediction_runner_called_while_busy(runner):
 
     assert runner.is_busy()
     with pytest.raises(RunnerBusyError):
-        await runner.predict(request)[1]
+        _, task = runner.predict(request)
+        await task
 
     # Await to ensure that the first prediction is scheduled before we
     # attempt to shut down the runner.
@@ -90,7 +92,7 @@ async def test_prediction_runner_called_while_busy_idempotent(runner):
     runner.predict(request)
     _, async_result = runner.predict(request)
 
-    response = await async_result
+    response = await asyncio.wait_for(async_result, timeout=1)
     assert response.id == "abcd1234"
     assert response.output == "done in 0.1 seconds"
     assert response.status == "succeeded"