diff --git a/python/llm/setup.py b/python/llm/setup.py index 9603cc2e..a8d91108 100644 --- a/python/llm/setup.py +++ b/python/llm/setup.py @@ -126,6 +126,17 @@ linux_binarys = [ "chatglm_C.cpython-39-x86_64-linux-gnu.so", ] +ext_lib_urls = [ + "https://github.com/analytics-zoo/jemalloc/releases/download/v5.3.0/libjemalloc.so", + "https://github.com/analytics-zoo/tcmalloc/releases/download/v2.10/libtcmalloc.so" +] + +ext_libs = [ + "libjemalloc.so", + "libtcmalloc.so" +] + + def get_llm_packages(): llm_packages = [] @@ -187,7 +198,7 @@ def obtain_lib_urls(): windows_binary_urls = get_urls_for_binary(date_urls, windows_binarys) lib_urls["Windows"] = list(windows_binary_urls.values()) linux_binary_urls = get_urls_for_binary(date_urls, linux_binarys) - lib_urls["Linux"] = list(linux_binary_urls.values()) + lib_urls["Linux"] = list(linux_binary_urls.values()) + ext_lib_urls return lib_urls @@ -208,7 +219,7 @@ def setup_package(): package_data["Windows"] = list(map(lambda x: os.path.join('libs', x), windows_binarys)) package_data["Linux"] = list(map(lambda x: os.path.join('libs', x), - linux_binarys)) + linux_binarys + ext_libs)) platform_name = None if "--win" in sys.argv: platform_name = "Windows" @@ -252,7 +263,8 @@ def setup_package(): raise FileNotFoundError( f'Could not find package dependency file: {file_path}') - all_requires = ['py-cpuinfo', 'protobuf'] + all_requires = ['py-cpuinfo', 'protobuf', + "intel-openmp; (platform_machine=='x86_64' or platform_machine == 'AMD64')"] all_requires += CONVERT_DEP # install with -f https://developer.intel.com/ipex-whl-stable-xpu