@@ -23,13 +23,19 @@ import (
2323 corev1 "k8s.io/api/core/v1"
2424 resourceapi "k8s.io/api/resource/v1"
2525 "k8s.io/apimachinery/pkg/api/resource"
26+ metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
2627 "k8s.io/apimachinery/pkg/runtime"
2728 "k8s.io/apimachinery/pkg/runtime/schema"
2829 "k8s.io/apiserver/pkg/admission"
2930 quota "k8s.io/apiserver/pkg/quota/v1"
3031 "k8s.io/apiserver/pkg/quota/v1/generic"
32+ utilfeature "k8s.io/apiserver/pkg/util/feature"
33+ corev1listers "k8s.io/client-go/listers/core/v1"
34+ "k8s.io/dynamic-resource-allocation/deviceclass/extendedresourcecache"
3135 resourceinternal "k8s.io/kubernetes/pkg/apis/resource"
3236 resourceversioned "k8s.io/kubernetes/pkg/apis/resource/v1"
37+ "k8s.io/kubernetes/pkg/features"
38+ "k8s.io/utils/clock"
3339)
3440
3541// The name used for object count quota. This evaluator takes over counting
@@ -38,21 +44,26 @@ import (
3844var ClaimObjectCountName = generic .ObjectCountQuotaResourceNameFor (resourceapi .SchemeGroupVersion .WithResource ("resourceclaims" ).GroupResource ())
3945
4046// V1ResourceByDeviceClass returns a quota resource name by device class.
47+ // gpuclass -> gpuclass.deviceclass.resource.k8s.io/devices
4148func V1ResourceByDeviceClass (className string ) corev1.ResourceName {
4249 return corev1 .ResourceName (className + corev1 .ResourceClaimsPerClass )
4350}
4451
4552// NewResourceClaimEvaluator returns an evaluator that can evaluate resource claims
46- func NewResourceClaimEvaluator (f quota.ListerForResourceFunc ) quota.Evaluator {
53+ func NewResourceClaimEvaluator (f quota.ListerForResourceFunc , m * extendedresourcecache. ExtendedResourceCache , podsGetter corev1listers. PodLister ) quota.Evaluator {
4754 listFuncByNamespace := generic .ListResourceUsingListerFunc (f , resourceapi .SchemeGroupVersion .WithResource ("resourceclaims" ))
48- claimEvaluator := & claimEvaluator {listFuncByNamespace : listFuncByNamespace }
55+ claimEvaluator := & claimEvaluator {listFuncByNamespace : listFuncByNamespace , deviceClassMapping : m , podsGetter : podsGetter }
4956 return claimEvaluator
5057}
5158
5259// claimEvaluator knows how to evaluate quota usage for resource claims
5360type claimEvaluator struct {
5461 // listFuncByNamespace knows how to list resource claims
5562 listFuncByNamespace generic.ListFuncByNamespace
63+ // a global cache of device class and extended resource mapping
64+ deviceClassMapping * extendedresourcecache.ExtendedResourceCache
65+ // podsGetter is used to get pods
66+ podsGetter corev1listers.PodLister
5667}
5768
5869// Constraints verifies that all required resources are present on the item.
@@ -98,11 +109,91 @@ func (p *claimEvaluator) MatchingResources(items []corev1.ResourceName) []corev1
98109 if item == ClaimObjectCountName /* object count quota fields */ ||
99110 strings .HasSuffix (string (item ), corev1 .ResourceClaimsPerClass /* by device class */ ) {
100111 result = append (result , item )
112+ continue
113+ }
114+ if utilfeature .DefaultFeatureGate .Enabled (features .DRAExtendedResource ) {
115+ if strings .HasPrefix (string (item ), corev1 .ResourceImplicitExtendedClaimsPerClass /* by implicit extended resource name */ ) {
116+ className := string (item [len (corev1 .ResourceImplicitExtendedClaimsPerClass ):])
117+ if p .deviceClassMapping .GetExtendedResource (className ) != "" {
118+ result = append (result , item )
119+ continue
120+ }
121+ }
122+ if isExtendedResourceNameForQuota (item ) /* by extended resource name */ {
123+ resourceName := string (item [len (corev1 .DefaultResourceRequestsPrefix ):])
124+ if p .deviceClassMapping .GetDeviceClass (corev1 .ResourceName (resourceName )) != nil {
125+ result = append (result , item )
126+ }
127+ }
101128 }
102129 }
103130 return result
104131}
105132
133+ func (p * claimEvaluator ) addExtendedResourceQuota (resourceClaimUsage map [corev1.ResourceName ]resource.Quantity , podUsage corev1.ResourceList ) {
134+ extendedResourceUsage := make (map [corev1.ResourceName ]resource.Quantity )
135+ for name , quantity := range resourceClaimUsage {
136+ // e.g. myclass
137+ deviceClassName , isDeviceClassUsage := strings .CutSuffix (string (name ), corev1 .ResourceClaimsPerClass )
138+ if ! isDeviceClassUsage || len (deviceClassName ) == 0 {
139+ continue
140+ }
141+
142+ // requests.deviceclass.resource.kubernetes.io/myclass
143+ extendedResourceUsage [corev1 .ResourceName (corev1 .ResourceImplicitExtendedClaimsPerClass + deviceClassName )] = quantity
144+
145+ // e.g. example.com/mygpu
146+ if extendedResourceName := p .deviceClassMapping .GetExtendedResource (deviceClassName ); len (extendedResourceName ) > 0 {
147+ // requests.example.com/gpu
148+ extendedResourceUsage [corev1 .ResourceName (corev1 .DefaultResourceRequestsPrefix + extendedResourceName )] = quantity
149+ }
150+ }
151+
152+ for name , quantity := range extendedResourceUsage {
153+ // Subtract any amount already accounted for in the pod
154+ if podQuantity , found := podUsage [name ]; found {
155+ quantity .Sub (podQuantity )
156+ }
157+ // Add any remaining amount to the resource claim resources
158+ if quantity .CmpInt64 (0 ) > 0 {
159+ resourceClaimUsage [name ] = quantity
160+ }
161+ }
162+ }
163+
164+ // Verify extended resource claim owning pod exists, and the pod's ExtendedResourceClaimStatus points
165+ // back to the claim if it's not nil, and returns the pod's quota usage. If any error is encountered, nil is returned.
166+ func (p * claimEvaluator ) getVerifiedPodUsage (claim * resourceapi.ResourceClaim ) corev1.ResourceList {
167+ if claim .Annotations [resourceapi .ExtendedResourceClaimAnnotation ] != "true" {
168+ return nil
169+ }
170+ controllerRef := metav1 .GetControllerOf (claim )
171+ if controllerRef == nil {
172+ return nil
173+ }
174+ if controllerRef .Kind != "Pod" || controllerRef .APIVersion != "v1" {
175+ return nil
176+ }
177+ if p .podsGetter == nil {
178+ return nil
179+ }
180+ pod , err := p .podsGetter .Pods (claim .Namespace ).Get (controllerRef .Name )
181+ if err != nil {
182+ return nil
183+ }
184+ if controllerRef .UID != pod .UID {
185+ return nil
186+ }
187+ if pod .Status .ExtendedResourceClaimStatus != nil && pod .Status .ExtendedResourceClaimStatus .ResourceClaimName != claim .Name {
188+ return nil
189+ }
190+ quotaReqs , err := PodUsageFunc (pod , clock.RealClock {})
191+ if err != nil {
192+ return nil
193+ }
194+ return quotaReqs
195+ }
196+
106197// Usage knows how to measure usage associated with item.
107198func (p * claimEvaluator ) Usage (item runtime.Object ) (corev1.ResourceList , error ) {
108199 result := corev1.ResourceList {}
@@ -168,6 +259,10 @@ func (p *claimEvaluator) Usage(item runtime.Object) (corev1.ResourceList, error)
168259 }
169260 }
170261
262+ if utilfeature .DefaultFeatureGate .Enabled (features .DRAExtendedResource ) {
263+ p .addExtendedResourceQuota (result , p .getVerifiedPodUsage (claim ))
264+ }
265+
171266 return result , nil
172267}
173268
0 commit comments