chore: import upstream snapshot with attribution

python/ray/llm/tests/serve/benchmark/conftest.py

@@ -0,0 +1,123 @@
+"""Shared fixtures for benchmark tests."""
+
+import asyncio
+import json
+import socket
+import threading
+from collections.abc import Generator
+from typing import Any
+
+import aiohttp.web
+import pytest
+
+
+@pytest.fixture()
+def free_port() -> int:
+    """Find and return a free TCP port on localhost."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def make_sse_chunk(data: dict[str, Any]) -> bytes:
+    """Encode a dict as an SSE data line."""
+    return f"data: {json.dumps(data)}\n\n".encode()
+
+
+async def mock_chat_handler(
+    request: aiohttp.web.Request,
+) -> aiohttp.web.StreamResponse:
+    """Mock OpenAI streaming chat completions endpoint."""
+    body = await request.json()
+    assert body["model"] is not None
+    assert body["stream"] is True
+
+    resp = aiohttp.web.StreamResponse(
+        status=200,
+        headers={"Content-Type": "text/event-stream"},
+    )
+    await resp.prepare(request)
+
+    words = ["this", " is", " a", " test"]
+    for word in words:
+        chunk = {
+            "id": "chatcmpl-test",
+            "object": "chat.completion.chunk",
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {"content": word},
+                    "finish_reason": None,
+                }
+            ],
+        }
+        await resp.write(make_sse_chunk(chunk))
+
+    finish_chunk = {
+        "id": "chatcmpl-test",
+        "object": "chat.completion.chunk",
+        "choices": [
+            {
+                "index": 0,
+                "delta": {},
+                "finish_reason": "stop",
+            }
+        ],
+    }
+    await resp.write(make_sse_chunk(finish_chunk))
+
+    usage_chunk = {
+        "id": "chatcmpl-test",
+        "object": "chat.completion.chunk",
+        "choices": [],
+        "usage": {
+            "prompt_tokens": 10,
+            "completion_tokens": 4,
+            "total_tokens": 14,
+        },
+    }
+    await resp.write(make_sse_chunk(usage_chunk))
+
+    await resp.write(b"data: [DONE]\n\n")
+
+    return resp
+
+
+def _run_server(
+    loop: asyncio.AbstractEventLoop,
+    runner: aiohttp.web.AppRunner,
+    port: int,
+    started: threading.Event,
+) -> None:
+    asyncio.set_event_loop(loop)
+    site = aiohttp.web.TCPSite(runner, "127.0.0.1", port)
+    loop.run_until_complete(site.start())
+    started.set()
+    loop.run_forever()
+
+
+@pytest.fixture()
+def mock_server(free_port: int) -> Generator[str, None, None]:
+    """Start a mock SSE server in a background thread and yield its base URL."""
+    port = free_port
+    app = aiohttp.web.Application()
+    app.router.add_post("/v1/chat/completions", mock_chat_handler)
+
+    loop = asyncio.new_event_loop()
+    runner = aiohttp.web.AppRunner(app)
+    loop.run_until_complete(runner.setup())
+
+    started = threading.Event()
+    thread = threading.Thread(
+        target=_run_server, args=(loop, runner, port, started), daemon=True
+    )
+    thread.start()
+    assert started.wait(timeout=5), "Mock server failed to start"
+
+    base_url = f"http://127.0.0.1:{port}"
+    yield base_url
+
+    loop.call_soon_threadsafe(loop.stop)
+    thread.join(timeout=5)
+    loop.run_until_complete(runner.cleanup())
+    loop.close()

python/ray/llm/tests/serve/benchmark/test_engine.py

@@ -0,0 +1,403 @@
+"""Tests for benchmark engine components: TextGenerator, Conversation, conversation_factory, BenchmarkState."""
+
+import asyncio
+import sys
+
+import numpy as np
+import pytest
+from transformers import AutoTokenizer
+
+from ray.llm._internal.serve.benchmark.models import TurnMetric, WorkloadSpec
+from ray.llm._internal.serve.benchmark.runners import BenchmarkState
+from ray.llm._internal.serve.benchmark.text_gen import (
+    Conversation,
+    TextGenerator,
+    conversation_factory,
+)
+
+# ============================================================================
+# Fixtures
+# ============================================================================
+
+
+@pytest.fixture(scope="module")
+def tokenizer():
+    return AutoTokenizer.from_pretrained("gpt2")
+
+
+@pytest.fixture(scope="module")
+def text_gen(tokenizer):
+    return TextGenerator(tokenizer)
+
+
+def _make_spec(**overrides) -> WorkloadSpec:
+    """Create a resolved WorkloadSpec with sensible defaults."""
+    defaults = dict(
+        isl=1000,
+        hit_rate=0.7,
+        num_turns=3,
+        osl=50,
+        shared_system_prompt_ratio=0.5,
+        concurrency=4,
+        num_sessions=10,
+    )
+    defaults.update(overrides)
+    return WorkloadSpec(**defaults).resolve()
+
+
+# ============================================================================
+# TextGenerator
+# ============================================================================
+
+
+class TestTextGenerator:
+    def test_generate_exact_token_count(self, text_gen, tokenizer):
+        """Generated text should tokenize to exactly the requested count."""
+        np.random.seed(42)
+        for target in [1, 10, 50, 200]:
+            text = text_gen.generate(target)
+            actual = len(tokenizer.encode(text, add_special_tokens=False))
+            assert actual == target, f"Expected {target} tokens, got {actual}"
+
+    def test_generate_zero_tokens(self, text_gen):
+        assert text_gen.generate(0) == ""
+
+    def test_generate_negative_tokens(self, text_gen):
+        assert text_gen.generate(-5) == ""
+
+    def test_generate_token_ids_length(self, text_gen):
+        """generate_token_ids should return exactly N IDs."""
+        np.random.seed(42)
+        ids = text_gen.generate_token_ids(100)
+        assert len(ids) == 100
+
+    def test_generate_token_ids_in_vocab_range(self, text_gen, tokenizer):
+        np.random.seed(42)
+        ids = text_gen.generate_token_ids(500)
+        assert all(0 <= i < tokenizer.vocab_size for i in ids)
+
+    def test_generate_token_ids_zero(self, text_gen):
+        assert text_gen.generate_token_ids(0) == []
+
+    def test_generate_token_ids_negative(self, text_gen):
+        assert text_gen.generate_token_ids(-1) == []
+
+
+# ============================================================================
+# Conversation
+# ============================================================================
+
+
+class TestConversation:
+    def test_turn_zero_messages(self):
+        conv = Conversation(
+            session_id="s0",
+            system_prompt="You are helpful.",
+            user_messages=["Hello", "Follow up"],
+            num_turns=2,
+        )
+        msgs = conv.get_turn_messages(0)
+        assert len(msgs) == 2
+        assert msgs[0] == {"role": "system", "content": "You are helpful."}
+        assert msgs[1] == {"role": "user", "content": "Hello"}
+
+    def test_turn_one_with_injected_response(self):
+        conv = Conversation(
+            session_id="s0",
+            system_prompt="sys",
+            user_messages=["u0", "u1"],
+            num_turns=2,
+        )
+        conv.inject_assistant_response(0, "a0")
+        msgs = conv.get_turn_messages(1)
+        assert len(msgs) == 4
+        assert msgs[0] == {"role": "system", "content": "sys"}
+        assert msgs[1] == {"role": "user", "content": "u0"}
+        assert msgs[2] == {"role": "assistant", "content": "a0"}
+        assert msgs[3] == {"role": "user", "content": "u1"}
+
+    def test_turn_one_without_response_uses_placeholder(self):
+        conv = Conversation(
+            session_id="s0",
+            system_prompt="sys",
+            user_messages=["u0", "u1"],
+            num_turns=2,
+        )
+        msgs = conv.get_turn_messages(1)
+        assert msgs[2] == {"role": "assistant", "content": "(placeholder)"}
+
+    def test_no_system_prompt(self):
+        conv = Conversation(
+            session_id="s0",
+            system_prompt="",
+            user_messages=["u0"],
+            num_turns=1,
+        )
+        msgs = conv.get_turn_messages(0)
+        assert len(msgs) == 1
+        assert msgs[0]["role"] == "user"
+
+    def test_inject_sequential(self):
+        conv = Conversation(
+            session_id="s0",
+            system_prompt="sys",
+            user_messages=["u0", "u1", "u2"],
+            num_turns=3,
+        )
+        conv.inject_assistant_response(0, "a0")
+        conv.inject_assistant_response(1, "a1")
+        msgs = conv.get_turn_messages(2)
+        roles = [m["role"] for m in msgs]
+        assert roles == ["system", "user", "assistant", "user", "assistant", "user"]
+
+    def test_inject_overwrite(self):
+        conv = Conversation(
+            session_id="s0",
+            system_prompt="sys",
+            user_messages=["u0"],
+            num_turns=1,
+        )
+        conv.inject_assistant_response(0, "first")
+        conv.inject_assistant_response(0, "second")
+        assert conv._assistant_responses[0] == "second"
+
+    def test_inject_out_of_order_raises(self):
+        conv = Conversation(
+            session_id="s0",
+            system_prompt="sys",
+            user_messages=["u0", "u1", "u2"],
+            num_turns=3,
+        )
+        with pytest.raises(ValueError, match="Cannot inject"):
+            conv.inject_assistant_response(2, "skipped turn 0 and 1")
+
+
+# ============================================================================
+# conversation_factory
+# ============================================================================
+
+
+class TestConversationFactory:
+    def test_creates_valid_conversation(self, text_gen):
+        np.random.seed(42)
+        spec = _make_spec()
+        shared_text = text_gen.generate(spec.shared_s) if spec.shared_s > 0 else ""
+        conv = conversation_factory(0, spec, shared_text, text_gen)
+
+        assert conv.session_id == "session-000000"
+        assert conv.num_turns == spec.num_turns
+        assert len(conv.user_messages) == spec.num_turns
+
+    def test_user_messages_have_correct_token_count(self, text_gen, tokenizer):
+        np.random.seed(42)
+        spec = _make_spec()
+        shared_text = text_gen.generate(spec.shared_s) if spec.shared_s > 0 else ""
+        conv = conversation_factory(0, spec, shared_text, text_gen)
+
+        for msg in conv.user_messages:
+            token_count = len(tokenizer.encode(msg, add_special_tokens=False))
+            assert (
+                token_count == spec.user_tokens
+            ), f"Expected {spec.user_tokens} tokens, got {token_count}"
+
+    def test_cross_session_shared_prefix(self, text_gen):
+        """Two sessions with cross_sharing > 0 should share the same prefix."""
+        np.random.seed(42)
+        spec = _make_spec(shared_system_prompt_ratio=0.8)
+        shared_text = text_gen.generate(spec.shared_s) if spec.shared_s > 0 else ""
+
+        conv0 = conversation_factory(0, spec, shared_text, text_gen)
+        conv1 = conversation_factory(1, spec, shared_text, text_gen)
+
+        assert conv0.system_prompt.startswith(shared_text)
+        assert conv1.system_prompt.startswith(shared_text)
+
+    def test_zero_cross_sharing_no_shared_prefix(self, text_gen):
+        np.random.seed(42)
+        spec = _make_spec(shared_system_prompt_ratio=0.0, isl=2000, hit_rate=0.6)
+        assert spec.shared_s == 0
+
+    def test_session_id_format(self, text_gen):
+        np.random.seed(42)
+        spec = _make_spec()
+        shared_text = text_gen.generate(spec.shared_s) if spec.shared_s > 0 else ""
+        conv = conversation_factory(42, spec, shared_text, text_gen)
+        assert conv.session_id == "session-000042"
+
+
+# ============================================================================
+# BenchmarkState
+# ============================================================================
+
+
+class TestBenchmarkState:
+    def _run(self, coro):
+        return asyncio.get_event_loop().run_until_complete(coro)
+
+    @pytest.fixture(autouse=True)
+    def _setup_loop(self):
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        yield
+        loop.close()
+
+    def test_get_next_session_increments(self):
+        spec = _make_spec(num_sessions=5)
+        state = BenchmarkState(spec)
+        ids = [self._run(state.get_next_session()) for _ in range(5)]
+        assert ids == [0, 1, 2, 3, 4]
+
+    def test_get_next_session_exhausted(self):
+        spec = _make_spec(num_sessions=2)
+        state = BenchmarkState(spec)
+        self._run(state.get_next_session())
+        self._run(state.get_next_session())
+        assert self._run(state.get_next_session()) is None
+
+    def test_get_next_session_unlimited(self):
+        spec = _make_spec(
+            num_sessions=None, duration_s=60.0, request_rate=1.0, concurrency=None
+        )
+        state = BenchmarkState(spec)
+        ids = [self._run(state.get_next_session()) for _ in range(100)]
+        assert ids == list(range(100))
+
+    def test_record_metric_during_warmup(self):
+        spec = _make_spec()
+        state = BenchmarkState(spec)
+        metric = TurnMetric(
+            session_id="s0",
+            turn=0,
+            ttft_ms=1.0,
+            fc_ms=1.0,
+            itl_ms=1.0,
+            e2e_latency_ms=10.0,
+            input_tokens=100,
+            output_tokens=50,
+            start_time_ms=0.0,
+        )
+        self._run(state.record_metric(metric))
+        assert len(state.warmup_metrics) == 1
+        assert len(state.measured_metrics) == 0
+
+    def test_record_metric_after_warmup(self):
+        spec = _make_spec()
+        state = BenchmarkState(spec)
+        self._run(state.mark_warmup_complete())
+
+        metric = TurnMetric(
+            session_id="s0",
+            turn=0,
+            ttft_ms=1.0,
+            fc_ms=1.0,
+            itl_ms=1.0,
+            e2e_latency_ms=10.0,
+            input_tokens=100,
+            output_tokens=50,
+            start_time_ms=0.0,
+        )
+        self._run(state.record_metric(metric))
+        assert len(state.warmup_metrics) == 0
+        assert len(state.measured_metrics) == 1
+
+    def test_inflight_tracking(self):
+        spec = _make_spec()
+        state = BenchmarkState(spec)
+
+        self._run(state.track_inflight_start())
+        self._run(state.track_inflight_start())
+        assert state.inflight == 2
+        assert state.max_inflight == 2
+
+        self._run(state.track_inflight_end())
+        assert state.inflight == 1
+        assert state.max_inflight == 2
+
+    def test_mark_session_complete(self):
+        spec = _make_spec()
+        state = BenchmarkState(spec)
+        state.active_turns["session-000000"] = 0
+        state.conversations[0] = Conversation(
+            session_id="session-000000",
+            system_prompt="sys",
+            user_messages=["u0"],
+            num_turns=1,
+        )
+
+        self._run(state.mark_session_complete(0, "session-000000"))
+        assert state.completed_sessions == 1
+        assert "session-000000" not in state.active_turns
+        assert 0 not in state.conversations
+
+    def test_warmup_entropy_check_single_turn_spec(self):
+        """With num_turns=1, max_entropy=0 so warmup completes immediately."""
+        spec = _make_spec(num_turns=1, hit_rate=0.3)
+        state = BenchmarkState(spec)
+        state.active_turns["s0"] = 0
+        result = self._run(state.check_and_complete_warmup())
+        assert result is True
+        assert state.warmup_complete is True
+
+    def test_warmup_entropy_not_reached(self):
+        """All sessions on same turn -> entropy=0 -> warmup not complete."""
+        spec = _make_spec(num_turns=5)
+        state = BenchmarkState(spec)
+        for i in range(10):
+            state.active_turns[f"s{i}"] = 0
+        result = self._run(state.check_and_complete_warmup())
+        assert result is False
+        assert state.warmup_complete is False
+
+    def test_warmup_entropy_reached(self):
+        """Spread sessions across all turns -> high entropy -> warmup complete."""
+        spec = _make_spec(num_turns=4)
+        state = BenchmarkState(spec)
+        for i in range(20):
+            state.active_turns[f"s{i}"] = i % 4
+        result = self._run(state.check_and_complete_warmup())
+        assert result is True
+        assert state.warmup_complete is True
+
+    def test_has_remaining_sessions_with_limit(self):
+        spec = _make_spec(num_sessions=3)
+        state = BenchmarkState(spec)
+        assert self._run(state.has_remaining_sessions()) is True
+        self._run(state.get_next_session())
+        self._run(state.get_next_session())
+        assert self._run(state.has_remaining_sessions()) is True
+        self._run(state.get_next_session())
+        assert self._run(state.has_remaining_sessions()) is False
+
+    def test_has_remaining_sessions_unlimited(self):
+        spec = _make_spec(
+            num_sessions=None, duration_s=60.0, request_rate=1.0, concurrency=None
+        )
+        state = BenchmarkState(spec)
+        for _ in range(100):
+            self._run(state.get_next_session())
+        assert self._run(state.has_remaining_sessions()) is True
+
+    def test_get_inflight(self):
+        spec = _make_spec()
+        state = BenchmarkState(spec)
+        assert self._run(state.get_inflight()) == 0
+        self._run(state.track_inflight_start())
+        self._run(state.track_inflight_start())
+        assert self._run(state.get_inflight()) == 2
+        self._run(state.track_inflight_end())
+        assert self._run(state.get_inflight()) == 1
+
+    def test_get_stats(self):
+        spec = _make_spec()
+        state = BenchmarkState(spec)
+        stats = self._run(state.get_stats())
+        assert "warmup_metrics" in stats
+        assert "measured_metrics" in stats
+        assert "completed_sessions" in stats
+        assert "max_inflight" in stats
+        assert "warmup_complete" in stats
+        assert stats["warmup_complete"] is False
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/benchmark/test_interactive.py

@@ -0,0 +1,429 @@
+"""Tests for interactive command handler, _build_spec, and helpers."""
+
+import argparse
+import asyncio
+import json
+import os
+import sys
+import tempfile
+
+import pytest
+
+from ray.llm._internal.serve.benchmark.interactive import (
+    CommandHandler,
+    RuntimeState,
+    _build_spec,
+    _save_window_result,
+)
+from ray.llm._internal.serve.benchmark.metrics import percentile, summarize_metrics
+from ray.llm._internal.serve.benchmark.models import TurnMetric
+
+# ============================================================================
+# Helpers
+# ============================================================================
+
+
+def _make_args(**overrides) -> argparse.Namespace:
+    """Create a minimal args namespace for interactive mode."""
+    defaults = dict(
+        base_url="http://127.0.0.1:8000",
+        model="test-model",
+        tokenizer=None,
+        first_chunk_threshold=16,
+        isl=1000,
+        hit_rate=0.5,
+        num_turns=1,
+        osl=50,
+        shared_system_prompt_ratio=1.0,
+        save_result=None,
+        save_dir=None,
+    )
+    defaults.update(overrides)
+    return argparse.Namespace(**defaults)
+
+
+def _make_spec_for_handler():
+    """Build a resolved spec for use in CommandHandler tests."""
+    args = _make_args()
+    return _build_spec(args)
+
+
+def _make_handler(**runtime_overrides) -> CommandHandler:
+    """Create a CommandHandler with default state for testing."""
+    args = _make_args()
+    spec = _build_spec(args)
+    workload = {"spec": spec, "shared_system_text": ""}
+    runtime = RuntimeState(save_dir="/tmp/test_bench", **runtime_overrides)
+    return CommandHandler(
+        runtime=runtime,
+        workload=workload,
+        args=args,
+    )
+
+
+def _run(coro):
+    """Run an async coroutine synchronously."""
+    loop = asyncio.new_event_loop()
+    try:
+        return loop.run_until_complete(coro)
+    finally:
+        loop.close()
+
+
+def _make_metric(**overrides) -> TurnMetric:
+    defaults = dict(
+        session_id="s0",
+        turn=0,
+        ttft_ms=10.0,
+        fc_ms=20.0,
+        itl_ms=5.0,
+        e2e_latency_ms=100.0,
+        input_tokens=50,
+        output_tokens=20,
+        start_time_ms=0.0,
+    )
+    defaults.update(overrides)
+    return TurnMetric(**defaults)
+
+
+# ============================================================================
+# _summarize_metrics
+# ============================================================================
+
+
+class TestSummarizeMetrics:
+    def test_empty_metrics(self):
+        result = summarize_metrics([], 10.0)
+        assert result["requests"] == 0
+
+    def test_basic_metrics(self):
+        metrics = [
+            _make_metric(ttft_ms=10.0, output_tokens=20),
+            _make_metric(
+                session_id="s1",
+                ttft_ms=15.0,
+                fc_ms=25.0,
+                itl_ms=6.0,
+                e2e_latency_ms=120.0,
+                input_tokens=60,
+                output_tokens=25,
+                start_time_ms=100.0,
+            ),
+        ]
+        result = summarize_metrics(metrics, 2.0)
+        assert result["requests"] == 2
+        assert result["avg_ttft_ms"] == pytest.approx(12.5, abs=0.1)
+        assert result["throughput_tok_s"] == pytest.approx(22.5, abs=0.1)
+
+
+# ============================================================================
+# RuntimeState
+# ============================================================================
+
+
+class TestRuntimeState:
+    def test_default_init(self):
+        state = RuntimeState()
+        assert state.current_qps == 0.0
+        assert state.measurement_metrics == []
+        assert state.last_window_metrics == []
+        assert state.measurement_active is False
+
+
+# ============================================================================
+# percentile (from multiturn_bench)
+# ============================================================================
+
+
+class TestPercentile:
+    def test_empty(self):
+        assert percentile([], 50) == 0.0
+
+    def test_basic(self):
+        assert percentile([1.0, 2.0, 3.0, 4.0, 5.0], 50) == pytest.approx(3.0, abs=0.1)
+
+
+# ============================================================================
+# _build_spec
+# ============================================================================
+
+
+class TestBuildSpec:
+    def test_basic_build(self):
+        args = _make_args()
+        spec = _build_spec(args)
+        assert spec.isl == 1000
+        assert spec.hit_rate == 0.5
+        assert spec.num_turns == 1
+        assert spec.osl == 50
+
+    def test_overrides(self):
+        args = _make_args()
+        spec = _build_spec(args, {"isl": 2000, "osl": 100})
+        assert spec.isl == 2000
+        assert spec.osl == 100
+        assert spec.num_turns == 1  # unchanged
+
+    def test_override_num_turns(self):
+        args = _make_args(
+            isl=5600, hit_rate=0.8, osl=140, shared_system_prompt_ratio=1.0
+        )
+        spec = _build_spec(args, {"num_turns": 5})
+        assert spec.num_turns == 5
+
+    def test_override_hit_rate(self):
+        args = _make_args()
+        spec = _build_spec(args, {"hit_rate": 0.8})
+        assert spec.hit_rate == 0.8
+
+    def test_always_has_request_rate(self):
+        """_build_spec always sets request_rate=1.0 for spec resolution."""
+        args = _make_args()
+        spec = _build_spec(args)
+        assert spec.request_rate == 1.0
+
+
+# ============================================================================
+# _save_window_result
+# ============================================================================
+
+
+class TestSaveWindowResult:
+    def test_saves_valid_json(self):
+        args = _make_args()
+        spec = _build_spec(args)
+        metrics = [_make_metric(), _make_metric(session_id="s1")]
+
+        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
+            path = f.name
+
+        try:
+            _save_window_result(path, args, spec, metrics, 2.0, runtime_qps=5.0)
+            with open(path) as f:
+                data = json.load(f)
+            assert data["mode"] == "interactive_rate"
+            assert data["config"]["runtime_qps"] == 5.0
+            assert data["config"]["first_chunk_threshold"] == 16
+            assert "chunk_size" not in data["config"]
+            assert data["window"]["requests"] == 2
+            assert len(data["raw_metrics"]) == 2
+            # spec should use summary() format, not asdict
+            assert "effective_isl" in data["spec"]
+            assert "_u" not in data["spec"]
+        finally:
+            os.unlink(path)
+
+
+# ============================================================================
+# CommandHandler
+# ============================================================================
+
+
+class TestCommandHandler:
+    def test_empty_command(self):
+        h = _make_handler()
+        assert _run(h.handle("")) == "empty command"
+        assert _run(h.handle("  ")) == "empty command"
+
+    def test_help(self):
+        h = _make_handler()
+        resp = _run(h.handle("help"))
+        assert "rate <qps>" in resp
+        assert "workload" in resp
+
+    def test_unknown_command(self):
+        h = _make_handler()
+        resp = _run(h.handle("foobar"))
+        assert "Unknown command" in resp
+
+    # --- rate ---
+    def test_rate_set(self):
+        h = _make_handler()
+        resp = _run(h.handle("rate 5.5"))
+        assert "5.500" in resp
+        assert h.runtime.current_qps == 5.5
+        assert h.rate_changed.is_set()
+
+    def test_rate_zero(self):
+        h = _make_handler()
+        resp = _run(h.handle("rate 0"))
+        assert "0.000" in resp
+        assert h.runtime.current_qps == 0.0
+
+    def test_rate_negative(self):
+        h = _make_handler()
+        resp = _run(h.handle("rate -1"))
+        assert "non-negative" in resp
+
+    def test_rate_invalid(self):
+        h = _make_handler()
+        resp = _run(h.handle("rate abc"))
+        assert "non-negative" in resp
+
+    def test_rate_missing_arg(self):
+        h = _make_handler()
+        resp = _run(h.handle("rate"))
+        assert "Usage" in resp
+
+    # --- start ---
+    def test_start(self):
+        h = _make_handler()
+        resp = _run(h.handle("start"))
+        assert "started" in resp.lower()
+        assert h.runtime.measurement_active is True
+        assert h.runtime.measurement_start_ns is not None
+        assert h.runtime.measurement_metrics == []
+
+    # --- measure ---
+    def test_measure(self):
+        h = _make_handler()
+        resp = _run(h.handle("measure 100"))
+        assert "100" in resp
+        assert h.runtime.measurement_active is True
+        assert h.runtime.measurement_target_requests == 100
+
+    def test_measure_invalid(self):
+        h = _make_handler()
+        assert "positive integer" in _run(h.handle("measure 0"))
+        assert "positive integer" in _run(h.handle("measure -5"))
+        assert "positive integer" in _run(h.handle("measure abc"))
+
+    def test_measure_missing_arg(self):
+        h = _make_handler()
+        assert "Usage" in _run(h.handle("measure"))
+
+    # --- stop ---
+    def test_stop_when_active(self):
+        h = _make_handler()
+        _run(h.handle("start"))
+        h.runtime.measurement_metrics = [_make_metric()]
+        resp = _run(h.handle("stop"))
+        assert "stopped" in resp.lower()
+        assert h.runtime.measurement_active is False
+        assert len(h.runtime.last_window_metrics) == 1
+
+    def test_stop_when_inactive(self):
+        h = _make_handler()
+        resp = _run(h.handle("stop"))
+        assert "not active" in resp.lower()
+
+    # --- status ---
+    def test_status(self):
+        h = _make_handler(current_qps=3.0)
+        h.runtime.total_completed = 42
+        resp = _run(h.handle("status"))
+        assert "qps=3.00" in resp
+        assert "completed=42" in resp
+        assert "workload:" in resp
+
+    def test_status_clears_notice(self):
+        h = _make_handler()
+        h.runtime.last_notice = "some notice"
+        resp = _run(h.handle("status"))
+        assert "some notice" in resp
+        # Second call should not have the notice
+        resp2 = _run(h.handle("status"))
+        assert "some notice" not in resp2
+
+    # --- save-dir ---
+    def test_save_dir(self):
+        h = _make_handler()
+        resp = _run(h.handle("save-dir /tmp/newdir"))
+        assert "/tmp/newdir" in resp
+        assert h.runtime.save_dir == "/tmp/newdir"
+
+    def test_save_dir_missing_arg(self):
+        h = _make_handler()
+        assert "Usage" in _run(h.handle("save-dir"))
+
+    # --- save ---
+    def test_save_no_data(self):
+        h = _make_handler()
+        resp = _run(h.handle("save"))
+        assert "No measured" in resp
+
+    def test_save_with_data(self):
+        h = _make_handler()
+        h.runtime.last_window_metrics = [_make_metric()]
+        h.runtime.last_window_elapsed_s = 1.0
+
+        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
+            path = f.name
+        try:
+            resp = _run(h.handle(f"save {path}"))
+            assert "Saved" in resp
+            with open(path) as f:
+                data = json.load(f)
+            assert data["window"]["requests"] == 1
+        finally:
+            os.unlink(path)
+
+    # --- workload (query) ---
+    def test_workload_query(self):
+        h = _make_handler()
+        resp = _run(h.handle("workload"))
+        assert "isl=1000" in resp
+        assert "osl=50" in resp
+
+    # --- workload (update) ---
+    def test_workload_update(self):
+        h = _make_handler()
+        resp = _run(h.handle("workload isl=2000 osl=100"))
+        assert "Workload updated" in resp
+        assert h.workload["spec"].isl == 2000
+        assert h.workload["spec"].osl == 100
+        assert h.workload_changed.is_set()
+
+    def test_workload_partial_update(self):
+        h = _make_handler()
+        _run(h.handle("workload isl=2000"))
+        assert h.workload["spec"].isl == 2000
+        assert h.workload["spec"].osl == 50  # unchanged
+
+    def test_workload_aliases(self):
+        h = _make_handler()
+        resp = _run(h.handle("workload hit-rate=0.8"))
+        assert "Workload updated" in resp
+        assert h.workload["spec"].hit_rate == 0.8
+
+    def test_workload_bad_token(self):
+        h = _make_handler()
+        resp = _run(h.handle("workload notaparam"))
+        assert "Error" in resp
+        assert "bad token" in resp
+
+    def test_workload_unknown_param(self):
+        h = _make_handler()
+        resp = _run(h.handle("workload foo=123"))
+        assert "Error" in resp
+        assert "unknown param" in resp
+
+    def test_workload_invalid_value(self):
+        h = _make_handler()
+        resp = _run(h.handle("workload isl=abc"))
+        assert "Error" in resp
+        assert "invalid value" in resp
+
+    # --- quit ---
+    def test_quit(self):
+        h = _make_handler()
+        resp = _run(h.handle("quit"))
+        assert "Stopping" in resp
+        assert h.stop_event.is_set()
+
+    def test_exit(self):
+        h = _make_handler()
+        resp = _run(h.handle("exit"))
+        assert "Stopping" in resp
+        assert h.stop_event.is_set()
+
+    # --- case insensitivity ---
+    def test_case_insensitive(self):
+        h = _make_handler()
+        assert "rate <qps>" in _run(h.handle("HELP"))
+        resp = _run(h.handle("RATE 1.0"))
+        assert "1.000" in resp
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/benchmark/test_smoke.py

@@ -0,0 +1,60 @@
+"""Tests for benchmark smoke mode."""
+
+import sys
+import types
+
+import aiohttp
+import pytest
+
+from ray.llm._internal.serve.benchmark.http_client import send_chat_completion
+from ray.llm._internal.serve.benchmark.models import TurnResult
+from ray.llm._internal.serve.benchmark.runners import run_smoke
+
+
+@pytest.mark.asyncio
+async def test_send_chat_completion(mock_server: str) -> None:
+    """send_chat_completion should parse SSE stream and return correct TurnResult."""
+    messages = [{"role": "user", "content": "hello"}]
+    async with aiohttp.ClientSession() as session:
+        result = await send_chat_completion(
+            session=session,
+            base_url=mock_server,
+            model="test-model",
+            messages=messages,
+            first_chunk_threshold=2,
+        )
+
+    assert isinstance(result, TurnResult)
+    assert result.generated_text == "this is a test"
+    assert result.input_tokens == 10
+    assert result.output_tokens == 4
+    assert result.ttft_ms > 0
+    assert result.e2e_latency_ms > 0
+    assert result.fc_ms > 0
+    assert isinstance(result.itl_ms, float)
+
+
+def test_run_smoke(mock_server: str) -> None:
+    """run_smoke should return 0 and print JSON metrics."""
+    args = types.SimpleNamespace(
+        base_url=mock_server,
+        model="test-model",
+        first_chunk_threshold=2,
+    )
+    exit_code = run_smoke(args)
+    assert exit_code == 0
+
+
+def test_run_smoke_connection_error(free_port: int) -> None:
+    """run_smoke should return 1 when the server is unreachable."""
+    args = types.SimpleNamespace(
+        base_url=f"http://127.0.0.1:{free_port}",
+        model="test-model",
+        first_chunk_threshold=16,
+    )
+    exit_code = run_smoke(args)
+    assert exit_code == 1
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/benchmark/test_workload.py

@@ -0,0 +1,200 @@
+"""Tests for WorkloadSpec solver."""
+import sys
+
+import pytest
+
+from ray.llm._internal.serve.benchmark.models import WorkloadSpec
+
+
+class TestWorkloadSpec:
+    """Test WorkloadSpec resolution and invariants."""
+
+    def test_single_turn(self):
+        """Single turn: ISL round-trip."""
+        spec = WorkloadSpec(
+            isl=1000,
+            hit_rate=0.5,
+            num_turns=1,
+            osl=100,
+            shared_system_prompt_ratio=1.0,
+            concurrency=1,
+            num_sessions=1,
+        ).resolve()
+        assert abs(spec.effective_isl - 1000) < 2
+        assert abs(spec.effective_h - 0.5) < 0.02
+
+    def test_multi_turn_full_sharing(self):
+        """Multi-turn with full cross-sharing."""
+        spec = WorkloadSpec(
+            isl=5600,
+            hit_rate=0.8,
+            num_turns=5,
+            osl=140,
+            shared_system_prompt_ratio=1.0,
+            concurrency=8,
+            num_sessions=100,
+        ).resolve()
+        assert abs(spec.effective_isl - 5600) < 50
+        assert abs(spec.effective_h - 0.8) < 0.02
+
+    def test_multi_turn_partial_sharing(self):
+        """Multi-turn with partial cross-sharing."""
+        spec = WorkloadSpec(
+            isl=3000,
+            hit_rate=0.6,
+            num_turns=3,
+            osl=100,
+            shared_system_prompt_ratio=0.5,
+            concurrency=4,
+            num_sessions=50,
+        ).resolve()
+        assert abs(spec.effective_isl - 3000) < 50
+        assert abs(spec.effective_h - 0.6) < 0.05
+
+    def test_rate_mode_with_duration(self):
+        """Rate mode with duration (no num_sessions required)."""
+        spec = WorkloadSpec(
+            isl=2000,
+            hit_rate=0.7,
+            num_turns=3,
+            osl=100,
+            shared_system_prompt_ratio=0.5,
+            request_rate=10.0,
+            duration_s=60.0,
+        ).resolve()
+        assert spec.request_rate == 10.0
+
+    def test_missing_isl_raises(self):
+        """Missing ISL should raise."""
+        with pytest.raises(ValueError, match="--isl.*--hit-rate"):
+            WorkloadSpec(
+                hit_rate=0.5,
+                num_turns=1,
+                osl=100,
+                concurrency=1,
+                num_sessions=1,
+            ).resolve()
+
+    def test_missing_traffic_raises(self):
+        """Missing both concurrency and request_rate should raise."""
+        with pytest.raises(ValueError, match="--concurrency.*--request-rate"):
+            WorkloadSpec(
+                isl=1000,
+                hit_rate=0.0,
+                num_turns=1,
+                osl=100,
+                num_sessions=10,
+            ).resolve()
+
+    def test_both_traffic_modes_raises(self):
+        """Setting both concurrency and request_rate should raise."""
+        with pytest.raises(ValueError, match="Cannot specify both"):
+            WorkloadSpec(
+                isl=1000,
+                hit_rate=0.0,
+                num_turns=1,
+                osl=100,
+                concurrency=4,
+                request_rate=10.0,
+                num_sessions=10,
+            ).resolve()
+
+    def test_per_turn_tokens_monotonic(self):
+        """Input tokens should increase monotonically across turns."""
+        spec = WorkloadSpec(
+            isl=5600,
+            hit_rate=0.8,
+            num_turns=5,
+            osl=140,
+            concurrency=8,
+            num_sessions=100,
+        ).resolve()
+        for k in range(2, spec.num_turns + 1):
+            assert spec.turn_input_tokens(k) > spec.turn_input_tokens(k - 1)
+
+    def test_summary_keys(self):
+        """Summary dict should have expected keys."""
+        spec = WorkloadSpec(
+            isl=1000,
+            hit_rate=0.5,
+            num_turns=2,
+            osl=100,
+            concurrency=4,
+            num_sessions=10,
+        ).resolve()
+        s = spec.summary()
+        expected_keys = {
+            "num_sessions",
+            "duration_s",
+            "num_turns",
+            "osl",
+            "think_time",
+            "concurrency",
+            "request_rate",
+            "shared_system_prompt_ratio",
+            "user_tokens_per_turn",
+            "system_prompt_tokens",
+            "shared_system_prompt",
+            "unique_system_prompt",
+            "effective_isl",
+            "effective_hit_rate",
+            "per_turn",
+        }
+        assert expected_keys.issubset(s.keys())
+        assert len(s["per_turn"]) == 2
+
+    def test_infeasible_user_tokens_suggests_lower_hit_rate_and_boundary_resolves(self):
+        """user_tokens < 0.5: error names user_tokens and suggests a feasible hit_rate cap."""
+        base = dict(
+            isl=1000,
+            hit_rate=1.0,
+            num_turns=5,
+            osl=100,
+            shared_system_prompt_ratio=1.0,
+            concurrency=1,
+            num_sessions=1,
+        )
+        with pytest.raises(ValueError) as exc_info:
+            WorkloadSpec(**base).resolve()
+        msg = str(exc_info.value)
+        assert "user_tokens" in msg
+        assert "hit_rate <=" in msg
+        # Suggested boundary (from _feasibility_suggestions) must resolve when applied.
+        WorkloadSpec(**{**base, "hit_rate": 0.74}).resolve()
+
+    def test_infeasible_sys_tokens_suggests_num_turns_and_boundary_resolves(self):
+        """sys_tokens < -0.5: error names sys_tokens; reducing num_turns can fix."""
+        base = dict(
+            isl=500,
+            hit_rate=0.99,
+            num_turns=8,
+            osl=200,
+            shared_system_prompt_ratio=0.9,
+            concurrency=1,
+            num_sessions=1,
+        )
+        with pytest.raises(ValueError) as exc_info:
+            WorkloadSpec(**base).resolve()
+        msg = str(exc_info.value)
+        assert "sys_tokens" in msg
+        assert "num_turns <=" in msg
+        WorkloadSpec(**{**base, "num_turns": 5}).resolve()
+
+    def test_infeasible_no_single_parameter_fix(self):
+        """Some parameter sets are infeasible with no one-dimensional remedy."""
+        kw = dict(
+            isl=442,
+            hit_rate=0.39,
+            num_turns=15,
+            osl=216,
+            shared_system_prompt_ratio=0.04,
+            concurrency=1,
+            num_sessions=1,
+        )
+        with pytest.raises(ValueError) as exc_info:
+            WorkloadSpec(**kw).resolve()
+        assert "no single-parameter fix found" in str(exc_info.value)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))