Add cude docker

label_studio_ml/examples/yolo/Dockerfile

@@ -1,48 +1,52 @@
-# syntax=docker/dockerfile:1
-ARG PYTHON_VERSION=3.12
-
-FROM python:${PYTHON_VERSION}-slim AS python-base
+FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime
+ARG DEBIAN_FRONTEND=noninteractive
 ARG TEST_ENV
 
 WORKDIR /app
 
+RUN conda update conda -y
+
+RUN --mount=type=cache,target="/var/cache/apt",sharing=locked \
+    --mount=type=cache,target="/var/lib/apt/lists",sharing=locked \
+    apt-get -y update \
+    && apt-get install -y git \
+    && apt-get install -y wget \
+    && apt-get install -y g++ freeglut3-dev build-essential libx11-dev \
+    libxmu-dev libxi-dev libglu1-mesa libglu1-mesa-dev libfreeimage-dev \
+    && apt-get -y install ffmpeg libsm6 libxext6 libffi-dev python3-dev python3-pip gcc
+
 ENV PYTHONUNBUFFERED=1 \
     PYTHONDONTWRITEBYTECODE=1 \
-    PORT=${PORT:-9090} \
     PIP_CACHE_DIR=/.cache \
-    WORKERS=1 \
-    THREADS=8
+    PORT=9090 \
+    WORKERS=2 \
+    THREADS=4 \
+    CUDA_HOME=/usr/local/cuda \
 
-# Update the base OS
-RUN --mount=type=cache,target="/var/cache/apt",sharing=locked \
-    --mount=type=cache,target="/var/lib/apt/lists",sharing=locked \
-    set -eux; \
-    apt-get update; \
-    apt-get upgrade -y; \
-    apt install --no-install-recommends -y  \
-        git; \
-    apt-get autoremove -y
+RUN conda install -c "nvidia/label/cuda-12.1.1" cuda -y
+ENV CUDA_HOME=/opt/conda \
+    TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6+PTX;8.9;9.0"
 
 # install base requirements
 COPY requirements-base.txt .
 RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
     pip install -r requirements-base.txt
 
-# install custom requirements
+# install model requirements
 COPY requirements.txt .
 RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
-    pip install -r requirements.txt
+    pip3 install -r requirements.txt
 
 # install test requirements if needed
 COPY requirements-test.txt .
 # build only when TEST_ENV="true"
 RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
     if [ "$TEST_ENV" = "true" ]; then \
-      pip install -r requirements-test.txt; \
+      pip3 install -r requirements-test.txt; \
     fi
 
-COPY . .
+WORKDIR /app
 
-EXPOSE 9090
+COPY . ./
 
-CMD gunicorn --preload --bind :$PORT --workers $WORKERS --threads $THREADS --timeout 0 _wsgi:app
+CMD ["/app/start.sh"]

label_studio_ml/examples/yolo/start.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+# Execute the gunicorn command
+exec gunicorn --bind :${PORT:-9090} --workers ${WORKERS:-1} --threads ${THREADS:-4} --timeout 0 _wsgi:app