Handle unset usage lister client gracefully (#388)

* Handle unset usage lister client gracefully

Author: itsomri

pkg/scheduler/cache/cache.go

@@ -158,7 +158,9 @@ func newSchedulerCache(schedulerCacheParams *SchedulerCacheParams) *SchedulerCac
 	sc.podLister = sc.informerFactory.Core().V1().Pods().Lister()
 	sc.podGroupLister = sc.kubeAiSchedulerInformerFactory.Scheduling().V2alpha2().PodGroups().Lister()
 
-	sc.usageLister = usagedb.NewUsageLister(schedulerCacheParams.UsageDBClient, nil, nil, nil)
+	if schedulerCacheParams.UsageDBClient != nil {
+		sc.usageLister = usagedb.NewUsageLister(schedulerCacheParams.UsageDBClient, nil, nil, nil)
+	}
 
 	clusterInfo, err := cluster_info.New(sc.informerFactory, sc.kubeAiSchedulerInformerFactory, sc.kueueInformerFactory, sc.usageLister, sc.schedulingNodePoolParams,
 		sc.restrictNodeScheduling, &sc.K8sClusterPodAffinityInfo, sc.scheduleCSIStorage, sc.fullHierarchyFairness, sc.StatusUpdater)
@@ -194,15 +196,19 @@ func (sc *SchedulerCache) Run(stopCh <-chan struct{}) {
 	sc.kueueInformerFactory.Start(stopCh)
 	sc.StatusUpdater.Run(stopCh)
 
-	sc.usageLister.Start(stopCh)
+	if sc.usageLister != nil {
+		sc.usageLister.Start(stopCh)
+	}
 }
 
 func (sc *SchedulerCache) WaitForCacheSync(stopCh <-chan struct{}) {
 	sc.informerFactory.WaitForCacheSync(stopCh)
 	sc.kubeAiSchedulerInformerFactory.WaitForCacheSync(stopCh)
 	sc.kueueInformerFactory.WaitForCacheSync(stopCh)
 
-	sc.usageLister.WaitForCacheSync(stopCh)
+	if sc.usageLister != nil {
+		sc.usageLister.WaitForCacheSync(stopCh)
+	}
 }
 
 func (sc *SchedulerCache) Evict(evictedPod *v1.Pod, evictedPodGroup *podgroup_info.PodGroupInfo,

pkg/scheduler/cache/cluster_info/cluster_info.go

@@ -59,6 +59,7 @@ type ClusterInfo struct {
 	includeCSIStorageObjects bool
 	nodePoolSelector         labels.Selector
 	fairnessLevelType        FairnessLevelType
+	collectUsageData         bool
 }
 
 type FairnessLevelType string
@@ -106,6 +107,7 @@ func New(
 		nodePoolSelector:         nodePoolSelector,
 		fairnessLevelType:        fairnessLevelType,
 		podGroupSync:             podGroupSync,
+		collectUsageData:         usageLister != nil,
 	}, nil
 }
 
@@ -150,10 +152,9 @@ func (c *ClusterInfo) Snapshot() (*api.ClusterInfo, error) {
 
 	usage, usageErr := c.snapshotQueueResourceUsage()
 	if usageErr != nil {
-		log.InfraLogger.Warningf("error snapshotting queue resource usage: %c", usageErr)
+		log.InfraLogger.V(2).Warnf("error snapshotting queue resource usage: %c", usageErr)
 	}
 	if usage == nil {
-		log.InfraLogger.Warningf("resource usage is nil, using 0 values for all queues")
 		usage = queue_info.NewClusterUsage()
 	}
 	snapshot.QueueResourceUsage = *usage

pkg/scheduler/cache/cluster_info/data_lister/kubernetes_lister.go

@@ -112,7 +112,7 @@ func (k *k8sLister) ListQueues() ([]*enginev2.Queue, error) {
 
 func (k *k8sLister) ListResourceUsage() (*queue_info.ClusterUsage, error) {
 	if k.usageLister == nil {
-		return nil, fmt.Errorf("usage lister is not set")
+		return queue_info.NewClusterUsage(), fmt.Errorf("usage lister is not set")
 	}
 
 	return k.usageLister.GetResourceUsage()

pkg/scheduler/cache/cluster_info/queue.go

@@ -78,6 +78,10 @@ func (c *ClusterInfo) snapshotQueues() (map[common_info.QueueID]*queue_info.Queu
 }
 
 func (c *ClusterInfo) snapshotQueueResourceUsage() (*queue_info.ClusterUsage, error) {
+	if !c.collectUsageData {
+		return nil, nil
+	}
+
 	return c.dataLister.ListResourceUsage()
 }
 

pkg/scheduler/cache/usagedb/usagedb.go

@@ -65,6 +65,10 @@ func (l *UsageLister) GetResourceUsage() (*queue_info.ClusterUsage, error) {
 	l.lastUsageDataMutex.RLock()
 	defer l.lastUsageDataMutex.RUnlock()
 
+	if l.client == nil {
+		return nil, fmt.Errorf("client is not set")
+	}
+
 	if l.lastUsageDataTime == nil {
 		return nil, fmt.Errorf("usage data is not available")
 	}