change: Allowing account overrides in special regions

Author: mnuyens

src/sagemaker/fw_utils.py

@@ -58,6 +58,7 @@
 VALID_ACCOUNTS_BY_REGION = {"us-gov-west-1": "246785580436", "us-iso-east-1": "744548109606"}
 OPT_IN_ACCOUNTS_BY_REGION = {"ap-east-1": "057415533634"}
 ASIMOV_OPT_IN_ACCOUNTS_BY_REGION = {"ap-east-1": "871362719292"}
+DEFAULT_ACCOUNT = "520713654638"
 
 MERGED_FRAMEWORKS_REPO_MAP = {
     "tensorflow-scriptmode": "tensorflow-training",
@@ -183,7 +184,7 @@ def create_image_uri(
     instance_type,
     framework_version,
     py_version=None,
-    account="520713654638",
+    account=None,
     accelerator_type=None,
     optimized_families=None,
 ):
@@ -218,13 +219,14 @@ def create_image_uri(
         framework += "-eia"
 
     # Handle Account Number for Gov Cloud and frameworks with DLC merged images
-    account = _registry_id(
-        region=region,
-        framework=framework,
-        py_version=py_version,
-        account=account,
-        framework_version=framework_version,
-    )
+    if account is None:
+        account = _registry_id(
+            region=region,
+            framework=framework,
+            py_version=py_version,
+            account=DEFAULT_ACCOUNT,
+            framework_version=framework_version,
+        )
 
     # Handle Local Mode
     if instance_type.startswith("local"):

tests/unit/test_fw_utils.py

@@ -68,14 +68,14 @@ def test_create_image_uri_cpu():
     assert image_uri == "23.dkr.ecr.mars-south-3.amazonaws.com/sagemaker-mlfw:1.0rc-cpu-py2"
 
     image_uri = fw_utils.create_image_uri(
-        "us-gov-west-1", MOCK_FRAMEWORK, "ml.c4.large", "1.0rc", "py2", "23"
+        "us-gov-west-1", MOCK_FRAMEWORK, "ml.c4.large", "1.0rc", "py2"
     )
     assert (
         image_uri == "246785580436.dkr.ecr.us-gov-west-1.amazonaws.com/sagemaker-mlfw:1.0rc-cpu-py2"
     )
 
     image_uri = fw_utils.create_image_uri(
-        "us-iso-east-1", MOCK_FRAMEWORK, "ml.c4.large", "1.0rc", "py2", "23"
+        "us-iso-east-1", MOCK_FRAMEWORK, "ml.c4.large", "1.0rc", "py2"
     )
     assert image_uri == "744548109606.dkr.ecr.us-iso-east-1.c2s.ic.gov/sagemaker-mlfw:1.0rc-cpu-py2"
 
@@ -188,6 +188,25 @@ def test_mxnet_eia_images():
         == "763104351884.dkr.ecr.us-east-1.amazonaws.com/mxnet-inference-eia:1.4.1-cpu-py3"
     )
 
+def test_create_image_uri_override_account():
+    image_uri = fw_utils.create_image_uri(
+        "us-west-1", MOCK_FRAMEWORK, "ml.p3.2xlarge", "1.0rc", "py3", account="fake"
+    )
+    assert image_uri == "fake.dkr.ecr.us-west-1.amazonaws.com/sagemaker-mlfw:1.0rc-gpu-py3"
+
+
+def test_create_image_uri_gov_cloud_override_account():
+    image_uri = fw_utils.create_image_uri(
+        "us-gov-west-1", MOCK_FRAMEWORK, "ml.p3.2xlarge", "1.0rc", "py3", account="fake"
+    )
+    assert image_uri == "fake.dkr.ecr.us-gov-west-1.amazonaws.com/sagemaker-mlfw:1.0rc-gpu-py3"
+
+
+def test_create_image_uri_hkg_override_account():
+    image_uri = fw_utils.create_image_uri(
+        MOCK_HKG_REGION, MOCK_FRAMEWORK, "ml.p3.2xlarge", "1.0rc", "py3", account="fake"
+    )
+    assert {image_uri == "fake.dkr.ecr.ap-east-1.amazonaws.com/sagemaker-mlfw:1.0rc-gpu-py3"}
 
 def test_create_image_uri_merged():
     image_uri = fw_utils.create_image_uri(