diff --git a/sagemaker-serve/src/sagemaker/serve/detector/dependency_manager.py b/sagemaker-serve/src/sagemaker/serve/detector/dependency_manager.py index c39b8c0e30..c455c1b168 100644 --- a/sagemaker-serve/src/sagemaker/serve/detector/dependency_manager.py +++ b/sagemaker-serve/src/sagemaker/serve/detector/dependency_manager.py @@ -66,11 +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") - # autodetect_depedencies = [] + # Pin sagemaker to 3.4.0+ to ensure SHA256 hashing is used for integrity checks + autodetect_depedencies.append("sagemaker[huggingface]>=3.4.0,<4.0.0") else: - autodetect_depedencies = [] - # autodetect_depedencies = ["sagemaker[huggingface]>=2.199"] + # Pin sagemaker to 3.4.0+ to ensure SHA256 hashing is used for integrity checks + autodetect_depedencies = ["sagemaker[huggingface]>=3.4.0,<4.0.0"] module_version_dict = _parse_dependency_list(autodetect_depedencies) diff --git a/sagemaker-serve/src/sagemaker/serve/local_resources.py b/sagemaker-serve/src/sagemaker/serve/local_resources.py index efd0c3fdc3..86d2217dd3 100644 --- a/sagemaker-serve/src/sagemaker/serve/local_resources.py +++ b/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, diff --git a/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py b/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py index bb8fdde960..7ab2f860b9 100644 --- a/sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py +++ b/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,7 +107,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, ) elif self.model_server == ModelServer.DJL_SERVING: @@ -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, ) elif self.model_server == ModelServer.TORCHSERVE: @@ -125,7 +121,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, ) elif self.model_server == ModelServer.TGI: @@ -133,7 +128,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, ) @@ -142,7 +136,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, ) elif self.model_server == ModelServer.TENSORFLOW_SERVING: @@ -150,7 +143,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, ) 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 diff --git a/sagemaker-serve/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py b/sagemaker-serve/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py index eac26a3748..30575f5b10 100644 --- a/sagemaker-serve/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py +++ b/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, @@ -77,7 +76,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, @@ -87,7 +85,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, @@ -106,7 +103,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, @@ -132,7 +128,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, ) @@ -150,7 +145,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, diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 2cfaaaca00..d89ae499b8 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/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']: diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py index 831c37ee14..ef18003783 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py +++ b/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, @@ -705,12 +697,6 @@ def _build_for_transformers(self) -> Model: self.s3_model_data_url, _ = self._prepare_for_mode() - # Clean up empty secret key - if ( - "SAGEMAKER_SERVE_SECRET_KEY" in self.env_vars - and not self.env_vars["SAGEMAKER_SERVE_SECRET_KEY"] - ): - del self.env_vars["SAGEMAKER_SERVE_SECRET_KEY"] # Instance type validation for SAGEMAKER_ENDPOINT mode if self.mode == Mode.SAGEMAKER_ENDPOINT and not self.instance_type: @@ -731,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 @@ -767,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 @@ -803,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 @@ -846,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( diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py index d486f24acb..213772dfe6 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py @@ -130,10 +130,7 @@ def build(self): from sagemaker.core.deserializers import JSONDeserializer from sagemaker.serve.detector.pickler import save_pkl from sagemaker.serve.builder.requirements_manager import RequirementsManager -from sagemaker.serve.validations.check_integrity import ( - generate_secret_key, - compute_hash, -) +from sagemaker.serve.validations.check_integrity import compute_hash from sagemaker.core.remote_function.core.serialization import _MetaData from sagemaker.serve.model_server.triton.config_template import CONFIG_TEMPLATE @@ -2882,20 +2879,6 @@ def _save_inference_spec(self) -> None: pkl_path = Path(self.model_path).joinpath("model_repository").joinpath("model") save_pkl(pkl_path, (self.inference_spec, self.schema_builder)) - def _hmac_signing(self): - """Perform HMAC signing on picke file for integrity check""" - secret_key = generate_secret_key() - pkl_path = Path(self.model_path).joinpath("model_repository").joinpath("model") - - with open(str(pkl_path.joinpath("serve.pkl")), "rb") as f: - buffer = f.read() - hash_value = compute_hash(buffer=buffer, secret_key=secret_key) - - with open(str(pkl_path.joinpath("metadata.json")), "wb") as metadata: - metadata.write(_MetaData(hash_value).to_json()) - - self.secret_key = secret_key - def _generate_config_pbtxt(self, pkl_path: Path): """Generate Triton config.pbtxt file.""" config_path = pkl_path.joinpath("config.pbtxt") @@ -3073,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 @@ -3097,7 +3078,12 @@ def _prepare_for_triton(self): self._pack_conda_env(pkl_path=pkl_path) - self._hmac_signing() + # Compute SHA256 hash for integrity check + with open(str(pkl_path.joinpath("serve.pkl")), "rb") as f: + buffer = f.read() + hash_value = compute_hash(buffer=buffer) + with open(str(pkl_path.joinpath("metadata.json")), "wb") as metadata: + metadata.write(_MetaData(hash_value).to_json()) return diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/djl_serving/server.py b/sagemaker-serve/src/sagemaker/serve/model_server/djl_serving/server.py index cfa88a4d4d..db8cbb423b 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/djl_serving/server.py +++ b/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) diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/prepare.py b/sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/prepare.py index 37ca745987..79fcb052f9 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/prepare.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/prepare.py @@ -25,10 +25,7 @@ from sagemaker.core.helper.session_helper 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.core.remote_function.core.serialization import _MetaData logger = logging.getLogger(__name__) @@ -84,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 @@ -94,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(): @@ -119,11 +116,8 @@ 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 diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/server.py b/sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/server.py index 9401dd74d9..531a290bc7 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/server.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/server.py @@ -28,14 +28,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: @@ -80,7 +78,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, @@ -127,15 +124,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) diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.py b/sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.py index b66de32bf7..681416ef99 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.py +++ b/sagemaker-serve/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.core.remote_function.core.serialization import _MetaData from sagemaker.serve.spec.inference_base import CustomOrchestrator, AsyncCustomOrchestrator @@ -64,11 +61,8 @@ 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 diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/smd/server.py b/sagemaker-serve/src/sagemaker/serve/model_server/smd/server.py index e40dc3aa61..6a63365f0c 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/smd/server.py +++ b/sagemaker-serve/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 diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/tei/server.py b/sagemaker-serve/src/sagemaker/serve/model_server/tei/server.py index 9f2f4b71b3..a36d51132c 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/tei/server.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/tei/server.py @@ -27,7 +27,7 @@ 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. @@ -35,12 +35,8 @@ def _start_tei_serving( 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, diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py b/sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py index 3525cc9b4a..abbef79724 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py +++ b/sagemaker-serve/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.core.remote_function.core.serialization import _MetaData @@ -57,11 +54,8 @@ 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 diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/server.py b/sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/server.py index 2f4a959528..e4881d242f 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/server.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/server.py @@ -20,7 +20,7 @@ class LocalTensorflowServing: """LocalTensorflowServing class.""" def _start_tensorflow_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 tensorflow serving container. @@ -28,7 +28,6 @@ def _start_tensorflow_serving( 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 """ self.container = client.containers.run( @@ -47,7 +46,6 @@ def _start_tensorflow_serving( environment={ "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code", "SAGEMAKER_PROGRAM": "inference.py", - "SAGEMAKER_SERVE_SECRET_KEY": secret_key, "LOCAL_PYTHON": platform.python_version(), **env_vars, }, @@ -81,7 +79,6 @@ def _upload_tensorflow_serving_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, @@ -91,7 +88,6 @@ def _upload_tensorflow_serving_artifacts( Args: model_path: Path to the model sagemaker_session: SageMaker session - secret_key: Secret key to use for authentication s3_model_data_url: S3 model data URL image: Image to use model_data_s3_path: S3 model data URI @@ -124,7 +120,6 @@ def _upload_tensorflow_serving_artifacts( "SAGEMAKER_PROGRAM": "inference.py", "SAGEMAKER_REGION": sagemaker_session.boto_region_name, "SAGEMAKER_CONTAINER_LOG_LEVEL": "10", - "SAGEMAKER_SERVE_SECRET_KEY": secret_key, "LOCAL_PYTHON": platform.python_version(), } return s3_upload_path, env_vars diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/tgi/server.py b/sagemaker-serve/src/sagemaker/serve/model_server/tgi/server.py index 6b38c20cda..544cae6b2a 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/tgi/server.py +++ b/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, ): diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/torchserve/prepare.py b/sagemaker-serve/src/sagemaker/serve/model_server/torchserve/prepare.py index 988acf646d..28a0c4c62e 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/torchserve/prepare.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/torchserve/prepare.py @@ -12,10 +12,7 @@ from sagemaker.core.helper.session_helper 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.serve.validations.check_image_uri import is_1p_image_uri from sagemaker.core.remote_function.core.serialization import _MetaData @@ -67,11 +64,8 @@ def prepare_for_torchserve( 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 \ No newline at end of file + metadata.write(_MetaData(hash_value).to_json()) \ No newline at end of file diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/torchserve/server.py b/sagemaker-serve/src/sagemaker/serve/model_server/torchserve/server.py index 0d237df987..b46c9a041c 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/torchserve/server.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/torchserve/server.py @@ -20,7 +20,7 @@ class LocalTorchServe: """Placeholder docstring""" def _start_torch_serve( - 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""" self.container = client.containers.run( @@ -39,7 +39,6 @@ def _start_torch_serve( environment={ "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code", "SAGEMAKER_PROGRAM": "inference.py", - "SAGEMAKER_SERVE_SECRET_KEY": secret_key, "LOCAL_PYTHON": platform.python_version(), **env_vars, }, @@ -69,7 +68,6 @@ def _upload_torchserve_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, @@ -103,7 +101,6 @@ def _upload_torchserve_artifacts( "SAGEMAKER_PROGRAM": "inference.py", "SAGEMAKER_REGION": sagemaker_session.boto_region_name, "SAGEMAKER_CONTAINER_LOG_LEVEL": "10", - "SAGEMAKER_SERVE_SECRET_KEY": secret_key, "LOCAL_PYTHON": platform.python_version(), } return s3_upload_path, env_vars diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/triton/model.py b/sagemaker-serve/src/sagemaker/serve/model_server/triton/model.py index a1c731b0d6..eab916cbd0 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/triton/model.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/triton/model.py @@ -29,8 +29,6 @@ def initialize(self, args: dict) -> None: with open(str(serve_path), mode="rb") as f: inference_spec, schema_builder = cloudpickle.load(f) - # TODO: HMAC signing for integrity check - self.inference_spec = inference_spec self.schema_builder = schema_builder self.model = inference_spec.load(model_dir=TRITON_MODEL_DIR) diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/triton/server.py b/sagemaker-serve/src/sagemaker/serve/model_server/triton/server.py index 134f12dd42..1798115bd7 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/triton/server.py +++ b/sagemaker-serve/src/sagemaker/serve/model_server/triton/server.py @@ -31,7 +31,6 @@ def _start_triton_server( self, docker_client: docker.DockerClient, model_path: str, - secret_key: str, image_uri: str, env_vars: dict, ): @@ -41,7 +40,6 @@ def _start_triton_server( env_vars.update( { "TRITON_MODEL_DIR": "/models/model", - "SAGEMAKER_SERVE_SECRET_KEY": secret_key, "LOCAL_PYTHON": platform.python_version(), } ) @@ -100,7 +98,6 @@ def _upload_triton_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, @@ -133,7 +130,6 @@ def _upload_triton_artifacts( env_vars = { "SAGEMAKER_TRITON_DEFAULT_MODEL_NAME": "model", "TRITON_MODEL_DIR": "/opt/ml/model/model", - "SAGEMAKER_SERVE_SECRET_KEY": secret_key, "LOCAL_PYTHON": platform.python_version(), } return s3_upload_path, env_vars diff --git a/sagemaker-serve/src/sagemaker/serve/validations/check_integrity.py b/sagemaker-serve/src/sagemaker/serve/validations/check_integrity.py index 4363d8d6ed..d6b28e6f31 100644 --- a/sagemaker-serve/src/sagemaker/serve/validations/check_integrity.py +++ b/sagemaker-serve/src/sagemaker/serve/validations/check_integrity.py @@ -1,35 +1,42 @@ -"""Validates the integrity of pickled file with HMAC signing.""" +"""Validates the integrity of pickled file with SHA256 hashing.""" from __future__ import absolute_import -import secrets -import hmac import hashlib -import os from pathlib import Path from sagemaker.core.remote_function.core.serialization import _MetaData -def generate_secret_key(nbytes: int = 32) -> str: - """Generates secret key""" - return secrets.token_hex(nbytes) - - -def compute_hash(buffer: bytes, secret_key: str) -> str: - """Compute hash value using HMAC""" - return hmac.new(secret_key.encode(), msg=buffer, digestmod=hashlib.sha256).hexdigest() +def compute_hash(buffer: bytes) -> str: + """Compute SHA256 hash value of buffer. + + Args: + buffer: Bytes to hash + + Returns: + Hexadecimal SHA256 hash string + """ + return hashlib.sha256(buffer).hexdigest() def perform_integrity_check(buffer: bytes, metadata_path: Path): - """Validates the integrity of bytes by comparing the hash value""" - secret_key = os.environ.get("SAGEMAKER_SERVE_SECRET_KEY") - actual_hash_value = compute_hash(buffer=buffer, secret_key=secret_key) - + """Validates the integrity of bytes by comparing SHA256 hash. + + Args: + buffer: Bytes to verify + metadata_path: Path to metadata.json file + + Raises: + ValueError: If metadata file doesn't exist or hash doesn't match + """ 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 + metadata = _MetaData.from_json(md.read()) + expected_hash_value = metadata.sha256_hash + + actual_hash_value = compute_hash(buffer=buffer) - if not hmac.compare_digest(expected_hash_value, actual_hash_value): + if expected_hash_value != actual_hash_value: raise ValueError("Integrity check for the serialized function or data failed.") diff --git a/sagemaker-serve/tests/unit/model_server/test_multi_model_server_prepare.py b/sagemaker-serve/tests/unit/model_server/test_multi_model_server_prepare.py index d6a571cd1a..a3486f9918 100644 --- a/sagemaker-serve/tests/unit/model_server/test_multi_model_server_prepare.py +++ b/sagemaker-serve/tests/unit/model_server/test_multi_model_server_prepare.py @@ -68,10 +68,9 @@ def test_prepare_mms_js_resources(self, mock_create_dir, mock_copy_js): @patch('builtins.input', return_value='') @patch('sagemaker.serve.model_server.multi_model_server.prepare.compute_hash') - @patch('sagemaker.serve.model_server.multi_model_server.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.multi_model_server.prepare.capture_dependencies') @patch('shutil.copy2') - def test_prepare_for_mms_creates_structure(self, mock_copy, mock_capture, mock_gen_key, mock_hash, mock_input): + def test_prepare_for_mms_creates_structure(self, mock_copy, mock_capture, mock_hash, mock_input): """Test prepare_for_mms creates directory structure and files.""" from sagemaker.serve.model_server.multi_model_server.prepare import prepare_for_mms @@ -83,31 +82,29 @@ def test_prepare_for_mms_creates_structure(self, mock_copy, mock_capture, mock_g serve_pkl = code_dir / "serve.pkl" serve_pkl.write_bytes(b"test data") - mock_gen_key.return_value = "test-secret-key" mock_hash.return_value = "test-hash" mock_session = Mock() mock_inference_spec = Mock() - with patch('builtins.open', mock_open(read_data=b"test data")): - secret_key = prepare_for_mms( - model_path=str(model_path), - shared_libs=[], - dependencies={}, - session=mock_session, - image_uri="test-image", - inference_spec=mock_inference_spec - ) + secret_key = prepare_for_mms( + model_path=str(model_path), + shared_libs=[], + dependencies={}, + session=mock_session, + image_uri="test-image", + inference_spec=mock_inference_spec + ) - self.assertEqual(secret_key, "test-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() @patch('builtins.input', return_value='') @patch('sagemaker.serve.model_server.multi_model_server.prepare.compute_hash') - @patch('sagemaker.serve.model_server.multi_model_server.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.multi_model_server.prepare.capture_dependencies') @patch('shutil.copy2') - def test_prepare_for_mms_raises_on_invalid_dir(self, mock_copy, mock_capture, mock_gen_key, mock_hash, mock_input): + def test_prepare_for_mms_raises_on_invalid_dir(self, mock_copy, mock_capture, mock_hash, mock_input): """Test prepare_for_mms raises exception for invalid directory.""" from sagemaker.serve.model_server.multi_model_server.prepare import prepare_for_mms @@ -128,10 +125,9 @@ def test_prepare_for_mms_raises_on_invalid_dir(self, mock_copy, mock_capture, mo @patch('builtins.input', return_value='') @patch('sagemaker.serve.model_server.multi_model_server.prepare.compute_hash') - @patch('sagemaker.serve.model_server.multi_model_server.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.multi_model_server.prepare.capture_dependencies') @patch('shutil.copy2') - def test_prepare_for_mms_copies_shared_libs(self, mock_copy, mock_capture, mock_gen_key, mock_hash, mock_input): + def test_prepare_for_mms_copies_shared_libs(self, mock_copy, mock_capture, mock_hash, mock_input): """Test prepare_for_mms copies shared libraries.""" from sagemaker.serve.model_server.multi_model_server.prepare import prepare_for_mms @@ -145,7 +141,6 @@ def test_prepare_for_mms_copies_shared_libs(self, mock_copy, mock_capture, mock_ shared_lib = Path(self.temp_dir) / "lib.so" shared_lib.touch() - mock_gen_key.return_value = "test-key" mock_hash.return_value = "test-hash" mock_session = Mock() diff --git a/sagemaker-serve/tests/unit/model_server/test_multi_model_server_server.py b/sagemaker-serve/tests/unit/model_server/test_multi_model_server_server.py index 02ae4dc596..e5fe81b40d 100644 --- a/sagemaker-serve/tests/unit/model_server/test_multi_model_server_server.py +++ b/sagemaker-serve/tests/unit/model_server/test_multi_model_server_server.py @@ -25,15 +25,13 @@ def test_start_serving_creates_container(self, mock_path): client=mock_client, image="test-image:latest", model_path="/path/to/model", - secret_key="test-secret", env_vars={"CUSTOM_VAR": "value"} ) self.assertEqual(server.container, mock_container) mock_client.containers.run.assert_called_once() call_kwargs = mock_client.containers.run.call_args[1] - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", call_kwargs["environment"]) - self.assertEqual(call_kwargs["environment"]["SAGEMAKER_SERVE_SECRET_KEY"], "test-secret") + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", call_kwargs["environment"]) @patch('sagemaker.serve.model_server.multi_model_server.server.Path') def test_start_serving_with_no_env_vars(self, mock_path): @@ -52,7 +50,6 @@ def test_start_serving_with_no_env_vars(self, mock_path): client=mock_client, image="test-image:latest", model_path="/path/to/model", - secret_key="test-secret", env_vars=None ) @@ -121,7 +118,6 @@ def test_upload_server_artifacts_with_s3_path(self, mock_is_s3, mock_fw_utils, m model_data, env_vars = server._upload_server_artifacts( model_path="s3://bucket/model", - secret_key="test-key", sagemaker_session=mock_session, should_upload_artifacts=False ) @@ -158,7 +154,6 @@ def test_upload_server_artifacts_uploads_to_s3(self, mock_path, mock_is_s3, mock model_data, env_vars = server._upload_server_artifacts( model_path="/local/model", - secret_key="test-key", sagemaker_session=mock_session, s3_model_data_url="s3://bucket/prefix", image="test-image", @@ -166,8 +161,7 @@ def test_upload_server_artifacts_uploads_to_s3(self, mock_path, mock_is_s3, mock ) self.assertIsNotNone(model_data) - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) - self.assertEqual(env_vars["SAGEMAKER_SERVE_SECRET_KEY"], "test-key") + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) @patch('sagemaker.serve.model_server.multi_model_server.server._is_s3_uri') def test_upload_server_artifacts_no_upload(self, mock_is_s3): @@ -181,13 +175,12 @@ def test_upload_server_artifacts_no_upload(self, mock_is_s3): model_data, env_vars = server._upload_server_artifacts( model_path="/local/model", - secret_key="test-key", sagemaker_session=mock_session, should_upload_artifacts=False ) self.assertIsNone(model_data) - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) class TestUpdateEnvVars(unittest.TestCase): diff --git a/sagemaker-serve/tests/unit/model_server/test_smd_prepare.py b/sagemaker-serve/tests/unit/model_server/test_smd_prepare.py index 4d5a0a7de8..c1986f3655 100644 --- a/sagemaker-serve/tests/unit/model_server/test_smd_prepare.py +++ b/sagemaker-serve/tests/unit/model_server/test_smd_prepare.py @@ -18,10 +18,9 @@ def tearDown(self): shutil.rmtree(self.temp_dir) @patch('sagemaker.serve.model_server.smd.prepare.compute_hash') - @patch('sagemaker.serve.model_server.smd.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.smd.prepare.capture_dependencies') @patch('shutil.copy2') - def test_prepare_for_smd_with_inference_spec(self, mock_copy, mock_capture, mock_gen_key, mock_hash): + def test_prepare_for_smd_with_inference_spec(self, mock_copy, mock_capture, mock_hash): """Test prepare_for_smd with InferenceSpec.""" from sagemaker.serve.model_server.smd.prepare import prepare_for_smd from sagemaker.serve.spec.inference_spec import InferenceSpec @@ -33,27 +32,24 @@ def test_prepare_for_smd_with_inference_spec(self, mock_copy, mock_capture, mock serve_pkl = code_dir / "serve.pkl" serve_pkl.write_bytes(b"test data") - mock_gen_key.return_value = "test-secret-key" mock_hash.return_value = "test-hash" mock_inference_spec = Mock(spec=InferenceSpec) - with patch('builtins.open', mock_open(read_data=b"test data")): - secret_key = prepare_for_smd( - model_path=str(model_path), - shared_libs=[], - dependencies={}, - inference_spec=mock_inference_spec - ) + secret_key = prepare_for_smd( + model_path=str(model_path), + shared_libs=[], + dependencies={}, + inference_spec=mock_inference_spec + ) - self.assertEqual(secret_key, "test-secret-key") + self.assertIsNone(secret_key) mock_inference_spec.prepare.assert_called_once_with(str(model_path)) @patch('os.rename') @patch('sagemaker.serve.model_server.smd.prepare.compute_hash') - @patch('sagemaker.serve.model_server.smd.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.smd.prepare.capture_dependencies') @patch('shutil.copy2') - def test_prepare_for_smd_with_custom_orchestrator(self, mock_copy, mock_capture, mock_gen_key, mock_hash, mock_rename): + def test_prepare_for_smd_with_custom_orchestrator(self, mock_copy, mock_capture, mock_hash, mock_rename): """Test prepare_for_smd with CustomOrchestrator.""" from sagemaker.serve.model_server.smd.prepare import prepare_for_smd from sagemaker.serve.spec.inference_base import CustomOrchestrator @@ -65,27 +61,24 @@ def test_prepare_for_smd_with_custom_orchestrator(self, mock_copy, mock_capture, serve_pkl = code_dir / "serve.pkl" serve_pkl.write_bytes(b"test data") - mock_gen_key.return_value = "test-secret-key" mock_hash.return_value = "test-hash" mock_orchestrator = Mock(spec=CustomOrchestrator) - with patch('builtins.open', mock_open(read_data=b"test data")): - secret_key = prepare_for_smd( - model_path=str(model_path), - shared_libs=[], - dependencies={}, - inference_spec=mock_orchestrator - ) + secret_key = prepare_for_smd( + model_path=str(model_path), + shared_libs=[], + dependencies={}, + inference_spec=mock_orchestrator + ) - self.assertEqual(secret_key, "test-secret-key") + self.assertIsNone(secret_key) # Verify custom_execution_inference.py was copied and renamed mock_rename.assert_called_once() @patch('sagemaker.serve.model_server.smd.prepare.compute_hash') - @patch('sagemaker.serve.model_server.smd.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.smd.prepare.capture_dependencies') @patch('shutil.copy2') - def test_prepare_for_smd_with_shared_libs(self, mock_copy, mock_capture, mock_gen_key, mock_hash): + def test_prepare_for_smd_with_shared_libs(self, mock_copy, mock_capture, mock_hash): """Test prepare_for_smd copies shared libraries.""" from sagemaker.serve.model_server.smd.prepare import prepare_for_smd @@ -99,7 +92,6 @@ def test_prepare_for_smd_with_shared_libs(self, mock_copy, mock_capture, mock_ge shared_lib = Path(self.temp_dir) / "lib.so" shared_lib.touch() - mock_gen_key.return_value = "test-key" mock_hash.return_value = "test-hash" with patch('builtins.open', mock_open(read_data=b"test data")): diff --git a/sagemaker-serve/tests/unit/model_server/test_smd_server.py b/sagemaker-serve/tests/unit/model_server/test_smd_server.py index 8bf7d4424e..f2aff4eb34 100644 --- a/sagemaker-serve/tests/unit/model_server/test_smd_server.py +++ b/sagemaker-serve/tests/unit/model_server/test_smd_server.py @@ -20,13 +20,11 @@ def test_upload_smd_artifacts_with_s3_path(self, mock_is_s3): s3_path, env_vars = server._upload_smd_artifacts( model_path="s3://bucket/model", sagemaker_session=mock_session, - secret_key="test-key", should_upload_artifacts=False ) self.assertEqual(s3_path, "s3://bucket/model") - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) - self.assertEqual(env_vars["SAGEMAKER_SERVE_SECRET_KEY"], "test-key") + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) self.assertIn("SAGEMAKER_INFERENCE_CODE_DIRECTORY", env_vars) @patch('sagemaker.serve.model_server.smd.server.upload') @@ -51,14 +49,13 @@ def test_upload_smd_artifacts_uploads_to_s3(self, mock_is_s3, mock_fw_utils, s3_path, env_vars = server._upload_smd_artifacts( model_path="/local/model", sagemaker_session=mock_session, - secret_key="test-key", s3_model_data_url="s3://bucket/prefix", image="test-image", should_upload_artifacts=True ) self.assertEqual(s3_path, "s3://bucket/code_prefix/model.tar.gz") - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) self.assertIn("SAGEMAKER_INFERENCE_CODE", env_vars) mock_upload.assert_called_once() @@ -75,12 +72,11 @@ def test_upload_smd_artifacts_no_upload(self, mock_is_s3): s3_path, env_vars = server._upload_smd_artifacts( model_path="/local/model", sagemaker_session=mock_session, - secret_key="test-key", should_upload_artifacts=False ) self.assertIsNone(s3_path) - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) if __name__ == "__main__": diff --git a/sagemaker-serve/tests/unit/model_server/test_tei_server.py b/sagemaker-serve/tests/unit/model_server/test_tei_server.py index c280e4b546..67ffb5c62f 100644 --- a/sagemaker-serve/tests/unit/model_server/test_tei_server.py +++ b/sagemaker-serve/tests/unit/model_server/test_tei_server.py @@ -29,7 +29,6 @@ def test_start_tei_serving(self, mock_device_req, mock_path, mock_update_env): client=mock_client, image="tei:latest", model_path="/path/to/model", - secret_key="test-secret", env_vars={"CUSTOM_VAR": "value"} ) @@ -40,7 +39,7 @@ def test_start_tei_serving(self, mock_device_req, mock_path, mock_update_env): @patch('sagemaker.serve.model_server.tei.server.Path') @patch('sagemaker.serve.model_server.tei.server.DeviceRequest') def test_start_tei_serving_adds_secret_key(self, mock_device_req, mock_path, mock_update_env): - """Test _start_tei_serving adds secret key to env vars.""" + """Test _start_tei_serving does not add secret key to env vars.""" from sagemaker.serve.model_server.tei.server import LocalTeiServing server = LocalTeiServing() @@ -58,12 +57,11 @@ def test_start_tei_serving_adds_secret_key(self, mock_device_req, mock_path, moc client=mock_client, image="tei:latest", model_path="/path/to/model", - secret_key="test-secret", env_vars=env_vars ) - # Verify secret key was added to env_vars - self.assertEqual(env_vars["SAGEMAKER_SERVE_SECRET_KEY"], "test-secret") + # Verify secret key was NOT added to env_vars + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) @patch('sagemaker.serve.model_server.tei.server.requests.post') @patch('sagemaker.serve.model_server.tei.server.get_docker_host') diff --git a/sagemaker-serve/tests/unit/model_server/test_tensorflow_serving_prepare.py b/sagemaker-serve/tests/unit/model_server/test_tensorflow_serving_prepare.py index e6ca1161dc..f7abc7932d 100644 --- a/sagemaker-serve/tests/unit/model_server/test_tensorflow_serving_prepare.py +++ b/sagemaker-serve/tests/unit/model_server/test_tensorflow_serving_prepare.py @@ -20,11 +20,10 @@ def tearDown(self): @patch('sagemaker.serve.model_server.tensorflow_serving.prepare._move_contents') @patch('sagemaker.serve.model_server.tensorflow_serving.prepare._get_saved_model_path_for_tensorflow_and_keras_flavor') @patch('sagemaker.serve.model_server.tensorflow_serving.prepare.compute_hash') - @patch('sagemaker.serve.model_server.tensorflow_serving.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.tensorflow_serving.prepare.capture_dependencies') @patch('shutil.copy2') - def test_prepare_for_tf_serving_success(self, mock_copy, mock_capture, mock_gen_key, - mock_hash, mock_get_saved, mock_move): + def test_prepare_for_tf_serving_success(self, mock_copy, mock_capture, mock_hash, + mock_get_saved, mock_move): """Test prepare_for_tf_serving creates structure successfully.""" from sagemaker.serve.model_server.tensorflow_serving.prepare import prepare_for_tf_serving @@ -35,28 +34,25 @@ def test_prepare_for_tf_serving_success(self, mock_copy, mock_capture, mock_gen_ serve_pkl = code_dir / "serve.pkl" serve_pkl.write_bytes(b"test data") - mock_gen_key.return_value = "test-secret-key" mock_hash.return_value = "test-hash" mock_get_saved.return_value = Path(self.temp_dir) / "saved_model" - with patch('builtins.open', mock_open(read_data=b"test data")): - secret_key = prepare_for_tf_serving( - model_path=str(model_path), - shared_libs=[], - dependencies={} - ) + secret_key = prepare_for_tf_serving( + model_path=str(model_path), + shared_libs=[], + dependencies={} + ) - self.assertEqual(secret_key, "test-secret-key") + self.assertIsNone(secret_key) mock_capture.assert_called_once() mock_move.assert_called_once() @patch('sagemaker.serve.model_server.tensorflow_serving.prepare._get_saved_model_path_for_tensorflow_and_keras_flavor') @patch('sagemaker.serve.model_server.tensorflow_serving.prepare.compute_hash') - @patch('sagemaker.serve.model_server.tensorflow_serving.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.tensorflow_serving.prepare.capture_dependencies') @patch('shutil.copy2') - def test_prepare_for_tf_serving_no_saved_model(self, mock_copy, mock_capture, mock_gen_key, - mock_hash, mock_get_saved): + def test_prepare_for_tf_serving_no_saved_model(self, mock_copy, mock_capture, mock_hash, + mock_get_saved): """Test prepare_for_tf_serving raises error when SavedModel not found.""" from sagemaker.serve.model_server.tensorflow_serving.prepare import prepare_for_tf_serving @@ -67,7 +63,6 @@ def test_prepare_for_tf_serving_no_saved_model(self, mock_copy, mock_capture, mo serve_pkl = code_dir / "serve.pkl" serve_pkl.write_bytes(b"test data") - mock_gen_key.return_value = "test-secret-key" mock_hash.return_value = "test-hash" mock_get_saved.return_value = None @@ -82,11 +77,10 @@ def test_prepare_for_tf_serving_no_saved_model(self, mock_copy, mock_capture, mo @patch('sagemaker.serve.model_server.tensorflow_serving.prepare._move_contents') @patch('sagemaker.serve.model_server.tensorflow_serving.prepare._get_saved_model_path_for_tensorflow_and_keras_flavor') @patch('sagemaker.serve.model_server.tensorflow_serving.prepare.compute_hash') - @patch('sagemaker.serve.model_server.tensorflow_serving.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.tensorflow_serving.prepare.capture_dependencies') @patch('shutil.copy2') - def test_prepare_for_tf_serving_with_shared_libs(self, mock_copy, mock_capture, mock_gen_key, - mock_hash, mock_get_saved, mock_move): + def test_prepare_for_tf_serving_with_shared_libs(self, mock_copy, mock_capture, mock_hash, + mock_get_saved, mock_move): """Test prepare_for_tf_serving copies shared libraries.""" from sagemaker.serve.model_server.tensorflow_serving.prepare import prepare_for_tf_serving @@ -100,27 +94,24 @@ def test_prepare_for_tf_serving_with_shared_libs(self, mock_copy, mock_capture, shared_lib = Path(self.temp_dir) / "lib.so" shared_lib.touch() - mock_gen_key.return_value = "test-key" mock_hash.return_value = "test-hash" mock_get_saved.return_value = Path(self.temp_dir) / "saved_model" - with patch('builtins.open', mock_open(read_data=b"test data")): - prepare_for_tf_serving( - model_path=str(model_path), - shared_libs=[str(shared_lib)], - dependencies={} - ) + prepare_for_tf_serving( + model_path=str(model_path), + shared_libs=[str(shared_lib)], + dependencies={} + ) # Verify copy2 was called for shared lib self.assertTrue(any(str(shared_lib) in str(call) for call in mock_copy.call_args_list)) @patch('sagemaker.serve.model_server.tensorflow_serving.prepare._get_saved_model_path_for_tensorflow_and_keras_flavor') @patch('sagemaker.serve.model_server.tensorflow_serving.prepare.compute_hash') - @patch('sagemaker.serve.model_server.tensorflow_serving.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.tensorflow_serving.prepare.capture_dependencies') @patch('shutil.copy2') - def test_prepare_for_tf_serving_invalid_dir(self, mock_copy, mock_capture, mock_gen_key, - mock_hash, mock_get_saved): + def test_prepare_for_tf_serving_invalid_dir(self, mock_copy, mock_capture, mock_hash, + mock_get_saved): """Test prepare_for_tf_serving raises exception for invalid directory.""" from sagemaker.serve.model_server.tensorflow_serving.prepare import prepare_for_tf_serving diff --git a/sagemaker-serve/tests/unit/model_server/test_tensorflow_serving_server.py b/sagemaker-serve/tests/unit/model_server/test_tensorflow_serving_server.py index d0bac2e5dc..c6a8b46890 100644 --- a/sagemaker-serve/tests/unit/model_server/test_tensorflow_serving_server.py +++ b/sagemaker-serve/tests/unit/model_server/test_tensorflow_serving_server.py @@ -25,15 +25,13 @@ def test_start_tensorflow_serving(self, mock_path): client=mock_client, image="tensorflow-serving:latest", model_path="/path/to/model", - secret_key="test-secret", env_vars={"CUSTOM_VAR": "value"} ) self.assertEqual(server.container, mock_container) mock_client.containers.run.assert_called_once() call_kwargs = mock_client.containers.run.call_args[1] - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", call_kwargs["environment"]) - self.assertEqual(call_kwargs["environment"]["SAGEMAKER_SERVE_SECRET_KEY"], "test-secret") + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", call_kwargs["environment"]) self.assertEqual(call_kwargs["environment"]["CUSTOM_VAR"], "value") @patch('sagemaker.serve.model_server.tensorflow_serving.server.requests.post') @@ -92,13 +90,11 @@ def test_upload_tensorflow_serving_artifacts_with_s3_path(self, mock_is_s3): s3_path, env_vars = server._upload_tensorflow_serving_artifacts( model_path="s3://bucket/model", sagemaker_session=mock_session, - secret_key="test-key", should_upload_artifacts=False ) self.assertEqual(s3_path, "s3://bucket/model") - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) - self.assertEqual(env_vars["SAGEMAKER_SERVE_SECRET_KEY"], "test-key") + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) @patch('sagemaker.serve.model_server.tensorflow_serving.server.upload') @patch('sagemaker.serve.model_server.tensorflow_serving.server.determine_bucket_and_prefix') @@ -122,14 +118,13 @@ def test_upload_tensorflow_serving_artifacts_uploads_to_s3(self, mock_is_s3, moc s3_path, env_vars = server._upload_tensorflow_serving_artifacts( model_path="/local/model", sagemaker_session=mock_session, - secret_key="test-key", s3_model_data_url="s3://bucket/prefix", image="test-image", should_upload_artifacts=True ) self.assertEqual(s3_path, "s3://bucket/code_prefix/model.tar.gz") - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) mock_upload.assert_called_once() @patch('sagemaker.serve.model_server.tensorflow_serving.server._is_s3_uri') @@ -145,12 +140,11 @@ def test_upload_tensorflow_serving_artifacts_no_upload(self, mock_is_s3): s3_path, env_vars = server._upload_tensorflow_serving_artifacts( model_path="/local/model", sagemaker_session=mock_session, - secret_key="test-key", should_upload_artifacts=False ) self.assertIsNone(s3_path) - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) if __name__ == "__main__": diff --git a/sagemaker-serve/tests/unit/model_server/test_tgi_server.py b/sagemaker-serve/tests/unit/model_server/test_tgi_server.py index 12f84a747b..02872f4c7e 100644 --- a/sagemaker-serve/tests/unit/model_server/test_tgi_server.py +++ b/sagemaker-serve/tests/unit/model_server/test_tgi_server.py @@ -27,7 +27,6 @@ def test_start_tgi_serving_jumpstart(self, mock_device_req, mock_path): client=mock_client, image="test-image:latest", model_path="/path/to/model", - secret_key="test-secret", env_vars={"CUSTOM_VAR": "value"}, jumpstart=True ) @@ -59,7 +58,6 @@ def test_start_tgi_serving_non_jumpstart(self, mock_device_req, mock_path, mock_ client=mock_client, image="test-image:latest", model_path="/path/to/model", - secret_key="test-secret", env_vars={"CUSTOM_VAR": "value"}, jumpstart=False ) diff --git a/sagemaker-serve/tests/unit/model_server/test_torchserve_prepare.py b/sagemaker-serve/tests/unit/model_server/test_torchserve_prepare.py index 1ae35eca6a..f8bd02b45c 100644 --- a/sagemaker-serve/tests/unit/model_server/test_torchserve_prepare.py +++ b/sagemaker-serve/tests/unit/model_server/test_torchserve_prepare.py @@ -18,12 +18,11 @@ def tearDown(self): shutil.rmtree(self.temp_dir) @patch('sagemaker.serve.model_server.torchserve.prepare.compute_hash') - @patch('sagemaker.serve.model_server.torchserve.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.torchserve.prepare.capture_dependencies') @patch('sagemaker.serve.model_server.torchserve.prepare.is_1p_image_uri') @patch('shutil.copy2') def test_prepare_for_torchserve_standard_image(self, mock_copy, mock_is_1p, mock_capture, - mock_gen_key, mock_hash): + mock_hash): """Test prepare_for_torchserve with standard image.""" from sagemaker.serve.model_server.torchserve.prepare import prepare_for_torchserve @@ -35,33 +34,30 @@ def test_prepare_for_torchserve_standard_image(self, mock_copy, mock_is_1p, mock serve_pkl.write_bytes(b"test data") mock_is_1p.return_value = True - mock_gen_key.return_value = "test-secret-key" mock_hash.return_value = "test-hash" mock_session = Mock() mock_inference_spec = Mock() - with patch('builtins.open', mock_open(read_data=b"test data")): - secret_key = prepare_for_torchserve( - model_path=str(model_path), - shared_libs=[], - dependencies={}, - session=mock_session, - image_uri="test-pytorch-image", - inference_spec=mock_inference_spec - ) + secret_key = prepare_for_torchserve( + model_path=str(model_path), + shared_libs=[], + dependencies={}, + session=mock_session, + image_uri="test-pytorch-image", + inference_spec=mock_inference_spec + ) - self.assertEqual(secret_key, "test-secret-key") + self.assertIsNone(secret_key) mock_inference_spec.prepare.assert_called_once_with(str(model_path)) mock_capture.assert_called_once() @patch('os.rename') @patch('sagemaker.serve.model_server.torchserve.prepare.compute_hash') - @patch('sagemaker.serve.model_server.torchserve.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.torchserve.prepare.capture_dependencies') @patch('sagemaker.serve.model_server.torchserve.prepare.is_1p_image_uri') @patch('shutil.copy2') def test_prepare_for_torchserve_xgboost_image(self, mock_copy, mock_is_1p, mock_capture, - mock_gen_key, mock_hash, mock_rename): + mock_hash, mock_rename): """Test prepare_for_torchserve with xgboost image.""" from sagemaker.serve.model_server.torchserve.prepare import prepare_for_torchserve @@ -73,31 +69,28 @@ def test_prepare_for_torchserve_xgboost_image(self, mock_copy, mock_is_1p, mock_ serve_pkl.write_bytes(b"test data") mock_is_1p.return_value = True - mock_gen_key.return_value = "test-secret-key" mock_hash.return_value = "test-hash" mock_session = Mock() - with patch('builtins.open', mock_open(read_data=b"test data")): - secret_key = prepare_for_torchserve( - model_path=str(model_path), - shared_libs=[], - dependencies={}, - session=mock_session, - image_uri="xgboost-image:latest", - inference_spec=None - ) + secret_key = prepare_for_torchserve( + model_path=str(model_path), + shared_libs=[], + dependencies={}, + session=mock_session, + image_uri="xgboost-image", + inference_spec=None + ) - self.assertEqual(secret_key, "test-secret-key") - # Verify xgboost_inference.py was copied and renamed + self.assertIsNone(secret_key) mock_rename.assert_called_once() + mock_capture.assert_called_once() @patch('sagemaker.serve.model_server.torchserve.prepare.compute_hash') - @patch('sagemaker.serve.model_server.torchserve.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.torchserve.prepare.capture_dependencies') @patch('sagemaker.serve.model_server.torchserve.prepare.is_1p_image_uri') @patch('shutil.copy2') def test_prepare_for_torchserve_with_shared_libs(self, mock_copy, mock_is_1p, mock_capture, - mock_gen_key, mock_hash): + mock_hash): """Test prepare_for_torchserve copies shared libraries.""" from sagemaker.serve.model_server.torchserve.prepare import prepare_for_torchserve @@ -112,7 +105,6 @@ def test_prepare_for_torchserve_with_shared_libs(self, mock_copy, mock_is_1p, mo shared_lib.touch() mock_is_1p.return_value = False - mock_gen_key.return_value = "test-key" mock_hash.return_value = "test-hash" mock_session = Mock() @@ -149,12 +141,11 @@ def test_prepare_for_torchserve_invalid_dir(self, mock_is_1p): self.assertIn("not a valid directory", str(context.exception)) @patch('sagemaker.serve.model_server.torchserve.prepare.compute_hash') - @patch('sagemaker.serve.model_server.torchserve.prepare.generate_secret_key') @patch('sagemaker.serve.model_server.torchserve.prepare.capture_dependencies') @patch('sagemaker.serve.model_server.torchserve.prepare.is_1p_image_uri') @patch('shutil.copy2') def test_prepare_for_torchserve_no_inference_spec(self, mock_copy, mock_is_1p, mock_capture, - mock_gen_key, mock_hash): + mock_hash): """Test prepare_for_torchserve without inference_spec.""" from sagemaker.serve.model_server.torchserve.prepare import prepare_for_torchserve @@ -166,21 +157,19 @@ def test_prepare_for_torchserve_no_inference_spec(self, mock_copy, mock_is_1p, m serve_pkl.write_bytes(b"test data") mock_is_1p.return_value = False - mock_gen_key.return_value = "test-key" mock_hash.return_value = "test-hash" mock_session = Mock() - with patch('builtins.open', mock_open(read_data=b"test data")): - secret_key = prepare_for_torchserve( - model_path=str(model_path), - shared_libs=[], - dependencies={}, - session=mock_session, - image_uri="test-image", - inference_spec=None - ) + secret_key = prepare_for_torchserve( + model_path=str(model_path), + shared_libs=[], + dependencies={}, + session=mock_session, + image_uri="test-image", + inference_spec=None + ) - self.assertEqual(secret_key, "test-key") + self.assertIsNone(secret_key) if __name__ == "__main__": diff --git a/sagemaker-serve/tests/unit/model_server/test_torchserve_server.py b/sagemaker-serve/tests/unit/model_server/test_torchserve_server.py index 95b0645076..24283c8168 100644 --- a/sagemaker-serve/tests/unit/model_server/test_torchserve_server.py +++ b/sagemaker-serve/tests/unit/model_server/test_torchserve_server.py @@ -25,15 +25,13 @@ def test_start_torch_serve(self, mock_path): client=mock_client, image="torchserve:latest", model_path="/path/to/model", - secret_key="test-secret", env_vars={"CUSTOM_VAR": "value"} ) self.assertEqual(server.container, mock_container) mock_client.containers.run.assert_called_once() call_kwargs = mock_client.containers.run.call_args[1] - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", call_kwargs["environment"]) - self.assertEqual(call_kwargs["environment"]["SAGEMAKER_SERVE_SECRET_KEY"], "test-secret") + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", call_kwargs["environment"]) self.assertEqual(call_kwargs["environment"]["CUSTOM_VAR"], "value") @patch('sagemaker.serve.model_server.torchserve.server.requests.post') @@ -92,13 +90,11 @@ def test_upload_torchserve_artifacts_with_s3_path(self, mock_is_s3): s3_path, env_vars = server._upload_torchserve_artifacts( model_path="s3://bucket/model", sagemaker_session=mock_session, - secret_key="test-key", should_upload_artifacts=False ) self.assertEqual(s3_path, "s3://bucket/model") - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) - self.assertEqual(env_vars["SAGEMAKER_SERVE_SECRET_KEY"], "test-key") + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) @patch('sagemaker.serve.model_server.torchserve.server.upload') @patch('sagemaker.serve.model_server.torchserve.server.determine_bucket_and_prefix') @@ -122,14 +118,13 @@ def test_upload_torchserve_artifacts_uploads_to_s3(self, mock_is_s3, mock_fw_uti s3_path, env_vars = server._upload_torchserve_artifacts( model_path="/local/model", sagemaker_session=mock_session, - secret_key="test-key", s3_model_data_url="s3://bucket/prefix", image="test-image", should_upload_artifacts=True ) self.assertEqual(s3_path, "s3://bucket/code_prefix/model.tar.gz") - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) mock_upload.assert_called_once() @patch('sagemaker.serve.model_server.torchserve.server._is_s3_uri') @@ -145,12 +140,11 @@ def test_upload_torchserve_artifacts_no_upload(self, mock_is_s3): s3_path, env_vars = server._upload_torchserve_artifacts( model_path="/local/model", sagemaker_session=mock_session, - secret_key="test-key", should_upload_artifacts=False ) self.assertIsNone(s3_path) - self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) if __name__ == "__main__": diff --git a/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py b/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py index 4a57147ae9..318d320650 100644 --- a/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py +++ b/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py @@ -233,13 +233,12 @@ def test_build_local_container_mode(self, mock_create, mock_prepare, mock_detect """Test building for LOCAL_CONTAINER mode.""" self.builder.mode = Mode.LOCAL_CONTAINER self.builder.model = Mock() - mock_ts_prepare.return_value = "secret123" + mock_ts_prepare.return_value = None mock_create.return_value = Mock() result = self.builder._build_for_torchserve() mock_ts_prepare.assert_called_once() - self.assertEqual(self.builder.secret_key, "secret123") mock_create.assert_called_once() @patch('sagemaker.serve.model_builder_servers.prepare_for_torchserve') @@ -251,14 +250,13 @@ def test_build_sagemaker_endpoint_mode(self, mock_create, mock_prepare, mock_det """Test building for SAGEMAKER_ENDPOINT mode.""" self.builder.mode = Mode.SAGEMAKER_ENDPOINT self.builder.model = Mock() - mock_ts_prepare.return_value = "secret456" + mock_ts_prepare.return_value = None mock_create.return_value = Mock() mock_prepare.return_value = ("s3://bucket/model.tar.gz", None) result = self.builder._build_for_torchserve() mock_ts_prepare.assert_called_once() - self.assertEqual(self.builder.secret_key, "secret456") mock_prepare.assert_called_with(should_upload_artifacts=True) @@ -527,13 +525,12 @@ def setUp(self): @patch.object(MockModelBuilderServers, '_create_model') def test_build_mlflow_model(self, mock_create, mock_prepare_mode, mock_tf_prepare, mock_save): """Test building MLflow model for TensorFlow Serving.""" - mock_tf_prepare.return_value = "secret789" + mock_tf_prepare.return_value = None mock_create.return_value = Mock() mock_prepare_mode.return_value = ("s3://bucket/model.tar.gz", None) result = self.builder._build_for_tensorflow_serving() - self.assertEqual(self.builder.secret_key, "secret789") mock_save.assert_called_once() mock_create.assert_called_once() @@ -622,7 +619,7 @@ def test_build_with_auto_image(self, mock_create, mock_prepare_mode, mock_get_im """Test building with auto-detected image.""" mock_get_unit.return_value = "gpu" mock_get_img.return_value = "smd-image-uri" - mock_smd_prepare.return_value = "secret999" + mock_smd_prepare.return_value = None mock_create.return_value = Mock() self.builder.mode = Mode.LOCAL_CONTAINER self.builder.image_uri = None @@ -631,7 +628,6 @@ def test_build_with_auto_image(self, mock_create, mock_prepare_mode, mock_get_im result = self.builder._build_for_smd() self.assertEqual(self.builder.image_uri, "smd-image-uri") - self.assertEqual(self.builder.secret_key, "secret999") mock_create.assert_called_once() @@ -655,7 +651,7 @@ def test_build_with_inference_spec_local_container(self, mock_create, mock_prepa mock_nb, mock_mms_prepare, mock_save): """Test building with inference_spec for LOCAL_CONTAINER.""" mock_nb.return_value = None - mock_mms_prepare.return_value = "secret111" + mock_mms_prepare.return_value = None mock_create.return_value = Mock() self.builder.mode = Mode.LOCAL_CONTAINER self.builder.inference_spec = Mock() @@ -664,7 +660,6 @@ def test_build_with_inference_spec_local_container(self, mock_create, mock_prepa mock_save.assert_called_once() mock_mms_prepare.assert_called_once() - self.assertEqual(self.builder.secret_key, "secret111") mock_create.assert_called_once() @patch('sagemaker.serve.model_builder_servers._get_model_config_properties_from_hf') @@ -715,12 +710,11 @@ def test_build_sagemaker_endpoint_missing_instance_type(self, mock_create, mock_ @patch.object(MockModelBuilderServers, '_create_model') def test_build_clean_empty_secret_key(self, mock_create, mock_prepare, mock_detect, mock_dir, mock_nb): - """Test cleaning empty secret key from env_vars.""" + """Test that secret key is not added to env_vars.""" mock_nb.return_value = None mock_create.return_value = Mock() mock_prepare.return_value = ("s3://bucket/model.tar.gz", None) self.builder.model = Mock() - self.builder.env_vars["SAGEMAKER_SERVE_SECRET_KEY"] = "" result = self.builder._build_for_transformers() diff --git a/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py b/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py index 1f9aa2f3d4..707e6a59e4 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py +++ b/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py @@ -219,7 +219,6 @@ def test_reset_build_state(self): mb._reset_build_state() self.assertIsNone(mb.built_model) - self.assertEqual(mb.secret_key, "") self.assertFalse(hasattr(mb, 'prepared_for_djl')) self.assertFalse(hasattr(mb, 'modes')) diff --git a/sagemaker-serve/tests/unit/test_model_builder_deploy.py b/sagemaker-serve/tests/unit/test_model_builder_deploy.py index 3a0fca3d8e..8a08208a4d 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_deploy.py +++ b/sagemaker-serve/tests/unit/test_model_builder_deploy.py @@ -544,7 +544,6 @@ def test_reset_build_state_clears_built_model(self): builder._reset_build_state() self.assertIsNone(builder.built_model) - self.assertEqual(builder.secret_key, "") def test_reset_build_state_clears_jumpstart_flags(self): """Test _reset_build_state clears JumpStart preparation flags.""" diff --git a/sagemaker-serve/tests/unit/test_model_builder_utils_triton.py b/sagemaker-serve/tests/unit/test_model_builder_utils_triton.py index bb0d1d874c..7cd4d94b20 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_utils_triton.py +++ b/sagemaker-serve/tests/unit/test_model_builder_utils_triton.py @@ -113,8 +113,7 @@ def test_prepare_for_triton_tensorflow(self, mock_export, mock_copy): @patch('shutil.copy2') @patch.object(_ModelBuilderUtils, '_generate_config_pbtxt') @patch.object(_ModelBuilderUtils, '_pack_conda_env') - @patch.object(_ModelBuilderUtils, '_hmac_signing') - def test_prepare_for_triton_inference_spec(self, mock_hmac, mock_pack, mock_config, mock_copy): + def test_prepare_for_triton_inference_spec(self, mock_pack, mock_config, mock_copy): """Test preparing inference spec for Triton.""" utils = _ModelBuilderUtils() utils.inference_spec = Mock() @@ -123,11 +122,17 @@ def test_prepare_for_triton_inference_spec(self, mock_hmac, mock_pack, mock_conf with tempfile.TemporaryDirectory() as tmpdir: utils.model_path = tmpdir + # Create the model_repository/model directory structure + model_dir = Path(tmpdir) / "model_repository" / "model" + model_dir.mkdir(parents=True) + # Create serve.pkl file + (model_dir / "serve.pkl").write_bytes(b"test data") + utils._prepare_for_triton() mock_config.assert_called_once() mock_pack.assert_called_once() - mock_hmac.assert_called_once() + # SHA256 hashing is now done inline, not via _hmac_signing class TestExportPytorchToOnnx(unittest.TestCase): @@ -262,10 +267,10 @@ def test_save_inference_spec(self): class TestHMACSignin(unittest.TestCase): - """Test _hmac_signing method.""" + """Test SHA256 hashing for integrity checks.""" - def test_hmac_signing(self): - """Test HMAC signing.""" + def test_sha256_hashing(self): + """Test SHA256 hashing for integrity.""" utils = _ModelBuilderUtils() with tempfile.TemporaryDirectory() as tmpdir: @@ -276,11 +281,12 @@ def test_hmac_signing(self): # Create dummy serve.pkl (pkl_path / "serve.pkl").write_bytes(b"dummy content") - utils._hmac_signing() - - # Secret key is generated, not mocked - self.assertIsNotNone(utils.secret_key) - self.assertTrue((pkl_path / "metadata.json").exists()) + # SHA256 hashing should be done inline now, not via _hmac_signing + # Just verify metadata.json can be created + from sagemaker.serve.validations.check_integrity import compute_hash + hash_value = compute_hash(b"dummy content") + self.assertIsNotNone(hash_value) + self.assertEqual(len(hash_value), 64) # SHA256 hex is 64 chars class TestAutoDetectImageForTriton(unittest.TestCase): diff --git a/sagemaker-serve/tests/unit/validations/test_check_integrity.py b/sagemaker-serve/tests/unit/validations/test_check_integrity.py index 11e66eb716..b4e4cbb7ae 100644 --- a/sagemaker-serve/tests/unit/validations/test_check_integrity.py +++ b/sagemaker-serve/tests/unit/validations/test_check_integrity.py @@ -3,37 +3,24 @@ from pathlib import Path from unittest.mock import patch, mock_open from sagemaker.serve.validations.check_integrity import ( - generate_secret_key, compute_hash, perform_integrity_check ) class TestCheckIntegrity(unittest.TestCase): - def test_generate_secret_key(self): - key = generate_secret_key() - self.assertIsInstance(key, str) - self.assertEqual(len(key), 64) - - def test_generate_secret_key_custom_bytes(self): - key = generate_secret_key(nbytes=16) - self.assertEqual(len(key), 32) - def test_compute_hash(self): buffer = b"test data" - secret_key = "test_secret" - hash_value = compute_hash(buffer, secret_key) + hash_value = compute_hash(buffer) self.assertIsInstance(hash_value, str) self.assertEqual(len(hash_value), 64) def test_compute_hash_consistency(self): buffer = b"test data" - secret_key = "test_secret" - hash1 = compute_hash(buffer, secret_key) - hash2 = compute_hash(buffer, secret_key) + hash1 = compute_hash(buffer) + hash2 = compute_hash(buffer) self.assertEqual(hash1, hash2) - @patch.dict("os.environ", {"SAGEMAKER_SERVE_SECRET_KEY": "test_key"}) @patch("pathlib.Path.exists") @patch("builtins.open", new_callable=mock_open, read_data=b'{"sha256_hash": "test_hash"}') @patch("sagemaker.serve.validations.check_integrity._MetaData.from_json")