Add time aware configs to scheduling shard

itsomri · itsomri · commit d62e5fecec41 · 2025-11-11T19:43:51.000+02:00
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/scheduler/resources_for_shard.go b/pkg/operator/operands/scheduler/resources_for_shard.go
@@ -12,12 +12,12 @@ 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"
@@ -136,11 +136,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 +192,10 @@ func (s *SchedulerForShard) configMapForShard(
 		innerConfig.QueueDepthPerAction = shard.Spec.QueueDepthPerAction
 	}
 
+	if shard.Spec.UsageDBConfig != nil {
+		innerConfig.UsageDBConfig = shard.Spec.UsageDBConfig
+	}
+
 	data, marshalErr := yaml.Marshal(&innerConfig)
 	if marshalErr != nil {
 		return nil, marshalErr
@@ -300,15 +311,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
@@ -8,15 +8,15 @@ import (
 	"fmt"
 	"strings"
 	"testing"
-
-	"gopkg.in/yaml.v3"
+	"time"
 
 	"github.com/spf13/pflag"
 
 	"github.com/NVIDIA/KAI-scheduler/cmd/scheduler/app/options"
 	kaiv1 "github.com/NVIDIA/KAI-scheduler/pkg/apis/kai/v1"
 	kaiv1qc "github.com/NVIDIA/KAI-scheduler/pkg/apis/kai/v1/queue_controller"
 	kaiv1scheduler "github.com/NVIDIA/KAI-scheduler/pkg/apis/kai/v1/scheduler"
+	usagedbapi "github.com/NVIDIA/KAI-scheduler/pkg/scheduler/cache/usagedb/api"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/conf"
 
 	"github.com/stretchr/testify/assert"
@@ -26,6 +26,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/utils/ptr"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+	"sigs.k8s.io/yaml"
 )
 
 func TestDeploymentForShard(t *testing.T) {
@@ -432,6 +433,59 @@ tiers:
 			},
 			expectedErr: true,
 		},
+		{
+			name: "usage DB configuration",
+			config: &kaiv1.Config{
+				Spec: kaiv1.ConfigSpec{},
+			},
+			shard: &kaiv1.SchedulingShard{
+				Spec: kaiv1.SchedulingShardSpec{
+					UsageDBConfig: &usagedbapi.UsageDBConfig{
+						ClientType:       "prometheus",
+						ConnectionString: "http://prometheus-operated.kai-scheduler.svc.cluster.local:9090",
+						UsageParams: &usagedbapi.UsageParams{
+							HalfLifePeriod: &metav1.Duration{Duration: 10 * time.Minute},
+							WindowSize:     &metav1.Duration{Duration: 10 * time.Minute},
+							WindowType:     ptr.To(usagedbapi.SlidingWindow),
+						},
+					},
+				},
+			},
+			expected: map[string]string{
+				"config.yaml": `actions: allocate,consolidation,reclaim,preempt,stalegangeviction
+tiers:
+- plugins:
+  - name: predicates
+  - name: proportion
+  - name: priority
+  - name: nodeavailability
+  - name: resourcetype
+  - name: podaffinity
+  - name: elastic
+  - name: kubeflow
+  - name: ray
+  - name: subgrouporder
+  - name: taskorder
+  - name: nominatednode
+  - name: dynamicresources
+  - name: minruntime
+  - name: topology
+  - name: snapshot
+  - name: gpupack
+  - name: nodeplacement
+    arguments:
+      cpu: binpack
+      gpu: binpack
+  - name: gpusharingorder
+usageDBConfig:
+  clientType: prometheus
+  connectionString: http://prometheus-operated.kai-scheduler.svc.cluster.local:9090
+  usageParams:
+    halfLifePeriod: 10m
+    windowSize: 10m
+    windowType: sliding`,
+			},
+		},
 	}
 
 	for _, tt := range tests {
@@ -607,3 +661,61 @@ func TestServiceAccountForScheduler(t *testing.T) {
 		})
 	}
 }
