Add DRA support to the topology-aware policy

[bart0sh]

This PR adds optional DRA integration to the topology-aware policy, letting workloads request high-priority PCT (Performance/Turbo Class) CPUs via a Kubernetes ResourceClaim instead of (or alongside) the existing NRI-based allocation path.

What's included

• A reusable Kubernetes client wrapper (pkg/kubernetes/client) and a generic DRA kubelet-plugin skeleton (pkg/resmgr/dra), independent of any specific policy.
• cpuClass.dra.publish config field, plus ValidateCPUClassesForDRA guarding against overcommitting a turbo-frequency tier when multiple classes are published.
  pct.PickHpCpus/ReleaseHpCpus/Punits, giving the DRA driver a way to allocate and release high-priority CPUs from the existing PCT/SST-TF punit accounting.
• A DRA device builder (cpuclass.dra.go) that publishes one device per high-priority cpuClass × SST-TF punit, exposing nri/pctPriority, nri/packageID, nri/punitID as selectable device attributes.
• Full kubelet-plugin lifecycle (registration, PrepareResourceClaims/UnprepareResourceClaims, CDI spec writing, claim-state persistence and restart recovery).
• Wiring into the topology-aware policy: Setup/Start/Stop/Reconfigure integration, pool-supply accounting for claimed CPUs (multi-container claims via a refcount), and safe refusal of dra.enabled flips or DRA-attribute changes on classes with live claims.
• Helm chart additions: config.dra.{enabled,sharedCounters} values, RBAC, host mounts, and the base DeviceClass.
• An e2e test (test20-dra) covering claim -> pod -> CLOS-association → scheduler-accounting → cleanup, with a feature-gate probe that skips cleanly on clusters without the KEP-5075/KEP-5517 gates.

Requirements

• Kubernetes 1.34+ for KEP-5075 (DRAConsumableCapacity) — required.
• Kubernetes 1.37+ for KEP-5517 (DRANodeAllocatableResources) — optional; without it DRA allocation still works but scheduler-side capacity accounting is skipped.

Known v1 limitations (by design, not oversights)

• Only high-priority PCT classes are published as DRA devices; non-HP classes and CEL selectors targeting them are rejected at Prepare time.
  config.dra.sharedCounters is accepted by the schema but rejected at validation time — it's reserved for (KEP-5941) shared-counter publication, which isn't implemented yet.
