Adapt fixture-kit to gsm8k mixin (#15599)

python/sglang/test/gsm8k_mixin.py

@@ -1,44 +0,0 @@
-from abc import ABC
-from types import SimpleNamespace
-
-from sglang.srt.utils import kill_process_tree
-from sglang.test.few_shot_gsm8k import run_eval
-from sglang.test.test_utils import (
-    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-    DEFAULT_URL_FOR_TEST,
-    popen_launch_server,
-)
-
-
-class GSM8KMixin(ABC):
-    accuracy: float
-    model: str
-    other_args: list[str] = []
-
-    @classmethod
-    def setUpClass(cls):
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=cls.other_args,
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def test_gsm8k(self):
-        args = SimpleNamespace(
-            num_shots=5,
-            data_path=None,
-            num_questions=200,
-            max_new_tokens=512,
-            parallel=128,
-            host="http://127.0.0.1",
-            port=int(self.base_url.split(":")[-1]),
-        )
-        metrics = run_eval(args)
-        print(f"{metrics=}")
-        self.assertGreaterEqual(metrics["accuracy"], self.accuracy)

python/sglang/test/kits/gsm8k_accuracy_kit.py

@@ -0,0 +1,22 @@
+from types import SimpleNamespace
+
+from sglang.test.few_shot_gsm8k import run_eval as run_eval_gsm8k
+from sglang.test.test_utils import CustomTestCase
+
+
+class GSM8KMixin:
+    accuracy: float
+
+    def test_gsm8k(self: CustomTestCase):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=200,
+            max_new_tokens=512,
+            parallel=128,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval_gsm8k(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["accuracy"], self.accuracy)

python/sglang/test/kits/mmmu_vlm_kit.py

@@ -2,13 +2,13 @@ import glob
 import json
 import os
 import subprocess
-from abc import ABC
 from types import SimpleNamespace
 
 from sglang.srt.utils import kill_process_tree
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
     popen_launch_server,
 )
 
@@ -16,7 +16,8 @@ from sglang.test.test_utils import (
 DEFAULT_MEM_FRACTION_STATIC = 0.8
 
 
-class MMMUVLMMixin(ABC):
+class MMMUVLMTestBase(CustomTestCase):
+    # TODO: split the MMMUVLMTestBase into a fixture and a mixin
     parsed_args = None  # Class variable to store args
     other_args = []
     mmmu_args = []

python/sglang/test/server_fixtures/default_fixture.py

@@ -0,0 +1,67 @@
+import logging
+import time
+from contextlib import contextmanager
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def openai_api_env(api_key: str):
+    """Context manager to set OpenAI API environment variables."""
+    import os
+
+    original_api_key = os.environ.get("OPENAI_API_KEY")
+    original_base_url = os.environ.get("OPENAI_API_BASE")
+
+    os.environ["OPENAI_API_KEY"] = api_key
+    os.environ["OPENAI_API_BASE"] = f"{DEFAULT_URL_FOR_TEST}/v1"
+
+    try:
+        yield
+    finally:
+        if original_api_key is not None:
+            os.environ["OPENAI_API_KEY"] = original_api_key
+        else:
+            del os.environ["OPENAI_API_KEY"]
+
+        if original_base_url is not None:
+            os.environ["OPENAI_API_BASE"] = original_base_url
+        else:
+            del os.environ["OPENAI_API_BASE"]
+
+
+class DefaultServerBase(CustomTestCase):
+    model = None
+    base_url = DEFAULT_URL_FOR_TEST
+    timeout = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
+    other_args: list[str] = []
+
+    # For OpenAI API settings
+    api_key = "sk-123456"
+
+    @classmethod
+    def setUpClass(cls):
+        assert cls.model is not None, "Please set cls.model in subclass"
+
+        # Set OpenAI API key and base URL environment variables.
+        # Needed for lmm-evals to work.
+        with openai_api_env(cls.api_key):
+            cls.process = popen_launch_server(
+                cls.model,
+                cls.base_url,
+                timeout=cls.timeout,
+                other_args=cls.other_args,
+            )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+        time.sleep(2)