feature: Apache Airflow integration for SageMaker Processing Jobs #1620

guptamyr · 2020-06-24T07:36:03Z

Issue #, if available:

Description of changes:
This change enables an Apache Airflow operator for SageMaker Processing Jobs through a processing_config that can be passed to an Airflow operator.
Testing done:
Added a unit test for this config. Also tested this with local changes to the Airflow repository that add the Operator and was able to create/run a DAG for ProcessingJobs similar to the examples here.

Merge Checklist

Put an x in the boxes that apply. You can also fill these out after creating the PR. If you're unsure about any of them, don't hesitate to ask. We're here to help! This is simply a reminder of what we are going to look for before merging your pull request.

General

I have read the CONTRIBUTING doc
I used the commit message format described in CONTRIBUTING
I have passed the region in to all S3 and STS clients that I've initialized as part of this change.
I have updated any necessary documentation, including READMEs and API docs (if appropriate)

Tests

I have added tests that prove my fix is effective or that my feature works (if appropriate)
[] I have checked that my tests are not configured for a specific region or account (if appropriate)
[] I have used unique_name_from_base to create resource names in integ tests (if appropriate)

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

sagemaker-bot · 2020-06-24T07:49:13Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: b6abe7c
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-24T08:03:30Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: b6abe7c
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-24T08:12:17Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-pr
Commit ID: b6abe7c
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-24T08:20:22Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: b6abe7c
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

metrizable

nice work! thanks for the contribution. mostly style comments.

metrizable · 2020-06-24T15:20:44Z

src/sagemaker/processing.py

+        config = {}
+        if container_arguments is not None:
+            config["ContainerArguments"] = container_arguments
+        if container_entrypoint is not None:
+            config["ContainerEntrypoint"] = container_entrypoint
+        config["ImageUri"] = image_uri


it's more pythonic if, since "ImageUri" is always a key of config, just initialize with it:

config = {"ImageUri": image_uri} if container_arguments ...

Fixed in the next revision.

metrizable · 2020-06-24T15:21:53Z

src/sagemaker/processing.py

+        config = {}
+        if kms_key_id is not None:
+            config["KmsKeyId"] = kms_key_id
+        config["Outputs"] = outputs


since "Outputs" is always present as a key, initialize config with it:

config = {"Outputs": outputs} if kms_key_id ...

Fixed in the next revision.

metrizable · 2020-06-24T15:23:50Z

src/sagemaker/processing.py

+        stopping_condition = {}
+        stopping_condition["MaxRuntimeInSeconds"] = max_runtime_in_seconds
+        return stopping_condition


just:

return {"MaxRuntimeInSeconds": max_runtime_in_seconds}

Fixed in the next revision.

metrizable · 2020-06-24T15:26:02Z

src/sagemaker/processing.py

@@ -739,6 +739,66 @@ def stop(self):
        """Stops the processing job."""
        self.sagemaker_session.stop_processing_job(self.name)

+    @staticmethod
+    def _prepare_app_specification(container_arguments, container_entrypoint, image_uri):


probably don't want the underscore to precede the name? although there is no practical limitation to accessibility, by convention of the preceding underscore, it typically indicates a "protected" method.

there are also some material ramifications: from the public names section here:

If __all__ is not defined, the set of public names includes all names found in the module’s namespace which do not begin with an underscore character ('_')

Removed the underscore in the next revision.

metrizable · 2020-06-24T15:41:15Z

src/sagemaker/processing.py

+
+    @staticmethod
+    def _prepare_output_config(kms_key_id, outputs):
+        """


nit-pep257: include summary line

(i know the preceding underscore indicates the method is not "public", but is used outside of this class/subclasses, and probably shouldn't have the underscore)

Fixed the underscore and added summary lines, descriptions of args and type for all of these functions.

metrizable · 2020-06-24T15:43:36Z

src/sagemaker/processing.py

+    def _prepare_output_config(kms_key_id, outputs):
+        """
+        Args:
+            kms_key_id:


use google docstring conventions and include brief descriptions of args. we typically also indicate type in parentheses.

Fixed the underscore and added summary lines, descriptions of args and type for all of these functions.

metrizable · 2020-06-24T15:51:34Z

src/sagemaker/workflow/airflow.py

+            else utils.base_name_from_image(processor.image_uri)
+        )
+
+    input_dict = []


maybe call it input_dicts, since it's actually a list of dicts?

Fixed in the next revision.

metrizable · 2020-06-24T15:53:46Z

src/sagemaker/workflow/airflow.py

+    for i in inputs:
+        input_dict.append(i._to_request_dict())


more pythonic (and performs better) to use list comprehension:

input_dicts = [i._to_request_dict() for i in inputs]

Fixed in the next revision.

metrizable · 2020-06-24T15:55:15Z

src/sagemaker/workflow/airflow.py

+    output_dict = []
+    for output in outputs:
+        output_dict.append(output._to_request_dict())


same:

output_dicts = [output._to_request_dict() for output in outputs]

and it's very similar to previous transformation on inputs. helper function to encapsulate same logic and invoke function?

Added a helper function in the next revision.

sagemaker-bot · 2020-06-25T00:18:20Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: aff3deb
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-25T00:39:57Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: aff3deb
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-25T00:53:49Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 95bd791
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-25T00:56:09Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: aff3deb
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-25T01:15:39Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: 95bd791
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-25T01:32:00Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: 95bd791
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-25T02:56:37Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-pr
Commit ID: 95bd791
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

