llmgw/internal/cli/facility.go

package cli

import (
	"fmt"
	"io"
	"maps"
	"net/http"
	"net/url"
	"slices"
	"strings"

	"github.com/spf13/cobra"

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

// newFacilityCmd shows the failure domains and holds them for maintenance.
func newFacilityCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "facility",
		Short: "Racks, PDUs and rooms: their power against their caps, and what is held",
		Long: `Where nodes are comes from their labels: llmgw.io/site, room, row, rack
and pdu, set in agent.yaml or with 'llmgw ctl label'. Placement spreads a
model's replicas across racks and PDUs, places nothing in a held domain, and
refuses a replica that would take a domain past its power cap (facility.budgets
in server.yaml).

USED is what counts against the cap: node by node, the larger of what the
replicas placed there are estimated to draw and what the node's sensors read.`,
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			var v adminapi.FacilityView
			if err := client().get(cmd.Context(), "/api/v1/facility", &v); err != nil {
				return err
			}
			if *asJSON {
				return emitJSON(v)
			}
			printFacility(cmd.OutOrStdout(), v)
			return nil
		},
	}
	cmd.AddCommand(newFacilityHoldCmd(client), newFacilityReleaseCmd(client))
	return cmd
}

// printFacility is the facility view as a table, then what is held, then
// the nodes nobody has said the place of.
func printFacility(out io.Writer, v adminapi.FacilityView) {
	if len(v.Domains) == 0 {
		fmt.Fprintln(out, "No node says where it is. Label them with llmgw.io/rack (and site, room, row, pdu):\n"+
			"  llmgw ctl label gpu01 llmgw.io/rack=r12-07 llmgw.io/pdu=r12-07-a,r12-07-b")
		return
	}
	w := tabTo(out)
	fmt.Fprintln(w, "DOMAIN\tNODES\tUSED\tLIMIT\tRESERVED\tMEASURED\tHELD")
	for _, d := range v.Domains {
		limit, used := "-", facility.KW(d.UsedWatts)
		if d.LimitWatts > 0 {
			limit = facility.KW(d.LimitWatts)
			used += fmt.Sprintf(" (%.0f%%)", 100*d.UsedWatts/d.LimitWatts)
		}
		held := "-"
		if d.Held != nil {
			held = "held"
			if d.Held.Reason != "" {
				held += ": " + d.Held.Reason
			}
			if d.Held.Drain {
				held += " (drained)"
			}
		}
		fmt.Fprintf(w, "%s\t%d\t%s\t%s\t%s\t%s\t%s\n", d.Domain, len(d.Nodes), used, limit,
			facility.KW(d.ReservedWatts), facility.KW(d.MeasuredWatts), held)
	}
	_ = w.Flush()
	for _, h := range v.Holds {
		fmt.Fprintf(out, "\n%s held by %s since %s", h.Domain, or(h.By, "someone"), h.At.Local().Format("2006-01-02 15:04"))
	}
	if len(v.Holds) > 0 {
		fmt.Fprintln(out)
	}
	if len(v.Unlocated) > 0 {
		fmt.Fprintf(out, "\nNo location: %s. Each is placed as a rack of its own.\n", strings.Join(v.Unlocated, ", "))
	}
}

func newFacilityHoldCmd(client func() *ctlClient) *cobra.Command {
	var req adminapi.HoldRequest
	c := &cobra.Command{
		Use:   "hold <level:name>",
		Short: "Place nothing new in a rack, row, PDU or room; with --drain, move its work off",
		Long: `For maintenance, a cooling fault, work on a PDU: like cordoning every node in
the domain, including nodes labelled into it later, and kept across restarts
when the control plane has a store. What runs there keeps running unless
--drain is given; then each of its nodes is drained, one replica at a time
with the replacement up elsewhere first.`,
		Example: `  llmgw ctl facility hold rack:r12-07 --reason "PDU A replacement"
  llmgw ctl facility hold pdu:r12-07-a --drain --reason "breaker work"`,
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			req.Domain = args[0]
			if _, err := facility.ParseDomain(req.Domain); err != nil {
				return err
			}
			var v adminapi.HoldView
			if err := client().do(cmd.Context(), http.MethodPost, "/api/v1/facility/holds", req, &v); err != nil {
				return err
			}
			fmt.Printf("%s is held: nothing new is placed in it.\n", v.Domain)
			if len(v.Nodes) > 0 {
				fmt.Printf("Draining %s; their replicas move one at a time.\n", strings.Join(v.Nodes, ", "))
			}
			fmt.Printf("Release it with: llmgw ctl facility release %s\n", v.Domain)
			return nil
		},
	}
	c.Flags().StringVar(&req.Reason, "reason", "", "why, as explain-placement and the node list will say")
	c.Flags().BoolVar(&req.Drain, "drain", false, "also drain the domain's nodes")
	return c
}

