feature: sts endpoint support

Author: Gray

src/sagemaker/session.py

@@ -63,7 +63,7 @@ class Session(object):
     a naming convention which includes the current AWS account ID.
     """
 
-    def __init__(self, boto_session=None, sagemaker_client=None, sagemaker_runtime_client=None):
+    def __init__(self, boto_session=None, sagemaker_client=None, sagemaker_runtime_client=None, sts_endpoint_url=None):
         """Initialize a SageMaker ``Session``.
 
         Args:
@@ -75,18 +75,19 @@ def __init__(self, boto_session=None, sagemaker_client=None, sagemaker_runtime_c
             sagemaker_runtime_client (boto3.SageMakerRuntime.Client): Client which makes ``InvokeEndpoint``
                 calls to Amazon SageMaker (default: None). Predictors created using this ``Session`` use this client.
                 If not provided, one will be created using this instance's ``boto_session``.
+            sts_endpoint_url (str): Endpoint URL for STS endpoint. If none provided, boto3 will default to use sts.amazonaws.com.
         """
         self._default_bucket = None
-
+        self.sts_endpoint_url = sts_endpoint_url
         sagemaker_config_file = os.path.join(os.path.expanduser('~'), '.sagemaker', 'config.yaml')
         if os.path.exists(sagemaker_config_file):
             self.config = yaml.load(open(sagemaker_config_file, 'r'))
         else:
             self.config = None
 
-        self._initialize(boto_session, sagemaker_client, sagemaker_runtime_client)
+        self._initialize(boto_session, sagemaker_client, sagemaker_runtime_client, sts_endpoint_url)
 
-    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
+    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client, sts_endpoint_url):
         """Initialize this SageMaker Session.
 
         Creates or uses a boto_session, sagemaker_client and sagemaker_runtime_client.
@@ -109,6 +110,8 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
 
         prepend_user_agent(self.sagemaker_runtime_client)
 
+        self.sts_endpoint_url = sts_endpoint_url
+
         self.local_mode = False
 
     @property
@@ -177,7 +180,11 @@ def default_bucket(self):
         if self._default_bucket:
             return self._default_bucket
 
-        account = self.boto_session.client('sts').get_caller_identity()['Account']
+        if self.sts_endpoint_url:
+            account = self.boto_session.client('sts', endpoint_url=self.sts_endpoint_url).get_caller_identity()['Account']
+        else:
+            account = self.boto_session.client('sts').get_caller_identity()['Account']
+
         region = self.boto_session.region_name
         default_bucket = 'sagemaker-{}-{}'.format(region, account)
 
@@ -1089,7 +1096,10 @@ def get_caller_identity_arn(self):
         Returns:
             (str): The ARN user or role
         """
-        assumed_role = self.boto_session.client('sts').get_caller_identity()['Arn']
+        if self.sts_endpoint_url:
+            assumed_role = self.boto_session.client('sts', endpoint_url=self.sts_endpoint_url).get_caller_identity()['Arn']
+        else:
+            assumed_role = self.boto_session.client('sts').get_caller_identity()['Arn']
 
         if 'AmazonSageMaker-ExecutionRole' in assumed_role:
             role = re.sub(r'^(.+)sts::(\d+):assumed-role/(.+?)/.*$', r'\1iam::\2:role/service-role/\3', assumed_role)

tests/unit/test_session.py

@@ -53,6 +53,15 @@ def test_get_execution_role():
     assert actual == 'arn:aws:iam::369233609183:role/SageMakerRole'
 
 
+def test_get_execution_role_with_sts_endpoint():
+    endpoint_url = "https://sts.{0}.amazonaws.com".format(REGION)
+    session = Session(sts_endpoint_url=endpoint_url)
+    session.get_caller_identity_arn.return_value = 'arn:aws:iam::369233609183:role/SageMakerRole'
+
+    actual = get_execution_role(session)
+    assert actual == 'arn:aws:iam::369233609183:role/SageMakerRole'
+
+
 def test_get_execution_role_works_with_service_role():
     session = Mock()
     session.get_caller_identity_arn.return_value = \