Skip to content

Dipowell/node readiness timing #1208

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
diamondpowell wants to merge 12 commits into
base: main
Choose a base branch
from
Open

Dipowell/node readiness timing #1208

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
16ba109
Add return_timestamp parameter to wait_for_nodes_ready
diamondpowell Jun 4, 2026
b00a147
Replace asyncio with ThreadPoolExecutor for concurrent timing
diamondpowell Jun 24, 2026
006c87d
Internalize ARM call and start_time in concurrent method
diamondpowell Jun 24, 2026
8658356
Rename _run_concurrent_arm_and_readiness to _instrument_nodepool_prov…
diamondpowell Jun 24, 2026
d2afe5a
Enhance timing logs with bottleneck analysis and total elapsed
diamondpowell Jun 24, 2026
49b64af
Add pipeline test config (to be reverted before merge)
diamondpowell Jun 24, 2026
30a6908
Extract _log_timing_metrics helper and trim _instrument_nodepool_prov…
diamondpowell Jun 24, 2026
0a1b010
Revert "Add pipeline test config (to be reverted before merge)"
diamondpowell Jun 24, 2026
666c1f1
Address review: simplify timing, restore retry, remove return_timestamp
diamondpowell Jul 1, 2026
d286aaf
Fix test: add time.time() side_effect values for concurrent timing
diamondpowell Jul 27, 2026
94bf515
Add pipeline test config (to be reverted before merge)
diamondpowell Jul 27, 2026
1b699f2
Revert "Add pipeline test config (to be reverted before merge)"
diamondpowell Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
156 changes: 108 additions & 48 deletions modules/python/clients/aks_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import os
import subprocess
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, Optional, Any

# Third party imports
Expand Down Expand Up @@ -61,6 +62,93 @@ 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,
) -> 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

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),
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,
Expand Down Expand Up @@ -519,23 +607,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"
Expand Down Expand Up @@ -688,24 +770,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(

@liyu-ma liyu-ma Jun 25, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The deletion here is a behavior change. Your _instrument_nodepool_provisioning calls begin_create_or_update which does not do retry. I would suggest using _begin_update_with_retry within your instrumentation as it retries on transient errors - but this may skew your metrics.

@diamondpowell diamondpowell Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, my initial refactor, I dropped it when I moved the ARM call inside _instrument_nodepool_provisioning. So went ahead and added a use_retry parameter so when True, the method calls _begin_update_with_retry instead of begin_create_or_update. scale_node_pool and _progressive_scale both pass use_retry=True now, so the transient error handling is restored.

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"
)
Expand Down Expand Up @@ -882,7 +957,6 @@ def _progressive_scale(

logger.info(f"Planned scaling steps: {list(steps)}")

result = None
completed_steps = []

# Execute scaling operation for each step
Expand Down Expand Up @@ -924,28 +998,14 @@ 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(

@liyu-ma liyu-ma Jun 25, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as above. Now the new codes lose the retry support.

@diamondpowell diamondpowell Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes this is fixed now.

node_pool_name=node_pool_name,
cluster_name=cluster_name,
node_pool=node_pool,
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,
# 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,
)
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
logger.info(f"All {step} nodes in pool {node_pool_name} are ready")

op.add_metadata(
"ready_nodes", len(ready_nodes) if ready_nodes else 0
Expand Down
4 changes: 2 additions & 2 deletions modules/python/clients/kubernetes_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
70 changes: 56 additions & 14 deletions modules/python/tests/clients/test_aks_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"""
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -366,14 +385,25 @@ 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"""
# Setup
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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading