Merge pull request #61 from yosefbs/master

Support china regions

Author: tomfaulhaber

sagemaker_run_notebook/container_build.py

@@ -38,6 +38,7 @@ def create_project(repo_name, role, zipfile, base_image=default_base):
     sts = session.client("sts")
     identity = sts.get_caller_identity()
     account = identity["Account"]
+    partition = identity["Arn"].split(':')[1]
     args = {
         "name": f"create-sagemaker-container-{repo_name}",
         "description": f"Build the container {repo_name} for running notebooks in SageMaker",
@@ -56,7 +57,7 @@ def create_project(repo_name, role, zipfile, base_image=default_base):
             ],
             "privilegedMode": True,
         },
-        "serviceRole": f"arn:aws:iam::{account}:role/{role}",
+        "serviceRole": f"arn:{partition}:iam::{account}:role/{role}",
     }
 
     response = client.create_project(**args)

sagemaker_run_notebook/lambda_function.py

@@ -20,19 +20,21 @@ def execute_notebook(
 ):
     session = ensure_session()
     region = session.region_name
-
-    account = session.client("sts").get_caller_identity()["Account"]
+    caller_id=session.client("sts").get_caller_identity()
+    partition = caller_id["Arn"].split(':')[1]
+    account = caller_id["Account"]
+    domain = domain_for_region(region)
     if not image:
         image = "notebook-runner"
     if "/" not in image:
-        image = f"{account}.dkr.ecr.{region}.amazonaws.com/{image}"
+        image = f"{account}.dkr.ecr.{region}.{domain}/{image}"
     if ":" not in image:
         image = image + ":latest"
 
     if not role:
         role = f"BasicExecuteNotebookRole-{region}"
     if "/" not in role:
-        role = f"arn:aws:iam::{account}:role/{role}"
+        role = f"arn:{partition}:iam::{account}:role/{role}"
 
     if output_prefix is None:
         output_prefix = os.path.dirname(input_path)
@@ -149,6 +151,21 @@ def ensure_session(session=None):
         session = boto3.session.Session()
     return session
 
