@@ -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