[CI] Optimize nixl dependency installation script (#28950)

Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
2026-06-23 01:58:22 +08:00
committed by GitHub
parent 99c18cceec
commit 43b1fb95f6
+15 -10
View File
@@ -443,18 +443,15 @@ stabilize_flashinfer_jit_paths() {
install_extra_deps() {
MOONCAKE_VERSION="0.3.11.post1"
NIXL_VERSION="1.3.0"
NIXL_PKG="nixl==${NIXL_VERSION}"
if [ "$CU_MAJOR" = "13" ]; then
MOONCAKE_PKG="mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"
MOONCAKE_STALE_PKG="mooncake-transfer-engine"
NIXL_BIN_PKG="nixl-cu13==${NIXL_VERSION}"
NIXL_STALE_BIN_PKG="nixl-cu12"
NIXL_BIN_NAME="nixl-cu13"
EXTRA_NVIDIA_SPECS="nvidia-cuda-nvrtc"
else
MOONCAKE_PKG="mooncake-transfer-engine==${MOONCAKE_VERSION}"
MOONCAKE_STALE_PKG="mooncake-transfer-engine-cuda13"
NIXL_BIN_PKG="nixl-cu12==${NIXL_VERSION}"
NIXL_STALE_BIN_PKG="nixl-cu13"
NIXL_BIN_NAME="nixl-cu12"
EXTRA_NVIDIA_SPECS="nvidia-cuda-nvrtc-cu12"
fi
# Both variants own the same mooncake/ package files and bin/ scripts
@@ -466,12 +463,20 @@ install_extra_deps() {
$PIP_UNINSTALL_CMD ${MOONCAKE_STALE_PKG} $PIP_UNINSTALL_SUFFIX || true
$PIP_CMD install ${MOONCAKE_PKG} --force-reinstall --no-deps $PIP_INSTALL_SUFFIX
fi
# Best-effort NIXL install for decode-radix disaggregation coverage.
if pip show ${NIXL_STALE_BIN_PKG} >/dev/null 2>&1; then
$PIP_UNINSTALL_CMD ${NIXL_STALE_BIN_PKG} $PIP_UNINSTALL_SUFFIX || true
$PIP_CMD install ${NIXL_PKG} ${NIXL_BIN_PKG} --force-reinstall --no-deps $PIP_INSTALL_SUFFIX
$PIP_CMD install ${MOONCAKE_PKG} ${EXTRA_NVIDIA_SPECS} py-spy scipy huggingface_hub[hf_xet] pytest $PIP_INSTALL_SUFFIX
NIXL_INSTALLED=$(pip show nixl 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
NIXL_BIN_INSTALLED=$(pip show "${NIXL_BIN_NAME}" 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
if [ "$NIXL_INSTALLED" = "$NIXL_VERSION" ] && [ "$NIXL_BIN_INSTALLED" = "$NIXL_VERSION" ]; then
echo "nixl==${NIXL_VERSION} and ${NIXL_BIN_NAME}==${NIXL_VERSION} already installed, keeping them"
else
echo "nixl mismatch (meta: ${NIXL_INSTALLED:-none}, ${NIXL_BIN_NAME}: ${NIXL_BIN_INSTALLED:-none}, required: ${NIXL_VERSION}); installing"
# Meta stub owns the nixl import path; install only the CUDA binary for
# this runner's torch CUDA major. --no-deps avoids pulling the other CUDA
# variant; leave any other variant already on the runner image untouched.
$PIP_CMD install "nixl==${NIXL_VERSION}" "${NIXL_BIN_NAME}==${NIXL_VERSION}" \
--no-deps --force-reinstall $PIP_INSTALL_SUFFIX
fi
$PIP_CMD install ${MOONCAKE_PKG} ${NIXL_PKG} ${NIXL_BIN_PKG} ${EXTRA_NVIDIA_SPECS} py-spy scipy huggingface_hub[hf_xet] pytest $PIP_INSTALL_SUFFIX
if [ "$IS_BLACKWELL" != "1" ]; then
git clone --branch v0.5 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git