Load external dependencies for workflows (deepjavalibrary#556)

Lokiiiiii · Mar 17, 2023 · 0a8d55b · 0a8d55b
1 parent fe08fd1
commit 0a8d55b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/serving/src/main/java/ai/djl/serving/ModelServer.java b/serving/src/main/java/ai/djl/serving/ModelServer.java
@@ -90,6 +90,7 @@ public class ModelServer {
     private ConfigManager configManager;
 
     private FolderScanPluginManager pluginManager;
+    private DependencyManager dependencyManager;
 
     /**
      * Creates a new {@code ModelServer} instance.
@@ -100,6 +101,7 @@ public ModelServer(ConfigManager configManager) {
         this.configManager = configManager;
         this.pluginManager = new FolderScanPluginManager(configManager);
         serverGroups = new ServerGroups(configManager);
+        dependencyManager = DependencyManager.getInstance();
     }
 
     /**
@@ -384,7 +386,7 @@ private void initModelStore() throws IOException {
                     continue;
                 }
             }
-            DependencyManager.getInstance().installEngine(engineName);
+            dependencyManager.installEngine(engineName);
             Engine engine = Engine.getEngine(engineName);
             String[] devices = parseDevices(deviceMapping, engine, pair.getValue());
 
@@ -461,7 +463,7 @@ private void initWorkflows() throws IOException, URISyntaxException, BadWorkflow
                 if (tokens.length > 1) {
                     Pair<String, Path> pair = ModelInfo.downloadModel(workflowUrlString);
                     String engineName = ModelInfo.inferEngine(pair.getValue(), pair.getKey());
-                    DependencyManager.getInstance().installEngine(engineName);
+                    dependencyManager.installEngine(engineName);
                     Engine engine = Engine.getEngine(engineName);
                     devices = parseDevices(tokens[1], engine, pair.getValue());
                 }

diff --git a/serving/src/main/java/ai/djl/serving/workflow/WorkflowDefinition.java b/serving/src/main/java/ai/djl/serving/workflow/WorkflowDefinition.java
@@ -14,6 +14,7 @@
 
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
+import ai.djl.serving.util.MutableClassLoader;
 import ai.djl.serving.wlm.ModelInfo;
 import ai.djl.serving.wlm.util.WlmConfigManager;
 import ai.djl.serving.workflow.WorkflowExpression.Item;
@@ -166,7 +167,8 @@ public Workflow toWorkflow() throws BadWorkflowException {
             for (Entry<String, String> f : funcs.entrySet()) {
                 try {
                     Class<? extends WorkflowFunction> clazz =
-                            Class.forName(f.getValue()).asSubclass(WorkflowFunction.class);
+                            Class.forName(f.getValue(), true, MutableClassLoader.getInstance())
+                                    .asSubclass(WorkflowFunction.class);
                     loadedFunctions.put(f.getKey(), clazz.getConstructor().newInstance());
                 } catch (Exception e) {
                     throw new BadWorkflowException("Could not load function " + f.getKey(), e);