llmgw/internal/config/config.go

// Package config loads llmgw's YAML configuration.
//
// Phase 1 only needs the router's half of it: the endpoint list is static here
// and will come from the control plane in phase 3 (plan §6.4). The shape is
// already the one the plan describes, so that later change stays additive.
package config

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/knadh/koanf/parsers/yaml"
	env "github.com/knadh/koanf/providers/env/v2"
	"github.com/knadh/koanf/providers/file"
	"github.com/knadh/koanf/v2"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// Server is the whole server configuration.
type Server struct {
	// Listen is the client-facing OpenAI API address.
	Listen string `koanf:"listen"`
	// AdminListen serves metrics, health and pprof. Kept separate so the
	// client-facing port can be exposed without exposing diagnostics.
	AdminListen string `koanf:"admin_listen"`

	// AgentListen is where node agents connect. Separate from the client port
	// so the fleet-facing surface can be firewalled independently.
	AgentListen string `koanf:"agent_listen"`

	// Auth turns on API keys on the client-facing port.
	Auth Auth `koanf:"auth"`

	// Landing configures the page served at "/".
	Landing Landing `koanf:"landing"`

	// AgentTLS authenticates the fleet-facing port. It deliberately does not
	// cover the client-facing one: that surface is fronted by a publicly
	// trusted certificate, while this one is mutual and private.
	AgentTLS tlsconf.Config `koanf:"agent_tls"`

	Router  Router          `koanf:"router"`
	Control Control         `koanf:"control"`
	Models  []catalog.Model `koanf:"models"`

	// Endpoints is the phase-1 static replica list. It is only consulted when
	// no model declares a profile, which is what makes a router-only smoke test
	// against sim-engines possible without a control plane.
	Endpoints []Endpoint `koanf:"endpoints"`
}

// Landing configures the public page at "/".
//
// It exists for one thing the gateway cannot know: who to ask for a key. A page
// that explains the API and then leaves the reader with no route to using it is
// where onboarding stops.
type Landing struct {
	// Contact is where someone without a key should go — an address, a URL, a
	// channel name. Rendered as a link when it is one, as text otherwise.
	Contact string `koanf:"contact"`
	// Notice is an optional line of your own, for whatever the reader needs to
	// know that this page cannot guess.
	Notice string `koanf:"notice"`
}

// Auth configures client authentication (plan §7.9).
//
// Off by default, because a gateway on a trusted network with no keys issued
// still has to serve. Turning it on is what makes the client-facing port safe
// to expose.
type Auth struct {
	Enabled bool `koanf:"enabled"`
	// Store is the SQLite file holding keys, tenancy and usage.
	Store string `koanf:"store"`
	// CacheTTL bounds how long a revoked key keeps working, traded against a
	// database read per key per TTL.
	CacheTTL time.Duration `koanf:"cache_ttl"`
	// NegativeTTL is how long a rejected key is remembered as rejected, so a
	// client looping with a bad key cannot become a database read per request.
	NegativeTTL time.Duration `koanf:"negative_ttl"`

	// DefaultRPM and DefaultTPM apply to a key that no scope gives a limit to.
	//
	// Unlimited used to be the default, and a project created without limits is
	// the easy path — the "default" project on this cluster still is one. The
	// first sign of that is the bill or the outage. Zero keeps the old
	// behaviour for anyone who wants it, but it should be a choice rather than
	// an omission. A fallback, not a ceiling: an explicit limit at any scope
	// wins, above or below.
	DefaultRPM int `koanf:"default_rpm"`
	DefaultTPM int `koanf:"default_tpm"`

	// SignupOpen serves the signup form. Accounts it creates are pending until
	// an operator approves them, so this is not a way in — but turning it off
	// means accounts are made with `llmgw users create` instead.
	SignupOpen bool `koanf:"signup_open"`

	// Pricing is what each model costs, for budgets. A model absent from the
	// table costs whatever DefaultPrice says, which is nothing unless set.
	Pricing      map[string]Price `koanf:"pricing"`
	DefaultPrice Price            `koanf:"default_price"`

	// FlushInterval and BatchSize decide when usage records are written.
	FlushInterval time.Duration `koanf:"flush_interval"`
	BatchSize     int           `koanf:"batch_size"`
	// BufferSize bounds the queue. Beyond it records are dropped rather than
	// blocking a request: losing accounting is bad, turning a stalled database
	// into an outage is worse.
	BufferSize int `koanf:"buffer_size"`
}

