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/alert-parser/node_updater.py b/src/alert-manager/src/alert-parser/node_updater.py index a6272a94..0f07c44e 100644 --- a/src/alert-manager/src/alert-parser/node_updater.py +++ b/src/alert-manager/src/alert-parser/node_updater.py @@ -76,7 +76,14 @@ 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='', + ) logger.info(f"Updated node action to {action} for node {node} on {timestamp}") status_updated = True break @@ -90,4 +97,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..84b58dd6 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,7 @@ 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, ) 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..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 @@ -58,7 +58,14 @@ 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, + ) 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 4ae993bb..6f9b58c8 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,184 @@ 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 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 = True + cls._layout_nodes_loaded = True + 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]: + """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() + if require_layout: + 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: + 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 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") + 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 +249,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}") @@ -78,13 +276,19 @@ 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, + ) 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, + ) if result and result.Action and result.Detail: action, detail = result.Action, result.Detail if action.endswith(from_state): @@ -228,7 +432,21 @@ 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) + ] + + 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}") @@ -324,7 +542,21 @@ 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) + ] + + 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 +751,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..d36f0428 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 @@ -72,14 +73,22 @@ 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" 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 + 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 +308,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 +342,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 +378,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 +454,130 @@ 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_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)): + 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() + + 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_cache_after_load_raises_runtime_error(self, recycler): + recycler._layout_nodes_loaded = True + recycler._layout_nodes_load_success = True + recycler._layout_nodes_cache = set() + + 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: + 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..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 @@ -157,8 +157,10 @@ def get_node_actions(self, node: str, start_time: str, end_time: str) -> List[NodeAction]: """Get action history for a node in a time range""" try: + escaped_endpoint = str(self.endpoint).replace("'", "''") query = f""" {self.table_name} + | where Endpoint == '{escaped_endpoint}' | where HostName == '{node}' | where Timestamp between (datetime({start_time}) .. datetime({end_time})) | order by Timestamp desc @@ -171,8 +173,10 @@ def get_node_actions(self, node: str, start_time: str, def get_latest_node_action(self, node: str) -> Optional[NodeAction]: """Get the most recent action for a node""" try: + escaped_endpoint = str(self.endpoint).replace("'", "''") query = f""" {self.table_name} + | 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 c78744ad..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,12 +138,14 @@ def get_status_group(self, status: str) -> str | None: def get_node_status(self, hostname: str, - timestamp: datetime = None) -> 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}'" + 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" @@ -187,16 +189,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') diff --git a/src/kusto-sdk/tests/test_node_action_client.py b/src/kusto-sdk/tests/test_node_action_client.py index 5629bb8c..aaa3a1a5 100644 --- a/src/kusto-sdk/tests/test_node_action_client.py +++ b/src/kusto-sdk/tests/test_node_action_client.py @@ -157,6 +157,16 @@ 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_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 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 +214,22 @@ 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_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 = [] + + client.get_node_actions( + node="test-node", + start_time=start_time.isoformat(), + end_time=end_time.isoformat(), + ) + + 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 = [] 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)