llmgw/internal/cli/node.go

package cli

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

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/firewall"
	"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/pgha"
	"gitea.mynodes.xyz/ilker/llmgw/internal/release"
	"gitea.mynodes.xyz/ilker/llmgw/internal/version"
)

func newNodeCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "node",
		Short: "Add machines to the cluster",
	}
	cmd.AddCommand(newNodeAddCmd(), newNodeRemoveCmd())
	return cmd
}

func newNodeRemoveCmd() *cobra.Command {
	var admin, serverCfg string
	cmd := &cobra.Command{
		Use:   "remove <node>",
		Short: "Forget a node that is gone for good",
		Long: `For a machine retired, or rebuilt under another name, which would otherwise
stay listed as lost until the control plane restarts.

Only a node that is not there is removed: stop its agent first, and wait for
it to be lost, a minute later. Its certificate stays valid, since the
authority keeps no revocation list: a machine that still has it can join
again, which is how a node rebuilt under the same name comes back.`,
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			na := nodeAdd{admin: admin, serverCfg: serverCfg}
			var out map[string]string
			if err := na.adminAPI().delete(cmd.Context(), "/api/v1/nodes/"+args[0], &out); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s removed\n", args[0])
			return nil
		},
	}
	cmd.Flags().StringVar(&admin, "admin", "", "the control plane's admin API (default: from server.yaml, else 127.0.0.1:9090)")
	cmd.Flags().StringVar(&serverCfg, "config", "/etc/llmgw/server.yaml", "the control plane's configuration")
	return cmd
}

// nodeAdd is what `llmgw node add` was asked to do.
type nodeAdd struct {
	opts      onboard.Options
	fabrics   []string
	servers   []string
	admin     string
	waitAll   bool
	noWait    bool
	plain     bool
	serverCfg string
	// named are the fabric names given in this run, by what the adapters
	// report the fabric to be.
	named map[string]string
}

