llmgw/internal/agent/engines/trtllm_test.go

package engines

import (
	"encoding/json"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"
	"testing"
)

func trtllmSpec() ShardSpec {
	s := singleNodeSpec()
	s.Profile.Engine = "trtllm"
	s.Profile.ServedModelName = "llama-3.1-70b"
	return s
}

func TestTRTLLMRunsTrtllmServeWithTheFlagsItDocuments(t *testing.T) {
	s := trtllmSpec()
	s.BindAddr = "10.0.0.147"
	s.Profile.Args = []string{"--max_batch_size", "64"}
	cmd := buildOne(t, &TRTLLM{cfg: TRTLLMConfig{Python: "/opt/trtllm/bin/python"}}, s)

	if cmd.Path != "/opt/trtllm/bin/python" {
		t.Errorf("Path=%q, expected the configured interpreter", cmd.Path)
	}
	if cmd.Args[0] != "-c" || !strings.Contains(cmd.Args[1], "from tensorrt_llm.commands.serve import main") {
		t.Fatalf("not started through the launcher: %v", cmd.Args[:2])
	}
	rest := cmd.Args[2:]
	if rest[0] != "serve" || rest[1] != "/models/llama" {
		t.Errorf("expected serve and the model first, got %v", rest[:2])
	}
	for flag, want := range map[string]string{
		"--host": "10.0.0.147", "--port": "20000", "--tp_size": "4",
		"--served_model_name":     "llama-3.1-70b",
		"--extra_llm_api_options": trtllmOptionsPlaceholder,
	} {
		if got := argValue(rest, flag); got != want {
			t.Errorf("%s=%q, want %q (args %v)", flag, got, want, rest)
		}
	}
	// One pipeline stage is the default, and saying so would only be one
	// more flag for a version to rename.
	if slices.Contains(rest, "--pp_size") {
		t.Errorf("--pp_size given for a single stage: %v", rest)
	}
	// The profile's own, last, so a tuning flag is the operator's to set.
	if got := rest[len(rest)-2:]; !slices.Equal(got, []string{"--max_batch_size", "64"}) {
		t.Errorf("profile args are not last: %v", rest)
	}

	s.Profile.TP, s.Profile.PP = 2, 2
	if got := argValue(buildOne(t, &TRTLLM{}, s).Args, "--pp_size"); got != "2" {
		t.Errorf("--pp_size=%q for two stages", got)
	}
	if got := buildOne(t, &TRTLLM{}, s).Path; got != "python3" {
		t.Errorf("unconfigured, the interpreter is %q", got)
	}
}

func TestTRTLLMsKeyIsNotOnItsCommandLine(t *testing.T) {
	s := trtllmSpec()
	s.APIKey = "k1"
	cmd := buildOne(t, &TRTLLM{}, s)
	for _, a := range cmd.Args {
		if strings.Contains(a, "k1") && a != cmd.Args[1] {
			t.Errorf("the key is in the arguments: %v", cmd.Args)
		}
	}
	if envMap(cmd.Env)[trtllmKeyEnv] != "k1" {
		t.Error("the key is not in the environment")
	}
}

func TestTRTLLMDoesNotReportUsageToNVIDIA(t *testing.T) {
	s := trtllmSpec()
	s.Profile.Env = map[string]string{"TRTLLM_NO_USAGE_STATS": "0", "TRTLLM_ENABLE_PDL": "1"}
	env := envMap(buildOne(t, &TRTLLM{}, s).Env)
	if env["TRTLLM_NO_USAGE_STATS"] != "1" || env["DO_NOT_TRACK"] != "1" {
		t.Errorf("usage reporting is not off: %q %q", env["TRTLLM_NO_USAGE_STATS"], env["DO_NOT_TRACK"])
	}
	// Tuning still reaches it.
	if env["TRTLLM_ENABLE_PDL"] != "1" {
		t.Error("a profile's TensorRT-LLM tuning was dropped")
	}
}

func TestTRTLLMIsReadFromItsPrometheusEndpoint(t *testing.T) {
	// /metrics on trtllm-serve is JSON, and each read drains what it returns.
	if got := (&TRTLLM{}).MetricsPath(); got != "/prometheus/metrics" {
		t.Errorf("MetricsPath=%q", got)
	}
	if !(&TRTLLM{}).ServesHTTP(trtllmSpec()) {
		t.Error("the one rank of a replica serves its API")
	}
}

