llmgw/internal/cli/system.go

package cli

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	"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/settings"
)

// systemClient is the admin API of the control plane on this machine, as its
// server.yaml says, with its token.
func systemClient(configPath string) func() *ctlClient {
	return func() *ctlClient {
		addr := "127.0.0.1:9090"
		var tokens []string
		if cfg, ok := loadServerConfig(configPath); ok {
			if cfg.AdminListen != "" {
				addr = cfg.AdminListen
				if strings.HasPrefix(addr, ":") {
					addr = "127.0.0.1" + addr
				}
			}
			tokens = append(tokens, adminTokenPath(configPath, cfg))
		}
		return newCtlClient(addr, tokens...)
	}
}

// newSystemCmds are llmgw status, stop, start and restart: the whole system,
// from a control plane.
func newSystemCmds() []*cobra.Command {
	var configPath string
	client := systemClient("/etc/llmgw/server.yaml")
	bind := func(c *cobra.Command) *cobra.Command {
		c.Flags().StringVar(&configPath, "config", "/etc/llmgw/server.yaml", "this control plane's configuration")
		c.PreRun = func(*cobra.Command, []string) { client = systemClient(configPath) }
		return c
	}
	call := func() *ctlClient { return client() }

	status := bind(&cobra.Command{
		Use:   "status",
		Short: "The whole system: whether it serves, its control planes, nodes, models and store",
		Long: "Run on a control plane. Says whether the service is stopped (llmgw stop),\n" +
			"and everything 'llmgw ctl status' does. One machine's services are\n" +
			"'llmgw service status'.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			var s adminapi.Summary
			if err := call().get(cmd.Context(), "/api/v1/summary", &s); err != nil {
				return notAControlPlane(err)
			}
			printSummary(cmd, s)
			return nil
		},
	})

	switchService := func(stop bool) *cobra.Command {
		use, short, long := "start", "Start the whole system serving again, after llmgw stop",
			"Every model comes back as its catalog entry asks: placed, loaded and\n"+
				"serving again, on every control plane within seconds. Run on a control\n"+
				"plane. The services themselves are never stopped by llmgw stop; a\n"+
				"machine's own are started with 'llmgw service start'."
		if stop {
			use, short, long = "stop", "Stop the whole system serving: every model taken down, requests answered 503",
				"Every model's replicas are taken down, freeing every GPU, and every\n"+
					"request for a model is answered 503 (system_stopped) until 'llmgw start'.\n"+
					"The control planes, agents and mirror keep running, so the console, the\n"+
					"admin pages and llmgw start still work, and no model's definition is\n"+
					"changed. In force on every control plane within seconds; kept in the\n"+
					"cluster's settings (serving.stopped), so a restart does not undo it."
		}
		var note string
		var wait bool
		c := bind(&cobra.Command{
			Use:   use,
			Short: short,
			Long:  long,
			Args:  cobra.NoArgs,
			RunE: func(cmd *cobra.Command, _ []string) error {
				if note == "" {
					note = "llmgw " + use
				}
				var s adminapi.Summary
				if err := call().get(cmd.Context(), "/api/v1/summary", &s); err != nil {
					return notAControlPlane(err)
				}
				if s.Stopped == stop {
					fmt.Fprintf(cmd.OutOrStdout(), "the system is already %s\n", map[bool]string{true: "stopped", false: "serving"}[stop])
					return nil
				}
				if err := changeSettings(cmd, call, note, func(doc *settings.Doc) error {
					return doc.SetYAML([]string{"serving", "stopped"}, fmt.Sprint(stop))
				}); err != nil {
					return err
				}
				out := cmd.OutOrStdout()
				if !wait {
					if stop {
						fmt.Fprintln(out, "stopping: every model's replicas are being taken down; llmgw status follows it")
					} else {
						fmt.Fprintln(out, "starting: every model is placed again as it asks; llmgw status follows it")
					}
					return nil
				}
				return waitForService(cmd, call, stop)
			},
		})
		c.Flags().StringVar(&note, "note", "", "why, kept with the settings' version and in the audit")
		c.Flags().BoolVar(&wait, "wait", false, "wait until every replica is down (stop) or the models have theirs (start)")
		return c
	}

	var sshUser string
	restart := bind(&cobra.Command{
		Use:   "restart",
		Short: "Restart every control plane and node agent; the engines keep serving",
		Long: "Every node's agent starts again, then the other control planes over ssh\n" +
			"as --ssh-user, then this one — the leader last, so it hands over to a\n" +
			"standby already restarted. The engines keep serving throughout and are\n" +
			"taken over again; nothing is placed anew. Run on a control plane, with\n" +
			"sudo. One machine's services are 'llmgw service restart'.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			if os.Geteuid() != 0 {
				return errors.New("restarting this control plane needs root: run it with sudo")
			}
			ctx, out := cmd.Context(), cmd.OutOrStdout()
			var s adminapi.Summary
			if err := call().get(ctx, "/api/v1/summary", &s); err != nil {
				return notAControlPlane(err)
			}
			var agents adminapi.AgentRestartView
			if err := call().post(ctx, "/api/v1/agents/restart", []byte("{}"), &agents); err != nil {
				return fmt.Errorf("restarting the agents: %w", err)
			}
			fmt.Fprintf(out, "agents     %d told to restart: %s\n", len(agents.Nodes), strings.Join(agents.Nodes, ", "))

			// The other control planes, standbys before a leader.
			var others []adminapi.ControlPlane
			for _, cp := range s.ControlPlanes.Members {
				if !cp.This {
					others = append(others, cp)
				}
			}
			for pass := 0; pass < 2; pass++ {
				for _, cp := range others {
					leading := cp.State == "leading"
					if leading != (pass == 1) {
						continue
					}
					r := &onboard.SSH{Addr: cp.Address, User: sshUser}
					if res, err := r.Run(ctx, "sudo systemctl restart llmgw-server", nil); err != nil {
						return fmt.Errorf("%s (%s) could not be restarted over ssh as %q: %v %s; "+
							"run `sudo llmgw service restart server` there", cp.Name, cp.Address, sshUser, err, strings.TrimSpace(string(res)))
					}
					fmt.Fprintf(out, "%-10s restarted (%s)\n", cp.Name, cp.State)
					time.Sleep(5 * time.Second)
				}
			}
			if _, err := os.Stat(unitDir + "/llmgw-server.service"); err == nil {
				fmt.Fprintln(out, "this       restarting")
				if b, err := exec.CommandContext(ctx, "systemctl", "restart", "llmgw-server").CombinedOutput(); err != nil { //nolint:gosec // fixed command
					return fmt.Errorf("restarting this control plane: %v: %s", err, strings.TrimSpace(string(b)))
				}
				fmt.Fprintln(out, "this       restarted")
			}
			fmt.Fprintln(out, "\nllmgw status shows the nodes coming back; their replicas are taken over as they are.")
			return nil
		},
	})
	restart.Flags().StringVar(&sshUser, "ssh-user", os.Getenv("SUDO_USER"), "who to log in to the other control planes as")

	return []*cobra.Command{status, switchService(true), switchService(false), restart}
}

