fix(serve): resolve Triton model.py and pack_conda_env.sh from model_server/triton

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

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)

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