From 0885e41aab0b267a8e4f526fc69f6a3cf576a044 Mon Sep 17 00:00:00 2001
From: Dmitry Meyer <me@undef.im>
Date: Fri, 16 May 2025 11:45:15 +0000
Subject: [PATCH] Replace conda with uv in dstackai/base images

Closes: https://github.com/dstackai/dstack/issues/2625
---
 docker/base/Dockerfile       | 30 +++++++++++++-----------------
 docker/base/devel.Dockerfile |  3 +--
 2 files changed, 14 insertions(+), 19 deletions(-)

diff --git a/docker/base/Dockerfile b/docker/base/Dockerfile
index 5687e56c3c..6acc57dd42 100644
--- a/docker/base/Dockerfile
+++ b/docker/base/Dockerfile
@@ -1,14 +1,20 @@
 FROM nvidia/cuda:12.1.0-base-ubuntu20.04
 
 ARG PYTHON
-ENV PYTHON=$PYTHON
+ARG _UV_HOME="/opt/uv"
+ENV UV_PYTHON="${PYTHON}"
+ENV UV_INSTALL_DIR="${_UV_HOME}/bin"
+ENV UV_PYTHON_INSTALL_DIR="${_UV_HOME}/python"
+ENV UV_PYTHON_BIN_DIR="${UV_PYTHON_INSTALL_DIR}/bin"
+ENV UV_MANAGED_PYTHON=1
 ENV LANG=C.UTF-8 LC_ALL=C.UTF-8
-ENV PIP_ROOT_USER_ACTION=ignore
 
-RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub &&  \
+ENV PATH="${UV_INSTALL_DIR}:${UV_PYTHON_BIN_DIR}:${PATH}"
+
+RUN export DEBIAN_FRONTEND=noninteractive && \
+    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub && \
     apt-get update --fix-missing && \
-    apt upgrade -y && \
-    export DEBIAN_FRONTEND=noninteractive && \
+    apt-get upgrade -y && \
     ln -fs /usr/share/zoneinfo/America/New_York /etc/localtime && \
     apt-get install -y tzdata && \
     dpkg-reconfigure --frontend noninteractive tzdata && \
@@ -16,15 +22,5 @@ RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/
     sed -i "s/.*PasswordAuthentication.*/PasswordAuthentication no/g" /etc/ssh/sshd_config && mkdir /run/sshd && \
     mkdir ~/.ssh && chmod 700 ~/.ssh && touch ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys && rm /etc/ssh/ssh_host_*
 
-RUN wget -O Miniforge3.sh "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh" && \
-    bash Miniforge3.sh -b -p "/opt/conda" && \
-    rm Miniforge3.sh && \
-    ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \
-    chmod +x /opt/conda/etc/profile.d/conda.sh && \
-    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \
-    /opt/conda/condabin/conda update conda --all -y && \
-    /opt/conda/condabin/conda create --name workflow python=${PYTHON} -y && \
-    /opt/conda/condabin/conda config --prepend channels "nvidia/label/cuda-12.1.0" && \
-    /opt/conda/condabin/conda config --set always_yes true && \
-    /opt/conda/condabin/conda clean --all && \
-    echo "conda activate workflow" >> ~/.bashrc
+RUN curl -LsSf https://astral.sh/uv/install.sh | INSTALLER_NO_MODIFY_PATH=1 sh && \
+    uv python install --preview --default
diff --git a/docker/base/devel.Dockerfile b/docker/base/devel.Dockerfile
index 39f7e350e9..eec0d6b6f0 100644
--- a/docker/base/devel.Dockerfile
+++ b/docker/base/devel.Dockerfile
@@ -3,5 +3,4 @@ ARG VERSION
 
 FROM dstackai/base:py$PYTHON-$VERSION-cuda-12.1
 
-RUN /opt/conda/condabin/conda install --name workflow cuda -y && \
-    /opt/conda/condabin/conda clean --all
+RUN DEBIAN_FRONTEND=noninteractive apt-get install -y cuda-12-1