Repository: caikit/caikit
Pull request #509 (2023-10-03 23:17:01)

🐛 Propagate sub-process errors properly
<!--  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**:

fix https://github.com/caikit/caikit/issues/501

There were 2 issues identified:
1. The tests in `test_model_train_servicer_impl.py` were not using `sub-process` correctly  since the `get_trainer` function always fetches a trainer that was initialized before our custom config (with `sub_process` set to `True`) comes into play. Hence we’re never actually checking a training in a `sub_process`. I verified this by running a failing test and seeing that we always hit `destroyable_thread.py` in the stacktrace instead of `destroyable_process` in both scenarios.
2. `CaikitRuntimeException` was not pickleable, which is a requirement for any object to be passed through sub-processes using `Pipe()`. 

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [x] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/runtime/types/caikit_runtime_exception.py b/caikit/runtime/types/caikit_runtime_exception.py
index 747bdb91..1f2e0552 100644
--- a/caikit/runtime/types/caikit_runtime_exception.py
+++ b/caikit/runtime/types/caikit_runtime_exception.py
@@ -13,8 +13,12 @@
 # limitations under the License.
 
 # Standard
+from typing import Dict
 import uuid
 
+# Third Party
+import grpc
+
 
 class CaikitRuntimeException(Exception):
     """Exception raised when something went wrong while running the Caikit Runtime.
@@ -28,7 +32,12 @@ class CaikitRuntimeException(Exception):
         Same as Args. See above.
     """
 
-    def __init__(self, status_code, message, metadata=None):
+    def __init__(
+        self,
+        status_code: grpc.StatusCode,
+        message: str,
+        metadata: Dict[any, any] = None,
+    ):
         super().__init__(message)
         self.status_code = status_code
         self.message = message
@@ -40,3 +49,12 @@ class CaikitRuntimeException(Exception):
         else:
             # metadata is of None type
             self.metadata = {"error_id": self.id}
+
+    # This is to make the CaikitRuntimeException pickleable
+    # this is needed to ensure we can pass it through sub-processes' pipe
+    # which requires objects to be pickleable
+    def __reduce__(self) -> "CaikitRuntimeException":
+        return (
+            CaikitRuntimeException,
+            (self.status_code, self.message, self.metadata),
+        )


