diff --git a/src/sagemaker/serve/builder/djl_builder.py b/src/sagemaker/serve/builder/djl_builder.py index 9b1ebf1257..90921467df 100644 --- a/src/sagemaker/serve/builder/djl_builder.py +++ b/src/sagemaker/serve/builder/djl_builder.py @@ -213,7 +213,6 @@ def _djl_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa self.modes[str(Mode.LOCAL_CONTAINER)].create_server( self.image_uri, timeout if timeout else 1800, - self.secret_key, predictor, self.pysdk_model.env, ) diff --git a/src/sagemaker/serve/builder/model_builder.py b/src/sagemaker/serve/builder/model_builder.py index 3c19e4aa43..21f0a75b25 100644 --- a/src/sagemaker/serve/builder/model_builder.py +++ b/src/sagemaker/serve/builder/model_builder.py @@ -448,7 +448,6 @@ def _prepare_for_mode( str(Mode.SAGEMAKER_ENDPOINT) ].prepare( (model_path or self.model_path), - self.secret_key, self.serve_settings.s3_model_data_url, self.sagemaker_session, self.image_uri, @@ -709,7 +708,7 @@ def _model_builder_deploy_wrapper( ) self.modes[str(Mode.LOCAL_CONTAINER)].create_server( - self.image_uri, container_timeout_in_second, self.secret_key, predictor + self.image_uri, container_timeout_in_second, predictor ) return predictor @@ -778,7 +777,7 @@ def _build_for_torchserve(self) -> Type[Model]: if self.mode != Mode.IN_PROCESS: self._auto_detect_container() - self.secret_key = prepare_for_torchserve( + prepare_for_torchserve( model_path=self.model_path, shared_libs=self.shared_libs, dependencies=self.dependencies, @@ -798,7 +797,7 @@ def _build_for_smd(self) -> Type[Model]: if self.mode != Mode.IN_PROCESS: self._auto_detect_container() - self.secret_key = prepare_for_smd( + prepare_for_smd( model_path=self.model_path, shared_libs=self.shared_libs, dependencies=self.dependencies, diff --git a/src/sagemaker/serve/builder/tf_serving_builder.py b/src/sagemaker/serve/builder/tf_serving_builder.py index 044e0460bc..563f2301f3 100644 --- a/src/sagemaker/serve/builder/tf_serving_builder.py +++ b/src/sagemaker/serve/builder/tf_serving_builder.py @@ -124,7 +124,7 @@ def _build_for_tensorflow_serving(self): if not self.image_uri: raise ValueError("image_uri is not set for tensorflow serving") - self.secret_key = prepare_for_tf_serving( + prepare_for_tf_serving( model_path=self.model_path, shared_libs=self.shared_libs, dependencies=self.dependencies, diff --git a/src/sagemaker/serve/builder/transformers_builder.py b/src/sagemaker/serve/builder/transformers_builder.py index 0388a9a05d..cada65ee78 100644 --- a/src/sagemaker/serve/builder/transformers_builder.py +++ b/src/sagemaker/serve/builder/transformers_builder.py @@ -267,12 +267,6 @@ def _transformers_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[Pr self.env_vars.update(env_vars) self.pysdk_model.env.update(self.env_vars) - if ( - "SAGEMAKER_SERVE_SECRET_KEY" in self.pysdk_model.env - and not self.pysdk_model.env["SAGEMAKER_SERVE_SECRET_KEY"] - ): - del self.pysdk_model.env["SAGEMAKER_SERVE_SECRET_KEY"] - if "endpoint_logging" not in kwargs: kwargs["endpoint_logging"] = True @@ -402,7 +396,7 @@ def _build_for_transformers(self): self._auto_detect_container() - self.secret_key = prepare_for_mms( + prepare_for_mms( model_path=self.model_path, shared_libs=self.shared_libs, dependencies=self.dependencies, diff --git a/src/sagemaker/serve/detector/dependency_manager.py b/src/sagemaker/serve/detector/dependency_manager.py index 8ff37c9185..9f301ff8f0 100644 --- a/src/sagemaker/serve/detector/dependency_manager.py +++ b/src/sagemaker/serve/detector/dependency_manager.py @@ -66,9 +66,11 @@ def capture_dependencies(dependencies: dict, work_dir: Path, capture_all: bool = with open(path, "r") as f: autodetect_depedencies = f.read().splitlines() - autodetect_depedencies.append("sagemaker[huggingface]>=2.199") + # Pin sagemaker to 2.257.0+ to ensure SHA256 hashing is used for integrity checks + autodetect_depedencies.append("sagemaker[huggingface]>=2.257.0,<3.0.0") else: - autodetect_depedencies = ["sagemaker[huggingface]>=2.199"] + # Pin sagemaker to 2.257.0+ to ensure SHA256 hashing is used for integrity checks + autodetect_depedencies = ["sagemaker[huggingface]>=2.257.0,<3.0.0"] module_version_dict = _parse_dependency_list(autodetect_depedencies) diff --git a/src/sagemaker/serve/mode/local_container_mode.py b/src/sagemaker/serve/mode/local_container_mode.py index f040c61c1d..94b39e99ee 100644 --- a/src/sagemaker/serve/mode/local_container_mode.py +++ b/src/sagemaker/serve/mode/local_container_mode.py @@ -68,7 +68,6 @@ def __init__( self.model_server = model_server self.client = None self.container = None - self.secret_key = None self._ping_container = None self._invoke_serving = None @@ -89,7 +88,6 @@ def create_server( self, image: str, container_timeout_seconds: int, - secret_key: str, predictor: PredictorBase, env_vars: Dict[str, str] = None, model_path: str = None, @@ -108,7 +106,6 @@ def create_server( docker_client=self.client, model_path=model_path if model_path else self.model_path, image_uri=image, - secret_key=secret_key, env_vars=env_vars if env_vars else self.env_vars, ) self._ping_container = self._triton_deep_ping @@ -117,7 +114,6 @@ def create_server( client=self.client, image=image, model_path=model_path if model_path else self.model_path, - secret_key=secret_key, env_vars=env_vars if env_vars else self.env_vars, ) self._ping_container = self._djl_deep_ping @@ -126,7 +122,6 @@ def create_server( client=self.client, image=image, model_path=model_path if model_path else self.model_path, - secret_key=secret_key, env_vars=env_vars if env_vars else self.env_vars, ) self._ping_container = self._torchserve_deep_ping @@ -135,7 +130,6 @@ def create_server( client=self.client, image=image, model_path=model_path if model_path else self.model_path, - secret_key=secret_key, env_vars=env_vars if env_vars else self.env_vars, jumpstart=jumpstart, ) @@ -145,7 +139,6 @@ def create_server( client=self.client, image=image, model_path=model_path if model_path else self.model_path, - secret_key=secret_key, env_vars=env_vars if env_vars else self.env_vars, ) self._ping_container = self._multi_model_server_deep_ping @@ -154,7 +147,6 @@ def create_server( client=self.client, image=image, model_path=model_path if model_path else self.model_path, - secret_key=secret_key, env_vars=env_vars if env_vars else self.env_vars, ) self._ping_container = self._tensorflow_serving_deep_ping @@ -164,7 +156,6 @@ def create_server( client=self.client, image=image, model_path=model_path if model_path else self.model_path, - secret_key=secret_key, env_vars=env_vars if env_vars else self.env_vars, ) tei_serving.schema_builder = self.schema_builder diff --git a/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py b/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py index 2b4473a706..e8f4e51f9e 100644 --- a/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py +++ b/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py @@ -58,7 +58,6 @@ def load(self, model_path: str): def prepare( self, model_path: str, - secret_key: str, s3_model_data_url: str = None, sagemaker_session: Session = None, image: str = None, @@ -79,7 +78,6 @@ def prepare( upload_artifacts = self._upload_torchserve_artifacts( model_path=model_path, sagemaker_session=sagemaker_session, - secret_key=secret_key, s3_model_data_url=s3_model_data_url, image=image, should_upload_artifacts=True, @@ -89,7 +87,6 @@ def prepare( upload_artifacts = self._upload_triton_artifacts( model_path=model_path, sagemaker_session=sagemaker_session, - secret_key=secret_key, s3_model_data_url=s3_model_data_url, image=image, should_upload_artifacts=True, @@ -108,7 +105,6 @@ def prepare( upload_artifacts = self._upload_tensorflow_serving_artifacts( model_path=model_path, sagemaker_session=sagemaker_session, - secret_key=secret_key, s3_model_data_url=s3_model_data_url, image=image, should_upload_artifacts=True, @@ -134,7 +130,6 @@ def prepare( model_path=model_path, sagemaker_session=sagemaker_session, s3_model_data_url=s3_model_data_url, - secret_key=secret_key, image=image, should_upload_artifacts=should_upload_artifacts, ) @@ -152,7 +147,6 @@ def prepare( upload_artifacts = self._upload_smd_artifacts( model_path=model_path, sagemaker_session=sagemaker_session, - secret_key=secret_key, s3_model_data_url=s3_model_data_url, image=image, should_upload_artifacts=True, diff --git a/src/sagemaker/serve/model_server/djl_serving/server.py b/src/sagemaker/serve/model_server/djl_serving/server.py index 4ba7dd227d..cd511c8591 100644 --- a/src/sagemaker/serve/model_server/djl_serving/server.py +++ b/src/sagemaker/serve/model_server/djl_serving/server.py @@ -31,7 +31,7 @@ class LocalDJLServing: """Placeholder docstring""" def _start_djl_serving( - self, client: object, image: str, model_path: str, secret_key: str, env_vars: dict + self, client: object, image: str, model_path: str, env_vars: dict ): """Placeholder docstring""" updated_env_vars = _update_env_vars(env_vars) diff --git a/src/sagemaker/serve/model_server/multi_model_server/prepare.py b/src/sagemaker/serve/model_server/multi_model_server/prepare.py index e3abc70dd6..12c7ad23e4 100644 --- a/src/sagemaker/serve/model_server/multi_model_server/prepare.py +++ b/src/sagemaker/serve/model_server/multi_model_server/prepare.py @@ -25,10 +25,7 @@ from sagemaker.session import Session from sagemaker.serve.spec.inference_spec import InferenceSpec from sagemaker.serve.detector.dependency_manager import capture_dependencies -from sagemaker.serve.validations.check_integrity import ( - generate_secret_key, - compute_hash, -) +from sagemaker.serve.validations.check_integrity import compute_hash from sagemaker.remote_function.core.serialization import _MetaData logger = logging.getLogger(__name__) @@ -85,7 +82,7 @@ def prepare_for_mms( inference_spec: InferenceSpec = None, ) -> str: """Prepares for InferenceSpec using model_path, writes inference.py, \ - and captures dependencies to generate secret_key. + and captures dependencies. Args:to model_path (str) : Argument @@ -95,7 +92,7 @@ def prepare_for_mms( inference_spec (InferenceSpec, optional) : Argument (default is None) Returns: - ( str ) : secret_key + ( str ) : Empty string for backward compatibility """ model_path = Path(model_path) if not model_path.exists(): @@ -120,11 +117,10 @@ def prepare_for_mms( capture_dependencies(dependencies=dependencies, work_dir=code_dir) - secret_key = generate_secret_key() with open(str(code_dir.joinpath("serve.pkl")), "rb") as f: buffer = f.read() - hash_value = compute_hash(buffer=buffer, secret_key=secret_key) + hash_value = compute_hash(buffer=buffer) with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata: metadata.write(_MetaData(hash_value).to_json()) - return secret_key + return "" diff --git a/src/sagemaker/serve/model_server/multi_model_server/server.py b/src/sagemaker/serve/model_server/multi_model_server/server.py index 2fab727c05..1358a3c0dc 100644 --- a/src/sagemaker/serve/model_server/multi_model_server/server.py +++ b/src/sagemaker/serve/model_server/multi_model_server/server.py @@ -29,14 +29,12 @@ def _start_serving( client: object, image: str, model_path: str, - secret_key: str, env_vars: dict, ): """Initializes the start of the server""" env = { "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code", "SAGEMAKER_PROGRAM": "inference.py", - "SAGEMAKER_SERVE_SECRET_KEY": secret_key, "LOCAL_PYTHON": platform.python_version(), } if env_vars: @@ -94,7 +92,6 @@ class SageMakerMultiModelServer: def _upload_server_artifacts( self, model_path: str, - secret_key: str, sagemaker_session: Session, s3_model_data_url: str = None, image: str = None, @@ -141,15 +138,16 @@ def _upload_server_artifacts( else None ) - if secret_key: - env_vars = { - "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code", - "SAGEMAKER_PROGRAM": "inference.py", - "SAGEMAKER_SERVE_SECRET_KEY": secret_key, - "SAGEMAKER_REGION": sagemaker_session.boto_region_name, - "SAGEMAKER_CONTAINER_LOG_LEVEL": "10", - "LOCAL_PYTHON": platform.python_version(), - } + if env_vars is None: + env_vars = {} + + env_vars.update({ + "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code", + "SAGEMAKER_PROGRAM": "inference.py", + "SAGEMAKER_REGION": sagemaker_session.boto_region_name, + "SAGEMAKER_CONTAINER_LOG_LEVEL": "10", + "LOCAL_PYTHON": platform.python_version(), + }) return model_data, _update_env_vars(env_vars) diff --git a/src/sagemaker/serve/model_server/smd/prepare.py b/src/sagemaker/serve/model_server/smd/prepare.py index 6461e4023f..67aa1e890f 100644 --- a/src/sagemaker/serve/model_server/smd/prepare.py +++ b/src/sagemaker/serve/model_server/smd/prepare.py @@ -11,10 +11,7 @@ from sagemaker.serve.spec.inference_spec import InferenceSpec from sagemaker.serve.detector.dependency_manager import capture_dependencies -from sagemaker.serve.validations.check_integrity import ( - generate_secret_key, - compute_hash, -) +from sagemaker.serve.validations.check_integrity import compute_hash from sagemaker.remote_function.core.serialization import _MetaData from sagemaker.serve.spec.inference_base import CustomOrchestrator, AsyncCustomOrchestrator @@ -64,11 +61,10 @@ def prepare_for_smd( capture_dependencies(dependencies=dependencies, work_dir=code_dir) - secret_key = generate_secret_key() with open(str(code_dir.joinpath("serve.pkl")), "rb") as f: buffer = f.read() - hash_value = compute_hash(buffer=buffer, secret_key=secret_key) + hash_value = compute_hash(buffer=buffer) with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata: metadata.write(_MetaData(hash_value).to_json()) - return secret_key + return "" diff --git a/src/sagemaker/serve/model_server/smd/server.py b/src/sagemaker/serve/model_server/smd/server.py index c700c39727..22471e4c62 100644 --- a/src/sagemaker/serve/model_server/smd/server.py +++ b/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/src/sagemaker/serve/model_server/tei/server.py b/src/sagemaker/serve/model_server/tei/server.py index 94265e224f..8d91eee520 100644 --- a/src/sagemaker/serve/model_server/tei/server.py +++ b/src/sagemaker/serve/model_server/tei/server.py @@ -28,7 +28,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. @@ -36,12 +36,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/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py b/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py index e9aa4aafff..450425d169 100644 --- a/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py +++ b/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py @@ -10,10 +10,7 @@ _move_contents, ) from sagemaker.serve.detector.dependency_manager import capture_dependencies -from sagemaker.serve.validations.check_integrity import ( - generate_secret_key, - compute_hash, -) +from sagemaker.serve.validations.check_integrity import compute_hash from sagemaker.remote_function.core.serialization import _MetaData @@ -57,11 +54,10 @@ def prepare_for_tf_serving( raise ValueError("SavedModel is not found for Tensorflow or Keras flavor.") _move_contents(src_dir=mlflow_saved_model_dir, dest_dir=saved_model_bundle_dir) - secret_key = generate_secret_key() with open(str(code_dir.joinpath("serve.pkl")), "rb") as f: buffer = f.read() - hash_value = compute_hash(buffer=buffer, secret_key=secret_key) + hash_value = compute_hash(buffer=buffer) with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata: metadata.write(_MetaData(hash_value).to_json()) - return secret_key + return "" diff --git a/src/sagemaker/serve/model_server/tensorflow_serving/server.py b/src/sagemaker/serve/model_server/tensorflow_serving/server.py index 45931e9afc..75f977b55b 100644 --- a/src/sagemaker/serve/model_server/tensorflow_serving/server.py +++ b/src/sagemaker/serve/model_server/tensorflow_serving/server.py @@ -22,7 +22,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. @@ -30,7 +30,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( @@ -48,7 +47,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, }, @@ -99,7 +97,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, @@ -109,7 +106,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 @@ -142,7 +138,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/src/sagemaker/serve/model_server/tgi/server.py b/src/sagemaker/serve/model_server/tgi/server.py index 8ccc8e7ddc..0147423664 100644 --- a/src/sagemaker/serve/model_server/tgi/server.py +++ b/src/sagemaker/serve/model_server/tgi/server.py @@ -32,7 +32,6 @@ def _start_tgi_serving( client: object, image: str, model_path: str, - secret_key: str, env_vars: dict, jumpstart: bool, ): diff --git a/src/sagemaker/serve/model_server/torchserve/prepare.py b/src/sagemaker/serve/model_server/torchserve/prepare.py index 530ff16f6d..5357aade98 100644 --- a/src/sagemaker/serve/model_server/torchserve/prepare.py +++ b/src/sagemaker/serve/model_server/torchserve/prepare.py @@ -12,10 +12,7 @@ from sagemaker.session import Session from sagemaker.serve.spec.inference_spec import InferenceSpec from sagemaker.serve.detector.dependency_manager import capture_dependencies -from sagemaker.serve.validations.check_integrity import ( - generate_secret_key, - compute_hash, -) +from sagemaker.serve.validations.check_integrity import compute_hash from sagemaker.serve.validations.check_image_uri import is_1p_image_uri from sagemaker.remote_function.core.serialization import _MetaData @@ -69,11 +66,10 @@ 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 + return "" diff --git a/src/sagemaker/serve/model_server/torchserve/server.py b/src/sagemaker/serve/model_server/torchserve/server.py index 74e37cd70b..8876da862f 100644 --- a/src/sagemaker/serve/model_server/torchserve/server.py +++ b/src/sagemaker/serve/model_server/torchserve/server.py @@ -22,7 +22,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( @@ -40,7 +40,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, }, @@ -82,7 +81,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, @@ -116,7 +114,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/src/sagemaker/serve/model_server/triton/server.py b/src/sagemaker/serve/model_server/triton/server.py index e2f3c20d7a..5674d1e9fd 100644 --- a/src/sagemaker/serve/model_server/triton/server.py +++ b/src/sagemaker/serve/model_server/triton/server.py @@ -33,7 +33,6 @@ def _start_triton_server( self, docker_client: docker.DockerClient, model_path: str, - secret_key: str, image_uri: str, env_vars: dict, ): @@ -43,7 +42,6 @@ def _start_triton_server( env_vars.update( { "TRITON_MODEL_DIR": "/models/model", - "SAGEMAKER_SERVE_SECRET_KEY": secret_key, "LOCAL_PYTHON": platform.python_version(), } ) @@ -113,7 +111,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, @@ -146,7 +143,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/src/sagemaker/serve/model_server/triton/triton_builder.py b/src/sagemaker/serve/model_server/triton/triton_builder.py index c47991fa09..274586a59f 100644 --- a/src/sagemaker/serve/model_server/triton/triton_builder.py +++ b/src/sagemaker/serve/model_server/triton/triton_builder.py @@ -22,10 +22,7 @@ from sagemaker.base_deserializers import JSONDeserializer from sagemaker.serve.detector.pickler import save_pkl from sagemaker.serve.model_server.triton.config_template import CONFIG_TEMPLATE -from sagemaker.serve.validations.check_integrity import ( - generate_secret_key, - compute_hash, -) +from sagemaker.serve.validations.check_integrity import compute_hash from sagemaker.remote_function.core.serialization import _MetaData @@ -213,8 +210,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 == "pytorch": self._export_pytorch_to_onnx( export_path=export_path, model=self.model, schema_builder=self.schema_builder @@ -237,26 +232,17 @@ 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 raise ValueError("Either model or inference_spec should be provided to ModelBuilder.") - 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): config_path = pkl_path.joinpath("config.pbtxt") diff --git a/src/sagemaker/serve/save_retrive/version_1_0_0/save/framework/pytorch_handler.py b/src/sagemaker/serve/save_retrive/version_1_0_0/save/framework/pytorch_handler.py index 7aa00f65b9..a61c173083 100644 --- a/src/sagemaker/serve/save_retrive/version_1_0_0/save/framework/pytorch_handler.py +++ b/src/sagemaker/serve/save_retrive/version_1_0_0/save/framework/pytorch_handler.py @@ -7,7 +7,6 @@ from pathlib import Path from typing import Type from sagemaker.serve.save_retrive.version_1_0_0.save.utils import ( - generate_secret_key, compute_hash, save_pkl, save_yaml, @@ -139,25 +138,25 @@ def save_model(self) -> None: def _pytorch_generate_hmac(self) -> None: """Placeholder docstring""" - logger.info("Generating Pytorch model hmac...") - self.secret_key = generate_secret_key() + logger.info("Generating Pytorch model hash...") + self.secret_key = "" if self.model: with open(Path(f"{self.model_path}.{self.model_format}").absolute(), "rb") as f: buffer = f.read() - self.model_hmac = compute_hash(buffer=buffer, secret_key=self.secret_key) + self.model_hmac = compute_hash(buffer=buffer) with open(Path(f"{self.schema_path}.{self.schema_format}").absolute(), "rb") as f: buffer = f.read() - self.schema_hmac = compute_hash(buffer=buffer, secret_key=self.secret_key) + self.schema_hmac = compute_hash(buffer=buffer) if self.inference_spec: with open( Path(f"{self.inference_spec_path}.{self.inference_spec_format}").absolute(), "rb" ) as f: buffer = f.read() - self.inference_spec_hmac = compute_hash(buffer=buffer, secret_key=self.secret_key) + self.inference_spec_hmac = compute_hash(buffer=buffer) - logger.info("Pytorch model hmac generated successfully") + logger.info("Pytorch model hash generated successfully") def get_pysdk_model( self, s3_path: str, role_arn: str, sagemaker_session: Session @@ -166,7 +165,6 @@ def get_pysdk_model( self.env_vars = { "SAGEMAKER_REGION": sagemaker_session.boto_region_name, "SAGEMAKER_CONTAINER_LOG_LEVEL": "10", - "SAGEMAKER_SERVE_SECRET_KEY": self.secret_key, "LOCAL_PYTHON": platform.python_version(), } diff --git a/src/sagemaker/serve/save_retrive/version_1_0_0/save/framework/xgboost_handler.py b/src/sagemaker/serve/save_retrive/version_1_0_0/save/framework/xgboost_handler.py index 19e83f58ce..0adffc79fe 100644 --- a/src/sagemaker/serve/save_retrive/version_1_0_0/save/framework/xgboost_handler.py +++ b/src/sagemaker/serve/save_retrive/version_1_0_0/save/framework/xgboost_handler.py @@ -6,7 +6,6 @@ from pathlib import Path from typing import Type from sagemaker.serve.save_retrive.version_1_0_0.save.utils import ( - generate_secret_key, compute_hash, save_pkl, save_yaml, @@ -124,26 +123,26 @@ def save_model(self) -> None: def _xgboost_generate_hmac(self) -> None: """Placeholder docstring""" - logger.info("Generating XGBoost model hmac...") - self.secret_key = generate_secret_key() + logger.info("Generating XGBoost model hash...") + self.secret_key = "" if self.model: with open(Path(f"{self.model_path}.{self.model_format}").absolute(), "rb") as f: buffer = f.read() - self.model_hmac = compute_hash(buffer=buffer, secret_key=self.secret_key) + self.model_hmac = compute_hash(buffer=buffer) with open(Path(f"{self.schema_path}.{self.schema_format}").absolute(), "rb") as f: buffer = f.read() - self.schema_hmac = compute_hash(buffer=buffer, secret_key=self.secret_key) + self.schema_hmac = compute_hash(buffer=buffer) if self.inference_spec: with open( Path(f"{self.inference_spec_path}.{self.inference_spec_format}").absolute(), "rb" ) as f: buffer = f.read() - self.inference_spec_hmac = compute_hash(buffer=buffer, secret_key=self.secret_key) + self.inference_spec_hmac = compute_hash(buffer=buffer) - logger.info("XGBoost model hmac generated successfully") + logger.info("XGBoost model hash generated successfully") def get_pysdk_model( self, s3_path: str, role_arn: str, sagemaker_session: Session @@ -152,7 +151,6 @@ def get_pysdk_model( self.env_vars = { "SAGEMAKER_REGION": sagemaker_session.boto_region_name, "SAGEMAKER_CONTAINER_LOG_LEVEL": "10", - "SAGEMAKER_SERVE_SECRET_KEY": self.secret_key, "LOCAL_PYTHON": platform.python_version(), } diff --git a/src/sagemaker/serve/save_retrive/version_1_0_0/save/utils.py b/src/sagemaker/serve/save_retrive/version_1_0_0/save/utils.py index 8df442d9f8..75ace8c5ad 100644 --- a/src/sagemaker/serve/save_retrive/version_1_0_0/save/utils.py +++ b/src/sagemaker/serve/save_retrive/version_1_0_0/save/utils.py @@ -127,14 +127,16 @@ def capture_optimization_metadata(model: Model, framework) -> dict: return optimizer_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): # pylint: disable=W0613 diff --git a/src/sagemaker/serve/validations/check_integrity.py b/src/sagemaker/serve/validations/check_integrity.py index 01b958611c..007d66cc36 100644 --- a/src/sagemaker/serve/validations/check_integrity.py +++ b/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.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/tests/unit/sagemaker/serve/model_server/multi_model_server/test_multi_model_server_prepare.py b/tests/unit/sagemaker/serve/model_server/multi_model_server/test_multi_model_server_prepare.py index 567a72182a..fcc129b63d 100644 --- a/tests/unit/sagemaker/serve/model_server/multi_model_server/test_multi_model_server_prepare.py +++ b/tests/unit/sagemaker/serve/model_server/multi_model_server/test_multi_model_server_prepare.py @@ -52,7 +52,6 @@ def test_start_invoke_destroy_local_multi_model_server(self): local_multi_model_server._start_serving( client=mock_docker_client, model_path=MODEL_PATH, - secret_key=SECRET_KEY, env_vars=ENV_VAR, image=CPU_TF_IMAGE, ) @@ -68,7 +67,6 @@ def test_start_invoke_destroy_local_multi_model_server(self): "KEY": "VALUE", "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code", "SAGEMAKER_PROGRAM": "inference.py", - "SAGEMAKER_SERVE_SECRET_KEY": "secret_key", "LOCAL_PYTHON": platform.python_version(), }, ) diff --git a/tests/unit/sagemaker/serve/model_server/tei/test_server.py b/tests/unit/sagemaker/serve/model_server/tei/test_server.py index 47399c1fad..b0444006b3 100644 --- a/tests/unit/sagemaker/serve/model_server/tei/test_server.py +++ b/tests/unit/sagemaker/serve/model_server/tei/test_server.py @@ -51,7 +51,6 @@ def test_start_invoke_destroy_local_tei_server(self, mock_requests): local_tei_server._start_tei_serving( client=mock_docker_client, model_path=MODEL_PATH, - secret_key=SECRET_KEY, image=TEI_IMAGE, env_vars=ENV_VAR, ) @@ -68,7 +67,6 @@ def test_start_invoke_destroy_local_tei_server(self, mock_requests): "HF_HOME": "/opt/ml/model/", "HUGGINGFACE_HUB_CACHE": "/opt/ml/model/", "KEY": "VALUE", - "SAGEMAKER_SERVE_SECRET_KEY": "secret_key", }, ) diff --git a/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_prepare.py b/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_prepare.py index 9915b19649..126f6b68f0 100644 --- a/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_prepare.py +++ b/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_prepare.py @@ -43,7 +43,6 @@ def setUp(self): ) @patch("sagemaker.serve.model_server.tensorflow_serving.prepare._MetaData") @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("sagemaker.serve.model_server.tensorflow_serving.prepare.shutil") @patch("sagemaker.serve.model_server.tensorflow_serving.prepare.Path") @@ -52,7 +51,6 @@ def test_prepare_happy( mock_path, mock_shutil, mock_capture_dependencies, - mock_generate_secret_key, mock_compute_hash, mock_metadata, mock_get_saved_model_path, @@ -65,8 +63,6 @@ def test_prepare_happy( mock_path_instance.joinpath.return_value = Mock() mock_get_saved_model_path.return_value = MODEL_PATH + "/1/" - mock_generate_secret_key.return_value = SECRET_KEY - secret_key = prepare_for_tf_serving( model_path=MODEL_PATH, shared_libs=SHARED_LIBS, @@ -74,7 +70,7 @@ def test_prepare_happy( ) mock_path_instance.mkdir.assert_not_called() - self.assertEqual(secret_key, SECRET_KEY) + self.assertEqual(secret_key, "") @patch("builtins.open", new_callable=mock_open, read_data=b"{}") @patch("sagemaker.serve.model_server.tensorflow_serving.prepare._move_contents") @@ -84,7 +80,6 @@ def test_prepare_happy( ) @patch("sagemaker.serve.model_server.tensorflow_serving.prepare._MetaData") @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("sagemaker.serve.model_server.tensorflow_serving.prepare.shutil") @patch("sagemaker.serve.model_server.tensorflow_serving.prepare.Path") @@ -93,7 +88,6 @@ def test_prepare_saved_model_not_found( mock_path, mock_shutil, mock_capture_dependencies, - mock_generate_secret_key, mock_compute_hash, mock_metadata, mock_get_saved_model_path, diff --git a/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_server.py b/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_server.py index b9cce13dbb..f76bd2d2c9 100644 --- a/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_server.py +++ b/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_server.py @@ -51,7 +51,6 @@ def test_start_invoke_destroy_local_tensorflow_serving_server(self): local_tensorflow_server._start_tensorflow_serving( client=mock_docker_client, model_path=MODEL_PATH, - secret_key=SECRET_KEY, env_vars=ENV_VAR, image=CPU_TF_IMAGE, ) @@ -66,7 +65,6 @@ def test_start_invoke_destroy_local_tensorflow_serving_server(self): environment={ "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code", "SAGEMAKER_PROGRAM": "inference.py", - "SAGEMAKER_SERVE_SECRET_KEY": "secret_key", "LOCAL_PYTHON": platform.python_version(), "KEY": "VALUE", }, @@ -89,7 +87,6 @@ def test_upload_artifacts_sagemaker_triton_server(self, mock_upload, mock_platfo ) = SageMakerTensorflowServing()._upload_tensorflow_serving_artifacts( model_path=MODEL_PATH, sagemaker_session=mock_session, - secret_key=SECRET_KEY, s3_model_data_url=S3_URI, image=CPU_TF_IMAGE, should_upload_artifacts=True, @@ -97,5 +94,5 @@ def test_upload_artifacts_sagemaker_triton_server(self, mock_upload, mock_platfo mock_upload.assert_called_once_with(mock_session, MODEL_PATH, "mock_model_data_uri", ANY) self.assertEqual(s3_upload_path, S3_URI) - self.assertEqual(env_vars.get("SAGEMAKER_SERVE_SECRET_KEY"), SECRET_KEY) + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) self.assertEqual(env_vars.get("LOCAL_PYTHON"), "3.8") diff --git a/tests/unit/sagemaker/serve/model_server/torchserve/test_prepare.py b/tests/unit/sagemaker/serve/model_server/torchserve/test_prepare.py index 3c5379b69d..46f53e6638 100644 --- a/tests/unit/sagemaker/serve/model_server/torchserve/test_prepare.py +++ b/tests/unit/sagemaker/serve/model_server/torchserve/test_prepare.py @@ -37,7 +37,6 @@ def setUp(self): @patch("builtins.open", new_callable=mock_open, read_data=b"{}") @patch("sagemaker.serve.model_server.torchserve.prepare._MetaData") @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.shutil") @patch("sagemaker.serve.model_server.torchserve.prepare.Path") @@ -46,7 +45,6 @@ def test_prepare_happy( mock_path, mock_shutil, mock_capture_dependencies, - mock_generate_secret_key, mock_compute_hash, mock_metadata, mock_open, @@ -56,8 +54,6 @@ def test_prepare_happy( mock_path_instance.exists.return_value = True mock_path_instance.joinpath.return_value = Mock() - mock_generate_secret_key.return_value = SECRET_KEY - secret_key = prepare_for_torchserve( model_path=MODEL_PATH, shared_libs=SHARED_LIBS, @@ -69,13 +65,12 @@ def test_prepare_happy( mock_path_instance.mkdir.assert_not_called() INFERENCE_SPEC.prepare.assert_called_once() - self.assertEqual(secret_key, SECRET_KEY) + self.assertEqual(secret_key, "") @patch("os.rename") @patch("builtins.open", new_callable=mock_open, read_data=b"{}") @patch("sagemaker.serve.model_server.torchserve.prepare._MetaData") @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.shutil") @patch("sagemaker.serve.model_server.torchserve.prepare.Path") @@ -84,7 +79,6 @@ def test_prepare_happy_xgboost( mock_path, mock_shutil, mock_capture_dependencies, - mock_generate_secret_key, mock_compute_hash, mock_metadata, mock_open, @@ -95,8 +89,6 @@ def test_prepare_happy_xgboost( mock_path_instance.exists.return_value = True mock_path_instance.joinpath.return_value = Mock() - mock_generate_secret_key.return_value = SECRET_KEY - secret_key = prepare_for_torchserve( model_path=MODEL_PATH, shared_libs=SHARED_LIBS, @@ -109,4 +101,4 @@ def test_prepare_happy_xgboost( mock_rename.assert_called_once() mock_path_instance.mkdir.assert_not_called() INFERENCE_SPEC.prepare.assert_called_once() - self.assertEqual(secret_key, SECRET_KEY) + self.assertEqual(secret_key, "") diff --git a/tests/unit/sagemaker/serve/model_server/triton/test_server.py b/tests/unit/sagemaker/serve/model_server/triton/test_server.py index 3f571424ed..5d13f5d5d7 100644 --- a/tests/unit/sagemaker/serve/model_server/triton/test_server.py +++ b/tests/unit/sagemaker/serve/model_server/triton/test_server.py @@ -60,7 +60,6 @@ def test_start_invoke_destroy_local_triton_server_gpu(self, mock_importlib): local_triton_server._start_triton_server( docker_client=mock_docker_client, model_path=MODEL_PATH, - secret_key=SECRET_KEY, image_uri=GPU_TRITON_IMAGE, env_vars=ENV_VAR, ) @@ -118,7 +117,6 @@ def test_start_invoke_destroy_local_triton_server_cpu(self, mock_importlib): local_triton_server._start_triton_server( docker_client=mock_docker_client, model_path=MODEL_PATH, - secret_key=SECRET_KEY, image_uri=CPU_TRITON_IMAGE, env_vars=ENV_VAR, ) @@ -169,7 +167,6 @@ def test_upload_artifacts_sagemaker_triton_server(self, mock_upload, mock_platfo s3_upload_path, env_vars = SageMakerTritonServer()._upload_triton_artifacts( model_path=MODEL_PATH, sagemaker_session=mock_session, - secret_key=SECRET_KEY, s3_model_data_url=S3_URI, image=GPU_TRITON_IMAGE, should_upload_artifacts=True, @@ -179,5 +176,5 @@ def test_upload_artifacts_sagemaker_triton_server(self, mock_upload, mock_platfo self.assertEqual(s3_upload_path, S3_URI) self.assertEqual(env_vars.get("SAGEMAKER_TRITON_DEFAULT_MODEL_NAME"), "model") self.assertEqual(env_vars.get("TRITON_MODEL_DIR"), "/opt/ml/model/model") - self.assertEqual(env_vars.get("SAGEMAKER_SERVE_SECRET_KEY"), SECRET_KEY) + self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars) self.assertEqual(env_vars.get("LOCAL_PYTHON"), "3.8")