Merge branch 'master' into issue_640

pytorch · Sep 30, 2020 · 7c38aab · 7c38aab
2 parents e1819bb + 0bc6de0
commit 7c38aab
Show file tree

Hide file tree

Showing 7 changed files with 166 additions and 59 deletions.
diff --git a/docs/configuration.md b/docs/configuration.md
@@ -206,11 +206,11 @@ By default, TorchServe uses all available GPUs for inference. Use `number_of_gpu
 Most of the following properties are designed for performance tuning. Adjusting these numbers will impact scalability and throughput.
 
 * `enable_envvars_config`: Enable configuring TorchServe through environment variables. When this option is set to "true", all the static configurations of TorchServe can come through environment variables as well. Default: false
-* `number_of_netty_threads`: number frontend netty thread. This specifies the numer of threads in the child [EventLoopGroup](https://livebook.manning.com/book/netty-in-action/chapter-8) of the frontend netty server. This group provides EventLoops for processing Netty Channel events (namely inference and management requests) from accepted connections. Default: number of logical processors available to the JVM.
-* `netty_client_threads`: number of backend netty thread. This specifies the number of threads in the WorkerThread [EventLoopGroup](https://livebook.manning.com/book/netty-in-action/chapter-8) which writes inference responses to the frontend. Default: number of logical processors available to the JVM.
-* `default_workers_per_model`: number of workers to create for each model that loaded at startup time. Default: available GPUs in system or number of logical processors available to the JVM.
-* `job_queue_size`: number inference jobs that frontend will queue before backend can serve. Default: 100.
-* `async_logging`: enable asynchronous logging for higher throughput, log output may be delayed if this is enabled. Default: false.
+* `number_of_netty_threads`: Number frontend netty thread. This specifies the numer of threads in the child [EventLoopGroup](https://livebook.manning.com/book/netty-in-action/chapter-8) of the frontend netty server. This group provides EventLoops for processing Netty Channel events (namely inference and management requests) from accepted connections. Default: number of logical processors available to the JVM.
+* `netty_client_threads`: Number of backend netty thread. This specifies the number of threads in the WorkerThread [EventLoopGroup](https://livebook.manning.com/book/netty-in-action/chapter-8) which writes inference responses to the frontend. Default: number of logical processors available to the JVM.
+* `default_workers_per_model`: Number of workers to create for each model that loaded at startup time. Default: available GPUs in system or number of logical processors available to the JVM.
+* `job_queue_size`: Number inference jobs that frontend will queue before backend can serve. Default: 100.
+* `async_logging`: Enable asynchronous logging for higher throughput, log output may be delayed if this is enabled. Default: false.
 * `default_response_timeout`: Timeout, in seconds, used for model's backend workers before they are deemed unresponsive and rebooted. Default: 120 seconds.
 * `unregister_model_timeout`: Timeout, in seconds, used when handling an unregister model request when cleaning a process before it is deemed unresponsive and an error response is sent. Default: 120 seconds.
 * `decode_input_request`: Configuration to let backend workers to decode requests, when the content type is known.
@@ -220,8 +220,14 @@ the backend workers convert "Bytearray to utf-8 string" when the Content-Type of
 * `model_server_home` : Torchserve home directory. 
 * `max_request_size` : The maximum allowable request size that the Torchserve accepts, in bytes. Default: 6553500
 * `max_response_size` : The maximum allowable response size that the Torchserve sends, in bytes. Default: 6553500
+* `allowed_urls` : Comma separated regex of allowed source URL(s) from where models can be registered. Default: "file://.*|http(s)?://.*" (all URLs and local file system)
+eg : To allow base URLs "https://s3.amazonaws.com/" and "https://torchserve.s3.amazonaws.com/" use following regex string
+```
+allowed_urls=https://s3.amazonaws.com/.*,https://torchserve.s3.amazonaws.com/.*
+```
 
 ---
+
 **NOTE**
 
 All the above config properties can be set using environment variable as follows.

diff --git a/docs/model_zoo.md b/docs/model_zoo.md
@@ -8,25 +8,25 @@ To propose a model for inclusion, please submit a [pull request](https://github.
 
 | Model | Type | Dataset | Size | Download | Sample Input| Model mode |
 | --- | --- | --- | --- | --- | --- | --- |
-| AlexNet | Image Classification | ImageNet | 216 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/alexnet.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Eager|
-| Densenet161 | Image Classification | ImageNet | 106 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/densenet161.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Eager|
-| Resnet18 | Image Classification | ImageNet | 41 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/resnet-18.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Eager|
-| VGG11 | Image Classification | ImageNet | 471 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/vgg11_v2.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Eager|
-| Squeezenet 1_1 | Image Classification | ImageNet | 4.4 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/squeezenet1_1.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Eager|
-| MNIST digit classifier | Image Classification | MNIST | 4.3 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/mnist_v2.mar) | [0.png](https://github.com/pytorch/serve/blob/master/examples/image_classifier/mnist/test_data/0.png) |Eager|
-| Resnet 152 |Image Classification | ImageNet | 214 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/resnet-152-batch_v2.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Eager|
-| Faster RCNN | Object Detection | COCO | 148 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/fastrcnn.mar) | [persons.jpg](https://github.com/pytorch/serve/blob/master/examples/object_detector/persons.jpg) |Eager|
-| MASK RCNN | Object Detection | COCO | 158 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/maskrcnn.mar) | [persons.jpg](https://github.com/pytorch/serve/blob/master/examples/object_detector/persons.jpg) |Eager|
-| Text classifier | Text Classification | AG_NEWS | 169 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/my_text_classifier_v2.mar) | [sample_text.txt](https://github.com/pytorch/serve/blob/master/examples/text_classification/sample_text.txt) |Eager|
-| FCN Resenet 101 | Image Segmentation | COCO | 193 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/fcn_resnet_101.mar) | [persons.jpg](https://github.com/pytorch/serve/blob/master/examples/image_segmenter/fcn/persons.jpg) |Eager|
-| AlexNet Scripted | Image Classification | ImageNet | 216 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/alexnet_scripted.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Torchscripted |
-| Densenet161 Scripted| Image Classification | ImageNet | 105 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/densenet161_scripted.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Torchscripted |
-| Resnet18 Scripted| Image Classification | ImageNet | 42 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/resnet-18_scripted.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Torchscripted |
-| VGG11 Scripted| Image Classification | ImageNet | 471 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/vgg11_scripted.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Torchscripted |
-| Squeezenet 1_1 Scripted | Image Classification | ImageNet | 4.4 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/squeezenet1_1_scripted.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Torchscripted |
-| MNIST digit classifier Scripted | Image Classification | MNIST | 4.3 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/mnist_scripted_v2.mar) | [0.png](https://github.com/pytorch/serve/blob/master/examples/image_classifier/mnist/test_data/0.png) |Torchscripted |
-| Resnet 152 Scripted |Image Classification | ImageNet | 215 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/resnet-152-scripted_v2.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Torchscripted |
-| Text classifier Scripted | Text Classification | AG_NEWS | 169 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/my_text_classifier_scripted_v2.mar) | [sample_text.txt](https://github.com/pytorch/serve/blob/master/examples/text_classification/sample_text.txt) |Torchscripted |
-| FCN Resenet 101 Scripted | Image Segmentation | COCO | 193 MB | [.mar](https://torchserve.s3.amazonaws.com/mar_files/fcn_resnet_101_scripted.mar) | [persons.jpg](https://github.com/pytorch/serve/blob/master/examples/image_segmenter/fcn/persons.jpg) |Torchscripted |
+| AlexNet | Image Classification | ImageNet | 216 MB | [.mar](https://torchserve.pytorch.org/mar_files/alexnet.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Eager|
+| Densenet161 | Image Classification | ImageNet | 106 MB | [.mar](https://torchserve.pytorch.org/mar_files/densenet161.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Eager|
+| Resnet18 | Image Classification | ImageNet | 41 MB | [.mar](https://torchserve.pytorch.org/mar_files/resnet-18.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Eager|
+| VGG11 | Image Classification | ImageNet | 471 MB | [.mar](https://torchserve.pytorch.org/mar_files/vgg11_v2.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Eager|
+| Squeezenet 1_1 | Image Classification | ImageNet | 4.4 MB | [.mar](https://torchserve.pytorch.org/mar_files/squeezenet1_1.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Eager|
+| MNIST digit classifier | Image Classification | MNIST | 4.3 MB | [.mar](https://torchserve.pytorch.org/mar_files/mnist_v2.mar) | [0.png](https://github.com/pytorch/serve/blob/master/examples/image_classifier/mnist/test_data/0.png) |Eager|
+| Resnet 152 |Image Classification | ImageNet | 214 MB | [.mar](https://torchserve.pytorch.org/mar_files/resnet-152-batch_v2.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Eager|
+| Faster RCNN | Object Detection | COCO | 148 MB | [.mar](https://torchserve.pytorch.org/mar_files/fastrcnn.mar) | [persons.jpg](https://github.com/pytorch/serve/blob/master/examples/object_detector/persons.jpg) |Eager|
+| MASK RCNN | Object Detection | COCO | 158 MB | [.mar](https://torchserve.pytorch.org/mar_files/maskrcnn.mar) | [persons.jpg](https://github.com/pytorch/serve/blob/master/examples/object_detector/persons.jpg) |Eager|
+| Text classifier | Text Classification | AG_NEWS | 169 MB | [.mar](https://torchserve.pytorch.org/mar_files/my_text_classifier_v2.mar) | [sample_text.txt](https://github.com/pytorch/serve/blob/master/examples/text_classification/sample_text.txt) |Eager|
+| FCN Resenet 101 | Image Segmentation | COCO | 193 MB | [.mar](https://torchserve.pytorch.org/mar_files/fcn_resnet_101.mar) | [persons.jpg](https://github.com/pytorch/serve/blob/master/examples/image_segmenter/fcn/persons.jpg) |Eager|
+| AlexNet Scripted | Image Classification | ImageNet | 216 MB | [.mar](https://torchserve.pytorch.org/mar_files/alexnet_scripted.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Torchscripted |
+| Densenet161 Scripted| Image Classification | ImageNet | 105 MB | [.mar](https://torchserve.pytorch.org/mar_files/densenet161_scripted.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Torchscripted |
+| Resnet18 Scripted| Image Classification | ImageNet | 42 MB | [.mar](https://torchserve.pytorch.org/mar_files/resnet-18_scripted.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Torchscripted |
+| VGG11 Scripted| Image Classification | ImageNet | 471 MB | [.mar](https://torchserve.pytorch.org/mar_files/vgg11_scripted.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Torchscripted |
+| Squeezenet 1_1 Scripted | Image Classification | ImageNet | 4.4 MB | [.mar](https://torchserve.pytorch.org/mar_files/squeezenet1_1_scripted.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Torchscripted |
+| MNIST digit classifier Scripted | Image Classification | MNIST | 4.3 MB | [.mar](https://torchserve.pytorch.org/mar_files/mnist_scripted_v2.mar) | [0.png](https://github.com/pytorch/serve/blob/master/examples/image_classifier/mnist/test_data/0.png) |Torchscripted |
+| Resnet 152 Scripted |Image Classification | ImageNet | 215 MB | [.mar](https://torchserve.pytorch.org/mar_files/resnet-152-scripted_v2.mar) | [kitten.jpg](https://github.com/pytorch/serve/blob/master/examples/image_classifier/kitten.jpg) |Torchscripted |
+| Text classifier Scripted | Text Classification | AG_NEWS | 169 MB | [.mar](https://torchserve.pytorch.org/mar_files/my_text_classifier_scripted_v2.mar) | [sample_text.txt](https://github.com/pytorch/serve/blob/master/examples/text_classification/sample_text.txt) |Torchscripted |
+| FCN Resenet 101 Scripted | Image Segmentation | COCO | 193 MB | [.mar](https://torchserve.pytorch.org/mar_files/fcn_resnet_101_scripted.mar) | [persons.jpg](https://github.com/pytorch/serve/blob/master/examples/image_segmenter/fcn/persons.jpg) |Torchscripted |
 
 Refer [example](../examples) for more details on above models.
diff --git a/frontend/modelarchive/src/main/java/org/pytorch/serve/archive/ModelArchive.java b/frontend/modelarchive/src/main/java/org/pytorch/serve/archive/ModelArchive.java
@@ -15,6 +15,7 @@
 import java.security.DigestInputStream;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
+import java.util.List;
 import java.util.UUID;
 import java.util.regex.Pattern;
 import org.apache.commons.io.FileUtils;
@@ -28,7 +29,7 @@ public class ModelArchive {
 
     public static final Gson GSON = new GsonBuilder().setPrettyPrinting().create();
 
-    private static final Pattern URL_PATTERN =
+    private static final Pattern VALID_URL_PATTERN =
             Pattern.compile("file?://.*|http(s)?://.*", Pattern.CASE_INSENSITIVE);
 
     private static final String MANIFEST_FILE = "MANIFEST.json";
@@ -45,7 +46,8 @@ public ModelArchive(Manifest manifest, String url, File modelDir, boolean extrac
         this.extracted = extracted;
     }
 
-    public static ModelArchive downloadModel(String modelStore, String url)
+    public static ModelArchive downloadModel(
+            List<String> allowedUrls, String modelStore, String url)
             throws ModelException, FileAlreadyExistsException, IOException {
 
         if (modelStore == null) {
@@ -54,7 +56,8 @@ public static ModelArchive downloadModel(String modelStore, String url)
 
         String marFileName = FilenameUtils.getName(url);
         File modelLocation = new File(modelStore, marFileName);
-        if (URL_PATTERN.matcher(url).matches()) {
+
+        if (checkAllowedUrl(allowedUrls, url)) {
             if (modelLocation.exists()) {
                 throw new FileAlreadyExistsException(marFileName);
             }
@@ -84,6 +87,23 @@ public static ModelArchive downloadModel(String modelStore, String url)
         throw new ModelNotFoundException("Model not found at: " + url);
     }
 
+    public static boolean checkAllowedUrl(List<String> allowedUrls, String url)
+            throws ModelNotFoundException {
+        boolean patternMatch = false;
+        for (String temp : allowedUrls) {
+            if (Pattern.compile(temp, Pattern.CASE_INSENSITIVE).matcher(url).matches()) {
+                patternMatch = true;
+                return patternMatch;
+            }
+        }
+        if (VALID_URL_PATTERN.matcher(url).matches()) {
+            // case when url is valid url but does not match valid hosts
+            throw new ModelNotFoundException(
+                    "Given URL " + url + " does not match any allowed URL(s)");
+        }
+        return patternMatch;
+    }
+
     private static ModelArchive load(String url, File dir, boolean extracted)
             throws InvalidModelException, IOException {
         boolean failed = true;
@@ -178,7 +198,7 @@ public void validate() throws InvalidModelException {
     }
 
     public static void removeModel(String modelStore, String marURL) {
-        if (URL_PATTERN.matcher(marURL).matches()) {
+        if (VALID_URL_PATTERN.matcher(marURL).matches()) {
             String marFileName = FilenameUtils.getName(marURL);
             File modelLocation = new File(modelStore, marFileName);
             FileUtils.deleteQuietly(modelLocation);