Merge pull request pinecone-io#135 from pinecone-io/grpc_retries

Use built-in gRPC retry configuration

Author: igiloh-pinecone

pinecone/core/grpc/index_grpc.py

@@ -13,6 +13,7 @@
 from google.protobuf import json_format
 from grpc._channel import _InactiveRpcError, _MultiThreadedRendezvous
 from tqdm import tqdm
+import json
 
 from pinecone import FetchResponse, QueryResponse, ScoredVector, SingleQueryResults, DescribeIndexStatsResponse
 from pinecone.config import Config
@@ -84,9 +85,38 @@ def __init__(self, index_name: str, channel=None, grpc_config: GRPCClientConfig
             "client-version": CLIENT_VERSION
         }
         self._endpoint_override = _endpoint_override
+
+        self.method_config = json.dumps(
+            {
+                "methodConfig": [
+                    {
+                        "name": [{"service": "VectorService.Upsert"}],
+                        "retryPolicy": {
+                            "maxAttempts": 5,
+                            "initialBackoff": "0.1s",
+                            "maxBackoff": "1s",
+                            "backoffMultiplier": 2,
+                            "retryableStatusCodes": ["UNAVAILABLE"],
+                        },
+                    },
+                    {
+                        "name": [{"service": "VectorService"}],
+                        "retryPolicy": {
+                            "maxAttempts": 5,
+                            "initialBackoff": "0.1s",
+                            "maxBackoff": "1s",
+                            "backoffMultiplier": 2,
+                            "retryableStatusCodes": ["UNAVAILABLE"],
+                        },
+                    }
+                ]
+            }
+        )
+
         self._channel = channel or self._gen_channel()
         self.stub = self.stub_class(self._channel)
 
+
     @property
     @abstractmethod
     def stub_class(self):
@@ -100,7 +130,9 @@ def _gen_channel(self, options=None):
         target = self._endpoint()
         default_options = {
             "grpc.max_send_message_length": MAX_MSG_SIZE,
-            "grpc.max_receive_message_length": MAX_MSG_SIZE
+            "grpc.max_receive_message_length": MAX_MSG_SIZE,
+            "grpc.service_config": self.method_config,
+            "grpc.enable_retries": True
         }
         if self.grpc_client_config.secure:
             default_options['grpc.ssl_target_name_override'] = target.split(':')[0]
@@ -114,8 +146,8 @@ def _gen_channel(self, options=None):
             root_cas = open(certifi.where(), "rb").read()
             tls = grpc.ssl_channel_credentials(root_certificates=root_cas)
             channel = grpc.secure_channel(target, tls, options=_options)
-        interceptor = RetryOnRpcErrorClientInterceptor(self.retry_config)
-        return grpc.intercept_channel(channel, interceptor)
+
+        return channel
 
     @property
     def channel(self):
@@ -246,7 +278,7 @@ def add_done_callback(self, fun):
 
     def result(self, timeout=None):
         try:
-            self._delegate.result(timeout=timeout)
+            return self._delegate.result(timeout=timeout)
         except _MultiThreadedRendezvous as e:
             raise PineconeException(e._state.debug_error_string) from e