func newNodeAddCmd() *cobra.Command {
	var na nodeAdd
	cmd := &cobra.Command{
		Use:   "add [user@]address...",
		Short: "Make a machine a node: check it, give it an identity, start its agent",
		Long: `Run on the machine holding the certificate authority — the control plane,
after 'llmgw server init' — as root. It logs in to the machine with ssh,
as you would by hand, and:

  1. installs llmgw there: the release for its architecture, checksum verified
  2. runs 'llmgw doctor' there, and stops if it finds a problem
  3. issues the node's certificate and installs it
  4. writes agent.yaml and the systemd unit, and starts the agent
  5. waits for the node to join and install llama.cpp

The node's id is its hostname, and the address its agent dials is this
machine's on the route to it, unless --node or --server say otherwise.
Nothing else needs giving, except --fabric for InfiniBand or RoCE: which
fabric an adapter is on has a name only you can choose, and it must be the
same on every node attached to it. Doctor suggests one when it sees adapters.

Several addresses add several nodes, one after another; a fabric named for
one is recognised on the next and not asked again.

The user needs sudo on the machine without a password. Run it again at any
time: the node gets a fresh certificate and configuration, and its running
engines are kept.`,
		Example: `  sudo llmgw node add 10.0.0.21
  sudo llmgw node add ubuntu@10.0.0.21 ubuntu@10.0.0.22 ubuntu@10.0.0.23
  sudo llmgw node add ubuntu@gpu01 --fabric ib-a:mlx5_0,mlx5_1,mlx5_2,mlx5_3`,
		Args: cobra.MinimumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			return na.run(cmd.Context(), args, cmd.OutOrStdout())
		},
	}
	f := cmd.Flags()
	f.StringVar(&na.opts.Node, "node", "", "node id (default: the machine's hostname)")
	f.StringVar(&na.opts.User, "ssh-user", "", "user to log in as (default: the one before @ in the address, else whoever ran sudo, else you)")
	f.StringArrayVar(&na.servers, "server", nil, "control plane address the agent dials, https://host:7443; repeatable for failover (default: this machine)")
	f.StringArrayVar(&na.fabrics, "fabric", nil, "InfiniBand or RoCE fabric and its adapters, name:mlx5_0,mlx5_1; repeatable")
	f.StringVar(&na.opts.MgmtAddr, "mgmt-addr", "", "address engines are reached at (default: found by the agent)")
	f.StringVar(&na.opts.MgmtIface, "mgmt-if", "", "interface NCCL bootstraps over (default: the one holding that address)")
	f.StringVar(&na.opts.LlamaCppBin, "llamacpp-bin", "", "a llama-server you installed, instead of the one the agent installs")
	f.StringVar(&na.opts.VLLMBin, "vllm-bin", "", "a vllm you installed, instead of the one the agent installs")
	f.StringVar(&na.opts.SGLangPython, "sglang-python", "", "python of an SGLang virtualenv you installed")
	f.StringVar(&na.opts.NvidiaSMI, "nvidia-smi", "", "where nvidia-smi is, when it is not on PATH")
	f.StringVar(&na.opts.IBRoot, "ib-root", "", "the sysfs InfiniBand directory, when not /sys/class/infiniband")
	f.StringVar(&na.opts.PKI, "pki", defaultPKIDir, "directory holding the certificate authority")
	f.StringVar(&na.opts.Releases, "releases", "", "where release files are downloaded from (default: releases in server.yaml, the mirror's copy on an offline cluster, else the project's)")
	f.BoolVar(&na.opts.IgnoreChecks, "ignore-checks", false, "add the node even when doctor finds a problem")
	f.BoolVar(&na.opts.NoNodeExporter, "no-node-exporter", false, "leave Prometheus' node exporter off the machine")
	f.BoolVar(&na.opts.NoFirewall, "no-firewall", false, "leave the machine's firewall alone, instead of keeping llmgw's ports to the cluster's networks")
	f.StringArrayVar(&na.opts.DisableEngines, "disable-engine", nil, "an engine this node never installs, e.g. vllm on GPUs it does not run on; repeatable")
	f.StringVar(&na.admin, "admin", "", "the control plane's admin API, to watch the node join (default: from server.yaml, else 127.0.0.1:9090)")
	f.StringVar(&na.serverCfg, "config", "/etc/llmgw/server.yaml", "the control plane's configuration, for the admin API's address")
	f.BoolVar(&na.waitAll, "wait", false, "wait for every engine, vLLM included, not only llama.cpp")
	f.BoolVar(&na.noWait, "no-wait", false, "do not wait for the node to join")
	f.BoolVar(&na.plain, "plain", false, "plain output, even in a terminal")
	return cmd
}

// prepare works out what every node added in this run shares.
func (na *nodeAdd) prepare() error {
	o := &na.opts
	for _, s := range na.fabrics {
		fab, err := onboard.ParseFabric(s)
		if err != nil {
			return err
		}
		if len(fab.HCAs) == 0 {
			return fmt.Errorf("--fabric %s names no adapters: give them as %s:mlx5_0,mlx5_1", s, fab.Name)
		}
		o.Fabrics = append(o.Fabrics, fab)
	}
	for _, s := range na.servers {
		for _, one := range strings.Split(s, ",") {
			if one = strings.TrimSpace(one); one != "" {
				o.Servers = append(o.Servers, one)
			}
		}
	}
	for _, e := range o.DisableEngines {
		if !slices.Contains(provision.Installable(), e) {
			return fmt.Errorf("--disable-engine %s: the engines are %s", e, strings.Join(provision.Installable(), ", "))
		}
	}
	// What the control plane will ask the node to install, so doctor checks
	// that vLLM rather than its own default. Its overrides for a GPU class
	// are not known until the node says what it has, and are not applied
	// here; doctor's verdict on vLLM is a warning either way.
	if o.VLLMVersion == "" {
		if cfg, err := loadEffectiveConfig(na.serverCfg); err == nil {
			o.VLLMVersion = cfg.WantedEngines()["vllm"]
			if o.VLLMVersion == "" {
				o.VLLMVersion = "none"
			}
		}
	}
	// The node's ports — its engines', its metrics — for the cluster only.
	if cfg, err := loadEffectiveConfig(na.serverCfg); err == nil && !o.NoFirewall {
		from, to := cfg.EnginePorts()
		o.FirewallPorts = []firewall.Port{firewall.Engines(from, to), firewall.NodeExporter}
		o.FirewallAllow, _ = cfg.FirewallAllow()
	}
	// Where the control plane upgrades nodes from, so a node joins on the
	// same release from the same place: an offline cluster's mirror.
	if o.Releases == "" {
		if cfg, err := loadEffectiveConfig(na.serverCfg); err == nil {
			o.Releases = cfg.ReleasesBase()
		}
		if o.Releases == "" {
			o.Releases = release.DefaultBase
		}
	}
	o.Version = version.Version
	if o.Binary == "" {
		exe, err := os.Executable()
		if err != nil {
			return err
		}
		o.Binary = exe
	}
	onboard.AgentUnit = agentUnit
	na.named = map[string]string{}
	return nil
}

