Repository: caikit/caikit
Pull request #297 (2023-07-03 14:40:02)

🐛 Fix multi-processing exception handling and train_location_handling fixture
### Changes
- Fix `set_train_location` fixture, which was not setting the `sample_train_servicer.use_subprocess` value correctly because:
    1. The format of config had changed
    2. The scope of `sample_train_servicer` was `session`, so the value of use_subprocess was not getting changed by test. Additionally, the GlobalTrainServicer interacts with service generation etc, changing the scope to `function` for this fixture would create errors.
- The `error` in `_ErrorCaptureProcess` was actually never getting set in case of exception. The reason for that is after forking another process, the instance with parent process of `_ErrorCaptureProcess` will not be same as with the child process. So even if child process sets the value of `self.error` it will not get propagated to the parent and thus the `SubProcessTrainExecutor` will never get any value in `self._worker.error`. In this fix, I am using python multiprocessing Pipe mechanism to communicate between child and parent.


Patch:
diff --git a/caikit/runtime/work_management/train_executors.py b/caikit/runtime/work_management/train_executors.py
index 799c514e..2779fbd9 100644
--- a/caikit/runtime/work_management/train_executors.py
+++ b/caikit/runtime/work_management/train_executors.py
@@ -199,23 +199,32 @@ class SubProcessTrainSaveExecutor(TrainSaveExecutorBase):
 
         def __init__(self, *args, event=None, **kwargs):
             super().__init__(*args, **kwargs)
-            self.error = None
+            self._error = None
+            self._parent_conn, self._child_conn = multiprocessing.Pipe()
             self._completion_event = event
 
         def run(self, *args, **kwargs):
             try:
-                return super().run(*args, **kwargs)
+                super().run(*args, **kwargs)
+                self._child_conn.send(None)
 
             # Catch any errors thrown within a subprocess so that they can be
             # forwarded to the parent
             # pylint: disable=broad-exception-caught
             except Exception as err:
-                log.error("<RUN69863806E>", "Caught exception in training", repr(err))
-                self.error = err
-
+                err_str = repr(err)
+                log.error("<RUN69863806E>", "Caught exception in training", err_str)
+                self._child_conn.send(err_str)
             finally:
                 self._completion_event.set()
 
+        @property
+        def error(self):
+            # NOTE: below is relying on child only sending error back and no other message
+            if self._parent_conn.poll():
+                self._error = self._parent_conn.recv()
+            return self._error
+
     def __init__(self, cancel_event) -> None:
 
         self.complete_event = multiprocessing.Event()
@@ -301,12 +310,10 @@ class SubProcessTrainSaveExecutor(TrainSaveExecutorBase):
         # If an error occurred, reraise it here
         # TODO: Make sure the stack trace is preserved
         if self._worker.error is not None:
-            if isinstance(self._worker.error, CaikitRuntimeException):
-                raise self._worker.error
             raise CaikitRuntimeException(
                 StatusCode.INTERNAL,
-                "Error caught in training subprocess",
-            ) from self._worker.error
+                f"Exception raised during training: {self._worker.error}",
+            )
 
         # If process exited with a non-zero exit code
         if self._worker.exitcode and self._worker.exitcode != os.EX_OK:


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 9ac047c2..9fc6d507 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -40,12 +40,15 @@ import caikit.core
 
 
 @pytest.fixture(autouse=True, params=[True, False])
-def set_train_location(request):
+def set_train_location(request, sample_train_servicer):
     """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}}):
-        yield
+    prev_value = sample_train_servicer.use_subprocess
+    sample_train_servicer.use_subprocess = request.param
+    yield
+    # Reset use_subprocess to previous value
+    sample_train_servicer.use_subprocess = prev_value
 
 
 # Train tests for the GlobalTrainServicer class ############################################################
@@ -78,6 +81,7 @@ def test_global_train_sample_task(
     training_response = sample_train_servicer.Train(
         train_request, Fixtures.build_context("foo")
     )
+
     assert training_response.model_name == model_name
 
     assert training_response.training_id is not None
@@ -340,6 +344,7 @@ def test_global_train_Edge_Case_Widget_should_raise_when_error_surfaces_from_mod
     training_data = stream_type(
         jsondata=stream_type.JsonData(data=[SampleTrainingType(1)])
     ).to_proto()
+
     train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
         model_name=random_test_id(),
         batch_size=999,
@@ -351,11 +356,11 @@ def test_global_train_Edge_Case_Widget_should_raise_when_error_surfaces_from_mod
             train_request, Fixtures.build_context("foo")
         )
 
-        training_result = sample_train_servicer.training_map.get(
+        _ = sample_train_servicer.training_map.get(
             training_response.training_id
         ).result()
 
-    assert f"This may be a problem with your input" in str(context.value.message)
+    assert f"Batch size of 999 is not allowed!" in str(context.value.message)
 
 
 def test_global_train_returns_exit_code_with_oom(