llmgw/internal/agent/engines/engines_test.go

package engines

import (
	"fmt"
	"slices"
	"strconv"
	"strings"
	"testing"
)

// buildOne returns the single command a shard produces, failing if it produced
// a different number.
func buildOne(t *testing.T, e Engine, s ShardSpec) Command {
	t.Helper()
	cmds := build(t, e, s)
	if len(cmds) != 1 {
		t.Fatalf("expected 1 command, got %d: %s", len(cmds), summarize(cmds))
	}
	return cmds[0]
}

// buildEngine returns the last command, which is the engine itself. The ones
// before it are sidecars.
func buildEngine(t *testing.T, e Engine, s ShardSpec) Command {
	t.Helper()
	cmds := build(t, e, s)
	return cmds[len(cmds)-1]
}

func build(t *testing.T, e Engine, s ShardSpec) []Command {
	t.Helper()
	cmds, err := e.Commands(s)
	if err != nil {
		t.Fatal(err)
	}
	if len(cmds) == 0 {
		t.Fatal("no commands produced")
	}
	return cmds
}

// summarize keeps failure output readable: a Command carries the whole
// environment, and printing it buries the actual problem.
func summarize(cmds []Command) string {
	var b strings.Builder
	for i, c := range cmds {
		if i > 0 {
			b.WriteString(" | ")
		}
		fmt.Fprintf(&b, "%s: %s %s", c.Name, c.Path, strings.Join(c.Args, " "))
	}
	return b.String()
}

// envMap turns exec-form environment back into a map for assertions.
func envMap(env []string) map[string]string {
	m := map[string]string{}
	for _, kv := range env {
		if k, v, ok := strings.Cut(kv, "="); ok {
			m[k] = v
		}
	}
	return m
}

// argValue returns the value following flag, or "" if the flag is absent.
func argValue(args []string, flag string) string {
	i := slices.Index(args, flag)
	if i < 0 || i+1 >= len(args) {
		return ""
	}
	return args[i+1]
}

func singleNodeSpec() ShardSpec {
	return ShardSpec{
		ReplicaID:     "r1",
		ShardID:       "r1-0",
		Rank:          0,
		NNodes:        1,
		GPUIndices:    []int{0, 1, 2, 3},
		MgmtInterface: "eth0",
		HTTPPort:      20000,
		Profile: Profile{
			Name:      "h100-tp4",
			Engine:    "vllm",
			TP:        4,
			ModelPath: "/models/llama",
		},
	}
}

func multiNodeSpec(rank int) ShardSpec {
	s := singleNodeSpec()
	s.Rank = rank
	s.NNodes = 2
	s.GPUIndices = []int{0, 1, 2, 3, 4, 5, 6, 7}
	s.HCAs = []string{"mlx5_0", "mlx5_1"}
	s.LeaderAddr = "10.0.0.10"
	s.LeaderPort = 20001
	s.Profile.TP = 16
	s.Profile.GPUsPerNode = 8
	s.Profile.Nodes = 2
	return s
}

// --- environment generation: the part that breaks multi-node runs ---

func TestSingleNodeDisablesIB(t *testing.T) {
	cmd := buildOne(t, &VLLM{}, singleNodeSpec())
	env := envMap(cmd.Env)

	if env["NCCL_IB_DISABLE"] != "1" {
		t.Error("NCCL_IB_DISABLE must be 1 without a fabric, or NCCL waits for one that is not there")
	}
	if _, set := env["NCCL_IB_HCA"]; set {
		t.Error("NCCL_IB_HCA was set without any HCAs")
	}
	if env["CUDA_VISIBLE_DEVICES"] != "0,1,2,3" {
		t.Errorf("CUDA_VISIBLE_DEVICES=%q", env["CUDA_VISIBLE_DEVICES"])
	}
}

func TestMultiNodePinsFabric(t *testing.T) {
	cmd := buildEngine(t, &VLLM{}, multiNodeSpec(0))
	env := envMap(cmd.Env)

	// Only the chosen fabric's adapters. Leaving this open lets NCCL pick an
	// adapter on another fabric, where the peer is unreachable, and the job
	// hangs rather than failing.
	if env["NCCL_IB_HCA"] != "mlx5_0,mlx5_1" {
		t.Errorf("NCCL_IB_HCA=%q", env["NCCL_IB_HCA"])
	}
	if env["NCCL_IB_DISABLE"] != "0" {
		t.Errorf("NCCL_IB_DISABLE=%q with a fabric present", env["NCCL_IB_DISABLE"])
	}
	// Bootstrap must stay on the management network; NCCL otherwise probes
	// every interface, including IB-over-IP that cannot route between fabrics.
	if env["NCCL_SOCKET_IFNAME"] != "eth0" || env["GLOO_SOCKET_IFNAME"] != "eth0" {
		t.Errorf("socket interfaces not pinned: %q / %q",
			env["NCCL_SOCKET_IFNAME"], env["GLOO_SOCKET_IFNAME"])
	}
}

func TestProfileEnvOverridesGenerated(t *testing.T) {
	s := multiNodeSpec(0)
	s.Profile.Env = map[string]string{
		"NCCL_IB_HCA":        "mlx5_9",
		"NCCL_NET_GDR_LEVEL": "SYS",
	}
	cmd := buildEngine(t, &VLLM{}, s)
	env := envMap(cmd.Env)

	if env["NCCL_IB_HCA"] != "mlx5_9" {
		t.Error("a profile override did not win; it is the escape hatch for per-cluster tuning")
	}
	if env["NCCL_NET_GDR_LEVEL"] != "SYS" {
		t.Error("an extra profile variable was dropped")
	}
}

func TestProcessEnvironmentIsInherited(t *testing.T) {
	// Engines need PATH, HF_HOME, LD_LIBRARY_PATH and whatever else the unit
	// file set up.
	t.Setenv("LLMGW_TEST_MARKER", "present")
	cmd := buildOne(t, &VLLM{}, singleNodeSpec())
	if envMap(cmd.Env)["LLMGW_TEST_MARKER"] != "present" {
		t.Error("the process environment was not inherited")
	}
}

func TestEnvIsSorted(t *testing.T) {
	cmd := buildOne(t, &VLLM{}, singleNodeSpec())
	if !slices.IsSorted(cmd.Env) {
		t.Error("environment is not sorted; launches would not be reproducible in logs")
	}
}

// --- vLLM ---

func TestVLLMSingleNodeArgs(t *testing.T) {
	cmd := buildOne(t, &VLLM{cfg: VLLMConfig{Bin: "/opt/vllm/bin/vllm"}}, singleNodeSpec())
	if cmd.Path != "/opt/vllm/bin/vllm" {
		t.Errorf("Path=%q", cmd.Path)
	}
	if cmd.Args[0] != "serve" || cmd.Args[1] != "/models/llama" {
		t.Errorf("args start wrong: %v", cmd.Args[:2])
	}
	if argValue(cmd.Args, "--tensor-parallel-size") != "4" {
		t.Errorf("TP not passed: %v", cmd.Args)
	}
	if argValue(cmd.Args, "--port") != "20000" {
		t.Errorf("port not passed: %v", cmd.Args)
	}
	// Single node means no distributed backend flag at all.
	if slices.Contains(cmd.Args, "--distributed-executor-backend") {
		t.Error("a distributed backend was selected for a single-node profile")
	}
}

func TestVLLMRayLeaderRunsTwoProcesses(t *testing.T) {
	// The leader starts a Ray head and then the server against it. A single
	// command would mean running vLLM against a cluster that does not exist.
	v := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}}
	cmds, err := v.Commands(multiNodeSpec(0))
	if err != nil {
		t.Fatal(err)
	}
	if len(cmds) != 2 {
		t.Fatalf("expected 2 commands for the Ray leader, got %d", len(cmds))
	}
	if cmds[0].Name != "ray" || !slices.Contains(cmds[0].Args, "--head") {
		t.Errorf("first command is not the Ray head: %+v", cmds[0])
	}
	if !cmds[0].WaitForLeaderPort {
		t.Error("the server would start before the head is listening")
	}
	if cmds[1].Name != "vllm" || argValue(cmds[1].Args, "--distributed-executor-backend") != "ray" {
		t.Errorf("second command is not the vLLM server: %+v", cmds[1])
	}
	if got := envMap(cmds[1].Env)["RAY_ADDRESS"]; got != "10.0.0.10:20001" {
		t.Errorf("RAY_ADDRESS=%q", got)
	}
}

func TestVLLMRayWorkerRunsOnlyRay(t *testing.T) {
	// A Ray worker runs no vLLM process: the cluster hands its GPUs to the
	// leader. Starting a second server there would bind a port nobody uses and
	// compete for the same cards.
	v := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}}
	cmds, err := v.Commands(multiNodeSpec(1))
	if err != nil {
		t.Fatal(err)
	}
	if len(cmds) != 1 {
		t.Fatalf("expected 1 command for a Ray worker, got %d: %+v", len(cmds), cmds)
	}
	if cmds[0].Name != "ray" || slices.Contains(cmds[0].Args, "--head") {
		t.Errorf("worker command wrong: %+v", cmds[0])
	}
	if argValue(cmds[0].Args, "--address") != "10.0.0.10:20001" {
		t.Errorf("worker is not joining the leader: %v", cmds[0].Args)
	}
}

func TestVLLMMPMode(t *testing.T) {
	// vLLM's multi-node launch: every node the same serve with its rank and
	// the leader's address, workers headless. Workers used to be given the
	// leader's address as their own VLLM_HOST_IP, and no rank at all.
	v := &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}
	for _, rank := range []int{0, 1} {
		cmd := buildOne(t, v, multiNodeSpec(rank))
		for flag, want := range map[string]string{
			"--distributed-executor-backend": "mp", "--nnodes": "2",
			"--node-rank": strconv.Itoa(rank), "--master-addr": "10.0.0.10", "--master-port": "20001",
		} {
			if got := argValue(cmd.Args, flag); got != want {
				t.Errorf("rank %d: %s = %q, want %q (%v)", rank, flag, got, want, cmd.Args)
			}
		}
		if headless := slices.Contains(cmd.Args, "--headless"); headless != (rank > 0) {
			t.Errorf("rank %d: headless = %v", rank, headless)
		}
		if ip, ok := envMap(cmd.Env)["VLLM_HOST_IP"]; ok && ip == "10.0.0.10" && rank > 0 {
			t.Errorf("a worker was told the leader's address is its own")
		}
	}
}

func TestRoCEGetsItsGIDIndex(t *testing.T) {
	s := multiNodeSpec(0)
	idx := 3
	s.GIDIndex = &idx
	env := envMap(buildOne(t, &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}, s).Env)
	if env["NCCL_IB_GID_INDEX"] != "3" {
		t.Errorf("NCCL_IB_GID_INDEX = %q", env["NCCL_IB_GID_INDEX"])
	}
	// InfiniBand needs none, and must not be given one.
	env = envMap(buildOne(t, &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}, multiNodeSpec(0)).Env)
	if _, ok := env["NCCL_IB_GID_INDEX"]; ok {
		t.Error("an InfiniBand shard was given a GID index")
	}
}

func TestVLLMUnknownLaunchMode(t *testing.T) {
	v := &VLLM{cfg: VLLMConfig{LaunchMode: "sorcery"}}
	if _, err := v.Commands(multiNodeSpec(0)); err == nil {
		t.Error("an unknown launch mode was accepted")
	}
}

func TestOnlyLeaderServesHTTP(t *testing.T) {
	// The fix for the bug that stopped multi-node replicas coming up at all:
	// health-checking a worker's port times out and fails a healthy replica.
	for _, e := range []Engine{&VLLM{}, &SGLang{}} {
		if !e.ServesHTTP(multiNodeSpec(0)) {
			t.Errorf("%s: the leader must serve HTTP", e.Name())
		}
		if e.ServesHTTP(multiNodeSpec(1)) {
			t.Errorf("%s: a worker must not be expected to serve HTTP", e.Name())
		}
		if !e.ServesHTTP(singleNodeSpec()) {
			t.Errorf("%s: a single-node shard must serve HTTP", e.Name())
		}
	}
}

func TestNoRayForSingleNodeOrMPMode(t *testing.T) {
	ray := &VLLM{cfg: VLLMConfig{LaunchMode: "ray"}}
	if cmds := buildOne(t, ray, singleNodeSpec()); cmds.Name != "vllm" {
		t.Errorf("a single-node profile was given a Ray process: %+v", cmds)
	}
	mp := &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}
	if cmds := buildOne(t, mp, multiNodeSpec(0)); cmds.Name != "vllm" {
		t.Errorf("mp mode was given a Ray process: %+v", cmds)
	}
}

// --- SGLang ---

func TestSGLangMultiNodeArgs(t *testing.T) {
	cmd := buildOne(t, &SGLang{cfg: SGLangConfig{Python: "/opt/sglang/bin/python"}}, multiNodeSpec(1))
	if cmd.Path != "/opt/sglang/bin/python" {
		t.Errorf("Path=%q", cmd.Path)
	}
	if cmd.Args[0] != "-c" || !strings.Contains(cmd.Args[1], `run_module("sglang.launch_server"`) {
		t.Errorf("not launched as the module: %v", cmd.Args[:2])
	}
	if argValue(cmd.Args, "--nnodes") != "2" || argValue(cmd.Args, "--node-rank") != "1" {
		t.Errorf("multi-node flags wrong: %v", cmd.Args)
	}
	if argValue(cmd.Args, "--dist-init-addr") != "10.0.0.10:20001" {
		t.Errorf("dist-init-addr wrong: %v", cmd.Args)
	}
}

func TestSGLangsKeyIsNotOnItsCommandLine(t *testing.T) {
	s := singleNodeSpec()
	s.APIKey = "k1"
	cmd := buildOne(t, &SGLang{}, s)
	if slices.Contains(cmd.Args, "k1") || slices.Contains(cmd.Args, "--api-key") {
		t.Errorf("the key is in the arguments: %v", cmd.Args)
	}
	if !slices.Contains(cmd.Env, sglangKeyEnv+"=k1") {
		t.Errorf("the key is not in the environment")
	}
	// A worker serves no HTTP, and is not given it.
	w := multiNodeSpec(1)
	w.APIKey = "k1"
	if cmd := buildOne(t, &SGLang{}, w); slices.Contains(cmd.Env, sglangKeyEnv+"=k1") {
		t.Error("a worker was given the key")
	}
}

func TestEnginesListenOnTheAddressTheyAreGiven(t *testing.T) {
	s := singleNodeSpec()
	for _, e := range []Engine{&VLLM{}, &SGLang{}} {
		s.BindAddr = ""
		if got := argValue(buildOne(t, e, s).Args, "--host"); got != "0.0.0.0" {
			t.Errorf("%s: without an address listens on %q", e.Name(), got)
		}
		s.BindAddr = "10.0.0.147"
		if got := argValue(buildOne(t, e, s).Args, "--host"); got != "10.0.0.147" {
			t.Errorf("%s: listens on %q, expected the management address", e.Name(), got)
		}
	}
}

func TestSGLangIsStartedServingItsMetrics(t *testing.T) {
	// Without --enable-metrics SGLang has no /metrics, and every scrape of
	// the replica fails as though it were unreachable.
	if cmd := buildOne(t, &SGLang{}, singleNodeSpec()); !slices.Contains(cmd.Args, "--enable-metrics") {
		t.Errorf("not started with its metrics: %v", cmd.Args)
	}
}

func TestSGLangIsAskedForItsHelpThroughItsModule(t *testing.T) {
	g := &SGLang{cfg: SGLangConfig{Python: "/opt/sglang/bin/python"}}
	g.UseInstalled("/opt/llmgw/engines/sglang/0.5.21-cu13/bin/python")
	path, args := g.HelpCommand(Command{})
	if path != "/opt/llmgw/engines/sglang/0.5.21-cu13/bin/python" {
		t.Errorf("asked %s, not the interpreter it was installed into", path)
	}
	if !slices.Equal(args, []string{"-m", "sglang.launch_server", "--help"}) {
		t.Errorf("asked %v", args)
	}
}

func TestSGLangSingleNodeOmitsMultiNodeFlags(t *testing.T) {
	cmd := buildOne(t, &SGLang{}, singleNodeSpec())
	for _, flag := range []string{"--nnodes", "--node-rank", "--dist-init-addr"} {
		if slices.Contains(cmd.Args, flag) {
			t.Errorf("%s was passed for a single-node profile", flag)
		}
	}
}

func TestProfileArgsComeLast(t *testing.T) {
	// Operators must be able to override anything the adapter generated.
	s := singleNodeSpec()
	s.Profile.Args = []string{"--max-model-len", "32768"}
	cmd := buildEngine(t, &VLLM{}, s)

	i := slices.Index(cmd.Args, "--max-model-len")
	j := slices.Index(cmd.Args, "--tensor-parallel-size")
	if i < j {
		t.Errorf("profile args came before generated ones: %v", cmd.Args)
	}
}

// --- validation ---

func h100(n int) []GPUInfo {
	out := make([]GPUInfo, n)
	for i := range out {
		out[i] = GPUInfo{Index: i, Name: "NVIDIA H100 80GB HBM3", ComputeCapability: "9.0", MemoryTotalMB: 81559}
	}
	return out
}

func TestValidateAcceptsMatchingProfile(t *testing.T) {
	p := Profile{Name: "ok", TP: 4, ModelPath: "/m"}
	if err := (&VLLM{}).Validate(p, h100(4)); err != nil {
		t.Errorf("a valid profile was rejected: %v", err)
	}
}

func TestValidateRejectsMixedGPUTypes(t *testing.T) {
	// The invariant from plan §1.1: a TP group never mixes GPU types.
	gpus := h100(2)
	gpus[1].Name = "NVIDIA A100-SXM4-40GB"

	err := (&VLLM{}).Validate(Profile{Name: "mixed", TP: 2, ModelPath: "/m"}, gpus)
	if err == nil || !strings.Contains(err.Error(), "may not mix GPU types") {
		t.Errorf("mixed GPU types accepted: %v", err)
	}
}