Test patch:
diff --git a/tests/runtime/servicers/test_global_train_servicer_impl.py b/tests/runtime/servicers/test_global_train_servicer_impl.py
index 2a1d047a..858c8b85 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -379,6 +379,33 @@ def test_local_trainer_rejects_s3_output_paths(
     assert ctx.value.status_code == grpc.StatusCode.INVALID_ARGUMENT
 
 
+def test_global_train_surfaces_caikit_errors(sample_train_servicer, sample_text_file):
+    """Test whether global trainer surfaces errors from Caikit using both sub-process and thread"""
+    stream_type = caikit.interfaces.common.data_model.DataStreamSourceSampleTrainingType
+    # we don't support .txt files yet, hence this should throw an error
+    training_data = stream_type(file=stream_type.File(filename=sample_text_file))
+
+    train_class = get_train_request(SampleModule)
+    train_request_params_class = get_train_params(SampleModule)
+    train_request = train_class(
+        model_name=random_test_id(),
+        parameters=train_request_params_class(
+            batch_size=42,
+            training_data=training_data,
+        ),
+    ).to_proto()
+
+    training_response = sample_train_servicer.Train(
+        train_request, Fixtures.build_context("foo")
+    )
+    MODEL_MANAGER.get_model_future(training_response.training_id).wait()
+
+    future_info = MODEL_MANAGER.get_model_future(
+        training_response.training_id
+    ).get_info()
+    assert "Extension not supported" in str(future_info.errors[0])
+
+
 #####################################################################
 
 
diff --git a/tests/runtime/servicers/test_model_train_servicer_impl.py b/tests/runtime/servicers/test_model_train_servicer_impl.py
index 4c93a7fd..bbbdd61e 100644
--- a/tests/runtime/servicers/test_model_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_model_train_servicer_impl.py
@@ -29,7 +29,7 @@ from caikit.core.data_model import TrainingStatus
 from caikit.runtime.protobufs import process_pb2
 from caikit.runtime.servicers.model_train_servicer import ModelTrainServicerImpl
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
-from tests.conftest import temp_config
+from tests.conftest import set_use_subprocess, temp_config
 from tests.fixtures import Fixtures
 import sample_lib
 
@@ -57,12 +57,19 @@ def clear_messages_from_servicer(servicer):
         servicer._training_service.messages = messages
 
 
+# We need to add reset_model_manager to make sure all trainers use the config that we provide.
+# Without this, it didn't actually work since the get_trainer function always fetches a trainer
+# that was initialized before this config comes into play. Hence we’re never actually checking
+# a training in a sub_process. I verified this by running a failing test and seeing that we always
+# hit destroyable_thread.py in the stacktrace instead of destroyable_process in both scenarios.
+
+
 @pytest.fixture(autouse=True, params=[True, False])
-def set_train_location(request):
+def set_train_location(request, reset_model_manager):
     """This fixture ensures that all tests in this file will be run with both
     subprocess and local training styles
     """
-    with temp_config({"training": {"use_subprocess": request.param}}):
+    with set_use_subprocess(request.param):
         yield
 
 
@@ -184,6 +191,46 @@ def test_model_train_validation_error_raises(sample_model_train_servicer, output
         sample_model_train_servicer.Run(model_train_request, context)
 
 
+def test_model_train_surfaces_caikit_errors(sample_model_train_servicer, output_dir):
+    """Test whether model train surfaces errors from Caikit using both sub-process and thread"""
+    training_id = str(uuid.uuid4())
+
+    training_output_dir = os.path.join(output_dir, training_id)
+    training_input_dir = os.path.join(output_dir, training_id, "inputs")
+    # we don't support .txt files yet, hence this should throw an error
+    input_file_name = "data.txt"
+
+    os.makedirs(training_input_dir, exist_ok=True)
+    with open(os.path.join(training_input_dir, input_file_name), "w") as f:
+        json.dump([sample_lib.data_model.SampleTrainingType(number=1).to_dict()], f)
+
+    model_train_request = process_pb2.ProcessRequest(
+        trainingID=training_id,
+        request_dict={
+            "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
+            "training_params": json.dumps(
+                {
+                    "model_name": "abc",
+                    "parameters": {
+                        "training_data": {
+                            "file": {
+                                "filename": input_file_name  # This is relative to training_input_dir
+                            },
+                        },
+                    },
+                }
+            ),
+        },
+        training_input_dir=training_input_dir,
+        training_output_dir=training_output_dir,
+    )
+    context = Fixtures.build_context("test-any-unresponsive-model")
+    with pytest.raises(CaikitRuntimeException) as e:
+        sample_model_train_servicer.Run(model_train_request, context)
+    assert isinstance(e.value, CaikitRuntimeException)
+    assert "Extension not supported" in e.value.message
+
+
 #####################################################################
 # Normal tests
 @pytest.mark.parametrize("output_in_config", (True, False))
diff --git a/tests/runtime/types/test_caikit_runtime_exception.py b/tests/runtime/types/test_caikit_runtime_exception.py
new file mode 100644
index 00000000..3f94a63b
--- /dev/null
+++ b/tests/runtime/types/test_caikit_runtime_exception.py
@@ -0,0 +1,54 @@
+# 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.
+
+# Standard
+import pickle
+
+# Third Party
+import grpc
+
+# Local
+from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
+
+
+def test_caikit_runtime_exception_initialization():
+    """Test that we can construct a CaikitRuntimeException with the right args"""
+
+    caikit_runtime_exception = CaikitRuntimeException(
+        status_code=grpc.StatusCode.INTERNAL,
+        message="This is a test message",
+        metadata={"key1": "val1"},
+    )
+
+    assert isinstance(caikit_runtime_exception, Exception)
+    assert caikit_runtime_exception.id
+    assert caikit_runtime_exception.message == "This is a test message"
+    assert caikit_runtime_exception.metadata == {
+        "key1": "val1",
+        "error_id": caikit_runtime_exception.id,
+    }
+
+
+def test_caikit_runtime_exception_is_pickleable():
+    """This is to ensure we can send caikit_runtime_exception through sub-processes"""
+
+    caikit_runtime_exception = CaikitRuntimeException(
+        status_code=grpc.StatusCode.INTERNAL, message="This is a test message"
+    )
+
+    caikit_runtime_exception_pickle = pickle.dumps(caikit_runtime_exception)
+    caikit_runtime_exception_loaded = pickle.loads(caikit_runtime_exception_pickle)
+
+    assert isinstance(caikit_runtime_exception_loaded, CaikitRuntimeException)
+    assert caikit_runtime_exception_loaded.message == caikit_runtime_exception.message