Cluster autoscaler adjustment for GPU sharing pods (#119)

* added node scale adjuster code

* added node scale adjuster deployment

* Added autoscaling flag to helm chart values

* Added autoscaling README doc

* Changed autoscaling to clusterAutoscaling

* Added node-scale-adjuster arguments reference to autoscaling doc

Author: romanbaron

Makefile

@@ -16,7 +16,7 @@ KUSTOMIZE ?= $(LOCALBIN)/kustomize
 
 # Space seperated list of services to build by default
 # SERVICE_NAMES := service1 service2 service3
-SERVICE_NAMES := podgrouper scheduler binder webhookmanager resourcereservation snapshot-tool
+SERVICE_NAMES := podgrouper scheduler binder webhookmanager resourcereservation snapshot-tool scalingpod nodescaleadjuster
 
 
 lint: fmt-go vet-go lint-go
@@ -62,6 +62,7 @@ manifests: controller-gen kustomize ## Generate ClusterRole and CustomResourceDe
 	$(CONTROLLER_GEN) rbac:roleName=kai-binder,headerFile="./hack/boilerplate.yaml.txt" paths="./pkg/binder/..." paths="./cmd/binder/..." output:stdout > deployments/kai-scheduler/templates/rbac/binder.yaml
 	$(CONTROLLER_GEN) rbac:roleName=kai-resource-reservation,headerFile="./hack/boilerplate.yaml.txt" paths="./pkg/resourcereservation/..." paths="./cmd/resourcereservation/..." output:stdout > deployments/kai-scheduler/templates/rbac/resourcereservation.yaml
 	$(CONTROLLER_GEN) rbac:roleName=kai-scheduler,headerFile="./hack/boilerplate.yaml.txt" paths="./pkg/scheduler/..." paths="./cmd/scheduler/..." output:stdout > deployments/kai-scheduler/templates/rbac/scheduler.yaml
+	$(CONTROLLER_GEN) rbac:roleName=kai-node-scale-adjuster,headerFile="./hack/boilerplate.yaml.txt" paths="./pkg/nodescaleadjuster/..." paths="./cmd/nodescaleadjuster/..." output:stdout > deployments/kai-scheduler/templates/rbac/nodescaleadjuster.yaml
 
 	$(CONTROLLER_GEN) rbac:roleName=kai-webhookmanager,headerFile="./hack/boilerplate.yaml.txt" paths="./pkg/webhookmanager/..." paths="./cmd/webhookmanager/..." output:stdout > deployments/kustomization/webhookmanager-clusterrole/resource.yaml
 	$(KUSTOMIZE) build deployments/kustomization/webhookmanager-clusterrole >  deployments/kai-scheduler/templates/rbac/webhookmanager.yaml

cmd/binder/app/app.go

@@ -122,7 +122,7 @@ func New() (*App, error) {
 	rrs := resourcereservation.NewService(options.FakeGPUNodes, clientWithWatch, options.ResourceReservationPodImage,
 		time.Duration(options.ResourceReservationAllocationTimeout)*time.Second,
 		options.ResourceReservationNamespace, options.ResourceReservationServiceAccount,
-		options.ResourceReservationAppLabel)
+		options.ResourceReservationAppLabel, options.ScalingPodNamespace)
 
 	reconcilerParams := &controllers.ReconcilerParams{
 		MaxConcurrentReconciles:     options.MaxConcurrentReconciles,

cmd/binder/app/options.go

@@ -16,6 +16,7 @@ type Options struct {
 	ResourceReservationPodImage          string
 	ResourceReservationAppLabel          string
 	ResourceReservationAllocationTimeout int
+	ScalingPodNamespace                  string
 	QPS                                  float64
 	Burst                                int
 	MaxConcurrentReconciles              int
@@ -54,6 +55,9 @@ func InitOptions() *Options {
 	fs.IntVar(&options.ResourceReservationAllocationTimeout,
 		"resource-reservation-allocation-timeout", 40,
 		"Resource reservation allocation timeout in seconds")
+	fs.StringVar(&options.ScalingPodNamespace,
+		"scale-adjust-namespace", "kai-scale-adjust",
+		"Scaling pods namespace")
 	fs.Float64Var(&options.QPS,
 		"qps", 50,
 		"Queries per second to the K8s API server")

cmd/nodescaleadjuster/app/app.go

@@ -0,0 +1,99 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package app
+
+import (
+	"flag"
+	"log"
+
+	"go.uber.org/zap/zapcore"
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
+	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/healthz"
+	"sigs.k8s.io/controller-runtime/pkg/log/zap"
+
+	"github.com/NVIDIA/KAI-scheduler/pkg/nodescaleadjuster/consts"
+	"github.com/NVIDIA/KAI-scheduler/pkg/nodescaleadjuster/controller"
+	"github.com/NVIDIA/KAI-scheduler/pkg/nodescaleadjuster/scale_adjuster"
+	"github.com/NVIDIA/KAI-scheduler/pkg/nodescaleadjuster/scaler"
+)
+
+var (
+	scheme   = runtime.NewScheme()
+	setupLog = ctrl.Log.WithName("setup")
+)
+
+func init() {
+	utilruntime.Must(clientgoscheme.AddToScheme(scheme))
+	utilruntime.Must(v1.AddToScheme(scheme))
+	// +kubebuilder:scaffold:scheme
+}
+
+func Run() error {
+	options := NewOptions()
+	options.AddFlags()
+
+	opts := zap.Options{
+		Development: true,
+		TimeEncoder: zapcore.ISO8601TimeEncoder,
+	}
+	opts.BindFlags(flag.CommandLine)
+	flag.Parse()
+	ctrl.SetLogger(zap.New(zap.UseFlagOptions(&opts)))
+
+	log.Println("Node scale adjuster started")
+
+	clientConfig := ctrl.GetConfigOrDie()
+	mgr, err := ctrl.NewManager(clientConfig, ctrl.Options{
+		Scheme: scheme,
+	})
+	if err != nil {
+		setupLog.Error(err, "unable to start manager")
+		return err
+	}
+
+	nodeScaler := scaler.NewScaler(mgr.GetClient(), options.ScalingPodImage, options.ScalingPodNamespace,
+		options.ScalingPodAppLabel, options.ScalingPodServiceAccount)
+
+	scaleAdjuster := scale_adjuster.NewScaleAdjuster(
+		mgr.GetClient(),
+		nodeScaler,
+		options.ScalingPodNamespace,
+		consts.DefaultCoolDownSeconds,
+		options.GPUMemoryToFractionRatio,
+		options.SchedulerName)
+
+	podReconciler := &controller.PodReconciler{
+		ScaleAdjuster:      scaleAdjuster,
+		SchedulerName:      options.SchedulerName,
+		NodeScaleNamespace: options.ScalingPodImage,
+		Client:             mgr.GetClient(),
+		Scheme:             mgr.GetScheme(),
+	}
+
+	if err = podReconciler.SetupWithManager(mgr); err != nil {
+		setupLog.Error(err, "unable to create controller", "controller", "Pod")
+		return err
+	}
+
+	if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil {
+		setupLog.Error(err, "unable to set up health check")
+		return err
+	}
+	if err := mgr.AddReadyzCheck("readyz", healthz.Ping); err != nil {
+		setupLog.Error(err, "unable to set up ready check")
+		return err
+	}
+
+	setupLog.Info("starting manager")
+	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
+		setupLog.Error(err, "problem running manager")
+		return err
+	}
+
+	return nil
+}

cmd/nodescaleadjuster/app/options.go

@@ -0,0 +1,48 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package app
+
+import (
+	"flag"
+	"fmt"
+
+	"github.com/NVIDIA/KAI-scheduler/pkg/nodescaleadjuster/consts"
+)
+
+type Options struct {
+	ScalingPodImage          string
+	ScalingPodNamespace      string
+	SchedulerName            string
+	GPUMemoryToFractionRatio float64
+	ScalingPodAppLabel       string
+	ScalingPodServiceAccount string
+}
+
+// NewOptions creates a new Options
+func NewOptions() *Options {
+	s := Options{}
+	return &s
+}
+
+// AddFlags adds flags for a specific CMServer to the specified FlagSet
+func (s *Options) AddFlags() {
+	flag.StringVar(&s.ScalingPodImage,
+		"scaling-pod-image", consts.DefaultScalingPodImage,
+		"The image to use for the scaling pod, defaults to "+consts.DefaultScalingPodImage)
+	flag.StringVar(&s.ScalingPodNamespace,
+		"scale-adjust-namespace", "kai-scale-adjust",
+		"The namespace to use for the scaling pods, defaults to kai-scale-adjust")
+	flag.StringVar(&s.SchedulerName,
+		"scheduler-name", "kai-scheduler",
+		"Scheduler name, defaults to kai-scheduler")
+	flag.StringVar(&s.ScalingPodAppLabel,
+		"scaling-pod-app-label", "scaling-pod",
+		"Scaling pod app label")
+	flag.StringVar(&s.ScalingPodServiceAccount,
+		"scaling-pod-service-account", "scaling-pod",
+		"Scaling pod service account name")
+	flag.Float64Var(&s.GPUMemoryToFractionRatio,
+		"gpu-memory-to-fraction-ratio", consts.DefaultGPUMemoryToFractionRatio,
+		fmt.Sprintf("The ratio of GPU memory to fraction, defaults to %f", consts.DefaultGPUMemoryToFractionRatio))
+}

cmd/nodescaleadjuster/main.go

@@ -0,0 +1,18 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package main
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/NVIDIA/KAI-scheduler/cmd/nodescaleadjuster/app"
+)
+
+func main() {
+	if err := app.Run(); err != nil {
+		fmt.Printf("Error while running the app: %v", err)
+		os.Exit(1)
+	}
+}

cmd/scalingpod/main.go

@@ -0,0 +1,7 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package main
+
+func main() {
+}

deployments/kai-scheduler/templates/rbac/nodescaleadjuster-binding.yaml

@@ -0,0 +1,15 @@
+# Copyright 2025 NVIDIA CORPORATION
+  # SPDX-License-Identifier: Apache-2.0
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRoleBinding
+metadata:
+  name: kai-node-scale-adjuster
+subjects:
+  - kind: ServiceAccount
+    name: node-scale-adjuster
+    namespace: {{ .Release.Namespace }}
+roleRef:
+  kind: ClusterRole
+  name: kai-node-scale-adjuster
+  apiGroup: rbac.authorization.k8s.io

deployments/kai-scheduler/templates/rbac/nodescaleadjuster.yaml

@@ -0,0 +1,36 @@
+# Copyright 2025 NVIDIA CORPORATION
+# SPDX-License-Identifier: Apache-2.0
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRole
+metadata:
+  name: kai-node-scale-adjuster
+rules:
+- apiGroups:
+  - ""
+  resources:
+  - pods
+  verbs:
+  - create
+  - delete
+  - get
+  - list
+  - patch
+  - update
+  - watch
+- apiGroups:
+  - ""
+  resources:
+  - pods/finalizers
+  verbs:
+  - create
+  - patch
+  - update
+- apiGroups:
+  - ""
+  resources:
+  - pods/status
+  verbs:
+  - get
+  - patch
+  - update