chore: import upstream snapshot with attribution

python/ray/llm/tests/serve/utils/__init__.py

@@ -0,0 +1 @@
+# Testing utilities for Ray LLM serve tests

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

@@ -0,0 +1,12 @@
+import pytest
+
+import ray
+from ray.cluster_utils import Cluster
+
+
+@pytest.fixture
+def ray_start_cluster():
+    cluster = Cluster()
+    yield cluster
+    ray.shutdown()
+    cluster.shutdown()

python/ray/llm/tests/serve/utils/test_error_log_ratelimiter.py

@@ -0,0 +1,233 @@
+"""Tests for the fatal-engine-error log rate limiter in server_utils."""
+
+import sys
+from contextlib import contextmanager
+from unittest.mock import MagicMock, patch
+
+import pytest
+from vllm.v1.engine.exceptions import EngineDeadError
+
+from ray.llm._internal.serve.utils import server_utils
+
+COOLDOWN = 10.0
+
+
+def _make_fatal_error() -> Exception:
+    return EngineDeadError("engine is dead")
+
+
+def _make_ray_wrapped_fatal_error() -> Exception:
+    """Simulate Ray's RayTaskError wrapping."""
+
+    class WrappedEngineDeadError(EngineDeadError):
+        pass
+
+    return WrappedEngineDeadError("engine is dead")
+
+
+@contextmanager
+def _patched_logger_and_time():
+    """Patch ``time.monotonic`` and the server_utils logger.
+
+    Yields (mock_logger, mock_time) via mock_time.return_value = <new_value>.
+    """
+    mock_time = MagicMock(return_value=100.0)
+    with (
+        patch.object(server_utils, "logger") as mock_logger,
+        patch("time.monotonic", mock_time),
+    ):
+        yield mock_logger, mock_time
+
+
+@pytest.fixture
+def handler():
+    return server_utils._FatalEngineErrorLogHandler(cooldown_s=COOLDOWN)
+
+
+class TestIsFatalEngineError:
+    """Tests if fatal engine errors are identified and other errors are untouched."""
+
+    def test_bare_engine_dead_error(self):
+        assert server_utils._is_fatal_engine_error(_make_fatal_error())
+
+    def test_ray_wrapped_engine_dead_error(self):
+        assert server_utils._is_fatal_engine_error(_make_ray_wrapped_fatal_error())
+
+    def test_non_fatal_exception(self):
+        assert not server_utils._is_fatal_engine_error(ValueError("hi"))
+
+    def test_runtime_error_is_not_fatal(self):
+        assert not server_utils._is_fatal_engine_error(RuntimeError("error"))
+
+
+class TestFatalEngineErrorLogHandler:
+    """Checks that the handler suppresses duplicate logs on fatal errors
+    within cooldown_s and keeps other logs undisturbed.
+    """
+
+    def test_first_fatal_logs_full_traceback(self, handler):
+        exc = _make_fatal_error()
+        with _patched_logger_and_time() as (mock_logger, _):
+            handler.log(exc, request_id="req-1", status_code=500)
+
+        mock_logger.error.assert_called_once()
+        assert mock_logger.error.call_args.kwargs["exc_info"] is exc
+
+    def test_suppressed_within_cooldown(self, handler):
+        exc = _make_fatal_error()
+        with _patched_logger_and_time() as (mock_logger, _):
+            handler.log(exc, request_id="req-1", status_code=500)
+            handler.log(exc, request_id="req-2", status_code=500)
+            handler.log(exc, request_id="req-3", status_code=500)
+            handler.log(exc, request_id="req-4", status_code=500)
+
+        mock_logger.error.assert_called_once()
+
+    def test_summary_emitted_after_cooldown(self, handler):
+        exc = _make_fatal_error()
+        with _patched_logger_and_time() as (mock_logger, mock_time):
+            handler.log(exc, request_id="req-1", status_code=500)
+            for i in range(5):
+                handler.log(exc, request_id=f"req-s{i}", status_code=500)
+
+            mock_time.return_value = 100.0 + COOLDOWN + 1
+            handler.log(exc, request_id="req-trigger", status_code=500)
+
+        assert mock_logger.error.call_count == 2
+        summary_call = mock_logger.error.call_args_list[1]
+        assert "Suppressed" in summary_call[0][0]
+        assert summary_call[0][1] == 6
+
+    def test_reset_after_quiet_period_logs_full_traceback(self, handler):
+        """Tests that after 2x cooldown of silence, the next fatal error should
+        be treated as a new event and log a full traceback again."""
+        exc = _make_fatal_error()
+        with _patched_logger_and_time() as (mock_logger, mock_time):
+            # First crash gives full traceback
+            handler.log(exc, request_id="req-1", status_code=500)
+            assert mock_logger.error.call_count == 1
+            assert mock_logger.error.call_args.kwargs["exc_info"] is exc
+
+            # Suppress a few within cooldown
+            for i in range(3):
+                handler.log(exc, request_id=f"req-s{i}", status_code=500)
+            assert mock_logger.error.call_count == 1
+
+            # Move time past 2x cooldown
+            mock_time.return_value = 100.0 + 2 * COOLDOWN + 1
+
+            # Second crash should get a fresh full traceback
+            handler.log(exc, request_id="req-new-crash", status_code=500)
+
+        assert mock_logger.error.call_count == 2
+        new_crash_call = mock_logger.error.call_args_list[1]
+        assert new_crash_call.kwargs["exc_info"] is exc
+
+    def test_non_fatal_500_logs_every_call(self, handler):
+        with _patched_logger_and_time() as (mock_logger, _):
+            for i in range(5):
+                handler.log(ValueError("a"), request_id=f"r{i}", status_code=500)
+
+        assert mock_logger.error.call_count == 5
+
+    def test_non_fatal_4xx_logs_every_call(self, handler):
+        with _patched_logger_and_time() as (mock_logger, _):
+            for i in range(5):
+                handler.log(ValueError("a"), request_id=f"r{i}", status_code=400)
+
+        assert mock_logger.warning.call_count == 5
+
+
+class TestFatalEngineErrorRateLimiting:
+    """
+    Verify internal state transitions and rate limiting invariants.
+    """
+
+    def test_initial_state(self, handler):
+        assert not handler._first_logged
+        assert handler._suppressed_count == 0
+        assert handler._last_summary_time == 0.0
+
+    def test_first_fatal_sets_state(self, handler):
+        with _patched_logger_and_time():
+            handler.log(
+                _make_fatal_error(),
+                request_id="r1",
+                status_code=500,
+            )
+
+        assert handler._first_logged
+        assert handler._suppressed_count == 0
+        assert handler._last_summary_time == 100.0
+
+    def test_suppressed_count_increments(self, handler):
+        exc = _make_fatal_error()
+        handler.log(exc, request_id="r1", status_code=500)
+        handler.log(exc, request_id="r2", status_code=500)
+        handler.log(exc, request_id="r3", status_code=500)
+        handler.log(exc, request_id="r4", status_code=500)
+
+        assert handler._suppressed_count == 3
+
+    def test_cooldown_resets_count(self, handler):
+        exc = _make_fatal_error()
+        with _patched_logger_and_time() as (_, mock_time):
+            handler.log(exc, request_id="r1", status_code=500)
+            for i in range(5):
+                handler.log(exc, request_id=f"r{i+2}", status_code=500)
+            assert handler._suppressed_count == 5
+            assert handler._last_summary_time == 100
+
+            mock_time.return_value = 100.0 + COOLDOWN + 1
+            handler.log(exc, request_id="r-trigger", status_code=500)
+
+        assert handler._suppressed_count == 0
+        assert handler._last_summary_time == 100.0 + COOLDOWN + 1
+
+    def test_multiple_cooldown_windows(self, handler):
+        exc = _make_fatal_error()
+        with _patched_logger_and_time() as (mock_logger, mock_time):
+            handler.log(exc, request_id="r1", status_code=500)
+            handler.log(exc, request_id="r2", status_code=500)
+            handler.log(exc, request_id="r3", status_code=500)
+            assert handler._suppressed_count == 2
+
+            mock_time.return_value = 100.0 + COOLDOWN + 1
+            handler.log(exc, request_id="r4", status_code=500)
+            assert handler._suppressed_count == 0
+
+            handler.log(exc, request_id="r5", status_code=500)
+            assert handler._suppressed_count == 1
+
+            mock_time.return_value = 100.0 + (COOLDOWN + 1) * 2
+            handler.log(exc, request_id="r6", status_code=500)
+            assert handler._suppressed_count == 0
+
+        assert mock_logger.error.call_count == 3
+
+    def test_non_fatal_does_not_affect_fatal_state(self, handler):
+        handler.log(ValueError("a"), request_id="r1", status_code=500)
+        handler.log(ValueError("b"), request_id="r2", status_code=400)
+
+        assert not handler._first_logged
+        assert handler._suppressed_count == 0
+        assert handler._last_summary_time == 0
+
+
+class TestGetResponseForError:
+    """Checks that get_response_for_error routes through the fatal error log handler."""
+
+    def test_returns_error_response_and_invokes_fatal_handler(self):
+        exc = _make_fatal_error()
+        with patch.object(
+            server_utils._fatal_error_log_handler, "log", autospec=True
+        ) as mock_log:
+            resp = server_utils.get_response_for_error(exc, request_id="rid")
+
+        mock_log.assert_called_once_with(exc, "rid", 500)
+        assert resp.error is not None
+        assert "rid" in resp.error.message
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/utils/test_pg_utils.py

