Merge pull request #6029 from kiryl-filatau:aks-nodepool-provisioning

PiperOrigin-RevId: 820775455

Author: copybara-github

perfkitbenchmarker/data/provision_node_pools/aks/nodepool.yaml.j2

@@ -0,0 +1,43 @@
+{% if id == '001' -%}
+---
+apiVersion: karpenter.azure.com/v1beta1
+kind: AKSNodeClass
+metadata:
+  name: aks-managed-class
+spec:
+  imageFamily: "Ubuntu2204"
+{%- endif %}
+---
+apiVersion: karpenter.sh/v1
+kind: NodePool
+metadata:
+  name: {{ batch }}-nodepool-{{ id }}
+spec:
+  disruption:
+    budgets:
+    - nodes: 30%
+    consolidateAfter: 0s
+    consolidationPolicy: WhenEmptyOrUnderutilized
+  template:
+    spec:
+      nodeClassRef:
+        group: karpenter.azure.com
+        kind: AKSNodeClass
+        name: aks-managed-class
+      requirements:
+        - key: kubernetes.io/arch
+          operator: In
+          values: ["amd64"]
+        - key: kubernetes.io/os
+          operator: In
+          values: ["linux"]
+        - key: karpenter.sh/capacity-type
+          operator: In
+          values: ["on-demand"]
+        - key: karpenter.azure.com/sku-family
+          operator: In
+          values: ["D"]
+        - key: "group"
+          operator: In
+          values:
+          - pool-{{ batch }}-{{ id }}

perfkitbenchmarker/linux_benchmarks/provisioning_benchmarks/provision_node_pools_benchmark.py

@@ -77,6 +77,8 @@
     gke_max_cpu: 520
     # (100 init_batch + 20 test_batch + 10 buffer) * 16 GB
     gke_max_memory: 2080
+    # Azure flag for AKS cluster with  node auto provisioning feature
+    azure_aks_auto_node_provisioning: True
 """
 
 INIT_BATCH_NAME = "init-batch"
@@ -102,6 +104,7 @@ def _CreateJobsAndWait(
       batch_name,
       jobs,
   )
+
   apply_start = time.monotonic()
   for i in range(jobs):
     cluster.AddNodepool(batch_name, pool_id="{:03d}".format(i + 1))

perfkitbenchmarker/providers/azure/azure_kubernetes_service.py

@@ -170,12 +170,16 @@ def _Create(self):
         '--nodepool-labels',
         f'pkb_nodepool={container_service.DEFAULT_NODEPOOL}',
     ] + self._GetNodeFlags(self.default_nodepool)
-    if self._IsAutoscalerEnabled():
-      cmd += [
-          '--enable-cluster-autoscaler',
-          f'--min-count={self.min_nodes}',
-          f'--max-count={self.max_nodes}',
-      ]
+    if FLAGS.azure_aks_auto_node_provisioning:
+      # For provision_node_pools benchmark, add auto provisioning mode
+      cmd.append('--node-provisioning-mode=auto')
+    else:
+      if self._IsAutoscalerEnabled():
+        cmd += [
+            '--enable-cluster-autoscaler',
+            f'--min-count={self.min_nodes}',
+            f'--max-count={self.max_nodes}',
+        ]
 
     # TODO(pclay): expose quota and capacity errors
     # Creating an AKS cluster with a fresh service principal usually fails due
@@ -435,17 +439,40 @@ def ResizeNodePool(
 
   def GetNodePoolNames(self) -> list[str]:
     """Get node pool names for the cluster."""
-    cmd = [
-        azure.AZURE_PATH,
-        'aks',
-        'nodepool',
-        'list',
-        '--cluster-name',
-        self.name,
-    ] + self.resource_group.args
-    stdout, _, _ = vm_util.IssueCommand(cmd)
-    nodepools = json.loads(stdout)
-    return [nodepool['name'] for nodepool in nodepools]
+    if FLAGS.azure_aks_auto_node_provisioning:
+      cmd = [
+          FLAGS.kubectl,
+          '--kubeconfig',
+          FLAGS.kubeconfig,
+          'get',
+          'nodepools',
+          '-o',
+          'json',
+      ]
+      stdout, _, _ = vm_util.IssueCommand(cmd)
+      nodepools = json.loads(stdout).get('items', [])
+      return [nodepool['metadata']['name'] for nodepool in nodepools]
+    else:
+      cmd = [
+          azure.AZURE_PATH,
+          'aks',
+          'nodepool',
+          'list',
+          '--cluster-name',
+          self.name,
+      ] + self.resource_group.args
+      stdout, _, _ = vm_util.IssueCommand(cmd)
+      nodepools = json.loads(stdout)
+      return [nodepool['name'] for nodepool in nodepools]
+
+  def AddNodepool(self, batch_name, pool_id):
+    """Add a Karpenter NodePool and AKSNodeClass to the AKS cluster."""
+    self.ApplyManifest(
+        'provision_node_pools/aks/nodepool.yaml.j2',
+        batch=batch_name,
+        id=pool_id,
+        cluster_name=self.name,
+    )
 
 
 class AksAutomaticCluster(AksCluster):

perfkitbenchmarker/providers/azure/flags.py

@@ -181,3 +181,9 @@
     'Enable/Disable secure boot of the VM to allow unsigned operating systems'
     ' and drivers. Defaults to None to fallback to Azure default behavior',
 )
+flags.DEFINE_boolean(
+    'azure_aks_auto_node_provisioning',
+    False,
+    'Enable automatic node provisioning for Azure AKS clusters. '
+    'This is automatically set to True for the provision_node_pools benchmark.',
+)