[CI Failure] Fix pre-commit issue for install_nixl_from_source_ubuntu.py (#26424)

Signed-off-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Michael Goin
2025-10-08 10:55:43 -04:00
committed by GitHub
parent 241b4cfe66
commit c1b06fc182

View File

@ -10,15 +10,15 @@ import sys
# --- Configuration --- # --- Configuration ---
WHEELS_CACHE_HOME = os.environ.get("WHEELS_CACHE_HOME", "/tmp/wheels_cache") WHEELS_CACHE_HOME = os.environ.get("WHEELS_CACHE_HOME", "/tmp/wheels_cache")
ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
UCX_DIR = os.path.join('/tmp', 'ucx_source') UCX_DIR = os.path.join("/tmp", "ucx_source")
NIXL_DIR = os.path.join('/tmp', 'nixl_source') NIXL_DIR = os.path.join("/tmp", "nixl_source")
UCX_INSTALL_DIR = os.path.join('/tmp', 'ucx_install') UCX_INSTALL_DIR = os.path.join("/tmp", "ucx_install")
UCX_REPO_URL = 'https://github.com/openucx/ucx.git' UCX_REPO_URL = "https://github.com/openucx/ucx.git"
NIXL_REPO_URL = 'https://github.com/ai-dynamo/nixl.git' NIXL_REPO_URL = "https://github.com/ai-dynamo/nixl.git"
# --- Helper Functions --- # --- Helper Functions ---
def run_command(command, cwd='.', env=None): def run_command(command, cwd=".", env=None):
"""Helper function to run a shell command and check for errors.""" """Helper function to run a shell command and check for errors."""
print(f"--> Running command: {' '.join(command)} in '{cwd}'", flush=True) print(f"--> Running command: {' '.join(command)} in '{cwd}'", flush=True)
subprocess.check_call(command, cwd=cwd, env=env) subprocess.check_call(command, cwd=cwd, env=env)
@ -27,9 +27,10 @@ def run_command(command, cwd='.', env=None):
def is_pip_package_installed(package_name): def is_pip_package_installed(package_name):
"""Checks if a package is installed via pip without raising an exception.""" """Checks if a package is installed via pip without raising an exception."""
result = subprocess.run( result = subprocess.run(
[sys.executable, '-m', 'pip', 'show', package_name], [sys.executable, "-m", "pip", "show", package_name],
stdout=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL) stderr=subprocess.DEVNULL,
)
return result.returncode == 0 return result.returncode == 0
@ -49,20 +50,24 @@ def install_system_dependencies():
"""Installs required system packages using apt-get if run as root.""" """Installs required system packages using apt-get if run as root."""
if os.geteuid() != 0: if os.geteuid() != 0:
print("\n---", flush=True) print("\n---", flush=True)
print("WARNING: Not running as root. \ print(
"WARNING: Not running as root. \
Skipping system dependency installation.", Skipping system dependency installation.",
flush=True) flush=True,
)
print( print(
"Please ensure the listed packages are installed on your system:", "Please ensure the listed packages are installed on your system:",
flush=True) flush=True,
print(" patchelf build-essential git cmake ninja-build \ )
print(
" patchelf build-essential git cmake ninja-build \
autotools-dev automake meson libtool libtool-bin", autotools-dev automake meson libtool libtool-bin",
flush=True) flush=True,
)
print("---\n", flush=True) print("---\n", flush=True)
return return
print("--- Running as root. Installing system dependencies... ---", print("--- Running as root. Installing system dependencies... ---", flush=True)
flush=True)
apt_packages = [ apt_packages = [
"patchelf", # <-- Add patchelf here "patchelf", # <-- Add patchelf here
"build-essential", "build-essential",
@ -73,39 +78,40 @@ def install_system_dependencies():
"automake", "automake",
"meson", "meson",
"libtool", "libtool",
"libtool-bin" "libtool-bin",
] ]
run_command(['apt-get', 'update']) run_command(["apt-get", "update"])
run_command(['apt-get', 'install', '-y'] + apt_packages) run_command(["apt-get", "install", "-y"] + apt_packages)
print("--- System dependencies installed successfully. ---\n", flush=True) print("--- System dependencies installed successfully. ---\n", flush=True)
def build_and_install_prerequisites(args): def build_and_install_prerequisites(args):
"""Builds UCX and NIXL from source, creating a self-contained wheel.""" """Builds UCX and NIXL from source, creating a self-contained wheel."""
if not args.force_reinstall and is_pip_package_installed('nixl'): if not args.force_reinstall and is_pip_package_installed("nixl"):
print("--> NIXL is already installed. Nothing to do.", flush=True) print("--> NIXL is already installed. Nothing to do.", flush=True)
return return
cached_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME) cached_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME)
if not args.force_reinstall and cached_wheel: if not args.force_reinstall and cached_wheel:
print(f"\n--> Found self-contained wheel: \ print(
f"\n--> Found self-contained wheel: \
{os.path.basename(cached_wheel)}.", {os.path.basename(cached_wheel)}.",
flush=True) flush=True,
print("--> Installing from cache, skipping all source builds.", )
flush=True) print("--> Installing from cache, skipping all source builds.", flush=True)
install_command = [ install_command = [sys.executable, "-m", "pip", "install", cached_wheel]
sys.executable, '-m', 'pip', 'install', cached_wheel
]
run_command(install_command) run_command(install_command)
print("\n--- Installation from cache complete. ---", flush=True) print("\n--- Installation from cache complete. ---", flush=True)
return return
print("\n--> No installed package or cached wheel found. \ print(
"\n--> No installed package or cached wheel found. \
Starting full build process...", Starting full build process...",
flush=True) flush=True,
)
print("\n--> Installing auditwheel...", flush=True) print("\n--> Installing auditwheel...", flush=True)
run_command([sys.executable, '-m', 'pip', 'install', 'auditwheel']) run_command([sys.executable, "-m", "pip", "install", "auditwheel"])
install_system_dependencies() install_system_dependencies()
ucx_install_path = os.path.abspath(UCX_INSTALL_DIR) ucx_install_path = os.path.abspath(UCX_INSTALL_DIR)
print(f"--> Using wheel cache directory: {WHEELS_CACHE_HOME}", flush=True) print(f"--> Using wheel cache directory: {WHEELS_CACHE_HOME}", flush=True)
@ -114,86 +120,91 @@ def build_and_install_prerequisites(args):
# -- Step 1: Build UCX from source -- # -- Step 1: Build UCX from source --
print("\n[1/3] Configuring and building UCX from source...", flush=True) print("\n[1/3] Configuring and building UCX from source...", flush=True)
if not os.path.exists(UCX_DIR): if not os.path.exists(UCX_DIR):
run_command(['git', 'clone', UCX_REPO_URL, UCX_DIR]) run_command(["git", "clone", UCX_REPO_URL, UCX_DIR])
ucx_source_path = os.path.abspath(UCX_DIR) ucx_source_path = os.path.abspath(UCX_DIR)
run_command(['git', 'checkout', 'v1.19.x'], cwd=ucx_source_path) run_command(["git", "checkout", "v1.19.x"], cwd=ucx_source_path)
run_command(['./autogen.sh'], cwd=ucx_source_path) run_command(["./autogen.sh"], cwd=ucx_source_path)
configure_command = [ configure_command = [
'./configure', "./configure",
f'--prefix={ucx_install_path}', f"--prefix={ucx_install_path}",
'--enable-shared', "--enable-shared",
'--disable-static', "--disable-static",
'--disable-doxygen-doc', "--disable-doxygen-doc",
'--enable-optimizations', "--enable-optimizations",
'--enable-cma', "--enable-cma",
'--enable-devel-headers', "--enable-devel-headers",
'--with-verbs', "--with-verbs",
'--enable-mt', "--enable-mt",
] ]
run_command(configure_command, cwd=ucx_source_path) run_command(configure_command, cwd=ucx_source_path)
run_command(['make', '-j', str(os.cpu_count() or 1)], cwd=ucx_source_path) run_command(["make", "-j", str(os.cpu_count() or 1)], cwd=ucx_source_path)
run_command(['make', 'install'], cwd=ucx_source_path) run_command(["make", "install"], cwd=ucx_source_path)
print("--- UCX build and install complete ---", flush=True) print("--- UCX build and install complete ---", flush=True)
# -- Step 2: Build NIXL wheel from source -- # -- Step 2: Build NIXL wheel from source --
print("\n[2/3] Building NIXL wheel from source...", flush=True) print("\n[2/3] Building NIXL wheel from source...", flush=True)
if not os.path.exists(NIXL_DIR): if not os.path.exists(NIXL_DIR):
run_command(['git', 'clone', NIXL_REPO_URL, NIXL_DIR]) run_command(["git", "clone", NIXL_REPO_URL, NIXL_DIR])
build_env = os.environ.copy() build_env = os.environ.copy()
build_env['PKG_CONFIG_PATH'] = os.path.join(ucx_install_path, 'lib', build_env["PKG_CONFIG_PATH"] = os.path.join(ucx_install_path, "lib", "pkgconfig")
'pkgconfig') ucx_lib_path = os.path.join(ucx_install_path, "lib")
ucx_lib_path = os.path.join(ucx_install_path, 'lib') ucx_plugin_path = os.path.join(ucx_lib_path, "ucx")
ucx_plugin_path = os.path.join(ucx_lib_path, 'ucx') existing_ld_path = os.environ.get("LD_LIBRARY_PATH", "")
existing_ld_path = os.environ.get('LD_LIBRARY_PATH', '') build_env["LD_LIBRARY_PATH"] = (
build_env['LD_LIBRARY_PATH'] = \ f"{ucx_lib_path}:{ucx_plugin_path}:{existing_ld_path}".strip(":")
f"{ucx_lib_path}:{ucx_plugin_path}:{existing_ld_path}".strip(':') )
print(f"--> Using LD_LIBRARY_PATH: {build_env['LD_LIBRARY_PATH']}", print(f"--> Using LD_LIBRARY_PATH: {build_env['LD_LIBRARY_PATH']}", flush=True)
flush=True)
temp_wheel_dir = os.path.join(ROOT_DIR, 'temp_wheelhouse') temp_wheel_dir = os.path.join(ROOT_DIR, "temp_wheelhouse")
run_command([ run_command(
sys.executable, '-m', 'pip', 'wheel', '.', '--no-deps', [
f'--wheel-dir={temp_wheel_dir}' sys.executable,
], "-m",
cwd=os.path.abspath(NIXL_DIR), "pip",
env=build_env) "wheel",
".",
"--no-deps",
f"--wheel-dir={temp_wheel_dir}",
],
cwd=os.path.abspath(NIXL_DIR),
env=build_env,
)
# -- Step 3: Repair the wheel by copying UCX libraries -- # -- Step 3: Repair the wheel by copying UCX libraries --
print("\n[3/3] Repairing NIXL wheel to include UCX libraries...", print("\n[3/3] Repairing NIXL wheel to include UCX libraries...", flush=True)
flush=True)
unrepaired_wheel = find_nixl_wheel_in_cache(temp_wheel_dir) unrepaired_wheel = find_nixl_wheel_in_cache(temp_wheel_dir)
if not unrepaired_wheel: if not unrepaired_wheel:
raise RuntimeError("Failed to find the NIXL wheel after building it.") raise RuntimeError("Failed to find the NIXL wheel after building it.")
# We tell auditwheel to ignore the plugin that mesonpy already handled. # We tell auditwheel to ignore the plugin that mesonpy already handled.
auditwheel_command = [ auditwheel_command = [
'auditwheel', "auditwheel",
'repair', "repair",
'--exclude', "--exclude",
'libplugin_UCX.so', # <-- Exclude because mesonpy already includes it "libplugin_UCX.so", # <-- Exclude because mesonpy already includes it
unrepaired_wheel, unrepaired_wheel,
f'--wheel-dir={WHEELS_CACHE_HOME}' f"--wheel-dir={WHEELS_CACHE_HOME}",
] ]
run_command(auditwheel_command, env=build_env) run_command(auditwheel_command, env=build_env)
# --- CLEANUP --- # --- CLEANUP ---
# No more temporary files to remove, just the temp wheelhouse # No more temporary files to remove, just the temp wheelhouse
run_command(['rm', '-rf', temp_wheel_dir]) run_command(["rm", "-rf", temp_wheel_dir])
# --- END CLEANUP --- # --- END CLEANUP ---
newly_built_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME) newly_built_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME)
if not newly_built_wheel: if not newly_built_wheel:
raise RuntimeError("Failed to find the repaired NIXL wheel.") raise RuntimeError("Failed to find the repaired NIXL wheel.")
print(f"--> Successfully built self-contained wheel: \ print(
f"--> Successfully built self-contained wheel: \
{os.path.basename(newly_built_wheel)}. Now installing...", {os.path.basename(newly_built_wheel)}. Now installing...",
flush=True) flush=True,
install_command = [ )
sys.executable, '-m', 'pip', 'install', newly_built_wheel install_command = [sys.executable, "-m", "pip", "install", newly_built_wheel]
]
if args.force_reinstall: if args.force_reinstall:
install_command.insert(-1, '--force-reinstall') install_command.insert(-1, "--force-reinstall")
run_command(install_command) run_command(install_command)
print("--- NIXL installation complete ---", flush=True) print("--- NIXL installation complete ---", flush=True)
@ -201,10 +212,13 @@ def build_and_install_prerequisites(args):
if __name__ == "__main__": if __name__ == "__main__":
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description="Build and install UCX and NIXL dependencies.") description="Build and install UCX and NIXL dependencies."
parser.add_argument('--force-reinstall', )
action='store_true', parser.add_argument(
help='Force rebuild and reinstall of UCX and NIXL \ "--force-reinstall",
even if they are already installed.') action="store_true",
help="Force rebuild and reinstall of UCX and NIXL \
even if they are already installed.",
)
args = parser.parse_args() args = parser.parse_args()
build_and_install_prerequisites(args) build_and_install_prerequisites(args)