Add Triton Client Plugin API for HTTP/GRPC clients (#301)

* Add Triton Client Plugin API for HTTP/GRPC clients

* Fix up

* Add async client plugin support

* Fix GRPC client plugin bugs

* Add docstring

* Review feedback

* Add request.py

* Add a method to fetch the registered plugin

src/python/library/build_wheel.py

@@ -103,7 +103,14 @@ def sed(pattern, replace, source, dest=None):
     print("Adding package files")
 
     mkdir(os.path.join(FLAGS.whl_dir, 'tritonclient'))
-    touch(os.path.join(FLAGS.whl_dir, 'tritonclient/__init__.py'))
+    shutil.copy('tritonclient/__init__.py',
+                os.path.join(FLAGS.whl_dir, 'tritonclient'))
+    shutil.copy('tritonclient/_client.py',
+                os.path.join(FLAGS.whl_dir, 'tritonclient'))
+    shutil.copy('tritonclient/_plugin.py',
+                os.path.join(FLAGS.whl_dir, 'tritonclient'))
+    shutil.copy('tritonclient/_request.py',
+                os.path.join(FLAGS.whl_dir, 'tritonclient'))
 
     # Needed for backwards-compatibility; remove when moving
     # completely to the new structure.
@@ -199,7 +206,7 @@ def sed(pattern, replace, source, dest=None):
         if os.uname().machine == "aarch64":
             platform_name = "manylinux2014_aarch64"
         elif os.uname().machine == "ppc64le":
-            platform_name = "manylinux2014_ppc64le" 
+            platform_name = "manylinux2014_ppc64le"
         else:
             platform_name = "manylinux1_x86_64"
         args = [

src/python/library/tests/test_inference_server_client.py

@@ -33,7 +33,6 @@
 from unittest.mock import patch, MagicMock
 import rapidjson
 from tritonclient.http import *
-from tritonclient.http import _raise_if_error
 from tritonclient.utils import *
 
 json_error_response = """{
@@ -73,7 +72,7 @@ def test_get_method_failure(self):
         """
 
         with self.assertRaises(InferenceServerException):
-            _raise_if_error(self.response)
+            raise_if_error(self.response)
 
     @patch('tritonclient.http.InferenceServerClient._post',
            MagicMock(return_value={"status_code": 200}))
@@ -98,4 +97,4 @@ def test_error_plain_text(self):
         """
 
         with self.assertRaises(InferenceServerException):
-            _raise_if_error(self.response)
+            raise_if_error(self.response)

src/python/library/tritonclient/CMakeLists.txt

@@ -34,4 +34,9 @@ if(${TRITON_ENABLE_PYTHON_GRPC})
   file(COPY grpc DESTINATION .)
 endif() # TRITON_ENABLE_PYTHON_GRPC
 
+file(COPY _client.py DESTINATION .)
+file(COPY _plugin.py DESTINATION .)
+file(COPY __init__.py DESTINATION .)
+file(COPY _request.py DESTINATION .)
+
 add_subdirectory(utils)

src/python/library/tritonclient/__init__.py

@@ -0,0 +1,25 @@
+# Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

src/python/library/tritonclient/_client.py

@@ -0,0 +1,82 @@
+# Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+from tritonclient.utils import raise_error
+
+
+class InferenceServerClientBase:
+
+    def __init__(self):
+        self._plugin = None
+
+    def _call_plugin(self, request):
+        """Called by the subclasses before sending a request to the
+        network.
+        """
+        if self._plugin != None:
+            self._plugin(request)
+
+    def register_plugin(self, plugin):
+        """Register a Client Plugin.
+
+        Parameters
+        ----------
+        plugin : InferenceServerClientPlugin
+            A client plugin
+        
+        Raises
+        ------
+        InferenceServerException
+            If a plugin is already registered.
+        """
+
+        if self._plugin is None:
+            self._plugin = plugin
+        else:
+            raise_error("A plugin is already registered. Please "
+                        "unregister the previous plugin first before"
+                        " registering a new plugin.")
+
+    def plugin(self):
+        """Retrieve the registered plugin if any.
+
+        Returns
+        ------
+        InferenceServerClientPlugin or None
+        """
+        return self._plugin
+
+    def unregister_plugin(self):
+        """Unregister a plugin.
+
+        Raises
+        ------
+        InferenceServerException
+            If no plugin has been registered.
+        """
+        if self._plugin is None:
+            raise_error("No plugin has been registered.")
+
+        self._plugin = None

src/python/library/tritonclient/_plugin.py

@@ -0,0 +1,44 @@
+# Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+from abc import ABC, abstractmethod
+
+
+class InferenceServerClientPlugin(ABC):
+    """Every Triton Client Plugin should extend this class.
+    Each plugin needs to implement the `__call__` method.
+    """
+
+    @abstractmethod
+    def __call__(self, request):
+        """This method will be called when any of the client functions are
+        invoked. Note that the request object must be modfied in-place.
+
+        Parameters
+        ----------
+        request : Request
+            The request object.
+        """
+        pass

src/python/library/tritonclient/_request.py

@@ -0,0 +1,36 @@
+# Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+class Request:
+    """A request object.
+
+    Attributes
+    ----------
+    headers : dict
+        A dictionary containing the request headers.
+    """
+
+    def __init__(self, headers):
+        self.headers = headers if headers is not None else {}

src/python/library/tritonclient/grpc/__init__.py

@@ -37,6 +37,7 @@
     from ._client import MAX_GRPC_MESSAGE_SIZE
     from tritonclient.utils import *
     from ._utils import raise_error, raise_error_grpc
+    from .._plugin import InferenceServerClientPlugin
 except ModuleNotFoundError as error:
     raise RuntimeError(
         'The installation does not include grpc support. '