func newFacilityReleaseCmd(client func() *ctlClient) *cobra.Command {
	return &cobra.Command{
		Use:   "release <level:name>",
		Short: "Allow placement in a held domain again",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			d, err := facility.ParseDomain(args[0])
			if err != nil {
				return err
			}
			var v adminapi.HoldView
			if err := client().delete(cmd.Context(), "/api/v1/facility/holds/"+url.PathEscape(string(d)), &v); err != nil {
				return err
			}
			fmt.Printf("%s is released.\n", v.Domain)
			if len(v.Nodes) > 0 {
				fmt.Printf("Ready again: %s.\n", strings.Join(v.Nodes, ", "))
			}
			return nil
		},
	}
}

// newLabelCmd sets labels on a node through the control plane.
func newLabelCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "label <node> key=value... [key-...]",
		Short: "Set or remove a node's labels, over those in its agent.yaml",
		Long: `Labels set here are kept by the control plane and win over the agent's own;
key- removes one, which puts the agent's value for the key back if it has
one. The node need not have joined yet.

The llmgw.io/ labels say where a node is, for facility-aware placement:
site, room, row, rack, and pdu (a comma separated list for a dual-fed rack).
Any other key is free for a profile's node_selector.`,
		Example: `  llmgw ctl label gpu01 llmgw.io/rack=r12-07 llmgw.io/pdu=r12-07-a,r12-07-b
  llmgw ctl label gpu01 llmgw.io/row=12 llmgw.io/room=hall-b
  llmgw ctl label gpu01 llmgw.io/rack-`,
		Args: cobra.MinimumNArgs(2),
		RunE: func(cmd *cobra.Command, args []string) error {
			req, err := parseLabelArgs(args[1:])
			if err != nil {
				return err
			}
			var v adminapi.NodeLabelsView
			if err := client().do(cmd.Context(), http.MethodPost,
				"/api/v1/nodes/"+url.PathEscape(args[0])+"/labels", req, &v); err != nil {
				return err
			}
			if *asJSON {
				return emitJSON(v)
			}
			if v.Labels == nil {
				fmt.Printf("%s has not joined yet; it takes these labels when it does.\n", v.ID)
			}
			keys := slices.Sorted(maps.Keys(v.OperatorLabels))
			if len(keys) == 0 {
				fmt.Printf("%s carries no labels set here, only its agent's.\n", v.ID)
			}
			for _, k := range keys {
				fmt.Printf("%s: %s=%s\n", v.ID, k, v.OperatorLabels[k])
			}
			if v.Location != nil {
				fmt.Printf("%s is at %s.\n", v.ID, v.Location.Short())
			}
			return nil
		},
	}
}

// parseLabelArgs reads key=value to set and key- to remove.
func parseLabelArgs(args []string) (adminapi.NodeLabelsRequest, error) {
	req := adminapi.NodeLabelsRequest{Set: map[string]string{}}
	for _, a := range args {
		if k, v, ok := strings.Cut(a, "="); ok {
			if k == "" || v == "" {
				return req, fmt.Errorf("%q: give key=value, or key- to remove a label", a)
			}
			req.Set[k] = v
			continue
		}
		if k, ok := strings.CutSuffix(a, "-"); ok && k != "" {
			req.Remove = append(req.Remove, k)
			continue
		}
		return req, fmt.Errorf("%q: give key=value, or key- to remove a label", a)
	}
	return req, nil
}