From e7f6905da153428ac8a2bdcdd766d614a879627a Mon Sep 17 00:00:00 2001 From: Calvin Combs Date: Mon, 15 Apr 2024 16:43:15 -0700 Subject: [PATCH] unit tests...still need integ --- .../aws-batch/lib/multinode-job-definition.ts | 47 +++++- .../test/multinode-job-definition.test.ts | 143 ++++++++++++++++++ 2 files changed, 187 insertions(+), 3 deletions(-) diff --git a/packages/aws-cdk-lib/aws-batch/lib/multinode-job-definition.ts b/packages/aws-cdk-lib/aws-batch/lib/multinode-job-definition.ts index a103c8819f1fd..49261c037390d 100644 --- a/packages/aws-cdk-lib/aws-batch/lib/multinode-job-definition.ts +++ b/packages/aws-cdk-lib/aws-batch/lib/multinode-job-definition.ts @@ -70,8 +70,11 @@ export interface MultiNodeContainer { * The index of the last node to run this container. * * The container is run on all nodes in the range [startNode, endNode] (inclusive) + * + * If this value is not specified, you must specify `numNodes`. + * @default - Batch assigns the highest possible node index. */ - readonly endNode: number; + readonly endNode?: number; /** * The container that this node range will run @@ -108,6 +111,14 @@ export interface MultiNodeJobDefinitionProps extends JobDefinitionProps { */ readonly mainNode?: number; + /** + * The total number of nodes used in this job. + * **Only specify if there is at least one container for which you have not specified `endNode`.** + * + * @default - computed based on the containers passed in. + */ + readonly numNodes?: number; + /** * Whether to propogate tags from the JobDefinition * to the ECS task that Batch spawns @@ -165,14 +176,14 @@ export class MultiNodeJobDefinition extends JobDefinitionBase implements IMultiN mainNode: this.mainNode ?? 0, nodeRangeProperties: Lazy.any({ produce: () => this.containers.map((container) => ({ - targetNodes: container.startNode + ':' + container.endNode, + targetNodes: container.startNode + ':' + (container.endNode ?? ''), container: { ...container.container._renderContainerDefinition(), instanceType: this._instanceType?.toString(), }, })), }), - numNodes: Lazy.number({ + numNodes: props?.numNodes ?? Lazy.number({ produce: () => computeNumNodes(this.containers), }), }, @@ -186,6 +197,7 @@ export class MultiNodeJobDefinition extends JobDefinitionBase implements IMultiN this.jobDefinitionName = this.getResourceNameAttribute(resource.ref); this.node.addValidation({ validate: () => validateContainers(this.containers) }); + this.node.addValidation({ validate: () => validateNumNodesPropConflicts(this.containers, this.node.id, props?.numNodes) }); } /** @@ -209,6 +221,9 @@ function computeNumNodes(containers: MultiNodeContainer[]) { let result = 0; for (const container of containers) { + if (!container.endNode) { + throw new Error(`The multinode container '${container.container.node.id}' does not specify an end node, and its Job Definition does not specify 'numNodes'. Please either specify 'numNodes' or specify 'endNode' for every container in this Job Definition.`); + } result += container.endNode - container.startNode + 1; } @@ -218,3 +233,29 @@ function computeNumNodes(containers: MultiNodeContainer[]) { function validateContainers(containers: MultiNodeContainer[]): string[] { return containers.length === 0 ? ['multinode job has no containers!'] : []; } + +function validateNumNodesPropConflicts(containers: MultiNodeContainer[], jobDefnName: string, numNodes?: number): string[] { + let allContainersSpecifyEndNode = true; + let noEndNodeContainers = []; + for (const container of containers ?? []) { + if (!container.endNode) { + allContainersSpecifyEndNode = false; + noEndNodeContainers.push(container.container.node.id); + } + } + + if (numNodes && allContainersSpecifyEndNode) { + return [`All containers of Multinode Job Definition '${jobDefnName}' specify 'endNode', but the job definition specifies 'numNodes'! Do not specify 'endNode' for every container with 'numNodes', the CDK will compute the correct value for 'numNodes' if all containers have 'endNode' specified.`]; + } + + const returnValue = []; + + if (!numNodes && !allContainersSpecifyEndNode) { + for (const containerId of noEndNodeContainers) { + returnValue.push(`The multinode container '${containerId}' does not specify an end node, and its Job Definition does not + specify 'numNodes'. Please either specify 'numNodes' or specify 'endNode' for every container in this Job Definition.`); + } + } + + return []; +} diff --git a/packages/aws-cdk-lib/aws-batch/test/multinode-job-definition.test.ts b/packages/aws-cdk-lib/aws-batch/test/multinode-job-definition.test.ts index 0e5fb4cd3395c..8833868245c32 100644 --- a/packages/aws-cdk-lib/aws-batch/test/multinode-job-definition.test.ts +++ b/packages/aws-cdk-lib/aws-batch/test/multinode-job-definition.test.ts @@ -93,6 +93,87 @@ test('MultiNodeJobDefinition respects instanceType', () => { }); }); +test('MultiNodeJobDefinition respects numNodes', () => { + // GIVEN + const stack = new Stack(); + + // WHEN + new MultiNodeJobDefinition(stack, 'ECSJobDefn', { + containers: [{ + container: new EcsEc2ContainerDefinition(stack, 'MultinodeContainer', { + cpu: 256, + memory: Size.mebibytes(2048), + image: ecs.ContainerImage.fromRegistry('amazon/amazon-ecs-sample'), + }), + startNode: 0, + }], + numNodes: 10, + instanceType: InstanceType.of(InstanceClass.R4, InstanceSize.LARGE), + }); + + // THEN + Template.fromStack(stack).hasResourceProperties('AWS::Batch::JobDefinition', { + NodeProperties: { + NodeRangeProperties: [{ + Container: { + }, + TargetNodes: '0:', + }], + NumNodes: 10, + }, + PlatformCapabilities: [Compatibility.EC2], + }); +}); + +test('MultiNodeJobDefinition respects numNodes with two containers', () => { + // GIVEN + const stack = new Stack(); + + // WHEN + const jobDefn = new MultiNodeJobDefinition(stack, 'ECSJobDefn', { + containers: [{ + container: new EcsEc2ContainerDefinition(stack, 'MultinodeContainer', { + cpu: 256, + memory: Size.mebibytes(2048), + image: ecs.ContainerImage.fromRegistry('amazon/amazon-ecs-sample'), + }), + startNode: 0, + endNode: 5, + }], + numNodes: 10, + instanceType: InstanceType.of(InstanceClass.R4, InstanceSize.LARGE), + }); + + jobDefn.addContainer({ + container: new EcsEc2ContainerDefinition(stack, 'MultinodeContainer2', { + cpu: 256, + memory: Size.mebibytes(2048), + image: ecs.ContainerImage.fromRegistry('amazon/amazon-ecs-sample'), + }), + startNode: 6, + }); + + // THEN + Template.fromStack(stack).hasResourceProperties('AWS::Batch::JobDefinition', { + NodeProperties: { + NodeRangeProperties: [ + { + Container: { + }, + TargetNodes: '0:5', + }, + { + Container: { + }, + TargetNodes: '6:', + }, + ], + NumNodes: 10, + }, + PlatformCapabilities: [Compatibility.EC2], + }); +}); + test('MultiNodeJobDefinition one container', () => { // GIVEN const stack = new Stack(); @@ -191,6 +272,68 @@ test('multinode job requires at least one container', () => { expect(() => Template.fromStack(stack)).toThrow(/multinode job has no containers!/); }); +test('multinode job does not allow specifying all containers with `endNode` if `numNodes` is specified', () => { + // GIVEN + const stack = new Stack(); + + // WHEN + const jobDefn = new MultiNodeJobDefinition(stack, 'ECSJobDefn', { + numNodes: 10, + containers: [{ + container: new EcsEc2ContainerDefinition(stack, 'MultinodeContainer1', { + cpu: 256, + memory: Size.mebibytes(2048), + image: ecs.ContainerImage.fromRegistry('amazon/amazon-ecs-sample'), + }), + startNode: 0, + endNode: 5, + }], + }); + jobDefn.addContainer({ + container: new EcsEc2ContainerDefinition(stack, 'MultinodeContainer2', { + cpu: 256, + memory: Size.mebibytes(2048), + image: ecs.ContainerImage.fromRegistry('amazon/amazon-ecs-sample'), + }), + startNode: 6, + endNode: 10, + }); + + // THEN + expect(() => Template.fromStack(stack)).toThrow(/All containers of Multinode Job Definition 'ECSJobDefn' specify 'endNode', but the job definition specifies 'numNodes'!/); +}); + +test('MultiNodeJobDefinition throws some dumb error somewhere', () => { + // GIVEN + const stack = new Stack(); + + // WHEN + const jobDefn = new MultiNodeJobDefinition(stack, 'ECSJobDefn', { + containers: [{ + container: new EcsEc2ContainerDefinition(stack, 'MultinodeContainer', { + cpu: 256, + memory: Size.mebibytes(2048), + image: ecs.ContainerImage.fromRegistry('amazon/amazon-ecs-sample'), + }), + startNode: 0, + }], + instanceType: InstanceType.of(InstanceClass.R4, InstanceSize.LARGE), + }); + + jobDefn.addContainer({ + container: new EcsEc2ContainerDefinition(stack, 'MultinodeContainer2', { + cpu: 256, + memory: Size.mebibytes(2048), + image: ecs.ContainerImage.fromRegistry('amazon/amazon-ecs-sample'), + }), + startNode: 6, + endNode: 10, + }); + + // THEN + expect(() => Template.fromStack(stack)).toThrow(/The multinode container 'MultinodeContainer' does not specify an end node, and its Job Definition does not specify 'numNodes'/); +}); + test('multinode job returns a dummy instance type when accessing `instanceType`', () => { // GIVEN const stack = new Stack();