fix(batch): allow endNode to be optional for multinode containers

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 [];
+}

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();