@@ -66,12 +66,14 @@ def create_namespace(self):
66
66
return RuntimeError (e )
67
67
68
68
69
- def create_new_resource_flavor (self , num_flavors ):
69
+ def create_new_resource_flavor (self , num_flavors , with_labels , with_tolerations ):
70
70
self .resource_flavors = []
71
71
for i in range (num_flavors ):
72
72
default = i < 1
73
73
resource_flavor = f"test-resource-flavor-{ random_choice ()} "
74
- create_resource_flavor (self , resource_flavor , default )
74
+ create_resource_flavor (
75
+ self , resource_flavor , default , with_labels , with_tolerations
76
+ )
75
77
self .resource_flavors .append (resource_flavor )
76
78
77
79
@@ -173,7 +175,9 @@ def create_cluster_queue(self, cluster_queue, flavor):
173
175
self .cluster_queue = cluster_queue
174
176
175
177
176
- def create_resource_flavor (self , flavor , default = True ):
178
+ def create_resource_flavor (
179
+ self , flavor , default = True , with_labels = False , with_tolerations = False
180
+ ):
177
181
worker_label , worker_value = os .getenv ("WORKER_LABEL" , "worker-1=true" ).split ("=" )
178
182
control_label , control_value = os .getenv (
179
183
"CONTROL_LABEL" , "ingress-ready=true"
@@ -182,23 +186,31 @@ def create_resource_flavor(self, flavor, default=True):
182
186
"TOLERATION_KEY" , "node-role.kubernetes.io/control-plane"
183
187
)
184
188
185
- node_labels = (
186
- {worker_label : worker_value } if default else {control_label : control_value }
187
- )
189
+ node_labels = {}
190
+ if with_labels :
191
+ node_labels = (
192
+ {worker_label : worker_value } if default else {control_label : control_value }
193
+ )
188
194
189
195
resource_flavor_json = {
190
196
"apiVersion" : "kueue.x-k8s.io/v1beta1" ,
191
197
"kind" : "ResourceFlavor" ,
192
198
"metadata" : {"name" : flavor },
193
199
"spec" : {
194
200
"nodeLabels" : node_labels ,
195
- "tolerations" : [
201
+ ** (
196
202
{
197
- "key" : toleration_key ,
198
- "operator" : "Exists" ,
199
- "effect" : "NoSchedule" ,
203
+ "tolerations" : [
204
+ {
205
+ "key" : toleration_key ,
206
+ "operator" : "Exists" ,
207
+ "effect" : "NoSchedule" ,
208
+ }
209
+ ]
200
210
}
201
- ],
211
+ if with_tolerations
212
+ else {}
213
+ ),
202
214
},
203
215
}
204
216
@@ -260,9 +272,11 @@ def create_local_queue(self, cluster_queue, local_queue, is_default=True):
260
272
self .local_queue = local_queue
261
273
262
274
263
- def create_kueue_resources (self , resource_ammount = 1 ):
275
+ def create_kueue_resources (
276
+ self , resource_ammount = 1 , with_labels = False , with_tolerations = False
277
+ ):
264
278
print ("creating Kueue resources ..." )
265
- create_new_resource_flavor (self , resource_ammount )
279
+ create_new_resource_flavor (self , resource_ammount , with_labels , with_tolerations )
266
280
create_new_cluster_queue (self , resource_ammount )
267
281
create_new_local_queue (self , resource_ammount )
268
282
0 commit comments