Skip to content

refactor(datalayer): key AttributeMap by DataKey instead of string - #2190

Open
satyamg1620 wants to merge 6 commits into
llm-d:mainfrom
satyamg1620:data-prod-cons
Open

refactor(datalayer): key AttributeMap by DataKey instead of string#2190
satyamg1620 wants to merge 6 commits into
llm-d:mainfrom
satyamg1620:data-prod-cons

Conversation

@satyamg1620

@satyamg1620 satyamg1620 commented Jul 26, 2026

Copy link
Copy Markdown
Contributor

What type of PR is this?
/kind cleanup
/kind api-change

What this PR does / why we need it:

Plugins declare the data they exchange as DataKey values, but AttributeMap access took a raw string. Nothing tied a plugin's reads and writes to its Produces()/Consumes() declaration, so a plugin could declare one key and use another.

This closes that in two layers:

  • Typing. AttributeMap and scheduling.Endpoint are keyed by DataKey, which removes the raw-string path into the attribute map. It does not enforce the declaration on its own: plugin.NewDataKey is exported, so a plugin can still construct a key it never declared. Runtime confinement is what makes the declaration binding. Custom scalar metrics derive their key from configuration via attrmetrics.ResolveConfiguredKey.

  • Runtime confinement. datalayer.Scope wraps the endpoints handed to a plugin so writes outside Produces() are dropped and reads outside Consumes() resolve as absent. It covers filters, scorers, and DataProducer.Produce. Every rejection increments llm_d_epp_plugin_data_scope_violations_total, labelled by extension point, plugin type, plugin name, and access kind. A rejected write additionally fails the request where the extension point has an error return, which today means DataProducer.

Runtime confinement lives next to ValidateAndOrderDataDependencies, which enforces the same contract statically.

The endpoint-attribute-filter and endpoint-attribute-scorer plugins previously declared Consumes() map[string]any, which never satisfied ConsumerPlugin, so they participated in no dependency ordering at all. They now declare plugin.DataDependencies and name their producer through a new producer parameter.

Not covered

  • The per-request store (InferenceRequest.PutAttribute/GetAttribute) is typed by DataKey but not confined: p2psource and disagg/topologyaffinity currently exchange request attributes they do not declare.
  • Datalayer extractors write through Endpoint.GetAttributes() rather than a scoped endpoint.
  • Admitter, Screener, and profile handlers receive endpoints unscoped. This is not only theoretical: LatencyAdmission declares a Required dependency in Consumes() and reads it off the endpoint outside the scope, and the disagg PrefixBasedPDDecider reads PrefixCacheMatchInfoDataKey the same way. Admit has an error return, so it could be scoped on the same terms as DataProducer in a follow-up.

Which issue(s) this PR fixes:
Fixes #1223

Release note:

Plugin endpoint attribute access is now confined at runtime to the DataKeys the plugin declares in Produces()/Consumes(). Rejected accesses are counted by the new `llm_d_epp_plugin_data_scope_violations_total` metric; a rejected write also fails the request for DataProducer plugins.

The `endpoint-attribute-filter` and `endpoint-attribute-scorer` plugins take a new `producer` parameter naming the plugin that publishes the attribute. The `attribute`/`attributeKey` parameter is now the attribute name alone. Configs using the combined "Name/Producer" form are rejected at startup with a message naming the split, rather than silently resolving to a key that matches nothing. An attribute name that itself contains "/" stays reachable by setting `producer` explicitly, including to the empty string for producer-agnostic attributes such as the multicluster pool aggregates.

AttributeMap and scheduling.Endpoint are keyed by DataKey rather than string. Out-of-tree plugins implementing these interfaces need updating.

@github-actions github-actions Bot added size/L Denotes a PR that changes 100-499 lines, ignoring generated files. kind/cleanup and removed kind/cleanup labels Jul 26, 2026
@satyamg1620
satyamg1620 marked this pull request as ready for review August 8, 2026 20:01
@satyamg1620
satyamg1620 marked this pull request as draft August 8, 2026 20:13
@satyamg1620
satyamg1620 marked this pull request as ready for review August 9, 2026 19:55
@github-actions github-actions Bot added size/XL Denotes a PR that changes 500-999 lines, ignoring generated files. and removed size/L Denotes a PR that changes 100-499 lines, ignoring generated files. labels Aug 9, 2026

@elevran elevran left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The datascope package adds actual runtime enforcement that a plugin's endpoint reads/writes match its declared Produces()/Consumes()

Two things worth addressing before merge:

  1. The PR description says enforcement against a plugin's own declaration is "not covered... worth a separate issue", but commits 2 and 4 (enforce(requestcontrol), enforce(framework)) add exactly that for endpoint attributes. Please update the description so it matches what's shipped.
  2. InferenceRequest.PutAttribute/GetAttribute (the request-side attribute store) got the DataKey type change but is not wrapped by any datascope-equivalent enforcement, unlike the endpoint side. If that asymmetry is intentional (e.g. deferred to a follow-up), a note in the description or a linked issue would help; if not, it is a real gap in what this PR claims to fix.

Comment thread pkg/epp/scheduling/scheduler_profile.go Outdated
// The write violation is dropped: Filter has no error return, and Scope
// has already logged and rejected the write. Producers run under
// executePluginsAsDAG, which does fail the request on one.
scoped, _ := datascope.Scope(logger, filter, filteredEndpoints)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit:
a filter's or scorer's undeclared write only surfaces as a dropped log line here and in runScorer below, since Filter/Score have no error return. Worth confirming there is a metric or otherwise-visible signal so a misbehaving plugin does not fail silently in production.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added. Rejections now increment llm_d_epp_plugin_data_scope_violations_total, labelled by extension point, plugin type, plugin name, and access kind (read/write), so a misdeclared filter or scorer is visible in production rather than only at log verbosity.

