Merge branch 'main' into training-step-with-dynamic-output-path

Author: shivlaks

CONTRIBUTING.md

@@ -11,7 +11,7 @@ information to effectively respond to your bug report or contribution.
 
 * [Table of Contents](#table-of-contents)
 * [Reporting Bugs/Feature Requests](#reporting-bugsfeature-requests)
-* [Contributing via Pull Requests (PRs)](#contributing-via-pull-requests-prs) 
+* [Contributing via Pull Requests (PRs)](#contributing-via-pull-requests-prs)
   * [Pulling Down the Code](#pulling-down-the-code)
   * [Running the Unit Tests](#running-the-unit-tests)
   * [Running the Integration Tests](#running-the-integration-tests)
@@ -61,20 +61,56 @@ Before sending us a pull request, please ensure that:
 1. Install test dependencies, including coverage, using `pip install ".[test]"`
 1. Run the following tox command and verify that all code checks and unit tests pass: `tox tests/unit`
 
-You can also run a single test with the following command: `tox -e py36 -- -s -vv <path_to_file><file_name>::<test_function_name>`  
+You can also run a single test with the following command: `tox -e py36 -- -s -vv <path_to_file><file_name>::<test_function_name>`
   * Note that the coverage test will fail if you only run a single test, so make sure to surround the command with `export IGNORE_COVERAGE=-` and `unset IGNORE_COVERAGE`
   * Example: `export IGNORE_COVERAGE=- ; tox -e py36 -- -s -vv tests/unit/test_sagemaker_steps.py::test_training_step_creation_with_model ; unset IGNORE_COVERAGE`
 
 
 ### Running the Integration Tests
 
-Our CI system runs integration tests (the ones in the `tests/integ` directory), in parallel, for every Pull Request.  
-You should only worry about manually running any new integration tests that you write, or integration tests that test an area of code that you've modified.  
+Our CI system runs integration tests (the ones in the `tests/integ` directory), in parallel, for every Pull Request.
+You should only worry about manually running any new integration tests that you write, or integration tests that test an area of code that you've modified.
+#### Setup
 
-1. Follow the instructions at [Set Up the AWS Command Line Interface (AWS CLI)](https://docs.aws.amazon.com/polly/latest/dg/setup-aws-cli.html).
+If you haven't done so already, install tox and test dependencies:
+1. `pip install tox`
+1. `pip install .[test]`
+
+#### AWS Credentials
+Follow the instructions at [Set Up the AWS Command Line Interface (AWS CLI)](https://docs.aws.amazon.com/polly/latest/dg/setup-aws-cli.html).
+#### Create IAM Roles
+
+The tests use two IAM roles to give Step Functions and SageMaker permissions to access AWS resources in your account. Use the following commands in the root directory of this repository:
+
+```bash
+aws iam create-role \
+   --role-name StepFunctionsMLWorkflowExecutionFullAccess \
+   --assume-role-policy-document file://tests/integ/resources/StepFunctionsMLWorkflowExecutionFullAccess-TrustPolicy.json
+```
+
+```bash
+aws iam put-role-policy \
+   --role-name StepFunctionsMLWorkflowExecutionFullAccess \
+   --policy-name StepFunctionsMLWorkflowExecutionFullAccess \
+   --policy-document file://tests/integ/resources/StepFunctionsMLWorkflowExecutionFullAccess-Policy.json
+```
+
+```bash
+aws iam create-role \
+   --role-name SageMakerRole \
+   --assume-role-policy-document file://tests/integ/resources/SageMaker-TrustPolicy.json
+```
+
+```bash
+aws iam attach-role-policy \
+   --role-name SageMakerRole \
+   --policy-arn arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
+```
+#### Execute the tests
 1. To run a test, specify the test file and method you want to run per the following command: `tox -e py36 -- -s -vv <path_to_file><file_name>::<test_function_name>`
    * Note that the coverage test will fail if you only run a single test, so make sure to surround the command with `export IGNORE_COVERAGE=-` and `unset IGNORE_COVERAGE`
    * Example: `export IGNORE_COVERAGE=- ; tox -e py36 -- -s -vv tests/integ/test_state_machine_definition.py::test_wait_state_machine_creation ; unset IGNORE_COVERAGE`
+1. To run all integration tests, run the following command: `tox tests/integ`
 
 ### Making and Testing Your Change
 
@@ -88,7 +124,7 @@ You should only worry about manually running any new integration tests that you
       1. Guard against future breaking changes to lower the maintenance cost.
    1. Please focus on the specific change you are contributing. If you also reformat all the code, it will be hard for us to focus on your change.
 1. Run all the unit tests as per [Running the Unit Tests](#running-the-unit-tests), and verify that all checks and tests pass.
-   1. Note that this also runs tools that may be necessary for the automated build to pass (ex: code reformatting by 'black').  
+   1. Note that this also runs tools that may be necessary for the automated build to pass (ex: code reformatting by 'black').
 
 
 ### Committing Your Change

README.rst

@@ -27,6 +27,7 @@ Table of Contents
 - `Getting Started With Sample Jupyter Notebooks <#getting-started-with-sample-jupyter-notebooks>`__
 - `Installing the AWS Step Functions Data Science SDK <#installing-the-aws-step-functions-data-science-sdk>`__
 - `Overview of SDK <#overview-of-sdk>`__
+- `Contributing <#contributing>`__
 - `AWS Permissions <#aws-permissions>`__
 - `Licensing <#licensing>`__
 - `Verifying the Signature <#verifying-the-signature>`__
@@ -309,6 +310,12 @@ the CloudFormation template in a different region, please make sure to update
 the region specific AWS resources (such as the Lambda ARN and Training Image)
 in the StateMachine definition.
 
+Contributing
+------------
+We welcome community contributions and pull requests. See
+`CONTRIBUTING.md <https://github.com/aws/aws-step-functions-data-science-sdk-python/blob/main/CONTRIBUTING.md>`__ for
+information on how to set up a development environment, run tests and submit code.
+
 AWS Permissions
 ---------------
 As a managed service, AWS Step Functions performs operations on your behalf on

doc/services.rst

@@ -8,12 +8,14 @@ This module provides classes to build steps that integrate with Amazon DynamoDB,
 
 - `Amazon DynamoDB <#amazon-dynamodb>`__
 
+- `Amazon EMR <#amazon-emr>`__
+
+- `Amazon EventBridge <#amazon-eventbridge>`__
+
 - `Amazon SNS <#amazon-sns>`__
 
 - `Amazon SQS <#amazon-sqs>`__
 
-- `Amazon EMR <#amazon-emr>`__
-
 
 Amazon DynamoDB
 ----------------
@@ -25,14 +27,6 @@ Amazon DynamoDB
 
 .. autoclass:: stepfunctions.steps.service.DynamoDBUpdateItemStep
 
-Amazon SNS
------------
-.. autoclass:: stepfunctions.steps.service.SnsPublishStep
-
-Amazon SQS
------------
-.. autoclass:: stepfunctions.steps.service.SqsSendMessageStep
-
 Amazon EMR
 -----------
 .. autoclass:: stepfunctions.steps.service.EmrCreateClusterStep
@@ -48,3 +42,15 @@ Amazon EMR
 .. autoclass:: stepfunctions.steps.service.EmrModifyInstanceFleetByNameStep
 
 .. autoclass:: stepfunctions.steps.service.EmrModifyInstanceGroupByNameStep
+
+Amazon EventBridge
+-----------
+.. autoclass:: stepfunctions.steps.service.EventBridgePutEventsStep
+
+Amazon SNS
+-----------
+.. autoclass:: stepfunctions.steps.service.SnsPublishStep
+
+Amazon SQS
+-----------
+.. autoclass:: stepfunctions.steps.service.SqsSendMessageStep

src/stepfunctions/steps/__init__.py

@@ -19,6 +19,6 @@
 from stepfunctions.steps.sagemaker import TrainingStep, TransformStep, ModelStep, EndpointConfigStep, EndpointStep, TuningStep, ProcessingStep
 from stepfunctions.steps.compute import LambdaStep, BatchSubmitJobStep, GlueStartJobRunStep, EcsRunTaskStep
 from stepfunctions.steps.service import DynamoDBGetItemStep, DynamoDBPutItemStep, DynamoDBUpdateItemStep, DynamoDBDeleteItemStep
-from stepfunctions.steps.service import SnsPublishStep, SqsSendMessageStep
 from stepfunctions.steps.service import EmrCreateClusterStep, EmrTerminateClusterStep, EmrAddStepStep, EmrCancelStepStep, EmrSetClusterTerminationProtectionStep, EmrModifyInstanceFleetByNameStep, EmrModifyInstanceGroupByNameStep
-
+from stepfunctions.steps.service import EventBridgePutEventsStep
+from stepfunctions.steps.service import SnsPublishStep, SqsSendMessageStep

src/stepfunctions/steps/service.py

@@ -18,9 +18,11 @@
 from stepfunctions.steps.integration_resources import IntegrationPattern, get_service_integration_arn
 
 DYNAMODB_SERVICE_NAME = "dynamodb"
+ELASTICMAPREDUCE_SERVICE_NAME = "elasticmapreduce"
+EVENTBRIDGE_SERVICE_NAME = "events"
 SNS_SERVICE_NAME = "sns"
 SQS_SERVICE_NAME = "sqs"
-ELASTICMAPREDUCE_SERVICE_NAME = "elasticmapreduce"
+
 
 
 class DynamoDBApi(Enum):
@@ -48,6 +50,10 @@ class ElasticMapReduceApi(Enum):
     ModifyInstanceGroupByName = "modifyInstanceGroupByName"
 
 
+class EventBridgeApi(Enum):
+    PutEvents = "putEvents"
+
+
 class DynamoDBGetItemStep(Task):
     """
     Creates a Task state to get an item from DynamoDB. See `Call DynamoDB APIs with Step Functions <https://docs.aws.amazon.com/step-functions/latest/dg/connect-ddb.html>`_ for more details.
@@ -77,6 +83,46 @@ def __init__(self, state_id, **kwargs):
         super(DynamoDBGetItemStep, self).__init__(state_id, **kwargs)
 
 
+class EventBridgePutEventsStep(Task):
+
+    """
+    Creates a Task to send custom events to Amazon EventBridge. See`Call EventBridge with Step Functions <https://docs.aws.amazon.com/step-functions/latest/dg/connect-eventbridge.html>`_ for more details.
+    """
+
+    def __init__(self, state_id, wait_for_callback=False, **kwargs):
+        """
+        Args:
+            state_id (str): State name whose length **must be** less than or equal to 128 unicode characters. State names **must be** unique within the scope of the whole state machine.
+            comment (str, optional): Human-readable comment or description. (default: None)
+            timeout_seconds (int, optional): Positive integer specifying timeout for the state in seconds. If the state runs longer than the specified timeout, then the interpreter fails the state with a `States.Timeout` Error Name. (default: 60)
+            timeout_seconds_path (str, optional): Path specifying the state's timeout value in seconds from the state input. When resolved, the path must select a field whose value is a positive integer.
+            heartbeat_seconds (int, optional): Positive integer specifying heartbeat timeout for the state in seconds. This value should be lower than the one specified for `timeout_seconds`. If more time than the specified heartbeat elapses between heartbeats from the task, then the interpreter fails the state with a `States.Timeout` Error Name.
+            heartbeat_seconds_path (str, optional): Path specifying the state's heartbeat value in seconds from the state input. When resolved, the path must select a field whose value is a positive integer.
+            input_path (str, optional): Path applied to the state’s raw input to select some or all of it; that selection is used by the state. (default: '$')
+            parameters (dict, optional): The value of this field becomes the effective input for the state.
+            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: '$')
+        """
+
+        if wait_for_callback:
+            """
+            Example resource arn: arn:aws:states:::events:putEvents.waitForTaskToken
+            """
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(EVENTBRIDGE_SERVICE_NAME,
+                                                                       EventBridgeApi.PutEvents,
+                                                                       IntegrationPattern.WaitForTaskToken)
+        else:
+            """
+            Example resource arn: arn:aws:states:::events:putEvents
+            """
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(EVENTBRIDGE_SERVICE_NAME,
+                                                                       EventBridgeApi.PutEvents)
+
+        super(EventBridgePutEventsStep, self).__init__(state_id, **kwargs)
+
+
 class DynamoDBPutItemStep(Task):
 
     """

tests/integ/resources/SageMaker-TrustPolicy.json

@@ -0,0 +1,13 @@
+{
+    "Version": "2012-10-17",
+    "Statement": [
+      {
+        "Sid": "",
+        "Effect": "Allow",
+        "Principal": {
+          "Service": "sagemaker.amazonaws.com"
+        },
+        "Action": "sts:AssumeRole"
+      }
+    ]
+  }

doc/StepFunctionsWorkflowExecutionPolicy.json renamed to tests/integ/resources/StepFunctionsMLWorkflowExecutionFullAccess-Policy.json

@@ -0,0 +1,13 @@
+{
+    "Version": "2012-10-17",
+    "Statement": [
+      {
+        "Sid": "",
+        "Effect": "Allow",
+        "Principal": {
+          "Service": "states.amazonaws.com"
+        },
+        "Action": "sts:AssumeRole"
+      }
+    ]
+  }

tests/integ/resources/StepFunctionsMLWorkflowExecutionFullAccess-TrustPolicy.json

@@ -19,6 +19,7 @@
 from stepfunctions.steps.service import DynamoDBGetItemStep, DynamoDBPutItemStep, DynamoDBUpdateItemStep, DynamoDBDeleteItemStep
 from stepfunctions.steps.service import SnsPublishStep, SqsSendMessageStep
 from stepfunctions.steps.service import EmrCreateClusterStep, EmrTerminateClusterStep, EmrAddStepStep, EmrCancelStepStep, EmrSetClusterTerminationProtectionStep, EmrModifyInstanceFleetByNameStep, EmrModifyInstanceGroupByNameStep
+from stepfunctions.steps.service import EventBridgePutEventsStep
 
 
 @patch.object(boto3.session.Session, 'region_name', 'us-east-1')
@@ -98,6 +99,70 @@ def test_sqs_send_message_step_creation():
         'End': True
     }
 
+@patch.object(boto3.session.Session, 'region_name', 'us-east-1')
+def test_eventbridge_put_events_step_creation():
+    step = EventBridgePutEventsStep('Send to EventBridge', parameters={
+        "Entries": [
+            {
+                "Detail": {
+                    "Message": "MyMessage"
+                },
+                "DetailType": "MyDetailType",
+                "EventBusName": "MyEventBus",
+                "Source": "my.source"
+            }
+        ]
+    })
+
+    assert step.to_dict() == {
+        "Type": "Task",
+        "Resource": 'arn:aws:states:::events:putEvents',
+        "Parameters": {
+            "Entries": [
+                {
+                    "Detail": {
+                        "Message": "MyMessage"
+                    },
+                    "DetailType": "MyDetailType",
+                    "EventBusName": "MyEventBus",
+                    "Source": "my.source"
+                }
+            ]
+        },
+        "End": True
+    }
+
+    step = EventBridgePutEventsStep('Send to EventBridge', wait_for_callback=True, parameters={
+        "Entries": [
+            {
+                "Detail": {
+                    "Message.$": "$.MyMessage"
+                },
+                "DetailType": "MyDetailType",
+                "EventBusName": "MyEventBus",
+                "Source": "my.source"
+            }
+        ]
+    })
+
+    assert step.to_dict() == {
+        "Type": "Task",
+        "Resource": "arn:aws:states:::events:putEvents.waitForTaskToken",
+        "Parameters": {
+            "Entries": [
+                {
+                    "Detail": {
+                        "Message.$": "$.MyMessage"
+                    },
+                    "DetailType": "MyDetailType",
+                    "EventBusName": "MyEventBus",
+                    "Source": "my.source"
+                }
+            ]
+        },
+        "End": True
+    }
+
 
 @patch.object(boto3.session.Session, 'region_name', 'us-east-1')
 def test_dynamodb_get_item_step_creation():

tests/unit/test_service_steps.py

@@ -9,7 +9,7 @@ envlist = python3.6
 skip_missing_interpreters = False
 
 [testenv]
-install_command = pip install {opts} {packages} {env:PWD}[test]
+deps = .[test]
 passenv =
     AWS_ACCESS_KEY_ID
     AWS_SECRET_ACCESS_KEY