Repository: caikit/caikit (python)

Issue:
Support "spawn" for training subprocesses
## Describe the bug

Some libraries such as `torch` have significant issues when running in a `subprocess` using the `"fork"` method. The current `DestroyableThread` [explicitly uses "fork"](https://github.com/caikit/caikit/blob/main/caikit/core/toolkit/destroyable_process.py#L18) to avoid problems with data model object not being recreated correctly, but this can break downstream usage of libraries such as `torch`.

Pull request description:
Use spawn for train subprocess 384
**What this PR does / why we need it**:

Closes #384 

This PR makes it safe to use "spawn" when performing training jobs.

**If applicable**:
- [x] this PR contains unit tests
- [x] this PR has been tested for backwards compatibility


Patch:
diff --git a/caikit/config/config.yml b/caikit/config/config.yml
index 2196bb74..74ae887f 100644
--- a/caikit/config/config.yml
+++ b/caikit/config/config.yml
@@ -38,6 +38,7 @@ model_management:
                 # Format: XdXhXmXs
                 retention_duration: 1d
                 use_subprocess: false
+                subprocess_start_method: spawn
     finders:
         default:
             type: LOCAL
diff --git a/caikit/core/model_management/local_model_trainer.py b/caikit/core/model_management/local_model_trainer.py
index 988751bf..4cdf16a2 100644
--- a/caikit/core/model_management/local_model_trainer.py
+++ b/caikit/core/model_management/local_model_trainer.py
@@ -18,7 +18,7 @@ The LocalModelTrainer uses a local thread to launch and manage each training job
 # Standard
 from concurrent.futures.thread import _threads_queues
 from datetime import datetime, timedelta
-from typing import Optional, Type, Union
+from typing import Any, Dict, Iterable, Optional, Type, Union
 import os
 import re
 import threading
@@ -35,6 +35,7 @@ from ..modules import ModuleBase
 from ..toolkit.destroyable_process import DestroyableProcess
 from ..toolkit.destroyable_thread import DestroyableThread
 from ..toolkit.errors import error_handler
+from ..toolkit.logging import configure as configure_logging
 from .model_trainer_base import ModelTrainerBase, TrainingInfo
 import caikit
 
@@ -66,13 +67,14 @@ class LocalModelTrainer(ModelTrainerBase):
             self,
             trainer_name: str,
             module_class: Type[ModuleBase],
-            *args,
             save_path: Optional[Union[str, S3Path]],
             save_with_id: bool,
             model_name: Optional[str],
             external_training_id: Optional[str],
             use_subprocess: bool,
-            **kwargs,
+            subprocess_start_method: str,
+            args: Iterable[Any],
+            kwargs: Dict[str, Any],
         ):
             super().__init__(
                 trainer_name=trainer_name,
@@ -98,11 +100,13 @@ class LocalModelTrainer(ModelTrainerBase):
 
             # Set up the worker and start it
             self._use_subprocess = use_subprocess
+            self._subprocess_start_method = subprocess_start_method
             if self._use_subprocess:
                 log.debug2("Running training %s as a SUBPROCESS", self.id)
                 self._worker = DestroyableProcess(
+                    start_method=self._subprocess_start_method,
                     target=self._train_and_save,
-                    return_result=True,
+                    return_result=False,
                     args=args,
                     kwargs={
                         **kwargs,
@@ -205,6 +209,9 @@ class LocalModelTrainer(ModelTrainerBase):
 
         def _train_and_save(self, *args, **kwargs):
             """Function that will run in the worker thread"""
+            # If running in a spawned subprocess, reconfigure logging
+            if self._use_subprocess and self._subprocess_start_method != "fork":
+                configure_logging()
             with alog.ContextTimer(log.debug, "Training %s finished in: ", self.id):
                 trained_model = self._module_class.train(*args, **kwargs)
             if self.save_path is not None:
@@ -225,6 +232,7 @@ class LocalModelTrainer(ModelTrainerBase):
         """Initialize with a shared dict of all trainings"""
         self._instance_name = instance_name
         self._use_subprocess = config.get("use_subprocess", False)
+        self._subprocess_start_method = config.get("subprocess_start_method", "spawn")
         self._retention_duration = config.get("retention_duration")
         if self._retention_duration is not None:
             try:
@@ -268,17 +276,28 @@ class LocalModelTrainer(ModelTrainerBase):
         # Always purge old futures
         self._purge_old_futures()
 
+        # Wrap any models in the kwargs for safe spawning if needed
+        if self._use_subprocess and self._subprocess_start_method != "fork":
+            wrapped_models = {
+                key: _SpawnProcessModelWrapper(val)
+                for key, val in kwargs.items()
+                if isinstance(val, ModuleBase)
+            }
+            log.debug2("Subprocess wrapped models: %s", wrapped_models.keys())
+            kwargs.update(wrapped_models)
+
         # Create the new future
         model_future = self.LocalModelFuture(
             self._instance_name,
             module_class,
-            *args,
             save_path=save_path,
             save_with_id=save_with_id,
             external_training_id=external_training_id,
             use_subprocess=self._use_subprocess,
+            subprocess_start_method=self._subprocess_start_method,
             model_name=model_name,
-            **kwargs,
+            args=args,
+            kwargs=kwargs,
         )
 
         # Lock the global futures dict and add it to the dict
@@ -322,3 +341,49 @@ class LocalModelTrainer(ModelTrainerBase):
                 # NOTE: Concurrent purges could have already done this, so don't
                 #   error if the id is already gone
                 self._futures.pop(fid, None)
+
+
+class _SpawnProcessModelWrapper(ModuleBase):
+    """This class wraps up a model to make it safe to pass to a spawned
+    subprocess. It will not be efficient, but it will be safe!
+    """
+
+    def __init__(self, model: ModuleBase):
+        super().__init__()
+        self._model = model
+
+    def __getattr__(self, name):
+        """Forward attributes that are not found on the base class to the model
+
+        NOTE: This does _not_ forward base class attributes since those are
+            resolved before __getattr__ is called.
+        """
+        return getattr(self._model, name)
+
+    def save(self, *args, **kwargs):
+        """Directly forward save to the model so that it is not called by the
+        base class
+        """
+        return self._model.save(*args, **kwargs)
+
+    def run(self, *args, **kwargs):
+        """Directly forward run to the model so that it is not called by the
+        base class
+        """
+        return self._model.run(*args, **kwargs)
+
+    def __getstate__(self) -> bytes:
+        """When pickling, only send the serialized model body for non-fork. This
+        is not a general-purpose pickle solution for models, but makes them safe
+        for training jobs that need to move models between processes.
+        """
+        return self._model.as_bytes()
+
+    def __setstate__(self, pickled: bytes):
+        """When unpickling, deserialize the body if the model is not already
+        loaded in the model manager. This must be used in conjunction with the
+        above __getstate__ across a process boundary and should not be used as a
+        general-purpose deserialization for models.
+        """
+        retrieved_model = caikit.core.load(pickled)
+        self._model = retrieved_model
diff --git a/caikit/core/toolkit/destroyable_process.py b/caikit/core/toolkit/destroyable_process.py
index cadd9693..a54fdec7 100644
--- a/caikit/core/toolkit/destroyable_process.py
+++ b/caikit/core/toolkit/destroyable_process.py
@@ -32,18 +32,26 @@ import alog
 
 # Local
 from .destroyable import Destroyable
+from .errors import error_handler
 
 log = alog.use_channel("DESTROY-PROC")
+error = error_handler.get(log)
 
-FORK_CTX = multiprocessing.get_context("fork")
 
 OOM_EXIT_CODE = 137
 
+FORK_CTX = multiprocessing.get_context("fork")
+SPAWN_CTX = multiprocessing.get_context("spawn")
+FORKSERVER_CTX = multiprocessing.get_context("forkserver")
+
 
-class DestroyableProcess(
-    FORK_CTX.Process, Destroyable
+class _DestroyableProcess(
+    multiprocessing.process.BaseProcess,
+    Destroyable,
 ):  # pylint: disable=too-many-instance-attributes
-    __doc__ = __doc__
+    """The _DestroyableProcess base class implements a context-agnostic process
+    class that manages the subprocess and allows it to be destroyed
+    """
 
     def __init__(
         self,
@@ -52,12 +60,12 @@ class DestroyableProcess(
         args: Optional[Tuple] = None,
         kwargs: Optional[dict] = None,
         destroy_grace_period: float = 10,
-        return_result: bool = True,
+        return_result: bool = False,
         **_kwargs,
     ):
         """Initialize with an event to use to signal completion"""
-        self._parent_conn, self._child_conn = FORK_CTX.Pipe()
-        self._completion_event = completion_event or FORK_CTX.Event()
+        self._parent_conn, self._child_conn = self.__class__._MP_CTX.Pipe()
+        self._completion_event = completion_event or self.__class__._MP_CTX.Event()
         self._destroy_grace_period = destroy_grace_period
         self._return_result = return_result
 
@@ -112,9 +120,9 @@ class DestroyableProcess(
             )
 
         # Update the result and throw if it's an error
-        error = self.error
-        if error is not None:
-            raise error
+        err = self.error
+        if err is not None:
+            raise err
         return self.__result
 
     def destroy(self):
@@ -203,3 +211,37 @@ class DestroyableProcess(
         if self._return_result:
             self._child_conn.send(result)
         self._completion_event.set()
+
+
+class _ForkDestroyableProcess(FORK_CTX.Process, _DestroyableProcess):
+    _MP_CTX = FORK_CTX
+
+
+class _SpawnDestroyableProcess(SPAWN_CTX.Process, _DestroyableProcess):
+    _MP_CTX = SPAWN_CTX
+
+
+class _ForkserverDestroyableProcess(FORKSERVER_CTX.Process, _DestroyableProcess):
+    _MP_CTX = FORKSERVER_CTX
+
+
+_PROCESS_TYPES = {
+    "fork": _ForkDestroyableProcess,
+    "forkserver": _ForkserverDestroyableProcess,
+    "spawn": _SpawnDestroyableProcess,
+}
+
+
+def DestroyableProcess(start_method: str, *args, **kwargs):
+    """Class wrapper that returns the appropriate process type based on the
+    requested start method.
+
+    NOTE: Naming intentionally looks like a class!
+    """
+    error.value_check(
+        "<COR16699811E>",
+        start_method in _PROCESS_TYPES,
+        "Unsupported start_method: {}",
+        start_method,
+    )
+    return _PROCESS_TYPES[start_method](*args, **kwargs)
diff --git a/caikit/runtime/service_generation/data_stream_source.py b/caikit/runtime/service_generation/data_stream_source.py
index 37ef65ec..607de0a8 100644
--- a/caikit/runtime/service_generation/data_stream_source.py
+++ b/caikit/runtime/service_generation/data_stream_source.py
@@ -78,6 +78,9 @@ class DataStreamSourceBase(DataStream):
         """
         new_inst = self.__class__.from_binary_buffer(pickle_bytes)
         setattr(self, new_inst.which_oneof("data_stream"), new_inst.data_stream)
+        self.generator_func = self._generator
+        self.generator_args = tuple()
+        self.generator_kwargs = {}
 
     # pylint: disable=too-many-return-statements
     def to_data_stream(self) -> DataStream:
diff --git a/caikit/runtime/utils/servicer_util.py b/caikit/runtime/utils/servicer_util.py
index 80a97c56..a0bc74cd 100644
--- a/caikit/runtime/utils/servicer_util.py
+++ b/caikit/runtime/utils/servicer_util.py
@@ -14,7 +14,9 @@
 """A generic module to help Predict and Train servicers
 """
 # Standard
-from typing import Any, Dict, Iterable, Iterator
+from functools import update_wrapper
+from typing import Any, Dict, Iterable, Iterator, Tuple
+import sys
 import traceback
 
 # Third Party
@@ -31,6 +33,10 @@ from caikit.core.data_model.base import DataBase
 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
+from caikit.runtime.service_generation.data_stream_source import (
+    DataStreamSourceBase,
+    make_data_stream_source,
+)
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 from caikit.runtime.utils.import_util import get_data_model
 
@@ -260,6 +266,95 @@ def validate_data_model(
         validate_caikit_library_class_method_exists(caikit_Library_class, "to_proto")
 
 
+class ServicePackageStreamWrapper(DataStreamSourceBase):
+    """This class wraps up a DataStreamSourceBase derived class so that it can
+    be safely pickled to a subprocess using "spawn".
+
+    When using "spawn" to start a subprocess, a clean python process is created
+    and none of the runtime code leading to the creation is re-executed. For
+    most classes that are created at import time, this is fine, but for
+    dynamically created classes like those created with make_data_stream_source,
+    the class is not recreated on spawn and cannot be looked up.
+
+    To get around this, when unpickling this class, it falls back to performing
+    service generation dynamically if the class doesn't already exist
+    """
+
+    def __init__(self, stream: DataStreamSourceBase):
+        """Initialize with the data stream source to wrap
+
+        NOTE: We _intentionally_ don't initialize the base class here and
+            instead we forward
+        """
+        self._stream = stream
+        super().__init__()
+        update_wrapper(self, stream)
+
+    def __getstate__(self) -> Tuple[str, str, type, bytes]:
+        """Pickle as a tuple of primitive types
+
+        Returns:
+            stream_module_name (str): The name of the module that holds the
+                stream source class
+            stream_class_name (str): The name of the stream source class itself
+            stream_class_element_type (type): The type of the elements for this
+                stream source
+            content (bytes): The serialized bytes content of the underlying
+                stream source as a dataobject
+        """
+        return (
+            self._stream.__class__.__module__,
+            self._stream.__class__.__name__,
+            self._stream.__class__.ELEMENT_TYPE,
+            self._stream.__getstate__(),
+        )
+
+    def __setstate__(self, pickled: Tuple[str, str, type, bytes]):
+        """Unpickle the values from pickling and fall back to dynamic service
+        generation if needed.
+
+        Args:
+            pickled (Tuple[str, str, type, bytes]): The four elements created
+                when pickling this instance
+        """
+        stream_type_module, stream_type_name, element_type_class, stream_bytes = pickled
+        try:
+            mod = sys.modules[stream_type_module]
+            stream_type = getattr(mod, stream_type_name)
+        except (KeyError, AttributeError):
+            # Local
+            from caikit.runtime.service_factory import (  # pylint: disable=import-outside-toplevel
+                ServicePackageFactory,
+            )
+
+            ServicePackageFactory.get_service_package(
+                ServicePackageFactory.ServiceType.INFERENCE
+            )
+            ServicePackageFactory.get_service_package(
+                ServicePackageFactory.ServiceType.TRAINING
+            )
+            stream_type = make_data_stream_source(element_type_class)
+        self._stream = stream_type.__new__(stream_type)
+        self._stream.__setstate__(stream_bytes)
+
+    def __getattr__(self, name: str) -> Any:
+        """Forward all getattr requests to the underlying stream"""
+        return getattr(self._stream, name)
+
+    def __setattr__(self, name: str, val: Any):
+        """Forward attribute setting for generic attrs to the stream"""
+        if name == "_stream":
+            super().__setattr__(name, val)
+        else:
+            stream_attrs = list(self._stream.__class__.__annotations__.keys()) + list(
+                self._stream.get_proto_class().DESCRIPTOR.fields_by_name.keys()
+            )
+            if name in stream_attrs:
+                setattr(self._stream, name, val)
+            else:
+                super().__setattr__(name, val)
+
+
 def build_caikit_library_request_dict(
     request: ProtoMessageType,
     module_signature: CaikitMethodSignature,
@@ -313,16 +408,25 @@ def build_caikit_library_request_dict(
         for absent_field_name in absent_field_names:
             kwargs_dict.pop(absent_field_name)
 
-        # 3. Model Pointers
+        # 3. Handle type conversions
+        updated_kwargs = {}
         for field_name, field_value in kwargs_dict.items():
+
+            # 3.1 Model Pointers
             if isinstance(field_value, ModelPointer):
-                log.debug2("field_value is a ModelPointer obj")
+                log.debug2("field %s value is a ModelPointer obj", field_name)
                 model_manager = ModelManager.get_instance()
                 model_retrieved = model_manager.retrieve_model(field_value.model_id)
-                kwargs_dict[field_name] = model_retrieved
+                updated_kwargs[field_name] = model_retrieved
 
-        log.debug2("caikit_library_request_dict returned is: %s", kwargs_dict)
+            # 3.2 Data streams
+            elif isinstance(field_value, DataStreamSourceBase):
+                log.debug2("field %s value is a DataStreamSourceBase", field_name)
+                updated_kwargs[field_name] = ServicePackageStreamWrapper(field_value)
 
+        # Apply the type conversion updates
+        kwargs_dict.update(updated_kwargs)
+        log.debug2("caikit_library_request_dict returned is: %s", kwargs_dict)
         return kwargs_dict
 
     except CaikitRuntimeException as e:


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 5831ac5d..f1c6a180 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -49,12 +49,29 @@ def test_environment():
     """The most important fixture: This runs caikit configuration with the base test config overrides"""
     test_config_path = os.path.join(FIXTURES_DIR, "config", "config.yml")
     caikit.configure(test_config_path)
-    # import the mock backend that is specified in the config
-    # This is required to run any test that loads a model
-    # Local
-    from tests.core.helpers import MockBackend
-
-    yield
+    with tempfile.TemporaryDirectory() as workdir:
+        # 🌶️ All tests need to clean up after themselves! The metering and
+        # training output directories are the one thing easily creates leftover
+        # files, so we explicitly update config here to point them at a temp dir
+        with temp_config(
+            {
+                "runtime": {
+                    "metering": {
+                        "log_dir": os.path.join(workdir, "metering_logs"),
+                    },
+                    "training": {
+                        "output_dir": os.path.join(workdir, "training_output"),
+                    },
+                }
+            },
+            "merge",
+        ):
+            # import the mock backend that is specified in the config
+            # This is required to run any test that loads a model
+            # Local
+            from tests.core.helpers import MockBackend
+
+            yield
     # No cleanup required...?
 
 
diff --git a/tests/core/model_management/test_local_model_trainer.py b/tests/core/model_management/test_local_model_trainer.py
index dfd9ce0d..ac360fb0 100644
--- a/tests/core/model_management/test_local_model_trainer.py
+++ b/tests/core/model_management/test_local_model_trainer.py
@@ -28,9 +28,9 @@ import pytest
 import aconfig
 
 # Local
+from caikit.config import get_config
 from caikit.core.data_model import DataStream, TrainingStatus
 from caikit.core.model_management.local_model_trainer import LocalModelTrainer
-from caikit.core.toolkit.destroyable_process import OOM_EXIT_CODE
 from sample_lib.modules import SampleModule
 
 ## Helpers #####################################################################
@@ -53,7 +53,12 @@ def save_path():
 
 
 def get_event(cfg: dict):
-    return cfg.get("use_subprocess") and multiprocessing.Event() or threading.Event()
+    if cfg.get("use_subprocess"):
+        start_method = (
+            get_config().model_management.trainers.default.config.subprocess_start_method
+        )
+        return multiprocessing.get_context(start_method).Event()
+    return threading.Event()
 
 
 ## Tests #######################################################################
diff --git a/tests/core/toolkit/test_destroyable_process.py b/tests/core/toolkit/test_destroyable_process.py
index 46570b1a..6c6976c4 100644
--- a/tests/core/toolkit/test_destroyable_process.py
+++ b/tests/core/toolkit/test_destroyable_process.py
@@ -13,21 +13,51 @@
 # limitations under the License.
 
 # Standard
+import multiprocessing
 import time
 
 # Third Party
 import pytest
 
 # Local
-from caikit.core.toolkit.destroyable_process import FORK_CTX, DestroyableProcess
+from caikit.core.toolkit.destroyable_process import DestroyableProcess
 
+## Helpers #####################################################################
 
-def test_processes_can_be_interrupted():
-    def infinite_wait():
-        while True:
-            time.sleep(0.1)
 
-    proc = DestroyableProcess(infinite_wait)
+EXPECTED_THROW = ValueError("test-any-error")
+EXPECTED_SUCCESS = "test-any-result"
+
+
+def infinite_wait():
+    while True:
+        time.sleep(0.1)
+
+
+def long_sleep():
+    time.sleep(1000)
+
+
+def thrower():
+    raise EXPECTED_THROW
+
+
+def succeeder():
+    return EXPECTED_SUCCESS
+
+
+@pytest.fixture(
+    params=["fork", "forkserver", "spawn"],
+)
+def process_type(request):
+    yield request.param
+
+
+## Tests #######################################################################
+
+
+def test_processes_can_be_interrupted(process_type):
+    proc = DestroyableProcess(process_type, infinite_wait)
     proc.start()
     assert not proc.destroyed
     assert not proc.canceled
@@ -42,20 +72,19 @@ def test_processes_can_be_interrupted():
     assert not proc.threw
 
 
-def test_processes_can_return_results():
-    expected = "test-any-result"
-    proc = DestroyableProcess(lambda: expected)
+def test_processes_can_return_results(process_type):
+    proc = DestroyableProcess(process_type, succeeder, return_result=True)
     proc.start()
     proc.join()
-    assert expected == proc.get_or_throw()
+    assert EXPECTED_SUCCESS == proc.get_or_throw()
     assert not proc.destroyed
     assert not proc.canceled
     assert proc.ran
     assert not proc.threw
 
 
-def test_process_not_canceled_after_success():
-    proc = DestroyableProcess(lambda: None)
+def test_process_not_canceled_after_success(process_type):
+    proc = DestroyableProcess(process_type, succeeder)
     proc.start()
     proc.join()
     assert not proc.canceled
@@ -63,9 +92,8 @@ def test_process_not_canceled_after_success():
     assert not proc.canceled
 
 
-def test_processes_can_be_set_to_not_return_results():
-    expected = "test-any-result"
-    proc = DestroyableProcess(lambda: expected, return_result=False)
+def test_processes_can_be_set_to_not_return_results(process_type):
+    proc = DestroyableProcess(process_type, succeeder, return_result=False)
     proc.start()
     proc.join()
     assert proc.get_or_throw() is None
@@ -75,13 +103,8 @@ def test_processes_can_be_set_to_not_return_results():
     assert not proc.threw
 
 
-def test_processes_can_throw():
-    expected = ValueError("test-any-error")
-
-    def thrower():
-        raise expected
-
-    proc = DestroyableProcess(thrower)
+def test_processes_can_throw(process_type):
+    proc = DestroyableProcess(process_type, thrower)
     proc.start()
     proc.join()
     assert not proc.destroyed
@@ -92,11 +115,11 @@ def test_processes_can_throw():
     with pytest.raises(ValueError) as ctx:
         proc.get_or_throw()
 
-    assert str(expected) == str(ctx.value)
+    assert str(EXPECTED_THROW) == str(ctx.value)
 
 
-def test_processes_will_not_execute_if_destroyed_before_starting():
-    proc = DestroyableProcess(lambda: time.sleep(1000))
+def test_processes_will_not_execute_if_destroyed_before_starting(process_type):
+    proc = DestroyableProcess(process_type, long_sleep)
     proc.destroy()
     proc.start()
     assert not proc.is_alive()
@@ -109,9 +132,9 @@ def test_processes_will_not_execute_if_destroyed_before_starting():
     assert proc.threw
 
 
-def test_event_is_set_on_completion():
-    event = FORK_CTX.Event()
-    proc = DestroyableProcess(lambda: None, completion_event=event)
+def test_event_is_set_on_completion(process_type):
+    event = multiprocessing.get_context(process_type).Event()
+    proc = DestroyableProcess(process_type, succeeder, completion_event=event)
     assert not event.is_set()
     proc.start()
     proc.join()
@@ -122,13 +145,9 @@ def test_event_is_set_on_completion():
     assert not proc.threw
 
 
-def test_event_is_set_on_exception():
-    event = FORK_CTX.Event()
-
-    def thrower():
-        raise ValueError("test-any-exception")
-
-    proc = DestroyableProcess(thrower, completion_event=event)
+def test_event_is_set_on_exception(process_type):
+    event = multiprocessing.get_context(process_type).Event()
+    proc = DestroyableProcess(process_type, thrower, completion_event=event)
     assert not event.is_set()
     proc.start()
     proc.join()
@@ -139,8 +158,8 @@ def test_event_is_set_on_exception():
     assert proc.threw
 
 
-def test_default_event_is_set_on_completion():
-    proc = DestroyableProcess(lambda: None)
+def test_default_event_is_set_on_completion(process_type):
+    proc = DestroyableProcess(process_type, succeeder)
     assert not proc.completion_event.is_set()
     proc.start()
     proc.join()
diff --git a/tests/core/toolkit/test_error_handler.py b/tests/core/toolkit/test_error_handler.py
index ada7e641..5cf56d2b 100644
--- a/tests/core/toolkit/test_error_handler.py
+++ b/tests/core/toolkit/test_error_handler.py
@@ -351,6 +351,9 @@ class TestErrorHandler(TestCaseBase):
         with pytest.raises(TypeError):
             self.error.subclass_check("<III>", 1, int, allow_none=True)
 
+    @pytest.mark.skipif(
+        os.environ.get("LOG_FORMATTER") == "pretty", reason="Pretty formatter"
+    )
     def test_value_check_raises(self):
         """Verify that `value_check` raises a `ValueError`."""
         bad_value = 1.1
@@ -365,6 +368,9 @@ class TestErrorHandler(TestCaseBase):
             "exception raised: ValueError('value check failed: invalid range `1.1`')",
         )
 
+    @pytest.mark.skipif(
+        os.environ.get("LOG_FORMATTER") == "pretty", reason="Pretty formatter"
+    )
     def test_value_check_no_args_raises(self):
         """Verify that `value_check` raises a `ValueError` if no args are given with empty msg."""
         bad_value = 1.1
@@ -377,6 +383,9 @@ class TestErrorHandler(TestCaseBase):
             "exception raised: ValueError('value check failed: ')",
         )
 
+    @pytest.mark.skipif(
+        os.environ.get("LOG_FORMATTER") == "pretty", reason="Pretty formatter"
+    )
     def test_value_check_empty_msg_raises(self):
         """Verify that `value_check` raises a `ValueError` if msg is empty."""
         bad_value = 1.1
@@ -389,6 +398,9 @@ class TestErrorHandler(TestCaseBase):
             "exception raised: ValueError('value check failed: 1.1')",
         )
 
+    @pytest.mark.skipif(
+        os.environ.get("LOG_FORMATTER") == "pretty", reason="Pretty formatter"
+    )
     def test_value_check_empty_args_raises(self):
         """Verify that `value_check` raises a `ValueError` if only msg is provided."""
         bad_value = 1.1
diff --git a/tests/runtime/servicers/test_global_train_servicer_impl.py b/tests/runtime/servicers/test_global_train_servicer_impl.py
index 2801436b..0c4afdb9 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -25,6 +25,7 @@ import grpc
 import pytest
 
 # Local
+from caikit.config import get_config
 from caikit.core import MODEL_MANAGER
 from caikit.interfaces.common.data_model.stream_sources import S3Path
 from caikit.runtime.servicers.global_train_servicer import GlobalTrainServicer
@@ -423,9 +424,11 @@ def test_global_train_aborts_long_running_trains(
         oom_exit=False,
     )
 
-    # sample_train_servicer.use_subprocess = True
     if sample_train_servicer.use_subprocess:
-        test_event = multiprocessing.Event()
+        start_method = (
+            get_config().model_management.trainers.default.config.subprocess_start_method
+        )
+        test_event = multiprocessing.get_context(start_method).Event()
     else:
         test_event = threading.Event()
 
diff --git a/tests/runtime/utils/test_servicer_util.py b/tests/runtime/utils/test_servicer_util.py
index 264ce6dd..fe7beb0c 100644
--- a/tests/runtime/utils/test_servicer_util.py
+++ b/tests/runtime/utils/test_servicer_util.py
@@ -21,6 +21,7 @@ import pytest
 
 # Local
 from caikit.runtime.protobufs import model_runtime_pb2
+from caikit.runtime.service_generation.data_stream_source import DataStreamSourceBase
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 from caikit.runtime.utils.import_util import get_data_model
 from caikit.runtime.utils.servicer_util import (
@@ -306,6 +307,7 @@ def test_global_train_build_caikit_library_request_dict_strips_empty_list_from_r
     expected_arguments = {"training_data", "union_list"}
 
     assert expected_arguments == set(caikit.core_request.keys())
+    assert isinstance(caikit.core_request["training_data"], DataStreamSourceBase)
 
 
 def test_global_train_build_caikit_library_request_dict_works_for_repeated_fields(