Skip to content

Commit 8443c0e

Browse files
Ishita-ChailIshita-Chail
committed
Unified Consumption Options in Deployment File for A3 Mega
1 parent 1c7ee91 commit 8443c0e

13 files changed

Lines changed: 756 additions & 27 deletions

File tree

‎examples/dws-queues/dws-queues.yaml.tftpl‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
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 obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
apiVersion: kueue.x-k8s.io/v1beta2
16+
kind: ResourceFlavor
17+
metadata:
18+
name: "default-flavor"
19+
---
20+
apiVersion: kueue.x-k8s.io/v1beta2
21+
kind: AdmissionCheck
22+
metadata:
23+
name: dws-prov
24+
spec:
25+
controllerName: kueue.x-k8s.io/provisioning-request
26+
parameters:
27+
apiGroup: kueue.x-k8s.io
28+
kind: ProvisioningRequestConfig
29+
name: dws-config
30+
---
31+
apiVersion: kueue.x-k8s.io/v1beta2
32+
kind: ProvisioningRequestConfig
33+
metadata:
34+
name: dws-config
35+
spec:
36+
provisioningClassName: queued-provisioning.gke.io
37+
managedResources:
38+
- nvidia.com/gpu
39+
---
40+
apiVersion: kueue.x-k8s.io/v1beta2
41+
kind: ClusterQueue
42+
metadata:
43+
name: "dws-cluster-queue"
44+
spec:
45+
namespaceSelector: {}
46+
resourceGroups:
47+
- coveredResources: ["nvidia.com/gpu"]
48+
flavors:
49+
- name: "default-flavor"
50+
resources:
51+
- name: "nvidia.com/gpu"
52+
nominalQuota: ${num_gpus}
53+
admissionChecks:
54+
- dws-prov
55+
---
56+
apiVersion: kueue.x-k8s.io/v1beta2
57+
kind: LocalQueue
58+
metadata:
59+
namespace: "default"
60+
name: "dws-local-queue"
61+
spec:
62+
clusterQueue: "dws-cluster-queue"
63+
---
64+
apiVersion: scheduling.k8s.io/v1
65+
kind: PriorityClass
66+
metadata:
67+
name: very-low
68+
value: 100
69+
globalDefault: false
70+
description: "Very Low"
71+
---
72+
apiVersion: scheduling.k8s.io/v1
73+
kind: PriorityClass
74+
metadata:
75+
name: low
76+
value: 250
77+
globalDefault: false
78+
description: "Low"
79+
---
80+
apiVersion: scheduling.k8s.io/v1
81+
kind: PriorityClass
82+
metadata:
83+
name: medium
84+
value: 500
85+
globalDefault: false
86+
description: "Medium"
87+
---
88+
apiVersion: scheduling.k8s.io/v1
89+
kind: PriorityClass
90+
metadata:
91+
name: high
92+
value: 750
93+
globalDefault: false
94+
description: "High"
95+
---
96+
apiVersion: scheduling.k8s.io/v1
97+
kind: PriorityClass
98+
metadata:
99+
name: very-high
100+
value: 1000
101+
globalDefault: false
102+
description: "Very High"

‎examples/dws-sample-workloads/sample-job-dws.yaml‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
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 obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
apiVersion: batch/v1
16+
kind: Job
17+
metadata:
18+
name: gpu-job-1
19+
spec:
20+
template:
21+
spec:
22+
nodeSelector:
23+
cloud.google.com/gke-flex-start: "true"
24+
containers:
25+
- name: gpu-container-1
26+
image: gcr.io/k8s-staging-perf-tests/sleep:latest
27+
args: ["10s"] # Sleep for 10 seconds
28+
resources:
29+
requests:
30+
nvidia.com/gpu: 1
31+
limits:
32+
nvidia.com/gpu: 1
33+
restartPolicy: OnFailure
34+
---
35+
apiVersion: batch/v1
36+
kind: Job
37+
metadata:
38+
name: gpu-job-2
39+
spec:
40+
template:
41+
spec:
42+
nodeSelector:
43+
cloud.google.com/gke-flex-start: "true"
44+
containers:
45+
- name: gpu-container-2
46+
image: gcr.io/k8s-staging-perf-tests/sleep:latest
47+
args: ["10s"] # Sleep for 10 seconds
48+
resources:
49+
requests:
50+
nvidia.com/gpu: 1
51+
limits:
52+
nvidia.com/gpu: 1
53+
restartPolicy: OnFailure

