Add tests for wheels in CI

.circleci/config.yml

@@ -32,7 +32,7 @@ update_and_install_python: &update_and_install_python
       command: |
         sudo add-apt-repository -y ppa:deadsnakes/ppa
         sudo apt-get update
-        sudo apt-get install -y python3.7 python3.7-dev python3.8 python3.8-dev
+        sudo apt-get install -y python3.9 python3.9-dev
 
 install_nox: &install_nox
   - run:
@@ -86,7 +86,7 @@ setup_environment: &setup_environment
       name: Setup virtualenv and tools
       working_directory: ~/project
       command: |
-        virtualenv ~/theseus_venv -p /usr/bin/python3
+        virtualenv ~/theseus_venv -p /usr/bin/python3.9
         echo ". ~/theseus_venv/bin/activate" >> $BASH_ENV
         . ~/theseus_venv/bin/activate
         pip install --progress-bar off --upgrade pip
@@ -176,6 +176,27 @@ run_tests: &run_tests
         pytest -s tests/test_theseus_layer.py
         pytest -s tests -m "cudaext"
 
+build_cuda11_wheel: &build_cuda11_wheel
+  - run:
+      name: Building wheel for CUDA 11
+      working_directory: /home/circleci
+      command: |
+        THESEUS_GIT_COMMIT=$(git --git-dir project/.git log --format="%H" -n 1)
+        THESEUS_VERSION=$(grep -Eo  "[0-9].[0-9].[0-9][.0-9a-z]*" project/theseus/__init__.py)
+        ./project/build_scripts/build_wheel.sh . ${THESEUS_GIT_COMMIT} 11.6 ${THESEUS_VERSION}
+        pip install $(ls */*.whl)
+        pip install -r project/requirements/dev.txt
+
+run_tests_from_wheel: &run_tests_from_wheel
+  - run:
+      name: Installing theseus from wheel and running tests
+      working_directory: ~/project
+      command: |
+          mv theseus theseus_tmp
+          python -m pytest tests/test_theseus_layer.py
+          pytest -s tests -m "cudaext"
+
+
 # -------------------------------------------------------------------------------------
 # Jobs
 # -------------------------------------------------------------------------------------
@@ -258,6 +279,19 @@ jobs:
       - <<: *setup_project
       - <<: *run_tests
 
+  test_cuda11_wheel:
+    executor: gpu_cuda11
+    steps:
+      - checkout
+      - <<: *update_and_install_python
+      - <<: *install_suitesparse
+      - <<: *setup_cuda11_libs
+      - <<: *setup_environment
+      - <<: *build_cuda11_wheel
+      - <<: *run_tests_from_wheel
+
+
+
 workflows:
   version: 2
   build:
@@ -267,3 +301,4 @@ workflows:
       - py39_linux
       - unittests_gpu17_cuda10
       - unittests_gpu17_cuda11
+      - test_cuda11_wheel

build_scripts/build_wheel.sh

@@ -5,31 +5,35 @@
 # to install theseus.
 #
 # To use this script, from root theesus folder run 
-#    ./build_scripts/build_wheel.sh ROOT_DIR TAG CUDA_VERSION
+#    ./build_scripts/build_wheel.sh ROOT_DIR COMMIT CUDA_VERSION THESEUS_VERSION(optional) 
 #
 # ROOT_DIR: is the directory where the Dockerfile, tar.gz and .whl files will be stored
 #   (under a new subdirectory named theseus_docker_3.9)
-# TAG: is a theseus tag (e.g., 0.1.0)
+# COMMIT: is a theseus commit hash or tag (e.g., 0.1.3).
 # CUDA_VERSION: the version of CUDA to use. We have tested 10.2, 11.3, 11.6, and 11.7.
-#   You can also pass "cpu" to compile without CUDA extensions. 
+#   You can also pass "cpu" to compile without CUDA extensions.
+# THESEUS_VERSION: defaults to COMMIT, otherwise it must match the version in the commit.
 #
 #   For example
 #    ./build_scripts/build_wheel.sh . 0.1.0 10.2
 #   
 #   will run and store results under ./theseus_docker_3.9
 # -----------------
 
-# Ensure that 3 arguments (ROOT_DIR, TAG, CUDA_VERSION) are provided
+# Ensure that 3 or 4 arguments 
+# (ROOT_DIR, COMMIT, CUDA_VERSION, THESEUS_VERSION - optional) are provided.
 die () {
     echo >&2 "$@"
     exit 1
 }
-[ "$#" -eq 3 ] || die "3 arguments required, $# provided"
+[ "$#" -eq 3 ] || [ "$#" -eq 4 ] || die "3 or 4 arguments required, $# provided"
 ROOT_DIR=$1
-TAG=$2
+COMMIT=$2
 CUDA_VERSION=$3
+TH_VERSION=${4:-${COMMIT}}
 
-SUPPORTED_CUDA_VERSIONS="10.2 11.1 11.3 11.5 11.6 11.7"
+
+SUPPORTED_CUDA_VERSIONS="10.2 11.3 11.6 11.7"
 CUDA_VERSION_IS_SUPPORTED=$(echo "cpu ${SUPPORTED_CUDA_VERSIONS}" | grep -w ${CUDA_VERSION})
 [ "${CUDA_VERSION_IS_SUPPORTED}" ] || die "CUDA_VERSION must be one of (cpu ${SUPPORTED_CUDA_VERSIONS})"
 
@@ -55,7 +59,7 @@ else
         BASPACHO_CUDA_ARCHS="${BASPACHO_CUDA_ARCHS};80"
         TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST};8.0"
     fi
-    DEPRECATED_TORCH_CUDA="10.2 11.1 11.3"
+    DEPRECATED_TORCH_CUDA="10.2 11.3"
     CUDA_IS_TORCH_DEPRECATED=$(echo "${DEPRECATED_TORCH_CUDA}" | grep -w ${CUDA_VERSION})
     if [[ ${CUDA_IS_TORCH_DEPRECATED} ]]
     then
@@ -72,21 +76,6 @@ for PYTHON_VERSION in 3.9; do
     echo """# ----------------
     FROM ${IMAGE_NAME}
 
-    # --- Install conda and environment
-    ENV CONDA_DIR /opt/conda
-    RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh && \
-        /bin/bash ~/miniconda.sh -b -p /opt/conda
-    ENV PATH \$CONDA_DIR/bin:\$PATH
-    RUN conda create --name theseus python=${PYTHON_VERSION}
-    RUN source activate theseus
-
-    # --- Install torch
-    ENV CUDA_HOME /usr/local/cuda-${CUDA_VERSION}
-    RUN pip install ${TORCH_VERSION} --extra-index-url https://download.pytorch.org/whl/${DEVICE_TAG}
-
-    # --- Install sparse suitesparse
-    RUN conda install -c conda-forge suitesparse
-
     # --- Install baspacho dependencies (cmake, BLAS)
     RUN wget --quiet https://github.com/Kitware/CMake/releases/download/v3.24.2/cmake-3.24.2-linux-x86_64.sh -O ~/cmake3.24.sh
     RUN mkdir /opt/cmake3.24
@@ -97,7 +86,6 @@ for PYTHON_VERSION in 3.9; do
     # --- Install baspacho
     RUN git clone https://github.com/facebookresearch/baspacho.git
     WORKDIR baspacho
-
     # Note: to use static BLAS the option is really BLA_STATIC (https://cmake.org/cmake/help/latest/module/FindBLAS.html)
     RUN /opt/cmake3.24/bin/cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DBLA_STATIC=ON \
         ${BASPACHO_CUDA_ARGS} \
@@ -106,12 +94,27 @@ for PYTHON_VERSION in 3.9; do
     RUN /opt/cmake3.24/bin/cmake --build build -- -j16
     WORKDIR ..
 
+    # --- Install conda and environment
+    ENV CONDA_DIR /opt/conda
+    RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh && \
+        /bin/bash ~/miniconda.sh -b -p /opt/conda
+    ENV PATH \$CONDA_DIR/bin:\$PATH
+    RUN conda create --name theseus python=${PYTHON_VERSION}
+    RUN source activate theseus
+
+    # --- Install torch
+    ENV CUDA_HOME /usr/local/cuda-${CUDA_VERSION}
+    RUN pip install ${TORCH_VERSION} --extra-index-url https://download.pytorch.org/whl/${DEVICE_TAG}
+
+    # --- Install sparse suitesparse
+    RUN conda install -c conda-forge suitesparse
+
     # --- Compile theseus wheel
     RUN pip install build wheel
     RUN git clone https://github.com/facebookresearch/theseus.git
     WORKDIR theseus
     RUN git fetch --all --tags
-    RUN git checkout tags/${TAG} -b tmp_build
+    RUN git checkout ${COMMIT} -b tmp_build
     CMD BASPACHO_ROOT_DIR=/baspacho THESEUS_FORCE_CUDA=${ENABLE_CUDA} TORCH_CUDA_ARCH_LIST='${TORCH_CUDA_ARCH_LIST}' python3 -m build --no-isolation
     """ > ${DOCKER_DIR}/Dockerfile
 
@@ -124,16 +127,16 @@ for PYTHON_VERSION in 3.9; do
 
     # Copy the wheel to host
     CP_STR="cp"$(echo ${PYTHON_VERSION} | sed 's/[.]//g')
-    DOCKER_WHL="theseus/dist/theseus_ai-${TAG}-${CP_STR}-${CP_STR}-linux_x86_64.whl"
-    if [[ ${CUDA_VERSION} == "10.2" ]]
+    DOCKER_WHL="theseus/dist/theseus_ai-${TH_VERSION}-${CP_STR}-${CP_STR}-linux_x86_64.whl"
+    if [[ ${CUDA_VERSION} == "11.6" ]]
     then
-        PLUS_CU_TAG=""  # 10.2 will be the pypi version, so don't add +cu102
+        PLUS_CU_TAG=""  # 11.6 will be the pypi version, so don't add +cu116
     else
         PLUS_CU_TAG="+${DEVICE_TAG}"
     fi
-    HOST_WHL="theseus_ai-${TAG}${PLUS_CU_TAG}-${CP_STR}-${CP_STR}-manylinux_2_17_x86_64.whl"
+    HOST_WHL="theseus_ai-${TH_VERSION}${PLUS_CU_TAG}-${CP_STR}-${CP_STR}-manylinux_2_17_x86_64.whl"
 
-    sudo docker cp "${DOCKER_NAME}:theseus/dist/theseus-ai-${TAG}.tar.gz" "theseus-ai-${TAG}.tar.gz"
+    sudo docker cp "${DOCKER_NAME}:theseus/dist/theseus-ai-${TH_VERSION}.tar.gz" "theseus-ai-${TH_VERSION}.tar.gz"
     sudo docker cp "${DOCKER_NAME}:${DOCKER_WHL}" ${HOST_WHL}
     sudo docker rm ${DOCKER_NAME}
     sudo docker image rm "${DOCKER_NAME}_img"

tests/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.