[workflow] updated release bdist workflow (#1318)

* [workflow] updated release bdist workflow

* polish workflow

* polish workflow

.github/workflows/build.yml

@@ -14,7 +14,7 @@ jobs:
         contains( github.event.pull_request.labels.*.name, 'Run Build and Test')
     runs-on: [self-hosted, gpu]
     container:
-      image: hpcaitech/pytorch-cuda:1.10.1-11.3.0
+      image: hpcaitech/pytorch-cuda:1.12.0-11.3.0
       options: --gpus all --rm -v /data/scratch/cifar-10:/data/scratch/cifar-10
     timeout-minutes: 40
     steps:

.github/workflows/release_bdist.yml

@@ -3,16 +3,15 @@ name: Release bdist wheel
 on:
   workflow_dispatch:
     inputs:
-      cuda_version:
-        type: choice
-        description: CUDA Version
-        default: 'all'
+      torch_version:
+        type: string
+        description: torch version, separated by comma
+        required: true
+        default: "all"
+      cuda_version:
+        type: string
+        description: cuda version, separated by comma
         required: true
-        options:
-        - all
-        - "11.3"
-        - "11.1"
-        - "10.2"
       github_ref:
         type: string
         description: Branch or Tag
@@ -27,12 +26,24 @@ jobs:
       matrix: ${{ steps.set-matrix.outputs.matrix }}
     steps:
     - id: set-matrix
+      env:
+        TORCH_VERSIONS: ${{ inputs.torch_version }}
+        CUDA_VERSIONS: ${{ inputs.cuda_version }}
       run: |
-        [ "${{github.event.inputs.cuda_version}}" != "" ] && matrix="[\"hpcaitech/cuda-conda:${{github.event.inputs.cuda_version}}\"]"
-        [ "${{github.event.inputs.cuda_version}}" == "" ] || [ "${{github.event.inputs.cuda_version}}" == "all" ] && \
-          matrix="[\"hpcaitech/cuda-conda:11.3\", \"hpcaitech/cuda-conda:11.1\", \"hpcaitech/cuda-conda:10.2\"]"
-        echo $matrix
-        echo "::set-output name=matrix::{\"container\":$(echo $matrix)}"
+        echo $TORCH_VERSIONS
+        echo $CUDA_VERSIONS
+        IFS=','
+        DOCKER_IMAGE=()
+
+        for cv in $CUDA_VERSIONS
+        do
+            DOCKER_IMAGE+=("\"hpcaitech/cuda-conda:${cv}\"")
+        done
+
+        container=$( IFS=',' ; echo "${DOCKER_IMAGE[*]}" )
+        container="[${container}]"
+        echo "$container"
+        echo "::set-output name=matrix::{\"container\":$(echo "$container")}"
 
   build:
     name: Release bdist wheels
@@ -62,7 +73,9 @@ jobs:
       - name: Build bdist wheel
         run: |
           pip install beautifulsoup4 requests packaging
-          python ./build_colossalai_wheel.py
+          python ./build_colossalai_wheel.py --torch_version $TORCH_VERSIONS
+        env:
+          TORCH_VERSIONS: ${{ inputs.torch_version }}
       - name: 🚀 Deploy
         uses: garygrossgarten/github-action-scp@release
         with:

.github/workflows/scripts/build_colossalai_wheel.py

@@ -15,6 +15,7 @@ CUDA_HOME = os.environ['CUDA_HOME']
 
 def parse_args():
     parser = argparse.ArgumentParser()
+    parser.add_argument('--torch_version', type=str)
     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()
@@ -81,29 +82,27 @@ def main():
     args = parse_args()
     wheel_info = all_wheel_info()
 
+    # filter wheels on condition
+    all_torch_versions = list(wheel_info.keys())
+    def _compare_version(a, b):
+        if version.parse(a) > version.parse(b):
+            return 1
+        else:
+            return -1
+
+    all_torch_versions.sort(key=cmp_to_key(_compare_version))
+
     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]:
+        for key in all_torch_versions[:-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())
+    elif args.torch_version != 'all':
+        torch_versions = args.torch_version.split(',')
+        # only keep the torch versions specified
+        for key in all_torch_versions:
+            if key not in torch_versions:
+                wheel_info.pop(key)
 
-                for key in python_versions:
-                    if key != '3.8':
-                        python_versioned_info.pop(key)
     build_colossalai(wheel_info)
 
 if __name__ == '__main__':