Repository: caikit/caikit
Pull request #411 (2023-08-08 21:43:07)

Add more exception codes
<!--  Thanks for sending a pull request!  Here are some tips for you:
1. Make sure to read the Contributing Guide before submitting your PR: https://github.com/caikit/caikit/blob/main/CONTRIBUTING.md
2. If this PR closes another issue, add 'closes #<issue number>' somewhere in the PR summary. GitHub will automatically close that issue when this PR gets merged. Alternatively, adding 'refs #<issue number>' will not close the issue, but help provide the reviewer more context.-->

**What this PR does / why we need it**:

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [ ] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/core/exceptions/caikit_core_exception.py b/caikit/core/exceptions/caikit_core_exception.py
index 9e662b7a..1af8d705 100644
--- a/caikit/core/exceptions/caikit_core_exception.py
+++ b/caikit/core/exceptions/caikit_core_exception.py
@@ -25,12 +25,14 @@ class CaikitCoreStatusCode(Enum):
     CONNECTION_ERROR = 3
     UNAUTHORIZED = 4
     FORBIDDEN = 5
+    UNKNOWN = 6
+    FATAL = 7
 
 
 class CaikitCoreException(Exception):
-    StatusCode: CaikitCoreStatusCode
-    Message: str
+    status_code: CaikitCoreStatusCode
+    message: str
 
-    def __init__(self, status_code: CaikitCoreStatusCode, msg: str) -> None:
-        self.StatusCode = status_code
-        self.Message = msg
+    def __init__(self, status_code: CaikitCoreStatusCode, message: str) -> None:
+        self.status_code = status_code
+        self.message = message
diff --git a/caikit/core/model_management/local_model_trainer.py b/caikit/core/model_management/local_model_trainer.py
index c215e47d..e1348c6b 100644
--- a/caikit/core/model_management/local_model_trainer.py
+++ b/caikit/core/model_management/local_model_trainer.py
@@ -324,7 +324,7 @@ class LocalModelTrainer(ModelTrainerBase):
             return model_future
         raise CaikitCoreException(
             status_code=CaikitCoreStatusCode.NOT_FOUND,
-            msg=f"Unknown training_id: {training_id}",
+            message=f"Unknown training_id: {training_id}",
         )
 
     ## Impl ##
diff --git a/caikit/runtime/servicers/training_management_servicer.py b/caikit/runtime/servicers/training_management_servicer.py
index 2f33d3d1..412bb56b 100644
--- a/caikit/runtime/servicers/training_management_servicer.py
+++ b/caikit/runtime/servicers/training_management_servicer.py
@@ -95,7 +95,7 @@ class TrainingManagementServicerImpl:
             # This is to handle stateful trainers that implement `cancel` by fully deleting
             # the training. NB: Future `GetTrainingStatus` calls for these canceled trainings
             # would raise a not found error to the user.
