llmgw/internal/cli/serverinit_cluster.go

package cli

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/sourcecheck"
	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// A control plane on its own serves nothing: it needs a node, weights, and a
// model, and the install guide made each of those a separate afternoon. init
// makes the machine all three — a mirror, a node, and a small model served on
// it — so the first thing anyone sees is an answer. Each is skipped when it is
// already there and can be turned off.

const (
	mirrorDir      = "/srv/models"
	mirrorAPI      = "127.0.0.1:8082"
	mirrorUnitPath = "/etc/systemd/system/llmgw-mirror.service"
	agentUnitPath  = "/etc/systemd/system/llmgw-agent.service"

	// The starter: small enough to download in a minute and run on any CPU,
	// real enough to answer a question. Q4_K_M is about 400 MB.
	starterName    = "qwen2.5-0.5b-instruct"
	starterAlias   = "qwen-0.5b"
	starterSource  = "hf:Qwen/Qwen2.5-0.5B-Instruct-GGUF/qwen2.5-0.5b-instruct-q4_k_m.gguf"
	starterVersion = "starter"

	// The guardrails' judge: Meta's Llama Guard 3 1B, which reads a prompt
	// or an answer and says whether it is unsafe and why. Q4_K_M is about
	// 900 MB and judges a short text in about half a second on a CPU.
	guardModelName    = "llama-guard-3-1b"
	guardModelSource  = "hf:QuantFactory/Llama-Guard-3-1B-GGUF/Llama-Guard-3-1B.Q4_K_M.gguf"
	guardModelVersion = "guard"

	// localLabel marks the control plane's own node, so the starter model is
	// served where it was asked for and does not wander to the first GPU node
	// that joins later.
	localLabelKey   = "llmgw.control-plane"
	localLabelValue = "true"
)

// agentUnit is deploy/systemd/llmgw-agent.service; a test holds them together.
const agentUnit = `[Unit]
Description=llmgw node agent
After=network-online.target
Wants=network-online.target

[Service]
ExecStart=/usr/local/bin/llmgw agent --config /etc/llmgw/agent.yaml
# Root, because the agent launches engines and needs device access. Engines
# outlive a restart of this unit and are adopted again on the next start, so
# KillMode must not take the whole process group with it.
KillMode=process
Restart=on-failure
RestartSec=2s
# RDMA pins the buffers it registers, and NCCL registers a great deal: with
# systemd's default 8 MB lock limit, an engine that drops the capability that
# lets root ignore it fails its first collective over InfiniBand.
LimitMEMLOCK=infinity
StateDirectory=llmgw-agent

[Install]
WantedBy=multi-user.target
`

// reachableAddr is the address other machines use for this one: the first
// name given with --host, else the first address it has, else its hostname.
// It is in the server certificate either way.
func (in *serverInit) reachableAddr() string {
	if len(in.hosts) > 0 {
		return in.hosts[0]
	}
	if a := primaryAddr(); a != "" {
		return a
	}
	name, _ := os.Hostname()
	return name
}

func (in *serverInit) systemctl(ctx context.Context, args ...string) error {
	out, err := exec.CommandContext(ctx, "systemctl", args...).CombinedOutput()
	if err != nil {
		return fmt.Errorf("systemctl %s: %v: %s", strings.Join(args, " "), err, strings.TrimSpace(string(out)))
	}
	return nil
}

// writeUnit installs a unit when it differs from what is there, keeping
// somebody's edited copy beside it, and says whether it changed.
func (in *serverInit) writeUnit(path, body string) (bool, error) {
	b, err := os.ReadFile(path)
	if err == nil && string(b) == body {
		return false, nil
	}
	if err == nil {
		_ = os.WriteFile(path+".bak", b, 0o644)
	}
	return true, os.WriteFile(path, []byte(body), 0o644)
}

// --- the mirror ---

func (in *serverInit) mirrorUnit() string {
	base := "http://" + net.JoinHostPort(in.reachableAddr(), "8081")
	return mirrorUnitFor(mirrorDir, mirrorAPI, base, nil,
		"'llmgw server init'. Bytes on :8081 for every node; the\n# management API on loopback only, for the control plane beside it.")
}

