llmgw/internal/cli/mirror_peer.go

package cli

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"regexp"
	"slices"
	"strings"

	"github.com/spf13/cobra"

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

// A second mirror, so models can still be added and fetched while the first
// machine is down. Set up from the machine of the first, which it is kept in
// step with: the first is told of it, it of the first, and the control
// planes of both.

func newMirrorAddPeerCmd(dir *string) *cobra.Command {
	var controlPlanes []string
	var base string
	cmd := &cobra.Command{
		Use:   "add-peer [user@]host",
		Short: "Set up a second mirror on another machine, kept in step with this one",
		Long: `Run on the machine of the cluster's mirror, as root. It logs in to HOST with
ssh, installs llmgw there if it is not, and sets up a mirror that keeps in
step with this one: each copies what the other holds, checksum checked, and
a model removed from one is removed from the other. This mirror is told of
it, and the cluster's settings name it (mirror.peers), so while either
machine is down the control planes use the other and nodes fetch from it.

HOST needs room for every model this mirror holds: it copies them all.`,
		Example: "  sudo llmgw mirror add-peer ubuntu@10.0.0.12 --control-plane 10.0.0.10 --control-plane 10.0.0.12",
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if os.Geteuid() != 0 {
				return errors.New("add-peer changes this mirror's service and the cluster's settings: run it as root")
			}
			user, host, ok := strings.Cut(args[0], "@")
			if !ok {
				user, host = os.Getenv("SUDO_USER"), args[0]
			}
			return addMirrorPeer(cmd.Context(), cmd.OutOrStdout(), &onboard.SSH{Addr: host, User: user}, host,
				*dir, base, controlPlanes, "/etc/llmgw/server.yaml")
		},
	}
	cmd.Flags().StringArrayVar(&controlPlanes, "control-plane", nil,
		"address of a control plane that manages the mirrors; repeatable (default: this machine)")
	cmd.Flags().StringVar(&base, "base", "", "the address nodes fetch from on HOST (default: http://HOST:8081)")
	return cmd
}

// addMirrorPeer sets up a mirror on a machine reached by remote, kept in
// step with the one on this machine.
func addMirrorPeer(ctx context.Context, out io.Writer, remote onboard.Remote, host, dir, base string,
	controlPlanes []string, configPath string,
) error {
	ip, err := ipv4Of(host)
	if err != nil {
		return err
	}
	here := primaryAddr()
	if len(controlPlanes) == 0 {
		controlPlanes = []string{here}
	}
	if base == "" {
		base = "http://" + net.JoinHostPort(ip, "8081")
	}
	thisAPI := "http://" + net.JoinHostPort(here, "8082")
	peerAPI := "http://" + net.JoinHostPort(ip, "8082")

	// llmgw there, the same as here, when it has none.
	if _, err := remote.Run(ctx, "command -v llmgw >/dev/null", nil); err != nil {
		exe, err := os.Executable()
		if err != nil {
			return err
		}
		bin, err := os.ReadFile(exe) //nolint:gosec // this program
		if err != nil {
			return err
		}
		fmt.Fprintf(out, "  ..    llmgw to %s\n", host)
		if _, err := remote.Run(ctx, "sudo install -m 755 /dev/stdin /usr/local/bin/llmgw", bin); err != nil {
			return fmt.Errorf("installing llmgw on %s: %w", host, err)
		}
	}

	// The mirror there: its management API for the control planes and this
	// mirror, which reads it; kept in step with this one.
	allow := append([]string{}, controlPlanes...)
	if !slices.Contains(allow, here) {
		allow = append(allow, here)
	}
	script := "sudo llmgw mirror init --dir " + shellQuote(dir) + " --base " + shellQuote(base) + " --peer " + shellQuote(thisAPI)
	for _, a := range allow {
		script += " --control-plane " + shellQuote(a)
	}
	fmt.Fprintf(out, "  ..    the mirror on %s\n", host)
	if b, err := remote.Run(ctx, script, nil); err != nil {
		return fmt.Errorf("the mirror on %s: %w\n%s", host, err, b)
	}

	// This mirror told of it, and to let it read this one's API.
	changed, err := addToMirrorUnit(mirrorUnitPath, peerAPI, ip)
	if err != nil {
		return err
	}
	if changed {
		in := &serverInit{out: out}
		if err := in.systemctl(ctx, "daemon-reload"); err != nil {
			return err
		}
		if err := in.systemctl(ctx, "restart", "llmgw-mirror"); err != nil {
			return err
		}
	}
	fmt.Fprintf(out, "  ok    this mirror keeps in step with %s\n", peerAPI)

	// The control planes told of both.
	if err := addMirrorPeerSetting(configPath, peerAPI); err != nil {
		return err
	}
	fmt.Fprintf(out, "  ok    mirror.peers names %s; the control planes use it while this one is down\n", peerAPI)
	return nil
}

