llmgw/internal/cli/agent.go

package cli

import (
	"encoding/json"
	"fmt"
	"os"
	"text/tabwriter"

	"github.com/spf13/cobra"

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

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

	cmd := &cobra.Command{
		Use:   "agent",
		Short: "Node agent: discovery and engine supervision (phase 2)",
		Long: "Runs on every GPU node. It discovers the hardware, starts and supervises\n" +
			"inference engines, and reports to the control plane.\n\n" +
			"The control plane connection is not built yet, so for now the useful\n" +
			"subcommand is `agent inventory`.",
		RunE: func(_ *cobra.Command, _ []string) error {
			return errNotImplemented
		},
	}
	cmd.Flags().StringVar(&configPath, "config", "agent.yaml", "configuration file")
	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, "  %d\t%s\t%s\t%d MiB\t%s\t%s\n",
				g.Index, g.Class(), g.Name, g.MemoryTotalMB, g.ComputeCapability, health)
		}
		fmt.Fprintf(w, "  classes: %v\n", inv.Classes())
		if !discovery.NVMLSupported() {
			fmt.Fprintln(w, "  NVLink topology unavailable: rebuild with -tags nvml.")
			fmt.Fprintln(w, "  Without it, placement cannot keep a tensor-parallel group inside")
			fmt.Fprintln(w, "  one NVLink domain (plan §6.2).")
		}
	}

	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()
}