From 1a76c88aba4ed1ba66e2420eec24122363ab30ba Mon Sep 17 00:00:00 2001
From: Frank Lee <somerlee.9@gmail.com>
Date: Tue, 24 May 2022 17:56:01 +0800
Subject: [PATCH] [ci] added nightly build (#1018) (#1019)

---
 .github/workflows/release_nightly.yml         | 81 +++++++++++++++++++
 .../scripts/build_colossalai_wheel.py         | 37 ++++++++-
 2 files changed, 117 insertions(+), 1 deletion(-)
 create mode 100644 .github/workflows/release_nightly.yml

diff --git a/.github/workflows/release_nightly.yml b/.github/workflows/release_nightly.yml
new file mode 100644
index 000000000..8c8f219fd
--- /dev/null
+++ b/.github/workflows/release_nightly.yml
@@ -0,0 +1,81 @@
+name: Release bdist wheel for Nightly versions
+
+on:
+  schedule:
+    # run at 00:00 of every Sunday 
+    - cron:  '0 0 * * 6'
+  workflow_dispatch:
+    inputs:
+      cuda_version:
+        type: choice
+        description: CUDA Version
+        default: 'all'
+        required: true
+        options:
+        - all
+        - "11.3"
+        - "10.2"
+jobs:
+  matrix_preparation:
+    name: Prepare Container List
+    runs-on: ubuntu-latest
+    outputs:
+      matrix: ${{ steps.set-matrix.outputs.matrix }}
+    steps:
+    - id: set-matrix
+      run: |
+        [ "${{github.event.inputs.cuda_version}}" != "" ] && matrix="[\"hpcaitech/cuda-conda:${{github.event.inputs.cuda_version}}\"]"
+        [ "${{github.event.inputs.cuda_version}}" == "" ] || [ "${{github.event.inputs.version}}" == "all" ] && \
+          matrix="[\"hpcaitech/cuda-conda:11.3\", \"hpcaitech/cuda-conda:10.2\"]"
+        echo $matrix
+        echo "::set-output name=matrix::{\"container\":$(echo $matrix)}"
+
+  build:
+    name: Release bdist wheels
+    needs: matrix_preparation
+    if: github.repository == 'hpcaitech/ColossalAI' && contains(fromJson('["FrankLeeeee", "ver217", "feifeibear", "kurisusnowdeng"]'), github.actor)
+    runs-on: [self-hosted, gpu]
+    strategy:
+      fail-fast: false
+      matrix: ${{fromJson(needs.matrix_preparation.outputs.matrix)}}
+    container:
+      image: ${{ matrix.container }}
+      options: --gpus all --rm
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+      # cub is for cuda 10.2
+      - name: Copy scripts and checkout
+        run: |
+          cp -r ./.github/workflows/scripts/* ./
+          ln -s /github/home/pip_wheels ./pip_wheels
+          wget https://github.com/NVIDIA/cub/archive/refs/tags/1.8.0.zip
+          unzip 1.8.0.zip
+      - name: Build bdist wheel
+        run: |
+          pip install beautifulsoup4 requests packaging
+          python ./build_colossalai_wheel.py --nightly
+      - name: 🚀 Deploy
+        uses: garygrossgarten/github-action-scp@release
+        with:
+          local: all_dist
+          remote: ${{ secrets.PRIVATE_PYPI_NIGHTLY_DIR }}
+          host: ${{ secrets.PRIVATE_PYPI_HOST }}
+          username: ${{ secrets.PRIVATE_PYPI_USER }}
+          password: ${{ secrets.PRIVATE_PYPI_PASSWD }}
+  remove_old_build:
+    name: Remove old nightly build
+    needs: build
+    steps:
+      - name: executing remote ssh commands using password
+        uses: appleboy/ssh-action@master
+        with:
+          host: ${{ secrets.PRIVATE_PYPI_HOST }}
+          username: ${{ secrets.PRIVATE_PYPI_USER }}
+          password: ${{ secrets.PRIVATE_PYPI_PASSWD }}
+          script: |
+            cd $build_dir
+            find . -type f -mtime +0 -exec rm -f {} +
+        env:
+          build_dir: ${{ secrets.PRIVATE_PYPI_NIGHTLY_DIR }}
diff --git a/.github/workflows/scripts/build_colossalai_wheel.py b/.github/workflows/scripts/build_colossalai_wheel.py
index c6b7acc25..7bd79fb64 100644
--- a/.github/workflows/scripts/build_colossalai_wheel.py
+++ b/.github/workflows/scripts/build_colossalai_wheel.py
@@ -1,14 +1,24 @@
+from filecmp import cmp
 import requests
 from bs4 import BeautifulSoup
-import re
+import argparse
 import os
 import subprocess
+from packaging import version
+from functools import cmp_to_key
 
 
 WHEEL_TEXT_ROOT_URL = 'https://github.com/hpcaitech/public_assets/tree/main/colossalai/torch_build/torch_wheels'
 RAW_TEXT_FILE_PREFIX = 'https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/torch_build/torch_wheels'
 CUDA_HOME = os.environ['CUDA_HOME']
 
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--nightly', action='store_true', 
+        help='whether this build is for nightly release, if True, will only build on the latest PyTorch version and Python 3.8')
+    return parser.parse_args()
+
 def get_cuda_bare_metal_version():
     raw_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"], universal_newlines=True)
     output = raw_output.split()
@@ -69,7 +79,32 @@ def build_colossalai(wheel_info):
                     os.system(cmd)
 
 def main():
+    args = parse_args()
     wheel_info = all_wheel_info()
+
+    if args.nightly:
+        latest_torch_version = list(wheel_info.keys())
+
+        def _compare_version(a, b):
+            if version.parse(a) > version.parse(b):
+                return 1
+            else:
+                return -1
+
+        latest_torch_version.sort(key=cmp_to_key(_compare_version))
+        
+        # only keep the latest version
+        for key in latest_torch_version[:-1]:
+            wheel_info.pop(key)
+        
+        # we only keep python 3.8 for nightly release
+        for torch_version, cuda_versioned_info in wheel_info.items():
+            for cuda_version, python_versioned_info in cuda_versioned_info.items():
+                python_versions = list(python_versioned_info.keys())
+
+                for key in python_versions:
+                    if key != '3.8':
+                        python_versioned_info.pop(key)
     build_colossalai(wheel_info)
 
 if __name__ == '__main__':