boardroom-map/orchestrator/serving.py

"""Bring up model serving on the Sparks over SSH, in WAVES.

Boardroom Map serves a user-defined catalog of local models. Two Sparks can't hold
unlimited distinct models, so the job runner loads them in waves: each wave brings
up the vLLM container(s) that fit (<= maxConcurrentModels on the head Spark),
points a LiteLLM router at them, lets the reviewers assigned to those models run,
then tears the wave down and loads the next.

Network topology (the confidentiality boundary):
  * A per-job user-defined Docker network on the head Spark (default
    "boardroom-net"). In `airgapped` mode it is created with --internal, so the
    reviewer containers attached to it can reach the model proxy but have ZERO
    internet egress. In `local_services` mode it is a normal bridge (egress
    possible) so reviewers can also reach LAN services / the second Spark.
  * Head-Spark vLLMs join this network; the proxy reaches them by container name
    (bm-vllm-<alias>). Reviewers reach the proxy by name (boardroom-proxy).
  * Second-Spark vLLMs publish a host port; the proxy reaches them over the LAN.
    This only works in local_services mode (an --internal network can't route to
    the LAN), so airgapped jobs must keep all models on the head Spark — enforced
    in preflight.

Models are served from a pre-populated HF cache mounted from the Spark work dir,
so airgapped serving needs no live download.
"""
from __future__ import annotations

import json
import shlex

import spark_client as sc

PROXY_NAME = "boardroom-proxy"


def net_name(cfg: dict) -> str:
    return cfg.get("networkName") or "boardroom-net"


def _vllm_name(alias: str) -> str:
    return f"bm-vllm-{alias}"


def _hf_cache(cfg: dict) -> str:
    return f"{cfg['remoteWorkDir'].rstrip('/')}/hf-cache"


# ------------------------------------------------------------------ network
def ensure_network(cfg: dict, log) -> None:
    head = sc.head(cfg)
    name = net_name(cfg)
    internal = "--internal " if cfg.get("networkMode") == "airgapped" else ""
    cmd = (
        f"docker network inspect {shlex.quote(name)} >/dev/null 2>&1 && echo EXISTS || "
        f"docker network create {internal}{shlex.quote(name)}"
    )
    r = sc.run(head, cmd, timeout=60)
    if r.returncode != 0:
        raise RuntimeError(f"could not create docker network {name}: {r.stderr or r.stdout}")
    mode = "airgapped (--internal)" if internal else "local_services"
    log(f"[serving] network {name} ready ({mode})")


def remove_network(cfg: dict, log) -> None:
    head = sc.head(cfg)
    sc.run(head, f"docker network rm {shlex.quote(net_name(cfg))} 2>/dev/null; true", timeout=30)


# ------------------------------------------------------------------ vLLM
def _vllm_run(cfg: dict, model: dict, hf_token: str | None) -> tuple[sc.Spark, str]:
    """Build the docker run command for one model on its assigned Spark."""
    alias, hf, port = model["alias"], model["hfModel"], int(model["port"])
    image = cfg["servingImage"]
    gpu_util = cfg["gpuMemoryUtilization"]
    max_len = int(cfg["maxModelLen"])
    parser = cfg.get("toolCallParser", "hermes")
    tools = (f"--enable-auto-tool-choice --tool-call-parser {shlex.quote(parser)} " if parser else "")
    env = f"-e HF_TOKEN={shlex.quote(hf_token)} " if hf_token else ""
    cache = _hf_cache(cfg)
    name = _vllm_name(alias)

    if model.get("spark") == "secondary":
        target = sc.by_role(cfg, "secondary")
        # 2nd Spark: publish the port so the head's proxy can reach it over the LAN.
        net = f"-p {port}:{port}"
    else:
        target = sc.head(cfg)
        # Head Spark: attach to the per-job network; proxy reaches it by name.
        net = f"--network {shlex.quote(net_name(cfg))}"

    cmd = (
        f"mkdir -p {shlex.quote(cache)}; "
        f"docker rm -f {name} >/dev/null 2>&1; "
        f"docker run -d --name {name} --gpus all --ipc=host --shm-size=16g "
        f"--restart unless-stopped {net} {env}"
        f"-v {shlex.quote(cache)}:/root/.cache/huggingface "
        f"{shlex.quote(image)} "
        f"vllm serve {shlex.quote(hf)} --host 0.0.0.0 --port {port} "
        f"--gpu-memory-utilization {gpu_util} --max-model-len {max_len} {tools}"
        f"--served-model-name {shlex.quote(alias)}"
    )
    return target, cmd


def _litellm_config(cfg: dict, wave: list[dict]) -> dict:
    """Router config exposing each model alias in the wave on one endpoint."""
    model_list = []
    for m in wave:
        alias, port = m["alias"], int(m["port"])
        if m.get("spark") == "secondary":
            api_base = f"http://{cfg['secondarySparkHost']}:{port}/v1"
        else:
            api_base = f"http://{_vllm_name(alias)}:{port}/v1"
        model_list.append({
            "model_name": alias,
            "litellm_params": {
                "model": f"openai/{alias}",
                "api_base": api_base,
                "api_key": "sk-local",
            },
        })
    return {"model_list": model_list, "general_settings": {"master_key": "sk-local"}}


