Repository: caikit/caikit
Pull request #385 (2023-07-31 21:04:06)

:goal_net: Tighter handling around training mgmt
<!--  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**:

This just tightens up the error handling in the training management handlers.

I'm getting a 13 response with no related logs from caikit even on debug4 so this decreases the scope of each try/catch to provide better visibility into what's going wrong.

**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/runtime/servicers/training_management_servicer.py b/caikit/runtime/servicers/training_management_servicer.py
index 81dfbded..b1941316 100644
--- a/caikit/runtime/servicers/training_management_servicer.py
+++ b/caikit/runtime/servicers/training_management_servicer.py
@@ -37,14 +37,13 @@ class TrainingManagementServicerImpl:
     service in Model Mesh as a Model-Runtime."""
 
     def GetTrainingStatus(self, request, context):  # pylint: disable=unused-argument
-        """Missing associated documentation comment in .proto file."""
+        """Get the status of a training by ID"""
         training_info = TrainingInfoRequest.from_proto(request)
+        model_future = self._get_model_future(
+            training_info.training_id, operation="get_status"
+        )
 
         try:
-            model_future = MODEL_MANAGER.get_model_future(
-                training_id=training_info.training_id
-            )
-
             reasons = []
             if model_future.get_info().errors:
                 reasons = [str(error) for error in model_future.get_info().errors]
@@ -54,10 +53,10 @@ class TrainingManagementServicerImpl:
                 state=model_future.get_info().status,
                 reasons=reasons,
             ).to_proto()
-        except ValueError as err:
+        except Exception as err:
             raise CaikitRuntimeException(
-                grpc.StatusCode.NOT_FOUND,
-                "{} not found in the list of currently running training jobs".format(
+                grpc.StatusCode.INTERNAL,
+                "Failed to get status for training id {}".format(
                     training_info.training_id,
                 ),
             ) from err
@@ -65,12 +64,11 @@ class TrainingManagementServicerImpl:
     def CancelTraining(self, request, context):  # pylint: disable=unused-argument
         """Cancel a training future."""
         training_info = TrainingInfoRequest.from_proto(request)
+        model_future = self._get_model_future(
+            training_info.training_id, operation="cancel"
+        )
 
         try:
-            model_future = MODEL_MANAGER.get_model_future(
-                training_id=training_info.training_id
-            )
-
             model_future.cancel()
 
             reasons = []
@@ -82,19 +80,46 @@ class TrainingManagementServicerImpl:
                 state=model_future.get_info().status,
                 reasons=reasons,
             ).to_proto()
+        except Exception as err:
+            log.debug2(
+                "Unexpected error trying to cancel training id %s: [%s]",
+                training_info.training_id,
+                err,
+            )
+            raise CaikitRuntimeException(
+                grpc.StatusCode.INTERNAL,
+                "Failed to cancel training id {}".format(
+                    training_info.training_id,
+                ),
+            ) from err
 
+    @staticmethod
+    def _get_model_future(training_id: str, operation: str):
+        """Returns a model future, or raises 404 caikit runtime exception on error.
+        Wrapped here so that we only catch errors directly in the `trainer.get_model_future` call
+        """
+        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. \
-                    Did not perform cancel.".format(
-                    training_info.training_id,
-                ),
+                "{} not found in the list of currently running training jobs. "
+                "Could not perform {}".format(training_id, operation),
             ) from err
         except Exception as err:
+            log.debug2(
+                "Caught unexpected exception while trying to look up model future for id %s: [%s]",
+                training_id,
+                err,
+            )
             raise CaikitRuntimeException(
                 grpc.StatusCode.INTERNAL,
-                "Failed to cancel training id {}".format(
-                    training_info.training_id,
+                "Unexpected error with training id {}. Could not perform {}".format(
+                    training_id, operation
                 ),
             ) from err


Test patch:
diff --git a/tests/runtime/servicers/test_training_management_servicer.py b/tests/runtime/servicers/test_training_management_servicer.py
index 7a541bd5..98e76122 100644
--- a/tests/runtime/servicers/test_training_management_servicer.py
+++ b/tests/runtime/servicers/test_training_management_servicer.py
@@ -209,7 +209,7 @@ def test_training_raises_when_cancel_on_incorrect_id(training_management_service
     assert (
         "some_random_id not found in the list of currently running training jobs."
         in context.value.message
-        and "Did not perform cancel" in context.value.message
+        and "Could not perform cancel" in context.value.message
     )