Skip to content

Commit d88f907

Browse files
shubpal07shubpal07
committed
DWS Flex start support in TPU v6e
Change-Id: I879423f59924badb4d28b4fb892e4701a4281364 Change-Id: Ic16fcbb5b0ed510960c97f80ac916e331026234e Change-Id: I436f5bb6cee13efd083965493d56dd12b03fbaca Change-Id: Id77229ac117ff851539372c1477e5665527bde16 Change-Id: Ifa27a86ad463a8d29ebdd927010c1dc0abbf644d Change-Id: I87dc2a3b2bc8d813496a85789c132096ef5d070c Change-Id: I1c1cc689b0a46d3376682877b65ed1283c05c5fe Change-Id: I72773411f5e569e3562c722eaa0bdbad7b1e3fcd Change-Id: I46912057ac6fb59dc1f106eaf01cc73db7a97432 Change-Id: Iada75244c08618606f84769cf3c29186d6ad1b87
1 parent 8be2f15 commit d88f907

14 files changed

Lines changed: 1192 additions & 6 deletions

File tree

‎.github/workflows/pr-precommit.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ jobs:
5353
check-latest: true
5454
- uses: hashicorp/setup-terraform@v3
5555
with:
56-
terraform_version: "1.5.7"
56+
terraform_version: "1.12.2"
5757
terraform_wrapper: false
5858
- run: make install-dev-deps
5959
- uses: terraform-linters/setup-tflint@v4

