🥅 Catch InactiveError for TGIS to modify error message

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>
caikit · Aug 2, 2024 · 23dc8e6 · 23dc8e6
1 parent 1e5aac6
commit 23dc8e6
Showing 1 changed file with 8 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,10 @@ def unary_generate(
                 batch_response = self.tgis_client.Generate(
                     request, timeout=self.tgis_req_timeout
                 )
+            except grpc._channel._InactiveRpcError as err:
+                log.error("<NLP30829218E>", str(err))
+                error_message = "The underlying TCP connection is closed"
+                raise CaikitCoreException(CaikitCoreStatusCode.INTERNAL, error_message)
             except grpc.RpcError as err:
                 raise_caikit_core_exception(err)
 
@@ -653,6 +657,10 @@ def stream_generate(
                     input_tokens=input_token_list,
                     details=details,
                 )
+        except grpc._channel._InactiveRpcError as err:
+            log.error("<NLP11829118E>", str(err))
+            error_message = "The underlying TCP connection is closed"
+            raise CaikitCoreException(CaikitCoreStatusCode.INTERNAL, error_message)
         except grpc.RpcError as err:
             raise_caikit_core_exception(err)