fix(aws-stepfunctions): refactor sagemaker tasks and fix default role issue

packages/@aws-cdk/aws-stepfunctions-tasks/lib/sagemaker-task-base-types.ts

@@ -184,20 +184,15 @@ export interface ResourceConfig {
  * @experimental
  */
 export interface VpcConfig {
-    /**
-     * VPC security groups.
-     */
-    readonly securityGroups: ec2.ISecurityGroup[];
-
     /**
      * VPC id
      */
-    readonly vpc: ec2.Vpc;
+    readonly vpc: ec2.IVpc;
 
     /**
      * VPC subnets.
      */
-    readonly subnets: ec2.ISubnet[];
+    readonly subnets?: ec2.ISubnet[];
 }
 
 /**

packages/@aws-cdk/aws-stepfunctions-tasks/lib/sagemaker-train-task.ts

@@ -1,7 +1,7 @@
 import ec2 = require('@aws-cdk/aws-ec2');
 import iam = require('@aws-cdk/aws-iam');
 import sfn = require('@aws-cdk/aws-stepfunctions');
-import { Construct, Duration, Stack } from '@aws-cdk/cdk';
+import { Construct, Duration, Lazy, Stack } from '@aws-cdk/cdk';
 import { AlgorithmSpecification, Channel, InputMode, OutputDataConfig, ResourceConfig,
          S3DataType, StoppingCondition, VpcConfig,  } from './sagemaker-task-base-types';
 
@@ -45,7 +45,7 @@ export interface SagemakerTrainProps {
     /**
      * Tags to be applied to the train job.
      */
-    readonly tags?: {[key: string]: any};
+    readonly tags?: {[key: string]: string};
 
     /**
      * Identifies the Amazon S3 location where you want Amazon SageMaker to save the results of model training.
@@ -76,7 +76,7 @@ export class SagemakerTrainTask implements ec2.IConnectable, sfn.IStepFunctionsT
     /**
      * Allows specify security group connections for instances of this fleet.
      */
-    public readonly connections: ec2.Connections = new ec2.Connections();
+    public readonly connections: ec2.Connections;
 
     /**
      * The execution role for the Sagemaker training job.
@@ -105,6 +105,11 @@ export class SagemakerTrainTask implements ec2.IConnectable, sfn.IStepFunctionsT
      */
     private readonly stoppingCondition: StoppingCondition;
 
+    private readonly vpc: ec2.IVpc;
+    private readonly securityGroup: ec2.ISecurityGroup;
+    private readonly securityGroups: ec2.ISecurityGroup[] = [];
+    private readonly subnets: string[];
+
     constructor(scope: Construct, private readonly props: SagemakerTrainProps) {
 
         // set the default resource config if not defined.
@@ -120,13 +125,18 @@ export class SagemakerTrainTask implements ec2.IConnectable, sfn.IStepFunctionsT
         };
 
         // set the sagemaker role or create new one
-        this.role = props.role || new iam.Role(scope, 'SagemakerRole', {
+        this.role = props.role || new iam.Role(scope, 'SagemakerTrainRole', {
             assumedBy: new iam.ServicePrincipal('sagemaker.amazonaws.com'),
             managedPolicies: [
                 iam.ManagedPolicy.fromAwsManagedPolicyName('AmazonSageMakerFullAccess')
             ]
         });
 
+        // check that either algorithm name or image is defined
+        if ((!props.algorithmSpecification.algorithmName) && (!props.algorithmSpecification.trainingImage)) {
+            throw new Error("Must define either an algorithm name or training image URI in the algorithm specification");
+        }
+
         // set the input mode to 'File' if not defined
         this.algorithmSpecification = ( props.algorithmSpecification.trainingInputMode ) ?
             ( props.algorithmSpecification ) :
@@ -143,11 +153,27 @@ export class SagemakerTrainTask implements ec2.IConnectable, sfn.IStepFunctionsT
         });
 
         // add the security groups to the connections object
-        if (this.props.vpcConfig) {
-            this.props.vpcConfig.securityGroups.forEach(sg => this.connections.addSecurityGroup(sg));
+        if (props.vpcConfig) {
+            this.vpc = props.vpcConfig.vpc;
+            this.securityGroup = new ec2.SecurityGroup(scope, 'TrainJobSecurityGroup', {
+                vpc: this.vpc
+            });
+            this.connections = new ec2.Connections({ securityGroups: [this.securityGroup] });
+            this.securityGroups.push(this.securityGroup);
+            this.subnets = (props.vpcConfig.subnets) ? (props.vpcConfig.subnets.map(s => (s.subnetId))) : this.vpc.selectSubnets().subnetIds;
         }
     }
 
+    /**
+     * Add the security group to all instances via the launch configuration
+     * security groups array.
+     *
+     * @param securityGroup: The security group to add
+     */
+    public addSecurityGroup(securityGroup: ec2.ISecurityGroup): void {
+        this.securityGroups.push(securityGroup);
+    }
+
     public bind(task: sfn.Task): sfn.StepFunctionsTaskConfig  {
         return {
           resourceArn: 'arn:aws:states:::sagemaker:createTrainingJob' + (this.props.synchronous ? '.sync' : ''),
@@ -244,8 +270,8 @@ export class SagemakerTrainTask implements ec2.IConnectable, sfn.IStepFunctionsT
 
     private renderVpcConfig(config: VpcConfig | undefined): {[key: string]: any} {
         return (config) ? { VpcConfig: {
-            SecurityGroupIds: config.securityGroups.map(sg => ( sg.securityGroupId )),
-            Subnets: config.subnets.map(subnet => ( subnet.subnetId )),
+            SecurityGroupIds: Lazy.listValue({ produce: () => (this.securityGroups.map(sg => (sg.securityGroupId))) }),
+            Subnets: this.subnets,
         }} : {};
     }
 

packages/@aws-cdk/aws-stepfunctions-tasks/lib/sagemaker-transform-task.ts

@@ -32,7 +32,7 @@ export interface SagemakerTransformProps {
     /**
      * Environment variables to set in the Docker container.
      */
-    readonly environment?: {[key: string]: any};
+    readonly environment?: {[key: string]: string};
 
     /**
      * Maximum number of parallel requests that can be sent to each instance in a transform job.
@@ -52,7 +52,7 @@ export interface SagemakerTransformProps {
     /**
      * Tags to be applied to the train job.
      */
-    readonly tags?: {[key: string]: any};
+    readonly tags?: {[key: string]: string};
 
     /**
      * Dataset to be transformed and the Amazon S3 location where it is stored.
@@ -97,7 +97,7 @@ export class SagemakerTransformTask implements sfn.IStepFunctionsTask {
     constructor(scope: Construct, private readonly props: SagemakerTransformProps) {
 
         // set the sagemaker role or create new one
-        this.role = props.role || new iam.Role(scope, 'SagemakerRole', {
+        this.role = props.role || new iam.Role(scope, 'SagemakerTransformRole', {
             assumedBy: new iam.ServicePrincipal('sagemaker.amazonaws.com'),
             managedPolicies: [
                 iam.ManagedPolicy.fromAwsManagedPolicyName('AmazonSageMakerFullAccess')

packages/@aws-cdk/aws-stepfunctions-tasks/test/sagemaker-training-job.test.ts

@@ -11,7 +11,7 @@ let stack: cdk.Stack;
 beforeEach(() => {
     // GIVEN
     stack = new cdk.Stack();
-  });
+});
 
 test('create basic training job', () => {
     // WHEN
@@ -64,7 +64,7 @@ test('create basic training job', () => {
             InstanceType: 'ml.m4.xlarge',
             VolumeSizeInGB: 10
         },
-        RoleArn: { "Fn::GetAtt": [ "SagemakerRole5FDB64E1", "Arn" ] },
+        RoleArn: { "Fn::GetAtt": [ "SagemakerTrainRoleCBF0A724", "Arn" ] },
         StoppingCondition: {
             MaxRuntimeInSeconds: 3600
         },
@@ -87,7 +87,7 @@ test('create complex training job', () => {
         ],
     });
 
-    const task = new sfn.Task(stack, 'TrainSagemaker', { task: new tasks.SagemakerTrainTask(stack, {
+    const trainTask = new tasks.SagemakerTrainTask(stack, {
         trainingJobName: "MyTrainJob",
         synchronous: true,
         role,
@@ -148,9 +148,10 @@ test('create complex training job', () => {
         vpcConfig: {
             vpc,
             subnets: vpc.privateSubnets,
-            securityGroups: [ securityGroup ]
         }
-    })});
+    });
+    trainTask.addSecurityGroup(securityGroup);
+    const task = new sfn.Task(stack, 'TrainSagemaker', { task: trainTask });
 
     // THEN
     expect(stack.resolve(task.toStateJson())).toEqual({
@@ -213,7 +214,10 @@ test('create complex training job', () => {
             { Key: "Project", Value: "MyProject" }
         ],
         VpcConfig: {
-            SecurityGroupIds: [ { "Fn::GetAtt": [ "SecurityGroupDD263621", "GroupId" ] } ],
+            SecurityGroupIds: [
+                { "Fn::GetAtt": [ "TrainJobSecurityGroupBECEDCDC", "GroupId" ] },
+                { "Fn::GetAtt": [ "SecurityGroupDD263621", "GroupId" ] },
+            ],
             Subnets: [
                 { Ref: "VPCPrivateSubnet1Subnet8BCA10E0" },
                 { Ref: "VPCPrivateSubnet2SubnetCFCDAA7A" },
@@ -300,3 +304,26 @@ test('pass param to training job', () => {
       },
     });
 });
+
+test('Cannot create a SageMaker train task with both algorithm name and image name missing', () => {
+
+    expect(() => new tasks.SagemakerTrainTask(stack, {
+        trainingJobName: 'myTrainJob',
+        algorithmSpecification: {},
+        inputDataConfig: [
+            {
+                channelName: 'train',
+                dataSource: {
+                    s3DataSource: {
+                        s3DataType: tasks.S3DataType.S3_PREFIX,
+                        s3Uri: sfn.Data.stringAt('$.S3Bucket')
+                    }
+                }
+            }
+        ],
+        outputDataConfig: {
+            s3OutputPath: 's3://mybucket/myoutputpath'
+        },
+    }))
+      .toThrowError(/Must define either an algorithm name or training image URI in the algorithm specification/);
+});

packages/@aws-cdk/aws-stepfunctions-tasks/test/sagemaker-transform-job.test.ts

@@ -19,7 +19,7 @@ beforeEach(() => {
             iam.ManagedPolicy.fromAwsManagedPolicyName('AmazonSageMakerFullAccess')
         ],
     });
-  });
+});
 
 test('create basic transform job', () => {
     // WHEN