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/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
	)
	c := &cobra.Command{
		Use:   "upgrade VERSION",
		Short: "Upgrade this control plane, then every node, one at a time",
		Long: "Moves the cluster to a release: v0.10.4, or a candidate like v0.10.4-rc.2.\n\n" +
			"This machine first — run it with sudo on the control plane, since it\n" +
			"replaces /usr/local/bin/llmgw and restarts llmgw-server. Then each node, in\n" +
			"turn: its agent fetches the release for its own platform, checks it\n" +
			"against the release's checksums, swaps it in and restarts. Engines keep\n" +
			"serving through all of it and are adopted again; no ssh into the nodes is\n" +
			"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. One that is not 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 !skipCP {
				if err := upgradeControlPlane(ctx, out, client, v, releases, serverSvc); err != nil {
					return err
				}
			}
			return upgradeNodes(ctx, out, client, v, nodes, perNode)
		},
	}
	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", release.DefaultBase, "where release binaries are fetched from, for this machine")
	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")
	return c
}

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

func upgradeNodes(ctx context.Context, out io.Writer, client func() *ctlClient, v string, only []string, perNode time.Duration) 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) })
	done, skipped := 0, 0
	for _, n := range nodes {
		if len(only) > 0 && !slices.Contains(only, n.ID) {
			continue
		}
		if n.AgentVersion == v {
			skipped++
			continue
		}
		start := time.Now()
		before := readyReplicasOn(ctx, client, n.ID)
		fmt.Fprintf(out, "== %s: %s -> %s\n", n.ID, n.AgentVersion, v)
		body, _ := json.Marshal(adminapi.UpgradeRequest{Version: v})
		if err := client().post(ctx, "/api/v1/nodes/"+n.ID+"/upgrade", body, nil); err != nil {
			return fmt.Errorf("%s: %w\nstopped here; `llmgw ctl logs %s` says more", n.ID, err, n.ID)
		}
		fmt.Fprintf(out, "   new binary in place; waiting for it to come back with its %d replica(s)\n", before)
		deadline := time.Now().Add(perNode)
		for {
			var now []adminapi.NodeView
			_ = client().get(ctx, "/api/v1/nodes", &now)
			var cur adminapi.NodeView
			for _, x := range now {
				if x.ID == n.ID {
					cur = x
				}
			}
			after := readyReplicasOn(ctx, client, n.ID)
			if cur.AgentVersion == v && cur.State == "ready" && after >= before {
				fmt.Fprintf(out, "   done in %s: ready, replicas %d/%d\n",
					time.Since(start).Round(time.Second), after, before)
				done++
				break
			}
			if time.Now().After(deadline) {
				return fmt.Errorf("%s did not come back within %s (version %q, state %q, replicas %d/%d)\n"+
					"stopped here, before the next node; `llmgw ctl logs %s` says why, and the node "+
					"keeps its previous binary beside the new one", n.ID, perNode, cur.AgentVersion,
					cur.State, after, before, n.ID)
			}
			time.Sleep(2 * time.Second)
		}
	}
	fmt.Fprintf(out, "%d node(s) upgraded to %s, %d already on it\n", done, v, skipped)
	return nil
}

// readyReplicasOn counts the ready replicas with a shard on a node.
func readyReplicasOn(ctx context.Context, client func() *ctlClient, node string) int {
	var reps []adminapi.ReplicaView
	if err := client().get(ctx, "/api/v1/replicas", &reps); err != nil {
		return 0
	}
	n := 0
	for _, r := range reps {
		if r.State != "ready" {
			continue
		}
		for _, s := range r.Shards {
			if s.Node == node {
				n++
				break
			}
		}
	}
	return n
}