Skip to content

Commit ff838db

Browse files
author
Pravali Uppugunduri
committed
fix: Update all model server prepare.py to use plain SHA-256
Remove generate_secret_key import and usage from TorchServe, MMS, TF Serving, and SMD prepare functions. Switch compute_hash calls from HMAC-SHA256 to plain SHA-256 (no secret_key parameter). This is required because generate_secret_key was removed from check_integrity.py in the previous commit. Without this change, all model server imports fail with ImportError.
1 parent fc535b4 commit ff838db

21 files changed

+60
-128
lines changed

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

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
from sagemaker.serve.spec.inference_spec import InferenceSpec
2727
from sagemaker.serve.detector.dependency_manager import capture_dependencies
2828
from sagemaker.serve.validations.check_integrity import (
29-
generate_secret_key,
3029
compute_hash,
3130
)
3231
from sagemaker.core.remote_function.core.serialization import _MetaData
@@ -119,11 +118,10 @@ def prepare_for_mms(
119118

120119
capture_dependencies(dependencies=dependencies, work_dir=code_dir)
121120

122-
secret_key = generate_secret_key()
123121
with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
124122
buffer = f.read()
125-
hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
123+
hash_value = compute_hash(buffer=buffer)
126124
with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
127125
metadata.write(_MetaData(hash_value).to_json())
128126

129-
return secret_key
127+
return ""

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,6 @@ def _start_serving(
3535
env = {
3636
"SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
3737
"SAGEMAKER_PROGRAM": "inference.py",
38-
"SAGEMAKER_SERVE_SECRET_KEY": secret_key,
3938
"LOCAL_PYTHON": platform.python_version(),
4039
}
4140
if env_vars:
@@ -131,7 +130,6 @@ def _upload_server_artifacts(
131130
env_vars = {
132131
"SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
133132
"SAGEMAKER_PROGRAM": "inference.py",
134-
"SAGEMAKER_SERVE_SECRET_KEY": secret_key,
135133
"SAGEMAKER_REGION": sagemaker_session.boto_region_name,
136134
"SAGEMAKER_CONTAINER_LOG_LEVEL": "10",
137135
"LOCAL_PYTHON": platform.python_version(),

‎sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
from sagemaker.serve.spec.inference_spec import InferenceSpec
1313
from sagemaker.serve.detector.dependency_manager import capture_dependencies
1414
from sagemaker.serve.validations.check_integrity import (
15-
generate_secret_key,
1615
compute_hash,
1716
)
1817
from sagemaker.core.remote_function.core.serialization import _MetaData
@@ -64,11 +63,10 @@ def prepare_for_smd(
6463

6564
capture_dependencies(dependencies=dependencies, work_dir=code_dir)
6665

67-
secret_key = generate_secret_key()
6866
with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
6967
buffer = f.read()
70-
hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
68+
hash_value = compute_hash(buffer=buffer)
7169
with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
7270
metadata.write(_MetaData(hash_value).to_json())
7371

74-
return secret_key
72+
return ""

‎sagemaker-serve/src/sagemaker/serve/model_server/smd/server.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,6 @@ def _upload_smd_artifacts(
5353
"SAGEMAKER_INFERENCE_CODE_DIRECTORY": "/opt/ml/model/code",
5454
"SAGEMAKER_INFERENCE_CODE": "inference.handler",
5555
"SAGEMAKER_REGION": sagemaker_session.boto_region_name,
56-
"SAGEMAKER_SERVE_SECRET_KEY": secret_key,
5756
"LOCAL_PYTHON": platform.python_version(),
5857
}
5958
return s3_upload_path, env_vars

‎sagemaker-serve/src/sagemaker/serve/model_server/tei/server.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,6 @@ def _start_tei_serving(
3838
secret_key: Secret key to use for authentication
3939
env_vars: Environment variables to set
4040
"""
41-
if env_vars and secret_key:
42-
env_vars["SAGEMAKER_SERVE_SECRET_KEY"] = secret_key
4341

4442
self.container = client.containers.run(
4543
image,

‎sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
)
1212
from sagemaker.serve.detector.dependency_manager import capture_dependencies
1313
from sagemaker.serve.validations.check_integrity import (
14-
generate_secret_key,
1514
compute_hash,
1615
)
1716
from sagemaker.core.remote_function.core.serialization import _MetaData
@@ -57,11 +56,10 @@ def prepare_for_tf_serving(
5756
raise ValueError("SavedModel is not found for Tensorflow or Keras flavor.")
5857
_move_contents(src_dir=mlflow_saved_model_dir, dest_dir=saved_model_bundle_dir)
5958

60-
secret_key = generate_secret_key()
6159
with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
6260
buffer = f.read()
63-
hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
61+
hash_value = compute_hash(buffer=buffer)
6462
with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
6563
metadata.write(_MetaData(hash_value).to_json())
6664

67-
return secret_key
65+
return ""

‎sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/server.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,6 @@ def _start_tensorflow_serving(
4747
environment={
4848
"SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
4949
"SAGEMAKER_PROGRAM": "inference.py",
50-
"SAGEMAKER_SERVE_SECRET_KEY": secret_key,
5150
"LOCAL_PYTHON": platform.python_version(),
5251
**env_vars,
5352
},
@@ -124,7 +123,6 @@ def _upload_tensorflow_serving_artifacts(
124123
"SAGEMAKER_PROGRAM": "inference.py",
125124
"SAGEMAKER_REGION": sagemaker_session.boto_region_name,
126125
"SAGEMAKER_CONTAINER_LOG_LEVEL": "10",
127-
"SAGEMAKER_SERVE_SECRET_KEY": secret_key,
128126
"LOCAL_PYTHON": platform.python_version(),
129127
}
130128
return s3_upload_path, env_vars

‎sagemaker-serve/src/sagemaker/serve/model_server/torchserve/prepare.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
from sagemaker.serve.spec.inference_spec import InferenceSpec
1414
from sagemaker.serve.detector.dependency_manager import capture_dependencies
1515
from sagemaker.serve.validations.check_integrity import (
16-
generate_secret_key,
1716
compute_hash,
1817
)
1918
from sagemaker.serve.validations.check_image_uri import is_1p_image_uri
@@ -67,11 +66,10 @@ def prepare_for_torchserve(
6766

6867
capture_dependencies(dependencies=dependencies, work_dir=code_dir)
6968

70-
secret_key = generate_secret_key()
7169
with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
7270
buffer = f.read()
73-
hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
71+
hash_value = compute_hash(buffer=buffer)
7472
with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
7573
metadata.write(_MetaData(hash_value).to_json())
7674

77-
return secret_key
75+
return ""

‎sagemaker-serve/src/sagemaker/serve/model_server/torchserve/server.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@ def _start_torch_serve(
3939
environment={
4040
"SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
4141
"SAGEMAKER_PROGRAM": "inference.py",
42-
"SAGEMAKER_SERVE_SECRET_KEY": secret_key,
4342
"LOCAL_PYTHON": platform.python_version(),
4443
**env_vars,
4544
},
@@ -103,7 +102,6 @@ def _upload_torchserve_artifacts(
103102
"SAGEMAKER_PROGRAM": "inference.py",
104103
"SAGEMAKER_REGION": sagemaker_session.boto_region_name,
105104
"SAGEMAKER_CONTAINER_LOG_LEVEL": "10",
106-
"SAGEMAKER_SERVE_SECRET_KEY": secret_key,
107105
"LOCAL_PYTHON": platform.python_version(),
108106
}
109107
return s3_upload_path, env_vars

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

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -68,10 +68,9 @@ def test_prepare_mms_js_resources(self, mock_create_dir, mock_copy_js):
6868

6969
@patch('builtins.input', return_value='')
7070
@patch('sagemaker.serve.model_server.multi_model_server.prepare.compute_hash')
71-
@patch('sagemaker.serve.model_server.multi_model_server.prepare.generate_secret_key')
7271
@patch('sagemaker.serve.model_server.multi_model_server.prepare.capture_dependencies')
7372
@patch('shutil.copy2')
74-
def test_prepare_for_mms_creates_structure(self, mock_copy, mock_capture, mock_gen_key, mock_hash, mock_input):
73+
def test_prepare_for_mms_creates_structure(self, mock_copy, mock_capture, mock_hash, mock_input):
7574
"""Test prepare_for_mms creates directory structure and files."""
7675
from sagemaker.serve.model_server.multi_model_server.prepare import prepare_for_mms
7776

@@ -83,7 +82,6 @@ def test_prepare_for_mms_creates_structure(self, mock_copy, mock_capture, mock_g
8382
serve_pkl = code_dir / "serve.pkl"
8483
serve_pkl.write_bytes(b"test data")
8584

86-
mock_gen_key.return_value = "test-secret-key"
8785
mock_hash.return_value = "test-hash"
8886
mock_session = Mock()
8987
mock_inference_spec = Mock()
@@ -98,16 +96,14 @@ def test_prepare_for_mms_creates_structure(self, mock_copy, mock_capture, mock_g
9896
inference_spec=mock_inference_spec
9997
)
10098

101-
self.assertEqual(secret_key, "test-secret-key")
10299
mock_inference_spec.prepare.assert_called_once_with(str(model_path))
103100
mock_capture.assert_called_once()
104101

105102
@patch('builtins.input', return_value='')
106103
@patch('sagemaker.serve.model_server.multi_model_server.prepare.compute_hash')
107-
@patch('sagemaker.serve.model_server.multi_model_server.prepare.generate_secret_key')
108104
@patch('sagemaker.serve.model_server.multi_model_server.prepare.capture_dependencies')
109105
@patch('shutil.copy2')
110-
def test_prepare_for_mms_raises_on_invalid_dir(self, mock_copy, mock_capture, mock_gen_key, mock_hash, mock_input):
106+
def test_prepare_for_mms_raises_on_invalid_dir(self, mock_copy, mock_capture, mock_hash, mock_input):
111107
"""Test prepare_for_mms raises exception for invalid directory."""
112108
from sagemaker.serve.model_server.multi_model_server.prepare import prepare_for_mms
113109

@@ -128,10 +124,9 @@ def test_prepare_for_mms_raises_on_invalid_dir(self, mock_copy, mock_capture, mo
128124

129125
@patch('builtins.input', return_value='')
130126
@patch('sagemaker.serve.model_server.multi_model_server.prepare.compute_hash')
131-
@patch('sagemaker.serve.model_server.multi_model_server.prepare.generate_secret_key')
132127
@patch('sagemaker.serve.model_server.multi_model_server.prepare.capture_dependencies')
133128
@patch('shutil.copy2')
134-
def test_prepare_for_mms_copies_shared_libs(self, mock_copy, mock_capture, mock_gen_key, mock_hash, mock_input):
129+
def test_prepare_for_mms_copies_shared_libs(self, mock_copy, mock_capture, mock_hash, mock_input):
135130
"""Test prepare_for_mms copies shared libraries."""
136131
from sagemaker.serve.model_server.multi_model_server.prepare import prepare_for_mms
137132

@@ -145,7 +140,6 @@ def test_prepare_for_mms_copies_shared_libs(self, mock_copy, mock_capture, mock_
145140
shared_lib = Path(self.temp_dir) / "lib.so"
146141
shared_lib.touch()
147142

148-
mock_gen_key.return_value = "test-key"
149143
mock_hash.return_value = "test-hash"
150144
mock_session = Mock()
151145

0 commit comments

Comments
 (0)