scripts/build_wheel.py (1)

461-464: Windows mutation of extra_cmake_vars may TypeError with tuple default

See prior comment. Using += on a tuple will raise a TypeError if main() is invoked with its default. Apply the fix to coerce to list before concatenation.

scripts/build_wheel.py (1)

123-124: Good: add return type annotation; minor docstring nit

Adding -> tuple[Path, Path] improves clarity. Consider aligning the docstring’s “Returns:” type with the annotation (tuple[Path, Path]) for consistency with Google-style docstrings.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

scripts/build_wheel.py (4)

284-317: Avoid using sudo in build scripts and handle user-writable symlink directory

Creating symlinks in system directories using sudo can fail in CI environments and is a security risk. The function should create symlinks in a user-writable temporary directory when permission is denied.

Apply this refactor to avoid sudo and return the symlink directory:

-def create_cuda_stub_links(cuda_stub_dir: str, missing_libs: list[str]):
+def create_cuda_stub_links(cuda_stub_dir: str, missing_libs: list[str]) -> Path:
     """
     Creates symbolic links for CUDA stub libraries in the provided directory.
 
     Args:
         cuda_stub_dir (str): Path to the directory containing CUDA stubs.
         missing_libs: Versioned names of the missing libraries.
+    
+    Returns:
+        Path: Directory containing the symlinks (either cuda_stub_dir or a temp dir)
     """
     cuda_stub_path = Path(cuda_stub_dir)
     if not cuda_stub_path.exists():
         raise RuntimeError(
             f"CUDA stub directory '{cuda_stub_dir}' does not exist.")
 
+    link_dir = cuda_stub_path
     version_pattern = r'\.\d+$'
-    for missing_lib in filter(lambda x: re.search(version_pattern, x),
-                              missing_libs):
-        # Define `so` as the first part of `missing_lib` with trailing '.' and digits removed
-        so = cuda_stub_path / re.sub(version_pattern, '', missing_lib)
-        so_versioned = cuda_stub_path / missing_lib
-
-        # Check if the library exists and the versioned link does not.
-        if so.exists() and not so_versioned.exists():
-            try:
-                # Attempt to create the symbolic link.
-                so_versioned.symlink_to(so)
-            except PermissionError:
-                # Handle permission errors by attempting to use `sudo` to create the link.
-                try:
-                    build_run(f"sudo ln -s {str(so)} {str(so_versioned)}")
-                except CalledProcessError as sudo_error:
-                    print(
-                        f"Failed to create symbolic link even with sudo: {sudo_error}"
-                    )
+    
+    try:
+        # Try creating symlinks in the original directory
+        for missing_lib in filter(lambda x: re.search(version_pattern, x),
+                                  missing_libs):
+            so = cuda_stub_path / re.sub(version_pattern, '', missing_lib)
+            so_versioned = cuda_stub_path / missing_lib
+            
+            if so.exists() and not so_versioned.exists():
+                so_versioned.symlink_to(so)
+    except PermissionError:
+        # Fallback to user-writable temp directory
+        import tempfile
+        link_dir = Path(tempfile.mkdtemp(prefix="tllm_cuda_stubs_"))
+        print(f"Creating symlinks in temporary directory: {link_dir}")
+        
+        for missing_lib in filter(lambda x: re.search(version_pattern, x),
+                                  missing_libs):
+            so = cuda_stub_path / re.sub(version_pattern, '', missing_lib)
+            target = link_dir / missing_lib
+            
+            if so.exists() and not target.exists():
+                # Also copy the original .so file to the temp dir
+                import shutil
+                shutil.copy2(so, link_dir / so.name)
+                target.symlink_to(so.name)
+    
+    return link_dir

---

333-367: Use the returned symlink directory in LD_LIBRARY_PATH

The function now needs to use the directory returned by create_cuda_stub_links instead of always using the original cuda_stub_dir.

