llmgw/internal/agent/engines/unsafeenv.go

package engines

import (
	"fmt"
	"os"
	"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_",
	"TRTLLM_", "TLLM_", "FREETOKEN_",
}

// envExact are single variables that are tuning too.
var envExact = map[string]bool{
	"PYTHONUNBUFFERED": true, "PYTHONHASHSEED": true, "TZ": true,
}

// envNeverExact are variables in an allowed family that no profile may set,
// whose names the patterns below do not catch: what the agent sets itself,
// and FreeToken's way of loading a projector file.
var envNeverExact = map[string]string{
	"FREETOKEN_API_KEY": "is set by the agent", "FREETOKEN_HOST": "is set by the agent",
	"FREETOKEN_MMPROJ": "loads a file",
}

// 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.
func UnsafeEnv(key string) (bool, string) {
	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, "*, ") + "*)"
	}
	if why, ok := envNeverExact[k]; ok {
		return true, "it " + why
	}
	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_<FLAG> 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",
		// The short spellings of the same.
		"-mm": "loads a file; the projector is fetched with the model", "--mmproj-url": "downloads a file",
		"-mmu": "downloads a file", "-hfr": "downloads a model", "-hff": "downloads a model",
		"--hf-repo-draft": "downloads a model", "-hfd": "downloads a model", "-hfrd": "downloads a model",
		"--hf-repo-v": "downloads a model", "-hfv": "downloads a model", "--hf-file-v": "downloads a model",
		"-hffv": "downloads a model", "-hft": "is a credential", "-jf": "reads a file", "-a": "is set by the agent",
		"--model-vocoder": "loads a file", "-mv": "loads a file",
	},
	"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",
		"--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.
	"comfyui": {
		"--listen": "is set by the agent", "--base-directory": "is set by the agent",
		"--extra-model-paths-config": "is set by the agent",
		"--output-directory":         "writes files", "--temp-directory": "writes files",
		"--input-directory": "reads files", "--user-directory": "writes files",
		"--database-url": "writes a database elsewhere", "--front-end-root": "serves a directory",
		"--enable-cors-header": "lets other sites' pages call it",
		"--enable-manager":     "installs nodes, which run Python",
		"--tls-keyfile":        "reads a file", "--tls-certfile": "reads a file",
		"--comfy-api-base": "sends work to another service",
	},
	// Strata's setup.py, which a profile's arguments go to: llmgw
	// strata-serve sets where the model is, where its files go and where
	// the server listens, and the engine is the one the agent installed.
	// Its own --model, the size, is spelt --size in a profile, since
	// --model is the source everywhere else.
	"strata": {
		"--gguf-dir": "is the model's source, set by the agent", "--models-dir": "is set by the agent",
		"--data-dir": "is set by the agent", "--move-data": "moves files", "--set-default-data-dir": "writes settings",
		"--feed": "downloads an engine", "--prebuilt": "loads an engine", "--build": "compiles an engine",
		"--update": "downloads an engine", "--rollback-engine": "changes the engine",
		"--dsv41-package": "loads a file", "--inspect": "reads another repository or path",
		"--source": "downloads from elsewhere", "--setup": "is not for a server",
		"--browser": "opens a browser", "--gpu": "is set by the agent", "--gpus": "is set by the agent",
		"--gpu-roles": "is set by the agent", "--backend": "is set by the agent",
		"--experimental-speed-projection": "changes what the model refuses",
	},
	// ft serve: where it listens, its key, its model and its card are the
	// agent's, and what reads files or fetches is refused.
	"freetoken": {
		"--gpu": "is set by the agent", "--served-model-name": "is set by the agent",
		"--mmproj": "loads a file", "--mm-encoder-weights": "loads a file",
		"--allowed-local-media-path": "serves local files to requests",
		"--allowed-media-domains":    "lets requests have the engine fetch from those sites",
		"--moe-cache-split":          "reads a profile file",
		"--hf-overrides":             "rewrites the checkpoint's configuration, which can name code to load",
	},
	// trtllm-serve spells its flags with underscores, so the shared list's
	// hyphenated names do not reach them.
	"trtllm": {
		"--trust_remote_code":     "runs Python the model repository ships",
		"--served_model_name":     "is set by the agent",
		"--config":                "reads a file; the agent writes the options file itself",
		"--extra_llm_api_options": "reads a file; the agent writes the options file itself",
		"--tokenizer":             "reads another repository or path",
		"--custom_tokenizer":      "imports Python", "--post_processor_hook": "imports Python",
		"--custom_module_dirs": "imports Python", "--middleware": "imports Python",
		"--chat_template":               "reads a file",
		"--metadata_server_config_file": "reads a file",
		"--visual_gen_args":             "reads a file",
		"--report_addr":                 "writes a file",
		"--otlp_traces_endpoint":        "sends request traces elsewhere",
		"--disagg_cluster_uri":          "joins the server to another cluster",
		"--server_role":                 "makes the server part of another cluster",
		"--allow_request_chat_template": "lets a request supply the chat template",
		"--grpc":                        "serves without the key the agent sets",
		"--num_serve_frontends":         "starts front ends that do not demand the key the agent sets",
		"--telemetry":                   "sends usage to NVIDIA",
		"--tensor_parallel_size":        "is set by the agent from tp",
		"--tp_size":                     "is set by the agent from tp",
		"--pipeline_parallel_size":      "is set by the agent from pp",
		"--pp_size":                     "is set by the agent from pp",
	},
}

