llmgw/internal/cli/ctl_upgrade.go

package cli

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"runtime"
	"slices"
	"strings"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/release"
	"gitea.mynodes.xyz/ilker/llmgw/internal/version"
)

func newUpgradeCmd(client func() *ctlClient) *cobra.Command {
	var (
		nodes     []string
		skipCP    bool
		releases  string
		perNode   time.Duration
		serverSvc string
		parallel  int
		sshUser   string
		othersCP  bool
	)
	c := &cobra.Command{
		Use:   "upgrade VERSION",
		Short: "Upgrade the control planes, then every node, in waves",
		Long: "Moves the cluster to a release: v0.10.4, or a candidate like v0.10.4-rc.2.\n\n" +
			"The other control planes first, over ssh as --ssh-user (the one sudo was\n" +
			"run by), then this machine — run it with sudo on a control plane, since it\n" +
			"replaces /usr/local/bin/llmgw and restarts llmgw-server. A leader restarted\n" +
			"last hands over to a standby already on the release.\n\n" +
			"Then the nodes, in waves: one first, then twice as many each time a wave\n" +
			"comes back whole, up to --parallel at once. Each agent fetches the release\n" +
			"for its own platform, checks it against the release's checksums, swaps it\n" +
			"in and restarts; engines keep serving through all of it and are adopted\n" +
			"again, and no ssh into the nodes is needed.\n\n" +
			"A node counts as done when it is back on the new version, ready, with as\n" +
			"many ready replicas as before. A wave not back within --timeout stops the\n" +
			"upgrade there, so a bad release reaches one node rather than all of them.\n" +
			"Nodes already on the version are skipped, so running it again resumes.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			v := args[0]
			if !release.Tag.MatchString(v) {
				return fmt.Errorf("%q is not a release version (v1.2.3 or v1.2.3-rc.4)", v)
			}
			out := cmd.OutOrStdout()
			ctx := cmd.Context()
			if releases == "" {
				// Where the control plane sends its nodes: an offline
				// cluster's mirror, so this machine needs no internet either.
				if cfg, ok := loadServerConfig("/etc/llmgw/server.yaml"); ok {
					releases = cfg.ReleasesBase()
				}
				if releases == "" {
					releases = release.DefaultBase
				}
			}
			if !skipCP {
				// The leader last, wherever this runs: every other control
				// plane is a standby, upgraded with nothing moving, and the
				// leader's restart hands over once, to a machine already on
				// the release.
				others := func() error {
					if !othersCP {
						return nil
					}
					return upgradeOtherControlPlanes(ctx, out, client, v, releases, sshUser)
				}
				self := func() error { return upgradeControlPlane(ctx, out, client, v, releases, serverSvc) }
				order := []func() error{others, self}
				if !thisLeads(ctx, client) {
					order = []func() error{self, others}
				}
				for _, step := range order {
					if err := step(); err != nil {
						return err
					}
				}
			}
			return upgradeNodes(ctx, out, client, v, nodes, perNode, parallel)
		},
	}
	c.Flags().StringSliceVar(&nodes, "nodes", nil, "only these nodes (default: all)")
	c.Flags().BoolVar(&skipCP, "skip-control-plane", false, "leave this machine as it is and upgrade only nodes")
	c.Flags().StringVar(&releases, "releases", "", "where release binaries are fetched from, for this machine "+
		"(default: as server.yaml says, the mirror's copy on an offline cluster, else the project's)")
	c.Flags().DurationVar(&perNode, "timeout", 3*time.Minute, "how long a node may take to come back")
	c.Flags().StringVar(&serverSvc, "service", "llmgw-server", "this machine's systemd unit")
	c.Flags().IntVar(&parallel, "parallel", 20, "the most nodes upgraded at once, once the first waves are back")
	c.Flags().StringVar(&sshUser, "ssh-user", os.Getenv("SUDO_USER"), "who to log in to the other control planes as")
	c.Flags().BoolVar(&othersCP, "other-control-planes", true, "upgrade the cluster's other control planes too, over ssh")
	return c
}

// thisLeads says whether this machine's control plane leads; one that does
// not say, alone or older, is taken to.
func thisLeads(ctx context.Context, client func() *ctlClient) bool {
	var sum adminapi.Summary
	if err := client().get(ctx, "/api/v1/summary", &sum); err != nil {
		return true
	}
	for _, cp := range sum.ControlPlanes.Members {
		if cp.This {
			return cp.State == "leading"
		}
	}
	return true
}

