llmgw/internal/cli/bench.go

package cli

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"time"

	"github.com/spf13/cobra"

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

func newBenchCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	cfg := bench.DefaultConfig()
	var (
		steps   []int
		target  string
		metrics string
		engine  string
		key     string
		mkey    string
		save    bool
	)

	cmd := &cobra.Command{
		Use:   "bench <model>",
		Short: "Measure what a replica can do, and what the router should be told",
		Long: "Drives real traffic. Point it at a replica to measure that hardware, or at\n" +
			"the gateway to measure the whole path — they answer different questions,\n" +
			"and only the first produces a capacity profile.\n\n" +
			"Concurrency climbs until the SLO breaks. The last level that still met it\n" +
			"is the capacity: peak throughput arrives long after latency has become\n" +
			"unusable, so throughput alone would overstate what the hardware can serve.\n\n" +
			"This costs real GPU time and competes with real traffic. Run it on a\n" +
			"replica you have cordoned, or during a window where that is acceptable.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			cfg.Model = args[0]
			cfg.Target = target
			cfg.MetricsURL = metrics
			cfg.Engine = engine
			cfg.APIKey = key
			cfg.MetricsKey = mkey
			// A replica's engine demands its own key, for its API and — on
			// llama.cpp and TensorRT-LLM — for its metrics. Asked for rather
			// than typed: the admin API knows which replica answers where, and
			// this runs wherever ctl does, which can reach it.
			needKeys := cfg.APIKey == "" || (metrics != "" && cfg.MetricsKey == "")
			if needKeys {
				reps, err := listReplicas(cmd, client)
				if cfg.APIKey == "" {
					cfg.APIKey = replicaKey(reps, target)
				}
				if metrics != "" && cfg.MetricsKey == "" {
					cfg.MetricsKey = replicaKey(reps, metrics)
					if err != nil {
						cmd.PrintErrf("The admin API could not be asked for the engine key of the replica at "+
							"--metrics (%v); an engine that guards its metrics will refuse them. "+
							"Pass --metrics-key to give it.\n", err)
					}
				}
			}
			if len(steps) > 0 {
				cfg.Concurrency = steps
			}

			if metrics == "" {
				cmd.PrintErrln("No --metrics endpoint: queue depth and KV usage will not be " +
					"measured, so target_waiting and kv_ceiling cannot be suggested.")
			}
			cmd.PrintErrf("benchmarking %s against %s\n\n", cfg.Model, cfg.Target)

			// Fixed columns rather than a tabwriter: rows are printed as each
			// level finishes, and a tabwriter cannot align a column it has not
			// seen the rest of yet — it lays out whatever has been flushed, so
			// every row after the first came out ragged.
			const row = "%-5s %-7s %-7s %-9s %-9s %-9s %-6s %-5s %s\n"
			fmt.Fprintf(cmd.ErrOrStderr(), row,
				"CONC", "REQ/S", "TOK/S", "TTFT p50", "TTFT p99", "TPOT p99", "WAIT", "KV", "SLO")

			progress := func(s bench.Step) {
				verdict := "ok"
				if !s.MeetsSLO {
					verdict = "MISSED"
				}
				if s.Failed > 0 {
					verdict = fmt.Sprintf("%d failed", s.Failed)
				}
				// Unread is not zero: a dash, and the reason under the row.
				wait, kv := fmt.Sprintf("%.1f", s.Waiting), fmt.Sprintf("%.0f%%", s.KVUsage*100)
				if s.MetricsUnread != "" {
					wait, kv = "-", "-"
				}
				fmt.Fprintf(cmd.ErrOrStderr(), row,
					fmt.Sprint(s.Concurrency),
					fmt.Sprintf("%.1f", s.RequestsPerSec),
					fmt.Sprintf("%.0f", s.TokensPerSec),
					ms(s.TTFT.P50), ms(s.TTFT.P99), ms(s.TPOT.P99),
					wait, kv, verdict)
				if s.MetricsUnread != "" {
					fmt.Fprintf(cmd.ErrOrStderr(), "      metrics not read: %s\n", s.MetricsUnread)
				}
			}

			res, err := bench.Run(cmd.Context(), cfg, progress)
			var refused *bench.MetricsRefusedError
			if errors.As(err, &refused) {
				return fmt.Errorf("%w. The engine's key is the admin API's, given for the replica "+
					"whose endpoint `llmgw ctl replicas` lists at --metrics' address; "+
					"or pass it with --metrics-key", err)
			}
			if err != nil {
				return err
			}
			if save {
				if err := saveCapacity(cmd, client, res, target); err != nil {
					// Reported, not fatal. The run has already cost the GPU
					// time; losing its conclusion to a failed write would mean
					// running it again, and the numbers are on screen either
					// way.
					cmd.PrintErrf("the measurement was not saved: %v\n", err)
				}
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(res)
			}
			printConclusion(cmd, res, cfg)
			return nil
		},
	}

	f := cmd.Flags()
	f.StringVar(&target, "target", "http://127.0.0.1:8080", "base URL: a replica, or the gateway")
	f.StringVar(&metrics, "metrics", "", "the replica's /metrics, to observe queue depth and KV usage")
	f.StringVar(&engine, "engine", "", "which metric names --metrics uses: vllm | sglang | llamacpp | trtllm (whose are at /prometheus/metrics)")
	f.StringVar(&key, "key", "", "API key, when the gateway requires one")
	f.StringVar(&mkey, "metrics-key", "",
		"the engine key --metrics is read with (default: the replica's, from the admin API)")
	f.IntSliceVar(&steps, "concurrency", nil, "levels to try (default 1,2,4,…,128)")
	f.DurationVar(&cfg.Duration, "duration", cfg.Duration, "how long to hold each level")
	f.IntVar(&cfg.PromptTokens, "prompt-tokens", cfg.PromptTokens, "approximate prompt size")
	f.IntVar(&cfg.MaxTokens, "max-tokens", cfg.MaxTokens, "tokens to generate per request")
	f.BoolVar(&cfg.SharedPrefix, "shared-prefix", false,
		"send every request the same prompt, measuring the cached case instead of the hardware")
	f.DurationVar(&cfg.SLO.TTFTp99, "ttft-p99", cfg.SLO.TTFTp99, "time-to-first-token target")
	f.DurationVar(&cfg.SLO.TPOTp99, "tpot-p99", cfg.SLO.TPOTp99, "per-output-token target")
	f.BoolVar(&save, "save", false,
		"record the capacity against the model and profile, so the number outlives this terminal")
	return cmd
}