// Price is a model's cost in micro-units per thousand tokens.
//
// What a unit means is the operator's choice — self-hosted models have no
// per-token price, so it may be GPU-seconds, a notional rate, or 1000
// everywhere, which turns a budget into a token quota (plan §7.9.4).
type Price struct {
	PromptPer1k     int64 `koanf:"prompt_per_1k"`
	CompletionPer1k int64 `koanf:"completion_per_1k"`
}

// Control configures the control plane.
type Control struct {
	// Interval is how often desired and actual state are compared.
	Interval time.Duration `koanf:"interval"`
	// BaseBackoff and MaxBackoff bound how fast a failing model is retried.
	BaseBackoff time.Duration `koanf:"base_backoff"`
	MaxBackoff  time.Duration `koanf:"max_backoff"`
	// UnreachableAfter and LostAfter stage node expiry. They are deliberately
	// far apart: a network blip must not tear down replicas that still serve,
	// and rebuilding a multi-node replica costs tens of minutes of loading.
	UnreachableAfter time.Duration `koanf:"unreachable_after"`
	LostAfter        time.Duration `koanf:"lost_after"`
	// PortRange is the per-node range for engine HTTP and dist-init ports.
	PortRange []int `koanf:"port_range"`
	// MgmtInterface is the ethernet interface NCCL bootstraps over.
	MgmtInterface string `koanf:"mgmt_interface"`
}

// Router holds routing behaviour.
type Router struct {
	ScrapeInterval time.Duration `koanf:"scrape_interval"`
	ScrapeTimeout  time.Duration `koanf:"scrape_timeout"`
	// StaleAfter is how long an endpoint may go unscraped before being scored
	// down (plan §7.6).
	StaleAfter time.Duration `koanf:"stale_after"`

	// Retries is how many *additional* endpoints may be tried after a failure,
	// and only while nothing has been written to the client (plan §7.7).
	Retries int `koanf:"retries"`
	// Cooldown is how long a failing endpoint leaves the candidate pool.
	Cooldown time.Duration `koanf:"cooldown"`
	// FailureThreshold is the consecutive failure count that opens a breaker.
	// 1 turns the breaker into a pure cooldown.
	FailureThreshold int `koanf:"failure_threshold"`

	// DefaultProfile names the routing profile used when a model picks none.
	DefaultProfile string `koanf:"default_profile"`
	// Prefix turns on the approximate prefix tree (plan §7.5 tier 1). Off, the
	// router uses session affinity alone, which is tier 0 and already most of
	// the win for multi-turn traffic.
	Prefix Prefix `koanf:"prefix"`

	// SessionHeader carries the session id for affinity (plan §7.5 tier 0).
	SessionHeader string `koanf:"session_header"`
	// InjectUsage asks engines for usage even when the client did not, so token
	// accounting works regardless of client behaviour.
	InjectUsage bool `koanf:"inject_usage"`

	// MaxBodyBytes bounds the request body the router will read.
	MaxBodyBytes int64 `koanf:"max_body_bytes"`

	Timeouts  Timeouts  `koanf:"timeouts"`
	Scoring   Scoring   `koanf:"scoring"`
	Admission Admission `koanf:"admission"`
}

// Admission configures the flow control layer (plan §7.2).
type Admission struct {
	// Enabled turns the layer on. Off, every request goes straight through,
	// which is the v0.2.0 behaviour.
	Enabled bool `koanf:"enabled"`

	Bands       []Band `koanf:"bands"`
	DefaultBand string `koanf:"default_band"`
	// BandHeader lets a client declare its band. Until authentication lands in
	// v1.0.0 this is trusted, so the port must not be exposed to callers who
	// should not be able to promote themselves.
	BandHeader string `koanf:"band_header"`

	QueueTTL             time.Duration `koanf:"queue_ttl"`
	QueueCapacityPerBand int           `koanf:"queue_capacity_per_band"`
	Interval             time.Duration `koanf:"interval"`
}