// upgradeOtherControlPlanes moves every other control plane of the cluster to
// v, over ssh with `llmgw upgrade`. Nothing is changed here if one cannot be
// reached.
func upgradeOtherControlPlanes(ctx context.Context, out io.Writer, client func() *ctlClient, v, releases, user string) error {
	var sum adminapi.Summary
	if err := client().get(ctx, "/api/v1/summary", &sum); err != nil {
		return err
	}
	for _, cp := range sum.ControlPlanes.Members {
		if cp.This {
			continue
		}
		r := &onboard.SSH{Addr: cp.Address, User: user}
		have, err := r.Run(ctx, "llmgw --version", nil)
		if err != nil {
			return fmt.Errorf("%s (%s) cannot be reached over ssh as %q: %v\n"+
				"run `sudo llmgw upgrade %s --service llmgw-server` there, then this again with "+
				"--other-control-planes=false", cp.Name, cp.Address, user, err, v)
		}
		if strings.Contains(string(have), " "+v+"\n") || strings.HasSuffix(strings.TrimSpace(string(have)), " "+v) {
			fmt.Fprintf(out, "== %s already runs %s\n", cp.Name, v)
			continue
		}
		fmt.Fprintf(out, "== %s (%s): %s -> %s\n", cp.Name, cp.State, strings.TrimSpace(string(have)), v)
		res, err := r.Run(ctx, "sudo llmgw upgrade "+shellQuote(v)+" --service llmgw-server --releases "+shellQuote(releases), nil)
		if err != nil {
			return fmt.Errorf("upgrading %s: %v: %s", cp.Name, err, strings.TrimSpace(string(res)))
		}
		fmt.Fprintf(out, "   upgraded and restarted\n")
	}
	return nil
}

func upgradeControlPlane(ctx context.Context, out io.Writer, client func() *ctlClient, v, releases, svc string) error {
	if version.Version == v {
		fmt.Fprintf(out, "this machine already runs %s\n", v)
		return nil
	}
	if os.Geteuid() != 0 {
		return errors.New("replacing this machine's binary needs root: run it with sudo, " +
			"or pass --skip-control-plane to upgrade only the nodes")
	}
	self, err := os.Executable()
	if err != nil {
		return err
	}
	fmt.Fprintf(out, "== this machine: %s -> %s\n", version.Version, v)
	backup, err := release.Install(ctx, nil, releases, v, runtime.GOOS, runtime.GOARCH, self)
	if err != nil {
		return err
	}
	fmt.Fprintf(out, "   installed, checksum verified; the old binary is %s\n", backup)
	var before []adminapi.NodeView
	_ = client().get(ctx, "/api/v1/nodes", &before)
	if b, err := exec.CommandContext(ctx, "systemctl", "restart", svc).CombinedOutput(); err != nil { //nolint:gosec // fixed command
		return fmt.Errorf("restarting %s: %v: %s", svc, err, strings.TrimSpace(string(b)))
	}
	fmt.Fprintf(out, "   %s restarted; waiting for its nodes to reconnect\n", svc)
	// Back when the admin API answers and every node that was ready is ready
	// again: agents reconnect within seconds of the control plane returning.
	want := 0
	for _, n := range before {
		if n.State == "ready" {
			want++
		}
	}
	deadline := time.Now().Add(2 * time.Minute)
	for {
		var now []adminapi.NodeView
		if err := client().get(ctx, "/api/v1/nodes", &now); err == nil {
			ready := 0
			for _, n := range now {
				if n.State == "ready" {
					ready++
				}
			}
			if ready >= want {
				fmt.Fprintf(out, "   back: %d of %d nodes ready\n", ready, len(now))
				return nil
			}
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("%s did not come back with its %d nodes within two minutes; "+
				"see journalctl -u %s, and %s holds the previous binary", svc, want, svc, backup)
		}
		time.Sleep(2 * time.Second)
	}
}

