From e2c5e9170108f9aefda55a0806bb0d274d2b973c Mon Sep 17 00:00:00 2001 From: kaenafwi Date: Fri, 27 Feb 2026 01:04:38 +0000 Subject: [PATCH] Updated Dockerfiles based on build artifacts uploaded to the Neuron DLC Builder prod account --- jax/training/0.7/Dockerfile.neuronx | 15 ++++--- pytorch/inference/2.9.0/Dockerfile.neuronx | 18 +++++---- pytorch/training/2.9.0/Dockerfile.neuronx | 17 +++++--- vllm/inference/0.13.0/Dockerfile.neuronx | 47 +++++++++------------- 4 files changed, 52 insertions(+), 45 deletions(-) diff --git a/jax/training/0.7/Dockerfile.neuronx b/jax/training/0.7/Dockerfile.neuronx index 7b2466e..f601d94 100644 --- a/jax/training/0.7/Dockerfile.neuronx +++ b/jax/training/0.7/Dockerfile.neuronx @@ -106,6 +106,7 @@ RUN wget -q https://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VER pip \ requests \ setuptools \ + uv \ && rm -rf ~/.cache/pip/* # U24 will not allow installation of pip packages outside of venv without this flag @@ -165,11 +166,13 @@ RUN mkdir -p /etc/apt/keyrings \ && curl $([ -n "${NEURON_APT_REPO_KEY}" ] && echo "-u ${NEURON_APT_REPO_KEY}") --retry 3 --retry-delay 1 --retry-all-errors -fSL "https://${NEURON_APT_REPO}/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB" | gpg --dearmor > /etc/apt/keyrings/neuron.gpg # Neuron SDK components version numbers -ARG NEURONX_RUNTIME_LIB_VERSION=2.29.40.0-f954cd7a5 -ARG NEURONX_COLLECTIVES_LIB_VERSION=2.29.41.0-681fef5f5 -ARG NEURONX_TOOLS_VERSION=2.27.33.0-5d9c0b901 -ARG NEURONX_CC_VERSION=2.22.12471.0+b4a00d10 -ARG NEURONX_JAX_TRAINING_VERSION=0.7.0.1.0.7377+5e6a4049 +ARG NEURONX_RUNTIME_LIB_VERSION=2.30.51.0-faafe26f0 +ARG NEURONX_COLLECTIVES_LIB_VERSION=2.30.59.0-f5cdefb39 +ARG NEURONX_TOOLS_VERSION=2.28.23.0-f1c114a9d + +ARG NEURONX_CC_VERSION=2.23.6484.0+3b612583 +ARG NEURONX_JAX_TRAINING_VERSION=0.7.0.1.0.7584+5c8f4c3c +ARG NKI_VERSION=0.2.0+g82fdb402 FROM base AS repo @@ -190,6 +193,7 @@ RUN PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP --trusted-host ${NEURON_PIP_REPO} \ "neuronx-cc>=2.0" \ jax-neuronx \ + nki \ && rm -rf ~/.cache/pip/* FROM base AS prod @@ -213,6 +217,7 @@ RUN PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP --extra-index-url ${PYPI_SIMPLE_URL} \ neuronx-cc==$NEURONX_CC_VERSION \ jax-neuronx==$NEURONX_JAX_TRAINING_VERSION \ + nki==$NKI_VERSION \ && rm -rf ~/.cache/pip/* FROM ${BUILD_STAGE} AS final diff --git a/pytorch/inference/2.9.0/Dockerfile.neuronx b/pytorch/inference/2.9.0/Dockerfile.neuronx index b656e68..244eedb 100644 --- a/pytorch/inference/2.9.0/Dockerfile.neuronx +++ b/pytorch/inference/2.9.0/Dockerfile.neuronx @@ -125,6 +125,7 @@ RUN ${PIP} install --upgrade pip --trusted-host pypi.org --trusted-host files.py "awscli<2" \ pandas \ boto3 \ + uv \ cryptography \ "protobuf>=3.18.3,<4" \ torchserve==${TORCHSERVE_VERSION} \ @@ -175,14 +176,15 @@ RUN mkdir -p /etc/apt/keyrings \ && curl $([ -n "${NEURON_APT_REPO_KEY}" ] && echo "-u ${NEURON_APT_REPO_KEY}") --retry 3 --retry-delay 1 --retry-all-errors -fSL "https://${NEURON_APT_REPO}/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB" | gpg --dearmor > /etc/apt/keyrings/neuron.gpg # Neuron SDK components version numbers -ARG NEURONX_COLLECTIVES_LIB_VERSION=2.29.41.0-681fef5f5 -ARG NEURONX_RUNTIME_LIB_VERSION=2.29.40.0-f954cd7a5 -ARG NEURONX_TOOLS_VERSION=2.27.33.0-5d9c0b901 +ARG NEURONX_COLLECTIVES_LIB_VERSION=2.30.59.0-f5cdefb39 +ARG NEURONX_RUNTIME_LIB_VERSION=2.30.51.0-faafe26f0 +ARG NEURONX_TOOLS_VERSION=2.28.23.0-f1c114a9d -ARG NEURONX_CC_VERSION=2.22.12471.0+b4a00d10 -ARG NEURONX_FRAMEWORK_VERSION=2.9.0.2.11.19912+e48cd891 -ARG NEURONX_DISTRIBUTED_VERSION=0.16.25997+f431c02e -ARG NEURONX_DISTRIBUTED_INFERENCE_VERSION=0.7.15063+bafa28d5 +ARG NEURONX_CC_VERSION=2.23.6484.0+3b612583 +ARG NEURONX_FRAMEWORK_VERSION=2.9.0.2.12.22436+0f1dac25 +ARG NEURONX_DISTRIBUTED_VERSION=0.17.26814+4b18de63 +ARG NEURONX_DISTRIBUTED_INFERENCE_VERSION=0.8.16251+f3ca5575 +ARG NKI_VERSION=0.2.0+g82fdb402 FROM base AS repo @@ -205,6 +207,7 @@ RUN PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP "torch-neuronx==2.9.*" \ neuronx_distributed \ neuronx_distributed_inference \ + nki \ && rm -rf ~/.cache/pip/* FROM base AS prod @@ -228,6 +231,7 @@ RUN PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP torch-neuronx==$NEURONX_FRAMEWORK_VERSION \ neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION \ neuronx_distributed_inference==$NEURONX_DISTRIBUTED_INFERENCE_VERSION \ + nki==$NKI_VERSION \ && rm -rf ~/.cache/pip/* FROM ${BUILD_STAGE} AS final diff --git a/pytorch/training/2.9.0/Dockerfile.neuronx b/pytorch/training/2.9.0/Dockerfile.neuronx index a69dc6b..cc0439b 100644 --- a/pytorch/training/2.9.0/Dockerfile.neuronx +++ b/pytorch/training/2.9.0/Dockerfile.neuronx @@ -128,6 +128,7 @@ RUN ${PIP} install --no-cache-dir -U \ psutil==5.6.7 \ dataset \ Pillow \ + uv \ && rm -rf ~/.cache/pip/* RUN mkdir -p /etc/pki/tls/certs && cp /etc/ssl/certs/ca-certificates.crt /etc/pki/tls/certs/ca-bundle.crt @@ -211,13 +212,15 @@ RUN mkdir -p /etc/apt/keyrings \ && curl $([ -n "${NEURON_APT_REPO_KEY}" ] && echo "-u ${NEURON_APT_REPO_KEY}") --retry 3 --retry-delay 1 --retry-all-errors -fSL "https://${NEURON_APT_REPO}/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB" | gpg --dearmor > /etc/apt/keyrings/neuron.gpg # Neuron SDK components -ARG NEURONX_COLLECTIVES_LIB_VERSION=2.29.41.0-681fef5f5 -ARG NEURONX_RUNTIME_LIB_VERSION=2.29.40.0-f954cd7a5 -ARG NEURONX_TOOLS_VERSION=2.27.33.0-5d9c0b901 -ARG NEURONX_FRAMEWORK_VERSION=2.9.0.2.11.19912+e48cd891 -ARG NEURONX_CC_VERSION=2.22.12471.0+b4a00d10 -ARG NEURONX_DISTRIBUTED_VERSION=0.16.25997+f431c02e +ARG NEURONX_COLLECTIVES_LIB_VERSION=2.30.59.0-f5cdefb39 +ARG NEURONX_RUNTIME_LIB_VERSION=2.30.51.0-faafe26f0 +ARG NEURONX_TOOLS_VERSION=2.28.23.0-f1c114a9d + +ARG NEURONX_FRAMEWORK_VERSION=2.9.0.2.12.22436+0f1dac25 +ARG NEURONX_CC_VERSION=2.23.6484.0+3b612583 +ARG NEURONX_DISTRIBUTED_VERSION=0.17.26814+4b18de63 ARG NEURONX_DISTRIBUTED_TRAINING_VERSION=1.7.0 +ARG NKI_VERSION=0.2.0+g82fdb402 FROM base AS repo @@ -240,6 +243,7 @@ RUN PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP "neuronx-cc>=2.0" \ neuronx_distributed \ neuronx_distributed_training \ + nki \ && rm -rf ~/.cache/pip/* FROM base AS prod @@ -263,6 +267,7 @@ RUN PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP neuronx-cc==$NEURONX_CC_VERSION \ neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION \ neuronx_distributed_training==$NEURONX_DISTRIBUTED_TRAINING_VERSION \ + nki==$NKI_VERSION \ && rm -rf ~/.cache/pip/* FROM ${BUILD_STAGE} AS final diff --git a/vllm/inference/0.13.0/Dockerfile.neuronx b/vllm/inference/0.13.0/Dockerfile.neuronx index 36a651b..c20975f 100644 --- a/vllm/inference/0.13.0/Dockerfile.neuronx +++ b/vllm/inference/0.13.0/Dockerfile.neuronx @@ -123,6 +123,7 @@ RUN ${PIP} install --no-cache-dir -U \ "cmake>=3.26" \ "setuptools-scm>=8" \ "jinja2" \ + uv \ torchserve==${TORCHSERVE_VERSION} \ torch-model-archiver==${TORCHSERVE_VERSION} \ && rm -rf ~/.cache/pip/* @@ -155,30 +156,28 @@ RUN mkdir -p /etc/apt/keyrings \ && curl $([ -n "${NEURON_APT_REPO_KEY}" ] && echo "-u ${NEURON_APT_REPO_KEY}") --retry 3 --retry-delay 1 --retry-all-errors -fSL "https://${NEURON_APT_REPO}/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB" | gpg --dearmor > /etc/apt/keyrings/neuron.gpg # Neuron SDK components version numbers -ARG NEURONX_COLLECTIVES_LIB_VERSION=2.29.41.0-681fef5f5 -ARG NEURONX_RUNTIME_LIB_VERSION=2.29.40.0-f954cd7a5 -ARG NEURONX_TOOLS_VERSION=2.27.33.0-5d9c0b901 +ARG NEURONX_COLLECTIVES_LIB_VERSION=2.30.59.0-f5cdefb39 +ARG NEURONX_RUNTIME_LIB_VERSION=2.30.51.0-faafe26f0 +ARG NEURONX_TOOLS_VERSION=2.28.23.0-f1c114a9d -ARG NEURONX_CC_VERSION=2.22.12471.0+b4a00d10 -ARG NEURONX_FRAMEWORK_VERSION=2.9.0.2.11.19912+e48cd891 -ARG NEURONX_DISTRIBUTED_VERSION=0.16.25997+f431c02e -ARG NEURONX_DISTRIBUTED_INFERENCE_VERSION=0.7.15063+bafa28d5 -ARG NKI_VERSION=0.1.0+g432b459e +ARG NEURONX_CC_VERSION=2.23.6484.0+3b612583 +ARG NEURONX_FRAMEWORK_VERSION=2.9.0.2.12.22436+0f1dac25 +ARG NEURONX_DISTRIBUTED_VERSION=0.17.26814+4b18de63 +ARG NEURONX_DISTRIBUTED_INFERENCE_VERSION=0.8.16251+f3ca5575 +ARG NKI_VERSION=0.2.0+g82fdb402 -# vLLM branch names -ARG VLLM_PRIVATE_BRANCH=release-0.3.0 -ARG VLLM_PUBLIC_BRANCH=release-0.3.0 +# GitHub repository and branch +ARG GITHUB_REPO=https://github.com/vllm-project/vllm-neuron.git +ARG GITHUB_REPO_BRANCH=release-0.4.1 -FROM base AS vllm-clone +# Configure SSH access +RUN mkdir -p /root/.ssh \ + && echo "StrictHostKeyChecking no" >> /root/.ssh/config \ + && ssh-keyscan -t rsa github.com >> /root/.ssh/known_hosts -RUN mkdir -p /root/.ssh && \ - echo "StrictHostKeyChecking no" >> /root/.ssh/config && \ - ssh-keyscan -t rsa github.com >> /root/.ssh/known_hosts - -WORKDIR /vllm - -RUN --mount=type=secret,id=ssh_key,target=/root/.ssh/id_ed25519,mode=0600 \ - git clone -b ${VLLM_PRIVATE_BRANCH} git@github.com:aws-neuron/private-vllm-neuron.git . +# Clone vllm-neuronx repository +RUN --mount=type=secret,id=containers_github_ssh_key,target=/root/.ssh/id_ed25519,mode=0600 \ + git clone -b ${GITHUB_REPO_BRANCH} ${GITHUB_REPO} /opt/vllm FROM base AS repo @@ -193,9 +192,6 @@ RUN apt-get update \ && rm -rf /tmp/tmp* \ && apt-get clean -# Install VLLM from source -COPY --from=vllm-clone /vllm /opt/vllm - RUN PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP_REPO_KEY}@${NEURON_PIP_REPO}" || echo "https://${NEURON_PIP_REPO}") \ && ${PIP} install --no-cache-dir \ --index-url ${PIP_REPO_URL} \ @@ -221,9 +217,6 @@ RUN apt-get update \ && rm -rf /tmp/tmp* \ && apt-get clean -# Clone VLLM source before pip installations -RUN git clone -b "${VLLM_PUBLIC_BRANCH}" https://github.com/vllm-project/vllm-neuron.git /opt/vllm - RUN PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP_REPO_KEY}@${NEURON_PIP_REPO}" || echo "https://${NEURON_PIP_REPO}") \ && ${PIP} install --no-cache-dir \ --index-url ${PIP_REPO_URL} \ @@ -243,4 +236,4 @@ EXPOSE 8080 8081 ENTRYPOINT ["python", "/usr/local/bin/vllm_entrypoint.py"] CMD ["/bin/bash"] -HEALTHCHECK CMD curl --fail http://localhost:8080/ping || exit 1 \ No newline at end of file +HEALTHCHECK CMD curl --fail http://localhost:8080/ping || exit 1