fix(serve): decouple the SMD orchestrator entrypoint from the container's sagemaker (#6200)

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

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,

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

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

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

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)

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