Apply this fix:

 def generate_python_stubs_linux(binding_type: str, venv_python: Path,
                                 deep_ep: bool):
     is_nanobind = binding_type == "nanobind"
     if is_nanobind:
         build_run(f"\"{venv_python}\" -m pip install nanobind")
     build_run(f"\"{venv_python}\" -m pip install pybind11-stubgen")
 
     env_stub_gen = os.environ.copy()
     cuda_home_dir = env_stub_gen.get("CUDA_HOME") or env_stub_gen.get(
         "CUDA_PATH") or "/usr/local/cuda"
     missing_libs = check_missing_libs("bindings")
     cuda_stub_dir = f"{cuda_home_dir}/lib64/stubs"
 
     if missing_libs and Path(cuda_stub_dir).exists():
         # Create symbolic links for the CUDA stubs
-        create_cuda_stub_links(cuda_stub_dir, missing_libs)
+        link_dir = create_cuda_stub_links(cuda_stub_dir, missing_libs)
         ld_library_path = env_stub_gen.get("LD_LIBRARY_PATH")
         env_stub_gen[
-            "LD_LIBRARY_PATH"] = f"{ld_library_path}:{cuda_stub_dir}" if ld_library_path else cuda_stub_dir
+            "LD_LIBRARY_PATH"] = f"{ld_library_path}:{link_dir}" if ld_library_path else str(link_dir)
 
     if is_nanobind:
-        build_run(f"\"{venv_python}\" -m nanobind.stubgen -m bindings -O .",
+        build_run(f"\"{venv_python}\" -m nanobind.stubgen -m bindings -O . --exit-code",
                   env=env_stub_gen)
     else:
         build_run(
             f"\"{venv_python}\" -m pybind11_stubgen -o . bindings --exit-code",
             env=env_stub_gen)

---

369-397: Fix invalid Python string escaping and use sys.exit

The Windows stub generation has several issues:

1. Uses exit(1) instead of sys.exit(1)
2. Invalid escaping r\"{lib_dir}\" in the generated Python code
3. Should use an f-string for cleaner interpolation

Apply this fix:

 def generate_python_stubs_windows(binding_type: str, venv_python: Path,
                                   pkg_dir: Path, lib_dir: Path):
     if binding_type == "nanobind":
         print("Windows not yet supported for nanobind stubs")
-        exit(1)
+        sys.exit(1)
     else:
         build_run(f"\"{venv_python}\" -m pip install pybind11-stubgen")
         stubgen = "stubgen.py"
-        stubgen_contents = """
+        stubgen_contents = f"""
                         # Loading torch, trt before bindings is required to avoid import errors on windows.
                         # isort: off
                         import torch
                         import tensorrt as trt
                         # isort: on
                         import os
                         import platform
 
                         from pybind11_stubgen import main
 
                         if __name__ == "__main__":
                             # Load dlls from `libs` directory before launching bindings.
                             if platform.system() == "Windows":
-                                os.add_dll_directory(r\"{lib_dir}\")
+                                os.add_dll_directory(r"{lib_dir}")
                             main()
-                        """.format(lib_dir=lib_dir)
+                        """
         (pkg_dir / stubgen).write_text(dedent(stubgen_contents))
-        build_run(f"\"{venv_python}\" {stubgen} -o . bindings")
+        build_run(f"\"{venv_python}\" {stubgen} -o . bindings --exit-code")
         (pkg_dir / stubgen).unlink()

---

406-407: Ensure safe mutation of extra_cmake_vars parameter

The main() function signature was changed to use Sequence[str] = tuple() for extra_cmake_vars, but line 481 attempts to mutate it with +=. This will fail if a tuple is passed.

Apply this fix at the beginning of the main() function:

 def main(*,
          build_type: str = "Release",
          generator: str = "",
          build_dir: Path = None,
          dist_dir: Path = None,
          cuda_architectures: str = None,
          job_count: int = None,
          extra_cmake_vars: Sequence[str] = tuple(),
          extra_make_targets: str = "",
          # ... other parameters ...
          ):
+    # Convert to list to allow mutation
+    extra_cmake_vars = list(extra_cmake_vars)
 
     if clean:
         clean_wheel = True

Or alternatively, use a safer approach at line 481:

     if on_windows:
         # Windows does not support multi-device currently.
-        extra_cmake_vars += ["ENABLE_MULTI_DEVICE=0"]
+        extra_cmake_vars = list(extra_cmake_vars) + ["ENABLE_MULTI_DEVICE=0"]

Also applies to: 481-481

scripts/build_wheel.py (1)

319-331: Consider error handling improvements in check_missing_libs

The function correctly identifies missing libraries but could benefit from better error handling if the ldd command fails or if the .so file doesn't exist.

Consider adding error handling:

 def check_missing_libs(so_prefix: str) -> list[str]:
+    import glob
+    so_files = glob.glob(f"{so_prefix}.cpython*.so")
+    if not so_files:
+        print(f"Warning: No shared library found matching pattern {so_prefix}.cpython*.so")
+        return []
+    
+    try:
         result = build_run(f"ldd {so_prefix}.cpython*.so",
                            capture_output=True,
                            text=True)
+    except CalledProcessError as e:
+        print(f"Warning: ldd command failed: {e}")
+        return []
+    
     missing = []
     for line in result.stdout.splitlines():
         if "not found" in line:
             lib_name = line.split()[
                 0]  # Extract the library name before "=> not found"
             if lib_name not in missing:
                 missing.append(lib_name)
     return missing

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

scripts/build_wheel.py (4)

284-317: Address previous review concerns and add error handling

This implementation still has the issues identified in the previous review:

1. Uses sudo which can be fragile in CI environments
2. Doesn't return a directory path for the caller to use
3. Missing proper error handling for edge cases

Additionally, new issues:

• Line 298: The filter and re.search usage is inefficient
• Line 315-316: The print statement spans multiple lines unnecessarily

Apply this refactored approach to address the previous review concerns:

-def create_cuda_stub_links(cuda_stub_dir: str, missing_libs: list[str]):
+def create_cuda_stub_links(cuda_stub_dir: str, missing_libs: list[str]) -> Path:
     """
-  Creates symbolic links for CUDA stub libraries in the provided directory.
+    Creates symbolic links for CUDA stub libraries, returning the effective directory.

     Args:
         cuda_stub_dir (str): Path to the directory containing CUDA stubs.
-        missing_libs: Versioned names of the missing libraries.
-  """
+        missing_libs: Versioned names of the missing libraries.
+
+    Returns:
+        Path: Directory containing the symlinks (may be original or fallback).
+    """
     cuda_stub_path = Path(cuda_stub_dir)
     if not cuda_stub_path.exists():
         raise RuntimeError(
             f"CUDA stub directory '{cuda_stub_dir}' does not exist.")

     version_pattern = r'\.\d+$'
-    for missing_lib in filter(lambda x: re.search(version_pattern, x),
-                              missing_libs):
-        # Define `so` as the first part of `missing_lib` with trailing '.' and digits removed
-        so = cuda_stub_path / re.sub(version_pattern, '', missing_lib)
-        so_versioned = cuda_stub_path / missing_lib
-
-        # Check if the library exists and the versioned link does not.
-        if so.exists() and not so_versioned.exists():
-            try:
-                # Attempt to create the symbolic link.
-                so_versioned.symlink_to(so)
-            except PermissionError:
-                # Handle permission errors by attempting to use `sudo` to create the link.
-                try:
-                    build_run(f"sudo ln -s {str(so)} {str(so_versioned)}")
-                except CalledProcessError as sudo_error:
-                    print(
-                        f"Failed to create symbolic link even with sudo: {sudo_error}"
-                    )
+    versioned_libs = [lib for lib in missing_libs if re.search(version_pattern, lib)]
+    
+    try:
+        # Try in-place links inside the stubs directory first.
+        for missing_lib in versioned_libs:
+            so = cuda_stub_path / re.sub(version_pattern, '', missing_lib)
+            so_versioned = cuda_stub_path / missing_lib
+            if so.exists() and not so_versioned.exists():
+                so_versioned.symlink_to(so)
+        return cuda_stub_path
+    except PermissionError:
+        # Fall back to a per-user temp directory (no sudo).
+        import tempfile
+        link_dir = Path(tempfile.mkdtemp(prefix="tllm_cuda_stubs_"))
+        for missing_lib in versioned_libs:
+            so = cuda_stub_path / re.sub(version_pattern, '', missing_lib)
+            target = link_dir / missing_lib
+            if so.exists() and not target.exists():
+                target.symlink_to(so)
+        return link_dir

---

369-397: Fix invalid string escaping and use sys.exit

The Windows stub generation has the same issues identified in previous reviews:

1. Line 373: exit(1) should be sys.exit(1)
2. Line 391: r\"{lib_dir}\" is invalid Python syntax
3. Line 393: Using .format() with a multi-line string is error-prone

Apply this diff to fix the issues:

     if binding_type == "nanobind":
         print("Windows not yet supported for nanobind stubs")
-        exit(1)
+        sys.exit(1)
     else:
         build_run(f"\"{venv_python}\" -m pip install pybind11-stubgen")
         stubgen = "stubgen.py"
-        stubgen_contents = """
+        stubgen_contents = f"""
                         # Loading torch, trt before bindings is required to avoid import errors on windows.
                         # isort: off
                         import torch
                         import tensorrt as trt
                         # isort: on
                         import os
                         import platform

                         from pybind11_stubgen import main

                         if __name__ == "__main__":
                             # Load dlls from `libs` directory before launching bindings.
                             if platform.system() == "Windows":
-                                os.add_dll_directory(r\"{lib_dir}\")
+                                os.add_dll_directory(r"{lib_dir}")
                             main()
-                        """.format(lib_dir=lib_dir)
+                        """
         (pkg_dir / stubgen).write_text(dedent(stubgen_contents))
-        build_run(f"\"{venv_python}\" {stubgen} -o . bindings")
+        build_run(f"\"{venv_python}\" {stubgen} -o . bindings --exit-code")
         (pkg_dir / stubgen).unlink()

---

406-406: Fix type annotation and mutation safety

The function signature uses Sequence[str] = tuple() but Line 481 tries to mutate it with list(extra_cmake_vars) + [...]. This is the exact issue identified in previous reviews.

The current code on Line 481 already handles this correctly with list(extra_cmake_vars) + [...], so the implementation is actually correct. The type annotation properly reflects that the function accepts an immutable sequence.

---

343-351: Update to use the returned path from create_cuda_stub_links

The current implementation doesn't use the return value from create_cuda_stub_links, which means it may set LD_LIBRARY_PATH to a directory where symlink creation failed.

Apply this diff to use the returned path:

     if missing_libs and Path(cuda_stub_dir).exists():
-        # Create symbolic links for the CUDA stubs
-        create_cuda_stub_links(cuda_stub_dir, missing_libs)
+        # Create symbolic links for the CUDA stubs  
+        effective_stub_dir = create_cuda_stub_links(cuda_stub_dir, missing_libs)
         ld_library_path = env_stub_gen.get("LD_LIBRARY_PATH")
         env_stub_gen[
-            "LD_LIBRARY_PATH"] = f"{ld_library_path}:{cuda_stub_dir}" if ld_library_path else cuda_stub_dir
+            "LD_LIBRARY_PATH"] = f"{ld_library_path}:{effective_stub_dir}" if ld_library_path else str(effective_stub_dir)

Note: This change depends on first implementing the return value fix in create_cuda_stub_links.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

scripts/build_wheel.py (4)

285-329: Use the returned link directory and avoid filesystem operations without error handling

The function create_cuda_stub_links returns a link directory but the caller doesn't consistently use it. Also, the function could fail silently in some edge cases.

Current issues:

1. The caller at lines 361-364 uses link_dir but also includes cuda_stub_dir in LD_LIBRARY_PATH
2. No validation that temporary directory creation succeeded
3. Missing cleanup on partial failures

Apply this improvement:

-        env_stub_gen["LD_LIBRARY_PATH"] = ":".join(
-            filter(None, [link_dir, cuda_stub_dir, ld_library_path]))
+        env_stub_gen["LD_LIBRARY_PATH"] = ":".join(
+            filter(None, [link_dir, ld_library_path]))

The function should only add the link_dir (where versioned symlinks exist) to LD_LIBRARY_PATH, not both directories.

---

332-344: Add error handling for ldd command failure

The check_missing_libs function doesn't handle the case where the glob pattern matches no files or when ldd fails.

Apply this fix to handle errors gracefully:

 def check_missing_libs(so_prefix: str) -> list[str]:
-    result = build_run(f"ldd {so_prefix}.cpython*.so",
-                       capture_output=True,
-                       text=True)
+    import glob
+    so_files = glob.glob(f"{so_prefix}.cpython*.so")
+    if not so_files:
+        return []
+    
+    missing = []
+    for so_file in so_files:
+        try:
+            result = build_run(f"ldd {so_file}",
+                              capture_output=True,
+                              text=True)
+        except CalledProcessError:
+            continue
-    missing = []
-    for line in result.stdout.splitlines():
-        if "not found" in line:
-            lib_name = line.split()[
-                0]  # Extract the library name before "=> not found"
-            if lib_name not in missing:
-                missing.append(lib_name)
+        for line in result.stdout.splitlines():
+            if "not found" in line:
+                parts = line.split()
+                if parts:  # Ensure line has content
+                    lib_name = parts[0]
+                    if lib_name not in missing:
+                        missing.append(lib_name)
     return missing

---

388-416: Fix invalid raw string escaping and use sys.exit

The Windows stub generation has incorrect string escaping that will produce invalid Python code.

Apply this fix:

 def generate_python_stubs_windows(binding_type: str, venv_python: Path,
                                   pkg_dir: Path, lib_dir: Path):
     if binding_type == "nanobind":
         print("Windows not yet supported for nanobind stubs")
-        exit(1)
+        sys.exit(1)
     else:
         build_run(f"\"{venv_python}\" -m pip install pybind11-stubgen")
         stubgen = "stubgen.py"
-        stubgen_contents = """
+        stubgen_contents = f"""
                         # Loading torch, trt before bindings is required to avoid import errors on windows.
                         # isort: off
                         import torch
                         import tensorrt as trt
                         # isort: on
                         import os
                         import platform
 
                         from pybind11_stubgen import main
 
                         if __name__ == "__main__":
                             # Load dlls from `libs` directory before launching bindings.
                             if platform.system() == "Windows":
-                                os.add_dll_directory(r\"{lib_dir}\")
+                                os.add_dll_directory(r"{lib_dir}")
                             main()
-                        """.format(lib_dir=lib_dir)
+                        """
         (pkg_dir / stubgen).write_text(dedent(stubgen_contents))
-        build_run(f"\"{venv_python}\" {stubgen} -o . bindings")
+        build_run(f"\"{venv_python}\" {stubgen} -o . bindings --exit-code")
         (pkg_dir / stubgen).unlink()

---

425-426: Safely handle sequence mutation on Windows

The function signature accepts extra_cmake_vars: Sequence[str] which could be a tuple, but line 500 attempts to mutate it on Windows.

Apply this fix to ensure safe mutation:

     if on_windows:
         # Windows does not support multi-device currently.
-        extra_cmake_vars = list(extra_cmake_vars) + ["ENABLE_MULTI_DEVICE=0"]
+        extra_cmake_vars = list(extra_cmake_vars)
+        extra_cmake_vars.append("ENABLE_MULTI_DEVICE=0")

scripts/build_wheel.py (2)

305-309: Simplify regex pattern extraction logic

The regex operations for extracting library names could be more robust and clearer.

Consider this improvement for better readability:

-    version_pattern = r'\.\d+'
-    for missing_lib in filter(lambda x: re.search(version_pattern, x),
-                              missing_libs):
-        # Define `so` as the first part of `missing_lib` with trailing '.' and digits removed
-        so = cuda_stub_path / re.sub(version_pattern, '', missing_lib)
+    # Match library names ending with version numbers (e.g., .1, .12)
+    version_pattern = re.compile(r'(lib\w+\.so)(\.\d+)+$')
+    for missing_lib in missing_libs:
+        match = version_pattern.match(missing_lib)
+        if not match:
+            continue
+        base_lib_name = match.group(1)
+        so = cuda_stub_path / base_lib_name

---

369-382: Add consistent exit-code handling for nanobind stubgen

The nanobind stubgen call doesn't use --exit-code while pybind11_stubgen calls do, leading to inconsistent error reporting.

For consistency with pybind11_stubgen calls, consider adding exit code handling:

         if is_nanobind:
-            build_run(f"\"{venv_python}\" -m nanobind.stubgen -m bindings -O .",
+            build_run(f"\"{venv_python}\" -m nanobind.stubgen -m bindings -O . --exit-code",
                       env=env_stub_gen)

Note: This assumes nanobind.stubgen supports the --exit-code flag. If not supported, this suggestion can be ignored.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.