[feat] added version suffix to include cuda and torch version

python/setup.py

@@ -16,42 +16,87 @@
 def use_cxx11_abi():
     try:
         import torch
+
         return torch._C._GLIBCXX_USE_CXX11_ABI
     except ImportError:
         return False
-
+
+
 def get_torch_root():
     try:
         import torch
+
         return str(Path(torch.__file__).parent)
     except ImportError:
         return None
-
+
+
 def get_nccl_root():
     try:
         from nvidia import nccl
+
         return str(Path(nccl.__file__).parent)
     except ImportError:
         return None
 
-# ---- cmake extension ----
+
 def get_base_dir():
     return os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))
 
+
+def join_path(*paths):
+    return os.path.join(get_base_dir(), *paths)
+
+
+def get_version_suffix():
+    try:
+        import torch
+
+        version, other = torch.__version__.split("+")
+        major, minor, _ = version.split(".")
+        return f"{other}torch{major}.{minor}"
+    except ImportError:
+        return None
+
+
+def extract_version(file_path):
+    with open(file_path, "r") as f:
+        for line in f:
+            match = re.match(r"^__version__ = ['\"]([^'\"]*)['\"]", line)
+            if match:
+                return match.group(1)
+    raise RuntimeError("Unable to find version string.")
+
+
+def get_scalellm_version():
+    init_file = join_path("python", "scalellm", "__init__.py")
+    version = extract_version(init_file)
+    version_suffix = get_version_suffix()
+    if version_suffix is not None:
+        version += f"+{version_suffix}"
+    return version
+
+
+def read_readme() -> str:
+    p = os.path.join(get_base_dir(), "README.md")
+    if os.path.isfile(p):
+        return io.open(p, "r", encoding="utf-8").read()
+    else:
+        return ""
+
+
+# ---- cmake extension ----
 def get_cmake_dir():
     plat_name = sysconfig.get_platform()
     python_version = sysconfig.get_python_version().replace(".", "")
     dir_name = f"cmake.{plat_name}-{sys.implementation.name}-{python_version}"
+    version_suffix = get_version_suffix()
+    if version_suffix is not None:
+        dir_name += f"-{version_suffix}"
     cmake_dir = Path(get_base_dir()) / "python" / "build" / dir_name
     cmake_dir.mkdir(parents=True, exist_ok=True)
     return cmake_dir
 
-def read_readme() -> str:
-    p = os.path.join(get_base_dir(), "README.md")
-    if os.path.isfile(p):
-        return io.open(p, "r", encoding="utf-8").read()
-    else:
-        return ""
 
 # A CMakeExtension needs a sourcedir instead of a file list.
 # The name must be the _single_ output extension from the CMake build.
@@ -62,39 +107,45 @@ def __init__(self, name: str, path: str, sourcedir: str = "") -> None:
         self.sourcedir = os.fspath(Path(sourcedir).resolve())
         self.path = path
 
+
 class CMakeBuild(build_ext):
-    user_options = build_ext.user_options + \
-        [('base-dir=', None, 'base directory of Triton')]
+    user_options = build_ext.user_options + [
+        ("base-dir=", None, "base directory of ScaleLLM project"),
+    ]
 
     def initialize_options(self):
         build_ext.initialize_options(self)
         self.base_dir = get_base_dir()
 
     def finalize_options(self):
         build_ext.finalize_options(self)
-        
+
     def run(self):
         # check if cmake is installed
         try:
             out = subprocess.check_output(["cmake", "--version"])
         except OSError:
-            raise RuntimeError("CMake must be installed to build the following extensions: " +
-                               ", ".join(e.name for e in self.extensions))
-
-        match = re.search(r"version\s*(?P<major>\d+)\.(?P<minor>\d+)([\d.]+)?", out.decode())
+            raise RuntimeError(
+                "CMake must be installed to build the following extensions: "
+                + ", ".join(e.name for e in self.extensions)
+            )
+
+        match = re.search(
+            r"version\s*(?P<major>\d+)\.(?P<minor>\d+)([\d.]+)?", out.decode()
+        )
         cmake_major, cmake_minor = int(match.group("major")), int(match.group("minor"))
         if (cmake_major, cmake_minor) < (3, 18):
             raise RuntimeError("CMake >= 3.18.0 is required")
 
         # build extensions
         for ext in self.extensions:
             self.build_extension(ext)
-        
+
     def build_extension(self, ext: CMakeExtension):
         ninja_dir = shutil.which("ninja")
         # the output dir for the extension
         extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.path)))
-        
+
         # create build directory
         os.makedirs(self.build_temp, exist_ok=True)
 
@@ -106,11 +157,12 @@ def build_extension(self, ext: CMakeExtension):
 
         # python directories
         python_include_dir = sysconfig.get_path("platinclude")
-        cuda_architectures="80;89;90"
+        cuda_architectures = "80;89;90"
         cmake_args = [
-            "-G", "Ninja", # Ninja is much faster than make
-            "-DUSE_CCACHE=ON", # use ccache if available
-            f"-DCMAKE_MAKE_PROGRAM={ninja_dir}",     # pass in the ninja build path
+            "-G",
+            "Ninja",  # Ninja is much faster than make
+            "-DUSE_CCACHE=ON",  # use ccache if available
+            f"-DCMAKE_MAKE_PROGRAM={ninja_dir}",  # pass in the ninja build path
             "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON",
             f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}",
             "-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON",
@@ -125,7 +177,7 @@ def build_extension(self, ext: CMakeExtension):
         # (needed e.g. to build for ARM OSx on conda-forge)
         if "CMAKE_ARGS" in os.environ:
             cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item]
-                
+
         # check if torch binary is built with cxx11 abi
         if use_cxx11_abi():
             cmake_args += ["-DUSE_CXX11_ABI=ON"]
@@ -134,48 +186,46 @@ def build_extension(self, ext: CMakeExtension):
 
         build_args = ["--config", build_type]
         max_jobs = os.getenv("MAX_JOBS", str(os.cpu_count()))
-        build_args += ['-j' + max_jobs]
-        
+        build_args += ["-j" + max_jobs]
+
         env = os.environ.copy()
         LIBTORCH_ROOT = get_torch_root()
         if LIBTORCH_ROOT is not None:
             env["LIBTORCH_ROOT"] = LIBTORCH_ROOT
-            
+
         NCCL_ROOT = get_nccl_root()
         if NCCL_ROOT is not None:
             env["NCCL_ROOT"] = NCCL_ROOT
             env["NCCL_VERSION"] = "2"
-            
+
         # print cmake args
         print("CMake Args: ", cmake_args)
         print("Env: ", env)
-            
+
         cmake_dir = get_cmake_dir()
-        subprocess.check_call(["cmake", self.base_dir] + cmake_args, cwd=cmake_dir, env=env)
-
+        subprocess.check_call(
+            ["cmake", self.base_dir] + cmake_args, cwd=cmake_dir, env=env
+        )
+
         # add build target to speed up the build process
         build_args += ["--target", ext.name]
         subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=cmake_dir)
 
+
 # The information here can also be placed in setup.cfg - better separation of
 # logic and declaration, and simpler if you include description/version in a file.
 
-scalellm_package_data = [
-]
+scalellm_package_data = []
 
 setup(
     name="scalellm",
-    version="0.0.9",
+    version=get_scalellm_version(),
     license="Apache 2.0",
     author="ScaleLLM Team",
     description="A high-performance inference system for large language models.",
     long_description=read_readme(),
     url="https://github.com/vectorch-ai/ScaleLLM",
-    packages=[
-        "scalellm", 
-        "scalellm/serve",
-        "examples"
-    ],
+    packages=["scalellm", "scalellm/serve", "examples"],
     ext_modules=[CMakeExtension("_C", "scalellm/")],
     cmdclass={"build_ext": CMakeBuild},
     zip_safe=False,
@@ -186,15 +236,4 @@ def build_extension(self, ext: CMakeExtension):
     install_requires=[
         "torch",
     ],
-    extras_require={
-        "build": [
-            "cmake>=3.18",
-            "ninja",
-            "ccache"
-        ],
-        "tests": [
-            "pytest>=6.0",
-        ],
-    },
 )
-