@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 66c079ae83

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 09:45 UTC

[Frontend][4/n] Improve pooling entrypoints | pooling. (#39153)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>

wang.yuqi committed Apr 9, 2026 at 18:09 UTC 66c079ae83f5d5200306da28acb8a0949a6a9a45

43 files changed +549 -728

tests/entrypoints/openai/chat_completion/test_chat_error.py

-3

@@ -87,7 +87,6 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat:

87 serving_render = OpenAIServingRender(

88 model_config=engine.model_config,

89 renderer=engine.renderer,

90 - io_processor=engine.io_processor,

90 model_registry=models.registry,

91 request_logger=None,

92 chat_template=None,

@@ -123,7 +122,6 @@ async def test_chat_error_non_stream():

122 mock_engine.errored = False

123 mock_engine.model_config = MockModelConfig()

124 mock_engine.input_processor = MagicMock()

126 - mock_engine.io_processor = MagicMock()

125 mock_engine.renderer = _build_renderer(mock_engine.model_config)

126

127 serving_chat = _build_serving_chat(mock_engine)

@@ -173,7 +171,6 @@ async def test_chat_error_stream():

171 mock_engine.errored = False

172 mock_engine.model_config = MockModelConfig()

173 mock_engine.input_processor = MagicMock()

176 - mock_engine.io_processor = MagicMock()

174 mock_engine.renderer = _build_renderer(mock_engine.model_config)

175

176 serving_chat = _build_serving_chat(mock_engine)

tests/entrypoints/openai/chat_completion/test_serving_chat.py

-15

@@ -567,7 +567,6 @@ def _build_serving_render(

567 return OpenAIServingRender(

568 model_config=engine.model_config,

569 renderer=engine.renderer,

570 - io_processor=engine.io_processor,

570 model_registry=model_registry,

571 request_logger=None,

572 chat_template=CHAT_TEMPLATE,

@@ -599,7 +598,6 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat:

598 class MockEngine:

599 model_config: MockModelConfig = field(default_factory=MockModelConfig)

600 input_processor: MagicMock = field(default_factory=MagicMock)

602 - io_processor: MagicMock = field(default_factory=MagicMock)

601 renderer: MagicMock = field(default_factory=MagicMock)

602

603

@@ -632,7 +630,6 @@ async def test_serving_chat_returns_correct_model_name():

630 mock_engine.errored = False

631 mock_engine.model_config = MockModelConfig()

632 mock_engine.input_processor = MagicMock()

635 - mock_engine.io_processor = MagicMock()

633 mock_engine.renderer = _build_renderer(mock_engine.model_config)

634

635 serving_chat = _build_serving_chat(mock_engine)

@@ -662,7 +659,6 @@ async def test_serving_chat_should_set_correct_max_tokens():

659 mock_engine.errored = False

660 mock_engine.model_config = MockModelConfig()

661 mock_engine.input_processor = MagicMock()

665 - mock_engine.io_processor = MagicMock()

662 mock_engine.renderer = _build_renderer(mock_engine.model_config)

663

664 serving_chat = _build_serving_chat(mock_engine)

@@ -693,7 +689,6 @@ async def test_serving_chat_should_set_correct_max_tokens():

689 mock_engine.errored = False

690 mock_engine.model_config = mock_model_config

691 mock_engine.input_processor = MagicMock()

696 - mock_engine.io_processor = MagicMock()

692 mock_engine.renderer = _build_renderer(mock_engine.model_config)

693

694 # Initialize the serving chat

@@ -737,7 +732,6 @@ async def test_serving_chat_should_set_correct_max_tokens():

732 mock_engine.errored = False

733 mock_engine.model_config = mock_model_config

734 mock_engine.input_processor = MagicMock()

740 - mock_engine.io_processor = MagicMock()

735 mock_engine.renderer = _build_renderer(mock_engine.model_config)

736

737 serving_chat = _build_serving_chat(mock_engine)

@@ -779,7 +773,6 @@ async def test_serving_chat_should_set_correct_max_tokens():

773 mock_engine.errored = False

774 mock_engine.model_config = mock_model_config

775 mock_engine.input_processor = MagicMock()

782 - mock_engine.io_processor = MagicMock()

776 mock_engine.renderer = _build_renderer(mock_engine.model_config)

777

778 # Initialize the serving chat

@@ -823,7 +816,6 @@ async def test_serving_chat_mistral_token_ids_prompt_is_validated():

816 mock_engine.errored = False

817 mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True)

818 mock_engine.input_processor = MagicMock()

826 - mock_engine.io_processor = MagicMock()

819

820 mock_tokenizer = MagicMock(spec=MistralTokenizer)

821 mock_renderer = MistralRenderer(

@@ -863,7 +855,6 @@ async def test_serving_chat_mistral_token_ids_prompt_too_long_is_rejected():

855 mock_engine.errored = False

856 mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True)

857 mock_engine.input_processor = MagicMock()

866 - mock_engine.io_processor = MagicMock()

858

859 mock_tokenizer = MagicMock(spec=MistralTokenizer)

860 mock_renderer = MistralRenderer(

@@ -906,7 +897,6 @@ async def test_serving_chat_could_load_correct_generation_config():

897 mock_engine.errored = False

898 mock_engine.model_config = mock_model_config

899 mock_engine.input_processor = MagicMock()

909 - mock_engine.io_processor = MagicMock()

900 mock_engine.renderer = _build_renderer(mock_engine.model_config)

901

902 # Initialize the serving chat

@@ -952,7 +942,6 @@ async def test_serving_chat_did_set_correct_cache_salt(model_type):

942 mock_engine.errored = False

943 mock_engine.model_config = mock_model_config

944 mock_engine.input_processor = MagicMock()

955 - mock_engine.io_processor = MagicMock()

945 mock_engine.renderer = _build_renderer(mock_engine.model_config)

946

947 serving_chat = _build_serving_chat(mock_engine)

@@ -1003,7 +992,6 @@ async def test_serving_chat_data_parallel_rank_extraction():

992 mock_engine.errored = False

993 mock_engine.model_config = MockModelConfig()

994 mock_engine.input_processor = MagicMock()

1006 - mock_engine.io_processor = MagicMock()

995 mock_engine.renderer = _build_renderer(mock_engine.model_config)

996

997 # Mock the generate method to return an async generator

@@ -1095,7 +1083,6 @@ class TestServingChatWithHarmony:

1083 mock_engine.errored = False

1084 mock_engine.model_config = MockModelConfig()

1085 mock_engine.input_processor = MagicMock()

1098 - mock_engine.io_processor = MagicMock()

1086 mock_engine.renderer = _build_renderer(mock_engine.model_config)

1087 return mock_engine

1088

@@ -1732,7 +1719,6 @@ async def test_tool_choice_validation_without_parser():

1719 mock_engine.errored = False

1720 mock_engine.model_config = MockModelConfig()

1721 mock_engine.input_processor = MagicMock()

1735 - mock_engine.io_processor = MagicMock()

1722 mock_engine.renderer = _build_renderer(mock_engine.model_config)

1723

1724 models = OpenAIServingModels(

@@ -1802,7 +1788,6 @@ async def test_streaming_n_gt1_independent_tool_parsers():

1788 mock_engine.errored = False

1789 mock_engine.model_config = MockModelConfig()

1790 mock_engine.input_processor = MagicMock()

1805 - mock_engine.io_processor = MagicMock()

1791 mock_engine.renderer = _build_renderer(mock_engine.model_config)

1792

1793 models = OpenAIServingModels(

tests/entrypoints/openai/completion/test_completion_error.py

-3

@@ -79,7 +79,6 @@ def _build_serving_completion(engine: AsyncLLM) -> OpenAIServingCompletion:

79 serving_render = OpenAIServingRender(

80 model_config=engine.model_config,

81 renderer=engine.renderer,

82 - io_processor=engine.io_processor,

82 model_registry=models.registry,

83 request_logger=None,

84 chat_template=None,

@@ -107,7 +106,6 @@ async def test_completion_error_non_stream():

106 mock_engine.errored = False

107 mock_engine.model_config = MockModelConfig()

108 mock_engine.input_processor = MagicMock()

110 - mock_engine.io_processor = MagicMock()

109 mock_engine.renderer = _build_renderer(mock_engine.model_config)

110

111 serving_completion = _build_serving_completion(mock_engine)

@@ -157,7 +155,6 @@ async def test_completion_error_stream():

155 mock_engine.errored = False

156 mock_engine.model_config = MockModelConfig()

157 mock_engine.input_processor = MagicMock()

160 - mock_engine.io_processor = MagicMock()

158 mock_engine.renderer = _build_renderer(mock_engine.model_config)

159

160 serving_completion = _build_serving_completion(mock_engine)

tests/entrypoints/openai/completion/test_lora_resolvers.py

-2

@@ -137,7 +137,6 @@ def mock_serving_setup():

137

138 mock_engine.model_config = MockModelConfig()

139 mock_engine.input_processor = MagicMock()

140 - mock_engine.io_processor = MagicMock()

140 mock_engine.renderer = _build_renderer(mock_engine.model_config)

141

142 models = OpenAIServingModels(

@@ -148,7 +147,6 @@ def mock_serving_setup():

147 serving_render = OpenAIServingRender(

148 model_config=mock_engine.model_config,

149 renderer=mock_engine.renderer,

151 - io_processor=mock_engine.io_processor,

150 model_registry=models.registry,

151 request_logger=None,

152 chat_template=None,

tests/entrypoints/openai/generative_scoring/test_generative_scoring.py

-1

@@ -77,7 +77,6 @@ def _create_mock_engine():

77 mock_engine.errored = False

78 mock_engine.model_config = MockModelConfig()

79 mock_engine.input_processor = MagicMock()

80 - mock_engine.io_processor = MagicMock()

81 # renderer is accessed by OpenAIServing.__init__ and serving.py

82 mock_renderer = MagicMock()

tests/entrypoints/openai/responses/test_serving_responses.py

-4

@@ -218,7 +218,6 @@ class TestInitializeToolSessions:

218 engine_client.model_config = model_config

219

220 engine_client.input_processor = MagicMock()

221 - engine_client.io_processor = MagicMock()

221 engine_client.renderer = MagicMock()

222

223 models = MagicMock()

@@ -307,7 +306,6 @@ class TestValidateGeneratorInput:

306 engine_client.model_config = model_config

307

308 engine_client.input_processor = MagicMock()

310 - engine_client.io_processor = MagicMock()

309 engine_client.renderer = MagicMock()

310

311 models = MagicMock()

@@ -369,7 +367,6 @@ async def test_reasoning_tokens_counted_for_text_reasoning_model(monkeypatch):

367 model_config.get_diff_sampling_param.return_value = {}

368 engine_client.model_config = model_config

369 engine_client.input_processor = MagicMock()

372 - engine_client.io_processor = MagicMock()

370 engine_client.renderer = MagicMock()

371

372 tokenizer = FakeTokenizer()

@@ -672,7 +669,6 @@ def _make_serving_instance_with_reasoning():

669 model_config.get_diff_sampling_param.return_value = {}

670 engine_client.model_config = model_config

671 engine_client.input_processor = MagicMock()

675 - engine_client.io_processor = MagicMock()

672 engine_client.renderer = MagicMock()

673

674 models = MagicMock()

tests/entrypoints/pooling/classify/test_offline.py

+5 -2

@@ -110,8 +110,11 @@ def test_score_api(llm: LLM):

110 llm.score("ping", "pong", use_tqdm=False)

111

112

113 -@pytest.mark.parametrize("task", ["embed", "token_embed"])

113 +@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])

114 def test_unsupported_tasks(llm: LLM, task: PoolingTask):

115 - err_msg = "Embedding API is not supported by this model.+"

115 + if task == "plugin":

116 + err_msg = "No IOProcessor plugin installed."

117 + else:

118 + err_msg = "Embedding API is not supported by this model.+"

119 with pytest.raises(ValueError, match=err_msg):

120 llm.encode(prompt, pooling_task=task, use_tqdm=False)

tests/entrypoints/pooling/classify/test_online.py

+5 -1

@@ -469,4 +469,8 @@ async def test_pooling_not_supported(

469 },

470 )

471 assert response.json()["error"]["type"] == "BadRequestError"

472 - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")

472 + if task == "plugin":

473 + err_msg = "No IOProcessor plugin installed."

474 + else:

475 + err_msg = f"Unsupported task: {task!r}"

476 + assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/embed/test_offline.py

+5 -2

@@ -107,8 +107,11 @@ def test_pooling_params(llm: LLM):

107 )

108

109

110 -@pytest.mark.parametrize("task", ["token_classify", "classify"])

110 +@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"])

111 def test_unsupported_tasks(llm: LLM, task: PoolingTask):

112 - err_msg = "Classification API is not supported by this model.+"

112 + if task == "plugin":

113 + err_msg = "No IOProcessor plugin installed."

114 + else:

115 + err_msg = "Classification API is not supported by this model.+"

116 with pytest.raises(ValueError, match=err_msg):

117 llm.encode(prompt, pooling_task=task, use_tqdm=False)

tests/entrypoints/pooling/embed/test_online.py

+5 -1

@@ -767,4 +767,8 @@ async def test_pooling_not_supported(

767 },

768 )

769 assert response.json()["error"]["type"] == "BadRequestError"

770 - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")

770 + if task == "plugin":

771 + err_msg = "No IOProcessor plugin installed."

772 + else:

773 + err_msg = f"Unsupported task: {task!r}"

774 + assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/scoring/test_bi_encoder_online.py

+5 -1

@@ -411,4 +411,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):

411 },

412 )

413 assert response.json()["error"]["type"] == "BadRequestError"

414 - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")

414 + if task == "plugin":

415 + err_msg = "No IOProcessor plugin installed."

416 + else:

417 + err_msg = f"Unsupported task: {task!r}"

418 + assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/scoring/test_cross_encoder_online.py

+5 -1

@@ -484,4 +484,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):

