build: build_wheel tools update (#82)

Author: CyCle1024

.github/workflows/main.yml

@@ -51,6 +51,7 @@ jobs:
         run: |
           set -ex
           source /usr/local/Ascend/ascend-toolkit/set_env.sh
+          source /usr/local/Ascend/nnal/atb/set_env.sh
           cd ${{ env.CI_PATH }}
           pip uninstall -y dlinfer-ascend
           pip install -r requirements/ascend/full.txt

CMakeLists.txt

@@ -25,8 +25,4 @@ elseif(NOT DEVICE IN_LIST SUPPORTED_DEVICE)
 endif()
 
 add_subdirectory(dlinfer/vendor/${DEVICE})
-
-if("$ENV{DICP}" STREQUAL "true")
-  add_subdirectory(dlinfer/graph/dicp/vendor/AtbGraph/codegen/runtime)
-endif()
-
+add_subdirectory(dlinfer/graph/dicp/vendor)

dlinfer/graph/dicp/vendor/AtbGraph/codegen/runtime/CMakeLists.txt

@@ -67,12 +67,14 @@ set(COMPILE_OPTIONS
 set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -O2")
 
 add_library(dicp_model SHARED ${SOURCES})
-set_target_properties(dicp_model PROPERTIES
-    OUTPUT_NAME "dicp_model"
-    LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/../
-)
 
 target_compile_options(dicp_model PRIVATE ${COMPILE_OPTIONS})
 
 target_link_libraries(dicp_model PRIVATE
     atb torch c10 torch_cpu torch_python torch_npu opapi)
+
+file(RELATIVE_PATH OUTPUT_LIB_RELATIVE_PATH "${CMAKE_SOURCE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}/../")
+install(
+    TARGETS dicp_model
+    DESTINATION ${OUTPUT_LIB_RELATIVE_PATH}
+)

dlinfer/graph/dicp/vendor/CMakeLists.txt

@@ -0,0 +1,3 @@
+if (DEVICE STREQUAL "ascend")
+    add_subdirectory(AtbGraph/codegen/runtime)
+endif()

dlinfer/vendor/ascend/CMakeLists.txt

@@ -8,7 +8,8 @@ execute_process(
 )
 
 execute_process(
-    COMMAND bash -c "pip show torch-npu | awk '/^Location:/{ORS=\"\"; print \$2\"/torch_npu\"}'"
+    COMMAND python -c "from importlib.metadata import distribution; \
+    print(str(distribution('torch_npu').locate_file('torch_npu')), end='')"
     OUTPUT_VARIABLE Torch_npu_ROOT
 )
 
@@ -62,7 +63,8 @@ target_link_libraries(
     PUBLIC ${HCCL_LIBRARY}
 )
 
+file(RELATIVE_PATH OUTPUT_LIB_RELATIVE_PATH "${CMAKE_SOURCE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}")
 install(
     TARGETS ${PROJECT_NAME}
-    DESTINATION ./dlinfer/vendor/ascend
+    DESTINATION ${OUTPUT_LIB_RELATIVE_PATH}
 )

pyproject.toml

@@ -2,5 +2,6 @@
 requires = [
     "pyyaml",
     "scikit-build",
+    "setuptools==69.5.1",
 ]
 build-backend = "setuptools.build_meta"

requirements/ascend/torch.txt

@@ -1,5 +1,5 @@
-torch==2.1.0
-torchvision==0.16.0
-torch-npu==2.1.0.post6
+torch==2.3.1
+torchvision==0.18.1
+torch-npu==2.3.1
 numpy<2.0.0
 pyyaml

scripts/build_wheel.sh

@@ -3,7 +3,7 @@ set -e
 
 REPO_ROOT=$(cd $(dirname $(dirname $0)); pwd)
 pip install -U build
-rm -rf ${REPO_ROOT}/_skbuild ${REPO_ROOT}/dlinfer.egg*
+rm -rf ${REPO_ROOT}/_skbuild ${REPO_ROOT}/dlinfer*.egg*
 export DEVICE=${DEVICE:-ascend}
 python -m build \
     -C="--build-option=--plat-name" \

scripts/build_wheel_allpy.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+set -e
+
+eval "$(conda shell.bash hook)"
+
+REPO_ROOT=$(cd $(dirname $(dirname $0)); pwd)
+cd ${REPO_ROOT}
+
+# PY_VERSION_LIST=("3.8" "3.9" "3.10")
+PY_VERSION_LIST=("3.10")
+for PY_VERSION in ${PY_VERSION_LIST[@]}; do
+    echo start building wheels for python${PY_VERSION}
+    PY_VERSION_NAME=${PY_VERSION/./}
+    ENV_NAME=dlinfer_build_py${PY_VERSION_NAME}
+    conda env remove -n ${ENV_NAME} -y
+    conda create -n ${ENV_NAME} python=${PY_VERSION} -y
+    conda activate ${ENV_NAME}
+    pip install -U build 
+    bash ${REPO_ROOT}/scripts/build_wheel.sh
+    conda deactivate
+    conda env remove -n ${ENV_NAME} -y
+    echo end building wheels for python${PY_VERSION}
+done

setup.py

@@ -5,7 +5,7 @@
 from skbuild import setup
 
 
-VERSION = "0.1.0post1"
+VERSION = "0.1.1"
 
 vendor_dispatch_key_map = {
     "ascend": "PrivateUse1",
@@ -71,6 +71,14 @@ def get_requirements(file_name):
     return requirements
 
 
+def get_vendor_excludes():
+    device = get_device()
+    exclude_vendors = [
+        name for name in vendor_dispatch_key_map.keys() if name != device
+    ]
+    return [f"dlinfer.vendor.{name}" for name in exclude_vendors]
+
+
 def get_entry_points():
     device = get_device()
     if device == "ascend":
@@ -91,7 +99,7 @@ def main():
         long_description=get_readme(),
         long_description_content_type="text/markdown",
         url="https://github.com/DeepLink-org/dlinfer",
-        packages=find_packages(),
+        packages=find_packages(exclude=get_vendor_excludes()),
         package_data=get_package_data(),
         exclude_package_data={"": ["tests/*"]},
         cmake_args=get_cmake_args(),