s

Signed-off-by: Omer Yahud <[email protected]>

Author: omeryahud

cmd/admission/app/app.go

@@ -13,6 +13,7 @@ import (
 	// to ensure that exec-entrypoint and run can make use of them.
 	_ "k8s.io/client-go/plugin/pkg/client/auth"
 
+	ocpconf "github.com/openshift/api/config/v1"
 	"github.com/spf13/pflag"
 	"go.uber.org/zap/zapcore"
 	corev1 "k8s.io/api/core/v1"
@@ -41,8 +42,8 @@ var (
 
 func init() {
 	utilruntime.Must(clientgoscheme.AddToScheme(scheme))
-
 	utilruntime.Must(schedulingv1alpha2.AddToScheme(scheme))
+	utilruntime.Must(ocpconf.AddToScheme(scheme))
 	// +kubebuilder:scaffold:scheme
 }
 
@@ -174,5 +175,6 @@ func (app *App) Run() error {
 		setupLog.Error(err, "problem running manager")
 		return err
 	}
+
 	return nil
 }

cmd/admission/app/options.go

@@ -4,6 +4,8 @@
 package app
 
 import (
+	"fmt"
+
 	"github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
 	"github.com/spf13/pflag"
 
@@ -22,6 +24,7 @@ type Options struct {
 	WebhookPort                 int
 	FakeGPUNodes                bool
 	GPUSharingEnabled           bool
+	GPUPodRuntimeClassName      string
 }
 
 func InitOptions() *Options {
@@ -63,6 +66,10 @@ func InitOptions() *Options {
 	fs.BoolVar(&options.GPUSharingEnabled,
 		"gpu-sharing-enabled", false,
 		"Specifies if the GPU sharing is enabled")
+	fs.StringVar(&options.GPUPodRuntimeClassName,
+		"gpu-pod-runtime-class-name", constants.DefaultRuntimeClassName,
+		fmt.Sprintf("Runtime class to be set for GPU pods (defaults to %s), ignored in Openshift deployments."+
+			"Set to empty string to disable", constants.DefaultRuntimeClassName))
 
 	utilfeature.DefaultMutableFeatureGate.AddFlag(fs)
 

cmd/admission/main.go

@@ -43,10 +43,13 @@ func registerPlugins(app *app.App) error {
 	admissionPlugins := plugins.New()
 
 	admissionGpuSharingPlugin := gpusharing.New(app.Client, app.Options.GPUSharingEnabled)
-	admissionRuntimeEnforcementPlugin := runtimeenforcement.New(app.Client)
-
 	admissionPlugins.RegisterPlugin(admissionGpuSharingPlugin)
-	admissionPlugins.RegisterPlugin(admissionRuntimeEnforcementPlugin)
+
+	if app.Options.GPUPodRuntimeClassName != "" {
+		admissionRuntimeEnforcementPlugin := runtimeenforcement.New(app.Client, app.Options.GPUPodRuntimeClassName)
+		admissionPlugins.RegisterPlugin(admissionRuntimeEnforcementPlugin)
+	}
+
 	app.RegisterPlugins(admissionPlugins)
 	return nil
 }

deployments/kai-scheduler/templates/kai-config.yaml

@@ -105,6 +105,7 @@ spec:
       targetPort: {{ .Values.admission.ports.webhookPort | default 9443 }}
       probePort: {{ .Values.admission.ports.probePort | default 8081 }}
       metricsPort: {{ .Values.admission.ports.metricsPort | default 8080 }}
+    gpuPodRuntimeClassName: {{ .Values.admission.gpuPodRuntimeClassName | default "nvidia" }}
 
   nodeScaleAdjuster:
     service:

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

@@ -38,3 +38,19 @@ rules:
   - create
   - patch
   - update
+- apiGroups:
+  - node.k8s.io
+  resources:
+  - runtimeclasses
+  verbs:
+  - get
+  - list
+  - watch
+- apiGroups:
+  - config.openshift.io
+  resources:
+  - clusterversions
+  verbs:
+  - get
+  - list
+  - watch

deployments/kai-scheduler/values.yaml

@@ -76,6 +76,7 @@ admission:
     metricsPort: 8080
     probePort: 8081
   cdi: false
+  gpuPodRuntimeClassName: nvidia
 
 nodescaleadjuster:
   image:

go.mod

@@ -18,6 +18,7 @@ require (
 	github.com/kubeflow/training-operator v1.9.3
 	github.com/onsi/ginkgo/v2 v2.25.3
 	github.com/onsi/gomega v1.38.2
+	github.com/openshift/api v0.0.0-20250602203052-b29811a290c7
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring v0.81.0
 	github.com/prometheus/client_golang v1.23.2

go.sum

@@ -246,6 +246,8 @@ github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8
 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
 github.com/opencontainers/selinux v1.11.1 h1:nHFvthhM0qY8/m+vfhJylliSshm8G1jJ2jDMcgULaH8=
 github.com/opencontainers/selinux v1.11.1/go.mod h1:E5dMC3VPuVvVHDYmi78qvhJp8+M586T4DlDRYpFkyec=
+github.com/openshift/api v0.0.0-20250602203052-b29811a290c7 h1:dZ9uBd0Cw3+l1RGpYRkWdrRjM9yvfxrjW/uPHKUwtIQ=
+github.com/openshift/api v0.0.0-20250602203052-b29811a290c7/go.mod h1:yk60tHAmHhtVpJQo3TwVYq2zpuP70iJIFDCmeKMIzPw=
 github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
 github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pelletier/go-toml/v2 v2.2.3 h1:YmeHyLY8mFWbdkNWwpr+qIL2bEqT0o95WSdkNHvL12M=

pkg/admission/webhook/v1alpha2/runtimeenforcement/__debug_bin2990338812

@@ -5,24 +5,26 @@ package runtimeenforcement
 
 import (
 	"context"
+	"fmt"
 
 	v1 "k8s.io/api/core/v1"
+	"k8s.io/utils/ptr"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	"github.com/NVIDIA/KAI-scheduler/pkg/binder/binding/resourcereservation"
-	"github.com/NVIDIA/KAI-scheduler/pkg/binder/common"
-	"github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
 	"github.com/NVIDIA/KAI-scheduler/pkg/common/k8s_utils"
 	"github.com/NVIDIA/KAI-scheduler/pkg/common/resources"
 )
 
 type RuntimeEnforcement struct {
-	kubeClient client.Client
+	kubeClient             client.Client
+	gpuPodRuntimeClassName string
 }
 
-func New(kubeClient client.Client) *RuntimeEnforcement {
+func New(kubeClient client.Client, gpuPodRuntimeClassName string) *RuntimeEnforcement {
 	return &RuntimeEnforcement{
-		kubeClient: kubeClient,
+		kubeClient:             kubeClient,
+		gpuPodRuntimeClassName: gpuPodRuntimeClassName,
 	}
 }
 
@@ -50,15 +52,23 @@ func (p *RuntimeEnforcement) Mutate(pod *v1.Pod) error {
 
 	if resources.RequestsGPU(pod) {
 		exists, err := k8s_utils.RuntimeClassExists(context.Background(),
-			p.kubeClient, constants.DefaultRuntimeClassName)
+			p.kubeClient, p.gpuPodRuntimeClassName)
 		if err != nil {
 			return err
 		} else if !exists {
-			return nil
+			return runtimeClassDoesNotExistError(p.gpuPodRuntimeClassName)
 		}
 
-		common.SetNVIDIARuntimeClass(pod)
+		setRuntimeClass(pod, p.gpuPodRuntimeClassName)
 	}
 
 	return nil
 }
+
+func setRuntimeClass(pod *v1.Pod, runtimeClassName string) {
+	pod.Spec.RuntimeClassName = ptr.To(runtimeClassName)
+}
+
+func runtimeClassDoesNotExistError(runtimeClassName string) error {
+	return fmt.Errorf("cannot set runtimeClassName: runtimeClass '%s' does not exist", runtimeClassName)
+}