[python] switch to setup.py to build package

.gitignore

@@ -1,5 +1,7 @@
 .vscode/
+.hypothesis/
 build/
+dist/
 __pycache__
 .*.swp
 .idea/

python/CMakeLists.txt

@@ -3,21 +3,21 @@ cmake_minimum_required(VERSION 3.18...3.22)
 find_package(Python 3.9 COMPONENTS Interpreter Development.Module REQUIRED)
 
 if (NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
-  set(CMAKE_BUILD_TYPE Release CACHE STRING "Choose the type of build." FORCE)
-  set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
-endif()
+    set(CMAKE_BUILD_TYPE Release CACHE STRING "Choose the type of build." FORCE)
+    set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
+endif ()
 
 # Create CMake targets for all Python components needed by nanobind
 if (CMAKE_VERSION VERSION_GREATER_EQUAL 3.26)
-  find_package(Python 3.8 COMPONENTS Interpreter Development.Module Development.SABIModule REQUIRED)
-else()
-  find_package(Python 3.8 COMPONENTS Interpreter Development.Module REQUIRED)
-endif()
+    find_package(Python 3.8 COMPONENTS Interpreter Development.Module Development.SABIModule REQUIRED)
+else ()
+    find_package(Python 3.8 COMPONENTS Interpreter Development.Module REQUIRED)
+endif ()
 
 # Detect the installed nanobind package and import it into CMake
 execute_process(
-  COMMAND "${Python_EXECUTABLE}" -c "import nanobind; print(nanobind.cmake_dir())"
-  OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE NB_DIR)
+        COMMAND "${Python_EXECUTABLE}" -c "import nanobind; print(nanobind.cmake_dir())"
+        OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE NB_DIR)
 list(APPEND CMAKE_PREFIX_PATH "${NB_DIR}")
 find_package(nanobind CONFIG REQUIRED)
 
@@ -26,47 +26,36 @@ set(CUDA_DIR "/usr/local/cuda")
 set(MSCCLPP_DIR ${CMAKE_CURRENT_LIST_DIR}/../build)
 
 nanobind_add_module(
-  _py_mscclpp
-  NOSTRIP
-  NB_STATIC
-  src/_py_mscclpp.cpp
+        _py_mscclpp
+        NOSTRIP
+        NB_STATIC
+        src/_py_mscclpp.cpp
 )
 
 target_include_directories(
-  _py_mscclpp
-  PUBLIC
-  ${CUDA_DIR}/include
-  ${MSCCLPP_DIR}/include
+        _py_mscclpp
+        PUBLIC
+        ${CUDA_DIR}/include
+        ${MSCCLPP_DIR}/include
 )
 target_link_directories(
-  _py_mscclpp
-  PUBLIC
-  ${CUDA_DIR}/lib
-  ${MSCCLPP_DIR}/lib
+        _py_mscclpp
+        PUBLIC
+        ${CUDA_DIR}/lib
+        ${MSCCLPP_DIR}/lib
 )
 
 target_link_libraries(
-  _py_mscclpp
-  PUBLIC
-  mscclpp
+        _py_mscclpp
+        PUBLIC
+        mscclpp
 )
 
 add_custom_target(build-package ALL DEPENDS _py_mscclpp)
-add_custom_command(
-  TARGET build-package POST_BUILD
-  COMMAND ${CMAKE_COMMAND} -E copy_directory
-  ${CMAKE_CURRENT_SOURCE_DIR}/src/mscclpp
-  ${CMAKE_CURRENT_BINARY_DIR}/mscclpp)
-
-add_custom_command(
-  TARGET build-package POST_BUILD
-  COMMAND ${CMAKE_COMMAND} -E copy
-  $<TARGET_FILE:_py_mscclpp>
-  ${CMAKE_CURRENT_BINARY_DIR}/mscclpp/)
 
 add_custom_command(
-  TARGET build-package POST_BUILD
-  COMMAND ${CMAKE_COMMAND} -E copy
-  ${MSCCLPP_DIR}/lib/libmscclpp.so
-  ${CMAKE_CURRENT_BINARY_DIR}/mscclpp/)
+        TARGET build-package POST_BUILD
+        COMMAND ${CMAKE_COMMAND} -E copy
+        ${MSCCLPP_DIR}/lib/libmscclpp.so
+        ${CMAKE_CURRENT_BINARY_DIR})
 

python/README.md

@@ -1,77 +1,30 @@
 # Python bindings
 
 Test instructions:
- * Compile the `libmscclpp.so` library.
- * Install `cmake` verion >= 3.18
- * setup a python virtual env
-   * `pip install -r requirements.txt`
- * `./tesh.sh`
 
-Rough build attemtps
-```
-# cd to this directory:
-
-# setup/enter pyenv environment for python 3.9
-
-# install nanabind and the test requirements.
-pip install -r requirements.txt
-
-# setup and build the CMake environments.
-# this requires nanobind, installed above.
-./setup.sh
-
-# test the module
-pytest build/mscclpp
-```
+* Compile the `libmscclpp.so` library.
+* Install `cmake` verion >= 3.18
+* setup a python virtual env
+    * `pip install -r dev-requirements.txt`
+* `./tesh.sh`
 
+## Run CI:
 
-## Installing `gdrcopy` and `mpi`
-This assumes that some things are built/installed
+```bash
+./ci.sh
 ```
