Add Setup ssh to conda and wheel builds

.github/actions/wait-for-ssh/action.yml

@@ -0,0 +1,30 @@
+name: Wait for SSH session on Linux
+
+description: Stuff that should always run at the end of a linux job
+# test
+inputs:
+  skip-wait-ssh:
+    description: If set, don't wait for ssh to drain before tearing down
+    required: false
+    default: ""
+
+runs:
+  using: composite
+  steps:
+    - name: Hold runner for 2 hours or until ssh sessions have drained
+      # TODO working-directory: !{{ pytorch_directory }}
+      # Always hold for active ssh sessions
+      shell: bash
+      if: inputs.skip-wait-ssh == ''
+      run: |
+        set -eou pipefail
+
+        echo "Holding runner for 2 hours until all ssh sessions have logged out"
+        for _ in $(seq 1440); do
+            # Break if no ssh session exists anymore
+            if [ "$(who)" = "" ]; then
+              break
+            fi
+            echo "."
+            sleep 5
+        done

.github/workflows/build_conda_linux.yml

@@ -94,6 +94,10 @@ jobs:
           repository: ${{ inputs.test-infra-repository }}
           ref: ${{ inputs.test-infra-ref }}
           path: test-infra
+      - name: Setup SSH
+        uses: ./test-infra/.github/actions/setup-ssh
+        with:
+          github-secret: ${{ github.token }}
       - uses: ./test-infra/.github/actions/set-channel
       - name: Set PYTORCH_VERSION
         if: ${{ env.CHANNEL == 'test' }}
@@ -218,7 +222,10 @@ jobs:
         with:
           name: ${{ env.ARTIFACT_NAME }}
           path: ${{ inputs.repository }}/distr/
-
+      - name: Wait for ssh session
+        if: ${{ always() }}
+        uses: ./test-infra/.github/actions/wait-for-ssh
+
   upload:
     needs: build
     uses: ./.github/workflows/_binary_conda_upload.yml

.github/workflows/build_conda_windows.yml

@@ -94,6 +94,12 @@ jobs:
           repository: ${{ inputs.test-infra-repository }}
           ref: ${{ inputs.test-infra-ref }}
           path: test-infra
+      - uses: ./test-infra/.github/actions/setup-ssh
+        name: Setup SSH
+        with:
+          github-secret: ${{ secrets.GITHUB_TOKEN }}
+          activate-with-label: false
+          instructions: "SSH with rdesktop using ssh -L 3389:localhost:3389 %%username%%@%%hostname%%"
       - uses: ./test-infra/.github/actions/set-channel
       - name: Set PYTORCH_VERSION
         if: ${{ env.CHANNEL == 'test' }}
@@ -249,6 +255,9 @@ jobs:
         with:
           name: ${{ env.ARTIFACT_NAME }}
           path: ${{ inputs.repository }}/distr/
+      - uses: ./test-infra/.github/actions/teardown-windows
+        if: always()
+        name: Teardown Windows
 
   upload:
     needs: build

.github/workflows/build_wheels_linux.yml

@@ -136,6 +136,10 @@ jobs:
           repository: "pytorch/builder"
           ref: "main"
           path: builder
+      - name: Setup SSH
+        uses: ./test-infra/.github/actions/setup-ssh
+        with:
+          github-secret: ${{ github.token }}
       - name: Set linux aarch64 CI
         if: ${{ inputs.architecture == 'aarch64' }}
         shell: bash -l {0}
@@ -260,6 +264,9 @@ jobs:
         with:
           name: ${{ env.ARTIFACT_NAME }}
           path: ${{ inputs.repository }}/dist/
+      - name: Wait for ssh session
+        if: ${{ always() }}
+        uses: ./test-infra/.github/actions/wait-for-ssh
 
   upload:
     needs: build