🥅 Catch InactiveError for TGIS to modify error message

Signed-off-by: gkumbhat <[email protected]>
caikit · Aug 2, 2024 · 285ae7c · 285ae7c
1 parent 1e5aac6
commit 285ae7c
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/caikit_nlp/toolkit/text_generation/tgis_utils.py b/caikit_nlp/toolkit/text_generation/tgis_utils.py
@@ -474,6 +474,13 @@ 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)
+                error_message = "The underlying TCP connection is closed"
+                caikit_status_code = GRPC_TO_CAIKIT_CORE_STATUS.get(
+                    err.code(), CaikitCoreStatusCode.UNKNOWN
+                )
+                raise CaikitCoreException(caikit_status_code, error_message)
             except grpc.RpcError as err:
                 raise_caikit_core_exception(err)
 
@@ -653,6 +660,13 @@ def stream_generate(
                     input_tokens=input_token_list,
                     details=details,
                 )
+        except grpc._channel._InactiveRpcError as err:
+            log.error("<NLP11829118E>", err.details)
+            error_message = "The underlying TCP connection is closed"
+            caikit_status_code = GRPC_TO_CAIKIT_CORE_STATUS.get(
+                err.code(), CaikitCoreStatusCode.UNKNOWN
+            )
+            raise CaikitCoreException(caikit_status_code, error_message)
         except grpc.RpcError as err:
             raise_caikit_core_exception(err)