fix(stepfunctions-tasks): batch properties accept JsonPaths

aws · Mar 24, 2022 · 5690ffd · 5690ffd
1 parent 8c60d4d
commit 5690ffd
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/batch/submit-job.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/batch/submit-job.ts
@@ -1,7 +1,7 @@
 import * as ec2 from '@aws-cdk/aws-ec2';
 import * as iam from '@aws-cdk/aws-iam';
 import * as sfn from '@aws-cdk/aws-stepfunctions';
-import { Size, Stack, withResolved } from '@aws-cdk/core';
+import { Size, Stack, Token, withResolved } from '@aws-cdk/core';
 import { Construct } from 'constructs';
 import { integrationResourceArn, validatePatternSupported } from '../private/task-utils';
 
@@ -296,23 +296,23 @@ export class BatchSubmitJob extends sfn.TaskStateBase {
       resources.push(
         {
           Type: 'GPU',
-          Value: `${containerOverrides.gpuCount}`,
+          Value: Token.isUnresolved(containerOverrides.gpuCount) ? containerOverrides.gpuCount : `${containerOverrides.gpuCount}`,
         },
       );
     }
     if (containerOverrides.memory) {
       resources.push(
         {
           Type: 'MEMORY',
-          Value: `${containerOverrides.memory.toMebibytes()}`,
+          Value: Token.isUnresolved(containerOverrides.memory) ? containerOverrides.memory : `${containerOverrides.memory.toMebibytes()}`,
         },
       );
     }
     if (containerOverrides.vcpus) {
       resources.push(
         {
           Type: 'VCPU',
-          Value: `${containerOverrides.vcpus}`,
+          Value: Token.isUnresolved(containerOverrides.vcpus) ? containerOverrides.vcpus : `${containerOverrides.vcpus}`,
         },
       );
     }

diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/test/batch/submit-job.test.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/test/batch/submit-job.test.ts
@@ -79,7 +79,7 @@ test('Task with all the parameters', () => {
       instanceType: new ec2.InstanceType('MULTI'),
       memory: cdk.Size.mebibytes(1024),
       gpuCount: 1,
-      vcpus: 10,
+      vcpus: sfn.JsonPath.numberAt('$.taskParameters.cpu'),
     },
     dependsOn: [{ jobId: '1234', type: 'some_type' }],
     payload: sfn.TaskInput.fromObject({
@@ -115,7 +115,7 @@ test('Task with all the parameters', () => {
         Command: ['sudo', 'rm'],
         Environment: [{ Name: 'key', Value: 'value' }],
         InstanceType: 'MULTI',
-        ResourceRequirements: [{ Type: 'GPU', Value: '1' }, { Type: 'MEMORY', Value: '1024' }, { Type: 'VCPU', Value: '10' }],
+        ResourceRequirements: [{ Type: 'GPU', Value: '1' }, { Type: 'MEMORY', Value: '1024' }, { 'Type': 'VCPU', 'Value.$': '$.taskParameters.cpu' }],
       },
       DependsOn: [{ JobId: '1234', Type: 'some_type' }],
       Parameters: { 'foo.$': '$.bar' },