Repository: substratusai/kubeai
Pull request #497 (2025-04-07 23:52:09)

support use of JSONPatch for model server pods
Adds support for applying arbitrary JSON 6902 / JSONPatch support for model server pods.

For example, our clusters have a number of mutating webhooks which inject special priority class names and various pre-emption policies based on various labels, etc. This gives operators a way to set the pod templates up which can work in constrained clusters (i.e. those with OPA gatekeeper constraints).

```yaml
modelServerPods:
  - op: add
    path: /spec/priorityClassName
    value: gpu

```

nstogner: Thanks for the PR! Will try to get back to you on this one by tomorrow.
alam0rt: > Thanks for the PR! Will try to get back to you on this one by tomorrow.

cheers mate! I have been using in our clusters and it does what it says on the box. Was originally going to just add a way to specify priority class / preemption policy but figured arbitrary patches would make for a cleaner change.
nstogner: I'll take another look tomorrow, thanks for updating
nstogner: I would like to update the error handling logic to follow the general pattern that exists in the project: return an error as far up the stack as possible until it can be handled or logged at the highest point possible. Instead of logging in the `calculatePodPlan()` func, let's update that func to return an error. Inside of `Reconcile()`, if that error is returned, log the error, then return a nil error - because returning a non-nil error in Reconcile will start a retry-loop - and retrying is not going to solve the issue until the Model is updated - in which case Reconcile() will be called again naturally.

We will add a followup PR to fire a k8s event from within Reconcile() to bubble the error up to the admin on the Model object - we have several other use cases for events, so I would worry about that part in this PR.
alam0rt: May need to, in a subsequent PR, think of ways to target specific engines/model server pods. i.e. if we want to inject patches into _just_ VLLM.

Patch:
diff --git a/charts/kubeai/templates/configmap.yaml b/charts/kubeai/templates/configmap.yaml
index 0350721..47bcfa1 100644
--- a/charts/kubeai/templates/configmap.yaml
+++ b/charts/kubeai/templates/configmap.yaml
@@ -27,6 +27,10 @@ data:
       podSecurityContext:
         {{- .Values.modelServerPods.podSecurityContext | toYaml | nindent 8}}
       {{- end}}
+      {{- if .Values.modelServerPods.jsonPatches }}
+      jsonPatches:
+        {{- .Values.modelServerPods.jsonPatches | toYaml | nindent 8}}
+      {{- end}}
       {{- if .Values.modelServerPods.securityContext }}
       securityContext:
         {{- .Values.modelServerPods.securityContext | toYaml | nindent 8}}
@@ -42,4 +46,4 @@ data:
       timeWindow: {{ .Values.modelAutoscaling.timeWindow }}
       stateConfigMapName: {{ include "models.autoscalerStateConfigMapName" . }}
     messaging:
-      {{- .Values.messaging | toYaml | nindent 6 }}
\ No newline at end of file
+      {{- .Values.messaging | toYaml | nindent 6 }}
diff --git a/charts/kubeai/values.yaml b/charts/kubeai/values.yaml
index 58e61ae..ad96139 100644
--- a/charts/kubeai/values.yaml
+++ b/charts/kubeai/values.yaml
@@ -40,12 +40,12 @@ modelServers:
       # The key is the image name (referenced from resourceProfiles) and the value is the image.
       # The "default" image should always be specified.
       # "default" is used when no imageName is specified or if a specific image is not found.
-      default: "vllm/vllm-openai:v0.8.2"
+      default: "vllm/vllm-openai:v0.8.3"
+      nvidia-gpu: "vllm/vllm-openai:v0.8.3"
       cpu: "substratusai/vllm:v0.6.3.post1-cpu"
       google-tpu: "substratusai/vllm:v0.6.4.post1-tpu"
-      nvidia-gpu: "vllm/vllm-openai:v0.8.2"
       # Source: https://github.com/substratusai/vllm-docker/blob/main/Dockerfile.cuda-arm
-      gh200: "substratusai/vllm-gh200:v0.8.2"
+      gh200: "substratusai/vllm-gh200:v0.8.3"
       # upstream vLLM seems to have broken ROCm support, so we are using a fork from AMD.
       # Source: https://hub.docker.com/r/rocm/vllm-dev
       # Source: https://github.com/ROCm/vllm
@@ -74,6 +74,12 @@ modelServerPods:
     capabilities:
       drop:
         - ALL