def bring_up_wave(cfg: dict, wave: list[dict], hf_token: str | None, log) -> None:
    """Start the vLLMs for `wave` and a router that exposes them. Idempotent."""
    head = sc.head(cfg)
    for m in wave:
        target, cmd = _vllm_run(cfg, m, hf_token)
        log(f"[serving] launching {m['alias']} ({m['hfModel']}) on {target.host}:{m['port']}")
        r = sc.run(target, cmd, timeout=240)
        if r.returncode != 0:
            raise RuntimeError(f"vLLM launch for {m['alias']} failed: {r.stderr or r.stdout}")

    # LiteLLM router on the head Spark, attached to the per-job network so the
    # reviewers reach it by name (boardroom-proxy). No published port — preflight
    # probes it from inside the network.
    cfg_json = json.dumps(_litellm_config(cfg, wave))
    workdir = cfg["remoteWorkDir"]
    remote_cfg = f"{workdir}/litellm.config.json"
    proxy_port = int(cfg["proxyPort"])
    log(f"[serving] launching router {PROXY_NAME} on network {net_name(cfg)}:{proxy_port}")
    setup = (
        f"mkdir -p {shlex.quote(workdir)} && "
        f"printf '%s' {shlex.quote(cfg_json)} > {shlex.quote(remote_cfg)} && "
        f"docker rm -f {PROXY_NAME} >/dev/null 2>&1; "
        f"docker run -d --name {PROXY_NAME} --restart unless-stopped "
        f"--network {shlex.quote(net_name(cfg))} "
        f"-v {shlex.quote(remote_cfg)}:/app/config.json "
        f"ghcr.io/berriai/litellm:main-stable "
        f"--config /app/config.json --port {proxy_port} --host 0.0.0.0"
    )
    r = sc.run(head, setup, timeout=180)
    if r.returncode != 0:
        raise RuntimeError(f"LiteLLM router launch failed: {r.stderr or r.stdout}")


def tear_down_wave(cfg: dict, wave: list[dict], log) -> None:
    names = " ".join(_vllm_name(m["alias"]) for m in wave)
    # vLLMs may be split across both Sparks; clear the names on each.
    for sp in sc.sparks(cfg):
        sc.run(sp, f"docker rm -f {names} 2>/dev/null; true", timeout=120)
    sc.run(sc.head(cfg), f"docker rm -f {PROXY_NAME} 2>/dev/null; true", timeout=60)
    log(f"[serving] wave torn down ({names})")


def tear_down_all(cfg: dict, log) -> None:
    """Best-effort: remove every Boardroom Map serving container + the network."""
    for sp in sc.sparks(cfg):
        sc.run(sp, "docker ps -aq --filter name=bm-vllm- | xargs -r docker rm -f; "
                   "docker ps -aq --filter name=bm-grader- | xargs -r docker rm -f; "
                   f"docker rm -f {PROXY_NAME} 2>/dev/null; true", timeout=120)
    remove_network(cfg, log)
    log("[serving] all serving torn down")


def plan_waves(cfg: dict, needed_aliases: set[str]) -> list[list[dict]]:
    """Group the needed models into waves that respect per-Spark concurrency.

    Head-Spark models are chunked into groups of `maxConcurrentModels`; secondary
    models likewise. Wave i runs head-group-i and secondary-group-i together (they
    sit on different GPUs)."""
    catalog = {m["alias"]: m for m in (cfg.get("models") or [])}
    cap = max(1, int(cfg.get("maxConcurrentModels", 1)))
    primary, secondary = [], []
    for alias in sorted(needed_aliases):
        m = catalog.get(alias)
        if not m:
            continue
        (secondary if m.get("spark") == "secondary" else primary).append(m)

    def chunk(lst):
        return [lst[i:i + cap] for i in range(0, len(lst), cap)]

    pg, sg = chunk(primary), chunk(secondary)
    waves = []
    for i in range(max(len(pg), len(sg))):
        wave = (pg[i] if i < len(pg) else []) + (sg[i] if i < len(sg) else [])
        if wave:
            waves.append(wave)
    return waves


def reviewer_proxy_base(cfg: dict) -> str:
    """The OpenAI-compatible base URL reviewers use (by container name on the net)."""
    return f"http://{PROXY_NAME}:{int(cfg['proxyPort'])}/v1"


def health(cfg: dict) -> dict:
    head = sc.head(cfg)
    r = sc.run(head, "docker ps --filter name=bm-vllm- --filter name=boardroom-proxy "
                     "--format '{{.Names}} {{.Status}}'", timeout=30)
    return {"running": (r.stdout or "").strip().splitlines()}