Removing final secret key references, pin sm dependency in container

Author: aviruthen

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

@@ -396,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,13 @@ 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
+        # (version where HMAC vulnerability was fixed). Update this version when 2.257.0 is released.
+        autodetect_depedencies.append("sagemaker[huggingface]>=2.257.0")
     else:
-        autodetect_depedencies = ["sagemaker[huggingface]>=2.199"]
+        # Pin sagemaker to 2.257.0+ to ensure SHA256 hashing is used for integrity checks
+        # (version where HMAC vulnerability was fixed). Update this version when 2.257.0 is released.
+        autodetect_depedencies = ["sagemaker[huggingface]>=2.257.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

@@ -82,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
@@ -92,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():

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

@@ -92,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,