• A DRA-claimed CPU can still collide with a CPU already exclusively granted to a reserved-type container (updateSharedAllocations doesn't re-pin/evict reserved grants); tracked as a follow-up.
  dra.enabled: true with local-file (non-Kubernetes) configuration degrades to DRA-disabled for the process lifetime rather than retrying, since the agent skips kube-client setup in that mode.
• Live Reconfigure() support for DRA config changes (validated cpuClass edits, live-claim-aware refusal, PostReconfigure-deferred republish) was implemented and then deferred to a future PR — not because it doesn't work, but to keep this PR's diff reviewable. Right now, any config change while DRA is (or was) enabled is refused outright with a restart-required error.

Testing

• Unit tests across all new/touched packages (go test ./... — 27 suites, all passing).
• test20-dra e2e test, run against a live Kubernetes 1.37 cluster with both feature gates enabled.

[Copilot]

🟡 Changes recommended

The e2e kubeadm config generation appears invalid for kubeadm (extraArgs as a list instead of a map), and the new file watch has error-path resource leaks that should be fixed before merging.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds optional Kubernetes Dynamic Resource Allocation (DRA) integration to the topology-aware policy, enabling CPUClass-based high-priority CPU requests via ResourceClaims while preserving the existing NRI allocation paths.

Changes:

• Introduces a policy-agnostic DRA kubelet-plugin framework (claim lifecycle, CDI spec writing, persisted claim state) and a reusable Kubernetes client wrapper.
• Extends the topology-aware policy and cpuclass/PCT allocator to publish and account for HP-class CPU capacity as DRA devices and to reconcile claim CPU usage with existing accounting.
• Updates Helm, CRDs, docs, and e2e harness to support enabling required K8s feature gates and to validate the end-to-end DRA claim → pod → accounting → cleanup flow.

File summaries

File	Description
test/e2e/run.sh	Adds env wiring/printing for K8s feature gates and klog verbosity.
test/e2e/run_tests.sh	Preserves caller-provided vm_name to avoid vagrant-dir collisions.
test/e2e/policies.test-suite/topology-aware/n4c16/test20-dra/code.var.sh	Adds an e2e test covering DRA ResourceClaim lifecycle and validation probes.
test/e2e/policies.test-suite/topology-aware/helm-config.yaml.in	Allows enabling topology-aware DRA via test Helm config templating.
test/e2e/playbook/provision.yaml	Adds kubeadm feature-gate/log-verbosity plumbing for e2e VM provisioning.
test/e2e/files/Vagrantfile.in	Passes feature-gates/log-verbosity vars into the provisioner.
pkg/resmgr/resource-manager.go	Adds post-reconfigure hook, lock helper, kube client wiring, and safer shutdown ordering.
pkg/resmgr/resource_manager_test.go	Adds unit tests for Stop ordering, post-reconfigure hook, lock helper, kube client wrapper.
pkg/resmgr/policy/policy.go	Extends policy backend interfaces and forwards Stop/PostReconfigure; wires kube client/node/lock into backend options.
pkg/resmgr/policy/policy_test.go	Adds tests for Stop/PostReconfigure forwarding and WithLock/KubeClientFn wiring.
pkg/resmgr/nri.go	Makes NRI stop nil-stub-safe.
pkg/resmgr/main/main.go	Ensures resource manager Stop runs after Start returns.
pkg/resmgr/main/main_test.go	Adds tests verifying Run() triggers Stop even on Start error.
pkg/resmgr/dra/doc.go	Adds package documentation and import-cycle constraint note.
pkg/resmgr/dra/deps.go	Defines DRA plugin dependency contracts (DeviceLister/Allocator/CDI/Store/WithLock/etc.).
pkg/resmgr/dra/logging.go	Adds logr bridge for kubelet plugin logging.
pkg/resmgr/dra/state.go	Persists prepared-claim state via resmgr cache (marshal/unmarshal + store).
pkg/resmgr/dra/state_test.go	Adds tests for claim-state persistence round-trips and empty-cache behavior.
pkg/resmgr/dra/cdi.go	Implements CDI spec writing/removal/listing and CDI device-name sanitization.
pkg/resmgr/dra/cdi_test.go	Adds unit tests for CDI naming and on-disk spec behaviors.
pkg/resmgr/cpuclass/internal/pct/pct.go	Adds DRA-aware HP CPU accounting, punit capacity snapshotting, and HP pick/release/account APIs.
pkg/resmgr/cpuclass/cpuclass.go	Adds DRA allocator pass-through methods and stores last-applied classes for device publication.
pkg/resmgr/cpuclass/dra.go	Adds CPUClass DRA validation and device publication builder (per class × punit).
pkg/resmgr/cpuclass/cpuclass_dra_test.go	Tests cpuclass handler’s DRA allocator pass-through behavior.
pkg/resmgr/cache/cache.go	Extends Container interface with CDI device-name access.
pkg/resmgr/cache/container.go	Implements CDI device-name extraction from NRI container metadata.
pkg/resmgr/cache/container_test.go	Adds tests for CDI device-name extraction and test helpers to populate CDI devices.
pkg/kubernetes/watch/watch.go	Rebrands logger scope to watch and re-exports watch aliases.
pkg/kubernetes/watch/object.go	Adds single-object watch wrapper with transparent reopen behavior.
pkg/kubernetes/watch/file.go	Adds fsnotify-based file watch that emits apimachinery-like events.
pkg/kubernetes/watch/watch_test.go	Adds tests for type/constant aliases and file/object watch behavior.
pkg/kubernetes/client/testdata/kubeconfig-example.yaml	Adds kubeconfig fixture for client tests.
pkg/kubernetes/client/client.go	Adds Kubernetes client wrapper exposing REST config, HTTP client, and clientset with options.
pkg/kubernetes/client/client_test.go	Adds tests for kube client wrapper option ordering, fallback behavior, and copy semantics.
pkg/apis/config/v1alpha1/resmgr/policy/cpuclass.go	Adds per-cpuClass DRA publish toggle and helper DRAPublish().
pkg/apis/config/v1alpha1/resmgr/policy/cpuclass_test.go	Adds tests for DRAPublish defaults and deepcopy behavior.
pkg/apis/config/v1alpha1/resmgr/policy/zz_generated.deepcopy.go	Adds deepcopy support for new CPUClass DRA fields.
pkg/apis/config/v1alpha1/resmgr/policy/topologyaware/config.go	Adds topology-aware DRA config struct + nil-safe getters.
pkg/apis/config/v1alpha1/resmgr/policy/topologyaware/config_test.go	Adds tests for topology-aware DRA getters and validation.
pkg/apis/config/v1alpha1/resmgr/policy/topologyaware/zz_generated.deepcopy.go	Adds deepcopy support for topology-aware DRA config.
cmd/plugins/topology-aware/policy/dra_adapter.go	Adds adapter bridging topology-aware policy cpuclass handler to DRA deps.
cmd/plugins/topology-aware/policy/dra_adapter_test.go	Tests adapter routing to current handler and nil-safety.
cmd/plugins/topology-aware/policy/dra.go	Adds topology-aware DRA plugin construction, publishing, and device-diff utilities.
cmd/plugins/topology-aware/policy/resources.go	Adds pool-supply claim CPU accounting (ClaimCPUs/UnclaimCPUs) and clone behavior.
cmd/plugins/topology-aware/policy/resources_test.go	Adds unit tests for supply claim accounting propagation, idempotency, and clone semantics.
cmd/plugins/topology-aware/policy/mocks_test.go	Extends mocks to support CDI devices and cpuset tracking for DRA-related tests.
cmd/plugins/topology-aware/policy/cache_test.go	Updates allocation marshalling expectations for new/changed fields.
cmd/plugins/template/policy/template-policy.go	Implements new backend Stop/PostReconfigure hooks (no-op).
cmd/plugins/balloons/policy/balloons-policy.go	Implements new backend Stop/PostReconfigure hooks (no-op).
pkg/agent/agent.go	Switches to shared Kubernetes client wrapper + new watch package; adds kube client accessors.
deployment/helm/topology-aware/values.yaml	Adds config.dra.{enabled,sharedCounters} values.
deployment/helm/topology-aware/templates/deviceclass.yaml	Installs base DeviceClass when DRA is enabled.
deployment/helm/topology-aware/templates/daemonset.yaml	Adds kubelet plugin and CDI host mounts when DRA is enabled.
deployment/helm/topology-aware/templates/clusterrole.yaml	Adds RBAC for ResourceSlices and ResourceClaims when DRA is enabled.
deployment/helm/topology-aware/crds/config.nri_topologyawarepolicies.yaml	Updates generated CRD schema to include DRA config and per-cpuClass DRA publish options.
deployment/helm/balloons/crds/config.nri_balloonspolicies.yaml	Updates generated CRD schema to include per-cpuClass DRA publish options.
config/crd/bases/config.nri_topologyawarepolicies.yaml	Updates base CRD schema for topology-aware DRA config and cpuClass DRA publish options.
config/crd/bases/config.nri_balloonspolicies.yaml	Updates base CRD schema for cpuClass DRA publish options.
docs/resource-policy/policy/topology-aware.md	Documents DRA prerequisites, enabling, and claim usage examples.
go.mod	Bumps Kubernetes libraries and adds CDI/DRA/logr dependencies.
.gitignore	Ignores extracted test fixtures and *.test artifacts.

Review details

• Files reviewed: 67/71 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

Assoc-only classes can over-publish CPU capacity, and failed reconfiguration publication is not retried.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 67/71 changed files
• Comments generated: 3
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

DRA validation incorrectly rejects non-published low-priority classes, and failed post-reconfigure publication can leave stale ResourceSlices without retry.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

cmd/plugins/topology-aware/policy/dra.go:152

• A transient publication failure here occurs after resmgr.reconfigure has already committed the new policy configuration, leaving the old ResourceSlices advertised. Unlike failures in republisherLoop, this path does not enqueue a retry, so the stale scheduler view can persist indefinitely unless an unrelated allocation triggers republication. Enqueue a background retry before returning the error.

func (p *policy) PostReconfigure() error {
	if p.draPlugin == nil {
		return nil
	}
	return p.draPlugin.PublishResources(p.draCtx)
}

• Files reviewed: 67/71 changed files
• Comments generated: 1
• Review effort level: Balanced

[fmuyassarov]

@bart0sh have you tested it out? The resourceSlice is created but empty.

[fmuyassarov]

Let me double check. I'm sure the gate was open during helm otherwise I wouldn't even see the resourceSlice creation. Maybe something I messed up. Will update soon

[Copilot]

🟡 Changes recommended

Claim accounting can leave CPUs double-bookable across partial multi-container claim consumption, and DRA currently blocks all live reconfiguration.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 67/71 changed files
• Comments generated: 2
• Review effort level: Balanced

[pravk03]

/cc

[klihub]

If this eventually becomes the CPU set a container is pinned to and runs on, then the picking logic should be CPU-topology aware. This is easiest done today with cpuallocator.Allocator, which both the topology-aware and balloons policy use also in this way: first filter available CPUs to an allowed/bounding set based on some filtering criteria (here HP/PCT CPUs), then use the allocator, so that if we have to and can pick CPUs worth one or more CPU-topological zone (full core, LLC-cluster, explicitly advertised cluster, shared NUMA node, die, socket), then we do try to pick in multiple chunks of those.

I would not use the cpuallocator directly here in all the places where need to pick CPUs... just to better prepare for potential changes to the cpuallocator, leaking through its API, which we might soon get sincesome of its offered functionality, such as picking CPUs by SST-based CPU classification, has been obsoleted and superseded by the this/the cpuclass package. So I would wrap it in some pickCpus(from cpuset.CPUSet) (picked cpuset.CPUSet, error) convenience function here in a single code site, an use that instead.

This might be (but I'm not sure) a moot comment right now from purely a technical point of view today, for instance if PCT-HP CPU cores are always bound to be limited in practice well below any other topological entity size than a full physical core/two hyperthreads. Nevertheless, I still think we should bake in more general topology-aware picking already from the beginning.

[klihub]

Same comment applies here as for picking CPUs. We should do it so that the remaining unreleased subset is picked in a CPU topology-aware fashion. Again, that's the easiest done today using the cpuallocator.

[klihub]

Shouldn't this more naturally belong to/logically fit better to the previous commit (pct: add PickHpCpus, ReleaseHpCpus, and Punits) ? This is fixing PCT bookkeeping internals, which is necessary due to the (IIUC second external allocation channel) changes put in by the previous commit.

[klihub]

ditto for most of the rest of the changes made here to cpuclass/PCT: it looks to me as they'd better fit logically together with the previous commit which introduced {Pick,Release}HpCpus().

[klihub]

This looks a bit like we'd have put essentially the whole full blown DRA driver, running almost fully contained within the policy implementation, as an completely independent second driver, the other allocation path (via the stock resmgr NRI layer) being otherwise oblivious to it than via the shared (resmgr-level) lock. That's pretty implicit awareness observable only via more variance in the NRI-triggered code paths timings (usually not visible probably), and potentially via deadlocking which is now more easy to achieve with locking responsibility to protect against the NRI code paths having been fully delegated to the policy-contained DRA driver and it's related logic/additions.

Is there a viable alternative to this, where we'd have a delegation of responsibilities with closer resemblance to how we do it on the NRI-triggered allocation/release paths ? There we basically have the (resmgr) infrastructure/machinery take care of most boilerplate, including most (or maybe all, I don't recall now by heart) of the locking/serialization. The generic resmgr/policy interface has then functions for allocation/reallocation/release and those are used to hook in the policy-specific allocation/release/CPU,memory,etc selection logic into the generic/boilerplate (NRI) driver at the resmgr level.

So going more in that direction would mean the resmgr/policy interface gaining extra functions specific to/for the DRA-based allocation/release code paths, which then would be used to hook the policy-specific DRA decision/allocation logic into the generic/boilerplate DRA machinery, which for instance could live at the resmgr level.

I admit that such a split is easier for the traditional allocation, resulting in a rather small and simple API footprint, which is one of the key reasons why such an abstraction suggested itself as an obvious choice. This might not be the case here, or at least not such a clear cut, due to (a) DRA (device) being more versatile in this regard via device names+attributes, etc. But I think we should give it a serious and honest consideration, including thinking about the next steps of 1) exposing something more via DRA from the T-A policy, maybe for instance isolated CPU cores, P/E-cores, etc., and 2) adding DRA support to the balloons policy. Then think about what are the consequences and see what kind of compromises we'd need to make with one and the other approach, so we could take an educated guess which one looks to be less painful in the future.