484 },

485 )

486 assert response.json()["error"]["type"] == "BadRequestError"

487 - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")

487 + if task == "plugin":

488 + err_msg = "No IOProcessor plugin installed."

489 + else:

490 + err_msg = f"Unsupported task: {task!r}"

491 + assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/token_classify/test_offline.py

+5 -2

@@ -65,14 +65,17 @@ def test_score_api(llm: LLM):

65 llm.score("ping", "pong", use_tqdm=False)

68 -@pytest.mark.parametrize("task", ["classify", "embed", "token_embed"])

68 +@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])

69 def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):

70 if task == "classify":

71 with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):

72 llm.encode(prompt, pooling_task=task, use_tqdm=False)

73 assert "deprecated" in caplog_vllm.text

74 else:

75 - err_msg = "Embedding API is not supported by this model.+"

75 + if task == "plugin":

76 + err_msg = "No IOProcessor plugin installed."

77 + else:

78 + err_msg = "Embedding API is not supported by this model.+"

80 with pytest.raises(ValueError, match=err_msg):

81 llm.encode(prompt, pooling_task=task, use_tqdm=False)

tests/entrypoints/pooling/token_classify/test_online.py

+5 -4

@@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st

51 @pytest.mark.asyncio

52 @pytest.mark.parametrize("model_name", [MODEL_NAME])

53 -@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])

53 +@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])

54 async def test_pooling_not_supported(

55 server: RemoteOpenAIServer, model_name: str, task: str

56 ):

@@ -64,7 +64,8 @@ async def test_pooling_not_supported(

64 },

65 )

67 - if task != "classify":

68 - assert response.json()["error"]["type"] == "BadRequestError"

67 + if task == "plugin":

68 + err_msg = "No IOProcessor plugin installed."

69 + else:

70 err_msg = f"Unsupported task: {task!r}"

70 - assert response.json()["error"]["message"].startswith(err_msg)

71 + assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/pooling/token_embed/test_offline.py

+5 -2

@@ -62,14 +62,17 @@ def test_token_ids_prompts(llm: LLM):

62 assert outputs[0].outputs.data.shape == (11, 384)

65 -@pytest.mark.parametrize("task", ["embed", "classify", "token_classify"])

65 +@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])

66 def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):

67 if task == "embed":

68 with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):

69 llm.encode(prompt, pooling_task=task, use_tqdm=False)

70 assert "deprecated" in caplog_vllm.text

71 else:

72 - err_msg = "Classification API is not supported by this model.+"

72 + if task == "plugin":

73 + err_msg = "No IOProcessor plugin installed."

74 + else:

75 + err_msg = "Classification API is not supported by this model.+"

77 with pytest.raises(ValueError, match=err_msg):

78 llm.encode(prompt, pooling_task=task, use_tqdm=False)

tests/entrypoints/pooling/token_embed/test_online.py

+5 -4

@@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):

74 @pytest.mark.asyncio

75 @pytest.mark.parametrize("model_name", [MODEL_NAME])

76 -@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])

76 +@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])

77 async def test_pooling_not_supported(

78 server: RemoteOpenAIServer, model_name: str, task: str

79 ):

@@ -87,7 +87,8 @@ async def test_pooling_not_supported(

87 },

88 )

90 - if task != "embed":

91 - assert response.json()["error"]["type"] == "BadRequestError"

90 + if task == "plugin":

91 + err_msg = "No IOProcessor plugin installed."

92 + else:

93 err_msg = f"Unsupported task: {task!r}"

93 - assert response.json()["error"]["message"].startswith(err_msg)

94 + assert response.json()["error"]["message"].startswith(err_msg)

tests/entrypoints/serve/disagg/test_generate_stream.py

-2

@@ -86,7 +86,6 @@ def _build_serving_tokens(engine: AsyncLLM, **kwargs) -> ServingTokens:

