Repository: caikit/caikit (python)

Issue:
Allow uvicorn config overrides
## Is your feature request related to a problem? Please describe.

The `uvicorn` server is [highly configurable](https://github.com/encode/uvicorn/blob/master/uvicorn/config.py#L188), but currently `caikit.runtime.http_server` allows [a small subset of those configs](https://github.com/caikit/caikit/blob/main/caikit/runtime/http_server/http_server.py#L228) to be set. This feature request is to allow additional configurations to be passed through from the config ecosystem to the `uvicorn.Config` setup for more customized server setups.

## Describe the solution you'd like

* Add a new `config` in `runtime.http.server_config` that takes an arbitrary map
* Use `runtime.http.server_config` when instantiating the `uvicorn.Config` in the HTTP server


Pull request:
Uvicorn config 609
**What this PR does / why we need it**:

Closes #609 

This PR makes it possible to set arbitrary `uvicorn` server config overrides as long as they don't overlap with the ones that `caikit` manages explicitly.

**If applicable**:
- [ ] this PR contains documentation
- [x] this PR contains unit tests
- [x] this PR has been tested for backwards compatibility


Patch:
diff --git a/caikit/config/config.yml b/caikit/config/config.yml
index 45284a57..008a81e2 100644
--- a/caikit/config/config.yml
+++ b/caikit/config/config.yml
@@ -153,6 +153,9 @@ runtime:
         server_shutdown_grace_period_seconds: 5
         # Timeout for health probe to receive a response
         probe_timeout: 0.01
+        # Additional uvicorn server configuration
+        # CITE: https://github.com/encode/uvicorn/blob/master/uvicorn/config.py#L188
+        server_config: {}
 
     # Configuration for the metrics server
     metrics:
diff --git a/caikit/runtime/http_server/http_server.py b/caikit/runtime/http_server/http_server.py
index 44e22f2d..4d723466 100644
--- a/caikit/runtime/http_server/http_server.py
+++ b/caikit/runtime/http_server/http_server.py
@@ -62,6 +62,7 @@ from .utils import convert_json_schema_to_multipart, flatten_json_schema
 from caikit.config import get_config
 from caikit.core.data_model import DataBase
 from caikit.core.data_model.dataobject import make_dataobject
+from caikit.core.exceptions import error_handler
 from caikit.core.toolkit.sync_to_async import async_wrap_iter
 from caikit.runtime.server_base import RuntimeServerBase
 from caikit.runtime.service_factory import ServicePackage
@@ -78,6 +79,7 @@ from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 ## Globals #####################################################################
 
 log = alog.use_channel("SERVR-HTTP")
+error = error_handler.get(log)
 
 
 # Mapping from GRPC codes to their corresponding HTTP codes
@@ -221,18 +223,37 @@ class RuntimeHTTPServer(RuntimeServerBase):
                 log.info("<RUN10539515I>", "Running INSECURE")
 
             # Start the server with a timeout_graceful_shutdown
-            # if not set in config, this is None and unvicorn accepts None or number of seconds
+            # if not set in config, this is None and unvicorn accepts None or
+            # number of seconds
             unvicorn_timeout_graceful_shutdown = (
                 get_config().runtime.http.server_shutdown_grace_period_seconds
             )
+            server_config = get_config().runtime.http.server_config
+            overlapping_tls_config = set(tls_kwargs).intersection(server_config)
+            error.value_check(
+                "<RUN30233180E>",
+                not overlapping_tls_config,
+                "Found overlapping config keys between TLS and server_config: %s",
+                overlapping_tls_config,
+            )
+            config_kwargs = {
+                "port": self.port,
+                "log_level": None,
+                "log_config": None,
+                "timeout_graceful_shutdown": unvicorn_timeout_graceful_shutdown,
+            }
+            overlapping_kwarg_config = set(config_kwargs).intersection(server_config)
+            error.value_check(
+                "<RUN99488934E>",
+                not overlapping_kwarg_config,
+                "Found caikit-managed uvicorn config in server_config: %s",
+                overlapping_kwarg_config,
+            )
             config = uvicorn.Config(
                 self.app,
-                host="0.0.0.0",
-                port=self.port,
-                log_level=None,
-                log_config=None,
-                timeout_graceful_shutdown=unvicorn_timeout_graceful_shutdown,
+                **config_kwargs,
                 **tls_kwargs,
+                **server_config,
             )
             # Make sure the config loads TLS files here so they can safely be
             # deleted if they're ephemeral


Test patch:
diff --git a/tests/runtime/http_server/test_http_server.py b/tests/runtime/http_server/test_http_server.py
index 0b78d86b..040c1ee7 100644
--- a/tests/runtime/http_server/test_http_server.py
+++ b/tests/runtime/http_server/test_http_server.py
@@ -1081,3 +1081,59 @@ def test_http_and_grpc_server_share_threadpool(
     runtime_http_server, runtime_grpc_server
 ):
     assert runtime_grpc_server.thread_pool is runtime_http_server.thread_pool
+
+
+def test_uvicorn_server_config_valid():
+    """Make sure that arbitrary uvicorn configs can be passed through from
+    runtime.http.server_config
+    """
+    timeout_keep_alive = 10
+    with temp_config(
+        {
+            "runtime": {
+                "http": {"server_config": {"timeout_keep_alive": timeout_keep_alive}}
+            }
+        },
+        "merge",
+    ):
+        server = http_server.RuntimeHTTPServer()
+        assert server.server.config.timeout_keep_alive == timeout_keep_alive
+
+
+def test_uvicorn_server_config_invalid_tls_overlap():
+    """Make sure uvicorn TLS arguments cannot be set if TLS is enabled in caikit
+    config
+    """
+    with temp_config(
+        {
+            "runtime": {
+                "http": {
+                    "server_config": {
+                        "ssl_keyfile": "/some/file.pem",
+                    }
+                }
+            }
+        },
+        "merge",
+    ):
+        with generate_tls_configs(port=1234, tls=True, mtls=True):
+            with pytest.raises(ValueError):
+                http_server.RuntimeHTTPServer()
+
+
+def test_uvicorn_server_config_invalid_kwarg_overlap():
+    """Make sure uvicorn config can't be set for configs that caikit manages"""
+    with temp_config(
+        {
+            "runtime": {
+                "http": {
+                    "server_config": {
+                        "log_level": "debug",
+                    }
+                }
+            }
+        },
+        "merge",
+    ):
+        with pytest.raises(ValueError):
+            http_server.RuntimeHTTPServer()