// Prefix configures approximate prefix routing.
type Prefix struct {
	Enabled bool `koanf:"enabled"`
	// MaxNodes bounds the tree, MinMatch is the shortest match worth acting on
	// in bytes, and TTL is how long an entry is believed — an engine evicts on
	// its own schedule, so an old record may describe a cache that is gone.
	MaxNodes int           `koanf:"max_nodes"`
	MinMatch int           `koanf:"min_match"`
	TTL      time.Duration `koanf:"ttl"`

	// MatchRatio is the share of the prompt that must match to be worth acting
	// on. BalanceAbs and BalanceRel are the imbalance guard, and both must trip
	// before affinity is abandoned.
	MatchRatio float64 `koanf:"match_ratio"`
	BalanceAbs float64 `koanf:"balance_abs"`
	BalanceRel float64 `koanf:"balance_rel"`
}

// Band is one priority class.
type Band struct {
	Name              string  `koanf:"name"`
	Priority          int     `koanf:"priority"`
	SaturationCeiling float64 `koanf:"saturation_ceiling"`
}

// Timeouts bounds each phase of an upstream request.
type Timeouts struct {
	Dial time.Duration `koanf:"dial"`
	// ResponseHeader bounds the wait for headers on a streaming request. It is
	// a liveness check there, because headers arrive before generation starts.
	ResponseHeader time.Duration `koanf:"response_header"`
	// Generation bounds a non-streaming request, where the engine sends no
	// headers until the completion is done. Setting response_header alone used
	// to cap both, so 60s read as a liveness guard and behaved as a limit of
	// about a hundred tokens on a slow model.
	Generation time.Duration `koanf:"generation"`
	FirstToken time.Duration `koanf:"first_token"`
	Idle       time.Duration `koanf:"idle"`
	Total      time.Duration `koanf:"total"`
}

// Scoring holds the thresholds the scorers normalize against. These want to be
// measured with ctl bench (plan §10), not guessed.
type Scoring struct {
	TargetWaiting     float64 `koanf:"target_waiting"`
	TargetInflight    float64 `koanf:"target_inflight"`
	KVCeiling         float64 `koanf:"kv_ceiling"`
	SessionLoadFactor float64 `koanf:"session_load_factor"`
}

// Endpoint is one engine replica. Static in phase 1.
type Endpoint struct {
	ID     string   `koanf:"id"`
	URL    string   `koanf:"url"`
	Models []string `koanf:"models"`
	// Engine names the inference engine, which decides what its metrics are
	// called. Empty means vLLM's names.
	Engine   string  `koanf:"engine"`
	Pool     string  `koanf:"pool"`
	Weight   float64 `koanf:"weight"`
	Priority int     `koanf:"priority"`
}

// Default returns the configuration used when nothing is specified.
func Default() Server {
	return Server{
		Listen:      ":8080",
		AdminListen: ":9090",
		AgentListen: ":7443",
		Control: Control{
			Interval:         5 * time.Second,
			BaseBackoff:      10 * time.Second,
			MaxBackoff:       10 * time.Minute,
			UnreachableAfter: 15 * time.Second,
			LostAfter:        60 * time.Second,
			PortRange:        []int{20000, 20999},
			MgmtInterface:    "eth0",
		},
		Router: Router{
			ScrapeInterval:   2 * time.Second,
			ScrapeTimeout:    time.Second,
			StaleAfter:       10 * time.Second,
			Retries:          2,
			Cooldown:         15 * time.Second,
			FailureThreshold: 3,
			DefaultProfile:   "chat",
			SessionHeader:    "x-session-id",
			InjectUsage:      true,
			MaxBodyBytes:     32 << 20,
			Timeouts: Timeouts{
				Dial:           5 * time.Second,
				ResponseHeader: 60 * time.Second,
				FirstToken:     120 * time.Second,
				Idle:           60 * time.Second,
				Total:          30 * time.Minute,
			},
			Scoring: Scoring{
				TargetWaiting:     4,
				TargetInflight:    8,
				KVCeiling:         0.9,
				SessionLoadFactor: 1.25,
			},
			Admission: Admission{
				Enabled:     true,
				BandHeader:  "x-llmgw-band",
				DefaultBand: "gold",
				Bands: []Band{
					{Name: "platinum", Priority: 400, SaturationCeiling: 1.00},
					{Name: "gold", Priority: 300, SaturationCeiling: 0.85},
					{Name: "silver", Priority: 200, SaturationCeiling: 0.60},
					{Name: "bronze", Priority: 100, SaturationCeiling: 0.30},
				},
				QueueTTL:             30 * time.Second,
				QueueCapacityPerBand: 2000,
				Interval:             10 * time.Millisecond,
			},
		},
	}
}

