llmgw/internal/cli/quota.go

package cli

import (
	"encoding/json"
	"fmt"
	"net/url"
	"strconv"
	"strings"

	"github.com/spf13/cobra"

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

func newQuotaCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	c := &cobra.Command{
		Use:   "quota",
		Short: "What parts of the tenancy tree are guaranteed of the cluster, and may hold",
		Long: "A capacity quota belongs to an organization, team or project, and counts the\n" +
			"GPUs, CPU cores and memory held by replicas of the models it owns (see\n" +
			"`llmgw ctl model owner`), and of those its teams and projects own.\n\n" +
			"The guarantee is kept for it: when it holds less and a model of its needs\n" +
			"room, replicas other scopes hold beyond their own guarantees are stopped to\n" +
			"make it — whatever their class. What it holds within its guarantee is\n" +
			"never taken by another scope's model. Beyond it, it borrows: free room is\n" +
			"used, and given back when asked for.\n\n" +
			"The limit is the most it may hold, borrowing included; a model that would\n" +
			"take it past is not given another replica.\n\n" +
			"Scopes are written org:Acme, team:Acme/ml or project:Acme/ml/training, or\n" +
			"with IDs. Changes take effect within half a minute.",
	}
	c.AddCommand(newQuotaListCmd(client, asJSON), newQuotaSetCmd(client), newQuotaClearCmd(client))
	return c
}

func newQuotaListCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "list",
		Short: "Every quota, with what its scope holds now",
		RunE: func(cmd *cobra.Command, _ []string) error {
			var qs []adminapi.QuotaView
			if err := client().get(cmd.Context(), "/api/v1/quotas", &qs); err != nil {
				return err
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(qs)
			}
			if len(qs) == 0 {
				fmt.Fprintln(cmd.OutOrStdout(), "No capacity quotas: everything shares the cluster by class alone.")
				return nil
			}
			w := tab()
			fmt.Fprintln(w, "SCOPE\tNAME\tGPUS\tCORES\tMEMORY GB")
			for _, q := range qs {
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", q.Scope, q.Name,
					quotaCell(q.Holds.GPUs, q.Guarantee.GPUs, q.Limit.GPUs),
					quotaCell(q.Holds.CPUCores, q.Guarantee.CPUCores, q.Limit.CPUCores),
					quotaCell(q.Holds.MemoryGB, q.Guarantee.MemoryGB, q.Limit.MemoryGB))
			}
			if err := w.Flush(); err != nil {
				return err
			}
			fmt.Fprintln(cmd.OutOrStdout(), "\nEach cell: holds / guaranteed (at most limit). - is not governed.")
			return nil
		},
	}
}

// quotaCell is "6 / 8 (≤ 16)": what is held, what is guaranteed, the limit.
func quotaCell(holds int, guarantee, limit *int) string {
	if guarantee == nil && limit == nil {
		return "-"
	}
	out := strconv.Itoa(holds) + " / "
	if guarantee != nil {
		out += strconv.Itoa(*guarantee)
	} else {
		out += "-"
	}
	if limit != nil {
		out += fmt.Sprintf(" (≤ %d)", *limit)
	}
	return out
}

func newQuotaSetCmd(client func() *ctlClient) *cobra.Command {
	var gpus, cores, mem, lgpus, lcores, lmem int
	c := &cobra.Command{
		Use:   "set SCOPE [--gpus N] [--cores N] [--memory GB] [--limit-gpus N] ...",
		Short: "Set a scope's guarantee and limit, replacing what it had",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			f := cmd.Flags()
			pick := func(name string, v int) *int {
				if !f.Changed(name) {
					return nil
				}
				return &v
			}
			req := adminapi.QuotaRequest{Scope: args[0],
				Guarantee: store.Resources{GPUs: pick("gpus", gpus), CPUCores: pick("cores", cores), MemoryGB: pick("memory", mem)},
				Limit:     store.Resources{GPUs: pick("limit-gpus", lgpus), CPUCores: pick("limit-cores", lcores), MemoryGB: pick("limit-memory", lmem)},
			}
			body, _ := json.Marshal(req)
			var res map[string]string
			if err := client().put(cmd.Context(), "/api/v1/quotas", body, &res); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s: set; it takes effect within half a minute\n", res["scope"])
			return nil
		},
	}
	f := c.Flags()
	f.IntVar(&gpus, "gpus", 0, "GPUs guaranteed")
	f.IntVar(&cores, "cores", 0, "CPU cores guaranteed")
	f.IntVar(&mem, "memory", 0, "memory guaranteed, in GB")
	f.IntVar(&lgpus, "limit-gpus", 0, "the most GPUs it may hold")
	f.IntVar(&lcores, "limit-cores", 0, "the most CPU cores it may hold")
	f.IntVar(&lmem, "limit-memory", 0, "the most memory it may hold, in GB")
	return c
}

