llmgw/internal/agent/engines/freetoken.go

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
}