@@ -400,3 +400,19 @@ def test_compile_model_for_cloud(sagemaker_session, tmpdir):
400
400
model .compile (target_instance_family = 'ml_c4' , input_shape = {'data' : [1 , 3 , 1024 , 1024 ]},
401
401
output_path = 's3://output' , role = 'role' , framework = 'tensorflow' , job_name = "compile-model" )
402
402
assert model ._is_compiled_model is True
403
+
404
+
405
+ def test_check_neo_region (sagemaker_session , tmpdir ):
406
+ sagemaker_session .wait_for_compilation_job = Mock (
407
+ return_value = DESCRIBE_COMPILATION_JOB_RESPONSE )
408
+ model = DummyFrameworkModel (sagemaker_session , source_dir = str (tmpdir ))
409
+ ec2_region_list = ['us-east-2' , 'us-east-1' , 'us-west-1' , 'us-west-2' , 'ap-east-1' , 'ap-south-1' ,
410
+ 'ap-northeast-3' , 'ap-northeast-2' , 'ap-southeast-1' , 'ap-southeast-2' , 'ap-northeast-1' ,
411
+ 'ca-central-1' , 'cn-north-1' , 'cn-northwest-1' , 'eu-central-1' , ' eu-west-1' , 'eu-west-2' ,
412
+ 'eu-west-3' , 'eu-north-1' , 'sa-east-1' , 'us-gov-east-1' , 'us-gov-west-1' ]
413
+ neo_support_region = ['us-west-2' , 'eu-west-1' , 'us-east-1' , 'us-east-2' ]
414
+ for region_name in ec2_region_list :
415
+ if region_name in neo_support_region :
416
+ assert model .check_neo_region (region_name ) is True
417
+ else :
418
+ assert model .check_neo_region (region_name ) is False
0 commit comments