Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / multicore / platform / torch / setup.py: 0%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Build script for hyper_parallel MoE-FFN PyTorch extension.""" 

16 

17import os 

18import glob 

19import tarfile 

20from setuptools import setup, find_packages 

21from torch.utils.cpp_extension import BuildExtension 

22 

23import torch_npu 

24from torch_npu.utils.cpp_extension import NpuExtension 

25 

26PYTORCH_NPU_INSTALL_PATH = os.path.dirname(os.path.abspath(torch_npu.__file__)) 

27USE_NINJA = os.getenv('USE_NINJA') == '1' 

28BASE_DIR = os.path.dirname(os.path.realpath(__file__)) 

29 

30source_files = glob.glob(os.path.join(BASE_DIR, "csrc", "*.cpp"), recursive=True) 

31 

32 

33def _find_prebuild_vendor_libdirs(): 

34 """Return (fwd_libdir, bwd_libdir) from env vars or the prebuild directory. 

35 

36 Resolution order: 

37 1. CANN_VENDOR_FWD_LIBDIR / CANN_VENDOR_BWD_LIBDIR (explicit per-op) 

38 2. CANN_VENDOR_LIBDIR (legacy single-lib; used for both) 

39 3. prebuild/multicore_moe_ffn/vendors/... (auto-detect + extract) 

40 """ 

41 fwd = os.environ.get("CANN_VENDOR_FWD_LIBDIR", 

42 os.environ.get("CANN_VENDOR_LIBDIR", "")) 

43 bwd = os.environ.get("CANN_VENDOR_BWD_LIBDIR", 

44 os.environ.get("CANN_VENDOR_LIBDIR", "")) 

45 if fwd and bwd: 

46 return fwd, bwd 

47 # Auto-detect: prebuild dir is 2 levels up from platform/torch/ (→ multicore/) 

48 prebuild_dir = os.path.normpath( 

49 os.path.join(BASE_DIR, "../../prebuild/multicore_moe_ffn")) 

50 tarball = prebuild_dir + ".tar.gz" 

51 if not os.path.isdir(prebuild_dir) and os.path.isfile(tarball): 

52 print(f"[setup] Extracting prebuild: {tarball}") 

53 with tarfile.open(tarball) as tf: 

54 tf.extractall(os.path.dirname(prebuild_dir)) 

55 vendors = os.path.join(prebuild_dir, "vendors") 

56 fwd = fwd or os.path.join(vendors, "multicore_moe_ffn_nn", "op_api", "lib") 

57 bwd = bwd or os.path.join(vendors, "multicore_moe_ffn_grad_nn", "op_api", "lib") 

58 return fwd, bwd 

59 

60 

61_FWD_LIBDIR, _BWD_LIBDIR = _find_prebuild_vendor_libdirs() 

62 

63 

64def _write_vendor_env_to_bashrc(libdirs): 

65 """Append ASCEND_CUSTOM_OPP_PATH / LD_LIBRARY_PATH exports to ~/.bashrc. 

66 

67 Idempotent: each line is only appended if not already present. 

68 Required because g_custom_lib_path in libopapi.so is a const file-scope 

69 global initialised once when libopapi.so loads (triggered by 

70 ``import torch_npu``). Setting ASCEND_CUSTOM_OPP_PATH from Python at 

71 import time is too late if torch_npu was already imported. Persisting the 

72 variable in ~/.bashrc ensures it is present before Python starts. 

73 """ 

74 bashrc = os.path.expanduser("~/.bashrc") 

75 try: 

76 with open(bashrc, "r", encoding="utf-8") as f: 

77 content = f.read() 

78 except FileNotFoundError: 

79 content = "" 

80 

81 added = [] 

82 for libdir in dict.fromkeys(libdirs): 

83 if not libdir: 

84 continue 

85 root = libdir 

86 for sfx in ("/op_api/lib/", "/op_api/lib"): 

87 if root.endswith(sfx): 

88 root = root[: -len(sfx)] 

89 break 

90 opp_line = f'export ASCEND_CUSTOM_OPP_PATH="{root}:${{ASCEND_CUSTOM_OPP_PATH}}"' 

91 ld_line = f'export LD_LIBRARY_PATH="{libdir}:${{LD_LIBRARY_PATH}}"' 

92 for line in (opp_line, ld_line): 

93 if line not in content: 

94 content += f"\n{line}" 

95 added.append(line) 

96 

97 if added: 

98 with open(bashrc, "w", encoding="utf-8") as f: 

99 f.write(content) 

100 print(f"[setup] Appended to {bashrc}:") 

101 for line in added: 

102 print(f" {line}") 

103 print(f"[setup] Run 'source {bashrc}' or start a new shell to apply.") 

104 else: 

105 print(f"[setup] {bashrc} already up to date.") 

106 

107 

108def _print_env_banner(libdirs): 

109 """Print a copy-pasteable environment setup banner.""" 

110 sep = "=" * 64 

111 print(f"\n{sep}") 

112 print(" hyper_parallel multicore: environment setup") 

113 print(sep) 

114 for libdir in dict.fromkeys(libdirs): 

115 if not libdir: 

116 continue 

117 root = libdir 

118 for sfx in ("/op_api/lib/", "/op_api/lib"): 

119 if root.endswith(sfx): 

120 root = root[: -len(sfx)] 

121 break 

122 print(f' export ASCEND_CUSTOM_OPP_PATH="{root}:${{ASCEND_CUSTOM_OPP_PATH}}"') 

123 print(f' export LD_LIBRARY_PATH="{libdir}:${{LD_LIBRARY_PATH}}"') 

124 print(sep + "\n") 

125 

126 

127# Add -Wl,-rpath for each vendor lib dir (deduplicated) so the dynamic linker 

128# can find libcust_opapi.so at runtime without LD_LIBRARY_PATH. 

129_rpath_args = [] 

130for _ldir in dict.fromkeys([_FWD_LIBDIR, _BWD_LIBDIR]): 

131 if _ldir and os.path.isdir(_ldir): 

132 _rpath_args.append(f"-Wl,-rpath,{_ldir}") 

133 

134exts = [] 

135ext = NpuExtension( 

136 name="hyper_parallel_multicore_moe_ffn_pta", 

137 sources=source_files, 

138 extra_compile_args=[ 

139 '-I' + os.path.join(PYTORCH_NPU_INSTALL_PATH, "include/third_party/acl/inc"), 

140 '-I' + os.path.join(PYTORCH_NPU_INSTALL_PATH, "include/third_party/op-plugin"), 

141 '-I' + os.path.join(PYTORCH_NPU_INSTALL_PATH, "include/third_party/op-plugin/op_plugin/include"), 

142 ], 

143 extra_link_args=_rpath_args, 

144) 

145exts.append(ext) 

146 

147setup( 

148 name="hyper_parallel_multicore_moe_ffn_pta", 

149 version='1.0', 

150 keywords='hyper_parallel_multicore_moe_ffn_pta', 

151 ext_modules=exts, 

152 packages=find_packages(), 

153 cmdclass={"build_ext": BuildExtension.with_options(use_ninja=USE_NINJA)}, 

154) 

155 

156_write_vendor_env_to_bashrc([_FWD_LIBDIR, _BWD_LIBDIR]) 

157_print_env_banner([_FWD_LIBDIR, _BWD_LIBDIR])