Updated workflows/tpu-tests.yaml

.github/workflows/tpu-tests.yaml

@@ -4,54 +4,46 @@ on: [pull_request]
 jobs:
   tpu-tests:
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        xla-version: [1.5, nightly]
+
     steps:
     - uses: actions/checkout@v2
-    - name: Setup Conda
-      run: |
-        wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh
-        bash miniconda.sh -b -p $HOME/miniconda
-        export PATH="$HOME/miniconda/bin:$PATH"
-        conda config --set always_yes yes --set changeps1 no
-        conda update -q conda
-        # Useful for debugging any issues with conda
-        conda info -a
-        conda create -y -n py36 python=3.6
-        source activate py36
-
+    - name: Set up Python 3.6
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.6'
+        architecture: 'x64'
 
     - name: Install Torch XLA and others
       run: |
-        export PATH="$HOME/miniconda/bin:$PATH"
-        source activate py36
-
+
         ## Install gsutil
         echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | sudo tee -a /etc/apt/sources.list.d/google-cloud-sdk.list
         sudo apt-get install -y apt-transport-https ca-certificates gnupg curl
         curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key --keyring /usr/share/keyrings/cloud.google.gpg add -
         sudo apt-get update && sudo apt-get install -y google-cloud-sdk
 
         ## Install openblas and mkl
-        sudo apt-get install -y libopenblas-dev
-        conda install -y mkl=2020.0=166  # if change the version, think about to change LD_LIBRARY_PATH below
+        sudo apt-get install -y libopenblas-dev libomp5
+        pip install mkl
 
         ## Download torch & xla
-        gsutil cp gs://tpu-pytorch/wheels/torch-1.5-cp36-cp36m-linux_x86_64.whl .
-        gsutil cp gs://tpu-pytorch/wheels/torch_xla-1.5-cp36-cp36m-linux_x86_64.whl .
+        gsutil cp gs://tpu-pytorch/wheels/torch-${{ matrix.xla-version }}-cp36-cp36m-linux_x86_64.whl .
+        gsutil cp gs://tpu-pytorch/wheels/torch_xla-${{ matrix.xla-version }}-cp36-cp36m-linux_x86_64.whl .
 
         ## Install torch & xla
-        pip install torch-1.5-cp36-cp36m-linux_x86_64.whl
-        pip install torch_xla-1.5-cp36-cp36m-linux_x86_64.whl
+        pip install torch-${{ matrix.xla-version }}-cp36-cp36m-linux_x86_64.whl
+        pip install torch_xla-${{ matrix.xla-version }}-cp36-cp36m-linux_x86_64.whl
 
         ## Install test deps and Ignite
         pip install -r requirements-dev.txt
         python setup.py install
 
     - name: Run Tests
       run: |
-        export PATH="$HOME/miniconda/bin:$PATH"
-        source activate py36
-
-        export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/miniconda/pkgs/mkl-2020.0-166/lib/
+        export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/hostedtoolcache/Python/3.6.10/x64/lib
         export XRT_DEVICE_MAP="CPU:0;/job:localservice/replica:0/task:0/device:XLA_CPU:0"
         export XRT_WORKERS="localservice:0;grpc://localhost:40934"