Updated Dockerfiles for 2.24 release (#109)

Updated Dockerfiles

---------

Co-authored-by: kaenafwi <[email protected]>
Co-authored-by: Eddy Varela <[email protected]>

Author: 3 people

docker/common/vllm_entrypoint.py

@@ -0,0 +1,4 @@
+import subprocess
+import sys
+
+subprocess.check_call(sys.argv[1:])

docker/common/vllm_requirements.txt

@@ -0,0 +1,39 @@
+psutil
+sentencepiece  # Required for LLaMA tokenizer.
+numpy < 2.0.0
+requests >= 2.26.0
+tqdm
+blake3
+py-cpuinfo
+transformers >= 4.48.2  # Required for Bamba model and Transformers backend.
+tokenizers >= 0.19.1  # Required for Llama 3.
+protobuf # Required by LlamaTokenizer.
+fastapi >= 0.107.0, < 0.113.0; python_version < '3.9'
+fastapi >= 0.107.0, != 0.113.*, != 0.114.0; python_version >= '3.9'
+aiohttp
+openai >= 1.52.0 # Ensure modern openai package (ensure types module present and max_completion_tokens field support)
+uvicorn[standard]
+pydantic >= 2.9  # Required for fastapi >= 0.113.0
+prometheus_client >= 0.18.0
+pillow  # Required for image processing
+prometheus-fastapi-instrumentator >= 7.0.0
+tiktoken >= 0.6.0  # Required for DBRX tokenizer
+lm-format-enforcer >= 0.10.9, < 0.11
+outlines == 0.1.11
+lark == 1.2.2 
+xgrammar >= 0.1.6; platform_machine == "x86_64"
+typing_extensions >= 4.10
+filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
+partial-json-parser # used for parsing partial JSON outputs
+pyzmq
+msgspec
+gguf == 0.10.0
+importlib_metadata
+mistral_common[opencv] >= 1.5.0
+pyyaml
+six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
+setuptools>=74.1.1; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
+einops # Required for Qwen2-VL.
+compressed-tensors == 0.9.1 # required for compressed-tensors
+depyf==0.18.0 # required for profiling and debugging with compilation config
+cloudpickle # allows pickling lambda functions in model_executor/models/registry.py

docker/jax/training/0.6/Dockerfile.neuronx

@@ -0,0 +1,276 @@
+ARG BUILD_STAGE=prod
+
+FROM public.ecr.aws/docker/library/ubuntu:22.04 AS base
+
+LABEL dlc_major_version="1"
+LABEL maintainer="Amazon AI"
+
+# This arg required to stop docker build waiting for region configuration while installing tz data from ubuntu 22
+ARG DEBIAN_FRONTEND=noninteractive
+ARG PYTHON=python3.10
+ARG PYTHON_VERSION=3.10.12
+ARG PIP=pip3
+ARG OMPI_VERSION=4.1.5
+
+# 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 LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/aws/neuron/lib"
+ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/efa/lib"
+ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/efa/lib64"
+ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/openmpi/lib64"
+ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/lib"
+ENV PATH="/opt/aws/neuron/bin:${PATH}"
+
+RUN apt-get update \
+ && apt-get upgrade -y \
+ && apt-get install -y --no-install-recommends \
+    build-essential \
+    ca-certificates \
+    cmake \
+    curl \
+    emacs \
+    git \
+    gnupg2 \
+    gpg-agent \
+    jq \
+    libopencv-dev \
+    libglib2.0-0 \
+    libgl1-mesa-glx \
+    libsm6 \
+    libxext6 \
+    libxrender-dev \
+    libssl-dev \
+    libsqlite3-dev \
+    libgdbm-dev \
+    libc6-dev \
+    libbz2-dev \
+    libncurses-dev \
+    libffi-dev \
+    libcap-dev \
+    libhwloc-dev \
+    openjdk-8-jdk-headless \
+    openjdk-8-jdk \
+    openjdk-8-jre \
+    openjdk-11-jdk \
+    openssl \
+    software-properties-common \
+    tk-dev \
+    unzip \
+    wget \
+    vim \
+    zlib1g-dev \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean
+
+# Install Open MPI
+RUN mkdir -p /tmp/openmpi \
+ && cd /tmp/openmpi \
+ && wget --quiet https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-${OMPI_VERSION}.tar.gz \
+ && tar zxf openmpi-${OMPI_VERSION}.tar.gz \
+ && cd openmpi-${OMPI_VERSION} \
+ && ./configure --enable-orterun-prefix-by-default \
+ && make -j $(nproc) all \
+ && make install \
+ && ldconfig \
+ && rm -rf /tmp/openmpi
+
+# Install packages and configure SSH for MPI operator in k8s
+RUN apt-get update \
+ && apt-get install -y openmpi-bin openssh-server \
+ && mkdir -p /var/run/sshd \
+ && echo "    UserKnownHostsFile /dev/null" >> /etc/ssh/ssh_config \
+ && echo "    StrictHostKeyChecking no" >> /etc/ssh/ssh_config \
+ && sed -i 's/#\(StrictModes \).*/\1no/g' /etc/ssh/sshd_config \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean
+
+# Install Python
+RUN wget -q https://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VERSION.tgz \
+ && tar -xzf Python-$PYTHON_VERSION.tgz \
+ && cd Python-$PYTHON_VERSION \
+ && ./configure --enable-shared --prefix=/usr/local \
+ && make -j $(nproc) && make install \
+ && cd .. && rm -rf ../Python-$PYTHON_VERSION* \
+ && ln -s /usr/local/bin/pip3 /usr/bin/pip \
+ && ln -s /usr/local/bin/$PYTHON /usr/local/bin/python \
+ && ${PIP} --no-cache-dir install --upgrade \
+    "awscli<2" \
+    pip \
+    requests \
+    setuptools \
+ && rm -rf ~/.cache/pip/*
+
+# Install EFA
+RUN apt-get update \
+ && cd $HOME \
+ && curl -O https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz \
+ && wget https://efa-installer.amazonaws.com/aws-efa-installer.key && gpg --import aws-efa-installer.key \
+ && cat aws-efa-installer.key | gpg --fingerprint \
+ && wget https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz.sig && gpg --verify ./aws-efa-installer-latest.tar.gz.sig \
+ && tar -xf aws-efa-installer-latest.tar.gz \
+ && cd aws-efa-installer \
+ && ./efa_installer.sh -y -g --skip-kmod --skip-limit-conf --no-verify \
+ && cd $HOME \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean
+
+WORKDIR /
+
+# The ENV variables declared below are changed in the previous section
+# Grouping these ENV variables in the first section causes
+# ompi_info to fail. This is only observed in CPU containers
+ENV PATH="$PATH:/home/.openmpi/bin"
+ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/home/.openmpi/lib/"
+RUN ompi_info --parsable --all | grep mpi_built_with_cuda_support:value
+
+RUN mkdir -p /etc/pki/tls/certs && cp /etc/ssl/certs/ca-certificates.crt /etc/pki/tls/certs/ca-bundle.crt
+
+# Copy workaround script for incorrect hostname
+COPY changehostname.c /
+COPY --chmod=755 start_with_right_hostname.sh deep_learning_container.py /usr/local/bin/
+
+RUN HOME_DIR=/root \
+ && curl -o ${HOME_DIR}/oss_compliance.zip https://aws-dlinfra-utilities.s3.amazonaws.com/oss_compliance.zip \
+ && unzip ${HOME_DIR}/oss_compliance.zip -d ${HOME_DIR}/ \
+ && cp ${HOME_DIR}/oss_compliance/test/testOSSCompliance /usr/local/bin/testOSSCompliance \
+ && chmod +x /usr/local/bin/testOSSCompliance \
+ && chmod +x ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh \
+ && ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh ${HOME_DIR} ${PYTHON} \
+ && rm -rf ${HOME_DIR}/oss_compliance* \
+ && rm -rf /tmp/tmp*
+
+# Setting up APT and PIP repo for neuron artifacts
+ARG NEURON_APT_REPO=https://apt.repos.neuron.amazonaws.com
+ARG NEURON_APT_REPO_KEY
+ARG NEURON_PIP_REPO=https://pip.repos.neuron.amazonaws.com
+ARG NEURON_PIP_REPO_KEY
+RUN mkdir -p /etc/apt/keyrings \
+ && APT_REPO_PREFIX=$([ -n "${NEURON_APT_REPO_KEY}" ] && echo "${NEURON_APT_REPO_KEY}@" || echo "") \
+ && echo "deb [signed-by=/etc/apt/keyrings/neuron.gpg] https://${APT_REPO_PREFIX}${NEURON_APT_REPO} focal main" > /etc/apt/sources.list.d/neuron.list \
+ && curl $([ -n "${NEURON_APT_REPO_KEY}" ] && echo "-u ${NEURON_APT_REPO_KEY}") -sSL "https://${NEURON_APT_REPO}/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB" | gpg --dearmor > /etc/apt/keyrings/neuron.gpg \
+ && PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP_REPO_KEY}@${NEURON_PIP_REPO}" || echo "https://${NEURON_PIP_REPO}") \
+ && ${PIP} config set global.extra-index-url "${PIP_REPO_URL}"
+
+# Neuron SDK components version numbers
+ARG NEURON_ARTIFACT_PATH=/root/neuron_artifacts
+ARG IGNORE_MISSING_NEURON_COMPONENTS=false
+RUN IGNORE_MISSING_NEURON_COMPONENTS=$(echo ${IGNORE_MISSING_NEURON_COMPONENTS} | tr '[:upper:]' '[:lower:]')
+
+ARG NEURONX_RUNTIME_LIB_VERSION=2.26.42.0-2ff3b5c7d
+ARG NEURONX_COLLECTIVES_LIB_VERSION=2.26.43.0-47cc904ea
+ARG NEURONX_TOOLS_VERSION=2.24.54.0
+
+ARG NEURONX_CC_VERSION=2.19.8089.0
+ARG NEURONX_JAX_TRAINING_VERSION=0.6.0.1.0.1296+1f770067
+
+FROM base AS dev
+
+RUN --mount=type=bind,source=apt,target=${NEURON_ARTIFACT_PATH}/apt \
+    install_apt_package() { \
+        pkg_name=$1; \
+        version_arg=$2; \
+        if [ -f "${NEURON_ARTIFACT_PATH}/apt/${version_arg}" ]; then \
+            apt-get install -y ${NEURON_ARTIFACT_PATH}/apt/${version_arg}; \
+        elif [ "${IGNORE_MISSING_NEURON_COMPONENTS}" = "false" ]; then \
+            apt-get install -y ${pkg_name}=${version_arg}; \
+        else \
+            echo "Ignoring package ${pkg_name}"; \
+        fi; \
+    } \
+ && apt-get update \
+ && install_apt_package "aws-neuronx-collectives" "${NEURONX_COLLECTIVES_LIB_VERSION}" \
+ && install_apt_package "aws-neuronx-runtime-lib" "${NEURONX_RUNTIME_LIB_VERSION}" \
+ && install_apt_package "aws-neuronx-tools" "${NEURONX_TOOLS_VERSION}" \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean
+
+RUN --mount=type=bind,source=pip,target=${NEURON_ARTIFACT_PATH}/pip \
+    install_pip_package() { \
+        packages=""; \
+        flags=""; \
+        while [ "$#" -gt 0 ]; do \
+            pkg_name=$(echo $1 | cut -d: -f1); \
+            version_arg=$(echo $1 | cut -d: -f2); \
+            extra_flags=$(echo $1 | cut -d: -f3); \
+            if [ -f "${NEURON_ARTIFACT_PATH}/pip/${version_arg}" ]; then \
+                packages="${packages} ${NEURON_ARTIFACT_PATH}/pip/${version_arg}"; \
+            else \
+                if [ "${IGNORE_MISSING_NEURON_COMPONENTS}" = "false" ]; then \
+                    packages="${packages} ${pkg_name}==${version_arg}"; \
+                else \
+                    echo "Ignoring package ${pkg_name}"; \
+                fi; \
+            fi; \
+            # Store unique flags
+            if [ ! -z "${extra_flags}" ]; then \
+                for flag in $(echo "${extra_flags}" | tr ' ' '\n'); do \
+                    case " ${flags} " in \
+                        *" ${flag} "*) ;; \
+                        *) flags="${flags} ${flag}" ;; \
+                    esac \
+                done; \
+            fi; \
+            shift; \
+        done; \
+        if [ ! -z "${packages}" ]; then \
+            echo "Installing packages: ${packages} with flags ${flags}"; \
+            ${PIP} install --no-cache-dir --force-reinstall \
+                --extra-index-url="file:///${NEURON_ARTIFACT_PATH}/pip" \
+                ${packages} ${flags}; \
+        fi; \
+    } \
+ && install_pip_package "neuronx-cc:${NEURONX_CC_VERSION}:" "jax-neuronx:${NEURONX_JAX_TRAINING_VERSION}:" \
+ && rm -rf ~/.cache/pip/*
+
+FROM base AS repo
+
+# Install Neuron components from the apt and pip repos 
+RUN apt-get update \
+ && apt-get install -y \
+    aws-neuronx-tools \
+    aws-neuronx-collectives \
+    aws-neuronx-runtime-lib \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean \
+ && ${PIP} install --no-cache-dir --force-reinstall \
+    neuronx-cc \
+    jax-neuronx \
+ && rm -rf ~/.cache/pip/*
+
+
+FROM base AS prod
+
+# Install Neuron components
+# Install Neuron Driver, Runtime and Tools
+RUN apt-get update \
+ && apt-get install -y \
+    aws-neuronx-tools=$NEURONX_TOOLS_VERSION \
+    aws-neuronx-collectives=$NEURONX_COLLECTIVES_LIB_VERSION \
+    aws-neuronx-runtime-lib=$NEURONX_RUNTIME_LIB_VERSION \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean
+
+# Install JAX & Neuron CC
+RUN ${PIP} install --no-cache-dir --force-reinstall \
+    neuronx-cc==$NEURONX_CC_VERSION \
+    jax-neuronx==$NEURONX_JAX_TRAINING_VERSION \
+ && rm -rf ~/.cache/pip/*
+
+FROM ${BUILD_STAGE} AS final
+
+# Starts framework
+ENTRYPOINT ["bash", "-m", "start_with_right_hostname.sh"]
+CMD ["/bin/bash"]
+
+HEALTHCHECK CMD curl --fail http://localhost:8080/ping || exit 1