llmgw/internal/bench/eval/service.go

package eval

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io/fs"
	"log/slog"
	"os"
	"path/filepath"
	"slices"
	"sort"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// The control plane's side: suites found on its disk or carried in the
// binary, runs started and followed, and kept in the store.

// Store keeps finished runs.
type Store interface {
	PutEvalRun(ctx context.Context, r store.EvalRun) error
	EvalRuns(ctx context.Context, model string, limit int) ([]store.EvalRun, error)
	EvalRunByID(ctx context.Context, id string) (store.EvalRun, error)
}

// Deps are what the service runs suites with.
type Deps struct {
	// Call asks a model: the router's Replay.
	Call Caller
	// Store keeps runs; without one they are kept here until a restart.
	Store Store
	// Dir is where suites are read from by name: control.eval_dir.
	Dir string
	// Builtin are the suites the binary carries, found when Dir has none
	// of the name.
	Builtin fs.FS
	// Origin is where a catalog model's weights came from, when known.
	Origin func(model string) (repo, revision string)
	// CanaryOf is the model a name is the canary of, or "": a run on a
	// canary is filed under the model it is a canary of.
	CanaryOf func(name string) string
	Log      *slog.Logger
}

// Request is a suite to run against a model: one on the control plane, by
// its name, or one sent whole.
type Request struct {
	// Suite is a suite's name in the eval directory, or carried in the
	// binary (arithmetic, json-format).
	Suite string `json:"suite,omitempty"`
	// Content is a suite's text instead, and File the name of the file it
	// was read from, which names the suite unless it names itself.
	Content string `json:"content,omitempty"`
	File    string `json:"file,omitempty"`
	// Concurrency, Judge and MinScore are as Options has them.
	Concurrency int     `json:"concurrency,omitempty"`
	Judge       string  `json:"judge,omitempty"`
	MinScore    float64 `json:"min_score,omitempty"`
}

// Named says which suite a request is, for a step's detail.
func (r Request) Named() string {
	if r.Suite != "" {
		return r.Suite
	}
	return r.File
}

// The states a run can be in.
const (
	Running = "running"
	Done    = "done"
	Failed  = "failed"
)

// Live is a run as it stands: one running, with how many of its cases are
// done, or one finished.
type Live struct {
	store.EvalRun
	State string `json:"state"`
	Done  int    `json:"done"`
}

// SuiteInfo is a suite the control plane can run by name.
type SuiteInfo struct {
	Name string `json:"name"`
	// From is the file it is read from, or "built in".
	From        string  `json:"from"`
	Description string  `json:"description,omitempty"`
	Cases       int     `json:"cases"`
	Threshold   float64 `json:"threshold"`
	Judge       bool    `json:"judge,omitempty"`
	Hash        string  `json:"hash,omitempty"`
	// Error is why the file cannot be run.
	Error string `json:"error,omitempty"`
}

// Service runs suites and remembers them.
type Service struct {
	deps Deps
	mu   sync.Mutex
	live map[string]*Live
	// order is the live runs oldest first, so the finished ones past
	// keepLive are forgotten — the store has them.
	order []string
}

const keepLive = 50

// NewService builds a service.
func NewService(d Deps) *Service {
	if d.Log == nil {
		d.Log = slog.Default()
	}
	return &Service{deps: d, live: map[string]*Live{}}
}

var suiteExts = []string{".yaml", ".yml", ".jsonl", ".json"}

