llmgw/internal/control/modelserve/canary.go

package modelserve

import (
	"context"
	"errors"
	"fmt"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/canary"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
)

// A canary: a model's repository's newer weights run beside it as another
// model, given a share of its requests, compared, and then promoted to its
// place or taken out — by hand, or by itself when it serves worse.

// canaryName is the candidate's name in the catalog.
func canaryName(model string) string { return model + "-canary" }

// StartCanary fetches a model's newer weights and runs them as its canary,
// with the share and limits given.
func (s *Server) StartCanary(name string, cn catalog.Canary, by string) (Job, error) {
	if s.deps.Models == nil || s.deps.Restart == nil {
		return Job{}, errors.New("a canary needs the store, which this control plane lacks")
	}
	m, err := s.current(name)
	if err != nil {
		return Job{}, err
	}
	if s.deps.Models.Declared(name) {
		return Job{}, fmt.Errorf("%s is declared in server.yaml; a canary is for a model added at runtime", name)
	}
	if m.Canary != nil {
		return Job{}, fmt.Errorf("%s already has a canary, %s at %d%%: change its share with --percent, "+
			"or promote or abort it first", name, m.Canary.Model, m.Canary.Percent)
	}
	if cn.Percent < 0 || cn.Percent > 100 {
		return Job{}, fmt.Errorf("percent %d is not 0 to 100", cn.Percent)
	}
	s.mu.Lock()
	for _, j := range s.jobs {
		if (j.Kind == "update" || j.Kind == "canary" || j.Kind == "promote") && j.Model == name && j.State == Running {
			s.mu.Unlock()
			return Job{}, fmt.Errorf("%s has a job running: %s", name, j.ID)
		}
	}
	s.seq++
	j := &Job{ID: fmt.Sprintf("canary-%d-%d", time.Now().Unix(), s.seq), Kind: "canary", Model: name,
		Request: Request{Name: name, By: by}, State: Running, Started: time.Now()}
	s.keepJob(j)
	s.mu.Unlock()
	go s.runCanary(j, cn)
	return s.snapshot(j), nil
}

func (s *Server) runCanary(j *Job, cn catalog.Canary) {
	ctx := context.Background()
	name := j.Model
	p, ok := s.prepare(ctx, j, name)
	if !ok {
		return
	}

	// 4. The candidate, beside the model: the same in everything but its
	// weights, and one replica.
	cand := canaryName(name)
	s.step(j, "catalog", Running, "adding "+cand)
	c := p.m
	c.Name, c.Aliases, c.Fallbacks, c.Canary, c.Adapters = cand, nil, nil, nil, nil
	c.Source = p.source
	c.Origin = &catalog.Origin{Repo: p.repo.ID, Revision: p.repo.Revision, Files: originFiles(p.set)}
	c.Scaling = catalog.Scaling{Min: 1, Max: 1}
	c.Description = fmt.Sprintf("Canary of %s at %s.", name, p.repo.ShortRevision())
	if err := s.deps.Models.Add(ctx, c, j.Request.By); err != nil {
		s.fail(j, "catalog", err)
		return
	}
	s.step(j, "catalog", Done, cand+" runs "+p.repo.ID+" at "+p.repo.ShortRevision())

	// 5. Serving, before it is given anyone's requests.
	r, err := s.waitReady(ctx, j, cand, false)
	if err == nil {
		s.step(j, "smoke test", Running, "asking "+r.ID+" to say hello")
		var reply string
		var tps float64
		if reply, tps, err = s.smoke(ctx, r, cand); err == nil {
			s.step(j, "smoke test", Done, fmt.Sprintf("answered at %.1f tokens/s: %q", tps, reply))
		} else {
			s.step(j, "smoke test", Failed, err.Error())
		}
	}
	if err != nil {
		_ = s.deps.Models.Remove(ctx, cand)
		s.fail(j, "placement", fmt.Errorf("%w; %s was taken out again", err, cand))
		return
	}

	// 6. Its share.
	cn.Model = cand
	if _, err := s.deps.Models.SetCanary(ctx, name, &cn, j.Request.By); err != nil {
		_ = s.deps.Models.Remove(ctx, cand)
		s.fail(j, "split", err)
		return
	}
	if s.deps.ResetCanary != nil {
		s.deps.ResetCanary(name)
	}
	s.step(j, "split", Done, fmt.Sprintf("%d%% of %s's requests go to %s", cn.Percent, name, cand))
	s.finish(j)
}

// SetCanaryPercent changes the share of a model's requests its canary is
// given; 0 keeps it running and gives it none.
func (s *Server) SetCanaryPercent(ctx context.Context, name string, percent int, by string) (string, error) {
	m, err := s.current(name)
	if err != nil {
		return "", err
	}
	if m.Canary == nil {
		return "", fmt.Errorf("%s has no canary; start one with `llmgw ctl model canary %s --percent N`", name, name)
	}
	cn := *m.Canary
	cn.Percent = percent
	if _, err := s.deps.Models.SetCanary(ctx, name, &cn, by); err != nil {
		return "", err
	}
	return fmt.Sprintf("%d%% of %s's requests go to %s", percent, name, cn.Model), nil
}

// AbortCanary takes a model's canary out: its share first, then the
// candidate.
func (s *Server) AbortCanary(ctx context.Context, name, by string) (string, error) {
	m, err := s.current(name)
	if err != nil {
		return "", err
	}
	if m.Canary == nil {
		return "", fmt.Errorf("%s has no canary", name)
	}
	cand := m.Canary.Model
	if _, err := s.deps.Models.SetCanary(ctx, name, nil, by); err != nil {
		return "", err
	}
	if err := s.deps.Models.Remove(ctx, cand); err != nil {
		return "", fmt.Errorf("%s gets no more requests, but could not be removed: %w", cand, err)
	}
	if s.deps.ResetCanary != nil {
		s.deps.ResetCanary(name)
	}
	return fmt.Sprintf("%s is out; %s serves all of %s's requests", cand, name, name), nil
}