// saveCapacity records a run's conclusion against the profile it was measured
// on.
//
// The profile is looked up from the replica rather than asked for. Capacity
// belongs to the hardware and the engine arguments as much as to the weights,
// so a measurement filed under the wrong profile is worse than none — and the
// control plane already knows which profile is answering on which endpoint, so
// nobody needs to be trusted to type it.
func saveCapacity(cmd *cobra.Command, client func() *ctlClient, res bench.Result, target string) error {
	var reps []adminapi.ReplicaView
	if err := client().get(cmd.Context(), "/api/v1/replicas", &reps); err != nil {
		return err
	}
	profile := ""
	for _, rep := range reps {
		if rep.Model == res.Model && sameEndpoint(rep.Endpoint, target) {
			profile = rep.Profile
			break
		}
	}
	if profile == "" {
		// Almost always because --target was the gateway. Benchmarking the
		// whole path is a useful thing to do and says nothing about what one
		// replica can serve, so there is nothing here worth filing as capacity.
		return fmt.Errorf("%s is not a replica of %s, so this run measures the path rather than the hardware; "+
			"point --target at a replica endpoint from `llmgw ctl replicas`", target, res.Model)
	}

	body, err := json.Marshal(res.Capacity(profile))
	if err != nil {
		return err
	}
	var saved bench.Capacity
	if err := client().put(cmd.Context(),
		"/api/v1/models/"+res.Model+"/capacity", body, &saved); err != nil {
		return err
	}
	cmd.PrintErrf("saved: %s on profile %s, %d concurrent at the SLO\n",
		saved.Model, saved.Profile, saved.MaxConcurrencyAtSLO)
	return nil
}

// sameEndpoint compares two base URLs, ignoring a trailing slash.
//
// Written out because the two come from different places: one is what the
// control plane recorded, the other is what somebody typed or pasted, and a
// trailing slash is the difference that would silently make every run look like
// it was aimed at the gateway.
func sameEndpoint(a, b string) bool {
	return strings.TrimSuffix(a, "/") == strings.TrimSuffix(b, "/")
}

