llmgw/internal/control/adminapi/weights.go

package adminapi

import (
	"context"
	"encoding/json"
	"net/http"
	"sort"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control"
)

// OrphanView is cached weights on one node that no catalog model accounts for.
type OrphanView struct {
	Node  string   `json:"node"`
	Paths []string `json:"paths"`
	// Bytes is what deleting them would free, as far as the node reported. A
	// file whose size was not reported counts as zero rather than being left
	// out, so the list and the number always describe the same set.
	Bytes int64 `json:"bytes"`
}

// Orphans finds cached weights nothing in the catalog points at.
//
// Computed here from what nodes already report, so listing costs no round trip
// and a node that is offline simply contributes nothing — its files are not
// deletable while it is gone anyway.
//
// Only files llmgw fetched are ever considered. Everything else in a model
// directory was put there by the operator and is none of this system's
// business, however unreferenced it looks.
func (a *API) Orphans() []OrphanView {
	var wanted []string
	for _, m := range a.ctl.Catalog().Models() {
		wanted = append(wanted, m.WeightSources()...)
		// A model's adapters are its weights too; collecting them would
		// leave replicas that cannot restart.
		for _, ad := range m.Adapters {
			wanted = append(wanted, ad.Source)
		}
	}

	out := make([]OrphanView, 0)
	for _, n := range a.ctl.Inventory().Nodes() {
		size := map[string]int64{}
		have := make([]string, 0, len(n.Weights))
		for _, w := range n.Weights {
			have = append(have, w.Path)
			size[w.Path] = w.SizeBytes
		}

		paths := weights.Orphans(have, wanted)
		if len(paths) == 0 {
			continue
		}
		v := OrphanView{Node: n.ID, Paths: paths}
		for _, p := range paths {
			v.Bytes += size[p]
		}
		out = append(out, v)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Node < out[j].Node })
	return out
}

func (a *API) listOrphans(w http.ResponseWriter, _ *http.Request) {
	writeJSON(w, http.StatusOK, a.Orphans())
}

// CollectRequest limits a sweep to one node. Empty means every node, which
// is the common case after removing a model from the catalog.
type CollectRequest struct {
	Node string `json:"node,omitempty"`
}

// CollectResult is what was deleted on one node, or why nothing was: what
// the node no longer lists once asked, not what it was asked to delete.
type CollectResult struct {
	Node  string `json:"node"`
	Files int    `json:"files"`
	Bytes int64  `json:"bytes"`
	// Left and LeftBytes are what the node still lists: refused, or not
	// gone by the time the sweep stopped waiting. Its agent's log says why.
	Left      int    `json:"left,omitempty"`
	LeftBytes int64  `json:"left_bytes,omitempty"`
	Error     string `json:"error,omitempty"`
}

// confirmWait is how long a sweep waits for the nodes to list their files
// without the ones deleted: they scan theirs every heartbeat, five seconds
// by default, and a large set takes a moment to delete.
var confirmWait = 30 * time.Second

// removal is what one node is asked to delete, with each file's size.
type removal struct {
	node  string
	paths []string
	size  map[string]int64
}

// removeConfirmed asks each node to delete its paths and says what each then
// no longer lists. Counted as asked, a sweep said "25 GB freed" of files the
// node had refused and kept.
func (a *API) removeConfirmed(ctx context.Context, rm control.WeightRemover, rs []removal) []CollectResult {
	out := make([]CollectResult, len(rs))
	pending := map[int]bool{}
	for i, r := range rs {
		out[i].Node = r.node
		if err := rm.RemoveWeights(ctx, r.node, r.paths); err != nil {
			out[i].Error = err.Error()
			for _, p := range r.paths {
				out[i].Left++
				out[i].LeftBytes += r.size[p]
			}
			continue
		}
		pending[i] = true
	}
	deadline := time.Now().Add(confirmWait)
	for {
		for i := range pending {
			if left := a.stillListed(rs[i]); len(left) == 0 {
				delete(pending, i)
			}
		}
		if len(pending) == 0 || !time.Now().Before(deadline) {
			break
		}
		select {
		case <-ctx.Done():
			deadline = time.Now()
		case <-time.After(time.Second):
		}
	}
	for i, r := range rs {
		if out[i].Error != "" {
			continue
		}
		left := a.stillListed(r)
		for _, p := range r.paths {
			if left[p] {
				out[i].Left++
				out[i].LeftBytes += r.size[p]
			} else {
				out[i].Files++
				out[i].Bytes += r.size[p]
			}
		}
	}
	return out
}

// stillListed is which of a removal's paths its node still lists.
func (a *API) stillListed(r removal) map[string]bool {
	left := map[string]bool{}
	n, ok := a.ctl.Inventory().Get(r.node)
	if !ok {
		return left
	}
	want := map[string]bool{}
	for _, p := range r.paths {
		want[p] = true
	}
	for _, w := range n.Weights {
		if want[w.Path] {
			left[w.Path] = true
		}
	}
	return left
}

// collectWeights deletes them.
//
// A proposal, not an instruction: each node checks every path again before
// unlinking it, and refuses anything it cannot prove it fetched. The two
// checks are deliberate duplication — the control plane knows the catalog and
// the node knows its disk, and only the node can be sure.
func (a *API) collectWeights(w http.ResponseWriter, r *http.Request) {
	rm := a.ctl.WeightRemover()
	if rm == nil {
		writeError(w, http.StatusNotImplemented,
			"this control plane cannot ask nodes to delete files")
		return
	}

	var body CollectRequest
	// An empty body is the "all nodes" case and must not be an error.
	_ = json.NewDecoder(r.Body).Decode(&body)

	// A node that has gone since the list was built is reported rather than
	// failing the whole sweep: the other nodes' space is still worth
	// freeing, and this one will be offered again next time.
	var rs []removal
	for _, o := range a.Orphans() {
		if body.Node != "" && o.Node != body.Node {
			continue
		}
		rs = append(rs, removal{node: o.Node, paths: o.Paths, size: a.sizesOn(o.Node)})
	}
	writeJSON(w, http.StatusOK, a.removeConfirmed(r.Context(), rm, rs))
}

// sizesOn is each file a node lists, by path, with its size.
func (a *API) sizesOn(node string) map[string]int64 {
	out := map[string]int64{}
	if n, ok := a.ctl.Inventory().Get(node); ok {
		for _, w := range n.Weights {
			out[w.Path] = w.SizeBytes
		}
	}
	return out
}