95 lines
2.3 KiB
Go
95 lines
2.3 KiB
Go
package worker
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"os"
|
|
"time"
|
|
|
|
persiststate "github.com/example/ollama-fair-gateway/internal/state"
|
|
)
|
|
|
|
type PersistentPerformanceState struct {
|
|
Version int `json:"version"`
|
|
SavedAt time.Time `json:"saved_at"`
|
|
Workers map[string]map[string]ModelPerformance `json:"workers"`
|
|
}
|
|
|
|
func (p *Pool) SnapshotPerformance() PersistentPerformanceState {
|
|
out := PersistentPerformanceState{Version: 1, SavedAt: time.Now().UTC(), Workers: map[string]map[string]ModelPerformance{}}
|
|
for _, w := range p.workers {
|
|
w.mu.RLock()
|
|
models := make(map[string]ModelPerformance, len(w.performance))
|
|
for model, perf := range w.performance {
|
|
models[model] = ModelPerformance{Model: model, PromptTPS: perf.PromptTPS, OutputTPS: perf.OutputTPS, Samples: perf.Samples}
|
|
}
|
|
w.mu.RUnlock()
|
|
if len(models) > 0 {
|
|
out.Workers[w.cfg.Name] = models
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (p *Pool) RestorePerformance(s PersistentPerformanceState) {
|
|
if s.Version != 1 {
|
|
return
|
|
}
|
|
for workerName, models := range s.Workers {
|
|
w := p.byName[workerName]
|
|
if w == nil {
|
|
continue
|
|
}
|
|
w.mu.Lock()
|
|
for model, perf := range models {
|
|
if model == "" || perf.Samples <= 0 {
|
|
continue
|
|
}
|
|
w.performance[canonicalModel(model)] = performanceState{PromptTPS: perf.PromptTPS, OutputTPS: perf.OutputTPS, Samples: perf.Samples}
|
|
}
|
|
w.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
func (p *Pool) LoadPerformance(path string) error {
|
|
var s PersistentPerformanceState
|
|
err := (persiststate.AtomicJSON{Path: path, Mode: 0640}).Load(&s)
|
|
if errors.Is(err, os.ErrNotExist) {
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
p.RestorePerformance(s)
|
|
return nil
|
|
}
|
|
|
|
func (p *Pool) SavePerformance(path string) error {
|
|
return (persiststate.AtomicJSON{Path: path, Mode: 0640}).Save(p.SnapshotPerformance())
|
|
}
|
|
|
|
func (p *Pool) StartPerformancePersistence(ctx context.Context, path string, interval time.Duration, onError func(error)) {
|
|
if interval < time.Second {
|
|
interval = 10 * time.Second
|
|
}
|
|
go func() {
|
|
t := time.NewTicker(interval)
|
|
defer t.Stop()
|
|
defer func() {
|
|
if err := p.SavePerformance(path); err != nil && onError != nil {
|
|
onError(err)
|
|
}
|
|
}()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-t.C:
|
|
if err := p.SavePerformance(path); err != nil && onError != nil {
|
|
onError(err)
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
}
|