// waitForService waits until no replica runs (stopped) or every model has
// the replicas it asks for (started).
func waitForService(cmd *cobra.Command, call func() *ctlClient, stop bool) error {
	deadline := time.Now().Add(15 * time.Minute)
	for time.Now().Before(deadline) {
		var s adminapi.Summary
		if err := call().get(cmd.Context(), "/api/v1/summary", &s); err == nil {
			if stop && s.Replicas.Ready+s.Replicas.Starting+s.Replicas.Draining == 0 {
				fmt.Fprintln(cmd.OutOrStdout(), "stopped: no replica runs")
				return nil
			}
			if !stop && s.Models.Satisfied >= s.Models.Total {
				fmt.Fprintf(cmd.OutOrStdout(), "serving: %d of %d models have the replicas they ask for\n",
					s.Models.Satisfied, s.Models.Total)
				return nil
			}
		}
		select {
		case <-cmd.Context().Done():
			return cmd.Context().Err()
		case <-time.After(5 * time.Second):
		}
	}
	return errors.New("not there within fifteen minutes; llmgw status says where it stands")
}

// notAControlPlane is a failure to reach the admin API, said as what it
// most often means.
func notAControlPlane(err error) error {
	return fmt.Errorf("%w\n(the whole system is stopped and started from a control plane; this machine's own "+
		"services are 'llmgw service status|start|stop|restart')", err)
}