llmgw/internal/cli/ctl.go

package cli

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
)

// ctlClient talks to a server's admin API.
type ctlClient struct {
	base  string
	http  *http.Client
	token string
}

// newCtlClient talks to the admin API at addr, with the admin token from
// LLMGW_ADMIN_TOKEN or the first of tokenFiles it can read, then the usual
// place.
func newCtlClient(addr string, tokenFiles ...string) *ctlClient {
	if !strings.HasPrefix(addr, "http://") && !strings.HasPrefix(addr, "https://") {
		addr = "http://" + addr
	}
	return &ctlClient{
		base:  strings.TrimRight(addr, "/"),
		http:  &http.Client{Timeout: 15 * time.Second},
		token: readAdminToken(tokenFiles...),
	}
}

// patient is the same client waiting up to d for an answer, for a call the
// server itself may spend minutes on.
func (c *ctlClient) patient(d time.Duration) *ctlClient {
	cp := *c
	cp.http = &http.Client{Timeout: d}
	return &cp
}

func (c *ctlClient) get(ctx context.Context, path string, out any) error {
	return c.do(ctx, http.MethodGet, path, nil, out)
}

// post sends an already-encoded body, so a caller that has marshalled a typed
// value keeps its field names rather than having them re-derived here.
func (c *ctlClient) post(ctx context.Context, path string, body []byte, out any) error {
	return c.do(ctx, http.MethodPost, path, json.RawMessage(body), out)
}

// put replaces a resource. Used where a second call is a better answer to the
// same question rather than another one: re-measuring a model's capacity.
func (c *ctlClient) put(ctx context.Context, path string, body []byte, out any) error {
	return c.do(ctx, http.MethodPut, path, json.RawMessage(body), out)
}

func (c *ctlClient) delete(ctx context.Context, path string, out any) error {
	return c.do(ctx, http.MethodDelete, path, nil, out)
}

func (c *ctlClient) do(ctx context.Context, method, path string, body, out any) error {
	var rdr io.Reader
	if body != nil {
		b, err := json.Marshal(body)
		if err != nil {
			return err
		}
		rdr = bytes.NewReader(b)
	}

	req, err := http.NewRequestWithContext(ctx, method, c.base+path, rdr)
	if err != nil {
		return err
	}
	// On every change, body or not: the admin listener refuses anything else,
	// because a browser can send a bodiless POST across origins unasked.
	if method != http.MethodGet {
		req.Header.Set("Content-Type", "application/json")
	}
	// Who made the change, kept with a stored model so a catalog that grew
	// over months can say where each entry came from.
	req.Header.Set("X-Llmgw-Actor", actor())
	if c.token != "" {
		req.Header.Set("Authorization", "Bearer "+c.token)
	}

	resp, err := c.http.Do(req)
	if err != nil {
		return fmt.Errorf("reaching the admin API at %s: %w\n"+
			"(is the server running, and is --server pointing at its admin_listen port?)", c.base, err)
	}
	defer resp.Body.Close()

	// Any 2xx. The API answers a change it has queued rather than made — a
	// restart — with 202, and treating that as a failure printed
	// "Error: 202 Accepted" over a request that had worked.
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		var e struct {
			Error string `json:"error"`
		}
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<12))
		if json.Unmarshal(b, &e) == nil && e.Error != "" {
			return fmt.Errorf("%s: %s", resp.Status, e.Error)
		}
		return fmt.Errorf("%s: %s", resp.Status, strings.TrimSpace(string(b)))
	}
	if out == nil {
		return nil
	}
	// A stream, as it comes: NDJSON a long run writes line by line.
	if w, ok := out.(io.Writer); ok {
		_, err := io.Copy(w, resp.Body)
		return err
	}
	// An older control plane answers a path it does not know with its
	// dashboard: a page, which read as JSON said only "invalid character '<'".
	if strings.HasPrefix(resp.Header.Get("Content-Type"), "text/html") {
		return fmt.Errorf("%s %s answered with a page, not the API: the control plane answering "+
			"(the leader, through a standby) does not know it — is it on an older release? "+
			"llmgw ctl status says which it runs", method, path)
	}
	return json.NewDecoder(resp.Body).Decode(out)
}

func newCtlCmd() *cobra.Command {
	var (
		server    string
		asJSON    bool
		tokenFile string
	)

	cmd := &cobra.Command{
		Use:   "ctl",
		Short: "Inspect and operate a running cluster",
		Long: "Talks to a server's admin API, which listens on admin_listen (:9090 by\n" +
			"default) rather than the client-facing port.",
	}
	cmd.PersistentFlags().StringVar(&server, "server", "127.0.0.1:9090", "admin API address")
	cmd.PersistentFlags().StringVar(&tokenFile, "token-file", "",
		"the admin token's file (default: LLMGW_ADMIN_TOKEN, else "+defaultAdminTokenFile+")")
	cmd.PersistentFlags().BoolVar(&asJSON, "json", false, "emit JSON")

	client := func() *ctlClient { return newCtlClient(server, tokenFile) }

	cmd.AddCommand(
		newNodesCmd(client, &asJSON),
		newLogsCmd(client, &asJSON),
		newRequestCmd(client, &asJSON),
		newModelsCmd(client, &asJSON),
		newReplicasCmd(client, &asJSON),
		newExplainCmd(client, &asJSON),
		newFabricTestCmd(client, &asJSON),
		newFacilityCmd(client, &asJSON),
		newLabelCmd(client, &asJSON),
		newNodeStateCmd(client, "cordon", "cordoned",
			"Stop placing new work on a node; what runs there keeps running"),
		newNodeStateCmd(client, "uncordon", "ready",
			"Allow placement on a node again"),
		newNodeStateCmd(client, "drain", "draining",
			"Move work off a node, one replica at a time"),
		// cert works on files rather than the admin API, so --server does not
		// apply to it; it lives here because this is where operators look.
		newBenchCmd(client, &asJSON),
		newEvalCmd(client, &asJSON),
		newStatusCmd(client, &asJSON),
		newModelCmd(client, &asJSON),
		newScaleCmd(client),
		newRestartCmd(client),
		newRetryCmd(client),
		newUpgradeCmd(client),
		newJobsCmd(client, &asJSON),
		newMeshJobsCmd(client, &asJSON),
		newEngineCheckCmd(),
		newEngineBuildCmd(client),
		newToolCheckCmd(client),
		newWeightsCmd(client, &asJSON),
		newCapacityCmd(client, &asJSON),
		newQuotaCmd(client, &asJSON),
		newGuardCmd(client, &asJSON),
		newKBCmd(client, &asJSON),
		newTrainCmd(client, &asJSON),
		newConfigCmd(client, &asJSON),
		newAuditCmd(client, &asJSON),
		newNotifyCmd(client, &asJSON),
		newUsageCmd(client, &asJSON),
		newFederationCmd(client, &asJSON),
		newCertCmd(),
	)
	return cmd
}

func newNodesCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	var f nodeFilter
	c := &cobra.Command{
		Use:   "nodes",
		Short: "List nodes and their capacity",
		Long: "On a fleet too large to read a whole table of, the filters are how a\n" +
			"question gets asked: which nodes on this fabric, which of this GPU class,\n" +
			"which are not ready. They combine, and they apply to --json too.",
		RunE: func(cmd *cobra.Command, _ []string) error {
			var nodes []adminapi.NodeView
			if err := client().get(cmd.Context(), "/api/v1/nodes", &nodes); err != nil {
				return err
			}
			total := len(nodes)
			nodes, err := f.apply(nodes)
			if err != nil {
				return err
			}
			if *asJSON {
				return emitJSON(nodes)
			}
			if len(nodes) == 0 {
				// A filter that matched nothing and an empty fleet call for
				// opposite next steps: one is a cluster to look at, the other
				// is a word to correct.
				if !f.empty() && total > 0 {
					fmt.Println(noMatches("nodes", f.describe(), total))
					return nil
				}
				fmt.Println("No nodes have registered.")
				return nil
			}

			w := tab()
			fmt.Fprintln(w, "NODE\tSTATE\tLOCATION\tGPUS\tMEM\tDISK\tENGINES\tCLASSES\tFABRICS\tREPLICAS\tLAST SEEN")
			for _, n := range nodes {
				fmt.Fprintf(w, "%s\t%s\t%s\t%d/%d free\t%s\t%s\t%s\t%s\t%s\t%d\t%s\n",
					n.ID, stateColumn(n), locationColumn(n), n.GPUFree, n.GPUTotal, memColumn(n.MemoryGB),
					diskColumn(n.DiskFreeGB, n.DiskTotalGB),
					// What the node can run, which is the first thing to look
					// at when a model will not place and the hardware looks
					// fine.
					enginesColumn(n),
					classesColumn(n), fabricsColumn(n), n.Replicas, n.LastSeenAgo)
			}
			_ = w.Flush()

			// Unhealthy GPUs are usually the answer when a node looks idle but
			// nothing schedules onto it. An engine still installing, or one
			// that failed to, is the other usual answer.
			said := false
			for _, n := range nodes {
				for _, p := range n.Unhealthy {
					fmt.Printf("\n%s: %s (kept out of placement)", n.ID, p)
					said = true
				}
				for _, p := range n.Provision {
					// Skipped is vLLM on a node with no GPU: as it should be.
					if p.State == "ready" || p.State == "skipped" {
						continue
					}
					fmt.Printf("\n%s: %s %s %s: %s", n.ID, p.Engine, p.Version, p.State, p.Detail)
					said = true
				}
				for _, w := range n.Warnings {
					fmt.Printf("\n%s: %s", n.ID, w)
					said = true
				}
				if n.Held != "" {
					fmt.Printf("\n%s: %s (kept out of placement)", n.ID, n.Held)
					said = true
				}
			}
			if said {
				fmt.Println()
			}
			return nil
		},
	}
	c.Flags().StringVar(&f.fabric, "fabric", "", "only nodes attached to this fabric")
	c.Flags().StringVar(&f.class, "class", "", "only nodes with this GPU class")
	c.Flags().StringVar(&f.state, "state", "", "only nodes in this state (prefix match)")
	c.Flags().StringArrayVar(&f.labels, "label", nil,
		"only nodes carrying key=value; repeatable, and key alone matches any value")
	c.Flags().StringVar(&f.domain, "domain", "", "only nodes in this rack, row, room, site or PDU, as level:name (rack:r12-07)")
	return c
}

// stateColumn is a node's state, with "held" when a held domain keeps it out
// of placement while its own state says ready.
func stateColumn(n adminapi.NodeView) string {
	if n.Held != "" && n.State == "ready" {
		return "ready, held"
	}
	return n.State
}

// locationColumn is where a node is, in a few characters; a dash when it
// carries no location label.
func locationColumn(n adminapi.NodeView) string {
	if n.Location == nil || n.Location.Empty() {
		return "-"
	}
	return n.Location.Short()
}

// modelStatus says what is happening to a model, in the words that decide what
// to do about it.
//
// The previous rendering distinguished "a replica is coming up" from "nothing
// could be placed" by a parenthetical — "0/1 ready (1 total)" against "0/1
// ready" — and those need opposite responses: one means wait, the other means
// look at explain-placement. On four models the difference is findable; on
// forty it is invisible, and a model that will never start looks exactly like
// one that is thirty seconds away.
//
// A model appearing here before it serves is deliberate: this is the operator's
// view, and a model coming up is precisely what someone wants to see. Clients
// are told separately and only when it is true — /v1/models lists what is
// actually served, never what is merely catalogued.
func modelStatus(m adminapi.ModelView) string {
	switch {
	case m.Parked:
		return "parked"
	case m.External != "":
		return "external: " + m.External
	case m.HeldFor != "":
		return "gave room to " + m.HeldFor
	case m.Unplaceable:
		return "cannot place"
	case m.Min == 0 && m.Total == 0:
		return "stopped"
	case m.Ready >= m.Min && m.Min > 0:
		return "ready"
	case m.Min == 0 && m.Ready > 0 && m.Ready >= m.Total:
		// Asked for, and up until it goes unused for long enough.
		return "ready (on demand)"
	case m.Total > m.Ready:
		return "starting"
	default:
		return "not placed"
	}
}

