| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import io |
| 5 | from unittest.mock import Mock |
| 6 | |
| 7 | # imports for structured outputs tests |
| 8 | import openai |
| 9 | import pybase64 |
| 10 | import pytest |
| 11 | import regex as re |
| 12 | import torch |
| 13 | |
| 14 | from tests.utils import RemoteOpenAIServer |
| 15 | from vllm.config import ModelConfig |
| 16 | from vllm.exceptions import VLLMValidationError |
| 17 | from vllm.renderers.embed_utils import safe_load_prompt_embeds |
| 18 | |
| 19 | |
| 20 | @pytest.mark.asyncio |
| 21 | async def test_empty_prompt(): |
| 22 | model_name = "openai-community/gpt2" |
| 23 | server_args = ["--enforce-eager"] |
| 24 | with RemoteOpenAIServer(model_name, server_args) as remote_server: |
| 25 | client = remote_server.get_async_client() |
| 26 | |
| 27 | with pytest.raises( |
| 28 | openai.BadRequestError, |
| 29 | match="Either prompt or prompt_embeds must be provided and non-empty.", |
| 30 | ): |
| 31 | await client.completions.create( |
| 32 | model=model_name, |
| 33 | prompt=None, |
| 34 | max_tokens=5, |
| 35 | temperature=0.0, |
| 36 | extra_body={"prompt_embeds": []}, |
| 37 | ) |
| 38 | |
| 39 | |
| 40 | @pytest.mark.asyncio |
| 41 | async def test_out_of_vocab_token_ids(): |
| 42 | model_name = "openai-community/gpt2" |
| 43 | server_args = ["--enforce-eager"] |
| 44 | with RemoteOpenAIServer(model_name, server_args) as remote_server: |
| 45 | client = remote_server.get_async_client() |
| 46 | |
| 47 | with pytest.raises( |
| 48 | openai.BadRequestError, match=re.compile(".*out of vocabulary.*").pattern |
| 49 | ): |
| 50 | await client.completions.create( |
| 51 | model=model_name, prompt=[999999], max_tokens=5, temperature=0.0 |
| 52 | ) |
| 53 | |
| 54 | |
| 55 | @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) |
| 56 | @pytest.mark.parametrize( |
| 57 | "layout", [torch.strided, torch.sparse_coo, torch.sparse_csc, torch.sparse_csr] |
| 58 | ) |
| 59 | @pytest.mark.parametrize("seq_len", [2, 10]) |
| 60 | @pytest.mark.parametrize("hidden_size", [2, 10]) |
| 61 | def test_load_prompt_embeds( |
| 62 | dtype: torch.dtype, layout: torch.layout, seq_len: int, hidden_size: int |
| 63 | ): |
| 64 | model_config = Mock(spec=ModelConfig) |
| 65 | model_config.enable_prompt_embeds = True |
| 66 | model_config.get_hidden_size.return_value = hidden_size |
| 67 | model_config.dtype = dtype |
| 68 | |
| 69 | # construct arbitrary tensors of various dtypes, layouts, and sizes. |
| 70 | # We need to check against different layouts to make sure that if a user |
| 71 | # uses sparse tensors to reduce the transmission size of prompt embeddings, |
| 72 | # we must cast them to dense/strided before passing them into the engine. |
| 73 | # We don't use non-CPU tensors in this test to avoid preemptively |
| 74 | # initializing cuda and break other tests in the suite that fork processes. |
| 75 | # We also need to make sure that we only use devices that are actually |
| 76 | # available in the environment the test is running on. For simplicity, |
| 77 | # we just test against CPU. |
| 78 | tensor = torch.randn((seq_len, hidden_size), dtype=dtype) |
| 79 | if layout == torch.strided: |
| 80 | tensor = tensor.contiguous() |
| 81 | elif layout == torch.sparse_coo: |
| 82 | tensor = tensor.to_sparse_coo() |
| 83 | elif layout == torch.sparse_csc: |
| 84 | tensor = tensor.to_sparse_csc() |
| 85 | elif layout == torch.sparse_csr: |
| 86 | tensor = tensor.to_sparse_csr() |
| 87 | |
| 88 | buffer = io.BytesIO() |
| 89 | torch.save(tensor, buffer) |
| 90 | buffer.seek(0) |
| 91 | encoded_tensor = pybase64.b64encode(buffer.getvalue()) |
| 92 | |
| 93 | loaded_tensor = safe_load_prompt_embeds(model_config, encoded_tensor) |
| 94 | assert loaded_tensor.device.type == "cpu" |
| 95 | assert loaded_tensor.layout == torch.strided |
| 96 | torch.testing.assert_close( |
| 97 | loaded_tensor, tensor.to("cpu").to_dense(), equal_nan=True |
| 98 | ) |
| 99 | |
| 100 | |
| 101 | @pytest.mark.parametrize("dtype", [torch.float32]) |
| 102 | @pytest.mark.parametrize("seq_len", [2]) |
| 103 | @pytest.mark.parametrize("hidden_size", [2]) |
| 104 | def test_disable_prompt_embeds(dtype: torch.dtype, seq_len: int, hidden_size: int): |
| 105 | model_config = Mock(spec=ModelConfig) |
| 106 | model_config.enable_prompt_embeds = False |
| 107 | |
| 108 | tensor = torch.randn((seq_len, hidden_size), dtype=dtype) |
| 109 | |
| 110 | buffer = io.BytesIO() |
| 111 | torch.save(tensor, buffer) |
| 112 | buffer.seek(0) |
| 113 | encoded_tensor = pybase64.b64encode(buffer.getvalue()) |
| 114 | |
| 115 | with pytest.raises(VLLMValidationError, match="--enable-prompt-embeds"): |
| 116 | safe_load_prompt_embeds(model_config, encoded_tensor) |