From 0372ed7951d605e965f143a0e09a995d533010ee Mon Sep 17 00:00:00 2001
From: Frank Lee <somerlee.9@gmail.com>
Date: Thu, 7 Apr 2022 17:53:03 +0800
Subject: [PATCH] [ci] update workflow trigger condition and support options
 (#691)

---
 .github/workflows/compatibility_test.yml | 48 ++++++++++++++++--------
 1 file changed, 32 insertions(+), 16 deletions(-)

diff --git a/.github/workflows/compatibility_test.yml b/.github/workflows/compatibility_test.yml
index e41ae9379..d88fa9f49 100644
--- a/.github/workflows/compatibility_test.yml
+++ b/.github/workflows/compatibility_test.yml
@@ -1,23 +1,42 @@
 name: Compatibility Test
 
-on: workflow_dispatch
+on:
+  workflow_dispatch:
+    version:
+        type: choice
+        description: which version to test
+        required: true
+        options:
+        - all
+        - pytorch-cuda:1.9.0-11.1.1 # python 3.8
+        - pytorch-cuda:1.8.1-11.1.1 # python 3.8
+        - pytorch-cuda:1.7.1-11.0.3 # python 3.8
+        - pytorch-cuda:1.6.0-10.2 # python 3.6
+        
 
 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.version}}" != "" ] && matrix="[\"frankleeeee/${{github.event.inputs.version}}\"]"
+        [ "${{github.event.inputs.version}}" == "" ] || [ "${{github.event.inputs.version}}" == "all" ] && \
+          matrix="[\"frankleeeee/pytorch-cuda:1.9.0-11.1.1\", \"frankleeeee/pytorch-cuda:1.8.1-11.1.1\", \"frankleeeee/pytorch-cuda:1.7.1-11.0.3\", \"frankleeeee/pytorch-cuda:1.6.0-10.2\"]"
+        echo $matrix
+        echo "::set-output name=matrix::{\"container\":$(echo $matrix)}"
+
   build:
-    name: Test for PyTorch compatibility
-    if: github.base_ref == 'main' && github.repository == 'hpcaitech/ColossalAI'
+    name: Test for PyTorch Compatibility
+    needs: matrix_preparation
+    if: github.repository == 'hpcaitech/ColossalAI'
     runs-on: [self-hosted, gpu]
     strategy:
       fail-fast: false
-      matrix:
-        # pytorch-cuda:1.9.0-11.1.1 + Python 3.8
-        # pytorch-cuda:1.8.1-11.1.1 + Python 3.8
-        # pytorch-cuda:1.7.1-11.0.3 + Python 3.8
-        # pytorch-cuda:1.6.0-10.2 + Python 3.6
-        container: ["frankleeeee/pytorch-cuda:1.9.0-11.1.1",
-                    "frankleeeee/pytorch-cuda:1.8.1-11.1.1",
-                    "frankleeeee/pytorch-cuda:1.7.1-11.0.3",
-                    "frankleeeee/pytorch-cuda:1.6.0-10.2"]
+      matrix: ${{fromJson(needs.matrix_preparation.outputs.matrix)}}
     container:
       image: ${{ matrix.container }}
       options: --gpus all --rm --ipc=host -v /data/scratch/cifar-10:/data/scratch/cifar-10
@@ -27,7 +46,7 @@ jobs:
         run: |
           pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
           pip install -U pip setuptools wheel --user
-          pip install pytest tensorboard
+          pip install pytest tensorboard transformers
       - uses: actions/checkout@v2
       - name: Install Colossal-AI      
         run: |
@@ -38,6 +57,3 @@ jobs:
           PYTHONPATH=$PWD pytest tests
         env:
           DATA: /data/scratch/cifar-10
-      
-
-