diff --git a/CHANGELOG.md b/CHANGELOG.md index 1cd0d85..d28ca7c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -46,6 +46,18 @@ APIs change between minors without a deprecation period. `agent_insecure: true` is set, and then listens on loopback unless `agent_listen` names an address. `server init` writes `agent_tls`, as before. +- **A model definition cannot run code on its nodes.** A profile's `env` and + `args` reach a process the agent runs as root. The environment was checked + against a short denylist that missed NCCL's and CUDA's plugin and + injection libraries; arguments were not checked at all, so + `--trust-remote-code`, llama.cpp's `--log-file` or `--slot-save-path`, or + vLLM's `--tool-parser-plugin` were taken. `env` is now an allowlist of the + families engines are tuned with (CUDA_, NCCL_, VLLM_, OMP_, HF_…) less any + name for a path, plugin, library, address or credential, and llama.cpp's + `LLAMA_ARG_*` are checked as the flags they are; `args` refuse what reads, + writes or loads files, loads code, or moves what the agent sets (host, + port, key, model). Both are refused where a model is defined and again on + the node before the engine starts. ### Everything that can be changed on a running cluster, from the console diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 2de863d..e26bde1 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -920,6 +920,12 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan< // Validate before launching. An engine that refuses a profile does so after // minutes of loading; catching it here costs nothing. + // Checked again here, whatever sent the profile: this is the process + // about to run as root. + if err := engines.UnsafeArgs(spec.Profile.Engine, spec.Profile.Args); err != nil { + report(supervisor.Failed, err.Error(), nil) + return + } if err := eng.Validate(spec.Profile, gpuInfo(inv, spec.GPUIndices)); err != nil { report(supervisor.Failed, err.Error(), nil) return diff --git a/internal/agent/engines/unsafeenv.go b/internal/agent/engines/unsafeenv.go index f60bc53..43ab6af 100644 --- a/internal/agent/engines/unsafeenv.go +++ b/internal/agent/engines/unsafeenv.go @@ -1,30 +1,147 @@ package engines -import "strings" +import ( + "fmt" + "strings" +) + +// A model profile's env and args are the operator's escape hatch for tuning +// an engine, and they reach a process the agent starts as root on every node +// the model is placed on. Some of what an engine accepts is not tuning but a +// way to run other code, or to read or write files, in that process: a +// loader preload, an NCCL network plugin, a CUDA injection library, vLLM's +// --trust-remote-code, llama.cpp's --log-file pointed at /etc/cron.d. Taken +// from a model definition, which the admin API or a config file supplies, +// they would turn "add a model" into "run this as root on every node". +// +// So env is an allowlist — the families of variables engines are tuned with +// — less any whose name says it names a path, a plugin or a library, and args +// are checked against what each engine takes that reads, writes or loads. +// Both are refused where a model is defined and again where an engine starts. + +// envFamilies are the prefixes of variables engines are tuned with. +var envFamilies = []string{ + "CUDA_", "NCCL_", "NVIDIA_", "NVTE_", "GLOO_", "UCX_", "TORCH_", "PYTORCH_", + "VLLM_", "SGLANG_", "GGML_", "LLAMA_", "OMP_", "MKL_", "OPENBLAS_", "KMP_", + "HF_", "TRANSFORMERS_", "TOKENIZERS_", "TRITON_", "FLASHINFER_", +} + +// envExact are single variables that are tuning too. +var envExact = map[string]bool{ + "PYTHONUNBUFFERED": true, "PYTHONHASHSEED": true, "TZ": true, +} + +// envNever are names in an allowed family that load code or move where the +// engine reads and writes. +var envNever = []string{ + "PLUGIN", "INJECTION", "PRELOAD", "LIBRARY", "_LIB", + "_PATH", "_DIR", "_HOME", "_FILE", "_ROOT", "_CACHE", "_CONFIG", + "_ADDRESS", "_ENDPOINT", "_TOKEN", "_URL", +} // UnsafeEnv reports whether a profile may not set an environment variable, // and why. -// -// A profile's env is the escape hatch for tuning no adapter can guess — -// CUDA_*, NCCL_*, VLLM_*, HF_* — and it reaches an engine the agent starts as -// root. Some variables are not tuning but a way to run other code in that -// process: the dynamic loader's LD_PRELOAD, Python's PYTHONSTARTUP, a shell's -// BASH_ENV. Set from a model definition, which an admin API or a config file -// supplies, they would turn "add a model" into "run this as root on every -// node". They are refused where a model is defined and again where an engine -// starts. func UnsafeEnv(key string) (bool, string) { - k := strings.ToUpper(key) - switch { - case strings.HasPrefix(k, "LD_"), strings.HasPrefix(k, "DYLD_"): - return true, "it changes what the dynamic loader loads" - case k == "PATH": - return true, "it changes which binaries run" - case strings.HasPrefix(k, "PYTHON") && k != "PYTHONUNBUFFERED" && k != "PYTHONHASHSEED": - return true, "it changes what Python imports or runs at start" - case k == "BASH_ENV" || k == "ENV" || k == "PERL5OPT" || k == "PERL5LIB" || - k == "NODE_OPTIONS" || k == "RUBYOPT" || k == "GCONV_PATH" || k == "OPENSSL_CONF": - return true, "it makes an interpreter or library load code" + k := strings.ToUpper(strings.TrimSpace(key)) + if k == "" { + return true, "it has no name" + } + allowed := envExact[k] + for _, f := range envFamilies { + if strings.HasPrefix(k, f) { + allowed = true + break + } + } + if !allowed { + return true, "only the variables engines are tuned with are allowed (" + + strings.Join(envFamilies, "*, ") + "*)" + } + for _, n := range envNever { + if strings.Contains(k, n) && !envExact[k] { + return true, "it names a path, a plugin, a library or an address, which is a way to make the engine " + + "read, write or load something other than its model" + } + } + // llama.cpp reads every flag from LLAMA_ARG_ too. + if rest, ok := strings.CutPrefix(k, "LLAMA_ARG_"); ok { + flag := "--" + strings.ToLower(strings.ReplaceAll(rest, "_", "-")) + if why := unsafeFlag("llamacpp", flag, ""); why != "" { + return true, "it is " + flag + ", which " + why + } } return false, "" } + +// flagsNever are flags no profile may give, by engine; "" is every engine. +// Each reads or writes files the agent did not choose, loads code, or moves +// what the agent itself sets: where the engine listens and the key it +// demands. +var flagsNever = map[string]map[string]string{ + "": { + "--host": "is set by the agent", "--port": "is set by the agent", + "--api-key": "is set by the agent", "--api-key-file": "is set by the agent", + "--model": "is the model's source, set by the agent", "-m": "is the model's source, set by the agent", + "--ssl-keyfile": "reads a file", "--ssl-certfile": "reads a file", + "--ssl-key-file": "reads a file", "--ssl-cert-file": "reads a file", + "--trust-remote-code": "runs Python the model repository ships", + }, + "llamacpp": { + "--log-file": "writes a file", "--slot-save-path": "writes files", + "--path": "serves a directory's files over HTTP", "--rpc": "sends the model's work to other machines", + "--chat-template-file": "reads a file", "--system-prompt-file": "reads a file", + "--grammar-file": "reads a file", "--json-schema-file": "reads a file", + "--lora": "loads a file; adapters are set in the model's adapters", + "--lora-scaled": "loads a file; adapters are set in the model's adapters", + "--control-vector": "loads a file", "--control-vector-scaled": "loads a file", + "--mmproj": "loads a file; the projector is fetched with the model", + "--model-url": "downloads a model", "-mu": "downloads a model", + "--hf-repo": "downloads a model", "-hf": "downloads a model", "--hf-file": "downloads a model", + "--hf-token": "is a credential", "--alias": "is set by the agent", + "--model-draft": "loads a file", "-md": "loads a file", + "--prompt-cache": "writes a file", "--logdir": "writes files", + }, + "vllm": { + "--download-dir": "writes files", "--tokenizer": "reads another repository or path", + "--chat-template": "reads a file", "--allowed-local-media-path": "serves local files to requests", + "--lora-modules": "loads files; adapters are set in the model's adapters", + "--tool-parser-plugin": "runs a Python file", "--middleware": "imports Python", + "--worker-extension-cls": "imports Python", "--worker-cls": "imports Python", + "--load-format": "changes how weights are loaded, some ways running code", + "--otlp-traces-endpoint": "sends request traces elsewhere", + "--served-model-name": "is set by the agent", + "--enable-log-requests": "logs prompts", + }, + "sglang": { + "--tokenizer-path": "reads another repository or path", "--chat-template": "reads a file", + "--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", + }, +} + +// unsafeFlag is why an engine may not be given a flag, or "". +func unsafeFlag(engine, flag, _ string) string { + name, _, _ := strings.Cut(flag, "=") + for _, set := range []string{"", engine} { + if why, ok := flagsNever[set][name]; ok && why != "" { + return why + } + } + return "" +} + +// UnsafeArgs refuses the first of a profile's arguments an engine may not be +// given, saying why. +func UnsafeArgs(engine string, args []string) error { + for _, a := range args { + if !strings.HasPrefix(a, "-") { + continue // a value + } + if why := unsafeFlag(engine, a, ""); why != "" { + name, _, _ := strings.Cut(a, "=") + return fmt.Errorf("argument %s is not allowed: it %s", name, why) + } + } + return nil +} diff --git a/internal/agent/engines/unsafeenv_test.go b/internal/agent/engines/unsafeenv_test.go new file mode 100644 index 0000000..cdf8e76 --- /dev/null +++ b/internal/agent/engines/unsafeenv_test.go @@ -0,0 +1,64 @@ +package engines + +import "testing" + +func TestTuningVariablesAreAllowedAndWaysToRunCodeAreNot(t *testing.T) { + for _, k := range []string{ + "CUDA_VISIBLE_DEVICES", "CUDA_MODULE_LOADING", "NCCL_IB_HCA", "NCCL_DEBUG", "NCCL_NET", + "VLLM_ATTENTION_BACKEND", "PYTORCH_CUDA_ALLOC_CONF", "OMP_NUM_THREADS", "HF_HUB_OFFLINE", + "GGML_VK_VISIBLE_DEVICES", "TOKENIZERS_PARALLELISM", "PYTHONUNBUFFERED", "LLAMA_ARG_CTX_SIZE", + } { + if bad, why := UnsafeEnv(k); bad { + t.Errorf("%s refused: %s", k, why) + } + } + for _, k := range []string{ + "LD_PRELOAD", "PATH", "PYTHONPATH", "PYTHONSTARTUP", "BASH_ENV", + // What the denylist missed: each loads a library or moves where the + // engine reads and writes. + "NCCL_NET_PLUGIN", "NCCL_TUNER_PLUGIN", "NCCL_PROFILER_PLUGIN", "CUDA_INJECTION64_PATH", + "NVTX_INJECTION64_PATH", "TRITON_CACHE_DIR", "HF_HOME", "HF_TOKEN", "NCCL_DEBUG_FILE", + "VLLM_PLUGINS", "UCX_MODULE_DIR", "TORCH_EXTENSIONS_DIR", + // llama.cpp's flags as variables. + "LLAMA_ARG_LOG_FILE", "LLAMA_ARG_HOST", "LLAMA_ARG_SLOT_SAVE_PATH", + // Anything outside the families. + "SOMETHING_ELSE", "", + } { + if bad, _ := UnsafeEnv(k); !bad { + t.Errorf("%s allowed", k) + } + } +} + +func TestArgumentsThatReadWriteOrLoadAreRefused(t *testing.T) { + ok := map[string][]string{ + "llamacpp": {"--ctx-size", "16384", "--embedding", "--parallel", "4", "--jinja", "--n-gpu-layers=999"}, + "vllm": {"--max-model-len", "8192", "--gpu-memory-utilization", "0.9", "--enable-auto-tool-choice"}, + "sglang": {"--mem-fraction-static", "0.8", "--tool-call-parser", "qwen25"}, + } + for eng, args := range ok { + if err := UnsafeArgs(eng, args); err != nil { + t.Errorf("%s %v: %v", eng, args, err) + } + } + for _, c := range []struct { + eng string + args []string + }{ + {"vllm", []string{"--trust-remote-code"}}, + {"sglang", []string{"--trust-remote-code"}}, + {"vllm", []string{"--tool-parser-plugin", "/tmp/x.py"}}, + {"vllm", []string{"--middleware=evil.mod"}}, + {"vllm", []string{"--allowed-local-media-path", "/"}}, + {"llamacpp", []string{"--log-file", "/etc/cron.d/x"}}, + {"llamacpp", []string{"--slot-save-path=/root"}}, + {"llamacpp", []string{"--path", "/"}}, + {"llamacpp", []string{"--system-prompt-file", "/etc/shadow"}}, + {"llamacpp", []string{"--host", "0.0.0.0"}}, + {"vllm", []string{"--api-key", "x"}}, + } { + if err := UnsafeArgs(c.eng, c.args); err == nil { + t.Errorf("%s %v allowed", c.eng, c.args) + } + } +} diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index 710e5c5..c74c06a 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -440,6 +440,9 @@ func (m Model) validate() error { m.Name, p.Name, k, why) } } + if err := engines.UnsafeArgs(p.Engine, p.Args); err != nil { + return fmt.Errorf("model %q profile %q: %w", m.Name, p.Name, err) + } if p.GPUClass == "" && !cpuCapable(p.Engine) { return fmt.Errorf("model %q profile %q: gpu_class is required for engine %q", diff --git a/internal/control/catalog/catalog_test.go b/internal/control/catalog/catalog_test.go index 255e117..4a62194 100644 --- a/internal/control/catalog/catalog_test.go +++ b/internal/control/catalog/catalog_test.go @@ -385,3 +385,21 @@ func TestARepositoryPageIsReadAsTheRepositoryForVLLM(t *testing.T) { t.Errorf("a file under /resolve/ is a source: %v", err) } } + +func TestAModelThatWouldRunCodeOnItsNodesIsRefusedWhereItIsDefined(t *testing.T) { + for _, p := range []Profile{ + {Name: "p", Engine: "vllm", GPUClass: "h100-80g", TP: 1, Args: []string{"--trust-remote-code"}}, + {Name: "p", Engine: "llamacpp", MinMemoryGB: 4, Args: []string{"--log-file", "/etc/cron.d/x"}}, + {Name: "p", Engine: "vllm", GPUClass: "h100-80g", TP: 1, Env: map[string]string{"NCCL_NET_PLUGIN": "/tmp/x.so"}}, + {Name: "p", Engine: "vllm", GPUClass: "h100-80g", TP: 1, Env: map[string]string{"CUDA_INJECTION64_PATH": "/tmp/x.so"}}, + } { + if _, err := New([]Model{model("m", p)}); err == nil { + t.Errorf("accepted %v %v", p.Args, p.Env) + } + } + ok := Profile{Name: "p", Engine: "llamacpp", MinMemoryGB: 4, Args: []string{"--ctx-size", "16384"}, + Env: map[string]string{"OMP_NUM_THREADS": "8"}} + if _, err := New([]Model{model("m", ok)}); err != nil { + t.Errorf("a tuned model was refused: %v", err) + } +}