ci/slurm: use ec2 instance connect + mssh instead of using SSH keys (#265)

Summary:
This switches the integration tests to use ec2 instance connect w/ an assumed role instead of embedding the slurm ssh key in GitHub secrets.

Pull Request resolved: #265

Test Plan:
```
$ env SLURM_INSTANCE_MASTER=ubuntu@i-01dd4b95724eb0b4b scripts/slurmint.sh
```

CI

Reviewed By: kiukchung, aivanou

Differential Revision: D31695261

Pulled By: d4l3k

fbshipit-source-id: 48a52e911e68bc9b18ed470a5f7e725ff58697b1

Author: d4l3k

.github/workflows/slurm-integration-tests.yaml

@@ -9,6 +9,9 @@ on:
 jobs:
   slurm:
     runs-on: ubuntu-18.04
+    permissions:
+      id-token: write
+      contents: read
     steps:
       - name: Setup Python
         uses: actions/setup-python@v2
@@ -17,21 +20,32 @@ jobs:
           architecture: x64
       - name: Checkout TorchX
         uses: actions/checkout@v2
+      - name: Configure AWS
+        env:
+          AWS_ROLE_ARN: ${{ secrets.AWS_ROLE_ARN }}
+        run: |
+          if [ -n "$AWS_ROLE_ARN" ]; then
+            export AWS_WEB_IDENTITY_TOKEN_FILE=/tmp/awscreds
+            export AWS_DEFAULT_REGION=us-west-2
+
+            echo AWS_WEB_IDENTITY_TOKEN_FILE=$AWS_WEB_IDENTITY_TOKEN_FILE >> $GITHUB_ENV
+            echo AWS_ROLE_ARN=$AWS_ROLE_ARN >> $GITHUB_ENV
+            echo AWS_DEFAULT_REGION=$AWS_DEFAULT_REGION >> $GITHUB_ENV
+
+            curl -H "Authorization: bearer $ACTIONS_ID_TOKEN_REQUEST_TOKEN" "$ACTIONS_ID_TOKEN_REQUEST_URL" | jq -r '.value' > $AWS_WEB_IDENTITY_TOKEN_FILE
+          fi
       - name: Install Dependencies
         run:
           set -ex
 
-          pip install wheel
+          pip install wheel ec2instanceconnectcli
       - name: Run Slurm Integration Tests
         env:
-          SLURM_SSH: ${{ secrets.SLURM_SSH }}
-          SLURM_MASTER: ${{ secrets.SLURM_MASTER }}
+          SLURM_INSTANCE_MASTER: ${{ secrets.SLURM_INSTANCE_MASTER }}
           SLURM_KNOWN_HOST: ${{ secrets.SLURM_KNOWN_HOST }}
-          SLURM_IDENT: id_rsa
         run: |
           set -e
-          echo "$SLURM_SSH" > "$SLURM_IDENT"
-          chmod 600 "$SLURM_IDENT"
+
           mkdir -p ~/.ssh
           echo "$SLURM_KNOWN_HOST" >> ~/.ssh/known_hosts
 

scripts/slurmint.sh

@@ -14,8 +14,8 @@ python setup.py bdist_wheel
 
 WHEEL="$DIST/$(ls $DIST)"
 
-if [[ -z "${SLURM_MASTER}" ]]; then
-    echo "slurm master is not set, skipping test..."
+if [[ -z "${SLURM_INSTANCE_MASTER}" ]]; then
+    echo "SLURM_INSTANCE_MASTER is not set, skipping test..."
     exit 0
 fi
 
@@ -25,11 +25,11 @@ VENV="$DIR/venv"
 
 function run_cmd {
     # shellcheck disable=SC2048,SC2086
-    ssh -o ServerAliveInterval=60 "$SLURM_MASTER" -i "$SLURM_IDENT" $*
+    mssh -o ServerAliveInterval=60 "$SLURM_INSTANCE_MASTER" -- $*
 }
 
 function run_scp {
-    scp -i "$SLURM_IDENT" "$1" "$SLURM_MASTER:$2"
+    rsync -rav -e mssh "$1" "$SLURM_INSTANCE_MASTER:$2"
 }
 
 function cleanup {