Adding model-control-mode (#3165)

* adding model-control-mode * newman test fixes * adding cmd line and fixing pytest * fixes to tests * fixing frontend tests * adding doc for model mode control * updating docs and moving check * adding test and fixing check * adding test * small changes * removing unused files * changing explicit to api enabled * updating tests * fixing test * changing config properties * changing enabled to true * Change test_util.start_torchserve interface --------- Co-authored-by: Matthias Reso <[email protected]>
pytorch · Jun 11, 2024 · 3d17a94 · 3d17a94
1 parent 36049cb
commit 3d17a94
Show file tree

Hide file tree

Showing 27 changed files with 292 additions and 19 deletions.
diff --git a/docs/model_control_mode.md b/docs/model_control_mode.md
@@ -0,0 +1,70 @@
+# Model Control Mode
+
+TorchServe now supports model control mode with two settings "none"(default) and "enabled"
+
+## Two ways to set Model Control
+1. Add `--model-api-enabled` to command line when running TorchServe to switch from none to enabled mode. Command line cannot be used to set mode to none, can only be used to set to enabled
+2. Add `model_api_enabled=false` or `model_api_enabled=true` to config.properties file
+    * `model_api_enabled=false` is default and prevents users from registering or deleting models once TorchServe is running
+    * `model_api_enabled=true` is not default and allows users to register and delete models using the TorchServe model load APIs
+
+Priority between cmd and config file follows the following [TorchServer standard](https://github.com/pytorch/serve/blob/c74a29e8144bc12b84196775076b0e8cf3c5a6fc/docs/configuration.md#advanced-configuration)
+* Example 1:
+  * Config file: `model_api_enabled=false`
+
+    cmd line: `torchserve --start --ncs --model-store model_store --model-api-enabled`
+
+    Result: Model api mode enabled
+* Example 2:
+  * Config file: `model_api_enabled=true`
+
+    cmd line: `torchserve --start --ncs --model-store model_store`
+
+    Result: Mode is enabled (no way to disable api mode through cmd)
+
+## Model Control Mode Default
+At startup TorchServe loads only those models specified explicitly with the `--models` command-line option. After startup users will be unable to register or delete models in this mode.
+
+### Example default
+```
+ubuntu@ip-172-31-11-32:~/serve$ torchserve --start --ncs --model-store model_store --models resnet-18=resnet-18.mar --ts-config config.properties
+...
+ubuntu@ip-172-31-11-32:~/serve$ curl -X POST  "http://localhost:8081/models?url=https://torchserve.pytorch.org/mar_files/squeezenet1_1.mar"
+2024-05-30T21:46:03,625 [INFO ] epollEventLoopGroup-3-2 ACCESS_LOG - /127.0.0.1:53514 "POST /models?url=https://torchserve.pytorch.org/mar_files/squeezenet1_1.mar HTTP/1.1" 405 0
+2024-05-30T21:46:03,626 [INFO ] epollEventLoopGroup-3-2 TS_METRICS - Requests4XX.Count:1.0|#Level:Host|#hostname:ip-172-31-11-32,timestamp:1717105563
+{
+  "code": 405,
+  "type": "MethodNotAllowedException",
+  "message": "Requested method is not allowed, please refer to API document."
+}
+```
+
+## Model Control API Enabled
+Setting model control to `enabled` allows users to load and unload models using the model load APIs.
+
+### Example using cmd line to set mode to enabled
+```
+ubuntu@ip-172-31-11-32:~/serve$ torchserve --start --ncs --model-store model_store --models resnet-18=resnet-18.mar --ts-config config.properties --model-api-enabled
+
+ubuntu@ip-172-31-11-32:~/serve$ curl -X POST  "http://localhost:8081/models?url=https://torchserve.pytorch.org/mar_files/squeezenet1_1.mar"
+{
+  "status": "Model \"squeezenet1_1\" Version: 1.0 registered with 0 initial workers. Use scale workers API to add workers for the model."
+}
+ubuntu@ip-172-31-11-32:~/serve$ curl http://localhost:8081/models
+2024-05-30T21:41:47,098 [INFO ] epollEventLoopGroup-3-2 ACCESS_LOG - /127.0.0.1:36270 "GET /models HTTP/1.1" 200 2
+2024-05-30T21:41:47,099 [INFO ] epollEventLoopGroup-3-2 TS_METRICS - Requests2XX.Count:1.0|#Level:Host|#hostname:ip-172-31-11-32,timestamp:1717105307
+{
+  "models": [
+    {
+      "modelName": "resnet-18",
+      "modelUrl": "resnet-18.mar"
+    },
+    {
+      "modelName": "squeezenet1_1",
+      "modelUrl": "https://torchserve.pytorch.org/mar_files/squeezenet1_1.mar"
+    }
+  ]
+}
+ubuntu@ip-172-31-11-32:~/serve$ torchserve --stop
+TorchServe has stopped.
+```
diff --git a/frontend/server/src/main/java/org/pytorch/serve/http/api/rest/ManagementRequestHandler.java b/frontend/server/src/main/java/org/pytorch/serve/http/api/rest/ManagementRequestHandler.java
@@ -32,6 +32,7 @@
 import org.pytorch.serve.openapi.OpenApiUtils;
 import org.pytorch.serve.servingsdk.ModelServerEndpoint;
 import org.pytorch.serve.util.ApiUtils;
+import org.pytorch.serve.util.ConfigManager;
 import org.pytorch.serve.util.JsonUtils;
 import org.pytorch.serve.util.NettyUtils;
 import org.pytorch.serve.util.messages.RequestInput;
@@ -48,9 +49,12 @@
  */
 public class ManagementRequestHandler extends HttpRequestHandlerChain {
 
+    private ConfigManager configManager;
+
     /** Creates a new {@code ManagementRequestHandler} instance. */
     public ManagementRequestHandler(Map<String, ModelServerEndpoint> ep) {
         endpointMap = ep;
+        configManager = ConfigManager.getInstance();
     }
 
     @Override
@@ -74,7 +78,7 @@ public void handleRequest(
                     if (HttpMethod.GET.equals(method)) {
                         handleListModels(ctx, decoder);
                         return;
-                    } else if (HttpMethod.POST.equals(method)) {
+                    } else if (HttpMethod.POST.equals(method) && isModeEnabled()) {
                         handleRegisterModel(ctx, decoder, req);
                         return;
                     }
@@ -93,7 +97,7 @@ public void handleRequest(
                     } else {
                         handleScaleModel(ctx, decoder, segments[2], modelVersion);
                     }
-                } else if (HttpMethod.DELETE.equals(method)) {
+                } else if (HttpMethod.DELETE.equals(method) && isModeEnabled()) {
                     handleUnregisterModel(ctx, segments[2], modelVersion);
                 } else if (HttpMethod.OPTIONS.equals(method)) {
                     ModelManager modelManager = ModelManager.getInstance();
@@ -129,6 +133,10 @@ private boolean isManagementReq(String[] segments) {
                 || endpointMap.containsKey(segments[1]);
     }
 
+    private boolean isModeEnabled() {
+        return configManager.getModelControlMode();
+    }
+
     private boolean isKFV1ManagementReq(String[] segments) {
         return segments.length == 4 && "v1".equals(segments[1]) && "models".equals(segments[2]);
     }

diff --git a/frontend/server/src/main/java/org/pytorch/serve/util/ConfigManager.java b/frontend/server/src/main/java/org/pytorch/serve/util/ConfigManager.java
@@ -133,6 +133,7 @@ public final class ConfigManager {
     private static final String MODEL_CONFIG = "models";
     private static final String VERSION = "version";
     private static final String SYSTEM_METRICS_CMD = "system_metrics_cmd";
+    private static final String MODEL_CONTROL_MODE = "model_api_enabled";
 
     // Configuration default values
     private static final String DEFAULT_TS_ALLOWED_URLS = "file://.*|http(s)?://.*";
@@ -254,6 +255,10 @@ private ConfigManager(Arguments args) throws IOException {
             prop.setProperty(TS_LOAD_MODELS, String.join(",", models));
         }
 
+        if (args.isModelEnabled().equals("true")) {
+            prop.setProperty(MODEL_CONTROL_MODE, args.isModelEnabled());
+        }
+
         prop.setProperty(
                 TS_NUMBER_OF_GPU,
                 String.valueOf(
@@ -478,6 +483,10 @@ public int getNumberOfGpu() {
         return getIntProperty(TS_NUMBER_OF_GPU, 0);
     }
 
+    public boolean getModelControlMode() {
+        return Boolean.parseBoolean(getProperty(MODEL_CONTROL_MODE, "false"));
+    }
+
     public String getMetricsConfigPath() {
         String path = getCanonicalPath(prop.getProperty(TS_METRICS_CONFIG));
         if (path == null) {
@@ -806,7 +815,9 @@ public String dumpConfigurations() {
                 + "\nModel config: "
                 + prop.getProperty(MODEL_CONFIG, "N/A")
                 + "\nSystem metrics command: "
-                + (getSystemMetricsCmd().isEmpty() ? "default" : getSystemMetricsCmd());
+                + (getSystemMetricsCmd().isEmpty() ? "default" : getSystemMetricsCmd())
+                + "\nModel API enabled: "
+                + (getModelControlMode() ? "true" : "false");
     }
 
     public boolean useNativeIo() {
@@ -1117,6 +1128,7 @@ public static final class Arguments {
         private boolean snapshotDisabled;
         private String workflowStore;
         private String cppLogConfigFile;
+        private boolean modelApiEnabled;
 
         public Arguments() {}
 
@@ -1128,6 +1140,7 @@ public Arguments(CommandLine cmd) {
             snapshotDisabled = cmd.hasOption("no-config-snapshot");
             workflowStore = cmd.getOptionValue("workflow-store");
             cppLogConfigFile = cmd.getOptionValue("cpp-log-config");
+            modelApiEnabled = cmd.hasOption("model-api-enabled");
         }
 
         public static Options getOptions() {
@@ -1180,6 +1193,12 @@ public static Options getOptions() {
                             .argName("CPP-LOG-CONFIG")
                             .desc("log configuration file for cpp backend.")
                             .build());
+            options.addOption(
+                    Option.builder("mapi")
+                            .longOpt("model-api-enabled")
+                            .argName("MODEL-API-ENABLED")
+                            .desc("sets model apis to enabled")
+                            .build());
             return options;
         }
 
@@ -1215,6 +1234,10 @@ public void setModels(String[] models) {
             this.models = models.clone();
         }
 
+        public String isModelEnabled() {
+            return String.valueOf(modelApiEnabled);
+        }
+
         public boolean isSnapshotDisabled() {
             return snapshotDisabled;
         }

diff --git a/frontend/server/src/test/resources/config.properties b/frontend/server/src/test/resources/config.properties
@@ -46,3 +46,4 @@ models={\
 # install_py_dep_per_model=false
 # enable_metrics_api=false
 workflow_store=../archive/src/test/resources/workflows
+model_api_enabled=true
diff --git a/frontend/server/src/test/resources/config_snapshot.properties b/frontend/server/src/test/resources/config_snapshot.properties
@@ -33,3 +33,4 @@ enable_envvars_config=true
 # enable_metrics_api=false
 workflow_store=../archive/src/test/resources/workflows
 metrics_config=src/test/resources/metrics_default.yaml
+model_api_enabled=true
diff --git a/frontend/server/src/test/resources/config_test_env.properties b/frontend/server/src/test/resources/config_test_env.properties
@@ -53,3 +53,4 @@ models={\
   }\
 }
 metrics_config=src/test/resources/metrics_default.yaml
+model_api_enabled=true
diff --git a/frontend/server/src/test/resources/config_test_workflow.properties b/frontend/server/src/test/resources/config_test_workflow.properties
@@ -53,3 +53,4 @@ models={\
   }\
 }
 metrics_config=src/test/resources/metrics_default.yaml
+model_api_enabled=true
diff --git a/frontend/server/src/test/resources/snapshots/snapshot1.cfg b/frontend/server/src/test/resources/snapshots/snapshot1.cfg
@@ -18,4 +18,5 @@ certificate_file=src/test/resources/certs.pem
 management_address=https\://127.0.0.1\:8444
 metrics_address=https\://127.0.0.1\:8445
 workflow_store=../archive/src/test/resources/workflows
-metrics_config=src/test/resources/metrics_default.yaml
+metrics_config=src/test/resources/metrics_default.yaml
+model_api_enabled=true
diff --git a/frontend/server/src/test/resources/snapshots/snapshot2.cfg b/frontend/server/src/test/resources/snapshots/snapshot2.cfg
@@ -18,4 +18,5 @@ certificate_file=src/test/resources/certs.pem
 management_address=https\://127.0.0.1\:8444
 metrics_address=https\://127.0.0.1\:8445
 workflow_store=../archive/src/test/resources/workflows
-metrics_config=src/test/resources/metrics_default.yaml
+metrics_config=src/test/resources/metrics_default.yaml
+model_api_enabled=true
diff --git a/frontend/server/src/test/resources/snapshots/snapshot3.cfg b/frontend/server/src/test/resources/snapshots/snapshot3.cfg
@@ -18,4 +18,5 @@ certificate_file=src/test/resources/certs.pem
 management_address=https\://127.0.0.1\:8444
 metrics_address=https\://127.0.0.1\:8445
 workflow_store=../archive/src/test/resources/workflows
-metrics_config=src/test/resources/metrics_default.yaml
+metrics_config=src/test/resources/metrics_default.yaml
+model_api_enabled=true
diff --git a/frontend/server/src/test/resources/snapshots/snapshot4.cfg b/frontend/server/src/test/resources/snapshots/snapshot4.cfg
@@ -18,4 +18,5 @@ certificate_file=src/test/resources/certs.pem
 management_address=https\://127.0.0.1\:8444
 metrics_address=https\://127.0.0.1\:8445
 workflow_store=../archive/src/test/resources/workflows
-metrics_config=src/test/resources/metrics_default.yaml
+metrics_config=src/test/resources/metrics_default.yaml
+model_api_enabled=true
diff --git a/frontend/server/src/test/resources/snapshots/snapshot5.cfg b/frontend/server/src/test/resources/snapshots/snapshot5.cfg
@@ -18,4 +18,5 @@ certificate_file=src/test/resources/certs.pem
 management_address=https\://127.0.0.1\:8444
 metrics_address=https\://127.0.0.1\:8445
 workflow_store=../archive/src/test/resources/workflows
-metrics_config=src/test/resources/metrics_default.yaml
+metrics_config=src/test/resources/metrics_default.yaml
+model_api_enabled=true
diff --git a/frontend/server/src/test/resources/snapshots/snapshot6.cfg b/frontend/server/src/test/resources/snapshots/snapshot6.cfg
@@ -18,4 +18,5 @@ certificate_file=src/test/resources/certs.pem
 management_address=https\://127.0.0.1\:8444
 metrics_address=https\://127.0.0.1\:8445
 workflow_store=../archive/src/test/resources/workflows
-metrics_config=src/test/resources/metrics_default.yaml
+metrics_config=src/test/resources/metrics_default.yaml
+model_api_enabled=true
diff --git a/frontend/server/src/test/resources/snapshots/snapshot7.cfg b/frontend/server/src/test/resources/snapshots/snapshot7.cfg
@@ -18,4 +18,5 @@ certificate_file=src/test/resources/certs.pem
 management_address=https\://127.0.0.1\:8444
 metrics_address=https\://127.0.0.1\:8445
 workflow_store=../archive/src/test/resources/workflows
-metrics_config=src/test/resources/metrics_default.yaml
+metrics_config=src/test/resources/metrics_default.yaml
+model_api_enabled=true
diff --git a/frontend/server/src/test/resources/snapshots/snapshot8.cfg b/frontend/server/src/test/resources/snapshots/snapshot8.cfg
@@ -18,4 +18,5 @@ certificate_file=src/test/resources/certs.pem
 management_address=https\://127.0.0.1\:8444
 metrics_address=https\://127.0.0.1\:8445
 workflow_store=../archive/src/test/resources/workflows
-metrics_config=src/test/resources/metrics_default.yaml
+metrics_config=src/test/resources/metrics_default.yaml
+model_api_enabled=true
diff --git a/frontend/server/src/test/resources/snapshots/snapshot9.cfg b/frontend/server/src/test/resources/snapshots/snapshot9.cfg
@@ -18,4 +18,5 @@ certificate_file=src/test/resources/certs.pem
 management_address=https\://127.0.0.1\:8444
 metrics_address=https\://127.0.0.1\:8445
 workflow_store=../archive/src/test/resources/workflows
-metrics_config=src/test/resources/metrics_default.yaml
+metrics_config=src/test/resources/metrics_default.yaml
+model_api_enabled=true
-Original file line number
+Diff line change
@@ Expand Up / @@ -53,3 +53,4 @@ models={\ @@
       }\
     }
     metrics_config=src/test/resources/metrics_default.yaml
+    model_api_enabled=true