fix: prepare_for_smd() missing return value causes CustomOrchestrator container h (6200)

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()

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

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

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