Dockerfile.gpu

ARG PYTHON=python3
ARG PYTHON_VERSION=3.12.10
ARG PYTHON_SHORT_VERSION=3.12
ARG PYTORCH_VERSION=2.8.0
ARG TORCHAUDIO_VERSION=2.8.0
ARG TORCHVISION_VERSION=0.23.0
ARG RAY_VERSION=2.49.0

#################################################################
#   ____
#  / ___| ___  _ __ ___  _ __ ___   ___  _ __
# | |    / _ \| '_ ` _ \| '_ ` _ \ / _ \| '_ \
# | |___  (_) | | | | | | | | | | | (_) | | | |
#  \____|\___/|_| |_| |_|_| |_| |_|\___/|_| |_|
#  ___                              ____           _
# |_ _|_ __ ___   __ _  __ _  ___  |  _ \ ___  ___(_)_ __   ___
#  | || '_ ` _ \ / _` |/ _` |/ _ \ | |_) / _ \/ __| | '_ \ / _ \
#  | || | | | | | (_| | (_| |  __/ |  _ <  __/ (__| | |_) |  __/
# |___|_| |_| |_|\__,_|\__, |\___| |_| \_\___|\___|_| .__/ \___|
#                      |___/                        |_|
#################################################################

FROM nvidia/cuda:12.9.1-cudnn-runtime-ubuntu22.04 AS common

LABEL maintainer="Amazon AI"
LABEL dlc_major_version="1"

ARG PYTHON
ARG PYTHON_VERSION
ARG PYTORCH_VERSION
ARG TORCHAUDIO_VERSION
ARG TORCHVISION_VERSION
ARG RAY_VERSION

ENV CUDA_HOME="/usr/local/cuda"
ENV PATH="${CUDA_HOME}/bin:${PATH}"

# Python won't try to write .pyc or .pyo files on the import of source modules
# Force stdin, stdout and stderr to be totally unbuffered. Good for logging
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
ENV PYTHONIOENCODING=UTF-8
ENV LANG=C.UTF-8
ENV LC_ALL=C.UTF-8

ENV DLC_CONTAINER_TYPE=inference
WORKDIR /

RUN apt-get update \
 && apt-get -y upgrade --only-upgrade systemd \
 && apt-get install -y --allow-change-held-packages --no-install-recommends \
    libgl1-mesa-glx \
    build-essential \
    ca-certificates \
    zlib1g-dev \
    openssl \
    libssl-dev \
    pkg-config \
    check \
    llvm \
    xz-utils \
    curl \
    wget \
    unzip \
    libffi-dev \
    libbz2-dev \
    liblzma-dev \
    libsqlite3-dev \
    libreadline-dev \
 && rm -rf /var/lib/apt/lists/* \
 && apt-get clean

COPY install_python.sh install_python.sh
RUN bash install_python.sh ${PYTHON_VERSION} && rm install_python.sh

# Python Path - nvidia/cuda base already has CUDA paths configured
ENV PATH="/usr/local/bin:${PATH}"

# Install PyTorch first (needs specific index-url)
RUN pip install --no-cache-dir \
    torch==${PYTORCH_VERSION} \
    torchvision==${TORCHVISION_VERSION} \
    torchaudio==${TORCHAUDIO_VERSION} \
    --index-url https://download.pytorch.org/whl/cu129 \
 && pip install --no-cache-dir "triton==3.2.0" \
 && pip uninstall -y dataclasses || true

# Install Ray Serve + all dependencies in one optimized command
RUN pip install --no-cache-dir \
    ray[serve]==${RAY_VERSION} \
    cryptography \
    pyOpenSSL \
    mkl \
    charset-normalizer \
    packaging \
    PyYAML \
    numpy \
    scipy \
    click \
    psutil \
    pillow \
    h5py \
    fsspec \
    "idna>=3.7" \
    "tqdm>=4.66.3" \
    "requests>=2.32.0" \
    "setuptools>=70.0.0" \
    "urllib3>=2.5.0" \
    opencv-python==4.11.0.86 \
    jinja2>=3.1.6 \
    tornado>=6.5.1

RUN curl -o /license.txt https://aws-dlc-licenses.s3.amazonaws.com/pytorch-2.8/license.txt

# Removing the cache as it is needed for security verification
RUN rm -rf /root/.cache | true

########################################################
#  _____ ____ ____    ___
# | ____/ ___|___ \  |_ _|_ __ ___   __ _  __ _  ___
# |  _|| |     __) |  | || '_ ` _ \ / _` |/ _` |/ _ \
# | |__| |___ / __/   | || | | | | | (_| | (_| |  __/
# |_____\____|_____| |___|_| |_| |_|\__,_|\__, |\___|
#                                         |___/
#  ____           _
# |  _ \ ___  ___(_)_ __   ___
# | |_) / _ \/ __| | '_ \ / _ \
# |  _ <  __/ (__| | |_) |  __/
# |_| \_\___|\___|_| .__/ \___|
#                  |_|
########################################################

FROM common AS ec2

ARG PYTHON

WORKDIR /

COPY setup_oss_compliance.sh setup_oss_compliance.sh
RUN bash setup_oss_compliance.sh ${PYTHON} && rm setup_oss_compliance.sh

# This arg required to stop docker build waiting for region configuration while installing tz data from ubuntu 20
ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update \
 && apt-get upgrade -y \
 && apt-get autoremove -y \
 && apt-get clean \
 && rm -rf /var/lib/apt/lists/*

COPY rayserve-entrypoint.py /usr/local/bin/dockerd-entrypoint.py
RUN chmod +x /usr/local/bin/dockerd-entrypoint.py

# Port 8000: Ray Serve default serving port, Port 8265: Ray dashboard
# References : https://docs.ray.io/en/latest/serve/production-guide/config.html
#            https://docs.ray.io/en/latest/ray-observability/getting-started.html
EXPOSE 8000 8265
ENTRYPOINT ["python", "/usr/local/bin/dockerd-entrypoint.py"]