feat: add public POST /tokenize endpoint #43

Merged

rcsheets merged 2 commits from fix/request-too-large-4xx-and-metrics into main 2026-07-22 05:07:07 +00:00

Conversation 2 Commits 2 Files changed 10 +265 -5

rcsheets commented 2026-07-22 01:49:52 +00:00

Owner

Copy link

Also fix a scheduler error handling bug and add size metrics.

Also fix a scheduler error handling bug and add size metrics.

rcsheets added 2 commits 2026-07-22 01:49:52 +00:00

fix(scheduler): return 4xx for oversized requests, add size metrics ... dc68b738c4

Requests larger than a hard capacity limit were surfaced as HTTP 500.
They are client-caused and unretryable as-is, so map them to 400.

- Add scheduler.ErrRequestTooLarge, a shared sentinel wrapping both the
  per-batch-budget and whole-KV-cache rejections. writeEngineError maps
  it to 400 before the ERROR log, so oversized input no longer logs as an
  internal failure.
- The rejection message now reports the tokenized size, the limit, and a
  proportional hint (shorten by ~N%), since clients send untokenized text
  and cannot predict the token count.
- Add gllm_requests_too_large_total{reason=batch|kv_cache}, a lagging
  counter split by the binding constraint (kv_cache signals cache
  under-provisioning). Both label values pre-registered at 0.
- Add gllm_request_prompt_tokens, a histogram of accepted-request prompt
  sizes observed once per completed request -- the leading indicator: a
  rising high percentile shows headroom shrinking before rejections begin.
- Add kvcache.Manager.TokenCapacity() for the whole-cache limit in the
  message and metric.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

feat(server): add public POST /tokenize endpoint ... 7b95dd3640

A vLLM-compatible tokenize-only endpoint so a client can size a request
against the batch budget before sending it. Complements the 4xx hint on
oversized requests: the server is the only authority on the token count
(it applies the model's chat template and special tokens), so a client
cannot reliably reproduce it locally.

- POST /tokenize takes exactly one of {prompt, messages} and returns
  {count, tokens}. messages applies the chat template as
  /v1/chat/completions does; prompt encodes raw text as /v1/completions
  does. Runs no forward pass -- pure CPU, off the scheduler/KV path.
- Engine.Tokenize/TokenizeChat wrap the exact Encode/EncodeChat calls
  generation uses, so the count is identical to what prefill will see.
- Consistency test asserts Tokenize(x) length equals a real generation's
  PromptTokens for the same input; server test covers routing and the
  exactly-one-of validation.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

rcsheets changed title from feat: add public POST /tokenize endpoint to WIP: feat: add public POST /tokenize endpoint 2026-07-22 04:15:03 +00:00

rcsheets changed title from WIP: feat: add public POST /tokenize endpoint to feat: add public POST /tokenize endpoint 2026-07-22 04:15:05 +00:00

rcsheets requested review from pr-reviewer-bot 2026-07-22 04:15:14 +00:00

pr-reviewer-bot commented 2026-07-22 04:16:09 +00:00

Collaborator

Copy link

Automated review by pr-reviewer v0.36.0 | Full Review | Claude Sonnet 5 | tracking id r-604409-a1ff98
This is an AI-generated review and may contain mistakes.

Status: ✅ Completed

---

⚠️ The primary review (Mistral Small) failed. The review below is from Claude Sonnet 5, which reviewed the same diff.

✅ Verdict: Approve — The PR is correct, well-tested for its primary behavior change (the /tokenize count matching generation, and the 4xx error mapping fix), and stylistically consistent; the only gap is missing direct tests for the two new metrics' wiring through the engine, which is minor enough not to block.

Overall this is a well-scoped, well-tested PR that does three related things: adds POST /tokenize, fixes the scheduler's "too large" error handling so it maps to a proper client error (4xx) instead of an opaque internal error, and adds a couple of Prometheus metrics (gllm_requests_too_large_total, gllm_request_prompt_tokens) to observe that condition. Going through each:

/tokenize endpoint

• Engine.Tokenize/TokenizeChat reuse the exact encode paths Generate/GenerateChat use, and TestTokenizeMatchesGenerate pins that the count is identical to what a real generation's prefill sees -- that's the right thing to test given the endpoint's whole value proposition is "trust this count."
• handleTokenize validates exactly-one-of prompt/messages, and routes correctly. checkModel gates it like the other endpoints. Good, consistent style with the rest of openai.go/server.go.
• Minor: TokenizeResponse.Tokens always returns the full token id array. For a large prompt (e.g. thousands of tokens) that's a nontrivial response payload just to get a count -- vLLM's /tokenize does return tokens too, so this matches precedent, but it's worth a beat of thought since this is explicitly a "check before sending the real thing" endpoint and could itself be heavy for huge prompts. Not a blocker, just a note; no request size cap here either, but that already exists at the scheduler level for the request proper.
• No new architecture doc/AGENTS.md needed beyond the README bullet already added, which reads well.

Scheduler error-handling fix

• Previously the two "too large" failure paths in Schedule() produced plain fmt.Errorf with no wrapped sentinel, so the server's generic writeEngineError mapped them to 500 Internal Server Error even though they're unambiguously client errors (oversized input, unrecoverable). The new ErrRequestTooLarge sentinel + errors.Is check in writeEngineError fixes that cleanly, matching the existing pattern for backend.ErrNotImplemented. This is a real bug fix and a sensible one.
• tooLargeErr computing a percentage hint is a nice touch, though the division over/n*100 looks like it should probably be over/limit or something proportional to what to trim off -- but on reflection, dividing by n (the request's own token count) does directly answer "shorten the input by X%", which is exactly what's promised in the message. Let me double check there's no divide-by-zero: n is the tokenized sequence length, which for the batch-budget path is seq.Len() (prompt+output) and for the KV-cache path is also seq.Len(); both call sites only reach tooLargeErr when n > limit >= 0... actually limit could be 0 in theory for the batch path if MaxTokensPerBatch is 0, but n itself being 0 would mean an empty sequence, which wouldn't trigger n > limit unless limit is negative. n == 0 seems structurally unreachable given Len() requires at least prompt tokens, and prompts are non-empty by construction upstream (tokenizer emits at least BOS). This is a minor latent risk but not one this PR introduces or worsens; it's consistent with prior behavior lacking any hint at all.
• Test updates (scheduler_test.go) properly switch from raw substring matching to errors.Is(seq.Err, ErrRequestTooLarge), which is more robust to message wording changes and directly enforces the sentinel-wrapping contract. Good.
• TokenCapacity() added to kvcache.Manager is a clean, minimal addition purely to support building a better error message; no side effects, trivially correct.

Metrics additions

• RequestsTooLarge (CounterVec by reason) and RequestPromptTokens (Histogram) are added to metrics.Metrics, pre-registering both reason label values up front (a nice touch per the existing style of "continuous series for alerting" seen elsewhere in the codebase, e.g. this mirrors how other counters are described in AGENTS.md).
• The engine mirrors the scheduler's cumulative too-large counts into the counter under e.mu in step(), following the exact same pattern already used for Preemptions -- consistent design, no new locking risk.
• RequestPromptTokens.Observe is called once per successfully completed generate in engine.go's generate() at the <-seq.Done branch, which is the right place (mirrors the existing "request complete" log call site) and only observes prompt tokens for requests that actually ran to completion (as documented). This does mean prompt token counts for rejected (too-large) requests are not captured in the histogram -- that's fine and intentional per the doc comment ("accepted requests"), but worth flagging as a slight gap: an operator watching this histogram to anticipate rejections won't see the rejected request's own size reflected in the percentile (only the accepted ones near the edge). That's a reasonable design tradeoff given accepted-request sizing is the intended signal, not a bug.
• No dedicated metrics test file changed for these two new metrics themselves (i.e., no direct unit test asserting RequestsTooLarge increments correctly end-to-end through the engine, e.g. verifying label values or the histogram observation). The PR's own summary flags internal/metrics/metrics.go, internal/engine/engine.go, internal/kvcache/kvcache.go, internal/scheduler/scheduler.go, and internal/server/openai.go as lacking direct tests for these changes. Some of this is indirectly covered (scheduler_test.go checks the sentinel wraps; server_test.go's TestErrorMapping checks the 400 mapping), but there's no test asserting the gllm_requests_too_large_total{reason=...} counter or the gllm_request_prompt_tokens histogram actually get bumped by a real end-to-end oversized-request or completed-request flow through the engine. Given TestPerfStatsCounters already demonstrates the pattern for existing counters, a similar test extending coverage to these two would be low-cost and valuable, especially since metric-plumbing bugs (wrong label, off-by-one in mirroring cumulative counts) are exactly the kind of thing that silently rot unnoticed. I'd call this a gap worth addressing, though not one that blocks merge given the underlying logic (scheduler counting, mirroring pattern) is proven correct by existing analogous tests.

Other observations

• server.go's Generator interface and fakeEngine test double are updated in lockstep for Tokenize/TokenizeChat, so any other real or fake implementer would fail to compile if forgotten -- good, no silent breakage risk here per Go's interface satisfaction.
• No breaking changes to existing wire types; TokenizeRequest/TokenizeResponse are new types, and the scheduler error message text changed slightly (e.g., "max batch tokens" -> "the max batch size") but no code outside the test suite depended on exact wording (confirmed by grep of the diff -- only the test assertions changed correspondingly).
• Style/conventions: ASCII only, ends of files match project comment density norms, consistent doc comments. No issues there.

Nothing here looks incorrect or unsafe; the fixes are genuine improvements (proper 4xx mapping, richer error messages, and useful new metrics), and the added test coverage for the primary contract (/tokenize count matches generation) is exactly the kind of test called for. The one real ask is a bit more direct test coverage of the new metrics wiring, which is a coverage gap rather than a functional problem.

📚 Context used in this review

Project guidance files:

• CLAUDE.md
• AGENTS.md

*Automated review by [pr-reviewer](https://git.brooktrails.org/brooktrails/pr-reviewer) v0.36.0 | Full Review | Claude Sonnet 5 | tracking id `r-604409-a1ff98`* *This is an AI-generated review and may contain mistakes.* **Status:** ✅ Completed --- > ⚠️ The primary review (**Mistral Small**) failed. The review below is from **Claude Sonnet 5**, which reviewed the same diff. **✅ Verdict: Approve** — The PR is correct, well-tested for its primary behavior change (the /tokenize count matching generation, and the 4xx error mapping fix), and stylistically consistent; the only gap is missing direct tests for the two new metrics' wiring through the engine, which is minor enough not to block. Overall this is a well-scoped, well-tested PR that does three related things: adds `POST /tokenize`, fixes the scheduler's "too large" error handling so it maps to a proper client error (4xx) instead of an opaque internal error, and adds a couple of Prometheus metrics (`gllm_requests_too_large_total`, `gllm_request_prompt_tokens`) to observe that condition. Going through each: ## `/tokenize` endpoint - `Engine.Tokenize`/`TokenizeChat` reuse the exact encode paths `Generate`/`GenerateChat` use, and `TestTokenizeMatchesGenerate` pins that the count is identical to what a real generation's prefill sees -- that's the right thing to test given the endpoint's whole value proposition is "trust this count." - `handleTokenize` validates exactly-one-of prompt/messages, and routes correctly. `checkModel` gates it like the other endpoints. Good, consistent style with the rest of `openai.go`/`server.go`. - Minor: `TokenizeResponse.Tokens` always returns the full token id array. For a large prompt (e.g. thousands of tokens) that's a nontrivial response payload just to get a count -- vLLM's `/tokenize` does return `tokens` too, so this matches precedent, but it's worth a beat of thought since this is explicitly a "check before sending the real thing" endpoint and could itself be heavy for huge prompts. Not a blocker, just a note; no request size cap here either, but that already exists at the scheduler level for the request proper. - No new architecture doc/AGENTS.md needed beyond the README bullet already added, which reads well. ## Scheduler error-handling fix - Previously the two "too large" failure paths in `Schedule()` produced plain `fmt.Errorf` with no wrapped sentinel, so the server's generic `writeEngineError` mapped them to 500 Internal Server Error even though they're unambiguously client errors (oversized input, unrecoverable). The new `ErrRequestTooLarge` sentinel + `errors.Is` check in `writeEngineError` fixes that cleanly, matching the existing pattern for `backend.ErrNotImplemented`. This is a real bug fix and a sensible one. - `tooLargeErr` computing a percentage hint is a nice touch, though the division `over/n*100` looks like it should probably be `over/limit` or something proportional to what to trim off -- but on reflection, dividing by `n` (the request's own token count) does directly answer "shorten the input by X%", which is exactly what's promised in the message. Let me double check there's no divide-by-zero: `n` is the tokenized sequence length, which for the batch-budget path is `seq.Len()` (prompt+output) and for the KV-cache path is also `seq.Len()`; both call sites only reach `tooLargeErr` when `n > limit >= 0`... actually `limit` could be 0 in theory for the batch path if `MaxTokensPerBatch` is 0, but `n` itself being 0 would mean an empty sequence, which wouldn't trigger `n > limit` unless limit is negative. `n == 0` seems structurally unreachable given `Len()` requires at least prompt tokens, and prompts are non-empty by construction upstream (tokenizer emits at least BOS). This is a minor latent risk but not one this PR introduces or worsens; it's consistent with prior behavior lacking any hint at all. - Test updates (`scheduler_test.go`) properly switch from raw substring matching to `errors.Is(seq.Err, ErrRequestTooLarge)`, which is more robust to message wording changes and directly enforces the sentinel-wrapping contract. Good. - `TokenCapacity()` added to `kvcache.Manager` is a clean, minimal addition purely to support building a better error message; no side effects, trivially correct. ## Metrics additions - `RequestsTooLarge` (CounterVec by reason) and `RequestPromptTokens` (Histogram) are added to `metrics.Metrics`, pre-registering both `reason` label values up front (a nice touch per the existing style of "continuous series for alerting" seen elsewhere in the codebase, e.g. this mirrors how other counters are described in AGENTS.md). - The engine mirrors the scheduler's cumulative too-large counts into the counter under `e.mu` in `step()`, following the exact same pattern already used for `Preemptions` -- consistent design, no new locking risk. - `RequestPromptTokens.Observe` is called once per successfully completed generate in `engine.go`'s `generate()` at the `<-seq.Done` branch, which is the right place (mirrors the existing "request complete" log call site) and only observes prompt tokens for requests that actually ran to completion (as documented). This does mean prompt token counts for *rejected* (too-large) requests are not captured in the histogram -- that's fine and intentional per the doc comment ("accepted requests"), but worth flagging as a slight gap: an operator watching this histogram to anticipate rejections won't see the rejected request's own size reflected in the percentile (only the accepted ones near the edge). That's a reasonable design tradeoff given accepted-request sizing is the intended signal, not a bug. - No dedicated metrics test file changed for these two new metrics themselves (i.e., no direct unit test asserting `RequestsTooLarge` increments correctly end-to-end through the engine, e.g. verifying label values or the histogram observation). The PR's own summary flags `internal/metrics/metrics.go`, `internal/engine/engine.go`, `internal/kvcache/kvcache.go`, `internal/scheduler/scheduler.go`, and `internal/server/openai.go` as lacking direct tests for these changes. Some of this is indirectly covered (scheduler_test.go checks the sentinel wraps; server_test.go's `TestErrorMapping` checks the 400 mapping), but there's no test asserting the `gllm_requests_too_large_total{reason=...}` counter or the `gllm_request_prompt_tokens` histogram actually get bumped by a real end-to-end oversized-request or completed-request flow through the engine. Given `TestPerfStatsCounters` already demonstrates the pattern for existing counters, a similar test extending coverage to these two would be low-cost and valuable, especially since metric-plumbing bugs (wrong label, off-by-one in mirroring cumulative counts) are exactly the kind of thing that silently rot unnoticed. I'd call this a gap worth addressing, though not one that blocks merge given the underlying logic (scheduler counting, mirroring pattern) is proven correct by existing analogous tests. ## Other observations - `server.go`'s `Generator` interface and `fakeEngine` test double are updated in lockstep for `Tokenize`/`TokenizeChat`, so any other real or fake implementer would fail to compile if forgotten -- good, no silent breakage risk here per Go's interface satisfaction. - No breaking changes to existing wire types; `TokenizeRequest`/`TokenizeResponse` are new types, and the scheduler error message text changed slightly (e.g., "max batch tokens" -> "the max batch size") but no code outside the test suite depended on exact wording (confirmed by grep of the diff -- only the test assertions changed correspondingly). - Style/conventions: ASCII only, ends of files match project comment density norms, consistent doc comments. No issues there. Nothing here looks incorrect or unsafe; the fixes are genuine improvements (proper 4xx mapping, richer error messages, and useful new metrics), and the added test coverage for the primary contract (`/tokenize` count matches generation) is exactly the kind of test called for. The one real ask is a bit more direct test coverage of the new metrics wiring, which is a coverage gap rather than a functional problem. <details> <summary>📚 Context used in this review</summary> **Project guidance files:** - `CLAUDE.md` - `AGENTS.md` </details>

rcsheets force-pushed fix/request-too-large-4xx-and-metrics from 7b95dd3640 to c67c2949f7 2026-07-22 05:06:11 +00:00 Compare

pr-reviewer-bot commented 2026-07-22 05:06:13 +00:00

Collaborator

Copy link

Automated review by pr-reviewer v0.36.0 | Safety Check | Claude Sonnet 5 | tracking id r-604fc4-c5cf45
This is an AI-generated review and may contain mistakes.

Status: ✅ Completed

---

⚠️ The primary review (Mistral Small) failed. The review below is from Claude Sonnet 5, which reviewed the same diff.

✅ Verdict: LGTM — Additive feature (new /tokenize endpoint) plus a sensible scheduler error-handling fix and new metrics; no security issues, accidental commits, or breaking changes spotted.

Reviewed the diff:

• New /tokenize endpoint: validates exactly-one-of prompt/messages, routes to Tokenize/TokenizeChat, returns count+tokens. Straightforward, tested.
• Scheduler bug fix: error messages now use a shared tooLargeErr helper with a wrapped ErrRequestTooLarge sentinel; server maps it to 400 instead of 500. Tests updated to use errors.Is. Looks correct.
• tooLargeErr computes pct := ceil((n-limit)/n * 100) — division by n (tokenized count), guarded implicitly since n > limit >= 0 in both call sites, and n can't be 0 when this path is hit (already exceeds a limit), so no realistic divide-by-zero.
• Metrics additions (RequestsTooLarge counter, RequestPromptTokens histogram) are additive, registered as new metrics, don't touch existing metric names.
• Engine mirrors scheduler's cumulative too-large counts into the metric, using a monotonic diff pattern like the existing Preemptions handling — consistent with existing code, correctly guarded under the same mutex.
• Server interface extended with Tokenize/TokenizeChat; fakeEngine in tests updated accordingly.
• No secrets, credentials, or unrelated binary files added. No obvious breaking change to existing endpoints — new route is additive, existing error-mapping logic extended with a new branch that doesn't disturb prior branches.

Nothing here looks like an accidental commit, security issue, or an obvious bug that would break existing behavior.

*Automated review by [pr-reviewer](https://git.brooktrails.org/brooktrails/pr-reviewer) v0.36.0 | Safety Check | Claude Sonnet 5 | tracking id `r-604fc4-c5cf45`* *This is an AI-generated review and may contain mistakes.* **Status:** ✅ Completed --- > ⚠️ The primary review (**Mistral Small**) failed. The review below is from **Claude Sonnet 5**, which reviewed the same diff. **✅ Verdict: LGTM** — Additive feature (new /tokenize endpoint) plus a sensible scheduler error-handling fix and new metrics; no security issues, accidental commits, or breaking changes spotted. Reviewed the diff: - New `/tokenize` endpoint: validates exactly-one-of prompt/messages, routes to Tokenize/TokenizeChat, returns count+tokens. Straightforward, tested. - Scheduler bug fix: error messages now use a shared `tooLargeErr` helper with a wrapped `ErrRequestTooLarge` sentinel; server maps it to 400 instead of 500. Tests updated to use `errors.Is`. Looks correct. - `tooLargeErr` computes `pct := ceil((n-limit)/n * 100)` — division by `n` (tokenized count), guarded implicitly since `n` > `limit` >= 0 in both call sites, and n can't be 0 when this path is hit (already exceeds a limit), so no realistic divide-by-zero. - Metrics additions (RequestsTooLarge counter, RequestPromptTokens histogram) are additive, registered as new metrics, don't touch existing metric names. - Engine mirrors scheduler's cumulative too-large counts into the metric, using a monotonic diff pattern like the existing Preemptions handling — consistent with existing code, correctly guarded under the same mutex. - Server interface extended with Tokenize/TokenizeChat; fakeEngine in tests updated accordingly. - No secrets, credentials, or unrelated binary files added. No obvious breaking change to existing endpoints — new route is additive, existing error-mapping logic extended with a new branch that doesn't disturb prior branches. Nothing here looks like an accidental commit, security issue, or an obvious bug that would break existing behavior.

rcsheets merged commit df51a84dfa into main 2026-07-22 05:07:07 +00:00

rcsheets deleted branch fix/request-too-large-4xx-and-metrics 2026-07-22 05:07:07 +00:00

rcsheets referenced this pull request from a commit 2026-07-22 05:07:09 +00:00

Merge pull request 'feat: add public POST /tokenize endpoint' (#43) from fix/request-too-large-4xx-and-metrics into main

Sign in to join this conversation.

Reviewers

No reviewers

pr-reviewer-bot

Labels

Clear labels

No items

No labels

Milestone

Clear milestone

No items

No milestone

Projects

Clear projects

No items

No project

Assignees

Clear assignees

No assignees

2 participants

Notifications

Subscribe

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set

Reference

brooktrails/gllm!43

No description provided.