Repository: kserve/kserve

Issue:
KFServing should have a consistent way of supporting model download across inference server implementations
/kind feature

**Describe the solution you'd like**
KFServing should expose a consistent way to download models across inference servers and clouds. The current implementation depends on the features of individual inference servers expose. E.g. see #137 

**Anything else you would like to add:**
Proposed solution design is documented here: https://docs.google.com/document/d/1xqBOkoQ6Vzc5gv4O5MgVVNE3qILbKuMkC-DN5zp5w28/edit?usp=sharing



Discussion:
@krabhishek8260 @ellis-bigelow @cliveseldon @yuzisun 
Apart from the needed knative workarounds, the init container solution looks best to me and can be tied into a separate caching layer easily.
Cool! Can you share this with the kubeflow community?
Also, i'd title this as model caching or something about unified storage layer instead of "referring to models"
> model caching

I was hoping to keep the scope to a minimum aka tackle caching once the download part is figured out. Made an attempt to update the title, if you have a better suggestion let me know...
Sounds good now

Resolving patch:
diff --git a/python/kfserving/kfserving/storage.py b/python/kfserving/kfserving/storage.py
--- a/python/kfserving/kfserving/storage.py
+++ b/python/kfserving/kfserving/storage.py
@@ -27,23 +27,25 @@
 
 class Storage(object): # pylint: disable=too-few-public-methods
     @staticmethod
-    def download(uri: str) -> str:
+    def download(uri: str, out_dir: str = None) -> str:
         logging.info("Copying contents of %s to local", uri)
         if uri.startswith(_LOCAL_PREFIX) or os.path.exists(uri):
             return Storage._download_local(uri)
 
-        temp_dir = tempfile.mkdtemp()
+        if out_dir is None:
+            out_dir = tempfile.mkdtemp()
+
         if uri.startswith(_GCS_PREFIX):
-            Storage._download_gcs(uri, temp_dir)
+            Storage._download_gcs(uri, out_dir)
         elif uri.startswith(_S3_PREFIX):
-            Storage._download_s3(uri, temp_dir)
+            Storage._download_s3(uri, out_dir)
         else:
             raise Exception("Cannot recognize storage type for " + uri +
                             "\n'%s', '%s', and '%s' are the current available storage type." %
                             (_GCS_PREFIX, _S3_PREFIX, _LOCAL_PREFIX))
 
-        logging.info("Successfully copied %s to %s", uri, temp_dir)
-        return temp_dir
+        logging.info("Successfully copied %s to %s", uri, out_dir)
+        return out_dir
 
     @staticmethod
     def _download_s3(uri, temp_dir: str):
@@ -68,16 +70,23 @@ def _download_gcs(uri, temp_dir: str):
         bucket_name = bucket_args[0]
         bucket_path = bucket_args[1] if len(bucket_args) > 1 else ""
         bucket = storage_client.bucket(bucket_name)
-        blobs = bucket.list_blobs(prefix=bucket_path)
+        prefix = bucket_path
+        if not prefix.endswith("/"):
+            prefix = prefix + "/"
+        blobs = bucket.list_blobs(prefix=prefix)
         for blob in blobs:
             # Replace any prefix from the object key with temp_dir
             subdir_object_key = blob.name.replace(bucket_path, "", 1).strip("/")
+
             # Create necessary subdirectory to store the object locally
             if "/" in subdir_object_key:
                 local_object_dir = os.path.join(temp_dir, subdir_object_key.rsplit("/", 1)[0])
                 if not os.path.isdir(local_object_dir):
                     os.makedirs(local_object_dir, exist_ok=True)
-            blob.download_to_filename(os.path.join(temp_dir, subdir_object_key))
+            if subdir_object_key.strip() != "":
+                dest_path = os.path.join(temp_dir, subdir_object_key)
+                logging.info("Downloading: %s", dest_path)
+                blob.download_to_filename(dest_path)
 
     @staticmethod
     def _download_local(uri):


