From 0a90fe0db8f2a1d95db0e856ed9bf06e0bfac0e2 Mon Sep 17 00:00:00 2001 From: Nuyens Date: Thu, 23 May 2019 14:48:36 -0700 Subject: [PATCH] change: Updating compilation to be current w/ Sagemaker Neo --- src/sagemaker/model.py | 9 ++++++--- tests/unit/test_model.py | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py index a0412e2ff0..a5c7a96ed3 100644 --- a/src/sagemaker/model.py +++ b/src/sagemaker/model.py @@ -23,15 +23,18 @@ LOGGER = logging.getLogger('sagemaker') -NEO_ALLOWED_TARGET_INSTANCE_FAMILY = set(['ml_c5', 'ml_m5', 'ml_c4', 'ml_m4', 'jetson_tx1', 'jetson_tx2', 'ml_p2', - 'ml_p3', 'deeplens', 'rasp3b']) +NEO_ALLOWED_TARGET_INSTANCE_FAMILY = set(['ml_c5', 'ml_m5', 'ml_c4', 'ml_m4', + 'jetson_tx1', 'jetson_tx2', 'jetson_nano', 'ml_p2', + 'ml_p3', 'deeplens', 'rasp3b', + 'rk3288', 'rk3399', 'sbe_c']) NEO_ALLOWED_FRAMEWORKS = set(['mxnet', 'tensorflow', 'pytorch', 'onnx', 'xgboost']) NEO_IMAGE_ACCOUNT = { 'us-west-2': '301217895009', 'us-east-1': '785573368785', 'eu-west-1': '802834080501', - 'us-east-2': '007439368137' + 'us-east-2': '007439368137', + 'ap-northeast-1': '941853720454' } diff --git a/tests/unit/test_model.py b/tests/unit/test_model.py index 3d5ee4e688..6811e6f3db 100644 --- a/tests/unit/test_model.py +++ b/tests/unit/test_model.py @@ -415,7 +415,7 @@ def test_check_neo_region(sagemaker_session, tmpdir): 'ap-northeast-3', 'ap-northeast-2', 'ap-southeast-1', 'ap-southeast-2', 'ap-northeast-1', 'ca-central-1', 'cn-north-1', 'cn-northwest-1', 'eu-central-1', ' eu-west-1', 'eu-west-2', 'eu-west-3', 'eu-north-1', 'sa-east-1', 'us-gov-east-1', 'us-gov-west-1'] - neo_support_region = ['us-west-2', 'eu-west-1', 'us-east-1', 'us-east-2'] + neo_support_region = ['us-west-2', 'eu-west-1', 'us-east-1', 'us-east-2', 'ap-northeast-1'] for region_name in ec2_region_list: if region_name in neo_support_region: assert model.check_neo_region(region_name) is True