Skip to content

Commit e3412cf

Browse files
aviruthenaviruthen
committed
Fix unit and integ tests
1 parent 74d7905 commit e3412cf

13 files changed

Lines changed: 123 additions & 188 deletions

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

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,6 @@ def prepare(
7777
upload_artifacts = self._upload_torchserve_artifacts(
7878
model_path=model_path,
7979
sagemaker_session=sagemaker_session,
80-
secret_key=secret_key,
8180
s3_model_data_url=s3_model_data_url,
8281
image=image,
8382
should_upload_artifacts=True,
@@ -87,7 +86,6 @@ def prepare(
8786
upload_artifacts = self._upload_triton_artifacts(
8887
model_path=model_path,
8988
sagemaker_session=sagemaker_session,
90-
secret_key=secret_key,
9189
s3_model_data_url=s3_model_data_url,
9290
image=image,
9391
should_upload_artifacts=True,
@@ -106,7 +104,6 @@ def prepare(
106104
upload_artifacts = self._upload_tensorflow_serving_artifacts(
107105
model_path=model_path,
108106
sagemaker_session=sagemaker_session,
109-
secret_key=secret_key,
110107
s3_model_data_url=s3_model_data_url,
111108
image=image,
112109
should_upload_artifacts=True,
@@ -132,7 +129,6 @@ def prepare(
132129
model_path=model_path,
133130
sagemaker_session=sagemaker_session,
134131
s3_model_data_url=s3_model_data_url,
135-
secret_key=secret_key,
136132
image=image,
137133
should_upload_artifacts=should_upload_artifacts,
138134
)
@@ -150,7 +146,6 @@ def prepare(
150146
upload_artifacts = self._upload_smd_artifacts(
151147
model_path=model_path,
152148
sagemaker_session=sagemaker_session,
153-
secret_key=secret_key,
154149
s3_model_data_url=s3_model_data_url,
155150
image=image,
156151
should_upload_artifacts=True,

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

Lines changed: 13 additions & 18 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,31 +82,29 @@ 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()
9088

91-
with patch('builtins.open', mock_open(read_data=b"test data")):
92-
secret_key = prepare_for_mms(
93-
model_path=str(model_path),
94-
shared_libs=[],
95-
dependencies={},
96-
session=mock_session,
97-
image_uri="test-image",
98-
inference_spec=mock_inference_spec
99-
)
89+
secret_key = prepare_for_mms(
90+
model_path=str(model_path),
91+
shared_libs=[],
92+
dependencies={},
93+
session=mock_session,
94+
image_uri="test-image",
95+
inference_spec=mock_inference_spec
96+
)
10097

101-
self.assertEqual(secret_key, "test-secret-key")
98+
# Should return empty string now (not secret key)
99+
self.assertEqual(secret_key, "")
102100
mock_inference_spec.prepare.assert_called_once_with(str(model_path))
103101
mock_capture.assert_called_once()
104102

105103
@patch('builtins.input', return_value='')
106104
@patch('sagemaker.serve.model_server.multi_model_server.prepare.compute_hash')
107-
@patch('sagemaker.serve.model_server.multi_model_server.prepare.generate_secret_key')
108105
@patch('sagemaker.serve.model_server.multi_model_server.prepare.capture_dependencies')
109106
@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):
107+
def test_prepare_for_mms_raises_on_invalid_dir(self, mock_copy, mock_capture, mock_hash, mock_input):
111108
"""Test prepare_for_mms raises exception for invalid directory."""
112109
from sagemaker.serve.model_server.multi_model_server.prepare import prepare_for_mms
113110

@@ -128,10 +125,9 @@ def test_prepare_for_mms_raises_on_invalid_dir(self, mock_copy, mock_capture, mo
128125

129126
@patch('builtins.input', return_value='')
130127
@patch('sagemaker.serve.model_server.multi_model_server.prepare.compute_hash')
131-
@patch('sagemaker.serve.model_server.multi_model_server.prepare.generate_secret_key')
132128
@patch('sagemaker.serve.model_server.multi_model_server.prepare.capture_dependencies')
133129
@patch('shutil.copy2')
134-
def test_prepare_for_mms_copies_shared_libs(self, mock_copy, mock_capture, mock_gen_key, mock_hash, mock_input):
130+
def test_prepare_for_mms_copies_shared_libs(self, mock_copy, mock_capture, mock_hash, mock_input):
135131
"""Test prepare_for_mms copies shared libraries."""
136132
from sagemaker.serve.model_server.multi_model_server.prepare import prepare_for_mms
137133

@@ -145,7 +141,6 @@ def test_prepare_for_mms_copies_shared_libs(self, mock_copy, mock_capture, mock_
145141
shared_lib = Path(self.temp_dir) / "lib.so"
146142
shared_lib.touch()
147143

148-
mock_gen_key.return_value = "test-key"
149144
mock_hash.return_value = "test-hash"
150145
mock_session = Mock()
151146

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

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,13 @@ def test_start_serving_creates_container(self, mock_path):
2525
client=mock_client,
2626
image="test-image:latest",
2727
model_path="/path/to/model",
28-
secret_key="test-secret",
2928
env_vars={"CUSTOM_VAR": "value"}
3029
)
3130

3231
self.assertEqual(server.container, mock_container)
3332
mock_client.containers.run.assert_called_once()
3433
call_kwargs = mock_client.containers.run.call_args[1]
35-
self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", call_kwargs["environment"])
36-
self.assertEqual(call_kwargs["environment"]["SAGEMAKER_SERVE_SECRET_KEY"], "test-secret")
34+
self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", call_kwargs["environment"])
3735

3836
@patch('sagemaker.serve.model_server.multi_model_server.server.Path')
3937
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):
5250
client=mock_client,
5351
image="test-image:latest",
5452
model_path="/path/to/model",
55-
secret_key="test-secret",
5653
env_vars=None
5754
)
5855

@@ -121,7 +118,6 @@ def test_upload_server_artifacts_with_s3_path(self, mock_is_s3, mock_fw_utils, m
121118

122119
model_data, env_vars = server._upload_server_artifacts(
123120
model_path="s3://bucket/model",
124-
secret_key="test-key",
125121
sagemaker_session=mock_session,
126122
should_upload_artifacts=False
127123
)
@@ -158,16 +154,14 @@ def test_upload_server_artifacts_uploads_to_s3(self, mock_path, mock_is_s3, mock
158154

159155
model_data, env_vars = server._upload_server_artifacts(
160156
model_path="/local/model",
161-
secret_key="test-key",
162157
sagemaker_session=mock_session,
163158
s3_model_data_url="s3://bucket/prefix",
164159
image="test-image",
165160
should_upload_artifacts=True
166161
)
167162

168163
self.assertIsNotNone(model_data)
169-
self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars)
170-
self.assertEqual(env_vars["SAGEMAKER_SERVE_SECRET_KEY"], "test-key")
164+
self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars)
171165

172166
@patch('sagemaker.serve.model_server.multi_model_server.server._is_s3_uri')
173167
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):
181175

182176
model_data, env_vars = server._upload_server_artifacts(
183177
model_path="/local/model",
184-
secret_key="test-key",
185178
sagemaker_session=mock_session,
186179
should_upload_artifacts=False
187180
)
188181

189182
self.assertIsNone(model_data)
190-
self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars)
183+
self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars)
191184

192185

193186
class TestUpdateEnvVars(unittest.TestCase):

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

Lines changed: 17 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,9 @@ def tearDown(self):
1818
shutil.rmtree(self.temp_dir)
1919

2020
@patch('sagemaker.serve.model_server.smd.prepare.compute_hash')
21-
@patch('sagemaker.serve.model_server.smd.prepare.generate_secret_key')
2221
@patch('sagemaker.serve.model_server.smd.prepare.capture_dependencies')
2322
@patch('shutil.copy2')
24-
def test_prepare_for_smd_with_inference_spec(self, mock_copy, mock_capture, mock_gen_key, mock_hash):
23+
def test_prepare_for_smd_with_inference_spec(self, mock_copy, mock_capture, mock_hash):
2524
"""Test prepare_for_smd with InferenceSpec."""
2625
from sagemaker.serve.model_server.smd.prepare import prepare_for_smd
2726
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
3332
serve_pkl = code_dir / "serve.pkl"
3433
serve_pkl.write_bytes(b"test data")
3534

36-
mock_gen_key.return_value = "test-secret-key"
3735
mock_hash.return_value = "test-hash"
3836
mock_inference_spec = Mock(spec=InferenceSpec)
3937

40-
with patch('builtins.open', mock_open(read_data=b"test data")):
41-
secret_key = prepare_for_smd(
42-
model_path=str(model_path),
43-
shared_libs=[],
44-
dependencies={},
45-
inference_spec=mock_inference_spec
46-
)
38+
secret_key = prepare_for_smd(
39+
model_path=str(model_path),
40+
shared_libs=[],
41+
dependencies={},
42+
inference_spec=mock_inference_spec
43+
)
4744

48-
self.assertEqual(secret_key, "test-secret-key")
45+
self.assertEqual(secret_key, "")
4946
mock_inference_spec.prepare.assert_called_once_with(str(model_path))
5047

5148
@patch('os.rename')
5249
@patch('sagemaker.serve.model_server.smd.prepare.compute_hash')
53-
@patch('sagemaker.serve.model_server.smd.prepare.generate_secret_key')
5450
@patch('sagemaker.serve.model_server.smd.prepare.capture_dependencies')
5551
@patch('shutil.copy2')
56-
def test_prepare_for_smd_with_custom_orchestrator(self, mock_copy, mock_capture, mock_gen_key, mock_hash, mock_rename):
52+
def test_prepare_for_smd_with_custom_orchestrator(self, mock_copy, mock_capture, mock_hash, mock_rename):
5753
"""Test prepare_for_smd with CustomOrchestrator."""
5854
from sagemaker.serve.model_server.smd.prepare import prepare_for_smd
5955
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,
6561
serve_pkl = code_dir / "serve.pkl"
6662
serve_pkl.write_bytes(b"test data")
6763

68-
mock_gen_key.return_value = "test-secret-key"
6964
mock_hash.return_value = "test-hash"
7065
mock_orchestrator = Mock(spec=CustomOrchestrator)
7166

72-
with patch('builtins.open', mock_open(read_data=b"test data")):
73-
secret_key = prepare_for_smd(
74-
model_path=str(model_path),
75-
shared_libs=[],
76-
dependencies={},
77-
inference_spec=mock_orchestrator
78-
)
67+
secret_key = prepare_for_smd(
68+
model_path=str(model_path),
69+
shared_libs=[],
70+
dependencies={},
71+
inference_spec=mock_orchestrator
72+
)
7973

80-
self.assertEqual(secret_key, "test-secret-key")
74+
self.assertEqual(secret_key, "")
8175
# Verify custom_execution_inference.py was copied and renamed
8276
mock_rename.assert_called_once()
8377

8478
@patch('sagemaker.serve.model_server.smd.prepare.compute_hash')
85-
@patch('sagemaker.serve.model_server.smd.prepare.generate_secret_key')
8679
@patch('sagemaker.serve.model_server.smd.prepare.capture_dependencies')
8780
@patch('shutil.copy2')
88-
def test_prepare_for_smd_with_shared_libs(self, mock_copy, mock_capture, mock_gen_key, mock_hash):
81+
def test_prepare_for_smd_with_shared_libs(self, mock_copy, mock_capture, mock_hash):
8982
"""Test prepare_for_smd copies shared libraries."""
9083
from sagemaker.serve.model_server.smd.prepare import prepare_for_smd
9184

@@ -99,7 +92,6 @@ def test_prepare_for_smd_with_shared_libs(self, mock_copy, mock_capture, mock_ge
9992
shared_lib = Path(self.temp_dir) / "lib.so"
10093
shared_lib.touch()
10194

102-
mock_gen_key.return_value = "test-key"
10395
mock_hash.return_value = "test-hash"
10496

10597
with patch('builtins.open', mock_open(read_data=b"test data")):

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

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,11 @@ def test_upload_smd_artifacts_with_s3_path(self, mock_is_s3):
2020
s3_path, env_vars = server._upload_smd_artifacts(
2121
model_path="s3://bucket/model",
2222
sagemaker_session=mock_session,
23-
secret_key="test-key",
2423
should_upload_artifacts=False
2524
)
2625

2726
self.assertEqual(s3_path, "s3://bucket/model")
28-
self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars)
29-
self.assertEqual(env_vars["SAGEMAKER_SERVE_SECRET_KEY"], "test-key")
27+
self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars)
3028
self.assertIn("SAGEMAKER_INFERENCE_CODE_DIRECTORY", env_vars)
3129

3230
@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,
5149
s3_path, env_vars = server._upload_smd_artifacts(
5250
model_path="/local/model",
5351
sagemaker_session=mock_session,
54-
secret_key="test-key",
5552
s3_model_data_url="s3://bucket/prefix",
5653
image="test-image",
5754
should_upload_artifacts=True
5855
)
5956

6057
self.assertEqual(s3_path, "s3://bucket/code_prefix/model.tar.gz")
61-
self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars)
58+
self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars)
6259
self.assertIn("SAGEMAKER_INFERENCE_CODE", env_vars)
6360
mock_upload.assert_called_once()
6461

