feat: install nccl deps in bagua core and add generated __version__ v…

…ariable

bagua-core-py/src/lib.rs

@@ -308,7 +308,7 @@ impl BaguaBucketPy {
 }
 
 #[pymodule]
-fn bagua_core_py(_py: Python, m: &PyModule) -> PyResult<()> {
+fn bagua_core(_py: Python, m: &PyModule) -> PyResult<()> {
     tracing_subscriber::fmt()
         .with_env_filter(tracing_subscriber::EnvFilter::from_env("LOG_LEVEL"))
         .init();

bagua-core-py/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.1.dev1"

bagua_install_library/install_library.py

@@ -94,7 +94,7 @@ def install_lib(cuda, prefix, library):
             )
         )
     if prefix is None:
-        prefix = os.path.expanduser("~/.bagua/cuda_lib")
+        prefix = os.path.expanduser("~/.bagua_core/cuda_lib")
     destination = calculate_destination(prefix, cuda, library, record[library])
 
     if os.path.exists(destination):

pyproject.toml

@@ -3,3 +3,5 @@ requires = ["setuptools", "wheel", "setuptools-rust", "colorama", "tqdm", "setup
 
 [tool.setuptools_scm]
 local_scheme = "no-local-version"
+write_to = "bagua-core-py/version.py"
+write_to_template = "__version__ = \"{version}\""

python/bagua_core/__init__.py

@@ -0,0 +1,4 @@
+from . import _environment
+_environment._preload_libraries()
+from .version import __version__
+from .bagua_core import *

python/bagua_core/_environment.py

@@ -0,0 +1,8 @@
+import ctypes
+import os
+
+
+def _preload_libraries():
+    cwd = os.path.dirname(os.path.abspath(__file__))
+    libnccl_path = os.path.join(cwd, ".data", "lib", "libnccl.so")
+    ctypes.CDLL(libnccl_path)

setup.py

@@ -39,7 +39,7 @@ def install_dependency_library():
         "--cuda",
         nvcc_version,
         "--prefix",
-        os.path.join(cwd, "bagua"),
+        os.path.join(cwd, "python/bagua_core"),
     ]
     install_library.main(args)
 
@@ -63,16 +63,19 @@ def install_dependency_library():
         url="https://github.com/BaguaSys/bagua-core",
         python_requires=">=3.6",
         description="Core communication lib for Bagua.",
+        package_dir={"": "python/"},
+        packages=find_packages("python/"),
+        package_data={"": [".data/lib/libnccl.so"]},
         rust_extensions=[
             RustExtension(
-                "bagua_core_py",
+                "bagua_core.bagua_core",
                 path="bagua-core-py/Cargo.toml",
                 binding=Binding.PyO3,
                 native=True,
             )
         ],
         author="Kuaishou AI Platform & DS3 Lab",
         author_email="admin@mail.xrlian.com",
-        package_data={"": [".data/lib/libnccl.so"]},
         install_requires=[],
+        zip_safe=False,
     )