| """CPU-only request/early-stop and no-output-disclosure regressions.""" |
| from __future__ import annotations |
|
|
| from contextlib import redirect_stdout |
| import importlib.util |
| import io |
| import json |
| from pathlib import Path |
| import unittest |
| from unittest.mock import patch |
|
|
| SPEC = importlib.util.spec_from_file_location("warmup", Path(__file__).parents[1] / "warmup.py") |
| assert SPEC and SPEC.loader |
| warmup = importlib.util.module_from_spec(SPEC) |
| SPEC.loader.exec_module(warmup) |
|
|
|
|
| def reply(count: int = 32, reason: str = "length") -> str: |
| return json.dumps({"usage": {"prompt_tokens": 5041, "completion_tokens": count}, |
| "choices": [{"text": "DO NOT DISCLOSE", "finish_reason": reason, |
| "stop_reason": None}]}) |
|
|
|
|
| class WarmupTest(unittest.TestCase): |
| def test_internal_budget(self) -> None: |
| body = warmup.payload('quoted"model', "salt", 1, 0) |
| self.assertEqual(json.loads(json.dumps(body)), body) |
| self.assertEqual(body["min_tokens"], body["max_tokens"]) |
| self.assertTrue(body["ignore_eos"]) |
| self.assertEqual(body["priority"], 20) |
| self.assertNotIn("messages", body) |
|
|
| def test_one_token_and_wrong_reason_fail(self) -> None: |
| for raw in (reply(1, "stop"), reply(32, "stop"), '{"error": "DO NOT DISCLOSE"}'): |
| with self.assertRaises(ValueError) as err: |
| warmup.summary(raw) |
| self.assertNotIn("DO NOT DISCLOSE", str(err.exception)) |
|
|
| def test_missing_metrics_are_unknown(self) -> None: |
| record = warmup.summary(reply()) |
| self.assertIsNone(record["speculative_metrics"]) |
| self.assertIsNone(record["reasoning_tokens"]) |
| self.assertNotIn("DO NOT DISCLOSE", json.dumps(record)) |
|
|
| def test_decode_waves_do_not_serialize_long_prefills(self) -> None: |
| long = warmup.payload("test", "salt", 1, 0) |
| short = warmup.payload("test", "salt", 4, 0) |
| self.assertGreater(len(str(long["prompt"])), 20000) |
| self.assertLess(len(str(short["prompt"])), 200) |
|
|
| def test_report_actual_spec_counts_without_text(self) -> None: |
| data = json.loads(reply()) |
| data["metrics"] = {"speculative_decoding": {"num_spec_steps": 12, |
| "num_accepted_draft_tokens": 19, "num_draft_tokens": 60}} |
| data["usage"]["prompt_tokens_details"] = {"cached_tokens": 0} |
| result = warmup.summary(json.dumps(data)) |
| self.assertEqual(result["speculative_metrics"]["num_spec_steps"], 12) |
| self.assertEqual(result["cached_tokens"], 0) |
|
|
| def test_idle_fails_closed(self) -> None: |
| for raw in ("", "vllm:num_requests_running 0", "vllm:num_requests_running NaN\nvllm:num_requests_waiting 0", |
| "vllm:num_requests_running 1\nvllm:num_requests_waiting 0"): |
| with self.assertRaises(ValueError): |
| warmup.idle(raw) |
|
|
| def test_all_shapes_without_network(self) -> None: |
| bodies: list[dict[str, object]] = [] |
|
|
| def read(path: str, body: dict[str, object] | None) -> str: |
| if body is None: |
| self.assertEqual(path, "/metrics") |
| return 'vllm:num_requests_running{model_name="test"} 0\nvllm:num_requests_waiting 0' |
| bodies.append(body) |
| return reply() |
|
|
| output = io.StringIO() |
| with patch.object(warmup, "urlopen", side_effect=AssertionError("network forbidden")), redirect_stdout(output): |
| warmup.run("test", read) |
| self.assertEqual(len(bodies), 10) |
| self.assertEqual(len({str(body["prompt"]) for body in bodies}), 10) |
| self.assertEqual([json.loads(line)["requested_concurrency"] for line in output.getvalue().splitlines()], [1, 2, 3, 4]) |
| self.assertNotIn("DO NOT DISCLOSE", output.getvalue()) |
|
|
| def test_foreign_work_stops_before_post(self) -> None: |
| def read(path: str, body: dict[str, object] | None) -> str: |
| self.assertIsNone(body) |
| return "vllm:num_requests_running 1\nvllm:num_requests_waiting 0" |
| with self.assertRaises(ValueError): |
| warmup.run("test", read) |
|
|
| def test_long_prefill_requires_real_uncached_full_chunk(self) -> None: |
| for prompt_tokens, cached_tokens in ((4096, 0), (5041, 64)): |
| def read(path: str, body: dict[str, object] | None) -> str: |
| if body is None: |
| return "vllm:num_requests_running 0\nvllm:num_requests_waiting 0" |
| data = json.loads(reply()) |
| data["usage"]["prompt_tokens"] = prompt_tokens |
| data["usage"]["prompt_tokens_details"] = {"cached_tokens": cached_tokens} |
| return json.dumps(data) |
| with self.assertRaises(ValueError): |
| warmup.run("test", read) |
|
|
|
|
| def test_deadline_stops_before_the_next_wave_and_bounds_requests(self) -> None: |
| posts: list[dict[str, object]] = [] |
|
|
| def read(path: str, body: dict[str, object] | None) -> str: |
| if body is None: |
| return "vllm:num_requests_running 0\nvllm:num_requests_waiting 0" |
| posts.append(body) |
| return reply() |
|
|
| with self.assertRaisesRegex(ValueError, "deadline"): |
| warmup.run("test", read, warmup.time.monotonic() - 1) |
| self.assertEqual(posts, []) |
| with self.assertRaisesRegex(ValueError, "deadline"): |
| warmup.request_timeout(warmup.time.monotonic() - 0.1) |
| self.assertEqual(warmup.request_timeout(warmup.time.monotonic() + 1000), warmup.TIMEOUT) |
| self.assertLessEqual(warmup.request_timeout(warmup.time.monotonic() + 5), 5) |
| self.assertLess(warmup.DEADLINE_S + 432, 600 - 60) |
|
|
| def test_late_final_wave_fails_and_timeout_floor_is_the_time_left(self) -> None: |
| clock = [0.0] |
|
|
| def read(path: str, body: dict[str, object] | None) -> str: |
| if body is None: |
| return "vllm:num_requests_running 0\nvllm:num_requests_waiting 0" |
| clock[0] += 30.0 |
| return reply() |
|
|
| with patch.object(warmup.time, "monotonic", side_effect=lambda: clock[0]), redirect_stdout(io.StringIO()): |
| with self.assertRaisesRegex(ValueError, "before the concurrency-3 wave"): |
| warmup.run("test", read, 89.0) |
| clock[0] = 0.0 |
| with patch.object(warmup.time, "monotonic", side_effect=lambda: clock[0]), redirect_stdout(io.StringIO()): |
| with self.assertRaisesRegex(ValueError, "finished after"): |
| warmup.run("test", read, 200.0) |
| with patch.object(warmup.time, "monotonic", return_value=100.0): |
| self.assertAlmostEqual(warmup.request_timeout(100.01), 0.01, places=3) |
|
|
| def test_serve_wrapper_kills_a_hung_helper(self) -> None: |
| import os |
| import subprocess |
| import tempfile |
| root = Path(__file__).parents[3] |
| with tempfile.TemporaryDirectory() as fake: |
| shim = Path(fake) / "python3" |
| shim.write_text("#!/bin/sh\ntrap '' TERM\nwhile :; do sleep 1; done\n") |
| shim.chmod(0o755) |
| env = {**os.environ, "PATH": f"{fake}:{os.environ['PATH']}", "WARMUP_HARD_S": "1", "SERVED": "x"} |
| start = warmup.time.monotonic() |
| done = subprocess.run(["bash", str(root / "offload/launch/warmup_serve.sh"), "8000", "x"], env=env, |
| cwd=root, capture_output=True, timeout=30) |
| self.assertNotEqual(done.returncode, 0) |
| self.assertLess(warmup.time.monotonic() - start, 15) |
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
