llmgw/internal/cli/judgekeeper.go

package cli

import (
	"context"
	"log/slog"
	"slices"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelserve"
	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
)

// judgeKeeper fetches and serves, on the leader, each judge llmgw knows that
// a policy names and the catalog does not have: naming qwen3guard-4b in a
// policy is all it takes to have one, and a judge removed by mistake comes
// back. A judge it does not know is left to llmgw ctl guard judge enable
// --repo, said in guard judge status.
type judgeKeeper struct {
	library func() guard.Config
	models  func() []catalog.Model
	serve   interface {
		Start(modelserve.Request) (modelserve.Job, error)
		Jobs() []modelserve.Job
	}
	leader func() bool
	gpus   func() []adminapi.ClassUse
	log    *slog.Logger
	now    func() time.Time

	// tried is when each judge was last asked for, so one that cannot be
	// fetched is not asked for again on every look.
	tried map[string]time.Time
	// grace is how long after its first look the keeper serves no judge
	// itself: init adds its own judge just after the control plane starts,
	// sized and placed as the install says, and one the leader added first
	// was what init then kept. Zero waits for nothing.
	grace time.Duration
	first time.Time

	// What refitting an earlier install's judge needs: the catalog to change
	// its profile in, the replicas to restart, and the nodes' CPUs. Any nil
	// leaves judges as they are.
	profiles interface {
		SetProfiles(ctx context.Context, name string, profiles []catalog.Profile, by string) (catalog.Model, error)
	}
	restart  func(ids ...string) error
	replicas func(model string) []string
	nodeCPUs func() []int
	// looked is each judge already refitted or found changed by hand, so
	// each is said once.
	looked map[string]bool
}

// judgeRetry is how long a judge whose serving failed waits before it is
// asked for again.
const judgeRetry = 10 * time.Minute

// judgeGrace is the keeper's grace on a control plane just started.
const judgeGrace = 10 * time.Minute

func (k *judgeKeeper) run(ctx context.Context, every time.Duration) {
	t := time.NewTicker(every)
	defer t.Stop()
	for {
		k.check()
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		}
	}
}

func (k *judgeKeeper) check() {
	if k.leader != nil && !k.leader() {
		return
	}
	lib := k.library()
	if !lib.On() {
		return
	}
	if k.tried == nil {
		k.tried = map[string]time.Time{}
	}
	now := time.Now
	if k.now != nil {
		now = k.now
	}
	have := map[string]bool{}
	for _, m := range k.models() {
		have[m.Name] = true
		k.refit(m)
	}
	if k.first.IsZero() {
		k.first = now()
	}
	if now().Sub(k.first) < k.grace {
		return
	}
	busy := map[string]bool{}
	for _, j := range k.serve.Jobs() {
		if j.State == modelserve.Running {
			busy[j.Model] = true
			busy[j.Request.Name] = true
		}
	}
	var wanted []string
	for _, p := range lib.Policies {
		for _, j := range p.Judge.Judges() {
			if !slices.Contains(wanted, j) {
				wanted = append(wanted, j)
			}
		}
	}
	slices.Sort(wanted)
	for _, name := range wanted {
		kj, known := knownJudgeNamed(name)
		if have[name] || !known || busy[name] {
			continue
		}
		if at, ok := k.tried[name]; ok && now().Sub(at) < judgeRetry {
			continue
		}
		k.tried[name] = now()
		var opts judgeServeOptions
		if !kj.GPU {
			// The cores init would give the same number of judges.
			var cpus []int
			if k.nodeCPUs != nil {
				cpus = k.nodeCPUs()
			}
			opts.Replicas = 2
			opts.Cores = refitCores(opts.Replicas, cpus)
		}
		if kj.GPU {
			if opts.GPUClass = roomiestGPUClass(k.gpus()); opts.GPUClass == "" {
				k.log.Warn("a policy names a judge that runs on a GPU, and this cluster has none",
					"judge", name, "instead", defaultJudge)
				continue
			}
		}
		req := kj.serveRequest(opts)
		req.By = "llmgw (a guard policy names it)"
		// init, or an operator, may add it while this fetches: theirs stays.
		req.IfAbsent = true
		job, err := k.serve.Start(req)
		if err != nil {
			k.log.Error("could not start serving a judge a policy names", "judge", name, "err", err)
			continue
		}
		k.log.Info("serving a judge a policy names, which the catalog did not have",
			"judge", name, "repo", kj.Repo, "file", kj.File, "job", job.ID)
	}
}

