fix(ci): apply MMMU retry logic to all affected test files (#17329)

python/sglang/test/kits/mmmu_vlm_kit.py

@@ -64,11 +64,38 @@ def _run_lmms_eval_with_retry(cmd: list[str], timeout: int = 3600) -> None:
             capture_output=True,
             text=True,
         )
-        # Print captured output to maintain visibility of successful runs
+        # Check for errors in output even if exit code is 0
-        if result.stdout:
+        # lmms_eval sometimes returns 0 even when errors occur
-            print(result.stdout, end="")
+        combined_output = result.stdout + result.stderr
-        if result.stderr:
+        if _is_mmmu_parquet_corruption(combined_output):
-            print(result.stderr, end="")
+            print(
+                "Detected MMMU parquet corruption error in output. Attempting recovery..."
+            )
+            if _cleanup_mmmu_dataset_cache():
+                print("Retrying lmms_eval with fresh download...")
+                with temp_set_env(
+                    HF_HUB_OFFLINE="0",
+                    HF_DATASETS_DOWNLOAD_MODE="force_redownload",
+                ):
+                    retry_result = subprocess.run(
+                        cmd, check=True, timeout=timeout, capture_output=True, text=True
+                    )
+                    # Print retry output
+                    if retry_result.stdout:
+                        print(retry_result.stdout, end="")
+                    if retry_result.stderr:
+                        print(retry_result.stderr, end="")
+            else:
+                print(
+                    f"Failed to cleanup corrupted MMMU cache. Output from lmms_eval:\nStdout:\n{result.stdout}\nStderr:\n{result.stderr}"
+                )
+                raise RuntimeError("Failed to cleanup corrupted MMMU cache")
+        else:
+            # Print captured output to maintain visibility of successful runs
+            if result.stdout:
+                print(result.stdout, end="")
+            if result.stderr:
+                print(result.stderr, end="")
     except subprocess.CalledProcessError as e:
         error_output = e.stderr + e.stdout
         if _is_mmmu_parquet_corruption(error_output):
@@ -79,7 +106,14 @@ def _run_lmms_eval_with_retry(cmd: list[str], timeout: int = 3600) -> None:
                     HF_HUB_OFFLINE="0",
                     HF_DATASETS_DOWNLOAD_MODE="force_redownload",
                 ):
-                    subprocess.run(cmd, check=True, timeout=timeout)
+                    retry_result = subprocess.run(
+                        cmd, check=True, timeout=timeout, capture_output=True, text=True
+                    )
+                    # Print retry output
+                    if retry_result.stdout:
+                        print(retry_result.stdout, end="")
+                    if retry_result.stderr:
+                        print(retry_result.stderr, end="")
             else:
                 print(
                     f"Failed to cleanup corrupted MMMU cache. Error from lmms_eval:\nStdout:\n{e.stdout}\nStderr:\n{e.stderr}"

test/manual/nightly/test_vlms_piecewise_cuda_graph.py

@@ -3,12 +3,12 @@ import glob
 import json
 import os
 import random
-import subprocess
 import sys
 import unittest
 from types import SimpleNamespace
 
 from sglang.srt.utils import kill_process_tree
+from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
@@ -85,11 +85,7 @@ class TestVLMPiecewiseCudaGraph(CustomTestCase):
             str(output_path),
         ]
 
-        subprocess.run(
+        _run_lmms_eval_with_retry(cmd, timeout=3600)
-            cmd,
-            check=True,
-            timeout=3600,
-        )
 
     def _run_vlm_mmmu_test(
         self,

test/manual/nightly/test_vlms_vit_cuda_graph.py

@@ -3,12 +3,12 @@ import glob
 import json
 import os
 import random
-import subprocess
 import sys
 import unittest
 from types import SimpleNamespace
 
 from sglang.srt.utils import kill_process_tree
+from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
@@ -88,11 +88,7 @@ class TestVLMViTCudaGraph(CustomTestCase):
             str(output_path),
         ]
 
-        subprocess.run(
+        _run_lmms_eval_with_retry(cmd, timeout=3600)
-            cmd,
-            check=True,
-            timeout=3600,
-        )
 
     def _run_vlm_mmmu_test(
         self,

test/registered/vlm/test_encoder_dp.py

@@ -3,13 +3,13 @@ import glob
 import json
 import os
 import random
-import subprocess
 import sys
 import unittest
 from types import SimpleNamespace
 
 from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
@@ -94,11 +94,7 @@ class TestVLMEncoderDP(CustomTestCase):
             str(output_path),
         ]
 
-        subprocess.run(
+        _run_lmms_eval_with_retry(cmd, timeout=3600)
-            cmd,
-            check=True,
-            timeout=3600,
-        )
 
     def _run_vlm_mmmu_test(
         self,

test/srt/test_epd_disaggregation.py

@@ -1,9 +1,9 @@
 import os
-import subprocess
 import threading
 import unittest
 
 from sglang.srt.utils import kill_process_tree
+from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
 from sglang.test.server_fixtures.disaggregation_fixture import (
     PDDisaggregationServerBase,
 )
@@ -177,7 +177,7 @@ class TestEPDDisaggregationOneEncoder(PDDisaggregationServerBase):
             limit,
         ]
 
-        subprocess.run(cmd, check=True, timeout=3600)
+        _run_lmms_eval_with_retry(cmd, timeout=3600)
 
     def test_mmmu(self):
         """Test MMMU evaluation with EPD disaggregation"""
@@ -393,7 +393,7 @@ class TestEPDDisaggregationMultiEncoders(PDDisaggregationServerBase):
             limit,
         ]
 
-        subprocess.run(cmd, check=True, timeout=3600)
+        _run_lmms_eval_with_retry(cmd, timeout=3600)
 
     def test_mmmu(self):
         """Test MMMU evaluation with EPD disaggregation (multiple encoders)"""