llmgw/internal/cli/mirrorinit.go

package cli

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/netip"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/firewall"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/pgha"
	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// mirror init sets a machine up as the cluster's model mirror, on its own:
// the bytes for every node on :8081, the management API on :8082 for the
// control plane alone. server init does the same beside the control plane;
// this is for when the mirror wants its own disk.

func newMirrorInitCmd(dir *string) *cobra.Command {
	var (
		controlPlanes []string
		base          string
		noFirewall    bool
	)
	cmd := &cobra.Command{
		Use:   "init",
		Short: "Set this machine up as the cluster's model mirror and start it",
		Long: `Creates the llmgw user and the mirror's directory, and a systemd unit that
serves the model files to every node on :8081 and the management API on
:8082, and starts it.

The management API can fetch and delete models, so it answers this machine
and the control planes given with --control-plane, and nobody else. Then on
the control plane:

  sudo llmgw server init --mirror http://<this machine>:8082

Safe to run again: it changes only what differs.`,
		Example: `  sudo llmgw mirror init --control-plane 10.0.0.10`,
		Args:    cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			if runtime.GOOS != "linux" {
				return fmt.Errorf("mirror init sets up a Linux machine; this is %s", runtime.GOOS)
			}
			if os.Geteuid() != 0 {
				return errors.New("mirror init creates a user, writes /etc and starts a service: run it as root")
			}
			if len(controlPlanes) == 0 {
				return errors.New("--control-plane: which control plane manages this mirror (its address)")
			}
			if _, err := parseAllow(controlPlanes); err != nil {
				return err
			}
			if err := mirrorInit(cmd.Context(), cmd.OutOrStdout(), *dir, base, controlPlanes); err != nil {
				return err
			}
			if !noFirewall {
				mirrorFirewall(cmd.Context(), cmd.OutOrStdout(), controlPlanes)
			}
			return nil
		},
	}
	cmd.Flags().StringArrayVar(&controlPlanes, "control-plane", nil,
		"address of the control plane that manages this mirror; repeatable")
	cmd.Flags().BoolVar(&noFirewall, "no-firewall", false,
		"leave the firewall alone, instead of keeping the mirror's ports to the cluster's networks")
	cmd.Flags().StringVar(&base, "base", "",
		"the address nodes fetch from, http://host:8081 (default: this machine's first address)")
	return cmd
}

func mirrorInit(ctx context.Context, out io.Writer, dir, base string, controlPlanes []string) error {
	in := &serverInit{user: "llmgw", stateDir: "/var/lib/llmgw", out: out}
	if err := in.ensureUser(ctx); err != nil {
		return err
	}
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	if err := in.chown(dir, true); err != nil {
		return err
	}
	addr := primaryAddr()
	if base == "" {
		if addr == "" {
			return errors.New("this machine has no address to give nodes; give one with --base")
		}
		base = "http://" + net.JoinHostPort(addr, "8081")
	}
	unit := mirrorUnitFor(dir, ":8082", base, controlPlanes,
		"'llmgw mirror init'. Bytes on :8081 for every node; the\n# management API on :8082, for the control plane only.")
	changed, err := in.writeUnit(mirrorUnitPath, unit)
	if err != nil {
		return err
	}
	for _, args := range [][]string{{"daemon-reload"}, {"enable", "--now", "llmgw-mirror"}} {
		if err := in.systemctl(ctx, args...); err != nil {
			return err
		}
	}
	if changed {
		if err := in.systemctl(ctx, "restart", "llmgw-mirror"); err != nil {
			return err
		}
		in.done("model mirror in %s: files on :8081 at %s, management on :8082 for %s", dir, base, strings.Join(controlPlanes, ", "))
	} else {
		in.kept("model mirror in %s", dir)
	}
	m, err := mirror.NewClient("http://127.0.0.1:8082", tlsconf.Config{}, 5*time.Second)
	if err != nil {
		return err
	}
	if err := waitFor(ctx, 30*time.Second, "the mirror to answer", func() (bool, error) {
		return m.Health(ctx) == nil, nil
	}); err != nil {
		return err
	}
	in.kept("answering")
	api := "http://<this machine>:8082"
	if addr != "" {
		api = "http://" + net.JoinHostPort(addr, "8082")
	}
	fmt.Fprintf(out, "\nThe mirror is up. On the control plane:\n\n  sudo llmgw server init --mirror %s\n\n", api)
	return nil
}

// mirrorFirewall keeps the mirror's ports to the networks it reaches the
// control planes on: its files are for the cluster's nodes, its API for the
// control planes, which the API checks again by address.
func mirrorFirewall(ctx context.Context, out io.Writer, controlPlanes []string) {
	allowed, _ := parseAllow(controlPlanes)
	var toward []string
	var nets []netip.Prefix
	for _, p := range allowed {
		if p.IsSingleIP() {
			toward = append(toward, p.Addr().String())
		} else {
			nets = append(nets, p)
		}
	}
	found, err := firewall.Networks(ctx, pgha.Local{}, toward)
	if err == nil {
		var changed bool
		changed, err = firewall.Install(ctx, pgha.Local{}, firewall.Spec{
			Ports: []firewall.Port{firewall.MirrorFiles, firewall.MirrorAPI, firewall.NodeExporter},
			Allow: append(found, nets...),
		})
		if err == nil {
			if changed {
				fmt.Fprintln(out, "  ok    firewall: the mirror's ports answer the cluster's networks only")
			}
			return
		}
	}
	fmt.Fprintf(out, "  !!    firewall: %v; the mirror's ports answer every network\n", err)
}