[LLM] Simplify the NIXL dependency on ray-llm images (ray-project#57706)

Signed-off-by: Kourosh Hakhamaneshi <kourosh@anyscale.com>

Author: kouroshHakha

ci/raydepsets/configs/rayllm.depsets.yaml

@@ -18,7 +18,6 @@ build_arg_sets:
     - --extra-index-url https://download.pytorch.org/whl/${CUDA_CODE}
   build_arg_sets:
     - cpu
-    - cu121
     - cu128
 
 depsets:

docker/ray-llm/Dockerfile

@@ -6,10 +6,6 @@ FROM "$BASE_IMAGE"
 COPY python/deplocks/llm/rayllm_*.lock ./
 
 ARG KVER="5.15.0-139-generic"
-ARG ROOT_DIR="/usr/local"
-ARG GDR_HOME="${ROOT_DIR}/gdrcopy"
-ARG UCX_HOME="${ROOT_DIR}/ucx"
-ARG NIXL_HOME="${ROOT_DIR}/nixl"
 
 RUN <<EOF
 #!/bin/bash
@@ -33,97 +29,9 @@ uv pip install --system --no-cache-dir --no-deps \
 # Export installed packages
 $HOME/anaconda3/bin/pip freeze > /home/ray/pip-freeze.txt
 
-# Begin NIXL installation
-
-mkdir -p "${ROOT_DIR}"
-
-CUDACXX=$(which nvcc)
-CUDA_HOME=$(dirname $(dirname ${CUDACXX}))
-
-TEMP_DIR="nixl_installer"
-mkdir -p "${TEMP_DIR}"
-
-sudo apt-get update
-# kmod needed by nvidia-installer, pkg-config needed by GDRCopy, librdmacm-dev needed by UCX
-sudo apt-get install -y kmod pkg-config librdmacm-dev cmake
-
-(
-    echo "Installing GDRCopy"
-    cd "${TEMP_DIR}"
-    [[ -d "/lib/modules/${KVER}" ]] || sudo apt-get install linux-headers-${KVER} -y
-    NV_DRIVER_VERSION="570.153.02"
-    wget "https://us.download.nvidia.com/XFree86/Linux-x86_64/${NV_DRIVER_VERSION}/NVIDIA-Linux-x86_64-${NV_DRIVER_VERSION}.run" -q
-    sh NVIDIA-Linux-x86_64-${NV_DRIVER_VERSION}.run -x
-    sudo NVIDIA-Linux-x86_64-${NV_DRIVER_VERSION}/nvidia-installer \
-        --silent \
-        --no-questions \
-        --no-install-compat32-libs \
-        --kernel-source-path="/lib/modules/${KVER}/build" \
-        --utility-prefix="/usr"
-
-    (
-        wget "https://github.com/NVIDIA/gdrcopy/archive/refs/tags/v2.5.tar.gz" -q
-        tar xzf v2.5.tar.gz; rm v2.5.tar.gz
-        cd gdrcopy-2.5
-        sudo make prefix=$GDR_HOME CUDA=$CUDA_HOME KVER=${KVER} all install
-    )
-
-    # Uninstall the driver, this driver might have conflict with the library
-    # version on host. Remove it from container.
-    sudo NVIDIA-Linux-x86_64-${NV_DRIVER_VERSION}/nvidia-installer \
-        --uninstall \
-        --silent \
-        --no-questions
-)
-
-UCX_VERSION="1.19.0"
-(
-    echo "Installing UCX ${UCX_VERSION}"
-    cd "${TEMP_DIR}"
-    wget "https://github.com/openucx/ucx/releases/download/v${UCX_VERSION}/ucx-${UCX_VERSION}.tar.gz" -q
-    tar xzf "ucx-${UCX_VERSION}.tar.gz"; rm "ucx-${UCX_VERSION}.tar.gz"
-    cd "ucx-${UCX_VERSION}"
-
-    # Additional options for Mellanox NICs, install by default
-    MLX_OPTS="--with-rdmacm \
-              --with-mlx5-dv \
-              --with-ib-hw-tm"
-
-    ./configure --prefix=${UCX_HOME}               \
-                --enable-shared                    \
-                --disable-static                   \
-                --disable-doxygen-doc              \
-                --enable-optimizations             \
-                --enable-cma                       \
-                --enable-devel-headers             \
-                --with-cuda=${CUDA_HOME}           \
-                --with-dm                          \
-                --with-gdrcopy=${GDR_HOME}         \
-                --with-verbs                       \
-                --enable-mt                        \
-                ${MLX_OPTS}
-    make -j
-    sudo make -j install-strip
-
-    sudo ldconfig
-)
-
-# Keep in sync with llm-requirements.txt
-NIXL_VERSION="0.6.0"
-(
-    echo "Installing NIXL ${NIXL_VERSION}"
-    # NIXL needs meson pybind11 ninja, but should have been included in requirements_*.txt
-    cd "${TEMP_DIR}"
-    wget "https://github.com/ai-dynamo/nixl/archive/refs/tags/${NIXL_VERSION}.tar.gz" -q
-    tar xzf "${NIXL_VERSION}.tar.gz"; rm "${NIXL_VERSION}.tar.gz"
-    cd "nixl-${NIXL_VERSION}"
-    meson setup build --prefix=${NIXL_HOME} -Ducx_path=${UCX_HOME}
-    cd build
-    ninja
-    sudo env "PATH=$PATH" ninja install
-)
-sudo rm -rf "${TEMP_DIR}"
+sudo apt-get update -y && sudo apt-get install -y kmod pkg-config librdmacm-dev cmake
 
+# Install DeepEP kernels
 EP_TEMP_DIR=$(pwd)/"ep_temp_dir"
 mkdir -p "${EP_TEMP_DIR}"
 
@@ -189,7 +97,3 @@ sudo rm -rf /var/lib/apt/lists/*
 sudo apt-get clean
 
 EOF
-
-ENV PATH="${PATH}:${UCX_HOME}/bin:${NIXL_HOME}/bin"
-ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${UCX_HOME}/lib:${NIXL_HOME}/lib/x86_64-linux-gnu"
-ENV NIXL_PLUGIN_DIR="${NIXL_HOME}/lib/x86_64-linux-gnu/plugins/"

python/deplocks/llm/ray_py311_cu121.lock

@@ -1710,15 +1710,15 @@ ninja==1.11.1.3 \
     #   -r python/requirements/llm/llm-requirements.txt
     #   vllm
     #   xgrammar
-nixl==0.6.0 \
-    --hash=sha256:1e703abfa56fb020831df8940539a4447205aa593632ace676cb2db93c8862ee \
-    --hash=sha256:4468619d1730aed076ad7de5aa206d6def4c61ae8a19067dbe31bc2846699e4a \
-    --hash=sha256:5368f36bc50ba9fb1aa1b6b7a9b083dfd19af0a45d1649142785ec9bee6b7919 \
-    --hash=sha256:6cff5aa014b80b3d17f364c21367dd055838ba6dcfaa4c50dc09b056276e8e14 \
-    --hash=sha256:ac733f706c80edf4fa9d3c2325cfb50671952b5aec5cf6a479c71e1ad175a2a0 \
-    --hash=sha256:b4a21d381cb78f92be695fb614b33fc8979a2d37f1ca6951121bd2fcccaf9daa \
-    --hash=sha256:b4dafa2b59d96c57ab7c036b6dcf41def4c28619415d6ff751190d94f321eef7 \
-    --hash=sha256:e5bc4b0a901a97e681b6f2b54a798cb8743c79bb5f4e4e4e68f06b913da2cfc9
+nixl==0.6.1 \
+    --hash=sha256:24e9e98a72839d762bedb8faca010c5878aa0b2d5624a1590d6a588aab1d223e \
+    --hash=sha256:2a9f29718e5dde20ee9e6e5fb25411d1950ab84733e0d4fceb8bb6ccf555a1e5 \
+    --hash=sha256:77eab96bef382bfb91b9d6222e5581e49b193fcf573b38dcaa7a296822a2894e \
+    --hash=sha256:7abbaccc88f0330d38e5344efa4a0768fe523e9a0083b785ea60da858d73b265 \
+    --hash=sha256:831affb62a6ff6199e41ffdccaab3430cb61bf3ca71e597ca214d2db26620955 \
+    --hash=sha256:8507c73d9bc044dd921edbef81ebae3e0750584a70a63ea90e5ade79233535d2 \
+    --hash=sha256:d28c348371045962b109d5ebf1ab054017fd9c89a6d9167902c62dc793465e2d \
+    --hash=sha256:f562139f23609336e5254b96e07b20b3298cca81ddc7549fa2da6dd788a80564
     # via
     #   -c python/deplocks/llm/rayllm_test_py311_cpu.lock
     #   -r python/requirements/llm/llm-requirements.txt