setup_cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""aquin setup machine — prepare this host as a (remote) compute box for desktop/CLI."""
 
from __future__ import annotations
 
import json
import os
import platform
import socket
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
 
 
def _utc_now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
 
 
def _gpu_summary() -> dict[str, Any]:
    out: dict[str, Any] = {"accelerator": "cpu", "detail": None}
    try:
        import torch
 
        if torch.cuda.is_available():
            out["accelerator"] = "cuda"
            out["detail"] = torch.cuda.get_device_name(0)
            out["device_count"] = torch.cuda.device_count()
        elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
            out["accelerator"] = "mps"
            out["detail"] = "Apple Metal"
        else:
            out["accelerator"] = "cpu"
            out["detail"] = "torch installed, no CUDA/MPS"
        out["torch"] = getattr(torch, "__version__", "?")
    except Exception as exc:
        out["accelerator"] = "unknown"
        out["detail"] = f"torch unavailable: {exc}"
    return out
 
 
def _sshd_listening() -> dict[str, Any]:
    """Best-effort: is something listening on port 22? (OpenSSH usually)."""
    result: dict[str, Any] = {"port_22_open": None, "notes": []}
    try:
        with socket.create_connection(("127.0.0.1", 22), timeout=0.4):
            result["port_22_open"] = True
    except OSError:
        result["port_22_open"] = False
        result["notes"].append(
            "Nothing on localhost:22 — install/start OpenSSH (sshd) if this box should accept SSH."
        )
 
    # Optional service name hints
    for cmd in (
        ["systemctl", "is-active", "ssh"],
        ["systemctl", "is-active", "sshd"],
        ["service", "ssh", "status"],
    ):
        try:
            r = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=2,
            )
            text = (r.stdout or r.stderr or "").strip().lower()
            if "active" in text or r.returncode == 0:
                result["service_hint"] = " ".join(cmd)
                break
        except Exception:
            continue
    return result
 
 
def _local_ips() -> list[str]:
    ips: list[str] = []
    try:
        hostname = socket.gethostname()
        for info in socket.getaddrinfo(hostname, None):
            ip = info[4][0]
            if ip and ip not in ips and not ip.startswith("127."):
                ips.append(ip)
    except Exception:
        pass
    return ips[:8]
 
 
def _ensure_aquin_home() -> Path:
    home = Path(os.environ.get("AQUIN_HOME") or (Path.home() / ".aquin"))
    home.mkdir(parents=True, exist_ok=True)
    return home
 
 
def _start_daemon(port: int, wait_s: float = 40.0) -> dict[str, Any]:
    from aquin.engine import model_daemon
    from aquin.engine.engine_info import resolve_engine_port
 
    if port:
        os.environ["AQUIN_ENGINE_PORT"] = str(port)
    effective = port or resolve_engine_port()
    already = model_daemon.is_running()
    if already:
        h = model_daemon.health() or {}
        return {
            "ok": True,
            "already": True,
            "port": effective,
            "health": h,
        }
    ok = model_daemon.ensure_running(wait_seconds=wait_s)
    h = model_daemon.health() if ok else None
    return {
        "ok": bool(ok),
        "already": False,
        "port": effective,
        "health": h or {},
        "error": None if ok else "daemon did not become healthy (see ~/.aquin/daemon.log)",
    }
 
 
def _write_machine_record(payload: dict[str, Any]) -> Path:
    home = _ensure_aquin_home()
    path = home / "machine-setup.json"
    path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
    return path
 
 
def _tunnel_block(port: int, user: str, host: str) -> str:
    target = f"{user}@{host}" if user and host else "USER@HOST"
    return "\n".join(
        [
            "# ── Desktop / laptop connection ─────────────────────────────────",
            "# On your laptop (keep tunnel open):",
            f"ssh -N -L {port}:127.0.0.1:{port} {target}",
            "# or from monorepo:",
            f"./ide/app/scripts/ssh-engine-tunnel.sh {target}",
            "",
            "# Then start Aquin desktop with remote engine:",
            f"export AQUIN_ENGINE_REMOTE=1",
            f"export AQUIN_ENGINE_HOST=127.0.0.1",
            f"export AQUIN_ENGINE_PORT={port}",
            "cd ide/app && npm start",
            "",
            "# Engine stays on *this* machine (127.0.0.1 only). Tunnel makes it look local.",
            "# Optional: aquin load model <id>  here  to put a model in VRAM.",
        ]
    )
 
 
def cmd_setup(args: list[str]) -> None:
    """
    aquin setup machine  — prepare this host for local or remote (desktop tunnel) compute.
    """
    if not args or args[0] in ("--help", "-h", "help"):
        print("Usage: aquin setup machine [options]")
        print("")
        print("  Probe GPU + SSH readiness, ensure ~/.aquin, start the model daemon")
        print("  (localhost only), and print SSH tunnel + desktop env for remote GPU.")
        print("")
        print("Options:")
        print("  --port N        Daemon port (default: 17832 / AQUIN_ENGINE_PORT)")
        print("  --no-daemon     Do not start the model daemon")
        print("  --check         Report only (no write / no start)")
        print("  --json          Machine-readable summary on stdout")
        print("  --ssh-user U    Hint for tunnel command (default: $USER)")
        print("  --ssh-host H    Hint for tunnel command (default: hostname)")
        return
 
    if args[0] != "machine":
        print(f"Unknown setup target: {args[0]}")
        print("Usage: aquin setup machine")
        sys.exit(1)
 
    rest = args[1:]
    port = 0
    no_daemon = False
    check_only = False
    as_json = False
    ssh_user = (os.environ.get("USER") or os.environ.get("LOGNAME") or "").strip() or "user"
    ssh_host = socket.gethostname()
 
    i = 0
    while i < len(rest):
        a = rest[i]
        if a in ("--help", "-h", "help"):
            cmd_setup([])
            return
        if a == "--port" and i + 1 < len(rest):
            port = int(rest[i + 1])
            i += 2
            continue
        if a.startswith("--port="):
            port = int(a.split("=", 1)[1])
            i += 1
            continue
        if a == "--no-daemon":
            no_daemon = True
            i += 1
            continue
        if a == "--check":
            check_only = True
            i += 1
            continue
        if a == "--json":
            as_json = True
            i += 1
            continue
        if a == "--ssh-user" and i + 1 < len(rest):
            ssh_user = rest[i + 1]
            i += 2
            continue
        if a == "--ssh-host" and i + 1 < len(rest):
            ssh_host = rest[i + 1]
            i += 2
            continue
        print(f"Unknown flag: {a}", file=sys.stderr)
        sys.exit(1)
 
    from aquin.engine.engine_info import DEFAULT_ENGINE_PORT, resolve_engine_port
 
    if not port:
        port = resolve_engine_port() or DEFAULT_ENGINE_PORT
 
    summary: dict[str, Any] = {
        "command": "setup machine",
        "check_only": check_only,
        "when": _utc_now(),
        "hostname": socket.gethostname(),
        "platform": {
            "system": platform.system(),
            "release": platform.release(),
            "machine": platform.machine(),
            "python": sys.version.split()[0],
            "executable": sys.executable,
        },
        "gpu": _gpu_summary(),
        "ssh": _sshd_listening(),
        "network": {"local_ips": _local_ips()},
        "port": port,
        "daemon": None,
        "ok": True,
        "next": {},
    }
 
    try:
        import aquin  # noqa: F401
 
        summary["aquin_import"] = "ok"
    except Exception as exc:
        summary["aquin_import"] = f"failed: {exc}"
        summary["ok"] = False
 
    if not no_daemon and not check_only and summary.get("aquin_import") == "ok":
        daemon_info = _start_daemon(port)
        summary["daemon"] = daemon_info
        if not daemon_info.get("ok"):
            summary["ok"] = False
    elif not no_daemon:
        # check path: probe only
        try:
            from aquin.engine import model_daemon
 
            h = model_daemon.health()
            summary["daemon"] = {
                "ok": h is not None,
                "already": h is not None,
                "port": port,
                "health": h or {},
            }
        except Exception as exc:
            summary["daemon"] = {"ok": False, "error": str(exc), "port": port}
 
    tunnel = _tunnel_block(port, ssh_user, ssh_host)
    summary["next"] = {
        "ssh_target": f"{ssh_user}@{ssh_host}",
        "tunnel": f"ssh -N -L {port}:127.0.0.1:{port} {ssh_user}@{ssh_host}",
        "desktop_env": {
            "AQUIN_ENGINE_REMOTE": "1",
            "AQUIN_ENGINE_HOST": "127.0.0.1",
            "AQUIN_ENGINE_PORT": str(port),
        },
        "load_model": "aquin load model <id>",
    }
 
    if not check_only:
        try:
            path = _write_machine_record(
                {
                    "setup_at": summary["when"],
                    "hostname": summary["hostname"],
                    "port": port,
                    "ssh_target": summary["next"]["ssh_target"],
                    "gpu": summary["gpu"],
                    "daemon_ok": (summary.get("daemon") or {}).get("ok"),
                }
            )
            summary["wrote"] = str(path)
        except Exception as exc:
            summary["wrote_error"] = str(exc)
 
    if as_json:
        print(json.dumps(summary, indent=2, ensure_ascii=False))
        sys.exit(0 if summary.get("ok") else 1)
 
    # Human output — agent-friendly, greppable markers
    print("[setup machine] host probes")
    print(f"  hostname   : {summary['hostname']}")
    print(f"  platform   : {summary['platform']['system']} {summary['platform']['machine']}")
    print(f"  python     : {summary['platform']['python']} ({summary['platform']['executable']})")
    print(f"  aquin      : {summary.get('aquin_import')}")
    gpu = summary["gpu"]
    print(f"  accelerator: {gpu.get('accelerator')} {gpu.get('detail') or ''}".rstrip())
    ssh = summary["ssh"]
    print(
        f"  ssh :22    : "
        f"{'open' if ssh.get('port_22_open') else 'closed/unavailable' if ssh.get('port_22_open') is False else 'unknown'}"
    )
    for n in ssh.get("notes") or []:
        print(f"    note: {n}")
 
    d = summary.get("daemon") or {}
    if no_daemon:
        print("[setup machine] daemon: skipped (--no-daemon)")
    elif check_only:
        print(f"[setup machine] daemon check: ok={d.get('ok')} port={d.get('port', port)}")
    else:
        if d.get("ok"):
            state = "already running" if d.get("already") else "started"
            print(f"[setup machine] daemon: {state} on 127.0.0.1:{d.get('port', port)}")
            h = d.get("health") or {}
            if h.get("model_id"):
                print(f"  model_id   : {h.get('model_id')}")
        else:
            print(f"[setup machine] daemon: FAILED — {d.get('error') or 'unknown'}", file=sys.stderr)
 
    if summary.get("wrote"):
        print(f"[setup machine] wrote {summary['wrote']}")
 
    print("")
    print(tunnel)
    print("")
    if summary.get("ok"):
        print("[setup machine] OK — engine ready for SSH tunnel from your laptop.")
    else:
        print("[setup machine] completed with issues (see above).", file=sys.stderr)
        sys.exit(1)