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 }