llmgw/internal/control/modelsvc/modelsvc.go

// Package modelsvc owns the catalog that is in force: the models declared in
// configuration, plus the ones added while the cluster is running.
//
// Two sources rather than one, deliberately. Configuration stays declarative —
// a model in the YAML is in the catalog on every start, whatever the database
// says — because that is how a team keeps its production set under review and
// in version control. The database holds what someone added afterwards, which
// is the case this exists for: trying a model should not mean a config change,
// a restart and an outage.
//
// Configuration wins every conflict, and the surprising direction is the
// dangerous one: a request for a reviewed production model must never be
// answered by something added from a laptop. A stored model that collides with
// a declared name is refused when it is added and skipped if it somehow
// survives into the database.
package modelsvc

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"time"

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

// Store is the persistence this needs.
type Store interface {
	PutModel(ctx context.Context, m store.StoredModel) error
	DeleteModel(ctx context.Context, name string) error
	Models(ctx context.Context) ([]store.StoredModel, error)
}

// Applier receives the rebuilt catalog. The controller implements it.
type Applier interface {
	ReplaceCatalog(ctx context.Context, c *catalog.Catalog)
}

// Service adds and removes models without a restart.
type Service struct {
	declared []catalog.Model
	store    Store
	apply    Applier
	log      *slog.Logger

	// check refuses a model that cannot run before it is added; see
	// SetCheck.
	check func(context.Context, catalog.Model) error

	// mu serialises rebuilds. They are rare — a person adding a model — and two
	// at once could apply catalogs in either order, leaving the one in force
	// disagreeing with the database that outlives it.
	mu sync.Mutex
}

// New builds the service around the models configuration declared.
func New(declared []catalog.Model, st Store, apply Applier, log *slog.Logger) *Service {
	if log == nil {
		log = slog.Default()
	}
	return &Service{declared: declared, store: st, apply: apply, log: log}
}

// Declared reports whether a name comes from configuration.
func (s *Service) Declared(name string) bool {
	for _, d := range s.declared {
		if d.Name == name {
			return true
		}
	}
	return false
}

// Sources maps each weights source in the catalog to the models that name it,
// so a page listing files can say which of them are in use and how they run.
func (s *Service) Sources(ctx context.Context) (map[string][]catalog.Model, error) {
	c, err := s.Load(ctx)
	if err != nil {
		return nil, err
	}
	out := map[string][]catalog.Model{}
	for _, m := range c.Models() {
		if m.Source != "" {
			out[m.Source] = append(out[m.Source], m)
		}
		for _, ad := range m.Adapters {
			out[ad.Source] = append(out[ad.Source], m)
		}
	}
	return out, nil
}

// Add puts a model in the catalog and keeps it across restarts.
//
// Validated by building the whole catalog before anything is written. A model
// that collides with an existing name, names a fallback that does not exist, or
// is malformed in any of the ways catalog.New refuses is rejected here rather
// than persisted and then skipped on every future start — which would be a
// model that looks added and never appears.
func (s *Service) Add(ctx context.Context, m catalog.Model, by string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.Declared(m.Name) {
		return fmt.Errorf("model %q is declared in configuration; edit that rather than adding a second definition", m.Name)
	}
	// A model with no GPU is all host memory, and one that declares none
	// fits everywhere as far as placement can tell — however many of it
	// there already are. Required of what is added now; models stored before
	// this still load, and say what they are missing.
	for _, p := range m.Profiles {
		if p.CPUOnly() && p.MinMemoryGB <= 0 {
			return fmt.Errorf("model %q profile %q runs without a GPU and declares no memory (min_memory_gb): "+
				"placement could not tell when a node is full", m.Name, p.Name)
		}
	}
	if s.check != nil {
		if err := s.check(ctx, m); err != nil {
			return err
		}
	}

	stored, err := s.storedModels(ctx)
	if err != nil {
		return err
	}
	next := append([]catalog.Model{}, s.declared...)
	for _, e := range stored {
		if e.Name == m.Name {
			continue // this call replaces it
		}
		next = append(next, e)
	}
	next = append(next, m)

	built, err := catalog.New(next)
	if err != nil {
		return fmt.Errorf("model %q would make the catalog invalid: %w", m.Name, err)
	}

	def, err := json.Marshal(m)
	if err != nil {
		return err
	}
	if err := s.store.PutModel(ctx, store.StoredModel{
		Name: m.Name, Definition: def, CreatedAt: time.Now(), CreatedBy: by,
	}); err != nil {
		return err
	}

	s.apply.ReplaceCatalog(ctx, built)
	s.log.Info("model added", "model", m.Name, "by", by)
	return nil
}