-# assumes WORKDIR has:
-# git clone [email protected]/NVIDIA/gdrcopy.git
-# git clone [email protected]:microsoft/mscclpp.git
-
-uname -r
-# 5.4.0-1090-azure
-
-# install
 
-# break /usr/sbin/policy-rc.d so we can install modules
-echo '#!/bin/sh
-exit 0' > /usr/sbin/policy-rc.d
+## Build a wheel:
 
-apt update
-apt install -y \
-  build-essential devscripts debhelper check \
-  libsubunit-dev fakeroot pkg-config dkms \
-  linux-headers-5.4.0-1090-azure 
-  
-apt install -y nvidia-dkms-525-server
+Setup dev environment, then:
 
+```bash
+python setup.py bdist_wheel
+```
 
-cd $WORKDIR/gdrcopy
-sed -i 's/\(-L \$(CUDA)\/lib64\)/\1 \1\/stubs/' tests/Makefile
-cd packages
-CUDA=/usr/local/cuda ./build-deb-packages.sh
-
-dpkg -i gdrdrv-dkms_2.3-1_amd64.Ubuntu20_04.deb
-dpkg -i libgdrapi_2.3-1_amd64.Ubuntu20_04.deb
-dpkg -i gdrcopy-tests_2.3-1_amd64.Ubuntu20_04+cuda11.6.deb
-dpkg -i gdrcopy_2.3-1_amd64.Ubuntu20_04.deb
-
-# validate:
-# $ sanity
-# Running suite(s): Sanity
-# 100%: Checks: 27, Failures: 0, Errors: 0
-
-# dkms install -m gdrdrv/2.3
-
-cd $WORKDIR/mscclpp
+## Installing mpi and numa libs.
 
+```
 ## numctl
 apt install -y numactl libnuma-dev libnuma1
-
-# if not mpi testing
-USE_MPI_FOR_TESTS=0 make -j
 ```

python/ci.sh

@@ -17,7 +17,7 @@ python -m venv .venv
 source .venv/bin/activate
 
 # install venv deps.
-pip install -r requirements.txt
+pip install -r dev-requirements.txt
 
 # run the build and test.
 ./test.sh

python/requirements.txt → python/dev-requirements.txt

@@ -9,4 +9,3 @@ PyHamcrest
 nanobind
 
 torch
-numpy

python/setup.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python
+
+import os
+import shutil
+import sys
+import logging
+from setuptools import Extension, find_packages, setup
+from setuptools.command.build_ext import build_ext
+import subprocess
+
+THIS_DIR = os.path.abspath(os.path.dirname(__file__))
+
+
+class CustomExt(Extension):
+    def __init__(self, name):
+        # don't invoke the original build_ext for this special extension
+        super().__init__(name, sources=[])
+
+
+class custom_build_ext(build_ext):
+    def run(self):
+        for ext in self.extensions:
+            if isinstance(ext, CustomExt):
+                self.build_extension(ext)
+            else:
+                super().run()
+
+    def build_extension(self, ext):
+        if sys.platform == "darwin":
+            return
+
+        # these dirs will be created in build_py, so if you don't have
+        # any python sources to bundle, the dirs will be missing
+        build_temp = os.path.abspath(self.build_temp)
+        os.makedirs(build_temp, exist_ok=True)
+
+        try:
+            subprocess.check_output(
+                ["cmake", "-S", THIS_DIR, "-B", build_temp],
+                stderr=subprocess.STDOUT,
+            )
+            subprocess.check_output(
+                ["cmake", "--build", build_temp],
+                stderr=subprocess.STDOUT,
+            )
+        except subprocess.CalledProcessError as e:
+            logging.error(e.output.decode())
+            raise
+
+        libname = os.path.basename(self.get_ext_fullpath(ext.name))
+
+        target_dir = os.path.join(
+            os.path.dirname(self.get_ext_fullpath(ext.name)),
+            "mscclpp",
+        )
+
+        shutil.copy(
+            os.path.join(build_temp, "libmscclpp.so"),
+            target_dir,
+        )
+
+        shutil.copy(
+            os.path.join(build_temp, libname),
+            target_dir,
+        )
+
+
+setup(
+    name='mscclpp',
+    version='0.1.0',
+    description='Python bindings for mscclpp',
+    # packages=['mscclpp'],
+    package_dir={'': 'src'},
+    packages=find_packages(where='./src'),
+    ext_modules=[CustomExt('_py_mscclpp')],
+    cmdclass={
+        "build_ext": custom_build_ext,
+    },
+    install_requires=[
+        'torch',
+        'nanobind',
+    ],
+)

python/src/mscclpp/__init__.py

@@ -180,7 +180,7 @@ def connect_rank_from_address(
             rank=rank,
             world_size=world_size,
         )
-    
+
         for i in range(world_size):
             if i == rank:
                 continue

python/test.sh

@@ -2,11 +2,7 @@
 
 set -ex
 
-if ! [ -d build ] ; then
-  ./setup.sh
-fi
+pip install -e .
 
-cmake --build build
-
-cd build
-pytest -s mscclpp
+cd src
+pytest -vs mscclpp

src/.gitignore

@@ -0,0 +1,2 @@
+# Python in-place installs move the .so files into the source directories.
+*.so