fix: make arns of all task resources aws-partition aware

Author: Bowen Yuan

src/stepfunctions/steps/compute.py

@@ -14,6 +14,7 @@
 
 from stepfunctions.steps.states import Task
 from stepfunctions.steps.fields import Field
+from stepfunctions.steps.utils import get_aws_partition
 
 
 class LambdaStep(Task):
@@ -38,9 +39,9 @@ def __init__(self, state_id, wait_for_callback=False, **kwargs):
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
         if wait_for_callback:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::lambda:invoke.waitForTaskToken'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::lambda:invoke.waitForTaskToken'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::lambda:invoke'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::lambda:invoke'
 
         super(LambdaStep, self).__init__(state_id, **kwargs)
 
@@ -67,9 +68,9 @@ def __init__(self, state_id, wait_for_completion=True, **kwargs):
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
         if wait_for_completion:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::glue:startJobRun.sync'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::glue:startJobRun.sync'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::glue:startJobRun'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::glue:startJobRun'
 
         super(GlueStartJobRunStep, self).__init__(state_id, **kwargs)
 
@@ -96,9 +97,9 @@ def __init__(self, state_id, wait_for_completion=True, **kwargs):
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
         if wait_for_completion:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::batch:submitJob.sync'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::batch:submitJob.sync'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::batch:submitJob'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::batch:submitJob'
 
         super(BatchSubmitJobStep, self).__init__(state_id, **kwargs)
 
@@ -125,8 +126,8 @@ def __init__(self, state_id, wait_for_completion=True, **kwargs):
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
         if wait_for_completion:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::ecs:runTask.sync'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::ecs:runTask.sync'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::ecs:runTask'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::ecs:runTask'
 
         super(EcsRunTaskStep, self).__init__(state_id, **kwargs)

src/stepfunctions/steps/sagemaker.py

@@ -15,7 +15,7 @@
 from stepfunctions.inputs import ExecutionInput, StepInput
 from stepfunctions.steps.states import Task
 from stepfunctions.steps.fields import Field
-from stepfunctions.steps.utils import tags_dict_to_kv_list
+from stepfunctions.steps.utils import tags_dict_to_kv_list, get_aws_partition
 
 from sagemaker.workflow.airflow import training_config, transform_config, model_config, tuning_config, processing_config
 from sagemaker.model import Model, FrameworkModel
@@ -58,9 +58,9 @@ def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=Non
         self.job_name = job_name
 
         if wait_for_completion:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sagemaker:createTrainingJob.sync'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sagemaker:createTrainingJob.sync'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sagemaker:createTrainingJob'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sagemaker:createTrainingJob'
 
         if isinstance(job_name, str):
             parameters = training_config(estimator=estimator, inputs=data, job_name=job_name, mini_batch_size=mini_batch_size)
@@ -141,9 +141,9 @@ def __init__(self, state_id, transformer, job_name, model_name, data, data_type=
             join_source (str): The source of data to be joined to the transform output. It can be set to ‘Input’ meaning the entire input record will be joined to the inference result. You can use OutputFilter to select the useful portion before uploading to S3. (default: None). Valid values: Input, None.
         """
         if wait_for_completion:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sagemaker:createTransformJob.sync'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sagemaker:createTransformJob.sync'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sagemaker:createTransformJob'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sagemaker:createTransformJob'
 
         if isinstance(job_name, str):
             parameters = transform_config(
@@ -225,7 +225,7 @@ def __init__(self, state_id, model, model_name=None, instance_type=None, tags=No
             parameters['Tags'] = tags_dict_to_kv_list(tags)
 
         kwargs[Field.Parameters.value] = parameters
-        kwargs[Field.Resource.value] = 'arn:aws:states:::sagemaker:createModel'
+        kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sagemaker:createModel'
 
         super(ModelStep, self).__init__(state_id, **kwargs)
 
@@ -266,7 +266,7 @@ def __init__(self, state_id, endpoint_config_name, model_name, initial_instance_
         if tags:
             parameters['Tags'] = tags_dict_to_kv_list(tags)
 
-        kwargs[Field.Resource.value] = 'arn:aws:states:::sagemaker:createEndpointConfig'
+        kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sagemaker:createEndpointConfig'
         kwargs[Field.Parameters.value] = parameters
 
         super(EndpointConfigStep, self).__init__(state_id, **kwargs)
@@ -298,9 +298,9 @@ def __init__(self, state_id, endpoint_name, endpoint_config_name, tags=None, upd
             parameters['Tags'] = tags_dict_to_kv_list(tags)
 
         if update:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sagemaker:updateEndpoint'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sagemaker:updateEndpoint'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sagemaker:createEndpoint'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sagemaker:createEndpoint'
 
         kwargs[Field.Parameters.value] = parameters
 
@@ -338,9 +338,9 @@ def __init__(self, state_id, tuner, job_name, data, wait_for_completion=True, ta
             tags (list[dict], optional): `List to tags <https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html>`_ to associate with the resource.
         """
         if wait_for_completion:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sagemaker:createHyperParameterTuningJob.sync'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sagemaker:createHyperParameterTuningJob.sync'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sagemaker:createHyperParameterTuningJob'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sagemaker:createHyperParameterTuningJob'
 
         parameters = tuning_config(tuner=tuner, inputs=data, job_name=job_name).copy()
 
@@ -387,9 +387,9 @@ def __init__(self, state_id, processor, job_name, inputs=None, outputs=None, exp
             tags (list[dict], optional): `List to tags <https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html>`_ to associate with the resource.
         """
         if wait_for_completion:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sagemaker:createProcessingJob.sync'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sagemaker:createProcessingJob.sync'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sagemaker:createProcessingJob'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sagemaker:createProcessingJob'
 
         if isinstance(job_name, str):
             parameters = processing_config(processor=processor, inputs=inputs, outputs=outputs, container_arguments=container_arguments, container_entrypoint=container_entrypoint, kms_key_id=kms_key_id, job_name=job_name)

src/stepfunctions/steps/service.py

@@ -14,6 +14,7 @@
 
 from stepfunctions.steps.states import Task
 from stepfunctions.steps.fields import Field
+from stepfunctions.steps.utils import get_aws_partition
 
 
 class DynamoDBGetItemStep(Task):
@@ -35,7 +36,7 @@ def __init__(self, state_id, **kwargs):
             result_path (str, optional): Path specifying the raw input’s combination with or replacement by the state’s result. (default: '$')
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
-        kwargs[Field.Resource.value] = 'arn:aws:states:::dynamodb:getItem'
+        kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::dynamodb:getItem'
         super(DynamoDBGetItemStep, self).__init__(state_id, **kwargs)
 
 
@@ -59,7 +60,7 @@ def __init__(self, state_id, **kwargs):
             result_path (str, optional): Path specifying the raw input’s combination with or replacement by the state’s result. (default: '$')
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
-        kwargs[Field.Resource.value] = 'arn:aws:states:::dynamodb:putItem'
+        kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::dynamodb:putItem'
         super(DynamoDBPutItemStep, self).__init__(state_id, **kwargs)
 
 
@@ -83,7 +84,7 @@ def __init__(self, state_id, **kwargs):
             result_path (str, optional): Path specifying the raw input’s combination with or replacement by the state’s result. (default: '$')
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
-        kwargs[Field.Resource.value] = 'arn:aws:states:::dynamodb:deleteItem'
+        kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::dynamodb:deleteItem'
         super(DynamoDBDeleteItemStep, self).__init__(state_id, **kwargs)
 
 
@@ -107,7 +108,7 @@ def __init__(self, state_id, **kwargs):
             result_path (str, optional): Path specifying the raw input’s combination with or replacement by the state’s result. (default: '$')
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
-        kwargs[Field.Resource.value] = 'arn:aws:states:::dynamodb:updateItem'
+        kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::dynamodb:updateItem'
         super(DynamoDBUpdateItemStep, self).__init__(state_id, **kwargs)
 
 
@@ -133,9 +134,9 @@ def __init__(self, state_id, wait_for_callback=False, **kwargs):
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
         if wait_for_callback:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sns:publish.waitForTaskToken'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sns:publish.waitForTaskToken'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sns:publish'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sns:publish'
 
         super(SnsPublishStep, self).__init__(state_id, **kwargs)
 
@@ -162,9 +163,9 @@ def __init__(self, state_id, wait_for_callback=False, **kwargs):
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
         if wait_for_callback:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sqs:sendMessage.waitForTaskToken'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sqs:sendMessage.waitForTaskToken'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::sqs:sendMessage'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::sqs:sendMessage'
 
         super(SqsSendMessageStep, self).__init__(state_id, **kwargs)
 
@@ -190,9 +191,9 @@ def __init__(self, state_id, wait_for_completion=True, **kwargs):
             wait_for_completion (bool, optional): Boolean value set to `True` if the Task state should wait to complete before proceeding to the next step in the workflow. (default: True)
         """
         if wait_for_completion:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::elasticmapreduce:createCluster.sync'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::elasticmapreduce:createCluster.sync'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::elasticmapreduce:createCluster'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::elasticmapreduce:createCluster'
 
         super(EmrCreateClusterStep, self).__init__(state_id, **kwargs)
 
@@ -218,9 +219,9 @@ def __init__(self, state_id, wait_for_completion=True, **kwargs):
             wait_for_completion (bool, optional): Boolean value set to `True` if the Task state should wait to complete before proceeding to the next step in the workflow. (default: True)
         """
         if wait_for_completion:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::elasticmapreduce:terminateCluster.sync'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::elasticmapreduce:terminateCluster.sync'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::elasticmapreduce:terminateCluster'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::elasticmapreduce:terminateCluster'
 
         super(EmrTerminateClusterStep, self).__init__(state_id, **kwargs)
 
@@ -246,9 +247,9 @@ def __init__(self, state_id, wait_for_completion=True, **kwargs):
             wait_for_completion (bool, optional): Boolean value set to `True` if the Task state should wait to complete before proceeding to the next step in the workflow. (default: True)
         """
         if wait_for_completion:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::elasticmapreduce:addStep.sync'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::elasticmapreduce:addStep.sync'
         else:
-            kwargs[Field.Resource.value] = 'arn:aws:states:::elasticmapreduce:addStep'
+            kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::elasticmapreduce:addStep'
 
         super(EmrAddStepStep, self).__init__(state_id, **kwargs)
 
@@ -272,7 +273,7 @@ def __init__(self, state_id, **kwargs):
             result_path (str, optional): Path specifying the raw input’s combination with or replacement by the state’s result. (default: '$')
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
-        kwargs[Field.Resource.value] = 'arn:aws:states:::elasticmapreduce:cancelStep'
+        kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::elasticmapreduce:cancelStep'
 
         super(EmrCancelStepStep, self).__init__(state_id, **kwargs)
 
@@ -296,7 +297,7 @@ def __init__(self, state_id, **kwargs):
             result_path (str, optional): Path specifying the raw input’s combination with or replacement by the state’s result. (default: '$')
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
-        kwargs[Field.Resource.value] = 'arn:aws:states:::elasticmapreduce:setClusterTerminationProtection'
+        kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::elasticmapreduce:setClusterTerminationProtection'
 
         super(EmrSetClusterTerminationProtectionStep, self).__init__(state_id, **kwargs)
 
@@ -320,7 +321,7 @@ def __init__(self, state_id, **kwargs):
             result_path (str, optional): Path specifying the raw input’s combination with or replacement by the state’s result. (default: '$')
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
-        kwargs[Field.Resource.value] = 'arn:aws:states:::elasticmapreduce:modifyInstanceFleetByName'
+        kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::elasticmapreduce:modifyInstanceFleetByName'
 
         super(EmrModifyInstanceFleetByNameStep, self).__init__(state_id, **kwargs)
 
@@ -344,7 +345,7 @@ def __init__(self, state_id, **kwargs):
             result_path (str, optional): Path specifying the raw input’s combination with or replacement by the state’s result. (default: '$')
             output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
         """
-        kwargs[Field.Resource.value] = 'arn:aws:states:::elasticmapreduce:modifyInstanceGroupByName'
+        kwargs[Field.Resource.value] = 'arn:' + get_aws_partition() + ':states:::elasticmapreduce:modifyInstanceGroupByName'
 
         super(EmrModifyInstanceGroupByNameStep, self).__init__(state_id, **kwargs)
 

src/stepfunctions/steps/utils.py

@@ -12,6 +12,31 @@
 # permissions and limitations under the License.
 from __future__ import absolute_import
 
+import boto3
+import logging
+
+logger = logging.getLogger('stepfunctions')
+
+
 def tags_dict_to_kv_list(tags_dict):
-    kv_list = [{"Key": k, "Value": v} for k,v in tags_dict.items()]
-    return kv_list
+    kv_list = [{"Key": k, "Value": v} for k, v in tags_dict.items()]
+    return kv_list
+
+
+# Obtain matching aws partition name based on region
+def get_aws_partition():
+    partitions = boto3.session.Session().get_available_partitions()
+    cur_region = boto3.session.Session().region_name
+    cur_partition = "aws"
+
+    if cur_region is None:
+        logger.warning("No region detected for the session, will use default partition: aws")
+        return cur_partition
+
+    for partition in partitions:
+        regions = boto3.session.Session().get_available_regions("stepfunctions", partition)
+        if cur_region in regions:
+            cur_partition = partition
+            return cur_partition
+
+    return cur_partition