llmgw/internal/agent/config.go

package agent

import (
	"errors"
	"fmt"
	"net"
	"net/url"
	"os"
	"slices"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// Config configures the agent.
type Config struct {
	// NodeID defaults to the hostname. It must equal the common name of the
	// agent's certificate, so set it only where the two differ.
	NodeID string `koanf:"node_id"`
	// Server is the control plane, and Servers is the same thing when there is
	// more than one. Both are accepted: every configuration written before
	// there could be a second control plane names one address, and none of
	// them should have to change.
	//
	// The order is tried in order, and a refusal moves to the next. Only the
	// leader accepts an agent, so a standby is skipped in the time it takes to
	// be told so — which is what makes failover a reconnect rather than an
	// edit to every node's configuration.
	Server  string   `koanf:"server"`
	Servers []string `koanf:"servers"`
	// Hostname defaults to the machine's.
	Hostname string `koanf:"hostname"`
	// MgmtAddr is the address engines are reached on. It defaults to this
	// machine's address on the route to the control plane, which is the one
	// the control plane and the router can reach.
	MgmtAddr string `koanf:"mgmt_addr"`
	// Labels are free-form facts about a node for a profile's node_selector.
	// Nothing needs them: memory, GPUs, engines and disk are measured, and a
	// label restating one of those goes stale the day the machine changes.
	Labels map[string]string `koanf:"labels"`

	// MgmtInterface is the ethernet interface NCCL bootstraps over. It
	// defaults to the interface holding MgmtAddr, then eth0.
	MgmtInterface string `koanf:"mgmt_interface"`
	// Fabrics labels this node's HCAs. It is explicit on purpose: OpenSM's
	// default subnet prefix is the same on every fabric, so nothing can be
	// inferred from the wire (plan §5.2).
	Fabrics []inventory.FabricAttachment `koanf:"fabrics"`

	HeartbeatInterval time.Duration `koanf:"heartbeat_interval"`
	// ReconnectBase is the first reconnect delay; it backs off from there.
	ReconnectBase time.Duration `koanf:"reconnect_base"`
	ReconnectMax  time.Duration `koanf:"reconnect_max"`

	// AdoptReadyTimeout bounds how long an adopted engine is given to answer a
	// health check before the agent stops waiting on it.
	//
	// Shorter than a startup timeout on purpose: this engine is already
	// running, so it is not loading a model, it is either serving or it is not.
	AdoptReadyTimeout time.Duration `koanf:"adopt_ready_timeout"`

	// TLS authenticates both ends of the control-plane stream. Leaving it unset
	// keeps the plaintext behaviour, which only belongs on a trusted network.
	TLS tlsconf.Config `koanf:"tls"`

	// ModelDirs are the directories this node keeps model weights in, beyond
	// WeightsCache, which is always looked in. Storage differs at every installation,
	// so llmgw observes it rather than owning it: a Lustre mount, an image
	// exported over the fabric, local NVMe and a copy made by hand all look
	// identical from here, and placement only needs to know whether the file
	// is present.
	ModelDirs []string `koanf:"model_dirs"`

	// EnginesDir is where the agent installs engines, one directory per
	// build. Defaults to /opt/llmgw/engines.
	EnginesDir string `koanf:"engines_dir"`

	// WeightsCache is where a remote source is downloaded to. Defaults to the
	// first entry of ModelDirs, because a node that says where its models live
	// has already answered this question, and otherwise to
	// /var/lib/llmgw-weights.
	//
	// Fetching is the simple arrangement: every node keeps its own copy. The
	// exotic ones need no code — weights on a shared filer or an image exported
	// over the fabric are simply present, and present means nothing is fetched.
	WeightsCache string `koanf:"weights_cache"`

	Discovery discovery.Options `koanf:"discovery"`
	Engines   engines.Config    `koanf:"engines"`
	Superv    supervisor.Config `koanf:"supervisor"`

	Version string
}

// DefaultConfig returns sensible defaults.
func DefaultConfig() Config {
	// No MgmtInterface: normalize finds the one holding the management
	// address, and falls back to eth0 only when it cannot. Defaulted here, it
	// was never found — every node said eth0 unless told otherwise, and a node
	// on ens3 or bond0 had NCCL bootstrap over an interface it does not have.
	return Config{
		HeartbeatInterval: 5 * time.Second,
		ReconnectBase:     time.Second,
		ReconnectMax:      30 * time.Second,
		AdoptReadyTimeout: 60 * time.Second,
		Superv:            supervisor.DefaultConfig(),
		// NVML gives NVLink topology, so prefer it when this build has it.
		Discovery: discovery.Options{PreferNVML: true},
	}
}

func (c *Config) normalize() error {
	d := DefaultConfig()
	if err := c.TLS.Validate(); err != nil {
		return err
	}
	if err := c.normalizeServerURL(); err != nil {
		return err
	}
	// Everything below is found rather than written. A file that has to say
	// which machine it is on, and at what address, is a file that differs on
	// every node — and a copy taken from one node to the next registers the
	// second as the first.
	host, _ := os.Hostname()
	if c.NodeID == "" {
		c.NodeID = host
	}
	if c.NodeID == "" {
		return errors.New("node_id is not set and the hostname could not be read")
	}
	if c.Hostname == "" {
		c.Hostname = host
	}
	if c.MgmtAddr == "" {
		c.MgmtAddr = addrToward(c.Server)
	}
	if c.MgmtInterface == "" {
		c.MgmtInterface = interfaceOf(c.MgmtAddr)
	}
	if c.MgmtInterface == "" {
		c.MgmtInterface = "eth0"
	}
	// The directory downloads land in is always one that is looked in. When
	// they were two settings, a node could fill its disk with fetched weights
	// that placement never saw — which is what every node here was doing.
	if c.WeightsCache == "" && len(c.ModelDirs) > 0 {
		c.WeightsCache = c.ModelDirs[0]
	}
	if c.WeightsCache == "" {
		c.WeightsCache = DefaultWeightsCache
	}
	if !slices.Contains(c.ModelDirs, c.WeightsCache) {
		c.ModelDirs = append(c.ModelDirs, c.WeightsCache)
	}
	if c.HeartbeatInterval <= 0 {
		c.HeartbeatInterval = d.HeartbeatInterval
	}
	if c.ReconnectBase <= 0 {
		c.ReconnectBase = d.ReconnectBase
	}
	if c.AdoptReadyTimeout <= 0 {
		c.AdoptReadyTimeout = d.AdoptReadyTimeout
	}
	if c.ReconnectMax < c.ReconnectBase {
		c.ReconnectMax = max(c.ReconnectBase, d.ReconnectMax)
	}
	return nil
}

// normalizeServerURL makes the configured address a URL connect can use, and
// refuses the combination that would quietly undo the certificates: a scheme
// that says plaintext while TLS is configured. Silently connecting anyway would
// mean a cluster believing it has mutual authentication when it has none.
func (c *Config) normalizeServerURL() error {
	// One list, however it was written. Server first when both are given,
	// because an operator adding a standby to an existing file expects the
	// one already there to stay preferred.
	all := make([]string, 0, 1+len(c.Servers))
	if c.Server != "" {
		all = append(all, c.Server)
	}
	all = append(all, c.Servers...)

	seen := map[string]struct{}{}
	out := make([]string, 0, len(all))
	for _, s := range all {
		fixed, err := normalizeOne(s, c.TLS.Enabled())
		if err != nil {
			return err
		}
		// A duplicate would be tried twice in a row on failover, which reads
		// as the rotation not working.
		if _, dup := seen[fixed]; dup {
			continue
		}
		seen[fixed] = struct{}{}
		out = append(out, fixed)
	}
	if len(out) == 0 {
		return errors.New("server is required")
	}
	c.Servers = out
	c.Server = out[0]
	return nil
}

// DefaultWeightsCache is where fetched weights go when nothing says otherwise.
const DefaultWeightsCache = "/var/lib/llmgw-weights"

// addrToward is this machine's address on the route to the control plane.
//
// The address engines are reached on is whichever one the control plane can
// reach, and the kernel already knows which that is: a UDP "connection" picks
// the source address for the route without sending a packet. Asking the
// interfaces instead would have to guess among docker bridges, VPNs and
// link-local addresses.
func addrToward(server string) string {
	u, err := url.Parse(server)
	if err != nil || u.Host == "" {
		return ""
	}
	host := u.Hostname()
	port := u.Port()
	if port == "" {
		port = "443"
	}
	conn, err := net.Dial("udp", net.JoinHostPort(host, port))
	if err != nil {
		return ""
	}
	defer conn.Close()
	if a, ok := conn.LocalAddr().(*net.UDPAddr); ok {
		return a.IP.String()
	}
	return ""
}

// interfaceOf names the interface that holds an address, or "" when none does.
func interfaceOf(addr string) string {
	ip := net.ParseIP(addr)
	if ip == nil {
		return ""
	}
	ifaces, err := net.Interfaces()
	if err != nil {
		return ""
	}
	for _, ifc := range ifaces {
		addrs, err := ifc.Addrs()
		if err != nil {
			continue
		}
		for _, a := range addrs {
			if n, ok := a.(*net.IPNet); ok && n.IP.Equal(ip) {
				return ifc.Name
			}
		}
	}
	return ""
}

// normalizeOne applies the scheme rules to a single address.
func normalizeOne(addr string, secure bool) (string, error) {
	switch {
	case strings.HasPrefix(addr, "https://"):
		if !secure {
			return "", fmt.Errorf("%s is an https:// address but no tls certificates are configured", addr)
		}
	case strings.HasPrefix(addr, "http://"):
		if secure {
			return "", fmt.Errorf("tls is configured but %s is an http:// address; use https://", addr)
		}
	case strings.Contains(addr, "://"):
		return "", fmt.Errorf("server %q must be an http:// or https:// address", addr)
	case secure:
		return "https://" + addr, nil
	default:
		return "http://" + addr, nil
	}
	return addr, nil
}