93 lines
2.0 KiB
Go
93 lines
2.0 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"flag"
|
|
"log"
|
|
"os/signal"
|
|
"syscall"
|
|
"time"
|
|
|
|
"git.govekar.net/LovroG05/nvllama-watchdog/config"
|
|
"git.govekar.net/LovroG05/nvllama-watchdog/core"
|
|
"git.govekar.net/LovroG05/nvllama-watchdog/exporters"
|
|
"git.govekar.net/LovroG05/nvllama-watchdog/readers"
|
|
)
|
|
|
|
func main() {
|
|
cfgPath := flag.String("config", "config/config.yaml", "path to the YAML config file")
|
|
flag.Parse()
|
|
|
|
cfg, err := config.Load(*cfgPath)
|
|
if err != nil {
|
|
log.Fatalf("config: %v", err)
|
|
}
|
|
|
|
// Build readers
|
|
var rs []core.Reader
|
|
for _, rc := range cfg.Readers {
|
|
switch rc.Type {
|
|
case "llamacpp":
|
|
rs = append(rs, readers.NewLLamaCPPReader(rc.Endpoint, rc.Model))
|
|
case "nvidia_smi":
|
|
rs = append(rs, readers.NewNvidiaSmiReader())
|
|
case "system":
|
|
rs = append(rs, readers.NewSystemReader(rc.Enabled))
|
|
}
|
|
}
|
|
|
|
// Build exporters
|
|
var es []core.Exporter
|
|
var prom *exporters.PrometheusExporter
|
|
for _, ec := range cfg.Exporters {
|
|
switch ec.Type {
|
|
case "prometheus":
|
|
prom = exporters.NewPrometheusExporter(ec.Addr)
|
|
es = append(es, prom)
|
|
case "dev":
|
|
es = append(es, exporters.NewDevExporter(ec.JSONFile))
|
|
case "influxdb":
|
|
es = append(es, exporters.NewInfluxDBExporter(ec.URL, ec.Token, ec.User, ec.Pass))
|
|
}
|
|
}
|
|
|
|
// Start Prometheus HTTP listener in background
|
|
if prom != nil {
|
|
go func() {
|
|
log.Printf("prometheus exporter listening on %s", prom.Addr())
|
|
if err := prom.Start(); err != nil {
|
|
log.Fatalf("prometheus: %v", err)
|
|
}
|
|
}()
|
|
}
|
|
|
|
// Scrape loop
|
|
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
|
defer stop()
|
|
|
|
ticker := time.NewTicker(cfg.ScrapeInterval)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-ticker.C:
|
|
var all []core.Metric
|
|
for _, r := range rs {
|
|
ms, err := r.Read(ctx)
|
|
if err != nil {
|
|
log.Printf("reader %s: %v", r.Name(), err)
|
|
continue
|
|
}
|
|
all = append(all, ms...)
|
|
}
|
|
for _, e := range es {
|
|
if err := e.Write(ctx, all); err != nil {
|
|
log.Printf("exporter %s: %v", e.Name(), err)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|