llmgw/internal/cli/agent.go

package cli

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"text/tabwriter"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
	"gitea.mynodes.xyz/ilker/llmgw/internal/version"
)

func newAgentCmd() *cobra.Command {
	var configPath string

	var (
		server   string
		nodeID   string
		mgmtAddr string
	)

	cmd := &cobra.Command{
		Use:   "agent",
		Short: "Node agent: discovery and engine supervision",
		Long: "Runs on every GPU node. It discovers the hardware, connects to the control\n" +
			"plane, and starts, supervises and stops inference engines on demand.\n\n" +
			"The agent dials out and keeps one stream open. Engines keep serving if the\n" +
			"agent restarts or the control plane goes away: the router talks to them\n" +
			"directly, and orphan adoption picks them up again.",
		RunE: func(cmd *cobra.Command, _ []string) error {
			log := logger()

			cfg, err := config.LoadAgent(configPath)
			if err != nil {
				return err
			}
			if server != "" {
				cfg.Server = server
			}
			if nodeID != "" {
				cfg.NodeID = nodeID
			}
			if mgmtAddr != "" {
				cfg.MgmtAddr = mgmtAddr
			}
			if cfg.NodeID == "" {
				host, err := os.Hostname()
				if err != nil {
					return fmt.Errorf("node_id is unset and the hostname is unavailable: %w", err)
				}
				cfg.NodeID = host
			}
			if cfg.Hostname == "" {
				cfg.Hostname, _ = os.Hostname()
			}
			cfg.Version = version.Version

			// Before the agent, so its tracer is the configured one.
			shutdownTracing, err := obs.SetupTracing(cmd.Context(), obs.TracingOptions{
				Endpoint: cfg.Tracing.Endpoint, SampleRatio: cfg.Tracing.SampleRatio,
				Service: "llmgw-agent", Version: version.Version, Instance: cfg.NodeID,
			})
			if err != nil {
				return err
			}
			defer func() {
				// Bounded: a collector that is gone must not hold up a shutdown.
				sctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
				defer cancel()
				_ = shutdownTracing(sctx)
			}()

			httpClient, err := agent.Client(cfg.TLS, 30*time.Second)
			if err != nil {
				return err
			}

			a, err := agent.New(cfg, httpClient, log)
			if err != nil {
				return err
			}

			ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
			defer stop()

			log.Info("agent starting",
				"node", cfg.NodeID, "server", cfg.Server,
				"fabrics", len(cfg.Fabrics), "nvml", discovery.NVMLSupported())

			err = a.Run(ctx)
			if errors.Is(err, context.Canceled) {
				return nil
			}
			return err
		},
	}
	cmd.Flags().StringVar(&configPath, "config", "", "configuration file")
	cmd.Flags().StringVar(&server, "server", "", "control plane address (overrides the config)")
	cmd.Flags().StringVar(&nodeID, "node-id", "", "node id (defaults to the hostname)")
	cmd.Flags().StringVar(&mgmtAddr, "mgmt-addr", "", "address engines are reachable on")
	cmd.AddCommand(newInventoryCmd())
	return cmd
}

func newInventoryCmd() *cobra.Command {
	var (
		asJSON  bool
		smiBin  string
		ibRoot  string
		skipGPU bool
		useNVML bool
	)

	cmd := &cobra.Command{
		Use:   "inventory",
		Short: "Show this node's GPUs and InfiniBand adapters",
		Long: "Reports what the scheduler would see on this node: GPU classes, health,\n" +
			"and the fabric each adapter is attached to.\n\n" +
			"Run it before adding a node to the cluster — it is the fastest way to find\n" +
			"a card another process is already using, or a port with no subnet manager.",
		RunE: func(cmd *cobra.Command, _ []string) error {
			inv, err := discovery.Discover(cmd.Context(), discovery.Options{
				NvidiaSMI:  smiBin,
				IBRoot:     ibRoot,
				PreferNVML: useNVML,
				SkipGPUs:   skipGPU,
			}, logger())
			if err != nil {
				// Not fatal: a node whose GPUs cannot be read still has IB
				// worth showing, and the error is itself the useful output.
				fmt.Fprintf(os.Stderr, "%v\n", err)
				inv.Source = "unavailable"
			}

			if asJSON {
				enc := json.NewEncoder(os.Stdout)
				enc.SetIndent("", "  ")
				return enc.Encode(inv)
			}
			printInventory(inv)
			return nil
		},
	}

	f := cmd.Flags()
	f.BoolVar(&asJSON, "json", false, "emit JSON")
	f.StringVar(&smiBin, "nvidia-smi", "", "path to nvidia-smi")
	f.StringVar(&ibRoot, "ib-root", discovery.SysfsRoot, "sysfs InfiniBand directory")
	f.BoolVar(&skipGPU, "no-gpu", false, "skip GPU discovery")
	f.BoolVar(&useNVML, "nvml", true, "prefer the NVML backend when this build has one (it is the only source of NVLink topology)")

	return cmd
}

func printInventory(inv discovery.Inventory) {
	w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)

	fmt.Fprintf(w, "GPUs (%d, via %s)\n", len(inv.GPUs), inv.Source)
	if len(inv.GPUs) == 0 {
		fmt.Fprintln(w, "  none")
	} else {
		fmt.Fprintln(w, "  IDX\tCLASS\tNAME\tMEM\tCC\tHEALTH")
		for _, g := range inv.GPUs {
			health := "ok"
			if !g.Health.OK() {
				health = g.Health.Problems[0]
			}
			fmt.Fprintf(w, "  %s\t%s\t%s\t%d MiB\t%s\t%s\n",
				discovery.DeviceLabel(g.Index), g.Class(), g.Name, g.MemoryTotalMB, g.ComputeCapability, health)
		}
		fmt.Fprintf(w, "  classes: %v\n", inv.Classes())
		for _, g := range inv.GPUs {
			if len(g.NVLinkPeers) > 0 {
				fmt.Fprintf(w, "  GPU %s NVLink peers: %v\n", discovery.DeviceLabel(g.Index), g.NVLinkPeers)
			}
		}
	}

	fmt.Fprintf(w, "\nInfiniBand adapters (%d)\n", len(inv.HCAs))
	if len(inv.HCAs) == 0 {
		// Worth stating plainly: such a node is its own single-node island and
		// can never host a multi-node replica.
		fmt.Fprintln(w, "  none — this node cannot host multi-node replicas")
	} else {
		fmt.Fprintln(w, "  HCA\tPORT\tSTATE\tRATE\tLINK\tGID PREFIX\tNUMA")
		for _, h := range inv.HCAs {
			for _, p := range h.Ports {
				fmt.Fprintf(w, "  %s\t%d\t%s\t%d Gb/s\t%s\t%s\t%d\n",
					h.Name, p.Num, p.State, p.RateGbps, p.LinkLayer, p.GIDPrefix, h.NUMANode)
			}
		}
		fmt.Fprintln(w, "\n  Fabric identity is not derived from the GID prefix by default:")
		fmt.Fprintln(w, "  OpenSM's default is fe80:: on every fabric, so a shared prefix proves")
		fmt.Fprintln(w, "  nothing. Label fabrics explicitly in agent.yaml (plan §5.2).")
	}

	_ = w.Flush()
}