[ci] cache cuda extension (#860)

.github/workflows/build.yml

@@ -18,18 +18,17 @@ jobs:
       options: --gpus all --rm -v /data/scratch/cifar-10:/data/scratch/cifar-10
     timeout-minutes: 40
     steps:
-      - name: Install dependencies
-        run: |
-          pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
-          pip install -U pip setuptools wheel --user
       - uses: actions/checkout@v2
         with:
           ssh-key: ${{ secrets.SSH_KEY_FOR_CI }}
       - name: Install Colossal-AI      
         run: |
+          [ ! -z "$(ls -A /github/home/cuda_ext_cache/)" ] && cp -r /github/home/cuda_ext_cache/* /__w/ColossalAI/ColossalAI/
           pip install -r requirements/requirements.txt
           pip install -r requirements/requirements-test.txt
-          pip install -v --no-cache-dir .
+          pip install -v -e .
+          cp -r /__w/ColossalAI/ColossalAI/build /github/home/cuda_ext_cache/
+          cp /__w/ColossalAI/ColossalAI/*.so /github/home/cuda_ext_cache/
       - name: Unit Testing
         run: |
           PYTHONPATH=$PWD pytest tests

tests/components_to_test/simple_net.py

@@ -1,12 +1,11 @@
-from zmq import device
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from colossalai.nn import CheckpointModule
 from .utils.dummy_data_generator import DummyDataGenerator
 from .registry import non_distributed_component_funcs
 from colossalai.utils.cuda import get_current_device
 
+
 class SimpleNet(CheckpointModule):
     """
     In this no-leaf module, it has subordinate nn.modules and a nn.Parameter.