@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2f2212e6cc
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 7, 06:30 UTC

Split generic IO Processor plugins tests from Terratorch specific ones (#35756)

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

Christian Pinto committed Mar 4, 2026 at 16:01 UTC 2f2212e6ccfc01d123879d635d19448f5cc3653c
4 files changed +239 -136
.buildkite/test_areas/plugins.yaml
+5 -2
@@ -15,9 +15,12 @@ steps:
15 - pytest -v -s plugins_tests/test_platform_plugins.py
16 - pip uninstall vllm_add_dummy_platform -y
17 # end platform plugin tests
18 - # begin io_processor plugins test, all the code in between uses the prithvi_io_processor plugin
18 + # begin io_processor plugins test
19 + # test generic io_processor plugins functions
20 + - pytest -v -s ./plugins_tests/test_io_processor_plugins.py
21 + # test Terratorch io_processor plugins
22 - pip install -e ./plugins/prithvi_io_processor_plugin
20 - - pytest -v -s plugins_tests/test_io_processor_plugins.py
23 + - pytest -v -s plugins_tests/test_terratorch_io_processor_plugins.py
24 - pip uninstall prithvi_io_processor_plugin -y
25 # test bge_m3_sparse io_processor plugin
26 - pip install -e ./plugins/bge_m3_sparse_plugin
tests/conftest.py
+9
@@ -1140,6 +1140,15 @@ class VllmRunner:
1140 return self
1141
1142 def __exit__(self, exc_type, exc_value, traceback):
1143 + # Explicitly shutdown the engine core to release GPU resources
1144 + # This is needed because when executing consecutive tests, the GC
1145 + # might not be fast enough in shutting down the llm engine. This can lead to OOMs
1146 + # because when the next test starts some GPU memory is still in use.
1147 + try:
1148 + self.llm.llm_engine.engine_core.shutdown()
1149 + except Exception:
1150 + # Ignore shutdown errors as cleanup will still proceed
1151 + pass
1152 del self.llm
1153 cleanup_dist_env_and_memory()
1154
tests/plugins_tests/test_io_processor_plugins.py
+78 -134
@@ -1,154 +1,98 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import base64
4 -import io
3 +from collections.abc import Sequence
4 +from unittest.mock import MagicMock, patch
5
6 -import imagehash
6 import pytest
8 -import requests
9 -from PIL import Image
7
11 -from tests.utils import RemoteOpenAIServer
8 from vllm.config import VllmConfig
13 -from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse
9 +from vllm.inputs.data import PromptType
10 +from vllm.outputs import PoolingRequestOutput
11 from vllm.plugins.io_processors import get_io_processor
15 -
16 -models_config = {
17 - "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11": {
18 - "image_url": "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff", # noqa: E501
19 - "out_hash": "aa6d92ad25926a5e",
20 - "plugin": "prithvi_to_tiff",
21 - },
22 - "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-BurnScars": {
23 - "image_url": "https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-BurnScars/resolve/main/examples/subsetted_512x512_HLS.S30.T10SEH.2018190.v1.4_merged.tif", # noqa: E501
24 - "out_hash": "c07f4f602da73552",
25 - "plugin": "prithvi_to_tiff",
26 - },
27 -}
28 -
29 -
30 -def _compute_image_hash(base64_data: str) -> str:
31 - # Decode the base64 output and create image from byte stream
32 - decoded_image = base64.b64decode(base64_data)
33 - image = Image.open(io.BytesIO(decoded_image))
34 -
35 - # Compute perceptual hash of the output image
36 - return str(imagehash.phash(image))
12 +from vllm.plugins.io_processors.interface import IOProcessor
13 +from vllm.renderers import BaseRenderer
14 +
15 +
16 +class DummyIOProcessor(IOProcessor):
17 + """Minimal IOProcessor used as the target of the mocked plugin entry point."""
18 +
19 + def pre_process(
20 + self,
21 + prompt: object,
22 + request_id: str | None = None,
23 + **kwargs,
24 + ) -> PromptType | Sequence[PromptType]:
25 + raise NotImplementedError
26 +
27 + def post_process(
28 + self,
29 + model_output: Sequence[PoolingRequestOutput],
30 + request_id: str | None = None,
31 + **kwargs,
32 + ) -> object:
33 + raise NotImplementedError
34 +
35 +
36 +@pytest.fixture
37 +def my_plugin_entry_points():
38 + """Patch importlib.metadata.entry_points to expose a single 'my_plugin'
39 + entry point backed by DummyIOProcessor, exercising the full plugin-loading
40 + code path: entry_points → plugin.load() → func() →
41 + resolve_obj_by_qualname → IOProcessor.__init__."""
42 + qualname = f"{DummyIOProcessor.__module__}.{DummyIOProcessor.__qualname__}"
43 + ep = MagicMock()
44 + ep.name = "my_plugin"
45 + ep.value = qualname
46 + ep.load.return_value = lambda: qualname
47 + with patch("importlib.metadata.entry_points", return_value=[ep]):
48 + yield
49
50
51 def test_loading_missing_plugin():
52 vllm_config = VllmConfig()
53 + renderer = MagicMock(spec=BaseRenderer)
54 with pytest.raises(ValueError):
42 - get_io_processor(vllm_config, None, "wrong_plugin")
43 -
44 -
45 -@pytest.fixture(scope="function")
46 -def server(model_name, plugin):
47 - args = [
48 - "--runner",
49 - "pooling",
50 - "--enforce-eager",
51 - "--skip-tokenizer-init",
52 - # Limit the maximum number of parallel requests
53 - # to avoid the model going OOM in CI.
54 - "--max-num-seqs",
55 - "32",
56 - "--io-processor-plugin",
57 - plugin,
58 - "--enable-mm-embeds",
59 - ]
60 -
61 - with RemoteOpenAIServer(model_name, args) as remote_server:
62 - yield remote_server
63 -
64 -
65 -@pytest.mark.asyncio
66 -@pytest.mark.parametrize(
67 - "model_name, image_url, plugin, expected_hash",
68 - [
69 - (model_name, config["image_url"], config["plugin"], config["out_hash"])
70 - for model_name, config in models_config.items()
71 - ],
72 -)
73 -async def test_prithvi_mae_plugin_online(
74 - server: RemoteOpenAIServer,
75 - model_name: str,
76 - image_url: str | dict,
77 - plugin: str,
78 - expected_hash: str,
79 -):
80 - request_payload_url = {
81 - "data": {
82 - "data": image_url,
83 - "data_format": "url",
84 - "image_format": "tiff",
85 - "out_data_format": "b64_json",
86 - },
87 - "priority": 0,
88 - "model": model_name,
89 - "softmax": False,
90 - }
91 -
92 - ret = requests.post(
93 - server.url_for("pooling"),
94 - json=request_payload_url,
55 + get_io_processor(
56 + vllm_config, renderer=renderer, plugin_from_init="wrong_plugin"
57 + )
58 +
59 +
60 +def test_loading_plugin(my_plugin_entry_points):
61 + # Plugin name supplied via plugin_from_init.
62 + vllm_config = MagicMock(spec=VllmConfig)
63 + renderer = MagicMock(spec=BaseRenderer)
64 +
65 + result = get_io_processor(
66 + vllm_config, renderer=renderer, plugin_from_init="my_plugin"
67 )
68
97 - response = ret.json()
69 + assert isinstance(result, DummyIOProcessor)
70
99 - # verify the request response is in the correct format
100 - assert (parsed_response := IOProcessorResponse(**response))
71
102 - # verify the output is formatted as expected for this plugin
103 - plugin_data = parsed_response.data
104 - assert all(plugin_data.get(attr) for attr in ["type", "format", "data"])
72 +def test_loading_missing_plugin_from_model_config():
73 + # Build a mock VllmConfig whose hf_config advertises a plugin name,
74 + # exercising the model-config code path without loading a real model.
75 + mock_hf_config = MagicMock()
76 + mock_hf_config.to_dict.return_value = {"io_processor_plugin": "wrong_plugin"}
77
106 - # Compute the output image hash and compare it against the expected hash
107 - image_hash = _compute_image_hash(plugin_data["data"])
108 - assert image_hash == expected_hash, (
109 - f"Image hash mismatch: expected {expected_hash}, got {image_hash}"
110 - )
78 + vllm_config = MagicMock(spec=VllmConfig)
79 + vllm_config.model_config.hf_config = mock_hf_config
80
81 + renderer = MagicMock(spec=BaseRenderer)
82 + with pytest.raises(ValueError):
83 + get_io_processor(vllm_config, renderer=renderer)
84
113 -@pytest.mark.parametrize(
114 - "model_name, image_url, plugin, expected_hash",
115 - [
116 - (model_name, config["image_url"], config["plugin"], config["out_hash"])
117 - for model_name, config in models_config.items()
118 - ],
119 -)
120 -def test_prithvi_mae_plugin_offline(
121 - vllm_runner, model_name: str, image_url: str | dict, plugin: str, expected_hash: str
122 -):
123 - img_data = dict(
124 - data=image_url,
125 - data_format="url",
126 - image_format="tiff",
127 - out_data_format="b64_json",
128 - )
85
130 - prompt = dict(data=img_data)
131 -
132 - with vllm_runner(
133 - model_name,
134 - runner="pooling",
135 - skip_tokenizer_init=True,
136 - enable_mm_embeds=True,
137 - enforce_eager=True,
138 - # Limit the maximum number of parallel requests
139 - # to avoid the model going OOM in CI.
140 - max_num_seqs=32,
141 - io_processor_plugin=plugin,
142 - default_torch_num_threads=1,
143 - ) as llm_runner:
144 - pooler_output = llm_runner.get_llm().encode(prompt, pooling_task="plugin")
145 - output = pooler_output[0].outputs
146 -
147 - # verify the output is formatted as expected for this plugin
148 - assert all(hasattr(output, attr) for attr in ["type", "format", "data"])
149 -
150 - # Compute the output image hash and compare it against the expected hash
151 - image_hash = _compute_image_hash(output.data)
152 - assert image_hash == expected_hash, (
153 - f"Image hash mismatch: expected {expected_hash}, got {image_hash}"
154 - )
86 +def test_loading_plugin_from_model_config(my_plugin_entry_points):
87 + # Plugin name supplied via the model's hf_config.
88 + mock_hf_config = MagicMock()
89 + mock_hf_config.to_dict.return_value = {"io_processor_plugin": "my_plugin"}
90 +
91 + vllm_config = MagicMock(spec=VllmConfig)
92 + vllm_config.model_config.hf_config = mock_hf_config
93 +
94 + renderer = MagicMock(spec=BaseRenderer)
95 +
96 + result = get_io_processor(vllm_config, renderer=renderer)
97 +
98 + assert isinstance(result, DummyIOProcessor)
tests/plugins_tests/test_terratorch_io_processor_plugins.py new
+147
@@ -0,0 +1,147 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import base64
4 +import io
5 +
6 +import imagehash
7 +import pytest
8 +import requests
9 +from PIL import Image
10 +
11 +from tests.utils import RemoteOpenAIServer
12 +from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse
13 +
14 +models_config = {
15 + "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11": {
16 + "image_url": "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff", # noqa: E501
17 + "out_hash": "aa6d92ad25926a5e",
18 + "plugin": "prithvi_to_tiff",
19 + },
20 + "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-BurnScars": {
21 + "image_url": "https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-BurnScars/resolve/main/examples/subsetted_512x512_HLS.S30.T10SEH.2018190.v1.4_merged.tif", # noqa: E501
22 + "out_hash": "c07f4f602da73552",
23 + "plugin": "prithvi_to_tiff",
24 + },
25 +}
26 +
27 +
28 +def _compute_image_hash(base64_data: str) -> str:
29 + # Decode the base64 output and create image from byte stream
30 + decoded_image = base64.b64decode(base64_data)
31 + image = Image.open(io.BytesIO(decoded_image))
32 +
33 + # Compute perceptual hash of the output image
34 + return str(imagehash.phash(image))
35 +
36 +
37 +@pytest.fixture(scope="function")
38 +def server(model_name, plugin):
39 + args = [
40 + "--runner",
41 + "pooling",
42 + "--enforce-eager",
43 + "--skip-tokenizer-init",
44 + # Limit the maximum number of parallel requests
45 + # to avoid the model going OOM in CI.
46 + "--max-num-seqs",
47 + "32",
48 + "--io-processor-plugin",
49 + plugin,
50 + "--enable-mm-embeds",
51 + ]
52 +
53 + with RemoteOpenAIServer(model_name, args) as remote_server:
54 + yield remote_server
55 +
56 +
57 +@pytest.mark.asyncio
58 +@pytest.mark.parametrize(
59 + "model_name, image_url, plugin, expected_hash",
60 + [
61 + (model_name, config["image_url"], config["plugin"], config["out_hash"])
62 + for model_name, config in models_config.items()
63 + ],
64 +)
65 +async def test_prithvi_mae_plugin_online(
66 + server: RemoteOpenAIServer,
67 + model_name: str,
68 + image_url: str | dict,
69 + plugin: str,
70 + expected_hash: str,
71 +):
72 + request_payload_url = {
73 + "data": {
74 + "data": image_url,
75 + "data_format": "url",
76 + "image_format": "tiff",
77 + "out_data_format": "b64_json",
78 + },
79 + "priority": 0,
80 + "model": model_name,
81 + "softmax": False,
82 + }
83 +
84 + ret = requests.post(
85 + server.url_for("pooling"),
86 + json=request_payload_url,
87 + )
88 +
89 + response = ret.json()
90 +
91 + # verify the request response is in the correct format
92 + assert (parsed_response := IOProcessorResponse(**response))
93 +
94 + # verify the output is formatted as expected for this plugin
95 + plugin_data = parsed_response.data
96 + assert all(plugin_data.get(attr) for attr in ["type", "format", "data"])
97 +
98 + # Compute the output image hash and compare it against the expected hash
99 + image_hash = _compute_image_hash(plugin_data["data"])
100 + assert image_hash == expected_hash, (
101 + f"Image hash mismatch: expected {expected_hash}, got {image_hash}"
102 + )
103 +
104 +
105 +@pytest.mark.parametrize(
106 + "model_name, image_url, plugin, expected_hash",
107 + [
108 + (model_name, config["image_url"], config["plugin"], config["out_hash"])
109 + for model_name, config in models_config.items()
110 + ],
111 +)
112 +def test_prithvi_mae_plugin_offline(
113 + vllm_runner, model_name: str, image_url: str | dict, plugin: str, expected_hash: str
114 +):
115 + img_data = dict(
116 + data=image_url,
117 + data_format="url",
118 + image_format="tiff",
119 + out_data_format="b64_json",
120 + )
121 +
122 + prompt = dict(data=img_data)
123 +
124 + with vllm_runner(
125 + model_name,
126 + runner="pooling",
127 + skip_tokenizer_init=True,
128 + enable_mm_embeds=True,
129 + enforce_eager=True,
130 + # Limit the maximum number of parallel requests
131 + # to avoid the model going OOM in CI.
132 + max_num_seqs=32,
133 + io_processor_plugin=plugin,
134 + default_torch_num_threads=1,
135 + ) as llm_runner:
136 + pooler_output = llm_runner.get_llm().encode(prompt, pooling_task="plugin")
137 +
138 + output = pooler_output[0].outputs
139 +
140 + # verify the output is formatted as expected for this plugin
141 + assert all(hasattr(output, attr) for attr in ["type", "format", "data"])
142 +
143 + # Compute the output image hash and compare it against the expected hash
144 + image_hash = _compute_image_hash(output.data)
145 + assert image_hash == expected_hash, (
146 + f"Image hash mismatch: expected {expected_hash}, got {image_hash}"
147 + )