// target is the options for one address, [user@]host.
func (na *nodeAdd) target(addr string) onboard.Options {
	o := na.opts
	o.Fabrics = slices.Clone(na.opts.Fabrics)
	o.Addr = addr
	if user, host, ok := strings.Cut(addr, "@"); ok {
		o.User, o.Addr = user, host
	}
	if o.User == "" {
		// Under sudo, the person who ran it — their account is usually the
		// one that exists on the nodes too. root's own keys still do the
		// logging in, since that is who is running ssh.
		o.User = os.Getenv("SUDO_USER")
	}
	return o
}

// rememberFabrics wraps the question of which fabric a node's adapters are
// on, so a name given for one node is used for the next whose adapters
// report the same fabric, without asking again. ask nil is a script, which
// is asked nothing and gets only what an earlier answer said.
func (na *nodeAdd) rememberFabrics(ask func([]onboard.Fabric) []onboard.Fabric) func([]onboard.Fabric) []onboard.Fabric {
	return func(s []onboard.Fabric) []onboard.Fabric {
		var known []onboard.Fabric
		all := true
		for i, f := range s {
			if name, ok := na.named[f.Key]; ok && f.Key != "" && f.Key != "unknown" {
				f.Name = name
				s[i].Name = name
				known = append(known, f)
			} else {
				all = false
			}
		}
		if all || ask == nil {
			return known
		}
		got := ask(s)
		for _, f := range got {
			if f.Key != "" && f.Key != "unknown" {
				na.named[f.Key] = f.Name
			}
		}
		return got
	}
}

func (na *nodeAdd) run(ctx context.Context, addrs []string, out io.Writer) error {
	if len(addrs) > 1 && na.opts.Node != "" {
		return errors.New("--node names one node; give one address with it")
	}
	if err := na.prepare(); err != nil {
		return err
	}
	tui := !na.plain && isTerminal(out)
	type outcome struct {
		addr, node string
		err        error
	}
	var results []outcome
	for i, addr := range addrs {
		o := na.target(addr)
		title := "llmgw · adding " + o.Addr
		if len(addrs) > 1 {
			title += fmt.Sprintf(" (%d of %d)", i+1, len(addrs))
		}
		if !o.NoFirewall {
			na.letIn(ctx, o.Addr, out)
		}
		var node string
		var err error
		if tui {
			node, err = na.runTUI(ctx, o, title)
		} else {
			node, err = na.runPlain(ctx, o, title, out)
		}
		results = append(results, outcome{addr, node, err})
		if ctx.Err() != nil {
			break
		}
	}
	if len(addrs) == 1 {
		return results[0].err
	}
	fmt.Fprintln(out)
	failed := 0
	for _, r := range results {
		if r.err != nil {
			failed++
			fmt.Fprintf(out, "  FAIL  %-20s %s\n", r.addr, r.err)
		} else {
			fmt.Fprintf(out, "  ok    %-20s %s\n", r.addr, r.node)
		}
	}
	if failed > 0 {
		return fmt.Errorf("%d of %d not added; running the same command again finishes them and leaves the rest as they are", failed, len(addrs))
	}
	return nil
}

