diff --git a/CHANGELOG.md b/CHANGELOG.md index 03c8596..eb552b8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -39,6 +39,17 @@ APIs change between minors without a deprecation period. ### Fixed +- **Weights that run code got past the refusal in four ways.** A mirror + version already on a node was used without being looked at again; a + `?query` or `#fragment` on an `hf:` source hid the file's name; `.pt2` was + not on the list; and a profile could write `--load_format`, + `--trust_remote_code` or an abbreviation such as `--trust-remote`, which + vLLM and argparse read as the flag, or give SGLang `--load-format`. All + refused now. TripoSR's preset says `allow_pickle`, without which it could + not be added, and a shared ComfyUI environment takes its first member's. +- **A temperature written as a string was refused.** Since the GenAI traces + read it, `"temperature": "0.7"` got a 400 before reaching the engine; it + is read loosely again, for the trace only. - **A reused answer could carry another caller's personal data.** With `restore: true` and the answer cache on, two callers' requests that redacted to the same text shared one answer — with the first caller's diff --git a/docs/install.md b/docs/install.md index de8cf47..55c2816 100644 --- a/docs/install.md +++ b/docs/install.md @@ -891,8 +891,9 @@ files were checked, can say so: allow_pickle: true ``` -A file put on the machine by hand, as a path, is the operator's own and is not -looked at. +A local path is judged by its name as any source is; a local directory is +not looked into, nor a model an engine downloads itself other than by its +Hugging Face id (an adapter by id, TensorRT-LLM's checkpoints). Put them where the node can read them. `scp`, a shared filesystem, or the engine's own downloader are all valid, and which you want depends on how many @@ -2326,6 +2327,7 @@ triposr`. What that adds is this entry, and its workflow is built in: ```yaml - name: triposr source: hf:stabilityai/TripoSR/model.ckpt + allow_pickle: true # published as this checkpoint only comfyui: workflow: triposr # built in; a model of your own names a JSON file on every control plane output: meshes diff --git a/internal/agent/engines/pickle_test.go b/internal/agent/engines/pickle_test.go index 3f91851..2a7c905 100644 --- a/internal/agent/engines/pickle_test.go +++ b/internal/agent/engines/pickle_test.go @@ -20,3 +20,25 @@ func TestAModelDownloadedByTheEngineIsReadAsSafetensorsOnly(t *testing.T) { t.Error("allow_pickle still forced safetensors") } } + +func TestAForbiddenFlagIsKnownHoweverItIsSpelled(t *testing.T) { + for engine, args := range map[string][]string{ + "vllm": {"--trust_remote_code"}, + "sglang": {"--load-format", "pt"}, + } { + if UnsafeArgs(engine, args) == nil { + t.Errorf("%s %v was let through", engine, args) + } + } + for _, a := range [][]string{{"--load_format", "pt"}, {"--trust-remote"}, {"--load-form=pt"}} { + if UnsafeArgs("vllm", a) == nil { + t.Errorf("vllm %v was let through", a) + } + } + // What is not one of them, nor the start of one, still passes. + for _, a := range [][]string{{"--max-model-len", "8192"}, {"--tokenizer-mode", "mistral"}, {"--gpu-memory-utilization=0.9"}} { + if err := UnsafeArgs("vllm", a); err != nil { + t.Errorf("vllm %v was refused: %v", a, err) + } + } +} diff --git a/internal/agent/engines/unsafeenv.go b/internal/agent/engines/unsafeenv.go index fb3a171..32fad5b 100644 --- a/internal/agent/engines/unsafeenv.go +++ b/internal/agent/engines/unsafeenv.go @@ -119,6 +119,7 @@ var flagsNever = map[string]map[string]string{ "--file-storage-path": "writes files", "--log-requests": "logs prompts", "--model-path": "is the model's source, set by the agent", "--served-model-name": "is set by the agent", "--lora-paths": "loads files", "--download-dir": "writes files", + "--load-format": "changes how weights are loaded, some ways running code", }, // ComfyUI's main.py, behind the gate that sets where it listens and // where its files are. @@ -164,9 +165,23 @@ var flagsNever = map[string]map[string]string{ // unsafeFlag is why an engine may not be given a flag, or "". func unsafeFlag(engine, flag, _ string) string { name, _, _ := strings.Cut(flag, "=") + if !strings.HasPrefix(name, "--") { + return "" + } + // Written as the engine would read it: vLLM takes --load_format for + // --load-format, TensorRT-LLM the other way round, and argparse takes + // any unambiguous beginning of a flag for the flag — --trust-remote is + // --trust-remote-code. + dashed := "--" + strings.ReplaceAll(name[2:], "_", "-") for _, set := range []string{"", engine} { - if why, ok := flagsNever[set][name]; ok && why != "" { - return why + for never, why := range flagsNever[set] { + if why == "" { + continue + } + n := "--" + strings.ReplaceAll(never[2:], "_", "-") + if n == dashed || (len(dashed) > 4 && strings.HasPrefix(n, dashed)) { + return why + } } } return "" diff --git a/internal/agent/weights/pickle.go b/internal/agent/weights/pickle.go index f66610c..5a0d12e 100644 --- a/internal/agent/weights/pickle.go +++ b/internal/agent/weights/pickle.go @@ -16,11 +16,13 @@ import ( // pickleExts are the extensions of PyTorch's pickle formats, and of pickle // itself. A .bin beside safetensors is the same model in pickle form. -var pickleExts = []string{".bin", ".pt", ".pth", ".ckpt", ".pkl", ".pickle"} +var pickleExts = []string{".bin", ".pt", ".pt2", ".pth", ".ckpt", ".pkl", ".pickle"} // Pickled reports whether a file name or URL is in a pickle format. func Pickled(name string) bool { - if i := strings.IndexAny(name, "?#"); i >= 0 && strings.Contains(name, "://") { + // A query or a fragment is not the file's name, on a URL or an hf: + // source alike: model.bin?download=true is still model.bin. + if i := strings.IndexAny(name, "?#"); i >= 0 { name = name[:i] } ext := strings.ToLower(path.Ext(name)) diff --git a/internal/agent/weights/pickle_test.go b/internal/agent/weights/pickle_test.go index eb9d50f..adc403d 100644 --- a/internal/agent/weights/pickle_test.go +++ b/internal/agent/weights/pickle_test.go @@ -10,6 +10,7 @@ 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, } { @@ -51,3 +52,15 @@ func TestASetWithAPickleFileIsRefusedWhole(t *testing.T) { 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) + } +} diff --git a/internal/agent/weights/set.go b/internal/agent/weights/set.go index d1fd369..6393acd 100644 --- a/internal/agent/weights/set.go +++ b/internal/agent/weights/set.go @@ -56,6 +56,16 @@ func (f *Fetcher) ensureSet(ctx context.Context, source string, allowPickle bool // exactly the failure the marker exists to prevent. done := filepath.Join(dir, ".complete") if _, err := os.Stat(done); err == nil { + // Already here, maybe fetched for a model that allowed pickles, or + // before they were refused: looked at again, not trusted. + if !allowPickle { + entries, _ := os.ReadDir(dir) + for _, e := range entries { + if Pickled(e.Name()) { + return "", refusePickle(source + e.Name()) + } + } + } return dir, nil } diff --git a/internal/control/catalog/pickle_test.go b/internal/control/catalog/pickle_test.go index ca0da0d..1eaaabc 100644 --- a/internal/control/catalog/pickle_test.go +++ b/internal/control/catalog/pickle_test.go @@ -28,3 +28,13 @@ func TestAModelInAPickleFormatIsRefusedUnlessItSaysSo(t *testing.T) { t.Errorf("a GGUF was refused: %v", err) } } + +func TestTheTripoSRPresetSaysItsCheckpointIsAllowed(t *testing.T) { + p, ok := PresetNamed("triposr") + if !ok { + t.Fatal("no triposr preset") + } + if err := p.Model.CheckPickle(); err != nil { + t.Errorf("the preset llmgw ships is refused by its own check: %v", err) + } +} diff --git a/internal/control/catalog/presets.go b/internal/control/catalog/presets.go index 1771a58..c9d1abb 100644 --- a/internal/control/catalog/presets.go +++ b/internal/control/catalog/presets.go @@ -52,6 +52,9 @@ func Presets() []Preset { Name: "triposr", Description: "TripoSR: a 3D model (GLB) from a picture of one object on a plain background, as a job", Source: "hf:stabilityai/TripoSR/model.ckpt", + // TripoSR is published as this checkpoint only; it is + // Stability AI's, pinned, and the one pickle llmgw ships. + AllowPickle: true, ComfyUI: &ComfyUI{ Workflow: "triposr", Output: OutputMeshes, diff --git a/internal/control/catalog/shared.go b/internal/control/catalog/shared.go index 3cb8823..d9e13b3 100644 --- a/internal/control/catalog/shared.go +++ b/internal/control/catalog/shared.go @@ -99,6 +99,7 @@ func envModel(env Environment, members []Model) Model { Role: RoleEnvironment, Description: "the shared ComfyUI replica of environment " + env.Name, Source: first.Source, + AllowPickle: first.AllowPickle, ComfyUI: &ComfyUI{Environment: env.Name, Nodes: env.Nodes, Output: OutputEnvironment}, Profiles: profs, Scaling: Scaling{Min: minN, Max: maxN, IdleTimeout: idle}, diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 069fd44..a4e8ca8 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -516,8 +516,10 @@ type inferenceBody struct { MaxTokens int `json:"max_tokens"` MaxCompletionTokens int `json:"max_completion_tokens"` // Read only to say them in the request's trace; sent on untouched. - Temperature *float64 `json:"temperature"` - TopP *float64 `json:"top_p"` + // Temperature and TopP are read loosely, for the trace only: a client + // that sends "0.7" is the engine's to accept or refuse, as before. + Temperature looseFloat `json:"temperature"` + TopP looseFloat `json:"top_p"` StreamOptions *struct { IncludeUsage bool `json:"include_usage"` } `json:"stream_options"` @@ -907,8 +909,8 @@ func (r *Router) parse(req *http.Request, path string) (*parsedRequest, error) { Band: req.Header.Get(r.cfg.Admission.BandHeader), ProfileOverride: req.Header.Get("x-llmgw-profile"), MaxTokens: max(b.MaxTokens, b.MaxCompletionTokens), - Temperature: b.Temperature, - TopP: b.TopP, + Temperature: b.Temperature.v, + TopP: b.TopP.v, PromptPrefix: promptPrefix(&b), ThreadID: req.Header.Get(threadHeader), } @@ -1536,3 +1538,15 @@ func withAdapter(body []byte, model string, index int, scale float64) []byte { } return out } + +// looseFloat is a number, or a string holding one, or anything else read as +// none: a field a client may write in more ways than one, read only to say. +type looseFloat struct{ v *float64 } + +func (f *looseFloat) UnmarshalJSON(b []byte) error { + s := strings.Trim(string(b), `"`) + if v, err := strconv.ParseFloat(s, 64); err == nil { + f.v = &v + } + return nil +} diff --git a/internal/router/trace_test.go b/internal/router/trace_test.go index 83d7638..79fb247 100644 --- a/internal/router/trace_test.go +++ b/internal/router/trace_test.go @@ -141,3 +141,14 @@ func TestEveryInferencePathHasItsGenAIOperation(t *testing.T) { } } } + +func TestATemperatureWrittenAsAStringIsTheEnginesToJudge(t *testing.T) { + // Read for the trace only: before it was, "0.7" reached the engine. + e := newFakeEngine(t, "m") + rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})) + rec := post(t, rt.Handler(), map[string]any{"model": "m", "temperature": "0.7", + "messages": []any{map[string]any{"role": "user", "content": "hi"}}}, nil) + if rec.Code != 200 { + t.Errorf("status %d: %s", rec.Code, rec.Body.String()) + } +}