llmgw/internal/cli/enginecheck.go

package cli

import (
	"context"
	"fmt"
	"os/exec"
	"sort"
	"strings"
	"time"

	"github.com/spf13/cobra"

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

func newEngineCheckCmd() *cobra.Command {
	var (
		help    string
		nodes   int
		tp      int
		perNode int
	)

	cmd := &cobra.Command{
		Use:   "engine-check <engine>",
		Short: "Check that an engine still accepts the flags llmgw builds for it",
		Long: "llmgw's launch arguments are read from each engine's documentation, and\n" +
			"engines rename and remove flags between versions. When that happens the\n" +
			"failure arrives minutes into a model load, in an engine's own error\n" +
			"message, on a node nobody was watching.\n\n" +
			"This asks the engine itself. It runs the engine's --help and checks every\n" +
			"flag llmgw would pass against it, which needs no GPU — a CLI prints its\n" +
			"help on any machine. Run it after upgrading an engine.\n\n" +
			"With --help-output it reads a saved --help instead of running anything,\n" +
			"so the check works from a laptop against help captured on a node.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			name := args[0]
			reg := engines.NewRegistry(engines.Config{})
			eng, err := reg.Get(name)
			if err != nil {
				return err
			}

			spec := representativeShard(name, nodes, tp, perNode)
			cmds, err := eng.Commands(spec)
			if err != nil {
				return fmt.Errorf("building a command for %s: %w", name, err)
			}

			text := help
			if text == "" {
				if text, err = runHelp(cmd.Context(), cmds); err != nil {
					return err
				}
			}

			// A truncated help produces confident nonsense: every flag below the
			// cut looks missing. Real help text documents --help itself, which
			// is the cheapest signal that we are looking at all of it.
			if !mentions(text, "--help") {
				cmd.PrintErrln("warning: this help text does not document --help, " +
					"which usually means it was truncated. Flags below the cut will " +
					"look missing.")
			}

			var unknown []string
			seen := map[string]bool{}
			for _, c := range cmds {
				for _, a := range c.Args {
					flag, _, _ := strings.Cut(a, "=")
					if !strings.HasPrefix(flag, "--") || seen[flag] {
						continue
					}
					seen[flag] = true
					if !mentions(text, flag) {
						unknown = append(unknown, flag)
					}
				}
			}
			sort.Strings(unknown)

			out := cmd.OutOrStdout()
			fmt.Fprintf(out, "%s: %d flags checked\n", name, len(seen))
			for _, c := range cmds {
				fmt.Fprintf(cmd.ErrOrStderr(), "  %s %s\n", c.Path, strings.Join(c.Args, " "))
			}
			if len(unknown) == 0 {
				fmt.Fprintln(out, "every flag llmgw builds is accepted by this engine")
				return nil
			}
			fmt.Fprintf(out, "\n%d flags this engine does not mention:\n", len(unknown))
			for _, f := range unknown {
				fmt.Fprintf(out, "  %s\n", f)
			}
			return fmt.Errorf("%s does not accept %d of the flags llmgw builds", name, len(unknown))
		},
	}

	f := cmd.Flags()
	f.StringVar(&help, "help-output", "", "a saved --help to check against, instead of running the engine")
	f.IntVar(&nodes, "nodes", 2, "shape the command for a replica of this many nodes")
	f.IntVar(&tp, "tp", 8, "tensor parallel degree")
	f.IntVar(&perNode, "gpus-per-node", 4, "GPUs each node contributes")
	return cmd
}

// representativeShard is the widest command llmgw would ever build, because a
// flag that only appears in the multi-node case is exactly the one nobody
// notices has been renamed.
func representativeShard(engine string, nodes, tp, perNode int) engines.ShardSpec {
	if nodes < 1 {
		nodes = 1
	}
	return engines.ShardSpec{
		ReplicaID: "check", ShardID: "check-r0", Rank: 0, NNodes: nodes,
		LeaderAddr: "10.0.0.1", LeaderPort: 20001,
		GPUIndices: seq(perNode), HCAs: []string{"mlx5_0"},
		MgmtInterface: "eth0", HTTPPort: 20000,
		Profile: engines.Profile{
			Name: "check", Engine: engine, TP: tp, Nodes: nodes, GPUsPerNode: perNode,
			ModelPath: "/models/x", ServedModelName: "x",
		},
	}
}

func seq(n int) []int {
	out := make([]int, n)
	for i := range out {
		out[i] = i
	}
	return out
}

// runHelp asks the engine to describe itself.
//
// The last command is the one that serves; the sidecars before it are a
// different program with a different vocabulary, so its help would be the wrong
// help to check against.
func runHelp(ctx context.Context, cmds []engines.Command) (string, error) {
	c := cmds[len(cmds)-1]

	// Everything up to the first flag, plus --help. Passing the real arguments
	// would start the engine.
	var args []string
	for _, a := range c.Args {
		if strings.HasPrefix(a, "-") {
			break
		}
		args = append(args, a)
	}
	args = append(args, "--help")

	ctx, cancel := context.WithTimeout(ctx, 2*time.Minute)
	defer cancel()

	out, err := exec.CommandContext(ctx, c.Path, args...).CombinedOutput()
	if err != nil && len(out) == 0 {
		return "", fmt.Errorf("running %s %s: %w\n"+
			"(is the engine installed here? use --help-output to check against a saved one)",
			c.Path, strings.Join(args, " "), err)
	}
	return string(out), nil
}

// mentions reports whether the help text documents a flag.
//
// Matched on a word boundary, so --port does not satisfy --port-range and a
// missing flag cannot hide behind a longer one that happens to contain it.
func mentions(help, flag string) bool {
	for i := 0; ; {
		j := strings.Index(help[i:], flag)
		if j < 0 {
			return false
		}
		end := i + j + len(flag)
		if end == len(help) || !isFlagChar(help[end]) {
			return true
		}
		i = end
	}
}

func isFlagChar(b byte) bool {
	return b == '-' || b == '_' || (b >= 'a' && b <= 'z') ||
		(b >= 'A' && b <= 'Z') || (b >= '0' && b <= '9')
}