[docker] Install onnxruntime in docker (#1790)

deepjavalibrary · Apr 23, 2024 · a585f9e · a585f9e
1 parent 8f8fea0
commit a585f9e
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 13 deletions.
diff --git a/serving/build.gradle b/serving/build.gradle
@@ -22,6 +22,7 @@ dependencies {
     runtimeOnly "ai.djl:model-zoo"
     runtimeOnly "ai.djl.tensorflow:tensorflow-model-zoo"
     runtimeOnly "ai.djl.pytorch:pytorch-model-zoo"
+    runtimeOnly "ai.djl.onnxruntime:onnxruntime-engine"
     runtimeOnly "ai.djl.huggingface:tokenizers"
     runtimeOnly "ai.djl.tensorrt:tensorrt"
     if (JavaVersion.current() >= JavaVersion.VERSION_15) {

diff --git a/serving/docker/deepspeed.Dockerfile b/serving/docker/deepspeed.Dockerfile
@@ -17,6 +17,7 @@ ARG djl_version=0.28.0~SNAPSHOT
 ARG python_version=3.10
 ARG torch_version=2.1.2
 ARG torch_vision_version=0.16.2
+ARG onnx_version=1.17.1
 ARG pydantic_version=2.6.1
 # HF Deps
 ARG protobuf_version=3.20.3
@@ -51,6 +52,11 @@ ENV MODEL_SERVER_HOME=/opt/djl
 ENV MODEL_LOADING_TIMEOUT=1200
 ENV PREDICT_TIMEOUT=240
 ENV DJL_CACHE_DIR=/tmp/.djl.ai
+ENV PYTORCH_LIBRARY_PATH=/usr/local/lib/python3.10/dist-packages/torch/lib
+ENV PYTORCH_PRECXX11=true
+ENV PYTORCH_VERSION=${torch_version}
+ENV PYTORCH_FLAVOR=cu121-precxx11
+
 ENV HF_HOME=/tmp/.cache/huggingface
 ENV PYTORCH_KERNEL_CACHE_PATH=/tmp/.cache
 ENV BITSANDBYTES_NOWELCOME=1
@@ -78,6 +84,8 @@ RUN mv *.deb djl-serving_all.deb || true
 
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -yq libaio-dev libopenmpi-dev g++ && \
     scripts/install_djl_serving.sh $djl_version && \
+    rm -f /usr/local/djl-serving-*/lib/onnxruntime-$onnx_version.jar && \
+    curl -o $(ls -d /usr/local/djl-serving-*/)lib/onnxruntime_gpu-$onnx_version.jar https://publish.djl.ai/onnxruntime/$onnx_version/onnxruntime_gpu-$onnx_version.jar && \
     mkdir -p /opt/djl/bin && cp scripts/telemetry.sh /opt/djl/bin && \
     echo "${djl_version} deepspeed" > /opt/djl/bin/telemetry && \
     scripts/install_python.sh ${python_version} && \

diff --git a/serving/docker/pytorch-gpu.Dockerfile b/serving/docker/pytorch-gpu.Dockerfile
@@ -15,8 +15,9 @@ FROM nvidia/cuda:$version as base
 
 ARG djl_version=0.28.0~SNAPSHOT
 ARG cuda_version=cu121
-ARG torch_version=2.1.1
-ARG torch_vision_version=0.16.1
+ARG torch_version=2.1.2
+ARG torch_vision_version=0.16.2
+ARG onnx_version=1.17.1
 ARG python_version=3.10
 
 RUN mkdir -p /opt/djl/conf && \
@@ -51,6 +52,8 @@ RUN chmod +x /usr/local/bin/dockerd-entrypoint.sh && \
     mkdir -p /opt/djl/bin && cp scripts/telemetry.sh /opt/djl/bin && \
     echo "${djl_version} pytorchgpu" > /opt/djl/bin/telemetry && \
     scripts/install_djl_serving.sh $djl_version ${torch_version} && \
+    rm -f /usr/local/djl-serving-*/lib/onnxruntime-$onnx_version.jar && \
+    curl -o $(ls -d /usr/local/djl-serving-*/)lib/onnxruntime_gpu-$onnx_version.jar https://publish.djl.ai/onnxruntime/$onnx_version/onnxruntime_gpu-$onnx_version.jar && \
     scripts/install_python.sh ${python_version} && \
     scripts/install_s5cmd.sh x64 && \
     pip3 install numpy && pip3 install torch==${torch_version} torchvision==${torch_vision_version} --extra-index-url https://download.pytorch.org/whl/cu121 && \

diff --git a/serving/src/main/java/ai/djl/serving/plugins/DependencyManager.java b/serving/src/main/java/ai/djl/serving/plugins/DependencyManager.java
@@ -19,7 +19,6 @@
 import ai.djl.serving.util.ConfigManager;
 import ai.djl.serving.util.MutableClassLoader;
 import ai.djl.util.Utils;
-import ai.djl.util.cuda.CudaUtils;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -76,15 +75,6 @@ public void installEngine(String engineName) throws IOException {
                 installDependency("ai.djl.mxnet:mxnet-engine:" + djlVersion);
                 installDependency("ai.djl.mxnet:mxnet-model-zoo:" + djlVersion);
                 break;
-            case "OnnxRuntime":
-                installDependency("ai.djl.onnxruntime:onnxruntime-engine:" + djlVersion);
-                String ortVersion = getOrtVersion(djlVersion);
-                if (CudaUtils.hasCuda()) {
-                    installDependency("com.microsoft.onnxruntime:onnxruntime_gpu:" + ortVersion);
-                } else {
-                    installDependency("com.microsoft.onnxruntime:onnxruntime:" + ortVersion);
-                }
-                break;
             case "PaddlePaddle":
                 installDependency("ai.djl.paddlepaddle:paddlepaddle-engine:" + djlVersion);
                 break;

diff --git a/serving/src/test/java/ai/djl/serving/plugins/DependencyManagerTest.java b/serving/src/test/java/ai/djl/serving/plugins/DependencyManagerTest.java
@@ -29,7 +29,6 @@ public void testInstallDependency() throws IOException {
         System.setProperty("MODEL_SERVER_HOME", "build/home");
         try {
             DependencyManager dm = DependencyManager.getInstance();
-            dm.installEngine("OnnxRuntime");
             dm.installEngine("XGBoost");
 
             dm.installDependency("ai.djl.pytorch:pytorch-jni:2.1.1-0.27.0");