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:
1 parent
30f7f3474a
commit
7940e8593c
28 files changed
+2376
-18
No files matched your search
+4
-1
@@ -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
@@ -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
|
||||
|
||||
@@ -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`
|
||||
|
||||
@@ -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 ---
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
`
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 "".
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -44,6 +44,8 @@ func newMirrorCmd() *cobra.Command {
|
||||
|
||||
c.AddCommand(
|
||||
newMirrorAddCmd(&dir),
|
||||
newMirrorImportCmd(&dir),
|
||||
newMirrorImportOllamaCmd(&dir),
|
||||
newMirrorListCmd(&dir),
|
||||
newMirrorUnusedCmd(&dir),
|
||||
newMirrorRemoveCmd(&dir),
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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{
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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{
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
Reference in new issue
Block a user