// mirrorUnitFor is a mirror's systemd unit. allow is who, besides the
// machine itself, may use the management API.
func mirrorUnitFor(dir, apiListen, base string, allow []string, about string) string {
	exec := "/usr/local/bin/llmgw mirror serve --dir " + dir + " --listen :8081 --api-listen " + apiListen + " --base " + base
	for _, a := range allow {
		exec += " --api-allow " + a
	}
	return `[Unit]
Description=llmgw model mirror
After=network-online.target
Wants=network-online.target

[Service]
# Written by ` + about + `
ExecStart=` + exec + `
User=llmgw
Group=llmgw
Restart=on-failure
RestartSec=2s
NoNewPrivileges=true
ProtectSystem=strict
ProtectHome=true
PrivateTmp=true
ReadWritePaths=` + dir + `

[Install]
WantedBy=multi-user.target
`
}

// mirrorAPIURL is the mirror's management API: one on its own machine, or
// the one beside this control plane.
func (in *serverInit) mirrorAPIURL() string {
	if in.mirrorURL != "" {
		return in.mirrorURL
	}
	return "http://" + mirrorAPI
}

// useMirror points the control plane at a mirror on another machine, set up
// there by 'llmgw mirror init'.
func (in *serverInit) useMirror(ctx context.Context) error {
	m, err := mirror.NewClient(in.mirrorURL, tlsconf.Config{}, 5*time.Second)
	if err != nil {
		return err
	}
	if err := m.Health(ctx); err != nil {
		return fmt.Errorf("the mirror at %s does not answer: %v. On that machine: sudo llmgw mirror init --control-plane <this machine's address>", in.mirrorURL, err)
	}
	cfg, err := config.Load(in.configPath())
	if err != nil {
		return err
	}
	switch cfg.Mirror.API {
	case in.mirrorURL:
		in.kept("mirror at %s", in.mirrorURL)
		return nil
	case "":
	default:
		in.kept("mirror at %s, as server.yaml says; change it there to use %s", cfg.Mirror.API, in.mirrorURL)
		return nil
	}
	if err := appendMirrorBlockFor(in.configPath(), in.mirrorURL,
		"# The model mirror, on its own machine, added by 'llmgw server init'."); err != nil {
		return err
	}
	in.done("server.yaml points at the mirror at %s", in.mirrorURL)
	return nil
}

// ensureMirror runs a mirror on this machine and points the control plane at
// it. Its management API is on loopback, which is why it needs no certificate:
// nothing but this machine can reach it.
func (in *serverInit) ensureMirror(ctx context.Context) error {
	cfg, err := config.Load(in.configPath())
	if err != nil {
		return err
	}
	if cfg.Mirror.API != "" && cfg.Mirror.API != "http://"+mirrorAPI {
		in.kept("mirror at %s, as server.yaml says", cfg.Mirror.API)
		return nil
	}

	if err := os.MkdirAll(mirrorDir, 0o755); err != nil {
		return err
	}
	if err := in.chown(mirrorDir, true); err != nil {
		return err
	}
	changed, err := in.writeUnit(mirrorUnitPath, in.mirrorUnit())
	if err != nil {
		return err
	}
	if err := in.systemctl(ctx, "daemon-reload"); err != nil {
		return err
	}
	if err := in.systemctl(ctx, "enable", "--now", "llmgw-mirror"); err != nil {
		return err
	}
	if changed {
		if err := in.systemctl(ctx, "restart", "llmgw-mirror"); err != nil {
			return err
		}
		in.done("model mirror in %s, serving on :8081", mirrorDir)
	} else {
		in.kept("model mirror in %s", mirrorDir)
	}

	if cfg.Mirror.API == "" {
		if err := appendMirrorBlockFor(in.configPath(), "http://"+mirrorAPI,
			"# The model mirror on this machine, added by 'llmgw server init'. Its\n"+
				"# management API is on loopback, so it needs no certificate."); err != nil {
			return err
		}
		in.done("server.yaml points at the mirror")
	}

	m, err := mirror.NewClient("http://"+mirrorAPI, tlsconf.Config{}, 5*time.Second)
	if err != nil {
		return err
	}
	return waitFor(ctx, 30*time.Second, "the mirror to answer", func() (bool, error) {
		return m.Health(ctx) == nil, nil
	})
}

// appendMirrorBlockFor adds the mirror to server.yaml without touching the
// rest, keeping the file as it was beside it.
func appendMirrorBlockFor(path, api, comment string) error {
	b, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	backup := fmt.Sprintf("%s.bak-%s", path, time.Now().Format("20060102-150405"))
	if err := os.WriteFile(backup, b, 0o640); err != nil {
		return err
	}
	block := "\n" + comment + "\nmirror:\n  api: " + api + "\n"
	info, err := os.Stat(path)
	if err != nil {
		return err
	}
	return os.WriteFile(path, append(b, []byte(block)...), info.Mode().Perm())
}

// --- this machine as a node ---

// gpuClasses asks nvidia-smi what cards this machine has. None is an answer,
// not an error: the node then serves on CPU.
func gpuClasses(ctx context.Context) ([]string, int) {
	gpus, err := discovery.DiscoverGPUsSMI(ctx, "")
	if err != nil || len(gpus) == 0 {
		return nil, 0
	}
	return discovery.Inventory{GPUs: gpus}.Classes(), len(gpus)
}

// ensureNode makes this machine an agent of its own control plane.
func (in *serverInit) ensureNode(ctx context.Context) (string, error) {
	nodeID, _ := os.Hostname()
	if nodeID == "" {
		return "", errors.New("this machine has no hostname to use as its node id")
	}

	if classes, n := gpuClasses(ctx); n > 0 {
		in.kept("%d NVIDIA GPU(s) here: %s — the node will use them", n, strings.Join(classes, ", "))
	} else if hasNVIDIADevice() {
		fmt.Fprintf(in.out, "  !!    an NVIDIA card is on the PCI bus but nvidia-smi does not answer: "+
			"install the driver and run init again to serve on it; until then this node is CPU only\n")
	} else {
		in.kept("no NVIDIA GPU here; the node serves on CPU")
	}

	certPath := filepath.Join(in.etcDir, "agent.pem")
	if _, err := os.Stat(certPath); err == nil {
		in.kept("agent certificate for %s", nodeID)
	} else {
		if _, err := os.Stat(filepath.Join(in.pkiDir, "ca-key.pem")); err != nil {
			return "", fmt.Errorf("the authority is not on this machine, so init cannot give it a node " +
				"identity; onboard it from where the authority is, or pass --no-node")
		}
		ca, err := loadCA(in.pkiDir)
		if err != nil {
			return "", err
		}
		leaf, err := ca.IssueClient(nodeID)
		if err != nil {
			return "", err
		}
		if err := leaf.Write(in.etcDir, "agent"); err != nil {
			return "", err
		}
		in.done("agent certificate for %s", nodeID)
	}

	agentConf := filepath.Join(in.etcDir, "agent.yaml")
	if _, err := os.Stat(agentConf); err == nil {
		in.kept("%s (never overwritten)", agentConf)
	} else {
		server := "https://" + net.JoinHostPort(in.reachableAddr(), "7443")
		body := fmt.Sprintf(`# This machine as a node of its own control plane, written by 'llmgw server init'.
server: %q
node_id: %q
tls:
  ca: %s/ca.pem
  cert: %s/agent.pem
  key: %s/agent-key.pem
labels:
  %s: %q
`, server, nodeID, in.etcDir, in.etcDir, in.etcDir, localLabelKey, localLabelValue)
		if err := os.WriteFile(agentConf, []byte(body), 0o644); err != nil {
			return "", err
		}
		in.done("%s", agentConf)
	}

	changed, err := in.writeUnit(agentUnitPath, agentUnit)
	if err != nil {
		return "", err
	}
	if err := in.systemctl(ctx, "daemon-reload"); err != nil {
		return "", err
	}
	if err := in.systemctl(ctx, "enable", "--now", "llmgw-agent"); err != nil {
		return "", err
	}
	if changed {
		_ = in.systemctl(ctx, "restart", "llmgw-agent")
		in.done("llmgw-agent started")
	}

	// Registered is quick. Having an engine is not: the node downloads and
	// checks its own llama.cpp, the CUDA build when it has a GPU, and a new
	// machine is often busy with its own package updates first.
	client := in.adminAPI()
	last := ""
	err = waitFor(ctx, 20*time.Minute, "this node to be ready with an engine", func() (bool, error) {
		var nodes []adminapi.NodeView
		if err := client.get(ctx, "/api/v1/nodes", &nodes); err != nil {
			return false, nil
		}
		for _, n := range nodes {
			if n.ID != nodeID {
				continue
			}
			if n.State == "ready" && slices.Contains(n.Engines, "llamacpp") {
				return true, nil
			}
			for _, p := range n.Provision {
				if line := fmt.Sprintf("%s %s: %s %s", p.Engine, p.Version, p.State, p.Detail); line != last {
					fmt.Fprintf(in.out, "  ..    %s\n", strings.TrimSpace(line))
					last = line
				}
				// Only llama.cpp is waited for: vLLM installs for minutes after,
				// and one this node cannot run is not a reason to stop here.
				if p.Engine == "llamacpp" && (p.State == "failed" || p.State == "unsupported") {
					return false, fmt.Errorf("installing %s on this node: %s", p.Engine, p.Detail)
				}
			}
		}
		return false, nil
	})
	if err != nil {
		return "", err
	}
	in.kept("node %s is ready", nodeID)
	return nodeID, nil
}

// hasNVIDIADevice reports an NVIDIA card on the PCI bus, driver or not, so a
// machine whose GPU is invisible for want of a driver says so instead of
// quietly serving on CPU.
func hasNVIDIADevice() bool {
	vendors, _ := filepath.Glob("/sys/bus/pci/devices/*/vendor")
	for _, v := range vendors {
		b, err := os.ReadFile(v)
		if err == nil && strings.TrimSpace(string(b)) == "0x10de" {
			class, _ := os.ReadFile(filepath.Join(filepath.Dir(v), "class"))
			// 0x03xxxx is a display or 3D controller; NVIDIA also makes bridges.
			if strings.HasPrefix(strings.TrimSpace(string(class)), "0x03") {
				return true
			}
		}
	}
	return false
}

// --- the starter model ---

// starterModel is the catalog entry: a GPU profile first when there is a card,
// and a CPU profile always, both pinned to this machine.
func starterModel(source string, gpuClasses []string, pinned bool) catalog.Model {
	var pin map[string]string
	if pinned {
		pin = map[string]string{localLabelKey: localLabelValue}
	}
	var profiles []catalog.Profile
	for _, class := range gpuClasses {
		profiles = append(profiles, catalog.Profile{
			Name: "gpu-" + class, Engine: "llamacpp", GPUClass: class,
			Priority: 200, MinMemoryGB: 2, NodeSelector: pin,
			Args: []string{"--ctx-size", "32768"},
		})
	}
	profiles = append(profiles, catalog.Profile{
		Name: "cpu", Engine: "llamacpp", Priority: 100, MinMemoryGB: 2, NodeSelector: pin,
		Args: []string{"--ctx-size", "16384"},
	})
	return catalog.Model{
		Name:           starterName,
		Aliases:        []string{starterAlias},
		Description:    "A small model init put here so there is something to try; fast, not clever.",
		Source:         source,
		RoutingProfile: "chat",
		ContextLength:  4096,
		MaxOutput:      1024,
		Scaling:        catalog.Scaling{Min: 1, Max: 1},
		Profiles:       profiles,
	}
}

