[docker] Install onnxruntime in docker

deepjavalibrary · Apr 19, 2024 · 7ac496d · 7ac496d
1 parent 812c164
commit 7ac496d
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 5 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/Dockerfile b/serving/docker/Dockerfile
@@ -60,14 +60,15 @@ LABEL djl-version=$djl_version
 
 FROM base AS cpu-full
 
-ARG torch_version=2.1.1
+ARG torch_version=2.1.2
+ARG onnx_version=1.17.1
 
 COPY scripts scripts/
 RUN scripts/install_python.sh && \
     scripts/install_djl_serving.sh $djl_version $torch_version && \
     echo "${djl_version} cpufull" > /opt/djl/bin/telemetry && \
     djl-serving -i ai.djl.pytorch:pytorch-native-cpu:$torch_version:linux-x86_64 && \
-    djl-serving -i ai.djl.tensorflow:tensorflow-native-cpu:2.10.1:linux-x86_64 && \
+    djl-serving -i com.microsoft.onnxruntime:onnxruntime:$onnx_version && \
     scripts/patch_oss_dlc.sh python && \
     rm -rf /opt/djl/logs && \
     chown -R djl:djl /opt/djl && \

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,7 @@ 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} && \
+    mkdir -p /opt/djl/deps && curl -o /opt/djl/deps/onnxruntime_gpu-$onnx_version.jar https://djl-ai/publish/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
@@ -77,7 +77,6 @@ public void installEngine(String engineName) throws IOException {
                 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);