code/federation/agillm41_lease_decide.py

#!/usr/bin/env python3
"""Adaptive AGILLM4.x lease sizer.

Prints three fields for shell callers:
    <batch> <block> <max_layers>

Batch/block size affect activation memory and useful tokens. max_layers controls
package size and resident model-slice memory, which is the important knob for
DiffusionBlock layer cycling on small machines.
"""
from __future__ import annotations

import json
import os
import re
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

STATE = Path(os.environ.get("AGILLM41_LEASE_STATE", "/workspace/agillm41_lease_state.json"))
MASTER_LOG = os.environ.get("AGILLM41_MASTER_LOG", "/workspace/agillm41_master_train.log")
GK = os.environ.get("AGILLM41_GETH_KEY", "/root/.ssh/agillm41_geth_ed25519")
GH = os.environ.get("AGILLM41_GETH_HOST", "root@5.75.217.57")
OPP = os.environ.get("AGILLM41_OPPORTUNISTIC_ROOT", "/root/agillm41_opportunistic")
FAILURE_TTL_SEC = int(os.environ.get("AGILLM41_FAILURE_TTL_SEC", "1800"))

GPU_VRAM = {
    "h100": 80, "a100-80": 80, "a100": 40, "l40": 48, "a6000": 48,
    "a40": 48, "v100-pcie-32": 32, "v100-sxm2-32": 32, "v100": 16,
    "rtx 5090": 32, "5090": 32, "rtx 4090": 24, "4090": 24,
    "rtx 3090": 24, "3090": 24, "a10": 24, "rtx 4080": 16,
    "4080": 16, "t4": 16, "rtx 4070": 12, "3060": 12,
    "rtx 3080": 10, "quadro m620": 2, "m620": 2,
}

# Trusted core. CPU RAM is intentionally authoritative for these names; stale
# opportunistic GPU heartbeats must not resize them.
TRUSTED_CPU: dict[str, dict[str, Any]] = {
    "geth": {"ram_gb": 30, "batch": 1, "block": 512, "max_layers": 4, "cap_layers": 7, "cap_batch": 2},
    "mcp": {"ram_gb": 3, "batch": 1, "block": 128, "max_layers": 1, "cap_layers": 2, "cap_batch": 1},
    "prime": {"ram_gb": 3, "batch": 1, "block": 128, "max_layers": 2, "cap_layers": 2, "cap_batch": 1},
    "communist-web": {"ram_gb": 3, "batch": 1, "block": 128, "max_layers": 2, "cap_layers": 2, "cap_batch": 1},
}

TRUSTED_LAPTOP: dict[str, dict[str, Any]] = {
    "laptop-cuda": {"batch": 1, "block": 128, "max_layers": 1, "cap_layers": 1, "vram": 2},
    "laptop-cpu": {"batch": 1, "block": 128, "max_layers": 1, "cap_layers": 2, "ram_gb": 16},
    "laptop-igpu": {"batch": 1, "block": 96, "max_layers": 1, "cap_layers": 1, "vram": 0.5},
}

GPU_DEFAULT_BLOCK = int(os.environ.get("AGILLM41_LEASE_BLOCK", "1300"))


def read_state() -> dict[str, Any]:
    if not STATE.exists():
        return {}
    try:
        return json.loads(STATE.read_text())
    except Exception:
        return {}


def write_state(st: dict[str, Any]) -> None:
    tmp = STATE.with_suffix(STATE.suffix + ".tmp")
    tmp.write_text(json.dumps(st, indent=2, sort_keys=True))
    tmp.replace(STATE)


def vram_for(gpu: str) -> float:
    g = (gpu or "").lower()
    for key, val in GPU_VRAM.items():
        if key in g:
            return float(val)
    return 16.0 if g else 0.0


