[docker] Fixes duplicated onnx jar file issue (#2017)

deepjavalibrary · Jun 4, 2024 · ee540be · ee540be
1 parent 7f54a57
commit ee540be
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/serving/docker/lmi.Dockerfile b/serving/docker/lmi.Dockerfile
@@ -85,7 +85,7 @@ 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 && \
+    rm -f /usr/local/djl-serving-*/lib/onnxruntime-1.*.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} lmi" > /opt/djl/bin/telemetry && \

diff --git a/wlm/src/main/java/ai/djl/serving/wlm/LmiUtils.java b/wlm/src/main/java/ai/djl/serving/wlm/LmiUtils.java
@@ -172,6 +172,9 @@ private static Path convertOnnx(String modelId) throws IOException {
             return repoDir;
         }
 
+        Engine onnx = Engine.getEngine("OnnxRuntime");
+        boolean hasCuda = onnx.getGpuCount() > 0;
+
         String[] cmd = {
             "djl-convert",
             "--output-dir",
@@ -181,9 +184,9 @@ private static Path convertOnnx(String modelId) throws IOException {
             "-m",
             modelId,
             "--optimize",
-            CudaUtils.hasCuda() ? "O4" : "O2",
+            hasCuda ? "O4" : "O2",
             "--device",
-            CudaUtils.hasCuda() ? "cuda" : "cpu"
+            hasCuda ? "cuda" : "cpu"
         };
         boolean success = false;
         try {