fix(batch): remove default optimal for arm based instance types and add error checks

packages/aws-cdk-lib/aws-batch/lib/managed-compute-environment.ts

@@ -5,7 +5,7 @@ import * as ec2 from '../../aws-ec2';
 import * as eks from '../../aws-eks';
 import * as iam from '../../aws-iam';
 import { IRole } from '../../aws-iam';
-import { ArnFormat, Duration, ITaggable, Lazy, Resource, Stack, TagManager, TagType } from '../../core';
+import { Annotations, ArnFormat, Duration, ITaggable, Lazy, Resource, Stack, TagManager, TagType } from '../../core';
 
 /**
  * Represents a Managed ComputeEnvironment. Batch will provision EC2 Instances to
@@ -685,7 +685,7 @@ export class ManagedEc2EcsComputeEnvironment extends ManagedComputeEnvironmentBa
         minvCpus: this.minvCpus,
         instanceRole: this.instanceProfile.attrArn, // this is not a typo; this property actually takes a profile, not a standard role
         instanceTypes: Lazy.list({
-          produce: () => renderInstances(this.instanceTypes, this.instanceClasses, props.useOptimalInstanceClasses),
+          produce: () => renderInstances(this, this.instanceTypes, this.instanceClasses, props.useOptimalInstanceClasses),
         }),
         type: this.spot ? 'SPOT' : 'EC2',
         spotIamFleetRole: this.spotFleetRole?.roleArn,
@@ -712,14 +712,18 @@ export class ManagedEc2EcsComputeEnvironment extends ManagedComputeEnvironmentBa
       resourceName: this.physicalName,
     });
 
-    this.node.addValidation({ validate: () => validateInstances(this.instanceTypes, this.instanceClasses, props.useOptimalInstanceClasses) });
+    this.node.addValidation({
+      validate: () => validateInstances.call(this, this.instanceTypes, this.instanceClasses, props.useOptimalInstanceClasses),
+    });
   }
 
   public addInstanceType(instanceType: ec2.InstanceType): void {
+    // TODO add check
     this.instanceTypes.push(instanceType);
   }
 
   public addInstanceClass(instanceClass: ec2.InstanceClass): void {
+    // TODO add check
     this.instanceClasses.push(instanceClass);
   }
 }
@@ -1034,7 +1038,7 @@ export class ManagedEc2EksComputeEnvironment extends ManagedComputeEnvironmentBa
         ...baseManagedResourceProperties(this, subnetIds).computeResources as CfnComputeEnvironment.ComputeResourcesProperty,
         minvCpus: this.minvCpus,
         instanceRole: this.instanceProfile.attrArn, // this is not a typo; this property actually takes a profile, not a standard role
-        instanceTypes: Lazy.list({ produce: () => renderInstances(this.instanceTypes, this.instanceClasses, props.useOptimalInstanceClasses) }),
+        instanceTypes: Lazy.list({ produce: () => renderInstances(this, this.instanceTypes, this.instanceClasses, props.useOptimalInstanceClasses) }),
         type: this.spot ? 'SPOT' : 'EC2',
         allocationStrategy: this.allocationStrategy,
         bidPercentage: this.spotBidPercentage,
@@ -1059,14 +1063,18 @@ export class ManagedEc2EksComputeEnvironment extends ManagedComputeEnvironmentBa
       resourceName: this.physicalName,
     });
 
-    this.node.addValidation({ validate: () => validateInstances(this.instanceTypes, this.instanceClasses, props.useOptimalInstanceClasses) });
+    this.node.addValidation({
+      validate: () => validateInstances.call(this, this.instanceTypes, this.instanceClasses, props.useOptimalInstanceClasses),
+    });  
   }
 
   public addInstanceType(instanceType: ec2.InstanceType): void {
+    // TODO: implement check if it is arm
     this.instanceTypes.push(instanceType);
   }
 
   public addInstanceClass(instanceClass: ec2.InstanceClass): void {
+    // TODO: implement check if it is arm
     this.instanceClasses.push(instanceClass);
   }
 }
@@ -1131,15 +1139,39 @@ export class FargateComputeEnvironment extends ManagedComputeEnvironmentBase imp
   }
 }
 
-function renderInstances(types?: ec2.InstanceType[], classes?: ec2.InstanceClass[], useOptimalInstanceClasses?: boolean): string[] {
+function renderInstances(scope: Construct, types?: ec2.InstanceType[], classes?: ec2.InstanceClass[], useOptimalInstanceClasses?: boolean): string[] {
   const instances = [];
-
+  let hasArmInstances = false;
+  let hasX86Instances = false;
+
   for (const instanceType of types ?? []) {
     instances.push(instanceType.toString());
+    if (instanceType.architecture === ec2.InstanceArchitecture.ARM_64) {
+      hasArmInstances = true;
+    } else {
+      hasX86Instances = true;
+    }
   }
+
   for (const instanceClass of classes ?? []) {
     instances.push(instanceClass);
+    const instanceType = new ec2.InstanceType(`${instanceClass.toString()}.large`);
+    if (instanceType.architecture === ec2.InstanceArchitecture.ARM_64) {
+      hasArmInstances = true;
+    } else {
+      hasX86Instances = true;
+    }
   }
+
+  if (hasArmInstances && hasX86Instances) {
+    throw new Error('Cannot mix ARM and x86 instance types or classes');
+  }
+
+  if (hasArmInstances && (useOptimalInstanceClasses || useOptimalInstanceClasses === undefined)) {
+    Annotations.of(scope).addWarningV2('@aws-cdk/aws-batch:optimalNotSupportedWithARM', '\'optimal\' instance types are not supported with ARM instance types or classes, setting useOptimalInstanceClasses to false');
+    useOptimalInstanceClasses = false;
+  }
+
   if (useOptimalInstanceClasses || useOptimalInstanceClasses === undefined) {
     instances.push('optimal');
   }
@@ -1181,8 +1213,8 @@ function determineAllocationStrategy(id: string, allocationStrategy?: Allocation
   return result;
 }
 
-function validateInstances(types?: ec2.InstanceType[], classes?: ec2.InstanceClass[], useOptimalInstanceClasses?: boolean): string[] {
-  if (renderInstances(types, classes, useOptimalInstanceClasses).length === 0) {
+function validateInstances(this: Construct, types?: ec2.InstanceType[], classes?: ec2.InstanceClass[], useOptimalInstanceClasses?: boolean): string[] {
+  if (renderInstances(this, types, classes, useOptimalInstanceClasses).length === 0) {
     return ["Specifies 'useOptimalInstanceClasses: false' without specifying any instance types or classes"];
   }