wehub-resource-sync
1061 lines
38 KiB
Python
1061 lines
38 KiB
Python
"""Coppied from https://github.com/fw-ai/benchmark/blob/main/llm_bench/load_test.py"""
|
|
|
|
import abc
|
|
import argparse
|
|
import csv
|
|
from dataclasses import dataclass
|
|
from functools import partial
|
|
import os
|
|
import random
|
|
import sys
|
|
import traceback
|
|
from typing import Optional
|
|
from locust import HttpUser, task, events, constant_pacing
|
|
import copy
|
|
import json
|
|
import time
|
|
import orjson
|
|
import threading
|
|
|
|
|
|
def add_custom_metric(name, value, length_value=0):
|
|
events.request.fire(
|
|
request_type="METRIC",
|
|
name=name,
|
|
response_time=value,
|
|
response_length=length_value,
|
|
exception=None,
|
|
context=None,
|
|
)
|
|
|
|
|
|
PROMPT_PREFIX_TOKEN = "Pad " # exactly one token
|
|
# "Lengthy" prompt borrowed from nat.dev
|
|
PROMPT_SUFFIX = """Generate a Django application with Authentication, JWT, Tests, DB support. Show docker-compose for python and postgres. Show the complete code for every file!"""
|
|
PROMPT_SUFFIX_TOKENS = 35 # from Llama tokenizer tool (so we don't import it here)
|
|
|
|
|
|
class FixedQPSPacer:
|
|
_instance = None
|
|
_lock = threading.Lock()
|
|
|
|
def __init__(self, qps, distribution):
|
|
self.qps = qps
|
|
self.distribution = distribution
|
|
|
|
# It's kind of thread safe thanks to GIL as the only state is `t` - good enough for a loadtest
|
|
def gen():
|
|
t = time.time()
|
|
mean_wait = 1 / self.qps
|
|
while True:
|
|
if self.distribution == "exponential":
|
|
wait = random.expovariate(1 / mean_wait)
|
|
elif self.distribution == "uniform":
|
|
wait = random.uniform(0, 2 * mean_wait)
|
|
elif self.distribution == "constant":
|
|
wait = mean_wait
|
|
else:
|
|
print("Unknown distribution {self.distribution}")
|
|
os._exit(1)
|
|
t += wait
|
|
yield t
|
|
|
|
self.iterator = gen()
|
|
|
|
@classmethod
|
|
def instance(cls, qps, distribution):
|
|
with cls._lock:
|
|
if cls._instance is None:
|
|
cls._instance = cls(qps, distribution)
|
|
else:
|
|
assert cls._instance.qps == qps
|
|
assert cls._instance.distribution == distribution
|
|
return cls._instance
|
|
|
|
def wait_time_till_next(self):
|
|
with self._lock:
|
|
t = next(self.iterator)
|
|
now = time.time()
|
|
if now > t:
|
|
print(
|
|
f"WARNING: not enough locust users to keep up with the desired QPS. Either the number of locust users is too low or the server is overloaded. Delay: {now-t:.3f}s"
|
|
)
|
|
return 0
|
|
return t - now
|
|
|
|
|
|
class LengthSampler:
|
|
def __init__(self, distribution: str, mean: int, cap: Optional[int], alpha: float):
|
|
self.distribution = distribution
|
|
self.mean = mean
|
|
self.cap = cap
|
|
self.alpha = alpha
|
|
|
|
if self.distribution == "exponential":
|
|
self.sample_func = lambda: int(random.expovariate(1 / self.mean))
|
|
elif self.distribution == "uniform":
|
|
mx = self.mean + int(self.alpha * self.mean)
|
|
if self.cap is not None:
|
|
mx = min(mx, self.cap)
|
|
self.sample_func = lambda: random.randint(
|
|
max(1, self.mean - int(self.alpha * self.mean)), mx
|
|
)
|
|
elif self.distribution == "constant":
|
|
self.sample_func = lambda: self.mean
|
|
elif self.distribution == "normal":
|
|
self.sample_func = lambda: int(
|
|
random.gauss(self.mean, self.mean * self.alpha)
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown distribution {self.distribution}")
|
|
|
|
def sample(self) -> int:
|
|
for _ in range(1000):
|
|
sample = self.sample_func()
|
|
if sample <= 0:
|
|
continue
|
|
if self.cap is not None and sample > self.cap:
|
|
continue
|
|
return sample
|
|
else:
|
|
raise ValueError(
|
|
"Can't sample a value after 1000 attempts, check distribution parameters"
|
|
)
|
|
|
|
def __str__(self):
|
|
r = int(self.mean * self.alpha)
|
|
if self.distribution == "constant":
|
|
s = str(self.mean)
|
|
elif self.distribution == "uniform":
|
|
s = f"uniform({self.mean} +/- {r})"
|
|
elif self.distribution == "normal":
|
|
s = f"normal({self.mean}, {r})"
|
|
elif self.distribution == "exponential":
|
|
s = f"exponential({self.mean})"
|
|
else:
|
|
assert False
|
|
if self.cap is not None:
|
|
s += f" capped at {self.cap}"
|
|
return s
|
|
|
|
|
|
class InitTracker:
|
|
lock = threading.Lock()
|
|
users = None
|
|
first_request_done = 0
|
|
logging_params = None
|
|
environment = None
|
|
tokenizer = None
|
|
|
|
@classmethod
|
|
def notify_init(cls, environment, logging_params):
|
|
with cls.lock:
|
|
if cls.environment is None:
|
|
cls.environment = environment
|
|
if cls.logging_params is None:
|
|
cls.logging_params = logging_params
|
|
else:
|
|
assert (
|
|
cls.logging_params == logging_params
|
|
), f"Inconsistent settings between workers: {cls.logging_params} != {logging_params}"
|
|
|
|
@classmethod
|
|
def notify_first_request(cls):
|
|
with cls.lock:
|
|
if (
|
|
cls.environment.parsed_options.qps is not None
|
|
and cls.first_request_done == 0
|
|
):
|
|
# if in QPS mode, reset after first successful request comes back
|
|
cls.reset_stats()
|
|
cls.first_request_done += 1
|
|
if (
|
|
cls.environment.parsed_options.qps is not None
|
|
and cls.first_request_done == 0
|
|
and cls.users == cls.first_request_done
|
|
):
|
|
# if in fixed load mode, reset after all users issued one request (we're in a steady state)
|
|
cls.reset_stats()
|
|
|
|
@classmethod
|
|
def notify_spawning_complete(cls, user_count):
|
|
with cls.lock:
|
|
cls.users = user_count
|
|
if cls.users == cls.first_request_done:
|
|
cls.reset_stats()
|
|
|
|
@classmethod
|
|
def reset_stats(cls):
|
|
assert cls.environment.runner, "only local mode is supported"
|
|
print("Resetting stats after traffic reach a steady state")
|
|
cls.environment.events.reset_stats.fire()
|
|
cls.environment.runner.stats.reset_all()
|
|
|
|
@classmethod
|
|
def load_tokenizer(cls, dir):
|
|
if not dir:
|
|
return None
|
|
with cls.lock:
|
|
if cls.tokenizer:
|
|
return cls.tokenizer
|
|
import transformers
|
|
|
|
cls.tokenizer = transformers.AutoTokenizer.from_pretrained(dir)
|
|
cls.tokenizer.add_bos_token = False
|
|
cls.tokenizer.add_eos_token = False
|
|
return cls.tokenizer
|
|
|
|
|
|
events.spawning_complete.add_listener(InitTracker.notify_spawning_complete)
|
|
|
|
|
|
@dataclass
|
|
class ChunkMetadata:
|
|
text: str
|
|
logprob_tokens: Optional[int]
|
|
usage_tokens: Optional[int]
|
|
prompt_usage_tokens: Optional[int]
|
|
|
|
|
|
class BaseProvider(abc.ABC):
|
|
DEFAULT_MODEL_NAME = None
|
|
|
|
def __init__(self, model, parsed_options):
|
|
self.model = model
|
|
self.parsed_options = parsed_options
|
|
|
|
@abc.abstractmethod
|
|
def get_url(self):
|
|
...
|
|
|
|
@abc.abstractmethod
|
|
def format_payload(self, prompt, max_tokens, images):
|
|
...
|
|
|
|
@abc.abstractmethod
|
|
def parse_output_json(self, json, prompt):
|
|
...
|
|
|
|
|
|
class OpenAIProvider(BaseProvider):
|
|
def get_url(self):
|
|
if self.parsed_options.chat:
|
|
return "/v1/chat/completions"
|
|
else:
|
|
return "/v1/completions"
|
|
|
|
def format_payload(self, prompt, max_tokens, images):
|
|
data = {
|
|
"model": self.model,
|
|
"max_tokens": max_tokens,
|
|
"stream": self.parsed_options.stream,
|
|
"temperature": self.parsed_options.temperature,
|
|
"n": self.parsed_options.n,
|
|
}
|
|
if self.parsed_options.chat:
|
|
if images is None:
|
|
data["messages"] = [{"role": "user", "content": prompt}]
|
|
else:
|
|
image_urls = []
|
|
for image in images:
|
|
image_urls.append(
|
|
{"type": "image_url", "image_url": {"url": image}}
|
|
)
|
|
data["messages"] = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": prompt},
|
|
*image_urls,
|
|
],
|
|
}
|
|
]
|
|
else:
|
|
data["prompt"] = prompt
|
|
if images is not None:
|
|
data["images"] = images
|
|
if self.parsed_options.logprobs is not None:
|
|
data["logprobs"] = self.parsed_options.logprobs
|
|
return data
|
|
|
|
def parse_output_json(self, data, prompt):
|
|
usage = data.get("usage", None)
|
|
|
|
assert len(data["choices"]) == 1, f"Too many choices {len(data['choices'])}"
|
|
choice = data["choices"][0]
|
|
if self.parsed_options.chat:
|
|
if self.parsed_options.stream:
|
|
text = choice["delta"].get("content", "")
|
|
else:
|
|
text = choice["message"]["content"]
|
|
else:
|
|
text = choice["text"]
|
|
|
|
logprobs = (choice.get("logprobs", {}) or {}).get("content", [])
|
|
return ChunkMetadata(
|
|
text=text,
|
|
logprob_tokens=len(logprobs["tokens"]) if logprobs else None,
|
|
usage_tokens=usage["completion_tokens"] if usage else None,
|
|
prompt_usage_tokens=(usage.get("prompt_tokens", None) if usage else None),
|
|
)
|
|
|
|
|
|
class FireworksProvider(OpenAIProvider):
|
|
def format_payload(self, prompt, max_tokens, images):
|
|
data = super().format_payload(prompt, max_tokens, images)
|
|
data["min_tokens"] = max_tokens
|
|
data["prompt_cache_max_len"] = self.parsed_options.prompt_cache_max_len
|
|
return data
|
|
|
|
|
|
class VllmProvider(OpenAIProvider):
|
|
def format_payload(self, prompt, max_tokens, images):
|
|
data = super().format_payload(prompt, max_tokens, images)
|
|
data["ignore_eos"] = True
|
|
return data
|
|
|
|
|
|
class TogetherProvider(OpenAIProvider):
|
|
def get_url(self):
|
|
assert not self.parsed_options.chat, "Chat is not supported"
|
|
return "/"
|
|
|
|
def format_payload(self, prompt, max_tokens, images):
|
|
data = super().format_payload(prompt, max_tokens, images)
|
|
data["ignore_eos"] = True
|
|
data["stream_tokens"] = data.pop("stream")
|
|
return data
|
|
|
|
def parse_output_json(self, data, prompt):
|
|
if not self.parsed_options.stream:
|
|
data = data["output"]
|
|
return super().parse_output_json(data, prompt)
|
|
|
|
|
|
class TritonInferProvider(BaseProvider):
|
|
DEFAULT_MODEL_NAME = "ensemble"
|
|
|
|
def get_url(self):
|
|
assert not self.parsed_options.chat, "Chat is not supported"
|
|
assert not self.parsed_options.stream, "Stream is not supported"
|
|
assert self.parsed_options.n == 1, "n > 1 is not supported"
|
|
return f"/v2/models/{self.model}/infer"
|
|
|
|
def format_payload(self, prompt, max_tokens, images):
|
|
assert images is None, "images are not supported"
|
|
# matching latest TRT-LLM example, your model configuration might be different
|
|
data = {
|
|
"inputs": [
|
|
{
|
|
"name": "text_input",
|
|
"datatype": "BYTES",
|
|
"shape": [1, 1],
|
|
"data": [[prompt]],
|
|
},
|
|
{
|
|
"name": "max_tokens",
|
|
"datatype": "UINT32",
|
|
"shape": [1, 1],
|
|
"data": [[max_tokens]],
|
|
},
|
|
{
|
|
"name": "bad_words",
|
|
"datatype": "BYTES",
|
|
"shape": [1, 1],
|
|
"data": [[""]],
|
|
},
|
|
{
|
|
"name": "stop_words",
|
|
"datatype": "BYTES",
|
|
"shape": [1, 1],
|
|
"data": [[""]],
|
|
},
|
|
{
|
|
"name": "temperature",
|
|
"datatype": "FP32",
|
|
"shape": [1, 1],
|
|
"data": [[self.parsed_options.temperature]],
|
|
},
|
|
]
|
|
}
|
|
assert self.parsed_options.logprobs is None, "logprobs are not supported"
|
|
return data
|
|
|
|
def parse_output_json(self, data, prompt):
|
|
for output in data["outputs"]:
|
|
if output["name"] == "text_output":
|
|
assert output["datatype"] == "BYTES"
|
|
assert output["shape"] == [1]
|
|
text = output["data"][0]
|
|
# Triton returns the original prompt in the output, cut it off
|
|
text = text.removeprefix("<s> ")
|
|
if text.startswith(prompt):
|
|
# HF tokenizers get confused by the leading space
|
|
text = text[len(prompt) :].removeprefix(" ")
|
|
else:
|
|
print("WARNING: prompt not found in the output")
|
|
return ChunkMetadata(
|
|
text=text,
|
|
logprob_tokens=None,
|
|
usage_tokens=None,
|
|
prompt_usage_tokens=None,
|
|
)
|
|
raise ValueError("text_output not found in the response")
|
|
|
|
|
|
class TritonGenerateProvider(BaseProvider):
|
|
DEFAULT_MODEL_NAME = "ensemble"
|
|
|
|
def get_url(self):
|
|
assert not self.parsed_options.chat, "Chat is not supported"
|
|
stream_suffix = "_stream" if self.parsed_options.stream else ""
|
|
return f"/v2/models/{self.model}/generate{stream_suffix}"
|
|
|
|
def format_payload(self, prompt, max_tokens, images):
|
|
assert images is None, "images are not supported"
|
|
assert self.parsed_options.n == 1, "n > 1 is not supported"
|
|
data = {
|
|
"text_input": prompt,
|
|
"max_tokens": max_tokens,
|
|
"stream": self.parsed_options.stream,
|
|
"temperature": self.parsed_options.temperature,
|
|
# for whatever reason these has to be provided
|
|
"bad_words": "",
|
|
"stop_words": "",
|
|
}
|
|
assert self.parsed_options.logprobs is None, "logprobs are not supported"
|
|
return data
|
|
|
|
def parse_output_json(self, data, prompt):
|
|
text = data["text_output"]
|
|
if not self.parsed_options.stream:
|
|
# Triton returns the original prompt in the output, cut it off
|
|
text = text.removeprefix("<s> ")
|
|
if text.startswith(prompt):
|
|
# HF tokenizers get confused by the leading space
|
|
text = text[len(prompt) :].removeprefix(" ")
|
|
else:
|
|
print("WARNING: prompt not found in the output")
|
|
return ChunkMetadata(
|
|
text=text,
|
|
logprob_tokens=None,
|
|
usage_tokens=None,
|
|
prompt_usage_tokens=None,
|
|
)
|
|
|
|
|
|
class TgiProvider(BaseProvider):
|
|
DEFAULT_MODEL_NAME = "<unused>"
|
|
|
|
def get_url(self):
|
|
assert self.parsed_options.n == 1, "n > 1 is not supported"
|
|
assert not self.parsed_options.chat, "Chat is not supported"
|
|
stream_suffix = "_stream" if self.parsed_options.stream else ""
|
|
return f"/generate{stream_suffix}"
|
|
|
|
def format_payload(self, prompt, max_tokens, images):
|
|
assert images is None, "images are not supported"
|
|
data = {
|
|
"inputs": prompt,
|
|
"parameters": {
|
|
"max_new_tokens": max_tokens,
|
|
"temperature": self.parsed_options.temperature,
|
|
"top_n_tokens": self.parsed_options.logprobs,
|
|
"details": self.parsed_options.logprobs is not None,
|
|
},
|
|
}
|
|
return data
|
|
|
|
def parse_output_json(self, data, prompt):
|
|
if "token" in data:
|
|
# streaming chunk
|
|
return ChunkMetadata(
|
|
text=data["token"]["text"],
|
|
logprob_tokens=1,
|
|
usage_tokens=None,
|
|
prompt_usage_tokens=None,
|
|
)
|
|
else:
|
|
# non-streaming response
|
|
return ChunkMetadata(
|
|
text=data["generated_text"],
|
|
logprob_tokens=(
|
|
len(data["details"]["tokens"]) if "details" in data else None
|
|
),
|
|
usage_tokens=(
|
|
data["details"]["generated_tokens"] if "details" in data else None
|
|
),
|
|
prompt_usage_tokens=None,
|
|
)
|
|
|
|
|
|
PROVIDER_CLASS_MAP = {
|
|
"fireworks": FireworksProvider,
|
|
"vllm": VllmProvider,
|
|
"sglang": VllmProvider,
|
|
"openai": OpenAIProvider,
|
|
"anyscale": OpenAIProvider,
|
|
"together": TogetherProvider,
|
|
"triton-infer": TritonInferProvider,
|
|
"triton-generate": TritonGenerateProvider,
|
|
"tgi": TgiProvider,
|
|
}
|
|
|
|
|
|
def _load_curl_like_data(text):
|
|
"""
|
|
Either use the passed string or load from a file if the string is `@filename`
|
|
"""
|
|
if text.startswith("@"):
|
|
try:
|
|
if text.endswith(".jsonl"):
|
|
with open(text[1:], "r") as f:
|
|
return [json.loads(line) for line in f]
|
|
else:
|
|
with open(text[1:], "r") as f:
|
|
return f.read()
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to read file {text[1:]}") from e
|
|
else:
|
|
return text
|
|
|
|
|
|
class LLMUser(HttpUser):
|
|
# no wait time, so every user creates a continuous load, sending requests as quickly as possible
|
|
|
|
def on_start(self):
|
|
try:
|
|
self._on_start()
|
|
except Exception as e:
|
|
print(f"Failed to initialize: {repr(e)}")
|
|
print(traceback.format_exc())
|
|
sys.exit(1)
|
|
|
|
def _guess_provider(self):
|
|
self.model = self.environment.parsed_options.model
|
|
self.provider = self.environment.parsed_options.provider
|
|
# guess based on URL
|
|
if self.provider is None:
|
|
if "fireworks.ai" in self.host:
|
|
self.provider = "fireworks"
|
|
elif "together" in self.host:
|
|
self.provider = "together"
|
|
elif "openai" in self.host:
|
|
self.provider = "openai"
|
|
elif "anyscale" in self.host:
|
|
self.provider = "anyscale"
|
|
|
|
if (
|
|
self.model is None
|
|
and self.provider is not None
|
|
and PROVIDER_CLASS_MAP[self.provider].DEFAULT_MODEL_NAME is not None
|
|
):
|
|
self.model = PROVIDER_CLASS_MAP[self.provider].DEFAULT_MODEL_NAME
|
|
|
|
if self.model and self.provider:
|
|
return
|
|
|
|
# vllm doesn't support /model/<name> endpoint, so iterate over all models
|
|
try:
|
|
resp = self.client.get("/v1/models")
|
|
resp.raise_for_status()
|
|
resp = resp.json()
|
|
except Exception as e:
|
|
raise ValueError(
|
|
"Argument --model or --provider was not specified and /v1/models failed"
|
|
) from e
|
|
|
|
models = resp["data"]
|
|
assert len(models) > 0, "No models found in /v1/models"
|
|
owned_by = None
|
|
# pick the first model
|
|
for m in models:
|
|
if self.model is None or m["id"] == self.model:
|
|
self.model = m["id"]
|
|
owned_by = m["owned_by"]
|
|
break
|
|
if self.provider is None:
|
|
if not owned_by:
|
|
raise ValueError(
|
|
f"Model {self.model} not found in /v1/models. Specify --provider explicitly"
|
|
)
|
|
if owned_by in PROVIDER_CLASS_MAP:
|
|
self.provider = owned_by
|
|
else:
|
|
raise ValueError(
|
|
f"Can't detect provider, specify it explicitly with --provider, owned_by={owned_by}"
|
|
)
|
|
|
|
def _on_start(self):
|
|
self.client.headers["Content-Type"] = "application/json"
|
|
if self.environment.parsed_options.api_key:
|
|
self.client.headers["Authorization"] = (
|
|
"Bearer " + self.environment.parsed_options.api_key
|
|
)
|
|
if self.environment.parsed_options.header:
|
|
for header in self.environment.parsed_options.header:
|
|
key, val = header.split(":", 1)
|
|
self.client.headers[key] = val
|
|
self._guess_provider()
|
|
print(f" Provider {self.provider} using model {self.model} ".center(80, "*"))
|
|
self.provider_formatter = PROVIDER_CLASS_MAP[self.provider](
|
|
self.model, self.environment.parsed_options
|
|
)
|
|
|
|
self.stream = self.environment.parsed_options.stream
|
|
prompt_chars = self.environment.parsed_options.prompt_chars
|
|
if self.environment.parsed_options.prompt_text:
|
|
self.input = _load_curl_like_data(
|
|
self.environment.parsed_options.prompt_text
|
|
)
|
|
elif prompt_chars:
|
|
self.input = (
|
|
PROMPT_PREFIX_TOKEN * (prompt_chars // len(PROMPT_PREFIX_TOKEN) + 1)
|
|
+ PROMPT_SUFFIX
|
|
)[:prompt_chars]
|
|
else:
|
|
assert (
|
|
self.environment.parsed_options.prompt_tokens >= PROMPT_SUFFIX_TOKENS
|
|
), f"Minimal prompt length is {PROMPT_SUFFIX_TOKENS}"
|
|
self.input = (
|
|
PROMPT_PREFIX_TOKEN
|
|
* (self.environment.parsed_options.prompt_tokens - PROMPT_SUFFIX_TOKENS)
|
|
+ PROMPT_SUFFIX
|
|
)
|
|
self.max_tokens_sampler = LengthSampler(
|
|
distribution=self.environment.parsed_options.max_tokens_distribution,
|
|
mean=self.environment.parsed_options.max_tokens,
|
|
cap=self.environment.parsed_options.max_tokens_cap,
|
|
alpha=self.environment.parsed_options.max_tokens_range,
|
|
)
|
|
self.temperature = self.environment.parsed_options.temperature
|
|
|
|
logging_params = {
|
|
# TODO: add some server info with git version
|
|
"provider": self.provider,
|
|
"model": self.model,
|
|
"prompt_tokens": self.environment.parsed_options.prompt_tokens, # might be overwritten based on metric
|
|
"generation_tokens": str(self.max_tokens_sampler),
|
|
"stream": self.stream,
|
|
"temperature": self.temperature,
|
|
"logprobs": self.environment.parsed_options.logprobs,
|
|
}
|
|
InitTracker.notify_init(self.environment, logging_params)
|
|
|
|
self.tokenizer = InitTracker.load_tokenizer(
|
|
self.environment.parsed_options.tokenizer
|
|
)
|
|
if self.tokenizer:
|
|
self.prompt_tokenizer_tokens = len(
|
|
self.tokenizer.encode(self._get_input()[0])
|
|
)
|
|
else:
|
|
self.prompt_tokenizer_tokens = None
|
|
|
|
if self.environment.parsed_options.qps is not None:
|
|
if self.environment.parsed_options.burst:
|
|
raise ValueError("Burst and QPS modes are mutually exclusive")
|
|
pacer = FixedQPSPacer.instance(
|
|
self.environment.parsed_options.qps,
|
|
self.environment.parsed_options.qps_distribution,
|
|
)
|
|
# it will be called by Locust after each task
|
|
self.wait_time = pacer.wait_time_till_next
|
|
self.wait()
|
|
elif self.environment.parsed_options.burst:
|
|
self.wait_time = partial(
|
|
constant_pacing(self.environment.parsed_options.burst), self
|
|
)
|
|
else:
|
|
# introduce initial delay to avoid all users hitting the service at the same time
|
|
time.sleep(random.random())
|
|
|
|
self.first_done = False
|
|
|
|
def _get_input(self):
|
|
def _maybe_randomize(prompt):
|
|
if not self.environment.parsed_options.prompt_randomize:
|
|
return prompt
|
|
|
|
# single letters are single tokens
|
|
num_random_tokens = (len(prompt) - len(PROMPT_SUFFIX)) // len(
|
|
PROMPT_PREFIX_TOKEN
|
|
)
|
|
return (
|
|
" ".join(
|
|
chr(ord("a") + random.randint(0, 25))
|
|
for _ in range(num_random_tokens)
|
|
)
|
|
+ " "
|
|
+ prompt[-len(PROMPT_SUFFIX) :]
|
|
)
|
|
|
|
if isinstance(self.input, str):
|
|
return _maybe_randomize(self.input), None
|
|
else:
|
|
item = self.input[random.randint(0, len(self.input) - 1)]
|
|
assert "prompt" in item
|
|
return _maybe_randomize(item["prompt"]), item.get("images", None)
|
|
|
|
@task
|
|
def generate_text(self):
|
|
max_tokens = self.max_tokens_sampler.sample()
|
|
prompt, images = self._get_input()
|
|
data = self.provider_formatter.format_payload(prompt, max_tokens, images)
|
|
t_start = time.perf_counter()
|
|
|
|
with self.client.post(
|
|
self.provider_formatter.get_url(),
|
|
data=json.dumps(data),
|
|
stream=True,
|
|
catch_response=True,
|
|
) as response:
|
|
combined_text = ""
|
|
done = False
|
|
prompt_usage_tokens = self.prompt_tokenizer_tokens
|
|
total_usage_tokens = None
|
|
total_logprob_tokens = None
|
|
try:
|
|
response.raise_for_status()
|
|
except Exception as e:
|
|
raise RuntimeError(f"Error in response: {response.text}") from e
|
|
t_first_token = None
|
|
for chunk in response.iter_lines(delimiter=b"\n\n"):
|
|
if len(chunk) == 0:
|
|
continue # come providers send empty lines between data chunks
|
|
if done:
|
|
if chunk != b"data: [DONE]":
|
|
print(f"WARNING: Received more chunks after [DONE]: {chunk}")
|
|
try:
|
|
now = time.perf_counter()
|
|
if self.stream:
|
|
assert chunk.startswith(
|
|
b"data:"
|
|
), f"Unexpected chunk not starting with 'data': {chunk}"
|
|
chunk = chunk[len(b"data:") :]
|
|
if chunk.strip() == b"[DONE]":
|
|
done = True
|
|
continue
|
|
data = orjson.loads(chunk)
|
|
out = self.provider_formatter.parse_output_json(data, prompt)
|
|
if out.usage_tokens:
|
|
total_usage_tokens = (
|
|
total_usage_tokens or 0
|
|
) + out.usage_tokens
|
|
if out.prompt_usage_tokens:
|
|
prompt_usage_tokens = out.prompt_usage_tokens
|
|
combined_text += out.text
|
|
|
|
# some providers (SGLang) send an empty chunk first skewing the TTFT
|
|
if combined_text and t_first_token is None:
|
|
t_first_token = now
|
|
|
|
if out.logprob_tokens:
|
|
total_logprob_tokens = (
|
|
total_logprob_tokens or 0
|
|
) + out.logprob_tokens
|
|
except Exception as e:
|
|
print(f"Failed to parse response: {chunk} with error {repr(e)}")
|
|
response.failure(e)
|
|
return
|
|
assert t_first_token is not None, "empty response received"
|
|
if (
|
|
(total_logprob_tokens is not None)
|
|
and (total_usage_tokens is not None)
|
|
and total_logprob_tokens != total_usage_tokens
|
|
):
|
|
print(
|
|
f"WARNING: usage_tokens {total_usage_tokens} != logprob_tokens {total_logprob_tokens}"
|
|
)
|
|
if total_logprob_tokens is not None:
|
|
num_tokens = total_logprob_tokens
|
|
else:
|
|
num_tokens = total_usage_tokens
|
|
if self.tokenizer:
|
|
num_tokenizer_tokens = len(self.tokenizer.encode(combined_text))
|
|
if num_tokens is None:
|
|
num_tokens = num_tokenizer_tokens
|
|
elif num_tokens != num_tokenizer_tokens:
|
|
print(
|
|
f"WARNING: tokenizer token count {num_tokenizer_tokens} != {num_tokens} received from server"
|
|
)
|
|
num_tokens = num_tokens or 0
|
|
num_chars = len(combined_text)
|
|
now = time.perf_counter()
|
|
dur_total = now - t_start
|
|
dur_generation = now - t_first_token
|
|
dur_first_token = t_first_token - t_start
|
|
print(
|
|
f"Response received: total {dur_total*1000:.2f} ms, first token {dur_first_token*1000:.2f} ms, {num_chars} chars, {num_tokens} tokens"
|
|
)
|
|
if self.environment.parsed_options.show_response:
|
|
print("---")
|
|
print(combined_text)
|
|
print("---")
|
|
if num_chars:
|
|
add_custom_metric(
|
|
"latency_per_char",
|
|
dur_generation / num_chars * 1000,
|
|
num_chars,
|
|
)
|
|
if self.stream:
|
|
add_custom_metric("time_to_first_token", dur_first_token * 1000)
|
|
add_custom_metric("total_latency", dur_total * 1000)
|
|
if num_tokens:
|
|
if num_tokens != max_tokens:
|
|
print(
|
|
f"WARNING: wrong number of tokens: {num_tokens}, expected {max_tokens}"
|
|
)
|
|
add_custom_metric("num_tokens", num_tokens)
|
|
add_custom_metric(
|
|
"latency_per_token",
|
|
dur_generation / num_tokens * 1000,
|
|
num_tokens,
|
|
)
|
|
add_custom_metric(
|
|
"overall_latency_per_token",
|
|
dur_total / num_tokens * 1000,
|
|
num_tokens,
|
|
)
|
|
if (
|
|
prompt_usage_tokens is not None
|
|
and self.prompt_tokenizer_tokens is not None
|
|
and prompt_usage_tokens != self.prompt_tokenizer_tokens
|
|
):
|
|
print(
|
|
f"WARNING: prompt usage tokens {prompt_usage_tokens} != {self.prompt_tokenizer_tokens} derived from local tokenizer"
|
|
)
|
|
prompt_tokens = prompt_usage_tokens or self.prompt_tokenizer_tokens
|
|
if prompt_tokens:
|
|
add_custom_metric("prompt_tokens", prompt_tokens)
|
|
|
|
if not self.first_done:
|
|
self.first_done = True
|
|
InitTracker.notify_first_request()
|
|
|
|
|
|
@events.init_command_line_parser.add_listener
|
|
def init_parser(parser):
|
|
parser.add_argument(
|
|
"--provider",
|
|
choices=list(PROVIDER_CLASS_MAP.keys()),
|
|
type=str,
|
|
help="Which flavor of API to use. If not specified, we'll try to guess based on the URL and /v1/models output",
|
|
)
|
|
parser.add_argument(
|
|
"-m",
|
|
"--model",
|
|
env_var="MODEL",
|
|
type=str,
|
|
help="The model to use for generating text. If not specified we will pick the first model from the service as returned by /v1/models",
|
|
)
|
|
parser.add_argument(
|
|
"--chat",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=False,
|
|
help="Use /v1/chat/completions API",
|
|
)
|
|
parser.add_argument(
|
|
"-p",
|
|
"--prompt-tokens",
|
|
env_var="PROMPT_TOKENS",
|
|
type=int,
|
|
default=512,
|
|
help="Length of the prompt in tokens. Default 512",
|
|
)
|
|
parser.add_argument(
|
|
"--prompt-chars",
|
|
env_var="PROMPT_CHARS",
|
|
type=int,
|
|
help="Length of the prompt in characters.",
|
|
)
|
|
parser.add_argument(
|
|
"--prompt-text",
|
|
env_var="PROMPT_TEXT",
|
|
type=str,
|
|
help="Prompt text to use instead of generating one. It can be a file reference starting with an ampersand, e.g. `@prompt.txt`",
|
|
)
|
|
parser.add_argument(
|
|
"--prompt-randomize",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=False,
|
|
help="Include a few random numbers in the generated prompt to avoid caching",
|
|
)
|
|
parser.add_argument(
|
|
"-o",
|
|
"--max-tokens",
|
|
env_var="MAX_TOKENS",
|
|
type=int,
|
|
default=64,
|
|
help="Max number of tokens to generate. If --max-tokens-distribution is non-constant this is going to be the mean. Defaults to 64",
|
|
)
|
|
parser.add_argument(
|
|
"--max-tokens-cap",
|
|
env_var="MAX_TOKENS_CAP",
|
|
type=int,
|
|
help="If --max-tokens-distribution is non-constant, this truncates the distribition at the specified limit",
|
|
)
|
|
parser.add_argument(
|
|
"--max-tokens-distribution",
|
|
env_var="MAX_TOKENS_DISTRIBUTION",
|
|
type=str,
|
|
choices=["constant", "uniform", "exponential", "normal"],
|
|
default="constant",
|
|
help="How to sample `max-tokens` on each request",
|
|
)
|
|
parser.add_argument(
|
|
"--max-tokens-range",
|
|
env_var="MAX_TOKENS_RANGE",
|
|
type=float,
|
|
default=0.3,
|
|
help="Specifies the width of the distribution. Specified value `alpha` is relative to `max-tokens`. For uniform distribution we'd sample from [max_tokens - max_tokens * alpha, max_tokens + max_tokens * alpha]. For normal distribution we'd sample from `N(max_tokens, max_tokens * alpha)`. Defaults to 0.3",
|
|
)
|
|
parser.add_argument(
|
|
"--stream",
|
|
dest="stream",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=True,
|
|
help="Use the streaming API",
|
|
)
|
|
parser.add_argument(
|
|
"-k",
|
|
"--api-key",
|
|
env_var="API_KEY",
|
|
help="Auth for the API",
|
|
)
|
|
parser.add_argument(
|
|
"--temperature",
|
|
env_var="TEMPERATURE",
|
|
type=float,
|
|
default=1.0,
|
|
help="Temperature parameter for the API",
|
|
)
|
|
parser.add_argument(
|
|
"--logprobs",
|
|
type=int,
|
|
default=None,
|
|
help="Whether to ask for logprobs, it makes things slower for some providers but is necessary for token count in streaming (unless it's Fireworks API that returns usage in streaming mode)",
|
|
)
|
|
parser.add_argument(
|
|
"--summary-file",
|
|
type=str,
|
|
help="Append the line with the summary to the specified CSV file. Useful for generating a spreadsheet with perf sweep results. If the file doesn't exist, writes out the header first",
|
|
)
|
|
parser.add_argument(
|
|
"--qps",
|
|
type=float,
|
|
default=None,
|
|
help="Enabled 'fixed QPS' mode where requests are issues at the specified rate regardless of how long the processing takes. In this case --users and --spawn-rate need to be set to a sufficiently high value (e.g. 100)",
|
|
)
|
|
parser.add_argument(
|
|
"--qps-distribution",
|
|
type=str,
|
|
choices=["constant", "uniform", "exponential"],
|
|
default="constant",
|
|
help="Must be used with --qps. Specifies how to space out requests: equally ('constant') or by sampling wait times from a distribution ('uniform' or 'exponential'). Expected QPS is going to match --qps",
|
|
)
|
|
parser.add_argument(
|
|
"--burst",
|
|
type=float,
|
|
default=None,
|
|
help="Makes requests to arrive in bursts every specified number of seconds. Note that burst duration has to be longer than maximum time of the response. Size of the burst is controlled by --users. The spawn rate -r is best set to a high value",
|
|
)
|
|
parser.add_argument(
|
|
"--tokenizer",
|
|
type=str,
|
|
help="Specify HF tokenizer to use for validating the output of the model. It's optional, we're going to rely on 'usage' or 'logprobs' field to get token count information",
|
|
)
|
|
parser.add_argument(
|
|
"--show-response",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=False,
|
|
help="Print the result of each generation",
|
|
)
|
|
parser.add_argument(
|
|
"-pcml",
|
|
"--prompt-cache-max-len",
|
|
env_var="PROMPT_CACHE_MAX_LEN",
|
|
type=int,
|
|
default=0,
|
|
help="Maximum length of the prompt cache to use. Defaults to 0 (no caching).",
|
|
)
|
|
parser.add_argument(
|
|
"--header",
|
|
action="append",
|
|
default=[],
|
|
help="Arbitrary headers to add to the inference request. Can be used multiple times. For example, --header header1:value1 --header header2:value2",
|
|
)
|
|
parser.add_argument(
|
|
"-n",
|
|
"--n",
|
|
default=1,
|
|
type=int,
|
|
help="How many sequences to generate (makes sense to use with non-zero temperature).",
|
|
)
|
|
|
|
|
|
@events.quitting.add_listener
|
|
# ADDED A NAME TO THE FUNCTION
|
|
def collect_metrics(environment, **kw):
|
|
total_latency = environment.stats.entries[("total_latency", "METRIC")]
|
|
if environment.stats.total.num_failures > 0 or total_latency.num_requests == 0:
|
|
print("Test failed due to failed requests")
|
|
environment.process_exit_code = 1
|
|
return
|
|
|
|
entries = copy.copy(InitTracker.logging_params)
|
|
if environment.parsed_options.qps is not None:
|
|
entries[
|
|
"concurrency"
|
|
] = f"QPS {environment.parsed_options.qps} {environment.parsed_options.qps_distribution}"
|
|
else:
|
|
entries["concurrency"] = InitTracker.users
|
|
for metric_name in [
|
|
"time_to_first_token",
|
|
"latency_per_token",
|
|
"num_tokens",
|
|
"total_latency",
|
|
"prompt_tokens", # might overwrite the static value based on server side tokenization
|
|
]:
|
|
entries[metric_name] = environment.stats.entries[
|
|
(metric_name, "METRIC")
|
|
].avg_response_time
|
|
if not environment.parsed_options.stream:
|
|
# if there's no streaming these metrics are meaningless
|
|
entries["time_to_first_token"] = ""
|
|
entries["latency_per_token"] = ""
|
|
entries["num_requests"] = total_latency.num_requests
|
|
entries["qps"] = total_latency.total_rps
|
|
percentile_to_report = [50, 90, 99, 99.9]
|
|
percentile_metrics = ["time_to_first_token", "total_latency"]
|
|
for percentile_metric in percentile_metrics:
|
|
metrics = environment.stats.entries[percentile_metric, "METRIC"]
|
|
for percentile in percentile_to_report:
|
|
name = f"P{percentile}_{percentile_metric}"
|
|
entries[name] = metrics.get_response_time_percentile(percentile / 100)
|
|
|
|
# Pretty print the entries
|
|
def pretty_name(s):
|
|
return " ".join([w.capitalize() for w in s.split("_")])
|
|
|
|
entries = {pretty_name(k): v for k, v in entries.items()}
|
|
|
|
# print in the final event handler to make sure our output is the last one
|
|
@events.quit.add_listener
|
|
def exit_printer(**kw):
|
|
entries = environment.stats.entries
|
|
max_width = max(len(k) for k in entries.keys())
|
|
print(" Summary ".center(80, "="))
|
|
for k, v in entries.items():
|
|
print(f"{k:<{max_width}}: {v}")
|
|
print("=" * 80)
|
|
|
|
if environment.parsed_options.summary_file:
|
|
with open(environment.parsed_options.summary_file, "a") as f:
|
|
writer = csv.DictWriter(f, fieldnames=entries.keys())
|
|
if f.tell() == 0:
|
|
writer.writeheader()
|
|
writer.writerow(entries)
|
|
|
|
return entries
|