Add the KV-canary mock-model end-to-end test harness (#26811)

python/sglang/test/mock_model/perturb_e2e_base.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import io
+import os
+from typing import ClassVar
+
+from sglang.test.kv_canary.e2e_base import CapturedServerE2EBase
+from sglang.test.kv_canary.utils import post_parallel_generate
+from sglang.test.mock_model.utils import (
+    MOCK_MODEL_PATH,
+    mock_model_server_args,
+    mock_model_server_env,
+)
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    popen_launch_server,
+)
+
+
+class MockModelPerturbE2EBase(CapturedServerE2EBase):
+    extra_env: ClassVar[dict[str, str]] = {}
+    extra_server_args: ClassVar[tuple[str, ...]] = ()
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        server_env = os.environ.copy()
+        server_env.update(mock_model_server_env())
+        server_env.update(cls.extra_env)
+
+        cls._stdout_buf = io.StringIO()
+        cls._stderr_buf = io.StringIO()
+
+        cls.process = popen_launch_server(
+            MOCK_MODEL_PATH,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=mock_model_server_args(
+                *cls.extra_server_args, canary_mode="log"
+            ),
+            env=server_env,
+            return_stdout_stderr=(cls._stdout_buf, cls._stderr_buf),
+        )
+
+    def send_parallel_requests(
+        self,
+        n: int = 4,
+        *,
+        max_new_tokens: int = 256,
+        timeout: float = 30.0,
+    ) -> list[dict]:
+        prompts = ["hello world " * 50] * n
+        return post_parallel_generate(
+            url=self.base_url + "/generate",
+            prompts=prompts,
+            max_new_tokens=max_new_tokens,
+            timeout=timeout,
+        )

python/sglang/test/mock_model/utils.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import io
+import subprocess
+from dataclasses import dataclass
+from typing import Any, Sequence
+
+from sglang.bench_serving import run_benchmark
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    get_benchmark_args,
+    popen_launch_server,
+)
+
+MOCK_MODEL_PATH = "Qwen/Qwen3-0.6B"
+
+_MOCK_MODEL_SERVER_ARGS_NO_CANARY: list[str] = [
+    "--load-format",
+    "dummy",
+    "--sampling-backend",
+    "pytorch",
+    "--disable-piecewise-cuda-graph",
+]
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class MockModelBenchResult:
+    result: dict[str, Any]
+    stdout: str
+    stderr: str
+    server_return_code: int | None
+
+    @property
+    def log_text(self) -> str:
+        return self.stdout + self.stderr
+
+    def log_tail(self, length: int = 2000) -> str:
+        return self.log_text[-length:]
+
+
+def mock_model_server_args(*extra_args: str, canary_mode: str = "raise") -> list[str]:
+    return [
+        *_MOCK_MODEL_SERVER_ARGS_NO_CANARY,
+        "--kv-canary",
+        canary_mode,
+        *extra_args,
+    ]
+
+
+def mock_model_server_env(*, input_check_enabled: bool = True) -> dict[str, str]:
+    """Return env overrides for popen_launch_server in mock-model + canary mode."""
+    return {}
+
+
+def run_mock_model_bench_serving(
+    *,
+    extra_server_args: Sequence[str],
+    input_check_enabled: bool = True,
+    num_prompts: int = 32,
+    random_input_len: int = 6144,
+    random_output_len: int = 1024,
+) -> MockModelBenchResult:
+    stdout_buf = io.StringIO()
+    stderr_buf = io.StringIO()
+
+    process: subprocess.Popen[Any] | None = None
+    try:
+        process = popen_launch_server(
+            MOCK_MODEL_PATH,
+            DEFAULT_URL_FOR_TEST,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=mock_model_server_args(*extra_server_args),
+            env=mock_model_server_env(input_check_enabled=input_check_enabled),
+            return_stdout_stderr=(stdout_buf, stderr_buf),
+        )
+
+        args = get_benchmark_args(
+            base_url=DEFAULT_URL_FOR_TEST,
+            dataset_name="random",
+            tokenizer=MOCK_MODEL_PATH,
+            num_prompts=num_prompts,
+            random_input_len=random_input_len,
+            random_output_len=random_output_len,
+            request_rate=float("inf"),
+            max_concurrency=num_prompts,
+        )
+        args.random_range_ratio = 1.0
+        args.warmup_requests = 0
+        args.disable_tqdm = True
+
+        result = run_benchmark(args)
+        server_return_code = process.poll()
+        bench_result = MockModelBenchResult(
+            result=result,
+            stdout=stdout_buf.getvalue(),
+            stderr=stderr_buf.getvalue(),
+            server_return_code=server_return_code,
+        )
+        _assert_mock_model_bench_succeeded(
+            bench_result=bench_result,
+            expected_completed=num_prompts,
+        )
+
+        return bench_result
+    finally:
+        if process is not None:
+            kill_process_tree(process.pid)
+
+
+def _assert_mock_model_bench_succeeded(
+    *,
+    bench_result: MockModelBenchResult,
+    expected_completed: int,
+) -> None:
+    completed = bench_result.result.get("completed")
+    if completed != expected_completed:
+        raise AssertionError(
+            f"Expected {expected_completed} completed requests, got {completed}"
+        )
+
+    if bench_result.server_return_code is not None:
+        raise AssertionError(
+            f"Mock-model server exited with code {bench_result.server_return_code}.\n{bench_result.log_tail()}"
+        )
+
+    if "kv_canary violation:" in bench_result.log_text:
+        raise AssertionError(
+            f"Unexpected kv_canary violation in mock-model server log.\n{bench_result.log_tail()}"
+        )

