diff --git a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py index bebe559b9e..ee56faf4af 100755 --- a/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py +++ b/deploy/mlflow-triton-plugin/mlflow_triton/deployments.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2021-2026, 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 @@ -90,7 +90,7 @@ def create_deployment(self, name, model_uri, flavor=None, config=None): """ Deploy the model at the model_uri to the Triton model repo. Associated config.pbtxt and *labels* files will be deployed. - :param name: Name of the of the model + :param name: Name of the model :param model_uri: Model uri in format model:// :param flavor: Flavor of the deployed model :param config: Configuration parameters @@ -99,6 +99,9 @@ def create_deployment(self, name, model_uri, flavor=None, config=None): """ self._validate_flavor(flavor) + # Validate model name + self._validate_model_name(name) + # Verify model does not already exist in Triton if self._model_exists(name): raise Exception( @@ -513,6 +516,18 @@ def _validate_flavor(self, flavor): if flavor not in self.supported_flavors: raise Exception("{} model flavor not supported by Triton".format(flavor)) + def _validate_model_name(self, name): + # Check if the model name is empty or only contains whitespace, tabs, or newlines + if name.strip() == "": + raise Exception( + "Model name cannot be empty. Please enter a valid name to deploy." + ) + # Path traversal protection + if "/" in name or name == "..": + raise Exception( + "Path traversal is not allowed in model's name: {}".format(name) + ) + def _model_exists(self, name): deploys = self.list_deployments() exists = False diff --git a/qa/L0_mlflow/plugin_test.py b/qa/L0_mlflow/plugin_test.py index a5d87a3c19..e225ed5881 100755 --- a/qa/L0_mlflow/plugin_test.py +++ b/qa/L0_mlflow/plugin_test.py @@ -1,6 +1,6 @@ #!/usr/bin/python -# Copyright 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2026, 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 @@ -33,7 +33,9 @@ import json import unittest +import mlflow.onnx import numpy as np +import onnx import test_util as tu from mlflow.deployments import get_deploy_client @@ -45,7 +47,7 @@ def setUp(self): def _validate_deployment(self, model_name): # create self.client_.create_deployment( - model_name, "models:/{}/1".format(model_name), flavor="onnx" + model_name, f"models:/{model_name}/1", flavor="onnx" ) # list @@ -79,8 +81,6 @@ def _validate_deployment(self, 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" @@ -92,8 +92,6 @@ def test_onnx_flavor(self): def test_onnx_flavor_with_files(self): # Log the ONNX model and additional Triton config file to MLFlow - import mlflow.onnx - import onnx model = onnx.load( "./mlflow-triton-plugin/examples/onnx_float32_int32_int32/1/model.onnx" @@ -116,6 +114,65 @@ def test_onnx_flavor_with_files(self): filecmp.cmp(config_path, "./models/onnx_model_with_files/config.pbtxt") ) + def test_model_name(self): + EMPTY_MODEL_NAMES = [ + "", + " ", + " ", + "\n", + "\t", + "\r", + "\v", + "\f", + ] + INVALID_PATH_TRAVERSAL_NAMES = [ + "/opt/sys/", + "../../etc/passwd", + "../outside/repo", + "test_models/../identity_py", + "..", + ] + VALID_MODEL_NAMES = [ + "model123", + # "model OAI", TRI-769: Fix this test case + "model.version1", + "...", + "..my_model", + "model..1", + "model....1", + ] + + for model_name in EMPTY_MODEL_NAMES: + model_uri = f"models:/{model_name}/1" + with self.assertRaises(Exception) as e: + self.client_.create_deployment(model_name, model_uri, flavor="onnx") + self.assertIn( + "Model name cannot be empty. Please enter a valid name to deploy.", + str(e.exception), + ) + + for model_name in INVALID_PATH_TRAVERSAL_NAMES: + model_uri = f"models:/{model_name}/1" + with self.assertRaises(Exception) as e: + self.client_.create_deployment(model_name, model_uri, flavor="onnx") + self.assertIn( + f"Path traversal is not allowed in model's name: {model_name}", + str(e.exception), + ) + + for model_name in VALID_MODEL_NAMES: + 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=f"{model_name}" + ) + + # Validate deployment functionalities - create, list, get, predict, delete + self._validate_deployment(model_name) + if __name__ == "__main__": unittest.main() diff --git a/qa/L0_mlflow/test.sh b/qa/L0_mlflow/test.sh index 1bf7bc18f8..5e863a8988 100755 --- a/qa/L0_mlflow/test.sh +++ b/qa/L0_mlflow/test.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2022-2026, 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 @@ -187,7 +187,7 @@ if [ $? -ne 0 ]; then echo -e "\n***\n*** Python Test Failed\n***" RET=1 else - check_test_results $TEST_RESULT_FILE 2 + check_test_results $TEST_RESULT_FILE 3 if [ $? -ne 0 ]; then cat $PY_LOG echo -e "\n***\n*** Test Result Verification Failed\n***" @@ -256,7 +256,7 @@ if [ $? -ne 0 ]; then echo -e "\n***\n*** Python Test Failed\n***" RET=1 else - check_test_results $TEST_RESULT_FILE 2 + check_test_results $TEST_RESULT_FILE 3 if [ $? -ne 0 ]; then cat $PY_LOG echo -e "\n***\n*** Test Result Verification Failed\n***"