Skip to content

Commit

Permalink
unit tests...still need integ
Browse files Browse the repository at this point in the history
  • Loading branch information
@comcalvi
comcalvi committed Apr 15, 2024
1 parent 7103fed commit e7f6905
Show file tree
Hide file tree
Showing 2 changed files with 187 additions and 3 deletions.
47 changes: 44 additions & 3 deletions packages/aws-cdk-lib/aws-batch/lib/multinode-job-definition.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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),
}),
},
Expand All @@ -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) });
}

/**
Expand All @@ -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;
}

Expand All @@ -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 [];
}
143 changes: 143 additions & 0 deletions packages/aws-cdk-lib/aws-batch/test/multinode-job-definition.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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();
Expand Down

0 comments on commit e7f6905

Please sign in to comment.