Dockerfile: # The vLLM Dockerfile is used to construct vLLM image that can be directly used
  # to run the OpenAI compatible server.
  
  #################### BASE BUILD IMAGE ####################
  FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 AS dev
  
  RUN apt-get update -y \
      && apt-get install -y python3-pip git
  
  WORKDIR /workspace
  
  # install build and runtime dependencies
  COPY requirements.txt requirements.txt
  RUN --mount=type=cache,target=/root/.cache/pip \
      pip install -r requirements.txt
  
  # install development dependencies
  COPY requirements-dev.txt requirements-dev.txt
  RUN --mount=type=cache,target=/root/.cache/pip \
      pip install -r requirements-dev.txt
  #################### BASE BUILD IMAGE ####################
  
  
  #################### EXTENSION BUILD IMAGE ####################
  FROM dev AS build
  
  # install build dependencies
  COPY requirements-build.txt requirements-build.txt
  RUN --mount=type=cache,target=/root/.cache/pip \
      pip install -r requirements-build.txt
  
  # copy input files
  COPY csrc csrc
  COPY setup.py setup.py
  COPY requirements.txt requirements.txt
  COPY pyproject.toml pyproject.toml
  COPY vllm/__init__.py vllm/__init__.py
  
  # cuda arch list used by torch
  ARG torch_cuda_arch_list='7.0 7.5 8.0 8.6 8.9 9.0+PTX'
  ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
  # max jobs used by Ninja to build extensions
  ARG max_jobs=2
  ENV MAX_JOBS=${max_jobs}
  # number of threads used by nvcc
  ARG nvcc_threads=8
  ENV NVCC_THREADS=$nvcc_threads
  # make sure punica kernels are built (for LoRA)
  ENV VLLM_INSTALL_PUNICA_KERNELS=1
  
  RUN python3 setup.py build_ext --inplace
  #################### EXTENSION Build IMAGE ####################
  
  
  #################### TEST IMAGE ####################
  # image to run unit testing suite
  FROM dev AS test
  
  # copy pytorch extensions separately to avoid having to rebuild
  # when python code changes
  WORKDIR /vllm-workspace
  # ADD is used to preserve directory structure
  ADD . /vllm-workspace/
  COPY --from=build /workspace/vllm/*.so /vllm-workspace/vllm/
  # ignore build dependencies installation because we are using pre-complied extensions
  RUN rm pyproject.toml
  RUN --mount=type=cache,target=/root/.cache/pip VLLM_USE_PRECOMPILED=1 pip install . --verbose
  #################### TEST IMAGE ####################
  
  
  #################### RUNTIME BASE IMAGE ####################
  # use CUDA base as CUDA runtime dependencies are already installed via pip
  FROM nvidia/cuda:12.1.0-base-ubuntu22.04 AS vllm-base
  
  # libnccl required for ray
  RUN apt-get update -y \
      && apt-get install -y python3-pip
  
  WORKDIR /workspace
  COPY requirements.txt requirements.txt
  RUN --mount=type=cache,target=/root/.cache/pip \
      pip install -r requirements.txt
  #################### RUNTIME BASE IMAGE ####################
  
  
  #################### OPENAI API SERVER ####################
  # openai api server alternative
  FROM vllm-base AS vllm-openai
  # install additional dependencies for openai api server
  RUN --mount=type=cache,target=/root/.cache/pip \
      pip install accelerate
  
  COPY --from=build /workspace/vllm/*.so /workspace/vllm/
  COPY vllm vllm
  
  ENTRYPOINT ["python3", "-m", "vllm.entrypoints.openai.api_server"]
  #################### OPENAI API SERVER ####################
  
Dockerfile.rocm: # default base image
  ARG BASE_IMAGE="rocm/pytorch:rocm6.0_ubuntu20.04_py3.9_pytorch_2.1.1"
  
  FROM $BASE_IMAGE
  
  ARG BASE_IMAGE="rocm/pytorch:rocm6.0_ubuntu20.04_py3.9_pytorch_2.1.1"
  
  RUN echo "Base image is $BASE_IMAGE"
  
  # BASE_IMAGE for ROCm_5.7: "rocm/pytorch:rocm5.7_ubuntu22.04_py3.10_pytorch_2.0.1"
  # BASE_IMAGE for ROCm_6.0: "rocm/pytorch:rocm6.0_ubuntu20.04_py3.9_pytorch_2.1.1"
  
  # this does not always work for all rocm versions
  RUN LLVM_GFX_ARCH=$(/opt/rocm/llvm/bin/amdgpu-offload-arch) && \
      echo "LLVM_GFX_ARCH is $LLVM_GFX_ARCH"
  
  ARG FA_GFX_ARCHS="gfx90a;gfx942"
  RUN echo "FA_GFX_ARCHS is $FA_GFX_ARCHS"
  
  ARG FA_BRANCH="3d2b6f5"
  RUN echo "FA_BRANCH is $FA_BRANCH"
  
  # Install some basic utilities
  RUN apt-get update && apt-get install python3 python3-pip -y
  
  # Install some basic utilities
  RUN apt-get update && apt-get install -y \
      curl \
      ca-certificates \
      sudo \
      git \
      bzip2 \
      libx11-6 \
      build-essential \
      wget \
      unzip \
      nvidia-cuda-toolkit \
      tmux \
   && rm -rf /var/lib/apt/lists/*
  
  ### Mount Point ###
  # When launching the container, mount the code directory to /app
  ARG APP_MOUNT=/app
  VOLUME [ ${APP_MOUNT} ]
  WORKDIR ${APP_MOUNT}
  
  RUN python3 -m pip install --upgrade pip
  RUN python3 -m pip install --no-cache-dir fastapi ninja tokenizers pandas
  
  ENV LLVM_SYMBOLIZER_PATH=/opt/rocm/llvm/bin/llvm-symbolizer
  ENV PATH=$PATH:/opt/rocm/bin:/libtorch/bin:
  ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/rocm/lib/:/libtorch/lib:
  ENV CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:/libtorch/include:/libtorch/include/torch/csrc/api/include/:/opt/rocm/include/:
  
  # Install ROCm flash-attention
  RUN mkdir libs \
      && cd libs \
      && git clone https://github.com/ROCmSoftwarePlatform/flash-attention.git \
      && cd flash-attention \
      && git checkout ${FA_BRANCH} \
      && git submodule update --init \
      && export GPU_ARCHS=${FA_GFX_ARCHS} \
      && if [ "$BASE_IMAGE" = "rocm/pytorch:rocm5.7_ubuntu22.04_py3.10_pytorch_2.0.1" ]; then \
          patch /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/utils/hipify/hipify_python.py hipify_patch.patch; fi \
      && python3 setup.py install \
      && cd ..
  
  COPY ./ /app/vllm
  
  RUN python3 -m pip install --upgrade pip
  RUN python3 -m pip install xformers==0.0.23 --no-deps
  
  # Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt.
  # Manually removed it so that later steps of numpy upgrade can continue
  RUN if [ "$BASE_IMAGE" = "rocm/pytorch:rocm6.0_ubuntu20.04_py3.9_pytorch_2.1.1" ]; then \
      rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/; fi
  
  RUN cd /app \
      && cd vllm \
      && pip install -U -r requirements-rocm.txt \
      && bash patch_xformers.rocm.sh \
      && python3 setup.py install \
      && cd ..
  
  RUN python3 -m pip install --upgrade pip
  RUN python3 -m pip install --no-cache-dir ray[all]
  
  CMD ["/bin/bash"]
  