Create wheels.yml for github actions (pytorch#244)

ai4co · Jul 5, 2022 · ad92dd7 · ad92dd7
1 parent 009cea6
commit ad92dd7
Show file tree

Hide file tree

Showing 19 changed files with 261 additions and 63 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -656,32 +656,102 @@ workflows:
 
   unittest:
     jobs:
+      - unittest_macos_cpu:
+          cu_version: cpu
+          name: unittest_macos_cpu_py3.7
+          python_version: '3.7'
+      - unittest_linux_cpu:
+          cu_version: cpu
+          name: unittest_linux_cpu_py3.7
+          python_version: '3.7'
+      - unittest_linux_gpu:
+          cu_version: cu113
+          name: unittest_linux_gpu_py3.7
+          python_version: '3.7'
+      - unittest_linux_optdeps_gpu:
+          cu_version: cu113
+          name: unittest_linux_optdeps_gpu_py3.7
+          python_version: '3.7'
+      - unittest_linux_stable_cpu:
+          cu_version: cpu
+          name: unittest_linux_stable_cpu_py3.7
+          python_version: '3.7'
+      - unittest_linux_stable_gpu:
+          cu_version: cu113
+          name: unittest_linux_stable_gpu_py3.7
+          python_version: '3.7'
+
       - unittest_macos_cpu:
           cu_version: cpu
           name: unittest_macos_cpu_py3.8
           python_version: '3.8'
-
       - unittest_linux_cpu:
           cu_version: cpu
           name: unittest_linux_cpu_py3.8
           python_version: '3.8'
-
       - unittest_linux_gpu:
           cu_version: cu113
           name: unittest_linux_gpu_py3.8
           python_version: '3.8'
-
       - unittest_linux_optdeps_gpu:
           cu_version: cu113
           name: unittest_linux_optdeps_gpu_py3.8
           python_version: '3.8'
-
       - unittest_linux_stable_cpu:
           cu_version: cpu
           name: unittest_linux_stable_cpu_py3.8
           python_version: '3.8'
-
       - unittest_linux_stable_gpu:
           cu_version: cu113
           name: unittest_linux_stable_gpu_py3.8
           python_version: '3.8'
+
+      - unittest_macos_cpu:
+          cu_version: cpu
+          name: unittest_macos_cpu_py3.9
+          python_version: '3.9'
+      - unittest_linux_cpu:
+          cu_version: cpu
+          name: unittest_linux_cpu_py3.9
+          python_version: '3.9'
+      - unittest_linux_gpu:
+          cu_version: cu113
+          name: unittest_linux_gpu_py3.9
+          python_version: '3.9'
+      - unittest_linux_optdeps_gpu:
+          cu_version: cu113
+          name: unittest_linux_optdeps_gpu_py3.9
+          python_version: '3.9'
+      - unittest_linux_stable_cpu:
+          cu_version: cpu
+          name: unittest_linux_stable_cpu_py3.9
+          python_version: '3.9'
+      - unittest_linux_stable_gpu:
+          cu_version: cu113
+          name: unittest_linux_stable_gpu_py3.9
+          python_version: '3.9'
+
+      - unittest_macos_cpu:
+          cu_version: cpu
+          name: unittest_macos_cpu_py3.10
+          python_version: '3.10'
+      - unittest_linux_cpu:
+          cu_version: cpu
+          name: unittest_linux_cpu_py3.10
+          python_version: '3.10'
+      - unittest_linux_gpu:
+          cu_version: cu113
+          name: unittest_linux_gpu_py3.10
+          python_version: '3.10'
+      - unittest_linux_optdeps_gpu:
+          cu_version: cu113
+          name: unittest_linux_optdeps_gpu_py3.10
+          python_version: '3.10'
+      - unittest_linux_stable_cpu:
+          cu_version: cpu
+          name: unittest_linux_stable_cpu_py3.10
+          python_version: '3.10'
+      - unittest_linux_stable_gpu:
+          cu_version: cu113
+          name: unittest_linux_stable_gpu_py3.10
+          python_version: '3.10'
diff --git a/.circleci/unittest/linux/scripts/environment.yml b/.circleci/unittest/linux/scripts/environment.yml
@@ -8,7 +8,6 @@ dependencies:
     - hypothesis
     - future
     - cloudpickle
-    - gym_retro
     - gym
     - pygame
     - gym[accept-rom-license]
@@ -24,5 +23,4 @@ dependencies:
     - dm_control
     - mujoco_py
     - hydra-core
-    - pyrender
     - tensorboard
diff --git a/.circleci/unittest/linux/scripts/install.sh b/.circleci/unittest/linux/scripts/install.sh
@@ -36,18 +36,14 @@ printf "Installing PyTorch with %s\n" "${CU_VERSION}"
 if [ "${CU_VERSION:-}" == cpu ] ; then
     # conda install -y pytorch torchvision cpuonly -c pytorch-nightly
     # use pip to install pytorch as conda can frequently pick older release
-    if [[ $OSTYPE == 'darwin'* ]]; then
-      pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu
-    else
-      pip3 install torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html --pre
-    fi
+    conda install -y pytorch torchvision torchaudio cpuonly -c pytorch-nightly
 else
-    pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cu113
+    conda install -y pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch-nightly
 fi
 
 printf "Installing functorch\n"
-pip install ninja  # Makes the build go faster
-pip install "git+https://github.com/pytorch/functorch.git"
+python -m pip install ninja  # Makes the build go faster
+python -m pip install "git+https://github.com/pytorch/functorch.git"
 
 # smoke test
 python -c "import functorch"

diff --git a/.circleci/unittest/linux_optdeps/scripts/install.sh b/.circleci/unittest/linux_optdeps/scripts/install.sh
@@ -36,18 +36,14 @@ printf "Installing PyTorch with %s\n" "${CU_VERSION}"
 if [ "${CU_VERSION:-}" == cpu ] ; then
     # conda install -y pytorch torchvision cpuonly -c pytorch-nightly
     # use pip to install pytorch as conda can frequently pick older release
-    if [[ $OSTYPE == 'darwin'* ]]; then
-      pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu
-    else
-      pip3 install torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html --pre
-    fi
+    conda install -y pytorch cpuonly -c pytorch-nightly
 else
-    pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cu113
+    conda install -y pytorch cudatoolkit=11.3 -c pytorch-nightly
 fi
 
 printf "Installing functorch\n"
-pip install ninja  # Makes the build go faster
-pip install "git+https://github.com/pytorch/functorch.git"
+python -m pip install ninja  # Makes the build go faster
+python -m pip install "git+https://github.com/pytorch/functorch.git"
 
 # smoke test
 python -c "import functorch"

diff --git a/.circleci/unittest/linux_stable/scripts/environment.yml b/.circleci/unittest/linux_stable/scripts/environment.yml
@@ -9,7 +9,6 @@ dependencies:
     - hypothesis
     - future
     - cloudpickle
-    - gym_retro
     - gym
     - pygame
     - gym[accept-rom-license]
@@ -25,5 +24,4 @@ dependencies:
     - dm_control
     - mujoco_py
     - hydra-core
-    - pyrender
     - tensorboard
diff --git a/.circleci/unittest/linux_stable/scripts/install.sh b/.circleci/unittest/linux_stable/scripts/install.sh
@@ -37,13 +37,13 @@ printf "Installing PyTorch with %s\n" "${CU_VERSION}"
 if [ "${CU_VERSION:-}" == cpu ] ; then
     # conda install -y pytorch torchvision cpuonly -c pytorch-nightly
     # use pip to install pytorch as conda can frequently pick older release
-    pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
+    conda install -y pytorch torchvision torchaudio cpuonly -c pytorch
 else
-    pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
+    conda install -y pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
 fi
 
 printf "Installing functorch\n"
-pip install git+https://github.com/pytorch/functorch.git@release/0.2
+python -m pip install git+https://github.com/pytorch/functorch.git@release/0.2
 
 # smoke test
 python -c "import functorch"

diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml
@@ -0,0 +1,130 @@
+name: Wheels
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
+  push:
+    branches:
+      - release/0.0.1
+
+jobs:
+
+  build-wheel-linux:
+    runs-on: ubuntu-18.04
+    strategy:
+      matrix:
+        python_version: [["3.7", "cp37-cp37m"], ["3.8", "cp38-cp38"], ["3.9", "cp39-cp39"], ["3.10", "cp310-cp310"]]
+        cuda_support: [["", "--extra-index-url https://download.pytorch.org/whl/cpu", "\"['cpu', '11.3', '11.6']\"", "cpu"], ["+cu102", "", "\"['10.2']\"", "cuda102"]]
+    container: pytorch/manylinux-${{ matrix.cuda_support[3] }}
+    steps:
+      - name: Checkout torchrl
+        uses: actions/checkout@v2
+      - name: Install PyTorch 1.12 RC
+        run: |
+          export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH"
+          python3 -mpip install torch==1.12 ${{ matrix.cuda_support[1] }}
+          python3 -mpip install "git+https://github.com/pytorch/functorch.git@release/0.2"
+      - name: Build wheel
+        run: |
+          export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH"
+          python3 -mpip install wheel
+          VERSION_TAG=${{ matrix.cuda_support[0] }} PYTORCH_CUDA_RESTRICTIONS=${{ matrix.cuda_support[2] }} python3 setup.py bdist_wheel
+          # NB: wheels have the linux_x86_64 tag so we rename to manylinux1
+          # find . -name 'dist/*whl' -exec bash -c ' mv $0 ${0/linux/manylinux1}' {} \;
+      # pytorch/pytorch binaries are also manylinux_2_17 compliant but they
+      # pretend that they're manylinux1 compliant so we do the same.
+      - name: Show auditwheel output; confirm 2-17
+        run: |
+          python3 -mpip install auditwheel
+          auditwheel show dist/*
+      - name: Upload wheel for the test-wheel job
+        uses: actions/upload-artifact@v2
+        with:
+          name: torchrl-linux-${{ matrix.python_version[0] }}.whl
+          path: dist/torchrl-0.0.1-*.whl
+      - name: Upload wheel for download
+        uses: actions/upload-artifact@v2
+        with:
+          name: torchrl-batch.whl
+          path: dist/*.whl
+
+  build-wheel-mac:
+    runs-on: macos-latest
+    strategy:
+      matrix:
+        python_version: [["3.7", "3.7"], ["3.8", "3.8"], ["3.9", "3.9"], ["3.10", "3.10.3"]]
+    steps:
+      - name: Setup Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python_version[1] }}
+          architecture: x64
+      - name: Checkout torchrl
+        uses: actions/checkout@v2
+      - name: Install PyTorch 1.12 RC
+        run: |
+          python3 -mpip install torch==1.12 --extra-index-url https://download.pytorch.org/whl/cpu
+          python3 -mpip install "git+https://github.com/pytorch/functorch.git@release/0.2"
+      - name: Build wheel
+        run: |
+          export CC=clang CXX=clang++
+          python3 -mpip install wheel
+          python3 setup.py bdist_wheel
+      - name: Upload wheel for the test-wheel job
+        uses: actions/upload-artifact@v2
+        with:
+          name: torchrl-mac-${{ matrix.python_version[0] }}.whl
+          path: dist/*.whl
+      - name: Upload wheel for download
+        uses: actions/upload-artifact@v2
+        with:
+          name: torchrl-batch.whl
+          path: dist/*.whl
+
+  test-wheel:
+    needs: [build-wheel-linux, build-wheel-mac]
+    strategy:
+      matrix:
+        os: [["linux", "ubuntu-18.04"], ["mac", "macos-latest"]]
+        python_version: [ "3.7", "3.8", "3.9", "3.10" ]
+    runs-on: ${{ matrix.os[1] }}
+    steps:
+      - name: Setup Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python_version }}
+          architecture: x64
+      - name: Checkout torchrl
+        uses: actions/checkout@v2
+      - name: Install PyTorch 1.12 RC
+        run: |
+          python3 -mpip install torch==1.12 torchvision --extra-index-url https://download.pytorch.org/whl/cpu
+          python3 -mpip install "git+https://github.com/pytorch/functorch.git@release/0.2"
+      - name: Upgrade pip
+        run: |
+          python3 -mpip install --upgrade pip
+      - name: Install test dependencies
+        run: |
+          python3 -mpip install numpy pytest pytest-cov codecov unittest-xml-reporting pillow>=4.1.1 scipy av networkx expecttest pyyaml
+      - name: Download built wheels
+        uses: actions/download-artifact@v2
+        with:
+          name: torchrl-${{ matrix.os[0] }}-${{ matrix.python_version }}.whl
+          path: /tmp/wheels
+      - name: Install built wheels
+        run: |
+          python3 -mpip install /tmp/wheels/*
+      - name: Log version string
+        run: |
+          # Avoid ambiguity of "import torchrl" by deleting the source files.
+          rm -rf torchrl/
+          python -c "import torchrl; print(torchrl.__version__)"
+      - name: Run tests
+        run: |
+          set -e
+          export IN_CI=1
+          mkdir test-reports
+          python -m torch.utils.collect_env
+          python -c "import torchrl; print(torchrl.__version__)"
+          EXIT_STATUS=0
+          pytest test/smoke_test.py -v --durations 20
+          exit $EXIT_STATUS
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
@@ -7,11 +7,11 @@
 import os.path
 import re
 
-import numpy as np
 import pytest
 import torch
 from _utils_internal import get_available_devices
 from torch import multiprocessing as mp
+from torchrl import prod
 from torchrl.data import SavedTensorDict, TensorDict, MemmapTensor
 from torchrl.data.tensordict.tensordict import (
     assert_allclose_td,
@@ -810,7 +810,7 @@ def test_view(self, td_name):
         td = getattr(self, td_name)
         td_view = td.view(-1)
         tensor = td.get("a")
-        tensor = tensor.view(-1, tensor.numel() // np.prod(td.batch_size))
+        tensor = tensor.view(-1, tensor.numel() // prod(td.batch_size))
         tensor = torch.ones_like(tensor)
         if td_name == "sub_td":
             td_view.set_("a", tensor)

diff --git a/test/test_transforms.py b/test/test_transforms.py
@@ -3,7 +3,6 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 import argparse
-import math
 from multiprocessing import Lock
 
 import pytest
@@ -12,6 +11,7 @@
 from mocking_classes import ContinuousActionVecMockEnv
 from torch import Tensor
 from torch import multiprocessing as mp
+from torchrl import prod
 from torchrl.data import NdBoundedTensorSpec, CompositeSpec
 from torchrl.data import TensorDict
 from torchrl.envs import EnvCreator, SerialEnv
@@ -448,7 +448,7 @@ def test_flatten(self, keys, size, nchannels, batch, device):
         )
         td.set("dont touch", dont_touch.clone())
         flatten(td)
-        expected_size = math.prod(size + [nchannels])
+        expected_size = prod(size + [nchannels])
         for key in keys:
             assert td.get(key).shape[-3] == expected_size
         assert (td.get("dont touch") == dont_touch).all()

diff --git a/torchrl/__init__.py b/torchrl/__init__.py
@@ -5,7 +5,10 @@
 
 import abc
 import collections
+import math
 import time
+import typing
+from typing import Optional, Type, Tuple
 from warnings import warn
 
 import numpy as np
@@ -112,3 +115,10 @@ def __init__(self, fun=lambda x: x):
     def __missing__(self, key):
         value = self.fun(key)
         return value
+
+
+def prod(sequence):
+    if hasattr(math, "prod"):
+        return math.prod(sequence)
+    else:
+        return int(np.prod(sequence))