metrizable · 2020-06-25T06:12:06Z

src/sagemaker/processing.py

+        processing_resources = {}
+        cluster_config = {}
+        if volume_kms_key_id is not None:
+            cluster_config["VolumeKmsKeyId"] = volume_kms_key_id
+        cluster_config["InstanceCount"] = instance_count
+        cluster_config["InstanceType"] = instance_type
+        cluster_config["VolumeSizeInGB"] = volume_size_in_gb
+        processing_resources["ClusterConfig"] = cluster_config
+        return processing_resources


nit-style: initialize dict with values?

cluster_config = { "InstanceCount": instance_count, ... } ... return {"ClusterConfig": cluster_config}

}

Fixed in the next revision.

metrizable · 2020-06-25T06:14:04Z

src/sagemaker/processing.py

+
+    @staticmethod
+    def prepare_output_config(kms_key_id, outputs):
+        """Prepares a dict that represents a ProcessingOutputConfig,


nit-pep257: single line summary, additional information after blank line in body.

Fixed in the next revision. (moved the extra information to describe what the function returns.

metrizable · 2020-06-25T06:16:44Z

src/sagemaker/processing.py

@@ -739,6 +739,94 @@ def stop(self):
        """Stops the processing job."""
        self.sagemaker_session.stop_processing_job(self.name)

+    @staticmethod
+    def prepare_app_specification(container_arguments, container_entrypoint, image_uri):
+        """Prepares a dict that represents a ProcessingJob's AppSpecification


Fixed in the next revision. (moved the extra information to describe what the function returns.)

sagemaker-bot · 2020-06-26T00:01:53Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 35ced0a
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T00:25:07Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: 35ced0a
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T00:38:39Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 190f09f
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T00:39:27Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: 35ced0a
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T00:40:04Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 85de82a
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T00:53:17Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: 190f09f
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T00:56:30Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: 85de82a
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T01:08:03Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: 190f09f
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T01:27:44Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: 85de82a
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T04:20:24Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-pr
Commit ID: 85de82a
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

metrizable

ship it

src/sagemaker/processing.py

src/sagemaker/workflow/airflow.py

metrizable

ship it

sagemaker-bot · 2020-06-26T04:37:54Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 120da98
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T04:37:56Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 120da98
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T04:38:12Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 334acea
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T04:38:35Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 334acea
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T04:56:28Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: 120da98
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T04:56:53Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: 334acea
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T04:57:12Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: 334acea
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T04:57:28Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-notebook-tests
Commit ID: 120da98
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T05:21:28Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 334acea
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T05:25:51Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: 120da98
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T05:26:18Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: 334acea
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T05:26:20Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: 120da98
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T05:26:45Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-local-mode-tests
Commit ID: 334acea
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T06:13:29Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-unit-tests
Commit ID: 334acea
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T06:18:46Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-pr
Commit ID: 334acea
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2020-06-26T08:16:46Z

AWS CodeBuild CI Report

CodeBuild project: sagemaker-python-sdk-pr
Commit ID: 334acea
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

guptamyr marked this pull request as ready for review June 24, 2020 15:09

metrizable suggested changes Jun 24, 2020

View reviewed changes

guptamyr force-pushed the airflow_processing branch from aff3deb to 95bd791 Compare June 25, 2020 00:48

metrizable suggested changes Jun 25, 2020

View reviewed changes

guptamyr force-pushed the airflow_processing branch from 95bd791 to 35ced0a Compare June 25, 2020 23:56

feature: Apache Airflow integration for SageMaker Processing Jobs

190f09f

guptamyr force-pushed the airflow_processing branch from 35ced0a to 190f09f Compare June 26, 2020 00:25

Merge branch 'master' into airflow_processing

85de82a

metrizable previously approved these changes Jun 26, 2020

View reviewed changes

src/sagemaker/processing.py Outdated Show resolved Hide resolved

src/sagemaker/workflow/airflow.py Outdated Show resolved Hide resolved

change: in docstring indent additional lines

120da98

metrizable dismissed their stale review via 120da98 June 26, 2020 04:27

change: in docstring indent additional lines

334acea

metrizable approved these changes Jun 26, 2020

View reviewed changes

metrizable merged commit 52d0189 into aws:master Jun 26, 2020

guptamyr mentioned this pull request Jun 30, 2020

Adding new SageMaker operator for ProcessingJobs apache/airflow#9594

Merged

6 tasks

feature: Apache Airflow integration for SageMaker Processing Jobs #1620

feature: Apache Airflow integration for SageMaker Processing Jobs #1620

Uh oh!

Conversation

guptamyr commented Jun 24, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge Checklist

General

Tests

Uh oh!

sagemaker-bot commented Jun 24, 2020

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Jun 24, 2020

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Jun 24, 2020

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Jun 24, 2020

AWS CodeBuild CI Report

Uh oh!

metrizable left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

metrizable Jun 24, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sagemaker-bot commented Jun 25, 2020

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Jun 25, 2020

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Jun 25, 2020

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Jun 25, 2020

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Jun 25, 2020

AWS CodeBuild CI Report

Uh oh!

sagemaker-bot commented Jun 25, 2020

AWS CodeBuild CI Report

guptamyr commented Jun 24, 2020 •

edited

Loading

metrizable Jun 24, 2020 •

edited

Loading