Test patch:
diff --git a/pkg/apis/serving/v1alpha1/kfservice_validation_test.go b/pkg/apis/serving/v1alpha1/kfservice_validation_test.go
--- a/pkg/apis/serving/v1alpha1/kfservice_validation_test.go
+++ b/pkg/apis/serving/v1alpha1/kfservice_validation_test.go
@@ -41,6 +41,42 @@ func makeTestKFService() KFService {
 	return kfservice
 }
 
+func TestValidModelURIPrefixOK(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+	for _, prefix := range SupportedModelSourceURIPrefixList {
+		kfsvc := makeTestKFService()
+		kfsvc.Spec.Default.Tensorflow.ModelURI = prefix + "foo/bar"
+		g.Expect(kfsvc.ValidateCreate()).Should(gomega.Succeed())
+	}
+}
+
+func TestEmptyModelURIPrefixOK(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+	kfsvc := makeTestKFService()
+	kfsvc.Spec.Default.Tensorflow.ModelURI = ""
+	g.Expect(kfsvc.ValidateCreate()).Should(gomega.Succeed())
+}
+
+func TestLocalPathModelURIPrefixOK(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+	kfsvc := makeTestKFService()
+	kfsvc.Spec.Default.Tensorflow.ModelURI = "some/relative/path"
+	g.Expect(kfsvc.ValidateCreate()).Should(gomega.Succeed())
+	kfsvc.Spec.Default.Tensorflow.ModelURI = "/some/absolute/path"
+	g.Expect(kfsvc.ValidateCreate()).Should(gomega.Succeed())
+	kfsvc.Spec.Default.Tensorflow.ModelURI = "/"
+	g.Expect(kfsvc.ValidateCreate()).Should(gomega.Succeed())
+	kfsvc.Spec.Default.Tensorflow.ModelURI = "foo"
+	g.Expect(kfsvc.ValidateCreate()).Should(gomega.Succeed())
+}
+
+func TestUnkownModelURIPrefixFails(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+	kfsvc := makeTestKFService()
+	kfsvc.Spec.Default.Tensorflow.ModelURI = "blob://foo/bar"
+	g.Expect(kfsvc.ValidateCreate()).ShouldNot(gomega.Succeed())
+}
+
 func TestRejectMultipleModelSpecs(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 	kfsvc := makeTestKFService()
diff --git a/pkg/controller/kfservice/kfservice_controller_test.go b/pkg/controller/kfservice/kfservice_controller_test.go
--- a/pkg/controller/kfservice/kfservice_controller_test.go
+++ b/pkg/controller/kfservice/kfservice_controller_test.go
@@ -17,13 +17,14 @@ limitations under the License.
 package service
 
 import (
+	"testing"
+	"time"
+
 	"github.com/knative/pkg/apis"
 	"github.com/knative/serving/pkg/apis/serving/v1beta1"
 	"github.com/kubeflow/kfserving/pkg/constants"
 	testutils "github.com/kubeflow/kfserving/pkg/testing"
 	v1 "k8s.io/api/core/v1"
-	"testing"
-	"time"
 
 	duckv1beta1 "github.com/knative/pkg/apis/duck/v1beta1"
 	knservingv1alpha1 "github.com/knative/serving/pkg/apis/serving/v1alpha1"
@@ -165,10 +166,11 @@ func TestReconcile(t *testing.T) {
 				ObjectMeta: metav1.ObjectMeta{
 					Labels: map[string]string{"serving.kubeflow.org/kfservice": "foo"},
 					Annotations: map[string]string{
-						"autoscaling.knative.dev/target":   "1",
-						"autoscaling.knative.dev/class":    "kpa.autoscaling.knative.dev",
-						"autoscaling.knative.dev/maxScale": "3",
-						"autoscaling.knative.dev/minScale": "1",
+						"autoscaling.knative.dev/target":                         "1",
+						"autoscaling.knative.dev/class":                          "kpa.autoscaling.knative.dev",
+						"autoscaling.knative.dev/maxScale":                       "3",
+						"autoscaling.knative.dev/minScale":                       "1",
+						constants.ModelInitializerSourceUriInternalAnnotationKey: instance.Spec.Default.Tensorflow.ModelURI,
 					},
 				},
 				Spec: knservingv1alpha1.RevisionSpec{
@@ -183,7 +185,7 @@ func TestReconcile(t *testing.T) {
 							"--port=" + servingv1alpha1.TensorflowServingGRPCPort,
 							"--rest_api_port=" + servingv1alpha1.TensorflowServingRestPort,
 							"--model_name=" + instance.Name,
-							"--model_base_path=" + instance.Spec.Default.Tensorflow.ModelURI,
+							"--model_base_path=" + constants.DefaultModelLocalMountPath,
 						},
 					},
 				},
@@ -301,10 +303,11 @@ func TestCanaryReconcile(t *testing.T) {
 				ObjectMeta: metav1.ObjectMeta{
 					Labels: map[string]string{"serving.kubeflow.org/kfservice": "bar"},
 					Annotations: map[string]string{
-						"autoscaling.knative.dev/target":   "1",
-						"autoscaling.knative.dev/class":    "kpa.autoscaling.knative.dev",
-						"autoscaling.knative.dev/maxScale": "3",
-						"autoscaling.knative.dev/minScale": "1",
+						"autoscaling.knative.dev/target":                         "1",
+						"autoscaling.knative.dev/class":                          "kpa.autoscaling.knative.dev",
+						"autoscaling.knative.dev/maxScale":                       "3",
+						"autoscaling.knative.dev/minScale":                       "1",
+						constants.ModelInitializerSourceUriInternalAnnotationKey: canary.Spec.Canary.Tensorflow.ModelURI,
 					},
 				},
 				Spec: knservingv1alpha1.RevisionSpec{
@@ -319,7 +322,7 @@ func TestCanaryReconcile(t *testing.T) {
 							"--port=" + servingv1alpha1.TensorflowServingGRPCPort,
 							"--rest_api_port=" + servingv1alpha1.TensorflowServingRestPort,
 							"--model_name=" + canary.Name,
-							"--model_base_path=" + canary.Spec.Canary.Tensorflow.ModelURI,
+							"--model_base_path=" + constants.DefaultModelLocalMountPath,
 						},
 					},
 				},
diff --git a/pkg/controller/kfservice/reconcilers/knative/configuration_reconciler_test.go b/pkg/controller/kfservice/reconcilers/knative/configuration_reconciler_test.go
--- a/pkg/controller/kfservice/reconcilers/knative/configuration_reconciler_test.go
+++ b/pkg/controller/kfservice/reconcilers/knative/configuration_reconciler_test.go
@@ -88,8 +88,9 @@ func TestKnativeConfigurationReconcile(t *testing.T) {
 						ObjectMeta: metav1.ObjectMeta{
 							Labels: map[string]string{"serving.kubeflow.org/kfservice": "mnist"},
 							Annotations: map[string]string{
-								"autoscaling.knative.dev/class":  "kpa.autoscaling.knative.dev",
-								"autoscaling.knative.dev/target": "1",
+								"autoscaling.knative.dev/class":                             "kpa.autoscaling.knative.dev",
+								"autoscaling.knative.dev/target":                            "1",
+								"internal.serving.kubeflow.org/model-initializer-sourceuri": "gs://testuri",
 							},
 						},
 						Spec: knservingv1alpha1.RevisionSpec{
@@ -103,7 +104,7 @@ func TestKnativeConfigurationReconcile(t *testing.T) {
 									"--port=" + v1alpha1.TensorflowServingGRPCPort,
 									"--rest_api_port=" + v1alpha1.TensorflowServingRestPort,
 									"--model_name=mnist",
-									"--model_base_path=gs://testuri",
+									"--model_base_path=" + constants.DefaultModelLocalMountPath,
 								},
 							},
 						},
@@ -120,8 +121,9 @@ func TestKnativeConfigurationReconcile(t *testing.T) {
 						ObjectMeta: metav1.ObjectMeta{
 							Labels: map[string]string{"serving.kubeflow.org/kfservice": "mnist"},
 							Annotations: map[string]string{
-								"autoscaling.knative.dev/class":  "kpa.autoscaling.knative.dev",
-								"autoscaling.knative.dev/target": "1",
+								"autoscaling.knative.dev/class":                             "kpa.autoscaling.knative.dev",
+								"autoscaling.knative.dev/target":                            "1",
+								"internal.serving.kubeflow.org/model-initializer-sourceuri": "gs://testuri2",
 							},
 						},
 						Spec: knservingv1alpha1.RevisionSpec{
@@ -135,7 +137,7 @@ func TestKnativeConfigurationReconcile(t *testing.T) {
 									"--port=" + v1alpha1.TensorflowServingGRPCPort,
 									"--rest_api_port=" + v1alpha1.TensorflowServingRestPort,
 									"--model_name=mnist",
-									"--model_base_path=gs://testuri2",
+									"--model_base_path=" + constants.DefaultModelLocalMountPath,
 								},
 							},
 						},
@@ -168,8 +170,9 @@ func TestKnativeConfigurationReconcile(t *testing.T) {
 						ObjectMeta: metav1.ObjectMeta{
 							Labels: map[string]string{"serving.kubeflow.org/kfservice": "mnist"},
 							Annotations: map[string]string{
-								"autoscaling.knative.dev/class":  "kpa.autoscaling.knative.dev",
-								"autoscaling.knative.dev/target": "1",
+								"autoscaling.knative.dev/class":                             "kpa.autoscaling.knative.dev",
+								"autoscaling.knative.dev/target":                            "1",
+								"internal.serving.kubeflow.org/model-initializer-sourceuri": "gs://testuri",
 							},
 						},
 						Spec: knservingv1alpha1.RevisionSpec{
@@ -183,7 +186,7 @@ func TestKnativeConfigurationReconcile(t *testing.T) {
 									"--port=" + v1alpha1.TensorflowServingGRPCPort,
 									"--rest_api_port=" + v1alpha1.TensorflowServingRestPort,
 									"--model_name=mnist",
-									"--model_base_path=gs://testuri",
+									"--model_base_path=" + constants.DefaultModelLocalMountPath,
 								},
 							},
 						},
diff --git a/pkg/controller/kfservice/resources/knative/configuration_test.go b/pkg/controller/kfservice/resources/knative/configuration_test.go
--- a/pkg/controller/kfservice/resources/knative/configuration_test.go
+++ b/pkg/controller/kfservice/resources/knative/configuration_test.go
@@ -74,11 +74,12 @@ var defaultConfiguration = &knservingv1alpha1.Configuration{
 			ObjectMeta: metav1.ObjectMeta{
 				Labels: map[string]string{"serving.kubeflow.org/kfservice": "mnist"},
 				Annotations: map[string]string{
-					"autoscaling.knative.dev/class":                "kpa.autoscaling.knative.dev",
-					"autoscaling.knative.dev/target":               "1",
-					"autoscaling.knative.dev/minScale":             "1",
-					"autoscaling.knative.dev/maxScale":             "3",
-					constants.KFServiceGKEAcceleratorAnnotationKey: "nvidia-tesla-t4",
+					"autoscaling.knative.dev/class":                          "kpa.autoscaling.knative.dev",
+					"autoscaling.knative.dev/target":                         "1",
+					"autoscaling.knative.dev/minScale":                       "1",
+					"autoscaling.knative.dev/maxScale":                       "3",
+					constants.KFServiceGKEAcceleratorAnnotationKey:           "nvidia-tesla-t4",
+					constants.ModelInitializerSourceUriInternalAnnotationKey: kfsvc.Spec.Default.Tensorflow.ModelURI,
 				},
 			},
 			Spec: knservingv1alpha1.RevisionSpec{
@@ -95,7 +96,7 @@ var defaultConfiguration = &knservingv1alpha1.Configuration{
 						"--port=" + v1alpha1.TensorflowServingGRPCPort,
 						"--rest_api_port=" + v1alpha1.TensorflowServingRestPort,
 						"--model_name=mnist",
-						"--model_base_path=" + kfsvc.Spec.Default.Tensorflow.ModelURI,
+						"--model_base_path=" + constants.DefaultModelLocalMountPath,
 					},
 				},
 			},
@@ -113,11 +114,12 @@ var canaryConfiguration = &knservingv1alpha1.Configuration{
 			ObjectMeta: metav1.ObjectMeta{
 				Labels: map[string]string{"serving.kubeflow.org/kfservice": "mnist"},
 				Annotations: map[string]string{
-					"autoscaling.knative.dev/class":                "kpa.autoscaling.knative.dev",
-					"autoscaling.knative.dev/target":               "1",
-					"autoscaling.knative.dev/minScale":             "1",
-					"autoscaling.knative.dev/maxScale":             "3",
-					constants.KFServiceGKEAcceleratorAnnotationKey: "nvidia-tesla-t4",
+					"autoscaling.knative.dev/class":                          "kpa.autoscaling.knative.dev",
+					"autoscaling.knative.dev/target":                         "1",
+					"autoscaling.knative.dev/minScale":                       "1",
+					"autoscaling.knative.dev/maxScale":                       "3",
+					constants.KFServiceGKEAcceleratorAnnotationKey:           "nvidia-tesla-t4",
+					constants.ModelInitializerSourceUriInternalAnnotationKey: "s3://test/mnist-2/export",
 				},
 			},
 			Spec: knservingv1alpha1.RevisionSpec{
@@ -131,7 +133,7 @@ var canaryConfiguration = &knservingv1alpha1.Configuration{
 						"--port=" + v1alpha1.TensorflowServingGRPCPort,
 						"--rest_api_port=" + v1alpha1.TensorflowServingRestPort,
 						"--model_name=mnist",
-						"--model_base_path=s3://test/mnist-2/export",
+						"--model_base_path=" + constants.DefaultModelLocalMountPath,
 					},
 				},
 			},
@@ -165,7 +167,7 @@ func TestKnativeConfiguration(t *testing.T) {
 						MinReplicas: 1,
 						MaxReplicas: 3,
 						Tensorflow: &v1alpha1.TensorflowSpec{
-							ModelURI:       "s3://test/mnist/export",
+							ModelURI:       kfsvc.Spec.Default.Tensorflow.ModelURI,
 							RuntimeVersion: "1.13.0",
 						},
 						ServiceAccountName: "testsvcacc",
@@ -214,8 +216,9 @@ func TestKnativeConfiguration(t *testing.T) {
 						ObjectMeta: metav1.ObjectMeta{
 							Labels: map[string]string{"serving.kubeflow.org/kfservice": "sklearn"},
 							Annotations: map[string]string{
-								"autoscaling.knative.dev/class":  "kpa.autoscaling.knative.dev",
-								"autoscaling.knative.dev/target": "1",
+								constants.ModelInitializerSourceUriInternalAnnotationKey: "s3://test/sklearn/export",
+								"autoscaling.knative.dev/class":                          "kpa.autoscaling.knative.dev",
+								"autoscaling.knative.dev/target":                         "1",
 							},
 						},
 						Spec: knservingv1alpha1.RevisionSpec{
@@ -226,7 +229,7 @@ func TestKnativeConfiguration(t *testing.T) {
 								Image: v1alpha1.SKLearnServerImageName + ":" + v1alpha1.DefaultSKLearnRuntimeVersion,
 								Args: []string{
 									"--model_name=sklearn",
-									"--model_dir=s3://test/sklearn/export",
+									"--model_dir=" + constants.DefaultModelLocalMountPath,
 								},
 							},
 						},
@@ -259,8 +262,9 @@ func TestKnativeConfiguration(t *testing.T) {
 						ObjectMeta: metav1.ObjectMeta{
 							Labels: map[string]string{"serving.kubeflow.org/kfservice": "xgboost"},
 							Annotations: map[string]string{
-								"autoscaling.knative.dev/class":  "kpa.autoscaling.knative.dev",
-								"autoscaling.knative.dev/target": "1",
+								constants.ModelInitializerSourceUriInternalAnnotationKey: "s3://test/xgboost/export",
+								"autoscaling.knative.dev/class":                          "kpa.autoscaling.knative.dev",
+								"autoscaling.knative.dev/target":                         "1",
 							},
 						},
 						Spec: knservingv1alpha1.RevisionSpec{
@@ -271,7 +275,7 @@ func TestKnativeConfiguration(t *testing.T) {
 								Image: v1alpha1.XGBoostServerImageName + ":" + v1alpha1.DefaultXGBoostRuntimeVersion,
 								Args: []string{
 									"--model_name=xgboost",
-									"--model_dir=s3://test/xgboost/export",
+									"--model_dir=" + constants.DefaultModelLocalMountPath,
 								},
 							},
 						},
@@ -305,8 +309,9 @@ func TestKnativeConfiguration(t *testing.T) {
 						ObjectMeta: metav1.ObjectMeta{
 							Labels: map[string]string{"serving.kubeflow.org/kfservice": "xgboost"},
 							Annotations: map[string]string{
-								"autoscaling.knative.dev/class":  "kpa.autoscaling.knative.dev",
-								"autoscaling.knative.dev/target": "1",
+								constants.ModelInitializerSourceUriInternalAnnotationKey: "s3://test/xgboost/export",
+								"autoscaling.knative.dev/class":                          "kpa.autoscaling.knative.dev",
+								"autoscaling.knative.dev/target":                         "1",
 							},
 						},
 						Spec: knservingv1alpha1.RevisionSpec{
@@ -317,7 +322,7 @@ func TestKnativeConfiguration(t *testing.T) {
 								Image: "kfserving/xgbserver:" + v1alpha1.DefaultXGBoostRuntimeVersion,
 								Args: []string{
 									"--model_name=xgboost",
-									"--model_dir=s3://test/xgboost/export",
+									"--model_dir=" + constants.DefaultModelLocalMountPath,
 								},
 							},
 						},
@@ -331,7 +336,7 @@ func TestKnativeConfiguration(t *testing.T) {
 		configurationBuilder := NewConfigurationBuilder(c, &v1.ConfigMap{
 			Data: scenario.configMapData,
 		})
-		defaultConfiguration, err := configurationBuilder.CreateKnativeConfiguration(
+		actualDefaultConfiguration, err := configurationBuilder.CreateKnativeConfiguration(
 			constants.DefaultConfigurationName(scenario.kfService.Name),
 			scenario.kfService.ObjectMeta,
 			&scenario.kfService.Spec.Default,
@@ -340,12 +345,12 @@ func TestKnativeConfiguration(t *testing.T) {
 			t.Errorf("Test %q unexpected error %s", name, err.Error())
 		}
 
-		if diff := cmp.Diff(scenario.expectedDefault, defaultConfiguration); diff != "" {
+		if diff := cmp.Diff(scenario.expectedDefault, actualDefaultConfiguration); diff != "" {
 			t.Errorf("Test %q unexpected default configuration (-want +got): %v", name, diff)
 		}
 
 		if scenario.kfService.Spec.Canary != nil {
-			canaryConfiguration, err := configurationBuilder.CreateKnativeConfiguration(
+			actualCanaryConfiguration, err := configurationBuilder.CreateKnativeConfiguration(
 				constants.CanaryConfigurationName(kfsvc.Name),
 				scenario.kfService.ObjectMeta,
 				scenario.kfService.Spec.Canary,
@@ -353,7 +358,7 @@ func TestKnativeConfiguration(t *testing.T) {
 			if err != nil {
 				t.Errorf("Test %q unexpected error %s", name, err.Error())
 			}
-			if diff := cmp.Diff(scenario.expectedCanary, canaryConfiguration); diff != "" {
+			if diff := cmp.Diff(scenario.expectedCanary, actualCanaryConfiguration); diff != "" {
 				t.Errorf("Test %q unexpected canary configuration (-want +got): %v", name, diff)
 			}
 		}
diff --git a/pkg/webhook/admission/deployment/accelerator_injector_test.go b/pkg/webhook/admission/deployment/accelerator_injector_test.go
--- a/pkg/webhook/admission/deployment/accelerator_injector_test.go
+++ b/pkg/webhook/admission/deployment/accelerator_injector_test.go
@@ -111,7 +111,7 @@ func TestAcceleratorInjector(t *testing.T) {
 	}
 
 	for name, scenario := range scenarios {
-		mutate(scenario.original)
+		InjectGKEAcceleratorSelector(scenario.original)
 		// cmd.Diff complains on ResourceList when Nvidia is key. Objects are explicitly compared
 		if diff := cmp.Diff(
 			scenario.expected.Spec.Template.Spec.NodeSelector,
diff --git a/pkg/webhook/admission/deployment/model_initializer_injector_test.go b/pkg/webhook/admission/deployment/model_initializer_injector_test.go
new file mode 100644
--- /dev/null
+++ b/pkg/webhook/admission/deployment/model_initializer_injector_test.go
@@ -0,0 +1,263 @@
+/*
+Copyright 2019 kubeflow.org.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package deployment
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/kubeflow/kfserving/pkg/constants"
+	appsv1 "k8s.io/api/apps/v1"
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func TestModelInitializerInjector(t *testing.T) {
+	scenarios := map[string]struct {
+		original *appsv1.Deployment
+		expected *appsv1.Deployment
+	}{
+		"MissingAnnotations": {
+			original: &appsv1.Deployment{
+				Spec: appsv1.DeploymentSpec{
+					Template: v1.PodTemplateSpec{
+						ObjectMeta: metav1.ObjectMeta{
+							Annotations: map[string]string{},
+						},
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Name: "user-container",
+								},
+							},
+						},
+					},
+				},
+			},
+			expected: &appsv1.Deployment{
+				Spec: appsv1.DeploymentSpec{
+					Template: v1.PodTemplateSpec{
+						ObjectMeta: metav1.ObjectMeta{
+							Annotations: map[string]string{},
+						},
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Name: "user-container",
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		"ModelInitializerInjected": {
+			original: &appsv1.Deployment{
+				Spec: appsv1.DeploymentSpec{
+					Template: v1.PodTemplateSpec{
+						ObjectMeta: metav1.ObjectMeta{
+							Annotations: map[string]string{
+								constants.ModelInitializerSourceUriInternalAnnotationKey: "gs://foo",
+							},
+						},
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Name: "user-container",
+								},
+							},
+						},
+					},
+				},
+			},
+			expected: &appsv1.Deployment{
+				Spec: appsv1.DeploymentSpec{
+					Template: v1.PodTemplateSpec{
+						ObjectMeta: metav1.ObjectMeta{
+							Annotations: map[string]string{
+								constants.ModelInitializerSourceUriInternalAnnotationKey: "gs://foo",
+							},
+						},
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Name: "user-container",
+									VolumeMounts: []v1.VolumeMount{
+										{
+											Name:      "kfserving-provision-location",
+											MountPath: constants.DefaultModelLocalMountPath,
+											ReadOnly:  true,
+										},
+									},
+								},
+							},
+							InitContainers: []v1.Container{
+								v1.Container{
+									Name:  "model-initializer",
+									Image: ModelInitializerContainerImage + ":" + ModelInitializerContainerVersion,
+									Args:  []string{"gs://foo", constants.DefaultModelLocalMountPath},
+									VolumeMounts: []v1.VolumeMount{
+										{
+											Name:      "kfserving-provision-location",
+											MountPath: constants.DefaultModelLocalMountPath,
+										},
+									},
+								},
+							},
+							Volumes: []v1.Volume{
+								v1.Volume{
+									Name: "kfserving-provision-location",
+									VolumeSource: v1.VolumeSource{
+										EmptyDir: &v1.EmptyDirVolumeSource{},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		"ModelInitializerInjectedAndMountsPvc": {
+			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",
+								},
+							},
+						},
+					},
+				},
+			},
+			expected: &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",
+									VolumeMounts: []v1.VolumeMount{
+										{
+											Name:      "kfserving-provision-location",
+											MountPath: constants.DefaultModelLocalMountPath,
+											ReadOnly:  true,
+										},
+									},
+								},
+							},
+							InitContainers: []v1.Container{
+								v1.Container{
+									Name:  "model-initializer",
+									Image: ModelInitializerContainerImage + ":" + ModelInitializerContainerVersion,
+									Args:  []string{"/mnt/pvc/some/path/on/pvc", constants.DefaultModelLocalMountPath},
+									VolumeMounts: []v1.VolumeMount{
+										{
+											Name:      "kfserving-pvc-source",
+											MountPath: "/mnt/pvc",
+											ReadOnly:  true,
+										},
+										{
+											Name:      "kfserving-provision-location",
+											MountPath: constants.DefaultModelLocalMountPath,
+										},
+									},
+								},
+							},
+							Volumes: []v1.Volume{
+								v1.Volume{
+									Name: "kfserving-pvc-source",
+									VolumeSource: v1.VolumeSource{
+										PersistentVolumeClaim: &v1.PersistentVolumeClaimVolumeSource{
+											ClaimName: "mypvcname",
+											ReadOnly:  false,
+										},
+									},
+								},
+								v1.Volume{
+									Name: "kfserving-provision-location",
+									VolumeSource: v1.VolumeSource{
+										EmptyDir: &v1.EmptyDirVolumeSource{},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	for name, scenario := range scenarios {
+		if err := InjectModelInitializer(scenario.original); err != nil {
+			t.Errorf("Test %q unexpected result: %s", name, err)
+		}
+		if diff := cmp.Diff(scenario.expected.Spec, scenario.original.Spec); diff != "" {
+			t.Errorf("Test %q unexpected result (-want +got): %v", name, diff)
+		}
+	}
+}
+
+func TestModelInitializerFailureCases(t *testing.T) {
+	scenarios := map[string]struct {
+		original            *appsv1.Deployment
+		expectedErrorPrefix string
+	}{
+		"MissingUserContainer": {
+			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: "random-container",
+								},
+							},
+						},
+					},
+				},
+			},
+			expectedErrorPrefix: "Invalid configuration: cannot find container",
+		},
+	}
+
+	for name, scenario := range scenarios {
+		if err := InjectModelInitializer(scenario.original); err != nil {
+			if !strings.HasPrefix(err.Error(), scenario.expectedErrorPrefix) {
+				t.Errorf("Test %q unexpected failure [%s], expected: %s", name, err.Error(), scenario.expectedErrorPrefix)
+			}
+		} else {
+			t.Errorf("Test %q should have failed with: %s", name, scenario.expectedErrorPrefix)
+		}
+	}
+}