feat(admission): Explicitly apply 'nvidia' runtimeClass to GPU pods (v0.9) (#625)

omeryahud · web-flow · commit 6ffb4252969d · 2025-11-10T12:46:56.000+01:00
Signed-off-by: Omer Yahud &lt;oyahud@nvidia.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -13,6 +13,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Added
 - Support DRA in kubernetes 1.34
+- Added enforcement of the `nvidia` runtime class for GPU pods, with the option to enforce a custom runtime class, or disable enforcement entirely.
 
 ### Fixed
 - Fixed a bug where the scheduler would not re-try updating podgroup status after failure
diff --git a/README.md b/README.md
@@ -49,6 +49,13 @@ helm upgrade -i kai-scheduler oci://ghcr.io/nvidia/kai-scheduler/kai-scheduler -
 #### Build from Source
 Follow the instructions [here](docs/developer/building-from-source.md)
 
+## Flavor Specific Instructions
+### Openshift
+When `gpu-operator` <v25.10.0 is installed, the following flag should be added to the installation command:
+```
+--set admission.gpuPodRuntimeClassName=null
+```
+
 ## Quick Start
 To start scheduling workloads with KAI Scheduler, please continue to [Quick Start example](docs/quickstart/README.md)
 
diff --git a/cmd/admission/app/app.go b/cmd/admission/app/app.go
@@ -41,7 +41,6 @@ var (
 
 func init() {
 	utilruntime.Must(clientgoscheme.AddToScheme(scheme))
-
 	utilruntime.Must(schedulingv1alpha2.AddToScheme(scheme))
 	// +kubebuilder:scaffold:scheme
 }
@@ -174,5 +173,6 @@ func (app *App) Run() error {
 		setupLog.Error(err, "problem running manager")
 		return err
 	}
+
 	return nil
 }
diff --git a/cmd/admission/app/options.go b/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,9 @@ 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) Set to empty string to disable", constants.DefaultRuntimeClassName))
 
 	utilfeature.DefaultMutableFeatureGate.AddFlag(fs)
 