-            if err.StatusCode == CaikitCoreStatusCode.NOT_FOUND:
+            if err.status_code == CaikitCoreStatusCode.NOT_FOUND:
                 return TrainingStatusResponse(
                     training_id=training_info_request.training_id,
                     state=TrainingStatus.CANCELED,
diff --git a/caikit/runtime/utils/servicer_util.py b/caikit/runtime/utils/servicer_util.py
index bf3deb30..f12a66e1 100644
--- a/caikit/runtime/utils/servicer_util.py
+++ b/caikit/runtime/utils/servicer_util.py
@@ -458,44 +458,58 @@ def build_caikit_library_request_dict(
 
 
 def raise_caikit_runtime_exception(exception: CaikitCoreException):
-    if exception.StatusCode == CaikitCoreStatusCode.NOT_FOUND:
-        log.debug2("Not Found Error: [%s]", exception.Message, exc_info=True)
+    if exception.status_code == CaikitCoreStatusCode.NOT_FOUND:
+        log.debug2("Not Found Error: [%s]", exception.message, exc_info=True)
         raise CaikitRuntimeException(
             grpc.StatusCode.NOT_FOUND,
-            exception.Message,
+            exception.message,
         ) from exception
 
-    if exception.StatusCode == CaikitCoreStatusCode.INVALID_ARGUMENT:
-        log.debug2("Invalid Argument Error: [%s]", exception.Message, exc_info=True)
+    if exception.status_code == CaikitCoreStatusCode.INVALID_ARGUMENT:
+        log.debug2("Invalid Argument Error: [%s]", exception.message, exc_info=True)
         raise CaikitRuntimeException(
             grpc.StatusCode.INVALID_ARGUMENT,
-            exception.Message,
+            exception.message,
         ) from exception
 
-    if exception.StatusCode == CaikitCoreStatusCode.FORBIDDEN:
-        log.debug2("Forbidden Error: [%s]", exception.Message, exc_info=True)
+    if exception.status_code == CaikitCoreStatusCode.FORBIDDEN:
+        log.debug2("Forbidden Error: [%s]", exception.message, exc_info=True)
         raise CaikitRuntimeException(
             grpc.StatusCode.PERMISSION_DENIED,
-            exception.Message,
+            exception.message,
         ) from exception
 
-    if exception.StatusCode == CaikitCoreStatusCode.CONNECTION_ERROR:
-        log.debug2("Connection Error: [%s]", exception.Message, exc_info=True)
+    if exception.status_code == CaikitCoreStatusCode.CONNECTION_ERROR:
+        log.debug2("Connection Error: [%s]", exception.message, exc_info=True)
         raise CaikitRuntimeException(
             grpc.StatusCode.UNAVAILABLE,
-            exception.Message,
+            exception.message,
         ) from exception
 
-    if exception.StatusCode == CaikitCoreStatusCode.UNAUTHORIZED:
+    if exception.status_code == CaikitCoreStatusCode.UNAUTHORIZED:
         log.debug2("Unauthorized Error: [%s]", exception.Message, exc_info=True)
         raise CaikitRuntimeException(
             grpc.StatusCode.UNAUTHENTICATED,
-            exception.Message,
+            exception.message,
+        ) from exception
+
+    if exception.status_code == CaikitCoreStatusCode.FATAL:
+        log.debug2("Fatal Error: [%s]", exception.message, exc_info=True)
+        raise CaikitRuntimeException(
+            grpc.StatusCode.INTERNAL,
+            exception.message,
+        ) from exception
+
+    if exception.status_code == CaikitCoreStatusCode.UNKNOWN:
+        log.debug2("Unknown Error: [%s]", exception.message, exc_info=True)
+        raise CaikitRuntimeException(
+            grpc.StatusCode.UNKNOWN,
+            exception.message,
         ) from exception
 
     log.debug2(
-        "Unknown CaikitCoreException Error: [%s]", exception.Message, exc_info=True
+        "Unhandled CaikitCoreException Error: [%s]", exception.message, exc_info=True
     )
     raise CaikitRuntimeException(
-        grpc.StatusCode.UNKNOWN, exception.Message
+        grpc.StatusCode.INTERNAL, exception.message
     ) from exception


Test patch:
diff --git a/tests/runtime/servicers/test_training_management_servicer.py b/tests/runtime/servicers/test_training_management_servicer.py
index 9b67d3a1..7f36510b 100644
--- a/tests/runtime/servicers/test_training_management_servicer.py
+++ b/tests/runtime/servicers/test_training_management_servicer.py
@@ -64,7 +64,7 @@ class MockModelFuture:
             return TrainingInfo(status=TrainingStatus.RUNNING)
         else:
             raise CaikitCoreException(
-                status_code=CaikitCoreStatusCode.NOT_FOUND, msg="Training not found"
+                status_code=CaikitCoreStatusCode.NOT_FOUND, message="Training not found"
             )
 
     def cancel(self):