Increase TPU CI node count from 2 to 32 (#8857)

Author: tengyifei

infra/terraform_modules/arc_v4_container_cluster/arc-values.yaml

@@ -1,6 +1,6 @@
 githubConfigUrl: ${github_repo_url}
 githubConfigSecret: github-pat
-minRunners: 1
+minRunners: ${min_tpu_nodes}
 maxRunners: ${max_tpu_nodes}
 template:
   spec:

infra/terraform_modules/arc_v4_container_cluster/main.tf

@@ -17,7 +17,7 @@ resource "google_container_cluster" "arc_v4_cluster" {
   location = "us-central2"
 
   remove_default_node_pool = true
-  initial_node_count       = 1
+  initial_node_count       = var.min_tpu_nodes
 
   release_channel {
     channel = "RAPID"
@@ -53,7 +53,7 @@ resource "google_container_node_pool" "arc_v4_tpu_nodes" {
   cluster            = google_container_cluster.arc_v4_cluster.name
   initial_node_count = 1
   autoscaling {
-    total_min_node_count = 1
+    total_min_node_count = var.min_tpu_nodes
     total_max_node_count = var.max_tpu_nodes
     location_policy      = "ANY"
   }

infra/terraform_modules/arc_v4_container_cluster/variables.tf

@@ -18,6 +18,11 @@ variable "tpu_nodepool_name" {
   type        = string
 }
 
+variable "min_tpu_nodes" {
+  description = "Minimum number of TPU nodes and runners"
+  type        = number
+}
+
 variable "max_tpu_nodes" {
   description = "Maximum number of TPU nodes and runners"
   type        = number

infra/tpu-pytorch/tpu_ci.tf

@@ -5,7 +5,8 @@ module "v4_arc_cluster" {
   cpu_nodepool_name = "cpu-nodepool"
   cpu_node_count    = 1
   tpu_nodepool_name = "tpu-nodepool"
-  max_tpu_nodes     = 2
+  min_tpu_nodes     = 32
+  max_tpu_nodes     = 32
   github_repo_url   = "https://github.com/pytorch/xla"
   # Dockerfile for this image can be found at test/tpu/Dockerfile
   runner_image      = "gcr.io/tpu-pytorch/tpu-ci-runner:latest"