llmgw/internal/agent/engines/strata_test.go

package engines

import (
	"os"
	"path/filepath"
	"slices"
	"strings"
	"testing"
)

func TestStrataStartsStrataServeWithTheKeyInTheEnvironment(t *testing.T) {
	s := &Strata{}
	s.UseInstalled("/opt/llmgw/engines/strata/0.2.7-cuda13/bin/python")
	spec := ShardSpec{
		ReplicaID: "r1", HTTPPort: 18090, BindAddr: "10.0.0.147", GPUIndices: []int{0},
		APIKey: "secret-key",
		Profile: Profile{Name: "p", Engine: "strata", TP: 1, ModelPath: "/var/lib/llmgw/models/q38",
			ServedModelName: "qwen3.8-flash-next", Args: []string{"--family", "qwen", "--size", "IQ2_XS"}},
	}
	cmds, err := s.Commands(spec)
	if err != nil {
		t.Fatal(err)
	}
	c := cmds[0]
	args := strings.Join(c.Args, " ")
	env := "/opt/llmgw/engines/strata/0.2.7-cuda13"
	for _, want := range []string{
		"strata-serve", "--listen 10.0.0.147:18090", "--model qwen3.8-flash-next",
		"--weights /var/lib/llmgw/models/q38", "--state " + filepath.Join("strata", "r1"),
		"--python " + env + "/bin/python", "--strata " + env + "/strata", "--engine " + env + "/strata-engine",
		"-- --family qwen --size IQ2_XS",
	} {
		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")
	}
	if !slices.Contains(c.Env, "STRATA_API_KEY=secret-key") {
		t.Error("the key is not in the environment, where Strata's server reads it")
	}
	if !slices.Contains(c.Env, "CUDA_VISIBLE_DEVICES=0") {
		t.Error("the shard is not pinned to its GPU")
	}
}

func TestStrataRunsOnOneCard(t *testing.T) {
	s := &Strata{}
	p := Profile{Name: "p", Engine: "strata", TP: 1, GPUsPerNode: 2, ModelPath: "/m"}
	if err := s.Validate(p, nil); err == nil {
		t.Error("a two-card profile was accepted")
	}
	p.GPUsPerNode = 1
	if err := s.Validate(p, []GPUInfo{{Name: "Quadro RTX 6000", ComputeCapability: "7.5"}}); err != nil {
		t.Errorf("one Turing card was refused: %v", err)
	}
}

func TestStrataSetupFlagsThatMoveFilesAreRefused(t *testing.T) {
	for _, bad := range []string{"--gguf-dir", "--gguf", "--data-dir=/etc", "--model", "--mod", "--feed", "--host", "--gpus"} {
		if err := UnsafeArgs("strata", []string{bad, "x"}); err == nil {
			t.Errorf("%s was accepted", bad)
		}
	}
	if err := UnsafeArgs("strata", []string{"--family", "qwen", "--size", "IQ2_XS", "--context", "32768",
		"--parallel", "2", "--vram-reserve-mib", "2048", "--vision", "no"}); err != nil {
		t.Errorf("an ordinary profile was refused: %v", err)
	}
}

func TestAStrataNamedByHandIsReportedWithItsVersion(t *testing.T) {
	dir := t.TempDir()
	for _, f := range []string{"setup.py", "serve/server.py", ".venv/bin/python", "engine/strata"} {
		p := filepath.Join(dir, f)
		if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(p, nil, 0o755); err != nil {
			t.Fatal(err)
		}
	}
	s := &Strata{cfg: StrataConfig{Dir: dir}}
	if v, err := s.Probe(t.Context()); err != nil || v != Unknown {
		t.Fatalf("%q %v", v, err)
	}
	_ = os.MkdirAll(filepath.Join(dir, "strata_cli"), 0o755)
	_ = os.WriteFile(filepath.Join(dir, "strata_cli", "__init__.py"), []byte("PRODUCT_VERSION = \"0.2.7\"\n"), 0o644)
	if v, _ := s.Probe(t.Context()); v != "0.2.7" {
		t.Errorf("version %q", v)
	}
	r := NewRegistry(Config{Strata: StrataConfig{Dir: dir}})
	if got := r.ProbeConfigured(t.Context()); got["strata"] != "0.2.7" {
		t.Errorf("not reported: %v", got)
	}
	_ = os.Remove(filepath.Join(dir, "engine", "strata"))
	if _, err := s.Probe(t.Context()); err == nil {
		t.Error("reported without its engine")
	}
}

func TestStrataRefusesDeepSeekOnACardItsEngineWasNotBuiltFor(t *testing.T) {
	s := &Strata{}
	p := Profile{Name: "p", Engine: "strata", ModelPath: "/m", TP: 1, GPUsPerNode: 1, Args: []string{"--family", "deepseek41", "--size", "IST-DAS-1.75bpw"}}
	gpus := []GPUInfo{{Index: 0, Name: "Quadro RTX 6000", ComputeCapability: "7.5"}}
	if err := s.Validate(p, gpus); err == nil || !strings.Contains(err.Error(), "8.6") {
		t.Errorf("DeepSeek on 7.5: %v", err)
	}
	p.Args = []string{"--family", "qwen", "--size", "IQ2_XS"}
	if err := s.Validate(p, gpus); err != nil {
		t.Errorf("Qwen on 7.5: %v", err)
	}
}