atc-agents/api/node_ops.py

"""Live probe + context for topology nodes."""

from __future__ import annotations

import asyncio
from typing import Any

import httpx

from agent_terminal import terminal_log
from lab_context import (
    AIRFLOW_URL,
    DOCKHAND_URL,
    GPU_URL,
    HDFS_NN_URL,
    KAFKA_CONNECT_URL,
    KAFKA_UI_URL,
    LAKEHOUSE_HOST,
    OBJECTSCALE_URL,
    SPARK_UI_URL,
    TRINO_URL,
    collect_databases,
    collect_docker_rack,
    collect_etl,
    collect_gpu_metrics,
    collect_hdfs,
    collect_lakehouse,
    collect_objectscale,
    dockhand_containers,
)
from node_registry import NODE_AGENT, NODE_REGISTRY


async def _log(node_id: str, level: str, phase: str, text: str) -> None:
    await terminal_log(node_id, text, level=level, phase=phase)


async def probe_node(node_id: str) -> dict[str, Any]:
    """Run live probes for a topology node; stream output to node terminal."""
    meta = NODE_REGISTRY.get(node_id)
    if not meta:
        return {"error": "unknown node"}

    await _log(node_id, "info", "shell", f"═══ Connecting to {meta['label']} ({meta['ip']}) ═══")
    await _log(node_id, "cmd", "shell", f"$ probe --node {node_id} --vm {meta['vm']}")

    result: dict[str, Any] = {"node_id": node_id, "ok": True}

    async with httpx.AsyncClient(timeout=10.0, verify=False) as client:
        if node_id == "airflow":
            etl = await collect_etl(client)
            result["data"] = etl
            healthy = etl.get("airflow_healthy")
            await _log(node_id, "ok" if healthy else "warn", "shell", f"Airflow scheduler: {'HEALTHY' if healthy else 'DEGRADED'}")
            for comp, st in (etl.get("airflow_components") or {}).items():
                await _log(node_id, "info", "shell", f"  · {comp}: {st}")

        elif node_id == "db":
            raw = await dockhand_containers(client, 5)
            db = await collect_databases(client, raw)
            result["data"] = db
            await _log(node_id, "ok", "shell", f"DB vault: {db['running']}/{db['total']} containers up")
            for engine, items in (db.get("by_engine") or {}).items():
                await _log(node_id, "info", "shell", f"  {engine}:")
                for item in items:
                    await _log(node_id, "info", "shell", f"    - {item}")

        elif node_id == "debezium":
            etl = await collect_etl(client)
            result["data"] = {"connectors": etl.get("connectors")}
            await _log(node_id, "ok", "shell", f"Kafka Connect @ {KAFKA_CONNECT_URL}")
            for c in etl.get("connectors") or []:
                await _log(node_id, "info", "shell", f"  ✓ {c}")

        elif node_id == "kafka":
            etl = await collect_etl(client)
            result["data"] = {"kafka_ui_ok": etl.get("kafka_ui_ok")}
            await _log(node_id, "ok" if etl.get("kafka_ui_ok") else "warn", "shell", f"Kafka UI {KAFKA_UI_URL}: {'UP' if etl.get('kafka_ui_ok') else 'DOWN'}")
            await _log(node_id, "info", "shell", f"  Broker: 10.0.21.36:9092")

        elif node_id == "lakehouse":
            raw = await dockhand_containers(client, 9)
            lh = await collect_lakehouse(client, raw)
            result["data"] = lh
            await _log(node_id, "ok", "shell", f"Lakehouse {lh['host']}: {lh['running']}/{lh['total']} containers")
            await _log(node_id, "info", "shell", f"  Trino {TRINO_URL}: {'UP' if lh.get('trino_ok') else 'DOWN'}")
            for c in lh.get("containers") or []:
                ports = ",".join(c.get("ports") or []) or "internal"
                await _log(node_id, "info", "shell", f"  · {c['name']}: {c['state']} ports={ports}")

        elif node_id == "s3":
            os_data = await collect_objectscale(client)
            raw = await dockhand_containers(client, 9)
            consumer = next((c for c in raw if "s3-kafka" in f"{c.get('name', '')} {c.get('image', '')}".lower()), None)
            result["data"] = {"objectscale": os_data, "consumer": consumer}
            await _log(node_id, "ok" if os_data.get("reachable") else "warn", "shell", f"ObjectScale {OBJECTSCALE_URL}: {'UP' if os_data.get('reachable') else 'DOWN'}")
            await _log(node_id, "info", "shell", f"  Bucket: data @ {os_data.get('host')}:{os_data.get('port')}")
            if consumer:
                await _log(node_id, "info", "shell", f"  s3-kafka-consumer: {consumer.get('state')}")

        elif node_id == "docker":
            raw = await dockhand_containers(client, 1)
            dk = await collect_docker_rack(client, raw)
            result["data"] = dk
            await _log(node_id, "ok", "shell", f"Docker rack: {dk['running']}/{dk['total']} running")
            for c in dk.get("containers") or []:
                ports = ",".join(c.get("ports") or []) or "internal"
                lvl = "info" if c.get("state") == "running" else "warn"
                await _log(node_id, lvl, "shell", f"  · {c['name']}: {c['state']} ports={ports}")

        elif node_id == "hadoop":
            hdfs = await collect_hdfs(client)
            result["data"] = hdfs
            if hdfs.get("reachable"):
                await _log(node_id, "ok", "shell", f"NameNode {HDFS_NN_URL}: UP")
                await _log(node_id, "info", "shell", f"  Capacity: {hdfs.get('capacity_used_gb')}GB / {hdfs.get('capacity_total_gb')}GB")
                await _log(node_id, "info", "shell", f"  DataNodes: {hdfs.get('live_datanodes')} live, RF=3")
                for dn in hdfs.get("datanodes") or []:
                    await _log(node_id, "info", "shell", f"    · {dn['host']}: {dn['used_gb']}GB used, {dn['blocks']} blocks")
            else:
                await _log(node_id, "err", "shell", "NameNode unreachable")

        elif node_id == "gpu":
            gpu = await collect_gpu_metrics(client)
            result["data"] = gpu
            await _log(node_id, "ok" if gpu.get("ok") else "warn", "shell", f"GPU Lab {GPU_URL}")
            await _log(node_id, "info", "shell", f"  Model: {gpu.get('active_model')} inference={'ON' if gpu.get('inference_active') else 'OFF'}")
            for g in gpu.get("gpus") or []:
                await _log(node_id, "info", "shell", f"  GPU{g['index']}: {g['util_gpu']:.0f}% VRAM {g['memory_used_mib']:.0f}/{g['memory_total_mib']:.0f} MiB")

        elif node_id == "command":
            result["data"] = {"agents": 9, "url": "http://10.0.21.33"}
            await _log(node_id, "ok", "shell", "Command Center online — 9 agents ready")
            await _log(node_id, "info", "shell", "  API: http://10.0.21.33/api")
            await _log(node_id, "info", "shell", "  WebSocket: /api/ws/ops")

        elif node_id in ("mo-commander", "bart-commander", "mcp-coordinator", "network-watcher"):
            meta = NODE_REGISTRY[node_id]
            await _log(node_id, "ok", "shell", f"{meta['label']} online — monitoring all agent comms")
            await _log(node_id, "info", "shell", meta.get("description", ""))
            result["data"] = {"role": meta.get("role")}

    await _log(node_id, "ok", "shell", "═══ Probe complete — type a question below ═══")
    return result


