llmgw/internal/web/admin_capacity.go

package web

import (
	"context"
	"fmt"
	"net/http"
	"strconv"
	"strings"

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

// CapacityAdmin sets what parts of the tenancy tree are guaranteed of the
// cluster and may hold. The admin API implements it; the page and the CLI go
// through the same checks.
type CapacityAdmin interface {
	QuotaViews(ctx context.Context) ([]adminapi.QuotaView, error)
	SetQuota(ctx context.Context, req adminapi.QuotaRequest, by string) (string, error)
	ClearQuota(ctx context.Context, scope string) (string, error)
	Scopes(ctx context.Context) ([]adminapi.ScopeName, error)
}

// SetCapacity gives the quotas page capacity quotas to show and set. Without
// it the section is left off.
func (p *AccountPages) SetCapacity(c CapacityAdmin) { p.capacity = c }

// capacityRow is one scope's capacity quota as the page shows it.
type capacityRow struct {
	adminapi.QuotaView
	GPUs, Cores, Memory string
}

// capacityForm is what the form starts with: a scope's quota to change, or
// nothing for a new one.
type capacityForm struct {
	Scope               string
	GPUs, Cores, Memory string
	LGPUs, LCores, LMem string
}

func (p *AccountPages) capacitySection(ctx context.Context, d *adminQuotasData, editing string) {
	if p.capacity == nil {
		return
	}
	d.HasCapacity = true
	views, err := p.capacity.QuotaViews(ctx)
	if err != nil {
		d.Error = err.Error()
		return
	}
	for _, v := range views {
		d.Capacity = append(d.Capacity, capacityRow{QuotaView: v,
			GPUs:   capCell(v.Holds.GPUs, v.Guarantee.GPUs, v.Limit.GPUs),
			Cores:  capCell(v.Holds.CPUCores, v.Guarantee.CPUCores, v.Limit.CPUCores),
			Memory: capCell(v.Holds.MemoryGB, v.Guarantee.MemoryGB, v.Limit.MemoryGB)})
		if v.Scope == editing {
			d.CapForm = capacityForm{Scope: v.Scope,
				GPUs: intStr(v.Guarantee.GPUs), Cores: intStr(v.Guarantee.CPUCores), Memory: intStr(v.Guarantee.MemoryGB),
				LGPUs: intStr(v.Limit.GPUs), LCores: intStr(v.Limit.CPUCores), LMem: intStr(v.Limit.MemoryGB)}
		}
	}
	if d.Scopes, err = p.capacity.Scopes(ctx); err != nil {
		d.Error = err.Error()
	}
}

// capCell is "6 of 8 (at most 16)": held, guaranteed, limit.
func capCell(holds int, g, l *int) string {
	if g == nil && l == nil {
		return ""
	}
	out := strconv.Itoa(holds)
	if g != nil {
		out += " of " + strconv.Itoa(*g)
	}
	if l != nil {
		out += fmt.Sprintf(" (at most %d)", *l)
	}
	return out
}

func intStr(v *int) string {
	if v == nil {
		return ""
	}
	return strconv.Itoa(*v)
}

// applyCapacity sets or clears one scope's capacity quota from the form.
func (p *AccountPages) applyCapacity(r *http.Request, admin store.User) (string, error) {
	scope := strings.TrimSpace(r.FormValue("scope"))
	if scope == "" {
		return "", fmt.Errorf("choose where the quota applies")
	}
	if r.FormValue("clear") != "" {
		key, err := p.capacity.ClearQuota(r.Context(), scope)
		if err != nil {
			return "", err
		}
		return "capacity quota removed from " + key, nil
	}
	field := func(name, what string) (*int, error) {
		v := strings.TrimSpace(r.FormValue(name))
		if v == "" {
			return nil, nil
		}
		n, err := strconv.Atoi(v)
		if err != nil || n < 0 {
			return nil, fmt.Errorf("%s must be a whole number, or blank for none", what)
		}
		return &n, nil
	}
	var req adminapi.QuotaRequest
	req.Scope = scope
	var err error
	for _, f := range []struct {
		dst  **int
		name string
		what string
	}{
		{&req.Guarantee.GPUs, "gpus", "guaranteed GPUs"},
		{&req.Guarantee.CPUCores, "cores", "guaranteed cores"},
		{&req.Guarantee.MemoryGB, "memory", "guaranteed memory"},
		{&req.Limit.GPUs, "limit_gpus", "the GPU limit"},
		{&req.Limit.CPUCores, "limit_cores", "the core limit"},
		{&req.Limit.MemoryGB, "limit_memory", "the memory limit"},
	} {
		if *f.dst, err = field(f.name, f.what); err != nil {
			return "", err
		}
	}
	key, err := p.capacity.SetQuota(r.Context(), req, admin.Email)
	if err != nil {
		return "", err
	}
	return "capacity quota set on " + key + "; it takes effect within half a minute", nil
}