@@ -66,12 +66,14 @@ def create_namespace(self):
6666 return RuntimeError (e )
6767
6868
69- def create_new_resource_flavor (self , num_flavors ):
69+ def create_new_resource_flavor (self , num_flavors , with_labels , with_tolerations ):
7070 self .resource_flavors = []
7171 for i in range (num_flavors ):
7272 default = i < 1
7373 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+ )
7577 self .resource_flavors .append (resource_flavor )
7678
7779
@@ -173,7 +175,9 @@ def create_cluster_queue(self, cluster_queue, flavor):
173175 self .cluster_queue = cluster_queue
174176
175177
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+ ):
177181 worker_label , worker_value = os .getenv ("WORKER_LABEL" , "worker-1=true" ).split ("=" )
178182 control_label , control_value = os .getenv (
179183 "CONTROL_LABEL" , "ingress-ready=true"
@@ -182,23 +186,31 @@ def create_resource_flavor(self, flavor, default=True):
182186 "TOLERATION_KEY" , "node-role.kubernetes.io/control-plane"
183187 )
184188
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+ )
188194
189195 resource_flavor_json = {
190196 "apiVersion" : "kueue.x-k8s.io/v1beta1" ,
191197 "kind" : "ResourceFlavor" ,
192198 "metadata" : {"name" : flavor },
193199 "spec" : {
194200 "nodeLabels" : node_labels ,
195- "tolerations" : [
201+ ** (
196202 {
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+ ]
200210 }
201- ],
211+ if with_tolerations
212+ else {}
213+ ),
202214 },
203215 }
204216
@@ -260,9 +272,11 @@ def create_local_queue(self, cluster_queue, local_queue, is_default=True):
260272 self .local_queue = local_queue
261273
262274
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+ ):
264278 print ("creating Kueue resources ..." )
265- create_new_resource_flavor (self , resource_ammount )
279+ create_new_resource_flavor (self , resource_ammount , with_labels , with_tolerations )
266280 create_new_cluster_queue (self , resource_ammount )
267281 create_new_local_queue (self , resource_ammount )
268282
0 commit comments