mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
2717b5f597
* test(e2e): rename the load suite to perf, one E2E_PERF_TESTS opt-in, and a mixed-provider Locust perf test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): drop the mixed-provider perf test, keep the rename and the E2E_PERF_TESTS consolidation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry <kerry@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import random
|
|
import uuid
|
|
from itertools import cycle
|
|
from typing import Final
|
|
|
|
from locust import FastHttpUser, constant, task
|
|
|
|
_MODEL: Final = os.environ["LOAD_MODEL"]
|
|
_API_KEYS: Final = tuple(os.environ["LOAD_API_KEYS"].split(","))
|
|
_NEXT_ENDPOINT: Final = cycle(os.environ["LOAD_ENDPOINTS"].split(","))
|
|
_FILLER: Final = "x" * 40_000
|
|
|
|
|
|
def _payload() -> dict[str, object]:
|
|
"""A prompt no other request sent, so the response cache never answers for the deployment.
|
|
|
|
Both endpoints take the same body: /v1/messages requires max_tokens, which /chat/completions
|
|
also accepts, so one payload serves the whole round robin. Padded to tens of KB so a
|
|
per-request bookkeeping cost that scales with body size (string formatting, hashing) shows
|
|
up in the CPU and log-size budgets instead of hiding behind a 40-byte prompt.
|
|
"""
|
|
return {
|
|
"model": _MODEL,
|
|
"messages": [{"role": "user", "content": f"load test ping {uuid.uuid4().hex} {_FILLER}"}],
|
|
"max_tokens": 16,
|
|
}
|
|
|
|
|
|
class GatewayUser(FastHttpUser):
|
|
"""One simulated user, pinned to one endpoint for its lifetime.
|
|
|
|
Endpoints are handed out round robin as users spawn, so a run spreads evenly over them
|
|
while each user's traffic stays on a single route, the way a real client behaves.
|
|
"""
|
|
|
|
wait_time = constant(0)
|
|
|
|
def on_start(self) -> None:
|
|
self.headers = {"Authorization": f"Bearer {random.choice(_API_KEYS)}"}
|
|
self.endpoint = next(_NEXT_ENDPOINT)
|
|
|
|
@task
|
|
def call(self) -> None:
|
|
self.client.post( # pyright: ignore[reportUnknownMemberType] # locust FastHttpSession.post types json/**kwargs as Any
|
|
self.endpoint,
|
|
json=_payload(),
|
|
headers=self.headers,
|
|
name=self.endpoint,
|
|
)
|