chore: import upstream snapshot with attribution

python/ray/llm/tests/serve/cpu/deployments/utils/direct_streaming_utils.py

@@ -0,0 +1,68 @@
+"""
+Shared helpers for direct-streaming session-affinity tests.
+"""
+
+import httpx
+import pytest
+
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray.llm._internal.serve.constants import RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING
+from ray.serve._private.constants import RAY_SERVE_ENABLE_HA_PROXY, SERVE_SESSION_ID
+from ray.serve._private.test_utils import check_running, get_application_url
+from ray.serve.config import RequestRouterConfig
+
+CONSISTENT_HASH_ROUTER = (
+    "ray.serve.experimental.consistent_hash_router:ConsistentHashRouter"
+)
+
+# Skip unless the direct-streaming + HAProxy env is set
+requires_direct_streaming = pytest.mark.skipif(
+    not (RAY_SERVE_ENABLE_HA_PROXY and RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING),
+    reason="Direct streaming requires RAY_SERVE_ENABLE_HA_PROXY=1 and "
+    "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING=1.",
+)
+
+
+def consistent_hash_deployment_config() -> dict:
+    return {
+        "num_replicas": 4,
+        "ray_actor_options": {"num_cpus": 0.1},
+        "request_router_config": RequestRouterConfig(
+            request_router_class=CONSISTENT_HASH_ROUTER,
+            request_router_kwargs={
+                "num_virtual_nodes": 100,
+                "num_fallback_replicas": 2,
+            },
+        ),
+    }
+
+
+def run_app_through_haproxy(app, timeout_s: int = 60) -> str:
+    """Run ``app`` and return its (HAProxy) URL once all replicas are RUNNING."""
+    serve.run(app)
+    wait_for_condition(check_running, timeout=timeout_s)
+    return get_application_url(use_localhost=True)
+
+
+def session_chat_response(base_url: str, session_id: str, model: str = "test-model"):
+    """POST a one-token chat request carrying ``session_id`` through HAProxy.
+
+    Asserts the request succeeded and the session id survived the HAProxy hop to
+    the serving replica. Returns the response so callers can read the serving
+    replica from the ``x-replica-id`` header (and, for P/D, the prefill replica
+    from ``kv_transfer_params.remote_engine_id``).
+    """
+    resp = httpx.post(
+        f"{base_url}/v1/chat/completions",
+        json={
+            "model": model,
+            "messages": [{"role": "user", "content": "hi"}],
+            "max_tokens": 1,
+        },
+        headers={SERVE_SESSION_ID: session_id},
+        timeout=30,
+    )
+    assert resp.status_code == 200, resp.text
+    assert resp.headers["x-serve-session-id"] == session_id
+    return resp

python/ray/llm/tests/serve/cpu/deployments/utils/test_batcher.py

@@ -0,0 +1,236 @@
+import asyncio
+import sys
+import time
+from typing import List, Optional
+
+import numpy as np
+import pytest
+
+from ray.llm._internal.serve.constants import MODEL_RESPONSE_BATCH_TIMEOUT_MS
+from ray.llm._internal.serve.utils.batcher import Batcher
+
+TEXT_VALUE = "foo"
+FINAL_TEXT_VALUE = "bar"
+
+
+async def fake_generator():
+    """Returns 100 responses with no delay"""
+    for _i in range(100):
+        yield dict(num_generated_tokens=1, generated_text=TEXT_VALUE)
+
+
+async def fake_generator_slow(num_batches: int):
+    """Returns 100 responses with small delay.
+
+    Delay is set such that the responses are batched into roughly num_batches
+    batches.
+    """
+
+    for _i in range(100):
+        await asyncio.sleep(MODEL_RESPONSE_BATCH_TIMEOUT_MS / 1000 / num_batches)
+        yield dict(num_generated_tokens=1, generated_text=TEXT_VALUE)
+
+
+async def fake_generator_slow_last_return_immediate():
+    """Returns 11 responses with small delay, aside from the last one which is immediate"""
+    for _i in range(10):
+        await asyncio.sleep(MODEL_RESPONSE_BATCH_TIMEOUT_MS / 1000)
+        yield dict(num_generated_tokens=1, generated_text=TEXT_VALUE)
+    yield dict(num_generated_tokens=1, generated_text=FINAL_TEXT_VALUE)
+
+
+async def count_interval_ms_from_stream(stream) -> list[float]:
+    output_intervals: list[float] = []
+    start = None
+    async for _ in stream:
+        if start is None:
+            start = time.perf_counter()
+        else:
+            end = time.perf_counter()
+            output_intervals.append((end - start) * 1e3)
+            start = end
+    return output_intervals
+
+
+class TestBatcher(Batcher):
+    def _merge_results(self, results: List[dict]) -> dict:
+        merged_result = {"num_generated_tokens": 0, "generated_text": ""}
+        for result in results:
+            for key, value in result.items():
+                merged_result[key] += value
+        return merged_result
+
+
+class TestBatching:
+    @pytest.mark.asyncio
+    async def test_batch(self):
+        count = 0
+        batcher = TestBatcher(fake_generator())
+        async for x in batcher.stream():
+            count += 1
+            assert x["num_generated_tokens"] == 100
+            assert x["generated_text"] == TEXT_VALUE * 100
+
+        # Should only have been called once
+        assert count == 1
+        assert batcher.queue.empty()
+
+    @pytest.mark.asyncio
+    async def test_batch_timing(self):
+        count = 0
+        batcher = TestBatcher(fake_generator_slow(num_batches=10))
+        async for _x in batcher.stream():
+            count += 1
+
+        assert 9 <= count <= 12, (
+            "Count should have been called between 9 and 12 times, "
+            "because each iteration takes 1/10th of an interval to yield."
+        )
+        assert batcher.queue.empty()
+
+    @pytest.mark.asyncio
+    async def test_batch_last_return_is_immediate(self):
+        """Test that we don't wait the entire interval for
+        the last response if it returns quickly."""
+        count = 0
+        token_count = 0
+        batcher = TestBatcher(fake_generator_slow_last_return_immediate())
+        last_response = None
+        async for _x in batcher.stream():
+            count += 1
+            token_count += _x["num_generated_tokens"]
+            last_response = _x
+
+        assert (
+            last_response["generated_text"] == TEXT_VALUE + FINAL_TEXT_VALUE
+        ), "the last generated response should be batched with previous one"
+        assert token_count == 11, "token_count should be exactly 11"
+        assert (
+            count == 10
+        ), "Count should have been called exactly 10 times (as many as we generated - 1)"
+        assert batcher.queue.empty()
+
+    @pytest.mark.asyncio
+    async def test_batch_no_interval(self):
+        """Check that the class creates only one batch if there's no interval."""
+
+        batcher = TestBatcher(fake_generator_slow(num_batches=10), interval_ms=None)
+
+        count = 0
+        async for _x in batcher.stream():
+            count += 1
+
+        assert count == 1
+        assert batcher.queue.empty()
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("interval_ms", [100, None])
+    async def test_exception_propagation(self, interval_ms: Optional[float]):
+        """Test that exceptions are propagated correctly to parent."""
+
+        async def generator_should_raise():
+            for _i in range(100):
+                await asyncio.sleep(0.01)
+                yield dict(num_generated_tokens=1, generated_text=TEXT_VALUE)
+                raise ValueError()
+
+        count = 0
+        batched = TestBatcher(generator_should_raise(), interval_ms=interval_ms)
+
+        async def parent():
+            nonlocal count
+            nonlocal batched
+            async for _x in batched.stream():
+                count += 1
+
+        task = asyncio.create_task(parent())
+        await asyncio.sleep(0.2)
+
+        with pytest.raises(ValueError):
+            task.result()
+        assert count == 1
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("interval_ms", [100, None])
+    @pytest.mark.parametrize("to_cancel", ["parent", "inner", "stream"])
+    async def test_cancellation(self, interval_ms: Optional[float], to_cancel: str):
+        """There are 3 ways cancellation can happen:
+        1. The parent is cancelled
+        2. The generator is cancelled
+        3. The stream task is directly cancelled.
+
+        Make sure all associated tasks are cancelled in each instance.
+        """
+
+        async def generator_should_raise():
+            with pytest.raises(asyncio.CancelledError):
+                for _i in range(100):
+                    await asyncio.sleep(0.01)
+                    yield dict(num_generated_tokens=1, generated_text=TEXT_VALUE)
+                    if to_cancel == "inner":
+                        raise asyncio.CancelledError()
+
+        batched = TestBatcher(generator_should_raise(), interval_ms=interval_ms)
+
+        async def parent():
+            nonlocal batched
+            async for _x in batched.stream():
+                pass
+
+        task = asyncio.create_task(parent())
+        await asyncio.sleep(0.2)
+
+        cancel_task = {
+            "parent": task,
+            "stream": batched.read_task,
+        }.get(to_cancel)
+
+        if cancel_task:
+            assert not task.done()
+            assert not batched.read_task.done()
+            cancel_task.cancel()
+
+        await asyncio.sleep(0.3)
+        assert batched.read_task.done(), "Read task should be completed"
+        assert task.done(), "All tasks should be done"
+
+        # Inner task is checked automatically with pytest.raises
+
+    @pytest.mark.asyncio
+    async def test_stable_streaming(self):
+        """Test that the batcher does not add jitter to the stream when interval_ms is 0"""
+
+        async def generator():
+            for i in range(100):
+                await asyncio.sleep(0.01)
+                yield i
+
+        concurrency = 10
+
+        output_intervals = await asyncio.gather(
+            *[
+                count_interval_ms_from_stream(
+                    Batcher(generator(), interval_ms=0).stream()
+                )
+                for _ in range(concurrency)
+            ]
+        )
+        mean_batcher_interval = np.mean(output_intervals)
+        std_batcher_interval = np.std(output_intervals)
+
+        generator_intervals = await asyncio.gather(
+            *[count_interval_ms_from_stream(generator()) for _ in range(concurrency)]
+        )
+        mean_generator_interval = np.mean(generator_intervals)
+        std_generator_interval = np.std(generator_intervals)
+
+        assert np.isclose(
+            mean_batcher_interval, mean_generator_interval, rtol=0.1
+        ), f"{mean_batcher_interval=}, {mean_generator_interval=}"
+        assert np.isclose(
+            std_batcher_interval, std_generator_interval, atol=0.1
+        ), f"{std_batcher_interval=}, {std_generator_interval=}"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/cpu/deployments/utils/test_broadcast.py

