diff --git a/.buildkite/scripts/install-kv-connectors.sh b/.buildkite/scripts/install-kv-connectors.sh index ef47c3625f..8733a21347 100755 --- a/.buildkite/scripts/install-kv-connectors.sh +++ b/.buildkite/scripts/install-kv-connectors.sh @@ -9,9 +9,20 @@ if python3 -c "import torch; raise SystemExit(0 if torch.version.hip is not None exit 0 fi -REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}" +# Select requirements before optional packages can change the installed Torch. +CUDA_MAJOR=$(python3 -c 'import torch; print(torch.version.cuda.split(".", 1)[0])') +DEFAULT_REQUIREMENTS=/vllm-workspace/requirements/kv_connectors.txt +if [ "${CUDA_MAJOR}" = "12" ]; then + DEFAULT_REQUIREMENTS=/vllm-workspace/requirements/kv_connectors_cu12.txt +fi +REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-${DEFAULT_REQUIREMENTS}}" -uv pip install --system -r "${REQUIREMENTS_FILE}" +CONSTRAINTS_DIR=$(mktemp -d) +trap 'rm -rf "${CONSTRAINTS_DIR}"' EXIT +uv pip freeze --system > "${CONSTRAINTS_DIR}/installed.txt" +grep -E '^(torch|torchaudio|torchvision|triton|cuda-bindings|cuda-python|cuda-toolkit)==' \ + "${CONSTRAINTS_DIR}/installed.txt" > "${CONSTRAINTS_DIR}/constraints.txt" +uv pip install --system -c "${CONSTRAINTS_DIR}/constraints.txt" -r "${REQUIREMENTS_FILE}" KV_METADATA=$(python3 - <<'PY' import importlib.metadata as metadata @@ -56,5 +67,5 @@ PY # build first to avoid a clash. if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true - uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}" + uv pip install --system -c "${CONSTRAINTS_DIR}/constraints.txt" "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}" fi diff --git a/docker/Dockerfile b/docker/Dockerfile index 9661f593da..895968601f 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1264,7 +1264,8 @@ ENV UV_HTTP_TIMEOUT=500 ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0' ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list} RUN --mount=type=cache,target=/opt/uv/cache \ - --mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \ + --mount=type=bind,source=requirements,target=/tmp/vllm-requirements,ro \ + set -eu; \ CUDA_MAJOR="${CUDA_VERSION%%.*}"; \ CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-'); \ CUDA_HOME=/usr/local/cuda; \ @@ -1273,11 +1274,19 @@ RUN --mount=type=cache,target=/opt/uv/cache \ libcublas-dev-${CUDA_VERSION_DASH} \ libcusolver-dev-${CUDA_VERSION_DASH}"; \ if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \ - uv pip install --system -r /tmp/kv_connectors.txt --no-build || ( \ + KV_REQUIREMENTS=/tmp/vllm-requirements/kv_connectors.txt; \ + if [ "${CUDA_MAJOR}" = "12" ]; then \ + KV_REQUIREMENTS=/tmp/vllm-requirements/kv_connectors_cu12.txt; \ + fi; \ + # Preserve the stack used to compile vLLM when resolving optional dependencies. \ + uv pip freeze --system > /tmp/kv-installed.txt; \ + grep -E '^(torch|torchaudio|torchvision|triton|cuda-bindings|cuda-python|cuda-toolkit)==' \ + /tmp/kv-installed.txt > /tmp/kv-constraints.txt; \ + uv pip install --system -c /tmp/kv-constraints.txt -r "${KV_REQUIREMENTS}" --no-build || ( \ # if the above fails, install from source apt-get update -y && \ apt-get install -y --no-install-recommends --allow-change-held-packages ${BUILD_PKGS} && \ - uv pip install --system -r /tmp/kv_connectors.txt --no-build-isolation && \ + uv pip install --system -c /tmp/kv-constraints.txt -r "${KV_REQUIREMENTS}" --no-build-isolation && \ apt-get purge -y ${BUILD_PKGS} && \ # clean up -dev packages, keep runtime libraries rm -rf /var/lib/apt/lists/* \ @@ -1288,8 +1297,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \ MOONCAKE_VERSION=$(python3 -c "import importlib.metadata as m; print(m.version('mooncake-transfer-engine'))" 2>/dev/null || true) && \ if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then \ uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true; \ - uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"; \ + uv pip install --system -c /tmp/kv-constraints.txt "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"; \ fi; \ + rm -f /tmp/kv-installed.txt /tmp/kv-constraints.txt; \ fi # Mooncake registers GPU memory for RDMA via nvidia-peermem by default. Hosts diff --git a/requirements/kv_connectors_cu12.txt b/requirements/kv_connectors_cu12.txt new file mode 100644 index 0000000000..5ae5fba5ff --- /dev/null +++ b/requirements/kv_connectors_cu12.txt @@ -0,0 +1,5 @@ +# LMCache 0.5.x uses CUDA 13 dependencies; retain the CUDA 12 release here. +lmcache == 0.4.4 +cupy-cuda12x != 14.1.0 +nixl == 1.4.1 +mooncake-transfer-engine >= 0.3.12