llmgw/internal/strata/serve.go

// 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-<size>.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)
}