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()
	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 [ "$*" = "-m tensorrt_llm.commands.serve serve --help" ]; 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)
	}
}

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