modelbeast/server/remote.py

"""Remote GPU worker dispatch (the unified worker pool).

The gpu lane is a pool of NODES: the local Mac plus any remote workers listed in
nodes.json. A remote node runs the SAME operator code (a git-synced checkout) with
repo-relative paths. The primary never lets a remote touch its DB — it rsyncs the
inputs out, runs the operator's run.py over ssh, and rsyncs the outputs back, then
registers them locally exactly as for a local job.

nodes.json (repo root, gitignored — present only on the primary), e.g.:
  [{"name": "m1", "ssh": "johnking@100.91.239.7", "root": "/Users/johnking/MODELBEAST",
    "operators": ["flux_local","mflux_image_edit","seedvr2_upscale","bg_remove_local","sf3d"]}]
Omit "operators" to allow every gpu op. The implicit local node is always present.
"""
import asyncio
import json
import shlex
import time
from pathlib import Path

from . import db

SSH = ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=8"]
_health: dict[str, tuple[float, bool]] = {}  # ssh -> (checked_at, ok)
_HEALTH_TTL = 30.0


def load_gpu_nodes() -> list[dict]:
    """Local node first, then remote workers from nodes.json."""
    nodes = [{"name": "local", "local": True, "busy": False}]
    f = db.ROOT / "nodes.json"
    if f.exists():
        try:
            for n in json.loads(f.read_text()):
                n["busy"] = False
                n["local"] = False
                nodes.append(n)
        except (ValueError, KeyError):
            pass
    return nodes


def node_supports(node: dict, op_id: str) -> bool:
    if node.get("local"):
        return True
    allow = node.get("operators")
    return (not allow) or (op_id in allow)


async def _run(cmd: list[str], timeout: float = 300) -> tuple[int, str]:
    proc = await asyncio.create_subprocess_exec(
        *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT)
    try:
        out, _ = await asyncio.wait_for(proc.communicate(), timeout)
    except asyncio.TimeoutError:
        proc.kill()
        return 124, "timeout"
    return proc.returncode, (out or b"").decode(errors="replace")


async def healthy(node: dict) -> bool:
    """Reachable + checkout present. Cached for _HEALTH_TTL to avoid per-job ssh."""
    if node.get("local"):
        return True
    key = node["ssh"]
    hit = _health.get(key)
    if hit and time.time() - hit[0] < _HEALTH_TTL:
        return hit[1]
    code, _ = await _run(SSH + [node["ssh"], f"test -d {shlex.quote(node['root'])}"], timeout=10)
    ok = code == 0
    _health[key] = (time.time(), ok)
    return ok


async def prepare(node: dict, job_id: str, local_inputs: list[str]) -> list[str]:
    """Make remote dirs + rsync inputs. Returns the remote input paths."""
    root = node["root"]
    rjob = f"{root}/data/remote/{job_id}"
    await _run(SSH + [node["ssh"], f"mkdir -p {shlex.quote(rjob)}/inputs {shlex.quote(rjob)}/out"])
    remote_inputs = []
    for lp in local_inputs:
        name = Path(lp).name
        dest = f"{node['ssh']}:{rjob}/inputs/{name}"
        code, out = await _run(["rsync", "-a", "-e", "ssh -o BatchMode=yes",
                                lp.rstrip("/") + ("/" if Path(lp).is_dir() else ""), dest], timeout=600)
        if code != 0:
            raise RuntimeError(f"rsync input failed: {out[:200]}")
        remote_inputs.append(f"{rjob}/inputs/{name}")
    return remote_inputs


def remote_cmd(node: dict, op_id: str, op_python_rel: str, entry_rel: str,
               job_id: str, params_json: str, remote_inputs: list[str]) -> list[str]:
    """ssh command that runs the operator on the remote node with repo-relative paths.
    HF_TOKEN etc. are sourced from the node's .env.remote (kept off the process table)."""
    root = node["root"]
    rout = f"{root}/data/remote/{job_id}/out"
    inner = (f"cd {shlex.quote(root)} && set -a; [ -f .env.remote ] && . ./.env.remote; set +a; "
             f"exec {shlex.quote(op_python_rel)} {shlex.quote(entry_rel)} "
             f"--outdir {shlex.quote(rout)} --params {shlex.quote(params_json)}")
    for ri in remote_inputs:
        inner += f" --input {shlex.quote(ri)}"
    return SSH + ["-T", node["ssh"], f"bash -lc {shlex.quote(inner)}"]


async def collect(node: dict, job_id: str, local_outdir: Path) -> None:
    """rsync the remote outputs back into the local job outdir, then clean up remote."""
    root = node["root"]
    rjob = f"{root}/data/remote/{job_id}"
    code, out = await _run(["rsync", "-a", "-e", "ssh -o BatchMode=yes",
                            f"{node['ssh']}:{rjob}/out/", f"{str(local_outdir).rstrip('/')}/"], timeout=600)
    if code != 0:
        raise RuntimeError(f"rsync output failed: {out[:200]}")
    await _run(SSH + [node["ssh"], f"rm -rf {shlex.quote(rjob)}"], timeout=30)