diff --git a/.gitea/workflows/engine-releases.yml b/.gitea/workflows/engine-releases.yml new file mode 100644 index 0000000..6dd42c0 --- /dev/null +++ b/.gitea/workflows/engine-releases.yml @@ -0,0 +1,20 @@ +# Fails when an engine llmgw follows release by release (FreeToken) has +# published a version newer than the one DefaultEngines installs: the failed +# run is the notice to check the new version and move to it. See +# scripts/engine_releases.py. +name: Engine releases + +on: + schedule: + # Mondays 06:00 UTC. + - cron: "0 6 * * 1" + workflow_dispatch: + +jobs: + check: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Compare with PyPI + run: python3 scripts/engine_releases.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e43a186..089fb85 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,19 @@ APIs change between minors without a deprecation period. ### Added +- **FreeToken, an engine for models larger than the card.** Mixture-of-experts + models several times a card's size, from one card, with the experts in RAM + and streamed over PCIe — read from Hugging Face safetensors checkpoints as + published (DeepSeek V4 Flash, GLM-5.x, Qwen3.6/3.8, gpt-oss, Gemma 4, + MiniMax), with batching. Installed like vLLM on every node it runs on: a + card of compute capability 8.0 or newer, driver 580+, x86-64, and the CUDA 13 + toolkit it compiles kernels with; older cards are passed over (tried on a + Turing card, where 0.2.1 has no kernels). Profiles are one card + (`engine: freetoken`); the key goes in its environment, and flags that read + files or fetch are refused. Readiness now also waits for an engine whose + `/health` answers 200 while saying `loading` in its body. FreeToken is + followed release by release: a weekly workflow fails when PyPI has a version + newer than `DefaultEngines` installs. - **How much of each GPU works, beside whether anything does.** The busy figure is nvidia-smi's `utilization.gpu` — the share of time any kernel ran, which a kernel on one SM fills — and it is now said to be that: on diff --git a/docs/install.md b/docs/install.md index fa49961..ffa6e36 100644 --- a/docs/install.md +++ b/docs/install.md @@ -703,6 +703,7 @@ Nothing to do: the agent installs them once the node joins, under | vLLM | nodes with an NVIDIA GPU it can run on | a Python environment built by uv from PyPI, about 8 GB | | SGLang | only when you ask for it | the same way as vLLM | | TensorRT-LLM | only nodes where you installed it | not installed by llmgw: see [TensorRT-LLM, which you bring](#tensorrt-llm-which-you-bring) | +| FreeToken | nodes with a GPU of compute capability 8.0 or newer (RTX 30, A100 and later), driver 580+, x86-64 and the CUDA 13 toolkit | a Python environment built by uv from PyPI (`freetoken[accel]`); see [FreeToken](#freetoken-models-larger-than-the-card) | | Strata | only when you ask for it (`engines: {strata: "0.2.7"}`) | its source, pinned by digest, in a Python environment, and its prebuilt engine for CUDA 13 (driver 580+) or CUDA 12.9 (Pascal to Ada, driver 575+); see [Strata](#strata-a-large-model-of-experts-on-one-card) | vLLM is resolved as PyPI stood a day after that vLLM was published, so every @@ -785,6 +786,61 @@ in (48 GB and the engine's 4 for IQ2_XS) unless it says `min_memory_gb`. Its first start prepares the model on the node, which takes minutes, and fetches a draft layer of about 5 GB from the internet. +#### FreeToken: models larger than the card + +[FreeToken](https://github.com/FlashML-org/FreeToken) serves mixture-of-experts +models several times a card's size from one card: the experts stay in the +node's RAM, the ones used most are cached in VRAM, and the rest stream over +PCIe. Unlike Strata it reads Hugging Face safetensors checkpoints as they are +published — DeepSeek V4 Flash, GLM-5.x, Qwen3.6/3.8 (Flash-Next among them), +gpt-oss, Gemma 4, MiniMax M2.5/M3 — in FP8, NVFP4, MXFP4 or BF16, and batches +several requests (`--max-running-requests`, 4 by default). + +Every node that can run it installs it, as with vLLM: an NVIDIA GPU of compute +capability 8.0 or newer, driver 580 or newer, x86-64. Older cards are passed +over without a word — on a Quadro RTX 6000 (7.5) version 0.2.1 had no kernel +image for the card's first prefill, and in fp16 its first request ended in an +illegal memory access. Its FP8 checkpoints need 8.9 (Ada, Hopper) and its NVFP4 +ones 12.0 (RTX 50, RTX PRO 6000 Blackwell). It compiles some kernels the first +time a model loads, so the node also needs the **CUDA 13 toolkit** (`nvcc` in +`CUDA_HOME`, `/usr/local/cuda` or `/usr/local/cuda-13.*`); a node of new cards +without one says so on the cluster page instead of installing an engine that +fails on its first model. + +A FreeToken profile runs on one card of one node (`llmgw ctl model add -f`): + +```yaml +name: glm-5.2 +source: /var/lib/llmgw-weights/GLM-5.2-NVFP4 # a safetensors checkpoint's directory +profiles: + - name: one-card + engine: freetoken + gpu_class: rtx-pro-6000-96g + min_memory_gb: 200 # the experts live in RAM: say how much + args: ["--max-running-requests", "8"] +``` + +A Hugging Face id as the source works too — FreeToken downloads the +repository itself — but it is not pinned to a commit or checked against +Hugging Face's checksums as vLLM's are, FreeToken having no `--revision`; a +directory the mirror fetched, or one put on the node, is. + +The agent gives it the model, its address, its card (`--gpu`) and the key (in +`FREETOKEN_API_KEY`, not on the command line); `--mmproj`, media paths and +domains, `--hf-overrides` and `--moe-cache-split` are refused as they read +files or fetch. `ctl model context` sets `--max-seq-len-override`. Readiness +waits for its `/health` to stop saying `loading`, since it answers 200 from the +moment its HTTP server is up. It serves no Prometheus metrics, so the router +counts its requests itself. + +llmgw follows FreeToken release by release: the version is +`DefaultEngines["freetoken"]`, and a weekly workflow +(`.gitea/workflows/engine-releases.yml`) fails when PyPI has a newer one, the +notice to check it and move to it. A version of your own is named under +`engines` like any other (`freetoken: "0.2.2"`), and `freetoken: ""` installs +none; one node is turned off with its build set to *none* (console, the node's +menu). + #### Choosing a node's vLLM build Which build a GPU node installs is chosen on the cluster page, in the **vLLM diff --git a/docs/reference.md b/docs/reference.md index 5ffe50d..47351a3 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -126,6 +126,7 @@ rather than repeated. | `name` | `string` | [Federation](#federation) | | `interval` | `time.Duration` | [Federation](#federation) | | `enabled` | `bool` | [FirewallSettings](#firewallsettings) | +| `binary` | `string` | [FreeTokenConfig](#freetokenconfig) | | `min_compute_capability` | `float64` | [GPUNeed](#gpuneed) | | `min_memory_gb` | `int` | [GPUNeed](#gpuneed) | | `format` | `string` | [GuardModel](#guardmodel) | @@ -482,6 +483,7 @@ rather than repeated. | `sim` | `SimConfig` | [engines.Config](#enginesconfig) | | `comfyui` | `ComfyUIConfig` | [engines.Config](#enginesconfig) | | `strata` | `StrataConfig` | [engines.Config](#enginesconfig) | +| `freetoken` | `FreeTokenConfig` | [engines.Config](#enginesconfig) | | `enabled` | `*bool` | [guard.Config](#guardconfig) | | `default` | `[]string` | [guard.Config](#guardconfig) | | `policies` | `map[string]Policy` | [guard.Config](#guardconfig) | @@ -1317,6 +1319,16 @@ site whose own firewall does that job; the rules are kept, and true puts them back. Changed in the cluster's settings, it is in force on every node and control plane within seconds. +## FreeTokenConfig + +### `binary` + +`string` + +Binary is its ft command, in the environment it was installed into. +The agent installs one itself on a node whose GPUs it runs on; this +is for one installed by hand. + ## GPUNeed ### `min_compute_capability` @@ -4042,6 +4054,16 @@ offered by a node whose Strata is named here, or that installed it. Takes a [StrataConfig](#strataconfig) block. +### `freetoken` + +`FreeTokenConfig` + +FreeToken serves mixture-of-experts models larger than a card from +RAM and one card of 8.0 or newer; installed by the agent where the +cards allow it, or named here. + +Takes a [FreeTokenConfig](#freetokenconfig) block. + ## guard.Config ### `enabled` diff --git a/internal/agent/discovery/cudatoolkit.go b/internal/agent/discovery/cudatoolkit.go new file mode 100644 index 0000000..c1c43a4 --- /dev/null +++ b/internal/agent/discovery/cudatoolkit.go @@ -0,0 +1,93 @@ +package discovery + +import ( + "encoding/json" + "os" + "os/exec" + "path/filepath" + "regexp" + "slices" + "strconv" + "strings" +) + +// CUDAToolkit is the root of an installed CUDA toolkit of this major version +// — the directory whose bin holds nvcc — or "" when there is none. Engines +// that compile their kernels where they run (FreeToken, ComfyUI's compiled +// nodes) need one matching the CUDA their torch was built for, and a +// toolkit of another major version compiles code that torch cannot load. +// +// CUDA_HOME is looked at first, then /usr/local/cuda, then every +// /usr/local/cuda-, newest first. +func CUDAToolkit(major int) string { + return cudaToolkitIn(major, os.Getenv("CUDA_HOME"), "/usr/local") +} + +func cudaToolkitIn(major int, home, root string) string { + var dirs []string + if home != "" { + dirs = append(dirs, home) + } + dirs = append(dirs, filepath.Join(root, "cuda")) + versioned, _ := filepath.Glob(filepath.Join(root, "cuda-*")) + slices.SortFunc(versioned, func(a, b string) int { return versionCompare(b, a) }) + dirs = append(dirs, versioned...) + for _, d := range dirs { + if _, err := os.Stat(filepath.Join(d, "bin", "nvcc")); err != nil { + continue + } + if toolkitMajor(d) == major { + return d + } + } + return "" +} + +// toolkitMajor is a toolkit's major version, from the version.json every +// toolkit since 11.1 carries, or from nvcc where it has none; 0 when neither +// says. +func toolkitMajor(dir string) int { + if b, err := os.ReadFile(filepath.Join(dir, "version.json")); err == nil { //nolint:gosec // a toolkit's own file + var v struct { + CUDA struct { + Version string `json:"version"` + } `json:"cuda"` + } + if json.Unmarshal(b, &v) == nil && v.CUDA.Version != "" { + n, _ := strconv.Atoi(strings.SplitN(v.CUDA.Version, ".", 2)[0]) + return n + } + } + //nolint:gosec // nvcc of a toolkit found on disk + out, err := exec.Command(filepath.Join(dir, "bin", "nvcc"), "--version").Output() + if err != nil { + return 0 + } + if m := nvccRelease.FindSubmatch(out); m != nil { + n, _ := strconv.Atoi(string(m[1])) + return n + } + return 0 +} + +var nvccRelease = regexp.MustCompile(`release (\d+)\.`) + +// versionCompare orders two paths ending in dotted versions: cuda-13.10 is +// newer than cuda-13.4, which comparing them as text gets wrong. +func versionCompare(a, b string) int { + pa := strings.Split(strings.TrimPrefix(filepath.Base(a), "cuda-"), ".") + pb := strings.Split(strings.TrimPrefix(filepath.Base(b), "cuda-"), ".") + for i := range max(len(pa), len(pb)) { + var x, y int + if i < len(pa) { + x, _ = strconv.Atoi(pa[i]) + } + if i < len(pb) { + y, _ = strconv.Atoi(pb[i]) + } + if x != y { + return x - y + } + } + return 0 +} diff --git a/internal/agent/discovery/cudatoolkit_test.go b/internal/agent/discovery/cudatoolkit_test.go new file mode 100644 index 0000000..5bc7bab --- /dev/null +++ b/internal/agent/discovery/cudatoolkit_test.go @@ -0,0 +1,48 @@ +package discovery + +import ( + "os" + "path/filepath" + "testing" +) + +func toolkit(t *testing.T, root, name, version string) string { + t.Helper() + d := filepath.Join(root, name) + put(t, d, "bin/nvcc", "#!/bin/sh") + if version != "" { + put(t, d, "version.json", `{"cuda":{"name":"CUDA SDK","version":"`+version+`"}}`) + } + return d +} + +func TestTheNewestToolkitOfTheMajorVersionAskedForIsFound(t *testing.T) { + root := t.TempDir() + toolkit(t, root, "cuda-12.8", "12.8.1") + toolkit(t, root, "cuda-13.4", "13.4.0") + want := toolkit(t, root, "cuda-13.10", "13.10.0") + // /usr/local/cuda is a 12: a toolkit of the wrong major version is not + // one torch built for 13 can load what it compiles. + toolkit(t, root, "cuda", "12.8.1") + if got := cudaToolkitIn(13, "", root); got != want { + t.Errorf("got %q, want %q", got, want) + } + if got := cudaToolkitIn(12, "", root); got != filepath.Join(root, "cuda") { + t.Errorf("12: got %q", got) + } + // CUDA_HOME first, when it is one. + home := toolkit(t, t.TempDir(), "mine", "13.0.2") + if got := cudaToolkitIn(13, home, root); got != home { + t.Errorf("CUDA_HOME: got %q", got) + } + if got := cudaToolkitIn(14, "", root); got != "" { + t.Errorf("14: got %q", got) + } + // A directory without nvcc is not a toolkit. + if err := os.MkdirAll(filepath.Join(root, "cuda-13.99"), 0o755); err != nil { + t.Fatal(err) + } + if got := cudaToolkitIn(13, "", root); got != want { + t.Errorf("with an empty cuda-13.99: got %q", got) + } +} diff --git a/internal/agent/engines/engine.go b/internal/agent/engines/engine.go index 5a86470..acba5a5 100644 --- a/internal/agent/engines/engine.go +++ b/internal/agent/engines/engine.go @@ -223,6 +223,7 @@ func NewRegistry(cfg Config) *Registry { r.Add(&Sim{cfg: cfg.Sim}) r.Add(&ComfyUI{cfg: cfg.ComfyUI}) r.Add(&Strata{cfg: cfg.Strata}) + r.Add(&FreeToken{cfg: cfg.FreeToken}) return r } @@ -276,6 +277,10 @@ type Config struct { // 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"` + // FreeToken serves mixture-of-experts models larger than a card from + // RAM and one card of 8.0 or newer; installed by the agent where the + // cards allow it, or named here. + FreeToken FreeTokenConfig `koanf:"freetoken"` } // --- shared helpers --- diff --git a/internal/agent/engines/engines_test.go b/internal/agent/engines/engines_test.go index b0548c8..88a5b57 100644 --- a/internal/agent/engines/engines_test.go +++ b/internal/agent/engines/engines_test.go @@ -477,7 +477,7 @@ 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", "strata", "trtllm", "vllm"}) { + if got := r.Names(); !slices.Equal(got, []string{"comfyui", "freetoken", "llamacpp", "sglang", "sim", "strata", "trtllm", "vllm"}) { t.Errorf("registry contents: %v", got) } if _, err := r.Get("vllm"); err != nil { diff --git a/internal/agent/engines/freetoken.go b/internal/agent/engines/freetoken.go new file mode 100644 index 0000000..421e6f3 --- /dev/null +++ b/internal/agent/engines/freetoken.go @@ -0,0 +1,137 @@ +package engines + +import ( + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "sync/atomic" + + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery" +) + +// FreeTokenConfig locates a FreeToken installation. +type FreeTokenConfig struct { + // Binary is its ft command, in the environment it was installed into. + // The agent installs one itself on a node whose GPUs it runs on; this + // is for one installed by hand. + Binary string `koanf:"binary"` +} + +// freetokenToolkit finds the CUDA 13 toolkit; a variable for tests. +var freetokenToolkit = func() string { return discovery.CUDAToolkit(13) } + +// FreeToken runs FreeToken (github.com/FlashML-org/FreeToken): a +// mixture-of-experts engine that keeps a model's experts in host RAM, the +// ones used most in VRAM, and streams the rest over PCIe, so a model several +// times a card's size serves from one card. It reads Hugging Face +// safetensors checkpoints as they are published (FP8, NVFP4, MXFP4, BF16) +// and serves the OpenAI and Anthropic APIs. One card of compute capability +// 8.0 or newer; it compiles some of its kernels on first start, with the +// CUDA 13 toolkit. +type FreeToken struct { + cfg FreeTokenConfig + installed atomic.Pointer[string] +} + +// Name identifies the engine in the catalog. +func (*FreeToken) Name() string { return "freetoken" } + +// HealthPath answers before the model is loaded too, with its progress in +// the body; the supervisor waits for the status code, so it counts as ready +// only once FreeToken's own startup — loading, warm-up, CUDA graphs — is +// over and the API server is accepting. +func (*FreeToken) HealthPath() string { return "/health" } + +// MetricsPath is empty: FreeToken serves its numbers as JSON (/v1/stats), +// not in Prometheus' text, so the router counts its requests itself. +func (*FreeToken) MetricsPath() string { return "" } + +// ServesHTTP is true: a FreeToken replica is one process on one node. +func (*FreeToken) ServesHTTP(ShardSpec) bool { return true } + +// UseInstalled points the engine at the ft the agent installed. +func (f *FreeToken) UseInstalled(path string) { f.installed.Store(&path) } + +// Configured reports whether the configuration names a FreeToken. +func (f *FreeToken) Configured() bool { return f.cfg.Binary != "" } + +func (f *FreeToken) binary() string { + if p := f.installed.Load(); p != nil { + return *p + } + return f.cfg.Binary +} + +// HelpCommand is ft serve's own help. +func (f *FreeToken) HelpCommand(Command) (string, []string) { + return f.binary(), []string{"serve", "--help"} +} + +// Probe reports the FreeToken a configuration names, by what it says its +// version is. +func (f *FreeToken) Probe(ctx context.Context) (string, error) { + bin := f.binary() + if bin == "" { + return "", fmt.Errorf("no FreeToken is configured") + } + out, err := exec.CommandContext(ctx, bin, "--version").Output() //nolint:gosec // the configured engine + if err != nil { + return "", fmt.Errorf("%s does not run: %w", bin, err) + } + // "freetoken version 0.2.1" + if fs := strings.Fields(string(out)); len(fs) > 0 { + return fs[len(fs)-1], nil + } + return Unknown, nil +} + +// Validate checks the profile: one card of 8.0 or newer, on one node. +// FreeToken has no tensor parallel; it picks one device. +func (*FreeToken) Validate(p Profile, gpus []GPUInfo) error { + if p.Nodes > 1 || p.TP > 1 || p.PP > 1 || p.GPUsPerNode > 1 { + return fmt.Errorf("profile %q: freetoken runs on one GPU of one node", p.Name) + } + if err := requireComputeCapability(p, gpus, 8.0, "freetoken"); err != nil { + return err + } + return validateCommon(p, gpus) +} + +// Commands is ft serve. The card is the shard's, by CUDA_VISIBLE_DEVICES, +// so it is device 0 to FreeToken; its key goes in the environment, where it +// reads it from, rather than on a command line every user can read. +func (f *FreeToken) Commands(sh ShardSpec) ([]Command, error) { + bin := f.binary() + if bin == "" { + return nil, fmt.Errorf("freetoken: no FreeToken is installed or configured (engines.freetoken.binary)") + } + toolkit := freetokenToolkit() + if toolkit == "" { + return nil, fmt.Errorf("freetoken compiles its kernels on first start and this node has no CUDA 13 toolkit " + + "(nvcc in CUDA_HOME, /usr/local/cuda or /usr/local/cuda-13.*)") + } + p := sh.Profile + args := []string{"serve", + "--model", p.ModelPath, + "--host", sh.Host(), + "--port", strconv.Itoa(sh.HTTPPort), + "--gpu", "0", + } + if p.ServedModelName != "" { + args = append(args, "--served-model-name", p.ServedModelName) + } + args = append(args, p.Args...) + env := baseEnv(sh) + env["CUDA_HOME"] = toolkit + // nvcc for the kernels it compiles, and the environment's own bin + // first, where ninja is. + env["PATH"] = strings.Join([]string{filepath.Dir(bin), filepath.Join(toolkit, "bin"), os.Getenv("PATH")}, ":") + if sh.APIKey != "" { + env["FREETOKEN_API_KEY"] = sh.APIKey + } + return []Command{{Name: "freetoken", Path: bin, Args: args, Env: materializeEnv(env)}}, nil +} diff --git a/internal/agent/engines/freetoken_test.go b/internal/agent/engines/freetoken_test.go new file mode 100644 index 0000000..c856ecd --- /dev/null +++ b/internal/agent/engines/freetoken_test.go @@ -0,0 +1,112 @@ +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) + } + } +} diff --git a/internal/agent/engines/unsafeenv.go b/internal/agent/engines/unsafeenv.go index cab4e62..d7ae8ef 100644 --- a/internal/agent/engines/unsafeenv.go +++ b/internal/agent/engines/unsafeenv.go @@ -25,7 +25,7 @@ var envFamilies = []string{ "CUDA_", "NCCL_", "NVIDIA_", "NVTE_", "GLOO_", "UCX_", "TORCH_", "PYTORCH_", "VLLM_", "SGLANG_", "GGML_", "LLAMA_", "OMP_", "MKL_", "OPENBLAS_", "KMP_", "HF_", "TRANSFORMERS_", "TOKENIZERS_", "TRITON_", "FLASHINFER_", - "TRTLLM_", "TLLM_", + "TRTLLM_", "TLLM_", "FREETOKEN_", } // envExact are single variables that are tuning too. @@ -33,6 +33,14 @@ var envExact = map[string]bool{ "PYTHONUNBUFFERED": true, "PYTHONHASHSEED": true, "TZ": true, } +// envNeverExact are variables in an allowed family that no profile may set, +// whose names the patterns below do not catch: what the agent sets itself, +// and FreeToken's way of loading a projector file. +var envNeverExact = map[string]string{ + "FREETOKEN_API_KEY": "is set by the agent", "FREETOKEN_HOST": "is set by the agent", + "FREETOKEN_MMPROJ": "loads a file", +} + // envNever are names in an allowed family that load code or move where the // engine reads and writes. var envNever = []string{ @@ -59,6 +67,9 @@ func UnsafeEnv(key string) (bool, string) { return true, "only the variables engines are tuned with are allowed (" + strings.Join(envFamilies, "*, ") + "*)" } + if why, ok := envNeverExact[k]; ok { + return true, "it " + why + } for _, n := range envNever { if strings.Contains(k, n) && !envExact[k] { return true, "it names a path, a plugin, a library or an address, which is a way to make the engine " + @@ -157,6 +168,16 @@ var flagsNever = map[string]map[string]string{ "--gpu-roles": "is set by the agent", "--backend": "is set by the agent", "--experimental-speed-projection": "changes what the model refuses", }, + // ft serve: where it listens, its key, its model and its card are the + // agent's, and what reads files or fetches is refused. + "freetoken": { + "--gpu": "is set by the agent", "--served-model-name": "is set by the agent", + "--mmproj": "loads a file", "--mm-encoder-weights": "loads a file", + "--allowed-local-media-path": "serves local files to requests", + "--allowed-media-domains": "lets requests have the engine fetch from those sites", + "--moe-cache-split": "reads a profile file", + "--hf-overrides": "rewrites the checkpoint's configuration, which can name code to load", + }, // trtllm-serve spells its flags with underscores, so the shared list's // hyphenated names do not reach them. "trtllm": { @@ -223,7 +244,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, "strata": true} +var argparseEngines = map[string]bool{"vllm": true, "sglang": true, "comfyui": true, "strata": true, "freetoken": true} // UnsafeArgs refuses the first of a profile's arguments an engine may not be // given, saying why. diff --git a/internal/agent/provision/freetoken_test.go b/internal/agent/provision/freetoken_test.go new file mode 100644 index 0000000..dc07c74 --- /dev/null +++ b/internal/agent/provision/freetoken_test.go @@ -0,0 +1,53 @@ +package provision + +import ( + "errors" + "slices" + "strings" + "testing" + + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery" +) + +func TestFreeTokenIsInstalledOnlyWhereItsKernelsRun(t *testing.T) { + old := freetokenToolkit + toolkit := "/usr/local/cuda-13.4" + freetokenToolkit = func() string { return toolkit } + t.Cleanup(func() { freetokenToolkit = old }) + + ampere := discovery.Host{Arch: "amd64", NVIDIADriver: "580.95.05", MinComputeCapability: 8.6} + got, err := freetokenCandidates("0.2.1", "", ampere) + if err != nil { + t.Fatal(err) + } + b := got[0] + if b.Binary != "ft" || !slices.Equal(b.Requirements, []string{"freetoken[accel]==0.2.1"}) || !b.GPUCheck || b.Python == "" { + t.Errorf("build %+v", b) + } + + // Not meant for these, and not reported as a fault: no GPU, a Turing + // card (cuda01, where 0.2.1 had no kernel image for it), an Arm node. + for name, h := range map[string]discovery.Host{ + "no gpu": {Arch: "amd64"}, + "turing": {Arch: "amd64", NVIDIADriver: "615.71.09", MinComputeCapability: 7.5}, + "aarch64": {Arch: "arm64", NVIDIADriver: "580.95.05", MinComputeCapability: 9.0}, + } { + _, err := freetokenCandidates("0.2.1", "", h) + var nh NotHere + if !errors.As(err, &nh) { + t.Errorf("%s: %v, want NotHere", name, err) + } + } + + // Meant for these, and said why not. + if _, err := freetokenCandidates("0.2.1", "", discovery.Host{Arch: "amd64", NVIDIADriver: "570.1", MinComputeCapability: 9.0}); err == nil || + !strings.Contains(err.Error(), "driver 580") { + t.Errorf("old driver: %v", err) + } + toolkit = "" + _, err = freetokenCandidates("0.2.1", "", ampere) + var nh NotHere + if err == nil || errors.As(err, &nh) || !strings.Contains(err.Error(), "CUDA 13 toolkit") { + t.Errorf("no toolkit: %v", err) + } +} diff --git a/internal/agent/provision/recipe.go b/internal/agent/provision/recipe.go index 3318138..e7f7156 100644 --- a/internal/agent/provision/recipe.go +++ b/internal/agent/provision/recipe.go @@ -126,12 +126,13 @@ type Recipe struct { // from PyPI by uv, which is itself installed like llama.cpp: a GitHub release // checked against its published checksum. var Recipes = map[string]Recipe{ - "llamacpp": {Repo: "ggml-org/llama.cpp", Candidates: llamaCppCandidates}, - "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}, + "llamacpp": {Repo: "ggml-org/llama.cpp", Candidates: llamaCppCandidates}, + "vllm": {PyPI: "vllm", Candidates: vllmCandidates}, + "sglang": {PyPI: "sglang", Candidates: sglangCandidates}, + "comfyui": {Candidates: comfyuiCandidates}, + "strata": {Repo: StrataRepo, Candidates: strataCandidates}, + "freetoken": {PyPI: "freetoken", Candidates: freetokenCandidates}, + "uv": {Repo: "astral-sh/uv", Candidates: uvCandidates}, } // Installable are the engines this agent can install, as it tells the @@ -758,3 +759,56 @@ print("strata version", "` + version + `", "engine", "` + rel.engine + `")`}, {build("cuda12"), gpuNeeds{575, 6.0}, "Strata " + version + " built for CUDA 12.9"}, }, h) } + +// --- FreeToken --- + +// FreeTokenRepo is where FreeToken is developed; its builds are on PyPI. +const FreeTokenRepo = "FlashML-org/FreeToken" + +// freetokenMinCompute is the oldest GPU FreeToken runs on. It says RTX 30 +// and newer, and its kernels agree: on a Quadro RTX 6000 (7.5) version 0.2.1 +// loaded gpt-oss-20b and failed its first prefill with no kernel image for +// the card, and in fp16 with fused experts its first request hit an illegal +// memory access. Its FP8 kernels need 8.9 and its NVFP4 ones 12.0, which +// the model's quantisation decides. +const freetokenMinCompute = 8.0 + +// freetokenToolkit finds the CUDA toolkit FreeToken compiles its kernels +// with; a variable for tests. +var freetokenToolkit = func() string { return discovery.CUDAToolkit(13) } + +// freetokenCandidates are the FreeToken builds for a machine: its PyPI +// wheels with the kernel extras (FlashInfer, SGLang's kernels), CUDA 13, +// x86-64 only. A node of older cards is not one it is meant for, and is not +// reported as a fault; one of new cards without the toolkit is. +func freetokenCandidates(version, variant string, h discovery.Host) ([]Build, error) { + if _, err := gpuArch("FreeToken", h); err != nil { + return nil, err + } + if h.Arch != "amd64" { + return nil, NotHere{"FreeToken publishes wheels for x86-64 only; this node is " + h.Arch} + } + if cc := h.MinComputeCapability; cc > 0 && cc < freetokenMinCompute { + return nil, NotHere{fmt.Sprintf("FreeToken runs on GPUs of compute capability %.1f and newer "+ + "(RTX 30, A100 and later); this node has %.1f", freetokenMinCompute, cc)} + } + // It compiles kernels the first time a model loads, with nvcc, as its + // own install guide says; without a toolkit the installation looks + // finished and the first replica fails. + if freetokenToolkit() == "" { + return nil, fmt.Errorf("FreeToken compiles its kernels where it runs and needs the CUDA 13 toolkit's nvcc " + + "(in CUDA_HOME, /usr/local/cuda or /usr/local/cuda-13.*); install cuda-toolkit-13 from NVIDIA's repository") + } + b := Build{ + Variant: "cu13", + Python: pythonForEngines, + Requirements: []string{"freetoken[accel]==" + version}, + Binary: "ft", + Packages: compilerPackages, + FreeGB: 20, + GPUCheck: true, + } + return pick("FreeToken "+version, variant, []candidate{ + {b, gpuNeeds{minDriverForCUDA13, freetokenMinCompute}, "FreeToken " + version}, + }, h) +} diff --git a/internal/agent/supervisor/health.go b/internal/agent/supervisor/health.go index 46f6288..01d1668 100644 --- a/internal/agent/supervisor/health.go +++ b/internal/agent/supervisor/health.go @@ -54,9 +54,13 @@ func NewHTTPHealth(timeout time.Duration, opts ...HTTPHealthOption) *HTTPHealth // Ready reports whether the engine at baseURL is serving model. func (h *HTTPHealth) Ready(ctx context.Context, baseURL, model, apiKey string) error { - if err := h.get(ctx, baseURL+h.healthPath, apiKey, nil); err != nil { + var body healthBody + if err := h.get(ctx, baseURL+h.healthPath, apiKey, &body); err != nil { return fmt.Errorf("health: %w", err) } + if body.notReady() { + return fmt.Errorf("health: the engine says %q", body.Status) + } if !h.requireName || model == "" { return nil } @@ -98,5 +102,29 @@ func (h *HTTPHealth) get(ctx context.Context, url, apiKey string, out any) error if out == nil { return nil } + if hb, ok := out.(*healthBody); ok { + // Most engines answer /health with nothing, or with text: only a + // JSON status says anything, and a body that is not one is not a + // failure. + _ = json.NewDecoder(io.LimitReader(resp.Body, 64<<10)).Decode(hb) + return nil + } return json.NewDecoder(resp.Body).Decode(out) } + +// healthBody is what an engine's /health may say in its body. Some answer +// 200 from the moment their HTTP server is up and say in the body that the +// model is still loading — FreeToken's {"status":"loading","phase":...} — +// and a model that lists itself in /v1/models meanwhile would be routed to. +type healthBody struct { + Status string `json:"status"` +} + +// notReady is a body that says the engine cannot serve yet, or has failed. +func (b healthBody) notReady() bool { + switch b.Status { + case "loading", "starting", "initializing", "warming_up", "error", "unhealthy": + return true + } + return false +} diff --git a/internal/agent/supervisor/health_test.go b/internal/agent/supervisor/health_test.go index 7186fe1..68b6a90 100644 --- a/internal/agent/supervisor/health_test.go +++ b/internal/agent/supervisor/health_test.go @@ -103,3 +103,29 @@ func TestTheHealthCheckSendsTheEnginesKey(t *testing.T) { t.Error("without the key the engine refused, and that read as ready") } } + +func TestNotReadyWhileTheHealthBodySaysLoading(t *testing.T) { + // FreeToken answers 200 from the moment its HTTP server is up and lists + // its model meanwhile; only the body says the weights are still loading. + body := `{"status":"loading","phase":"weights","progress":{"done_bytes":1,"total_bytes":9},"model":"glm"}` + mux := http.NewServeMux() + mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) { io.WriteString(w, body) }) + mux.HandleFunc("GET /v1/models", func(w http.ResponseWriter, _ *http.Request) { + io.WriteString(w, `{"data":[{"id":"glm"}]}`) + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + h := NewHTTPHealth(time.Second) + if err := h.Ready(t.Context(), srv.URL, "glm", ""); err == nil || !strings.Contains(err.Error(), "loading") { + t.Errorf("a loading engine was reported ready, or not said to be loading: %v", err) + } + body = `{"status":"ok","model":"glm","uptime_s":3}` + if err := h.Ready(t.Context(), srv.URL, "glm", ""); err != nil { + t.Errorf("a serving engine was reported unready: %v", err) + } + // Text, as some engines answer: not a status, and not a failure. + body = "OK" + if err := h.Ready(t.Context(), srv.URL, "glm", ""); err != nil { + t.Errorf("a plain-text health answer was taken for a failure: %v", err) + } +} diff --git a/internal/cli/enginecheck.go b/internal/cli/enginecheck.go index 461e07f..e60c380 100644 --- a/internal/cli/enginecheck.go +++ b/internal/cli/enginecheck.go @@ -72,6 +72,8 @@ func newEngineCheckCmd() *cobra.Command { cfg.VLLM.Bin = bin case "sglang": cfg.SGLang.Python = bin + case "freetoken": + cfg.FreeToken.Binary = bin case "trtllm": cfg.TRTLLM.Python = bin case "sim": diff --git a/internal/cli/tui_actions.go b/internal/cli/tui_actions.go index 2439820..c30eeb2 100644 --- a/internal/cli/tui_actions.go +++ b/internal/cli/tui_actions.go @@ -555,7 +555,7 @@ func (m consoleModel) nodeMenu(n adminapi.NodeView) *menuDialog { } items = append(items, m.locationItems(n)...) if n.GPUTotal > 0 { - for _, engine := range []string{"vllm", "sglang"} { + for _, engine := range []string{"vllm", "sglang", "freetoken"} { items = append(items, menuItem{engine + " build…", func(m *consoleModel) tea.Cmd { m.openForm(&formDialog{title: engine + " on " + n.ID, note: "Which build of " + engine + " this node installs: auto follows its hardware, none installs none.", diff --git a/internal/config/config.go b/internal/config/config.go index 6d129d7..506feed 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -1069,12 +1069,18 @@ func (s Server) ControlPlaneMode() bool { // DefaultEngines are the engine versions this release was checked against: // every flag llmgw builds for them is one they take. // -// vLLM is installed on every node with an NVIDIA GPU it can run on and -// skipped elsewhere. SGLang is not here: it is opt-in, by naming a version -// under engines. +// vLLM and FreeToken are installed on every node with an NVIDIA GPU they +// can run on and skipped elsewhere — FreeToken on cards of 8.0 and newer +// with the CUDA 13 toolkit — so a node that can serve a model too large for +// its cards has the engine for it before anyone asks. SGLang is not here: +// it is opt-in, by naming a version under engines. +// +// FreeToken is followed release by release: a new one on PyPI is moved to +// here once ft serve has been checked to take the flags llmgw gives it. var DefaultEngines = map[string]string{ - "llamacpp": "b10991", - "vllm": "0.30.0", + "llamacpp": "b10991", + "vllm": "0.30.0", + "freetoken": "0.2.1", } // EngineOverride gives the nodes it picks different engines. Every selector diff --git a/internal/control/enginechoice.go b/internal/control/enginechoice.go index 4cbcbb8..d7481b6 100644 --- a/internal/control/enginechoice.go +++ b/internal/control/enginechoice.go @@ -96,7 +96,7 @@ func EngineBuilds(engine string) []string { switch engine { case "vllm": return append(append([]string{"auto"}, provision.VLLMVariants...), "none") - case "sglang": + case "sglang", "freetoken": return []string{"auto", "none"} } return nil diff --git a/internal/control/modelserve/context.go b/internal/control/modelserve/context.go index 17fd1bd..6b30d1f 100644 --- a/internal/control/modelserve/context.go +++ b/internal/control/modelserve/context.go @@ -174,6 +174,13 @@ func resizeProfile(p catalog.Profile, req ContextRequest, sh shape) (catalog.Pro case "sglang": pool = req.Tokens p.Args = setArg(p.Args, "--context-length", strconv.Itoa(pool)) + case "freetoken": + // Its experts live in RAM and stream to the card, so the weights + // are not what has to fit there, and the card check below would + // refuse what it serves. + p.Args = setArg(p.Args, "--max-seq-len-override", strconv.Itoa(req.Tokens)) + change.Args = p.Args + return p, change, nil default: return p, change, fmt.Errorf("a %s profile (%s) does not take a context this way", p.Engine, p.Name) } diff --git a/internal/control/modelserve/gpus.go b/internal/control/modelserve/gpus.go index 70e90f3..e353b04 100644 --- a/internal/control/modelserve/gpus.go +++ b/internal/control/modelserve/gpus.go @@ -208,6 +208,10 @@ func contextOf(m catalog.Model, p catalog.Profile) int { if n := argInt(p.Args, "--context-length"); n > 0 { return n } + case "freetoken": + if n := argInt(p.Args, "--max-seq-len-override"); n > 0 { + return n + } } return m.ContextLength } diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index 93e20db..19d5a75 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -157,6 +157,9 @@ var Reference = []Block{ {Package: "config", Name: "FirewallSettings", Title: "FirewallSettings", Doc: "", Fields: []Field{ {Key: "enabled", Type: "bool", Doc: "Enabled keeps llmgw's internal ports — engines, node exporter, the\nagents' stream, the mirror, the store — to the cluster's networks on\nevery machine. False removes llmgw's nftables table everywhere, for a\nsite whose own firewall does that job; the rules are kept, and true\nputs them back. Changed in the cluster's settings, it is in force on\nevery node and control plane within seconds.", Nested: ""}, }}, + {Package: "engines", Name: "FreeTokenConfig", Title: "FreeTokenConfig", Doc: "", Fields: []Field{ + {Key: "binary", Type: "string", Doc: "Binary is its ft command, in the environment it was installed into.\nThe agent installs one itself on a node whose GPUs it runs on; this\nis for one installed by hand.", Nested: ""}, + }}, {Package: "catalog", Name: "GPUNeed", Title: "GPUNeed", Doc: "", Fields: []Field{ {Key: "min_compute_capability", Type: "float64", Doc: "MinComputeCapability is the oldest architecture that runs it: 7.5 is\nTuring, 8.0 Ampere (bf16, flash-attention 2), 8.9 Ada (fp8), 9.0 Hopper.", Nested: ""}, {Key: "min_memory_gb", Type: "int", Doc: "MinMemoryGB is the card's own memory it needs.", Nested: ""}, @@ -622,6 +625,7 @@ var Reference = []Block{ {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"}, + {Key: "freetoken", Type: "FreeTokenConfig", Doc: "FreeToken serves mixture-of-experts models larger than a card from\nRAM and one card of 8.0 or newer; installed by the agent where the\ncards allow it, or named here.", Nested: "engines.FreeTokenConfig"}, }}, {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 98edc8a..18afb32 100644 --- a/internal/docgen/render.go +++ b/internal/docgen/render.go @@ -24,6 +24,7 @@ var Sources = []string{ "internal/agent/engines/sim.go", "internal/agent/engines/comfyui.go", "internal/agent/engines/strata.go", + "internal/agent/engines/freetoken.go", "internal/agent/supervisor/supervisor.go", "internal/tlsconf/tlsconf.go", "internal/control/catalog/catalog.go", diff --git a/internal/web/admin_serve.go b/internal/web/admin_serve.go index 9d2c636..f703dab 100644 --- a/internal/web/admin_serve.go +++ b/internal/web/admin_serve.go @@ -606,7 +606,8 @@ func describeServe(l *Lang, m catalog.Model) string { return "" } p := m.Profiles[0] - engine := map[string]string{"llamacpp": "llama.cpp", "vllm": "vLLM", "sglang": "SGLang", "trtllm": "TensorRT-LLM"}[p.Engine] + engine := map[string]string{"llamacpp": "llama.cpp", "vllm": "vLLM", "sglang": "SGLang", "trtllm": "TensorRT-LLM", + "freetoken": "FreeToken"}[p.Engine] if engine == "" { engine = p.Engine } diff --git a/scripts/engine_releases.py b/scripts/engine_releases.py new file mode 100755 index 0000000..0107c37 --- /dev/null +++ b/scripts/engine_releases.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +"""Say which followed engines have published a version newer than the one llmgw installs. + +The versions llmgw installs are DefaultEngines in internal/config/config.go. An +engine followed here is one whose releases are taken up as they come: each new +one is checked against the flags llmgw gives it (llmgw ctl engine-check, and a +replica on a card it runs on) and then moved to in DefaultEngines. + +Exits 1 when one is behind, so the scheduled workflow's failed run is the +notice; 0 when all are current. No token needed: PyPI's JSON API is public. + +Usage: + scripts/engine_releases.py [--config internal/config/config.go] +""" + +import argparse +import json +import re +import sys +import urllib.request + +# Engine name in DefaultEngines -> (PyPI project, where its changes are read). +FOLLOWED = { + "freetoken": ("freetoken", "https://github.com/FlashML-org/FreeToken/releases"), +} + + +def pinned(config_path): + """The DefaultEngines map, as written in the Go source.""" + src = open(config_path, encoding="utf-8").read() + block = re.search(r"var DefaultEngines = map\[string\]string\{(.*?)\n\}", src, re.S) + if not block: + sys.exit(f"no DefaultEngines in {config_path}") + return dict(re.findall(r'"([^"]+)":\s*"([^"]+)"', block.group(1))) + + +def key(version): + """A version as numbers to compare; a pre-release sorts before its release.""" + nums = [int(n) for n in re.findall(r"\d+", version.split("+")[0])] + pre = bool(re.search(r"(a|b|rc|dev)\d*$", version)) + return (nums, 0 if pre else 1) + + +def latest(project): + """The newest version PyPI has of a project that is not yanked or a pre-release.""" + with urllib.request.urlopen(f"https://pypi.org/pypi/{project}/json", timeout=30) as r: + data = json.load(r) + versions = [ + v for v, files in data["releases"].items() + if files and not all(f.get("yanked") for f in files) and not re.search(r"(a|b|rc|dev)\d*$", v) + ] + return max(versions, key=key) if versions else data["info"]["version"] + + +def main(): + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--config", default="internal/config/config.go") + args = ap.parse_args() + pins = pinned(args.config) + behind = False + for engine, (project, changes) in FOLLOWED.items(): + have = pins.get(engine) + if have is None: + print(f"{engine}: not in DefaultEngines") + behind = True + continue + new = latest(project) + if key(new) > key(have): + print(f"{engine}: {new} is out, llmgw installs {have} — {changes}") + behind = True + else: + print(f"{engine}: {have} is current") + return 1 if behind else 0 + + +if __name__ == "__main__": + sys.exit(main())