Bug fix for integrity check using hmac key in serve V2

src/sagemaker/serve/builder/djl_builder.py

@@ -213,7 +213,6 @@ def _djl_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
             self.modes[str(Mode.LOCAL_CONTAINER)].create_server(
                 self.image_uri,
                 timeout if timeout else 1800,
-                self.secret_key,
                 predictor,
                 self.pysdk_model.env,
             )

src/sagemaker/serve/builder/model_builder.py

@@ -448,7 +448,6 @@ def _prepare_for_mode(
                 str(Mode.SAGEMAKER_ENDPOINT)
             ].prepare(
                 (model_path or self.model_path),
-                self.secret_key,
                 self.serve_settings.s3_model_data_url,
                 self.sagemaker_session,
                 self.image_uri,
@@ -709,7 +708,7 @@ def _model_builder_deploy_wrapper(
             )
 
             self.modes[str(Mode.LOCAL_CONTAINER)].create_server(
-                self.image_uri, container_timeout_in_second, self.secret_key, predictor
+                self.image_uri, container_timeout_in_second, predictor
             )
             return predictor
 
@@ -778,7 +777,7 @@ def _build_for_torchserve(self) -> Type[Model]:
         if self.mode != Mode.IN_PROCESS:
             self._auto_detect_container()
 
-            self.secret_key = prepare_for_torchserve(
+            prepare_for_torchserve(
                 model_path=self.model_path,
                 shared_libs=self.shared_libs,
                 dependencies=self.dependencies,
@@ -798,7 +797,7 @@ def _build_for_smd(self) -> Type[Model]:
         if self.mode != Mode.IN_PROCESS:
             self._auto_detect_container()
 
-            self.secret_key = prepare_for_smd(
+            prepare_for_smd(
                 model_path=self.model_path,
                 shared_libs=self.shared_libs,
                 dependencies=self.dependencies,

src/sagemaker/serve/builder/tf_serving_builder.py

@@ -124,7 +124,7 @@ def _build_for_tensorflow_serving(self):
         if not self.image_uri:
             raise ValueError("image_uri is not set for tensorflow serving")
 
-        self.secret_key = prepare_for_tf_serving(
+        prepare_for_tf_serving(
             model_path=self.model_path,
             shared_libs=self.shared_libs,
             dependencies=self.dependencies,

src/sagemaker/serve/builder/transformers_builder.py

@@ -267,12 +267,6 @@ def _transformers_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[Pr
             self.env_vars.update(env_vars)
             self.pysdk_model.env.update(self.env_vars)
 
-        if (
-            "SAGEMAKER_SERVE_SECRET_KEY" in self.pysdk_model.env
-            and not self.pysdk_model.env["SAGEMAKER_SERVE_SECRET_KEY"]
-        ):
-            del self.pysdk_model.env["SAGEMAKER_SERVE_SECRET_KEY"]
-
         if "endpoint_logging" not in kwargs:
             kwargs["endpoint_logging"] = True
 
@@ -402,7 +396,7 @@ def _build_for_transformers(self):
 
             self._auto_detect_container()
 
-            self.secret_key = prepare_for_mms(
+            prepare_for_mms(
                 model_path=self.model_path,
                 shared_libs=self.shared_libs,
                 dependencies=self.dependencies,

src/sagemaker/serve/detector/dependency_manager.py

@@ -66,9 +66,11 @@ def capture_dependencies(dependencies: dict, work_dir: Path, capture_all: bool =
 
         with open(path, "r") as f:
             autodetect_depedencies = f.read().splitlines()
-        autodetect_depedencies.append("sagemaker[huggingface]>=2.199")
+        # Pin sagemaker to 2.257.0+ to ensure SHA256 hashing is used for integrity checks
+        autodetect_depedencies.append("sagemaker[huggingface]>=2.257.0,<3.0.0")
     else:
-        autodetect_depedencies = ["sagemaker[huggingface]>=2.199"]
+        # Pin sagemaker to 2.257.0+ to ensure SHA256 hashing is used for integrity checks
+        autodetect_depedencies = ["sagemaker[huggingface]>=2.257.0,<3.0.0"]
 
     module_version_dict = _parse_dependency_list(autodetect_depedencies)
 

src/sagemaker/serve/mode/local_container_mode.py

@@ -68,7 +68,6 @@ def __init__(
         self.model_server = model_server
         self.client = None
         self.container = None
-        self.secret_key = None
         self._ping_container = None
         self._invoke_serving = None
 
@@ -89,7 +88,6 @@ def create_server(
         self,
         image: str,
         container_timeout_seconds: int,
-        secret_key: str,
         predictor: PredictorBase,
         env_vars: Dict[str, str] = None,
         model_path: str = None,
@@ -108,7 +106,6 @@ def create_server(
                 docker_client=self.client,
                 model_path=model_path if model_path else self.model_path,
                 image_uri=image,
-                secret_key=secret_key,
                 env_vars=env_vars if env_vars else self.env_vars,
             )
             self._ping_container = self._triton_deep_ping
@@ -117,7 +114,6 @@ def create_server(
                 client=self.client,
                 image=image,
                 model_path=model_path if model_path else self.model_path,
-                secret_key=secret_key,
                 env_vars=env_vars if env_vars else self.env_vars,
             )
             self._ping_container = self._djl_deep_ping
@@ -126,7 +122,6 @@ def create_server(
                 client=self.client,
                 image=image,
                 model_path=model_path if model_path else self.model_path,
-                secret_key=secret_key,
                 env_vars=env_vars if env_vars else self.env_vars,
             )
             self._ping_container = self._torchserve_deep_ping
@@ -135,7 +130,6 @@ def create_server(
                 client=self.client,
                 image=image,
                 model_path=model_path if model_path else self.model_path,
-                secret_key=secret_key,
                 env_vars=env_vars if env_vars else self.env_vars,
                 jumpstart=jumpstart,
             )
@@ -145,7 +139,6 @@ def create_server(
                 client=self.client,
                 image=image,
                 model_path=model_path if model_path else self.model_path,
-                secret_key=secret_key,
                 env_vars=env_vars if env_vars else self.env_vars,
             )
             self._ping_container = self._multi_model_server_deep_ping
@@ -154,7 +147,6 @@ def create_server(
                 client=self.client,
                 image=image,
                 model_path=model_path if model_path else self.model_path,
-                secret_key=secret_key,
                 env_vars=env_vars if env_vars else self.env_vars,
             )
             self._ping_container = self._tensorflow_serving_deep_ping
@@ -164,7 +156,6 @@ def create_server(
                 client=self.client,
                 image=image,
                 model_path=model_path if model_path else self.model_path,
-                secret_key=secret_key,
                 env_vars=env_vars if env_vars else self.env_vars,
             )
             tei_serving.schema_builder = self.schema_builder

src/sagemaker/serve/mode/sagemaker_endpoint_mode.py

@@ -58,7 +58,6 @@ def load(self, model_path: str):
     def prepare(
         self,
         model_path: str,
-        secret_key: str,
         s3_model_data_url: str = None,
         sagemaker_session: Session = None,
         image: str = None,
@@ -79,7 +78,6 @@ def prepare(
             upload_artifacts = self._upload_torchserve_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
-                secret_key=secret_key,
                 s3_model_data_url=s3_model_data_url,
                 image=image,
                 should_upload_artifacts=True,
@@ -89,7 +87,6 @@ def prepare(
             upload_artifacts = self._upload_triton_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
-                secret_key=secret_key,
                 s3_model_data_url=s3_model_data_url,
                 image=image,
                 should_upload_artifacts=True,
@@ -108,7 +105,6 @@ def prepare(
             upload_artifacts = self._upload_tensorflow_serving_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
-                secret_key=secret_key,
                 s3_model_data_url=s3_model_data_url,
                 image=image,
                 should_upload_artifacts=True,
@@ -134,7 +130,6 @@ def prepare(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 s3_model_data_url=s3_model_data_url,
-                secret_key=secret_key,
                 image=image,
                 should_upload_artifacts=should_upload_artifacts,
             )
@@ -152,7 +147,6 @@ def prepare(
             upload_artifacts = self._upload_smd_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
-                secret_key=secret_key,
                 s3_model_data_url=s3_model_data_url,
                 image=image,
                 should_upload_artifacts=True,

src/sagemaker/serve/model_server/djl_serving/server.py

@@ -31,7 +31,7 @@ class LocalDJLServing:
     """Placeholder docstring"""
 
     def _start_djl_serving(
-        self, client: object, image: str, model_path: str, secret_key: str, env_vars: dict
+        self, client: object, image: str, model_path: str, env_vars: dict
     ):
         """Placeholder docstring"""
         updated_env_vars = _update_env_vars(env_vars)

src/sagemaker/serve/model_server/multi_model_server/prepare.py

@@ -25,10 +25,7 @@
 from sagemaker.session import Session
 from sagemaker.serve.spec.inference_spec import InferenceSpec
 from sagemaker.serve.detector.dependency_manager import capture_dependencies
-from sagemaker.serve.validations.check_integrity import (
-    generate_secret_key,
-    compute_hash,
-)
+from sagemaker.serve.validations.check_integrity import compute_hash
 from sagemaker.remote_function.core.serialization import _MetaData
 
 logger = logging.getLogger(__name__)
@@ -85,7 +82,7 @@ def prepare_for_mms(
     inference_spec: InferenceSpec = None,
 ) -> str:
     """Prepares for InferenceSpec using model_path, writes inference.py, \
-        and captures dependencies to generate secret_key.
+        and captures dependencies.
 
     Args:to
         model_path (str) : Argument
@@ -95,7 +92,7 @@ def prepare_for_mms(
         inference_spec (InferenceSpec, optional) : Argument
             (default is None)
     Returns:
-        ( str ) : secret_key
+        ( str ) : Empty string for backward compatibility
     """
     model_path = Path(model_path)
     if not model_path.exists():
@@ -120,11 +117,10 @@ def prepare_for_mms(
 
     capture_dependencies(dependencies=dependencies, work_dir=code_dir)
 
-    secret_key = generate_secret_key()
     with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
         buffer = f.read()
-    hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
+    hash_value = compute_hash(buffer=buffer)
     with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
         metadata.write(_MetaData(hash_value).to_json())
 
-    return secret_key
+    return ""

src/sagemaker/serve/model_server/multi_model_server/server.py

@@ -29,14 +29,12 @@ def _start_serving(
         client: object,
         image: str,
         model_path: str,
-        secret_key: str,
         env_vars: dict,
     ):
         """Initializes the start of the server"""
         env = {
             "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
             "SAGEMAKER_PROGRAM": "inference.py",
-            "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
             "LOCAL_PYTHON": platform.python_version(),
         }
         if env_vars:
@@ -94,7 +92,6 @@ class SageMakerMultiModelServer:
     def _upload_server_artifacts(
         self,
         model_path: str,
-        secret_key: str,
         sagemaker_session: Session,
         s3_model_data_url: str = None,
         image: str = None,
@@ -141,15 +138,16 @@ def _upload_server_artifacts(
             else None
         )
 
-        if secret_key:
-            env_vars = {
-                "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
-                "SAGEMAKER_PROGRAM": "inference.py",
-                "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
-                "SAGEMAKER_REGION": sagemaker_session.boto_region_name,
-                "SAGEMAKER_CONTAINER_LOG_LEVEL": "10",
-                "LOCAL_PYTHON": platform.python_version(),
-            }
+        if env_vars is None:
+            env_vars = {}
+
+        env_vars.update({
+            "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
+            "SAGEMAKER_PROGRAM": "inference.py",
+            "SAGEMAKER_REGION": sagemaker_session.boto_region_name,
+            "SAGEMAKER_CONTAINER_LOG_LEVEL": "10",
+            "LOCAL_PYTHON": platform.python_version(),
+        })
 
         return model_data, _update_env_vars(env_vars)
 

src/sagemaker/serve/model_server/smd/prepare.py

@@ -11,10 +11,7 @@
 
 from sagemaker.serve.spec.inference_spec import InferenceSpec
 from sagemaker.serve.detector.dependency_manager import capture_dependencies
-from sagemaker.serve.validations.check_integrity import (
-    generate_secret_key,
-    compute_hash,
-)
+from sagemaker.serve.validations.check_integrity import compute_hash
 from sagemaker.remote_function.core.serialization import _MetaData
 from sagemaker.serve.spec.inference_base import CustomOrchestrator, AsyncCustomOrchestrator
 
@@ -64,11 +61,10 @@ def prepare_for_smd(
 
     capture_dependencies(dependencies=dependencies, work_dir=code_dir)
 
-    secret_key = generate_secret_key()
     with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
         buffer = f.read()
-    hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
+    hash_value = compute_hash(buffer=buffer)
     with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
         metadata.write(_MetaData(hash_value).to_json())
 
-    return secret_key
+    return ""

src/sagemaker/serve/model_server/smd/server.py

@@ -20,7 +20,6 @@ def _upload_smd_artifacts(
         self,
         model_path: str,
         sagemaker_session: Session,
-        secret_key: str,
         s3_model_data_url: str = None,
         image: str = None,
         should_upload_artifacts: bool = False,
@@ -53,7 +52,6 @@ def _upload_smd_artifacts(
             "SAGEMAKER_INFERENCE_CODE_DIRECTORY": "/opt/ml/model/code",
             "SAGEMAKER_INFERENCE_CODE": "inference.handler",
             "SAGEMAKER_REGION": sagemaker_session.boto_region_name,
-            "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
             "LOCAL_PYTHON": platform.python_version(),
         }
         return s3_upload_path, env_vars

src/sagemaker/serve/model_server/tei/server.py

@@ -28,20 +28,16 @@ class LocalTeiServing:
     """LocalTeiServing class"""
 
     def _start_tei_serving(
-        self, client: object, image: str, model_path: str, secret_key: str, env_vars: dict
+        self, client: object, image: str, model_path: str, env_vars: dict
     ):
         """Starts a local tei serving container.
 
         Args:
             client: Docker client
             image: Image to use
             model_path: Path to the model
-            secret_key: Secret key to use for authentication
             env_vars: Environment variables to set
         """
-        if env_vars and secret_key:
-            env_vars["SAGEMAKER_SERVE_SECRET_KEY"] = secret_key
-
         self.container = client.containers.run(
             image,
             shm_size=_SHM_SIZE,

src/sagemaker/serve/model_server/tensorflow_serving/prepare.py

@@ -10,10 +10,7 @@
     _move_contents,
 )
 from sagemaker.serve.detector.dependency_manager import capture_dependencies
-from sagemaker.serve.validations.check_integrity import (
-    generate_secret_key,
-    compute_hash,
-)
+from sagemaker.serve.validations.check_integrity import compute_hash
 from sagemaker.remote_function.core.serialization import _MetaData
 
 
@@ -57,11 +54,10 @@ def prepare_for_tf_serving(
         raise ValueError("SavedModel is not found for Tensorflow or Keras flavor.")
     _move_contents(src_dir=mlflow_saved_model_dir, dest_dir=saved_model_bundle_dir)
 
-    secret_key = generate_secret_key()
     with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
         buffer = f.read()
-    hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
+    hash_value = compute_hash(buffer=buffer)
     with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
         metadata.write(_MetaData(hash_value).to_json())
 
-    return secret_key
+    return ""