Repository: SeldonIO/MLServer

Issue:
Rasing an exception during load of a model.py custom model shuts down processing loop

```
Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/site-packages/mlserver/parallel/worker.py", line 140, in _process_model_update
    await self._model_registry.load(model_settings)
  File "/opt/conda/lib/python3.8/site-packages/mlserver/registry.py", line 273, in load
    return await self._models[model_settings.name].load(model_settings)
  File "/opt/conda/lib/python3.8/site-packages/mlserver/registry.py", line 141, in load
    await self._load_model(new_model)
  File "/opt/conda/lib/python3.8/site-packages/mlserver/registry.py", line 157, in _load_model
    await model.load()
  File "/mnt/agent/models/choice-is-two_1/1/model.py", line 25, in load
    raise ModelParametersMissing(self.name, "no settings.parameters.extra.init_parameters found")
model.ModelParametersMissing: Parameters missing for model choice-is-two_1 no settings.parameters.extra.init_parameters found
2022-12-03 16:11:47,354 [mlserver.parallel] ERROR - Response processing loop crashed. Restarting the loop...
Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/site-packages/mlserver/parallel/dispatcher.py", line 40, in _process_responses_cb
    process_responses.result()
  File "/opt/conda/lib/python3.8/site-packages/mlserver/parallel/dispatcher.py", line 54, in _process_responses
    response = await loop.run_in_executor(self._executor, self._responses.get)
  File "/opt/conda/lib/python3.8/concurrent/futures/thread.py", line 57, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/opt/conda/lib/python3.8/multiprocessing/queues.py", line 116, in get
    return _ForkingPickler.loads(res)
TypeError: __init__() missing 1 required positional argument: 'reason'
```


Discussion:
The above error happened during a grpc loadModel call and that loadModel call did not return with an error but was left waiting (presently no timeout).
Just as an update on this one, this seems to be an edge case triggered by raising custom exceptions. It doesn't affect regular exceptions (e.g. like `MLServerError`), or custom exceptions built within a custom image, but rather exceptions defined in [custom code loaded on-the-fly](https://mlserver.readthedocs.io/en/latest/user-guide/custom.html#loading-a-custom-mlserver-runtime).

It's quite a strange one. I think the best fix would be to stop propagating exceptions from the parallel workers back to the main process. We can log the stacktrace and original exception on the workers, but then just return a dummy exception with the message (and status code, if it inherits from `MLServerError`).  

Resolving patch:
diff --git a/mlserver/parallel/errors.py b/mlserver/parallel/errors.py
new file mode 100644
--- /dev/null
+++ b/mlserver/parallel/errors.py
@@ -0,0 +1,27 @@
+from fastapi import status
+
+from ..utils import get_import_path
+from ..errors import MLServerError
+
+
+class WorkerError(MLServerError):
+    """
+    Class used to wrap exceptions raised from the workers.
+
+    All stacktrace details will be hidden, and the original class won't be
+    returned. This is to avoid issues with custom exceptions, like:
+
+        https://github.com/SeldonIO/MLServer/issues/881
+    """
+
+    def __init__(self, exc: BaseException):
+        msg = str(exc)
+        if isinstance(exc, BaseException):
+            import_path = get_import_path(exc.__class__)
+            msg = f"{import_path}: {exc}"
+
+        status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
+        if isinstance(exc, MLServerError):
+            status_code = exc.status_code
+
+        super().__init__(msg, status_code)
diff --git a/mlserver/parallel/worker.py b/mlserver/parallel/worker.py
--- a/mlserver/parallel/worker.py
+++ b/mlserver/parallel/worker.py
@@ -19,6 +19,7 @@
 )
 from .utils import terminate_queue, END_OF_QUEUE
 from .logging import logger
+from .errors import WorkerError
 
 IGNORED_SIGNALS = [signal.SIGINT, signal.SIGTERM, signal.SIGQUIT]
 
@@ -125,7 +126,8 @@ async def _process_request(self, request) -> ModelResponseMessage:
                 f"An error occurred calling method '{request.method_name}' "
                 f"from model '{request.model_name}'."
             )
