wehub-resource-sync
157 lines
4.6 KiB
Python
157 lines
4.6 KiB
Python
from typing import List, Optional
|
|
from enum import Enum
|
|
from pydantic import BaseModel, Field
|
|
import argparse
|
|
|
|
|
|
class DistributionType(str, Enum):
|
|
CONSTANT = "constant"
|
|
UNIFORM = "uniform"
|
|
EXPONENTIAL = "exponential"
|
|
NORMAL = "normal"
|
|
|
|
|
|
class TokensDistributionType(str, Enum):
|
|
CONSTANT = "constant"
|
|
UNIFORM = "uniform"
|
|
EXPONENTIAL = "exponential"
|
|
|
|
|
|
class LoadTestConfig(BaseModel):
|
|
provider: Optional[str] = Field(
|
|
None,
|
|
description="Which flavor of API to use. If not specified, we'll try to guess based on the URL and /v1/models output",
|
|
)
|
|
|
|
model: Optional[str] = Field(
|
|
None,
|
|
description="The model to use for generating text. If not specified we will pick the first model from the service as returned by /v1/models",
|
|
)
|
|
|
|
chat: bool = Field(True, description="Use /v1/chat/completions API")
|
|
|
|
prompt_tokens: int = Field(
|
|
512,
|
|
description="Length of the prompt in tokens",
|
|
)
|
|
|
|
prompt_chars: Optional[int] = Field(
|
|
None,
|
|
description="Length of the prompt in characters",
|
|
)
|
|
|
|
prompt_text: Optional[str] = Field(
|
|
None,
|
|
description="Prompt text to use instead of generating one. It can be a file reference starting with an ampersand, e.g. `@prompt.txt`",
|
|
)
|
|
|
|
prompt_randomize: bool = Field(
|
|
False,
|
|
description="Include a few random numbers in the generated prompt to avoid caching",
|
|
)
|
|
|
|
max_tokens: int = Field(
|
|
64,
|
|
description="Max number of tokens to generate. If max_tokens_distribution is non-constant this is going to be the mean",
|
|
)
|
|
|
|
max_tokens_cap: Optional[int] = Field(
|
|
None,
|
|
description="If max_tokens_distribution is non-constant, this truncates the distribition at the specified limit",
|
|
)
|
|
|
|
max_tokens_distribution: TokensDistributionType = Field(
|
|
TokensDistributionType.CONSTANT,
|
|
description="How to sample max_tokens on each request",
|
|
)
|
|
|
|
max_tokens_range: float = Field(
|
|
0.3,
|
|
description="Specifies the width of the distribution. Specified value `alpha` is relative to `max_tokens`",
|
|
)
|
|
|
|
stream: bool = Field(True, description="Use the streaming API")
|
|
|
|
api_key: Optional[str] = Field(
|
|
None,
|
|
description="Auth for the API",
|
|
)
|
|
|
|
temperature: float = Field(0.1, description="Temperature parameter for the API")
|
|
|
|
logprobs: Optional[int] = Field(
|
|
None,
|
|
description="Whether to ask for logprobs, it makes things slower for some providers but is necessary for token count in streaming",
|
|
)
|
|
|
|
summary_file: Optional[str] = Field(
|
|
None,
|
|
description="Append the line with the summary to the specified CSV file",
|
|
)
|
|
|
|
qps: Optional[float] = Field(
|
|
None,
|
|
description="Enabled 'fixed QPS' mode where requests are issues at the specified rate regardless of how long the processing takes",
|
|
)
|
|
|
|
qps_distribution: DistributionType = Field(
|
|
DistributionType.CONSTANT,
|
|
description="Must be used with qps. Specifies how to space out requests",
|
|
)
|
|
|
|
burst: Optional[float] = Field(
|
|
None,
|
|
description="Makes requests to arrive in bursts every specified number of seconds",
|
|
)
|
|
|
|
tokenizer: Optional[str] = Field(
|
|
None,
|
|
description="Specify HF tokenizer to use for validating the output of the model",
|
|
)
|
|
|
|
show_response: bool = Field(
|
|
False,
|
|
description="Print the result of each generation",
|
|
)
|
|
|
|
prompt_cache_max_len: int = Field(
|
|
0,
|
|
description="Maximum length of the prompt cache to use",
|
|
)
|
|
|
|
header: List[str] = Field(
|
|
default_factory=list,
|
|
description="Arbitrary headers to add to the inference request",
|
|
)
|
|
|
|
n: int = Field(
|
|
1,
|
|
description="How many sequences to generate (makes sense to use with non-zero temperature)",
|
|
)
|
|
|
|
host: Optional[str] = Field(
|
|
default=None,
|
|
description="Host to load test in the following format: http://10.21.32.33",
|
|
)
|
|
|
|
reset_stats: bool = Field(
|
|
default=True,
|
|
description="Determines if stats should be reset once hatching is complete",
|
|
)
|
|
|
|
users: int = Field(
|
|
default=None,
|
|
description="Number of concurrent users to spawn for benchmarking.",
|
|
)
|
|
|
|
run_time: str = Field(
|
|
default="30s",
|
|
description="The runtime it is in form of Ns, Nm, or Nh, for seconds, minutes, and hours.",
|
|
)
|
|
|
|
def to_namespace(self) -> argparse.Namespace:
|
|
"""
|
|
Convert the model to an argparse.Namespace object
|
|
"""
|
|
return argparse.Namespace(**self.dict())
|