llmgw/internal/store/quotas.go

package store

import (
	"context"
	"database/sql"
	"fmt"
	"time"
)

// Resources are amounts of the three things a replica holds of a cluster.
// Nil is not governed: nothing guaranteed, or nothing limited.
type Resources struct {
	GPUs     *int `json:"gpus,omitempty"`
	CPUCores *int `json:"cpu_cores,omitempty"`
	MemoryGB *int `json:"memory_gb,omitempty"`
}

// Any reports whether any resource is governed.
func (r Resources) Any() bool { return r.GPUs != nil || r.CPUCores != nil || r.MemoryGB != nil }

// CapacityQuota is what one part of the tenancy tree is guaranteed of the
// cluster, and the most it may hold.
type CapacityQuota struct {
	Scope   Scope  `json:"scope"`
	ScopeID string `json:"scope_id"`
	// Guarantee is held for the scope: while it holds less, what others
	// borrowed beyond their own guarantees is taken back for it.
	Guarantee Resources `json:"guarantee"`
	// Limit is the most it may hold, borrowing included.
	Limit     Resources `json:"limit"`
	UpdatedAt time.Time `json:"updated_at"`
	UpdatedBy string    `json:"updated_by"`
}

// PutCapacityQuota sets a scope's quota, replacing what was there.
func (s *DB) PutCapacityQuota(ctx context.Context, q CapacityQuota) error {
	if !validMembershipScope(q.Scope) || q.ScopeID == "" {
		return fmt.Errorf("a capacity quota belongs to an organization, team or project")
	}
	for _, v := range []*int{q.Guarantee.GPUs, q.Guarantee.CPUCores, q.Guarantee.MemoryGB,
		q.Limit.GPUs, q.Limit.CPUCores, q.Limit.MemoryGB} {
		if v != nil && *v < 0 {
			return fmt.Errorf("a quota cannot be negative")
		}
	}
	if q.UpdatedAt.IsZero() {
		q.UpdatedAt = time.Now()
	}
	_, err := s.db.ExecContext(ctx,
		`INSERT INTO capacity_quotas (scope_type, scope_id, gpus, cpu_cores, memory_gb,
		   limit_gpus, limit_cpu_cores, limit_memory_gb, updated_at, updated_by)
		 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
		 ON CONFLICT (scope_type, scope_id) DO UPDATE SET
		   gpus = excluded.gpus, cpu_cores = excluded.cpu_cores, memory_gb = excluded.memory_gb,
		   limit_gpus = excluded.limit_gpus, limit_cpu_cores = excluded.limit_cpu_cores,
		   limit_memory_gb = excluded.limit_memory_gb,
		   updated_at = excluded.updated_at, updated_by = excluded.updated_by`,
		string(q.Scope), q.ScopeID,
		nullInt(q.Guarantee.GPUs), nullInt(q.Guarantee.CPUCores), nullInt(q.Guarantee.MemoryGB),
		nullInt(q.Limit.GPUs), nullInt(q.Limit.CPUCores), nullInt(q.Limit.MemoryGB),
		q.UpdatedAt.UnixMilli(), q.UpdatedBy)
	return err
}

// CapacityQuotas lists every quota.
func (s *DB) CapacityQuotas(ctx context.Context) ([]CapacityQuota, error) {
	return queryAll(ctx, s.db,
		`SELECT scope_type, scope_id, gpus, cpu_cores, memory_gb,
		   limit_gpus, limit_cpu_cores, limit_memory_gb, updated_at, updated_by
		 FROM capacity_quotas ORDER BY scope_type, scope_id`, nil,
		func(rows *sql.Rows) (CapacityQuota, error) {
			var q CapacityQuota
			var scope string
			var g [6]sql.NullInt64
			var at int64
			err := rows.Scan(&scope, &q.ScopeID, &g[0], &g[1], &g[2], &g[3], &g[4], &g[5], &at, &q.UpdatedBy)
			q.Scope = Scope(scope)
			q.Guarantee = Resources{GPUs: intOf(g[0]), CPUCores: intOf(g[1]), MemoryGB: intOf(g[2])}
			q.Limit = Resources{GPUs: intOf(g[3]), CPUCores: intOf(g[4]), MemoryGB: intOf(g[5])}
			q.UpdatedAt = time.UnixMilli(at)
			return q, err
		})
}

// DeleteCapacityQuota removes a scope's quota.
func (s *DB) DeleteCapacityQuota(ctx context.Context, scope Scope, id string) error {
	_, err := s.db.ExecContext(ctx,
		`DELETE FROM capacity_quotas WHERE scope_type = ? AND scope_id = ?`, string(scope), id)
	return err
}

// ScopeParents maps every team and project to the scopes above it, nearest
// first: a project to its team and organization, a team to its organization.
// Keys and values are written scope:id.
func (s *DB) ScopeParents(ctx context.Context) (map[string][]string, error) {
	out := map[string][]string{}
	teams, err := s.AllTeams(ctx)
	if err != nil {
		return nil, err
	}
	orgOf := map[string]string{}
	for _, t := range teams {
		orgOf[t.ID] = t.OrgID
		out["team:"+t.ID] = []string{"org:" + t.OrgID}
	}
	type pt struct{ project, team string }
	rows, err := queryAll(ctx, s.db, `SELECT id, team_id FROM projects`, nil,
		func(r *sql.Rows) (pt, error) {
			var p pt
			return p, r.Scan(&p.project, &p.team)
		})
	if err != nil {
		return nil, err
	}
	for _, p := range rows {
		out["project:"+p.project] = []string{"team:" + p.team, "org:" + orgOf[p.team]}
	}
	return out, nil
}

func intOf(n sql.NullInt64) *int {
	if !n.Valid {
		return nil
	}
	v := int(n.Int64)
	return &v
}