From a1016284092c51fb3b5d8fb71bfee345e61ef213 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pablo=20Barb=C3=A1chano?= <pablob@amazon.com>
Date: Wed, 20 Nov 2024 13:35:36 +0100
Subject: [PATCH] ci: infer instance architecture from a heuristic
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This helps the script work for instances it doesn't know about, which is
helpful while onboarding new instances.

Only Graviton metal instances are aarch64 at the moment.

Fixes: c33bc6c6f32a886cf186d4c4130cb0f6bd66bdca
Signed-off-by: Pablo Barbáchano <pablob@amazon.com>
---
 .buildkite/common.py                | 27 ++++++++++++++++-----------
 .buildkite/pipeline_cpu_template.py |  2 +-
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/.buildkite/common.py b/.buildkite/common.py
index 196435bb2e0..a8913d5d6ef 100644
--- a/.buildkite/common.py
+++ b/.buildkite/common.py
@@ -13,14 +13,14 @@
 import subprocess
 from pathlib import Path
 
-DEFAULT_INSTANCES = {
-    "c5n.metal": "x86_64",  # Intel Skylake
-    "m5n.metal": "x86_64",  # Intel Cascade Lake
-    "m6i.metal": "x86_64",  # Intel Icelake
-    "m6a.metal": "x86_64",  # AMD Milan
-    "m6g.metal": "aarch64",  # Graviton2
-    "m7g.metal": "aarch64",  # Graviton3
-}
+DEFAULT_INSTANCES = [
+    "c5n.metal",  # Intel Skylake
+    "m5n.metal",  # Intel Cascade Lake
+    "m6i.metal",  # Intel Icelake
+    "m6a.metal",  # AMD Milan
+    "m6g.metal",  # Graviton2
+    "m7g.metal",  # Graviton3
+]
 
 DEFAULT_PLATFORMS = [
     ("al2", "linux_5.10"),
@@ -28,6 +28,11 @@
 ]
 
 
+def get_arch_for_instance(instance):
+    """Return instance architecture"""
+    return "x86_64" if instance[2] != "g" else "aarch64"
+
+
 def overlay_dict(base: dict, update: dict):
     """Overlay a dict over a base one"""
     base = base.copy()
@@ -145,7 +150,7 @@ def __call__(self, parser, namespace, value, option_string=None):
     "--instances",
     required=False,
     nargs="+",
-    default=DEFAULT_INSTANCES.keys(),
+    default=DEFAULT_INSTANCES,
 )
 COMMON_PARSER.add_argument(
     "--platforms",
@@ -288,7 +293,7 @@ def _adapt_group(self, group):
             step["command"] = prepend + step["command"]
             if self.shared_build is not None:
                 step["depends_on"] = self.build_key(
-                    DEFAULT_INSTANCES[step["agents"]["instance"]]
+                    get_arch_for_instance(step["agents"]["instance"])
                 )
         return group
 
@@ -323,7 +328,7 @@ def build_group_per_arch(self, label, *args, **kwargs):
         if set_key:
             for step in grp["steps"]:
                 step["key"] = self.build_key(
-                    DEFAULT_INSTANCES[step["agents"]["instance"]]
+                    get_arch_for_instance(step["agents"]["instance"])
                 )
         return self.add_step(grp, depends_on_build=depends_on_build)
 
diff --git a/.buildkite/pipeline_cpu_template.py b/.buildkite/pipeline_cpu_template.py
index 72a67ea5e7d..c8c5f75fa09 100755
--- a/.buildkite/pipeline_cpu_template.py
+++ b/.buildkite/pipeline_cpu_template.py
@@ -34,7 +34,7 @@ class BkStep(str, Enum):
             "tools/devtool -y test --no-build -- -m no_block_pr integration_tests/functional/test_cpu_template_helper.py -k test_guest_cpu_config_change",
         ],
         BkStep.LABEL: "🖐️ fingerprint",
-        "instances": DEFAULT_INSTANCES.keys(),
+        "instances": DEFAULT_INSTANCES,
         "platforms": DEFAULT_PLATFORMS,
     },
     "cpuid_wrmsr": {