// addToMirrorUnit adds a peer and the address allowed to read this mirror's
// API to its systemd unit, and says whether it changed.
func addToMirrorUnit(path, peer, allow string) (bool, error) {
	b, err := os.ReadFile(path) //nolint:gosec // the mirror's unit
	if err != nil {
		return false, fmt.Errorf("this machine has no mirror to keep in step (%s): run add-peer on the mirror's machine", path)
	}
	unit := string(b)
	re := regexp.MustCompile(`(?m)^ExecStart=.*$`)
	line := re.FindString(unit)
	if line == "" {
		return false, fmt.Errorf("%s has no ExecStart", path)
	}
	next := line
	if !strings.Contains(next, "--peer "+peer) {
		next += " --peer " + peer
	}
	if !strings.Contains(next, "--api-allow "+allow+" ") && !strings.HasSuffix(next, "--api-allow "+allow) &&
		strings.Contains(next, "--api-listen :") {
		next += " --api-allow " + allow
	}
	if next == line {
		return false, nil
	}
	return true, os.WriteFile(path, []byte(strings.Replace(unit, line, next, 1)), 0o644)
}

// addMirrorPeerSetting names a peer in mirror.peers: in the cluster's
// settings when the store holds them, else in server.yaml.
func addMirrorPeerSetting(path, peer string) error {
	add := func(doc *settings.Doc) error {
		var peers []string
		if n := doc.Node([]string{"mirror", "peers"}); n != nil {
			_ = n.Decode(&peers)
		}
		if slices.Contains(peers, peer) {
			return nil
		}
		list := "[" + strings.Join(append(peers, peer), ", ") + "]"
		return doc.SetYAML([]string{"mirror", "peers"}, list)
	}
	changed, err := changeStoredSettings(path, "a second mirror (llmgw mirror add-peer)", add)
	if err != nil || changed {
		return err
	}
	doc, err := settings.Open(path)
	if err != nil {
		return err
	}
	if err := add(doc); err != nil {
		return err
	}
	_, err = doc.Save(nil)
	return err
}

// ensureSecondMirror sets up the mirror kept in step with this machine's on
// the store machine chosen, reached as the store's setup reached it.
func (in *serverInit) ensureSecondMirror(ctx context.Context) error {
	if !in.standbys {
		return errors.New("it needs --standbys: without them this mirror's API answers this machine alone, which the second cannot read")
	}
	i := in.secondMirror - 1
	if i < 0 || i >= len(in.patroniNodes) {
		return fmt.Errorf("there is no database machine %d to put it on", in.secondMirror)
	}
	if in.noMirror || in.mirrorURL != "" {
		return errors.New("it keeps in step with a mirror on this machine, and there is none")
	}
	n := in.patroniNodes[i]
	fmt.Fprintf(in.out, "  ..    a second mirror on %s\n", n.Name)
	if err := addMirrorPeer(ctx, in.out, n.Remote, n.IP, in.mirrorRoot(), "", in.mirrorAllow(), in.configPath()); err != nil {
		return err
	}
	in.done("a second mirror on %s (%s), kept in step with this one", n.Name, n.IP)
	return nil
}