@@ -147,7 +147,7 @@ def test_model_server_override_djl_without_model_or_mlflow(self, mock_serve_sett
147
147
)
148
148
self .assertRaisesRegex (
149
149
Exception ,
150
- "Missing required parameter `model` or 'ml_flow' path" ,
150
+ "Missing required parameter `model` or 'ml_flow' path or inf_spec " ,
151
151
builder .build ,
152
152
Mode .SAGEMAKER_ENDPOINT ,
153
153
mock_role_arn ,
@@ -168,12 +168,26 @@ def test_model_server_override_torchserve_with_model(
168
168
169
169
mock_build_for_ts .assert_called_once ()
170
170
171
+ @patch ("sagemaker.serve.builder.model_builder._ServeSettings" )
172
+ @patch ("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_torchserve" )
173
+ def test_model_server_override_torchserve_with_inf_spec (
174
+ self , mock_build_for_ts , mock_serve_settings
175
+ ):
176
+ mock_setting_object = mock_serve_settings .return_value
177
+ mock_setting_object .role_arn = mock_role_arn
178
+ mock_setting_object .s3_model_data_url = mock_s3_model_data_url
179
+
180
+ builder = ModelBuilder (model_server = ModelServer .TORCHSERVE , inference_spec = "some value" )
181
+ builder .build (sagemaker_session = mock_session )
182
+
183
+ mock_build_for_ts .assert_called_once ()
184
+
171
185
@patch ("sagemaker.serve.builder.model_builder._ServeSettings" )
172
186
def test_model_server_override_torchserve_without_model_or_mlflow (self , mock_serve_settings ):
173
187
builder = ModelBuilder (model_server = ModelServer .TORCHSERVE )
174
188
self .assertRaisesRegex (
175
189
Exception ,
176
- "Missing required parameter `model` or 'ml_flow' path" ,
190
+ "Missing required parameter `model` or 'ml_flow' path or inf_spec " ,
177
191
builder .build ,
178
192
Mode .SAGEMAKER_ENDPOINT ,
179
193
mock_role_arn ,
0 commit comments