feat: add GPU support

kallebysantos · kallebysantos · commit 57d6ccd2cab0 · 2024-10-31T14:51:10.000Z
Signed-off-by: kallebysantos &lt;kalleby_santos@hotmail.com&gt;
diff --git a/Dockerfile b/Dockerfile
@@ -26,20 +26,53 @@ RUN objcopy --strip-debug \
   --add-gnu-debuglink=/root/edge-runtime.debug \
   /root/edge-runtime
 
-RUN ./scripts/install_onnx.sh $ONNXRUNTIME_VERSION $TARGETPLATFORM /root/onnxruntime
-RUN ./scripts/download_models.sh
 
-FROM debian:bookworm-slim
+# Application runtime without ONNX
+FROM debian:bookworm-slim as edge-runtime-base
 
 RUN apt-get update && apt-get install -y libssl-dev && rm -rf /var/lib/apt/lists/*
 RUN apt-get remove -y perl && apt-get autoremove -y
 
 COPY --from=builder /root/edge-runtime /usr/local/bin/edge-runtime
 COPY --from=builder /root/edge-runtime.debug /usr/local/bin/edge-runtime.debug
-COPY --from=builder /root/onnxruntime /usr/local/bin/onnxruntime
-COPY --from=builder /usr/src/edge-runtime/models /etc/sb_ai/models
 
 ENV ORT_DYLIB_PATH=/usr/local/bin/onnxruntime/lib/libonnxruntime.so
-ENV SB_AI_MODELS_DIR=/etc/sb_ai/models
+
+
+# ONNX Runtime provider
+# Application runtime with ONNX
+FROM builder as ort
+RUN ./scripts/install_onnx.sh $ONNXRUNTIME_VERSION $TARGETPLATFORM /root/onnxruntime
+
+
+# ONNX Runtime CUDA provider
+# Application runtime with ONNX CUDA
+FROM builder as ort-cuda
+RUN ./scripts/install_onnx.sh $ONNXRUNTIME_VERSION $TARGETPLATFORM /root/onnxruntime --gpu
+
+
+FROM builder as preload-models
+RUN ./scripts/download_models.sh
+
+
+# With CUDA
+FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 as edge-runtime-cuda
+
+COPY --from=edge-runtime-base /usr/local/bin/edge-runtime /usr/local/bin/edge-runtime
+COPY --from=builder /root/edge-runtime.debug /usr/local/bin/edge-runtime.debug
+COPY --from=ort-cuda /root/onnxruntime /usr/local/bin/onnxruntime
+COPY --from=preload-models /usr/src/edge-runtime/models /etc/sb_ai/models
+
+ENV ORT_DYLIB_PATH=/usr/local/bin/onnxruntime/lib/libonnxruntime.so
+ENV NVIDIA_VISIBLE_DEVICES=all
+ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
+
+ENTRYPOINT ["edge-runtime"]
+
+
+# Base
+FROM edge-runtime-base as edge-runtime
+COPY --from=ort /root/onnxruntime /usr/local/bin/onnxruntime
+COPY --from=preload-models /usr/src/edge-runtime/models /etc/sb_ai/models
 
 ENTRYPOINT ["edge-runtime"]
diff --git a/crates/sb_ai/session.rs b/crates/sb_ai/session.rs
@@ -3,11 +3,12 @@ use once_cell::sync::Lazy;
 use std::collections::HashMap;
 use std::sync::Mutex;
 use std::{path::PathBuf, sync::Arc};
+use tracing::debug;
 
 use anyhow::{anyhow, Error};
 use ort::{
-    CPUExecutionProvider, ExecutionProviderDispatch, GraphOptimizationLevel, Session,
-    SessionBuilder,
+    CPUExecutionProvider, CUDAExecutionProvider, ExecutionProvider, ExecutionProviderDispatch,
+    GraphOptimizationLevel, Session, SessionBuilder,
 };
 
 use crate::onnx::ensure_onnx_env_init;
@@ -49,14 +50,32 @@ fn cpu_execution_provider() -> Box<dyn Iterator<Item = ExecutionProviderDispatch
     )
 }
 
+fn cuda_execution_provider() -> Box<dyn Iterator<Item = ExecutionProviderDispatch>> {
+    let cuda = CUDAExecutionProvider::default();
+    let providers = match cuda.is_available() {
+        Ok(is_cuda_available) => {
+            debug!(cuda_support = is_cuda_available);
+            if is_cuda_available {
+                vec![cuda.build()]
+            } else {
+                vec![]
+            }
+        }
+
+        _ => vec![],
+    };
+
+    Box::new(providers.into_iter().chain(cpu_execution_provider()))
+}
+
 fn create_session(model_bytes: &[u8]) -> Result<Arc<Session>, Error> {
     let session = {
         if let Some(err) = ensure_onnx_env_init() {
             return Err(anyhow!("failed to create onnx environment: {err}"));
         }
 
         get_session_builder()?
-            .with_execution_providers(cpu_execution_provider())?
+            .with_execution_providers(cuda_execution_provider())?
             .commit_from_memory(model_bytes)?
     };
 