-            return ModelResponseMessage(id=request.id, exception=e)
+            worker_error = WorkerError(e)
+            return ModelResponseMessage(id=request.id, exception=worker_error)
 
     def _handle_response(self, process_task: Task):
         response_message = process_task.result()
@@ -148,7 +150,8 @@ async def _process_model_update(
             return ModelResponseMessage(id=update.id)
         except (Exception, CancelledError) as e:
             logger.exception("An error occurred processing a model update.")
-            return ModelResponseMessage(id=update.id, exception=e)
+            worker_error = WorkerError(e)
+            return ModelResponseMessage(id=update.id, exception=worker_error)
 
     def send_request(self, request_message: ModelRequestMessage):
         """


Test patch:
diff --git a/tests/parallel/test_errors.py b/tests/parallel/test_errors.py
new file mode 100644
--- /dev/null
+++ b/tests/parallel/test_errors.py
@@ -0,0 +1,44 @@
+import pytest
+
+from fastapi import status
+from asyncio import CancelledError
+
+from mlserver.errors import MLServerError, ModelNotFound
+from mlserver.parallel.errors import WorkerError
+
+
+class CustomError(MLServerError):
+    def __init__(self):
+        super().__init__("foo")
+
+
+@pytest.mark.parametrize(
+    "exc, message, status_code",
+    [
+        (
+            CustomError(),
+            "tests.parallel.test_errors.CustomError: foo",
+            status.HTTP_400_BAD_REQUEST,
+        ),
+        (
+            ModelNotFound("sum-model"),
+            "mlserver.errors.ModelNotFound: Model sum-model not found",
+            status.HTTP_404_NOT_FOUND,
+        ),
+        (
+            Exception("bar"),
+            "builtins.Exception: bar",
+            status.HTTP_500_INTERNAL_SERVER_ERROR,
+        ),
+        (
+            CancelledError("bad error"),
+            "asyncio.exceptions.CancelledError: bad error",
+            status.HTTP_500_INTERNAL_SERVER_ERROR,
+        ),
+    ],
+)
+def test_worker_error(exc, message, status_code):
+    err = WorkerError(exc)
+
+    assert str(err) == message
+    assert err.status_code == status_code
diff --git a/tests/parallel/test_model.py b/tests/parallel/test_model.py
--- a/tests/parallel/test_model.py
+++ b/tests/parallel/test_model.py
@@ -28,7 +28,8 @@ async def test_predict_error(
     with pytest.raises(MLServerError) as excinfo:
         await error_model.predict(inference_request)
 
-    assert str(excinfo.value) == ErrorModel.error_message
+    expected_msg = f"mlserver.errors.MLServerError: {ErrorModel.error_message}"
+    assert str(excinfo.value) == expected_msg
 
 
 async def test_metadata(
diff --git a/tests/parallel/test_pool.py b/tests/parallel/test_pool.py
--- a/tests/parallel/test_pool.py
+++ b/tests/parallel/test_pool.py
@@ -65,4 +65,5 @@ async def test_load_error(
     with pytest.raises(MLServerError) as excinfo:
         await inference_pool.load_model(load_error_model)
 
-    assert str(excinfo.value) == ErrorModel.error_message
+    expected_msg = f"mlserver.errors.MLServerError: {ErrorModel.error_message}"
+    assert str(excinfo.value) == expected_msg
diff --git a/tests/parallel/test_worker.py b/tests/parallel/test_worker.py
--- a/tests/parallel/test_worker.py
+++ b/tests/parallel/test_worker.py
@@ -1,6 +1,7 @@
 from multiprocessing import Queue
 
 from mlserver.settings import ModelSettings
+from mlserver.parallel.errors import WorkerError
 from mlserver.parallel.worker import Worker
 from mlserver.parallel.messages import ModelUpdateMessage, ModelRequestMessage
 
@@ -112,4 +113,5 @@ async def _async_exception(*args, **kwargs):
     assert response.id == inference_request_message.id
     assert response.return_value is None
     assert response.exception is not None
-    assert str(response.exception) == error_msg
+    assert response.exception.__class__ == WorkerError
+    assert str(response.exception) == f"builtins.Exception: {error_msg}"