prole/knoe/core/policy.py

from __future__ import annotations

import json
import subprocess
from typing import Sequence

# Central scheduling policy:
# optional_workloads_allowed = ready_schedulable_nodes >= 2

POLICY_CFG_KEY = "OPTIONAL_WORKLOADS_MIN_READY_SCHEDULABLE_NODES"
OPTIONAL_WORKLOADS_MIN_READY_SCHEDULABLE_NODES = 2


def _coerce_int(val: object, default: int) -> int:
    try:
        s = str(val).strip()
        if not s:
            return default
        return int(s)
    except Exception:
        return default


def evaluate_optional_workloads_allowed(
    *,
    kubectl_cmd: Sequence[str] = ("kubectl",),
    env: dict | None = None,
    min_nodes: int = OPTIONAL_WORKLOADS_MIN_READY_SCHEDULABLE_NODES,
    timeout_s: int = 15,
) -> tuple[bool, int, str]:
    """Return (allowed, ready_schedulable_nodes, reason).

    Definition:
    - ready_schedulable_nodes: nodes that are Ready=True and not spec.unschedulable.
    - optional_workloads_allowed: ready_schedulable_nodes >= min_nodes
    """

    min_nodes = max(0, _coerce_int(min_nodes, OPTIONAL_WORKLOADS_MIN_READY_SCHEDULABLE_NODES))

    # If we don't have a kube context to query, default-allow (this keeps unit tests
    # deterministic and avoids blocking optional workloads when kube access isn't configured).
    if min_nodes <= 0:
        return (
            True,
            0,
            "optional_workloads_allowed=True (min_required=0)",
        )

    if env is not None:
        has_ctx = bool(
            (env.get("KUBECONFIG") or "").strip()
            or (env.get("KUBECONTEXT") or "").strip()
            or (env.get("KUBE_CONTEXT") or "").strip()
        )
        if not has_ctx:
            return (
                True,
                0,
                "optional_workloads_allowed=True (no kube context provided; skipping node check)",
            )

    try:
        res = subprocess.run(
            [*kubectl_cmd, "get", "nodes", "-o", "json"],
            env=env,
            capture_output=True,
            text=True,
            timeout=timeout_s,
        )
    except Exception as e:
        default_allow = env is not None
        return (
            default_allow,
            0,
            f"optional_workloads_allowed={default_allow} (failed to query nodes: {e})",
        )

    if res.returncode != 0:
        err = (res.stderr or "").strip() or (res.stdout or "").strip()
        msg = err if err else f"kubectl exited with code {res.returncode}"
        default_allow = env is not None
        return (
            default_allow,
            0,
            f"optional_workloads_allowed={default_allow} (failed to query nodes: {msg})",
        )

    try:
        payload = json.loads(res.stdout or "{}")
    except Exception as e:
        default_allow = env is not None
        return (
            default_allow,
            0,
            f"optional_workloads_allowed={default_allow} (failed to parse kubectl output: {e})",
        )

    ready_schedulable = 0
    for node in (payload.get("items") or []):
        spec = node.get("spec") or {}
        if spec.get("unschedulable") is True:
            continue
        conditions = (node.get("status") or {}).get("conditions") or []
        is_ready = any(
            (c.get("type") == "Ready" and str(c.get("status")).strip() == "True")
            for c in conditions
            if isinstance(c, dict)
        )
        if is_ready:
            ready_schedulable += 1

    allowed = ready_schedulable >= min_nodes
    return (
        allowed,
        ready_schedulable,
        f"optional_workloads_allowed={allowed} (ready_schedulable_nodes={ready_schedulable} min_required={min_nodes})",
    )