package worker import ( "bytes" "context" "crypto/rand" "encoding/hex" "encoding/json" "errors" "fmt" "io" "math" "net/http" "net/url" "sort" "strconv" "strings" "sync" "sync/atomic" "time" "github.com/example/ollama-fair-gateway/internal/config" "github.com/example/ollama-fair-gateway/internal/hoststats" ) var ( ErrModelPlacementBlocked = errors.New("model blocked by placement policy") ErrModelNotInstalled = errors.New("model not installed on eligible worker") ) type state struct { cfg config.WorkerConfig url *url.URL active atomic.Int64 healthy atomic.Bool mu sync.RWMutex models map[string]bool installed map[string]bool installedKnown bool installedModels []string inventoryError string modelActive map[string]int modelMaintenance map[string]bool metadata map[string]metadataEntry performance map[string]performanceState baselinePlacement config.ModelPlacementRule baselineModelConcurrency map[string]int placement config.ModelPlacementRule placementOverride bool loadedModels []LoadedModel telemetry ResourceTelemetry lastError string lastCheck time.Time maintenance string circuitState string circuitFailures int circuitOpenUntil time.Time halfOpenInFlight bool lastCircuitError string } type LoadedModel struct { Name string `json:"name"` Model string `json:"model,omitempty"` Size int64 `json:"size,omitempty"` SizeVRAM int64 `json:"size_vram,omitempty"` ContextLength int64 `json:"context_length,omitempty"` ExpiresAt time.Time `json:"expires_at,omitempty"` } type ResourceTelemetry struct { MemoryUsedBytes int64 `json:"memory_used_bytes,omitempty"` MemoryTotalBytes int64 `json:"memory_total_bytes,omitempty"` VRAMUsedBytes int64 `json:"vram_used_bytes,omitempty"` VRAMTotalBytes int64 `json:"vram_total_bytes,omitempty"` GPUUtilizationPct float64 `json:"gpu_utilization_percent,omitempty"` GPUTemperatureC float64 `json:"gpu_temperature_c,omitempty"` GPUPowerWatts float64 `json:"gpu_power_watts,omitempty"` Source string `json:"source,omitempty"` UpdatedAt time.Time `json:"updated_at,omitempty"` Error string `json:"error,omitempty"` } // externalTelemetry uses pointers so an omitted field does not accidentally // overwrite a previously collected local value with the JSON zero value. type externalTelemetry struct { MemoryUsedBytes *int64 `json:"memory_used_bytes"` MemoryTotalBytes *int64 `json:"memory_total_bytes"` VRAMUsedBytes *int64 `json:"vram_used_bytes"` VRAMTotalBytes *int64 `json:"vram_total_bytes"` GPUUtilizationPct *float64 `json:"gpu_utilization_percent"` GPUTemperatureC *float64 `json:"gpu_temperature_c"` GPUPowerWatts *float64 `json:"gpu_power_watts"` Source string `json:"source,omitempty"` UpdatedAt *time.Time `json:"updated_at,omitempty"` Error string `json:"error,omitempty"` } type ModelMetadata struct { Model string `json:"model"` Capabilities []string `json:"capabilities,omitempty"` ContextLength int64 `json:"context_length,omitempty"` ConfiguredContextLength int64 `json:"configured_context_length,omitempty"` Details ModelDetails `json:"details,omitempty"` UpdatedAt time.Time `json:"updated_at,omitempty"` Error string `json:"error,omitempty"` } type ContextWindow struct { Worker string `json:"worker"` Model string `json:"model"` ModelMaxTokens int64 `json:"model_max_tokens,omitempty"` ConfiguredTokens int64 `json:"configured_tokens,omitempty"` LoadedTokens int64 `json:"loaded_tokens,omitempty"` WorkerDefaultTokens int64 `json:"worker_default_tokens,omitempty"` WorkerLimitTokens int64 `json:"worker_limit_tokens,omitempty"` EffectiveTokens int64 `json:"effective_tokens,omitempty"` EffectiveSource string `json:"effective_source,omitempty"` MetadataError string `json:"metadata_error,omitempty"` } type metadataEntry struct { Data ModelMetadata FetchedAt time.Time } type ModelPerformance struct { Model string `json:"model"` PromptTPS float64 `json:"prompt_tps,omitempty"` OutputTPS float64 `json:"output_tps,omitempty"` Samples int64 `json:"samples"` } type performanceState struct { PromptTPS float64 OutputTPS float64 Samples int64 } type Snapshot struct { Name string `json:"name"` URL string `json:"url"` Healthy bool `json:"healthy"` Active int64 `json:"active"` MaxConcurrent int `json:"max_concurrent"` Models []string `json:"models"` LoadedModels []LoadedModel `json:"loaded_models,omitempty"` ResidentBytes int64 `json:"resident_bytes,omitempty"` VRAMBytes int64 `json:"vram_bytes,omitempty"` MemoryCapacityBytes int64 `json:"memory_capacity_bytes,omitempty"` VRAMCapacityBytes int64 `json:"vram_capacity_bytes,omitempty"` MemoryUsedBytes int64 `json:"memory_used_bytes,omitempty"` MemoryTotalBytes int64 `json:"memory_total_bytes,omitempty"` VRAMUsedBytes int64 `json:"vram_used_bytes,omitempty"` VRAMTotalBytes int64 `json:"vram_total_bytes,omitempty"` GPUUtilizationPct float64 `json:"gpu_utilization_percent,omitempty"` GPUTemperatureC float64 `json:"gpu_temperature_c,omitempty"` GPUPowerWatts float64 `json:"gpu_power_watts,omitempty"` ModelActive map[string]int `json:"model_active,omitempty"` ModelLimits map[string]int `json:"model_limits,omitempty"` Performance []ModelPerformance `json:"performance,omitempty"` ModelPlacement config.ModelPlacementRule `json:"model_placement"` PlacementOverride bool `json:"placement_override,omitempty"` Maintenance string `json:"maintenance"` AcceptingNew bool `json:"accepting_new"` CircuitState string `json:"circuit_state"` CircuitFailures int `json:"circuit_failures,omitempty"` CircuitOpenUntil time.Time `json:"circuit_open_until,omitempty"` LastCircuitError string `json:"last_circuit_error,omitempty"` TelemetrySource string `json:"telemetry_source,omitempty"` TelemetryError string `json:"telemetry_error,omitempty"` Labels map[string]string `json:"labels,omitempty"` LastError string `json:"last_error,omitempty"` LastCheck time.Time `json:"last_check"` } type PlacementDecision struct { Allowed bool `json:"allowed"` Source string `json:"source"` Pattern string `json:"pattern,omitempty"` ExactOverride bool `json:"exact_override,omitempty"` } type PlacementSnapshot struct { Worker string `json:"worker"` Baseline config.ModelPlacementRule `json:"baseline"` Effective config.ModelPlacementRule `json:"effective"` Override bool `json:"override"` InventoryKnown bool `json:"inventory_known"` InstalledModels []string `json:"installed_models,omitempty"` LoadedModels []string `json:"loaded_models,omitempty"` InventoryError string `json:"inventory_error,omitempty"` } type Pool struct { workers []*state byName map[string]*state client *http.Client notify chan struct{} control string routing config.RoutingConfig capCfg config.ModelCapabilitiesConfig reliability config.ReliabilityConfig } type Lease struct { State *state model string once sync.Once pool *Pool } func (l *Lease) Name() string { return l.State.cfg.Name } func (l *Lease) URL() *url.URL { return l.State.url } func (l *Lease) Release() { if l == nil { return } l.once.Do(func() { l.State.abandonCircuitProbe() l.State.active.Add(-1) if l.model != "" { l.State.mu.Lock() if l.State.modelActive[l.model] > 1 { l.State.modelActive[l.model]-- } else { delete(l.State.modelActive, l.model) } l.State.mu.Unlock() } l.pool.signal() }) } func New(cfgs []config.WorkerConfig, control string) *Pool { p := &Pool{byName: map[string]*state{}, notify: make(chan struct{}, 1), control: control, client: &http.Client{Timeout: 4 * time.Second}} p.SetRoutingConfig(config.RoutingConfig{}) p.SetModelCapabilitiesConfig(config.ModelCapabilitiesConfig{}) for _, c := range cfgs { u, _ := url.Parse(strings.TrimRight(c.URL, "/")) placement := normalizePlacementRule(c.ModelPlacement) baselineConc := map[string]int{} for k, v := range c.ModelConcurrency { baselineConc[k] = v } s := &state{cfg: c, url: u, models: map[string]bool{}, installed: map[string]bool{}, modelActive: map[string]int{}, modelMaintenance: map[string]bool{}, metadata: map[string]metadataEntry{}, performance: map[string]performanceState{}, baselinePlacement: placement, baselineModelConcurrency: baselineConc, placement: placement, maintenance: "active", circuitState: "closed"} s.healthy.Store(true) p.workers = append(p.workers, s) p.byName[c.Name] = s } return p } func (p *Pool) SetRoutingConfig(c config.RoutingConfig) { if c.LoadedBonus == 0 { c.LoadedBonus = 60 } if c.InstalledBonus == 0 { c.InstalledBonus = 30 } if c.ThroughputBonus == 0 { c.ThroughputBonus = 20 } if c.VRAMPressurePenalty == 0 { c.VRAMPressurePenalty = 35 } if c.GPUUtilizationPenalty == 0 { c.GPUUtilizationPenalty = 10 } if c.AvoidVRAMPercent == 0 { c.AvoidVRAMPercent = 97 } p.routing = c } func (p *Pool) SetModelCapabilitiesConfig(c config.ModelCapabilitiesConfig) { if c.Mode == "" { c.Mode = "enforce" } if c.CacheTTL == 0 { c.CacheTTL = config.Duration(10 * time.Minute) } if c.ContextGuard == "" { c.ContextGuard = "reject" } if c.Context.MaxRequestedTokens == 0 { c.Context.MaxRequestedTokens = 32768 } if c.Context.DefaultWorkerTokens == 0 { c.Context.DefaultWorkerTokens = 4096 } if c.Context.EstimationMarginPercent == 0 { c.Context.EstimationMarginPercent = 15 } if c.Context.VisionReserveTokensPerImage == 0 { c.Context.VisionReserveTokensPerImage = 2048 } p.capCfg = c } func (p *Pool) SetReliabilityConfig(c config.ReliabilityConfig) { if c.FailureThreshold <= 0 { c.FailureThreshold = 3 } if c.OpenDuration == 0 { c.OpenDuration = config.Duration(30 * time.Second) } p.reliability = c } func (p *Pool) Start(ctx context.Context) { for _, w := range p.workers { p.refresh(ctx, w) go p.healthLoop(ctx, w) } } func (p *Pool) Health(context.Context) error { for _, w := range p.workers { if w.healthy.Load() { return nil } } return fmt.Errorf("no healthy Ollama workers") } func (p *Pool) signal() { select { case p.notify <- struct{}{}: default: } } func (p *Pool) Acquire(ctx context.Context, model string) (*Lease, error) { return p.AcquireExcluding(ctx, model, nil) } func (p *Pool) AcquireExcluding(ctx context.Context, model string, excluded map[string]bool) (*Lease, error) { return p.AcquireAllowedExcluding(ctx, model, nil, excluded, 0) } // AcquireAllowed limits worker selection to an admission-approved set. A nil // set preserves the legacy behavior. requestedContext is used only to avoid a // misleading loaded-model bonus when an explicit num_ctx would require Ollama // to reload the model with a larger KV cache. func (p *Pool) AcquireAllowed(ctx context.Context, model string, allowed map[string]bool, requestedContext int64) (*Lease, error) { return p.AcquireAllowedExcluding(ctx, model, allowed, nil, requestedContext) } func (p *Pool) AcquireAllowedExcluding(ctx context.Context, model string, allowed, excluded map[string]bool, requestedContext int64) (*Lease, error) { for { candidates := p.candidatesExcludingAllowed(model, excluded, allowed, requestedContext) for _, w := range candidates { for { cur := w.active.Load() if cur >= int64(w.cfg.MaxConcurrent) { break } if !w.active.CompareAndSwap(cur, cur+1) { continue } if !w.claimCircuitProbe(p.reliability) { w.active.Add(-1) continue } if model != "" && !w.acquireModelSlot(model) { w.abandonCircuitProbe() w.active.Add(-1) break } return &Lease{State: w, model: canonicalModel(model), pool: p}, nil } } if len(candidates) == 0 { return nil, p.noCandidateError(model) } select { case <-ctx.Done(): return nil, ctx.Err() case <-p.notify: case <-time.After(50 * time.Millisecond): } } } func (p *Pool) Control() (*url.URL, string, error) { if p.control != "" { if w := p.byName[p.control]; w != nil && w.healthy.Load() { return w.url, w.cfg.Name, nil } } for _, w := range p.workers { if w.healthy.Load() { return w.url, w.cfg.Name, nil } } return nil, "", fmt.Errorf("no healthy control worker") } func (p *Pool) ControlForModel(model string) (*url.URL, string, error) { if model != "" { candidates := p.candidates(model) if len(candidates) == 0 { return nil, "", p.noCandidateError(model) } return candidates[0].url, candidates[0].cfg.Name, nil } return p.Control() } func (p *Pool) noCandidateError(model string) error { healthy, allowed, knownInventories := 0, 0, 0 for _, w := range p.workers { if !w.healthy.Load() { continue } healthy++ if model != "" && !w.placementDecision(model).Allowed { continue } allowed++ w.mu.RLock() if w.installedKnown { knownInventories++ } w.mu.RUnlock() } if healthy == 0 { return fmt.Errorf("no healthy Ollama workers") } if model != "" && allowed == 0 { return fmt.Errorf("%w: model %q is blocked on all healthy workers", ErrModelPlacementBlocked, model) } if model != "" && allowed > 0 && knownInventories == allowed { return fmt.Errorf("%w: model %q is not installed on any eligible worker", ErrModelNotInstalled, model) } return fmt.Errorf("no eligible Ollama worker for model %q", model) } func (p *Pool) Snapshots() []Snapshot { out := make([]Snapshot, 0, len(p.workers)) for _, w := range p.workers { w.mu.RLock() models := make([]string, 0, len(w.models)) for m := range w.models { models = append(models, m) } sort.Strings(models) loaded := append([]LoadedModel(nil), w.loadedModels...) tel := w.telemetry var resident, vram int64 for _, m := range loaded { resident += m.Size vram += m.SizeVRAM } memoryTotal := tel.MemoryTotalBytes if memoryTotal <= 0 { memoryTotal = w.cfg.MemoryCapacityBytes } vramTotal := tel.VRAMTotalBytes if vramTotal <= 0 { vramTotal = w.cfg.VRAMCapacityBytes } activeByModel := make(map[string]int, len(w.modelActive)) for k, v := range w.modelActive { activeByModel[k] = v } limits := make(map[string]int, len(w.cfg.ModelConcurrency)) for k, v := range w.cfg.ModelConcurrency { limits[k] = v } perf := performanceSnapshotLocked(w) maintenance := w.maintenance circuitState := w.circuitState if circuitState == "open" && !w.circuitOpenUntil.IsZero() && !time.Now().Before(w.circuitOpenUntil) { circuitState = "half_open" } accepting := maintenance == "active" && circuitState != "open" x := Snapshot{Name: w.cfg.Name, URL: w.url.String(), Healthy: w.healthy.Load(), Active: w.active.Load(), MaxConcurrent: w.cfg.MaxConcurrent, Models: models, LoadedModels: loaded, ResidentBytes: resident, VRAMBytes: vram, MemoryCapacityBytes: w.cfg.MemoryCapacityBytes, VRAMCapacityBytes: w.cfg.VRAMCapacityBytes, MemoryUsedBytes: tel.MemoryUsedBytes, MemoryTotalBytes: memoryTotal, VRAMUsedBytes: tel.VRAMUsedBytes, VRAMTotalBytes: vramTotal, GPUUtilizationPct: tel.GPUUtilizationPct, GPUTemperatureC: tel.GPUTemperatureC, GPUPowerWatts: tel.GPUPowerWatts, ModelActive: activeByModel, ModelLimits: limits, Performance: perf, ModelPlacement: clonePlacementRule(w.placement), PlacementOverride: w.placementOverride, Maintenance: maintenance, AcceptingNew: accepting, CircuitState: circuitState, CircuitFailures: w.circuitFailures, CircuitOpenUntil: w.circuitOpenUntil, LastCircuitError: w.lastCircuitError, TelemetrySource: tel.Source, TelemetryError: tel.Error, Labels: w.cfg.Labels, LastError: w.lastError, LastCheck: w.lastCheck} w.mu.RUnlock() out = append(out, x) } return out } func (p *Pool) candidates(model string) []*state { return p.candidatesExcluding(model, nil) } func (p *Pool) candidatesExcluding(model string, excluded map[string]bool) []*state { return p.candidatesExcludingAllowed(model, excluded, nil, 0) } func (p *Pool) candidatesExcludingAllowed(model string, excluded, allowed map[string]bool, requestedContext int64) []*state { eligible := make([]*state, 0, len(p.workers)) maxTPS := 0.0 for _, w := range p.workers { if !w.healthy.Load() || (excluded != nil && excluded[w.cfg.Name]) || (allowed != nil && !allowed[w.cfg.Name]) || !w.acceptingNew(p.reliability) { continue } if model != "" && !w.placementDecision(model).Allowed { continue } if model != "" && w.modelMaintenanceActive(model) { continue } eligible = append(eligible, w) if t := w.outputTPS(model); t > maxTPS { maxTPS = t } } out := eligible if model != "" { installed := make([]*state, 0, len(eligible)) unknown := make([]*state, 0, len(eligible)) for _, w := range eligible { known, has := w.inventoryState(model) if !known { unknown = append(unknown, w) continue } if has { installed = append(installed, w) } } if len(installed) > 0 { out = installed } else if len(unknown) > 0 { out = unknown } else { out = nil } } sort.SliceStable(out, func(i, j int) bool { return p.scoreForContext(out[i], model, maxTPS, requestedContext) < p.scoreForContext(out[j], model, maxTPS, requestedContext) }) return out } func clonePlacementRule(r config.ModelPlacementRule) config.ModelPlacementRule { r.AllowedModels = append([]string(nil), r.AllowedModels...) r.DeniedModels = append([]string(nil), r.DeniedModels...) return r } func normalizePlacementRule(r config.ModelPlacementRule) config.ModelPlacementRule { if strings.TrimSpace(r.Mode) == "" { r.Mode = "allow_all" } r.Mode = strings.TrimSpace(r.Mode) norm := func(in []string) []string { out := make([]string, 0, len(in)) seen := map[string]bool{} for _, x := range in { x = strings.TrimSpace(x) if x == "" || seen[x] { continue } seen[x] = true out = append(out, x) } sort.Strings(out) return out } r.AllowedModels = norm(r.AllowedModels) r.DeniedModels = norm(r.DeniedModels) return r } func modelPatternSpecificity(pattern, model string) (int, bool, bool) { pattern = strings.TrimSpace(pattern) model = strings.TrimSpace(model) plain := canonicalModel(model) if pattern == "" { return 0, false, false } if pattern == model || pattern == plain { return 100000 + len(pattern), true, true } if pattern == "*" { return 0, true, false } if strings.HasSuffix(pattern, "*") { prefix := strings.TrimSuffix(pattern, "*") if strings.HasPrefix(model, prefix) || strings.HasPrefix(plain, prefix) { return len(prefix), true, false } } return 0, false, false } func evaluatePlacement(r config.ModelPlacementRule, model string) PlacementDecision { r = normalizePlacementRule(r) bestAllow, bestDeny := -1, -1 allowPattern, denyPattern := "", "" allowExact, denyExact := false, false for _, p := range r.AllowedModels { if spec, ok, exact := modelPatternSpecificity(p, model); ok && spec > bestAllow { bestAllow, allowPattern, allowExact = spec, p, exact } } for _, p := range r.DeniedModels { if spec, ok, exact := modelPatternSpecificity(p, model); ok && spec > bestDeny { bestDeny, denyPattern, denyExact = spec, p, exact } } if bestAllow >= 0 || bestDeny >= 0 { if bestAllow > bestDeny { return PlacementDecision{Allowed: true, Source: "allow_rule", Pattern: allowPattern, ExactOverride: allowExact} } return PlacementDecision{Allowed: false, Source: "deny_rule", Pattern: denyPattern, ExactOverride: denyExact} } if r.Mode == "whitelist" { return PlacementDecision{Allowed: false, Source: "whitelist_default"} } return PlacementDecision{Allowed: true, Source: "allow_all_default"} } func (w *state) placementDecision(model string) PlacementDecision { w.mu.RLock() r := clonePlacementRule(w.placement) w.mu.RUnlock() return evaluatePlacement(r, model) } func (w *state) inventoryState(model string) (known, has bool) { if model == "" { return true, true } plain := canonicalModel(model) w.mu.RLock() known = w.installedKnown has = w.installed[model] || w.installed[plain] || w.models[model] || w.models[plain] w.mu.RUnlock() return known, has } func (p *Pool) PlacementDecision(workerName, model string) (PlacementDecision, bool) { w := p.byName[workerName] if w == nil { return PlacementDecision{}, false } return w.placementDecision(model), true } func (p *Pool) PlacementSnapshots() []PlacementSnapshot { out := make([]PlacementSnapshot, 0, len(p.workers)) for _, w := range p.workers { w.mu.RLock() loaded := make([]string, 0, len(w.loadedModels)) seenLoaded := map[string]bool{} for _, m := range w.loadedModels { name := strings.TrimSpace(m.Model) if name == "" { name = strings.TrimSpace(m.Name) } if name != "" && !seenLoaded[name] { seenLoaded[name] = true loaded = append(loaded, name) } } sort.Strings(loaded) out = append(out, PlacementSnapshot{Worker: w.cfg.Name, Baseline: clonePlacementRule(w.baselinePlacement), Effective: clonePlacementRule(w.placement), Override: w.placementOverride, InventoryKnown: w.installedKnown, InstalledModels: append([]string(nil), w.installedModels...), LoadedModels: loaded, InventoryError: w.inventoryError}) w.mu.RUnlock() } sort.Slice(out, func(i, j int) bool { return out[i].Worker < out[j].Worker }) return out } func (p *Pool) SetPlacement(workerName string, rule config.ModelPlacementRule, override bool) error { if err := config.ValidateModelPlacementRule(rule); err != nil { return err } w := p.byName[workerName] if w == nil { return fmt.Errorf("unknown worker %q", workerName) } rule = normalizePlacementRule(rule) w.mu.Lock() w.placement = rule w.placementOverride = override w.mu.Unlock() p.signal() return nil } func (p *Pool) ResetPlacement(workerName string) error { w := p.byName[workerName] if w == nil { return fmt.Errorf("unknown worker %q", workerName) } w.mu.Lock() w.placement = clonePlacementRule(w.baselinePlacement) w.placementOverride = false w.mu.Unlock() p.signal() return nil } func normalizeMaintenance(mode string) (string, error) { mode = strings.ToLower(strings.TrimSpace(mode)) if mode == "" { mode = "active" } switch mode { case "active", "draining", "disabled": return mode, nil } return "", fmt.Errorf("maintenance mode must be active, draining, or disabled") } func (p *Pool) SetMaintenance(workerName, mode string) error { w := p.byName[workerName] if w == nil { return fmt.Errorf("unknown worker %q", workerName) } m, err := normalizeMaintenance(mode) if err != nil { return err } w.mu.Lock() w.maintenance = m w.mu.Unlock() p.signal() return nil } func (p *Pool) Maintenance(workerName string) (string, bool) { w := p.byName[workerName] if w == nil { return "", false } w.mu.RLock() m := w.maintenance w.mu.RUnlock() return m, true } func (w *state) acceptingNew(rel config.ReliabilityConfig) bool { w.mu.RLock() defer w.mu.RUnlock() if w.maintenance != "active" { return false } if !rel.Enabled { return true } if w.circuitState == "open" && time.Now().Before(w.circuitOpenUntil) { return false } if w.circuitState == "half_open" && w.halfOpenInFlight { return false } return true } func (w *state) claimCircuitProbe(rel config.ReliabilityConfig) bool { if !rel.Enabled { return true } w.mu.Lock() defer w.mu.Unlock() if w.maintenance != "active" { return false } now := time.Now() if w.circuitState == "open" { if now.Before(w.circuitOpenUntil) { return false } if w.halfOpenInFlight { return false } w.circuitState = "half_open" w.halfOpenInFlight = true return true } if w.circuitState == "half_open" { if w.halfOpenInFlight { return false } w.halfOpenInFlight = true return true } return true } func (w *state) abandonCircuitProbe() { w.mu.Lock() if w.circuitState == "half_open" && w.halfOpenInFlight { w.halfOpenInFlight = false } w.mu.Unlock() } func (p *Pool) ReportResult(workerName string, failed bool, errText string) bool { w := p.byName[workerName] if w == nil || !p.reliability.Enabled { return false } w.mu.Lock() defer w.mu.Unlock() if !failed { w.circuitFailures = 0 w.circuitState = "closed" w.circuitOpenUntil = time.Time{} w.halfOpenInFlight = false w.lastCircuitError = "" return false } wasOpen := w.circuitState == "open" w.circuitFailures++ w.lastCircuitError = errText if w.circuitState == "half_open" || w.circuitFailures >= p.reliability.FailureThreshold { w.circuitState = "open" w.circuitOpenUntil = time.Now().Add(p.reliability.OpenDuration.Value()) } w.halfOpenInFlight = false p.signal() return !wasOpen && w.circuitState == "open" } func (p *Pool) CircuitReset(workerName string) error { w := p.byName[workerName] if w == nil { return fmt.Errorf("unknown worker %q", workerName) } w.mu.Lock() w.circuitState = "closed" w.circuitFailures = 0 w.circuitOpenUntil = time.Time{} w.halfOpenInFlight = false w.lastCircuitError = "" w.mu.Unlock() p.signal() return nil } func (p *Pool) CanRoute(model string) bool { return len(p.candidates(model)) > 0 } func (p *Pool) SetModelConcurrency(workerName, model string, limit int) error { w := p.byName[workerName] if w == nil { return fmt.Errorf("unknown worker %q", workerName) } model = strings.TrimSpace(model) if model == "" || limit <= 0 { return fmt.Errorf("model and positive limit are required") } w.mu.Lock() if w.cfg.ModelConcurrency == nil { w.cfg.ModelConcurrency = map[string]int{} } w.cfg.ModelConcurrency[model] = minInt(limit, w.cfg.MaxConcurrent) w.mu.Unlock() p.signal() return nil } func (p *Pool) ResetModelConcurrency(workerName, model string) error { w := p.byName[workerName] if w == nil { return fmt.Errorf("unknown worker %q", workerName) } model = strings.TrimSpace(model) if model == "" { return fmt.Errorf("model required") } w.mu.Lock() if v, ok := w.baselineModelConcurrency[model]; ok { if w.cfg.ModelConcurrency == nil { w.cfg.ModelConcurrency = map[string]int{} } w.cfg.ModelConcurrency[model] = v } else { delete(w.cfg.ModelConcurrency, model) } w.mu.Unlock() p.signal() return nil } func (p *Pool) WorkerConfig(workerName string) (config.WorkerConfig, bool) { w := p.byName[workerName] if w == nil { return config.WorkerConfig{}, false } w.mu.RLock() defer w.mu.RUnlock() c := w.cfg c.ModelConcurrency = map[string]int{} for k, v := range w.cfg.ModelConcurrency { c.ModelConcurrency[k] = v } c.ContextLimits = map[string]int64{} for k, v := range w.cfg.ContextLimits { c.ContextLimits[k] = v } return c, true } func canonicalModel(model string) string { model = strings.TrimSpace(model) if model == "" { return "" } return strings.TrimSuffix(model, ":latest") } func (w *state) hasModel(model string) bool { if model == "" { return true } plain := canonicalModel(model) w.mu.RLock() ok := w.installed[model] || w.installed[plain] || w.models[model] || w.models[plain] w.mu.RUnlock() return ok } func (w *state) loadedContext(model string) int64 { plain := canonicalModel(model) w.mu.RLock() defer w.mu.RUnlock() for _, m := range w.loadedModels { id := strings.TrimSpace(m.Model) if id == "" { id = strings.TrimSpace(m.Name) } if id == model || canonicalModel(id) == plain { return m.ContextLength } } return 0 } func (w *state) contextLimit(model string) int64 { w.mu.RLock() limits := make(map[string]int64, len(w.cfg.ContextLimits)) for pattern, limit := range w.cfg.ContextLimits { limits[pattern] = limit } w.mu.RUnlock() bestSpec := -1 var best int64 for pattern, limit := range limits { if spec, ok, _ := modelPatternSpecificity(pattern, model); ok && spec > bestSpec { bestSpec, best = spec, limit } } return best } func minPositive(values ...int64) int64 { var out int64 for _, v := range values { if v <= 0 { continue } if out == 0 || v < out { out = v } } return out } func (p *Pool) contextWindowForWorker(ctx context.Context, w *state, model string) ContextWindow { workerDefault := w.cfg.DefaultContextTokens if workerDefault == 0 { workerDefault = p.capCfg.Context.DefaultWorkerTokens } x := ContextWindow{Worker: w.cfg.Name, Model: model, WorkerDefaultTokens: workerDefault, WorkerLimitTokens: w.contextLimit(model)} meta, err := p.metadataForWorker(ctx, w, model) if err != nil { x.MetadataError = err.Error() } else { x.ModelMaxTokens = meta.ContextLength x.ConfiguredTokens = meta.ConfiguredContextLength } x.LoadedTokens = w.loadedContext(model) // The loaded context is the strongest evidence because it is what Ollama is // actually using right now. Otherwise prefer a Modelfile num_ctx, then the // operator/default worker context. -1 means explicitly fall back to the // theoretical model maximum. switch { case x.LoadedTokens > 0: x.EffectiveTokens, x.EffectiveSource = x.LoadedTokens, "loaded" case x.ConfiguredTokens > 0: x.EffectiveTokens, x.EffectiveSource = x.ConfiguredTokens, "modelfile" case x.WorkerDefaultTokens > 0: x.EffectiveTokens, x.EffectiveSource = x.WorkerDefaultTokens, "worker_default" case x.WorkerDefaultTokens == -1 && x.ModelMaxTokens > 0: x.EffectiveTokens, x.EffectiveSource = x.ModelMaxTokens, "model_max" } x.EffectiveTokens = minPositive(x.EffectiveTokens, x.ModelMaxTokens, x.WorkerLimitTokens) if x.EffectiveTokens == 0 && x.EffectiveSource != "" { x.EffectiveSource = "" } return x } // ContextWindows returns the effective context evidence for every currently // routable worker that owns the model. It is intentionally independent of // concurrency so admission can reject an impossible context before waiting in // the fair scheduler queue. func (p *Pool) ContextWindows(ctx context.Context, model string) []ContextWindow { candidates := p.candidates(model) out := make([]ContextWindow, 0, len(candidates)) for _, w := range candidates { out = append(out, p.contextWindowForWorker(ctx, w, model)) } return out } func (w *state) modelLimit(model string) int { raw := strings.TrimSpace(model) plain := canonicalModel(raw) limit := w.cfg.MaxConcurrent best := "" for pattern, n := range w.cfg.ModelConcurrency { if pattern == raw || pattern == plain { return minInt(n, w.cfg.MaxConcurrent) } if pattern == "*" && best == "" { best = pattern limit = n continue } if strings.HasSuffix(pattern, "*") { prefix := strings.TrimSuffix(pattern, "*") if (strings.HasPrefix(raw, prefix) || strings.HasPrefix(plain, prefix)) && len(pattern) > len(best) { best = pattern limit = n } } } if limit <= 0 { limit = w.cfg.MaxConcurrent } return minInt(limit, w.cfg.MaxConcurrent) } func (w *state) acquireModelSlot(model string) bool { key := canonicalModel(model) limit := w.modelLimit(model) w.mu.Lock() defer w.mu.Unlock() if w.modelMaintenance[key] || w.modelActive[key] >= limit { return false } w.modelActive[key]++ return true } func (w *state) modelMaintenanceActive(model string) bool { key := canonicalModel(model) w.mu.RLock() busy := w.modelMaintenance[key] w.mu.RUnlock() return busy } // BeginModelMaintenance reserves a worker/model pair for a short control-plane // operation such as preload or unload. New inference leases for the same model // are excluded until the returned release function is called. This closes the // race where an idle-unload could otherwise start exactly as inference arrives. func (p *Pool) BeginModelMaintenance(workerName, model string) (func(), error) { w := p.byName[workerName] if w == nil { return nil, fmt.Errorf("unknown worker %q", workerName) } key := canonicalModel(model) if key == "" { return nil, errors.New("model required") } if !w.healthy.Load() { return nil, fmt.Errorf("worker %s is unhealthy", workerName) } w.mu.Lock() if w.maintenance != "active" { w.mu.Unlock() return nil, fmt.Errorf("worker %s is %s", workerName, w.maintenance) } if w.circuitState != "" && w.circuitState != "closed" { w.mu.Unlock() return nil, fmt.Errorf("worker %s circuit is %s", workerName, w.circuitState) } if w.modelMaintenance[key] { w.mu.Unlock() return nil, fmt.Errorf("model %s already has a maintenance operation on worker %s", model, workerName) } if w.modelActive[key] > 0 { w.mu.Unlock() return nil, fmt.Errorf("model %s is active on worker %s", model, workerName) } w.modelMaintenance[key] = true w.mu.Unlock() p.signal() var once sync.Once return func() { once.Do(func() { w.mu.Lock() delete(w.modelMaintenance, key) w.mu.Unlock() p.signal() }) }, nil } func (w *state) outputTPS(model string) float64 { w.mu.RLock() defer w.mu.RUnlock() if p, ok := w.performance[canonicalModel(model)]; ok { return p.OutputTPS } return 0 } func (p *Pool) score(w *state, model string, maxTPS float64) float64 { return p.scoreForContext(w, model, maxTPS, 0) } func (p *Pool) scoreForContext(w *state, model string, maxTPS float64, requestedContext int64) float64 { active := float64(w.active.Load()) / float64(max(1, w.cfg.MaxConcurrent)) s := active * 100 w.mu.RLock() plain := canonicalModel(model) loaded := w.models[model] || w.models[plain] installed := w.installed[model] || w.installed[plain] loadedContext := int64(0) if requestedContext > 0 && loaded { for _, m := range w.loadedModels { id := strings.TrimSpace(m.Model) if id == "" { id = strings.TrimSpace(m.Name) } if id == model || canonicalModel(id) == plain { loadedContext = m.ContextLength break } } } tel := w.telemetry modelActive := w.modelActive[plain] perf := w.performance[plain] w.mu.RUnlock() if requestedContext > 0 && loadedContext > 0 && loadedContext < requestedContext { // An explicit larger num_ctx can still be routed here, but Ollama must // reload/resize the model, so do not pretend the current loaded state is // a warm-context advantage. loaded = false installed = true } if loaded { s -= p.routing.LoadedBonus } else if installed { s -= p.routing.InstalledBonus } if lim := w.modelLimit(model); lim > 0 { s += 25 * float64(modelActive) / float64(lim) } if maxTPS > 0 && perf.OutputTPS > 0 { s -= p.routing.ThroughputBonus * math.Min(1, perf.OutputTPS/maxTPS) } vramTotal := tel.VRAMTotalBytes if vramTotal <= 0 { vramTotal = w.cfg.VRAMCapacityBytes } if vramTotal > 0 && tel.VRAMUsedBytes > 0 { pct := 100 * float64(tel.VRAMUsedBytes) / float64(vramTotal) s += p.routing.VRAMPressurePenalty * math.Min(1, pct/100) if !loaded && p.routing.AvoidVRAMPercent > 0 && pct >= p.routing.AvoidVRAMPercent { s += 500 } } if tel.GPUUtilizationPct > 0 { s += p.routing.GPUUtilizationPenalty * math.Min(1, tel.GPUUtilizationPct/100) } return s } // RoutingExplanation describes why a worker is or is not eligible for a model. // It is intended for the admin policy simulator and never acquires a worker slot. type RoutingExplanation struct { Worker string `json:"worker"` Eligible bool `json:"eligible"` Reason string `json:"reason,omitempty"` Healthy bool `json:"healthy"` Maintenance string `json:"maintenance"` CircuitState string `json:"circuit_state"` Placement PlacementDecision `json:"placement"` InventoryKnown bool `json:"inventory_known"` Installed bool `json:"installed"` Loaded bool `json:"loaded"` Active int64 `json:"active"` MaxConcurrent int `json:"max_concurrent"` ModelActive int `json:"model_active"` ModelLimit int `json:"model_limit"` OutputTPS float64 `json:"output_tps,omitempty"` VRAMPercent float64 `json:"vram_percent,omitempty"` GPUUtilizationPercent float64 `json:"gpu_utilization_percent,omitempty"` Score float64 `json:"score,omitempty"` ScoreComponents map[string]float64 `json:"score_components,omitempty"` } // ExplainRouting snapshots the exact eligibility gates and adaptive score used // for a model without reserving scheduler or worker capacity. func (p *Pool) ExplainRouting(model string) []RoutingExplanation { return p.ExplainRoutingAllowed(model, nil, 0) } // ExplainRoutingAllowed additionally applies an admission-approved context // worker set so the policy simulator can mirror real context-aware routing. func (p *Pool) ExplainRoutingAllowed(model string, allowed map[string]bool, requestedContext int64) []RoutingExplanation { plain := canonicalModel(model) maxTPS := 0.0 for _, w := range p.workers { if t := w.outputTPS(model); t > maxTPS { maxTPS = t } } out := make([]RoutingExplanation, 0, len(p.workers)) for _, w := range p.workers { w.mu.RLock() tel := w.telemetry loaded := w.models[model] || w.models[plain] installed := w.installed[model] || w.installed[plain] || loaded loadedContext := int64(0) if requestedContext > 0 && loaded { for _, lm := range w.loadedModels { id := strings.TrimSpace(lm.Model) if id == "" { id = strings.TrimSpace(lm.Name) } if id == model || canonicalModel(id) == plain { loadedContext = lm.ContextLength break } } } known := w.installedKnown modelActive := w.modelActive[plain] perf := w.performance[plain] maintenance := w.maintenance circuit := w.circuitState openUntil := w.circuitOpenUntil w.mu.RUnlock() if requestedContext > 0 && loadedContext > 0 && loadedContext < requestedContext { loaded = false installed = true } if circuit == "open" && !openUntil.IsZero() && !time.Now().Before(openUntil) { circuit = "half_open" } pl := w.placementDecision(model) e := RoutingExplanation{Worker: w.cfg.Name, Healthy: w.healthy.Load(), Maintenance: maintenance, CircuitState: circuit, Placement: pl, InventoryKnown: known, Installed: installed, Loaded: loaded, Active: w.active.Load(), MaxConcurrent: w.cfg.MaxConcurrent, ModelActive: modelActive, ModelLimit: w.modelLimit(model), OutputTPS: perf.OutputTPS, GPUUtilizationPercent: tel.GPUUtilizationPct} total := tel.VRAMTotalBytes if total <= 0 { total = w.cfg.VRAMCapacityBytes } if total > 0 && tel.VRAMUsedBytes > 0 { e.VRAMPercent = 100 * float64(tel.VRAMUsedBytes) / float64(total) } switch { case allowed != nil && !allowed[w.cfg.Name]: e.Reason = "context_window" case !e.Healthy: e.Reason = "worker_unhealthy" case maintenance != "active": e.Reason = "worker_" + maintenance case circuit == "open": e.Reason = "circuit_open" case !pl.Allowed: e.Reason = "placement_denied" case known && !installed: e.Reason = "model_not_installed" case e.Active >= int64(max(1, w.cfg.MaxConcurrent)): e.Reason = "worker_concurrency_full" case e.ModelActive >= e.ModelLimit: e.Reason = "model_concurrency_full" default: e.Eligible = true } if e.Eligible { components := map[string]float64{} components["worker_load"] = float64(e.Active) / float64(max(1, w.cfg.MaxConcurrent)) * 100 if loaded { components["loaded_bonus"] = -p.routing.LoadedBonus } else if installed { components["installed_bonus"] = -p.routing.InstalledBonus } if e.ModelLimit > 0 { components["model_load"] = 25 * float64(e.ModelActive) / float64(e.ModelLimit) } if maxTPS > 0 && perf.OutputTPS > 0 { components["throughput_bonus"] = -p.routing.ThroughputBonus * math.Min(1, perf.OutputTPS/maxTPS) } if e.VRAMPercent > 0 { components["vram_pressure"] = p.routing.VRAMPressurePenalty * math.Min(1, e.VRAMPercent/100) if !loaded && p.routing.AvoidVRAMPercent > 0 && e.VRAMPercent >= p.routing.AvoidVRAMPercent { components["vram_avoid"] = 500 } } if tel.GPUUtilizationPct > 0 { components["gpu_pressure"] = p.routing.GPUUtilizationPenalty * math.Min(1, tel.GPUUtilizationPct/100) } for _, v := range components { e.Score += v } e.ScoreComponents = components } out = append(out, e) } sort.Slice(out, func(i, j int) bool { if out[i].Eligible != out[j].Eligible { return out[i].Eligible } if out[i].Eligible && out[i].Score != out[j].Score { return out[i].Score < out[j].Score } return out[i].Worker < out[j].Worker }) return out } func minInt(a, b int) int { if a < b { return a } return b } func (p *Pool) healthLoop(ctx context.Context, w *state) { t := time.NewTicker(w.cfg.HealthInterval.Value()) defer t.Stop() for { select { case <-ctx.Done(): return case <-t.C: p.refresh(ctx, w) } } } func (p *Pool) refresh(ctx context.Context, w *state) { cctx, cancel := context.WithTimeout(ctx, 3*time.Second) defer cancel() req, _ := http.NewRequestWithContext(cctx, http.MethodGet, w.url.String()+"/api/ps", nil) resp, err := p.client.Do(req) if err != nil { p.setHealth(w, false, err.Error(), nil, nil) return } defer resp.Body.Close() if resp.StatusCode/100 != 2 { p.setHealth(w, false, fmt.Sprintf("HTTP %d", resp.StatusCode), nil, nil) return } b, err := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) if err != nil { p.setHealth(w, false, err.Error(), nil, nil) return } var doc struct { Models []LoadedModel `json:"models"` } if err := json.Unmarshal(b, &doc); err != nil { p.setHealth(w, false, err.Error(), nil, nil) return } models := map[string]bool{} for _, m := range doc.Models { if m.Name != "" { models[m.Name] = true models[strings.TrimSuffix(m.Name, ":latest")] = true } if m.Model != "" { models[m.Model] = true models[strings.TrimSuffix(m.Model, ":latest")] = true } } p.setHealth(w, true, "", models, doc.Models) p.refreshInstalled(ctx, w) p.refreshTelemetry(ctx, w) } func (p *Pool) refreshInstalled(ctx context.Context, w *state) { cctx, cancel := context.WithTimeout(ctx, 3*time.Second) defer cancel() req, _ := http.NewRequestWithContext(cctx, http.MethodGet, w.url.String()+"/api/tags", nil) resp, err := p.client.Do(req) if err != nil { w.mu.Lock() w.inventoryError = err.Error() w.mu.Unlock() return } defer resp.Body.Close() if resp.StatusCode/100 != 2 { w.mu.Lock() w.inventoryError = fmt.Sprintf("HTTP %d", resp.StatusCode) w.mu.Unlock() return } var doc struct { Models []ModelInfo `json:"models"` } if err := json.NewDecoder(io.LimitReader(resp.Body, 8<<20)).Decode(&doc); err != nil { w.mu.Lock() w.inventoryError = err.Error() w.mu.Unlock() return } installed := make(map[string]bool, len(doc.Models)*2) display := make([]string, 0, len(doc.Models)) seenDisplay := map[string]bool{} for _, m := range doc.Models { preferred := strings.TrimSpace(m.Model) if preferred == "" { preferred = strings.TrimSpace(m.Name) } if preferred != "" && !seenDisplay[preferred] { seenDisplay[preferred] = true display = append(display, preferred) } for _, id := range []string{m.Model, m.Name} { id = strings.TrimSpace(id) if id == "" { continue } installed[id] = true installed[strings.TrimSuffix(id, ":latest")] = true } } sort.Strings(display) w.mu.Lock() w.installed = installed w.installedKnown = true w.installedModels = display w.inventoryError = "" w.mu.Unlock() } func (p *Pool) refreshTelemetry(ctx context.Context, w *state) { if !w.cfg.LocalSystemStats && !w.cfg.NVIDIASMI && w.cfg.TelemetryURL == "" { return } t := ResourceTelemetry{UpdatedAt: time.Now().UTC()} if w.cfg.LocalSystemStats { cctx, cancel := context.WithTimeout(ctx, 2*time.Second) m, err := hoststats.ReadMemory(cctx) cancel() if err != nil { t.Error = err.Error() } else { t.MemoryTotalBytes = m.TotalBytes t.MemoryUsedBytes = m.UsedBytes t.Source = "local-system" } } if w.cfg.NVIDIASMI { cctx, cancel := context.WithTimeout(ctx, 2*time.Second) n, err := hoststats.ReadNVIDIA(cctx, w.cfg.NVIDIAGPU) cancel() if err != nil { if t.Error == "" { t.Error = err.Error() } else { t.Error += "; " + err.Error() } } else { t.VRAMUsedBytes = n.MemoryUsedBytes t.VRAMTotalBytes = n.MemoryTotalBytes t.GPUUtilizationPct = n.UtilizationPercent t.GPUTemperatureC = n.TemperatureC t.GPUPowerWatts = n.PowerWatts if t.Source == "" { t.Source = "nvidia-smi" } else { t.Source += "+nvidia-smi" } } } if w.cfg.TelemetryURL != "" { cctx, cancel := context.WithTimeout(ctx, 2*time.Second) req, _ := http.NewRequestWithContext(cctx, http.MethodGet, w.cfg.TelemetryURL, nil) resp, err := p.client.Do(req) if err != nil { appendTelemetryError(&t, err.Error()) } else { func() { defer resp.Body.Close() if resp.StatusCode/100 != 2 { appendTelemetryError(&t, fmt.Sprintf("telemetry HTTP %d", resp.StatusCode)) return } var ext externalTelemetry if err := json.NewDecoder(io.LimitReader(resp.Body, 64<<10)).Decode(&ext); err != nil { appendTelemetryError(&t, err.Error()) return } if err := validateExternalTelemetryTimestamp(time.Now().UTC(), ext.UpdatedAt, telemetryMaxAge(w.cfg.HealthInterval.Value())); err != nil { appendTelemetryError(&t, err.Error()) return } mergeExternalTelemetry(&t, ext) }() } cancel() } w.mu.Lock() w.telemetry = t w.mu.Unlock() } func telemetryMaxAge(healthInterval time.Duration) time.Duration { maxAge := 6 * healthInterval if maxAge < 30*time.Second { maxAge = 30 * time.Second } return maxAge } func validateExternalTelemetryTimestamp(now time.Time, updatedAt *time.Time, maxAge time.Duration) error { // updated_at was optional before checkpoint 21. Preserve compatibility with // existing exporters, while the shipped agent always supplies it. if updatedAt == nil { return nil } if maxAge <= 0 { maxAge = 30 * time.Second } age := now.Sub(updatedAt.UTC()) if age < -30*time.Second { return fmt.Errorf("telemetry timestamp is %.0fs in the future", -age.Seconds()) } if age > maxAge { return fmt.Errorf("telemetry sample is stale: age %s exceeds %s", age.Round(time.Second), maxAge) } return nil } func mergeExternalTelemetry(t *ResourceTelemetry, ext externalTelemetry) { if ext.MemoryTotalBytes != nil { t.MemoryTotalBytes = *ext.MemoryTotalBytes } if ext.MemoryUsedBytes != nil { t.MemoryUsedBytes = *ext.MemoryUsedBytes } if ext.VRAMTotalBytes != nil { t.VRAMTotalBytes = *ext.VRAMTotalBytes } if ext.VRAMUsedBytes != nil { t.VRAMUsedBytes = *ext.VRAMUsedBytes } if ext.GPUUtilizationPct != nil { t.GPUUtilizationPct = *ext.GPUUtilizationPct } if ext.GPUTemperatureC != nil { t.GPUTemperatureC = *ext.GPUTemperatureC } if ext.GPUPowerWatts != nil { t.GPUPowerWatts = *ext.GPUPowerWatts } source := "telemetry-url" if extSource := strings.TrimSpace(ext.Source); extSource != "" { source += ":" + extSource } if t.Source == "" { t.Source = source } else { t.Source += "+" + source } if strings.TrimSpace(ext.Error) != "" { appendTelemetryError(t, ext.Error) } } func appendTelemetryError(t *ResourceTelemetry, message string) { message = strings.TrimSpace(message) if message == "" { return } if t.Error == "" { t.Error = message } else { t.Error += "; " + message } } // Observe updates an in-memory EWMA of model throughput for adaptive routing. // Exact Ollama eval durations are preferred; wall time is only used when an // OpenAI-compatible response did not expose token evaluation timing. func (p *Pool) Observe(workerName, model string, promptTokens, completionTokens, promptEvalNS, evalNS int64, service time.Duration) { w := p.byName[workerName] model = canonicalModel(model) if w == nil || model == "" { return } promptTPS, outputTPS := 0.0, 0.0 if promptTokens > 0 && promptEvalNS > 0 { promptTPS = float64(promptTokens) / (float64(promptEvalNS) / 1e9) } if completionTokens > 0 { den := time.Duration(evalNS) if den <= 0 { den = service } if den > 0 { outputTPS = float64(completionTokens) / den.Seconds() } } if promptTPS <= 0 && outputTPS <= 0 { return } w.mu.Lock() ps := w.performance[model] const alpha = 0.25 if ps.Samples == 0 { ps.PromptTPS, ps.OutputTPS = promptTPS, outputTPS } else { if promptTPS > 0 { ps.PromptTPS = alpha*promptTPS + (1-alpha)*ps.PromptTPS } if outputTPS > 0 { ps.OutputTPS = alpha*outputTPS + (1-alpha)*ps.OutputTPS } } ps.Samples++ w.performance[model] = ps w.mu.Unlock() } func performanceSnapshotLocked(w *state) []ModelPerformance { out := make([]ModelPerformance, 0, len(w.performance)) for model, p := range w.performance { out = append(out, ModelPerformance{Model: model, PromptTPS: p.PromptTPS, OutputTPS: p.OutputTPS, Samples: p.Samples}) } sort.Slice(out, func(i, j int) bool { return out[i].Model < out[j].Model }) return out } func (p *Pool) setHealth(w *state, ok bool, errText string, models map[string]bool, loaded []LoadedModel) { was := w.healthy.Swap(ok) w.mu.Lock() w.lastCheck = time.Now() w.lastError = errText if models != nil { w.models = models w.loadedModels = append([]LoadedModel(nil), loaded...) } w.mu.Unlock() if was != ok { p.signal() } } func newID() string { b := make([]byte, 12); _, _ = rand.Read(b); return hex.EncodeToString(b) } type ModelDetails struct { ParentModel string `json:"parent_model,omitempty"` Format string `json:"format,omitempty"` Family string `json:"family,omitempty"` Families []string `json:"families,omitempty"` ParameterSize string `json:"parameter_size,omitempty"` QuantizationLevel string `json:"quantization_level,omitempty"` } type ModelInfo struct { Name string `json:"name"` Model string `json:"model,omitempty"` ModifiedAt time.Time `json:"modified_at,omitempty"` Size int64 `json:"size,omitempty"` Digest string `json:"digest,omitempty"` Details ModelDetails `json:"details"` Loaded bool `json:"loaded"` Capabilities []string `json:"capabilities,omitempty"` ContextLength int64 `json:"context_length,omitempty"` ConfiguredContextLength int64 `json:"configured_context_length,omitempty"` LoadedContextLength int64 `json:"loaded_context_length,omitempty"` MetadataError string `json:"metadata_error,omitempty"` } // Metadata returns cached /api/show metadata for a model from a healthy // worker that has the model installed. The second return value is the worker // used for discovery. func (p *Pool) Metadata(ctx context.Context, model string) (ModelMetadata, string, error) { if p.capCfg.Mode == "off" || strings.TrimSpace(model) == "" { return ModelMetadata{Model: model}, "", nil } for _, w := range p.candidates(model) { if !w.hasModel(model) { continue } m, err := p.metadataForWorker(ctx, w, model) return m, w.cfg.Name, err } return ModelMetadata{Model: model}, "", fmt.Errorf("model %q is not installed on a healthy worker", model) } func (p *Pool) metadataForWorker(ctx context.Context, w *state, model string) (ModelMetadata, error) { key := canonicalModel(model) w.mu.RLock() entry, ok := w.metadata[key] w.mu.RUnlock() if ok { ttl := p.capCfg.CacheTTL.Value() if entry.Data.Error != "" && ttl > 15*time.Second { ttl = 15 * time.Second } if ttl > 0 && time.Since(entry.FetchedAt) < ttl { if entry.Data.Error != "" { return entry.Data, errors.New(entry.Data.Error) } return entry.Data, nil } } cctx, cancel := context.WithTimeout(ctx, 4*time.Second) defer cancel() body, _ := json.Marshal(map[string]any{"model": model, "verbose": false}) req, _ := http.NewRequestWithContext(cctx, http.MethodPost, w.url.String()+"/api/show", strings.NewReader(string(body))) req.Header.Set("Content-Type", "application/json") resp, err := p.client.Do(req) if err != nil { m := ModelMetadata{Model: model, Error: err.Error(), UpdatedAt: time.Now().UTC()} w.mu.Lock() w.metadata[key] = metadataEntry{Data: m, FetchedAt: time.Now()} w.mu.Unlock() return m, err } defer resp.Body.Close() if resp.StatusCode/100 != 2 { b, _ := io.ReadAll(io.LimitReader(resp.Body, 16<<10)) err := fmt.Errorf("/api/show HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(b))) m := ModelMetadata{Model: model, Error: err.Error(), UpdatedAt: time.Now().UTC()} w.mu.Lock() w.metadata[key] = metadataEntry{Data: m, FetchedAt: time.Now()} w.mu.Unlock() return m, err } var doc struct { Capabilities []string `json:"capabilities"` Details ModelDetails `json:"details"` ModelInfo map[string]any `json:"model_info"` Parameters json.RawMessage `json:"parameters"` } if err := json.NewDecoder(io.LimitReader(resp.Body, 4<<20)).Decode(&doc); err != nil { return ModelMetadata{Model: model}, err } caps := make([]string, 0, len(doc.Capabilities)) seen := map[string]bool{} for _, c := range doc.Capabilities { c = strings.ToLower(strings.TrimSpace(c)) if c != "" && !seen[c] { seen[c] = true caps = append(caps, c) } } sort.Strings(caps) var contextLength int64 for k, v := range doc.ModelInfo { if k != "context_length" && !strings.HasSuffix(k, ".context_length") { continue } var n int64 switch x := v.(type) { case float64: n = int64(x) case int64: n = x case json.Number: n, _ = x.Int64() } if n > contextLength { contextLength = n } } m := ModelMetadata{Model: model, Capabilities: caps, ContextLength: contextLength, ConfiguredContextLength: parseNumCtxParameter(doc.Parameters), Details: doc.Details, UpdatedAt: time.Now().UTC()} w.mu.Lock() w.metadata[key] = metadataEntry{Data: m, FetchedAt: time.Now()} w.mu.Unlock() return m, nil } func parseNumCtxParameter(raw json.RawMessage) int64 { raw = bytes.TrimSpace(raw) if len(raw) == 0 || bytes.Equal(raw, []byte("null")) { return 0 } var text string if json.Unmarshal(raw, &text) == nil { for _, line := range strings.Split(text, "\n") { fields := strings.Fields(strings.TrimSpace(line)) if len(fields) < 2 || !strings.EqualFold(fields[0], "num_ctx") { continue } n, err := strconv.ParseInt(strings.Trim(fields[1], "\"'"), 10, 64) if err == nil && n > 0 { return n } } return 0 } var obj map[string]any dec := json.NewDecoder(bytes.NewReader(raw)) dec.UseNumber() if dec.Decode(&obj) != nil { return 0 } v, ok := obj["num_ctx"] if !ok { return 0 } switch x := v.(type) { case json.Number: n, _ := x.Int64() if n > 0 { return n } case float64: if x > 0 { return int64(x) } case string: n, _ := strconv.ParseInt(strings.TrimSpace(x), 10, 64) if n > 0 { return n } } return 0 } func HasCapability(m ModelMetadata, capability string) bool { capability = strings.ToLower(strings.TrimSpace(capability)) for _, c := range m.Capabilities { if strings.EqualFold(c, capability) { return true } } return false } type Inventory struct { Worker string `json:"worker"` URL string `json:"url"` Models []ModelInfo `json:"models"` Error string `json:"error,omitempty"` } func (p *Pool) URLFor(name string) (*url.URL, bool) { w := p.byName[name] if w == nil { return nil, false } u := *w.url return &u, true } // TagModel is the native Ollama /api/tags model shape used for client-facing // discovery. It intentionally excludes gateway-only fields such as Loaded. type TagModel struct { Name string `json:"name"` Model string `json:"model"` ModifiedAt time.Time `json:"modified_at,omitempty"` Size int64 `json:"size,omitempty"` Digest string `json:"digest,omitempty"` Details ModelDetails `json:"details"` } // Tags queries every healthy worker, merges installed models and normalizes the // modern Ollama discovery contract so clients can rely on both name and model. // A failed worker does not hide models served by the remaining healthy workers. func (p *Pool) Tags(ctx context.Context) ([]TagModel, []string) { inventories := p.Inventories(ctx) byID := make(map[string]TagModel) errs := make([]string, 0) for _, inv := range inventories { if inv.Error != "" { errs = append(errs, inv.Worker+": "+inv.Error) continue } for _, m := range inv.Models { name := strings.TrimSpace(m.Name) model := strings.TrimSpace(m.Model) if model == "" { model = name } if name == "" { name = model } if model == "" { continue } if decision, ok := p.PlacementDecision(inv.Worker, model); ok && !decision.Allowed { continue } cur, exists := byID[model] if !exists || m.ModifiedAt.After(cur.ModifiedAt) { byID[model] = TagModel{Name: name, Model: model, ModifiedAt: m.ModifiedAt, Size: m.Size, Digest: m.Digest, Details: m.Details} } } } out := make([]TagModel, 0, len(byID)) for _, m := range byID { out = append(out, m) } sort.Slice(out, func(i, j int) bool { return out[i].Model < out[j].Model }) sort.Strings(errs) return out, errs } // Loaded returns a de-duplicated native Ollama /api/ps view across workers. func (p *Pool) Loaded() []LoadedModel { byID := make(map[string]LoadedModel) for _, snap := range p.Snapshots() { if !snap.Healthy { continue } for _, m := range snap.LoadedModels { id := strings.TrimSpace(m.Model) if id == "" { id = strings.TrimSpace(m.Name) } if id == "" { continue } if m.Model == "" { m.Model = id } if m.Name == "" { m.Name = id } if cur, ok := byID[id]; !ok || m.SizeVRAM > cur.SizeVRAM { byID[id] = m } } } out := make([]LoadedModel, 0, len(byID)) for _, m := range byID { out = append(out, m) } sort.Slice(out, func(i, j int) bool { return out[i].Model < out[j].Model }) return out } func (p *Pool) Inventories(ctx context.Context) []Inventory { out := make([]Inventory, len(p.workers)) var wg sync.WaitGroup for i, w := range p.workers { wg.Add(1) go func(i int, w *state) { defer wg.Done() inv := Inventory{Worker: w.cfg.Name, URL: w.url.String()} cctx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() req, _ := http.NewRequestWithContext(cctx, http.MethodGet, w.url.String()+"/api/tags", nil) resp, err := p.client.Do(req) if err != nil { inv.Error = err.Error() out[i] = inv return } defer resp.Body.Close() if resp.StatusCode/100 != 2 { inv.Error = fmt.Sprintf("HTTP %d", resp.StatusCode) out[i] = inv return } var doc struct { Models []ModelInfo `json:"models"` } if err := json.NewDecoder(io.LimitReader(resp.Body, 8<<20)).Decode(&doc); err != nil { inv.Error = err.Error() out[i] = inv return } installed := make(map[string]bool, len(doc.Models)*2) for _, m := range doc.Models { for _, id := range []string{m.Model, m.Name} { id = strings.TrimSpace(id) if id != "" { installed[id] = true installed[strings.TrimSuffix(id, ":latest")] = true } } } w.mu.Lock() w.installed = installed w.installedKnown = true display := make([]string, 0, len(doc.Models)) seenDisplay := map[string]bool{} for _, m := range doc.Models { name := strings.TrimSpace(m.Model) if name == "" { name = strings.TrimSpace(m.Name) } if name != "" && !seenDisplay[name] { seenDisplay[name] = true display = append(display, name) } } sort.Strings(display) w.installedModels = display w.inventoryError = "" w.mu.Unlock() w.mu.RLock() loaded := make(map[string]bool, len(w.models)) loadedCtx := make(map[string]int64, len(w.loadedModels)*2) for k, v := range w.models { loaded[k] = v } for _, lm := range w.loadedModels { id := strings.TrimSpace(lm.Model) if id == "" { id = strings.TrimSpace(lm.Name) } if id != "" && lm.ContextLength > 0 { loadedCtx[id] = lm.ContextLength loadedCtx[canonicalModel(id)] = lm.ContextLength } } w.mu.RUnlock() for j := range doc.Models { name := doc.Models[j].Name if name == "" { name = doc.Models[j].Model } doc.Models[j].Loaded = loaded[name] || loaded[strings.TrimSuffix(name, ":latest")] doc.Models[j].LoadedContextLength = loadedCtx[name] if doc.Models[j].LoadedContextLength == 0 { doc.Models[j].LoadedContextLength = loadedCtx[canonicalModel(name)] } } if p.capCfg.Mode != "off" { sem := make(chan struct{}, 4) var mwg sync.WaitGroup for j := range doc.Models { j := j model := doc.Models[j].Model if model == "" { model = doc.Models[j].Name } mwg.Add(1) go func() { defer mwg.Done() select { case sem <- struct{}{}: case <-ctx.Done(): return } defer func() { <-sem }() m, err := p.metadataForWorker(ctx, w, model) if err != nil { doc.Models[j].MetadataError = err.Error() return } doc.Models[j].Capabilities = append([]string(nil), m.Capabilities...) doc.Models[j].ContextLength = m.ContextLength doc.Models[j].ConfiguredContextLength = m.ConfiguredContextLength }() } mwg.Wait() } sort.Slice(doc.Models, func(a, b int) bool { return doc.Models[a].Name < doc.Models[b].Name }) inv.Models = doc.Models out[i] = inv }(i, w) } wg.Wait() return out }