func printConclusion(cmd *cobra.Command, res bench.Result, cfg bench.Config) {
	out := cmd.OutOrStdout()
	fmt.Fprintln(out)

	if res.MaxConcurrencyAtSLO == 0 {
		fmt.Fprintf(out, "No level met the SLO (TTFT p99 ≤ %s, TPOT p99 ≤ %s).\n",
			cfg.SLO.TTFTp99, cfg.SLO.TPOTp99)
		if len(res.Steps) > 0 && res.Steps[0].FirstError != "" {
			fmt.Fprintf(out, "First error: %s\n", res.Steps[0].FirstError)
		}
		fmt.Fprintln(out, "Either the hardware cannot serve this model at this SLO, or the")
		fmt.Fprintln(out, "target is stricter than the model can meet at any concurrency.")
		return
	}

	fmt.Fprintf(out, "Capacity: %d concurrent requests at %.0f tokens/s\n",
		res.MaxConcurrencyAtSLO, res.ThroughputTokPerSec)
	fmt.Fprintf(out, "  TTFT p99 %s · TPOT p99 %s · measured over %s\n\n",
		ms(res.TTFTp99), ms(res.TPOTp99), res.Elapsed.Round(time.Second))

	fmt.Fprintln(out, "Put this in the model's profile:")
	fmt.Fprintf(out, "  weight: %d\n", res.MaxConcurrencyAtSLO)
	if res.SuggestedTargetWaiting > 0 || res.SuggestedKVCeiling > 0 {
		fmt.Fprintln(out, "\nAnd in router.scoring, from what the engine reported while still")
		fmt.Fprintln(out, "meeting the target:")
		if res.SuggestedTargetWaiting > 0 {
			fmt.Fprintf(out, "  target_waiting: %d\n", res.SuggestedTargetWaiting)
		}
		if res.SuggestedKVCeiling > 0 {
			fmt.Fprintf(out, "  kv_ceiling: %.2f\n", res.SuggestedKVCeiling)
		}
	}

	if last := res.Steps[len(res.Steps)-1]; !last.MeetsSLO {
		fmt.Fprintf(out, "\n%d concurrent was the first level to miss it (TTFT p99 %s).\n",
			last.Concurrency, ms(last.TTFT.P99))
	} else {
		fmt.Fprintf(out, "\nEvery level tried met the SLO, so %d is a floor rather than a\n",
			res.MaxConcurrencyAtSLO)
		fmt.Fprintln(out, "ceiling — raise --concurrency to find where it actually breaks.")
	}
	if cfg.SharedPrefix {
		fmt.Fprintln(out, "\nRun with --shared-prefix: this measures the engine's prefix cache,")
		fmt.Fprintln(out, "not its hardware. Do not use these numbers as capacity.")
	}
}

func ms(d time.Duration) string {
	switch {
	case d == 0:
		return "-"
	case d < time.Millisecond:
		return d.Round(10 * time.Microsecond).String()
	case d < time.Second:
		return d.Round(time.Millisecond).String()
	default:
		return d.Round(10 * time.Millisecond).String()
	}
}

func listReplicas(cmd *cobra.Command, client func() *ctlClient) ([]adminapi.ReplicaView, error) {
	var reps []adminapi.ReplicaView
	err := client().get(cmd.Context(), "/api/v1/replicas", &reps)
	return reps, err
}

// replicaKey is the engine key of the replica serving at addr, or "" when
// addr is not a replica's (the gateway's, say).
//
// Matched by scheme, host and port rather than by the whole URL, so the same
// lookup serves the target, which is a replica's base URL, and --metrics,
// which is a path under it: /metrics, or /prometheus/metrics on TensorRT-LLM.
func replicaKey(reps []adminapi.ReplicaView, addr string) string {
	for _, rep := range reps {
		if sameOrigin(rep.Endpoint, addr) {
			return rep.EngineKey
		}
	}
	return ""
}

func sameOrigin(a, b string) bool {
	ua, err := url.Parse(a)
	if err != nil || ua.Host == "" {
		return false
	}
	ub, err := url.Parse(b)
	if err != nil {
		return false
	}
	return ua.Scheme == ub.Scheme && ua.Host == ub.Host
}