aks: support nodepool labels and taints

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/azure/azure_kubernetes_service.py

@@ -237,6 +237,9 @@ def _RunCreateClusterCmd(self, cmd: list[str]):
 
   def _CreateNodePool(self, nodepool_config: container.BaseNodePoolConfig):
     """Creates a node pool."""
+    labels = {'pkb_nodepool': nodepool_config.name}
+    if nodepool_config.node_labels:
+      labels.update(nodepool_config.node_labels)
     cmd = [
         azure.AZURE_PATH,
         'aks',
@@ -247,8 +250,10 @@ def _CreateNodePool(self, nodepool_config: container.BaseNodePoolConfig):
         '--name',
         _AzureNodePoolName(nodepool_config.name),
         '--labels',
-        f'pkb_nodepool={nodepool_config.name}',
+        ' '.join(f'{k}={v}' for k, v in labels.items()),
     ] + self._GetNodeFlags(nodepool_config)
+    if nodepool_config.node_taints:
+      cmd += ['--node-taints', ','.join(nodepool_config.node_taints)]
     _, stderr, retcode = vm_util.IssueCommand(
         cmd, timeout=600, raise_on_failure=False
     )

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/azure/azure_kubernetes_service_test.py

@@ -285,6 +285,32 @@ def testFullCreateAksAutomatic(self):
         mock_cmd.all_commands,
     )
 
+  def testCreateNodepoolLabels(self):
+    mock_cmd = self.MockIssueCommand({
+        'az aks nodepool add': [('', '', 0)],
+    })
+    nodepool_config = self.aks.default_nodepool
+    nodepool_config.node_labels = {'env': 'prod', 'team': 'ml'}
+    self.aks._CreateNodePool(nodepool_config)
+    call_args = str(mock_cmd.func_to_mock.mock_calls[0])
+    self.assertIn('env=prod', call_args)
+    self.assertIn('team=ml', call_args)
+
+  def testCreateNodepoolTaints(self):
+    mock_cmd = self.MockIssueCommand({
+        'az aks nodepool add': [('', '', 0)],
+    })
+    nodepool_config = self.aks.default_nodepool
+    nodepool_config.node_taints = [
+        'sandbox.gke.io/runtime=runsc:NoSchedule',
+        'dedicated:NoExecute',
+    ]
+    self.aks._CreateNodePool(nodepool_config)
+    call_args = str(mock_cmd.func_to_mock.mock_calls[0])
+    self.assertIn('--node-taints', call_args)
+    self.assertIn('sandbox.gke.io/runtime=runsc:NoSchedule', call_args)
+    self.assertIn('dedicated:NoExecute', call_args)
+
   def testGetNodePoolNames(self):
     self.MockIssueCommand(
         {