llmgw/internal/control/enginechoice.go

package control

import (
	"context"
	"fmt"
	"log/slog"
	"maps"
	"slices"
	"strings"
	"sync"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
)

// EngineChoiceStore is where the engine builds chosen per node are kept.
type EngineChoiceStore interface {
	NodeEngines(ctx context.Context) (map[string]map[string]string, error)
	SetNodeEngine(ctx context.Context, node, engine, variant, by string) error
}

// EngineChoices are the engine builds chosen for nodes, rather than left to
// their hardware: vLLM's CUDA 12.9 build on a node its CUDA 13 build would
// also install on, or no vLLM at all. Kept in the store, held here for the
// agent service to ask on every registration, and pushed to a node the
// moment its choice changes.
type EngineChoices struct {
	st  EngineChoiceStore
	log *slog.Logger

	mu     sync.RWMutex
	byNode map[string]map[string]string
	notify func(node string) error
}

// NewEngineChoices reads what was chosen before.
func NewEngineChoices(ctx context.Context, st EngineChoiceStore, log *slog.Logger) *EngineChoices {
	c := &EngineChoices{st: st, log: log, byNode: map[string]map[string]string{}}
	if got, err := st.NodeEngines(ctx); err != nil {
		log.Error("could not read the engine builds chosen per node; every node installs what its hardware runs", "err", err)
	} else {
		c.byNode = got
	}
	return c
}

// Reload reads the choices again, as a standby becoming the leader does, and
// tells each node whose choice changed.
func (c *EngineChoices) Reload(ctx context.Context) {
	got, err := c.st.NodeEngines(ctx)
	if err != nil {
		c.log.Error("could not read the engine builds chosen per node again; keeping the ones read before", "err", err)
		return
	}
	c.mu.Lock()
	var changed []string
	for node := range got {
		if !maps.Equal(got[node], c.byNode[node]) {
			changed = append(changed, node)
		}
	}
	for node := range c.byNode {
		if _, ok := got[node]; !ok {
			changed = append(changed, node)
		}
	}
	c.byNode = got
	notify := c.notify
	c.mu.Unlock()
	if notify == nil {
		return
	}
	for _, node := range changed {
		if err := notify(node); err != nil {
			c.log.Warn("could not tell a node its engine builds", "node", node, "err", err)
		}
	}
}

// OnChange is called with a node whose choice changed, to tell it.
func (c *EngineChoices) OnChange(f func(node string) error) {
	c.mu.Lock()
	c.notify = f
	c.mu.Unlock()
}

// For is what was chosen for one node, by engine: a build, or "none".
func (c *EngineChoices) For(node string) map[string]string {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return maps.Clone(c.byNode[node])
}

// EngineBuilds are what can be chosen for each engine that has a choice:
// "auto" is the node's hardware deciding, "none" no engine at all.
func EngineBuilds(engine string) []string {
	switch engine {
	case "vllm":
		return append(append([]string{"auto"}, provision.VLLMVariants...), "none")
	case "sglang", "freetoken":
		return []string{"auto", "none"}
	}
	return nil
}

// Set chooses an engine's build for a node, keeps it, and tells the node.
func (c *EngineChoices) Set(ctx context.Context, node, engine, build, by string) error {
	builds := EngineBuilds(engine)
	if builds == nil {
		return fmt.Errorf("%s has no builds to choose between", engine)
	}
	if !slices.Contains(builds, build) {
		return fmt.Errorf("%s has no build %q; choose %s", engine, build, strings.Join(builds, ", "))
	}
	variant := build
	if build == "auto" {
		variant = ""
	}
	if err := c.st.SetNodeEngine(ctx, node, engine, variant, by); err != nil {
		return err
	}
	c.mu.Lock()
	if variant == "" {
		delete(c.byNode[node], engine)
	} else {
		if c.byNode[node] == nil {
			c.byNode[node] = map[string]string{}
		}
		c.byNode[node][engine] = variant
	}
	notify := c.notify
	c.mu.Unlock()
	c.log.Info("engine build chosen for a node", "node", node, "engine", engine, "build", build, "by", by)
	if notify != nil {
		if err := notify(node); err != nil {
			// Kept and applied on its next connection; said, since a node that
			// is not connected is the usual reason.
			c.log.Warn("the node was not told yet; it is when it next connects", "node", node, "err", err)
		}
	}
	return nil
}