@@ -285,6 +285,47 @@ def is_autoprovisioning_enabled(
285
285
return False , 1
286
286
287
287
288
+ def get_capacity_type_str_with_cluster_default (args ) -> tuple [str , int ]:
289
+ """Determine the capacity type based on user arguments or cluster default.
290
+
291
+ Args:
292
+ args: user provided arguments for running the command.
293
+
294
+ Returns:
295
+ Tuple with string with the system characteristics and
296
+ int of 0 if successful and 1 otherwise.
297
+ """
298
+ # If the user doesn't specify args, then use the cluster settings.
299
+ capacity_type , return_code = get_capacity_type (args )
300
+ if return_code != 0 :
301
+ xpk_print ('Unable to get capacity type.' )
302
+ return CapacityType .UNKNOWN .name , return_code
303
+
304
+ if capacity_type != CapacityType .UNKNOWN :
305
+ return capacity_type .name , 0
306
+
307
+ # Use default settings from cluster creation.
308
+ metadata_configmap_name = f'{ args .cluster } -{ CLUSTER_METADATA_CONFIGMAP } '
309
+ cluster_config_map = get_cluster_configmap (args , metadata_configmap_name )
310
+
311
+ # Error out if the metadata config map doesn't exist, and is attempting to use
312
+ # autoprovisioning.
313
+ if cluster_config_map is None :
314
+ xpk_print (
315
+ 'Unable to find config map. Please specify a capacity type'
316
+ ' --on-demand, --spot, --reservation=$RESERVATION_ID) to continue'
317
+ ' to use autoprovisioning (--enable-autoprovisioning).'
318
+ )
319
+ return CapacityType .UNKNOWN .name , 1
320
+
321
+ return_code , capacity_type_str = get_value_from_map (
322
+ CAPACITY_TYPE_CONFIG_KEY , cluster_config_map
323
+ )
324
+ if return_code != 0 :
325
+ return CapacityType .UNKNOWN .name , return_code
326
+
327
+ return capacity_type_str , 0
328
+
288
329
def get_autoprovisioning_node_selector_args (args ) -> tuple [str , int ]:
289
330
"""Determine the capacity type when autoprovisioning is enabled.
290
331
@@ -297,44 +338,20 @@ def get_autoprovisioning_node_selector_args(args) -> tuple[str, int]:
297
338
"""
298
339
return_code = 0
299
340
node_selector_args = ''
300
- # If the user doesn't specify args, then use the cluster settings.
301
- capacity_type , return_code = get_capacity_type (args )
302
- capacity_type_str = capacity_type .name
341
+ capacity_type_str , return_code = get_capacity_type_str_with_cluster_default (args )
303
342
if return_code != 0 :
304
- xpk_print ('Unable to get capacity type.' )
305
343
return node_selector_args , return_code
306
344
307
- if capacity_type_str == CapacityType .UNKNOWN .name :
308
- # Use default settings from cluster creation.
309
- metadata_configmap_name = f'{ args .cluster } -{ CLUSTER_METADATA_CONFIGMAP } '
310
- cluster_config_map = get_cluster_configmap (args , metadata_configmap_name )
311
-
312
- # Error out if the metadata config map doesn't exist, and is attempting to use
313
- # autoprovisioning.
314
- if cluster_config_map is None :
315
- xpk_print (
316
- 'Unable to find config map. Please specify a capacity type'
317
- ' --on-demand, --spot, --reservation=$RESERVATION_ID) to continue'
318
- ' to use autoprovisioning (--enable-autoprovisioning).'
319
- )
320
- return node_selector_args , 1
321
-
322
- return_code , capacity_type_str = get_value_from_map (
323
- CAPACITY_TYPE_CONFIG_KEY , cluster_config_map
345
+ if capacity_type_str == CapacityType .RESERVATION .name :
346
+ return_code , args .reservation = get_value_from_map (
347
+ RESERVATION_CONFIG_KEY , cluster_config_map
324
348
)
325
349
if return_code != 0 :
326
350
return node_selector_args , return_code
327
-
328
- if capacity_type_str == CapacityType .RESERVATION .name :
329
- return_code , args .reservation = get_value_from_map (
330
- RESERVATION_CONFIG_KEY , cluster_config_map
331
- )
332
- if return_code != 0 :
333
- return node_selector_args , return_code
334
- return_code = verify_reservation_exists (args )
335
- if return_code > 0 :
336
- xpk_print ('Unable to verify reservation name saved in config map.' )
337
- return node_selector_args , return_code
351
+ return_code = verify_reservation_exists (args )
352
+ if return_code > 0 :
353
+ xpk_print ('Unable to verify reservation name saved in config map.' )
354
+ return node_selector_args , return_code
338
355
339
356
# Check if reservation id is valid. Shared function with cluster creation.
340
357
node_selector_args , return_code = (
@@ -345,3 +362,31 @@ def get_autoprovisioning_node_selector_args(args) -> tuple[str, int]:
345
362
return node_selector_args , return_code
346
363
347
364
return node_selector_args , return_code
365
+
366
+
367
+ def get_autoprovisioning_tolerations (args ) -> tuple [str , int ]:
368
+ """Determine the pod tolerations when autoprovisioning is enabled.
369
+
370
+ Args:
371
+ args: user provided arguments for running the command.
372
+
373
+ Returns:
374
+ Tuple with string of autoprovisioning tolerations and
375
+ int of 0 if successful and 1 otherwise.
376
+ """
377
+ capacity_type_str , return_code = get_capacity_type_str_with_cluster_default (args )
378
+ if return_code != 0 :
379
+ return '' , return_code
380
+
381
+ if capacity_type_str == CapacityType .SPOT .name :
382
+ # https://cloud.google.com/kubernetes-engine/docs/concepts/node-auto-provisioning#support_for_spot_vms
383
+ #
384
+ # > Creating node pools based on Spot VMs is only considered if
385
+ # > unschedulable pods with a toleration for the
386
+ # > cloud.google.com/gke-spot="true":NoSchedule taint exist
387
+ return '''- key: "cloud.google.com/gke-spot"
388
+ operator: "Equal"
389
+ value: "true"
390
+ effect: "NoSchedule"''' , 0
391
+
392
+ return '' , 0
0 commit comments