Merge branch 'master' into change/dw_image_uri

Author: chenliu0831

src/sagemaker/image_uri_config/xgboost.json

@@ -154,6 +154,35 @@
                 "us-west-2": "246618743249"
             },
             "repository": "sagemaker-xgboost"
+        },
+        "1.2-2": {
+            "registries": {
+                "af-south-1": "510948584623",
+                "ap-east-1": "651117190479",
+                "ap-northeast-1": "354813040037",
+                "ap-northeast-2": "366743142698",
+                "ap-south-1": "720646828776",
+                "ap-southeast-1": "121021644041",
+                "ap-southeast-2": "783357654285",
+                "ca-central-1": "341280168497",
+                "cn-north-1": "450853457545",
+                "cn-northwest-1": "451049120500",
+                "eu-central-1": "492215442770",
+                "eu-north-1": "662702820516",
+                "eu-west-1": "141502667606",
+                "eu-west-2": "764974769150",
+                "eu-west-3": "659782779980",
+                "eu-south-1": "978288397137",
+                "me-south-1": "801668240914",
+                "sa-east-1": "737474898029",
+                "us-east-1": "683313688378",
+                "us-east-2": "257758044811",
+                "us-gov-west-1": "414596584902",
+                "us-iso-east-1": "833128469047",
+                "us-west-1": "746614075791",
+                "us-west-2": "246618743249"
+            },
+            "repository": "sagemaker-xgboost"
         }
     }
 }

src/sagemaker/session.py

@@ -3490,12 +3490,20 @@ def get_caller_identity_arn(self):
         """
         if os.path.exists(NOTEBOOK_METADATA_FILE):
             with open(NOTEBOOK_METADATA_FILE, "rb") as f:
-                instance_name = json.loads(f.read())["ResourceName"]
+                metadata = json.loads(f.read())
+                instance_name = metadata["ResourceName"]
+                domain_id = metadata.get("DomainId")
+                user_profile_name = metadata.get("UserProfileName")
             try:
-                instance_desc = self.sagemaker_client.describe_notebook_instance(
-                    NotebookInstanceName=instance_name
+                if domain_id is None:
+                    instance_desc = self.sagemaker_client.describe_notebook_instance(
+                        NotebookInstanceName=instance_name
+                    )
+                    return instance_desc["RoleArn"]
+                user_profile_desc = self.sagemaker_client.describe_user_profile(
+                    DomainId=domain_id, UserProfileName=user_profile_name
                 )
-                return instance_desc["RoleArn"]
+                return user_profile_desc["UserSettings"]["ExecutionRole"]
             except ClientError:
                 LOGGER.debug(
                     "Couldn't call 'describe_notebook_instance' to get the Role "

tests/unit/sagemaker/image_uris/test_xgboost.py

@@ -45,7 +45,7 @@
 }
 ALGO_VERSIONS = ("1", "latest")
 XGBOOST_FRAMEWORK_CPU_ONLY_VERSIONS = ("0.90-2", "0.90-1", "1.0-1")
-XGBOOST_FRAMEWORK_CPU_GPU_VERSIONS = ("1.2-1",)
+XGBOOST_FRAMEWORK_CPU_GPU_VERSIONS = ("1.2-1", "1.2-2")
 
 FRAMEWORK_REGISTRIES = {
     "af-south-1": "510948584623",

tests/unit/test_session.py

@@ -315,6 +315,31 @@ def test_get_caller_identity_arn_from_describe_notebook_instance(boto_session):
     )
 
 
+@patch(
+    "six.moves.builtins.open",
+    mock_open(
+        read_data='{"ResourceName": "SageMakerInstance", '
+        '"DomainId": "d-kbnw5yk6tg8j", '
+        '"UserProfileName": "default-1617915559064"}'
+    ),
+)
+@patch("os.path.exists", side_effect=mock_exists(NOTEBOOK_METADATA_FILE, True))
+def test_get_caller_identity_arn_from_describe_user_profile(boto_session):
+    sess = Session(boto_session)
+    expected_role = "arn:aws:iam::369233609183:role/service-role/SageMakerRole-20171129T072388"
+    sess.sagemaker_client.describe_user_profile.return_value = {
+        "UserSettings": {"ExecutionRole": expected_role}
+    }
+
+    actual = sess.get_caller_identity_arn()
+
+    assert actual == expected_role
+    sess.sagemaker_client.describe_user_profile.assert_called_once_with(
+        DomainId="d-kbnw5yk6tg8j",
+        UserProfileName="default-1617915559064",
+    )
+
+
 @patch("six.moves.builtins.open", mock_open(read_data='{"ResourceName": "SageMakerInstance"}'))
 @patch("os.path.exists", side_effect=mock_exists(NOTEBOOK_METADATA_FILE, True))
 @patch("sagemaker.session.sts_regional_endpoint", return_value=STS_ENDPOINT)