Repository: caikit/caikit (python)

Issue:
Refactor: Extract protocol logic from global predict servicer
## Is your feature request related to a problem? Please describe.

For REST support, we're building a native rest server that would also like to use the `GlobalPredictServicer` to run model inference, instead of rewriting another inference handler.

But, the global predict handler currently assumes the request object is a `ProtobufMessage` and assumes that it should also serialize responses as `ProtobufMessage(s)`.

## Describe the solution you'd like

The protocol-specific (de)serialization logic should be extracted from the global predict servicer so that it can instead accept something more along the lines of a model id and a dictionary of kwargs to pass the model, and output the response of the model directly. This will allow us to plug in a REST translator alongside the gRPC translator so that we can run both servers using the same inference handler.

## Describe alternatives you've considered

Copy-pasting code into a REST specific handler, but we're on team re-use over here

## Additional context



Pull request:
Split global predict servicer
## Description

Closes #223 

This PR splits out the core inferencing logic from the `GlobalPredictServicer`'s `Predict` function so that it can be called directly with the data model objects. This is needed to allow a REST server to reuse the same inferencing logic as the grpc server.

Patch:
diff --git a/caikit/runtime/servicers/global_predict_servicer.py b/caikit/runtime/servicers/global_predict_servicer.py
index 30dd2004..1eca5493 100644
--- a/caikit/runtime/servicers/global_predict_servicer.py
+++ b/caikit/runtime/servicers/global_predict_servicer.py
@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # Standard
+from contextlib import contextmanager
 from importlib.metadata import version
-from typing import Set
+from typing import Iterable, Optional, Set, Union
 import traceback
 
 # Third Party
 from google.protobuf.descriptor import FieldDescriptor
-from grpc import StatusCode
+from google.protobuf.message import Message as ProtobufMessage
+from grpc import ServicerContext, StatusCode
 from prometheus_client import Counter, Summary
 
 # First Party
@@ -27,6 +29,7 @@ import alog
 # Local
 from caikit import get_config
 from caikit.core import ModuleBase
+from caikit.core.data_model import DataBase
 from caikit.runtime.metrics.rpc_meter import RPCMeter
 from caikit.runtime.model_management.model_manager import ModelManager
 from caikit.runtime.service_factory import ServicePackage
@@ -125,35 +128,34 @@ class GlobalPredictServicer:
         )
         super()
 
-    def Predict(self, request, context):
+    def Predict(
+        self,
+        request: ProtobufMessage,
+        context: ServicerContext,
+    ) -> Union[ProtobufMessage, Iterable[ProtobufMessage]]:
         """Global predict RPC -- Mocks the invocation of a Caikit Library module.run()
         method for a loaded Caikit Library model
 
         Args:
-            request(object):
+            request (ProtobufMessage):
                 A deserialized RPC request message
-            context(grpc.ServicerContext): Context object (contains request metadata, etc)
+            context (ServicerContext):
+                Context object (contains request metadata, etc)
 
         Returns:
-            response (object):
+            response (Union[ProtobufMessage, Iterable[ProtobufMessage]]):
                 A Caikit Library data model response object
         """
+        # Make sure the request has a model before doing anything
         request_name = request.DESCRIPTOR.name
-        outer_scope_name = "GlobalPredictServicer.Predict:%s" % request_name
-        inner_scope_name = (
-            "GlobalPredictServicer.Predict.caikit_library_run:%s" % request_name
-        )
-
-        try:
-            with alog.ContextLog(log.debug, outer_scope_name):
-                # Make sure the request has a model before doing anything
-                model_id = get_metadata(context, self.MODEL_MESH_MODEL_ID_KEY)
+        model_id = get_metadata(context, self.MODEL_MESH_MODEL_ID_KEY)
+        with self._handle_predict_exceptions(model_id, request_name):
+            with alog.ContextLog(
+                log.debug, "GlobalPredictServicer.Predict:%s", request_name
+            ):
                 # Retrieve the model from the model manager
                 log.debug("<RUN52259029D>", "Retrieving model '%s'", model_id)
                 model = self._model_manager.retrieve_model(model_id)
-
-                self._verify_model_task(model, request_name)
-
                 model_class = type(model)
                 # Unmarshall the request object into the required module run argument(s)
                 with PREDICT_FROM_PROTO_SUMMARY.labels(
@@ -163,22 +165,14 @@ class GlobalPredictServicer:
                         request,
                         model_class.RUN_SIGNATURE,
                     )