@@ -0,0 +1,163 @@
+import sys
+import time
+
+import pytest
+
+import ray
+from ray import serve
+from ray.llm._internal.serve.utils.broadcast import broadcast
+
+
+# Define a simple deployment for testing
+@serve.deployment(num_replicas=2)
+class MockLLMDeployment:
+    def __init__(self):
+        self.reset_count = 0
+        self.id = id(self)
+
+    async def reset_prefix_cache(self):
+        self.reset_count += 1
+        return self.id, self.reset_count
+
+    async def get_reset_count(self):
+        return self.id, self.reset_count
+
+    async def echo(self, msg, repeat=1):
+        return f"{self.id}:{msg * repeat}"
+
+    async def self_destruct(self):
+        """Kill this replica's actor. Used for testing dead replica handling."""
+        import os
+
+        os._exit(1)
+
+
+@pytest.fixture(scope="module")
+def serve_instance():
+    # Start ray and serve once for the module
+    if not ray.is_initialized():
+        ray.init()
+    yield
+    serve.shutdown()
+    ray.shutdown()
+
+
+@pytest.fixture
+def mock_handle(serve_instance, request):
+    # Ensure deployment is up and running.
+    # serve.run waits for the deployment to be ready by default unless _blocking=False.
+    app_name = f"mock-llm-{request.node.name}"
+    route_prefix = f"/{app_name}"
+    handle = serve.run(
+        MockLLMDeployment.bind(), name=app_name, route_prefix=route_prefix
+    )
+    yield handle
+    serve.delete(app_name, _blocking=True)
+
+
+@pytest.mark.asyncio
+async def test_dispatch_basic(mock_handle):
+    """Test basic dispatch without combine."""
+    # We can use get_reset_count which doesn't modify state
+    results = broadcast(mock_handle, "get_reset_count")
+
+    assert len(results) == 2
+    # Verify we got unique IDs back
+    ids = {r[0] for r in results}
+    assert len(ids) == 2
+
+
+@pytest.mark.asyncio
+async def test_dispatch_with_combine(mock_handle):
+    """Test dispatch with a combine function."""
+    # First, increment count so we have something to sum
+    broadcast(mock_handle, "reset_prefix_cache")
+
+    def sum_counts(results):
+        # results is list of (id, count)
+        return sum(r[1] for r in results)
+
+    # Get counts using dispatch and combine
+    total_count = broadcast(mock_handle, "get_reset_count", combine=sum_counts)
+
+    # We have 2 replicas, each should have reset_count=1 after one reset call
+    assert total_count == 2
+    assert isinstance(total_count, int)
+
+
+@pytest.mark.asyncio
+async def test_dispatch_args_kwargs(mock_handle):
+    """Test dispatch passing args and kwargs."""
+    results = broadcast(mock_handle, "echo", args=("hello",), kwargs={"repeat": 2})
+
+    assert len(results) == 2
+    for r in results:
+        # Format is "id:msg"
+        msg_part = r.split(":")[1]
+        assert msg_part == "hellohello"
+
+
+@pytest.mark.asyncio
+async def test_dispatch_callable_args(mock_handle):
+    """Test dispatch with callable args generator."""
+
+    def arg_gen(replica):
+        # replica has unique_id or similar
+        return (f"msg-{replica.unique_id}",)
+
+    results = broadcast(mock_handle, "echo", args=arg_gen)
+
+    assert len(results) == 2
+    msgs = set()
+    for r in results:
+        msg_part = r.split(":")[1]
+        msgs.add(msg_part)
+
+    assert len(msgs) == 2
+    for msg in msgs:
+        assert msg.startswith("msg-")
+
+
+@pytest.mark.asyncio
+async def test_dispatch_handles_dead_replica(serve_instance, request):
+    """Test that dispatch gracefully handles a dead replica.
+
+    This test verifies that if one replica dies, dispatch still completes
+    successfully and returns results from the remaining live replicas.
+    """
+    app_name = f"mock-llm-{request.node.name}"
+    route_prefix = f"/{app_name}"
+
+    # Deploy with 2 replicas
+    handle = serve.run(
+        MockLLMDeployment.bind(), name=app_name, route_prefix=route_prefix
+    )
+
+    # First, verify dispatch works with all replicas alive
+    results_before = broadcast(handle, "get_reset_count")
+    assert len(results_before) == 2, "Should have 2 results from 2 replicas"
+
+    # Kill one replica by calling self_destruct through the handle.
+    # This sends an RPC to one replica which will kill itself.
+    # We use options to not wait for response since the actor will die.
+    try:
+        handle.self_destruct.remote()
+    except Exception:
+        # The call may raise if the actor dies mid-request
+        pass
+
+    # Give Serve a moment to detect the dead replica
+    time.sleep(2)
+
+    # Dispatch should still work with the remaining replica(s)
+    # The dead replica will be skipped (ValueError caught in dispatch)
+    results_after = broadcast(handle, "get_reset_count")
+
+    # Should get at least 1 result from the surviving replica
+    # (The killed replica may or may not be in the replica set depending
+    # on timing of Serve's failure detection)
+    assert len(results_after) >= 1, "Should have at least 1 result from live replica"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))