Skip to content

Commit 4a3787a

Browse files
FIoannidesFIoannides
authored
Add Kueue default config for Dynamic Slicing (#5693)
1 parent 863faa9 commit 4a3787a

7 files changed

Lines changed: 257 additions & 14 deletions

File tree

‎examples/gke-tpu-7x/gke-tpu-7x-advanced.yaml‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,9 @@ vars:
5252
# Accelerator type for Kueue configuration
5353
accelerator_type: tpu7x
5454

55+
# Enable Dynamic Slicing for TPUs
56+
enable_dynamic_slicing_for_tpus: false
57+
5558
# Kueue configuration
5659
kueue_configuration_path: $(ghpc_stage("./kueue-configuration.yaml.tftpl"))
5760

@@ -192,6 +195,7 @@ deployment_groups:
192195
start_time: "2025-12-01T00:00:00Z"
193196
exclusion_scope: NO_MINOR_OR_NODE_UPGRADES
194197
exclusion_end_time_behavior: UNTIL_END_OF_SUPPORT
198+
enable_slice_controller: $(vars.enable_dynamic_slicing_for_tpus)
195199
outputs: [instructions]
196200

197201
- id: workload_policy
@@ -219,7 +223,7 @@ deployment_groups:
219223
consume_reservation_type: SPECIFIC_RESERVATION
220224
specific_reservations:
221225
- name: $(vars.reservation)
222-
outputs: [instructions]
226+
outputs: [instructions, machine_type, accelerator_topology_mode]
223227

224228
# # --- MANAGED LUSTRE ADDITIONS ---
225229
# # Private Service Access (PSA) requires the compute.networkAdmin role which is
@@ -272,6 +276,9 @@ deployment_groups:
272276
settings:
273277
kueue:
274278
install: true
279+
enable_dynamic_slicing_for_tpus: $(vars.enable_dynamic_slicing_for_tpus)
280+
accelerator_topology_mode: $(gke-tpu-7x-pool.accelerator_topology_mode)
281+
machine_type: $(gke-tpu-7x-pool.machine_type)
275282
config_path: $(vars.kueue_configuration_path)
276283
config_template_vars:
277284
tpu_quota: $(vars.num_slices * gke-tpu-7x-pool.node_count_static * gke-tpu-7x-pool.tpu_chips_per_node)

‎examples/gke-tpu-7x/gke-tpu-7x.yaml‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,9 @@ vars:
5252
# Accelerator type for Kueue configuration
5353
accelerator_type: tpu7x
5454

55+
# Enable Dynamic Slicing for TPUs
56+
enable_dynamic_slicing_for_tpus: false
57+
5558
# Kueue configuration
5659
kueue_configuration_path: $(ghpc_stage("./kueue-configuration.yaml.tftpl"))
5760

@@ -135,6 +138,7 @@ deployment_groups:
135138
start_time: "2025-12-01T00:00:00Z"
136139
exclusion_scope: NO_MINOR_OR_NODE_UPGRADES
137140
exclusion_end_time_behavior: UNTIL_END_OF_SUPPORT
141+
enable_slice_controller: $(vars.enable_dynamic_slicing_for_tpus)
138142
outputs: [instructions]
139143

140144
- id: workload_policy
@@ -161,14 +165,17 @@ deployment_groups:
161165
consume_reservation_type: SPECIFIC_RESERVATION
162166
specific_reservations:
163167
- name: $(vars.reservation)
164-
outputs: [instructions]
168+
outputs: [instructions, machine_type, accelerator_topology_mode]
165169

166170
- id: workload-manager-install
167171
source: modules/management/kubectl-apply
168172
use: [gke-tpu-7x-cluster]
169173
settings:
170174
kueue:
171175
install: true
176+
enable_dynamic_slicing_for_tpus: $(vars.enable_dynamic_slicing_for_tpus)
177+
accelerator_topology_mode: $(gke-tpu-7x-pool.accelerator_topology_mode)
178+
machine_type: $(gke-tpu-7x-pool.machine_type)
172179
config_path: $(vars.kueue_configuration_path)
173180
config_template_vars:
174181
tpu_quota: $(vars.num_slices * gke-tpu-7x-pool.node_count_static * gke-tpu-7x-pool.tpu_chips_per_node)

‎modules/management/kubectl-apply/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -275,7 +275,7 @@ limitations under the License.
275275
| <a name="input_gke_cluster_exists"></a> [gke\_cluster\_exists](#input\_gke\_cluster\_exists) | A static flag that signals to downstream modules that a cluster has been created. | `bool` | `false` | no |
276276
| <a name="input_gpu_operator"></a> [gpu\_operator](#input\_gpu\_operator) | Install [GPU Operator](https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/getting-started.html) which uses the [Kubernetes operator](https://kubernetes.io/docs/concepts/extend-kubernetes/operator/) to automate the management of all NVIDIA software components needed to provision GPU. | <pre>object({<br/> install = optional(bool, false)<br/> version = optional(string, "v25.3.0")<br/> })</pre> | `{}` | no |
277277
| <a name="input_jobset"></a> [jobset](#input\_jobset) | Install [Jobset](https://github.com/kubernetes-sigs/jobset) which manages a group of K8s [jobs](https://kubernetes.io/docs/concepts/workloads/controllers/job/) as a unit. | <pre>object({<br/> install = optional(bool, false)<br/> version = optional(string, "0.10.1")<br/> controller_cpu = optional(string, null)<br/> controller_memory = optional(string, null)<br/> })</pre> | `{}` | no |
278-
| <a name="input_kueue"></a> [kueue](#input\_kueue) | Install and configure [Kueue](https://kueue.sigs.k8s.io/docs/overview/) workload scheduler. A configuration yaml/template file can be provided with config\_path to be applied right after kueue installation. If a template file provided, its variables can be set to config\_template\_vars. | <pre>object({<br/> # ATTENTION: If you update the KUEUE's default version below, please also update the corresponding<br/> # defaultKueueVersion constant in pkg/orchestrator/gke/infra_manager.go. (note the 'v' prefix there)<br/> version = optional(string, "0.17.1")<br/> install = optional(bool, false)<br/> config_path = optional(string, null)<br/> config_template_vars = optional(map(any), null)<br/> enable_pathways_for_tpus = optional(bool, false)<br/> controller_cpu = optional(string, null)<br/> controller_memory = optional(string, null)<br/> controller_replicas = optional(number, null)<br/> })</pre> | `{}` | no |
278+
| <a name="input_kueue"></a> [kueue](#input\_kueue) | Install and configure [Kueue](https://kueue.sigs.k8s.io/docs/overview/) workload scheduler. A configuration yaml/template file can be provided with config\_path to be applied right after kueue installation. If a template file provided, its variables can be set to config\_template\_vars. | <pre>object({<br/> # ATTENTION: If you update the KUEUE's default version below, please also update the corresponding<br/> # defaultKueueVersion constant in pkg/orchestrator/gke/infra_manager.go. (note the 'v' prefix there)<br/> version = optional(string, "0.17.1")<br/> install = optional(bool, false)<br/> config_path = optional(string, null)<br/> config_template_vars = optional(map(any), null)<br/> enable_pathways_for_tpus = optional(bool, false)<br/> enable_dynamic_slicing_for_tpus = optional(bool, false)<br/> accelerator_topology_mode = optional(string, null)<br/> machine_type = optional(string, null)<br/> controller_cpu = optional(string, null)<br/> controller_memory = optional(string, null)<br/> controller_replicas = optional(number, null)<br/> })</pre> | `{}` | no |
279279
| <a name="input_module_id"></a> [module\_id](#input\_module\_id) | The ID of the module as defined in the blueprint. Injected by ghpc. | `string` | `"kubectl-apply"` | no |
280280
| <a name="input_nvidia_dra_driver"></a> [nvidia\_dra\_driver](#input\_nvidia\_dra\_driver) | Installs [Nvidia DRA driver](https://github.com/NVIDIA/k8s-dra-driver-gpu) which supports Dynamic Resource Allocation for NVIDIA GPUs in Kubernetes | <pre>object({<br/> install = optional(bool, false)<br/> version = optional(string, "v25.3.0")<br/> accelerator_type = optional(string, "nvidia-gb200")<br/> })</pre> | `{}` | no |
281281
| <a name="input_project_id"></a> [project\_id](#input\_project\_id) | The project ID that hosts the gke cluster. | `string` | n/a | yes |

‎modules/management/kubectl-apply/kueue/kueue-configuration-dynamic-slicing-pathways.yaml.tftpl‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
apiVersion: kueue.x-k8s.io/v1beta1
2+
kind: Topology
3+
metadata:
4+
name: dynamic-slicing-topology
5+
spec:
6+
levels:
7+
- nodeLabel: cloud.google.com/gce-topology-block
8+
- nodeLabel: cloud.google.com/gke-tpu-partition-4x4x4-id
9+
- nodeLabel: kubernetes.io/hostname
10+
---
11+
apiVersion: kueue.x-k8s.io/v1beta1
12+
kind: AdmissionCheck
13+
metadata:
14+
name: dynamic-slicing-ac
15+
spec:
16+
controllerName: accelerator.gke.io/slice
17+
---
18+
apiVersion: kueue.x-k8s.io/v1beta2
19+
kind: ResourceFlavor
20+
metadata:
21+
name: tpu-flavor
22+
spec:
23+
nodeLabels:
24+
cloud.google.com/gke-tpu-accelerator: ${accelerator_type}
25+
topologyName: dynamic-slicing-topology
26+
tolerations:
27+
- effect: NoSchedule
28+
key: google.com/tpu
29+
value: "present"
30+
---
31+
apiVersion: kueue.x-k8s.io/v1beta2
32+
kind: ResourceFlavor
33+
metadata:
34+
name: pathways-flavor
35+
spec:
36+
nodeLabels:
37+
cloud.google.com/gke-nodepool: cpu-np
38+
---
39+
apiVersion: kueue.x-k8s.io/v1beta2
40+
kind: ClusterQueue
41+
metadata:
42+
name: cluster-queue
43+
spec:
44+
preemption:
45+
reclaimWithinCohort: Never
46+
withinClusterQueue: LowerPriority
47+
namespaceSelector: {}
48+
admissionChecks:
49+
- dynamic-slicing-ac
50+
resourceGroups:
51+
- coveredResources: ["google.com/tpu"]
52+
flavors:
53+
- name: "tpu-flavor"
54+
resources:
55+
- name: "google.com/tpu"
56+
nominalQuota: ${tpu_quota}
57+
- coveredResources: ["cpu", "memory"]
58+
flavors:
59+
- name: pathways-flavor
60+
resources:
61+
- name: cpu
62+
nominalQuota: ${pathways_cpu_quota}
63+
- name: memory
64+
nominalQuota: ${pathways_memory_quota}
65+
---
66+
apiVersion: kueue.x-k8s.io/v1beta2
67+
kind: LocalQueue
68+
metadata:
69+
namespace: default
70+
name: user-queue
71+
spec:
72+
clusterQueue: cluster-queue
73+
---
74+
apiVersion: scheduling.k8s.io/v1
75+
kind: PriorityClass
76+
metadata:
77+
name: very-low
78+
value: 100
79+
globalDefault: false
80+
description: "Very Low"
81+
---
82+
apiVersion: scheduling.k8s.io/v1
83+
kind: PriorityClass
84+
metadata:
85+
name: low
86+
value: 250
87+
globalDefault: false
88+
description: "Low"
89+
---
90+
apiVersion: scheduling.k8s.io/v1
91+
kind: PriorityClass
92+
metadata:
93+
name: medium
94+
value: 500
95+
globalDefault: false
96+
description: "Medium"
97+
---
98+
apiVersion: scheduling.k8s.io/v1
99+
kind: PriorityClass
100+
metadata:
101+
name: high
102+
value: 750
103+
globalDefault: false
104+
description: "High"
105+
---
106+
apiVersion: scheduling.k8s.io/v1
107+
kind: PriorityClass
108+
metadata:
109+
name: very-high
110+
value: 1000
111+
globalDefault: false
112+
description: "Very High"

‎modules/management/kubectl-apply/kueue/kueue-configuration-dynamic-slicing.yaml.tftpl‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
apiVersion: kueue.x-k8s.io/v1beta1
2+
kind: Topology
3+
metadata:
4+
name: dynamic-slicing-topology
5+
spec:
6+
levels:
7+
- nodeLabel: cloud.google.com/gce-topology-block
8+
- nodeLabel: cloud.google.com/gke-tpu-partition-4x4x4-id
9+
- nodeLabel: kubernetes.io/hostname
10+
---
11+
apiVersion: kueue.x-k8s.io/v1beta1
12+
kind: AdmissionCheck
13+
metadata:
14+
name: dynamic-slicing-ac
15+
spec:
16+
controllerName: accelerator.gke.io/slice
17+
---
18+
apiVersion: kueue.x-k8s.io/v1beta2
19+
kind: ResourceFlavor
20+
metadata:
21+
name: tpu-flavor
22+
spec:
23+
nodeLabels:
24+
cloud.google.com/gke-tpu-accelerator: ${accelerator_type}
25+
topologyName: dynamic-slicing-topology
26+
tolerations:
27+
- effect: NoSchedule
28+
key: google.com/tpu
29+
value: "present"
30+
---
31+
apiVersion: kueue.x-k8s.io/v1beta2
32+
kind: ClusterQueue
33+
metadata:
34+
name: cluster-queue
35+
spec:
36+
preemption:
37+
reclaimWithinCohort: Never
38+
withinClusterQueue: LowerPriority
39+
namespaceSelector: {}
40+
admissionChecks:
41+
- dynamic-slicing-ac
42+
resourceGroups:
43+
- coveredResources: ["google.com/tpu"]
44+
flavors:
45+
- name: "tpu-flavor"
46+
resources:
47+
- name: "google.com/tpu"
48+
nominalQuota: ${tpu_quota}
49+
---
50+
apiVersion: kueue.x-k8s.io/v1beta2
51+
kind: LocalQueue
52+
metadata:
53+
namespace: default
54+
name: user-queue
55+
spec:
56+
clusterQueue: cluster-queue
57+
---
58+
apiVersion: scheduling.k8s.io/v1
59+
kind: PriorityClass
60+
metadata:
61+
name: very-low
62+
value: 100
63+
globalDefault: false
64+
description: "Very Low"
65+
---
66+
apiVersion: scheduling.k8s.io/v1
67+
kind: PriorityClass
68+
metadata:
69+
name: low
70+
value: 250
71+
globalDefault: false
72+
description: "Low"
73+
---
74+
apiVersion: scheduling.k8s.io/v1
75+
kind: PriorityClass
76+
metadata:
77+
name: medium
78+
value: 500
79+
globalDefault: false
80+
description: "Medium"
81+
---
82+
apiVersion: scheduling.k8s.io/v1
83+
kind: PriorityClass
84+
metadata:
85+
name: high
86+
value: 750
87+
globalDefault: false
88+
description: "High"
89+
---
90+
apiVersion: scheduling.k8s.io/v1
91+
kind: PriorityClass
92+
metadata:
93+
name: very-high
94+
value: 1000
95+
globalDefault: false
96+
description: "Very High"

‎modules/management/kubectl-apply/main.tf‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,14 @@ locals {
2020
cluster_location = local.cluster_id_parts[3]
2121
project_id = var.project_id != null ? var.project_id : local.cluster_id_parts[1]
2222
enable_pathways = var.enable_pathways_for_tpus || var.kueue.enable_pathways_for_tpus
23+
enable_slicing = var.kueue.enable_dynamic_slicing_for_tpus
24+
25+
kueue_default_config_template = lookup({
26+
"true-true" = "${path.module}/kueue/kueue-configuration-dynamic-slicing-pathways.yaml.tftpl",
27+
"false-true" = "${path.module}/kueue/kueue-configuration-dynamic-slicing.yaml.tftpl",
28+
"true-false" = "${path.module}/kueue/kueue-configuration-pathways.yaml.tftpl",
29+
}, "${local.enable_pathways}-${local.enable_slicing}", "")
2330

24-
kueue_default_config_template = local.enable_pathways ? "${path.module}/kueue/kueue-configuration-pathways.yaml.tftpl" : ""
2531

2632
kueue_config_template_vars = merge(
2733
{
@@ -116,7 +122,7 @@ locals {
116122
# Step A: Use the fetched body if it's a URL
117123
contains(keys(local.url_manifests), tostring(index)) ? data.http.manifest_from_url[tostring(index)].body :
118124
119-
# Step B: Process directory files
125+
# Step B: Process directory files
120126
contains(keys(local.directory_manifests), index) ? (
121127
join("\n---\n", [
122128
# Use union() to combine the results of fileset (which are sets)
@@ -227,7 +233,7 @@ module "install_kueue" {
227233
depends_on = [var.gke_cluster_exists]
228234
}
229235

230-
# This sleep ensures that subsequent configuration of Kueue custom resources
236+
# This sleep ensures that subsequent configuration of Kueue custom resources
231237
# do not fail due to the webhook not being available.
232238
resource "time_sleep" "wait_for_webhook" {
233239
count = local.install_kueue ? 1 : 0

‎modules/management/kubectl-apply/variables.tf‎

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,18 @@ resource "terraform_data" "kueue_validations" {
5555
condition = !var.kueue.install || !(var.enable_pathways_for_tpus || try(var.kueue.enable_pathways_for_tpus, false)) || try(var.kueue.config_path, "") != "" || contains(keys(coalesce(var.kueue.config_template_vars, {})), "accelerator_type")
5656
error_message = "accelerator_type must be set in kueue.config_template_vars when using the default pathways configuration."
5757
}
58+
precondition {
59+
condition = !var.kueue.enable_dynamic_slicing_for_tpus || var.kueue.accelerator_topology_mode == "PROVISION_ONLY"
60+
error_message = "When enable_dynamic_slicing_for_tpus is true, accelerator_topology_mode must be 'PROVISION_ONLY'."
61+
}
62+
precondition {
63+
condition = !var.kueue.enable_dynamic_slicing_for_tpus || (var.kueue.machine_type != null && length(regexall("^(tpu|ct)", var.kueue.machine_type)) > 0)
64+
error_message = "When enable_dynamic_slicing_for_tpus is true, machine_type must be a TPU machine type."
65+
}
66+
precondition {
67+
condition = !var.kueue.enable_dynamic_slicing_for_tpus || coalesce(var.kueue.config_path, "") != "" || (var.kueue.config_template_vars != null && contains(keys(var.kueue.config_template_vars), "accelerator_type"))
68+
error_message = "accelerator_type must be set in kueue.config_template_vars when using the default dynamic slicing configuration."
69+
}
5870
}
5971
}
6072

@@ -128,14 +140,17 @@ variable "kueue" {
128140
type = object({
129141
# ATTENTION: If you update the KUEUE's default version below, please also update the corresponding
130142
# defaultKueueVersion constant in pkg/orchestrator/gke/infra_manager.go. (note the 'v' prefix there)
131-
version = optional(string, "0.17.1")
132-
install = optional(bool, false)
133-
config_path = optional(string, null)
134-
config_template_vars = optional(map(any), null)
135-
enable_pathways_for_tpus = optional(bool, false)
136-
controller_cpu = optional(string, null)
137-
controller_memory = optional(string, null)
138-
controller_replicas = optional(number, null)
143+
version = optional(string, "0.17.1")
144+
install = optional(bool, false)
145+
config_path = optional(string, null)
146+
config_template_vars = optional(map(any), null)
147+
enable_pathways_for_tpus = optional(bool, false)
148+
enable_dynamic_slicing_for_tpus = optional(bool, false)
149+
accelerator_topology_mode = optional(string, null)
150+
machine_type = optional(string, null)
151+
controller_cpu = optional(string, null)
152+
controller_memory = optional(string, null)
153+
controller_replicas = optional(number, null)
139154
})
140155
default = {}
141156
}

0 commit comments

Comments
 (0)