llmgw/internal/control/adminapi/quotas.go

package adminapi

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"sort"
	"strings"
	"time"

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

// Quotas is where capacity quotas and the tenancy tree they name are kept.
// The store implements it.
type Quotas interface {
	CapacityQuotas(ctx context.Context) ([]store.CapacityQuota, error)
	PutCapacityQuota(ctx context.Context, q store.CapacityQuota) error
	DeleteCapacityQuota(ctx context.Context, scope store.Scope, id string) error
	ScopeParents(ctx context.Context) (map[string][]string, error)
	AllOrgs(ctx context.Context) ([]store.Org, error)
	AllTeams(ctx context.Context) ([]store.TeamPath, error)
	AllProjects(ctx context.Context) ([]store.ProjectPath, error)
}

// SetQuotas supplies the quota store, which mounts the quota endpoints'
// behaviour; without it they answer 501.
func (a *API) SetQuotas(q Quotas) { a.quotas = q }

// QuotaView is one scope's quota with what it holds now.
type QuotaView struct {
	// Scope is team:ID; Name is how people know it, "Acme / ml".
	Scope     string          `json:"scope"`
	Name      string          `json:"name"`
	Guarantee store.Resources `json:"guarantee"`
	Limit     store.Resources `json:"limit"`
	Holds     quota.Amount    `json:"holds"`
	UpdatedBy string          `json:"updated_by,omitempty"`
	UpdatedAt time.Time       `json:"updated_at"`
}

// QuotaRequest sets a scope's quota. Scope may be written with an ID or a
// name: org:Acme, team:Acme/ml, project:Acme/ml/training.
type QuotaRequest struct {
	Scope     string          `json:"scope"`
	Guarantee store.Resources `json:"guarantee"`
	Limit     store.Resources `json:"limit"`
}

// QuotaSet is the scope a quota was set on, as org:ID, team:ID or
// project:ID whatever name it was asked by.
type QuotaSet struct {
	Scope string `json:"scope"`
}

// OwnerRequest sets the part of the tenancy tree a model's capacity counts
// against: org:NAME, team:ORG/TEAM or project:ORG/TEAM/PROJECT, or their IDs.
// Empty makes the model shared.
type OwnerRequest struct {
	Owner string `json:"owner"`
}

// OwnerView is a model's owner now, as scope:ID.
type OwnerView struct {
	Model string `json:"model"`
	Owner string `json:"owner"`
}

func (a *API) registerQuotas(mux Mux) {
	mux.HandleFunc("GET /api/v1/quotas", func(w http.ResponseWriter, r *http.Request) {
		if a.quotas == nil {
			writeError(w, http.StatusNotImplemented, "this control plane has no store")
			return
		}
		views, err := a.QuotaViews(r.Context())
		if err != nil {
			writeError(w, http.StatusInternalServerError, err.Error())
			return
		}
		writeJSON(w, http.StatusOK, views)
	})
	mux.HandleFunc("PUT /api/v1/quotas", func(w http.ResponseWriter, r *http.Request) {
		if a.quotas == nil {
			writeError(w, http.StatusNotImplemented, "this control plane has no store")
			return
		}
		var req QuotaRequest
		if err := json.NewDecoder(io.LimitReader(r.Body, 1<<14)).Decode(&req); err != nil {
			writeError(w, http.StatusBadRequest, "the request is not valid JSON: "+err.Error())
			return
		}
		key, err := a.SetQuota(r.Context(), req, actorOf(r))
		if err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		writeJSON(w, http.StatusOK, QuotaSet{Scope: key})
	})
	mux.HandleFunc("DELETE /api/v1/quotas/{scope}", func(w http.ResponseWriter, r *http.Request) {
		if a.quotas == nil {
			writeError(w, http.StatusNotImplemented, "this control plane has no store")
			return
		}
		key, err := a.ClearQuota(r.Context(), r.PathValue("scope"))
		if err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		writeJSON(w, http.StatusOK, Removed{Removed: key})
	})
	mux.HandleFunc("POST /api/v1/models/{name}/owner", func(w http.ResponseWriter, r *http.Request) {
		if a.models == nil || a.quotas == nil {
			writeError(w, http.StatusNotImplemented, "this control plane has no store")
			return
		}
		var req OwnerRequest
		if err := json.NewDecoder(io.LimitReader(r.Body, 1<<12)).Decode(&req); err != nil {
			writeError(w, http.StatusBadRequest, "the request is not valid JSON: "+err.Error())
			return
		}
		owner := ""
		if req.Owner != "" {
			scope, id, _, err := a.resolveScope(r.Context(), req.Owner)
			if err != nil {
				writeError(w, http.StatusBadRequest, err.Error())
				return
			}
			owner = quota.Key(scope, id)
		}
		m, err := a.models.SetOwner(r.Context(), r.PathValue("name"), owner, actorOf(r))
		if err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		writeJSON(w, http.StatusOK, OwnerView{Model: m.Name, Owner: m.Owner})
	})
}

