377 lines
9.3 KiB
Go
377 lines
9.3 KiB
Go
package telemetry
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/example/ollama-fair-gateway/internal/config"
|
|
)
|
|
|
|
type Context struct {
|
|
TraceID string
|
|
RootSpanID string
|
|
ParentSpanID string
|
|
Sampled bool
|
|
}
|
|
|
|
type Interval struct {
|
|
Name string
|
|
Start time.Time
|
|
End time.Time
|
|
Attrs map[string]any
|
|
}
|
|
|
|
type Record struct {
|
|
Trace Context
|
|
RequestID string
|
|
API string
|
|
Path string
|
|
Tenant string
|
|
Actor string
|
|
Application string
|
|
ServiceClass string
|
|
Model string
|
|
Alias string
|
|
Worker string
|
|
Started time.Time
|
|
Finished time.Time
|
|
FirstByte time.Duration
|
|
Status int
|
|
PromptTokens int64
|
|
OutputTokens int64
|
|
CachedTokens int64
|
|
Credits float64
|
|
Error string
|
|
Intervals []Interval
|
|
}
|
|
|
|
type Exporter struct {
|
|
cfg config.OpenTelemetryConfig
|
|
client *http.Client
|
|
ch chan []span
|
|
stop chan struct{}
|
|
done chan struct{}
|
|
closed atomic.Bool
|
|
dropped atomic.Uint64
|
|
failed atomic.Uint64
|
|
exported atomic.Uint64
|
|
}
|
|
|
|
type span struct {
|
|
TraceID string
|
|
SpanID string
|
|
ParentSpanID string
|
|
Name string
|
|
Kind int
|
|
Start time.Time
|
|
End time.Time
|
|
Attrs map[string]any
|
|
Error string
|
|
}
|
|
|
|
func New(cfg config.OpenTelemetryConfig) *Exporter {
|
|
if !cfg.Enabled {
|
|
return nil
|
|
}
|
|
if cfg.BatchSize <= 0 {
|
|
cfg.BatchSize = 128
|
|
}
|
|
if cfg.FlushInterval == 0 {
|
|
cfg.FlushInterval = config.Duration(2 * time.Second)
|
|
}
|
|
e := &Exporter{cfg: cfg, client: &http.Client{Timeout: 10 * time.Second}, ch: make(chan []span, cfg.BatchSize*4), stop: make(chan struct{}), done: make(chan struct{})}
|
|
go e.loop()
|
|
return e
|
|
}
|
|
|
|
func (e *Exporter) NewTrace(traceparent string) Context {
|
|
if e == nil {
|
|
return Context{}
|
|
}
|
|
ratio := e.cfg.SampleRatio
|
|
if ratio <= 0 || ratio > 1 {
|
|
ratio = 1
|
|
}
|
|
sampled := ratio >= 1 || randomFraction() < ratio
|
|
traceID, parent := parseTraceParent(traceparent)
|
|
if traceID == "" {
|
|
traceID = randomHex(16)
|
|
}
|
|
return Context{TraceID: traceID, RootSpanID: randomHex(8), ParentSpanID: parent, Sampled: sampled}
|
|
}
|
|
func (e *Exporter) TraceParent(c Context) string {
|
|
if e == nil || !c.Sampled {
|
|
return ""
|
|
}
|
|
return "00-" + c.TraceID + "-" + c.RootSpanID + "-01"
|
|
}
|
|
func (e *Exporter) Dropped() uint64 {
|
|
if e == nil {
|
|
return 0
|
|
}
|
|
return e.dropped.Load()
|
|
}
|
|
func (e *Exporter) Failed() uint64 {
|
|
if e == nil {
|
|
return 0
|
|
}
|
|
return e.failed.Load()
|
|
}
|
|
func (e *Exporter) Exported() uint64 {
|
|
if e == nil {
|
|
return 0
|
|
}
|
|
return e.exported.Load()
|
|
}
|
|
|
|
func (e *Exporter) Record(r Record) {
|
|
if e == nil || !r.Trace.Sampled || e.closed.Load() {
|
|
return
|
|
}
|
|
if r.Finished.IsZero() {
|
|
r.Finished = time.Now().UTC()
|
|
}
|
|
attrs := map[string]any{
|
|
"gen_ai.operation.name": operationName(r.Path),
|
|
"gen_ai.provider.name": "ollama",
|
|
"gen_ai.request.model": r.Model,
|
|
"gen_ai.usage.input_tokens": r.PromptTokens,
|
|
"gen_ai.usage.output_tokens": r.OutputTokens,
|
|
"gen_ai.usage.cache_read.input_tokens": r.CachedTokens,
|
|
"http.request.method": "POST",
|
|
"http.response.status_code": r.Status,
|
|
"server.address": r.Worker,
|
|
"ollama.gateway.request_id": r.RequestID,
|
|
"ollama.gateway.api": r.API,
|
|
"ollama.gateway.tenant": r.Tenant,
|
|
"ollama.gateway.actor": r.Actor,
|
|
"ollama.gateway.application": r.Application,
|
|
"ollama.gateway.service_class": r.ServiceClass,
|
|
"ollama.gateway.credits": r.Credits,
|
|
}
|
|
if r.Alias != "" {
|
|
attrs["ollama.gateway.model_alias"] = r.Alias
|
|
}
|
|
if r.FirstByte > 0 {
|
|
attrs["gen_ai.response.time_to_first_chunk"] = r.FirstByte.Seconds()
|
|
}
|
|
root := span{TraceID: r.Trace.TraceID, SpanID: r.Trace.RootSpanID, ParentSpanID: r.Trace.ParentSpanID, Name: operationName(r.Path) + " " + r.Model, Kind: 2, Start: r.Started, End: r.Finished, Attrs: attrs, Error: r.Error}
|
|
spans := []span{root}
|
|
for _, iv := range r.Intervals {
|
|
if iv.Start.IsZero() || iv.End.IsZero() || iv.End.Before(iv.Start) {
|
|
continue
|
|
}
|
|
a := map[string]any{"ollama.gateway.request_id": r.RequestID}
|
|
for k, v := range iv.Attrs {
|
|
a[k] = v
|
|
}
|
|
spans = append(spans, span{TraceID: r.Trace.TraceID, SpanID: randomHex(8), ParentSpanID: r.Trace.RootSpanID, Name: iv.Name, Kind: 1, Start: iv.Start, End: iv.End, Attrs: a})
|
|
}
|
|
select {
|
|
case e.ch <- spans:
|
|
default:
|
|
e.dropped.Add(uint64(len(spans)))
|
|
}
|
|
}
|
|
|
|
func (e *Exporter) Close(ctx context.Context) error {
|
|
if e == nil || !e.closed.CompareAndSwap(false, true) {
|
|
return nil
|
|
}
|
|
close(e.stop)
|
|
select {
|
|
case <-e.done:
|
|
return nil
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
}
|
|
}
|
|
|
|
func (e *Exporter) loop() {
|
|
defer close(e.done)
|
|
tick := time.NewTicker(e.cfg.FlushInterval.Value())
|
|
defer tick.Stop()
|
|
batch := make([]span, 0, e.cfg.BatchSize)
|
|
flush := func() {
|
|
if len(batch) == 0 {
|
|
return
|
|
}
|
|
if err := e.export(batch); err != nil {
|
|
e.failed.Add(uint64(len(batch)))
|
|
} else {
|
|
e.exported.Add(uint64(len(batch)))
|
|
}
|
|
batch = batch[:0]
|
|
}
|
|
for {
|
|
select {
|
|
case xs := <-e.ch:
|
|
batch = append(batch, xs...)
|
|
if len(batch) >= e.cfg.BatchSize {
|
|
flush()
|
|
}
|
|
case <-tick.C:
|
|
flush()
|
|
case <-e.stop:
|
|
for {
|
|
select {
|
|
case xs := <-e.ch:
|
|
batch = append(batch, xs...)
|
|
default:
|
|
flush()
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (e *Exporter) export(spans []span) error {
|
|
endpoint := strings.TrimSpace(e.cfg.Endpoint)
|
|
if endpoint == "" {
|
|
return fmt.Errorf("empty OTLP endpoint")
|
|
}
|
|
u, err := url.Parse(endpoint)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if u.Path == "" || u.Path == "/" {
|
|
u.Path = "/v1/traces"
|
|
} else if !strings.HasSuffix(u.Path, "/v1/traces") {
|
|
u.Path = strings.TrimRight(u.Path, "/") + "/v1/traces"
|
|
}
|
|
payload := map[string]any{"resourceSpans": []any{map[string]any{"resource": map[string]any{"attributes": attrsToOTLP(map[string]any{"service.name": e.cfg.ServiceName, "service.version": e.cfg.ServiceVersion, "telemetry.sdk.name": "ollama-fair-gateway", "telemetry.sdk.language": "go"})}, "scopeSpans": []any{map[string]any{"scope": map[string]any{"name": "ollama-fair-gateway"}, "spans": spansToOTLP(spans)}}}}}
|
|
b, _ := json.Marshal(payload)
|
|
req, err := http.NewRequest(http.MethodPost, u.String(), bytes.NewReader(b))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
for k, v := range e.cfg.Headers {
|
|
req.Header.Set(k, v)
|
|
}
|
|
resp, err := e.client.Do(req)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer resp.Body.Close()
|
|
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<20))
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return fmt.Errorf("OTLP HTTP %d", resp.StatusCode)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func spansToOTLP(in []span) []any {
|
|
out := make([]any, 0, len(in))
|
|
for _, s := range in {
|
|
x := map[string]any{"traceId": s.TraceID, "spanId": s.SpanID, "name": s.Name, "kind": s.Kind, "startTimeUnixNano": fmt.Sprintf("%d", s.Start.UnixNano()), "endTimeUnixNano": fmt.Sprintf("%d", s.End.UnixNano()), "attributes": attrsToOTLP(s.Attrs)}
|
|
if s.ParentSpanID != "" {
|
|
x["parentSpanId"] = s.ParentSpanID
|
|
}
|
|
if s.Error != "" {
|
|
x["status"] = map[string]any{"code": 2, "message": s.Error}
|
|
}
|
|
out = append(out, x)
|
|
}
|
|
return out
|
|
}
|
|
func attrsToOTLP(in map[string]any) []any {
|
|
keys := make([]string, 0, len(in))
|
|
for k, v := range in {
|
|
if !empty(v) {
|
|
keys = append(keys, k)
|
|
}
|
|
}
|
|
sortStrings(keys)
|
|
out := make([]any, 0, len(keys))
|
|
for _, k := range keys {
|
|
out = append(out, map[string]any{"key": k, "value": otlpValue(in[k])})
|
|
}
|
|
return out
|
|
}
|
|
func otlpValue(v any) map[string]any {
|
|
switch x := v.(type) {
|
|
case string:
|
|
return map[string]any{"stringValue": x}
|
|
case bool:
|
|
return map[string]any{"boolValue": x}
|
|
case int:
|
|
return map[string]any{"intValue": fmt.Sprintf("%d", x)}
|
|
case int64:
|
|
return map[string]any{"intValue": fmt.Sprintf("%d", x)}
|
|
case uint64:
|
|
return map[string]any{"intValue": fmt.Sprintf("%d", x)}
|
|
case float64:
|
|
return map[string]any{"doubleValue": x}
|
|
default:
|
|
return map[string]any{"stringValue": fmt.Sprint(x)}
|
|
}
|
|
}
|
|
func empty(v any) bool {
|
|
switch x := v.(type) {
|
|
case string:
|
|
return x == ""
|
|
case int:
|
|
return x == 0
|
|
case int64:
|
|
return x == 0
|
|
case uint64:
|
|
return x == 0
|
|
case float64:
|
|
return x == 0
|
|
}
|
|
return false
|
|
}
|
|
func operationName(path string) string {
|
|
p := strings.ToLower(path)
|
|
switch {
|
|
case strings.Contains(p, "embed"):
|
|
return "embeddings"
|
|
case strings.Contains(p, "chat") || strings.Contains(p, "messages"):
|
|
return "chat"
|
|
default:
|
|
return "text_completion"
|
|
}
|
|
}
|
|
func randomHex(n int) string { b := make([]byte, n); _, _ = rand.Read(b); return hex.EncodeToString(b) }
|
|
func randomFraction() float64 {
|
|
b := make([]byte, 8)
|
|
_, _ = rand.Read(b)
|
|
var x uint64
|
|
for _, z := range b {
|
|
x = x<<8 | uint64(z)
|
|
}
|
|
return float64(x>>11) / float64(uint64(1)<<53)
|
|
}
|
|
func parseTraceParent(v string) (string, string) {
|
|
parts := strings.Split(strings.TrimSpace(v), "-")
|
|
if len(parts) != 4 || len(parts[1]) != 32 || len(parts[2]) != 16 {
|
|
return "", ""
|
|
}
|
|
if _, e := hex.DecodeString(parts[1]); e != nil {
|
|
return "", ""
|
|
}
|
|
if _, e := hex.DecodeString(parts[2]); e != nil {
|
|
return "", ""
|
|
}
|
|
return strings.ToLower(parts[1]), strings.ToLower(parts[2])
|
|
}
|
|
func sortStrings(x []string) {
|
|
for i := 1; i < len(x); i++ {
|
|
for j := i; j > 0 && x[j] < x[j-1]; j-- {
|
|
x[j], x[j-1] = x[j-1], x[j]
|
|
}
|
|
}
|
|
}
|