Implement DeviceClassResolver interface

Author: gnufied

cluster-autoscaler/simulator/clustersnapshot/store/basic.go

@@ -360,6 +360,10 @@ func (snapshot *BasicSnapshotStore) DeviceClasses() fwk.DeviceClassLister {
 	return snapshot.DraSnapshot().DeviceClasses()
 }
 
+func (snapshot *BasicSnapshotStore) DeviceClassResolver() fwk.DeviceClassResolver {
+	return snapshot.DraSnapshot().DeviceClassResolver()
+}
+
 // List returns the list of nodes in the snapshot.
 func (snapshot *basicSnapshotStoreNodeLister) List() ([]fwk.NodeInfo, error) {
 	return (*BasicSnapshotStore)(snapshot).getInternalData().listNodeInfos(), nil

cluster-autoscaler/simulator/clustersnapshot/store/delta.go

@@ -443,6 +443,10 @@ func (snapshot *DeltaSnapshotStore) CSINodes() fwk.CSINodeLister {
 	return snapshot.csiSnapshot.CSINodes()
 }
 
+func (snapshot *DeltaSnapshotStore) DeviceClassResolver() fwk.DeviceClassResolver {
+	return snapshot.DraSnapshot().DeviceClassResolver()
+}
+
 // AddSchedulerNodeInfo adds a NodeInfo.
 func (snapshot *DeltaSnapshotStore) AddSchedulerNodeInfo(nodeInfo fwk.NodeInfo) error {
 	if _, err := snapshot.data.addNode(nodeInfo.Node()); err != nil {

cluster-autoscaler/simulator/dynamicresources/snapshot/snapshot.go

@@ -92,7 +92,7 @@ func (s *Snapshot) ResourceClaims() schedulerframework.ResourceClaimTracker {
 }
 
 func (s *Snapshot) DeviceClassResolver() schedulerframework.DeviceClassResolver {
-	return snapshotClassResolver{snapshot: s}
+	return newSnapshotDeviceClassResolver(s)
 }
 
 // ResourceSlices exposes the Snapshot as schedulerframework.ResourceSliceLister, in order to interact with

cluster-autoscaler/simulator/dynamicresources/snapshot/snapshot_class_resolver.go

@@ -3,12 +3,37 @@ package snapshot
 import (
 	v1 "k8s.io/api/core/v1"
 	resourceapi "k8s.io/api/resource/v1"
+	schedulerframework "k8s.io/kube-scheduler/framework"
 )
 
-type snapshotClassResolver struct {
-	snapshot *Snapshot
+type snapshotDeviceClassResolver struct {
+	// resourceName2class maps extended resource name to device class
+	resourceName2class map[v1.ResourceName]*resourceapi.DeviceClass
 }
 
-func (s snapshotClassResolver) GetDeviceClass(resourceName v1.ResourceName) *resourceapi.DeviceClass {
-	return s.snapshot.getDeviceClass(resourceName)
+var _ schedulerframework.DeviceClassResolver = &snapshotDeviceClassResolver{}
+
+func newSnapshotDeviceClassResolver(snapshot *Snapshot) schedulerframework.DeviceClassResolver {
+	resourceName2class := make(map[v1.ResourceName]*resourceapi.DeviceClass)
+	for _, class := range snapshot.listDeviceClasses() {
+		classResourceName := class.Name
+		extendedResourceName := class.Spec.ExtendedResourceName
+		if extendedResourceName != nil {
+			resourceName2class[v1.ResourceName(*extendedResourceName)] = class
+		}
+
+		// Also add the default extended resource name
+		defaultResourceName := v1.ResourceName(resourceapi.ResourceDeviceClassPrefix + classResourceName)
+		resourceName2class[defaultResourceName] = class
+	}
+	return snapshotDeviceClassResolver{resourceName2class: resourceName2class}
+}
+
+// GetDeviceClass returns the device class for the given extended resource name
+func (s snapshotDeviceClassResolver) GetDeviceClass(resourceName v1.ResourceName) *resourceapi.DeviceClass {
+	class, found := s.resourceName2class[resourceName]
+	if !found {
+		return nil
+	}
+	return class
 }