mirror of
https://github.com/Haidra-Org/AI-Horde.git
synced 2026-11-06 05:17:31 +00:00
test(stress): chat front scale demo and chat contract mix
- `tests/stress/run_chat_scale_demo.sh` brings up the test stack, N Horde instances and the front, drives them with `locustfile_chat.py` through client stages, samples the waitress gauges, the front's in-flight count, Redis and Postgres, and writes `samples.csv`, `summary.md` and `summary.json`; it stops every process it started; `chat_scale_demo.py` summarises a run directory per stage without gating it - An opt-in chat contract mix in the Locust suite (`--chat-mix-*` counts; a class with a zero count is left out): front clients, native-route clients and simulated chat workers with conformant, limited and flaky personas - The front client sends a catalogue of request permutations and checks each response against the outcome the contract promises (stream and body shape, tool calls, JSON schema, nonce echo, error envelope, retry and CORS headers); failures are labelled contract violations or environment conditions, and operating conditions such as a 429 or a full front are benign - `check_smoke_results.py` gates on contract and environment buckets; a coverage gate at test stop checks every requested worker and scenario group, and distributed runs merge worker reports on the master, with unique simulated-worker slots and shared persona availability - Generated worker names avoid substrings the Horde refuses (colab, tpu, google); aborting scenarios share one hourly budget, so worker names stay out of maintenance - `--chat-mix-worker-progress` and its interval make simulated workers post replies as deltas to `/api/v2/generate/text/progress` for jobs whose payload carries `stream_progress`, spread over the generation time; 404 and 413 stop posting, 403 and 409 fail the post, and 429, 503 and transport errors resend at the next interval; a `cancelled` acknowledgement stops posting and skips the submit - `--chat-mix-fail-fast-share` sends a share of the queued scenarios with `availability_mode: fail_fast` and accepts a matching 503 `queue_unavailable` as benign, which stays a contract violation for a request that did not ask for it; `--chat-mix-stream-abandon-share` and `--chat-mix-stream-abandon-chunks` close streams after that many chunks and log the request id and characters for the settlement check - The scenarios expect a declaration gap before a stream's headers and 400 `no_worker_supports_limit` for `max_tokens` above every worker's limit; `docs/how-to/send_chat_completions.md` states the same limits behaviour - Unit tests cover the mix settings and scenarios against the front in process, the progress mode, option bounds, the `fail_fast` verdicts and a raw-socket front for the abandon count; `tests/stress/README.md` documents the populations, stages, knobs, shares and the database checks for the cancel settlement
This commit is contained in:
32 files changed
+11350
-42
No files matched your search
@@ -476,14 +476,15 @@ curl -s -i "$OAI_BASE_URL/chat/completions" -H "Authorization: Bearer $HORDE_API
|
||||
| Tools, tool choice and response format | 256 KiB of compact JSON together; 128 tools; schemas nested at most 10 levels | 400 `UnsupportedChatParameter` |
|
||||
| Tool calls replayed in one assistant message | 128 | 400 |
|
||||
| `max_tokens` | 16 to 4096 | 400 `UnsupportedChatParameter` |
|
||||
| `max_tokens` against workers | At most a worker's declared length; the card's `max_completion_tokens` is the largest | 503 `no_eligible_worker` after the grace period |
|
||||
| `max_tokens` against workers | At most a worker's declared length; the card's `max_completion_tokens` is the largest | 400 `no_worker_supports_limit` on `max_tokens`, at once |
|
||||
| Non-streamed and streamed wait | 300 and 600 seconds by default; the deployment sets them | 408 `ChatRequestTimeout`, or an in-stream error |
|
||||
|
||||
`max_tokens` is the limit clients most often exceed: a client default above every online worker's declared length
|
||||
makes the request ineligible. Omitting `max_tokens` sends 512, which is also ineligible on a model whose workers all
|
||||
declare less, so send an explicit value at or below the card's `max_completion_tokens`.
|
||||
|
||||
Verification: a request at the card's `max_completion_tokens` is served; one above it ends with `no_eligible_worker`.
|
||||
Verification: a request at the card's `max_completion_tokens` is served; one above it ends with 400
|
||||
`no_worker_supports_limit` and `param` `max_tokens`.
|
||||
|
||||
## Deliberate deviations
|
||||
|
||||
|
||||
+401
-1
@@ -32,6 +32,10 @@ helpers, and Locust event hooks stay in coherent modules:
|
||||
and queue-pressure populations).
|
||||
- `locustsuite/shapes.py`, `helpers.py`, and `ground_truth.py` hold the staged
|
||||
load profiles, shared request helpers, and the attribution oracle.
|
||||
- `locustsuite/chat_*.py` hold the chat contract mix: request permutations,
|
||||
response checks, simulated worker behaviour, settings and one module per
|
||||
population; `locustsuite/users/chat_mix.py` exports its users. See
|
||||
[Chat contract mix](#chat-contract-mix).
|
||||
|
||||
The `locustfile*.py` files at the top of `tests/stress/` are thin entrypoints
|
||||
that re-export the User classes Locust should discover:
|
||||
@@ -50,6 +54,10 @@ that re-export the User classes Locust should discover:
|
||||
nonsense tokens, reviews the evidence listing as a moderator (following
|
||||
`text_url` links and checking digests), and probes the moderation surface the
|
||||
way an attacker would. See [Moderation evidence scenario](#moderation-evidence-scenario).
|
||||
- `locustfile_chat.py` drives the OpenAI-compatible chat front with streamed,
|
||||
non-streamed and impatient clients while simulated chat workers serve the
|
||||
Horde instances behind it. `run_chat_scale_demo.sh` wraps it with the
|
||||
processes it needs. See [Chat front scale demo](#chat-front-scale-demo).
|
||||
|
||||
## Prerequisites
|
||||
|
||||
@@ -57,7 +65,8 @@ that re-export the User classes Locust should discover:
|
||||
uv, `pip install locust`).
|
||||
- Bring up a running AI Horde deployment for Locust to target. The suite never
|
||||
starts the server or its backends itself; it only sends HTTP requests to the
|
||||
`--host` you give it.
|
||||
`--host` you give it. `run_chat_scale_demo.sh` is the exception: it starts its
|
||||
own Horde instances and front (see [Chat front scale demo](#chat-front-scale-demo)).
|
||||
- Stand up the server's backends. `tests/docker-compose.yml` provides an
|
||||
optional long-lived Postgres, Redis, and Garage (S3) stack for local work:
|
||||
|
||||
@@ -567,3 +576,394 @@ before the response that set a timeout arrived may still be rejected. A timeout
|
||||
that follows a filter rejection says nothing about the model threshold, so the
|
||||
onset check then tests only its upper bound.
|
||||
|
||||
## Chat contract mix
|
||||
|
||||
The chat contract mix adds chat users to the default `locustfile.py` (and
|
||||
`locustfile_shaped.py`) mix, beside the image, text and interrogation users. It
|
||||
is opt-in: `events.on_test_start` pins each population to its count like every
|
||||
other class and leaves out one whose count is zero, since the clients need a
|
||||
running OpenAI-compatible front and the workers need worker names their key
|
||||
owns. With every count at zero, the default, the mix is unchanged.
|
||||
|
||||
| Population | Count flag | What it does |
|
||||
| --- | --- | --- |
|
||||
| `ChatMixFrontClient` | `--chat-mix-front-clients` | Posts the permutations in `chat_scenarios.py` to the front, streamed or not, and checks every response against the outcome the contract promises: body and stream shape, `finish_reason`, tool calls, JSON against the schema sent, the request's nonce echoed back, usage, `x_horde`, a non-null `system_fingerprint`, error envelope and `type`, `param`, `x-should-retry`, `Retry-After` and CORS. Also reads the model cards and sends Lite's CORS preflight. |
|
||||
| `ChatMixNativeClient` | `--chat-mix-native-clients` | Drives the native routes: polls a chat request to the end, reuses an operation key unchanged and changed, asks for a dry run, reads the declaration gap from the status, and sends text requests to a chat worker's model and to a chat-only worker's model. |
|
||||
| `ChatMixWorker` | `--chat-mix-workers` | Pops chat jobs as a persona, checks each payload against its own declaration (routing and dispatch narrowing), and answers as the client's reply marker asks, echoing the request's nonce. With `--chat-mix-worker-progress`, a conforming text reply to a job whose payload carries `stream_progress` is first posted to `/api/v2/generate/text/progress` as deltas, one per `--chat-mix-worker-progress-interval` seconds (default 1) of the generation time, under the same rules as `ChatWorkerSimulator`'s progress mode; faults, tool calls and replies that break a contract are not posted. |
|
||||
|
||||
The modules are `locustsuite/chat_scenarios.py` (permutations and expected
|
||||
outcomes), `chat_checks.py` (response checks), `chat_worker_script.py`
|
||||
(personas, replies, payload checks), `chat_mix_settings.py` (arguments,
|
||||
settings, coverage gate), and one module per population
|
||||
(`chat_front_client.py`, `chat_native_client.py`, `chat_worker_user.py`);
|
||||
`users/chat_mix.py` exports the classes.
|
||||
|
||||
### Workers and their names
|
||||
|
||||
The run's chat workers take worker slots in order, and slot `n` takes persona
|
||||
`n` modulo `--chat-mix-worker-personas` (default `conformant,limited,flaky`)
|
||||
and, when names are configured, the `n`-th name, so a name always pops as the
|
||||
same persona. Each persona serves its own model (`--chat-mix-model`,
|
||||
`--chat-mix-limited-model`, `--chat-mix-flaky-model`), so a scenario reaches the
|
||||
declaration it needs and no other worker answers it:
|
||||
|
||||
- `conformant` declares every feature and guarantee and takes chat jobs only.
|
||||
- `limited` declares `seed` and `tools` with only `prompt_overflow_refused` and
|
||||
`tools_max_temperature` 0.5, and takes text jobs too; forced tool choices,
|
||||
`json_object` and hot tool temperatures on its model end in 404
|
||||
`no_worker_supports_features`.
|
||||
- `flaky` declares every feature but `vision` and `tools_with_response_format`
|
||||
and breaks an unscripted reply's contract at `--chat-mix-flaky-rate`; a
|
||||
`strict` reply that reaches the client still has to honour the request.
|
||||
|
||||
On a deployment with anti-abuse limits, reuse worker names a key already owns:
|
||||
`--chat-mix-worker-names` holds one name per worker slot of the whole run, and
|
||||
`--chat-mix-worker-api-key` is the key that owns every one of them. The run
|
||||
refuses to start with names but no key, fewer names than
|
||||
`--chat-mix-workers`, or duplicate names among the active worker slots. Without names the workers generate names under the worker key pool.
|
||||
No worker takes a new identity to get past a refusal: a configured name whose
|
||||
pop is refused (wrong owner, too many workers, maintenance) backs off and stops
|
||||
after three refusals in a row, and a generated name stops at the first. A
|
||||
stopped worker fails the run through the coverage gate.
|
||||
|
||||
A scenario that needs a worker behaviour writes `[reply:<script>]` into its last
|
||||
user message (`ReplyScript` in `chat_worker_script.py`), so the client knows the
|
||||
outcome it caused, such as a 502 `ChatRequestFaulted` after three refused
|
||||
replies. The aborting scenarios (an undeclared function, duplicate call ids, a
|
||||
worker fault) fault three jobs each, and every fault counts as the worker's
|
||||
aborted job; the Horde puts a worker into maintenance past 20 aborted jobs an
|
||||
hour, 10 while the deployment is in raid mode, as dev is. They are capped by
|
||||
`--chat-mix-aborting-per-hour`, across the whole run: 2 by default (6 aborted
|
||||
jobs), at most 6. In a distributed run only worker process 0 sends them, so the
|
||||
run's total stays the configured one; the other processes log that they skip
|
||||
them. As a breaker,
|
||||
a worker whose replies meant to conform are refused three times stops without
|
||||
faulting the last job, so a Horde regression cannot push a worker name into
|
||||
maintenance.
|
||||
|
||||
### Groups, tags and classification
|
||||
|
||||
`--chat-mix-group-weights` weights the front's scenario groups (`ordinary`,
|
||||
`tools`, `structured`, `conformance`, `gap`, `misuse`, `lite`, `vision`); a
|
||||
group left out weighs 0. Give a group weight 0 when the run has no worker of the
|
||||
persona it needs. The `vision` group runs only with `--chat-mix-vision`, which
|
||||
makes the conformant persona declare `vision` with the `base64` transport; the
|
||||
worker then checks the image reached it as a JPEG data URL (prefix and JPEG
|
||||
magic bytes) and answers `vision-ok`, which the scenario requires; the Horde has
|
||||
to store the image at submit. Every chat task carries the `chat` tag, plus
|
||||
`chat-front`, `chat-native` or `chat-worker`, so `--tags chat` runs the chat mix
|
||||
alone and `--exclude-tags chat` leaves it out.
|
||||
|
||||
Outcomes are classified in the users, and `check_smoke_results.py` gates on
|
||||
them:
|
||||
|
||||
- A response that breaks its scenario's expected outcome fails with a message
|
||||
starting `Contract violation`, whatever its status, as does an unexpected
|
||||
refusal of a worker's pop or submit, a native text or status call, or the
|
||||
model list. The checker gates on it.
|
||||
- A deployment condition fails with a message starting `Environment
|
||||
condition`: 503 `ChatAdmissionDisabled`, `MaintenanceMode`, 502
|
||||
`HordeUnavailable`, 408 `ChatRequestTimeout`, 403 `TimeoutIP`, and any 5xx
|
||||
without an error envelope (a proxy's error page). The checker reports these in
|
||||
their own section and gates on them: the run did not exercise what it meant to.
|
||||
- An intended outcome, such as a 404 declaration gap or a 502 after refused
|
||||
replies, is a success of the scenario's row. An intended 502 is also
|
||||
cross-checked on the native status (`faulted`, `fault_rc`) unless
|
||||
`--chat-mix-skip-fault-verify` is set.
|
||||
- Operating conditions are successes recorded under `<row> [benign <kind>]`:
|
||||
429, 503 `FrontAtCapacity`, 403 `KudosUpfront`, 503 `queue_unavailable` to a
|
||||
`fail_fast` request (see [Fail-fast share](#fail-fast-share)), and an
|
||||
oversized upload the front cut off before the client read its 413 (`upload
|
||||
reset before 413`: the front answers from the `Content-Length` without reading
|
||||
the body and closes the connection).
|
||||
|
||||
At test stop a coverage gate fails the run (`[coverage] chat mix`) when a
|
||||
configured persona never came online or its worker stopped itself, when no
|
||||
front scenario outside the misuse group completed, or, in a run longer than 5
|
||||
minutes, when a weighted group recorded no attempt.
|
||||
|
||||
### Anonymous share
|
||||
|
||||
`--chat-mix-anonymous-share` (0 to 1, default 0.1) sends that share of the
|
||||
`ordinary` and `lite` scenarios without `Authorization`, as anonymous users,
|
||||
under rows named `<scenario> anonymous`. The front treats a missing header as
|
||||
the anonymous key and nothing else differs on the wire
|
||||
(`docs/decisions/0035-keep-anonymous-access-on-the-chat-front.md`), so the
|
||||
expected outcome stays the scenario's. The Horde's side differs: the anonymous
|
||||
user holds no kudos and ranks last, so its queue wait is recorded apart
|
||||
(`[queue_wait_seconds] <persona> anonymous`); a request above 256 tokens, the
|
||||
floor of the Horde's upfront-kudos threshold (`require_upfront_kudos`), may get
|
||||
403 `KudosUpfront` (benign), while one at or below it never may, so that refusal
|
||||
is a contract violation there; and all anonymous traffic shares one concurrency
|
||||
limit, whose 429 is benign like any rate limit. `anonymous_plain` stays as an
|
||||
always-anonymous scenario.
|
||||
|
||||
### Fail-fast share
|
||||
|
||||
`--chat-mix-fail-fast-share` (0 to 1, default 0) sends that share of the
|
||||
`ordinary`, `tools`, `structured`, `conformance`, `gap` and `vision` scenarios
|
||||
with `horde: {"availability_mode": "fail_fast"}`, under rows named `<scenario>
|
||||
fail_fast`. The `misuse` group is refused before anything is queued, and the
|
||||
`lite` group keeps the body shape Lite sends, so neither gets the mode. It
|
||||
measures the front's admission decision at the first status read
|
||||
(`docs/reference/chat_requests.md`): the scenario's own outcome stays expected,
|
||||
and a 503 `queue_unavailable` is allowed beside it as an operating condition. The
|
||||
refusal has to carry `param: horde.availability_mode`, `Retry-After` and
|
||||
`x-should-retry: false`, as a plain HTTP error for streamed and non-streamed
|
||||
requests alike, or it is a contract violation; a `queue_unavailable` to a request
|
||||
that did not ask `fail_fast` is one too. The benign rows count the refusals by
|
||||
the reason the message states:
|
||||
|
||||
- `<row> [benign queue_unavailable no_eligible_threads]`: the pool had zero
|
||||
eligible worker threads and nothing processed the request.
|
||||
- `<row> [benign queue_unavailable wait_exceeds_budget]`: the Horde's wait
|
||||
estimate exceeded the time left in the request's budget.
|
||||
|
||||
A refusal whose message states neither reason is counted under `<row> [benign
|
||||
queue_unavailable]`.
|
||||
|
||||
### Stream-abandon share
|
||||
|
||||
`--chat-mix-stream-abandon-share` (0 to 1, default 0) makes that share of the
|
||||
streamed front requests that expect a reply close the connection after
|
||||
`--chat-mix-stream-abandon-chunks` (default 2) chunks carrying content or tool
|
||||
calls. The client closes the socket, so the front sees the disconnect and
|
||||
cancels the native request; a client that only stopped reading would not
|
||||
trigger the cancel. Such a request's row is `<row> [abandoned]`, and its response
|
||||
length is the body bytes received before the close. The client checks only what
|
||||
it read: the line framing, the role chunk first, the values every chunk shares,
|
||||
and no chunk carrying both content and tool calls. A stream that ends or carries
|
||||
an error before the count is judged whole, as without the share. Each abandoned
|
||||
stream is logged at INFO with its request id, content chunks, content
|
||||
characters and bytes received.
|
||||
|
||||
The front's cancel settles a generation in progress on the content its worker
|
||||
posted, so the share measures the settlement only with workers that post
|
||||
progress: `--chat-mix-worker-progress` for the simulated workers, with a
|
||||
generation time long enough for more deltas than the chunk count, that is a
|
||||
`--chat-mix-gen-time-min` of at least the chunk count plus one, times
|
||||
`--chat-mix-worker-progress-interval`. A close that lands after the worker
|
||||
submitted its reply ends the request without a cancelled generation, so the
|
||||
Horde can record fewer settlements than the run has `[abandoned]` rows.
|
||||
|
||||
The settlement is checked on the Horde, against the request ids the client
|
||||
logged:
|
||||
|
||||
- The Horde log has a `New Cancelled Generation <generation id> worth <kudos>
|
||||
kudos, delivered by worker: <worker> for wp <request id>` line for the
|
||||
request: the worker was paid prorated kudos and the requesting user was
|
||||
charged the same.
|
||||
- `processing_gens.settled_character_count` of that generation equals the
|
||||
content characters the client logged, or exceeds them by the content the
|
||||
worker posted after the last chunk the client read.
|
||||
- The worker's `aborted_jobs` stays unchanged: a cancelled streamed generation
|
||||
carries no dropped-job penalty.
|
||||
|
||||
### Distributed runs
|
||||
|
||||
The chat mix runs under Locust's master and worker mode, with workers on several
|
||||
hosts:
|
||||
|
||||
- Set every `--chat-mix-*` option, and the suite's other options such as the
|
||||
keys, on the master. Locust sends the master's custom options with each spawn
|
||||
message and they replace a worker's own command line and environment values,
|
||||
so a worker's `HORDE_CHAT_MIX_*` variables have no effect. The keys therefore
|
||||
travel to the workers over Locust's unencrypted connection; tunnel it.
|
||||
- Locust's dispatcher spreads a fixed count over the worker processes unevenly
|
||||
(six chat workers over two processes can land four and two), so each chat
|
||||
worker claims its slot, name and persona from the master when it starts.
|
||||
`--chat-mix-worker-names` needs one name per chat worker of the whole run.
|
||||
The master broadcasts current persona availability from the workers' stats
|
||||
reports, so front and native clients can use personas served by another
|
||||
process. Offline workers are removed, and availability expires after 15
|
||||
seconds without fresh reports.
|
||||
- Each worker process adds its coverage to its stats report every few seconds;
|
||||
the master merges the latest report of each worker and runs the coverage gate
|
||||
at test stop, where it also writes the statistics files.
|
||||
- Model card checks run once per 30 seconds per process, whatever the number of
|
||||
front clients; the budgets and the ledger are per process and bounded by the
|
||||
run's worker names and limits, not by its users.
|
||||
- A master cannot remove a user class once its dispatcher exists, so it gives
|
||||
unrequested opt-in classes and classes a tag filter emptied a weight of 0; a
|
||||
worker keeps every class, since the master's spawn messages name them all.
|
||||
- The front allows 5 open requests per client address: with one address per
|
||||
load generator, front clients beyond 5 per host mostly measure 503
|
||||
`FrontAtCapacity`. Raise `HORDE_OAI_MAX_INFLIGHT_PER_SOURCE` on the front for a
|
||||
load test, or put the volume on native clients.
|
||||
|
||||
A master and a worker on one host, and a second worker on another host that
|
||||
reaches the master's port 5557 through an SSH tunnel (both checkouts at the same
|
||||
revision):
|
||||
|
||||
```
|
||||
# master (runs no users; writes the CSV and HTML)
|
||||
locust -f tests/stress/locustfile.py --master --master-bind-host 127.0.0.1 --master-bind-port 5557 \
|
||||
--expect-workers 2 --host https://dev.example --headless --only-summary -u 220 -r 10 --run-time 30m \
|
||||
--bootstrap-requestors 0 --bootstrap-workers 0 --preflight-fail-hard \
|
||||
--requestor-api-keys "$REQUESTOR_KEYS" --worker-api-keys "$WORKER_KEYS" \
|
||||
--chat-mix-front-host https://oai.dev.example/v1 --chat-mix-front-clients 10 --chat-mix-native-clients 20 \
|
||||
--chat-mix-workers 6 --chat-mix-worker-names N0,N1,N2,N3,N4,N5 --chat-mix-worker-api-key "$WORKER_OWNER_KEY" \
|
||||
--chat-mix-anonymous-share 0.2 --csv run/chat --html run/chat.html
|
||||
|
||||
# worker on the master's host
|
||||
locust -f tests/stress/locustfile.py --worker --master-host 127.0.0.1 --master-port 5557
|
||||
|
||||
# worker on another host, after: ssh -N -L 5557:127.0.0.1:5557 <master-host>
|
||||
locust -f tests/stress/locustfile.py --worker --master-host 127.0.0.1 --master-port 5557
|
||||
|
||||
python tests/stress/check_smoke_results.py --stats run/chat_stats.csv
|
||||
```
|
||||
|
||||
### Measurements
|
||||
|
||||
The run records `CHAT-METRIC` rows for correlation with the front's and the
|
||||
Horde's telemetry: `[queue_wait_seconds] <persona>` from `x_horde`,
|
||||
`[time-to-first-keepalive] stream`, `[time-to-first-content] stream`,
|
||||
`[queue_wait_seconds] native chat` and `[submit-to-done] native chat`. A
|
||||
streamed request's own row times the response headers; its content timing is
|
||||
the metric row. At test stop each worker's submits are logged beside the
|
||||
Horde's `nonconformant_replies` and `uncompleted_jobs` counters.
|
||||
|
||||
The front allows 5 open requests per client address by default, so more than 5
|
||||
front clients from one load generator see 503 `FrontAtCapacity` (benign);
|
||||
native clients are not limited by it.
|
||||
|
||||
The chat counts are pinned, so `-u` has to cover them plus the users the rest of
|
||||
the mix should spawn; the example leaves 6 for the default classes:
|
||||
|
||||
```
|
||||
locust -f tests/stress/locustfile.py --host https://dev.example --headless -u 11 -r 5 --run-time 10m \
|
||||
--chat-mix-front-host https://oai.dev.example/v1 \
|
||||
--chat-mix-front-clients 1 --chat-mix-native-clients 1 --chat-mix-workers 3 \
|
||||
--chat-mix-worker-names NameA,NameB,NameC --chat-mix-worker-api-key "$WORKER_OWNER_KEY" --csv run/chat
|
||||
python tests/stress/check_smoke_results.py --stats run/chat_stats.csv
|
||||
```
|
||||
|
||||
## Chat front scale demo
|
||||
|
||||
The chat front (`python -m horde.oai`) holds each open `/v1/chat/completions`
|
||||
request as an idle coroutine and wakes it on the Horde's Redis completion
|
||||
signal, so a waiting client should cost the Horde one submit, one or two status
|
||||
reads and a publish, and no waitress thread while it waits. The demo measures
|
||||
that: it holds thousands of clients open against the front and records the
|
||||
Horde instances' waitress occupancy, the front's in-flight count and memory,
|
||||
and the calls the front makes. It reports numbers and passes or fails nothing.
|
||||
|
||||
Unlike the other scenarios, the demo starts its own processes.
|
||||
`run_chat_scale_demo.sh` brings up the backends through
|
||||
`tests/run_with_test_stack.sh`, then starts:
|
||||
|
||||
- N Horde instances (`server.py --waitress_threads 45`, the production default,
|
||||
on consecutive ports from 7101; the first holds `--quorum`), each logging to
|
||||
`horde-<port>.log` in the run directory;
|
||||
- the front on port 7002, listing every instance in `HORDE_OAI_HORDE_URLS` (the
|
||||
front round-robins, so no load balancer is needed), logging to `front.log`
|
||||
with one `httpx` line per call it makes to a Horde;
|
||||
- `chat_scale_demo.py sample`, which reads every few seconds each instance's
|
||||
`GET /api/v2/status/heartbeat` (the route serving the waitress `active_count`,
|
||||
`queue` and `threads` gauges), the front's `GET /healthz` `inflight` count,
|
||||
Redis `INFO`, the Postgres client backend count, and the RSS and CPU of each
|
||||
process listening on a sampled port, into `samples.csv`;
|
||||
- Locust with `locustfile_chat.py`, headless.
|
||||
|
||||
It stops every process it started on exit. The backends stay up unless `--down`
|
||||
is given. Pass the run directory as the last argument:
|
||||
|
||||
```
|
||||
bash tests/stress/run_chat_scale_demo.sh path/to/run
|
||||
bash tests/stress/run_chat_scale_demo.sh --stages 50,200 --stage-seconds 60 path/to/smaller-run
|
||||
bash tests/stress/run_chat_scale_demo.sh --instances 1 --sample-only 20 path/to/idle-run
|
||||
```
|
||||
|
||||
`--sample-only S` starts everything except Locust and samples the idle stack for
|
||||
S seconds, which checks the wiring without load.
|
||||
|
||||
The Horde instances create the schema in the stack's database and write to its
|
||||
Redis. When another test run is using the default stack, give the demo its own
|
||||
Compose project and ports so the two never share data:
|
||||
|
||||
```
|
||||
AI_HORDE_TEST_POSTGRES_PORT=25432 AI_HORDE_TEST_REDIS_PORT=26379 \
|
||||
AI_HORDE_TEST_GARAGE_S3_PORT=23900 AI_HORDE_TEST_GARAGE_RPC_PORT=23901 \
|
||||
AI_HORDE_TEST_GARAGE_ADMIN_PORT=23903 \
|
||||
bash tests/stress/run_chat_scale_demo.sh --compose-project chat-scale-demo --down path/to/run
|
||||
```
|
||||
|
||||
### Populations
|
||||
|
||||
| User | Flag | Behavior |
|
||||
| --- | --- | --- |
|
||||
| `ChatWorkerSimulator` | `--chat-workers` (fixed count) | Pops `/api/v2/generate/text/pop` from the instances round-robin with a chat declaration (`chat_jobs`, `chat_features` `seed` and `presence_penalty`, `chat_limits`, `model_id`), built by `tests/integration/_fake_chat_worker.py`. Fails the run if a chat job carries a `prompt` or no `messages`, waits `--sim-gen-time-min` to `--sim-gen-time-max` seconds (3 to 8 in this scenario), and submits a structured reply: `message`, `finish_reason`, `prompt_tokens`, `completion_tokens` and `model_id`, without `generation`. `--chat-unreported-usage-chance` of submits leave the counts out so the front's estimate runs. The reply carries the time the submit was sent. With `--chat-worker-progress`, a chat job whose payload carries `stream_progress` also posts the reply text to `/api/v2/generate/text/progress` while it generates: one delta per `--chat-worker-progress-interval` seconds (default 1) of the generation time, at least one, with `seq` counting from 1. The submitted reply starts with the posted text. Posts are recorded as `/api/v2/generate/text/progress [<status>]` rows; 404 and 413 stop the job's posting, 403 and 409 stop it and fail the post, and 429, 503 and transport errors resend the delta at the next interval. |
|
||||
| `ChatStreamClient` | `--chat-stream-clients` (share) | Posts `stream: true` with `stream_options.include_usage` and reads the events to `[DONE]`. Records time to the first keepalive, time to the content chunk and total time as `SSE` rows. An error chunk or a stream without `[DONE]` fails the request. |
|
||||
| `ChatClient` | `--chat-clients` (share) | Posts non-streamed requests and waits for the body, up to 30 seconds past `--chat-nonstream-timeout`, so the front's 408 arrives first. |
|
||||
| `ChatImpatientClient` | `--chat-impatient-clients` (share) | Posts a streamed (`--chat-impatient-stream-chance`) or non-streamed request and disconnects after `--chat-impatient-timeout` seconds, so the front cancels the job on the Horde. |
|
||||
|
||||
The client flags are shares of each stage's client count (Locust weights);
|
||||
0 leaves a population out. Clients send `Authorization: Bearer <key>` with the
|
||||
requestor keys in round-robin order, so no key reaches the Horde's per-user
|
||||
concurrency cap before the others. The runner registers one key per 25 clients
|
||||
of the largest stage and one worker key per three workers. Like an OpenAI SDK
|
||||
client, each client reuses its connection between requests and waits 2 to 6
|
||||
seconds after a refused or failed connection before its next request.
|
||||
|
||||
408, 429 and 503 from the front count as successes and appear under their own
|
||||
`[expected]` rows, as do deliberate disconnects; 5xx, transport errors and
|
||||
malformed streams fail.
|
||||
|
||||
### Stages and knobs
|
||||
|
||||
`ChatStagesShape` runs the workers alone for `--chat-warmup-seconds`, then holds
|
||||
`--chat-workers` plus each `--chat-stages` count for `--chat-stage-seconds`,
|
||||
moving between stages at `--chat-spawn-rate` users a second, and stops after the
|
||||
last stage. The stage clock starts after the test-start hooks have registered
|
||||
the keys, so a slow registration does not shorten a stage.
|
||||
|
||||
| Runner option | Environment variable | Default |
|
||||
| --- | --- | --- |
|
||||
| `--instances` | `CHAT_DEMO_HORDE_INSTANCES` | `3` |
|
||||
| `--horde-base-port` | `CHAT_DEMO_HORDE_BASE_PORT` | `7101` |
|
||||
| `--front-port` | `CHAT_DEMO_FRONT_PORT` | `7002` |
|
||||
| `--waitress-threads` | `CHAT_DEMO_WAITRESS_THREADS` | `45` |
|
||||
| `--workers` | `CHAT_DEMO_WORKERS` | `50` |
|
||||
| `--stages` | `CHAT_DEMO_STAGES` | `500,2000,5000` |
|
||||
| `--stage-seconds` | `CHAT_DEMO_STAGE_SECONDS` | `120` |
|
||||
| `--warmup-seconds` | `CHAT_DEMO_WARMUP_SECONDS` | `20` |
|
||||
| `--spawn-rate` | `CHAT_DEMO_SPAWN_RATE` | `100` |
|
||||
| `--sample-interval` | `CHAT_DEMO_SAMPLE_INTERVAL` | `5` |
|
||||
| `--compose-project` | `CHAT_DEMO_COMPOSE_PROJECT` | unset (the default stack) |
|
||||
| | `CHAT_DEMO_FRONT_MAX_INFLIGHT` | 1.2 times the largest stage plus 100, above every stage so a 503 marks the front's own limit |
|
||||
| | `CHAT_DEMO_RATELIMIT_DISABLED` | `1`; every client reaches the Horde from the front's one address, so per-IP limits would cap the run |
|
||||
| | `CHAT_DEMO_SIM_GEN_TIME_MIN`, `CHAT_DEMO_SIM_GEN_TIME_MAX` | `3`, `8` |
|
||||
| | `CHAT_DEMO_PYTHON` | the repository's `.venv` interpreter |
|
||||
|
||||
Arguments after `--` go to Locust unchanged. Locust runs as one process.
|
||||
|
||||
### Reading the summary
|
||||
|
||||
`chat_scale_demo.py summarize --run-dir <run>` (run by the script at the end)
|
||||
writes `summary.md` and `summary.json` with one column per stage:
|
||||
|
||||
| Row | Source | Reading |
|
||||
| --- | --- | --- |
|
||||
| Front in-flight, peak / mean | `samples.csv`, `/healthz` | Open client requests on the front. Close to the stage's client count when clients spend little time between requests. |
|
||||
| Horde waitress active and queue | `samples.csv`, heartbeat | Threads busy per instance and tasks waiting for one. Flat across stages means waiting clients cost no threads; a non-zero queue means an instance ran out of threads. |
|
||||
| Front to Horde calls per second, by route | `front.log` `httpx` lines | Submits (`POST /v2/generate/chat/async`), status reads (`GET .../chat/status/<id>`) and cancels (`DELETE`) the front sent. |
|
||||
| Workers pop and submit per second | `locust_stats_history.csv` | The worker side of the Horde load. |
|
||||
| Status reads per finished request | derived | Status reads divided by the front's finished lines; near 1 to 2 when the Redis signal wakes requests, higher when they fall back to polling. |
|
||||
| Woken by | `front.log` finished lines | `signal`, `poll` or `immediate` (finished before the wait began). |
|
||||
| Worker submit to front wake | evidence plus `front.log` | The front's `woke_at` minus the time the worker sent the submit, joined by request id. Includes the submit's own handling up to the publish. |
|
||||
| Stream timings | `chat_evidence.jsonl` | Time to the first keepalive (only requests that waited past a heartbeat have one), to the content, and total. |
|
||||
| Usage reported / estimated | `chat_evidence.jsonl` | `x_horde.estimated_usage` on completed requests. |
|
||||
| Client statuses, refusal codes, outcomes | `chat_evidence.jsonl` | Counts of 408, 429, 503 and the rest, the error codes behind them, and completed, refused, disconnected and failed requests. |
|
||||
| Front RSS | `samples.csv` | Per stage and the run's peak. |
|
||||
| Postgres connections | `samples.csv` | Client backends against `max_connections`, the usual first limit for several instances. |
|
||||
|
||||
The log pattern counts below the table show where a stage that fell short of
|
||||
its client count stopped. Locust's own `CPU usage above` warnings and socket
|
||||
errors (`10055` no buffer space and `10048` address in use on Windows, `Too many
|
||||
open files` elsewhere) mean the load generator was the limit. `QueuePool limit`
|
||||
or `too many clients already` in a Horde log mean Postgres connections were.
|
||||
`refused the connection` and status `unreachable` in the front log mean a Horde
|
||||
instance stopped answering. `run.env` records the settings and the open file
|
||||
limit the run had.
|
||||
|
||||
@@ -0,0 +1,862 @@
|
||||
# SPDX-FileCopyrightText: 2026 Tazlin
|
||||
#
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
"""Process helpers for the chat front scale demo driven by ``run_chat_scale_demo.sh``.
|
||||
|
||||
Three subcommands:
|
||||
|
||||
- ``front`` runs the OpenAI-compatible front (``horde.oai``) with the ``httpx`` request log enabled, one line per
|
||||
call the front makes to a Horde, so the summary can count submits, status reads and cancels by route. The
|
||||
repository root must be on ``PYTHONPATH``.
|
||||
- ``sample`` polls each Horde instance's heartbeat (which carries the waitress gauges), the front's ``/healthz``
|
||||
in-flight count, Redis ``INFO`` and, optionally, Postgres backend counts every ``--interval`` seconds, and appends
|
||||
one CSV row per sample, with the resident memory and CPU of each process found listening on the sampled ports.
|
||||
- ``summarize`` reads a run directory (samples, Locust CSVs, the client evidence file and the logs) and writes
|
||||
``summary.md`` and ``summary.json`` with per-stage numbers. It reports measurements and never passes or fails a
|
||||
run.
|
||||
|
||||
Usage:
|
||||
python tests/stress/chat_scale_demo.py sample --horde-urls http://127.0.0.1:7101 --front-url http://127.0.0.1:7002 \\
|
||||
--redis-host 127.0.0.1 --redis-port 6379 --out run/samples.csv
|
||||
python tests/stress/chat_scale_demo.py summarize --run-dir run
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import concurrent.futures
|
||||
import csv
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import re
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
from collections import Counter
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import psutil
|
||||
import requests
|
||||
|
||||
HEARTBEAT_PATH: str = "/api/v2/status/heartbeat"
|
||||
"""The Horde route that reports the waitress gauges (``queue``, ``threads``, ``active_count``)."""
|
||||
|
||||
HEALTHZ_PATH: str = "/healthz"
|
||||
"""The front route that reports its in-flight request count."""
|
||||
|
||||
SAMPLE_REQUEST_TIMEOUT_SECONDS: float = 4.0
|
||||
"""How long one gauge read may take; below the default interval, so a saturated Horde cannot stall the sampler."""
|
||||
|
||||
BYTES_PER_MEGABYTE: int = 1024 * 1024
|
||||
"""Bytes in a mebibyte, the unit RSS is reported in."""
|
||||
|
||||
HTTPX_LOGGER_NAME: str = "httpx"
|
||||
"""The logger ``httpx`` writes one INFO line per request to."""
|
||||
|
||||
HTTPX_LOG_FORMAT: str = "%(created).3f httpx %(message)s"
|
||||
"""The front's request log line: Unix time, then httpx's ``HTTP Request: METHOD URL "STATUS"``."""
|
||||
|
||||
HTTPX_LINE_PATTERN: re.Pattern[str] = re.compile(r"^(\d+\.\d+) httpx HTTP Request: (\w+) (\S+) \"HTTP/[\d.]+ (\d+)")
|
||||
"""Parse a front request log line into time, method, URL and status."""
|
||||
|
||||
FINISHED_LINE_PATTERN: re.Pattern[str] = re.compile(
|
||||
r"Chat request (\S+) finished: streamed=(\w+) submitted_at=([\d.]+) woke_at=([\d.]+) woke_by=(\S+) waited=([\d.]+)s",
|
||||
)
|
||||
"""Parse the front's one line per finished request (``horde.oai.app.RequestTrace.finish``)."""
|
||||
|
||||
SAMPLE_HORDE_COLUMN_PATTERN: re.Pattern[str] = re.compile(r"^horde_(\d+)_active$")
|
||||
"""Find the Horde ports a samples file covers from its column names."""
|
||||
|
||||
FRONT_LOG_PATTERNS: dict[str, str] = {
|
||||
"cancelled": "Cancelled chat request",
|
||||
"cancel_failed": "Cancelling chat request",
|
||||
"client_disconnected": "disconnected while waiting",
|
||||
"horde_refused_connection": "refused the connection",
|
||||
"status_unavailable": "unavailable, retrying at the next poll",
|
||||
"status_unreachable": "unreachable, retrying at the next poll",
|
||||
"signal_subscription_failed": "Chat signal subscription failed",
|
||||
"traceback": "Traceback",
|
||||
}
|
||||
"""Front log substrings counted in the summary: the cancel path and the failure modes that point at a bottleneck."""
|
||||
|
||||
HORDE_LOG_PATTERNS: dict[str, str] = {
|
||||
"queue_pool_limit": "QueuePool limit",
|
||||
"too_many_postgres_clients": "too many clients already",
|
||||
"operational_error": "OperationalError",
|
||||
"redis_error": "redis.exceptions",
|
||||
"traceback": "Traceback",
|
||||
}
|
||||
"""Horde log substrings counted in the summary as evidence of database, Redis or crash bottlenecks."""
|
||||
|
||||
LOCUST_LOG_PATTERNS: dict[str, str] = {
|
||||
"cpu_warning": "CPU usage above",
|
||||
"too_many_open_files": "Too many open files",
|
||||
"no_buffer_space": "10055",
|
||||
"address_in_use": "10048",
|
||||
"connection_refused": "10061",
|
||||
"connection_reset": "Connection reset",
|
||||
}
|
||||
"""Load generator log substrings that show the generator, not the system under test, was the limit."""
|
||||
|
||||
PERCENTILES: tuple[int, ...] = (50, 95, 99)
|
||||
"""The percentiles the summary reports for latency distributions."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProcessGauge:
|
||||
"""Track the process listening on one port, for RSS and CPU sampling."""
|
||||
|
||||
port: int
|
||||
"""The listening port the process is found by."""
|
||||
process: psutil.Process | None = None
|
||||
"""The process last found listening on ``port``."""
|
||||
|
||||
def read(self) -> tuple[float | None, float | None]:
|
||||
"""Return the process's resident memory in MB and CPU percent since the last read, or Nones when not found."""
|
||||
if self.process is None or not self.process.is_running():
|
||||
self.process = _find_listening_process(self.port)
|
||||
if self.process is None:
|
||||
return None, None
|
||||
self.process.cpu_percent(None)
|
||||
try:
|
||||
with self.process.oneshot():
|
||||
return self.process.memory_info().rss / BYTES_PER_MEGABYTE, self.process.cpu_percent(None)
|
||||
except psutil.Error as error:
|
||||
logging.warning("Reading process on port %d failed: %s", self.port, error)
|
||||
self.process = None
|
||||
return None, None
|
||||
|
||||
|
||||
def _find_listening_process(port: int) -> psutil.Process | None:
|
||||
"""Return the process listening on TCP ``port``, or None when there is none or it is not visible."""
|
||||
try:
|
||||
connections = psutil.net_connections(kind="tcp")
|
||||
except psutil.AccessDenied as error:
|
||||
logging.warning("Listing connections to find port %d is not permitted: %s", port, error)
|
||||
return None
|
||||
for connection in connections:
|
||||
is_listener = connection.status == psutil.CONN_LISTEN and connection.laddr and connection.laddr.port == port
|
||||
if is_listener and connection.pid:
|
||||
try:
|
||||
return psutil.Process(connection.pid)
|
||||
except psutil.NoSuchProcess:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _port_of(url: str) -> int:
|
||||
"""Return the port of an ``http://host:port`` URL."""
|
||||
return int(url.rstrip("/").rsplit(":", 1)[1].split("/", 1)[0])
|
||||
|
||||
|
||||
def _timed_get_json(url: str) -> tuple[dict[str, Any] | None, float]:
|
||||
"""Return the JSON body of a GET and its latency in milliseconds; the body is None on any failure."""
|
||||
started_at = time.perf_counter()
|
||||
try:
|
||||
response = requests.get(url, timeout=SAMPLE_REQUEST_TIMEOUT_SECONDS)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
except (requests.RequestException, ValueError) as error:
|
||||
logging.warning("Sampling %s failed: %s", url, error)
|
||||
return None, (time.perf_counter() - started_at) * 1000
|
||||
return body, (time.perf_counter() - started_at) * 1000
|
||||
|
||||
|
||||
def _read_redis(redis_host: str, redis_port: int) -> dict[str, Any]:
|
||||
"""Return the Redis ``INFO`` fields the demo samples, or an empty dict when Redis does not answer."""
|
||||
import redis
|
||||
|
||||
try:
|
||||
info = redis.Redis(host=redis_host, port=redis_port, socket_timeout=SAMPLE_REQUEST_TIMEOUT_SECONDS).info()
|
||||
except redis.RedisError as error:
|
||||
logging.warning("Sampling Redis failed: %s", error)
|
||||
return {}
|
||||
return {
|
||||
"redis_connected_clients": info.get("connected_clients"),
|
||||
"redis_blocked_clients": info.get("blocked_clients"),
|
||||
"redis_pubsub_channels": info.get("pubsub_channels"),
|
||||
"redis_pubsub_patterns": info.get("pubsub_patterns"),
|
||||
"redis_ops_per_sec": info.get("instantaneous_ops_per_sec"),
|
||||
"redis_used_memory_mb": round(info.get("used_memory", 0) / BYTES_PER_MEGABYTE, 1),
|
||||
}
|
||||
|
||||
|
||||
def _read_postgres(postgres_dsn: str | None) -> dict[str, Any]:
|
||||
"""Return Postgres backend counts, or an empty dict when no DSN is given or the query fails."""
|
||||
if not postgres_dsn:
|
||||
return {}
|
||||
import psycopg2
|
||||
|
||||
try:
|
||||
with psycopg2.connect(postgres_dsn, connect_timeout=int(SAMPLE_REQUEST_TIMEOUT_SECONDS)) as connection:
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(
|
||||
"SELECT count(*), count(*) FILTER (WHERE state = 'active'), "
|
||||
"(SELECT setting::int FROM pg_settings WHERE name = 'max_connections') "
|
||||
"FROM pg_stat_activity WHERE backend_type = 'client backend'",
|
||||
)
|
||||
total, active, max_connections = cursor.fetchone()
|
||||
connection.close()
|
||||
except psycopg2.Error as error:
|
||||
logging.warning("Sampling Postgres failed: %s", error)
|
||||
return {}
|
||||
return {"pg_connections": total, "pg_active": active, "pg_max_connections": max_connections}
|
||||
|
||||
|
||||
def _sample_columns(horde_ports: list[int]) -> list[str]:
|
||||
"""Return the samples CSV header for the given Horde ports."""
|
||||
columns = [
|
||||
"ts",
|
||||
"front_inflight",
|
||||
"front_signal_subscribed",
|
||||
"front_healthz_ms",
|
||||
"front_rss_mb",
|
||||
"front_cpu_percent",
|
||||
"redis_connected_clients",
|
||||
"redis_blocked_clients",
|
||||
"redis_pubsub_channels",
|
||||
"redis_pubsub_patterns",
|
||||
"redis_ops_per_sec",
|
||||
"redis_used_memory_mb",
|
||||
"pg_connections",
|
||||
"pg_active",
|
||||
"pg_max_connections",
|
||||
]
|
||||
for port in horde_ports:
|
||||
columns.extend(
|
||||
[
|
||||
f"horde_{port}_active",
|
||||
f"horde_{port}_queue",
|
||||
f"horde_{port}_threads",
|
||||
f"horde_{port}_heartbeat_ms",
|
||||
f"horde_{port}_rss_mb",
|
||||
f"horde_{port}_cpu_percent",
|
||||
],
|
||||
)
|
||||
return columns
|
||||
|
||||
|
||||
def _take_sample(
|
||||
*,
|
||||
executor: concurrent.futures.ThreadPoolExecutor,
|
||||
horde_urls: list[str],
|
||||
front_url: str,
|
||||
redis_host: str,
|
||||
redis_port: int,
|
||||
postgres_dsn: str | None,
|
||||
gauges: dict[int, ProcessGauge],
|
||||
) -> dict[str, Any]:
|
||||
"""Return one sample row; the HTTP, Redis and Postgres reads run concurrently so one slow source cannot skew the
|
||||
others' timestamps."""
|
||||
row: dict[str, Any] = {"ts": round(time.time(), 3)}
|
||||
heartbeat_futures = {horde_url: executor.submit(_timed_get_json, horde_url + HEARTBEAT_PATH) for horde_url in horde_urls}
|
||||
healthz_future = executor.submit(_timed_get_json, front_url + HEALTHZ_PATH)
|
||||
redis_future = executor.submit(_read_redis, redis_host, redis_port)
|
||||
postgres_future = executor.submit(_read_postgres, postgres_dsn)
|
||||
|
||||
healthz_body, healthz_ms = healthz_future.result()
|
||||
row["front_healthz_ms"] = round(healthz_ms, 1)
|
||||
if healthz_body is not None:
|
||||
row["front_inflight"] = healthz_body.get("inflight")
|
||||
row["front_signal_subscribed"] = healthz_body.get("signal_subscribed")
|
||||
front_rss_mb, front_cpu_percent = gauges[_port_of(front_url)].read()
|
||||
row["front_rss_mb"] = round(front_rss_mb, 1) if front_rss_mb is not None else None
|
||||
row["front_cpu_percent"] = front_cpu_percent
|
||||
row.update(redis_future.result())
|
||||
row.update(postgres_future.result())
|
||||
for horde_url, heartbeat_future in heartbeat_futures.items():
|
||||
port = _port_of(horde_url)
|
||||
heartbeat_body, heartbeat_ms = heartbeat_future.result()
|
||||
row[f"horde_{port}_heartbeat_ms"] = round(heartbeat_ms, 1)
|
||||
if heartbeat_body is not None:
|
||||
row[f"horde_{port}_active"] = heartbeat_body.get("active_count")
|
||||
row[f"horde_{port}_queue"] = heartbeat_body.get("queue")
|
||||
row[f"horde_{port}_threads"] = heartbeat_body.get("threads")
|
||||
horde_rss_mb, horde_cpu_percent = gauges[port].read()
|
||||
row[f"horde_{port}_rss_mb"] = round(horde_rss_mb, 1) if horde_rss_mb is not None else None
|
||||
row[f"horde_{port}_cpu_percent"] = horde_cpu_percent
|
||||
return row
|
||||
|
||||
|
||||
def run_sampler(arguments: argparse.Namespace) -> int:
|
||||
"""Append a sample row every interval until the duration passes or the process is stopped.
|
||||
|
||||
Args:
|
||||
arguments: The ``sample`` subcommand's arguments.
|
||||
|
||||
Returns:
|
||||
The process exit code.
|
||||
"""
|
||||
horde_urls = [horde_url.strip().rstrip("/") for horde_url in arguments.horde_urls.split(",") if horde_url.strip()]
|
||||
front_url = arguments.front_url.rstrip("/")
|
||||
horde_ports = [_port_of(horde_url) for horde_url in horde_urls]
|
||||
gauges = {port: ProcessGauge(port) for port in [*horde_ports, _port_of(front_url)]}
|
||||
columns = _sample_columns(horde_ports)
|
||||
output_path = Path(arguments.out)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
deadline = time.monotonic() + arguments.duration if arguments.duration else math.inf
|
||||
with (
|
||||
output_path.open("w", newline="", encoding="utf-8") as output_file,
|
||||
concurrent.futures.ThreadPoolExecutor(max_workers=len(horde_urls) + 3) as executor,
|
||||
):
|
||||
writer = csv.DictWriter(output_file, fieldnames=columns)
|
||||
writer.writeheader()
|
||||
while time.monotonic() < deadline:
|
||||
tick_started_at = time.monotonic()
|
||||
row = _take_sample(
|
||||
executor=executor,
|
||||
horde_urls=horde_urls,
|
||||
front_url=front_url,
|
||||
redis_host=arguments.redis_host,
|
||||
redis_port=arguments.redis_port,
|
||||
postgres_dsn=arguments.postgres_dsn,
|
||||
gauges=gauges,
|
||||
)
|
||||
writer.writerow(row)
|
||||
# Each row is flushed so a sampler stopped by the runner's teardown leaves every sample on disk.
|
||||
output_file.flush()
|
||||
time.sleep(max(0.0, arguments.interval - (time.monotonic() - tick_started_at)))
|
||||
return 0
|
||||
|
||||
|
||||
def run_front(_arguments: argparse.Namespace) -> int:
|
||||
"""Run the front with the ``httpx`` request log written to stderr; returns when uvicorn stops."""
|
||||
httpx_handler = logging.StreamHandler(sys.stderr)
|
||||
httpx_handler.setFormatter(logging.Formatter(HTTPX_LOG_FORMAT))
|
||||
httpx_logger = logging.getLogger(HTTPX_LOGGER_NAME)
|
||||
httpx_logger.addHandler(httpx_handler)
|
||||
httpx_logger.setLevel(logging.INFO)
|
||||
httpx_logger.propagate = False
|
||||
|
||||
from horde.oai.__main__ import main as run_front_main
|
||||
|
||||
run_front_main()
|
||||
return 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class StageWindow:
|
||||
"""Represents one stage of the run on the wall clock."""
|
||||
|
||||
index: int
|
||||
"""The stage's position, from 1."""
|
||||
target_clients: int
|
||||
"""The client users the stage holds, on top of the workers."""
|
||||
started_at: float
|
||||
"""When the stage started, in Unix seconds."""
|
||||
ended_at: float
|
||||
"""When the stage ended, in Unix seconds."""
|
||||
|
||||
def contains(self, timestamp: float | None) -> bool:
|
||||
"""Return whether ``timestamp`` falls in the stage."""
|
||||
return timestamp is not None and self.started_at <= timestamp < self.ended_at
|
||||
|
||||
@property
|
||||
def seconds(self) -> float:
|
||||
"""Return the stage's length."""
|
||||
return self.ended_at - self.started_at
|
||||
|
||||
|
||||
@dataclass
|
||||
class RunFiles:
|
||||
"""Represents the parsed contents of a run directory."""
|
||||
|
||||
run_start: dict[str, Any]
|
||||
"""The evidence file's ``run_start`` record."""
|
||||
run_stop_at: float | None
|
||||
"""The evidence file's ``run_stop`` time, when the run ended cleanly."""
|
||||
client_records: list[dict[str, Any]]
|
||||
"""Every client request record."""
|
||||
samples: list[dict[str, str]]
|
||||
"""Every sampler row."""
|
||||
horde_ports: list[int]
|
||||
"""The Horde ports the samples cover."""
|
||||
finished_lines: list[dict[str, Any]]
|
||||
"""The front's finished-request lines."""
|
||||
front_calls: list[tuple[float, str, str, int]]
|
||||
"""The front's Horde calls as (time, method, route, status)."""
|
||||
stats_history: list[dict[str, str]]
|
||||
"""The Locust full stats history rows."""
|
||||
stats_totals: list[dict[str, str]]
|
||||
"""The Locust final stats rows."""
|
||||
log_counts: dict[str, dict[str, int]] = field(default_factory=dict)
|
||||
"""Pattern counts per log file."""
|
||||
|
||||
|
||||
def _percentile(values: list[float], percentile: int) -> float | None:
|
||||
"""Return the nearest-rank percentile of ``values``, or None when empty."""
|
||||
if not values:
|
||||
return None
|
||||
ordered = sorted(values)
|
||||
rank = max(1, math.ceil(percentile / 100 * len(ordered)))
|
||||
return ordered[rank - 1]
|
||||
|
||||
|
||||
def _distribution(values: list[float]) -> dict[str, float | int | None]:
|
||||
"""Return the count, the reported percentiles and the maximum of ``values``, in seconds rounded to milliseconds."""
|
||||
summary: dict[str, float | int | None] = {"count": len(values)}
|
||||
for percentile in PERCENTILES:
|
||||
value = _percentile(values, percentile)
|
||||
summary[f"p{percentile}"] = round(value, 3) if value is not None else None
|
||||
summary["max"] = round(max(values), 3) if values else None
|
||||
return summary
|
||||
|
||||
|
||||
def _float_or_none(text: str | None) -> float | None:
|
||||
"""Return ``text`` as a float, or None when it is blank or not a number."""
|
||||
if text is None or text == "":
|
||||
return None
|
||||
try:
|
||||
return float(text)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def _route_of(url: str) -> str:
|
||||
"""Return a Horde call's route with the request id replaced, for grouping."""
|
||||
path = url.split("/api", 1)[-1].split("?", 1)[0]
|
||||
return re.sub(r"/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", "/<id>", path)
|
||||
|
||||
|
||||
def _count_patterns(path: Path, patterns: dict[str, str]) -> dict[str, int]:
|
||||
"""Return how many lines of ``path`` contain each pattern; an absent file counts nothing."""
|
||||
counts = dict.fromkeys(patterns, 0)
|
||||
if not path.exists():
|
||||
return counts
|
||||
with path.open(encoding="utf-8", errors="replace") as log_file:
|
||||
for line in log_file:
|
||||
for pattern_name, pattern_text in patterns.items():
|
||||
if pattern_text in line:
|
||||
counts[pattern_name] += 1
|
||||
return counts
|
||||
|
||||
|
||||
def _read_csv(path: Path) -> list[dict[str, str]]:
|
||||
"""Return the rows of a CSV file, or an empty list when it does not exist."""
|
||||
if not path.exists():
|
||||
return []
|
||||
with path.open(newline="", encoding="utf-8") as csv_file:
|
||||
return list(csv.DictReader(csv_file))
|
||||
|
||||
|
||||
def _load_run(run_dir: Path) -> RunFiles:
|
||||
"""Parse a run directory's files.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: The evidence file or its ``run_start`` record is missing.
|
||||
"""
|
||||
run_start: dict[str, Any] | None = None
|
||||
run_stop_at: float | None = None
|
||||
client_records: list[dict[str, Any]] = []
|
||||
with (run_dir / "chat_evidence.jsonl").open(encoding="utf-8") as evidence_file:
|
||||
for line in evidence_file:
|
||||
if not line.strip():
|
||||
continue
|
||||
record = json.loads(line)
|
||||
if record.get("kind") == "run_start":
|
||||
run_start = record
|
||||
elif record.get("kind") == "run_stop":
|
||||
run_stop_at = record.get("ts")
|
||||
elif record.get("kind") == "chat_request":
|
||||
client_records.append(record)
|
||||
if run_start is None:
|
||||
raise FileNotFoundError(f"No run_start record in {run_dir / 'chat_evidence.jsonl'}")
|
||||
|
||||
samples = _read_csv(run_dir / "samples.csv")
|
||||
horde_ports: list[int] = []
|
||||
if samples:
|
||||
for column in samples[0]:
|
||||
match = SAMPLE_HORDE_COLUMN_PATTERN.match(column)
|
||||
if match:
|
||||
horde_ports.append(int(match.group(1)))
|
||||
|
||||
finished_lines: list[dict[str, Any]] = []
|
||||
front_calls: list[tuple[float, str, str, int]] = []
|
||||
front_log_path = run_dir / "front.log"
|
||||
if front_log_path.exists():
|
||||
with front_log_path.open(encoding="utf-8", errors="replace") as front_log:
|
||||
for line in front_log:
|
||||
httpx_match = HTTPX_LINE_PATTERN.match(line)
|
||||
if httpx_match:
|
||||
front_calls.append(
|
||||
(float(httpx_match.group(1)), httpx_match.group(2), _route_of(httpx_match.group(3)), int(httpx_match.group(4))),
|
||||
)
|
||||
continue
|
||||
finished_match = FINISHED_LINE_PATTERN.search(line)
|
||||
if finished_match:
|
||||
finished_lines.append(
|
||||
{
|
||||
"request_id": finished_match.group(1),
|
||||
"streamed": finished_match.group(2) == "True",
|
||||
"submitted_at": float(finished_match.group(3)),
|
||||
"woke_at": float(finished_match.group(4)),
|
||||
"woke_by": finished_match.group(5),
|
||||
},
|
||||
)
|
||||
|
||||
log_counts = {
|
||||
"front.log": _count_patterns(front_log_path, FRONT_LOG_PATTERNS),
|
||||
"locust.log": _count_patterns(run_dir / "locust.log", LOCUST_LOG_PATTERNS),
|
||||
"locust.out": _count_patterns(run_dir / "locust.out", LOCUST_LOG_PATTERNS),
|
||||
}
|
||||
for horde_log_path in sorted(run_dir.glob("horde-*.log")):
|
||||
log_counts[horde_log_path.name] = _count_patterns(horde_log_path, HORDE_LOG_PATTERNS)
|
||||
|
||||
return RunFiles(
|
||||
run_start=run_start,
|
||||
run_stop_at=run_stop_at,
|
||||
client_records=client_records,
|
||||
samples=samples,
|
||||
horde_ports=horde_ports,
|
||||
finished_lines=finished_lines,
|
||||
front_calls=front_calls,
|
||||
stats_history=_read_csv(run_dir / "locust_stats_history.csv"),
|
||||
stats_totals=_read_csv(run_dir / "locust_stats.csv"),
|
||||
log_counts=log_counts,
|
||||
)
|
||||
|
||||
|
||||
def _stage_windows(run_start: dict[str, Any], run_stop_at: float | None) -> list[StageWindow]:
|
||||
"""Return the stage windows the ``run_start`` record describes; the last one ends at the run's stop when earlier."""
|
||||
first_stage_at = float(run_start["ts"]) + float(run_start["warmup_seconds"])
|
||||
stage_seconds = float(run_start["stage_seconds"])
|
||||
windows = []
|
||||
for stage_position, target_clients in enumerate(run_start["stages"]):
|
||||
started_at = first_stage_at + stage_position * stage_seconds
|
||||
ended_at = started_at + stage_seconds
|
||||
if run_stop_at is not None:
|
||||
ended_at = min(ended_at, run_stop_at)
|
||||
if ended_at <= started_at:
|
||||
break
|
||||
windows.append(StageWindow(index=stage_position + 1, target_clients=int(target_clients), started_at=started_at, ended_at=ended_at))
|
||||
return windows
|
||||
|
||||
|
||||
def _column_values(samples: list[dict[str, str]], column: str, window: StageWindow) -> list[float]:
|
||||
"""Return a sample column's numeric values within ``window``."""
|
||||
values = []
|
||||
for sample in samples:
|
||||
if not window.contains(_float_or_none(sample.get("ts"))):
|
||||
continue
|
||||
value = _float_or_none(sample.get(column))
|
||||
if value is not None:
|
||||
values.append(value)
|
||||
return values
|
||||
|
||||
|
||||
def _peak_and_mean(values: list[float]) -> dict[str, float | None]:
|
||||
"""Return the maximum and the mean of ``values``."""
|
||||
if not values:
|
||||
return {"peak": None, "mean": None}
|
||||
return {"peak": round(max(values), 1), "mean": round(statistics.fmean(values), 1)}
|
||||
|
||||
|
||||
def _history_delta(stats_history: list[dict[str, str]], name: str, window: StageWindow) -> tuple[int, int]:
|
||||
"""Return how many requests and failures Locust counted for ``name`` within ``window``, from the cumulative history."""
|
||||
in_window = [row for row in stats_history if row.get("Name") == name and window.contains(_float_or_none(row.get("Timestamp")))]
|
||||
before_window = [
|
||||
row for row in stats_history if row.get("Name") == name and (_float_or_none(row.get("Timestamp")) or 0) < window.started_at
|
||||
]
|
||||
if not in_window:
|
||||
return 0, 0
|
||||
last_row = in_window[-1]
|
||||
baseline_row = before_window[-1] if before_window else None
|
||||
requests_total = int(float(last_row.get("Total Request Count") or 0))
|
||||
failures_total = int(float(last_row.get("Total Failure Count") or 0))
|
||||
if baseline_row is not None:
|
||||
requests_total -= int(float(baseline_row.get("Total Request Count") or 0))
|
||||
failures_total -= int(float(baseline_row.get("Total Failure Count") or 0))
|
||||
return requests_total, failures_total
|
||||
|
||||
|
||||
def _peak_users(stats_history: list[dict[str, str]], window: StageWindow) -> int | None:
|
||||
"""Return the most Locust users seen within ``window``."""
|
||||
user_counts = [
|
||||
int(float(row["User Count"]))
|
||||
for row in stats_history
|
||||
if row.get("Name") == "Aggregated" and window.contains(_float_or_none(row.get("Timestamp"))) and row.get("User Count")
|
||||
]
|
||||
return max(user_counts) if user_counts else None
|
||||
|
||||
|
||||
def _summarize_stage(run_files: RunFiles, window: StageWindow) -> dict[str, Any]:
|
||||
"""Return one stage's numbers."""
|
||||
samples = run_files.samples
|
||||
stage: dict[str, Any] = {
|
||||
"stage": window.index,
|
||||
"target_clients": window.target_clients,
|
||||
"started_at": round(window.started_at, 3),
|
||||
"seconds": round(window.seconds, 1),
|
||||
"peak_locust_users": _peak_users(run_files.stats_history, window),
|
||||
"front_inflight": _peak_and_mean(_column_values(samples, "front_inflight", window)),
|
||||
"front_rss_mb": _peak_and_mean(_column_values(samples, "front_rss_mb", window)),
|
||||
"front_cpu_percent": _peak_and_mean(_column_values(samples, "front_cpu_percent", window)),
|
||||
"front_healthz_ms": _peak_and_mean(_column_values(samples, "front_healthz_ms", window)),
|
||||
"redis_connected_clients": _peak_and_mean(_column_values(samples, "redis_connected_clients", window)),
|
||||
"redis_ops_per_sec": _peak_and_mean(_column_values(samples, "redis_ops_per_sec", window)),
|
||||
"pg_connections": _peak_and_mean(_column_values(samples, "pg_connections", window)),
|
||||
"pg_active": _peak_and_mean(_column_values(samples, "pg_active", window)),
|
||||
"horde": {},
|
||||
}
|
||||
for port in run_files.horde_ports:
|
||||
stage["horde"][str(port)] = {
|
||||
metric: _peak_and_mean(_column_values(samples, f"horde_{port}_{metric}", window))
|
||||
for metric in ("active", "queue", "threads", "heartbeat_ms", "rss_mb", "cpu_percent")
|
||||
}
|
||||
|
||||
calls_by_route: Counter[str] = Counter(
|
||||
f"{method} {route}" for called_at, method, route, _status in run_files.front_calls if window.contains(called_at)
|
||||
)
|
||||
status_by_route: Counter[str] = Counter(
|
||||
f"{method} {route} {status}" for called_at, method, route, status in run_files.front_calls if window.contains(called_at)
|
||||
)
|
||||
stage["front_horde_calls_per_second"] = {route: round(count / window.seconds, 2) for route, count in sorted(calls_by_route.items())}
|
||||
stage["front_horde_call_statuses"] = dict(sorted(status_by_route.items()))
|
||||
finished_in_stage = [finished for finished in run_files.finished_lines if window.contains(finished["woke_at"])]
|
||||
status_reads = sum(count for route, count in calls_by_route.items() if route.startswith("GET /v2/generate/chat/status"))
|
||||
stage["finished_on_front"] = len(finished_in_stage)
|
||||
stage["status_reads_per_finished_request"] = round(status_reads / len(finished_in_stage), 2) if finished_in_stage else None
|
||||
submits = sum(count for route, count in calls_by_route.items() if route.startswith("POST /v2/generate/chat/async"))
|
||||
stage["status_reads_per_submit"] = round(status_reads / submits, 2) if submits else None
|
||||
stage["woke_by"] = dict(Counter(finished["woke_by"] for finished in finished_in_stage))
|
||||
|
||||
worker_rates: dict[str, Any] = {}
|
||||
for route_name in ("/api/v2/generate/text/pop", "/api/v2/generate/text/submit"):
|
||||
request_count, failure_count = _history_delta(run_files.stats_history, route_name, window)
|
||||
worker_rates[route_name] = {"per_second": round(request_count / window.seconds, 2), "failures": failure_count}
|
||||
stage["worker_calls"] = worker_rates
|
||||
|
||||
records = [record for record in run_files.client_records if window.contains(record.get("finished_at"))]
|
||||
stage["client_outcomes"] = dict(Counter(f"{record['client']} {record['outcome']}" for record in records))
|
||||
stage["client_statuses"] = dict(Counter(str(record.get("status")) for record in records))
|
||||
stage["refusal_codes"] = dict(Counter(str(record.get("error_code")) for record in records if record.get("outcome") == "refused"))
|
||||
completed = [record for record in records if record["outcome"] == "completed"]
|
||||
stage["usage"] = {
|
||||
"reported": sum(1 for record in completed if record.get("estimated_usage") is False),
|
||||
"estimated": sum(1 for record in completed if record.get("estimated_usage") is True),
|
||||
"unknown": sum(1 for record in completed if record.get("estimated_usage") is None),
|
||||
}
|
||||
streamed = [record for record in completed if record["client"] == "stream"]
|
||||
stage["stream_seconds"] = {
|
||||
"time_to_first_heartbeat": _distribution(
|
||||
[record["first_heartbeat_at"] - record["sent_at"] for record in streamed if record.get("first_heartbeat_at")],
|
||||
),
|
||||
"time_to_content": _distribution([record["content_at"] - record["sent_at"] for record in streamed if record.get("content_at")]),
|
||||
"total": _distribution([record["finished_at"] - record["sent_at"] for record in streamed]),
|
||||
}
|
||||
stage["nonstream_total_seconds"] = _distribution(
|
||||
[record["finished_at"] - record["sent_at"] for record in completed if record["client"] == "nonstream"],
|
||||
)
|
||||
|
||||
woke_at_by_request = {finished["request_id"]: finished["woke_at"] for finished in run_files.finished_lines}
|
||||
wake_latencies = [
|
||||
woke_at_by_request[record["request_id"]] - record["worker_submitted_at"]
|
||||
for record in completed
|
||||
if record.get("worker_submitted_at") and record.get("request_id") in woke_at_by_request
|
||||
]
|
||||
stage["submit_to_front_wake_seconds"] = _distribution(wake_latencies)
|
||||
return stage
|
||||
|
||||
|
||||
def _summarize_run(run_files: RunFiles) -> dict[str, Any]:
|
||||
"""Return the run's summary: settings, per-stage numbers, whole-run totals and the log pattern counts."""
|
||||
windows = _stage_windows(run_files.run_start, run_files.run_stop_at)
|
||||
stages = [_summarize_stage(run_files, window) for window in windows]
|
||||
front_rss_values = [
|
||||
value for value in (_float_or_none(sample.get("front_rss_mb")) for sample in run_files.samples) if value is not None
|
||||
]
|
||||
totals = {
|
||||
row["Name"]: {
|
||||
"requests": int(float(row.get("Request Count") or 0)),
|
||||
"failures": int(float(row.get("Failure Count") or 0)),
|
||||
"requests_per_second": round(float(row.get("Requests/s") or 0), 2),
|
||||
"median_ms": _float_or_none(row.get("Median Response Time")),
|
||||
"p95_ms": _float_or_none(row.get("95%")),
|
||||
}
|
||||
for row in run_files.stats_totals
|
||||
if row.get("Name")
|
||||
}
|
||||
return {
|
||||
"settings": run_files.run_start,
|
||||
"stages": stages,
|
||||
"front_peak_rss_mb": round(max(front_rss_values), 1) if front_rss_values else None,
|
||||
"locust_totals": totals,
|
||||
"log_pattern_counts": run_files.log_counts,
|
||||
}
|
||||
|
||||
|
||||
def _format_number(value: Any) -> str:
|
||||
"""Return a table cell for a number or None."""
|
||||
if value is None:
|
||||
return "-"
|
||||
return str(value)
|
||||
|
||||
|
||||
def _peak_mean_cell(values: dict[str, Any]) -> str:
|
||||
"""Return a ``peak / mean`` table cell."""
|
||||
return f"{_format_number(values.get('peak'))} / {_format_number(values.get('mean'))}"
|
||||
|
||||
|
||||
def _distribution_cell(values: dict[str, Any]) -> str:
|
||||
"""Return a ``p50 / p95 / max (n)`` table cell."""
|
||||
percentiles = " / ".join(_format_number(values.get(key)) for key in ("p50", "p95", "max"))
|
||||
return f"{percentiles} (n={values.get('count')})"
|
||||
|
||||
|
||||
def _render_markdown(summary: dict[str, Any]) -> str:
|
||||
"""Return the summary as Markdown tables, one column per stage."""
|
||||
stages = summary["stages"]
|
||||
lines = ["# Chat front scale demo summary", ""]
|
||||
settings = summary["settings"]
|
||||
lines.append(
|
||||
f"Workers {settings.get('workers')}, stages {settings.get('stages')} of {settings.get('stage_seconds')} s after a "
|
||||
f"{settings.get('warmup_seconds')} s warm-up, client shares {settings.get('client_shares')}, simulated generation "
|
||||
f"{settings.get('sim_gen_time_seconds')} s, Horde instances {settings.get('horde_urls')}.",
|
||||
)
|
||||
lines.append("")
|
||||
lines.append(f"Front peak RSS over the run: {_format_number(summary['front_peak_rss_mb'])} MB.")
|
||||
lines.append("")
|
||||
if not stages:
|
||||
lines.append("No stage started; see the run's logs.")
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
header = "| Metric | " + " | ".join(f"Stage {stage['stage']} ({stage['target_clients']} clients)" for stage in stages) + " |"
|
||||
separator = "| --- | " + " | ".join("---" for _ in stages) + " |"
|
||||
rows: list[tuple[str, list[str]]] = [
|
||||
("Stage length (s)", [_format_number(stage["seconds"]) for stage in stages]),
|
||||
("Peak Locust users (workers included)", [_format_number(stage["peak_locust_users"]) for stage in stages]),
|
||||
("Front in-flight, peak / mean", [_peak_mean_cell(stage["front_inflight"]) for stage in stages]),
|
||||
("Front RSS MB, peak / mean", [_peak_mean_cell(stage["front_rss_mb"]) for stage in stages]),
|
||||
("Front CPU %, peak / mean", [_peak_mean_cell(stage["front_cpu_percent"]) for stage in stages]),
|
||||
("Front /healthz ms, peak / mean", [_peak_mean_cell(stage["front_healthz_ms"]) for stage in stages]),
|
||||
("Redis clients, peak / mean", [_peak_mean_cell(stage["redis_connected_clients"]) for stage in stages]),
|
||||
("Postgres connections, peak / mean", [_peak_mean_cell(stage["pg_connections"]) for stage in stages]),
|
||||
("Postgres active, peak / mean", [_peak_mean_cell(stage["pg_active"]) for stage in stages]),
|
||||
]
|
||||
horde_ports = list(stages[0]["horde"])
|
||||
for port in horde_ports:
|
||||
rows.extend(
|
||||
[
|
||||
(f"Horde {port} waitress active, peak / mean", [_peak_mean_cell(stage["horde"][port]["active"]) for stage in stages]),
|
||||
(f"Horde {port} waitress queue, peak / mean", [_peak_mean_cell(stage["horde"][port]["queue"]) for stage in stages]),
|
||||
(f"Horde {port} heartbeat ms, peak / mean", [_peak_mean_cell(stage["horde"][port]["heartbeat_ms"]) for stage in stages]),
|
||||
(f"Horde {port} CPU %, peak / mean", [_peak_mean_cell(stage["horde"][port]["cpu_percent"]) for stage in stages]),
|
||||
(f"Horde {port} RSS MB, peak / mean", [_peak_mean_cell(stage["horde"][port]["rss_mb"]) for stage in stages]),
|
||||
],
|
||||
)
|
||||
all_routes = sorted({route for stage in stages for route in stage["front_horde_calls_per_second"]})
|
||||
for route in all_routes:
|
||||
rows.append((f"Front to Horde {route} /s", [_format_number(stage["front_horde_calls_per_second"].get(route)) for stage in stages]))
|
||||
for route_name in ("/api/v2/generate/text/pop", "/api/v2/generate/text/submit"):
|
||||
rows.append(
|
||||
(
|
||||
f"Workers {route_name} /s (failures)",
|
||||
[
|
||||
f"{stage['worker_calls'][route_name]['per_second']} ({stage['worker_calls'][route_name]['failures']})"
|
||||
for stage in stages
|
||||
],
|
||||
),
|
||||
)
|
||||
rows.extend(
|
||||
[
|
||||
("Requests finished on the front", [_format_number(stage["finished_on_front"]) for stage in stages]),
|
||||
("Status reads per finished request", [_format_number(stage["status_reads_per_finished_request"]) for stage in stages]),
|
||||
("Status reads per submit", [_format_number(stage["status_reads_per_submit"]) for stage in stages]),
|
||||
("Woken by", [_format_number(stage["woke_by"]) for stage in stages]),
|
||||
(
|
||||
"Worker submit to front wake s, p50 / p95 / max",
|
||||
[_distribution_cell(stage["submit_to_front_wake_seconds"]) for stage in stages],
|
||||
),
|
||||
(
|
||||
"Stream time to first keepalive s",
|
||||
[_distribution_cell(stage["stream_seconds"]["time_to_first_heartbeat"]) for stage in stages],
|
||||
),
|
||||
("Stream time to content s", [_distribution_cell(stage["stream_seconds"]["time_to_content"]) for stage in stages]),
|
||||
("Stream total s", [_distribution_cell(stage["stream_seconds"]["total"]) for stage in stages]),
|
||||
("Non-stream total s", [_distribution_cell(stage["nonstream_total_seconds"]) for stage in stages]),
|
||||
("Usage reported / estimated", [f"{stage['usage']['reported']} / {stage['usage']['estimated']}" for stage in stages]),
|
||||
("Client statuses", [_format_number(stage["client_statuses"]) for stage in stages]),
|
||||
("Refusal codes", [_format_number(stage["refusal_codes"]) for stage in stages]),
|
||||
("Client outcomes", [_format_number(stage["client_outcomes"]) for stage in stages]),
|
||||
],
|
||||
)
|
||||
lines.append(header)
|
||||
lines.append(separator)
|
||||
for metric_name, cells in rows:
|
||||
lines.append(f"| {metric_name} | " + " | ".join(cells) + " |")
|
||||
lines.append("")
|
||||
lines.append("## Log pattern counts")
|
||||
lines.append("")
|
||||
for log_name, counts in summary["log_pattern_counts"].items():
|
||||
lines.append(f"- {log_name}: {counts}")
|
||||
lines.append("")
|
||||
lines.append("## Locust totals")
|
||||
lines.append("")
|
||||
lines.append("| Name | Requests | Failures | Req/s | Median ms | p95 ms |")
|
||||
lines.append("| --- | --- | --- | --- | --- | --- |")
|
||||
for name, totals in summary["locust_totals"].items():
|
||||
lines.append(
|
||||
f"| {name} | {totals['requests']} | {totals['failures']} | {totals['requests_per_second']} | "
|
||||
f"{_format_number(totals['median_ms'])} | {_format_number(totals['p95_ms'])} |",
|
||||
)
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def run_summarize(arguments: argparse.Namespace) -> int:
|
||||
"""Write ``summary.json`` and ``summary.md`` into the run directory and print the Markdown.
|
||||
|
||||
Args:
|
||||
arguments: The ``summarize`` subcommand's arguments.
|
||||
|
||||
Returns:
|
||||
The process exit code.
|
||||
"""
|
||||
run_dir = Path(arguments.run_dir)
|
||||
summary = _summarize_run(_load_run(run_dir))
|
||||
(run_dir / "summary.json").write_text(json.dumps(summary, indent=2, sort_keys=True), encoding="utf-8")
|
||||
markdown = _render_markdown(summary)
|
||||
(run_dir / "summary.md").write_text(markdown, encoding="utf-8")
|
||||
sys.stdout.write(markdown)
|
||||
return 0
|
||||
|
||||
|
||||
def _build_parser() -> argparse.ArgumentParser:
|
||||
"""Create the command-line parser."""
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
front_parser = subparsers.add_parser("front", help="Run the chat front with the Horde call log enabled.")
|
||||
front_parser.set_defaults(handler=run_front)
|
||||
|
||||
sample_parser = subparsers.add_parser("sample", help="Sample the gauges into a CSV.")
|
||||
sample_parser.add_argument("--horde-urls", required=True, help="Comma-separated Horde base URLs (without /api).")
|
||||
sample_parser.add_argument("--front-url", required=True, help="The front's base URL.")
|
||||
sample_parser.add_argument("--redis-host", default="127.0.0.1", help="Redis host for INFO.")
|
||||
sample_parser.add_argument("--redis-port", type=int, default=6379, help="Redis port for INFO.")
|
||||
sample_parser.add_argument("--postgres-dsn", default=None, help="Optional libpq DSN for backend counts.")
|
||||
sample_parser.add_argument("--interval", type=float, default=5.0, help="Seconds between samples.")
|
||||
sample_parser.add_argument("--duration", type=float, default=0.0, help="Seconds to sample; 0 samples until stopped.")
|
||||
sample_parser.add_argument("--out", required=True, help="CSV file to write.")
|
||||
sample_parser.set_defaults(handler=run_sampler)
|
||||
|
||||
summarize_parser = subparsers.add_parser("summarize", help="Summarize a run directory.")
|
||||
summarize_parser.add_argument("--run-dir", required=True, help="The run directory run_chat_scale_demo.sh wrote.")
|
||||
summarize_parser.set_defaults(handler=run_summarize)
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""Run the requested subcommand and return its exit code."""
|
||||
arguments = _build_parser().parse_args()
|
||||
if arguments.command != "front":
|
||||
# The front's logging is configured by uvicorn and loguru; a root handler there would repeat their lines.
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
return int(arguments.handler(arguments))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -13,10 +13,25 @@ failure is therefore either:
|
||||
example the demand/kudos threshold 403 returned for large anonymous image
|
||||
requests). These are normal API behaviour, not a server fault.
|
||||
|
||||
A third kind is a **contract violation**: a chat mix user (``locustsuite/chat_*_client.py``, ``chat_worker_user.py``)
|
||||
found that a response breaks the chat contract its scenario expects, such as a
|
||||
stream without ``[DONE]``, a reply that ignores the requested JSON schema, or an
|
||||
expected 404 that came back as a 400. Its message contains
|
||||
``CONTRACT_VIOLATION_PREFIX`` whatever HTTP status the response had, and it gates
|
||||
the run like a crash signal, since a 4xx where the contract promises a reply is a
|
||||
regression rather than normal API behaviour.
|
||||
|
||||
A fourth kind is an **environment condition**: a chat mix user met a condition of
|
||||
the deployment, not of the contract (chat admission off, maintenance, an
|
||||
unreachable Horde, a 408 at the front's deadline, an address in timeout, a proxy
|
||||
error page). Its message contains ``ENVIRONMENT_CONDITION_PREFIX``. It is reported
|
||||
in its own section and gates the run, since a run that hit one did not exercise
|
||||
what it meant to, but it is labelled apart so nobody chases it as a regression.
|
||||
|
||||
This checker reads the ``--csv`` output produced by ``locust ... --csv <prefix>``
|
||||
and **fails only on crash signals**. It deliberately asserts nothing about
|
||||
latency or throughput -- that is the job of the baselines in
|
||||
``tests/stress/BASELINE.md``, gated manually. Benign 4xx failures are reported
|
||||
and **fails only on crash signals, contract violations and environment
|
||||
conditions**. It deliberately asserts nothing about latency or throughput --
|
||||
that is the job of the baselines in ``tests/stress/BASELINE.md``, gated manually. Benign 4xx failures are reported
|
||||
as warnings so regressions in request shaping stay visible without flaking the
|
||||
smoke gate.
|
||||
|
||||
@@ -30,11 +45,38 @@ import argparse
|
||||
import csv
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
_AGGREGATED_ROW_NAME = "Aggregated"
|
||||
_STATUS_RE = re.compile(r"Status (\d{3})")
|
||||
|
||||
CONTRACT_VIOLATION_PREFIX = "Contract violation"
|
||||
"""The text a chat contract failure carries. LOG CONTRACT: ``tests/stress/locustsuite/chat_checks.py`` defines the same
|
||||
``CONTRACT_VIOLATION_PREFIX``; change both together. Locust writes a failure's exception ``repr``, so the text is matched
|
||||
anywhere in the error, not at its start."""
|
||||
|
||||
ENVIRONMENT_CONDITION_PREFIX = "Environment condition"
|
||||
"""The text a chat environment condition carries. LOG CONTRACT: ``tests/stress/locustsuite/chat_checks.py`` defines the
|
||||
same ``ENVIRONMENT_CONDITION_PREFIX``; change both together. Matched anywhere in the error, as above."""
|
||||
|
||||
FailureEntry = tuple[str, str, str, int]
|
||||
"""One failures CSV row: method, name, error and occurrences."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClassifiedFailures:
|
||||
"""Represents a run's recorded failures sorted into the checker's buckets."""
|
||||
|
||||
contract: list[FailureEntry] = field(default_factory=list)
|
||||
"""Chat contract violations, whatever their status; gating."""
|
||||
environment: list[FailureEntry] = field(default_factory=list)
|
||||
"""Deployment conditions a chat user met; gating, reported apart."""
|
||||
crash: list[FailureEntry] = field(default_factory=list)
|
||||
"""5xx responses and transport errors; gating."""
|
||||
client: list[FailureEntry] = field(default_factory=list)
|
||||
"""4xx responses the suite did not classify; reported, not gating."""
|
||||
|
||||
|
||||
def _read_aggregated(stats_path: Path) -> dict[str, str]:
|
||||
"""Return the ``Aggregated`` summary row from a Locust ``*_stats.csv``."""
|
||||
@@ -46,21 +88,22 @@ def _read_aggregated(stats_path: Path) -> dict[str, str]:
|
||||
raise SystemExit(f"No '{_AGGREGATED_ROW_NAME}' row found in {stats_path}; did Locust write CSV output?")
|
||||
|
||||
|
||||
def _classify_failures(failures_path: Path) -> tuple[list[tuple[str, str, str, int]], list[tuple[str, str, str, int]]]:
|
||||
"""Split recorded failures into ``(crash, client)`` buckets.
|
||||
def _classify_failures(failures_path: Path) -> ClassifiedFailures:
|
||||
"""Sort recorded failures into the contract, environment, crash and client buckets.
|
||||
|
||||
A failure is a *crash* signal when its error carries a 5xx HTTP status or
|
||||
no HTTP status at all (a transport-level exception). A failure is a *client*
|
||||
rejection when it carries a 4xx HTTP status.
|
||||
A failure is a *contract* violation when its error contains
|
||||
``CONTRACT_VIOLATION_PREFIX``, whatever status it mentions, and an
|
||||
*environment* condition when it contains ``ENVIRONMENT_CONDITION_PREFIX``. Otherwise it is
|
||||
a *crash* signal when its error carries a 5xx HTTP status or no HTTP status
|
||||
at all (a transport-level exception), and a *client* rejection when it
|
||||
carries a 4xx HTTP status.
|
||||
|
||||
Returns:
|
||||
A ``(crash, client)`` tuple of lists. Each entry is
|
||||
``(method, name, error, occurrences)``.
|
||||
The buckets; each entry is ``(method, name, error, occurrences)``.
|
||||
"""
|
||||
crash: list[tuple[str, str, str, int]] = []
|
||||
client: list[tuple[str, str, str, int]] = []
|
||||
classified = ClassifiedFailures()
|
||||
if not failures_path.is_file():
|
||||
return crash, client
|
||||
return classified
|
||||
with failures_path.open(newline="") as handle:
|
||||
for row in csv.DictReader(handle):
|
||||
error = row.get("Error", "") or ""
|
||||
@@ -69,14 +112,18 @@ def _classify_failures(failures_path: Path) -> tuple[list[tuple[str, str, str, i
|
||||
occurrences = int(row.get("Occurrences", "0") or "0")
|
||||
match = _STATUS_RE.search(error)
|
||||
entry = (method, name, error, occurrences)
|
||||
if match is not None and 400 <= int(match.group(1)) <= 499:
|
||||
client.append(entry)
|
||||
if CONTRACT_VIOLATION_PREFIX in error:
|
||||
classified.contract.append(entry)
|
||||
elif ENVIRONMENT_CONDITION_PREFIX in error:
|
||||
classified.environment.append(entry)
|
||||
elif match is not None and 400 <= int(match.group(1)) <= 499:
|
||||
classified.client.append(entry)
|
||||
else:
|
||||
crash.append(entry)
|
||||
return crash, client
|
||||
classified.crash.append(entry)
|
||||
return classified
|
||||
|
||||
|
||||
def _print_bucket(label: str, bucket: list[tuple[str, str, str, int]]) -> int:
|
||||
def _print_bucket(label: str, bucket: list[FailureEntry]) -> int:
|
||||
"""Print a failure bucket to stderr and return its total occurrence count."""
|
||||
total = sum(occurrences for *_rest, occurrences in bucket)
|
||||
if not bucket:
|
||||
@@ -115,19 +162,22 @@ def main(argv: list[str] | None = None) -> int:
|
||||
if request_count <= 0:
|
||||
raise SystemExit("Smoke run drove zero requests; the workload never reached the target.")
|
||||
|
||||
crash, client = _classify_failures(failures_path)
|
||||
client_total = _print_bucket("benign client (4xx) failures [non-gating]", client)
|
||||
crash_total = _print_bucket("CRASH-class failures (5xx / transport errors)", crash)
|
||||
classified = _classify_failures(failures_path)
|
||||
client_total = _print_bucket("benign client (4xx) failures [non-gating]", classified.client)
|
||||
contract_total = _print_bucket("CONTRACT violations (chat responses breaking their expected outcome)", classified.contract)
|
||||
environment_total = _print_bucket("ENVIRONMENT conditions (deployment not ready for the run)", classified.environment)
|
||||
crash_total = _print_bucket("CRASH-class failures (5xx / transport errors)", classified.crash)
|
||||
|
||||
if client_total:
|
||||
print(f"::warning::Locust smoke saw {client_total} benign client (4xx) failure(s); not gating.")
|
||||
|
||||
if crash_total:
|
||||
if crash_total or contract_total or environment_total:
|
||||
raise SystemExit(
|
||||
f"Smoke run recorded {crash_total} crash-class failure(s) (5xx or transport-level errors); see above.",
|
||||
f"Smoke run recorded {crash_total} crash-class failure(s) (5xx or transport-level errors), "
|
||||
f"{contract_total} contract violation(s) and {environment_total} environment condition(s); see above.",
|
||||
)
|
||||
|
||||
print("Smoke run passed: no 5xx or transport-level failures.")
|
||||
print("Smoke run passed: no 5xx, transport-level failures, contract violations or environment conditions.")
|
||||
return 0
|
||||
|
||||
|
||||
|
||||
@@ -109,10 +109,29 @@ spawn-rate = 5
|
||||
# --anon-chance-poller / HORDE_ANON_CHANCE_POLLER (default: 0.75)
|
||||
# --anon-chance-requester / HORDE_ANON_CHANCE_REQUESTER (default: 0.65)
|
||||
#
|
||||
# Chat contract mix (opt-in; every count 0 leaves the chat users out; see README "Chat contract mix"):
|
||||
# --chat-mix-front-host / HORDE_CHAT_MIX_FRONT_HOST (front base URL, with or without /v1)
|
||||
# --chat-mix-front-clients / HORDE_CHAT_MIX_FRONT_CLIENTS (default: 0)
|
||||
# --chat-mix-native-clients / HORDE_CHAT_MIX_NATIVE_CLIENTS (default: 0)
|
||||
# --chat-mix-workers / HORDE_CHAT_MIX_WORKERS (default: 0)
|
||||
# --chat-mix-worker-names / HORDE_CHAT_MIX_WORKER_NAMES (default: generated names; one per chat worker of the whole run)
|
||||
# --chat-mix-worker-api-key / HORDE_CHAT_MIX_WORKER_API_KEY (the key owning every configured name; required with them)
|
||||
# --chat-mix-worker-personas / HORDE_CHAT_MIX_WORKER_PERSONAS (default: conformant,limited,flaky; taken in turn by the worker slots)
|
||||
# --chat-mix-model, --chat-mix-limited-model, --chat-mix-flaky-model
|
||||
# (default: stress/chat-conformant, -limited, -flaky)
|
||||
# --chat-mix-flaky-rate / HORDE_CHAT_MIX_FLAKY_RATE (default: 0.3)
|
||||
# --chat-mix-group-weights / HORDE_CHAT_MIX_GROUP_WEIGHTS (default: ordinary=6,tools=4,structured=3,
|
||||
# conformance=2,gap=2,misuse=3,lite=2,vision=1)
|
||||
# --chat-mix-stream-chance, --chat-mix-include-usage-chance (default: 0.5, 0.5)
|
||||
# --chat-mix-anonymous-share / HORDE_CHAT_MIX_ANONYMOUS_SHARE (default: 0.1 of ordinary and lite scenarios)
|
||||
# --chat-mix-gen-time-min, --chat-mix-gen-time-max (default: 1.0, 3.0 seconds)
|
||||
# --chat-mix-aborting-per-hour, --chat-mix-oversized-per-hour (default: 2 in the whole run, at most 6; 2)
|
||||
# --chat-mix-vision, --chat-mix-skip-fault-verify (flags)
|
||||
#
|
||||
# Staged load shape (only when using locustfile_shaped.py):
|
||||
# --stress-shape-profile / HORDE_STRESS_SHAPE_PROFILE (smoke|baseline|spike; default: baseline)
|
||||
# --stress-shape-scale / HORDE_STRESS_SHAPE_SCALE (default: 1.0)
|
||||
#
|
||||
# Task tags:
|
||||
# image, text, interrogation, meta, misuse, hot, cold, status, worker,
|
||||
# requestor, sampler-features
|
||||
# requestor, sampler-features, chat, chat-front, chat-native, chat-worker
|
||||
@@ -18,6 +18,9 @@ For staged load profiles, use ``locustfile_shaped.py`` instead.
|
||||
|
||||
from locustsuite import (
|
||||
BaselineFeatureRequester,
|
||||
ChatMixFrontClient,
|
||||
ChatMixNativeClient,
|
||||
ChatMixWorker,
|
||||
ExtendedWorkerSimulator,
|
||||
HotPathRequester,
|
||||
InterrogationRequester,
|
||||
@@ -34,6 +37,9 @@ from locustsuite import (
|
||||
|
||||
__all__ = [
|
||||
"BaselineFeatureRequester",
|
||||
"ChatMixFrontClient",
|
||||
"ChatMixNativeClient",
|
||||
"ChatMixWorker",
|
||||
"ExtendedWorkerSimulator",
|
||||
"HotPathRequester",
|
||||
"InterrogationRequester",
|
||||
|
||||
@@ -0,0 +1,271 @@
|
||||
# SPDX-FileCopyrightText: 2026 Tazlin
|
||||
#
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
"""Locust entrypoint for the chat front scale demo.
|
||||
|
||||
Spawns ``ChatWorkerSimulator`` (pinned to ``--chat-workers``) against the Horde instances, and ``ChatStreamClient``,
|
||||
``ChatClient`` and ``ChatImpatientClient`` against the front, split by the ``--chat-*-clients`` shares. The
|
||||
``ChatStagesShape`` holds the workers alone for ``--chat-warmup-seconds``, then steps the client population through
|
||||
``--chat-stages``, each for ``--chat-stage-seconds``. It reuses the suite's key-parsing, key-bootstrap and
|
||||
target-preflight ``test_start`` handler by importing ``locustsuite.events`` for its side effects; that import must
|
||||
precede the configuration handler below so keys are registered before the stage clock starts.
|
||||
|
||||
``--host`` is a Horde instance: keys are registered there and the preflight checks its heartbeat. Workers pop from
|
||||
every instance in ``--chat-horde-urls`` round-robin, and clients talk only to ``--chat-front-host``.
|
||||
|
||||
Usage:
|
||||
locust -f tests/stress/locustfile_chat.py --host http://127.0.0.1:7101 --headless \\
|
||||
--chat-front-host http://127.0.0.1:7002 \\
|
||||
--chat-horde-urls http://127.0.0.1:7101,http://127.0.0.1:7102,http://127.0.0.1:7103 \\
|
||||
--chat-workers 50 --chat-stages 500,2000,5000 --chat-stage-seconds 120 \\
|
||||
--bootstrap-requestors 210 --bootstrap-workers 20 \\
|
||||
--chat-evidence-path run/chat_evidence.jsonl --csv run/locust --csv-full-history
|
||||
|
||||
``tests/stress/run_chat_scale_demo.sh`` starts the Horde instances, the front and a gauge sampler around this
|
||||
entrypoint, and ``tests/stress/chat_scale_demo.py summarize`` turns the run's files into the summary.
|
||||
"""
|
||||
|
||||
# ruff: noqa: I001
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
from locust import events
|
||||
from locust.env import Environment
|
||||
|
||||
from locustsuite import events as _suite_events # noqa: F401
|
||||
from locustsuite.shapes import ChatStagesShape, chat_stage_clock
|
||||
from locustsuite.users.chat import (
|
||||
DEFAULT_CHAT_MODEL,
|
||||
ChatClient,
|
||||
ChatImpatientClient,
|
||||
ChatStreamClient,
|
||||
ChatWorkerSimulator,
|
||||
chat_evidence_recorder,
|
||||
chat_settings,
|
||||
)
|
||||
|
||||
__all__ = ["ChatClient", "ChatImpatientClient", "ChatStagesShape", "ChatStreamClient", "ChatWorkerSimulator"]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_SIM_GEN_TIME_MIN_SECONDS: float = 3.0
|
||||
"""The shortest simulated generation time for this scenario; a short chat reply, unlike the suite's text default."""
|
||||
|
||||
DEFAULT_SIM_GEN_TIME_MAX_SECONDS: float = 8.0
|
||||
"""The longest simulated generation time for this scenario."""
|
||||
|
||||
FRONT_PREFLIGHT_TIMEOUT_SECONDS: float = 10.0
|
||||
"""How long the start hook waits for the front's ``/healthz``."""
|
||||
|
||||
|
||||
@events.init_command_line_parser.add_listener
|
||||
def _add_chat_arguments(parser: Any) -> None:
|
||||
group = parser.add_argument_group("AI Horde Chat Front Scale Demo")
|
||||
group.add_argument(
|
||||
"--chat-front-host",
|
||||
type=str,
|
||||
env_var="HORDE_CHAT_FRONT_HOST",
|
||||
default="http://127.0.0.1:7002",
|
||||
help="Base URL of the OpenAI-compatible front the clients post to.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-horde-urls",
|
||||
type=str,
|
||||
env_var="HORDE_CHAT_HORDE_URLS",
|
||||
default="",
|
||||
help="Comma-separated Horde base URLs (without /api) the workers pop from round-robin; empty uses --host.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-model",
|
||||
type=str,
|
||||
env_var="HORDE_CHAT_MODEL",
|
||||
default=DEFAULT_CHAT_MODEL,
|
||||
help="Model the clients request and the workers serve.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-workers",
|
||||
type=int,
|
||||
env_var="HORDE_CHAT_WORKERS",
|
||||
default=50,
|
||||
help="Concurrent ChatWorkerSimulator users, held through every stage.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-stream-clients",
|
||||
type=int,
|
||||
env_var="HORDE_CHAT_STREAM_CLIENTS",
|
||||
default=70,
|
||||
help="Share (Locust weight) of each stage's clients that are ChatStreamClient users; 0 leaves them out.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-clients",
|
||||
type=int,
|
||||
env_var="HORDE_CHAT_CLIENTS",
|
||||
default=25,
|
||||
help="Share (Locust weight) of each stage's clients that are non-streamed ChatClient users; 0 leaves them out.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-impatient-clients",
|
||||
type=int,
|
||||
env_var="HORDE_CHAT_IMPATIENT_CLIENTS",
|
||||
default=5,
|
||||
help="Share (Locust weight) of each stage's clients that are ChatImpatientClient users; 0 leaves them out.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-stages",
|
||||
type=str,
|
||||
env_var="HORDE_CHAT_STAGES",
|
||||
default="500,2000,5000",
|
||||
help="Comma-separated concurrent client counts, one per stage, on top of the workers.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-stage-seconds",
|
||||
type=float,
|
||||
env_var="HORDE_CHAT_STAGE_SECONDS",
|
||||
default=120.0,
|
||||
help="How long each stage is held, including its ramp.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-warmup-seconds",
|
||||
type=float,
|
||||
env_var="HORDE_CHAT_WARMUP_SECONDS",
|
||||
default=20.0,
|
||||
help="How long the workers run alone before the first stage, so every Horde knows them before clients arrive.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-spawn-rate",
|
||||
type=float,
|
||||
env_var="HORDE_CHAT_SPAWN_RATE",
|
||||
default=100.0,
|
||||
help="Users started (or stopped) per second when moving between stages.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-max-tokens",
|
||||
type=int,
|
||||
env_var="HORDE_CHAT_MAX_TOKENS",
|
||||
default=64,
|
||||
help="max_tokens of every client request.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-no-include-usage",
|
||||
action="store_true",
|
||||
env_var="HORDE_CHAT_NO_INCLUDE_USAGE",
|
||||
default=False,
|
||||
help="Leave stream_options.include_usage out of streamed requests.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-impatient-timeout",
|
||||
type=float,
|
||||
env_var="HORDE_CHAT_IMPATIENT_TIMEOUT",
|
||||
default=2.0,
|
||||
help="Seconds an impatient client waits before disconnecting.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-impatient-stream-chance",
|
||||
type=float,
|
||||
env_var="HORDE_CHAT_IMPATIENT_STREAM_CHANCE",
|
||||
default=0.5,
|
||||
help="Probability (0.0-1.0) that an impatient request is streamed rather than non-streamed.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-nonstream-timeout",
|
||||
type=float,
|
||||
env_var="HORDE_CHAT_NONSTREAM_TIMEOUT",
|
||||
default=300.0,
|
||||
help="The front's HORDE_OAI_NONSTREAM_TIMEOUT_SECONDS; non-streamed clients wait 30 seconds longer than it.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-unreported-usage-chance",
|
||||
type=float,
|
||||
env_var="HORDE_CHAT_UNREPORTED_USAGE_CHANCE",
|
||||
default=0.1,
|
||||
help="Probability (0.0-1.0) that a worker submits without token counts, so the front's estimate path runs.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-worker-progress",
|
||||
action="store_true",
|
||||
env_var="HORDE_CHAT_WORKER_PROGRESS",
|
||||
default=False,
|
||||
help="Workers post the reply as content deltas to /api/v2/generate/text/progress for jobs whose client streams.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-worker-progress-interval",
|
||||
type=float,
|
||||
env_var="HORDE_CHAT_WORKER_PROGRESS_INTERVAL",
|
||||
default=1.0,
|
||||
help="Seconds between two progress posts of one job; a job posts one delta per interval of its generation time.",
|
||||
)
|
||||
group.add_argument(
|
||||
"--chat-evidence-path",
|
||||
type=str,
|
||||
env_var="HORDE_CHAT_EVIDENCE_PATH",
|
||||
default="chat_evidence.jsonl",
|
||||
help="JSONL file of every client request, truncated at test start; chat_scale_demo.py summarize reads it.",
|
||||
)
|
||||
# The suite's text default (4 to 120 seconds) models long completions; a chat reply of max_tokens 64 is short.
|
||||
parser.set_defaults(sim_gen_time_min=DEFAULT_SIM_GEN_TIME_MIN_SECONDS, sim_gen_time_max=DEFAULT_SIM_GEN_TIME_MAX_SECONDS)
|
||||
|
||||
|
||||
def _check_front(front_host: str) -> None:
|
||||
"""Log whether the front's ``/healthz`` answers, so a wrong ``--chat-front-host`` shows before the first stage."""
|
||||
healthz_url = f"{front_host}/healthz"
|
||||
try:
|
||||
response = requests.get(healthz_url, timeout=FRONT_PREFLIGHT_TIMEOUT_SECONDS)
|
||||
except requests.RequestException as error:
|
||||
logger.warning("Front preflight failed for %s: %s", healthz_url, error)
|
||||
return
|
||||
if response.status_code >= 400:
|
||||
logger.warning("Front preflight returned HTTP %d for %s", response.status_code, healthz_url)
|
||||
return
|
||||
logger.info("Front preflight OK: %s %s", healthz_url, response.text[:200])
|
||||
|
||||
|
||||
@events.test_start.add_listener
|
||||
def _configure_chat_run(environment: Environment, **_kwargs: Any) -> None:
|
||||
options = environment.parsed_options
|
||||
chat_settings.configure(options, fallback_horde_url=environment.host or "")
|
||||
|
||||
ChatWorkerSimulator.fixed_count = max(int(options.chat_workers), 0)
|
||||
client_shares = (
|
||||
(ChatStreamClient, int(options.chat_stream_clients)),
|
||||
(ChatClient, int(options.chat_clients)),
|
||||
(ChatImpatientClient, int(options.chat_impatient_clients)),
|
||||
)
|
||||
excluded_classes: set[type] = {ChatWorkerSimulator} if ChatWorkerSimulator.fixed_count == 0 else set()
|
||||
for client_class, share in client_shares:
|
||||
if share <= 0:
|
||||
excluded_classes.add(client_class)
|
||||
continue
|
||||
client_class.weight = share
|
||||
environment.user_classes[:] = [user_class for user_class in environment.user_classes if user_class not in excluded_classes]
|
||||
|
||||
_check_front(chat_settings.front_host)
|
||||
chat_evidence_recorder.configure(options.chat_evidence_path)
|
||||
chat_stage_clock.mark_started()
|
||||
chat_evidence_recorder.write(
|
||||
{
|
||||
"kind": "run_start",
|
||||
"ts": chat_stage_clock.started_at,
|
||||
"warmup_seconds": float(options.chat_warmup_seconds),
|
||||
"stage_seconds": float(options.chat_stage_seconds),
|
||||
"stages": [int(stage) for stage in str(options.chat_stages).split(",") if stage.strip()],
|
||||
"workers": ChatWorkerSimulator.fixed_count,
|
||||
"client_shares": {client_class.__name__: share for client_class, share in client_shares},
|
||||
"front_host": chat_settings.front_host,
|
||||
"horde_urls": chat_settings.horde_urls,
|
||||
"model": chat_settings.model,
|
||||
"sim_gen_time_seconds": [chat_settings.gen_time_min_seconds, chat_settings.gen_time_max_seconds],
|
||||
"worker_progress": chat_settings.worker_progress,
|
||||
"worker_progress_interval_seconds": chat_settings.worker_progress_interval_seconds,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@events.test_stop.add_listener
|
||||
def _close_chat_evidence(environment: Environment, **_kwargs: Any) -> None:
|
||||
chat_evidence_recorder.write({"kind": "run_stop", "ts": time.time()})
|
||||
chat_evidence_recorder.close()
|
||||
@@ -10,6 +10,9 @@ This imports the same users/events as ``locustfile.py`` and additionally exposes
|
||||
"""
|
||||
|
||||
from locustsuite import (
|
||||
ChatMixFrontClient,
|
||||
ChatMixNativeClient,
|
||||
ChatMixWorker,
|
||||
ExtendedWorkerSimulator,
|
||||
HotPathRequester,
|
||||
InterrogationRequester,
|
||||
@@ -26,6 +29,9 @@ from locustsuite import (
|
||||
from locustsuite.shapes import HordeStagesShape
|
||||
|
||||
__all__ = [
|
||||
"ChatMixFrontClient",
|
||||
"ChatMixNativeClient",
|
||||
"ChatMixWorker",
|
||||
"ExtendedWorkerSimulator",
|
||||
"HordeStagesShape",
|
||||
"HotPathRequester",
|
||||
|
||||
@@ -11,6 +11,9 @@ the concrete ``HttpUser`` classes that Locust should discover.
|
||||
from . import events as _events # noqa: F401 - imported for Locust event registration side effects
|
||||
from .users import (
|
||||
BaselineFeatureRequester,
|
||||
ChatMixFrontClient,
|
||||
ChatMixNativeClient,
|
||||
ChatMixWorker,
|
||||
ExtendedWorkerSimulator,
|
||||
HotPathRequester,
|
||||
InterrogationRequester,
|
||||
@@ -27,6 +30,9 @@ from .users import (
|
||||
|
||||
__all__ = [
|
||||
"BaselineFeatureRequester",
|
||||
"ChatMixFrontClient",
|
||||
"ChatMixNativeClient",
|
||||
"ChatMixWorker",
|
||||
"ExtendedWorkerSimulator",
|
||||
"HotPathRequester",
|
||||
"InterrogationRequester",
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,580 @@
|
||||
# SPDX-FileCopyrightText: 2026 Tazlin
|
||||
#
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
"""The chat mix's front client: sends the scenarios in [`chat_scenarios`][tests.stress.locustsuite.chat_scenarios] to
|
||||
the OpenAI-compatible front and records each response as the contract judges it.
|
||||
|
||||
A response that breaks its scenario's expected outcome fails as a contract violation; a deployment condition fails as
|
||||
an environment condition; an operating condition (rate limit, full front, too little kudos) is a ``[benign ...]``
|
||||
success; the expected outcome, a refusal included, is a plain success of the scenario's row.
|
||||
|
||||
Two shares change how a request is sent. Under ``--chat-mix-fail-fast-share`` a queued scenario asks
|
||||
``horde.availability_mode: fail_fast``, and its 503 ``queue_unavailable`` is a benign outcome labelled by the reason the
|
||||
front states. Under ``--chat-mix-stream-abandon-share`` a streamed request's client closes the connection after
|
||||
``--chat-mix-stream-abandon-chunks`` content chunks, which makes the front cancel the native request; the row is
|
||||
``<row> [abandoned]`` and its response length is the bytes received before the close.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
from locust import HttpUser, between, tag, task
|
||||
from locust.exception import RescheduleTask
|
||||
|
||||
from .chat_checks import (
|
||||
COMPLETION_ID_PREFIX,
|
||||
BenignRefusal,
|
||||
CompletionExpectation,
|
||||
ErrorExpectation,
|
||||
FrontVerdict,
|
||||
StreamTranscript,
|
||||
classify_environment_condition,
|
||||
error_code_of,
|
||||
find_abandoned_stream_violations,
|
||||
find_cors_violations,
|
||||
judge_front_response,
|
||||
queue_unavailable_reason,
|
||||
)
|
||||
from .chat_mix_settings import (
|
||||
back_off,
|
||||
chat_mix_settings,
|
||||
fail_contract,
|
||||
fail_environment,
|
||||
record_benign,
|
||||
record_measurement,
|
||||
)
|
||||
from .chat_scenarios import (
|
||||
ANONYMOUS_SHARE_GROUPS,
|
||||
FAIL_FAST_SHARE_GROUPS,
|
||||
LITE_ORIGIN,
|
||||
STRESS_CLIENT_TITLE,
|
||||
AuthorizationMode,
|
||||
ChatScenario,
|
||||
ScenarioContext,
|
||||
ScenarioGroup,
|
||||
StreamMode,
|
||||
as_anonymous,
|
||||
as_fail_fast,
|
||||
build_authorization_headers,
|
||||
pick_scenario_builder,
|
||||
)
|
||||
from .chat_wire import CHAT_COMPLETIONS_PATH, CONNECT_TIMEOUT_SECONDS, STREAM_READ_TIMEOUT_SECONDS
|
||||
from .chat_worker_script import DEFAULT_PERSONA_MAX_LENGTH, ChatPersonaRole, build_persona, find_request_nonce
|
||||
from .config import _config
|
||||
from .helpers import _headers, _pick_requestor_key, _safe_json
|
||||
|
||||
__all__ = ["ChatMixFrontClient", "build_request_body"]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CHAT_TAG: str = "chat"
|
||||
"""The tag every chat mix task carries, so ``--tags chat`` runs only chat and ``--exclude-tags chat`` leaves it out."""
|
||||
|
||||
MODELS_PATH: str = "/v1/models"
|
||||
"""The front's model list."""
|
||||
|
||||
NATIVE_CHAT_STATUS_PATH: str = "/api/v2/generate/chat/status"
|
||||
"""The native chat status route, followed by the request id."""
|
||||
|
||||
NONSTREAM_READ_TIMEOUT_SECONDS: float = 330.0
|
||||
"""How long a non-streamed client waits: the front's default 300 second deadline plus a margin, so its 408 comes first."""
|
||||
|
||||
MODEL_CARD_GRACE_SECONDS: float = 60.0
|
||||
"""How long after a persona's first pop its model card is checked: the model cache and the model list cache lag."""
|
||||
|
||||
UPLOAD_RESET_BEFORE_413: str = "upload reset before 413"
|
||||
"""The benign label of an oversized upload the front cut off before the client read its 413."""
|
||||
|
||||
EXCERPT_CHARACTERS: int = 200
|
||||
"""How much of a body a failure quotes."""
|
||||
|
||||
ABANDONED_LABEL: str = "abandoned"
|
||||
"""The row label of a stream its client closed after the configured content chunks."""
|
||||
|
||||
STREAM_READ_CHUNK_BYTES: int = 512
|
||||
"""The most bytes one read of an abandoning client's stream takes, the size ``requests.iter_lines`` reads by default, so
|
||||
a line is seen as soon as it arrives."""
|
||||
|
||||
CONFORMANT_CARD_PARAMETERS: frozenset[str] = frozenset(
|
||||
{
|
||||
"max_tokens",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"stop",
|
||||
"seed",
|
||||
"frequency_penalty",
|
||||
"presence_penalty",
|
||||
"tools",
|
||||
"tool_choice",
|
||||
"parallel_tool_calls",
|
||||
"response_format",
|
||||
},
|
||||
)
|
||||
"""The ``supported_parameters`` a model served by the conformant or flaky persona lists."""
|
||||
|
||||
LIMITED_CARD_PARAMETERS: frozenset[str] = frozenset({"max_tokens", "temperature", "top_p", "stop", "seed", "tools"})
|
||||
"""The ``supported_parameters`` a model served only by the limited persona lists, and no more."""
|
||||
|
||||
|
||||
class ChatMixFrontClient(HttpUser):
|
||||
"""Send the chat scenarios to the OpenAI-compatible front and check each response against its promised outcome."""
|
||||
|
||||
weight = 1
|
||||
fixed_count = 0
|
||||
wait_time = between(1, 3)
|
||||
|
||||
def on_start(self) -> None:
|
||||
"""Point the client at the front and take a requestor key."""
|
||||
self.client.base_url = chat_mix_settings.front_host
|
||||
self.api_key = _pick_requestor_key()
|
||||
self.rng = random.Random()
|
||||
|
||||
@tag(CHAT_TAG, "chat-front")
|
||||
@task(10)
|
||||
def run_scenario(self) -> None:
|
||||
"""Send one scenario from a group picked by the run's weights, skipping one whose worker is not online or whose
|
||||
budget is spent."""
|
||||
weighted_groups = [(group, weight) for group, weight in chat_mix_settings.group_weights.items() if weight > 0]
|
||||
if not weighted_groups:
|
||||
return
|
||||
groups, weights = zip(*weighted_groups, strict=True)
|
||||
context = ScenarioContext(
|
||||
models=chat_mix_settings.models, worker_max_length=DEFAULT_PERSONA_MAX_LENGTH, nonce=uuid.uuid4().hex[:12]
|
||||
)
|
||||
scenario = pick_scenario_builder(self.rng.choices(groups, weights=weights)[0], rng=self.rng)(context)
|
||||
if self._sends_anonymously(scenario):
|
||||
scenario = as_anonymous(scenario)
|
||||
if self._sends_fail_fast(scenario):
|
||||
scenario = as_fail_fast(scenario)
|
||||
if not chat_mix_settings.is_persona_online(scenario.persona):
|
||||
return
|
||||
if scenario.aborts_worker_jobs and not chat_mix_settings.aborting_budget.try_take():
|
||||
return
|
||||
if scenario.oversized_body and not chat_mix_settings.oversized_budget.try_take():
|
||||
return
|
||||
self.send_scenario(scenario)
|
||||
|
||||
def _sends_anonymously(self, scenario: ChatScenario) -> bool:
|
||||
"""Return whether this request of an ordinary or Lite scenario goes without ``Authorization``, at the run's
|
||||
anonymous share."""
|
||||
is_eligible = scenario.group in ANONYMOUS_SHARE_GROUPS and scenario.authorization == AuthorizationMode.REQUESTOR_KEY
|
||||
return is_eligible and self.rng.random() < chat_mix_settings.anonymous_share
|
||||
|
||||
def _sends_fail_fast(self, scenario: ChatScenario) -> bool:
|
||||
"""Return whether this request of a queued scenario asks ``horde.availability_mode: fail_fast``, at the run's
|
||||
fail-fast share."""
|
||||
if chat_mix_settings.fail_fast_share <= 0 or scenario.group not in FAIL_FAST_SHARE_GROUPS:
|
||||
return False
|
||||
return self.rng.random() < chat_mix_settings.fail_fast_share
|
||||
|
||||
@tag(CHAT_TAG, "chat-front")
|
||||
@task(1)
|
||||
def check_model_cards(self) -> None:
|
||||
"""Read the model list as a browser would and check each online persona's card against its declaration.
|
||||
|
||||
The check runs once per ``MODEL_CARD_CHECK_INTERVAL_SECONDS`` per process, however many front clients it runs,
|
||||
so the model list sees a fixed rate of reads rather than one growing with the users.
|
||||
"""
|
||||
if not chat_mix_settings.model_card_checks.try_open():
|
||||
return
|
||||
name = f"{MODELS_PATH} [cards]"
|
||||
headers = {**self._headers_for(AuthorizationMode.REQUESTOR_KEY), "Origin": LITE_ORIGIN}
|
||||
with self.client.get(MODELS_PATH, headers=headers, catch_response=True, name=name) as resp:
|
||||
if resp.status_code == 0:
|
||||
resp.failure(f"Transport error: {resp.error}")
|
||||
return
|
||||
if resp.status_code != 200:
|
||||
self._fail_unexpected_status(resp, "model cards")
|
||||
return
|
||||
violations = find_cors_violations(resp.headers)
|
||||
cards = {card.get("id"): card for card in (_safe_json(resp) or {}).get("data") or [] if isinstance(card, dict)}
|
||||
for persona, online_since in chat_mix_settings.persona_online_since().items():
|
||||
if time.time() - online_since < MODEL_CARD_GRACE_SECONDS:
|
||||
continue
|
||||
violations.extend(find_card_violations(persona, cards.get(chat_mix_settings.model_of(persona))))
|
||||
if violations:
|
||||
fail_contract(resp, "model cards", violations)
|
||||
return
|
||||
resp.success()
|
||||
|
||||
@tag(CHAT_TAG, "chat-front")
|
||||
@task(1)
|
||||
def lite_preflight(self) -> None:
|
||||
"""Send the CORS preflight a browser sends before Lite's POST, and check it allows the call."""
|
||||
if chat_mix_settings.group_weights.get(ScenarioGroup.LITE, 0) <= 0:
|
||||
return
|
||||
preflight_headers = {
|
||||
"Origin": LITE_ORIGIN,
|
||||
"Access-Control-Request-Method": "POST",
|
||||
"Access-Control-Request-Headers": "authorization,content-type",
|
||||
}
|
||||
name = f"{CHAT_COMPLETIONS_PATH} [lite preflight]"
|
||||
with self.client.options(CHAT_COMPLETIONS_PATH, headers=preflight_headers, catch_response=True, name=name) as resp:
|
||||
if resp.status_code == 0:
|
||||
resp.failure(f"Transport error: {resp.error}")
|
||||
return
|
||||
violations = [] if resp.status_code in (200, 204) else [f"the preflight answered {resp.status_code}"]
|
||||
violations.extend(find_cors_violations(resp.headers))
|
||||
allowed_headers = (resp.headers.get("access-control-allow-headers") or "").lower()
|
||||
if "authorization" not in allowed_headers and allowed_headers != "*":
|
||||
violations.append(f"the preflight does not allow the authorization header: {allowed_headers!r}")
|
||||
if violations:
|
||||
fail_contract(resp, "lite preflight", violations)
|
||||
return
|
||||
resp.success()
|
||||
|
||||
def send_scenario(self, scenario: ChatScenario) -> None:
|
||||
"""Send one scenario, judge the response and record the outcome.
|
||||
|
||||
Raises:
|
||||
RescheduleTask: After a back-off for a rate limit or a full front, so the client starts afresh.
|
||||
"""
|
||||
streamed = self._decides_to_stream(scenario)
|
||||
include_usage = (
|
||||
streamed and scenario.stream_mode == StreamMode.EITHER and self.rng.random() < chat_mix_settings.include_usage_chance
|
||||
)
|
||||
abandon_after = self._abandon_after(scenario, streamed=streamed)
|
||||
request_body = build_request_body(scenario, streamed=streamed, include_usage=include_usage)
|
||||
headers = {**self._headers_for(scenario.authorization), **scenario.extra_headers, "Content-Type": "application/json"}
|
||||
name = f"{CHAT_COMPLETIONS_PATH} [{scenario.name}] [{'stream' if streamed else 'non-stream'}]"
|
||||
payload_bytes = chat_mix_settings.oversized_body() if scenario.oversized_body else json.dumps(request_body).encode()
|
||||
read_timeout = STREAM_READ_TIMEOUT_SECONDS if streamed else NONSTREAM_READ_TIMEOUT_SECONDS
|
||||
chat_mix_settings.coverage.group_attempts[scenario.group.value] += 1
|
||||
sent_at = time.time()
|
||||
verdict: FrontVerdict | None = None
|
||||
with self.client.post(
|
||||
CHAT_COMPLETIONS_PATH,
|
||||
data=payload_bytes,
|
||||
headers=headers,
|
||||
catch_response=True,
|
||||
name=name,
|
||||
stream=streamed,
|
||||
timeout=(CONNECT_TIMEOUT_SECONDS, read_timeout),
|
||||
) as resp:
|
||||
if resp.status_code == 0:
|
||||
self._record_transport_error(resp, scenario, name)
|
||||
return
|
||||
if abandon_after is not None and resp.status_code == 200 and isinstance(scenario.expectation, CompletionExpectation):
|
||||
transcript, received_bytes = self._read_until_abandoned(resp, abandon_after)
|
||||
body = None
|
||||
if transcript is None:
|
||||
return
|
||||
if received_bytes is not None:
|
||||
self._record_abandoned(resp, scenario, name, scenario.expectation, transcript, received_bytes=received_bytes)
|
||||
return
|
||||
else:
|
||||
transcript, body = self._read_response(resp, streamed=streamed)
|
||||
if transcript is None and streamed and resp.status_code == 200:
|
||||
return
|
||||
verdict = judge_front_response(
|
||||
scenario.expectation,
|
||||
scenario.alternative_error,
|
||||
nonce=find_request_nonce(scenario.body.get("messages") or []),
|
||||
kudos_upfront_benign=scenario.kudos_upfront_benign,
|
||||
status_code=resp.status_code,
|
||||
headers=resp.headers,
|
||||
body=body,
|
||||
transcript=transcript,
|
||||
streamed=streamed,
|
||||
include_usage=include_usage,
|
||||
cors_expected="Origin" in scenario.extra_headers,
|
||||
fail_fast=scenario.fail_fast,
|
||||
)
|
||||
self._record_verdict(resp, scenario, name, verdict, sent_at=sent_at, transcript=transcript, body=body)
|
||||
if verdict.benign_refusal in (BenignRefusal.RATE_LIMITED, BenignRefusal.FRONT_AT_CAPACITY):
|
||||
back_off(resp)
|
||||
raise RescheduleTask()
|
||||
if (
|
||||
not verdict.violations
|
||||
and verdict.matched_error is not None
|
||||
and verdict.matched_error.rc is not None
|
||||
and verdict.error_envelope is not None
|
||||
):
|
||||
self._cross_check_fault(scenario, verdict.matched_error, verdict.error_envelope)
|
||||
|
||||
def _record_verdict(
|
||||
self,
|
||||
resp: Any,
|
||||
scenario: ChatScenario,
|
||||
name: str,
|
||||
verdict: FrontVerdict,
|
||||
*,
|
||||
sent_at: float,
|
||||
transcript: StreamTranscript | None,
|
||||
body: Any,
|
||||
) -> None:
|
||||
"""Mark the response by its verdict: a violation, an environment condition, a benign refusal, or the outcome."""
|
||||
if verdict.violations:
|
||||
fail_contract(resp, scenario.name, verdict.violations)
|
||||
return
|
||||
if verdict.environment_condition is not None:
|
||||
fail_environment(resp, verdict.environment_condition, scenario.name)
|
||||
return
|
||||
if verdict.benign_refusal is not None:
|
||||
record_benign(self.environment, resp, name, _benign_label(verdict.benign_refusal, verdict.error_envelope))
|
||||
return
|
||||
resp.success()
|
||||
if scenario.group != ScenarioGroup.MISUSE:
|
||||
chat_mix_settings.coverage.completed_front_scenarios += 1
|
||||
self._record_timings(scenario, sent_at=sent_at, transcript=transcript, body=body)
|
||||
|
||||
def _record_transport_error(self, resp: Any, scenario: ChatScenario, name: str) -> None:
|
||||
"""Record a request that got no HTTP response.
|
||||
|
||||
An oversized upload is the exception: the front answers 413 from the Content-Length and closes the connection
|
||||
without reading the body, so the client can be reset while it is still sending and never read the 413. That is
|
||||
the expected end of the scenario, so it is a benign outcome of its own; any other transport error is a crash.
|
||||
"""
|
||||
if scenario.oversized_body and isinstance(resp.error, requests.ConnectionError):
|
||||
record_benign(self.environment, resp, name, UPLOAD_RESET_BEFORE_413)
|
||||
return
|
||||
resp.failure(f"Transport error: {resp.error}")
|
||||
|
||||
def _record_abandoned(
|
||||
self,
|
||||
resp: Any,
|
||||
scenario: ChatScenario,
|
||||
name: str,
|
||||
expectation: CompletionExpectation,
|
||||
transcript: StreamTranscript,
|
||||
*,
|
||||
received_bytes: int,
|
||||
) -> None:
|
||||
"""Record a stream this client closed mid-relay under ``<name> [abandoned]``, with the bytes received as the
|
||||
row's response length.
|
||||
|
||||
Only what the client read is checked here. What the Horde settled the cancelled generation on is read from its
|
||||
log and database, so the request id and the content characters received are logged for that comparison.
|
||||
"""
|
||||
resp.request_meta["name"] = f"{name} [{ABANDONED_LABEL}]"
|
||||
resp.request_meta["response_length"] = received_bytes
|
||||
violations = find_abandoned_stream_violations(transcript, expectation)
|
||||
if violations:
|
||||
fail_contract(resp, scenario.name, violations)
|
||||
return
|
||||
resp.success()
|
||||
chat_mix_settings.coverage.completed_front_scenarios += 1
|
||||
completion_id = str(transcript.chunks[0].get("id") or "")
|
||||
logger.info(
|
||||
"Abandoned the stream of chat request %s after %d content chunks: %d content characters, %d bytes received",
|
||||
completion_id.removeprefix(COMPLETION_ID_PREFIX),
|
||||
transcript.reply_chunks,
|
||||
transcript.content_characters,
|
||||
received_bytes,
|
||||
)
|
||||
|
||||
def _abandon_after(self, scenario: ChatScenario, *, streamed: bool) -> int | None:
|
||||
"""Return after how many content chunks this request's client closes the stream, at the run's abandon share, or
|
||||
None when it reads the stream to its end. Only a streamed request that expects a reply has content to abandon."""
|
||||
if not streamed or chat_mix_settings.stream_abandon_share <= 0:
|
||||
return None
|
||||
if not isinstance(scenario.expectation, CompletionExpectation):
|
||||
return None
|
||||
if self.rng.random() >= chat_mix_settings.stream_abandon_share:
|
||||
return None
|
||||
return chat_mix_settings.stream_abandon_chunks
|
||||
|
||||
def _decides_to_stream(self, scenario: ChatScenario) -> bool:
|
||||
"""Return whether this request of the scenario is streamed."""
|
||||
if scenario.stream_mode == StreamMode.LITE_STREAM:
|
||||
return True
|
||||
if scenario.stream_mode == StreamMode.NEVER:
|
||||
return False
|
||||
return self.rng.random() < chat_mix_settings.stream_chance
|
||||
|
||||
def _headers_for(self, authorization: AuthorizationMode) -> dict[str, str]:
|
||||
"""Return the authorization and attribution headers of a request."""
|
||||
headers = {"X-Title": STRESS_CLIENT_TITLE, "User-Agent": _config.get("client_agent", "aihorde_stress_test")}
|
||||
return {**headers, **build_authorization_headers(authorization, api_key=self.api_key)}
|
||||
|
||||
def _read_response(self, resp: Any, *, streamed: bool) -> tuple[StreamTranscript | None, Any]:
|
||||
"""Return the stream transcript of a 200 stream, or the parsed body of any other response.
|
||||
|
||||
A stream that fails mid-read fails the response as a transport error and returns ``(None, None)``.
|
||||
"""
|
||||
if not (streamed and resp.status_code == 200):
|
||||
return None, _safe_json(resp)
|
||||
transcript = StreamTranscript()
|
||||
try:
|
||||
for raw_line in resp.iter_lines(decode_unicode=True):
|
||||
transcript.feed_line(raw_line or "", time.time())
|
||||
except requests.RequestException as error:
|
||||
resp.failure(f"Stream read failed: {error}")
|
||||
return None, None
|
||||
finally:
|
||||
resp.close()
|
||||
return transcript, None
|
||||
|
||||
def _read_until_abandoned(self, resp: Any, abandon_after: int) -> tuple[StreamTranscript | None, int | None]:
|
||||
"""Read a 200 stream until ``abandon_after`` chunks carrying a reply arrived, then close the connection.
|
||||
|
||||
The close is what the front reacts to: it sees the client disconnect and cancels the native request, which
|
||||
settles a generation in progress on the text its worker posted. A response closed before its body is consumed
|
||||
closes its socket rather than returning it to the pool. A client that only stopped reading would hold the
|
||||
connection open, and the front would relay the rest of the reply into it.
|
||||
|
||||
Returns:
|
||||
The transcript and the body bytes received before the close; the whole transcript and None when the stream
|
||||
ended or carried an error first; ``(None, None)`` after a failed read, which fails the response.
|
||||
"""
|
||||
transcript = StreamTranscript()
|
||||
try:
|
||||
for line, received_bytes in _iter_stream_lines(resp):
|
||||
transcript.feed_line(line, time.time())
|
||||
if transcript.done or transcript.error_envelope is not None:
|
||||
continue
|
||||
if transcript.reply_chunks >= abandon_after:
|
||||
return transcript, received_bytes
|
||||
except requests.RequestException as error:
|
||||
resp.failure(f"Stream read failed: {error}")
|
||||
return None, None
|
||||
finally:
|
||||
resp.close()
|
||||
return transcript, None
|
||||
|
||||
def _fail_unexpected_status(self, resp: Any, subject: str) -> None:
|
||||
"""Fail a response with an unexpected status: an environment condition when it reports one, else a contract
|
||||
violation."""
|
||||
body = _safe_json(resp)
|
||||
condition = classify_environment_condition(resp.status_code, error_code_of(body), has_envelope=isinstance(body, dict))
|
||||
if condition is not None:
|
||||
fail_environment(resp, condition, subject)
|
||||
return
|
||||
fail_contract(resp, subject, [f"answered {resp.status_code}: {resp.text[:EXCERPT_CHARACTERS]}"])
|
||||
|
||||
def _record_timings(self, scenario: ChatScenario, *, sent_at: float, transcript: StreamTranscript | None, body: Any) -> None:
|
||||
"""Record the queue wait the Horde reports and, for a stream, the time to the first keepalive and content."""
|
||||
x_horde = body.get("x_horde") if isinstance(body, dict) else None
|
||||
if transcript is not None:
|
||||
record_measurement(self.environment, "[time-to-first-keepalive] stream", _elapsed(sent_at, transcript.first_keepalive_at))
|
||||
record_measurement(self.environment, "[time-to-first-content] stream", _elapsed(sent_at, transcript.first_content_at))
|
||||
finish_chunks = [chunk for chunk in transcript.chunks if chunk.get("x_horde")]
|
||||
x_horde = finish_chunks[-1]["x_horde"] if finish_chunks else None
|
||||
if isinstance(x_horde, dict) and scenario.persona is not None:
|
||||
# Anonymous requests rank last by kudos, so their queue wait is recorded apart.
|
||||
requester = " anonymous" if scenario.authorization == AuthorizationMode.ANONYMOUS else ""
|
||||
record_measurement(
|
||||
self.environment, f"[queue_wait_seconds] {scenario.persona.value}{requester}", x_horde.get("queue_wait_seconds")
|
||||
)
|
||||
|
||||
def _cross_check_fault(self, scenario: ChatScenario, matched_error: ErrorExpectation, envelope: dict[str, Any]) -> None:
|
||||
"""Read the native status of a request the front reported faulted and check it reports the same fault."""
|
||||
if not chat_mix_settings.verify_faults:
|
||||
return
|
||||
request_id = (envelope["error"].get("metadata") or {}).get("request_id")
|
||||
if not request_id:
|
||||
return
|
||||
name = f"{NATIVE_CHAT_STATUS_PATH}/[id] [fault cross-check]"
|
||||
with self.client.get(
|
||||
f"{chat_mix_settings.horde_host}{NATIVE_CHAT_STATUS_PATH}/{request_id}",
|
||||
headers=_headers(self.api_key),
|
||||
catch_response=True,
|
||||
name=name,
|
||||
) as resp:
|
||||
if resp.status_code == 0:
|
||||
resp.failure(f"Transport error: {resp.error}")
|
||||
return
|
||||
if resp.status_code == 429:
|
||||
record_benign(self.environment, resp, name, BenignRefusal.RATE_LIMITED.value)
|
||||
return
|
||||
if resp.status_code != 200:
|
||||
self._fail_unexpected_status(resp, f"{scenario.name} fault cross-check")
|
||||
return
|
||||
status_body = _safe_json(resp) or {}
|
||||
if status_body.get("faulted") is not True or status_body.get("fault_rc") != matched_error.rc:
|
||||
reported = f"faulted={status_body.get('faulted')!r} fault_rc={status_body.get('fault_rc')!r}"
|
||||
fail_contract(
|
||||
resp,
|
||||
f"{scenario.name} fault cross-check",
|
||||
[f"the native status reports {reported}, expected fault_rc {matched_error.rc!r}"],
|
||||
)
|
||||
return
|
||||
resp.success()
|
||||
|
||||
|
||||
def build_request_body(scenario: ChatScenario, *, streamed: bool, include_usage: bool) -> dict[str, Any]:
|
||||
"""Return the body sent for one request of a scenario: a copy of the scenario's body, deep enough that the stream
|
||||
keys never touch it, plus ``stream`` and ``stream_options`` when streamed.
|
||||
|
||||
Args:
|
||||
scenario: The scenario.
|
||||
streamed: Whether the request is streamed.
|
||||
include_usage: Whether a streamed request asks for the usage chunk.
|
||||
"""
|
||||
request_body = copy.deepcopy(scenario.body)
|
||||
if streamed:
|
||||
request_body["stream"] = True
|
||||
if include_usage:
|
||||
request_body["stream_options"] = {"include_usage": True}
|
||||
return request_body
|
||||
|
||||
|
||||
def find_card_violations(persona: ChatPersonaRole, card: Any) -> list[str]:
|
||||
"""Return how a persona's model card breaks what its declaration advertises.
|
||||
|
||||
Args:
|
||||
persona: The persona whose worker serves the model.
|
||||
card: The model's entry in ``/v1/models``, or None when it is missing.
|
||||
"""
|
||||
model = chat_mix_settings.model_of(persona)
|
||||
if not isinstance(card, dict):
|
||||
return [f"no model card for {model}, though its {persona.value} worker has been online for a minute"]
|
||||
violations: list[str] = []
|
||||
declared_max_length = build_persona(persona, vision=chat_mix_settings.vision).chat_max_length
|
||||
card_max_length = (card.get("top_provider") or {}).get("max_completion_tokens")
|
||||
if card_max_length != declared_max_length:
|
||||
violations.append(f"{model} card max_completion_tokens is {card_max_length!r}, declared {declared_max_length}")
|
||||
supported = set(card.get("supported_parameters") or [])
|
||||
if persona == ChatPersonaRole.LIMITED and supported != LIMITED_CARD_PARAMETERS:
|
||||
violations.append(f"{model} card supported_parameters is {sorted(supported)}, expected {sorted(LIMITED_CARD_PARAMETERS)}")
|
||||
if persona != ChatPersonaRole.LIMITED and not CONFORMANT_CARD_PARAMETERS <= supported:
|
||||
violations.append(f"{model} card lacks {sorted(CONFORMANT_CARD_PARAMETERS - supported)}")
|
||||
advertises_image = "image" in ((card.get("architecture") or {}).get("input_modalities") or [])
|
||||
declares_vision = persona == ChatPersonaRole.CONFORMANT and chat_mix_settings.vision
|
||||
if advertises_image != declares_vision:
|
||||
violations.append(f"{model} card input_modalities image={advertises_image}, declared vision={declares_vision}")
|
||||
return violations
|
||||
|
||||
|
||||
def _benign_label(benign_refusal: BenignRefusal, error_envelope: dict[str, Any] | None) -> str:
|
||||
"""Return the row label of a benign outcome. A ``queue_unavailable`` refusal carries the reason its message states,
|
||||
so a run counts the refusals for zero eligible threads apart from those for a wait over the budget."""
|
||||
if benign_refusal == BenignRefusal.QUEUE_UNAVAILABLE:
|
||||
reason = queue_unavailable_reason(error_envelope)
|
||||
if reason is not None:
|
||||
return f"{benign_refusal.value} {reason}"
|
||||
return benign_refusal.value
|
||||
|
||||
|
||||
def _iter_stream_lines(resp: Any) -> Iterator[tuple[str, int]]:
|
||||
"""Yield each line of a streamed body, without its line ending, with the body bytes received up to it.
|
||||
|
||||
The lines split as ``requests.iter_lines`` splits them; the byte count is why the body is read here instead.
|
||||
"""
|
||||
received_bytes = 0
|
||||
pending = b""
|
||||
for data in resp.iter_content(chunk_size=STREAM_READ_CHUNK_BYTES):
|
||||
received_bytes += len(data)
|
||||
lines = (pending + data).splitlines(keepends=True)
|
||||
pending = lines.pop() if lines and not lines[-1].endswith((b"\n", b"\r")) else b""
|
||||
for line in lines:
|
||||
yield line.rstrip(b"\r\n").decode("utf-8", errors="replace"), received_bytes
|
||||
if pending:
|
||||
yield pending.decode("utf-8", errors="replace"), received_bytes
|
||||
|
||||
|
||||
def _elapsed(started_at: float, ended_at: float | None) -> float | None:
|
||||
"""Return the seconds between two times, or None when the second never happened."""
|
||||
if ended_at is None:
|
||||
return None
|
||||
return ended_at - started_at
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,417 @@
|
||||
# SPDX-FileCopyrightText: 2026 Tazlin
|
||||
#
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
"""The chat mix's native client: drives the Horde's chat routes where they differ from the front.
|
||||
|
||||
It polls a chat request to the end and checks its structured reply, reuses an operation key unchanged and changed,
|
||||
asks for a dry run, reads the declaration gap from the status, and sends text requests to a chat worker's model and to
|
||||
a chat-only worker's model. Every task waits for the persona it needs to be online. An unexpected status is a contract
|
||||
violation unless it reports a deployment condition, which fails as an environment condition.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from locust import HttpUser, between, tag, task
|
||||
from locust.exception import RescheduleTask
|
||||
|
||||
from .chat_checks import CONTRACT_VIOLATION_PREFIX, BenignRefusal, classify_benign_refusal, classify_environment_condition, error_code_of
|
||||
from .chat_mix_settings import (
|
||||
back_off,
|
||||
chat_mix_settings,
|
||||
fail_contract,
|
||||
fail_environment,
|
||||
fire_check_failure,
|
||||
fire_row,
|
||||
record_benign,
|
||||
record_measurement,
|
||||
)
|
||||
from .chat_worker_script import ChatPersonaRole, nonce_marker
|
||||
from .helpers import _headers, _pick_requestor_key, _safe_json
|
||||
|
||||
__all__ = ["ChatMixNativeClient"]
|
||||
|
||||
CHAT_TAG: str = "chat"
|
||||
"""The tag every chat mix task carries."""
|
||||
|
||||
NATIVE_CHAT_ASYNC_PATH: str = "/api/v2/generate/chat/async"
|
||||
"""The native route a chat request is queued on."""
|
||||
|
||||
NATIVE_CHAT_STATUS_PATH: str = "/api/v2/generate/chat/status"
|
||||
"""The native chat status and cancel route, followed by the request id."""
|
||||
|
||||
TEXT_ASYNC_PATH: str = "/api/v2/generate/text/async"
|
||||
"""The native text request route."""
|
||||
|
||||
TEXT_STATUS_PATH: str = "/api/v2/generate/text/status"
|
||||
"""The native text status and cancel route, followed by the request id."""
|
||||
|
||||
NATIVE_STATUS_POLL_SECONDS: float = 7.0
|
||||
"""The pause between two native status reads; the status routes allow 10 reads a minute per path."""
|
||||
|
||||
NATIVE_COMPLETION_DEADLINE_SECONDS: float = 300.0
|
||||
"""How long a native client polls before it gives up and cancels."""
|
||||
|
||||
NATIVE_GAP_POLLS: int = 3
|
||||
"""How many status reads a native client makes looking for the declaration gap before it cancels."""
|
||||
|
||||
NATIVE_SUBMIT_SPACING_SECONDS: float = 0.6
|
||||
"""The pause between back-to-back native submits; the generate routes allow 2 a second per key."""
|
||||
|
||||
EXCERPT_CHARACTERS: int = 200
|
||||
"""How much of a body a failure quotes."""
|
||||
|
||||
|
||||
class ChatMixNativeClient(HttpUser):
|
||||
"""Drive the native chat routes where they differ from the front: polling, operation keys, dry runs, the
|
||||
declaration gap on the status, and text requests beside chat workers."""
|
||||
|
||||
weight = 1
|
||||
fixed_count = 0
|
||||
wait_time = between(2, 5)
|
||||
|
||||
def on_start(self) -> None:
|
||||
"""Take a requestor key; the client talks to the run's ``--host``."""
|
||||
self.api_key = _pick_requestor_key()
|
||||
|
||||
@tag(CHAT_TAG, "chat-native")
|
||||
@task(4)
|
||||
def native_round_trip(self) -> None:
|
||||
"""Queue a chat request, poll its status to the end, and check the structured reply and its accounting."""
|
||||
if not chat_mix_settings.is_persona_online(ChatPersonaRole.CONFORMANT):
|
||||
return
|
||||
scenario_name = "native_round_trip"
|
||||
nonce = uuid.uuid4().hex[:12]
|
||||
body = {
|
||||
"model": chat_mix_settings.models.conformant,
|
||||
"messages": [{"role": "user", "content": f"Say hello. {nonce_marker(nonce)}"}],
|
||||
"max_tokens": 64,
|
||||
}
|
||||
submitted_at = time.time()
|
||||
request_id = self._submit_chat(scenario_name, body, expected_status=202)
|
||||
if request_id is None:
|
||||
return
|
||||
self._check_route_separation(request_id)
|
||||
status_body = self._poll_until_finished(NATIVE_CHAT_STATUS_PATH, request_id, scenario_name)
|
||||
if status_body is None:
|
||||
return
|
||||
record_measurement(self.environment, "[submit-to-done] native chat", time.time() - submitted_at)
|
||||
generations = status_body.get("generations") or []
|
||||
if generations:
|
||||
record_measurement(self.environment, "[queue_wait_seconds] native chat", generations[0].get("queue_wait_seconds"))
|
||||
violations = find_finished_chat_violations(request_id, status_body, nonce=nonce)
|
||||
if violations:
|
||||
self._fail_check(scenario_name, violations)
|
||||
|
||||
@tag(CHAT_TAG, "chat-native")
|
||||
@task(2)
|
||||
def native_operation_key(self) -> None:
|
||||
"""Submit with an operation key, repeat it unchanged (same request, 200, nothing new queued) and changed
|
||||
(``OperationKeyConflict``), then cancel."""
|
||||
if not chat_mix_settings.is_persona_online(ChatPersonaRole.CONFORMANT):
|
||||
return
|
||||
scenario_name = "native_operation_key"
|
||||
operation_key = str(uuid.uuid4())
|
||||
body = {
|
||||
"model": chat_mix_settings.models.conformant,
|
||||
"messages": [{"role": "user", "content": f"Say hello. {nonce_marker(operation_key[:8])}"}],
|
||||
"max_tokens": 64,
|
||||
"horde": {"operation_key": operation_key},
|
||||
}
|
||||
request_id = self._submit_chat(scenario_name, body, expected_status=202)
|
||||
if request_id is None:
|
||||
return
|
||||
time.sleep(NATIVE_SUBMIT_SPACING_SECONDS)
|
||||
repeated_id = self._submit_chat(f"{scenario_name} repeat", body, expected_status=200)
|
||||
if repeated_id is not None and repeated_id != request_id:
|
||||
self._fail_check(f"{scenario_name} repeat", [f"the repeated operation key returned request {repeated_id}, not {request_id}"])
|
||||
time.sleep(NATIVE_SUBMIT_SPACING_SECONDS)
|
||||
self._submit_expecting_refusal(f"{scenario_name} changed", {**body, "max_tokens": 65}, expected_rc="OperationKeyConflict")
|
||||
self._cancel(NATIVE_CHAT_STATUS_PATH, request_id, scenario_name)
|
||||
|
||||
@tag(CHAT_TAG, "chat-native")
|
||||
@task(1)
|
||||
def native_dry_run(self) -> None:
|
||||
"""Ask for a quote: 200 with the kudos and no request id, so nothing is queued. The model has to be known, which
|
||||
an unlisted model is only while its worker is online."""
|
||||
if not chat_mix_settings.is_persona_online(ChatPersonaRole.CONFORMANT):
|
||||
return
|
||||
name = f"{NATIVE_CHAT_ASYNC_PATH} [native_dry_run]"
|
||||
body = {
|
||||
"model": chat_mix_settings.models.conformant,
|
||||
"messages": [{"role": "user", "content": "Quote this."}],
|
||||
"max_tokens": 64,
|
||||
"horde": {"dry_run": True},
|
||||
}
|
||||
with self.client.post(NATIVE_CHAT_ASYNC_PATH, json=body, headers=_headers(self.api_key), catch_response=True, name=name) as resp:
|
||||
if self._handled_as_benign_or_environment(resp, name, "native_dry_run"):
|
||||
return
|
||||
quote = _safe_json(resp) or {}
|
||||
if resp.status_code != 200:
|
||||
fail_contract(resp, "native_dry_run", [f"a dry run answered {resp.status_code}: {resp.text[:EXCERPT_CHARACTERS]}"])
|
||||
return
|
||||
if "id" in quote or not isinstance(quote.get("kudos"), int | float):
|
||||
fail_contract(
|
||||
resp, "native_dry_run", [f"a dry run returned {json.dumps(quote)[:EXCERPT_CHARACTERS]}, expected kudos and no id"]
|
||||
)
|
||||
return
|
||||
resp.success()
|
||||
|
||||
@tag(CHAT_TAG, "chat-native")
|
||||
@task(1)
|
||||
def native_declaration_gap(self) -> None:
|
||||
"""Queue a forced function on the limited model under ``strict``: the status reports the declaration gap
|
||||
naming ``tool_choice_enforced``. Then cancel."""
|
||||
if not chat_mix_settings.is_persona_online(ChatPersonaRole.LIMITED):
|
||||
return
|
||||
scenario_name = "native_declaration_gap"
|
||||
body = {
|
||||
"model": chat_mix_settings.models.limited,
|
||||
"messages": [{"role": "user", "content": f"Hello. {nonce_marker(uuid.uuid4().hex[:12])}"}],
|
||||
"max_tokens": 64,
|
||||
"temperature": 0.2,
|
||||
"tools": [{"type": "function", "function": {"name": "get_time", "parameters": {"type": "object", "properties": {}}}}],
|
||||
"tool_choice": {"type": "function", "function": {"name": "get_time"}},
|
||||
}
|
||||
request_id = self._submit_chat(scenario_name, body, expected_status=202)
|
||||
if request_id is None:
|
||||
return
|
||||
name = f"{NATIVE_CHAT_STATUS_PATH}/[id] [{scenario_name}]"
|
||||
found_gap = False
|
||||
last_reason: Any = None
|
||||
for _poll in range(NATIVE_GAP_POLLS):
|
||||
time.sleep(NATIVE_STATUS_POLL_SECONDS)
|
||||
status_body = self._read_status(NATIVE_CHAT_STATUS_PATH, request_id, name)
|
||||
if status_body is None:
|
||||
continue
|
||||
last_reason = status_body.get("impossible_reason")
|
||||
if _reports_tool_choice_gap(last_reason):
|
||||
found_gap = status_body.get("is_possible") is False
|
||||
break
|
||||
if not found_gap:
|
||||
self._fail_check(
|
||||
scenario_name,
|
||||
[f"the status never reported the declaration gap; last impossible_reason {json.dumps(last_reason)[:EXCERPT_CHARACTERS]}"],
|
||||
)
|
||||
self._cancel(NATIVE_CHAT_STATUS_PATH, request_id, scenario_name)
|
||||
|
||||
@tag(CHAT_TAG, "chat-native", "text")
|
||||
@task(1)
|
||||
def native_text_beside_chat(self) -> None:
|
||||
"""Queue a text request on the limited model: its chat worker also accepts text jobs, pops it without a
|
||||
``chat`` key and submits plain generation text, so text is unaffected by chat."""
|
||||
if not chat_mix_settings.is_persona_online(ChatPersonaRole.LIMITED):
|
||||
return
|
||||
scenario_name = "native_text_beside_chat"
|
||||
request_id = self._submit_text(scenario_name, chat_mix_settings.models.limited)
|
||||
if request_id is None:
|
||||
return
|
||||
status_body = self._poll_until_finished(TEXT_STATUS_PATH, request_id, scenario_name)
|
||||
if status_body is None:
|
||||
return
|
||||
generations = status_body.get("generations") or []
|
||||
if status_body.get("faulted") or not generations or not generations[0].get("text"):
|
||||
self._fail_check(
|
||||
scenario_name, [f"the text request finished without a generation: {json.dumps(status_body)[:EXCERPT_CHARACTERS]}"]
|
||||
)
|
||||
elif "message" in generations[0]:
|
||||
self._fail_check(scenario_name, ["a text generation carries a chat message"])
|
||||
|
||||
@tag(CHAT_TAG, "chat-native", "text")
|
||||
@task(1)
|
||||
def native_text_to_chat_only_model(self) -> None:
|
||||
"""Queue a text request on the conformant model, whose only worker sets ``prompt_jobs: false``: no worker can
|
||||
serve it, so the status reports it impossible. Then cancel."""
|
||||
if not chat_mix_settings.is_persona_online(ChatPersonaRole.CONFORMANT):
|
||||
return
|
||||
scenario_name = "native_text_to_chat_only_model"
|
||||
request_id = self._submit_text(scenario_name, chat_mix_settings.models.conformant)
|
||||
if request_id is None:
|
||||
return
|
||||
time.sleep(NATIVE_STATUS_POLL_SECONDS)
|
||||
status_body = self._read_status(TEXT_STATUS_PATH, request_id, f"{TEXT_STATUS_PATH}/[id] [{scenario_name}]")
|
||||
if status_body is not None and (status_body.get("is_possible") is not False or status_body.get("done")):
|
||||
reported = f"is_possible={status_body.get('is_possible')!r} done={status_body.get('done')!r}"
|
||||
self._fail_check(scenario_name, [f"a text request for a model only a chat-only worker serves reports {reported}"])
|
||||
self._cancel(TEXT_STATUS_PATH, request_id, scenario_name)
|
||||
|
||||
def _submit_chat(self, scenario_name: str, body: dict[str, Any], *, expected_status: int) -> str | None:
|
||||
"""Submit a chat request and return its id, or None after a benign refusal or a failure."""
|
||||
name = f"{NATIVE_CHAT_ASYNC_PATH} [{scenario_name}]"
|
||||
with self.client.post(NATIVE_CHAT_ASYNC_PATH, json=body, headers=_headers(self.api_key), catch_response=True, name=name) as resp:
|
||||
if self._handled_as_benign_or_environment(resp, name, scenario_name):
|
||||
return None
|
||||
submit_body = _safe_json(resp) or {}
|
||||
request_id = submit_body.get("id")
|
||||
if resp.status_code != expected_status:
|
||||
fail_contract(
|
||||
resp,
|
||||
scenario_name,
|
||||
[f"the submit answered {resp.status_code}, expected {expected_status}: {resp.text[:EXCERPT_CHARACTERS]}"],
|
||||
)
|
||||
return None
|
||||
if not request_id or not isinstance(submit_body.get("kudos"), int | float):
|
||||
fail_contract(
|
||||
resp, scenario_name, [f"the submit returned {json.dumps(submit_body)[:EXCERPT_CHARACTERS]}, expected an id and kudos"]
|
||||
)
|
||||
return None
|
||||
resp.success()
|
||||
return str(request_id)
|
||||
|
||||
def _submit_expecting_refusal(self, scenario_name: str, body: dict[str, Any], *, expected_rc: str) -> None:
|
||||
"""Submit a chat request the Horde has to refuse with ``expected_rc``."""
|
||||
name = f"{NATIVE_CHAT_ASYNC_PATH} [{scenario_name}]"
|
||||
with self.client.post(NATIVE_CHAT_ASYNC_PATH, json=body, headers=_headers(self.api_key), catch_response=True, name=name) as resp:
|
||||
if self._handled_as_benign_or_environment(resp, name, scenario_name):
|
||||
return
|
||||
if resp.status_code != 400 or error_code_of(_safe_json(resp)) != expected_rc:
|
||||
fail_contract(
|
||||
resp, scenario_name, [f"expected 400 {expected_rc}, got {resp.status_code}: {resp.text[:EXCERPT_CHARACTERS]}"]
|
||||
)
|
||||
return
|
||||
resp.success()
|
||||
|
||||
def _submit_text(self, scenario_name: str, model: str) -> str | None:
|
||||
"""Submit a short text request for ``model`` and return its id, or None after a refusal."""
|
||||
name = f"{TEXT_ASYNC_PATH} [{scenario_name}]"
|
||||
body = {
|
||||
"prompt": f"Once upon a time ({uuid.uuid4().hex[:12]})",
|
||||
"params": {"max_length": 32, "max_context_length": 1024},
|
||||
"models": [model],
|
||||
}
|
||||
with self.client.post(TEXT_ASYNC_PATH, json=body, headers=_headers(self.api_key), catch_response=True, name=name) as resp:
|
||||
if self._handled_as_benign_or_environment(resp, name, scenario_name):
|
||||
return None
|
||||
request_id = (_safe_json(resp) or {}).get("id")
|
||||
if resp.status_code != 202 or not request_id:
|
||||
fail_contract(resp, scenario_name, [f"the text submit answered {resp.status_code}: {resp.text[:EXCERPT_CHARACTERS]}"])
|
||||
return None
|
||||
resp.success()
|
||||
return str(request_id)
|
||||
|
||||
def _check_route_separation(self, chat_request_id: str) -> None:
|
||||
"""Check a chat request's id is not found on the text status route."""
|
||||
name = f"{TEXT_STATUS_PATH}/[id] [chat id]"
|
||||
with self.client.get(
|
||||
f"{TEXT_STATUS_PATH}/{chat_request_id}", headers=_headers(self.api_key), catch_response=True, name=name
|
||||
) as resp:
|
||||
if self._handled_as_benign_or_environment(resp, name, "route separation", back_off_after_rate_limit=False):
|
||||
return
|
||||
if resp.status_code != 404:
|
||||
fail_contract(resp, "route separation", [f"a chat id on the text status route answered {resp.status_code}, expected 404"])
|
||||
return
|
||||
resp.success()
|
||||
|
||||
def _poll_until_finished(self, status_path: str, request_id: str, scenario_name: str) -> dict[str, Any] | None:
|
||||
"""Poll a status until it is done or faulted and return it; cancel and return None at the deadline."""
|
||||
name = f"{status_path}/[id] [{scenario_name}]"
|
||||
deadline = time.time() + NATIVE_COMPLETION_DEADLINE_SECONDS
|
||||
while time.time() < deadline:
|
||||
time.sleep(NATIVE_STATUS_POLL_SECONDS)
|
||||
status_body = self._read_status(status_path, request_id, name)
|
||||
if status_body is not None and (status_body.get("done") or status_body.get("faulted")):
|
||||
return status_body
|
||||
fire_row(
|
||||
self.environment,
|
||||
request_type="GET",
|
||||
name=f"{name} [benign gave up]",
|
||||
response_time_ms=NATIVE_COMPLETION_DEADLINE_SECONDS * 1000,
|
||||
)
|
||||
self._cancel(status_path, request_id, scenario_name)
|
||||
return None
|
||||
|
||||
def _read_status(self, status_path: str, request_id: str, name: str) -> dict[str, Any] | None:
|
||||
"""Read one status and return it, or None after a rate limit or a failure."""
|
||||
with self.client.get(f"{status_path}/{request_id}", headers=_headers(self.api_key), catch_response=True, name=name) as resp:
|
||||
if self._handled_as_benign_or_environment(resp, name, name, back_off_after_rate_limit=False):
|
||||
return None
|
||||
if resp.status_code != 200:
|
||||
fail_contract(resp, name, [f"the status answered {resp.status_code}: {resp.text[:EXCERPT_CHARACTERS]}"])
|
||||
return None
|
||||
resp.success()
|
||||
return _safe_json(resp)
|
||||
|
||||
def _cancel(self, status_path: str, request_id: str, scenario_name: str) -> None:
|
||||
"""Cancel a request; a request already gone (404) is fine."""
|
||||
name = f"{status_path}/[id] [cancel {scenario_name}]"
|
||||
with self.client.delete(f"{status_path}/{request_id}", headers=_headers(self.api_key), catch_response=True, name=name) as resp:
|
||||
if self._handled_as_benign_or_environment(resp, name, f"cancel {scenario_name}", back_off_after_rate_limit=False):
|
||||
return
|
||||
if resp.status_code not in (200, 404):
|
||||
fail_contract(
|
||||
resp, f"cancel {scenario_name}", [f"the cancel answered {resp.status_code}: {resp.text[:EXCERPT_CHARACTERS]}"]
|
||||
)
|
||||
return
|
||||
resp.success()
|
||||
|
||||
def _fail_check(self, scenario_name: str, violations: list[str]) -> None:
|
||||
"""Record a contract violation found across several requests on a row of its own."""
|
||||
fire_check_failure(
|
||||
self.environment,
|
||||
name=f"[native] {scenario_name}",
|
||||
message=f"{CONTRACT_VIOLATION_PREFIX} [{scenario_name}]: {'; '.join(violations[:4])}",
|
||||
)
|
||||
|
||||
def _handled_as_benign_or_environment(self, resp: Any, name: str, subject: str, *, back_off_after_rate_limit: bool = True) -> bool:
|
||||
"""Return True after recording a transport failure, a deployment condition or a benign refusal for ``resp``.
|
||||
|
||||
Raises:
|
||||
RescheduleTask: After a rate limit's back-off, when ``back_off_after_rate_limit`` is set.
|
||||
"""
|
||||
if resp.status_code == 0:
|
||||
resp.failure(f"Transport error: {resp.error}")
|
||||
return True
|
||||
body = _safe_json(resp)
|
||||
error_code = error_code_of(body)
|
||||
condition = classify_environment_condition(resp.status_code, error_code, has_envelope=isinstance(body, dict))
|
||||
if condition is not None:
|
||||
fail_environment(resp, condition, subject)
|
||||
return True
|
||||
benign_refusal = classify_benign_refusal(resp.status_code, error_code)
|
||||
if benign_refusal is None:
|
||||
return False
|
||||
record_benign(self.environment, resp, name, benign_refusal.value)
|
||||
if back_off_after_rate_limit and benign_refusal == BenignRefusal.RATE_LIMITED:
|
||||
back_off(resp)
|
||||
raise RescheduleTask()
|
||||
return True
|
||||
|
||||
|
||||
def _reports_tool_choice_gap(impossible_reason: Any) -> bool:
|
||||
"""Return whether an ``impossible_reason`` is a declaration gap naming ``tool_choice_enforced``."""
|
||||
if not isinstance(impossible_reason, dict) or impossible_reason.get("kind") != "declaration_gap":
|
||||
return False
|
||||
return "tool_choice_enforced" in (impossible_reason.get("missing_declarations") or [])
|
||||
|
||||
|
||||
def find_finished_chat_violations(request_id: str, status_body: dict[str, Any], *, nonce: str) -> list[str]:
|
||||
"""Return how a finished native chat status breaks the contract: one generation whose structured reply answers this
|
||||
request, finished with ``stop``, charged on the output estimate, with its queue wait.
|
||||
|
||||
Args:
|
||||
request_id: The request's id, for the message.
|
||||
status_body: The finished status.
|
||||
nonce: The request's nonce, which the reply has to carry.
|
||||
"""
|
||||
generations = status_body.get("generations") or []
|
||||
if status_body.get("faulted") or len(generations) != 1:
|
||||
return [f"request {request_id} finished faulted={status_body.get('faulted')!r} with {len(generations)} generations"]
|
||||
generation = generations[0]
|
||||
message = generation.get("message") or {}
|
||||
violations: list[str] = []
|
||||
content = message.get("content")
|
||||
if message.get("role") != "assistant" or not isinstance(content, str) or nonce not in content:
|
||||
violations.append(f"the generation's message is not this request's reply: {json.dumps(message)[:EXCERPT_CHARACTERS]}")
|
||||
if generation.get("finish_reason") != "stop":
|
||||
violations.append(f"the generation's finish_reason is {generation.get('finish_reason')!r}")
|
||||
if generation.get("charge_basis") != "output_estimate":
|
||||
violations.append(f"the generation's charge_basis is {generation.get('charge_basis')!r}")
|
||||
queue_wait = generation.get("queue_wait_seconds")
|
||||
if not isinstance(queue_wait, int | float) or queue_wait < 0:
|
||||
violations.append(f"the generation's queue_wait_seconds is {queue_wait!r}")
|
||||
return violations
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,38 @@
|
||||
# SPDX-FileCopyrightText: 2026 Tazlin
|
||||
#
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
"""Routes, timeouts and the bridge agent every chat population of the suite shares: the scale demo in
|
||||
``users/chat.py`` and the chat contract mix."""
|
||||
|
||||
__all__ = [
|
||||
"CHAT_COMPLETIONS_PATH",
|
||||
"CONNECT_TIMEOUT_SECONDS",
|
||||
"STREAM_READ_TIMEOUT_SECONDS",
|
||||
"TEXT_POP_PATH",
|
||||
"TEXT_PROGRESS_PATH",
|
||||
"TEXT_SUBMIT_PATH",
|
||||
"WORKER_BRIDGE_AGENT",
|
||||
]
|
||||
|
||||
CHAT_COMPLETIONS_PATH: str = "/v1/chat/completions"
|
||||
"""The front's chat completions route."""
|
||||
|
||||
TEXT_POP_PATH: str = "/api/v2/generate/text/pop"
|
||||
"""The Horde route a text worker pops jobs from; chat jobs come through it when the pop opts in."""
|
||||
|
||||
TEXT_SUBMIT_PATH: str = "/api/v2/generate/text/submit"
|
||||
"""The Horde route a text worker submits generations to."""
|
||||
|
||||
TEXT_PROGRESS_PATH: str = "/api/v2/generate/text/progress"
|
||||
"""The Horde route a chat worker posts content deltas to while it generates a job whose client streams."""
|
||||
|
||||
CONNECT_TIMEOUT_SECONDS: float = 10.0
|
||||
"""How long a client waits to connect; thousands of clients connecting at once can fill the listen backlog."""
|
||||
|
||||
STREAM_READ_TIMEOUT_SECONDS: float = 60.0
|
||||
"""How long a streamed client waits between two lines; the front sends a keepalive every few seconds, so a minute
|
||||
of silence is a stalled stream."""
|
||||
|
||||
WORKER_BRIDGE_AGENT: str = "KoboldAI Client:1.19.2-stress-chat:https://github.com/koboldai/koboldai-client"
|
||||
"""The bridge agent simulated chat workers report."""
|
||||
@@ -0,0 +1,543 @@
|
||||
# SPDX-FileCopyrightText: 2026 Tazlin
|
||||
#
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
"""Decide what a simulated chat worker declares and replies, and check the jobs it receives, without a GPU.
|
||||
|
||||
A client scenario that needs a particular worker behaviour writes a reply marker, ``[reply:<script>]``, into its last
|
||||
user message; the worker reads it from the popped payload and answers as scripted, so the client knows the outcome it
|
||||
caused. Without a marker the worker answers as a well-behaved backend would. Every reply echoes the request's nonce
|
||||
(``[nonce:<hex>]`` in a user message) into its text, JSON document or call arguments, so a client detects a reply
|
||||
delivered to the wrong request. The pop and submit bodies come from ``tests/integration/_fake_chat_worker.py``, the fake
|
||||
worker the integration suite uses, so both speak one contract.
|
||||
|
||||
- [`ChatPersonaRole`][tests.stress.locustsuite.chat_worker_script.ChatPersonaRole] and
|
||||
[`build_persona`][tests.stress.locustsuite.chat_worker_script.build_persona]: the three declarations the suite runs, a
|
||||
conformant worker, a limited one that leaves out guarantees and features and caps the tool temperature, and a flaky
|
||||
one that sometimes breaks the request's contract.
|
||||
- [`nonce_marker`][tests.stress.locustsuite.chat_worker_script.nonce_marker] and
|
||||
[`find_request_nonce`][tests.stress.locustsuite.chat_worker_script.find_request_nonce]: how a client tags a request
|
||||
and the worker finds the tag.
|
||||
- [`build_scripted_reply`][tests.stress.locustsuite.chat_worker_script.build_scripted_reply] and
|
||||
[`build_misbehaving_reply`][tests.stress.locustsuite.chat_worker_script.build_misbehaving_reply]: the reply to a popped
|
||||
chat payload.
|
||||
- [`find_job_payload_violations`][tests.stress.locustsuite.chat_worker_script.find_job_payload_violations]: what a popped
|
||||
payload breaks of the routing and dispatch rules, checked against the worker's own declaration.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
from horde_sdk.generation_parameters.text.chat_limits import CHAT_SAMPLING_DEFAULTS
|
||||
from horde_sdk.generation_parameters.text.chat_policy import ChatFeature, ChatGuarantee
|
||||
from horde_sdk.worker.chat_declaration import ChatFaultReason, ChatFinishReason, ChatImageTransport
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_PERSONA_MAX_LENGTH",
|
||||
"ChatPersona",
|
||||
"ChatPersonaRole",
|
||||
"ReplyScript",
|
||||
"WorkerReply",
|
||||
"build_misbehaving_reply",
|
||||
"build_persona",
|
||||
"build_scripted_reply",
|
||||
"find_job_payload_violations",
|
||||
"find_reply_script",
|
||||
"find_request_nonce",
|
||||
"nonce_marker",
|
||||
"reply_marker",
|
||||
]
|
||||
|
||||
|
||||
class ReplyScript(StrEnum):
|
||||
"""A reply a client asks the simulated worker for with ``[reply:<value>]`` in its last user message."""
|
||||
|
||||
TWO_CALLS = "two_calls"
|
||||
"""Two calls to the first offered tool, with distinct ids; one survives ``parallel_tool_calls: false``."""
|
||||
UNDECLARED = "undeclared"
|
||||
"""A call to a function the request did not offer: refused as ``ChatSubmitInvalid`` under every mode."""
|
||||
DUP_IDS = "dup_ids"
|
||||
"""Two calls sharing one id: refused as ``ChatSubmitInvalid`` under every mode."""
|
||||
PROSE = "prose"
|
||||
"""Prose with no call, which breaks a forced ``tool_choice`` under ``strict``."""
|
||||
NONCONFORMANT = "nonconformant"
|
||||
"""``{"x": 0}``, which breaks any schema with other required keys under ``strict``."""
|
||||
TEXT_AND_CALL = "text_and_call"
|
||||
"""Text and one call in the same reply."""
|
||||
FAULT_GENERATION_ERROR = "fault_generation_error"
|
||||
"""A faulted submit with ``fault_reason`` ``generation_error`` instead of a reply."""
|
||||
|
||||
|
||||
REPLY_MARKER_PATTERN: re.Pattern[str] = re.compile(r"\[reply:([a-z_]+)\]")
|
||||
"""Find the reply marker in a user message."""
|
||||
|
||||
NONCE_MARKER_PATTERN: re.Pattern[str] = re.compile(r"\[nonce:([0-9a-f]+)\]")
|
||||
"""Find a request's nonce in a user message."""
|
||||
|
||||
PLAIN_REPLY: str = "Sunny and mild, a stress test reply."
|
||||
"""The text of a well-behaved reply to a request without tools or a response format, before the nonce."""
|
||||
|
||||
PROSE_REPLY: str = "Sunny."
|
||||
"""The text of a scripted prose reply, before the nonce."""
|
||||
|
||||
TEXT_BEFORE_CALL: str = "Checking."
|
||||
"""The text a ``text_and_call`` reply carries beside its call, before the nonce."""
|
||||
|
||||
NOT_JSON_REPLY: str = "This is not JSON."
|
||||
"""The content a misbehaving worker returns for a response format, before the nonce."""
|
||||
|
||||
UNOFFERED_FUNCTION_NAME: str = "not_offered"
|
||||
"""The function name an ``undeclared`` reply calls."""
|
||||
|
||||
VISION_OK_PREFIX: str = "vision-ok"
|
||||
"""The start of the reply to a job whose images all arrived as JPEG data URLs; the vision scenario requires it."""
|
||||
|
||||
VISION_UNEXPECTED_PREFIX: str = "vision-unexpected"
|
||||
"""The start of the reply to a job whose images arrived in another shape."""
|
||||
|
||||
JPEG_DATA_URL_PREFIX: str = "data:image/jpeg;base64,"
|
||||
"""The data URL prefix a ``base64`` transport worker receives images under; the Horde stores every image as JPEG."""
|
||||
|
||||
JPEG_MAGIC: bytes = b"\xff\xd8\xff"
|
||||
"""The first bytes of every JPEG file."""
|
||||
|
||||
WORKER_MESSAGE_ROLES: frozenset[str] = frozenset({"system", "user", "assistant", "tool"})
|
||||
"""The roles a chat job carries; ``developer`` arrives as ``system``."""
|
||||
|
||||
DISPATCHED_TOOL_CHOICES: frozenset[str] = frozenset({"auto", "required"})
|
||||
"""The ``tool_choice`` values a worker may receive: ``none`` is never sent and a function choice arrives as ``required``."""
|
||||
|
||||
DEFAULT_PERSONA_MAX_LENGTH: int = 512
|
||||
"""The chat output limit every persona declares; a ``max_tokens`` above it has no worker."""
|
||||
|
||||
DEFAULT_PERSONA_MAX_CONTEXT_LENGTH: int = 8192
|
||||
"""The chat context limit every persona declares."""
|
||||
|
||||
LIMITED_TOOLS_MAX_TEMPERATURE: float = 0.5
|
||||
"""The tool temperature limit of the limited persona; a tools request above it has no worker under either mode."""
|
||||
|
||||
|
||||
class ChatPersonaRole(StrEnum):
|
||||
"""Which declaration a simulated chat worker pops with."""
|
||||
|
||||
CONFORMANT = "conformant"
|
||||
"""Every feature and guarantee, chat jobs only; answers every scenario as scripted."""
|
||||
LIMITED = "limited"
|
||||
"""``seed`` and ``tools`` with only ``prompt_overflow_refused``, a tool temperature limit, and text jobs as well."""
|
||||
FLAKY = "flaky"
|
||||
"""Every feature but ``vision`` and ``tools_with_response_format`` and every guarantee, chat jobs only; sometimes breaks a reply's
|
||||
contract so refusal accounting and the fault cap run at volume."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ChatPersona:
|
||||
"""Represents the chat declaration one simulated worker pops with."""
|
||||
|
||||
role: ChatPersonaRole
|
||||
"""The persona."""
|
||||
features: tuple[str, ...]
|
||||
"""The declared features."""
|
||||
guarantees: tuple[str, ...]
|
||||
"""The declared guarantees."""
|
||||
tools_max_temperature: float | None
|
||||
"""The declared tool temperature limit, or None for none."""
|
||||
prompt_jobs: bool
|
||||
"""Whether the worker also accepts text jobs."""
|
||||
image_transports: tuple[str, ...]
|
||||
"""The declared image transports; empty without vision."""
|
||||
chat_max_length: int = DEFAULT_PERSONA_MAX_LENGTH
|
||||
"""The declared chat output limit."""
|
||||
chat_max_context_length: int = DEFAULT_PERSONA_MAX_CONTEXT_LENGTH
|
||||
"""The declared chat context limit."""
|
||||
|
||||
|
||||
def build_persona(role: ChatPersonaRole, *, vision: bool) -> ChatPersona:
|
||||
"""Create the declaration of a persona.
|
||||
|
||||
Args:
|
||||
role: The persona.
|
||||
vision: Whether the conformant persona also declares ``vision`` with the ``base64`` transport. The base64
|
||||
transport lets a worker without a GPU or object storage access check the image it receives.
|
||||
"""
|
||||
every_guarantee = tuple(guarantee.value for guarantee in ChatGuarantee)
|
||||
if role == ChatPersonaRole.LIMITED:
|
||||
return ChatPersona(
|
||||
role=role,
|
||||
features=(ChatFeature.SEED.value, ChatFeature.TOOLS.value),
|
||||
guarantees=(ChatGuarantee.PROMPT_OVERFLOW_REFUSED.value,),
|
||||
tools_max_temperature=LIMITED_TOOLS_MAX_TEMPERATURE,
|
||||
prompt_jobs=True,
|
||||
image_transports=(),
|
||||
)
|
||||
if role == ChatPersonaRole.FLAKY:
|
||||
flaky_features = tuple(
|
||||
feature.value for feature in ChatFeature if feature not in (ChatFeature.VISION, ChatFeature.TOOLS_WITH_RESPONSE_FORMAT)
|
||||
)
|
||||
return ChatPersona(
|
||||
role=role,
|
||||
features=flaky_features,
|
||||
guarantees=every_guarantee,
|
||||
tools_max_temperature=None,
|
||||
prompt_jobs=False,
|
||||
image_transports=(),
|
||||
)
|
||||
conformant_features = tuple(feature.value for feature in ChatFeature if vision or feature != ChatFeature.VISION)
|
||||
return ChatPersona(
|
||||
role=role,
|
||||
features=conformant_features,
|
||||
guarantees=every_guarantee,
|
||||
tools_max_temperature=None,
|
||||
prompt_jobs=False,
|
||||
image_transports=(ChatImageTransport.BASE64.value,) if vision else (),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkerReply:
|
||||
"""Represents what a simulated worker submits for one chat job."""
|
||||
|
||||
content: str | None
|
||||
"""The reply text; None for a reply that only calls tools, and for a fault."""
|
||||
tool_calls: list[dict[str, Any]] | None
|
||||
"""The calls, or None."""
|
||||
finish_reason: str
|
||||
"""The submitted ``finish_reason``."""
|
||||
fault_reason: str | None = None
|
||||
"""When set, the worker faults the job with this ``fault_reason`` instead of submitting a reply."""
|
||||
breaks_strict_contract: bool = False
|
||||
"""True when a ``strict`` request's reply verification refuses it as ``ChatSubmitNonconformant``."""
|
||||
breaks_every_contract: bool = False
|
||||
"""True when reply verification refuses it as ``ChatSubmitInvalid`` under every conformance mode."""
|
||||
|
||||
|
||||
def reply_marker(script: ReplyScript) -> str:
|
||||
"""Return the marker a client writes into its last user message to ask for ``script``."""
|
||||
return f"[reply:{script.value}]"
|
||||
|
||||
|
||||
def nonce_marker(nonce: str) -> str:
|
||||
"""Return the tag a client writes into a user message so the worker can echo ``nonce`` into its reply.
|
||||
|
||||
Args:
|
||||
nonce: Lowercase hexadecimal, unique to the request.
|
||||
"""
|
||||
return f"[nonce:{nonce}]"
|
||||
|
||||
|
||||
def find_request_nonce(messages: Sequence[Any]) -> str:
|
||||
"""Return the nonce of the last user message that carries one, or an empty string when none does."""
|
||||
for message in reversed(messages):
|
||||
if isinstance(message, dict) and message.get("role") == "user":
|
||||
match = NONCE_MARKER_PATTERN.search(_message_text(message))
|
||||
if match is not None:
|
||||
return match.group(1)
|
||||
return ""
|
||||
|
||||
|
||||
def build_scripted_reply(payload: Mapping[str, Any]) -> WorkerReply:
|
||||
"""Create the reply to a popped chat payload: the marker's script when the last user message has one, else the
|
||||
reply a well-behaved backend gives. Every reply carries the request's nonce.
|
||||
|
||||
Args:
|
||||
payload: The popped job's ``payload``.
|
||||
"""
|
||||
messages = payload.get("messages") or []
|
||||
script = find_reply_script(messages)
|
||||
nonce = find_request_nonce(messages)
|
||||
first_tool_name = _first_offered_tool_name(payload)
|
||||
if script == ReplyScript.TWO_CALLS:
|
||||
return WorkerReply(
|
||||
content=None,
|
||||
tool_calls=[_tool_call(first_tool_name, nonce), _tool_call(first_tool_name, f"{nonce} again")],
|
||||
finish_reason=ChatFinishReason.TOOL_CALLS.value,
|
||||
)
|
||||
if script == ReplyScript.UNDECLARED:
|
||||
return WorkerReply(
|
||||
content=None,
|
||||
tool_calls=[_tool_call(UNOFFERED_FUNCTION_NAME, nonce)],
|
||||
finish_reason=ChatFinishReason.TOOL_CALLS.value,
|
||||
breaks_every_contract=True,
|
||||
)
|
||||
if script == ReplyScript.DUP_IDS:
|
||||
shared_call_id = _new_call_id()
|
||||
return WorkerReply(
|
||||
content=None,
|
||||
tool_calls=[
|
||||
_tool_call(first_tool_name, nonce, call_id=shared_call_id),
|
||||
_tool_call(first_tool_name, f"{nonce} again", call_id=shared_call_id),
|
||||
],
|
||||
finish_reason=ChatFinishReason.TOOL_CALLS.value,
|
||||
breaks_every_contract=True,
|
||||
)
|
||||
if script == ReplyScript.PROSE:
|
||||
return WorkerReply(
|
||||
content=f"{PROSE_REPLY} {nonce}", tool_calls=None, finish_reason=ChatFinishReason.STOP.value, breaks_strict_contract=True
|
||||
)
|
||||
if script == ReplyScript.NONCONFORMANT:
|
||||
# Well-formed JSON that misses every key a test schema requires, and carries a key none allows.
|
||||
nonconformant_document = json.dumps({"x": 0, "nonce": nonce})
|
||||
return WorkerReply(
|
||||
content=nonconformant_document, tool_calls=None, finish_reason=ChatFinishReason.STOP.value, breaks_strict_contract=True
|
||||
)
|
||||
if script == ReplyScript.TEXT_AND_CALL:
|
||||
return WorkerReply(
|
||||
content=f"{TEXT_BEFORE_CALL} {nonce}",
|
||||
tool_calls=[_tool_call(first_tool_name, nonce)],
|
||||
finish_reason=ChatFinishReason.TOOL_CALLS.value,
|
||||
)
|
||||
if script == ReplyScript.FAULT_GENERATION_ERROR:
|
||||
return WorkerReply(
|
||||
content=None, tool_calls=None, finish_reason=ChatFinishReason.STOP.value, fault_reason=ChatFaultReason.GENERATION_ERROR.value
|
||||
)
|
||||
return _build_well_behaved_reply(payload, nonce)
|
||||
|
||||
|
||||
def build_misbehaving_reply(payload: Mapping[str, Any]) -> WorkerReply:
|
||||
"""Create the reply a flaky backend gives when it breaks the request's contract: prose where a tool call is required,
|
||||
text that is not JSON where a response format is set, and the well-behaved reply where nothing can be broken.
|
||||
|
||||
Args:
|
||||
payload: The popped job's ``payload``.
|
||||
"""
|
||||
nonce = find_request_nonce(payload.get("messages") or [])
|
||||
if payload.get("tools") and payload.get("tool_choice") == "required":
|
||||
return WorkerReply(
|
||||
content=f"{PROSE_REPLY} {nonce}", tool_calls=None, finish_reason=ChatFinishReason.STOP.value, breaks_strict_contract=True
|
||||
)
|
||||
response_format = payload.get("response_format") or {}
|
||||
if response_format.get("type") in ("json_object", "json_schema") and not _expects_tool_call(payload):
|
||||
return WorkerReply(
|
||||
content=f"{NOT_JSON_REPLY} {nonce}", tool_calls=None, finish_reason=ChatFinishReason.STOP.value, breaks_strict_contract=True
|
||||
)
|
||||
return _build_well_behaved_reply(payload, nonce)
|
||||
|
||||
|
||||
def _build_well_behaved_reply(payload: Mapping[str, Any], nonce: str) -> WorkerReply:
|
||||
"""Create the reply a conforming backend gives, carrying ``nonce``: a call when tools are offered to a user turn, a
|
||||
document that fits the response format, a description of the images, or plain text."""
|
||||
if _expects_tool_call(payload):
|
||||
return WorkerReply(
|
||||
content=None,
|
||||
tool_calls=[_tool_call(_first_offered_tool_name(payload), nonce)],
|
||||
finish_reason=ChatFinishReason.TOOL_CALLS.value,
|
||||
)
|
||||
response_format = payload.get("response_format") or {}
|
||||
if response_format.get("type") == "json_schema":
|
||||
schema = (response_format.get("json_schema") or {}).get("schema") or {}
|
||||
return WorkerReply(
|
||||
content=json.dumps(build_schema_sample(schema, text=nonce)), tool_calls=None, finish_reason=ChatFinishReason.STOP.value
|
||||
)
|
||||
if response_format.get("type") == "json_object":
|
||||
return WorkerReply(content=json.dumps({"x": 1, "nonce": nonce}), tool_calls=None, finish_reason=ChatFinishReason.STOP.value)
|
||||
image_urls = find_image_urls(payload.get("messages") or [])
|
||||
if image_urls:
|
||||
return WorkerReply(content=f"{describe_image_urls(image_urls)} {nonce}", tool_calls=None, finish_reason=ChatFinishReason.STOP.value)
|
||||
return WorkerReply(content=f"{PLAIN_REPLY} {nonce}", tool_calls=None, finish_reason=ChatFinishReason.STOP.value)
|
||||
|
||||
|
||||
def _expects_tool_call(payload: Mapping[str, Any]) -> bool:
|
||||
"""Return whether a conforming backend calls a tool: tools arrived and the last message is the user's turn.
|
||||
|
||||
After a tool result the backend answers in prose, which is how a tool round ends.
|
||||
"""
|
||||
messages = payload.get("messages") or []
|
||||
last_role = messages[-1].get("role") if messages else None
|
||||
return bool(payload.get("tools")) and last_role == "user"
|
||||
|
||||
|
||||
def build_schema_sample(schema: Mapping[str, Any], *, text: str) -> Any:
|
||||
"""Return a document that satisfies a JSON Schema built from the keywords the suite's schemas use.
|
||||
|
||||
Covers ``type`` object (every property), integer and number (``minimum`` honoured), string, boolean, array
|
||||
(``minItems`` honoured), ``enum`` and ``anyOf``. A schema using anything else gets a document that may not fit it.
|
||||
|
||||
Args:
|
||||
schema: The schema.
|
||||
text: The value of every string, such as the request's nonce.
|
||||
"""
|
||||
if "enum" in schema:
|
||||
return schema["enum"][0]
|
||||
if "anyOf" in schema:
|
||||
return build_schema_sample(schema["anyOf"][0], text=text)
|
||||
schema_type = schema.get("type")
|
||||
if schema_type == "object":
|
||||
return {
|
||||
property_name: build_schema_sample(property_schema, text=text)
|
||||
for property_name, property_schema in (schema.get("properties") or {}).items()
|
||||
}
|
||||
if schema_type in ("integer", "number"):
|
||||
return max(5, schema.get("minimum", 5))
|
||||
if schema_type == "string":
|
||||
return text
|
||||
if schema_type == "boolean":
|
||||
return True
|
||||
if schema_type == "array":
|
||||
return [build_schema_sample(schema.get("items") or {}, text=text) for _ in range(schema.get("minItems", 0))]
|
||||
return None
|
||||
|
||||
|
||||
def find_image_urls(messages: Sequence[Mapping[str, Any]]) -> list[str]:
|
||||
"""Return the ``url`` of every image part in the messages, in order."""
|
||||
image_urls: list[str] = []
|
||||
for message in messages:
|
||||
content = message.get("content")
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
for part in content:
|
||||
if isinstance(part, dict) and part.get("type") == "image_url":
|
||||
image_urls.append(str((part.get("image_url") or {}).get("url", "")))
|
||||
return image_urls
|
||||
|
||||
|
||||
def describe_image_urls(image_urls: Sequence[str]) -> str:
|
||||
"""Return the vision reply: ``vision-ok`` when every image arrived as a JPEG data URL, else what arrived instead."""
|
||||
unexpected = [image_url[:40] for image_url in image_urls if not _is_jpeg_data_url(image_url)]
|
||||
if unexpected:
|
||||
return f"{VISION_UNEXPECTED_PREFIX} {unexpected}"
|
||||
return f"{VISION_OK_PREFIX} {len(image_urls)} jpeg"
|
||||
|
||||
|
||||
def _is_jpeg_data_url(image_url: str) -> bool:
|
||||
"""Return whether a URL is a base64 JPEG data URL whose bytes start like a JPEG file."""
|
||||
if not image_url.startswith(JPEG_DATA_URL_PREFIX):
|
||||
return False
|
||||
try:
|
||||
image_bytes = base64.b64decode(image_url[len(JPEG_DATA_URL_PREFIX) :], validate=True)
|
||||
except (binascii.Error, ValueError):
|
||||
return False
|
||||
return image_bytes.startswith(JPEG_MAGIC)
|
||||
|
||||
|
||||
def find_job_payload_violations(payload: Mapping[str, Any], persona: ChatPersona) -> list[str]:
|
||||
"""Return how a popped chat payload breaks the payload contract or the worker's own declaration; empty when it
|
||||
honours both.
|
||||
|
||||
The checks are the routing layer (nothing the worker did not declare reaches it) and the dispatch narrowing (no
|
||||
``tool_choice`` ``none``, a function choice arrives as ``required`` with one tool), plus the payload shape: messages
|
||||
and no prompt, every sampling field present, ``developer`` rewritten to ``system``, budgets within the declaration.
|
||||
|
||||
Args:
|
||||
payload: The popped job's ``payload``.
|
||||
persona: The declaration the worker popped with.
|
||||
"""
|
||||
violations: list[str] = []
|
||||
messages = payload.get("messages")
|
||||
if not isinstance(messages, list) or not messages:
|
||||
violations.append("the chat payload carries no messages")
|
||||
messages = []
|
||||
if "prompt" in payload:
|
||||
violations.append("the chat payload carries a prompt; a chat job never does")
|
||||
missing_sampling = [sampling_key for sampling_key in CHAT_SAMPLING_DEFAULTS if sampling_key not in payload]
|
||||
if missing_sampling:
|
||||
violations.append(f"the chat payload lacks sampling fields {missing_sampling}; every one is always sent")
|
||||
roles = {message.get("role") for message in messages if isinstance(message, dict)}
|
||||
if not roles <= WORKER_MESSAGE_ROLES:
|
||||
violations.append(f"the chat payload carries roles {sorted(str(role) for role in roles - WORKER_MESSAGE_ROLES)}")
|
||||
violations.extend(_find_tool_dispatch_violations(payload, persona))
|
||||
violations.extend(_find_declaration_violations(payload, persona, messages))
|
||||
max_length = payload.get("max_length")
|
||||
if not isinstance(max_length, int) or max_length > persona.chat_max_length:
|
||||
violations.append(f"max_length {max_length!r} exceeds the declared chat_limits.max_length {persona.chat_max_length}")
|
||||
max_context_length = payload.get("max_context_length")
|
||||
if not isinstance(max_context_length, int) or max_context_length > persona.chat_max_context_length:
|
||||
violations.append(f"max_context_length {max_context_length!r} exceeds the declared {persona.chat_max_context_length}")
|
||||
return violations
|
||||
|
||||
|
||||
def _find_tool_dispatch_violations(payload: Mapping[str, Any], persona: ChatPersona) -> list[str]:
|
||||
"""Return the dispatch narrowing violations of a payload's tool keys."""
|
||||
violations: list[str] = []
|
||||
tools = payload.get("tools")
|
||||
tool_choice = payload.get("tool_choice")
|
||||
if tool_choice is not None and not (isinstance(tool_choice, str) and tool_choice in DISPATCHED_TOOL_CHOICES):
|
||||
violations.append(f"tool_choice {json.dumps(tool_choice)} reached the worker; only auto and required are dispatched")
|
||||
if (tool_choice is not None or "parallel_tool_calls" in payload) and not tools:
|
||||
violations.append("tool_choice or parallel_tool_calls arrived without tools")
|
||||
tool_temperature_limit = persona.tools_max_temperature
|
||||
temperature = payload.get("temperature")
|
||||
if tools and tool_temperature_limit is not None and isinstance(temperature, int | float) and temperature > tool_temperature_limit:
|
||||
violations.append(
|
||||
f"a tools job at temperature {temperature} reached a worker whose tools_max_temperature is {tool_temperature_limit}"
|
||||
)
|
||||
return violations
|
||||
|
||||
|
||||
def _find_declaration_violations(payload: Mapping[str, Any], persona: ChatPersona, messages: Sequence[Any]) -> list[str]:
|
||||
"""Return the routing violations of a payload: a feature the worker did not declare reached it."""
|
||||
declared = set(persona.features) | set(persona.guarantees)
|
||||
violations: list[str] = []
|
||||
response_format_type = (payload.get("response_format") or {}).get("type")
|
||||
needed: list[str] = []
|
||||
if payload.get("tools") or any(
|
||||
isinstance(message, dict) and (message.get("role") == "tool" or message.get("tool_calls")) for message in messages
|
||||
):
|
||||
needed.append(ChatFeature.TOOLS.value)
|
||||
if response_format_type in (ChatFeature.JSON_OBJECT.value, ChatFeature.JSON_SCHEMA.value):
|
||||
needed.append(response_format_type)
|
||||
if payload.get("tools") and response_format_type is not None:
|
||||
needed.append(ChatFeature.TOOLS_WITH_RESPONSE_FORMAT.value)
|
||||
if "seed" in payload:
|
||||
needed.append(ChatFeature.SEED.value)
|
||||
if payload.get("presence_penalty"):
|
||||
needed.append(ChatFeature.PRESENCE_PENALTY.value)
|
||||
if payload.get("frequency_penalty"):
|
||||
needed.append(ChatFeature.FREQUENCY_PENALTY.value)
|
||||
if find_image_urls(messages):
|
||||
needed.append(ChatFeature.VISION.value)
|
||||
if not persona.image_transports:
|
||||
violations.append("an image reached a worker that declares no image transport")
|
||||
undeclared = [feature for feature in needed if feature not in declared]
|
||||
if undeclared:
|
||||
violations.append(f"the job needs {undeclared}, which the worker did not declare")
|
||||
return violations
|
||||
|
||||
|
||||
def find_reply_script(messages: Sequence[Any]) -> ReplyScript | None:
|
||||
"""Return the script named by the marker in the last user message, or None without a known one."""
|
||||
for message in reversed(messages):
|
||||
if isinstance(message, dict) and message.get("role") == "user":
|
||||
match = REPLY_MARKER_PATTERN.search(_message_text(message))
|
||||
if match is None:
|
||||
return None
|
||||
try:
|
||||
return ReplyScript(match.group(1))
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _message_text(message: Mapping[str, Any]) -> str:
|
||||
"""Return a message's text, joining text parts."""
|
||||
content = message.get("content")
|
||||
if isinstance(content, list):
|
||||
return " ".join(str(part.get("text", "")) for part in content if isinstance(part, dict))
|
||||
return str(content or "")
|
||||
|
||||
|
||||
def _first_offered_tool_name(payload: Mapping[str, Any]) -> str:
|
||||
"""Return the first offered tool's name, or ``get_weather`` when none arrived."""
|
||||
for tool in payload.get("tools") or []:
|
||||
name = (tool.get("function") or {}).get("name")
|
||||
if name:
|
||||
return str(name)
|
||||
return "get_weather"
|
||||
|
||||
|
||||
def _new_call_id() -> str:
|
||||
"""Create a call id unique to this reply."""
|
||||
return f"call_{uuid.uuid4().hex[:12]}"
|
||||
|
||||
|
||||
def _tool_call(name: str, city: str, *, call_id: str | None = None) -> dict[str, Any]:
|
||||
"""Create one call in the shape a backend mints it, with ``city`` as its only argument; both offered tools take a
|
||||
``city`` string, which is where the request's nonce goes."""
|
||||
return {"id": call_id or _new_call_id(), "type": "function", "function": {"name": name, "arguments": json.dumps({"city": city})}}
|
||||
@@ -0,0 +1,423 @@
|
||||
# SPDX-FileCopyrightText: 2026 Tazlin
|
||||
#
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
"""The chat mix's simulated chat worker: pops as one persona, checks each payload against its declaration, and
|
||||
replies as the client's reply marker asks.
|
||||
|
||||
The worker never takes a new identity to get past a refusal. A configured name, owned by ``--chat-mix-worker-api-key``,
|
||||
backs off after a refused pop and stops after three refusals in a row; a generated name stops at the first identity
|
||||
refusal, since minting a new worker per refusal would evade the Horde's anti-abuse limits. A worker that stops records
|
||||
why, and the coverage gate fails the run on it.
|
||||
|
||||
A breaker protects the worker name: after three refusals of replies the worker meant to conform, it stops faulting
|
||||
(each fault would count as an aborted job) and stops, so a Horde regression cannot drive a deployment's worker into
|
||||
maintenance.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from locust import HttpUser, between, tag, task
|
||||
from locust.exception import RescheduleTask, StopUser
|
||||
|
||||
from tests.integration._fake_chat_worker import build_chat_submit_body, build_faulted_chat_submit_body, build_pop_body
|
||||
|
||||
from .chat_checks import classify_environment_condition, error_code_of
|
||||
from .chat_mix_settings import (
|
||||
chat_mix_settings,
|
||||
claim_worker_slot,
|
||||
fail_contract,
|
||||
fail_environment,
|
||||
read_worker_counters,
|
||||
record_benign,
|
||||
)
|
||||
from .chat_wire import TEXT_POP_PATH, TEXT_SUBMIT_PATH, WORKER_BRIDGE_AGENT
|
||||
from .chat_worker_script import (
|
||||
ChatPersona,
|
||||
ChatPersonaRole,
|
||||
WorkerReply,
|
||||
build_misbehaving_reply,
|
||||
build_persona,
|
||||
build_scripted_reply,
|
||||
find_job_payload_violations,
|
||||
find_reply_script,
|
||||
)
|
||||
from .config import _EXPECTED_RC_RECOVER
|
||||
from .helpers import _headers, _is_expected_rc, _is_too_many_workers, _pick_worker_key, _random_worker_name, _safe_json
|
||||
from .users.chat import STREAM_PROGRESS_FIELD, post_progress_while_generating
|
||||
|
||||
__all__ = ["ChatMixWorker", "refusal_was_intended", "reply_posts_progress"]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CHAT_TAG: str = "chat"
|
||||
"""The tag every chat mix task carries."""
|
||||
|
||||
STALE_JOB_RCS: frozenset[str] = frozenset({"InvalidJobID", "InvalidProcGen", "AbortedGen", "DuplicateGen"})
|
||||
"""The return codes of a submit for a job that is no longer the worker's, such as one the client cancelled."""
|
||||
|
||||
WORKER_POLL_MIN_SECONDS: float = 0.5
|
||||
"""The shortest pause between two pops."""
|
||||
|
||||
WORKER_POLL_MAX_SECONDS: float = 1.5
|
||||
"""The longest pause between two pops."""
|
||||
|
||||
IDENTITY_REFUSAL_BACKOFF_SECONDS: float = 30.0
|
||||
"""How long a configured name waits after its pop was refused, before trying again."""
|
||||
|
||||
IDENTITY_REFUSALS_BEFORE_STOP: int = 3
|
||||
"""How many consecutive pop refusals a configured name takes before its worker stops."""
|
||||
|
||||
UNINTENDED_REFUSALS_BEFORE_STOP: int = 3
|
||||
"""How many refusals of replies the worker meant to conform it takes before its breaker stops it."""
|
||||
|
||||
WORKER_TEXT_MAX_LENGTH: int = 512
|
||||
"""The text output limit a simulated chat worker declares for text jobs."""
|
||||
|
||||
WORKER_TEXT_MAX_CONTEXT_LENGTH: int = 2048
|
||||
"""The text context limit a simulated chat worker declares for text jobs."""
|
||||
|
||||
CHARACTERS_PER_TOKEN: int = 4
|
||||
"""Characters per token in the simulated workers' reported counts."""
|
||||
|
||||
UNREPORTED_USAGE_CHANCE: float = 0.1
|
||||
"""The chance a worker leaves out its token counts, so the front's usage estimate runs."""
|
||||
|
||||
TEXT_JOB_GENERATION: str = "Once upon a time there was a stress test that completed successfully."
|
||||
"""The generation a simulated chat worker submits for a text job."""
|
||||
|
||||
REPLY_REFUSED_FAULT_REASON: str = "reply_refused"
|
||||
"""The ``fault_reason`` a worker sends after the Horde refused its reply."""
|
||||
|
||||
EXCERPT_CHARACTERS: int = 200
|
||||
"""How much of a body a failure quotes."""
|
||||
|
||||
|
||||
class ChatMixWorker(HttpUser):
|
||||
"""Pop chat jobs as one persona, check each payload against the declaration, and reply as scripted."""
|
||||
|
||||
weight = 1
|
||||
fixed_count = 0
|
||||
wait_time = between(WORKER_POLL_MIN_SECONDS, WORKER_POLL_MAX_SECONDS)
|
||||
|
||||
def on_start(self) -> None:
|
||||
"""Take the next worker slot: a persona and a configured name with its key, or a generated name.
|
||||
|
||||
In a distributed run the slot comes from the master, so a name is never popped by two processes.
|
||||
|
||||
Raises:
|
||||
StopUser: The run has no slot left or the master did not grant one; the coverage gate fails the run on it.
|
||||
"""
|
||||
self.has_slot: bool = False
|
||||
slot = claim_worker_slot(self.environment)
|
||||
if slot is None:
|
||||
reason = "a chat worker got no worker slot (more workers than --chat-mix-workers, or no answer from the master)"
|
||||
logger.error("Chat mix: %s; it stops at once.", reason)
|
||||
chat_mix_settings.coverage.stopped_workers.append(reason)
|
||||
raise StopUser()
|
||||
self.persona_role: ChatPersonaRole = slot.persona
|
||||
self.uses_configured_name: bool = slot.configured_name is not None
|
||||
self.worker_name: str = slot.configured_name or self._generate_name()
|
||||
self.api_key: str = slot.api_key or _pick_worker_key()
|
||||
self.persona: ChatPersona = build_persona(self.persona_role, vision=chat_mix_settings.vision)
|
||||
self.model: str = chat_mix_settings.model_of(self.persona_role)
|
||||
self.rng = random.Random()
|
||||
self.consecutive_pop_refusals: int = 0
|
||||
self.unintended_refusals: int = 0
|
||||
chat_mix_settings.mark_worker_started(self.worker_name, self.persona_role)
|
||||
self.has_slot = True
|
||||
_snapshot_worker_counters(self.worker_name)
|
||||
|
||||
def on_stop(self) -> None:
|
||||
"""Take the worker offline, so scenarios needing its persona pause."""
|
||||
if self.has_slot:
|
||||
chat_mix_settings.mark_worker_offline(self.worker_name)
|
||||
|
||||
def _generate_name(self) -> str:
|
||||
"""Create a worker name unlikely to collide with another simulated worker's."""
|
||||
return _random_worker_name("StressChatMix-", 6)
|
||||
|
||||
@tag(CHAT_TAG, "chat-worker", "worker")
|
||||
@task
|
||||
def pop_and_serve(self) -> None:
|
||||
"""Pop one job and serve it: a chat job as the persona and its script say, a text job with plain text.
|
||||
|
||||
With ``--chat-mix-worker-progress``, a chat job whose client streams and whose reply conforms posts the reply
|
||||
text as content deltas while it generates, and the submit carries the same text. An acknowledgement reporting
|
||||
the request cancelled ends the job without a submit.
|
||||
"""
|
||||
job = self._pop()
|
||||
if job is None:
|
||||
return
|
||||
generation_seconds = self.rng.uniform(chat_mix_settings.gen_time_min_seconds, chat_mix_settings.gen_time_max_seconds)
|
||||
if job.get("chat") is not True:
|
||||
time.sleep(generation_seconds)
|
||||
self._submit_text_job(job)
|
||||
return
|
||||
payload = job.get("payload") or {}
|
||||
reply = self._choose_reply(payload)
|
||||
if chat_mix_settings.worker_progress and reply_posts_progress(payload, reply):
|
||||
request_cancelled = post_progress_while_generating(
|
||||
self.client,
|
||||
api_key=self.api_key,
|
||||
generation_id=job["id"],
|
||||
reply_text=reply.content or "",
|
||||
generation_seconds=generation_seconds,
|
||||
interval_seconds=chat_mix_settings.worker_progress_interval_seconds,
|
||||
)
|
||||
if request_cancelled:
|
||||
return
|
||||
else:
|
||||
time.sleep(generation_seconds)
|
||||
if reply.fault_reason is not None:
|
||||
self._submit_fault(job, reply.fault_reason, name=f"{TEXT_SUBMIT_PATH} [chat fault scripted]")
|
||||
return
|
||||
self._submit_reply(job, reply)
|
||||
|
||||
def _pop(self) -> dict[str, Any] | None:
|
||||
"""Pop with the persona's declaration and return the job, or None when there is none or the pop was refused."""
|
||||
pop_body = build_pop_body(
|
||||
name=self.worker_name,
|
||||
models=[self.model],
|
||||
bridge_agent=WORKER_BRIDGE_AGENT,
|
||||
chat_jobs=True,
|
||||
prompt_jobs=self.persona.prompt_jobs,
|
||||
chat_features=list(self.persona.features),
|
||||
chat_guarantees=list(self.persona.guarantees),
|
||||
chat_max_context_length=self.persona.chat_max_context_length,
|
||||
chat_max_length=self.persona.chat_max_length,
|
||||
tools_max_temperature=self.persona.tools_max_temperature,
|
||||
max_context_length=WORKER_TEXT_MAX_CONTEXT_LENGTH,
|
||||
max_length=WORKER_TEXT_MAX_LENGTH,
|
||||
image_transports=list(self.persona.image_transports) or None,
|
||||
)
|
||||
pop_body.update({"nsfw": True, "softprompts": [], "threads": 1})
|
||||
name = f"{TEXT_POP_PATH} [chat {self.persona_role.value}]"
|
||||
with self.client.post(TEXT_POP_PATH, json=pop_body, headers=_headers(self.api_key), catch_response=True, name=name) as resp:
|
||||
if resp.status_code == 0:
|
||||
resp.failure(f"Transport error: {resp.error}")
|
||||
return None
|
||||
body = _safe_json(resp)
|
||||
if not resp.ok:
|
||||
self.handle_pop_refusal(resp, body, name)
|
||||
return None
|
||||
self.consecutive_pop_refusals = 0
|
||||
chat_mix_settings.mark_worker_online(self.worker_name, self.persona_role)
|
||||
job = body or {}
|
||||
if not job.get("id"):
|
||||
resp.success()
|
||||
return None
|
||||
if job.get("chat") is True:
|
||||
violations = find_job_payload_violations(job.get("payload") or {}, self.persona)
|
||||
if violations:
|
||||
fail_contract(resp, f"{self.persona_role.value} worker payload", violations)
|
||||
return job
|
||||
elif not self.persona.prompt_jobs:
|
||||
fail_contract(resp, f"{self.persona_role.value} worker", ["a worker that set prompt_jobs false popped a text job"])
|
||||
return job
|
||||
resp.success()
|
||||
return job
|
||||
|
||||
def handle_pop_refusal(self, resp: Any, body: Any, name: str) -> None:
|
||||
"""Record a refused pop and decide whether the worker waits, tries again or stops.
|
||||
|
||||
A rate limit is benign. A refusal of the worker's identity (wrong owner, too many workers, maintenance, the
|
||||
other anti-abuse codes) takes the persona offline; a configured name waits and stops after three in a row, a
|
||||
generated name stops at once. A deployment condition fails as one; anything else is a contract violation.
|
||||
|
||||
Raises:
|
||||
RescheduleTask: After a rate limit, or a configured name's refusal short of the limit.
|
||||
StopUser: When the worker gives up its identity.
|
||||
"""
|
||||
return_code = error_code_of(body)
|
||||
if resp.status_code == 429:
|
||||
record_benign(self.environment, resp, name, "rate_limited")
|
||||
raise RescheduleTask()
|
||||
condition = classify_environment_condition(resp.status_code, return_code, has_envelope=isinstance(body, dict))
|
||||
if condition is not None:
|
||||
fail_environment(resp, condition, f"{self.persona_role.value} worker pop")
|
||||
raise RescheduleTask()
|
||||
refuses_identity = resp.status_code in (400, 403) and (_is_expected_rc(body, _EXPECTED_RC_RECOVER) or _is_too_many_workers(body))
|
||||
if not refuses_identity:
|
||||
fail_contract(
|
||||
resp, f"{self.persona_role.value} worker pop", [f"the pop answered {resp.status_code}: {resp.text[:EXCERPT_CHARACTERS]}"]
|
||||
)
|
||||
return
|
||||
record_benign(self.environment, resp, name, str(return_code))
|
||||
self.consecutive_pop_refusals += 1
|
||||
chat_mix_settings.mark_worker_offline(self.worker_name)
|
||||
logger.warning(
|
||||
"Chat mix: %s worker %s pop refused with %s (%d in a row)",
|
||||
self.persona_role.value,
|
||||
self.worker_name,
|
||||
return_code,
|
||||
self.consecutive_pop_refusals,
|
||||
)
|
||||
if not self.uses_configured_name or self.consecutive_pop_refusals >= IDENTITY_REFUSALS_BEFORE_STOP:
|
||||
refusals = f"{self.consecutive_pop_refusals} pop refusals ({return_code})"
|
||||
self._stop(f"the {self.persona_role.value} worker {self.worker_name} stopped after {refusals}")
|
||||
time.sleep(IDENTITY_REFUSAL_BACKOFF_SECONDS)
|
||||
raise RescheduleTask()
|
||||
|
||||
def _stop(self, reason: str) -> None:
|
||||
"""Record why the worker stops, for the coverage gate, and stop it.
|
||||
|
||||
Raises:
|
||||
StopUser: Always.
|
||||
"""
|
||||
logger.error("Chat mix: %s", reason)
|
||||
chat_mix_settings.coverage.stopped_workers.append(reason)
|
||||
chat_mix_settings.mark_worker_offline(self.worker_name)
|
||||
raise StopUser()
|
||||
|
||||
def _choose_reply(self, payload: dict[str, Any]) -> WorkerReply:
|
||||
"""Return the reply: the client's script when it wrote one, else the persona's own behaviour."""
|
||||
is_unscripted = find_reply_script(payload.get("messages") or []) is None
|
||||
if self.persona_role == ChatPersonaRole.FLAKY and is_unscripted and self.rng.random() < chat_mix_settings.flaky_rate:
|
||||
return build_misbehaving_reply(payload)
|
||||
return build_scripted_reply(payload)
|
||||
|
||||
def _submit_reply(self, job: dict[str, Any], reply: WorkerReply) -> None:
|
||||
"""Submit a structured reply and check the Horde accepts or refuses it as the contract says; fault a refused
|
||||
reply, unless the breaker trips."""
|
||||
reported_usage = self.rng.random() >= UNREPORTED_USAGE_CHANCE
|
||||
reply_characters = len(reply.content or "") + len(json.dumps(reply.tool_calls or []))
|
||||
submit_body = build_chat_submit_body(
|
||||
generation_id=job["id"],
|
||||
content=reply.content,
|
||||
model_id=job.get("model") or self.model,
|
||||
finish_reason=reply.finish_reason,
|
||||
prompt_tokens=20 if reported_usage else None,
|
||||
completion_tokens=max(1, reply_characters // CHARACTERS_PER_TOKEN) if reported_usage else None,
|
||||
tool_calls=reply.tool_calls,
|
||||
)
|
||||
name = f"{TEXT_SUBMIT_PATH} [chat {self.persona_role.value}]"
|
||||
with self.client.post(TEXT_SUBMIT_PATH, json=submit_body, headers=_headers(self.api_key), catch_response=True, name=name) as resp:
|
||||
if resp.status_code == 0:
|
||||
resp.failure(f"Transport error: {resp.error}")
|
||||
return
|
||||
return_code = error_code_of(_safe_json(resp))
|
||||
if resp.ok:
|
||||
if reply.breaks_every_contract:
|
||||
fail_contract(
|
||||
resp, f"{self.persona_role.value} worker submit", ["the Horde stored a reply every conformance mode refuses"]
|
||||
)
|
||||
return
|
||||
chat_mix_settings.ledger.accepted_replies[self.worker_name] += 1
|
||||
resp.success()
|
||||
return
|
||||
if resp.status_code in (404, 429) or return_code in STALE_JOB_RCS:
|
||||
record_benign(self.environment, resp, name, str(return_code or resp.status_code))
|
||||
return
|
||||
intended = refusal_was_intended(reply, return_code)
|
||||
if intended:
|
||||
record_benign(self.environment, resp, name, f"refused as intended {return_code}")
|
||||
else:
|
||||
fail_contract(
|
||||
resp,
|
||||
f"{self.persona_role.value} worker submit",
|
||||
[f"a reply the worker meant to conform was refused: {resp.text[:EXCERPT_CHARACTERS]}"],
|
||||
)
|
||||
self._count_refusal(return_code)
|
||||
if not intended:
|
||||
self.unintended_refusals += 1
|
||||
if self.unintended_refusals >= UNINTENDED_REFUSALS_BEFORE_STOP:
|
||||
# Faulting would count one more aborted job against a worker name the run cannot replace.
|
||||
refusals = f"{self.unintended_refusals} replies it meant to conform were refused"
|
||||
self._stop(f"the {self.persona_role.value} worker {self.worker_name} stopped: {refusals}")
|
||||
# A refusal is definitive: fault the job, as koboldcpp does, so it goes back to the queue at once.
|
||||
self._submit_fault(job, REPLY_REFUSED_FAULT_REASON, name=f"{TEXT_SUBMIT_PATH} [chat fault after refusal]")
|
||||
|
||||
def _count_refusal(self, return_code: str | None) -> None:
|
||||
"""Count a refused reply in the ledger by its return code."""
|
||||
if return_code == "ChatSubmitNonconformant":
|
||||
chat_mix_settings.ledger.refused_nonconformant[self.worker_name] += 1
|
||||
elif return_code == "ChatSubmitInvalid":
|
||||
chat_mix_settings.ledger.refused_invalid[self.worker_name] += 1
|
||||
|
||||
def _submit_fault(self, job: dict[str, Any], fault_reason: str, *, name: str) -> None:
|
||||
"""Fault a chat job with a reason; the Horde answers with a reward of 0."""
|
||||
submit_body = build_faulted_chat_submit_body(
|
||||
generation_id=job["id"], model_id=job.get("model") or self.model, fault_reason=fault_reason
|
||||
)
|
||||
with self.client.post(TEXT_SUBMIT_PATH, json=submit_body, headers=_headers(self.api_key), catch_response=True, name=name) as resp:
|
||||
self._judge_plain_submit(resp, name, "fault submit")
|
||||
if resp.ok:
|
||||
chat_mix_settings.ledger.faults_submitted[self.worker_name] += 1
|
||||
|
||||
def _submit_text_job(self, job: dict[str, Any]) -> None:
|
||||
"""Submit a text job the way a text worker submits it."""
|
||||
name = f"{TEXT_SUBMIT_PATH} [text on chat {self.persona_role.value}]"
|
||||
submit_body = {"id": job["id"], "generation": TEXT_JOB_GENERATION, "state": "ok", "seed": self.rng.randint(0, 999_999_999)}
|
||||
with self.client.post(TEXT_SUBMIT_PATH, json=submit_body, headers=_headers(self.api_key), catch_response=True, name=name) as resp:
|
||||
self._judge_plain_submit(resp, name, "text submit")
|
||||
if resp.ok:
|
||||
chat_mix_settings.ledger.text_jobs[self.worker_name] += 1
|
||||
|
||||
def _judge_plain_submit(self, resp: Any, name: str, subject: str) -> None:
|
||||
"""Mark a fault or text submit: success, a stale job (benign), a deployment condition, or a contract violation."""
|
||||
if resp.status_code == 0:
|
||||
resp.failure(f"Transport error: {resp.error}")
|
||||
return
|
||||
body = _safe_json(resp)
|
||||
return_code = error_code_of(body)
|
||||
if resp.ok:
|
||||
resp.success()
|
||||
return
|
||||
if resp.status_code in (404, 429) or return_code in STALE_JOB_RCS:
|
||||
record_benign(self.environment, resp, name, str(return_code or resp.status_code))
|
||||
return
|
||||
condition = classify_environment_condition(resp.status_code, return_code, has_envelope=isinstance(body, dict))
|
||||
if condition is not None:
|
||||
fail_environment(resp, condition, f"{self.persona_role.value} worker {subject}")
|
||||
return
|
||||
fail_contract(
|
||||
resp, f"{self.persona_role.value} worker {subject}", [f"answered {resp.status_code}: {resp.text[:EXCERPT_CHARACTERS]}"]
|
||||
)
|
||||
|
||||
|
||||
def reply_posts_progress(payload: dict[str, Any], reply: WorkerReply) -> bool:
|
||||
"""Return whether a chat job's reply is posted as content deltas while it generates.
|
||||
|
||||
Only a text reply the worker means to conform is posted, for a job whose client streams. A fault, a reply that
|
||||
breaks a contract and a reply with tool calls are submitted without progress; koboldcpp does not stream tool jobs
|
||||
either, since a streamed tool reply can differ from the submitted one.
|
||||
|
||||
Args:
|
||||
payload: The popped chat job's payload.
|
||||
reply: The reply the worker will submit.
|
||||
"""
|
||||
client_streams = payload.get(STREAM_PROGRESS_FIELD) is True
|
||||
is_conforming_text = bool(reply.content) and not reply.tool_calls and reply.fault_reason is None
|
||||
breaks_a_contract = reply.breaks_strict_contract or reply.breaks_every_contract
|
||||
return client_streams and is_conforming_text and not breaks_a_contract
|
||||
|
||||
|
||||
def refusal_was_intended(reply: WorkerReply, return_code: str | None) -> bool:
|
||||
"""Return whether the Horde refused a reply for the reason the worker broke it: nonconformant replies under
|
||||
``strict``, invalid ones under every mode.
|
||||
|
||||
Args:
|
||||
reply: The reply the worker submitted.
|
||||
return_code: The refusal's return code.
|
||||
"""
|
||||
if return_code == "ChatSubmitNonconformant":
|
||||
return reply.breaks_strict_contract
|
||||
if return_code == "ChatSubmitInvalid":
|
||||
return reply.breaks_every_contract
|
||||
return False
|
||||
|
||||
|
||||
def _snapshot_worker_counters(worker_name: str) -> None:
|
||||
"""Keep a worker's counters at the start of the run, once per name."""
|
||||
if worker_name in chat_mix_settings.ledger.counters_at_start:
|
||||
return
|
||||
counters = read_worker_counters(worker_name)
|
||||
if counters is not None:
|
||||
chat_mix_settings.ledger.counters_at_start[worker_name] = counters
|
||||
@@ -9,6 +9,7 @@ import logging
|
||||
import uuid
|
||||
|
||||
from locust import events
|
||||
from locust.runners import MasterRunner, WorkerRunner
|
||||
import requests
|
||||
|
||||
from .config import _config
|
||||
@@ -412,6 +413,64 @@ def _remove_tag_filtered_user_classes(environment) -> None:
|
||||
logger.info("Excluding user classes with no tasks after tag filtering: %s", ", ".join(removed))
|
||||
|
||||
|
||||
def _leave_out_user_classes(environment, opt_in_counts) -> None:
|
||||
"""Keep unrequested opt-in classes and classes a tag filter emptied from spawning, in the way each runner allows.
|
||||
|
||||
A local run removes them from ``environment.user_classes`` before its dispatcher exists. A master built its
|
||||
dispatcher before the test start, from every class, so removing a class there changes nothing; it sets the class's
|
||||
weight to 0 instead, which the dispatcher reads when it dispatches. A worker process keeps every class: the master's
|
||||
spawn messages name each class it dispatches, zero counts included, and a class the worker no longer has fails the
|
||||
spawn with a ``KeyError``.
|
||||
|
||||
Args:
|
||||
environment: The Locust environment.
|
||||
opt_in_counts: The requested count of each opt-in class, such as the chat mix users.
|
||||
"""
|
||||
if isinstance(environment.runner, WorkerRunner):
|
||||
return
|
||||
if isinstance(environment.runner, MasterRunner):
|
||||
_zero_weight_of_classes_left_out(environment, opt_in_counts)
|
||||
return
|
||||
_remove_unrequested_opt_in_classes(environment, opt_in_counts)
|
||||
_remove_tag_filtered_user_classes(environment)
|
||||
|
||||
|
||||
def _zero_weight_of_classes_left_out(environment, opt_in_counts) -> None:
|
||||
"""Mutate the weight of unrequested opt-in classes and of classes without tasks to 0, on a master.
|
||||
|
||||
Raises:
|
||||
RuntimeError: A class with a fixed count lost all its tasks to the tag filter.
|
||||
"""
|
||||
left_out = []
|
||||
for user_class in environment.user_classes:
|
||||
is_unrequested = user_class in opt_in_counts and opt_in_counts[user_class] <= 0
|
||||
if user_class.tasks and not is_unrequested:
|
||||
continue
|
||||
if user_class.fixed_count > 0:
|
||||
raise RuntimeError(
|
||||
f"{user_class.__name__} requests {user_class.fixed_count} fixed users, but tag filtering removed all of its tasks.",
|
||||
)
|
||||
user_class.weight = 0
|
||||
left_out.append(user_class.__name__)
|
||||
if left_out:
|
||||
logger.info("Dispatching no users of: %s", ", ".join(left_out))
|
||||
|
||||
|
||||
def _remove_unrequested_opt_in_classes(environment, opt_in_counts) -> None:
|
||||
"""Mutate ``environment.user_classes`` to leave out the opt-in classes whose requested count is zero.
|
||||
|
||||
Args:
|
||||
environment: The Locust environment.
|
||||
opt_in_counts: The requested count of each opt-in class, such as the chat mix users.
|
||||
"""
|
||||
removed = [user_class.__name__ for user_class, count in opt_in_counts.items() if count <= 0 and user_class in environment.user_classes]
|
||||
environment.user_classes[:] = [
|
||||
user_class for user_class in environment.user_classes if not (user_class in opt_in_counts and opt_in_counts[user_class] <= 0)
|
||||
]
|
||||
if removed:
|
||||
logger.info("Leaving out opt-in user classes with no requested count: %s", ", ".join(removed))
|
||||
|
||||
|
||||
@events.test_start.add_listener
|
||||
def on_test_start(environment, **kw):
|
||||
opts = environment.parsed_options
|
||||
@@ -463,6 +522,9 @@ def on_test_start(environment, **kw):
|
||||
# CLI arguments are guaranteed to be parsed.
|
||||
from .users import (
|
||||
BaselineFeatureRequester,
|
||||
ChatMixFrontClient,
|
||||
ChatMixNativeClient,
|
||||
ChatMixWorker,
|
||||
ExtendedWorkerSimulator,
|
||||
HotPathRequester,
|
||||
InterrogationRequester,
|
||||
@@ -491,7 +553,17 @@ def on_test_start(environment, **kw):
|
||||
(BaselineFeatureRequester, "baseline_feature_requestors"),
|
||||
(MetaBrowser, "meta_browsers"),
|
||||
(MisuseUser, "misuse_users"),
|
||||
(ChatMixFrontClient, "chat_mix_front_clients"),
|
||||
(ChatMixNativeClient, "chat_mix_native_clients"),
|
||||
(ChatMixWorker, "chat_mix_workers"),
|
||||
)
|
||||
# The chat mix needs a running front and worker names its key owns, so its classes join the run only when asked
|
||||
# for by count, never through the weight distribution.
|
||||
_opt_in_counts = {
|
||||
ChatMixFrontClient: int(getattr(opts, "chat_mix_front_clients", 0) or 0),
|
||||
ChatMixNativeClient: int(getattr(opts, "chat_mix_native_clients", 0) or 0),
|
||||
ChatMixWorker: int(getattr(opts, "chat_mix_workers", 0) or 0),
|
||||
}
|
||||
total_image_workers = int(getattr(opts, "image_workers", 0) or 0)
|
||||
if total_image_workers and not opts.legacy_image_workers and not opts.extended_image_workers:
|
||||
if total_image_workers < 2:
|
||||
@@ -514,4 +586,4 @@ def on_test_start(environment, **kw):
|
||||
fixed_total,
|
||||
)
|
||||
|
||||
_remove_tag_filtered_user_classes(environment)
|
||||
_leave_out_user_classes(environment, _opt_in_counts)
|
||||
@@ -5,12 +5,18 @@
|
||||
"""Shared helpers used by the AI Horde Locust workloads."""
|
||||
|
||||
import random
|
||||
import re
|
||||
import string
|
||||
import time
|
||||
|
||||
from locust.exception import RescheduleTask
|
||||
|
||||
from .config import _config
|
||||
|
||||
# The Horde refuses a new worker whose name contains any of these, case-insensitively, with a 400 asking to avoid
|
||||
# unwanted attention (``colab_search`` in ``horde/apis/v2/base.py``); keep the two lists equal.
|
||||
_REFUSED_WORKER_NAME_PATTERN = re.compile(r"colab|tpu|google", re.IGNORECASE)
|
||||
|
||||
|
||||
def _is_expected_rc(body_json: dict | None, rc_set: set[str]) -> bool:
|
||||
return bool(body_json and body_json.get("rc") in rc_set)
|
||||
@@ -85,6 +91,15 @@ def _handle_async_generate(resp, environment) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def _random_worker_name(prefix: str, suffix_length: int) -> str:
|
||||
"""Return ``prefix`` followed by random lowercase letters, drawn again until the whole name holds none of the
|
||||
substrings the Horde refuses in a worker name."""
|
||||
while True:
|
||||
worker_name = prefix + "".join(random.choices(string.ascii_lowercase, k=suffix_length))
|
||||
if not _REFUSED_WORKER_NAME_PATTERN.search(worker_name):
|
||||
return worker_name
|
||||
|
||||
|
||||
def _parse_csv(value: str) -> list[str]:
|
||||
if not value:
|
||||
return []
|
||||
|
||||
@@ -4,6 +4,8 @@
|
||||
|
||||
"""Opt-in staged load shapes for the AI Horde Locust suite."""
|
||||
|
||||
import time
|
||||
|
||||
from locust import LoadTestShape
|
||||
|
||||
_STAGE_PROFILES = {
|
||||
@@ -43,3 +45,53 @@ class HordeStagesShape(LoadTestShape):
|
||||
spawn_rate = max(1.0, stage["spawn_rate"] * scale)
|
||||
return users, spawn_rate
|
||||
return None
|
||||
|
||||
|
||||
class ChatStageClock:
|
||||
"""Hold the time the chat scenario finished its test-start work, which the chat stages are measured from.
|
||||
|
||||
Locust starts a shape's clock before the test-start hooks run, and the hooks register the run's API keys, which
|
||||
takes seconds per hundred keys. Measuring the stages from the end of that work keeps every stage its full length.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Create the clock, not yet started."""
|
||||
self.started_at: float | None = None
|
||||
"""The Unix time the chat scenario was ready, or None before then."""
|
||||
|
||||
def mark_started(self) -> None:
|
||||
"""Mutate the clock to start now."""
|
||||
self.started_at = time.time()
|
||||
|
||||
def elapsed_seconds(self) -> float | None:
|
||||
"""Return the seconds since the clock started, or None when it has not."""
|
||||
if self.started_at is None:
|
||||
return None
|
||||
return time.time() - self.started_at
|
||||
|
||||
|
||||
chat_stage_clock = ChatStageClock()
|
||||
"""The process-wide clock the chat locustfile starts and ``ChatStagesShape`` reads."""
|
||||
|
||||
|
||||
class ChatStagesShape(LoadTestShape):
|
||||
"""Hold the chat workers alone for a warm-up, then step the chat clients through ``--chat-stages``.
|
||||
|
||||
Each stage holds ``--chat-workers`` plus that stage's client count for ``--chat-stage-seconds``; the run stops
|
||||
after the last stage. The workers are pinned with ``fixed_count``, so Locust spreads only the clients across the
|
||||
client classes by weight.
|
||||
"""
|
||||
|
||||
def tick(self) -> tuple[int, float] | None:
|
||||
"""Return the user count and spawn rate for the current moment, or None once the last stage is over."""
|
||||
options = self.runner.environment.parsed_options
|
||||
workers = int(options.chat_workers)
|
||||
spawn_rate = float(options.chat_spawn_rate)
|
||||
elapsed_seconds = chat_stage_clock.elapsed_seconds()
|
||||
if elapsed_seconds is None or elapsed_seconds < float(options.chat_warmup_seconds):
|
||||
return max(workers, 1), spawn_rate
|
||||
stage_client_counts = [int(stage) for stage in str(options.chat_stages).split(",") if stage.strip()]
|
||||
stage_index = int((elapsed_seconds - float(options.chat_warmup_seconds)) // float(options.chat_stage_seconds))
|
||||
if stage_index >= len(stage_client_counts):
|
||||
return None
|
||||
return workers + stage_client_counts[stage_index], spawn_rate
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
"""User classes exported for Locust discovery."""
|
||||
|
||||
from .chat_mix import ChatMixFrontClient, ChatMixNativeClient, ChatMixWorker
|
||||
from .image import (
|
||||
BaselineFeatureRequester,
|
||||
ExtendedWorkerSimulator,
|
||||
@@ -20,6 +21,9 @@ from .text import TextRequester, TextWorkerSimulator
|
||||
|
||||
__all__ = [
|
||||
"BaselineFeatureRequester",
|
||||
"ChatMixFrontClient",
|
||||
"ChatMixNativeClient",
|
||||
"ChatMixWorker",
|
||||
"ExtendedWorkerSimulator",
|
||||
"HotPathRequester",
|
||||
"InterrogationRequester",
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,22 @@
|
||||
# SPDX-FileCopyrightText: 2026 Tazlin
|
||||
#
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
"""The chat contract mix's user classes, exported for Locust discovery.
|
||||
|
||||
The populations are opt-in: ``events.on_test_start`` pins each to its ``--chat-mix-*`` count and leaves out a class
|
||||
whose count is zero, because the clients need a running front and the workers need worker names their key owns. With
|
||||
every count at zero (the default) the mix runs as before, which keeps the CI smoke job, which has no front, unchanged.
|
||||
|
||||
- ``ChatMixFrontClient`` ([`chat_front_client`][tests.stress.locustsuite.chat_front_client]) sends the front scenarios.
|
||||
- ``ChatMixNativeClient`` ([`chat_native_client`][tests.stress.locustsuite.chat_native_client]) drives the native routes.
|
||||
- ``ChatMixWorker`` ([`chat_worker_user`][tests.stress.locustsuite.chat_worker_user]) serves chat jobs as a persona.
|
||||
|
||||
The settings, arguments and run hooks they share are in [`chat_mix_settings`][tests.stress.locustsuite.chat_mix_settings].
|
||||
"""
|
||||
|
||||
from ..chat_front_client import ChatMixFrontClient
|
||||
from ..chat_native_client import ChatMixNativeClient
|
||||
from ..chat_worker_user import ChatMixWorker
|
||||
|
||||
__all__ = ["ChatMixFrontClient", "ChatMixNativeClient", "ChatMixWorker"]
|
||||
@@ -43,7 +43,6 @@ from __future__ import annotations
|
||||
|
||||
import collections
|
||||
import random
|
||||
import string
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
@@ -58,6 +57,7 @@ from ..helpers import (
|
||||
_headers,
|
||||
_is_expected_rc,
|
||||
_is_too_many_workers,
|
||||
_random_worker_name,
|
||||
_record_expected,
|
||||
_safe_json,
|
||||
)
|
||||
@@ -365,7 +365,7 @@ class ConvoyWorker(HttpUser):
|
||||
wait_time = constant_pacing(0.1)
|
||||
|
||||
def on_start(self) -> None:
|
||||
self.worker_name = f"HCWorker-{''.join(random.choices(string.ascii_lowercase, k=8))}"
|
||||
self.worker_name = _random_worker_name("HCWorker-", 8)
|
||||
keys = convoy_state.worker_keys
|
||||
self.api_key = random.choice(keys) if keys else _ANON_API_KEY
|
||||
|
||||
|
||||
@@ -5,7 +5,6 @@
|
||||
"""Image generation requestor and worker Locust users."""
|
||||
|
||||
import random
|
||||
import string
|
||||
import time
|
||||
from collections import deque
|
||||
from collections.abc import Sequence
|
||||
@@ -37,6 +36,7 @@ from ..helpers import (
|
||||
_pick_requestor_key,
|
||||
_pick_worker_key,
|
||||
_random_prompt,
|
||||
_random_worker_name,
|
||||
_record_expected,
|
||||
_safe_json,
|
||||
)
|
||||
@@ -390,7 +390,7 @@ class WorkerSimulator(HttpUser):
|
||||
wait_time = between(1, 4)
|
||||
|
||||
def create_worker_name(self):
|
||||
return f"StressWorker-{self.bridge_generation}-{''.join(random.choices(string.ascii_lowercase, k=4))}"
|
||||
return _random_worker_name(f"StressWorker-{self.bridge_generation}-", 4)
|
||||
|
||||
def on_start(self):
|
||||
self.worker_name = self.create_worker_name()
|
||||
@@ -506,7 +506,7 @@ class WorkerSimulator(HttpUser):
|
||||
)
|
||||
time.sleep(random.uniform(2.0, 6.0))
|
||||
raise RescheduleTask()
|
||||
resp.failure(f"Pop failed: {resp.status_code}: {resp.text[:200]}")
|
||||
resp.failure(f"Pop failed: Status {resp.status_code}: {resp.text[:200]}")
|
||||
return
|
||||
pop_data = body or {}
|
||||
job_id = pop_data.get("id", None)
|
||||
|
||||
@@ -5,7 +5,6 @@
|
||||
"""Interrogation requestor and worker Locust users."""
|
||||
|
||||
import random
|
||||
import string
|
||||
import time
|
||||
|
||||
from locust import HttpUser, between, tag, task
|
||||
@@ -18,6 +17,7 @@ from ..helpers import (
|
||||
_is_too_many_workers,
|
||||
_pick_requestor_key,
|
||||
_pick_worker_key,
|
||||
_random_worker_name,
|
||||
_record_expected,
|
||||
_safe_json,
|
||||
)
|
||||
@@ -112,7 +112,7 @@ class InterrogationWorkerSimulator(HttpUser):
|
||||
wait_time = between(2, 5)
|
||||
|
||||
def create_worker_name(self):
|
||||
return f"StressAlchemist-{''.join(random.choices(string.ascii_lowercase, k=4))}"
|
||||
return _random_worker_name("StressAlchemist-", 4)
|
||||
|
||||
def on_start(self):
|
||||
self.worker_name = self.create_worker_name()
|
||||
@@ -151,7 +151,7 @@ class InterrogationWorkerSimulator(HttpUser):
|
||||
if resp.status_code == 429:
|
||||
resp.success()
|
||||
raise RescheduleTask()
|
||||
resp.failure(f"Interrogate pop failed: {resp.status_code}: {resp.text[:200]}")
|
||||
resp.failure(f"Interrogate pop failed: Status {resp.status_code}: {resp.text[:200]}")
|
||||
return
|
||||
data = body or {}
|
||||
forms = data.get("forms") or []
|
||||
|
||||
@@ -37,7 +37,6 @@ run's elapsed wall-clock time against durations configured at test start.
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import string
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
@@ -52,6 +51,7 @@ from ..helpers import (
|
||||
_headers,
|
||||
_is_expected_rc,
|
||||
_is_too_many_workers,
|
||||
_random_worker_name,
|
||||
_record_expected,
|
||||
_safe_json,
|
||||
)
|
||||
@@ -270,7 +270,7 @@ class ServingWorker(HttpUser):
|
||||
wait_time = constant_pacing(0.1)
|
||||
|
||||
def on_start(self) -> None:
|
||||
self.worker_name = f"QPWorker-{''.join(random.choices(string.ascii_lowercase, k=6))}"
|
||||
self.worker_name = _random_worker_name("QPWorker-", 6)
|
||||
keys = _config.get("worker_api_keys", [])
|
||||
# Each worker owns a stable key for its lifetime so worker identity does
|
||||
# not churn mid-run and confound the pop attribution.
|
||||
|
||||
@@ -5,7 +5,6 @@
|
||||
"""Text generation requestor and worker Locust users."""
|
||||
|
||||
import random
|
||||
import string
|
||||
import time
|
||||
import uuid
|
||||
|
||||
@@ -20,6 +19,7 @@ from ..helpers import (
|
||||
_pick_requestor_key,
|
||||
_pick_worker_key,
|
||||
_random_prompt,
|
||||
_random_worker_name,
|
||||
_record_expected,
|
||||
_safe_json,
|
||||
)
|
||||
@@ -143,7 +143,7 @@ class TextWorkerSimulator(HttpUser):
|
||||
wait_time = between(1, 4)
|
||||
|
||||
def create_worker_name(self):
|
||||
return f"StressTextWorker-{''.join(random.choices(string.ascii_lowercase, k=4))}"
|
||||
return _random_worker_name("StressTextWorker-", 4)
|
||||
|
||||
def on_start(self):
|
||||
self.worker_name = self.create_worker_name()
|
||||
@@ -188,7 +188,7 @@ class TextWorkerSimulator(HttpUser):
|
||||
self.environment, "POST", "/api/v2/generate/text/pop", resp.elapsed.total_seconds() * 1000, len(resp.content or b"")
|
||||
)
|
||||
raise RescheduleTask()
|
||||
resp.failure(f"Text pop failed: {resp.status_code}: {resp.text[:200]}")
|
||||
resp.failure(f"Text pop failed: Status {resp.status_code}: {resp.text[:200]}")
|
||||
return
|
||||
data = body or {}
|
||||
job_id = data.get("id")
|
||||
|
||||
@@ -0,0 +1,402 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-FileCopyrightText: 2026 Tazlin
|
||||
#
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
repo_root="$(cd -- "$script_dir/../.." && pwd)"
|
||||
|
||||
usage() {
|
||||
cat <<EOF
|
||||
Usage: bash tests/stress/run_chat_scale_demo.sh [options] <run-dir> [-- <extra locust args>]
|
||||
|
||||
Runs the chat front scale demo: brings up the test backends through
|
||||
tests/run_with_test_stack.sh, starts N Horde instances and the chat front,
|
||||
drives them with tests/stress/locustfile_chat.py through its client stages,
|
||||
samples the waitress gauges, the front's in-flight count, Redis and Postgres
|
||||
every few seconds, and writes samples.csv, summary.md and summary.json into
|
||||
<run-dir>. Every process it started is stopped on exit.
|
||||
|
||||
Options (each also read from the environment variable shown):
|
||||
--instances N Horde instances (CHAT_DEMO_HORDE_INSTANCES, 3)
|
||||
--horde-base-port P First Horde port; the rest follow (CHAT_DEMO_HORDE_BASE_PORT, 7101)
|
||||
--front-port P Front port (CHAT_DEMO_FRONT_PORT, 7002)
|
||||
--waitress-threads N Threads per Horde instance (CHAT_DEMO_WAITRESS_THREADS, 45)
|
||||
--workers N Simulated chat workers (CHAT_DEMO_WORKERS, 50)
|
||||
--stages LIST Client counts per stage (CHAT_DEMO_STAGES, 500,2000,5000)
|
||||
--stage-seconds S Length of each stage (CHAT_DEMO_STAGE_SECONDS, 120)
|
||||
--warmup-seconds S Workers-only warm-up (CHAT_DEMO_WARMUP_SECONDS, 20)
|
||||
--spawn-rate R Users started per second (CHAT_DEMO_SPAWN_RATE, 100)
|
||||
--sample-interval S Seconds between gauge samples (CHAT_DEMO_SAMPLE_INTERVAL, 5)
|
||||
--sample-only S Start the stack, Horde instances and front, sample for
|
||||
S seconds without Locust, then stop (no summary)
|
||||
--compose-project NAME Run the backends as their own Compose project, so the
|
||||
demo never touches the default test stack's data
|
||||
(CHAT_DEMO_COMPOSE_PROJECT, unset). Give it its own
|
||||
AI_HORDE_TEST_*_PORT values.
|
||||
--down Stop the backends when the run ends
|
||||
-h, --help Show this help
|
||||
|
||||
Other environment variables:
|
||||
CHAT_DEMO_PYTHON Interpreter (default: the repository's .venv)
|
||||
CHAT_DEMO_FRONT_MAX_INFLIGHT HORDE_OAI_MAX_INFLIGHT (default: 1.2 x the largest stage)
|
||||
CHAT_DEMO_RATELIMIT_DISABLED HORDE_TEST_RATELIMIT_DISABLED for the Horde instances (default 1;
|
||||
every client reaches the Horde from one address through the front)
|
||||
CHAT_DEMO_HORDE_VERBOSITY Horde log verbosity flag (default -vvvv)
|
||||
CHAT_DEMO_SIM_GEN_TIME_MIN/MAX Simulated generation seconds (default 3 and 8)
|
||||
|
||||
Example:
|
||||
bash tests/stress/run_chat_scale_demo.sh --stages 50,200 --stage-seconds 60 /tmp/chat-demo
|
||||
EOF
|
||||
}
|
||||
|
||||
instances="${CHAT_DEMO_HORDE_INSTANCES:-3}"
|
||||
horde_base_port="${CHAT_DEMO_HORDE_BASE_PORT:-7101}"
|
||||
front_port="${CHAT_DEMO_FRONT_PORT:-7002}"
|
||||
waitress_threads="${CHAT_DEMO_WAITRESS_THREADS:-45}"
|
||||
workers="${CHAT_DEMO_WORKERS:-50}"
|
||||
stages="${CHAT_DEMO_STAGES:-500,2000,5000}"
|
||||
stage_seconds="${CHAT_DEMO_STAGE_SECONDS:-120}"
|
||||
warmup_seconds="${CHAT_DEMO_WARMUP_SECONDS:-20}"
|
||||
spawn_rate="${CHAT_DEMO_SPAWN_RATE:-100}"
|
||||
sample_interval="${CHAT_DEMO_SAMPLE_INTERVAL:-5}"
|
||||
sample_only_seconds=""
|
||||
compose_project="${CHAT_DEMO_COMPOSE_PROJECT:-}"
|
||||
stack_down=0
|
||||
run_dir=""
|
||||
extra_locust_args=()
|
||||
original_args=("$@")
|
||||
|
||||
parse_args() {
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--instances) instances="$2"; shift ;;
|
||||
--horde-base-port) horde_base_port="$2"; shift ;;
|
||||
--front-port) front_port="$2"; shift ;;
|
||||
--waitress-threads) waitress_threads="$2"; shift ;;
|
||||
--workers) workers="$2"; shift ;;
|
||||
--stages) stages="$2"; shift ;;
|
||||
--stage-seconds) stage_seconds="$2"; shift ;;
|
||||
--warmup-seconds) warmup_seconds="$2"; shift ;;
|
||||
--spawn-rate) spawn_rate="$2"; shift ;;
|
||||
--sample-interval) sample_interval="$2"; shift ;;
|
||||
--sample-only) sample_only_seconds="$2"; shift ;;
|
||||
--compose-project) compose_project="$2"; shift ;;
|
||||
--down) stack_down=1 ;;
|
||||
-h|--help) usage; exit 0 ;;
|
||||
--) shift; extra_locust_args=("$@"); break ;;
|
||||
-*) echo "Unknown option: $1" >&2; usage >&2; exit 1 ;;
|
||||
*)
|
||||
if [[ -n "$run_dir" ]]; then
|
||||
echo "Only one run directory may be given." >&2
|
||||
exit 1
|
||||
fi
|
||||
run_dir="$1"
|
||||
;;
|
||||
esac
|
||||
shift
|
||||
done
|
||||
if [[ -z "$run_dir" ]]; then
|
||||
echo "A run directory is required." >&2
|
||||
usage >&2
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
parse_args "$@"
|
||||
|
||||
# The backends and their env come from run_with_test_stack.sh, which runs this
|
||||
# script again inside the injected environment.
|
||||
if [[ -z "${CHAT_DEMO_INSIDE_TEST_STACK:-}" ]]; then
|
||||
# run_with_test_stack.sh changes to the repository root, so a relative run
|
||||
# directory is resolved here, against the caller's directory.
|
||||
mkdir -p "$run_dir"
|
||||
CHAT_DEMO_RUN_DIR="$(cd -- "$run_dir" && pwd)"
|
||||
export CHAT_DEMO_RUN_DIR
|
||||
stack_args=()
|
||||
if [[ "$stack_down" == "1" ]]; then
|
||||
stack_args+=(--down)
|
||||
fi
|
||||
if [[ -n "$compose_project" ]]; then
|
||||
export COMPOSE_PROJECT_NAME="$compose_project"
|
||||
fi
|
||||
case "$(uname -s)" in
|
||||
MINGW*|MSYS*)
|
||||
# Git Bash rewrites arguments that look like POSIX paths, which turns the
|
||||
# in-container paths of the Garage bootstrap's `compose exec` into Windows
|
||||
# paths; those two are excluded, the host paths still convert.
|
||||
export MSYS2_ARG_CONV_EXCL="/garage;/etc/garage.toml"
|
||||
;;
|
||||
esac
|
||||
exec bash "$repo_root/tests/run_with_test_stack.sh" "${stack_args[@]}" -- \
|
||||
env CHAT_DEMO_INSIDE_TEST_STACK=1 bash "${BASH_SOURCE[0]}" "${original_args[@]}"
|
||||
fi
|
||||
|
||||
# Native Windows interpreters launched from Git Bash need Windows-form paths in
|
||||
# environment variables; elsewhere paths pass through unchanged.
|
||||
native_path() {
|
||||
if command -v cygpath >/dev/null 2>&1; then
|
||||
cygpath -m "$1"
|
||||
else
|
||||
printf '%s\n' "$1"
|
||||
fi
|
||||
}
|
||||
|
||||
resolve_python() {
|
||||
if [[ -n "${CHAT_DEMO_PYTHON:-}" ]]; then
|
||||
printf '%s\n' "$CHAT_DEMO_PYTHON"
|
||||
elif [[ -x "$repo_root/.venv/Scripts/python.exe" ]]; then
|
||||
printf '%s\n' "$repo_root/.venv/Scripts/python.exe"
|
||||
elif [[ -x "$repo_root/.venv/bin/python" ]]; then
|
||||
printf '%s\n' "$repo_root/.venv/bin/python"
|
||||
else
|
||||
command -v python
|
||||
fi
|
||||
}
|
||||
|
||||
started_pids=()
|
||||
started_names=()
|
||||
|
||||
# On Windows a venv's python.exe is a launcher that starts the interpreter as a
|
||||
# child process; ending only the launcher leaves the interpreter serving, so the
|
||||
# whole process tree is ended there. Elsewhere the process gets 15 seconds to
|
||||
# exit after SIGTERM before SIGKILL.
|
||||
stop_pid() {
|
||||
local pid="$1"
|
||||
local name="$2"
|
||||
local waited=0
|
||||
|
||||
if ! kill -0 "$pid" 2>/dev/null; then
|
||||
return 0
|
||||
fi
|
||||
echo "Stopping $name (pid $pid)"
|
||||
if [[ -r "/proc/$pid/winpid" ]]; then
|
||||
taskkill //F //T //PID "$(cat "/proc/$pid/winpid")" >/dev/null 2>&1 || true
|
||||
else
|
||||
kill -TERM "$pid" 2>/dev/null || true
|
||||
while kill -0 "$pid" 2>/dev/null && [[ "$waited" -lt 15 ]]; do
|
||||
sleep 1
|
||||
waited=$((waited + 1))
|
||||
done
|
||||
kill -KILL "$pid" 2>/dev/null || true
|
||||
fi
|
||||
wait "$pid" 2>/dev/null || true
|
||||
}
|
||||
|
||||
cleanup() {
|
||||
local index
|
||||
for ((index = ${#started_pids[@]} - 1; index >= 0; index--)); do
|
||||
stop_pid "${started_pids[$index]}" "${started_names[$index]}"
|
||||
done
|
||||
}
|
||||
|
||||
wait_for_url() {
|
||||
local url="$1"
|
||||
local name="$2"
|
||||
local timeout_seconds="$3"
|
||||
local pid="$4"
|
||||
local waited=0
|
||||
|
||||
until curl -sf -o /dev/null "$url"; do
|
||||
if ! kill -0 "$pid" 2>/dev/null; then
|
||||
echo "$name exited before answering at $url; see its log in $run_dir." >&2
|
||||
return 1
|
||||
fi
|
||||
if [[ "$waited" -ge "$timeout_seconds" ]]; then
|
||||
echo "$name did not answer at $url within ${timeout_seconds}s; see its log in $run_dir." >&2
|
||||
return 1
|
||||
fi
|
||||
sleep 2
|
||||
waited=$((waited + 2))
|
||||
done
|
||||
echo "$name is up at $url"
|
||||
}
|
||||
|
||||
port_is_free() {
|
||||
local port="$1"
|
||||
! curl -s -o /dev/null --max-time 2 "http://127.0.0.1:$port/"
|
||||
}
|
||||
|
||||
max_stage() {
|
||||
local largest=0
|
||||
local stage
|
||||
IFS=',' read -r -a stage_list <<<"$stages"
|
||||
for stage in "${stage_list[@]}"; do
|
||||
if [[ "$stage" -gt "$largest" ]]; then
|
||||
largest="$stage"
|
||||
fi
|
||||
done
|
||||
printf '%s\n' "$largest"
|
||||
}
|
||||
|
||||
require_stack_env() {
|
||||
local name
|
||||
local missing=()
|
||||
for name in POSTGRES_URL REDIS_IP REDIS_PORT R2_TRANSIENT_ACCOUNT; do
|
||||
if [[ -z "${!name:-}" ]]; then
|
||||
missing+=("$name")
|
||||
fi
|
||||
done
|
||||
if [[ "${#missing[@]}" -gt 0 ]]; then
|
||||
# Without the stack's env the Horde falls back to localhost defaults, which
|
||||
# can be a different stack's Redis; starting it there would mix the runs.
|
||||
echo "The test stack env is missing ${missing[*]} (the Garage bootstrap may have timed out); rerun once the stack is healthy." >&2
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
main() {
|
||||
trap cleanup EXIT
|
||||
trap 'exit 130' INT TERM
|
||||
require_stack_env
|
||||
|
||||
run_dir="${CHAT_DEMO_RUN_DIR:-$run_dir}"
|
||||
mkdir -p "$run_dir"
|
||||
run_dir="$(cd -- "$run_dir" && pwd)"
|
||||
local python
|
||||
python="$(resolve_python)"
|
||||
local largest_stage
|
||||
largest_stage="$(max_stage)"
|
||||
local front_max_inflight="${CHAT_DEMO_FRONT_MAX_INFLIGHT:-$((largest_stage + largest_stage / 5 + 100))}"
|
||||
local horde_verbosity="${CHAT_DEMO_HORDE_VERBOSITY:--vvvv}"
|
||||
|
||||
export HORDE_TEST_RATELIMIT_DISABLED="${CHAT_DEMO_RATELIMIT_DISABLED:-1}"
|
||||
export REDIS_SERVERS="${REDIS_SERVERS:-[\"${REDIS_IP:-127.0.0.1}\"]}"
|
||||
|
||||
# Thousands of concurrent client sockets exceed common default descriptor
|
||||
# limits; raise the soft limit where the platform allows and record it.
|
||||
ulimit -n 65535 2>/dev/null || ulimit -n "$(ulimit -Hn)" 2>/dev/null || true
|
||||
|
||||
local horde_ports=()
|
||||
local horde_urls=()
|
||||
local horde_api_urls=()
|
||||
local index
|
||||
for ((index = 0; index < instances; index++)); do
|
||||
horde_ports+=("$((horde_base_port + index))")
|
||||
horde_urls+=("http://127.0.0.1:$((horde_base_port + index))")
|
||||
horde_api_urls+=("http://127.0.0.1:$((horde_base_port + index))/api")
|
||||
done
|
||||
local horde_urls_csv horde_api_urls_csv
|
||||
horde_urls_csv="$(IFS=,; printf '%s' "${horde_urls[*]}")"
|
||||
horde_api_urls_csv="$(IFS=,; printf '%s' "${horde_api_urls[*]}")"
|
||||
local front_url="http://127.0.0.1:$front_port"
|
||||
|
||||
for port in "${horde_ports[@]}" "$front_port"; do
|
||||
if ! port_is_free "$port"; then
|
||||
echo "Port $port already answers HTTP; stop that process or choose other ports." >&2
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
|
||||
cat >"$run_dir/run.env" <<EOF
|
||||
instances=$instances
|
||||
horde_ports=${horde_ports[*]}
|
||||
front_port=$front_port
|
||||
waitress_threads=$waitress_threads
|
||||
workers=$workers
|
||||
stages=$stages
|
||||
stage_seconds=$stage_seconds
|
||||
warmup_seconds=$warmup_seconds
|
||||
spawn_rate=$spawn_rate
|
||||
sample_interval=$sample_interval
|
||||
front_max_inflight=$front_max_inflight
|
||||
ratelimit_disabled=$HORDE_TEST_RATELIMIT_DISABLED
|
||||
compose_project=${COMPOSE_PROJECT_NAME:-}
|
||||
open_file_limit=$(ulimit -n)
|
||||
uname=$(uname -s)
|
||||
EOF
|
||||
|
||||
local quorum_flag=()
|
||||
local -A horde_pids=()
|
||||
for ((index = 0; index < instances; index++)); do
|
||||
local port="${horde_ports[$index]}"
|
||||
quorum_flag=()
|
||||
if [[ "$index" == "0" ]]; then
|
||||
# One instance holds the quorum and runs the background threads, as in
|
||||
# the two-node CI job; the others serve requests only.
|
||||
quorum_flag=(--quorum)
|
||||
fi
|
||||
mkdir -p "$run_dir/horde-$port"
|
||||
# Each instance runs in its own directory so its rotating horde.log file
|
||||
# is not shared with the others.
|
||||
(cd "$run_dir/horde-$port" && exec "$python" "$(native_path "$repo_root/server.py")" "$horde_verbosity" -i \
|
||||
--horde stable --allow_all_ips --port "$port" --waitress_threads "$waitress_threads" "${quorum_flag[@]}") \
|
||||
>"$run_dir/horde-$port.log" 2>&1 &
|
||||
started_pids+=("$!")
|
||||
started_names+=("horde-$port")
|
||||
horde_pids[$port]="$!"
|
||||
done
|
||||
for port in "${horde_ports[@]}"; do
|
||||
wait_for_url "http://127.0.0.1:$port/api/v2/status/heartbeat" "Horde $port" 300 "${horde_pids[$port]}"
|
||||
done
|
||||
|
||||
mkdir -p "$run_dir/front"
|
||||
(cd "$run_dir/front" && exec env \
|
||||
PYTHONPATH="$(native_path "$repo_root")" \
|
||||
HORDE_OAI_HORDE_URLS="$horde_api_urls_csv" \
|
||||
HORDE_OAI_LISTEN=127.0.0.1 \
|
||||
HORDE_OAI_PORT="$front_port" \
|
||||
HORDE_OAI_REDIS_URL="redis://${REDIS_IP:-127.0.0.1}:${REDIS_PORT:-6379}/7" \
|
||||
HORDE_OAI_MAX_INFLIGHT="$front_max_inflight" \
|
||||
"$python" "$(native_path "$script_dir/chat_scale_demo.py")" front) \
|
||||
>"$run_dir/front.log" 2>&1 &
|
||||
local front_pid="$!"
|
||||
started_pids+=("$front_pid")
|
||||
started_names+=("front")
|
||||
wait_for_url "$front_url/healthz" "Front" 120 "$front_pid"
|
||||
|
||||
local postgres_dsn="postgresql://${POSTGRES_USER:-postgres}@${POSTGRES_URL:-localhost:5432/postgres}"
|
||||
local sampler_args=(
|
||||
sample
|
||||
--horde-urls "$horde_urls_csv"
|
||||
--front-url "$front_url"
|
||||
--redis-host "${REDIS_IP:-127.0.0.1}"
|
||||
--redis-port "${REDIS_PORT:-6379}"
|
||||
--postgres-dsn "$postgres_dsn"
|
||||
--interval "$sample_interval"
|
||||
--out "$(native_path "$run_dir/samples.csv")"
|
||||
)
|
||||
if [[ -n "$sample_only_seconds" ]]; then
|
||||
sampler_args+=(--duration "$sample_only_seconds")
|
||||
fi
|
||||
"$python" "$(native_path "$script_dir/chat_scale_demo.py")" "${sampler_args[@]}" >"$run_dir/sampler.log" 2>&1 &
|
||||
local sampler_pid="$!"
|
||||
started_pids+=("$sampler_pid")
|
||||
started_names+=("sampler")
|
||||
|
||||
if [[ -n "$sample_only_seconds" ]]; then
|
||||
echo "Sampling for ${sample_only_seconds}s without load."
|
||||
wait "$sampler_pid" || true
|
||||
echo "Samples written to $run_dir/samples.csv"
|
||||
return 0
|
||||
fi
|
||||
|
||||
local bootstrap_requestors=$(((largest_stage + 24) / 25 + 2))
|
||||
local bootstrap_workers=$(((workers + 2) / 3 + 2))
|
||||
local locust_status=0
|
||||
(cd "$repo_root" && "$python" -m locust -f tests/stress/locustfile_chat.py \
|
||||
--host "${horde_urls[0]}" --headless \
|
||||
--chat-front-host "$front_url" \
|
||||
--chat-horde-urls "$horde_urls_csv" \
|
||||
--chat-workers "$workers" \
|
||||
--chat-stages "$stages" \
|
||||
--chat-stage-seconds "$stage_seconds" \
|
||||
--chat-warmup-seconds "$warmup_seconds" \
|
||||
--chat-spawn-rate "$spawn_rate" \
|
||||
--sim-gen-time-min "${CHAT_DEMO_SIM_GEN_TIME_MIN:-3}" \
|
||||
--sim-gen-time-max "${CHAT_DEMO_SIM_GEN_TIME_MAX:-8}" \
|
||||
--bootstrap-requestors "$bootstrap_requestors" \
|
||||
--bootstrap-workers "$bootstrap_workers" \
|
||||
--chat-evidence-path "$(native_path "$run_dir/chat_evidence.jsonl")" \
|
||||
--csv "$(native_path "$run_dir/locust")" --csv-full-history \
|
||||
--logfile "$(native_path "$run_dir/locust.log")" \
|
||||
"${extra_locust_args[@]}") >"$run_dir/locust.out" 2>&1 || locust_status=$?
|
||||
echo "Locust exited with status $locust_status (non-zero when any request failed; see summary.md)."
|
||||
|
||||
stop_pid "$sampler_pid" "sampler"
|
||||
"$python" "$(native_path "$script_dir/chat_scale_demo.py")" summarize --run-dir "$(native_path "$run_dir")" | tee "$run_dir/summary.log"
|
||||
}
|
||||
|
||||
main
|
||||
@@ -0,0 +1,520 @@
|
||||
# SPDX-FileCopyrightText: 2026 Tazlin
|
||||
#
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
"""The stress suite's chat scenarios, response checks and front client against the front itself, run in process
|
||||
against a mocked Horde.
|
||||
|
||||
Every refusal a scenario expects is the one the front sends, param and type included. Every request a scenario expects
|
||||
queued is forwarded as sent, and the reply the suite's worker script gives for the forwarded body, delivered through
|
||||
the front, meets the scenario's expectation. The front client's own judgement runs against a front served over HTTP, so
|
||||
its stream goes through ``requests.iter_lines`` with the front's real keepalive comments.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import fakeredis
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
import uvicorn
|
||||
from opentelemetry.sdk.metrics import MeterProvider
|
||||
|
||||
os.environ.setdefault("LOCUST_SKIP_MONKEY_PATCH", "1")
|
||||
|
||||
from locust.env import Environment
|
||||
|
||||
from horde.oai.app import create_app
|
||||
from horde.oai.metrics import METER_NAME
|
||||
from horde.oai.settings import FrontSettings
|
||||
from tests.stress.locustsuite.chat_checks import ErrorExpectation, FrontVerdict, StreamTranscript, judge_front_response
|
||||
from tests.stress.locustsuite.chat_front_client import ChatMixFrontClient, build_request_body
|
||||
from tests.stress.locustsuite.chat_mix_settings import RunnerRole, RunTopology, chat_mix_settings
|
||||
from tests.stress.locustsuite.chat_scenarios import (
|
||||
SCENARIO_BUILDERS,
|
||||
ChatModelNames,
|
||||
ChatScenario,
|
||||
ScenarioContext,
|
||||
ScenarioGroup,
|
||||
build_authorization_headers,
|
||||
gap_forced_strict,
|
||||
nonconformant_schema_strict,
|
||||
oversized_body,
|
||||
plain_prompt,
|
||||
tools_auto,
|
||||
)
|
||||
from tests.stress.locustsuite.chat_worker_script import build_scripted_reply
|
||||
from tests.unit.oai.conftest import (
|
||||
FAST_TEST_SETTINGS,
|
||||
HORDE_URL,
|
||||
MODEL_NAME,
|
||||
REQUEST_ID,
|
||||
FakeHorde,
|
||||
StartFront,
|
||||
done_status,
|
||||
waiting_status,
|
||||
)
|
||||
from tests.unit.test_stress_chat_mix import chat_mix_options
|
||||
|
||||
CONTEXT: ScenarioContext = ScenarioContext(
|
||||
models=ChatModelNames(conformant=MODEL_NAME, limited=MODEL_NAME, flaky=MODEL_NAME),
|
||||
worker_max_length=512,
|
||||
nonce="0123456789ab",
|
||||
)
|
||||
"""Every persona's model is the one the mocked Horde serves, so a completion names the requested model."""
|
||||
|
||||
API_KEY: str = "stress-test-key"
|
||||
"""The requestor key the scenarios authenticate with."""
|
||||
|
||||
ANSWERED_BY_THE_HORDE: frozenset[str] = frozenset({"unknown_model", "unknown_api_key"})
|
||||
"""Refusals the Horde, not the front, decides; the mocked Horde accepts everything, so they are left out here."""
|
||||
|
||||
ALL_SCENARIOS: list[ChatScenario] = [builder(CONTEXT) for builders in SCENARIO_BUILDERS.values() for builder in builders]
|
||||
"""Every scenario of the catalogue."""
|
||||
|
||||
FRONT_REFUSED_SCENARIOS: list[ChatScenario] = [
|
||||
scenario
|
||||
for scenario in ALL_SCENARIOS
|
||||
if scenario.persona is None and not scenario.oversized_body and scenario.name not in ANSWERED_BY_THE_HORDE
|
||||
]
|
||||
"""The scenarios the front refuses before submitting."""
|
||||
|
||||
QUEUED_SCENARIOS: list[ChatScenario] = [scenario for scenario in ALL_SCENARIOS if scenario.persona is not None]
|
||||
"""The scenarios the front submits to the Horde."""
|
||||
|
||||
NATIVE_REQUEST_KEYS: frozenset[str] = frozenset(
|
||||
{
|
||||
"model",
|
||||
"messages",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"frequency_penalty",
|
||||
"presence_penalty",
|
||||
"stop",
|
||||
"seed",
|
||||
"n",
|
||||
"tools",
|
||||
"tool_choice",
|
||||
"parallel_tool_calls",
|
||||
"response_format",
|
||||
"horde",
|
||||
},
|
||||
)
|
||||
"""The keys the native chat route accepts; nothing else may be forwarded."""
|
||||
|
||||
FORWARDED_AS_SENT: tuple[str, ...] = ("tools", "tool_choice", "parallel_tool_calls", "response_format", "seed", "stop", "temperature")
|
||||
"""Keys the front forwards unchanged whenever the scenario sends them at a value that changes the answer."""
|
||||
|
||||
HORDE_JPEG_DATA_URL: str = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQ=="
|
||||
"""What the Horde hands a ``base64`` worker for an image: the stored JPEG as a data URL."""
|
||||
|
||||
FAST_GRACE_SETTINGS: FrontSettings = FAST_TEST_SETTINGS.model_copy(update={"fail_fast_grace_seconds": 0.2})
|
||||
"""Test settings whose no-eligible-worker grace is short enough for a test to wait it out."""
|
||||
|
||||
|
||||
def scenario_headers(scenario: ChatScenario) -> dict[str, str]:
|
||||
"""Return the headers the front client sends with a scenario."""
|
||||
return {**build_authorization_headers(scenario.authorization, api_key=API_KEY), **scenario.extra_headers}
|
||||
|
||||
|
||||
def streamed_body(scenario: ChatScenario, *, include_usage: bool) -> dict[str, Any]:
|
||||
"""Return a scenario's body as a streamed request."""
|
||||
stream_options = {"stream_options": {"include_usage": True}} if include_usage else {}
|
||||
return {**scenario.body, "stream": True, **stream_options}
|
||||
|
||||
|
||||
def read_stream(response: httpx.Response) -> StreamTranscript:
|
||||
"""Return the transcript of a whole streamed response."""
|
||||
transcript = StreamTranscript()
|
||||
for line_index, line in enumerate(response.text.splitlines()):
|
||||
transcript.feed_line(line, float(line_index))
|
||||
return transcript
|
||||
|
||||
|
||||
def judge(scenario: ChatScenario, response: httpx.Response, *, streamed: bool, include_usage: bool = False) -> FrontVerdict:
|
||||
"""Return the suite's verdict on a front response to ``scenario``."""
|
||||
stream_started = streamed and response.status_code == 200
|
||||
return judge_front_response(
|
||||
scenario.expectation,
|
||||
scenario.alternative_error,
|
||||
status_code=response.status_code,
|
||||
headers=response.headers,
|
||||
body=None if stream_started else response.json(),
|
||||
transcript=read_stream(response) if stream_started else None,
|
||||
streamed=streamed,
|
||||
include_usage=include_usage,
|
||||
cors_expected="Origin" in scenario.extra_headers,
|
||||
)
|
||||
|
||||
|
||||
def dispatch_like_the_horde(native_body: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return the chat job payload the Horde builds from a forwarded body: tools narrowed for ``none`` and a function
|
||||
choice, and image parts replaced by the stored JPEG."""
|
||||
payload = {
|
||||
payload_key: native_body[payload_key]
|
||||
for payload_key in ("messages", "tools", "tool_choice", "parallel_tool_calls", "response_format")
|
||||
if payload_key in native_body
|
||||
}
|
||||
payload["messages"] = copy.deepcopy(payload["messages"])
|
||||
for message in payload["messages"]:
|
||||
for part in message.get("content") if isinstance(message.get("content"), list) else []:
|
||||
if part.get("type") == "image_url":
|
||||
part["image_url"] = {**part["image_url"], "url": HORDE_JPEG_DATA_URL}
|
||||
tool_choice = payload.get("tool_choice")
|
||||
if tool_choice == "none":
|
||||
for narrowed_key in ("tools", "tool_choice", "parallel_tool_calls"):
|
||||
payload.pop(narrowed_key, None)
|
||||
elif isinstance(tool_choice, dict):
|
||||
chosen_name = tool_choice["function"]["name"]
|
||||
payload["tools"] = [tool for tool in payload["tools"] if tool["function"]["name"] == chosen_name]
|
||||
payload["tool_choice"] = "required"
|
||||
return payload
|
||||
|
||||
|
||||
def stored_generation(native_body: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return the generation the Horde stores for the suite worker's reply to a forwarded body, as the native status
|
||||
shows it: a single call kept under ``strict`` ``parallel_tool_calls: false``, and ``stop`` for a forced function."""
|
||||
reply = build_scripted_reply(dispatch_like_the_horde(native_body))
|
||||
conformance = (native_body.get("horde") or {}).get("conformance", "strict")
|
||||
tool_calls = reply.tool_calls
|
||||
if tool_calls and native_body.get("parallel_tool_calls") is False and conformance == "strict":
|
||||
tool_calls = tool_calls[:1]
|
||||
finish_reason = reply.finish_reason
|
||||
if tool_calls and isinstance(native_body.get("tool_choice"), dict):
|
||||
finish_reason = "stop"
|
||||
message: dict[str, Any] = {"role": "assistant", "content": reply.content}
|
||||
if tool_calls:
|
||||
message["tool_calls"] = tool_calls
|
||||
status = done_status(content=reply.content or "", finish_reason=finish_reason)
|
||||
status["generations"][0]["message"] = message
|
||||
return status
|
||||
|
||||
|
||||
def refusal_status(expectation: ErrorExpectation) -> dict[str, Any]:
|
||||
"""Return the native status that leads the front to the refusal a queued scenario expects."""
|
||||
if expectation.rc is not None:
|
||||
reason = " ".join(expectation.message_fragments) or "The reply broke the request."
|
||||
return {**waiting_status(), "faulted": True, "fault_reason": reason, "fault_rc": expectation.rc}
|
||||
if "no_worker_supports_features" in expectation.codes:
|
||||
(named,) = expectation.message_fragments
|
||||
exceeded = [named] if named == "tools_max_temperature" else []
|
||||
missing = [] if exceeded else [named]
|
||||
gap = {"kind": "declaration_gap", "missing_declarations": missing, "exceeded_limits": exceeded}
|
||||
return {**waiting_status(is_possible=False), "impossible_reason": gap}
|
||||
if "no_worker_supports_limit" in expectation.codes:
|
||||
gap = {"kind": "declaration_gap", "missing_declarations": [], "exceeded_limits": ["max_length"]}
|
||||
return {**waiting_status(is_possible=False), "impossible_reason": gap}
|
||||
return waiting_status(is_possible=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scenario", FRONT_REFUSED_SCENARIOS, ids=lambda scenario: scenario.name)
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_refusal_scenario_expects_the_refusal_the_front_sends(start_front: StartFront, scenario: ChatScenario) -> None:
|
||||
assert isinstance(scenario.expectation, ErrorExpectation)
|
||||
async with start_front() as front:
|
||||
response = await front.client.post("/v1/chat/completions", json=scenario.body, headers=scenario_headers(scenario))
|
||||
|
||||
assert judge(scenario, response, streamed=False).violations == []
|
||||
assert front.horde.submit_route.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scenario", QUEUED_SCENARIOS, ids=lambda scenario: scenario.name)
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_queued_scenario_is_forwarded_as_sent_and_its_reply_meets_the_expectation(
|
||||
start_front: StartFront,
|
||||
scenario: ChatScenario,
|
||||
) -> None:
|
||||
async with start_front(FAST_GRACE_SETTINGS) as front:
|
||||
if isinstance(scenario.expectation, ErrorExpectation):
|
||||
front.horde.status_payload = refusal_status(scenario.expectation)
|
||||
response = await front.client.post("/v1/chat/completions", json=scenario.body, headers=scenario_headers(scenario))
|
||||
else:
|
||||
front.horde.status_payload = waiting_status()
|
||||
front.horde.submit_route.side_effect = lambda request: _answer_submit_with_reply(front.horde, request)
|
||||
response = await front.client.post("/v1/chat/completions", json=scenario.body, headers=scenario_headers(scenario))
|
||||
forwarded = front.horde.submitted_body()
|
||||
|
||||
assert set(forwarded) <= NATIVE_REQUEST_KEYS, sorted(set(forwarded) - NATIVE_REQUEST_KEYS)
|
||||
assert forwarded["model"] == scenario.body["model"]
|
||||
assert len(forwarded["messages"]) == len(scenario.body["messages"])
|
||||
assert forwarded["horde"]["operation_key"]
|
||||
assert CONTEXT.nonce in json.dumps(forwarded["messages"])
|
||||
for request_key in FORWARDED_AS_SENT:
|
||||
if scenario.body.get(request_key) not in (None, [], {"type": "text"}):
|
||||
assert forwarded[request_key] == scenario.body[request_key], request_key
|
||||
assert judge(scenario, response, streamed=False).violations == []
|
||||
|
||||
|
||||
CONFORMANCE_SCENARIOS: list[ChatScenario] = [scenario for scenario in QUEUED_SCENARIOS if scenario.group == ScenarioGroup.CONFORMANCE]
|
||||
"""The scenarios that send a request under one conformance mode, some of them twins under the other."""
|
||||
|
||||
|
||||
@pytest.mark.parametrize("streamed", [False, True])
|
||||
@pytest.mark.parametrize("scenario", CONFORMANCE_SCENARIOS, ids=lambda scenario: scenario.name)
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_front_forwards_the_scenarios_conformance_streamed_or_not(
|
||||
start_front: StartFront,
|
||||
scenario: ChatScenario,
|
||||
streamed: bool,
|
||||
) -> None:
|
||||
body = build_request_body(scenario, streamed=streamed, include_usage=streamed)
|
||||
async with start_front(FAST_GRACE_SETTINGS) as front:
|
||||
front.horde.status_payload = (
|
||||
refusal_status(scenario.expectation) if isinstance(scenario.expectation, ErrorExpectation) else done_status()
|
||||
)
|
||||
await front.client.post("/v1/chat/completions", json=body, headers=scenario_headers(scenario))
|
||||
forwarded = front.horde.submitted_body()
|
||||
|
||||
expected_conformance = (scenario.body.get("horde") or {}).get("conformance", "strict")
|
||||
assert forwarded["horde"].get("conformance", "strict") == expected_conformance
|
||||
|
||||
|
||||
def _answer_submit_with_reply(horde: FakeHorde, request: httpx.Request) -> httpx.Response:
|
||||
"""Answer a submit and make the status report the suite worker's reply to the submitted body."""
|
||||
horde.status_payload = stored_generation(json.loads(request.content))
|
||||
return httpx.Response(202, json={"id": REQUEST_ID, "kudos": 12.5})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("streamed", "include_usage"), [(False, False), (True, False), (True, True)])
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_checks_accept_the_fronts_text_completion(start_front: StartFront, streamed: bool, include_usage: bool) -> None:
|
||||
scenario = plain_prompt(CONTEXT)
|
||||
async with start_front() as front:
|
||||
front.horde.status_payload = done_status(content=f"Hello {CONTEXT.nonce}", finish_reason="stop")
|
||||
body = streamed_body(scenario, include_usage=include_usage) if streamed else scenario.body
|
||||
response = await front.client.post("/v1/chat/completions", json=body, headers=scenario_headers(scenario))
|
||||
|
||||
assert judge(scenario, response, streamed=streamed, include_usage=include_usage).violations == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("streamed", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_checks_accept_the_fronts_tool_call_completion(start_front: StartFront, streamed: bool) -> None:
|
||||
scenario = tools_auto(CONTEXT)
|
||||
weather_call = {
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": json.dumps({"city": CONTEXT.nonce})},
|
||||
}
|
||||
status = done_status(content="", finish_reason="tool_calls")
|
||||
status["generations"][0]["message"] = {"role": "assistant", "content": None, "tool_calls": [weather_call]}
|
||||
async with start_front() as front:
|
||||
front.horde.status_payload = status
|
||||
body = streamed_body(scenario, include_usage=True) if streamed else scenario.body
|
||||
response = await front.client.post("/v1/chat/completions", json=body, headers=scenario_headers(scenario))
|
||||
|
||||
assert judge(scenario, response, streamed=streamed, include_usage=streamed).violations == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("streamed", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_checks_match_the_fronts_fault_for_a_refused_reply(start_front: StartFront, streamed: bool) -> None:
|
||||
scenario = nonconformant_schema_strict(CONTEXT)
|
||||
async with start_front() as front:
|
||||
front.horde.status_payload = refusal_status(scenario.expectation)
|
||||
body = streamed_body(scenario, include_usage=False) if streamed else scenario.body
|
||||
response = await front.client.post("/v1/chat/completions", json=body, headers=scenario_headers(scenario))
|
||||
|
||||
verdict = judge(scenario, response, streamed=streamed)
|
||||
assert verdict.violations == []
|
||||
assert verdict.matched_error is scenario.expectation
|
||||
assert verdict.error_envelope is not None
|
||||
assert verdict.error_envelope["error"]["metadata"]["request_id"] == REQUEST_ID
|
||||
|
||||
|
||||
@pytest.mark.parametrize("streamed", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_checks_match_the_fronts_declaration_gap(start_front: StartFront, streamed: bool) -> None:
|
||||
scenario = gap_forced_strict(CONTEXT)
|
||||
async with start_front() as front:
|
||||
front.horde.status_payload = refusal_status(scenario.expectation)
|
||||
body = streamed_body(scenario, include_usage=False) if streamed else scenario.body
|
||||
response = await front.client.post("/v1/chat/completions", json=body, headers=scenario_headers(scenario))
|
||||
|
||||
assert judge(scenario, response, streamed=streamed).violations == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_checks_refuse_a_completion_where_a_refusal_was_promised(start_front: StartFront) -> None:
|
||||
scenario = nonconformant_schema_strict(CONTEXT)
|
||||
async with start_front() as front:
|
||||
front.horde.status_payload = done_status(finish_reason="stop")
|
||||
response = await front.client.post("/v1/chat/completions", json=scenario.body, headers=scenario_headers(scenario))
|
||||
|
||||
violations = judge(scenario, response, streamed=False).violations
|
||||
assert violations and "got a completion" in violations[0]
|
||||
|
||||
|
||||
SERVED_FRONT_SETTINGS: FrontSettings = FAST_TEST_SETTINGS.model_copy(update={"fallback_poll_seconds": 0.1, "max_body_bytes": 1024})
|
||||
"""Settings for a front served over HTTP without Redis: a fast fallback poll replaces the completion signal, and a small
|
||||
body limit lets the oversized scenario run with a small body."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServedFront:
|
||||
"""Represents a front served over HTTP on a local port, with its mocked Horde."""
|
||||
|
||||
url: str
|
||||
"""The front's base URL."""
|
||||
horde: FakeHorde
|
||||
"""The mocked Horde."""
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
"""Return a local port no one listens on."""
|
||||
with socket.socket() as probe:
|
||||
probe.bind(("127.0.0.1", 0))
|
||||
port: int = probe.getsockname()[1]
|
||||
return port
|
||||
|
||||
|
||||
@contextmanager
|
||||
def serve_front() -> Iterator[ServedFront]:
|
||||
"""Run a front with uvicorn on a thread against a mocked Horde, for a client that speaks real HTTP."""
|
||||
with respx.mock(assert_all_called=False) as router:
|
||||
horde = FakeHorde(router, fakeredis.FakeAsyncRedis())
|
||||
app = create_app(
|
||||
SERVED_FRONT_SETTINGS,
|
||||
http_client=httpx.AsyncClient(),
|
||||
cancel_http_client=httpx.AsyncClient(),
|
||||
redis_client=None,
|
||||
meter=MeterProvider().get_meter(METER_NAME),
|
||||
)
|
||||
port = _free_port()
|
||||
server = uvicorn.Server(uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning"))
|
||||
server_thread = threading.Thread(target=server.run, daemon=True)
|
||||
server_thread.start()
|
||||
deadline = time.monotonic() + 10
|
||||
while not server.started:
|
||||
assert time.monotonic() < deadline, "The served front did not start."
|
||||
time.sleep(0.02)
|
||||
try:
|
||||
yield ServedFront(url=f"http://127.0.0.1:{port}", horde=horde)
|
||||
finally:
|
||||
server.should_exit = True
|
||||
server_thread.join(timeout=10)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def served_front() -> Iterator[ServedFront]:
|
||||
"""Serve a front for one test and configure the chat mix to talk to it."""
|
||||
with serve_front() as front:
|
||||
chat_mix_settings.configure(
|
||||
chat_mix_options(front_host=front.url, front_clients=1), horde_host=HORDE_URL, topology=RunTopology(role=RunnerRole.LOCAL)
|
||||
)
|
||||
chat_mix_settings.verify_faults = False
|
||||
chat_mix_settings.oversized_body_bytes = 4096
|
||||
yield front
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecordedClient:
|
||||
"""Represents a front client and the request events it fired."""
|
||||
|
||||
client: ChatMixFrontClient
|
||||
"""The Locust user."""
|
||||
events: list[dict[str, Any]]
|
||||
"""Every request event, in order."""
|
||||
|
||||
def failures(self) -> list[str]:
|
||||
"""Return the messages of the events that failed."""
|
||||
return [str(event["exception"]) for event in self.events if event["exception"] is not None]
|
||||
|
||||
def names(self) -> list[str]:
|
||||
"""Return the names of every event."""
|
||||
return [event["name"] for event in self.events]
|
||||
|
||||
|
||||
def build_front_client(front: ServedFront) -> RecordedClient:
|
||||
"""Create a front client for the served front, recording its events."""
|
||||
# Locust reads the host from the user class, which a run sets from --host.
|
||||
hosted_client_class = type("HostedChatMixFrontClient", (ChatMixFrontClient,), {"host": front.url})
|
||||
environment = Environment(user_classes=[hosted_client_class])
|
||||
events: list[dict[str, Any]] = []
|
||||
environment.events.request.add_listener(lambda **event: events.append(event))
|
||||
front_client = hosted_client_class(environment)
|
||||
front_client.api_key = API_KEY
|
||||
front_client.rng = random.Random(7)
|
||||
return RecordedClient(client=front_client, events=events)
|
||||
|
||||
|
||||
def served_scenario(builder: Any) -> ChatScenario:
|
||||
"""Build a scenario for the mocked Horde's model."""
|
||||
return builder(CONTEXT)
|
||||
|
||||
|
||||
def test_a_streamed_reply_read_through_iter_lines_with_keepalives_passes(served_front: ServedFront) -> None:
|
||||
chat_mix_settings.stream_chance = 1.0
|
||||
chat_mix_settings.include_usage_chance = 1.0
|
||||
recorded = build_front_client(served_front)
|
||||
finished = threading.Timer(
|
||||
0.4, lambda: setattr(served_front.horde, "status_payload", done_status(content=f"Hello {CONTEXT.nonce}", finish_reason="stop"))
|
||||
)
|
||||
finished.start()
|
||||
|
||||
recorded.client.send_scenario(served_scenario(plain_prompt))
|
||||
finished.join()
|
||||
|
||||
assert recorded.failures() == []
|
||||
assert "/v1/chat/completions [plain_prompt] [stream]" in recorded.names()
|
||||
assert "[time-to-first-keepalive] stream" in recorded.names()
|
||||
assert "[time-to-first-content] stream" in recorded.names()
|
||||
assert chat_mix_settings.coverage.completed_front_scenarios == 1
|
||||
|
||||
|
||||
def test_a_reply_for_another_request_is_a_contract_violation(served_front: ServedFront) -> None:
|
||||
chat_mix_settings.stream_chance = 0.0
|
||||
served_front.horde.status_payload = done_status(content="Hello from another request", finish_reason="stop")
|
||||
recorded = build_front_client(served_front)
|
||||
|
||||
recorded.client.send_scenario(served_scenario(plain_prompt))
|
||||
|
||||
(failure,) = recorded.failures()
|
||||
assert "Contract violation [plain_prompt]" in failure
|
||||
assert "nonce" in failure
|
||||
|
||||
|
||||
def test_an_expected_fault_is_a_success_of_the_scenario(served_front: ServedFront) -> None:
|
||||
chat_mix_settings.stream_chance = 0.0
|
||||
scenario = served_scenario(nonconformant_schema_strict)
|
||||
served_front.horde.status_payload = refusal_status(scenario.expectation)
|
||||
recorded = build_front_client(served_front)
|
||||
|
||||
recorded.client.send_scenario(scenario)
|
||||
|
||||
assert recorded.failures() == []
|
||||
assert recorded.names() == ["/v1/chat/completions [nonconformant_schema_strict] [non-stream]"]
|
||||
|
||||
|
||||
def test_an_unreachable_horde_is_an_environment_condition(served_front: ServedFront) -> None:
|
||||
chat_mix_settings.stream_chance = 0.0
|
||||
served_front.horde.submit_route.side_effect = httpx.ConnectError("refused")
|
||||
recorded = build_front_client(served_front)
|
||||
|
||||
recorded.client.send_scenario(served_scenario(plain_prompt))
|
||||
|
||||
(failure,) = recorded.failures()
|
||||
assert "Environment condition [horde_unavailable]" in failure
|
||||
|
||||
|
||||
def test_an_oversized_body_ends_in_a_413_or_a_benign_upload_reset(served_front: ServedFront) -> None:
|
||||
recorded = build_front_client(served_front)
|
||||
|
||||
recorded.client.send_scenario(served_scenario(oversized_body))
|
||||
|
||||
assert recorded.failures() == []
|
||||
assert served_front.horde.submit_route.call_count == 0
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,245 @@
|
||||
# SPDX-FileCopyrightText: 2026 Tazlin
|
||||
#
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
|
||||
"""The chat scale demo's simulated worker in its progress mode: the delta split, the posting cadence and how each
|
||||
progress status ends or continues a job's posting."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from datetime import timedelta
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("LOCUST_SKIP_MONKEY_PATCH", "1")
|
||||
|
||||
from locust.env import Environment
|
||||
|
||||
from tests.stress.locustsuite.chat_wire import TEXT_POP_PATH, TEXT_PROGRESS_PATH, TEXT_SUBMIT_PATH
|
||||
from tests.stress.locustsuite.users import chat as chat_users
|
||||
from tests.stress.locustsuite.users.chat import (
|
||||
WORKER_GENERATION_TEXT,
|
||||
ChatWorkerSimulator,
|
||||
chat_settings,
|
||||
count_progress_deltas,
|
||||
progress_request_name,
|
||||
split_reply_into_deltas,
|
||||
)
|
||||
|
||||
GENERATION_ID: str = "00000000-0000-0000-0000-000000000001"
|
||||
"""The generation id every popped test job carries."""
|
||||
|
||||
|
||||
def scenario_options(**overrides: Any) -> SimpleNamespace:
|
||||
"""Return parsed scale demo options with a fixed three-second generation, with ``overrides`` by attribute name."""
|
||||
defaults: dict[str, Any] = {
|
||||
"chat_front_host": "http://front.test",
|
||||
"chat_horde_urls": "http://horde.test",
|
||||
"chat_model": "stress/chat",
|
||||
"chat_max_tokens": 64,
|
||||
"chat_no_include_usage": False,
|
||||
"chat_impatient_timeout": 2.0,
|
||||
"chat_impatient_stream_chance": 0.5,
|
||||
"chat_nonstream_timeout": 300.0,
|
||||
"sim_gen_time_min": 3.0,
|
||||
"sim_gen_time_max": 3.0,
|
||||
"chat_unreported_usage_chance": 0.0,
|
||||
"chat_worker_progress": True,
|
||||
"chat_worker_progress_interval": 1.0,
|
||||
}
|
||||
return SimpleNamespace(**{**defaults, **overrides})
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
"""A stand-in for Locust's response context manager, recording how the user marked it and the name it set."""
|
||||
|
||||
def __init__(self, status_code: int, body: Any) -> None:
|
||||
self.status_code = status_code
|
||||
self.ok = 0 < status_code < 400
|
||||
self.text = json.dumps(body)
|
||||
self.content = self.text.encode()
|
||||
self.error = "connection reset" if status_code == 0 else None
|
||||
self.elapsed = timedelta(0)
|
||||
self.request_meta: dict[str, Any] = {}
|
||||
self.marked: str | None = None
|
||||
self._body = body
|
||||
|
||||
def json(self) -> Any:
|
||||
return self._body
|
||||
|
||||
def success(self) -> None:
|
||||
self.marked = "success"
|
||||
|
||||
def failure(self, message: str) -> None:
|
||||
self.marked = message
|
||||
|
||||
def __enter__(self) -> FakeResponse:
|
||||
return self
|
||||
|
||||
def __exit__(self, *_exception: object) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
class FakeClient:
|
||||
"""A stand-in for Locust's HTTP session, answering posts from a queue and keeping each path and body sent."""
|
||||
|
||||
def __init__(self, responses: list[FakeResponse]) -> None:
|
||||
self.base_url = ""
|
||||
self.responses = responses
|
||||
self.sent: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
def post(self, path: str, *, json: dict[str, Any], **_options: Any) -> FakeResponse:
|
||||
self.sent.append((path, json))
|
||||
return self.responses.pop(0)
|
||||
|
||||
def bodies_sent_to(self, path: str) -> list[dict[str, Any]]:
|
||||
return [body for sent_path, body in self.sent if sent_path == path]
|
||||
|
||||
|
||||
def popped_chat_job(*, stream_progress: bool | None) -> FakeResponse:
|
||||
"""Return a pop answer assigning one chat job, with ``stream_progress`` left out when None."""
|
||||
payload: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}], "max_length": 64}
|
||||
if stream_progress is not None:
|
||||
payload["stream_progress"] = stream_progress
|
||||
return FakeResponse(200, {"id": GENERATION_ID, "chat": True, "model": "stress/chat", "payload": payload})
|
||||
|
||||
|
||||
def accepted(seq: int) -> FakeResponse:
|
||||
return FakeResponse(200, {"accepted_seq": seq})
|
||||
|
||||
|
||||
def submitted() -> FakeResponse:
|
||||
return FakeResponse(200, {"reward": 1})
|
||||
|
||||
|
||||
class TestDeltaSplit:
|
||||
def test_the_deltas_concatenate_to_the_reply_in_near_equal_pieces(self) -> None:
|
||||
deltas = split_reply_into_deltas(WORKER_GENERATION_TEXT, 5)
|
||||
|
||||
assert "".join(deltas) == WORKER_GENERATION_TEXT
|
||||
assert len(deltas) == 5
|
||||
assert max(map(len, deltas)) - min(map(len, deltas)) <= 1
|
||||
|
||||
def test_a_reply_shorter_than_the_count_gives_one_delta_per_character(self) -> None:
|
||||
assert split_reply_into_deltas("abc", 10) == ["a", "b", "c"]
|
||||
|
||||
def test_a_count_below_one_gives_the_whole_reply(self) -> None:
|
||||
assert split_reply_into_deltas("abc", 0) == ["abc"]
|
||||
|
||||
def test_an_empty_reply_is_refused(self) -> None:
|
||||
with pytest.raises(ValueError):
|
||||
split_reply_into_deltas("", 3)
|
||||
|
||||
def test_a_generation_posts_one_delta_per_whole_interval_and_at_least_one(self) -> None:
|
||||
assert count_progress_deltas(generation_seconds=3.7, interval_seconds=1.0) == 3
|
||||
assert count_progress_deltas(generation_seconds=0.4, interval_seconds=1.0) == 1
|
||||
assert count_progress_deltas(generation_seconds=3.0, interval_seconds=0.5) == 6
|
||||
|
||||
|
||||
class TestProgressMode:
|
||||
@pytest.fixture(autouse=True)
|
||||
def no_sleep(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(chat_users.time, "sleep", lambda _seconds: None)
|
||||
|
||||
def build_worker(self, responses: list[FakeResponse], **overrides: Any) -> ChatWorkerSimulator:
|
||||
chat_settings.configure(scenario_options(**overrides), fallback_horde_url="http://horde.test")
|
||||
hosted_worker_class = type("HostedChatWorkerSimulator", (ChatWorkerSimulator,), {"host": "http://horde.test"})
|
||||
worker = hosted_worker_class(Environment(user_classes=[hosted_worker_class]))
|
||||
worker.client = FakeClient(responses)
|
||||
worker.on_start()
|
||||
return worker
|
||||
|
||||
def test_a_streamed_job_posts_the_reply_once_a_second_and_submits_it_after(self) -> None:
|
||||
progress_answers = [accepted(1), accepted(2), accepted(3)]
|
||||
worker = self.build_worker([popped_chat_job(stream_progress=True), *progress_answers, submitted()])
|
||||
|
||||
worker.pop_and_submit_chat()
|
||||
|
||||
progress_bodies = worker.client.bodies_sent_to(TEXT_PROGRESS_PATH)
|
||||
assert [body["seq"] for body in progress_bodies] == [1, 2, 3]
|
||||
assert all(body["id"] == GENERATION_ID for body in progress_bodies)
|
||||
assert "".join(body["delta"]["content"] for body in progress_bodies) == WORKER_GENERATION_TEXT
|
||||
(submit_body,) = worker.client.bodies_sent_to(TEXT_SUBMIT_PATH)
|
||||
assert submit_body["message"]["content"].startswith(WORKER_GENERATION_TEXT)
|
||||
assert [answer.request_meta["name"] for answer in progress_answers] == [progress_request_name(200)] * 3
|
||||
assert all(answer.marked == "success" for answer in progress_answers)
|
||||
assert worker.client.responses == []
|
||||
|
||||
def test_the_interval_sets_the_delta_count(self) -> None:
|
||||
progress_answers = [accepted(seq) for seq in range(1, 7)]
|
||||
worker = self.build_worker(
|
||||
[popped_chat_job(stream_progress=True), *progress_answers, submitted()], chat_worker_progress_interval=0.5
|
||||
)
|
||||
|
||||
worker.pop_and_submit_chat()
|
||||
|
||||
assert len(worker.client.bodies_sent_to(TEXT_PROGRESS_PATH)) == 6
|
||||
|
||||
@pytest.mark.parametrize("stop_status", [404, 413])
|
||||
def test_a_job_without_a_lease_or_at_its_cap_stops_posting_and_still_submits(self, stop_status: int) -> None:
|
||||
refusal = FakeResponse(stop_status, {"message": "stop", "rc": "ChatProgressNotStreaming"})
|
||||
worker = self.build_worker([popped_chat_job(stream_progress=True), refusal, submitted()])
|
||||
|
||||
worker.pop_and_submit_chat()
|
||||
|
||||
assert len(worker.client.bodies_sent_to(TEXT_PROGRESS_PATH)) == 1
|
||||
assert len(worker.client.bodies_sent_to(TEXT_SUBMIT_PATH)) == 1
|
||||
assert refusal.request_meta["name"] == f"{TEXT_PROGRESS_PATH} [{stop_status}]"
|
||||
assert refusal.marked == "success"
|
||||
|
||||
def test_a_cancelled_acknowledgement_stops_posting_and_skips_the_submit(self) -> None:
|
||||
cancelled_answer = FakeResponse(200, {"accepted_seq": 2, "cancelled": True})
|
||||
worker = self.build_worker([popped_chat_job(stream_progress=True), accepted(1), cancelled_answer])
|
||||
|
||||
worker.pop_and_submit_chat()
|
||||
|
||||
assert len(worker.client.bodies_sent_to(TEXT_PROGRESS_PATH)) == 2
|
||||
assert worker.client.bodies_sent_to(TEXT_SUBMIT_PATH) == []
|
||||
assert cancelled_answer.marked == "success"
|
||||
assert worker.client.responses == []
|
||||
|
||||
def test_an_out_of_sequence_refusal_stops_posting_and_fails_the_post(self) -> None:
|
||||
refusal = FakeResponse(409, {"message": "out of sequence", "rc": "ChatProgressOutOfSequence", "expected_seq": 3})
|
||||
worker = self.build_worker([popped_chat_job(stream_progress=True), accepted(1), refusal, submitted()])
|
||||
|
||||
worker.pop_and_submit_chat()
|
||||
|
||||
assert len(worker.client.bodies_sent_to(TEXT_PROGRESS_PATH)) == 2
|
||||
assert len(worker.client.bodies_sent_to(TEXT_SUBMIT_PATH)) == 1
|
||||
assert refusal.marked is not None and refusal.marked != "success"
|
||||
|
||||
@pytest.mark.parametrize("retry_status", [0, 429, 503])
|
||||
def test_an_unaccepted_post_is_resent_unchanged_and_the_next_carries_the_rest(self, retry_status: int) -> None:
|
||||
unaccepted = FakeResponse(retry_status, {"message": "later"})
|
||||
worker = self.build_worker([popped_chat_job(stream_progress=True), unaccepted, accepted(1), accepted(2), submitted()])
|
||||
|
||||
worker.pop_and_submit_chat()
|
||||
|
||||
first, resent, last = worker.client.bodies_sent_to(TEXT_PROGRESS_PATH)
|
||||
assert resent == first
|
||||
assert last["seq"] == 2
|
||||
assert first["delta"]["content"] + last["delta"]["content"] == WORKER_GENERATION_TEXT
|
||||
assert unaccepted.request_meta["name"] == progress_request_name(retry_status)
|
||||
|
||||
def test_the_mode_off_posts_no_progress(self) -> None:
|
||||
worker = self.build_worker([popped_chat_job(stream_progress=True), submitted()], chat_worker_progress=False)
|
||||
|
||||
worker.pop_and_submit_chat()
|
||||
|
||||
assert [path for path, _body in worker.client.sent] == [TEXT_POP_PATH, TEXT_SUBMIT_PATH]
|
||||
|
||||
@pytest.mark.parametrize("stream_progress", [None, False])
|
||||
def test_a_job_whose_client_does_not_stream_posts_no_progress(self, stream_progress: bool | None) -> None:
|
||||
worker = self.build_worker([popped_chat_job(stream_progress=stream_progress), submitted()])
|
||||
|
||||
worker.pop_and_submit_chat()
|
||||
|
||||
assert [path for path, _body in worker.client.sent] == [TEXT_POP_PATH, TEXT_SUBMIT_PATH]
|
||||
|
||||
def test_a_non_positive_interval_is_refused(self) -> None:
|
||||
with pytest.raises(ValueError):
|
||||
chat_settings.configure(scenario_options(chat_worker_progress_interval=0), fallback_horde_url="http://horde.test")
|
||||
Reference in new issue
Block a user