| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # install_prerequisites.py |
| 4 | import argparse |
| 5 | import glob |
| 6 | import json |
| 7 | import os |
| 8 | import subprocess |
| 9 | import sys |
| 10 | import urllib.request |
| 11 | |
| 12 | # --- Configuration --- |
| 13 | WHEELS_CACHE_HOME = os.environ.get("WHEELS_CACHE_HOME", "/tmp/wheels_cache") |
| 14 | ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) |
| 15 | UCX_DIR = os.path.join("/tmp", "ucx_source") |
| 16 | NIXL_DIR = os.path.join("/tmp", "nixl_source") |
| 17 | UCX_INSTALL_DIR = os.path.join("/tmp", "ucx_install") |
| 18 | UCX_REPO_URL = "https://github.com/openucx/ucx.git" |
| 19 | NIXL_REPO_URL = "https://github.com/ai-dynamo/nixl.git" |
| 20 | |
| 21 | |
| 22 | # --- Helper Functions --- |
| 23 | def get_latest_nixl_version(): |
| 24 | """Helper function to get latest release version of NIXL.""" |
| 25 | try: |
| 26 | nixl_release_url = "https://api.github.com/repos/ai-dynamo/nixl/releases/latest" |
| 27 | with urllib.request.urlopen(nixl_release_url) as response: |
| 28 | data = json.load(response) |
| 29 | return data.get("tag_name", "0.7.0") |
| 30 | except Exception: |
| 31 | return "0.7.0" |
| 32 | |
| 33 | |
| 34 | NIXL_VERSION = os.environ.get("NIXL_VERSION", get_latest_nixl_version()) |
| 35 | |
| 36 | |
| 37 | def run_command(command, cwd=".", env=None): |
| 38 | """Helper function to run a shell command and check for errors.""" |
| 39 | print(f"--> Running command: {' '.join(command)} in '{cwd}'", flush=True) |
| 40 | subprocess.check_call(command, cwd=cwd, env=env) |
| 41 | |
| 42 | |
| 43 | def is_pip_package_installed(package_name): |
| 44 | """Checks if a package is installed via pip without raising an exception.""" |
| 45 | result = subprocess.run( |
| 46 | [sys.executable, "-m", "pip", "show", package_name], |
| 47 | stdout=subprocess.DEVNULL, |
| 48 | stderr=subprocess.DEVNULL, |
| 49 | ) |
| 50 | return result.returncode == 0 |
| 51 | |
| 52 | |
| 53 | def find_nixl_wheel_in_cache(cache_dir): |
| 54 | """Finds a nixl wheel file in the specified cache directory.""" |
| 55 | # The repaired wheel will have a 'manylinux' tag, but this glob still works. |
| 56 | search_pattern = os.path.join(cache_dir, f"nixl*{NIXL_VERSION}*.whl") |
| 57 | wheels = glob.glob(search_pattern) |
| 58 | if wheels: |
| 59 | # Sort to get the most recent/highest version if multiple exist |
| 60 | wheels.sort() |
| 61 | return wheels[-1] |
| 62 | return None |
| 63 | |
| 64 | |
| 65 | def install_system_dependencies(): |
| 66 | """Installs required system packages using apt-get if run as root.""" |
| 67 | if os.geteuid() != 0: |
| 68 | print("\n---", flush=True) |
| 69 | print( |
| 70 | "WARNING: Not running as root. \ |
| 71 | Skipping system dependency installation.", |
| 72 | flush=True, |
| 73 | ) |
| 74 | print( |
| 75 | "Please ensure the listed packages are installed on your system:", |
| 76 | flush=True, |
| 77 | ) |
| 78 | print( |
| 79 | " patchelf build-essential git cmake ninja-build \ |
| 80 | autotools-dev automake meson libtool libtool-bin", |
| 81 | flush=True, |
| 82 | ) |
| 83 | print("---\n", flush=True) |
| 84 | return |
| 85 | |
| 86 | print("--- Running as root. Installing system dependencies... ---", flush=True) |
| 87 | apt_packages = [ |
| 88 | "patchelf", # <-- Add patchelf here |
| 89 | "build-essential", |
| 90 | "git", |
| 91 | "cmake", |
| 92 | "ninja-build", |
| 93 | "autotools-dev", |
| 94 | "automake", |
| 95 | "meson", |
| 96 | "libtool", |
| 97 | "libtool-bin", |
| 98 | "pkg-config", |
| 99 | ] |
| 100 | run_command(["apt-get", "update"]) |
| 101 | run_command(["apt-get", "install", "-y"] + apt_packages) |
| 102 | print("--- System dependencies installed successfully. ---\n", flush=True) |
| 103 | |
| 104 | |
| 105 | def build_and_install_prerequisites(args): |
| 106 | """Builds UCX and NIXL from source, creating a self-contained wheel.""" |
| 107 | if not args.force_reinstall and is_pip_package_installed("nixl"): |
| 108 | print("--> NIXL is already installed. Nothing to do.", flush=True) |
| 109 | return |
| 110 | |
| 111 | cached_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME) |
| 112 | if not args.force_reinstall and cached_wheel: |
| 113 | print( |
| 114 | f"\n--> Found self-contained wheel: \ |
| 115 | {os.path.basename(cached_wheel)}.", |
| 116 | flush=True, |
| 117 | ) |
| 118 | print("--> Installing from cache, skipping all source builds.", flush=True) |
| 119 | install_command = [sys.executable, "-m", "pip", "install", cached_wheel] |
| 120 | run_command(install_command) |
| 121 | print("\n--- Installation from cache complete. ---", flush=True) |
| 122 | return |
| 123 | |
| 124 | print( |
| 125 | "\n--> No installed package or cached wheel found. \ |
| 126 | Starting full build process...", |
| 127 | flush=True, |
| 128 | ) |
| 129 | print("\n--> Installing auditwheel...", flush=True) |
| 130 | run_command([sys.executable, "-m", "pip", "install", "auditwheel"]) |
| 131 | install_system_dependencies() |
| 132 | ucx_install_path = os.path.abspath(UCX_INSTALL_DIR) |
| 133 | print(f"--> Using wheel cache directory: {WHEELS_CACHE_HOME}", flush=True) |
| 134 | os.makedirs(WHEELS_CACHE_HOME, exist_ok=True) |
| 135 | |
| 136 | # -- Step 1: Build UCX from source -- |
| 137 | print("\n[1/3] Configuring and building UCX from source...", flush=True) |
| 138 | if not os.path.exists(UCX_DIR): |
| 139 | run_command(["git", "clone", UCX_REPO_URL, UCX_DIR]) |
| 140 | ucx_source_path = os.path.abspath(UCX_DIR) |
| 141 | run_command(["git", "checkout", "v1.19.x"], cwd=ucx_source_path) |
| 142 | run_command(["./autogen.sh"], cwd=ucx_source_path) |
| 143 | configure_command = [ |
| 144 | "./configure", |
| 145 | f"--prefix={ucx_install_path}", |
| 146 | "--enable-shared", |
| 147 | "--disable-static", |
| 148 | "--disable-doxygen-doc", |
| 149 | "--enable-optimizations", |
| 150 | "--enable-cma", |
| 151 | "--enable-devel-headers", |
| 152 | "--with-verbs", |
| 153 | "--enable-mt", |
| 154 | "--with-ze=no", |
| 155 | ] |
| 156 | run_command(configure_command, cwd=ucx_source_path) |
| 157 | run_command(["make", "-j", str(os.cpu_count() or 1)], cwd=ucx_source_path) |
| 158 | run_command(["make", "install"], cwd=ucx_source_path) |
| 159 | print("--- UCX build and install complete ---", flush=True) |
| 160 | |
| 161 | # -- Step 2: Build NIXL wheel from source -- |
| 162 | print("\n[2/3] Building NIXL wheel from source...", flush=True) |
| 163 | if not os.path.exists(NIXL_DIR): |
| 164 | run_command(["git", "clone", NIXL_REPO_URL, NIXL_DIR]) |
| 165 | else: |
| 166 | run_command(["git", "fetch", "--tags"], cwd=NIXL_DIR) |
| 167 | run_command(["git", "checkout", NIXL_VERSION], cwd=NIXL_DIR) |
| 168 | print(f"--> Checked out NIXL version: {NIXL_VERSION}", flush=True) |
| 169 | |
| 170 | build_env = os.environ.copy() |
| 171 | build_env["PKG_CONFIG_PATH"] = os.path.join(ucx_install_path, "lib", "pkgconfig") |
| 172 | ucx_lib_path = os.path.join(ucx_install_path, "lib") |
| 173 | ucx_plugin_path = os.path.join(ucx_lib_path, "ucx") |
| 174 | existing_ld_path = os.environ.get("LD_LIBRARY_PATH", "") |
| 175 | build_env["LD_LIBRARY_PATH"] = ( |
| 176 | f"{ucx_lib_path}:{ucx_plugin_path}:{existing_ld_path}".strip(":") |
| 177 | ) |
| 178 | build_env["LDFLAGS"] = "-Wl,-rpath,$ORIGIN" |
| 179 | print(f"--> Using LD_LIBRARY_PATH: {build_env['LD_LIBRARY_PATH']}", flush=True) |
| 180 | |
| 181 | temp_wheel_dir = os.path.join(ROOT_DIR, "temp_wheelhouse") |
| 182 | run_command( |
| 183 | [ |
| 184 | sys.executable, |
| 185 | "-m", |
| 186 | "pip", |
| 187 | "wheel", |
| 188 | ".", |
| 189 | "--no-deps", |
| 190 | f"--wheel-dir={temp_wheel_dir}", |
| 191 | ], |
| 192 | cwd=os.path.abspath(NIXL_DIR), |
| 193 | env=build_env, |
| 194 | ) |
| 195 | |
| 196 | # -- Step 3: Repair the wheel by copying UCX libraries -- |
| 197 | print("\n[3/3] Repairing NIXL wheel to include UCX libraries...", flush=True) |
| 198 | unrepaired_wheel = find_nixl_wheel_in_cache(temp_wheel_dir) |
| 199 | if not unrepaired_wheel: |
| 200 | raise RuntimeError("Failed to find the NIXL wheel after building it.") |
| 201 | |
| 202 | # We tell auditwheel to ignore the plugin that mesonpy already handled. |
| 203 | auditwheel_command = [ |
| 204 | "auditwheel", |
| 205 | "repair", |
| 206 | "--exclude", |
| 207 | "libplugin_UCX.so", # <-- Exclude because mesonpy already includes it |
| 208 | unrepaired_wheel, |
| 209 | f"--wheel-dir={WHEELS_CACHE_HOME}", |
| 210 | ] |
| 211 | run_command(auditwheel_command, env=build_env) |
| 212 | |
| 213 | # --- CLEANUP --- |
| 214 | # No more temporary files to remove, just the temp wheelhouse |
| 215 | run_command(["rm", "-rf", temp_wheel_dir]) |
| 216 | # --- END CLEANUP --- |
| 217 | |
| 218 | newly_built_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME) |
| 219 | if not newly_built_wheel: |
| 220 | raise RuntimeError("Failed to find the repaired NIXL wheel.") |
| 221 | |
| 222 | print( |
| 223 | f"--> Successfully built self-contained wheel: \ |
| 224 | {os.path.basename(newly_built_wheel)}. Now installing...", |
| 225 | flush=True, |
| 226 | ) |
| 227 | install_command = [ |
| 228 | sys.executable, |
| 229 | "-m", |
| 230 | "pip", |
| 231 | "install", |
| 232 | "--no-deps", # w/o "no-deps", it will install cuda-torch |
| 233 | newly_built_wheel, |
| 234 | ] |
| 235 | if args.force_reinstall: |
| 236 | install_command.insert(-1, "--force-reinstall") |
| 237 | |
| 238 | run_command(install_command) |
| 239 | print("--- NIXL installation complete ---", flush=True) |
| 240 | |
| 241 | |
| 242 | if __name__ == "__main__": |
| 243 | parser = argparse.ArgumentParser( |
| 244 | description="Build and install UCX and NIXL dependencies." |
| 245 | ) |
| 246 | parser.add_argument( |
| 247 | "--force-reinstall", |
| 248 | action="store_true", |
| 249 | help="Force rebuild and reinstall of UCX and NIXL \ |
| 250 | even if they are already installed.", |
| 251 | ) |
| 252 | args = parser.parse_args() |
| 253 | build_and_install_prerequisites(args) |