67 lines
1.5 KiB
Go
67 lines
1.5 KiB
Go
package readers
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"time"
|
|
|
|
"git.govekar.net/LovroG05/nvllama-watchdog/core"
|
|
"git.govekar.net/LovroG05/nvllama-watchdog/parsers"
|
|
)
|
|
|
|
type LLamaCPPReader struct {
|
|
endpoint string
|
|
model string
|
|
client *http.Client
|
|
}
|
|
|
|
func NewLLamaCPPReader(endpoint, model string) *LLamaCPPReader {
|
|
return &LLamaCPPReader{
|
|
endpoint: endpoint,
|
|
model: model,
|
|
client: &http.Client{Timeout: 5 * time.Second},
|
|
}
|
|
}
|
|
|
|
func (r *LLamaCPPReader) Name() string { return "llamacpp:" + r.model }
|
|
|
|
func (r *LLamaCPPReader) Read(ctx context.Context) ([]core.Metric, error) {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, r.endpoint, nil)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("llamacpp %s: %w", r.model, err)
|
|
}
|
|
|
|
resp, err := r.client.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("llamacpp %s: %w", r.model, err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
body, _ := io.ReadAll(io.LimitReader(resp.Body, 256))
|
|
return nil, fmt.Errorf("llamacpp %s: HTTP %d: %s", r.model, resp.StatusCode, string(body))
|
|
}
|
|
|
|
raw, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("llamacpp %s: read body: %w", r.model, err)
|
|
}
|
|
|
|
metrics, err := parsers.ParsePrometheusText(string(raw))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("llamacpp %s: %w", r.model, err)
|
|
}
|
|
|
|
// Tag every metric with the model instance label
|
|
for i := range metrics {
|
|
if metrics[i].Labels == nil {
|
|
metrics[i].Labels = map[string]string{}
|
|
}
|
|
metrics[i].Labels["model_instance"] = r.model
|
|
}
|
|
|
|
return metrics, nil
|
|
}
|