feat(stepfunctions-tasks): add enableNetworkIsolation property to SageMakerCreateTrainingJobProps

packages/@aws-cdk/aws-stepfunctions-tasks/README.md

@@ -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.

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

@@ -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
@@ -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),

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

@@ -142,6 +142,7 @@ test('create complex training job', () => {
         },
       ],
     },
+    enableNetworkIsolation: true,
     hyperparameters: {
       lr: '0.1',
     },
@@ -221,6 +222,7 @@ test('create complex training job', () => {
           { Name: 'mymetric', Regex: 'regex_pattern' },
         ],
       },
+      EnableNetworkIsolation: true,
       HyperParameters: {
         lr: '0.1',
       },