Add model initialization for CustomSpec (kubeflow#222)

* Add model initialization for CustomSpec * Change model path enc varibale to STORAGE_URI
magdalenakuhn17 · Aug 7, 2019 · 8e7aeaa · 8e7aeaa
1 parent 7ebb5a4
commit 8e7aeaa
Show file tree

Hide file tree

Showing 4 changed files with 135 additions and 0 deletions.
diff --git a/pkg/apis/serving/v1alpha1/framework_custom.go b/pkg/apis/serving/v1alpha1/framework_custom.go
@@ -17,11 +17,18 @@ import (
 	"fmt"
 
 	knserving "github.com/knative/serving/pkg/apis/serving"
+	"github.com/kubeflow/kfserving/pkg/constants"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/util/sets"
 )
 
 func (c *CustomSpec) GetModelSourceUri() string {
+	// return the CustomSpecModelUri env variable value if set on the spec
+	for _, envVar := range c.Container.Env {
+		if envVar.Name == constants.CustomSpecModelUriEnvVarKey {
+			return envVar.Value
+		}
+	}
 	return ""
 }
 

diff --git a/pkg/constants/constants.go b/pkg/constants/constants.go
@@ -77,6 +77,11 @@ const (
 // DefaultModelLocalMountPath is where models will be mounted by the model-initializer
 const DefaultModelLocalMountPath = "/mnt/models"
 
+// KFService Environment Variables
+const (
+	CustomSpecModelUriEnvVarKey = "STORAGE_URI"
+)
+
 func getEnvOrDefault(key string, fallback string) string {
 	if value, ok := os.LookupEnv(key); ok {
 		return value

diff --git a/pkg/webhook/admission/deployment/model_initializer_injector.go b/pkg/webhook/admission/deployment/model_initializer_injector.go
@@ -153,6 +153,13 @@ func (mi *ModelInitializerInjector) InjectModelInitializer(deployment *appsv1.De
 	}
 	userContainer.VolumeMounts = append(userContainer.VolumeMounts, sharedVolumeReadMount)
 
+	// Change the CustomSpecModelUri env variable value to the default model path if present
+	for index, envVar := range userContainer.Env {
+		if envVar.Name == constants.CustomSpecModelUriEnvVarKey && envVar.Value != "" {
+			userContainer.Env[index].Value = constants.DefaultModelLocalMountPath
+		}
+	}
+
 	// Add volumes to the PodSpec
 	podSpec.Volumes = append(podSpec.Volumes, podVolumes...)
 

diff --git a/pkg/webhook/admission/deployment/model_initializer_injector_test.go b/pkg/webhook/admission/deployment/model_initializer_injector_test.go
@@ -325,6 +325,122 @@ func TestModelInitializerFailureCases(t *testing.T) {
 	}
 }
 
+func TestCustomSpecModelUriInjection(t *testing.T) {
+	scenarios := map[string]struct {
+		original                    *appsv1.Deployment
+		expectedModelUriEnvVariable *v1.EnvVar
+	}{
+		"CustomSpecModelUriSet": {
+			original: &appsv1.Deployment{
+				Spec: appsv1.DeploymentSpec{
+					Template: v1.PodTemplateSpec{
+						ObjectMeta: metav1.ObjectMeta{
+							Annotations: map[string]string{
+								constants.ModelInitializerSourceUriInternalAnnotationKey: "pvc://mypvcname/some/path/on/pvc",
+							},
+						},
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Name: "user-container",
+									Env: []v1.EnvVar{
+										v1.EnvVar{
+											Name:  constants.CustomSpecModelUriEnvVarKey,
+											Value: "pvc://mypvcname/some/path/on/pvc",
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			expectedModelUriEnvVariable: &v1.EnvVar{
+				Name:  constants.CustomSpecModelUriEnvVarKey,
+				Value: constants.DefaultModelLocalMountPath,
+			},
+		},
+		"CustomSpecModelUriEmpty": {
+			original: &appsv1.Deployment{
+				Spec: appsv1.DeploymentSpec{
+					Template: v1.PodTemplateSpec{
+						ObjectMeta: metav1.ObjectMeta{
+							Annotations: map[string]string{
+								constants.ModelInitializerSourceUriInternalAnnotationKey: "pvc://mypvcname/some/path/on/pvc",
+							},
+						},
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Name: "user-container",
+									Env: []v1.EnvVar{
+										v1.EnvVar{
+											Name:  constants.CustomSpecModelUriEnvVarKey,
+											Value: "",
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			expectedModelUriEnvVariable: &v1.EnvVar{
+				Name:  constants.CustomSpecModelUriEnvVarKey,
+				Value: "",
+			},
+		},
+		"CustomSpecModelUriNotSet": {
+			original: &appsv1.Deployment{
+				Spec: appsv1.DeploymentSpec{
+					Template: v1.PodTemplateSpec{
+						ObjectMeta: metav1.ObjectMeta{
+							Annotations: map[string]string{
+								constants.ModelInitializerSourceUriInternalAnnotationKey: "pvc://mypvcname/some/path/on/pvc",
+							},
+						},
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Name: "user-container",
+									Env: []v1.EnvVar{
+										v1.EnvVar{
+											Name:  "TestRandom",
+											Value: "val",
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			expectedModelUriEnvVariable: nil,
+		},
+	}
+
+	for name, scenario := range scenarios {
+		injector := &ModelInitializerInjector{
+			credentialBuilder: credentials.NewCredentialBulder(c, &v1.ConfigMap{
+				Data: map[string]string{},
+			}),
+		}
+		if err := injector.InjectModelInitializer(scenario.original); err != nil {
+			t.Errorf("Test %q unexpected result: %s", name, err)
+		}
+
+		var originalEnvVar *v1.EnvVar
+		for _, envVar := range scenario.original.Spec.Template.Spec.Containers[0].Env {
+			if envVar.Name == constants.CustomSpecModelUriEnvVarKey {
+				originalEnvVar = &envVar
+			}
+		}
+		if diff := cmp.Diff(scenario.expectedModelUriEnvVariable, originalEnvVar); diff != "" {
+			t.Errorf("Test %q unexpected result (-want +got): %v", name, diff)
+		}
+	}
+}
+
 func makeDeployment() *appsv1.Deployment {
 	return &appsv1.Deployment{
 		ObjectMeta: metav1.ObjectMeta{