Repository: caikit/caikit
Pull request #393 (2023-08-02 21:31:13)

Add CaikitCoreException 
<!--  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**:
Related to #390 

**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
new file mode 100644
index 00000000..9e662b7a
--- /dev/null
+++ b/caikit/core/exceptions/caikit_core_exception.py
@@ -0,0 +1,36 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Caikit Core Exception enum used for reporting Exception status raised in caikit core
+"""
+
+# Standard
+from enum import Enum
+
+
+class CaikitCoreStatusCode(Enum):
+    NOT_FOUND = 1
+    INVALID_ARGUMENT = 2
+    CONNECTION_ERROR = 3
+    UNAUTHORIZED = 4
+    FORBIDDEN = 5
+
+
+class CaikitCoreException(Exception):
+    StatusCode: CaikitCoreStatusCode
+    Message: str
+
+    def __init__(self, status_code: CaikitCoreStatusCode, msg: str) -> None:
+        self.StatusCode = status_code
+        self.Message = msg
diff --git a/caikit/core/model_management/local_model_trainer.py b/caikit/core/model_management/local_model_trainer.py
index 11345e08..1b9a6a14 100644
--- a/caikit/core/model_management/local_model_trainer.py
+++ b/caikit/core/model_management/local_model_trainer.py
@@ -37,6 +37,10 @@ from ..toolkit.destroyable_process import DestroyableProcess
 from ..toolkit.destroyable_thread import DestroyableThread
 from ..toolkit.logging import configure as configure_logging
 from .model_trainer_base import ModelTrainerBase, TrainingInfo
+from caikit.core.exceptions.caikit_core_exception import (
+    CaikitCoreException,
+    CaikitCoreStatusCode,
+)
 import caikit
 
 log = alog.use_channel("LOC-TRNR")
@@ -315,7 +319,10 @@ class LocalModelTrainer(ModelTrainerBase):
         self._purge_old_futures()
         if model_future := self._futures.get(training_id):
             return model_future
-        raise ValueError(f"Unknown training_id: {training_id}")
+        raise CaikitCoreException(
+            status_code=CaikitCoreStatusCode.NOT_FOUND,
+            msg=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 b1941316..56772a2f 100644
--- a/caikit/runtime/servicers/training_management_servicer.py
+++ b/caikit/runtime/servicers/training_management_servicer.py
@@ -23,11 +23,13 @@ import alog
 
 # Local
 from caikit.core import MODEL_MANAGER
+from caikit.core.exceptions.caikit_core_exception import CaikitCoreException
 from caikit.interfaces.runtime.data_model import (
     TrainingInfoRequest,
     TrainingStatusResponse,
 )
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
+from caikit.runtime.utils.servicer_util import raise_caikit_runtime_exception
 
 log = alog.use_channel("MR-SERVICR-I")
 
@@ -42,7 +44,6 @@ class TrainingManagementServicerImpl:
         model_future = self._get_model_future(
             training_info.training_id, operation="get_status"
         )
-
         try:
             reasons = []
             if model_future.get_info().errors:
@@ -53,6 +54,8 @@ class TrainingManagementServicerImpl:
                 state=model_future.get_info().status,
                 reasons=reasons,
             ).to_proto()
+        except CaikitCoreException as err:
+            raise_caikit_runtime_exception(exception=err)
         except Exception as err:
             raise CaikitRuntimeException(
                 grpc.StatusCode.INTERNAL,
@@ -67,7 +70,6 @@ class TrainingManagementServicerImpl:
         model_future = self._get_model_future(
             training_info.training_id, operation="cancel"
         )
-
         try:
             model_future.cancel()
 
@@ -80,6 +82,8 @@ class TrainingManagementServicerImpl:
                 state=model_future.get_info().status,
                 reasons=reasons,
             ).to_proto()
+        except CaikitCoreException as err:
+            raise_caikit_runtime_exception(exception=err)
         except Exception as err:
             log.debug2(
                 "Unexpected error trying to cancel training id %s: [%s]",
@@ -100,17 +104,8 @@ class TrainingManagementServicerImpl:
         """
         try:
             return MODEL_MANAGER.get_model_future(training_id=training_id)
