test: Python models filtering outputs based on requested outputs (#7338)

triton-inference-server · Jun 17, 2024 · a2a6b6d · a2a6b6d
1 parent 833f2b9
commit a2a6b6d
Show file tree

Hide file tree

Showing 2 changed files with 130 additions and 2 deletions.
diff --git a/qa/L0_backend_python/io/io_test.py b/qa/L0_backend_python/io/io_test.py
@@ -162,6 +162,103 @@ def test_variable_gpu_output(self):
                 self.assertEqual(output.size, i + 1)
                 np.testing.assert_almost_equal(output, np.ones(i + 1) * (i + 1))
 
+    # Non-decoupled models should filter outputs base on requested outputs.
+    def test_requested_output_default(self):
+        model_name = "add_sub"
+        shape = [16]
+
+        input0_data = np.random.rand(*shape).astype(np.float32)
+        input1_data = np.random.rand(*shape).astype(np.float32)
+        inputs = [
+            grpcclient.InferInput(
+                "INPUT0", input0_data.shape, np_to_triton_dtype(input0_data.dtype)
+            ),
+            grpcclient.InferInput(
+                "INPUT1", input1_data.shape, np_to_triton_dtype(input1_data.dtype)
+            ),
+        ]
+        inputs[0].set_data_from_numpy(input0_data)
+        inputs[1].set_data_from_numpy(input1_data)
+
+        # request for output 1, among output 0 and 1.
+        requested_outputs = [grpcclient.InferRequestedOutput("OUTPUT1")]
+        with self._shm_leak_detector.Probe():
+            response = self._client.infer(
+                model_name=model_name,
+                inputs=inputs,
+                outputs=requested_outputs,
+            )
+        outputs = response.get_response().outputs
+        self.assertEqual(len(outputs), len(requested_outputs))
+        output1_data = response.as_numpy("OUTPUT1")
+        self.assertTrue(np.allclose(input0_data - input1_data, output1_data))
+
+        # without requested output should return all outputs
+        with self._shm_leak_detector.Probe():
+            response = self._client.infer(model_name=model_name, inputs=inputs)
+        outputs = response.get_response().outputs
+        self.assertEqual(len(outputs), len(inputs))
+        output0_data = response.as_numpy("OUTPUT0")
+        output1_data = response.as_numpy("OUTPUT1")
+        self.assertTrue(np.allclose(input0_data + input1_data, output0_data))
+        self.assertTrue(np.allclose(input0_data - input1_data, output1_data))
+
+    # Decoupled models should filter outputs base on requested outputs.
+    def test_requested_output_decoupled(self):
+        model_name = "dlpack_io_identity_decoupled"
+        shape = [4]
+        expected_response_repeat = 2
+
+        input0_data = np.random.rand(*shape).astype(np.float32)
+        gpu_output_data = np.random.rand(*shape).astype(np.bool_)
+        inputs = [
+            grpcclient.InferInput(
+                "INPUT0", input0_data.shape, np_to_triton_dtype(input0_data.dtype)
+            ),
+            grpcclient.InferInput(
+                "GPU_OUTPUT",
+                gpu_output_data.shape,
+                np_to_triton_dtype(gpu_output_data.dtype),
+            ),
+        ]
+        inputs[0].set_data_from_numpy(input0_data)
+        inputs[1].set_data_from_numpy(gpu_output_data)
+
+        # request for output 0, among output 0 and next gpu output.
+        requested_outputs = [grpcclient.InferRequestedOutput("OUTPUT0")]
+        user_data = UserData()
+        with grpcclient.InferenceServerClient(f"{_tritonserver_ipaddr}:8001") as client:
+            client.start_stream(callback=partial(callback, user_data))
+            client.async_stream_infer(
+                model_name=model_name, inputs=inputs, outputs=requested_outputs
+            )
+            client.stop_stream()
+        for _ in range(expected_response_repeat):
+            self.assertFalse(user_data._completed_requests.empty())
+            response = user_data._completed_requests.get()
+            outputs = response.get_response().outputs
+            self.assertEqual(len(outputs), len(requested_outputs))
+            output0_data = response.as_numpy("OUTPUT0")
+            self.assertTrue(np.allclose(input0_data, output0_data))
+        self.assertTrue(user_data._completed_requests.empty())
+
+        # without requested output should return all outputs
+        user_data = UserData()
+        with grpcclient.InferenceServerClient(f"{_tritonserver_ipaddr}:8001") as client:
+            client.start_stream(callback=partial(callback, user_data))
+            client.async_stream_infer(model_name=model_name, inputs=inputs)
+            client.stop_stream()
+        for _ in range(expected_response_repeat):
+            self.assertFalse(user_data._completed_requests.empty())
+            response = user_data._completed_requests.get()
+            outputs = response.get_response().outputs
+            self.assertEqual(len(outputs), len(inputs))
+            output0_data = response.as_numpy("OUTPUT0")
+            next_gpu_output_data = response.as_numpy("NEXT_GPU_OUTPUT")
+            self.assertTrue(np.allclose(input0_data, output0_data))
+            self.assertTrue(np.allclose(gpu_output_data[1:], next_gpu_output_data))
+        self.assertTrue(user_data._completed_requests.empty())
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/qa/L0_backend_python/io/test.sh b/qa/L0_backend_python/io/test.sh
@@ -106,7 +106,7 @@ fi
 
 set +e
 SUBTEST="test_empty_gpu_output"
-python3 -m pytest --junitxml=${SUBTEST}.${TRIAL}.report.xml ${UNITTEST_PY}::IOTest::${SUBTEST} > ${CLIENT_LOG}.${SUBTEST}
+python3 -m pytest --junitxml=${SUBTEST}.report.xml ${UNITTEST_PY}::IOTest::${SUBTEST} > ${CLIENT_LOG}.${SUBTEST}
 
 if [ $? -ne 0 ]; then
     echo -e "\n***\n*** IOTest.${SUBTEST} FAILED. \n***"
@@ -133,7 +133,7 @@ fi
 
 set +e
 SUBTEST="test_variable_gpu_output"
-python3 -m pytest --junitxml=${SUBTEST}.${TRIAL}.report.xml ${UNITTEST_PY}::IOTest::${SUBTEST} > ${CLIENT_LOG}.${SUBTEST}
+python3 -m pytest --junitxml=${SUBTEST}.report.xml ${UNITTEST_PY}::IOTest::${SUBTEST} > ${CLIENT_LOG}.${SUBTEST}
 
 if [ $? -ne 0 ]; then
     echo -e "\n***\n*** IOTest.${SUBTEST} FAILED. \n***"
@@ -145,6 +145,37 @@ set -e
 kill $SERVER_PID
 wait $SERVER_PID
 
+# IOTest.test_requested_output_default & IOTest.test_requested_output_decoupled
+rm -rf models && mkdir models
+mkdir -p models/add_sub/1/
+cp ../../python_models/add_sub/model.py ./models/add_sub/1/
+cp ../../python_models/add_sub/config.pbtxt ./models/add_sub/
+mkdir -p models/dlpack_io_identity_decoupled/1/
+cp ../../python_models/dlpack_io_identity_decoupled/model.py ./models/dlpack_io_identity_decoupled/1/
+cp ../../python_models/dlpack_io_identity_decoupled/config.pbtxt ./models/dlpack_io_identity_decoupled/
+
+run_server
+if [ "$SERVER_PID" == "0" ]; then
+    echo -e "\n***\n*** Failed to start $SERVER\n***"
+    cat $SERVER_LOG
+    RET=1
+fi
+
+SUBTESTS="test_requested_output_default test_requested_output_decoupled"
+for SUBTEST in $SUBTESTS; do
+    set +e
+    python3 -m pytest --junitxml=${SUBTEST}.report.xml ${UNITTEST_PY}::IOTest::${SUBTEST} > ${CLIENT_LOG}.${SUBTEST}
+    if [ $? -ne 0 ]; then
+        echo -e "\n***\n*** IOTest.${SUBTEST} FAILED. \n***"
+        cat $CLIENT_LOG.${SUBTEST}
+        RET=1
+    fi
+    set -e
+done
+
+kill $SERVER_PID
+wait $SERVER_PID
+
 if [ $RET -eq 0 ]; then
     echo -e "\n***\n*** IO test PASSED.\n***"
 else