llmgw/internal/cli/ctlweights.go

package cli

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/spf13/cobra"

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

// newWeightsCmd reports and reclaims model files nothing needs.
//
// The dangerous half of storage management, so its shape is defensive: listing
// is the default and deleting takes a flag, only files llmgw fetched are ever
// candidates, and the node checks every path again before unlinking it.
//
// What it will never touch is the more important half of the contract. Weights
// an operator put on a node — by scp, by a shared mount, by hand — are none of
// this system's business however unreferenced they look, and a model that took
// a week to copy must not be removable by a command someone ran to tidy up.
func newWeightsCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "weights",
		Short: "Show cached model files no model in the catalog needs",
		Long: "Lists weights llmgw downloaded for models that are no longer in the\n" +
			"catalog, per node, with what deleting them would free.\n\n" +
			"Only files llmgw fetched are considered. Anything an operator put on a\n" +
			"node is left alone, whatever the catalog says — llmgw observes storage\n" +
			"rather than owning it, and only takes back what it put there itself.\n\n" +
			"Nothing is deleted without --collect.",
		Args: cobra.NoArgs,
	}
	var collect bool
	var node string

	cmd.Flags().BoolVar(&collect, "collect", false, "delete them, rather than only listing")
	cmd.Flags().StringVar(&node, "node", "", "limit to one node")

	cmd.RunE = func(cmd *cobra.Command, _ []string) error {
		if collect {
			return collectOrphans(cmd, client(), node, *asJSON)
		}
		return listOrphanWeights(cmd, client(), node, *asJSON)
	}
	return cmd
}

func listOrphanWeights(cmd *cobra.Command, c *ctlClient, node string, asJSON bool) error {
	var found []adminapi.OrphanView
	if err := c.get(cmd.Context(), "/api/v1/weights/orphans", &found); err != nil {
		return err
	}
	if node != "" {
		found = onlyNode(found, node)
	}
	if asJSON {
		return emitJSON(found)
	}
	if len(found) == 0 {
		fmt.Fprintln(cmd.OutOrStdout(),
			"Nothing to reclaim: every file llmgw fetched is still wanted by a model.")
		return nil
	}

	w := tab()
	fmt.Fprintln(w, "NODE\tFILES\tWOULD FREE")
	var total int64
	for _, o := range found {
		fmt.Fprintf(w, "%s\t%d\t%s\n", o.Node, len(o.Paths), fileSize(o.Bytes))
		total += o.Bytes
	}
	_ = w.Flush()

	for _, o := range found {
		for _, p := range o.Paths {
			fmt.Fprintf(cmd.OutOrStdout(), "\n  %s: %s", o.Node, p)
		}
	}
	fmt.Fprintf(cmd.OutOrStdout(), "\n\n%s in total. Delete with --collect.\n", fileSize(total))
	return nil
}

func collectOrphans(cmd *cobra.Command, c *ctlClient, node string, asJSON bool) error {
	body, err := json.Marshal(map[string]string{"node": node})
	if err != nil {
		return err
	}
	var results []adminapi.CollectResult
	// The control plane waits for the nodes to list their files without the
	// deleted ones: half a minute at most.
	if err := c.patient(time.Minute).post(cmd.Context(), "/api/v1/weights/collect", body, &results); err != nil {
		return err
	}
	if asJSON {
		return emitJSON(results)
	}
	if len(results) == 0 {
		fmt.Fprintln(cmd.OutOrStdout(), "Nothing to reclaim.")
		return nil
	}
	var freed, left int64
	for _, r := range results {
		if r.Error != "" {
			// Reported rather than fatal: the other nodes' space is still worth
			// freeing, and this one will be offered again next time.
			fmt.Fprintf(cmd.OutOrStdout(), "%s: %s\n", r.Node, r.Error)
			continue
		}
		line := fmt.Sprintf("%s: removed %d files, %s", r.Node, r.Files, fileSize(r.Bytes))
		if r.Left > 0 {
			line += fmt.Sprintf("; %d still there, %s", r.Left, fileSize(r.LeftBytes))
		}
		fmt.Fprintln(cmd.OutOrStdout(), line)
		freed, left = freed+r.Bytes, left+r.LeftBytes
	}
	fmt.Fprintf(cmd.OutOrStdout(), "\n%s freed.\n", fileSize(freed))
	if left > 0 {
		// Each node decides: one that keeps a path it cannot prove it
		// fetched is the feature working. Its agent's log names each.
		fmt.Fprintf(cmd.OutOrStdout(), "%s the nodes kept; each node's agent log says why "+
			"(journalctl -u llmgw-agent).\n", fileSize(left))
	}
	return nil
}

func onlyNode(in []adminapi.OrphanView, node string) []adminapi.OrphanView {
	out := make([]adminapi.OrphanView, 0, len(in))
	for _, o := range in {
		if o.Node == node {
			out = append(out, o)
		}
	}
	return out
}

// fileSize renders a size the way someone reads one when deciding whether a
// file is worth deleting.
//
// Separate from the per-token cache size renderer in modelplan, which never
// leaves KiB and MiB. Model files are tens of gigabytes, and "8601.6 MiB" is
// not a number anyone weighs a decision against.
func fileSize(b int64) string {
	switch {
	case b >= 1<<30:
		return fmt.Sprintf("%.1f GB", float64(b)/(1<<30))
	case b >= 1<<20:
		return fmt.Sprintf("%.0f MB", float64(b)/(1<<20))
	case b > 0:
		return fmt.Sprintf("%d B", b)
	default:
		// The node reported no size. Saying "0 B" would suggest deleting it
		// gains nothing, which is a different claim from not knowing.
		return "size unknown"
	}
}

// fetchColumn renders a download's progress.
//
// A percentage when the size is known and a byte count when it is not, because
// a server that does not send Content-Length is common enough — and an
// invented percentage that stops moving is worse than a number that rises.
func fetchColumn(f *adminapi.FetchView) string {
	if f.Percent != nil {
		return fmt.Sprintf("%s  %d%% of %s", shortSource(f.Source), *f.Percent, fileSize(f.BytesTotal))
	}
	return fmt.Sprintf("%s  %s so far", shortSource(f.Source), fileSize(f.BytesDone))
}

// shortSource keeps the end of a source, which is the part that identifies it.
func shortSource(s string) string {
	if i := strings.LastIndex(s, "/"); i >= 0 && i < len(s)-1 {
		return s[i+1:]
	}
	return s
}