[CI] This PR enhances testing of the CI procedures on both v6e and v7x.

Signed-off-by: dennis yeh <[email protected]>

Author: dennisYehCienet

.buildkite/features/Collective_Communication_Matmul.yml

@@ -1,44 +1,48 @@
 # Collective Communication Matmul
 # kernel support matrix
 steps:
-  - label: "Correctness tests for Collective Communication Matmul"
-    key: "Collective_Communication_Matmul_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Correctness tests for Collective Communication Matmul"
+    key: "${TEST_KEY_PREFIX}Collective_Communication_Matmul_CorrectnessTest"
     soft_fail: true
     agents:
-      queue: tpu_v6e_queue
+      queue: "${TPU_SMALL_CORE_QUEUE}"
+    env:
+      IS_FOR_V7X: "${IS_FOR_V7X}"
     commands:
       - .buildkite/scripts/run_in_docker.sh python3 -m pytest -s -v /workspace/tpu_inference/tests/kernels/collectives/all_gather_matmul_kernel_test.py
-  - label: "Record correctness test result for Collective Communication Matmul"
-    key: "record_Collective_Communication_Matmul_CorrectnessTest"
-    depends_on: "Collective_Communication_Matmul_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Record correctness test result for Collective Communication Matmul"
+    key: "${TEST_KEY_PREFIX}record_Collective_Communication_Matmul_CorrectnessTest"
+    depends_on: "${TEST_KEY_PREFIX}Collective_Communication_Matmul_CorrectnessTest"
     env:
-      CI_TARGET: "Collective Communication Matmul"
+      CI_TARGET: "${TEST_LABEL_PREFIX}Collective Communication Matmul"
       CI_STAGE: "CorrectnessTest"
       CI_CATEGORY: "kernel support matrix"
     agents:
       queue: cpu
     commands:
       - |
-        .buildkite/scripts/record_step_result.sh Collective_Communication_Matmul_CorrectnessTest
+        .buildkite/scripts/record_step_result.sh ${TEST_KEY_PREFIX}Collective_Communication_Matmul_CorrectnessTest
 
-  - label: "Performance tests for Collective Communication Matmul"
-    key: "Collective_Communication_Matmul_PerformanceTest"
-    depends_on: "record_Collective_Communication_Matmul_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Performance tests for Collective Communication Matmul"
+    key: "${TEST_KEY_PREFIX}Collective_Communication_Matmul_PerformanceTest"
+    depends_on: "${TEST_KEY_PREFIX}record_Collective_Communication_Matmul_CorrectnessTest"
     soft_fail: true
     agents:
-      queue: tpu_v6e_queue
+      queue: "${TPU_SMALL_CORE_QUEUE}"
+    env:
+      IS_FOR_V7X: "${IS_FOR_V7X}"
     commands:
       - |
         buildkite-agent meta-data set "Collective_Communication_Matmul_PerformanceTest" "unverified"
-  - label: "Record performance test result for Collective Communication Matmul"
-    key: "record_Collective_Communication_Matmul_PerformanceTest"
-    depends_on: "Collective_Communication_Matmul_PerformanceTest"
+  - label: "${TEST_LABEL_PREFIX} Record performance test result for Collective Communication Matmul"
+    key: "${TEST_KEY_PREFIX}record_Collective_Communication_Matmul_PerformanceTest"
+    depends_on: "${TEST_KEY_PREFIX}Collective_Communication_Matmul_PerformanceTest"
     env:
-      CI_TARGET: "Collective Communication Matmul"
+      CI_TARGET: "${TEST_LABEL_PREFIX}Collective Communication Matmul"
       CI_STAGE: "PerformanceTest"
       CI_CATEGORY: "kernel support matrix"
     agents:
       queue: cpu
     commands:
       - |
-        .buildkite/scripts/record_step_result.sh Collective_Communication_Matmul_PerformanceTest
+        .buildkite/scripts/record_step_result.sh ${TEST_KEY_PREFIX}Collective_Communication_Matmul_PerformanceTest

.buildkite/features/DCN-Based_P-D_disaggregation.yml

@@ -1,45 +1,49 @@
 # DCN-based P/D disaggregation
 # feature support matrix
 steps:
-  - label: "Correctness tests for DCN-based P/D disaggregation"
-    key: "DCN_based_P-D_disaggregation_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Correctness tests for DCN-based P/D disaggregation"
+    key: "${TEST_KEY_PREFIX}DCN_based_P-D_disaggregation_CorrectnessTest"
     soft_fail: true
     agents:
