From 7b1b998e0ab825f736971cf144254f2238be6964 Mon Sep 17 00:00:00 2001
From: Robert Steiner <robertt.debug@gmail.com>
Date: Fri, 21 Feb 2020 10:16:24 +0100
Subject: [PATCH] PB-439 Make thread pool workers configurable

---
 configs/example-config.toml | 2 ++
 tests/test_config.py        | 3 ++-
 xain_fl/config/schema.py    | 3 +++
 xain_fl/serve.py            | 3 ++-
 4 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/configs/example-config.toml b/configs/example-config.toml
index e696bc683..d0ff1d6a6 100644
--- a/configs/example-config.toml
+++ b/configs/example-config.toml
@@ -7,6 +7,8 @@
 host = "localhost"
 # (Optional) Port to listen on for incoming gRPC connections
 port = 50051
+# (Optional) The maximum number of gRPC thread pool workers
+thread_pool_workers = 10
 
 # (Optional)
 [server.grpc_options]
diff --git a/tests/test_config.py b/tests/test_config.py
index ddb75ae45..42150df78 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -20,6 +20,7 @@ def server_sample():
             "grpc.max_receive_message_length": -1,
             "grpc.max_send_message_length": -1,
         },
+        "thread_pool_workers": 11,
     }
 
 
@@ -122,11 +123,11 @@ def test_load_valid_config(config_sample):  # pylint: disable=redefined-outer-na
 
     assert config.server.host == "localhost"
     assert config.server.port == 50051
-
     assert config.server.grpc_options == [
         ("grpc.max_receive_message_length", -1),
         ("grpc.max_send_message_length", -1),
     ]
+    assert config.server.thread_pool_workers == 11
 
     assert config.ai.rounds == 1
     assert config.ai.epochs == 1
diff --git a/xain_fl/config/schema.py b/xain_fl/config/schema.py
index e3c5ae2cf..bec3e578d 100644
--- a/xain_fl/config/schema.py
+++ b/xain_fl/config/schema.py
@@ -179,6 +179,9 @@ def log_level(key: str) -> Schema:
             lambda opt: list(opt.items()),
             error=error("server.grpc_options", "valid gRPC options"),
         ),
+        Optional("thread_pool_workers", default=10): positive_integer(
+            "server.thread_pool_workers"
+        ),
     }
 )
 
diff --git a/xain_fl/serve.py b/xain_fl/serve.py
index e4867a9e9..45f0d2ba1 100644
--- a/xain_fl/serve.py
+++ b/xain_fl/serve.py
@@ -31,7 +31,8 @@ def serve(coordinator: Coordinator, server_config: ServerConfig) -> None:
 
     """
     server = grpc.server(
-        futures.ThreadPoolExecutor(max_workers=10), options=server_config.grpc_options
+        futures.ThreadPoolExecutor(max_workers=server_config.thread_pool_workers),
+        options=server_config.grpc_options,
     )
     coordinator_pb2_grpc.add_CoordinatorServicer_to_server(
         CoordinatorGrpc(coordinator), server