fix: make Uvicorn event loop configurable (kserve#4971)

Signed-off-by: Joshy Joseph <joshyjoseph.act@gmail.com>
Co-authored-by: Jooho Lee <jlee@redhat.com>
Co-authored-by: Sivanantham <90966311+sivanantha321@users.noreply.github.com>

Author: 3 people

python/kserve/docs/EventLoop.md

@@ -0,0 +1,36 @@
+# HTTP Server Event Loop Configuration
+
+KServe allows configuring the event loop implementation used by the HTTP server.
+This can be useful for performance tuning or for ensuring compatibility with
+specific runtime environments.
+
+## Configuration
+
+The event loop is configured using the `--event-loop` command-line argument.
+
+Supported values:
+
+- `auto` (default): Automatically select the event loop. If `uvloop` is installed,
+  it will be used; otherwise, the standard `asyncio` event loop is used.
+- `asyncio`: Force the use of Python’s built-in `asyncio` event loop.
+- `uvloop`: Force the use of `uvloop` (requires `uvloop` to be installed).
+
+## Example
+
+```bash
+kserve start \
+  --event-loop uvloop \
+  --http_port 8080
+```
+
+```python
+from kserve import ModelServer
+
+server = ModelServer(
+    http_port=8080,
+    event_loop="uvloop",  # "auto", "asyncio", or "uvloop"
+)
+
+# Register models and start the server
+server.start(models=[])
+```

python/kserve/kserve/model_server.py

@@ -116,6 +116,18 @@
     help="The asgi access logging format. It allows to override only the `uvicorn.access`'s format configuration "
     "with a richer set of fields",
 )