Logging is deduplicated per invocation: the first offence of each kind logs at Error, the rest at DEBUG, so a plugin that misreaches on every candidate does not emit one error line per endpoint per request. The metric counts every rejection.

@elevran elevran self-assigned this Aug 11, 2026
Comment thread pkg/epp/framework/datascope/endpoint.go Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. datascope is a plausible name but it's vague: "scope" could mean many things in Go code (variable scope, request scope, tracing scope). Better options, given what it actually does (wraps an Endpoint and enforces the plugin's own Produces()/Consumes() declaration) might be confine.
    The use if confine matches the doc comment's own language ("confines a plugin's attribute access"), is short, verb-based, and reads naturally as confine.Scope(...). It doesn't collide with other Go-ecosystem naming conventions the way "scope" does or acl would. This is a naming preference, not a blocking issue.
  2. Instead of new package, consider pkg/epp/datalayer/endpoint_scope.go (or similar) alongside data_graph.go instead of a new framework/* package. The new code here and existing code in datalayer enforce the same Produces()/Consumes() contract, one statically and one at runtime, and no import cycle blocks it.Using framework/datascope (or framework/confine) as its own package is defensible too (keeps runtime enforcement separate from static DAG validation), but I would lean towards reusing the datalayer package..

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Took the placement suggestion, which moots the naming one. This file is gone — it's now pkg/epp/datalayer/endpoint_scope.go, in package datalayer beside data_graph.go, so the static and runtime halves sit together. Type is ScopedEndpoint; call sites read datalayer.Scope(...).

Confirmed no cycle: pkg/epp/datalayer already depends on the three framework/interface/* packages and on pkg/epp/metrics, with nothing importing back.

@elevran

elevran commented Aug 11, 2026

Copy link
Copy Markdown
Collaborator

@satyamg1620 please address conflict

@satyamg1620

Copy link
Copy Markdown
Contributor Author

The datascope package adds actual runtime enforcement that a plugin's endpoint reads/writes match its declared Produces()/Consumes()

Two things worth addressing before merge:

  1. The PR description says enforcement against a plugin's own declaration is "not covered... worth a separate issue", but commits 2 and 4 (enforce(requestcontrol), enforce(framework)) add exactly that for endpoint attributes. Please update the description so it matches what's shipped.
  2. InferenceRequest.PutAttribute/GetAttribute (the request-side attribute store) got the DataKey type change but is not wrapped by any datascope-equivalent enforcement, unlike the endpoint side. If that asymmetry is intentional (e.g. deferred to a follow-up), a note in the description or a linked issue would help; if not, it is a real gap in what this PR claims to fix.

Thanks @elevran for the review. I will update the PR description and code accordingly

Signed-off-by: satyamg1620 <Satyam.Gupta.3@ibm.com>
Signed-off-by: satyamg1620 <Satyam.Gupta.3@ibm.com>
Signed-off-by: satyamg1620 <Satyam.Gupta.3@ibm.com>
Signed-off-by: satyamg1620 <Satyam.Gupta.3@ibm.com>
Move runtime confinement alongside the static dependency validation it
mirrors, so both halves of the Produces()/Consumes() contract live in one
package.

A rejected access surfaced only as a log line, and a rejected read logged
one per endpoint per request. Count every rejection under
llm_d_epp_plugin_data_scope_violations_total and log the first offence of
each kind per invocation, so a misdeclared plugin stays visible in
production without depending on log verbosity.

Signed-off-by: satyamg1620 <Satyam.Gupta.3@ibm.com>
@satyamg1620

Copy link
Copy Markdown
Contributor Author

The datascope package adds actual runtime enforcement that a plugin's endpoint reads/writes match its declared Produces()/Consumes()

Two things worth addressing before merge:

  1. The PR description says enforcement against a plugin's own declaration is "not covered... worth a separate issue", but commits 2 and 4 (enforce(requestcontrol), enforce(framework)) add exactly that for endpoint attributes. Please update the description so it matches what's shipped.
  2. InferenceRequest.PutAttribute/GetAttribute (the request-side attribute store) got the DataKey type change but is not wrapped by any datascope-equivalent enforcement, unlike the endpoint side. If that asymmetry is intentional (e.g. deferred to a follow-up), a note in the description or a linked issue would help; if not, it is a real gap in what this PR claims to fix.
  1. PR Description updated. It now describes the runtime confinement as shipped rather than deferred.

  2. The per-request store asymmetry is intentional and is now called out under "Not covered", with the reason: p2psource and disagg/topologyaffinity exchange request attributes they do not declare, so confining that store today would fail those paths. I can file a follow-up issue to track it if you would prefer

@satyamg1620

Copy link
Copy Markdown
Contributor Author

@satyamg1620 please address conflict

Resolved and rebased PR

@satyamg1620
satyamg1620 requested a review from elevran August 11, 2026 16:38
The attribute name and its producer are separate parameters, but they were
once a single "Attribute/Producer" string. A key built from that spelling
matches nothing: the read resolves as absent, so the filter keeps every
endpoint and the scorer returns zero, with no error anywhere.

Reject it at construction and name the split in the message. A producer
pointer tells an omitted producer from one set to the empty string, so an
attribute whose own name contains a slash stays reachable.

Signed-off-by: satyamg1620 <Satyam.Gupta.3@ibm.com>
@github-actions github-actions Bot added size/XXL Denotes a PR that changes 1000+ lines, ignoring generated files. and removed size/XL Denotes a PR that changes 500-999 lines, ignoring generated files. labels Aug 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

kind/cleanup size/XXL Denotes a PR that changes 1000+ lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Data producer/consumer behavior should be enforced with their declaration

2 participants