86 serving_render = OpenAIServingRender(

87 model_config=engine.model_config,

88 renderer=engine.renderer,

89 - io_processor=engine.io_processor,

89 model_registry=models.registry,

90 request_logger=None,

91 chat_template=None,

@@ -148,7 +147,6 @@ def _mock_engine() -> MagicMock:

147 engine.errored = False

148 engine.model_config = MockModelConfig()

149 engine.input_processor = MagicMock()

151 - engine.io_processor = MagicMock()

150 engine.renderer = _build_renderer(engine.model_config)

151 return engine

152

tests/entrypoints/serve/lora/test_serving_models.py

-1

@@ -34,7 +34,6 @@ async def _async_serving_models_init() -> OpenAIServingModels:

34 mock_model_config.max_model_len = 2048

35 mock_engine_client.model_config = mock_model_config

36 mock_engine_client.input_processor = MagicMock()

37 - mock_engine_client.io_processor = MagicMock()

37 mock_engine_client.renderer = MagicMock()

39 serving_models = OpenAIServingModels(

tests/v1/engine/test_async_llm.py

-1

@@ -514,7 +514,6 @@ async def test_header_dp_rank_argument():

514 serving_render = OpenAIServingRender(

515 model_config=engine.model_config,

516 renderer=engine.renderer,

517 - io_processor=engine.io_processor,

517 model_registry=models.registry,

518 request_logger=None,

519 chat_template=None,

vllm/engine/protocol.py

-2

@@ -14,7 +14,6 @@ from vllm.distributed.weight_transfer.base import (

14 from vllm.inputs import EngineInput, PromptType

15 from vllm.lora.request import LoRARequest

16 from vllm.outputs import PoolingRequestOutput, RequestOutput

17 -from vllm.plugins.io_processors import IOProcessor

17 from vllm.pooling_params import PoolingParams

18 from vllm.renderers import BaseRenderer

19 from vllm.sampling_params import SamplingParams

@@ -44,7 +43,6 @@ class EngineClient(ABC):

43 vllm_config: VllmConfig

44 model_config: ModelConfig

45 renderer: BaseRenderer

47 - io_processor: IOProcessor | None

46 input_processor: InputProcessor

48 @property

vllm/entrypoints/llm.py

+59 -118

@@ -49,9 +49,7 @@ from vllm.entrypoints.chat_utils import (

49 load_chat_template,

50 )

51 from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors

52 -from vllm.entrypoints.pooling.scoring.io_processor import (

53 - ScoringIOProcessor,

54 -)

52 +from vllm.entrypoints.pooling.scoring.io_processor import ScoringIOProcessor

53 from vllm.entrypoints.pooling.scoring.typing import ScoreInput

54 from vllm.entrypoints.pooling.typing import OfflineInputsContext, OfflineOutputsContext

55 from vllm.entrypoints.utils import log_non_default_args

@@ -398,12 +396,11 @@ class LLM:

396 self.runner_type = self.model_config.runner_type

397 self.renderer = self.llm_engine.renderer

398 self.chat_template = load_chat_template(chat_template)

401 - self.io_processor = self.llm_engine.io_processor

399 self.input_processor = self.llm_engine.input_processor

400 self.chat_template_config = ChatTemplateConfig(chat_template=self.chat_template)

401 self.pooling_io_processors = init_pooling_io_processors(

402 supported_tasks=supported_tasks,

406 - model_config=self.model_config,

403 + vllm_config=self.llm_engine.vllm_config,

404 renderer=self.renderer,

405 chat_template_config=self.chat_template_config,

406 )

@@ -1081,118 +1078,55 @@ class LLM:

1078 pooled hidden states in the same order as the input prompts.

1079 """

1080

1081 + if isinstance(prompts, dict) and "data" in prompts and pooling_task != "plugin":

1082 + raise ValueError(

1083 + "The 'data' field is only supported for the 'plugin' pooling task."

1084 + )

1085 self._verify_pooling_task(pooling_task)

1086 + assert pooling_task is not None and pooling_task in self.pooling_io_processors

1087

1086 - if isinstance(prompts, dict) and "data" in prompts:

1087 - if self.io_processor is None:

1088 - raise ValueError(

1089 - "No IOProcessor plugin installed. Please refer "

1090 - "to the documentation and to the "

1091 - "'prithvi_geospatial_mae_io_processor' "

1092 - "offline inference example for more details."

1093 - )

1088 + io_processor = self.pooling_io_processors[pooling_task]

1089

1095 - # Validate the request data is valid for the loaded plugin

1096 - prompt_data = prompts.get("data")

1097 - if prompt_data is None:

1098 - raise ValueError(

1099 - "The 'data' field of the prompt is expected to contain "

1100 - "the prompt data and it cannot be None. "

1101 - "Refer to the documentation of the IOProcessor "

1102 - "in use for more details."

1103 - )

1104 - validated_prompt = self.io_processor.parse_data(prompt_data)

1090 + if pooling_params is None:

1091 + pooling_params = PoolingParams()

1092

1106 - # obtain the actual model prompts from the pre-processor

1107 - prompts = self.io_processor.pre_process(prompt=validated_prompt)

1108 - prompts_seq = prompt_to_seq(prompts)

1093 + ctx = OfflineInputsContext(

1094 + prompts=prompts,

1095 + pooling_params=pooling_params,

1096 + tokenization_kwargs=tokenization_kwargs,

1097 + )

1098

1110 - params_seq: Sequence[PoolingParams] = [

1111 - self.io_processor.merge_pooling_params(param)

1112 - for param in self._params_to_seq(

1113 - pooling_params,

1114 - len(prompts_seq),

1115 - )

1116 - ]

1117 - for p in params_seq:

1118 - if p.task is None:

1119 - p.task = "plugin"

1120 -

1121 - outputs = self._run_completion(

1122 - prompts=prompts_seq,

1123 - params=params_seq,

1124 - output_type=PoolingRequestOutput,

1125 - use_tqdm=use_tqdm,

1126 - lora_request=lora_request,

1127 - tokenization_kwargs=tokenization_kwargs,

1128 - )

1099 + engine_inputs = io_processor.pre_process_offline(ctx)

1100 + n_inputs = len(engine_inputs)

1101 + assert ctx.pooling_params is not None

1102

1130 - # get the post-processed model outputs

1131 - assert self.io_processor is not None

1132 - processed_outputs = self.io_processor.post_process(outputs)

1103 + params_seq = self._params_to_seq(ctx.pooling_params, n_inputs)

1104

1134 - return [

1135 - PoolingRequestOutput[Any](

1136 - request_id="",

1137 - outputs=processed_outputs,

1138 - num_cached_tokens=getattr(

1139 - processed_outputs, "num_cached_tokens", 0

1140 - ),

1141 - prompt_token_ids=[],

1142 - finished=True,

1143 - )

1144 - ]

1145 - else:

1146 - if pooling_params is None:

1147 - # Use default pooling params.

1148 - pooling_params = PoolingParams()

1149 -

1150 - prompts_seq = prompt_to_seq(prompts)

1151 - params_seq = self._params_to_seq(pooling_params, len(prompts_seq))

1152 -

1153 - for param in params_seq:

1154 - if param.task is None:

1155 - param.task = pooling_task

1156 - elif param.task != pooling_task:

1157 - msg = (

1158 - f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"

1159 - )

1160 - raise ValueError(msg)

1105 + for param in params_seq:

1106 + if param.task is None:

1107 + param.task = pooling_task

1108 + elif pooling_task == "plugin":

1109 + # `plugin` task uses io_processor.parse_request to verify inputs.

1110 + # We actually allow plugin to overwrite pooling_task.

1111 + pass

1112 + elif param.task != pooling_task:

1113 + msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"

1114 + raise ValueError(msg)

1115

1162 - if pooling_task in self.pooling_io_processors:

1163 - io_processor = self.pooling_io_processors[pooling_task]

1164 - processor_inputs = io_processor.pre_process_offline(

1165 - ctx=OfflineInputsContext(

1166 - prompts=prompts_seq, tokenization_kwargs=tokenization_kwargs

1167 - )

1168 - )

1169 - seq_lora_requests = self._lora_request_to_seq(

1170 - lora_request, len(prompts_seq)

1171 - )

1172 - seq_priority = self._priority_to_seq(None, len(prompts))

1116 + seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs)

1117 + seq_priority = self._priority_to_seq(None, n_inputs)

1118

1174 - self._render_and_add_requests(

1175 - prompts=processor_inputs,

1176 - params=params_seq,

1177 - lora_requests=seq_lora_requests,

1178 - priorities=seq_priority,

1179 - )

1119 + self._render_and_add_requests(

1120 + prompts=engine_inputs,

1121 + params=params_seq,

1122 + lora_requests=seq_lora_requests,

1123 + priorities=seq_priority,

1124 + )

1125

1181 - outputs = self._run_engine(

1182 - use_tqdm=use_tqdm, output_type=PoolingRequestOutput

1183 - )

1184 - outputs = io_processor.post_process_offline(

1185 - ctx=OfflineOutputsContext(outputs=outputs)

1186 - )

1187 - else:

1188 - outputs = self._run_completion(

1189 - prompts=prompts_seq,

1190 - params=params_seq,

1191 - output_type=PoolingRequestOutput,

1192 - use_tqdm=use_tqdm,

1193 - lora_request=lora_request,

1194 - tokenization_kwargs=tokenization_kwargs,

1195 - )

1126 + outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput)

1127 + outputs = io_processor.post_process_offline(

1128 + ctx=OfflineOutputsContext(outputs=outputs)

1129 + )

1130 return outputs

1131

1132 def _verify_pooling_task(self, pooling_task: PoolingTask | None):

@@ -1254,6 +1188,14 @@ class LLM:

1188 pooling_task,

1189 )

1190

1191 + if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors:

1192 + raise ValueError(

1193 + "No IOProcessor plugin installed. Please refer "

1194 + "to the documentation and to the "

1195 + "'prithvi_geospatial_mae_io_processor' "

1196 + "offline inference example for more details."

1197 + )

1198 +

1199 def embed(

1200 self,

1201 prompts: PromptType | Sequence[PromptType],

@@ -1458,6 +1400,9 @@ class LLM:

1400 scoring_data = io_processor.valid_inputs(data_1, data_2)

1401 n_queries = len(scoring_data.data_1)

1402

1403 + if pooling_params is None:

1404 + pooling_params = PoolingParams()

1405 +

1406 ctx = OfflineInputsContext(

1407 prompts=scoring_data,

1408 pooling_params=pooling_params,

@@ -1466,15 +1411,11 @@ class LLM:

1411 n_queries=n_queries,

1412 )

1413

1469 - processor_inputs = io_processor.pre_process_offline(ctx)

1414 + engine_inputs = io_processor.pre_process_offline(ctx)

1415 + n_inputs = len(engine_inputs)

1416

1471 - seq_lora_requests = self._lora_request_to_seq(

1472 - lora_request, len(processor_inputs)

1473 - )

1474 -

1475 - if ctx.pooling_params is None:

1476 - ctx.pooling_params = PoolingParams()

1477 - params_seq = self._params_to_seq(ctx.pooling_params, len(processor_inputs))

1417 + seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs)

1418 + params_seq = self._params_to_seq(ctx.pooling_params, n_inputs)

1419

1420 for param in params_seq:

1421 if param.task is None:

@@ -1483,10 +1424,10 @@ class LLM:

1424 msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"

1425 raise ValueError(msg)

1426

1486 - seq_priority = self._priority_to_seq(None, len(processor_inputs))

1427 + seq_priority = self._priority_to_seq(None, n_inputs)

1428

1429 self._render_and_add_requests(

1489 - prompts=processor_inputs,

1430 + prompts=engine_inputs,

1431 params=params_seq,

1432 lora_requests=seq_lora_requests,

1433 priorities=seq_priority,

@@ -1579,7 +1520,7 @@ class LLM:

1520 if isinstance(params, Sequence):

1521 if len(params) != num_requests:

1522 raise ValueError(

1582 - f"The lengths of prompts ({params}) "

1523 + f"The lengths of prompts ({num_requests}) "

1524 f"and params ({len(params)}) must be the same."

1525 )

1526

vllm/entrypoints/openai/api_server.py

+1 -7

@@ -370,7 +370,6 @@ async def init_app_state(

370 state.openai_serving_render = OpenAIServingRender(

371 model_config=engine_client.model_config,

372 renderer=engine_client.renderer,

373 - io_processor=engine_client.io_processor,

373 model_registry=state.openai_serving_models.registry,

374 request_logger=request_logger,

375 chat_template=resolved_chat_template,

@@ -441,13 +440,12 @@ async def init_render_app_state(

440

441 Unlike :func:`init_app_state` this function does not require an

442 :class:`~vllm.engine.protocol.EngineClient`; it bootstraps the

444 - preprocessing pipeline (renderer, io_processor, input_processor)

443 + preprocessing pipeline (renderer, input_processor)

444 directly from the :class:`~vllm.config.VllmConfig`.

445 """

446 from vllm.entrypoints.chat_utils import load_chat_template

447 from vllm.entrypoints.openai.models.serving import OpenAIModelRegistry

448 from vllm.entrypoints.serve.render.serving import OpenAIServingRender

450 - from vllm.plugins.io_processors import get_io_processor

449 from vllm.renderers import renderer_from_config

450

451 served_model_names = args.served_model_name or [args.model]

@@ -465,15 +463,11 @@ async def init_render_app_state(

463 request_logger = None

464

465 renderer = renderer_from_config(vllm_config)

468 - io_processor = get_io_processor(

469 - vllm_config, renderer, vllm_config.model_config.io_processor_plugin

470 - )

466 resolved_chat_template = load_chat_template(args.chat_template)

467

468 state.openai_serving_render = OpenAIServingRender(

469 model_config=vllm_config.model_config,

470 renderer=renderer,

476 - io_processor=io_processor,

471 model_registry=model_registry,

472 request_logger=request_logger,

473 chat_template=resolved_chat_template,

vllm/entrypoints/openai/engine/serving.py

+5 -180

@@ -44,12 +44,6 @@ from vllm.entrypoints.openai.speech_to_text.protocol import (

44 TranscriptionResponse,

45 TranslationRequest,

46 )

47 -from vllm.entrypoints.pooling.pooling.protocol import (

48 - IOProcessorRequest,

49 - PoolingChatRequest,

50 - PoolingCompletionRequest,

51 - PoolingResponse,

52 -)

47 from vllm.entrypoints.serve.disagg.protocol import GenerateRequest, GenerateResponse

48 from vllm.entrypoints.serve.tokenize.protocol import (

49 DetokenizeRequest,

@@ -62,8 +56,7 @@ from vllm.inputs import EngineInput, PromptType

56 from vllm.logger import init_logger

57 from vllm.logprobs import Logprob, PromptLogprobs

58 from vllm.lora.request import LoRARequest

65 -from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput

66 -from vllm.pooling_params import PoolingParams

59 +from vllm.outputs import CompletionOutput, RequestOutput

60 from vllm.renderers import ChatParams, TokenizeParams

61 from vllm.renderers.inputs.preprocess import (

62 extract_prompt_components,

@@ -78,10 +71,7 @@ from vllm.tracing import (

71 log_tracing_disabled_warning,

72 )

73 from vllm.utils import random_uuid

81 -from vllm.utils.async_utils import (

82 - collect_from_async_generator,

83 - merge_async_iterators,

84 -)

74 +from vllm.utils.async_utils import collect_from_async_generator

76 logger = init_logger(__name__)

@@ -101,17 +91,11 @@ class RendererChatRequest(RendererRequest, Protocol):

93 CompletionLikeRequest: TypeAlias = (

104 - CompletionRequest

105 - | TokenizeCompletionRequest

106 - | DetokenizeRequest

107 - | PoolingCompletionRequest

94 + CompletionRequest | TokenizeCompletionRequest | DetokenizeRequest

95 )

97 ChatLikeRequest: TypeAlias = (

111 - ChatCompletionRequest

112 - | BatchChatCompletionRequest

113 - | TokenizeChatRequest

114 - | PoolingChatRequest

98 + ChatCompletionRequest | BatchChatCompletionRequest | TokenizeChatRequest

99 )

100

101 SpeechToTextRequest: TypeAlias = TranscriptionRequest | TranslationRequest

@@ -121,7 +105,6 @@ AnyRequest: TypeAlias = (

105 | ChatLikeRequest

106 | SpeechToTextRequest

107 | ResponsesRequest

124 - | IOProcessorRequest

108 | GenerateRequest

109 )

110

@@ -130,7 +113,6 @@ AnyResponse: TypeAlias = (

113 | ChatCompletionResponse

114 | TranscriptionResponse

115 | TokenizeResponse

133 - | PoolingResponse

116 | GenerateResponse

117 )

118

@@ -146,12 +128,6 @@ class ServeContext(Generic[RequestT]):

128 created_time: int = field(default_factory=lambda: int(time.time()))

129 lora_request: LoRARequest | None = None

130 engine_inputs: list[EngineInput] | None = None

149 -

150 - result_generator: AsyncGenerator[tuple[int, PoolingRequestOutput], None] | None = (

151 - None

152 - )

153 - final_res_batch: list[PoolingRequestOutput] = field(default_factory=list)

154 -

131 model_config = ConfigDict(arbitrary_types_allowed=True)

132

133

@@ -171,7 +147,6 @@ class OpenAIServing:

147 super().__init__()

148

149 self.engine_client = engine_client

174 -

150 self.models = models

151

152 self.request_logger = request_logger

@@ -179,7 +154,6 @@ class OpenAIServing:

154

155 self.model_config = engine_client.model_config

156 self.renderer = engine_client.renderer

182 - self.io_processor = engine_client.io_processor

157 self.input_processor = engine_client.input_processor

158

159 async def beam_search(

@@ -381,155 +355,6 @@ class OpenAIServing:

355 prompt_logprobs=None,

356 )

357

384 - async def _preprocess(

385 - self,

386 - ctx: ServeContext,

387 - ) -> ErrorResponse | None:

388 - """

389 - Default preprocessing hook. Subclasses may override to prepare `ctx`.

390 - """

391 - return None

392 -

393 - def _build_response(

394 - self,

395 - ctx: ServeContext,

396 - ) -> AnyResponse | ErrorResponse:

397 - """

398 - Default response builder. Subclass may override this method

399 - to return the appropriate response object.

400 - """

401 - return self.create_error_response("unimplemented endpoint")

402 -

403 - async def handle(

404 - self,

405 - ctx: ServeContext,

406 - ) -> AnyResponse | ErrorResponse:

407 - async for response in self._pipeline(ctx):

408 - return response

409 -

410 - return self.create_error_response("No response yielded from pipeline")

411 -

412 - async def _pipeline(

413 - self,

414 - ctx: ServeContext,

415 - ) -> AsyncGenerator[AnyResponse | ErrorResponse, None]:

416 - """Execute the request processing pipeline yielding responses."""

417 - if error := await self._check_model(ctx.request):

418 - yield error

419 - if error := self._validate_request(ctx):

420 - yield error

421 -

422 - preprocess_ret = await self._preprocess(ctx)

423 - if isinstance(preprocess_ret, ErrorResponse):

424 - yield preprocess_ret

425 -

426 - generators_ret = await self._prepare_generators(ctx)

427 - if isinstance(generators_ret, ErrorResponse):

428 - yield generators_ret

429 -

430 - collect_ret = await self._collect_batch(ctx)

431 - if isinstance(collect_ret, ErrorResponse):

432 - yield collect_ret

433 -

434 - yield self._build_response(ctx)

435 -

436 - def _validate_request(self, ctx: ServeContext) -> ErrorResponse | None:

437 - truncate_prompt_tokens = getattr(ctx.request, "truncate_prompt_tokens", None)

438 -

439 - if (

440 - truncate_prompt_tokens is not None

441 - and truncate_prompt_tokens > self.model_config.max_model_len

442 - ):

443 - return self.create_error_response(

444 - "truncate_prompt_tokens value is "

445 - "greater than max_model_len."

446 - " Please request a smaller truncation size."

447 - )

448 - return None

449 -

450 - def _create_pooling_params(

451 - self,

452 - ctx: ServeContext,

453 - ) -> PoolingParams | ErrorResponse:

454 - if not hasattr(ctx.request, "to_pooling_params"):

455 - return self.create_error_response(

456 - "Request type does not support pooling parameters"

457 - )

458 -

459 - return ctx.request.to_pooling_params()

460 -

461 - async def _prepare_generators(

462 - self,

463 - ctx: ServeContext,

464 - ) -> ErrorResponse | None:

465 - """Schedule the request and get the result generator."""

466 - generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []

467 -

468 - trace_headers = (

469 - None

470 - if ctx.raw_request is None

471 - else await self._get_trace_headers(ctx.raw_request.headers)

472 - )

473 -

474 - pooling_params = self._create_pooling_params(ctx)

475 - if isinstance(pooling_params, ErrorResponse):

476 - return pooling_params

477 -

478 - if ctx.engine_inputs is None:

479 - return self.create_error_response("Engine prompts not available")

480 -

481 - for i, engine_input in enumerate(ctx.engine_inputs):

482 - request_id_item = f"{ctx.request_id}-{i}"

483 -

484 - self._log_inputs(

485 - request_id_item,

486 - engine_input,

487 - params=pooling_params,

488 - lora_request=ctx.lora_request,

489 - )

490 -

491 - generator = self.engine_client.encode(

492 - engine_input,

493 - pooling_params,

494 - request_id_item,

495 - lora_request=ctx.lora_request,

496 - trace_headers=trace_headers,

497 - priority=getattr(ctx.request, "priority", 0),

498 - )

499 -

500 - generators.append(generator)

501 -

502 - ctx.result_generator = merge_async_iterators(*generators)

503 -

504 - return None

505 -

506 - async def _collect_batch(

507 - self,

508 - ctx: ServeContext,

509 - ) -> ErrorResponse | None:

510 - """Collect batch results from the result generator."""

511 - if ctx.engine_inputs is None:

512 - return self.create_error_response("Engine prompts not available")

513 -

514 - num_prompts = len(ctx.engine_inputs)

515 - final_res_batch: list[PoolingRequestOutput | None]

516 - final_res_batch = [None] * num_prompts

517 -

518 - if ctx.result_generator is None:

519 - return self.create_error_response("Result generator not available")

520 -

521 - async for i, res in ctx.result_generator:

522 - final_res_batch[i] = res

523 -

524 - if None in final_res_batch:

525 - return self.create_error_response(

526 - "Failed to generate results for all prompts"

527 - )

528 -

529 - ctx.final_res_batch = [res for res in final_res_batch if res is not None]

530 -

531 - return None

532 -

358 @staticmethod

359 def create_error_response(

360 message: str | Exception,

@@ -719,7 +544,7 @@ class OpenAIServing:

544 self,

545 request_id: str,

546 inputs: PromptType | EngineInput,

722 - params: SamplingParams | PoolingParams | BeamSearchParams | None,

547 + params: SamplingParams | BeamSearchParams | None,

548 lora_request: LoRARequest | None,

549 ) -> None:

550 if self.request_logger is None:

vllm/entrypoints/openai/models/serving.py

-1

@@ -112,7 +112,6 @@ class OpenAIServingModels:

112

113 self.model_config = self.engine_client.model_config

114 self.renderer = self.engine_client.renderer

115 - self.io_processor = self.engine_client.io_processor

115 self.input_processor = self.engine_client.input_processor

116

117 async def init_static_loras(self):

vllm/entrypoints/pooling/__init__.py

+2 -4

@@ -67,20 +67,18 @@ def init_pooling_state(

67 from vllm.entrypoints.chat_utils import load_chat_template

68 from vllm.entrypoints.pooling.classify.serving import ServingClassification

69 from vllm.entrypoints.pooling.embed.serving import ServingEmbedding

70 - from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling

70 + from vllm.entrypoints.pooling.pooling.serving import ServingPooling

71 from vllm.entrypoints.pooling.scoring.serving import ServingScores

72 from vllm.tasks import POOLING_TASKS

74 model_config = engine_client.model_config

75 -

75 resolved_chat_template = load_chat_template(args.chat_template)

77 state.serving_pooling = (

78 (

80 - OpenAIServingPooling(

79 + ServingPooling(

80 engine_client,

81 state.openai_serving_models,

83 - state.openai_serving_render,

82 supported_tasks=supported_tasks,

83 request_logger=request_logger,

84 chat_template=resolved_chat_template,

vllm/entrypoints/pooling/base/io_processor.py

+30 -9

@@ -4,8 +4,8 @@

4 from collections.abc import Sequence

5 from typing import Any, Final

7 -from vllm import PoolingRequestOutput, PromptType

8 -from vllm.config import ModelConfig

7 +from vllm import PoolingParams, PoolingRequestOutput, PromptType

8 +from vllm.config import VllmConfig

9 from vllm.entrypoints.chat_utils import (

10 ChatCompletionMessageParam,

11 ChatTemplateConfig,

@@ -33,11 +33,12 @@ class PoolingIOProcessor:

34 def __init__(

35 self,

36 - model_config: ModelConfig,

36 + vllm_config: VllmConfig,

37 renderer: BaseRenderer,

38 chat_template_config: ChatTemplateConfig,

39 ):

40 - self.model_config = model_config

40 + self.vllm_config = vllm_config

41 + self.model_config = vllm_config.model_config

42 self.renderer = renderer

44 self.chat_template = chat_template_config.chat_template

@@ -48,12 +49,12 @@ class PoolingIOProcessor:

49 chat_template_config.trust_request_chat_template

50 )

51 - def create_pooling_params(self, request):

52 - return request.to_pooling_params()

53 -

52 #######################################

53 # online APIs

55 + def create_pooling_params(self, request):

56 + return request.to_pooling_params()

57 +

58 def pre_process_online(self, ctx: PoolingServeContext):

59 request = ctx.request

@@ -100,12 +101,16 @@ class PoolingIOProcessor:

101 # offline APIs

102

103 def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:

103 - assert not isinstance(ctx.prompts, ScoringData)

104 + assert not isinstance(ctx.prompts, ScoringData) and not (

105 + isinstance(ctx.prompts, dict) and "data" in ctx.prompts

106 + )

107 +

108 + prompts_seq = prompt_to_seq(ctx.prompts)

109 tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(

110 **(ctx.tokenization_kwargs or {})

111 )

112 return self._preprocess_completion_offline(

108 - prompts=ctx.prompts, tok_params=tok_params

113 + prompts=prompts_seq, tok_params=tok_params

114 )

115

116 async def pre_process_offline_async(self, ctx: OfflineInputsContext):

@@ -243,3 +248,19 @@ class PoolingIOProcessor:

248 "Refused request with untrusted chat template."

249 )

250 return None

251 +

252 + def _params_to_seq(

253 + self,

254 + params: PoolingParams | Sequence[PoolingParams],

255 + num_requests: int,

256 + ) -> Sequence[PoolingParams]:

257 + if isinstance(params, Sequence):

258 + if len(params) != num_requests:

259 + raise ValueError(

260 + f"The lengths of prompts ({num_requests}) "

261 + f"and params ({len(params)}) must be the same."

262 + )

263 +

264 + return params

265 +

266 + return [params] * num_requests

vllm/entrypoints/pooling/base/serving.py

+47 -26

@@ -1,5 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from abc import ABC, abstractmethod

5 from collections.abc import AsyncGenerator, Mapping

6 from http import HTTPStatus

7 from typing import ClassVar

@@ -9,7 +11,7 @@ from fastapi.responses import Response

11 from starlette.datastructures import Headers

13 from vllm import PoolingParams, PoolingRequestOutput, envs

12 -from vllm.config import ModelConfig

14 +from vllm.config import VllmConfig

15 from vllm.engine.protocol import EngineClient

16 from vllm.entrypoints.chat_utils import (

17 ChatTemplateConfig,

@@ -35,7 +37,7 @@ from vllm.utils.async_utils import merge_async_iterators

37 from .io_processor import PoolingIOProcessor

38 -class PoolingServing:

40 +class PoolingServingBase(ABC):

41 request_id_prefix: ClassVar[str]

43 def __init__(

@@ -50,10 +52,11 @@ class PoolingServing:

52 return_tokens_as_token_ids: bool = False,

53 log_error_stack: bool = False,

54 ):

53 - super().__init__()

55 self.engine_client = engine_client

56 self.models = models

57 self.model_config = models.model_config

58 + self.renderer = models.renderer

59 + self.vllm_config = engine_client.vllm_config

60 self.max_model_len = self.model_config.max_model_len

61 self.request_logger = request_logger

62 self.return_tokens_as_token_ids = return_tokens_as_token_ids

@@ -63,31 +66,14 @@ class PoolingServing:

66 chat_template_content_format=chat_template_content_format,

67 trust_request_chat_template=trust_request_chat_template,

68 )

66 - self.io_processor = self.init_io_processor(

67 - model_config=models.model_config,

68 - renderer=models.renderer,

69 - chat_template_config=self.chat_template_config,

70 - )

71 -

72 - def init_io_processor(

73 - self,

74 - model_config: ModelConfig,

75 - renderer: BaseRenderer,

76 - chat_template_config: ChatTemplateConfig,

77 - ) -> PoolingIOProcessor:

78 - raise NotImplementedError

70 + @abstractmethod

71 async def __call__(

72 self,

73 request: AnyPoolingRequest,

74 raw_request: Request | None = None,

75 ) -> Response:

85 - ctx = await self._init_ctx(request, raw_request)

86 - await self.io_processor.pre_process_online_async(ctx)

87 - await self._prepare_generators(ctx)

88 - await self._collect_batch(ctx)

89 - await self.io_processor.post_process_online_async(ctx)

90 - return await self._build_response(ctx)

76 + raise NotImplementedError

78 async def _init_ctx(

79 self,

@@ -124,10 +110,8 @@ class PoolingServing:

110 else await self._get_trace_headers(ctx.raw_request.headers)

111 )

112

127 - if ctx.pooling_params is None:

128 - pooling_params = self.io_processor.create_pooling_params(ctx.request)

129 - else:

130 - pooling_params = ctx.pooling_params

113 + assert ctx.pooling_params is not None

114 + pooling_params = ctx.pooling_params

115

116 if isinstance(pooling_params, list):

117 for params in pooling_params:

@@ -190,6 +174,7 @@ class PoolingServing:

174

175 ctx.final_res_batch = [res for res in final_res_batch if res is not None]

176

177 + @abstractmethod

178 async def _build_response(

179 self,

180 ctx: PoolingServeContext,

@@ -355,3 +340,39 @@ class PoolingServing:

340 params=params,

341 lora_request=lora_request,

342 )

343 +

344 +

345 +class PoolingServing(PoolingServingBase, ABC):

346 + def __init__(self, *args, **kwargs):

347 + super().__init__(*args, **kwargs)

348 +

349 + self.io_processor = self.init_io_processor(

350 + vllm_config=self.vllm_config,

351 + renderer=self.renderer,

352 + chat_template_config=self.chat_template_config,

353 + )

354 +

355 + @abstractmethod

356 + def init_io_processor(

357 + self,

358 + vllm_config: VllmConfig,

359 + renderer: BaseRenderer,

360 + chat_template_config: ChatTemplateConfig,

361 + ) -> PoolingIOProcessor:

362 + raise NotImplementedError

363 +

364 + async def __call__(

365 + self,

366 + request: AnyPoolingRequest,

367 + raw_request: Request | None = None,

368 + ) -> Response:

369 + ctx = await self._init_ctx(request, raw_request)

370 + await self.io_processor.pre_process_online_async(ctx)

371 +

372 + if ctx.pooling_params is None:

373 + ctx.pooling_params = self.io_processor.create_pooling_params(request)

374 +

375 + await self._prepare_generators(ctx)

376 + await self._collect_batch(ctx)

377 + await self.io_processor.post_process_online_async(ctx)

378 + return await self._build_response(ctx)

vllm/entrypoints/pooling/classify/io_processor.py

+5 -1

@@ -5,4 +5,8 @@ from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor

7 class ClassifyIOProcessor(PoolingIOProcessor):

8 - name = "classification"

8 + name = "classify"

9 +

10 +

11 +class TokenClassifyIOProcessor(PoolingIOProcessor):

12 + name = "token_classify"

vllm/entrypoints/pooling/classify/serving.py

+2 -14

@@ -6,14 +6,11 @@ from typing import TypeAlias

6 import numpy as np

7 from fastapi.responses import JSONResponse

9 -from vllm.config import ModelConfig

10 -from vllm.entrypoints.chat_utils import ChatTemplateConfig

9 from vllm.entrypoints.openai.engine.protocol import UsageInfo

10 from vllm.entrypoints.pooling.base.serving import PoolingServing

11 from vllm.entrypoints.pooling.typing import PoolingServeContext

12 from vllm.logger import init_logger

13 from vllm.outputs import ClassificationOutput

16 -from vllm.renderers import BaseRenderer

15 from .io_processor import ClassifyIOProcessor

16 from .protocol import (

@@ -31,17 +28,8 @@ ClassificationServeContext: TypeAlias = PoolingServeContext[ClassificationReques

28 class ServingClassification(PoolingServing):

29 request_id_prefix = "classify"

34 - def init_io_processor(

35 - self,

36 - model_config: ModelConfig,

37 - renderer: BaseRenderer,

38 - chat_template_config: ChatTemplateConfig,

39 - ) -> ClassifyIOProcessor:

40 - return ClassifyIOProcessor(

41 - model_config=model_config,

42 - renderer=renderer,

43 - chat_template_config=chat_template_config,

44 - )

31 + def init_io_processor(self, *args, **kwargs) -> ClassifyIOProcessor:

32 + return ClassifyIOProcessor(*args, **kwargs)

34 async def _build_response(

35 self,

vllm/entrypoints/pooling/embed/io_processor.py

+5 -1

@@ -37,7 +37,7 @@ logger = init_logger(__name__)

39 class EmbedIOProcessor(PoolingIOProcessor):

40 - name = "embedding"

40 + name = "embed"

42 def __init__(self, *args, **kwargs):

43 super().__init__(*args, **kwargs)

@@ -549,3 +549,7 @@ class EmbedIOProcessor(PoolingIOProcessor):

549 request = ctx.request

550 if request.truncate == "NONE" and request.max_tokens is not None:

551 self._check_cohere_max_tokens(ctx.final_res_batch, request.max_tokens)

552 +

553 +

554 +class TokenEmbedIOProcessor(PoolingIOProcessor):

555 + name = "token_embed"

vllm/entrypoints/pooling/embed/serving.py

+10 -18

@@ -8,8 +8,6 @@ from typing import Literal, TypeAlias, cast

8 from fastapi.responses import JSONResponse, Response, StreamingResponse

9 from typing_extensions import assert_never

11 -from vllm.config import ModelConfig

12 -from vllm.entrypoints.chat_utils import ChatTemplateConfig

11 from vllm.entrypoints.openai.engine.protocol import UsageInfo

12 from vllm.entrypoints.pooling.base.serving import PoolingServing

13 from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor

@@ -33,12 +31,10 @@ from vllm.entrypoints.pooling.utils import (

31 )

32 from vllm.logger import init_logger

33 from vllm.outputs import PoolingRequestOutput

36 -from vllm.renderers import BaseRenderer

34 from vllm.utils.serial_utils import EmbedDType, Endianness

36 logger = init_logger(__name__)

41 -JSONResponseCLS = get_json_response_cls()

39 EmbeddingServeContext: TypeAlias = PoolingServeContext[EmbeddingRequest]

@@ -49,17 +45,13 @@ class ServingEmbedding(PoolingServing):

45 request_id_prefix = "embd"

46 io_processor: EmbedIOProcessor

52 - def init_io_processor(

53 - self,

54 - model_config: ModelConfig,

55 - renderer: BaseRenderer,

56 - chat_template_config: ChatTemplateConfig,

57 - ) -> EmbedIOProcessor:

58 - return EmbedIOProcessor(

59 - model_config=model_config,

60 - renderer=renderer,

61 - chat_template_config=chat_template_config,

62 - )

48 + def __init__(self, *args, **kwargs):

49 + super().__init__(*args, **kwargs)

50 +

51 + self.json_response_cls = get_json_response_cls()

52 +

53 + def init_io_processor(self, *args, **kwargs) -> EmbedIOProcessor:

54 + return EmbedIOProcessor(*args, **kwargs)

56 async def _build_response(

57 self,

@@ -149,7 +141,7 @@ class ServingEmbedding(PoolingServing):

141 data=items,

142 usage=usage,

143 )

152 - return JSONResponseCLS(content=response.model_dump())

144 + return self.json_response_cls(content=response.model_dump())

145

146 def _openai_bytes_response(

147 self,

@@ -190,8 +182,8 @@ class ServingEmbedding(PoolingServing):

182 media_type=response.media_type,

183 )

184

193 - @staticmethod

185 def _build_cohere_response_from_ctx(

186 + self,

187 ctx: PoolingServeContext,

188 ) -> JSONResponse:

189 request = ctx.request

@@ -218,4 +210,4 @@ class ServingEmbedding(PoolingServing):

210 ),

211 ),

212 )

221 - return JSONResponse(content=response.model_dump(exclude_none=True))

213 + return self.json_response_cls(content=response.model_dump(exclude_none=True))

vllm/entrypoints/pooling/io_processor_factories.py

+41 -14

@@ -1,42 +1,69 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -

5 -from vllm.config import ModelConfig

4 +from vllm.config import VllmConfig

5 from vllm.entrypoints.chat_utils import ChatTemplateConfig

7 -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor

8 -from vllm.entrypoints.pooling.scoring.io_processor import ScoringIOProcessors

9 -from vllm.entrypoints.pooling.utils import enable_scoring_api

6 +from vllm.plugins.io_processors import has_io_processor

7 from vllm.renderers import BaseRenderer

8 from vllm.tasks import SupportedTask

10 +from .base.io_processor import PoolingIOProcessor

11 +from .utils import enable_scoring_api

12 +

14 def init_pooling_io_processors(

15 supported_tasks: tuple[SupportedTask, ...],

16 - model_config: ModelConfig,

16 + vllm_config: VllmConfig,

17 renderer: BaseRenderer,

18 chat_template_config: ChatTemplateConfig,

19 ) -> dict[str, PoolingIOProcessor]:

20 - processors: list[tuple[str, type[PoolingIOProcessor]]] = []

20 + model_config = vllm_config.model_config

21 + processors: dict[str, type[PoolingIOProcessor]] = {}

22 +

23 if "classify" in supported_tasks:

22 - from vllm.entrypoints.pooling.classify.io_processor import ClassifyIOProcessor

24 + from .classify.io_processor import ClassifyIOProcessor

25 +

26 + processors["classify"] = ClassifyIOProcessor

27 +

28 + if "token_classify" in supported_tasks:

29 + from .classify.io_processor import TokenClassifyIOProcessor

30 +

31 + processors["token_classify"] = TokenClassifyIOProcessor

24 - processors.append(("classify", ClassifyIOProcessor))

33 if "embed" in supported_tasks:

26 - from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor

34 + from .embed.io_processor import EmbedIOProcessor

28 - processors.append(("embed", EmbedIOProcessor))

36 + processors["embed"] = EmbedIOProcessor

37 +

38 + if "token_embed" in supported_tasks:

39 + from .embed.io_processor import TokenEmbedIOProcessor

40 +

41 + processors["token_embed"] = TokenEmbedIOProcessor

42 +

43 + if has_io_processor(

44 + vllm_config,

45 + model_config.io_processor_plugin,

46 + ):

47 + from .pooling.io_processor import PluginWithIOProcessorPlugins

48 +

49 + processors["plugin"] = PluginWithIOProcessorPlugins

50 + elif "plugin" in supported_tasks:

51 + from .pooling.io_processor import PluginWithoutIOProcessorPlugins

52 +

53 + processors["plugin"] = PluginWithoutIOProcessorPlugins

55 if enable_scoring_api(supported_tasks, model_config):

56 score_type = model_config.score_type

57 + from .scoring.io_processor import ScoringIOProcessors

58 +

59 if score_type is not None and score_type in ScoringIOProcessors:

33 - processors.append((score_type, ScoringIOProcessors[score_type]))

60 + processors[score_type] = ScoringIOProcessors[score_type]

62 return {

63 task: processor_cls(

37 - model_config=model_config,

64 + vllm_config=vllm_config,

65 renderer=renderer,

66 chat_template_config=chat_template_config,

67 )

41 - for task, processor_cls in processors

68 + for task, processor_cls in processors.items()

69 }

vllm/entrypoints/pooling/pooling/api_router.py

+4 -26

@@ -3,24 +3,17 @@

3 from http import HTTPStatus

5 from fastapi import APIRouter, Depends, Request

6 -from fastapi.responses import JSONResponse, StreamingResponse

7 -from typing_extensions import assert_never

7 from vllm.entrypoints.openai.engine.protocol import ErrorResponse

8 from vllm.entrypoints.openai.utils import validate_json_request

11 -from vllm.entrypoints.pooling.pooling.protocol import (

12 - IOProcessorResponse,

13 - PoolingBytesResponse,

14 - PoolingRequest,

15 - PoolingResponse,

16 -)

17 -from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling

9 +from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest

10 +from vllm.entrypoints.pooling.pooling.serving import ServingPooling

11 from vllm.entrypoints.utils import load_aware_call, with_cancellation

13 router = APIRouter()

23 -def pooling(request: Request) -> OpenAIServingPooling | None:

16 +def pooling(request: Request) -> ServingPooling | None:

17 return request.app.state.serving_pooling

@@ -39,19 +32,4 @@ async def create_pooling(request: PoolingRequest, raw_request: Request):

32 if handler is None:

33 raise NotImplementedError("The model does not support Pooling API")

42 - generator = await handler.create_pooling(request, raw_request)

43 -

44 - if isinstance(generator, ErrorResponse):

45 - return JSONResponse(

46 - content=generator.model_dump(), status_code=generator.error.code

47 - )

48 - elif isinstance(generator, (PoolingResponse, IOProcessorResponse)):

49 - return JSONResponse(content=generator.model_dump())

50 - elif isinstance(generator, PoolingBytesResponse):

51 - return StreamingResponse(

52 - content=generator.content,

53 - headers=generator.headers,

54 - media_type=generator.media_type,

55 - )

56 -

57 - assert_never(generator)

35 + return await handler(request, raw_request)

vllm/entrypoints/pooling/pooling/io_processor.py new

+156

@@ -0,0 +1,156 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from collections.abc import Sequence

4 +from typing import Any

5 +

6 +from vllm import PoolingParams, PoolingRequestOutput

7 +from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor

8 +from vllm.inputs import EngineInput

9 +from vllm.logger import init_logger

10 +from vllm.plugins.io_processors import get_io_processor

11 +from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq

12 +

13 +from ..typing import OfflineInputsContext, OfflineOutputsContext, PoolingServeContext

14 +from .protocol import IOProcessorRequest, IOProcessorResponse

15 +

16 +logger = init_logger(__name__)

17 +

18 +

19 +class PluginWithoutIOProcessorPlugins(PoolingIOProcessor):

20 + name = "plugin"

21 +

22 +

23 +class PluginWithIOProcessorPlugins(PoolingIOProcessor):

24 + """IO Processor plugins are a feature that allows pre- and post-processing

25 + of the model input and output for pooling models."""

26 +

27 + name = "plugin"

28 +

29 + def __init__(self, *args, **kwargs):

30 + super().__init__(*args, **kwargs)

31 +

32 + io_processor = get_io_processor(

33 + self.vllm_config,

34 + self.renderer,

35 + self.model_config.io_processor_plugin,

36 + )

37 +

38 + assert io_processor is not None

39 + self.io_processor = io_processor

40 +

41 + #######################################

42 + # online APIs

43 +

44 + def pre_process_online(self, ctx: PoolingServeContext):

45 + assert isinstance(ctx.request, IOProcessorRequest)

46 +

47 + validated_prompt = self.io_processor.parse_data(ctx.request.data)

48 +

49 + raw_prompts = self.io_processor.pre_process(

50 + prompt=validated_prompt, request_id=ctx.request_id

51 + )

52 +

53 + parsed_prompts = [

54 + (

55 + prompt

56 + if isinstance(prompt, bytes)

57 + else parse_model_prompt(self.model_config, prompt)

58 + )

59 + for prompt in prompt_to_seq(raw_prompts)

60 + ]

61 +

62 + tok_params = ctx.request.build_tok_params(self.model_config)

63 +

64 + ctx.engine_inputs = self.renderer.render_cmpl(

65 + parsed_prompts,

66 + tok_params,

67 + prompt_extras={

68 + k: v

69 + for k in ("mm_processor_kwargs", "cache_salt")

70 + if (v := getattr(ctx.request, k, None)) is not None

71 + },

72 + )

73 +

74 + pooling_params = self.io_processor.merge_pooling_params()

75 + if pooling_params.task is None:

76 + pooling_params.task = "plugin"

77 + ctx.pooling_params = pooling_params

78 +

79 + def post_process_online(

80 + self,

81 + ctx: PoolingServeContext,

82 + ):

83 + output = self.io_processor.post_process(

84 + ctx.final_res_batch,

85 + request_id=ctx.request_id,

86 + )

87 +

88 + if callable(

89 + output_to_response := getattr(self.io_processor, "output_to_response", None)

90 + ):

91 + logger.warning_once(

92 + "`IOProcessor.output_to_response` is deprecated. To ensure "

93 + "consistency between offline and online APIs, "

94 + "`IOProcessorResponse` will become a transparent wrapper "

95 + "around output data from v0.19 onwards.",

96 + )

97 +

98 + if hasattr(output, "request_id") and output.request_id is None:

99 + output.request_id = ctx.request_id # type: ignore

100 +

101 + ctx.response = output_to_response(output) # type: ignore

102 + else:

103 + ctx.response = IOProcessorResponse(request_id=ctx.request_id, data=output)

104 +

105 + #######################################

106 + # offline APIs

107 +

108 + def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:

109 + assert isinstance(ctx.prompts, dict) and "data" in ctx.prompts

110 + assert ctx.pooling_params is not None

111 +

112 + # Validate the request data is valid for the loaded plugin

113 + prompt_data = ctx.prompts.get("data")

114 + if prompt_data is None:

115 + raise ValueError(

116 + "The 'data' field of the prompt is expected to contain "

117 + "the prompt data and it cannot be None. "

118 + "Refer to the documentation of the IOProcessor "

119 + "in use for more details."

120 + )

121 + validated_prompt = self.io_processor.parse_data(prompt_data)

122 +

123 + # obtain the actual model prompts from the pre-processor

124 + prompts = self.io_processor.pre_process(prompt=validated_prompt)

125 + prompts_seq = prompt_to_seq(prompts)

126 +

127 + params_seq: list[PoolingParams] = [

128 + self.io_processor.merge_pooling_params(param)

129 + for param in self._params_to_seq(

130 + ctx.pooling_params,

131 + len(prompts_seq),

132 + )

133 + ]

134 + for p in params_seq:

135 + if p.task is None:

136 + p.task = "plugin"

137 +

138 + ctx.pooling_params = params_seq

139 + ctx.prompts = prompts_seq

140 + return super().pre_process_offline(ctx)

141 +

142 + def post_process_offline(

143 + self,

144 + ctx: OfflineOutputsContext,

145 + ) -> list[PoolingRequestOutput]:

146 + processed_outputs = self.io_processor.post_process(ctx.outputs)

147 +

148 + return [

149 + PoolingRequestOutput[Any](

150 + request_id="",

151 + outputs=processed_outputs,

152 + num_cached_tokens=getattr(processed_outputs, "num_cached_tokens", 0),

153 + prompt_token_ids=[],

154 + finished=True,

155 + )

156 + ]

vllm/entrypoints/pooling/pooling/serving.py

+101 -223

@@ -1,252 +1,157 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -import asyncio

4 +

5 import json

6 -import time

7 -from collections.abc import AsyncGenerator, Callable, Sequence

6 +from collections.abc import Callable

7 from functools import partial

9 -from typing import Final, Literal, cast

8 +from typing import Literal, cast

10 from fastapi import Request

11 +from fastapi.responses import JSONResponse, Response, StreamingResponse

12 from typing_extensions import assert_never

14 -from vllm.engine.protocol import EngineClient

15 -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption

16 -from vllm.entrypoints.logger import RequestLogger

17 -from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo

18 -from vllm.entrypoints.openai.engine.serving import OpenAIServing

19 -from vllm.entrypoints.openai.models.serving import OpenAIServingModels

14 +from vllm.entrypoints.openai.engine.protocol import UsageInfo

15 +from vllm.entrypoints.pooling.base.serving import PoolingServingBase

16 +from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors

17 from vllm.entrypoints.pooling.pooling.protocol import (

18 IOProcessorRequest,

22 - IOProcessorResponse,

19 PoolingBytesResponse,

24 - PoolingChatRequest,

25 - PoolingCompletionRequest,

20 PoolingRequest,

21 PoolingResponse,

22 PoolingResponseData,

23 )

24 +from vllm.entrypoints.pooling.typing import AnyPoolingRequest, PoolingServeContext

25 from vllm.entrypoints.pooling.utils import (

26 encode_pooling_bytes,

27 encode_pooling_output_base64,

28 encode_pooling_output_float,

29 + get_json_response_cls,

30 )

35 -from vllm.entrypoints.serve.render.serving import OpenAIServingRender

36 -from vllm.inputs import EngineInput

31 from vllm.logger import init_logger

32 from vllm.outputs import PoolingRequestOutput

39 -from vllm.renderers.inputs.preprocess import prompt_to_seq

33 from vllm.tasks import SupportedTask

41 -from vllm.utils.async_utils import merge_async_iterators

42 -from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness

34 +from vllm.utils.serial_utils import EmbedDType, Endianness

36 logger = init_logger(__name__)

47 -class OpenAIServingPooling(OpenAIServing):

39 +class ServingPooling(PoolingServingBase):

40 + request_id_prefix = "pooling"

41 +

42 def __init__(

43 self,

50 - engine_client: EngineClient,

51 - models: OpenAIServingModels,

52 - openai_serving_render: OpenAIServingRender,

44 + *args,

45 supported_tasks: tuple[SupportedTask, ...],

54 - *,

55 - request_logger: RequestLogger | None,

56 - chat_template: str | None,

57 - chat_template_content_format: ChatTemplateContentFormatOption,

58 - trust_request_chat_template: bool = False,

59 - ) -> None:

60 - super().__init__(

61 - engine_client=engine_client,

62 - models=models,

63 - request_logger=request_logger,

64 - )

46 + **kwargs,

47 + ):

48 + super().__init__(*args, **kwargs)

49 +

50 self.supported_tasks = supported_tasks

51 self.pooling_task = self.model_config.get_pooling_task(supported_tasks)

67 - self.openai_serving_render = openai_serving_render

68 - self.chat_template = chat_template

69 - self.chat_template_content_format: Final = chat_template_content_format

70 - self.trust_request_chat_template = trust_request_chat_template

52 + self.io_processors = init_pooling_io_processors(

53 + supported_tasks=supported_tasks,

54 + vllm_config=self.vllm_config,

55 + renderer=self.renderer,

56 + chat_template_config=self.chat_template_config,

57 + )

58 + self.json_response_cls = get_json_response_cls()

72 - async def create_pooling(

60 + async def __call__(

61 self,

74 - request: PoolingRequest,

62 + request: AnyPoolingRequest,

63 raw_request: Request | None = None,

76 - ) -> PoolingResponse | IOProcessorResponse | PoolingBytesResponse | ErrorResponse:

77 - """

78 - See https://platform.openai.com/docs/api-reference/embeddings/create

79 - for the API specification. This API mimics the OpenAI Embedding API.

80 - """

81 - error_check_ret = await self._check_model(request)

82 - if error_check_ret is not None:

83 - return error_check_ret

64 + ) -> Response:

65 + assert isinstance(request, PoolingRequest)

66 + pooling_task = self._verify_pooling_task(request)

85 - model_name = self.models.model_name()

68 + io_processor = self.io_processors[pooling_task]

69 + ctx = await self._init_ctx(request, raw_request)

87 - request_id = f"pool-{self._base_request_id(raw_request)}"

88 - created_time = int(time.time())

71 + await io_processor.pre_process_online_async(ctx)

90 - lora_request = self._maybe_get_adapters(request)

73 + if ctx.pooling_params is None:

74 + ctx.pooling_params = io_processor.create_pooling_params(request)

75 +

76 + await self._prepare_generators(ctx)

77 + await self._collect_batch(ctx)

78 +

79 + await io_processor.post_process_online_async(ctx)

80 + return await self._build_response(ctx)

81 +

82 + def _verify_pooling_task(self, request: PoolingRequest) -> str:

83 + if getattr(request, "dimensions", None) is not None:

84 + raise ValueError("dimensions is currently not supported")

86 if request.task is None:

87 request.task = self.pooling_task

95 - if getattr(request, "dimensions", None) is not None:

96 - return self.create_error_response("dimensions is currently not supported")

89 + if isinstance(request, IOProcessorRequest):

90 + request.task = "plugin"

91 +

92 + assert request.task is not None

93 + pooling_task = request.task

95 # plugin task uses io_processor.parse_request to verify inputs

99 - if request.task != "plugin" and request.task != self.pooling_task:

100 - if request.task not in self.supported_tasks:

96 + if pooling_task != "plugin" and pooling_task != self.pooling_task:

97 + if pooling_task not in self.io_processors:

98 raise ValueError(

102 - f"Unsupported task: {request.task!r} "

99 + f"Unsupported task: {pooling_task!r} "

100 f"Supported tasks: {self.supported_tasks}"

101 )

102 else:

103 logger.warning_once(

104 "Pooling multitask support is deprecated and will be removed "

105 "in v0.20. When the default pooling task is not what you want, you "

109 - 'need to manually specify it via --pooler-config.task "%s". ',

110 - request.task,

111 - )

112 -

113 - engine_inputs: Sequence[EngineInput]

114 - if use_io_processor := isinstance(request, IOProcessorRequest):

115 - if self.io_processor is None:

116 - raise ValueError(

117 - "No IOProcessor plugin installed. Please refer "

118 - "to the documentation and to the "

119 - "'prithvi_geospatial_mae_io_processor' "

120 - "offline inference example for more details."

106 + "need to manually specify it via --pooler-config.task %s. ",

107 + pooling_task,

108 )

109

123 - validated_prompt = self.io_processor.parse_data(request.data)

124 -

125 - raw_prompts = await self.io_processor.pre_process_async(

126 - prompt=validated_prompt, request_id=request_id

127 - )

128 - engine_inputs = await self.openai_serving_render.preprocess_cmpl(

129 - request,

130 - prompt_to_seq(raw_prompts),

131 - )

132 - elif isinstance(request, PoolingChatRequest):

133 - error_check_ret = self.openai_serving_render.validate_chat_template(

134 - request_chat_template=request.chat_template,

135 - chat_template_kwargs=request.chat_template_kwargs,

136 - trust_request_chat_template=self.trust_request_chat_template,

137 - )

138 - if error_check_ret is not None:

139 - return error_check_ret

140 -

141 - _, engine_inputs = await self.openai_serving_render.preprocess_chat(

142 - request,

143 - request.messages,

144 - default_template=self.chat_template,

145 - default_template_content_format=self.chat_template_content_format,

146 - default_template_kwargs=None,

147 - )

148 - elif isinstance(request, PoolingCompletionRequest):

149 - engine_inputs = await self.openai_serving_render.preprocess_completion(

150 - request,

151 - prompt_input=request.input,

152 - prompt_embeds=None,

153 - )

154 - else:

155 - raise ValueError(f"Unsupported request of type {type(request)}")

156 -

157 - # Schedule the request and get the result generator.

158 - generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []

159 - if use_io_processor:

160 - assert self.io_processor is not None

161 -

162 - pooling_params = self.io_processor.merge_pooling_params()

163 - if pooling_params.task is None:

164 - pooling_params.task = "plugin"

165 - else:

166 - pooling_params = request.to_pooling_params() # type: ignore

167 -

168 - for i, engine_input in enumerate(engine_inputs):

169 - request_id_item = f"{request_id}-{i}"

170 -

171 - self._log_inputs(

172 - request_id_item,

173 - engine_input,

174 - params=pooling_params,

175 - lora_request=lora_request,

110 + if pooling_task == "plugin" and "plugin" not in self.io_processors:

111 + raise ValueError(

112 + "No IOProcessor plugin installed. Please refer "

113 + "to the documentation and to the "

114 + "'prithvi_geospatial_mae_io_processor' "

115 + "offline inference example for more details."

116 )

117

178 - trace_headers = (

179 - None

180 - if raw_request is None

181 - else await self._get_trace_headers(raw_request.headers)

182 - )

183 -

184 - generator = self.engine_client.encode(

185 - engine_input,

186 - pooling_params,

187 - request_id_item,

188 - lora_request=lora_request,

189 - trace_headers=trace_headers,

190 - priority=request.priority,

191 - )

118 + return pooling_task

119

193 - generators.append(generator)

120 + async def _build_response(

121 + self,

122 + ctx: PoolingServeContext,

123 + ) -> Response:

124 + if ctx.response is not None:

125 + # for IOProcessorResponse

126 + return self.json_response_cls(content=ctx.response.model_dump())

127

195 - result_generator = merge_async_iterators(*generators)

128 + encoding_format = ctx.request.encoding_format

129 + embed_dtype = ctx.request.embed_dtype

130 + endianness = ctx.request.endianness

131

197 - if use_io_processor:

198 - assert self.io_processor is not None

199 - output = await self.io_processor.post_process_async(

200 - result_generator,

201 - request_id=request_id,

132 + if encoding_format == "float" or encoding_format == "base64":

133 + return self.request_output_to_pooling_json_response(

134 + ctx.final_res_batch,

135 + ctx.request_id,

136 + ctx.created_time,

137 + ctx.model_name,

138 + encoding_format,

139 + embed_dtype,

140 + endianness,

141 )

142

204 - if callable(

205 - output_to_response := getattr(

206 - self.io_processor, "output_to_response", None

207 - )

208 - ):

209 - logger.warning_once(

210 - "`IOProcessor.output_to_response` is deprecated. To ensure "

211 - "consistency between offline and online APIs, "

212 - "`IOProcessorResponse` will become a transparent wrapper "

213 - "around output data from v0.19 onwards.",

214 - )

215 -

216 - if hasattr(output, "request_id") and output.request_id is None:

217 - output.request_id = request_id # type: ignore

218 -

219 - return output_to_response(output) # type: ignore

220 -

221 - return IOProcessorResponse(request_id=request_id, data=output)

222 -

223 - assert isinstance(request, (PoolingCompletionRequest, PoolingChatRequest))

224 - num_prompts = len(engine_inputs)

225 -

226 - # Non-streaming response

227 - final_res_batch: list[PoolingRequestOutput | None]

228 - final_res_batch = [None] * num_prompts

229 - try:

230 - async for i, res in result_generator:

231 - final_res_batch[i] = res

232 -

233 - assert all(final_res is not None for final_res in final_res_batch)

234 -

235 - final_res_batch_checked = cast(list[PoolingRequestOutput], final_res_batch)

236 -

237 - response = self.request_output_to_pooling_response(

238 - final_res_batch_checked,

239 - request_id,

240 - created_time,

241 - model_name,

242 - request.encoding_format,

243 - request.embed_dtype,

244 - request.endianness,

143 + if encoding_format == "bytes" or encoding_format == "bytes_only":

144 + return self.request_output_to_pooling_bytes_response(

145 + ctx.final_res_batch,

146 + ctx.request_id,

147 + ctx.created_time,

148 + ctx.model_name,

149 + encoding_format,

150 + embed_dtype,

151 + endianness,

152 )

246 - except asyncio.CancelledError:

247 - return self.create_error_response("Client disconnected")

153

249 - return response

154 + assert_never(encoding_format)

155

156 def request_output_to_pooling_json_response(

157 self,

@@ -257,7 +162,7 @@ class OpenAIServingPooling(OpenAIServing):

162 encoding_format: Literal["float", "base64"],

163 embed_dtype: EmbedDType,

164 endianness: Endianness,

260 - ) -> PoolingResponse:

165 + ) -> JSONResponse:

166 encode_fn = cast(

167 Callable[[PoolingRequestOutput], list[float] | str],

168 (

@@ -289,13 +194,14 @@ class OpenAIServingPooling(OpenAIServing):

194 total_tokens=num_prompt_tokens,

195 )

196

292 - return PoolingResponse(

197 + response = PoolingResponse(

198 id=request_id,

199 created=created_time,

200 model=model_name,

201 data=items,

202 usage=usage,

203 )

204 + return self.json_response_cls(content=response.model_dump())

205

206 def request_output_to_pooling_bytes_response(

207 self,

@@ -306,7 +212,7 @@ class OpenAIServingPooling(OpenAIServing):

212 encoding_format: Literal["bytes", "bytes_only"],

213 embed_dtype: EmbedDType,

214 endianness: Endianness,

309 - ) -> PoolingBytesResponse:

215 + ) -> StreamingResponse:

216 content, items, usage = encode_pooling_bytes(

217 pooling_outputs=final_res_batch,

218 embed_dtype=embed_dtype,

@@ -329,38 +235,10 @@ class OpenAIServingPooling(OpenAIServing):

235 }

236 )

237

332 - return PoolingBytesResponse(content=content, headers=headers)

333 -

334 - def request_output_to_pooling_response(

335 - self,

336 - final_res_batch: list[PoolingRequestOutput],

337 - request_id: str,

338 - created_time: int,

339 - model_name: str,

340 - encoding_format: EncodingFormat,

341 - embed_dtype: EmbedDType,

342 - endianness: Endianness,

343 - ) -> PoolingResponse | PoolingBytesResponse:

344 - if encoding_format == "float" or encoding_format == "base64":

345 - return self.request_output_to_pooling_json_response(

346 - final_res_batch,

347 - request_id,

348 - created_time,

349 - model_name,

350 - encoding_format,

351 - embed_dtype,

352 - endianness,

353 - )

354 -

355 - if encoding_format == "bytes" or encoding_format == "bytes_only":

356 - return self.request_output_to_pooling_bytes_response(

357 - final_res_batch,

358 - request_id,

359 - created_time,

360 - model_name,

361 - encoding_format,

362 - embed_dtype,

363 - endianness,

364 - )

238 + response = PoolingBytesResponse(content=content, headers=headers)

239

366 - assert_never(encoding_format)

240 + return StreamingResponse(

241 + content=response.content,

242 + headers=response.headers,

243 + media_type=response.media_type,

244 + )

vllm/entrypoints/pooling/scoring/io_processor.py

+1 -1

@@ -278,7 +278,7 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):

278

279 def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:

280 assert isinstance(ctx.prompts, ScoringData)

281 - assert not isinstance(ctx.pooling_params, list)

281 + assert not isinstance(ctx.pooling_params, Sequence)

282

283 tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(

284 **(ctx.tokenization_kwargs or {})

vllm/entrypoints/pooling/scoring/serving.py

+5 -12

@@ -4,15 +4,13 @@

4 from fastapi.responses import JSONResponse, Response

6 from vllm import PoolingParams

7 -from vllm.config import ModelConfig

7 +from vllm.config import VllmConfig

8 from vllm.engine.protocol import EngineClient

9 -from vllm.entrypoints.chat_utils import ChatTemplateConfig

9 from vllm.entrypoints.openai.engine.protocol import UsageInfo

10 from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor

11 from vllm.entrypoints.pooling.base.serving import PoolingServing

12 from vllm.logger import init_logger

13 from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput

15 -from vllm.renderers import BaseRenderer

14 from vllm.v1.pool.late_interaction import (

15 build_late_interaction_doc_params,

16 build_late_interaction_query_params,

@@ -52,22 +50,17 @@ class ServingScores(PoolingServing):

50 super().__init__(engine_client, *args, **kwargs)

52 def init_io_processor(

55 - self,

56 - model_config: ModelConfig,

57 - renderer: BaseRenderer,

58 - chat_template_config: ChatTemplateConfig,

53 + self, vllm_config: VllmConfig, *args, **kwargs

54 ) -> PoolingIOProcessor:

55 + model_config = vllm_config.model_config

56 +

57 score_type: str = model_config.score_type

58 if self.enable_flash_late_interaction:

59 score_type = "flash-late-interaction"

61 assert score_type in ScoringIOProcessors

62 processor_cls = ScoringIOProcessors[score_type]

66 - return processor_cls(

67 - model_config=model_config,

68 - renderer=renderer,

69 - chat_template_config=chat_template_config,

70 - )

63 + return processor_cls(vllm_config, *args, **kwargs)

65 async def __call__(self, *args, **kwargs) -> Response:

66 if not self.enable_flash_late_interaction:

vllm/entrypoints/pooling/typing.py

+6 -3

@@ -30,7 +30,7 @@ from vllm.entrypoints.pooling.pooling.protocol import (

30 )

31 from vllm.entrypoints.pooling.scoring.protocol import ScoringRequest, ScoringResponse

32 from vllm.entrypoints.pooling.scoring.typing import ScoringData

33 -from vllm.inputs import EngineInput

33 +from vllm.inputs import DataPrompt, EngineInput

34 from vllm.lora.request import LoRARequest

36 PoolingCompletionLikeRequest: TypeAlias = (

@@ -86,11 +86,14 @@ class PoolingServeContext(Generic[PoolingRequestT]):

86 ## for bi-encoder & late-interaction

87 n_queries: int | None = None

89 + ## for IOProcessorResponse

90 + response: Any | None = None

91 +

93 @dataclass

94 class OfflineInputsContext:

92 - prompts: PromptType | Sequence[PromptType] | ScoringData

93 - pooling_params: PoolingParams | list[PoolingParams] | None = None

95 + prompts: PromptType | Sequence[PromptType] | DataPrompt | ScoringData

96 + pooling_params: PoolingParams | Sequence[PoolingParams]

97 tokenization_kwargs: dict[str, Any] | None = None

98 chat_template: str | None = None

vllm/entrypoints/sagemaker/api_router.py

+2 -2

@@ -14,7 +14,7 @@ from vllm.config import ModelConfig

14 from vllm.entrypoints.openai.engine.protocol import ErrorResponse

15 from vllm.entrypoints.openai.engine.serving import OpenAIServing

16 from vllm.entrypoints.openai.utils import validate_json_request

17 -from vllm.entrypoints.pooling.base.serving import PoolingServing

17 +from vllm.entrypoints.pooling.base.serving import PoolingServingBase

18 from vllm.entrypoints.pooling.utils import enable_scoring_api

19 from vllm.entrypoints.serve.instrumentator.basic import base

20 from vllm.entrypoints.serve.instrumentator.health import health

@@ -23,7 +23,7 @@ from vllm.tasks import POOLING_TASKS, SupportedTask

23 # TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers

24 # (requires typing_extensions >= 4.13)

25 RequestType = Any

26 -GetHandlerFn = Callable[[Request], OpenAIServing | PoolingServing | None]

26 +GetHandlerFn = Callable[[Request], OpenAIServing | PoolingServingBase | None]

27 EndpointFn = Callable[[RequestType, Request], Awaitable[Any]]

vllm/entrypoints/serve/render/serving.py

-2

@@ -65,7 +65,6 @@ class OpenAIServingRender:

65 self,

66 model_config: ModelConfig,

67 renderer: BaseRenderer,

68 - io_processor: Any,

68 model_registry: OpenAIModelRegistry,

69 *,

70 request_logger: RequestLogger | None,

@@ -81,7 +80,6 @@ class OpenAIServingRender:

80 ) -> None:

81 self.model_config = model_config

82 self.renderer = renderer

84 - self.io_processor = io_processor

83 self.model_registry = model_registry

84 self.request_logger = request_logger

85 self.chat_template = chat_template

vllm/plugins/io_processors/__init__.py

+17

@@ -12,6 +12,23 @@ from vllm.utils.import_utils import resolve_obj_by_qualname

12 logger = logging.getLogger(__name__)

15 +def has_io_processor(

16 + vllm_config: VllmConfig,

17 + plugin_from_init: str | None = None,

18 +):

19 + if plugin_from_init:

20 + model_plugin = plugin_from_init

21 + else:

22 + # A plugin can be specified via the model config

23 + # Retrieve the model specific plugin if available

24 + # This is using a custom field in the hf_config for the model

25 + hf_config = vllm_config.model_config.hf_config.to_dict()

26 + config_plugin = hf_config.get("io_processor_plugin")

27 + model_plugin = config_plugin

28 +

29 + return model_plugin is not None

30 +

31 +

32 def get_io_processor(

33 vllm_config: VllmConfig,

34 renderer: BaseRenderer,

vllm/v1/engine/async_llm.py

-6

@@ -26,7 +26,6 @@ from vllm.logger import init_logger

26 from vllm.lora.request import LoRARequest

27 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry

28 from vllm.outputs import STREAM_FINISHED, PoolingRequestOutput, RequestOutput

29 -from vllm.plugins.io_processors import get_io_processor

29 from vllm.pooling_params import PoolingParams

30 from vllm.renderers import renderer_from_config

31 from vllm.renderers.inputs.preprocess import extract_prompt_components

@@ -133,11 +132,6 @@ class AsyncLLM(EngineClient):

132 )

133

134 self.renderer = renderer = renderer_from_config(self.vllm_config)

136 - self.io_processor = get_io_processor(

137 - self.vllm_config,

138 - self.renderer,

139 - self.model_config.io_processor_plugin,

140 - )

135

136 # Convert EngineInput --> EngineCoreRequest.

137 self.input_processor = InputProcessor(self.vllm_config, renderer)

vllm/v1/engine/llm_engine.py

-6

@@ -19,7 +19,6 @@ from vllm.logger import init_logger

19 from vllm.lora.request import LoRARequest

20 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry

21 from vllm.outputs import PoolingRequestOutput, RequestOutput

22 -from vllm.plugins.io_processors import get_io_processor

22 from vllm.pooling_params import PoolingParams

23 from vllm.renderers import renderer_from_config

24 from vllm.renderers.inputs.preprocess import extract_prompt_components

@@ -90,11 +89,6 @@ class LLMEngine:

89 self.should_execute_dummy_batch = False

91 self.renderer = renderer = renderer_from_config(self.vllm_config)

93 - self.io_processor = get_io_processor(

94 - self.vllm_config,

95 - self.renderer,

96 - self.model_config.io_processor_plugin,

97 - )

93 # Convert EngineInput --> EngineCoreRequest.

94 self.input_processor = InputProcessor(self.vllm_config, renderer)