-
-                # NB: we previously recorded the size of the request, and timed this module to
-                # provide a rudimentary throughput metric of size / time
-                with alog.ContextLog(log.debug, inner_scope_name):
-                    with PREDICT_CAIKIT_LIBRARY_SUMMARY.labels(
-                        grpc_request=request_name, model_id=model_id
-                    ).time():
-                        if self.use_abortable_threads:
-                            work = AbortableAction(
-                                CallAborter(context),
-                                model.run,
-                                **caikit_library_request,
-                            )
-                            response = work.do()
-                        else:
-                            response = model.run(**caikit_library_request)
+                response = self.predict_model(
+                    request_name,
+                    model_id,
+                    aborter=CallAborter(context)
+                    if self.use_abortable_threads
+                    else None,
+                    **caikit_library_request,
+                )
 
                 # Marshall the response to the necessary return type
                 with PREDICT_TO_PROTO_SUMMARY.labels(
@@ -188,6 +182,53 @@ class GlobalPredictServicer:
                         response_proto = build_proto_stream(response)
                     else:
                         response_proto = build_proto_response(response)
+                return response_proto
+
+    def predict_model(
+        self,
+        request_name: str,
+        model_id: str,
+        aborter: Optional[CallAborter] = None,
+        **kwargs,
+    ) -> Union[DataBase, Iterable[DataBase]]:
+        """Run a prediction against the given model using the raw arguments to
+        the model's run function.
+
+        Args:
+            request_name (str):
+                The name of the request message to validate the model's task
+            model_id (str):
+                The ID of the loaded model
+            aborter (Optional[CallAborter]):
+                If using abortable calls, this is the aborter to use
+            **kwargs:
+                Keyword arguments to pass to the model's run function
+
+        Returns:
+            response (Union[DataBase, Iterable[DataBase]]):
+                The object (unary) or objects (output stream) produced by the
+                inference request
+        """
+
+        with self._handle_predict_exceptions(model_id, request_name):
+            model = self._model_manager.retrieve_model(model_id)
+            self._verify_model_task(model, request_name)
+
+            # NB: we previously recorded the size of the request, and timed this module to
+            # provide a rudimentary throughput metric of size / time
+            with alog.ContextLog(
+                log.debug,
+                "GlobalPredictServicer.Predict.caikit_library_run:%s",
+                request_name,
+            ):
+                with PREDICT_CAIKIT_LIBRARY_SUMMARY.labels(
+                    grpc_request=request_name, model_id=model_id
+                ).time():
+                    if aborter is not None:
+                        work = AbortableAction(aborter, model.run, **kwargs)
+                        response = work.do()
+                    else:
+                        response = model.run(**kwargs)
 
             # Update Prometheus metrics
             PREDICT_RPC_COUNTER.labels(
@@ -195,13 +236,20 @@ class GlobalPredictServicer:
             ).inc()
             if get_config().runtime.metering.enabled:
                 self.rpc_meter.update_metrics(str(type(model)))
-            return response_proto
+            return response
+
+    ## Implementation Details ##################################################
+
+    @contextmanager
+    def _handle_predict_exceptions(self, model_id: str, request_name: str):
+        try:
+            yield
 
         except CaikitRuntimeException as e:
             log_dict = {
                 "log_code": "<RUN50530380W>",
                 "message": e.message,
-                "model_id": get_metadata(context, self.MODEL_MESH_MODEL_ID_KEY),
+                "model_id": model_id,
                 "error_id": e.id,
             }
             log.warning({**log_dict, **e.metadata})
@@ -216,7 +264,7 @@ class GlobalPredictServicer:
             log_dict = {
                 "log_code": "<RUN490439039W>",
                 "message": repr(e),
-                "model_id": get_metadata(context, self.MODEL_MESH_MODEL_ID_KEY),
+                "model_id": model_id,
                 "stack_trace": traceback.format_exc(),
             }
             log.warning(log_dict)
@@ -234,7 +282,7 @@ class GlobalPredictServicer:
             log_dict = {
                 "log_code": "<RUN49049070W>",
                 "message": repr(e),
-                "model_id": get_metadata(context, self.MODEL_MESH_MODEL_ID_KEY),
+                "model_id": model_id,
                 "stack_trace": traceback.format_exc(),
             }
             log.warning(log_dict)


Test patch:
diff --git a/tests/runtime/servicers/test_global_predict_servicer_impl.py b/tests/runtime/servicers/test_global_predict_servicer_impl.py
index 7b58a21f..f107107b 100644
--- a/tests/runtime/servicers/test_global_predict_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_predict_servicer_impl.py
@@ -22,11 +22,8 @@ except (NameError, ModuleNotFoundError):
 
 # Standard
 import json
-import os
-import tempfile
 import threading
 import time
-import uuid
 
 # Third Party
 import grpc
@@ -39,10 +36,11 @@ from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 from sample_lib.data_model import SampleInputType, SampleOutputType
 from sample_lib.modules.sample_task import SampleModule
 from tests.fixtures import Fixtures
-import sample_lib
 
-HAPPY_PATH_INPUT = SampleInputType(name="Gabe").to_proto()
-HAPPY_PATH_RESPONSE = SampleOutputType(greeting="Hello Gabe").to_proto()
+HAPPY_PATH_INPUT_DM = SampleInputType(name="Gabe")
+HAPPY_PATH_RESPONSE_DM = SampleOutputType(greeting="Hello Gabe")
+HAPPY_PATH_INPUT = HAPPY_PATH_INPUT_DM.to_proto()
+HAPPY_PATH_RESPONSE = HAPPY_PATH_RESPONSE_DM.to_proto()
 
 
 def test_calling_predict_should_raise_if_module_raises(
@@ -89,6 +87,20 @@ def test_global_predict_works_on_good_inputs(
     assert response == HAPPY_PATH_RESPONSE
 
 
+def test_global_predict_predict_model_direct(
+    sample_inference_service, sample_predict_servicer, sample_task_model_id
+):
+    """Test that the direct predict_model function can be called to bypass the
+    gRPC handler
+    """
+    response = sample_predict_servicer.predict_model(
+        request_name="SampleTaskRequest",
+        model_id=sample_task_model_id,
+        sample_input=HAPPY_PATH_INPUT_DM,
+    )
+    assert response == HAPPY_PATH_RESPONSE_DM
+
+
 def test_global_predict_aborts_long_running_predicts(
     sample_inference_service, sample_predict_servicer
 ):