‎examples/gke-a3-megagpu/README.md‎

Lines changed: 103 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ The blueprint automatically configures the following components to enable optima
88
- **Multi-networking**: Configures 8 secondary interfaces (VPC networks) for dedicated GPU-to-GPU traffic.
99
- **NRI Device Injector**: Automatically injects required networking and GPU configurations into your ML containers.
1010
- **Kueue and JobSet**: Kubernetes-native tools for managing large-scale, multi-node training jobs with Topology Aware Scheduling (TAS).
11+
- **Consumption Options**: Supports On-Demand, GCE Reservations, Spot VMs, Dynamic Workload Scheduler (DWS) Flex Start, and DWS Flex Start with Queued Provisioning.
1112

1213
## Prerequisites
1314

@@ -27,10 +28,27 @@ Before deploying, fill out the `gke-a3-megagpu-deployment.yaml` file with your p
2728
| `deployment_name` | A unique name for this Cluster Toolkit deployment. |
2829
| `region` / `zone` | The GCP region and zone (e.g., `us-east5`, `us-east5-a`). |
2930
| `authorized_cidr` | Your public IP address in CIDR notation (e.g., `1.2.3.4/32`). |
30-
| `static_node_count` | Number of A3 Mega nodes to provision. |
31-
| `reservation` | (Optional) The name of a GCE reservation to use. |
3231
| `bucket` | Name of the GCS bucket to store Terraform state. |
3332

