Skip to content

Commit

Permalink
feat(stepfunctions-tasks): add enableNetworkIsolation property to `…
Browse files Browse the repository at this point in the history
…SageMakerCreateTrainingJobProps` (#16792)

Closes #16779.

This PR adds an option of `EnableNetworkIsolation` to `SageMakerCreateTrainingJob` class and enables a user to use AWS Marketplace.

[Doc](https://docs.aws.amazon.com/sagemaker/latest/dg/mkt-algo-model-internet-free.html) says a user must turn the option to True when they attempt to use a machine learning model from AWS Marketplace. 

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
  • Loading branch information
TakaakiFuruse authored Oct 18, 2021
1 parent 999e999 commit 69ac520
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 0 deletions.
4 changes: 4 additions & 0 deletions packages/@aws-cdk/aws-stepfunctions-tasks/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -969,6 +969,10 @@ disable this behavior.

Step Functions supports [AWS SageMaker](https://docs.aws.amazon.com/step-functions/latest/dg/connect-sagemaker.html) through the service integration pattern.

If your training job or model uses resources from AWS Marketplace,
[network isolation is required](https://docs.aws.amazon.com/sagemaker/latest/dg/mkt-algo-model-internet-free.html).
To do so, set the `enableNetworkIsolation` property to `true` for `SageMakerCreateModel` or `SageMakerCreateTrainingJob`.

### Create Training Job

You can call the [`CreateTrainingJob`](https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateTrainingJob.html) API from a `Task` state.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,13 @@ export interface SageMakerCreateTrainingJobProps extends sfn.TaskStateBaseProps
*/
readonly algorithmSpecification: AlgorithmSpecification;

/**
* Isolates the training container. No inbound or outbound network calls can be made to or from the training container.
*
* @default false
*/
readonly enableNetworkIsolation?: boolean;

/**
* Algorithm-specific parameters that influence the quality of the model. Set hyperparameters before you start the learning process.
* For a list of hyperparameters provided by Amazon SageMaker
Expand Down Expand Up @@ -217,6 +224,7 @@ export class SageMakerCreateTrainingJob extends sfn.TaskStateBase implements iam
private renderParameters(): { [key: string]: any } {
return {
TrainingJobName: this.props.trainingJobName,
EnableNetworkIsolation: this.props.enableNetworkIsolation,
RoleArn: this._role!.roleArn,
...this.renderAlgorithmSpecification(this.algorithmSpecification),
...this.renderInputDataConfig(this.inputDataConfig),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,7 @@ test('create complex training job', () => {
},
],
},
enableNetworkIsolation: true,
hyperparameters: {
lr: '0.1',
},
Expand Down Expand Up @@ -221,6 +222,7 @@ test('create complex training job', () => {
{ Name: 'mymetric', Regex: 'regex_pattern' },
],
},
EnableNetworkIsolation: true,
HyperParameters: {
lr: '0.1',
},
Expand Down

0 comments on commit 69ac520

Please sign in to comment.