| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import argparse |
| 4 | import json |
| 5 | import os |
| 6 | import time |
| 7 | |
| 8 | import openai |
| 9 | import requests |
| 10 | |
| 11 | MAX_OUTPUT_LEN = 30 |
| 12 | |
| 13 | SAMPLE_PROMPTS = ( |
| 14 | "Red Hat is the best company in the world to work for because it works on " |
| 15 | "open source software, which means that all the contributions are " |
| 16 | "delivered to the community. As a result, when working on projects like " |
| 17 | "vLLM we are able to meet many amazing people from various organizations " |
| 18 | "like AMD, Google, NVIDIA, ", |
| 19 | "We hold these truths to be self-evident, that all men are created equal, " |
| 20 | "that they are endowed by their Creator with certain unalienable Rights, " |
| 21 | "that among these are Life, Liberty and the pursuit of Happiness.--That " |
| 22 | "to secure these rights, Governments are instituted among Men, deriving " |
| 23 | "their just powers from the consent of the governed, ", |
| 24 | ) |
| 25 | |
| 26 | |
| 27 | def check_vllm_server(url: str, timeout=5, retries=3) -> bool: |
| 28 | """Checks if the vLLM server is ready by sending a GET request to the |
| 29 | /health endpoint. |
| 30 | |
| 31 | Args: |
| 32 | url (str): The base URL of the vLLM server. |
| 33 | timeout (int): Timeout in seconds for the request. |
| 34 | retries (int): Number of retries if the server is not ready. |
| 35 | |
| 36 | Returns: |
| 37 | bool: True if the server is ready, False otherwise. |
| 38 | |
| 39 | """ |
| 40 | for attempt in range(retries): |
| 41 | try: |
| 42 | response = requests.get(url, timeout=timeout) |
| 43 | if response.status_code == 200: |
| 44 | return True |
| 45 | else: |
| 46 | print( |
| 47 | f"Attempt {attempt + 1}: Server returned status code " |
| 48 | "{response.status_code}" |
| 49 | ) |
| 50 | except requests.exceptions.RequestException as e: |
| 51 | print(f"Attempt {attempt + 1}: Error connecting to server: {e}") |
| 52 | time.sleep(1) # Wait before retrying |
| 53 | return False |
| 54 | |
| 55 | |
| 56 | def run_simple_prompt( |
| 57 | base_url: str, model_name: str, input_prompt: str, use_chat_endpoint: bool |
| 58 | ) -> str: |
| 59 | client = openai.OpenAI(api_key="EMPTY", base_url=base_url) |
| 60 | if use_chat_endpoint: |
| 61 | completion = client.chat.completions.create( |
| 62 | model=model_name, |
| 63 | messages=[ |
| 64 | {"role": "user", "content": [{"type": "text", "text": input_prompt}]} |
| 65 | ], |
| 66 | max_completion_tokens=MAX_OUTPUT_LEN, |
| 67 | temperature=0.0, |
| 68 | seed=42, |
| 69 | ) |
| 70 | return completion.choices[0].message.content |
| 71 | else: |
| 72 | completion = client.completions.create( |
| 73 | model=model_name, |
| 74 | prompt=input_prompt, |
| 75 | max_tokens=MAX_OUTPUT_LEN, |
| 76 | temperature=0.0, |
| 77 | seed=42, |
| 78 | ) |
| 79 | |
| 80 | return completion.choices[0].text |
| 81 | |
| 82 | |
| 83 | def main(): |
| 84 | """This script demonstrates how to accept two optional string arguments |
| 85 | ("service_url" and "file_name") from the command line, each with a |
| 86 | default value of an empty string, using the argparse module. |
| 87 | """ |
| 88 | parser = argparse.ArgumentParser(description="vLLM client script") |
| 89 | |
| 90 | parser.add_argument( |
| 91 | "--service_url", # Name of the first argument |
| 92 | type=str, |
| 93 | required=True, |
| 94 | help="The vLLM service URL.", |
| 95 | ) |
| 96 | |
| 97 | parser.add_argument( |
| 98 | "--model_name", # Name of the first argument |
| 99 | type=str, |
| 100 | required=True, |
| 101 | help="model_name", |
| 102 | ) |
| 103 | |
| 104 | parser.add_argument( |
| 105 | "--mode", # Name of the second argument |
| 106 | type=str, |
| 107 | default="baseline", |
| 108 | help="mode: baseline==non-disagg, or disagg", |
| 109 | ) |
| 110 | |
| 111 | parser.add_argument( |
| 112 | "--file_name", # Name of the second argument |
| 113 | type=str, |
| 114 | default=".vllm_output.txt", |
| 115 | help="the file that saves the output tokens ", |
| 116 | ) |
| 117 | |
| 118 | args = parser.parse_args() |
| 119 | |
| 120 | for arg in vars(args): |
| 121 | print(f"{arg}: {getattr(args, arg)}") |
| 122 | |
| 123 | if args.mode == "baseline": |
| 124 | # non-disagg |
| 125 | health_check_url = f"{args.service_url}/health" |
| 126 | else: |
| 127 | # disagg proxy |
| 128 | health_check_url = f"{args.service_url}/healthcheck" |
| 129 | if not os.path.exists(args.file_name): |
| 130 | raise ValueError( |
| 131 | f"In disagg mode, the output file {args.file_name} from " |
| 132 | "non-disagg. baseline does not exist." |
| 133 | ) |
| 134 | |
| 135 | service_url = f"{args.service_url}/v1" |
| 136 | |
| 137 | if not check_vllm_server(health_check_url): |
| 138 | raise RuntimeError(f"vllm server: {args.service_url} is not ready yet!") |
| 139 | |
| 140 | output_strs = dict() |
| 141 | for i, prompt in enumerate(SAMPLE_PROMPTS): |
| 142 | use_chat_endpoint = i % 2 == 1 |
| 143 | output_str = run_simple_prompt( |
| 144 | base_url=service_url, |
| 145 | model_name=args.model_name, |
| 146 | input_prompt=prompt, |
| 147 | use_chat_endpoint=use_chat_endpoint, |
| 148 | ) |
| 149 | print(f"Prompt: {prompt}, output: {output_str}") |
| 150 | output_strs[prompt] = output_str |
| 151 | |
| 152 | if args.mode == "baseline": |
| 153 | # baseline: save outputs |
| 154 | try: |
| 155 | with open(args.file_name, "w") as json_file: |
| 156 | json.dump(output_strs, json_file, indent=4) |
| 157 | except OSError as e: |
| 158 | print(f"Error writing to file: {e}") |
| 159 | raise |
| 160 | else: |
| 161 | # disagg. verify outputs |
| 162 | baseline_outputs = None |
| 163 | try: |
| 164 | with open(args.file_name) as json_file: |
| 165 | baseline_outputs = json.load(json_file) |
| 166 | except OSError as e: |
| 167 | print(f"Error writing to file: {e}") |
| 168 | raise |
| 169 | assert isinstance(baseline_outputs, dict) |
| 170 | assert len(baseline_outputs) == len(output_strs) |
| 171 | for prompt, output in baseline_outputs.items(): |
| 172 | assert prompt in output_strs, f"{prompt} not included" |
| 173 | assert output == output_strs[prompt], ( |
| 174 | f"baseline_output: {output} != PD output: {output_strs[prompt]}" |
| 175 | ) |
| 176 | |
| 177 | |
| 178 | if __name__ == "__main__": |
| 179 | main() |