diff --git a/.buildkite/image_build/image_build.yaml b/.buildkite/image_build/image_build.yaml index e0ef7d59242..5d6ce7309e8 100644 --- a/.buildkite/image_build/image_build.yaml +++ b/.buildkite/image_build/image_build.yaml @@ -98,3 +98,21 @@ steps: limit: 2 - exit_status: -10 # Agent was lost limit: 2 + + - label: ":docker: Build arm64 image" + key: arm64-image-build + depends_on: [] + source_file_dependencies: + - ".buildkite/image_build/image_build.yaml" + - ".buildkite/image_build/image_build_arm64.sh" + - "docker/Dockerfile" + commands: + - .buildkite/image_build/image_build_arm64.sh $REGISTRY $REPO $BUILDKITE_COMMIT + env: + DOCKER_BUILDKIT: "1" + retry: + automatic: + - exit_status: -1 # Agent was lost + limit: 2 + - exit_status: -10 # Agent was lost + limit: 2 diff --git a/.buildkite/image_build/image_build_arm64.sh b/.buildkite/image_build/image_build_arm64.sh new file mode 100755 index 00000000000..5baa55a1965 --- /dev/null +++ b/.buildkite/image_build/image_build_arm64.sh @@ -0,0 +1,37 @@ +#!/bin/bash +set -e + +if [[ $# -lt 3 ]]; then + echo "Usage: $0 " + exit 1 +fi + +REGISTRY=$1 +REPO=$2 +BUILDKITE_COMMIT=$3 + +# authenticate with AWS ECR +aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true + +# skip build if image already exists +if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64) ]]; then + echo "Image not found, proceeding with build..." +else + echo "Image found" + exit 0 +fi + +# build (Grace/GH200 is the arm64 GPU target; sm_90) +docker build --file docker/Dockerfile \ + --platform linux/arm64 \ + --build-arg max_jobs=16 \ + --build-arg nvcc_threads=4 \ + --build-arg torch_cuda_arch_list="9.0" \ + --build-arg USE_SCCACHE=1 \ + --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ + --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64 \ + --target test \ + --progress plain . + +# push +docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64 diff --git a/docker/Dockerfile b/docker/Dockerfile index 47b17fc2d10..06cdc0b667f 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -868,6 +868,7 @@ FROM vllm-base AS test ADD . /vllm-workspace/ ARG PYTHON_VERSION +ARG TARGETPLATFORM ARG PIP_INDEX_URL UV_INDEX_URL ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL @@ -907,6 +908,11 @@ RUN --mount=type=cache,target=/opt/uv/cache \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ else \ echo "Installing dev requirements..." \ + && if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \ + echo "Recompiling test requirements for arm64..." \ + && uv pip compile requirements/test/cuda.in -o requirements/test/cuda.txt --index-strategy unsafe-best-match \ + --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + fi \ && uv pip install --system -r requirements/dev.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ fi \ diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 1a71e3ee0e4..6c786491603 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -53,12 +53,12 @@ tritonclient>=2.51.0 grpcio==1.78.0 grpcio-reflection==1.78.0 -arctic-inference == 0.1.1 # Required for suffix decoding test +arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix decoding test numba == 0.65.0 # Required for N-gram speculative decoding numpy runai-model-streamer[s3,gcs,azure]==0.15.7 -fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage -instanttensor>=0.1.5 +fastsafetensors>=0.2.2; platform_machine == "x86_64" # 0.2.2 contains important fixes for multi-GPU mem usage +instanttensor>=0.1.5; platform_machine == "x86_64" pydantic>=2.12 # 2.11 leads to error on python 3.13 decord==0.6.0; platform_machine == "x86_64" # terratorch is temporarily disabled while PyPI has the `lightning` package