Skip to content

Commit 9e6f7d9

Browse files
JunAr7112JunAr7112
authored
Merge pull request #866 from JunAr7112/priority_class
Add priorityClass to NIMService
2 parents 363d80d + 8be2434 commit 9e6f7d9

13 files changed

Lines changed: 105 additions & 6 deletions

‎api/apps/v1alpha1/nimservice_types.go‎

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -125,12 +125,15 @@ type NIMServiceSpec struct {
125125
SchedulerName string `json:"schedulerName,omitempty"`
126126
Metrics Metrics `json:"metrics,omitempty"`
127127
// +kubebuilder:validation:Minimum=0
128-
Replicas *int32 `json:"replicas,omitempty"`
129-
UserID *int64 `json:"userID,omitempty"`
130-
GroupID *int64 `json:"groupID,omitempty"`
131-
RuntimeClassName string `json:"runtimeClassName,omitempty"`
132-
Proxy *ProxySpec `json:"proxy,omitempty"`
133-
MultiNode *NimServiceMultiNodeConfig `json:"multiNode,omitempty"`
128+
Replicas *int32 `json:"replicas,omitempty"`
129+
UserID *int64 `json:"userID,omitempty"`
130+
GroupID *int64 `json:"groupID,omitempty"`
131+
RuntimeClassName string `json:"runtimeClassName,omitempty"`
132+
// PriorityClassName is the name of the PriorityClass to assign to pods created by this NIMService.
133+
// If set, the Kubernetes scheduler uses the priority for preemption decisions.
134+
PriorityClassName string `json:"priorityClassName,omitempty"`
135+
Proxy *ProxySpec `json:"proxy,omitempty"`
136+
MultiNode *NimServiceMultiNodeConfig `json:"multiNode,omitempty"`
134137
// InferencePlatform specifies the inference platform to use for this NIMService.
135138
// Valid values are "standalone" (default) and "kserve".
136139
// +kubebuilder:validation:Enum=standalone;kserve
@@ -1249,6 +1252,9 @@ func (n *NIMService) GetDeploymentParams() *rendertypes.DeploymentParams {
12491252
// Set scheduler
12501253
params.SchedulerName = n.GetSchedulerName()
12511254

1255+
// Set priority class
1256+
params.PriorityClassName = n.GetPriorityClassName()
1257+
12521258
// Setup container ports for nimservice
12531259
params.Ports = []corev1.ContainerPort{
12541260
{
@@ -1333,6 +1339,7 @@ func (n *NIMService) GetLWSParams() *rendertypes.LeaderWorkerSetParams {
13331339
params.ServiceAccountName = n.GetServiceAccountName()
13341340
params.SchedulerName = n.GetSchedulerName()
13351341
params.RuntimeClassName = n.GetRuntimeClassName()
1342+
params.PriorityClassName = n.GetPriorityClassName()
13361343

13371344
params.InitContainers = n.GetInitContainers()
13381345
for idx := range params.InitContainers {
@@ -1350,6 +1357,11 @@ func (n *NIMService) GetSchedulerName() string {
13501357
return n.Spec.SchedulerName
13511358
}
13521359

1360+
// GetPriorityClassName returns the priority class for the NIMService deployment.
1361+
func (n *NIMService) GetPriorityClassName() string {
1362+
return n.Spec.PriorityClassName
1363+
}
1364+
13531365
// GetStatefulSetParams returns params to render StatefulSet from templates.
13541366
func (n *NIMService) GetStatefulSetParams() *rendertypes.StatefulSetParams {
13551367

@@ -1393,6 +1405,9 @@ func (n *NIMService) GetStatefulSetParams() *rendertypes.StatefulSetParams {
13931405

13941406
// Set runtime class
13951407
params.RuntimeClassName = n.GetRuntimeClassName()
1408+
1409+
// Set priority class
1410+
params.PriorityClassName = n.GetPriorityClassName()
13961411
return params
13971412
}
13981413

@@ -1846,6 +1861,9 @@ func (n *NIMService) GetInferenceServiceParams(
18461861
// Set scheduler
18471862
params.SchedulerName = n.GetSchedulerName()
18481863

1864+
// Set priority class
1865+
params.PriorityClassName = n.GetPriorityClassName()
1866+
18491867
params.Ports = n.GetInferenceServicePorts(deploymentMode)
18501868

18511869
params.InitContainers = n.GetInitContainers()

‎api/apps/v1alpha1/nimservice_types_test.go‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -999,3 +999,35 @@ func TestStandardModeReplicasWithoutAutoscaling(t *testing.T) {
999999
t.Errorf("MaxReplicas = %d, want 3", *params.MaxReplicas)
10001000
}
10011001
}
1002+
1003+
func TestGetPriorityClassName(t *testing.T) {
1004+
nimService := &NIMService{
1005+
Spec: NIMServiceSpec{
1006+
PriorityClassName: "gpu-critical",
1007+
Image: Image{
1008+
Repository: "test-repo",
1009+
Tag: "test-tag",
1010+
},
1011+
AuthSecret: "test-secret",
1012+
Expose: Expose{
1013+
Service: Service{
1014+
Port: ptr.To[int32](8000),
1015+
},
1016+
},
1017+
},
1018+
}
1019+
1020+
if got := nimService.GetPriorityClassName(); got != "gpu-critical" {
1021+
t.Errorf("GetPriorityClassName() = %q, want %q", got, "gpu-critical")
1022+
}
1023+
1024+
deploymentParams := nimService.GetDeploymentParams()
1025+
if deploymentParams.PriorityClassName != "gpu-critical" {
1026+
t.Errorf("GetDeploymentParams().PriorityClassName = %q, want %q", deploymentParams.PriorityClassName, "gpu-critical")
1027+
}
1028+
1029+
inferenceParams := nimService.GetInferenceServiceParams("Standard")
1030+
if inferenceParams.PriorityClassName != "gpu-critical" {
1031+
t.Errorf("GetInferenceServiceParams().PriorityClassName = %q, want %q", inferenceParams.PriorityClassName, "gpu-critical")
1032+
}
1033+
}

‎bundle/manifests/apps.nvidia.com_nimpipelines.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3749,6 +3749,11 @@ spec:
37493749
type: array
37503750
x-kubernetes-list-type: atomic
37513751
type: object
3752+
priorityClassName:
3753+
description: |-
3754+
PriorityClassName is the name of the PriorityClass to assign to pods created by this NIMService.
3755+
If set, the Kubernetes scheduler uses the priority for preemption decisions.
3756+
type: string
37523757
proxy:
37533758
description: ProxySpec defines the proxy configuration for
37543759
NIMService.

‎bundle/manifests/apps.nvidia.com_nimservices.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3627,6 +3627,11 @@ spec:
36273627
type: array
36283628
x-kubernetes-list-type: atomic
36293629
type: object
3630+
priorityClassName:
3631+
description: |-
3632+
PriorityClassName is the name of the PriorityClass to assign to pods created by this NIMService.
3633+
If set, the Kubernetes scheduler uses the priority for preemption decisions.
3634+
type: string
36303635
proxy:
36313636
description: ProxySpec defines the proxy configuration for NIMService.
36323637
properties:

‎config/crd/bases/apps.nvidia.com_nimpipelines.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3749,6 +3749,11 @@ spec:
37493749
type: array
37503750
x-kubernetes-list-type: atomic
37513751
type: object
3752+
priorityClassName:
3753+
description: |-
3754+
PriorityClassName is the name of the PriorityClass to assign to pods created by this NIMService.
3755+
If set, the Kubernetes scheduler uses the priority for preemption decisions.
3756+
type: string
37523757
proxy:
37533758
description: ProxySpec defines the proxy configuration for
37543759
NIMService.

‎config/crd/bases/apps.nvidia.com_nimservices.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3627,6 +3627,11 @@ spec:
36273627
type: array
36283628
x-kubernetes-list-type: atomic
36293629
type: object
3630+
priorityClassName:
3631+
description: |-
3632+
PriorityClassName is the name of the PriorityClass to assign to pods created by this NIMService.
3633+
If set, the Kubernetes scheduler uses the priority for preemption decisions.
3634+
type: string
36303635
proxy:
36313636
description: ProxySpec defines the proxy configuration for NIMService.
36323637
properties:

‎deployments/helm/k8s-nim-operator/crds/apps.nvidia.com_nimpipelines.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3749,6 +3749,11 @@ spec:
37493749
type: array
37503750
x-kubernetes-list-type: atomic
37513751
type: object
3752+
priorityClassName:
3753+
description: |-
3754+
PriorityClassName is the name of the PriorityClass to assign to pods created by this NIMService.
3755+
If set, the Kubernetes scheduler uses the priority for preemption decisions.
3756+
type: string
37523757
proxy:
37533758
description: ProxySpec defines the proxy configuration for
37543759
NIMService.

‎deployments/helm/k8s-nim-operator/crds/apps.nvidia.com_nimservices.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3627,6 +3627,11 @@ spec:
36273627
type: array
36283628
x-kubernetes-list-type: atomic
36293629
type: object
3630+
priorityClassName:
3631+
description: |-
3632+
PriorityClassName is the name of the PriorityClass to assign to pods created by this NIMService.
3633+
If set, the Kubernetes scheduler uses the priority for preemption decisions.
3634+
type: string
36303635
proxy:
36313636
description: ProxySpec defines the proxy configuration for NIMService.
36323637
properties:

‎internal/render/types/types.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ type DeploymentParams struct {
6868
ImagePullSecrets []string
6969
ImagePullPolicy string
7070
SchedulerName string
71+
PriorityClassName string
7172
Volumes []corev1.Volume
7273
VolumeMounts []corev1.VolumeMount
7374
Env []corev1.EnvVar
@@ -109,6 +110,7 @@ type LeaderWorkerSetParams struct {
109110
ImagePullSecrets []string
110111
ImagePullPolicy string
111112
SchedulerName string
113+
PriorityClassName string
112114
WorkerVolumes []corev1.Volume
113115
LeaderVolumes []corev1.Volume
114116
WorkerVolumeMounts []corev1.VolumeMount
@@ -164,6 +166,7 @@ type StatefulSetParams struct {
164166
NIMCachePVC string
165167
RuntimeClassName string
166168
OrchestratorType string
169+
PriorityClassName string
167170
InitContainers []corev1.Container
168171
}
169172

@@ -317,6 +320,7 @@ type InferenceServiceParams struct {
317320
ImagePullSecrets []string
318321
ImagePullPolicy string
319322
SchedulerName string
323+
PriorityClassName string
320324
Volumes []corev1.Volume
321325
VolumeMounts []corev1.VolumeMount
322326
Env []corev1.EnvVar

‎manifests/deployment.yaml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,9 @@ spec:
3838
{{- if .SchedulerName }}
3939
schedulerName: {{ .SchedulerName }}
4040
{{- end }}
41+
{{- if .PriorityClassName }}
42+
priorityClassName: {{ .PriorityClassName }}
43+
{{- end }}
4144
{{- if .Affinity }}
4245
affinity:
4346
{{ .Affinity | yaml | nindent 8 }}

0 commit comments

Comments
 (0)