fix: replace image tag (kubeflow#2074)

* fix: replace image tag Signed-off-by: Thang Vu <[email protected]> * chore: remove redundant condition Signed-off-by: Thang Vu <[email protected]>
magdalenakuhn17 · Mar 6, 2022 · d1c24cd · d1c24cd
1 parent 89d25e1
commit d1c24cd
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 2 deletions.
diff --git a/pkg/controller/v1beta1/inferenceservice/utils/utils.go b/pkg/controller/v1beta1/inferenceservice/utils/utils.go
@@ -21,6 +21,7 @@ import (
 	"context"
 	"encoding/json"
 	"html/template"
+	"regexp"
 	"strings"
 
 	"github.com/kserve/kserve/pkg/apis/serving/v1alpha1"
@@ -185,8 +186,13 @@ func ReplacePlaceholders(container *v1.Container, meta metav1.ObjectMeta) error
 // UpdateImageTag Update image tag if GPU is enabled or runtime version is provided
 func UpdateImageTag(container *v1.Container, runtimeVersion *string, isvcConfig *v1beta1.InferenceServicesConfig) {
 	image := container.Image
-	if runtimeVersion != nil && len(strings.Split(image, ":")) > 0 {
-		container.Image = strings.Split(image, ":")[0] + ":" + *runtimeVersion
+	if runtimeVersion != nil {
+		re := regexp.MustCompile(`(:([\w.\-_]*))$`)
+		if len(re.FindString(image)) == 0 {
+			container.Image = image + ":" + *runtimeVersion
+		} else {
+			container.Image = re.ReplaceAllString(image, ":" + *runtimeVersion)
+		}
 		return
 	}
 	if utils.IsGPUEnabled(container.Resources) && len(strings.Split(image, ":")) > 0 {

diff --git a/pkg/controller/v1beta1/inferenceservice/utils/utils_test.go b/pkg/controller/v1beta1/inferenceservice/utils/utils_test.go
@@ -1551,6 +1551,62 @@ func TestUpdateImageTag(t *testing.T) {
 			isvcConfig:     &config,
 			expected:       "tfserving:1.14.0-gpu",
 		},
+		"UpdateRuntimeVersionWithProxy": {
+			container: &v1.Container{
+				Name:  "kserve-container",
+				Image: "localhost:8888/tfserving",
+				Args: []string{
+					"--foo=bar",
+					"--test=dummy",
+					"--new-arg=baz",
+				},
+				Env: []v1.EnvVar{
+					{Name: "PORT", Value: "8080"},
+					{Name: "MODELS_DIR", Value: "/mnt/models"},
+				},
+				Resources: v1.ResourceRequirements{
+					Limits: v1.ResourceList{
+						v1.ResourceCPU:    resource.MustParse("2"),
+						v1.ResourceMemory: resource.MustParse("4Gi"),
+					},
+					Requests: v1.ResourceList{
+						v1.ResourceCPU:    resource.MustParse("1"),
+						v1.ResourceMemory: resource.MustParse("2Gi"),
+					},
+				},
+			},
+			runtimeVersion: proto.String("2.6.2"),
+			isvcConfig:     &config,
+			expected:       "localhost:8888/tfserving:2.6.2",
+		},
+		"UpdateRuntimeVersionWithProxyAndTag": {
+			container: &v1.Container{
+				Name:  "kserve-container",
+				Image: "localhost:8888/tfserving:1.2.3",
+				Args: []string{
+					"--foo=bar",
+					"--test=dummy",
+					"--new-arg=baz",
+				},
+				Env: []v1.EnvVar{
+					{Name: "PORT", Value: "8080"},
+					{Name: "MODELS_DIR", Value: "/mnt/models"},
+				},
+				Resources: v1.ResourceRequirements{
+					Limits: v1.ResourceList{
+						v1.ResourceCPU:    resource.MustParse("2"),
+						v1.ResourceMemory: resource.MustParse("4Gi"),
+					},
+					Requests: v1.ResourceList{
+						v1.ResourceCPU:    resource.MustParse("1"),
+						v1.ResourceMemory: resource.MustParse("2Gi"),
+					},
+				},
+			},
+			runtimeVersion: proto.String("2.6.2"),
+			isvcConfig:     &config,
+			expected:       "localhost:8888/tfserving:2.6.2",
+		},
 	}
 	for name, scenario := range scenarios {
 		t.Run(name, func(t *testing.T) {