From 384fc46585a83a33eb6dadcefa6eaac991730fb5 Mon Sep 17 00:00:00 2001 From: Houjiang Chen Date: Sat, 28 Dec 2019 20:53:21 +0800 Subject: [PATCH] Update setup py (#2532) * Update tensorflow cmakefile * Update setup.py --- CMakeLists.txt | 1 - cmake/third_party.cmake | 1 + cmake/third_party/tensorflow.cmake | 24 +++++++---- setup.py | 64 ++++++++++++++++++++++++++++++ setup.py.in | 51 ------------------------ 5 files changed, 82 insertions(+), 59 deletions(-) create mode 100644 setup.py delete mode 100644 setup.py.in diff --git a/CMakeLists.txt b/CMakeLists.txt index 5579b404f7..4702270f56 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -108,5 +108,4 @@ if (THIRD_PARTY) set(THIRD_PARTY OFF CACHE BOOL "" FORCE) else() include(oneflow) - configure_file(${PROJECT_SOURCE_DIR}/setup.py.in ${PROJECT_BINARY_DIR}/setup.py) endif() diff --git a/cmake/third_party.cmake b/cmake/third_party.cmake index 1ee8d40aad..e280e5c155 100644 --- a/cmake/third_party.cmake +++ b/cmake/third_party.cmake @@ -103,6 +103,7 @@ set(oneflow_third_party_dependencies cocoapi_copy_libs_to_destination half_copy_headers_to_destination json_copy_headers_to_destination + tensorflow_copy_libs_to_destination ) include_directories( diff --git a/cmake/third_party/tensorflow.cmake b/cmake/third_party/tensorflow.cmake index 39bf6191cb..45b491c6cd 100644 --- a/cmake/third_party/tensorflow.cmake +++ b/cmake/third_party/tensorflow.cmake @@ -28,10 +28,12 @@ set(TENSORFLOW_PROJECT tensorflow) set(TENSORFLOW_GIT_URL https://github.com/tensorflow/tensorflow.git) #set(TENSORFLOW_GIT_TAG master) set(TENSORFLOW_GIT_TAG 80c04b80ad66bf95aa3f41d72a6bba5e84a99622) -set(TENSORFLOW_SOURCES_DIR ${THIRD_PARTY_DIR}/tensorflow) +set(TENSORFLOW_SOURCES_DIR ${CMAKE_CURRENT_BINARY_DIR}/third_party/tensorflow) set(TENSORFLOW_SRCS_DIR ${TENSORFLOW_SOURCES_DIR}/src/tensorflow) set(TENSORFLOW_INC_DIR ${TENSORFLOW_SOURCES_DIR}/src/tensorflow) +set(TENSORFLOW_INSTALL_DIR ${THIRD_PARTY_DIR}/tensorflow) + set(PATCHES_DIR ${PROJECT_SOURCE_DIR}/oneflow/xrt/patches) set(TENSORFLOW_JIT_DIR ${TENSORFLOW_SRCS_DIR}/tensorflow/compiler/jit) @@ -51,13 +53,9 @@ list(APPEND TENSORFLOW_XLA_INCLUDE_DIR ${THIRD_SNAPPY_DIR} ) include_directories(${TENSORFLOW_XLA_INCLUDE_DIR}) - list(APPEND TENSORFLOW_XLA_LIBRARIES libtensorflow_framework.so.1) list(APPEND TENSORFLOW_XLA_LIBRARIES libxla_core.so) -link_directories( - ${TENSORFLOW_SRCS_DIR}/bazel-bin/tensorflow - ${TENSORFLOW_SRCS_DIR}/bazel-bin/tensorflow/compiler/jit/xla_lib -) +link_directories(${TENSORFLOW_INSTALL_DIR}/lib) if (THIRD_PARTY) ExternalProject_Add(${TENSORFLOW_PROJECT} @@ -70,9 +68,21 @@ if (THIRD_PARTY) bazel build ${TENSORFLOW_BUILD_CMD} -j 20 //tensorflow/compiler/jit/xla_lib:libxla_core.so INSTALL_COMMAND "" ) -endif(THIRD_PARTY) set(TENSORFLOW_XLA_FRAMEWORK_LIB ${TENSORFLOW_SRCS_DIR}/bazel-bin/tensorflow/libtensorflow_framework.so.1) set(TENSORFLOW_XLA_CORE_LIB ${TENSORFLOW_SRCS_DIR}/bazel-bin/tensorflow/compiler/jit/xla_lib/libxla_core.so) +add_custom_target(tensorflow_create_library_dir + COMMAND ${CMAKE_COMMAND} -E make_directory ${TENSORFLOW_INSTALL_DIR}/lib + DEPENDS ${TENSORFLOW_PROJECT}) + +add_custom_target(tensorflow_copy_libs_to_destination + COMMAND ${CMAKE_COMMAND} -E copy_if_different + ${TENSORFLOW_XLA_FRAMEWORK_LIB} ${TENSORFLOW_XLA_CORE_LIB} ${TENSORFLOW_INSTALL_DIR}/lib + COMMAND ${CMAKE_COMMAND} -E create_symlink + ${TENSORFLOW_INSTALL_DIR}/lib/libtensorflow_framework.so.1 + ${TENSORFLOW_INSTALL_DIR}/lib/libtensorflow_framework.so + DEPENDS tensorflow_create_library_dir) +endif(THIRD_PARTY) + endif(WITH_XLA) diff --git a/setup.py b/setup.py new file mode 100644 index 0000000000..fd75712e2b --- /dev/null +++ b/setup.py @@ -0,0 +1,64 @@ +from __future__ import absolute_import + +import os +import re +import sys +import argparse +import shutil + +from setuptools import find_packages +from setuptools import setup +from setuptools.dist import Distribution + +parser = argparse.ArgumentParser() +parser.register("type", "bool", lambda v: v.lower() == "true") +parser.add_argument( + "--with_xla", + type='bool', + default=False, + help="Package xla libraries if true, otherwise not." +) +args, remain_args = parser.parse_known_args() +sys.argv = ['setup.py'] + remain_args + + +REQUIRED_PACKAGES = [ + 'numpy', + 'protobuf', +] + + +class BinaryDistribution(Distribution): + + def has_ext_modules(self): + return True + +packages = find_packages("build/python_scripts") +package_dir = { + '':'build/python_scripts', +} +package_data = {'oneflow': ['_oneflow_internal.so']} + +if args.with_xla: + packages += ['oneflow.libs'] + package_dir['oneflow.libs'] = 'third_party/tensorflow/lib' + package_data['oneflow.libs'] = ['libtensorflow_framework.so.1', 'libxla_core.so'] + # Patchelf >= 0.9 is required. + oneflow_internal_so = "build/python_scripts/oneflow/_oneflow_internal.so" + rpath = os.popen("patchelf --print-rpath " + oneflow_internal_so).read() + command = "patchelf --set-rpath '$ORIGIN/:$ORIGIN/libs/:%s' %s" % \ + (rpath.strip(), oneflow_internal_so) + if os.system(command) != 0: + raise Exception("Patchelf set rpath failed. command is: %s" % command) + +setup( + name='oneflow', + version='0.0.1', + url='https://www.oneflow.org/', + install_requires=REQUIRED_PACKAGES, + packages=packages, + package_dir=package_dir, + package_data=package_data, + zip_safe=False, + distclass=BinaryDistribution, +) diff --git a/setup.py.in b/setup.py.in deleted file mode 100644 index fbf5ad1d3d..0000000000 --- a/setup.py.in +++ /dev/null @@ -1,51 +0,0 @@ -from __future__ import absolute_import - -import os -import re -import sys -import shutil - -from setuptools import find_packages -from setuptools import setup -from setuptools.dist import Distribution - -REQUIRED_PACKAGES = [ - 'numpy', - 'protobuf', -] - - -class BinaryDistribution(Distribution): - - def has_ext_modules(self): - return True - -packages = find_packages("build/python_scripts") -package_dir = { - '':'build/python_scripts', -} -package_data['oneflow'] = ['_oneflow_internal.so'] - -if '${WITH_XLA}' == 'ON': - packages += ['oneflow.libs'] - libs_path = 'python_scripts/oneflow/libs' - package_dir['oneflow.libs'] = libs_path - package_data['oneflow.libs'] = ['libtensorflow_framework.so.1', 'libxla_core.so'] - - shutil.copy('${TENSORFLOW_XLA_FRAMEWORK_LIB}', libs_path) - shutil.copy('${TENSORFLOW_XLA_CORE_LIB}', libs_path) - command = "patchelf --set-rpath '$ORIGIN/' ${TENSORFLOW_XLA_CORE_LIB}" - if os.system(command) != 0: - raise Exception("patch xla failed, command: %s" % command) - -setup( - name='oneflow', - version='0.0.1', - url='https://www.oneflow.org/', - install_requires=REQUIRED_PACKAGES, - packages=packages, - package_dir=package_dir, - package_data=package_data, - zip_safe=False, - distclass=BinaryDistribution, -) -- GitLab