// letIn opens this control plane's firewall to the node, when it has one
// and the node is on a network it does not already let in: otherwise a node
// on another subnet is set up and then never heard from. Said, not fatal.
func (na *nodeAdd) letIn(ctx context.Context, host string, out io.Writer) {
	var addrs []string
	if a, err := netip.ParseAddr(host); err == nil {
		addrs = []string{a.String()}
	} else {
		ips, err := net.DefaultResolver.LookupNetIP(ctx, "ip4", host)
		if err != nil {
			fmt.Fprintf(out, "  !!    firewall: cannot resolve %s (%v); if it is on another network, "+
				"list it in firewall_networks and run 'llmgw server init'\n", host, err)
			return
		}
		for _, ip := range ips {
			addrs = append(addrs, ip.Unmap().String())
		}
	}
	added, err := firewall.Admit(ctx, pgha.Local{}, addrs)
	switch {
	case err != nil:
		fmt.Fprintf(out, "  !!    firewall: could not let %s in: %v\n", host, err)
	case len(added) > 0:
		shown := make([]string, len(added))
		for i, p := range added {
			shown[i] = p.String()
		}
		fmt.Fprintf(out, "  ok    firewall: this control plane now lets in %s\n", strings.Join(shown, ", "))
	}
}

func (na *nodeAdd) runPlain(ctx context.Context, o onboard.Options, title string, out io.Writer) (string, error) {
	p := &plainProgress{out: out}
	fmt.Fprintf(out, "%s\n", title)
	o.AskFabrics = na.rememberFabrics(nil)
	res, err := onboard.New(o, p.event).Run(ctx)
	if err != nil {
		return "", err
	}
	if na.noWait {
		fmt.Fprintf(out, "\n%s is added. 'llmgw ctl nodes' shows it once it joins.\n\n", res.Node)
		return res.Node, nil
	}
	if err := na.watch(ctx, res.Node, p.event); err != nil {
		return res.Node, err
	}
	fmt.Fprintf(out, "\n%s is in the cluster.\n\n", res.Node)
	return res.Node, nil
}

// adminAPI is the control plane's admin API, from its configuration when
// this machine has it.
func (na *nodeAdd) adminAPI() *ctlClient {
	addr := na.admin
	token := ""
	cfg, err := config.Load(na.serverCfg)
	if err == nil {
		token = adminTokenPath(na.serverCfg, cfg)
	}
	if addr == "" {
		addr = "127.0.0.1:9090"
		if err == nil && cfg.AdminListen != "" {
			addr = cfg.AdminListen
			if strings.HasPrefix(addr, ":") {
				addr = "127.0.0.1" + addr
			}
		}
	}
	return newCtlClient(addr, token)
}

// errNoAdmin is the watch giving up because the admin API is not here, which
// is not a failure of onboarding.
var errNoAdmin = errors.New("no admin API")

// watch follows the node from the control plane: joining, then installing
// its engines. llama.cpp is waited for, since without it the node serves
// nothing; vLLM takes minutes more and is waited for only with --wait.
func (na *nodeAdd) watch(ctx context.Context, node string, emit func(onboard.Event)) error {
	client := na.adminAPI()
	var nodes []adminapi.NodeView
	if err := client.get(ctx, "/api/v1/nodes", &nodes); err != nil {
		emit(onboard.Event{Step: onboard.StepJoin, Level: onboard.Warn,
			Detail: "cannot reach the control plane's admin API here, so the node's joining is not followed",
			Fix:    "llmgw ctl nodes on the control plane shows it"})
		return nil
	}
	emit(onboard.Event{Step: onboard.StepJoin, Level: onboard.Running, Detail: "waiting for " + node + " to join"})
	var view adminapi.NodeView
	started := time.Now()
	err := waitFor(ctx, 2*time.Minute, node+" to join", func() (bool, error) {
		if err := client.get(ctx, "/api/v1/nodes", &nodes); err != nil {
			return false, nil
		}
		for _, n := range nodes {
			if n.ID == node && heardSince(n, started) {
				view = n
				return true, nil
			}
		}
		return false, nil
	})
	if err != nil {
		return fmt.Errorf("%s did not join: its agent's log says why (journalctl -u llmgw-agent on the node); a certificate whose name differs from the node id is refused", node)
	}
	emit(onboard.Event{Step: onboard.StepJoin, Level: onboard.Done, Detail: node + " joined: " + describeNode(view)})

	emit(onboard.Event{Step: onboard.StepEngines, Level: onboard.Running, Detail: "installing engines"})
	seen := map[string]string{}
	limit := 20 * time.Minute
	if na.waitAll {
		limit = time.Hour
	}
	err = waitFor(ctx, limit, "the engines on "+node, func() (bool, error) {
		if err := client.get(ctx, "/api/v1/nodes", &nodes); err != nil {
			return false, nil
		}
		for _, n := range nodes {
			if n.ID != node {
				continue
			}
			finished := true
			for _, p := range n.Provision {
				line := p.State + " " + p.Detail
				if seen[p.Engine] != line {
					seen[p.Engine] = line
					emit(engineEvent(p.Engine, p.Version, p.State, p.Detail))
				}
				switch p.State {
				case "installing":
					if p.Engine == "llamacpp" || na.waitAll {
						finished = false
					}
				case "failed":
					if p.Engine == "llamacpp" {
						return false, fmt.Errorf("llama.cpp did not install: %s", p.Detail)
					}
				}
			}
			return finished && len(n.Provision) > 0, nil
		}
		return false, nil
	})
	if err != nil {
		return err
	}
	msg := "engines ready"
	if v, ok := seen["vllm"]; ok && strings.HasPrefix(v, "installing") && !na.waitAll {
		msg = "llama.cpp ready; vLLM is still installing, which takes a few minutes more (watch /admin, or add --wait)"
	}
	emit(onboard.Event{Step: onboard.StepEngines, Level: onboard.Done, Detail: msg})
	return nil
}

// heardSince reports whether the node has been heard from since t, by the
// agent just started. A node added again is already listed, from its agent
// before the restart, and a cordoned one is not "ready": neither says whether
// the new agent connected.
func heardSince(n adminapi.NodeView, t time.Time) bool {
	switch n.State {
	case "unreachable", "lost", "unknown":
		return false
	}
	ago, err := time.ParseDuration(n.LastSeenAgo)
	if err != nil {
		return false
	}
	// A second's margin: the age is rounded to one.
	return ago+time.Second < time.Since(t)
}

func engineEvent(engine, version, state, detail string) onboard.Event {
	level := onboard.Info
	switch state {
	case "failed", "unsupported":
		level = onboard.Warn
	}
	return onboard.Event{Step: onboard.StepEngines, Level: level,
		Detail: fmt.Sprintf("%s %s: %s%s", engine, version, state, prefixed(": ", detail))}
}

func prefixed(p, s string) string {
	if s == "" {
		return ""
	}
	return p + s
}

func describeNode(n adminapi.NodeView) string {
	parts := []string{}
	if n.GPUTotal > 0 {
		parts = append(parts, fmt.Sprintf("%d GPU(s) %s", n.GPUTotal, strings.Join(n.Classes, ", ")))
	} else {
		parts = append(parts, "no GPU")
	}
	if len(n.Fabrics) > 0 {
		parts = append(parts, "fabrics "+strings.Join(n.Fabrics, ", "))
	}
	return strings.Join(parts, ", ")
}

// plainProgress prints events as lines, for a log or a pipe.
type plainProgress struct {
	out io.Writer
}

func (p *plainProgress) event(e onboard.Event) {
	mark := map[string]string{
		onboard.Running: "  ..  ", onboard.Done: "  ok  ", onboard.Info: "      ",
		onboard.Warn: "  !!  ", onboard.Fail: "  FAIL",
	}[e.Level]
	fmt.Fprintf(p.out, "%s %-13s %s\n", mark, e.Step, e.Detail)
	if e.Fix != "" {
		fmt.Fprintf(p.out, "       %-13s → %s\n", "", e.Fix)
	}
}