feat(plugin-state): Allow customizing staleness threshold

In decode heavy with non-streaming, it is fairly common that a request
may take over 5 minutes (default staleness threshold), which causes
plugin state to remove those from accounting, when used with inflight flow
control this results in prematurly releasing request or token budget
which results in EPP admitting more request than configured.

This PR allows us to customize this staleness threshold, it adds a new epp
command line argument `--plugin-state-staleness-threshold` defaulting to
5m.

Signed-off-by: mohammadkhan <mohammadkhan@digitalocean.com>

Author: asharkhan3101

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 {