prime-rl/Dockerfile.cuda at main · ob1-s/prime-rl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
#inspired from https://github.com/astral-sh/uv-docker-example/blob/main/multistage.Dockerfile

# Build stage
FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel AS builder
LABEL maintainer="prime intellect"
LABEL repository="prime-rl"

# Set en_US.UTF-8 locale by default
RUN echo "LC_ALL=en_US.UTF-8" >> /etc/environment

# Set CUDA_HOME and update PATH
ENV CUDA_HOME=/usr/local/cuda
ENV PATH=$PATH:/usr/local/cuda/bin

# Install packages
RUN apt-get update && apt-get install -y --no-install-recommends --force-yes \
    build-essential \
    curl \
    sudo \
    git \
    && apt-get clean autoclean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*

# Download the latest installer
ADD https://astral.sh/uv/install.sh /uv-installer.sh

# Set install dir to location accessible to non-root users
RUN INSTALLER_NO_MODIFY_PATH=1 UV_INSTALL_DIR="/usr/local/bin" sh /uv-installer.sh && rm /uv-installer.sh
ENV PATH="/usr/local/bin:$PATH"
ENV UV_PYTHON_INSTALL_DIR="/usr/local/share/uv/python"
ENV UV_CACHE_DIR="/usr/local/share/uv/cache"

# Install Python dependencies (The gradual copies help with caching)
WORKDIR /app

COPY ./pyproject.toml ./pyproject.toml
COPY ./uv.lock ./uv.lock
COPY ./README.md ./README.md
COPY ./src/ ./src/

ENV UV_COMPILE_BYTECODE=1 UV_LINK_MODE=copy


COPY src/ /app/src/
COPY pyproject.toml /app/pyproject.toml
COPY uv.lock /app/uv.lock
COPY README.md /app/README.md
COPY configs /app/configs

RUN --mount=type=cache,target=/app/.cache/uv \
    uv sync --locked --no-dev

RUN --mount=type=cache,target=/app/.cache/uv \
    uv sync --all-extras --locked --no-dev

FROM python:3.12-slim

RUN apt-get update && apt-get install -y \
    --no-install-recommends \
    --force-yes \
    build-essential \
    wget \
    clang \
    tmux \
    iperf \
    openssh-server \
    git-lfs \
    gpg \
    && apt-get clean autoclean \
    && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*


ARG USER_ID=1000
ARG GROUP_ID=1000
RUN groupadd --gid $GROUP_ID appuser && \
    useradd --uid $USER_ID --gid appuser --create-home --shell /bin/bash appuser

USER appuser
WORKDIR /app
# Copy the application from the builder
COPY --from=builder --chown=appuser:appuser /app /app

RUN rm /app/.venv/bin/python && ln -s /usr/local/bin/python /app/.venv/bin/python
RUN rm /app/.venv/bin/python3 && ln -s /usr/local/bin/python /app/.venv/bin/python3
RUN rm /app/.venv/bin/python3.12 && ln -s /usr/local/bin/python /app/.venv/bin/python3.12


# Place executables in the environment at the front of the path
ENV PATH="/app/.venv/bin:$PATH"
ENTRYPOINT ["python", "src/prime_rl/inference/server.py"]