From 7ac496da39b7d516139c5a1b8c05231e1dc1153e Mon Sep 17 00:00:00 2001
From: Frank Liu <frankfliu2000@gmail.com>
Date: Fri, 19 Apr 2024 14:51:18 -0700
Subject: [PATCH] [docker] Install onnxruntime in docker

---
 serving/build.gradle                                        | 1 +
 serving/docker/Dockerfile                                   | 5 +++--
 serving/docker/pytorch-gpu.Dockerfile                       | 6 ++++--
 .../main/java/ai/djl/serving/plugins/DependencyManager.java | 1 -
 4 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/serving/build.gradle b/serving/build.gradle
index 3b9ff487c5..993c98aca4 100644
--- 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
index 0fc95b2385..0d0c14f108 100644
--- 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
index ca89339906..3909d9cc5c 100644
--- 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
index 8c5828fad1..a2f6af418e 100644
--- 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);