@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 323a29d8f2

[Bugfix] Unreadable `prompt_embeds` payload should be a 400, not a 500 (#55451)

Signed-off-by: lzhan011 <35493221+lzhan011@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

lzhan011 committed Sep 16, 2026 at 00:53 UTC 323a29d8f2b2a50ed77a6173f120822376a8b14a
4 files changed +230 -17
tests/renderers/test_prompt_embeds_payload_validation.py new
+163
@@ -0,0 +1,163 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""`prompt_embeds` payloads that are not tensor files must be client errors.
4 +
5 +`prompt_embeds` is a base64 blob the caller writes, and every other way of
6 +getting it wrong -- wrong rank, wrong hidden_size, wrong dtype, not a tensor at
7 +all -- already answers 400. A payload that `torch.load` cannot read at all was
8 +the exception: torch has no single error type for that (the zip reader raises
9 +`RuntimeError`, the legacy pickle path raises an unpickling error, `KeyError` or
10 +`IndexError`, an empty payload raises `EOFError`), none of them is a
11 +`ValueError`, and the entrypoints' fallback therefore mapped them to 500.
12 +
13 +The assertions go through `create_error_response`, which is what actually turns
14 +an exception into a status code, so they pin the HTTP answer rather than the
15 +exception class alone.
16 +"""
17 +
18 +import io
19 +import zipfile
20 +from unittest.mock import patch
21 +
22 +import pybase64 as base64
23 +import pytest
24 +import torch
25 +
26 +from vllm.entrypoints.serve.exception_handling.error_response import (
27 + create_error_response,
28 +)
29 +from vllm.exceptions import VLLMValidationError
30 +from vllm.renderers.embed_utils import (
31 + _MAX_EMBED_ERROR_REASON_CHARS,
32 + safe_load_prompt_embeds,
33 +)
34 +
35 +
36 +@pytest.fixture
37 +def model_config():
38 + from vllm.config import ModelConfig
39 +
40 + return ModelConfig(
41 + model="facebook/opt-125m",
42 + tokenizer="facebook/opt-125m",
43 + tokenizer_mode="auto",
44 + trust_remote_code=False,
45 + dtype="float32",
46 + seed=0,
47 + enable_prompt_embeds=True,
48 + )
49 +
50 +
51 +def _encode(raw: bytes) -> bytes:
52 + return base64.b64encode(raw)
53 +
54 +
55 +def _zip_without_a_tensor() -> bytes:
56 + buffer = io.BytesIO()
57 + with zipfile.ZipFile(buffer, "w") as archive:
58 + archive.writestr("junk.txt", "hello")
59 + return buffer.getvalue()
60 +
61 +
62 +def _saved_tensor(tensor: torch.Tensor) -> bytes:
63 + buffer = io.BytesIO()
64 + torch.save(tensor, buffer)
65 + return buffer.getvalue()
66 +
67 +
68 +UNPARSABLE_PAYLOADS = [
69 + pytest.param(b"", id="empty"),
70 + pytest.param(b"\x80", id="single-byte"),
71 + pytest.param(b"\x00\x01\x02\x03not-a-tensor", id="random-bytes"),
72 + pytest.param(b"hello world this is not a torch file", id="ascii-text"),
73 + pytest.param(b"PK\x03\x04" + b"\x00" * 20, id="truncated-zip"),
74 + pytest.param(_zip_without_a_tensor(), id="zip-without-a-tensor"),
75 + # pickle.dumps(42): a valid pickle stream that is not a torch checkpoint.
76 + # Spelled out rather than built, so this file does not import pickle.
77 + pytest.param(b"\x80\x04K*.", id="legacy-pickle-of-an-int"),
78 +]
79 +
80 +
81 +@pytest.mark.parametrize("payload", UNPARSABLE_PAYLOADS)
82 +def test_unreadable_prompt_embeds_payload_is_a_client_error(model_config, payload):
83 + with pytest.raises(VLLMValidationError) as excinfo:
84 + safe_load_prompt_embeds(model_config, _encode(payload))
85 +
86 + assert excinfo.value.parameter == "prompt_embeds"
87 + assert int(create_error_response(excinfo.value).error.code) == 400
88 +
89 +
90 +@pytest.mark.parametrize("payload", UNPARSABLE_PAYLOADS)
91 +def test_the_error_body_stays_bounded(model_config, payload):
92 + """torch's reason is built from the caller's bytes, so it must be truncated.
93 +
94 + An unpickler that reports the offending global, for instance, quotes a name
95 + the caller chose; without a bound a large one would be reflected whole.
96 + """
97 + padding = b"A" * 100_000
98 + with pytest.raises(VLLMValidationError) as excinfo:
99 + safe_load_prompt_embeds(model_config, _encode(payload + padding))
100 +
101 + assert len(str(excinfo.value)) < 500
102 +
103 +
104 +def test_a_long_reason_is_truncated_and_says_how_much_was_left_out(model_config):
105 + """The cap is pinned against a reason of known length, not a real payload.
106 +
107 + Which of torch's messages happen to exceed the cap is torch's business and
108 + changes between releases, so the reason is injected here: that keeps the
109 + assertion about *this* code's truncation rule rather than about whichever
110 + error text the installed torch produces.
111 + """
112 + reason = "E" * (_MAX_EMBED_ERROR_REASON_CHARS + 37)
113 +
114 + with (
115 + patch(
116 + "vllm.renderers.embed_utils.torch.load", side_effect=RuntimeError(reason)
117 + ),
118 + pytest.raises(VLLMValidationError) as excinfo,
119 + ):
120 + safe_load_prompt_embeds(model_config, _encode(b"ignored"))
121 +
122 + message = str(excinfo.value)
123 + assert "E" * _MAX_EMBED_ERROR_REASON_CHARS in message
124 + assert "E" * (_MAX_EMBED_ERROR_REASON_CHARS + 1) not in message
125 + assert "... (37 more characters truncated)" in message
126 +
127 +
128 +def test_a_short_reason_is_passed_through_whole(model_config):
129 + """A reason that fits is not annotated -- the count only appears when it is
130 + telling the reader something."""
131 + reason = "invalid magic number; corrupt file?"
132 + assert len(reason) <= _MAX_EMBED_ERROR_REASON_CHARS
133 +
134 + with (
135 + patch(
136 + "vllm.renderers.embed_utils.torch.load", side_effect=RuntimeError(reason)
137 + ),
138 + pytest.raises(VLLMValidationError) as excinfo,
139 + ):
140 + safe_load_prompt_embeds(model_config, _encode(b"ignored"))
141 +
142 + message = str(excinfo.value)
143 + assert reason in message
144 + assert "truncated" not in message
145 +
146 +
147 +def test_a_well_formed_payload_still_loads(model_config):
148 + """Positive control: the healthy path is untouched."""
149 + hidden_size = model_config.get_hidden_size()
150 + tensor = torch.zeros(3, hidden_size, dtype=torch.float32)
151 +
152 + loaded = safe_load_prompt_embeds(model_config, _encode(_saved_tensor(tensor)))
153 +
154 + assert loaded.shape == (3, hidden_size)
155 + assert loaded.dtype == model_config.dtype
156 +
157 +
158 +def test_a_payload_that_is_not_base64_is_still_a_client_error(model_config):
159 + """Positive control: the base64 stage already answered 400 and still does."""
160 + with pytest.raises(ValueError) as excinfo:
161 + safe_load_prompt_embeds(model_config, b"!!!!not base64!!!!")
162 +
163 + assert int(create_error_response(excinfo.value).error.code) == 400
tests/renderers/test_sparse_tensor_concurrent_race.py
+6 -5
@@ -11,6 +11,7 @@ import pytest
11 import torch
12
13 from vllm.config import ModelConfig
14 +from vllm.exceptions import VLLMValidationError
15 from vllm.multimodal.media import AudioEmbeddingMediaIO, ImageEmbeddingMediaIO
16 from vllm.renderers.embed_utils import safe_load_prompt_embeds
17 from vllm.utils.sparse_utils import check_sparse_tensor_invariants_threadsafe
@@ -54,7 +55,7 @@ class TestNegativeControlWithoutRace:
55
56 def test_malicious_sparse_rejected_by_prompt_loader(self, model_config):
57 encoded = _encode_tensor(_create_malicious_sparse_tensor())
57 - with pytest.raises((RuntimeError, ValueError)):
58 + with pytest.raises((RuntimeError, ValueError, VLLMValidationError)):
59 safe_load_prompt_embeds(model_config, encoded)
60
61 def test_malicious_sparse_rejected_by_image_loader(self):
@@ -147,7 +148,7 @@ class TestConcurrentRaceProtection:
148 try:
149 safe_load_prompt_embeds(model_config, malicious_encoded)
150 bypassed.append(True)
150 - except (RuntimeError, ValueError):
151 + except (RuntimeError, ValueError, VLLMValidationError):
152 rejected.append(True)
153
154 with ThreadPoolExecutor(max_workers=num_workers) as pool:
@@ -192,7 +193,7 @@ class TestGlobalFlagRestoration:
193 initial = torch.sparse.check_sparse_tensor_invariants.is_enabled()
194 malicious_encoded = _encode_tensor(_create_malicious_sparse_tensor())
195
195 - with pytest.raises((RuntimeError, ValueError)):
196 + with pytest.raises((RuntimeError, ValueError, VLLMValidationError)):
197 safe_load_prompt_embeds(model_config, malicious_encoded)
198
199 assert torch.sparse.check_sparse_tensor_invariants.is_enabled() == initial
@@ -203,7 +204,7 @@ class TestGlobalFlagRestoration:
204 malicious_encoded = _encode_tensor(_create_malicious_sparse_tensor())
205
206 def attempt_load(_):
206 - with contextlib.suppress(RuntimeError, ValueError):
207 + with contextlib.suppress(RuntimeError, ValueError, VLLMValidationError):
208 safe_load_prompt_embeds(model_config, malicious_encoded)
209
210 with ThreadPoolExecutor(max_workers=4) as pool:
@@ -244,7 +245,7 @@ class TestCrossLoaderLockSharing:
245 # serialization — one finishes before the other starts), AND invalid
246 # tensors are still rejected.
247 malicious = _encode_tensor(_create_malicious_sparse_tensor())
247 - with pytest.raises((RuntimeError, ValueError)):
248 + with pytest.raises((RuntimeError, ValueError, VLLMValidationError)):
249 safe_load_prompt_embeds(model_config, malicious)
250 with pytest.raises((RuntimeError, ValueError)):
251 io_handler.load_base64("", malicious.decode("utf-8"))
tests/renderers/test_sparse_tensor_validation.py
+8 -6
@@ -117,8 +117,10 @@ class TestPromptEmbedsValidation:
117 malicious_tensor = _create_malicious_sparse_tensor()
118 encoded = _encode_tensor(malicious_tensor)
119
120 - # Should raise RuntimeError due to invalid sparse tensor
121 - with pytest.raises((RuntimeError, ValueError)) as exc_info:
120 + # The invariant check raises RuntimeError from inside torch.load;
121 + # safe_load_prompt_embeds now reports that as a client error so the
122 + # caller gets a 400 rather than the 500 a bare RuntimeError produced.
123 + with pytest.raises(VLLMValidationError) as exc_info:
124 safe_load_prompt_embeds(model_config, encoded)
125
126 # Error should indicate sparse tensor validation failure
@@ -137,7 +139,7 @@ class TestPromptEmbedsValidation:
139 )
140 encoded = _encode_tensor(malicious_tensor)
141
140 - with pytest.raises((RuntimeError, ValueError)):
142 + with pytest.raises(VLLMValidationError):
143 safe_load_prompt_embeds(model_config, encoded)
144
145 def test_negative_indices_rejected(self, model_config):
@@ -152,7 +154,7 @@ class TestPromptEmbedsValidation:
154 )
155 encoded = _encode_tensor(malicious_tensor)
156
155 - with pytest.raises((RuntimeError, ValueError)):
157 + with pytest.raises(VLLMValidationError):
158 safe_load_prompt_embeds(model_config, encoded)
159
160 def test_hidden_size_mismatch_rejected(self, model_config):
@@ -379,8 +381,8 @@ class TestSparseTensorValidationIntegration:
381 # Step 1-2: Attacker creates malicious payload
382 attack_payload = _encode_tensor(_create_malicious_sparse_tensor())
383
382 - # Step 3-4: Server processes and should reject
383 - with pytest.raises((RuntimeError, ValueError)):
384 + # Step 3-4: Server processes and should reject, as a client error
385 + with pytest.raises(VLLMValidationError):
386 safe_load_prompt_embeds(model_config, attack_payload)
387
388 def test_attack_scenario_chat_api_image(self):
vllm/renderers/embed_utils.py
+53 -6
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from io import BytesIO
4 -from typing import TYPE_CHECKING
4 +from typing import TYPE_CHECKING, Final
5
6 import pybase64
7 import torch
@@ -17,6 +17,27 @@ if TYPE_CHECKING:
17 from vllm.config import ModelConfig
18
19
20 +# How much of torch's own reason is repeated back to the caller. The reason is
21 +# built from bytes the caller sent, so it is bounded rather than echoed whole.
22 +_MAX_EMBED_ERROR_REASON_CHARS: Final = 200
23 +
24 +
25 +def _truncated_reason(exc: Exception) -> str:
26 + """`exc`'s message, capped, saying how much was left out when it is capped.
27 +
28 + Without the count a capped reason reads like the whole one, so a caller
29 + debugging a large payload cannot tell that torch said more than this.
30 + """
31 + reason: Final = str(exc).strip() or type(exc).__name__
32 + omitted: Final = len(reason) - _MAX_EMBED_ERROR_REASON_CHARS
33 + if omitted <= 0:
34 + return reason
35 + return (
36 + f"{reason[:_MAX_EMBED_ERROR_REASON_CHARS]}"
37 + f"... ({omitted} more characters truncated)"
38 + )
39 +
40 +
41 def safe_load_prompt_embeds(
42 model_config: "ModelConfig",
43 embed: bytes,
@@ -27,12 +48,38 @@ def safe_load_prompt_embeds(
48 parameter="prompt_embeds",
49 )
50
51 + # Decoded outside the try below so that a body which is not base64 keeps
52 + # raising binascii.Error, which is a ValueError and already answered 400.
53 + payload: Final = pybase64.b64decode(embed, validate=True)
54 +
55 with check_sparse_tensor_invariants_threadsafe():
31 - tensor = torch.load(
32 - BytesIO(pybase64.b64decode(embed, validate=True)),
33 - weights_only=True,
34 - map_location=torch.device("cpu"),
35 - )
56 + try:
57 + tensor = torch.load(
58 + BytesIO(payload),
59 + weights_only=True,
60 + map_location=torch.device("cpu"),
61 + )
62 + except (torch.OutOfMemoryError, MemoryError):
63 + # Running out of memory is the server's condition, not the caller's.
64 + raise
65 + except Exception as exc:
66 + # Everything else here is about the bytes the caller sent, and torch
67 + # has no single error type for "this is not a tensor file": the zip
68 + # reader raises `RuntimeError`, the legacy pickle path raises
69 + # `UnpicklingError`, `KeyError` or `IndexError`, an empty payload
70 + # raises `EOFError`. None of them is a `ValueError`, so the
71 + # entrypoints' fallback mapped them all to 500. Catching by
72 + # behaviour rather than by type also survives torch changing which
73 + # error it raises for a given malformed payload.
74 + #
75 + # torch's reason is worth keeping -- for a malformed sparse tensor
76 + # it names the offending index -- but it is built from the caller's
77 + # own bytes, so it is truncated rather than echoed whole.
78 + raise VLLMValidationError(
79 + "`prompt_embeds` could not be deserialized as a torch tensor: "
80 + f"{_truncated_reason(exc)}",
81 + parameter="prompt_embeds",
82 + ) from exc
83 tensor = safe_to_dense(tensor, parameter="prompt_embeds")
84
85 if tensor.dim() > 2: