llmgw/internal/cli/enginecheck_test.go

package cli

import (
	"bytes"
	"context"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestAFlagIsNotSatisfiedByALongerOne(t *testing.T) {
	// The failure this guards against is silent: --port would be reported as
	// present because --port-range is documented, and a renamed flag would sail
	// through the check that exists to catch renamed flags.
	help := "  --port-range N   a range\n  --model-name S   a name\n"

	if mentions(help, "--port") {
		t.Error("--port was satisfied by --port-range")
	}
	if mentions(help, "--model") {
		t.Error("--model was satisfied by --model-name")
	}
	if !mentions(help, "--port-range") {
		t.Error("--port-range was not found")
	}
}

func TestFlagsAreFoundInTheShapesHelpTextUses(t *testing.T) {
	help := `
  -m, --model FNAME         model path
      --host HOST           ip address
      --parallel, -np N     slots
      --metrics             enable prometheus
  --alias=NAME              served name
`
	for _, flag := range []string{"--model", "--host", "--parallel", "--metrics", "--alias"} {
		if !mentions(help, flag) {
			t.Errorf("%s was not found in a help text that documents it", flag)
		}
	}
	for _, flag := range []string{"--nonsense", "--mode", "--para"} {
		if mentions(help, flag) {
			t.Errorf("%s was found in a help text that does not document it", flag)
		}
	}
}

func TestTheRepresentativeShardIsTheWidestOne(t *testing.T) {
	// A flag that only appears in the multi-node case is exactly the one nobody
	// notices has been renamed, so the command built for checking has to be the
	// widest llmgw would ever produce.
	spec := representativeShard("vllm", 2, 8, 4)
	if spec.NNodes != 2 || spec.Profile.Nodes != 2 {
		t.Errorf("not a multi-node shard: %+v", spec)
	}
	if len(spec.GPUIndices) != 4 {
		t.Errorf("%d GPUs, want 4", len(spec.GPUIndices))
	}
	if spec.LeaderAddr == "" || spec.LeaderPort == 0 {
		t.Error("no rendezvous point, so the flags that carry it will not appear")
	}
}

func TestATruncatedHelpIsSuspected(t *testing.T) {
	// Found by truncating llama.cpp's 725-line help to 400 lines and being told,
	// confidently, that five flags the running cluster uses do not exist.
	full := "  --help    show this\n  --model F   path\n"
	if !mentions(full, "--help") {
		t.Fatal("the guard's own signal is missing from a complete help")
	}
	truncated := "  --model F   path\n"
	if mentions(truncated, "--help") {
		t.Error("a truncated help was taken for a complete one")
	}
}

// fakeInterpreter stands in for TensorRT-LLM's Python: asked for
// trtllm-serve's help through its module, it prints the help given; asked
// anything else, it fails, as running the launcher here would.
func fakeInterpreter(t *testing.T, help string) string {
	t.Helper()
	return fakePython(t, "-m tensorrt_llm.commands.serve serve --help", help)
}

// fakePython is an interpreter that prints help when asked exactly ask, and
// fails on anything else — the interpreter's own --help included.
func fakePython(t *testing.T, ask, help string) string {
	t.Helper()
	dir := t.TempDir()
	if err := os.WriteFile(filepath.Join(dir, "help.txt"), []byte(help), 0o644); err != nil {
		t.Fatal(err)
	}
	p := filepath.Join(dir, "python")
	body := "#!/bin/sh\n" +
		`if [ "$*" = "` + ask + `" ]; then cat "` + filepath.Join(dir, "help.txt") + `"; exit 0; fi` + "\n" +
		"echo \"unexpected: $*\" >&2; exit 2\n"
	if err := os.WriteFile(p, []byte(body), 0o755); err != nil { //nolint:gosec // an executable stand-in
		t.Fatal(err)
	}
	return p
}

// trtllmHelp is the part of trtllm-serve serve --help (1.3) that holds the
// flags llmgw builds.
const trtllmHelp = `Usage: trtllm-serve serve [OPTIONS] MODEL

Options:
  --host TEXT                     Hostname of the server.
  --port INTEGER                  Port of the server.
  --tensor_parallel_size, --tp_size INTEGER
                                  Tensor parallelism size.
  --pipeline_parallel_size, --pp_size INTEGER
                                  Pipeline parallelism size.
  --config, --extra_llm_api_options TEXT
                                  Path to a YAML configuration file.
  --served_model_name TEXT        The model name used in the API.
  --help                          Show this message and exit.
`

func runEngineCheck(t *testing.T, args ...string) (string, error) {
	t.Helper()
	c := newEngineCheckCmd()
	var out bytes.Buffer
	c.SetOut(&out)
	c.SetErr(&out)
	c.SetArgs(args)
	err := c.ExecuteContext(context.Background())
	return out.String(), err
}

func TestEngineCheckAsksTrtllmServeThroughItsModule(t *testing.T) {
	out, err := runEngineCheck(t, "trtllm", "--bin", fakeInterpreter(t, trtllmHelp))
	if err != nil {
		t.Fatalf("%v\n%s", err, out)
	}
	if !strings.Contains(out, "every flag llmgw builds is accepted") {
		t.Errorf("unexpected report:\n%s", out)
	}
	// The interpreter's own help would describe Python, not trtllm-serve; and
	// one node is what trtllm runs, so no multi-node flag is in the command.
	if strings.Contains(out, "unexpected:") || strings.Contains(out, "skipped") {
		t.Errorf("the help was not asked of trtllm-serve:\n%s", out)
	}
}

// sglangHelp is the part of python -m sglang.launch_server --help (0.5.21)
// that holds the flags llmgw builds, in argparse's layout.
const sglangHelp = `usage: sglang serve [-h] --model-path MODEL_PATH [--host HOST] [--port PORT]

options:
  -h, --help            show this help message and exit
  --model-path MODEL_PATH, --model MODEL_PATH
                        Path to model weights.
  --host HOST           The host of the HTTP server.
  --port PORT           The port of the HTTP server.
  --served-model-name SERVED_MODEL_NAME
                        Override the model name returned by the v1/models endpoint.
  --enable-metrics      Enable log prometheus metrics.
  --dist-init-addr DIST_INIT_ADDR, --nccl-init-addr DIST_INIT_ADDR
                        The host address for initializing distributed backend.
  --nnodes NNODES       The number of nodes.
  --node-rank NODE_RANK
                        The node rank.
  --tp-size TP_SIZE, --tensor-parallel-size TP_SIZE
                        The tensor parallelism size.
`

func TestEngineCheckAsksSGLangThroughItsModule(t *testing.T) {
	// The command is the interpreter running llmgw's launcher; asked for
	// --help it would describe Python, and every SGLang flag would look
	// missing — or, against a real python3, be checked against nothing.
	py := fakePython(t, "-m sglang.launch_server --help", sglangHelp)
	out, err := runEngineCheck(t, "sglang", "--bin", py)
	if err != nil {
		t.Fatalf("%v\n%s", err, out)
	}
	if !strings.Contains(out, "every flag llmgw builds is accepted") {
		t.Errorf("unexpected report:\n%s", out)
	}
	if strings.Contains(out, "unexpected:") || strings.Contains(out, "skipped") {
		t.Errorf("the help was not asked of SGLang:\n%s", out)
	}
	// Two nodes by default, so the rendezvous flags were among those checked.
	if !strings.Contains(out, "sglang: 9 flags checked") {
		t.Errorf("not every flag was checked:\n%s", out)
	}
}

func TestEngineCheckCatchesAnSGLangWithoutTheFlagForItsMetrics(t *testing.T) {
	old := strings.ReplaceAll(sglangHelp, "  --enable-metrics      Enable log prometheus metrics.\n", "")
	out, err := runEngineCheck(t, "sglang", "--bin", fakePython(t, "-m sglang.launch_server --help", old))
	if err == nil || !strings.Contains(out, "sglang: --enable-metrics") {
		t.Errorf("an SGLang without --enable-metrics passed: %v\n%s", err, out)
	}
}

func TestEngineCheckCatchesATensorRTLLMTooOldForTheServedName(t *testing.T) {
	// TensorRT-LLM 1.2 has no --served_model_name, and a replica of it would
	// never be seen serving the name the agent waits for.
	old := strings.ReplaceAll(trtllmHelp, "  --served_model_name TEXT        The model name used in the API.\n", "")
	out, err := runEngineCheck(t, "trtllm", "--bin", fakeInterpreter(t, old))
	if err == nil || !strings.Contains(out, "trtllm-serve: --served_model_name") {
		t.Errorf("an old trtllm-serve passed: %v\n%s", err, out)
	}
}