Merge branch 'main' into erez/support-gpu-operator-25-10-0-better

enoodle · web-flow · commit 33ccb26a78ea · 2025-11-13T11:43:03.000+01:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 - Added enforcement of the `nvidia` runtime class for GPU pods, with the option to enforce a custom runtime class, or disable enforcement entirely.
 - Added a preferred podAntiAffinity term by default for all services, can be set to required instead by setting `global.requireDefaultPodAffinityTerm`
 - Added support for service-level affinities
+- Added time aware scheduling configurations in scheduling shard
 
 ### Fixed
 - (Openshift only) - High CPU usage for the operator pod due to continues reconciles
diff --git a/cmd/podgrouper/app/app.go b/cmd/podgrouper/app/app.go
@@ -24,7 +24,7 @@ import (
 	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"
 	"sigs.k8s.io/controller-runtime/pkg/webhook"
 
-	"github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2"
+	v2 "github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2"
 	kubeAiSchedulerV2alpha2 "github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"
 	controllers "github.com/NVIDIA/KAI-scheduler/pkg/podgrouper"
 	pluginshub "github.com/NVIDIA/KAI-scheduler/pkg/podgrouper/podgrouper/hub"
@@ -51,7 +51,7 @@ func init() {
 
 type App struct {
 	Mgr               manager.Manager
-	DefaultPluginsHub pluginshub.PluginsHub
+	DefaultPluginsHub *pluginshub.DefaultPluginsHub
 
 	configs    controllers.Configs
 	pluginsHub pluginshub.PluginsHub
diff --git a/deployments/kai-scheduler/crds/kai.scheduler_configs.yaml b/deployments/kai-scheduler/crds/kai.scheduler_configs.yaml
@@ -3256,8 +3256,6 @@ spec:
                     description: PodLabelSelector filters pods for webhooks and pod
                       grouper
                     type: object
-                  prometheusEnabled:
-                    type: boolean
                   queueLabelKey:
                     description: QueueLabelKey specifies the pod label key whose value
                       will be the queue name of the pod.
diff --git a/deployments/kai-scheduler/crds/kai.scheduler_schedulingshards.yaml b/deployments/kai-scheduler/crds/kai.scheduler_schedulingshards.yaml
@@ -55,6 +55,10 @@ spec:
                     * Only valid flags defined in the scheduler's flag set will be accepted
                     * Duplicated flags will override the behavior of flags generated by other fields
                 type: object
+              kValue:
+                description: KValue specifies the kValue for the proportion plugin.
+                  Default is 1.0.
+                type: number
               minRuntime:
                 description: MinRuntime specifies the minimum runtime of a jobs in
                   the shard
@@ -87,6 +91,54 @@ spec:
                 description: QueueDepthPerAction max number of jobs to try for action
                   per queue
                 type: object
+              usageDBConfig:
+                description: UsageDBConfig defines configuration for the usage db
+                  client
+                properties:
+                  clientType:
+                    type: string
+                  connectionString:
+                    type: string
+                  connectionStringEnvVar:
+                    type: string
+                  usageParams:
+                    description: UsageParams defines common params for all usage db
+                      clients. Some clients may not support all the params.
+                    properties:
+                      extraParams:
+                        additionalProperties:
+                          type: string
+                        description: ExtraParams are extra parameters for the usage
+                          db client, which are client specific.
+                        type: object
+                      fetchInterval:
+                        description: Fetch interval of the usage. Default is 1 minute.
+                        type: string
+                      halfLifePeriod:
+                        description: Half life period of the usage. If not set, or
+                          set to 0, the usage will not be decayed.
+                        type: string
+                      stalenessPeriod:
+                        description: Staleness period of the usage. Default is 5 minutes.
+                        type: string
+                      tumblingWindowCronString:
+                        description: A cron string used to determine when to reset
+                          resource usage for all queues.
+                        type: string
+                      waitTimeout:
+                        description: Wait timeout of the usage. Default is 1 minute.
+                        type: string
+                      windowSize:
+                        description: Window size of the usage. Default is 1 week.
+                        type: string
+                      windowType:
+                        description: Window type for time-series aggregation. If not
+                          set, defaults to sliding.
+                        type: string
+                    type: object
+                required:
+                - clientType
+                type: object
             type: object
           status:
             description: SchedulingShardStatus defines the observed state of SchedulingShard
diff --git a/deployments/kai-scheduler/templates/rbac/prometheus-binding.yaml b/deployments/kai-scheduler/templates/rbac/prometheus-binding.yaml
@@ -7,7 +7,7 @@ metadata:
   name: kai-prometheus
 subjects:
   - kind: ServiceAccount
-    name: kai-prometheus
+    name: prometheus
     namespace: kai-scheduler
 roleRef:
   kind: ClusterRole
diff --git a/pkg/apis/kai/v1/global.go b/pkg/apis/kai/v1/global.go
@@ -68,9 +68,6 @@ type GlobalConfig struct {
 	// +kubebuilder:validation:Optional
 	PodLabelSelector map[string]string `json:"podLabelSelector,omitempty"`
 
-	// +kubebuilder:validation:Optional
-	PrometheusEnabled *bool `json:"prometheusEnabled,omitempty"`
-
 	// Connection defines the connection configuration for TSDB
 	// +kubebuilder:validation:Optional
 	ExternalTSDBConnection *Connection `json:"connection,omitempty"`
diff --git a/pkg/apis/kai/v1/schedulingshard_types.go b/pkg/apis/kai/v1/schedulingshard_types.go
@@ -21,6 +21,7 @@ import (
 	"k8s.io/utils/ptr"
 
 	"github.com/NVIDIA/KAI-scheduler/pkg/apis/kai/v1/common"
+	usagedbapi "github.com/NVIDIA/KAI-scheduler/pkg/scheduler/cache/usagedb/api"
 )
 
 const (
@@ -58,6 +59,14 @@ type SchedulingShardSpec struct {
 	// MinRuntime specifies the minimum runtime of a jobs in the shard
 	// +kubebuilder:validation:Optional
 	MinRuntime *MinRuntime `json:"minRuntime,omitempty"`
+
+	// KValue specifies the kValue for the proportion plugin. Default is 1.0.
+	// +kubebuilder:validation:Optional
+	KValue *float64 `json:"kValue,omitempty"`
+
+	// UsageDBConfig defines configuration for the usage db client
+	// +kubebuilder:validation:Optional
+	UsageDBConfig *usagedbapi.UsageDBConfig `yaml:"usageDBConfig,omitempty" json:"usageDBConfig,omitempty"`
 }
 
 func (s *SchedulingShardSpec) SetDefaultsWhereNeeded() {
diff --git a/pkg/apis/kai/v1/zz_generated.deepcopy.go b/pkg/apis/kai/v1/zz_generated.deepcopy.go
diff --git a/pkg/operator/operands/common/common.go b/pkg/operator/operands/common/common.go
@@ -31,8 +31,13 @@ var KaiServicesForServiceMonitor = []struct {
 	Name     string
 	Port     string
 	JobLabel string
+
+	LabelSelector map[string]string // optional, if not provided, "app": NAME will be used
+	Namespaces    []string          // optional, if not provided, the service will be monitored in the namespace of the KAI config
 }{
-	{"queue-controller", "metrics", "queue-controller"},
+	{"queue-controller", "metrics", "queue-controller", nil, nil},
+	{"kube-state-metrics", "http", "kube-state-metrics",
+		map[string]string{"app.kubernetes.io/name": "kube-state-metrics"}, []string{"monitoring", "default"}},
 }
 
 func AllControllersAvailable(
diff --git a/pkg/operator/operands/prometheus/resources.go b/pkg/operator/operands/prometheus/resources.go
@@ -24,7 +24,9 @@ import (
 )
 
 const (
-	mainResourceName = "prometheus"
+	mainResourceName              = "prometheus"
+	serviceMonitorAccountingLabel = "accounting"
+	serviceMonitorAccountingValue = "kai"
 )
 
 func prometheusForKAIConfig(
@@ -96,7 +98,7 @@ func prometheusForKAIConfig(
 	if config.ServiceMonitor != nil && *config.ServiceMonitor.Enabled {
 		prometheusSpec.ServiceMonitorSelector = &metav1.LabelSelector{
 			MatchLabels: map[string]string{
-				"accounting": mainResourceName,
+				serviceMonitorAccountingLabel: serviceMonitorAccountingValue,
 			},
 		}
 		prometheusSpec.ServiceMonitorNamespaceSelector = &metav1.LabelSelector{}
@@ -136,18 +138,26 @@ func serviceMonitorsForKAIConfig(
 			return nil, err
 		}
 
-		serviceMonitorObj.GetLabels()["accounting"] = mainResourceName
+		serviceMonitorObj.GetLabels()[serviceMonitorAccountingLabel] = serviceMonitorAccountingValue
+
+		namespaces := []string{kaiConfig.Spec.Namespace}
+		if kaiService.Namespaces != nil {
+			namespaces = kaiService.Namespaces
+		}
+
+		labelSelector := map[string]string{"app": kaiService.Name}
+		if kaiService.LabelSelector != nil {
+			labelSelector = kaiService.LabelSelector
+		}
 
 		// Set the ServiceMonitor spec from configuration
 		serviceMonitorSpec := monitoringv1.ServiceMonitorSpec{
 			JobLabel: kaiService.JobLabel,
 			NamespaceSelector: monitoringv1.NamespaceSelector{
-				MatchNames: []string{kaiConfig.Spec.Namespace},
+				MatchNames: namespaces,
 			},
 			Selector: metav1.LabelSelector{
-				MatchLabels: map[string]string{
-					"app": kaiService.Name,
-				},
+				MatchLabels: labelSelector,
 			},
 			Endpoints: []monitoringv1.Endpoint{
 				{
@@ -172,35 +182,6 @@ func serviceMonitorsForKAIConfig(
 		serviceMonitors = append(serviceMonitors, serviceMonitorObj)
 	}
 
-	kubeStateMetric := &monitoringv1.ServiceMonitor{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      "kube-state-metrics",
-			Namespace: kaiConfig.Spec.Namespace,
-			Labels: map[string]string{
-				"accounting": mainResourceName,
-			},
-		},
-		Spec: monitoringv1.ServiceMonitorSpec{
-			JobLabel: "kube-state-metrics",
-			NamespaceSelector: monitoringv1.NamespaceSelector{
-				MatchNames: []string{"monitoring", "default"},
-			},
-			Selector: metav1.LabelSelector{
-				MatchLabels: map[string]string{
-					"app.kubernetes.io/name": "kube-state-metrics",
-				},
-			},
-			Endpoints: []monitoringv1.Endpoint{
-				{
-					Port:            "http",
-					BearerTokenFile: "/var/run/secrets/kubernetes.io/serviceaccount/token",
-					Interval:        "30s",
-				},
-			},
-		},
-	}
-
-	serviceMonitors = append(serviceMonitors, kubeStateMetric)
 	return serviceMonitors, nil
 }
 
diff --git a/pkg/operator/operands/scheduler/resources_for_shard.go b/pkg/operator/operands/scheduler/resources_for_shard.go
@@ -12,17 +12,18 @@ import (
 	"github.com/spf13/pflag"
 	"golang.org/x/exp/slices"
 
-	"gopkg.in/yaml.v3"
 	v1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/yaml"
 
 	"github.com/NVIDIA/KAI-scheduler/cmd/scheduler/app/options"
 	kaiv1 "github.com/NVIDIA/KAI-scheduler/pkg/apis/kai/v1"
 	kaiConfigUtils "github.com/NVIDIA/KAI-scheduler/pkg/operator/config"
 	"github.com/NVIDIA/KAI-scheduler/pkg/operator/operands/common"
+	usagedbapi "github.com/NVIDIA/KAI-scheduler/pkg/scheduler/cache/usagedb/api"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/conf"
 )
 
@@ -136,11 +137,18 @@ func (s *SchedulerForShard) configMapForShard(
 
 	innerConfig.Actions = strings.Join(actions, ", ")
 
+	var proportionArgs map[string]string
+	if shard.Spec.KValue != nil {
+		proportionArgs = map[string]string{
+			"kValue": strconv.FormatFloat(*shard.Spec.KValue, 'f', -1, 64),
+		}
+	}
+
 	innerConfig.Tiers = []conf.Tier{
 		{
 			Plugins: []conf.PluginOption{
 				{Name: "predicates"},
-				{Name: "proportion"},
+				{Name: "proportion", Arguments: proportionArgs},
 				{Name: "priority"},
 				{Name: "nodeavailability"},
 				{Name: "resourcetype"},
@@ -185,6 +193,12 @@ func (s *SchedulerForShard) configMapForShard(
 		innerConfig.QueueDepthPerAction = shard.Spec.QueueDepthPerAction
 	}
 
+	usageDBConfig, err := getUsageDBConfig(shard, kaiConfig)
+	if err != nil {
+		return nil, err
+	}
+	innerConfig.UsageDBConfig = usageDBConfig
+
 	data, marshalErr := yaml.Marshal(&innerConfig)
 	if marshalErr != nil {
 		return nil, marshalErr
@@ -205,6 +219,41 @@ func validateJobDepthMap(shard *kaiv1.SchedulingShard, innerConfig conf.Schedule
 	return nil
 }
 
+func getUsageDBConfig(shard *kaiv1.SchedulingShard, kaiConfig *kaiv1.Config) (*usagedbapi.UsageDBConfig, error) {
+	// Check for nil inputs
+	if shard == nil {
+		return nil, fmt.Errorf("shard cannot be nil")
+	}
+	if kaiConfig == nil {
+		return nil, fmt.Errorf("kaiConfig cannot be nil")
+	}
+
+	if shard.Spec.UsageDBConfig == nil {
+		return nil, nil
+	}
+
+	usageDBConfig := shard.Spec.UsageDBConfig.DeepCopy()
+
+	if usageDBConfig.ClientType != "prometheus" {
+		return usageDBConfig, nil
+	}
+
+	if usageDBConfig.ConnectionString == "" && usageDBConfig.ConnectionStringEnvVar == "" {
+		// Use prometheus from config
+		if kaiConfig.Spec.Prometheus != nil &&
+			kaiConfig.Spec.Prometheus.Enabled != nil &&
+			*kaiConfig.Spec.Prometheus.Enabled {
+			usageDBConfig.ConnectionString = fmt.Sprintf("http://prometheus-operated.%s.svc.cluster.local:9090", kaiConfig.Spec.Namespace)
+		} else if kaiConfig.Spec.Global != nil && kaiConfig.Spec.Global.ExternalTSDBConnection != nil && kaiConfig.Spec.Global.ExternalTSDBConnection.URL != nil {
+			usageDBConfig.ConnectionString = *kaiConfig.Spec.Global.ExternalTSDBConnection.URL
+		} else {
+			return nil, fmt.Errorf("prometheus connection string not configured: either enable internal prometheus or configure external TSDB connection URL")
+		}
+	}
+
+	return usageDBConfig, nil
+}
+
 func (s *SchedulerForShard) serviceForShard(
 	ctx context.Context, readerClient client.Reader,
 	kaiConfig *kaiv1.Config, shard *kaiv1.SchedulingShard,
@@ -300,15 +349,17 @@ func addMinRuntimePluginIfNeeded(plugins *[]conf.PluginOption, minRuntime *kaiv1
 		return
 	}
 
-	minRuntimePlugin := conf.PluginOption{Name: "minruntime", Arguments: map[string]string{}}
+	minRuntimeArgs := make(map[string]string)
 
 	if minRuntime.PreemptMinRuntime != nil {
-		minRuntimePlugin.Arguments["defaultPreemptMinRuntime"] = *minRuntime.PreemptMinRuntime
+		minRuntimeArgs["defaultPreemptMinRuntime"] = *minRuntime.PreemptMinRuntime
 	}
 	if minRuntime.ReclaimMinRuntime != nil {
-		minRuntimePlugin.Arguments["defaultReclaimMinRuntime"] = *minRuntime.ReclaimMinRuntime
+		minRuntimeArgs["defaultReclaimMinRuntime"] = *minRuntime.ReclaimMinRuntime
 	}
 
+	minRuntimePlugin := conf.PluginOption{Name: "minruntime", Arguments: minRuntimeArgs}
+
 	*plugins = append(*plugins, minRuntimePlugin)
 }
 
diff --git a/pkg/operator/operands/scheduler/resources_test.go b/pkg/operator/operands/scheduler/resources_test.go
diff --git a/pkg/podgrouper/podgrouper/hub/hub.go b/pkg/podgrouper/podgrouper/hub/hub.go
diff --git a/pkg/scheduler/cache/usagedb/api/interface.go b/pkg/scheduler/cache/usagedb/api/interface.go
diff --git a/pkg/scheduler/cache/usagedb/prometheus/prometheus.go b/pkg/scheduler/cache/usagedb/prometheus/prometheus.go
