104 lines
2.3 KiB
Go
104 lines
2.3 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sort"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/example/ollama-fair-gateway/internal/liveflow"
|
|
)
|
|
|
|
var errJobCancelled = errors.New("job cancelled by administrator")
|
|
|
|
type jobEntry struct {
|
|
ID string `json:"id"`
|
|
Tenant string `json:"tenant"`
|
|
Actor string `json:"actor"`
|
|
Application string `json:"application,omitempty"`
|
|
ServiceClass string `json:"service_class,omitempty"`
|
|
Model string `json:"model,omitempty"`
|
|
Path string `json:"path"`
|
|
API string `json:"api"`
|
|
Worker string `json:"worker,omitempty"`
|
|
CreatedAt time.Time `json:"created_at"`
|
|
CancelledAt *time.Time `json:"cancelled_at,omitempty"`
|
|
Cancelling bool `json:"cancelling,omitempty"`
|
|
}
|
|
|
|
type jobView struct {
|
|
liveflow.Request
|
|
Cancellable bool `json:"cancellable"`
|
|
Cancelling bool `json:"cancelling,omitempty"`
|
|
}
|
|
|
|
type jobManager struct {
|
|
mu sync.RWMutex
|
|
jobs map[string]jobEntry
|
|
cancel map[string]context.CancelCauseFunc
|
|
}
|
|
|
|
func newJobManager() *jobManager {
|
|
return &jobManager{jobs: make(map[string]jobEntry), cancel: make(map[string]context.CancelCauseFunc)}
|
|
}
|
|
|
|
func (m *jobManager) register(j jobEntry, cancel context.CancelCauseFunc) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
m.jobs[j.ID] = j
|
|
m.cancel[j.ID] = cancel
|
|
}
|
|
|
|
func (m *jobManager) setWorker(id, worker string) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
j, ok := m.jobs[id]
|
|
if !ok {
|
|
return
|
|
}
|
|
j.Worker = worker
|
|
m.jobs[id] = j
|
|
}
|
|
|
|
func (m *jobManager) finish(id string) {
|
|
m.mu.Lock()
|
|
delete(m.jobs, id)
|
|
delete(m.cancel, id)
|
|
m.mu.Unlock()
|
|
}
|
|
|
|
func (m *jobManager) cancelJob(id string) bool {
|
|
m.mu.Lock()
|
|
cancel := m.cancel[id]
|
|
j, ok := m.jobs[id]
|
|
if cancel == nil || !ok {
|
|
m.mu.Unlock()
|
|
return false
|
|
}
|
|
if !j.Cancelling {
|
|
now := time.Now().UTC()
|
|
j.Cancelling = true
|
|
j.CancelledAt = &now
|
|
m.jobs[id] = j
|
|
}
|
|
m.mu.Unlock()
|
|
cancel(errJobCancelled)
|
|
return true
|
|
}
|
|
|
|
func (m *jobManager) list() []jobEntry {
|
|
m.mu.RLock()
|
|
out := make([]jobEntry, 0, len(m.jobs))
|
|
for _, j := range m.jobs {
|
|
out = append(out, j)
|
|
}
|
|
m.mu.RUnlock()
|
|
sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt.Before(out[j].CreatedAt) })
|
|
return out
|
|
}
|
|
|
|
func isAdminJobCancel(ctx context.Context) bool {
|
|
return errors.Is(context.Cause(ctx), errJobCancelled)
|
|
}
|