Dockerfile.gpu

# Multi-stage build for PyTorch Ray inference container (SageMaker + EC2)
# GPU variant - NVIDIA CUDA 12.9 base on Amazon Linux 2023

#############
# Stage 1: FFmpeg builder - needs CUDA devel image for nvcc + CUDA headers
#############
FROM nvidia/cuda:12.9.1-devel-amzn2023 AS ffmpeg-builder

RUN dnf install -y --setopt=install_weak_deps=False \
  gcc \
  gcc-c++ \
  git \
  make \
  nasm \
  yasm \
  zlib-devel \
  pkgconfig \
  diffutils \
  && dnf clean all && rm -rf /var/cache/dnf

# Install NVIDIA codec headers for hardware-accelerated video encode/decode (NVENC/NVDEC)
# Ref: https://docs.nvidia.com/video-technologies/video-codec-sdk/12.2/ffmpeg-with-nvidia-gpu/index.html
RUN git clone --depth 1 --branch n12.2.72.0 https://github.com/FFmpeg/nv-codec-headers.git /tmp/nv-codec-headers \
  && cd /tmp/nv-codec-headers \
  && make install \
  && rm -rf /tmp/nv-codec-headers

# Build FFmpeg with NVIDIA GPU acceleration
# --nvccflags targets compute_75 (Turing/T4) as minimum -- forward-compatible with all newer GPUs
ENV PATH="/usr/local/cuda/bin:${PATH}"
RUN git clone --depth 1 --branch n8.0.1 https://github.com/FFmpeg/FFmpeg.git /tmp/ffmpeg \
  && cd /tmp/ffmpeg \
  && PKG_CONFIG_PATH="/usr/local/lib/pkgconfig" \
    ./configure --prefix=/opt/ffmpeg \
    --enable-shared --disable-static --disable-doc \
    --enable-nonfree --enable-cuda-nvcc --enable-libnpp \
    --extra-cflags=-I/usr/local/cuda/include \
    --extra-ldflags="-L/usr/local/cuda/lib64 -L/usr/local/cuda-12.9/targets/x86_64-linux/lib" \
    --nvccflags="-gencode arch=compute_75,code=sm_75 -O2" \
  && make -j$(nproc) \
  && make install \
  && rm -rf /tmp/ffmpeg


#############
# Stage 2: Python builder - no CUDA toolkit needed, uses lightweight base
#############
FROM amazonlinux:2023 AS python-builder

RUN dnf install -y --setopt=install_weak_deps=False \
  gcc \
  gcc-c++ \
  make \
  tar \
  gzip \
  wget \
  openssl-devel \
  libffi-devel \
  bzip2-devel \
  xz-devel \
  zlib-devel \
  && dnf clean all && rm -rf /var/cache/dnf

# Build Python from source with security hardening
ARG PYTHON_VERSION=3.13.12
ARG PYTHON_SHORT_VERSION=3.13
COPY ./scripts/common/install_python.sh /tmp/install_python.sh
RUN bash /tmp/install_python.sh ${PYTHON_VERSION} && rm /tmp/install_python.sh

# Install pinned dependencies
COPY --from=ghcr.io/astral-sh/uv:0.11.10 /uv /usr/local/bin/uv
ENV UV_PROJECT_ENVIRONMENT=/usr/local
COPY ./docker/ray/pyproject.toml ./docker/ray/uv.lock /tmp/deps/
RUN --mount=type=cache,target=/root/.cache/uv cd /tmp/deps \
  && uv sync --frozen --no-dev --no-install-project \
  && rm -rf /tmp/deps


#############
# Stage 3: Runtime base - shared by SageMaker and EC2
#############
FROM nvidia/cuda:12.9.1-runtime-amzn2023 AS base

ARG CACHE_REFRESH=0
RUN dnf upgrade -y --security --releasever latest && dnf clean all && rm -rf /var/cache/dnf

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

ARG PYTHON="python"
ARG FRAMEWORK="ray"
ARG FRAMEWORK_VERSION="2.55.1"
ARG CONTAINER_TYPE="inference"

# Enable video capability to mount NVENC/NVDEC driver libraries
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility,video

ARG PYTHON_SHORT_VERSION=3.13

# Copy Python and packages from python-builder
COPY --from=python-builder /usr/local/bin/ /usr/local/bin/
COPY --from=python-builder /usr/local/lib/python${PYTHON_SHORT_VERSION} /usr/local/lib/python${PYTHON_SHORT_VERSION}
COPY --from=python-builder /usr/local/lib/libpython${PYTHON_SHORT_VERSION}* /usr/local/lib/
COPY --from=python-builder /usr/local/include/python${PYTHON_SHORT_VERSION} /usr/local/include/python${PYTHON_SHORT_VERSION}