// Remove takes a model out of the catalog.
//
// A model declared in configuration cannot be removed this way: it would come
// back on the next restart, so the cluster would disagree with its own config
// file until then. Refusing and saying where to edit is the smaller surprise.
func (s *Service) Remove(ctx context.Context, name string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.Declared(name) {
		return fmt.Errorf("model %q is declared in configuration; remove it there, or it comes back on the next restart", name)
	}

	stored, err := s.storedModels(ctx)
	if err != nil {
		return err
	}
	next := append([]catalog.Model{}, s.declared...)
	found := false
	for _, e := range stored {
		if e.Name == name {
			found = true
			continue
		}
		next = append(next, e)
	}
	if !found {
		return fmt.Errorf("model %q was not added at runtime", name)
	}

	built, err := catalog.New(next)
	if err != nil {
		// Removing a model can invalidate a catalog: another may name it as a
		// fallback. Refusing beats leaving a fallback pointing at nothing,
		// which is discovered during the outage it was meant to cover.
		return fmt.Errorf("removing %q would make the catalog invalid: %w", name, err)
	}
	if err := s.store.DeleteModel(ctx, name); err != nil {
		return err
	}

	s.apply.ReplaceCatalog(ctx, built)
	s.log.Info("model removed", "model", name)
	return nil
}

// Scale changes how many replicas a model added at runtime runs.
func (s *Service) Scale(ctx context.Context, name string, sc catalog.Scaling, by string) (catalog.Model, error) {
	final, err := s.update(ctx, name, by, "change its scaling", func(m *catalog.Model) error {
		m.Scaling = sc
		return nil
	})
	if err == nil {
		s.log.Info("model scaled", "model", name,
			"min", final.Scaling.Min, "max", final.Scaling.Max, "by", by)
	}
	return final, err
}

// Resources sets what a model's profile reserves of a node: host memory, and
// the CPU cores a replica without a GPU computes on. An empty profile name
// sets every profile.
func (s *Service) Resources(ctx context.Context, name, profile string, memoryGB, cpuCores int, by string) (catalog.Model, error) {
	if memoryGB < 0 || cpuCores < 0 {
		return catalog.Model{}, errors.New("memory and cores cannot be negative")
	}
	final, err := s.update(ctx, name, by, "change its resources", func(m *catalog.Model) error {
		found := false
		for i := range m.Profiles {
			if profile != "" && m.Profiles[i].Name != profile {
				continue
			}
			found = true
			m.Profiles[i].MinMemoryGB = memoryGB
			m.Profiles[i].CPUCores = cpuCores
		}
		if !found {
			return fmt.Errorf("model %q has no profile %q", name, profile)
		}
		return nil
	})
	if err == nil {
		s.log.Info("model resources set", "model", name, "profile", profile,
			"memory_gb", memoryGB, "cpu_cores", cpuCores, "by", by)
	}
	return final, err
}

// SetClass sets a model's capacity class: critical, standard or
// opportunistic, or empty to go back to what its scaling implies.
func (s *Service) SetClass(ctx context.Context, name, class, by string) (catalog.Model, error) {
	final, err := s.update(ctx, name, by, "change its class", func(m *catalog.Model) error {
		m.Class = class
		return nil
	})
	if err == nil {
		s.log.Info("model class set", "model", name, "class", final.EffectiveClass(), "by", by)
	}
	return final, err
}

// SetOwner sets the part of the tenancy tree a model's capacity counts
// against, org:ID, team:ID or project:ID; empty makes it shared.
func (s *Service) SetOwner(ctx context.Context, name, owner, by string) (catalog.Model, error) {
	final, err := s.update(ctx, name, by, "change its owner", func(m *catalog.Model) error {
		m.Owner = owner
		return nil
	})
	if err == nil {
		s.log.Info("model owner set", "model", name, "owner", owner, "by", by)
	}
	return final, err
}

// SetAdapters replaces the LoRA adapters a model serves. Replicas load the
// new set when they next start.
func (s *Service) SetAdapters(ctx context.Context, name string, adapters []catalog.Adapter, by string) (catalog.Model, error) {
	final, err := s.update(ctx, name, by, "change its adapters", func(m *catalog.Model) error {
		m.Adapters = adapters
		return nil
	})
	if err == nil {
		names := make([]string, 0, len(adapters))
		for _, a := range adapters {
			names = append(names, a.Name)
		}
		s.log.Info("model adapters set", "model", name, "adapters", names, "by", by)
	}
	return final, err
}

