Repository: deepjavalibrary/djl
Pull request #2922 (2024-01-06 01:57:51)

[api] Adds ModelZooResolver interface
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/.github/workflows/native_jni_s3_pytorch.yml b/.github/workflows/native_jni_s3_pytorch.yml
index 6258413b1..82d580b37 100644
--- a/.github/workflows/native_jni_s3_pytorch.yml
+++ b/.github/workflows/native_jni_s3_pytorch.yml
@@ -56,7 +56,7 @@ jobs:
     steps:
       - uses: actions/checkout@v3
       - name: Set up JDK 17
-        uses: actions/setup-java@v4
+        uses: actions/setup-java@v3
         with:
           distribution: 'corretto'
           java-version: 17
@@ -269,14 +269,14 @@ jobs:
         run: |
           yum -y update
           yum -y groupinstall "Development Tools"
-          yum -y install patch git cmake3 python3-devel java-17-amazon-corretto
+          yum -y install patch git cmake3 python3-devel java-17-amazon-corretto-devel
           ln -sf /usr/bin/cmake3 /usr/bin/cmake
           pip3 install awscli --upgrade
       - uses: actions/checkout@v3
       - name: Release JNI prep
         run: |
           export JAVA_HOME=/usr/lib/jvm/java-17-amazon-corretto.aarch64
-          export PATH=$PATH:$JAVA_HOME
+          export PATH=$PATH:$JAVA_HOME/bin
           PYTORCH_VERSION=${{ github.event.inputs.pt_version }}
           export PYTORCH_VERSION=${PYTORCH_VERSION:-$(cat gradle.properties | awk -F '=' '/pytorch_version/ {print $2}')}
           echo $PYTORCH_VERSION
diff --git a/.github/workflows/native_s3_tflite.yml b/.github/workflows/native_s3_tflite.yml
index 890dfc707..7f5f92a8c 100644
--- a/.github/workflows/native_s3_tflite.yml
+++ b/.github/workflows/native_s3_tflite.yml
@@ -53,7 +53,7 @@ jobs:
         run: |
           yum -y update
           yum -y groupinstall "Development Tools"
