Add testing for Triton Client Plugin API

qa/L0_grpc/client_plugin_models/client_plugin_test/1/model.py

@@ -0,0 +1,58 @@
+# Copyright (c) 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.
+
+import json
+import triton_python_backend_utils as pb_utils
+import numpy as np
+
+
+class TritonPythonModel:
+
+    def execute(self, requests):
+        responses = []
+
+        for request in requests:
+            json_string = pb_utils.get_input_tensor_by_name(
+                request, "EXPECTED_HEADERS").as_numpy()[0].decode("utf-8")
+            expected_headers = json.loads(json_string)
+
+            success = True
+            if request.parameters() != '':
+                parameters = json.loads(request.parameters())
+                for key, value in expected_headers.items():
+                    if key in parameters:
+                        if parameters[key] != value:
+                            success = False
+                    else:
+                        success = False
+
+            test_success = pb_utils.Tensor("TEST_SUCCESS",
+                                           np.array([success], dtype=bool))
+            inference_response = pb_utils.InferenceResponse(
+                output_tensors=[test_success])
+            responses.append(inference_response)
+
+        return responses

qa/L0_grpc/client_plugin_models/client_plugin_test/config.pbtxt

@@ -0,0 +1,45 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. 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.
+
+name: "client_plugin_test"
+backend: "python"
+
+input [
+  {
+    name: "EXPECTED_HEADERS"
+    data_type: TYPE_STRING
+    dims: [ 1 ]
+  }
+]
+output [
+  {
+    name: "TEST_SUCCESS"
+    data_type: TYPE_BOOL
+    dims: [ 1 ]
+  }
+]
+
+instance_group [{ kind: KIND_CPU }]

qa/L0_grpc/grpc_client_plugin_test.py

@@ -0,0 +1,120 @@
+#!/usr/bin/python
+# 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.
+import sys
+import json
+
+sys.path.append("../common")
+
+import unittest
+import numpy as np
+import test_util as tu
+import tritonclient.grpc as tritongrpcclient
+from tritonclient.grpc import InferenceServerClientPlugin
+from tritonclient.utils import np_to_triton_dtype
+import tritonclient.grpc.aio as asynctritongrpcclient
+
+
+# A simple plugin that adds headers to the inference request.
+class TestPlugin(InferenceServerClientPlugin):
+
+    def __init__(self, headers):
+        self._headers = headers
+
+    def __call__(self, request):
+        request.headers.update(self._headers)
+
+
+def prepare_infer_inputs(headers):
+    expected_headers = np.array([json.dumps(headers)], dtype=object)
+    inputs = []
+    inputs.append(
+        tritongrpcclient.InferInput('EXPECTED_HEADERS', expected_headers.shape,
+                                    np_to_triton_dtype(expected_headers.dtype)))
+    inputs[0].set_data_from_numpy(expected_headers)
+
+    return inputs
+
+
+class GRPCClientPluginAsyncTest(unittest.IsolatedAsyncioTestCase):
+
+    async def asyncSetUp(self):
+        self._headers = {'my-key': 'my-value'}
+        self._plugin = TestPlugin(self._headers)
+        self._client = asynctritongrpcclient.InferenceServerClient(
+            url='localhost:8001')
+
+    async def test_simple_infer(self):
+        model = "client_plugin_test"
+        inputs = prepare_infer_inputs(self._headers)
+        self._client.register_plugin(self._plugin)
+        response = await self._client.infer(model_name=model, inputs=inputs)
+        test_success = response.as_numpy('TEST_SUCCESS')
+        self.assertEqual(test_success, True)
+
+        self._client.unregister_plugin()
+        inputs = prepare_infer_inputs({})
+        response = await self._client.infer(model_name=model, inputs=inputs)
+        test_success = response.as_numpy('TEST_SUCCESS')
+        self.assertEqual(test_success, True)
+
+    async def asyncTearDown(self):
+        await self._client.close()
+
+
+class GRPCClientPluginTest(tu.TestResultCollector):
+
+    def setUp(self):
+        self._headers = {'my-key': 'my-value'}
+        self._plugin = TestPlugin(self._headers)
+        self._client = tritongrpcclient.InferenceServerClient(
+            url='localhost:8001')
+
+    def test_simple_infer(self):
+        # Set the binary data to False so that 'Inference-Header-Length' is not
+        # added to the headers.
+        model = "client_plugin_test"
+        inputs = prepare_infer_inputs(self._headers)
+        self._client.register_plugin(self._plugin)
+        self.assertEqual(self._plugin, self._client.plugin())
+        response = self._client.infer(model_name=model, inputs=inputs)
+        test_success = response.as_numpy('TEST_SUCCESS')
+        self.assertEqual(test_success, True)
+
+        # Unregister the plugin
+        inputs = prepare_infer_inputs({})
+        self._client.unregister_plugin()
+        self.assertEqual(None, self._client.plugin())
+        response = self._client.infer(model_name=model, inputs=inputs)
+        test_success = response.as_numpy('TEST_SUCCESS')
+        self.assertEqual(test_success, True)
+
+    def tearDown(self):
+        self._client.close()
+
+
+if __name__ == '__main__':
+    unittest.main()

