test: test with pre-release of PyTorch

metaopt · Aug 29, 2022 · 1fded5d · 1fded5d
1 parent 3eee243
commit 1fded5d
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 15 deletions.
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
@@ -42,13 +42,13 @@ jobs:
       - run: |
           CUDA_VERSION="${{steps.cuda-toolkit.outputs.cuda}}"
           echo "CUDA_VERSION=${CUDA_VERSION}" >> "${GITHUB_ENV}"
-          TORCH_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')"
-          echo "TORCH_INDEX_URL=${TORCH_INDEX_URL}" >> "${GITHUB_ENV}"
+          PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')"
+          echo "PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}" >> "${GITHUB_ENV}"
 
           echo "Installed CUDA version is: ${CUDA_VERSION}"
           echo "CUDA install location: ${{steps.cuda-toolkit.outputs.CUDA_PATH}}"
           nvcc -V
-          echo "Torch index URL: ${TORCH_INDEX_URL}"
+          echo "Torch index URL: ${PIP_EXTRA_INDEX_URL}"
 
       - name: Upgrade pip
         run: |
@@ -92,8 +92,7 @@ jobs:
 
       - name: Install dependencies
         run: |
-          python -m pip install --extra-index-url "${TORCH_INDEX_URL}" \
-            -r docs/requirements.txt
+          python -m pip install -r docs/requirements.txt
 
       - name: docstyle
         run: |

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -53,25 +53,31 @@ jobs:
       - run: |
           CUDA_VERSION="${{steps.cuda-toolkit.outputs.cuda}}"
           echo "CUDA_VERSION=${CUDA_VERSION}" >> "${GITHUB_ENV}"
-          TORCH_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')"
-          echo "TORCH_INDEX_URL=${TORCH_INDEX_URL}" >> "${GITHUB_ENV}"
+          PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')"
+          echo "PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}" >> "${GITHUB_ENV}"
 
           echo "Installed CUDA version is: ${CUDA_VERSION}"
           echo "CUDA install location: ${{steps.cuda-toolkit.outputs.CUDA_PATH}}"
           nvcc -V
-          echo "Torch index URL: ${TORCH_INDEX_URL}"
+          echo "Torch index URL: ${PIP_EXTRA_INDEX_URL}"
 
       - name: Upgrade pip
         run: |
           python -m pip install --upgrade pip setuptools
 
+      - name: Install PyTorch and FuncTorch nightly
+        run: |
+          export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL//whl/whl\/nightly}"
+          python -m pip install 'torch >= 1.13.0dev' ninja
+          python -m pip install git+https://github.com/pytorch/functorch.git
+
       - name: Install dependencies
         run: |
-          python -m pip install --extra-index-url "${TORCH_INDEX_URL}" \
-            -r tests/requirements.txt
+          python -m pip install -r tests/requirements.txt
 
       - name: Install TorchOpt
         run: |
+          export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL//whl/whl\/nightly}"
           python -m pip install -vvv -e .
 
       - name: Test with pytest

diff --git a/Dockerfile b/Dockerfile
@@ -31,15 +31,15 @@ RUN echo "export PS1='[\[\e[1;33m\]\u\[\e[0m\]:\[\e[1;35m\]\w\[\e[0m\]]\$ '" >>
 
 # Setup virtual environment
 RUN /usr/bin/python3.9 -m venv --upgrade-deps ~/venv && rm -rf ~/.pip/cache
-RUN TORCH_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')" && \
-    echo "export TORCH_INDEX_URL='${TORCH_INDEX_URL}'" >> ~/venv/bin/activate && \
+RUN PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu$(echo "${CUDA_VERSION}" | cut -d'.' -f-2  | tr -d '.')" && \
+    echo "export PIP_EXTRA_INDEX_URL='${PIP_EXTRA_INDEX_URL}'" >> ~/venv/bin/activate && \
     echo "source /home/torchopt/venv/bin/activate" >> ~/.bashrc
 
 # Install dependencies
 WORKDIR /home/torchopt/torchopt
 COPY --chown=torchopt requirements.txt requirements.txt
 RUN source ~/venv/bin/activate && \
-    python -m pip install --extra-index-url "${TORCH_INDEX_URL}" -r requirements.txt && \
+    python -m pip install -r requirements.txt && \
     rm -rf ~/.pip/cache ~/.cache/pip
 
 ####################################################################################################
@@ -63,8 +63,7 @@ RUN go install github.com/google/addlicense@latest
 COPY --chown=torchopt tests/requirements.txt tests/requirements.txt
 COPY --chown=torchopt tutorials/requirements.txt tutorials/requirements.txt
 RUN source ~/venv/bin/activate && \
-    python -m pip install --extra-index-url "${TORCH_INDEX_URL}" \
-        -r tests/requirements.txt -r tutorials/requirements.txt && \
+    python -m pip install -r tests/requirements.txt -r tutorials/requirements.txt && \
     rm -rf ~/.pip/cache ~/.cache/pip
 
 ####################################################################################################