Address comments

Author: yangaws

src/sagemaker/amazon/amazon_estimator.py

@@ -23,7 +23,7 @@
 from sagemaker.amazon.common import write_numpy_to_dense_tensor
 from sagemaker.estimator import EstimatorBase, _TrainingJob
 from sagemaker.session import s3_input
-from sagemaker.utils import sagemaker_timestamp
+from sagemaker.utils import sagemaker_timestamp, get_ecr_image_uri_prefix
 
 logger = logging.getLogger(__name__)
 
@@ -283,79 +283,79 @@ def registry(region_name, algorithm=None):
 
     https://github.com/aws/sagemaker-python-sdk/tree/master/src/sagemaker/amazon
     """
-    if algorithm in [None, "pca", "kmeans", "linear-learner", "factorization-machines", "ntm",
-                     "randomcutforest", "knn", "object2vec", "ipinsights"]:
+    if algorithm in [None, 'pca', 'kmeans', 'linear-learner', 'factorization-machines', 'ntm',
+                     'randomcutforest', 'knn', 'object2vec', 'ipinsights']:
         account_id = {
-            "us-east-1": "382416733822",
-            "us-east-2": "404615174143",
-            "us-west-2": "174872318107",
-            "eu-west-1": "438346466558",
-            "eu-central-1": "664544806723",
-            "ap-northeast-1": "351501993468",
-            "ap-northeast-2": "835164637446",
-            "ap-southeast-2": "712309505854",
-            "us-gov-west-1": "226302683700",
-            "ap-southeast-1": "475088953585",
-            "ap-south-1": "991648021394",
-            "ca-central-1": "469771592824",
-            "eu-west-2": "644912444149",
-            "us-west-1": "632365934929",
-            "us-iso-east-1": "490574956308",
+            'us-east-1': '382416733822',
+            'us-east-2': '404615174143',
+            'us-west-2': '174872318107',
+            'eu-west-1': '438346466558',
+            'eu-central-1': '664544806723',
+            'ap-northeast-1': '351501993468',
+            'ap-northeast-2': '835164637446',
+            'ap-southeast-2': '712309505854',
+            'us-gov-west-1': '226302683700',
+            'ap-southeast-1': '475088953585',
+            'ap-south-1': '991648021394',
+            'ca-central-1': '469771592824',
+            'eu-west-2': '644912444149',
+            'us-west-1': '632365934929',
+            'us-iso-east-1': '490574956308',
         }[region_name]
-    elif algorithm in ["lda"]:
+    elif algorithm in ['lda']:
         account_id = {
-            "us-east-1": "766337827248",
-            "us-east-2": "999911452149",
-            "us-west-2": "266724342769",
-            "eu-west-1": "999678624901",
-            "eu-central-1": "353608530281",
-            "ap-northeast-1": "258307448986",
-            "ap-northeast-2": "293181348795",
-            "ap-southeast-2": "297031611018",
-            "us-gov-west-1": "226302683700",
-            "ap-southeast-1": "475088953585",
-            "ap-south-1": "991648021394",
-            "ca-central-1": "469771592824",
-            "eu-west-2": "644912444149",
-            "us-west-1": "632365934929",
-            "us-iso-east-1": "490574956308",
+            'us-east-1': '766337827248',
+            'us-east-2': '999911452149',
+            'us-west-2': '266724342769',
+            'eu-west-1': '999678624901',
+            'eu-central-1': '353608530281',
+            'ap-northeast-1': '258307448986',
+            'ap-northeast-2': '293181348795',
+            'ap-southeast-2': '297031611018',
+            'us-gov-west-1': '226302683700',
+            'ap-southeast-1': '475088953585',
+            'ap-south-1': '991648021394',
+            'ca-central-1': '469771592824',
+            'eu-west-2': '644912444149',
+            'us-west-1': '632365934929',
+            'us-iso-east-1': '490574956308',
         }[region_name]
-    elif algorithm in ["forecasting-deepar"]:
+    elif algorithm in ['forecasting-deepar']:
         account_id = {
-            "us-east-1": "522234722520",
-            "us-east-2": "566113047672",
-            "us-west-2": "156387875391",
-            "eu-west-1": "224300973850",
-            "eu-central-1": "495149712605",
-            "ap-northeast-1": "633353088612",
-            "ap-northeast-2": "204372634319",
-            "ap-southeast-2": "514117268639",
-            "us-gov-west-1": "226302683700",
-            "ap-southeast-1": "475088953585",
-            "ap-south-1": "991648021394",
-            "ca-central-1": "469771592824",
-            "eu-west-2": "644912444149",
-            "us-west-1": "632365934929",
-            "us-iso-east-1": "490574956308",
+            'us-east-1': '522234722520',
+            'us-east-2': '566113047672',
+            'us-west-2': '156387875391',
+            'eu-west-1': '224300973850',
+            'eu-central-1': '495149712605',
+            'ap-northeast-1': '633353088612',
+            'ap-northeast-2': '204372634319',
+            'ap-southeast-2': '514117268639',
+            'us-gov-west-1': '226302683700',
+            'ap-southeast-1': '475088953585',
+            'ap-south-1': '991648021394',
+            'ca-central-1': '469771592824',
+            'eu-west-2': '644912444149',
+            'us-west-1': '632365934929',
+            'us-iso-east-1': '490574956308',
         }[region_name]
-    elif algorithm in ["xgboost", "seq2seq", "image-classification", "blazingtext",
-                       "object-detection", "semantic-segmentation"]:
+    elif algorithm in ['xgboost', 'seq2seq', 'image-classification', 'blazingtext',
+                       'object-detection', 'semantic-segmentation']:
         account_id = {
-            "us-east-1": "811284229777",
-            "us-east-2": "825641698319",
-            "us-west-2": "433757028032",
-            "eu-west-1": "685385470294",
-            "eu-central-1": "813361260812",
-            "ap-northeast-1": "501404015308",
-            "ap-northeast-2": "306986355934",
-            "ap-southeast-2": "544295431143",
-            "us-gov-west-1": "226302683700",
-            "ap-southeast-1": "475088953585",
-            "ap-south-1": "991648021394",
-            "ca-central-1": "469771592824",
-            "eu-west-2": "644912444149",
-            "us-west-1": "632365934929",
-            "us-iso-east-1": "490574956308",
+            'us-east-1': '811284229777',
+            'us-east-2': '825641698319',
+            'us-west-2': '433757028032',
+            'eu-west-1': '685385470294',
+            'eu-central-1': '813361260812',
+            'ap-northeast-1': '501404015308',
+            'ap-northeast-2': '306986355934',
+            'ap-southeast-2': '544295431143',
+            'us-gov-west-1': '226302683700',
+            'ap-southeast-1': '475088953585',
+            'ap-south-1': '991648021394',
+            'ca-central-1': '469771592824',
+            'eu-west-2': '644912444149',
+            'us-west-1': '632365934929',
+            'us-iso-east-1': '490574956308',
         }[region_name]
     elif algorithm in ['image-classification-neo', 'xgboost-neo']:
         account_id = {
@@ -365,10 +365,9 @@ def registry(region_name, algorithm=None):
             'us-east-2': '007439368137',
         }[region_name]
     else:
-        raise ValueError("Algorithm class:{} doesn't have mapping to account_id with images".format(algorithm))
+        raise ValueError('Algorithm class:{} does not have mapping to account_id with images'.format(algorithm))
 
-    domain_name = "c2s.ic.gov" if region_name == "us-iso-east-1" else "amazonaws.com"
-    return "{}.dkr.ecr.{}.{}".format(account_id, region_name, domain_name)
+    return get_ecr_image_uri_prefix(account_id, region_name)
 
 
 def get_image_uri(region_name, repo_name, repo_version=1):

src/sagemaker/fw_registry.py

@@ -13,6 +13,8 @@
 from __future__ import absolute_import
 import logging
 
+from sagemaker.utils import get_ecr_image_uri_prefix
+
 image_registry_map = {
     "us-west-1": {
         "sparkml-serving": "746614075791",
@@ -84,8 +86,7 @@ def registry(region_name, framework=None):
     """
     try:
         account_id = image_registry_map[region_name][framework]