+parser.add_argument(
+    "--event-loop",
+    dest="event_loop",
+    default="auto",
+    type=str,
+    choices=["auto", "asyncio", "uvloop"],
+    help=(
+        "Event loop implementation used by the HTTP server. "
+        "Valid values are 'auto' (default; use uvloop if available), "
+        "'asyncio', or 'uvloop'."
+    ),
+)
 
 # Model arguments: The arguments are passed to the kserve.Model object
 parser.add_argument(
@@ -206,6 +218,7 @@ def __init__(
         enable_docs_url: bool = args.enable_docs_url,
         enable_latency_logging: bool = args.enable_latency_logging,
         access_log_format: str = args.access_log_format,
+        event_loop: str = args.event_loop,
         grace_period: int = 30,
         predictor_config: Optional[PredictorConfig] = None,
     ):
@@ -227,6 +240,7 @@ def __init__(
                                ASGI specs that don't describe how access logging should be implemented in detail
                                (please refer to this Uvicorn
                                [github issue](https://github.com/encode/uvicorn/issues/527) for more info).
+            event_loop: Uvicorn event loop. Default: ``'auto'``. It supports "auto", "asyncio", "uvloop".
             grace_period: The grace period in seconds to wait for the server to stop. Default: ``30``.
             predictor_config: Optional configuration for the predictor. Default: ``None``.
         """
@@ -236,6 +250,7 @@ def __init__(
         self.http_port = http_port
         self.grpc_port = grpc_port
         self.workers = workers
+        self.event_loop = event_loop
         self.max_threads = max_threads
         self.max_asyncio_workers = max_asyncio_workers
         self.enable_grpc = enable_grpc
@@ -317,6 +332,7 @@ def start(self, models: List[BaseKServeModel]) -> None:
                 workers=self.workers,
                 grace_period=self.grace_period,
                 log_config_file=args.log_config_file,
+                event_loop=self.event_loop,
             )
             self.servers.append(self._rest_multiprocess_server.start())
         else:
@@ -328,6 +344,7 @@ def start(self, models: List[BaseKServeModel]) -> None:
                 access_log_format=self.access_log_format,
                 workers=self.workers,
                 grace_period=self.grace_period,
+                event_loop=self.event_loop,
             )
             self.servers.append(self._rest_server.start())
         if self.enable_grpc:

python/kserve/kserve/protocol/rest/multiprocess/server.py

@@ -132,6 +132,7 @@ def __init__(
         workers: int = 1,
         grace_period: int = 30,
         log_config_file: Optional[str] = None,
+        event_loop: str = "auto",
     ) -> None:
         self.log_config_file = log_config_file
         self._rest_server = RESTServer(
@@ -142,6 +143,7 @@ def __init__(
             access_log_format,
             workers,
             grace_period,
+            event_loop,
         )
         self._processes: List[RESTServerProcess] = []
         self.should_exit = asyncio.Event()

python/kserve/kserve/protocol/rest/server.py

@@ -63,6 +63,9 @@ def timing(self, metric_name, timing, tags):
         trace_logger.info(f"{metric_name}: {timing} {tags}")
 
 
+VALID_UVICORN_LOOPS = {"auto", "asyncio", "uvloop"}
+
+
 class RESTServer:
     def __init__(
         self,
@@ -73,10 +76,18 @@ def __init__(
         access_log_format: Optional[str] = None,
         workers: int = 1,
         grace_period: int = 30,
+        event_loop: str = "auto",
     ):
         self.dataplane = data_plane
         self.model_repository_extension = model_repository_extension
         self.access_log_format = access_log_format
+        if event_loop not in VALID_UVICORN_LOOPS:
+            logger.error(
+                f"Invalid event loop specified: '{event_loop}'. "
+                f"Supported values are {VALID_UVICORN_LOOPS}, "
+                f"Falling back to 'auto'."
+            )
+            event_loop = "auto"
         self.config = uvicorn.Config(
             app,
             host="0.0.0.0",
@@ -86,7 +97,7 @@ def __init__(
             # configured by kserve.
             log_config=None,
             timeout_graceful_shutdown=grace_period,
-            loop="asyncio",
+            loop=event_loop,
         )
         self._server = uvicorn.Server(self.config)
 

python/kserve/test/test_rest_server.py

@@ -0,0 +1,44 @@
+# Copyright 2023 The KServe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import Mock
+
+import pytest
+
+from kserve.protocol.rest import server as rest_mod
+
+
+@pytest.mark.parametrize(
+    "loop_value,expected",
+    [
+        ("auto", "auto"),
+        ("asyncio", "asyncio"),
+        ("uvloop", "uvloop"),
+        ("invalid-value", "auto"),  # invalid falls back to 'auto'
+    ],
+)
+def test_config_loop_value(loop_value, expected, monkeypatch):
+    monkeypatch.setattr(rest_mod.RESTServer, "create_application", lambda self: None)
+    data_plane = Mock()
+    model_repo_ext = Mock()
+
+    rs = rest_mod.RESTServer(
+        app="dummy:app",
+        data_plane=data_plane,
+        model_repository_extension=model_repo_ext,
+        http_port=8080,
+        event_loop=loop_value,
+    )
+
+    assert rs.config.loop == expected

python/kserve/test/test_server.py

@@ -20,6 +20,7 @@
 import re
 from typing import Dict
 from unittest import mock
+from unittest.mock import Mock
 
 import avro.io
 import avro.schema
@@ -33,7 +34,7 @@
 from fastapi.testclient import TestClient
 from ray import serve
 
-from kserve import Model, ModelRepository, ModelServer
+from kserve import Model, ModelRepository, ModelServer, model_server
 from kserve.constants.constants import (
     FASTAPI_APP_IMPORT_STRING,
     INFERENCE_CONTENT_LENGTH_HEADER,
@@ -1412,3 +1413,64 @@ def test_model_server_backwards_compatibility(self):
 
         # Context should handle None predictor config gracefully
         # The DataPlane should still be functional even without predictor config
+
+
+@pytest.mark.asyncio
+class TestModelServerEventLoopWithDummy:
+    async def test_start_passes_event_loop_to_rest_server(self, monkeypatch):
+        created = {}
+
+        class DummyREST:
+            def __init__(self, *args, **kwargs):
+                created["instance"] = self
+                self.event_loop = kwargs.get("event_loop")
+
+            def start(self):
+                return None
+
+        # Patch to capture constructor and avoid side effects
+        monkeypatch.setattr(model_server, "RESTServer", DummyREST)
+        monkeypatch.setattr(model_server.asyncio, "run", Mock())
+        monkeypatch.setattr(
+            model_server.ModelServer, "setup_event_loop", lambda _: None
+        )
+        monkeypatch.setattr(
+            model_server.ModelServer, "register_signal_handler", lambda _: None
+        )
+
+        ms = model_server.ModelServer(workers=1, event_loop="uvloop")
+        m = DummyModel("TestModel")
+        m.load()
+        ms.start(models=[m])
+
+        assert isinstance(created.get("instance"), DummyREST)
+        assert created["instance"].event_loop == "uvloop"
+
+    async def test_start_passes_event_loop_to_rest_multiprocess(self, monkeypatch):
+        created = {}
+
+        class DummyMulti:
+            def __init__(self, *args, **kwargs):
+                created["instance"] = self
+                self.event_loop = kwargs.get("event_loop")
+
+            def start(self):
+                return None
+
+        # Patch multiprocess REST server and side effects
+        monkeypatch.setattr(model_server, "RESTServerMultiProcess", DummyMulti)
+        monkeypatch.setattr(model_server.asyncio, "run", Mock())
+        monkeypatch.setattr(
+            model_server.ModelServer, "setup_event_loop", lambda _: None
+        )
+        monkeypatch.setattr(
+            model_server.ModelServer, "register_signal_handler", lambda _: None
+        )
+
+        ms = model_server.ModelServer(workers=4, event_loop="asyncio")
+        m = DummyModel("TestModel")
+        m.load()
+        ms.start(models=[m])
+
+        assert isinstance(created.get("instance"), DummyMulti)
+        assert created["instance"].event_loop == "asyncio"