+
+func TestMarshalingShardVsConfig(t *testing.T) {
+	shardSpecString := `
+spec:
+  partitionLabelValue: ""
+  placementStrategy:
+    cpu: binpack
+    gpu: binpack
+  usageDBConfig:
+    clientType: prometheus
+    connectionString: http://prometheus-operated.kai-scheduler.svc.cluster.local:9090
+    usageParams:
+      halfLifePeriod: 10m
+      windowSize: 10m
+      windowType: sliding
+`
+
+	shardSpec := &kaiv1.SchedulingShardSpec{}
+	err := yaml.Unmarshal([]byte(shardSpecString), shardSpec)
+	assert.NoError(t, err)
+
+	configString := `actions: allocate,consolidation,reclaim,preempt,stalegangeviction
+tiers:
+- plugins:
+  - name: predicates
+  - name: proportion
+  - name: priority
+  - name: nodeavailability
+  - name: resourcetype
+  - name: podaffinity
+  - name: elastic
+  - name: kubeflow
+  - name: ray
+  - name: subgrouporder
+  - name: taskorder
+  - name: nominatednode
+  - name: dynamicresources
+  - name: minruntime
+  - name: topology
+  - name: snapshot
+  - name: gpupack
+  - name: nodeplacement
+    arguments:
+      cpu: binpack
+      gpu: binpack
+  - name: gpusharingorder
+usageDBConfig:
+  clientType: prometheus
+  connectionString: http://prometheus-operated.kai-scheduler.svc.cluster.local:9090
+  usageParams:
+    halfLifePeriod: 10m
+    windowSize: 10m
+    windowType: sliding
+`
+	config := &conf.SchedulerConfiguration{}
+	err = yaml.Unmarshal([]byte(configString), config)
+	assert.NoError(t, err)
+}
diff --git a/pkg/scheduler/cache/usagedb/prometheus/prometheus.go b/pkg/scheduler/cache/usagedb/prometheus/prometheus.go
@@ -18,6 +18,7 @@ import (
 	promv1 "github.com/prometheus/client_golang/api/prometheus/v1"
 	"github.com/prometheus/common/model"
 	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 const (
@@ -139,7 +140,7 @@ func (p *PrometheusClient) GetResourceUsage() (*queue_info.ClusterUsage, error)
 func (p *PrometheusClient) queryResourceCapacity(ctx context.Context, capacityMetric string, queryByWindow usageWindowQueryFunction) (float64, error) {
 	decayedCapacityMetric := capacityMetric
 	if p.usageParams.HalfLifePeriod != nil {
-		decayedCapacityMetric = fmt.Sprintf("((%s) * (%s))", capacityMetric, getExponentialDecayQuery(&p.usageParams.HalfLifePeriod.Duration))
+		decayedCapacityMetric = fmt.Sprintf("((%s) * (%s))", capacityMetric, getExponentialDecayQuery(p.usageParams.HalfLifePeriod))
 	}
 
 	capacityResult, warnings, err := queryByWindow(ctx, decayedCapacityMetric)
@@ -170,7 +171,7 @@ func (p *PrometheusClient) queryResourceUsage(
 
 	decayedAllocationMetric := allocationMetric
 	if p.usageParams.HalfLifePeriod != nil {
-		decayedAllocationMetric = fmt.Sprintf("((%s) * (%s))", allocationMetric, getExponentialDecayQuery(&p.usageParams.HalfLifePeriod.Duration))
+		decayedAllocationMetric = fmt.Sprintf("((%s) * (%s))", allocationMetric, getExponentialDecayQuery(p.usageParams.HalfLifePeriod))
 	}
 
 	usageResult, warnings, err := queryByWindow(ctx, decayedAllocationMetric)
@@ -189,7 +190,8 @@ func (p *PrometheusClient) queryResourceUsage(
 
 	usageVector := usageResult.(model.Vector)
 	if len(usageVector) == 0 {
-		return nil, fmt.Errorf("no data returned for cluster usage metric %s", decayedAllocationMetric)
+		log.InfraLogger.V(3).Warnf("No data returned for cluster usage metric %s", decayedAllocationMetric)
+		return queueUsage, nil
 	}
 
 	for _, usageSample := range usageVector {
@@ -205,11 +207,14 @@ func (p *PrometheusClient) queryResourceUsage(
 func (p *PrometheusClient) querySlidingTimeWindow(ctx context.Context, decayedAllocationMetric string) (model.Value, promv1.Warnings, error) {
 	usageQuery := fmt.Sprintf("sum_over_time((%s)[%s:%s])",
 		decayedAllocationMetric,
-		p.usageParams.WindowSize.String(),
+		p.usageParams.WindowSize.Duration.String(),
 		p.queryResolution.String(),
 	)
 
 	usageResult, warnings, err := p.client.Query(ctx, usageQuery, time.Now())
+	if err != nil {
+		err = fmt.Errorf("%w, full query: %s", err, usageQuery)
+	}
 	return usageResult, warnings, err
 }
 
@@ -222,11 +227,16 @@ func (p *PrometheusClient) queryTumblingTimeWindow(ctx context.Context, decayedA
 		End:   time.Now(),
 		Step:  p.queryResolution,
 	})
+
+	if err != nil {
+		err = fmt.Errorf("%w, full query: %s", err, usageQuery)
+	}
+
 	return usageResult, warnings, err
 }
 
 func (p *PrometheusClient) getLatestUsageResetTime() time.Time {
-	maxWindowStartingPoint := time.Now().Add(-*&p.usageParams.WindowSize.Duration)
+	maxWindowStartingPoint := time.Now().Add(-p.usageParams.WindowSize.Duration)
 	lastUsageReset := maxWindowStartingPoint
 	nextInWindowReset := maxWindowStartingPoint
 
@@ -237,12 +247,12 @@ func (p *PrometheusClient) getLatestUsageResetTime() time.Time {
 	return lastUsageReset
 }
 
-func getExponentialDecayQuery(halfLifePeriod *time.Duration) string {
+func getExponentialDecayQuery(halfLifePeriod *metav1.Duration) string {
 	if halfLifePeriod == nil {
 		return ""
 	}
 
-	halfLifeSeconds := halfLifePeriod.Seconds()
+	halfLifeSeconds := halfLifePeriod.Duration.Seconds()
 	now := time.Now().Unix()
 
 	return fmt.Sprintf("0.5^((%d - time()) / %f)", now, halfLifeSeconds)