func newQuotaClearCmd(client func() *ctlClient) *cobra.Command {
	return &cobra.Command{
		Use:   "clear SCOPE",
		Short: "Remove a scope's quota",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var res map[string]string
			if err := client().delete(cmd.Context(), "/api/v1/quotas/"+url.PathEscape(args[0]), &res); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s: quota removed\n", res["removed"])
			return nil
		},
	}
}

func newOwnerCmd(client func() *ctlClient) *cobra.Command {
	return &cobra.Command{
		Use:   "owner <model> SCOPE|shared",
		Short: "Set whose capacity quota a model counts against",
		Long: "SCOPE is org:Acme, team:Acme/ml or project:Acme/ml/training. A shared\n" +
			"model belongs to the cluster and counts against nobody's quota.",
		Args: cobra.ExactArgs(2),
		RunE: func(cmd *cobra.Command, args []string) error {
			owner := args[1]
			if owner == "shared" {
				owner = ""
			}
			body, _ := json.Marshal(map[string]string{"owner": owner})
			var res map[string]string
			if err := client().post(cmd.Context(), "/api/v1/models/"+args[0]+"/owner", body, &res); err != nil {
				return err
			}
			if res["owner"] == "" {
				fmt.Fprintf(cmd.OutOrStdout(), "%s is shared\n", res["model"])
			} else {
				fmt.Fprintf(cmd.OutOrStdout(), "%s counts against %s\n", res["model"], res["owner"])
			}
			return nil
		},
	}
}

func newTargetCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	var tps float64
	var apply bool
	c := &cobra.Command{
		Use:   "target <model> --tokens-per-sec N [--apply]",
		Short: "How many replicas a model needs for a throughput, and whether they fit",
		Long: "Sized from what one replica was measured to serve while meeting its SLO\n" +
			"(`llmgw ctl bench MODEL --save`), on the profile its replicas run on.\n" +
			"--apply makes that the number it always runs, if the cluster has room.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if tps <= 0 {
				return fmt.Errorf("--tokens-per-sec is required")
			}
			var v adminapi.TargetView
			path := "/api/v1/models/" + url.PathEscape(args[0]) + "/target?tokens_per_sec=" +
				strconv.FormatFloat(tps, 'f', -1, 64)
			if err := client().get(cmd.Context(), path, &v); err != nil {
				return err
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(v)
			}
			out := cmd.OutOrStdout()
			fmt.Fprintf(out, "%s: %.0f tokens/s needs %d replicas at %.1f each (measured on %s); %d running\n",
				v.Model, v.TokPerSec, v.Needed, v.PerReplica, v.Profile, v.Current)
			switch {
			case v.More == 0:
				fmt.Fprintln(out, "It already has them.")
			case v.Fits && len(v.Preempts) > 0:
				fmt.Fprintf(out, "%d more; there is room once replicas of %s are stopped, which its class may do.\n",
					v.More, strings.Join(v.Preempts, ", "))
			case v.Fits:
				fmt.Fprintf(out, "%d more, and there is room for them.\n", v.More)
			default:
				fmt.Fprintf(out, "%d more, and there is room for %d:\n", v.More, v.Room)
				for _, b := range v.Blocked {
					fmt.Fprintln(out, "  "+b)
				}
			}
			if !apply || v.More == 0 {
				return nil
			}
			if !v.Fits {
				return fmt.Errorf("not applied: the cluster has no room for them")
			}
			var models []adminapi.ModelView
			if err := client().get(cmd.Context(), "/api/v1/models", &models); err != nil {
				return err
			}
			top := v.Needed
			for _, m := range models {
				if m.Name == v.Model {
					top = max(top, m.Max)
				}
			}
			body, _ := json.Marshal(adminapi.ScaleRequest{Min: v.Needed, Max: top})
			if err := client().post(cmd.Context(), "/api/v1/models/"+url.PathEscape(v.Model)+"/scale", body, nil); err != nil {
				return err
			}
			fmt.Fprintf(out, "Applied: %s runs %d (up to %d under load).\n", v.Model, v.Needed, top)
			return nil
		},
	}
	c.Flags().Float64Var(&tps, "tokens-per-sec", 0, "the throughput to size for")
	c.Flags().BoolVar(&apply, "apply", false, "set it as the number of replicas the model always runs")
	return c
}