Removing all secret key references, updating unit tests

Author: aviruthen

sagemaker-serve/src/sagemaker/serve/local_resources.py

@@ -68,7 +68,6 @@ def __init__(
         local_container_mode_obj=None,
         in_process_mode_obj=None,
         model_server=None,
-        secret_key=None,
         serializer=None,
         deserializer=None,
         container_config="auto",
@@ -89,7 +88,6 @@ def __init__(
         self.local_container_mode_obj=local_container_mode_obj
         self.in_process_mode_obj=in_process_mode_obj
         self.model_server=model_server
-        self.secret_key=secret_key
         self.serializer=serializer
         self.deserializer=deserializer
         self.container_config=container_config
@@ -295,7 +293,6 @@ def create(
         local_container_mode_obj=None,
         in_process_mode_obj=None,
         model_server=None,
-        secret_key=None,
         serializer=None,
         deserializer=None,
         container_config="auto",
@@ -317,7 +314,6 @@ def create(
                 local_container_mode_obj=local_container_mode_obj,
                 in_process_mode_obj=in_process_mode_obj,
                 model_server=model_server,
-                secret_key=secret_key,
                 serializer=serializer,
                 deserializer=deserializer,
                 container_config=container_config,
@@ -342,7 +338,6 @@ def create(
                 local_container_mode_obj=local_container_mode_obj,
                 in_process_mode_obj=in_process_mode_obj,
                 model_server=model_server,
-                secret_key=secret_key,
                 serializer=serializer,
                 deserializer=deserializer,
                 container_config=container_config,
@@ -353,7 +348,6 @@ def create(
             endpoint.local_container_mode_obj.create_server(
                 image=local_model.primary_container.image,
                 container_timeout_seconds=kwargs.get("container_timeout_seconds", 300),
-                secret_key=endpoint.secret_key,
                 ping_fn=endpoint._universal_deep_ping,
                 env_vars=local_model.primary_container.environment or {},
                 model_path=endpoint.local_container_mode_obj.model_path,

sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py

@@ -67,7 +67,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
 
@@ -88,7 +87,6 @@ def create_server(
         self,
         image: str,
         container_timeout_seconds: int,
-        secret_key: str,
         container_config: Dict,
         ping_fn = None,
         env_vars: Dict[str, str] = None,
@@ -109,31 +107,27 @@ 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,
             )
         elif self.model_server == ModelServer.DJL_SERVING:
             self._start_djl_serving(
                 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,
             )
         elif self.model_server == ModelServer.TORCHSERVE:
             self._start_torch_serve(
                 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,
             )
         elif self.model_server == ModelServer.TGI:
             self._start_tgi_serving(
                 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,
             )
@@ -142,15 +136,13 @@ 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,
             )
         elif self.model_server == ModelServer.TENSORFLOW_SERVING:
             self._start_tensorflow_serving(
                 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,
             )
         elif self.model_server == ModelServer.TEI:
@@ -159,7 +151,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

sagemaker-serve/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py

@@ -56,7 +56,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,

sagemaker-serve/src/sagemaker/serve/model_builder.py

@@ -779,7 +779,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,
@@ -1902,7 +1901,6 @@ def _deploy_local_endpoint(self, **kwargs):
                 local_model=self.built_model,
                 local_session=local_session,
                 container_timeout_seconds=kwargs.get("container_timeout_in_seconds", 300),
-                secret_key=self.secret_key,
                 local_container_mode_obj=self.modes[str(Mode.LOCAL_CONTAINER)],
                 serializer=self._serializer,
                 deserializer=self._deserializer,
@@ -2373,7 +2371,6 @@ def _deploy(self, **kwargs):
                 in_process_mode=True,
                 local_model=self.built_model,
                 container_timeout_seconds=kwargs.get("container_timeout_in_seconds", 300),
-                secret_key=self.secret_key,
                 in_process_mode_obj=self.modes[str(Mode.IN_PROCESS)],
                 serializer=self._serializer,
                 deserializer=self._deserializer,
@@ -2469,7 +2466,6 @@ def _reset_build_state(self):
         """Reset all dynamically added build-related state."""
         # Core build state
         self.built_model = None
-        self.secret_key = ""
 
         # JumpStart preparation flags
         for attr in ['prepared_for_djl', 'prepared_for_tgi', 'prepared_for_mms']:

sagemaker-serve/src/sagemaker/serve/model_builder_servers.py

@@ -128,7 +128,6 @@ def _build_for_torchserve(self) -> Model:
         Returns:
             Model: Configured model ready for TorchServe deployment.
         """
-        self.secret_key = ""
 
         # Save inference spec if we have local artifacts
         self._save_model_inference_spec()
@@ -150,7 +149,7 @@ def _build_for_torchserve(self) -> Model:
 
             # Prepare TorchServe artifacts for local container mode
             if self.mode == Mode.LOCAL_CONTAINER and self.model_path:
-                self.secret_key = prepare_for_torchserve(
+                prepare_for_torchserve(
                     model_path=self.model_path,
                     shared_libs=self.shared_libs,
                     dependencies=self.dependencies,
@@ -159,7 +158,7 @@ def _build_for_torchserve(self) -> Model:
                     inference_spec=self.inference_spec,
                 )
             if self.mode == Mode.SAGEMAKER_ENDPOINT and self.model_path:
-                self.secret_key = prepare_for_torchserve(
+                prepare_for_torchserve(
                     model_path=self.model_path,
                     shared_libs=self.shared_libs,
                     dependencies=self.dependencies,
@@ -187,7 +186,6 @@ def _build_for_tgi(self) -> Model:
         Returns:
             Model: Configured model ready for TGI deployment.
         """
-        self.secret_key = ""
 
         # Initialize TGI-specific configuration
         if self.model_server != ModelServer.TGI:
@@ -299,7 +297,6 @@ def _build_for_djl(self) -> Model:
         Returns:
             Model: Configured model ready for DJL Serving deployment.
         """
-        self.secret_key = ""
         self.model_server = ModelServer.DJL_SERVING
 
         # Set MODEL_LOADING_TIMEOUT from instance variable
@@ -408,7 +405,6 @@ def _build_for_triton(self) -> Model:
         Returns:
             Model: Configured model ready for Triton deployment.
         """
-        self.secret_key = ""
         self._validate_for_triton()
 
         if isinstance(self.model, str):
@@ -467,7 +463,6 @@ def _build_for_tensorflow_serving(self) -> Model:
         Raises:
             ValueError: If image_uri is not provided for TensorFlow Serving.
         """
-        self.secret_key = ""
         if not getattr(self, "_is_mlflow_model", False):
             raise ValueError("Tensorflow Serving is currently only supported for mlflow models.")
 
@@ -481,7 +476,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
-        self.secret_key = prepare_for_tf_serving(
+        prepare_for_tf_serving(
             model_path=self.model_path,
             shared_libs=self.shared_libs,
             dependencies=self.dependencies,
@@ -506,7 +501,6 @@ def _build_for_tei(self) -> Model:
         Returns:
             Model: Configured model ready for TEI deployment.
         """
-        self.secret_key = ""
 
         # Set MODEL_LOADING_TIMEOUT from instance variable
         if self.model_data_download_timeout:
@@ -592,7 +586,6 @@ def _build_for_smd(self) -> Model:
         Returns:
             Model: Configured model ready for SMD deployment.
         """
-        self.secret_key = ""
 
         self._save_model_inference_spec()
 
@@ -602,7 +595,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)
 
-            self.secret_key = prepare_for_smd(
+            prepare_for_smd(
                 model_path=self.model_path,
                 shared_libs=self.shared_libs,
                 dependencies=self.dependencies,
@@ -626,7 +619,6 @@ def _build_for_transformers(self) -> Model:
         Returns:
             Model: Configured model ready for Transformers deployment.
         """
-        self.secret_key = ""
         self.model_server = ModelServer.MMS
 
         # Set MODEL_LOADING_TIMEOUT from instance variable
@@ -646,7 +638,7 @@ def _build_for_transformers(self) -> Model:
                 self._create_conda_env()
 
             if self.mode in [Mode.LOCAL_CONTAINER] and self.model_path:
-                self.secret_key = prepare_for_mms(
+                prepare_for_mms(
                     model_path=self.model_path,
                     shared_libs=self.shared_libs,
                     dependencies=self.dependencies,
@@ -655,7 +647,7 @@ def _build_for_transformers(self) -> Model:
                     inference_spec=self.inference_spec,
                 )
             if self.mode == Mode.SAGEMAKER_ENDPOINT and self.model_path:
-                self.secret_key = prepare_for_mms(
+                prepare_for_mms(
                     model_path=self.model_path,
                     shared_libs=self.shared_libs,
                     dependencies=self.dependencies,
@@ -725,7 +717,6 @@ def _build_for_djl_jumpstart(self, init_kwargs) -> Model:
         Returns:
             Model: Configured DJL model for JumpStart deployment.
         """
-        self.secret_key = ""
         self.model_server = ModelServer.DJL_SERVING
 
         from sagemaker.serve.model_server.djl_serving.prepare import _create_dir_structure
@@ -761,7 +752,6 @@ def _build_for_tgi_jumpstart(self, init_kwargs) -> Model:
         Returns:
             Model: Configured TGI model for JumpStart deployment.
         """
-        self.secret_key = ""
         self.model_server = ModelServer.TGI
 
         from sagemaker.serve.model_server.tgi.prepare import _create_dir_structure
@@ -797,7 +787,6 @@ def _build_for_mms_jumpstart(self, init_kwargs) -> Model:
         Returns:
             Model: Configured MMS model for JumpStart deployment.
         """
-        self.secret_key = ""
         self.model_server = ModelServer.MMS
 
         from sagemaker.serve.model_server.multi_model_server.prepare import _create_dir_structure
@@ -840,7 +829,6 @@ def _build_for_jumpstart(self) -> Model:
         """
         from sagemaker.core.jumpstart.factory.utils import get_init_kwargs
 
-        self.secret_key = ""
 
         # Get JumpStart model configuration
         init_kwargs = get_init_kwargs(

sagemaker-serve/src/sagemaker/serve/model_builder_utils.py

@@ -3056,8 +3056,6 @@ def _prepare_for_triton(self):
             export_path.mkdir(parents=True)
 
         if self.model:
-            self.secret_key = "dummy secret key for onnx backend"
-
             if self.framework == Framework.PYTORCH:
                 self._export_pytorch_to_onnx(
                     export_path=export_path, model=self.model, schema_builder=self.schema_builder

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

@@ -30,7 +30,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)

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

@@ -81,7 +81,7 @@ def prepare_for_mms(
     image_uri: str,
     inference_spec: InferenceSpec = None,
 ) -> str:
-    """Prepares for InferenceSpec using model_path, writes inference.py, and captures dependencies to generate secret_key.
+    """Prepares for InferenceSpec using model_path, writes inference.py, and captures dependencies.
 
     Args:to
         model_path (str) : Argument
@@ -91,7 +91,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():

sagemaker-serve/src/sagemaker/serve/model_server/tgi/server.py

@@ -31,7 +31,6 @@ def _start_tgi_serving(
         client: object,
         image: str,
         model_path: str,
-        secret_key: str,
         env_vars: dict,
         jumpstart: bool,
     ):

sagemaker-serve/tests/unit/model_server/test_multi_model_server_prepare.py

@@ -95,8 +95,8 @@ def test_prepare_for_mms_creates_structure(self, mock_copy, mock_capture, mock_h
             inference_spec=mock_inference_spec
         )
 
-        # Should return empty string now (not secret key)
-        self.assertEqual(secret_key, "")
+        # Should return None now (no longer returns secret key)
+        self.assertIsNone(secret_key)
         mock_inference_spec.prepare.assert_called_once_with(str(model_path))
         mock_capture.assert_called_once()