// QuotaViews lists every quota with what its scope holds now.
func (a *API) QuotaViews(ctx context.Context) ([]QuotaView, error) {
	if a.quotas == nil {
		return nil, nil
	}
	qs, err := a.quotas.CapacityQuotas(ctx)
	if err != nil {
		return nil, err
	}
	parents, err := a.quotas.ScopeParents(ctx)
	if err != nil {
		return nil, err
	}
	names, err := a.scopeNames(ctx)
	if err != nil {
		return nil, err
	}
	usage := quota.New(qs, parents).Usage(a.ctl.Reconciler().Holdings())
	out := make([]QuotaView, 0, len(qs))
	for _, q := range qs {
		key := quota.Key(q.Scope, q.ScopeID)
		out = append(out, QuotaView{Scope: key, Name: names[key], Guarantee: q.Guarantee, Limit: q.Limit,
			Holds: usage[key], UpdatedBy: q.UpdatedBy, UpdatedAt: q.UpdatedAt})
	}
	return out, nil
}

// scopeNames maps org:ID, team:ID and project:ID to their paths.
func (a *API) scopeNames(ctx context.Context) (map[string]string, error) {
	out := map[string]string{}
	orgs, err := a.quotas.AllOrgs(ctx)
	if err != nil {
		return nil, err
	}
	for _, o := range orgs {
		out["org:"+o.ID] = o.Name
	}
	teams, err := a.quotas.AllTeams(ctx)
	if err != nil {
		return nil, err
	}
	for _, t := range teams {
		out["team:"+t.ID] = t.OrgName + " / " + t.Name
	}
	projects, err := a.quotas.AllProjects(ctx)
	if err != nil {
		return nil, err
	}
	for _, p := range projects {
		out["project:"+p.ID] = p.Path
	}
	return out, nil
}

// resolveScope reads org:X, team:X or project:X, where X is an ID or the
// scope's path — Acme, Acme/ml, Acme/ml/training — and returns its ID and
// path. A path that could mean two scopes is refused rather than guessed.
func (a *API) resolveScope(ctx context.Context, s string) (store.Scope, string, string, error) {
	kind, what, ok := strings.Cut(strings.TrimSpace(s), ":")
	scope := store.Scope(kind)
	if !ok || what == "" || (scope != store.ScopeOrg && scope != store.ScopeTeam && scope != store.ScopeProject) {
		return "", "", "", fmt.Errorf("%q is not org:NAME, team:ORG/TEAM or project:ORG/TEAM/PROJECT", s)
	}
	names, err := a.scopeNames(ctx)
	if err != nil {
		return "", "", "", err
	}
	norm := func(p string) string {
		parts := strings.Split(p, "/")
		for i := range parts {
			parts[i] = strings.ToLower(strings.TrimSpace(parts[i]))
		}
		return strings.Join(parts, "/")
	}
	if name, ok := names[kind+":"+what]; ok {
		return scope, what, name, nil
	}
	match := func(same func(name string) bool) []string {
		var found []string
		for key, name := range names {
			if strings.HasPrefix(key, kind+":") && same(name) {
				found = append(found, key)
			}
		}
		return found
	}
	found := match(func(name string) bool { return norm(name) == norm(what) })
	if len(found) == 0 {
		// Its own name alone — ml rather than Acme/ml — when that is unique.
		found = match(func(name string) bool {
			parts := strings.Split(norm(name), "/")
			return parts[len(parts)-1] == norm(what)
		})
	}
	switch len(found) {
	case 0:
		return "", "", "", fmt.Errorf("there is no %s %q", kind, what)
	case 1:
		_, id, _ := strings.Cut(found[0], ":")
		return scope, id, names[found[0]], nil
	}
	return "", "", "", fmt.Errorf("%q could be any of %d %ss; give its ID", what, len(found), kind)
}

// SetQuota sets a scope's capacity quota, named by path or ID, and returns
// the scope as org:ID.
func (a *API) SetQuota(ctx context.Context, req QuotaRequest, by string) (string, error) {
	if a.quotas == nil {
		return "", fmt.Errorf("this control plane has no store")
	}
	scope, id, _, err := a.resolveScope(ctx, req.Scope)
	if err != nil {
		return "", err
	}
	if !req.Guarantee.Any() && !req.Limit.Any() {
		return "", fmt.Errorf("a quota needs a guarantee or a limit; to remove one, clear it")
	}
	q := store.CapacityQuota{Scope: scope, ScopeID: id, Guarantee: req.Guarantee, Limit: req.Limit, UpdatedBy: by}
	if err := a.quotas.PutCapacityQuota(ctx, q); err != nil {
		return "", err
	}
	return quota.Key(scope, id), nil
}

// ClearQuota removes a scope's capacity quota.
func (a *API) ClearQuota(ctx context.Context, scope string) (string, error) {
	if a.quotas == nil {
		return "", fmt.Errorf("this control plane has no store")
	}
	s, id, _, err := a.resolveScope(ctx, scope)
	if err != nil {
		return "", err
	}
	if err := a.quotas.DeleteCapacityQuota(ctx, s, id); err != nil {
		return "", err
	}
	return quota.Key(s, id), nil
}

// ScopeName is a part of the tenancy tree: org:ID and its path.
type ScopeName struct {
	Key  string `json:"scope"`
	Name string `json:"name"`
}

// Scopes lists every organization, team and project, by path.
func (a *API) Scopes(ctx context.Context) ([]ScopeName, error) {
	if a.quotas == nil {
		return nil, nil
	}
	names, err := a.scopeNames(ctx)
	if err != nil {
		return nil, err
	}
	out := make([]ScopeName, 0, len(names))
	for k, n := range names {
		out = append(out, ScopeName{Key: k, Name: n})
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name })
	return out, nil
}