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"

	// 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 ips := localAddrs(); len(ips) > 0 {
		return ips[0].String()
	}
	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 `[Unit]
Description=llmgw model mirror
After=network-online.target
Wants=network-online.target

[Service]
# Written by 'llmgw server init'. Bytes on :8081 for every node; the
# management API on loopback only, for the control plane beside it.
ExecStart=/usr/local/bin/llmgw mirror serve --dir ` + mirrorDir + ` --listen :8081 --api-listen ` + mirrorAPI + ` --base ` + base + `
User=llmgw
Group=llmgw
Restart=on-failure
RestartSec=2s
NoNewPrivileges=true
ProtectSystem=strict
ProtectHome=true
PrivateTmp=true
ReadWritePaths=` + mirrorDir + `

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

// 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 := appendMirrorBlock(in.configPath()); 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
	})
}

// appendMirrorBlock adds the mirror to server.yaml without touching the rest,
// keeping the file as it was beside it.
func appendMirrorBlock(path 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# The model mirror on this machine, added by 'llmgw server init'. Its\n" +
		"# management API is on loopback, so it needs no certificate.\n" +
		"mirror:\n  api: http://" + mirrorAPI + "\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
				}
				if 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) catalog.Model {
	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,
	}
}

// ensureStarter fetches the starter model into the mirror and serves it here.
func (in *serverInit) ensureStarter(ctx context.Context) 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 == starterName {
			in.kept("%s in the catalog", starterName)
			return in.waitForStarter(ctx)
		}
	}

	upstream, ok := sourcecheck.HTTPSource(starterSource)
	if !ok {
		return fmt.Errorf("cannot turn %s into a download", starterSource)
	}
	m, err := mirror.NewClient("http://"+mirrorAPI, tlsconf.Config{}, 15*time.Second)
	if err != nil {
		return err
	}
	st, err := m.Status(ctx, starterName, starterVersion)
	var apiErr *mirror.APIError
	switch {
	case err == nil:
	case errors.As(err, &apiErr) && apiErr.Status == 404:
		if st, err = m.Fetch(ctx, mirror.FetchRequest{
			Source: upstream, Model: starterName, Version: starterVersion,
		}); err != nil {
			return fmt.Errorf("asking the mirror for %s: %w", starterName, err)
		}
		fmt.Fprintf(in.out, "  ..    downloading %s into the mirror\n", starterName)
	default:
		return err
	}

	lastPct := -10
	err = waitFor(ctx, 30*time.Minute, starterName+" to download", func() (bool, error) {
		cur, err := m.Status(ctx, starterName, starterVersion)
		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", starterName, 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", starterName)

	classes, _ := gpuClasses(ctx)
	body, err := json.Marshal(starterModel(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", starterName, err)
	}
	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", starterName, where)
	return in.waitForStarter(ctx)
}

func (in *serverInit) waitForStarter(ctx context.Context) error {
	client := in.adminAPI()
	last := ""
	return waitFor(ctx, 10*time.Minute, starterName+" to start", func() (bool, error) {
		var reps []adminapi.ReplicaView
		if err := client.get(ctx, "/api/v1/replicas?model="+starterName, &reps); err != nil {
			return false, nil
		}
		for _, r := range reps {
			if r.State == "ready" {
				in.kept("%s is serving (profile %s)", starterName, 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", starterName, line)
				last = line
			}
		}
		return false, nil
	})
}

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

// 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):
		}
	}
}