Repository: substratusai/kubeai
Pull request #506 (2025-04-21 20:07:11)

Add priorityClassName to the model spec and pass to generated pods
Enable use of  Kubernetes [Pod Priority and Preemption](https://kubernetes.io/docs/concepts/scheduling-eviction/pod-priority-preemption/) to configure your models with different priority levels. This is useful when you have limited resources and want to ensure that high-priority models can preempt lower-priority models when necessary.

liebman: Just rebased on latest main.
samos123: Thank you! PR looks great to me. Will give @nstogner a chance to review too.
nstogner: I would prefer to avoid exposing Pod spec options on a field-by-field basis as we will trend towards replicating the entire Pod spec in the Model spec. Instead, I would prefer to support patching via the Model spec. Similar to what we do at the system config-level: https://github.com/substratusai/kubeai/blob/d6e393ca76f11da76b3a6db74b737b94d1a4f057/internal/config/system.go#L258
liebman: This won't work either as an int or string won't be accepted as an `object` right?
```
	// Value is the value to be used in the operation.
	// Required for "add", "replace", and "test" operations.
	// +kubebuilder:validation:Optional
	// +kubebuilder:validation:Schemaless
	// +kubebuilder:validation:Type=object
	Value interface{} `json:"value,omitempty"`
```

samos123: Thanks for giving it a quick try. Let's keep this PR open because I do see value in exposing priorityClassName in a more straight forward way with defaulting to it not being set.

Patch:
diff --git a/api/k8s/v1/model_types.go b/api/k8s/v1/model_types.go
index 4c03044..e09dc85 100644
--- a/api/k8s/v1/model_types.go
+++ b/api/k8s/v1/model_types.go
@@ -134,6 +134,12 @@ type ModelSpec struct {
 	// Files to be mounted in the model Pods.
 	// +kubebuilder:validation:MaxItems=10
 	Files []File `json:"files,omitempty"`
+
+	// PriorityClassName sets the priority class for all pods created for this model.
+	// If specified, the PriorityClass must exist before the model is created.
+	// This is useful for implementing priority and preemption for models.
+	// +kubebuilder:validation:Optional
+	PriorityClassName string `json:"priorityClassName,omitempty"`
 }
 
 // +kubebuilder:validation:Enum=TextGeneration;TextEmbedding;SpeechToText
diff --git a/charts/kubeai/Chart.yaml b/charts/kubeai/Chart.yaml
index 08d2fe5..6e594dd 100644
--- a/charts/kubeai/Chart.yaml
+++ b/charts/kubeai/Chart.yaml
@@ -21,7 +21,7 @@ dependencies:
   - name: open-webui
     condition: open-webui.enabled
     repository: https://helm.openwebui.com/
-    version: 5.4.0
+    version: 6.4.0
 
 keywords: ["LLM", "AI"]
 
diff --git a/charts/kubeai/templates/crds/kubeai.org_models.yaml b/charts/kubeai/templates/crds/kubeai.org_models.yaml
index 3c7c1cf..98a30b2 100644
--- a/charts/kubeai/templates/crds/kubeai.org_models.yaml
+++ b/charts/kubeai/templates/crds/kubeai.org_models.yaml
@@ -243,6 +243,12 @@ spec:
                   OpenAI /v1/models endpoint.
                   DEPRECATED.
                 type: string
+              priorityClassName:
+                description: |-
+                  PriorityClassName sets the priority class for all pods created for this model.
+                  If specified, the PriorityClass must exist before the model is created.
+                  This is useful for implementing priority and preemption for models.
+                type: string
               replicas:
                 description: |-
                   Replicas is the number of Pod replicas that should be actively
diff --git a/docs/how-to/install-models.md b/docs/how-to/install-models.md
index 15cad9b..0dea3ee 100644
--- a/docs/how-to/install-models.md
+++ b/docs/how-to/install-models.md
@@ -102,6 +102,66 @@ See the [examples](https://github.com/substratusai/kubeai/tree/main/examples/k8s
 
 You can inference a model by calling the KubeAI OpenAI compatible API. The model name should match the KubeAI model name.
 
+## Using Pod Priority Classes for Model Preemption
+
+You can use Kubernetes [Pod Priority and Preemption](https://kubernetes.io/docs/concepts/scheduling-eviction/pod-priority-preemption/) to configure your models with different priority levels. This is useful when you have limited resources and want to ensure that high-priority models can preempt lower-priority models when necessary.
+
+First, create the necessary PriorityClasses in your cluster:
+
+```yaml
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: high-priority
+value: 1000000  # Higher value means higher priority
+globalDefault: false
+description: "This priority class should be used for critical inference models only."
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: medium-priority
+value: 100000
+globalDefault: false
+description: "This priority class should be used for medium priority inference models."
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: low-priority
+value: 10000
+globalDefault: false
+description: "This priority class should be used for low priority inference models."
+```
+
+Then, assign these priority classes to your models:
+
+```yaml
+apiVersion: kubeai.org/v1
+kind: Model
+metadata:
+  name: critical-service-model
+spec:
+  features: [TextGeneration]
+  url: hf://neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8
+  engine: VLLM
+  resourceProfile: nvidia-gpu-l4:1
+  priorityClassName: high-priority
+---
+apiVersion: kubeai.org/v1
+kind: Model
+metadata:
+  name: background-research-model
+spec:
+  features: [TextGeneration]
+  url: ollama://gemma2:2b
+  engine: OLlama
+  resourceProfile: cpu:2
+  priorityClassName: low-priority
+```
+
+When the cluster is under resource pressure, Kubernetes will evict lower-priority model pods to make room for higher-priority model pods. This ensures that your most critical models remain available even during resource constraints.
+
 ## Feedback welcome: A model management UI
 
 We are considering adding a UI for managing models in a running KubeAI instance. Give the [GitHub Issue](https://github.com/substratusai/kubeai/issues/148) a thumbs up if you would be interested in this feature.
diff --git a/docs/reference/kubernetes-api.md b/docs/reference/kubernetes-api.md
index 6947a50..a488670 100644
--- a/docs/reference/kubernetes-api.md
+++ b/docs/reference/kubernetes-api.md
@@ -147,6 +147,7 @@ _Appears in:_
 | `owner` _string_ | Owner of the model. Used solely to populate the owner field in the<br />OpenAI /v1/models endpoint.<br />DEPRECATED. |  | Optional: \{\} <br /> |
 | `loadBalancing` _[LoadBalancing](#loadbalancing)_ | LoadBalancing configuration for the model.<br />If not specified, a default is used based on the engine and request. | \{  \} |  |
 | `files` _[File](#file) array_ | Files to be mounted in the model Pods. |  | MaxItems: 10 <br /> |
+| `priorityClassName` _string_ | PriorityClassName sets the priority class for all pods created for this model.<br />If specified, the PriorityClass must exist before the model is created.<br />This is useful for implementing priority and preemption for models. |  | Optional: \{\} <br /> |
 
 
 #### ModelStatus
diff --git a/examples/priority-examples/README.md b/examples/priority-examples/README.md
new file mode 100644
index 0000000..1b84537
--- /dev/null
+++ b/examples/priority-examples/README.md
@@ -0,0 +1,50 @@
+# Model Priority Examples
+
+This directory contains examples of how to use Kubernetes Pod Priority and Preemption with KubeAI models.
+
+## Understanding Pod Priority and Preemption
+
+Kubernetes [Pod Priority and Preemption](https://kubernetes.io/docs/concepts/scheduling-eviction/pod-priority-preemption/) is a feature that allows you to assign different priority levels to pods. When resources are constrained, Kubernetes will prioritize scheduling high-priority pods over low-priority pods, and may even evict low-priority pods to make room for high-priority pods.
+
+This is particularly useful in GPU-constrained environments where you want to ensure that critical AI services always have access to resources, even if it means temporarily removing less important tasks.
+
+## Files in this Directory
+
+- `priority-classes.yaml`: Defines the PriorityClass resources needed to implement pod priority
+- `critical-service-model.yaml`: An example of a high-priority model that should preempt other models when resources are limited
+- `background-research-model.yaml`: An example of a low-priority model that can be preempted when resources are needed for higher priority models
+
+## Usage
+
+1. First, apply the priority classes to your cluster:
+
+```bash
+kubectl apply -f priority-classes.yaml
+```
+
+2. Deploy your models with appropriate priority classes:
+
+```bash
+kubectl apply -f critical-service-model.yaml
+kubectl apply -f background-research-model.yaml
+```
+
+## Behavior
+
+When the cluster is under resource pressure:
+
+1. The `critical-service-model` will be scheduled preferentially
+2. If necessary, the `background-research-model` may be evicted to make room for the `critical-service-model`
+
+This ensures that your most important models remain available even during resource constraints.
+
+## Notes on Pod Disruption
+
+When a low-priority pod is preempted:
+
+1. The pod receives a SIGTERM signal
+2. After the grace period (30s by default), if the pod hasn't terminated, it receives a SIGKILL
+3. The pod is removed from the cluster
+4. KubeAI's controller will attempt to reschedule the pod when resources become available again
+
+For more details on configuring priority and preemption, see the [Kubernetes documentation](https://kubernetes.io/docs/concepts/scheduling-eviction/pod-priority-preemption/). 
\ No newline at end of file
diff --git a/examples/priority-examples/background-research-model.yaml b/examples/priority-examples/background-research-model.yaml
new file mode 100644
index 0000000..5c3d351
--- /dev/null
+++ b/examples/priority-examples/background-research-model.yaml
@@ -0,0 +1,11 @@
+apiVersion: kubeai.org/v1
+kind: Model
+metadata:
+  name: background-research-model
+spec:
+  features: [TextGeneration]
+  url: ollama://gemma2:2b
+  engine: OLlama
+  # Background tasks with low priority will be preempted when resources are needed for higher priority models
+  priorityClassName: low-priority
+  resourceProfile: cpu:2 
\ No newline at end of file
diff --git a/examples/priority-examples/critical-service-model.yaml b/examples/priority-examples/critical-service-model.yaml
new file mode 100644
index 0000000..b4168db
--- /dev/null
+++ b/examples/priority-examples/critical-service-model.yaml
@@ -0,0 +1,11 @@
+apiVersion: kubeai.org/v1
+kind: Model
+metadata:
+  name: critical-service-model
+spec:
+  features: [TextGeneration]
+  url: hf://neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8
+  engine: VLLM
+  # Critical service model gets high priority to preempt other models when resources are limited
+  priorityClassName: high-priority
+  resourceProfile: nvidia-gpu-l4:1 
\ No newline at end of file
diff --git a/examples/priority-examples/hello-world-llm.yaml b/examples/priority-examples/hello-world-llm.yaml
new file mode 100644
index 0000000..af4e767
--- /dev/null
+++ b/examples/priority-examples/hello-world-llm.yaml
@@ -0,0 +1,11 @@
+apiVersion: kubeai.org/v1
+kind: Model
+metadata:
+  name: hello-world-llm
+spec:
+  features: [TextGeneration]
+  url: hf://neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8
+  engine: VLLM
+  # Uncomment to set the priority class for this model
+  # priorityClassName: high-priority
+  resourceProfile: nvidia-gpu-l4:1 
\ No newline at end of file
diff --git a/examples/priority-examples/priority-classes.yaml b/examples/priority-examples/priority-classes.yaml
new file mode 100644
index 0000000..d5574bd
--- /dev/null
+++ b/examples/priority-examples/priority-classes.yaml
@@ -0,0 +1,23 @@
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: high-priority
+value: 1000000  # Higher value means higher priority
+globalDefault: false
+description: "This priority class should be used for critical inference models only."
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: medium-priority
+value: 100000
+globalDefault: false
+description: "This priority class should be used for medium priority inference models."
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: low-priority
+value: 10000
+globalDefault: false
+description: "This priority class should be used for low priority inference models." 
\ No newline at end of file
diff --git a/internal/modelcontroller/engine_fasterwhisper.go b/internal/modelcontroller/engine_fasterwhisper.go
index 0de8b23..ae1da21 100644
--- a/internal/modelcontroller/engine_fasterwhisper.go
+++ b/internal/modelcontroller/engine_fasterwhisper.go
@@ -58,6 +58,7 @@ func (r *ModelReconciler) fasterWhisperPodForModel(m *kubeaiv1.Model, c ModelCon
 			Affinity:           c.Affinity,
 			Tolerations:        c.Tolerations,
 			RuntimeClassName:   c.RuntimeClassName,
+			PriorityClassName:  m.Spec.PriorityClassName,
 			ServiceAccountName: r.ModelServerPods.ModelServiceAccountName,
 			SecurityContext:    r.ModelServerPods.ModelPodSecurityContext,
 			ImagePullSecrets:   r.ModelServerPods.ImagePullSecrets,
diff --git a/internal/modelcontroller/engine_infinity.go b/internal/modelcontroller/engine_infinity.go
index 267da1e..f3ce4ca 100644
--- a/internal/modelcontroller/engine_infinity.go
+++ b/internal/modelcontroller/engine_infinity.go
@@ -76,6 +76,7 @@ func (r *ModelReconciler) infinityPodForModel(m *kubeaiv1.Model, c ModelConfig)
 			Affinity:           c.Affinity,
 			Tolerations:        c.Tolerations,
 			RuntimeClassName:   c.RuntimeClassName,
+			PriorityClassName:  m.Spec.PriorityClassName,
 			ServiceAccountName: r.ModelServerPods.ModelServiceAccountName,
 			SecurityContext:    r.ModelServerPods.ModelPodSecurityContext,
 			ImagePullSecrets:   r.ModelServerPods.ImagePullSecrets,
diff --git a/internal/modelcontroller/engine_ollama.go b/internal/modelcontroller/engine_ollama.go
index 32df66a..4be86a0 100644
--- a/internal/modelcontroller/engine_ollama.go
+++ b/internal/modelcontroller/engine_ollama.go
@@ -71,6 +71,7 @@ func (r *ModelReconciler) oLlamaPodForModel(m *kubeaiv1.Model, c ModelConfig) *c
 			Affinity:           c.Affinity,
 			Tolerations:        c.Tolerations,
 			RuntimeClassName:   c.RuntimeClassName,
+			PriorityClassName:  m.Spec.PriorityClassName,
 			ServiceAccountName: r.ModelServerPods.ModelServiceAccountName,
 			SecurityContext:    r.ModelServerPods.ModelPodSecurityContext,
 			ImagePullSecrets:   r.ModelServerPods.ImagePullSecrets,
diff --git a/internal/modelcontroller/engine_vllm.go b/internal/modelcontroller/engine_vllm.go
index 63e7a0e..b3382a8 100644
--- a/internal/modelcontroller/engine_vllm.go
+++ b/internal/modelcontroller/engine_vllm.go
@@ -74,6 +74,7 @@ func (r *ModelReconciler) vLLMPodForModel(m *kubeaiv1.Model, c ModelConfig) *cor
 			Affinity:           c.Affinity,
 			Tolerations:        c.Tolerations,
 			RuntimeClassName:   c.RuntimeClassName,
+			PriorityClassName:  m.Spec.PriorityClassName,
 			ServiceAccountName: r.ModelServerPods.ModelServiceAccountName,
 			SecurityContext:    r.ModelServerPods.ModelPodSecurityContext,
 			ImagePullSecrets:   r.ModelServerPods.ImagePullSecrets,


Test patch:
diff --git a/test/integration/model_priority_test.go b/test/integration/model_priority_test.go
new file mode 100644
index 0000000..28371f6
--- /dev/null
+++ b/test/integration/model_priority_test.go
@@ -0,0 +1,48 @@
+package integration
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	corev1 "k8s.io/api/core/v1"
+	schedulingv1 "k8s.io/api/scheduling/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+// TestModelPriorityClassName tests that priorityClassName is set properly on model pods.
+func TestModelPriorityClassName(t *testing.T) {
+	initTest(t, baseSysCfg(t))
+	priorityClass := &schedulingv1.PriorityClass{
+		ObjectMeta: metav1.ObjectMeta{
+			Name: "test-priority-class",
+		},
+	}
+	require.NoError(t, testK8sClient.Create(testCtx, priorityClass))
+
+	// Create a Model object with priorityClassName set
+	m := modelForTest(t)
+	const expectedPriorityClassName = "test-priority-class"
+	m.Spec.PriorityClassName = expectedPriorityClassName
+	m.Spec.MinReplicas = 1
+
+	// Create the Model in the Kubernetes cluster
+	require.NoError(t, testK8sClient.Create(testCtx, m))
+
+	// Check that the pod has the correct priorityClassName
+	require.EventuallyWithT(t, func(t *assert.CollectT) {
+		podList := &corev1.PodList{}
+		if !assert.NoError(t, testK8sClient.List(testCtx, podList, client.InNamespace(testNS), client.MatchingLabels{"model": m.Name})) {
+			return
+		}
+		if !assert.Len(t, podList.Items, 1) {
+			return
+		}
+		pod := &podList.Items[0]
+
+		// Verify the priorityClassName is set on the pod
+		assert.Equal(t, expectedPriorityClassName, pod.Spec.PriorityClassName)
+	}, 5*time.Second, time.Second/10, "Pod should have the correct priorityClassName")
+}