llmgw/internal/cli/ctlweights.go

package cli

import (
	"encoding/json"
	"fmt"

	"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 []struct {
		Node  string `json:"node"`
		Files int    `json:"files"`
		Bytes int64  `json:"bytes"`
		Error string `json:"error,omitempty"`
	}
	if err := c.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 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
		}
		fmt.Fprintf(cmd.OutOrStdout(), "%s: asked to remove %d files, %s\n",
			r.Node, r.Files, fileSize(r.Bytes))
		freed += r.Bytes
	}
	// "asked", not "removed": each node decides, and one that refuses a path it
	// cannot prove it fetched is the feature working rather than failing.
	fmt.Fprintf(cmd.OutOrStdout(), "\n%s requested. The nodes verify each path before deleting it.\n",
		fileSize(freed))
	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"
	}
}