def build_node_detail(node_id: str, snap: dict[str, Any], workload_node: dict | None = None) -> dict[str, Any]:
    """Rich context payload for a single node."""
    meta = dict(NODE_REGISTRY.get(node_id, {}))
    if not meta:
        return {"error": "unknown node"}

    wn = workload_node or {}
    agent_id = NODE_AGENT.get(node_id, "infra-sentinel")

    if node_id == "gpu":
        try:
            from gpu_config import resolve_gpu_identity
            gid = resolve_gpu_identity(snap.get("gpu") if isinstance(snap, dict) else None)
            meta["ip"] = gid["ip"]
            meta["vm"] = gid["vm"]
            meta["vmid"] = gid["vmid"]
            meta["ssh"] = f"ssh root@{gid['ip']}"
            meta["links"] = [
                {"label": "GPU Lab UI", "url": gid["ui_url"]},
                {"label": "vLLM API", "url": gid["llm_url"]},
            ]
            meta["endpoints"] = [
                {"name": "gpu-lab", "host": gid["ip"], "port": "9000", "proto": "http"},
                {"name": "vllm", "host": gid["ip"], "port": "8001", "proto": "http"},
            ]
        except Exception:
            pass

    detail: dict[str, Any] = {
        "id": node_id,
        "agent_id": agent_id,
        **meta,
        "level": wn.get("level", "unknown"),
        "running": wn.get("running", 0),
        "total": wn.get("total", 0),
        "apps": wn.get("apps", []),
        "connectors": wn.get("connectors"),
        "bucket": wn.get("bucket") or meta.get("bucket"),
        "port": wn.get("port"),
        "model": wn.get("model"),
        "util": wn.get("util"),
        "hdfs_used_gb": wn.get("hdfs_used_gb"),
        "hdfs_total_gb": wn.get("hdfs_total_gb"),
        "trino_ok": wn.get("trino_ok"),
        "consumer_ok": wn.get("consumer_ok"),
    }

    edges = (snap.get("_edges") or []) if False else []
    _ = edges  # reserved for future edge context from workload

    return detail


async def run_node_probe_task(node_id: str) -> None:
    try:
        await probe_node(node_id)
    except Exception as exc:
        await _log(node_id, "err", "shell", f"Probe failed: {exc}")