llmgw/internal/cli/service.go

package cli

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"strings"

	"github.com/spf13/cobra"
)

// llmgwServices are the units llmgw runs as on a machine, in the order they
// start: the store's clients after the store, the server before the agent
// that dials it.
var llmgwServices = []string{"llmgw-mirror", "llmgw-server", "llmgw-agent"}

// unitDir is where llmgw's units are installed.
var unitDir = "/etc/systemd/system"

// serviceRunner runs systemctl; a variable so tests stand in for it.
var serviceRunner = func(ctx context.Context, args ...string) (string, error) {
	out, err := exec.CommandContext(ctx, "systemctl", args...).CombinedOutput() //nolint:gosec // fixed command
	return strings.TrimSpace(string(out)), err
}

// installedServices are the llmgw units this machine has, or those named.
func installedServices(ctx context.Context, names []string) ([]string, error) {
	var want []string
	for _, n := range names {
		if !strings.HasPrefix(n, "llmgw-") {
			n = "llmgw-" + n
		}
		want = append(want, n)
	}
	var have []string
	for _, u := range llmgwServices {
		if _, err := os.Stat(unitDir + "/" + u + ".service"); err != nil {
			continue
		}
		if len(want) > 0 && !hasUnit(want, u) {
			continue
		}
		have = append(have, u)
	}
	for _, w := range want {
		if !hasUnit(have, w) {
			return nil, fmt.Errorf("%s is not installed here (this machine has: %s)", w, strings.Join(present(), ", "))
		}
	}
	if len(have) == 0 {
		return nil, errors.New("no llmgw service is installed on this machine")
	}
	return have, nil
}

func present() []string {
	var out []string
	for _, u := range llmgwServices {
		if _, err := os.Stat(unitDir + "/" + u + ".service"); err == nil {
			out = append(out, strings.TrimPrefix(u, "llmgw-"))
		}
	}
	if len(out) == 0 {
		return []string{"none"}
	}
	return out
}

func hasUnit(list []string, s string) bool {
	for _, x := range list {
		if x == s {
			return true
		}
	}
	return false
}

// newServiceCmd is llmgw service: this machine's llmgw services, one by one.
// The whole system is llmgw status, start, stop and restart.
func newServiceCmd() *cobra.Command {
	c := &cobra.Command{
		Use:   "service",
		Short: "This machine's llmgw services: status, start, stop, restart",
		Long: "The control plane (server), node agent (agent) and model mirror (mirror) on\n" +
			"this machine alone, by systemd. The whole system is 'llmgw status', 'llmgw\n" +
			"stop', 'llmgw start' and 'llmgw restart', run on a control plane.",
	}
	c.AddCommand(serviceVerbs()...)
	return c
}

func serviceVerbs() []*cobra.Command {
	status := &cobra.Command{
		Use:   "status [server|agent|mirror]...",
		Short: "Whether llmgw's services on this machine run",
		Long: "Each llmgw service this machine has — the control plane (server), the\n" +
			"node agent (agent), the model mirror (mirror) — and whether it runs.",
		RunE: func(cmd *cobra.Command, args []string) error {
			units, err := localUnits(cmd, args)
			if err != nil {
				return err
			}
			out := cmd.OutOrStdout()
			for _, u := range units {
				state, _ := serviceRunner(cmd.Context(), "is-active", u)
				since, _ := serviceRunner(cmd.Context(), "show", u, "-p", "ActiveEnterTimestamp", "--value")
				fmt.Fprintf(out, "%-14s %-10s %s\n", strings.TrimPrefix(u, "llmgw-"), state, since)
			}
			return nil
		},
	}
	verb := func(use, short, long, action string, reverse bool) *cobra.Command {
		var engines bool
		c := &cobra.Command{
			Use:   use + " [server|agent|mirror]...",
			Short: short,
			Long:  long,
			RunE: func(cmd *cobra.Command, args []string) error {
				if os.Geteuid() != 0 {
					return errors.New("starting and stopping services needs root: run it with sudo")
				}
				units, err := localUnits(cmd, args)
				if err != nil {
					return err
				}
				if reverse {
					// The agent before the server it dials, the server before
					// the mirror it reads.
					for i, j := 0, len(units)-1; i < j; i, j = i+1, j-1 {
						units[i], units[j] = units[j], units[i]
					}
				}
				out := cmd.OutOrStdout()
				for _, u := range units {
					if msg, err := serviceRunner(cmd.Context(), action, u); err != nil {
						return fmt.Errorf("%s %s: %v %s", action, u, err, msg)
					}
					state, _ := serviceRunner(cmd.Context(), "is-active", u)
					fmt.Fprintf(out, "%-14s %s\n", strings.TrimPrefix(u, "llmgw-"), state)
				}
				if action == "stop" && hasUnit(units, "llmgw-agent") {
					if engines {
						_, _ = exec.CommandContext(cmd.Context(), "pkill", "-u", "llmgw-engine").CombinedOutput() //nolint:gosec // fixed command
						fmt.Fprintln(out, "engines   stopped (every process of llmgw-engine)")
					} else {
						fmt.Fprintln(out, "\nThe engines keep serving: they outlive the agent, which takes them over again when it starts.\n"+
							"--engines stops them too; the control plane then places their replicas elsewhere.")
					}
				}
				if action == "stop" && hasUnit(units, "llmgw-server") {
					fmt.Fprintln(out, "A standby control plane, where there is one, takes over within seconds.")
				}
				return nil
			},
		}
		if action == "stop" {
			c.Flags().BoolVar(&engines, "engines", false, "also stop the engines this node runs")
		}
		return c
	}
	return []*cobra.Command{
		status,
		verb("start", "Start llmgw's services on this machine",
			"Starts each llmgw service this machine has, or those named: the mirror, the\n"+
				"control plane, the node agent, in that order.", "start", false),
		verb("stop", "Stop llmgw's services on this machine",
			"Stops each llmgw service this machine has, or those named: the agent, the\n"+
				"control plane, the mirror, in that order. The engines a node runs keep\n"+
				"serving unless --engines is given; a standby control plane takes over from\n"+
				"a leader stopped. They stay stopped until 'llmgw start', or the machine\n"+
				"boots.", "stop", true),
		verb("restart", "Restart llmgw's services on this machine",
			"Restarts each llmgw service this machine has, or those named. Engines keep\n"+
				"serving and are taken over again by the agent; a control plane's settings\n"+
				"and server.yaml are read anew.", "restart", false),
	}
}

func localUnits(cmd *cobra.Command, args []string) ([]string, error) {
	if runtime.GOOS != "linux" {
		return nil, fmt.Errorf("llmgw's services run under systemd on Linux; this is %s", runtime.GOOS)
	}
	return installedServices(cmd.Context(), args)
}