From 065af8a9fa20e5ac74d83275cfdd9ad9919d727e Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Fri, 9 Oct 2026 13:08:18 -0500 Subject: [PATCH] fix(serve): resolve Triton model.py and pack_conda_env.sh from model_server/triton When this code moved from model_server/triton/triton_builder.py into model_builder_utils.py, the Path(__file__).parent lookups for the Triton python-backend model.py and pack_conda_env.sh were not updated. They now resolve to sagemaker/serve/model.py and sagemaker/serve/pack_conda_env.sh, which do not exist, so building a Triton model with an inference_spec failed with FileNotFoundError when copying model.py (and the conda env pack script would not be found either). Point both lookups at model_server/triton, and add the .sh script to the package data: it is currently not included in the published sagemaker-serve wheel (1.24.0), so the path fix alone would not help installed users. --- sagemaker-serve/pyproject.toml | 2 +- .../sagemaker/serve/model_builder_utils.py | 4 +-- .../unit/test_model_builder_utils_triton.py | 34 +++++++++++++++++++ 3 files changed, 37 insertions(+), 3 deletions(-) diff --git a/sagemaker-serve/pyproject.toml b/sagemaker-serve/pyproject.toml index 34c27f755b..764833bb8c 100644 --- a/sagemaker-serve/pyproject.toml +++ b/sagemaker-serve/pyproject.toml @@ -63,7 +63,7 @@ include = ["sagemaker*"] namespaces = true [tool.setuptools.package-data] -"sagemaker.serve" = ["schema/*.json"] +"sagemaker.serve" = ["schema/*.json", "model_server/triton/*.sh"] [tool.setuptools.dynamic] version = { file = "VERSION"} diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py index 0cf1940bd1..82a19b15bf 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py @@ -3003,7 +3003,7 @@ def _pack_conda_env(self, pkl_path: Path): "Please note that you may need to restart your runtime after installation." ) - script_path = Path(__file__).parent.joinpath("pack_conda_env.sh") + script_path = Path(__file__).parent.joinpath("model_server", "triton", "pack_conda_env.sh") env_tar_path = pkl_path.joinpath("triton_env.tar.gz") conda_env_name = os.getenv("CONDA_DEFAULT_ENV") @@ -3162,7 +3162,7 @@ def _prepare_for_triton(self): raise ValueError("%s is not supported" % self.framework) if self.inference_spec: - triton_model_path = Path(__file__).parent.joinpath("model.py") + triton_model_path = Path(__file__).parent.joinpath("model_server", "triton", "model.py") shutil.copy2(str(triton_model_path), str(export_path)) self._generate_config_pbtxt(pkl_path=pkl_path) diff --git a/sagemaker-serve/tests/unit/test_model_builder_utils_triton.py b/sagemaker-serve/tests/unit/test_model_builder_utils_triton.py index 8ff75ad5de..c578dc74d2 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_utils_triton.py +++ b/sagemaker-serve/tests/unit/test_model_builder_utils_triton.py @@ -130,6 +130,27 @@ def test_prepare_for_triton_inference_spec(self, mock_hmac, mock_pack, mock_conf mock_pack.assert_called_once() mock_hmac.assert_called_once() + @patch("shutil.copy2") + @patch.object(_ModelBuilderUtils, "_generate_config_pbtxt") + @patch.object(_ModelBuilderUtils, "_pack_conda_env") + @patch.object(_ModelBuilderUtils, "_compute_integrity_hash") + def test_prepare_for_triton_inference_spec_copies_triton_model_py( + self, mock_hmac, mock_pack, mock_config, mock_copy + ): + """The Triton python backend model.py must be copied from model_server/triton.""" + utils = _ModelBuilderUtils() + utils.inference_spec = Mock() + utils.model = None + utils.schema_builder = Mock() + + with tempfile.TemporaryDirectory() as tmpdir: + utils.model_path = tmpdir + utils._prepare_for_triton() + + source = Path(mock_copy.call_args[0][0]) + self.assertTrue(source.is_file(), f"{source} does not exist") + self.assertEqual(source.parent.name, "triton") + class TestExportPytorchToOnnx(unittest.TestCase): """Test _export_pytorch_to_onnx method.""" @@ -242,6 +263,19 @@ def test_pack_conda_env_no_conda_pack_real(self): with self.assertRaises(ImportError): utils._pack_conda_env(Path(tmpdir)) + @patch("sagemaker.serve.model_builder_utils.subprocess.run") + def test_pack_conda_env_runs_triton_pack_script(self, mock_run): + """The pack_conda_env.sh script must be resolved from model_server/triton.""" + utils = _ModelBuilderUtils() + + with patch.dict("sys.modules", {"conda_pack": Mock(__version__="0.8.0")}): + with tempfile.TemporaryDirectory() as tmpdir: + utils._pack_conda_env(Path(tmpdir)) + + script = Path(mock_run.call_args[0][0][1]) + self.assertTrue(script.is_file(), f"{script} does not exist") + self.assertEqual(script.name, "pack_conda_env.sh") + class TestSaveInferenceSpec(unittest.TestCase): """Test _save_inference_spec method."""