From bd4f9acfd30456cb8fa98af53d32f7adc34e03a0 Mon Sep 17 00:00:00 2001 From: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com> Date: Sat, 22 Aug 2026 23:57:09 +0900 Subject: [PATCH] [Bugfix] Preserve pipeline sampling constraints (#6182) Signed-off-by: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- tests/entrypoints/test_omni_entrypoints.py | 32 +++++++++++++-- vllm_omni/config/stage_config.py | 3 ++ vllm_omni/entrypoints/omni_base.py | 45 +++++++++++++++++++++- 3 files changed, 76 insertions(+), 4 deletions(-) diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 161b3fe2a..f6649e328 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -12,12 +12,13 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError +from vllm_omni.config.stage_config import StageConfig from vllm_omni.engine.async_omni_engine import StageRuntimeInfo from vllm_omni.engine.messages import ErrorMessage, OutputMessage from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.omni import Omni -from vllm_omni.entrypoints.omni_base import OmniEngineDeadError +from vllm_omni.entrypoints.omni_base import OmniBase, OmniEngineDeadError from vllm_omni.errors import ( OmniClientError, client_error_from_metadata, @@ -188,14 +189,39 @@ def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine) def _make_base(): - from vllm_omni.entrypoints.omni_base import OmniBase - obj = object.__new__(OmniBase) obj.engine = MagicMock() obj.request_states = {} return obj +def test_resolve_sampling_params_list_preserves_stage_constraints(): + base = _make_base() + base.engine.num_stages = 1 + base.default_sampling_params_list = [SamplingParams(max_tokens=1000, detokenize=False, stop_token_ids=[42])] + base.engine.stage_configs = [ + StageConfig( + stage_id=0, + model_stage="dummy-model", + sampling_constraints={"detokenize": False, "stop_token_ids": [42]}, + ).to_omegaconf() + ] + base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs) + assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": [42]}] + caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[7]) + + resolved = base.resolve_sampling_params_list(caller_params) + + assert resolved[0] is not caller_params + assert resolved[0].seed == 1234 + assert resolved[0].max_tokens == 7 + assert resolved[0].detokenize is False + assert resolved[0].stop_token_ids == [42] + assert 42 in resolved[0]._all_stop_token_ids + assert caller_params.detokenize is True + assert caller_params.stop_token_ids == [7] + + def _stage_spec( stage_id: int, *, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 2b1a9f479..70ea7726a 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -959,6 +959,7 @@ def merge_pipeline_deploy( scheduler_cls=ps.scheduler_cls or _scheduler_path(sched_cls), hf_config_name=ps.hf_config_name, is_comprehension=ps.owns_tokenizer, + sampling_constraints=dict(ps.sampling_constraints), yaml_engine_args=engine_args, yaml_runtime=runtime, yaml_extras=extras, @@ -986,6 +987,7 @@ class StageConfig: scheduler_cls: str | None = None hf_config_name: str | None = None is_comprehension: bool = False + sampling_constraints: dict[str, Any] = field(default_factory=dict) yaml_engine_args: dict[str, Any] = field(default_factory=dict) yaml_runtime: dict[str, Any] = field(default_factory=dict) yaml_extras: dict[str, Any] = field(default_factory=dict) @@ -1046,6 +1048,7 @@ class StageConfig: "final_output": self.final_output, "final_output_type": self.final_output_type, "is_comprehension": self.is_comprehension, + "sampling_constraints": dict(self.sampling_constraints), } if self.custom_process_input_func: diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index ab2bee01b..ff21f55a2 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -3,11 +3,13 @@ from __future__ import annotations import os import time import weakref -from collections.abc import Sequence +from collections.abc import Mapping, Sequence +from dataclasses import fields, is_dataclass from typing import TYPE_CHECKING, Any, Literal import huggingface_hub import vllm.envs as envs +from omegaconf import OmegaConf from vllm.logger import init_logger from vllm.transformers_utils.repo_utils import file_or_path_exists from vllm.transformers_utils.runai_utils import is_runai_obj_uri @@ -223,6 +225,7 @@ class OmniBase(PDDisaggregationMixin): self.mod_metrics = OmniModalityMetrics(model_name=model, log_stats=log_stats) self.default_sampling_params_list = self.engine.default_sampling_params_list + self.sampling_constraints_list = self._get_sampling_constraints_list(self.engine.stage_configs) if not self.output_modalities: self.output_modalities = [ self.engine.get_stage_metadata(i).final_output_type for i in range(self.engine.num_stages) @@ -325,6 +328,7 @@ class OmniBase(PDDisaggregationMixin): sampling_params_list: Sequence[Any] | Any | None, allow_delta_coercion: bool = False, ) -> Sequence[Any]: + """Resolve request parameters; pipeline sampling constraints override caller values.""" if sampling_params_list is None: normalized = self.default_sampling_params_list # Set the output kind to delta since no params were specified @@ -339,8 +343,47 @@ class OmniBase(PDDisaggregationMixin): raise ValueError(f"Expected {self.num_stages} sampling params, got a single sampling params object") if len(normalized) != self.num_stages: raise ValueError(f"Expected {self.num_stages} sampling params, got {len(normalized)}") + + if sampling_params_list is not None: + normalized = [ + self._apply_sampling_constraints(params, constraints) + for params, constraints in zip(normalized, self.sampling_constraints_list, strict=True) + ] return normalized + @staticmethod + def _get_sampling_constraints_list(stage_configs: Sequence[Any]) -> list[dict[str, Any]]: + """Extract each stage's required sampling settings from runtime configs.""" + constraints_list = [] + for stage_config in stage_configs: + constraints = getattr(stage_config, "sampling_constraints", {}) + if not isinstance(constraints, Mapping): + constraints_list.append({}) + continue + if OmegaConf.is_config(constraints): + constraints = OmegaConf.to_container(constraints, resolve=True) + constraints_list.append(dict(constraints)) + return constraints_list + + @staticmethod + def _apply_sampling_constraints(params: Any, constraints: Mapping[str, Any]) -> Any: + """Rebuild params with pipeline-required settings without mutating caller input.""" + if not constraints: + return params + if isinstance(params, Mapping): + return {**params, **constraints} + if is_dataclass(params): + values = {field.name: getattr(params, field.name) for field in fields(params) if field.init} + elif struct_fields := getattr(params, "__struct_fields__", None): + values = { + name: getattr(params, name) + for name in struct_fields + if not name.startswith("_") and name != "output_text_buffer_length" + } + else: + raise TypeError(f"Expected a mapping, dataclass, or msgspec struct, got {type(params).__name__}") + return type(params)(**{**values, **constraints}) + def _fire_failure_counter_if_alive(self, request_id: str) -> None: """Fire the abort/exception bucket of requests_success_total.