api/load_balancer.py

"""
backend/api/load_balancer.py — Dynamic Load Balancer (S951)
Analizza il carico dei nodi (A, B, C, D) e decide a chi assegnare il task.
Supporta il Cross-Node Offloading: se il nodo target è saturo, delega al nodo più scarico.
M1-FIX: timeout Redis configurabile via env REDIS_TIMEOUT_S (default 5.0s, allineato a job_queue e state_sync).
"""
import json
import logging
import os
import time
from typing import Optional, Dict, Any, List
import httpx

_logger = logging.getLogger("agente_ai.balancer")

# M1-FIX: timeout Redis allineato agli altri moduli (job_queue: 5.0s, state_sync: 5.0s)
_REDIS_TIMEOUT = float(os.getenv("REDIS_TIMEOUT_S", "5.0"))

# Ruoli del cluster
ROLES = ["brain", "hands", "memory", "audit"]

class LoadBalancer:
    def __init__(self):
        self.redis_url = os.getenv("UPSTASH_REDIS_REST_URL")
        self.redis_token = os.getenv("UPSTASH_REDIS_REST_TOKEN")
        self.internal_token = os.getenv("INTERNAL_TOKEN", "")

    async def _rcmd(self, cmd: List[Any]) -> Optional[Dict[str, Any]]:
        if not self.redis_url or not self.redis_token: return None
        try:
            async with httpx.AsyncClient() as client:
                r = await client.post(
                    self.redis_url,
                    headers={"Authorization": f"Bearer {self.redis_token}"},
                    json=cmd,
                    timeout=_REDIS_TIMEOUT,  # M1-FIX: era hardcodato 2.0
                )
                return r.json()
        except Exception as e:
            _logger.error("[S951] Redis error: %s", e)
            return None

    async def get_node_load(self, role: str) -> Optional[Dict[str, Any]]:
        """Recupera le metriche di carico per un nodo specifico."""
        res = await self._rcmd(["GET", f"jq:load:{role}"])
        if res and res.get("result"):
            try:
                return json.loads(res["result"])
            except (json.JSONDecodeError, TypeError):
                return None
        return None

    async def get_best_node(self, preferred_role: str = "hands") -> str:
        """
        Determina il miglior nodo per un task.
        Se il preferred_role è sotto soglia di carico, lo mantiene.
        Altrimenti cerca il nodo con meno active_tasks.
        """
        pref_load = await self.get_node_load(preferred_role)
        if pref_load:
            active = pref_load.get("active_tasks", 0)
            if active < 3:
                return preferred_role

        _logger.info("[S951] Nodo %s saturo o offline. Ricerca miglior nodo alternativo...", preferred_role)

        best_role = preferred_role
        min_load = 999

        for role in ROLES:
            load = await self.get_node_load(role)
            if load:
                active = load.get("active_tasks", 0)
                ts = load.get("ts", 0)
                if (time.time() * 1000) - ts < 90000:
                    if active < min_load:
                        min_load = active
                        best_role = role

        if best_role != preferred_role:
            _logger.info("[S951] Offloading dinamico: %s -> %s (carico: %d)", preferred_role, best_role, min_load)

        return best_role

    async def get_cluster_status(self) -> Dict[str, Any]:
        """Ritorna lo stato di salute di tutto il cluster."""
        status = {}
        for role in ROLES:
            load = await self.get_node_load(role)
            if load:
                ts = load.get("ts", 0)
                is_online = (time.time() * 1000) - ts < 90000
                status[role] = {
                    "online":       is_online,
                    "active_tasks": load.get("active_tasks", 0),
                    "mem_mb":       load.get("mem", 0),
                    "last_seen":    ts,
                }
            else:
                status[role] = {"online": False}
        return status

# Singleton instance
balancer = LoadBalancer()