diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 556525a45e..07e7c18268 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -1944,7 +1944,14 @@ def _save_model_inference_spec(self) -> None: code_path = Path(self.model_path).joinpath("code") if self.inference_spec: - save_pkl(code_path, (self.inference_spec, self.schema_builder)) + if isinstance(self.inference_spec, (CustomOrchestrator, AsyncCustomOrchestrator)): + # The SMD orchestrator entrypoint discards the second element of the + # tuple, and a SchemaBuilder pickles with references to sagemaker.core + # modules that an older sagemaker in the serving container cannot + # import. Leave it out so serve.pkl stays loadable there. + save_pkl(code_path, (self.inference_spec, None)) + else: + save_pkl(code_path, (self.inference_spec, self.schema_builder)) elif self.model: if isinstance(self.model, str): self.framework = None diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py index d4f2e0cc37..995205e288 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py @@ -180,9 +180,7 @@ def _build_for_torchserve(self) -> Model: # Prepare TorchServe artifacts for local container mode if self.mode == Mode.LOCAL_CONTAINER and self.model_path: - # NOTE: prepare_* annotated -> str but returns None (latent bug; fix out of scope) - # pylint: disable-next=assignment-from-no-return - self.secret_key = prepare_for_torchserve( + prepare_for_torchserve( model_path=self.model_path, shared_libs=self.shared_libs, dependencies=self.dependencies, @@ -191,9 +189,7 @@ def _build_for_torchserve(self) -> Model: inference_spec=self.inference_spec, ) if self.mode == Mode.SAGEMAKER_ENDPOINT and self.model_path: - # NOTE: prepare_* annotated -> str but returns None (latent bug; fix out of scope) - # pylint: disable-next=assignment-from-no-return - self.secret_key = prepare_for_torchserve( + prepare_for_torchserve( model_path=self.model_path, shared_libs=self.shared_libs, dependencies=self.dependencies, @@ -666,9 +662,7 @@ def _build_for_tensorflow_serving(self) -> Model: raise ValueError("image_uri is required for TensorFlow Serving deployment") # Prepare TensorFlow Serving artifacts for local container mode - # NOTE: prepare_* annotated -> str but returns None (latent bug; fix out of scope) - # pylint: disable-next=assignment-from-no-return - self.secret_key = prepare_for_tf_serving( + prepare_for_tf_serving( model_path=self.model_path, shared_libs=self.shared_libs, dependencies=self.dependencies, @@ -788,9 +782,7 @@ def _build_for_smd(self) -> Model: cpu_or_gpu = self._get_processing_unit() self.image_uri = self._get_smd_image_uri(processing_unit=cpu_or_gpu) - # NOTE: prepare_* annotated -> str but returns None (latent bug; fix out of scope) - # pylint: disable-next=assignment-from-no-return - self.secret_key = prepare_for_smd( + prepare_for_smd( model_path=self.model_path, shared_libs=self.shared_libs, dependencies=self.dependencies, @@ -834,9 +826,7 @@ def _build_for_transformers(self) -> Model: self._create_conda_env() if self.mode in [Mode.LOCAL_CONTAINER] and self.model_path: - # NOTE: prepare_* annotated -> str but returns None (latent bug; fix out of scope) - # pylint: disable-next=assignment-from-no-return - self.secret_key = prepare_for_mms( + prepare_for_mms( model_path=self.model_path, shared_libs=self.shared_libs, dependencies=self.dependencies, @@ -845,9 +835,7 @@ def _build_for_transformers(self) -> Model: inference_spec=self.inference_spec, ) if self.mode == Mode.SAGEMAKER_ENDPOINT and self.model_path: - # NOTE: prepare_* annotated -> str but returns None (latent bug; fix out of scope) - # pylint: disable-next=assignment-from-no-return - self.secret_key = prepare_for_mms( + prepare_for_mms( model_path=self.model_path, shared_libs=self.shared_libs, dependencies=self.dependencies, diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/prepare.py b/sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/prepare.py index 51b987c0ea..2895c16015 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/prepare.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/prepare.py @@ -81,18 +81,20 @@ def prepare_for_mms( session: Session, image_uri: str, inference_spec: InferenceSpec = None, -) -> str: - """Prepares for InferenceSpec, writes inference.py, and captures dependencies for secret_key. +) -> None: + """Prepares for InferenceSpec, writes inference.py, and captures dependencies. - Args:to + Nothing is returned: the SHA-256 digest of ``serve.pkl`` is recorded in + ``metadata.json`` and verified against it in the container, so there is no secret to + hand back to the caller. + + Args: model_path (str) : Argument shared_libs (List[]) : Argument dependencies (dict) : Argument session (Session) : Argument inference_spec (InferenceSpec, optional) : Argument (default is None) - Returns: - ( str ) : secret_key """ model_path = Path(model_path) if not model_path.exists(): diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/smd/custom_execution_inference.py b/sagemaker-serve/src/sagemaker/serve/model_server/smd/custom_execution_inference.py index f53677fc69..f8b4bf1660 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/smd/custom_execution_inference.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/smd/custom_execution_inference.py @@ -10,24 +10,40 @@ # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. -"""This module is for SageMaker inference.py.""" +"""This module is for SageMaker inference.py. + +This file is copied into the model artifacts as ``code/inference.py`` and runs inside +the SageMaker Distribution serving container, whose installed ``sagemaker`` version is +set by the image rather than by the SDK that built the artifacts. It therefore must not +import from ``sagemaker``: the integrity check is inlined below instead of imported from +``sagemaker.serve.validations.check_integrity`` so that it cannot drift against the +container's copy of that module. +""" from __future__ import absolute_import import asyncio +import hashlib +import hmac +import json +import logging import os import platform -import cloudpickle -import logging from pathlib import Path -from sagemaker.serve.validations.check_integrity import perform_integrity_check +import cloudpickle logger = LOGGER = logging.getLogger("sagemaker") +DEFAULT_CODE_DIRECTORY = "/opt/ml/model/code" + + +def _get_code_directory() -> Path: + """Returns the directory holding serve.pkl and metadata.json.""" + return Path(os.getenv("SAGEMAKER_INFERENCE_CODE_DIRECTORY") or DEFAULT_CODE_DIRECTORY) + def initialize_custom_orchestrator(): """Initializes the custom orchestrator.""" - code_dir = os.getenv("SAGEMAKER_INFERENCE_CODE_DIRECTORY", None) - serve_path = Path(code_dir).joinpath("serve.pkl") + serve_path = _get_code_directory().joinpath("serve.pkl") with open(str(serve_path), mode="rb") as pkl_file: return cloudpickle.load(pkl_file) @@ -48,12 +64,32 @@ def _py_vs_parity_check(): ) +def _perform_integrity_check(buffer: bytes, metadata_path: Path): + """Validates the integrity of bytes by comparing the SHA-256 hash value.""" + if not Path.exists(metadata_path): + raise ValueError("Path to metadata.json does not exist") + + with open(str(metadata_path), "rb") as md: + try: + expected_hash_value = json.loads(md.read()).get("sha256_hash") + except json.JSONDecodeError: + raise ValueError("Corrupt metadata file. It is not a valid json file.") + + if not expected_hash_value: + raise ValueError("metadata.json does not contain a sha256_hash value") + + actual_hash_value = hashlib.sha256(buffer).hexdigest() + + if not hmac.compare_digest(expected_hash_value, actual_hash_value): + raise ValueError("Integrity check for the serialized function or data failed.") + + def _pickle_file_integrity_check(): - with open("/opt/ml/model/code/serve.pkl", "rb") as f: + code_directory = _get_code_directory() + with open(str(code_directory.joinpath("serve.pkl")), "rb") as f: buffer = f.read() - metadata_path = Path("/opt/ml/model/code/metadata.json") - perform_integrity_check(buffer=buffer, metadata_path=metadata_path) + _perform_integrity_check(buffer=buffer, metadata_path=code_directory.joinpath("metadata.json")) _run_preflight_diagnostics() diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.py b/sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.py index f29b8ebcbd..f2ccd5d3d2 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.py @@ -23,19 +23,20 @@ def prepare_for_smd( shared_libs: List[str], dependencies: dict, inference_spec: InferenceSpec = None, -) -> str: +) -> None: """Prepares artifacts for SageMaker model deployment. - Args:to + Writes the inference entrypoint, dependencies and a ``metadata.json`` holding the + SHA-256 digest of ``serve.pkl``. Nothing is returned: the digest is verified in the + container against ``metadata.json``, so there is no secret to hand back to the + caller. + + Args: model_path (str) : Argument shared_libs (List[]) : Argument dependencies (dict) : Argument inference_spec (InferenceSpec, optional) : Argument (default is None) - - Returns: - ( str ) : - """ model_path = Path(model_path) if not model_path.exists(): diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py b/sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py index d56d0ec7bd..0244939b96 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py @@ -20,16 +20,17 @@ def prepare_for_tf_serving( model_path: str, shared_libs: List[str], dependencies: Dict[str, Any], -) -> str: +) -> None: """Prepares the model for serving. + Nothing is returned: the SHA-256 digest of ``serve.pkl`` is recorded in + ``metadata.json`` and verified against it in the container, so there is no secret to + hand back to the caller. + Args: model_path (str): Path to the model directory. shared_libs (List[str]): List of shared libraries. dependencies (Dict[str, Any]): Dictionary of dependencies. - - Returns: - str: Secret key. """ _model_path = Path(model_path) diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/torchserve/prepare.py b/sagemaker-serve/src/sagemaker/serve/model_server/torchserve/prepare.py index ad053d25c9..c92840dd79 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/torchserve/prepare.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/torchserve/prepare.py @@ -26,20 +26,20 @@ def prepare_for_torchserve( session: Session, image_uri: str, inference_spec: InferenceSpec = None, -) -> str: +) -> None: """This is a one-line summary of the function. - Args:to + Nothing is returned: the SHA-256 digest of ``serve.pkl`` is recorded in + ``metadata.json`` and verified against it in the container, so there is no secret to + hand back to the caller. + + Args: model_path (str) : Argument shared_libs (List[]) : Argument dependencies (dict) : Argument session (Session) : Argument inference_spec (InferenceSpec, optional) : Argument (default is None) - - Returns: - ( str ) : - """ model_path = Path(model_path) if not model_path.exists(): diff --git a/sagemaker-serve/tests/unit/model_server/test_smd_custom_execution_inference.py b/sagemaker-serve/tests/unit/model_server/test_smd_custom_execution_inference.py new file mode 100644 index 0000000000..dffaf70e52 --- /dev/null +++ b/sagemaker-serve/tests/unit/model_server/test_smd_custom_execution_inference.py @@ -0,0 +1,215 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for the SMD custom orchestrator entrypoint (custom_execution_inference.py). + +The module under test is copied into the model artifacts as ``code/inference.py`` and +executed inside the serving container, so these tests exec its source the same way the +container does rather than importing it as a package module. +""" + +from __future__ import absolute_import + +import ast +import asyncio +import hashlib +import importlib +import importlib.util +import json +import os +import shutil +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import cloudpickle + +from sagemaker.serve.spec.inference_base import AsyncCustomOrchestrator, CustomOrchestrator + +MODULE_NAME = "sagemaker.serve.model_server.smd.custom_execution_inference" +INFERENCE_SOURCE_PATH = Path(importlib.util.find_spec(MODULE_NAME).origin) + + +class EchoOrchestrator(CustomOrchestrator): + """Minimal orchestrator used to populate serve.pkl.""" + + def handle(self, data, context=None): + return {"echo": data} + + +class AsyncEchoOrchestrator(AsyncCustomOrchestrator): + """Minimal async orchestrator, to exercise the awaited handler branch.""" + + async def handle(self, data, context=None): + return {"async_echo": data} + + +class _Request: + def __init__(self, body): + self.body = body + + +def _write_artifacts(code_dir: Path, obj=None): + """Write a serve.pkl/metadata.json pair the way prepare_for_smd does.""" + code_dir.mkdir(parents=True, exist_ok=True) + buffer = cloudpickle.dumps((EchoOrchestrator(), None) if obj is None else obj) + code_dir.joinpath("serve.pkl").write_bytes(buffer) + code_dir.joinpath("metadata.json").write_bytes( + json.dumps({"sha256_hash": hashlib.sha256(buffer).hexdigest()}).encode() + ) + + +def _load_inference(code_dir: Path): + """Import the entrypoint against code_dir, running its import-time preflight. + + The model server imports this file as ``inference``; importing the module for real + (rather than exec-ing its source) keeps the import-time work attributable to the + module under test. + """ + with patch.dict(os.environ, {"SAGEMAKER_INFERENCE_CODE_DIRECTORY": str(code_dir)}): + module = importlib.import_module(MODULE_NAME) + return importlib.reload(module) + + +class TestCustomExecutionInferenceIsSelfContained(unittest.TestCase): + """The entrypoint must not depend on the container's bundled sagemaker version.""" + + def test_no_sagemaker_imports(self): + """Regression test for issue #6200. + + The entrypoint runs against whichever sagemaker version the SMD image happens to + bundle, which is not the SDK version that produced the artifacts. Importing + anything from sagemaker -- the integrity check in particular -- made startup fail + on images carrying a pre-plain-SHA-256 sagemaker. + """ + tree = ast.parse(INFERENCE_SOURCE_PATH.read_text()) + + imported = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported.update(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module: + imported.add(node.module) + + offenders = sorted( + name for name in imported if name == "sagemaker" or name.startswith("sagemaker.") + ) + self.assertEqual(offenders, []) + + +class TestCodeDirectoryResolution(unittest.TestCase): + """Test _get_code_directory.""" + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + self.code_dir = Path(self.temp_dir) / "code" + _write_artifacts(self.code_dir) + self.module = _load_inference(self.code_dir) + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_uses_env_var_when_set(self): + with patch.dict(os.environ, {"SAGEMAKER_INFERENCE_CODE_DIRECTORY": "/some/code"}): + self.assertEqual(str(self.module._get_code_directory()), "/some/code") + + def test_falls_back_to_default(self): + with patch.dict(os.environ, {}, clear=True): + self.assertEqual(str(self.module._get_code_directory()), "/opt/ml/model/code") + + +class TestPerformIntegrityCheck(unittest.TestCase): + """Test the inlined _perform_integrity_check.""" + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + self.code_dir = Path(self.temp_dir) / "code" + _write_artifacts(self.code_dir) + self.module = _load_inference(self.code_dir) + self.check = self.module._perform_integrity_check + self.metadata_path = self.code_dir / "metadata.json" + self.buffer = (self.code_dir / "serve.pkl").read_bytes() + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_matching_digest_passes(self): + self.check(buffer=self.buffer, metadata_path=self.metadata_path) + + def test_no_secret_key_needed(self): + """The check must not read SAGEMAKER_SERVE_SECRET_KEY (removed for CVE-2026-8596).""" + with patch.dict(os.environ, {}, clear=True): + self.check(buffer=self.buffer, metadata_path=self.metadata_path) + + def test_tampered_buffer_fails(self): + with self.assertRaises(ValueError) as ctx: + self.check(buffer=self.buffer + b"x", metadata_path=self.metadata_path) + self.assertIn("Integrity check", str(ctx.exception)) + + def test_missing_metadata_fails(self): + with self.assertRaises(ValueError) as ctx: + self.check(buffer=self.buffer, metadata_path=self.code_dir / "absent.json") + self.assertIn("does not exist", str(ctx.exception)) + + def test_corrupt_metadata_fails(self): + self.metadata_path.write_bytes(b"not json") + with self.assertRaises(ValueError) as ctx: + self.check(buffer=self.buffer, metadata_path=self.metadata_path) + self.assertIn("not a valid json file", str(ctx.exception)) + + def test_metadata_without_digest_fails(self): + self.metadata_path.write_bytes(json.dumps({"version": "2023-04-24"}).encode()) + with self.assertRaises(ValueError) as ctx: + self.check(buffer=self.buffer, metadata_path=self.metadata_path) + self.assertIn("sha256_hash", str(ctx.exception)) + + +class TestEntrypointStartup(unittest.TestCase): + """Test the import-time preflight plus handler dispatch.""" + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + self.code_dir = Path(self.temp_dir) / "code" + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_startup_and_handle(self): + _write_artifacts(self.code_dir) + + module = _load_inference(self.code_dir) + + self.assertIsInstance(module.custom_orchestrator, EchoOrchestrator) + result = asyncio.run(module.handler(_Request(b"hello"))) + self.assertEqual(result, {"echo": b"hello"}) + + def test_startup_and_handle_async(self): + _write_artifacts(self.code_dir, obj=(AsyncEchoOrchestrator(), None)) + + module = _load_inference(self.code_dir) + + self.assertIsInstance(module.custom_orchestrator, AsyncEchoOrchestrator) + result = asyncio.run(module.handler(_Request(b"hello"))) + self.assertEqual(result, {"async_echo": b"hello"}) + + def test_startup_fails_on_tampered_pickle(self): + _write_artifacts(self.code_dir) + serve_pkl = self.code_dir / "serve.pkl" + serve_pkl.write_bytes(serve_pkl.read_bytes() + b"x") + + with self.assertRaises(ValueError): + _load_inference(self.code_dir) + + +if __name__ == "__main__": + unittest.main() diff --git a/sagemaker-serve/tests/unit/model_server/test_smd_prepare.py b/sagemaker-serve/tests/unit/model_server/test_smd_prepare.py index eeb9b6288a..3436fb7e5c 100644 --- a/sagemaker-serve/tests/unit/model_server/test_smd_prepare.py +++ b/sagemaker-serve/tests/unit/model_server/test_smd_prepare.py @@ -104,6 +104,70 @@ def test_prepare_for_smd_with_shared_libs(self, mock_copy, mock_capture, mock_ha # Verify copy2 was called for shared lib self.assertTrue(any(str(shared_lib) in str(call) for call in mock_copy.call_args_list)) + @patch("sagemaker.serve.model_server.smd.prepare.capture_dependencies") + def test_prepare_for_smd_returns_none(self, mock_capture): + """prepare_for_smd records the digest in metadata.json and returns nothing. + + The secret key it used to return was removed for CVE-2026-8596; the digest is + verified against metadata.json in the container instead. See issue #6200 -- + callers must not resurrect ``secret_key = prepare_for_smd(...)``. + """ + import hashlib + import json + + from sagemaker.serve.model_server.smd.prepare import prepare_for_smd + + model_path = Path(self.temp_dir) / "model" + code_dir = model_path / "code" + code_dir.mkdir(parents=True) + (code_dir / "serve.pkl").write_bytes(b"test data") + + result = prepare_for_smd(model_path=str(model_path), shared_libs=[], dependencies={}) + + self.assertIsNone(result) + metadata = json.loads((code_dir / "metadata.json").read_bytes()) + self.assertEqual(metadata["sha256_hash"], hashlib.sha256(b"test data").hexdigest()) + + @patch("sagemaker.serve.model_server.smd.prepare.capture_dependencies") + def test_prepare_for_smd_writes_sagemaker_free_entrypoint(self, mock_capture): + """The copied inference.py must not import the container's sagemaker package. + + Regression test for issue #6200: the container's bundled sagemaker version is set + by the SMD image, not by the SDK that built the artifacts. + """ + import ast + + from sagemaker.serve.model_server.smd.prepare import prepare_for_smd + from sagemaker.serve.spec.inference_base import CustomOrchestrator + + model_path = Path(self.temp_dir) / "model" + code_dir = model_path / "code" + code_dir.mkdir(parents=True) + (code_dir / "serve.pkl").write_bytes(b"test data") + + class _Orchestrator(CustomOrchestrator): + def handle(self, data, context=None): + return data + + prepare_for_smd( + model_path=str(model_path), + shared_libs=[], + dependencies={}, + inference_spec=_Orchestrator(), + ) + + tree = ast.parse((code_dir / "inference.py").read_text()) + imported = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported.update(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module: + imported.add(node.module) + + self.assertEqual( + sorted(n for n in imported if n == "sagemaker" or n.startswith("sagemaker.")), [] + ) + def test_prepare_for_smd_invalid_dir(self): """Test prepare_for_smd raises exception for invalid directory.""" from sagemaker.serve.model_server.smd.prepare import prepare_for_smd diff --git a/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py b/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py index e4509c6af8..b2f4be5703 100644 --- a/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py +++ b/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py @@ -1309,6 +1309,41 @@ def test_build_with_auto_image( self.assertEqual(self.builder.secret_key, "") mock_create.assert_called_once() + @patch("sagemaker.serve.model_builder_servers.prepare_for_smd") + @patch.object(MockModelBuilderServers, "_save_model_inference_spec") + @patch.object(MockModelBuilderServers, "_get_processing_unit") + @patch.object(MockModelBuilderServers, "_get_smd_image_uri") + @patch.object(MockModelBuilderServers, "_prepare_for_mode") + @patch.object(MockModelBuilderServers, "_create_model") + def test_build_does_not_adopt_prepare_return_as_secret_key( + self, + mock_create, + mock_prepare_mode, + mock_get_img, + mock_get_unit, + mock_save, + mock_smd_prepare, + ): + """secret_key stays empty regardless of what prepare_for_smd returns. + + Regression test for issue #6200: the secret key was removed for CVE-2026-8596, + so prepare_for_smd returns nothing and its result must not be published as + SAGEMAKER_SERVE_SECRET_KEY. + """ + mock_get_unit.return_value = "cpu" + mock_get_img.return_value = "smd-image-uri" + mock_smd_prepare.return_value = "a-secret-key" + mock_create.return_value = Mock() + self.builder.mode = Mode.LOCAL_CONTAINER + self.builder.image_uri = None + self.builder.model = Mock() + + self.builder._build_for_smd() + + mock_smd_prepare.assert_called_once() + self.assertEqual(self.builder.secret_key, "") + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", self.builder.env_vars) + class TestBuildForTransformers(unittest.TestCase): """Test _build_for_transformers method.""" diff --git a/sagemaker-serve/tests/unit/test_model_builder_build.py b/sagemaker-serve/tests/unit/test_model_builder_build.py index d0d5eebc28..9077847c03 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_build.py +++ b/sagemaker-serve/tests/unit/test_model_builder_build.py @@ -57,6 +57,35 @@ def test_save_model_inference_spec_with_inference_spec(self, mock_save_pkl): self.assertIn("code", str(args[0])) self.assertEqual(args[1], (mock_inference_spec, mock_schema)) + @patch("sagemaker.serve.model_builder.save_pkl") + def test_save_model_inference_spec_omits_schema_for_custom_orchestrator(self, mock_save_pkl): + """Custom orchestrators pickle without the SchemaBuilder. + + Regression test for issue #6200: the SMD entrypoint discards the second element + of the tuple, and a SchemaBuilder pickles with references to sagemaker.core + modules that an older sagemaker in the serving container cannot import, so + including it made cloudpickle.load() fail at container startup. + """ + from sagemaker.serve.builder.schema_builder import SchemaBuilder + from sagemaker.serve.spec.inference_base import CustomOrchestrator + + mock_orchestrator = Mock(spec=CustomOrchestrator) + mock_schema = Mock(spec=SchemaBuilder) + + builder = ModelBuilder( + inference_spec=mock_orchestrator, + schema_builder=mock_schema, + model_path=self.temp_dir, + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + ) + builder.env_vars = {} + + builder._save_model_inference_spec() + + mock_save_pkl.assert_called_once() + self.assertEqual(mock_save_pkl.call_args[0][1], (mock_orchestrator, None)) + @patch("sagemaker.serve.model_builder.save_pkl") @patch("sagemaker.serve.model_builder._detect_framework_and_version") @patch("sagemaker.serve.model_builder._get_model_base")