runtime server stop with drain calls

launchflow · Nov 27, 2023 · 96c7ce1 · 96c7ce1
1 parent d951778
commit 96c7ce1
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/buildflow/core/app/runtime/server.py b/buildflow/core/app/runtime/server.py
@@ -7,7 +7,6 @@
 import uvicorn
 from fastapi import APIRouter, FastAPI
 from fastapi.responses import HTMLResponse, JSONResponse
-from ray import kill
 
 from buildflow.core.app.flow_state import FlowState
 from buildflow.core.app.infra.actors.infra import InfraActor
@@ -115,7 +114,9 @@ async def runtime_snapshot(self):
         return JSONResponse(snapshot.as_dict())
 
     async def runtime_stop(self):
-        kill(self.runtime_actor)
+        # Send two drain requests to stop the runtime.
+        self.runtime_actor.drain.remote()
+        self.runtime_actor.drain.remote()
 
     async def runtime_status(self):
         status = await self.runtime_actor.status.remote()