support config request resources in kserve runtime (#1050)

Signed-off-by: zibai <[email protected]>
kubeflow · Mar 7, 2024 · b7f0ecf · b7f0ecf
1 parent 57093a2
commit b7f0ecf
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/charts/kserve/templates/inferenceservice.yaml b/charts/kserve/templates/inferenceservice.yaml
@@ -103,6 +103,12 @@ spec:
           aliyun.com/gpu-core.percentage: {{ .Values.gpuCore }}
           {{- end }}
         requests:
+          {{- if .Values.cpu }}
+          cpu: {{ .Values.cpu }}
+          {{- end }}
+          {{- if .Values.memory }}
+          memory: {{ .Values.memory }}
+          {{- end }}
           {{- if gt (int $gpuCount) 0}}
           nvidia.com/gpu: {{ .Values.gpuCount }}
           {{- end }}
@@ -171,6 +177,12 @@ spec:
           {{- end }}
         resources:
           requests:
+            {{- if .Values.cpu }}
+            cpu: {{ .Values.cpu }}
+            {{- end }}
+            {{- if .Values.memory }}
+            memory: {{ .Values.memory }}
+            {{- end }}
             {{- if gt (int $gpuCount) 0}}
             nvidia.com/gpu: {{ .Values.gpuCount }}
             {{- end }}

diff --git a/pkg/serving/update.go b/pkg/serving/update.go
@@ -563,6 +563,19 @@ func setInferenceServiceForCustomModel(args *types.UpdateKServeArgs, inferenceSe
 		inferenceService.Spec.Predictor.Containers[0].Image = args.Image
 	}
 
+	// set resources requests
+	resourceRequests := inferenceService.Spec.Predictor.Containers[0].Resources.Requests
+	if resourceRequests == nil {
+		resourceRequests = make(map[v1.ResourceName]resource.Quantity)
+	}
+	if args.Cpu != "" {
+		resourceRequests[v1.ResourceCPU] = resource.MustParse(args.Cpu)
+	}
+	if args.Memory != "" {
+		resourceRequests[v1.ResourceMemory] = resource.MustParse(args.Memory)
+	}
+	inferenceService.Spec.Predictor.Containers[0].Resources.Requests = resourceRequests
+
 	// set resources limits
 	resourceLimits := inferenceService.Spec.Predictor.Containers[0].Resources.Limits
 	if resourceLimits == nil {