From 4a366a23865d4888e90b8e98d9b1673ebe7e198d Mon Sep 17 00:00:00 2001 From: Frank Liu Date: Fri, 4 Mar 2022 19:38:05 -0800 Subject: [PATCH] [serving] Make startup model default devices configurable --- serving/src/main/java/ai/djl/serving/ModelServer.java | 2 +- .../main/java/ai/djl/serving/util/ConfigManager.java | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/serving/src/main/java/ai/djl/serving/ModelServer.java b/serving/src/main/java/ai/djl/serving/ModelServer.java index 5a2f556a4..a3ca9b3e9 100644 --- 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 index 91887f99d..8603b7ced 100644 --- 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. *