From c812309518154064dad5803d68c38eb155dddce7 Mon Sep 17 00:00:00 2001
From: Jan Hamacek <8758068+hamacekh@users.noreply.github.com>
Date: Thu, 20 Aug 2026 14:26:56 +0200
Subject: [PATCH] [Bug] Derive cache service identity from an explicit flag,
 not EnableKVSync

InitWithOptions inferred the calling component from opts.EnableKVSync, so a
gateway running with prefix-cache KV event sync disabled was classified as
"metadata" and had enableGPUOptimizerTracing and enableModelGPUProfileCaching
forced off. AIBRIX_GPU_OPTIMIZER_TRACING_FLAG=true was therefore silently
discarded, initTraceCache never ran, and no aibrix:<model>_request_trace_<ts>
keys were written -- leaving the GPU optimizer with no workload signal.

EnableKVSync is a feature toggle, not an identity. Add an explicit IsGateway
field to InitOptions, set it in cmd/plugins, and extract the classification
into serviceIdentity() so it can be unit tested. The intended suppression for
metadata-service and controller-manager is preserved.

Note this also restores GPU profile caching on the gateway, which defaults to
true but was being disabled by the same block.

Signed-off-by: Jan Hamacek <8758068+hamacekh@users.noreply.github.com>
---
 cmd/plugins/main.go          |  1 +
 pkg/cache/cache_init.go      | 24 +++++++++++------
 pkg/cache/cache_init_test.go | 51 ++++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 8 deletions(-)

diff --git a/cmd/plugins/main.go b/cmd/plugins/main.go
index d593712aa..9d17b8228 100644
--- a/cmd/plugins/main.go
+++ b/cmd/plugins/main.go
@@ -155,6 +155,7 @@ func main() {
 	remoteTokenizerEnabled := utils.LoadEnvBool(constants.EnvPrefixCacheUseRemoteTokenizer, false)
 
 	cache.InitWithOptions(config, stopCh, cache.InitOptions{
+		IsGateway:           true,
 		EnableKVSync:        kvSyncEnabled && remoteTokenizerEnabled,
 		RedisClient:         redisClient,
 		ModelRouterProvider: routing.ModelRouterFactory,
diff --git a/pkg/cache/cache_init.go b/pkg/cache/cache_init.go
index b1d57d41a..2add9cb93 100644
--- a/pkg/cache/cache_init.go
+++ b/pkg/cache/cache_init.go
@@ -46,6 +46,9 @@ var (
 
 // InitOptions configures the cache initialization behavior
 type InitOptions struct {
+	// IsGateway marks the caller as the gateway-plugins service.
+	IsGateway bool
+
 	// EnableKVSync configures whether to start the ZMQ KV event sync
 	EnableKVSync bool
 
@@ -325,14 +328,7 @@ func InitForTest() *Store {
 func InitWithOptions(config *rest.Config, stopCh <-chan struct{}, opts InitOptions) *Store {
 	once.Do(func() {
 		// Log initialization based on configuration
-		var service string
-		if opts.EnableKVSync {
-			service = "gateway"
-		} else if opts.RedisClient != nil {
-			service = "metadata"
-		} else {
-			service = "controllers"
-		}
+		service := serviceIdentity(opts)
 
 		klog.InfoS("initialize cache",
 			"service", service,
@@ -394,6 +390,18 @@ func InitWithOptions(config *rest.Config, stopCh <-chan struct{}, opts InitOptio
 	return store
 }
 
+// serviceIdentity returns the component identity implied by opts.
+func serviceIdentity(opts InitOptions) string {
+	switch {
+	case opts.IsGateway:
+		return "gateway"
+	case opts.RedisClient != nil:
+		return "metadata"
+	default:
+		return "controllers"
+	}
+}
+
 // initMetricsCache initializes metrics cache update loop
 // Parameters:
 //
diff --git a/pkg/cache/cache_init_test.go b/pkg/cache/cache_init_test.go
index 6ffe73c6a..c74adbb23 100644
--- a/pkg/cache/cache_init_test.go
+++ b/pkg/cache/cache_init_test.go
@@ -378,3 +378,54 @@ func TestInitWithOptions_KVSyncBehavior(t *testing.T) {
 		})
 	}
 }
+
+func TestServiceIdentity(t *testing.T) {
+	scenarios := []struct {
+		name     string
+		opts     InitOptions
+		expected string
+	}{
+		{
+			name: "gateway without redis",
+			opts: InitOptions{
+				IsGateway: true,
+			},
+			expected: "gateway",
+		},
+		{
+			name: "gateway with KV sync disabled is still the gateway",
+			opts: InitOptions{
+				IsGateway:    true,
+				EnableKVSync: false,
+				RedisClient:  &redis.Client{},
+			},
+			expected: "gateway",
+		},
+		{
+			name: "metadata service",
+			opts: InitOptions{
+				RedisClient: &redis.Client{},
+			},
+			expected: "metadata",
+		},
+		{
+			name: "EnableKVSync alone does not imply gateway",
+			opts: InitOptions{
+				EnableKVSync: true,
+				RedisClient:  &redis.Client{},
+			},
+			expected: "metadata",
+		},
+		{
+			name:     "controllers",
+			opts:     InitOptions{},
+			expected: "controllers",
+		},
+	}
+
+	for _, sc := range scenarios {
+		t.Run(sc.name, func(t *testing.T) {
+			assert.Equal(t, sc.expected, serviceIdentity(sc.opts))
+		})
+	}
+}