@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1cb39dbcdd
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 14, 05:30 UTC

[Misc] IO Processor plugins for pooling models (#22820)

Signed-off-by: Christian Pinto <christian.pinto@ibm.com> Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Co-authored-by: Max de Bayser <mbayser@br.ibm.com>

Christian Pinto committed Sep 1, 2025 at 07:07 UTC 1cb39dbcddd9bf98d459cb4da81e3a9cc942088e
25 files changed +1183 -43
.buildkite/test-pipeline.yaml
+5
@@ -769,6 +769,11 @@ steps:
769 - pytest -v -s plugins_tests/test_platform_plugins.py
770 - pip uninstall vllm_add_dummy_platform -y
771 # end platform plugin tests
772 + # begin io_processor plugins test, all the code in between uses the prithvi_io_processor plugin
773 + - pip install -e ./plugins/prithvi_io_processor_plugin
774 + - pytest -v -s plugins_tests/test_io_processor_plugins.py
775 + - pip uninstall prithvi_io_processor_plugin -y
776 + # end io_processor plugins test
777 # other tests continue here:
778 - pytest -v -s plugins_tests/test_scheduler_plugins.py
779 - pip install -e ./plugins/vllm_add_dummy_model
docs/design/io_processor_plugins.md new
+78
@@ -0,0 +1,78 @@
1 +# IO Processor Plugins
2 +
3 +IO Processor plugins are a feature that allows pre and post processing of the model input and output for pooling models. The idea is that users are allowed to pass a custom input to vLLM that is converted into one or more model prompts and fed to the model `encode` method. One potential use-case of such plugins is that of using vLLM for generating multi-modal data. Say users feed an image to vLLM and get an image in output.
4 +
5 +When performing an inference with IO Processor plugins, the prompt type is defined by the plugin and the same is valid for the final request output. vLLM does not perform any validation of input/output data, and it is up to the plugin to ensure the correct data is being fed to the model and returned to the user. As of now these plugins support only pooling models and can be triggerd via the `encode` method in `LLM` and `AsyncLLM`, or in online serving mode via the `/pooling` endpoint.
6 +
7 +## Writing an IO Processor Plugin
8 +
9 +IO Processor plugins implement the `IOProcessor` interface (<gh-file:vllm/plugins/io_processors/interface.py>):
10 +
11 +```python
12 +IOProcessorInput = TypeVar('IOProcessorInput')
13 +IOProcessorOutput = TypeVar('IOProcessorOutput')
14 +
15 +class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):
16 +
17 + def __init__(self, vllm_config: VllmConfig):
18 + self.vllm_config = vllm_config
19 +
20 + @abstractmethod
21 + def pre_process(
22 + self,
23 + prompt: IOProcessorInput,
24 + request_id: Optional[str] = None,
25 + **kwargs,
26 + ) -> Union[PromptType, Sequence[PromptType]]:
27 + raise NotImplementedError
28 +
29 + async def pre_process_async(
30 + self,
31 + prompt: IOProcessorInput,
32 + request_id: Optional[str] = None,
33 + **kwargs,
34 + ) -> Union[PromptType, Sequence[PromptType]]:
35 + return self.pre_process(prompt, request_id, **kwargs)
36 +
37 + @abstractmethod
38 + def post_process(self,
39 + model_output: Sequence[PoolingRequestOutput],
40 + request_id: Optional[str] = None,
41 + **kwargs) -> IOProcessorOutput:
42 + raise NotImplementedError
43 +
44 + async def post_process_async(
45 + self,
46 + model_output: AsyncGenerator[tuple[int, PoolingRequestOutput]],
47 + request_id: Optional[str] = None,
48 + **kwargs,
49 + ) -> IOProcessorOutput:
50 + collected_output = [item async for i, item in model_output]
51 + return self.post_process(collected_output, request_id, **kwargs)
52 +
53 + @abstractmethod
54 + def parse_request(self, request: Any) -> IOProcessorInput:
55 + raise NotImplementedError
56 +
57 + @abstractmethod
58 + def output_to_response(
59 + self, plugin_output: IOProcessorOutput) -> IOProcessorResponse:
60 + raise NotImplementedError
61 +```
62 +
63 +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.
64 +The `pre_process*` methods take the validated plugin input to generate vLLM's model prompts for regular inference.
65 +The `post_process*` methods take `PoolingRequestOutput` objects as input and generate a custom plugin output.
66 +
67 +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 `/io_processor_pooling` serving endpoint is [here](../../vllm/entrypoints/openai/serving_pooling_with_io_plugin.py).
68 +
69 +An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/christian-pinto/prithvi_io_processor_plugin). Please, also refer to our [online](../../examples/online_serving/prithvi_geospatial_mae.py) and [offline](../../examples/offline_inference/prithvi_geospatial_mae_io_processor.py) inference examples.
70 +
71 +## Using an IO Processor plugin
72 +
73 +IO Processor plugins are loaded at engine startup and there are two methods for specifying the name of the plugin to be loaded:
74 +
75 +1. Via vLLM's `EngineArgs`: setting the `io_processor_plugin` argument in the `EngineArgs` used to initialize the `AsyncLLM`. The same can be achieved by passing the `io_processor_plugin` argument to `LLM` in offline mode, or by passing the `--io-processor-plugin` argument in serving mode.
76 +2. Via the model HF configuration: adding an `io_processor_plugin` field to the model config (config.json).
77 +
78 +The order also determines method priority. i.e., setting the plugin name via `EngineArgs` will override any plugin name specified in the model HF config (config.json).
docs/design/plugin_system.md
+2
@@ -49,6 +49,8 @@ Every plugin has three parts:
49
50 - **Platform plugins** (with group name `vllm.platform_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree platforms into vLLM. The plugin function should return `None` when the platform is not supported in the current environment, or the platform class's fully qualified name when the platform is supported.
51
52 +- **IO Processor plugins** (with group name `vllm.io_processor_plugins`): The primary use case for these plugins is to register custom pre/post processing of the model prompt and model output for poling models. The plugin function returns the IOProcessor's class fully qualified name.
53 +
54 ## Guidelines for Writing Plugins
55
56 - **Being re-entrant**: The function specified in the entry point should be re-entrant, meaning it can be called multiple times without causing issues. This is necessary because the function might be called multiple times in some processes.
examples/offline_inference/prithvi_geospatial_mae_io_processor.py new
+60
@@ -0,0 +1,60 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import base64
4 +import os
5 +
6 +import torch
7 +
8 +from vllm import LLM
9 +from vllm.pooling_params import PoolingParams
10 +
11 +# This example shows how to perform an offline inference that generates
12 +# multimodal data. In this specific case this example will take a geotiff
13 +# image as input, process it using the multimodal data processor, and
14 +# perform inference.
15 +# Reuirement - install plugin at:
16 +# https://github.com/christian-pinto/prithvi_io_processor_plugin
17 +
18 +
19 +def main():
20 + torch.set_default_dtype(torch.float16)
21 + image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/India_900498_S2Hand.tif" # noqa: E501
22 +
23 + img_prompt = dict(
24 + data=image_url,
25 + data_format="url",
26 + image_format="tiff",
27 + out_data_format="b64_json",
28 + )
29 +
30 + llm = LLM(
31 + model="christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
32 + skip_tokenizer_init=True,
33 + trust_remote_code=True,
34 + enforce_eager=True,
35 + # Limit the maximum number of parallel requests
36 + # to avoid the model going OOM.
37 + # The maximum number depends on the available GPU memory
38 + max_num_seqs=32,
39 + io_processor_plugin="prithvi_to_tiff_india",
40 + )
41 +
42 + pooling_params = PoolingParams(task="encode", softmax=False)
43 + pooler_output = llm.encode(
44 + img_prompt,
45 + pooling_params=pooling_params,
46 + )
47 + output = pooler_output[0].outputs
48 +
49 + print(output)
50 + decoded_data = base64.b64decode(output.data)
51 +
52 + file_path = os.path.join(os.getcwd(), "offline_prediction.tiff")
53 + with open(file_path, "wb") as f:
54 + f.write(decoded_data)
55 +
56 + print(f"Output file path: {file_path}")
57 +
58 +
59 +if __name__ == "__main__":
60 + main()
examples/online_serving/prithvi_geospatial_mae.py new
+54
@@ -0,0 +1,54 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import base64
5 +import os
6 +
7 +import requests
8 +
9 +# This example shows how to perform an online inference that generates
10 +# multimodal data. In this specific case this example will take a geotiff
11 +# image as input, process it using the multimodal data processor, and
12 +# perform inference.
13 +# Reuirements :
14 +# - install plugin at:
15 +# https://github.com/christian-pinto/prithvi_io_processor_plugin
16 +# - start vllm in serving mode with the below args
17 +# --model='christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM'
18 +# --task embed --trust-remote-code
19 +# --skip-tokenizer-init --enforce-eager
20 +# --io-processor-plugin prithvi_to_tiff_india
21 +
22 +
23 +def main():
24 + image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/India_900498_S2Hand.tif" # noqa: E501
25 + server_endpoint = "http://localhost:8000/pooling"
26 +
27 + request_payload_url = {
28 + "data": {
29 + "data": image_url,
30 + "data_format": "url",
31 + "image_format": "tiff",
32 + "out_data_format": "b64_json",
33 + },
34 + "priority": 0,
35 + "model": "christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
36 + }
37 +
38 + ret = requests.post(server_endpoint, json=request_payload_url)
39 +
40 + print(f"response.status_code: {ret.status_code}")
41 + print(f"response.reason:{ret.reason}")
42 +
43 + response = ret.json()
44 +
45 + decoded_image = base64.b64decode(response["data"]["data"])
46 +
47 + out_path = os.path.join(os.getcwd(), "online_prediction.tiff")
48 +
49 + with open(out_path, "wb") as f:
50 + f.write(decoded_image)
51 +
52 +
53 +if __name__ == "__main__":
54 + main()
tests/conftest.py
+3
@@ -1120,6 +1120,9 @@ class VllmRunner:
1120
1121 return self.llm.llm_engine.collective_rpc(_apply_model)
1122
1123 + def get_llm(self) -> LLM:
1124 + return self.llm
1125 +
1126 def __enter__(self):
1127 return self
1128
tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/__init__.py new
+8
@@ -0,0 +1,8 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +def register_prithvi_india():
4 + return "prithvi_io_processor.prithvi_processor.PrithviMultimodalDataProcessorIndia" # noqa: E501
5 +
6 +
7 +def register_prithvi_valencia():
8 + return "prithvi_io_processor.prithvi_processor.PrithviMultimodalDataProcessorValencia" # noqa: E501
tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py new
+449
@@ -0,0 +1,449 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from __future__ import annotations
5 +
6 +import base64
7 +import datetime
8 +import os
9 +import tempfile
10 +import urllib.request
11 +from collections.abc import AsyncGenerator, Sequence
12 +from typing import Any, Optional, Union
13 +
14 +import albumentations
15 +import numpy as np
16 +import rasterio
17 +import regex as re
18 +import torch
19 +from einops import rearrange
20 +from terratorch.datamodules import Sen1Floods11NonGeoDataModule
21 +
22 +from vllm.config import VllmConfig
23 +from vllm.entrypoints.openai.protocol import (IOProcessorRequest,
24 + IOProcessorResponse)
25 +from vllm.inputs.data import PromptType
26 +from vllm.logger import init_logger
27 +from vllm.outputs import PoolingRequestOutput
28 +from vllm.plugins.io_processors.interface import (IOProcessor,
29 + IOProcessorInput,
30 + IOProcessorOutput)
31 +
32 +from .types import DataModuleConfig, ImagePrompt, ImageRequestOutput
33 +
34 +logger = init_logger(__name__)
35 +
36 +NO_DATA = -9999
37 +NO_DATA_FLOAT = 0.0001
38 +OFFSET = 0
39 +PERCENTILE = 99
40 +
41 +DEFAULT_INPUT_INDICES = [0, 1, 2, 3, 4, 5]
42 +
43 +datamodule_config: DataModuleConfig = {
44 + "bands": ["BLUE", "GREEN", "RED", "NIR_NARROW", "SWIR_1", "SWIR_2"],
45 + "batch_size":
46 + 16,
47 + "constant_scale":
48 + 0.0001,
49 + "data_root":
50 + "/dccstor/geofm-finetuning/datasets/sen1floods11",
51 + "drop_last":
52 + True,
53 + "no_data_replace":
54 + 0.0,
55 + "no_label_replace":
56 + -1,
57 + "num_workers":
58 + 8,
59 + "test_transform": [
60 + albumentations.Resize(always_apply=False,
61 + height=448,
62 + interpolation=1,
63 + p=1,
64 + width=448),
65 + albumentations.pytorch.ToTensorV2(transpose_mask=False,
66 + always_apply=True,
67 + p=1.0),
68 + ],
69 +}
70 +
71 +
72 +def save_geotiff(image: torch.Tensor, meta: dict,
73 + out_format: str) -> str | bytes:
74 + """Save multi-band image in Geotiff file.
75 +
76 + Args:
77 + image: np.ndarray with shape (bands, height, width)
78 + output_path: path where to save the image
79 + meta: dict with meta info.
80 + """
81 + if out_format == "path":
82 + # create temp file
83 + file_path = os.path.join(os.getcwd(), "prediction.tiff")
84 + with rasterio.open(file_path, "w", **meta) as dest:
85 + for i in range(image.shape[0]):
86 + dest.write(image[i, :, :], i + 1)
87 +
88 + return file_path
89 + elif out_format == "b64_json":
90 + with tempfile.NamedTemporaryFile() as tmpfile:
91 + with rasterio.open(tmpfile.name, "w", **meta) as dest:
92 + for i in range(image.shape[0]):
93 + dest.write(image[i, :, :], i + 1)
94 +
95 + file_data = tmpfile.read()
96 + return base64.b64encode(file_data)
97 +
98 + else:
99 + raise ValueError("Unknown output format")
100 +
101 +
102 +def _convert_np_uint8(float_image: torch.Tensor):
103 + image = float_image.numpy() * 255.0
104 + image = image.astype(dtype=np.uint8)
105 +
106 + return image
107 +
108 +
109 +def read_geotiff(
110 + file_path: Optional[str] = None,
111 + path_type: Optional[str] = None,
112 + file_data: Optional[bytes] = None,
113 +) -> tuple[torch.Tensor, dict, tuple[float, float] | None]:
114 + """Read all bands from *file_path* and return image + meta info.
115 +
116 + Args:
117 + file_path: path to image file.
118 +
119 + Returns:
120 + np.ndarray with shape (bands, height, width)
121 + meta info dict
122 + """
123 +
124 + if all([x is None for x in [file_path, path_type, file_data]]):
125 + raise Exception("All input fields to read_geotiff are None")
126 + write_to_file: Optional[bytes] = None
127 + path: Optional[str] = None
128 + if file_data is not None:
129 + # with tempfile.NamedTemporaryFile() as tmpfile:
130 + # tmpfile.write(file_data)
131 + # path = tmpfile.name
132 +
133 + write_to_file = file_data
134 + elif file_path is not None and path_type == "url":
135 + resp = urllib.request.urlopen(file_path)
136 + # with tempfile.NamedTemporaryFile() as tmpfile:
137 + # tmpfile.write(resp.read())
138 + # path = tmpfile.name
139 + write_to_file = resp.read()
140 + elif file_path is not None and path_type == "path":
141 + path = file_path
142 + elif file_path is not None and path_type == "b64_json":
143 + image_data = base64.b64decode(file_path)
144 + # with tempfile.NamedTemporaryFile() as tmpfile:
145 + # tmpfile.write(image_data)
146 + # path = tmpfile.name
147 + write_to_file = image_data
148 + else:
149 + raise Exception("Wrong combination of parameters to read_geotiff")
150 +
151 + with tempfile.NamedTemporaryFile() as tmpfile:
152 + path_to_use = None
153 + if write_to_file:
154 + tmpfile.write(write_to_file)
155 + path_to_use = tmpfile.name
156 + elif path:
157 + path_to_use = path
158 +
159 + with rasterio.open(path_to_use) as src:
160 + img = src.read()
161 + meta = src.meta
162 + try:
163 + coords = src.lnglat()
164 + except Exception:
165 + # Cannot read coords
166 + coords = None
167 +
168 + return img, meta, coords
169 +
170 +
171 +def load_image(
172 + data: Union[list[str]],
173 + path_type: str,
174 + mean: Optional[list[float]] = None,
175 + std: Optional[list[float]] = None,
176 + indices: Optional[Union[list[int], None]] = None,
177 +):
178 + """Build an input example by loading images in *file_paths*.
179 +
180 + Args:
181 + file_paths: list of file paths .
182 + mean: list containing mean values for each band in the
183 + images in *file_paths*.
184 + std: list containing std values for each band in the
185 + images in *file_paths*.
186 +
187 + Returns:
188 + np.array containing created example
189 + list of meta info for each image in *file_paths*
190 + """
191 +
192 + imgs = []
193 + metas = []
194 + temporal_coords = []
195 + location_coords = []
196 +
197 + for file in data:
198 + # if isinstance(file, bytes):
199 + # img, meta, coords = read_geotiff(file_data=file)
200 + # else:
201 + img, meta, coords = read_geotiff(file_path=file, path_type=path_type)
202 + # Rescaling (don't normalize on nodata)
203 + img = np.moveaxis(img, 0, -1) # channels last for rescaling
204 + if indices is not None:
205 + img = img[..., indices]
206 + if mean is not None and std is not None:
207 + img = np.where(img == NO_DATA, NO_DATA_FLOAT, (img - mean) / std)
208 +
209 + imgs.append(img)
210 + metas.append(meta)
211 + if coords is not None:
212 + location_coords.append(coords)
213 +
214 + try:
215 + match = re.search(r"(\d{7,8}T\d{6})", file)
216 + if match:
217 + year = int(match.group(1)[:4])
218 + julian_day = match.group(1).split("T")[0][4:]
219 + if len(julian_day) == 3:
220 + julian_day = int(julian_day)
221 + else:
222 + julian_day = (datetime.datetime.strptime(
223 + julian_day, "%m%d").timetuple().tm_yday)
224 + temporal_coords.append([year, julian_day])
225 + except Exception:
226 + logger.exception("Could not extract timestamp for %s", file)
227 +
228 + imgs = np.stack(imgs, axis=0) # num_frames, H, W, C
229 + imgs = np.moveaxis(imgs, -1, 0).astype("float32") # C, num_frames, H, W
230 + imgs = np.expand_dims(imgs, axis=0) # add batch di
231 +
232 + return imgs, temporal_coords, location_coords, metas
233 +
234 +
235 +class PrithviMultimodalDataProcessor(IOProcessor):
236 +
237 + def __init__(self, vllm_config: VllmConfig):
238 +
239 + super().__init__(vllm_config)
240 +
241 + self.datamodule = Sen1Floods11NonGeoDataModule(
242 + data_root=datamodule_config["data_root"],
243 + batch_size=datamodule_config["batch_size"],
244 + num_workers=datamodule_config["num_workers"],
245 + bands=datamodule_config["bands"],
246 + drop_last=datamodule_config["drop_last"],
247 + test_transform=datamodule_config["test_transform"],
248 + )
249 + self.img_size = 512
250 + self.h1 = 1
251 + self.w1 = 1
252 + self.original_h = 512
253 + self.original_w = 512
254 + self.batch_size = 1
255 + self.meta_data = None
256 + self.requests_cache: dict[str, dict[str, Any]] = {}
257 + self.indices = DEFAULT_INPUT_INDICES
258 +
259 + def parse_request(self, request: Any) -> IOProcessorInput:
260 + if type(request) is dict:
261 + image_prompt = ImagePrompt(**request)
262 + return image_prompt
263 + if isinstance(request, IOProcessorRequest):
264 + if not hasattr(request, "data"):
265 + raise ValueError(
266 + "missing 'data' field in OpenAIBaseModel Request")
267 +
268 + request_data = request.data
269 +
270 + if type(request_data) is dict:
271 + return ImagePrompt(**request_data)
272 + else:
273 + raise ValueError("Unable to parse the request data")
274 +
275 + raise ValueError("Unable to parse request")
276 +
277 + def output_to_response(
278 + self, plugin_output: IOProcessorOutput) -> IOProcessorResponse:
279 + return IOProcessorResponse(
280 + request_id=plugin_output.request_id,
281 + data=plugin_output,
282 + )
283 +
284 + def pre_process(
285 + self,
286 + prompt: IOProcessorInput,
287 + request_id: Optional[str] = None,
288 + **kwargs,
289 + ) -> Union[PromptType, Sequence[PromptType]]:
290 +
291 + image_data = dict(prompt)
292 +
293 + if request_id:
294 + self.requests_cache[request_id] = {
295 + "out_format": image_data["out_data_format"],
296 + }
297 +
298 + input_data, temporal_coords, location_coords, meta_data = load_image(
299 + data=[image_data["data"]],
300 + indices=self.indices,
301 + path_type=image_data["data_format"],
302 + )
303 +
304 + self.meta_data = meta_data[0]
305 +
306 + if input_data.mean() > 1:
307 + input_data = input_data / 10000 # Convert to range 0-1
308 +
309 + self.original_h, self.original_w = input_data.shape[-2:]
310 + pad_h = (self.img_size -
311 + (self.original_h % self.img_size)) % self.img_size
312 + pad_w = (self.img_size -
313 + (self.original_w % self.img_size)) % self.img_size
314 + input_data = np.pad(
315 + input_data,
316 + ((0, 0), (0, 0), (0, 0), (0, pad_h), (0, pad_w)),
317 + mode="reflect",
318 + )
319 +
320 + batch = torch.tensor(input_data)
321 + windows = batch.unfold(3, self.img_size,
322 + self.img_size).unfold(4, self.img_size,
323 + self.img_size)
324 + self.h1, self.w1 = windows.shape[3:5]
325 + windows = rearrange(
326 + windows,
327 + "b c t h1 w1 h w -> (b h1 w1) c t h w",
328 + h=self.img_size,
329 + w=self.img_size,
330 + )
331 +
332 + # Split into batches if number of windows > batch_size
333 + num_batches = (windows.shape[0] // self.batch_size
334 + if windows.shape[0] > self.batch_size else 1)
335 + windows = torch.tensor_split(windows, num_batches, dim=0)
336 +
337 + if temporal_coords:
338 + temporal_coords = torch.tensor(temporal_coords).unsqueeze(0)
339 + else:
340 + temporal_coords = None
341 + if location_coords:
342 + location_coords = torch.tensor(location_coords[0]).unsqueeze(0)
343 + else:
344 + location_coords = None
345 +
346 + prompts = []
347 + for window in windows:
348 + # Apply standardization
349 + window = self.datamodule.test_transform(
350 + image=window.squeeze().numpy().transpose(1, 2, 0))
351 + window = self.datamodule.aug(window)["image"]
352 + prompts.append({
353 + "prompt_token_ids": [1],
354 + "multi_modal_data": {
355 + "pixel_values": window.to(torch.float16)[0],
356 + "location_coords": location_coords.to(torch.float16),
357 + },
358 + })
359 +
360 + return prompts
361 +
362 + async def pre_process_async(
363 + self,
364 + prompt: IOProcessorInput,
365 + request_id: Optional[str] = None,
366 + **kwargs,
367 + ) -> Union[PromptType, Sequence[PromptType]]:
368 + return self.pre_process(prompt, request_id, **kwargs)
369 +
370 + def post_process(
371 + self,
372 + model_output: Sequence[PoolingRequestOutput],
373 + request_id: Optional[str] = None,
374 + **kwargs,
375 + ) -> IOProcessorOutput:
376 +
377 + pred_imgs_list = []
378 +
379 + if request_id and (request_id in self.requests_cache):
380 + out_format = self.requests_cache[request_id]["out_format"]
381 + else:
382 + out_format = "b64_json"
383 +
384 + for output in model_output:
385 + y_hat = output.outputs.data.argmax(dim=1)
386 + pred = torch.nn.functional.interpolate(
387 + y_hat.unsqueeze(1).float(),
388 + size=self.img_size,
389 + mode="nearest",
390 + )
391 + pred_imgs_list.append(pred)
392 +
393 + pred_imgs: torch.Tensor = torch.concat(pred_imgs_list, dim=0)
394 +
395 + # Build images from patches
396 + pred_imgs = rearrange(
397 + pred_imgs,
398 + "(b h1 w1) c h w -> b c (h1 h) (w1 w)",
399 + h=self.img_size,
400 + w=self.img_size,
401 + b=1,
402 + c=1,
403 + h1=self.h1,
404 + w1=self.w1,
405 + )
406 +
407 + # Cut padded area back to original size
408 + pred_imgs = pred_imgs[..., :self.original_h, :self.original_w]
409 +
410 + # Squeeze (batch size 1)
411 + pred_imgs = pred_imgs[0]
412 +
413 + if not self.meta_data:
414 + raise ValueError("No metadata available for the current task")
415 + self.meta_data.update(count=1, dtype="uint8", compress="lzw", nodata=0)
416 + out_data = save_geotiff(_convert_np_uint8(pred_imgs), self.meta_data,
417 + out_format)
418 +
419 + return ImageRequestOutput(type=out_format,
420 + format="tiff",
421 + data=out_data,
422 + request_id=request_id)
423 +
424 + async def post_process_async(
425 + self,
426 + model_output: AsyncGenerator[tuple[int, PoolingRequestOutput]],
427 + request_id: Optional[str] = None,
428 + **kwargs,
429 + ) -> IOProcessorOutput:
430 + collected_output = [item async for i, item in model_output]
431 + return self.post_process(collected_output, request_id, **kwargs)
432 +
433 +
434 +class PrithviMultimodalDataProcessorIndia(PrithviMultimodalDataProcessor):
435 +
436 + def __init__(self, vllm_config: VllmConfig):
437 +
438 + super().__init__(vllm_config)
439 +
440 + self.indices = [1, 2, 3, 8, 11, 12]
441 +
442 +
443 +class PrithviMultimodalDataProcessorValencia(PrithviMultimodalDataProcessor):
444 +
445 + def __init__(self, vllm_config: VllmConfig):
446 +
447 + super().__init__(vllm_config)
448 +
449 + self.indices = [0, 1, 2, 3, 4, 5]
tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/types.py new
+59
@@ -0,0 +1,59 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Any, Literal, Optional, TypedDict, Union
5 +
6 +import albumentations
7 +from pydantic import BaseModel
8 +
9 +
10 +class DataModuleConfig(TypedDict):
11 + bands: list[str]
12 + batch_size: int
13 + constant_scale: float
14 + data_root: str
15 + drop_last: bool
16 + no_data_replace: float
17 + no_label_replace: int
18 + num_workers: int
19 + test_transform: list[
20 + albumentations.core.transforms_interface.BasicTransform]
21 +
22 +
23 +class ImagePrompt(BaseModel):
24 +
25 + data_format: Literal["b64_json", "bytes", "url"]
26 + """
27 + This is the data type for the input image
28 + """
29 +
30 + image_format: str
31 + """
32 + This is the image format (e.g., jpeg, png, etc.)
33 + """
34 +
35 + out_data_format: Literal["b64_json", "url"]
36 +
37 + data: Any
38 + """
39 + Input image data
40 + """
41 +
42 +
43 +MultiModalPromptType = Union[ImagePrompt]
44 +
45 +
46 +class ImageRequestOutput(BaseModel):
47 + """
48 + The output data of an image request to vLLM.
49 +
50 + Args:
51 + type (str): The data content type [path, object]
52 + format (str): The image format (e.g., jpeg, png, etc.)
53 + data (Any): The resulting data.
54 + """
55 +
56 + type: Literal["path", "b64_json"]
57 + format: str
58 + data: str
59 + request_id: Optional[str] = None
tests/plugins/prithvi_io_processor_plugin/setup.py new
+16
@@ -0,0 +1,16 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from setuptools import setup
5 +
6 +setup(
7 + name="prithvi_io_processor_plugin",
8 + version="0.1",
9 + packages=["prithvi_io_processor"],
10 + entry_points={
11 + "vllm.io_processor_plugins": [
12 + "prithvi_to_tiff_india = prithvi_io_processor:register_prithvi_india", # noqa: E501
13 + "prithvi_to_tiff_valencia = prithvi_io_processor:register_prithvi_valencia", # noqa: E501
14 + ]
15 + },
16 +)
tests/plugins_tests/conftest.py deleted
-12
@@ -1,12 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import pytest
4 -
5 -
6 -@pytest.fixture(scope="function", autouse=True)
7 -def use_v0_only(monkeypatch):
8 - """
9 - Since this module is V0 only, set VLLM_USE_V1=0 for
10 - all tests in the module.
11 - """
12 - monkeypatch.setenv('VLLM_USE_V1', '0')
\ No newline at end of file
tests/plugins_tests/test_io_processor_plugins.py new
+137
@@ -0,0 +1,137 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import base64
4 +
5 +import pytest
6 +import requests
7 +
8 +from tests.utils import RemoteOpenAIServer
9 +from vllm.config import VllmConfig
10 +from vllm.entrypoints.llm import LLM
11 +from vllm.entrypoints.openai.protocol import IOProcessorResponse
12 +from vllm.plugins.io_processors import get_io_processor
13 +from vllm.pooling_params import PoolingParams
14 +
15 +MODEL_NAME = "christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM"
16 +
17 +image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff" # noqa: E501
18 +
19 +
20 +def test_loading_missing_plugin():
21 + vllm_config = VllmConfig()
22 + with pytest.raises(ValueError):
23 + get_io_processor(vllm_config, "wrong_plugin")
24 +
25 +
26 +def test_loading_engine_with_wrong_plugin():
27 +
28 + with pytest.raises(ValueError):
29 + LLM(
30 + model=MODEL_NAME,
31 + skip_tokenizer_init=True,
32 + trust_remote_code=True,
33 + enforce_eager=True,
34 + # Limit the maximum number of parallel requests
35 + # to avoid the model going OOM in CI.
36 + max_num_seqs=32,
37 + io_processor_plugin="wrong_plugin",
38 + )
39 +
40 +
41 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
42 +def test_prithvi_mae_plugin_offline(vllm_runner, model_name: str):
43 +
44 + img_prompt = dict(
45 + data=image_url,
46 + data_format="url",
47 + image_format="tiff",
48 + out_data_format="b64_json",
49 + )
50 +
51 + pooling_params = PoolingParams(task="encode", softmax=False)
52 +
53 + with vllm_runner(
54 + model_name,
55 + runner="pooling",
56 + skip_tokenizer_init=True,
57 + trust_remote_code=True,
58 + enforce_eager=True,
59 + # Limit the maximum number of parallel requests
60 + # to avoid the model going OOM in CI.
61 + max_num_seqs=1,
62 + io_processor_plugin="prithvi_to_tiff_valencia",
63 + ) as llm_runner:
64 + pooler_output = llm_runner.get_llm().encode(
65 + img_prompt,
66 + pooling_params=pooling_params,
67 + )
68 + output = pooler_output[0].outputs
69 +
70 + # verify the output is formatted as expected for this plugin
71 + assert all(
72 + hasattr(output, attr)
73 + for attr in ["type", "format", "data", "request_id"])
74 +
75 + # We just check that the output is a valid base64 string.
76 + # Raises an exception and fails the test if the string is corrupted.
77 + base64.b64decode(output.data)
78 +
79 +
80 +@pytest.fixture(scope="module")
81 +def server():
82 + args = [
83 + "--runner",
84 + "pooling",
85 + "--enforce-eager",
86 + "--trust-remote-code",
87 + "--skip-tokenizer-init",
88 + # Limit the maximum number of parallel requests
89 + # to avoid the model going OOM in CI.
90 + "--max-num-seqs",
91 + "32",
92 + "--io-processor-plugin",
93 + "prithvi_to_tiff_valencia"
94 + ]
95 +
96 + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
97 + yield remote_server
98 +
99 +
100 +@pytest.mark.asyncio
101 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
102 +async def test_prithvi_mae_plugin_online(
103 + server: RemoteOpenAIServer,
104 + model_name: str,
105 +):
106 +
107 + request_payload_url = {
108 + "data": {
109 + "data": image_url,
110 + "data_format": "url",
111 + "image_format": "tiff",
112 + "out_data_format": "b64_json",
113 + },
114 + "priority": 0,
115 + "model": model_name,
116 + }
117 +
118 + ret = requests.post(
119 + server.url_for("pooling"),
120 + json=request_payload_url,
121 + )
122 +
123 + response = ret.json()
124 +
125 + # verify the request response is in the correct format
126 + assert (parsed_response := IOProcessorResponse(**response))
127 +
128 + # verify the output is formatted as expected for this plugin
129 + plugin_data = parsed_response.data
130 +
131 + assert all(
132 + plugin_data.get(attr)
133 + for attr in ["type", "format", "data", "request_id"])
134 +
135 + # We just check that the output is a valid base64 string.
136 + # Raises an exception and fails the test if the string is corrupted.
137 + base64.b64decode(plugin_data["data"])
tests/plugins_tests/test_platform_plugins.py
+9
@@ -7,6 +7,15 @@ import torch
7 from vllm.plugins import load_general_plugins
8
9
10 +@pytest.fixture(scope="function", autouse=True)
11 +def use_v0_only(monkeypatch):
12 + """
13 + Since this module is V0 only, set VLLM_USE_V1=0 for
14 + all tests in the module.
15 + """
16 + monkeypatch.setenv('VLLM_USE_V1', '0')
17 +
18 +
19 def test_platform_plugins():
20 # simulate workload by running an example
21 import runpy
vllm/config/__init__.py
+2
@@ -501,6 +501,8 @@ class ModelConfig:
501 logits_processors: Optional[list[Union[str, type[LogitsProcessor]]]] = None
502 """One or more logits processors' fully-qualified class names or class
503 definitions"""
504 + io_processor_plugin: Optional[str] = None
505 + """IOProcessor plugin name to load at model startup"""
506
507 def compute_hash(self) -> str:
508 """
vllm/engine/arg_utils.py
+4
@@ -364,6 +364,7 @@ class EngineArgs:
364 disable_mm_preprocessor_cache: bool = False # DEPRECATED
365 mm_processor_cache_gb: float = MultiModalConfig.mm_processor_cache_gb
366 mm_encoder_tp_mode: MMEncoderTPMode = MultiModalConfig.mm_encoder_tp_mode
367 + io_processor_plugin: Optional[str] = None
368 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling
369 # LoRA fields
370 enable_lora: bool = False
@@ -577,6 +578,8 @@ class EngineArgs:
578 **model_kwargs["override_attention_dtype"])
579 model_group.add_argument("--logits-processors",
580 **model_kwargs["logits_processors"])
581 + model_group.add_argument("--io-processor-plugin",
582 + **model_kwargs["io_processor_plugin"])
583
584 # Model loading arguments
585 load_kwargs = get_kwargs(LoadConfig)
@@ -993,6 +996,7 @@ class EngineArgs:
996 model_impl=self.model_impl,
997 override_attention_dtype=self.override_attention_dtype,
998 logits_processors=self.logits_processors,
999 + io_processor_plugin=self.io_processor_plugin,
1000 )
1001
1002 def validate_tensorizer_args(self):
vllm/engine/protocol.py
+4
@@ -15,6 +15,7 @@ from vllm.logger import init_logger
15 from vllm.lora.request import LoRARequest
16 from vllm.model_executor.layers.sampler import SamplerOutput
17 from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput
18 +from vllm.plugins.io_processors.interface import IOProcessor
19 from vllm.pooling_params import PoolingParams
20 from vllm.sampling_params import BeamSearchParams, SamplingParams
21 from vllm.transformers_utils.tokenizer import AnyTokenizer
@@ -267,6 +268,9 @@ class EngineClient(ABC):
268 """Get the appropriate tokenizer for the request"""
269 ...
270
271 + async def get_io_processor(self) -> IOProcessor:
272 + raise NotImplementedError
273 +
274 @abstractmethod
275 async def is_tracing_enabled(self) -> bool:
276 ...
vllm/entrypoints/llm.py
+43 -4
@@ -37,13 +37,15 @@ from vllm.entrypoints.score_utils import (ScoreContentPartParam,
37 # yapf: enable
38 from vllm.entrypoints.utils import (_validate_truncation_size,
39 log_non_default_args)
40 -from vllm.inputs import PromptType, SingletonPrompt, TextPrompt, TokensPrompt
40 +from vllm.inputs import (DataPrompt, PromptType, SingletonPrompt, TextPrompt,
41 + TokensPrompt)
42 from vllm.logger import init_logger
43 from vllm.lora.request import LoRARequest
44 from vllm.model_executor.layers.quantization import QuantizationMethods
45 from vllm.outputs import (ClassificationRequestOutput, EmbeddingRequestOutput,
46 PoolingRequestOutput, RequestOutput,
47 ScoringRequestOutput)
48 +from vllm.plugins.io_processors import get_io_processor
49 from vllm.pooling_params import PoolingParams
50 from vllm.sampling_params import (BeamSearchParams, RequestOutputKind,
51 SamplingParams)
@@ -284,6 +286,11 @@ class LLM:
286
287 self.supported_tasks = supported_tasks
288
289 + # Load the Input/Output processor plugin if any
290 + io_processor_plugin = self.llm_engine.model_config.io_processor_plugin
291 + self.io_processor = get_io_processor(self.llm_engine.vllm_config,
292 + io_processor_plugin)
293 +
294 def get_tokenizer(
295 self,
296 lora_request: Optional[LoRARequest] = None,
@@ -833,7 +840,7 @@ class LLM:
840
841 def encode(
842 self,
836 - prompts: Union[PromptType, Sequence[PromptType]],
843 + prompts: Union[PromptType, Sequence[PromptType], DataPrompt],
844 pooling_params: Optional[Union[PoolingParams,
845 Sequence[PoolingParams]]] = None,
846 *,
@@ -915,6 +922,22 @@ class LLM:
922 if truncate_prompt_tokens is not None:
923 param.truncate_prompt_tokens = truncate_prompt_tokens
924
925 + io_processor_prompt = False
926 + if isinstance(prompts, dict) and "data" in prompts:
927 + io_processor_prompt = True
928 + if self.io_processor is None:
929 + raise ValueError(
930 + "No IOProcessor plugin installed. Please refer "
931 + "to the documentation and to the "
932 + "'prithvi_geospatial_mae_io_processor' "
933 + "offline inference example for more details.")
934 +
935 + # Validate the request data is valid for the loaded plugin
936 + validated_prompt = self.io_processor.parse_request(prompts)
937 +
938 + # obtain the actual model prompts from the pre-processor
939 + prompts = self.io_processor.pre_process(prompt=validated_prompt)
940 +
941 self._validate_and_add_requests(
942 prompts=prompts,
943 params=pooling_params,
@@ -923,8 +946,24 @@ class LLM:
946 )
947
948 outputs = self._run_engine(use_tqdm=use_tqdm)
926 - return self.engine_class.validate_outputs(outputs,
927 - PoolingRequestOutput)
949 +
950 + model_outputs = self.engine_class.validate_outputs(
951 + outputs, PoolingRequestOutput)
952 +
953 + if io_processor_prompt:
954 + # get the post-processed model outputs
955 + assert self.io_processor is not None
956 + processed_outputs = self.io_processor.post_process(
957 + model_output=model_outputs)
958 +
959 + return [
960 + PoolingRequestOutput[Any](request_id="",
961 + outputs=processed_outputs,
962 + prompt_token_ids=[],
963 + finished=True)
964 + ]
965 + else:
966 + return model_outputs
967
968 def embed(
969 self,
vllm/entrypoints/openai/api_server.py
+3 -2
@@ -64,6 +64,7 @@ from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
64 EmbeddingRequest,
65 EmbeddingResponse, ErrorInfo,
66 ErrorResponse,
67 + IOProcessorResponse,
68 LoadLoRAAdapterRequest,
69 PoolingRequest, PoolingResponse,
70 RerankRequest, RerankResponse,
@@ -795,7 +796,7 @@ async def create_pooling(request: PoolingRequest, raw_request: Request):
796 if isinstance(generator, ErrorResponse):
797 return JSONResponse(content=generator.model_dump(),
798 status_code=generator.error.code)
798 - elif isinstance(generator, PoolingResponse):
799 + elif isinstance(generator, (PoolingResponse, IOProcessorResponse)):
800 return JSONResponse(content=generator.model_dump())
801
802 assert_never(generator)
@@ -1782,7 +1783,7 @@ async def init_app_state(
1783 ) if "generate" in supported_tasks else None
1784 state.openai_serving_pooling = OpenAIServingPooling(
1785 engine_client,
1785 - model_config,
1786 + vllm_config,
1787 state.openai_serving_models,
1788 request_logger=request_logger,
1789 chat_template=resolved_chat_template,
vllm/entrypoints/openai/protocol.py
+42 -2
@@ -6,7 +6,8 @@
6 import json
7 import time
8 from http import HTTPStatus
9 -from typing import Annotated, Any, ClassVar, Literal, Optional, Union
9 +from typing import (Annotated, Any, ClassVar, Generic, Literal, Optional,
10 + TypeVar, Union)
11
12 import regex as re
13 import torch
@@ -1405,7 +1406,46 @@ EmbeddingRequest = Union[EmbeddingCompletionRequest, EmbeddingChatRequest]
1406
1407 PoolingCompletionRequest = EmbeddingCompletionRequest
1408 PoolingChatRequest = EmbeddingChatRequest
1408 -PoolingRequest = Union[PoolingCompletionRequest, PoolingChatRequest]
1409 +
1410 +T = TypeVar("T")
1411 +
1412 +
1413 +class IOProcessorRequest(OpenAIBaseModel, Generic[T]):
1414 + model: Optional[str] = None
1415 +
1416 + priority: int = Field(default=0)
1417 + """
1418 + The priority of the request (lower means earlier handling;
1419 + default: 0). Any priority other than 0 will raise an error
1420 + if the served model does not use priority scheduling.
1421 + """
1422 + data: T
1423 + """
1424 + When using plugins IOProcessor plugins, the actual input is processed
1425 + by the plugin itself. Hence, we use a generic type for the request data
1426 + """
1427 +
1428 + def to_pooling_params(self):
1429 + return PoolingParams(task="encode")
1430 +
1431 +
1432 +class IOProcessorResponse(OpenAIBaseModel, Generic[T]):
1433 +
1434 + request_id: Optional[str] = None
1435 + """
1436 + The request_id associated with this response
1437 + """
1438 + created_at: int = Field(default_factory=lambda: int(time.time()))
1439 +
1440 + data: T
1441 + """
1442 + When using plugins IOProcessor plugins, the actual output is generated
1443 + by the plugin itself. Hence, we use a generic type for the response data
1444 + """
1445 +
1446 +
1447 +PoolingRequest = Union[PoolingCompletionRequest, PoolingChatRequest,
1448 + IOProcessorRequest]
1449
1450
1451 class ScoreRequest(OpenAIBaseModel):
vllm/entrypoints/openai/serving_engine.py
+6 -4
@@ -49,9 +49,11 @@ from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
49 EmbeddingCompletionRequest,
50 EmbeddingRequest,
51 EmbeddingResponse, ErrorInfo,
52 - ErrorResponse, PoolingResponse,
53 - RerankRequest, ResponsesRequest,
54 - ScoreRequest, ScoreResponse,
52 + ErrorResponse,
53 + IOProcessorRequest,
54 + PoolingResponse, RerankRequest,
55 + ResponsesRequest, ScoreRequest,
56 + ScoreResponse,
57 TokenizeChatRequest,
58 TokenizeCompletionRequest,
59 TokenizeResponse,
@@ -89,7 +91,7 @@ ChatLikeRequest = Union[ChatCompletionRequest, EmbeddingChatRequest,
91 TokenizeChatRequest]
92 SpeechToTextRequest = Union[TranscriptionRequest, TranslationRequest]
93 AnyRequest = Union[CompletionLikeRequest, ChatLikeRequest, SpeechToTextRequest,
92 - ResponsesRequest]
94 + ResponsesRequest, IOProcessorRequest]
95
96 AnyResponse = Union[
97 CompletionResponse,
vllm/entrypoints/openai/serving_pooling.py
+57 -18
@@ -4,7 +4,7 @@
4 import asyncio
5 import base64
6 import time
7 -from collections.abc import AsyncGenerator
7 +from collections.abc import AsyncGenerator, Sequence
8 from typing import Final, Literal, Optional, Union, cast
9
10 import jinja2
@@ -13,19 +13,25 @@ import torch
13 from fastapi import Request
14 from typing_extensions import assert_never
15
16 -from vllm.config import ModelConfig
16 +from vllm.config import VllmConfig
17 from vllm.engine.protocol import EngineClient
18 from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
19 from vllm.entrypoints.logger import RequestLogger
20 +# yapf: disable
21 from vllm.entrypoints.openai.protocol import (ErrorResponse,
22 + IOProcessorRequest,
23 + IOProcessorResponse,
24 PoolingChatRequest,
25 + PoolingCompletionRequest,
26 PoolingRequest, PoolingResponse,
27 PoolingResponseData, UsageInfo)
24 -from vllm.entrypoints.openai.serving_engine import OpenAIServing
28 +# yapf: enable
29 +from vllm.entrypoints.openai.serving_engine import OpenAIServing, RequestPrompt
30 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
31 from vllm.entrypoints.utils import _validate_truncation_size
32 from vllm.logger import init_logger
33 from vllm.outputs import PoolingOutput, PoolingRequestOutput
34 +from vllm.plugins.io_processors import get_io_processor
35 from vllm.utils import merge_async_iterators
36
37 logger = init_logger(__name__)
@@ -52,7 +58,7 @@ class OpenAIServingPooling(OpenAIServing):
58 def __init__(
59 self,
60 engine_client: EngineClient,
55 - model_config: ModelConfig,
61 + vllm_config: VllmConfig,
62 models: OpenAIServingModels,
63 *,
64 request_logger: Optional[RequestLogger],
@@ -61,19 +67,21 @@ class OpenAIServingPooling(OpenAIServing):
67 log_error_stack: bool = False,
68 ) -> None:
69 super().__init__(engine_client=engine_client,
64 - model_config=model_config,
70 + model_config=vllm_config.model_config,
71 models=models,
72 request_logger=request_logger,
73 log_error_stack=log_error_stack)
74
75 self.chat_template = chat_template
76 self.chat_template_content_format: Final = chat_template_content_format
77 + io_processor_plugin = self.model_config.io_processor_plugin
78 + self.io_processor = get_io_processor(vllm_config, io_processor_plugin)
79
80 async def create_pooling(
81 self,
82 request: PoolingRequest,
83 raw_request: Optional[Request] = None,
76 - ) -> Union[PoolingResponse, ErrorResponse]:
84 + ) -> Union[PoolingResponse, IOProcessorResponse, ErrorResponse]:
85 """
86 See https://platform.openai.com/docs/api-reference/embeddings/create
87 for the API specification. This API mimics the OpenAI Embedding API.
@@ -82,20 +90,13 @@ class OpenAIServingPooling(OpenAIServing):
90 if error_check_ret is not None:
91 return error_check_ret
92
85 - encoding_format = request.encoding_format
86 - if request.dimensions is not None:
87 - return self.create_error_response(
88 - "dimensions is currently not supported")
89 -
93 model_name = self._get_model_name(request.model)
94 +
95 request_id = f"pool-{self._base_request_id(raw_request)}"
96 created_time = int(time.time())
97
94 - truncate_prompt_tokens = request.truncate_prompt_tokens
95 -
98 + is_io_processor_request = isinstance(request, IOProcessorRequest)
99 try:
97 - truncate_prompt_tokens = _validate_truncation_size(
98 - self.max_model_len, truncate_prompt_tokens)
100 lora_request = self._maybe_get_adapters(request)
101
102 if self.model_config.skip_tokenizer_init:
@@ -104,7 +105,32 @@ class OpenAIServingPooling(OpenAIServing):
105 tokenizer = await self.engine_client.get_tokenizer(lora_request
106 )
107
107 - if isinstance(request, PoolingChatRequest):
108 + if getattr(request, "dimensions", None) is not None:
109 + return self.create_error_response(
110 + "dimensions is currently not supported")
111 +
112 + truncate_prompt_tokens = getattr(request, "truncate_prompt_tokens",
113 + None)
114 + truncate_prompt_tokens = _validate_truncation_size(
115 + self.max_model_len, truncate_prompt_tokens)
116 +
117 + if is_io_processor_request:
118 + if self.io_processor is None:
119 + raise ValueError(
120 + "No IOProcessor plugin installed. Please refer "
121 + "to the documentation and to the "
122 + "'prithvi_geospatial_mae_io_processor' "
123 + "offline inference example for more details.")
124 +
125 + validated_prompt = self.io_processor.parse_request(request)
126 +
127 + engine_prompts = await self.io_processor.pre_process_async(
128 + prompt=validated_prompt, request_id=request_id)
129 + request_prompts: Sequence[RequestPrompt] = [
130 + ""
131 + ] * len(engine_prompts)
132 +
133 + elif isinstance(request, PoolingChatRequest):
134 (
135 _,
136 request_prompts,
@@ -122,7 +148,7 @@ class OpenAIServingPooling(OpenAIServing):
148 continue_final_message=False,
149 add_special_tokens=request.add_special_tokens,
150 )
125 - else:
151 + elif isinstance(request, PoolingCompletionRequest):
152 (request_prompts,
153 engine_prompts) = await self._preprocess_completion(
154 request,
@@ -130,6 +156,9 @@ class OpenAIServingPooling(OpenAIServing):
156 request.input,
157 add_special_tokens=request.add_special_tokens,
158 )
159 + else:
160 + raise ValueError(
161 + f"Unsupported request of type {type(request)}")
162 except (ValueError, TypeError, jinja2.TemplateError) as e:
163 logger.exception("Error in preprocessing prompt inputs")
164 return self.create_error_response(str(e))
@@ -171,6 +200,16 @@ class OpenAIServingPooling(OpenAIServing):
200
201 result_generator = merge_async_iterators(*generators)
202
203 + if is_io_processor_request:
204 + assert self.io_processor is not None
205 + output = await self.io_processor.post_process_async(
206 + model_output=result_generator,
207 + request_id=request_id,
208 + )
209 + return self.io_processor.output_to_response(output)
210 +
211 + assert isinstance(request,
212 + (PoolingCompletionRequest, PoolingChatRequest))
213 num_prompts = len(engine_prompts)
214
215 # Non-streaming response
@@ -190,7 +229,7 @@ class OpenAIServingPooling(OpenAIServing):
229 request_id,
230 created_time,
231 model_name,
193 - encoding_format,
232 + request.encoding_format,
233 )
234 except asyncio.CancelledError:
235 return self.create_error_response("Client disconnected")
vllm/inputs/__init__.py
+2 -1
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from .data import (DecoderOnlyInputs, EmbedsInputs, EmbedsPrompt,
4 +from .data import (DataPrompt, DecoderOnlyInputs, EmbedsInputs, EmbedsPrompt,
5 EncoderDecoderInputs, ExplicitEncoderDecoderPrompt,
6 ProcessorInputs, PromptType, SingletonInputs,
7 SingletonPrompt, TextPrompt, TokenInputs, TokensPrompt,
@@ -18,6 +18,7 @@ target model.
18 """
19
20 __all__ = [
21 + "DataPrompt",
22 "TextPrompt",
23 "TokensPrompt",
24 "PromptType",
vllm/inputs/data.py
+10
@@ -95,6 +95,16 @@ class EmbedsPrompt(TypedDict):
95 """
96
97
98 +class DataPrompt(TypedDict):
99 + """Represents generic inputs handled by IO processor plugins."""
100 +
101 + data: Any
102 + """The input data"""
103 +
104 + data_format: str
105 + """The input data format"""
106 +
107 +
108 SingletonPrompt = Union[str, TextPrompt, TokensPrompt, EmbedsPrompt]
109 """
110 Set of possible schemas for a single prompt:
vllm/plugins/io_processors/__init__.py new
+68
@@ -0,0 +1,68 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from __future__ import annotations
5 +
6 +import logging
7 +from typing import Optional
8 +
9 +from vllm.config import VllmConfig
10 +from vllm.plugins import load_plugins_by_group
11 +from vllm.plugins.io_processors.interface import IOProcessor
12 +from vllm.utils import resolve_obj_by_qualname
13 +
14 +logger = logging.getLogger(__name__)
15 +
16 +
17 +def get_io_processor(
18 + vllm_config: VllmConfig,
19 + plugin_from_init: Optional[str] = None) -> IOProcessor | None:
20 + # Input.Output processors are loaded as plugins under the
21 + # 'vllm.io_processor_plugins' group. Similar to platform
22 + # plugins, these plugins register a function that returns the class
23 + # name for the processor to install.
24 +
25 + if plugin_from_init:
26 + model_plugin = plugin_from_init
27 + else:
28 + # A plugin can be specified via the model config
29 + # Retrieve the model specific plugin if available
30 + # This is using a custom field in the hf_config for the model
31 + hf_config = vllm_config.model_config.hf_config.to_dict()
32 + config_plugin = hf_config.get("io_processor_plugin")
33 + model_plugin = config_plugin
34 +
35 + if model_plugin is None:
36 + logger.info("No IOProcessor plugins requested by the model")
37 + return None
38 +
39 + logger.debug("IOProcessor plugin to be loaded %s", model_plugin)
40 +
41 + # Load all installed plugin in the group
42 + multimodal_data_processor_plugins = \
43 + load_plugins_by_group('vllm.io_processor_plugins')
44 +
45 + loadable_plugins = {}
46 + for name, func in multimodal_data_processor_plugins.items():
47 + try:
48 + assert callable(func)
49 + processor_cls_qualname = func()
50 + if processor_cls_qualname is not None:
51 + loadable_plugins[name] = processor_cls_qualname
52 + except Exception:
53 + logger.warning("Failed to load plugin %s.", name, exc_info=True)
54 +
55 + num_available_plugins = len(loadable_plugins.keys())
56 + if num_available_plugins == 0:
57 + raise ValueError("No IOProcessor plugins installed"
58 + f" but one is required ({model_plugin}).")
59 +
60 + if model_plugin not in loadable_plugins:
61 + raise ValueError(
62 + f"The model requires the '{model_plugin}' IO Processor plugin "
63 + "but it is not installed. "
64 + f"Available plugins: {list(loadable_plugins.keys())}")
65 +
66 + activated_plugin_cls = loadable_plugins[model_plugin]
67 +
68 + return resolve_obj_by_qualname(activated_plugin_cls)(vllm_config)
vllm/plugins/io_processors/interface.py new
+62
@@ -0,0 +1,62 @@
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, Sequence
6 +from typing import Any, Generic, Optional, TypeVar, Union
7 +
8 +from vllm.config import VllmConfig
9 +from vllm.entrypoints.openai.protocol import IOProcessorResponse
10 +from vllm.inputs.data import PromptType
11 +from vllm.outputs import PoolingRequestOutput
12 +
13 +IOProcessorInput = TypeVar('IOProcessorInput')
14 +IOProcessorOutput = TypeVar('IOProcessorOutput')
15 +
16 +
17 +class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):
18 +
19 + def __init__(self, vllm_config: VllmConfig):
20 + self.vllm_config = vllm_config
21 +
22 + @abstractmethod
23 + def pre_process(
24 + self,
25 + prompt: IOProcessorInput,
26 + request_id: Optional[str] = None,
27 + **kwargs,
28 + ) -> Union[PromptType, Sequence[PromptType]]:
29 + raise NotImplementedError
30 +
31 + async def pre_process_async(
32 + self,
33 + prompt: IOProcessorInput,
34 + request_id: Optional[str] = None,
35 + **kwargs,
36 + ) -> Union[PromptType, Sequence[PromptType]]:
37 + return self.pre_process(prompt, request_id, **kwargs)
38 +
39 + @abstractmethod
40 + def post_process(self,
41 + model_output: Sequence[PoolingRequestOutput],
42 + request_id: Optional[str] = None,
43 + **kwargs) -> IOProcessorOutput:
44 + raise NotImplementedError
45 +
46 + async def post_process_async(
47 + self,
48 + model_output: AsyncGenerator[tuple[int, PoolingRequestOutput]],
49 + request_id: Optional[str] = None,
50 + **kwargs,
51 + ) -> IOProcessorOutput:
52 + collected_output = [item async for i, item in model_output]
53 + return self.post_process(collected_output, request_id, **kwargs)
54 +
55 + @abstractmethod
56 + def parse_request(self, request: Any) -> IOProcessorInput:
57 + raise NotImplementedError
58 +
59 + @abstractmethod
60 + def output_to_response(
61 + self, plugin_output: IOProcessorOutput) -> IOProcessorResponse:
62 + raise NotImplementedError
\ No newline at end of file