Added lambda step example (#2847)

* Added lambda step example * Updated notes on IAM policy for Lambda * Updated IAM notes * Added pip install * Updated pip install * Added iam helper function * Updated role description * Fixed typos * Updated md cells * updated formatting Co-authored-by: kthadaka <kthadaka!>
aws · Aug 6, 2021 · 50c356f · 50c356f
1 parent 006f229
commit 50c356f
Show file tree

Hide file tree

Showing 2 changed files with 933 additions and 0 deletions.
diff --git a/sagemaker-pipelines/tabular/lambda-step/iam_helper.py b/sagemaker-pipelines/tabular/lambda-step/iam_helper.py
@@ -0,0 +1,42 @@
+import boto3
+import json
+
+iam = boto3.client('iam')
+
+def create_lambda_role(role_name):
+    try:
+        response = iam.create_role(
+            RoleName = role_name,
+            AssumeRolePolicyDocument = json.dumps({
+                "Version": "2012-10-17",
+                "Statement": [
+                    {
+                        "Effect": "Allow",
+                        "Principal": {
+                            "Service": "lambda.amazonaws.com"
+                        },
+                        "Action": "sts:AssumeRole"
+                    }
+                ]
+            }),
+            Description='Role for Lambda to call SageMaker functions'
+        )
+
+        role_arn = response['Role']['Arn']
+
+        response = iam.attach_role_policy(
+            RoleName=role_name,
+            PolicyArn='arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
+        )
+
+        response = iam.attach_role_policy(
+            PolicyArn='arn:aws:iam::aws:policy/AmazonSageMakerFullAccess',
+            RoleName=role_name
+        )
+
+        return role_arn
+
+    except iam.exceptions.EntityAlreadyExistsException:
+        print(f'Using ARN from existing role: {role_name}')
+        response = iam.get_role(RoleName=role_name)
+        return response['Role']['Arn']