fix(clients): fix message lenght

instill-ai · May 8, 2024 · 50411ef · 50411ef
1 parent bafa292
commit 50411ef
Showing 1 changed file with 18 additions and 4 deletions.
diff --git a/instill/clients/instance.py b/instill/clients/instance.py
@@ -6,30 +6,44 @@
 import instill.protogen.model.model.v1alpha.model_public_service_pb2_grpc as model_service
 import instill.protogen.vdp.pipeline.v1beta.pipeline_public_service_pb2_grpc as pipeline_service
 
+MB = 1024**2
+
 
 class InstillInstance:
     def __init__(self, stub, url: str, token: str, secure: bool, async_enabled: bool):
         self.url: str = url
         self.token: str = token
         self.async_enabled: bool = async_enabled
         self.metadata: Union[str, tuple] = ""
+
+        channel_options = (
+            [
+                ("grpc.max_send_message_length", 10 * MB),
+                ("grpc.max_receive_message_length", 10 * MB),
+            ],
+        )
+
         if not secure:
-            channel = grpc.insecure_channel(url)
+            channel = grpc.insecure_channel(url, options=channel_options)
             self.metadata = (
                 (
                     "authorization",
                     f"Bearer {token}",
                 ),
             )
             if async_enabled:
-                async_channel = grpc.aio.insecure_channel(url)
+                async_channel = grpc.aio.insecure_channel(url, options=channel_options)
         else:
             ssl_creds = grpc.ssl_channel_credentials()
             call_creds = grpc.access_token_call_credentials(token)
             creds = grpc.composite_channel_credentials(ssl_creds, call_creds)
-            channel = grpc.secure_channel(target=url, credentials=creds)
+            channel = grpc.secure_channel(
+                target=url, credentials=creds, options=channel_options
+            )
             if async_enabled:
-                async_channel = grpc.aio.secure_channel(target=url, credentials=creds)
+                async_channel = grpc.aio.secure_channel(
+                    target=url, credentials=creds, options=channel_options
+                )
         self.channel: grpc.Channel = channel
         self.client: Union[
             model_service.ModelPublicServiceStub,