From 0a8d55beb037b250d002ca87ac9d478557095946 Mon Sep 17 00:00:00 2001
From: Xin Yang <105740670+xyang16@users.noreply.github.com>
Date: Fri, 17 Mar 2023 09:46:29 -0700
Subject: [PATCH] Load external dependencies for workflows (#556)

---
 serving/src/main/java/ai/djl/serving/ModelServer.java       | 6 ++++--
 .../java/ai/djl/serving/workflow/WorkflowDefinition.java    | 4 +++-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/serving/src/main/java/ai/djl/serving/ModelServer.java b/serving/src/main/java/ai/djl/serving/ModelServer.java
index 09d614111fa..01e22ecdf37 100644
--- 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
index 43333acd281..71d06eb6040 100644
--- 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);