// SetSource moves a model added at runtime to other weights, recording where
// they came from. Running replicas keep what they loaded until they are
// replaced; the caller restarts them.
func (s *Service) SetSource(ctx context.Context, name, source string, origin *catalog.Origin, by string) (catalog.Model, error) {
	final, err := s.update(ctx, name, by, "change its source", func(m *catalog.Model) error {
		m.Source, m.Origin = source, origin
		return nil
	})
	if err == nil {
		rev := ""
		if origin != nil {
			rev = origin.Revision
		}
		s.log.Info("model source set", "model", name, "source", source, "revision", rev, "by", by)
	}
	return final, err
}

// SetOrigin records where a model's weights came from without changing them:
// for a model added before origins were kept. A model from configuration is
// left alone, its origin worked out again when it is needed.
func (s *Service) SetOrigin(ctx context.Context, name string, origin *catalog.Origin) error {
	if s.Declared(name) {
		return nil
	}
	_, err := s.update(ctx, name, "llmgw", "record its origin", func(m *catalog.Model) error {
		m.Origin = origin
		return nil
	})
	return err
}

// update changes one model added at runtime and keeps the change.
//
// Built before anything is written, like Add: catalog.New is what applies the
// defaults and refuses what does not validate, so the model that comes back
// is the one that will run rather than the one that was asked for.
func (s *Service) update(ctx context.Context, name, by, what string, change func(*catalog.Model) error) (catalog.Model, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.Declared(name) {
		return catalog.Model{}, fmt.Errorf(
			"model %q is declared in configuration; %s there, or it reverts on the next restart", name, what)
	}
	stored, err := s.storedModels(ctx)
	if err != nil {
		return catalog.Model{}, err
	}
	next := append([]catalog.Model{}, s.declared...)
	found := false
	for _, e := range stored {
		if e.Name == name {
			if err := change(&e); err != nil {
				return catalog.Model{}, err
			}
			found = true
		}
		next = append(next, e)
	}
	if !found {
		return catalog.Model{}, fmt.Errorf("model %q was not added at runtime", name)
	}
	built, err := catalog.New(next)
	if err != nil {
		return catalog.Model{}, fmt.Errorf("that would make the catalog invalid: %w", err)
	}
	final, ok := built.Resolve(name)
	if !ok {
		return catalog.Model{}, fmt.Errorf("model %q vanished from the catalog it was just built into", name)
	}
	def, err := json.Marshal(final)
	if err != nil {
		return catalog.Model{}, err
	}
	if err := s.store.PutModel(ctx, store.StoredModel{
		Name: name, Definition: def, CreatedAt: time.Now(), CreatedBy: by,
	}); err != nil {
		return catalog.Model{}, err
	}
	s.apply.ReplaceCatalog(ctx, built)
	return final, nil
}

// Load builds the catalog at startup: configuration plus what was stored.
//
// A stored model that no longer validates — its engine left the build, a
// field changed meaning — is reported and skipped. Refusing to start over one
// bad row would let a model somebody added from a laptop take down the cluster,
// and the models in configuration are precisely the ones that must survive it.
func (s *Service) Load(ctx context.Context) (*catalog.Catalog, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	stored, err := s.storedModels(ctx)
	if err != nil {
		return nil, err
	}

	next := append([]catalog.Model{}, s.declared...)
	for _, m := range stored {
		// Checked one at a time and against what is already accepted, so one
		// bad model cannot take a good one with it.
		candidate := append(append([]catalog.Model{}, next...), m)
		if _, err := catalog.New(candidate); err != nil {
			s.log.Warn("a stored model was skipped; it does not make a valid catalog",
				"model", m.Name, "err", err)
			continue
		}
		next = append(next, m)
	}
	return catalog.New(next)
}

// storedModels reads and decodes the runtime models, skipping any that no
// longer parse rather than failing every operation around them.
func (s *Service) storedModels(ctx context.Context) ([]catalog.Model, error) {
	if s.store == nil {
		return nil, nil
	}
	rows, err := s.store.Models(ctx)
	if err != nil {
		return nil, err
	}
	out := make([]catalog.Model, 0, len(rows))
	for _, r := range rows {
		var m catalog.Model
		if err := json.Unmarshal(r.Definition, &m); err != nil {
			s.log.Warn("a stored model could not be decoded and was skipped",
				"model", r.Name, "err", err)
			continue
		}
		out = append(out, m)
	}
	return out, nil
}

// SetCheck runs before a model is added, and refuses it by returning an error.
// Used to catch weights the engine cannot load, which otherwise are placed and
// then fail on every node they are tried on.
func (s *Service) SetCheck(check func(context.Context, catalog.Model) error) { s.check = check }