diff --git a/perfkitbenchmarker/providers/azure/azure_kubernetes_service.py b/perfkitbenchmarker/providers/azure/azure_kubernetes_service.py index 8d3864c920..a6d46d8d3f 100644 --- a/perfkitbenchmarker/providers/azure/azure_kubernetes_service.py +++ b/perfkitbenchmarker/providers/azure/azure_kubernetes_service.py @@ -16,6 +16,7 @@ import base64 import json +import time import logging from typing import Any, List @@ -573,6 +574,390 @@ def ApplyFusePVC(self, pvc_name: str) -> None: ) logging.info('Successfully applied BlobFuse PVC.') + def _BuildNodePoolAddCmd( + self, + nodepool_config: container.BaseNodePoolConfig, + node_flags: list[str], + no_wait: bool = False, + ) -> list[str]: + """Builds the az aks nodepool add command. + + Shared by CreateNodePool and CreateNodePoolAsync so the argument list + is never duplicated. + + Args: + nodepool_config: The nodepool configuration. + node_flags: Flags produced by _GetNodeFlags (and version override). + no_wait: If True, appends --no-wait for the async variant. + + Returns: + The complete az CLI command list. + """ + cmd = [ + azure.AZURE_PATH, + 'aks', + 'nodepool', + 'add', + '--cluster-name', + self.name, + '--name', + _AzureNodePoolName(nodepool_config.name), + '--labels', + f'pkb_nodepool={nodepool_config.name}', + ] + if no_wait: + cmd.append('--no-wait') + return cmd + node_flags + + def CreateNodePool( + self, + nodepool_config: container.BaseNodePoolConfig, + node_version: str | None = None, + ) -> None: + """Creates a single named node pool on the cluster.""" + node_flags = self._GetNodeFlags(nodepool_config) + if node_version: + # Remove any --kubernetes-version _GetNodeFlags may have added from + # self.cluster_version, then always append the caller-supplied version. + while '--kubernetes-version' in node_flags: + idx = node_flags.index('--kubernetes-version') + node_flags.pop(idx) + node_flags.pop(idx) + node_flags += ['--kubernetes-version', node_version] + cmd = self._BuildNodePoolAddCmd(nodepool_config, node_flags) + _, stderr, retcode = vm_util.IssueCommand( + cmd, timeout=1800, raise_on_failure=False + ) + if retcode: + raise errors.Resource.CreationError(stderr) + + def DeleteNodePool(self, name: str) -> None: + """Deletes the named node pool.""" + cmd = [ + azure.AZURE_PATH, + 'aks', + 'nodepool', + 'delete', + '--cluster-name', + self.name, + '--name', + _AzureNodePoolName(name), + ] + self.resource_group.args + self._RunCreateClusterCmd(cmd) + + def UpgradeNodePool(self, name: str, target_version: str) -> None: + """Upgrades the named node pool to target_version.""" + cmd = [ + azure.AZURE_PATH, + 'aks', + 'nodepool', + 'upgrade', + '--cluster-name', + self.name, + '--name', + _AzureNodePoolName(name), + '--kubernetes-version', + target_version, + ] + self.resource_group.args + vm_util.IssueCommand(cmd, timeout=1800) + + def UpdateCluster(self) -> None: + """Real cluster-level update via a unique-timestamp tag change. + + Triggers a control-plane operation (cluster-scoped, not pool-scoped) by + updating the cluster tags. Always succeeds because the tag value changes + every call. + """ + cmd = [ + azure.AZURE_PATH, + 'aks', + 'update', + '--name', + self.name, + '--tags', + f'k8s-mgmt-ts={int(time.time())}', + ] + self.resource_group.args + vm_util.IssueCommand(cmd, timeout=1800) + + # ---- Async variants (return opaque handles) ------------------------------- + + def CreateNodePoolAsync( + self, + nodepool_config: container.BaseNodePoolConfig, + node_version: str | None = None, + ) -> str: + """Initiates AKS nodepool create; returns np_succeeded handle.""" + node_flags = self._GetNodeFlags(nodepool_config) + if node_version: + # Remove any --kubernetes-version _GetNodeFlags may have added from + # self.cluster_version, then always append the caller-supplied version. + while '--kubernetes-version' in node_flags: + idx = node_flags.index('--kubernetes-version') + node_flags.pop(idx) + node_flags.pop(idx) + node_flags += ['--kubernetes-version', node_version] + cmd = self._BuildNodePoolAddCmd(nodepool_config, node_flags, no_wait=True) + # fix: raise timeout to 600s (AKS can take >300s to accept a + # --no-wait request under concurrent load) and retry on transient errors + # that indicate the cluster is temporarily at its concurrent-op or + # pool-count limit. + retryable_errors = ( + 'OperationNotAllowed', + 'ConflictingOperationInProgress', + 'MaxAgentPoolCountReached', + ) + + @vm_util.Retry( + retryable_exceptions=(errors.Resource.RetryableCreationError,), + max_retries=5, + poll_interval=30, + log_errors=True, + ) + def IssueWithRetry(): + """Issues create-nodepool command with retry on transient errors.""" + _, stderr, retcode = vm_util.IssueCommand( + cmd, timeout=600, raise_on_failure=False + ) + if retcode: + if any(e in stderr for e in retryable_errors): + raise errors.Resource.RetryableCreationError(stderr) + raise errors.Resource.CreationError(stderr) + + IssueWithRetry() + return f'np_succeeded:{_AzureNodePoolName(nodepool_config.name)}' + + def UpgradeNodePoolAsync(self, name: str, target_version: str) -> str: + """Initiates AKS nodepool upgrade; returns np_succeeded handle.""" + cmd = [ + azure.AZURE_PATH, + 'aks', + 'nodepool', + 'upgrade', + '--cluster-name', + self.name, + '--name', + _AzureNodePoolName(name), + '--kubernetes-version', + target_version, + '--no-wait', + ] + self.resource_group.args + # fix: raise timeout to 600s — az aks nodepool upgrade --no-wait + # can take >300s to be accepted by Azure under concurrent load. + _, stderr, retcode = vm_util.IssueCommand( + cmd, timeout=600, raise_on_failure=False + ) + if retcode: + raise errors.Resource.CreationError(stderr) + return f'np_succeeded:{_AzureNodePoolName(name)}' + + def DeleteNodePoolAsync(self, name: str) -> str: + """Initiates AKS nodepool delete; returns np_gone handle.""" + cmd = [ + azure.AZURE_PATH, + 'aks', + 'nodepool', + 'delete', + '--cluster-name', + self.name, + '--name', + _AzureNodePoolName(name), + '--no-wait', + ] + self.resource_group.args + # fix: raise timeout to 600s and treat NotFound as success. + # A pool that never existed or was already removed is the desired end-state + # for a delete — raising CreationError here caused all delete phases to + # fail for any pool whose create had previously failed. + _, stderr, retcode = vm_util.IssueCommand( + cmd, timeout=600, raise_on_failure=False + ) + if retcode: + if 'NotFound' in stderr or 'not found' in stderr.lower(): + logging.info( + '[AKS] DeleteNodePoolAsync: %s already gone — treating as success', + _AzureNodePoolName(name), + ) + return f'np_gone:{_AzureNodePoolName(name)}' + raise errors.Resource.CreationError(stderr) + return f'np_gone:{_AzureNodePoolName(name)}' + + def UpdateClusterAsync(self) -> str: + """Triggers a tag update on the cluster; returns 'cluster_succeeded'. + + Updating a cluster tag is a lightweight cluster-scoped control-plane + operation. The tag value changes on every call (unix timestamp) so it + is always a real change, giving a meaningful overlap window for + Scenario B concurrent NodePool creates. + """ + cmd = [ + azure.AZURE_PATH, + 'aks', + 'update', + '--name', + self.name, + '--tags', + f'k8s-mgmt-ts={int(time.time())}', + '--no-wait', + ] + self.resource_group.args + _, stderr, retcode = vm_util.IssueCommand( + cmd, timeout=300, raise_on_failure=False + ) + if retcode: + raise errors.Resource.CreationError(stderr) + return 'cluster_succeeded' + + def ResolveNodePoolVersions(self) -> tuple[str, str]: + """Returns (initial, target) AKS node pool versions. + + Uses cluster_version (already set) rather than querying kubectl. + initial = N-1 (adjacent minor below cluster version) + target = N (cluster version = latest) + """ + cluster_ver = self.cluster_version or self.k8s_version + parts = cluster_ver.lstrip('v').split('.') + major, minor = int(parts[0]), int(parts[1]) + target = f'{major}.{minor}' + initial = f'{major}.{minor - 1}' + logging.info( + '[AKS] ResolveNodePoolVersions: cluster=%s initial=%s target=%s', + cluster_ver, + initial, + target, + ) + return initial, target + + def _WaitForProvisioningState( + self, + cmd: list[str], + resource_desc: str, + not_found_is_success: bool = False, + retryable_exception_type=errors.Resource.RetryableCreationError, + ) -> None: + """Polls an AKS resource until it reaches Succeeded or a terminal error. + + Issues cmd on each poll (with a 120 s per-call timeout) and inspects + the provisioningState TSV output. Terminal outcomes: + + Succeeded → return normally. + Failed → raise errors.Resource.CreationError. + NotFound (rc!=0) → return if not_found_is_success, else raise + errors.Resource.CreationError. + Other state → raise retryable_exception_type (retried by Retry). + + Args: + cmd: The az CLI command to run on each poll. + resource_desc: Human-readable label for log / error messages. + not_found_is_success: True when the desired end-state is gone (delete). + retryable_exception_type: Exception class used for in-progress states; + defaults to RetryableCreationError, pass RetryableDeletionError for + delete waits. + """ + + @vm_util.Retry( + poll_interval=5, + fuzz=0, + timeout=3600, + retryable_exceptions=(retryable_exception_type,), + ) + def _Poll(): + out, err, rc = vm_util.IssueCommand( + cmd, raise_on_failure=False, timeout=120 + ) + if rc: + is_not_found = ( + 'NotFound' in (err or '') or 'not found' in (err or '').lower() + ) + if is_not_found and not_found_is_success: + return + if is_not_found: + raise errors.Resource.CreationError( + f'{resource_desc} not found: {err}' + ) + raise retryable_exception_type(err) + status = out.strip() + if status == 'Succeeded': + return + if status == 'Failed': + raise errors.Resource.CreationError( + f'{resource_desc} ended in Failed' + ) + raise retryable_exception_type(f'{resource_desc} state={status}') + + _Poll() + + def WaitForOperation(self, op_handle: str) -> None: + """Polls AKS resources until the expected terminal state is observed. + + Args: + op_handle: Opaque string returned by the *Async methods. Format: + + 'np_succeeded:' + Wait for the nodepool to reach provisioningState=Succeeded. + Example: 'np_succeeded:pkbma001' + + 'np_gone:' + Wait for the nodepool to be deleted (NotFound = success). + Example: 'np_gone:pkbma001' + + 'cluster_succeeded' + Wait for the cluster to reach provisioningState=Succeeded. + """ + kind, _, name = op_handle.partition(':') + + if kind == 'np_succeeded': + self._WaitForProvisioningState( + cmd=[ + azure.AZURE_PATH, + 'aks', + 'nodepool', + 'show', + '--cluster-name', + self.name, + '--name', + name, + '--query', + 'provisioningState', + '--output', + 'tsv', + ] + + self.resource_group.args, + resource_desc=f'nodepool {name}', + ) + elif kind == 'np_gone': + self._WaitForProvisioningState( + cmd=[ + azure.AZURE_PATH, + 'aks', + 'nodepool', + 'show', + '--cluster-name', + self.name, + '--name', + name, + ] + + self.resource_group.args, + resource_desc=f'nodepool {name}', + not_found_is_success=True, + retryable_exception_type=errors.Resource.RetryableDeletionError, + ) + elif kind == 'cluster_succeeded': + self._WaitForProvisioningState( + cmd=[ + azure.AZURE_PATH, + 'aks', + 'show', + '--name', + self.name, + '--query', + 'provisioningState', + '--output', + 'tsv', + ] + + self.resource_group.args, + resource_desc='cluster', + ) + else: + raise ValueError(f'Unknown AKS op handle: {op_handle!r}') + class AksAutomaticCluster(AksCluster): """Class representing an AKS Automatic cluster, which has managed node pools. diff --git a/perfkitbenchmarker/resources/container_service/kubernetes_cluster.py b/perfkitbenchmarker/resources/container_service/kubernetes_cluster.py index 9e20f57583..50b56c3692 100644 --- a/perfkitbenchmarker/resources/container_service/kubernetes_cluster.py +++ b/perfkitbenchmarker/resources/container_service/kubernetes_cluster.py @@ -1,5 +1,6 @@ """Classes related to KubernetesCluster.""" +import abc import functools import json import logging @@ -49,9 +50,25 @@ def Create(self, restore: bool = False) -> None: self.inference_server.Create() def _PostCreate(self): + """Starts the event poller after the cluster has been created.""" super()._PostCreate() - if self.cluster_spec.poll_for_events: - self.event_poller.StartPolling() + if self.event_poller: + try: + self.event_poller.StartPolling() + except Exception as exc: # pylint: disable=broad-except + # Python 3.14 tightened pickling rules for multiprocessing — local + # functions passed to Process cannot be pickled. Rather than crashing + # PKB entirely (which prevents cleanup and orphans cloud resources), + # log a warning and continue without the event poller. + # Impact: no Kubernetes event streaming during the run — benchmark + # metrics are unaffected. + logging.warning( + 'Event poller failed to start (non-fatal, continuing without ' + + 'event polling): %s. This is a known Python 3.14 pickling ' + + 'issue — switch to Python 3.13 to enable event polling.', + exc, + ) + self.event_poller = None def Delete(self, freeze: bool = False) -> None: if self.inference_server: @@ -297,9 +314,132 @@ def _GetAddressFromIngress(self, ingress_out: str): ) return 'http://' + ip.strip() - def AddNodepool(self, batch_name: str, pool_id: str): - """Adds an additional nodepool with the given name to the cluster.""" - pass + + def AddNodepool(self, batch_name: str, pool_id: str) -> None: + """Adds a node pool; delegates to CreateNodePool for standard clusters. + + Karpenter-based subclasses override this to apply a manifest instead. + """ + nodepool_config = container_lib.BaseNodePoolConfig( + self.nodepools[container_cluster.DEFAULT_NODEPOOL].vm_spec, + name=f'{batch_name}-{pool_id}', + ) + self.CreateNodePool(nodepool_config) + + def CreateNodePool( + self, + nodepool_config: container_lib.BaseNodePoolConfig, + node_version: str | None = None, + ) -> None: + """Creates a node pool and blocks until ready (sync wrapper). + + Args: + nodepool_config: Node pool definition (name, machine type, node count). + node_version: Optional Kubernetes version to pin the node pool to. None + means use the cluster default. + """ + op = self.CreateNodePoolAsync(nodepool_config, node_version) + self.WaitForOperation(op) + + def DeleteNodePool(self, name: str) -> None: + """Deletes the named node pool and blocks until removed (sync wrapper).""" + op = self.DeleteNodePoolAsync(name) + self.WaitForOperation(op) + + def UpgradeNodePool(self, name: str, target_version: str) -> None: + """Upgrades the named node pool and blocks until done (sync wrapper).""" + op = self.UpgradeNodePoolAsync(name, target_version) + self.WaitForOperation(op) + + def UpdateCluster(self) -> None: + """Performs a cluster-level update and blocks until done (sync wrapper). + + Intended for management-plane benchmarks that need to overlap a real + cluster-level operation with a node-pool operation. The implementation + should issue a control-plane mutation (so an actual operation runs) that + is non-destructive and idempotent across repeated invocations. + """ + op = self.UpdateClusterAsync() + self.WaitForOperation(op) + + def CreateNodePoolAsync( + self, + nodepool_config: container_lib.BaseNodePoolConfig, + node_version: str | None = None, + ) -> str: + """Initiates node-pool create; returns opaque op handle. Does NOT wait.""" + raise NotImplementedError + + def UpgradeNodePoolAsync(self, name: str, target_version: str) -> str: + """Initiates node-pool upgrade; returns opaque op handle. Does NOT wait.""" + raise NotImplementedError + + def DeleteNodePoolAsync(self, name: str) -> str: + """Initiates node-pool delete; returns opaque op handle. Does NOT wait.""" + raise NotImplementedError + + def UpdateClusterAsync(self) -> str: + """Initiates cluster-level update. Returns op handle; does NOT wait.""" + raise NotImplementedError + + @abc.abstractmethod + def GetNodePoolNames(self) -> list[str]: + """Returns the names of all node pools currently in the cluster. + + Used by the kubernetes_management benchmark to: + - Sweep stale pkbm* pools before each run (clean-start spec requirement) + - Re-list live pools after creates before deleting (avoids stale names) + """ + + def WaitForOperation(self, op_handle: str) -> None: + """Blocks until the operation identified by op_handle completes. + + Args: + op_handle: provider-specific opaque string from one of the *Async + methods above. + + Raises: + errors.Resource.RetryableCreationError or similar on timeout/failure. + """ + raise NotImplementedError + + def ResolveNodePoolVersions(self) -> tuple[str, str]: + """Returns (initial, target) K8s versions per benchmark spec. + + Spec contract: + target = cluster's current K8s version (the latest available) + initial = the adjacent minor below target (e.g., target=1.35 -> 1.34) + Default implementation returns bare-minor strings ("1.34", "1.35") which + EKS and AKS accept directly. Providers requiring fully-qualified versions + (notably GKE) must override. + """ + target = BareMinor(self.k8s_version) + initial = AdjacentMinorBelow(self.k8s_version) + return initial, target + + +def BareMinor(version: str) -> str: + """Returns the 'major.minor' part of a K8s version string. + + Accepts and normalizes formats like 'v1.35.4', '1.35.4-gke.1234', '1.35'. + """ + if version.startswith('v'): + version = version[1:] + bare = version.split('-', 1)[0] + parts = bare.split('.') + if len(parts) < 2 or not parts[0].isdigit() or not parts[1].isdigit(): + raise ValueError(f'Cannot parse K8s version: {version!r}') + return f'{parts[0]}.{parts[1]}' + + +def AdjacentMinorBelow(version: str) -> str: + """Returns the bare minor one below the given version: '1.35.4' -> '1.34'.""" + bare = BareMinor(version) + major_s, minor_s = bare.split('.') + minor = int(minor_s) + if minor <= 0: + raise ValueError(f'No adjacent minor below {version!r}') + return f'{major_s}.{minor - 1}' def _DeleteAllFromDefaultNamespace(): diff --git a/tests/container_service_mock.py b/tests/container_service_mock.py index 5fb9829c38..98960c848d 100644 --- a/tests/container_service_mock.py +++ b/tests/container_service_mock.py @@ -26,6 +26,9 @@ def _Create(self): def _Delete(self): pass + def GetNodePoolNames(self) -> list[str]: + return [] + def CreateTestKubernetesCluster( container_cluster_spec: container_spec.ContainerClusterSpec | None = None, diff --git a/tests/linux_benchmarks/kubernetes_management_benchmark_test.py b/tests/linux_benchmarks/kubernetes_management_benchmark_test.py index a58ad03497..044459c280 100644 --- a/tests/linux_benchmarks/kubernetes_management_benchmark_test.py +++ b/tests/linux_benchmarks/kubernetes_management_benchmark_test.py @@ -12,6 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +# pylint: disable=invalid-name,protected-access + +import threading +import time import unittest from unittest import mock diff --git a/tests/linux_benchmarks/provision_container_cluster_benchmark_test.py b/tests/linux_benchmarks/provision_container_cluster_benchmark_test.py index 7b64d89124..e0ab7ccf31 100644 --- a/tests/linux_benchmarks/provision_container_cluster_benchmark_test.py +++ b/tests/linux_benchmarks/provision_container_cluster_benchmark_test.py @@ -26,6 +26,9 @@ def __init__(self): kubernetes_cluster.kubernetes_commands.GetEvents ) + def GetNodePoolNames(self) -> list[str]: + return [] + _CONTAINER_START_YAML = """ - apiVersion: v1 diff --git a/tests/providers/azure/azure_kubernetes_service_test.py b/tests/providers/azure/azure_kubernetes_service_test.py index 7ca09fb29c..68e3848f1b 100644 --- a/tests/providers/azure/azure_kubernetes_service_test.py +++ b/tests/providers/azure/azure_kubernetes_service_test.py @@ -308,5 +308,55 @@ def testGetNodePoolNames(self): self.assertEqual(self.aks.GetNodePoolNames(), ['default', 'nodepool1']) + def testWaitForOperationNpSucceeded(self): + """np_succeeded handle polls nodepool show until Succeeded.""" + self.MockIssueCommand({ + 'az aks nodepool show': [('Succeeded', '', 0)], + }) + self.aks.WaitForOperation('np_succeeded:pkbma001') # should not raise + + def testWaitForOperationNpSucceededFailed(self): + """np_succeeded raises CreationError when nodepool reaches Failed.""" + self.MockIssueCommand({ + 'az aks nodepool show': [('Failed', '', 0)], + }) + with self.assertRaises(errors.Resource.CreationError): + self.aks.WaitForOperation('np_succeeded:pkbma001') + + def testWaitForOperationNpSucceededNotFound(self): + """np_succeeded raises CreationError when nodepool is not found.""" + self.MockIssueCommand({ + 'az aks nodepool show': [('', 'ResourceNotFound', 1)], + }) + with self.assertRaises(errors.Resource.CreationError): + self.aks.WaitForOperation('np_succeeded:pkbma001') + + def testWaitForOperationNpGone(self): + """np_gone handle returns once nodepool is NotFound.""" + self.MockIssueCommand({ + 'az aks nodepool show': [('', 'ResourceNotFound', 1)], + }) + self.aks.WaitForOperation('np_gone:pkbma001') # should not raise + + def testWaitForOperationClusterSucceeded(self): + """cluster_succeeded handle polls cluster show until Succeeded.""" + self.MockIssueCommand({ + 'az aks show': [('Succeeded', '', 0)], + }) + self.aks.WaitForOperation('cluster_succeeded') # should not raise + + def testWaitForOperationClusterFailed(self): + """cluster_succeeded raises CreationError when cluster reaches Failed.""" + self.MockIssueCommand({ + 'az aks show': [('Failed', '', 0)], + }) + with self.assertRaises(errors.Resource.CreationError): + self.aks.WaitForOperation('cluster_succeeded') + + def testWaitForOperationInvalidHandle(self): + """Unknown op_handle kind raises ValueError.""" + with self.assertRaises(ValueError): + self.aks.WaitForOperation('unknown:foo') + if __name__ == '__main__': unittest.main() diff --git a/tests/traces/kubernetes_tracker_test.py b/tests/traces/kubernetes_tracker_test.py index 6034ff2765..96613d28aa 100644 --- a/tests/traces/kubernetes_tracker_test.py +++ b/tests/traces/kubernetes_tracker_test.py @@ -264,6 +264,9 @@ def _Create(self): def _Delete(self): pass + def GetNodePoolNames(self) -> list[str]: + return [] + def _CreateEvent( name: str, reason: str, timestamp: float