llmgw/internal/preflight/preflight.go

// Package preflight checks machines before anything is installed on them:
// that they can be logged into and administered, run a supported system, have
// room for their role, keep time, have the ports free, and can reach the
// internet and each other. Everything wrong is found at once, with what to
// do about it, rather than one failure at a time halfway through an install.
package preflight

import (
	"context"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
)

// Role is what a machine is going to be.
type Role string

// The roles a setup gives machines.
const (
	RoleControl  Role = "control plane"
	RoleDatabase Role = "database"
	RoleWitness  Role = "etcd witness"
	RoleNode     Role = "node"
)

// Need is the least a role asks of a machine.
type Need struct {
	CPUs     int
	MemoryMB int
	DiskGB   int
	Ports    []int
}

// Needs are what each role asks for; a machine with several roles needs the
// largest of each and the ports of all.
var Needs = map[Role]Need{
	RoleControl:  {CPUs: 2, MemoryMB: 3500, DiskGB: 15, Ports: []int{7443, 8080}},
	RoleDatabase: {CPUs: 2, MemoryMB: 3500, DiskGB: 30, Ports: []int{5432, 8008, 2379, 2380}},
	RoleWitness:  {CPUs: 1, MemoryMB: 900, DiskGB: 5, Ports: []int{2379, 2380}},
	RoleNode:     {CPUs: 2, MemoryMB: 3500, DiskGB: 20},
}

// Machine is one machine to check.
type Machine struct {
	Name   string // how the screen calls it
	Addr   string // the address the others will use
	Roles  []Role
	Remote onboard.Remote
	// Local is this machine, where sudo is not needed (setup runs as root).
	Local bool
	// ServesMB is memory it needs beyond its roles' for the models it will
	// serve itself, and Serves names them. Short of it is a failure, not a
	// warning: the install would finish and then have nowhere to run them.
	ServesMB int
	Serves   string
}

// Level is how bad a finding is.
type Level int

// Levels, mildest first.
const (
	OK Level = iota
	Warn
	Fail
)

func (l Level) String() string { return [...]string{"ok", "warn", "fail"}[l] }

// Finding is one thing checked on one machine.
type Finding struct {
	Machine string
	Check   string
	Level   Level
	Detail  string
	// Fix says what to do, for anything that is not OK.
	Fix string
}

// Report is everything found.
type Report struct {
	Findings []Finding
}

// Worst is the most serious level found.
func (r Report) Worst() Level {
	w := OK
	for _, f := range r.Findings {
		if f.Level > w {
			w = f.Level
		}
	}
	return w
}

// For is one machine's findings, in check order.
func (r Report) For(machine string) []Finding {
	var out []Finding
	for _, f := range r.Findings {
		if f.Machine == machine {
			out = append(out, f)
		}
	}
	return out
}

// probeScript gathers everything in one round trip, as key=value lines, so
// the checks are decided here and can be tested without a machine.
const probeScript = `
. /etc/os-release 2>/dev/null; echo "os=$ID"; echo "version=$VERSION_ID"; echo "pretty=$PRETTY_NAME"
echo "arch=$(uname -m)"
echo "cpus=$(nproc)"
echo "mem_mb=$(awk '/MemTotal/{print int($2/1024)}' /proc/meminfo)"
echo "disk_gb=$(df -BG --output=avail / | tail -1 | tr -dc 0-9)"
echo "now=$(date +%s.%N)"
echo "ntp=$(timedatectl show -p NTPSynchronized --value 2>/dev/null)"
echo "listening=$(ss -Hltn 2>/dev/null | awk '{n=split($4,a,":"); print a[n]}' | sort -un | tr '\n' ' ')"
if sudo -n true 2>/dev/null; then echo "sudo=yes"; else echo "sudo=no"; fi
echo "apt=$(command -v apt-get >/dev/null && echo yes || echo no)"
echo "internet=$(timeout 6 bash -c 'exec 3<>/dev/tcp/deb.debian.org/80' 2>/dev/null && echo yes || (command -v curl >/dev/null && curl -s -m 6 -o /dev/null http://deb.debian.org && echo yes) || echo no)"
`

