Python: Add version string & wheel build command (#1985)

* add version string & wheel build command * Update Python version in CI script --------- Co-authored-by: Georgy Evtushenko <[email protected]>
NVIDIA · Jul 16, 2024 · cb7e845 · cb7e845
1 parent 99e7d60
commit cb7e845
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 2 deletions.
diff --git a/ci/update_version.sh b/ci/update_version.sh
@@ -19,6 +19,8 @@ done
 major="$1"
 minor="$2"
 patch="$3"
+pymajor="0"
+pyminor="1"
 
 if [ -z "$major" ] || [ -z "$minor" ] || [ -z "$patch" ]; then
     echo "Usage: $0 [--dry-run] <major> <minor> <patch>"
@@ -34,6 +36,7 @@ CUB_CMAKE_VERSION_FILE="cub/cub/cmake/cub-config-version.cmake"
 LIBCUDACXX_CMAKE_VERSION_FILE="libcudacxx/lib/cmake/libcudacxx/libcudacxx-config-version.cmake"
 THRUST_CMAKE_VERSION_FILE="thrust/thrust/cmake/thrust-config-version.cmake"
 CUDAX_CMAKE_VERSION_FILE="cudax/lib/cmake/cudax/cudax-config-version.cmake"
+PYCUDA_VERSION_FILE="python/cuda/cuda/cooperative/_version.py"
 
 # Calculated version codes
 new_cccl_version=$((major * 1000000 + minor * 1000 + patch))     # MMMmmmppp
@@ -99,6 +102,8 @@ update_file "$CUDAX_CMAKE_VERSION_FILE" "set(cudax_VERSION_MAJOR \([0-9]\+\))" "
 update_file "$CUDAX_CMAKE_VERSION_FILE" "set(cudax_VERSION_MINOR \([0-9]\+\))" "set(cudax_VERSION_MINOR $minor)"
 update_file "$CUDAX_CMAKE_VERSION_FILE" "set(cudax_VERSION_PATCH \([0-9]\+\))" "set(cudax_VERSION_PATCH $patch)"
 
+update_file "$PYCUDA_VERSION_FILE" "^__version__ = \"\([0-9.]\+\)\"" "__version__ = \"$pymajor.$pyminor.$major.$minor.$patch\""
+
 if [ "$DRY_RUN" = true ]; then
     echo "Dry run completed. No changes made."
 else

diff --git a/python/cuda/cuda/cooperative/__init__.py b/python/cuda/cuda/cooperative/__init__.py
@@ -3,3 +3,4 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 import cuda.cooperative.experimental
+from cuda.cooperative._version import __version__
diff --git a/python/cuda/cuda/cooperative/_version.py b/python/cuda/cuda/cooperative/_version.py
@@ -0,0 +1,7 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# This file is generated by ci/update_version.sh
+# Do not edit this file manually.
+__version__ = "0.1.2.6.0"
diff --git a/python/cuda/setup.py b/python/cuda/setup.py
@@ -6,8 +6,10 @@
 import os
 import glob
 import shutil
+
 from setuptools import Command, setup, find_packages, find_namespace_packages
 from setuptools.command.build_py import build_py
+from wheel.bdist_wheel import bdist_wheel
 
 
 project_path = os.path.abspath(os.path.dirname(__file__))
@@ -17,13 +19,25 @@
     ['libcudacxx', 'include'],
     ['thrust', 'thrust']
 ]
+with open(os.path.join(project_path, 'cuda', 'cooperative', '_version.py')) as f:
+    exec(f.read())
+ver = __version__
+del __version__
+
 
 class CustomBuildCommand(build_py):
     def run(self):
         self.run_command('package_cccl')
         build_py.run(self)
 
 
+class CustomWheelBuild(bdist_wheel):
+
+    def run(self):
+        self.run_command('package_cccl')
+        super().run()
+
+
 class PackageCCCLCommand(Command):
     description = 'Generate additional files'
     user_options = []
@@ -46,7 +60,7 @@ def run(self):
 
 setup(
     name="cuda-cooperative",
-    version="0.1.0",  # TODO Read from CCCL version
+    version=ver,
     description="Experimental Core Library for CUDA Python",
     author="NVIDIA Corporation",
     classifiers=[
@@ -68,7 +82,8 @@ def run(self):
     },
     cmdclass={
         'package_cccl': PackageCCCLCommand,
-        'build_py': CustomBuildCommand
+        'build_py': CustomBuildCommand,
+        'bdist_wheel': CustomWheelBuild,
     },
     include_package_data=True
 )