Skip to content

perf(epp): skip the RawBody copy for bodies above the pool cap - #3883

Merged
ahg-g merged 2 commits into
llm-d:mainfrom
thomas-fahrner-parasail:epp-large-body-no-copy
Nov 5, 2026
Merged

ahg-g merged 2 commits into
llm-d:mainfrom
thomas-fahrner-parasail:epp-large-body-no-copy

Conversation

@thomas-fahrner-parasail

@thomas-fahrner-parasail thomas-fahrner-parasail commented Nov 2, 2026 •

Copy link
Copy Markdown
Contributor

What type of PR is this?
/kind performance

What this PR does / why we need it:

putBuffer drops a request buffer larger than the pool cap (--grpc-max-recv-msg-size, 4 MiB by default). The handler still copies that buffer into RawBody, so a large body takes 2 times its size while it is parsed and scheduled.

With this PR, RawBody uses the memory of a buffer above the pool cap, and the handler does not return that buffer to the pool. A buffer at or below the cap keeps the copy and returns to the pool. Parsers and the Director replace RawBody when they rewrite the body and do not write into it, so the inbound bytes stay valid for the tool-calling snapshot comparison.

For a 20 MB body, this removes 1 copy of 20 MB for the duration of the request phase.

Test plan:

  • A buffer above the pool cap: RawBody shares its memory
  • A buffer at or below the pool cap, and a pool without a cap: RawBody is a copy
  • Tool-calling telemetry with a buffer above the pool cap: the snapshot comparison runs on the shared memory and reports tools as preserved
  • make presubmit

Which issue(s) this PR fixes:
Part of #3881

Depends on #3884: merge this PR after #3884. RawBody holds the buffer's whole backing array, which is up to 2 times the body size (32 MiB for a 20 MB body). Without #3884, RawBody stays live until the stream ends, so a long streamed response would hold more memory than the copy it replaces.

Release note (write NONE if no user-facing change):

Lower EPP memory use for request bodies larger than --grpc-max-recv-msg-size: the EPP keeps 1 copy of such a body during the request phase.

@ahg-g ahg-g left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Change is correct; please land it after #3884 (inline), plus a helper for the pool-cap rule and one missing test case.