-          yum -y install patch cmake3 unzip which java-17-amazon-corretto
+          yum -y install patch cmake3 unzip which java-17-amazon-corretto-devel
           ln -sf /usr/bin/cmake3 /usr/bin/cmake
           pip3 install awscli --upgrade
           pip3 install numpy --upgrade
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/SeqBatchScheduler.java b/api/src/main/java/ai/djl/modality/nlp/generate/SeqBatchScheduler.java
index 78f40c5b2..07c63428f 100644
--- a/api/src/main/java/ai/djl/modality/nlp/generate/SeqBatchScheduler.java
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/SeqBatchScheduler.java
@@ -35,15 +35,13 @@ import java.util.concurrent.ConcurrentHashMap;
  * policy is setting several thresholds.
  */
 public abstract class SeqBatchScheduler {
+
     private static final Logger logger = LoggerFactory.getLogger(SeqBatchScheduler.class);
 
     Predictor<NDList, CausalLMOutput> predictor;
     SeqBatcher seqBatcher;
-
     NDManager manager;
-
     SearchConfig config;
-
     Map<Long, NDArray> results;
 
     /**
@@ -101,7 +99,7 @@ public abstract class SeqBatchScheduler {
      * @return the output token ids
      * @throws TranslateException if forward fails
      */
-    abstract NDArray inferenceCall() throws TranslateException;
+    protected abstract NDArray inferenceCall() throws TranslateException;
 
     /**
      * Adds new batch.
diff --git a/api/src/main/java/ai/djl/repository/zoo/ModelZoo.java b/api/src/main/java/ai/djl/repository/zoo/ModelZoo.java
index 50b219be5..e903a1677 100644
--- a/api/src/main/java/ai/djl/repository/zoo/ModelZoo.java
+++ b/api/src/main/java/ai/djl/repository/zoo/ModelZoo.java
@@ -32,6 +32,7 @@ import java.util.concurrent.ConcurrentHashMap;
 public abstract class ModelZoo {
 
     private static final Map<String, ModelZoo> MODEL_ZOO_MAP = new ConcurrentHashMap<>();
+    private static ModelZooResolver resolver;
 
     private Map<String, ModelLoader> modelLoaders = new ConcurrentHashMap<>();
 
@@ -86,6 +87,15 @@ public abstract class ModelZoo {
         modelLoaders.put(loader.getArtifactId(), loader);
     }
 
+    /**
+     * Sets the {@code ModelZooResolver}.
+     *
+     * @param resolver the {@code ModelZooResolver}
+     */
+    public static void setModelZooResolver(ModelZooResolver resolver) {
+        ModelZoo.resolver = resolver;
+    }
+
     /**
      * Refreshes model zoo.
      *
@@ -112,7 +122,14 @@ public abstract class ModelZoo {
      * @return the {@code ModelZoo} with the {@code groupId}
      */
     public static ModelZoo getModelZoo(String groupId) {
-        return MODEL_ZOO_MAP.get(groupId);
+        ModelZoo zoo = MODEL_ZOO_MAP.get(groupId);
+        if (zoo == null && resolver != null) {
+            zoo = resolver.resolve(groupId);
+            if (zoo != null) {
+                MODEL_ZOO_MAP.putIfAbsent(groupId, zoo);
+            }
+        }
+        return zoo;
     }
 
     /**
diff --git a/api/src/main/java/ai/djl/repository/zoo/ModelZooResolver.java b/api/src/main/java/ai/djl/repository/zoo/ModelZooResolver.java
new file mode 100644
index 000000000..897e122f1
--- /dev/null
+++ b/api/src/main/java/ai/djl/repository/zoo/ModelZooResolver.java
@@ -0,0 +1,25 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.repository.zoo;
+
+/** An interface that resolves external ModelZoo. */
+public interface ModelZooResolver {
+
+    /**
+     * Returns {@link ModelZoo} based on model zoo group ID.
+     *
+     * @param groupId the model zoo group ID.
+     * @return the resolved {@code ModelZoo}
+     */
+    ModelZoo resolve(String groupId);
+}
diff --git a/engines/pytorch/pytorch-native/CMakeLists.txt b/engines/pytorch/pytorch-native/CMakeLists.txt
index 4eddbfc35..c53d71dc9 100644
--- a/engines/pytorch/pytorch-native/CMakeLists.txt
+++ b/engines/pytorch/pytorch-native/CMakeLists.txt
@@ -60,11 +60,12 @@ if(USE_CUDA)
 endif()
 
 add_library(djl_torch SHARED ${SOURCE_FILES})
+set_property(TARGET djl_torch PROPERTY CXX_STANDARD 17)
+
 # build host
 if(NOT BUILD_ANDROID)
     target_link_libraries(djl_torch "${TORCH_LIBRARIES}")
     target_include_directories(djl_torch PUBLIC build/include ${JNI_INCLUDE_DIRS} ${UTILS_INCLUDE_DIR})
-    set_property(TARGET djl_torch PROPERTY CXX_STANDARD 17)
     # We have to kill the default rpath and use current dir
     set(CMAKE_SKIP_RPATH TRUE)
     if(${CMAKE_SYSTEM_NAME} MATCHES "Linux")
diff --git a/engines/pytorch/pytorch-native/build.sh b/engines/pytorch/pytorch-native/build.sh
index b3cc2f117..ae0456bec 100755
--- a/engines/pytorch/pytorch-native/build.sh
+++ b/engines/pytorch/pytorch-native/build.sh
@@ -29,16 +29,16 @@ if [[ ! -d "libtorch" ]]; then
     fi
 
     if [[ $ARCH == 'aarch64' ]]; then
-      curl -s https://djl-ai.s3.amazonaws.com/publish/pytorch/${VERSION}/libtorch${AARCH64_CXX11ABI}-shared-with-deps-${VERSION}-aarch64.zip | jar xv
+      curl -s https://djl-ai.s3.amazonaws.com/publish/pytorch/${VERSION}/libtorch${AARCH64_CXX11ABI}-shared-with-deps-${VERSION}-aarch64.zip | jar xv > /dev/null
     else
-      curl -s https://download.pytorch.org/libtorch/${FLAVOR}/libtorch${CXX11ABI}-shared-with-deps-${VERSION}%2B${FLAVOR}.zip | jar xv
+      curl -s https://download.pytorch.org/libtorch/${FLAVOR}/libtorch${CXX11ABI}-shared-with-deps-${VERSION}%2B${FLAVOR}.zip | jar xv > /dev/null
     fi
 
   elif [[ $PLATFORM == 'darwin' ]]; then
     if [[ $ARCH == 'aarch64' ]]; then
-      curl -s https://djl-ai.s3.amazonaws.com/publish/pytorch/${VERSION}/libtorch-macos-${VERSION}-aarch64.zip | jar xv
+      curl -s https://djl-ai.s3.amazonaws.com/publish/pytorch/${VERSION}/libtorch-macos-${VERSION}-aarch64.zip | jar xv > /dev/null
     else
-      curl -s https://download.pytorch.org/libtorch/cpu/libtorch-macos-${VERSION}.zip | jar xv
+      curl -s https://download.pytorch.org/libtorch/cpu/libtorch-macos-${VERSION}.zip | jar xv > /dev/null
     fi
   else
     echo "$PLATFORM is not supported."
diff --git a/engines/pytorch/pytorch-native/build_android.sh b/engines/pytorch/pytorch-native/build_android.sh
index b37dd96a8..72050b20a 100755
--- a/engines/pytorch/pytorch-native/build_android.sh
+++ b/engines/pytorch/pytorch-native/build_android.sh
@@ -20,7 +20,7 @@ if [[ ! -d libtorch_android/"$FLAVOR" ]]; then
     mkdir -p libtorch_android/"$FLAVOR"
     cd libtorch_android/"$FLAVOR"
     echo "Downloading https://publish.djl.ai/pytorch/$VERSION/android_native/${FLAVOR}_native.zip"
-    curl -s "https://publish.djl.ai/pytorch/$VERSION/android_native/${FLAVOR}_native.zip" | jar xv
+    curl -s "https://publish.djl.ai/pytorch/$VERSION/android_native/${FLAVOR}_native.zip" | jar xv > /dev/null
     mv install/include include
     cd -
   fi


Test patch:
diff --git a/api/src/test/java/ai/djl/repository/ZooTest.java b/api/src/test/java/ai/djl/repository/ZooTest.java
index 2b44f9671..29fc10391 100644
--- a/api/src/test/java/ai/djl/repository/ZooTest.java
+++ b/api/src/test/java/ai/djl/repository/ZooTest.java
@@ -17,6 +17,7 @@ import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ModelNotFoundException;
+import ai.djl.repository.zoo.ModelZoo;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -48,4 +49,11 @@ public class ZooTest {
         Criteria<?, ?> criteria = Criteria.builder().build();
         criteria.loadModel();
     }
+
+    @Test
+    public void testModelZooResolver() {
+        ModelZoo.setModelZooResolver(groupId -> null);
+        ModelZoo zoo = ModelZoo.getModelZoo("unknown");
+        Assert.assertNull(zoo);
+    }
 }