@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 163 lines 5.8 KB
Raw
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