From 23a2e861a0334857befebea974dedb833288fb82 Mon Sep 17 00:00:00 2001 From: narmstrong-berry Date: Tue, 3 Nov 2026 03:44:01 +0000 Subject: [PATCH] feat(loadtest): portkey gateway routing for the perf lane genconfig's portkey key set writes portkey-headers.json, one x-portkey-config routing header per model group that load-balances over the group's deployments on the mock (custom_host, override_params.model, SigV4 keys for Bedrock), plus keys-portkey.txt. loadgen -model-headers sends each request's header. Basic features only: the open-source Portkey gateway has no keys, limits, budgets, logs or metrics. --- litellm-rust/loadtest/README.md | 28 +++++- litellm-rust/loadtest/cmd/genconfig/main.go | 12 ++- .../loadtest/cmd/genconfig/portkey.go | 91 +++++++++++++++++++ litellm-rust/loadtest/cmd/loadgen/main.go | 48 +++++++++- 4 files changed, 175 insertions(+), 4 deletions(-) create mode 100644 litellm-rust/loadtest/cmd/genconfig/portkey.go diff --git a/litellm-rust/loadtest/README.md b/litellm-rust/loadtest/README.md index 16c5fd3b73c..d9b0255a778 100644 --- a/litellm-rust/loadtest/README.md +++ b/litellm-rust/loadtest/README.md @@ -7,8 +7,8 @@ machines. | Tool | What it does | | --- | --- | | `mockllm` | Fake OpenAI chat completions, Anthropic messages and Bedrock invoke upstream. Streams or not, and takes its timing from a directive in the prompt (`[[mock ttft_ms=800 dur_ms=300000 out=2000]]`). Counts hits per deployment (`/d//` path prefix) at `GET /stats`; `-faults "chat-0-3=429:1.0"` makes chosen deployments fail. | -| `loadgen` | Open-loop arrivals from a workload file, Poisson or evenly spaced (`"arrivals": "constant"`): weighted request classes with their own endpoint, stream share, context size (10k to 1M tokens), time to first token and total time. Writes one JSON line per request, with `lag_ms`, how late it left against its planned arrival (when that grows, the generator is the limit). Shard across machines with `-shard i -shards n`. | -| `genconfig` | Gateway `config.yaml` with many model groups and deployments pointing at the mocks, plus `keys.txt` and `keys.sql` (rows for `LiteLLM_VerificationToken`). `-key-seed` makes every machine generate the same keys. The `bifrost` key set is written as a Bifrost config; see [Bifrost](#bifrost). | +| `loadgen` | Open-loop arrivals from a workload file, Poisson or evenly spaced (`"arrivals": "constant"`): weighted request classes with their own endpoint, stream share, context size (10k to 1M tokens), time to first token and total time. Writes one JSON line per request, with `lag_ms`, how late it left against its planned arrival (when that grows, the generator is the limit). Shard across machines with `-shard i -shards n`. `-model-headers` adds per-model-group headers (Portkey's routing config). | +| `genconfig` | Gateway `config.yaml` with many model groups and deployments pointing at the mocks, plus `keys.txt` and `keys.sql` (rows for `LiteLLM_VerificationToken`). `-key-seed` makes every machine generate the same keys. The `bifrost` key set is written as a Bifrost config; see [Bifrost](#bifrost). The `portkey` key set is written as Portkey routing headers; see [Portkey](#portkey). | | `summarize` | Merges result files and prints proxy overhead (observed minus planned time) by stage, class and context size. | | `procmon` | Samples the gateway process's CPU, RSS, fds and TCP connections from `/proc`. | | `verify` | Checks what each gateway recorded for a run against what loadgen sent: spend logs, key/team/user/member spend, daily tables, Prometheus counters, S3 request logs and the headers the provider received. See [Full-feature runs](#full-feature-runs). | @@ -181,6 +181,30 @@ tables, `bifrost_*` metrics, its S3 log objects) and reports what Bifrost has no n/a. With three sides it compares each pair's fitted deployment prices; metric and S3 field lists are only compared between the LiteLLM sides, since Bifrost's schema differs throughout. +## Portkey + +A key set named `portkey` (`-key-sets ,portkey`, basic features only) is written as routing for the +[open-source Portkey gateway](https://github.com/Portkey-AI/gateway) instead of key rows. That +gateway keeps no model list, keys or database: every request carries its routing in an +`x-portkey-config` header. So genconfig writes `portkey-headers.json`, one header per model group, +and `keys-portkey.txt`, keys of the same shape as the other sets that Portkey never checks (it has +no client auth). Each group's config load-balances evenly over the group's deployments (Portkey's +counterpart of simple-shuffle) with `retry.attempts` 2, and each target points at the mock's +`/d/` with `custom_host` and `override_params.model`. Bedrock targets carry static AWS keys, so +requests are SigV4 signed; `custom_host` replaces the Bedrock URL, so they reach the mock over plain +HTTP, and Portkey sends Anthropic models' messages to Converse. Portkey has no limits, budgets, logs +or metrics, so `-features full` with the `portkey` set is refused. + +Portkey refuses a `custom_host` on a private or `.local` name unless `TRUSTED_CUSTOM_HOSTS` lists +it, so start it with that set to the mock's host name. loadgen sends each request's header with +`-model-headers`, and fails at start when a model the workload can pick has none: + +```sh +TRUSTED_CUSTOM_HOSTS=mockllm node build/start-server.js --headless --port=8787 +loadgen -workload workloads/smoke.json -target http://portkey:8787 -keys out/keys-portkey.txt \ + -model-headers out/portkey-headers.json ... +``` + ## Bedrock deployments `genconfig` serves half of each messages group through Bedrock invoke by default diff --git a/litellm-rust/loadtest/cmd/genconfig/main.go b/litellm-rust/loadtest/cmd/genconfig/main.go index ad92b7ca252..42461f3b0e1 100644 --- a/litellm-rust/loadtest/cmd/genconfig/main.go +++ b/litellm-rust/loadtest/cmd/genconfig/main.go @@ -15,6 +15,9 @@ // custom_callbacks.py. headers.txt (headers-.txt) holds the request headers those hooks act // on. S3 settings come from LOADTEST_S3_* environment variables. // +// -key-sets ,portkey adds Portkey gateway routing (portkey.go): one x-portkey-config header per +// model group, in portkey-headers.json, for loadgen -model-headers. Basic features only. +// // -key-sets python,rust writes one key set, with its own users and teams, per name, so each // gateway's spend can be told apart in a shared database. package main @@ -45,7 +48,7 @@ func main() { masterKey := flag.String("master-key", "sk-loadtest-master", "gateway master key") outDir := flag.String("out", ".", "output directory") features := flag.String("features", "basic", "basic (virtual keys only) or full (limits, budgets, callbacks and hooks)") - keySets := flag.String("key-sets", "", "comma-separated key set names, each written to keys-.txt with its own users and teams (default: one set in keys.txt); the set named bifrost becomes bifrost-config.json") + keySets := flag.String("key-sets", "", "comma-separated key set names, each written to keys-.txt with its own users and teams (default: one set in keys.txt); the set named bifrost becomes bifrost-config.json, the set named portkey portkey-headers.json") bifrostConcurrency := flag.Int("bifrost-concurrency", 500, "Bifrost workers per deployment (each deployment is a provider with its own worker pool)") bifrostBuffer := flag.Int("bifrost-buffer", 5000, "Bifrost queue length per deployment, and its client pool size") bifrostBedrock := flag.String("bifrost-bedrock-tls", "127.0.0.1:9443", "host:port of mockllm -tls-addr, where Bifrost's Bedrock deployments go") @@ -138,6 +141,13 @@ general_settings: sw := bufio.NewWriter(sf) sw.WriteString("-- Load-test virtual keys. Hash = sha256 hex of the key, as the proxy stores it.\n") for _, set := range sets { + if set == portkeySet { + if full { + log.Fatal("the portkey key set needs -features basic: the open-source Portkey gateway has no keys, limits, budgets or logs to turn on") + } + writePortkey(bifrostDeps, portkeyKeys(*keys, *keySeed), *outDir) + continue + } if set == bifrostSet { writeBifrost(bifrostDeps, bifrostOptions{full: full, keys: *keys, keySeed: *keySeed, concurrency: *bifrostConcurrency, bufferSize: *bifrostBuffer, pricingFile: *bifrostPricing, bedrockTLS: *bifrostBedrock, outDir: *outDir}) diff --git a/litellm-rust/loadtest/cmd/genconfig/portkey.go b/litellm-rust/loadtest/cmd/genconfig/portkey.go new file mode 100644 index 00000000000..3a2c5bd1eed --- /dev/null +++ b/litellm-rust/loadtest/cmd/genconfig/portkey.go @@ -0,0 +1,91 @@ +package main + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "log" + "path/filepath" + "strings" +) + +// portkeySet is the key set genconfig writes as Portkey gateway routing instead of LiteLLM key rows. +// The open-source Portkey gateway (github.com/Portkey-AI/gateway) keeps no model list, keys or +// database: every request carries its routing in an x-portkey-config header. So the set becomes +// portkey-headers.json, the header loadgen sends for each model group (-model-headers), and +// keys-portkey.txt, which the gateway never checks (it has no client auth) and exists only so the +// request bytes match the other sides. +const portkeySet = "portkey" + +// portkeyHeader is the routing header name. Its value is a Portkey config (JSON). +const portkeyHeader = "x-portkey-config" + +// writePortkey writes portkey-headers.json and keys-portkey.txt beside config.yaml. Each model +// group's config load-balances evenly over the group's deployments, Portkey's counterpart of +// simple-shuffle, and retries twice like router_settings.num_retries. Every target points at its +// deployment's /d/ path on the mock with custom_host, so the mock counts hits per deployment as +// it does for the other sides. Bedrock targets carry static AWS keys, so every request is SigV4 +// signed; custom_host replaces the Bedrock runtime URL, so they reach the mock over plain HTTP +// (Portkey sends Anthropic-model messages to Bedrock Converse, which the mock serves). +// +// The gateway refuses a custom_host on a private or .local name unless it is listed in +// TRUSTED_CUSTOM_HOSTS, so the pod sets that to the mock's host name. +func writePortkey(deps []bifrostDeployment, keys []string, outDir string) { + byGroup := map[string][]any{} + var groups []string + for _, d := range deps { + if _, ok := byGroup[d.group]; !ok { + groups = append(groups, d.group) + } + var t map[string]any + switch d.kind { + case "bedrock": + t = map[string]any{ + "provider": "bedrock", + "aws_access_key_id": "mock-loadtest-access-key", + "aws_secret_access_key": "mock-loadtest-secret", + "aws_region": "us-east-1", + "custom_host": d.base, + } + case "openai", "anthropic": + // Both providers append their endpoint (/chat/completions, /messages) to a base that + // ends in /v1, as their public base URLs do. + t = map[string]any{"provider": d.kind, "api_key": "sk-mock", "custom_host": d.base + "/v1"} + default: + log.Fatalf("portkey: no mapping for deployment %s of kind %q", d.id, d.kind) + } + t["override_params"] = map[string]any{"model": d.model} + t["weight"] = 1 + byGroup[d.group] = append(byGroup[d.group], t) + } + headers := map[string]map[string]string{} + for _, g := range groups { + cfg := map[string]any{ + "strategy": map[string]any{"mode": "loadbalance"}, + "retry": map[string]any{"attempts": 2}, + "targets": byGroup[g], + } + b, err := json.Marshal(cfg) + if err != nil { + log.Fatal(err) + } + headers[g] = map[string]string{portkeyHeader: string(b)} + } + writeJSONFile(filepath.Join(outDir, "portkey-headers.json"), headers) + write(filepath.Join(outDir, "keys-"+portkeySet+".txt"), strings.Join(keys, "\n")+"\n") +} + +// portkeyKeys derives the set's keys like writeKeys does (same sk-lt- shape and length), so each +// request carries the same bytes as on the other sides. +func portkeyKeys(n int, keySeed string) []string { + if keySeed == "" { + keySeed = "loadtest" + } + keys := make([]string, 0, n) + for i := range n { + derived := sha256.Sum256(fmt.Appendf(nil, "%s/%s/%d", keySeed, portkeySet, i)) + keys = append(keys, "sk-lt-"+hex.EncodeToString(derived[:16])) + } + return keys +} diff --git a/litellm-rust/loadtest/cmd/loadgen/main.go b/litellm-rust/loadtest/cmd/loadgen/main.go index 0b308484b52..dc36a81a920 100644 --- a/litellm-rust/loadtest/cmd/loadgen/main.go +++ b/litellm-rust/loadtest/cmd/loadgen/main.go @@ -25,6 +25,7 @@ import ( "net/http" "os" "os/signal" + "sort" "strings" "sync" "sync/atomic" @@ -78,12 +79,32 @@ func main() { messagesPath := flag.String("messages-path", "/v1/messages", "path of the messages endpoint under -target (Bifrost serves it at /anthropic/v1/messages)") var headers headerFlags flag.Var(&headers, "header", "extra request header as 'Name: value' (repeatable)") + modelHeadersFile := flag.String("model-headers", "", "JSON file of {model: {header: value}} sent with each request for that model, on top of -header (Portkey's per-group x-portkey-config, from genconfig's portkey key set)") flag.Parse() + modelHeaders := map[string]headerFlags{} + if *modelHeadersFile != "" { + mh, err := readModelHeaders(*modelHeadersFile) + modelHeaders = mh + if err != nil { + log.Fatalf("model headers: %v", err) + } + } w, err := loadWorkload(*workloadPath) if err != nil { log.Fatal(err) } + // With -model-headers, every model the workload can pick needs its headers (a Portkey request + // without its config is refused, which would read as a gateway failure). + if *modelHeadersFile != "" { + for _, m := range w.Models { + for i := range m.Count { + if model := fmt.Sprintf("%s%d", m.Prefix, i); len(modelHeaders[model]) == 0 { + log.Fatalf("model headers: %s has no entry for model %s", *modelHeadersFile, model) + } + } + } + } keys := []string{*key} if *key == "" { file := w.Keys.File @@ -241,7 +262,7 @@ schedule: go func() { defer wg.Done() defer c.inflight.Add(-1) - do(reqCtx, transport, *target, paths{*chatPath, *messagesPath}, keys[keyIdx], *messagesAuth, headers, body, &res) + do(reqCtx, transport, *target, paths{*chatPath, *messagesPath}, keys[keyIdx], *messagesAuth, append(headers[:len(headers):len(headers)], modelHeaders[res.Model]...), body, &res) if res.OK { c.ok.Add(1) } else { @@ -286,6 +307,31 @@ func (h *headerFlags) Set(value string) error { return nil } +// readModelHeaders reads a {model: {header: value}} file into each model's extra headers, in a +// stable order. +func readModelHeaders(path string) (map[string]headerFlags, error) { + b, err := os.ReadFile(path) + if err != nil { + return nil, err + } + var raw map[string]map[string]string + if err := json.Unmarshal(b, &raw); err != nil { + return nil, err + } + out := make(map[string]headerFlags, len(raw)) + for model, hs := range raw { + names := make([]string, 0, len(hs)) + for name := range hs { + names = append(names, name) + } + sort.Strings(names) + for _, name := range names { + out[model] = append(out[model], [2]string{name, hs[name]}) + } + } + return out, nil +} + // paths are the chat and messages endpoints under the target. type paths struct{ chat, messages string }