func TestValidateRejectsFP8OnOldHardware(t *testing.T) {
	gpus := h100(4)
	for i := range gpus {
		gpus[i].Name = "NVIDIA A100-SXM4-80GB"
		gpus[i].ComputeCapability = "8.0"
	}
	p := Profile{Name: "fp8", TP: 4, ModelPath: "/m", Args: []string{"--quantization", "fp8"}}

	err := (&VLLM{}).Validate(p, gpus)
	if err == nil || !strings.Contains(err.Error(), "compute capability") {
		t.Errorf("fp8 accepted on compute capability 8.0: %v", err)
	}

	// The same profile is fine on Hopper.
	if err := (&VLLM{}).Validate(p, h100(4)); err != nil {
		t.Errorf("fp8 rejected on H100: %v", err)
	}
}

func TestValidateRejectsInconsistentParallelism(t *testing.T) {
	// 2 nodes x 8 GPUs cannot serve TP 4.
	p := Profile{Name: "bad", TP: 4, Nodes: 2, GPUsPerNode: 8, ModelPath: "/m"}
	if err := (&VLLM{}).Validate(p, h100(8)); err == nil {
		t.Error("inconsistent node/GPU/TP arithmetic accepted")
	}
}

func TestValidateRejectsTooFewGPUs(t *testing.T) {
	if err := (&VLLM{}).Validate(Profile{Name: "big", TP: 8, ModelPath: "/m"}, h100(4)); err == nil {
		t.Error("a TP 8 profile was accepted on a 4-GPU node")
	}
}

func TestValidateRequiresModelPath(t *testing.T) {
	if err := (&VLLM{}).Validate(Profile{Name: "nopath", TP: 1}, h100(1)); err == nil {
		t.Error("a profile without a model path was accepted")
	}
}

func TestRegistry(t *testing.T) {
	r := NewRegistry(Config{})
	if got := r.Names(); !slices.Equal(got, []string{"comfyui", "llamacpp", "sglang", "sim", "trtllm", "vllm"}) {
		t.Errorf("registry contents: %v", got)
	}
	if _, err := r.Get("vllm"); err != nil {
		t.Errorf("vllm missing: %v", err)
	}
	err := func() error { _, err := r.Get("nope"); return err }()
	if err == nil || !strings.Contains(err.Error(), "llamacpp, sglang, sim, trtllm, vllm") {
		t.Errorf("an unknown engine should list what is available: %v", err)
	}
}

func TestQuantizationOf(t *testing.T) {
	for _, args := range [][]string{
		{"--quantization", "fp8"},
		{"--quantization=fp8"},
		{"--max-model-len", "100", "--quantization", "fp8"},
	} {
		if got := quantizationOf(args); got != "fp8" {
			t.Errorf("quantizationOf(%v)=%q", args, got)
		}
	}
	if got := quantizationOf([]string{"--quantization"}); got != "" {
		t.Errorf("a dangling flag should yield nothing, got %q", got)
	}
}

func TestAReplicaInAnNVLinkDomainIsToldToUseIt(t *testing.T) {
	s := multiNodeSpec(0)
	s.Fabric = "nvl:76b1ad76-7c24-4b6b-9f6c-2a1e0d8c11aa/7"
	env := envMap(buildOne(t, &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}, s).Env)
	if env["NCCL_MNNVL_ENABLE"] != "1" {
		t.Errorf("NCCL_MNNVL_ENABLE = %q", env["NCCL_MNNVL_ENABLE"])
	}
	if env := envMap(buildOne(t, &VLLM{cfg: VLLMConfig{LaunchMode: "mp"}}, multiNodeSpec(0)).Env); env["NCCL_MNNVL_ENABLE"] != "" {
		t.Error("a replica on InfiniBand was told to use multi-node NVLink")
	}
}

func TestMultiNodeVLLMDefaultsToItsOwnLaunch(t *testing.T) {
	// mp is what vLLM documents and needs nothing but vLLM; Ray was the
	// default and was never checked against a real cluster.
	cmds, err := (&VLLM{}).Commands(multiNodeSpec(1))
	if err != nil {
		t.Fatal(err)
	}
	if len(cmds) != 1 || cmds[0].Name != "vllm" || argValue(cmds[0].Args, "--distributed-executor-backend") != "mp" ||
		!slices.Contains(cmds[0].Args, "--headless") {
		t.Errorf("default worker = %+v", cmds)
	}
}