Skip to content

Commit 185bb7a

Browse files
authored
perf(epp): derive Scope allowed-key sets once at datalayer init (#2471)
* perf(epp): cache Scope allowed-key sets per plugin instance Scope rebuilt the allowedPut/allowedGet maps from Produces()/Consumes() on every invocation, though the declarations are fixed at plugin construction and Scope runs for every filter, scorer, and DataProducer on every request. Derive the sets once per plugin instance and share them read-only across invocations. BenchmarkScope: -45% time / -40% allocs at 10 endpoints, -6% / -50% at 100. Confinement behavior is unchanged. Signed-off-by: Luke Van Drie <lukevandrie@google.com> * perf(epp): derive Scope allowed-key sets at startup registration Scope looks up a spec registered by RegisterScopeSpecs, keyed by the plugin's typed name, instead of caching lazily by plugin identity. The runner registers every plugin once the full set, including auto-created producers, is known. An unregistered plugin is confined to nothing and the miss is logged once per typed name. Signed-off-by: Luke Van Drie <lukevandrie@google.com> --------- Signed-off-by: Luke Van Drie <lukevandrie@google.com>
1 parent bc5d0be commit 185bb7a

10 files changed

Lines changed: 163 additions & 21 deletions

File tree

cmd/epp/runner/runner.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -823,6 +823,11 @@ func (r *Runner) parseConfigurationPhaseTwo(ctx context.Context, rawConfig *conf
823823
// The plugins will be executed in topologically sorted order to ensure that data is produced before it is consumed.
824824
r.requestControlConfig.OrderDataProducerPlugins(dag)
825825

826+
// Derive the endpoint-scope allowed-key sets while the full plugin set,
827+
// including auto-created producers, is known. A plugin missing here is
828+
// confined to nothing at request time.
829+
datalayer.RegisterScopeSpecs(handle.GetAllPlugins())
830+
826831
r.parserRegistry = cfg.ParserRegistry
827832
logger.Info("loaded configuration from file/text successfully")
828833

pkg/epp/datalayer/endpoint_scope.go

Lines changed: 84 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -189,47 +189,113 @@ func (e *ScopedEndpoint) Clone() fwkdl.AttributeMap {
189189
return clone
190190
}
191191

192-
// Scope confines endpoints to what plugin declares. extensionPoint labels the
193-
// violation counter, so a misdeclared plugin is attributable to the point that
194-
// ran it. The returned Violations accumulates what the plugin did outside its
195-
// declarations, for callers whose extension point can report it.
196-
//
197-
// A plugin that declares nothing reaches nothing: an absent declaration is a
198-
// statement that the plugin exchanges no data, not a request to be exempt.
199-
func Scope(logger logr.Logger, extensionPoint string, plugin fwkplugin.Plugin, endpoints []fwksched.Endpoint) ([]fwksched.Endpoint, *Violations) {
192+
// scopeSpec holds the allowed-key sets derived from a plugin's declarations.
193+
// The maps are shared read-only by every ScopedEndpoint wrapper across all
194+
// invocations.
195+
type scopeSpec struct {
196+
allowedPut map[fwkplugin.DataKey]struct{}
197+
allowedGet map[fwkplugin.DataKey]struct{}
198+
}
199+
200+
// denyAllSpec confines an unregistered plugin the same way as one that
201+
// declares nothing.
202+
var denyAllSpec = &scopeSpec{}
203+
204+
// scopeSpecs maps a plugin's TypedName().String() to the spec derived from its
205+
// declarations. RegisterScopeSpecs writes it during startup; Scope only reads.
206+
// ValidateAndOrderDataDependencies keys plugins the same way, so within the
207+
// datalayer contract the typed name identifies the plugin.
208+
var (
209+
scopeSpecsMu sync.RWMutex
210+
scopeSpecs = map[string]*scopeSpec{}
211+
212+
// unregisteredReported dedups the error log for plugins missing from
213+
// scopeSpecs, which Scope would otherwise emit on every invocation.
214+
unregisteredReported sync.Map // string -> struct{}
215+
)
216+
217+
// RegisterScopeSpecs derives the allowed-key sets from each plugin's
218+
// Produces() and Consumes() declarations and stores them for Scope to look up.
219+
// Declarations are fixed at plugin construction, so the sets are derived once
220+
// here rather than on every Scope invocation. Call it after all plugins are
221+
// instantiated; registering a typed name again replaces its spec.
222+
func RegisterScopeSpecs(plugins []fwkplugin.Plugin) {
223+
scopeSpecsMu.Lock()
224+
defer scopeSpecsMu.Unlock()
225+
for _, plugin := range plugins {
226+
scopeSpecs[plugin.TypedName().String()] = buildScopeSpec(plugin)
227+
}
228+
}
229+
230+
func buildScopeSpec(plugin fwkplugin.Plugin) *scopeSpec {
200231
produces := map[fwkplugin.DataKey]any{}
201232
if producer, ok := plugin.(fwkplugin.ProducerPlugin); ok {
202233
produces = producer.Produces()
203234
}
204-
205-
allowedPut := make(map[fwkplugin.DataKey]struct{}, len(produces))
206-
allowedGet := make(map[fwkplugin.DataKey]struct{}, len(produces))
235+
spec := &scopeSpec{
236+
allowedPut: make(map[fwkplugin.DataKey]struct{}, len(produces)),
237+
allowedGet: make(map[fwkplugin.DataKey]struct{}, len(produces)),
238+
}
207239
for key := range produces {
208-
allowedPut[key] = struct{}{}
240+
spec.allowedPut[key] = struct{}{}
209241
// A producer may read back its own output.
210-
allowedGet[key] = struct{}{}
242+
spec.allowedGet[key] = struct{}{}
211243
}
212244
if consumer, ok := plugin.(fwkplugin.ConsumerPlugin); ok {
213245
deps := consumer.Consumes()
214246
for key := range deps.Required {
215-
allowedGet[key] = struct{}{}
247+
spec.allowedGet[key] = struct{}{}
216248
}
217249
for key := range deps.Optional {
218-
allowedGet[key] = struct{}{}
250+
spec.allowedGet[key] = struct{}{}
219251
}
220252
}
253+
return spec
254+
}
221255

222-
violations := &Violations{}
256+
// scopeSpecFor returns the registered spec for a plugin's typed name, or the
257+
// deny-all spec when none was registered. The miss is logged once per typed
258+
// name: it indicates a wiring bug, and the resulting confinement also shows up
259+
// through the violation counter as soon as the plugin touches an attribute.
260+
func scopeSpecFor(logger logr.Logger, name string) *scopeSpec {
261+
scopeSpecsMu.RLock()
262+
spec, ok := scopeSpecs[name]
263+
scopeSpecsMu.RUnlock()
264+
if ok {
265+
return spec
266+
}
267+
if _, reported := unregisteredReported.LoadOrStore(name, struct{}{}); !reported {
268+
logger.Error(fmt.Errorf("plugin %q has no registered scope spec; pass it to RegisterScopeSpecs", name),
269+
"Confining an unregistered plugin to nothing")
270+
}
271+
return denyAllSpec
272+
}
273+
274+
// Scope confines endpoints to what plugin declares. extensionPoint labels the
275+
// violation counter, so a misdeclared plugin is attributable to the point that
276+
// ran it. The returned Violations accumulates what the plugin did outside its
277+
// declarations, for callers whose extension point can report it.
278+
//
279+
// A plugin that declares nothing reaches nothing: an absent declaration is a
280+
// statement that the plugin exchanges no data, not a request to be exempt.
281+
//
282+
// The allowed-key sets come from the registry populated by RegisterScopeSpecs
283+
// at startup, treating Produces() and Consumes() as immutable after
284+
// construction. A plugin that was never registered is confined to nothing.
285+
func Scope(logger logr.Logger, extensionPoint string, plugin fwkplugin.Plugin, endpoints []fwksched.Endpoint) ([]fwksched.Endpoint, *Violations) {
223286
typedName := plugin.TypedName()
287+
spec := scopeSpecFor(logger, typedName.String())
288+
289+
violations := &Violations{}
224290
// One backing array rather than an allocation per endpoint: this runs for
225291
// every filter and scorer on every request, over the whole candidate set.
226292
wrappers := make([]ScopedEndpoint, len(endpoints))
227293
scoped := make([]fwksched.Endpoint, len(endpoints))
228294
for i, endpoint := range endpoints {
229295
wrappers[i] = ScopedEndpoint{
230296
inner: endpoint,
231-
allowedPut: allowedPut,
232-
allowedGet: allowedGet,
297+
allowedPut: spec.allowedPut,
298+
allowedGet: spec.allowedGet,
233299
typedName: typedName,
234300
extensionPoint: extensionPoint,
235301
logger: logger,

pkg/epp/datalayer/endpoint_scope_bench_test.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ func benchPlugin() fwkplugin.Plugin {
4545
p := &producerConsumerPlugin{}
4646
p.produces = map[fwkplugin.DataKey]any{producedKey: nil}
4747
p.consumes = &fwkplugin.DataDependencies{Optional: map[fwkplugin.DataKey]any{consumedKey: nil}}
48+
RegisterScopeSpecs([]fwkplugin.Plugin{p})
4849
return p
4950
}
5051

pkg/epp/datalayer/endpoint_scope_test.go

Lines changed: 55 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,14 +42,21 @@ var (
4242
)
4343

4444
// testPlugin implements Plugin plus whichever of Produces/Consumes the test
45-
// populates, mirroring how real plugins opt into each role.
45+
// populates, mirroring how real plugins opt into each role. The scope registry
46+
// is keyed by typed name, so tests that must not see another test's spec set a
47+
// distinct name.
4648
type testPlugin struct {
49+
name string
4750
produces map[fwkplugin.DataKey]any
4851
consumes *fwkplugin.DataDependencies
4952
}
5053

5154
func (p *testPlugin) TypedName() fwkplugin.TypedName {
52-
return fwkplugin.TypedName{Type: "testPlugin", Name: "mock"}
55+
name := p.name
56+
if name == "" {
57+
name = "mock"
58+
}
59+
return fwkplugin.TypedName{Type: "testPlugin", Name: name}
5360
}
5461

5562
type producerPlugin struct{ testPlugin }
@@ -79,6 +86,7 @@ func scopeProducerConsumer(t *testing.T, endpoint fwksched.Endpoint) (fwksched.E
7986
plug := &producerConsumerPlugin{}
8087
plug.produces = map[fwkplugin.DataKey]any{producedKey: nil}
8188
plug.consumes = &fwkplugin.DataDependencies{Optional: map[fwkplugin.DataKey]any{consumedKey: nil}}
89+
RegisterScopeSpecs([]fwkplugin.Plugin{plug})
8290
scoped, violation := Scope(testLogger(), "test-extension-point", plug, []fwksched.Endpoint{endpoint})
8391
require.Len(t, scoped, 1)
8492
return scoped[0], violation
@@ -115,6 +123,7 @@ func TestScope_ViolationsAreSharedAndReportTheFirstWrite(t *testing.T) {
115123

116124
plug := &producerPlugin{}
117125
plug.produces = map[fwkplugin.DataKey]any{producedKey: nil}
126+
RegisterScopeSpecs([]fwkplugin.Plugin{plug})
118127
scoped, violations := Scope(testLogger(), "test-extension-point", plug, endpoints)
119128

120129
first := fwkplugin.NewDataKey("first-undeclared", "otherPlugin")
@@ -162,7 +171,9 @@ func TestScope_GetHonoursConsumesAndProduces(t *testing.T) {
162171
// the case that previously fell through to unrestricted reads.
163172
func TestScope_PluginDeclaringNothingReachesNothing(t *testing.T) {
164173
endpoint := newEndpoint(t)
165-
scoped, violation := Scope(testLogger(), "test-extension-point", &testPlugin{}, []fwksched.Endpoint{endpoint})
174+
plug := &testPlugin{}
175+
RegisterScopeSpecs([]fwkplugin.Plugin{plug})
176+
scoped, violation := Scope(testLogger(), "test-extension-point", plug, []fwksched.Endpoint{endpoint})
166177
require.Len(t, scoped, 1)
167178

168179
_, ok := scoped[0].Get(consumedKey)
@@ -180,6 +191,7 @@ func TestScope_ProducerWithoutConsumesReadsOnlyItsOwnOutput(t *testing.T) {
180191
endpoint := newEndpoint(t)
181192
plug := &producerPlugin{}
182193
plug.produces = map[fwkplugin.DataKey]any{producedKey: nil}
194+
RegisterScopeSpecs([]fwkplugin.Plugin{plug})
183195

184196
scoped, _ := Scope(testLogger(), "test-extension-point", plug, []fwksched.Endpoint{endpoint})
185197

@@ -247,3 +259,43 @@ func TestUnscope_LeavesUnwrappedEndpointsAlone(t *testing.T) {
247259
endpoints := []fwksched.Endpoint{newEndpoint(t)}
248260
assert.Equal(t, endpoints, Unscope(endpoints))
249261
}
262+
263+
// Declarations grant nothing on their own: the allowed sets come from the
264+
// registry, and a plugin nobody registered is confined like one that declares
265+
// nothing.
266+
func TestScope_UnregisteredPluginIsConfinedToNothing(t *testing.T) {
267+
plug := &producerConsumerPlugin{}
268+
plug.name = "never-registered"
269+
plug.produces = map[fwkplugin.DataKey]any{producedKey: nil}
270+
plug.consumes = &fwkplugin.DataDependencies{Optional: map[fwkplugin.DataKey]any{consumedKey: nil}}
271+
272+
endpoint := newEndpoint(t)
273+
scoped, violations := Scope(testLogger(), "test-extension-point", plug, []fwksched.Endpoint{endpoint})
274+
require.Len(t, scoped, 1)
275+
276+
_, ok := scoped[0].Get(consumedKey)
277+
assert.False(t, ok, "a declared read must not resolve without registration")
278+
scoped[0].Put(producedKey, cloneableStr("nope"))
279+
_, ok = endpoint.Get(producedKey)
280+
assert.False(t, ok, "a declared write must not reach the endpoint without registration")
281+
assert.Error(t, violations.Write())
282+
}
283+
284+
// The registry is keyed by typed name; registering a name again replaces its
285+
// spec rather than accumulating.
286+
func TestRegisterScopeSpecs_ReregistrationReplacesTheSpec(t *testing.T) {
287+
declaring := &producerConsumerPlugin{}
288+
declaring.name = "replaced"
289+
declaring.produces = map[fwkplugin.DataKey]any{producedKey: nil}
290+
declaring.consumes = &fwkplugin.DataDependencies{Optional: map[fwkplugin.DataKey]any{consumedKey: nil}}
291+
RegisterScopeSpecs([]fwkplugin.Plugin{declaring})
292+
293+
silent := &testPlugin{name: "replaced"}
294+
RegisterScopeSpecs([]fwkplugin.Plugin{silent})
295+
296+
endpoint := newEndpoint(t)
297+
scoped, _ := Scope(testLogger(), "test-extension-point", declaring, []fwksched.Endpoint{endpoint})
298+
require.Len(t, scoped, 1)
299+
_, ok := scoped[0].Get(consumedKey)
300+
assert.False(t, ok, "the registry must serve the last spec registered for a typed name")
301+
}

pkg/epp/framework/plugins/scheduling/profilehandler/disagg/scheduler_test.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@ import (
1212
k8stypes "k8s.io/apimachinery/pkg/types"
1313
"sigs.k8s.io/controller-runtime/pkg/log" // Import config for thresholds
1414

15+
"github.com/llm-d/llm-d-router/pkg/epp/datalayer"
1516
fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
17+
fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
1618
fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling"
1719
fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
1820
attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix"
@@ -224,6 +226,7 @@ func TestPDSchedule(t *testing.T) {
224226
// initialize scheduler with config
225227
prefixScorer, err := prefix.New(ctx, prefix.PrefixCacheScorerPluginType, "")
226228
assert.NoError(t, err, "Prefix plugin creation returned unexpected error")
229+
datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{prefixScorer})
227230

228231
prefillSchedulerProfile := scheduling.NewSchedulerProfile().
229232
WithFilters(bylabel.NewPrefillRole()).

pkg/epp/requestcontrol/director_test.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1111,6 +1111,7 @@ func TestDirector_HandleRequest(t *testing.T) {
11111111
}
11121112
config := NewConfig()
11131113
if test.dataProducerPlugin != nil {
1114+
datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{test.dataProducerPlugin})
11141115
config = config.WithDataProducerPlugins(test.dataProducerPlugin)
11151116
}
11161117
if test.screener != nil {

pkg/epp/requestcontrol/plugin_executor_test.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import (
2626
"github.com/stretchr/testify/assert"
2727
"github.com/stretchr/testify/require"
2828

29+
"github.com/llm-d/llm-d-router/pkg/epp/datalayer"
2930
fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
3031
fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
3132
fwkrc "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requestcontrol"
@@ -428,6 +429,7 @@ func TestExecutePluginsAsDAG_EnforcesProducesDeclaration(t *testing.T) {
428429
t.Run("declared write reaches the endpoint", func(t *testing.T) {
429430
endpoint := newEndpoint()
430431
producer := &writingProducer{name: "p", declares: map[fwkplugin.DataKey]any{declared: nil}, writes: declared}
432+
datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{producer})
431433

432434
err := executePluginsAsDAG(context.Background(), []fwkrc.DataProducer{producer},
433435
&fwksched.InferenceRequest{}, []fwksched.Endpoint{endpoint})
@@ -440,6 +442,7 @@ func TestExecutePluginsAsDAG_EnforcesProducesDeclaration(t *testing.T) {
440442
t.Run("undeclared write fails the producer and leaves the endpoint untouched", func(t *testing.T) {
441443
endpoint := newEndpoint()
442444
producer := &writingProducer{name: "p", declares: map[fwkplugin.DataKey]any{declared: nil}, writes: undeclared}
445+
datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{producer})
443446

444447
err := executePluginsAsDAG(context.Background(), []fwkrc.DataProducer{producer},
445448
&fwksched.InferenceRequest{}, []fwksched.Endpoint{endpoint})

pkg/epp/scheduling/scheduler_bench_test.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,9 @@ import (
2525
"github.com/google/uuid"
2626
k8stypes "k8s.io/apimachinery/pkg/types"
2727

28+
"github.com/llm-d/llm-d-router/pkg/epp/datalayer"
2829
fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
30+
fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
2931
fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling"
3032
fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
3133
"github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/scheduling/picker"
@@ -64,6 +66,9 @@ func BenchmarkSchedule(b *testing.B) {
6466
b.Fatalf("prefix scorer setup: %v", err)
6567
}
6668
loraAffinityScorer := loraaffinity.NewLoraAffinityScorer()
69+
datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{
70+
kvCacheUtilizationScorer, queueingScorer, prefixCacheScorer, loraAffinityScorer,
71+
})
6772

6873
profile := NewSchedulerProfile().
6974
WithScorers(

pkg/epp/scheduling/scheduler_profile_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import (
3131
tracenoop "go.opentelemetry.io/otel/trace/noop"
3232
k8stypes "k8s.io/apimachinery/pkg/types"
3333

34+
"github.com/llm-d/llm-d-router/pkg/epp/datalayer"
3435
fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
3536
fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
3637
fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
@@ -976,6 +977,7 @@ func TestRunScorer_ScopesTheWrappedPluginDeclarations(t *testing.T) {
976977
&fwkdl.Metrics{}, attrs)
977978

978979
scorer := &declaringScorer{key: key, reads: map[string]bool{}}
980+
datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{scorer})
979981
scores := runScorer(context.Background(), nil, false,
980982
NewWeightedScorer(scorer, 1), &fwksched.InferenceRequest{}, []fwksched.Endpoint{endpoint})
981983

pkg/epp/scheduling/scheduler_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ import (
2828
k8stypes "k8s.io/apimachinery/pkg/types"
2929

3030
errcommon "github.com/llm-d/llm-d-router/pkg/common/error"
31+
"github.com/llm-d/llm-d-router/pkg/epp/datalayer"
3132
fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
3233
fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
3334
fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
@@ -47,6 +48,9 @@ func TestSchedule(t *testing.T) {
4748
prefixCacheScorer, err := schedprefix.New(context.Background(), schedprefix.PrefixCacheScorerPluginType, "approx-prefix-cache-producer")
4849
assert.NoError(t, err)
4950
loraAffinityScorer := loraaffinity.NewLoraAffinityScorer()
51+
datalayer.RegisterScopeSpecs([]fwkplugin.Plugin{
52+
kvCacheUtilizationScorer, queueingScorer, prefixCacheScorer, loraAffinityScorer,
53+
})
5054

5155
defaultProfile := NewSchedulerProfile().
5256
WithScorers(NewWeightedScorer(kvCacheUtilizationScorer, 1),

0 commit comments

Comments
 (0)