refactor(sidecar): encapsulate code better in order to share protocol implementation between different connectors by kyanokashi · Pull Request #566 · llm-d/llm-d-inference-scheduler

kyanokashi · 2026-01-15T01:20:47Z

Updated the proxy server to utilize a new ProxyManager for handling decoder and prefill proxies.
Introduced runners for SGLang and a decode-first flows (covers SharedStorage & NIXL) to streamline request handling.
Refactored existing code to improve encapsulation and reduce duplication, enhancing maintainability.
Removed obsolete files related to previous connector implementations.

Signed-off-by: kyano kyanokashi2@gmail.com

- if cache_hit_threshold field is present in completion request, then we perform a decode first flow Signed-off-by: kyano <kyanokashi2@gmail.com>

Signed-off-by: kyanokashi <71283892+kyanokashi@users.noreply.github.com>

Signed-off-by: kyano <kyanokashi2@gmail.com>

…w decode first flow Signed-off-by: kyano <kyanokashi2@gmail.com>

- decrease verbosity for common log - add cache_hit_threshold attribute Signed-off-by: kyano <kyanokashi2@gmail.com>

…marshal decode response Signed-off-by: kyano <kyanokashi2@gmail.com>

Signed-off-by: kyano <kyanokashi2@gmail.com>

…regardless of cache condition Signed-off-by: kyano <kyanokashi2@gmail.com>

Signed-off-by: kyano <kyanokashi2@gmail.com>

…xecution Signed-off-by: kyano <kyanokashi2@gmail.com>

Signed-off-by: kyano <kyanokashi2@gmail.com>

kyanokashi · 2026-01-16T04:28:58Z

sure thing - thanks @vMaroon. @kyanokashi the link checker failure is unrelated to your changes and can be resolved by rebasing on the latest. The test failures are likely more relavent.

Addressed

kfirwolfson · 2026-01-18T15:03:05Z

Hi. I didn't review the code yet, just a thought on naming.
Sounds like we're going to have 2 connectors, one connector_sglang for SGLang and the other one replacing both nixl and shared-storage connectors.
How about naming the second one connector_vllm?
Thinking that we would want a general name, as it supports both prefill-first and decode-first flows and either nixl or other data-movement options (shared-storage or lmcache)

kyanokashi · 2026-01-20T17:18:29Z

Hi. I didn't review the code yet, just a thought on naming. Sounds like we're going to have 2 connectors, one connector_sglang for SGLang and the other one replacing both nixl and shared-storage connectors. How about naming the second one connector_vllm? Thinking that we would want a general name, as it supports both prefill-first and decode-first flows and either nixl or other data-movement options (shared-storage or lmcache)

I think I'll rename to runner_vllm that's because nixl or lmcache are the connectors. Will rename sglang to that as well

- renaming various components and packages - further encapsulation in submodules to avoid import cycles - improved encapsulation under well defined types Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

pkg/sidecar/proxy/allowlist.go

pkg/sidecar/proxy/manager/manager.go

pkg/sidecar/proxy/keys/keys.go

pkg/sidecar/proxy/http_errors/http_errors.go

pkg/sidecar/proxy/proxy.go

pkg/sidecar/proxy/runners/runner_nixlv2_test.go

pkg/sidecar/proxy/runners/runner_sglang.go

pkg/sidecar/proxy/runners/runner_test.go

Co-authored-by: Shmuel Kallner <kallner@il.ibm.com> Signed-off-by: kyanokashi <71283892+kyanokashi@users.noreply.github.com>

Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

…ilder Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

elevran · 2026-02-05T17:28:41Z

@kyanokashi could you please remove the pd-sidecar binary from the PR?

Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

Incorporates main changes (GIE update, models extractor plugin, dependency bumps, decode-first flow, lmcache→shared-storage rename) into the refactored sidecar architecture. Conflicts resolved by keeping refactor branch versions; superseded connector_shared_storage.go and duplicate isHTTPError() removed.

- Add FIXME comments for context.TODO() in allowlist pod watchers - Document dual-mutex acquisition pattern in rebuildAllowlist - Clarify Clone() shared state for data parallel workers Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

