cmake: python packages now install to the cannonical directory

Summary:
Addresses issue pytorch#1676

Now when `make install` is run, the `caffe2` (and `caffe`) python modules will be installed into the correct site-packages directory (relative to the prefix) instead of directly in the prefix.
Closes facebookarchive/caffe2#1677

Reviewed By: pietern

Differential Revision: D6710247

Pulled By: bddppq

fbshipit-source-id: b49167d48fd94d87f7b7c1ebf0f187ec6a203470

Author: Erotemic

.jenkins/test.sh

@@ -2,6 +2,21 @@
 
 set -e
 
+# Figure out which Python to use
+PYTHON="python"
+if [ -n "$BUILD_ENVIRONMENT" ]; then
+  if [[ "$BUILD_ENVIRONMENT" == py2* ]]; then
+    PYTHON="python2"
+  elif [[ "$BUILD_ENVIRONMENT" == py3* ]]; then
+    PYTHON="python3"
+  fi
+fi
+
+# The prefix must mirror the setting from build.sh
+INSTALL_PREFIX="/usr/local/caffe2"
+# Add the site-packages in the caffe2 install prefix to the PYTHONPATH
+SITE_DIR=$($PYTHON -c "from distutils import sysconfig; print(sysconfig.get_python_lib(prefix=''))")
+
 LOCAL_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
 ROOT_DIR=$(cd "$LOCAL_DIR"/.. && pwd)
 
@@ -11,8 +26,8 @@ if [[ "${BUILD_ENVIRONMENT}" == *-android* ]]; then
   exit 0
 fi
 
-export PYTHONPATH="${PYTHONPATH}:/usr/local/caffe2"
-export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/caffe2/lib"
+export PYTHONPATH="${PYTHONPATH}:${INSTALL_PREFIX}/${SITE_DIR}"
+export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${INSTALL_PREFIX}/lib"
 
 exit_code=0
 
@@ -38,8 +53,7 @@ fi
 mkdir -p ./test/{cpp,python}
 TEST_DIR="$PWD/test"
 
-
-cd /usr/local/caffe2
+cd ${INSTALL_PREFIX}
 
 # Commands below may exit with non-zero status
 set +e
