chore: import upstream snapshot with attribution

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__]))