[serving] Adds deps folder to classpath in MutableClassLoader constru…

…ctor
deepjavalibrary · Apr 9, 2023 · ae672a5 · ae672a5
1 parent 0d29cb4
commit ae672a5
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 34 deletions.
diff --git a/serving/src/main/java/ai/djl/serving/plugins/DependencyManager.java b/serving/src/main/java/ai/djl/serving/plugins/DependencyManager.java
@@ -30,49 +30,25 @@
 
 import java.io.IOException;
 import java.io.InputStream;
-import java.net.MalformedURLException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ServiceLoader;
-import java.util.stream.Stream;
 
 import javax.xml.parsers.DocumentBuilder;
 import javax.xml.parsers.DocumentBuilderFactory;
 import javax.xml.parsers.ParserConfigurationException;
 
 /** {@code DependencyManager} is responsible to manage extra maven dependencies. */
-public class DependencyManager {
+public final class DependencyManager {
 
     private static final Logger logger = LoggerFactory.getLogger(DependencyManager.class);
 
     private static final DependencyManager INSTANCE = new DependencyManager();
     private static final String OSS_URL =
             "https://oss.sonatype.org/service/local/repositories/snapshots/content/";
 
-    private Path depDir;
-
-    DependencyManager() {
-        String serverHome = ConfigManager.getModelServerHome();
-        depDir = Paths.get(serverHome, "deps");
-        if (Files.isDirectory(depDir)) {
-            MutableClassLoader mc = MutableClassLoader.getInstance();
-            try (Stream<Path> stream = Files.list(depDir)) {
-                stream.forEach(
-                        p -> {
-                            if (p.toString().endsWith(".jar")) {
-                                try {
-                                    mc.addURL(p.toUri().toURL());
-                                } catch (MalformedURLException e) {
-                                    logger.warn("Invalid file system path: " + p, e);
-                                }
-                            }
-                        });
-            } catch (IOException e) {
-                logger.warn("Failed to load dependencies from deps folder.", e);
-            }
-        }
-    }
+    private DependencyManager() {}
 
     /**
      * Returns the singleton instance of {@code DependencyManager}.
@@ -141,6 +117,8 @@ public synchronized void installDependency(String dependency) throws IOException
         if (tokens.length < 3) {
             throw new IllegalArgumentException("Invalid dependency: " + dependency);
         }
+        String serverHome = ConfigManager.getModelServerHome();
+        Path depDir = Paths.get(serverHome, "deps");
         Files.createDirectories(depDir);
 
         logger.info("Loading dependency: {}", dependency);

diff --git a/serving/src/main/java/ai/djl/serving/util/MutableClassLoader.java b/serving/src/main/java/ai/djl/serving/util/MutableClassLoader.java
@@ -12,14 +12,25 @@
  */
 package ai.djl.serving.util;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.net.MalformedURLException;
 import java.net.URL;
 import java.net.URLClassLoader;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
+import java.util.stream.Stream;
 
 /** A {@code URLClassLoader} that can add new class at runtime. */
 public class MutableClassLoader extends URLClassLoader {
 
+    private static final Logger logger = LoggerFactory.getLogger(MutableClassLoader.class);
+
     private static final MutableClassLoader INSTANCE =
             AccessController.doPrivileged(
                     (PrivilegedAction<MutableClassLoader>) MutableClassLoader::new);
@@ -40,6 +51,24 @@ public class MutableClassLoader extends URLClassLoader {
      */
     public MutableClassLoader() {
         super(new URL[0]);
+        String serverHome = ConfigManager.getModelServerHome();
+        Path depDir = Paths.get(serverHome, "deps");
+        if (Files.isDirectory(depDir)) {
+            try (Stream<Path> stream = Files.list(depDir)) {
+                stream.forEach(
+                        p -> {
+                            if (p.toString().endsWith(".jar")) {
+                                try {
+                                    addURL(p.toUri().toURL());
+                                } catch (MalformedURLException e) {
+                                    logger.warn("Invalid file system path: " + p, e);
+                                }
+                            }
+                        });
+            } catch (IOException e) {
+                logger.warn("Failed to load dependencies from deps folder.", e);
+            }
+        }
     }
 
     /**
@@ -53,7 +82,7 @@ public static MutableClassLoader getInstance() {
 
     /** {@inheritDoc} */
     @Override
-    public void addURL(URL url) {
+    public final void addURL(URL url) {
         super.addURL(url);
     }
 }
diff --git a/serving/src/main/java/ai/djl/serving/workflow/WorkflowDefinition.java b/serving/src/main/java/ai/djl/serving/workflow/WorkflowDefinition.java
@@ -168,14 +168,13 @@ public Workflow toWorkflow() throws BadWorkflowException {
         Map<String, WorkflowFunction> loadedFunctions = new ConcurrentHashMap<>();
         if (funcs != null) {
             for (Entry<String, String> f : funcs.entrySet()) {
-                try {
-                    Class<? extends WorkflowFunction> clazz =
-                            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);
+                ClassLoader cl = MutableClassLoader.getInstance();
+                WorkflowFunction func =
+                        ClassLoaderUtils.initClass(cl, WorkflowFunction.class, f.getValue());
+                if (func == null) {
+                    throw new BadWorkflowException("Could not load function " + f.getKey());
                 }
+                loadedFunctions.put(f.getKey(), func);
             }
         }