A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[CI Failure] Fix pre-commit issue for install_nixl_from_source_ubuntu.py (#26424)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Oct 8, 2025 at 10:55 UTC
c1b06fc18282db18f231fb088c9113f2942f82ee
1 file changed
+95
-81
tools/install_nixl_from_source_ubuntu.py
+95
-81
@@ -10,15 +10,15 @@ import sys
10
# --- Configuration ---
11
WHEELS_CACHE_HOME = os.environ.get("WHEELS_CACHE_HOME", "/tmp/wheels_cache")
12
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
13
-UCX_DIR = os.path.join('/tmp', 'ucx_source')
14
-NIXL_DIR = os.path.join('/tmp', 'nixl_source')
15
-UCX_INSTALL_DIR = os.path.join('/tmp', 'ucx_install')
16
-UCX_REPO_URL = 'https://github.com/openucx/ucx.git'
17
-NIXL_REPO_URL = 'https://github.com/ai-dynamo/nixl.git'
13
+UCX_DIR = os.path.join("/tmp", "ucx_source")
14
+NIXL_DIR = os.path.join("/tmp", "nixl_source")
15
+UCX_INSTALL_DIR = os.path.join("/tmp", "ucx_install")
16
+UCX_REPO_URL = "https://github.com/openucx/ucx.git"
17
+NIXL_REPO_URL = "https://github.com/ai-dynamo/nixl.git"
18
19
20
# --- Helper Functions ---
21
-def run_command(command, cwd='.', env=None):
21
+def run_command(command, cwd=".", env=None):
22
"""Helper function to run a shell command and check for errors."""
23
print(f"--> Running command: {' '.join(command)} in '{cwd}'", flush=True)
24
subprocess.check_call(command, cwd=cwd, env=env)
@@ -27,9 +27,10 @@ def run_command(command, cwd='.', env=None):
27
def is_pip_package_installed(package_name):
28
"""Checks if a package is installed via pip without raising an exception."""
29
result = subprocess.run(
30
- [sys.executable, '-m', 'pip', 'show', package_name],
30
+ [sys.executable, "-m", "pip", "show", package_name],
31
stdout=subprocess.DEVNULL,
32
- stderr=subprocess.DEVNULL)
32
+ stderr=subprocess.DEVNULL,
33
+ )
34
return result.returncode == 0
35
36
@@ -49,20 +50,24 @@ def install_system_dependencies():
50
"""Installs required system packages using apt-get if run as root."""
51
if os.geteuid() != 0:
52
print("\n---", flush=True)
52
- print("WARNING: Not running as root. \
53
+ print(
54
+ "WARNING: Not running as root. \
55
Skipping system dependency installation.",
54
- flush=True)
56
+ flush=True,
57
+ )
58
print(
59
"Please ensure the listed packages are installed on your system:",
57
- flush=True)
58
- print(" patchelf build-essential git cmake ninja-build \
60
+ flush=True,
61
+ )
62
+ print(
63
+ " patchelf build-essential git cmake ninja-build \
64
autotools-dev automake meson libtool libtool-bin",
60
- flush=True)
65
+ flush=True,
66
+ )
67
print("---\n", flush=True)
68
return
69
64
- print("--- Running as root. Installing system dependencies... ---",
65
- flush=True)
70
+ print("--- Running as root. Installing system dependencies... ---", flush=True)
71
apt_packages = [
72
"patchelf", # <-- Add patchelf here
73
"build-essential",
@@ -73,39 +78,40 @@ def install_system_dependencies():
78
"automake",
79
"meson",
80
"libtool",
76
- "libtool-bin"
81
+ "libtool-bin",
82
]
78
- run_command(['apt-get', 'update'])
79
- run_command(['apt-get', 'install', '-y'] + apt_packages)
83
+ run_command(["apt-get", "update"])
84
+ run_command(["apt-get", "install", "-y"] + apt_packages)
85
print("--- System dependencies installed successfully. ---\n", flush=True)
86
87
88
def build_and_install_prerequisites(args):
89
"""Builds UCX and NIXL from source, creating a self-contained wheel."""
90
86
- if not args.force_reinstall and is_pip_package_installed('nixl'):
91
+ if not args.force_reinstall and is_pip_package_installed("nixl"):
92
print("--> NIXL is already installed. Nothing to do.", flush=True)
93
return
94
95
cached_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME)
96
if not args.force_reinstall and cached_wheel:
92
- print(f"\n--> Found self-contained wheel: \
97
+ print(
98
+ f"\n--> Found self-contained wheel: \
99
{os.path.basename(cached_wheel)}.",
94
- flush=True)
95
- print("--> Installing from cache, skipping all source builds.",
96
- flush=True)
97
- install_command = [
98
- sys.executable, '-m', 'pip', 'install', cached_wheel
99
- ]
100
+ flush=True,
101
+ )
102
+ print("--> Installing from cache, skipping all source builds.", flush=True)
103
+ install_command = [sys.executable, "-m", "pip", "install", cached_wheel]
104
run_command(install_command)
105
print("\n--- Installation from cache complete. ---", flush=True)
106
return
107
104
- print("\n--> No installed package or cached wheel found. \
108
+ print(
109
+ "\n--> No installed package or cached wheel found. \
110
Starting full build process...",
106
- flush=True)
111
+ flush=True,
112
+ )
113
print("\n--> Installing auditwheel...", flush=True)
108
- run_command([sys.executable, '-m', 'pip', 'install', 'auditwheel'])
114
+ run_command([sys.executable, "-m", "pip", "install", "auditwheel"])
115
install_system_dependencies()
116
ucx_install_path = os.path.abspath(UCX_INSTALL_DIR)
117
print(f"--> Using wheel cache directory: {WHEELS_CACHE_HOME}", flush=True)
@@ -114,86 +120,91 @@ def build_and_install_prerequisites(args):
120
# -- Step 1: Build UCX from source --
121
print("\n[1/3] Configuring and building UCX from source...", flush=True)
122
if not os.path.exists(UCX_DIR):
117
- run_command(['git', 'clone', UCX_REPO_URL, UCX_DIR])
123
+ run_command(["git", "clone", UCX_REPO_URL, UCX_DIR])
124
ucx_source_path = os.path.abspath(UCX_DIR)
119
- run_command(['git', 'checkout', 'v1.19.x'], cwd=ucx_source_path)
120
- run_command(['./autogen.sh'], cwd=ucx_source_path)
125
+ run_command(["git", "checkout", "v1.19.x"], cwd=ucx_source_path)
126
+ run_command(["./autogen.sh"], cwd=ucx_source_path)
127
configure_command = [
122
- './configure',
123
- f'--prefix={ucx_install_path}',
124
- '--enable-shared',
125
- '--disable-static',
126
- '--disable-doxygen-doc',
127
- '--enable-optimizations',
128
- '--enable-cma',
129
- '--enable-devel-headers',
130
- '--with-verbs',
131
- '--enable-mt',
128
+ "./configure",
129
+ f"--prefix={ucx_install_path}",
130
+ "--enable-shared",
131
+ "--disable-static",
132
+ "--disable-doxygen-doc",
133
+ "--enable-optimizations",
134
+ "--enable-cma",
135
+ "--enable-devel-headers",
136
+ "--with-verbs",
137
+ "--enable-mt",
138
]
139
run_command(configure_command, cwd=ucx_source_path)
134
- run_command(['make', '-j', str(os.cpu_count() or 1)], cwd=ucx_source_path)
135
- run_command(['make', 'install'], cwd=ucx_source_path)
140
+ run_command(["make", "-j", str(os.cpu_count() or 1)], cwd=ucx_source_path)
141
+ run_command(["make", "install"], cwd=ucx_source_path)
142
print("--- UCX build and install complete ---", flush=True)
143
144
# -- Step 2: Build NIXL wheel from source --
145
print("\n[2/3] Building NIXL wheel from source...", flush=True)
146
if not os.path.exists(NIXL_DIR):
141
- run_command(['git', 'clone', NIXL_REPO_URL, NIXL_DIR])
147
+ run_command(["git", "clone", NIXL_REPO_URL, NIXL_DIR])
148
149
build_env = os.environ.copy()
144
- build_env['PKG_CONFIG_PATH'] = os.path.join(ucx_install_path, 'lib',
145
- 'pkgconfig')
146
- ucx_lib_path = os.path.join(ucx_install_path, 'lib')
147
- ucx_plugin_path = os.path.join(ucx_lib_path, 'ucx')
148
- existing_ld_path = os.environ.get('LD_LIBRARY_PATH', '')
149
- build_env['LD_LIBRARY_PATH'] = \
150
- f"{ucx_lib_path}:{ucx_plugin_path}:{existing_ld_path}".strip(':')
151
- print(f"--> Using LD_LIBRARY_PATH: {build_env['LD_LIBRARY_PATH']}",
152
- flush=True)
153
-
154
- temp_wheel_dir = os.path.join(ROOT_DIR, 'temp_wheelhouse')
155
- run_command([
156
- sys.executable, '-m', 'pip', 'wheel', '.', '--no-deps',
157
- f'--wheel-dir={temp_wheel_dir}'
158
- ],
159
- cwd=os.path.abspath(NIXL_DIR),
160
- env=build_env)
150
+ build_env["PKG_CONFIG_PATH"] = os.path.join(ucx_install_path, "lib", "pkgconfig")
151
+ ucx_lib_path = os.path.join(ucx_install_path, "lib")
152
+ ucx_plugin_path = os.path.join(ucx_lib_path, "ucx")
153
+ existing_ld_path = os.environ.get("LD_LIBRARY_PATH", "")
154
+ build_env["LD_LIBRARY_PATH"] = (
155
+ f"{ucx_lib_path}:{ucx_plugin_path}:{existing_ld_path}".strip(":")
156
+ )
157
+ print(f"--> Using LD_LIBRARY_PATH: {build_env['LD_LIBRARY_PATH']}", flush=True)
158
+
159
+ temp_wheel_dir = os.path.join(ROOT_DIR, "temp_wheelhouse")
160
+ run_command(
161
+ [
162
+ sys.executable,
163
+ "-m",
164
+ "pip",
165
+ "wheel",
166
+ ".",
167
+ "--no-deps",
168
+ f"--wheel-dir={temp_wheel_dir}",
169
+ ],
170
+ cwd=os.path.abspath(NIXL_DIR),
171
+ env=build_env,
172
+ )
173
174
# -- Step 3: Repair the wheel by copying UCX libraries --
163
- print("\n[3/3] Repairing NIXL wheel to include UCX libraries...",
164
- flush=True)
175
+ print("\n[3/3] Repairing NIXL wheel to include UCX libraries...", flush=True)
176
unrepaired_wheel = find_nixl_wheel_in_cache(temp_wheel_dir)
177
if not unrepaired_wheel:
178
raise RuntimeError("Failed to find the NIXL wheel after building it.")
179
180
# We tell auditwheel to ignore the plugin that mesonpy already handled.
181
auditwheel_command = [
171
- 'auditwheel',
172
- 'repair',
173
- '--exclude',
174
- 'libplugin_UCX.so', # <-- Exclude because mesonpy already includes it
182
+ "auditwheel",
183
+ "repair",
184
+ "--exclude",
185
+ "libplugin_UCX.so", # <-- Exclude because mesonpy already includes it
186
unrepaired_wheel,
176
- f'--wheel-dir={WHEELS_CACHE_HOME}'
187
+ f"--wheel-dir={WHEELS_CACHE_HOME}",
188
]
189
run_command(auditwheel_command, env=build_env)
190
191
# --- CLEANUP ---
192
# No more temporary files to remove, just the temp wheelhouse
182
- run_command(['rm', '-rf', temp_wheel_dir])
193
+ run_command(["rm", "-rf", temp_wheel_dir])
194
# --- END CLEANUP ---
195
196
newly_built_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME)
197
if not newly_built_wheel:
198
raise RuntimeError("Failed to find the repaired NIXL wheel.")
199
189
- print(f"--> Successfully built self-contained wheel: \
200
+ print(
201
+ f"--> Successfully built self-contained wheel: \
202
{os.path.basename(newly_built_wheel)}. Now installing...",
191
- flush=True)
192
- install_command = [
193
- sys.executable, '-m', 'pip', 'install', newly_built_wheel
194
- ]
203
+ flush=True,
204
+ )
205
+ install_command = [sys.executable, "-m", "pip", "install", newly_built_wheel]
206
if args.force_reinstall:
196
- install_command.insert(-1, '--force-reinstall')
207
+ install_command.insert(-1, "--force-reinstall")
208
209
run_command(install_command)
210
print("--- NIXL installation complete ---", flush=True)
@@ -201,10 +212,13 @@ def build_and_install_prerequisites(args):
212
213
if __name__ == "__main__":
214
parser = argparse.ArgumentParser(
204
- description="Build and install UCX and NIXL dependencies.")
205
- parser.add_argument('--force-reinstall',
206
- action='store_true',
207
- help='Force rebuild and reinstall of UCX and NIXL \
208
- even if they are already installed.')
215
+ description="Build and install UCX and NIXL dependencies."
216
+ )
217
+ parser.add_argument(
218
+ "--force-reinstall",
219
+ action="store_true",
220
+ help="Force rebuild and reinstall of UCX and NIXL \
221
+ even if they are already installed.",
222
+ )
223
args = parser.parse_args()
224
build_and_install_prerequisites(args)