add aime25 test case

Author: zhengkezhou1

.github/workflows/gke-connectivity-smoke.yml

@@ -192,7 +192,6 @@ jobs:
           set -euo pipefail
 
           HEADLESS_SERVICE_NAME="${WORKLOAD_NAME}-headless"
-          JOB_COMPLETE_TIMEOUT="30m"
           TPU_PROCESS_ADDRESSES="${WORKLOAD_NAME}-0.${HEADLESS_SERVICE_NAME}:8471,${WORKLOAD_NAME}-1.${HEADLESS_SERVICE_NAME}:8471,${WORKLOAD_NAME}-2.${HEADLESS_SERVICE_NAME}:8471,${WORKLOAD_NAME}-3.${HEADLESS_SERVICE_NAME}:8471"
           TPU_WORKER_HOSTNAMES="${WORKLOAD_NAME}-0.${HEADLESS_SERVICE_NAME},${WORKLOAD_NAME}-1.${HEADLESS_SERVICE_NAME},${WORKLOAD_NAME}-2.${HEADLESS_SERVICE_NAME},${WORKLOAD_NAME}-3.${HEADLESS_SERVICE_NAME}"
 
@@ -346,7 +345,7 @@ jobs:
           # Wait for Job to reach a terminal state (Complete or Failed). Plain
           # `kubectl wait --for=condition=complete` only matches success and
           # would block until --timeout even if the Job has already Failed.
-          deadline=$(($(date +%s) + 1800))
+          deadline=$(($(date +%s) + 5400))
           while [[ $(date +%s) -lt $deadline ]]; do
             conds=$(kubectl get job "${WORKLOAD_NAME}" --namespace "${NAMESPACE}" \
               -o jsonpath='{.status.conditions[?(@.type=="Complete")].status}|{.status.conditions[?(@.type=="Failed")].status}')

test/srt/mulit_host/multi_host_suite.py

@@ -26,6 +26,7 @@ class AccuracyCase:
     eval_batch_size: int = 32
     generation_config: dict[str, Any] = field(default_factory=dict)
     limit: int | None = None
+    timeout: int | None = None
     dry_run_result: Literal["success", "failed"] = "success"
 
 

test/srt/mulit_host/run_suite.py

@@ -189,12 +189,15 @@ def run_accuracy_case(case: AccuracyCase, port: int) -> None:
         cmd.extend(["--generation-config", json.dumps(case.generation_config)])
     if case.limit is not None:
         cmd.extend(["--limit", str(case.limit)])
+    if case.timeout is not None:
+        cmd.extend(["--timeout", str(case.timeout)])
 
     _log(
         "Running accuracy case "
         f"name={case.name}, dataset={case.dataset}, "
         f"eval_batch_size={case.eval_batch_size}, "
-        f"generation_config={case.generation_config}, limit={case.limit}"
+        f"generation_config={case.generation_config}, limit={case.limit}, "
+        f"timeout={case.timeout}"
     )
     _log(f"Command: {' '.join(cmd)}")
     completed = subprocess.run(cmd, check=False)

test/srt/mulit_host/test_mimo_flash.py

@@ -58,12 +58,25 @@ def get_suites() -> list[MultiHostSuite]:
                         #     flush_cache=True,
                         # ),
                         AccuracyCase(
-                            name="mimo-flash-accuracy",
+                            name="mimo-flash-gsm8k",
                             dataset="gsm8k",
                             model_id="XiaomiMiMo/MiMo-V2-Flash",
                             eval_batch_size=32,
                             generation_config={"temperature": 0.8, "top_p": 0.95},
                         ),
+                        AccuracyCase(
+                            name="mimo-flash-aime25",
+                            dataset="aime25",
+                            model_id="XiaomiMiMo/MiMo-V2-Flash",
+                            eval_batch_size=16,
+                            timeout=6000000,
+                            generation_config={
+                                "temperature": 1,
+                                "top_p": 0.95,
+                                "max_tokens": 131072,
+                                "chat_template_kwargs": {"enable_thinking": True},
+                            },
+                        ),
                     ],
                 )
             ],