// idleColumn shows how long a replica has gone unused.
//
// "-" for a replica nothing has reported use of, which is emphatically not the
// same as idle forever: it is what every replica looks like before its first
// request, and what they all look like after a router restart. Showing that as
// a very large number would make an untouched fleet read as an abandoned one.
func idleColumn(sec int64) string {
	if sec < 0 {
		return "-"
	}
	return (time.Duration(sec) * time.Second).Round(time.Second).String()
}

// memColumn shows a node's memory, distinguishing "did not say" from a number.
// An older agent reports nothing, and that is not the same as a small machine.
func memColumn(gb int) string {
	if gb <= 0 {
		return "-"
	}
	return fmt.Sprintf("%dG", gb)
}

// diskColumn renders free space against the size it is free out of.
//
// Both, because they answer different questions and the first is useless
// without the second: 8 GB free is comfortable on a 10 TB filer and nearly
// full on a 10 GB root partition, and it is the second case that stops a
// model from being placed.
//
// A dash for a node that reported nothing, which an older agent and a node
// with no model directories both do. It has to read as "not known" rather
// than as "no space", or every unconfigured node looks like it is full.
func diskColumn(freeGB, totalGB int) string {
	if totalGB <= 0 {
		return "-"
	}
	return fmt.Sprintf("%dG/%dG", freeGB, totalGB)
}

func newModelsCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "models",
		Short: "List the catalog and how many replicas each model has",
		RunE: func(cmd *cobra.Command, _ []string) error {
			var models []adminapi.ModelView
			if err := client().get(cmd.Context(), "/api/v1/models", &models); err != nil {
				return err
			}
			if *asJSON {
				return emitJSON(models)
			}
			if len(models) == 0 {
				fmt.Println("The catalog is empty.")
				return nil
			}

			w := tab()
			fmt.Fprintln(w, "MODEL\tSERVING\tREPLICAS\tSTATUS\tCLASS\tRESERVES\tALIASES\tPROFILES\tFALLBACKS")
			var warnings []string
			for _, m := range models {
				serving := "no"
				if m.Ready > 0 {
					serving = "yes"
				}
				name := m.Name
				if m.Canary != nil {
					name += fmt.Sprintf(" (%d%% to canary)", m.Canary.Percent)
				}
				if m.CanaryOf != "" {
					name += " (canary of " + m.CanaryOf + ")"
				}
				fmt.Fprintf(w, "%s\t%s\t%d/%d\t%s\t%s\t%s\t%s\t%s\t%s\n",
					name, serving, m.Ready, m.Min, modelStatus(m), m.Class, reserves(m),
					join(m.Aliases), join(m.Profiles), join(m.Fallbacks))
				for _, x := range m.Warnings {
					warnings = append(warnings, m.Name+": "+x)
				}
			}
			if err := w.Flush(); err != nil {
				return err
			}
			for _, x := range warnings {
				fmt.Println("warning:", x)
			}
			return nil
		},
	}
}

func newReplicasCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	var model string
	var f replicaFilter
	c := &cobra.Command{
		Use:   "replicas",
		Short: "List replicas and the shards they are made of",
		RunE: func(cmd *cobra.Command, _ []string) error {
			path := "/api/v1/replicas"
			if model != "" {
				path += "?model=" + model
			}
			var reps []adminapi.ReplicaView
			if err := client().get(cmd.Context(), path, &reps); err != nil {
				return err
			}
			total := len(reps)
			reps = f.apply(reps)
			if *asJSON {
				return emitJSON(reps)
			}
			if len(reps) == 0 {
				if !f.empty() && total > 0 {
					fmt.Println(noMatches("replicas", f.describe(), total))
					return nil
				}
				fmt.Println("No replicas are running.")
				return nil
			}

			w := tab()
			fmt.Fprintln(w, "REPLICA\tMODEL\tSTATE\tPOOL\tON\tNODES\tAGE\tIDLE\tENDPOINT")
			for _, r := range reps {
				var nodes []string
				for _, s := range r.Shards {
					// "[0/mig2]" for a MIG instance, not the 1002 it is
					// numbered by.
					nodes = append(nodes, fmt.Sprintf("%s:%v", s.Node, discovery.DeviceLabels(s.GPUs)))
				}
				// The state says stale rather than a column of its own: most
				// replicas are not, a column would be empty for all of them,
				// and when one is it is the first thing worth knowing about it.
				state := r.State
				if r.Stale {
					state += " (stale profile)"
				}
				if r.Restart != "" {
					state += " (restart " + r.Restart + ")"
				}
				// What the engine computes on, by its own account: the first
				// shard's, which serves.
				on := "-"
				if len(r.Shards) > 0 {
					on = r.Shards[0].Compute.Short()
				}
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
					r.ID, r.Model, state, r.Pool, on, join(nodes),
					(time.Duration(r.AgeSec) * time.Second).String(),
					idleColumn(r.IdleSec), r.Endpoint)
				// A line rather than a column, because most replicas are not
				// downloading anything and a column would be empty for all of
				// them. When one is, it is the only thing worth reading.
				for _, s := range r.Shards {
					if s.Fetch == nil {
						continue
					}
					fmt.Fprintf(w, "  %s: fetching %s\t\t\t\t\t\t\t\t\n",
						s.Node, fetchColumn(s.Fetch))
				}
			}
			_ = w.Flush()

			for _, r := range reps {
				if r.Err != "" {
					fmt.Printf("\n%s: %s\n", r.ID, r.Err)
				}
				for _, s := range r.Shards {
					if s.Err != "" {
						fmt.Printf("  %s on %s: %s\n", s.ID, s.Node, s.Err)
					}
				}
			}
			return nil
		},
	}
	c.Flags().StringVar(&model, "model", "", "only this model's replicas")
	c.Flags().StringVar(&f.node, "node", "", "only replicas with a shard on this node")
	c.Flags().StringVar(&f.state, "state", "", "only replicas in this state (prefix match)")
	c.Flags().StringVar(&f.pool, "pool", "", "only this pool; matches either half of fabric/class")
	return c
}

func newExplainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "explain-placement <model>",
		Short: "Explain why a model has not been placed",
		Long: "Prints every candidate placement considered and the reason it was\n" +
			"rejected. \"Unschedulable\" on its own says nothing about which of six\n" +
			"constraints bit.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var v adminapi.ExplainView
			if err := client().get(cmd.Context(),
				"/api/v1/explain-placement?model="+args[0], &v); err != nil {
				return err
			}
			if *asJSON {
				return emitJSON(v)
			}

			if v.Placed && len(v.Rejections) == 0 {
				fmt.Printf("%s is placed; nothing was rejected.\n", v.Model)
				return nil
			}
			if len(v.Rejections) == 0 {
				fmt.Printf("No placement has been attempted for %s yet.\n", v.Model)
				return nil
			}

			if v.Placed {
				fmt.Printf("%s is placed. The candidates rejected along the way:\n\n", v.Model)
			} else {
				fmt.Printf("%s could not be placed. Every candidate and why:\n\n", v.Model)
			}
			w := tab()
			fmt.Fprintln(w, "PROFILE\tPOOL\tNODE\tREASON")
			for _, r := range v.Rejections {
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", r.Profile, dash(r.Pool), dash(r.Node), r.Reason)
			}
			return w.Flush()
		},
	}
}

