[Feature] Support standard OpenAI API sampling parameters for thinker stage. (max_tokens...) (#450)

Signed-off-by: wonjae.lee0 <wonjae.lee0@navercorp.com>
Co-authored-by: wonjae.lee0 <wonjae.lee0@navercorp.com>

.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:

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

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()

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,