fix memory token issues

aws · Mar 25, 2022 · d075b81 · d075b81
1 parent 5690ffd
commit d075b81
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 5 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
@@ -197,6 +197,13 @@ export class BatchSubmitJob extends sfn.TaskStateBase {
  }
  });
 
+ // validate memory
+ props.containerOverrides?.memory !== undefined && withResolved(props.containerOverrides.memory.toMebibytes(), (memory) => {
+ if (memory < 4) {
+ throw new Error('memory must be at least 4MiB');
+ }
+ });
+
  // This is required since environment variables must not start with AWS_BATCH;
  // this naming convention is reserved for variables that are set by the AWS Batch service.
  if (props.containerOverrides?.environment) {
@@ -296,23 +303,23 @@ export class BatchSubmitJob extends sfn.TaskStateBase {
  resources.push(
  {
  Type: 'GPU',
- Value: Token.isUnresolved(containerOverrides.gpuCount) ? containerOverrides.gpuCount : `${containerOverrides.gpuCount}`,
+ Value: tokenOrString(containerOverrides.gpuCount),
  },
  );
  }
  if (containerOverrides.memory) {
  resources.push(
  {
  Type: 'MEMORY',
- Value: Token.isUnresolved(containerOverrides.memory) ? containerOverrides.memory : `${containerOverrides.memory.toMebibytes()}`,
+ Value: tokenOrString(containerOverrides.memory.toMebibytes()),
  },
  );
  }
  if (containerOverrides.vcpus) {
  resources.push(
  {
  Type: 'VCPU',
- Value: Token.isUnresolved(containerOverrides.vcpus) ? containerOverrides.vcpus : `${containerOverrides.vcpus}`,
+ Value: tokenOrString(containerOverrides.vcpus),
  },
  );
  }
@@ -325,3 +332,7 @@ export class BatchSubmitJob extends sfn.TaskStateBase {
  };
  }
 }
+
+function tokenOrString(value: any) {
+ return Token.isUnresolved(value) ? value : value.toString();
+}
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: sfn.JsonPath.numberAt('$.taskParameters.cpu'),
+ vcpus: 10,
  },
  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.$': '$.taskParameters.cpu' }],
+ ResourceRequirements: [{ Type: 'GPU', Value: '1' }, { Type: 'MEMORY', Value: '1024' }, { Type: 'VCPU', Value: '10' }],
  },
  DependsOn: [{ JobId: '1234', Type: 'some_type' }],
  Parameters: { 'foo.$': '$.bar' },
@@ -134,6 +134,11 @@ test('supports tokens', () => {
  arraySize: sfn.JsonPath.numberAt('$.arraySize'),
  timeout: cdk.Duration.seconds(sfn.JsonPath.numberAt('$.timeout')),
  attempts: sfn.JsonPath.numberAt('$.attempts'),
+ containerOverrides: {
+ gpuCount: sfn.JsonPath.numberAt('$.gpuCount'),
+ memory: cdk.Size.mebibytes(sfn.JsonPath.numberAt('$.memory')),
+ vcpus: sfn.JsonPath.numberAt('$.vcpus'),
+ },
  });
 
  // THEN
@@ -165,6 +170,22 @@ test('supports tokens', () => {
  'Timeout': {
  'AttemptDurationSeconds.$': '$.timeout',
  },
+ 'ContainerOverrides': {
+ ResourceRequirements: [
+ {
+ 'Type': 'GPU',
+ 'Value.$': '$.gpuCount',
+ },
+ {
+ 'Type': 'MEMORY',
+ 'Value.$': '$.memory',
+ },
+ {
+ 'Type': 'VCPU',
+ 'Value.$': '$.vcpus',
+ },
+ ],
+ },
  },
  });
 });