// Load reads the suite a request names or carries.
func (s *Service) Load(r Request) (Suite, error) {
	if r.Content != "" {
		file := r.File
		if file == "" {
			file = "suite"
		}
		return Parse(file, []byte(r.Content))
	}
	name := strings.TrimSpace(r.Suite)
	if name == "" {
		return Suite{}, errors.New("no suite: name one on the control plane, or send one's text")
	}
	// A name, not a path: this reads the control plane's disk on behalf of
	// whoever holds the admin API, and only inside the eval directory.
	if strings.ContainsAny(name, `/\`) || name == "." || name == ".." {
		return Suite{}, fmt.Errorf("%q is not a suite's name: the control plane reads suites from %s by name; "+
			"send a file from elsewhere with its text", name, s.dirOrNone())
	}
	tries := []string{name}
	if !slices.Contains(suiteExts, filepath.Ext(name)) {
		tries = nil
		for _, e := range suiteExts {
			tries = append(tries, name+e)
		}
	}
	if s.deps.Dir != "" {
		for _, t := range tries {
			b, err := os.ReadFile(filepath.Join(s.deps.Dir, t))
			if err == nil {
				return Parse(t, b)
			}
			if !errors.Is(err, fs.ErrNotExist) {
				return Suite{}, err
			}
		}
	}
	if s.deps.Builtin != nil {
		for _, t := range tries {
			if b, err := fs.ReadFile(s.deps.Builtin, t); err == nil {
				return Parse(t, b)
			}
		}
	}
	return Suite{}, fmt.Errorf("no suite called %s in %s or built in: `llmgw ctl eval --suites` lists them", name, s.dirOrNone())
}

func (s *Service) dirOrNone() string {
	if s.deps.Dir == "" {
		return "no eval directory"
	}
	return s.deps.Dir
}

// Suites are those the control plane can run by name: its directory's, and
// the binary's that the directory does not replace.
func (s *Service) Suites() []SuiteInfo {
	var out []SuiteInfo
	seen := map[string]bool{}
	add := func(fsys fs.FS, from func(string) string) {
		entries, err := fs.ReadDir(fsys, ".")
		if err != nil {
			return
		}
		for _, e := range entries {
			ext := filepath.Ext(e.Name())
			if e.IsDir() || !slices.Contains(suiteExts, ext) {
				continue
			}
			name := strings.TrimSuffix(e.Name(), ext)
			if seen[name] {
				continue
			}
			seen[name] = true
			info := SuiteInfo{Name: name, From: from(e.Name())}
			b, err := fs.ReadFile(fsys, e.Name())
			var su Suite
			if err == nil {
				su, err = Parse(e.Name(), b)
			}
			if err != nil {
				info.Error = err.Error()
			} else {
				info.Description, info.Cases, info.Threshold = su.Description, len(su.Cases), su.Threshold
				info.Judge, info.Hash = su.NeedsJudge(), su.Hash
			}
			out = append(out, info)
		}
	}
	if s.deps.Dir != "" {
		add(os.DirFS(s.deps.Dir), func(n string) string { return filepath.Join(s.deps.Dir, n) })
	}
	if s.deps.Builtin != nil {
		add(s.deps.Builtin, func(string) string { return "built in" })
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name })
	return out
}

func newID() string {
	b := make([]byte, 4)
	_, _ = rand.Read(b)
	return fmt.Sprintf("eval-%s-%s", time.Now().UTC().Format("20060102-150405"), hex.EncodeToString(b))
}

// prepare reads the suite and checks it can be run against model, before
// anything is asked.
func (s *Service) prepare(model string, r Request) (Suite, error) {
	if s.deps.Call == nil {
		return Suite{}, errors.New("this control plane does not route requests, so it cannot ask a model")
	}
	if strings.TrimSpace(model) == "" {
		return Suite{}, errors.New("no model to ask")
	}
	su, err := s.Load(r)
	if err != nil {
		return Suite{}, err
	}
	if r.MinScore < 0 || r.MinScore > 1 {
		return Suite{}, fmt.Errorf("the minimum score %v is not 0 to 1", r.MinScore)
	}
	if su.NeedsJudge() && r.Judge == "" && su.Judge == "" {
		for _, c := range su.Cases {
			if c.Judge != "" && c.JudgeModel == "" {
				return Suite{}, fmt.Errorf("case %s of %s is graded by a judge and the suite names none: give one "+
					"with --judge MODEL", c.ID, su.Name)
			}
		}
	}
	return su, nil
}

// Check says whether a suite can be run against a model, without asking it
// anything: for a gate, refused before its job starts rather than after the
// new weights are fetched.
func (s *Service) Check(model string, r Request) error {
	_, err := s.prepare(model, r)
	return err
}

// file fills in what a run is filed under: the model it is about, the one
// that answered, and where its weights came from.
func (s *Service) file(run *store.EvalRun, asked string) {
	run.Model = asked
	if s.deps.CanaryOf != nil {
		if parent := s.deps.CanaryOf(asked); parent != "" {
			run.Model, run.Target = parent, asked
		}
	}
	if s.deps.Origin != nil {
		run.Repo, run.Revision = s.deps.Origin(asked)
	}
}

// Start runs a suite against a model in the background and returns at once;
// Get follows it.
func (s *Service) Start(model string, r Request, by string) (Live, error) {
	su, err := s.prepare(model, r)
	if err != nil {
		return Live{}, err
	}
	l := &Live{State: Running}
	l.ID, l.Suite, l.SuiteHash, l.Cases, l.By, l.At = newID(), su.Name, su.Hash, len(su.Cases), by, time.Now().UTC()
	s.file(&l.EvalRun, model)
	s.mu.Lock()
	s.keep(l)
	s.mu.Unlock()
	go func() {
		_, _ = s.run(context.Background(), l, model, su, r)
	}()
	return s.snapshot(l), nil
}

// Run runs a suite against a model and waits for it: for a gate, which
// cannot go on until it knows. The run is kept like any other; job is the
// update or promotion it gates.
func (s *Service) Run(ctx context.Context, model string, r Request, by, job string) (store.EvalRun, error) {
	su, err := s.prepare(model, r)
	if err != nil {
		return store.EvalRun{}, err
	}
	l := &Live{State: Running}
	l.ID, l.Suite, l.SuiteHash, l.Cases, l.By, l.Job, l.At = newID(), su.Name, su.Hash, len(su.Cases), by, job, time.Now().UTC()
	s.file(&l.EvalRun, model)
	s.mu.Lock()
	s.keep(l)
	s.mu.Unlock()
	return s.run(ctx, l, model, su, r)
}

func (s *Service) run(ctx context.Context, l *Live, model string, su Suite, r Request) (store.EvalRun, error) {
	run, err := Evaluate(ctx, s.deps.Call, su, Options{Model: model, Concurrency: r.Concurrency, Judge: r.Judge,
		MinScore: r.MinScore, Progress: func(done, _ int) {
			s.mu.Lock()
			l.Done = done
			s.mu.Unlock()
		}})
	s.mu.Lock()
	if err != nil {
		l.State, l.Error = Failed, err.Error()
		out := l.EvalRun
		s.mu.Unlock()
		return out, err
	}
	run.ID, run.By, run.Job = l.ID, l.By, l.Job
	run.Model, run.Target, run.Repo, run.Revision = l.Model, l.Target, l.Repo, l.Revision
	s.mu.Unlock()

	// Kept before it is said to be done, so a history read in between does
	// not find it in neither place.
	if s.deps.Store != nil {
		// Its own context: a gate's caller giving up does not unkeep a run
		// that finished.
		sctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		if err := s.deps.Store.PutEvalRun(sctx, run); err != nil {
			s.deps.Log.Warn("could not keep an eval run", "run", run.ID, "model", run.Model, "err", err)
		}
	}
	s.mu.Lock()
	l.EvalRun, l.State, l.Done = run, Done, run.Cases
	s.mu.Unlock()
	s.deps.Log.Info("eval run finished", "run", run.ID, "model", run.Model, "target", run.Target,
		"suite", run.Suite, "score", fmt.Sprintf("%.3f", run.Score), "threshold", run.Threshold, "passed", run.Passed)
	return run, nil
}

// keep remembers a run, and forgets the oldest finished ones past keepLive.
// Called with mu held.
func (s *Service) keep(l *Live) {
	s.live[l.ID] = l
	s.order = append(s.order, l.ID)
	for len(s.order) > keepLive {
		oldest := s.live[s.order[0]]
		if oldest != nil && oldest.State == Running {
			break
		}
		delete(s.live, s.order[0])
		s.order = s.order[1:]
	}
}

func (s *Service) snapshot(l *Live) Live {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := *l
	out.Results = slices.Clone(l.Results)
	return out
}

// Get is a run by its id: one running or finished here, else one kept.
func (s *Service) Get(ctx context.Context, id string) (Live, error) {
	s.mu.Lock()
	l, ok := s.live[id]
	s.mu.Unlock()
	if ok {
		return s.snapshot(l), nil
	}
	if s.deps.Store == nil {
		return Live{}, store.ErrNotFound
	}
	r, err := s.deps.Store.EvalRunByID(ctx, id)
	if err != nil {
		return Live{}, err
	}
	return Live{EvalRun: r, State: Done, Done: r.Cases}, nil
}

// History is a model's runs, newest first, without their cases: those
// running here, then those kept.
func (s *Service) History(ctx context.Context, model string, limit int) ([]Live, error) {
	if limit <= 0 {
		limit = 20
	}
	var out []Live
	seen := map[string]bool{}
	s.mu.Lock()
	for i := len(s.order) - 1; i >= 0; i-- {
		l := s.live[s.order[i]]
		if l == nil || (model != "" && l.Model != model && l.Target != model) {
			continue
		}
		// Finished ones are listed from the store, when there is one, in
		// its order.
		if l.State != Running && s.deps.Store != nil {
			continue
		}
		v := *l
		v.Results = nil
		out = append(out, v)
		seen[l.ID] = true
	}
	s.mu.Unlock()
	if s.deps.Store != nil {
		runs, err := s.deps.Store.EvalRuns(ctx, model, limit)
		if err != nil {
			return out, err
		}
		for _, r := range runs {
			if !seen[r.ID] {
				out = append(out, Live{EvalRun: r, State: Done, Done: r.Cases})
			}
		}
	}
	if len(out) > limit {
		out = out[:limit]
	}
	return out, nil
}