llmgw/internal/store/capacity.go

package store

import (
	"context"
	"encoding/json"
	"fmt"
	"time"
)

// StoredCapacity is a benchmark's conclusion, kept for whoever has to size the
// model later.
type StoredCapacity struct {
	Model   string
	Profile string
	// Measurement is bench.Capacity as JSON. Stored whole for the reason a
	// model definition is: a measurement is a nested shape that is still
	// growing, and as columns every field added would be a migration.
	Measurement []byte
	MeasuredAt  time.Time
	MeasuredBy  string
}

// PutCapacity records what a replica of a model was measured to serve.
//
// Replacing rather than accumulating. A second measurement of the same model on
// the same profile is a better answer to the same question, not another data
// point: hardware and engine builds change, and keeping every run would leave a
// reader choosing between numbers with no basis for the choice. The run itself
// is printed when it happens, for anyone who wants the ramp.
func (s *DB) PutCapacity(ctx context.Context, c StoredCapacity) error {
	if c.Model == "" || c.Profile == "" {
		return fmt.Errorf("a capacity measurement needs a model and a profile")
	}
	if !json.Valid(c.Measurement) {
		return fmt.Errorf("capacity for %s/%s: the measurement is not valid JSON", c.Model, c.Profile)
	}
	if c.MeasuredAt.IsZero() {
		c.MeasuredAt = time.Now()
	}
	_, err := s.db.ExecContext(ctx,
		`INSERT INTO capacity (model, profile, measurement, measured_at, measured_by)
		 VALUES (?, ?, ?, ?, ?)
		 ON CONFLICT(model, profile) DO UPDATE SET
		     measurement = excluded.measurement,
		     measured_at = excluded.measured_at,
		     measured_by = excluded.measured_by`,
		c.Model, c.Profile, string(c.Measurement), c.MeasuredAt, c.MeasuredBy)
	return err
}

// Capacities lists every measurement, newest first.
//
// Newest first because the question a reader usually has is "what do we know
// lately", and an old measurement is the one most likely to be describing
// hardware that has since changed.
func (s *DB) Capacities(ctx context.Context) ([]StoredCapacity, error) {
	rows, err := s.db.QueryContext(ctx,
		`SELECT model, profile, measurement, measured_at, measured_by
		 FROM capacity ORDER BY measured_at DESC, model, profile`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var out []StoredCapacity
	for rows.Next() {
		var c StoredCapacity
		var m string
		if err := rows.Scan(&c.Model, &c.Profile, &m, &c.MeasuredAt, &c.MeasuredBy); err != nil {
			return nil, err
		}
		c.Measurement = []byte(m)
		out = append(out, c)
	}
	return out, rows.Err()
}

// DeleteCapacity forgets a measurement.
//
// Wanted when a model is removed, and wanted on its own: a profile whose
// hardware has been replaced is worse than no measurement, because it looks
// like knowledge. Deleting one that is not there is not an error — the caller
// wanted it gone, and it is.
func (s *DB) DeleteCapacity(ctx context.Context, model, profile string) error {
	_, err := s.db.ExecContext(ctx,
		`DELETE FROM capacity WHERE model = ? AND profile = ?`, model, profile)
	return err
}