fix runner cancel issue

.github/workflows/_build_plugin.yml

@@ -11,14 +11,19 @@ on:
         type: string
         description: Runner type for the test
         default: linux.12xlarge
-
+      timeout-minutes:
+        required: false
+        type: number
+        description: Timeout in minutes for the build job
+        default: 120
     secrets:
       gcloud-service-key:
         required: true
         description: Secret to access Bazel build cache
 jobs:
   build:
     runs-on: ${{ inputs.runner }}
+    timeout-minutes: ${{ inputs.timeout-minutes }}
     container:
       image: ${{ inputs.dev-image }}
     env:

.github/workflows/_build_torch_with_cuda.yml

@@ -15,9 +15,15 @@ on:
         type: string
         description: Runner type for the test
         default: linux.12xlarge
+      timeout-minutes:
+        required: false
+        type: number
+        description: Timeout in minutes for the build job
+        default: 120
 jobs:
   build:
     runs-on: ${{ inputs.runner }}
+    timeout-minutes: ${{ inputs.timeout-minutes }}
     container:
       image: ${{ inputs.dev-image }}
     env:

.github/workflows/_build_torch_xla.yml

@@ -15,14 +15,19 @@ on:
         type: string
         description: Runner type for the test
         default: linux.12xlarge
-
+      timeout-minutes:
+        required: false
+        type: number
+        description: Timeout in minutes for the build job
+        default: 120
     secrets:
       gcloud-service-key:
         required: true
         description: Secret to access Bazel build cache
 jobs:
   build:
     runs-on: ${{ inputs.runner }}
+    timeout-minutes: ${{ inputs.timeout-minutes }}
     container:
       image: ${{ inputs.dev-image }}
     env:

.github/workflows/_docs.yml

@@ -16,7 +16,7 @@ on:
         required: true
 jobs:
   build-docs:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     timeout-minutes: 45
     container:
       image: ${{ inputs.dev-image }}

.github/workflows/_tpu_ci.yml

@@ -1,9 +1,16 @@
 name: TPU Integration Test
 on:
   workflow_call:
+    inputs:
+      timeout-minutes:
+        required: false
+        type: number
+        description: Timeout in minutes for the job run
+        default: 120
 jobs:
   tpu-test:
     runs-on: v4-runner-set
+    timeout-minutes: ${{ inputs.timeout-minutes }}
     steps:
       - name: Checkout actions
         uses: actions/checkout@v4

.github/workflows/build_and_test.yml

@@ -17,7 +17,7 @@ concurrency:
 jobs:
 
   get-torch-commit:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     outputs:
       torch_commit: ${{ steps.commit.outputs.torch_commit }}
     steps:
@@ -33,6 +33,7 @@ jobs:
     with:
       dev-image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:3.10_tpuvm
       torch-commit: ${{needs.get-torch-commit.outputs.torch_commit}}
+      timeout-minutes: 240
     secrets:
       gcloud-service-key: ${{ secrets.GCLOUD_SERVICE_KEY }}
 
@@ -47,6 +48,7 @@ jobs:
       torch-commit: ${{needs.get-torch-commit.outputs.torch_commit}}
       # note that to build a torch wheel with CUDA enabled, we do not need a GPU runner.
       runner: linux.24xlarge
+      timeout-minutes: 120
 
   build-cuda-plugin:
     name: "Build XLA CUDA plugin"
@@ -97,6 +99,8 @@ jobs:
     name: "TPU tests"
     uses: ./.github/workflows/_tpu_ci.yml
     needs: build-torch-xla
+    with:
+      timeout-minutes: 300 
     if: github.event_name == 'push' || github.event_name == 'pull_request'
 
   push-docs:

.github/workflows/lintercheck.yml

@@ -9,7 +9,7 @@ on:
 
 jobs:
   linter_check:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout repo
         uses: actions/checkout@v3

.github/workflows/openxla_pin_update_weekly.yml

@@ -6,7 +6,7 @@ on:
 
 jobs:
   update_and_create_pr:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     steps:
       - name: Checkout code
         uses: actions/checkout@v3

.github/workflows/torchax.yml

@@ -16,7 +16,7 @@ concurrency:
 
 jobs:
   torchax-cpu:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     strategy:
       matrix:
         python-version: ['3.10']