sskj/src/sskj/deploy/runtime.py

"""Docker/native server lifecycle for deployment profiles.

Supports single-node and multi-node (NODE_HOSTS) docker deployments. For
multi-node profiles, start/stop/status distribute docker commands to each
node over ssh; health checks only target the head (rank 0) node.
"""

from __future__ import annotations

import base64
import re
import shlex
import subprocess
import time
from pathlib import Path

from sskj.bench.runner import split_env_assignments, wait_health

_NODE_RANK_RE = re.compile(r"\$\{NODE_RANK(?::-([^}]*))?\}")
_TEMPLATED_KEYS = (
    "CONTAINER_NAME",
    "LAUNCH_ARGS",
    "DEVICE_VARS",
    "ENGINE_ENV",
    "DOCKER_FLAGS",
    "VOLUMES",
    "PATCH_MOUNTS",
    "BOOTSTRAP",
)


def _log_dir(root: Path, log_dir: str | None) -> Path:
    path = Path(log_dir) if log_dir else root / "deploy" / "logs"
    path.mkdir(parents=True, exist_ok=True)
    return path


def _health_url(profile: dict[str, str]) -> str:
    host = profile.get("HEALTH_HOST", "127.0.0.1")
    return f"http://{host}:{profile['PORT']}"


def _node_hosts(profile: dict[str, str]) -> list[str]:
    return [h for h in shlex.split(profile.get("NODE_HOSTS", "")) if h]


def _ssh_host(profile: dict[str, str], host: str) -> str:
    user = profile.get("NODE_SSH_USER", "root")
    if not user or "@" in host:
        return host
    return f"{user}@{host}"


def _sub_node_rank(profile: dict[str, str], rank: int) -> dict[str, str]:
    """Return a profile copy with NODE_RANK rendered for the given rank."""
    rendered = dict(profile)
    container_name = str(rendered.get("CONTAINER_NAME", "server"))
    has_template = _NODE_RANK_RE.search(container_name) is not None

    def _sub(value: str | None) -> str:
        if not value:
            return value or ""
        return _NODE_RANK_RE.sub(lambda m: m.group(1) or str(rank), str(value))

    for key in _TEMPLATED_KEYS:
        if key in rendered:
            rendered[key] = _sub(rendered[key])
    if not has_template:
        rendered["CONTAINER_NAME"] = f"{container_name}_node{rank}"
    rendered["NODE_RANK"] = str(rank)
    return rendered


def _server_cmd(profile: dict[str, str]) -> str:
    """Launch command prefix (without LAUNCH_ARGS)."""
    if profile.get("SERVER_CMD"):
        return profile["SERVER_CMD"]
    engine = (profile.get("ENGINE") or "sglang").lower()
    python_bin = profile.get("CONTAINER_PYTHON", "python")
    if engine == "vllm":
        return f"{python_bin} -m vllm.entrypoints.openai.api_server"
    return f"{python_bin} -m sglang.launch_server"


def _docker_run_cmd(profile: dict[str, str]) -> list[str]:
    container = profile["CONTAINER_NAME"]
    cmd = ["docker", "run", "-d", "--name", container]
    if profile.get("DOCKER_FLAGS"):
        cmd += shlex.split(profile["DOCKER_FLAGS"])
    if profile.get("NETWORK_MODE") == "bridge" and profile.get("PORT_MAP") == "1":
        cmd += ["-p", f"{profile['PORT']}:{profile['PORT']}"]
    for volume in split_env_assignments(profile.get("VOLUMES", "") + " " + profile.get("PATCH_MOUNTS", "")):
        cmd += ["-v", volume]
    for assignment in split_env_assignments(
        profile.get("DEVICE_VARS", "") + " " + profile.get("ENGINE_ENV", "")
    ):
        cmd += ["-e", assignment]
    cmd.append(profile["DOCKER_IMAGE"])

    bootstrap = profile.get("BOOTSTRAP", "")
    if bootstrap:
        encoded = base64.b64encode(bootstrap.encode("utf-8")).decode("ascii")
        cmd += ["bash", "-c", f"echo {encoded} | base64 -d | bash"]
    else:
        cmd += ["bash", "-c", f"{_server_cmd(profile)} {profile['LAUNCH_ARGS']}"]
    return cmd


def _is_multinode(profile: dict[str, str], node_rank: str | None) -> bool:
    return len(_node_hosts(profile)) > 1 and node_rank is None


def _local_rank(profile: dict[str, str]) -> int:
    """Rank that runs on this host (no ssh); defaults to 0 (head)."""
    try:
        return int(profile.get("LOCAL_NODE_RANK", "0"))
    except ValueError:
        return 0


def _node_target(profile: dict[str, str], rank: int) -> str | None:
    """ssh host for the rank, or None when the rank runs locally."""
    if rank == _local_rank(profile):
        return None
    return _ssh_host(profile, _node_hosts(profile)[rank])


def _run_docker_cmd(
    profile: dict[str, str],
    cmd: list[str],
    host: str | None,
    dry_run: bool,
    log_path: Path | None = None,
) -> int:
    joined = shlex.join(cmd)
    if log_path is not None:
        log_path.write_text(joined + "\n", encoding="utf-8")
    if dry_run:
        print(joined)
        return 0
    if host:
        remote = f"bash -c {shlex.quote(joined)}"
        result = subprocess.run(
            ["ssh", "-o", "StrictHostKeyChecking=no", host, remote],
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            print(result.stdout, result.stderr)
        return result.returncode
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(result.stdout, result.stderr)
    return result.returncode


def _start_docker(
    profile: dict[str, str],
    root: Path,
    dry_run: bool,
    logs: Path,
) -> int:
    if not profile.get("DOCKER_IMAGE"):
        raise SystemExit("docker runtime requires DOCKER_IMAGE")
    if not profile.get("CONTAINER_NAME"):
        raise SystemExit("docker runtime requires CONTAINER_NAME")

    hosts = _node_hosts(profile)
    node_rank = profile.get("NODE_RANK")
    health_wait = int(profile.get("HEALTH_WAIT_S", "600") or 600)

    if _is_multinode(profile, node_rank):
        # Workers first (rank > 0), then head (rank 0), mirroring the legacy
        # multi-node scripts so NCCL bootstrap finds the head up. The
        # LOCAL_NODE_RANK rank runs locally (no ssh).
        order = [r for r in range(len(hosts)) if r != 0] + [0]
        for rank in order:
            host = hosts[rank]
            target = _node_target(profile, rank)
            node_profile = _sub_node_rank(profile, rank)
            cmd = _docker_run_cmd(node_profile)
            if not dry_run:
                if target:
                    subprocess.run(
                        ["ssh", "-o", "StrictHostKeyChecking=no", target, f"docker rm -f {shlex.quote(node_profile['CONTAINER_NAME'])}"],
                        capture_output=True,
                        check=False,
                    )
                else:
                    subprocess.run(
                        ["docker", "rm", "-f", node_profile["CONTAINER_NAME"]],
                        capture_output=True,
                        check=False,
                    )
            print(f"starting container {node_profile['CONTAINER_NAME']} on {host} (rank {rank}) ...")
            rc = _run_docker_cmd(
                profile,
                cmd,
                target,
                dry_run,
                logs / f"{node_profile['CONTAINER_NAME']}.cmd.txt",
            )
            if rc != 0:
                return rc
            if rank != 0:
                print("  sleeping 5s before head ...")
                if not dry_run:
                    time.sleep(5)

        head_url = f"http://{profile.get('HEALTH_HOST', hosts[0])}:{profile['PORT']}"
        if dry_run:
            print(f"[dry-run] health: {head_url}{profile.get('HEALTH_PATH', '/health')}")
            return 0
        if wait_health(head_url, profile.get("HEALTH_PATH", "/health"), health_wait):
            print(f"server on {hosts[0]} is healthy")
            return 0
        print("ERROR: head failed health check; inspect per-node logs")
        return 1

    # Single-node or explicit --node-rank: run locally on this host.
    if not dry_run:
        subprocess.run(["docker", "rm", "-f", profile["CONTAINER_NAME"]], capture_output=True, check=False)

    cmd = _docker_run_cmd(profile)
    container = profile["CONTAINER_NAME"]
    print(f"starting container {container} ...")
    rc = _run_docker_cmd(profile, cmd, None, dry_run, logs / f"{container}.cmd.txt")
    if rc != 0:
        return rc
    if node_rank is not None and node_rank != "0":
        print(f"node rank {node_rank} started; skipping health check")
        return 0
    if dry_run:
        print(f"[dry-run] health: {_health_url(profile)}{profile.get('HEALTH_PATH', '/health')}")
        return 0
    if wait_health(_health_url(profile), profile.get("HEALTH_PATH", "/health"), health_wait):
        print(f"container {container} is healthy")
        return 0
    subprocess.run(["docker", "logs", "--tail", "100", container], check=False)
    print(f"ERROR: container {container} failed health check")
    return 1


def start(profile: dict[str, str], root: Path, dry_run: bool = False, log_dir: str | None = None) -> int:
    validate_runtime(profile)
    logs = _log_dir(root, log_dir)
    if profile["RUNTIME"] == "docker":
        return _start_docker(profile, root, dry_run, logs)

    # native runtime
    python_bin = profile.get("PYTHON_BIN") or profile.get("CONTAINER_PYTHON") or "python3"
    if profile.get("SERVER_CMD"):
        cmd = shlex.split(profile["SERVER_CMD"]) + shlex.split(profile.get("LAUNCH_ARGS", ""))
    else:
        engine = (profile.get("ENGINE") or "sglang").lower()
        module = "vllm.entrypoints.openai.api_server" if engine == "vllm" else "sglang.launch_server"
        cmd = [python_bin, "-m", module, *shlex.split(profile.get("LAUNCH_ARGS", ""))]
    log_path = logs / f"{profile.get('MODEL_NAME', 'server')}.log"
    pid_file = logs / f"{profile.get('MODEL_NAME', 'server')}.pid"
    if dry_run:
        print(shlex.join(cmd))
        return 0
    with open(log_path, "wb") as f:
        proc = subprocess.Popen(cmd, stdout=f, stderr=subprocess.STDOUT)
    pid_file.write_text(str(proc.pid), encoding="utf-8")
    print(f"started native server pid={proc.pid} log={log_path}")
    health_wait = int(profile.get("HEALTH_WAIT_S", "600") or 600)
    if wait_health(_health_url(profile), profile.get("HEALTH_PATH", "/health"), health_wait):
        print("native server is healthy")
        return 0
    print(f"ERROR: native server failed health check; see {log_path}")
    return 1


def _docker_stop(profile: dict[str, str], dry_run: bool = False) -> int:
    node_rank = profile.get("NODE_RANK")
    if _is_multinode(profile, node_rank):
        for rank, host in enumerate(_node_hosts(profile)):
            node_profile = _sub_node_rank(profile, rank)
            name = node_profile["CONTAINER_NAME"]
            target = _node_target(profile, rank)
            if dry_run:
                if target:
                    print(f"[dry-run] ssh {target} docker rm -f {name}")
                else:
                    print(f"[dry-run] docker rm -f {name} (local)")
                continue
            if target:
                subprocess.run(
                    ["ssh", "-o", "StrictHostKeyChecking=no", target, f"docker rm -f {shlex.quote(name)}"],
                    capture_output=True,
                    check=False,
                )
                print(f"container {name} removed on {host}")
            else:
                subprocess.run(["docker", "rm", "-f", name], capture_output=True, check=False)
                print(f"container {name} removed (local)")
        return 0
    container = profile.get("CONTAINER_NAME")
    if not container:
        raise SystemExit("docker runtime requires CONTAINER_NAME")
    if dry_run:
        print(f"[dry-run] docker rm -f {container}")
        return 0
    subprocess.run(["docker", "rm", "-f", container], capture_output=True, check=False)
    print(f"container {container} removed")
    return 0


def stop(profile: dict[str, str], root: Path) -> int:
    validate_runtime(profile)
    if profile["RUNTIME"] == "docker":
        return _docker_stop(profile)
    logs = _log_dir(root, None)
    pid_file = logs / f"{profile.get('MODEL_NAME', 'server')}.pid"
    if pid_file.exists():
        pid = pid_file.read_text(encoding="utf-8").strip()
        if pid:
            subprocess.run(["kill", pid], capture_output=True, check=False)
            print(f"killed native server pid={pid}")
        pid_file.unlink(missing_ok=True)
    return 0


def _docker_status(profile: dict[str, str]) -> int:
    node_rank = profile.get("NODE_RANK")
    if _is_multinode(profile, node_rank):
        all_ok = 0
        for rank, host in enumerate(_node_hosts(profile)):
            node_profile = _sub_node_rank(profile, rank)
            name = node_profile["CONTAINER_NAME"]
            target = _node_target(profile, rank)
            if target:
                result = subprocess.run(
                    ["ssh", "-o", "StrictHostKeyChecking=no", target, f"docker inspect -f '{{{{.State.Status}}}}' {shlex.quote(name)}"],
                    capture_output=True,
                    text=True,
                    check=False,
                )
                state = result.stdout.strip() or "missing"
                print(f"{host}: {name} -> {state}")
                if result.returncode != 0 or state != "running":
                    all_ok = 1
            else:
                result = subprocess.run(
                    ["docker", "inspect", "-f", "{{.State.Status}}", name],
                    capture_output=True,
                    text=True,
                    check=False,
                )
                state = result.stdout.strip() or "missing"
                print(f"{host}: {name} -> {state}")
                if result.returncode != 0 or state != "running":
                    all_ok = 1
        return all_ok
    container = profile.get("CONTAINER_NAME", "")
    result = subprocess.run(
        ["docker", "inspect", "-f", "{{.State.Status}}", container],
        capture_output=True,
        text=True,
        check=False,
    )
    print(result.stdout.strip() or "missing")
    return 0 if result.returncode == 0 else 1


def status(profile: dict[str, str], root: Path) -> int:
    validate_runtime(profile)
    if profile["RUNTIME"] == "docker":
        return _docker_status(profile)
    logs = _log_dir(root, None)
    pid_file = logs / f"{profile.get('MODEL_NAME', 'server')}.pid"
    if not pid_file.exists():
        print("stopped")
        return 1
    pid = pid_file.read_text(encoding="utf-8").strip()
    result = subprocess.run(["kill", "-0", pid], capture_output=True, check=False)
    print("running" if result.returncode == 0 else "stopped")
    return 0 if result.returncode == 0 else 1


def validate_runtime(profile: dict[str, str]) -> None:
    runtime = profile.get("RUNTIME", "docker" if profile.get("DOCKER_IMAGE") else "native")
    if runtime not in ("docker", "native"):
        raise SystemExit(f"unsupported runtime: {runtime}")
    profile["RUNTIME"] = runtime
    if runtime == "docker" and not profile.get("DOCKER_IMAGE"):
        raise SystemExit("docker runtime requires DOCKER_IMAGE")
    if runtime == "native" and not profile.get("PYTHON_BIN") and not profile.get("CONTAINER_PYTHON"):
        raise SystemExit("native runtime requires PYTHON_BIN or CONTAINER_PYTHON")