Skip to content

Commit

Permalink
feat(stepfunctions-tasks): task constructs for creating and transform…
Browse files Browse the repository at this point in the history
…ing SageMaker jobs (#8391)

replacement for the current implementation of `SageMaker` service
integration and state level properties are merged and represented
as a construct.

The previous implementation that implemented `IStepFunctionsTask` has
been removed. The previously existing classes were directly converted to constructs
as these were marked **experimental** and still require further iterations
as they are not backed by a SageMaker L2 (does not exist yet).

In the interest of pragmatism, I decided to move them to leverage the newer pattern
so we can deprecate the `Task` construct.

Note that I have left the unit and integration tests verbatim. The integration test
requires some additional steps as there are pre-requisites to running a training job
such as creating and configuring input data that are not currently included.

BREAKING CHANGE: constructs for `SageMakerCreateTrainingJob` and
`SageMakerCreateTransformJob` replace previous implementation that
implemented `IStepFunctionsTask`.
* **stepfunctions-tasks:** `volumeSizeInGB` property in `ResourceConfig` for
SageMaker tasks are now type `core.Size`
* **stepfunctions-tasks:** `maxPayload` property in `SagemakerTransformProps`
is now type `core.Size`
* **stepfunctions-tasks:** `volumeKmsKeyId` property in `SageMakerCreateTrainingJob` is now `volumeEncryptionKey`

----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
  • Loading branch information
shivlaks authored Jun 9, 2020
1 parent ae90cba commit 480d4c0
Show file tree
Hide file tree
Showing 10 changed files with 387 additions and 373 deletions.
94 changes: 44 additions & 50 deletions packages/@aws-cdk/aws-stepfunctions-tasks/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -617,37 +617,33 @@ Step Functions supports [AWS SageMaker](https://docs.aws.amazon.com/step-functio
You can call the [`CreateTrainingJob`](https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateTrainingJob.html) API from a `Task` state.

```ts
new sfn.Task(stack, 'TrainSagemaker', {
task: new tasks.SagemakerTrainTask({
trainingJobName: sfn.Data.stringAt('$.JobName'),
role,
algorithmSpecification: {
algorithmName: 'BlazingText',
trainingInputMode: tasks.InputMode.FILE,
},
inputDataConfig: [
{
channelName: 'train',
dataSource: {
s3DataSource: {
s3DataType: tasks.S3DataType.S3_PREFIX,
s3Location: tasks.S3Location.fromJsonExpression('$.S3Bucket'),
},
},
new sfn.SagemakerTrainTask(this, 'TrainSagemaker', {
trainingJobName: sfn.Data.stringAt('$.JobName'),
role,
algorithmSpecification: {
algorithmName: 'BlazingText',
trainingInputMode: tasks.InputMode.FILE,
},
inputDataConfig: [{
channelName: 'train',
dataSource: {
s3DataSource: {
s3DataType: tasks.S3DataType.S3_PREFIX,
s3Location: tasks.S3Location.fromJsonExpression('$.S3Bucket'),
},
],
outputDataConfig: {
s3OutputLocation: tasks.S3Location.fromBucket(s3.Bucket.fromBucketName(stack, 'Bucket', 'mybucket'), 'myoutputpath'),
},
resourceConfig: {
instanceCount: 1,
instanceType: ec2.InstanceType.of(ec2.InstanceClass.P3, ec2.InstanceSize.XLARGE2),
volumeSizeInGB: 50,
},
stoppingCondition: {
maxRuntime: cdk.Duration.hours(1),
},
}),
}],
outputDataConfig: {
s3OutputLocation: tasks.S3Location.fromBucket(s3.Bucket.fromBucketName(stack, 'Bucket', 'mybucket'), 'myoutputpath'),
},
resourceConfig: {
instanceCount: 1,
instanceType: ec2.InstanceType.of(ec2.InstanceClass.P3, ec2.InstanceSize.XLARGE2),
volumeSize: cdk.Size.gibibytes(50),
},
stoppingCondition: {
maxRuntime: cdk.Duration.hours(1),
},
});
```

Expand All @@ -656,29 +652,27 @@ new sfn.Task(stack, 'TrainSagemaker', {
You can call the [`CreateTransformJob`](https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateTransformJob.html) API from a `Task` state.

```ts
const transformJob = new tasks.SagemakerTransformTask(
transformJobName: "MyTransformJob",
modelName: "MyModelName",
role,
transformInput: {
transformDataSource: {
s3DataSource: {
s3Uri: 's3://inputbucket/train',
s3DataType: S3DataType.S3Prefix,
}
}
},
transformOutput: {
s3OutputPath: 's3://outputbucket/TransformJobOutputPath',
},
transformResources: {
instanceCount: 1,
instanceType: ec2.InstanceType.of(ec2.InstanceClass.M4, ec2.InstanceSize.XLarge),
new sfn.SagemakerTransformTask(this, 'Batch Inference', {
transformJobName: 'MyTransformJob',
modelName: 'MyModelName',
role,
transformInput: {
transformDataSource: {
s3DataSource: {
s3Uri: 's3://inputbucket/train',
s3DataType: S3DataType.S3Prefix,
}
}
},
transformOutput: {
s3OutputPath: 's3://outputbucket/TransformJobOutputPath',
},
transformResources: {
instanceCount: 1,
instanceType: ec2.InstanceType.of(ec2.InstanceClass.M4, ec2.InstanceSize.XLarge),
}
});

const task = new sfn.Task(this, 'Batch Inference', {
task: transformJob
});
```

## SNS
Expand Down
6 changes: 3 additions & 3 deletions packages/@aws-cdk/aws-stepfunctions-tasks/lib/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,9 @@ export * from './sqs/send-to-queue';
export * from './sqs/send-message';
export * from './ecs/run-ecs-ec2-task';
export * from './ecs/run-ecs-fargate-task';
export * from './sagemaker/sagemaker-task-base-types';
export * from './sagemaker/sagemaker-train-task';
export * from './sagemaker/sagemaker-transform-task';
export * from './sagemaker/base-types';
export * from './sagemaker/create-training-job';
export * from './sagemaker/create-transform-job';
export * from './start-execution';
export * from './stepfunctions/start-execution';
export * from './evaluate-expression';
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,13 @@ import * as iam from '@aws-cdk/aws-iam';
import * as kms from '@aws-cdk/aws-kms';
import * as s3 from '@aws-cdk/aws-s3';
import * as sfn from '@aws-cdk/aws-stepfunctions';
import { Construct, Duration } from '@aws-cdk/core';
import { Construct, Duration, Size } from '@aws-cdk/core';

/**
* Task to train a machine learning model using Amazon SageMaker
* @experimental
*/
export interface ISageMakerTask extends sfn.IStepFunctionsTask, iam.IGrantable {}
export interface ISageMakerTask extends iam.IGrantable {}

/**
* Specify the training algorithm and algorithm-specific metadata
Expand Down Expand Up @@ -230,7 +230,7 @@ export interface ResourceConfig {
*
* @default 10 GB EBS volume.
*/
readonly volumeSizeInGB: number;
readonly volumeSize: Size;
}

/**
Expand Down Expand Up @@ -622,7 +622,7 @@ export interface TransformResources {
*
* @default - None
*/
readonly volumeKmsKeyId?: kms.Key;
readonly volumeEncryptionKey?: kms.IKey;
}

/**
Expand Down
Loading

0 comments on commit 480d4c0

Please sign in to comment.