Files
llmgw/internal/agent/weights/pickle_test.go
T
Ilker ManapandClaude Opus 5.5 fd94265e08
CI / test (push) Failing after 7m17s
Pickles refused however they come; a string temperature passes again
From review. A mirror set already on the node is looked at again, not
trusted for having been fetched once; Pickled cuts a query or a fragment
from any source, hf: too; .pt2 is a pickle. The flag check reads a flag as
the engine does: underscores as dashes and any beginning of a refused flag
as that flag, since argparse takes --trust-remote for --trust-remote-code;
SGLang refuses --load-format, which would undo the agent's safetensors.

TripoSR's preset says allow_pickle: without it the preset could not be
added and the presets pass failed on it each time. A shared ComfyUI
environment's hidden entry takes its first member's allow_pickle, as it
takes its source.

The GenAI trace attributes made temperature and top_p typed fields of the
request's parse, so "0.7" as a string was a 400; they are read loosely,
as they are read only for the trace.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-28 22:19:06 +01:00

67 lines
2.3 KiB
Go

package weights
import (
"errors"
"sync/atomic"
"testing"
)
func TestPickleFormatsAreKnownByTheirNames(t *testing.T) {
for name, want := range map[string]bool{
"pytorch_model.bin": true, "model.ckpt": true, "x.PT": true, "w.pth": true, "a.pkl": true,
"https://h/r/resolve/main/model.ckpt?download=true": true,
"hf:o/r/model.bin?download=true": true, "hf:o/r/model.ckpt#x": true, "m.pt2": true,
"model.safetensors": false, "m-Q4_K_M.gguf": false, "Qwen/Qwen2.5-7B-Instruct": false,
"config.json": false, "": false,
} {
if got := Pickled(name); got != want {
t.Errorf("Pickled(%q) = %v, want %v", name, got, want)
}
}
}
func TestAPickleFileIsRefusedBeforeItIsFetched(t *testing.T) {
var hits atomic.Int64
srv := serveBytes(t, []byte("not really a pickle"), &hits)
f := &Fetcher{CacheDir: t.TempDir()}
_, err := f.Ensure(t.Context(), srv.URL+"/model.ckpt", false)
if !errors.Is(err, ErrPickle) {
t.Fatalf("a .ckpt was not refused: %v", err)
}
if hits.Load() != 0 {
t.Errorf("the refused file was fetched %d times", hits.Load())
}
if _, err := f.Ensure(t.Context(), srv.URL+"/model.ckpt", true); err != nil {
t.Errorf("allowed, the .ckpt was still refused: %v", err)
}
}
func TestASetWithAPickleFileIsRefusedWhole(t *testing.T) {
f, srv, hits := mirrorOf(t, map[string]string{
"pytorch_model.bin": "pickle",
"config.json": "{}",
})
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); !errors.Is(err, ErrPickle) {
t.Fatalf("a set holding pytorch_model.bin was not refused: %v", err)
}
if *hits != 0 {
t.Errorf("%d files of the refused set were fetched", *hits)
}
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", true); err != nil {
t.Errorf("allowed, the set was still refused: %v", err)
}
}
func TestASetAlreadyHereIsLookedAtAgain(t *testing.T) {
// Fetched once for a model that allowed pickles, or before they were
// refused: another asking for it without leave is refused all the same.
f, srv, _ := mirrorOf(t, map[string]string{"pytorch_model.bin": "pickle", "config.json": "{}"})
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", true); err != nil {
t.Fatal(err)
}
if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); !errors.Is(err, ErrPickle) {
t.Errorf("a set already on the node was not looked at: %v", err)
}
}