diff --git a/modules/python/clients/aks_client.py b/modules/python/clients/aks_client.py index 4e571b2840..c1ededce95 100644 --- a/modules/python/clients/aks_client.py +++ b/modules/python/clients/aks_client.py @@ -16,6 +16,7 @@ import os import subprocess import time +from concurrent.futures import ThreadPoolExecutor from typing import Dict, Optional, Any # Third party imports @@ -61,6 +62,95 @@ def _get_operation_context(self): return OperationContext + def _instrument_nodepool_provisioning( + self, + node_pool_name: str, + cluster_name: str, + parameters: Any, + node_count: int, + op, + use_retry: bool = False, + label: str = "", + ) -> list: + """ + Run ARM operation and K8s node readiness check concurrently using threads. + + Args: + node_pool_name: Name of the node pool being provisioned + cluster_name: Name of the AKS cluster + parameters: Parameters for begin_create_or_update (dict or node pool object) + node_count: Expected number of nodes to be ready + op: Operation context for recording timing metadata + use_retry: If True, use _begin_update_with_retry for transient error handling + label: Optional label for retry log messages (e.g. "step 2 ") + + Returns: + List of ready nodes + + Raises: + Exception: If either the ARM operation or K8s readiness check fails. + """ + start_time = time.time() + label_selector = f"agentpool={node_pool_name}" + + with ThreadPoolExecutor(max_workers=2) as executor: + if use_retry: + arm_future = executor.submit( + lambda: ( + self._begin_update_with_retry(node_pool_name, cluster_name, parameters, label=label), + time.time(), + ) + ) + else: + poller = self.aks_client.agent_pools.begin_create_or_update( + resource_group_name=self.resource_group, + resource_name=cluster_name, + agent_pool_name=node_pool_name, + parameters=parameters, + ) + arm_future = executor.submit(lambda: (poller.result(), time.time())) + k8s_future = executor.submit( + lambda: ( + self.k8s_client.wait_for_nodes_ready( + node_count=node_count, + operation_timeout_in_minutes=self.operation_timeout_minutes, + label_selector=label_selector, + ), + time.time(), + ) + ) + + arm_exc = arm_future.exception() + k8s_exc = k8s_future.exception() + + if arm_exc or k8s_exc: + elapsed = time.time() - start_time + arm_status = f"FAILED: {arm_exc}" if arm_exc else "succeeded" + k8s_status = f"FAILED: {k8s_exc}" if k8s_exc else "succeeded" + logger.error( + "Concurrent operation failed after %.2fs - ARM: %s, K8s readiness: %s", + elapsed, arm_status, k8s_status + ) + if arm_exc: + raise arm_exc + raise k8s_exc + + _, arm_timestamp = arm_future.result() + ready_nodes, ready_timestamp = k8s_future.result() + + node_readiness_time = ready_timestamp - start_time + command_execution_time = arm_timestamp - start_time + + op.add_metadata("node_readiness_time", node_readiness_time) + op.add_metadata("command_execution_time", command_execution_time) + logger.info( + "[%s] ARM completed in %.2fs, K8s nodes ready in %.2fs | Delta: %.2fs", + node_pool_name, command_execution_time, node_readiness_time, + abs(command_execution_time - node_readiness_time) + ) + + return ready_nodes + def __init__( self, subscription_id: Optional[str] = None, @@ -519,23 +609,17 @@ def create_node_pool( gpu_instance_profile=gpu_instance_profile, gpu_mig_strategy=gpu_mig_strategy, ) + label_selector = f"agentpool={node_pool_name}" + ready_nodes = self.k8s_client.wait_for_nodes_ready( + node_count=node_count, + operation_timeout_in_minutes=self.operation_timeout_minutes, + label_selector=label_selector, + ) else: - # Create the node pool via SDK - self.aks_client.agent_pools.begin_create_or_update( - resource_group_name=self.resource_group, - resource_name=cluster_name, - agent_pool_name=node_pool_name, - parameters=parameters, - ).result() - - label_selector = f"agentpool={node_pool_name}" - - # Wait for nodes to be ready - ready_nodes = self.k8s_client.wait_for_nodes_ready( - node_count=node_count, - operation_timeout_in_minutes=self.operation_timeout_minutes, - label_selector=label_selector, - ) + # Run ARM and K8s readiness concurrently to capture both timings + ready_nodes = self._instrument_nodepool_provisioning( + node_pool_name, cluster_name, parameters, node_count, op + ) logger.info( f"All {node_count} nodes in pool {node_pool_name} are ready" @@ -688,24 +772,17 @@ def scale_node_pool( node_pool.count = node_count logger.info(f"Scaling node pool {node_pool_name} to {node_count} nodes") - self._begin_update_with_retry( - node_pool_name=node_pool_name, - cluster_name=cluster_name, - node_pool=node_pool, - ) logger.info( f"Waiting for {node_count} nodes in pool {node_pool_name} to be ready..." ) - # Use agentpool=node_pool_name as default label if not specified - label_selector = f"agentpool={node_pool_name}" - - ready_nodes = self.k8s_client.wait_for_nodes_ready( - node_count=node_count, - operation_timeout_in_minutes=self.operation_timeout_minutes, - label_selector=label_selector, + # Run ARM and K8s readiness concurrently to capture both timings + ready_nodes = self._instrument_nodepool_provisioning( + node_pool_name, cluster_name, node_pool, node_count, op, + use_retry=True, ) + logger.info( f"All {node_count} nodes in pool {node_pool_name} are ready" ) @@ -882,7 +959,6 @@ def _progressive_scale( logger.info(f"Planned scaling steps: {list(steps)}") - result = None completed_steps = [] # Execute scaling operation for each step @@ -924,29 +1000,16 @@ def _progressive_scale( "cluster_info", self.get_cluster_data(cluster_name) ) node_pool.count = step # Update node count in the node pool object - self._begin_update_with_retry( - node_pool_name=node_pool_name, - cluster_name=cluster_name, - node_pool=node_pool, + + # Run ARM and K8s readiness concurrently to capture both timings + ready_nodes = self._instrument_nodepool_provisioning( + node_pool_name, cluster_name, node_pool, step, op, + use_retry=True, label=f"step {step} ", ) - result = node_pool - # Use agentpool=node_pool_name as default label if not specified - label_selector = f"agentpool={node_pool_name}" - - ready_nodes = self.k8s_client.wait_for_nodes_ready( - node_count=step, - operation_timeout_in_minutes=self.operation_timeout_minutes, - label_selector=label_selector, - ) logger.info(f"All {step} nodes in pool {node_pool_name} are ready") - if result is None: - logger.error(f"Progressive scaling failed at step {step}") - op.add_metadata("error", "Scaling operation returned None") - return None - op.add_metadata( "ready_nodes", len(ready_nodes) if ready_nodes else 0 ) diff --git a/modules/python/clients/kubernetes_client.py b/modules/python/clients/kubernetes_client.py index 9d7b171353..a564c503f9 100644 --- a/modules/python/clients/kubernetes_client.py +++ b/modules/python/clients/kubernetes_client.py @@ -242,10 +242,10 @@ def wait_for_nodes_ready(self, node_count, operation_timeout_in_minutes, label_s Waits for a specific number of nodes with a given label to be ready within a specified timeout. Raises an exception if the expected number of nodes are not ready within the timeout. - :param node_label: The label to filter nodes. :param node_count: The expected number of nodes to be ready. :param operation_timeout_in_minutes: The timeout in minutes to wait for the nodes to be ready. - :return: None + :param label_selector: The label to filter nodes. + :return: List of ready nodes. """ ready_nodes = [] ready_node_count = 0 diff --git a/modules/python/tests/clients/test_aks_client.py b/modules/python/tests/clients/test_aks_client.py index d9f6a6f11f..c621cae849 100644 --- a/modules/python/tests/clients/test_aks_client.py +++ b/modules/python/tests/clients/test_aks_client.py @@ -174,7 +174,11 @@ def test_create_node_pool_success(self, mock_time): vm_size = "Standard_DS2_v2" node_count = 2 - mock_time.time.side_effect = [100, 150] # Start and end times + # Define timestamps for clarity + start_time = 100 + arm_done_time = 150 + nodes_ready_time = 130 + mock_time.time.side_effect = [start_time, arm_done_time, nodes_ready_time] mock_operation = mock.MagicMock() self.mock_agent_pools.begin_create_or_update.return_value = mock_operation @@ -219,6 +223,12 @@ def mock_get_node_pool_side_effect(name, cluster=None): label_selector=f"agentpool={node_pool_name}", ) + # Verify timing measurements are calculated and stored correctly + # node_readiness_time = nodes_ready_time - start_time = 130 - 100 = 30 + # command_execution_time = arm_done_time - start_time = 150 - 100 = 50 + self.mock_operation.add_metadata.assert_any_call("node_readiness_time", 30) + self.mock_operation.add_metadata.assert_any_call("command_execution_time", 50) + @mock.patch("clients.aks_client.time") def test_create_node_pool_gpu(self, mock_time): """Test creating a GPU node pool with driver verification""" @@ -227,7 +237,11 @@ def test_create_node_pool_gpu(self, mock_time): vm_size = "Standard_NC6s_v3" # GPU VM size node_count = 1 - mock_time.time.side_effect = [100, 150] # Start and end times + # Define timestamps for clarity + start_time = 100 + arm_done_time = 150 + nodes_ready_time = 130 + mock_time.time.side_effect = [start_time, arm_done_time, nodes_ready_time] mock_operation = mock.MagicMock() self.mock_agent_pools.begin_create_or_update.return_value = mock_operation @@ -329,7 +343,12 @@ def test_scale_node_pool_up(self, mock_time): node_pool_name = "test-pool" node_count = 3 - mock_time.time.side_effect = [100, 150] # Start and end times + # Define timestamps for clarity + start_time = 100 + arm_done_time = 150 + nodes_ready_time = 130 + mock_time.time.side_effect = [start_time, arm_done_time, nodes_ready_time] + mock_time.sleep = mock.MagicMock() mock_node_pool = mock.MagicMock() mock_node_pool.count = 1 # Current count @@ -366,6 +385,12 @@ def test_scale_node_pool_up(self, mock_time): ) self.assertEqual(mock_node_pool.count, node_count) + # Verify timing measurements are calculated and stored correctly + # node_readiness_time = nodes_ready_time - start_time = 130 - 100 = 30 + # command_execution_time = arm_done_time - start_time = 150 - 100 = 50 + self.mock_operation.add_metadata.assert_any_call("node_readiness_time", 30) + self.mock_operation.add_metadata.assert_any_call("command_execution_time", 50) + @mock.patch("clients.aks_client.time") def test_scale_node_pool_down(self, mock_time): """Test scaling a node pool down""" @@ -373,7 +398,12 @@ def test_scale_node_pool_down(self, mock_time): node_pool_name = "test-pool" node_count = 1 - mock_time.time.side_effect = [100, 150] # Start and end times + # Define timestamps for clarity + start_time = 100 + arm_done_time = 150 + nodes_ready_time = 130 + mock_time.time.side_effect = [start_time, arm_done_time, nodes_ready_time] + mock_time.sleep = mock.MagicMock() mock_node_pool = mock.MagicMock() mock_node_pool.count = 3 # Current count @@ -454,7 +484,12 @@ def test_scale_gpu_node_pool_up_final_target(self, mock_time): # Set VM size that triggers NVIDIA verification self.aks_client.vm_size = "Standard_NC40ads_H100_v5" - mock_time.time.side_effect = [100, 150] # Start and end times + # Define timestamps for clarity + start_time = 100 + arm_done_time = 150 + nodes_ready_time = 130 + mock_time.time.side_effect = [start_time, arm_done_time, nodes_ready_time] + mock_time.sleep = mock.MagicMock() mock_node_pool = mock.MagicMock() mock_node_pool.count = 1 # Current count @@ -514,12 +549,10 @@ def test_scale_gpu_node_pool_up_progressive_final_step(self, mock_time): # Set VM size that triggers NVIDIA verification self.aks_client.vm_size = "Standard_NC40ads_H100_v5" - mock_time.time.side_effect = [ - 100, - 150, - 200, - 250, - ] # Start and end times including progressive scaling + # Define timestamps for clarity (progressive has multiple steps) + # Each step calls time.time() 3x: start, arm_done, nodes_ready + mock_time.time.side_effect = [100, 150, 130, 200, 250, 230] + mock_time.sleep = mock.MagicMock() mock_node_pool = mock.MagicMock() mock_node_pool.count = 1 # Current count @@ -548,7 +581,10 @@ def test_scale_gpu_node_pool_up_progressive_final_step(self, mock_time): mock.MagicMock(), mock.MagicMock(), ] # Second step to 3 nodes - self.mock_k8s.wait_for_nodes_ready.side_effect = [ready_nodes1, ready_nodes2] + self.mock_k8s.wait_for_nodes_ready.side_effect = [ + ready_nodes1, + ready_nodes2, + ] # Add nvidia-smi verification mock self.mock_k8s.verify_nvidia_smi_on_node = mock.MagicMock( @@ -584,7 +620,12 @@ def test_scale_gpu_node_pool_down_no_verification(self, mock_time): node_pool_name = "gpu-pool" node_count = 1 - mock_time.time.side_effect = [100, 150] # Start and end times + # Define timestamps for clarity + start_time = 100 + arm_done_time = 150 + nodes_ready_time = 130 + mock_time.time.side_effect = [start_time, arm_done_time, nodes_ready_time] + mock_time.sleep = mock.MagicMock() mock_node_pool = mock.MagicMock() mock_node_pool.count = 3 # Current count @@ -701,7 +742,8 @@ def test_scale_node_pool_records_gpu_mode_metadata(self, mock_time): node_pool_name = "h100fullmgd" node_count = 3 - mock_time.time.side_effect = [100, 150] + mock_time.time.side_effect = [100, 150, 130] + mock_time.sleep = mock.MagicMock() mock_node_pool = mock.MagicMock() mock_node_pool.count = 1