From 78b0f3d260b6cc83421531e6a7e67b6d9c0f8b36 Mon Sep 17 00:00:00 2001 From: SageMaker Bot <49924207+sagemaker-bot@users.noreply.github.com> Date: Fri, 28 Aug 2026 12:10:56 -0700 Subject: [PATCH] fix: `prepare_for_smd()` missing return value causes `CustomOrchestrator` container h (6200) --- .../sagemaker/serve/model_builder_servers.py | 6 ++ .../serve/model_server/smd/prepare.py | 16 +++- .../serve/model_server/smd/server.py | 2 + .../serve/validations/check_integrity.py | 80 +++++++++++++++++-- 4 files changed, 94 insertions(+), 10 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py index 349bee9552..214d5390b6 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py @@ -751,6 +751,12 @@ def _build_for_smd(self) -> Model: inference_spec=self.inference_spec, ) + # Propagate secret key to container environment + if self.secret_key: + self.env_vars["SAGEMAKER_SERVE_SECRET_KEY"] = self.secret_key + else: + self.env_vars.pop("SAGEMAKER_SERVE_SECRET_KEY", None) + # Prepare deployment artifacts if self.mode in LOCAL_MODES: self._prepare_for_mode() 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..dc4faf0fcb 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.py @@ -13,10 +13,15 @@ from sagemaker.serve.detector.dependency_manager import capture_dependencies from sagemaker.serve.validations.check_integrity import ( compute_hash, + generate_secret_key, ) from sagemaker.core.remote_function.core.serialization import _MetaData from sagemaker.serve.spec.inference_base import CustomOrchestrator, AsyncCustomOrchestrator +import logging + +logger = logging.getLogger(__name__) + def prepare_for_smd( model_path: str, @@ -34,7 +39,9 @@ def prepare_for_smd( (default is None) Returns: - ( str ) : + str: A generated secret key used to compute the HMAC hash stored in + metadata.json. Callers should propagate this value to the container + environment as SAGEMAKER_SERVE_SECRET_KEY. """ model_path = Path(model_path) @@ -63,8 +70,13 @@ def prepare_for_smd( capture_dependencies(dependencies=dependencies, work_dir=code_dir) + secret_key = generate_secret_key() + logger.debug("Generated secret key for SMD artifact integrity check.") + with open(str(code_dir.joinpath("serve.pkl")), "rb") as f: buffer = f.read() - hash_value = compute_hash(buffer=buffer) + hash_value = compute_hash(buffer=buffer, secret_key=secret_key) with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata: metadata.write(_MetaData(hash_value).to_json()) + + return secret_key diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/smd/server.py b/sagemaker-serve/src/sagemaker/serve/model_server/smd/server.py index ecb68406c1..abb47014d2 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/smd/server.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/smd/server.py @@ -55,4 +55,6 @@ def _upload_smd_artifacts( "SAGEMAKER_REGION": sagemaker_session.boto_region_name, "LOCAL_PYTHON": platform.python_version(), } + if secret_key: + env_vars["SAGEMAKER_SERVE_SECRET_KEY"] = secret_key return s3_upload_path, env_vars diff --git a/sagemaker-serve/src/sagemaker/serve/validations/check_integrity.py b/sagemaker-serve/src/sagemaker/serve/validations/check_integrity.py index 880ca5b602..daa1a0d1c2 100644 --- a/sagemaker-serve/src/sagemaker/serve/validations/check_integrity.py +++ b/sagemaker-serve/src/sagemaker/serve/validations/check_integrity.py @@ -1,27 +1,91 @@ -"""Validates the integrity of pickled file with SHA-256 hash.""" +"""Validates the integrity of pickled file with SHA-256 hash. + +Supports two modes: +- Plain SHA-256 (default): Used when no secret key is provided. +- HMAC-SHA256 (keyed): Used when a secret key is provided, for backward + compatibility with older container images that perform HMAC-based checks. +""" from __future__ import absolute_import import hmac import hashlib +import os +import secrets from pathlib import Path from sagemaker.core.remote_function.core.serialization import _MetaData +SAGEMAKER_SERVE_SECRET_KEY = "SAGEMAKER_SERVE_SECRET_KEY" + + +def generate_secret_key(nbytes: int = 32) -> str: + """Generate a cryptographically secure secret key. + + Args: + nbytes: Number of random bytes (the returned hex string will be + twice this length). Defaults to 32 (256-bit key). + + Returns: + A hex-encoded random string suitable for use as an HMAC key. + """ + return secrets.token_hex(nbytes) + + +def compute_hash(buffer: bytes, secret_key: str = None) -> str: + """Compute hash of the given buffer. -def compute_hash(buffer: bytes) -> str: - """Compute SHA-256 hash of the given buffer.""" + When *secret_key* is provided the hash is an HMAC-SHA256 keyed digest; + otherwise a plain SHA-256 digest is returned. + + Args: + buffer: The bytes to hash. + secret_key: Optional HMAC key. When ``None`` (default) a plain + SHA-256 hash is computed. + + Returns: + Hex-encoded hash string. + """ + if secret_key: + return hmac.new(secret_key.encode(), msg=buffer, digestmod=hashlib.sha256).hexdigest() return hashlib.sha256(buffer).hexdigest() -def perform_integrity_check(buffer: bytes, metadata_path: Path): - """Validates the integrity of bytes by comparing the hash value.""" - actual_hash_value = compute_hash(buffer=buffer) +def perform_integrity_check(buffer: bytes, metadata_path: Path, secret_key: str = None): + """Validates the integrity of bytes by comparing the hash value. + Computes both the plain SHA-256 digest and (when a secret key is + available) the HMAC-SHA256 digest, then checks whether the expected + hash stored in *metadata_path* matches either one. This provides + backward compatibility between SDK versions that write plain hashes + and container images that expect HMAC hashes (or vice-versa). + + Args: + buffer: The serialized bytes to verify. + metadata_path: Path to the ``metadata.json`` file containing the + expected hash. + secret_key: Optional HMAC key. When ``None`` the function falls + back to the ``SAGEMAKER_SERVE_SECRET_KEY`` environment variable. + """ if not Path.exists(metadata_path): raise ValueError("Path to metadata.json does not exist") with open(str(metadata_path), "rb") as md: expected_hash_value = _MetaData.from_json(md.read()).sha256_hash - if not hmac.compare_digest(expected_hash_value, actual_hash_value): - raise ValueError("Integrity check for the serialized function or data failed.") + # Resolve secret key: explicit arg > environment variable > None + effective_secret_key = secret_key or os.environ.get(SAGEMAKER_SERVE_SECRET_KEY) + + # Compute candidate digests + plain_hash = hashlib.sha256(buffer).hexdigest() + + if hmac.compare_digest(expected_hash_value, plain_hash): + return + + if effective_secret_key: + hmac_hash = hmac.new( + effective_secret_key.encode(), msg=buffer, digestmod=hashlib.sha256 + ).hexdigest() + if hmac.compare_digest(expected_hash_value, hmac_hash): + return + + raise ValueError("Integrity check for the serialized function or data failed.")