Files
llmgw/internal/control/catalog/pickle.go
T
Ilker ManapandClaude Opus 5.5 1948a787ae
CI / test (push) Successful in 9m15s
Fixes from the second review: grants, flags, updates, setup
- Revoking a server's tools takes them from every level beneath too, and a
  level left naming no tool stays closed ([]) while any level above or
  beneath names tools: an organization's revoke left its teams the copies
  their first grants took, and going back to no list opened the level to
  what was above it. A closed level is shown on the page and in the
  console with an action to open it (store.OpenTools). A key's owner's
  list is among the levels above a key.
- The flag check matches short flags as written again (-m, -hf, -mu were
  let through), takes a long flag's beginning for it only for the engines
  that parse with argparse (vLLM, SGLang, ComfyUI: --li is --listen), so
  llama.cpp's own --chat-template is not refused, and refuses llama.cpp's
  short spellings (-hfr, -hff, -mm, -jf, …).
- A stored model with an adapter by Hugging Face id can be changed again:
  what it already had is not asked about anew.
- The catalog reads a Hugging Face id as the agent does (engines.IsHubID).
- The page's form cannot move llmgw:admin to another address.
- server init skips the memory check only where the gateway already runs,
  not wherever a half-finished first run left server.yaml.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-29 12:31:07 +01:00

92 lines
2.7 KiB
Go

package catalog
import (
"fmt"
"path"
"slices"
"strings"
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
)
// PickleSources are the sources of a model's weights and adapters that are
// in a pickle format, which run code as they are loaded (weights.Pickled).
func (m Model) PickleSources() []string {
var out []string
add := func(s string) {
if s != "" && weights.Pickled(s) && !slices.Contains(out, s) {
out = append(out, s)
}
}
add(m.Source)
for _, p := range m.Profiles {
add(p.Source)
}
for _, a := range m.Adapters {
add(a.Source)
}
return out
}
// CheckPickle refuses a model whose weights are in a pickle format, unless
// it says allow_pickle. Asked when a model is added or its sources change,
// not when the catalog is built: a model already running is refused by its
// node, with the same words, rather than dropped from the catalog.
// Sources in except were there before and are not asked about again.
func (m Model) CheckPickle(except ...string) error {
if m.AllowPickle {
return nil
}
for _, s := range m.PickleSources() {
if slices.Contains(except, s) {
continue
}
return fmt.Errorf("model %q: %s is in a pickle format, which runs code when it is loaded; "+
"use the model's safetensors or GGUF, or set allow_pickle: true if it was checked and has no other form",
m.Name, s)
}
// What an engine downloads itself, by a Hugging Face id, no node looks
// at: vLLM reads a model by id as safetensors only, but not an adapter,
// and TensorRT-LLM neither.
for _, s := range m.UnlookedSources() {
if slices.Contains(except, s) {
continue
}
return fmt.Errorf("model %q: %s is a Hugging Face id its engine downloads itself, and may load in a pickle "+
"format, which runs code; mirror its safetensors and give that, or set allow_pickle: true if it was checked",
m.Name, s)
}
return nil
}
// UnlookedSources are those an engine fetches by a Hugging Face id, past
// the node's check: adapters, and a TensorRT-LLM model's.
func (m Model) UnlookedSources() []string {
var out []string
for _, a := range m.Adapters {
if hubID(a.Source) {
out = append(out, a.Source)
}
}
for _, p := range m.Profiles {
if p.Engine == "trtllm" {
if src := m.SourceOf(p); hubID(src) && !slices.Contains(out, src) {
out = append(out, src)
}
}
}
return out
}
// hubID reports whether a source is a Hugging Face id, as the agent reads
// one (engines.IsHubID) — gpt2 too — and not a file named so, nor a GGUF
// in a repository (org/name:Q4_K_M).
func hubID(s string) bool {
if !engines.IsHubID(s) || strings.Contains(s, ":") {
return false
}
ext := strings.ToLower(path.Ext(s))
return ext != ".gguf" && ext != ".safetensors" && !weights.Pickled(s)
}