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

Signed-off-by: mgoin <mgoin64@gmail.com>

tools/install_nixl_from_source_ubuntu.py

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