Bindings: Add Python wrapper in meta-package (#954)

ovidiusm · web-flow · commit a92772cb40ba · 2025-10-27T16:59:36.000-04:00
* Add Python wrapper in meta-package

Signed-off-by: Ovidiu Mara &lt;ovidium@nvidia.com&gt;

* Clarify installation procedure

Signed-off-by: Ovidiu Mara &lt;ovidium@nvidia.com&gt;

* Remove shebang

Signed-off-by: Ovidiu Mara &lt;ovidium@nvidia.com&gt;

* Revert quote change

Signed-off-by: Ovidiu Mara &lt;ovidium@nvidia.com&gt;

* Silence mypy warning in precommit

Signed-off-by: Ovidiu Mara &lt;ovidium@nvidia.com&gt;

* Fix `pip install .` it should build nixl-cu12 by default

Signed-off-by: Ovidiu Mara &lt;ovidium@nvidia.com&gt;

* Make nixl-cu12 a required dependency of nixl

Signed-off-by: Ovidiu Mara &lt;ovidium@nvidia.com&gt;

* Install uv build dependency in CI

Signed-off-by: Ovidiu Mara &lt;ovidium@nvidia.com&gt;

* Fix wheel build issues in CI environment

Signed-off-by: Ovidiu Mara &lt;ovidium@nvidia.com&gt;

* Fix gitlab CI pipeline

Signed-off-by: Ovidiu Mara &lt;ovidium@nvidia.com&gt;

---------

Signed-off-by: Ovidiu Mara &lt;ovidium@nvidia.com&gt;
diff --git a/.gitlab/build.sh b/.gitlab/build.sh
@@ -122,6 +122,11 @@ chmod +x rustup-init
 ./rustup-init -y --default-toolchain 1.86.0
 export PATH="$HOME/.cargo/bin:$PATH"
 
+wget --tries=3 --waitretry=5 "https://astral.sh/uv/install.sh" -O install_uv.sh
+chmod +x install_uv.sh
+./install_uv.sh
+export PATH="$HOME/.local/bin:$PATH"
+
 curl -fSsL "https://github.com/openucx/ucx/tarball/${UCX_VERSION}" | tar xz
 ( \
   cd openucx-ucx* && \
@@ -211,6 +216,7 @@ export UCX_TLS=^cuda_ipc
 # shellcheck disable=SC2086
 meson setup nixl_build --prefix=${INSTALL_DIR} -Ducx_path=${UCX_INSTALL_DIR} -Dbuild_docs=true -Drust=false ${EXTRA_BUILD_ARGS} -Dlibfabric_path="${LIBFABRIC_INSTALL_DIR}"
 ninja -C nixl_build && ninja -C nixl_build install
+mkdir -p dist && cp nixl_build/src/bindings/python/nixl-meta/nixl-*.whl dist/
 
 # TODO(kapila): Copy the nixl.pc file to the install directory if needed.
 # cp ${BUILD_DIR}/nixl.pc ${INSTALL_DIR}/lib/pkgconfig/nixl.pc
diff --git a/.gitlab/test_python.sh b/.gitlab/test_python.sh
@@ -41,6 +41,7 @@ export NIXL_PREFIX=${INSTALL_DIR}
 export NIXL_DEBUG_LOGGING=yes
 
 pip3 install --break-system-packages .
+pip3 install --break-system-packages dist/nixl-*none-any.whl
 pip3 install --break-system-packages pytest
 pip3 install --break-system-packages pytest-timeout
 pip3 install --break-system-packages zmq
diff --git a/README.md b/README.md
@@ -162,6 +162,7 @@ Or build from source:
 ```bash
 # From the root nixl directory
 pip install .
+pip install <name_of_build_dir>/src/bindings/python/nixl-meta/nixl-*.whl
 ```
 
 For Python examples, see [examples/python/](examples/python/).
diff --git a/contrib/Dockerfile b/contrib/Dockerfile
@@ -222,4 +222,5 @@ RUN export PATH=$VIRTUAL_ENV/bin:$PATH && \
     --nixl-plugins-dir $NIXL_PLUGIN_DIR \
     --output-dir /workspace/nixl/dist
 
-RUN uv pip install dist/nixl*cp${DEFAULT_PYTHON_VERSION//./}*.whl
+RUN cp build/src/bindings/python/nixl-meta/nixl-*.whl dist/
+RUN uv pip install dist/nixl*cp${DEFAULT_PYTHON_VERSION//./}*.whl dist/nixl-*-none-any.whl
diff --git a/contrib/build-wheel.sh b/contrib/build-wheel.sh
@@ -83,7 +83,8 @@ if [ "$CUDA_MAJOR" -ne 12 ] && [ "$CUDA_MAJOR" -ne 13 ]; then
     exit 1
 fi
 PKG_NAME="nixl-cu${CUDA_MAJOR}"
-./contrib/tomlutil.py --set-name $PKG_NAME pyproject.toml
+PKG_DIR="nixl_cu${CUDA_MAJOR}"
+./contrib/tomlutil.py --wheel-name $PKG_NAME --wheel-dir $PKG_DIR pyproject.toml
 uv build --wheel --out-dir $TMP_DIR --python $PYTHON_VERSION
 
 # Bundle libraries
diff --git a/contrib/tomlutil.py b/contrib/tomlutil.py
@@ -20,15 +20,38 @@
 import tomlkit
 
 parser = argparse.ArgumentParser()
-parser.add_argument("--set-name", type=str, help="Set the project name")
+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()
 
 with open(args.file) as f:
     doc = tomlkit.parse(f.read())
 
-if args.set_name:
-    doc["project"]["name"] = args.set_name
+if args.wheel_name:
+    # Set the wheel name
+    # Example:
+    # ```toml
+    # [project]
+    # name = "<wheel_name>"
+    # ```
+    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))
diff --git a/meson_options.txt b/meson_options.txt
@@ -28,6 +28,7 @@ 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')
diff --git a/pyproject.toml b/pyproject.toml
@@ -18,7 +18,7 @@ requires = ["meson-python", "pybind11", "patchelf", "pyyaml", "types-PyYAML", "p
 build-backend = "mesonpy"
 
 [project]
-name = 'nixl'
+name = 'nixl-cu12'
 version = '0.6.1'
 description = 'NIXL Python API'
 readme = 'README.md'
@@ -34,4 +34,4 @@ dependencies = ["torch", "numpy"]
 profile = "black"
 
 [tool.meson-python.args]
-setup = ['-Dinstall_headers=false']
+setup = ['-Dinstall_headers=false', '-Dwheel_dir=nixl_cu12']
diff --git a/src/api/python/__init__.py b/src/api/python/__init__.py
@@ -15,17 +15,14 @@
 
 import os
 
-import nixl
-
-__all__ = ["nixl"]
-
 if "NIXL_PLUGIN_DIR" not in os.environ:
     # name for local installation
-    plugin_dir = nixl.__file__[:-16] + ".nixl.mesonpy.libs/plugins/"
+    package_root = os.path.dirname(__file__)
+    plugin_dir = os.path.join(package_root, ".nixl.mesonpy.libs/plugins/")
 
-    # name for pypi installation
+    # fallback for PyPI installation
     if not os.path.isdir(plugin_dir):
-        plugin_dir = nixl.__file__[:-16] + ".nixl_pybind.mesonpy.libs/plugins/"
+        plugin_dir = os.path.join(package_root, ".nixl_pybind.mesonpy.libs/plugins/")
 
     if os.path.isdir(plugin_dir):
         os.environ["NIXL_PLUGIN_DIR"] = plugin_dir
diff --git a/src/api/python/_api.py b/src/api/python/_api.py
@@ -19,8 +19,8 @@
 import numpy as np
 import torch
 
-import nixl._bindings as nixlBind
-from nixl.logging import get_logger
+from . import _bindings as nixlBind  # type: ignore
+from .logging import get_logger
 
 # Get logger using centralized configuration
 logger = get_logger(__name__)
diff --git a/src/api/python/meson.build b/src/api/python/meson.build
@@ -15,7 +15,9 @@
 
 py = import('python').find_installation('python3', pure: false)
 
-py.install_sources('_api.py', subdir: ('nixl'))
-py.install_sources('__init__.py', subdir: ('nixl'))
-py.install_sources('py.typed', subdir: ('nixl'))
-py.install_sources('logging.py', subdir: ('nixl'))
+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))
diff --git a/src/bindings/python/meson.build b/src/bindings/python/meson.build
@@ -19,16 +19,18 @@ 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: ('nixl'),
+           subdir: (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: ('nixl'),
+           subdir: (wheel_dir),
            dependencies: [pybind_dep],
            install: true)
 
diff --git a/src/bindings/python/nixl-meta/README.md b/src/bindings/python/nixl-meta/README.md
@@ -7,3 +7,5 @@ Install one of:
 pip install "nixl[cu12]"   # for CUDA 12
 pip install "nixl[cu13]"   # for CUDA 13
 ```
+
+If both are installed, the higher version will be used.
diff --git a/src/bindings/python/nixl-meta/meson.build b/src/bindings/python/nixl-meta/meson.build
@@ -13,20 +13,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-source_dir = meson.current_source_dir()
+fs = import('fs')
+
 build_dir = meson.current_build_dir()
 
-configure_file(
+pyproject_toml = configure_file(
   input: 'pyproject.toml.in',
   output: 'pyproject.toml',
   configuration: { 'VERSION': meson.project_version() }
 )
 
+readme_md = fs.copyfile('README.md')
+
+nixl_sources = files('nixl/__init__.py')
+
+subdir('nixl')
+
 uv = find_program('uv', required: false)
 if uv.found()
   wheel_name = 'nixl-@0@-py3-none-any.whl'.format(meson.project_version())
   meta_wheel = custom_target(
     'build_nixl_meta',
+    input: [pyproject_toml, readme_md] + nixl_sources,
     output: [wheel_name],
     command: [uv, 'build', '--wheel', '--out-dir', build_dir, build_dir],
     install: false,
diff --git a/src/bindings/python/nixl-meta/nixl/__init__.py b/src/bindings/python/nixl-meta/nixl/__init__.py
@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib
+import sys
+
+# Try packages in order
+candidates = ["nixl_cu13", "nixl_cu12"]
+_pkg = None
+for pkg in candidates:
+    try:
+        _pkg = importlib.import_module(pkg)
+        break
+    except ModuleNotFoundError:
+        continue
+
+if _pkg is None:
+    raise ImportError(
+        "Could not find CUDA-specific NIXL package. Please install NIXL with `pip install nixl[cu12]` or `pip install nixl[cu13]`"
+    )
+
+submodules = ["_api", "_bindings", "_utils", "logging"]
+for sub_name in submodules:
+    # Import submodule from actual wheel
+    module = importlib.import_module(f"{pkg}.{sub_name}")
+    # Make it accessible as nixl._api, nixl._utils, nixl.logging
+    sys.modules[f"nixl.{sub_name}"] = module
+    # Also add the submodule itself to the nixl namespace
+    setattr(sys.modules[__name__], sub_name, module)
+
+    # Expose all symbols from the submodule under the nixl namespace
+    for attr in dir(module):
+        if not attr.startswith("_"):
+            setattr(sys.modules[__name__], attr, getattr(module, attr))
diff --git a/src/bindings/python/nixl-meta/nixl/meson.build b/src/bindings/python/nixl-meta/nixl/meson.build
@@ -0,0 +1,17 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+fs = import('fs')
+fs.copyfile('__init__.py')
diff --git a/src/bindings/python/nixl-meta/pyproject.toml.in b/src/bindings/python/nixl-meta/pyproject.toml.in
@@ -24,8 +24,11 @@ description = "Meta package for nixl CUDA variants"
 readme = "README.md"
 requires-python = ">=3.9"
 
-dependencies = []
+dependencies = ["nixl-cu12>=@VERSION@"]
 
 [project.optional-dependencies]
 cu12 = ["nixl-cu12>=@VERSION@"]
 cu13 = ["nixl-cu13>=@VERSION@"]
+
+[tool.setuptools]
+packages = ["nixl"]
