From 341457a1219d1f227549ba68d20b0b9a50ff86a3 Mon Sep 17 00:00:00 2001 From: Wonjae Lee Date: Wed, 24 Dec 2025 21:39:01 +0900 Subject: [PATCH] [Feature] Support standard OpenAI API sampling parameters for thinker stage. (max_tokens...) (#450) Signed-off-by: wonjae.lee0 Co-authored-by: wonjae.lee0 --- .buildkite/pipeline.yml | 2 +- docs/serving/image_generation_api.md | 4 +- tests/entrypoints/openai_api/__init__.py | 0 .../test_image_server.py | 0 .../test_serving_chat_sampling_params.py | 333 ++++++++++++++++++ vllm_omni/entrypoints/openai/serving_chat.py | 81 ++++- 6 files changed, 416 insertions(+), 4 deletions(-) create mode 100644 tests/entrypoints/openai_api/__init__.py rename tests/entrypoints/{openai => openai_api}/test_image_server.py (100%) create mode 100644 tests/entrypoints/openai_api/test_serving_chat_sampling_params.py diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 5d3219ff9..3ba45afc1 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -27,7 +27,7 @@ steps: - label: "Image Generation API Test" depends_on: image-build commands: - - pytest -s -v tests/entrypoints/openai/ + - pytest -s -v tests/entrypoints/openai_api/ agents: queue: "cpu_queue_premerge" plugins: diff --git a/docs/serving/image_generation_api.md b/docs/serving/image_generation_api.md index 113524efa..747cef995 100644 --- a/docs/serving/image_generation_api.md +++ b/docs/serving/image_generation_api.md @@ -233,10 +233,10 @@ Run the test suite to verify functionality: ```bash # All image generation tests -pytest tests/entrypoints/openai/test_image_server.py -v +pytest tests/entrypoints/openai_api/test_image_server.py -v # Specific test -pytest tests/entrypoints/openai/test_image_server.py::test_generate_single_image -v +pytest tests/entrypoints/openai_api/test_image_server.py::test_generate_single_image -v ``` ## Development diff --git a/tests/entrypoints/openai_api/__init__.py b/tests/entrypoints/openai_api/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/entrypoints/openai/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py similarity index 100% rename from tests/entrypoints/openai/test_image_server.py rename to tests/entrypoints/openai_api/test_image_server.py diff --git a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py new file mode 100644 index 000000000..240fd2051 --- /dev/null +++ b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py @@ -0,0 +1,333 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Unit tests for OmniOpenAIServingChat sampling params handling. + +Tests that standard OpenAI API parameters (max_tokens, temperature, etc.) +are correctly applied to the comprehension stage while preserving YAML defaults. +""" + +from unittest.mock import MagicMock + +import pytest +from vllm.sampling_params import SamplingParams + + +@pytest.fixture +def mock_comprehension_stage(): + """Create a mock comprehension stage with is_comprehension=True.""" + stage = MagicMock() + stage.is_comprehension = True + stage.model_stage = "comprehension" + return stage + + +@pytest.fixture +def mock_other_stage(): + """Create a mock non-comprehension stage.""" + stage = MagicMock() + stage.is_comprehension = False + stage.model_stage = "other" + return stage + + +@pytest.fixture +def default_comprehension_params(): + """Default sampling params for comprehension stage (from YAML).""" + return SamplingParams( + temperature=0.4, + top_p=0.9, + top_k=1, + max_tokens=2048, + seed=42, + repetition_penalty=1.05, + ) + + +@pytest.fixture +def default_other_params(): + """Default sampling params for non-comprehension stage (from YAML).""" + return SamplingParams( + temperature=0.9, + top_k=50, + max_tokens=4096, + seed=42, + ) + + +@pytest.fixture +def mock_engine_client(mock_comprehension_stage, mock_other_stage, default_comprehension_params, default_other_params): + """Create mock engine client with stage_list and default_sampling_params_list.""" + engine_client = MagicMock() + engine_client.stage_list = [mock_comprehension_stage, mock_other_stage] + engine_client.default_sampling_params_list = [ + default_comprehension_params, + default_other_params, + ] + return engine_client + + +@pytest.fixture +def serving_chat(mock_engine_client): + """Create OmniOpenAIServingChat instance with mocked dependencies.""" + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + # Create instance without calling __init__ + instance = object.__new__(OmniOpenAIServingChat) + instance.engine_client = mock_engine_client + return instance + + +@pytest.fixture +def mock_request(): + """Create a mock request with all OpenAI sampling params set to None.""" + request = MagicMock() + # OpenAI standard sampling fields + request.temperature = None + request.top_p = None + request.max_tokens = None + request.seed = None + request.stop = None + request.frequency_penalty = None + request.presence_penalty = None + return request + + +# ============================================================================= +# Tests for _OPENAI_SAMPLING_FIELDS constant +# ============================================================================= + + +def test_openai_sampling_fields_contains_expected_fields(): + """Test that _OPENAI_SAMPLING_FIELDS contains all expected OpenAI params.""" + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + expected_fields = { + "temperature", + "top_p", + "max_tokens", + "seed", + "stop", + "frequency_penalty", + "presence_penalty", + } + assert OmniOpenAIServingChat._OPENAI_SAMPLING_FIELDS == expected_fields + + +# ============================================================================= +# Tests for _build_sampling_params_list_from_request +# ============================================================================= + + +def test_preserves_yaml_defaults_when_no_request_params(serving_chat, mock_request): + """Test that YAML defaults are preserved when request has no params.""" + result = serving_chat._build_sampling_params_list_from_request(mock_request) + + assert len(result) == 2 + comprehension_params = result[0] + assert comprehension_params.temperature == 0.4 + assert comprehension_params.top_p == 0.9 + assert comprehension_params.top_k == 1 # YAML custom param preserved + assert comprehension_params.max_tokens == 2048 + assert comprehension_params.seed == 42 + assert comprehension_params.repetition_penalty == 1.05 # YAML custom param preserved + + +def test_request_temperature_overrides_yaml_default(serving_chat, mock_request): + """Test that request temperature overrides YAML default.""" + mock_request.temperature = 0.8 + + result = serving_chat._build_sampling_params_list_from_request(mock_request) + + comprehension_params = result[0] + assert comprehension_params.temperature == 0.8 # Overridden + assert comprehension_params.seed == 42 # Preserved from YAML + assert comprehension_params.top_k == 1 # YAML custom param preserved + + +def test_request_top_p_overrides_yaml_default(serving_chat, mock_request): + """Test that request top_p overrides YAML default.""" + mock_request.top_p = 0.95 + + result = serving_chat._build_sampling_params_list_from_request(mock_request) + + comprehension_params = result[0] + assert comprehension_params.top_p == 0.95 # Overridden + assert comprehension_params.temperature == 0.4 # Preserved from YAML + + +def test_request_max_tokens_overrides_yaml_default(serving_chat, mock_request): + """Test that request max_tokens overrides YAML default.""" + mock_request.max_tokens = 100 + + result = serving_chat._build_sampling_params_list_from_request(mock_request) + + assert result[0].max_tokens == 100 + + +def test_max_tokens_uses_yaml_default_when_not_specified(serving_chat, mock_request): + """Test that max_tokens falls back to YAML default when not in request.""" + result = serving_chat._build_sampling_params_list_from_request(mock_request) + + assert result[0].max_tokens == 2048 + + +def test_request_seed_overrides_yaml_default(serving_chat, mock_request): + """Test that request seed overrides YAML default.""" + mock_request.seed = 123 + + result = serving_chat._build_sampling_params_list_from_request(mock_request) + + comprehension_params = result[0] + assert comprehension_params.seed == 123 # Overridden + assert comprehension_params.temperature == 0.4 # Preserved from YAML + + +def test_request_frequency_penalty_overrides(serving_chat, mock_request): + """Test that request frequency_penalty is applied.""" + mock_request.frequency_penalty = 0.5 + + result = serving_chat._build_sampling_params_list_from_request(mock_request) + + assert result[0].frequency_penalty == 0.5 + + +def test_request_presence_penalty_overrides(serving_chat, mock_request): + """Test that request presence_penalty is applied.""" + mock_request.presence_penalty = 0.3 + + result = serving_chat._build_sampling_params_list_from_request(mock_request) + + assert result[0].presence_penalty == 0.3 + + +def test_non_comprehension_stages_use_cloned_defaults(serving_chat, mock_request): + """Test that non-comprehension stages always use cloned YAML defaults.""" + mock_request.max_tokens = 50 + mock_request.temperature = 0.1 + + result = serving_chat._build_sampling_params_list_from_request(mock_request) + + other_params = result[1] + assert other_params.temperature == 0.9 # YAML default (not affected by request) + assert other_params.max_tokens == 4096 # YAML default (not affected by request) + assert other_params.top_k == 50 # YAML default + assert other_params.seed == 42 # YAML default + + +def test_multiple_params_override_together(serving_chat, mock_request): + """Test that multiple request params can override together.""" + mock_request.max_tokens = 200 + mock_request.temperature = 0.7 + mock_request.top_p = 0.85 + mock_request.seed = 999 + + result = serving_chat._build_sampling_params_list_from_request(mock_request) + + comprehension_params = result[0] + # Overridden by request + assert comprehension_params.temperature == 0.7 + assert comprehension_params.top_p == 0.85 + assert comprehension_params.max_tokens == 200 + assert comprehension_params.seed == 999 + # Preserved from YAML (not in _OPENAI_SAMPLING_FIELDS) + assert comprehension_params.top_k == 1 + assert comprehension_params.repetition_penalty == 1.05 + + +def test_yaml_custom_params_not_overridden_by_request(serving_chat, mock_request): + """Test that YAML custom params (top_k, repetition_penalty) are not affected.""" + # Even if request has these attributes, they should not override YAML + # because they're not in _OPENAI_SAMPLING_FIELDS + mock_request.top_k = 100 # Not in allowlist + mock_request.repetition_penalty = 2.0 # Not in allowlist + + result = serving_chat._build_sampling_params_list_from_request(mock_request) + + comprehension_params = result[0] + assert comprehension_params.top_k == 1 # YAML default preserved + assert comprehension_params.repetition_penalty == 1.05 # YAML default preserved + + +# ============================================================================= +# Tests for _apply_request_overrides +# ============================================================================= + + +def test_apply_request_overrides_clones_params(serving_chat, mock_request, default_comprehension_params): + """Test that _apply_request_overrides returns a cloned object.""" + result = serving_chat._apply_request_overrides(default_comprehension_params, mock_request) + + assert result is not default_comprehension_params # Different object + + +def test_apply_request_overrides_preserves_defaults(serving_chat, mock_request, default_comprehension_params): + """Test that _apply_request_overrides preserves defaults when request has None.""" + result = serving_chat._apply_request_overrides(default_comprehension_params, mock_request) + + assert result.temperature == 0.4 + assert result.top_p == 0.9 + assert result.seed == 42 + assert result.top_k == 1 # YAML custom param + + +def test_apply_request_overrides_applies_values(serving_chat, mock_request, default_comprehension_params): + """Test that _apply_request_overrides applies non-None request values.""" + mock_request.temperature = 0.8 + mock_request.seed = 123 + + result = serving_chat._apply_request_overrides(default_comprehension_params, mock_request) + + assert result.temperature == 0.8 # Overridden + assert result.seed == 123 # Overridden + assert result.top_p == 0.9 # Preserved from default + assert result.top_k == 1 # YAML custom param preserved + + +# ============================================================================= +# Tests for _get_comprehension_stage_index +# ============================================================================= + + +def test_get_comprehension_stage_index_finds_first_stage(mock_engine_client): + """Test finding comprehension stage when it's at index 0.""" + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + instance = object.__new__(OmniOpenAIServingChat) + instance.engine_client = mock_engine_client + + assert instance._get_comprehension_stage_index() == 0 + + +def test_get_comprehension_stage_index_finds_second_stage(): + """Test finding comprehension stage when it's at index 1.""" + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + instance = object.__new__(OmniOpenAIServingChat) + + other = MagicMock() + other.is_comprehension = False + comprehension = MagicMock() + comprehension.is_comprehension = True + + instance.engine_client = MagicMock() + instance.engine_client.stage_list = [other, comprehension] + + assert instance._get_comprehension_stage_index() == 1 + + +def test_get_comprehension_stage_index_raises_when_not_found(): + """Test that ValueError is raised when no comprehension stage exists.""" + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + instance = object.__new__(OmniOpenAIServingChat) + + stage1 = MagicMock() + stage1.is_comprehension = False + stage2 = MagicMock() + stage2.is_comprehension = False + + instance.engine_client = MagicMock() + instance.engine_client.stage_list = [stage1, stage2] + + with pytest.raises(ValueError, match="No comprehension stage"): + instance._get_comprehension_stage_index() diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 9e8e8dcc7..422a38788 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -232,7 +232,8 @@ class OmniOpenAIServingChat(OpenAIServingChat): if hasattr(request, "sampling_params_list"): sampling_params_list = self._to_sampling_params_list(request.sampling_params_list) else: - sampling_params_list = None + # Use standard OpenAI API parameters for comprehension stage + sampling_params_list = self._build_sampling_params_list_from_request(request) self._log_inputs( request_id, @@ -407,6 +408,84 @@ class OmniOpenAIServingChat(OpenAIServingChat): raise ValueError(f"Invalid sampling params: {sampling_params}") return final_sampling_params_list + def _get_comprehension_stage_index(self) -> int: + for idx, stage in enumerate(self.engine_client.stage_list): + if stage.is_comprehension: + return idx + raise ValueError("No comprehension stage (is_comprehension=True) found in stage_list") + + # OpenAI API standard sampling parameters that can be safely overridden. + # These are the most commonly used parameters with compatible types + # between ChatCompletionRequest and SamplingParams. + # Users who need more control can use sampling_params_list in extra_body. + _OPENAI_SAMPLING_FIELDS: set[str] = { + "temperature", + "top_p", + "max_tokens", + "seed", + "stop", + "frequency_penalty", + "presence_penalty", + } + + def _apply_request_overrides( + self, + default_params: SamplingParams, + request: ChatCompletionRequest, + ) -> SamplingParams: + """Clone default params and override with user-provided request values. + + Starts with YAML defaults and only overrides fields that the user + explicitly provided (non-None values) in the request. + + Args: + default_params: Default SamplingParams from stage config YAML. + request: The chat completion request containing user-provided values. + + Returns: + New SamplingParams with YAML defaults overridden by request values. + """ + params = default_params.clone() + + for field_name in self._OPENAI_SAMPLING_FIELDS: + value = getattr(request, field_name, None) + if value is not None: + setattr(params, field_name, value) + + return params + + def _build_sampling_params_list_from_request( + self, + request: ChatCompletionRequest, + ) -> list[SamplingParams]: + """Build sampling_params_list using standard OpenAI API parameters. + + For the comprehension stage, starts with YAML defaults and overrides with + user-provided request values. For other stages, uses cloned YAML defaults. + + This approach ensures all YAML defaults (including seed, detokenize, etc.) + are preserved while allowing users to override specific parameters. + + Args: + request: The chat completion request containing OpenAI API parameters. + + Returns: + List of SamplingParams, one for each stage. + """ + default_params_list = self.engine_client.default_sampling_params_list + comprehension_idx = self._get_comprehension_stage_index() + + sampling_params_list = [] + for idx, default_params in enumerate(default_params_list): + if idx == comprehension_idx: + params = self._apply_request_overrides(default_params, request) + sampling_params_list.append(params) + else: + # For other stages, clone default params + sampling_params_list.append(default_params.clone()) + + return sampling_params_list + def _log_inputs( self, request_id: str,