qa/L0_grpc/test.sh

@@ -42,12 +42,14 @@ export CUDA_VISIBLE_DEVICES=0
 
 RET=0
 
+CLIENT_PLUGIN_TEST="./grpc_client_plugin_test.py"
 # On windows the paths invoked by the script (running in WSL) must use
 # /mnt/c when needed but the paths on the tritonserver command-line
 # must be C:/ style.
 if [[ "$(< /proc/sys/kernel/osrelease)" == *microsoft* ]]; then
     SDKDIR=${SDKDIR:=C:/sdk}
     MODELDIR=${MODELDIR:=C:/models}
+    CLIENT_PLUGIN_MODELDIR=${MODELDIR:=C:/client_plugin_models}
     DATADIR=${DATADIR:="/mnt/c/data/inferenceserver/${REPO_VERSION}"}
     BACKEND_DIR=${BACKEND_DIR:=C:/tritonserver/backends}
     SERVER=${SERVER:=/mnt/c/tritonserver/bin/tritonserver.exe}
@@ -93,6 +95,7 @@ if [[ "$(< /proc/sys/kernel/osrelease)" == *microsoft* ]]; then
     CC_UNIT_TEST=${SDKDIR}/python/cc_client_test
 else
     MODELDIR=${MODELDIR:=`pwd`/models}
+    CLIENT_PLUGIN_MODELDIR=${CLIENTPLUGINMODELDIR:=`pwd`/client_plugin_models}
     DATADIR=${DATADIR:="/data/inferenceserver/${REPO_VERSION}"}
     TRITON_DIR=${TRITON_DIR:="/opt/tritonserver"}
     SERVER=${TRITON_DIR}/bin/tritonserver
@@ -334,6 +337,24 @@ set -e
 kill $SERVER_PID
 wait $SERVER_PID
 
+SERVER_ARGS="--backend-directory=${BACKEND_DIR} --model-repository=${CLIENT_PLUGIN_MODELDIR} --http-header-forward-pattern=.* --grpc-header-forward-pattern=.*"
+run_server
+if [ "$SERVER_PID" == "0" ]; then
+    echo -e "\n***\n*** Failed to start $SERVER\n***"
+    cat $SERVER_LOG
+    exit 1
+fi
+
+set +e
+python3 $CLIENT_PLUGIN_TEST >> ${CLIENT_LOG}.python.plugin 2>&1
+if [ $? -ne 0 ]; then
+    cat ${CLIENT_LOG}.python.plugin
+    RET=1
+fi
+set -e
+kill $SERVER_PID
+wait $SERVER_PID
+
 export GRPC_TRACE=compression, channel
 export GRPC_VERBOSITY=DEBUG
 SERVER_ARGS="--backend-directory=${BACKEND_DIR} --model-repository=${MODELDIR} --grpc-infer-response-compression-level=high"