feat: rust sglang server openai apis (#33103)

Co-authored-by: Rain Jiang <rain-jiang@outlook.com>

test/registered/openai_server/basic/test_openai_completion_rust.py

@@ -0,0 +1,104 @@
+import math
+import unittest
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_rust_server_built,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=80, stage="base-b", runner_config="1-gpu-small")
+
+
+@unittest.skipUnless(
+    is_rust_server_built(),
+    "embedded rust server extension not built",
+)
+class TestOpenAICompletionRustParity(CustomTestCase):
+    model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+    api_key = "sk-123456"
+
+    def _get_logprobs(self, *, rust_frontend):
+        process = popen_launch_server(
+            self.model,
+            DEFAULT_URL_FOR_TEST,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            api_key=self.api_key,
+            env={"SGLANG_RUST_SERVER": "1" if rust_frontend else "0"},
+            other_args=["--random-seed", "42"],
+        )
+        try:
+            response = requests.post(
+                DEFAULT_URL_FOR_TEST + "/v1/completions",
+                headers={"Authorization": f"Bearer {self.api_key}"},
+                json={
+                    "model": self.model,
+                    "prompt": "The capital of France is",
+                    "temperature": 0,
+                    "max_tokens": 8,
+                    "logprobs": 5,
+                },
+                timeout=30,
+            )
+            response.raise_for_status()
+            return response.json()["choices"][0]["logprobs"]
+        finally:
+            kill_process_tree(process.pid)
+
+    @staticmethod
+    def _kl_divergence(reference, candidate):
+        assert reference.keys() == candidate.keys()
+        keys = sorted(reference)
+        reference_max = max(reference.values())
+        candidate_max = max(candidate.values())
+        reference_weights = [math.exp(reference[key] - reference_max) for key in keys]
+        candidate_weights = [math.exp(candidate[key] - candidate_max) for key in keys]
+        reference_sum = sum(reference_weights)
+        candidate_sum = sum(candidate_weights)
+        reference_probabilities = [
+            weight / reference_sum for weight in reference_weights
+        ]
+        candidate_probabilities = [
+            weight / candidate_sum for weight in candidate_weights
+        ]
+        return sum(
+            reference_probability
+            * math.log(reference_probability / candidate_probability)
+            for reference_probability, candidate_probability in zip(
+                reference_probabilities,
+                candidate_probabilities,
+                strict=True,
+            )
+        )
+
+    def test_logprobs_have_zero_kl_against_python_frontend(self):
+        python_logprobs = self._get_logprobs(rust_frontend=False)
+        rust_logprobs = self._get_logprobs(rust_frontend=True)
+
+        self.assertEqual(rust_logprobs["tokens"], python_logprobs["tokens"])
+        self.assertEqual(
+            rust_logprobs["token_logprobs"],
+            python_logprobs["token_logprobs"],
+        )
+        self.assertEqual(
+            len(rust_logprobs["top_logprobs"]),
+            len(python_logprobs["top_logprobs"]),
+        )
+        for python_top, rust_top in zip(
+            python_logprobs["top_logprobs"],
+            rust_logprobs["top_logprobs"],
+            strict=True,
+        ):
+            self.assertEqual(rust_top, python_top)
+            self.assertEqual(self._kl_divergence(python_top, rust_top), 0.0)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/openai_server/basic/test_openai_server.py

@@ -25,10 +25,11 @@ from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
+    is_rust_server_built,
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=182, stage="base-b", runner_config="1-gpu-small")
+register_cuda_ci(est_time=240, stage="base-b", runner_config="1-gpu-small")
 register_amd_ci(est_time=200, suite="stage-b-test-1-gpu-small-amd")
 
 
@@ -443,6 +444,57 @@ The SmartHome Mini is a compact smart home assistant available in black or white
             client.models.retrieve("non-existent-model")
 
 
+@unittest.skipUnless(
+    is_rust_server_built(),
+    "embedded rust server extension not built",
+)
+class TestOpenAICompletionWithRust(CustomTestCase):
+    """Run the existing Completion matrix unchanged through the Rust frontend."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.api_key = "sk-123456"
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            api_key=cls.api_key,
+            env={"SGLANG_RUST_SERVER": "1"},
+        )
+        cls.base_url += "/v1"
+        cls.tokenizer = get_tokenizer(cls.model)
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    # Reuse the reference test methods directly so Python and Rust coverage
+    # cannot drift into separate matrices.
+    run_completion = TestOpenAIServer.run_completion
+    run_completion_stream = TestOpenAIServer.run_completion_stream
+    test_completion = TestOpenAIServer.test_completion
+    test_completion_stream = TestOpenAIServer.test_completion_stream
+
+
+@unittest.skipUnless(
+    is_rust_server_built(),
+    "embedded rust server extension not built",
+)
+class TestOpenAIChatWithRust(TestOpenAICompletionWithRust):
+    """Run the existing Chat matrix unchanged through the Rust frontend."""
+
+    run_chat_completion = TestOpenAIServer.run_chat_completion
+    run_chat_completion_stream = TestOpenAIServer.run_chat_completion_stream
+    test_chat_completion = TestOpenAIServer.test_chat_completion
+    test_chat_completion_stream = TestOpenAIServer.test_chat_completion_stream
+
+    # This class is a Chat gate; Completion already has its own Rust matrix.
+    test_completion = None
+    test_completion_stream = None
+
+
 class TestOpenAIServerv1Responses(CustomTestCase):
     @classmethod
     def setUpClass(cls):