Fix Python path finding for benchmark tests

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/16022

Differential Revision: D13673792

Pulled By: pjh5

fbshipit-source-id: 177a823ef343b7f60e26ad9ef51415332045438d

diff --git a/.jenkins/caffe2/bench.sh b/.jenkins/caffe2/bench.sh
index f4f5f2a..03720a7 100755 (executable)
--- a/.jenkins/caffe2/bench.sh
+++ b/.jenkins/caffe2/bench.sh
@@ -13,7 +13,8 @@ else
     num_gpus=0
 fi
 
-cmd="$PYTHON $CAFFE2_PYPATH/python/examples/resnet50_trainer.py --train_data null --batch_size 64 --epoch_size 6400 --num_epochs 2"
+caffe2_pypath="$(python -c 'import os; import caffe2; print(os.path.dirname(os.path.realpath(caffe2.__file__)))')"
+cmd="$PYTHON $caffe2_pypath/python/examples/resnet50_trainer.py --train_data null --batch_size 64 --epoch_size 6400 --num_epochs 2"
 if (( $num_gpus == 0 )); then
     cmd="$cmd --use_cpu"
 else

diff --git a/.jenkins/caffe2/test.sh b/.jenkins/caffe2/test.sh
index 6800393..66207c5 100755 (executable)
--- a/.jenkins/caffe2/test.sh
+++ b/.jenkins/caffe2/test.sh
@@ -55,38 +55,9 @@ if [[ "$BUILD_ENVIRONMENT" == *cmake* ]]; then
   exit 0
 fi
 
-# Ideally this would be where the Python bits get installed to from using
-# setup.py. However on our dockers this is not correct for two reasons.
-# 1. This lies in /usr/local/lib/pythonM.m, but the dockers don't have the
-#    right permissions setup so the build doesn't have write access to this
-#    dir.  For this reason we use --user flag in all pip install instructions
-#    which install into the $HOME/.local directory instead.
-# 2. This returns lib/pythonM.m/dist-packages, but we install in site-packages.
-# We use this same way of getting the install directory in other places in our
-# build, so not really sure why it is not correct here
-INSTALL_SITE_DIR=$($PYTHON -c "from distutils import sysconfig; print(sysconfig.get_python_lib())")
-if [[ -n "$(find $INSTALL_SITE_DIR -name caffe2 2>/dev/null)" ]]; then
-  # Caffe2 will probably be found here if using a Python from a virtualenv or
-  # from conda
-  CAFFE2_PYPATH="$INSTALL_SITE_DIR/caffe2"
-elif [[ -n "$(find $HOME/.local/lib -name caffe2 2>/dev/null)" ]]; then
-  # Caffe2 will be found here in the case of using no env and adding --user to
-  # the setup.py call
-  pyver=($(python --version 2>&1))
-  pyver=${pyver[1]}
-  pyver=${pyver:0:3}
-  CAFFE2_PYPATH="$HOME/.local/lib/python$pyver/site-packages/caffe2"
-else
-  echo "I do not know where Caffe2 is installed"
-  find / -name caffe2 2>/dev/null
-  exit 1
-fi
-if [[ ! -d "$CAFFE2_PYPATH" ]]; then
-  echo "Failed to find where Caffe2 Python bits are installed"
-  find / -name caffe2 2>/dev/null
-  exit 1
-fi
-
+# Anywhere except $ROOT_DIR should work
+cd "$INSTALL_PREFIX"
+caffe2_pypath="$(python -c 'import os; import caffe2; print(os.path.dirname(os.path.realpath(caffe2.__file__)))')"
 
 if [[ "$BUILD_ENVIRONMENT" == *ubuntu14.04* ]]; then
   # Hotfix, use hypothesis 3.44.6 on Ubuntu 14.04
@@ -109,7 +80,7 @@ EXTRA_TESTS=()
 
 # CUDA builds always include NCCL support
 if [[ "$BUILD_ENVIRONMENT" == *-cuda* ]]; then
-  EXTRA_TESTS+=("$CAFFE2_PYPATH/contrib/nccl")
+  EXTRA_TESTS+=("$caffe2_pypath/contrib/nccl")
 fi
 
 rocm_ignore_test=()
@@ -117,14 +88,14 @@ if [[ $BUILD_ENVIRONMENT == *-rocm* ]]; then
   # Currently these tests are failing on ROCM platform:
 
   # Unknown reasons, need to debug
-  rocm_ignore_test+=("--ignore $CAFFE2_PYPATH/python/operator_test/arg_ops_test.py")
-  rocm_ignore_test+=("--ignore $CAFFE2_PYPATH/python/operator_test/piecewise_linear_transform_test.py")
-  rocm_ignore_test+=("--ignore $CAFFE2_PYPATH/python/operator_test/softmax_ops_test.py")
-  rocm_ignore_test+=("--ignore $CAFFE2_PYPATH/python/operator_test/unique_ops_test.py")
+  rocm_ignore_test+=("--ignore $caffe2_pypath/python/operator_test/arg_ops_test.py")
+  rocm_ignore_test+=("--ignore $caffe2_pypath/python/operator_test/piecewise_linear_transform_test.py")
+  rocm_ignore_test+=("--ignore $caffe2_pypath/python/operator_test/softmax_ops_test.py")
+  rocm_ignore_test+=("--ignore $caffe2_pypath/python/operator_test/unique_ops_test.py")
 
   # On ROCm, RCCL (distributed) development isn't complete.
   # https://github.com/ROCmSoftwarePlatform/rccl
-  rocm_ignore_test+=("--ignore $CAFFE2_PYPATH/python/data_parallel_model_test.py")
+  rocm_ignore_test+=("--ignore $caffe2_pypath/python/data_parallel_model_test.py")
 fi
 
 # NB: Warnings are disabled because they make it harder to see what
@@ -137,20 +108,18 @@ pip install --user pytest-sugar
   -v \
   --disable-warnings \
   --junit-xml="$pytest_reports_dir/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" \
+  --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" \
   ${rocm_ignore_test[@]} \
-  "$CAFFE2_PYPATH/python" \
+  "$caffe2_pypath/python" \
   "${EXTRA_TESTS[@]}"
 
 #####################
 # torchvision tests #
 #####################
 
-cd ${INSTALL_PREFIX}
-
 if [[ "$BUILD_ENVIRONMENT" == *onnx* ]]; then
   pip install --user torchvision
   "$ROOT_DIR/scripts/onnx/test.sh"

diff --git a/setup.py b/setup.py
index 4fd8dca..4727bdd 100644 (file)
--- a/setup.py
+++ b/setup.py
@@ -910,6 +910,7 @@ if __name__ == '__main__':
             ],
             'caffe2': [
                 'cpp_test/*',
+                'python/serialized_test/data/operator_test/*.zip',
             ]
         },
     )