chore: add option to install py module into scalellm folder (#438)

Author: guocuimi

CMakeLists.txt

@@ -5,6 +5,7 @@ option(USE_CCACHE "Attempt using CCache to wrap the compilation" ON)
 option(USE_MANYLINUX "Build for manylinux" OFF)
 
 option(BUILD_NVBENCH "Build the nvbench binary" OFF)
+option(INSTALL_PY_MODULE "Install python module to scalellm directory" OFF)
 
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)

scalellm/CMakeLists.txt

@@ -11,7 +11,7 @@ pybind_extension(
     csrc/llm_handler.cpp
     csrc/module.cpp
     csrc/kernels.cu
-  DEFINES 
+  DEFINES
     PYBIND11_DETAILED_ERROR_MESSAGES=1
   LINKDIRS
     ${TORCH_INSTALL_PREFIX}/lib
@@ -25,3 +25,12 @@ pybind_extension(
     glog::glog
     Python::Module
 )
+
+if (INSTALL_PY_MODULE)
+  add_custom_command(TARGET _C POST_BUILD
+      COMMENT "Copying $<TARGET_FILE_NAME:_C> to scalellm directory"
+      COMMAND ${CMAKE_COMMAND} -E create_symlink
+              $<TARGET_FILE:_C>
+              ${CMAKE_SOURCE_DIR}/scalellm/$<TARGET_FILE_NAME:_C>
+  )
+endif()