@@ -0,0 +1,158 @@
+"""Tests for get_bundle_indices_sorted_by_node."""
+
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import ray
+from ray._private.test_utils import placement_group_assert_no_leak
+from ray.llm._internal.serve.utils.pg_utils import get_bundle_indices_sorted_by_node
+from ray.serve._private.utils import get_head_node_id
+
+# Realistic 56-char hex node IDs
+NODE_A = "ab7c1d2e3f4a5b6c7d8e9f0a1b2c3d4e5f6a7b8c9d0e1f2a3b4c"  # driver node
+NODE_B = "1234567890abcdef1234567890abcdef1234567890abcdef12345678"
+NODE_C = "fedcba0987654321fedcba0987654321fedcba0987654321fedcba09"
+
+
+def _assert_grouped_by_node(result, bundles_to_node_id):
+    seen_nodes = set()
+    prev_node = None
+    for idx in result:
+        node = bundles_to_node_id[idx]
+        if node != prev_node:
+            assert (
+                node not in seen_nodes
+            ), f"Node {node} appeared non-contiguously in {result}"
+            seen_nodes.add(node)
+            prev_node = node
+
+
+@pytest.mark.parametrize(
+    "bundles_to_node_id,expected_sorted_bundle_indices",
+    [
+        pytest.param(
+            {0: NODE_C, 1: NODE_A, 2: NODE_B, 3: NODE_C, 4: NODE_A, 5: NODE_B},
+            [1, 4, 2, 5, 0, 3],
+        ),
+        pytest.param(
+            {0: NODE_B, 1: NODE_B, 2: NODE_A, 3: NODE_A},
+            [2, 3, 0, 1],
+        ),
+        pytest.param(
+            # Driver has no bundles
+            {0: NODE_C, 1: NODE_B, 2: NODE_C, 3: NODE_B},
+            [1, 3, 0, 2],
+        ),
+        pytest.param(
+            {0: NODE_A, 1: NODE_A, 2: NODE_A},
+            [0, 1, 2],
+        ),
+        pytest.param(
+            {0: NODE_A},
+            [0],
+        ),
+        pytest.param(
+            {},
+            [],
+        ),
+    ],
+)
+def test_sort_bundle_indices_by_node(
+    bundles_to_node_id, expected_sorted_bundle_indices
+):
+    mock_pg = MagicMock()
+
+    with (
+        patch(
+            "ray.llm._internal.serve.utils.pg_utils.placement_group_table",
+            return_value={"bundles_to_node_id": bundles_to_node_id},
+        ),
+        patch(
+            "ray.llm._internal.serve.utils.pg_utils.get_head_node_id",
+            return_value=NODE_A,
+        ),
+    ):
+        result = get_bundle_indices_sorted_by_node(mock_pg)
+
+    assert result == expected_sorted_bundle_indices
+
+
+@pytest.mark.parametrize(
+    "bundles_to_node_id,driver_node_id,expected_sorted_bundle_indices",
+    [
+        # Explicit driver node orders that node's bundles first, even though the
+        # head node (NODE_A) also owns bundles.
+        pytest.param(
+            {0: NODE_A, 1: NODE_B, 2: NODE_A, 3: NODE_B},
+            NODE_B,
+            [1, 3, 0, 2],
+        ),
+        # Head node (NODE_A) owns no bundles; explicit driver still wins over
+        # lexicographic fallback (which would otherwise pick NODE_B first).
+        pytest.param(
+            {0: NODE_B, 1: NODE_C, 2: NODE_B, 3: NODE_C},
+            NODE_C,
+            [1, 3, 0, 2],
+        ),
+    ],
+)
+def test_explicit_driver_node_ordered_first(
+    bundles_to_node_id, driver_node_id, expected_sorted_bundle_indices
+):
+    mock_pg = MagicMock()
+
+    with (
+        patch(
+            "ray.llm._internal.serve.utils.pg_utils.placement_group_table",
+            return_value={"bundles_to_node_id": bundles_to_node_id},
+        ),
+        patch(
+            "ray.llm._internal.serve.utils.pg_utils.get_head_node_id",
+            return_value=NODE_A,
+        ),
+    ):
+        result = get_bundle_indices_sorted_by_node(
+            mock_pg, driver_node_id=driver_node_id
+        )
+
+    assert result == expected_sorted_bundle_indices
+
+
+# Integration test: verifies the full path through a real cluster, though
+# bundle ordering across nodes is non-deterministic and may already be sorted.
+@pytest.mark.parametrize("strategy", ["SPREAD", "PACK"])
+def test_sort_bundles(ray_start_cluster, strategy):
+    cluster = ray_start_cluster
+    for _ in range(2):
+        cluster.add_node(num_cpus=2)
+    ray.init(address=cluster.address)
+    cluster.wait_for_nodes()
+
+    driver_node_id = get_head_node_id()
+
+    pg = ray.util.placement_group(
+        name=f"test_sort_bundles_{strategy}",
+        strategy=strategy,
+        bundles=[{"CPU": 1}] * 4,
+    )
+    ray.get(pg.ready())
+
+    result = get_bundle_indices_sorted_by_node(pg)
+
+    table = ray.util.placement_group_table(pg)
+    bundles_to_node_id = table["bundles_to_node_id"]
+
+    assert sorted(result) == list(range(4))
+    _assert_grouped_by_node(result, bundles_to_node_id)
+
+    driver_bundles = [i for i in result if bundles_to_node_id[i] == driver_node_id]
+    if driver_bundles:
+        assert result[: len(driver_bundles)] == driver_bundles
+
+    placement_group_assert_no_leak([pg])
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/serve/utils/testing_utils.py