- Add PDProxyManager interface to decouple runners from concrete manager - Move manager.go to pd_proxy_manager.go - Add GetDecoderProxy() method to ProxyManager - Update go:generate directive to include PDProxyManager in mocks Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

- Update VLLMRunner and SGLangRunner to use PDProxyManager interface - Rename field from proxyManager to proxyHandler for clarity - Update profile handlers (chat_completions, data_parallel) to use interface - Update tests to use PDProxyManager mock Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

- Export BufferedResponseWriter and ResponseWriterWithBuffer - Export helper methods and constructor - Add documentation comments for all exported items Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

WriteHeader must be called before Write per http.ResponseWriter contract. Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

- Test prefill request modifications (stream=false, max_tokens, etc.) - Test decode request state restoration - Verify original request immutability - Cover both SharedStorage and NIXLV2 builders Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

- Test JSON error response formatting - Verify correct HTTP status codes (400, 502, 500) - Ensure error messages are propagated correctly - Rename http_errors.go to errors.go for consistency Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

- Test BufferedResponseWriter append and status code handling - Test ResponseWriterWithBuffer buffering and direct mode transition - Test SSE event delimiter counting for FirstChunkReady signal - Verify GetStatusCode and Buffered accessors Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

Generated by moq to include PDProxyManager interface mock. Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

elevran

@kyanokashi thank you for this work.
It is a significant refactoring, restructuring the sidecar proxy code to improve modularity, testability, and maintainability.
Left some suggestions and nits. Nothing major from my side - all can go into future/followup PRs.
Leaving lgtm and approval to @shmuelk (he has a changes-requested on the PR)

elevran · 2026-02-19T16:36:41Z

pkg/sidecar/proxy/runners/runner_sglang.go

+	sr.proxyHandler.GetDecoderProxy().ServeHTTP(w, decodeReq)
+}
+
+func cloneWithJSONBody(r *http.Request, body []byte) *http.Request {


nit: consider sharing with runner_vllm.go's implementation of cloneRequestWithBody?

elevran · 2026-02-19T16:40:55Z

pkg/sidecar/proxy/allowlist.go

 				Version:  "v1",
 				Resource: "pods",
-			}).Namespace(av.namespace).List(context.TODO(), options)
+			}).Namespace(av.namespace).List(context.TODO(), options) // FIXME: use stored context from Start() to respect shutdown


nit: can you please extract TODO/FIXME from here and other files to a cleanup issue for awareness instead of keeping in the code?

elevran · 2026-02-19T16:43:22Z

cmd/pd-sidecar/main.go

 	vLLMPort := flag.String("vllm-port", "8001", "the port vLLM is listening on")
 	vLLMDataParallelSize := flag.Int("data-parallel-size", 1, "the vLLM DATA-PARALLEL-SIZE value")
-	connector := flag.String("connector", proxy.ConnectorNIXLV2, "the P/D connector being used. Supported: "+strings.Join(supportedConnectors, ", "))
+	connectorStr := flag.String("connector", types.ConnectorNIXLV2.String(), "the P/D connector being used. Supported: "+types.AllConnectorStrings())


nit: could benefit from a specialized flag type (e.g., with validation and conversion to the type instead of "raw" string.

elevran · 2026-02-19T16:47:29Z

pkg/sidecar/proxy/runners/types/types.go

+	return strings.Join(strs, ", ")
+}
+
+//go:generate moq -stub -out mock/types.go  -pkg mock . ProtocolRunner PDProxyManager


Q: will generate run automatically on changes to this file or should the makefile add a dependency of mocks file on this file so it is run on change?

elevran · 2026-02-19T16:49:51Z

pkg/sidecar/proxy/runners/runner_sglang.go

