llmgw/internal/agent/engines/freetoken_test.go

package engines

import (
	"slices"
	"strings"
	"testing"
)

func withToolkit(t *testing.T, dir string) {
	t.Helper()
	old := freetokenToolkit
	freetokenToolkit = func() string { return dir }
	t.Cleanup(func() { freetokenToolkit = old })
}

func TestFreeTokenServesOnTheShardsCardWithTheKeyInTheEnvironment(t *testing.T) {
	withToolkit(t, "/usr/local/cuda-13.4")
	f := &FreeToken{}
	f.UseInstalled("/opt/llmgw/engines/freetoken/0.2.1-cu13/bin/ft")
	spec := ShardSpec{
		ReplicaID: "r1", HTTPPort: 18091, BindAddr: "10.0.0.160", GPUIndices: []int{3},
		APIKey: "secret-key",
		Profile: Profile{Name: "p", Engine: "freetoken", TP: 1, ModelPath: "/var/lib/llmgw-weights/glm-5.2-nvfp4",
			ServedModelName: "glm-5.2", Args: []string{"--max-running-requests", "8"}},
	}
	cmds, err := f.Commands(spec)
	if err != nil {
		t.Fatal(err)
	}
	c := cmds[0]
	if c.Path != "/opt/llmgw/engines/freetoken/0.2.1-cu13/bin/ft" {
		t.Errorf("path %s", c.Path)
	}
	args := strings.Join(c.Args, " ")
	for _, want := range []string{
		"serve --model /var/lib/llmgw-weights/glm-5.2-nvfp4", "--host 10.0.0.160", "--port 18091",
		// The shard's card is device 0 once CUDA_VISIBLE_DEVICES pins it.
		"--gpu 0", "--served-model-name glm-5.2", "--max-running-requests 8",
	} {
		if !strings.Contains(args, want) {
			t.Errorf("args lack %q: %s", want, args)
		}
	}
	if strings.Contains(args, "secret-key") {
		t.Error("the key is on the command line, which every user on the node can read")
	}
	for _, want := range []string{"FREETOKEN_API_KEY=secret-key", "CUDA_VISIBLE_DEVICES=3", "CUDA_HOME=/usr/local/cuda-13.4"} {
		if !slices.Contains(c.Env, want) {
			t.Errorf("env lacks %s", want)
		}
	}
	var path string
	for _, e := range c.Env {
		if p, ok := strings.CutPrefix(e, "PATH="); ok {
			path = p
		}
	}
	if !strings.HasPrefix(path, "/opt/llmgw/engines/freetoken/0.2.1-cu13/bin:/usr/local/cuda-13.4/bin:") {
		t.Errorf("nvcc and the environment's tools are not first on PATH: %s", path)
	}
}

func TestFreeTokenWithoutTheToolkitSaysSoInsteadOfFailingOnItsFirstModel(t *testing.T) {
	withToolkit(t, "")
	f := &FreeToken{cfg: FreeTokenConfig{Binary: "/x/bin/ft"}}
	_, err := f.Commands(ShardSpec{HTTPPort: 1, Profile: Profile{Name: "p", ModelPath: "/m"}})
	if err == nil || !strings.Contains(err.Error(), "CUDA 13 toolkit") {
		t.Errorf("got %v", err)
	}
}

func TestFreeTokenRunsOnOneAmpereOrNewerCard(t *testing.T) {
	f := &FreeToken{}
	p := Profile{Name: "p", Engine: "freetoken", TP: 1, ModelPath: "/m"}
	if err := f.Validate(p, []GPUInfo{{Name: "NVIDIA RTX A6000", ComputeCapability: "8.6", MemoryTotalMB: 49140}}); err != nil {
		t.Errorf("an Ampere card was refused: %v", err)
	}
	// Tried on cuda01: no kernel image for the card, then an illegal memory
	// access in fp16.
	if err := f.Validate(p, []GPUInfo{{Name: "Quadro RTX 6000", ComputeCapability: "7.5", MemoryTotalMB: 24576}}); err == nil {
		t.Error("a Turing card was accepted")
	}
	p.TP = 2
	if err := f.Validate(p, h100(2)); err == nil {
		t.Error("a two-card profile was accepted; FreeToken has no tensor parallel")
	}
}

func TestFreeTokenFlagsThatReadFilesOrMoveWhatTheAgentSetsAreRefused(t *testing.T) {
	for _, a := range []string{
		"--gpu", "--served-model-name=x", "--mmproj", "--allowed-local-media-path",
		// argparse takes a beginning for the whole flag.
		"--allowed-local", "--api-key", "--port", "--hf-overrides",
	} {
		if UnsafeArgs("freetoken", []string{a, "v"}) == nil {
			t.Errorf("%s was allowed", a)
		}
	}
	for _, a := range []string{"--max-running-requests", "--moe-strategy", "--dtype", "--max-seq-len-override", "--moe-cache-auto"} {
		if err := UnsafeArgs("freetoken", []string{a, "1"}); err != nil {
			t.Errorf("%s was refused: %v", a, err)
		}
	}
	for k, refuse := range map[string]bool{
		"FREETOKEN_PIN_BUDGET_GB": false, "FREETOKEN_API_KEY": true, "FREETOKEN_API_LOG_DIR": true,
		"FREETOKEN_MMPROJ": true, "FREETOKEN_HOST": true,
	} {
		if got, why := UnsafeEnv(k); got != refuse {
			t.Errorf("%s: refused %v (%s), want %v", k, got, why, refuse)
		}
	}
}