141 lines
3.8 KiB
Go
141 lines
3.8 KiB
Go
package health
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
)
|
|
|
|
// modelProps is the part of a /props body that carries a model's context size.
|
|
type modelProps struct {
|
|
Generation struct {
|
|
NCtx *int `json:"n_ctx"`
|
|
} `json:"default_generation_settings"`
|
|
TotalSlots *int `json:"total_slots"`
|
|
}
|
|
|
|
// props reads <base>/props best-effort. A request that fails because ctx is
|
|
// done yields a cancelled result so the caller records nothing; any other
|
|
// outcome (status, body, or missing fields) leaves context unknown without
|
|
// failing the poll. When the body says "role":"router" the host carries no
|
|
// context of its own, so NCtx and Slots stay 0 whatever else it reports.
|
|
func (t *Table) props(ctx context.Context, base string) (int, int, pollResult) {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+"/props", nil)
|
|
if err != nil {
|
|
if ctx.Err() != nil {
|
|
return 0, 0, pollResult{cancelled: true}
|
|
}
|
|
return 0, 0, pollResult{}
|
|
}
|
|
resp, err := t.client.Do(req)
|
|
if err != nil {
|
|
if ctx.Err() != nil {
|
|
return 0, 0, pollResult{cancelled: true}
|
|
}
|
|
return 0, 0, pollResult{}
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return 0, 0, pollResult{}
|
|
}
|
|
var p struct {
|
|
Role string `json:"role"`
|
|
Generation struct {
|
|
NCtx *int `json:"n_ctx"`
|
|
} `json:"default_generation_settings"`
|
|
TotalSlots *int `json:"total_slots"`
|
|
}
|
|
if err := json.NewDecoder(io.LimitReader(resp.Body, MaxModelsBody)).Decode(&p); err != nil {
|
|
return 0, 0, pollResult{}
|
|
}
|
|
if p.Role == "router" {
|
|
return 0, 0, pollResult{}
|
|
}
|
|
nctx, slots := 0, 0
|
|
if p.Generation.NCtx != nil {
|
|
nctx = *p.Generation.NCtx
|
|
}
|
|
if p.TotalSlots != nil {
|
|
slots = *p.TotalSlots
|
|
}
|
|
if nctx < 0 {
|
|
nctx = 0
|
|
}
|
|
if slots < 0 {
|
|
slots = 0
|
|
}
|
|
return nctx, slots, pollResult{}
|
|
}
|
|
|
|
// propsModels asks /props?model= for each loaded model and returns a fresh,
|
|
// non-nil map of what each answered. A failed or malformed answer leaves that
|
|
// id absent; a request that fails because ctx is done cancels the whole poll.
|
|
// No model is ever asked that is not in loaded.
|
|
func (t *Table) propsModels(ctx context.Context, base string, loaded []string) (map[string]ModelCtx, pollResult) {
|
|
models := make(map[string]ModelCtx, len(loaded))
|
|
for _, id := range loaded {
|
|
mc, ok, r := t.propsModel(ctx, base, id)
|
|
if r.cancelled || r.reason != "" {
|
|
return nil, r
|
|
}
|
|
if !ok {
|
|
continue
|
|
}
|
|
models[id] = mc
|
|
}
|
|
return models, pollResult{}
|
|
}
|
|
|
|
// propsModel reads <base>/props?model=<id> best-effort. ok is false when the
|
|
// answer failed or was malformed (the model stays unknown); a request that
|
|
// fails because ctx is done yields a cancelled result so the caller records
|
|
// nothing.
|
|
func (t *Table) propsModel(ctx context.Context, base, id string) (ModelCtx, bool, pollResult) {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+"/props?model="+url.QueryEscape(id), nil)
|
|
if err != nil {
|
|
if ctx.Err() != nil {
|
|
return ModelCtx{}, false, pollResult{cancelled: true}
|
|
}
|
|
return ModelCtx{}, false, pollResult{}
|
|
}
|
|
resp, err := t.client.Do(req)
|
|
if err != nil {
|
|
if ctx.Err() != nil {
|
|
return ModelCtx{}, false, pollResult{cancelled: true}
|
|
}
|
|
return ModelCtx{}, false, pollResult{}
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return ModelCtx{}, false, pollResult{}
|
|
}
|
|
var p modelProps
|
|
if err := json.NewDecoder(io.LimitReader(resp.Body, MaxModelsBody)).Decode(&p); err != nil {
|
|
return ModelCtx{}, false, pollResult{}
|
|
}
|
|
mc, r := clampCtx(p.Generation.NCtx, p.TotalSlots)
|
|
return mc, true, r
|
|
}
|
|
|
|
// clampCtx turns the two optional fields into a non-negative context and slots.
|
|
func clampCtx(nctx, slots *int) (ModelCtx, pollResult) {
|
|
c := ModelCtx{}
|
|
if nctx != nil {
|
|
c.NCtx = *nctx
|
|
}
|
|
if slots != nil {
|
|
c.Slots = *slots
|
|
}
|
|
if c.NCtx < 0 {
|
|
c.NCtx = 0
|
|
}
|
|
if c.Slots < 0 {
|
|
c.Slots = 0
|
|
}
|
|
return c, pollResult{}
|
|
}
|