+// SGLangRunner implements P/D disaggregation for SGLang backends.
+// Unlike VLLMRunner, SGLangRunner sends prefill and decode requests concurrently
+// and uses SGLang's built-in bootstrap coordination mechanism.
+type SGLangRunner struct {


nit: consider using builders for consistency with vllm runner?

elevran · 2026-02-19T16:58:05Z

pkg/sidecar/proxy/runners/runner_sglang.go

+	sglangBootstrapPort = 8998
+
+	// Override from environment variable if set
+	if portStr := os.Getenv("SGLANG_BOOTSTRAP_PORT"); portStr != "" {


nit: perhaps use a command line option for this? environment variables are less visible to users (eg. not in --help). You can still the envvar as the default value.

elevran · 2026-02-19T17:02:51Z

pkg/sidecar/proxy/runners/runner_vllm.go

+
+	r.Header.Add(keys.RequestHeaderRequestID, uuidStr)
+
+	// If "cache_hit_threshold" is present in the request, we try to decode first. The decode node must meet the cache hit threshold in order to execute.


nit: might also want to keep decode first enablement behind a feature/cli flag.

elevran · 2026-02-19T17:18:51Z

pkg/sidecar/proxy/data_parallel.go

nit: we shuold be removing the data_parallel implementation from the sidecar altogether (not in this PR)

elevran · 2026-02-19T17:21:47Z

pkg/sidecar/proxy/pd_proxy_manager.go

+	prefillerInsecureSkipVerify bool
+
+	decoderProxy     http.Handler                     // decoder proxy handler
+	prefillerProxies *lru.Cache[string, http.Handler] // cached prefiller proxy handlers


nit: consider adding note regarding Clone() sharing these as shallow copy (if I'm not mistaken there's a note on the server.Clone, but maybe highlight here)?

elevran · 2026-02-19T17:24:28Z

pkg/sidecar/proxy/proxy_helpers.go

 	"syscall"
 	"time"
+
+	httperrors "github.com/llm-d/llm-d-inference-scheduler/pkg/sidecar/proxy/http_errors"


nit: consider placement of all helper functions (e.g., from this file and http_errors?) in http package.

shmuelk · 2026-02-19T16:23:25Z

pkg/sidecar/proxy/runners/runner_nixlv2_test.go

@@ -1,5 +1,5 @@
 /*
-Copyright 2025 The llm-d Authors.
+Copyright 2026 The llm-d Authors.


Suggested change

Copyright 2026 The llm-d Authors.

Copyright 2025, 2026 The llm-d Authors.

shmuelk · 2026-02-19T16:25:55Z

pkg/sidecar/proxy/runners/runner_sglang.go

+	sglangBootstrapPort = 8998
+
+	// Override from environment variable if set
+	if portStr := os.Getenv("SGLANG_BOOTSTRAP_PORT"); portStr != "" {


This should come from a configuration arameter, just as the vLLM port does.

shmuelk · 2026-02-19T16:29:04Z

pkg/sidecar/proxy/runners/runner_sglang.go

+func (sr *SGLangRunner) sendSGLangConcurrentRequests(w http.ResponseWriter, r *http.Request, body []byte, prefillHost string, logger logr.Logger) {
+	// Create separate requests for prefill and decode
+	prefillReq := cloneWithJSONBody(r, body)
+	decodeReq := cloneWithJSONBody(r, body)


This needs the fix for issue #632

shmuelk · 2026-02-19T16:29:39Z

pkg/sidecar/proxy/runners/runner_test.go

@@ -1,5 +1,5 @@
 /*
-Copyright 2025 The llm-d Authors.
+Copyright 2026 The llm-d Authors.


Suggested change

Copyright 2026 The llm-d Authors.

Copyright 2025, 2026 The llm-d Authors.

shmuelk · 2026-02-19T16:32:21Z

pkg/sidecar/proxy/runners/runner_test.go

 }

-var connectors = []string{ConnectorSharedStorage, ConnectorNIXLV2}
+var connectorsArr = []types.Connector{types.ConnectorSharedStorage, types.ConnectorNIXLV2}


I don't think there is a need to call the variable connectorsArr

Suggested change

var connectorsArr = []types.Connector{types.ConnectorSharedStorage, types.ConnectorNIXLV2}

var connectors = []types.Connector{types.ConnectorSharedStorage, types.ConnectorNIXLV2}

shmuelk · 2026-02-19T16:32:35Z

pkg/sidecar/proxy/runners/runner_test.go


-	for _, connector := range connectors {
-		When(fmt.Sprintf("running with the %s connector", connector), func() {
+	for _, connector := range connectorsArr {


Suggested change

for _, connector := range connectorsArr {

for _, connector := range connectors {

shmuelk · 2026-02-19T17:29:10Z

pkg/sidecar/proxy/runners/runner_vllm.go

+		if vr.hasCacheThresholdFinishReason(response) {
+			return true
+		}
+	}


There are Go SSE parsing and writing packages. Any reason why one wasn't used?

There's this one, but it isn't a standard Go endorsed package. Is this a dependency we'd want to introduce? https://github.com/tmaxmax/go-sse

shmuelk · 2026-02-19T17:31:11Z

pkg/sidecar/proxy/runners/status_response_writer.go

@@ -1,5 +1,5 @@
 /*
-Copyright 2025 The llm-d Authors.
+Copyright 2026 The llm-d Authors.


Suggested change

Copyright 2026 The llm-d Authors.

Copyright 2025, 2026 The llm-d Authors.

shmuelk · 2026-02-19T17:33:36Z

pkg/sidecar/proxy/runners/status_response_writer_test.go

+	})
+
+	It("should signal FirstChunkReady after 2 SSE events", func() {
+		rw.Write([]byte("data: {\"choices\":[{\"finish_reason\":null}]}\n\n")) //nolint:errcheck


There are Go SSE packages which might make this code cleaner.

kyanokashi and others added 30 commits December 9, 2025 18:45

feat: implement decode first flow on lmcache connector

b228bb7

- if cache_hit_threshold field is present in completion request, then we perform a decode first flow Signed-off-by: kyano <kyanokashi2@gmail.com>

Merge branch 'main' into feat/sidecar/lmcache-connector/decode-first

a436b50

Signed-off-by: kyanokashi <71283892+kyanokashi@users.noreply.github.com>

fix: error handling

a6ae771

Signed-off-by: kyano <kyanokashi2@gmail.com>

Merge branch 'main' into feat/sidecar/lmcache-connector/decode-first

58388eb

chore: add back todo comment

04b7ffd

Signed-off-by: kyano <kyanokashi2@gmail.com>

refactor: reduce code complexity and duplication

ce74f50

Signed-off-by: kyano <kyanokashi2@gmail.com>

refactor: improve header copying

1de6035

Signed-off-by: kyano <kyanokashi2@gmail.com>

chore: add comment explaning the cache_hit_threshold field and the ne…

c0ac69e

…w decode first flow Signed-off-by: kyano <kyanokashi2@gmail.com>

refactor: enhance logging for cache hit threshold in decode flow

7ce5e19

- decrease verbosity for common log - add cache_hit_threshold attribute Signed-off-by: kyano <kyanokashi2@gmail.com>

refactor: improve error handling and observability when failing to un…

6430a02

…marshal decode response Signed-off-by: kyano <kyanokashi2@gmail.com>

chore: add deleted informational comments

4c15d95

Signed-off-by: kyano <kyanokashi2@gmail.com>

typo

cac084f

Signed-off-by: kyano <kyanokashi2@gmail.com>

refactor: make error logs more descriptive of the failure reason

91c7a06

Signed-off-by: kyano <kyanokashi2@gmail.com>

feat: add cache hit threshold to prefill request so prefill executes …

69d30b5

…regardless of cache condition Signed-off-by: kyano <kyanokashi2@gmail.com>

fix: typo

1ed1d89

Signed-off-by: kyano <kyanokashi2@gmail.com>

refactor: assign 0 cache_hit_threshold before final decode attempt

515b385

Signed-off-by: kyano <kyanokashi2@gmail.com>

chore: update comment according to feedback

4c8659e

Signed-off-by: kyano <kyanokashi2@gmail.com>

Merge branch 'main' into feat/sidecar/lmcache-connector/decode-first

878585b

chore: remove istio workaround

722b58e

Signed-off-by: kyano <kyanokashi2@gmail.com>

fix: set cache hit threshold to 0 in prefill request for consistent e…

487d333

…xecution Signed-off-by: kyano <kyanokashi2@gmail.com>

refactor: update the log

cb00b52

Signed-off-by: kyano <kyanokashi2@gmail.com>

feat: support online decoding

88739c6

Signed-off-by: kyano <kyanokashi2@gmail.com>

fix: preserve request body in lmcache connector

9fbb2d1

Signed-off-by: kyano <kyanokashi2@gmail.com>

Merge branch 'main' into feat/sidecar/lmcache-connector/decode-first

7b18827

fix: support sse format for streamed decode

460b9c8

Signed-off-by: kyano <kyanokashi2@gmail.com>

chore: add and improve log descriptions

a722510

Signed-off-by: kyano <kyanokashi2@gmail.com>

fix: typo

e2b3380

Signed-off-by: kyano <kyanokashi2@gmail.com>

Merge branch 'main' into feat/sidecar/lmcache-connector/decode-first

548e6c8

nit: undo capitalization

160fb72

Signed-off-by: kyano <kyanokashi2@gmail.com>

fix: typos

9a08dfb

Signed-off-by: kyano <kyanokashi2@gmail.com>

Merge branch 'main' into refactor/sidecar/abstract-connector

758fb78

elevran added this to the v0.6 milestone Jan 22, 2026

kyanokashi and others added 2 commits January 22, 2026 14:53

Merge branch 'main' into refactor/sidecar/abstract-connector

aa6270a

refactor: various refactorings

1e71083

- renaming various components and packages - further encapsulation in submodules to avoid import cycles - improved encapsulation under well defined types Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

shmuelk requested changes Jan 26, 2026

View reviewed changes

github-project-automation bot moved this to In review in llm-d-inference-scheduler Jan 26, 2026

kyanokashi and others added 5 commits January 26, 2026 14:15

Apply suggestions from code review

5fbab01

Co-authored-by: Shmuel Kallner <kallner@il.ibm.com> Signed-off-by: kyanokashi <71283892+kyanokashi@users.noreply.github.com>

refactor: rewrite function only used by tests as a test util

aef9df1

Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

refactor: rename default request builder to shared storage request bu…

36f9b2e

…ilder Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

chore: update copyright comments to year 2026

0b8944a

Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

Merge branch 'main' into refactor/sidecar/abstract-connector

050b1d7

kyanokashi added 11 commits February 5, 2026 14:06

chore: remove pd-sidecar binary

2490cf0

Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

refactor(sidecar): export response writer types for testing

c150ae9

- Export BufferedResponseWriter and ResponseWriterWithBuffer - Export helper methods and constructor - Add documentation comments for all exported items Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

fix(sidecar): correct Write/WriteHeader order in mock handler

fab28a1

WriteHeader must be called before Write per http.ResponseWriter contract. Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

chore(sidecar): regenerate mocks for PDProxyManager interface

3c81eef

Generated by moq to include PDProxyManager interface mock. Signed-off-by: kyanokashi <kyanokashi2@gmail.com>

kyanokashi force-pushed the refactor/sidecar/abstract-connector branch from c50aa4c to 3c81eef Compare February 9, 2026 18:18

Merge branch 'main' into refactor/sidecar/abstract-connector

c854728

elevran reviewed Feb 19, 2026

View reviewed changes

shmuelk reviewed Feb 19, 2026

View reviewed changes


		r.Header.Add(keys.RequestHeaderRequestID, uuidStr)

		// If "cache_hit_threshold" is present in the request, we try to decode first. The decode node must meet the cache hit threshold in order to execute.

	Copyright 2026 The llm-d Authors.
	Copyright 2025, 2026 The llm-d Authors.

	var connectorsArr = []types.Connector{types.ConnectorSharedStorage, types.ConnectorNIXLV2}
	var connectors = []types.Connector{types.ConnectorSharedStorage, types.ConnectorNIXLV2}

	for _, connector := range connectorsArr {
	for _, connector := range connectors {

Conversation

kyanokashi commented Jan 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kyanokashi commented Jan 16, 2026

Uh oh!

kfirwolfson commented Jan 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kyanokashi commented Jan 20, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

elevran commented Feb 5, 2026

Uh oh!

elevran left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

kyanokashi commented Jan 15, 2026 •

edited

Loading

kfirwolfson commented Jan 18, 2026 •

edited

Loading