[serving] Make yaml support optional

serving/build.gradle

@@ -12,7 +12,6 @@ dependencies {
     api "io.netty:netty-transport-native-kqueue:${netty_version}:osx-aarch_64"
     api "io.netty:netty-transport-native-kqueue:${netty_version}:osx-x86_64"
 
-    implementation "org.yaml:snakeyaml:${snakeyaml_version}"
     //noinspection GradlePackageUpdate
     implementation "commons-cli:commons-cli:${commons_cli_version}"
     implementation "org.apache.logging.log4j:log4j-slf4j-impl:${log4j_slf4j_version}"
@@ -30,6 +29,7 @@ dependencies {
     }
     runtimeOnly project(":engines:python")
 
+    testRuntimeOnly "org.yaml:snakeyaml:${snakeyaml_version}"
     testImplementation("org.testng:testng:${testng_version}") {
         exclude group: "junit", module: "junit"
     }

serving/src/main/java/ai/djl/serving/workflow/WorkflowDefinition.java

@@ -18,6 +18,7 @@
 import ai.djl.serving.wlm.util.WlmConfigManager;
 import ai.djl.serving.workflow.WorkflowExpression.Item;
 import ai.djl.serving.workflow.function.WorkflowFunction;
+import ai.djl.util.ClassLoaderUtils;
 import ai.djl.util.JsonUtils;
 
 import com.google.gson.Gson;
@@ -28,12 +29,12 @@
 import com.google.gson.JsonParseException;
 import com.google.gson.annotations.SerializedName;
 
-import org.yaml.snakeyaml.Yaml;
-
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.io.Reader;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Method;
 import java.lang.reflect.Type;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
@@ -69,7 +70,6 @@ public class WorkflowDefinition {
     int maxBatchDelay;
     int batchSize;
 
-    private static final Yaml YAML = new Yaml();
     public static final Gson GSON =
             JsonUtils.builder()
                     .registerTypeAdapter(ModelInfo.class, new ModelDefinitionDeserializer())
@@ -109,9 +109,18 @@ public static WorkflowDefinition parse(URI uri, InputStream input) {
     public static WorkflowDefinition parse(URI uri, Reader input) {
         String fileName = Objects.requireNonNull(uri.toString());
         if (fileName.endsWith(".yml") || fileName.endsWith(".yaml")) {
-            Object yaml = YAML.load(input);
-            String asJson = GSON.toJson(yaml);
-            return GSON.fromJson(asJson, WorkflowDefinition.class);
+            try {
+                ClassLoader cl = ClassLoaderUtils.getContextClassLoader();
+                Class<?> clazz = Class.forName("org.yaml.snakeyaml.Yaml", true, cl);
+                Constructor<?> constructor = clazz.getConstructor();
+                Method method = clazz.getMethod("load", Reader.class);
+                Object obj = constructor.newInstance();
+                Object yaml = method.invoke(obj, input);
+                String asJson = GSON.toJson(yaml);
+                return GSON.fromJson(asJson, WorkflowDefinition.class);
+            } catch (ReflectiveOperationException e) {
+                throw new IllegalArgumentException("Yaml parsing is not supported.", e);
+            }
         } else if (fileName.endsWith(".json")) {
             return GSON.fromJson(input, WorkflowDefinition.class);
         } else {

serving/src/test/java/ai/djl/serving/WorkflowTest.java

@@ -32,7 +32,6 @@
 import java.net.URL;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.concurrent.ExecutionException;
 
 public class WorkflowTest {
 
@@ -52,29 +51,25 @@ public void beforeAll() throws IOException {
     }
 
     @Test
-    public void testJson()
-            throws IOException, BadWorkflowException, ExecutionException, InterruptedException {
+    public void testJson() throws IOException, BadWorkflowException {
         Path workflowFile = Paths.get("src/test/resources/workflows/basic.json");
         runWorkflow(workflowFile, zeroInput);
     }
 
     @Test
-    public void testYaml()
-            throws IOException, BadWorkflowException, ExecutionException, InterruptedException {
+    public void testYaml() throws IOException, BadWorkflowException {
         Path workflowFile = Paths.get("src/test/resources/workflows/basic.yaml");
         runWorkflow(workflowFile, zeroInput);
     }
 
     @Test
-    public void testCriteria()
-            throws IOException, BadWorkflowException, ExecutionException, InterruptedException {
+    public void testCriteria() throws IOException, BadWorkflowException {
         Path workflowFile = Paths.get("src/test/resources/workflows/criteria.json");
         runWorkflow(workflowFile, zeroInput);
     }
 
     @Test
-    public void testFunctions()
-            throws IOException, BadWorkflowException, ExecutionException, InterruptedException {
+    public void testFunctions() throws IOException, BadWorkflowException {
         Path workflowFile = Paths.get("src/test/resources/workflows/functions.json");
         runWorkflow(workflowFile, zeroInput);
     }