// 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. Short flags (-m,
	// -hf) are matched as written.
	norm := func(f string) string {
		if strings.HasPrefix(f, "--") {
			return "--" + strings.ReplaceAll(f[2:], "_", "-")
		}
		return f
	}
	given := norm(name)
	// argparse also takes any unambiguous beginning of a long flag for the
	// flag — --trust-remote is --trust-remote-code, --li is ComfyUI's
	// --listen — so for the engines that parse with it, a beginning of a
	// refused flag is refused. llama.cpp and TensorRT-LLM take no such
	// thing, and --chat-template there is a flag of its own.
	abbreviates := argparseEngines[engine] && strings.HasPrefix(given, "--") && len(given) >= 4
	for _, set := range []string{"", engine} {
		for never, why := range flagsNever[set] {
			if why == "" {
				continue
			}
			n := norm(never)
			if n == given || (abbreviates && strings.HasPrefix(n, "--") && strings.HasPrefix(n, given)) {
				return why
			}
		}
	}
	return ""
}

// argparseEngines read their flags with Python's argparse, which takes a
// long flag's beginning for the flag.
var argparseEngines = map[string]bool{"vllm": true, "sglang": true, "comfyui": true, "strata": true, "freetoken": true}

// 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
}

// safetensorsOnly has vLLM or SGLang read only safetensors from a model it
// downloads itself, by its Hugging Face id: left to choose, either takes
// the repository's pickle files (pytorch_model.bin) when it has no
// safetensors, and loading those runs whatever they hold. A path was
// fetched by the agent, which checked it, or put there by the operator.
func safetensorsOnly(p Profile) []string {
	if p.AllowPickle || !hubID(p.ModelPath) {
		return nil
	}
	return []string{"--load-format", "safetensors"}
}

// hubID reports whether a model path is a Hugging Face id, org/name, rather
// than a file or directory.
func hubID(s string) bool {
	if s == "" || strings.HasPrefix(s, "/") || strings.HasPrefix(s, ".") || strings.Contains(s, "://") {
		return false
	}
	// org/name:Q4_K_M is a GGUF in a repository, which vLLM reads as one.
	if strings.Count(s, "/") != 1 || strings.Contains(s, ":") {
		return false
	}
	_, err := os.Stat(s)
	return err != nil
}