// upgradeNodes moves nodes to v in waves: one node first, then twice as
// many each time one comes back whole, up to parallel at once.
//
// Waves rather than one at a time, because one at a time is hours on a
// fleet of hundreds, and an agent's restart costs nothing served — the
// engines keep running and are adopted again. One first, because a release
// that does not come back should reach one node rather than a wave of them.
func upgradeNodes(ctx context.Context, out io.Writer, client func() *ctlClient, v string, only []string,
	perNode time.Duration, parallel int,
) error {
	var nodes []adminapi.NodeView
	if err := client().get(ctx, "/api/v1/nodes", &nodes); err != nil {
		return err
	}
	slices.SortFunc(nodes, func(a, b adminapi.NodeView) int { return strings.Compare(a.ID, b.ID) })
	var todo []adminapi.NodeView
	skipped := 0
	for _, n := range nodes {
		if len(only) > 0 && !slices.Contains(only, n.ID) {
			continue
		}
		if n.AgentVersion == v {
			skipped++
			continue
		}
		todo = append(todo, n)
	}
	parallel = max(parallel, 1)
	start, size, done := time.Now(), 1, 0
	for len(todo) > 0 {
		wave := todo[:min(size, len(todo))]
		todo = todo[len(wave):]
		if err := upgradeWave(ctx, out, client, v, wave, perNode); err != nil {
			return fmt.Errorf("%w\nstopped after this wave: %d upgraded, %d not tried", err, done, len(todo))
		}
		done += len(wave)
		size = min(size*2, parallel)
	}
	fmt.Fprintf(out, "%d node(s) upgraded to %s in %s, %d already on it\n",
		done, v, time.Since(start).Round(time.Second), skipped)
	// The one machine this does not reach. Said every time, because it is
	// exactly the kind of thing that stays three versions behind unnoticed.
	fmt.Fprintf(out, "The model mirror's host is not upgraded by this; on it, run:\n"+
		"  sudo llmgw upgrade %s --service llmgw-mirror\n", v)
	return nil
}

// upgradeWave tells every node in the wave at once, and waits until each is
// back on v, ready, with as many ready replicas as it had.
func upgradeWave(ctx context.Context, out io.Writer, client func() *ctlClient, v string,
	wave []adminapi.NodeView, perNode time.Duration,
) error {
	start := time.Now()
	before := readyReplicasByNode(ctx, client)
	ids := make([]string, len(wave))
	for i, n := range wave {
		ids[i] = n.ID
	}
	fmt.Fprintf(out, "== %d node(s): %s\n", len(wave), abbreviate(ids, 8))
	body, _ := json.Marshal(adminapi.UpgradeRequest{Version: v})
	var failed []string
	for _, n := range wave {
		if err := client().post(ctx, "/api/v1/nodes/"+n.ID+"/upgrade", body, nil); err != nil {
			failed = append(failed, fmt.Sprintf("%s: %v", n.ID, err))
		}
	}
	if len(failed) > 0 {
		return fmt.Errorf("not told to upgrade:\n  %s\n`llmgw ctl logs NODE` says more", strings.Join(failed, "\n  "))
	}
	waiting := slices.Clone(ids)
	deadline := time.Now().Add(perNode)
	for {
		var now []adminapi.NodeView
		_ = client().get(ctx, "/api/v1/nodes", &now)
		after := readyReplicasByNode(ctx, client)
		state := map[string]adminapi.NodeView{}
		for _, x := range now {
			state[x.ID] = x
		}
		waiting = slices.DeleteFunc(waiting, func(id string) bool {
			cur := state[id]
			return cur.AgentVersion == v && cur.State == "ready" && after[id] >= before[id]
		})
		if len(waiting) == 0 {
			fmt.Fprintf(out, "   done in %s\n", time.Since(start).Round(time.Second))
			return nil
		}
		if time.Now().After(deadline) {
			var why []string
			for _, id := range waiting {
				cur := state[id]
				why = append(why, fmt.Sprintf("%s (version %q, state %q, replicas %d/%d)",
					id, cur.AgentVersion, cur.State, after[id], before[id]))
			}
			return fmt.Errorf("not back within %s:\n  %s\n`llmgw ctl logs NODE` says why, and each keeps its "+
				"previous binary beside the new one", perNode, strings.Join(why, "\n  "))
		}
		time.Sleep(2 * time.Second)
	}
}

// abbreviate lists up to n names, and says how many more.
func abbreviate(names []string, n int) string {
	if len(names) <= n {
		return strings.Join(names, ", ")
	}
	return strings.Join(names[:n], ", ") + fmt.Sprintf(" and %d more", len(names)-n)
}

// readyReplicasByNode counts, for each node, the ready replicas with a
// shard on it.
func readyReplicasByNode(ctx context.Context, client func() *ctlClient) map[string]int {
	var reps []adminapi.ReplicaView
	out := map[string]int{}
	if err := client().get(ctx, "/api/v1/replicas", &reps); err != nil {
		return out
	}
	for _, r := range reps {
		if r.State != "ready" {
			continue
		}
		seen := map[string]bool{}
		for _, sh := range r.Shards {
			if !seen[sh.Node] {
				seen[sh.Node] = true
				out[sh.Node]++
			}
		}
	}
	return out
}