@@ -0,0 +1,232 @@
+"""Shared testing utilities for Ray LLM serve tests.
+
+This is written with assumptions around how mocks for testing are expected to behave.
+"""
+
+import json
+import re
+from typing import List, Optional, Union
+
+from ray.llm._internal.serve.core.configs.openai_api_models import (
+    ChatCompletionResponse,
+    CompletionResponse,
+    DetokenizeResponse,
+    EmbeddingResponse,
+    ScoreResponse,
+    TokenizeResponse,
+    TranscriptionResponse,
+)
+
+
+class LLMResponseValidator:
+    """Reusable validation logic for LLM responses."""
+
+    @staticmethod
+    def get_expected_content(
+        api_type: str, max_tokens: int, lora_model_id: str = ""
+    ) -> str:
+        """Get expected content based on API type."""
+        expected_content = " ".join(f"test_{i}" for i in range(max_tokens))
+        if lora_model_id:
+            expected_content = f"[lora_model] {lora_model_id}: {expected_content}"
+        return expected_content
+
+    @staticmethod
+    def validate_non_streaming_response(
+        response: Union[ChatCompletionResponse, CompletionResponse],
+        api_type: str,
+        max_tokens: int,
+        lora_model_id: str = "",
+    ):
+        """Validate non-streaming responses."""
+        expected_content = LLMResponseValidator.get_expected_content(
+            api_type, max_tokens, lora_model_id
+        )
+
+        if api_type == "chat":
+            assert isinstance(response, ChatCompletionResponse)
+            assert response.choices[0].message.content == expected_content
+        elif api_type == "completion":
+            assert isinstance(response, CompletionResponse)
+            assert response.choices[0].text == expected_content
+
+    @staticmethod
+    def validate_streaming_chunks(
+        chunks: List[str], api_type: str, max_tokens: int, lora_model_id: str = ""
+    ):
+        """Validate streaming response chunks."""
+        # Should have max_tokens + 1 chunks (tokens + [DONE])
+        assert len(chunks) == max_tokens + 1
+
+        # Validate each chunk except the last [DONE] chunk
+        for chunk_iter, chunk in enumerate(chunks[:-1]):
+            pattern = r"data: (.*)\n\n"
+            match = re.match(pattern, chunk)
+            assert match is not None
+            chunk_data = json.loads(match.group(1))
+
+            expected_chunk = f"test_{chunk_iter}"
+            if lora_model_id and chunk_iter == 0:
+                expected_chunk = f"[lora_model] {lora_model_id}: {expected_chunk}"
+
+            if api_type == "chat":
+                delta = chunk_data["choices"][0]["delta"]
+                if chunk_iter == 0:
+                    assert delta["role"] == "assistant"
+                else:
+                    assert delta["role"] is None
+                assert delta["content"].strip() == expected_chunk
+            elif api_type == "completion":
+                text = chunk_data["choices"][0]["text"]
+                assert text.strip() == expected_chunk
+
+    @staticmethod
+    def validate_embedding_response(
+        response: EmbeddingResponse, expected_dimensions: Optional[int] = None
+    ):
+        """Validate embedding responses."""
+        assert isinstance(response, EmbeddingResponse)
+        assert response.object == "list"
+        assert len(response.data) == 1
+        assert response.data[0].object == "embedding"
+        assert isinstance(response.data[0].embedding, list)
+        assert (
+            len(response.data[0].embedding) > 0
+        )  # Should have some embedding dimensions
+        assert response.data[0].index == 0
+
+        # Check dimensions if specified
+        if expected_dimensions:
+            assert len(response.data[0].embedding) == expected_dimensions
+
+    @staticmethod
+    def validate_score_response(response: ScoreResponse):
+        """Validate score responses."""
+        assert isinstance(response, ScoreResponse)
+        assert response.object == "list"
+        assert len(response.data) >= 1
+
+        # Validate each score data element
+        for i, score_data in enumerate(response.data):
+            assert score_data.object == "score"
+            assert isinstance(score_data.score, float)
+            assert score_data.index == i  # Index should match position in list
+
+    @staticmethod
+    def validate_tokenize_response(
+        response: TokenizeResponse,
+        expected_prompt: str,
+        return_token_strs: bool = False,
+    ):
+        """Validate tokenize responses."""
+        assert isinstance(response, TokenizeResponse)
+        assert response.count == len(expected_prompt)
+        assert response.max_model_len > 0
+        assert isinstance(response.tokens, list)
+        assert len(response.tokens) == len(expected_prompt)
+
+        # Validate tokens are the character codes of the prompt
+        expected_tokens = [ord(c) for c in expected_prompt]
+        assert response.tokens == expected_tokens
+
+        # Validate token strings if requested
+        if return_token_strs:
+            assert response.token_strs is not None
+            assert len(response.token_strs) == len(expected_prompt)
+            assert response.token_strs == list(expected_prompt)
+        else:
+            assert response.token_strs is None
+
+    @staticmethod
+    def validate_detokenize_response(
+        response: DetokenizeResponse,
+        expected_text: str,
+    ):
+        """Validate detokenize responses."""
+        assert isinstance(response, DetokenizeResponse)
+        assert response.prompt == expected_text
+
+    @staticmethod
+    def validate_transcription_response(
+        response: Union[TranscriptionResponse, List[str]],
+        temperature: float,
+        language: Optional[str] = None,
+        lora_model_id: str = "",
+    ):
+        """Validate transcription responses for both streaming and non-streaming."""
+        if isinstance(response, list):
+            # Streaming response - validate chunks
+            LLMResponseValidator.validate_transcription_streaming_chunks(
+                response, temperature, language, lora_model_id
+            )
+        else:
+            # Non-streaming response
+            assert isinstance(response, TranscriptionResponse)
+            assert hasattr(response, "text")
+            assert isinstance(response.text, str)
+            assert len(response.text) > 0
+
+            # Check that the response contains expected language and temperature info
+            expected_text = f"Mock transcription in {language} language with temperature {temperature}"
+            if lora_model_id:
+                expected_text = f"[lora_model] {lora_model_id}: {expected_text}"
+            assert response.text == expected_text
+
+            # Validate usage information
+            if hasattr(response, "usage"):
+                assert hasattr(response.usage, "seconds")
+                assert hasattr(response.usage, "type")
+                assert response.usage.seconds > 0
+                assert response.usage.type == "duration"
+
+    @staticmethod
+    def validate_transcription_streaming_chunks(
+        chunks: List[str],
+        temperature: float,
+        language: Optional[str] = None,
+        lora_model_id: str = "",
+    ):
+        """Validate streaming transcription response chunks."""
+        # Should have at least one chunk (transcription text) + final chunk + [DONE]
+        assert len(chunks) >= 3
+
+        # Validate each chunk except the last [DONE] chunk
+        transcription_chunks = []
+        for chunk in chunks[:-1]:  # Exclude the final [DONE] chunk
+            pattern = r"data: (.*)\n\n"
+            match = re.match(pattern, chunk)
+            assert match is not None
+            chunk_data = json.loads(match.group(1))
+
+            # Validate chunk structure
+            assert "id" in chunk_data
+            assert "object" in chunk_data
+            assert chunk_data["object"] == "transcription.chunk"
+            assert "delta" in chunk_data
+            assert chunk_data["delta"] is None
+            assert "type" in chunk_data
+            assert chunk_data["type"] is None
+            assert "logprobs" in chunk_data
+            assert chunk_data["logprobs"] is None
+            assert "choices" in chunk_data
+            assert len(chunk_data["choices"]) == 1
+
+            choice = chunk_data["choices"][0]
+            assert "delta" in choice
+            assert "content" in choice["delta"]
+
+            # Collect text for final validation
+            if choice["delta"]["content"]:
+                transcription_chunks.append(choice["delta"]["content"])
+
+        # Validate final transcription text
+        full_transcription = "".join(transcription_chunks)
+        expected_text = (
+            f"Mock transcription in {language} language with temperature {temperature}"
+        )
+        if lora_model_id:
+            expected_text = f"[lora_model] {lora_model_id}: {expected_text}"
+        assert full_transcription.strip() == expected_text.strip()
+
+        # Validate final [DONE] chunk
+        assert chunks[-1] == "data: [DONE]\n\n"