# Copy FFmpeg from ffmpeg-builder
COPY --from=ffmpeg-builder /opt/ffmpeg /opt/ffmpeg
RUN echo "/usr/local/lib" >/etc/ld.so.conf.d/python.conf && ldconfig

# Set PATH for FFmpeg
ENV PATH="/opt/ffmpeg/bin:$PATH" \
  LD_LIBRARY_PATH="/opt/ffmpeg/lib:${LD_LIBRARY_PATH}" \
  DLC_CONTAINER_TYPE=inference

# Create model directories
RUN mkdir -p /opt/ml/model /opt/ml/input /opt/ml/output

WORKDIR /app

# Telemetry and OSS compliance
COPY ./scripts/telemetry/deep_learning_container.py /usr/local/bin/deep_learning_container.py
COPY ./scripts/telemetry/bash_telemetry.sh.template /tmp/bash_telemetry.sh.template
COPY ./scripts/common/setup_oss_compliance.sh setup_oss_compliance.sh
RUN chmod +x /usr/local/bin/deep_learning_container.py \
  && sed -e "s/{{FRAMEWORK}}/${FRAMEWORK}/g" \
    -e "s/{{FRAMEWORK_VERSION}}/${FRAMEWORK_VERSION}/g" \
    -e "s/{{CONTAINER_TYPE}}/${CONTAINER_TYPE}/g" \
    /tmp/bash_telemetry.sh.template >/usr/local/bin/bash_telemetry.sh \
  && chmod +x /usr/local/bin/bash_telemetry.sh \
  && rm /tmp/bash_telemetry.sh.template \
  && echo 'source /usr/local/bin/bash_telemetry.sh' >>/etc/bash.bashrc \
  && echo 'source /usr/local/bin/bash_telemetry.sh' >>/etc/bashrc \
  && echo 'source /usr/local/bin/bash_telemetry.sh' >>/root/.bashrc \
  && bash setup_oss_compliance.sh ${PYTHON} \
  && printf '\n** FFmpeg; version n8.0.1 -- https://github.com/FFmpeg/FFmpeg/tree/n8.0.1\n' >> /root/BUILD_FROM_SOURCE_PACKAGES_LICENCES \
  && printf '\nFFmpeg n8.0.1 https://github.com/FFmpeg/FFmpeg/tree/n8.0.1\n' >> /root/THIRD_PARTY_SOURCE_CODE_URLS \
  && printf '\n** FFprobe; version n8.0.1 -- https://github.com/FFmpeg/FFmpeg/tree/n8.0.1\n' >> /root/BUILD_FROM_SOURCE_PACKAGES_LICENCES \
  && printf '\nFFprobe n8.0.1 https://github.com/FFmpeg/FFmpeg/tree/n8.0.1\n' >> /root/THIRD_PARTY_SOURCE_CODE_URLS \
  && rm setup_oss_compliance.sh \
  && rm -rf /root/oss_compliance* \
  && rm -rf /tmp/tmp* \
  && rm -rf /root/.cache | true


#############
# Stage 4: SageMaker image
#############
FROM base AS ray-sagemaker-gpu

# SageMaker environment variables
ENV SM_MODEL_DIR=/opt/ml/model \
  SM_INPUT=/opt/ml/input \
  SM_INPUT_TRAINING_CONFIG_FILE=/opt/ml/input/config/hyperparameters.json \
  SM_INPUT_DATA_CONFIG_FILE=/opt/ml/input/config/inputdataconfig.json \
  SM_CHECKPOINT_CONFIG_FILE=/opt/ml/input/config/checkpointconfig.json

COPY ./scripts/ray/sagemaker_serve.py /app/
COPY ./scripts/ray/sagemaker_entrypoint.sh /usr/local/bin/
RUN chmod +x /usr/local/bin/sagemaker_entrypoint.sh

EXPOSE 8080

LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
LABEL com.amazonaws.sagemaker.capabilities.multi-models=false

ENTRYPOINT ["/usr/local/bin/sagemaker_entrypoint.sh"]
CMD ["serve"]


#############
# Stage 5: EC2 image
#############
FROM base AS ray-ec2-gpu

COPY ./scripts/ray/dockerd_entrypoint.sh /usr/local/bin/
RUN chmod +x /usr/local/bin/dockerd_entrypoint.sh

ENTRYPOINT ["/usr/local/bin/dockerd_entrypoint.sh"]