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>
109 lines
3.2 KiB
Go
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
|
|
}
|