-        domain_name = "c2s.ic.gov" if region_name == "us-iso-east-1" else "amazonaws.com"
-        return "{}.dkr.ecr.{}.{}".format(account_id, region_name, domain_name)
+        return get_ecr_image_uri_prefix(account_id, region_name)
     except KeyError:
         logging.error("The specific image or region does not exist")
         raise

src/sagemaker/fw_utils.py

@@ -21,6 +21,8 @@
 import tempfile
 from six.moves.urllib.parse import urlparse
 
+from sagemaker.utils import get_ecr_image_uri_prefix
+
 _TAR_SOURCE_FILENAME = 'source.tar.gz'
 
 UploadedCode = namedtuple('UserCode', ['s3_prefix', 'script_name'])
@@ -97,10 +99,8 @@ def create_image_uri(region, framework, instance_type, framework_version, py_ver
                                              optimized_families=optimized_families):
         framework += '-eia'
 
-    domain_name = "c2s.ic.gov" if region == "us-iso-east-1" else "amazonaws.com"
-
-    return "{}.dkr.ecr.{}.{}/sagemaker-{}:{}" \
-        .format(account, region, domain_name, framework, tag)
+    return "{}/sagemaker-{}:{}" \
+        .format(get_ecr_image_uri_prefix(account, region), framework, tag)
 
 
 def _accelerator_type_valid_for_framework(framework, accelerator_type=None, optimized_families=None):

src/sagemaker/utils.py

@@ -292,6 +292,20 @@ def download_file(bucket_name, path, target, sagemaker_session):
     bucket.download_file(path, target)
 
 
+def get_ecr_image_uri_prefix(account, region):
+    """get prefix of ECR image URI
+
+    Args:
+        account (str): AWS account number
+        region (str): AWS region name
+
+    Returns:
+        (str): URI prefix of ECR image
+    """
+    domain = 'c2s.ic.gov' if region == 'us-iso-east-1' else 'amazonaws.com'
+    return '{}.dkr.ecr.{}.{}'.format(account, region, domain)
+
+
 class DeferredError(object):
     """Stores an exception and raises it at a later time if this
     object is accessed in any way.  Useful to allow soft-dependencies on imports,