+def domain_for_region(region):
+    """Get the DNS suffix for the given region.
+    Args:
+        region (str): AWS region name
+    Returns:
+        str: the DNS suffix
+    """
+    if region.startswith("us-iso-"):
+        return "c2s.ic.gov"
+    if region.startswith("us-isob-"):
+        return "sc2s.sgov.gov"
+    if region.startswith("cn-"):
+        return "amazonaws.com.cn"    
+    return  "amazonaws.com"
+
 
 def lambda_handler(event, context):
     job = execute_notebook(

sagemaker_run_notebook/run_notebook.py

@@ -27,10 +27,10 @@
 import botocore
 import boto3
 
-from .utils import default_bucket, get_execution_role
+from .utils import default_bucket, domain_for_region, get_execution_role
 
 abbrev_image_pat = re.compile(
-    r"(?P<account>\d+).dkr.ecr.(?P<region>[^.]+).amazonaws.com/(?P<image>[^:/]+)(?P<tag>:[^:]+)?"
+    r"(?P<account>\d+).dkr.ecr.(?P<region>[^.]+).(amazonaws.com|amazonaws.com.cn)/(?P<image>[^:/]+)(?P<tag>:[^:]+)?"
 )
 
 
@@ -46,7 +46,8 @@ def abbreviate_image(image):
         return image
 
 
-abbrev_role_pat = re.compile(r"arn:aws:iam::(?P<account>\d+):role/(?P<name>[^/]+)")
+abbrev_role_pat = re.compile(
+    r"arn:aws([^:]*):iam::(?P<account>\d+):role/(?P<name>[^/]+)")
 
 
 def abbreviate_role(role):
@@ -124,13 +125,17 @@ def execute_notebook(
     if not role:
         role = get_execution_role(session)
     elif "/" not in role:
-        account = session.client("sts").get_caller_identity()["Account"]
-        role = "arn:aws:iam::{}:role/{}".format(account, role)
+        identity = session.client("sts").get_caller_identity()
+        account = identity["Account"]
+        partition = identity["Arn"].split(':')[1]
+        role = "arn:{}:iam::{}:role/{}".format(partition, account, role)
 
     if "/" not in image:
         account = session.client("sts").get_caller_identity()["Account"]
         region = session.region_name
-        image = "{}.dkr.ecr.{}.amazonaws.com/{}:latest".format(account, region, image)
+        domain = domain_for_region(region)
+        image = "{}.dkr.ecr.{}.{}/{}:latest".format(
+            account, region, domain, image)
 
     if notebook == None:
         notebook = input_path
@@ -140,7 +145,8 @@ def execute_notebook(
     timestamp = time.strftime("%Y-%m-%d-%H-%M-%S", time.gmtime())
 
     job_name = (
-        ("papermill-" + re.sub(r"[^-a-zA-Z0-9]", "-", nb_name))[: 62 - len(timestamp)]
+        ("papermill-" + re.sub(r"[^-a-zA-Z0-9]",
+         "-", nb_name))[: 62 - len(timestamp)]
         + "-"
         + timestamp
     )
@@ -628,8 +634,10 @@ def create_lambda(role=None, session=None):
         # time.sleep(30) # wait for eventual consistency, we hope
 
     if "/" not in role:
-        account = session.client("sts").get_caller_identity()["Account"]
-        role = "arn:aws:iam::{}:role/{}".format(account, role)
+        identity = session.client("sts").get_caller_identity()
+        account = identity["Account"]
+        partition = identity["Arn"].split(':')[1]
+        role = "arn:{}:iam::{}:role/{}".format(partition, account, role)
 
     code_bytes = zip_bytes(code_file)
 
@@ -780,7 +788,9 @@ def proc(extras):
     if "/" not in image:
         account = session.client("sts").get_caller_identity()["Account"]
         region = session.region_name
-        image = "{}.dkr.ecr.{}.amazonaws.com/{}:latest".format(account, region, image)
+        domain =  domain_for_region(region)
+        image = "{}.dkr.ecr.{}.{}/{}:latest".format(
+            account, region, domain, image)
 
     if not role:
         try:
@@ -789,8 +799,10 @@ def proc(extras):
             role = "BasicExecuteNotebookRole-{}".format(session.region_name)
 
     if "/" not in role:
-        account = session.client("sts").get_caller_identity()["Account"]
-        role = "arn:aws:iam::{}:role/{}".format(account, role)
+        identity = session.client("sts").get_caller_identity()
+        account = identity["Account"]
+        partition = identity["Arn"].split(':')[1]
+        role = "arn:{}:iam::{}:role/{}".format(partition, account, role)
 
     if input_path is None:
         input_path = upload_notebook(notebook)
@@ -849,7 +861,7 @@ def schedule(
 
     Creates a CloudWatch Event rule to invoke the installed Lambda either on the provided schedule or in response
     to the provided event. \
-  
+
     :meth:`schedule` can upload a local notebook file to run or use one previously uploaded to S3. 
     To find jobs run by the schedule, see :meth:`list_runs` using the `rule` argument to filter to 
     a specific rule. To download the results, see :meth:`download_notebook` (or :meth:`download_all` 
@@ -905,7 +917,9 @@ def proc(extras):
     if "/" not in image:
         account = session.client("sts").get_caller_identity()["Account"]
         region = session.region_name
-        image = "{}.dkr.ecr.{}.amazonaws.com/{}:latest".format(account, region, image)
+        domain = domain_for_region(region)
+        image = "{}.dkr.ecr.{}.{}/{}:latest".format(
+            account, region, domain, image)
 
     if not role:
         try:
@@ -914,8 +928,10 @@ def proc(extras):
             role = "BasicExecuteNotebookRole-{}".format(session.region_name)
 
     if "/" not in role:
-        account = session.client("sts").get_caller_identity()["Account"]
-        role = "arn:aws:iam::{}:role/{}".format(account, role)
+        identity = session.client("sts").get_caller_identity()
+        account = identity["Account"]
+        partition = identity["Arn"].split(':')[1]
+        role = "arn:{}:iam::{}:role/{}".format(partition, account, role)
 
     if input_path is None:
         input_path = upload_notebook(notebook)
@@ -945,11 +961,12 @@ def proc(extras):
         Description='Rule to run the Jupyter notebook "{}"'.format(notebook),
         **kwargs,
     )
-
-    account = session.client("sts").get_caller_identity()["Account"]
+    identity = session.client("sts").get_caller_identity()
+    account = identity["Account"]
+    partition = identity["Arn"].split(':')[1]
     region = session.region_name
-    target_arn = "arn:aws:lambda:{}:{}:function:{}".format(
-        region, account, lambda_function_name
+    target_arn = "arn:{}:lambda:{}:{}:function:{}".format(
+        partition, region, account, lambda_function_name
     )
 
     result = events.put_targets(
@@ -1082,7 +1099,7 @@ def base_image(s):
         return s
 
 
-role_pat = re.compile(r"arn:aws:iam::([0-9]+):role/(.*)$")
+role_pat = re.compile(r"arn:aws([^:]*):iam::([0-9]+):role/(.*)$")
 
 
 def base_role(s):

sagemaker_run_notebook/utils.py

@@ -116,11 +116,11 @@ def sts_regional_endpoint(region):
     Returns:
         str: AWS STS regional endpoint
     """
-    domain = _domain_for_region(region)
+    domain = domain_for_region(region)
     return "https://sts.{}.{}".format(region, domain)
 
 
-def _domain_for_region(region):
+def domain_for_region(region):
     """Get the DNS suffix for the given region.
 
     Args:
@@ -129,7 +129,13 @@ def _domain_for_region(region):
     Returns:
         str: the DNS suffix
     """
-    return "c2s.ic.gov" if region == "us-iso-east-1" else "amazonaws.com"
+    if region.startswith("us-iso-"):
+        return "c2s.ic.gov"
+    if region.startswith("us-isob-"):
+        return "sc2s.sgov.gov"
+    if region.startswith("cn-"):
+        return "amazonaws.com.cn"
+    return "amazonaws.com"
 
 
 def get_execution_role(session):