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

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()