336 lines
8.1 KiB
Go
336 lines
8.1 KiB
Go
package liveflow
|
|
|
|
import (
|
|
"sort"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/example/ollama-fair-gateway/internal/cost"
|
|
)
|
|
|
|
const (
|
|
StateQueued = "queued"
|
|
StateRouting = "routing"
|
|
StateRunning = "running"
|
|
StateStreaming = "streaming"
|
|
StateCompleted = "completed"
|
|
StateCancelled = "cancelled"
|
|
StateFailed = "failed"
|
|
)
|
|
|
|
type Request struct {
|
|
ID string `json:"id"`
|
|
Tenant string `json:"tenant"`
|
|
Actor string `json:"actor"`
|
|
Application string `json:"application,omitempty"`
|
|
ServiceClass string `json:"service_class,omitempty"`
|
|
API string `json:"api"`
|
|
Path string `json:"path"`
|
|
Model string `json:"model,omitempty"`
|
|
Worker string `json:"worker,omitempty"`
|
|
State string `json:"state"`
|
|
EstimatedCredits float64 `json:"estimated_credits"`
|
|
ActualCredits float64 `json:"actual_credits,omitempty"`
|
|
EstimatedPromptTokens int64 `json:"estimated_prompt_tokens,omitempty"`
|
|
PromptTokens int64 `json:"prompt_tokens,omitempty"`
|
|
CompletionTokens int64 `json:"completion_tokens,omitempty"`
|
|
BytesOut int64 `json:"bytes_out,omitempty"`
|
|
Status int `json:"status,omitempty"`
|
|
QueuedAt time.Time `json:"queued_at"`
|
|
StartedAt *time.Time `json:"started_at,omitempty"`
|
|
FinishedAt *time.Time `json:"finished_at,omitempty"`
|
|
UpdatedAt time.Time `json:"updated_at"`
|
|
QueueMS int64 `json:"queue_ms,omitempty"`
|
|
ServiceMS int64 `json:"service_ms,omitempty"`
|
|
}
|
|
|
|
type Counts struct {
|
|
Total int `json:"total"`
|
|
Active int `json:"active"`
|
|
Queued int `json:"queued"`
|
|
Routing int `json:"routing"`
|
|
Running int `json:"running"`
|
|
Streaming int `json:"streaming"`
|
|
Completed int `json:"completed"`
|
|
Cancelled int `json:"cancelled"`
|
|
Failed int `json:"failed"`
|
|
}
|
|
|
|
type Snapshot struct {
|
|
GeneratedAt time.Time `json:"generated_at"`
|
|
Version uint64 `json:"version"`
|
|
Counts Counts `json:"counts"`
|
|
Truncated bool `json:"truncated,omitempty"`
|
|
Requests []Request `json:"requests"`
|
|
}
|
|
|
|
type Tracker struct {
|
|
mu sync.RWMutex
|
|
active map[string]Request
|
|
recent []Request
|
|
recentTTL time.Duration
|
|
maxRecent int
|
|
version uint64
|
|
notify chan struct{}
|
|
}
|
|
|
|
func New(recentTTL time.Duration, maxRecent int) *Tracker {
|
|
if recentTTL <= 0 {
|
|
recentTTL = 8 * time.Second
|
|
}
|
|
if maxRecent <= 0 {
|
|
maxRecent = 256
|
|
}
|
|
return &Tracker{active: make(map[string]Request), recentTTL: recentTTL, maxRecent: maxRecent, notify: make(chan struct{})}
|
|
}
|
|
|
|
func (t *Tracker) Begin(r Request) {
|
|
now := time.Now().UTC()
|
|
r.State = StateQueued
|
|
r.QueuedAt = now
|
|
r.UpdatedAt = now
|
|
t.mu.Lock()
|
|
t.active[r.ID] = r
|
|
t.bumpLocked()
|
|
t.mu.Unlock()
|
|
}
|
|
|
|
func (t *Tracker) MarkRouting(id, worker string, queue time.Duration) {
|
|
t.update(id, func(r *Request, now time.Time) {
|
|
r.State = StateRouting
|
|
r.Worker = worker
|
|
r.QueueMS = queue.Milliseconds()
|
|
r.UpdatedAt = now
|
|
})
|
|
}
|
|
|
|
func (t *Tracker) MarkRunning(id string) {
|
|
t.update(id, func(r *Request, now time.Time) {
|
|
r.State = StateRunning
|
|
if r.StartedAt == nil {
|
|
x := now
|
|
r.StartedAt = &x
|
|
}
|
|
r.UpdatedAt = now
|
|
})
|
|
}
|
|
|
|
func (t *Tracker) Progress(id string, bytesOut int64, u cost.Usage) {
|
|
t.update(id, func(r *Request, now time.Time) {
|
|
if bytesOut > 0 {
|
|
r.BytesOut = bytesOut
|
|
if r.State == StateRunning || r.State == StateRouting {
|
|
r.State = StateStreaming
|
|
}
|
|
}
|
|
if u.PromptTokens > 0 {
|
|
r.PromptTokens = u.PromptTokens
|
|
}
|
|
if u.CompletionTokens > 0 {
|
|
r.CompletionTokens = u.CompletionTokens
|
|
} else if bytesOut > 0 {
|
|
// Purely for live visualization. Final accounting still comes from
|
|
// Ollama/OpenAI usage metadata in the normal request recorder.
|
|
approx := (bytesOut + 15) / 16
|
|
if approx > r.CompletionTokens {
|
|
r.CompletionTokens = approx
|
|
}
|
|
}
|
|
r.UpdatedAt = now
|
|
})
|
|
}
|
|
|
|
func (t *Tracker) Finish(id string, status int, actualCredits float64, u cost.Usage, service time.Duration) {
|
|
now := time.Now().UTC()
|
|
t.mu.Lock()
|
|
r, ok := t.active[id]
|
|
if !ok {
|
|
t.mu.Unlock()
|
|
return
|
|
}
|
|
delete(t.active, id)
|
|
if status >= 200 && status < 400 {
|
|
r.State = StateCompleted
|
|
} else {
|
|
r.State = StateFailed
|
|
}
|
|
r.Status = status
|
|
r.ActualCredits = actualCredits
|
|
r.ServiceMS = service.Milliseconds()
|
|
if u.PromptTokens > 0 {
|
|
r.PromptTokens = u.PromptTokens
|
|
}
|
|
if u.CompletionTokens > 0 {
|
|
r.CompletionTokens = u.CompletionTokens
|
|
}
|
|
x := now
|
|
r.FinishedAt = &x
|
|
r.UpdatedAt = now
|
|
t.pruneLocked(now)
|
|
t.recent = append(t.recent, r)
|
|
if len(t.recent) > t.maxRecent {
|
|
t.recent = append([]Request(nil), t.recent[len(t.recent)-t.maxRecent:]...)
|
|
}
|
|
t.bumpLocked()
|
|
t.mu.Unlock()
|
|
}
|
|
|
|
func (t *Tracker) Cancel(id string, status int, actualCredits float64, u cost.Usage, service time.Duration) {
|
|
now := time.Now().UTC()
|
|
t.mu.Lock()
|
|
r, ok := t.active[id]
|
|
if !ok {
|
|
t.mu.Unlock()
|
|
return
|
|
}
|
|
delete(t.active, id)
|
|
r.State = StateCancelled
|
|
r.Status = status
|
|
r.ActualCredits = actualCredits
|
|
r.ServiceMS = service.Milliseconds()
|
|
if u.PromptTokens > 0 {
|
|
r.PromptTokens = u.PromptTokens
|
|
}
|
|
if u.CompletionTokens > 0 {
|
|
r.CompletionTokens = u.CompletionTokens
|
|
}
|
|
x := now
|
|
r.FinishedAt = &x
|
|
r.UpdatedAt = now
|
|
t.pruneLocked(now)
|
|
t.recent = append(t.recent, r)
|
|
if len(t.recent) > t.maxRecent {
|
|
t.recent = append([]Request(nil), t.recent[len(t.recent)-t.maxRecent:]...)
|
|
}
|
|
t.bumpLocked()
|
|
t.mu.Unlock()
|
|
}
|
|
|
|
func (t *Tracker) Drop(id string, status int) {
|
|
t.Finish(id, status, 0, cost.Usage{}, 0)
|
|
}
|
|
|
|
func (t *Tracker) Snapshot() Snapshot {
|
|
now := time.Now().UTC()
|
|
t.mu.Lock()
|
|
t.pruneLocked(now)
|
|
out := make([]Request, 0, len(t.active)+len(t.recent))
|
|
counts := Counts{}
|
|
for _, r := range t.active {
|
|
out = append(out, r)
|
|
addCount(&counts, r.State)
|
|
}
|
|
for _, r := range t.recent {
|
|
out = append(out, r)
|
|
addCount(&counts, r.State)
|
|
}
|
|
counts.Total = len(t.active) + len(t.recent)
|
|
counts.Active = len(t.active)
|
|
version := t.version
|
|
limit := t.maxRecent
|
|
t.mu.Unlock()
|
|
|
|
truncated := limit > 0 && len(out) > limit
|
|
if truncated {
|
|
sort.SliceStable(out, func(i, j int) bool {
|
|
pi, pj := statePriority(out[i].State), statePriority(out[j].State)
|
|
if pi != pj {
|
|
return pi < pj
|
|
}
|
|
if out[i].State == StateCompleted || out[i].State == StateCancelled || out[i].State == StateFailed {
|
|
return out[i].UpdatedAt.After(out[j].UpdatedAt)
|
|
}
|
|
return out[i].QueuedAt.Before(out[j].QueuedAt)
|
|
})
|
|
out = out[:limit]
|
|
}
|
|
sort.SliceStable(out, func(i, j int) bool {
|
|
if out[i].QueuedAt.Equal(out[j].QueuedAt) {
|
|
return out[i].ID < out[j].ID
|
|
}
|
|
return out[i].QueuedAt.Before(out[j].QueuedAt)
|
|
})
|
|
return Snapshot{GeneratedAt: now, Version: version, Counts: counts, Truncated: truncated, Requests: out}
|
|
}
|
|
|
|
func addCount(c *Counts, state string) {
|
|
switch state {
|
|
case StateQueued:
|
|
c.Queued++
|
|
case StateRouting:
|
|
c.Routing++
|
|
case StateRunning:
|
|
c.Running++
|
|
case StateStreaming:
|
|
c.Streaming++
|
|
case StateCompleted:
|
|
c.Completed++
|
|
case StateCancelled:
|
|
c.Cancelled++
|
|
case StateFailed:
|
|
c.Failed++
|
|
}
|
|
}
|
|
|
|
func statePriority(state string) int {
|
|
switch state {
|
|
case StateStreaming:
|
|
return 0
|
|
case StateRunning:
|
|
return 1
|
|
case StateRouting:
|
|
return 2
|
|
case StateQueued:
|
|
return 3
|
|
case StateCompleted:
|
|
return 4
|
|
case StateCancelled:
|
|
return 5
|
|
default:
|
|
return 5
|
|
}
|
|
}
|
|
|
|
func (t *Tracker) Changed() <-chan struct{} {
|
|
t.mu.RLock()
|
|
ch := t.notify
|
|
t.mu.RUnlock()
|
|
return ch
|
|
}
|
|
|
|
func (t *Tracker) update(id string, fn func(*Request, time.Time)) {
|
|
now := time.Now().UTC()
|
|
t.mu.Lock()
|
|
r, ok := t.active[id]
|
|
if ok {
|
|
fn(&r, now)
|
|
t.active[id] = r
|
|
t.bumpLocked()
|
|
}
|
|
t.mu.Unlock()
|
|
}
|
|
|
|
func (t *Tracker) pruneLocked(now time.Time) {
|
|
if len(t.recent) == 0 {
|
|
return
|
|
}
|
|
cutoff := now.Add(-t.recentTTL)
|
|
first := 0
|
|
for first < len(t.recent) {
|
|
r := t.recent[first]
|
|
if r.FinishedAt == nil || r.FinishedAt.After(cutoff) {
|
|
break
|
|
}
|
|
first++
|
|
}
|
|
if first > 0 {
|
|
t.recent = append([]Request(nil), t.recent[first:]...)
|
|
}
|
|
}
|
|
|
|
func (t *Tracker) bumpLocked() {
|
|
t.version++
|
|
close(t.notify)
|
|
t.notify = make(chan struct{})
|
|
}
|