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/control/adminapi"
)

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

func newCtlClient(addr 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},
	}
}

func (c *ctlClient) get(ctx context.Context, path string, out any) error {
	return c.do(ctx, http.MethodGet, 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
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}

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

	if resp.StatusCode != http.StatusOK {
		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
	}
	return json.NewDecoder(resp.Body).Decode(out)
}

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

	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().BoolVar(&asJSON, "json", false, "emit JSON")

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

	cmd.AddCommand(
		newNodesCmd(client, &asJSON),
		newModelsCmd(client, &asJSON),
		newReplicasCmd(client, &asJSON),
		newExplainCmd(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"),
	)
	return cmd
}

func newNodesCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "nodes",
		Short: "List nodes and their capacity",
		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
			}
			if *asJSON {
				return emitJSON(nodes)
			}
			if len(nodes) == 0 {
				fmt.Println("No nodes have registered.")
				return nil
			}

			w := tab()
			fmt.Fprintln(w, "NODE\tSTATE\tGPUS\tCLASSES\tFABRICS\tREPLICAS\tLAST SEEN")
			for _, n := range nodes {
				fmt.Fprintf(w, "%s\t%s\t%d/%d free\t%s\t%s\t%d\t%s\n",
					n.ID, n.State, n.GPUFree, n.GPUTotal,
					join(n.Classes), join(n.Fabrics), n.Replicas, n.LastSeenAgo)
			}
			_ = w.Flush()

			// Unhealthy GPUs are usually the answer when a node looks idle but
			// nothing schedules onto it.
			for _, n := range nodes {
				for _, p := range n.Unhealthy {
					fmt.Printf("\n%s: %s (kept out of placement)", n.ID, p)
				}
			}
			if hasUnhealthy(nodes) {
				fmt.Println()
			}
			return nil
		},
	}
}

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\tREPLICAS\tALIASES\tPROFILES\tFALLBACKS")
			for _, m := range models {
				replicas := fmt.Sprintf("%d/%d ready", m.Ready, m.Min)
				if m.Total > m.Ready {
					replicas += fmt.Sprintf(" (%d total)", m.Total)
				}
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n",
					m.Name, replicas, join(m.Aliases), join(m.Profiles), join(m.Fallbacks))
			}
			return w.Flush()
		},
	}
}

func newReplicasCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	var model string
	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
			}
			if *asJSON {
				return emitJSON(reps)
			}
			if len(reps) == 0 {
				fmt.Println("No replicas are running.")
				return nil
			}

			w := tab()
			fmt.Fprintln(w, "REPLICA\tMODEL\tSTATE\tPOOL\tNODES\tAGE\tENDPOINT")
			for _, r := range reps {
				var nodes []string
				for _, s := range r.Shards {
					nodes = append(nodes, fmt.Sprintf("%s:%v", s.Node, s.GPUs))
				}
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
					r.ID, r.Model, r.State, r.Pool, join(nodes),
					(time.Duration(r.AgeSec) * time.Second).String(), r.Endpoint)
			}
			_ = 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")
	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 tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
}

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

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

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

func hasUnhealthy(nodes []adminapi.NodeView) bool {
	for _, n := range nodes {
		if len(n.Unhealthy) > 0 {
			return true
		}
	}
	return false
}