Merge pull request #372 from gkumbhat/add_tgis_inactive_error

🥅 Catch InactiveError for TGIS to modify error message
caikit · Aug 2, 2024 · 79ac58e · 79ac58e
2 parents 1e5aac6 + d6adb79
commit 79ac58e
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/caikit_nlp/toolkit/text_generation/tgis_utils.py b/caikit_nlp/toolkit/text_generation/tgis_utils.py
@@ -91,6 +91,7 @@
 # HTTP Header / gRPC Metadata key used to identify a route override
 # (forwarded for API compatibility)
 ROUTE_INFO_HEADER_KEY = TGISBackend.ROUTE_INFO_HEADER_KEY
+INACTIVE_RPC_CONN_ERR_MESSAGE = "The underlying TCP connection is closed"
 get_route_info = TGISBackend.get_route_info
 
 
@@ -474,6 +475,14 @@ def unary_generate(
                 batch_response = self.tgis_client.Generate(
                     request, timeout=self.tgis_req_timeout
                 )
+            except grpc._channel._InactiveRpcError as err:
+                log.error("<NLP30829218E>", err.details)
+                caikit_status_code = GRPC_TO_CAIKIT_CORE_STATUS.get(
+                    err.code(), CaikitCoreStatusCode.UNKNOWN
+                )
+                raise CaikitCoreException(
+                    caikit_status_code, INACTIVE_RPC_CONN_ERR_MESSAGE
+                ) from err
             except grpc.RpcError as err:
                 raise_caikit_core_exception(err)
 
@@ -653,6 +662,14 @@ def stream_generate(
                     input_tokens=input_token_list,
                     details=details,
                 )
+        except grpc._channel._InactiveRpcError as err:
+            log.error("<NLP11829118E>", err.details)
+            caikit_status_code = GRPC_TO_CAIKIT_CORE_STATUS.get(
+                err.code(), CaikitCoreStatusCode.UNKNOWN
+            )
+            raise CaikitCoreException(
+                caikit_status_code, INACTIVE_RPC_CONN_ERR_MESSAGE
+            ) from err
         except grpc.RpcError as err:
             raise_caikit_core_exception(err)