func TestTRTLLMRunsAReplicaOnOneNode(t *testing.T) {
	p := trtllmSpec().Profile
	p.Nodes, p.GPUsPerNode, p.TP = 2, 4, 8
	err := (&TRTLLM{}).Validate(p, nil)
	if err == nil || !strings.Contains(err.Error(), "one node") {
		t.Errorf("a two-node profile was accepted: %v", err)
	}
	s := multiNodeSpec(0)
	s.Profile.Engine = "trtllm"
	if _, err := (&TRTLLM{}).Commands(s); err == nil {
		t.Error("a command was built for a shard of a two-node replica")
	}
}

func TestTRTLLMNeedsAmpereOrNewer(t *testing.T) {
	p := trtllmSpec().Profile
	p.TP = 2
	v100 := []GPUInfo{{Index: 0, Name: "Tesla V100", ComputeCapability: "7.0"}, {Index: 1, Name: "Tesla V100", ComputeCapability: "7.0"}}
	if err := (&TRTLLM{}).Validate(p, v100); err == nil {
		t.Error("a V100 was accepted")
	}
	h100 := []GPUInfo{{Index: 0, Name: "H100", ComputeCapability: "9.0"}, {Index: 1, Name: "H100", ComputeCapability: "9.0"}}
	if err := (&TRTLLM{}).Validate(p, h100); err != nil {
		t.Errorf("two H100s were refused: %v", err)
	}
}

func TestTRTLLMProfilesMayNotOpenWhatTheAgentCloses(t *testing.T) {
	for _, args := range [][]string{
		{"--trust_remote_code"},
		{"--num_serve_frontends", "4"},
		{"--grpc"},
		{"--config", "/etc/shadow"},
		{"--extra_llm_api_options=/tmp/x.yaml"},
		{"--served_model_name", "other"},
		{"--middleware", "evil.mw"},
		{"--tp_size", "8"},
	} {
		if err := UnsafeArgs("trtllm", args); err == nil {
			t.Errorf("%v was allowed", args)
		}
	}
	if err := UnsafeArgs("trtllm", []string{"--max_batch_size", "64", "--set", "kv_cache_config.enable_block_reuse=false",
		"--free_gpu_memory_fraction", "0.85"}); err != nil {
		t.Errorf("tuning was refused: %v", err)
	}
}

func TestTRTLLMIsOfferedOnlyWhenItsPackageIsThere(t *testing.T) {
	there := script(t, "python3", `echo "1.3.0rc29"`)
	r := NewRegistry(Config{})
	r.engines = map[string]Engine{"trtllm": &TRTLLM{cfg: TRTLLMConfig{Python: there}}}
	if got := r.ProbeConfigured(t.Context()); got["trtllm"] != "1.3.0rc29" {
		t.Errorf("expected the version, got %v", got)
	}

	// An interpreter without the package, which every machine has.
	absent := script(t, "python3", "exit 1")
	r.engines = map[string]Engine{"trtllm": &TRTLLM{cfg: TRTLLMConfig{Python: absent}}}
	if got := r.ProbeConfigured(t.Context()); len(got) != 0 {
		t.Errorf("reported without its package: %v", got)
	}

	// Not named, not offered: python3 on PATH is nobody's TensorRT-LLM.
	r.engines = map[string]Engine{"trtllm": &TRTLLM{}}
	if got := r.ProbeConfigured(t.Context()); len(got) != 0 {
		t.Errorf("offered without being configured: %v", got)
	}
}

