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."""