| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import os |
| 4 | |
| 5 | import openai |
| 6 | |
| 7 | PREFILL_HOST = os.getenv("PREFILL_HOST", "localhost") |
| 8 | PREFILL_PORT = os.getenv("PREFILL_PORT", None) |
| 9 | DECODE_HOST = os.getenv("DECODE_HOST", "localhost") |
| 10 | DECODE_PORT = os.getenv("DECODE_PORT", None) |
| 11 | PROXY_HOST = os.getenv("PROXY_HOST", "localhost") |
| 12 | PROXY_PORT = os.getenv("PROXY_PORT", None) |
| 13 | |
| 14 | if PREFILL_PORT is None or DECODE_PORT is None or PROXY_PORT is None: |
| 15 | raise ValueError("Please set the PREFILL_PORT, DECODE_PORT, and PROXY_PORT.") |
| 16 | |
| 17 | LONG_PROMPT = "Red Hat is the best company in the world to work for because it works on open source software, which means that all the contributions are delivered to the community. As a result, when working on projects like vLLM we are able to meet many amazing people from various organizations like AMD, Google, NVIDIA, " # noqa: E501 |
| 18 | PROMPT = "Red Hat is the best company in the world to work for because it works on open source software, which means that all the contributions are delivered to the community. As a result," # noqa: E501 |
| 19 | SHORT_PROMPT = "Red Hat is " |
| 20 | |
| 21 | |
| 22 | def test_edge_cases(): |
| 23 | # Set the OpenAI API key and base URL |
| 24 | decode_client = openai.OpenAI( |
| 25 | api_key="MY_KEY", |
| 26 | base_url=f"http://{DECODE_HOST}:{DECODE_PORT}/v1", |
| 27 | ) |
| 28 | prefill_client = openai.OpenAI( |
| 29 | api_key="MY_KEY", |
| 30 | base_url=f"http://{PREFILL_HOST}:{PREFILL_PORT}/v1", |
| 31 | ) |
| 32 | proxy_client = openai.OpenAI( |
| 33 | api_key="MY_KEY", |
| 34 | base_url=f"http://{PROXY_HOST}:{PROXY_PORT}/v1", |
| 35 | ) |
| 36 | |
| 37 | # Get the list of models |
| 38 | models = decode_client.models.list() |
| 39 | MODEL = models.data[0].id |
| 40 | |
| 41 | # (1) Check that we can handle a very short prompt, |
| 42 | # less than the length of the block size. |
| 43 | completion = proxy_client.completions.create( |
| 44 | model=MODEL, prompt=SHORT_PROMPT, temperature=0 |
| 45 | ) |
| 46 | proxy_response = completion.choices[0].text |
| 47 | completion = prefill_client.completions.create( |
| 48 | model=MODEL, prompt=SHORT_PROMPT, temperature=0 |
| 49 | ) |
| 50 | prefill_response = completion.choices[0].text |
| 51 | print(f"SMALL PROMPT: {proxy_response=}") |
| 52 | assert proxy_response == prefill_response |
| 53 | |
| 54 | # (2) Check that we can handle a full prefix cache |
| 55 | # hit on the D worker but not on the P worker. |
| 56 | # (2a): prime the D worker. |
| 57 | completion = decode_client.completions.create( |
| 58 | model=MODEL, prompt=PROMPT, temperature=0 |
| 59 | ) |
| 60 | decode_response = completion.choices[0].text |
| 61 | # (2b): send via the P/D setup |
| 62 | completion = proxy_client.completions.create( |
| 63 | model=MODEL, prompt=PROMPT, temperature=0 |
| 64 | ) |
| 65 | proxy_response = completion.choices[0].text |
| 66 | print(f"FULL CACHE HIT: {proxy_response=}") |
| 67 | assert proxy_response == decode_response |
| 68 | |
| 69 | # (3) Check that we can handle a partial prefix cache |
| 70 | # hit on the D worker. |
| 71 | completion = proxy_client.completions.create( |
| 72 | model=MODEL, prompt=LONG_PROMPT, temperature=0 |
| 73 | ) |
| 74 | proxy_response = completion.choices[0].text |
| 75 | completion = prefill_client.completions.create( |
| 76 | model=MODEL, prompt=LONG_PROMPT, temperature=0 |
| 77 | ) |
| 78 | prefill_response = completion.choices[0].text |
| 79 | print(f"PARTIAL CACHE HIT: {proxy_response=}") |
| 80 | assert proxy_response == prefill_response |