func newNodeStateCmd(client func() *ctlClient, use, state, short string) *cobra.Command {
	return &cobra.Command{
		Use:   use + " <node>",
		Short: short,
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var out map[string]string
			err := client().do(cmd.Context(), http.MethodPost,
				"/api/v1/nodes/"+args[0]+"/state",
				map[string]string{"state": state}, &out)
			if err != nil {
				return err
			}
			fmt.Printf("%s is now %s\n", out["id"], out["state"])
			if state == "draining" {
				fmt.Println("Replicas will be moved off one at a time, each replacement " +
					"brought up before the original is stopped.")
			}
			return nil
		},
	}
}

// --- output helpers ---

func tab() *tabwriter.Writer { return tabTo(os.Stdout) }

// tabTo is tab writing where it is told: a summary drawn into the console
// went to the terminal underneath it instead, every three seconds.
func tabTo(w io.Writer) *tabwriter.Writer {
	return tabwriter.NewWriter(w, 0, 0, 2, ' ', 0)
}

func emitJSON(v any) error {
	enc := json.NewEncoder(os.Stdout)
	enc.SetIndent("", "  ")
	return enc.Encode(v)
}

// classesColumn is a node's GPU classes, with how many of each when it has
// more than one or MIG instances among them: a card cut into two 3g.40gb
// instances beside a whole card is "a100-80g-mig-3g.40gb×2,a100-80g×1", where
// the GPUS column's 3 alone would read as three of something.
// fabricsColumn is the node's fabrics, and when they are not all up, which
// adapters are down: what keeps it out of replicas spanning nodes.
func fabricsColumn(n adminapi.NodeView) string {
	s := join(n.Fabrics)
	if n.MultiNode != "" && n.MultiNode != "ready" {
		s += " (" + n.MultiNode + ")"
	}
	return s
}

// enginesColumn is each engine with its version: llamacpp b10991.
func enginesColumn(n adminapi.NodeView) string {
	if len(n.EngineVersions) == 0 {
		return join(n.Engines)
	}
	parts := make([]string, 0, len(n.Engines))
	for _, e := range n.Engines {
		if v := n.EngineVersions[e]; v != "" && v != "unknown" {
			e += " " + v
		}
		parts = append(parts, e)
	}
	return strings.Join(parts, ",")
}

// classesColumn is the node's GPU classes, with their compute capability
// when they report one: 6000-24g (7.5).
func classesColumn(n adminapi.NodeView) string {
	s := classesOnly(n)
	if len(n.ComputeCapabilities) > 0 && s != "-" {
		s += " (" + strings.Join(n.ComputeCapabilities, ",") + ")"
	}
	return s
}

func classesOnly(n adminapi.NodeView) string {
	counted := len(n.ClassGPUs) > 1
	for _, c := range n.ClassGPUs {
		counted = counted || c.MIG
	}
	if !counted {
		return join(n.Classes)
	}
	parts := make([]string, len(n.ClassGPUs))
	for i, c := range n.ClassGPUs {
		parts[i] = fmt.Sprintf("%s×%d", c.Class, c.Total)
	}
	return strings.Join(parts, ",")
}

func join(v []string) string {
	if len(v) == 0 {
		return "-"
	}
	return strings.Join(v, ",")
}

func dash(s string) string {
	if s == "" {
		return "-"
	}
	return s
}

// reserves is what one replica of a model holds of a node, with what it was
// measured using beside it: "6G 2c (5.1G used)".
func reserves(m adminapi.ModelView) string {
	if m.MemoryGB == 0 && m.CPUCores == 0 {
		return "-"
	}
	out := fmt.Sprintf("%dG", m.MemoryGB)
	if m.CPUCores > 0 {
		out += fmt.Sprintf(" %dc", m.CPUCores)
	}
	if m.MeasuredMemoryGB > 0 {
		out += fmt.Sprintf(" (%.1fG used)", m.MeasuredMemoryGB)
	}
	if len(m.Warnings) > 0 {
		out += " !"
	}
	return out
}