def gpu_profile(vram: float) -> dict[str, Any]:
    if vram >= 70:
        return {"batch": 2, "block": GPU_DEFAULT_BLOCK, "max_layers": 4, "cap_batch": 4, "cap_layers": 7}
    if vram >= 40:
        return {"batch": 2, "block": GPU_DEFAULT_BLOCK, "max_layers": 3, "cap_batch": 3, "cap_layers": 5}
    if vram >= 30:
        return {"batch": 1, "block": GPU_DEFAULT_BLOCK, "max_layers": 3, "cap_batch": 2, "cap_layers": 4}
    if vram >= 20:
        return {"batch": 1, "block": 768, "max_layers": 2, "cap_batch": 1, "cap_layers": 3}
    if vram >= 10:
        return {"batch": 1, "block": 512, "max_layers": 1, "cap_batch": 1, "cap_layers": 2}
    if vram > 0:
        return {"batch": 1, "block": 128, "max_layers": 1, "cap_batch": 1, "cap_layers": 1}
    return {"batch": 1, "block": 128, "max_layers": 1, "cap_batch": 1, "cap_layers": 1}


def ssh_geth(cmd: str, timeout: int = 15) -> str:
    try:
        return subprocess.run(
            [
                "ssh", "-i", GK, "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=no",
                "-o", "ConnectTimeout=8", GH, cmd,
            ],
            capture_output=True,
            text=True,
            timeout=timeout,
        ).stdout
    except Exception:
        return ""


def heartbeat(worker: str) -> dict[str, Any]:
    out = ssh_geth(f"cat {OPP}/heartbeats/{worker}.json 2>/dev/null", timeout=15)
    try:
        return json.loads(out)
    except Exception:
        return {}


def latest_master_tokps(worker: str) -> float | None:
    try:
        tail = subprocess.run(["tail", "-n", "8000", MASTER_LOG], capture_output=True, text=True, timeout=10).stdout
    except Exception:
        return None
    best = None
    for line in tail.splitlines():
        if "async_side_update_applied" not in line or worker not in line:
            continue
        try:
            data = json.loads(line[line.index("{"):])
        except Exception:
            continue
        if data.get("worker_id") == worker and data.get("tok_per_sec") is not None:
            best = float(data["tok_per_sec"])
    return best


def latest_worker_tokps(worker: str) -> float | None:
    # Immediate side-worker logs live on GETH before the master has consumed the update.
    safe = re.sub(r"[^A-Za-z0-9_.-]", "", worker)
    out = ssh_geth(
        "grep -h '\"tok_per_sec\"' /root/agillm41_worker/logs/*_"
        + safe
        + "_* 2>/dev/null | tail -1",
        timeout=15,
    )
    m = re.search(r"([-+]?[0-9]*\.?[0-9]+)", out)
    if m:
        try:
            return float(m.group(1))
        except Exception:
            pass
    return latest_master_tokps(worker)



def parse_utc_ts(value: Any) -> float | None:
    if not value:
        return None
    if isinstance(value, (int, float)):
        return float(value)
    text = str(value).strip().replace("Z", "+00:00")
    try:
        return datetime.fromisoformat(text).astimezone(timezone.utc).timestamp()
    except Exception:
        return None


def fresh_event(ts: float | None) -> bool:
    return ts is not None and (time.time() - ts) <= FAILURE_TTL_SEC

def recent_failure(worker: str) -> str:
    hb = heartbeat(worker)
    hb_state = str(hb.get("state") or hb.get("status") or "").lower()
    hb_ts = parse_utc_ts(hb.get("at") or hb.get("updated_at") or hb.get("ts"))
    if hb_state in {"error", "failed"} and fresh_event(hb_ts):
        return str(hb.get("error") or hb.get("err") or "heartbeat_failed")[:160]
    safe = re.sub(r"[^A-Za-z0-9_.-]", "", worker)
    out = ssh_geth(
        "find /root/agillm41_worker/logs -maxdepth 1 -type f -mmin -120 -name '*_" + safe + "_*' "
        "-exec tail -80 {} + 2>/dev/null | grep -Ei 'out of memory|killed|traceback|runtimeerror|failed' | tail -1",
        timeout=15,
    ).strip()
    return out[:160]