‎examples/gke-consumption-options/dws-flex-start/gke-tpu-7x/README.md‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
## Create a TPU 7x Flex cluster
2+
These steps guide you through the cluster creation process for TPUs using DWS Flex Start.
3+
4+
1. Complete the common setup steps (1-4) in the [Create a cluster](https://github.com/GoogleCloudPlatform/cluster-toolkit/tree/main/examples/gke-consumption-options/dws-flex-start#create-a-cluster) section.
5+
6+
1. In the `examples/gke-consumption-options/dws-flex-start/gke-tpu-7x/gke-tpu-7x-deployment.yaml` file, fill in the following settings in the terraform_backend_defaults and vars sections to match the specific values for your deployment:
7+
8+
`bucket`: the name of the Cloud Storage bucket you created in the previous step.
9+
`deployment_name`: the name of the deployment.
10+
`project_id`: your Google Cloud project ID.
11+
`region`: the compute region for the cluster.
12+
`zone`: the compute zone for the node pool of TPU 7x machines.
13+
**`enable_flex_start`**: set to `true` to enable DWS Flex Start.
14+
**`autoscaling_min_node_count`**: set to `0` (required for Flex Start).
15+
**`autoscaling_max_node_count`**: set to the required node count for your topology (e.g., `2` for a `2x2x2` topology).
16+
`authorized_cidr`: The IP address range that you want to allow to connect with the cluster.
17+
To modify advanced settings, edit `examples/gke-consumption-options/dws-flex-start/gke-tpu-7x/gke-tpu-7x.yaml`.
18+
19+
1. Generate [Application Default Credentials (ADC)](https://cloud.google.com/docs/authentication/provide-credentials-adc#google-idp) to provide access to Terraform.
20+
21+
```sh
22+
gcloud auth application-default login
23+
```
24+
25+
1. Deploy the blueprint to provision the GKE infrastructure using TPU 7x machine types:
26+
27+
```sh
28+
cd ~/cluster-toolkit
29+
./gcluster deploy -d \
30+
examples/gke-consumption-options/dws-flex-start/gke-tpu-7x/gke-tpu-7x-deployment.yaml \
31+
examples/gke-consumption-options/dws-flex-start/gke-tpu-7x/gke-tpu-7x.yaml
32+
```
33+
34+
1. When prompted, select (A)pply to deploy the blueprint.
35+
36+
## Note
37+
38+
* DWS Flex Start does not work with static nodes. So, static_node_count cannot be set.
39+
* To use DWS Flex Start, `auto_repair` should be set to `false`.
40+
41+
## Running Workloads with DWS Flex Start
42+
43+
The Cluster Toolkit automatically generates a pre-configured, DWS-compliant job file located in your deployment folder (e.g., `gke-tpu-7x-flex/primary/my-job-xxxx.yaml`).
44+
45+
If you wish to create your own custom job, ensure it includes the following critical settings:
46+
47+
### 1. Node Selectors & Tolerations
48+
The job must target the TPU 7x nodes and tolerate the standard TPU taint:
49+
50+
```yaml
51+
nodeSelector:
52+
cloud.google.com/gke-tpu-accelerator: "tpu7x"
53+
cloud.google.com/gke-tpu-topology: "2x2x2"
54+
tolerations:
55+
- key: google.com/tpu
56+
operator: Equal
57+
value: "present"
58+
effect: NoSchedule
59+
```
60+
61+
### 2. Resource Requests
62+
63+
Each pod in the job should request the full amount of TPU chips available on the node (typically 4 for `tpu7x-standard-4t`):
64+
65+
```yaml
66+
resources:
67+
limits:
68+
google.com/tpu: 4
69+
requests:
70+
google.com/tpu: 4
71+
```
72+
73+
### 3. Parallelism
74+
75+
The `parallelism` and `completions` count in your Kubernetes Job (or ReplicatedJob replicas in a JobSet) must exactly match the number of nodes in your TPU slice (e.g., `2` for a `2x2x2` topology).
76+
77+
## Testing TPU Flex Start (Dynamic Scaling)
78+
79+
When using the TPU Flex Start model, the cluster begins with **0 nodes** in the TPU node pool. You can verify the dynamic scaling by following these steps:
80+
81+
1. **Monitor the cluster status:**
82+
Open two terminals. In the first, watch the pods:
83+
84+
```sh
85+
kubectl get pods -w
86+
```
87+
88+
In the second, watch the nodes:
89+
90+
```sh
91+
kubectl get nodes -w
92+
```
93+
94+
2. **Submit the TPU job:**
95+
Submit the generated job file:
96+
97+
```sh
98+
kubectl apply -f <deployment_folder>/primary/my-job-xxxx.yaml
99+
```
100+
101+
3. **Observe Scale-Up:**
102+
* **Initial State:** Pods will show as `Pending`.
103+
104+
```text
105+
NAME READY STATUS RESTARTS AGE
106+
my-job-2932-0-q2ksv 0/2 Pending 0 10s
107+
```
108+
109+
* **Autoscaling Triggered:** Check events to see the scale-up trigger: `kubectl get events`. You will see `TriggeredScaleUp`.
110+
111+
* **Nodes Joining:** After a few minutes, nodes will appear and transition to `Ready`.
112+
113+
```text
114+
NAME STATUS ROLES AGE VERSION
115+
gke-tpu-7f1325ce-8hwg Ready <none> 10s v1.34.1-gke.3971000
116+
```
117+
118+
* **Running:** Once nodes are ready, pods will transition to `Running`.
119+
120+
4. **Observe Scale-Down:**
121+
* **Completion:** Once the workload finishes, pods will move to `Completed`.
122+
123+
```text
124+
NAME READY STATUS RESTARTS AGE
125+
my-job-2932-0-q2ksv 0/2 Completed 0 5m
126+
```
127+
128+
* **Automatic Removal:** After a short idle period (typically 1-10 minutes), the Cluster Autoscaler will delete the nodes.
129+
130+
```text
131+
gke-tpu-7f1325ce-8hwg NotReady,SchedulingDisabled <none> 6m v1.34.1-gke.3971000
132+
```
133+
134+
* **Final State:** The node pool will return to **0 nodes**.

‎examples/gke-consumption-options/dws-flex-start/gke-tpu-7x/gke-tpu-7x-deployment.yaml‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
# Copyright 2026 "Google LLC"
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
terraform_backend_defaults:
17+
type: gcs
18+
configuration:
19+
# The GCS bucket used for storing terraform state
20+
bucket:
21+
22+
vars:
23+
# Your GCP Project ID
24+
project_id:
25+
26+
# This should be unique across all of your Cluster
27+
# Toolkit Deployments.
28+
deployment_name: gke-tpu-7x-flex
29+
30+
# The GCP Region used for this deployment.
31+
# example: us-central1
32+
region:
33+
34+
# The GCP Zone used for this deployment.
35+
# example: us-central1-c
36+
zone:
37+
38+
# The number of identical, independent TPU slices (node pools) to create.
39+
# For a standard single-slice deployment, this value will be 1.
40+
num_slices: 1
41+
42+
# Machine type
43+
machine_type: tpu7x-standard-4t
44+
45+
# The physical arrangement of the TPU chips in the slice.
46+
tpu_topology: 2x2x2
47+
48+
# --- FLEX START SETTINGS ---
49+
# Flex Start (Dynamic Provisioning)
50+
# For TPUs, max node count should be num of chips based on topology divided by
51+
# num chips for the machine type (e.g. 4 for tpu7x-standard-4t).
52+
# Reference: https://cloud.google.com/tpu/docs/v6e
53+
enable_flex_start: true
54+
autoscaling_min_node_count: 0
55+
autoscaling_max_node_count: 2
56+
57+
# Cidr block containing the IP of the machine calling terraform.
58+
# To allow all (IAM restrictions still enforced), use 0.0.0.0/0
59+
# To allow only your IP address, use <YOUR-IP-ADDRESS>/32
60+
authorized_cidr:

0 commit comments

Comments
 (0)