diff --git a/pkg/apis/serving/v1alpha1/framework_custom.go b/pkg/apis/serving/v1alpha1/framework_custom.go index d8bf94e356f..41f27e20314 100644 --- 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 index a63c323a191..232c0fd4249 100644 --- 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 index a85909485d8..e3a365804f9 100644 --- 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 index 3185f6830c7..f768db1848b 100644 --- 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{