33+
### Consumption Options
34+
35+
You can select a consumption model by uncommenting the corresponding section in `gke-a3-megagpu-deployment.yaml`. For certain options, specific variable values must be configured:
36+
37+
- **DWS Flex Start**: Dynamic Workload Scheduler (DWS) Flex Start improves access to AI/ML resources and optimizes spend by scheduling required accelerators simultaneously.
38+
- Note: DWS Flex Start does not work with static nodes. So, `static_node_count` cannot be set.
39+
- Note: To use DWS Flex Start, `auto_repair` should be set to `false`.
40+
41+
- **DWS Flex Start + Queued Provisioning**: Enables queue-based provisioning on top of DWS Flex Start.
42+
- Note: Along with these flex start requirements, there are a few queue-provisioning specific requirements.
43+
- Note: Queued provisioning does not work with `static_node_count` and requires `autoscaling_total_min_nodes` be set to `0`.
44+
- Note: Set `kueue_configuration_path: $(ghpc_stage("../dws-queues/dws-queues.yaml.tftpl"))`.
45+
46+
- **Spot**: Requires setting `static_node_count`.
47+
48+
- **Specific Reservation**: Requires setting `static_node_count`. Refer to [Using GCE Reservations](https://github.com/GoogleCloudPlatform/cluster-toolkit/tree/main/modules/compute/gke-node-pool#using-gce-reservations) for details.
49+
50+
- **On-Demand (Default)**: Requires setting `static_node_count`.
51+
3452
## Deploy the Cluster
3553

3654
1. Switch to the toolkit directory:
@@ -154,6 +172,88 @@ After deployment, you can verify the GPU and networking performance using the in
154172
kubectl delete -f examples/gke-a3-megagpu/nccl-test-latest.yaml
155173
```
156174

175+
## DWS Flex Start
176+
177+
### Submit a job using DWS Flex Start
178+
179+
1. Submit the DWS Flex Start job:
180+
181+
```bash
182+
kubectl apply -f examples/dws-sample-workloads/sample-job-dws.yaml
183+
```
184+
185+
2. Consider using `kubectl get jobs` and `kubectl describe job <job-name>` to get information about the jobs.\
186+
You can also use `kubectl get pods` and `kubectl describe pod <pod-name>` to get pod information.
187+
188+
3. Clean up the job:
189+
190+
```bash
191+
kubectl delete -f examples/dws-sample-workloads/sample-job-dws.yaml
192+
```
193+
194+
*Note: DWS Flex Start workloads require `nodeSelector: cloud.google.com/gke-flex-start: "true"`.*
195+
196+
### Deploy the NCCL test JobSet
197+
198+
1. Deploy the NCCL test JobSet:
199+
200+
```bash
201+
kubectl create -f examples/gke-a3-megagpu/nccl-jobset-dws.yaml
202+
```
203+
204+
2. Monitor pods (`kubectl get pods`) and check results in the primary pod logs:
205+
206+
```bash
207+
kubectl logs <jobset-pod-name>
208+
```
209+
210+
3. Clean up test resources:
211+
212+
```bash
213+
kubectl delete -f examples/gke-a3-megagpu/nccl-jobset-dws.yaml
214+
```
215+
216+
## DWS Flex Start + Queued Provisioning
217+
218+
### Submit a job using Queued Provisioning
219+
220+
1. Submit the Queued Provisioning job:
221+
222+
```bash
223+
kubectl apply -f examples/gke-a3-megagpu/sample-job-dws-queue.yaml
224+
```
225+
226+
2. Consider using `kubectl get jobs` and `kubectl describe job <job-name>` to get information about the jobs.\
227+
You can also use `kubectl get pods` and `kubectl describe pod <pod-name>` to get pod information.
228+
229+
3. Clean up the job:
230+
231+
```bash
232+
kubectl delete -f examples/gke-a3-megagpu/sample-job-dws-queue.yaml
233+
```
234+
235+
*Note: Queued Provisioning workloads require the label `kueue.x-k8s.io/queue-name: dws-local-queue` and annotation `provreq.kueue.x-k8s.io/maxRunDurationSeconds`.*
236+
237+
### Deploy the NCCL test JobSet
238+
239+
1. Deploy the NCCL test JobSet:
240+
241+
```bash
242+
kubectl create -f examples/gke-a3-megagpu/nccl-jobset-dws-queue.yaml
243+
```
244+
245+
2. Monitor pods (`kubectl get pods`) and check results in the primary pod logs:
246+
247+
```bash
248+
kubectl logs <jobset-pod-name>
249+
```
250+
251+
3. Clean up test resources:
252+
253+
```bash
254+
kubectl delete -f examples/gke-a3-megagpu/nccl-jobset-dws-queue.yaml
255+
```
256+
157257
## Clean Up
158258

159259
To avoid incurring charges for the resources created, destroy the deployment:
@@ -168,6 +268,4 @@ Refer to [Deploy an A3 Mega GKE cluster for ML training](https://cloud.google.co
168268

169269
Refer to [Deploy and run NCCL test with Topology Aware Scheduling (TAS)](https://docs.cloud.google.com/ai-hypercomputer/docs/nccl/test-gke#a3-mega) for more instructions on running a NCCL test on the GKE-A3M cluster.
170270

171-
### Additional Consumption Options
172-
The Cluster Toolkit supports alternative consumption options such as Spot VMs or Dynamic Workload Scheduler (DWS) Flex-start.
173-
Refer to step 5 of [Create a cluster using Cluster Toolkit](https://docs.cloud.google.com/ai-hypercomputer/docs/create/gke-ai-hypercompute#use-cluster-toolkit) for general instructions on other consumption options. Similar configuration settings can be used for GKE-A3M cluster as well.
271+
Refer to [Consumption options](https://cloud.google.com/ai-hypercomputer/docs/create/gke-ai-hypercompute#use-cluster-toolkit) for more instructions on alternative consumption options.

‎examples/gke-a3-megagpu/gke-a3-megagpu-deployment.yaml‎

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -38,11 +38,30 @@ vars:
3838
# To allow only your IP address, use <YOUR-IP-ADDRESS>/32
3939
authorized_cidr:
4040

41-
# The number of nodes to be created
42-
static_node_count:
41+
## Choose ONE consumption model from the options below (default is On-Demand):
42+
43+
## --- Option 1: DWS Flex Start ---
44+
#enable_flex_start: true
45+
#static_node_count: 0
46+
47+
## --- Option 2: DWS Flex Start + Queued Provisioning ---
48+
## Change this variable if you want to have a custom kueue config file
49+
#kueue_configuration_path: $(ghpc_stage("../dws-queues/dws-queues.yaml.tftpl"))
50+
#enable_flex_start: true
51+
#enable_queued_provisioning: true
52+
#static_node_count: 0
4353

44-
# The name of the compute engine reservation in the form of
45-
# <reservation-name>
46-
# To target a BLOCK_NAME, the name of the extended reservation
47-
# can be inputted as <reservation-name>/reservationBlocks/<reservation-block-name>
48-
reservation:
54+
## --- Option 3: Spot ---
55+
#spot: true
56+
#static_node_count:
57+
58+
## --- Option 4: Specific Reservation ---
59+
## See for more information: https://github.com/GoogleCloudPlatform/cluster-toolkit/tree/main/modules/compute/gke-node-pool#using-gce-reservations
60+
#reservation_affinity:
61+
# consume_reservation_type: SPECIFIC_RESERVATION
62+
# specific_reservations:
63+
# - name:
64+
#static_node_count:
65+
66+
## --- Option 5: On-Demand (Default) ---
67+
static_node_count:

‎examples/gke-a3-megagpu/gke-a3-megagpu.yaml‎

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -44,11 +44,13 @@ vars:
4444
# The number of nodes to be created
4545
static_node_count:
4646

47-
# The name of the compute engine reservation in the form of
48-
# <reservation-name>
49-
# To target a BLOCK_NAME, the name of the extended reservation
50-
# can be inputted as <reservation-name>/reservationBlocks/<reservation-block-name>
51-
reservation:
47+
# Consumption model variables
48+
spot: false
49+
reservation_affinity: {consume_reservation_type: "NO_RESERVATION", specific_reservations: []}
50+
enable_flex_start: false
51+
auto_repair: true
52+
enable_queued_provisioning: false
53+
autoscaling_total_min_nodes: 0
5254

5355
accelerator_type: nvidia-h100-mega-80gb
5456
version_prefix: "1.35."
@@ -130,15 +132,17 @@ deployment_groups:
130132
install_gpu_direct_manifests: false
131133
run_workload_script: false
132134
machine_type: a3-megagpu-8g
133-
static_node_count: $(vars.static_node_count)
135+
static_node_count: "$((vars.enable_flex_start) ? null : vars.static_node_count)"
134136
zones: [$(vars.zone)]
135137
guest_accelerator:
136138
- type: $(vars.accelerator_type)
137139
count: 8
138-
reservation_affinity:
139-
consume_reservation_type: SPECIFIC_RESERVATION
140-
specific_reservations:
141-
- name: $(vars.reservation)
140+
spot: $(vars.spot)
141+
reservation_affinity: $(vars.reservation_affinity)
142+
enable_flex_start: $(vars.enable_flex_start)
143+
auto_repair: "$((vars.enable_flex_start) ? false : vars.auto_repair)"
144+
enable_queued_provisioning: $(vars.enable_queued_provisioning)
145+
autoscaling_total_min_nodes: $(vars.autoscaling_total_min_nodes)
142146
outputs: [instructions]
143147

144148
- id: workload_manager_install

0 commit comments

Comments
 (0)