// llamaGuard3Categories name Llama Guard 3's S1 to S13, so a verdict is
// recorded as "Violent Crimes" rather than "S1".
var llamaGuard3Categories = []string{"Violent Crimes", "Non-Violent Crimes", "Sex-Related Crimes",
	"Child Sexual Exploitation", "Defamation", "Specialized Advice", "Privacy", "Intellectual Property",
	"Indiscriminate Weapons", "Hate", "Suicide & Self-Harm", "Sexual Content", "Elections"}

// guardModel is the judge's catalog entry: CPU only, so it takes no card from
// the models people use, on whichever node has room.
func guardModel(source string) catalog.Model {
	return catalog.Model{
		Name:        guardModelName,
		Description: "The guardrails' judge: reads prompts and answers for unsafe content. Not offered to clients.",
		Source:      source,
		Role:        catalog.RoleGuard,
		Guard:       &catalog.GuardModel{Format: "llamaguard3", Categories: llamaGuard3Categories},
		// Two: every judged request waits on it, and with fail: closed a
		// judge that is not there refuses everything.
		Scaling: catalog.Scaling{Min: 2, Max: 2},
		Profiles: []catalog.Profile{{
			Name: "cpu", Engine: "llamacpp", Priority: 100, MinMemoryGB: 2,
			Args: []string{"--ctx-size", "8192"},
		}},
	}
}

// ensureStarter fetches the starter model into the mirror and serves it here.
func (in *serverInit) ensureStarter(ctx context.Context) error {
	// Served here when this machine is a node, pinned so it does not wander
	// to the first GPU node that joins; otherwise on whichever node joins.
	pinned := !in.noNode
	return in.ensureServed(ctx, starterName, starterSource, starterVersion, pinned,
		func(url string, classes []string) catalog.Model { return starterModel(url, classes, pinned) })
}

// ensureGuardModel fetches the guardrails' judge model into the mirror and
// adds it to the catalog, to run on any node with room.
func (in *serverInit) ensureGuardModel(ctx context.Context) error {
	return in.ensureServed(ctx, guardModelName, guardModelSource, guardModelVersion, false,
		func(url string, _ []string) catalog.Model { return guardModel(url) })
}

// ensureServed fetches a model into the mirror, adds it to the catalog, and
// when it runs here, waits for it to answer.
func (in *serverInit) ensureServed(ctx context.Context, name, source, version string, here bool,
	entry func(url string, gpuClasses []string) catalog.Model,
) error {
	client := in.adminAPI()
	var models []adminapi.ModelView
	if err := client.get(ctx, "/api/v1/models", &models); err != nil {
		return fmt.Errorf("reading the catalog: %w", err)
	}
	for _, m := range models {
		if m.Name == name {
			in.kept("%s in the catalog", name)
			if !here {
				return nil
			}
			return in.waitServing(ctx, name)
		}
	}

	upstream, ok := sourcecheck.HTTPSource(source)
	if !ok {
		return fmt.Errorf("cannot turn %s into a download", source)
	}
	m, err := mirror.NewClient(in.mirrorAPIURL(), tlsconf.Config{}, 15*time.Second)
	if err != nil {
		return err
	}
	st, err := m.Status(ctx, name, version)
	var apiErr *mirror.APIError
	switch {
	case err == nil:
	case errors.As(err, &apiErr) && apiErr.Status == 404 && in.offline():
		// Asked of the internet, it would fail after a timeout and stop init
		// over a model nobody carried in.
		fmt.Fprintf(in.out, "  skip  %s: not in the mirror, and an offline cluster does not download it\n", name)
		return nil
	case errors.As(err, &apiErr) && apiErr.Status == 404:
		if st, err = m.Fetch(ctx, mirror.FetchRequest{
			Source: upstream, Model: name, Version: version,
		}); err != nil {
			return fmt.Errorf("asking the mirror for %s: %w", name, err)
		}
		fmt.Fprintf(in.out, "  ..    downloading %s into the mirror\n", name)
	default:
		return err
	}

	lastPct := -10
	err = waitFor(ctx, 30*time.Minute, name+" to download", func() (bool, error) {
		cur, err := m.Status(ctx, name, version)
		if err != nil {
			return false, nil
		}
		st = cur
		switch cur.Status.State {
		case mirror.Present:
			return true, nil
		case mirror.Failed:
			return false, fmt.Errorf("the mirror could not download %s: %s", name, cur.Status.Err)
		}
		if pct, ok := cur.Status.Percent(); ok && pct >= lastPct+10 {
			fmt.Fprintf(in.out, "  ..    %d%%\n", pct)
			lastPct = pct
		}
		return false, nil
	})
	if err != nil {
		return err
	}
	in.done("%s in the mirror", name)

	var classes []string
	if here {
		classes, _ = gpuClasses(ctx)
	}
	body, err := json.Marshal(entry(st.URL, classes))
	if err != nil {
		return err
	}
	if err := client.post(ctx, "/api/v1/models", body, &map[string]string{}); err != nil {
		return fmt.Errorf("adding %s to the catalog: %w", name, err)
	}
	if !here {
		in.done("%s in the catalog; it runs on a node with room for it", name)
		return nil
	}
	where := "on CPU"
	if len(classes) > 0 {
		where = "on " + strings.Join(classes, ", ") + ", or CPU if the card is busy"
	}
	in.done("%s in the catalog, served here %s", name, where)
	return in.waitServing(ctx, name)
}

