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
}

// judgeRetry is how long a judge whose serving failed waits before it is
// asked for again.
const judgeRetry = 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
	}
	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 {
			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)"
		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)
	}
}