chore(ml): improve shutdown (#5689)

machine-learning/app/config.py

@@ -1,12 +1,16 @@
 import logging
 import os
+import sys
 from pathlib import Path
+from socket import socket
 
-import gunicorn
 import starlette
+from gunicorn.arbiter import Arbiter
 from pydantic import BaseSettings
 from rich.console import Console
 from rich.logging import RichHandler
+from uvicorn import Server
+from uvicorn.workers import UvicornWorker
 
 from .schemas import ModelType
 
@@ -69,10 +73,26 @@ log_settings = LogSettings()
 class CustomRichHandler(RichHandler):
     def __init__(self) -> None:
         console = Console(color_system="standard", no_color=log_settings.no_color)
-        super().__init__(
-            show_path=False, omit_repeated_times=False, console=console, tracebacks_suppress=[gunicorn, starlette]
-        )
+        super().__init__(show_path=False, omit_repeated_times=False, console=console, tracebacks_suppress=[starlette])
 
 
 log = logging.getLogger("gunicorn.access")
 log.setLevel(LOG_LEVELS.get(log_settings.log_level.lower(), logging.INFO))
+
+
+# patches this issue https://github.com/encode/uvicorn/discussions/1803
+class CustomUvicornServer(Server):
+    async def shutdown(self, sockets: list[socket] | None = None) -> None:
+        for sock in sockets or []:
+            sock.close()
+        await super().shutdown()
+
+
+class CustomUvicornWorker(UvicornWorker):
+    async def _serve(self) -> None:
+        self.config.app = self.wsgi
+        server = CustomUvicornServer(config=self.config)
+        self._install_sigquit_handler()
+        await server.serve(sockets=self.sockets)
+        if not server.started:
+            sys.exit(Arbiter.WORKER_BOOT_ERROR)