python/sglang/test/server_fixtures/disaggregation_fixture.py

@@ -1,8 +1,10 @@
+import io
 import logging
 import os
 import shlex
 import time
 import warnings
+from typing import ClassVar, Optional
 from urllib.parse import urlparse
 
 from sglang.srt.environ import envs
@@ -22,6 +24,14 @@ logger = logging.getLogger(__name__)
 
 
 class PDDisaggregationServerBase(CustomTestCase):
+    capture_per_side_logs: ClassVar[bool] = False
+    extra_prefill_env: ClassVar[dict[str, str]] = {}
+    extra_decode_env: ClassVar[dict[str, str]] = {}
+    _prefill_stdout_buf: ClassVar[Optional[io.StringIO]] = None
+    _prefill_stderr_buf: ClassVar[Optional[io.StringIO]] = None
+    _decode_stdout_buf: ClassVar[Optional[io.StringIO]] = None
+    _decode_stderr_buf: ClassVar[Optional[io.StringIO]] = None
+
     @classmethod
     def setUpClass(cls):
         os.environ["MC_TCP_ENABLE_CONNECTION_POOL"] = "true"
@@ -40,6 +50,11 @@ class PDDisaggregationServerBase(CustomTestCase):
             f"{cls.base_host=} {cls.lb_port=} {cls.prefill_port=} {cls.decode_port=} {cls.bootstrap_port=}"
         )
         cls.process_lb, cls.process_decode, cls.process_prefill = None, None, None
+        if cls.capture_per_side_logs:
+            cls._prefill_stdout_buf = io.StringIO()
+            cls._prefill_stderr_buf = io.StringIO()
+            cls._decode_stdout_buf = io.StringIO()
+            cls._decode_stderr_buf = io.StringIO()
         cls._fail_fast_stop = None
 
         # config transfer backend and rdma devices
@@ -81,6 +96,12 @@ class PDDisaggregationServerBase(CustomTestCase):
             cls.prefill_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=prefill_args,
+            env=dict(cls.extra_prefill_env),
+            return_stdout_stderr=(
+                (cls._prefill_stdout_buf, cls._prefill_stderr_buf)
+                if cls.capture_per_side_logs
+                else None
+            ),
         )
 
     @classmethod
@@ -102,6 +123,12 @@ class PDDisaggregationServerBase(CustomTestCase):
             cls.decode_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=decode_args,
+            env=dict(cls.extra_decode_env),
+            return_stdout_stderr=(
+                (cls._decode_stdout_buf, cls._decode_stderr_buf)
+                if cls.capture_per_side_logs
+                else None
+            ),
         )
 
     @classmethod
@@ -163,6 +190,20 @@ class PDDisaggregationServerBase(CustomTestCase):
                 except Exception as e:
                     print(f"Error killing process {process.pid}: {e}")
 
+        if cls.capture_per_side_logs:
+            for buf in (
+                cls._prefill_stdout_buf,
+                cls._prefill_stderr_buf,
+                cls._decode_stdout_buf,
+                cls._decode_stderr_buf,
+            ):
+                if buf is not None:
+                    buf.close()
+            cls._prefill_stdout_buf = None
+            cls._prefill_stderr_buf = None
+            cls._decode_stdout_buf = None
+            cls._decode_stderr_buf = None
+
         # wait for 5 seconds
         time.sleep(5)
 

