Enable TPU V0 and TPU V1 test on CI

Author: khluu

scripts/test-template-aws.j2

@@ -357,25 +357,20 @@ steps:
     command: nvidia-smi && bash .buildkite/run-gh200-test.sh
   {% endif %}
 
-  {% if branch == "main" %}
-  - label: "TPU Test"
+  - label: "TPU V0 Test"
     depends_on: ~
     soft_fail: True
     agents:
       queue: tpu_queue
     commands: 
     - yes | docker system prune -a
     - if [[ -f ".buildkite/run-tpu-test.sh" ]]; then bash .buildkite/run-tpu-test.sh; fi
-  {% else %}
-  - block: "Run TPU Test"
+
+  - label: "TPU V1 Test"
     depends_on: ~
-    key: block-tpu-test
-  
-  - label: "TPU Test"
-    depends_on: block-tpu-test
+    soft_fail: true
     agents:
       queue: tpu_queue
-    commands: 
-    - yes | docker system prune -a
-    - if [[ -f ".buildkite/run-tpu-test.sh" ]]; then bash .buildkite/run-tpu-test.sh; fi
-  {% endif %}
+    commands:
+      - if [[ -f ".buildkite/run-tpu-v1-test.sh" ]]; then bash .buildkite/run-tpu-v1-test.sh; fi
+      - yes | docker system prune -a