// Load reads configuration from path, overlaying it on the defaults, then
// applies LLMGW_-prefixed environment overrides.
//
// An empty path loads the defaults, which is what makes `llmgw server` with a
// couple of --endpoint flags usable for a smoke test.
func Load(path string) (Server, error) {
	cfg := Default()
	k := koanf.New(".")

	if path != "" {
		if _, err := os.Stat(path); err != nil {
			return cfg, fmt.Errorf("config file: %w", err)
		}
		if err := k.Load(file.Provider(path), yaml.Parser()); err != nil {
			return cfg, fmt.Errorf("parsing %s: %w", path, err)
		}
	}

	// LLMGW_ROUTER__RETRIES=3 maps onto router.retries.
	err := k.Load(env.Provider(".", env.Opt{
		Prefix: "LLMGW_",
		TransformFunc: func(key, value string) (string, any) {
			key = strings.ToLower(strings.TrimPrefix(key, "LLMGW_"))
			return strings.ReplaceAll(key, "__", "."), value
		},
	}), nil)
	if err != nil {
		return cfg, fmt.Errorf("environment overrides: %w", err)
	}

	if err := k.UnmarshalWithConf("", &cfg, koanf.UnmarshalConf{Tag: "koanf"}); err != nil {
		return cfg, fmt.Errorf("decoding configuration: %w", err)
	}
	return cfg, cfg.Validate()
}

// Validate rejects configurations that would otherwise fail confusingly at
// request time.
func (s Server) Validate() error {
	if s.Listen == "" {
		return fmt.Errorf("listen is empty")
	}
	if err := s.AgentTLS.Validate(); err != nil {
		return err
	}
	if s.Auth.Enabled && s.Auth.Store == "" {
		return fmt.Errorf("auth.store is required when auth is enabled")
	}

	seen := make(map[string]struct{}, len(s.Endpoints))
	for i, e := range s.Endpoints {
		switch {
		case e.ID == "":
			return fmt.Errorf("endpoints[%d]: id is required", i)
		case e.URL == "":
			return fmt.Errorf("endpoint %q: url is required", e.ID)
		case !strings.HasPrefix(e.URL, "http://") && !strings.HasPrefix(e.URL, "https://"):
			return fmt.Errorf("endpoint %q: url must start with http:// or https://", e.ID)
		case len(e.Models) == 0:
			return fmt.Errorf("endpoint %q: at least one model is required", e.ID)
		}
		if _, dup := seen[e.ID]; dup {
			return fmt.Errorf("endpoint id %q is used twice", e.ID)
		}
		seen[e.ID] = struct{}{}
	}

	// The catalog enforces the rest — aliases, fallbacks, parallelism
	// arithmetic — so there is one place those rules live.
	if _, err := s.Catalog(); err != nil {
		return err
	}
	return nil
}

// Catalog builds the model catalog from the configuration.
func (s Server) Catalog() (*catalog.Catalog, error) {
	return catalog.New(s.Models)
}

// ControlPlaneMode reports whether the server should run a control plane.
//
// A model with no profiles cannot be placed, so a configuration of bare names
// plus a static endpoint list is the router-only mode used for testing against
// sim-engines.
func (s Server) ControlPlaneMode() bool {
	for _, m := range s.Models {
		if len(m.Profiles) > 0 {
			return true
		}
	}
	return false
}