@@ -100,18 +100,17 @@ def test_js_model_with_optimize_speculative_decoding_config_gated_requests_are_e
100100def test_js_model_with_optimize_sharding_and_resource_requirements_requests_are_expected (
101101 sagemaker_session ,
102102):
103- with (
104- patch .object (
105- Session ,
106- "wait_for_optimization_job" ,
107- return_value = {"OptimizationJobName" : "mock_optimization_job" },
108- ),
109- patch .object (Session , "create_model" , return_value = "mock_model" ) as mock_create_model ,
110- patch .object (
111- Session , "endpoint_from_production_variants" , return_value = "mock_endpoint_name"
112- ) as mock_endpoint_from_production_variants ,
113- patch .object (Session , "create_inference_component" ) as mock_create_inference_component ,
114- ):
103+ with patch .object (
104+ Session ,
105+ "wait_for_optimization_job" ,
106+ return_value = {"OptimizationJobName" : "mock_optimization_job" },
107+ ), patch .object (
108+ Session , "create_model" , return_value = "mock_model"
109+ ) as mock_create_model , patch .object (
110+ Session , "endpoint_from_production_variants" , return_value = "mock_endpoint_name"
111+ ) as mock_endpoint_from_production_variants , patch .object (
112+ Session , "create_inference_component"
113+ ) as mock_create_inference_component :
115114 iam_client = sagemaker_session .boto_session .client ("iam" )
116115 role_arn = iam_client .get_role (RoleName = ROLE_NAME )["Role" ]["Arn" ]
117116
@@ -186,17 +185,15 @@ def test_js_model_with_optimize_sharding_and_resource_requirements_requests_are_
186185def test_js_model_with_optimize_quantization_on_pre_optimized_model_requests_are_expected (
187186 sagemaker_session ,
188187):
189- with (
190- patch .object (
191- Session ,
192- "wait_for_optimization_job" ,
193- return_value = {"OptimizationJobName" : "mock_optimization_job" },
194- ),
195- patch .object (Session , "create_model" , return_value = "mock_model" ) as mock_create_model ,
196- patch .object (
197- Session , "endpoint_from_production_variants" , return_value = "mock_endpoint_name"
198- ) as mock_endpoint_from_production_variants ,
199- ):
188+ with patch .object (
189+ Session ,
190+ "wait_for_optimization_job" ,
191+ return_value = {"OptimizationJobName" : "mock_optimization_job" },
192+ ), patch .object (
193+ Session , "create_model" , return_value = "mock_model"
194+ ) as mock_create_model , patch .object (
195+ Session , "endpoint_from_production_variants" , return_value = "mock_endpoint_name"
196+ ) as mock_endpoint_from_production_variants :
200197 iam_client = sagemaker_session .boto_session .client ("iam" )
201198 role_arn = iam_client .get_role (RoleName = ROLE_NAME )["Role" ]["Arn" ]
202199
0 commit comments