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) <noreply@anthropic.com>
This commit is contained in:
Ilker ManapandClaude Opus 5.5 committed 2026-10-05 00:02:01 +02:00
1 parent 30f7f3474a
commit 7940e8593c
28 files changed
+2376 -18

No files matched your search

+4 -1
View File
@@ -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.
+119 -2
View File
@@ -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 <the checksum its publisher gave>
```
#### 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 <model> --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
+26
View File
@@ -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: <venv>/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`
+12 -1
View File
@@ -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 ---
+2 -2
View File
@@ -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)
}
}
+218
View File
@@ -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: <venv>/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()
`
+316
View File
@@ -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)
}
}
+27
View File
@@ -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 "".
+1 -1
View File
@@ -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")
+15 -6
View File
@@ -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
}
+77
View File
@@ -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)
}
}
+2
View File
@@ -44,6 +44,8 @@ func newMirrorCmd() *cobra.Command {
c.AddCommand(
newMirrorAddCmd(&dir),
newMirrorImportCmd(&dir),
newMirrorImportOllamaCmd(&dir),
newMirrorListCmd(&dir),
newMirrorUnusedCmd(&dir),
newMirrorRemoveCmd(&dir),
+216
View File
@@ -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://<this-host>:8081"
func newMirrorImportCmd(dir *string) *cobra.Command {
var model, version, as, sum, base string
c := &cobra.Command{
Use: "import <file>... --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))
}
}
+144
View File
@@ -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://<this-host>: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")
}
}
+4
View File
@@ -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")
+34 -2
View File
@@ -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
}
+13
View File
@@ -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) {
+53
View File
@@ -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)
}
}
+4
View File
@@ -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: <venv>/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{
+1
View File
@@ -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",
+131
View File
@@ -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)
}
+352
View File
@@ -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/<host>/<namespace>/<model>/<tag> under its models directory —
// registry.ollama.ai/library/llama3.2/3b for "llama3.2:3b" — and each layer
// is a file blobs/sha256-<hex>, 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)
}
+369
View File
@@ -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-<hex> named by their content's digest, and a manifest per
// model and tag under manifests/<host>/<namespace>/<model>/<tag>.
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("<html>rate limited</html>"), 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
}
+34
View File
@@ -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{
+7 -2
View File
@@ -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)
}
+133
View File
@@ -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())
}
+1 -1
View File
@@ -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
}
+61
View File
@@ -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 <python>` 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.