Comment thread pkg/epp/handlers/server.go Outdated
Comment on lines +680 to +685
if s.maxPoolBufferSize > 0 && buf.Cap() > s.maxPoolBufferSize {
// putBuffer drops a buffer above the pool cap, so RawBody takes its
// memory. Parsers and the Director replace RawBody on a rewrite and
// never write into it, so inboundBody stays the comparison baseline.
reqCtx.Request.RawBody = inboundBody
buf = nil

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Merge this after #3884, and say in the description that it depends on it.

Why: RawBody now holds the buffer's whole underlying array, not just the body. bytes.Buffer roughly doubles as it grows, so that array is 1.6 to 1.9 times the body size. A 20 MB body written in 64 KiB chunks ends up in a 33.5 MiB array. On main, RawBody is held until the stream ends, so without #3884 each long streamed response keeps 33.5 MiB alive instead of a 20 MB copy. Peak memory while the request is processed goes down, but memory held for the rest of the generation goes up.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed. The description now says this PR depends on #3884 and must merge after it. I will rebase after #3884 merges; both PRs add tests at the same place in server_request_buffer_test.go.

Comment thread pkg/epp/handlers/server.go Outdated
inboundBody := buf.Bytes()
reqCtx.Request.RawBody = make([]byte, buf.Len())
copy(reqCtx.Request.RawBody, inboundBody)
if s.maxPoolBufferSize > 0 && buf.Cap() > s.maxPoolBufferSize {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Move the pool-cap rule into one helper and use it in both putBuffer and this branch.

func (s *StreamingServer) poolKeeps(buf *bytes.Buffer) bool {
	return s.maxPoolBufferSize == 0 || buf.Cap() <= s.maxPoolBufferSize
}

Why: this branch is only safe because putBuffer never returns these buffers to the pool, and the rule is written out twice. If putBuffer changes, for example to keep larger buffers, the pool hands out a buffer that RawBody still points into. The next request overwrites it, and the earlier request forwards a corrupted body.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done. poolKeeps holds the rule, and putBuffer and the RawBody branch both call it.

return reqCtx, nil
}

func TestProcessRawBodyUsesBufferAbovePoolCap(t *testing.T) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add a case with tool-calling telemetry enabled and a buffer above the cap.

Why: that path now passes buf == nil to the new if buf != nil guard, while the snapshot and compareEPPToolCallingSnapshotToBody read the same memory RawBody uses. No test covers this combination, so a nil dereference or a broken comparison there would not be caught.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done. TestProcessToolCallingSnapshotWithBufferAbovePoolCap enables telemetry with a buffer above the cap. It checks that RawBody shares the buffer memory and that the comparison reports tools as preserved.

@ahg-g ahg-g self-assigned this Nov 3, 2026
@github-actions github-actions Bot added size/L Denotes a PR that changes 100-499 lines, ignoring generated files. kind/performance and removed size/M Denotes a PR that changes 30-99 lines, ignoring generated files. kind/performance labels Nov 3, 2026
ahg-g
ahg-g previously approved these changes Nov 4, 2026
@ahg-g
ahg-g added this pull request to the merge queue Nov 4, 2026
@ahg-g
ahg-g removed this pull request from the merge queue due to a manual request Nov 4, 2026
@ahg-g

ahg-g commented Nov 4, 2026

Copy link
Copy Markdown
Collaborator

@thomas-fahrner-parasail pls rebase, I think your other PR that got merged cause a conflict on this one

@github-actions github-actions Bot added the needs-rebase PR has merge conflicts and needs a rebase label Nov 4, 2026
@github-actions

github-actions Bot commented Nov 4, 2026

Copy link
Copy Markdown
Contributor

@thomas-fahrner-parasail This PR has merge conflicts with main. Please rebase on main and push to resolve them:

gh pr checkout 3883
git fetch origin main
git rebase origin/main
# resolve conflicts
git push --force-with-lease

A request buffer larger than the pool cap (the gRPC max receive size,
4 MiB by default) is dropped after the request phase, so copying it
into RawBody only doubles the memory a large body takes while it is
parsed and scheduled. RawBody takes the buffer's memory instead. Smaller
bodies keep the copy so their buffers return to the pool.

Signed-off-by: Thomas Fahrner <thomas.fahrner@parasail.io>
RawBody can take a buffer's memory only because putBuffer never returns
that buffer to the pool. poolKeeps now holds that rule for both, so a
change to which buffers the pool keeps cannot leave RawBody pointing into
a reused buffer.

A new test covers tool-calling telemetry with a buffer above the pool
cap: the snapshot and RawBody share memory, and the comparison reports
the tools field as preserved.

Signed-off-by: Thomas Fahrner <thomas.fahrner@parasail.io>
@thomas-fahrner-parasail

Copy link
Copy Markdown
Contributor Author

Rebased on main. The only conflict was in server_request_buffer_test.go: the #3884 tests and these tests were added at the same place, and I kept both. server.go has no changes. go test -race ./pkg/epp/handlers/ passes.

@github-actions github-actions Bot added re-approve PR was rebased after prior approval and needs re-approval and removed needs-rebase PR has merge conflicts and needs a rebase labels Nov 4, 2026
@github-actions

github-actions Bot commented Nov 4, 2026

Copy link
Copy Markdown
Contributor

@ahg-g @thomas-fahrner-parasail has rebased this PR and it is now conflict-free. Please take another look and re-approve when ready.

@ahg-g
ahg-g added this pull request to the merge queue Nov 5, 2026
Merged via the queue into llm-d:main with commit 9a9a8ad Nov 5, 2026
53 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area/epp kind/performance re-approve PR was rebased after prior approval and needs re-approval size/L Denotes a PR that changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants