From 2874ad96a5227abbaaa1594a030ed38ebbd2ddde Mon Sep 17 00:00:00 2001 From: GuanLuo Date: Fri, 3 Jun 2022 17:50:02 -0700 Subject: [PATCH 1/3] Add test for different variance of ONNX flavor --- .../onnx_float32_int32_int32/config.pbtxt | 3 +- qa/L0_mlflow/plugin_test.py | 61 +++++++++++++------ qa/L0_mlflow/test.sh | 11 +++- 3 files changed, 53 insertions(+), 22 deletions(-) diff --git a/deploy/mlflow-triton-plugin/examples/onnx_float32_int32_int32/config.pbtxt b/deploy/mlflow-triton-plugin/examples/onnx_float32_int32_int32/config.pbtxt index 74cc2a6c78..75ea016cfa 100755 --- a/deploy/mlflow-triton-plugin/examples/onnx_float32_int32_int32/config.pbtxt +++ b/deploy/mlflow-triton-plugin/examples/onnx_float32_int32_int32/config.pbtxt @@ -1,5 +1,5 @@ -# Copyright 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2021-2022, 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 @@ -24,7 +24,6 @@ # 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: "onnx_float32_int32_int32" platform: "onnxruntime_onnx" max_batch_size: 8 version_policy: { latest { num_versions: 1 }} diff --git a/qa/L0_mlflow/plugin_test.py b/qa/L0_mlflow/plugin_test.py index bed892bb54..3522fad563 100644 --- a/qa/L0_mlflow/plugin_test.py +++ b/qa/L0_mlflow/plugin_test.py @@ -42,31 +42,20 @@ class PluginTest(tu.TestResultCollector): def setUp(self): self.client_ = get_deploy_client('triton') - def test_onnx_flavor(self): - # Log the ONNX model to MLFlow - import mlflow.onnx - import onnx - model = onnx.load( - "./mlflow-triton-plugin/examples/onnx_float32_int32_int32/1/model.onnx" - ) - # Use a different name to ensure the plugin operates on correct model - mlflow.onnx.log_model(model, - "triton", - registered_model_name="onnx_model") - + def _validate_deployment(self, model_name): # create - self.client_.create_deployment("onnx_model", - "models:/onnx_model/1", + self.client_.create_deployment(model_name, + "models:/{}/1".format(model_name), flavor="onnx") # list deployment_list = self.client_.list_deployments() self.assertEqual(len(deployment_list), 1) - self.assertEqual(deployment_list[0]['name'], "onnx_model") + self.assertEqual(deployment_list[0]['name'], model_name) # get - deployment = self.client_.get_deployment("onnx_model") - self.assertEqual(deployment['name'], "onnx_model") + deployment = self.client_.get_deployment(model_name) + self.assertEqual(deployment['name'], model_name) # predict inputs = {} @@ -75,7 +64,7 @@ def test_onnx_flavor(self): for key, value in input_json['inputs'].items(): inputs[key] = np.array(value, dtype=np.float32) - output = self.client_.predict("onnx_model", inputs) + output = self.client_.predict(model_name, inputs) with open("./mlflow-triton-plugin/examples/expected_output.json", "r") as f: output_json = json.load(f) @@ -86,7 +75,41 @@ def test_onnx_flavor(self): err_msg='Inference result is not correct') # delete - self.client_.delete_deployment("onnx_model") + self.client_.delete_deployment(model_name) + + def test_onnx_flavor(self): + # Log the ONNX model to MLFlow + import mlflow.onnx + import onnx + model = onnx.load( + "./mlflow-triton-plugin/examples/onnx_float32_int32_int32/1/model.onnx" + ) + # Use a different name to ensure the plugin operates on correct model + mlflow.onnx.log_model(model, + "triton", + registered_model_name="onnx_model") + + self._validate_deployment("onnx_model") + + def test_onnx_flavor_with_files(self): + # Log the ONNX model to MLFlow + import mlflow.onnx + import onnx + model = onnx.load( + "./mlflow-triton-plugin/examples/onnx_float32_int32_int32/1/model.onnx" + ) + config_path = "./mlflow-triton-plugin/examples/onnx_float32_int32_int32/config.pbtxt" + # Use a different name to ensure the plugin operates on correct model + mlflow.onnx.log_model(model, + "triton", + registered_model_name="onnx_model_with_files") + mlflow.log_artifact(config_path, "triton") + + self._validate_deployment("onnx_model_with_files") + + # Check if the additional files are properly copied + import filecmp + filecmp.cmp(config_path, "./models/onnx_model_with_files/config.pbtxt") if __name__ == '__main__': diff --git a/qa/L0_mlflow/test.sh b/qa/L0_mlflow/test.sh index 13f04a8d1a..74c9348f1d 100644 --- a/qa/L0_mlflow/test.sh +++ b/qa/L0_mlflow/test.sh @@ -45,6 +45,14 @@ mkdir -p ./mlflow/artifacts pip install ./mlflow-triton-plugin/ +# Clear mlflow registered models if any +python - << EOF +from mlflow.tracking import MlflowClient +c = MlflowClient() +for m in c.list_registered_models(): + c.delete_registered_model(m.name) +EOF + rm -rf ./models mkdir -p ./models SERVER=/opt/tritonserver/bin/tritonserver @@ -137,6 +145,7 @@ if [ $CLI_RET -ne 0 ]; then fi set -e +# ONNX flavor with Python package set +e PY_LOG=plugin_py.log PY_TEST=plugin_test.py @@ -147,7 +156,7 @@ if [ $? -ne 0 ]; then echo -e "\n***\n*** Python Test Failed\n***" RET=1 else - check_test_results $TEST_RESULT_FILE 1 + check_test_results $TEST_RESULT_FILE 2 if [ $? -ne 0 ]; then cat $PY_LOG echo -e "\n***\n*** Test Result Verification Failed\n***" From 08f7a646721546456135daff8e5d8275cd8fb2d0 Mon Sep 17 00:00:00 2001 From: GuanLuo Date: Fri, 3 Jun 2022 17:51:55 -0700 Subject: [PATCH 2/3] Fix up --- qa/L0_mlflow/plugin_test.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/qa/L0_mlflow/plugin_test.py b/qa/L0_mlflow/plugin_test.py index 3522fad563..43a18d80a8 100644 --- a/qa/L0_mlflow/plugin_test.py +++ b/qa/L0_mlflow/plugin_test.py @@ -109,7 +109,9 @@ def test_onnx_flavor_with_files(self): # Check if the additional files are properly copied import filecmp - filecmp.cmp(config_path, "./models/onnx_model_with_files/config.pbtxt") + self.assertTrue( + filecmp.cmp(config_path, + "./models/onnx_model_with_files/config.pbtxt")) if __name__ == '__main__': From 2e6f14cde852fb4b75023d80e373e5222e010257 Mon Sep 17 00:00:00 2001 From: GuanLuo Date: Mon, 6 Jun 2022 14:38:10 -0700 Subject: [PATCH 3/3] Address comment --- qa/L0_mlflow/plugin_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/qa/L0_mlflow/plugin_test.py b/qa/L0_mlflow/plugin_test.py index 43a18d80a8..8dbf9d9146 100644 --- a/qa/L0_mlflow/plugin_test.py +++ b/qa/L0_mlflow/plugin_test.py @@ -92,7 +92,7 @@ def test_onnx_flavor(self): self._validate_deployment("onnx_model") def test_onnx_flavor_with_files(self): - # Log the ONNX model to MLFlow + # Log the ONNX model and additional Triton config file to MLFlow import mlflow.onnx import onnx model = onnx.load(