diff --git a/.dockerignore b/.dockerignore index 01b4c993..39c692c3 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,6 +1,9 @@ .git .github target +target-pgo-gen +target-pgo +bench/pgo-training/trainer/target docs tests/fixtures/tmp diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml index 1cf71caa..c4b18c15 100644 --- a/.github/workflows/docker-image.yml +++ b/.github/workflows/docker-image.yml @@ -33,6 +33,7 @@ on: - rust-toolchain.toml - "crates/**" - "schemas/**" + - "bench/pgo-training/**" workflow_dispatch: inputs: tag: @@ -55,10 +56,45 @@ concurrency: jobs: build: runs-on: ubuntu-latest - timeout-minutes: 45 + # PGO builds (#967) compile the workspace twice (instrumented + optimized) + # with a training run in between — roughly 2.5x the plain fat-LTO build + # that used to fit in 45 minutes. + timeout-minutes: 120 steps: - uses: actions/checkout@v6 + # PGO build mode (#967). Push builds (dev/poc/tags/dispatch) are always + # PGO=on: shipped artifacts are PGO'd, fail-closed, no fallback. PR + # builds default to PGO=off so review latency stays at one compile — + # EXCEPT when the PR touches the build pipeline or the training assets, + # in which case the full three-phase build must prove itself pre-merge. + - name: Decide PGO build mode + id: pgomode + env: + BASE_SHA: ${{ github.event.pull_request.base.sha || '' }} + run: | + set -eu + mode=on + if [ "${{ github.event_name }}" = "pull_request" ]; then + mode=off + git fetch --no-tags --depth=1 origin "$BASE_SHA" + # Capture the diff as a CHECKED command before matching: inside + # an `if` condition a git failure is exempt from `set -e` and + # would silently leave mode=off (fail-open). As an assignment, + # a diff failure kills the step instead. Captured text also + # avoids grep -q closing the pipe under git (SIGPIPE). + changed="$(git diff --name-only "$BASE_SHA" HEAD)" + # Root build inputs (workspace manifest with [profile.release], + # lockfile, toolchain pin) are PGO inputs too: a bump can break + # the instrumented/profile-use build or the training run while + # the plain PGO=off compile stays green. + if printf '%s\n' "$changed" \ + | grep -qE '^(Dockerfile|\.dockerignore|Cargo\.(toml|lock)|rust-toolchain\.toml|bench/pgo-training/|\.github/workflows/docker-image\.ya?ml)'; then + mode=on + fi + fi + echo "value=$mode" >> "$GITHUB_OUTPUT" + - name: Set up Docker Buildx uses: docker/setup-buildx-action@v4 @@ -145,6 +181,7 @@ jobs: build-args: | BUILD_SHA=${{ steps.shortsha.outputs.value }} BUILD_VERSION=${{ steps.relver.outputs.value }} + PGO=${{ steps.pgomode.outputs.value }} cache-from: type=gha cache-to: type=gha,mode=max provenance: false # keep manifest format compatible with older clients @@ -163,6 +200,35 @@ jobs: docker run --rm --entrypoint /usr/local/bin/aisix "$IMG" --version | tee /tmp/aisix-version grep -qx "aisix ${VER}" /tmp/aisix-version + # Fail-closed release gate (#967): every image built with PGO=on must + # carry the proof marker the Dockerfile writes only after the + # profile-optimized build succeeds. An absent marker, a short shape + # list, or an undersized merged profile fails the workflow here — + # before signing, and before any human treats the artifact as good. + # The `if` is deliberately NOT just `pgomode == 'on'`: every push build + # asserts unconditionally, so a future bug in the pgomode step itself + # (empty output, renamed id) cannot skip both the PGO build AND its + # guard in the same breath. + - name: Assert PGO proof marker (#967) + if: github.event_name != 'pull_request' || steps.pgomode.outputs.value == 'on' + env: + TAGS: ${{ steps.meta.outputs.tags }} + DIGEST: ${{ steps.build.outputs.digest }} + run: | + set -eux + if [ "${{ github.event_name }}" = "pull_request" ]; then + IMG="$(printf '%s\n' "$TAGS" | head -n1)" + else + IMG="${REGISTRY}/${IMAGE_NAME}@${DIGEST}" + docker pull "$IMG" + fi + docker run --rm --entrypoint cat "$IMG" /usr/local/share/aisix/pgo-verified.json \ + | tee /tmp/pgo-verified.json + shapes="$(jq -r '.shapes | length' /tmp/pgo-verified.json)" + test "$shapes" -ge 12 + test "$(jq -r '.profraw_count' /tmp/pgo-verified.json)" -ge "$shapes" + test "$(jq -r '.profdata_bytes' /tmp/pgo-verified.json)" -ge 524288 + # The image contract customers rely on for hostNetwork/:80 # deployments: the default non-root user (uid 10001) must be able # to bind privileged ports via the CAP_NET_BIND_SERVICE file @@ -208,11 +274,21 @@ jobs: # RUSTFLAGS change, or base-image swap would break flame graphs # silently, so pin it here — and print the size so every PR # records the real Linux artifact cost. + # Also enforced on tag builds since #967: the PGO'd binary that ships + # must provably keep the profiling contract, not just the PR variant. - name: Verify shipped binary keeps its symbol table (#847) - if: github.event_name == 'pull_request' + if: github.event_name == 'pull_request' || startsWith(github.ref, 'refs/tags/v') + env: + TAGS: ${{ steps.meta.outputs.tags }} + DIGEST: ${{ steps.build.outputs.digest }} run: | set -eux - IMAGE="$(printf '%s\n' "${{ steps.meta.outputs.tags }}" | head -n1)" + if [ "${{ github.event_name }}" = "pull_request" ]; then + IMAGE="$(printf '%s\n' "$TAGS" | head -n1)" + else + IMAGE="${REGISTRY}/${IMAGE_NAME}@${DIGEST}" + docker pull "$IMAGE" + fi cid="$(docker create "$IMAGE")" docker cp "$cid:/usr/local/bin/aisix" /tmp/aisix-shipped docker rm "$cid" diff --git a/.gitignore b/.gitignore index e458f699..59caf342 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,9 @@ # Rust /target +# PGO phase target dirs (native runs of the three-phase release recipe) +/target-pgo-gen +/target-pgo +/bench/pgo-training/trainer/target Cargo.lock.bak **/*.rs.bk *.profraw diff --git a/Cargo.toml b/Cargo.toml index c879ffac..777723e3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,9 @@ members = [ "crates/aisix-guardrails", "crates/aisix-server", ] +# Build-time tool for the PGO release pipeline (#967): own manifest + lockfile, +# never part of the product build graph. +exclude = ["bench/pgo-training/trainer"] [workspace.package] version = "0.3.0" diff --git a/Dockerfile b/Dockerfile index 80e69ed3..fd1af97d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -12,6 +12,7 @@ # # Build: # docker build -t aisix:dev . +# docker build --build-arg PGO=off -t aisix:dev . # quick local build, skips PGO # # Run, standalone (mount your own config): # docker run --rm -v $(pwd)/config.example.yaml:/etc/aisix/config.yaml \ @@ -68,18 +69,66 @@ COPY crates ./crates # Docker context must carry this directory or the release build fails. COPY schemas ./schemas +# PGO training assets (#967): trainer tool + train.sh. Copied separately from +# crates/ so editing training assets doesn't invalidate the dependency layers +# above. +COPY bench/pgo-training ./bench/pgo-training + +# Profile-guided optimization gate. Default ON: release artifacts are always +# PGO-built, and a forgotten build-arg ships a PGO'd image — never a silently +# un-optimized one. CI passes PGO=off only for pull-request smoke builds. +ARG PGO=on + # `--locked` forces the build to use the exact versions in Cargo.lock — # fails fast if the lockfile is stale rather than silently resolving # fresh deps in CI. # +# PGO=on runs the three-phase build (#967): +# A. instrumented build (-Cprofile-generate) in its own target dir; +# B. train.sh drives the committed 12-shape matrix through the +# instrumented gateway against the trainer's local mock, then merges +# the .profraw files with the pinned toolchain's own llvm-profdata +# (llvm-tools-preview — exact LLVM match with rustc, no extra deps); +# C. optimized build (-Cprofile-use) in a third target dir, so profile +# builds never share cargo fingerprints with plain builds. +# FAIL-CLOSED: any phase failing fails this RUN and nothing is shipped. +# The proof marker (pgo-verified.json) is written only after phase C +# succeeds; the push workflows assert it before trusting the image. +# The merged profile is content-addressed (merged-.profdata) because +# cargo fingerprints the -Cprofile-use PATH, not the file content — a +# retrained profile at a fixed path would silently reuse stale artifacts +# from the persistent target cache mount. +# # If this ever builds for linux/arm64: jemalloc bakes the build host's # page size into the binary, and QEMU reports 4K — set # JEMALLOC_SYS_WITH_LG_PAGE=16 here or the image aborts at startup on -# 64K-page kernels (see crates/aisix-server/src/main.rs). +# 64K-page kernels (see crates/aisix-server/src/main.rs). PGO training +# additionally requires a native arm64 builder: an instrumented binary +# cannot self-train under QEMU emulation. RUN --mount=type=cache,target=/usr/local/cargo/registry \ --mount=type=cache,target=/src/target \ - cargo build --locked --release --bin aisix \ - && cp target/release/aisix /usr/local/bin/aisix + --mount=type=cache,target=/src/target-pgo-gen \ + --mount=type=cache,target=/src/target-pgo \ + set -eu; \ + mkdir -p /usr/local/share/aisix; \ + if [ "$PGO" = "on" ]; then \ + RUSTFLAGS="-Cprofile-generate=/tmp/pgo-data" CARGO_TARGET_DIR=/src/target-pgo-gen \ + cargo build --locked --release --bin aisix; \ + cargo build --locked --release \ + --manifest-path bench/pgo-training/trainer/Cargo.toml; \ + bash bench/pgo-training/train.sh /src/target-pgo-gen/release/aisix /tmp/pgo-data; \ + PROFDATA="$(ls /tmp/pgo-data/merged-*.profdata)"; \ + RUSTFLAGS="-Cprofile-use=$PROFDATA" CARGO_TARGET_DIR=/src/target-pgo \ + cargo build --locked --release --bin aisix; \ + cp /src/target-pgo/release/aisix /usr/local/bin/aisix; \ + cp /tmp/pgo-data/train-manifest.json /usr/local/share/aisix/pgo-verified.json; \ + elif [ "$PGO" = "off" ]; then \ + cargo build --locked --release --bin aisix; \ + cp target/release/aisix /usr/local/bin/aisix; \ + else \ + echo "unsupported PGO value: '$PGO' (use on|off)" >&2; \ + exit 2; \ + fi # --- Stage 2: runtime -------------------------------------------------------- FROM debian:bookworm-slim AS runtime @@ -100,11 +149,19 @@ RUN apt-get update \ # missing from the container's bounding set — it is in the default # Docker/containerd cap set, but `capabilities: {drop: [ALL]}` pod # specs must add NET_BIND_SERVICE back. +# The PGO proof marker (#967) ships with the image: written by the builder +# only after a successful profile-optimized build, asserted by the push +# workflows before an image is trusted. Absent on PGO=off (PR smoke) builds. RUN --mount=type=bind,from=builder,source=/usr/local/bin/aisix,target=/mnt/aisix \ + --mount=type=bind,from=builder,source=/usr/local/share/aisix,target=/mnt/aisix-share \ apt-get update \ && apt-get install -y --no-install-recommends libcap2-bin \ && install -m 0755 /mnt/aisix /usr/local/bin/aisix \ && setcap 'cap_net_bind_service=+ep' /usr/local/bin/aisix \ + && mkdir -p /usr/local/share/aisix \ + && if [ -f /mnt/aisix-share/pgo-verified.json ]; then \ + install -m 0644 /mnt/aisix-share/pgo-verified.json /usr/local/share/aisix/pgo-verified.json; \ + fi \ && apt-get purge -y --auto-remove libcap2-bin \ && rm -rf /var/lib/apt/lists/* diff --git a/RELEASING.md b/RELEASING.md index 2c0d28d2..24e95686 100644 --- a/RELEASING.md +++ b/RELEASING.md @@ -27,6 +27,27 @@ Pushing the tag triggers two workflows: curated-notes scaffold to fill in, then GitHub's auto-generated **What's Changed** list as a starting skeleton. +### PGO is mandatory and fail-closed + +Published images are profile-guided-optimized (#967): the Docker build +compiles an instrumented gateway, drives the committed training matrix +(`bench/pgo-training/`) against it, and rebuilds with the merged profile. +Any phase failing — instrumented build, training, profile merge, optimized +build — fails the image build; there is no fallback to a plain build. After +the push, the workflow asserts the `pgo-verified.json` proof marker inside +the image (shape count, profile size) before signing. If a release build +fails in a PGO phase, fix the cause; never ship around it. To inspect a +shipped image's marker: + +```bash +docker run --rm --entrypoint cat ghcr.io/api7/aisix:X.Y.Z \ + /usr/local/share/aisix/pgo-verified.json +``` + +Local note: each retrained profile is content-addressed, so repeated local +PGO builds accumulate build artifacts in the persistent BuildKit cache +mounts; reclaim with `docker builder prune`. + ## 2. Polish the release notes Edit the draft before publishing. The Get-started/Download header and the diff --git a/bench/pgo-training/train.sh b/bench/pgo-training/train.sh new file mode 100755 index 00000000..5ba9c874 --- /dev/null +++ b/bench/pgo-training/train.sh @@ -0,0 +1,195 @@ +#!/usr/bin/env bash +# PGO training run — phase B of the three-phase release build (#967). +# +# Drives the v1 shape matrix (defined in trainer/src/main.rs, single source +# of truth) through an INSTRUMENTED gateway binary, one gateway process +# lifetime per shape, so every shape leaves its own .profraw and the merged +# profile is a union of hotness. Then merges the profraws with the +# llvm-profdata that ships in the pinned rustup toolchain (llvm-tools-preview +# — exact LLVM match with rustc, zero external toolchain dependencies) and +# writes a content-addressed merged-.profdata plus train-manifest.json. +# +# The content-addressed profdata name is load-bearing: -Cprofile-use= +# enters cargo's fingerprint via RUSTFLAGS, but cargo does NOT fingerprint the +# file's CONTENT — a retrained profile at an unchanged path would silently +# reuse stale phase-C artifacts from a persistent target dir. Hashing the name +# makes every retrain a fresh RUSTFLAGS value. +# +# FAIL-CLOSED (#967 hard gate 2): every failure path here exits non-zero and +# the caller must treat that as fatal. Never add a fallback that lets a +# release continue with a partial or empty profile. +# +# Usage: train.sh +# +# Knobs (env): +# PGO_TRAINER_BIN trainer binary (default: trainer/target/release/pgo-trainer) +# PGO_TRAIN_REQUESTS requests per shape (default 3000) +# PGO_TRAIN_CONCURRENCY driver connections (default 8) +# PGO_GW_PORT / PGO_MOCK_PORT / PGO_METRICS_PORT (defaults 13000/18001/19090) +set -euo pipefail + +BIN="${1:?usage: train.sh }" +PGO_DIR="${2:?usage: train.sh }" + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TRAINER="${PGO_TRAINER_BIN:-$HERE/trainer/target/release/pgo-trainer}" +REQS="${PGO_TRAIN_REQUESTS:-3000}" +CONC="${PGO_TRAIN_CONCURRENCY:-8}" +GW_PORT="${PGO_GW_PORT:-13000}" +MOCK_PORT="${PGO_MOCK_PORT:-18001}" +METRICS_PORT="${PGO_METRICS_PORT:-19090}" + +# A profraw for this binary is megabytes; an empty one (no counters flushed) +# is a few KB. The floor catches "gateway died before flushing" without ever +# false-failing a real profile. +PROFRAW_FLOOR_BYTES=262144 +PROFDATA_FLOOR_BYTES=524288 + +[ -x "$BIN" ] || { echo "FATAL: instrumented binary missing: $BIN" >&2; exit 1; } +[ -x "$TRAINER" ] || { echo "FATAL: trainer binary missing: $TRAINER (build trainer/ first)" >&2; exit 1; } + +mkdir -p "$PGO_DIR" +rm -f "$PGO_DIR"/*.profraw "$PGO_DIR"/merged-*.profdata "$PGO_DIR"/train-manifest.json + +# Shape list comes from the trainer itself — one source of truth; a drift +# between this loop and the trainer's table is impossible by construction. +mapfile -t SHAPES < <("$TRAINER" --list-shapes) +[ "${#SHAPES[@]}" -ge 12 ] || { echo "FATAL: trainer lists ${#SHAPES[@]} shapes, expected >= 12" >&2; exit 1; } + +# ---- gateway training config ------------------------------------------------- +# Modeled on the committed bench topology (bench/onthebench/run-baseline.sh): +# standalone resources_file mode, no etcd, one mock upstream, one api key from +# env. Default observability posture stays ON (the metrics recording path must +# train hot); only the prometheus listener address moves off :9090 so a +# training run can never collide with a gateway already running on the host. +cat > "$PGO_DIR/config.yaml" < "$PGO_DIR/resources.yaml" <&2 + PGO_TRAIN_KEY=pgo-token \ + LLVM_PROFILE_FILE="$PGO_DIR/aisix-$shape-%p.profraw" \ + "$BIN" --config "$PGO_DIR/config.yaml" > "$PGO_DIR/gateway-$shape.log" 2>&1 & + GW_PID=$! + + if ! "$TRAINER" --mock-port "$MOCK_PORT" --gateway "127.0.0.1:$GW_PORT" \ + --api-key pgo-token --shape "$shape" --requests "$REQS" --concurrency "$CONC"; then + echo "FATAL: training shape '$shape' failed; gateway log tail:" >&2 + tail -n 40 "$PGO_DIR/gateway-$shape.log" >&2 || true + kill -KILL "$GW_PID" 2>/dev/null || true + exit 1 + fi + + # Graceful shutdown is what flushes the profile counters to disk; a + # non-zero exit means the flush cannot be trusted. The watchdog bounds a + # hung shutdown at 120s instead of burning the whole CI job timeout — + # SIGKILL forces rc=137 and the FATAL path below. + kill -TERM "$GW_PID" + ( sleep 120; kill -KILL "$GW_PID" 2>/dev/null ) & WATCHDOG=$! + rc=0; wait "$GW_PID" || rc=$? + kill "$WATCHDOG" 2>/dev/null || true + if [ "$rc" -ne 0 ]; then + echo "FATAL: gateway exited rc=$rc after shape '$shape'; log tail:" >&2 + tail -n 40 "$PGO_DIR/gateway-$shape.log" >&2 || true + exit 1 + fi + + raw_count=$(find "$PGO_DIR" -maxdepth 1 -name "aisix-$shape-*.profraw" | wc -l) + [ "$raw_count" -ge 1 ] || { echo "FATAL: shape '$shape' left no .profraw" >&2; exit 1; } + while read -r raw; do + sz=$(stat -c%s "$raw") + [ "$sz" -ge "$PROFRAW_FLOOR_BYTES" ] || + { echo "FATAL: $raw is ${sz}B (< ${PROFRAW_FLOOR_BYTES}B floor) - counters did not flush" >&2; exit 1; } + done < <(find "$PGO_DIR" -maxdepth 1 -name "aisix-$shape-*.profraw") +done + +TOTAL_RAW=$(find "$PGO_DIR" -maxdepth 1 -name '*.profraw' | wc -l) +[ "$TOTAL_RAW" -ge "${#SHAPES[@]}" ] || + { echo "FATAL: $TOTAL_RAW profraw files for ${#SHAPES[@]} shapes" >&2; exit 1; } + +# ---- merge with the toolchain's own llvm-profdata ------------------------------ +HOST_TUPLE=$(rustc -vV | sed -n 's/^host: //p') +LLVM_PROFDATA="${PGO_LLVM_PROFDATA:-$(rustc --print sysroot)/lib/rustlib/$HOST_TUPLE/bin/llvm-profdata}" +[ -x "$LLVM_PROFDATA" ] || + { echo "FATAL: llvm-profdata not found at $LLVM_PROFDATA (llvm-tools-preview component missing?)" >&2; exit 1; } + +"$LLVM_PROFDATA" merge -o "$PGO_DIR/merged.profdata" "$PGO_DIR"/*.profraw + +PROFDATA_SHA=$(sha256sum "$PGO_DIR/merged.profdata" | cut -d' ' -f1) +PROFDATA_BYTES=$(stat -c%s "$PGO_DIR/merged.profdata") +[ "$PROFDATA_BYTES" -ge "$PROFDATA_FLOOR_BYTES" ] || + { echo "FATAL: merged profile is ${PROFDATA_BYTES}B (< ${PROFDATA_FLOOR_BYTES}B floor)" >&2; exit 1; } +mv "$PGO_DIR/merged.profdata" "$PGO_DIR/merged-${PROFDATA_SHA:0:16}.profdata" + +SHAPES_JSON=$(printf '"%s",' "${SHAPES[@]}") +cat > "$PGO_DIR/train-manifest.json" < merged-${PROFDATA_SHA:0:16}.profdata (${PROFDATA_BYTES}B) ==" >&2 diff --git a/bench/pgo-training/trainer/Cargo.lock b/bench/pgo-training/trainer/Cargo.lock new file mode 100644 index 00000000..19dfc978 --- /dev/null +++ b/bench/pgo-training/trainer/Cargo.lock @@ -0,0 +1,7 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "pgo-trainer" +version = "0.1.0" diff --git a/bench/pgo-training/trainer/Cargo.toml b/bench/pgo-training/trainer/Cargo.toml new file mode 100644 index 00000000..ed624126 --- /dev/null +++ b/bench/pgo-training/trainer/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "pgo-trainer" +version = "0.1.0" +edition = "2021" +rust-version = "1.93" +publish = false +description = "Mock upstream + deterministic load driver for the PGO release build (#967). Std-only on purpose: no network fetch, no workspace lockfile impact." + +# Standalone workspace root: stops cargo's upward workspace discovery at this +# manifest, so the tool builds identically whether the checkout is the repo +# root (Docker/CI) or a git worktree nested under another checkout. +[workspace] + +[[bin]] +name = "pgo-trainer" +path = "src/main.rs" diff --git a/bench/pgo-training/trainer/src/main.rs b/bench/pgo-training/trainer/src/main.rs new file mode 100644 index 00000000..32bc1f2d --- /dev/null +++ b/bench/pgo-training/trainer/src/main.rs @@ -0,0 +1,666 @@ +//! pgo-trainer — deterministic training-traffic generator for the PGO release +//! build (api7/aisix#967). +//! +//! One process, two roles: +//! +//! 1. **Mock upstream** — canned OpenAI / Anthropic dialect responses (JSON + +//! SSE) on a local port; just enough surface for the gateway to dispatch +//! every training shape against it. Streaming responses are written frame +//! by frame over chunked transfer encoding so the gateway's SSE relay loop +//! trains on incremental reads, not one buffered blob. +//! 2. **Load driver** — a fixed number of requests per shape through the +//! gateway, on a small keep-alive HTTP/1.1 client pool. Deterministic by +//! construction: fixed bodies, fixed counts, no randomness, no clocks in +//! the payloads. +//! +//! A PGO profile is a union of hotness, not a traffic-mix replica: each shape +//! is driven in its own gateway process lifetime (see train.sh) and the +//! resulting `.profraw` files are merged. Missing an entire hot path is the +//! failure mode; skewed proportions are not (#967 measured the three-dialect +//! dilution at −1.09%). +//! +//! Std-only on purpose: the release Docker build compiles this tool right +//! before training, and a dependency-free build keeps that phase off the +//! network and off the workspace lockfile. +//! +//! Exit code 0 means every request of the shape succeeded. Anything else is +//! a training failure and MUST fail the calling pipeline — fail-closed, #967 +//! hard gate 2. Never weaken a failure here into a warning. + +use std::io::{BufRead, BufReader, Read, Write}; +use std::net::{TcpListener, TcpStream}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::thread; +use std::time::{Duration, Instant}; + +/// Frames per streaming response. Real LLM streams run tens to hundreds of +/// chunks; a 3-frame canned stream would leave the relay loop body lukewarm. +const STREAM_FRAMES: usize = 48; + +/// Pause between streamed frames so the gateway sees separate reads instead +/// of one coalesced buffer. Kept tiny: it paces training, it does not +/// simulate latency. +const FRAME_PACING: Duration = Duration::from_millis(1); + +/// Per-socket read timeout. Generous because the instrumented gateway runs +/// several times slower than a release build. +const READ_TIMEOUT: Duration = Duration::from_secs(60); + +// ---- shapes ----------------------------------------------------------------- + +/// Byte-identical to `bench/onthebench/lib.sh` `DEFAULT_BODY` — the public +/// benchmark board's default workload. This shape is load-bearing for the +/// program's published numbers and must NEVER be removed or reworded. +const BOARD_BODY: &str = + r#"{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hello"}],"max_tokens":16}"#; + +struct Shape { + name: &'static str, + path: &'static str, + /// `None` = body built at runtime (the padded mid-size one). + body: Option<&'static str>, + anthropic_headers: bool, +} + +/// The v1 training matrix (#967, extended at the 2026-08-13 scenario review): +/// 3 dialects × {stream, non-stream} + mid body + quota gate + native +/// anthropic passthrough + routing kind + embeddings. Error paths are +/// deliberately untrained — they are cold in production and PGO treating +/// them as cold is correct. Maintenance rule: this list mirrors the bench +/// suites; when the product grows a new hot path, add a shape in the same PR. +const SHAPES: &[Shape] = &[ + Shape { + name: "chat-board", + path: "/v1/chat/completions", + body: Some(BOARD_BODY), + anthropic_headers: false, + }, + Shape { + name: "chat-stream", + path: "/v1/chat/completions", + body: Some( + r#"{"model":"gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"hello"}],"max_tokens":64}"#, + ), + anthropic_headers: false, + }, + Shape { + name: "chat-mid", + path: "/v1/chat/completions", + body: None, + anthropic_headers: false, + }, + Shape { + name: "bridge-messages", + path: "/v1/messages", + body: Some( + r#"{"model":"gpt-4o-mini","max_tokens":32,"messages":[{"role":"user","content":"hello"}]}"#, + ), + anthropic_headers: true, + }, + Shape { + name: "bridge-messages-stream", + path: "/v1/messages", + body: Some( + r#"{"model":"gpt-4o-mini","max_tokens":64,"stream":true,"messages":[{"role":"user","content":"hello"}]}"#, + ), + anthropic_headers: true, + }, + Shape { + name: "responses", + path: "/v1/responses", + body: Some(r#"{"model":"gpt-4o-mini","input":"hello"}"#), + anthropic_headers: false, + }, + Shape { + name: "responses-stream", + path: "/v1/responses", + body: Some(r#"{"model":"gpt-4o-mini","input":"hello","stream":true}"#), + anthropic_headers: false, + }, + Shape { + name: "chat-ratelimit", + path: "/v1/chat/completions", + body: Some( + r#"{"model":"gpt-4o-mini-rl","messages":[{"role":"user","content":"hello"}],"max_tokens":16}"#, + ), + anthropic_headers: false, + }, + Shape { + name: "native-messages", + path: "/v1/messages", + body: Some( + r#"{"model":"claude-pgo","max_tokens":32,"messages":[{"role":"user","content":"hello"}]}"#, + ), + anthropic_headers: true, + }, + Shape { + name: "native-messages-stream", + path: "/v1/messages", + body: Some( + r#"{"model":"claude-pgo","max_tokens":64,"stream":true,"messages":[{"role":"user","content":"hello"}]}"#, + ), + anthropic_headers: true, + }, + Shape { + name: "chat-routing", + path: "/v1/chat/completions", + body: Some( + r#"{"model":"gpt-router","messages":[{"role":"user","content":"hello"}],"max_tokens":16}"#, + ), + anthropic_headers: false, + }, + Shape { + name: "embeddings", + path: "/v1/embeddings", + body: Some(r#"{"model":"text-embedding-mock","input":"hello world"}"#), + anthropic_headers: false, + }, +]; + +/// ~4 KB user content: the mid-size tier of the body-size axis. +fn mid_body() -> String { + let filler = "The quick brown fox jumps over the lazy dog. ".repeat(88); + format!( + r#"{{"model":"gpt-4o-mini","messages":[{{"role":"user","content":"{filler}"}}],"max_tokens":16}}"# + ) +} + +// ---- args ------------------------------------------------------------------- + +struct Args { + mock_port: u16, + gateway: String, + api_key: String, + shape: String, + requests: u64, + concurrency: u64, + wait_secs: u64, +} + +fn usage() -> ! { + eprintln!( + "usage: pgo-trainer --mock-port