// refit moves a judge an earlier `server init` installed to what init
// installs now, and restarts its replicas onto it. Those judges reserved no
// cores or one, 2 or 4 GB, and ran llama.cpp's prompt cache: on fewer than
// eight cores every judgment missed its 5 s budget and fail: closed refused
// every request, until fixed by hand. A judge whose profile differs in any
// other way was set by somebody, and is left as it is, said once.
func (k *judgeKeeper) refit(m catalog.Model) {
	if k.profiles == nil || k.restart == nil || k.replicas == nil || m.Name != defaultJudge || !m.IsGuard() {
		return
	}
	if k.looked == nil {
		k.looked = map[string]bool{}
	}
	if k.looked[m.Name] {
		return
	}
	k.looked[m.Name] = true
	if !earlierInitJudge(m) {
		if len(m.Profiles) != 1 || !slices.Contains(m.Profiles[0].Args, "--cache-ram") {
			k.log.Info("the judge's profile is not one init wrote; left as it is",
				"judge", m.Name, "see", "llmgw ctl model resources "+m.Name)
		}
		return
	}
	var cpus []int
	if k.nodeCPUs != nil {
		cpus = k.nodeCPUs()
	}
	p := m.Profiles[0]
	p.MinMemoryGB = judgeMemoryGB
	p.CPUCores = refitCores(m.Scaling.Min, cpus)
	p.Args = []string{"--ctx-size", "8192", "--cache-ram", "0"}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	if _, err := k.profiles.SetProfiles(ctx, m.Name, []catalog.Profile{p}, "llmgw (the judge an earlier init installed)"); err != nil {
		k.log.Warn("could not refit the judge an earlier init installed", "judge", m.Name, "err", err)
		k.looked[m.Name] = false
		return
	}
	ids := k.replicas(m.Name)
	if len(ids) > 0 {
		if err := k.restart(ids...); err != nil {
			k.log.Warn("refitted the judge, but could not restart it", "judge", m.Name, "err", err)
		}
	}
	k.log.Info("refitted the judge an earlier init installed", "judge", m.Name,
		"memory_gb", p.MinMemoryGB, "cpu_cores", p.CPUCores, "replicas_restarted", len(ids))
}

// earlierInitJudge is whether a judge's entry is exactly what `server init`
// before v0.14.7 wrote: one llama.cpp profile on CPUs, an 8192 context and
// no other args, 2 or 4 GB, none to two cores.
func earlierInitJudge(m catalog.Model) bool {
	if len(m.Profiles) != 1 {
		return false
	}
	p := m.Profiles[0]
	return p.Name == "cpu" && p.Engine == "llamacpp" && p.GPUClass == "" && p.GPU == nil &&
		len(p.NodeSelector) == 0 && slices.Equal(p.Args, []string{"--ctx-size", "8192"}) &&
		(p.MinMemoryGB == 2 || p.MinMemoryGB == 4) && p.CPUCores >= 0 && p.CPUCores <= 2
}

// refitCores is a judge replica's cores by init's rule: up to four for the
// one judge of a cluster, two each where there are more, fewer where the
// smallest node cannot hold them all.
func refitCores(replicas int, cpus []int) int {
	smallest := 0
	for _, c := range cpus {
		if c > 0 && (smallest == 0 || c < smallest) {
			smallest = c
		}
	}
	if replicas <= 1 {
		if smallest == 0 {
			return judgeCores
		}
		return min(smallest, judgeCoresAlone)
	}
	cores := judgeCores
	if len(cpus) == 1 && smallest > 0 && replicas*cores > smallest {
		cores = max(1, smallest/replicas)
	}
	return cores
}