def base_profile(worker: str) -> dict[str, Any]:
    if worker in TRUSTED_CPU:
        return {"tier": "trusted-cpu", **TRUSTED_CPU[worker]}
    if worker in TRUSTED_LAPTOP:
        return {"tier": "trusted-laptop", **TRUSTED_LAPTOP[worker]}
    hb = heartbeat(worker)
    gpu = str(hb.get("gpu") or hb.get("device_name") or "")
    # Friendly aliases for opportunistic lanes whose heartbeat may be stale/missing.
    if not gpu and "v100" in worker.lower():
        gpu = "Tesla V100-PCIE-32GB"
    if not gpu and any(x in worker.lower() for x in ("gpu", "cuda", "4090", "3090")):
        gpu = "RTX 4090"
    vram = float(hb.get("vram_gb") or hb.get("vram") or vram_for(gpu) or 0.0)
    prof = gpu_profile(vram)
    return {"tier": "gpu" if vram else "unknown", "gpu": gpu, "vram": vram, **prof}


def clamp(v: int, lo: int, hi: int) -> int:
    return max(lo, min(hi, int(v)))


def decide(worker: str) -> tuple[int, int, int]:
    st = read_state()
    prev = st.get(worker, {}) if isinstance(st.get(worker), dict) else {}
    prof = base_profile(worker)
    tokps = latest_worker_tokps(worker)
    fail = recent_failure(worker)

    batch = int(prev.get("batch") or prof.get("batch", 1))
    block = int(prev.get("block") or prof.get("block", 128))
    layers = int(prev.get("max_layers") or prof.get("max_layers", 1))
    cap_batch = int(prof.get("cap_batch", batch))
    cap_layers = int(prof.get("cap_layers", layers))
    floor_block = 64 if prof.get("tier") == "trusted-laptop" else 128

    prev_tokps = prev.get("tokps")
    decision_tokps = prev.get("decision_tokps")
    failure_seen = prev.get("failure_seen")
    new_failure = bool(fail) and fail != failure_seen
    new_measurement = tokps is not None and repr(tokps) != repr(decision_tokps)

    if new_failure:
        batch = 1
        block = max(floor_block, block // 2)
        layers = max(1, layers - 1)
    elif not fail and prev.get("failure_seen") and tokps is None:
        # Recover from stale/cleared failures by returning to the base profile.
        batch = int(prof.get("batch", batch))
        block = int(prof.get("block", block))
        layers = int(prof.get("max_layers", layers))
    elif new_measurement:
        if prev_tokps is None:
            # First clean observation: grow the layer window before batch. That
            # improves useful model coverage without exploding activation memory.
            layers = min(cap_layers, layers + (1 if cap_layers > layers else 0))
        else:
            try:
                good = float(tokps) >= float(prev_tokps) * 0.92
            except Exception:
                good = True
            if good:
                if layers < cap_layers:
                    layers += 1
                elif batch < cap_batch:
                    batch += 1
            else:
                if batch > int(prof.get("batch", 1)):
                    batch -= 1
                else:
                    layers = max(1, layers - 1)

    batch = clamp(batch, 1, cap_batch)
    block = clamp(block, floor_block, int(prof.get("block_cap", prof.get("block", block))))
    layers = clamp(layers, 1, cap_layers)

    rec = {
        "batch": batch,
        "block": block,
        "max_layers": layers,
        "tokps": tokps,
        "prev_tokps": prev_tokps,
        "failure": fail,
        "failure_seen": fail or failure_seen,
        "decision_tokps": tokps if new_measurement else decision_tokps,
        "profile": prof,
        "ts": time.time(),
    }
    st[worker] = rec
    write_state(st)
    return batch, block, layers


if __name__ == "__main__":
    wid = sys.argv[1] if len(sys.argv) > 1 else "laptop-cuda"
    b, blk, ml = decide(wid)
    print(f"{b} {blk} {ml}")