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 8, 12:30 UTC
[Plugin] Simplify IO Processor Plugin interface (#34236)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Feb 11, 2026 at 11:47 UTC
c9a1923bb470f79a33963ad80cc8ad12bab2ad52
9 files changed
+167
-151
docs/design/io_processor_plugins.md
+21
-19
@@ -14,8 +14,26 @@ IOProcessorOutput = TypeVar("IOProcessorOutput")
14
15
class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):
16
def __init__(self, vllm_config: VllmConfig):
17
+ super().__init__()
18
+
19
self.vllm_config = vllm_config
20
21
+ @abstractmethod
22
+ def parse_data(self, data: object) -> IOProcessorInput:
23
+ raise NotImplementedError
24
+
25
+ def merge_sampling_params(
26
+ self,
27
+ params: SamplingParams | None = None,
28
+ ) -> SamplingParams:
29
+ return params or SamplingParams()
30
+
31
+ def merge_pooling_params(
32
+ self,
33
+ params: PoolingParams | None = None,
34
+ ) -> PoolingParams:
35
+ return params or PoolingParams()
36
+
37
@abstractmethod
38
def pre_process(
39
self,
@@ -55,29 +73,13 @@ class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):
73
[(i, item) async for i, item in model_output], key=lambda output: output[0]
74
)
75
collected_output = [output[1] for output in sorted_output]
58
- return self.post_process(collected_output, request_id, **kwargs)
59
-
60
- @abstractmethod
61
- def parse_request(self, request: Any) -> IOProcessorInput:
62
- raise NotImplementedError
63
-
64
- def validate_or_generate_params(
65
- self, params: SamplingParams | PoolingParams | None = None
66
- ) -> SamplingParams | PoolingParams:
67
- return params or PoolingParams()
68
-
69
- @abstractmethod
70
- def output_to_response(
71
- self, plugin_output: IOProcessorOutput
72
- ) -> IOProcessorResponse:
73
- raise NotImplementedError
76
+ return self.post_process(collected_output, request_id=request_id, **kwargs)
77
```
78
76
-The `parse_request` method is used for validating the user prompt and converting it into the input expected by the `pre_process`/`pre_process_async` methods.
79
+The `parse_data` method is used for validating the user data and converting it into the input expected by the `pre_process*` methods.
80
+The `merge_sampling_params` and `merge_pooling_params` methods merge input `SamplingParams` or `PoolingParams` (if any) with the default one.
81
The `pre_process*` methods take the validated plugin input to generate vLLM's model prompts for regular inference.
82
The `post_process*` methods take `PoolingRequestOutput` objects as input and generate a custom plugin output.
79
-The `validate_or_generate_params` method is used for validating with the plugin any `SamplingParameters`/`PoolingParameters` received with the user request, or to generate new ones if none are specified. The function always returns the validated/generated parameters.
80
-The `output_to_response` method is used only for online serving and converts the plugin output to the `IOProcessorResponse` type that is then returned by the API Server. The implementation of the `/pooling` serving endpoint is available here [vllm/entrypoints/openai/serving_pooling.py](../../vllm/entrypoints/pooling/pooling/serving.py).
83
84
An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/IBM/terratorch/tree/main/terratorch/vllm/plugins/segmentation). Please, also refer to our online ([examples/pooling/plugin/prithvi_geospatial_mae_online.py](../../examples/pooling/plugin/prithvi_geospatial_mae_online.py)) and offline ([examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py](../../examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py)) inference examples.
85
tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py
+12
-37
@@ -18,18 +18,10 @@ from einops import rearrange
18
from terratorch.datamodules import Sen1Floods11NonGeoDataModule
19
20
from vllm.config import VllmConfig
21
-from vllm.entrypoints.pooling.pooling.protocol import (
22
- IOProcessorRequest,
23
- IOProcessorResponse,
24
-)
21
from vllm.inputs.data import PromptType
22
from vllm.logger import init_logger
23
from vllm.outputs import PoolingRequestOutput
28
-from vllm.plugins.io_processors.interface import (
29
- IOProcessor,
30
- IOProcessorInput,
31
- IOProcessorOutput,
32
-)
24
+from vllm.plugins.io_processors.interface import IOProcessor
25
26
from .types import DataModuleConfig, ImagePrompt, ImageRequestOutput
27
@@ -227,7 +219,7 @@ def load_image(
219
return imgs, temporal_coords, location_coords, metas
220
221
230
-class PrithviMultimodalDataProcessor(IOProcessor):
222
+class PrithviMultimodalDataProcessor(IOProcessor[ImagePrompt, ImageRequestOutput]):
223
indices = [0, 1, 2, 3, 4, 5]
224
225
def __init__(self, vllm_config: VllmConfig):
@@ -251,34 +243,15 @@ class PrithviMultimodalDataProcessor(IOProcessor):
243
self.requests_cache: dict[str, dict[str, Any]] = {}
244
self.indices = DEFAULT_INPUT_INDICES
245
254
- def parse_request(self, request: Any) -> IOProcessorInput:
255
- if type(request) is dict:
256
- image_prompt = ImagePrompt(**request)
257
- return image_prompt
258
- if isinstance(request, IOProcessorRequest):
259
- if not hasattr(request, "data"):
260
- raise ValueError("missing 'data' field in OpenAIBaseModel Request")
261
-
262
- request_data = request.data
263
-
264
- if type(request_data) is dict:
265
- return ImagePrompt(**request_data)
266
- else:
267
- raise ValueError("Unable to parse the request data")
268
-
269
- raise ValueError("Unable to parse request")
270
-
271
- def output_to_response(
272
- self, plugin_output: IOProcessorOutput
273
- ) -> IOProcessorResponse:
274
- return IOProcessorResponse(
275
- request_id=plugin_output.request_id,
276
- data=plugin_output,
277
- )
246
+ def parse_data(self, data: object) -> ImagePrompt:
247
+ if isinstance(data, dict):
248
+ return ImagePrompt(**data)
249
+
250
+ raise ValueError("Prompt data should be an `ImagePrompt`")
251
252
def pre_process(
253
self,
281
- prompt: IOProcessorInput,
254
+ prompt: ImagePrompt,
255
request_id: str | None = None,
256
**kwargs,
257
) -> PromptType | Sequence[PromptType]:
@@ -364,7 +337,7 @@ class PrithviMultimodalDataProcessor(IOProcessor):
337
model_output: Sequence[PoolingRequestOutput],
338
request_id: str | None = None,
339
**kwargs,
367
- ) -> IOProcessorOutput:
340
+ ) -> ImageRequestOutput:
341
pred_imgs_list = []
342
343
if request_id and (request_id in self.requests_cache):
@@ -409,5 +382,7 @@ class PrithviMultimodalDataProcessor(IOProcessor):
382
)
383
384
return ImageRequestOutput(
412
- type=out_format, format="tiff", data=out_data, request_id=request_id
385
+ type=out_format,
386
+ format="tiff",
387
+ data=out_data,
388
)
tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/types.py
-4
@@ -38,9 +38,6 @@ class ImagePrompt(BaseModel):
38
"""
39
40
41
-MultiModalPromptType = ImagePrompt
42
-
43
-
41
class ImageRequestOutput(BaseModel):
42
"""
43
The output data of an image request to vLLM.
@@ -54,4 +51,3 @@ class ImageRequestOutput(BaseModel):
51
type: Literal["path", "b64_json"]
52
format: str
53
data: str
57
- request_id: str | None = None
tests/plugins_tests/test_io_processor_plugins.py
+2
-6
@@ -75,9 +75,7 @@ async def test_prithvi_mae_plugin_online(
75
# verify the output is formatted as expected for this plugin
76
plugin_data = parsed_response.data
77
78
- assert all(
79
- plugin_data.get(attr) for attr in ["type", "format", "data", "request_id"]
80
- )
78
+ assert all(plugin_data.get(attr) for attr in ["type", "format", "data"])
79
80
# We just check that the output is a valid base64 string.
81
# Raises an exception and fails the test if the string is corrupted.
@@ -110,9 +108,7 @@ def test_prithvi_mae_plugin_offline(vllm_runner, model_name: str):
108
output = pooler_output[0].outputs
109
110
# verify the output is formatted as expected for this plugin
113
- assert all(
114
- hasattr(output, attr) for attr in ["type", "format", "data", "request_id"]
115
- )
111
+ assert all(hasattr(output, attr) for attr in ["type", "format", "data"])
112
113
# We just check that the output is a valid base64 string.
114
# Raises an exception and fails the test if the string is corrupted.
vllm/entrypoints/llm.py
+35
-39
@@ -85,7 +85,6 @@ from vllm.tasks import PoolingTask
85
from vllm.tokenizers import TokenizerLike
86
from vllm.tokenizers.mistral import MistralTokenizer
87
from vllm.usage.usage_lib import UsageContext
88
-from vllm.utils.collection_utils import as_iter, is_list_of
88
from vllm.utils.counter import Counter
89
from vllm.v1.engine.llm_engine import LLMEngine
90
from vllm.v1.sample.logits_processor import LogitsProcessor
@@ -95,6 +94,7 @@ if TYPE_CHECKING:
94
95
logger = init_logger(__name__)
96
97
+_P = TypeVar("_P", bound=SamplingParams | PoolingParams | None)
98
_R = TypeVar("_R", default=Any)
99
100
@@ -1056,9 +1056,7 @@ class LLM:
1056
dict(truncate_prompt_tokens=truncate_prompt_tokens),
1057
)
1058
1059
- io_processor_prompt = False
1060
- if isinstance(prompts, dict) and "data" in prompts:
1061
- io_processor_prompt = True
1059
+ if use_io_processor := (isinstance(prompts, dict) and "data" in prompts):
1060
if self.io_processor is None:
1061
raise ValueError(
1062
"No IOProcessor plugin installed. Please refer "
@@ -1068,40 +1066,42 @@ class LLM:
1066
)
1067
1068
# Validate the request data is valid for the loaded plugin
1071
- validated_prompt = self.io_processor.parse_request(prompts)
1069
+ validated_prompt = self.io_processor.parse_data(prompts)
1070
1071
# obtain the actual model prompts from the pre-processor
1072
prompts = self.io_processor.pre_process(prompt=validated_prompt)
1073
+ prompts_seq = prompt_to_seq(prompts)
1074
1076
- if io_processor_prompt:
1077
- assert self.io_processor is not None
1078
- if is_list_of(pooling_params, PoolingParams):
1079
- validated_pooling_params: list[PoolingParams] = []
1080
- for param in as_iter(pooling_params):
1081
- validated_pooling_params.append(
1082
- self.io_processor.validate_or_generate_params(param)
1083
- )
1084
- pooling_params = validated_pooling_params
1085
- else:
1086
- assert not isinstance(pooling_params, Sequence)
1087
- pooling_params = self.io_processor.validate_or_generate_params(
1088
- pooling_params
1075
+ params_seq: Sequence[PoolingParams] = [
1076
+ self.io_processor.merge_pooling_params(param)
1077
+ for param in self._params_to_seq(
1078
+ pooling_params,
1079
+ len(prompts_seq),
1080
)
1090
-
1091
- if pooling_params is None:
1092
- # Use default pooling params.
1093
- pooling_params = PoolingParams()
1094
-
1095
- for param in as_iter(pooling_params):
1096
- if param.task is None:
1097
- param.task = pooling_task
1098
- elif param.task != pooling_task:
1099
- msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
1100
- raise ValueError(msg)
1081
+ ]
1082
+ for p in params_seq:
1083
+ if p.task is None:
1084
+ p.task = "plugin"
1085
+ else:
1086
+ if pooling_params is None:
1087
+ # Use default pooling params.
1088
+ pooling_params = PoolingParams()
1089
+
1090
+ prompts_seq = prompt_to_seq(prompts)
1091
+ params_seq = self._params_to_seq(pooling_params, len(prompts_seq))
1092
+
1093
+ for param in params_seq:
1094
+ if param.task is None:
1095
+ param.task = pooling_task
1096
+ elif param.task != pooling_task:
1097
+ msg = (
1098
+ f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
1099
+ )
1100
+ raise ValueError(msg)
1101
1102
outputs = self._run_completion(
1103
- prompts=prompts,
1104
- params=pooling_params,
1103
+ prompts=prompts_seq,
1104
+ params=params_seq,
1105
use_tqdm=use_tqdm,
1106
lora_request=lora_request,
1107
tokenization_kwargs=tokenization_kwargs,
@@ -1111,12 +1111,10 @@ class LLM:
1111
outputs, PoolingRequestOutput
1112
)
1113
1114
- if io_processor_prompt:
1114
+ if use_io_processor:
1115
# get the post-processed model outputs
1116
assert self.io_processor is not None
1117
- processed_outputs = self.io_processor.post_process(
1118
- model_output=model_outputs
1119
- )
1117
+ processed_outputs = self.io_processor.post_process(model_outputs)
1118
1119
return [
1120
PoolingRequestOutput[Any](
@@ -1662,11 +1660,9 @@ class LLM:
1660
1661
def _params_to_seq(
1662
self,
1665
- params: SamplingParams
1666
- | PoolingParams
1667
- | Sequence[SamplingParams | PoolingParams],
1663
+ params: _P | Sequence[_P],
1664
num_requests: int,
1669
- ) -> Sequence[SamplingParams | PoolingParams]:
1665
+ ) -> Sequence[_P]:
1666
if isinstance(params, Sequence):
1667
if len(params) != num_requests:
1668
raise ValueError(
vllm/entrypoints/pooling/pooling/protocol.py
-3
@@ -100,9 +100,6 @@ class IOProcessorRequest(PoolingBasicRequestMixin, EncodingRequestMixin, Generic
100
data: T
101
task: PoolingTask = "plugin"
102
103
- def to_pooling_params(self):
104
- return PoolingParams(task=self.task)
105
-
103
104
class IOProcessorResponse(OpenAIBaseModel, Generic[T]):
105
request_id: str | None = None
vllm/entrypoints/pooling/pooling/serving.py
+34
-18
@@ -85,7 +85,6 @@ class OpenAIServingPooling(OpenAIServing):
85
request_id = f"pool-{self._base_request_id(raw_request)}"
86
created_time = int(time.time())
87
88
- is_io_processor_request = isinstance(request, IOProcessorRequest)
88
try:
89
lora_request = self._maybe_get_adapters(request)
90
@@ -95,7 +94,7 @@ class OpenAIServingPooling(OpenAIServing):
94
)
95
96
engine_prompts: Sequence[PromptType | TokPrompt]
98
- if is_io_processor_request:
97
+ if use_io_processor := isinstance(request, IOProcessorRequest):
98
if self.io_processor is None:
99
raise ValueError(
100
"No IOProcessor plugin installed. Please refer "
@@ -104,7 +103,7 @@ class OpenAIServingPooling(OpenAIServing):
103
"offline inference example for more details."
104
)
105
107
- validated_prompt = self.io_processor.parse_request(request)
106
+ validated_prompt = self.io_processor.parse_data(request.data)
107
108
raw_prompts = await self.io_processor.pre_process_async(
109
prompt=validated_prompt, request_id=request_id
@@ -141,13 +140,18 @@ class OpenAIServingPooling(OpenAIServing):
140
# Schedule the request and get the result generator.
141
generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
142
try:
144
- if is_io_processor_request:
145
- assert self.io_processor is not None and isinstance(
146
- request, IOProcessorRequest
147
- )
148
- pooling_params = self.io_processor.validate_or_generate_params()
143
+ if use_io_processor:
144
+ assert self.io_processor is not None
145
+
146
+ pooling_params = self.io_processor.merge_pooling_params()
147
+ if pooling_params.task is None:
148
+ pooling_params.task = "plugin"
149
+
150
+ tokenization_kwargs: dict[str, Any] = {}
151
else:
150
- pooling_params = request.to_pooling_params()
152
+ pooling_params = request.to_pooling_params() # type: ignore
153
+ tok_params = request.build_tok_params(self.model_config) # type: ignore
154
+ tokenization_kwargs = tok_params.get_encode_kwargs()
155
156
for i, engine_prompt in enumerate(engine_prompts):
157
request_id_item = f"{request_id}-{i}"
@@ -165,12 +169,6 @@ class OpenAIServingPooling(OpenAIServing):
169
else await self._get_trace_headers(raw_request.headers)
170
)
171
168
- if is_io_processor_request:
169
- tokenization_kwargs: dict[str, Any] = {}
170
- else:
171
- tok_params = request.build_tok_params(self.model_config) # type: ignore
172
- tokenization_kwargs = tok_params.get_encode_kwargs()
173
-
172
generator = self.engine_client.encode(
173
engine_prompt,
174
pooling_params,
@@ -187,13 +185,31 @@ class OpenAIServingPooling(OpenAIServing):
185
186
result_generator = merge_async_iterators(*generators)
187
190
- if is_io_processor_request:
188
+ if use_io_processor:
189
assert self.io_processor is not None
190
output = await self.io_processor.post_process_async(
193
- model_output=result_generator,
191
+ result_generator,
192
request_id=request_id,
193
)
196
- return self.io_processor.output_to_response(output)
194
+
195
+ if callable(
196
+ output_to_response := getattr(
197
+ self.io_processor, "output_to_response", None
198
+ )
199
+ ):
200
+ logger.warning_once(
201
+ "`IOProcessor.output_to_response` is deprecated. To ensure "
202
+ "consistency between offline and online APIs, "
203
+ "`IOProcessorResponse` will become a transparent wrapper "
204
+ "around output data from v0.19 onwards.",
205
+ )
206
+
207
+ if hasattr(output, "request_id") and output.request_id is None:
208
+ output.request_id = request_id # type: ignore
209
+
210
+ return output_to_response(output) # type: ignore
211
+
212
+ return IOProcessorResponse(request_id=request_id, data=output)
213
214
assert isinstance(request, (PoolingCompletionRequest, PoolingChatRequest))
215
num_prompts = len(engine_prompts)
vllm/plugins/io_processors/interface.py
+63
-19
@@ -1,12 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
+import warnings
4
from abc import ABC, abstractmethod
5
from collections.abc import AsyncGenerator, Sequence
6
-from typing import Any, Generic, TypeVar
6
+from typing import Generic, TypeVar
7
8
from vllm.config import VllmConfig
9
-from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse
9
from vllm.inputs.data import PromptType
10
from vllm.outputs import PoolingRequestOutput
11
from vllm.pooling_params import PoolingParams
@@ -18,8 +17,68 @@ IOProcessorOutput = TypeVar("IOProcessorOutput")
17
18
class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):
19
def __init__(self, vllm_config: VllmConfig):
20
+ super().__init__()
21
+
22
self.vllm_config = vllm_config
23
24
+ def parse_data(self, data: object) -> IOProcessorInput:
25
+ if callable(parse_request := getattr(self, "parse_request", None)):
26
+ warnings.warn(
27
+ "`parse_request` has been renamed to `parse_data`. "
28
+ "Please update your IO Processor Plugin to use the new name. "
29
+ "The old name will be removed in v0.19.",
30
+ DeprecationWarning,
31
+ stacklevel=2,
32
+ )
33
+
34
+ return parse_request(data) # type: ignore
35
+
36
+ raise NotImplementedError
37
+
38
+ def merge_sampling_params(
39
+ self,
40
+ params: SamplingParams | None = None,
41
+ ) -> SamplingParams:
42
+ if callable(
43
+ validate_or_generate_params := getattr(
44
+ self, "validate_or_generate_params", None
45
+ )
46
+ ):
47
+ warnings.warn(
48
+ "`validate_or_generate_params` has been split into "
49
+ "`merge_sampling_params` and `merge_pooling_params`."
50
+ "Please update your IO Processor Plugin to use the new methods. "
51
+ "The old name will be removed in v0.19.",
52
+ DeprecationWarning,
53
+ stacklevel=2,
54
+ )
55
+
56
+ return validate_or_generate_params(params) # type: ignore
57
+
58
+ return params or SamplingParams()
59
+
60
+ def merge_pooling_params(
61
+ self,
62
+ params: PoolingParams | None = None,
63
+ ) -> PoolingParams:
64
+ if callable(
65
+ validate_or_generate_params := getattr(
66
+ self, "validate_or_generate_params", None
67
+ )
68
+ ):
69
+ warnings.warn(
70
+ "`validate_or_generate_params` has been split into "
71
+ "`merge_sampling_params` and `merge_pooling_params`."
72
+ "Please update your IO Processor Plugin to use the new methods. "
73
+ "The old name will be removed in v0.19.",
74
+ DeprecationWarning,
75
+ stacklevel=2,
76
+ )
77
+
78
+ return validate_or_generate_params(params) # type: ignore
79
+
80
+ return params or PoolingParams(task="plugin")
81
+
82
@abstractmethod
83
def pre_process(
84
self,
@@ -59,19 +118,4 @@ class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):
118
[(i, item) async for i, item in model_output], key=lambda output: output[0]
119
)
120
collected_output = [output[1] for output in sorted_output]
62
- return self.post_process(collected_output, request_id, **kwargs)
63
-
64
- @abstractmethod
65
- def parse_request(self, request: Any) -> IOProcessorInput:
66
- raise NotImplementedError
67
-
68
- def validate_or_generate_params(
69
- self, params: SamplingParams | PoolingParams | None = None
70
- ) -> SamplingParams | PoolingParams:
71
- return params or PoolingParams()
72
-
73
- @abstractmethod
74
- def output_to_response(
75
- self, plugin_output: IOProcessorOutput
76
- ) -> IOProcessorResponse:
77
- raise NotImplementedError
121
+ return self.post_process(collected_output, request_id=request_id, **kwargs)
vllm/utils/collection_utils.py
-6
@@ -51,12 +51,6 @@ def as_list(maybe_list: Iterable[T]) -> list[T]:
51
return maybe_list if isinstance(maybe_list, list) else list(maybe_list)
52
53
54
-def as_iter(obj: T | Iterable[T]) -> Iterable[T]:
55
- if isinstance(obj, str) or not isinstance(obj, Iterable):
56
- return [obj] # type: ignore[list-item]
57
- return obj
58
-
59
-
54
def is_list_of(
55
value: object,
56
typ: type[T] | tuple[type[T], ...],