feat: Support --timeout option in serve and start cmd (#4061)

src/bentoml/_internal/server/http_app.py

@@ -106,14 +106,14 @@ def __init__(
         enable_metrics: bool = Provide[
             BentoMLContainer.api_server_config.metrics.enabled
         ],
-        timeout: int = Provide[BentoMLContainer.api_server_config.traffic.timeout],
         max_concurrency: int
         | None = Provide[BentoMLContainer.api_server_config.traffic.max_concurrency],
     ):
         self.bento_service = bento_service
         self.enable_access_control = enable_access_control
         self.access_control_options = access_control_options
         self.enable_metrics = enable_metrics
+        timeout = BentoMLContainer.api_server_config.traffic.timeout.get()
         super().__init__(timeout=timeout, max_concurrency=max_concurrency)
 
     @property

src/bentoml/serve.py

@@ -250,6 +250,7 @@ def serve_http_production(
     host: str = Provide[BentoMLContainer.http.host],
     backlog: int = Provide[BentoMLContainer.api_server_config.backlog],
     api_workers: int = Provide[BentoMLContainer.api_server_workers],
+    timeout: int | None = None,
     ssl_certfile: str | None = Provide[BentoMLContainer.ssl.certfile],
     ssl_keyfile: str | None = Provide[BentoMLContainer.ssl.keyfile],
     ssl_keyfile_password: str | None = Provide[BentoMLContainer.ssl.keyfile_password],
@@ -279,6 +280,7 @@ def serve_http_production(
     circus_socket_map: t.Dict[str, CircusSocket] = {}
     runner_bind_map: t.Dict[str, str] = {}
     uds_path = None
+    timeout_args = ["--timeout", str(timeout)] if timeout else []
 
     if psutil.POSIX:
         # use AF_UNIX sockets for Circus
@@ -317,6 +319,7 @@ def serve_http_production(
                             json.dumps(runner.scheduled_worker_env_map),
                             "--prometheus-dir",
                             prometheus_dir,
+                            *timeout_args,
                         ],
                         working_dir=working_dir,
                         numprocesses=runner.scheduled_worker_count,
@@ -377,6 +380,7 @@ def serve_http_production(
                                 "$(CIRCUS.WID)",
                                 "--worker-env-map",
                                 json.dumps(runner.scheduled_worker_env_map),
+                                *timeout_args,
                             ],
                             working_dir=working_dir,
                             numprocesses=runner.scheduled_worker_count,
@@ -442,6 +446,7 @@ def serve_http_production(
         "--prometheus-dir",
         prometheus_dir,
         *ssl_args,
+        *timeout_args,
     ]
 
     if development_mode:

src/bentoml/start.py

@@ -30,6 +30,7 @@ def start_runner_server(
     runner_name: str,
     port: int | None = None,
     host: str | None = None,
+    timeout: int | None = None,
     backlog: int = Provide[BentoMLContainer.api_server_config.backlog],
 ) -> None:
     """
@@ -48,6 +49,7 @@ def start_runner_server(
 
     working_dir = os.path.realpath(os.path.expanduser(working_dir))
     svc = load(bento_identifier, working_dir=working_dir, standalone_load=True)
+    timeout_args = ["--timeout", str(timeout)] if timeout else []
 
     from circus.sockets import CircusSocket  # type: ignore
     from circus.watcher import Watcher  # type: ignore
@@ -91,6 +93,7 @@ def start_runner_server(
                                 "$(circus.wid)",
                                 "--prometheus-dir",
                                 prometheus_dir,
+                                *timeout_args,
                             ],
                             working_dir=working_dir,
                             numprocesses=runner.scheduled_worker_count,
@@ -144,6 +147,7 @@ def start_http_server(
     host: str = Provide[BentoMLContainer.api_server_config.host],
     backlog: int = Provide[BentoMLContainer.api_server_config.backlog],
     api_workers: int = Provide[BentoMLContainer.api_server_workers],
+    timeout: int | None = None,
     ssl_certfile: str | None = Provide[BentoMLContainer.ssl.certfile],
     ssl_keyfile: str | None = Provide[BentoMLContainer.ssl.keyfile],
     ssl_keyfile_password: str | None = Provide[BentoMLContainer.ssl.keyfile_password],
@@ -169,6 +173,7 @@ def start_http_server(
 
     working_dir = os.path.realpath(os.path.expanduser(working_dir))
     svc = load(bento_identifier, working_dir=working_dir, standalone_load=True)
+    timeout_args = ["--timeout", str(timeout)] if timeout else []
     runner_requirements = {runner.name for runner in svc.runners}
     if not runner_requirements.issubset(set(runner_map)):
         raise ValueError(
@@ -213,6 +218,7 @@ def start_http_server(
                 "--prometheus-dir",
                 prometheus_dir,
                 *ssl_args,
+                *timeout_args,
             ],
             working_dir=working_dir,
             numprocesses=api_workers,

src/bentoml_cli/serve.py

@@ -96,6 +96,12 @@ def add_serve_command(cli: click.Group) -> None:
         envvar="BENTOML_API_WORKERS",
         show_default=True,
     )
+    @click.option(
+        "--timeout",
+        type=click.INT,
+        help="Specify the timeout (seconds) for API server and runners",
+        envvar="BENTOML_TIMEOUT",
+    )
     @click.option(
         "--backlog",
         type=click.INT,
@@ -174,6 +180,7 @@ def serve(  # type: ignore (unused warning)
         port: int,
         host: str,
         api_workers: int | None,
+        timeout: int | None,
         backlog: int,
         reload: bool,
         working_dir: str | None,
@@ -238,6 +245,7 @@ def serve(  # type: ignore (unused warning)
                 host=DEFAULT_DEV_SERVER_HOST if not host else host,
                 backlog=backlog,
                 api_workers=1,
+                timeout=timeout,
                 ssl_keyfile=ssl_keyfile,
                 ssl_certfile=ssl_certfile,
                 ssl_keyfile_password=ssl_keyfile_password,
@@ -255,6 +263,7 @@ def serve(  # type: ignore (unused warning)
                 port=port,
                 host=host,
                 api_workers=api_workers,
+                timeout=timeout,
                 ssl_keyfile=ssl_keyfile,
                 ssl_certfile=ssl_certfile,
                 ssl_keyfile_password=ssl_keyfile_password,

src/bentoml_cli/start.py

@@ -68,6 +68,12 @@ def add_start_command(cli: click.Group) -> None:
         help="Specify the number of API server workers to start. Default to number of available CPU cores in production mode",
         envvar="BENTOML_API_WORKERS",
     )
+    @click.option(
+        "--timeout",
+        type=click.INT,
+        help="Specify the timeout (seconds) for API server",
+        envvar="BENTOML_TIMEOUT",
+    )
     @click.option(
         "--working-dir",
         type=click.Path(),
@@ -128,6 +134,7 @@ def start_http_server(  # type: ignore (unused warning)
         backlog: int,
         working_dir: str | None,
         api_workers: int | None,
+        timeout: int | None,
         ssl_certfile: str | None,
         ssl_keyfile: str | None,
         ssl_keyfile_password: str | None,
@@ -173,6 +180,7 @@ def start_http_server(  # type: ignore (unused warning)
             host=host,
             backlog=backlog,
             api_workers=api_workers,
+            timeout=timeout,
             ssl_keyfile=ssl_keyfile,
             ssl_certfile=ssl_certfile,
             ssl_keyfile_password=ssl_keyfile_password,
@@ -227,6 +235,12 @@ def start_http_server(  # type: ignore (unused warning)
         default=None,
         show_default=True,
     )
+    @click.option(
+        "--timeout",
+        type=click.INT,
+        help="Specify the timeout (seconds) for runners",
+        envvar="BENTOML_TIMEOUT",
+    )
     @add_experimental_docstring
     def start_runner_server(  # type: ignore (unused warning)
         bento: str,
@@ -236,6 +250,7 @@ def start_runner_server(  # type: ignore (unused warning)
         host: str,
         backlog: int,
         working_dir: str,
+        timeout: int | None,
     ) -> None:
         """
         Start Runner server standalone. This will be used inside Yatai.
@@ -260,6 +275,7 @@ def start_runner_server(  # type: ignore (unused warning)
             bento,
             runner_name=runner_name,
             working_dir=working_dir,
+            timeout=timeout,
             port=port,
             host=host,
             backlog=backlog,

src/bentoml_cli/worker/http_api_server.py

@@ -91,6 +91,11 @@
     default=False,
     show_default=True,
 )
+@click.option(
+    "--timeout",
+    type=click.INT,
+    help="Specify the timeout for API server",
+)
 def main(
     bento_identifier: str,
     fd: int,
@@ -107,6 +112,7 @@ def main(
     ssl_ca_certs: str | None,
     ssl_ciphers: str | None,
     development_mode: bool,
+    timeout: int | None,
 ):
     """
     Start a HTTP api server worker.
@@ -133,6 +139,8 @@ def main(
 
     if runner_map is not None:
         BentoMLContainer.remote_runner_mapping.set(json.loads(runner_map))
+    if timeout is not None:
+        BentoMLContainer.api_server_config.traffic.timeout.set(timeout)
     svc = bentoml.load(bento_identifier, working_dir=working_dir, standalone_load=True)
 
     # setup context

src/bentoml_cli/worker/runner.py

@@ -39,6 +39,11 @@
     default=None,
     help="The environment variables to pass to the worker process. The format is a JSON string, e.g. '{0: {\"CUDA_VISIBLE_DEVICES\": 0}}'.",
 )
+@click.option(
+    "--timeout",
+    type=click.INT,
+    help="Specify the timeout (seconds) for runner",
+)
 @click.option(
     "--prometheus-dir",
     type=click.Path(exists=True),
@@ -51,6 +56,7 @@ def main(
     working_dir: t.Optional[str],
     no_access_log: bool,
     worker_id: int,
+    timeout: int | None,
     worker_env_map: str | None,
     prometheus_dir: str | None,
 ) -> None:
@@ -123,7 +129,10 @@ def main(
     else:
         raise ValueError(f"Runner {runner_name} not found")
 
-    app = RunnerAppFactory(runner, worker_index=worker_id)()
+    factory = RunnerAppFactory(runner, worker_index=worker_id)
+    if timeout is not None:
+        factory.timeout = timeout
+    app = factory()
 
     uvicorn_options: dict[str, int | None | str] = {
         "log_config": None,