// StartPromote moves a model to its canary's weights — the update a canary
// was the trial of — and takes the canary out.
func (s *Server) StartPromote(name, by string) (Job, error) {
	m, err := s.current(name)
	if err != nil {
		return Job{}, err
	}
	if m.Canary == nil {
		return Job{}, fmt.Errorf("%s has no canary to promote", name)
	}
	s.mu.Lock()
	s.seq++
	j := &Job{ID: fmt.Sprintf("promote-%d-%d", time.Now().Unix(), s.seq), Kind: "promote", Model: name,
		Request: Request{Name: name, By: by}, State: Running, Started: time.Now()}
	s.keepJob(j)
	s.mu.Unlock()
	go s.runPromote(j)
	return s.snapshot(j), nil
}

func (s *Server) runPromote(j *Job) {
	ctx := context.Background()
	name := j.Model
	m, err := s.current(name)
	if err != nil || m.Canary == nil {
		s.fail(j, "catalog", fmt.Errorf("%s no longer has a canary", name))
		return
	}
	c, err := s.current(m.Canary.Model)
	if err != nil {
		s.fail(j, "catalog", err)
		return
	}
	s.step(j, "catalog", Running, fmt.Sprintf("moving %s to %s's weights", name, c.Name))
	next := &catalog.Origin{}
	if c.Origin != nil {
		*next = *c.Origin
	}
	next.Previous = &catalog.PreviousOrigin{Source: m.Source}
	if m.Origin != nil {
		next.Previous.Revision, next.Previous.Files = m.Origin.Revision, m.Origin.Files
	}
	if _, err := s.deps.Models.SetSource(ctx, name, c.Source, next, j.Request.By); err != nil {
		s.fail(j, "catalog", err)
		return
	}
	// Its share back to the model, which is about to run the same weights.
	if _, err := s.deps.Models.SetCanary(ctx, name, nil, j.Request.By); err != nil {
		s.fail(j, "catalog", err)
		return
	}
	s.step(j, "catalog", Done, "the catalog names "+c.Source)

	if err := s.replace(ctx, j, name); err != nil {
		s.fail(j, "replicas", fmt.Errorf("%w. `llmgw ctl model rollback %s` goes back", err, name))
		return
	}
	if err := s.deps.Models.Remove(ctx, c.Name); err != nil {
		s.step(j, "canary", Failed, fmt.Sprintf("%s could not be removed: %v", c.Name, err))
	} else {
		s.step(j, "canary", Done, c.Name+" removed")
	}
	if s.deps.ResetCanary != nil {
		s.deps.ResetCanary(name)
	}
	if c.Origin != nil {
		s.mu.Lock()
		s.updates[name] = UpdateInfo{Model: name, Repo: c.Origin.Repo, Current: c.Origin.Revision,
			Latest: c.Origin.Revision, Previous: next.Previous.Revision, CheckedAt: time.Now()}
		s.mu.Unlock()
	}
	s.finish(j)
}

// WatchCanaries takes out, every so often, each canary serving worse than
// its model by the limits it was given. Until ctx ends.
func (s *Server) WatchCanaries(ctx context.Context, every time.Duration) {
	if s.deps.CanaryStats == nil || s.deps.Current == nil {
		return
	}
	if every <= 0 {
		every = 30 * time.Second
	}
	t := time.NewTicker(every)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			s.checkCanaries(ctx)
		}
	}
}

func (s *Server) checkCanaries(ctx context.Context) {
	stats := s.deps.CanaryStats()
	for _, m := range s.deps.Current() {
		if m.Canary == nil || m.Canary.Percent == 0 {
			continue
		}
		st, ok := stats[m.Name]
		if !ok || st.Canary != m.Canary.Model {
			continue
		}
		bad, why := canary.Verdict(st, *m.Canary)
		if !bad {
			continue
		}
		cand := m.Canary.Model
		if _, err := s.AbortCanary(ctx, m.Name, "llmgw"); err != nil {
			s.deps.Log.Error("a canary serving worse could not be taken out", "model", m.Name, "canary", cand, "err", err)
			continue
		}
		s.deps.Log.Warn("canary rolled back: it served worse than the model", "model", m.Name, "canary", cand, "why", why)
		s.deps.Notify.Notify(notify.CanaryRolledBack, m.Name+"/"+cand,
			fmt.Sprintf("the canary of %s was rolled back: %s, over %d requests", m.Name, why, st.Candidate.Requests),
			map[string]any{"model": m.Name, "canary": cand, "reason": why,
				"canary_requests": st.Candidate.Requests, "stable_requests": st.Stable.Requests})
	}
}

// CanaryView is one model's canary as it stands.
type CanaryView struct {
	Model   string         `json:"model"`
	Canary  catalog.Canary `json:"canary"`
	Serving *canary.Stat   `json:"serving,omitempty"`
}

// Canaries are the models with a canary, and how each side has served.
func (s *Server) Canaries() []CanaryView {
	if s.deps.Current == nil {
		return nil
	}
	var stats map[string]canary.Stat
	if s.deps.CanaryStats != nil {
		stats = s.deps.CanaryStats()
	}
	var out []CanaryView
	for _, m := range s.deps.Current() {
		if m.Canary == nil {
			continue
		}
		v := CanaryView{Model: m.Name, Canary: *m.Canary}
		if st, ok := stats[m.Name]; ok && st.Canary == m.Canary.Model {
			v.Serving = &st
		}
		out = append(out, v)
	}
	return out
}