llmgw/internal/cli/federation.go

package cli

import (
	"encoding/json"
	"fmt"
	"net/url"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/spf13/cobra"

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

func newFederationCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	c := &cobra.Command{
		Use:   "federation",
		Short: "Offer another llmgw cluster's models here, as PEER/MODEL",
		Long: "A peer is another llmgw gateway, reached with a key it issued to this\n" +
			"cluster. Its models join this catalog as PEER/MODEL and follow its own\n" +
			"list; each is sent requests only while the peer's /ready answers, which\n" +
			"is checked every federation.interval (10s) and timed. Name one as a\n" +
			"local model's fallback to send overflow to the other site:\n" +
			"  llmgw ctl model fallbacks qwen-7b east/qwen-7b\n\n" +
			"A project whose requests stay in the cluster (keys residency\n" +
			"--in-cluster), or a policy that keeps nothing, is never sent to a peer.\n" +
			"Requests carry x-llmgw-hop, so two clusters federating each other do\n" +
			"not send one back and forth.",
	}
	c.AddCommand(newFederationAddCmd(client), newFederationListCmd(client, asJSON),
		newFederationRemoveCmd(client), newFederationSyncCmd(client, asJSON))
	return c
}

func newFederationAddCmd(client func() *ctlClient) *cobra.Command {
	var req federation.AddRequest
	c := &cobra.Command{
		Use:   "add NAME --url URL --key-file FILE",
		Short: "Add another llmgw cluster as a peer, or change one",
		Long: "The peer is checked first — its /ready, and its /v1/models with the\n" +
			"key — and refused if either fails. The key file is read on the control\n" +
			"plane (or env:NAME, vault:PATH#FIELD), at every check, so a rotated\n" +
			"key is used without a change here; every control plane needs it.\n\n" +
			"--send-tags makes each request sent to the peer carry x-llmgw-tags\n" +
			"federated_from=<this cluster>,org=,team=,project=<ids>, so the peer's\n" +
			"usage can say who here asked. Off unless given: the peer otherwise sees\n" +
			"one key, this cluster's, and nothing about who is behind it.",
		Example: "  llmgw ctl federation add east --url https://llm.east.example --key-file /etc/llmgw/east.key\n" +
			"  llmgw ctl federation add east --url https://llm.east.example --key-file /etc/llmgw/east.key --send-tags",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			req.Name = args[0]
			body, _ := json.Marshal(req)
			var st federation.Status
			if err := client().post(cmd.Context(), "/api/v1/federation", body, &st); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s added: %d models, %s\n", st.Name, len(st.Models), healthText(st))
			for _, m := range st.Models {
				fmt.Fprintln(cmd.OutOrStdout(), "  "+m.Name)
			}
			return nil
		},
	}
	f := c.Flags()
	f.StringVar(&req.URL, "url", "", "the peer gateway's client address, without /v1")
	f.StringVar(&req.KeyFile, "key-file", "", "file on the control plane holding the key the peer issued to this cluster")
	f.BoolVar(&req.SendTags, "send-tags", false, "tell the peer which cluster and project each request is from")
	_ = c.MarkFlagRequired("url")
	_ = c.MarkFlagRequired("key-file")
	return c
}

func newFederationListCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "list",
		Short: "Peers, whether each is up, how far away, and the models it offers",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			var v adminapi.FederationView
			if err := client().get(cmd.Context(), "/api/v1/federation", &v); err != nil {
				return err
			}
			if *asJSON {
				return emitJSON(v)
			}
			out := cmd.OutOrStdout()
			fmt.Fprintf(out, "this cluster: %s\n", v.Cluster)
			if len(v.Peers) == 0 {
				fmt.Fprintln(out, "no peers: llmgw ctl federation add NAME --url URL --key-file FILE")
				return nil
			}
			w := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
			fmt.Fprintln(w, "PEER\tURL\tHEALTH\tLATENCY\tTAGS\tMODELS")
			for _, p := range v.Peers {
				tags := "off"
				if p.SendTags {
					tags = "on"
				}
				lat := "-"
				if p.Latency > 0 {
					lat = fmt.Sprintf("%.1fms", p.LatencyMs)
				}
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n", p.Name, p.URL, healthText(p), lat, tags, modelList(p))
			}
			return w.Flush()
		},
	}
}

func newFederationRemoveCmd(client func() *ctlClient) *cobra.Command {
	return &cobra.Command{
		Use:   "remove NAME",
		Short: "Remove a peer and its models",
		Long: "Refused while a model here falls back to one of the peer's: change its\n" +
			"fallbacks first (llmgw ctl model fallbacks).",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var out adminapi.Removed
			if err := client().delete(cmd.Context(), "/api/v1/federation/"+url.PathEscape(args[0]), &out); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s removed\n", out.Removed)
			return nil
		},
	}
}

func newFederationSyncCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "sync NAME",
		Short: "Check a peer now and take its current models, rather than at the next interval",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var st federation.Status
			if err := client().post(cmd.Context(), "/api/v1/federation/"+url.PathEscape(args[0])+"/sync", nil, &st); err != nil {
				return err
			}
			if *asJSON {
				return emitJSON(st)
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s: %s, %s\n", st.Name, healthText(st), modelList(st))
			return nil
		},
	}
}

// healthText is "up" or "down: why", with how long ago it was checked.
func healthText(st federation.Status) string {
	ago := ""
	if !st.Checked.IsZero() {
		ago = fmt.Sprintf(" (%s ago)", time.Since(st.Checked).Round(time.Second))
	}
	if st.Up {
		return "up" + ago
	}
	return "down" + ago + ": " + st.Error
}

// modelList is the peer's models by their names here, a withdrawn one
// marked as such.
func modelList(st federation.Status) string {
	if len(st.Models) == 0 {
		return "-"
	}
	names := make([]string, 0, len(st.Models))
	for _, m := range st.Models {
		n := m.Name
		if m.Withdrawn {
			n += " (withdrawn)"
		}
		names = append(names, n)
	}
	return strings.Join(names, ", ")
}

func newModelFallbacksCmd(client func() *ctlClient) *cobra.Command {
	var clear bool
	c := &cobra.Command{
		Use:   "fallbacks MODEL [FALLBACK...]",
		Short: "Set the models a model falls back to when it has no capacity",
		Long: "Replaces the model's fallbacks, tried in order when nothing serving it\n" +
			"has room. A federation peer's model, PEER/MODEL, sends the overflow to\n" +
			"the other cluster. A model added at runtime only; one declared in\n" +
			"server.yaml has its fallbacks there. --clear takes them all away.",
		Example: "  llmgw ctl model fallbacks qwen-7b east/qwen-7b\n" +
			"  llmgw ctl model fallbacks qwen-7b qwen-3b east/qwen-7b\n" +
			"  llmgw ctl model fallbacks qwen-7b --clear",
		Args: cobra.MinimumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) == 1 && !clear {
				return fmt.Errorf("name the fallbacks, or give --clear to take them away")
			}
			body, _ := json.Marshal(adminapi.FallbacksRequest{Fallbacks: args[1:]})
			var v adminapi.FallbacksView
			if err := client().put(cmd.Context(), "/api/v1/models/"+args[0]+"/fallbacks", body, &v); err != nil {
				return err
			}
			if len(v.Fallbacks) == 0 {
				fmt.Fprintf(cmd.OutOrStdout(), "%s has no fallbacks\n", v.Model)
				return nil
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s falls back to %s\n", v.Model, strings.Join(v.Fallbacks, ", then "))
			return nil
		},
	}
	c.Flags().BoolVar(&clear, "clear", false, "take every fallback away")
	return c
}