Repository: caikit/caikit
Pull request #466 (2023-09-11 22:08:24)

GrpcServerOptions: Add runtime.grpc.options to support server options
Closes: #465 

**What this PR does / why we need it**:

This PR adds a new config section `runtime.grpc.options` that allows arbitrary key/value overrides with `grpc` server options.

**If applicable**:
- [ ] this PR contains documentation
- [X] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/config/config.yml b/caikit/config/config.yml
index c7954cf6..28862e77 100644
--- a/caikit/config/config.yml
+++ b/caikit/config/config.yml
@@ -111,6 +111,9 @@ runtime:
         server_shutdown_grace_period_seconds: 45
         # Listen on a unix socket (model mesh feature)
         unix_socket_path: /tmp/mmesh/grpc.sock
+        # Additional server options as key/value pairs
+        # CITE: https://github.com/grpc/grpc/blob/master/include/grpc/impl/channel_arg_names.h#L22
+        options: {}
 
     # Configuration for the http server
     http:
diff --git a/caikit/runtime/grpc_server.py b/caikit/runtime/grpc_server.py
index a913ab8d..b7b38af9 100644
--- a/caikit/runtime/grpc_server.py
+++ b/caikit/runtime/grpc_server.py
@@ -72,6 +72,7 @@ class RuntimeGRPCServer(RuntimeServerBase):
                 max_workers=self.config.runtime.grpc.server_thread_pool_size
             ),
             interceptors=(PROMETHEUS_METRICS_INTERCEPTOR,),
+            options=(self.config.runtime.grpc.options or {}).items(),
         )
 
         # Start metrics server


Test patch:
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index 2c6303ef..c11af4e4 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -1214,6 +1214,45 @@ def test_grpc_sever_shutdown_with_model_poll(open_port):
         assert not server_proc.killed
 
 
+def test_construct_with_options(open_port, sample_train_service, sample_int_file):
+    """Make sure that the server can be booted with config options"""
+    with temp_config(
+        {
+            "runtime": {
+                "grpc": {
+                    "port": open_port,
+                    "options": {
+                        "grpc.max_receive_message_length": 1,
+                    },
+                }
+            }
+        },
+        "merge",
+    ):
+        server = RuntimeGRPCServer()
+        with server:
+            # Send a request with a payload that's too big and make sure it gets
+            # rejected
+            stream_type = caikit.interfaces.common.data_model.DataStreamSourceInt
+            training_data = stream_type(
+                file=stream_type.File(filename=sample_int_file)
+            ).to_proto()
+            train_request = sample_train_service.messages.OtherTaskOtherModuleTrainRequest(
+                model_name="Bar Training",
+                parameters=sample_train_service.messages.OtherTaskOtherModuleTrainParameters(
+                    sample_input_sampleinputtype=SampleInputType(
+                        name="Gabe"
+                    ).to_proto(),
+                    batch_size=100,
+                    training_data=training_data,
+                ),
+            )
+            train_stub = sample_train_service.stub_class(server.make_local_channel())
+            with pytest.raises(grpc.RpcError) as context:
+                train_stub.OtherTaskOtherModuleTrain(train_request)
+            assert context.value.code() == grpc.StatusCode.RESOURCE_EXHAUSTED
+
+
 # Test implementation details #########################
 @dataclass
 class KeyPair: