Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
May 21, 2021
Merged

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
7a38877
fix: make arns of all task resources aws-partition aware
May 5, 2021
723214e
Update utils.py
yuan-bwn May 7, 2021
bee2953
fix: use boto3 mock for utils test and create arn builder for task in…
May 12, 2021
0c9b0f0
fix: update arn builder method and its usage
May 14, 2021
162f9db
fix: move arn builder method into integration_resources module
May 17, 2021
ae85706
Update src/stepfunctions/steps/integration_resources.py
yuan-bwn May 19, 2021
33d624b
Merge branch 'main' into aws_partition_main
yuan-bwn May 20, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 51 additions & 21 deletions src/stepfunctions/steps/compute.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,32 @@
# permissions and limitations under the License.
from __future__ import absolute_import

from enum import Enum
from stepfunctions.steps.states import Task
from stepfunctions.steps.fields import Field
from stepfunctions.steps.utils import resource_integration_arn_builder
from stepfunctions.steps.integration_resources import IntegrationPattern, IntegrationServices, LambdaApi, GlueApi, BatchApi, EcsApi
from stepfunctions.steps.utils import get_service_integration_arn
from stepfunctions.steps.integration_resources import IntegrationPattern

Lambda = "lambda"
Glue = "glue"
Ecs = "ecs"
Batch = "batch"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Lambda = "lambda"
Glue = "glue"
Ecs = "ecs"
Batch = "batch"
LAMBDA_SERVICE_NAME = "lambda"
GLUE_SERVICE_NAME = "glue"
ECS_SERVICE_NAME = "ecs"
BATCH_SERVICE_NAME = "batch"

nit: Variable names should written be in snake_case. Class names are PascalCase. Constants in CAPITAL_SNAKE_CASE. Repeats.



class LambdaApi(Enum):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit (not blocking): Not necessary to be a class

Invoke = "invoke"


class GlueApi(Enum):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not blocking: also agree that these don't need to be classes.

If you do keep them as classes, we should include docstrings that explain what they are and their purpose :)

StartJobRun = "startJobRun"


class EcsApi(Enum):
RunTask = "runTask"


class BatchApi(Enum):
SubmitJob = "submitJob"


class LambdaStep(Task):
Expand Down Expand Up @@ -44,14 +66,16 @@ def __init__(self, state_id, wait_for_callback=False, **kwargs):
"""
Example resource arn: arn:aws:states:::lambda:invoke.waitForTaskToken
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.Lambda,
LambdaApi.Invoke,
IntegrationPattern.WaitForTaskToken)

kwargs[Field.Resource.value] = get_service_integration_arn(Lambda,
LambdaApi.Invoke,
IntegrationPattern.WaitForTaskToken)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: indentation is off

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will fix

else:
"""
Example resource arn: arn:aws:states:::lambda:invoke
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.Lambda, LambdaApi.Invoke)

kwargs[Field.Resource.value] = get_service_integration_arn(Lambda, LambdaApi.Invoke)


super(LambdaStep, self).__init__(state_id, **kwargs)
Expand Down Expand Up @@ -82,15 +106,17 @@ def __init__(self, state_id, wait_for_completion=True, **kwargs):
"""
Example resource arn: arn:aws:states:::glue:startJobRun.sync
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.Glue,
GlueApi.StartJobRun,
IntegrationPattern.WaitForCompletion)

kwargs[Field.Resource.value] = get_service_integration_arn(Glue,
GlueApi.StartJobRun,
IntegrationPattern.WaitForCompletion)
else:
"""
Example resource arn: arn:aws:states:::glue:startJobRun
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.Glue,
GlueApi.StartJobRun)

kwargs[Field.Resource.value] = get_service_integration_arn(Glue,
GlueApi.StartJobRun)

super(GlueStartJobRunStep, self).__init__(state_id, **kwargs)

Expand Down Expand Up @@ -120,15 +146,17 @@ def __init__(self, state_id, wait_for_completion=True, **kwargs):
"""
Example resource arn: arn:aws:states:::batch:submitJob.sync
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.Batch,
BatchApi.SubmitJob,
IntegrationPattern.WaitForCompletion)

kwargs[Field.Resource.value] = get_service_integration_arn(Batch,
BatchApi.SubmitJob,
IntegrationPattern.WaitForCompletion)
else:
"""
Example resource arn: arn:aws:states:::batch:submitJob
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.Batch,
BatchApi.SubmitJob)

kwargs[Field.Resource.value] = get_service_integration_arn(Batch,
BatchApi.SubmitJob)

super(BatchSubmitJobStep, self).__init__(state_id, **kwargs)

Expand Down Expand Up @@ -158,14 +186,16 @@ def __init__(self, state_id, wait_for_completion=True, **kwargs):
"""
Example resource arn: arn:aws:states:::ecs:runTask.sync
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.ECS,
EcsApi.RunTask,
IntegrationPattern.WaitForCompletion)

kwargs[Field.Resource.value] = get_service_integration_arn(Ecs,
EcsApi.RunTask,
IntegrationPattern.WaitForCompletion)
else:
"""
Example resource arn: arn:aws:states:::ecs:runTask
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.ECS,
EcsApi.RunTask)

kwargs[Field.Resource.value] = get_service_integration_arn(Ecs,
EcsApi.RunTask)

super(EcsRunTaskStep, self).__init__(state_id, **kwargs)
69 changes: 5 additions & 64 deletions src/stepfunctions/steps/integration_resources.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,75 +15,16 @@

from enum import Enum

"""
Enum classes for task integration resource arn builder
"""


class IntegrationPattern(Enum):
"""
Integration pattern enum classes for task integration resource arn builder
"""

WaitForTaskToken = "waitForTaskToken"
WaitForCompletion = "sync"
RequestResponse = ""


class IntegrationServices(Enum):
Lambda = "lambda"
SageMaker = "sagemaker"
Glue = "glue"
ECS = "ecs"
Batch = "batch"
DynamoDB = "dynamodb"
SNS = "sns"
SQS = "sqs"
ElasticMapReduce = "elasticmapreduce"


class LambdaApi(Enum):
Invoke = "invoke"


class SageMakerApi(Enum):
CreateTrainingJob = "createTrainingJob"
CreateTransformJob = "createTransformJob"
CreateModel = "createModel"
CreateEndpointConfig = "createEndpointConfig"
UpdateEndpoint = "updateEndpoint"
CreateEndpoint = "createEndpoint"
CreateHyperParameterTuningJob = "createHyperParameterTuningJob"
CreateProcessingJob = "createProcessingJob"


class GlueApi(Enum):
StartJobRun = "startJobRun"


class EcsApi(Enum):
RunTask = "runTask"


class BatchApi(Enum):
SubmitJob = "submitJob"


class DynamoDBApi(Enum):
GetItem = "getItem"
PutItem = "putItem"
DeleteItem = "deleteItem"
UpdateItem = "updateItem"


class SnsApi(Enum):
Publish = "publish"


class SqsApi(Enum):
SendMessage = "sendMessage"


class ElasticMapReduceApi(Enum):
CreateCluster = "createCluster"
TerminateCluster = "terminateCluster"
AddStep = "addStep"
CancelStep = "cancelStep"
SetClusterTerminationProtection = "setClusterTerminationProtection"
ModifyInstanceFleetByName = "modifyInstanceFleetByName"
ModifyInstanceGroupByName = "modifyInstanceGroupByName"
87 changes: 57 additions & 30 deletions src/stepfunctions/steps/sagemaker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,16 +12,31 @@
# permissions and limitations under the License.
from __future__ import absolute_import

from enum import Enum
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, resource_integration_arn_builder
from stepfunctions.steps.integration_resources import IntegrationPattern, IntegrationServices, SageMakerApi
from stepfunctions.steps.utils import tags_dict_to_kv_list, get_service_integration_arn
from stepfunctions.steps.integration_resources import IntegrationPattern

from sagemaker.workflow.airflow import training_config, transform_config, model_config, tuning_config, processing_config
from sagemaker.model import Model, FrameworkModel
from sagemaker.model_monitor import DataCaptureConfig

SageMaker = "sagemaker"


class SageMakerApi(Enum):
CreateTrainingJob = "createTrainingJob"
CreateTransformJob = "createTransformJob"
CreateModel = "createModel"
CreateEndpointConfig = "createEndpointConfig"
UpdateEndpoint = "updateEndpoint"
CreateEndpoint = "createEndpoint"
CreateHyperParameterTuningJob = "createHyperParameterTuningJob"
CreateProcessingJob = "createProcessingJob"


class TrainingStep(Task):

"""
Expand Down Expand Up @@ -62,15 +77,17 @@ def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=Non
"""
Example resource arn: arn:aws:states:::sagemaker:createTrainingJob.sync
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
SageMakerApi.CreateTrainingJob,
IntegrationPattern.WaitForCompletion)

kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
SageMakerApi.CreateTrainingJob,
IntegrationPattern.WaitForCompletion)
else:
"""
Example resource arn: arn:aws:states:::sagemaker:createTrainingJob
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
SageMakerApi.CreateTrainingJob)

kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
SageMakerApi.CreateTrainingJob)

if isinstance(job_name, str):
parameters = training_config(estimator=estimator, inputs=data, job_name=job_name, mini_batch_size=mini_batch_size)
Expand Down Expand Up @@ -154,15 +171,17 @@ def __init__(self, state_id, transformer, job_name, model_name, data, data_type=
"""
Example resource arn: arn:aws:states:::sagemaker:createTransformJob.sync
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
SageMakerApi.CreateTransformJob,
IntegrationPattern.WaitForCompletion)

kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
SageMakerApi.CreateTransformJob,
IntegrationPattern.WaitForCompletion)
else:
"""
Example resource arn: arn:aws:states:::sagemaker:createTransformJob
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
SageMakerApi.CreateTransformJob)

kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
SageMakerApi.CreateTransformJob)

if isinstance(job_name, str):
parameters = transform_config(
Expand Down Expand Up @@ -248,8 +267,9 @@ def __init__(self, state_id, model, model_name=None, instance_type=None, tags=No
"""
Example resource arn: arn:aws:states:::sagemaker:createModel
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
SageMakerApi.CreateModel)

kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
SageMakerApi.CreateModel)

super(ModelStep, self).__init__(state_id, **kwargs)

Expand Down Expand Up @@ -293,8 +313,9 @@ def __init__(self, state_id, endpoint_config_name, model_name, initial_instance_
"""
Example resource arn: arn:aws:states:::sagemaker:createEndpointConfig
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
SageMakerApi.CreateEndpointConfig)

kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
SageMakerApi.CreateEndpointConfig)

kwargs[Field.Parameters.value] = parameters

Expand Down Expand Up @@ -330,14 +351,16 @@ def __init__(self, state_id, endpoint_name, endpoint_config_name, tags=None, upd
"""
Example resource arn: arn:aws:states:::sagemaker:updateEndpoint
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
SageMakerApi.UpdateEndpoint)

kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
SageMakerApi.UpdateEndpoint)
else:
"""
Example resource arn: arn:aws:states:::sagemaker:createEndpoint
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
SageMakerApi.CreateEndpoint)

kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
SageMakerApi.CreateEndpoint)

kwargs[Field.Parameters.value] = parameters

Expand Down Expand Up @@ -378,15 +401,17 @@ def __init__(self, state_id, tuner, job_name, data, wait_for_completion=True, ta
"""
Example resource arn: arn:aws:states:::sagemaker:createHyperParameterTuningJob.sync
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
SageMakerApi.CreateHyperParameterTuningJob,
IntegrationPattern.WaitForCompletion)

kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
SageMakerApi.CreateHyperParameterTuningJob,
IntegrationPattern.WaitForCompletion)
else:
"""
Example resource arn: arn:aws:states:::sagemaker:createHyperParameterTuningJob
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
SageMakerApi.CreateHyperParameterTuningJob)

kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
SageMakerApi.CreateHyperParameterTuningJob)

parameters = tuning_config(tuner=tuner, inputs=data, job_name=job_name).copy()

Expand Down Expand Up @@ -436,15 +461,17 @@ def __init__(self, state_id, processor, job_name, inputs=None, outputs=None, exp
"""
Example resource arn: arn:aws:states:::sagemaker:createProcessingJob.sync
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
SageMakerApi.CreateProcessingJob,
IntegrationPattern.WaitForCompletion)

kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
SageMakerApi.CreateProcessingJob,
IntegrationPattern.WaitForCompletion)
else:
"""
Example resource arn: arn:aws:states:::sagemaker:createProcessingJob
"""
kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
SageMakerApi.CreateProcessingJob)

kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
SageMakerApi.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)
Expand Down
Loading