--gateway --api-key \ + --shape [--requests N] [--concurrency C] [--wait-secs S]\n\ + \x20 pgo-trainer --list-shapes" + ); + std::process::exit(2); +} + +fn parse_args() -> Args { + let mut a = Args { + mock_port: 0, + gateway: String::new(), + api_key: String::new(), + shape: String::new(), + requests: 3000, + concurrency: 8, + wait_secs: 120, + }; + let argv: Vec = std::env::args().skip(1).collect(); + if argv.iter().any(|s| s == "--list-shapes") { + for s in SHAPES { + println!("{}", s.name); + } + std::process::exit(0); + } + let mut i = 0; + while i < argv.len() { + let need = |i: usize| argv.get(i + 1).cloned().unwrap_or_else(|| usage()); + match argv[i].as_str() { + "--mock-port" => a.mock_port = need(i).parse().unwrap_or_else(|_| usage()), + "--gateway" => a.gateway = need(i), + "--api-key" => a.api_key = need(i), + "--shape" => a.shape = need(i), + "--requests" => a.requests = need(i).parse().unwrap_or_else(|_| usage()), + "--concurrency" => a.concurrency = need(i).parse().unwrap_or_else(|_| usage()), + "--wait-secs" => a.wait_secs = need(i).parse().unwrap_or_else(|_| usage()), + _ => usage(), + } + i += 2; + } + if a.mock_port == 0 || a.gateway.is_empty() || a.api_key.is_empty() || a.shape.is_empty() { + usage(); + } + if a.concurrency == 0 || a.requests == 0 { + usage(); + } + a +} + +fn main() { + let args = parse_args(); + let shape = SHAPES + .iter() + .find(|s| s.name == args.shape) + .unwrap_or_else(|| { + eprintln!("FATAL: unknown shape '{}' (see --list-shapes)", args.shape); + std::process::exit(2); + }); + + // Bind before anything else: a bind failure means the port is not ours + // and the profile would be collected against a stranger. Hard exit. + let listener = TcpListener::bind(("127.0.0.1", args.mock_port)).unwrap_or_else(|e| { + eprintln!("FATAL: mock bind 127.0.0.1:{} failed: {e}", args.mock_port); + std::process::exit(2); + }); + thread::spawn(move || { + for conn in listener.incoming().flatten() { + thread::spawn(move || mock_conn(conn)); + } + }); + + if let Err(e) = wait_gateway(&args.gateway, args.wait_secs) { + eprintln!("FATAL: gateway not ready: {e}"); + std::process::exit(3); + } + + let body: Arc = Arc::new(match shape.body { + Some(b) => b.to_string(), + None => mid_body(), + }); + let ok = Arc::new(AtomicU64::new(0)); + let started = Instant::now(); + let mut workers = Vec::new(); + for w in 0..args.concurrency { + // Deterministic split: the first (requests % concurrency) workers + // take one extra request. + let n = args.requests / args.concurrency + u64::from(w < args.requests % args.concurrency); + if n == 0 { + continue; + } + let gateway = args.gateway.clone(); + let api_key = args.api_key.clone(); + let body = Arc::clone(&body); + let ok = Arc::clone(&ok); + let path = shape.path; + let anthropic = shape.anthropic_headers; + workers.push(thread::spawn(move || -> Result<(), String> { + drive(&gateway, &api_key, path, anthropic, &body, n, &ok) + })); + } + + let mut failures = Vec::new(); + for w in workers { + match w.join() { + Ok(Ok(())) => {} + Ok(Err(e)) => failures.push(e), + Err(_) => failures.push("worker panicked".to_string()), + } + } + let done = ok.load(Ordering::Relaxed); + println!( + "shape={} ok={}/{} elapsed={:.1}s", + shape.name, + done, + args.requests, + started.elapsed().as_secs_f64() + ); + if !failures.is_empty() || done != args.requests { + for f in failures.iter().take(4) { + eprintln!("FATAL: {f}"); + } + std::process::exit(1); + } +} + +// ---- driver (client side) ----------------------------------------------------- + +fn wait_gateway(addr: &str, secs: u64) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(secs); + let mut last = String::from("never connected"); + while Instant::now() < deadline { + match TcpStream::connect(addr) { + Ok(mut s) => { + s.set_read_timeout(Some(Duration::from_secs(5))).ok(); + let req = + format!("GET /livez HTTP/1.1\r\nhost: {addr}\r\nconnection: close\r\n\r\n"); + if s.write_all(req.as_bytes()).is_ok() { + let mut buf = [0u8; 64]; + if let Ok(n) = s.read(&mut buf) { + let head = String::from_utf8_lossy(&buf[..n]).to_string(); + if head.starts_with("HTTP/1.1 200") || head.starts_with("HTTP/1.0 200") { + return Ok(()); + } + last = head.lines().next().unwrap_or("").to_string(); + } + } + } + Err(e) => last = e.to_string(), + } + thread::sleep(Duration::from_millis(200)); + } + Err(format!("gave up after {secs}s (last: {last})")) +} + +fn connect(gateway: &str) -> Result<(BufReader, TcpStream), String> { + let stream = TcpStream::connect(gateway).map_err(|e| format!("connect {gateway}: {e}"))?; + stream.set_read_timeout(Some(READ_TIMEOUT)).ok(); + stream.set_nodelay(true).ok(); + let write = stream.try_clone().map_err(|e| format!("clone: {e}"))?; + Ok((BufReader::new(stream), write)) +} + +fn drive( + gateway: &str, + api_key: &str, + path: &str, + anthropic: bool, + body: &str, + n: u64, + ok: &AtomicU64, +) -> Result<(), String> { + let extra = if anthropic { + "anthropic-version: 2023-06-01\r\n" + } else { + "" + }; + let request = format!( + "POST {path} HTTP/1.1\r\nhost: {gateway}\r\nauthorization: Bearer {api_key}\r\n\ + {extra}content-type: application/json\r\ncontent-length: {}\r\n\r\n{body}", + body.len() + ); + let (mut reader, mut writer) = connect(gateway)?; + for i in 0..n { + writer + .write_all(request.as_bytes()) + .map_err(|e| format!("{path} req {i}: write: {e}"))?; + let (status, reused) = + read_response(&mut reader).map_err(|e| format!("{path} req {i}: {e}"))?; + if status != 200 { + return Err(format!("{path} req {i}: HTTP {status}")); + } + ok.fetch_add(1, Ordering::Relaxed); + if !reused { + let (r, w) = connect(gateway)?; + reader = r; + writer = w; + } + } + Ok(()) +} + +/// Read one HTTP/1.1 response. Returns (status, connection-reusable). +/// Handles both framings the gateway emits: content-length (JSON) and +/// chunked (SSE) — reading to the zero chunk is dialect-agnostic, so one +/// client covers every streaming shape. +fn read_response(r: &mut BufReader) -> Result<(u16, bool), String> { + let mut line = String::new(); + r.read_line(&mut line).map_err(|e| format!("status: {e}"))?; + if line.is_empty() { + return Err("connection closed before status line".into()); + } + let status: u16 = line + .split_whitespace() + .nth(1) + .and_then(|s| s.parse().ok()) + .ok_or_else(|| format!("bad status line: {line:?}"))?; + + let mut content_length: Option = None; + let mut chunked = false; + let mut close = false; + loop { + let mut h = String::new(); + r.read_line(&mut h).map_err(|e| format!("header: {e}"))?; + let h = h.trim_end(); + if h.is_empty() { + break; + } + let lower = h.to_ascii_lowercase(); + if let Some(v) = lower.strip_prefix("content-length:") { + content_length = Some( + v.trim() + .parse() + .map_err(|_| format!("bad content-length: {h:?}"))?, + ); + } else if lower.starts_with("transfer-encoding:") && lower.contains("chunked") { + chunked = true; + } else if lower.starts_with("connection:") && lower.contains("close") { + close = true; + } + } + + if chunked { + loop { + let mut sz = String::new(); + r.read_line(&mut sz) + .map_err(|e| format!("chunk size: {e}"))?; + let sz = usize::from_str_radix(sz.trim().split(';').next().unwrap_or(""), 16) + .map_err(|_| format!("bad chunk size: {sz:?}"))?; + let mut chunk = vec![0u8; sz + 2]; // data + CRLF + r.read_exact(&mut chunk) + .map_err(|e| format!("chunk body: {e}"))?; + if sz == 0 { + break; + } + } + } else if let Some(len) = content_length { + let mut body = vec![0u8; len]; + r.read_exact(&mut body).map_err(|e| format!("body: {e}"))?; + } else { + // No length and not chunked: body runs to EOF; connection is done. + let mut sink = Vec::new(); + r.read_to_end(&mut sink).map_err(|e| format!("body: {e}"))?; + return Ok((status, false)); + } + Ok((status, !close)) +} + +// ---- mock upstream (server side) ---------------------------------------------- + +fn mock_conn(stream: TcpStream) { + stream.set_read_timeout(Some(READ_TIMEOUT)).ok(); + stream.set_nodelay(true).ok(); + let write = match stream.try_clone() { + Ok(w) => w, + Err(_) => return, + }; + let mut reader = BufReader::new(stream); + let mut writer = write; + loop { + let (path, body) = match read_mock_request(&mut reader) { + Ok(Some(v)) => v, + _ => return, + }; + let stream_requested = + body.contains(r#""stream":true"#) || body.contains(r#""stream": true"#); + let result = match path.as_str() { + "/v1/chat/completions" => { + if stream_requested { + write_sse(&mut writer, &chat_sse_frames()) + } else { + write_json(&mut writer, CHAT_JSON) + } + } + "/v1/responses" => { + if stream_requested { + write_sse(&mut writer, &responses_sse_frames()) + } else { + write_json(&mut writer, RESPONSES_JSON) + } + } + "/v1/messages" => { + if stream_requested { + write_sse(&mut writer, &anthropic_sse_frames()) + } else { + write_json(&mut writer, ANTHROPIC_JSON) + } + } + "/v1/embeddings" => write_json(&mut writer, EMBEDDINGS_JSON), + // Anything else is a routing bug in the training setup: answer + // 404 so the driving request fails loudly instead of training a + // wrong path. + _ => { + let resp = format!( + "HTTP/1.1 404 Not Found\r\ncontent-length: 0\r\nconnection: close\r\n\r\n" + ); + let _ = writer.write_all(resp.as_bytes()); + return; + } + }; + if result.is_err() { + return; + } + } +} + +fn read_mock_request(r: &mut BufReader) -> Result, String> { + let mut line = String::new(); + let n = r.read_line(&mut line).map_err(|e| e.to_string())?; + if n == 0 { + return Ok(None); // clean EOF between keep-alive requests + } + let path = line + .split_whitespace() + .nth(1) + .unwrap_or("") + .split('?') + .next() + .unwrap_or("") + .to_string(); + let mut content_length = 0usize; + loop { + let mut h = String::new(); + r.read_line(&mut h).map_err(|e| e.to_string())?; + let h = h.trim_end(); + if h.is_empty() { + break; + } + let lower = h.to_ascii_lowercase(); + if let Some(v) = lower.strip_prefix("content-length:") { + content_length = v.trim().parse().unwrap_or(0); + } + } + let mut body = vec![0u8; content_length]; + if content_length > 0 { + r.read_exact(&mut body).map_err(|e| e.to_string())?; + } + Ok(Some((path, String::from_utf8_lossy(&body).into_owned()))) +} + +fn write_json(w: &mut TcpStream, body: &str) -> std::io::Result<()> { + let resp = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\n\r\n{body}", + body.len() + ); + w.write_all(resp.as_bytes()) +} + +/// Stream SSE frames over chunked transfer encoding, one chunk per frame with +/// a short pause, so the gateway's relay loop sees many small reads — the +/// production streaming profile, not one buffered blob. +fn write_sse(w: &mut TcpStream, frames: &[String]) -> std::io::Result<()> { + w.write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n\ + cache-control: no-cache\r\ntransfer-encoding: chunked\r\n\r\n", + )?; + for f in frames { + let chunk = format!("{:x}\r\n{f}\r\n", f.len()); + w.write_all(chunk.as_bytes())?; + w.flush()?; + thread::sleep(FRAME_PACING); + } + w.write_all(b"0\r\n\r\n") +} + +// ---- canned upstream bodies ----------------------------------------------------- + +const CHAT_JSON: &str = r#"{"id":"chatcmpl-pgo","object":"chat.completion","created":0,"model":"gpt-4o-mini","choices":[{"index":0,"message":{"role":"assistant","content":"PGO training reply."},"finish_reason":"stop"}],"usage":{"prompt_tokens":9,"completion_tokens":12,"total_tokens":21}}"#; + +const RESPONSES_JSON: &str = r#"{"id":"resp-pgo","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg-pgo","role":"assistant","content":[{"type":"output_text","text":"PGO training reply."}]}],"usage":{"input_tokens":9,"output_tokens":12,"total_tokens":21}}"#; + +const ANTHROPIC_JSON: &str = r#"{"id":"msg-pgo","type":"message","role":"assistant","model":"claude-pgo","content":[{"type":"text","text":"PGO training reply."}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":9,"output_tokens":12}}"#; + +const EMBEDDINGS_JSON: &str = r#"{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.01,-0.02,0.03,-0.04,0.05,-0.06,0.07,-0.08,0.09,-0.1,0.11,-0.12,0.13,-0.14,0.15,-0.16]}],"model":"text-embedding-mock","usage":{"prompt_tokens":3,"total_tokens":3}}"#; + +/// OpenAI chat SSE: role delta, N content deltas, a final chunk carrying +/// `finish_reason` + `usage`, then the `[DONE]` sentinel. +fn chat_sse_frames() -> Vec { + let mut f = Vec::with_capacity(STREAM_FRAMES + 1); + f.push(concat!( + r#"data: {"id":"chatcmpl-pgo","object":"chat.completion.chunk","created":0,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}"#, + "\n\n" + ) + .to_string()); + for _ in 0..STREAM_FRAMES - 2 { + f.push(concat!( + r#"data: {"id":"chatcmpl-pgo","object":"chat.completion.chunk","created":0,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"chunk "},"finish_reason":null}]}"#, + "\n\n" + ) + .to_string()); + } + f.push(concat!( + r#"data: {"id":"chatcmpl-pgo","object":"chat.completion.chunk","created":0,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":9,"completion_tokens":46,"total_tokens":55}}"#, + "\n\n" + ) + .to_string()); + f.push("data: [DONE]\n\n".to_string()); + f +} + +/// OpenAI Responses SSE: text deltas, then the terminal `response.completed` +/// event whose `usage` block the gateway parses in-flight (#808), then +/// `[DONE]`. +fn responses_sse_frames() -> Vec { + let mut f = Vec::with_capacity(STREAM_FRAMES + 1); + for _ in 0..STREAM_FRAMES - 1 { + f.push( + concat!( + "event: response.output_text.delta\n", + r#"data: {"type":"response.output_text.delta","delta":"chunk "}"#, + "\n\n" + ) + .to_string(), + ); + } + f.push(concat!( + "event: response.completed\n", + r#"data: {"type":"response.completed","response":{"id":"resp-pgo","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg-pgo","role":"assistant","content":[{"type":"output_text","text":"PGO training reply."}]}],"usage":{"input_tokens":9,"output_tokens":47,"total_tokens":56}}}"#, + "\n\n" + ) + .to_string()); + f.push("data: [DONE]\n\n".to_string()); + f +} + +/// Anthropic Messages SSE: message_start → content_block deltas → +/// message_delta (usage) → message_stop, per the Messages streaming format. +fn anthropic_sse_frames() -> Vec { + let mut f = Vec::with_capacity(STREAM_FRAMES); + f.push(concat!( + "event: message_start\n", + r#"data: {"type":"message_start","message":{"id":"msg-pgo","type":"message","role":"assistant","model":"claude-pgo","content":[],"stop_reason":null,"usage":{"input_tokens":9,"output_tokens":1}}}"#, + "\n\n" + ) + .to_string()); + f.push(concat!( + "event: content_block_start\n", + r#"data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}"#, + "\n\n" + ) + .to_string()); + for _ in 0..STREAM_FRAMES - 5 { + f.push(concat!( + "event: content_block_delta\n", + r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"chunk "}}"#, + "\n\n" + ) + .to_string()); + } + f.push( + concat!( + "event: content_block_stop\n", + r#"data: {"type":"content_block_stop","index":0}"#, + "\n\n" + ) + .to_string(), + ); + f.push(concat!( + "event: message_delta\n", + r#"data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":43}}"#, + "\n\n" + ) + .to_string()); + f.push("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n".to_string()); + f +}