@@ -75,10 +72,12 @@ def test_upload_smd_artifacts_no_upload(self, mock_is_s3):
7572
s3_path, env_vars = server._upload_smd_artifacts(
7673
model_path="/local/model",
7774
sagemaker_session=mock_session,
78-
secret_key="test-key",
7975
should_upload_artifacts=False
8076
)
8177

78+
self.assertIsNone(s3_path)
79+
self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars)
80+
8281
self.assertIsNone(s3_path)
8382
self.assertIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars)
8483

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@ def test_start_tei_serving(self, mock_device_req, mock_path, mock_update_env):
2929
client=mock_client,
3030
image="tei:latest",
3131
model_path="/path/to/model",
32-
secret_key="test-secret",
3332
env_vars={"CUSTOM_VAR": "value"}
3433
)
3534

@@ -40,7 +39,7 @@ def test_start_tei_serving(self, mock_device_req, mock_path, mock_update_env):
4039
@patch('sagemaker.serve.model_server.tei.server.Path')
4140
@patch('sagemaker.serve.model_server.tei.server.DeviceRequest')
4241
def test_start_tei_serving_adds_secret_key(self, mock_device_req, mock_path, mock_update_env):
43-
"""Test _start_tei_serving adds secret key to env vars."""
42+
"""Test _start_tei_serving does not add secret key to env vars."""
4443
from sagemaker.serve.model_server.tei.server import LocalTeiServing
4544

4645
server = LocalTeiServing()
@@ -58,12 +57,11 @@ def test_start_tei_serving_adds_secret_key(self, mock_device_req, mock_path, moc
5857
client=mock_client,
5958
image="tei:latest",
6059
model_path="/path/to/model",
61-
secret_key="test-secret",
6260
env_vars=env_vars
6361
)
6462

65-
# Verify secret key was added to env_vars
66-
self.assertEqual(env_vars["SAGEMAKER_SERVE_SECRET_KEY"], "test-secret")
63+
# Verify secret key was NOT added to env_vars
64+
self.assertNotIn("SAGEMAKER_SERVE_SECRET_KEY", env_vars)
6765

6866
@patch('sagemaker.serve.model_server.tei.server.requests.post')
6967
@patch('sagemaker.serve.model_server.tei.server.get_docker_host')

0 commit comments

Comments
 (0)