feat(plugin-state): Allow customizing staleness threshold

cmd/epp/runner/runner.go

@@ -222,6 +222,9 @@ func (r *Runner) Run(ctx context.Context) error {
 		return err
 	}
 
+	// Apply the process-wide plugin-state staleness threshold before any plugin is instantiated.
+	fwkplugin.SetDefaultStalenessThreshold(opts.PluginStateStalenessThreshold)
+
 	// Print flag values, skipping deprecated metric flags configured via engineConfigs
 	flags := make(map[string]any)
 	pflag.VisitAll(func(f *pflag.Flag) {

pkg/epp/framework/interface/plugin/plugin_state.go

@@ -28,17 +28,32 @@ import (
 )
 
 const (
-	// stalenessThreshold defines the threshold for considering data as stale.
-	// if data of a request hasn't been read/write in the last "stalenessThreshold", it is considered as stale data
-	// and will be cleaned in the next cleanup cycle.
-	stalenessThreshold = time.Minute * 5
+	// DefaultStalenessThreshold is the built-in threshold for considering data as stale: if a
+	// request's data has not been read/written within this duration it is reaped in the next cleanup
+	// cycle. It applies when no process-wide (SetDefaultStalenessThreshold) override is set.
+	DefaultStalenessThreshold = time.Minute * 5
 	// cleanupInterval defines the periodic interval that the cleanup go routine uses to check for stale data.
 	cleanupInterval = time.Minute
 )
 
-// NewPluginState initializes a new PluginState and returns its pointer.
+// defaultStalenessThreshold is the process-wide default applied to PluginState instances.
+// It may be overridden once during startup via SetDefaultStalenessThreshold.
+var defaultStalenessThreshold = DefaultStalenessThreshold
+
+// SetDefaultStalenessThreshold overrides the process-wide default staleness threshold used by
+// PluginState instances that do not set an explicit WithStalenessThreshold. It is intended to be
+// called once at startup, before any plugin is instantiated. Non-positive values are ignored.
+func SetDefaultStalenessThreshold(d time.Duration) {
+	if d > 0 {
+		defaultStalenessThreshold = d
+	}
+}
+
+// NewPluginState initializes a new PluginState and returns its pointer. Each instance captures the
+// process-wide staleness threshold at construction; configure it via SetDefaultStalenessThreshold
+// before instantiating plugins.
 func NewPluginState(ctx context.Context) *PluginState {
-	pluginState := &PluginState{}
+	pluginState := &PluginState{stalenessThreshold: defaultStalenessThreshold}
 	go pluginState.cleanup(ctx)
 	return pluginState
 }
@@ -60,6 +75,8 @@ type PluginState struct {
 	storage sync.Map
 	// key: RequestID, value: time.Time
 	requestToLastAccessTime sync.Map
+	// stalenessThreshold is the inactivity duration after which a request's data is reaped.
+	stalenessThreshold time.Duration
 }
 
 // Read retrieves data with the given "key" in the context of "requestID" from PluginState.
@@ -167,14 +184,14 @@ func (s *PluginState) cleanup(ctx context.Context) {
 }
 
 // cleanStaleRequests iterates through all requests and removes those that haven't been
-// accessed for longer than stalenessThreshold. This operation is safe to run concurrently
+// accessed for longer than the staleness threshold. This operation is safe to run concurrently
 // with other operations on the PluginState.
 func (s *PluginState) cleanStaleRequests() {
 	s.requestToLastAccessTime.Range(func(k, v any) bool {
 		requestID := k.(string)
 		lastAccessTime := v.(time.Time)
-		if time.Since(lastAccessTime) > stalenessThreshold {
-			log.Log.V(logutil.DEBUG).Info("Cleaning up stale request from PluginState", "requestID", requestID, "lastAccessTime", lastAccessTime)
+		if time.Since(lastAccessTime) > s.stalenessThreshold {
+			log.Log.V(logutil.DEBUG).Info("Cleaning up stale request from PluginState", "requestID", requestID, "lastAccessTime", lastAccessTime, "stalenessThreshold", s.stalenessThreshold)
 			s.Delete(requestID) // cleanup stale requests (this is safe in sync.Map)
 		}
 		return true

pkg/epp/framework/interface/plugin/plugin_state_test.go

@@ -109,7 +109,7 @@ func TestPluginState_EvictionCallback(t *testing.T) {
 	data.evictedID = ""
 	data.evictedKey = ""
 	state.Write(requestID, key, data)
-	state.requestToLastAccessTime.Store(requestID, time.Now().Add(-2*stalenessThreshold))
+	state.requestToLastAccessTime.Store(requestID, time.Now().Add(-2*defaultStalenessThreshold))
 	state.cleanStaleRequests()
 	assert.Equal(t, requestID, data.evictedID)
 	assert.Equal(t, key, data.evictedKey)
@@ -128,7 +128,7 @@ func TestPluginState_Touch(t *testing.T) {
 	state.Write(requestID, key, data)
 
 	// Set last access time to near-stale
-	nearStale := time.Now().Add(-stalenessThreshold + time.Second*10)
+	nearStale := time.Now().Add(-defaultStalenessThreshold + time.Second*10)
 	state.requestToLastAccessTime.Store(requestID, nearStale)
 
 	// Touch it
@@ -220,7 +220,7 @@ func TestReadPluginStateKey(t *testing.T) {
 }
 
 // TestPluginState_Cleanup verifies the automatic cleanup of stale data.
-// It tests that data which hasn't been accessed for longer than stalenessThreshold
+// It tests that data which hasn't been accessed for longer than the staleness threshold
 // is properly removed from the storage.
 func TestPluginState_Cleanup(t *testing.T) {
 	ctx, cancel := context.WithCancel(logutil.NewTestLoggerIntoContext(context.Background()))
@@ -235,14 +235,42 @@ func TestPluginState_Cleanup(t *testing.T) {
 	state.Write(requestID, key, data)
 
 	// Manually set last access time to far in the past
-	state.requestToLastAccessTime.Store(requestID, time.Now().Add(-2*stalenessThreshold))
+	state.requestToLastAccessTime.Store(requestID, time.Now().Add(-2*defaultStalenessThreshold))
 	// Manually CleanUp
 	state.cleanStaleRequests()
 
 	_, err := state.Read(requestID, key)
 	assert.Equal(t, ErrNotFound, err)
 }
 
+// TestSetDefaultStalenessThreshold verifies that the process-wide default is applied to new
+// PluginState instances, that a configured value drives reaping, and that non-positive values are
+// ignored.
+func TestSetDefaultStalenessThreshold(t *testing.T) {
+	orig := defaultStalenessThreshold
+	t.Cleanup(func() { defaultStalenessThreshold = orig })
+
+	SetDefaultStalenessThreshold(30 * time.Second)
+	assert.Equal(t, 30*time.Second, defaultStalenessThreshold)
+
+	ctx, cancel := context.WithCancel(logutil.NewTestLoggerIntoContext(context.Background()))
+	t.Cleanup(cancel)
+	state := NewPluginState(ctx)
+	assert.Equal(t, 30*time.Second, state.stalenessThreshold, "new state should inherit the configured default")
+
+	// A one-minute-old request is stale under the 30s configured threshold (but would survive the 5m built-in default).
+	requestID := "req-configured-threshold"
+	key := StateKey("foo")
+	state.Write(requestID, key, &pluginTestData{value: "bar"})
+	state.requestToLastAccessTime.Store(requestID, time.Now().Add(-time.Minute))
+	state.cleanStaleRequests()
+	_, err := state.Read(requestID, key)
+	assert.Equal(t, ErrNotFound, err, "request older than the configured threshold should be reaped")
+
+	SetDefaultStalenessThreshold(0)
+	assert.Equal(t, 30*time.Second, defaultStalenessThreshold, "non-positive value must be ignored")
+}
+
 // TestPluginState_DeleteKey verifies that DeleteKey correctly removes only the specified key for a request.
 func TestPluginState_DeleteKey(t *testing.T) {
 	ctx, cancel := context.WithCancel(logutil.NewTestLoggerIntoContext(context.Background()))

pkg/epp/server/options.go

@@ -35,6 +35,7 @@ import (
 
 	"github.com/llm-d/llm-d-router/pkg/common/observability/logging"
 	"github.com/llm-d/llm-d-router/pkg/common/routing"
+	fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
 )
 
 const (
@@ -95,6 +96,10 @@ type Options struct {
 	LoRAInfoMetric                   string        // Prometheus metric specification for the LoRA info metrics.
 	CacheInfoMetric                  string        // Prometheus metric specification for the cache info metrics.
 	//
+	// Plugin state.
+	//
+	PluginStateStalenessThreshold time.Duration // Inactivity duration after which a request's plugin state is reaped.
+	//
 	// Diagnostics.
 	//
 	logging.LoggingOptions         // Logging configuration.
@@ -137,6 +142,7 @@ func NewOptions() *Options {
 		KVCacheUsagePercentageMetric:     "vllm:kv_cache_usage_perc",
 		LoRAInfoMetric:                   "vllm:lora_requests_info",
 		CacheInfoMetric:                  "vllm:cache_config_info",
+		PluginStateStalenessThreshold:    fwkplugin.DefaultStalenessThreshold,
 		LoggingOptions:                   *logging.NewOptions(),
 		Tracing:                          true,
 		MetricsPort:                      9090,
@@ -195,6 +201,10 @@ func (opts *Options) AddFlags(fs *pflag.FlagSet) {
 	fs.StringVar(&opts.CacheInfoMetric, "cache-info-metric", opts.CacheInfoMetric, "Prometheus metric for the cache info metrics.")
 	_ = fs.MarkDeprecated("cache-info-metric", "use engineConfigs in EndpointPickerConfig instead")
 
+	fs.DurationVar(&opts.PluginStateStalenessThreshold, "plugin-state-staleness-threshold", opts.PluginStateStalenessThreshold,
+		"Inactivity duration after which a request's plugin state (e.g. in-flight load accounting) is reaped. "+
+			"Set this above the maximum expected request duration to avoid releasing in-flight accounting for long-running requests.")
+
 	opts.LoggingOptions.AddFlags(fs) // Add logging flags.
 
 	fs.BoolVar(&opts.Tracing, "tracing", opts.Tracing, "Enables emitting traces.")
@@ -360,6 +370,10 @@ func (opts *Options) Validate() error {
 		return fmt.Errorf("grpc-max-send-msg-size must be non-negative, got %d", opts.GRPCMaxSendMsgSize)
 	}
 
+	if opts.PluginStateStalenessThreshold <= 0 {
+		return fmt.Errorf("plugin-state-staleness-threshold must be positive, got %v", opts.PluginStateStalenessThreshold)
+	}
+
 	// Validate deprecated metric flags are not explicitly set
 	for flagName := range deprecatedMetricFlags {
 		if f := opts.fs.Lookup(flagName); f != nil && f.Changed {