+  # JSONPatch to apply to the model server pods
+  # Invalid patches will be ignored.
+  # jsonPatches:
+  # - op: add
+  #   path: /spec/priorityClassName
+  #   value: kubeai-model-server
 
 modelRollouts:
   # The number of replicas to add when rolling out a new model.
diff --git a/charts/models/values.yaml b/charts/models/values.yaml
index bd1987f..ad64266 100644
--- a/charts/models/values.yaml
+++ b/charts/models/values.yaml
@@ -495,6 +495,19 @@ catalog:
       - --disable-log-requests
     resourceProfile: nvidia-gpu-gh200:1
     targetRequests: 200
+  llama-4-maverick-430k-h100:
+    enabled: false
+    features: [TextGeneration]
+    url: hf://meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8
+    engine: VLLM
+    env:
+      VLLM_DISABLE_COMPILE_CACHE: "1"
+    args:
+      - --max-model-len=430000
+      - --tensor-parallel-size=8
+      - --enable-prefix-caching
+      - --disable-log-requests
+    resourceProfile: nvidia-gpu-h100:8
   deepseek-r1-70b-gh200-fp8:
     enabled: false
     features: [TextGeneration]
diff --git a/go.mod b/go.mod
index e53b648..109bc68 100644
--- a/go.mod
+++ b/go.mod
@@ -152,6 +152,7 @@ require (
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240812133136-8ffd90a71988 // indirect
 	google.golang.org/grpc v1.65.0 // indirect
 	google.golang.org/protobuf v1.36.3 // indirect
+	gopkg.in/evanphx/json-patch.v4 v4.12.0
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
diff --git a/go.sum b/go.sum
index 3d2e619..2672098 100644
--- a/go.sum
+++ b/go.sum
@@ -501,6 +501,8 @@ google.golang.org/protobuf v1.36.3/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojt
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
+gopkg.in/evanphx/json-patch.v4 v4.12.0 h1:n6jtcsulIzXPJaxegRbvFNNrZDjbij7ny3gmSPG+6V4=
+gopkg.in/evanphx/json-patch.v4 v4.12.0/go.mod h1:p8EYWUEYMpynmqDbY58zCKCFZw8pRWMG4EsWvDvM72M=
 gopkg.in/inf.v0 v0.9.1 h1:73M5CoZyi3ZLMOyDlQh031Cx6N9NDJ2Vvfl76EDAgDc=
 gopkg.in/inf.v0 v0.9.1/go.mod h1:cWUDdTG/fYaXco+Dcufb5Vnc6Gp2YChqWtbxRZE0mXw=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
diff --git a/internal/config/system.go b/internal/config/system.go
index 40909e2..ee2e44b 100644
--- a/internal/config/system.go
+++ b/internal/config/system.go
@@ -233,6 +233,12 @@ type ModelLoading struct {
 	Image string `json:"image" validate:"required"`
 }
 
+type JSONPatch struct {
+	Op    string      `json:"op"`
+	Path  string      `json:"path"`
+	Value interface{} `json:"value"`
+}
+
 type ModelServerPods struct {
 	// The service account to use for all model pods
 	ModelServiceAccountName string `json:"serviceAccountName,omitempty"`
@@ -245,4 +251,9 @@ type ModelServerPods struct {
 
 	// ImagePullSecrets is a list of references to secrets in the same namespace to use for pulling any of the images
 	ImagePullSecrets []corev1.LocalObjectReference `json:"imagePullSecrets,omitempty"`
+
+	// JSONPatches is a list of patches to apply to the model pod template.
+	// This is a JSON Patch as defined in RFC 6902.
+	// https://datatracker.ietf.org/doc/html/rfc6902
+	JSONPatches []JSONPatch `json:"jsonPatches,omitempty"`
 }
diff --git a/internal/modelcontroller/model_controller.go b/internal/modelcontroller/model_controller.go
index f9e4285..a007d4c 100644
--- a/internal/modelcontroller/model_controller.go
+++ b/internal/modelcontroller/model_controller.go
@@ -173,7 +173,12 @@ func (r *ModelReconciler) Reconcile(ctx context.Context, req ctrl.Request) (res
 		}
 	}()
 
-	plan := r.calculatePodPlan(allPods, model, modelConfig)
+	plan, err := r.calculatePodPlan(allPods, model, modelConfig)
+	if err != nil {
+		log.Error(err, "Failed to calculate pod plan")
+		return ctrl.Result{}, nil
+	}
+
 	if plan.containsActions() {
 		var err error
 		scaled, err = plan.execute(ctx, r.Client, r.Scheme)
diff --git a/internal/modelcontroller/patch.go b/internal/modelcontroller/patch.go
new file mode 100644
index 0000000..517b48b
--- /dev/null
+++ b/internal/modelcontroller/patch.go
@@ -0,0 +1,43 @@
+package modelcontroller
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/substratusai/kubeai/internal/config"
+	jsonpatch "gopkg.in/evanphx/json-patch.v4"
+	corev1 "k8s.io/api/core/v1"
+)
+
+func applyJSONPatchToPod(patches []config.JSONPatch, pod *corev1.Pod) error {
+	if len(patches) == 0 {
+		return nil
+	}
+
+	pb, err := json.Marshal(patches)
+	if err != nil {
+		return fmt.Errorf("marshal pod patch: %w", err)
+	}
+
+	patch, err := jsonpatch.DecodePatch(pb)
+	if err != nil {
+		return fmt.Errorf("decode pod patch: %w", err)
+	}
+
+	podJson, err := json.Marshal(pod)
+	if err != nil {
+		return fmt.Errorf("marshal pod: %w", err)
+	}
+
+	patchedPodJson, err := patch.Apply(podJson)
+	if err != nil {
+		return fmt.Errorf("apply pod patch: %w", err)
+	}
+
+	patchedPod := &corev1.Pod{}
+	if err := json.Unmarshal(patchedPodJson, patchedPod); err != nil {
+		return fmt.Errorf("unmarshal patched pod: %w", err)
+	}
+	*pod = *patchedPod
+	return nil
+}
diff --git a/internal/modelcontroller/pod_plan.go b/internal/modelcontroller/pod_plan.go
index c7a0cf7..d29df57 100644
--- a/internal/modelcontroller/pod_plan.go
+++ b/internal/modelcontroller/pod_plan.go
@@ -25,8 +25,9 @@ import (
 // - Adds a surge Pod
 // - Recreates any out-of-date Pod that is not Ready immediately
 // - Waits for all Pods to be Ready before recreating any out-of-date Pods that are Ready
-func (r *ModelReconciler) calculatePodPlan(allPods *corev1.PodList, model *kubeaiv1.Model, modelConfig ModelConfig) *podPlan {
+func (r *ModelReconciler) calculatePodPlan(allPods *corev1.PodList, model *kubeaiv1.Model, modelConfig ModelConfig) (*podPlan, error) {
 	var podForModel *corev1.Pod
+
 	switch model.Spec.Engine {
 	case kubeaiv1.OLlamaEngine:
 		podForModel = r.oLlamaPodForModel(model, modelConfig)
@@ -37,6 +38,11 @@ func (r *ModelReconciler) calculatePodPlan(allPods *corev1.PodList, model *kubea
 	default:
 		podForModel = r.vLLMPodForModel(model, modelConfig)
 	}
+
+	if err := applyJSONPatchToPod(r.ModelServerPods.JSONPatches, podForModel); err != nil {
+		return nil, err
+	}
+
 	expectedHash := k8sutils.PodHash(podForModel.Spec)
 	podForModel.GenerateName = fmt.Sprintf("model-%s-%s-", model.Name, expectedHash)
 	k8sutils.SetLabel(podForModel, kubeaiv1.PodHashLabel, expectedHash)
@@ -146,7 +152,7 @@ func (r *ModelReconciler) calculatePodPlan(allPods *corev1.PodList, model *kubea
 		toDelete: toDelete,
 		toRemain: toRemain,
 		details:  details,
-	}
+	}, nil
 }
 
 type podPlan struct {
diff --git a/manifests/models/llama-4-maverick-430k-h100.yaml b/manifests/models/llama-4-maverick-430k-h100.yaml
new file mode 100644
index 0000000..bd9d308
--- /dev/null
+++ b/manifests/models/llama-4-maverick-430k-h100.yaml
@@ -0,0 +1,18 @@
+# Source: models/templates/models.yaml
+apiVersion: kubeai.org/v1
+kind: Model
+metadata:
+  name: llama-4-maverick-430k-h100
+spec:
+  features: [TextGeneration]
+  url: hf://meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8
+  engine: VLLM
+  args:
+    - --max-model-len=430000
+    - --tensor-parallel-size=8
+    - --enable-prefix-caching
+    - --disable-log-requests
+  env:
+    VLLM_DISABLE_COMPILE_CACHE: "1"
+  minReplicas: 0
+  resourceProfile: nvidia-gpu-h100:8


Test patch:
diff --git a/internal/modelcontroller/patch_test.go b/internal/modelcontroller/patch_test.go
new file mode 100644
index 0000000..cd9ff12
--- /dev/null
+++ b/internal/modelcontroller/patch_test.go
@@ -0,0 +1,117 @@
+package modelcontroller
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/substratusai/kubeai/internal/config"
+	"github.com/substratusai/kubeai/internal/k8sutils"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/utils/ptr"
+)
+
+func Test_applyJSONPatchToPod(t *testing.T) {
+	t.Parallel()
+
+	emptyPod := &corev1.Pod{}
+	emptyPodHash := k8sutils.PodHash(emptyPod.Spec)
+
+	cases := map[string]struct {
+		modelPod     *corev1.Pod
+		patches      []config.JSONPatch
+		want         *corev1.Pod
+		errorString  *string
+		expectedHash *string
+	}{
+		"no-patch": {
+			modelPod:     emptyPod,
+			patches:      nil,
+			want:         emptyPod,
+			expectedHash: ptr.To(emptyPodHash), // default hash
+		},
+		"error-patch": {
+			modelPod: emptyPod,
+			patches: []config.JSONPatch{
+				{Op: "invalid", Path: "/spec/containers/0/image", Value: "new-image"},
+			},
+			want:         emptyPod,
+			expectedHash: ptr.To(emptyPodHash), // since the patch failed it should not change the hash
+			errorString:  ptr.To("apply pod patch: Unexpected kind: invalid"),
+		},
+		"replace-image": {
+			modelPod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-pod",
+					Namespace: "test-namespace",
+				},
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Name:  "test-container",
+							Image: "test-image",
+						},
+					},
+				},
+			},
+			patches: []config.JSONPatch{
+				{Op: "replace", Path: "/spec/containers/0/image", Value: "new-image"},
+			},
+			want: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-pod",
+					Namespace: "test-namespace",
+				},
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Name:  "test-container",
+							Image: "new-image",
+						},
+					},
+				},
+			},
+		},
+		"add-preemption-policy": {
+			modelPod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-pod",
+					Namespace: "test-namespace",
+				},
+				Spec: corev1.PodSpec{},
+			},
+			patches: []config.JSONPatch{
+				{Op: "add", Path: "/spec/preemptionPolicy", Value: "Never"},
+			},
+			want: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-pod",
+					Namespace: "test-namespace",
+				},
+				Spec: corev1.PodSpec{
+					PreemptionPolicy: ptr.To(corev1.PreemptNever),
+				},
+			},
+		},
+	}
+
+	for name, c := range cases {
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+			gotHash := k8sutils.PodHash(c.modelPod.Spec)
+
+			err := applyJSONPatchToPod(c.patches, c.modelPod)
+			if c.errorString != nil {
+				require.ErrorContains(t, err, *c.errorString)
+			} else {
+				require.NoError(t, err)
+			}
+
+			if c.expectedHash != nil {
+				require.Equal(t, *c.expectedHash, gotHash, "expected pod hash to be %s", *c.expectedHash)
+			}
+
+			require.Equal(t, c.want, c.modelPod, "expected pod to be patched correctly")
+		})
+	}
+}
diff --git a/internal/modelcontroller/pod_plan_test.go b/internal/modelcontroller/pod_plan_test.go
index cbeb798..cda1ca8 100644
--- a/internal/modelcontroller/pod_plan_test.go
+++ b/internal/modelcontroller/pod_plan_test.go
@@ -91,6 +91,7 @@ func Test_calculatePodPlan(t *testing.T) {
 		pods           []corev1.Pod
 		wantNCreations int
 		wantDeletions  []string
+		jsonPatches    []config.JSONPatch
 	}{
 		{
 			name: "do nothing",
@@ -154,7 +155,12 @@ func Test_calculatePodPlan(t *testing.T) {
 
 	for _, c := range cases {
 		t.Run(c.name, func(t *testing.T) {
-			plan := r.calculatePodPlan(&corev1.PodList{Items: c.pods}, model, modelConfig)
+			r.ModelServerPods.JSONPatches = nil
+			if c.jsonPatches != nil {
+				r.ModelServerPods.JSONPatches = c.jsonPatches
+			}
+			plan, err := r.calculatePodPlan(&corev1.PodList{Items: c.pods}, model, modelConfig)
+			require.NoError(t, err)
 			detailsCSV := strings.Join(plan.details, ", ")
 			require.Lenf(t, plan.toCreate, c.wantNCreations, "Unexpected creation count, details: %v", detailsCSV)
 			var deletionNames []string