feat: add uvicorn timeouts options (#4682)

* feat: add uvicorn timeouts options

* style: run pre-commit

* fix: fox docs format

* fix: delete for bento service

* ci: auto fixes from pre-commit.ci

For more information, see https://pre-commit.ci

* fix: set default value to None

* style: restore yaml format

* fix: set default value to None

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

src/_bentoml_impl/server/serving.py

@@ -156,6 +156,8 @@ def serve_http(
     bentoml_home: str = Provide[BentoMLContainer.bentoml_home],
     development_mode: bool = False,
     reload: bool = False,
+    timeout_keep_alive: int | None = None,
+    timeout_graceful_shutdown: int | None = None,
     dependency_map: dict[str, str] | None = None,
     service_name: str = "",
     threaded: bool = False,
@@ -167,6 +169,7 @@ def serve_http(
     from bentoml._internal.utils.analytics.usage_stats import track_serve
     from bentoml._internal.utils.circus import create_standalone_arbiter
     from bentoml.serve import construct_ssl_args
+    from bentoml.serve import construct_timeouts_args
     from bentoml.serve import create_watcher
     from bentoml.serve import ensure_prometheus_dir
     from bentoml.serve import make_reload_plugin
@@ -257,6 +260,10 @@ def serve_http(
             ssl_ca_certs=ssl_ca_certs,
             ssl_ciphers=ssl_ciphers,
         )
+        timeouts_args = construct_timeouts_args(
+            timeout_keep_alive=timeout_keep_alive,
+            timeout_graceful_shutdown=timeout_graceful_shutdown,
+        )
         timeout_args = ["--timeout", str(timeout)] if timeout else []
 
         server_args = [
@@ -274,6 +281,7 @@ def serve_http(
             "--prometheus-dir",
             prometheus_dir,
             *ssl_args,
+            *timeouts_args,
             *timeout_args,
         ]
         if worker_envs:

src/_bentoml_impl/worker/service.py

@@ -82,6 +82,18 @@
     default=None,
     help="Ciphers to use (see stdlib 'ssl' module)",
 )
+@click.option(
+    "--timeout-keep-alive",
+    type=int,
+    default=None,
+    help="Close Keep-Alive connections if no new data is received within this timeout.",
+)
+@click.option(
+    "--timeout-graceful-shutdown",
+    type=int,
+    default=None,
+    help="Maximum number of seconds to wait for graceful shutdown. After this timeout, the server will start terminating requests.",
+)
 @click.option(
     "--development-mode",
     type=click.BOOL,
@@ -111,6 +123,8 @@ def main(
     ssl_cert_reqs: int | None,
     ssl_ca_certs: str | None,
     ssl_ciphers: str | None,
+    timeout_keep_alive: int | None,
+    timeout_graceful_shutdown: int | None,
     development_mode: bool,
     timeout: int,
 ):
@@ -193,6 +207,8 @@ def main(
         ssl_keyfile=ssl_keyfile,
         ssl_keyfile_password=ssl_keyfile_password,
         ssl_ca_certs=ssl_ca_certs,
+        timeout_keep_alive=timeout_keep_alive,
+        timeout_graceful_shutdown=timeout_graceful_shutdown,
         server_header=False,
         **uvicorn_extra_options,
     )

src/bentoml/serve.py

@@ -170,6 +170,20 @@ def construct_ssl_args(
     return args
 
 
+def construct_timeouts_args(
+    timeout_keep_alive: int | None,
+    timeout_graceful_shutdown: int | None,
+) -> list[str]:
+    args: list[str] = []
+
+    if timeout_keep_alive:
+        args.extend(["--timeout-keep-alive", str(timeout_keep_alive)])
+    if timeout_graceful_shutdown:
+        args.extend(["--timeout-graceful-shutdown", str(timeout_graceful_shutdown)])
+
+    return args
+
+
 def find_triton_binary():
     binary = shutil.which("tritonserver")
     if binary is None:
@@ -219,6 +233,8 @@ def serve_http_development(
     ssl_ca_certs: str | None = Provide[BentoMLContainer.ssl.ca_certs],
     ssl_ciphers: str | None = Provide[BentoMLContainer.ssl.ciphers],
     reload: bool = False,
+    timeout_keep_alive: int | None = None,
+    timeout_graceful_shutdown: int | None = None,
 ) -> None:
     logger.warning(
         "serve_http_development is deprecated. Please use serve_http_production with api_workers=1 and development_mode=True"
@@ -241,6 +257,8 @@ def serve_http_development(
         reload=reload,
         api_workers=1,
         development_mode=True,
+        timeout_keep_alive=timeout_keep_alive,
+        timeout_graceful_shutdown=timeout_graceful_shutdown,
     )
 
 
@@ -301,6 +319,8 @@ def serve_http_production(
     bentoml_home: str = Provide[BentoMLContainer.bentoml_home],
     development_mode: bool = False,
     reload: bool = False,
+    timeout_keep_alive: int | None = None,
+    timeout_graceful_shutdown: int | None = None,
 ) -> None:
     prometheus_dir = ensure_prometheus_dir()
 
@@ -428,6 +448,11 @@ def serve_http_production(
         ssl_ciphers=ssl_ciphers,
     )
 
+    timeouts_args = construct_timeouts_args(
+        timeout_keep_alive=timeout_keep_alive,
+        timeout_graceful_shutdown=timeout_graceful_shutdown,
+    )
+
     api_server_args = [
         "-m",
         SCRIPT_API_SERVER,
@@ -445,6 +470,7 @@ def serve_http_production(
         "--prometheus-dir",
         prometheus_dir,
         *ssl_args,
+        *timeouts_args,
         *timeout_args,
     ]
 

src/bentoml/server.py

@@ -337,9 +337,12 @@ def __init__(
         ssl_cert_reqs: int | None = Provide[BentoMLContainer.ssl.cert_reqs],
         ssl_ca_certs: str | None = Provide[BentoMLContainer.ssl.ca_certs],
         ssl_ciphers: str | None = Provide[BentoMLContainer.ssl.ciphers],
+        timeout_keep_alive: int | None = None,
+        timeout_graceful_shutdown: int | None = None,
     ):
         # hacky workaround to prevent bentoml.serve being overwritten immediately
         from .serve import construct_ssl_args
+        from .serve import construct_timeouts_args
 
         super().__init__(
             bento,
@@ -369,6 +372,13 @@ def __init__(
 
         self.args.extend(construct_ssl_args(**ssl_args))
 
+        timeouts_args = {
+            "timeout_keep_alive": timeout_keep_alive,
+            "timeout_graceful_shutdown": timeout_graceful_shutdown,
+        }
+
+        self.args.extend(construct_timeouts_args(**timeouts_args))
+
     def get_client(self) -> HTTPClient:
         return super().get_client()
 

src/bentoml/start.py

@@ -159,6 +159,8 @@ def start_http_server(
     ssl_cert_reqs: int | None = Provide[BentoMLContainer.ssl.cert_reqs],
     ssl_ca_certs: str | None = Provide[BentoMLContainer.ssl.ca_certs],
     ssl_ciphers: str | None = Provide[BentoMLContainer.ssl.ciphers],
+    timeout_keep_alive: int | None = None,
+    timeout_graceful_shutdown: int | None = None,
 ) -> None:
     from .serve import ensure_prometheus_dir
 
@@ -173,6 +175,7 @@ def start_http_server(
     from .serve import API_SERVER_NAME
     from .serve import PROMETHEUS_MESSAGE
     from .serve import construct_ssl_args
+    from .serve import construct_timeouts_args
     from .serve import create_watcher
 
     working_dir = os.path.realpath(os.path.expanduser(working_dir))
@@ -201,6 +204,10 @@ def start_http_server(
         ssl_ca_certs=ssl_ca_certs,
         ssl_ciphers=ssl_ciphers,
     )
+    timeouts_args = construct_timeouts_args(
+        timeout_keep_alive=timeout_keep_alive,
+        timeout_graceful_shutdown=timeout_graceful_shutdown,
+    )
     scheme = "https" if BentoMLContainer.ssl.enabled.get() else "http"
     watchers.append(
         create_watcher(
@@ -222,6 +229,7 @@ def start_http_server(
                 "--prometheus-dir",
                 prometheus_dir,
                 *ssl_args,
+                *timeouts_args,
                 *timeout_args,
             ],
             working_dir=working_dir,

src/bentoml_cli/serve.py

@@ -189,6 +189,22 @@ def cli():
         show_default=True,
         hidden=True,
     )
+    @click.option(
+        "--timeout-keep-alive",
+        type=int,
+        default=None,
+        help="Close Keep-Alive connections if no new data is received within this timeout.",
+        show_default=True,
+        hidden=True,
+    )
+    @click.option(
+        "--timeout-graceful-shutdown",
+        type=int,
+        default=None,
+        help="Maximum number of seconds to wait for graceful shutdown. After this timeout, the server will start terminating requests.",
+        show_default=True,
+        hidden=True,
+    )
     @env_manager
     def serve(  # type: ignore (unused warning)
         bento: str,
@@ -207,6 +223,8 @@ def serve(  # type: ignore (unused warning)
         ssl_cert_reqs: int | None,
         ssl_ca_certs: str | None,
         ssl_ciphers: str | None,
+        timeout_keep_alive: int | None,
+        timeout_graceful_shutdown: int | None,
         **attrs: t.Any,
     ) -> None:
         """Start a HTTP BentoServer from a given 🍱
@@ -276,6 +294,8 @@ def serve(  # type: ignore (unused warning)
                     ssl_ciphers=ssl_ciphers,
                     reload=reload,
                     development_mode=True,
+                    timeout_keep_alive=timeout_keep_alive,
+                    timeout_graceful_shutdown=timeout_graceful_shutdown,
                 )
             else:
                 serve_http_production(
@@ -294,6 +314,8 @@ def serve(  # type: ignore (unused warning)
                     ssl_ciphers=ssl_ciphers,
                     reload=reload,
                     development_mode=False,
+                    timeout_keep_alive=timeout_keep_alive,
+                    timeout_graceful_shutdown=timeout_graceful_shutdown,
                 )
         else:
             # bentoml>=1.2
@@ -317,6 +339,8 @@ def serve(  # type: ignore (unused warning)
                 ssl_ciphers=ssl_ciphers,
                 development_mode=development,
                 reload=reload,
+                timeout_keep_alive=timeout_keep_alive,
+                timeout_graceful_shutdown=timeout_graceful_shutdown,
             )
 
     @cli.command(name="serve-grpc", hidden=True)

src/bentoml_cli/start.py

@@ -136,6 +136,18 @@ def cli():
         default=BentoMLContainer.ssl.ciphers.get(),
         help="Ciphers to use (see stdlib 'ssl' module)",
     )
+    @click.option(
+        "--timeout-keep-alive",
+        type=int,
+        default=None,
+        help="Close Keep-Alive connections if no new data is received within this timeout.",
+    )
+    @click.option(
+        "--timeout-graceful-shutdown",
+        type=int,
+        default=None,
+        help="Maximum number of seconds to wait for graceful shutdown. After this timeout, the server will start terminating requests.",
+    )
     @add_experimental_docstring
     def start_http_server(  # type: ignore (unused warning)
         bento: str,
@@ -156,6 +168,8 @@ def start_http_server(  # type: ignore (unused warning)
         ssl_cert_reqs: int | None,
         ssl_ca_certs: str | None,
         ssl_ciphers: str | None,
+        timeout_keep_alive: int | None,
+        timeout_graceful_shutdown: int | None,
     ) -> None:
         """
         Start a HTTP API server standalone. This will be used inside Yatai.
@@ -207,6 +221,8 @@ def start_http_server(  # type: ignore (unused warning)
                     ssl_cert_reqs=ssl_cert_reqs,
                     ssl_ca_certs=ssl_ca_certs,
                     ssl_ciphers=ssl_ciphers,
+                    timeout_keep_alive=timeout_keep_alive,
+                    timeout_graceful_shutdown=timeout_graceful_shutdown,
                 )
             else:
                 from bentoml.start import start_runner_server
@@ -245,6 +261,8 @@ def start_http_server(  # type: ignore (unused warning)
                 ssl_cert_reqs=ssl_cert_reqs,
                 ssl_ca_certs=ssl_ca_certs,
                 ssl_ciphers=ssl_ciphers,
+                timeout_keep_alive=timeout_keep_alive,
+                timeout_graceful_shutdown=timeout_graceful_shutdown,
                 dependency_map=runner_map_dict,
                 service_name=service_name,
             )

src/bentoml_cli/worker/http_api_server.py

@@ -83,6 +83,18 @@
     default=None,
     help="Ciphers to use (see stdlib 'ssl' module)",
 )
+@click.option(
+    "--timeout-keep-alive",
+    type=click.INT,
+    default=None,
+    help="Close Keep-Alive connections if no new data is received within this timeout.",
+)
+@click.option(
+    "--timeout-graceful-shutdown",
+    type=click.INT,
+    default=None,
+    help="Maximum number of seconds to wait for graceful shutdown. After this timeout, the server will start terminating requests.",
+)
 @click.option(
     "--development-mode",
     type=click.BOOL,
@@ -111,6 +123,8 @@ def main(
     ssl_cert_reqs: int | None,
     ssl_ca_certs: str | None,
     ssl_ciphers: str | None,
+    timeout_keep_alive: int | None,
+    timeout_graceful_shutdown: int | None,
     development_mode: bool,
     timeout: int | None,
 ):
@@ -181,6 +195,11 @@ def main(
             ssl_ciphers = "TLSv1"
             uvicorn_options["ssl_ciphers"] = ssl_ciphers
 
+    if timeout_keep_alive:
+        uvicorn_options["timeout_keep_alive"] = timeout_keep_alive
+    if timeout_graceful_shutdown:
+        uvicorn_options["timeout_graceful_shutdown"] = timeout_graceful_shutdown
+
     if psutil.WINDOWS:
         uvicorn_options["loop"] = "asyncio"
         import asyncio