Skip to content

Bug fix for integrity check using hmac key in serve #5482

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aviruthen wants to merge 5 commits into
base: master
Choose a base branch
from
Open

Bug fix for integrity check using hmac key in serve #5482

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
8d55ad5
Bug fix for integrity check using hmac key in serve
aviruthen Jan 13, 2026
74d7905
Removing unneeded comment
aviruthen Jan 13, 2026
e3412cf
Fix unit and integ tests
aviruthen Jan 14, 2026
998cd6a
Removing all secret key references, updating unit tests
aviruthen Jan 15, 2026
47befe3
Pin sm version as auto-dependency
aviruthen Jan 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions sagemaker-serve/src/sagemaker/serve/detector/dependency_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
6 changes: 0 additions & 6 deletions sagemaker-serve/src/sagemaker/serve/local_resources.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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
Expand Down Expand Up @@ -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",
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
9 changes: 0 additions & 9 deletions sagemaker-serve/src/sagemaker/serve/mode/local_container_mode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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,
Expand All @@ -109,31 +107,27 @@ def create_server(
docker_client=self.client,
model_path=model_path if model_path else self.model_path,
image_uri=image,
secret_key=secret_key,
env_vars=env_vars if env_vars else self.env_vars,
)
elif self.model_server == ModelServer.DJL_SERVING:
self._start_djl_serving(
client=self.client,
image=image,
model_path=model_path if model_path else self.model_path,
secret_key=secret_key,
env_vars=env_vars if env_vars else self.env_vars,
)
elif self.model_server == ModelServer.TORCHSERVE:
self._start_torch_serve(
client=self.client,
image=image,
model_path=model_path if model_path else self.model_path,
secret_key=secret_key,
env_vars=env_vars if env_vars else self.env_vars,
)
elif self.model_server == ModelServer.TGI:
self._start_tgi_serving(
client=self.client,
image=image,
model_path=model_path if model_path else self.model_path,
secret_key=secret_key,
env_vars=env_vars if env_vars else self.env_vars,
jumpstart=jumpstart,
)
Expand All @@ -142,15 +136,13 @@ def create_server(
client=self.client,
image=image,
model_path=model_path if model_path else self.model_path,
secret_key=secret_key,
env_vars=env_vars if env_vars else self.env_vars,
)
elif self.model_server == ModelServer.TENSORFLOW_SERVING:
self._start_tensorflow_serving(
client=self.client,
image=image,
model_path=model_path if model_path else self.model_path,
secret_key=secret_key,
env_vars=env_vars if env_vars else self.env_vars,
)
elif self.model_server == ModelServer.TEI:
Expand All @@ -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
Expand Down
6 changes: 0 additions & 6 deletions sagemaker-serve/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
)
Expand All @@ -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,
Expand Down
4 changes: 0 additions & 4 deletions sagemaker-serve/src/sagemaker/serve/model_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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']:
Expand Down
30 changes: 6 additions & 24 deletions sagemaker-serve/src/sagemaker/serve/model_builder_servers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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.")

Expand All @@ -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,
Expand All @@ -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:
Expand Down Expand Up @@ -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()

Expand All @@ -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,
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down
Loading
Loading