// reachScript tries one TCP port on each address: whether the machines can
// reach one another at all. Port 22 because it is open everywhere before
// anything is installed; the ports llmgw will use have nothing listening yet.
func reachScript(addrs []string) string {
	var b strings.Builder
	for _, a := range addrs {
		fmt.Fprintf(&b, "if timeout 4 bash -c 'exec 3<>/dev/tcp/%s/22' 2>/dev/null; then echo \"reach %s yes\"; else echo \"reach %s no\"; fi\n", a, a, a)
	}
	return b.String()
}

// Run checks every machine, concurrently.
func Run(ctx context.Context, machines []Machine) Report {
	var (
		mu  sync.Mutex
		rep Report
		wg  sync.WaitGroup
	)
	add := func(fs ...Finding) {
		mu.Lock()
		rep.Findings = append(rep.Findings, fs...)
		mu.Unlock()
	}
	var addrs []string
	for _, m := range machines {
		addrs = append(addrs, m.Addr)
	}
	order := map[string]int{}
	for i, m := range machines {
		order[m.Name] = i
		wg.Add(1)
		go func(m Machine) {
			defer wg.Done()
			add(check(ctx, m, addrs)...)
		}(m)
	}
	wg.Wait()
	// Machines in the order given, checks in the order made.
	sort.SliceStable(rep.Findings, func(i, j int) bool {
		return order[rep.Findings[i].Machine] < order[rep.Findings[j].Machine]
	})
	return rep
}

func check(ctx context.Context, m Machine, all []string) []Finding {
	sent := time.Now()
	out, err := m.Remote.Run(ctx, probeScript, nil)
	took := time.Since(sent)
	if err != nil {
		return []Finding{{Machine: m.Name, Check: "login", Level: Fail,
			Detail: strings.TrimSpace(firstLine(err.Error() + " " + string(out))),
			Fix:    "check the address, the user and the key or password given for it"}}
	}
	facts := parse(out)
	fs := Evaluate(m, facts, sent.Add(took/2))
	var others []string
	for _, a := range all {
		if a != m.Addr {
			others = append(others, a)
		}
	}
	if len(others) > 0 {
		rout, _ := m.Remote.Run(ctx, reachScript(others), nil)
		fs = append(fs, evaluateReach(m, string(rout))...)
	}
	return fs
}

func firstLine(s string) string {
	s = strings.TrimSpace(s)
	if i := strings.IndexByte(s, '\n'); i >= 0 {
		return s[:i]
	}
	return s
}

func parse(out []byte) map[string]string {
	facts := map[string]string{}
	for _, line := range strings.Split(string(out), "\n") {
		if k, v, ok := strings.Cut(strings.TrimSpace(line), "="); ok {
			facts[k] = strings.TrimSpace(v)
		}
	}
	return facts
}

// supported are the systems the setup installs packages on with apt.
var supported = map[string][]string{
	"ubuntu": {"22.04", "24.04"},
	"debian": {"12"},
}

// Evaluate turns one machine's facts into findings. at is when the machine's
// clock was read, by this machine's clock.
func Evaluate(m Machine, facts map[string]string, at time.Time) []Finding {
	var fs []Finding
	f := func(check string, l Level, detail, fix string) {
		fs = append(fs, Finding{Machine: m.Name, Check: check, Level: l, Detail: detail, Fix: fix})
	}
	f("login", OK, "logged in", "")

	// The system.
	os, ver := facts["os"], facts["version"]
	switch vers, ok := supported[os]; {
	case !ok:
		f("system", Fail, facts["pretty"], "use Ubuntu 22.04 or 24.04, or Debian 12")
	case !contains(vers, ver):
		f("system", Warn, facts["pretty"]+" (tried on "+strings.Join(vers, ", ")+")", "it may work; Ubuntu 24.04 is what this is tested on")
	default:
		f("system", OK, facts["pretty"], "")
	}
	switch facts["arch"] {
	case "x86_64", "aarch64":
		f("architecture", OK, facts["arch"], "")
	default:
		f("architecture", Fail, facts["arch"], "llmgw is built for amd64 and arm64")
	}

	// Administering it.
	if !m.Local {
		if facts["sudo"] == "yes" {
			f("sudo", OK, "without a password", "")
		} else {
			// Not a failure: the setup asks for the password and grants
			// itself sudo for the length of the install.
			f("sudo", Warn, "asks for a password", "give the sudo password on the access screen")
		}
	}
	if facts["apt"] != "yes" {
		f("packages", Fail, "no apt-get", "the setup installs PostgreSQL, Patroni and etcd with apt")
	}

	// Room for its roles.
	need := needOf(m.Roles)
	cpus, _ := strconv.Atoi(facts["cpus"])
	mem, _ := strconv.Atoi(facts["mem_mb"])
	disk, _ := strconv.Atoi(facts["disk_gb"])
	room := func(check string, have, want int, unit string) {
		d := fmt.Sprintf("%d %s (needs %d)", have, unit, want)
		switch {
		case have >= want:
			f(check, OK, d, "")
		case have*2 >= want:
			f(check, Warn, d, "it will run, short of what the role is sized for")
		default:
			f(check, Fail, d, "give the machine more")
		}
	}
	room("cpus", cpus, need.CPUs, "CPUs")
	if m.ServesMB > 0 {
		want := need.MemoryMB + m.ServesMB
		d := fmt.Sprintf("%d MB (needs %d with %s)", mem, want, m.Serves)
		if mem >= want {
			f("memory", OK, d, "")
		} else {
			f("memory", Fail, d, "give it more memory, or turn off "+m.Serves+" or This machine as a node")
		}
	} else {
		room("memory", mem, need.MemoryMB, "MB")
	}
	room("disk", disk, need.DiskGB, "GB free")

	// Time. Leases and etcd's elections go wrong on clocks that disagree.
	if now, err := strconv.ParseFloat(facts["now"], 64); err == nil {
		skew := time.Duration((now - float64(at.UnixNano())/1e9) * float64(time.Second))
		if skew < 0 {
			skew = -skew
		}
		switch {
		case skew > 5*time.Second:
			f("clock", Fail, fmt.Sprintf("%s off from this machine", skew.Round(time.Second)), "turn on NTP (timedatectl set-ntp true) on both")
		case skew > time.Second:
			f("clock", Warn, fmt.Sprintf("%s off from this machine", skew.Round(100*time.Millisecond)), "turn on NTP on both")
		default:
			f("clock", OK, "within a second of this machine", "")
		}
	}
	if facts["ntp"] != "yes" {
		f("time sync", Warn, "NTP is not synchronised", "timedatectl set-ntp true")
	}

	// Ports its roles will listen on.
	listening := map[string]bool{}
	for _, p := range strings.Fields(facts["listening"]) {
		listening[p] = true
	}
	var busy []string
	for _, p := range need.Ports {
		if listening[strconv.Itoa(p)] {
			busy = append(busy, strconv.Itoa(p))
		}
	}
	if len(busy) > 0 {
		// A warning: a re-run finds its own services listening.
		f("ports", Warn, "already listening: "+strings.Join(busy, ", "),
			"fine if it is llmgw's own from an earlier run; otherwise stop what holds them")
	} else if len(need.Ports) > 0 {
		f("ports", OK, "free", "")
	}

	if facts["internet"] == "yes" {
		f("internet", OK, "reaches the package mirrors", "")
	} else {
		f("internet", Fail, "cannot reach deb.debian.org", "packages and releases are downloaded; allow outbound HTTP/HTTPS or set a proxy")
	}
	return fs
}

func evaluateReach(m Machine, out string) []Finding {
	var no []string
	for _, line := range strings.Split(out, "\n") {
		parts := strings.Fields(line)
		if len(parts) == 3 && parts[0] == "reach" && parts[2] == "no" {
			no = append(no, parts[1])
		}
	}
	if len(no) > 0 {
		return []Finding{{Machine: m.Name, Check: "network", Level: Fail,
			Detail: "cannot reach " + strings.Join(no, ", "),
			Fix:    "the machines must reach one another; check routes and firewalls"}}
	}
	return []Finding{{Machine: m.Name, Check: "network", Level: OK, Detail: "reaches every other machine"}}
}

func needOf(roles []Role) Need {
	var n Need
	seen := map[int]bool{}
	for _, r := range roles {
		x := Needs[r]
		n.CPUs, n.MemoryMB, n.DiskGB = max(n.CPUs, x.CPUs), max(n.MemoryMB, x.MemoryMB), max(n.DiskGB, x.DiskGB)
		for _, p := range x.Ports {
			if !seen[p] {
				seen[p] = true
				n.Ports = append(n.Ports, p)
			}
		}
	}
	return n
}

func contains(xs []string, x string) bool {
	for _, y := range xs {
		if y == x {
			return true
		}
	}
	return false
}