diff --git a/osdc/modules/zombie-cleanup/deploy.sh b/osdc/modules/zombie-cleanup/deploy.sh index 23e447ec..71fffd58 100755 --- a/osdc/modules/zombie-cleanup/deploy.sh +++ b/osdc/modules/zombie-cleanup/deploy.sh @@ -49,6 +49,9 @@ RUNNING_MAX_AGE=$(uv run "$CFG" "$CLUSTER" zombie_cleanup.running_max_age_hours DRY_RUN=$(uv run "$CFG" "$CLUSTER" zombie_cleanup.dry_run "false") PUSHGATEWAY_URL=$(uv run "$CFG" "$CLUSTER" zombie_cleanup.pushgateway_url "http://prometheus-pushgateway.monitoring.svc.cluster.local:9091") +NUMA_CORDON_DRY_RUN=$(uv run "$CFG" "$CLUSTER" zombie_cleanup.numa_cordon_dry_run "false") +NUMA_CORDON_UNCORDON=$(uv run "$CFG" "$CLUSTER" zombie_cleanup.numa_cordon_uncordon "true") + # --- Compute content-based image tag --- TAG=$(find "$MODULE_DIR/docker" "$MODULE_DIR/scripts/python" \ \( -name '*.py' -o -name 'Dockerfile' -o -name 'pyproject.toml' \) \ @@ -131,6 +134,7 @@ echo "Using ${IMAGE}:${TAG}" # --- Apply RBAC --- echo " Applying RBAC..." kubectl_apply_if_changed -f "$MODULE_DIR/kubernetes/rbac.yaml" +kubectl_apply_if_changed -f "$MODULE_DIR/kubernetes/numa-cordon-rbac.yaml" # --- Apply CronJob with config substitution --- echo " Applying CronJob..." @@ -142,4 +146,11 @@ sed \ -e "s|PUSHGATEWAY_URL_PLACEHOLDER|${PUSHGATEWAY_URL}|" \ "$MODULE_DIR/kubernetes/cronjob.yaml" | kubectl_apply_if_changed -f - +echo " Applying numa-cordon CronJob..." +sed \ + -e "s|ZOMBIE_CLEANUP_IMAGE_PLACEHOLDER|${IMAGE}:${TAG}|" \ + -e "s|NUMA_CORDON_DRY_RUN_PLACEHOLDER|${NUMA_CORDON_DRY_RUN}|" \ + -e "s|NUMA_CORDON_UNCORDON_PLACEHOLDER|${NUMA_CORDON_UNCORDON}|" \ + "$MODULE_DIR/kubernetes/numa-cordon-cronjob.yaml" | kubectl_apply_if_changed -f - + echo " zombie-cleanup deployed." diff --git a/osdc/modules/zombie-cleanup/kubernetes/numa-cordon-cronjob.yaml b/osdc/modules/zombie-cleanup/kubernetes/numa-cordon-cronjob.yaml new file mode 100644 index 00000000..8961cc80 --- /dev/null +++ b/osdc/modules/zombie-cleanup/kubernetes/numa-cordon-cronjob.yaml @@ -0,0 +1,66 @@ +apiVersion: batch/v1 +kind: CronJob +metadata: + name: numa-cordon + namespace: arc-runners + labels: + app.kubernetes.io/name: numa-cordon + app.kubernetes.io/part-of: osdc-zombie-cleanup +spec: + # Run every 5 minutes to break the fragmentation livelock quickly. + # The zombie-cleanup CronJob runs hourly; this needs to be faster + # because 4 consecutive failures can happen in ~90s on a single node. + schedule: "*/5 * * * *" + concurrencyPolicy: Forbid + successfulJobsHistoryLimit: 3 + failedJobsHistoryLimit: 3 + jobTemplate: + spec: + backoffLimit: 0 + activeDeadlineSeconds: 120 + template: + metadata: + labels: + app.kubernetes.io/name: numa-cordon + app.kubernetes.io/part-of: osdc-zombie-cleanup + spec: + serviceAccountName: numa-cordon + restartPolicy: Never + nodeSelector: + role: base-infrastructure + tolerations: + - key: CriticalAddonsOnly + operator: Exists + effect: NoSchedule + securityContext: + runAsNonRoot: true + runAsUser: 65534 + runAsGroup: 65534 + seccompProfile: + type: RuntimeDefault + containers: + - name: numa-cordon + image: ZOMBIE_CLEANUP_IMAGE_PLACEHOLDER + command: ["python", "numa_cordon.py"] + env: + - name: PYTHONDONTWRITEBYTECODE + value: "1" + - name: TARGET_NAMESPACE + value: "arc-runners" + - name: DRY_RUN + value: "NUMA_CORDON_DRY_RUN_PLACEHOLDER" + - name: UNCORDON_ENABLED + value: "NUMA_CORDON_UNCORDON_PLACEHOLDER" + resources: + requests: + cpu: 100m + memory: 128Mi + limits: + cpu: 500m + memory: 256Mi + securityContext: + readOnlyRootFilesystem: true + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL diff --git a/osdc/modules/zombie-cleanup/kubernetes/numa-cordon-rbac.yaml b/osdc/modules/zombie-cleanup/kubernetes/numa-cordon-rbac.yaml new file mode 100644 index 00000000..7c64b630 --- /dev/null +++ b/osdc/modules/zombie-cleanup/kubernetes/numa-cordon-rbac.yaml @@ -0,0 +1,68 @@ +apiVersion: v1 +kind: ServiceAccount +metadata: + name: numa-cordon + namespace: arc-runners + labels: + app.kubernetes.io/name: numa-cordon + app.kubernetes.io/part-of: osdc-zombie-cleanup +--- +# Namespace-scoped: list/delete Failed pods in arc-runners. +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: numa-cordon + namespace: arc-runners + labels: + app.kubernetes.io/name: numa-cordon + app.kubernetes.io/part-of: osdc-zombie-cleanup +rules: + - apiGroups: [""] + resources: ["pods"] + verbs: ["get", "list", "delete"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: numa-cordon + namespace: arc-runners + labels: + app.kubernetes.io/name: numa-cordon + app.kubernetes.io/part-of: osdc-zombie-cleanup +subjects: + - kind: ServiceAccount + name: numa-cordon + namespace: arc-runners +roleRef: + kind: Role + name: numa-cordon + apiGroup: rbac.authorization.k8s.io +--- +# Cluster-scoped: get/list/patch nodes for cordon/uncordon. +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: numa-cordon + labels: + app.kubernetes.io/name: numa-cordon + app.kubernetes.io/part-of: osdc-zombie-cleanup +rules: + - apiGroups: [""] + resources: ["nodes"] + verbs: ["get", "list", "patch"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: numa-cordon + labels: + app.kubernetes.io/name: numa-cordon + app.kubernetes.io/part-of: osdc-zombie-cleanup +subjects: + - kind: ServiceAccount + name: numa-cordon + namespace: arc-runners +roleRef: + kind: ClusterRole + name: numa-cordon + apiGroup: rbac.authorization.k8s.io diff --git a/osdc/modules/zombie-cleanup/scripts/python/numa_cordon.py b/osdc/modules/zombie-cleanup/scripts/python/numa_cordon.py new file mode 100644 index 00000000..394c4508 --- /dev/null +++ b/osdc/modules/zombie-cleanup/scripts/python/numa_cordon.py @@ -0,0 +1,311 @@ +#!/usr/bin/env python3 +"""NUMA trap-node cordoner for ARC GPU runner namespaces. + +Detects pods that failed with TopologyAffinityError (kubelet rejected the +pod because the requested GPUs could not fit on a single NUMA socket) and +cordons the offending node so the scheduler stops routing new GPU pods to +it. This breaks the fragmentation livelock described in +https://github.com/pytorch/ci-infra/issues/696. + +Cordoned nodes drain naturally as existing pods finish. Karpenter will +provision replacement capacity if needed. A separate uncordon pass +(run with UNCORDON_ENABLED=true) removes the cordon once all GPU pods +on the node have finished, returning it to the schedulable pool. +""" + +import logging +import os +import sys +import time +from datetime import UTC, datetime + +from lightkube import Client +from lightkube.core.exceptions import ApiError +from lightkube.models.meta_v1 import ObjectMeta +from lightkube.resources.core_v1 import Node, Pod +from lightkube.types import PatchType + +log = logging.getLogger("numa-cordon") + +# Annotation set on nodes we cordon so we only uncordon our own work. +CORDON_ANNOTATION = "osdc.io/numa-cordoned" + + +def get_config() -> dict: + return { + "namespace": os.environ.get("TARGET_NAMESPACE", "arc-runners"), + "dry_run": os.environ.get("DRY_RUN", "false").lower() in ("true", "1", "yes"), + "uncordon_enabled": os.environ.get("UNCORDON_ENABLED", "true").lower() + in ("true", "1", "yes"), + } + + +def find_topology_failed_pods(client: Client, namespace: str) -> list[Pod]: + """Find pods that failed with TopologyAffinityError.""" + failed = [] + for pod in client.list(Pod, namespace=namespace): + if pod.status is None: + continue + if pod.status.phase != "Failed": + continue + reason = getattr(pod.status, "reason", None) or "" + if reason == "TopologyAffinityError": + failed.append(pod) + return failed + + +def cordon_nodes( + client: Client, pods: list[Pod], dry_run: bool +) -> tuple[set[str], int, int]: + """Cordon nodes that rejected pods with TopologyAffinityError. + + Returns (cordoned_node_names, cordon_count, fail_count). + """ + # Collect unique node names from failed pods. + node_names: dict[str, list[str]] = {} + for pod in pods: + node = getattr(pod.spec, "nodeName", None) + if not node: + log.warning( + "Pod %s has TopologyAffinityError but no nodeName, skipping", + pod.metadata.name, + ) + continue + node_names.setdefault(node, []).append(pod.metadata.name) + + cordoned: set[str] = set() + failed = 0 + + for node_name, pod_names in node_names.items(): + # Check if already unschedulable. + try: + node_obj = client.get(Node, name=node_name) + except ApiError as e: + if e.status.code == 404: + log.info("Node %s no longer exists, skipping", node_name) + continue + log.exception("Failed to get node %s", node_name) + failed += 1 + continue + + if getattr(node_obj.spec, "unschedulable", False): + log.info( + "Node %s already unschedulable, skipping cordon (pods: %s)", + node_name, + ", ".join(pod_names), + ) + cordoned.add(node_name) + continue + + if dry_run: + log.info( + "DRY RUN: would cordon node %s (pods: %s)", + node_name, + ", ".join(pod_names), + ) + cordoned.add(node_name) + continue + + # Cordon: set spec.unschedulable = true and annotate. + patch = { + "spec": {"unschedulable": True}, + "metadata": { + "annotations": { + CORDON_ANNOTATION: datetime.now(UTC).isoformat(), + } + }, + } + try: + client.patch(Node, name=node_name, obj=patch, patch_type=PatchType.MERGE) + log.info( + "Cordoned node %s (pods: %s)", node_name, ", ".join(pod_names) + ) + cordoned.add(node_name) + except Exception: + log.exception("Failed to cordon node %s", node_name) + failed += 1 + + return cordoned, len(cordoned), failed + + +def cleanup_failed_pods( + client: Client, + pods: list[Pod], + cordoned_nodes: set[str], + namespace: str, + dry_run: bool, +) -> tuple[int, int]: + """Delete failed TopologyAffinityError pods on cordoned nodes. + + Returns (deleted_count, fail_count). + """ + deleted = 0 + failed = 0 + for pod in pods: + node = getattr(pod.spec, "nodeName", None) + if node not in cordoned_nodes: + continue + name = pod.metadata.name + if dry_run: + log.info("DRY RUN: would delete failed pod %s", name) + deleted += 1 + continue + try: + client.delete(Pod, name=name, namespace=namespace) + log.info("Deleted failed pod %s", name) + deleted += 1 + except ApiError as e: + if e.status.code == 404: + log.info("Pod %s already gone", name) + deleted += 1 + else: + log.exception("Failed to delete pod %s", name) + failed += 1 + except Exception: + log.exception("Failed to delete pod %s", name) + failed += 1 + return deleted, failed + + +def uncordon_drained_nodes(client: Client, dry_run: bool) -> tuple[int, int]: + """Uncordon nodes we previously cordoned that no longer have GPU pods. + + Only touches nodes bearing our CORDON_ANNOTATION. + Returns (uncordoned_count, fail_count). + """ + uncordoned = 0 + failed = 0 + + for node in client.list(Node): + annotations = node.metadata.annotations or {} + if CORDON_ANNOTATION not in annotations: + continue + if not getattr(node.spec, "unschedulable", False): + # Already schedulable — clean up stale annotation. + _remove_annotation(client, node.metadata.name, dry_run) + continue + + # Check if any GPU pods remain on this node. + has_gpu_pods = False + field_sel = f"spec.nodeName={node.metadata.name}" + for pod in client.list(Pod, field_selector=field_sel): + if pod.status and pod.status.phase in ("Succeeded", "Failed"): + continue + containers = [] + if pod.spec and pod.spec.containers: + containers = pod.spec.containers + for c in containers: + limits = {} + if c.resources and c.resources.limits: + limits = c.resources.limits + if "nvidia.com/gpu" in limits: + has_gpu_pods = True + break + if has_gpu_pods: + break + + if has_gpu_pods: + log.info( + "Node %s still has active GPU pods, keeping cordoned", + node.metadata.name, + ) + continue + + if dry_run: + log.info("DRY RUN: would uncordon node %s", node.metadata.name) + uncordoned += 1 + continue + + patch = { + "spec": {"unschedulable": False}, + "metadata": {"annotations": {CORDON_ANNOTATION: None}}, + } + try: + client.patch( + Node, name=node.metadata.name, obj=patch, patch_type=PatchType.MERGE + ) + log.info("Uncordoned node %s (GPU pods drained)", node.metadata.name) + uncordoned += 1 + except Exception: + log.exception("Failed to uncordon node %s", node.metadata.name) + failed += 1 + + return uncordoned, failed + + +def _remove_annotation(client: Client, node_name: str, dry_run: bool) -> None: + """Remove our cordon annotation from a node that's already schedulable.""" + if dry_run: + return + patch = {"metadata": {"annotations": {CORDON_ANNOTATION: None}}} + try: + client.patch(Node, name=node_name, obj=patch, patch_type=PatchType.MERGE) + except Exception: + log.exception("Failed to remove annotation from node %s", node_name) + + +def main() -> int: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + datefmt="%Y-%m-%dT%H:%M:%S", + ) + + config = get_config() + log.info( + "Starting numa-cordon: namespace=%s dry_run=%s uncordon_enabled=%s", + config["namespace"], + config["dry_run"], + config["uncordon_enabled"], + ) + + client = Client() + start = time.monotonic() + + try: + # Phase 1: Find and cordon. + pods = find_topology_failed_pods(client, config["namespace"]) + if pods: + log.info("Found %d pod(s) with TopologyAffinityError", len(pods)) + cordoned, cordon_ok, cordon_fail = cordon_nodes( + client, pods, config["dry_run"] + ) + del_ok, del_fail = cleanup_failed_pods( + client, pods, cordoned, config["namespace"], config["dry_run"] + ) + log.info( + "Cordon phase: nodes_cordoned=%d cordon_failed=%d " + "pods_deleted=%d pod_delete_failed=%d", + cordon_ok, + cordon_fail, + del_ok, + del_fail, + ) + else: + log.info("No TopologyAffinityError pods found") + cordon_fail = 0 + del_fail = 0 + + # Phase 2: Uncordon drained nodes. + uncordon_fail = 0 + if config["uncordon_enabled"]: + uncordon_ok, uncordon_fail = uncordon_drained_nodes( + client, config["dry_run"] + ) + log.info( + "Uncordon phase: nodes_uncordoned=%d uncordon_failed=%d", + uncordon_ok, + uncordon_fail, + ) + + elapsed = time.monotonic() - start + log.info("Completed in %.1fs", elapsed) + return 1 if (cordon_fail + del_fail + uncordon_fail) > 0 else 0 + + except Exception as e: + log.exception("numa-cordon failed: %s", e) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/osdc/modules/zombie-cleanup/scripts/python/test_numa_cordon.py b/osdc/modules/zombie-cleanup/scripts/python/test_numa_cordon.py new file mode 100644 index 00000000..fb8009cc --- /dev/null +++ b/osdc/modules/zombie-cleanup/scripts/python/test_numa_cordon.py @@ -0,0 +1,462 @@ +"""Tests for numa_cordon module.""" + +from datetime import UTC, datetime, timedelta +from unittest.mock import MagicMock, call, patch + +import pytest +from lightkube.core.exceptions import ApiError +from numa_cordon import ( + CORDON_ANNOTATION, + cleanup_failed_pods, + cordon_nodes, + find_topology_failed_pods, + get_config, + main, + uncordon_drained_nodes, +) + + +def _make_pod( + name: str, + phase: str = "Running", + reason: str | None = None, + node_name: str | None = None, + gpu: int = 0, +): + pod = MagicMock() + pod.metadata.name = name + pod.status.phase = phase + pod.status.reason = reason + pod.spec.nodeName = node_name + + if gpu > 0: + container = MagicMock() + container.resources.limits = {"nvidia.com/gpu": str(gpu)} + pod.spec.containers = [container] + else: + container = MagicMock() + container.resources.limits = {"cpu": "1"} + pod.spec.containers = [container] + + return pod + + +def _make_node(name: str, unschedulable: bool = False, cordoned_by_us: bool = False): + node = MagicMock() + node.metadata.name = name + node.spec.unschedulable = unschedulable + annotations = {} + if cordoned_by_us: + annotations[CORDON_ANNOTATION] = datetime.now(UTC).isoformat() + node.metadata.annotations = annotations + return node + + +# --- find_topology_failed_pods --- + + +class TestFindTopologyFailedPods: + def test_no_pods(self): + client = MagicMock() + client.list.return_value = [] + assert find_topology_failed_pods(client, "arc-runners") == [] + + def test_finds_topology_error_pods(self): + client = MagicMock() + client.list.return_value = [ + _make_pod("good-pod", phase="Running"), + _make_pod( + "topo-fail", + phase="Failed", + reason="TopologyAffinityError", + node_name="node-1", + ), + _make_pod("other-fail", phase="Failed", reason="OOMKilled"), + ] + result = find_topology_failed_pods(client, "arc-runners") + assert len(result) == 1 + assert result[0].metadata.name == "topo-fail" + + def test_skips_non_failed_phases(self): + client = MagicMock() + client.list.return_value = [ + _make_pod("pending", phase="Pending"), + _make_pod("running", phase="Running"), + _make_pod("succeeded", phase="Succeeded"), + ] + assert find_topology_failed_pods(client, "arc-runners") == [] + + def test_skips_pods_without_status(self): + client = MagicMock() + pod = _make_pod("no-status") + pod.status = None + client.list.return_value = [pod] + assert find_topology_failed_pods(client, "arc-runners") == [] + + def test_multiple_topology_errors(self): + client = MagicMock() + client.list.return_value = [ + _make_pod("fail-1", phase="Failed", reason="TopologyAffinityError", node_name="n1"), + _make_pod("fail-2", phase="Failed", reason="TopologyAffinityError", node_name="n1"), + _make_pod("fail-3", phase="Failed", reason="TopologyAffinityError", node_name="n2"), + ] + result = find_topology_failed_pods(client, "arc-runners") + assert len(result) == 3 + + +# --- cordon_nodes --- + + +class TestCordonNodes: + def test_cordons_node(self): + client = MagicMock() + node = _make_node("node-1", unschedulable=False) + client.get.return_value = node + pods = [ + _make_pod("fail-1", phase="Failed", reason="TopologyAffinityError", node_name="node-1") + ] + + cordoned, ok, fail = cordon_nodes(client, pods, dry_run=False) + assert cordoned == {"node-1"} + assert ok == 1 + assert fail == 0 + client.patch.assert_called_once() + + def test_skips_already_unschedulable(self): + client = MagicMock() + node = _make_node("node-1", unschedulable=True) + client.get.return_value = node + pods = [ + _make_pod("fail-1", phase="Failed", reason="TopologyAffinityError", node_name="node-1") + ] + + cordoned, ok, fail = cordon_nodes(client, pods, dry_run=False) + assert cordoned == {"node-1"} + assert ok == 1 + assert fail == 0 + client.patch.assert_not_called() + + def test_dry_run_does_not_patch(self): + client = MagicMock() + node = _make_node("node-1", unschedulable=False) + client.get.return_value = node + pods = [ + _make_pod("fail-1", phase="Failed", reason="TopologyAffinityError", node_name="node-1") + ] + + cordoned, ok, fail = cordon_nodes(client, pods, dry_run=True) + assert cordoned == {"node-1"} + assert ok == 1 + assert fail == 0 + client.patch.assert_not_called() + + def test_deduplicates_nodes(self): + client = MagicMock() + node = _make_node("node-1", unschedulable=False) + client.get.return_value = node + pods = [ + _make_pod("f1", phase="Failed", reason="TopologyAffinityError", node_name="node-1"), + _make_pod("f2", phase="Failed", reason="TopologyAffinityError", node_name="node-1"), + ] + + cordoned, ok, fail = cordon_nodes(client, pods, dry_run=False) + assert cordoned == {"node-1"} + assert ok == 1 + # Only one patch call for the node, not two. + assert client.patch.call_count == 1 + + def test_handles_node_not_found(self): + client = MagicMock() + not_found = ApiError.__new__(ApiError) + not_found.status = MagicMock(code=404) + client.get.side_effect = not_found + pods = [ + _make_pod("f1", phase="Failed", reason="TopologyAffinityError", node_name="gone-node") + ] + + cordoned, ok, fail = cordon_nodes(client, pods, dry_run=False) + assert cordoned == set() + assert ok == 0 + assert fail == 0 + + def test_handles_patch_failure(self): + client = MagicMock() + node = _make_node("node-1", unschedulable=False) + client.get.return_value = node + client.patch.side_effect = Exception("API error") + pods = [ + _make_pod("f1", phase="Failed", reason="TopologyAffinityError", node_name="node-1") + ] + + cordoned, ok, fail = cordon_nodes(client, pods, dry_run=False) + assert fail == 1 + + def test_skips_pod_without_node_name(self): + client = MagicMock() + pods = [ + _make_pod("f1", phase="Failed", reason="TopologyAffinityError", node_name=None) + ] + + cordoned, ok, fail = cordon_nodes(client, pods, dry_run=False) + assert cordoned == set() + assert ok == 0 + assert fail == 0 + client.get.assert_not_called() + + def test_handles_get_node_api_error(self): + client = MagicMock() + forbidden = ApiError.__new__(ApiError) + forbidden.status = MagicMock(code=403) + client.get.side_effect = forbidden + pods = [ + _make_pod("f1", phase="Failed", reason="TopologyAffinityError", node_name="node-1") + ] + + cordoned, ok, fail = cordon_nodes(client, pods, dry_run=False) + assert fail == 1 + + +# --- cleanup_failed_pods --- + + +class TestCleanupFailedPods: + def test_deletes_pods_on_cordoned_nodes(self): + client = MagicMock() + pods = [ + _make_pod("f1", phase="Failed", reason="TopologyAffinityError", node_name="node-1"), + _make_pod("f2", phase="Failed", reason="TopologyAffinityError", node_name="node-2"), + ] + + deleted, failed = cleanup_failed_pods( + client, pods, {"node-1", "node-2"}, "arc-runners", dry_run=False + ) + assert deleted == 2 + assert failed == 0 + assert client.delete.call_count == 2 + + def test_skips_pods_not_on_cordoned_nodes(self): + client = MagicMock() + pods = [ + _make_pod("f1", phase="Failed", reason="TopologyAffinityError", node_name="other-node"), + ] + + deleted, failed = cleanup_failed_pods( + client, pods, {"node-1"}, "arc-runners", dry_run=False + ) + assert deleted == 0 + assert failed == 0 + client.delete.assert_not_called() + + def test_dry_run(self): + client = MagicMock() + pods = [ + _make_pod("f1", phase="Failed", reason="TopologyAffinityError", node_name="node-1"), + ] + + deleted, failed = cleanup_failed_pods( + client, pods, {"node-1"}, "arc-runners", dry_run=True + ) + assert deleted == 1 + assert failed == 0 + client.delete.assert_not_called() + + def test_404_counted_as_success(self): + client = MagicMock() + not_found = ApiError.__new__(ApiError) + not_found.status = MagicMock(code=404) + client.delete.side_effect = not_found + pods = [ + _make_pod("gone", phase="Failed", reason="TopologyAffinityError", node_name="node-1"), + ] + + deleted, failed = cleanup_failed_pods( + client, pods, {"node-1"}, "arc-runners", dry_run=False + ) + assert deleted == 1 + assert failed == 0 + + def test_api_error_non_404(self): + client = MagicMock() + forbidden = ApiError.__new__(ApiError) + forbidden.status = MagicMock(code=403) + client.delete.side_effect = forbidden + pods = [ + _make_pod("f1", phase="Failed", reason="TopologyAffinityError", node_name="node-1"), + ] + + deleted, failed = cleanup_failed_pods( + client, pods, {"node-1"}, "arc-runners", dry_run=False + ) + assert deleted == 0 + assert failed == 1 + + +# --- uncordon_drained_nodes --- + + +class TestUncordonDrainedNodes: + def test_uncordons_node_with_no_gpu_pods(self): + client = MagicMock() + node = _make_node("node-1", unschedulable=True, cordoned_by_us=True) + client.list.side_effect = [ + [node], # Node list + [], # Pod list for node-1 (no pods) + ] + + ok, fail = uncordon_drained_nodes(client, dry_run=False) + assert ok == 1 + assert fail == 0 + client.patch.assert_called_once() + + def test_keeps_cordon_when_gpu_pods_remain(self): + client = MagicMock() + node = _make_node("node-1", unschedulable=True, cordoned_by_us=True) + gpu_pod = _make_pod("gpu-runner", phase="Running", gpu=4) + client.list.side_effect = [ + [node], # Node list + [gpu_pod], # Pod list for node-1 + ] + + ok, fail = uncordon_drained_nodes(client, dry_run=False) + assert ok == 0 + assert fail == 0 + client.patch.assert_not_called() + + def test_ignores_nodes_not_cordoned_by_us(self): + client = MagicMock() + node = _make_node("node-1", unschedulable=True, cordoned_by_us=False) + client.list.return_value = [node] + + ok, fail = uncordon_drained_nodes(client, dry_run=False) + assert ok == 0 + assert fail == 0 + + def test_cleans_annotation_on_already_schedulable_node(self): + client = MagicMock() + node = _make_node("node-1", unschedulable=False, cordoned_by_us=True) + client.list.return_value = [node] + + ok, fail = uncordon_drained_nodes(client, dry_run=False) + assert ok == 0 + assert fail == 0 + # Should patch to remove annotation. + client.patch.assert_called_once() + + def test_dry_run(self): + client = MagicMock() + node = _make_node("node-1", unschedulable=True, cordoned_by_us=True) + client.list.side_effect = [ + [node], + [], # No pods + ] + + ok, fail = uncordon_drained_nodes(client, dry_run=True) + assert ok == 1 + assert fail == 0 + client.patch.assert_not_called() + + def test_ignores_completed_gpu_pods(self): + client = MagicMock() + node = _make_node("node-1", unschedulable=True, cordoned_by_us=True) + done_pod = _make_pod("done-gpu", phase="Succeeded", gpu=4) + client.list.side_effect = [ + [node], + [done_pod], # Only completed GPU pods + ] + + ok, fail = uncordon_drained_nodes(client, dry_run=False) + assert ok == 1 + assert fail == 0 + + def test_handles_uncordon_failure(self): + client = MagicMock() + node = _make_node("node-1", unschedulable=True, cordoned_by_us=True) + client.list.side_effect = [ + [node], + [], + ] + client.patch.side_effect = Exception("API error") + + ok, fail = uncordon_drained_nodes(client, dry_run=False) + assert ok == 0 + assert fail == 1 + + +# --- get_config --- + + +class TestGetConfig: + def test_defaults(self): + with patch.dict("os.environ", {}, clear=True): + config = get_config() + assert config["namespace"] == "arc-runners" + assert config["dry_run"] is False + assert config["uncordon_enabled"] is True + + def test_custom_values(self): + env = { + "TARGET_NAMESPACE": "custom-ns", + "DRY_RUN": "true", + "UNCORDON_ENABLED": "false", + } + with patch.dict("os.environ", env, clear=True): + config = get_config() + assert config["namespace"] == "custom-ns" + assert config["dry_run"] is True + assert config["uncordon_enabled"] is False + + +# --- main --- + + +class TestMain: + def test_no_failed_pods(self): + with ( + patch("numa_cordon.Client") as mock_cls, + patch.dict("os.environ", {}, clear=True), + ): + mock_client = MagicMock() + mock_client.list.return_value = [] + mock_cls.return_value = mock_client + + assert main() == 0 + + def test_cordons_and_cleans_up(self): + pod = _make_pod( + "fail-1", + phase="Failed", + reason="TopologyAffinityError", + node_name="node-1", + ) + node = _make_node("node-1", unschedulable=False) + + with ( + patch("numa_cordon.Client") as mock_cls, + patch.dict("os.environ", {}, clear=True), + ): + mock_client = MagicMock() + # 1st list: find_topology_failed_pods, 2nd list: uncordon_drained_nodes (nodes), + # 3rd list: uncordon_drained_nodes (pods on node) + mock_client.list.side_effect = [ + [pod], # find_topology_failed_pods + [], # uncordon_drained_nodes: no nodes with our annotation + ] + mock_client.get.return_value = node + mock_cls.return_value = mock_client + + assert main() == 0 + # patch called for cordon + assert mock_client.patch.call_count >= 1 + # delete called for failed pod cleanup + assert mock_client.delete.call_count == 1 + + def test_returns_1_on_failure(self): + with ( + patch("numa_cordon.Client") as mock_cls, + patch.dict("os.environ", {}, clear=True), + ): + mock_client = MagicMock() + mock_client.list.side_effect = Exception("boom") + mock_cls.return_value = mock_client + + assert main() == 1