// unplacedGrace is how long init lets placement say a model fits nowhere
// before believing it: long enough for a node still installing its engine.
var unplacedGrace = 90 * time.Second

func (in *serverInit) waitServing(ctx context.Context, name string) error {
	client := in.adminAPI()
	last := ""
	var unplaced time.Time
	return waitFor(ctx, 10*time.Minute, name+" to start", func() (bool, error) {
		var reps []adminapi.ReplicaView
		if err := client.get(ctx, "/api/v1/replicas?model="+name, &reps); err != nil {
			return false, nil
		}
		// Nowhere to put it is not waited out for ten minutes: once
		// placement has said why for unplacedGrace, that is the answer.
		if len(reps) == 0 {
			var ex adminapi.ExplainView
			if err := client.get(ctx, "/api/v1/explain-placement?model="+name, &ex); err == nil &&
				!ex.Placed && len(ex.Rejections) > 0 {
				if unplaced.IsZero() {
					unplaced = time.Now()
				}
				if time.Since(unplaced) > unplacedGrace {
					var why []string
					for _, r := range ex.Rejections {
						why = append(why, r.Node+": "+r.Reason)
					}
					return false, fmt.Errorf("%s has nowhere to run (%s); add a node with room, "+
						"or free some ('llmgw ctl explain-placement %s')", name, strings.Join(why, "; "), name)
				}
				return false, nil
			}
		}
		unplaced = time.Time{}
		for _, r := range reps {
			if r.State == "ready" {
				in.kept("%s is serving (profile %s)", name, r.Profile)
				return true, nil
			}
		}
		for _, r := range reps {
			line := r.State
			if r.Err != "" {
				line += ": " + r.Err
			}
			if line != last {
				fmt.Fprintf(in.out, "  ..    %s %s\n", name, line)
				last = line
			}
		}
		return false, nil
	})
}

func (in *serverInit) adminAPI() *ctlClient {
	addr, token := "127.0.0.1:9090", ""
	if cfg, err := config.Load(in.configPath()); err == nil {
		token = adminTokenPath(in.configPath(), cfg)
		if cfg.AdminListen != "" {
			addr = cfg.AdminListen
		}
	}
	return newCtlClient(addr, token)
}

// waitFor polls until done says so, fails, or the time runs out.
func waitFor(ctx context.Context, limit time.Duration, what string, done func() (bool, error)) error {
	deadline := time.Now().Add(limit)
	for {
		ok, err := done()
		if err != nil {
			return err
		}
		if ok {
			return nil
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("gave up waiting for %s after %s; run init again to keep waiting", what, limit)
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(2 * time.Second):
		}
	}
}