Commit 26298255 authored by lijian6's avatar lijian6
Browse files

Add flag for nvshmem/rocshmem on release name.



End of hash, 'a' means rocshmem, 'b' means nvshmem.
Signed-off-by: lijian6's avatarlijian <lijian6@sugon.com>
parent 3c57a9d8
...@@ -139,6 +139,11 @@ hipcc -Wno-unused-result -Wsign-compare -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack ...@@ -139,6 +139,11 @@ hipcc -Wno-unused-result -Wsign-compare -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack
# build whl # build whl
echo "Using Python: $(which python3)" echo "Using Python: $(which python3)"
python3 --version python3 --version
python setup.py bdist_wheel if [ "$USE_NVSHMEM" == "ON" ]; then
python setup.py bdist_wheel --shmem=nv
fi
if [ "$USE_ROCSHMEM" == "ON" ]; then
python setup.py bdist_wheel --shmem=rocm
fi
echo "✅ Build complete:" echo "✅ Build complete:"
ls -lh dist/ ls -lh dist/
import os import os, sys
import subprocess import subprocess
import setuptools import setuptools
from torch.utils.cpp_extension import BuildExtension, CUDAExtension from torch.utils.cpp_extension import BuildExtension, CUDAExtension
...@@ -6,6 +6,23 @@ from datetime import datetime ...@@ -6,6 +6,23 @@ from datetime import datetime
date_tag = datetime.now().strftime("%Y%m%d") date_tag = datetime.now().strftime("%Y%m%d")
# -----------------------
# 解析自定义参数
# -----------------------
shmem = None
other = []
for arg in sys.argv:
if arg.startswith("--shmem="):
shmem = arg.split("=", 1)[1]
if shmem == "rocm":
shmem = "a"
elif shmem == "nv":
shmem = "b"
else:
other.append(arg)
sys.argv = other
if __name__ == '__main__': if __name__ == '__main__':
try: try:
cmd = ['git', 'rev-parse', '--short', 'HEAD'] cmd = ['git', 'rev-parse', '--short', 'HEAD']
...@@ -15,7 +32,7 @@ if __name__ == '__main__': ...@@ -15,7 +32,7 @@ if __name__ == '__main__':
setuptools.setup( setuptools.setup(
name='deep_ep', name='deep_ep',
version='1.0.0' + revision + '.' + date_tag, version='1.0.0' + revision + shmem + '.' + date_tag,
packages=setuptools.find_packages(include=['deep_ep']), packages=setuptools.find_packages(include=['deep_ep']),
include_package_data=True, include_package_data=True,
package_data={"deep_ep": ["deep_ep_cpp.cpython-310-x86_64-linux-gnu.so"]}, package_data={"deep_ep": ["deep_ep_cpp.cpython-310-x86_64-linux-gnu.so"]},
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment