执行摘要

• 一句话：添加 arm64 GPU 镜像的 CI 构建流程
• 推荐动作：该 PR 是标准的基础设施扩展，值得关注其平台标记模式和 Dockerfile 条件编译技巧。建议在后续 arm64 测试增强中参考本 PR 的依赖管理方式。

功能与动机

需要 arm64 CI 镜像以在 Grace/GH200 硬件上运行测试，从而覆盖 arm64 架构的 CI 验证。PR body 明确说明“Add arm64 image to CI”。

实现拆解

1. 新增构建脚本：在 .buildkite/image_build/image_build_arm64.sh 中实现 arm64 Docker 镜像的构建与推送逻辑，使用 --platform linux/arm64 并指定 torch_cuda_arch_list="9.0" 针对 Grace/GH200。
2. 配置流水线步骤：在 .buildkite/image_build/image_build.yaml 中添加名为 :docker: Build arm64 image 的新步骤，标记为 optional: true，并声明源文件依赖以仅在这些文件变更时触发。
3. 标记 x86_64 独占依赖：在 requirements/test/cuda.in 中为 arctic-inference、fastsafetensors、instanttensor 等无 arm64 wheel 的包追加 platform_machine == "x86_64" 标记，使其在 arm64 上跳过安装。
4. Dockerfile 条件编译：在 docker/Dockerfile 的 test 阶段，若 TARGETPLATFORM 为 linux/arm64，则使用 uv pip compile 重新生成针对 arm64 的 requirements/test/cuda.txt，避免预编译的 x86_64 版本导致安装失败。

关键文件：

• .buildkite/image_build/image_build_arm64.sh（模块 构建脚本；类别 other；类型 core-logic）: 新增的 arm64 镜像构建脚本，封装了 docker build 与 push 逻辑，是整个 PR 的核心交付物。
• .buildkite/image_build/image_build.yaml（模块 构建配置；类别 config；类型 configuration）: 流水线配置，新增 arm64 构建步骤，声明源文件依赖，是触发构建的入口。
• requirements/test/cuda.in（模块 测试依赖；类别 test；类型 test-coverage）: 为 arm64 兼容性调整依赖标记，避免在 arm64 上安装无 wheel 的包。
• docker/Dockerfile（模块 Dockerfile；类别 infra；类型 infrastructure）: 在 arm64 构建时重新编译测试需求，解决预编译需求文件与平台不匹配问题。

关键符号：未识别

关键源码片段

.buildkite/image_build/image_build_arm64.sh

新增的 arm64 镜像构建脚本，封装了 docker build 与 push 逻辑，是整个 PR 的核心交付物。

#!/bin/bash
set -e
​
if [[ $# -lt 3 ]]; then
    echo "Usage: $0 <registry> <repo> <commit>"
    exit 1
fi
​
REGISTRY=$1
REPO=$2
BUILDKITE_COMMIT=$3
​
# 认证：若失败，|| true 会阻止脚本退出（gemini-code-assist 建议移除 || true）
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
​
# 检查镜像是否存在：若不存在，$(...) 在 set -e 下会触发 exit（gemini-code-assist 建议改用 if ! docker manifest inspect）
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
​
# 构建 arm64 镜像，目标架构 sm_90（Grace/GH200）
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 .
​
# 推送镜像
docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64

requirements/test/cuda.in

为 arm64 兼容性调整依赖标记，避免在 arm64 上安装无 wheel 的包。

# 原版本（所有平台安装）
# arctic-inference == 0.1.1
# fastsafetensors>=0.2.2
# instanttensor>=0.1.5
​
# 修改后：仅 x86_64 安装，arm64 跳过
arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix decoding test
fastsafetensors>=0.2.2; platform_machine == "x86_64" # C++ 扩展，无 arm64 wheel
instanttensor>=0.1.5; platform_machine == "x86_64" # C++/CUDA 扩展，无 arm64 wheel

docker/Dockerfile

在 arm64 构建时重新编译测试需求，解决预编译需求文件与平台不匹配问题。

# 在 test 阶段的 RUN 命令中插入条件块
RUN --mount=type=cache,target=/opt/uv/cache \
    if [ -n "$PYTORCH_CUDA_INDEX_BASE_URL" ]; then \
        # 原有 CUDA nightlies 安装逻辑 \
    else \
        echo "Installing dev requirements..." \
        # 新增：arm64 下重新编译需求文件，确保平台标记生效
        && 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 \
        # ... \
    fi

评论区精华

gemini-code-assist 审查指出了两个关键问题：

• 认证错误处理：docker login 后的 || true 会隐藏认证失败，建议移除以尽早失败。

• 镜像检查逻辑：$(docker manifest inspect ...) 在镜像不存在时因 set -e 导致脚本退出，建议使用 if ! docker manifest inspect 模式。
  这些建议在合并时未采纳，脚本保留原样。

• 构建脚本中 docker login 错误处理 (correctness): 未采纳该建议，合并时仍保留 || true。

• 构建脚本中镜像检查逻辑 (correctness): 未采纳该建议，合并时仍保留原有写法。

风险与影响

• 风险：

  1. 脚本健壮性：构建脚本未采纳审查建议，存在因认证失败或镜像检查逻辑浪费构建时间的风险；但步骤为 optional: true，失败不会阻塞流水线。
  2. 依赖标记遗漏：requirements/test/cuda.in 中仅标记了已知无 arm64 支持的包，可能遗漏其他 x86_64 独占依赖，导致 arm64 构建时安装失败。
  3. Dockerfile 条件编译：TARGETPLATFORM 依赖 Docker BuildKit，若环境未启用 BuildKit 可能导致条件不生效。
     - 影响：用户影响：无直接影响。
     系统影响：CI 流水线新增可选构建步骤，仅在特定文件变更或手动触发时运行，不增加普通 PR 等待时间。
     团队影响：为 arm64 CI 测试奠定基础，未来可在此基础上添加 arm64 测试步骤。

• 风险标记：构建脚本可能因认证或镜像检查问题浪费资源, arm64 依赖标记可能遗漏其他独占包

关联脉络

• PR #43709 [CI] Soft-fail AMD entrypoints mirror tests: 同为 CI 基础设施调整，涉及 Buildkite 配置，反映 CI 维护趋势。