Optional model for predictor (#549)

vishalbollu · web-flow · commit 119bca2ce9aa · 2019-11-04T16:38:52.000-05:00
* Introduce optional model path to predictor API

* Fix sklearn mpg model path

* Allow folders to be accessed correctly after being downloaded from s3

* Update example readme and docs

* Address PR comments
diff --git a/docs/deployments/apis.md b/docs/deployments/apis.md
@@ -15,6 +15,7 @@ In addition to supporting Python models via the Predictor interface, Cortex can
   endpoint: <string>  # the endpoint for the API (default: /<deployment_name>/<api_name>)
   predictor:
     path: <string>  # path to the predictor Python file, relative to the Cortex root (required)
+    model: <string>  # S3 path to a file or directory (e.g. s3://my-bucket/exported_model) (optional)
     python_path: <string>  # path to the root of your Python folder that will be appended to PYTHONPATH (default: folder containing cortex.yaml)
     metadata: <string: value>  # dictionary that can be used to configure custom values (optional)
   tracker:
diff --git a/docs/deployments/predictor.md b/docs/deployments/predictor.md
@@ -11,12 +11,13 @@ Global variables can be shared across functions safely because each replica hand
 ```python
 # initialization code and variables can be declared here in global scope
 
-def init(metadata):
+def init(model_path, metadata):
     """Called once before the API is made available. Setup for model serving such
     as downloading/initializing the model or downloading vocabulary can be done here.
     Optional.
 
     Args:
+        model_path: Local path to model file or directory if specified by user in API configuration, otherwise None.
         metadata: Custom dictionary specified by the user in API configuration.
     """
     pass
@@ -45,14 +46,9 @@ from my_model import IrisNet
 labels = ["iris-setosa", "iris-versicolor", "iris-virginica"]
 model = IrisNet()
 
-def init(metadata):
-    # Download model from S3 (location specified in the API's metadata)
-    s3 = boto3.client("s3")
-    bucket, key = re.match(r"s3:\/\/(.+?)\/(.+)", metadata["model"]).groups()
-    s3.download_file(bucket, key, "iris_model.pth")
-
+def init(model_path, metadata):
     # Initialize the model
-    model.load_state_dict(torch.load("iris_model.pth"))
+    model.load_state_dict(torch.load(model_path))
     model.eval()
 
 
diff --git a/examples/pytorch/iris-classifier/README.md b/examples/pytorch/iris-classifier/README.md
@@ -8,26 +8,22 @@ We implement Cortex's Predictor interface to load the model and make predictions
 
 ### Initialization
 
-We can place our code to download and initialize the model in the `init()` function. The PyTorch model class is defined in [src/my_model.py](./src/my_model.py), and we assume that we've already trained the model and uploaded the state_dict (weights) to S3.
+We can place our code to download and initialize the model in the `init()` function. The PyTorch model class is defined in [src/model.py](./src/model.py), and we assume that we've already trained the model and uploaded the state_dict (weights) to S3.
 
 ```python
 # predictor.py
 
-from my_model import IrisNet
+from model import IrisNet
 
 # instantiate the model
 model = IrisNet()
 
 # define the labels
 labels = ["iris-setosa", "iris-versicolor", "iris-virginica"]
 
-def init(metadata):
-    # download the model from S3 (location specified in the metadata field of our api configuration)
-    s3 = boto3.client("s3")
-    bucket, key = re.match(r"s3:\/\/(.+?)\/(.+)", metadata["model"]).groups()
-    s3.download_file(bucket, key, "weights.pth")
-
-    model.load_state_dict(torch.load("weights.pth"))
+def init(model_path, metadata):
+    # model_path is a local path pointing to your model weights file
+    model.load_state_dict(torch.load(model_path))
     model.eval()
 ```
 
@@ -71,8 +67,7 @@ A `deployment` specifies a set of resources that are deployed together. An `api`
   predictor:
     path: src/predictor.py
     python_path: src/
-    metadata:
-      model: s3://cortex-examples/pytorch/iris-classifier/weights.pth
+    model: s3://cortex-examples/pytorch/iris-classifier/weights.pth
   tracker:
     model_type: classification
 ```
diff --git a/examples/pytorch/iris-classifier/cortex.yaml b/examples/pytorch/iris-classifier/cortex.yaml
@@ -6,7 +6,6 @@
   predictor:
     path: src/predictor.py
     python_path: src/
-    metadata:
-      model: s3://cortex-examples/pytorch/iris-classifier/weights.pth
+    model: s3://cortex-examples/pytorch/iris-classifier/weights.pth
   tracker:
     model_type: classification
diff --git a/examples/pytorch/iris-classifier/src/predictor.py b/examples/pytorch/iris-classifier/src/predictor.py
@@ -1,4 +1,3 @@
-import boto3
 import re
 import torch
 from model import IrisNet
@@ -8,11 +7,8 @@
 model = IrisNet()
 
 
-def init(metadata):
-    s3 = boto3.client("s3")
-    bucket, key = re.match(r"s3:\/\/(.+?)\/(.+)", metadata["model"]).groups()
-    s3.download_file(bucket, key, "weights.pth")
-    model.load_state_dict(torch.load("weights.pth"))
+def init(model_path, metadata):
+    model.load_state_dict(torch.load(model_path))
     model.eval()
 
 
diff --git a/examples/pytorch/text-generator/README.md b/examples/pytorch/text-generator/README.md
@@ -20,7 +20,7 @@ model.eval()
 # download the tokenizer
 tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
 
-def init(metadata):
+def init(model_path, metadata):
     # load the model onto the device specified in the metadata field of our api configuration
     model.to(metadata["device"])
 ```
diff --git a/examples/pytorch/text-generator/predictor.py b/examples/pytorch/text-generator/predictor.py
@@ -46,7 +46,7 @@ def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float("Inf")
     return logits
 
 
-def init(metadata):
+def init(model_path, metadata):
     model.to(metadata["device"])
 
 
diff --git a/examples/sklearn/mpg-estimation/README.md b/examples/sklearn/mpg-estimation/README.md
@@ -8,21 +8,17 @@ We implement Cortex's Predictor interface to load the model and make predictions
 
 ### Initialization
 
-We can place our code to download and initialize the model in the `init()` function:
+We can place our code to initialize the model in the `init()` function:
 
 ```python
 # predictor.py
 
 model = None
 
-def init(metadata):
+def init(model_path, metadata):
     global model
-
-    # download the model from S3 (path specified in the metadata field of our api configuration)
-    s3 = boto3.client("s3")
-    bucket, key = re.match(r"s3:\/\/(.+?)\/(.+)", metadata["model"]).groups()
-    s3.download_file(bucket, key, "linreg.joblib")
-    model = load("linreg.joblib")
+    # model_path is a local path pointing to your model
+    model = load(model_path)
 ```
 
 ### Predict
@@ -61,8 +57,7 @@ A `deployment` specifies a set of resources that are deployed together. An `api`
   name: mpg
   predictor:
     path: predictor.py
-    metadata:
-      model: s3://cortex-examples/sklearn/mpg-estimation/linreg.joblib
+    model: s3://cortex-examples/sklearn/mpg-estimation/linreg.joblib
   tracker:
     model_type: regression
 ```
diff --git a/examples/sklearn/mpg-estimation/cortex.yaml b/examples/sklearn/mpg-estimation/cortex.yaml
@@ -5,7 +5,6 @@
   name: mpg
   predictor:
     path: predictor.py
-    metadata:
-      model: s3://cortex-examples/sklearn/mpg-estimation/linreg.joblib
+    model: s3://cortex-examples/sklearn/mpg-estimation/linreg.joblib
   tracker:
     model_type: regression
diff --git a/examples/sklearn/mpg-estimation/predictor.py b/examples/sklearn/mpg-estimation/predictor.py
@@ -1,18 +1,13 @@
-import boto3
-import re
-import numpy as np
 from joblib import load
+import numpy as np
 
 
 model = None
 
 
-def init(metadata):
+def init(model_path, metadata):
     global model
-    s3 = boto3.client("s3")
-    bucket, key = re.match(r"s3:\/\/(.+?)\/(.+)", metadata["model"]).groups()
-    s3.download_file(bucket, key, "linreg.joblib")
-    model = load("linreg.joblib")
+    model = load(model_path)
 
 
 def predict(sample, metadata):
diff --git a/pkg/operator/api/userconfig/apis.go b/pkg/operator/api/userconfig/apis.go
@@ -89,6 +89,7 @@ type ONNX struct {
 
 type Predictor struct {
 	Path       string                 `json:"path" yaml:"path"`
+	Model      *string                `json:"model" yaml:"model"`
 	PythonPath *string                `json:"python_path" yaml:"python_path"`
 	Metadata   map[string]interface{} `json:"metadata" yaml:"metadata"`
 }
@@ -188,6 +189,12 @@ var apiValidation = &cr.StructValidation{
 							Required: true,
 						},
 					},
+					{
+						StructField: "Model",
+						StringPtrValidation: &cr.StringPtrValidation{
+							Validator: cr.S3PathValidator(),
+						},
+					},
 					pythonPathValidation,
 					metadataValidation,
 				},
