Files
llmgw/internal/control/enginechoice.go
T
Ilker ManapandClaude Opus 5.5 4f280a3546
Release / build (push) Successful in 1m12s
Release / notes (push) Successful in 3s
CI / test (push) Successful in 3m55s
Give each GPU node a vLLM build it can run, chosen on the cluster page
A node that cannot run the CUDA 13 build installs the same version's
CUDA 12.9 build from vLLM's GitHub release, with PyTorch's cu129 torch;
the install computes on the GPU before it counts. The build can be
chosen per node — auto, cu13, cu129, none — on the cluster page, with
ctl engine-build, or over the admin API; the control plane keeps the
choice and tells the node at once.

Refs #2.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-01 12:39:14 +02:00

109 lines
3.2 KiB
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
}
// 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":
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
}