Support Max vCpu in project contexts

packages/cdk/lib/constructs/batch.ts

@@ -1,9 +1,8 @@
 import { Construct, Fn, Names, Stack } from "monocdk";
-import { CfnComputeEnvironment, ComputeEnvironment, ComputeResourceType, IComputeEnvironment, IJobQueue, JobQueue } from "monocdk/aws-batch";
+import { ComputeEnvironment, ComputeResourceType, IComputeEnvironment, IJobQueue, JobQueue } from "monocdk/aws-batch";
 import { CfnLaunchTemplate, InstanceType, IVpc } from "monocdk/aws-ec2";
 import { CfnInstanceProfile, IManagedPolicy, IRole, ManagedPolicy, PolicyDocument, PolicyStatement, Role, ServicePrincipal } from "monocdk/aws-iam";
 import { getInstanceTypesForBatch } from "../util/instance-types";
-import { ComputeType } from "../types";
 
 export interface ComputeOptions {
   /**
@@ -21,7 +20,7 @@ export interface ComputeOptions {
    *
    * @default ON_DEMAND
    */
-  computeType?: ComputeType;
+  computeType?: ComputeResourceType;
   /**
    * The types of EC2 instances that may be launched in the compute environment.
    *
@@ -31,6 +30,15 @@ export interface ComputeOptions {
    */
   instanceTypes?: InstanceType[];
 
+  /**
+   * The maximum number of EC2 vCPUs that a compute environment can reach.
+   *
+   * Each vCPU is equivalent to 1,024 CPU shares.
+   *
+   * @default aws-batch:{@link ComputeResources#maxvCpus}
+   */
+  maxVCpus?: number;
+
   /**
    * The tags to apply to any compute resources
    * @default none
@@ -51,7 +59,7 @@ export interface BatchProps extends ComputeOptions {
   awsPolicyNames?: string[];
 }
 
-const defaultComputeType = ComputeType.ON_DEMAND;
+const defaultComputeType = ComputeResourceType.ON_DEMAND;
 
 export class Batch extends Construct {
   public readonly role: IRole;
@@ -64,13 +72,6 @@ export class Batch extends Construct {
     this.role = this.renderRole(props.computeType, props.awsPolicyNames);
     this.computeEnvironment = this.renderComputeEnvironment(props);
 
-    // TODO: Remove once /~https://github.com/aws/aws-cdk/pull/13591 is merged
-    if (props.computeType == ComputeType.FARGATE || props.computeType == ComputeType.FARGATE_SPOT) {
-      ["AllocationStrategy", "InstanceTypes", "MinvCpus", "InstanceRole"].forEach((property) => {
-        (this.computeEnvironment.node.defaultChild as CfnComputeEnvironment).addPropertyDeletionOverride(`ComputeResources.${property}`);
-      });
-    }
-
     this.jobQueue = new JobQueue(this, "JobQueue", {
       computeEnvironments: [
         {
@@ -81,9 +82,9 @@ export class Batch extends Construct {
     });
   }
 
-  private renderRole(computeType?: ComputeType, awsPolicyNames?: string[]): IRole {
+  private renderRole(computeType?: ComputeResourceType, awsPolicyNames?: string[]): IRole {
     const awsPolicies = awsPolicyNames?.map((policyName) => ManagedPolicy.fromAwsManagedPolicyName(policyName));
-    if (computeType == ComputeType.FARGATE || computeType == ComputeType.FARGATE_SPOT) {
+    if (computeType == ComputeResourceType.FARGATE || computeType == ComputeResourceType.FARGATE_SPOT) {
       return this.renderEcsRole(awsPolicies);
     }
     return this.renderEc2Role(awsPolicies);
@@ -114,12 +115,13 @@ export class Batch extends Construct {
   }
 
   private renderComputeEnvironment(options: ComputeOptions): IComputeEnvironment {
-    options.computeType = options.computeType || defaultComputeType;
-    if (options.computeType == ComputeType.FARGATE || options.computeType == ComputeType.FARGATE_SPOT) {
+    const computeType = options.computeType || defaultComputeType;
+    if (computeType == ComputeResourceType.FARGATE || computeType == ComputeResourceType.FARGATE_SPOT) {
       return new ComputeEnvironment(this, "ComputeEnvironment", {
         computeResources: {
           vpc: options.vpc,
-          type: options.computeType as any as ComputeResourceType,
+          type: computeType,
+          maxvCpus: options.maxVCpus,
         },
       });
     }
@@ -136,13 +138,13 @@ export class Batch extends Construct {
     const instanceProfile = new CfnInstanceProfile(this, "ComputeProfile", {
       roles: [this.role.roleName],
     });
-
     return new ComputeEnvironment(this, "ComputeEnvironment", {
       computeResources: {
         vpc: options.vpc,
-        type: options.computeType as any as ComputeResourceType,
+        type: computeType,
+        maxvCpus: options.maxVCpus,
         instanceRole: instanceProfile.attrArn,
-        instanceTypes: getInstanceTypesForBatch(options.instanceTypes, options.computeType, Stack.of(this).region),
+        instanceTypes: getInstanceTypesForBatch(options.instanceTypes, computeType, Stack.of(this).region),
         launchTemplate: launchTemplate && {
           launchTemplateName: launchTemplate.launchTemplateName!,
         },

packages/cdk/lib/constructs/index.ts

@@ -1,4 +1,3 @@
 export { ApiProxy } from "./api-proxy";
 export { Batch } from "./batch";
-export { ComputeType } from "../types/index";
 export { SecureService } from "./secure-service";

packages/cdk/lib/env/context-app-parameters.ts

@@ -1,4 +1,4 @@
-import { getEnvBoolOrDefault, getEnvString, getEnvStringListOrDefault, getEnvStringOrDefault } from "./";
+import { getEnvNumber, getEnvBoolOrDefault, getEnvString, getEnvStringListOrDefault, getEnvStringOrDefault } from "./";
 import { InstanceType } from "monocdk/aws-ec2";
 import { ConstructNode } from "monocdk";
 import { ServiceContainer } from "../types";
@@ -67,6 +67,10 @@ export class ContextAppParameters {
    */
   public readonly adapterDesignation: string;
 
+  /**
+   * The maximum number of Amazon EC2 vCPUs that an environment can reach.
+   */
+  public readonly maxVCpus?: number;
   /**
    * Property to specify if the compute environment uses On-Demand or Spot compute resources.
    */
@@ -97,6 +101,7 @@ export class ContextAppParameters {
     this.adapterName = getEnvStringOrDefault(node, "ADAPTER_NAME", "wesAdapter")!;
     this.adapterDesignation = getEnvStringOrDefault(node, "ADAPTER_DESIGNATION", "wes")!;
 
+    this.maxVCpus = getEnvNumber(node, "MAX_V_CPUS");
     this.requestSpotInstances = getEnvBoolOrDefault(node, "REQUEST_SPOT_INSTANCES", false)!;
     this.instanceTypes = instanceTypeStrings ? instanceTypeStrings.map((instanceType) => new InstanceType(instanceType.trim())) : undefined;
   }

packages/cdk/lib/env/index.ts

@@ -22,6 +22,11 @@ export const getEnvBoolOrDefault = (node: ConstructNode, key: string, defaultVal
   return value ? valueToBoolean(value) : defaultValue;
 };
 
+export const getEnvNumber = (node: ConstructNode, key: string): Maybe<number> => {
+  const value = node.tryGetContext(key);
+  return value ? Number(value) : undefined;
+};
+
 export const getEnvStringList = (node: ConstructNode, key: string): string[] => {
   return valueToList(getValue(node, key));
 };

packages/cdk/lib/stacks/nested/batch-stack.ts

@@ -2,10 +2,12 @@ import { NestedStack, NestedStackProps } from "monocdk";
 import { InstanceType, IVpc } from "monocdk/aws-ec2";
 import { Construct } from "constructs";
 import { LAUNCH_TEMPLATE } from "../../constants";
-import { Batch, ComputeType } from "../../constructs";
+
+import { Batch } from "../../constructs";
 import { ContextAppParameters } from "../../env";
 import { BucketOperations } from "../../../common/BucketOperations";
 import { IRole } from "monocdk/aws-iam";
+import { ComputeResourceType } from "monocdk/aws-batch";
 
 export interface BatchStackProps extends NestedStackProps {
   /**
@@ -34,18 +36,17 @@ export class BatchStack extends NestedStack {
     super(scope, id, props);
 
     const { vpc, contextParameters, createSpotBatch, createOnDemandBatch } = props;
-
+    const { artifactBucketName, outputBucketName, readBucketArns = [], readWriteBucketArns = [] } = contextParameters;
     if (createSpotBatch) {
-      this.batchSpot = this.renderBatch("TaskBatchSpot", vpc, contextParameters.instanceTypes, ComputeType.SPOT);
+      this.batchSpot = this.renderBatch("TaskBatchSpot", vpc, contextParameters, ComputeResourceType.SPOT);
     }
     if (createOnDemandBatch) {
-      this.batchOnDemand = this.renderBatch("TaskBatch", vpc, contextParameters.instanceTypes, ComputeType.ON_DEMAND);
+      this.batchOnDemand = this.renderBatch("TaskBatch", vpc, contextParameters, ComputeResourceType.ON_DEMAND);
     }
 
-    const artifactBucket = BucketOperations.importBucket(this, "ArtifactBucket", contextParameters.artifactBucketName);
-    const outputBucket = BucketOperations.importBucket(this, "OutputBucket", contextParameters.outputBucketName);
+    const artifactBucket = BucketOperations.importBucket(this, "ArtifactBucket", artifactBucketName);
+    const outputBucket = BucketOperations.importBucket(this, "OutputBucket", outputBucketName);
 
-    const { readBucketArns = [], readWriteBucketArns = [] } = contextParameters;
     readBucketArns.push(artifactBucket.bucketArn);
     readWriteBucketArns.push(outputBucket.bucketArn);
 
@@ -56,11 +57,12 @@ export class BatchStack extends NestedStack {
     }
   }
 
-  private renderBatch(id: string, vpc: IVpc, instanceTypes?: InstanceType[], computeType?: ComputeType): Batch {
+  private renderBatch(id: string, vpc: IVpc, appParams: ContextAppParameters, computeType?: ComputeResourceType): Batch {
     return new Batch(this, id, {
       vpc,
-      instanceTypes,
       computeType,
+      instanceTypes: appParams.instanceTypes,
+      maxVCpus: appParams.maxVCpus,
       launchTemplateData: LAUNCH_TEMPLATE,
       awsPolicyNames: ["AmazonSSMManagedInstanceCore", "CloudWatchAgentServerPolicy"],
       resourceTags: this.nestedStackParent?.tags.tagValues(),

packages/cdk/lib/types/index.ts

@@ -1,5 +1,4 @@
 export type Maybe<T> = T | undefined;
 
-export { ComputeType } from "./compute-type";
 export { EngineOptions } from "./engine-options";
 export { ServiceContainer } from "./service-container";

packages/cdk/lib/util/instance-types.ts

@@ -1,10 +1,10 @@
 import { InstanceType } from "monocdk/aws-ec2";
-import { ComputeType } from "../types/index";
+import { ComputeResourceType } from "monocdk/aws-batch";
 
 // /~https://github.com/aws-samples/aws-genomics-workflows/blob/master/src/templates/gwfcore/gwfcore-batch.template.yaml#L145-L180
 // batch 'optimal' isn't optimal for genomics computation, these types have been tuned over several customer engagements
-export const optimalInstanceTypes: { [key in ComputeType]: string[] } = {
-  [ComputeType.ON_DEMAND]: [
+export const optimalInstanceTypes: { [key in ComputeResourceType]: string[] } = {
+  [ComputeResourceType.ON_DEMAND]: [
     "c5.large",
     "c5.xlarge",
     "c5.2xlarge",
@@ -42,7 +42,7 @@ export const optimalInstanceTypes: { [key in ComputeType]: string[] } = {
     "r5n.2xlarge",
     "r5n.4xlarge",
   ],
-  [ComputeType.SPOT]: [
+  [ComputeResourceType.SPOT]: [
     "c4.large",
     "c4.xlarge",
     "c4.2xlarge",
@@ -93,8 +93,8 @@ export const optimalInstanceTypes: { [key in ComputeType]: string[] } = {
     "r5n.4xlarge",
   ],
   // TODO: Determine optimal instances for fargate when it's supported
-  [ComputeType.FARGATE]: [],
-  [ComputeType.FARGATE_SPOT]: [],
+  [ComputeResourceType.FARGATE]: [],
+  [ComputeResourceType.FARGATE_SPOT]: [],
 };
 
 // These need to be manually updated as regions release new instance types
@@ -174,7 +174,7 @@ const unLaunchedInstanceTypesByRegion: { [key in string]: { [key in string]: boo
   "sa-east-1": { "m5n.large": true, "m5n.xlarge": true, "m5n.2xlarge": true, "m5n.4xlarge": true },
 };
 
-export const getInstanceTypesForBatch = (instanceTypes: InstanceType[] | undefined, computeType: ComputeType, region?: string): InstanceType[] => {
+export const getInstanceTypesForBatch = (instanceTypes: InstanceType[] | undefined, computeType: ComputeResourceType, region?: string): InstanceType[] => {
   if (instanceTypes && instanceTypes.length > 0) {
     return instanceTypes;
   }

packages/cli/internal/pkg/cli/context/context_environment.go

@@ -27,6 +27,7 @@ type contextEnvironment struct {
 	ReadWriteBucketArns  string
 	InstanceTypes        string
 	ResourceType         string
+	MaxVCpus             int
 	RequestSpotInstances bool
 }
 
@@ -51,6 +52,7 @@ func (input contextEnvironment) ToEnvironmentList() []string {
 		"READ_BUCKET_ARNS":             input.ReadBucketArns,
 		"READ_WRITE_BUCKET_ARNS":       input.ReadWriteBucketArns,
 		"BATCH_COMPUTE_INSTANCE_TYPES": input.InstanceTypes,
+		"MAX_V_CPUS":                   strconv.Itoa(input.MaxVCpus),
 		"REQUEST_SPOT_INSTANCES":       strconv.FormatBool(input.RequestSpotInstances),
 	})
 }

packages/cli/internal/pkg/cli/context/info.go

@@ -4,6 +4,7 @@ import "reflect"
 
 type Summary struct {
 	Name          string
+	MaxVCpus      int
 	IsSpot        bool
 	InstanceTypes []string
 }

packages/cli/internal/pkg/cli/context/manager.go

@@ -173,6 +173,7 @@ func (m *Manager) setTaskContext(contextName string) {
 		ReadBucketArns:       strings.Join(m.readBuckets, listDelimiter),
 		ReadWriteBucketArns:  strings.Join(m.readWriteBuckets, listDelimiter),
 		InstanceTypes:        strings.Join(m.contextSpec.InstanceTypes, listDelimiter),
+		MaxVCpus:             m.contextSpec.MaxVCpus,
 		RequestSpotInstances: m.contextSpec.RequestSpotInstances,
 	}
 }
@@ -198,6 +199,7 @@ func (m *Manager) setContextEnv(contextName string) {
 		ReadBucketArns:       strings.Join(m.readBuckets, listDelimiter),
 		ReadWriteBucketArns:  strings.Join(m.readWriteBuckets, listDelimiter),
 		InstanceTypes:        strings.Join(m.contextSpec.InstanceTypes, listDelimiter),
+		MaxVCpus:             m.contextSpec.MaxVCpus,
 		RequestSpotInstances: m.contextSpec.RequestSpotInstances,
 		// TODO: we default to a single engine in a context for now
 		// need to allow for multiple engines in the same context

packages/cli/internal/pkg/cli/context/manager_info.go

@@ -56,6 +56,7 @@ func (m *Manager) buildContextInfo(contextName string) (Detail, error) {
 		Summary: Summary{
 			Name:          contextName,
 			IsSpot:        m.projectSpec.Contexts[contextName].RequestSpotInstances,
+			MaxVCpus:      m.projectSpec.Contexts[contextName].MaxVCpus,
 			InstanceTypes: m.projectSpec.Contexts[contextName].InstanceTypes,
 		},
 		Status:             m.contextStatus,

packages/cli/internal/pkg/cli/context_describe.go

@@ -44,6 +44,7 @@ func (o *describeContextOpts) Execute() (types.Context, error) {
 		Status:               info.Status.ToString(),
 		StatusReason:         info.StatusReason,
 		InstanceTypes:        buildInstanceTypes(info.InstanceTypes),
+		MaxVCpus:             info.MaxVCpus,
 		RequestSpotInstances: info.IsSpot,
 		Output:               types.OutputLocation{Url: info.BucketLocation},
 	}, nil

packages/cli/internal/pkg/cli/describe_output_test.go

@@ -30,7 +30,8 @@ func TestDescribeOutput(t *testing.T) {
 		},
 		"Context": {
 			output:              types.Context{},
-			expectedDescription: "Output of the command has following format:\nCONTEXT: Name RequestSpotInstances Status StatusReason\nINSTANCETYPE: Value\nOUTPUTLOCATION: Url\n",
+			expectedDescription: "Output of the command has following format:\nCONTEXT: MaxVCpus Name RequestSpotInstances Status" +
+				" StatusReason\nINSTANCETYPE: Value\nOUTPUTLOCATION: Url\n",
 		},
 	}
 

packages/cli/internal/pkg/cli/spec/context.go

@@ -8,5 +8,18 @@ type Engine struct {
 type Context struct {
 	InstanceTypes        []string `yaml:"instanceTypes,omitempty"`
 	RequestSpotInstances bool     `yaml:"requestSpotInstances,omitempty"`
+	MaxVCpus             int      `yaml:"maxVCpus,omitempty"`
 	Engines              []Engine `yaml:"engines"`
 }
+
+func (context *Context) UnmarshalYAML(unmarshal func(interface{}) error) error {
+
+	type defValContext Context
+	defaults := defValContext{MaxVCpus: 256}
+	if err := unmarshal(&defaults); err != nil {
+		return err
+	}
+
+	*context = Context(defaults)
+	return nil
+}

packages/cli/internal/pkg/cli/spec/project_schema.json

@@ -67,6 +67,10 @@
             "requestSpotInstances":{
               "type":"boolean"
             },
+            "maxVCpus":{
+              "type":"integer",
+              "minimum": 1
+            },
             "instanceTypes":{
               "type":[
                 "array",