revert to single stage image

Signed-off-by: Peter St. John <[email protected]>

Author: pstjohn

bionemo-recipes/recipes/esm2_native_te/Dockerfile.cuda

@@ -1,18 +1,16 @@
 # An example, minimal Dockerfile to install dependencies in a fresh python environment with CUDA support.
 
-FROM nvcr.io/nvidia/cuda:13.0.2-cudnn-devel-ubuntu24.04 AS build
-COPY --from=ghcr.io/astral-sh/uv:0.6.13 /uv /usr/local/bin/uv
+FROM nvcr.io/nvidia/cuda:13.0.2-cudnn-devel-ubuntu24.04
 
-ENV UV_PYTHON_INSTALL_DIR=/python
-ENV UV_PYTHON_PREFERENCE=only-managed
 ENV UV_LINK_MODE=copy
 SHELL ["/bin/bash", "-c"]
 
-RUN uv python install 3.12
-
+RUN mkdir -p /workspace && chown -R ubuntu:ubuntu /workspace
+USER ubuntu
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=cache,target=/root/.cache/pip \
     --mount=type=bind,source=requirements.txt,target=/requirements.txt \
+    --mount=from=ghcr.io/astral-sh/uv,source=/uv,target=/bin/uv \
     <<EOF
 uv venv --python 3.12 --seed /workspace/.venv
 source /workspace/.venv/bin/activate
@@ -22,10 +20,5 @@ pip install --no-build-isolation transformer-engine[pytorch]==2.9.0
 uv pip install -r /requirements.txt
 EOF
 
-# We still need the devel image to be able to use torch.compile; otherwise, we can use the runtime image.
-FROM nvcr.io/nvidia/cuda:13.0.2-cudnn-devel-ubuntu24.04
-COPY --from=build --chown=ubuntu:ubuntu /python /python
-COPY --from=build --chown=ubuntu:ubuntu --chmod=777 /workspace/.venv /workspace/.venv
 ENV PATH="/workspace/.venv/bin:$PATH"
-USER ubuntu
 WORKDIR /workspace/bionemo