atc-agents/workload.py

"""Build UI workload payload from lab snapshot."""

from __future__ import annotations

from typing import Any


def _level(running: int, total: int) -> str:
    if total == 0:
        return "unknown"
    ratio = running / total
    if ratio >= 0.9:
        return "ok"
    if ratio >= 0.5:
        return "warn"
    return "down"


def _app_row(c: dict[str, Any]) -> dict[str, Any]:
    img = c.get("image") or ""
    short_img = img.split("/")[-1].split(":")[0][:20]
    return {
        "name": c.get("name", "?"),
        "state": c.get("state", "unknown"),
        "image": short_img,
        "ports": c.get("ports") or [],
    }


def build_workload_payload(snap: dict[str, Any]) -> dict[str, Any]:
    docker = snap.get("docker", {})
    databases = snap.get("databases", {})
    lakehouse = snap.get("lakehouse", {})
    etl = snap.get("etl", {})
    hadoop = snap.get("hadoop", {})
    gpu = snap.get("gpu", {})

    docker_apps = [_app_row(c) for c in docker.get("containers", [])]
    db_apps = [_app_row(c) for c in databases.get("containers", [])]
    lake_apps = [_app_row(c) for c in lakehouse.get("containers", [])]

    hdfs_ok = hadoop.get("reachable", False)
    etl_ok = etl.get("airflow_healthy") and etl.get("kafka_ui_ok")

    zones = [
        {
            "id": "docker",
            "label": "DOCKER RACK",
            "x": 8,
            "color": "#b366ff",
            "level": _level(docker.get("running", 0), docker.get("total", 1) or 1),
            "running": docker.get("running", 0),
            "total": docker.get("total", 0),
            "apps": docker_apps,
        },
        {
            "id": "db",
            "label": "DB VAULT",
            "x": 28,
            "color": "#ffaa00",
            "level": _level(databases.get("running", 0), databases.get("total", 1) or 1),
            "running": databases.get("running", 0),
            "total": databases.get("total", 0),
            "apps": db_apps,
        },
        {
            "id": "lakehouse",
            "label": "LAKEHOUSE HUB",
            "x": 50,
            "color": "#ff00aa",
            "level": _level(lakehouse.get("running", 0), lakehouse.get("total", 1) or 1),
            "running": lakehouse.get("running", 0),
            "total": lakehouse.get("total", 0),
            "apps": lake_apps,
            "trino_ok": lakehouse.get("trino_ok"),
        },
        {
            "id": "hadoop",
            "label": "HADOOP CLUSTER",
            "x": 72,
            "color": "#39ff14",
            "level": "ok" if hdfs_ok else "warn",
            "running": hadoop.get("live_datanodes", 0),
            "total": (hadoop.get("live_datanodes") or 0) + (hadoop.get("dead_datanodes") or 0),
            "apps": [
                {"name": "NameNode", "state": "running" if hdfs_ok else "down", "image": "hdfs-nn", "ports": ["9870"]},
                *[
                    {"name": dn.get("host", "?").split(".")[0], "state": "running", "image": "datanode", "ports": ["9866"]}
                    for dn in hadoop.get("datanodes", [])
                ],
            ],
            "hdfs_used_gb": hadoop.get("capacity_used_gb"),
            "hdfs_total_gb": hadoop.get("capacity_total_gb"),
        },
        {
            "id": "etl",
            "label": "ETL PIPE",
            "x": 92,
            "color": "#00f0ff",
            "level": "ok" if etl_ok else "warn",
            "running": sum(1 for s in [
                etl.get("airflow_healthy"),
                etl.get("kafka_ui_ok"),
                etl.get("spark_ui_ok"),
            ] if s),
            "total": 3,
            "apps": [
                {"name": "Airflow", "state": "running" if etl.get("airflow_healthy") else "down", "image": "airflow", "ports": ["8080"]},
                {"name": "Kafka UI", "state": "running" if etl.get("kafka_ui_ok") else "down", "image": "kafka", "ports": ["9000"]},
                {"name": "Spark UI", "state": "running" if etl.get("spark_ui_ok") else "down", "image": "spark", "ports": ["8080"]},
                *[
                    {"name": c, "state": "running", "image": "connect", "ports": ["8083"]}
                    for c in etl.get("connectors", [])
                ],
            ],
        },
    ]

    return {
        "ts": snap.get("ts"),
        "zones": zones,
        "gpu": {
            "level": "ok" if gpu.get("ok") and gpu.get("inference_active") else ("warn" if gpu.get("ok") else "down"),
            "model": gpu.get("active_model"),
            "inference_active": gpu.get("inference_active"),
            "gpu_count": gpu.get("gpu_count", 0),
            "avg_util": round(
                sum(g.get("util_gpu", 0) for g in gpu.get("gpus", [])) / max(len(gpu.get("gpus", [])), 1),
                1,
            ),
            "gpus": gpu.get("gpus", []),
        },
        "totals": {
            "apps_running": sum(z["running"] for z in zones if z["id"] != "hadoop") + (hadoop.get("live_datanodes") or 0),
            "apps_total": sum(z["total"] for z in zones),
            "connectors": len(etl.get("connectors", [])),
        },
    }