-
This commit is contained in:
@@ -0,0 +1,335 @@
|
||||
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{})
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package liveflow
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/example/ollama-fair-gateway/internal/cost"
|
||||
)
|
||||
|
||||
func TestLifecycle(t *testing.T) {
|
||||
tr := New(time.Second, 10)
|
||||
tr.Begin(Request{ID: "r1", Tenant: "a", Actor: "u", Model: "m", EstimatedCredits: 2})
|
||||
snap := tr.Snapshot()
|
||||
if len(snap.Requests) != 1 || snap.Requests[0].State != StateQueued {
|
||||
t.Fatalf("unexpected initial snapshot: %#v", snap)
|
||||
}
|
||||
tr.MarkRouting("r1", "w1", 25*time.Millisecond)
|
||||
tr.MarkRunning("r1")
|
||||
tr.Progress("r1", 160, cost.Usage{PromptTokens: 12})
|
||||
snap = tr.Snapshot()
|
||||
if got := snap.Requests[0]; got.State != StateStreaming || got.Worker != "w1" || got.PromptTokens != 12 || got.CompletionTokens == 0 {
|
||||
t.Fatalf("unexpected progress: %#v", got)
|
||||
}
|
||||
tr.Finish("r1", 200, 1.25, cost.Usage{PromptTokens: 12, CompletionTokens: 8}, 50*time.Millisecond)
|
||||
snap = tr.Snapshot()
|
||||
if got := snap.Requests[0]; got.State != StateCompleted || got.Status != 200 || got.ActualCredits != 1.25 || got.CompletionTokens != 8 {
|
||||
t.Fatalf("unexpected final: %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChangedNotifies(t *testing.T) {
|
||||
tr := New(time.Second, 10)
|
||||
ch := tr.Changed()
|
||||
tr.Begin(Request{ID: "r"})
|
||||
select {
|
||||
case <-ch:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("tracker did not notify")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSnapshotIsBoundedButCountsRemainExact(t *testing.T) {
|
||||
tr := New(time.Second, 2)
|
||||
tr.Begin(Request{ID: "a"})
|
||||
tr.Begin(Request{ID: "b"})
|
||||
tr.Begin(Request{ID: "c"})
|
||||
s := tr.Snapshot()
|
||||
if !s.Truncated || len(s.Requests) != 2 || s.Counts.Active != 3 || s.Counts.Queued != 3 {
|
||||
t.Fatalf("unexpected bounded snapshot: %#v", s)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user