Avoid load startup model on multiple GPUs

deepjavalibrary · Mar 11, 2022 · f96ebd6 · f96ebd6
1 parent 5c27dba
commit f96ebd6
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/serving/src/main/java/ai/djl/serving/ModelServer.java b/serving/src/main/java/ai/djl/serving/ModelServer.java
@@ -423,7 +423,7 @@ String mapModelUrl(Path path) {
             if (engine == null) {
                 return null;
             }
-            return modelName + "::" + engine + ":*=" + url;
+            return modelName + "::" + engine + ':' + configManager.getLoadOnDevices() + '=' + url;
         } catch (MalformedURLException e) {
             throw new AssertionError("Invalid path: " + path, e);
         } catch (IOException e) {

diff --git a/serving/src/main/java/ai/djl/serving/util/ConfigManager.java b/serving/src/main/java/ai/djl/serving/util/ConfigManager.java
@@ -66,6 +66,7 @@ public final class ConfigManager {
     private static final String MAX_REQUEST_SIZE = "max_request_size";
     private static final String MODEL_STORE = "model_store";
     private static final String MODEL_URL_PATTERN = "model_url_pattern";
+    private static final String LOAD_ON_DEVICES = "load_on_devices";
     private static final String PLUGIN_FOLDER = "plugin_folder";
 
     // Configuration which are not documented or enabled through environment variables
@@ -266,6 +267,15 @@ public String getLoadModels() {
         return prop.getProperty(LOAD_MODELS);
     }
 
+    /**
+     * Returns the devices the default model will be loaded on at startup.
+     *
+     * @return the devices the default model will be loaded on at startup
+     */
+    public String getLoadOnDevices() {
+        return prop.getProperty(LOAD_ON_DEVICES, "*");
+    }
+
     /**
      * Returns the CORS allowed origin setting.
      *