Reorgnize and deduplicate files (#4)

README.md

@@ -41,10 +41,10 @@ FLUX_FINAL_RELEASE=0 ./build.sh --arch 80 --package
 PYTHONPATH=./python:$PYTHONPATH python3 test/test_gemm_only.py 4096 12288 6144 --dtype=float16
 
 # gemm fused with reduce-scatter
-./launch.sh test/test_gemm_rs.py 4096 12288 49152 --dtype=float16 --iters=10
+./scripts/launch.sh test/test_gemm_rs.py 4096 12288 49152 --dtype=float16 --iters=10
 
 # all-gather fused with gemm
-./launch.sh test/test_ag_kernel.py 4096 49152 12288 --dtype=float16 --iters=10
+./scripts/launch.sh test/test_ag_kernel.py 4096 49152 12288 --dtype=float16 --iters=10
 ```
 
 

python/flux/__init__.py

@@ -15,6 +15,9 @@
 #
 ################################################################################
 
+# version info
+__version__ = "1.0.0"
+
 from .cpp_mod import *
 from .ag_gemm import *
 

setup.py

@@ -18,20 +18,101 @@
 import glob
 import os
 import re
+import shutil
+import subprocess
 from pathlib import Path
+from typing import Optional, Tuple
 
 import setuptools
 from torch.utils.cpp_extension import BuildExtension
 
-from gen_version import generate_versoin_file, check_final_release
+CUR_DIR = os.path.dirname(os.path.realpath(__file__))
+
+
+def _check_env_option(opt, default=""):
+    return os.getenv(opt, default).upper() in ["ON", "1", "YES", "TRUE"]
+
+
+def check_final_release():
+    return _check_env_option("FLUX_FINAL_RELEASE", "1")
+
+
+def get_git_commit(src_dir):
+    try:
+        return (
+            subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=src_dir)
+            .decode("ascii")
+            .strip()
+        )
+    except Exception:
+        return "unknown"
+
+
+def cuda_version() -> Tuple[int, ...]:
+    """CUDA Toolkit version as a (major, minor) by nvcc --version"""
+
+    # Try finding NVCC
+    nvcc_bin: Optional[Path] = None
+    if nvcc_bin is None and os.getenv("CUDA_HOME"):
+        # Check in CUDA_HOME
+        cuda_home = Path(os.getenv("CUDA_HOME"))
+        nvcc_bin = cuda_home / "bin" / "nvcc"
+    if nvcc_bin is None:
+        # Check if nvcc is in path
+        nvcc_bin = shutil.which("nvcc")
+        if nvcc_bin is not None:
+            nvcc_bin = Path(nvcc_bin)
+    if nvcc_bin is None:
+        # Last-ditch guess in /usr/local/cuda
+        cuda_home = Path("/usr/local/cuda")
+        nvcc_bin = cuda_home / "bin" / "nvcc"
+    if not nvcc_bin.is_file():
+        raise FileNotFoundError(f"Could not find NVCC at {nvcc_bin}")
+
+    # Query NVCC for version info
+    output = subprocess.run(
+        [nvcc_bin, "-V"],
+        capture_output=True,
+        check=True,
+        universal_newlines=True,
+    )
+    match = re.search(r"release\s*([\d.]+)", output.stdout)
+    version = match.group(1).split(".")
+    return tuple(int(v) for v in version)
+
+
+def get_flux_version(version_info, *, dev=False):
+    with open(version_info) as f:
+        (version,) = re.findall('__version__ = "(.*)"', f.read())
+    cuda_version_major, cuda_version_minor = cuda_version()
+    version = version + f"+cu{cuda_version_major}{cuda_version_minor}"
+    if dev:
+        commit_id = get_git_commit(CUR_DIR)
+
+        version += ".dev{}".format(commit_id[:8])
+    # version = version + (f'.{os.getenv("ARCH")}' if os.getenv("ARCH") else "")
+    return version
+
+
+def generate_versoin_file(version, version_file, *, dev=False):
+    flux_ver = get_flux_version(version, dev=dev)
+
+    with open(version_file, "w") as f:
+        f.write("__version__ = '{}'\n".format(flux_ver))
+        f.write("git_version = {}\n".format(repr(get_git_commit(CUR_DIR))))
+        cuda_version_major, cuda_version_minor = cuda_version()
+        f.write("cuda = {}.{}\n".format(cuda_version_major, cuda_version_minor))
+
+    return flux_ver
+
 
 # Project directory root
 root_path: Path = Path(__file__).resolve().parent
 
-version_txt = os.path.join(root_path, "version.txt")
+version = os.path.join(root_path, "python/flux/__init__.py")
 version_file = os.path.join(root_path, "python/flux/version.py")
 is_dev = not check_final_release()
-flux_version = generate_versoin_file(version_txt, version_file, dev=is_dev)
+flux_version = generate_versoin_file(version, version_file, dev=is_dev)
 
 
 def pathlib_wrapper(func):