@@ -61,36 +75,29 @@ for test in ./test/*; do
   fi
 done
 
-# Figure out which Python to use
-PYTHON="python"
-if [ -n "$BUILD_ENVIRONMENT" ]; then
-  if [[ "$BUILD_ENVIRONMENT" == py2* ]]; then
-    PYTHON="python2"
-  elif [[ "$BUILD_ENVIRONMENT" == py3* ]]; then
-    PYTHON="python3"
-  fi
-fi
+# Get the relative path to where the caffe2 python module was installed
+CAFFE2_PYPATH="$SITE_DIR/caffe2"
 
 # Collect additional tests to run (outside caffe2/python)
 EXTRA_TESTS=()
 
 # CUDA builds always include NCCL support
 if [[ "$BUILD_ENVIRONMENT" == *-cuda* ]]; then
-  EXTRA_TESTS+=(caffe2/contrib/nccl)
+  EXTRA_TESTS+=("$CAFFE2_PYPATH/contrib/nccl")
 fi
 
 # Python tests
 echo "Running Python tests.."
 "$PYTHON" \
   -m pytest \
   -v \
-  --junit-xml="$TEST_DIR"/python/result.xml \
-  --ignore caffe2/python/test/executor_test.py \
-  --ignore caffe2/python/operator_test/matmul_op_test.py \
-  --ignore caffe2/python/operator_test/pack_ops_test.py \
-  --ignore caffe2/python/mkl/mkl_sbn_speed_test.py \
-  caffe2/python/ \
-  ${EXTRA_TESTS[@]}
+  --junit-xml="$TEST_DIR/python/result.xml" \
+  --ignore "$CAFFE2_PYPATH/python/test/executor_test.py" \
+  --ignore "$CAFFE2_PYPATH/python/operator_test/matmul_op_test.py" \
+  --ignore "$CAFFE2_PYPATH/python/operator_test/pack_ops_test.py" \
+  --ignore "$CAFFE2_PYPATH/python/mkl/mkl_sbn_speed_test.py" \
+  "$CAFFE2_PYPATH/python" \
+  "${EXTRA_TESTS[@]}"
 
 tmp_exit_code="$?"
 if [ "$exit_code" -eq 0 ]; then

caffe2/CMakeLists.txt

@@ -197,6 +197,13 @@ endif()
 
 
 if (BUILD_PYTHON)
+  # Python site-packages
+  # Get canonical directory for python site packages (relative to install
+  # location).  It varys from system to system.
+  pycmd(python_site_packages "
+      from distutils import sysconfig
+      print(sysconfig.get_python_lib(prefix=''))
+  ")
   # ---[ Python.
   add_library(caffe2_pybind11_state MODULE ${Caffe2_CPU_PYTHON_SRCS})
   set_target_properties(caffe2_pybind11_state PROPERTIES COMPILE_FLAGS "-fvisibility=hidden")
@@ -214,7 +221,7 @@ if (BUILD_PYTHON)
   target_link_libraries(
       caffe2_pybind11_state ${Caffe2_CPU_LINK} ${Caffe2_DEPENDENCY_LIBS}
       ${Caffe2_PYTHON_DEPENDENCY_LIBS})
-  install(TARGETS caffe2_pybind11_state DESTINATION caffe2/python)
+  install(TARGETS caffe2_pybind11_state DESTINATION "${python_site_packages}/caffe2/python")
 
   if(USE_CUDA)
     add_library(caffe2_pybind11_state_gpu MODULE ${Caffe2_GPU_PYTHON_SRCS})
@@ -233,7 +240,7 @@ if (BUILD_PYTHON)
     target_link_libraries(
         caffe2_pybind11_state_gpu ${Caffe2_CPU_LINK} ${Caffe2_GPU_LINK} ${Caffe2_DEPENDENCY_LIBS}
         ${Caffe2_CUDA_DEPENDENCY_LIBS} ${Caffe2_PYTHON_DEPENDENCY_LIBS})
-    install(TARGETS caffe2_pybind11_state_gpu DESTINATION caffe2/python)
+    install(TARGETS caffe2_pybind11_state_gpu DESTINATION "${python_site_packages}/caffe2/python")
   endif()
 
   if (MSVC AND CMAKE_GENERATOR MATCHES "Visual Studio")
@@ -292,13 +299,13 @@ if (BUILD_PYTHON)
 
   # Install commands
   # Pick up static python files
-  install(DIRECTORY ${CMAKE_BINARY_DIR}/caffe2 DESTINATION .
+  install(DIRECTORY ${CMAKE_BINARY_DIR}/caffe2 DESTINATION ${python_site_packages}
           FILES_MATCHING PATTERN "*.py")
   # Caffe proto files
-  install(DIRECTORY ${CMAKE_BINARY_DIR}/caffe DESTINATION .
+  install(DIRECTORY ${CMAKE_BINARY_DIR}/caffe DESTINATION ${python_site_packages}
           FILES_MATCHING PATTERN "*.py")
   # Caffe2 proto files
-  install(DIRECTORY ${CMAKE_BINARY_DIR}/caffe2 DESTINATION .
+  install(DIRECTORY ${CMAKE_BINARY_DIR}/caffe2 DESTINATION ${python_site_packages}
           FILES_MATCHING PATTERN "*.py")
 endif()
 

cmake/Utils.cmake

@@ -285,3 +285,78 @@ function(caffe2_include_directories)
     endif()
   endforeach()
 endfunction()
+
+
+###
+# Removes common indentation from a block of text to produce code suitable for
+# setting to `python -c`, or using with pycmd. This allows multiline code to be
+# nested nicely in the surrounding code structure.
+#
+# This function respsects PYTHON_EXECUTABLE if it defined, otherwise it uses
+# `python` and hopes for the best. An error will be thrown if it is not found.
+#
+# Args:
+#     outvar : variable that will hold the stdout of the python command
+#     text   : text to remove indentation from
+#
+function(dedent outvar text)
+  # Use PYTHON_EXECUTABLE if it is defined, otherwise default to python
+  if ("${PYTHON_EXECUTABLE}" STREQUAL "")
+    set(_python_exe "python")
+  else()
+    set(_python_exe "${PYTHON_EXECUTABLE}")
+  endif()
+  set(_fixup_cmd "import sys; from textwrap import dedent; print(dedent(sys.stdin.read()))")
+  # Use echo to pipe the text to python's stdinput. This prevents us from
+  # needing to worry about any sort of special escaping.
+  execute_process(
+    COMMAND echo "${text}"
+    COMMAND "${_python_exe}" -c "${_fixup_cmd}"
+    RESULT_VARIABLE _dedent_exitcode
+    OUTPUT_VARIABLE _dedent_text)
+  if(NOT ${_dedent_exitcode} EQUAL 0)
+    message(ERROR " Failed to remove indentation from: \n\"\"\"\n${text}\n\"\"\"
+    Python dedent failed with error code: ${_dedent_exitcode}")
+    message(FATAL_ERROR " Python dedent failed with error code: ${_dedent_exitcode}")
+  endif()
+  # Remove supurflous newlines (artifacts of print)
+  string(STRIP "${_dedent_text}" _dedent_text)
+  set(${outvar} "${_dedent_text}" PARENT_SCOPE)
+endfunction()
+
+
+###
+# Helper function to run `python -c "<cmd>"` and capture the results of stdout
+#
+# Runs a python command and populates an outvar with the result of stdout.
+# Common indentation in the text of `cmd` is removed before the command is
+# executed, so the caller does not need to worry about indentation issues.
+#
+# This function respsects PYTHON_EXECUTABLE if it defined, otherwise it uses
+# `python` and hopes for the best. An error will be thrown if it is not found.
+#
+# Args:
+#     outvar : variable that will hold the stdout of the python command
+#     cmd    : text representing a (possibly multiline) block of python code
+#
+function(pycmd outvar cmd)
+  dedent(_dedent_cmd "${cmd}")
+  # Use PYTHON_EXECUTABLE if it is defined, otherwise default to python
+  if ("${PYTHON_EXECUTABLE}" STREQUAL "")
+    set(_python_exe "python")
+  else()
+    set(_python_exe "${PYTHON_EXECUTABLE}")
+  endif()
+  # run the actual command
+  execute_process(
+    COMMAND "${_python_exe}" -c "${_dedent_cmd}"
+    RESULT_VARIABLE _exitcode
+    OUTPUT_VARIABLE _output)
+  if(NOT ${_exitcode} EQUAL 0)
+    message(ERROR " Failed when running python code: \"\"\"\n${_dedent_cmd}\n\"\"\"")
+    message(FATAL_ERROR " Python command failed with error code: ${_exitcode}")
+  endif()
+  # Remove supurflous newlines (artifacts of print)
+  string(STRIP "${_output}" _output)
+  set(${outvar} "${_output}" PARENT_SCOPE)
+endfunction()

conda/no_cuda/build.sh

@@ -53,8 +53,3 @@ else
 fi
 
 make install/fast
-
-# Python libraries got installed to wrong place, so move them
-# to the right place. See https://github.com/caffe2/caffe2/issues/1015
-echo "Installing Python to $SP_DIR"
-mv $PREFIX/caffe2 $SP_DIR

scripts/build_local.sh

@@ -39,14 +39,20 @@ else
   cd "$BUILD_ROOT"
   echo "Building Caffe2 in: $BUILD_ROOT"
 
-  # Now, actually build the target.
   cmake "$CAFFE2_ROOT" \
         "${CMAKE_ARGS[@]}" \
         "$@"
 
-  if [ "$(uname)" == 'Darwin' ]; then
-    cmake --build . -- "-j$(sysctl -n hw.ncpu)"
+  # Determine the number of CPUs to build with.
+  # If the `CAFFE_MAKE_NCPUS` variable is not specified, use them all.
+  if [ -n "${CAFFE_MAKE_NCPUS}" ]; then
+      CAFFE_MAKE_NCPUS="$CAFFE_MAKE_NCPUS"
+  elif [ "$(uname)" == 'Darwin' ]; then
+      CAFFE_MAKE_NCPUS="$(sysctl -n hw.ncpu)"
   else
-    cmake --build . -- "-j$(nproc)"
+      CAFFE_MAKE_NCPUS="$(nproc)"
   fi
+
+  # Now, actually build the target.
+  cmake --build . -- "-j$CAFFE_MAKE_NCPUS"
 fi

setup.py

@@ -83,7 +83,35 @@ def run(self):
 
 
 class build_ext(setuptools.command.build_ext.build_ext):
+    """
+    Compiles everything when `python setup.py build` is run using cmake.
+
+    Custom args can be passed to cmake by specifying the `CMAKE_ARGS`
+    environment variable. E.g. to build without cuda support run:
+        `CMAKE_ARGS=-DUSE_CUDA=Off python setup.py build`
+
+    The number of CPUs used by `make` can be specified by passing `-j<ncpus>`
+    to `setup.py build`.  By default all CPUs are used.
+    """
+    user_options = [
+        ('jobs=', 'j', 'Specifies the number of jobs to use with make')
+    ]
+
+    def initialize_options(self):
+        setuptools.command.build_ext.build_ext.initialize_options(self)
+        self.jobs = None
+
+    def finalize_options(self):
+        setuptools.command.build_ext.build_ext.finalize_options(self)
+        # Check for the -j argument to make with a specific number of cpus
+        try:
+            self.jobs = int(self.jobs)
+        except Exception:
+            self.jobs = None
+
     def _build_with_cmake(self):
+        # build_temp resolves to something like: build/temp.linux-x86_64-3.5
+        # build_lib resolves to something like: build/lib.linux-x86_64-3.5
         build_temp = os.path.realpath(self.build_temp)
         build_lib = os.path.realpath(self.build_lib)
 
@@ -101,6 +129,10 @@ def _build_with_cmake(self):
                 cmake_args = []
             log.info('CMAKE_ARGS: {}'.format(cmake_args))
 
+            if self.jobs is not None:
+                # use envvars to pass information to `build_local.sh`
+                os.environ['CAFFE_MAKE_NCPUS'] = str(self.jobs)
+
             self.compiler.spawn([
                 os.path.join(TOP_DIR, 'scripts', 'build_local.sh'),
                 '-DBUILD_SHARED_LIBS=OFF',
@@ -114,8 +146,7 @@ def _build_with_cmake(self):
                 '-DBUILD_TEST=OFF',
                 '-BUILD_BENCHMARK=OFF',
                 '-DBUILD_BINARY=OFF',
-                TOP_DIR
-            ] + cmake_args)
+            ] + cmake_args + [TOP_DIR])
             # This is assuming build_local.sh will use TOP_DIR/build
             # as the cmake build directory
             self.compiler.spawn([
@@ -124,11 +155,17 @@ def _build_with_cmake(self):
                 'install'
             ])
         else:
+            # if `CMAKE_INSTALL_DIR` is specified in the environment, assume
+            # cmake has been run and skip the build step.
             cmake_install_dir = os.environ['CMAKE_INSTALL_DIR']
 
+        # CMake will install the python package to a directory that mirrors the
+        # standard site-packages name. This will vary slightly depending on the
+        # OS and python version.  (e.g. `lib/python3.5/site-packages`)
+        python_site_packages = sysconfig.get_python_lib(prefix='')
         for d in ['caffe', 'caffe2']:
-            self.copy_tree(os.path.join(cmake_install_dir, d),
-                           os.path.join(build_lib, d))
+            src = os.path.join(cmake_install_dir, python_site_packages, d)
+            self.copy_tree(src, os.path.join(build_lib, d))
 
     def get_outputs(self):
         return [os.path.join(self.build_lib, d)