-      queue: tpu_v6e_queue
+      queue: "${TPU_SMALL_CORE_QUEUE}"
+    env:
+      IS_FOR_V7X: "${IS_FOR_V7X}"
     commands:
       - |
         buildkite-agent meta-data set "DCN_based_P-D_disaggregation_CorrectnessTest" "unverified"
-  - label: "Record correctness test result for DCN-based P/D disaggregation"
-    key: "record_DCN_based_P-D_disaggregation_CorrectnessTest"
-    depends_on: "DCN_based_P-D_disaggregation_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Record correctness test result for DCN-based P/D disaggregation"
+    key: "${TEST_KEY_PREFIX}record_DCN_based_P-D_disaggregation_CorrectnessTest"
+    depends_on: "${TEST_KEY_PREFIX}DCN_based_P-D_disaggregation_CorrectnessTest"
     env:
-      CI_TARGET: "DCN-based P/D disaggregation"
+      CI_TARGET: "${TEST_LABEL_PREFIX}DCN-based P/D disaggregation"
       CI_STAGE: "CorrectnessTest"
       CI_CATEGORY: "feature support matrix"
     agents:
       queue: cpu
     commands:
       - |
-        .buildkite/scripts/record_step_result.sh DCN_based_P-D_disaggregation_CorrectnessTest
+        .buildkite/scripts/record_step_result.sh ${TEST_KEY_PREFIX}DCN_based_P-D_disaggregation_CorrectnessTest
 
-  - label: "Performance tests for DCN-based P/D disaggregation"
-    key: "DCN_based_P-D_disaggregation_PerformanceTest"
-    depends_on: "record_DCN_based_P-D_disaggregation_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Performance tests for DCN-based P/D disaggregation"
+    key: "${TEST_KEY_PREFIX}DCN_based_P-D_disaggregation_PerformanceTest"
+    depends_on: "${TEST_KEY_PREFIX}record_DCN_based_P-D_disaggregation_CorrectnessTest"
     soft_fail: true
     agents:
-      queue: tpu_v6e_queue
+      queue: "${TPU_SMALL_CORE_QUEUE}"
+    env:
+      IS_FOR_V7X: "${IS_FOR_V7X}"
     commands:
       - |
         buildkite-agent meta-data set "DCN_based_P-D_disaggregation_PerformanceTest" "unverified"
-  - label: "Record performance test result for DCN-based P/D disaggregation"
-    key: "record_DCN_based_P-D_disaggregation_PerformanceTest"
-    depends_on: "DCN_based_P-D_disaggregation_PerformanceTest"
+  - label: "${TEST_LABEL_PREFIX} Record performance test result for DCN-based P/D disaggregation"
+    key: "${TEST_KEY_PREFIX}record_DCN_based_P-D_disaggregation_PerformanceTest"
+    depends_on: "${TEST_KEY_PREFIX}DCN_based_P-D_disaggregation_PerformanceTest"
     env:
-      CI_TARGET: "DCN-based P/D disaggregation"
+      CI_TARGET: "${TEST_LABEL_PREFIX}DCN-based P/D disaggregation"
       CI_STAGE: "PerformanceTest"
       CI_CATEGORY: "feature support matrix"
     agents:
       queue: cpu
     commands:
       - |
-        .buildkite/scripts/record_step_result.sh DCN_based_P-D_disaggregation_PerformanceTest
+        .buildkite/scripts/record_step_result.sh ${TEST_KEY_PREFIX}DCN_based_P-D_disaggregation_PerformanceTest

.buildkite/features/KV_Cache_Host_Offloading.yml

@@ -1,45 +1,49 @@
 # KV cache host offloading
 # feature support matrix
 steps:
-  - label: "Correctness tests for KV cache host offloading"
-    key: "KV_Cache_Host_Offloading_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Correctness tests for KV cache host offloading"
+    key: "${TEST_KEY_PREFIX}KV_Cache_Host_Offloading_CorrectnessTest"
     soft_fail: true
     agents:
-      queue: tpu_v6e_queue
+      queue: "${TPU_SMALL_CORE_QUEUE}"
+    env:
+      IS_FOR_V7X: "${IS_FOR_V7X}"
     commands:
       - |
         buildkite-agent meta-data set "KV_Cache_Host_Offloading_CorrectnessTest" "unverified"
-  - label: "Record correctness test result for KV cache host offloading"
-    key: "record_KV_Cache_Host_Offloading_CorrectnessTest"
-    depends_on: "KV_Cache_Host_Offloading_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Record correctness test result for KV cache host offloading"
+    key: "${TEST_KEY_PREFIX}record_KV_Cache_Host_Offloading_CorrectnessTest"
+    depends_on: "${TEST_KEY_PREFIX}KV_Cache_Host_Offloading_CorrectnessTest"
     env:
-      CI_TARGET: "KV cache host offloading"
+      CI_TARGET: "${TEST_LABEL_PREFIX}KV cache host offloading"
       CI_STAGE: "CorrectnessTest"
       CI_CATEGORY: "feature support matrix"
     agents:
       queue: cpu
     commands:
       - |
-        .buildkite/scripts/record_step_result.sh KV_Cache_Host_Offloading_CorrectnessTest
+        .buildkite/scripts/record_step_result.sh ${TEST_KEY_PREFIX}KV_Cache_Host_Offloading_CorrectnessTest
 
-  - label: "Performance tests for KV cache host offloading"
-    key: "KV_Cache_Host_Offloading_PerformanceTest"
-    depends_on: "record_KV_Cache_Host_Offloading_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Performance tests for KV cache host offloading"
+    key: "${TEST_KEY_PREFIX}KV_Cache_Host_Offloading_PerformanceTest"
+    depends_on: "${TEST_KEY_PREFIX}record_KV_Cache_Host_Offloading_CorrectnessTest"
     soft_fail: true
     agents:
-      queue: tpu_v6e_queue
+      queue: "${TPU_SMALL_CORE_QUEUE}"
+    env:
+      IS_FOR_V7X: "${IS_FOR_V7X}"
     commands:
       - |
         buildkite-agent meta-data set "KV_Cache_Host_Offloading_PerformanceTest" "unverified"
-  - label: "Record performance test result for KV cache host offloading"
-    key: "record_KV_Cache_Host_Offloading_PerformanceTest"
-    depends_on: "KV_Cache_Host_Offloading_PerformanceTest"
+  - label: "${TEST_LABEL_PREFIX} Record performance test result for KV cache host offloading"
+    key: "${TEST_KEY_PREFIX}record_KV_Cache_Host_Offloading_PerformanceTest"
+    depends_on: "${TEST_KEY_PREFIX}KV_Cache_Host_Offloading_PerformanceTest"
     env:
-      CI_TARGET: "KV cache host offloading"
+      CI_TARGET: "${TEST_LABEL_PREFIX}KV cache host offloading"
       CI_STAGE: "PerformanceTest"
       CI_CATEGORY: "feature support matrix"
     agents:
       queue: cpu
     commands:
       - |
-        .buildkite/scripts/record_step_result.sh KV_Cache_Host_Offloading_PerformanceTest
+        .buildkite/scripts/record_step_result.sh ${TEST_KEY_PREFIX}KV_Cache_Host_Offloading_PerformanceTest

.buildkite/features/LoRA_Torch.yml

@@ -1,47 +1,51 @@
 # LoRA_Torch
 # feature support matrix
 steps:
-  - label: "Correctness tests for LoRA_Torch"
-    key: "LoRA_Torch_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Correctness tests for LoRA_Torch"
+    key: "${TEST_KEY_PREFIX}LoRA_Torch_CorrectnessTest"
     soft_fail: true
     agents:
-      queue: tpu_v6e_queue
+      queue: "${TPU_SMALL_CORE_QUEUE}"
+    env:
+      IS_FOR_V7X: "${IS_FOR_V7X}"
     commands:
       - |
         .buildkite/scripts/run_in_docker.sh \
           bash -c 'MODEL_IMPL_TYPE=vllm TPU_BACKEND_TYPE=jax python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_lora.py'
-  - label: "Record correctness test result for LoRA_Torch"
-    key: "record_LoRA_Torch_CorrectnessTest"
-    depends_on: "LoRA_Torch_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Record correctness test result for LoRA_Torch"
+    key: "${TEST_KEY_PREFIX}record_LoRA_Torch_CorrectnessTest"
+    depends_on: "${TEST_KEY_PREFIX}LoRA_Torch_CorrectnessTest"
     env:
-      CI_TARGET: "LoRA_Torch"
+      CI_TARGET: "${TEST_LABEL_PREFIX}LoRA_Torch"
       CI_STAGE: "CorrectnessTest"
       CI_CATEGORY: "feature support matrix"
     agents:
       queue: cpu
     commands:
       - |
-        .buildkite/scripts/record_step_result.sh LoRA_Torch_CorrectnessTest
+        .buildkite/scripts/record_step_result.sh ${TEST_KEY_PREFIX}LoRA_Torch_CorrectnessTest
 
-  - label: "Performance tests for LoRA_Torch"
-    key: "LoRA_Torch_PerformanceTest"
-    depends_on: "record_LoRA_Torch_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Performance tests for LoRA_Torch"
+    key: "${TEST_KEY_PREFIX}LoRA_Torch_PerformanceTest"
+    depends_on: "${TEST_KEY_PREFIX}record_LoRA_Torch_CorrectnessTest"
     soft_fail: true
     agents:
-      queue: tpu_v6e_queue
+      queue: "${TPU_SMALL_CORE_QUEUE}"
+    env:
+      IS_FOR_V7X: "${IS_FOR_V7X}"
     commands:
       - |
         .buildkite/scripts/run_in_docker.sh \
           bash -c 'MODEL_IMPL_TYPE=vllm TPU_BACKEND_TYPE=jax python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_lora_perf.py'
-  - label: "Record performance test result for LoRA_Torch"
-    key: "record_LoRA_Torch_PerformanceTest"
-    depends_on: "LoRA_Torch_PerformanceTest"
+  - label: "${TEST_LABEL_PREFIX} Record performance test result for LoRA_Torch"
+    key: "${TEST_KEY_PREFIX}record_LoRA_Torch_PerformanceTest"
+    depends_on: "${TEST_KEY_PREFIX}LoRA_Torch_PerformanceTest"
     env:
-      CI_TARGET: "LoRA_Torch"
+      CI_TARGET: "${TEST_LABEL_PREFIX}LoRA_Torch"
       CI_STAGE: "PerformanceTest"
       CI_CATEGORY: "feature support matrix"
     agents:
       queue: cpu
     commands:
       - |
-        .buildkite/scripts/record_step_result.sh LoRA_Torch_PerformanceTest
+        .buildkite/scripts/record_step_result.sh ${TEST_KEY_PREFIX}LoRA_Torch_PerformanceTest

.buildkite/features/MLA.yml

@@ -1,45 +1,49 @@
 # MLA
 # kernel support matrix
 steps:
-  - label: "Correctness tests for MLA"
-    key: "MLA_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Correctness tests for MLA"
+    key: "${TEST_KEY_PREFIX}MLA_CorrectnessTest"
     soft_fail: true
     agents:
-      queue: tpu_v6e_queue
+      queue: "${TPU_SMALL_CORE_QUEUE}"
+    env:
+      IS_FOR_V7X: "${IS_FOR_V7X}"
     commands:
       - |
         buildkite-agent meta-data set "MLA_CorrectnessTest" "unverified"
-  - label: "Record correctness test result for MLA"
-    key: "record_MLA_CorrectnessTest"
-    depends_on: "MLA_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Record correctness test result for MLA"
+    key: "${TEST_KEY_PREFIX}record_MLA_CorrectnessTest"
+    depends_on: "${TEST_KEY_PREFIX}MLA_CorrectnessTest"
     env:
-      CI_TARGET: "MLA"
+      CI_TARGET: "${TEST_LABEL_PREFIX}MLA"
       CI_STAGE: "CorrectnessTest"
       CI_CATEGORY: "kernel support matrix"
     agents:
       queue: cpu
     commands:
       - |
-        .buildkite/scripts/record_step_result.sh MLA_CorrectnessTest
+        .buildkite/scripts/record_step_result.sh ${TEST_KEY_PREFIX}MLA_CorrectnessTest
 
-  - label: "Performance tests for MLA"
-    key: "MLA_PerformanceTest"
-    depends_on: "record_MLA_CorrectnessTest"
+  - label: "${TEST_LABEL_PREFIX} Performance tests for MLA"
+    key: "${TEST_KEY_PREFIX}MLA_PerformanceTest"
+    depends_on: "${TEST_KEY_PREFIX}record_MLA_CorrectnessTest"
     soft_fail: true
     agents:
-      queue: tpu_v6e_queue
+      queue: "${TPU_SMALL_CORE_QUEUE}"
+    env:
+      IS_FOR_V7X: "${IS_FOR_V7X}"
     commands:
       - |
         buildkite-agent meta-data set "MLA_PerformanceTest" "unverified"
-  - label: "Record performance test result for MLA"
-    key: "record_MLA_PerformanceTest"
-    depends_on: "MLA_PerformanceTest"
+  - label: "${TEST_LABEL_PREFIX} Record performance test result for MLA"
+    key: "${TEST_KEY_PREFIX}record_MLA_PerformanceTest"
+    depends_on: "${TEST_KEY_PREFIX}MLA_PerformanceTest"
     env:
-      CI_TARGET: "MLA"
+      CI_TARGET: "${TEST_LABEL_PREFIX}MLA"
       CI_STAGE: "PerformanceTest"
       CI_CATEGORY: "kernel support matrix"
     agents:
       queue: cpu
     commands:
       - |
-        .buildkite/scripts/record_step_result.sh MLA_PerformanceTest
+        .buildkite/scripts/record_step_result.sh ${TEST_KEY_PREFIX}MLA_PerformanceTest