llmgw/internal/cli/bench.go

package cli

import (
	"encoding/json"
	"fmt"
	"text/tabwriter"
	"time"

	"github.com/spf13/cobra"

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

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

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

			w := tabwriter.NewWriter(cmd.ErrOrStderr(), 0, 0, 2, ' ', 0)
			fmt.Fprintln(w, "CONC\tREQ/S\tTOK/S\tTTFT p50\tTTFT p99\tTPOT p99\tWAIT\tKV\tSLO")
			progress := func(s bench.Step) {
				verdict := "ok"
				if !s.MeetsSLO {
					verdict = "MISSED"
				}
				if s.Failed > 0 {
					verdict = fmt.Sprintf("%d failed", s.Failed)
				}
				fmt.Fprintf(w, "%d\t%.1f\t%.0f\t%s\t%s\t%s\t%.1f\t%.0f%%\t%s\n",
					s.Concurrency, s.RequestsPerSec, s.TokensPerSec,
					ms(s.TTFT.P50), ms(s.TTFT.P99), ms(s.TPOT.P99),
					s.Waiting, s.KVUsage*100, verdict)
				_ = w.Flush()
			}

			res, err := bench.Run(cmd.Context(), cfg, progress)
			if err != nil {
				return 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")
	f.StringVar(&key, "key", "", "API key, when the gateway requires one")
	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")
	return cmd
}

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