diff --git a/perfkitbenchmarker/configs/container_spec.py b/perfkitbenchmarker/configs/container_spec.py index 1f808ad066..5071eba87d 100644 --- a/perfkitbenchmarker/configs/container_spec.py +++ b/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 diff --git a/perfkitbenchmarker/providers/azure/azure_kubernetes_service.py b/perfkitbenchmarker/providers/azure/azure_kubernetes_service.py index 8d3864c920..a7bcce6282 100644 --- a/perfkitbenchmarker/providers/azure/azure_kubernetes_service.py +++ b/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 ) diff --git a/perfkitbenchmarker/resources/container_service/container.py b/perfkitbenchmarker/resources/container_service/container.py index 3e05a1ec2b..2850ded4ce 100644 --- a/perfkitbenchmarker/resources/container_service/container.py +++ b/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 diff --git a/perfkitbenchmarker/resources/container_service/container_cluster.py b/perfkitbenchmarker/resources/container_service/container_cluster.py index 9458662c98..176d4ce25c 100644 --- a/perfkitbenchmarker/resources/container_service/container_cluster.py +++ b/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: diff --git a/tests/providers/azure/azure_kubernetes_service_test.py b/tests/providers/azure/azure_kubernetes_service_test.py index 7ca09fb29c..6466b937f0 100644 --- a/tests/providers/azure/azure_kubernetes_service_test.py +++ b/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( {