From 7940e8593c428b39b44be7cab708286ae2ff43c0 Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Mon, 5 Oct 2026 00:02:01 +0200 Subject: [PATCH] TensorRT-LLM as an engine, and models from Ollama into the mirror Engine trtllm runs trtllm-serve on one node, through a launcher of llmgw's that demands the replica's key and turns on the Prometheus endpoint its queue and cache gauges are read from; the router reads it at /prometheus/metrics. Brought by the operator (engines.trtllm.python), not installed; checked by ctl engine-check trtllm. Written against the documented CLI, server and metrics and tested with stand-ins: not yet run on a datacenter GPU. mirror import-ollama copies a model's GGUF, and a vision model's projector, from Ollama's directory into the mirror, checked against the digest Ollama named each blob by; mirror import does the same for any local file. The console's Mirror view imports from Ollama with I. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/console.md | 5 +- docs/install.md | 121 +++++++- docs/reference.md | 26 ++ internal/agent/engines/engine.go | 13 +- internal/agent/engines/engines_test.go | 4 +- internal/agent/engines/trtllm.go | 218 ++++++++++++++ internal/agent/engines/trtllm_test.go | 316 ++++++++++++++++++++ internal/agent/engines/unsafeenv.go | 27 ++ internal/cli/bench.go | 2 +- internal/cli/enginecheck.go | 21 +- internal/cli/enginecheck_test.go | 77 +++++ internal/cli/mirror.go | 2 + internal/cli/mirror_import.go | 216 ++++++++++++++ internal/cli/mirror_import_test.go | 144 +++++++++ internal/cli/tui_console.go | 4 + internal/cli/tui_mirror.go | 36 ++- internal/control/catalog/catalog.go | 13 + internal/control/catalog/trtllm_test.go | 53 ++++ internal/docgen/reference_gen.go | 4 + internal/docgen/render.go | 1 + internal/mirror/import.go | 131 +++++++++ internal/mirror/ollama.go | 352 ++++++++++++++++++++++ internal/mirror/ollama_test.go | 369 ++++++++++++++++++++++++ internal/router/scrape/parse.go | 34 +++ internal/router/scrape/poller.go | 9 +- internal/router/scrape/trtllm_test.go | 133 +++++++++ internal/web/admin_serve.go | 2 +- notes/c23.md | 61 ++++ 28 files changed, 2376 insertions(+), 18 deletions(-) create mode 100644 internal/agent/engines/trtllm.go create mode 100644 internal/agent/engines/trtllm_test.go create mode 100644 internal/cli/mirror_import.go create mode 100644 internal/cli/mirror_import_test.go create mode 100644 internal/control/catalog/trtllm_test.go create mode 100644 internal/mirror/import.go create mode 100644 internal/mirror/ollama.go create mode 100644 internal/mirror/ollama_test.go create mode 100644 internal/router/scrape/trtllm_test.go create mode 100644 notes/c23.md diff --git a/docs/console.md b/docs/console.md index f891558..8ae8195 100644 --- a/docs/console.md +++ b/docs/console.md @@ -138,7 +138,10 @@ team or project is guaranteed and may at most hold, and what it holds now. **Mirror** is what the model mirror holds, one row for each version of each model — a file, or the set of files a model of several is: shards, a vision projector, a config and tokenizer — with how far its download has got and -which models use it. **A** downloads a URL into it; on a row, its files with +which models use it. **A** downloads a URL into it; **I** imports a model +from an Ollama on this machine into a mirror directory on this machine — the +mirror beside the control plane — and shows the model definition to add for +it (a mirror elsewhere: `llmgw mirror import-ollama` there); on a row, its files with their checksums and where each was fetched from, the source a model definition names for it, or remove it — with a warning when a model still uses it. diff --git a/docs/install.md b/docs/install.md index c5923ce..f609f39 100644 --- a/docs/install.md +++ b/docs/install.md @@ -383,6 +383,46 @@ llmgw mirror fetch --model qwen2.5-14b-instruct --version 2026-03 \ or from the admin page, or by copying files into the directory by hand — the mirror is a directory, and anything that puts files in it works. +Files already on the mirror machine go in with `import`, which copies them the +way a download is kept — under a temporary name until whole, the checksum +beside them for nodes to check against, a version never replaced — and checks +that a `.gguf` is one: + +``` +llmgw mirror import ./qwen2.5-7b-instruct-q4_k_m.gguf --model qwen2.5-7b-instruct \ + --version 2026-10 --sha256 +``` + +#### From Ollama + +Models you already run in Ollama are GGUFs llama.cpp serves as they are. +On the mirror machine, with Ollama's models readable there: + +``` +llmgw mirror import-ollama # lists what Ollama has +llmgw mirror import-ollama llama3.2:3b +``` + +It finds Ollama's models directory where Ollama keeps it — `OLLAMA_MODELS`, +`~/.ollama/models`, or `/usr/share/ollama/.ollama/models` for Ollama's Linux +service, which needs `sudo` to read — or takes `--ollama-dir`. It copies the +model's GGUF, and a vision model's projector beside it, each checked against +the sha256 Ollama named its file by; a file that does not match is not +mirrored, and nothing of the model is. The version is the tag and the start of +the digest (`llama3.2/3b-a80c4f17acd5/llama3.2-3b.gguf`), because a tag like +`latest` moves in Ollama and a mirror version must not; `--model` and +`--version` choose otherwise. It ends with a model definition to save and +`llmgw ctl model add -f`, with `--ctx-size` set to the `num_ctx` Ollama ran it +with. + +Ollama's prompt template is not carried: llama.cpp renders the chat template +inside the GGUF, which for any recent model is the one it was trained on. Its +system prompt and adapters are not carried either, and the import says so. +Models Ollama keeps as separate tensors rather than one GGUF, or as several +weight layers, are refused with a sentence saying so — mirror that model's GGUF +from its publisher instead. When the mirror runs beside the control plane, the +console's Mirror view does the same with **I**. + ## 4. Preparing a node **Do this before onboarding.** A node missing its driver or engine registers @@ -458,6 +498,7 @@ Nothing to do: the agent installs them once the node joins, under | llama.cpp | every node | the release build for its processor and GPU — CUDA 13.3, CUDA 12.8 (x64), or Vulkan for GPUs and drivers the CUDA builds do not cover; the CPU build only where there is no NVIDIA GPU — checked against GitHub's checksum and seen to list the GPU | | vLLM | nodes with an NVIDIA GPU it can run on | a Python environment built by uv from PyPI, about 8 GB | | SGLang | only when you ask for it | the same way as vLLM | +| TensorRT-LLM | only nodes where you installed it | not installed by llmgw: see [TensorRT-LLM, which you bring](#tensorrt-llm-which-you-bring) | vLLM is resolved as PyPI stood a day after that vLLM was published, so every node gets the same dependencies whenever it is built, and installed with @@ -492,8 +533,8 @@ engines: A changed version reaches every node when it next connects. To use an engine you installed yourself on a node, name its path in that node's `agent.yaml` -(`engines.llamacpp.bin`, `engines.vllm.bin`, `engines.sglang.python`): the agent -then installs none of its own beside it. +(`engines.llamacpp.bin`, `engines.vllm.bin`, `engines.sglang.python`, +`engines.trtllm.python`): the agent then installs none of its own beside it. #### Choosing a node's vLLM build @@ -544,6 +585,82 @@ which drift between releases: llmgw ctl engine-check vllm --bin /opt/llmgw/engines/vllm/0.30.0-cu13/bin/vllm ``` +#### TensorRT-LLM, which you bring + +llmgw does not install TensorRT-LLM. It is a pip package tied to one CUDA, one +PyTorch and an MPI the machine has to provide, and NVIDIA's own advice is its +container; an installer guessing at all of that would fail on most machines it +met, a long way into a large download. Install it on the node yourself — a +virtual environment, or NVIDIA's TensorRT-LLM container — and name its +interpreter in that node's `agent.yaml`: + +```yaml +engines: + trtllm: + python: /opt/trtllm/venv/bin/python # python3 inside NVIDIA's container +``` + +The node then offers the engine `trtllm` whenever that interpreter has the +`tensorrt_llm` package; without the line it offers nothing, whatever is on +`PATH`. It needs **TensorRT-LLM 1.3** (1.3.0rc13 or later): 1.2 has no +`--served_model_name`, and publishes none of the queue figures the router +reads. Check a node's against the flags llmgw builds: + +``` +llmgw ctl engine-check trtllm --bin /opt/trtllm/venv/bin/python +``` + +A model for it names a Hugging Face checkpoint, as for vLLM — trtllm-serve +downloads a repository itself, or loads one from the mirror: + +```yaml +name: llama-3.1-70b-instruct +source: meta-llama/Llama-3.1-70B-Instruct +profiles: + - name: h100x4 + engine: trtllm + gpu_class: h100-80g + tp: 4 + args: ["--max_batch_size", "64", "--free_gpu_memory_fraction", "0.85"] +``` + +The agent starts `trtllm-serve serve --host … --port … --tp_size … +--served_model_name …`, through a few lines of Python of its own that do what +trtllm-serve has no flag for: + +- **the key.** trtllm-serve checks none; every path but `/health` is made to + demand the one the gateway gives the replica, so the port on the node is not + a way around the gateway. +- **the metrics.** Its options file turns on `return_perf_metrics` and + `enable_iter_perf_stats`, without which `/prometheus/metrics` does not exist + and its queue figures do not move. The router reads + `trtllm_num_requests_running`, `trtllm_num_requests_waiting` and + `trtllm_kv_cache_utilization` there — not `/metrics`, which on trtllm-serve + is JSON. They are set as the engine iterates, so an engine gone idle shows + its last batch until the next request. +- **no usage reports to NVIDIA**: `TRTLLM_NO_USAGE_STATS` and `DO_NOT_TRACK` + are set, after the profile's own `env`. + +What it does not do: + +- **span nodes.** A TensorRT-LLM replica larger than one machine is launched + under MPI or Slurm from one of them, which the agents, each starting its own + node's part, do not do. A profile with `nodes` above 1 is refused; use vLLM + or SGLang for one. +- **run before Ampere.** Compute capability 8.0 or newer, which is what current + TensorRT-LLM builds kernels for. +- **take an options file from a profile.** `--config` and + `--extra_llm_api_options` are llmgw's; change a setting with `--set + PATH=VALUE` instead (`--set kv_cache_config.enable_block_reuse=false`). Also + refused, like the other engines' equivalents: `--trust_remote_code`, + `--tokenizer`, `--middleware`, `--chat_template`, `--grpc` and + `--num_serve_frontends` (both would serve without the key), and the flags the + agent sets. + +This engine was written against TensorRT-LLM's documented command line, +server and metrics, and tested with stand-ins for them; it has not yet run on a +datacenter GPU. Run `engine-check` on the first node before placing a model on it. + ### Weights Nothing to do here if you built the mirror in step 3 — the agent fetches what a diff --git a/docs/reference.md b/docs/reference.md index 67325c4..c9413d7 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -264,6 +264,7 @@ rather than repeated. | `models` | `[]catalog.Model` | [Server](#server) | | `endpoints` | `[]Endpoint` | [Server](#server) | | `bin` | `string` | [SimConfig](#simconfig) | +| `python` | `string` | [TRTLLMConfig](#trtllmconfig) | | `store` | `string` | [Threads](#threads) | | `retain_hours` | `int` | [Threads](#threads) | | `dial` | `time.Duration` | [Timeouts](#timeouts) | @@ -331,6 +332,7 @@ rather than repeated. | `vllm` | `VLLMConfig` | [engines.Config](#enginesconfig) | | `sglang` | `SGLangConfig` | [engines.Config](#enginesconfig) | | `llamacpp` | `LlamaCppConfig` | [engines.Config](#enginesconfig) | +| `trtllm` | `TRTLLMConfig` | [engines.Config](#enginesconfig) | | `sim` | `SimConfig` | [engines.Config](#enginesconfig) | | `default` | `[]string` | [guard.Config](#guardconfig) | | `policies` | `map[string]Policy` | [guard.Config](#guardconfig) | @@ -2175,6 +2177,21 @@ against sim-engines possible without a control plane. Bin is the llmgw binary itself; the engine is one of its subcommands. Empty means whatever started this agent, which is right almost always. +## TRTLLMConfig + +### `python` + +`string` + +Python is the interpreter of the environment TensorRT-LLM is installed +in: /bin/python for a virtual environment, python3 in NVIDIA's +container. Naming it is what makes this node offer the engine; empty +offers nothing. + +The interpreter rather than trtllm-serve, because llmgw starts the +server through a few lines of its own that make it demand the +gateway's key, which trtllm-serve has no option for. + ## Threads ### `store` @@ -2633,6 +2650,15 @@ Takes a [SGLangConfig](#sglangconfig) block. Takes a [LlamaCppConfig](#llamacppconfig) block. +### `trtllm` + +`TRTLLMConfig` + +TRTLLM is TensorRT-LLM, which the agent does not install: a node +offers it only when its interpreter is named here. + +Takes a [TRTLLMConfig](#trtllmconfig) block. + ### `sim` `SimConfig` diff --git a/internal/agent/engines/engine.go b/internal/agent/engines/engine.go index 484ffed..343d6ac 100644 --- a/internal/agent/engines/engine.go +++ b/internal/agent/engines/engine.go @@ -168,6 +168,13 @@ type CPUFallback interface { Devices(ctx context.Context) ([]string, error) } +// HelpCommand is an engine whose program is not asked for its help by +// running the command's path with --help, because the path is an interpreter +// starting it through a launcher. +type HelpCommand interface { + HelpCommand(c Command) (path string, args []string) +} + // GPUInfo is the subset of GPU facts validation needs. type GPUInfo struct { Index int @@ -187,6 +194,7 @@ func NewRegistry(cfg Config) *Registry { r.Add(&VLLM{cfg: cfg.VLLM}) r.Add(&SGLang{cfg: cfg.SGLang}) r.Add(&LlamaCpp{cfg: cfg.LlamaCpp}) + r.Add(&TRTLLM{cfg: cfg.TRTLLM}) r.Add(&Sim{cfg: cfg.Sim}) return r } @@ -232,7 +240,10 @@ type Config struct { VLLM VLLMConfig `koanf:"vllm"` SGLang SGLangConfig `koanf:"sglang"` LlamaCpp LlamaCppConfig `koanf:"llamacpp"` - Sim SimConfig `koanf:"sim"` + // TRTLLM is TensorRT-LLM, which the agent does not install: a node + // offers it only when its interpreter is named here. + TRTLLM TRTLLMConfig `koanf:"trtllm"` + Sim SimConfig `koanf:"sim"` } // --- shared helpers --- diff --git a/internal/agent/engines/engines_test.go b/internal/agent/engines/engines_test.go index 049f0a2..c795bcf 100644 --- a/internal/agent/engines/engines_test.go +++ b/internal/agent/engines/engines_test.go @@ -457,14 +457,14 @@ func TestValidateRequiresModelPath(t *testing.T) { func TestRegistry(t *testing.T) { r := NewRegistry(Config{}) - if got := r.Names(); !slices.Equal(got, []string{"llamacpp", "sglang", "sim", "vllm"}) { + if got := r.Names(); !slices.Equal(got, []string{"llamacpp", "sglang", "sim", "trtllm", "vllm"}) { t.Errorf("registry contents: %v", got) } if _, err := r.Get("vllm"); err != nil { t.Errorf("vllm missing: %v", err) } err := func() error { _, err := r.Get("nope"); return err }() - if err == nil || !strings.Contains(err.Error(), "llamacpp, sglang, sim, vllm") { + if err == nil || !strings.Contains(err.Error(), "llamacpp, sglang, sim, trtllm, vllm") { t.Errorf("an unknown engine should list what is available: %v", err) } } diff --git a/internal/agent/engines/trtllm.go b/internal/agent/engines/trtllm.go new file mode 100644 index 0000000..c3f50ad --- /dev/null +++ b/internal/agent/engines/trtllm.go @@ -0,0 +1,218 @@ +package engines + +import ( + "context" + "fmt" + "strconv" +) + +// TRTLLMConfig locates a TensorRT-LLM the operator installed. +// +// llmgw does not install TensorRT-LLM. It is a pip package pinned to one +// CUDA, one PyTorch and an MPI the machine must provide, and NVIDIA's own +// advice is its container; an installer that guessed at all of that would +// fail on most machines it met, a long way into a large download. So it is +// brought by the operator — a virtual environment, or NVIDIA's container — +// and named here. +type TRTLLMConfig struct { + // Python is the interpreter of the environment TensorRT-LLM is installed + // in: /bin/python for a virtual environment, python3 in NVIDIA's + // container. Naming it is what makes this node offer the engine; empty + // offers nothing. + // + // The interpreter rather than trtllm-serve, because llmgw starts the + // server through a few lines of its own that make it demand the + // gateway's key, which trtllm-serve has no option for. + Python string `koanf:"python"` +} + +// TRTLLM builds TensorRT-LLM launch commands: its OpenAI-compatible server, +// trtllm-serve, on one node. +// +// Written against TensorRT-LLM's documented command line and server, and not +// yet run on hardware by anyone working on llmgw: no datacenter GPU was at +// hand. `llmgw ctl engine-check trtllm` asks an installed one whether it +// takes the flags built here. +type TRTLLM struct{ cfg TRTLLMConfig } + +// Name identifies the engine in the catalog. +func (*TRTLLM) Name() string { return "trtllm" } + +// HealthPath is trtllm-serve's liveness endpoint: 200 once the executor is +// up, 503 while it is not. The supervisor's /v1/models step is what says the +// model is being served. +func (*TRTLLM) HealthPath() string { return "/health" } + +// MetricsPath is trtllm-serve's Prometheus endpoint. Its /metrics is a +// different thing — a JSON list of iteration statistics, drained by each +// read — and a scrape there would both fail to parse and steal the numbers +// from anyone else reading them. +func (*TRTLLM) MetricsPath() string { return "/prometheus/metrics" } + +// ServesHTTP reports whether this rank answers the API; a TensorRT-LLM +// replica is one node, and so one rank. +func (*TRTLLM) ServesHTTP(s ShardSpec) bool { return s.Leader() } + +// Validate checks the profile against the GPUs it would run on. +// +// One node: trtllm-serve spans machines only when started under MPI or +// Slurm from one of them (trtllm-llmapi-launch), and the agents each start +// their own node's part, which is not that. Ampere or newer: the oldest +// generation current TensorRT-LLM builds kernels for. +func (*TRTLLM) Validate(p Profile, gpus []GPUInfo) error { + if p.Nodes > 1 { + return fmt.Errorf("profile %q: trtllm runs a replica on one node; "+ + "use vllm or sglang for a %d-node replica", p.Name, p.Nodes) + } + if err := validateCommon(p, gpus); err != nil { + return err + } + return requireComputeCapability(p, gpus, 8.0, "TensorRT-LLM") +} + +// Commands returns the processes for one shard, which is one process. Tensor +// and pipeline parallel ranks on the node are trtllm-serve's own MPI workers, +// started by it. +func (t *TRTLLM) Commands(s ShardSpec) ([]Command, error) { + if s.NNodes > 1 { + return nil, fmt.Errorf("trtllm: a replica spanning %d nodes is not something llmgw starts; "+ + "TensorRT-LLM runs one node per replica here", s.NNodes) + } + p := s.Profile + args := []string{"-c", trtllmLauncher, + // The subcommand by name: trtllm-serve takes a bare model as "serve", + // but saying it costs nothing and does not depend on that default. + "serve", p.ModelPath, + "--host", s.Host(), + "--port", strconv.Itoa(s.HTTPPort), + "--tp_size", strconv.Itoa(p.TP), + } + if p.PP > 1 { + args = append(args, "--pp_size", strconv.Itoa(p.PP)) + } + if p.ServedModelName != "" { + // Without it the model is named after its directory, and the agent's + // readiness check, which looks for this name in /v1/models, never + // finds it. + args = append(args, "--served_model_name", p.ServedModelName) + } + // What makes the queue and the cache measurable: the options file the + // launcher writes turns on the Prometheus endpoint and the iteration + // statistics its gauges are set from. A placeholder here, so the flag is + // one engine-check sees and checks; the launcher replaces it. + args = append(args, "--extra_llm_api_options", trtllmOptionsPlaceholder) + args = append(args, p.Args...) + + env := baseEnv(s) + if s.APIKey != "" { + env[trtllmKeyEnv] = s.APIKey + } + // TensorRT-LLM reports usage to NVIDIA unless told not to. A gateway on + // somebody's own network does not phone home on their behalf, and a + // model's profile cannot make it: these are set after the profile's env. + env["TRTLLM_NO_USAGE_STATS"] = "1" + env["DO_NOT_TRACK"] = "1" + return []Command{{Name: "trtllm-serve", Path: t.python(), Args: args, Env: materializeEnv(env)}}, nil +} + +// python is the interpreter to start trtllm-serve with. +func (t *TRTLLM) python() string { + if t.cfg.Python != "" { + return t.cfg.Python + } + return "python3" +} + +// Configured reports whether the configuration names TensorRT-LLM's +// interpreter, which is the only way a node offers it. +func (t *TRTLLM) Configured() bool { return t.cfg.Python != "" } + +// Probe reports the installed TensorRT-LLM's version. +// +// From the package's metadata rather than by importing it: the import loads +// CUDA, MPI and PyTorch and takes longer than a probe may, while what +// placement needs is whether the package is there. Whether it runs is +// engine-check's question, and the replica's. +func (t *TRTLLM) Probe(ctx context.Context) (string, error) { + return versionOfModule(ctx, t.python(), "-c", + "import importlib.metadata as m, importlib.util as u, sys\n"+ + "if u.find_spec('tensorrt_llm') is None: sys.exit(1)\n"+ + "sys.stdout.write(m.version('tensorrt_llm'))") +} + +// HelpCommand is how to ask trtllm-serve for its help: through its module, +// since the command itself is llmgw's launcher. +func (t *TRTLLM) HelpCommand(Command) (string, []string) { + return t.python(), []string{"-m", "tensorrt_llm.commands.serve", "serve", "--help"} +} + +const ( + trtllmKeyEnv = "LLMGW_ENGINE_KEY" + trtllmOptionsPlaceholder = "@llmgw-options" +) + +// trtllmLauncher starts trtllm-serve inside the process, after two things +// trtllm-serve has no flag for. +// +// The key. trtllm-serve checks none, and an engine listens on the node's +// management address, where without one anybody who can reach the port talks +// to the model past the gateway's authentication, quotas and accounting. So +// every FastAPI application the process makes — trtllm-serve's server is one +// — is given a middleware demanding it, before trtllm-serve imports anything. +// /health is left open: it says only up or down, and a supervisor of +// TensorRT-LLM's own may ask it without a key. The key comes in the +// environment and is taken out of it, so the MPI workers and anything else +// started from here never hold it. +// +// The options. The Prometheus endpoint exists only with return_perf_metrics, +// and its queue gauges move only with enable_iter_perf_stats; both are +// settings of TensorRT-LLM's options file, not flags. The file is written to +// an unnamed temporary file and passed by descriptor: nothing is left on disk +// however the server ends. +// +// trtllm-serve can start further HTTP front ends as copies of its own command +// line (--num_serve_frontends); those would not run this, so the flag is +// refused in a profile, and a copy started anyway would be given a program +// path that cannot exist — /dev/null is no directory — and fail rather than +// serve without the key. Click names itself by the path's last element, so +// trtllm-serve's messages still say trtllm-serve. +const trtllmLauncher = `import hmac, os, sys, tempfile +key = os.environ.pop("` + trtllmKeyEnv + `", "") +if key: + import fastapi + want = ("Bearer " + key).encode() + denied = b'{"error":{"message":"a valid API key is required","type":"invalid_request_error","code":"invalid_api_key"}}' + class RequireKey: + def __init__(self, app): + self.app = app + async def __call__(self, scope, receive, send): + if scope["type"] in ("http", "websocket") and scope.get("path") != "/health": + got = b"" + for name, value in scope.get("headers") or (): + if name == b"authorization": + got = value + if not hmac.compare_digest(got, want): + if scope["type"] == "websocket": + await send({"type": "websocket.close", "code": 1008}) + return + await send({"type": "http.response.start", "status": 401, + "headers": [(b"content-type", b"application/json")]}) + await send({"type": "http.response.body", "body": denied}) + return + await self.app(scope, receive, send) + plain_init = fastapi.FastAPI.__init__ + def init_with_key(self, *a, **kw): + plain_init(self, *a, **kw) + self.add_middleware(RequireKey) + fastapi.FastAPI.__init__ = init_with_key +args = sys.argv[1:] +if "` + trtllmOptionsPlaceholder + `" in args: + options = tempfile.TemporaryFile() + options.write(b"return_perf_metrics: true\nenable_iter_perf_stats: true\n") + options.flush() + options.seek(0) + args[args.index("` + trtllmOptionsPlaceholder + `")] = "/dev/fd/%d" % options.fileno() +sys.argv = ["/dev/null/trtllm-serve"] + args +from tensorrt_llm.commands.serve import main +main() +` diff --git a/internal/agent/engines/trtllm_test.go b/internal/agent/engines/trtllm_test.go new file mode 100644 index 0000000..e59bea1 --- /dev/null +++ b/internal/agent/engines/trtllm_test.go @@ -0,0 +1,316 @@ +package engines + +import ( + "encoding/json" + "os" + "os/exec" + "path/filepath" + "slices" + "strings" + "testing" +) + +func trtllmSpec() ShardSpec { + s := singleNodeSpec() + s.Profile.Engine = "trtllm" + s.Profile.ServedModelName = "llama-3.1-70b" + return s +} + +func TestTRTLLMRunsTrtllmServeWithTheFlagsItDocuments(t *testing.T) { + s := trtllmSpec() + s.BindAddr = "10.0.0.147" + s.Profile.Args = []string{"--max_batch_size", "64"} + cmd := buildOne(t, &TRTLLM{cfg: TRTLLMConfig{Python: "/opt/trtllm/bin/python"}}, s) + + if cmd.Path != "/opt/trtllm/bin/python" { + t.Errorf("Path=%q, expected the configured interpreter", cmd.Path) + } + if cmd.Args[0] != "-c" || !strings.Contains(cmd.Args[1], "from tensorrt_llm.commands.serve import main") { + t.Fatalf("not started through the launcher: %v", cmd.Args[:2]) + } + rest := cmd.Args[2:] + if rest[0] != "serve" || rest[1] != "/models/llama" { + t.Errorf("expected serve and the model first, got %v", rest[:2]) + } + for flag, want := range map[string]string{ + "--host": "10.0.0.147", "--port": "20000", "--tp_size": "4", + "--served_model_name": "llama-3.1-70b", + "--extra_llm_api_options": trtllmOptionsPlaceholder, + } { + if got := argValue(rest, flag); got != want { + t.Errorf("%s=%q, want %q (args %v)", flag, got, want, rest) + } + } + // One pipeline stage is the default, and saying so would only be one + // more flag for a version to rename. + if slices.Contains(rest, "--pp_size") { + t.Errorf("--pp_size given for a single stage: %v", rest) + } + // The profile's own, last, so a tuning flag is the operator's to set. + if got := rest[len(rest)-2:]; !slices.Equal(got, []string{"--max_batch_size", "64"}) { + t.Errorf("profile args are not last: %v", rest) + } + + s.Profile.TP, s.Profile.PP = 2, 2 + if got := argValue(buildOne(t, &TRTLLM{}, s).Args, "--pp_size"); got != "2" { + t.Errorf("--pp_size=%q for two stages", got) + } + if got := buildOne(t, &TRTLLM{}, s).Path; got != "python3" { + t.Errorf("unconfigured, the interpreter is %q", got) + } +} + +func TestTRTLLMsKeyIsNotOnItsCommandLine(t *testing.T) { + s := trtllmSpec() + s.APIKey = "k1" + cmd := buildOne(t, &TRTLLM{}, s) + for _, a := range cmd.Args { + if strings.Contains(a, "k1") && a != cmd.Args[1] { + t.Errorf("the key is in the arguments: %v", cmd.Args) + } + } + if envMap(cmd.Env)[trtllmKeyEnv] != "k1" { + t.Error("the key is not in the environment") + } +} + +func TestTRTLLMDoesNotReportUsageToNVIDIA(t *testing.T) { + s := trtllmSpec() + s.Profile.Env = map[string]string{"TRTLLM_NO_USAGE_STATS": "0", "TRTLLM_ENABLE_PDL": "1"} + env := envMap(buildOne(t, &TRTLLM{}, s).Env) + if env["TRTLLM_NO_USAGE_STATS"] != "1" || env["DO_NOT_TRACK"] != "1" { + t.Errorf("usage reporting is not off: %q %q", env["TRTLLM_NO_USAGE_STATS"], env["DO_NOT_TRACK"]) + } + // Tuning still reaches it. + if env["TRTLLM_ENABLE_PDL"] != "1" { + t.Error("a profile's TensorRT-LLM tuning was dropped") + } +} + +func TestTRTLLMIsReadFromItsPrometheusEndpoint(t *testing.T) { + // /metrics on trtllm-serve is JSON, and each read drains what it returns. + if got := (&TRTLLM{}).MetricsPath(); got != "/prometheus/metrics" { + t.Errorf("MetricsPath=%q", got) + } + if !(&TRTLLM{}).ServesHTTP(trtllmSpec()) { + t.Error("the one rank of a replica serves its API") + } +} + +func TestTRTLLMRunsAReplicaOnOneNode(t *testing.T) { + p := trtllmSpec().Profile + p.Nodes, p.GPUsPerNode, p.TP = 2, 4, 8 + err := (&TRTLLM{}).Validate(p, nil) + if err == nil || !strings.Contains(err.Error(), "one node") { + t.Errorf("a two-node profile was accepted: %v", err) + } + s := multiNodeSpec(0) + s.Profile.Engine = "trtllm" + if _, err := (&TRTLLM{}).Commands(s); err == nil { + t.Error("a command was built for a shard of a two-node replica") + } +} + +func TestTRTLLMNeedsAmpereOrNewer(t *testing.T) { + p := trtllmSpec().Profile + p.TP = 2 + v100 := []GPUInfo{{Index: 0, Name: "Tesla V100", ComputeCapability: "7.0"}, {Index: 1, Name: "Tesla V100", ComputeCapability: "7.0"}} + if err := (&TRTLLM{}).Validate(p, v100); err == nil { + t.Error("a V100 was accepted") + } + h100 := []GPUInfo{{Index: 0, Name: "H100", ComputeCapability: "9.0"}, {Index: 1, Name: "H100", ComputeCapability: "9.0"}} + if err := (&TRTLLM{}).Validate(p, h100); err != nil { + t.Errorf("two H100s were refused: %v", err) + } +} + +func TestTRTLLMProfilesMayNotOpenWhatTheAgentCloses(t *testing.T) { + for _, args := range [][]string{ + {"--trust_remote_code"}, + {"--num_serve_frontends", "4"}, + {"--grpc"}, + {"--config", "/etc/shadow"}, + {"--extra_llm_api_options=/tmp/x.yaml"}, + {"--served_model_name", "other"}, + {"--middleware", "evil.mw"}, + {"--tp_size", "8"}, + } { + if err := UnsafeArgs("trtllm", args); err == nil { + t.Errorf("%v was allowed", args) + } + } + if err := UnsafeArgs("trtllm", []string{"--max_batch_size", "64", "--set", "kv_cache_config.enable_block_reuse=false", + "--free_gpu_memory_fraction", "0.85"}); err != nil { + t.Errorf("tuning was refused: %v", err) + } +} + +func TestTRTLLMIsOfferedOnlyWhenItsPackageIsThere(t *testing.T) { + there := script(t, "python3", `echo "1.3.0rc29"`) + r := NewRegistry(Config{}) + r.engines = map[string]Engine{"trtllm": &TRTLLM{cfg: TRTLLMConfig{Python: there}}} + if got := r.ProbeConfigured(t.Context()); got["trtllm"] != "1.3.0rc29" { + t.Errorf("expected the version, got %v", got) + } + + // An interpreter without the package, which every machine has. + absent := script(t, "python3", "exit 1") + r.engines = map[string]Engine{"trtllm": &TRTLLM{cfg: TRTLLMConfig{Python: absent}}} + if got := r.ProbeConfigured(t.Context()); len(got) != 0 { + t.Errorf("reported without its package: %v", got) + } + + // Not named, not offered: python3 on PATH is nobody's TensorRT-LLM. + r.engines = map[string]Engine{"trtllm": &TRTLLM{}} + if got := r.ProbeConfigured(t.Context()); len(got) != 0 { + t.Errorf("offered without being configured: %v", got) + } +} + +// fakeTRTLLM is a Python path holding stand-ins for fastapi and +// tensorrt_llm: enough of each for the launcher to do what it does in front +// of the real server, and for the stand-in server to report what it saw. +func fakeTRTLLM(t *testing.T) string { + t.Helper() + dir := t.TempDir() + files := map[string]string{ + "fastapi/__init__.py": ` +class FastAPI: + def __init__(self, *a, **kw): + self.user_middleware = [] + def add_middleware(self, cls, **kw): + self.user_middleware.append(cls) +`, + "tensorrt_llm/__init__.py": "", + "tensorrt_llm/commands/__init__.py": "", + "tensorrt_llm/commands/serve.py": ` +import asyncio, json, os, sys +import fastapi + +def main(): + args = sys.argv[1:] + with open(args[args.index("--extra_llm_api_options") + 1]) as f: + options = f.read() + app = fastapi.FastAPI() + + async def server(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok"}) + + handler = server + for mw in reversed(app.user_middleware): + handler = mw(handler) + + async def ask(path, auth=None): + headers = [(b"authorization", auth.encode())] if auth else [] + sent = [] + async def send(msg): + sent.append(msg) + async def receive(): + return {"type": "http.request"} + await handler({"type": "http", "path": path, "headers": headers}, receive, send) + return sent[0]["status"] + + async def run(): + return { + "health": await ask("/health"), + "anonymous": await ask("/v1/chat/completions"), + "wrong": await ask("/v1/models", "Bearer nope"), + "keyed": await ask("/prometheus/metrics", "Bearer k1"), + } + + out = asyncio.run(run()) + out.update(options=options, argv0=sys.argv[0], args=args, + key_left=os.environ.get("LLMGW_ENGINE_KEY")) + print(json.dumps(out)) +`, + } + for name, body := range files { + p := filepath.Join(dir, name) + if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(p, []byte(body), 0o644); err != nil { + t.Fatal(err) + } + } + return dir +} + +type launched struct { + Health, Anonymous, Wrong, Keyed int + Options, Argv0 string + Args []string + KeyLeft *string `json:"key_left"` +} + +func launch(t *testing.T, s ShardSpec) launched { + t.Helper() + python, err := exec.LookPath("python3") + if err != nil { + t.Skip("no python3 here to run the launcher with") + } + cmd := buildOne(t, &TRTLLM{cfg: TRTLLMConfig{Python: python}}, s) + c := exec.Command(cmd.Path, cmd.Args...) //nolint:gosec // the command under test + c.Env = append(cmd.Env, "PYTHONPATH="+fakeTRTLLM(t)) + out, err := c.Output() + if err != nil { + var stderr string + if ee, ok := err.(*exec.ExitError); ok { + stderr = string(ee.Stderr) + } + t.Fatalf("the launcher failed: %v\n%s", err, stderr) + } + var got launched + if err := json.Unmarshal(out, &got); err != nil { + t.Fatalf("%v: %s", err, out) + } + return got +} + +func TestTheTRTLLMLauncherDemandsTheKey(t *testing.T) { + s := trtllmSpec() + s.APIKey = "k1" + got := launch(t, s) + + if got.Anonymous != 401 || got.Wrong != 401 { + t.Errorf("served without the key: anonymous %d, wrong key %d", got.Anonymous, got.Wrong) + } + if got.Keyed != 200 { + t.Errorf("refused the key: %d", got.Keyed) + } + // Up or down is all /health says, and it is asked without one. + if got.Health != 200 { + t.Errorf("/health answered %d", got.Health) + } + // Taken out of the environment, so nothing the server starts holds it. + if got.KeyLeft != nil { + t.Errorf("the key was left in the environment: %q", *got.KeyLeft) + } + // A copy of this command line, as trtllm-serve's extra front ends are, + // names no program that exists and so cannot serve without the key. + if _, err := os.Stat(got.Argv0); err == nil || filepath.Base(got.Argv0) != "trtllm-serve" { + t.Errorf("argv[0] is %q", got.Argv0) + } +} + +func TestTheTRTLLMLauncherTurnsOnTheMetricsTheRouterReads(t *testing.T) { + got := launch(t, trtllmSpec()) + for _, want := range []string{"return_perf_metrics: true", "enable_iter_perf_stats: true"} { + if !strings.Contains(got.Options, want) { + t.Errorf("the options file lacks %q: %q", want, got.Options) + } + } + path := argValue(got.Args, "--extra_llm_api_options") + if !strings.HasPrefix(path, "/dev/fd/") { + t.Errorf("the options were not passed by descriptor: %q", path) + } + if slices.Contains(got.Args, trtllmOptionsPlaceholder) { + t.Error("the placeholder reached trtllm-serve") + } + // Without a key there is nothing to demand. + if got.Anonymous != 200 { + t.Errorf("with no key set, a request was refused: %d", got.Anonymous) + } +} diff --git a/internal/agent/engines/unsafeenv.go b/internal/agent/engines/unsafeenv.go index 277b8a8..26c8cc8 100644 --- a/internal/agent/engines/unsafeenv.go +++ b/internal/agent/engines/unsafeenv.go @@ -24,6 +24,7 @@ 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_", } // envExact are single variables that are tuning too. @@ -118,6 +119,32 @@ var flagsNever = map[string]map[string]string{ "--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", }, + // 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 "". diff --git a/internal/cli/bench.go b/internal/cli/bench.go index 5c59691..b17f7f1 100644 --- a/internal/cli/bench.go +++ b/internal/cli/bench.go @@ -106,7 +106,7 @@ func newBenchCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { f := cmd.Flags() f.StringVar(&target, "target", "http://127.0.0.1:8080", "base URL: a replica, or the gateway") f.StringVar(&metrics, "metrics", "", "the replica's /metrics, to observe queue depth and KV usage") - f.StringVar(&engine, "engine", "", "which metric names --metrics uses: vllm | sglang | llamacpp") + f.StringVar(&engine, "engine", "", "which metric names --metrics uses: vllm | sglang | llamacpp | trtllm (whose are at /prometheus/metrics)") f.StringVar(&key, "key", "", "API key, when the gateway requires one") f.IntSliceVar(&steps, "concurrency", nil, "levels to try (default 1,2,4,…,128)") f.DurationVar(&cfg.Duration, "duration", cfg.Duration, "how long to hold each level") diff --git a/internal/cli/enginecheck.go b/internal/cli/enginecheck.go index 7726570..bed0d0b 100644 --- a/internal/cli/enginecheck.go +++ b/internal/cli/enginecheck.go @@ -57,6 +57,8 @@ func newEngineCheckCmd() *cobra.Command { cfg.VLLM.Bin = bin case "sglang": cfg.SGLang.Python = bin + case "trtllm": + cfg.TRTLLM.Python = bin case "sim": cfg.Sim.Bin = bin } @@ -97,7 +99,7 @@ func newEngineCheckCmd() *cobra.Command { continue default: var err error - if text, err = runHelp(cmd.Context(), c); err != nil { + if text, err = runHelp(cmd.Context(), eng, c); err != nil { skipped++ fmt.Fprintf(errOut, " skipped: %v\n", err) continue @@ -167,9 +169,10 @@ func newEngineCheckCmd() *cobra.Command { // representativeShard is the widest command llmgw would ever build, because a // flag that only appears in the multi-node case is exactly the one nobody -// notices has been renamed. +// notices has been renamed. For an engine that runs a replica on one node, +// that is one node. func representativeShard(engine string, nodes, tp, perNode int) engines.ShardSpec { - if nodes < 1 { + if nodes < 1 || engine == "trtllm" { nodes = 1 } return engines.ShardSpec{ @@ -203,9 +206,10 @@ var helpArgs = map[string][]string{ } // runHelp asks one program to describe itself. -func runHelp(ctx context.Context, c engines.Command) (string, error) { +func runHelp(ctx context.Context, eng engines.Engine, c engines.Command) (string, error) { // Everything up to the first flag, plus the help request. Passing the real // arguments would start the engine. + path := c.Path var args []string for _, a := range c.Args { if strings.HasPrefix(a, "-") { @@ -218,11 +222,16 @@ func runHelp(ctx context.Context, c engines.Command) (string, error) { ask = []string{"--help"} } args = append(args, ask...) + // An engine started through a launcher of llmgw's is asked through its + // own entry point: the interpreter's --help describes the interpreter. + if h, ok := eng.(engines.HelpCommand); ok { + path, args = h.HelpCommand(c) + } ctx, cancel := context.WithTimeout(ctx, 2*time.Minute) defer cancel() - ec := exec.CommandContext(ctx, c.Path, args...) + ec := exec.CommandContext(ctx, path, args...) // vLLM refuses to do anything at all on a machine with no accelerator, // including print its help, unless it is told what to assume. Asking an // engine about its flags should not require the hardware it runs on. @@ -231,7 +240,7 @@ func runHelp(ctx context.Context, c engines.Command) (string, error) { raw, err := ec.CombinedOutput() if err != nil && len(raw) == 0 { return "", fmt.Errorf("could not run %s: %w (install it here, or pass --help-output)", - c.Path, err) + path, err) } return string(raw), nil } diff --git a/internal/cli/enginecheck_test.go b/internal/cli/enginecheck_test.go index becd8e0..1b4d85f 100644 --- a/internal/cli/enginecheck_test.go +++ b/internal/cli/enginecheck_test.go @@ -1,6 +1,11 @@ package cli import ( + "bytes" + "context" + "os" + "path/filepath" + "strings" "testing" ) @@ -69,3 +74,75 @@ func TestATruncatedHelpIsSuspected(t *testing.T) { t.Error("a truncated help was taken for a complete one") } } + +// fakeInterpreter stands in for TensorRT-LLM's Python: asked for +// trtllm-serve's help through its module, it prints the help given; asked +// anything else, it fails, as running the launcher here would. +func fakeInterpreter(t *testing.T, help string) string { + t.Helper() + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "help.txt"), []byte(help), 0o644); err != nil { + t.Fatal(err) + } + p := filepath.Join(dir, "python") + body := "#!/bin/sh\n" + + `if [ "$*" = "-m tensorrt_llm.commands.serve serve --help" ]; then cat "` + filepath.Join(dir, "help.txt") + `"; exit 0; fi` + "\n" + + "echo \"unexpected: $*\" >&2; exit 2\n" + if err := os.WriteFile(p, []byte(body), 0o755); err != nil { //nolint:gosec // an executable stand-in + t.Fatal(err) + } + return p +} + +// trtllmHelp is the part of trtllm-serve serve --help (1.3) that holds the +// flags llmgw builds. +const trtllmHelp = `Usage: trtllm-serve serve [OPTIONS] MODEL + +Options: + --host TEXT Hostname of the server. + --port INTEGER Port of the server. + --tensor_parallel_size, --tp_size INTEGER + Tensor parallelism size. + --pipeline_parallel_size, --pp_size INTEGER + Pipeline parallelism size. + --config, --extra_llm_api_options TEXT + Path to a YAML configuration file. + --served_model_name TEXT The model name used in the API. + --help Show this message and exit. +` + +func runEngineCheck(t *testing.T, args ...string) (string, error) { + t.Helper() + c := newEngineCheckCmd() + var out bytes.Buffer + c.SetOut(&out) + c.SetErr(&out) + c.SetArgs(args) + err := c.ExecuteContext(context.Background()) + return out.String(), err +} + +func TestEngineCheckAsksTrtllmServeThroughItsModule(t *testing.T) { + out, err := runEngineCheck(t, "trtllm", "--bin", fakeInterpreter(t, trtllmHelp)) + if err != nil { + t.Fatalf("%v\n%s", err, out) + } + if !strings.Contains(out, "every flag llmgw builds is accepted") { + t.Errorf("unexpected report:\n%s", out) + } + // The interpreter's own help would describe Python, not trtllm-serve; and + // one node is what trtllm runs, so no multi-node flag is in the command. + if strings.Contains(out, "unexpected:") || strings.Contains(out, "skipped") { + t.Errorf("the help was not asked of trtllm-serve:\n%s", out) + } +} + +func TestEngineCheckCatchesATensorRTLLMTooOldForTheServedName(t *testing.T) { + // TensorRT-LLM 1.2 has no --served_model_name, and a replica of it would + // never be seen serving the name the agent waits for. + old := strings.ReplaceAll(trtllmHelp, " --served_model_name TEXT The model name used in the API.\n", "") + out, err := runEngineCheck(t, "trtllm", "--bin", fakeInterpreter(t, old)) + if err == nil || !strings.Contains(out, "trtllm-serve: --served_model_name") { + t.Errorf("an old trtllm-serve passed: %v\n%s", err, out) + } +} diff --git a/internal/cli/mirror.go b/internal/cli/mirror.go index db6ef70..c63ea54 100644 --- a/internal/cli/mirror.go +++ b/internal/cli/mirror.go @@ -44,6 +44,8 @@ func newMirrorCmd() *cobra.Command { c.AddCommand( newMirrorAddCmd(&dir), + newMirrorImportCmd(&dir), + newMirrorImportOllamaCmd(&dir), newMirrorListCmd(&dir), newMirrorUnusedCmd(&dir), newMirrorRemoveCmd(&dir), diff --git a/internal/cli/mirror_import.go b/internal/cli/mirror_import.go new file mode 100644 index 0000000..fe177a9 --- /dev/null +++ b/internal/cli/mirror_import.go @@ -0,0 +1,216 @@ +package cli + +import ( + "context" + "fmt" + "io" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/spf13/cobra" + + "gitea.mynodes.xyz/ilker/llmgw/internal/mirror" +) + +// placeholderBase is how the mirror's address is written when the operator +// has not said it, as `mirror add` writes it. +const placeholderBase = "http://:8081" + +func newMirrorImportCmd(dir *string) *cobra.Command { + var model, version, as, sum, base string + + c := &cobra.Command{ + Use: "import ... --model NAME", + Short: "Copy model files on this machine into the mirror", + Long: "For weights already here: a GGUF you converted or downloaded by hand,\n" + + "or the files of a safetensors model. The same discipline as add — copied\n" + + "under a temporary name and renamed when whole, the checksum kept beside\n" + + "it for nodes to check against, an existing version never replaced.\n\n" + + "A file named .gguf is checked to be one. Several files go into one\n" + + "version, as the files of one model; its source is then the version,\n" + + "ending in a slash. Run it on the mirror; for a model in Ollama, use\n" + + "import-ollama.", + Example: " llmgw mirror import ./qwen2.5-7b-instruct-q4_k_m.gguf --model qwen2.5-7b-instruct", + Args: cobra.MinimumNArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + if model == "" { + return fmt.Errorf("--model is required: it is the mirror's directory for the model") + } + if as != "" && len(args) > 1 { + return fmt.Errorf("--as names one file; with %d files each keeps its own name", len(args)) + } + if sum != "" && len(args) > 1 { + return fmt.Errorf("--sha256 is one file's checksum; with %d files, import them one at a time", len(args)) + } + if version == "" { + version = time.Now().Format("2006-01-02") + } + out := cmd.OutOrStdout() + var added []mirror.Added + for _, src := range args { + abs, err := filepath.Abs(src) + if err != nil { + return err + } + file := as + if file == "" { + file = filepath.Base(src) + } + e := mirror.Entry{Model: model, Version: version, File: file} + fmt.Fprintf(out, "copying %s\n", src) + a, err := mirror.Import(cmd.Context(), *dir, abs, e, sum, "file://"+abs, progressTo(out)) + if err != nil { + return err + } + fmt.Fprintf(out, "%s/%s/%s %s\n sha256 %s\n", model, version, file, fileSize(a.Entry.Bytes), a.SHA256) + added = append(added, a) + } + fmt.Fprintf(out, "\nPoint a model at it with:\n source: %s\n", importedSource(base, added)) + return nil + }, + } + c.Flags().StringVar(&model, "model", "", "the model's name, one directory level") + c.Flags().StringVar(&version, "version", "", "the version, one directory level (default today's date)") + c.Flags().StringVar(&as, "as", "", "the filename to keep it as, for one file; its own name when omitted") + c.Flags().StringVar(&sum, "sha256", "", "the checksum its publisher gave, to refuse a file that does not match") + c.Flags().StringVar(&base, "base", placeholderBase, "the address nodes reach this mirror at, for the source printed") + return c +} + +func newMirrorImportOllamaCmd(dir *string) *cobra.Command { + var o ollamaImport + + c := &cobra.Command{ + Use: "import-ollama [NAME[:TAG]]", + Short: "Put a model you already have in Ollama into the mirror, for llama.cpp", + Long: "Ollama keeps each model as a GGUF under its models directory, named by\n" + + "its sha256. This copies that file into the mirror — checked against the\n" + + "digest it is named by, and with a vision model's projector beside it —\n" + + "so llama.cpp can serve it on any node, from your own network.\n\n" + + "With no name it lists the models Ollama has. The mirror version is the\n" + + "tag and the start of the digest, since a tag such as latest moves in\n" + + "Ollama and a mirror version must not; --model and --version choose\n" + + "otherwise.\n\n" + + "Ollama's prompt template is not carried: llama.cpp renders the chat\n" + + "template inside the GGUF. Run it on the mirror, where Ollama's directory\n" + + "is readable — ~/.ollama/models, /usr/share/ollama/.ollama/models for\n" + + "Ollama's Linux service, or OLLAMA_MODELS.", + Example: " llmgw mirror import-ollama\n" + + " sudo llmgw mirror import-ollama llama3.2:3b --ollama-dir /usr/share/ollama/.ollama/models", + Args: cobra.MaximumNArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + o.mirrorDir = *dir + if len(args) == 1 { + o.name = args[0] + } + return o.run(cmd.Context(), cmd.OutOrStdout()) + }, + } + c.Flags().StringVar(&o.ollamaDir, "ollama-dir", "", "Ollama's models directory (default: found where Ollama keeps it)") + c.Flags().StringVar(&o.model, "model", "", "the mirror's directory for it (default the Ollama model's name)") + c.Flags().StringVar(&o.version, "version", "", "the mirror version (default its tag and the start of its digest)") + c.Flags().StringVar(&o.base, "base", placeholderBase, "the address nodes reach this mirror at, for the source printed") + return c +} + +// ollamaImport is one import-ollama, from the command line or the console. +type ollamaImport struct { + mirrorDir, ollamaDir, name, model, version, base string + // quiet leaves out the copy's progress, for the console, which shows + // the output once the import is over. + quiet bool +} + +func (o ollamaImport) run(ctx context.Context, out io.Writer) error { + dir, err := mirror.OllamaDir(o.ollamaDir) + if err != nil { + return err + } + if o.name == "" { + names, err := mirror.ListOllama(dir) + if err != nil { + return err + } + if len(names) == 0 { + fmt.Fprintf(out, "Ollama in %s has no models.\n", dir) + return nil + } + fmt.Fprintf(out, "Ollama's models in %s:\n", dir) + for _, n := range names { + fmt.Fprintf(out, " %s\n", n) + } + fmt.Fprintln(out, "\nImport one with: llmgw mirror import-ollama NAME:TAG") + return nil + } + + n, err := mirror.ParseOllamaName(o.name) + if err != nil { + return err + } + m, err := mirror.ReadOllama(dir, n) + if err != nil { + return err + } + fmt.Fprintf(out, "importing %s from %s\n", n, dir) + last := "" + added, err := mirror.ImportOllama(ctx, o.mirrorDir, dir, m, o.model, o.version, func(file string, done, total int64) { + if file != last { + fmt.Fprintf(out, "copying %s\n", file) + last = file + } + if !o.quiet { + progressTo(out)(done, total) + } + }) + if err != nil { + return err + } + for _, a := range added { + fmt.Fprintf(out, "%s/%s/%s %s\n sha256 %s, as Ollama named it\n", + a.Entry.Model, a.Entry.Version, a.Entry.File, fileSize(a.Entry.Bytes), a.SHA256) + } + for _, l := range m.Left { + fmt.Fprintf(out, "Not carried: %s.\n", l) + } + + name := n.Model + if n.Tag != "latest" { + name += "-" + n.Tag + } + fmt.Fprintf(out, "\nA model definition for it, for llmgw ctl model add -f:\n\n"+ + "name: %s\nsource: %s\nprofiles:\n - name: cpu\n engine: llamacpp\n", + name, importedSource(o.base, added)) + if m.ContextLength > 0 { + // Ollama's context, which llama.cpp's own default would not match; + // llama.cpp divides it between its parallel slots. + fmt.Fprintf(out, " args: [\"--ctx-size\", %q] # Ollama ran it with num_ctx %d\n", + strconv.Itoa(m.ContextLength), m.ContextLength) + } + return nil +} + +// importedSource is what a model's source says for what was just imported: +// the file, or the version's directory when it is several. +func importedSource(base string, added []mirror.Added) string { + if base == "" { + base = placeholderBase + } + if len(added) == 1 { + return added[0].Entry.URL(base) + } + e := added[0].Entry + return strings.TrimSuffix(mirror.Entry{Model: e.Model, Version: e.Version}.URL(base), "/") + "/" +} + +// progressTo reports a copy's progress the way add reports a download's. +func progressTo(out io.Writer) func(done, total int64) { + return func(done, total int64) { + if total > 0 { + fmt.Fprintf(out, " %s of %s (%d%%)\n", fileSize(done), fileSize(total), done*100/total) + return + } + fmt.Fprintf(out, " %s so far\n", fileSize(done)) + } +} diff --git a/internal/cli/mirror_import_test.go b/internal/cli/mirror_import_test.go new file mode 100644 index 0000000..b1e4f41 --- /dev/null +++ b/internal/cli/mirror_import_test.go @@ -0,0 +1,144 @@ +package cli + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "os" + "path/filepath" + "strconv" + "strings" + "testing" + + tea "github.com/charmbracelet/bubbletea" +) + +// runMirror runs `llmgw mirror ...` as an operator would and returns what it +// printed. +func runMirror(t *testing.T, args ...string) (string, error) { + t.Helper() + c := Root() + var out bytes.Buffer + c.SetOut(&out) + c.SetErr(&out) + c.SetArgs(append([]string{"mirror"}, args...)) + err := c.ExecuteContext(t.Context()) + return out.String(), err +} + +// ollamaWith lays out an Ollama models directory holding one model, as +// Ollama writes it, and returns the directory and the weights' sha256. +func ollamaWith(t *testing.T, name, tag string, weights []byte, params string) (string, string) { + t.Helper() + dir := filepath.Join(t.TempDir(), "models") + blob := func(content []byte) string { + s := sha256.Sum256(content) + h := hex.EncodeToString(s[:]) + if err := os.MkdirAll(filepath.Join(dir, "blobs"), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "blobs", "sha256-"+h), content, 0o644); err != nil { + t.Fatal(err) + } + return h + } + w := blob(weights) + p := blob([]byte(params)) + manifest := `{"schemaVersion":2,"mediaType":"application/vnd.docker.distribution.manifest.v2+json",` + + `"config":{"mediaType":"application/vnd.docker.container.image.v1+json","digest":"sha256:` + blob([]byte("{}")) + `","size":2},` + + `"layers":[{"mediaType":"application/vnd.ollama.image.model","digest":"sha256:` + w + `","size":` + strconv.Itoa(len(weights)) + `},` + + `{"mediaType":"application/vnd.ollama.image.params","digest":"sha256:` + p + `","size":` + strconv.Itoa(len(params)) + `}]}` + mdir := filepath.Join(dir, "manifests", "registry.ollama.ai", "library", name) + if err := os.MkdirAll(mdir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(mdir, tag), []byte(manifest), 0o644); err != nil { + t.Fatal(err) + } + return dir, w +} + +func TestImportOllamaListsThenImportsAModelForLlamaCpp(t *testing.T) { + weights := []byte("GGUF\x03\x00\x00\x00qwen weights") + ollama, sum := ollamaWith(t, "qwen2.5", "7b", weights, `{"num_ctx":16384}`) + root := t.TempDir() + + out, err := runMirror(t, "--dir", root, "import-ollama", "--ollama-dir", ollama) + if err != nil || !strings.Contains(out, " qwen2.5:7b\n") { + t.Fatalf("listing: %v\n%s", err, out) + } + + out, err = runMirror(t, "--dir", root, "import-ollama", "qwen2.5:7b", "--ollama-dir", ollama, + "--base", "http://models.internal:8081") + if err != nil { + t.Fatalf("%v\n%s", err, out) + } + file := filepath.Join(root, "qwen2.5", "7b-"+sum[:12], "qwen2.5-7b.gguf") + if b, err := os.ReadFile(file); err != nil || !bytes.Equal(b, weights) { + t.Fatalf("the GGUF is not in the mirror at %s: %v", file, err) + } + for _, want := range []string{ + "sha256 " + sum, + "source: http://models.internal:8081/qwen2.5/7b-" + sum[:12] + "/qwen2.5-7b.gguf", + "engine: llamacpp", + `args: ["--ctx-size", "16384"]`, + } { + if !strings.Contains(out, want) { + t.Errorf("the output does not say %q:\n%s", want, out) + } + } +} + +func TestTheConsoleImportsFromOllamaIntoTheMirrorBesideIt(t *testing.T) { + weights := []byte("GGUF\x03\x00\x00\x00llama weights") + ollama, sum := ollamaWith(t, "llama3.2", "3b", weights, `{}`) + root := t.TempDir() + + m, _ := consoleWith(t, cMirror) + var tm tea.Model = press(m, "i") + f := tm.(consoleModel).formD + if f == nil || f.title != "Import from Ollama" { + t.Fatalf("I opened %+v", f) + } + f.fields[0].text.SetValue("llama3.2:3b") + f.fields[1].text.SetValue(ollama) + f.fields[2].text.SetValue(root) + for range len(f.fields) { + tm = press(tm, "tab") + } + tm = pressRun(tm, "enter") + + if _, err := os.Stat(filepath.Join(root, "llama3.2", "3b-"+sum[:12], "llama3.2-3b.gguf")); err != nil { + t.Fatalf("not imported: %v", err) + } + text := tm.(consoleModel).dialogText + if !strings.Contains(text, "engine: llamacpp") || strings.Contains(text, "(100%)") { + t.Errorf("the console said:\n%s", text) + } +} + +func TestImportCopiesALocalGGUFIntoTheMirror(t *testing.T) { + src := filepath.Join(t.TempDir(), "q4_k_m.gguf") + content := []byte("GGUF\x03\x00\x00\x00converted") + if err := os.WriteFile(src, content, 0o644); err != nil { + t.Fatal(err) + } + root := t.TempDir() + out, err := runMirror(t, "--dir", root, "import", src, "--model", "qwen", "--version", "2026-10") + if err != nil { + t.Fatalf("%v\n%s", err, out) + } + if b, err := os.ReadFile(filepath.Join(root, "qwen", "2026-10", "q4_k_m.gguf")); err != nil || !bytes.Equal(b, content) { + t.Fatalf("not imported: %v", err) + } + if !strings.Contains(out, "source: http://:8081/qwen/2026-10/q4_k_m.gguf") { + t.Errorf("no source given:\n%s", out) + } + // The layout is the point: a version is never replaced. + if _, err := runMirror(t, "--dir", root, "import", src, "--model", "qwen", "--version", "2026-10"); err == nil { + t.Error("an existing version was replaced") + } + if _, err := runMirror(t, "--dir", root, "import", src); err == nil { + t.Error("imported without --model") + } +} diff --git a/internal/cli/tui_console.go b/internal/cli/tui_console.go index e96af5e..8cbef3d 100644 --- a/internal/cli/tui_console.go +++ b/internal/cli/tui_console.go @@ -536,6 +536,10 @@ func (m consoleModel) key(k tea.KeyMsg) (tea.Model, tea.Cmd) { case cOrgs: m.createOrgForm() } + case "i": + if m.menu == cMirror { + m.mirrorImportOllamaForm() + } case "n": if m.menu == cMonitor { return m, m.runCLI("Test the webhooks", "ctl", "--server="+m.client.base, "notify", "test") diff --git a/internal/cli/tui_mirror.go b/internal/cli/tui_mirror.go index 68d0f69..7dade93 100644 --- a/internal/cli/tui_mirror.go +++ b/internal/cli/tui_mirror.go @@ -1,6 +1,7 @@ package cli import ( + "bytes" "context" "errors" "fmt" @@ -178,10 +179,41 @@ func (m *consoleModel) mirrorFetchForm() { }}) } +// mirrorImportOllamaForm puts a model from an Ollama on this machine into a +// mirror on this machine: the control plane's own, when the mirror runs +// beside it. A mirror elsewhere cannot read this machine's files, and its +// API is given none; there it is `llmgw mirror import-ollama`, on the mirror. +func (m *consoleModel) mirrorImportOllamaForm() { + m.openForm(&formDialog{title: "Import from Ollama", + note: "From an Ollama on this machine into the mirror's directory on this machine.\n" + + "A mirror elsewhere: run llmgw mirror import-ollama there.", + fields: []formField{ + textF("model", "", "as Ollama names it, like llama3.2:3b"), + textF("ollama dir", "", "optional: found where Ollama keeps it"), + textF("mirror dir", "/srv/models", "the mirror's directory on this machine"), + }, + submit: func(m *consoleModel, v []string) (tea.Cmd, error) { + if v[0] == "" { + return nil, fmt.Errorf("a model, as Ollama names it") + } + if st, err := os.Stat(v[2]); err != nil || !st.IsDir() { + return nil, fmt.Errorf("%s is not a directory on this machine", v[2]) + } + title := "Import " + v[0] + m.busy = title + o := ollamaImport{name: v[0], ollamaDir: v[1], mirrorDir: v[2], base: m.mirror.res.Base, quiet: true} + return func() tea.Msg { + var out bytes.Buffer + err := o.run(context.Background(), &out) + return actionDoneMsg{title: title, text: strings.TrimSpace(out.String()), err: err} + }, nil + }}) +} + func (m consoleModel) mirrorPanel(pw int) (string, string, string) { keys := "→ or Enter open · ↑↓ choose · Q leave" if m.inPanel { - keys = "↑↓ row · Enter actions · A download one · ← back · Q leave" + keys = "↑↓ row · Enter actions · A download one · I import from Ollama · ← back · Q leave" } if m.mirror.err != nil { return "Mirror", "\n " + sBad.Render(m.mirror.err.Error()), keys @@ -216,7 +248,7 @@ func (m consoleModel) mirrorPanel(pw int) (string, string, string) { rows = append(rows, []string{g.model, g.version, what, state, bytesShort(total), used}) } if len(rows) == 0 { - return "Mirror", head + " " + sDim.Render("Empty. A downloads one."), keys + return "Mirror", head + " " + sDim.Render("Empty. A downloads one; I imports one from Ollama."), keys } return "Mirror", head + table([]string{"MODEL", "VERSION", "FILE", "STATE", "SIZE", "USED BY"}, rows, m.row, pw, m.inPanel), keys } diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index 602082b..66379f6 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -717,6 +717,19 @@ func (m Model) validate() error { } continue } + // TensorRT-LLM spans machines only under MPI started from one of them, + // which the agents, each starting its own node, do not do; and it + // loads a Hugging Face checkpoint, not llama.cpp's format. + if p.Engine == "trtllm" { + if p.Nodes > 1 { + return fmt.Errorf("model %q profile %q: trtllm runs a replica on one node; "+ + "use vllm or sglang for a %d-node replica", m.Name, p.Name, p.Nodes) + } + if strings.HasSuffix(strings.ToLower(strings.TrimSuffix(m.Source, "/")), ".gguf") { + return fmt.Errorf("model %q profile %q: trtllm loads a Hugging Face checkpoint, not a .gguf; "+ + "serve a GGUF with llamacpp", m.Name, p.Name) + } + } // The arithmetic has to add up, or the engine is handed a shard spec it // cannot satisfy and fails minutes later during model loading. if want := p.GPUsPerNode * p.Nodes; want != p.TP*max(p.PP, 1) { diff --git a/internal/control/catalog/trtllm_test.go b/internal/control/catalog/trtllm_test.go new file mode 100644 index 0000000..e9a8ef5 --- /dev/null +++ b/internal/control/catalog/trtllm_test.go @@ -0,0 +1,53 @@ +package catalog + +import ( + "strings" + "testing" +) + +func TestATensorRTLLMProfileIsAGPUProfileOnOneNode(t *testing.T) { + ok := Model{Name: "llama", Source: "meta-llama/Llama-3.1-70B-Instruct", Profiles: []Profile{{ + Name: "h100", Engine: "trtllm", GPUClass: "h100-80g", TP: 4, + Args: []string{"--max_batch_size", "64"}, + }}} + c, err := New([]Model{ok}) + if err != nil { + t.Fatalf("a four-GPU TensorRT-LLM profile was refused: %v", err) + } + m, _ := c.Resolve("llama") + if p := m.PlacementProfiles()[0]; p.GPUsPerNode != 4 || p.Nodes != 1 { + t.Errorf("placed as %+v, expected four GPUs of one node", p) + } + + for want, bad := range map[string]Model{ + // Two nodes need MPI started from one of them. + "one node": {Name: "x", Source: "org/x", Profiles: []Profile{{ + Name: "p", Engine: "trtllm", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8}}}, + // No CPU build exists. + "gpu_class is required": {Name: "x", Source: "org/x", Profiles: []Profile{{Name: "p", Engine: "trtllm"}}}, + // llama.cpp's format. + "not a .gguf": {Name: "x", Source: "http://mirror:8081/x/v1/x-q4.gguf", Profiles: []Profile{{ + Name: "p", Engine: "trtllm", GPUClass: "h100-80g"}}}, + // What its launcher closes, a profile may not open. + "--trust_remote_code is not allowed": {Name: "x", Source: "org/x", Profiles: []Profile{{ + Name: "p", Engine: "trtllm", GPUClass: "h100-80g", Args: []string{"--trust_remote_code"}}}}, + } { + if _, err := New([]Model{bad}); err == nil || !strings.Contains(err.Error(), want) { + t.Errorf("expected %q, got %v", want, err) + } + } +} + +func TestATensorRTLLMModelIsGivenTheRepositoryNotItsPage(t *testing.T) { + // trtllm-serve downloads a repository itself, as vLLM does. + m := Model{Name: "q", Source: "https://huggingface.co/Qwen/Qwen2.5-7B-Instruct", Profiles: []Profile{{ + Name: "p", Engine: "trtllm", GPUClass: "l40s-48g"}}} + c, err := New([]Model{m}) + if err != nil { + t.Fatal(err) + } + got, _ := c.Resolve("q") + if !strings.HasSuffix(got.Source, "Qwen/Qwen2.5-7B-Instruct") || strings.Contains(got.Source, "huggingface.co") { + t.Errorf("source %q", got.Source) + } +} diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index 4503a72..9321242 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -329,6 +329,9 @@ var Reference = []Block{ {Package: "engines", Name: "SimConfig", Title: "SimConfig", Doc: "", Fields: []Field{ {Key: "bin", Type: "string", Doc: "Bin is the llmgw binary itself; the engine is one of its subcommands.\nEmpty means whatever started this agent, which is right almost always.", Nested: ""}, }}, + {Package: "engines", Name: "TRTLLMConfig", Title: "TRTLLMConfig", Doc: "", Fields: []Field{ + {Key: "python", Type: "string", Doc: "Python is the interpreter of the environment TensorRT-LLM is installed\nin: /bin/python for a virtual environment, python3 in NVIDIA's\ncontainer. Naming it is what makes this node offer the engine; empty\noffers nothing.\n\nThe interpreter rather than trtllm-serve, because llmgw starts the\nserver through a few lines of its own that make it demand the\ngateway's key, which trtllm-serve has no option for.", Nested: ""}, + }}, {Package: "config", Name: "Threads", Title: "Threads", Doc: "", Fields: []Field{ {Key: "store", Type: "string", Doc: "Store is the SQLite file. Empty means nothing is kept.", Nested: ""}, {Key: "retain_hours", Type: "int", Doc: "RetainHours is how long idle conversations and notes are kept. Required\nwith Store: a store with no expiry keeps everything anyone typed, by\nomission.", Nested: ""}, @@ -423,6 +426,7 @@ var Reference = []Block{ {Key: "vllm", Type: "VLLMConfig", Doc: "", Nested: "engines.VLLMConfig"}, {Key: "sglang", Type: "SGLangConfig", Doc: "", Nested: "engines.SGLangConfig"}, {Key: "llamacpp", Type: "LlamaCppConfig", Doc: "", Nested: "engines.LlamaCppConfig"}, + {Key: "trtllm", Type: "TRTLLMConfig", Doc: "TRTLLM is TensorRT-LLM, which the agent does not install: a node\noffers it only when its interpreter is named here.", Nested: "engines.TRTLLMConfig"}, {Key: "sim", Type: "SimConfig", Doc: "", Nested: "engines.SimConfig"}, }}, {Package: "guard", Name: "Config", Title: "guard.Config", Doc: "", Fields: []Field{ diff --git a/internal/docgen/render.go b/internal/docgen/render.go index 89abf06..d70f501 100644 --- a/internal/docgen/render.go +++ b/internal/docgen/render.go @@ -20,6 +20,7 @@ var Sources = []string{ "internal/agent/engines/llamacpp.go", "internal/agent/engines/vllm.go", "internal/agent/engines/sglang.go", + "internal/agent/engines/trtllm.go", "internal/agent/engines/sim.go", "internal/agent/supervisor/supervisor.go", "internal/tlsconf/tlsconf.go", diff --git a/internal/mirror/import.go b/internal/mirror/import.go new file mode 100644 index 0000000..f49cd63 --- /dev/null +++ b/internal/mirror/import.go @@ -0,0 +1,131 @@ +package mirror + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "os" + "path/filepath" + "strings" +) + +// Import copies a file on this machine into the mirror at model/version/file: +// weights somebody already has — a GGUF they converted, Ollama's blob of +// one — put where every node can fetch them, with the checksum nodes check +// them against. +// +// The same discipline as Add: written under a temporary name and renamed +// when whole, the checksum and the source beside it first, and an existing +// version never replaced. wantSHA256, when given, is what the file must hash +// to — Ollama names each blob by its digest — and a file that does not is +// refused rather than copied to every node. +// +// source is what is recorded as where the file came from. +func Import(ctx context.Context, root, src string, e Entry, wantSHA256, source string, + progress func(done, total int64), +) (Added, error) { + if e.Model == "" || e.Version == "" || e.File == "" { + return Added{}, fmt.Errorf("a model, a version and a filename are all required") + } + for _, n := range []string{e.Model, e.Version, e.File} { + if err := ValidName(n); err != nil { + return Added{}, err + } + } + in, err := os.Open(src) //nolint:gosec // a file the operator named + if err != nil { + return Added{}, err + } + defer in.Close() + st, err := in.Stat() + if err != nil { + return Added{}, err + } + if !st.Mode().IsRegular() { + return Added{}, fmt.Errorf("%s is not a file", src) + } + if strings.HasSuffix(strings.ToLower(e.File), ".gguf") { + if err := isGGUF(in, src); err != nil { + return Added{}, err + } + } + + final := e.Path(root) + if _, err := os.Stat(final); err == nil { + return Added{}, fmt.Errorf("%s/%s/%s: %w", e.Model, e.Version, e.File, ErrExists) + } + if err := os.MkdirAll(filepath.Dir(final), 0o755); err != nil { //nolint:gosec // served to every node + return Added{}, err + } + + part := final + partSuffix + // Truncated rather than refused when there: a .part is what an import cut + // short left, and the next attempt is the one that should finish it. + out, err := os.OpenFile(part, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644) //nolint:gosec // inside the mirror root + if err != nil { + return Added{}, err + } + // Hashed on the way through: a forty gigabyte file is read once, not + // copied and then read again. + h := sha256.New() + written, err := copyWithProgress(io.MultiWriter(out, h), ctxReader{ctx, in}, st.Size(), progress) + closeErr := out.Close() + if err == nil { + err = closeErr + } + if err == nil && written != st.Size() { + err = fmt.Errorf("%s: copied %d bytes of %d; it changed while being read", src, written, st.Size()) + } + sum := hex.EncodeToString(h.Sum(nil)) + if err == nil && wantSHA256 != "" && !strings.EqualFold(sum, wantSHA256) { + err = fmt.Errorf("%s is damaged: it should hash to sha256 %s and hashes to %s; "+ + "nothing was added", src, wantSHA256, sum) + } + if err == nil { + err = writeChecksum(final, sum) + } + if err == nil { + err = writeSource(final, source) + } + if err == nil { + err = os.Rename(part, final) + } + if err != nil { + _ = os.Remove(part) + return Added{}, err + } + e.Bytes = written + return Added{Entry: e, SHA256: sum}, nil +} + +// ggufMagic opens every GGUF file. +const ggufMagic = "GGUF" + +// isGGUF checks a file named .gguf is one, and leaves it read from the start. +// A file that is not — an HTML error page saved under the name, Ollama's +// newer safetensors-based layers — would otherwise reach every node and fail +// there, in llama.cpp's words. +func isGGUF(f *os.File, name string) error { + head := make([]byte, len(ggufMagic)) + if _, err := io.ReadFull(f, head); err != nil || string(head) != ggufMagic { + return fmt.Errorf("%s is not a GGUF file: it does not begin with %q", name, ggufMagic) + } + _, err := f.Seek(0, io.SeekStart) + return err +} + +// ctxReader stops a copy when its context ends: an operator's ^C on a long +// import leaves no partial file behind. +type ctxReader struct { + ctx context.Context + r io.Reader +} + +func (c ctxReader) Read(p []byte) (int, error) { + if err := c.ctx.Err(); err != nil { + return 0, err + } + return c.r.Read(p) +} diff --git a/internal/mirror/ollama.go b/internal/mirror/ollama.go new file mode 100644 index 0000000..77984b4 --- /dev/null +++ b/internal/mirror/ollama.go @@ -0,0 +1,352 @@ +package mirror + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "regexp" + "sort" + "strings" +) + +// Ollama keeps a model as a manifest and the blobs it names. The manifest is +// at manifests//// under its models directory — +// registry.ollama.ai/library/llama3.2/3b for "llama3.2:3b" — and each layer +// is a file blobs/sha256-, named by its own digest. The layer whose media +// type is application/vnd.ollama.image.model is the weights as a GGUF, which +// is what llama.cpp runs; a vision model has its projector in a second GGUF. +// +// What an operator already has in Ollama is weights they chose and have +// tried. Importing them puts that file in the mirror, checked against the +// digest Ollama named it by, rather than downloading it again from somewhere +// that may not be the same file. + +// Ollama's defaults for the parts of a name left out. +const ( + ollamaHost = "registry.ollama.ai" + ollamaNamespace = "library" + ollamaTag = "latest" +) + +// The layer media types an import reads, or says it leaves. +const ( + ollamaModelLayer = "application/vnd.ollama.image.model" + ollamaProjectorLayer = "application/vnd.ollama.image.projector" + ollamaParamsLayer = "application/vnd.ollama.image.params" + ollamaTensorLayer = "application/vnd.ollama.image.tensor" +) + +// OllamaName is a model as Ollama names it. +type OllamaName struct { + Host, Namespace, Model, Tag string +} + +// ParseOllamaName reads a name as Ollama does: model, namespace/model or +// host/namespace/model, with an optional :tag, and its defaults for whatever +// is left out — "llama3.2" is registry.ollama.ai/library/llama3.2:latest. +func ParseOllamaName(s string) (OllamaName, error) { + s = strings.TrimSpace(s) + var n OllamaName + // A digest pins a name in Ollama; the manifest on disk is by tag. + s, _, _ = strings.Cut(s, "@") + // "/" cannot be in a tag, so a colon after the last slash is the tag's. + if i := strings.LastIndex(s, ":"); i > strings.LastIndex(s, "/") { + s, n.Tag = s[:i], s[i+1:] + } + parts := strings.Split(s, "/") + switch len(parts) { + case 1: + n.Model = parts[0] + case 2: + n.Namespace, n.Model = parts[0], parts[1] + case 3: + n.Host, n.Namespace, n.Model = parts[0], parts[1], parts[2] + default: + return OllamaName{}, fmt.Errorf("%q is not an Ollama model name: model[:tag], namespace/model[:tag] "+ + "or host/namespace/model[:tag]", s) + } + if n.Host == "" { + n.Host = ollamaHost + } + if n.Namespace == "" { + n.Namespace = ollamaNamespace + } + if n.Tag == "" { + n.Tag = ollamaTag + } + for _, p := range []string{n.Host, n.Namespace, n.Model, n.Tag} { + // Each part is a directory of Ollama's; one that is not a single path + // element would read a manifest from somewhere else. + if p == "" || p == "." || p == ".." || strings.ContainsAny(p, `\`+"\x00") { + return OllamaName{}, fmt.Errorf("%q is not an Ollama model name", s) + } + } + return n, nil +} + +// String is the name the way Ollama shows it: without the parts that are its +// defaults. +func (n OllamaName) String() string { + name := n.Model + ":" + n.Tag + switch { + case !strings.EqualFold(n.Host, ollamaHost): + return n.Host + "/" + n.Namespace + "/" + name + case !strings.EqualFold(n.Namespace, ollamaNamespace): + return n.Namespace + "/" + name + } + return name +} + +func (n OllamaName) manifestPath(dir string) string { + return filepath.Join(dir, "manifests", n.Host, n.Namespace, n.Model, n.Tag) +} + +// OllamaDir finds Ollama's models directory: the one given (or its models +// subdirectory, for a ~/.ollama given whole), else OLLAMA_MODELS, else the +// places Ollama keeps it — this user's ~/.ollama/models, and on Linux the +// home of the ollama user its installer makes. +func OllamaDir(given string) (string, error) { + var candidates []string + if given != "" { + candidates = []string{given, filepath.Join(given, "models")} + } else { + if env := os.Getenv("OLLAMA_MODELS"); env != "" { + candidates = append(candidates, env) + } + if home, err := os.UserHomeDir(); err == nil { + candidates = append(candidates, filepath.Join(home, ".ollama", "models")) + } + candidates = append(candidates, + "/usr/share/ollama/.ollama/models", // Ollama's Linux installer + "/var/lib/ollama/.ollama/models", // some distributions' packages + ) + } + for _, c := range candidates { + if st, err := os.Stat(filepath.Join(c, "manifests")); err == nil && st.IsDir() { + return c, nil + } + } + if given != "" { + return "", fmt.Errorf("%s holds no Ollama models: there is no manifests directory in it or in its models", given) + } + return "", fmt.Errorf("no Ollama models found in %s; name the directory with --ollama-dir", + strings.Join(candidates, ", ")) +} + +// ListOllama lists the models in an Ollama models directory. +func ListOllama(dir string) ([]OllamaName, error) { + root := filepath.Join(dir, "manifests") + var out []OllamaName + err := filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { + if err != nil || d.IsDir() { + return err + } + rel, err := filepath.Rel(root, path) + if err != nil { + return err + } + parts := strings.Split(rel, string(filepath.Separator)) + if len(parts) != 4 || strings.HasPrefix(parts[3], ".") { + return nil // not a manifest: something else left in Ollama's directory + } + out = append(out, OllamaName{Host: parts[0], Namespace: parts[1], Model: parts[2], Tag: parts[3]}) + return nil + }) + if err != nil { + return nil, fmt.Errorf("reading Ollama's models in %s: %w", dir, err) + } + sort.Slice(out, func(i, j int) bool { return out[i].String() < out[j].String() }) + return out, nil +} + +// OllamaLayer is one layer of an Ollama manifest. +type OllamaLayer struct { + MediaType string `json:"mediaType"` + Digest string `json:"digest"` + Size int64 `json:"size"` +} + +// ollamaManifest is the manifest's shape: Docker's image manifest, with +// Ollama's media types for the layers. +type ollamaManifest struct { + SchemaVersion int `json:"schemaVersion"` + Layers []OllamaLayer `json:"layers"` +} + +// OllamaModel is what an import takes from one of Ollama's models. +type OllamaModel struct { + Name OllamaName + // Model is the weights, a GGUF. + Model OllamaLayer + // Projectors are a vision model's image encoders, GGUFs beside it. + Projectors []OllamaLayer + // ContextLength is the num_ctx Ollama was told to run it with, 0 when + // it was not: worth carrying over, since llama.cpp's own default differs. + ContextLength int + // Left says what of the Ollama model an import does not carry. + Left []string +} + +var digestPattern = regexp.MustCompile(`^sha256:([0-9a-f]{64})$`) + +// ollamaBlob is where a layer's file is, and the sha256 its name promises. +func ollamaBlob(dir, digest string) (path, sum string, err error) { + m := digestPattern.FindStringSubmatch(digest) + if m == nil { + return "", "", fmt.Errorf("the manifest names a layer by %q, which is not a sha256 digest", digest) + } + return filepath.Join(dir, "blobs", "sha256-"+m[1]), m[1], nil +} + +// ReadOllama reads one model's manifest and the layers an import takes. +func ReadOllama(dir string, name OllamaName) (OllamaModel, error) { + b, err := os.ReadFile(name.manifestPath(dir)) //nolint:gosec // a manifest under Ollama's directory + if errors.Is(err, os.ErrNotExist) { + return OllamaModel{}, fmt.Errorf("no %s among the Ollama models in %s; `llmgw mirror import-ollama` with no name lists them", + name, dir) + } + if err != nil { + return OllamaModel{}, err + } + var mf ollamaManifest + if err := json.Unmarshal(b, &mf); err != nil { + return OllamaModel{}, fmt.Errorf("%s's manifest cannot be read: %w", name, err) + } + + m := OllamaModel{Name: name} + var models []OllamaLayer + for _, l := range mf.Layers { + switch l.MediaType { + case ollamaModelLayer: + models = append(models, l) + case ollamaProjectorLayer: + m.Projectors = append(m.Projectors, l) + case ollamaTensorLayer: + return OllamaModel{}, fmt.Errorf("%s is kept by Ollama as separate tensors rather than a GGUF; "+ + "llama.cpp cannot run that form: mirror the model's GGUF from its publisher instead", name) + case ollamaParamsLayer: + m.ContextLength = ollamaContext(dir, l) + case "application/vnd.ollama.image.template", "application/vnd.ollama.image.prompt": + m.Left = append(m.Left, "its prompt template: llama.cpp renders the chat template inside the GGUF") + case "application/vnd.ollama.image.system": + m.Left = append(m.Left, "its system prompt") + case "application/vnd.ollama.image.adapter": + m.Left = append(m.Left, "a LoRA adapter: give it to the model as one of its adapters") + case "application/vnd.ollama.image.draft": + m.Left = append(m.Left, "its draft model for speculative decoding") + } + } + switch len(models) { + case 0: + return OllamaModel{}, fmt.Errorf("%s has no GGUF layer (%s) to import", name, ollamaModelLayer) + case 1: + m.Model = models[0] + default: + return OllamaModel{}, fmt.Errorf("%s is made of %d weight layers, which an import does not join yet; "+ + "mirror the model's GGUF from its publisher instead", name, len(models)) + } + return m, nil +} + +// ollamaContext is the num_ctx in a params layer, 0 when there is none or it +// cannot be read: a hint, not worth failing an import over. +func ollamaContext(dir string, l OllamaLayer) int { + path, _, err := ollamaBlob(dir, l.Digest) + if err != nil { + return 0 + } + b, err := os.ReadFile(path) //nolint:gosec // a blob under Ollama's directory + if err != nil { + return 0 + } + var p struct { + NumCtx int `json:"num_ctx"` + } + _ = json.Unmarshal(b, &p) + return p.NumCtx +} + +// OllamaEntries are where an Ollama model's files go in the mirror, by +// default: the model's own name, and a version that is its tag with the start +// of its digest — a tag such as "latest" moves in Ollama, and a mirror +// version must not. +func OllamaEntries(m OllamaModel, model, version string) (weights Entry, projectors []Entry) { + if model == "" { + model = m.Name.Model + } + _, sum, _ := ollamaBlob("", m.Model.Digest) + if version == "" { + version = m.Name.Tag + "-" + sum[:min(12, len(sum))] + } + base := m.Name.Model + "-" + m.Name.Tag + weights = Entry{Model: model, Version: version, File: base + ".gguf"} + for i := range m.Projectors { + // "mmproj" in the name is what llama.cpp's engine knows a projector by, + // beside the model in one directory. + file := "mmproj-" + base + ".gguf" + if i > 0 { + file = fmt.Sprintf("mmproj-%s-%d.gguf", base, i+1) + } + projectors = append(projectors, Entry{Model: model, Version: version, File: file}) + } + return weights, projectors +} + +// ImportOllama copies an Ollama model's GGUF, and its projectors, into one +// version of the mirror, each checked against the digest Ollama named it by. +// progress is told which file it is about. +// +// All or nothing: a version left holding a projector and no model is one a +// node would fetch and fail to start, so the files already copied are taken +// out again when a later one fails. +func ImportOllama(ctx context.Context, root, dir string, m OllamaModel, model, version string, + progress func(file string, done, total int64), +) ([]Added, error) { + weights, projectors := OllamaEntries(m, model, version) + type job struct { + layer OllamaLayer + entry Entry + } + jobs := []job{{m.Model, weights}} + for i, l := range m.Projectors { + jobs = append(jobs, job{l, projectors[i]}) + } + // Refused before anything is copied, rather than after the first file. + if _, err := os.Stat(filepath.Join(root, weights.Model, weights.Version)); err == nil { + return nil, fmt.Errorf("%s/%s: %w", weights.Model, weights.Version, ErrExists) + } + + source := "ollama://" + m.Name.Host + "/" + m.Name.Namespace + "/" + m.Name.Model + ":" + m.Name.Tag + var done []Added + for _, j := range jobs { + a, err := importLayer(ctx, root, dir, j.layer, j.entry, source, progress) + if err != nil { + for _, prev := range done { + _ = Remove(root, prev.Entry) + } + return nil, fmt.Errorf("%s: %w", m.Name, err) + } + done = append(done, a) + } + return done, nil +} + +func importLayer(ctx context.Context, root, dir string, l OllamaLayer, e Entry, source string, + progress func(file string, done, total int64), +) (Added, error) { + src, sum, err := ollamaBlob(dir, l.Digest) + if err != nil { + return Added{}, err + } + if _, err := os.Stat(src); err != nil { + return Added{}, fmt.Errorf("its layer %s is not in %s: %w", l.Digest, filepath.Join(dir, "blobs"), err) + } + var report func(int64, int64) + if progress != nil { + report = func(d, t int64) { progress(e.File, d, t) } + } + return Import(ctx, root, src, e, sum, source, report) +} diff --git a/internal/mirror/ollama_test.go b/internal/mirror/ollama_test.go new file mode 100644 index 0000000..79fac6f --- /dev/null +++ b/internal/mirror/ollama_test.go @@ -0,0 +1,369 @@ +package mirror + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "os" + "path/filepath" + "slices" + "strings" + "testing" + + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines" +) + +// ollamaFixture is an Ollama models directory laid out as Ollama writes it: +// blobs/sha256- named by their content's digest, and a manifest per +// model and tag under manifests////. +type ollamaFixture struct { + t *testing.T + dir string +} + +func newOllamaFixture(t *testing.T) *ollamaFixture { + t.Helper() + dir := filepath.Join(t.TempDir(), ".ollama", "models") + for _, d := range []string{"blobs", "manifests"} { + if err := os.MkdirAll(filepath.Join(dir, d), 0o755); err != nil { + t.Fatal(err) + } + } + return &ollamaFixture{t: t, dir: dir} +} + +// blob stores content as Ollama does and returns the layer naming it. +func (f *ollamaFixture) blob(mediaType string, content []byte) OllamaLayer { + f.t.Helper() + sum := sha256.Sum256(content) + hexsum := hex.EncodeToString(sum[:]) + if err := os.WriteFile(filepath.Join(f.dir, "blobs", "sha256-"+hexsum), content, 0o644); err != nil { + f.t.Fatal(err) + } + return OllamaLayer{MediaType: mediaType, Digest: "sha256:" + hexsum, Size: int64(len(content))} +} + +func (f *ollamaFixture) manifest(name string, layers ...OllamaLayer) { + f.t.Helper() + n, err := ParseOllamaName(name) + if err != nil { + f.t.Fatal(err) + } + config := f.blob("application/vnd.docker.container.image.v1+json", + []byte(`{"model_format":"gguf","model_family":"llama","file_type":"Q4_K_M"}`)) + b, _ := json.Marshal(map[string]any{ + "schemaVersion": 2, + "mediaType": "application/vnd.docker.distribution.manifest.v2+json", + "config": config, + "layers": layers, + }) + p := n.manifestPath(f.dir) + if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil { + f.t.Fatal(err) + } + if err := os.WriteFile(p, b, 0o644); err != nil { + f.t.Fatal(err) + } +} + +// gguf is a small file that begins as every GGUF does. +func gguf(body string) []byte { return []byte("GGUF\x03\x00\x00\x00" + body) } + +// llama32 is llama3.2:3b as Ollama pulls it: the weights, a template, a +// license and parameters. +func (f *ollamaFixture) llama32() OllamaLayer { + weights := f.blob(ollamaModelLayer, gguf("llama 3.2 3b weights")) + f.manifest("llama3.2:3b", + weights, + f.blob("application/vnd.ollama.image.template", []byte(`<|start_header_id|>{{ .Prompt }}`)), + f.blob("application/vnd.ollama.image.license", []byte("LLAMA 3.2 COMMUNITY LICENSE")), + f.blob(ollamaParamsLayer, []byte(`{"num_ctx":8192,"stop":["<|eot_id|>"]}`)), + ) + return weights +} + +func TestOllamaNamesAreReadAsOllamaReadsThem(t *testing.T) { + for in, want := range map[string]OllamaName{ + "llama3.2": {ollamaHost, ollamaNamespace, "llama3.2", "latest"}, + "llama3.2:3b": {ollamaHost, ollamaNamespace, "llama3.2", "3b"}, + "jmorgan/llava:7b-q4": {ollamaHost, "jmorgan", "llava", "7b-q4"}, + "hf.co/bartowski/Qwen-GGUF:Q4": {"hf.co", "bartowski", "Qwen-GGUF", "Q4"}, + "localhost:5000/me/m:v1": {"localhost:5000", "me", "m", "v1"}, + "qwen3:8b@sha256:abcd": {ollamaHost, ollamaNamespace, "qwen3", "8b"}, + } { + got, err := ParseOllamaName(in) + if err != nil || got != want { + t.Errorf("%q: got %+v, %v; want %+v", in, got, err, want) + } + } + for _, bad := range []string{"", "a/b/c/d", "../etc:passwd", "x:.."} { + if _, err := ParseOllamaName(bad); err == nil { + t.Errorf("%q was taken for a name", bad) + } + } + // Shown the way Ollama shows them: without its defaults. + for in, want := range map[string]string{ + "llama3.2": "llama3.2:latest", "jmorgan/llava:7b": "jmorgan/llava:7b", + "hf.co/bartowski/Qwen-GGUF:Q4": "hf.co/bartowski/Qwen-GGUF:Q4", + } { + n, _ := ParseOllamaName(in) + if n.String() != want { + t.Errorf("%q shows as %q", in, n.String()) + } + } +} + +func TestOllamasDirectoryIsFoundWhereOllamaKeepsIt(t *testing.T) { + f := newOllamaFixture(t) + // Given as the models directory, or as the ~/.ollama it is in. + for _, given := range []string{f.dir, filepath.Dir(f.dir)} { + if got, err := OllamaDir(given); err != nil || got != f.dir { + t.Errorf("%s: got %q, %v", given, got, err) + } + } + t.Setenv("OLLAMA_MODELS", f.dir) + if got, err := OllamaDir(""); err != nil || got != f.dir { + t.Errorf("OLLAMA_MODELS: got %q, %v", got, err) + } + if _, err := OllamaDir(t.TempDir()); err == nil { + t.Error("a directory with no manifests was taken for Ollama's") + } +} + +func TestOllamasModelsAreListedFromItsManifests(t *testing.T) { + f := newOllamaFixture(t) + f.llama32() + f.manifest("hf.co/bartowski/Qwen2.5-7B-Instruct-GGUF:Q4_K_M", f.blob(ollamaModelLayer, gguf("qwen"))) + // Ollama's own leftovers are not models. + _ = os.WriteFile(filepath.Join(f.dir, "manifests", ".DS_Store"), nil, 0o644) + + got, err := ListOllama(f.dir) + if err != nil { + t.Fatal(err) + } + var names []string + for _, n := range got { + names = append(names, n.String()) + } + if !slices.Equal(names, []string{"hf.co/bartowski/Qwen2.5-7B-Instruct-GGUF:Q4_K_M", "llama3.2:3b"}) { + t.Errorf("listed %v", names) + } +} + +func TestAnOllamaManifestYieldsItsGGUFAndWhatIsLeftBehind(t *testing.T) { + f := newOllamaFixture(t) + weights := f.llama32() + m, err := ReadOllama(f.dir, mustName(t, "llama3.2:3b")) + if err != nil { + t.Fatal(err) + } + if m.Model != weights { + t.Errorf("model layer %+v, want %+v", m.Model, weights) + } + // What Ollama ran it with, to carry over as --ctx-size. + if m.ContextLength != 8192 { + t.Errorf("num_ctx=%d", m.ContextLength) + } + // The template is Ollama's own language; saying it is left is the honest + // part, since llama.cpp will use the GGUF's. + if len(m.Left) != 1 || !strings.Contains(m.Left[0], "template") { + t.Errorf("left: %v", m.Left) + } +} + +func TestOllamaModelsLlamaCppCannotRunAreRefused(t *testing.T) { + f := newOllamaFixture(t) + // Ollama's newer layout keeps some models as separate tensors. + f.manifest("gpt-oss:20b", f.blob(ollamaTensorLayer, []byte("tensor")), f.blob(ollamaTensorLayer, []byte("other"))) + // Two weight layers make a model this import does not join. + f.manifest("split:1", f.blob(ollamaModelLayer, gguf("a")), f.blob(ollamaModelLayer, gguf("b"))) + f.manifest("empty:1", f.blob("application/vnd.ollama.image.license", []byte("MIT"))) + + for name, want := range map[string]string{ + "gpt-oss:20b": "separate tensors", "split:1": "2 weight layers", "empty:1": "no GGUF layer", + "absent:1": "no absent:1", + } { + _, err := ReadOllama(f.dir, mustName(t, name)) + if err == nil || !strings.Contains(err.Error(), want) { + t.Errorf("%s: expected %q, got %v", name, want, err) + } + } +} + +func TestAnOllamaModelIsImportedIntoTheMirrorWithItsChecksum(t *testing.T) { + f := newOllamaFixture(t) + weights := f.llama32() + root := t.TempDir() + m, err := ReadOllama(f.dir, mustName(t, "llama3.2:3b")) + if err != nil { + t.Fatal(err) + } + + var reported []string + added, err := ImportOllama(t.Context(), root, f.dir, m, "", "", func(file string, _, _ int64) { + reported = append(reported, file) + }) + if err != nil { + t.Fatal(err) + } + // The tag alone moves in Ollama ("latest" is whatever was pulled last); + // the start of the digest is what pins the version. + sum := strings.TrimPrefix(weights.Digest, "sha256:") + want := Entry{Model: "llama3.2", Version: "3b-" + sum[:12], File: "llama3.2-3b.gguf", Bytes: weights.Size} + if len(added) != 1 || added[0].Entry != want { + t.Fatalf("added %+v, want %+v", added, want) + } + if added[0].SHA256 != sum { + t.Errorf("sha256 %s, want the digest Ollama named it by, %s", added[0].SHA256, sum) + } + if len(reported) == 0 || reported[0] != "llama3.2-3b.gguf" { + t.Errorf("progress said %v", reported) + } + + // What a node is given: the file, with the checksum it checks it against. + mf, err := ManifestFor(root, want.Model, want.Version) + if err != nil { + t.Fatal(err) + } + if len(mf.Files) != 1 || mf.Files[0].SHA256 != sum { + t.Errorf("manifest %+v", mf) + } + b, _ := os.ReadFile(want.Path(root)) + if string(b) != string(gguf("llama 3.2 3b weights")) { + t.Error("the mirrored file is not the blob") + } + if got := readSource(want.Path(root)); got != "ollama://registry.ollama.ai/library/llama3.2:3b" { + t.Errorf("source recorded as %q", got) + } + + // The same model again is the same version, and the mirror never + // replaces one. + if _, err := ImportOllama(t.Context(), root, f.dir, m, "", "", nil); !errors.Is(err, ErrExists) { + t.Errorf("a second import was not refused: %v", err) + } + // Under names the operator chose, it goes where they say. + added, err = ImportOllama(t.Context(), root, f.dir, m, "llama-3.2-3b-instruct", "2026-10", nil) + if err != nil || added[0].Entry.Path(root) != filepath.Join(root, "llama-3.2-3b-instruct", "2026-10", "llama3.2-3b.gguf") { + t.Errorf("named import: %+v, %v", added, err) + } +} + +func TestAVisionModelsProjectorGoesBesideItForLlamaCpp(t *testing.T) { + f := newOllamaFixture(t) + f.manifest("llava:7b", + f.blob(ollamaModelLayer, gguf("llava weights")), + f.blob(ollamaProjectorLayer, gguf("clip projector")), + ) + root := t.TempDir() + m, err := ReadOllama(f.dir, mustName(t, "llava:7b")) + if err != nil { + t.Fatal(err) + } + added, err := ImportOllama(t.Context(), root, f.dir, m, "", "v1", nil) + if err != nil { + t.Fatal(err) + } + if len(added) != 2 || added[1].Entry.File != "mmproj-llava-7b.gguf" { + t.Fatalf("added %+v", added) + } + + // The version directory is the model's source, and llama.cpp's engine + // finds in it the model to load and the projector to load beside it. + dir := filepath.Join(root, "llava", "v1") + s := engines.ShardSpec{HTTPPort: 20000, Profile: engines.Profile{Name: "cpu", Engine: "llamacpp", ModelPath: dir}} + cmds, err := (&engines.LlamaCpp{}).Commands(s) + if err != nil { + t.Fatal(err) + } + args := cmds[0].Args + if got := args[slices.Index(args, "--model")+1]; got != filepath.Join(dir, "llava-7b.gguf") { + t.Errorf("--model %s", got) + } + if got := args[slices.Index(args, "--mmproj")+1]; got != filepath.Join(dir, "mmproj-llava-7b.gguf") { + t.Errorf("--mmproj %s", got) + } +} + +func TestADamagedOllamaBlobIsNotMirrored(t *testing.T) { + f := newOllamaFixture(t) + f.manifest("llava:7b", + f.blob(ollamaModelLayer, gguf("llava weights")), + f.blob(ollamaProjectorLayer, gguf("clip projector")), + ) + m, err := ReadOllama(f.dir, mustName(t, "llava:7b")) + if err != nil { + t.Fatal(err) + } + // A bit gone wrong on Ollama's disk: the projector no longer hashes to + // the name it was stored under. + path, _, _ := ollamaBlob(f.dir, m.Projectors[0].Digest) + if err := os.WriteFile(path, gguf("clip projectoR"), 0o644); err != nil { + t.Fatal(err) + } + + root := t.TempDir() + _, err = ImportOllama(t.Context(), root, f.dir, m, "", "v1", nil) + if err == nil || !strings.Contains(err.Error(), "damaged") { + t.Fatalf("expected the damage to be found, got %v", err) + } + // All or nothing: the weights copied before it are taken out again, so + // no node is offered a vision model without its projector. + if entries, _ := List(root); len(entries) != 0 { + t.Errorf("left in the mirror: %+v", entries) + } + if _, err := os.Stat(filepath.Join(root, "llava")); !errors.Is(err, os.ErrNotExist) { + t.Errorf("the version's directory was left behind: %v", err) + } + + // A blob gone altogether says which. + _ = os.Remove(path) + if _, err := ImportOllama(t.Context(), root, f.dir, m, "", "v1", nil); err == nil || + !strings.Contains(err.Error(), m.Projectors[0].Digest) { + t.Errorf("a missing blob: %v", err) + } +} + +func TestOnlyAGGUFIsImportedAsOne(t *testing.T) { + // An HTML error page saved as a .gguf would reach every node otherwise. + src := filepath.Join(t.TempDir(), "model.gguf") + _ = os.WriteFile(src, []byte("rate limited"), 0o644) + _, err := Import(t.Context(), t.TempDir(), src, Entry{Model: "m", Version: "v", File: "model.gguf"}, "", "file://"+src, nil) + if err == nil || !strings.Contains(err.Error(), "not a GGUF") { + t.Errorf("expected a refusal, got %v", err) + } +} + +func TestALocalFileIsImportedWithItsChecksum(t *testing.T) { + src := filepath.Join(t.TempDir(), "q4_k_m.gguf") + content := gguf("converted by hand") + _ = os.WriteFile(src, content, 0o644) + root := t.TempDir() + e := Entry{Model: "qwen", Version: "2026-10-04", File: "q4_k_m.gguf"} + + a, err := Import(t.Context(), root, src, e, "", "file://"+src, nil) + if err != nil { + t.Fatal(err) + } + sum := sha256.Sum256(content) + if a.SHA256 != hex.EncodeToString(sum[:]) || readChecksum(e.Path(root)) != a.SHA256 { + t.Errorf("checksum %s, recorded %s", a.SHA256, readChecksum(e.Path(root))) + } + // A checksum the operator was given is checked too. + e.Version = "again" + if _, err := Import(t.Context(), root, src, e, strings.Repeat("0", 64), "file://"+src, nil); err == nil { + t.Error("a file that does not hash to the sum given was imported") + } + if _, err := os.Stat(e.Path(root) + partSuffix); !errors.Is(err, os.ErrNotExist) { + t.Error("a refused import left its partial file") + } +} + +func mustName(t *testing.T, s string) OllamaName { + t.Helper() + n, err := ParseOllamaName(s) + if err != nil { + t.Fatal(err) + } + return n +} diff --git a/internal/router/scrape/parse.go b/internal/router/scrape/parse.go index 379bba8..ec35a65 100644 --- a/internal/router/scrape/parse.go +++ b/internal/router/scrape/parse.go @@ -28,6 +28,11 @@ type Mapping struct { PromptTokens string GenerationTokens string SuccessTotal string + + // Path is where the engine serves these metrics, when it is not the + // poller's configured path: trtllm-serve's /metrics is JSON, and its + // Prometheus text is under /prometheus/metrics. + Path string } // MappingFor returns the metric names for an engine. @@ -41,6 +46,8 @@ func MappingFor(engine string) Mapping { switch engine { case "llamacpp": return LlamaCppMapping() + case "trtllm": + return TRTLLMMapping() default: return VLLMMapping() } @@ -72,6 +79,33 @@ func LlamaCppMapping() Mapping { } } +// TRTLLMMapping returns TensorRT-LLM's metric names, as trtllm-serve +// publishes them under /prometheus/metrics. +// +// They are there only when the server runs with return_perf_metrics, and +// running and waiting move only with enable_iter_perf_stats; llmgw's launcher +// turns both on. The gauges are set from each iteration's statistics, so +// they say what the engine was doing when it last iterated: an engine that +// has gone idle keeps showing its last batch until the next request moves it. +// Running and waiting first appeared in the 1.3 release candidates; an older +// server serves none of the three and is reported as breaking the contract, +// which is what it does. +// +// KV cache utilization is used blocks over the pool, a fraction like vLLM's. +// The cache's configuration is published under other label names than the +// block counts read from vLLM, so it is left unread. +func TRTLLMMapping() Mapping { + return Mapping{ + Running: "trtllm_num_requests_running", + Waiting: "trtllm_num_requests_waiting", + KVCacheUsage: "trtllm_kv_cache_utilization", + PromptTokens: "trtllm_prompt_tokens_total", + GenerationTokens: "trtllm_generation_tokens_total", + SuccessTotal: "trtllm_request_success_total", + Path: "/prometheus/metrics", + } +} + // VLLMMapping returns vLLM's current metric names. func VLLMMapping() Mapping { return Mapping{ diff --git a/internal/router/scrape/poller.go b/internal/router/scrape/poller.go index 096c023..e0fb2f2 100644 --- a/internal/router/scrape/poller.go +++ b/internal/router/scrape/poller.go @@ -208,11 +208,16 @@ func (p *Poller) fetch(ctx context.Context, base, apiKey string, mapping Mapping ctx, cancel := context.WithTimeout(ctx, p.cfg.Timeout) defer cancel() - req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+p.cfg.Path, nil) + path := p.cfg.Path + if mapping.Path != "" { + path = mapping.Path + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+path, nil) if err != nil { return endpoint.Stats{}, err } - // llama.cpp guards /metrics with its key; vLLM does not, and ignores one. + // llama.cpp guards /metrics with its key, and llmgw's TensorRT-LLM + // launcher guards everything but /health; vLLM does not, and ignores one. if apiKey != "" { req.Header.Set("Authorization", "Bearer "+apiKey) } diff --git a/internal/router/scrape/trtllm_test.go b/internal/router/scrape/trtllm_test.go new file mode 100644 index 0000000..deb44b4 --- /dev/null +++ b/internal/router/scrape/trtllm_test.go @@ -0,0 +1,133 @@ +package scrape + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +// trtllmMetrics is what trtllm-serve's /prometheus/metrics serves, in the +// shape prometheus_client's multiprocess collector writes it: every family +// labelled with the model and the backend, gauges with the pid of the +// process that set them, counters named with their _total. Trimmed to the +// families that matter and one histogram, which the parser has to step over. +const trtllmMetrics = `# HELP trtllm_request_success_total Count of successfully processed requests. +# TYPE trtllm_request_success_total counter +trtllm_request_success_total{engine_type="pytorch",finished_reason="stop",model_name="llama-3.1-70b"} 40.0 +trtllm_request_success_total{engine_type="pytorch",finished_reason="length",model_name="llama-3.1-70b"} 2.0 +# HELP trtllm_e2e_request_latency_seconds Histogram of end to end request latency in seconds. +# TYPE trtllm_e2e_request_latency_seconds histogram +trtllm_e2e_request_latency_seconds_bucket{engine_type="pytorch",le="0.5",model_name="llama-3.1-70b"} 10.0 +trtllm_e2e_request_latency_seconds_bucket{engine_type="pytorch",le="+Inf",model_name="llama-3.1-70b"} 42.0 +trtllm_e2e_request_latency_seconds_count{engine_type="pytorch",model_name="llama-3.1-70b"} 42.0 +trtllm_e2e_request_latency_seconds_sum{engine_type="pytorch",model_name="llama-3.1-70b"} 61.5 +# HELP trtllm_prompt_tokens_total Number of prefill tokens processed. +# TYPE trtllm_prompt_tokens_total counter +trtllm_prompt_tokens_total{engine_type="pytorch",model_name="llama-3.1-70b"} 18234.0 +# HELP trtllm_generation_tokens_total Number of generation tokens processed. +# TYPE trtllm_generation_tokens_total counter +trtllm_generation_tokens_total{engine_type="pytorch",model_name="llama-3.1-70b"} 9120.0 +# HELP trtllm_kv_cache_utilization KV cache utilization +# TYPE trtllm_kv_cache_utilization gauge +trtllm_kv_cache_utilization{engine_type="pytorch",model_name="llama-3.1-70b",pid="4121"} 0.375 +# HELP trtllm_kv_cache_config_info KV cache configuration info +# TYPE trtllm_kv_cache_config_info gauge +trtllm_kv_cache_config_info{enable_block_reuse="True",engine_type="pytorch",model_name="llama-3.1-70b",page_size="32",pid="4121"} 1.0 +# HELP trtllm_num_requests_running Number of active requests +# TYPE trtllm_num_requests_running gauge +trtllm_num_requests_running{engine_type="pytorch",model_name="llama-3.1-70b",pid="4121"} 5.0 +# HELP trtllm_num_requests_waiting Number of queued requests +# TYPE trtllm_num_requests_waiting gauge +trtllm_num_requests_waiting{engine_type="pytorch",model_name="llama-3.1-70b",pid="4121"} 2.0 +` + +func TestTRTLLMMetricsSatisfyTheContract(t *testing.T) { + s, err := Parse(strings.NewReader(trtllmMetrics), MappingFor("trtllm")) + if err != nil { + t.Fatal(err) + } + if s.Running != 5 || s.Waiting != 2 { + t.Errorf("running=%d waiting=%d", s.Running, s.Waiting) + } + // A fraction of the pool, as vLLM's is: the scorer compares the two. + if s.KVCacheUsage != 0.375 { + t.Errorf("kv=%v", s.KVCacheUsage) + } + if s.PromptTokens != 18234 || s.GenerationTokens != 9120 { + t.Errorf("tokens=%v/%v", s.PromptTokens, s.GenerationTokens) + } + // Split by finished_reason, summed. + if s.SuccessTotal != 42 { + t.Errorf("success=%v", s.SuccessTotal) + } + // Its cache configuration carries no block counts to read. + if s.BlockSize != 0 || s.NumGPUBlocks != 0 { + t.Errorf("block size %d and blocks %d read from labels that do not hold them", s.BlockSize, s.NumGPUBlocks) + } +} + +func TestATRTLLMWithoutItsQueueGaugesIsReportedNotIdle(t *testing.T) { + // TensorRT-LLM 1.2 publishes the cache's utilization and none of the queue + // gauges. Read as zeros, it would look like the emptiest endpoint in the + // cluster and take every request. + only := `# TYPE trtllm_kv_cache_utilization gauge +trtllm_kv_cache_utilization{engine_type="pytorch",model_name="m"} 0.1 +` + var missing *MissingError + _, err := Parse(strings.NewReader(only), TRTLLMMapping()) + if !errors.As(err, &missing) || len(missing.Missing) != 2 { + t.Fatalf("expected running and waiting reported missing, got %v", err) + } +} + +func TestTRTLLMIsScrapedWhereItPublishesPrometheusText(t *testing.T) { + // trtllm-serve's /metrics answers JSON — iteration statistics, drained by + // each read — and the launcher puts every path but /health behind the key. + var gotKey atomic.Value + gotKey.Store("") + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/prometheus/metrics": + gotKey.Store(r.Header.Get("Authorization")) + if r.Header.Get("Authorization") != "Bearer k1" { + http.Error(w, "a valid API key is required", http.StatusUnauthorized) + return + } + _, _ = w.Write([]byte(trtllmMetrics)) + case "/metrics": + _, _ = w.Write([]byte(`[{"iter": 154, "numActiveRequests": 1}]`)) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + + reg := endpoint.NewRegistry(endpoint.Config{}) + reg.Replace([]endpoint.Info{{ID: "t1", URL: srv.URL, Engine: "trtllm", APIKey: "k1", Models: []string{"m"}}}) + p := NewPoller(reg, Config{Interval: 20 * time.Millisecond, Timeout: 2 * time.Second}, nil) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + p.Sync(ctx) + defer p.Stop() + + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + for _, snap := range reg.Snapshots("m") { + if snap.ID == "t1" && snap.Stats.Running == 5 { + if snap.Stats.Waiting != 2 || snap.Stats.KVCacheUsage != 0.375 { + t.Errorf("stats read wrong: %+v", snap.Stats) + } + return + } + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("the endpoint was never measured (key sent: %q)", gotKey.Load()) +} diff --git a/internal/web/admin_serve.go b/internal/web/admin_serve.go index 6933240..9476fec 100644 --- a/internal/web/admin_serve.go +++ b/internal/web/admin_serve.go @@ -545,7 +545,7 @@ func describeServe(m catalog.Model) string { return "" } p := m.Profiles[0] - engine := map[string]string{"llamacpp": "llama.cpp", "vllm": "vLLM", "sglang": "SGLang"}[p.Engine] + engine := map[string]string{"llamacpp": "llama.cpp", "vllm": "vLLM", "sglang": "SGLang", "trtllm": "TensorRT-LLM"}[p.Engine] if engine == "" { engine = p.Engine } diff --git a/notes/c23.md b/notes/c23.md new file mode 100644 index 0000000..01db2bd --- /dev/null +++ b/notes/c23.md @@ -0,0 +1,61 @@ +## CHANGELOG + +### Engines and the mirror + +- **TensorRT-LLM as an engine** (`engine: trtllm`), on nodes where you + installed it: name its interpreter in the node's `agent.yaml` + (`engines.trtllm.python`) and the node offers it while that interpreter + has the `tensorrt_llm` package. llmgw does not install it. The agent starts + `trtllm-serve serve` with the replica's host, port, `--tp_size`, + `--pp_size` and served name, through a few lines of its own that make it + demand the gateway's key on every path but `/health` (trtllm-serve has no + key option), turn on the Prometheus endpoint and iteration statistics, and + keep usage reports to NVIDIA off. The router reads it at + `/prometheus/metrics` — `trtllm_num_requests_running`, + `trtllm_num_requests_waiting`, `trtllm_kv_cache_utilization`, tokens and + successes. Needs TensorRT-LLM 1.3 (rc13 or later) and an Ampere or newer + GPU; one node per replica. Profiles may not give it an options file, + `--trust_remote_code`, `--grpc`, `--num_serve_frontends` or the flags the + agent sets; `--set PATH=VALUE` changes a setting instead. `llmgw ctl + engine-check trtllm --bin ` checks an installed one. Written + against TensorRT-LLM's documented command line, server and metrics and + tested with stand-ins for them; not yet run on a datacenter GPU. +- **`llmgw mirror import-ollama [NAME[:TAG]]`** puts a model you already + have in Ollama into the mirror for llama.cpp: its GGUF, and a vision + model's projector beside it, each checked against the sha256 Ollama named + its file by. With no name it lists Ollama's models; it finds Ollama's + directory itself (`OLLAMA_MODELS`, `~/.ollama/models`, Ollama's Linux + service) or takes `--ollama-dir`. The version is the tag and the start of + the digest, since a tag moves in Ollama. It prints a model definition to + add, with Ollama's `num_ctx` as `--ctx-size`. The console's **Mirror** + view does the same with **I**, for the mirror beside the control plane. +- **`llmgw mirror import FILE... --model NAME`** copies files already on the + mirror machine into it the way a download is kept: whole or not at all, the + checksum beside it, a version never replaced, a `.gguf` checked to be one, + and `--sha256` checked when given. + +## BACKLOG + +**Done (branch `features-c2`):** engine `trtllm` (`internal/agent/engines/trtllm.go`): +`trtllm-serve` through a launcher of llmgw's that demands the replica's key +(trtllm-serve has none) and writes the options file that turns on +`/prometheus/metrics` and the iteration statistics; the scrape contract for +it (`scrape.TRTLLMMapping`, with a per-engine metrics path); found, not +installed — `engines.trtllm.python` in `agent.yaml`, probed from the +package's metadata; catalog validation (GPU class, one node, no GGUF source, +the flags a profile may not give); `ctl engine-check trtllm`, which asks +`trtllm-serve serve --help` through its module. Ollama/GGUF import: `mirror +import-ollama` (manifest parsing, blob digest check, projector beside the +model, all or nothing, a model definition printed), `mirror import` for any +local file, and the console's **I** on the Mirror view. Left: TensorRT-LLM has +not run on real hardware — the first node with one should run `engine-check` +and serve a model before anyone relies on it; no installer for it; no +multi-node TensorRT-LLM replica (it needs an MPI or Slurm launch from one +node); Triton Inference Server as an engine is not done; trtllm's queue gauges +are set per iteration, so an idle engine shows its last batch until the next +request; `--num_serve_frontends` is refused because extra front ends would not +demand the key. The Ollama import runs on the mirror machine only — the +mirror's API takes no files, so an Ollama elsewhere means copying the blob or +running the import there; Ollama models kept as separate tensors or as several +weight layers, and its LoRA adapters, are refused or left with a sentence; +files an import run with `sudo` writes are root's, as `mirror add`'s are.