eks: support nodepool labels and taints

Add _ParseTaint helper and update _RenderNodeGroupJson in BaseEksCluster
to pass node_labels and node_taints through to eksctl's node group JSON.
Taints are parsed from 'key=value:Effect' or 'key:Effect' strings into
the dict format eksctl expects.

Author: geojaz

perfkitbenchmarker/configs/container_spec.py

@@ -243,6 +243,8 @@ def __init__(
     self.vm_spec: virtual_machine_spec.BaseVmSpec
     self.machine_families: list[str] | None
     self.sandbox_config: SandboxSpec | None
+    self.node_labels: dict[str, str] | None
+    self.node_taints: list[str] | None
 
   @classmethod
   def _GetOptionDecoderConstructions(cls):
@@ -273,6 +275,14 @@ def _GetOptionDecoderConstructions(cls):
         ),
         'vm_spec': (spec.PerCloudConfigDecoder, {}),
         'sandbox_config': (_SandboxDecoder, {'default': None}),
+        'node_labels': (
+            option_decoders.TypeVerifier,
+            {'valid_types': (dict,), 'default': None, 'none_ok': True},
+        ),
+        'node_taints': (
+            option_decoders.TypeVerifier,
+            {'valid_types': (list,), 'default': None, 'none_ok': True},
+        ),
     })
     return result
 

perfkitbenchmarker/providers/aws/elastic_kubernetes_service.py

@@ -101,6 +101,22 @@ def ApplyInferenceS3PvAndPvc() -> None:
   logging.info('Successfully applied S3 PVC.')
 
 
+def _ParseTaint(taint_str: str) -> dict[str, str]:
+  """Parses a taint string into eksctl's dict format.
+
+  Args:
+    taint_str: Taint in 'key=value:Effect' or 'key:Effect' format.
+
+  Returns:
+    Dict with 'key', 'value' (optional), and 'effect' keys for eksctl.
+  """
+  key_value, effect = taint_str.rsplit(':', 1)
+  if '=' in key_value:
+    key, value = key_value.split('=', 1)
+    return {'key': key, 'value': value, 'effect': effect}
+  return {'key': key_value, 'effect': effect}
+
+
 class BaseEksCluster(kubernetes_cluster.KubernetesCluster):
   """Shared base class for Elastic Kubernetes Service cluster auto mode & not."""
 
@@ -195,16 +211,19 @@ def _RenderNodeGroupJson(
       self, nodepool: container.BaseNodePoolConfig
   ) -> dict[str, Any]:
     """Constructs the node group json dictionary."""
+    labels = {'pkb_nodepool': nodepool.name}
+    if nodepool.node_labels:
+      labels.update(nodepool.node_labels)
     group_json = {
         'name': nodepool.name,
         'instanceType': nodepool.machine_type,
         'desiredCapacity': nodepool.num_nodes,
         'amiFamily': 'AmazonLinux2023',
         'tags': util.MakeDefaultTags(),
-        'labels': {
-            'pkb_nodepool': nodepool.name,
-        },
+        'labels': labels,
     }
+    if nodepool.node_taints:
+      group_json['taints'] = [_ParseTaint(t) for t in nodepool.node_taints]
     if nodepool.min_nodes != nodepool.max_nodes:
       group_json['minSize'] = nodepool.min_nodes
       group_json['maxSize'] = nodepool.max_nodes

perfkitbenchmarker/resources/container_service/container.py

@@ -184,6 +184,8 @@ def __init__(
     self.disk_size: int = vm_spec.boot_disk_size
     self.gpu_type: str | None = vm_spec.gpu_type
     self.gpu_count: int | None = vm_spec.gpu_count
+    self.node_labels: dict[str, str] | None = None
+    self.node_taints: list[str] | None = None
     # Defined by GceVirtualMachineConfig. Used by google_kubernetes_engine
     # pylint: disable=g-missing-from-attributes
     self.sandbox_config: container_spec_lib.SandboxSpec | None = None

perfkitbenchmarker/resources/container_service/container_cluster.py

@@ -116,6 +116,8 @@ def _InitializeNodePool(
         nodepool_spec.machine_families,
     )
     nodepool_config.sandbox_config = nodepool_spec.sandbox_config
+    nodepool_config.node_labels = nodepool_spec.node_labels
+    nodepool_config.node_taints = nodepool_spec.node_taints
     nodepool_config.zone = zone
     nodepool_config.num_nodes = nodepool_spec.vm_count
     if nodepool_spec.min_vm_count is None:

tests/providers/aws/elastic_kubernetes_service_test.py

@@ -237,6 +237,47 @@ def testEksClusterNodepoolsAutoscaling(self):
     self.assertEqual(node_groups[1]['maxSize'], 10)
     self.assertEqual(node_groups[1]['desiredCapacity'], 3)
 
+  def testEksClusterNodepoolLabels(self):
+    cluster = elastic_kubernetes_service.EksCluster(EKS_SPEC)
+    nodepool = cluster.default_nodepool
+    nodepool.node_labels = {'env': 'prod', 'team': 'ml'}
+    actual = cluster._RenderNodeGroupJson(nodepool)
+    self.assertEqual(
+        actual['labels'],
+        {'pkb_nodepool': nodepool.name, 'env': 'prod', 'team': 'ml'},
+    )
+
+  def testEksClusterNodepoolTaints(self):
+    cluster = elastic_kubernetes_service.EksCluster(EKS_SPEC)
+    nodepool = cluster.default_nodepool
+    nodepool.node_taints = [
+        'sandbox.gke.io/runtime=runsc:NoSchedule',
+        'dedicated:NoExecute',
+    ]
+    actual = cluster._RenderNodeGroupJson(nodepool)
+    self.assertEqual(
+        actual['taints'],
+        [
+            {
+                'key': 'sandbox.gke.io/runtime',
+                'value': 'runsc',
+                'effect': 'NoSchedule',
+            },
+            {'key': 'dedicated', 'effect': 'NoExecute'},
+        ],
+    )
+
+
+  def testParseTaint(self):
+    self.assertEqual(
+        elastic_kubernetes_service._ParseTaint('key=value:NoSchedule'),
+        {'key': 'key', 'value': 'value', 'effect': 'NoSchedule'},
+    )
+    self.assertEqual(
+        elastic_kubernetes_service._ParseTaint('dedicated:NoExecute'),
+        {'key': 'dedicated', 'effect': 'NoExecute'},
+    )
+
   def testGetNodePoolNames(self):
     # Mock the output of the aws cli command
     cluster = elastic_kubernetes_service.EksCluster(EKS_SPEC)