test/registered/mock_model/test_e2e_pd.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import time
+import unittest
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import ClassVar, Dict, List
+
+import requests
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kv_canary.violation_log_utils import assert_no_violation_in_log
+from sglang.test.mock_model.utils import (
+    MOCK_MODEL_PATH,
+    mock_model_server_args,
+    mock_model_server_env,
+)
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+
+register_cuda_ci(est_time=600, stage="extra-a", runner_config="2-gpu-large")
+
+# DO NOT pass --disable-cuda-graph in canary e2e tests.  The canary kernel
+# must run inside the cuda graph alongside the real attn kernel; disabling the
+# full graph silently bypasses the only path that exercises that invariant
+# end-to-end.
+#
+# --disable-piecewise-cuda-graph is REQUIRED by canary: install_canary
+# (api.py) asserts it, and the SingleForwardManager design depends on it.
+# mock_model_server_args() already passes it; do not remove it.
+_NUM_PROMPTS = 32
+_INPUT_LEN = 6144
+_OUTPUT_LEN = 1024
+
+
+def _send_parallel_requests(
+    base_url: str,
+    *,
+    n: int,
+    max_new_tokens: int,
+    timeout: float = 60.0,
+    max_workers: int = 16,
+) -> List[Dict[str, object]]:
+    """Fire N /generate requests concurrently; return raw response dicts."""
+
+    def _one(i: int) -> Dict[str, object]:
+        payload = {
+            "input_ids": _make_input_ids(seed=i, length=_INPUT_LEN),
+            "sampling_params": {"max_new_tokens": max_new_tokens, "temperature": 0.0},
+        }
+        try:
+            resp = requests.post(base_url + "/generate", json=payload, timeout=timeout)
+            return {"index": i, "status_code": resp.status_code, "text": resp.text}
+        except requests.exceptions.RequestException as exc:
+            return {"index": i, "error": repr(exc)}
+
+    results: List[Dict[str, object]] = []
+    with ThreadPoolExecutor(max_workers=max_workers) as pool:
+        futures = [pool.submit(_one, i) for i in range(n)]
+        for fut in as_completed(futures):
+            results.append(fut.result())
+    results.sort(key=lambda r: r["index"])
+    return results
+
+
+def _make_input_ids(*, seed: int, length: int) -> List[int]:
+    return [((seed + i) % 2048) + 1 for i in range(length)]
+
+
+class _MockModelPDBase(PDDisaggregationServerBase):
+    """PD fixture for mock-model + canary e2e tests."""
+
+    capture_per_side_logs = True
+    model: ClassVar[str] = MOCK_MODEL_PATH
+    extra_prefill_args: ClassVar[List[str]] = mock_model_server_args(
+        "--skip-server-warmup"
+    )
+    extra_decode_args: ClassVar[List[str]] = mock_model_server_args(
+        "--skip-server-warmup"
+    )
+    extra_prefill_env: ClassVar[Dict[str, str]] = mock_model_server_env(
+        input_check_enabled=True
+    )
+    extra_decode_env: ClassVar[Dict[str, str]] = mock_model_server_env(
+        input_check_enabled=True
+    )
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        super().setUpClass()
+        cls.launch_all()
+
+    def assert_no_canary_violation(self) -> None:
+        time.sleep(2)
+        log_text = "".join(
+            buf.getvalue()
+            for buf in (
+                self._prefill_stdout_buf,
+                self._prefill_stderr_buf,
+                self._decode_stdout_buf,
+                self._decode_stderr_buf,
+            )
+            if buf is not None
+        )
+        assert_no_violation_in_log(log_text)
+
+
+class TestPdTransferCanaryClean(_MockModelPDBase, unittest.TestCase):
+    """PD standard scenario + baseline canary (input-check, no real-KV checksum); no violation expected."""
+
+    def test_pd_transfer_canary_clean(self) -> None:
+        # Step 1: send parallel requests through the LB to exercise PD transfer path.
+        results = _send_parallel_requests(
+            self.lb_url,
+            n=_NUM_PROMPTS,
+            max_new_tokens=_OUTPUT_LEN,
+            timeout=240.0,
+            max_workers=_NUM_PROMPTS,
+        )
+
+        # Step 2: every request must complete with status 200.
+        for result in results:
+            self.assertEqual(result.get("status_code"), 200, result)
+
+        # Step 3: servers must stay alive.
+        self.assertIsNone(self.process_prefill.poll(), "Prefill server died")
+        self.assertIsNone(self.process_decode.poll(), "Decode server died")
+        self.assert_no_canary_violation()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/mock_model/test_e2e_pp.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.mock_model.utils import run_mock_model_bench_serving
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=600, stage="extra-a", runner_config="2-gpu-large")
+
+
+class TestE2EPipelineParallel(CustomTestCase):
+    def test_pp_no_canary_violation(self) -> None:
+        run_mock_model_bench_serving(
+            extra_server_args=["--pp-size", "2"],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/mock_model/test_e2e_tp.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.mock_model.utils import run_mock_model_bench_serving
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=600, stage="extra-a", runner_config="2-gpu-large")
+
+
+class TestE2ETensorParallel(CustomTestCase):
+    def test_tp_no_canary_violation(self) -> None:
+        run_mock_model_bench_serving(
+            extra_server_args=["--tp", "2", "--mem-fraction-static", "0.88"],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()