feat(ai.triton.server): added an option to control the max. message s…

…ize limit for the GRPC calls (#4035) * feat: added an option to control the max. message size limit for the GRPC calls to the Triton server Signed-off-by: Marcello Martina <[email protected]> * feat: added default value in metatype for grpc.max.size Signed-off-by: Marcello Martina <[email protected]>
eclipse-kura · Jun 29, 2022 · 5a218b3 · 5a218b3
1 parent 7eb5916
commit 5a218b3
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 1 deletion.
diff --git a/...triton.server/OSGI-INF/metatype/org.eclipse.kura.ai.triton.server.TritonServerService.xml b/...triton.server/OSGI-INF/metatype/org.eclipse.kura.ai.triton.server.TritonServerService.xml
@@ -102,6 +102,17 @@
             default="3" 
             description="Timeout (in seconds) for time consuming tasks like server startup, shutdown or model load. If the task exceeds the timeout, the operation will be terminated with an error.">
         </AD>
+
+        <AD id="grpc.max.size"
+        	name="Max. GRPC message size (bytes)"
+        	type="Integer"
+        	description="Maximum accepted input size for the GRPC calls.
+        	Increase this value if the model input size is bigger than the default."
+        	cardinality="0"
+        	required="true"
+        	default="4194304"
+        	min="1">
+        </AD>        	
 
     </OCD>
     <Designate factoryPid="org.eclipse.kura.ai.triton.server.TritonServerService">

diff --git a/...riton.server/src/main/java/org/eclipse/kura/ai/triton/server/TritonServerServiceImpl.java b/...riton.server/src/main/java/org/eclipse/kura/ai/triton/server/TritonServerServiceImpl.java
@@ -169,7 +169,7 @@ private void stopLocalInstance() {
 
     private void setGrpcResources() {
         this.grpcChannel = ManagedChannelBuilder.forAddress(this.options.getAddress(), this.options.getGrpcPort())
-                .usePlaintext().build();
+                .usePlaintext().maxInboundMessageSize(this.options.getGrpcMaxMessageSize()).build();
         setGrpcStub(GRPCInferenceServiceGrpc.newBlockingStub(this.grpcChannel));
     }
 

diff --git a/...on.server/src/main/java/org/eclipse/kura/ai/triton/server/TritonServerServiceOptions.java b/...on.server/src/main/java/org/eclipse/kura/ai/triton/server/TritonServerServiceOptions.java
@@ -36,16 +36,19 @@ public class TritonServerServiceOptions {
     private static final String PROPERTY_MODELS = "models";
     private static final String PROPERTY_LOCAL = "enable.local";
     private static final String PROPERTY_TIMEOUT = "timeout";
+    private static final String PROPERTY_MAX_GRPC_MESSAGE_SIZE = "grpc.max.size";
     private final Map<String, Object> properties;
 
     private static final int RETRY_INTERVAL = 500; // ms
+    private static final int DEFAULT_MAX_GRPC_MESSAGE_SIZE = 4194304; // bytes
 
     private final int httpPort;
     private final int grpcPort;
     private final int metricsPort;
     private final boolean isLocal;
     private final int timeout;
     private final int nRetries;
+    private final int grpcMaxMessageSize;
 
     public TritonServerServiceOptions(final Map<String, Object> properties) {
         requireNonNull(properties, "Properties cannot be null");
@@ -86,6 +89,13 @@ public TritonServerServiceOptions(final Map<String, Object> properties) {
             this.timeout = 3;
         }
         this.nRetries = (this.timeout * 1000) / RETRY_INTERVAL;
+
+        final Object propertyGrpcMaxMessageSize = properties.get(PROPERTY_MAX_GRPC_MESSAGE_SIZE);
+        if (propertyGrpcMaxMessageSize instanceof Integer) {
+            this.grpcMaxMessageSize = (int) propertyGrpcMaxMessageSize;
+        } else {
+            this.grpcMaxMessageSize = DEFAULT_MAX_GRPC_MESSAGE_SIZE;
+        }
     }
 
     public String getAddress() {
@@ -163,6 +173,10 @@ public int getRetryInterval() {
         return RETRY_INTERVAL;
     }
 
+    public int getGrpcMaxMessageSize() {
+        return this.grpcMaxMessageSize;
+    }
+
     private String getStringProperty(String propertyName) {
         String stringProperty = "";
         final Object stringPropertyObj = this.properties.get(propertyName);

diff --git a/....test/src/test/java/org/eclipse/kura/ai/triton/server/TritonServerServiceOptionsTest.java b/....test/src/test/java/org/eclipse/kura/ai/triton/server/TritonServerServiceOptionsTest.java
@@ -19,6 +19,7 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import org.eclipse.kura.ai.triton.server.TritonServerServiceOptions;
 import org.junit.Test;
 
 public class TritonServerServiceOptionsTest {
@@ -205,6 +206,28 @@ public void hashCodeMethodShouldWorkWithOptionsBuiltWithDifferentProperties() {
         thenHashCodesShouldNotMatch();
     }
 
+    @Test
+    public void shouldReturnInputGrpcMaxMessageSize() {
+        givenPropertyWith("server.address", "localhost");
+        givenPropertyWith("server.ports", new Integer[] { 4000, 4001, 4002 });
+        givenPropertyWith("enable.local", Boolean.FALSE);
+        givenPropertyWith("grpc.max.size", 2000);
+        givenServiceOptionsBuiltWith(this.properties);
+
+        thenGrpcMaxMessageSizeIsEqualTo(2000);
+    }
+
+    @Test
+    public void shouldReturnDefaultGrpcMaxMessageSize() {
+        givenPropertyWith("server.address", "localhost");
+        givenPropertyWith("server.ports", new Integer[] { 4000, 4001, 4002 });
+        givenPropertyWith("enable.local", Boolean.FALSE);
+        givenPropertyWith("grpc.max.size", null);
+        givenServiceOptionsBuiltWith(this.properties);
+
+        thenGrpcMaxMessageSizeIsEqualTo(4194304);
+    }
+
     /*
      * Given
      */
@@ -266,6 +289,10 @@ private void thenLocalConfigIsEqualTo(boolean value) {
         assertEquals(value, this.options.isLocalEnabled());
     }
 
+    private void thenGrpcMaxMessageSizeIsEqualTo(int expectedValue) {
+        assertEquals(expectedValue, this.options.getGrpcMaxMessageSize());
+    }
+
     private void thenEqualsMethodShouldReturn(boolean value) {
         assertEquals(value, this.equalsResult);
     }
@@ -277,4 +304,5 @@ private void thenHashCodesShouldMatch() {
     private void thenHashCodesShouldNotMatch() {
         assertNotEquals(this.hashCode, this.otherHashCode);
     }
+
 }