diff --git a/cmd/admission/main.go b/cmd/admission/main.go
@@ -12,6 +12,7 @@ import (
 
 	"github.com/NVIDIA/KAI-scheduler/pkg/admission/plugins"
 	"github.com/NVIDIA/KAI-scheduler/pkg/admission/webhook/v1alpha2/gpusharing"
+	"github.com/NVIDIA/KAI-scheduler/pkg/admission/webhook/v1alpha2/runtimeenforcement"
 )
 
 var (
@@ -42,8 +43,13 @@ func registerPlugins(app *app.App) error {
 	admissionPlugins := plugins.New()
 
 	admissionGpuSharingPlugin := gpusharing.New(app.Client, app.Options.GPUSharingEnabled)
-
 	admissionPlugins.RegisterPlugin(admissionGpuSharingPlugin)
+
+	if app.Options.GPUPodRuntimeClassName != "" {
+		admissionRuntimeEnforcementPlugin := runtimeenforcement.New(app.Options.GPUPodRuntimeClassName)
+		admissionPlugins.RegisterPlugin(admissionRuntimeEnforcementPlugin)
+	}
+
 	app.RegisterPlugins(admissionPlugins)
 	return nil
 }
diff --git a/deployments/kai-scheduler/crds/kai.scheduler_configs.yaml b/deployments/kai-scheduler/crds/kai.scheduler_configs.yaml
@@ -44,6 +44,11 @@ spec:
               admission:
                 description: Admission holds KAI admission webhooks
                 properties:
+                  gpuPodRuntimeClassName:
+                    description: |-
+                      GPUPodRuntimeClassName specifies the runtime class to be set for GPU pods
+                      set to empty string to disable
+                    type: string
                   gpuSharing:
                     description: GPUSharing enables GPU sharing functionality for
                       the admission service
diff --git a/deployments/kai-scheduler/templates/kai-config.yaml b/deployments/kai-scheduler/templates/kai-config.yaml
@@ -85,6 +85,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:
diff --git a/deployments/kai-scheduler/values.yaml b/deployments/kai-scheduler/values.yaml
@@ -77,6 +77,7 @@ admission:
     metricsPort: 8080
     probePort: 8081
   cdi: false
+  gpuPodRuntimeClassName: nvidia
 
 nodescaleadjuster:
   image:
diff --git a/go.mod b/go.mod
@@ -16,6 +16,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/client_golang v1.23.2
 	github.com/prometheus/common v0.66.1
diff --git a/go.sum b/go.sum
@@ -226,6 +226,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=
diff --git a/pkg/admission/webhook/v1alpha2/runtimeenforcement/runtime_enforcement.go b/pkg/admission/webhook/v1alpha2/runtimeenforcement/runtime_enforcement.go
@@ -0,0 +1,52 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package runtimeenforcement
+
+import (
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/utils/ptr"
+
+	"github.com/NVIDIA/KAI-scheduler/pkg/binder/binding/resourcereservation"
+	"github.com/NVIDIA/KAI-scheduler/pkg/common/resources"
+)
+
+type RuntimeEnforcement struct {
+	gpuPodRuntimeClassName string
+}
+
+func New(gpuPodRuntimeClassName string) *RuntimeEnforcement {
+	return &RuntimeEnforcement{
+		gpuPodRuntimeClassName: gpuPodRuntimeClassName,
+	}
+}
+
+func (p *RuntimeEnforcement) Name() string {
+	return "runtimeenforcement"
+}
+
+func (p *RuntimeEnforcement) Validate(pod *v1.Pod) error {
+	return nil
+}
+
+func (p *RuntimeEnforcement) Mutate(pod *v1.Pod) error {
+	// in order to no collide with custom reservation pods runtimeClass
+	if resourcereservation.IsGPUReservationPod(pod) {
+		return nil
+	}
+
+	if !resources.RequestsGPU(pod) {
+		return nil
+	}
+
+	if pod.Spec.RuntimeClassName == nil || *pod.Spec.RuntimeClassName == "" {
+		setRuntimeClass(pod, p.gpuPodRuntimeClassName)
+		return nil
+	}
+
+	return nil
+}
+
+func setRuntimeClass(pod *v1.Pod, runtimeClassName string) {
+	pod.Spec.RuntimeClassName = ptr.To(runtimeClassName)
+}
diff --git a/pkg/admission/webhook/v1alpha2/runtimeenforcement/runtime_enforcement_test.go b/pkg/admission/webhook/v1alpha2/runtimeenforcement/runtime_enforcement_test.go
@@ -0,0 +1,128 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package runtimeenforcement
+
+import (
+	"testing"
+
+	"github.com/NVIDIA/KAI-scheduler/pkg/apis/client/clientset/versioned/scheme"
+	"github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
+	ocpconf "github.com/openshift/api/config/v1"
+	"github.com/stretchr/testify/assert"
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
+	"k8s.io/utils/ptr"
+)
+
+func TestMutate(t *testing.T) {
+	sc := scheme.Scheme
+	utilruntime.Must(ocpconf.AddToScheme(sc))
+
+	tests := []struct {
+		name                   string
+		gpuPodRuntimeClassName string
+		incomingPod            *v1.Pod
+		expectedOutboundPod    *v1.Pod
+		expectedError          error
+	}{
+		{
+			name:                   "pod without GPU requests",
+			gpuPodRuntimeClassName: constants.DefaultRuntimeClassName,
+			incomingPod:            &v1.Pod{},
+			expectedOutboundPod:    &v1.Pod{},
+			expectedError:          nil,
+		},
+		{
+			name:                   "pod with a fractional GPU request",
+			gpuPodRuntimeClassName: constants.DefaultRuntimeClassName,
+			incomingPod: &v1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{constants.GpuFraction: "0.5"},
+				},
+			},
+			expectedOutboundPod: &v1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{constants.GpuFraction: "0.5"},
+				},
+				Spec: v1.PodSpec{
+					RuntimeClassName: ptr.To(constants.DefaultRuntimeClassName),
+				},
+			},
+			expectedError: nil,
+		},
+		{
+			name:                   "pod with a whole GPU request",
+			gpuPodRuntimeClassName: constants.DefaultRuntimeClassName,
+			incomingPod: &v1.Pod{
+				Spec: v1.PodSpec{
+					Containers: []v1.Container{
+						{
+							Resources: v1.ResourceRequirements{
+								Limits: v1.ResourceList{constants.GpuResource: resource.MustParse("1")},
+							},
+						},
+					},
+				},
+			},
+			expectedOutboundPod: &v1.Pod{
+				Spec: v1.PodSpec{
+					RuntimeClassName: ptr.To(constants.DefaultRuntimeClassName),
+					Containers: []v1.Container{
+						{
+							Resources: v1.ResourceRequirements{
+								Limits: v1.ResourceList{constants.GpuResource: resource.MustParse("1")},
+							},
+						},
+					},
+				},
+			},
+			expectedError: nil,
+		},
+		{
+			name:                   "pod with GPU request and runtimeClassName set",
+			gpuPodRuntimeClassName: constants.DefaultRuntimeClassName,
+			incomingPod: &v1.Pod{
+				Spec: v1.PodSpec{
+					RuntimeClassName: ptr.To("custom-runtime"),
+					Containers: []v1.Container{
+						{
+							Resources: v1.ResourceRequirements{
+								Limits: v1.ResourceList{constants.GpuResource: resource.MustParse("1")},
+							},
+						},
+					},
+				},
+			},
+			expectedOutboundPod: &v1.Pod{
+				Spec: v1.PodSpec{
+					RuntimeClassName: ptr.To("custom-runtime"),
+					Containers: []v1.Container{
+						{
+							Resources: v1.ResourceRequirements{
+								Limits: v1.ResourceList{constants.GpuResource: resource.MustParse("1")},
+							},
+						},
+					},
+				},
+			},
+			expectedError: nil,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			p := New(tt.gpuPodRuntimeClassName)
+			err := p.Mutate(tt.incomingPod)
+			if tt.expectedError != nil {
+				assert.Error(t, err)
+				assert.Equal(t, tt.expectedError.Error(), err.Error())
+			} else {
+				assert.NoError(t, err)
+			}
+			assert.Equal(t, tt.expectedOutboundPod, tt.incomingPod)
+		})
+	}
+}
diff --git a/pkg/apis/kai/v1/admission/admission.go b/pkg/apis/kai/v1/admission/admission.go
@@ -8,6 +8,7 @@ import (
 	"k8s.io/utils/ptr"
 
 	"github.com/NVIDIA/KAI-scheduler/pkg/apis/kai/v1/common"
+	"github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
 )
 
 const (
@@ -32,6 +33,11 @@ type Admission struct {
 	// QueueLabelSelector enables the queue label MatchExpression in webhooks
 	// +kubebuilder:validation:Optional
 	QueueLabelSelector *bool `json:"queueLabelSelector,omitempty"`
+
+	// GPUPodRuntimeClassName specifies the runtime class to be set for GPU pods
+	// set to empty string to disable
+	// +kubebuilder:validation:Optional
+	GPUPodRuntimeClassName *string `json:"gpuPodRuntimeClassName,omitempty"`
 }
 
 func (b *Admission) SetDefaultsWhereNeeded(replicaCount *int32) {
@@ -47,6 +53,8 @@ func (b *Admission) SetDefaultsWhereNeeded(replicaCount *int32) {
 	b.Replicas = common.SetDefault(b.Replicas, ptr.To(ptr.Deref(replicaCount, 1)))
 	b.GPUSharing = common.SetDefault(b.GPUSharing, ptr.To(false))
 	b.QueueLabelSelector = common.SetDefault(b.QueueLabelSelector, ptr.To(false))
+
+	b.GPUPodRuntimeClassName = common.SetDefault(b.GPUPodRuntimeClassName, ptr.To(constants.DefaultRuntimeClassName))
 }
 
 // Webhook defines configuration for the admission webhook
diff --git a/pkg/apis/kai/v1/admission/admission_test.go b/pkg/apis/kai/v1/admission/admission_test.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 )
@@ -25,6 +26,7 @@ var _ = Describe("Admission", func() {
 		Expect(*Admission.Service.Enabled).To(Equal(true))
 		Expect(*Admission.Service.Image.Name).To(Equal("admission"))
 		Expect(*Admission.Replicas).To(Equal(int32(1)))
+		Expect(*Admission.GPUPodRuntimeClassName).To(Equal(constants.DefaultRuntimeClassName))
 	})
 	It("Set Defaults with replica count", func(ctx context.Context) {
 		Admission := &Admission{}
diff --git a/pkg/apis/kai/v1/admission/zz_generated.deepcopy.go b/pkg/apis/kai/v1/admission/zz_generated.deepcopy.go
diff --git a/pkg/binder/binding/resourcereservation/resource_reservation.go b/pkg/binder/binding/resourcereservation/resource_reservation.go
@@ -527,3 +527,7 @@ func (rsc *service) isScalingUp(ctx context.Context) bool {
 	}
 	return false
 }
+
+func IsGPUReservationPod(pod *v1.Pod) bool {
+	return strings.HasPrefix(pod.Name, gpuReservationPodPrefix)
+}
diff --git a/pkg/common/resources/gpu_sharing.go b/pkg/common/resources/gpu_sharing.go
diff --git a/pkg/operator/operands/admission/resources.go b/pkg/operator/operands/admission/resources.go
diff --git a/pkg/operator/operands/admission/resources_test.go b/pkg/operator/operands/admission/resources_test.go

Original file line number	Diff line number	Diff line change
`@@ -41,7 +41,6 @@ var (`
`41`	`41`
`42`	`42`	`func init() {`
`43`	`43`	`utilruntime.Must(clientgoscheme.AddToScheme(scheme))`
`44`		`-`
`45`	`44`	`utilruntime.Must(schedulingv1alpha2.AddToScheme(scheme))`
`46`	`45`	`// +kubebuilder:scaffold:scheme`
`47`	`46`	`}`
`@@ -174,5 +173,6 @@ func (app *App) Run() error {`
`174`	`173`	`setupLog.Error(err, "problem running manager")`
`175`	`174`	`return err`
`176`	`175`	`}`
	`176`	`+`
`177`	`177`	`return nil`
`178`	`178`	`}`