Add gov cloud account number for framework (aws#400)

Author: yangaws

src/sagemaker/fw_utils.py

@@ -49,6 +49,10 @@ def create_image_uri(region, framework, instance_type, framework_version, py_ver
         str: The appropriate image URI based on the given parameters.
     """
 
+    # Handle Account Number for Gov Cloud
+    if region == 'us-gov-west-1':
+        account = '246785580436'
+
     # Handle Local Mode
     if instance_type.startswith('local'):
         device_type = 'cpu' if instance_type == 'local' else 'gpu'

tests/unit/test_fw_utils.py

@@ -62,6 +62,11 @@ def test_create_image_uri_default_account():
     assert image_uri == '520713654638.dkr.ecr.mars-south-3.amazonaws.com/sagemaker-mlfw:1.0rc-gpu-py3'
 
 
+def test_create_image_uri_gov_cloud():
+    image_uri = create_image_uri('us-gov-west-1', 'mlfw', 'ml.p3.2xlarge', '1.0rc', 'py3')
+    assert image_uri == '246785580436.dkr.ecr.us-gov-west-1.amazonaws.com/sagemaker-mlfw:1.0rc-gpu-py3'
+
+
 def test_invalid_instance_type():
     # instance type is missing 'ml.' prefix
     with pytest.raises(ValueError):