Bindings: fix missing wheel dependency in CUDA 13 container builds (#966)

* Bindings: fix missing wheel dependency in CUDA 13 container builds

Signed-off-by: Ovidiu Mara <[email protected]>

* Fix nixlbench container installation

Signed-off-by: Ovidiu Mara <[email protected]>

---------

Signed-off-by: Ovidiu Mara <[email protected]>

Author: ovidiusm

benchmark/nixlbench/contrib/Dockerfile

@@ -205,7 +205,10 @@ RUN echo "/usr/local/nixl/lib/$ARCH-linux-gnu" > /etc/ld.so.conf.d/nixl.conf &&
     echo "/usr/local/nixl/lib/$ARCH-linux-gnu/plugins" >> /etc/ld.so.conf.d/nixl.conf && \
     ldconfig
 
-RUN uv pip install .
+RUN CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
+    ./contrib/tomlutil.py --wheel-name "nixl-cu${CUDA_MAJOR}" pyproject.toml && \
+    uv pip install . && \
+    uv pip install build/src/bindings/python/nixl-meta/nixl-*.whl
 
 WORKDIR /workspace/nixlbench
 

contrib/Dockerfile.manylinux

@@ -280,7 +280,10 @@ RUN IFS=',' read -ra PYTHON_VERSIONS <<< "$WHL_PYTHON_VERSIONS" && \
             --output-dir dist ; \
     done
 
-# Copy the meta package wheel to the output directory
-RUN cp build/src/bindings/python/nixl-meta/nixl*.whl dist/
+# Copy the meta package wheel to the dist directory, which will be used to push to PyPI.
+# Only do this for the CUDA 12 builds, since by default nixl depends on nixl-cu12.
+RUN if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
+      cp build/src/bindings/python/nixl-meta/nixl*.whl dist/; \
+    fi
 
-RUN uv pip install dist/nixl*cp${DEFAULT_PYTHON_VERSION//./}*.whl
+RUN uv pip install dist/nixl*cp${DEFAULT_PYTHON_VERSION//./}*.whl build/src/bindings/python/nixl-meta/nixl*.whl

contrib/build-wheel.sh

@@ -83,8 +83,7 @@ if [ "$CUDA_MAJOR" -ne 12 ] && [ "$CUDA_MAJOR" -ne 13 ]; then
     exit 1
 fi
 PKG_NAME="nixl-cu${CUDA_MAJOR}"
-PKG_DIR="nixl_cu${CUDA_MAJOR}"
-./contrib/tomlutil.py --wheel-name $PKG_NAME --wheel-dir $PKG_DIR pyproject.toml
+./contrib/tomlutil.py --wheel-name $PKG_NAME pyproject.toml
 uv build --wheel --out-dir $TMP_DIR --python $PYTHON_VERSION
 
 # Bundle libraries

contrib/tomlutil.py

@@ -21,7 +21,6 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--wheel-name", type=str, help="Set the project name")
-parser.add_argument("--wheel-dir", type=str, help="Set the wheel dir")
 parser.add_argument("file", type=str, help="The toml file to modify")
 args = parser.parse_args()
 
@@ -37,21 +36,5 @@
     # ```
     doc["project"]["name"] = args.wheel_name
 
-if args.wheel_dir:
-    # Set the wheel dir
-    # Example:
-    # ```toml
-    # [tool.meson-python.args]
-    # setup = ["-Dinstall_headers=false", "-Dwheel_dir=<wheel_dir>"]
-    # ```
-    if "meson-python" not in doc["tool"]:
-        doc["tool"]["meson-python"] = tomlkit.table()
-    if "args" not in doc["tool"]["meson-python"]:
-        doc["tool"]["meson-python"]["args"] = tomlkit.table()
-    setup = doc["tool"]["meson-python"]["args"].get("setup", [])
-    setup = [s for s in setup if not s.startswith("-Dwheel_dir=")]
-    setup.append(f"-Dwheel_dir={args.wheel_dir}")
-    doc["tool"]["meson-python"]["args"]["setup"] = setup
-
 with open(args.file, "w") as f:
     f.write(tomlkit.dumps(doc))

meson.build

@@ -107,9 +107,21 @@ if cuda_dep.found()
         warning('GPUNETIO plugin not supported in CUDA version: ' + nvcc.version())
         doca_gpunetio_dep = disabler()
     endif
+
+    # Set the Python CUDA-specific wheel directory
+    cuda_version_major = nvcc.version().split('.')[0]
+    if cuda_version_major == '12'
+        cuda_wheel_dir = 'nixl_cu12'
+    elif cuda_version_major == '13'
+        cuda_wheel_dir = 'nixl_cu13'
+    else
+        error('Unsupported CUDA version: ' + cuda_version_major)
+    endif
 else
     warning('CUDA not found. UCX backend will be built without CUDA support, and some plugins will be disabled.')
     doca_gpunetio_dep = disabler()
+    warning('CUDA not found, cannot autodetect wheel dir; defaulting to nixl_cu12')
+    cuda_wheel_dir = 'nixl_cu12'
 endif
 
 # DOCA GPUNETIO

meson_options.txt

@@ -28,7 +28,6 @@ option('static_plugins', type: 'string', value: '', description: 'Plugins to be
 option('build_docs', type: 'boolean', value: false, description: 'Build Doxygen documentation')
 option('log_level', type: 'combo', choices: ['trace', 'debug', 'info', 'warning', 'error', 'fatal', 'auto'], value: 'auto', description: 'Log Level (auto: auto-detect based on build type: trace for debug builds, info for release builds)')
 option('rust', type: 'boolean', value: false, description: 'Build Rust bindings')
-option('wheel_dir', type : 'string', value : 'nixl', description : 'Set a custom Python wheel dir (e.g. nixl_cu12)')
 
 # Tests
 option('build_tests', type: 'boolean', value: true, description: 'Build all tests')

pyproject.toml

@@ -34,4 +34,4 @@ dependencies = ["torch", "numpy"]
 profile = "black"
 
 [tool.meson-python.args]
-setup = ['-Dinstall_headers=false', '-Dwheel_dir=nixl_cu12']
+setup = ['-Dinstall_headers=false']

src/api/python/meson.build

@@ -15,9 +15,7 @@
 
 py = import('python').find_installation('python3', pure: false)
 
-wheel_dir = get_option('wheel_dir')
-
-py.install_sources('_api.py', subdir: (wheel_dir))
-py.install_sources('__init__.py', subdir: (wheel_dir))
-py.install_sources('py.typed', subdir: (wheel_dir))
-py.install_sources('logging.py', subdir: (wheel_dir))
+py.install_sources('_api.py', subdir: (cuda_wheel_dir))
+py.install_sources('__init__.py', subdir: (cuda_wheel_dir))
+py.install_sources('py.typed', subdir: (cuda_wheel_dir))
+py.install_sources('logging.py', subdir: (cuda_wheel_dir))

src/bindings/python/meson.build

@@ -19,18 +19,16 @@ nixl_dep = declare_dependency(link_with: nixl_lib, include_directories: nixl_inc
 
 py = import('python').find_installation('python3', pure: false)
 
-wheel_dir = get_option('wheel_dir')
-
 py.extension_module('_bindings',
            'nixl_bindings.cpp',
-           subdir: (wheel_dir),
+           subdir: (cuda_wheel_dir),
            dependencies: [nixl_dep, serdes_interface, pybind_dep],
            include_directories: [nixl_inc_dirs, utils_inc_dirs],
            install: true)
 
 py.extension_module('_utils',
            'nixl_utils.cpp',
-           subdir: (wheel_dir),
+           subdir: (cuda_wheel_dir),
            dependencies: [pybind_dep],
            install: true)
 

src/bindings/python/nixl-meta/meson.build

@@ -20,7 +20,7 @@ build_dir = meson.current_build_dir()
 pyproject_toml = configure_file(
   input: 'pyproject.toml.in',
   output: 'pyproject.toml',
-  configuration: { 'VERSION': meson.project_version() }
+  configuration: { 'VERSION': meson.project_version(), 'WHEEL_DEP': cuda_wheel_dir.replace('_', '-') }
 )
 
 readme_md = fs.copyfile('README.md')