Roman/podgroup controller (#215)

* Added PodGroupController code

* Added PodGroupController deployment

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 scalingpod nodescaleadjuster
+SERVICE_NAMES := podgrouper scheduler binder webhookmanager resourcereservation snapshot-tool scalingpod nodescaleadjuster podgroupcontroller
 
 
 lint: fmt-go vet-go lint-go
@@ -63,6 +63,7 @@ manifests: controller-gen kustomize ## Generate ClusterRole and CustomResourceDe
 	$(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-podgroup-controller,headerFile="./hack/boilerplate.yaml.txt" paths="./pkg/podgroupcontroller/..." paths="./cmd/podgroupcontroller/..." output:stdout > deployments/kai-scheduler/templates/rbac/podgroupcontroller.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/podgroupcontroller/app/app.go

@@ -0,0 +1,124 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package app
+
+import (
+	"flag"
+
+	"github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"
+	"github.com/NVIDIA/KAI-scheduler/pkg/podgroupcontroller/controllers"
+
+	"go.uber.org/zap/zapcore"
+	v1 "k8s.io/api/core/v1"
+	schedulingv1 "k8s.io/api/scheduling/v1"
+	"k8s.io/apimachinery/pkg/fields"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
+	// to ensure that exec-entrypoint and run can make use of them.
+	_ "k8s.io/client-go/plugin/pkg/client/auth"
+
+	"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"
+
+	"sigs.k8s.io/controller-runtime/pkg/cache"
+	// +kubebuilder:scaffold:imports
+)
+
+const (
+	schedulerNameField = "spec.schedulerName"
+)
+
+var (
+	scheme   = runtime.NewScheme()
+	setupLog = ctrl.Log.WithName("setup")
+)
+
+func init() {
+	utilruntime.Must(clientgoscheme.AddToScheme(scheme))
+	utilruntime.Must(v2alpha2.AddToScheme(scheme))
+
+	// +kubebuilder:scaffold:scheme
+}
+
+func Run() error {
+	options := InitOptions()
+	opts := zap.Options{
+		Development: true,
+		TimeEncoder: zapcore.ISO8601TimeEncoder,
+		Level:       zapcore.Level(-1 * options.LogLevel),
+	}
+	opts.BindFlags(flag.CommandLine)
+	flag.Parse()
+	ctrl.SetLogger(zap.New(zap.UseFlagOptions(&opts)))
+
+	clientConfig := ctrl.GetConfigOrDie()
+	clientConfig.QPS = float32(options.Qps)
+	clientConfig.Burst = options.Burst
+
+	schedulerSelector := fields.Set{schedulerNameField: options.SchedulerName}.AsSelector()
+	cacheOptions := cache.Options{}
+	cacheOptions.ByObject = map[client.Object]cache.ByObject{
+		&v1.Pod{}:                     {Field: schedulerSelector},
+		&v1.Node{}:                    {}, // TODO: filter by strict/non-strict runai nodes
+		&schedulingv1.PriorityClass{}: {},
+		&v2alpha2.PodGroup{}:          {},
+	}
+
+	mgr, err := ctrl.NewManager(clientConfig, ctrl.Options{
+		Scheme:                 scheme,
+		Cache:                  cacheOptions,
+		HealthProbeBindAddress: options.ProbeAddr,
+		LeaderElection:         options.EnableLeaderElection,
+		LeaderElectionID:       "3f770c00.run.ai",
+		// LeaderElectionReleaseOnCancel defines if the leader should step down voluntarily
+		// when the Manager ends. This requires the binary to immediately end when the
+		// Manager is stopped, otherwise, this setting is unsafe. Setting this significantly
+		// speeds up voluntary leader transitions as the new leader don't have to wait
+		// LeaseDuration time first.
+		//
+		// In the default scaffold provided, the program ends immediately after
+		// the manager stops, so would be fine to enable this option. However,
+		// if you are doing or is intended to do any operation such as perform cleanups
+		// after the manager stops then its usage might be unsafe.
+		// LeaderElectionReleaseOnCancel: true,
+	})
+	if err != nil {
+		setupLog.Error(err, "unable to start manager")
+		return err
+	}
+
+	configs := controllers.Configs{
+		MaxConcurrentReconciles: options.MaxConcurrentReconciles,
+	}
+	if err = (&controllers.PodGroupReconciler{
+		Client: mgr.GetClient(),
+		Scheme: mgr.GetScheme(),
+	}).SetupWithManager(mgr, configs); err != nil {
+		setupLog.Error(err, "unable to create controller", "controller", "Pod")
+		return err
+	}
+	// +kubebuilder:scaffold:builder
+
+	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/podgroupcontroller/app/options.go

@@ -0,0 +1,43 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package app
+
+import (
+	"flag"
+)
+
+type Options struct {
+	MetricsAddr             string
+	EnableLeaderElection    bool
+	ProbeAddr               string
+	Qps                     int
+	Burst                   int
+	MaxConcurrentReconciles int
+	LogLevel                int
+	SchedulerName           string
+}
+
+func InitOptions() *Options {
+	options := &Options{}
+
+	flag.StringVar(&options.MetricsAddr, "metrics-bind-address", ":8080",
+		"The address the metric endpoint binds to.")
+	flag.StringVar(&options.ProbeAddr, "health-probe-bind-address", ":8081",
+		"The address the probe endpoint binds to.")
+	flag.BoolVar(&options.EnableLeaderElection, "leader-elect", false,
+		"Enable leader election for controller manager. "+
+			"Enabling this will ensure there is only one active controller manager.")
+	flag.IntVar(&options.Qps, "qps", 50,
+		"Queries per second to the K8s API server")
+	flag.IntVar(&options.Burst, "burst", 300,
+		"Burst to the K8s API server")
+	flag.IntVar(&options.MaxConcurrentReconciles, "max-concurrent-reconciles", 10,
+		"Max concurrent reconciles")
+	flag.IntVar(&options.LogLevel, "log-level", 3,
+		"Log level")
+	flag.StringVar(&options.SchedulerName, "scheduler-name", "kai-scheduler",
+		"The name of the scheduler used to schedule pod groups")
+
+	return options
+}

cmd/podgroupcontroller/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/podgroupcontroller/app"
+)
+
+func main() {
+	if err := app.Run(); err != nil {
+		fmt.Printf("Error while running the app: %v", err)
+		os.Exit(1)
+	}
+}

deployments/kai-scheduler/templates/rbac/podgroupcontroller-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-podgroup-controller
+subjects:
+  - kind: ServiceAccount
+    name: podgroup-controller
+    namespace: {{ .Release.Namespace }}
+roleRef:
+  kind: ClusterRole
+  name: kai-podgroup-controller
+  apiGroup: rbac.authorization.k8s.io

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

@@ -0,0 +1,44 @@
+# Copyright 2025 NVIDIA CORPORATION
+# SPDX-License-Identifier: Apache-2.0
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRole
+metadata:
+  name: kai-podgroup-controller
+rules:
+- apiGroups:
+  - ""
+  resources:
+  - nodes
+  - pods
+  - pods/status
+  verbs:
+  - get
+  - list
+  - watch
+- apiGroups:
+  - scheduling.k8s.io
+  resources:
+  - priorityclasses
+  verbs:
+  - get
+  - list
+  - watch
+- apiGroups:
+  - scheduling.run.ai
+  resources:
+  - podgroups
+  verbs:
+  - get
+  - list
+  - watch
+- apiGroups:
+  - scheduling.run.ai
+  resources:
+  - podgroups/status
+  verbs:
+  - get
+  - list
+  - patch
+  - update
+  - watch

deployments/kai-scheduler/templates/services/podgroup-controller-serviceaccount.yaml

@@ -0,0 +1,7 @@
+# Copyright 2025 NVIDIA CORPORATION
+# SPDX-License-Identifier: Apache-2.0
+
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  name: podgroup-controller

deployments/kai-scheduler/templates/services/podgroup-controller.yaml

@@ -0,0 +1,50 @@
+# Copyright 2025 NVIDIA CORPORATION
+# SPDX-License-Identifier: Apache-2.0
+
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: podgroup-controller
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: podgroup-controller
+  template:
+    metadata:
+      labels:
+        app: podgroup-controller
+    spec:
+      serviceAccountName: podgroup-controller
+      containers:
+        - name: podgroup-controller
+          image: "{{ .Values.global.registry }}/{{ .Values.podgroupcontroller.image.name }}:{{ .Chart.Version }}"
+          imagePullPolicy: {{ .Values.podgroupcontroller.image.pullPolicy }}
+          {{- if .Values.podgroupcontroller.additionalArgs }}
+          args:
+            {{- toYaml .Values.podgroupcontroller.additionalArgs | nindent 12 }}
+          {{- end }}
+          {{- if .Values.podgroupcontroller.resources }}
+          resources:
+            {{- toYaml .Values.podgroupcontroller.resources | nindent 12 }}
+          {{- end }}
+          {{- if .Values.global.securityContext }}
+          securityContext:
+            {{- toYaml .Values.global.securityContext | nindent 12 }}
+          {{- end }}
+      {{- if .Values.global.imagePullSecrets }}
+      imagePullSecrets:
+        {{- toYaml .Values.global.imagePullSecrets | nindent 8 }}
+      {{- end }}
+      {{- if .Values.global.nodeSelector }}
+      nodeSelector:
+        {{- toYaml .Values.global.nodeSelector | nindent 8 }}
+      {{- end }}
+      {{- if .Values.global.affinity }}
+      affinity:
+        {{- toYaml .Values.global.affinity | nindent 8 }}
+      {{- end }}
+      {{- if .Values.global.tolerations }}
+      tolerations:
+        {{- toYaml .Values.global.tolerations | nindent 8 }}
+      {{- end }}

deployments/kai-scheduler/values.yaml

@@ -25,6 +25,19 @@ podgrouper:
       cpu: "250m"
       memory: "128Mi"
 
+podgroupcontroller:
+  image:
+    name: podgroupcontroller
+    pullPolicy: IfNotPresent
+  additionalArgs: []
+  resources:
+    limits:
+      cpu: "500m"
+      memory: "256Mi"
+    requests:
+      cpu: "250m"
+      memory: "128Mi"
+
 binder:
   name: binder
   image:

pkg/common/constants/constants.go

@@ -6,6 +6,7 @@ package constants
 const (
 	AppLabelName              = "app"
 	GpuResource               = "nvidia.com/gpu"
+	NvidiaGpuMemory           = "nvidia.com/gpu.memory"
 	UnlimitedResourceQuantity = float64(-1)
 	DefaultQueuePriority      = 100
 	DefaultNodePoolName       = "default"