@@ -431,6 +438,33 @@ func (predictor *Predictor) Validate(projectFileMap map[string][]byte) error {
 	if _, ok := projectFileMap[predictor.Path]; !ok {
 		return errors.Wrap(ErrorImplDoesNotExist(predictor.Path), PredictorKey, PathKey)
 	}
+	if predictor.Model != nil {
+		var err error
+		awsClient, err := aws.NewFromS3Path(*predictor.Model, false)
+		if err != nil {
+			return err
+		}
+
+		isFile := false
+		if isFile, err = awsClient.IsS3PathFile(*predictor.Model); err != nil {
+			return errors.Wrap(err, PredictorKey, ModelKey)
+		}
+
+		if !isFile {
+			isDir := false
+			modelPath := s.EnsureSuffix(*predictor.Model, "/")
+
+			if isDir, err = awsClient.IsS3PathDir(modelPath); err != nil {
+				return errors.Wrap(err, PredictorKey, ModelKey)
+			}
+
+			if !isDir {
+				return errors.Wrap(ErrorExternalNotFound(*predictor.Model), PredictorKey, ModelKey)
+			}
+
+			predictor.Model = pointer.String(modelPath)
+		}
+	}
 	if predictor.PythonPath != nil {
 		if err := ValidatePythonPath(*predictor.PythonPath, projectFileMap); err != nil {
 			return errors.Wrap(err, PredictorKey)
@@ -442,6 +476,9 @@ func (predictor *Predictor) Validate(projectFileMap map[string][]byte) error {
 func (predictor *Predictor) UserConfigStr() string {
 	var sb strings.Builder
 	sb.WriteString(fmt.Sprintf("%s: %s\n", PathKey, predictor.Path))
+	if predictor.Model != nil {
+		sb.WriteString(fmt.Sprintf("%s: %s\n", ModelKey, *predictor.Model))
+	}
 	if predictor.PythonPath != nil {
 		sb.WriteString(fmt.Sprintf("%s: %s\n", PythonPathKey, *predictor.PythonPath))
 	}
diff --git a/pkg/operator/workloads/api_workload.go b/pkg/operator/workloads/api_workload.go
@@ -461,6 +461,14 @@ func predictorAPISpec(
 		},
 	}
 
+	if api.Predictor.Model != nil {
+		downloadArgs = append(downloadArgs, downloadContainerArg{
+			From:     *api.Predictor.Model,
+			To:       path.Join(consts.EmptyDirMountPath, "model"),
+			ItemName: "model",
+		})
+	}
+
 	downloadArgsBytes, _ := json.Marshal(downloadArgs)
 	downloadArgsStr := base64.URLEncoding.EncodeToString(downloadArgsBytes)
 
@@ -534,6 +542,7 @@ func predictorAPISpec(
 							"--port=" + defaultPortStr,
 							"--context=" + config.AWS.S3Path(ctx.Key),
 							"--api=" + ctx.APIs[api.Name].ID,
+							"--model-dir=" + path.Join(consts.EmptyDirMountPath, "model"),
 							"--cache-dir=" + consts.ContextCacheDir,
 							"--project-dir=" + path.Join(consts.EmptyDirMountPath, "project"),
 						},
diff --git a/pkg/workloads/cortex/lib/context.py b/pkg/workloads/cortex/lib/context.py
@@ -168,7 +168,7 @@ def get_predictor_impl(self, api_name, project_dir):
             raise
 
         try:
-            _validate_impl(impl, INFERENCE_IMPL_VALIDATION)
+            _validate_impl(impl, PREDICTOR_IMPL_VALIDATION)
         except CortexException as e:
             e.wrap("api " + api_name, "predictor " + api["predictor"]["path"])
             raise
@@ -245,8 +245,8 @@ def publish_metrics(self, metrics):
     ]
 }
 
-INFERENCE_IMPL_VALIDATION = {
-    "optional": [{"name": "init", "args": ["metadata"]}],
+PREDICTOR_IMPL_VALIDATION = {
+    "optional": [{"name": "init", "args": ["model_path", "metadata"]}],
     "required": [{"name": "predict", "args": ["sample", "metadata"]}],
 }
 
diff --git a/pkg/workloads/cortex/lib/storage/s3.py b/pkg/workloads/cortex/lib/storage/s3.py
@@ -189,7 +189,7 @@ def download_file(self, key, local_path):
         except Exception as e:
             raise CortexException(
                 'key "{}" in bucket "{}" could not be accessed; '.format(key, self.bucket)
-                + "it may not exist, or you may not have suffienct permissions"
+                + "it may not exist, or you may not have sufficient permissions"
             ) from e
 
     def download_dir(self, prefix, local_dir):
@@ -200,6 +200,8 @@ def download_dir_contents(self, prefix, local_dir):
         util.mkdir_p(local_dir)
         prefix = util.ensure_suffix(prefix, "/")
         for key in self._get_matching_s3_keys_generator(prefix):
+            if key.endswith("/"):
+                continue
             rel_path = util.trim_prefix(key, prefix)
             local_dest_path = os.path.join(local_dir, rel_path)
             self.download_file(key, local_dest_path)
diff --git a/pkg/workloads/cortex/predictor_serve/api.py b/pkg/workloads/cortex/predictor_serve/api.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import sys
 import argparse
 import time
@@ -122,7 +123,14 @@ def start(args):
 
         if util.has_function(local_cache["predictor"], "init"):
             try:
-                local_cache["predictor"].init(api["predictor"]["metadata"])
+                model_path = None
+                if api["predictor"].get("model") is not None:
+                    _, prefix = ctx.storage.deconstruct_s3_path(api["predictor"]["model"])
+
+                    model_path = os.path.join(
+                        args.model_dir, os.path.basename(os.path.normpath(prefix))
+                    )
+                local_cache["predictor"].init(model_path, api["predictor"]["metadata"])
             except Exception as e:
                 raise UserRuntimeException(api["predictor"]["path"], "init", str(e)) from e
         logger.info("init ran successfully")
@@ -150,6 +158,7 @@ def main():
         help="S3 path to context (e.g. s3://bucket/path/to/context.json)",
     )
     na.add_argument("--api", required=True, help="Resource id of api to serve")
+    na.add_argument("--model-dir", required=True, help="Directory to download the model to")
     na.add_argument("--cache-dir", required=True, help="Local path for the context cache")
     na.add_argument("--project-dir", required=True, help="Local path for the project zip file")
 