// fakeTRTLLM is a Python path holding stand-ins for fastapi and
// tensorrt_llm: enough of each for the launcher to do what it does in front
// of the real server, and for the stand-in server to report what it saw.
func fakeTRTLLM(t *testing.T) string {
	t.Helper()
	dir := t.TempDir()
	files := map[string]string{
		"fastapi/__init__.py": `
class FastAPI:
    def __init__(self, *a, **kw):
        self.user_middleware = []
    def add_middleware(self, cls, **kw):
        self.user_middleware.append(cls)
`,
		"tensorrt_llm/__init__.py":          "",
		"tensorrt_llm/commands/__init__.py": "",
		"tensorrt_llm/commands/serve.py": `
import asyncio, json, os, sys
import fastapi

def main():
    args = sys.argv[1:]
    with open(args[args.index("--extra_llm_api_options") + 1]) as f:
        options = f.read()
    app = fastapi.FastAPI()

    async def server(scope, receive, send):
        await send({"type": "http.response.start", "status": 200, "headers": []})
        await send({"type": "http.response.body", "body": b"ok"})

    handler = server
    for mw in reversed(app.user_middleware):
        handler = mw(handler)

    async def ask(path, auth=None):
        headers = [(b"authorization", auth.encode())] if auth else []
        sent = []
        async def send(msg):
            sent.append(msg)
        async def receive():
            return {"type": "http.request"}
        await handler({"type": "http", "path": path, "headers": headers}, receive, send)
        return sent[0]["status"]

    async def run():
        return {
            "health": await ask("/health"),
            "anonymous": await ask("/v1/chat/completions"),
            "wrong": await ask("/v1/models", "Bearer nope"),
            "keyed": await ask("/prometheus/metrics", "Bearer k1"),
        }

    out = asyncio.run(run())
    out.update(options=options, argv0=sys.argv[0], args=args,
               key_left=os.environ.get("LLMGW_ENGINE_KEY"))
    print(json.dumps(out))
`,
	}
	for name, body := range files {
		p := filepath.Join(dir, name)
		if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(p, []byte(body), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	return dir
}

type launched struct {
	Health, Anonymous, Wrong, Keyed int
	Options, Argv0                  string
	Args                            []string
	KeyLeft                         *string `json:"key_left"`
}

func launch(t *testing.T, s ShardSpec) launched {
	t.Helper()
	python, err := exec.LookPath("python3")
	if err != nil {
		t.Skip("no python3 here to run the launcher with")
	}
	cmd := buildOne(t, &TRTLLM{cfg: TRTLLMConfig{Python: python}}, s)
	c := exec.Command(cmd.Path, cmd.Args...) //nolint:gosec // the command under test
	c.Env = append(cmd.Env, "PYTHONPATH="+fakeTRTLLM(t))
	out, err := c.Output()
	if err != nil {
		var stderr string
		if ee, ok := err.(*exec.ExitError); ok {
			stderr = string(ee.Stderr)
		}
		t.Fatalf("the launcher failed: %v\n%s", err, stderr)
	}
	var got launched
	if err := json.Unmarshal(out, &got); err != nil {
		t.Fatalf("%v: %s", err, out)
	}
	return got
}

func TestTheTRTLLMLauncherDemandsTheKey(t *testing.T) {
	s := trtllmSpec()
	s.APIKey = "k1"
	got := launch(t, s)

	if got.Anonymous != 401 || got.Wrong != 401 {
		t.Errorf("served without the key: anonymous %d, wrong key %d", got.Anonymous, got.Wrong)
	}
	if got.Keyed != 200 {
		t.Errorf("refused the key: %d", got.Keyed)
	}
	// Up or down is all /health says, and it is asked without one.
	if got.Health != 200 {
		t.Errorf("/health answered %d", got.Health)
	}
	// Taken out of the environment, so nothing the server starts holds it.
	if got.KeyLeft != nil {
		t.Errorf("the key was left in the environment: %q", *got.KeyLeft)
	}
	// A copy of this command line, as trtllm-serve's extra front ends are,
	// names no program that exists and so cannot serve without the key.
	if _, err := os.Stat(got.Argv0); err == nil || filepath.Base(got.Argv0) != "trtllm-serve" {
		t.Errorf("argv[0] is %q", got.Argv0)
	}
}

func TestTheTRTLLMLauncherTurnsOnTheMetricsTheRouterReads(t *testing.T) {
	got := launch(t, trtllmSpec())
	for _, want := range []string{"return_perf_metrics: true", "enable_iter_perf_stats: true"} {
		if !strings.Contains(got.Options, want) {
			t.Errorf("the options file lacks %q: %q", want, got.Options)
		}
	}
	path := argValue(got.Args, "--extra_llm_api_options")
	if !strings.HasPrefix(path, "/dev/fd/") {
		t.Errorf("the options were not passed by descriptor: %q", path)
	}
	if slices.Contains(got.Args, trtllmOptionsPlaceholder) {
		t.Error("the placeholder reached trtllm-serve")
	}
	// Without a key there is nothing to demand.
	if got.Anonymous != 200 {
		t.Errorf("with no key set, a request was refused: %d", got.Anonymous)
	}
}