-        except ValueError as err:
-            log.debug2(
-                "Caught ValueError while trying to look up model future for id %s: [%s]",
-                training_id,
-                err,
-            )
-            raise CaikitRuntimeException(
-                grpc.StatusCode.NOT_FOUND,
-                "{} not found in the list of currently running training jobs. "
-                "Could not perform {}".format(training_id, operation),
-            ) from err
+        except CaikitCoreException as err:
+            raise_caikit_runtime_exception(exception=err)
         except Exception as err:
             log.debug2(
                 "Caught unexpected exception while trying to look up model future for id %s: [%s]",
diff --git a/caikit/runtime/utils/servicer_util.py b/caikit/runtime/utils/servicer_util.py
index a0bc74cd..bf3deb30 100644
--- a/caikit/runtime/utils/servicer_util.py
+++ b/caikit/runtime/utils/servicer_util.py
@@ -30,6 +30,10 @@ import alog
 # Local
 from caikit.core.data_model import DataStream
 from caikit.core.data_model.base import DataBase
+from caikit.core.exceptions.caikit_core_exception import (
+    CaikitCoreException,
+    CaikitCoreStatusCode,
+)
 from caikit.core.signature_parsing import CaikitMethodSignature
 from caikit.interfaces.runtime.data_model.training_management import ModelPointer
 from caikit.runtime.model_management.model_manager import ModelManager
@@ -451,3 +455,47 @@ def build_caikit_library_request_dict(
         raise CaikitRuntimeException(
             grpc.StatusCode.INVALID_ARGUMENT, "Could not deserialize request"
         ) from e
+
+
+def raise_caikit_runtime_exception(exception: CaikitCoreException):
+    if exception.StatusCode == CaikitCoreStatusCode.NOT_FOUND:
+        log.debug2("Not Found Error: [%s]", exception.Message, exc_info=True)
+        raise CaikitRuntimeException(
+            grpc.StatusCode.NOT_FOUND,
+            exception.Message,
+        ) from exception
+
+    if exception.StatusCode == CaikitCoreStatusCode.INVALID_ARGUMENT:
+        log.debug2("Invalid Argument Error: [%s]", exception.Message, exc_info=True)
+        raise CaikitRuntimeException(
+            grpc.StatusCode.INVALID_ARGUMENT,
+            exception.Message,
+        ) from exception
+
+    if exception.StatusCode == CaikitCoreStatusCode.FORBIDDEN:
+        log.debug2("Forbidden Error: [%s]", exception.Message, exc_info=True)
+        raise CaikitRuntimeException(
+            grpc.StatusCode.PERMISSION_DENIED,
+            exception.Message,
+        ) from exception
+
+    if exception.StatusCode == CaikitCoreStatusCode.CONNECTION_ERROR:
+        log.debug2("Connection Error: [%s]", exception.Message, exc_info=True)
+        raise CaikitRuntimeException(
+            grpc.StatusCode.UNAVAILABLE,
+            exception.Message,
+        ) from exception
+
+    if exception.StatusCode == CaikitCoreStatusCode.UNAUTHORIZED:
+        log.debug2("Unauthorized Error: [%s]", exception.Message, exc_info=True)
+        raise CaikitRuntimeException(
+            grpc.StatusCode.UNAUTHENTICATED,
+            exception.Message,
+        ) from exception
+
+    log.debug2(
+        "Unknown CaikitCoreException Error: [%s]", exception.Message, exc_info=True
+    )
+    raise CaikitRuntimeException(
+        grpc.StatusCode.UNKNOWN, exception.Message
+    ) from exception


Test patch:
diff --git a/tests/core/model_management/test_local_model_trainer.py b/tests/core/model_management/test_local_model_trainer.py
index ac360fb0..c84c96d6 100644
--- a/tests/core/model_management/test_local_model_trainer.py
+++ b/tests/core/model_management/test_local_model_trainer.py
@@ -30,6 +30,7 @@ import aconfig
 # Local
 from caikit.config import get_config
 from caikit.core.data_model import DataStream, TrainingStatus
+from caikit.core.exceptions.caikit_core_exception import CaikitCoreException
 from caikit.core.model_management.local_model_trainer import LocalModelTrainer
 from sample_lib.modules import SampleModule
 
@@ -228,7 +229,7 @@ def test_purge_retention_time(trainer_type_cfg):
     retrieved_future = trainer.get_model_future(model_future.id)
     assert retrieved_future is model_future
     model_future._completion_time = model_future._completion_time - timedelta(days=2)
-    with pytest.raises(ValueError):
+    with pytest.raises(CaikitCoreException):
         trainer.get_model_future(model_future.id)
 
 
diff --git a/tests/runtime/servicers/test_training_management_servicer.py b/tests/runtime/servicers/test_training_management_servicer.py
index 98e76122..3eb65013 100644
--- a/tests/runtime/servicers/test_training_management_servicer.py
+++ b/tests/runtime/servicers/test_training_management_servicer.py
@@ -24,6 +24,7 @@ import pytest
 # Local
 from caikit.core import MODEL_MANAGER
 from caikit.core.data_model import DataStream, TrainingStatus
+from caikit.core.exceptions.caikit_core_exception import CaikitCoreException
 from caikit.interfaces.runtime.data_model import TrainingInfoRequest
 from caikit.runtime.servicers.training_management_servicer import (
     TrainingManagementServicerImpl,
@@ -178,9 +179,8 @@ def test_training_status_incorrect_id(training_management_servicer):
 
     assert context.value.status_code == grpc.StatusCode.NOT_FOUND
     assert (
-        "some_random_id not found in the list of currently running training jobs"
-        in context.value.message
-    )
+        f"Unknown training_id: some_random_id" in context.value.message
+    )  # message set by local_model_trainer.get_model_future
 
 
 def test_training_raises_when_cancel_on_incorrect_id(training_management_servicer):
@@ -207,10 +207,8 @@ def test_training_raises_when_cancel_on_incorrect_id(training_management_service
 
     assert context.value.status_code == grpc.StatusCode.NOT_FOUND
     assert (
-        "some_random_id not found in the list of currently running training jobs."
-        in context.value.message
-        and "Could not perform cancel" in context.value.message
-    )
+        f"Unknown training_id: some_random_id" in context.value.message
+    )  # message set by local_model_trainer.get_model_future
 
 
 def test_training_fails(training_management_servicer, training_pool):