llmgw/internal/agent/engines/strata.go

package engines

import (
	"context"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"sync/atomic"

	"gitea.mynodes.xyz/ilker/llmgw/internal/strata"
)

// 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 unless asked for this format.
func (*Strata) MetricsPath() string { return "/metrics?format=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")
}

var strataVersion = regexp.MustCompile(`PRODUCT_VERSION\s*=\s*"([^"]+)"`)

// Probe reports the Strata a configuration names, by the version its source
// states: a directory with its setup and server, an interpreter, and the
// engine binary. Without it a Strata named by hand was never reported, and
// so never placed on.
func (s *Strata) Probe(context.Context) (string, error) {
	dir := s.dir()
	if dir == "" {
		return "", fmt.Errorf("no Strata directory is configured")
	}
	for _, f := range []string{filepath.Join(dir, "setup.py"), filepath.Join(dir, "serve", "server.py"),
		s.python(), filepath.Join(s.engineDir(), "strata")} {
		if _, err := os.Stat(f); err != nil {
			return "", fmt.Errorf("%s is not a whole Strata: %w", dir, err)
		}
	}
	b, err := os.ReadFile(filepath.Join(dir, "strata_cli", "__init__.py"))
	if err != nil {
		return Unknown, nil //nolint:nilerr // whole, of a version it does not say
	}
	if m := strataVersion.FindSubmatch(b); m != nil {
		return string(m[1]), nil
	}
	return Unknown, nil
}

// 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)
	}
	// Pascal and Volta run its CUDA 12 build, which the install picks for
	// them.
	if err := requireComputeCapability(p, gpus, 6.0, "strata"); err != nil {
		return err
	}
	// A family whose engine needs more: said now, rather than by Strata's
	// setup after the replica was placed here.
	if fam, ok := strata.FamilyNamed(strata.ArgValue(p.Args, "--family")); ok && fam.MinCC > 0 {
		if err := requireComputeCapability(p, gpus, fam.MinCC, "strata's "+fam.Title+" engine"); 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
}