diff --git a/docs/reference.md b/docs/reference.md index c12ffbc..b75371c 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -396,6 +396,9 @@ rather than repeated. | `timeout` | `time.Duration` | [Service](#service) | | `stopped` | `bool` | [ServingSettings](#servingsettings) | | `bin` | `string` | [SimConfig](#simconfig) | +| `dir` | `string` | [StrataConfig](#strataconfig) | +| `python` | `string` | [StrataConfig](#strataconfig) | +| `engine` | `string` | [StrataConfig](#strataconfig) | | `python` | `string` | [TRTLLMConfig](#trtllmconfig) | | `store` | `string` | [Threads](#threads) | | `retain_hours` | `int` | [Threads](#threads) | @@ -476,6 +479,7 @@ rather than repeated. | `trtllm` | `TRTLLMConfig` | [engines.Config](#enginesconfig) | | `sim` | `SimConfig` | [engines.Config](#enginesconfig) | | `comfyui` | `ComfyUIConfig` | [engines.Config](#enginesconfig) | +| `strata` | `StrataConfig` | [engines.Config](#enginesconfig) | | `enabled` | `*bool` | [guard.Config](#guardconfig) | | `default` | `[]string` | [guard.Config](#guardconfig) | | `policies` | `map[string]Policy` | [guard.Config](#guardconfig) | @@ -3415,6 +3419,27 @@ entries ask. In force on every control plane within seconds. Bin is the llmgw binary itself; the engine is one of its subcommands. Empty means whatever started this agent, which is right almost always. +## StrataConfig + +### `dir` + +`string` + +Dir is Strata's source directory, where setup.py and serve/ are. + +### `python` + +`string` + +Python is the interpreter of its environment; Dir's .venv when unset. + +### `engine` + +`string` + +Engine is the directory of its prebuilt engine, where the strata +binary and the CUDA libraries it ships are; Dir's engine/ when unset. + ## TRTLLMConfig ### `python` @@ -3987,6 +4012,15 @@ ComfyUI makes pictures; offered by a node whose ComfyUI is named here. Takes a [ComfyUIConfig](#comfyuiconfig) block. +### `strata` + +`StrataConfig` + +Strata serves a few mixture-of-experts models from RAM and one card; +offered by a node whose Strata is named here, or that installed it. + +Takes a [StrataConfig](#strataconfig) block. + ## guard.Config ### `enabled` diff --git a/internal/agent/engines/engine.go b/internal/agent/engines/engine.go index cf08612..160e96c 100644 --- a/internal/agent/engines/engine.go +++ b/internal/agent/engines/engine.go @@ -218,6 +218,7 @@ func NewRegistry(cfg Config) *Registry { r.Add(&TRTLLM{cfg: cfg.TRTLLM}) r.Add(&Sim{cfg: cfg.Sim}) r.Add(&ComfyUI{cfg: cfg.ComfyUI}) + r.Add(&Strata{cfg: cfg.Strata}) return r } @@ -268,6 +269,9 @@ type Config struct { Sim SimConfig `koanf:"sim"` // ComfyUI makes pictures; offered by a node whose ComfyUI is named here. ComfyUI ComfyUIConfig `koanf:"comfyui"` + // Strata serves a few mixture-of-experts models from RAM and one card; + // offered by a node whose Strata is named here, or that installed it. + Strata StrataConfig `koanf:"strata"` } // --- shared helpers --- diff --git a/internal/agent/engines/engines_test.go b/internal/agent/engines/engines_test.go index 2f58000..b0548c8 100644 --- a/internal/agent/engines/engines_test.go +++ b/internal/agent/engines/engines_test.go @@ -477,14 +477,14 @@ func TestValidateRequiresModelPath(t *testing.T) { func TestRegistry(t *testing.T) { r := NewRegistry(Config{}) - if got := r.Names(); !slices.Equal(got, []string{"comfyui", "llamacpp", "sglang", "sim", "trtllm", "vllm"}) { + if got := r.Names(); !slices.Equal(got, []string{"comfyui", "llamacpp", "sglang", "sim", "strata", "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") { + if err == nil || !strings.Contains(err.Error(), "llamacpp, sglang, sim, strata, trtllm, vllm") { t.Errorf("an unknown engine should list what is available: %v", err) } } diff --git a/internal/agent/engines/strata.go b/internal/agent/engines/strata.go new file mode 100644 index 0000000..ff98817 --- /dev/null +++ b/internal/agent/engines/strata.go @@ -0,0 +1,144 @@ +package engines + +import ( + "fmt" + "net" + "os" + "path/filepath" + "strconv" + "sync/atomic" +) + +// StrataConfig locates a Strata installation. +type StrataConfig struct { + // Dir is Strata's source directory, where setup.py and serve/ are. + Dir string `koanf:"dir"` + // Python is the interpreter of its environment; Dir's .venv when unset. + Python string `koanf:"python"` + // Engine is the directory of its prebuilt engine, where the strata + // binary and the CUDA libraries it ships are; Dir's engine/ when unset. + Engine string `koanf:"engine"` +} + +// Strata runs Strata (github.com/Niko1221/Strata): a mixture-of-experts +// engine that keeps a model's experts in RAM and the ones used most in VRAM, +// so a 125B model serves from one 24 GB card. It runs a few model families +// only, each prepared from its GGUF files by Strata's own setup, so the agent +// starts llmgw strata-serve, which prepares the model and then becomes +// Strata's server. +type Strata struct { + cfg StrataConfig + installed atomic.Pointer[string] +} + +// Name identifies the engine in the catalog. +func (*Strata) Name() string { return "strata" } + +// HealthPath answers 200 only once the model is loaded. +func (*Strata) HealthPath() string { return "/health" } + +// MetricsPath is Strata's Prometheus text, in vLLM's names; its /metrics is +// JSON for its own monitor. +func (*Strata) MetricsPath() string { return "/metrics/prometheus" } + +// ServesHTTP is true: a Strata replica is one process on one node. +func (*Strata) ServesHTTP(ShardSpec) bool { return true } + +// UseInstalled points the engine at the environment the agent installed: +// its interpreter, with Strata's code beside its bin in strata/ and the +// engine in strata-engine/. +func (s *Strata) UseInstalled(python string) { s.installed.Store(&python) } + +// Configured reports whether the configuration names a Strata. +func (s *Strata) Configured() bool { return s.cfg.Dir != "" } + +// StrataSourceDir and StrataEngineDir are where an installed environment +// keeps Strata's code and its prebuilt engine, beside bin/. +const ( + StrataSourceDir = "strata" + StrataEngineDir = "strata-engine" +) + +func (s *Strata) env() string { + if p := s.installed.Load(); p != nil { + return filepath.Dir(filepath.Dir(*p)) + } + return "" +} + +func (s *Strata) dir() string { + if e := s.env(); e != "" { + return filepath.Join(e, StrataSourceDir) + } + return s.cfg.Dir +} + +func (s *Strata) engineDir() string { + if e := s.env(); e != "" { + return filepath.Join(e, StrataEngineDir) + } + if s.cfg.Engine != "" { + return s.cfg.Engine + } + return filepath.Join(s.cfg.Dir, "engine") +} + +func (s *Strata) python() string { + if p := s.installed.Load(); p != nil { + return *p + } + if s.cfg.Python != "" { + return s.cfg.Python + } + return filepath.Join(s.cfg.Dir, ".venv", "bin", "python") +} + +// Validate checks the profile: one card, on one node. Strata splits a model +// over a node's cards by layer, but that is not driven from here yet. +func (*Strata) Validate(p Profile, gpus []GPUInfo) error { + if p.Nodes > 1 || p.TP > 1 || p.PP > 1 || p.GPUsPerNode > 1 { + return fmt.Errorf("profile %q: strata runs on one GPU of one node", p.Name) + } + if err := requireComputeCapability(p, gpus, 7.0, "strata"); err != nil { + return err + } + return validateCommon(p, gpus) +} + +// Commands is llmgw strata-serve, which prepares the model with Strata's +// setup and then execs Strata's server. +func (s *Strata) Commands(sh ShardSpec) ([]Command, error) { + dir := s.dir() + if dir == "" { + return nil, fmt.Errorf("strata: no Strata directory is configured (engines.strata.dir)") + } + bin, err := os.Executable() + if err != nil { + return nil, err + } + p := sh.Profile + args := []string{"strata-serve", + "--listen", net.JoinHostPort(sh.Host(), strconv.Itoa(sh.HTTPPort)), + "--model", p.ServedModelName, + "--weights", p.ModelPath, + // Relative to where engines run — their user's home — so they are + // theirs to write. The data is shared by the node's replicas: a + // prepared model and its draft layer are gigabytes, the same for + // every replica of it. + "--state", filepath.Join("strata", sh.ReplicaID), + "--data", filepath.Join("strata", "data"), + "--python", s.python(), + "--strata", dir, + "--engine", s.engineDir(), + "--", + } + args = append(args, p.Args...) + env := baseEnv(sh) + if sh.APIKey != "" { + // Strata's server reads its key from here, and strata-serve passes + // it on in the environment: not on a command line every user on the + // node can read. + env["STRATA_API_KEY"] = sh.APIKey + } + return []Command{{Name: "strata", Path: bin, Args: args, Env: materializeEnv(env)}}, nil +} diff --git a/internal/agent/engines/strata_test.go b/internal/agent/engines/strata_test.go new file mode 100644 index 0000000..788685f --- /dev/null +++ b/internal/agent/engines/strata_test.go @@ -0,0 +1,69 @@ +package engines + +import ( + "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) + } +} diff --git a/internal/agent/engines/unsafeenv.go b/internal/agent/engines/unsafeenv.go index 157d1d1..cab4e62 100644 --- a/internal/agent/engines/unsafeenv.go +++ b/internal/agent/engines/unsafeenv.go @@ -141,6 +141,22 @@ var flagsNever = map[string]map[string]string{ "--tls-keyfile": "reads a file", "--tls-certfile": "reads a file", "--comfy-api-base": "sends work to another service", }, + // Strata's setup.py, which a profile's arguments go to: llmgw + // strata-serve sets where the model is, where its files go and where + // the server listens, and the engine is the one the agent installed. + // Its own --model, the size, is spelt --size in a profile, since + // --model is the source everywhere else. + "strata": { + "--gguf-dir": "is the model's source, set by the agent", "--models-dir": "is set by the agent", + "--data-dir": "is set by the agent", "--move-data": "moves files", "--set-default-data-dir": "writes settings", + "--feed": "downloads an engine", "--prebuilt": "loads an engine", "--build": "compiles an engine", + "--update": "downloads an engine", "--rollback-engine": "changes the engine", + "--dsv41-package": "loads a file", "--inspect": "reads another repository or path", + "--source": "downloads from elsewhere", "--setup": "is not for a server", + "--browser": "opens a browser", "--gpu": "is set by the agent", "--gpus": "is set by the agent", + "--gpu-roles": "is set by the agent", "--backend": "is set by the agent", + "--experimental-speed-projection": "changes what the model refuses", + }, // trtllm-serve spells its flags with underscores, so the shared list's // hyphenated names do not reach them. "trtllm": { @@ -207,7 +223,7 @@ func unsafeFlag(engine, flag, _ string) string { // argparseEngines read their flags with Python's argparse, which takes a // long flag's beginning for the flag. -var argparseEngines = map[string]bool{"vllm": true, "sglang": true, "comfyui": true} +var argparseEngines = map[string]bool{"vllm": true, "sglang": true, "comfyui": true, "strata": true} // UnsafeArgs refuses the first of a profile's arguments an engine may not be // given, saying why. diff --git a/internal/agent/provision/install.go b/internal/agent/provision/install.go index a1feed9..d34a155 100644 --- a/internal/agent/provision/install.go +++ b/internal/agent/provision/install.go @@ -174,7 +174,7 @@ func (in *Installer) one(ctx context.Context, w Want) Status { } return st } - if w.Offline != nil && recipe.Repo == "" { + if w.Offline != nil && (recipe.Repo == "" || builds[0].Python != "") { // Said, not attempted: uv would spend minutes failing to reach PyPI, // and the failure would read like a network fault rather than the // design of an offline site. diff --git a/internal/agent/provision/python.go b/internal/agent/provision/python.go index 0477f9d..d756462 100644 --- a/internal/agent/provision/python.go +++ b/internal/agent/provision/python.go @@ -109,6 +109,11 @@ func (in *Installer) installPython(ctx context.Context, r Recipe, w Want, b Buil } reqs = append(reqs, more...) } + if len(b.Assets) > 0 { + if err := in.unpackAssets(ctx, r, w, b, tmp, st); err != nil { + return err + } + } if b.Wheel != nil { // The engine from its GitHub release, pinned to the digest GitHub // publishes for the file: the same trust as an archive, and uv @@ -372,6 +377,48 @@ func (in *Installer) unpackSource(ctx context.Context, w Want, b Build, env stri return requirementsIn(string(f), b.IndexPackages) } +// unpackAssets fetches a Python engine's release binaries, each checked +// against the digest GitHub publishes for it, and unpacks them into the +// environment's AssetDir. +func (in *Installer) unpackAssets(ctx context.Context, r Recipe, w Want, b Build, env string, st *Status) error { + tag := b.AssetTag + if tag == "" { + tag = w.Version + } + in.progress(st, "asking GitHub for the checksums of %s %s", r.Repo, tag) + sums, err := in.checksums(ctx, r.Repo, tag) + if err != nil { + return err + } + dst := filepath.Join(env, b.AssetDir) + if err := os.MkdirAll(dst, 0o755); err != nil { + return err + } + for _, asset := range b.Assets { + want, ok := sums[asset] + if !ok { + // Refused rather than installed unchecked: this is a binary that + // will run on every node. + return fmt.Errorf("%s %s publishes no checksum for %s, so it cannot be verified", r.Repo, tag, asset) + } + url := strings.TrimSuffix(in.downloadBase(), "/") + "/" + r.Repo + "/releases/download/" + tag + "/" + asset + archive := filepath.Join(env+".source", asset) + if err := os.MkdirAll(filepath.Dir(archive), 0o755); err != nil { + return err + } + if err := in.download(ctx, url, asset, archive, want, st); err != nil { + return err + } + in.progress(st, "unpacking %s", asset) + err := untarGz(archive, dst) + _ = os.RemoveAll(filepath.Dir(archive)) + if err != nil { + return fmt.Errorf("unpacking %s: %w", asset, err) + } + } + return nil +} + // requirementsIn reads a requirements file's lines. Options (-r, -e, --index-url // and the like) are refused: they would fetch from somewhere other than PyPI, // or install what the lock never saw. diff --git a/internal/agent/provision/recipe.go b/internal/agent/provision/recipe.go index 85d6495..3318138 100644 --- a/internal/agent/provision/recipe.go +++ b/internal/agent/provision/recipe.go @@ -76,6 +76,11 @@ type Build struct { // was published, which PyPI cannot be asked for. Source *Source AsOf time.Time + // AssetTag and AssetDir, for a Python engine that also ships a binary + // in its releases, are the release tag Assets are taken from, checked + // against the digests GitHub publishes, and the directory in the + // environment they are unpacked into. + AssetTag, AssetDir string // FreeGB is the disk space the installation needs before it starts: a // Python engine is several gigabytes, and one that fills the disk halfway // leaves the node worse off than one that never began. @@ -125,6 +130,7 @@ var Recipes = map[string]Recipe{ "vllm": {PyPI: "vllm", Candidates: vllmCandidates}, "sglang": {PyPI: "sglang", Candidates: sglangCandidates}, "comfyui": {Candidates: comfyuiCandidates}, + "strata": {Repo: StrataRepo, Candidates: strataCandidates}, "uv": {Repo: "astral-sh/uv", Candidates: uvCandidates}, } @@ -172,6 +178,9 @@ func Builds(engine, version, arch string) ([]Build, error) { continue } for _, b := range bs { + if b.Python != "" { + return nil, fmt.Errorf("%s is a Python environment from PyPI, which an offline bundle does not carry", engine) + } if !seen[b.Variant] { seen[b.Variant] = true out = append(out, b) @@ -672,3 +681,80 @@ print("comfyui version", "` + version + `", "on torch", torch.__version__)`}, {cu126, gpuNeeds{minDriverForCUDA12, 5.0}, "ComfyUI " + version + " on CUDA 12.6's torch"}, }, h) } + +// --- Strata --- + +// StrataRepo is where Strata is published: its source, and in each release +// the prebuilt engine for each CUDA. +const StrataRepo = "Niko1221/Strata" + +// strataReleases are the Strata versions this agent installs: the source +// archive's digest, taken when it was added here and confirmed from a second +// machine, when the version was published, and the engine its release +// carries, whose files GitHub publishes digests for. +var strataReleases = map[string]struct { + sha256 string + published string + engine string +}{ + "0.2.7": {"41ae861478f5e5d9506864f97dfb5dd7cfe6b83193300e2b0d4f04c0380e3af9", "2026-10-27T14:54:42Z", "0.1.52"}, +} + +// strataCandidates are the Strata builds for a machine: its source, in a +// Python environment, with the prebuilt engine for the node's driver and +// cards — built for CUDA 13, Turing to Blackwell; or for CUDA 12.9, Pascal +// to Ada, which is what a V100, or a driver older than 580, gets. +func strataCandidates(version, variant string, h discovery.Host) ([]Build, error) { + if _, err := gpuArch("Strata", h); err != nil { + return nil, err + } + if h.Arch != "amd64" { + return nil, NotHere{"Strata publishes its Linux CUDA engine for x86-64 only"} + } + rel, ok := strataReleases[version] + if !ok { + known := make([]string, 0, len(strataReleases)) + for v := range strataReleases { + known = append(known, v) + } + slices.Sort(known) + return nil, fmt.Errorf("this agent knows the source digest of Strata %s, not of %s; "+ + "install that one by hand and name it in agent.yaml (engines.strata.dir)", + strings.Join(known, ", "), version) + } + published, err := time.Parse(time.RFC3339, rel.published) + if err != nil { + return nil, err + } + build := func(cuda string) Build { + return Build{ + Variant: cuda, + Python: pythonForEngines, + Source: &Source{ + URL: "/" + StrataRepo + "/archive/refs/tags/v" + version + ".tar.gz", + SHA256: rel.sha256, Top: "Strata-" + version, Dir: "strata", + Requirements: "requirements.txt", + }, + AsOf: published.Add(24 * time.Hour), + Assets: []string{"strata-engine-" + rel.engine + "-linux-x86_64-cuda-" + cuda + ".tar.gz"}, + AssetTag: "v" + version, + AssetDir: "strata-engine", + // Strata is started as a script, so what the agent is given is + // the environment's interpreter; the proof is importing what its + // server needs, and finding its engine beside it. + Binary: "python", + VersionArgs: []string{"-c", `import os, sys, numpy, jinja2, regex +e = os.path.dirname(os.path.dirname(os.path.abspath(sys.executable))) +for f in ("strata/setup.py", "strata/serve/server.py", "strata-engine/strata"): + assert os.path.isfile(os.path.join(e, f)), f + " is missing" +print("strata version", "` + version + `", "engine", "` + rel.engine + `")`}, + FreeGB: 5, + } + } + return pick("Strata "+version, variant, []candidate{ + {build("cuda13"), gpuNeeds{minDriverForCUDA13, 7.5}, "Strata " + version + " built for CUDA 13"}, + // 12.9 rather than 12: its machine code is for 6.0 to 8.9 and its + // PTX was compiled by 12.9, so the driver must be 575. + {build("cuda12"), gpuNeeds{575, 6.0}, "Strata " + version + " built for CUDA 12.9"}, + }, h) +} diff --git a/internal/agent/provision/strata_test.go b/internal/agent/provision/strata_test.go new file mode 100644 index 0000000..20c241a --- /dev/null +++ b/internal/agent/provision/strata_test.go @@ -0,0 +1,118 @@ +package provision + +import ( + "crypto/sha256" + "encoding/hex" + "os" + "path/filepath" + "strings" + "testing" + + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery" +) + +// strataRelease is Strata's source archive as GitHub makes one, and its +// release's engine for CUDA 13, with the recipe pointed at the archive. +func strataRelease(t *testing.T, rel *release) { + t.Helper() + archive := tarGz(t, map[string]string{ + "Strata-0.2.7/setup.py": "print('setup')\n", + "Strata-0.2.7/serve/server.py": "print('server')\n", + "Strata-0.2.7/requirements.txt": "numpy==2.5.3; python_version >= \"3.12\"\njinja2==3.1.6\n", + }, nil) + rel.assets["v0.2.7.tar.gz"] = archive + sum := sha256.Sum256(archive) + engine := tarGz(t, map[string]string{"strata": "\x7fELF", "libcudart.so.13": "lib"}, nil) + name := "strata-engine-0.1.52-linux-x86_64-cuda-cuda13.tar.gz" + rel.assets[name] = engine + esum := sha256.Sum256(engine) + rel.digests[name] = "sha256:" + hex.EncodeToString(esum[:]) + + old := strataReleases + strataReleases = map[string]struct { + sha256 string + published string + engine string + }{"0.2.7": {hex.EncodeToString(sum[:]), "2026-10-27T14:54:42Z", "0.1.52"}} + t.Cleanup(func() { strataReleases = old }) +} + +func TestStrataIsInstalledWithItsEngineBesideItsCode(t *testing.T) { + sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}} + in, rel := pythonInstaller(t, turingNode, sys) + strataRelease(t, rel) + + got := in.Converge(t.Context(), []Want{{Engine: "strata", Version: "0.2.7"}}) + if got[0].State != Ready { + t.Fatalf("not ready: %+v", got[0]) + } + dir := filepath.Join(in.Root, "strata", "0.2.7-cuda13") + if got[0].Binary != filepath.Join(dir, "bin", "python") { + t.Errorf("binary %s, want the environment's interpreter", got[0].Binary) + } + for _, f := range []string{"strata/setup.py", "strata/serve/server.py", "strata-engine/strata", "strata-engine/libcudart.so.13"} { + if _, err := os.Stat(filepath.Join(dir, f)); err != nil { + t.Errorf("%s is not in the environment: %v", f, err) + } + } + all := strings.Join(sys.commands, "\n") + for _, want := range []string{"--exclude-newer 2026-10-28T14:54:42Z", "jinja2==3.1.6"} { + if !strings.Contains(all, want) { + t.Errorf("the installation never used %q:\n%s", want, all) + } + } +} + +func TestStrataRefusesAnEngineWhoseDigestDoesNotMatch(t *testing.T) { + sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}} + in, rel := pythonInstaller(t, turingNode, sys) + strataRelease(t, rel) + rel.digests["strata-engine-0.1.52-linux-x86_64-cuda-cuda13.tar.gz"] = "sha256:" + strings.Repeat("0", 64) + + got := in.Converge(t.Context(), []Want{{Engine: "strata", Version: "0.2.7", Variant: "cuda13"}}) + if got[0].State != Failed || !strings.Contains(got[0].Detail, "does not match the checksum") { + t.Fatalf("an engine that does not match its digest was not refused: %+v", got[0]) + } +} + +func TestStrataPicksTheEngineForTheCardsAndDriver(t *testing.T) { + for _, tc := range []struct { + name string + host discovery.Host + want []string + err string + }{ + {"turing on 580", turingNode, []string{"cuda13", "cuda12"}, ""}, + {"volta", discovery.Host{Arch: "amd64", NVIDIADriver: "580.82.07", MinComputeCapability: 7.0}, []string{"cuda12"}, ""}, + {"driver 575", discovery.Host{Arch: "amd64", NVIDIADriver: "575.51.03", MinComputeCapability: 8.6}, []string{"cuda12"}, ""}, + {"driver 570", discovery.Host{Arch: "amd64", NVIDIADriver: "570.86.10", MinComputeCapability: 8.6}, nil, "575"}, + {"no gpu", discovery.Host{Arch: "amd64"}, nil, "no"}, + {"arm", discovery.Host{Arch: "arm64", NVIDIADriver: "580.82.07", MinComputeCapability: 9.0}, nil, "x86-64"}, + } { + t.Run(tc.name, func(t *testing.T) { + bs, err := strataCandidates("0.2.7", "", tc.host) + if tc.err != "" { + if err == nil || !strings.Contains(err.Error(), tc.err) { + t.Fatalf("got %v, %v; want an error saying %q", bs, err, tc.err) + } + return + } + if err != nil { + t.Fatal(err) + } + var got []string + for _, b := range bs { + got = append(got, b.Variant) + } + if strings.Join(got, ",") != strings.Join(tc.want, ",") { + t.Errorf("builds %v, want %v", got, tc.want) + } + }) + } +} + +func TestStrataIsNotInAnOfflineBundle(t *testing.T) { + if _, err := Builds("strata", "0.2.7", "amd64"); err == nil { + t.Fatal("an offline bundle would carry Strata, which needs PyPI") + } +} diff --git a/internal/cli/root.go b/internal/cli/root.go index 48c4ee7..b93d82e 100644 --- a/internal/cli/root.go +++ b/internal/cli/root.go @@ -55,6 +55,7 @@ func Root() *cobra.Command { newBundleCmd(), newSimEngineCmd(), newComfyUIGateCmd(), + newStrataServeCmd(), newDoctorCmd(), newNodeCmd(), newSelfUpgradeCmd(), diff --git a/internal/cli/strataserve.go b/internal/cli/strataserve.go new file mode 100644 index 0000000..5f5c2f8 --- /dev/null +++ b/internal/cli/strataserve.go @@ -0,0 +1,43 @@ +package cli + +import ( + "os/signal" + "syscall" + + "github.com/spf13/cobra" + + "gitea.mynodes.xyz/ilker/llmgw/internal/strata" +) + +func newStrataServeCmd() *cobra.Command { + var cfg strata.Config + cmd := &cobra.Command{ + Use: "strata-serve [flags] [-- setup.py flags]", + Short: "Run Strata as an engine: prepare the model, then serve it (started by the agent)", + Hidden: true, + Long: "Started by the node agent for a replica whose engine is strata. Prepares the\n" + + "model from the fetched GGUF files with Strata's setup, which skips what it has\n" + + "done before, then becomes Strata's server on the replica's address. The key is\n" + + "read from " + strata.KeyEnv + ".", + RunE: func(cmd *cobra.Command, args []string) error { + ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + cfg.Args = args + cfg.Log = logger() + return strata.Run(ctx, cfg) + }, + } + f := cmd.Flags() + f.StringVar(&cfg.Listen, "listen", "", "the replica's address") + f.StringVar(&cfg.Model, "model", "", "the model's name, for /v1/models") + f.StringVar(&cfg.Weights, "weights", "", "the fetched GGUF files' directory") + f.StringVar(&cfg.State, "state", "", "the replica's own directory") + f.StringVar(&cfg.Data, "data", "", "the node's directory for prepared models, shared by its replicas") + f.StringVar(&cfg.Python, "python", "python3", "the interpreter of Strata's environment") + f.StringVar(&cfg.Dir, "strata", "", "Strata's directory, where setup.py is") + f.StringVar(&cfg.Engine, "engine", "", "Strata's engine directory, where the strata binary is") + for _, n := range []string{"listen", "weights", "state", "data", "strata", "engine"} { + _ = cmd.MarkFlagRequired(n) + } + return cmd +} diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index 964863a..28e899c 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -1191,6 +1191,22 @@ func (m Model) validate() error { } continue } + // Strata prepares one of the few families it runs from their GGUF + // files, on one card, and has to be told which. + if p.Engine == "strata" { + if p.Nodes > 1 || p.TP > 1 || p.PP > 1 || p.GPUsPerNode > 1 { + return fmt.Errorf("model %q profile %q: strata runs a replica on one GPU of one node", m.Name, p.Name) + } + if !slices.ContainsFunc(p.Args, func(a string) bool { return a == "--family" || strings.HasPrefix(a, "--family=") }) { + return fmt.Errorf("model %q profile %q: strata needs the model's family and size in args "+ + "(--family qwen --size IQ2_XS)", m.Name, p.Name) + } + if strings.HasSuffix(strings.ToLower(strings.TrimSuffix(m.SourceOf(p), "/")), ".gguf") { + return fmt.Errorf("model %q profile %q: strata prepares a model from all of its GGUF files; "+ + "give the directory they are in as the source", m.Name, p.Name) + } + continue + } // TensorRT-LLM spans machines only under MPI started from one of them, // which the agents, each starting its own node, do not do; and it // loads a Hugging Face checkpoint, not llama.cpp's format. diff --git a/internal/control/catalog/strata_test.go b/internal/control/catalog/strata_test.go new file mode 100644 index 0000000..199bf8c --- /dev/null +++ b/internal/control/catalog/strata_test.go @@ -0,0 +1,31 @@ +package catalog + +import ( + "strings" + "testing" +) + +func TestAStrataProfileNamesItsFamilyAndRunsOnOneCard(t *testing.T) { + ok := Model{Name: "q38", Source: "http://mirror:8081/q38/ab12cd34/", Profiles: []Profile{{ + Name: "p", Engine: "strata", GPUClass: "6000-24g", + Args: []string{"--family", "qwen", "--size", "IQ2_XS", "--context", "32768"}, + }}} + if _, err := New([]Model{ok}); err != nil { + t.Fatalf("a one-card Strata profile was refused: %v", err) + } + for want, bad := range map[string]Model{ + "one GPU of one node": {Name: "x", Source: "http://mirror:8081/x/1/", Profiles: []Profile{{ + Name: "p", Engine: "strata", GPUClass: "6000-24g", GPUsPerNode: 2, TP: 2, + Args: []string{"--family", "qwen"}}}}, + "family and size": {Name: "x", Source: "http://mirror:8081/x/1/", Profiles: []Profile{{ + Name: "p", Engine: "strata", GPUClass: "6000-24g", Args: []string{"--size", "IQ2_XS"}}}}, + "all of its GGUF files": {Name: "x", Source: "http://mirror:8081/x/1/x-00001-of-00002.gguf", Profiles: []Profile{{ + Name: "p", Engine: "strata", GPUClass: "6000-24g", Args: []string{"--family", "qwen"}}}}, + "--gguf-dir": {Name: "x", Source: "http://mirror:8081/x/1/", Profiles: []Profile{{ + Name: "p", Engine: "strata", GPUClass: "6000-24g", Args: []string{"--family", "qwen", "--gguf-dir", "/etc"}}}}, + } { + if _, err := New([]Model{bad}); err == nil || !strings.Contains(err.Error(), want) { + t.Errorf("expected %q, got %v", want, err) + } + } +} diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index d8c8618..b36d0b7 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -505,6 +505,11 @@ var Reference = []Block{ {Package: "engines", Name: "SimConfig", Title: "SimConfig", Doc: "", Fields: []Field{ {Key: "bin", Type: "string", Doc: "Bin is the llmgw binary itself; the engine is one of its subcommands.\nEmpty means whatever started this agent, which is right almost always.", Nested: ""}, }}, + {Package: "engines", Name: "StrataConfig", Title: "StrataConfig", Doc: "", Fields: []Field{ + {Key: "dir", Type: "string", Doc: "Dir is Strata's source directory, where setup.py and serve/ are.", Nested: ""}, + {Key: "python", Type: "string", Doc: "Python is the interpreter of its environment; Dir's .venv when unset.", Nested: ""}, + {Key: "engine", Type: "string", Doc: "Engine is the directory of its prebuilt engine, where the strata\nbinary and the CUDA libraries it ships are; Dir's engine/ when unset.", Nested: ""}, + }}, {Package: "engines", Name: "TRTLLMConfig", Title: "TRTLLMConfig", Doc: "", Fields: []Field{ {Key: "python", Type: "string", Doc: "Python is the interpreter of the environment TensorRT-LLM is installed\nin: /bin/python for a virtual environment, python3 in NVIDIA's\ncontainer. Naming it is what makes this node offer the engine; empty\noffers nothing.\n\nThe interpreter rather than trtllm-serve, because llmgw starts the\nserver through a few lines of its own that make it demand the\ngateway's key, which trtllm-serve has no option for.", Nested: ""}, }}, @@ -614,6 +619,7 @@ var Reference = []Block{ {Key: "trtllm", Type: "TRTLLMConfig", Doc: "TRTLLM is TensorRT-LLM, which the agent does not install: a node\noffers it only when its interpreter is named here.", Nested: "engines.TRTLLMConfig"}, {Key: "sim", Type: "SimConfig", Doc: "", Nested: "engines.SimConfig"}, {Key: "comfyui", Type: "ComfyUIConfig", Doc: "ComfyUI makes pictures; offered by a node whose ComfyUI is named here.", Nested: "engines.ComfyUIConfig"}, + {Key: "strata", Type: "StrataConfig", Doc: "Strata serves a few mixture-of-experts models from RAM and one card;\noffered by a node whose Strata is named here, or that installed it.", Nested: "engines.StrataConfig"}, }}, {Package: "guard", Name: "Config", Title: "guard.Config", Doc: "", Fields: []Field{ {Key: "enabled", Type: "*bool", Doc: "Enabled turns every policy on or off at once: false checks nothing,\nwhatever is defined or attached. Unset is on.", Nested: ""}, diff --git a/internal/docgen/render.go b/internal/docgen/render.go index a10e3e2..98edc8a 100644 --- a/internal/docgen/render.go +++ b/internal/docgen/render.go @@ -23,6 +23,7 @@ var Sources = []string{ "internal/agent/engines/trtllm.go", "internal/agent/engines/sim.go", "internal/agent/engines/comfyui.go", + "internal/agent/engines/strata.go", "internal/agent/supervisor/supervisor.go", "internal/tlsconf/tlsconf.go", "internal/control/catalog/catalog.go", diff --git a/internal/router/scrape/parse.go b/internal/router/scrape/parse.go index 6b4d616..144a001 100644 --- a/internal/router/scrape/parse.go +++ b/internal/router/scrape/parse.go @@ -50,6 +50,12 @@ func MappingFor(engine string) Mapping { return TRTLLMMapping() case "sglang": return SGLangMapping() + case "strata": + // vLLM's names, under a path of its own: Strata's /metrics is the + // JSON its monitor reads. + m := VLLMMapping() + m.Path = "/metrics/prometheus" + return m default: return VLLMMapping() } diff --git a/internal/strata/serve.go b/internal/strata/serve.go new file mode 100644 index 0000000..5280e74 --- /dev/null +++ b/internal/strata/serve.go @@ -0,0 +1,204 @@ +// Package strata runs Strata as an llmgw engine: it prepares a model with +// Strata's own setup, from the GGUF files the agent fetched, and then becomes +// Strata's server. +// +// Strata serves a few model families, each prepared from its GGUF files into +// a pack and a draft layer before it can load: work its setup does, and +// which takes minutes the first time and seconds after. Setup runs the way +// Strata's desktop app runs it (STRATA_CONTROL_INSTALL): it installs nothing +// and builds nothing, and writes only under the directories it is given, so +// the code the agent installed stays read-only. +package strata + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "net" + "os" + "os/exec" + "path/filepath" + "slices" + "strings" + "syscall" +) + +// Config is one replica's. +type Config struct { + // Listen is the replica's address. + Listen string + // Model is the name clients use. + Model string + // Weights is the directory of the model's GGUF files, or, when the + // catalog named a repository the agent does not fetch, that name, and + // Strata's setup downloads the files itself. + Weights string + // State is the replica's own directory; Data the node's, shared by its + // replicas, where prepared models and draft layers are kept. + State, Data string + // Python is the interpreter of Strata's environment, Dir its code and + // Engine its prebuilt engine. + Python, Dir, Engine string + // Args go to Strata's setup, which turns them into the server's + // configuration; --size is its --model. + Args []string + Log *slog.Logger +} + +// KeyEnv is where the replica's key comes in, and where Strata's server +// reads it from. +const KeyEnv = "STRATA_API_KEY" + +// Run prepares the model and execs Strata's server, which then is this +// process: the agent's supervisor watches it, and stops it, as the engine. +func Run(ctx context.Context, c Config) error { + host, port, err := net.SplitHostPort(c.Listen) + if err != nil { + return fmt.Errorf("--listen %q: %w", c.Listen, err) + } + for _, d := range []*string{&c.State, &c.Data} { + if *d, err = filepath.Abs(*d); err != nil { + return err + } + } + for _, d := range []string{c.State, c.Data} { + if err := os.MkdirAll(d, 0o700); err != nil { + return err + } + } + cfgPath, err := prepare(ctx, c) + if err != nil { + return err + } + if err := adopt(cfgPath, c, host); err != nil { + return err + } + argv := []string{c.Python, filepath.Join(c.Dir, "serve", "server.py"), + "--engine", "strata", "--config", cfgPath, "--host", host, "--port", port} + c.Log.Info("strata: starting the server", "config", cfgPath, "listen", c.Listen) + env := append(os.Environ(), controlEnv(c)...) + return syscall.Exec(c.Python, argv, env) //nolint:gosec // the installed interpreter, on Strata's own server +} + +// controlEnv has Strata's setup and server write only under the replica's +// directories, never the engine user's home, and look for no engine but +// the one named in the configuration. +func controlEnv(c Config) []string { + return []string{ + "STRATA_CONTROL_INSTALL=1", + "STRATA_INSTALL_DIR=" + c.State, + "STRATA_SETTINGS_FILE=" + filepath.Join(c.State, "settings.json"), + "STRATA_COMPONENTS_DIR=" + filepath.Join(c.State, "components"), + } +} + +// prepare runs Strata's setup on the fetched files, and returns the run +// configuration it wrote. Each step setup has already done is skipped, so a +// replica started again is ready in seconds. +func prepare(ctx context.Context, c Config) (string, error) { + args, err := SetupArgs(c.Args) + if err != nil { + return "", err + } + argv := []string{filepath.Join(c.Dir, "setup.py"), "--data-dir", c.Data, "--no-start", "--yes", "--no-browser"} + if st, err := os.Stat(c.Weights); err == nil && st.IsDir() { + argv = append(argv, "--gguf-dir", c.Weights) + } else { + // A repository rather than files the agent fetched: setup downloads + // the files it knows for the family and size, each pinned to a + // commit and checked against its SHA-256, into the node's data. + argv = append(argv, "--models-dir", filepath.Join(c.Data, "models")) + } + argv = append(argv, args...) + cmd := exec.CommandContext(ctx, c.Python, argv...) //nolint:gosec // the installed interpreter, on Strata's own setup + cmd.Dir = c.Dir + // Setup is given no key: it would write one into the configuration. + cmd.Env = append(slices.DeleteFunc(os.Environ(), func(kv string) bool { + return strings.HasPrefix(kv, KeyEnv+"=") + }), controlEnv(c)...) + cmd.Stdout, cmd.Stderr = os.Stderr, os.Stderr + c.Log.Info("strata: preparing the model", "weights", c.Weights, "args", args) + if err := cmd.Run(); err != nil { + return "", fmt.Errorf("strata's setup did not prepare the model: %w", err) + } + return newestConfig(c.State) +} + +// SetupArgs are a profile's arguments as Strata's setup takes them: --size +// is its --model, which a profile cannot name, being the source everywhere +// else. +func SetupArgs(profile []string) ([]string, error) { + out := make([]string, 0, len(profile)) + for _, a := range profile { + switch { + case a == "--size": + out = append(out, "--model") + case strings.HasPrefix(a, "--size="): + out = append(out, "--model="+strings.TrimPrefix(a, "--size=")) + default: + out = append(out, a) + } + } + if !slices.ContainsFunc(out, func(a string) bool { return a == "--family" || strings.HasPrefix(a, "--family=") }) { + return nil, errors.New("strata needs the model's family in the profile's args (--family qwen, qwen36, ...)") + } + return out, nil +} + +// newestConfig is the run configuration setup wrote last: one per model +// size, strata-.json. +func newestConfig(dir string) (string, error) { + matches, err := filepath.Glob(filepath.Join(dir, "strata-*.json")) + if err != nil { + return "", err + } + var best string + var bestMod int64 + for _, m := range matches { + st, err := os.Stat(m) + if err != nil { + continue + } + if t := st.ModTime().UnixNano(); best == "" || t > bestMod { + best, bestMod = m, t + } + } + if best == "" { + return "", fmt.Errorf("strata's setup wrote no run configuration in %s", dir) + } + return best, nil +} + +// adopt points setup's configuration at the engine the agent installed, and +// at the replica's name and address. Setup in this mode names the engine a +// component of Strata's own app, which a node does not have. +func adopt(path string, c Config, host string) error { + b, err := os.ReadFile(path) + if err != nil { + return err + } + var cfg map[string]any + if err := json.Unmarshal(b, &cfg); err != nil { + return fmt.Errorf("%s: %w", path, err) + } + exe := filepath.Join(c.Engine, "strata") + if _, err := os.Stat(exe); err != nil { + return fmt.Errorf("strata's engine is not where it was installed: %w", err) + } + cfg["exe"] = exe + cfg["lib_dirs"] = []string{c.Engine} + cfg["host"] = host + cfg["open_browser"] = false + if c.Model != "" { + cfg["model_name"] = c.Model + } + // The key comes in the environment only. + delete(cfg, "api_key") + out, err := json.MarshalIndent(cfg, "", " ") + if err != nil { + return err + } + return os.WriteFile(path, out, 0o600) +} diff --git a/internal/strata/serve_test.go b/internal/strata/serve_test.go new file mode 100644 index 0000000..710f589 --- /dev/null +++ b/internal/strata/serve_test.go @@ -0,0 +1,76 @@ +package strata + +import ( + "encoding/json" + "os" + "path/filepath" + "slices" + "testing" + "time" +) + +func TestSizeIsSetupsModel(t *testing.T) { + got, err := SetupArgs([]string{"--family", "qwen", "--size", "IQ2_XS", "--context", "32768"}) + if err != nil { + t.Fatal(err) + } + if want := []string{"--family", "qwen", "--model", "IQ2_XS", "--context", "32768"}; !slices.Equal(got, want) { + t.Errorf("got %v, want %v", got, want) + } + if got, _ := SetupArgs([]string{"--family=qwen36", "--size=UD-Q4_K_XL"}); !slices.Equal(got, []string{"--family=qwen36", "--model=UD-Q4_K_XL"}) { + t.Errorf("the = forms: %v", got) + } + if _, err := SetupArgs([]string{"--size", "IQ2_XS"}); err == nil { + t.Error("a profile with no family was accepted") + } +} + +func TestTheConfigurationRunsTheInstalledEngineWithNoKey(t *testing.T) { + state, engine := t.TempDir(), t.TempDir() + if err := os.WriteFile(filepath.Join(engine, "strata"), []byte("\x7fELF"), 0o755); err != nil { + t.Fatal(err) + } + old := filepath.Join(state, "strata-q2_0.json") + path := filepath.Join(state, "strata-iq2_xs.json") + for _, p := range []string{old, path} { + b, _ := json.Marshal(map[string]any{"exe": "component:strata-engine", "args": []string{"--pack", "/p"}, + "model_name": "qwen3.8-flash-next-iq2_xs", "api_key": "from-setup", "host": "127.0.0.1"}) + if err := os.WriteFile(p, b, 0o600); err != nil { + t.Fatal(err) + } + } + past := time.Now().Add(-time.Hour) + _ = os.Chtimes(old, past, past) + + got, err := newestConfig(state) + if err != nil || got != path { + t.Fatalf("newest configuration %q, %v; want %q", got, err, path) + } + c := Config{Engine: engine, Model: "qwen3.8-flash-next"} + if err := adopt(path, c, "10.0.0.147"); err != nil { + t.Fatal(err) + } + var cfg map[string]any + b, _ := os.ReadFile(path) + if err := json.Unmarshal(b, &cfg); err != nil { + t.Fatal(err) + } + if cfg["exe"] != filepath.Join(engine, "strata") || cfg["host"] != "10.0.0.147" || cfg["model_name"] != "qwen3.8-flash-next" { + t.Errorf("configuration not adopted: %v", cfg) + } + if _, ok := cfg["api_key"]; ok { + t.Error("the configuration keeps a key; it comes in the environment only") + } + if args, _ := cfg["args"].([]any); len(args) != 2 { + t.Errorf("setup's engine arguments were lost: %v", cfg["args"]) + } +} + +func TestAMissingEngineIsSaid(t *testing.T) { + state := t.TempDir() + path := filepath.Join(state, "strata-x.json") + _ = os.WriteFile(path, []byte(`{"exe":"component:strata-engine"}`), 0o600) + if err := adopt(path, Config{Engine: filepath.Join(state, "nowhere")}, "0.0.0.0"); err == nil { + t.Error("a configuration was pointed at an engine that is not there") + } +}