Skip to content

Commit d4e5e3e

Browse files
romanlutzCopilot
andauthored
FIX: Drain requests on shutdown (#2965)
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 21bf73f6-4b26-4f5f-bb22-01fa72e49186
1 parent 9bdc3a4 commit d4e5e3e

5 files changed

Lines changed: 376 additions & 20 deletions

File tree

‎doc/gui/0_gui.md‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -460,6 +460,13 @@ prevents newly admitted work from overlapping replacement.
460460
Operation status survives a browser disconnect or navigation; other connected clients detect runtime generation
461461
changes and refresh catalogs without discarding chat or configuration drafts.
462462

463+
Backend shutdown closes runtime and management admission, then waits for accepted requests and any live apply
464+
to finish before stopping the scheduler and closing shared resources. This includes requests retained after a
465+
client disconnect and their offloaded writes. Cancelling the shutdown caller does not interrupt that cleanup;
466+
a single request or cleanup failure is re-raised unchanged, while multiple failures are reported together.
467+
Shutdown can therefore wait for an outstanding operation. Runtime readiness reports `ready: false` and
468+
`state: stopping` as soon as shutdown closes admission, including while an accepted live apply finishes.
469+
463470
If validation fails, PyRIT does not change the live runtime. Repair the saved source and retry. If startup fails, or
464471
if live initialization fails after replacement starts, runtime operations stay unavailable until you restart the
465472
backend. The administration UI stays available for configuration repair. Authentication and authorization retain

‎pyrit/backend/middleware/runtime.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,8 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
4646
apply_route = path.startswith("/api/config/runtime")
4747
write = management and scope["method"] not in ("GET", "HEAD", "OPTIONS") and not apply_route
4848
if (
49-
(write and (runtime.edit_lock.locked() or (runtime.apply_task and not runtime.apply_task.done())))
49+
runtime.is_stopping
50+
or (write and (runtime.edit_lock.locked() or (runtime.apply_task and not runtime.apply_task.done())))
5051
or (path.startswith("/api/initializers") and runtime.edit_lock.locked())
5152
or (path.startswith("/api/initializers") and runtime.state == "initializing")
5253
or (not management and runtime.state != "ready")

‎pyrit/backend/routes/health.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,12 @@ async def runtime_readiness_async(request: Request) -> dict[str, str | bool]:
2121
dict[str, str | bool]: Lightweight runtime readiness.
2222
"""
2323
runtime = getattr(request.app.state, "runtime_lifecycle", None)
24+
if runtime is None:
25+
return {"ready": False, "state": "failed", "generation": ""}
2426
return {
25-
"ready": runtime is not None and runtime.state == "ready",
26-
"state": runtime.state if runtime else "failed",
27-
"generation": runtime.generation if runtime else "",
27+
"ready": runtime.reported_state == "ready",
28+
"state": runtime.reported_state,
29+
"generation": runtime.generation,
2830
}
2931

3032

‎pyrit/backend/services/runtime_lifecycle.py‎

Lines changed: 49 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -55,15 +55,26 @@ def __init__(self, *, app: FastAPI, source: ConfigurationFileService) -> None:
5555
self.operations: set[asyncio.Task[None]] = set()
5656
self.management_operations: set[asyncio.Task[None]] = set()
5757
self.apply_task: asyncio.Task[None] | None = None
58+
self._shutdown_task: asyncio.Task[None] | None = None
5859
self.topology_supported = all(
5960
os.getenv(key, "1") == "1"
6061
for key in ("WEB_CONCURRENCY", "UVICORN_WORKERS", "PYRIT_API_WORKERS", "PYRIT_REPLICAS")
6162
)
6263

64+
@property
65+
def is_stopping(self) -> bool:
66+
"""Whether admission is closed even if an accepted apply is still finishing."""
67+
return self._shutdown_task is not None
68+
69+
@property
70+
def reported_state(self) -> str:
71+
"""Externally visible state, including admission closure before cleanup finishes."""
72+
return "stopping" if self.is_stopping else self.state
73+
6374
def status(self) -> dict[str, Any]:
6475
"""Return status recoverable after a disconnected apply."""
6576
return {
66-
"state": self.state,
77+
"state": self.reported_state,
6778
"generation": self.generation,
6879
"version": self.version,
6980
"outcome": self.outcome,
@@ -134,6 +145,8 @@ def begin_apply(self, *, version: str) -> dict[str, Any]:
134145
Returns:
135146
dict[str, Any]: Accepted operation or admission rejection.
136147
"""
148+
if self.is_stopping:
149+
return {**self.status(), "outcome": "stopping"}
137150
if not self.topology_supported:
138151
return {
139152
**self.status(),
@@ -232,17 +245,43 @@ def _has_active_work(self) -> bool:
232245
)
233246

234247
async def shutdown_async(self) -> None:
235-
"""Stop the current scheduler and close services owned by this process."""
236-
if self.apply_task and not self.apply_task.done():
237-
await asyncio.shield(self.apply_task)
248+
"""Drain admitted work before closure, deferring caller cancellation until cleanup finishes."""
249+
if self._shutdown_task is None:
250+
self._shutdown_task = asyncio.create_task(self._shutdown_runtime_async())
251+
cancellation: asyncio.CancelledError | None = None
252+
while not self._shutdown_task.done():
253+
try:
254+
await asyncio.shield(self._shutdown_task)
255+
except asyncio.CancelledError as error:
256+
cancellation = error
257+
self._shutdown_task.result()
258+
if cancellation is not None:
259+
raise cancellation
260+
261+
async def _shutdown_runtime_async(self) -> None:
262+
"""Finish retained requests without cancelling their offloaded writes."""
263+
pending = self.operations | self.management_operations
264+
if self.apply_task is not None:
265+
pending.add(self.apply_task)
266+
results = await asyncio.gather(*pending, return_exceptions=True)
267+
errors = [result for result in results if isinstance(result, BaseException)]
238268
self.state = "stopping"
239269
try:
240270
service = peek_scenario_run_service()
241271
if service:
242272
await service.shutdown_async()
243-
finally:
244-
try:
245-
await close_services_async()
246-
finally:
247-
if CentralMemory._memory_instance is not None:
248-
await CentralMemory.get_memory_instance().dispose_engine_async()
273+
except (Exception, asyncio.CancelledError) as error:
274+
errors.append(error)
275+
try:
276+
await close_services_async()
277+
except (Exception, asyncio.CancelledError) as error:
278+
errors.append(error)
279+
try:
280+
if CentralMemory._memory_instance is not None:
281+
await CentralMemory.get_memory_instance().dispose_engine_async()
282+
except (Exception, asyncio.CancelledError) as error:
283+
errors.append(error)
284+
if len(errors) == 1:
285+
raise errors[0]
286+
if errors:
287+
raise BaseExceptionGroup("Runtime shutdown failed after draining admitted work.", errors)

0 commit comments

Comments
 (0)