Merge pull request #827 from pyathena-dev/feat/644-benchmark-arm64

Allow Graviton benchmark hosts with an architecture parameter

Author: laughingman7743

benchmarks/README.md

@@ -92,6 +92,8 @@ There is no inbound security-group rule or SSH key; use Session Manager.
 Internet access permits package installation and calls to AWS APIs without a NAT gateway.
 
 The defaults are Amazon Linux 2023 x86_64, `r7i.2xlarge` (8 vCPUs, 64 GiB), and an encrypted 100 GiB gp3 root volume.
+Set `Architecture=arm64` with an `r7g` or `r8g` instance type to run on AWS Graviton; the template rejects an instance type that does not match the architecture.
+Measurements depend on the CPU architecture; `environment.json` records it, and runs on different architectures are not directly comparable.
 Instance size and disk size are parameters, so a later run can deliberately test another memory budget.
 The AMI parameter resolves the current AL2023 image at deployment; record the resulting AMI when comparing environments.
 Source bucket access is read-only and restricted to the configured prefix.

benchmarks/cloudformation/benchmark.yaml

@@ -19,10 +19,15 @@ Parameters:
     Type: String
     AllowedPattern: '[0-9a-f]{40}'
     Description: Full commit SHA containing the benchmark implementation.
+  Architecture:
+    Type: String
+    Default: x86_64
+    AllowedValues: [x86_64, arm64]
+    Description: CPU architecture of the benchmark hosts; it must match the instance type.
   InstanceType:
     Type: String
     Default: r7i.2xlarge
-    AllowedValues: [r7i.large, r7i.xlarge, r7i.2xlarge, r7i.4xlarge, r7i.8xlarge]
+    AllowedValues: [r7i.large, r7i.xlarge, r7i.2xlarge, r7i.4xlarge, r7i.8xlarge, r7g.large, r7g.xlarge, r7g.2xlarge, r7g.4xlarge, r7g.8xlarge, r8g.large, r8g.xlarge, r8g.2xlarge, r8g.4xlarge, r8g.8xlarge]
   VolumeSize:
     Type: Number
     Default: 100
@@ -31,6 +36,11 @@ Parameters:
   ImageId:
     Type: AWS::SSM::Parameter::Value<AWS::EC2::Image::Id>
     Default: /aws/service/ami-amazon-linux-latest/al2023-ami-kernel-default-x86_64
+    Description: Image for x86_64 hosts.
+  Arm64ImageId:
+    Type: AWS::SSM::Parameter::Value<AWS::EC2::Image::Id>
+    Default: /aws/service/ami-amazon-linux-latest/al2023-ami-kernel-default-arm64
+    Description: Image for arm64 hosts.
   UvVersion:
     Type: String
     Default: '0.12.5'
@@ -60,6 +70,21 @@ Parameters:
     Default: pypi_file_downloads
     AllowedPattern: '[a-z][a-z0-9_]*'
 
+Rules:
+  X86InstanceType:
+    RuleCondition: !Equals [!Ref Architecture, x86_64]
+    Assertions:
+      - Assert: !Contains [[r7i.large, r7i.xlarge, r7i.2xlarge, r7i.4xlarge, r7i.8xlarge], !Ref InstanceType]
+        AssertDescription: Use an r7i instance type with the x86_64 architecture.
+  Arm64InstanceType:
+    RuleCondition: !Equals [!Ref Architecture, arm64]
+    Assertions:
+      - Assert: !Contains [[r7g.large, r7g.xlarge, r7g.2xlarge, r7g.4xlarge, r7g.8xlarge, r8g.large, r8g.xlarge, r8g.2xlarge, r8g.4xlarge, r8g.8xlarge], !Ref InstanceType]
+        AssertDescription: Use an r7g or r8g instance type with the arm64 architecture.
+
+Conditions:
+  IsArm64: !Equals [!Ref Architecture, arm64]
+
 Resources:
   Vpc:
     Type: AWS::EC2::VPC
@@ -205,7 +230,7 @@ Resources:
         MetadataOptions:
           HttpTokens: required
           HttpEndpoint: enabled
-        ImageId: !Ref ImageId
+        ImageId: !If [IsArm64, !Ref Arm64ImageId, !Ref ImageId]
         InstanceType: !Ref InstanceType
         IamInstanceProfile:
           Arn: !GetAtt InstanceProfile.Arn

benchmarks/tests/test_cloudformation.py

@@ -47,3 +47,24 @@ def test_fleet_size_controls_identical_hosts_and_signals():
         template["Resources"]["LaunchTemplate"]["Properties"]["LaunchTemplateData"]["UserData"]
     )
     assert "AutoScalingGroup" in template["Outputs"]
+
+
+def test_architecture_selects_the_image_and_rules_match_instance_families():
+    template, errors = decode(
+        str(Path(__file__).resolve().parents[1] / "cloudformation/benchmark.yaml")
+    )
+    assert not errors
+    launch = template["Resources"]["LaunchTemplate"]["Properties"]["LaunchTemplateData"]
+    assert launch["ImageId"] == {"Fn::If": ["IsArm64", {"Ref": "Arm64ImageId"}, {"Ref": "ImageId"}]}
+    assert "arm64" in template["Parameters"]["Arm64ImageId"]["Default"]
+    assert "x86_64" in template["Parameters"]["ImageId"]["Default"]
+    allowed = set(template["Parameters"]["InstanceType"]["AllowedValues"])
+    families = {}
+    for name, architecture in (("X86InstanceType", "x86_64"), ("Arm64InstanceType", "arm64")):
+        rule = template["Rules"][name]
+        assert rule["RuleCondition"] == {"Fn::Equals": [{"Ref": "Architecture"}, architecture]}
+        families[architecture] = set(rule["Assertions"][0]["Assert"]["Fn::Contains"][0])
+    assert families["x86_64"] | families["arm64"] == allowed
+    assert not families["x86_64"] & families["arm64"]
+    assert all(t.startswith("r7i.") for t in families["x86_64"])
+    assert all(t.startswith(("r7g.", "r8g.")) for t in families["arm64"])