From 9189f4c4e4a656be7b151732d2f887a9b3a1f301 Mon Sep 17 00:00:00 2001 From: Rui Gao Date: Fri, 21 Aug 2026 02:19:48 +0000 Subject: [PATCH 1/6] Fix node-recycler filtering with cached layout and live-node retries --- .../alert-manager-deployment.yaml.template | 5 + .../src/node-recycler/recycler.py | 218 +++++++++++++++++- .../src/node-recycler/requirements.txt | 1 + .../src/node-recycler/tests/test_recycler.py | 157 ++++++++++++- .../features/node_action/client.py | 4 + .../features/node_status/client.py | 25 +- .../tests/test_node_status_client.py | 1 + 7 files changed, 386 insertions(+), 25 deletions(-) diff --git a/src/alert-manager/deploy/alert-manager-deployment.yaml.template b/src/alert-manager/deploy/alert-manager-deployment.yaml.template index 21d4ff40..56875fdb 100755 --- a/src/alert-manager/deploy/alert-manager-deployment.yaml.template +++ b/src/alert-manager/deploy/alert-manager-deployment.yaml.template @@ -312,6 +312,8 @@ spec: volumeMounts: - name: config-volume mountPath: /app/config + - name: pai-configuration + mountPath: /pai-cluster-config {% if "icm" in cluster_cfg["alert-manager"]["node-recycler"] %} - name: icm-certs mountPath: /etc/icm/certs @@ -456,6 +458,9 @@ spec: - name: config-volume configMap: name: alertmanager-configmap + - name: pai-configuration + configMap: + name: pai-configuration {% if cluster_cfg["alert-manager"]["alert-handler"]["configured"] %} {% if 'email-admin' in cluster_cfg["alert-manager"]["actions-available"] %} - name: templates-volume diff --git a/src/alert-manager/src/node-recycler/recycler.py b/src/alert-manager/src/node-recycler/recycler.py index 4ae993bb..47475c69 100644 --- a/src/alert-manager/src/node-recycler/recycler.py +++ b/src/alert-manager/src/node-recycler/recycler.py @@ -21,6 +21,7 @@ import icm import requests +import yaml from azure.identity import DefaultAzureCredential from azure.mgmt.compute import ComputeManagementClient @@ -34,6 +35,7 @@ class NodeRecycler: + _layout_config_path = os.getenv("PAI_LAYOUT_PATH", "/pai-cluster-config/layout.yaml") _icm_host = os.getenv("ICM_HOST", "prod.microsofticm.com") _icm_cert_path = os.getenv("ICM_CERT_PATH", "cert.pem") _icm_key_path = os.getenv("ICM_KEY_PATH", "key.pem") @@ -48,6 +50,177 @@ class NodeRecycler: _ltp_validation_image = os.getenv("LTP_VALIDATION_IMAGE") _ltp_vmss_ids = os.getenv("LTP_VMSS_IDS", "") _validation_skip_vmss_ids = set(filter(None, os.getenv("VALIDATION_SKIP_VMSS_IDS", "").split(","))) + _live_nodes_retry_attempts = int(os.getenv("LIVE_NODES_RETRY_ATTEMPTS", "3")) + _live_nodes_retry_interval_seconds = int(os.getenv("LIVE_NODES_RETRY_INTERVAL_SECONDS", "1")) + _layout_nodes_cache = set() + _layout_nodes_loaded = False + _layout_nodes_load_success = False + + @classmethod + def _initialize_layout_nodes(cls) -> bool: + """Load layout nodes once and cache the result for process lifetime.""" + if cls._layout_nodes_loaded: + return cls._layout_nodes_load_success + + layout_nodes, layout_ok = cls._load_layout_nodes() + cls._layout_nodes_cache = layout_nodes + cls._layout_nodes_load_success = layout_ok + cls._layout_nodes_loaded = True + + if not layout_ok: + logger.error("Failed to initialize layout nodes from %s", cls._layout_config_path) + return cls._layout_nodes_load_success + + @classmethod + def _load_layout_nodes(cls) -> tuple[set[str], bool]: + """Load node names from cluster layout file. + + Returns: + tuple[set[str], bool]: (node names, success flag) + """ + try: + with open(cls._layout_config_path, "r") as f: + layout = yaml.safe_load(f) or {} + machine_list = layout.get("machine-list", []) + nodes = { + str(machine.get("nodename", "")).strip().lower() + for machine in machine_list + if machine.get("nodename") + } + logger.info("Loaded %d nodes from layout file %s", len(nodes), cls._layout_config_path) + return nodes, True + except Exception as e: + logger.error("Failed to load layout nodes from %s: %s", cls._layout_config_path, e) + return set(), False + + @classmethod + def _load_live_nodes(cls) -> tuple[set[str], set[str], bool]: + """Load live and ready node names from rest-server kubernetes API. + + Returns: + tuple[set[str], set[str], bool]: (live nodes, ready nodes, success flag) + """ + if not cls._ltp_rest_server_uri or not cls._ltp_rest_server_token: + logger.error("REST_SERVER_URI or REST_SERVER_TOKEN is not configured") + return set(), set(), False + + url = f"{cls._ltp_rest_server_uri}/api/v1/kubernetes/nodes" + headers = {"Authorization": f"Bearer {cls._ltp_rest_server_token}"} + + try: + response = requests.get(url, headers=headers, timeout=30) + response.raise_for_status() + payload = response.json() or {} + items = payload.get("items", []) + + live_nodes = set() + ready_nodes = set() + for node in items: + name = str(node.get("metadata", {}).get("name", "")).strip().lower() + if not name: + continue + live_nodes.add(name) + conditions = node.get("status", {}).get("conditions", []) + ready_condition = next((c for c in conditions if c.get("type") == "Ready"), None) + if ready_condition and ready_condition.get("status") == "True": + ready_nodes.add(name) + + logger.info("Loaded %d live nodes (%d ready) from %s", len(live_nodes), len(ready_nodes), url) + return live_nodes, ready_nodes, True + except Exception as e: + logger.error("Failed to load live nodes from %s: %s", url, e) + return set(), set(), False + + @classmethod + def _load_live_nodes_with_retry(cls, stage: str) -> tuple[set[str], set[str], bool]: + """Load live nodes with retries. Returns failure after max attempts.""" + attempts = max(1, cls._live_nodes_retry_attempts) + interval_seconds = max(0, cls._live_nodes_retry_interval_seconds) + + for attempt in range(1, attempts + 1): + live_nodes, ready_nodes, ok = cls._load_live_nodes() + if ok: + return live_nodes, ready_nodes, True + + if attempt < attempts: + wait_seconds = interval_seconds * (2 ** (attempt - 1)) + logger.warning( + "[%s] failed to load live nodes (attempt %d/%d), retrying in %ss", + stage, + attempt, + attempts, + wait_seconds, + ) + time.sleep(wait_seconds) + + logger.error("[%s] failed to load live nodes after %d attempts", stage, attempts) + return set(), set(), False + + @classmethod + def _filter_nodes_by_policy( + cls, + hostnames: list[str], + stage: str, + require_layout: bool, + require_live: bool, + require_ready: bool = False, + ) -> tuple[list[str], bool]: + """Filter nodes by stage policy and return filtered list. + + Returns: + tuple[list[str], bool]: (filtered hostnames, success flag) + """ + normalized = [h.strip().lower() for h in hostnames if h and h.strip()] + if not normalized: + return [], True + + layout_nodes = set() + layout_filter_enabled = require_layout + if require_layout: + layout_ok = cls._initialize_layout_nodes() + if not layout_ok: + logger.warning("[%s] layout nodes unavailable, skipping layout filter and continuing", stage) + layout_filter_enabled = False + else: + layout_nodes = cls._layout_nodes_cache + if not layout_nodes: + logger.warning("[%s] layout nodes is empty, skipping layout filter and continuing", stage) + layout_filter_enabled = False + + live_nodes, ready_nodes, live_ok = set(), set(), True + if require_live or require_ready: + live_nodes, ready_nodes, live_ok = cls._load_live_nodes_with_retry(stage) + if not live_ok: + logger.error("[%s] skipping stage because live nodes cannot be loaded", stage) + return [], False + + kept = [] + skipped = [] + for hostname in normalized: + reasons = [] + if layout_filter_enabled and hostname not in layout_nodes: + reasons.append("not_in_layout") + if require_live and hostname not in live_nodes: + reasons.append("not_live") + if require_ready and hostname not in ready_nodes: + reasons.append("not_ready") + + if reasons: + skipped.append((hostname, ",".join(reasons))) + else: + kept.append(hostname) + + logger.info( + "[%s] candidate filtering: raw=%d kept=%d skipped=%d", + stage, + len(normalized), + len(kept), + len(skipped), + ) + for hostname, reason in skipped[:20]: + logger.info("[%s] skipped node %s: %s", stage, hostname, reason) + + return kept, True @classmethod def ofr(cls, node_faults=None, status_client=None, action_client=None): @@ -69,7 +242,25 @@ def ofr(cls, node_faults=None, status_client=None, action_client=None): created = [] if not node_faults and status_client and action_client: node_faults = [] - for node in status_client.get_nodes_by_status(from_state): + status_nodes = status_client.get_nodes_by_status(from_state) + candidates = [n.HostName for n in status_nodes] + filtered_candidates, ok = cls._filter_nodes_by_policy( + candidates, + stage="ofr-triaged_hardware", + require_layout=True, + require_live=True, + require_ready=False, + ) + if not ok: + logger.error("[ofr-triaged_hardware] skipping OFR stage due to filtering prerequisite failure") + return + + nodes_by_hostname = {str(n.HostName).strip().lower(): n for n in status_nodes} + for hostname in filtered_candidates: + node = nodes_by_hostname.get(hostname) + if node is None: + logger.warning("[%s] node %s disappeared from status query after filtering", "ofr-triaged_hardware", hostname) + continue try: hostname, node_id = node.HostName, node.NodeId logger.info(f"INFO: Querying node {hostname} with node id {node_id} in state {from_state}") @@ -229,6 +420,18 @@ def operate(cls, vmss_id, operation="start", hostnames=None, status_client=None, if not hostnames and status_client: hostnames = [n.HostName for n in status_client.get_nodes_by_status(from_state)] + + hostnames, ok = cls._filter_nodes_by_policy( + hostnames or [], + stage=f"operate-{from_state}-to-{to_state}", + require_layout=(from_state in (NodeStatus.UA.value, NodeStatus.DEALLOCATED_UA.value)), + require_live=False, + require_ready=False, + ) + if not ok: + logger.error("[operate-%s-to-%s] skipping VM operation due to filtering prerequisite failure", from_state, to_state) + return [] + if not hostnames: return [] logger.info(f"Operating on {hostnames} hostnames in VMSS {vmss_id} with operation {op}") @@ -325,6 +528,18 @@ def validate(cls, hostnames=None, filter_state='', status_client=None, action_cl filter_state = NodeStatus.ALLOCATED_UA.value if not hostnames and status_client: hostnames = [n.HostName for n in status_client.get_nodes_by_status(filter_state)] + + hostnames, ok = cls._filter_nodes_by_policy( + hostnames or [], + stage=f"validate-{filter_state}", + require_layout=True, + require_live=True, + require_ready=True, + ) + if not ok: + logger.error("[validate-%s] skipping validation stage due to filtering prerequisite failure", filter_state) + return + if not hostnames: return @@ -519,6 +734,7 @@ def node_recycle_pipeline_loop(cls, interval=600): """ status_client = create_node_status_client() action_client = create_node_action_client() + cls._initialize_layout_nodes() logger.info("Created storage clients for node status and action tables") while True: logger.info(f"{datetime.now()} Starting to UA and Deallocate Nodes") diff --git a/src/alert-manager/src/node-recycler/requirements.txt b/src/alert-manager/src/node-recycler/requirements.txt index b2a807ae..e0aa404b 100644 --- a/src/alert-manager/src/node-recycler/requirements.txt +++ b/src/alert-manager/src/node-recycler/requirements.txt @@ -5,3 +5,4 @@ azure-identity azure-mgmt-compute requests +PyYAML diff --git a/src/alert-manager/src/node-recycler/tests/test_recycler.py b/src/alert-manager/src/node-recycler/tests/test_recycler.py index 984644b3..3d3607bb 100644 --- a/src/alert-manager/src/node-recycler/tests/test_recycler.py +++ b/src/alert-manager/src/node-recycler/tests/test_recycler.py @@ -72,6 +72,7 @@ def action_client(): def recycler(mock_env): """Configure NodeRecycler for testing.""" cls = recycler_module.NodeRecycler + default_nodes = {"node-a", "node-b", "test-node", "gpu-node-1", "cpu-node-1"} cls._ltp_rest_server_uri = "http://test-server" cls._ltp_rest_server_token = "test-token" cls._ltp_validation_image = "test-image:latest" @@ -80,6 +81,12 @@ def recycler(mock_env): cls._validation_skip_vmss_ids = {"vmss-cpu-1"} cls._validation_max_retries = 3 cls._validation_retries = {} + cls._live_nodes_retry_attempts = 3 + cls._live_nodes_retry_interval_seconds = 0 + cls._layout_nodes_cache = set(default_nodes) + cls._layout_nodes_loaded = True + cls._layout_nodes_load_success = True + cls._load_live_nodes = classmethod(lambda _cls: (set(default_nodes), set(default_nodes), True)) return cls @@ -299,12 +306,15 @@ def test_skip_vmss_calls_skip_validation(self, recycler, status_client, action_c patch.object(recycler, "validate") as mock_validate, \ patch.object(recycler, "skip_validation") as mock_skip: - # operate is called 4 times (2 VMSS x 2 states) - # vmss-gpu-1 DEALLOCATED_UA -> returns gpu_vms - # vmss-gpu-1 DEALLOCATED_PLATFORM -> returns [] - # vmss-cpu-1 DEALLOCATED_UA -> returns cpu_vms - # vmss-cpu-1 DEALLOCATED_PLATFORM -> returns [] - mock_operate.side_effect = [gpu_vms, [], cpu_vms, []] + def operate_side_effect(vmss_id, **kwargs): + from_state = kwargs.get("from_state") + if vmss_id == "vmss-gpu-1" and from_state == "deallocated_ua": + return gpu_vms + if vmss_id == "vmss-cpu-1" and from_state == "deallocated_ua": + return cpu_vms + return [] + + mock_operate.side_effect = operate_side_effect recycler.start_and_validate_pipeline(status_client, action_client) @@ -330,7 +340,13 @@ def test_no_skip_vmss_all_go_to_validate(self, recycler, status_client, action_c patch.object(recycler, "validate") as mock_validate, \ patch.object(recycler, "skip_validation") as mock_skip: - mock_operate.side_effect = [vms, [], [], []] + def operate_side_effect(vmss_id, **kwargs): + from_state = kwargs.get("from_state") + if vmss_id == "vmss-gpu-1" and from_state == "deallocated_ua": + return vms + return [] + + mock_operate.side_effect = operate_side_effect recycler.start_and_validate_pipeline(status_client, action_client) mock_skip.assert_not_called() @@ -360,9 +376,13 @@ def test_skip_vmss_with_platform_state(self, recycler, status_client, action_cli patch.object(recycler, "validate") as mock_validate, \ patch.object(recycler, "skip_validation") as mock_skip: - # vmss-gpu-1: both return empty - # vmss-cpu-1: DEALLOCATED_UA returns [], DEALLOCATED_PLATFORM returns cpu_vms - mock_operate.side_effect = [[], [], [], cpu_vms] + def operate_side_effect(vmss_id, **kwargs): + from_state = kwargs.get("from_state") + if vmss_id == "vmss-cpu-1" and from_state == "deallocated_platform": + return cpu_vms + return [] + + mock_operate.side_effect = operate_side_effect recycler.start_and_validate_pipeline(status_client, action_client) @@ -432,3 +452,120 @@ def test_creates_ticket_for_new_node(self, recycler, status_client, action_clien # Should create exactly one incident mock_icm_api.create_incident.assert_called_once() + + +class TestNodeFilterPolicy: + def test_layout_only_loads_once(self, recycler): + recycler._layout_nodes_loaded = False + recycler._layout_nodes_load_success = False + recycler._layout_nodes_cache = set() + + with patch.object(recycler, "_load_layout_nodes", return_value=({"node-a"}, True)) as mock_load_layout: + filtered_1, ok_1 = recycler._filter_nodes_by_policy( + ["node-a"], + stage="test-stage-1", + require_layout=True, + require_live=False, + ) + filtered_2, ok_2 = recycler._filter_nodes_by_policy( + ["node-a"], + stage="test-stage-2", + require_layout=True, + require_live=False, + ) + + assert ok_1 is True + assert ok_2 is True + assert filtered_1 == ["node-a"] + assert filtered_2 == ["node-a"] + assert mock_load_layout.call_count == 1 + + def test_layout_load_failure_falls_back_to_original_behavior(self, recycler): + recycler._layout_nodes_loaded = False + recycler._layout_nodes_load_success = False + recycler._layout_nodes_cache = set() + + with patch.object(recycler, "_load_layout_nodes", return_value=(set(), False)): + filtered, ok = recycler._filter_nodes_by_policy( + ["node-a", "node-b"], + stage="test-layout-fail-fallback", + require_layout=True, + require_live=False, + ) + + assert ok is True + assert filtered == ["node-a", "node-b"] + + def test_empty_layout_falls_back_to_original_behavior(self, recycler): + recycler._layout_nodes_loaded = True + recycler._layout_nodes_load_success = True + recycler._layout_nodes_cache = set() + + filtered, ok = recycler._filter_nodes_by_policy( + ["node-a", "node-b"], + stage="test-layout-empty-fallback", + require_layout=True, + require_live=False, + ) + + assert ok is True + assert filtered == ["node-a", "node-b"] + + def test_no_layout_load_when_not_required(self, recycler): + with patch.object(recycler, "_load_layout_nodes") as mock_load_layout: + filtered, ok = recycler._filter_nodes_by_policy( + ["node-a"], + stage="test-no-layout", + require_layout=False, + require_live=False, + ) + + assert ok is True + assert filtered == ["node-a"] + mock_load_layout.assert_not_called() + + def test_live_nodes_retry_success(self, recycler): + recycler._live_nodes_retry_attempts = 3 + recycler._live_nodes_retry_interval_seconds = 0 + + with patch.object( + recycler, + "_load_live_nodes", + side_effect=[ + (set(), set(), False), + (set(), set(), False), + ({"node-a"}, {"node-a"}, True), + ], + ) as mock_load_live, patch("recycler.time.sleep") as mock_sleep: + filtered, ok = recycler._filter_nodes_by_policy( + ["node-a"], + stage="test-live-retry-success", + require_layout=True, + require_live=True, + ) + + assert ok is True + assert filtered == ["node-a"] + assert mock_load_live.call_count == 3 + assert mock_sleep.call_count == 2 + + def test_live_nodes_retry_exhausted_skip_stage(self, recycler): + recycler._live_nodes_retry_attempts = 3 + recycler._live_nodes_retry_interval_seconds = 0 + + with patch.object( + recycler, + "_load_live_nodes", + return_value=(set(), set(), False), + ) as mock_load_live, patch("recycler.time.sleep") as mock_sleep: + filtered, ok = recycler._filter_nodes_by_policy( + ["node-a"], + stage="test-live-retry-fail", + require_layout=True, + require_live=True, + ) + + assert ok is False + assert filtered == [] + assert mock_load_live.call_count == 3 + assert mock_sleep.call_count == 2 diff --git a/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py b/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py index 206e1a51..3fd00b02 100644 --- a/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py +++ b/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py @@ -132,6 +132,7 @@ def update_node_action(self, node: str, action: str, timestamp: str, # Check for existing record to avoid duplicates check_query = f""" {self.table_name} + | where Endpoint == '{self.endpoint}' | where HostName == '{node}' and Timestamp == datetime('{timestamp}') and Action == '{action}' | count """ @@ -159,6 +160,7 @@ def get_node_actions(self, node: str, start_time: str, try: query = f""" {self.table_name} + | where Endpoint == '{self.endpoint}' | where HostName == '{node}' | where Timestamp between (datetime({start_time}) .. datetime({end_time})) | order by Timestamp desc @@ -173,6 +175,7 @@ def get_latest_node_action(self, node: str) -> Optional[NodeAction]: try: query = f""" {self.table_name} + | where Endpoint == '{self.endpoint}' | where HostName == '{node}' | top 1 by Timestamp desc """ @@ -224,6 +227,7 @@ def get_latest_action_by_state( try: query = f""" {self.table_name} + | where Endpoint == '{self.endpoint}' | where Action endswith '{state}' | where HostName == '{hostname}' and NodeId == '{node_id}' | top 1 by Timestamp desc diff --git a/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py b/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py index c78744ad..f70d15e5 100644 --- a/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py +++ b/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py @@ -143,7 +143,11 @@ def get_node_status(self, timestamp_str = None if timestamp is not None: timestamp_str = convert_timestamp(timestamp, format="str") - query = f"{self.table_name} | where HostName == '{hostname}'" + query = ( + f"{self.table_name}" + f" | where HostName == '{hostname}'" + f" | where Endpoint == '{self.endpoint}'" + ) if timestamp_str is not None: query += f" | where Timestamp <= datetime({timestamp_str})" query += " | summarize arg_max(Timestamp, *) by HostName" @@ -187,16 +191,16 @@ def get_nodes_by_status( status: str, as_of_time: Optional[datetime] = None) -> List[NodeStatusRecord]: """Get all nodes whose latest/current status is exactly the specified status. - + Args: status (str): The status to filter nodes by as_of_time (datetime, optional): The reference time to check status. If not provided, uses current time. - + Returns: List[NodeStatusRecord]: List of node records whose latest status matches. Each record contains Timestamp, HostName, Status, NodeId, and Endpoint. - + Example: >>> client = NodeStatusClient() >>> current_cordoned_nodes = client.get_nodes_by_status('cordoned') @@ -209,18 +213,11 @@ def get_nodes_by_status( timestamp_condition = f"| where Timestamp <= datetime({timestamp_str})" query = f""" - let latest_status = {self.table_name} + {self.table_name} + | where Endpoint == '{self.endpoint}' {timestamp_condition} - | summarize arg_max(Timestamp, Status) by HostName; - latest_status + | summarize arg_max(Timestamp, *) by HostName | where Status == '{status}' - | join kind=inner ( - {self.table_name} - | where Status == '{status}' - | where Endpoint == '{self.endpoint}' - | summarize arg_max(Timestamp, *) by HostName - ) on HostName - | project Timestamp=Timestamp1, HostName, Status=Status1, NodeId, Endpoint """ results = self.execute_query(query) diff --git a/src/kusto-sdk/tests/test_node_status_client.py b/src/kusto-sdk/tests/test_node_status_client.py index 39756cc3..336fcf47 100644 --- a/src/kusto-sdk/tests/test_node_status_client.py +++ b/src/kusto-sdk/tests/test_node_status_client.py @@ -153,6 +153,7 @@ def test_get_node_status_existing(self, client, mock_kusto_client): query = mock_kusto_client.execute_command.call_args[0][0] assert TEST_STATUS_TABLE in query assert test_node in query + assert f"Endpoint == '{TEST_ENDPOINT}'" in query # Verify result assert isinstance(result, NodeStatusRecord) From 9e08edc2bc2ecdc7ab36a369f27c8fda061f3885 Mon Sep 17 00:00:00 2001 From: Rui Gao Date: Fri, 21 Aug 2026 07:41:02 +0000 Subject: [PATCH 2/6] Address PR feedback: fail-fast layout init and endpoint scoping --- .../src/node-recycler/recycler.py | 41 +++++++++------- .../src/node-recycler/tests/test_recycler.py | 48 +++++++++++-------- .../features/node_action/client.py | 4 ++ .../tests/test_node_action_client.py | 18 +++++++ 4 files changed, 75 insertions(+), 36 deletions(-) diff --git a/src/alert-manager/src/node-recycler/recycler.py b/src/alert-manager/src/node-recycler/recycler.py index 47475c69..804132b6 100644 --- a/src/alert-manager/src/node-recycler/recycler.py +++ b/src/alert-manager/src/node-recycler/recycler.py @@ -58,18 +58,33 @@ class NodeRecycler: @classmethod def _initialize_layout_nodes(cls) -> bool: - """Load layout nodes once and cache the result for process lifetime.""" + """Load and validate layout nodes once for process lifetime.""" if cls._layout_nodes_loaded: + if not cls._layout_nodes_load_success: + raise RuntimeError( + f"layout initialization previously failed from {cls._layout_config_path}" + ) + if not cls._layout_nodes_cache: + raise RuntimeError( + f"layout cache is empty from {cls._layout_config_path}" + ) return cls._layout_nodes_load_success layout_nodes, layout_ok = cls._load_layout_nodes() + if not layout_ok: + raise RuntimeError( + f"failed to initialize layout nodes from {cls._layout_config_path}" + ) + if not layout_nodes: + raise RuntimeError( + f"layout contains no usable nodes at {cls._layout_config_path}" + ) + cls._layout_nodes_cache = layout_nodes - cls._layout_nodes_load_success = layout_ok + cls._layout_nodes_load_success = True cls._layout_nodes_loaded = True - - if not layout_ok: - logger.error("Failed to initialize layout nodes from %s", cls._layout_config_path) - return cls._layout_nodes_load_success + logger.info("Initialized %d layout nodes from %s", len(layout_nodes), cls._layout_config_path) + return True @classmethod def _load_layout_nodes(cls) -> tuple[set[str], bool]: @@ -175,17 +190,9 @@ def _filter_nodes_by_policy( return [], True layout_nodes = set() - layout_filter_enabled = require_layout if require_layout: - layout_ok = cls._initialize_layout_nodes() - if not layout_ok: - logger.warning("[%s] layout nodes unavailable, skipping layout filter and continuing", stage) - layout_filter_enabled = False - else: - layout_nodes = cls._layout_nodes_cache - if not layout_nodes: - logger.warning("[%s] layout nodes is empty, skipping layout filter and continuing", stage) - layout_filter_enabled = False + cls._initialize_layout_nodes() + layout_nodes = cls._layout_nodes_cache live_nodes, ready_nodes, live_ok = set(), set(), True if require_live or require_ready: @@ -198,7 +205,7 @@ def _filter_nodes_by_policy( skipped = [] for hostname in normalized: reasons = [] - if layout_filter_enabled and hostname not in layout_nodes: + if require_layout and hostname not in layout_nodes: reasons.append("not_in_layout") if require_live and hostname not in live_nodes: reasons.append("not_live") diff --git a/src/alert-manager/src/node-recycler/tests/test_recycler.py b/src/alert-manager/src/node-recycler/tests/test_recycler.py index 3d3607bb..b623c76c 100644 --- a/src/alert-manager/src/node-recycler/tests/test_recycler.py +++ b/src/alert-manager/src/node-recycler/tests/test_recycler.py @@ -480,36 +480,46 @@ def test_layout_only_loads_once(self, recycler): assert filtered_2 == ["node-a"] assert mock_load_layout.call_count == 1 - def test_layout_load_failure_falls_back_to_original_behavior(self, recycler): + def test_layout_load_failure_raises_runtime_error(self, recycler): recycler._layout_nodes_loaded = False recycler._layout_nodes_load_success = False recycler._layout_nodes_cache = set() with patch.object(recycler, "_load_layout_nodes", return_value=(set(), False)): - filtered, ok = recycler._filter_nodes_by_policy( - ["node-a", "node-b"], - stage="test-layout-fail-fallback", - require_layout=True, - require_live=False, - ) + with pytest.raises(RuntimeError, match="failed to initialize layout nodes"): + recycler._filter_nodes_by_policy( + ["node-a", "node-b"], + stage="test-layout-fail", + require_layout=True, + require_live=False, + ) + + def test_empty_layout_raises_runtime_error(self, recycler): + recycler._layout_nodes_loaded = False + recycler._layout_nodes_load_success = False + recycler._layout_nodes_cache = set() - assert ok is True - assert filtered == ["node-a", "node-b"] + with patch.object(recycler, "_load_layout_nodes", return_value=(set(), True)): + with pytest.raises(RuntimeError, match="layout contains no usable nodes"): + recycler._filter_nodes_by_policy( + ["node-a", "node-b"], + stage="test-layout-empty", + require_layout=True, + require_live=False, + ) - def test_empty_layout_falls_back_to_original_behavior(self, recycler): + def test_empty_layout_cache_after_load_raises_runtime_error(self, recycler): recycler._layout_nodes_loaded = True recycler._layout_nodes_load_success = True recycler._layout_nodes_cache = set() - filtered, ok = recycler._filter_nodes_by_policy( - ["node-a", "node-b"], - stage="test-layout-empty-fallback", - require_layout=True, - require_live=False, - ) - - assert ok is True - assert filtered == ["node-a", "node-b"] + with pytest.raises(RuntimeError, match="layout cache is empty"): + recycler._filter_nodes_by_policy( + ["node-a"], + stage="test-layout-empty-cache", + require_layout=True, + require_live=False, + ) def test_no_layout_load_when_not_required(self, recycler): with patch.object(recycler, "_load_layout_nodes") as mock_load_layout: diff --git a/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py b/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py index 3fd00b02..f66fab2f 100644 --- a/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py +++ b/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py @@ -255,6 +255,7 @@ def find_triaged_failure(self, node_name: str, completed_time_ms: int, launched_ List of NodeAction records for triaged actions, or empty list if none found """ try: + endpoint = self.endpoint.replace("'", "''") # Get node actions in time range start_time = convert_timestamp(launched_time_ms / 1000, "str") end_time = convert_timestamp(completed_time_ms / 1000, "str") @@ -275,16 +276,19 @@ def find_triaged_failure(self, node_name: str, completed_time_ms: int, launched_ # Query triaged actions using KQL query = f""" let node_name = '{node_name}'; + let endpoint = '{endpoint}'; let cordoned_ts = datetime({cordoned_timestamp}); let next_available_ts = toscalar( {self.table_name} | where HostName == node_name + | where Endpoint == endpoint | where Action endswith '-available' and Action != 'available-cordoned' | where Timestamp > cordoned_ts | summarize min(Timestamp) ); {self.table_name} | where HostName == node_name + | where Endpoint == endpoint | where Timestamp >= cordoned_ts | where isnull(next_available_ts) or Timestamp <= next_available_ts | where Action in ('cordoned-triaged_platform', 'cordoned-triaged_hardware', 'cordoned-triaged_user', 'cordoned-triaged_unknown') diff --git a/src/kusto-sdk/tests/test_node_action_client.py b/src/kusto-sdk/tests/test_node_action_client.py index 5629bb8c..2a286309 100644 --- a/src/kusto-sdk/tests/test_node_action_client.py +++ b/src/kusto-sdk/tests/test_node_action_client.py @@ -233,3 +233,21 @@ def test_validate_action_format(self, client): assert not client.is_valid_action("InvalidAction") assert not client.is_valid_action("Cordoned") assert not client.is_valid_action("available") + + def test_find_triaged_failure_scopes_endpoint(self, client, mock_kusto_client): + """find_triaged_failure should scope all table scans by endpoint.""" + action = MagicMock() + action.Action = "available-cordoned" + action.Timestamp = "2026-01-01T00:00:00Z" + + with patch.object(client, "get_node_actions", return_value=[action]): + mock_kusto_client.execute_command.return_value = [] + client.find_triaged_failure( + node_name="test-node", + completed_time_ms=1704067500000, + launched_time_ms=1704067200000, + ) + + query = mock_kusto_client.execute_command.call_args[0][0] + assert "let endpoint = 'test-wcu'" in query + assert query.count("| where Endpoint == endpoint") == 2 From 248588f286069fda5fde65b0016f3e9b1a5b8713 Mon Sep 17 00:00:00 2001 From: Rui Gao Date: Fri, 21 Aug 2026 09:08:26 +0000 Subject: [PATCH 3/6] Refactor endpoint filtering to explicit optional parameters --- .../src/alert-parser/node_updater.py | 26 +++++- .../src/job-data-recorder/kusto_util.py | 3 +- .../node_recorder_helper.py | 26 +++++- .../src/node-recycler/recycler.py | 37 +++++++- .../src/node-recycler/tests/test_recycler.py | 3 + .../features/node_action/client.py | 71 +++++++++++---- .../features/node_status/client.py | 24 +++-- .../tests/test_node_action_client.py | 91 ++++++++++++++++++- .../tests/test_node_status_client.py | 26 +++++- 9 files changed, 265 insertions(+), 42 deletions(-) diff --git a/src/alert-manager/src/alert-parser/node_updater.py b/src/alert-manager/src/alert-parser/node_updater.py index a6272a94..6cc9dec0 100644 --- a/src/alert-manager/src/alert-parser/node_updater.py +++ b/src/alert-manager/src/alert-parser/node_updater.py @@ -39,7 +39,11 @@ def __init__(self): self.retries = 3 def get_node_latest_status(self, node, as_of_time=None): - node_status = self.node_status_client.get_node_status(node, as_of_time) + node_status = self.node_status_client.get_node_status( + node, + as_of_time, + endpoint=self.endpoint, + ) if not node_status: logger.info(f"No status found for node {node} as of {as_of_time}") return None @@ -47,11 +51,15 @@ def get_node_latest_status(self, node, as_of_time=None): return node_status def get_nodes_by_status(self, status, as_of_time=None): - nodes = self.node_status_client.get_nodes_by_status(status, as_of_time) + nodes = self.node_status_client.get_nodes_by_status( + status, + as_of_time, + endpoint=self.endpoint, + ) return nodes def get_last_actions_update_time(self): - time = self.node_action_client.get_last_update_time() + time = self.node_action_client.get_last_update_time(endpoint=self.endpoint) return time @@ -76,7 +84,15 @@ def update_status_action(self, node, from_status, to_status, timestamp, reason, action = self.node_status_client.get_transition_action(from_status, to_status) for i in range(self.retries): try: - self.node_action_client.update_node_action(node, action, timestamp, reason, detail, category='') + self.node_action_client.update_node_action( + node, + action, + timestamp, + reason, + detail, + category='', + endpoint=self.endpoint, + ) logger.info(f"Updated node action to {action} for node {node} on {timestamp}") status_updated = True break @@ -90,4 +106,4 @@ def update_status_action(self, node, from_status, to_status, timestamp, reason, return False logger.info(f"Successfully updated node status and action for node {node} on {timestamp}") return True - return False \ No newline at end of file + return False diff --git a/src/alert-manager/src/job-data-recorder/kusto_util.py b/src/alert-manager/src/job-data-recorder/kusto_util.py index d4de3dc5..9d683d62 100644 --- a/src/alert-manager/src/job-data-recorder/kusto_util.py +++ b/src/alert-manager/src/job-data-recorder/kusto_util.py @@ -72,7 +72,8 @@ def find_node_triaged_failure_in_kusto(self, node_name, completedTime, triaged_actions = self.node_action_client.find_triaged_failure( node_name=node_name, completed_time_ms=completedTime, - launched_time_ms=launchedTime + launched_time_ms=launchedTime, + endpoint=self.endpoint, ) if not triaged_actions: diff --git a/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py b/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py index afb59527..ef910142 100644 --- a/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py +++ b/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py @@ -37,15 +37,25 @@ def __init__(self): self.retries = 3 def get_node_latest_status(self, node): - node_status = self.node_status_client.get_node_status(node) + node_status = self.node_status_client.get_node_status( + node, + endpoint=self.endpoint, + ) return node_status def get_nodes_by_status(self, status, as_of_time=None): - nodes = self.node_status_client.get_nodes_by_status(status, as_of_time) + nodes = self.node_status_client.get_nodes_by_status( + status, + as_of_time, + endpoint=self.endpoint, + ) return nodes def get_node_latest_action(self, node): - node_action = self.node_action_client.get_latest_node_action(node) + node_action = self.node_action_client.get_latest_node_action( + node, + endpoint=self.endpoint, + ) return node_action def update_status_action(self, node, from_status, to_status, timestamp, reason, detail, category=''): @@ -58,7 +68,15 @@ def update_status_action(self, node, from_status, to_status, timestamp, reason, action = self.node_status_client.get_transition_action(from_status, to_status) for i in range(self.retries): try: - self.node_action_client.update_node_action(node, action, timestamp, reason, detail, category=category) + self.node_action_client.update_node_action( + node, + action, + timestamp, + reason, + detail, + category=category, + endpoint=self.endpoint, + ) logger.info(f"Updated node action to {action} for node {node} on {timestamp} with category {category}") status_updated = True break diff --git a/src/alert-manager/src/node-recycler/recycler.py b/src/alert-manager/src/node-recycler/recycler.py index 804132b6..f26607ab 100644 --- a/src/alert-manager/src/node-recycler/recycler.py +++ b/src/alert-manager/src/node-recycler/recycler.py @@ -50,6 +50,7 @@ class NodeRecycler: _ltp_validation_image = os.getenv("LTP_VALIDATION_IMAGE") _ltp_vmss_ids = os.getenv("LTP_VMSS_IDS", "") _validation_skip_vmss_ids = set(filter(None, os.getenv("VALIDATION_SKIP_VMSS_IDS", "").split(","))) + _cluster_id = os.getenv("CLUSTER_ID") _live_nodes_retry_attempts = int(os.getenv("LIVE_NODES_RETRY_ATTEMPTS", "3")) _live_nodes_retry_interval_seconds = int(os.getenv("LIVE_NODES_RETRY_INTERVAL_SECONDS", "1")) _layout_nodes_cache = set() @@ -249,7 +250,10 @@ def ofr(cls, node_faults=None, status_client=None, action_client=None): created = [] if not node_faults and status_client and action_client: node_faults = [] - status_nodes = status_client.get_nodes_by_status(from_state) + status_nodes = status_client.get_nodes_by_status( + from_state, + endpoint=cls._cluster_id, + ) candidates = [n.HostName for n in status_nodes] filtered_candidates, ok = cls._filter_nodes_by_policy( candidates, @@ -276,13 +280,21 @@ def ofr(cls, node_faults=None, status_client=None, action_client=None): # action for this node. get_latest_action_by_state uses # "endswith" which never returns triaged_hardware-ua, so we # query the latest action separately to detect prior OFR. - latest = action_client.get_latest_node_action(hostname) + latest = action_client.get_latest_node_action( + hostname, + endpoint=cls._cluster_id, + ) if latest and latest.Action == f"{from_state}-{to_state}": logger.info(f"OFR already submitted for {hostname}, ticket_id={latest.Detail}") created.append({"hostname": hostname, "node_id": node_id, "ticket_id": latest.Detail}) continue - result = action_client.get_latest_action_by_state(hostname, node_id, from_state) + result = action_client.get_latest_action_by_state( + hostname, + node_id, + from_state, + endpoint=cls._cluster_id, + ) if result and result.Action and result.Detail: action, detail = result.Action, result.Detail if action.endswith(from_state): @@ -342,6 +354,7 @@ def ofr(cls, node_faults=None, status_client=None, action_client=None): action_client.update_node_action( hostname, f"{from_state}-{to_state}", time.time(), ofr_str, ticket_id, "", + endpoint=cls._cluster_id, ) except Exception as e: logger.error(f"Error occured when creating OFR ticket for node {node_id} ({hostname}): {e}") @@ -426,7 +439,12 @@ def operate(cls, vmss_id, operation="start", hostnames=None, status_client=None, raise ValueError(f"Unsupported operation: {operation}") if not hostnames and status_client: - hostnames = [n.HostName for n in status_client.get_nodes_by_status(from_state)] + hostnames = [ + n.HostName for n in status_client.get_nodes_by_status( + from_state, + endpoint=cls._cluster_id, + ) + ] hostnames, ok = cls._filter_nodes_by_policy( hostnames or [], @@ -488,6 +506,7 @@ def is_vm_succeed_in_target(vm): action_client.update_node_action( name, f"{from_state}-{to_state}", time.time(), f"{op.title()}ing VM", "", "", + endpoint=cls._cluster_id, ) except Exception as e: logger.error(f"List instances in VMSS failed due to: {e}") @@ -534,7 +553,12 @@ def validate(cls, hostnames=None, filter_state='', status_client=None, action_cl if not filter_state: filter_state = NodeStatus.ALLOCATED_UA.value if not hostnames and status_client: - hostnames = [n.HostName for n in status_client.get_nodes_by_status(filter_state)] + hostnames = [ + n.HostName for n in status_client.get_nodes_by_status( + filter_state, + endpoint=cls._cluster_id, + ) + ] hostnames, ok = cls._filter_nodes_by_policy( hostnames or [], @@ -578,6 +602,7 @@ def validate(cls, hostnames=None, filter_state='', status_client=None, action_cl hostname, f"{filter_state}-{NodeStatus.VALIDATING.value}", time.time(), "Submitting validation job for VM", "", "", + endpoint=cls._cluster_id, ) res.raise_for_status() logger.info(f"Submitted validation job for {hostname} with response: {res.json()}") @@ -603,6 +628,7 @@ def validate(cls, hostnames=None, filter_state='', status_client=None, action_cl action_client.update_node_action( hostname, f"{filter_state}-{NodeStatus.CORDONED.value}", time.time(), f"Validation job submission failed after {cls._validation_retries[hostname]} attempts", str(e), "", + endpoint=cls._cluster_id, ) del cls._validation_retries[hostname] @@ -663,6 +689,7 @@ def skip_validation(cls, hostnames, filter_state, status_client=None, action_cli action_client.update_node_action( hostname, f"{filter_state}-{NodeStatus.AVAILABLE.value}", time.time(), "Skipping GPU validation per VMSS config", "", "", + endpoint=cls._cluster_id, ) @classmethod diff --git a/src/alert-manager/src/node-recycler/tests/test_recycler.py b/src/alert-manager/src/node-recycler/tests/test_recycler.py index b623c76c..a89f68cd 100644 --- a/src/alert-manager/src/node-recycler/tests/test_recycler.py +++ b/src/alert-manager/src/node-recycler/tests/test_recycler.py @@ -56,6 +56,7 @@ def mock_env(monkeypatch): monkeypatch.setenv("LTP_VMSS_IDS", "vmss-gpu-1,vmss-cpu-1") monkeypatch.setenv("VALIDATION_SKIP_VMSS_IDS", "vmss-cpu-1") monkeypatch.setenv("VALIDATION_MAX_RETRIES", "3") + monkeypatch.setenv("CLUSTER_ID", "test-endpoint") @pytest.fixture @@ -79,6 +80,7 @@ def recycler(mock_env): cls._azure_client_id = "test-client-id" cls._ltp_vmss_ids = "vmss-gpu-1,vmss-cpu-1" cls._validation_skip_vmss_ids = {"vmss-cpu-1"} + cls._cluster_id = "test-endpoint" cls._validation_max_retries = 3 cls._validation_retries = {} cls._live_nodes_retry_attempts = 3 @@ -137,6 +139,7 @@ def test_records_action(self, recycler, status_client, action_client): "node-a", "allocated_ua-available", pytest.approx(time.time(), abs=5), "Skipping GPU validation per VMSS config", "", "", + endpoint="test-endpoint", ) def test_log_contains_uncordon_hint(self, recycler, status_client, action_client, caplog): diff --git a/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py b/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py index f66fab2f..a378f0a0 100644 --- a/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py +++ b/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py @@ -103,7 +103,8 @@ def create_attribute_table(self) -> None: raise RuntimeError(f"Failed to create attribute table: {str(e)}") def update_node_action(self, node: str, action: str, timestamp: str, - reason: str, detail: str, category: str) -> None: + reason: str, detail: str, category: str, + endpoint: Optional[str] = None) -> None: """ Updates or inserts a node action record in the Kusto table. @@ -130,9 +131,13 @@ def update_node_action(self, node: str, action: str, timestamp: str, node_id = Node(node).get_vm_node_id_by_hostname(timestamp) # Check for existing record to avoid duplicates + endpoint_condition = "" + if endpoint: + escaped_endpoint = str(endpoint).replace("'", "''") + endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" check_query = f""" {self.table_name} - | where Endpoint == '{self.endpoint}' + {endpoint_condition} | where HostName == '{node}' and Timestamp == datetime('{timestamp}') and Action == '{action}' | count """ @@ -155,12 +160,17 @@ def update_node_action(self, node: str, action: str, timestamp: str, raise RuntimeError(f"Failed to update node action: {str(e)}") def get_node_actions(self, node: str, start_time: str, - end_time: str) -> List[NodeAction]: + end_time: str, + endpoint: Optional[str] = None) -> List[NodeAction]: """Get action history for a node in a time range""" try: + endpoint_condition = "" + if endpoint: + escaped_endpoint = str(endpoint).replace("'", "''") + endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" {self.table_name} - | where Endpoint == '{self.endpoint}' + {endpoint_condition} | where HostName == '{node}' | where Timestamp between (datetime({start_time}) .. datetime({end_time})) | order by Timestamp desc @@ -170,12 +180,18 @@ def get_node_actions(self, node: str, start_time: str, except Exception as e: raise RuntimeError(f"Failed to get node actions: {str(e)}") - def get_latest_node_action(self, node: str) -> Optional[NodeAction]: + def get_latest_node_action(self, + node: str, + endpoint: Optional[str] = None) -> Optional[NodeAction]: """Get the most recent action for a node""" try: + endpoint_condition = "" + if endpoint: + escaped_endpoint = str(endpoint).replace("'", "''") + endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" {self.table_name} - | where Endpoint == '{self.endpoint}' + {endpoint_condition} | where HostName == '{node}' | top 1 by Timestamp desc """ @@ -184,12 +200,17 @@ def get_latest_node_action(self, node: str) -> Optional[NodeAction]: except Exception as e: raise RuntimeError(f"Failed to get latest node action: {str(e)}") - def get_last_update_time(self) -> Optional[datetime]: + def get_last_update_time(self, + endpoint: Optional[str] = None) -> Optional[datetime]: """Get the last update time for the node action table""" try: + endpoint_condition = "" + if endpoint: + escaped_endpoint = str(endpoint).replace("'", "''") + endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" {self.table_name} - | where Endpoint == '{self.endpoint}' + {endpoint_condition} | summarize arg_max(Timestamp, *) by HostName | top 1 by Timestamp desc """ @@ -202,7 +223,8 @@ def get_latest_action_by_state( self, hostname: str, node_id: str, - state: str + state: str, + endpoint: Optional[str] = None ) -> Optional[Dict[str, Any]]: """ Get the latest action that ends with the specified state for a given hostname and node_id. @@ -225,9 +247,13 @@ def get_latest_action_by_state( ... print(f"Action: {result['Action']}, Detail: {result['Detail']}") """ try: + endpoint_condition = "" + if endpoint: + escaped_endpoint = str(endpoint).replace("'", "''") + endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" {self.table_name} - | where Endpoint == '{self.endpoint}' + {endpoint_condition} | where Action endswith '{state}' | where HostName == '{hostname}' and NodeId == '{node_id}' | top 1 by Timestamp desc @@ -240,7 +266,11 @@ def get_latest_action_by_state( except Exception as e: raise RuntimeError(f"Failed to get latest action by state: {str(e)}") - def find_triaged_failure(self, node_name: str, completed_time_ms: int, launched_time_ms: int) -> List[NodeAction]: + def find_triaged_failure(self, + node_name: str, + completed_time_ms: int, + launched_time_ms: int, + endpoint: Optional[str] = None) -> List[NodeAction]: """ Find triaged actions for a node between job launch and completion. @@ -255,12 +285,14 @@ def find_triaged_failure(self, node_name: str, completed_time_ms: int, launched_ List of NodeAction records for triaged actions, or empty list if none found """ try: - endpoint = self.endpoint.replace("'", "''") # Get node actions in time range start_time = convert_timestamp(launched_time_ms / 1000, "str") end_time = convert_timestamp(completed_time_ms / 1000, "str") - node_actions = self.get_node_actions(node_name, start_time, end_time) + node_actions = self.get_node_actions(node_name, + start_time, + end_time, + endpoint=endpoint) # Check if there's available-cordoned action cordoned_timestamp = None @@ -274,21 +306,28 @@ def find_triaged_failure(self, node_name: str, completed_time_ms: int, launched_ return [] # Query triaged actions using KQL + endpoint_declare = "" + endpoint_condition = "" + if endpoint: + escaped_endpoint = str(endpoint).replace("'", "''") + endpoint_declare = f"let endpoint = '{escaped_endpoint}';" + endpoint_condition = "| where Endpoint == endpoint" + query = f""" let node_name = '{node_name}'; - let endpoint = '{endpoint}'; + {endpoint_declare} let cordoned_ts = datetime({cordoned_timestamp}); let next_available_ts = toscalar( {self.table_name} | where HostName == node_name - | where Endpoint == endpoint + {endpoint_condition} | where Action endswith '-available' and Action != 'available-cordoned' | where Timestamp > cordoned_ts | summarize min(Timestamp) ); {self.table_name} | where HostName == node_name - | where Endpoint == endpoint + {endpoint_condition} | where Timestamp >= cordoned_ts | where isnull(next_available_ts) or Timestamp <= next_available_ts | where Action in ('cordoned-triaged_platform', 'cordoned-triaged_hardware', 'cordoned-triaged_user', 'cordoned-triaged_unknown') diff --git a/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py b/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py index f70d15e5..2636748b 100644 --- a/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py +++ b/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py @@ -138,16 +138,16 @@ def get_status_group(self, status: str) -> str | None: def get_node_status(self, hostname: str, - timestamp: datetime = None) -> NodeStatusRecord: + timestamp: datetime = None, + endpoint: Optional[str] = None) -> Optional[NodeStatusRecord]: """Get node status at a specific time""" timestamp_str = None if timestamp is not None: timestamp_str = convert_timestamp(timestamp, format="str") - query = ( - f"{self.table_name}" - f" | where HostName == '{hostname}'" - f" | where Endpoint == '{self.endpoint}'" - ) + query = f"{self.table_name} | where HostName == '{hostname}'" + if endpoint: + escaped_endpoint = str(endpoint).replace("'", "''") + query += f" | where Endpoint == '{escaped_endpoint}'" if timestamp_str is not None: query += f" | where Timestamp <= datetime({timestamp_str})" query += " | summarize arg_max(Timestamp, *) by HostName" @@ -162,7 +162,8 @@ def update_node_status(self, hostname: str, to_status: str, timestamp: datetime | str | int) -> str: """Update node status""" timestamp = convert_timestamp(timestamp, format="datetime") - current_record = self.get_node_status(hostname, timestamp) + current_record = self.get_node_status(hostname, timestamp, + endpoint=self.endpoint) # Validate status transition if current_record and not NodeStatus.can_transition( @@ -189,7 +190,8 @@ def update_node_status(self, hostname: str, to_status: str, def get_nodes_by_status( self, status: str, - as_of_time: Optional[datetime] = None) -> List[NodeStatusRecord]: + as_of_time: Optional[datetime] = None, + endpoint: Optional[str] = None) -> List[NodeStatusRecord]: """Get all nodes whose latest/current status is exactly the specified status. Args: @@ -208,13 +210,17 @@ def get_nodes_by_status( """ try: timestamp_condition = "" + endpoint_condition = "" if as_of_time: timestamp_str = convert_timestamp(as_of_time, format="str") timestamp_condition = f"| where Timestamp <= datetime({timestamp_str})" + if endpoint: + escaped_endpoint = str(endpoint).replace("'", "''") + endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" {self.table_name} - | where Endpoint == '{self.endpoint}' + {endpoint_condition} {timestamp_condition} | summarize arg_max(Timestamp, *) by HostName | where Status == '{status}' diff --git a/src/kusto-sdk/tests/test_node_action_client.py b/src/kusto-sdk/tests/test_node_action_client.py index 2a286309..bb64775e 100644 --- a/src/kusto-sdk/tests/test_node_action_client.py +++ b/src/kusto-sdk/tests/test_node_action_client.py @@ -119,6 +119,39 @@ def test_update_node_action(self, client, mock_kusto_client, mock_node): "Test details", "Test", TEST_ENDPOINT ]) + def test_update_node_action_endpoint_filter_is_optional(self, client, + mock_kusto_client, + mock_node): + """update_node_action duplicate-check should scope endpoint only when provided.""" + timestamp = datetime.utcnow().isoformat() + mock_kusto_client.execute_command.reset_mock() + mock_kusto_client.execute_command.return_value = [{"Count": 0}] + + client.update_node_action( + node="test-node", + action="available-cordoned", + timestamp=timestamp, + reason="r", + detail="d", + category="c", + ) + first_query = mock_kusto_client.execute_command.call_args_list[0][0][0] + assert "Endpoint ==" not in first_query + + mock_kusto_client.execute_command.reset_mock() + mock_kusto_client.execute_command.return_value = [{"Count": 0}] + client.update_node_action( + node="test-node", + action="available-cordoned", + timestamp=timestamp, + reason="r", + detail="d", + category="c", + endpoint=TEST_ENDPOINT, + ) + second_query = mock_kusto_client.execute_command.call_args_list[0][0][0] + assert f"Endpoint == '{TEST_ENDPOINT}'" in second_query + def test_update_node_action_invalid_action(self, client): """Test update_node_action method with invalid action""" with pytest.raises(RuntimeError) as exc_info: @@ -157,6 +190,26 @@ def test_get_latest_node_action(self, client, mock_kusto_client): assert action.Action == "available-cordoned" assert action.Endpoint == TEST_ENDPOINT + def test_get_latest_node_action_without_endpoint_filter(self, client, + mock_kusto_client): + """get_latest_node_action should not filter endpoint by default.""" + mock_kusto_client.execute_command.return_value = [] + + client.get_latest_node_action("test-node") + + query = mock_kusto_client.execute_command.call_args[0][0] + assert "Endpoint ==" not in query + + def test_get_latest_node_action_with_endpoint_filter(self, client, + mock_kusto_client): + """get_latest_node_action should filter endpoint when provided.""" + mock_kusto_client.execute_command.return_value = [] + + client.get_latest_node_action("test-node", endpoint=TEST_ENDPOINT) + + query = mock_kusto_client.execute_command.call_args[0][0] + assert f"Endpoint == '{TEST_ENDPOINT}'" in query + def test_get_latest_node_action_no_actions(self, client, mock_kusto_client): """Test get_latest_node_action method when no actions exist""" @@ -204,6 +257,23 @@ def test_get_node_actions(self, client, mock_kusto_client): assert actions[0].Action == "available-cordoned" assert actions[1].Action == "cordoned-triaged_hardware" + def test_get_node_actions_with_endpoint_filter(self, client, + mock_kusto_client): + """get_node_actions should filter endpoint when provided.""" + start_time = datetime.utcnow() - timedelta(hours=1) + end_time = datetime.utcnow() + mock_kusto_client.execute_command.return_value = [] + + client.get_node_actions( + node="test-node", + start_time=start_time.isoformat(), + end_time=end_time.isoformat(), + endpoint=TEST_ENDPOINT, + ) + + query = mock_kusto_client.execute_command.call_args[0][0] + assert f"Endpoint == '{TEST_ENDPOINT}'" in query + def test_get_node_actions_empty_result(self, client, mock_kusto_client): """Test get_node_actions method with no results""" mock_kusto_client.execute_command.return_value = [] @@ -235,7 +305,7 @@ def test_validate_action_format(self, client): assert not client.is_valid_action("available") def test_find_triaged_failure_scopes_endpoint(self, client, mock_kusto_client): - """find_triaged_failure should scope all table scans by endpoint.""" + """find_triaged_failure should scope endpoint only when provided.""" action = MagicMock() action.Action = "available-cordoned" action.Timestamp = "2026-01-01T00:00:00Z" @@ -246,8 +316,27 @@ def test_find_triaged_failure_scopes_endpoint(self, client, mock_kusto_client): node_name="test-node", completed_time_ms=1704067500000, launched_time_ms=1704067200000, + endpoint=TEST_ENDPOINT, ) query = mock_kusto_client.execute_command.call_args[0][0] assert "let endpoint = 'test-wcu'" in query assert query.count("| where Endpoint == endpoint") == 2 + + def test_find_triaged_failure_without_endpoint(self, client, mock_kusto_client): + """find_triaged_failure should not inject endpoint filter by default.""" + action = MagicMock() + action.Action = "available-cordoned" + action.Timestamp = "2026-01-01T00:00:00Z" + + with patch.object(client, "get_node_actions", return_value=[action]): + mock_kusto_client.execute_command.return_value = [] + client.find_triaged_failure( + node_name="test-node", + completed_time_ms=1704067500000, + launched_time_ms=1704067200000, + ) + + query = mock_kusto_client.execute_command.call_args[0][0] + assert "let endpoint =" not in query + assert "| where Endpoint == endpoint" not in query diff --git a/src/kusto-sdk/tests/test_node_status_client.py b/src/kusto-sdk/tests/test_node_status_client.py index 336fcf47..3a8eccd3 100644 --- a/src/kusto-sdk/tests/test_node_status_client.py +++ b/src/kusto-sdk/tests/test_node_status_client.py @@ -153,7 +153,7 @@ def test_get_node_status_existing(self, client, mock_kusto_client): query = mock_kusto_client.execute_command.call_args[0][0] assert TEST_STATUS_TABLE in query assert test_node in query - assert f"Endpoint == '{TEST_ENDPOINT}'" in query + assert "Endpoint ==" not in query # Verify result assert isinstance(result, NodeStatusRecord) @@ -161,6 +161,17 @@ def test_get_node_status_existing(self, client, mock_kusto_client): assert result.HostName == test_node assert result.Endpoint == TEST_ENDPOINT + def test_get_node_status_with_explicit_endpoint_filter(self, client, + mock_kusto_client): + """get_node_status should apply endpoint filter only when provided.""" + timestamp = datetime.utcnow() + mock_kusto_client.execute_command.return_value = [] + + client.get_node_status("test-node", timestamp.timestamp(), endpoint=TEST_ENDPOINT) + + query = mock_kusto_client.execute_command.call_args[0][0] + assert f"Endpoint == '{TEST_ENDPOINT}'" in query + def test_get_node_status_new(self, client, mock_kusto_client, mock_node): """Test get_node_status method for new node""" test_node = "test-node" @@ -235,3 +246,16 @@ def test_update_node_status_invalid_transition(self, client, with pytest.raises(ValueError) as exc_info: client.update_node_status(test_node, NodeStatus.TRIAGED_HARDWARE.value, timestamp.timestamp()) + + def test_get_nodes_by_status_with_optional_endpoint(self, client, + mock_kusto_client): + """get_nodes_by_status should only scope endpoint when provided.""" + mock_kusto_client.execute_command.return_value = [] + + client.get_nodes_by_status(NodeStatus.CORDONED.value) + first_query = mock_kusto_client.execute_command.call_args[0][0] + assert "Endpoint ==" not in first_query + + client.get_nodes_by_status(NodeStatus.CORDONED.value, endpoint=TEST_ENDPOINT) + second_query = mock_kusto_client.execute_command.call_args[0][0] + assert f"Endpoint == '{TEST_ENDPOINT}'" in second_query From 622c8df6eeceb0244775741b78b281d71a008786 Mon Sep 17 00:00:00 2001 From: Rui Gao Date: Fri, 21 Aug 2026 09:30:36 +0000 Subject: [PATCH 4/6] Use boolean endpoint scope for get queries only --- .../src/alert-parser/node_updater.py | 8 ++-- .../src/job-data-recorder/kusto_util.py | 2 +- .../node_recorder_helper.py | 7 ++-- .../src/node-recycler/recycler.py | 16 +++----- .../src/node-recycler/tests/test_recycler.py | 1 - .../features/node_action/client.py | 41 ++++++++----------- .../features/node_status/client.py | 14 +++---- .../tests/test_node_action_client.py | 30 ++++---------- .../tests/test_node_status_client.py | 11 ++++- 9 files changed, 55 insertions(+), 75 deletions(-) diff --git a/src/alert-manager/src/alert-parser/node_updater.py b/src/alert-manager/src/alert-parser/node_updater.py index 6cc9dec0..f7246f0a 100644 --- a/src/alert-manager/src/alert-parser/node_updater.py +++ b/src/alert-manager/src/alert-parser/node_updater.py @@ -42,7 +42,7 @@ def get_node_latest_status(self, node, as_of_time=None): node_status = self.node_status_client.get_node_status( node, as_of_time, - endpoint=self.endpoint, + use_current_endpoint=True, ) if not node_status: logger.info(f"No status found for node {node} as of {as_of_time}") @@ -54,12 +54,13 @@ def get_nodes_by_status(self, status, as_of_time=None): nodes = self.node_status_client.get_nodes_by_status( status, as_of_time, - endpoint=self.endpoint, + use_current_endpoint=True, ) return nodes def get_last_actions_update_time(self): - time = self.node_action_client.get_last_update_time(endpoint=self.endpoint) + time = self.node_action_client.get_last_update_time( + use_current_endpoint=True) return time @@ -91,7 +92,6 @@ def update_status_action(self, node, from_status, to_status, timestamp, reason, reason, detail, category='', - endpoint=self.endpoint, ) logger.info(f"Updated node action to {action} for node {node} on {timestamp}") status_updated = True diff --git a/src/alert-manager/src/job-data-recorder/kusto_util.py b/src/alert-manager/src/job-data-recorder/kusto_util.py index 9d683d62..26d0393f 100644 --- a/src/alert-manager/src/job-data-recorder/kusto_util.py +++ b/src/alert-manager/src/job-data-recorder/kusto_util.py @@ -73,7 +73,7 @@ def find_node_triaged_failure_in_kusto(self, node_name, completedTime, node_name=node_name, completed_time_ms=completedTime, launched_time_ms=launchedTime, - endpoint=self.endpoint, + use_current_endpoint=True, ) if not triaged_actions: diff --git a/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py b/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py index ef910142..80aaa4dc 100644 --- a/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py +++ b/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py @@ -39,7 +39,7 @@ def __init__(self): def get_node_latest_status(self, node): node_status = self.node_status_client.get_node_status( node, - endpoint=self.endpoint, + use_current_endpoint=True, ) return node_status @@ -47,14 +47,14 @@ def get_nodes_by_status(self, status, as_of_time=None): nodes = self.node_status_client.get_nodes_by_status( status, as_of_time, - endpoint=self.endpoint, + use_current_endpoint=True, ) return nodes def get_node_latest_action(self, node): node_action = self.node_action_client.get_latest_node_action( node, - endpoint=self.endpoint, + use_current_endpoint=True, ) return node_action @@ -75,7 +75,6 @@ def update_status_action(self, node, from_status, to_status, timestamp, reason, reason, detail, category=category, - endpoint=self.endpoint, ) logger.info(f"Updated node action to {action} for node {node} on {timestamp} with category {category}") status_updated = True diff --git a/src/alert-manager/src/node-recycler/recycler.py b/src/alert-manager/src/node-recycler/recycler.py index f26607ab..925a405a 100644 --- a/src/alert-manager/src/node-recycler/recycler.py +++ b/src/alert-manager/src/node-recycler/recycler.py @@ -50,7 +50,6 @@ class NodeRecycler: _ltp_validation_image = os.getenv("LTP_VALIDATION_IMAGE") _ltp_vmss_ids = os.getenv("LTP_VMSS_IDS", "") _validation_skip_vmss_ids = set(filter(None, os.getenv("VALIDATION_SKIP_VMSS_IDS", "").split(","))) - _cluster_id = os.getenv("CLUSTER_ID") _live_nodes_retry_attempts = int(os.getenv("LIVE_NODES_RETRY_ATTEMPTS", "3")) _live_nodes_retry_interval_seconds = int(os.getenv("LIVE_NODES_RETRY_INTERVAL_SECONDS", "1")) _layout_nodes_cache = set() @@ -252,7 +251,7 @@ def ofr(cls, node_faults=None, status_client=None, action_client=None): node_faults = [] status_nodes = status_client.get_nodes_by_status( from_state, - endpoint=cls._cluster_id, + use_current_endpoint=True, ) candidates = [n.HostName for n in status_nodes] filtered_candidates, ok = cls._filter_nodes_by_policy( @@ -282,7 +281,7 @@ def ofr(cls, node_faults=None, status_client=None, action_client=None): # query the latest action separately to detect prior OFR. latest = action_client.get_latest_node_action( hostname, - endpoint=cls._cluster_id, + use_current_endpoint=True, ) if latest and latest.Action == f"{from_state}-{to_state}": logger.info(f"OFR already submitted for {hostname}, ticket_id={latest.Detail}") @@ -293,7 +292,7 @@ def ofr(cls, node_faults=None, status_client=None, action_client=None): hostname, node_id, from_state, - endpoint=cls._cluster_id, + use_current_endpoint=True, ) if result and result.Action and result.Detail: action, detail = result.Action, result.Detail @@ -354,7 +353,6 @@ def ofr(cls, node_faults=None, status_client=None, action_client=None): action_client.update_node_action( hostname, f"{from_state}-{to_state}", time.time(), ofr_str, ticket_id, "", - endpoint=cls._cluster_id, ) except Exception as e: logger.error(f"Error occured when creating OFR ticket for node {node_id} ({hostname}): {e}") @@ -442,7 +440,7 @@ def operate(cls, vmss_id, operation="start", hostnames=None, status_client=None, hostnames = [ n.HostName for n in status_client.get_nodes_by_status( from_state, - endpoint=cls._cluster_id, + use_current_endpoint=True, ) ] @@ -506,7 +504,6 @@ def is_vm_succeed_in_target(vm): action_client.update_node_action( name, f"{from_state}-{to_state}", time.time(), f"{op.title()}ing VM", "", "", - endpoint=cls._cluster_id, ) except Exception as e: logger.error(f"List instances in VMSS failed due to: {e}") @@ -556,7 +553,7 @@ def validate(cls, hostnames=None, filter_state='', status_client=None, action_cl hostnames = [ n.HostName for n in status_client.get_nodes_by_status( filter_state, - endpoint=cls._cluster_id, + use_current_endpoint=True, ) ] @@ -602,7 +599,6 @@ def validate(cls, hostnames=None, filter_state='', status_client=None, action_cl hostname, f"{filter_state}-{NodeStatus.VALIDATING.value}", time.time(), "Submitting validation job for VM", "", "", - endpoint=cls._cluster_id, ) res.raise_for_status() logger.info(f"Submitted validation job for {hostname} with response: {res.json()}") @@ -628,7 +624,6 @@ def validate(cls, hostnames=None, filter_state='', status_client=None, action_cl action_client.update_node_action( hostname, f"{filter_state}-{NodeStatus.CORDONED.value}", time.time(), f"Validation job submission failed after {cls._validation_retries[hostname]} attempts", str(e), "", - endpoint=cls._cluster_id, ) del cls._validation_retries[hostname] @@ -689,7 +684,6 @@ def skip_validation(cls, hostnames, filter_state, status_client=None, action_cli action_client.update_node_action( hostname, f"{filter_state}-{NodeStatus.AVAILABLE.value}", time.time(), "Skipping GPU validation per VMSS config", "", "", - endpoint=cls._cluster_id, ) @classmethod diff --git a/src/alert-manager/src/node-recycler/tests/test_recycler.py b/src/alert-manager/src/node-recycler/tests/test_recycler.py index a89f68cd..d36f0428 100644 --- a/src/alert-manager/src/node-recycler/tests/test_recycler.py +++ b/src/alert-manager/src/node-recycler/tests/test_recycler.py @@ -139,7 +139,6 @@ def test_records_action(self, recycler, status_client, action_client): "node-a", "allocated_ua-available", pytest.approx(time.time(), abs=5), "Skipping GPU validation per VMSS config", "", "", - endpoint="test-endpoint", ) def test_log_contains_uncordon_hint(self, recycler, status_client, action_client, caplog): diff --git a/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py b/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py index a378f0a0..73c448f5 100644 --- a/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py +++ b/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py @@ -103,8 +103,7 @@ def create_attribute_table(self) -> None: raise RuntimeError(f"Failed to create attribute table: {str(e)}") def update_node_action(self, node: str, action: str, timestamp: str, - reason: str, detail: str, category: str, - endpoint: Optional[str] = None) -> None: + reason: str, detail: str, category: str) -> None: """ Updates or inserts a node action record in the Kusto table. @@ -131,13 +130,9 @@ def update_node_action(self, node: str, action: str, timestamp: str, node_id = Node(node).get_vm_node_id_by_hostname(timestamp) # Check for existing record to avoid duplicates - endpoint_condition = "" - if endpoint: - escaped_endpoint = str(endpoint).replace("'", "''") - endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" check_query = f""" {self.table_name} - {endpoint_condition} + | where Endpoint == '{self.endpoint}' | where HostName == '{node}' and Timestamp == datetime('{timestamp}') and Action == '{action}' | count """ @@ -161,12 +156,12 @@ def update_node_action(self, node: str, action: str, timestamp: str, def get_node_actions(self, node: str, start_time: str, end_time: str, - endpoint: Optional[str] = None) -> List[NodeAction]: + use_current_endpoint: bool = False) -> List[NodeAction]: """Get action history for a node in a time range""" try: endpoint_condition = "" - if endpoint: - escaped_endpoint = str(endpoint).replace("'", "''") + if use_current_endpoint: + escaped_endpoint = str(self.endpoint).replace("'", "''") endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" {self.table_name} @@ -182,12 +177,12 @@ def get_node_actions(self, node: str, start_time: str, def get_latest_node_action(self, node: str, - endpoint: Optional[str] = None) -> Optional[NodeAction]: + use_current_endpoint: bool = False) -> Optional[NodeAction]: """Get the most recent action for a node""" try: endpoint_condition = "" - if endpoint: - escaped_endpoint = str(endpoint).replace("'", "''") + if use_current_endpoint: + escaped_endpoint = str(self.endpoint).replace("'", "''") endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" {self.table_name} @@ -201,12 +196,12 @@ def get_latest_node_action(self, raise RuntimeError(f"Failed to get latest node action: {str(e)}") def get_last_update_time(self, - endpoint: Optional[str] = None) -> Optional[datetime]: + use_current_endpoint: bool = False) -> Optional[datetime]: """Get the last update time for the node action table""" try: endpoint_condition = "" - if endpoint: - escaped_endpoint = str(endpoint).replace("'", "''") + if use_current_endpoint: + escaped_endpoint = str(self.endpoint).replace("'", "''") endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" {self.table_name} @@ -224,7 +219,7 @@ def get_latest_action_by_state( hostname: str, node_id: str, state: str, - endpoint: Optional[str] = None + use_current_endpoint: bool = False ) -> Optional[Dict[str, Any]]: """ Get the latest action that ends with the specified state for a given hostname and node_id. @@ -248,8 +243,8 @@ def get_latest_action_by_state( """ try: endpoint_condition = "" - if endpoint: - escaped_endpoint = str(endpoint).replace("'", "''") + if use_current_endpoint: + escaped_endpoint = str(self.endpoint).replace("'", "''") endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" {self.table_name} @@ -270,7 +265,7 @@ def find_triaged_failure(self, node_name: str, completed_time_ms: int, launched_time_ms: int, - endpoint: Optional[str] = None) -> List[NodeAction]: + use_current_endpoint: bool = False) -> List[NodeAction]: """ Find triaged actions for a node between job launch and completion. @@ -292,7 +287,7 @@ def find_triaged_failure(self, node_actions = self.get_node_actions(node_name, start_time, end_time, - endpoint=endpoint) + use_current_endpoint=use_current_endpoint) # Check if there's available-cordoned action cordoned_timestamp = None @@ -308,8 +303,8 @@ def find_triaged_failure(self, # Query triaged actions using KQL endpoint_declare = "" endpoint_condition = "" - if endpoint: - escaped_endpoint = str(endpoint).replace("'", "''") + if use_current_endpoint: + escaped_endpoint = str(self.endpoint).replace("'", "''") endpoint_declare = f"let endpoint = '{escaped_endpoint}';" endpoint_condition = "| where Endpoint == endpoint" diff --git a/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py b/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py index 2636748b..9baae4c0 100644 --- a/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py +++ b/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py @@ -139,14 +139,14 @@ def get_status_group(self, status: str) -> str | None: def get_node_status(self, hostname: str, timestamp: datetime = None, - endpoint: Optional[str] = None) -> Optional[NodeStatusRecord]: + use_current_endpoint: bool = False) -> Optional[NodeStatusRecord]: """Get node status at a specific time""" timestamp_str = None if timestamp is not None: timestamp_str = convert_timestamp(timestamp, format="str") query = f"{self.table_name} | where HostName == '{hostname}'" - if endpoint: - escaped_endpoint = str(endpoint).replace("'", "''") + if use_current_endpoint: + escaped_endpoint = str(self.endpoint).replace("'", "''") query += f" | where Endpoint == '{escaped_endpoint}'" if timestamp_str is not None: query += f" | where Timestamp <= datetime({timestamp_str})" @@ -163,7 +163,7 @@ def update_node_status(self, hostname: str, to_status: str, """Update node status""" timestamp = convert_timestamp(timestamp, format="datetime") current_record = self.get_node_status(hostname, timestamp, - endpoint=self.endpoint) + use_current_endpoint=True) # Validate status transition if current_record and not NodeStatus.can_transition( @@ -191,7 +191,7 @@ def get_nodes_by_status( self, status: str, as_of_time: Optional[datetime] = None, - endpoint: Optional[str] = None) -> List[NodeStatusRecord]: + use_current_endpoint: bool = False) -> List[NodeStatusRecord]: """Get all nodes whose latest/current status is exactly the specified status. Args: @@ -214,8 +214,8 @@ def get_nodes_by_status( if as_of_time: timestamp_str = convert_timestamp(as_of_time, format="str") timestamp_condition = f"| where Timestamp <= datetime({timestamp_str})" - if endpoint: - escaped_endpoint = str(endpoint).replace("'", "''") + if use_current_endpoint: + escaped_endpoint = str(self.endpoint).replace("'", "''") endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" diff --git a/src/kusto-sdk/tests/test_node_action_client.py b/src/kusto-sdk/tests/test_node_action_client.py index bb64775e..22bdbd3b 100644 --- a/src/kusto-sdk/tests/test_node_action_client.py +++ b/src/kusto-sdk/tests/test_node_action_client.py @@ -119,10 +119,9 @@ def test_update_node_action(self, client, mock_kusto_client, mock_node): "Test details", "Test", TEST_ENDPOINT ]) - def test_update_node_action_endpoint_filter_is_optional(self, client, - mock_kusto_client, - mock_node): - """update_node_action duplicate-check should scope endpoint only when provided.""" + def test_update_node_action_scopes_duplicate_check_by_client_endpoint( + self, client, mock_kusto_client, mock_node): + """update_node_action should use client endpoint for duplicate check.""" timestamp = datetime.utcnow().isoformat() mock_kusto_client.execute_command.reset_mock() mock_kusto_client.execute_command.return_value = [{"Count": 0}] @@ -135,22 +134,9 @@ def test_update_node_action_endpoint_filter_is_optional(self, client, detail="d", category="c", ) - first_query = mock_kusto_client.execute_command.call_args_list[0][0][0] - assert "Endpoint ==" not in first_query - mock_kusto_client.execute_command.reset_mock() - mock_kusto_client.execute_command.return_value = [{"Count": 0}] - client.update_node_action( - node="test-node", - action="available-cordoned", - timestamp=timestamp, - reason="r", - detail="d", - category="c", - endpoint=TEST_ENDPOINT, - ) - second_query = mock_kusto_client.execute_command.call_args_list[0][0][0] - assert f"Endpoint == '{TEST_ENDPOINT}'" in second_query + query = mock_kusto_client.execute_command.call_args_list[0][0][0] + assert f"Endpoint == '{TEST_ENDPOINT}'" in query def test_update_node_action_invalid_action(self, client): """Test update_node_action method with invalid action""" @@ -205,7 +191,7 @@ def test_get_latest_node_action_with_endpoint_filter(self, client, """get_latest_node_action should filter endpoint when provided.""" mock_kusto_client.execute_command.return_value = [] - client.get_latest_node_action("test-node", endpoint=TEST_ENDPOINT) + client.get_latest_node_action("test-node", use_current_endpoint=True) query = mock_kusto_client.execute_command.call_args[0][0] assert f"Endpoint == '{TEST_ENDPOINT}'" in query @@ -268,7 +254,7 @@ def test_get_node_actions_with_endpoint_filter(self, client, node="test-node", start_time=start_time.isoformat(), end_time=end_time.isoformat(), - endpoint=TEST_ENDPOINT, + use_current_endpoint=True, ) query = mock_kusto_client.execute_command.call_args[0][0] @@ -316,7 +302,7 @@ def test_find_triaged_failure_scopes_endpoint(self, client, mock_kusto_client): node_name="test-node", completed_time_ms=1704067500000, launched_time_ms=1704067200000, - endpoint=TEST_ENDPOINT, + use_current_endpoint=True, ) query = mock_kusto_client.execute_command.call_args[0][0] diff --git a/src/kusto-sdk/tests/test_node_status_client.py b/src/kusto-sdk/tests/test_node_status_client.py index 3a8eccd3..8d79bc43 100644 --- a/src/kusto-sdk/tests/test_node_status_client.py +++ b/src/kusto-sdk/tests/test_node_status_client.py @@ -167,7 +167,11 @@ def test_get_node_status_with_explicit_endpoint_filter(self, client, timestamp = datetime.utcnow() mock_kusto_client.execute_command.return_value = [] - client.get_node_status("test-node", timestamp.timestamp(), endpoint=TEST_ENDPOINT) + client.get_node_status( + "test-node", + timestamp.timestamp(), + use_current_endpoint=True, + ) query = mock_kusto_client.execute_command.call_args[0][0] assert f"Endpoint == '{TEST_ENDPOINT}'" in query @@ -256,6 +260,9 @@ def test_get_nodes_by_status_with_optional_endpoint(self, client, first_query = mock_kusto_client.execute_command.call_args[0][0] assert "Endpoint ==" not in first_query - client.get_nodes_by_status(NodeStatus.CORDONED.value, endpoint=TEST_ENDPOINT) + client.get_nodes_by_status( + NodeStatus.CORDONED.value, + use_current_endpoint=True, + ) second_query = mock_kusto_client.execute_command.call_args[0][0] assert f"Endpoint == '{TEST_ENDPOINT}'" in second_query From 9c04a05991b65dd8a11a7483b1bd8073928fa95d Mon Sep 17 00:00:00 2001 From: Rui Gao Date: Fri, 21 Aug 2026 09:56:41 +0000 Subject: [PATCH 5/6] Narrow endpoint scoping to selected get methods --- .../src/alert-parser/node_updater.py | 9 +-- .../src/job-data-recorder/kusto_util.py | 1 - .../node_recorder_helper.py | 6 +- .../src/node-recycler/recycler.py | 16 +----- .../features/node_action/client.py | 39 ++----------- .../features/node_status/client.py | 23 ++++---- .../tests/test_node_action_client.py | 56 ------------------- .../tests/test_node_status_client.py | 16 ------ 8 files changed, 23 insertions(+), 143 deletions(-) diff --git a/src/alert-manager/src/alert-parser/node_updater.py b/src/alert-manager/src/alert-parser/node_updater.py index f7246f0a..322b578e 100644 --- a/src/alert-manager/src/alert-parser/node_updater.py +++ b/src/alert-manager/src/alert-parser/node_updater.py @@ -51,16 +51,11 @@ def get_node_latest_status(self, node, as_of_time=None): return node_status def get_nodes_by_status(self, status, as_of_time=None): - nodes = self.node_status_client.get_nodes_by_status( - status, - as_of_time, - use_current_endpoint=True, - ) + nodes = self.node_status_client.get_nodes_by_status(status, as_of_time) return nodes def get_last_actions_update_time(self): - time = self.node_action_client.get_last_update_time( - use_current_endpoint=True) + time = self.node_action_client.get_last_update_time() return time diff --git a/src/alert-manager/src/job-data-recorder/kusto_util.py b/src/alert-manager/src/job-data-recorder/kusto_util.py index 26d0393f..84b58dd6 100644 --- a/src/alert-manager/src/job-data-recorder/kusto_util.py +++ b/src/alert-manager/src/job-data-recorder/kusto_util.py @@ -73,7 +73,6 @@ def find_node_triaged_failure_in_kusto(self, node_name, completedTime, node_name=node_name, completed_time_ms=completedTime, launched_time_ms=launchedTime, - use_current_endpoint=True, ) if not triaged_actions: diff --git a/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py b/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py index 80aaa4dc..d41af9fe 100644 --- a/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py +++ b/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py @@ -44,11 +44,7 @@ def get_node_latest_status(self, node): return node_status def get_nodes_by_status(self, status, as_of_time=None): - nodes = self.node_status_client.get_nodes_by_status( - status, - as_of_time, - use_current_endpoint=True, - ) + nodes = self.node_status_client.get_nodes_by_status(status, as_of_time) return nodes def get_node_latest_action(self, node): diff --git a/src/alert-manager/src/node-recycler/recycler.py b/src/alert-manager/src/node-recycler/recycler.py index 925a405a..f625c1d2 100644 --- a/src/alert-manager/src/node-recycler/recycler.py +++ b/src/alert-manager/src/node-recycler/recycler.py @@ -249,10 +249,7 @@ def ofr(cls, node_faults=None, status_client=None, action_client=None): created = [] if not node_faults and status_client and action_client: node_faults = [] - status_nodes = status_client.get_nodes_by_status( - from_state, - use_current_endpoint=True, - ) + status_nodes = status_client.get_nodes_by_status(from_state) candidates = [n.HostName for n in status_nodes] filtered_candidates, ok = cls._filter_nodes_by_policy( candidates, @@ -292,7 +289,6 @@ def ofr(cls, node_faults=None, status_client=None, action_client=None): hostname, node_id, from_state, - use_current_endpoint=True, ) if result and result.Action and result.Detail: action, detail = result.Action, result.Detail @@ -438,10 +434,7 @@ def operate(cls, vmss_id, operation="start", hostnames=None, status_client=None, if not hostnames and status_client: hostnames = [ - n.HostName for n in status_client.get_nodes_by_status( - from_state, - use_current_endpoint=True, - ) + n.HostName for n in status_client.get_nodes_by_status(from_state) ] hostnames, ok = cls._filter_nodes_by_policy( @@ -551,10 +544,7 @@ def validate(cls, hostnames=None, filter_state='', status_client=None, action_cl filter_state = NodeStatus.ALLOCATED_UA.value if not hostnames and status_client: hostnames = [ - n.HostName for n in status_client.get_nodes_by_status( - filter_state, - use_current_endpoint=True, - ) + n.HostName for n in status_client.get_nodes_by_status(filter_state) ] hostnames, ok = cls._filter_nodes_by_policy( diff --git a/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py b/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py index 73c448f5..731c6e9e 100644 --- a/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py +++ b/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py @@ -132,7 +132,6 @@ def update_node_action(self, node: str, action: str, timestamp: str, # Check for existing record to avoid duplicates check_query = f""" {self.table_name} - | where Endpoint == '{self.endpoint}' | where HostName == '{node}' and Timestamp == datetime('{timestamp}') and Action == '{action}' | count """ @@ -195,17 +194,12 @@ def get_latest_node_action(self, except Exception as e: raise RuntimeError(f"Failed to get latest node action: {str(e)}") - def get_last_update_time(self, - use_current_endpoint: bool = False) -> Optional[datetime]: + def get_last_update_time(self) -> Optional[datetime]: """Get the last update time for the node action table""" try: - endpoint_condition = "" - if use_current_endpoint: - escaped_endpoint = str(self.endpoint).replace("'", "''") - endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" {self.table_name} - {endpoint_condition} + | where Endpoint == '{self.endpoint}' | summarize arg_max(Timestamp, *) by HostName | top 1 by Timestamp desc """ @@ -218,8 +212,7 @@ def get_latest_action_by_state( self, hostname: str, node_id: str, - state: str, - use_current_endpoint: bool = False + state: str ) -> Optional[Dict[str, Any]]: """ Get the latest action that ends with the specified state for a given hostname and node_id. @@ -242,13 +235,8 @@ def get_latest_action_by_state( ... print(f"Action: {result['Action']}, Detail: {result['Detail']}") """ try: - endpoint_condition = "" - if use_current_endpoint: - escaped_endpoint = str(self.endpoint).replace("'", "''") - endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" {self.table_name} - {endpoint_condition} | where Action endswith '{state}' | where HostName == '{hostname}' and NodeId == '{node_id}' | top 1 by Timestamp desc @@ -261,11 +249,7 @@ def get_latest_action_by_state( except Exception as e: raise RuntimeError(f"Failed to get latest action by state: {str(e)}") - def find_triaged_failure(self, - node_name: str, - completed_time_ms: int, - launched_time_ms: int, - use_current_endpoint: bool = False) -> List[NodeAction]: + def find_triaged_failure(self, node_name: str, completed_time_ms: int, launched_time_ms: int) -> List[NodeAction]: """ Find triaged actions for a node between job launch and completion. @@ -284,10 +268,7 @@ def find_triaged_failure(self, start_time = convert_timestamp(launched_time_ms / 1000, "str") end_time = convert_timestamp(completed_time_ms / 1000, "str") - node_actions = self.get_node_actions(node_name, - start_time, - end_time, - use_current_endpoint=use_current_endpoint) + node_actions = self.get_node_actions(node_name, start_time, end_time) # Check if there's available-cordoned action cordoned_timestamp = None @@ -301,28 +282,18 @@ def find_triaged_failure(self, return [] # Query triaged actions using KQL - endpoint_declare = "" - endpoint_condition = "" - if use_current_endpoint: - escaped_endpoint = str(self.endpoint).replace("'", "''") - endpoint_declare = f"let endpoint = '{escaped_endpoint}';" - endpoint_condition = "| where Endpoint == endpoint" - query = f""" let node_name = '{node_name}'; - {endpoint_declare} let cordoned_ts = datetime({cordoned_timestamp}); let next_available_ts = toscalar( {self.table_name} | where HostName == node_name - {endpoint_condition} | where Action endswith '-available' and Action != 'available-cordoned' | where Timestamp > cordoned_ts | summarize min(Timestamp) ); {self.table_name} | where HostName == node_name - {endpoint_condition} | where Timestamp >= cordoned_ts | where isnull(next_available_ts) or Timestamp <= next_available_ts | where Action in ('cordoned-triaged_platform', 'cordoned-triaged_hardware', 'cordoned-triaged_user', 'cordoned-triaged_unknown') diff --git a/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py b/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py index 9baae4c0..086ca42c 100644 --- a/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py +++ b/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py @@ -162,8 +162,7 @@ def update_node_status(self, hostname: str, to_status: str, timestamp: datetime | str | int) -> str: """Update node status""" timestamp = convert_timestamp(timestamp, format="datetime") - current_record = self.get_node_status(hostname, timestamp, - use_current_endpoint=True) + current_record = self.get_node_status(hostname, timestamp) # Validate status transition if current_record and not NodeStatus.can_transition( @@ -190,8 +189,7 @@ def update_node_status(self, hostname: str, to_status: str, def get_nodes_by_status( self, status: str, - as_of_time: Optional[datetime] = None, - use_current_endpoint: bool = False) -> List[NodeStatusRecord]: + as_of_time: Optional[datetime] = None) -> List[NodeStatusRecord]: """Get all nodes whose latest/current status is exactly the specified status. Args: @@ -210,20 +208,23 @@ def get_nodes_by_status( """ try: timestamp_condition = "" - endpoint_condition = "" if as_of_time: timestamp_str = convert_timestamp(as_of_time, format="str") timestamp_condition = f"| where Timestamp <= datetime({timestamp_str})" - if use_current_endpoint: - escaped_endpoint = str(self.endpoint).replace("'", "''") - endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" query = f""" - {self.table_name} - {endpoint_condition} + let latest_status = {self.table_name} {timestamp_condition} - | summarize arg_max(Timestamp, *) by HostName + | summarize arg_max(Timestamp, Status) by HostName; + latest_status | where Status == '{status}' + | join kind=inner ( + {self.table_name} + | where Status == '{status}' + | where Endpoint == '{self.endpoint}' + | summarize arg_max(Timestamp, *) by HostName + ) on HostName + | project Timestamp=Timestamp1, HostName, Status=Status1, NodeId, Endpoint """ results = self.execute_query(query) diff --git a/src/kusto-sdk/tests/test_node_action_client.py b/src/kusto-sdk/tests/test_node_action_client.py index 22bdbd3b..be15fe10 100644 --- a/src/kusto-sdk/tests/test_node_action_client.py +++ b/src/kusto-sdk/tests/test_node_action_client.py @@ -119,25 +119,6 @@ def test_update_node_action(self, client, mock_kusto_client, mock_node): "Test details", "Test", TEST_ENDPOINT ]) - def test_update_node_action_scopes_duplicate_check_by_client_endpoint( - self, client, mock_kusto_client, mock_node): - """update_node_action should use client endpoint for duplicate check.""" - timestamp = datetime.utcnow().isoformat() - mock_kusto_client.execute_command.reset_mock() - mock_kusto_client.execute_command.return_value = [{"Count": 0}] - - client.update_node_action( - node="test-node", - action="available-cordoned", - timestamp=timestamp, - reason="r", - detail="d", - category="c", - ) - - query = mock_kusto_client.execute_command.call_args_list[0][0][0] - assert f"Endpoint == '{TEST_ENDPOINT}'" in query - def test_update_node_action_invalid_action(self, client): """Test update_node_action method with invalid action""" with pytest.raises(RuntimeError) as exc_info: @@ -289,40 +270,3 @@ def test_validate_action_format(self, client): assert not client.is_valid_action("InvalidAction") assert not client.is_valid_action("Cordoned") assert not client.is_valid_action("available") - - def test_find_triaged_failure_scopes_endpoint(self, client, mock_kusto_client): - """find_triaged_failure should scope endpoint only when provided.""" - action = MagicMock() - action.Action = "available-cordoned" - action.Timestamp = "2026-01-01T00:00:00Z" - - with patch.object(client, "get_node_actions", return_value=[action]): - mock_kusto_client.execute_command.return_value = [] - client.find_triaged_failure( - node_name="test-node", - completed_time_ms=1704067500000, - launched_time_ms=1704067200000, - use_current_endpoint=True, - ) - - query = mock_kusto_client.execute_command.call_args[0][0] - assert "let endpoint = 'test-wcu'" in query - assert query.count("| where Endpoint == endpoint") == 2 - - def test_find_triaged_failure_without_endpoint(self, client, mock_kusto_client): - """find_triaged_failure should not inject endpoint filter by default.""" - action = MagicMock() - action.Action = "available-cordoned" - action.Timestamp = "2026-01-01T00:00:00Z" - - with patch.object(client, "get_node_actions", return_value=[action]): - mock_kusto_client.execute_command.return_value = [] - client.find_triaged_failure( - node_name="test-node", - completed_time_ms=1704067500000, - launched_time_ms=1704067200000, - ) - - query = mock_kusto_client.execute_command.call_args[0][0] - assert "let endpoint =" not in query - assert "| where Endpoint == endpoint" not in query diff --git a/src/kusto-sdk/tests/test_node_status_client.py b/src/kusto-sdk/tests/test_node_status_client.py index 8d79bc43..feacc185 100644 --- a/src/kusto-sdk/tests/test_node_status_client.py +++ b/src/kusto-sdk/tests/test_node_status_client.py @@ -250,19 +250,3 @@ def test_update_node_status_invalid_transition(self, client, with pytest.raises(ValueError) as exc_info: client.update_node_status(test_node, NodeStatus.TRIAGED_HARDWARE.value, timestamp.timestamp()) - - def test_get_nodes_by_status_with_optional_endpoint(self, client, - mock_kusto_client): - """get_nodes_by_status should only scope endpoint when provided.""" - mock_kusto_client.execute_command.return_value = [] - - client.get_nodes_by_status(NodeStatus.CORDONED.value) - first_query = mock_kusto_client.execute_command.call_args[0][0] - assert "Endpoint ==" not in first_query - - client.get_nodes_by_status( - NodeStatus.CORDONED.value, - use_current_endpoint=True, - ) - second_query = mock_kusto_client.execute_command.call_args[0][0] - assert f"Endpoint == '{TEST_ENDPOINT}'" in second_query From 28013d9c94ab949e38eef859b14b65f127ad643c Mon Sep 17 00:00:00 2001 From: Rui Gao Date: Fri, 21 Aug 2026 10:29:56 +0000 Subject: [PATCH 6/6] Simplify endpoint scoping in kusto get methods --- .../src/alert-parser/node_updater.py | 6 +---- .../node_recorder_helper.py | 10 ++------ .../src/node-recycler/recycler.py | 1 - .../features/node_action/client.py | 21 +++++------------ .../features/node_status/client.py | 8 +++---- .../tests/test_node_action_client.py | 23 +++++-------------- .../tests/test_node_status_client.py | 17 +------------- 7 files changed, 19 insertions(+), 67 deletions(-) diff --git a/src/alert-manager/src/alert-parser/node_updater.py b/src/alert-manager/src/alert-parser/node_updater.py index 322b578e..0f07c44e 100644 --- a/src/alert-manager/src/alert-parser/node_updater.py +++ b/src/alert-manager/src/alert-parser/node_updater.py @@ -39,11 +39,7 @@ def __init__(self): self.retries = 3 def get_node_latest_status(self, node, as_of_time=None): - node_status = self.node_status_client.get_node_status( - node, - as_of_time, - use_current_endpoint=True, - ) + node_status = self.node_status_client.get_node_status(node, as_of_time) if not node_status: logger.info(f"No status found for node {node} as of {as_of_time}") return None diff --git a/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py b/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py index d41af9fe..a7e03525 100644 --- a/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py +++ b/src/alert-manager/src/node-issue-classifier/node_recorder_helper.py @@ -37,10 +37,7 @@ def __init__(self): self.retries = 3 def get_node_latest_status(self, node): - node_status = self.node_status_client.get_node_status( - node, - use_current_endpoint=True, - ) + node_status = self.node_status_client.get_node_status(node) return node_status def get_nodes_by_status(self, status, as_of_time=None): @@ -48,10 +45,7 @@ def get_nodes_by_status(self, status, as_of_time=None): return nodes def get_node_latest_action(self, node): - node_action = self.node_action_client.get_latest_node_action( - node, - use_current_endpoint=True, - ) + node_action = self.node_action_client.get_latest_node_action(node) return node_action def update_status_action(self, node, from_status, to_status, timestamp, reason, detail, category=''): diff --git a/src/alert-manager/src/node-recycler/recycler.py b/src/alert-manager/src/node-recycler/recycler.py index f625c1d2..6f9b58c8 100644 --- a/src/alert-manager/src/node-recycler/recycler.py +++ b/src/alert-manager/src/node-recycler/recycler.py @@ -278,7 +278,6 @@ def ofr(cls, node_faults=None, status_client=None, action_client=None): # query the latest action separately to detect prior OFR. latest = action_client.get_latest_node_action( hostname, - use_current_endpoint=True, ) if latest and latest.Action == f"{from_state}-{to_state}": logger.info(f"OFR already submitted for {hostname}, ticket_id={latest.Detail}") diff --git a/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py b/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py index 731c6e9e..db58c786 100644 --- a/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py +++ b/src/kusto-sdk/ltp_kusto_sdk/features/node_action/client.py @@ -154,17 +154,13 @@ def update_node_action(self, node: str, action: str, timestamp: str, raise RuntimeError(f"Failed to update node action: {str(e)}") def get_node_actions(self, node: str, start_time: str, - end_time: str, - use_current_endpoint: bool = False) -> List[NodeAction]: + end_time: str) -> List[NodeAction]: """Get action history for a node in a time range""" try: - endpoint_condition = "" - if use_current_endpoint: - escaped_endpoint = str(self.endpoint).replace("'", "''") - endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" + escaped_endpoint = str(self.endpoint).replace("'", "''") query = f""" {self.table_name} - {endpoint_condition} + | where Endpoint == '{escaped_endpoint}' | where HostName == '{node}' | where Timestamp between (datetime({start_time}) .. datetime({end_time})) | order by Timestamp desc @@ -174,18 +170,13 @@ def get_node_actions(self, node: str, start_time: str, except Exception as e: raise RuntimeError(f"Failed to get node actions: {str(e)}") - def get_latest_node_action(self, - node: str, - use_current_endpoint: bool = False) -> Optional[NodeAction]: + def get_latest_node_action(self, node: str) -> Optional[NodeAction]: """Get the most recent action for a node""" try: - endpoint_condition = "" - if use_current_endpoint: - escaped_endpoint = str(self.endpoint).replace("'", "''") - endpoint_condition = f"| where Endpoint == '{escaped_endpoint}'" + escaped_endpoint = str(self.endpoint).replace("'", "''") query = f""" {self.table_name} - {endpoint_condition} + | where Endpoint == '{escaped_endpoint}' | where HostName == '{node}' | top 1 by Timestamp desc """ diff --git a/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py b/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py index 086ca42c..b00ca463 100644 --- a/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py +++ b/src/kusto-sdk/ltp_kusto_sdk/features/node_status/client.py @@ -138,16 +138,14 @@ def get_status_group(self, status: str) -> str | None: def get_node_status(self, hostname: str, - timestamp: datetime = None, - use_current_endpoint: bool = False) -> Optional[NodeStatusRecord]: + timestamp: datetime = None) -> Optional[NodeStatusRecord]: """Get node status at a specific time""" timestamp_str = None if timestamp is not None: timestamp_str = convert_timestamp(timestamp, format="str") query = f"{self.table_name} | where HostName == '{hostname}'" - if use_current_endpoint: - escaped_endpoint = str(self.endpoint).replace("'", "''") - query += f" | where Endpoint == '{escaped_endpoint}'" + escaped_endpoint = str(self.endpoint).replace("'", "''") + query += f" | where Endpoint == '{escaped_endpoint}'" if timestamp_str is not None: query += f" | where Timestamp <= datetime({timestamp_str})" query += " | summarize arg_max(Timestamp, *) by HostName" diff --git a/src/kusto-sdk/tests/test_node_action_client.py b/src/kusto-sdk/tests/test_node_action_client.py index be15fe10..aaa3a1a5 100644 --- a/src/kusto-sdk/tests/test_node_action_client.py +++ b/src/kusto-sdk/tests/test_node_action_client.py @@ -157,23 +157,13 @@ def test_get_latest_node_action(self, client, mock_kusto_client): assert action.Action == "available-cordoned" assert action.Endpoint == TEST_ENDPOINT - def test_get_latest_node_action_without_endpoint_filter(self, client, - mock_kusto_client): - """get_latest_node_action should not filter endpoint by default.""" + def test_get_latest_node_action_scopes_by_endpoint(self, client, + mock_kusto_client): + """get_latest_node_action should filter by client endpoint.""" mock_kusto_client.execute_command.return_value = [] client.get_latest_node_action("test-node") - query = mock_kusto_client.execute_command.call_args[0][0] - assert "Endpoint ==" not in query - - def test_get_latest_node_action_with_endpoint_filter(self, client, - mock_kusto_client): - """get_latest_node_action should filter endpoint when provided.""" - mock_kusto_client.execute_command.return_value = [] - - client.get_latest_node_action("test-node", use_current_endpoint=True) - query = mock_kusto_client.execute_command.call_args[0][0] assert f"Endpoint == '{TEST_ENDPOINT}'" in query @@ -224,9 +214,9 @@ def test_get_node_actions(self, client, mock_kusto_client): assert actions[0].Action == "available-cordoned" assert actions[1].Action == "cordoned-triaged_hardware" - def test_get_node_actions_with_endpoint_filter(self, client, - mock_kusto_client): - """get_node_actions should filter endpoint when provided.""" + def test_get_node_actions_scopes_by_endpoint(self, client, + mock_kusto_client): + """get_node_actions should filter by client endpoint.""" start_time = datetime.utcnow() - timedelta(hours=1) end_time = datetime.utcnow() mock_kusto_client.execute_command.return_value = [] @@ -235,7 +225,6 @@ def test_get_node_actions_with_endpoint_filter(self, client, node="test-node", start_time=start_time.isoformat(), end_time=end_time.isoformat(), - use_current_endpoint=True, ) query = mock_kusto_client.execute_command.call_args[0][0] diff --git a/src/kusto-sdk/tests/test_node_status_client.py b/src/kusto-sdk/tests/test_node_status_client.py index feacc185..336fcf47 100644 --- a/src/kusto-sdk/tests/test_node_status_client.py +++ b/src/kusto-sdk/tests/test_node_status_client.py @@ -153,7 +153,7 @@ def test_get_node_status_existing(self, client, mock_kusto_client): query = mock_kusto_client.execute_command.call_args[0][0] assert TEST_STATUS_TABLE in query assert test_node in query - assert "Endpoint ==" not in query + assert f"Endpoint == '{TEST_ENDPOINT}'" in query # Verify result assert isinstance(result, NodeStatusRecord) @@ -161,21 +161,6 @@ def test_get_node_status_existing(self, client, mock_kusto_client): assert result.HostName == test_node assert result.Endpoint == TEST_ENDPOINT - def test_get_node_status_with_explicit_endpoint_filter(self, client, - mock_kusto_client): - """get_node_status should apply endpoint filter only when provided.""" - timestamp = datetime.utcnow() - mock_kusto_client.execute_command.return_value = [] - - client.get_node_status( - "test-node", - timestamp.timestamp(), - use_current_endpoint=True, - ) - - query = mock_kusto_client.execute_command.call_args[0][0] - assert f"Endpoint == '{TEST_ENDPOINT}'" in query - def test_get_node_status_new(self, client, mock_kusto_client, mock_node): """Test get_node_status method for new node""" test_node = "test-node"