esh-pfi-infrastructure/services/lora-training-worker/worker/gpu.py

"""GPU status for arbo's device-aware scheduler (GET /gpu-status).

arbo steers a lean train OFF the 3090 when TTS is live there (the 3090/TTS co-OOM concern,
contract §7). The worker can't decide that — it just reports the raw per-device VRAM plus a
`tts_on_3090` boolean derived from nvidia-smi compute-apps. Everything is nvidia-smi CSV
parsing; no torch import.

Device indices are reported under CUDA_DEVICE_ORDER=PCI_BUS_ID (index 0 = RTX 3090,
index 1 = RTX A6000) — the same ordering the training recipe pins, so `device_index` in a
/train request and the indices here mean the same physical card.
"""

from __future__ import annotations

import subprocess

from . import config

_SMI_ENV = {"CUDA_DEVICE_ORDER": "PCI_BUS_ID"}


def _run(args: list[str]) -> str:
    return subprocess.run(
        ["nvidia-smi", *args], capture_output=True, text=True, timeout=15, env={**_SMI_ENV}
    ).stdout


def _query_devices() -> list[dict]:
    # index,name,memory.total,memory.used,memory.free (MiB) under PCI_BUS_ID ordering
    out = _run(["--query-gpu=index,name,memory.total,memory.used,memory.free",
                "--format=csv,noheader,nounits"])
    devices = []
    for line in out.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) != 5:
            continue
        idx, name, total, used, free = parts
        devices.append({
            "index": int(idx), "name": name,
            "total_mb": int(total), "used_mb": int(used), "free_mb": int(free),
            "procs": [],
        })
    return devices


def _query_compute_apps() -> list[dict]:
    # pid,used_memory,gpu_bus_id — map each compute app to a device by bus id
    out = _run(["--query-compute-apps=pid,used_memory,gpu_bus_id", "--format=csv,noheader,nounits"])
    apps = []
    for line in out.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) != 3:
            continue
        pid, mem, bus = parts
        if not pid.isdigit():
            continue
        apps.append({"pid": int(pid), "used_mb": int(mem) if mem.isdigit() else 0, "bus_id": bus})
    return apps


def _bus_id_by_index() -> dict[int, str]:
    out = _run(["--query-gpu=index,gpu_bus_id", "--format=csv,noheader"])
    mapping = {}
    for line in out.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) == 2 and parts[0].isdigit():
            mapping[int(parts[0])] = parts[1]
    return mapping


def _cmdline(pid: int) -> str:
    try:
        with open(f"/proc/{pid}/cmdline", "rb") as fh:
            return fh.read().replace(b"\x00", b" ").decode("utf-8", "replace").lower()
    except (OSError, ValueError):
        return ""


def gpu_status() -> dict:
    """Return `{devices: [...], tts_on_3090: bool}`. Degrades to an error field on nvidia-smi failure."""
    try:
        devices = _query_devices()
        apps = _query_compute_apps()
        bus_by_idx = _bus_id_by_index()
    except (subprocess.SubprocessError, OSError, ValueError) as exc:
        return {"devices": [], "tts_on_3090": False, "error": f"nvidia-smi failed: {exc}"}

    idx_by_bus = {bus: idx for idx, bus in bus_by_idx.items()}
    dev_by_idx = {d["index"]: d for d in devices}

    # Attach compute apps (with a marker flag) to their device.
    for app in apps:
        idx = idx_by_bus.get(app["bus_id"])
        cmd = _cmdline(app["pid"])
        app["is_tts"] = any(m in cmd for m in config.TTS_CMDLINE_MARKERS)
        if idx is not None and idx in dev_by_idx:
            dev_by_idx[idx]["procs"].append(
                {"pid": app["pid"], "used_mb": app["used_mb"], "is_tts": app["is_tts"]}
            )

    # tts_on_3090 — the co-OOM-risk signal for arbo's device-aware steering. Honest derivation
    # (transparent via `reason`), because the worker can't reliably fingerprint TTS: the TTS
    # containers run generic cmdlines (`uvicorn app:app`, `python app.py`) so only the container
    # NAME identifies them, and the worker's llmuser isn't in the docker group. So:
    #   (a) proc_marker — a compute app on device 0 whose cmdline matches a TTS marker (best effort), else
    #   (b) used_mb_floor — device 0 used >= floor (a busy 3090 = lean-train co-OOM risk, TTS/STT/audio alike).
    # The PRECISE lever arbo should key on is `devices[0].free_mb`; `tts_on_3090` is the convenience boolean.
    dev0 = dev_by_idx.get(config.DEVICE_3090)
    tts_on_3090, reason = False, "idle"
    free_3090 = None
    if dev0 is not None:
        free_3090 = dev0["free_mb"]
        if any(p["is_tts"] for p in dev0["procs"]):
            tts_on_3090, reason = True, "proc_marker"
        elif dev0["used_mb"] >= config.TTS_3090_USED_MB_FLOOR:
            tts_on_3090, reason = True, "used_mb_floor"
    return {
        "devices": devices,
        "tts_on_3090": tts_on_3090,
        "tts_on_3090_reason": reason,
        "gpu3090_free_mb": free_3090,  # the precise co-OOM lever; steer lean off the 3090 when low
    }