Compare commits
No files matched your search
+179
@@ -11,11 +11,190 @@ is only released once its **exit criterion** is met, not when its code lands.
|
||||
| v0.4.0 | 2 | Agent: GPU/IB discovery, single-node engines | released as part of v0.5.0 |
|
||||
| v0.5.0 | 3 | Control plane: catalog, placement, reconciler | released |
|
||||
| v0.6.0 | 4 | Multi-node replicas and fabric awareness | released |
|
||||
| v0.7.0 | — | Conversations, agent memory, disk management | released |
|
||||
| v0.8.0 | — | PostgreSQL, leader election, agent failover | released |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
|
||||
Until v1.0.0 the minor version is the breaking one: configuration formats and Go
|
||||
APIs change between minors without a deprecation period.
|
||||
|
||||
## v0.8.0 — 2026-09-27
|
||||
|
||||
The store stops being a file on one machine's disk. PostgreSQL, a way to move
|
||||
an existing installation onto it without reissuing every key, and the two
|
||||
things a second control plane needs so that running one is safe rather than
|
||||
clever.
|
||||
|
||||
### A second control plane, and a way to fail over to it
|
||||
|
||||
Leadership is a lease in the store, so two control planes against one database
|
||||
cannot both place work — which they otherwise would, both seeing a model
|
||||
wanting replicas and none running, leaving twice the replicas each holding GPUs
|
||||
the other believes are free.
|
||||
|
||||
A standby also refuses agent streams. Commands travel down the stream the agent
|
||||
opened, so a standby holding one is a node the leader cannot command, and the
|
||||
leader's placements would fail with "no agent connected" for a node that is
|
||||
plainly connected. Agents therefore take a list of addresses rather than one,
|
||||
trying each in turn: a standby refusing, a stopped control plane refusing the
|
||||
connection and a crashed one never answering are the same thing from a node.
|
||||
|
||||
Measured on the cluster. Leadership moved three seconds after the leader was
|
||||
stopped — it released the lease on the way out rather than letting it lapse —
|
||||
and all five agents had registered with the standby six seconds later. The
|
||||
standby adopted every replica rather than restarting any: six replicas, all
|
||||
ready, engines that never stopped serving.
|
||||
|
||||
Everything is unset-means-leader, so a single control plane and every
|
||||
deployment that existed before this behave exactly as they did.
|
||||
|
||||
What this does not do is move client traffic. The standby served inference
|
||||
correctly on its own port while the public address returned 502, because
|
||||
nothing in front of them moved — that needs a load balancer or a shared
|
||||
address, and it is the remaining piece.
|
||||
|
||||
### PostgreSQL
|
||||
|
||||
The store runs on PostgreSQL, chosen by what the address looks like: a path is
|
||||
still SQLite, a `postgres://` URL is not, and `@` followed by a path reads the
|
||||
address out of a file so a DSN carrying a password stays out of a configuration
|
||||
that gets copied and pasted into issues. Every existing installation keeps
|
||||
working untouched.
|
||||
|
||||
One set of queries rather than one per database, because the SQL here is almost
|
||||
entirely portable and two copies would drift one query at a time. The rewrite
|
||||
of placeholders is a wrapper around the handle rather than a change at sixty
|
||||
call sites — which found six places reaching past it, including all three
|
||||
manual transactions and two prepared statements, every one of which worked on
|
||||
SQLite and would have failed on PostgreSQL.
|
||||
|
||||
Exactly one query is not the same sentence in both: the note search, MATCH
|
||||
against `@@`. SQLite gets full-text search from an FTS5 table kept in step by
|
||||
three triggers; PostgreSQL derives a tsvector from the row, so there is nothing
|
||||
to keep in step and no trigger to forget. It uses the 'simple' configuration
|
||||
rather than English — these are an agent's own notes, in whatever language it
|
||||
chose, and stemming English rules over Turkish would quietly stop finding
|
||||
things.
|
||||
|
||||
`llmgw keys move-store` copies an existing store rather than starting a fresh
|
||||
one, because keys are stored hashed: a migration that drops them cannot
|
||||
recreate them, and every person and every script would need a new one.
|
||||
|
||||
Testing it against a real database found three things nothing else would have.
|
||||
`schema_migrations` is created in Go rather than in a migration file, so it
|
||||
escaped the port and still said INTEGER — 64 bits in SQLite, 32 in PostgreSQL,
|
||||
and a millisecond timestamp does not fit. The copy order followed the order the
|
||||
migrations create tables in, and `users` is created after `api_keys` even
|
||||
though `api_keys` references it. And `key_claims.sealed` is emptied the moment
|
||||
a claim is used, so a claimed row holds an empty blob in a NOT NULL column — an
|
||||
empty blob scans as a typed nil slice and a real NULL as an untyped one, a
|
||||
difference pgx cannot see, so it arrived as NULL.
|
||||
|
||||
The driver is pgx, which is pure Go: the whole binary is built with
|
||||
CGO_ENABLED=0 so that it runs on a node whose libc is older than the builder's,
|
||||
and a cgo driver would take that away from every deployment, including the ones
|
||||
staying on SQLite.
|
||||
|
||||
## v0.7.0 — 2026-09-25
|
||||
|
||||
What a cluster needs to be lived with rather than demonstrated: somewhere to
|
||||
keep a conversation, somewhere to keep a note, and an honest answer about the
|
||||
disk underneath both.
|
||||
|
||||
### Memory notes for agents, and threads without API keys
|
||||
|
||||
Notes an agent writes for itself and reads back in later conversations, offered
|
||||
as REST (`/v1/memory`) and as MCP tools at `/mcp` (`remember`, `recall`,
|
||||
`forget`). A note is a key, text and tags, searched by words or tag, private to
|
||||
a project and to the agent id in `x-llmgw-agent`. llmgw does not inject the
|
||||
tools into chat requests: an agent that is not given the server never sees
|
||||
memory.
|
||||
|
||||
Off for every project until someone runs `llmgw keys memory --retain-hours N`,
|
||||
and the retention is required, as it is for threads.
|
||||
|
||||
On a gateway with no authentication, a `threads:` block (`store`,
|
||||
`retain_hours`) now keeps conversations and notes in one shared pool. Before,
|
||||
`x-llmgw-thread` was silently ignored there; now a gateway configured to keep
|
||||
nothing refuses it with an explanation instead. The shared pool is for one
|
||||
trusted machine: anyone who can reach the port and knows an id can read it.
|
||||
|
||||
### Conversations that can move between models
|
||||
|
||||
A thread a client names with `x-llmgw-thread`, carrying only its new turn.
|
||||
llmgw replays the stored history to whichever model is serving, which is what
|
||||
lets a conversation begun on a large model continue on a small one — including
|
||||
when a spent budget moved it there and the caller never chose.
|
||||
|
||||
Off for every project until someone runs `llmgw keys threads`. Until then llmgw
|
||||
stores no message content anywhere, as it never has, and that default is worth
|
||||
more than the feature. The retention is required rather than optional, because
|
||||
a conversation store with no expiry is a transcript of everything a team ever
|
||||
typed kept forever by omission, and it is measured from the last turn so a
|
||||
conversation in progress does not expire underneath it.
|
||||
|
||||
History that does not fit the serving model is trimmed oldest-first and the
|
||||
reply says how many turns went, in `x-llmgw-thread-trimmed`. System messages
|
||||
are never dropped — losing one silently changes what the model is for while it
|
||||
carries on answering — and neither is the turn the caller just sent.
|
||||
|
||||
Threads are scoped to a project, and reading, writing or deleting one that
|
||||
belongs to another answers exactly as one that never existed: two answers would
|
||||
let a caller map which ids are real in projects they cannot read.
|
||||
|
||||
Capturing the reply is opt-in and bounded, so a request without a thread
|
||||
accumulates nothing on the hottest path in the system, and one with a thread
|
||||
stops at a cap rather than growing to it.
|
||||
|
||||
`DELETE /v1/threads/{id}` forgets a conversation, as a client rather than as an
|
||||
operator: the person who said the thing should not have to ask someone to read
|
||||
it before it can go.
|
||||
|
||||
### Disk, and taking back only what llmgw put there
|
||||
|
||||
Disk was the one resource nothing reported. `statfs` was already being called
|
||||
on every model directory, but only to ask whether the filesystem was a network
|
||||
one — the free space in the same struct was discarded, so the only way to find
|
||||
out why a node was filling up was to ssh to it.
|
||||
|
||||
Nodes now report it, per filesystem rather than per configured directory: free
|
||||
space belongs to the disk, and two model directories on one disk are one pool
|
||||
of bytes. Shared filesystems are kept out of a node's own figure, because space
|
||||
on a filer belongs to everyone mounting it and twenty nodes would otherwise
|
||||
each report the same free terabyte. It shows in `ctl nodes`, the dashboard and
|
||||
Prometheus, and a node that reported nothing reads as unknown everywhere — an
|
||||
older agent must not look like a full disk.
|
||||
|
||||
A profile may declare `min_disk_gb`, and a node without the room is no longer a
|
||||
candidate, with a reason `explain-placement` prints. Skipped on a node that
|
||||
already holds the weights: it needs no room for them, and without that the node
|
||||
best placed to run a model is the one most likely to be rejected.
|
||||
|
||||
`ctl weights` lists model files llmgw downloaded for models no longer in the
|
||||
catalog, and `--collect` deletes them. The design is mostly about what it will
|
||||
not do. Only files llmgw fetched are candidates, and that is a calculation
|
||||
rather than a heuristic — the cache name is a hash of the source, so the wanted
|
||||
set is computable exactly. Anything an operator put on a node is left alone
|
||||
however unreferenced it looks. Each node checks every path again before
|
||||
unlinking it, refusing anything outside its cache directory or without the
|
||||
cache name, so a control plane that is wrong cannot talk a node into losing a
|
||||
model that took a week to copy.
|
||||
|
||||
### Fixes
|
||||
|
||||
- **A capability was lost at an indirection.** The controller is handed a
|
||||
forwarding dispatcher, and reaches past its interface by type assertion to
|
||||
find what lies beyond scheduling. Deleting weights was implemented on the
|
||||
agent service and not on the forwarder, so it was simply absent at runtime —
|
||||
with the whole suite green, because the dispatchers the tests use do not
|
||||
implement it either and its absence there is the expected case.
|
||||
- **`ctl tool-check` blamed the engine for what the model did.** It had one
|
||||
verdict for a failed attempt and told operators not to try another model.
|
||||
Measured: the engine renders the tool prompt correctly, the coder models do
|
||||
not use the `<tool_call>` wrapper it asks for, and `qwen2.5-3b-instruct` on
|
||||
the same engine returns a populated `tool_calls` array. A failed attempt is
|
||||
now two verdicts, told apart by whether the reply used the wrapper.
|
||||
|
||||
## v0.6.0 — 2026-09-18
|
||||
|
||||
Phase 4: a replica can be bigger than one machine, and the fabric it spans
|
||||
|
||||
@@ -30,6 +30,8 @@ Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md).
|
||||
| v0.4.0 | 2 | Agent: GPU/IB discovery, engine supervision | released |
|
||||
| v0.5.0 | 3 | Control plane: catalog, placement, reconciler, ctl | released |
|
||||
| v0.6.0 | 4 | Multi-node placement and fabric awareness | released |
|
||||
| v0.7.0 | — | Conversations, agent memory, disk management | released |
|
||||
| v0.8.0 | — | PostgreSQL, leader election, agent failover | released |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
|
||||
As of v0.6.0 a cluster runs unattended: agents report their hardware, the
|
||||
@@ -37,8 +39,14 @@ control plane places replicas — including ones larger than a single machine
|
||||
models load when someone asks for them and stop when nobody does, and keys,
|
||||
budgets and rate limits decide who may use any of it.
|
||||
|
||||
Before v1.0.0: routing profiles selectable per request, spillover to a fallback
|
||||
pool, tracing, and reference documentation. One thing is worth saying plainly:
|
||||
Before v1.0.0: tracing, reference documentation, and filtering on `ctl nodes`
|
||||
and `ctl replicas` for fleets too large to read a whole table of. Two items
|
||||
that were on this list are not gaps: a model falls through to its fallbacks
|
||||
when there is no capacity as well as when a budget is spent, and routing
|
||||
profiles are chosen by the key rather than by a request header on purpose —
|
||||
a header would let any client promote itself into the top priority band.
|
||||
|
||||
One thing is worth saying plainly:
|
||||
multi-node replicas are proven on a simulated cluster and on real hardware
|
||||
without a fabric. The InfiniBand, NVML and NCCL paths are written against the
|
||||
formats and failure modes those tools actually have, but they have not yet run
|
||||
|
||||
+301
-84
@@ -90,7 +90,7 @@ func (x ShardStatus_State) Number() protoreflect.EnumNumber {
|
||||
|
||||
// Deprecated: Use ShardStatus_State.Descriptor instead.
|
||||
func (ShardStatus_State) EnumDescriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{8, 0}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{9, 0}
|
||||
}
|
||||
|
||||
// AgentMessage is anything a node tells the control plane.
|
||||
@@ -201,6 +201,7 @@ type ControlMessage struct {
|
||||
// *ControlMessage_Stop
|
||||
// *ControlMessage_Drain
|
||||
// *ControlMessage_Ack
|
||||
// *ControlMessage_RemoveWeights
|
||||
Msg isControlMessage_Msg `protobuf_oneof:"msg"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
@@ -279,6 +280,15 @@ func (x *ControlMessage) GetAck() *Ack {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *ControlMessage) GetRemoveWeights() *RemoveWeights {
|
||||
if x != nil {
|
||||
if x, ok := x.Msg.(*ControlMessage_RemoveWeights); ok {
|
||||
return x.RemoveWeights
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type isControlMessage_Msg interface {
|
||||
isControlMessage_Msg()
|
||||
}
|
||||
@@ -299,6 +309,10 @@ type ControlMessage_Ack struct {
|
||||
Ack *Ack `protobuf:"bytes,4,opt,name=ack,proto3,oneof"`
|
||||
}
|
||||
|
||||
type ControlMessage_RemoveWeights struct {
|
||||
RemoveWeights *RemoveWeights `protobuf:"bytes,5,opt,name=remove_weights,json=removeWeights,proto3,oneof"`
|
||||
}
|
||||
|
||||
func (*ControlMessage_Start) isControlMessage_Msg() {}
|
||||
|
||||
func (*ControlMessage_Stop) isControlMessage_Msg() {}
|
||||
@@ -307,6 +321,8 @@ func (*ControlMessage_Drain) isControlMessage_Msg() {}
|
||||
|
||||
func (*ControlMessage_Ack) isControlMessage_Msg() {}
|
||||
|
||||
func (*ControlMessage_RemoveWeights) isControlMessage_Msg() {}
|
||||
|
||||
// Register is the first message on every stream. It carries the node's full
|
||||
// hardware inventory, so the control plane needs no separate discovery call and
|
||||
// a reconnecting agent automatically refreshes what the scheduler knows.
|
||||
@@ -356,7 +372,16 @@ type Register struct {
|
||||
// Empty means the node was not asked to look, which is not the same as having
|
||||
// nothing — an agent with no model directories configured behaves exactly as
|
||||
// every agent did before this existed, and constrains nothing.
|
||||
Weights []*LocalWeights `protobuf:"bytes,11,rep,name=weights,proto3" json:"weights,omitempty"`
|
||||
Weights []*LocalWeights `protobuf:"bytes,11,rep,name=weights,proto3" json:"weights,omitempty"`
|
||||
// storage is the disk behind those weights, one entry per filesystem.
|
||||
//
|
||||
// Reported so an operator can see a node filling up before a fetch dies at
|
||||
// ninety percent, and so placement can refuse a node that cannot hold a
|
||||
// model rather than discovering it during the download.
|
||||
//
|
||||
// Empty means the node was not asked to look, exactly as for weights, and
|
||||
// constrains nothing.
|
||||
Storage []*Storage `protobuf:"bytes,12,rep,name=storage,proto3" json:"storage,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
@@ -468,6 +493,91 @@ func (x *Register) GetWeights() []*LocalWeights {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *Register) GetStorage() []*Storage {
|
||||
if x != nil {
|
||||
return x.Storage
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Storage is one filesystem a node keeps weights on.
|
||||
//
|
||||
// Per filesystem rather than per configured directory: free space belongs to
|
||||
// the disk, and two model directories on one disk are one pool of bytes.
|
||||
type Storage struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
// path is one directory on this filesystem, so the operator recognises it.
|
||||
Path string `protobuf:"bytes,1,opt,name=path,proto3" json:"path,omitempty"`
|
||||
// free_bytes is what a non-root process can write — not the root reserve,
|
||||
// which is space the agent cannot have.
|
||||
FreeBytes int64 `protobuf:"varint,2,opt,name=free_bytes,json=freeBytes,proto3" json:"free_bytes,omitempty"`
|
||||
TotalBytes int64 `protobuf:"varint,3,opt,name=total_bytes,json=totalBytes,proto3" json:"total_bytes,omitempty"`
|
||||
// shared marks a filesystem the node did not fill. Free space on one is
|
||||
// shared with every other node using it, so two nodes each reporting 500 GB
|
||||
// free on the same filer have 500 GB between them and not a terabyte.
|
||||
Shared bool `protobuf:"varint,4,opt,name=shared,proto3" json:"shared,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *Storage) Reset() {
|
||||
*x = Storage{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[3]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *Storage) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*Storage) ProtoMessage() {}
|
||||
|
||||
func (x *Storage) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[3]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use Storage.ProtoReflect.Descriptor instead.
|
||||
func (*Storage) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{3}
|
||||
}
|
||||
|
||||
func (x *Storage) GetPath() string {
|
||||
if x != nil {
|
||||
return x.Path
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *Storage) GetFreeBytes() int64 {
|
||||
if x != nil {
|
||||
return x.FreeBytes
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *Storage) GetTotalBytes() int64 {
|
||||
if x != nil {
|
||||
return x.TotalBytes
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *Storage) GetShared() bool {
|
||||
if x != nil {
|
||||
return x.Shared
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// RunningShard describes an engine already running on a node.
|
||||
//
|
||||
// It carries enough for the control plane to rebuild the replica it belongs to,
|
||||
@@ -498,7 +608,7 @@ type RunningShard struct {
|
||||
|
||||
func (x *RunningShard) Reset() {
|
||||
*x = RunningShard{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[3]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[4]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -510,7 +620,7 @@ func (x *RunningShard) String() string {
|
||||
func (*RunningShard) ProtoMessage() {}
|
||||
|
||||
func (x *RunningShard) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[3]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[4]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -523,7 +633,7 @@ func (x *RunningShard) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use RunningShard.ProtoReflect.Descriptor instead.
|
||||
func (*RunningShard) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{3}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{4}
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetReplicaId() string {
|
||||
@@ -653,7 +763,7 @@ type GPU struct {
|
||||
|
||||
func (x *GPU) Reset() {
|
||||
*x = GPU{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[4]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[5]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -665,7 +775,7 @@ func (x *GPU) String() string {
|
||||
func (*GPU) ProtoMessage() {}
|
||||
|
||||
func (x *GPU) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[4]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[5]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -678,7 +788,7 @@ func (x *GPU) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use GPU.ProtoReflect.Descriptor instead.
|
||||
func (*GPU) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{4}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{5}
|
||||
}
|
||||
|
||||
func (x *GPU) GetIndex() int32 {
|
||||
@@ -800,7 +910,7 @@ type FabricAttachment struct {
|
||||
|
||||
func (x *FabricAttachment) Reset() {
|
||||
*x = FabricAttachment{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[5]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[6]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -812,7 +922,7 @@ func (x *FabricAttachment) String() string {
|
||||
func (*FabricAttachment) ProtoMessage() {}
|
||||
|
||||
func (x *FabricAttachment) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[5]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[6]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -825,7 +935,7 @@ func (x *FabricAttachment) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use FabricAttachment.ProtoReflect.Descriptor instead.
|
||||
func (*FabricAttachment) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{5}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{6}
|
||||
}
|
||||
|
||||
func (x *FabricAttachment) GetFabric() string {
|
||||
@@ -867,7 +977,7 @@ type LocalWeights struct {
|
||||
|
||||
func (x *LocalWeights) Reset() {
|
||||
*x = LocalWeights{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[6]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[7]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -879,7 +989,7 @@ func (x *LocalWeights) String() string {
|
||||
func (*LocalWeights) ProtoMessage() {}
|
||||
|
||||
func (x *LocalWeights) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[6]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[7]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -892,7 +1002,7 @@ func (x *LocalWeights) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use LocalWeights.ProtoReflect.Descriptor instead.
|
||||
func (*LocalWeights) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{6}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{7}
|
||||
}
|
||||
|
||||
func (x *LocalWeights) GetPath() string {
|
||||
@@ -923,13 +1033,27 @@ type Heartbeat struct {
|
||||
// a card after registration.
|
||||
GpuMemoryUsedMb []int32 `protobuf:"varint,2,rep,packed,name=gpu_memory_used_mb,json=gpuMemoryUsedMb,proto3" json:"gpu_memory_used_mb,omitempty"`
|
||||
NewProblems []string `protobuf:"bytes,3,rep,name=new_problems,json=newProblems,proto3" json:"new_problems,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
// storage for the same reason the GPU memory is here: it changes while the
|
||||
// agent runs, and a figure sampled once at registration is a figure that is
|
||||
// wrong on exactly the node worth watching — the one busy filling up.
|
||||
//
|
||||
// Always sent, because it is a few numbers per filesystem.
|
||||
Storage []*Storage `protobuf:"bytes,4,rep,name=storage,proto3" json:"storage,omitempty"`
|
||||
// weights only when they changed, because the list is long enough that
|
||||
// sending it every few seconds would be paying for nothing most of the time.
|
||||
//
|
||||
// weights_changed is what says the list is meant: without it an empty list
|
||||
// would be ambiguous between "this node now holds nothing" and "nothing to
|
||||
// report", and the first would quietly empty a node's inventory.
|
||||
WeightsChanged bool `protobuf:"varint,5,opt,name=weights_changed,json=weightsChanged,proto3" json:"weights_changed,omitempty"`
|
||||
Weights []*LocalWeights `protobuf:"bytes,6,rep,name=weights,proto3" json:"weights,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *Heartbeat) Reset() {
|
||||
*x = Heartbeat{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[7]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[8]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -941,7 +1065,7 @@ func (x *Heartbeat) String() string {
|
||||
func (*Heartbeat) ProtoMessage() {}
|
||||
|
||||
func (x *Heartbeat) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[7]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[8]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -954,7 +1078,7 @@ func (x *Heartbeat) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use Heartbeat.ProtoReflect.Descriptor instead.
|
||||
func (*Heartbeat) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{7}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{8}
|
||||
}
|
||||
|
||||
func (x *Heartbeat) GetNodeId() string {
|
||||
@@ -978,6 +1102,27 @@ func (x *Heartbeat) GetNewProblems() []string {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *Heartbeat) GetStorage() []*Storage {
|
||||
if x != nil {
|
||||
return x.Storage
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *Heartbeat) GetWeightsChanged() bool {
|
||||
if x != nil {
|
||||
return x.WeightsChanged
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (x *Heartbeat) GetWeights() []*LocalWeights {
|
||||
if x != nil {
|
||||
return x.Weights
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ShardStatus reports a shard's lifecycle transition.
|
||||
type ShardStatus struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
@@ -996,7 +1141,7 @@ type ShardStatus struct {
|
||||
|
||||
func (x *ShardStatus) Reset() {
|
||||
*x = ShardStatus{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[8]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[9]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1008,7 +1153,7 @@ func (x *ShardStatus) String() string {
|
||||
func (*ShardStatus) ProtoMessage() {}
|
||||
|
||||
func (x *ShardStatus) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[8]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[9]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1021,7 +1166,7 @@ func (x *ShardStatus) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use ShardStatus.ProtoReflect.Descriptor instead.
|
||||
func (*ShardStatus) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{8}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{9}
|
||||
}
|
||||
|
||||
func (x *ShardStatus) GetNodeId() string {
|
||||
@@ -1102,7 +1247,7 @@ type StartShard struct {
|
||||
|
||||
func (x *StartShard) Reset() {
|
||||
*x = StartShard{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[9]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[10]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1114,7 +1259,7 @@ func (x *StartShard) String() string {
|
||||
func (*StartShard) ProtoMessage() {}
|
||||
|
||||
func (x *StartShard) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[9]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[10]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1127,7 +1272,7 @@ func (x *StartShard) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use StartShard.ProtoReflect.Descriptor instead.
|
||||
func (*StartShard) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{9}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{10}
|
||||
}
|
||||
|
||||
func (x *StartShard) GetReplicaId() string {
|
||||
@@ -1247,7 +1392,7 @@ type EngineProfile struct {
|
||||
|
||||
func (x *EngineProfile) Reset() {
|
||||
*x = EngineProfile{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[10]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[11]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1259,7 +1404,7 @@ func (x *EngineProfile) String() string {
|
||||
func (*EngineProfile) ProtoMessage() {}
|
||||
|
||||
func (x *EngineProfile) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[10]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[11]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1272,7 +1417,7 @@ func (x *EngineProfile) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use EngineProfile.ProtoReflect.Descriptor instead.
|
||||
func (*EngineProfile) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{10}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{11}
|
||||
}
|
||||
|
||||
func (x *EngineProfile) GetName() string {
|
||||
@@ -1352,6 +1497,57 @@ func (x *EngineProfile) GetServedModelName() string {
|
||||
return ""
|
||||
}
|
||||
|
||||
// RemoveWeights deletes cached model files a node no longer needs.
|
||||
//
|
||||
// The paths are computed by the control plane from the catalog, but the agent
|
||||
// checks them again before unlinking anything: it deletes only files inside its
|
||||
// own cache directory whose names this system generated. A control plane that
|
||||
// is wrong, compromised or simply newer than the agent must not be able to
|
||||
// talk it into deleting a model an operator spent a week copying there.
|
||||
type RemoveWeights struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
Paths []string `protobuf:"bytes,1,rep,name=paths,proto3" json:"paths,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *RemoveWeights) Reset() {
|
||||
*x = RemoveWeights{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[12]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *RemoveWeights) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*RemoveWeights) ProtoMessage() {}
|
||||
|
||||
func (x *RemoveWeights) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[12]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
return ms
|
||||
}
|
||||
return mi.MessageOf(x)
|
||||
}
|
||||
|
||||
// Deprecated: Use RemoveWeights.ProtoReflect.Descriptor instead.
|
||||
func (*RemoveWeights) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{12}
|
||||
}
|
||||
|
||||
func (x *RemoveWeights) GetPaths() []string {
|
||||
if x != nil {
|
||||
return x.Paths
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// StopShard terminates a shard and frees its GPUs.
|
||||
type StopShard struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
@@ -1362,7 +1558,7 @@ type StopShard struct {
|
||||
|
||||
func (x *StopShard) Reset() {
|
||||
*x = StopShard{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[11]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[13]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1374,7 +1570,7 @@ func (x *StopShard) String() string {
|
||||
func (*StopShard) ProtoMessage() {}
|
||||
|
||||
func (x *StopShard) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[11]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[13]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1387,7 +1583,7 @@ func (x *StopShard) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use StopShard.ProtoReflect.Descriptor instead.
|
||||
func (*StopShard) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{11}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{13}
|
||||
}
|
||||
|
||||
func (x *StopShard) GetShardId() string {
|
||||
@@ -1407,7 +1603,7 @@ type DrainShard struct {
|
||||
|
||||
func (x *DrainShard) Reset() {
|
||||
*x = DrainShard{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[12]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[14]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1419,7 +1615,7 @@ func (x *DrainShard) String() string {
|
||||
func (*DrainShard) ProtoMessage() {}
|
||||
|
||||
func (x *DrainShard) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[12]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[14]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1432,7 +1628,7 @@ func (x *DrainShard) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use DrainShard.ProtoReflect.Descriptor instead.
|
||||
func (*DrainShard) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{12}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{14}
|
||||
}
|
||||
|
||||
func (x *DrainShard) GetShardId() string {
|
||||
@@ -1454,7 +1650,7 @@ type Ack struct {
|
||||
|
||||
func (x *Ack) Reset() {
|
||||
*x = Ack{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[13]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[15]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1466,7 +1662,7 @@ func (x *Ack) String() string {
|
||||
func (*Ack) ProtoMessage() {}
|
||||
|
||||
func (x *Ack) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[13]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[15]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1479,7 +1675,7 @@ func (x *Ack) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use Ack.ProtoReflect.Descriptor instead.
|
||||
func (*Ack) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{13}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{15}
|
||||
}
|
||||
|
||||
func (x *Ack) GetMessageId() string {
|
||||
@@ -1512,13 +1708,14 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\bregister\x18\x01 \x01(\v2\x12.llmgw.v1.RegisterH\x00R\bregister\x123\n" +
|
||||
"\theartbeat\x18\x02 \x01(\v2\x13.llmgw.v1.HeartbeatH\x00R\theartbeat\x12-\n" +
|
||||
"\x05shard\x18\x03 \x01(\v2\x15.llmgw.v1.ShardStatusH\x00R\x05shardB\x05\n" +
|
||||
"\x03msg\"\xc1\x01\n" +
|
||||
"\x03msg\"\x83\x02\n" +
|
||||
"\x0eControlMessage\x12,\n" +
|
||||
"\x05start\x18\x01 \x01(\v2\x14.llmgw.v1.StartShardH\x00R\x05start\x12)\n" +
|
||||
"\x04stop\x18\x02 \x01(\v2\x13.llmgw.v1.StopShardH\x00R\x04stop\x12,\n" +
|
||||
"\x05drain\x18\x03 \x01(\v2\x14.llmgw.v1.DrainShardH\x00R\x05drain\x12!\n" +
|
||||
"\x03ack\x18\x04 \x01(\v2\r.llmgw.v1.AckH\x00R\x03ackB\x05\n" +
|
||||
"\x03msg\"\xfa\x04\n" +
|
||||
"\x03ack\x18\x04 \x01(\v2\r.llmgw.v1.AckH\x00R\x03ack\x12@\n" +
|
||||
"\x0eremove_weights\x18\x05 \x01(\v2\x17.llmgw.v1.RemoveWeightsH\x00R\rremoveWeightsB\x05\n" +
|
||||
"\x03msg\"\xa7\x05\n" +
|
||||
"\bRegister\x12\x17\n" +
|
||||
"\anode_id\x18\x01 \x01(\tR\x06nodeId\x12\x1a\n" +
|
||||
"\bhostname\x18\x02 \x01(\tR\bhostname\x12\x1b\n" +
|
||||
@@ -1531,13 +1728,21 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\x0erunning_shards\x18\t \x03(\v2\x16.llmgw.v1.RunningShardR\rrunningShards\x12&\n" +
|
||||
"\x0fmemory_total_mb\x18\n" +
|
||||
" \x01(\x03R\rmemoryTotalMb\x120\n" +
|
||||
"\aweights\x18\v \x03(\v2\x16.llmgw.v1.LocalWeightsR\aweights\x1a9\n" +
|
||||
"\aweights\x18\v \x03(\v2\x16.llmgw.v1.LocalWeightsR\aweights\x12+\n" +
|
||||
"\astorage\x18\f \x03(\v2\x11.llmgw.v1.StorageR\astorage\x1a9\n" +
|
||||
"\vLabelsEntry\x12\x10\n" +
|
||||
"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\x1aA\n" +
|
||||
"\x13EngineVersionsEntry\x12\x10\n" +
|
||||
"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\x97\x03\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"u\n" +
|
||||
"\aStorage\x12\x12\n" +
|
||||
"\x04path\x18\x01 \x01(\tR\x04path\x12\x1d\n" +
|
||||
"\n" +
|
||||
"free_bytes\x18\x02 \x01(\x03R\tfreeBytes\x12\x1f\n" +
|
||||
"\vtotal_bytes\x18\x03 \x01(\x03R\n" +
|
||||
"totalBytes\x12\x16\n" +
|
||||
"\x06shared\x18\x04 \x01(\bR\x06shared\"\x97\x03\n" +
|
||||
"\fRunningShard\x12\x1d\n" +
|
||||
"\n" +
|
||||
"replica_id\x18\x01 \x01(\tR\treplicaId\x12\x19\n" +
|
||||
@@ -1585,11 +1790,14 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\x04path\x18\x01 \x01(\tR\x04path\x12\x1d\n" +
|
||||
"\n" +
|
||||
"size_bytes\x18\x02 \x01(\x03R\tsizeBytes\x12\x16\n" +
|
||||
"\x06shared\x18\x03 \x01(\bR\x06shared\"t\n" +
|
||||
"\x06shared\x18\x03 \x01(\bR\x06shared\"\xfc\x01\n" +
|
||||
"\tHeartbeat\x12\x17\n" +
|
||||
"\anode_id\x18\x01 \x01(\tR\x06nodeId\x12+\n" +
|
||||
"\x12gpu_memory_used_mb\x18\x02 \x03(\x05R\x0fgpuMemoryUsedMb\x12!\n" +
|
||||
"\fnew_problems\x18\x03 \x03(\tR\vnewProblems\"\x8c\x03\n" +
|
||||
"\fnew_problems\x18\x03 \x03(\tR\vnewProblems\x12+\n" +
|
||||
"\astorage\x18\x04 \x03(\v2\x11.llmgw.v1.StorageR\astorage\x12'\n" +
|
||||
"\x0fweights_changed\x18\x05 \x01(\bR\x0eweightsChanged\x120\n" +
|
||||
"\aweights\x18\x06 \x03(\v2\x16.llmgw.v1.LocalWeightsR\aweights\"\x8c\x03\n" +
|
||||
"\vShardStatus\x12\x17\n" +
|
||||
"\anode_id\x18\x01 \x01(\tR\x06nodeId\x12\x1d\n" +
|
||||
"\n" +
|
||||
@@ -1646,7 +1854,9 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\x11served_model_name\x18\v \x01(\tR\x0fservedModelName\x1a6\n" +
|
||||
"\bEnvEntry\x12\x10\n" +
|
||||
"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"&\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"%\n" +
|
||||
"\rRemoveWeights\x12\x14\n" +
|
||||
"\x05paths\x18\x01 \x03(\tR\x05paths\"&\n" +
|
||||
"\tStopShard\x12\x19\n" +
|
||||
"\bshard_id\x18\x01 \x01(\tR\ashardId\"'\n" +
|
||||
"\n" +
|
||||
@@ -1675,52 +1885,58 @@ func file_llmgw_v1_agent_proto_rawDescGZIP() []byte {
|
||||
}
|
||||
|
||||
var file_llmgw_v1_agent_proto_enumTypes = make([]protoimpl.EnumInfo, 1)
|
||||
var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 17)
|
||||
var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 19)
|
||||
var file_llmgw_v1_agent_proto_goTypes = []any{
|
||||
(ShardStatus_State)(0), // 0: llmgw.v1.ShardStatus.State
|
||||
(*AgentMessage)(nil), // 1: llmgw.v1.AgentMessage
|
||||
(*ControlMessage)(nil), // 2: llmgw.v1.ControlMessage
|
||||
(*Register)(nil), // 3: llmgw.v1.Register
|
||||
(*RunningShard)(nil), // 4: llmgw.v1.RunningShard
|
||||
(*GPU)(nil), // 5: llmgw.v1.GPU
|
||||
(*FabricAttachment)(nil), // 6: llmgw.v1.FabricAttachment
|
||||
(*LocalWeights)(nil), // 7: llmgw.v1.LocalWeights
|
||||
(*Heartbeat)(nil), // 8: llmgw.v1.Heartbeat
|
||||
(*ShardStatus)(nil), // 9: llmgw.v1.ShardStatus
|
||||
(*StartShard)(nil), // 10: llmgw.v1.StartShard
|
||||
(*EngineProfile)(nil), // 11: llmgw.v1.EngineProfile
|
||||
(*StopShard)(nil), // 12: llmgw.v1.StopShard
|
||||
(*DrainShard)(nil), // 13: llmgw.v1.DrainShard
|
||||
(*Ack)(nil), // 14: llmgw.v1.Ack
|
||||
nil, // 15: llmgw.v1.Register.LabelsEntry
|
||||
nil, // 16: llmgw.v1.Register.EngineVersionsEntry
|
||||
nil, // 17: llmgw.v1.EngineProfile.EnvEntry
|
||||
(*Storage)(nil), // 4: llmgw.v1.Storage
|
||||
(*RunningShard)(nil), // 5: llmgw.v1.RunningShard
|
||||
(*GPU)(nil), // 6: llmgw.v1.GPU
|
||||
(*FabricAttachment)(nil), // 7: llmgw.v1.FabricAttachment
|
||||
(*LocalWeights)(nil), // 8: llmgw.v1.LocalWeights
|
||||
(*Heartbeat)(nil), // 9: llmgw.v1.Heartbeat
|
||||
(*ShardStatus)(nil), // 10: llmgw.v1.ShardStatus
|
||||
(*StartShard)(nil), // 11: llmgw.v1.StartShard
|
||||
(*EngineProfile)(nil), // 12: llmgw.v1.EngineProfile
|
||||
(*RemoveWeights)(nil), // 13: llmgw.v1.RemoveWeights
|
||||
(*StopShard)(nil), // 14: llmgw.v1.StopShard
|
||||
(*DrainShard)(nil), // 15: llmgw.v1.DrainShard
|
||||
(*Ack)(nil), // 16: llmgw.v1.Ack
|
||||
nil, // 17: llmgw.v1.Register.LabelsEntry
|
||||
nil, // 18: llmgw.v1.Register.EngineVersionsEntry
|
||||
nil, // 19: llmgw.v1.EngineProfile.EnvEntry
|
||||
}
|
||||
var file_llmgw_v1_agent_proto_depIdxs = []int32{
|
||||
3, // 0: llmgw.v1.AgentMessage.register:type_name -> llmgw.v1.Register
|
||||
8, // 1: llmgw.v1.AgentMessage.heartbeat:type_name -> llmgw.v1.Heartbeat
|
||||
9, // 2: llmgw.v1.AgentMessage.shard:type_name -> llmgw.v1.ShardStatus
|
||||
10, // 3: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard
|
||||
12, // 4: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard
|
||||
13, // 5: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard
|
||||
14, // 6: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack
|
||||
15, // 7: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry
|
||||
5, // 8: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU
|
||||
6, // 9: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment
|
||||
16, // 10: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry
|
||||
4, // 11: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard
|
||||
7, // 12: llmgw.v1.Register.weights:type_name -> llmgw.v1.LocalWeights
|
||||
0, // 13: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
0, // 14: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
11, // 15: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile
|
||||
17, // 16: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry
|
||||
1, // 17: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage
|
||||
2, // 18: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage
|
||||
18, // [18:19] is the sub-list for method output_type
|
||||
17, // [17:18] is the sub-list for method input_type
|
||||
17, // [17:17] is the sub-list for extension type_name
|
||||
17, // [17:17] is the sub-list for extension extendee
|
||||
0, // [0:17] is the sub-list for field type_name
|
||||
9, // 1: llmgw.v1.AgentMessage.heartbeat:type_name -> llmgw.v1.Heartbeat
|
||||
10, // 2: llmgw.v1.AgentMessage.shard:type_name -> llmgw.v1.ShardStatus
|
||||
11, // 3: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard
|
||||
14, // 4: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard
|
||||
15, // 5: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard
|
||||
16, // 6: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack
|
||||
13, // 7: llmgw.v1.ControlMessage.remove_weights:type_name -> llmgw.v1.RemoveWeights
|
||||
17, // 8: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry
|
||||
6, // 9: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU
|
||||
7, // 10: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment
|
||||
18, // 11: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry
|
||||
5, // 12: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard
|
||||
8, // 13: llmgw.v1.Register.weights:type_name -> llmgw.v1.LocalWeights
|
||||
4, // 14: llmgw.v1.Register.storage:type_name -> llmgw.v1.Storage
|
||||
0, // 15: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
4, // 16: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage
|
||||
8, // 17: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights
|
||||
0, // 18: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
12, // 19: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile
|
||||
19, // 20: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry
|
||||
1, // 21: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage
|
||||
2, // 22: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage
|
||||
22, // [22:23] is the sub-list for method output_type
|
||||
21, // [21:22] is the sub-list for method input_type
|
||||
21, // [21:21] is the sub-list for extension type_name
|
||||
21, // [21:21] is the sub-list for extension extendee
|
||||
0, // [0:21] is the sub-list for field type_name
|
||||
}
|
||||
|
||||
func init() { file_llmgw_v1_agent_proto_init() }
|
||||
@@ -1738,6 +1954,7 @@ func file_llmgw_v1_agent_proto_init() {
|
||||
(*ControlMessage_Stop)(nil),
|
||||
(*ControlMessage_Drain)(nil),
|
||||
(*ControlMessage_Ack)(nil),
|
||||
(*ControlMessage_RemoveWeights)(nil),
|
||||
}
|
||||
type x struct{}
|
||||
out := protoimpl.TypeBuilder{
|
||||
@@ -1745,7 +1962,7 @@ func file_llmgw_v1_agent_proto_init() {
|
||||
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
|
||||
RawDescriptor: unsafe.Slice(unsafe.StringData(file_llmgw_v1_agent_proto_rawDesc), len(file_llmgw_v1_agent_proto_rawDesc)),
|
||||
NumEnums: 1,
|
||||
NumMessages: 17,
|
||||
NumMessages: 19,
|
||||
NumExtensions: 0,
|
||||
NumServices: 1,
|
||||
},
|
||||
|
||||
@@ -33,6 +33,7 @@ message ControlMessage {
|
||||
StopShard stop = 2;
|
||||
DrainShard drain = 3;
|
||||
Ack ack = 4;
|
||||
RemoveWeights remove_weights = 5;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -90,6 +91,33 @@ message Register {
|
||||
// nothing — an agent with no model directories configured behaves exactly as
|
||||
// every agent did before this existed, and constrains nothing.
|
||||
repeated LocalWeights weights = 11;
|
||||
|
||||
// storage is the disk behind those weights, one entry per filesystem.
|
||||
//
|
||||
// Reported so an operator can see a node filling up before a fetch dies at
|
||||
// ninety percent, and so placement can refuse a node that cannot hold a
|
||||
// model rather than discovering it during the download.
|
||||
//
|
||||
// Empty means the node was not asked to look, exactly as for weights, and
|
||||
// constrains nothing.
|
||||
repeated Storage storage = 12;
|
||||
}
|
||||
|
||||
// Storage is one filesystem a node keeps weights on.
|
||||
//
|
||||
// Per filesystem rather than per configured directory: free space belongs to
|
||||
// the disk, and two model directories on one disk are one pool of bytes.
|
||||
message Storage {
|
||||
// path is one directory on this filesystem, so the operator recognises it.
|
||||
string path = 1;
|
||||
// free_bytes is what a non-root process can write — not the root reserve,
|
||||
// which is space the agent cannot have.
|
||||
int64 free_bytes = 2;
|
||||
int64 total_bytes = 3;
|
||||
// shared marks a filesystem the node did not fill. Free space on one is
|
||||
// shared with every other node using it, so two nodes each reporting 500 GB
|
||||
// free on the same filer have 500 GB between them and not a terabyte.
|
||||
bool shared = 4;
|
||||
}
|
||||
|
||||
// RunningShard describes an engine already running on a node.
|
||||
@@ -171,6 +199,22 @@ message Heartbeat {
|
||||
// a card after registration.
|
||||
repeated int32 gpu_memory_used_mb = 2;
|
||||
repeated string new_problems = 3;
|
||||
|
||||
// storage for the same reason the GPU memory is here: it changes while the
|
||||
// agent runs, and a figure sampled once at registration is a figure that is
|
||||
// wrong on exactly the node worth watching — the one busy filling up.
|
||||
//
|
||||
// Always sent, because it is a few numbers per filesystem.
|
||||
repeated Storage storage = 4;
|
||||
|
||||
// weights only when they changed, because the list is long enough that
|
||||
// sending it every few seconds would be paying for nothing most of the time.
|
||||
//
|
||||
// weights_changed is what says the list is meant: without it an empty list
|
||||
// would be ambiguous between "this node now holds nothing" and "nothing to
|
||||
// report", and the first would quietly empty a node's inventory.
|
||||
bool weights_changed = 5;
|
||||
repeated LocalWeights weights = 6;
|
||||
}
|
||||
|
||||
// ShardStatus reports a shard's lifecycle transition.
|
||||
@@ -239,6 +283,17 @@ message EngineProfile {
|
||||
string served_model_name = 11;
|
||||
}
|
||||
|
||||
// RemoveWeights deletes cached model files a node no longer needs.
|
||||
//
|
||||
// The paths are computed by the control plane from the catalog, but the agent
|
||||
// checks them again before unlinking anything: it deletes only files inside its
|
||||
// own cache directory whose names this system generated. A control plane that
|
||||
// is wrong, compromised or simply newer than the agent must not be able to
|
||||
// talk it into deleting a model an operator spent a week copying there.
|
||||
message RemoveWeights {
|
||||
repeated string paths = 1;
|
||||
}
|
||||
|
||||
// StopShard terminates a shard and frees its GPUs.
|
||||
message StopShard {
|
||||
string shard_id = 1;
|
||||
|
||||
+175
-10
@@ -59,7 +59,7 @@ a certificate and changing it later means reissuing one.
|
||||
One static binary, no runtime dependencies, the same on every machine.
|
||||
|
||||
```
|
||||
VERSION=v0.6.0 # check the releases page for the current tag
|
||||
VERSION=v0.8.0 # check the releases page for the current tag
|
||||
curl -fsSLO https://gitea.mynodes.xyz/ilker/llmgw/releases/download/$VERSION/llmgw-$VERSION-linux-amd64
|
||||
curl -fsSL https://gitea.mynodes.xyz/ilker/llmgw/releases/download/$VERSION/checksums.txt | sha256sum -c --ignore-missing
|
||||
sudo install -m0755 llmgw-$VERSION-linux-amd64 /usr/local/bin/llmgw
|
||||
@@ -305,9 +305,10 @@ llmgw ctl engine-check vllm --bin /opt/venvs/vllm/bin/vllm
|
||||
|
||||
### Weights
|
||||
|
||||
Put them where the node can read them. There is no distribution mechanism yet;
|
||||
`scp`, a shared filesystem, or the engine's own downloader are all valid, and
|
||||
which you want depends on how many nodes you have.
|
||||
Put them where the node can read them. `scp`, a shared filesystem, or the
|
||||
engine's own downloader are all valid, and which you want depends on how many
|
||||
nodes you have. If you would rather not choose, skip to the next section and
|
||||
let llmgw fetch them.
|
||||
|
||||
```
|
||||
sudo install -d -o root -g root /var/lib/llmgw-models
|
||||
@@ -501,8 +502,8 @@ scaling: { min: 0, max: 1, idle_timeout: 15m }
|
||||
model starts it and waits briefly; if it is not serving by then the client gets
|
||||
a 503 with `Retry-After` and an honest "the model is loading", rather than being
|
||||
told nothing serves it — a client told to retry succeeds where one told the
|
||||
model does not exist gives up. Once nobody has asked for `idle_timeout`, the
|
||||
replicas are stopped and the cards go back to the pool.
|
||||
model does not exist gives up. Once nobody has asked for it in `idle_timeout`,
|
||||
the replicas are stopped and the cards go back to the pool.
|
||||
|
||||
`min: 0, max: 0` is different and equally useful: configured, and deliberately
|
||||
not running. A model with no `scaling` stanza at all gets that, because making
|
||||
@@ -566,9 +567,9 @@ temperature, ECC faults and utilisation.
|
||||
|
||||
**The dashboard is the answer for most clusters.** `/admin` shows capacity per
|
||||
GPU class, free whole nodes per fabric, which replicas are holding hardware
|
||||
nothing is asking for, and the same node and model tables `ctl` prints — with
|
||||
no monitoring stack to stand up beside it. One binary, no dependencies, which
|
||||
is the point.
|
||||
nothing is asking for, who is using the cluster and what they are running, and
|
||||
the same node and model tables `ctl` prints — with no monitoring stack to stand
|
||||
up beside it. One binary, no dependencies, which is the point.
|
||||
|
||||
Prometheus is worth adding when you want history, alerts, or one view across
|
||||
several clusters. The dashboard shows what is true now; it keeps nothing.
|
||||
@@ -594,6 +595,161 @@ fleet means the measurement is missing, not that the fleet is idle. And
|
||||
`llmgw_gpu_nvlink_peers` at zero everywhere means the topology is unknown,
|
||||
which is what a binary built without the nvml tag reports.
|
||||
|
||||
## PostgreSQL, when SQLite is no longer the right shape
|
||||
|
||||
SQLite is the default and stays the right answer for most installations: one
|
||||
file, no server, and a request path that barely touches it — usage is written
|
||||
in batches, budgets are counted in memory, authentication is cached.
|
||||
|
||||
Two things change that. Conversations and agent notes read and write on every
|
||||
turn, which is the first per-request load this store has had; and SQLite takes
|
||||
one writer at a time, so those turns serialise. And a second control plane is
|
||||
impossible while the store is a file on one machine's disk.
|
||||
|
||||
The setting is the same one, and the address decides:
|
||||
|
||||
```yaml
|
||||
auth:
|
||||
store: "@/etc/llmgw/secrets/store-dsn"
|
||||
```
|
||||
|
||||
A path is SQLite, a `postgres://` URL is PostgreSQL, and `@` followed by a path
|
||||
reads the address out of that file — which is how a DSN carrying a password
|
||||
stays out of a configuration that gets copied, diffed and pasted into issues.
|
||||
Keep that file at `0600` and owned by the user llmgw runs as.
|
||||
|
||||
Moving an existing installation is a copy, not a fresh start. Keys are stored
|
||||
hashed, so starting fresh means reissuing every one of them and being unable
|
||||
to say which was which:
|
||||
|
||||
```
|
||||
sudo systemctl stop llmgw-server
|
||||
llmgw keys move-store --from /var/lib/llmgw/llmgw.db --to "@/etc/llmgw/secrets/store-dsn"
|
||||
# point auth.store at the new address, then
|
||||
sudo systemctl start llmgw-server
|
||||
```
|
||||
|
||||
Stop the gateway first. The copy makes no attempt to follow a store being
|
||||
written to: rows landing after the table they belong to has been passed would
|
||||
simply be absent, and nothing would say so. The destination must be empty, and
|
||||
the source is left untouched — so the way back is to point the configuration at
|
||||
the file again.
|
||||
|
||||
On a machine that also runs the control plane, PostgreSQL's defaults assume it
|
||||
owns the box. A gateway swapped out to give the database cache is a worse
|
||||
outcome than a smaller cache, so cap it: `shared_buffers = 96MB` and
|
||||
`max_connections = 40` are enough for a store whose working set is measured in
|
||||
megabytes.
|
||||
|
||||
## Conversations across models
|
||||
|
||||
Off everywhere until a project asks. llmgw otherwise stores no message content
|
||||
at all — usage rows carry token counts, costs and latencies, never text, and
|
||||
even its own chat page keeps the conversation in the browser. A gateway that
|
||||
keeps nothing anyone typed is a much easier thing to have in front of a team,
|
||||
so keeping it is a decision someone makes rather than a default they discover.
|
||||
|
||||
```
|
||||
llmgw keys threads --project default --retain-hours 24
|
||||
```
|
||||
|
||||
The retention is required. A conversation store with no expiry is a transcript
|
||||
of everything a team ever typed, kept forever because nobody chose a number.
|
||||
The window is measured from the last turn, so a conversation still in progress
|
||||
does not expire underneath it.
|
||||
|
||||
A client then sends a thread id and only its new turn:
|
||||
|
||||
```
|
||||
curl $BASE/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LLMGW_API_KEY" \
|
||||
-H "x-llmgw-thread: my-conversation" \
|
||||
-d '{"model": "<model-name>", "messages": [{"role":"user","content":"and then?"}]}'
|
||||
```
|
||||
|
||||
llmgw puts the stored history in front of that turn and sends it to whichever
|
||||
model is serving — **including a different one from the one the conversation
|
||||
started on**. That is the point: a thread begun on a large model can continue
|
||||
on a small one, and a caller dropped to a cheaper model by a spent budget keeps
|
||||
their conversation.
|
||||
|
||||
When the history does not fit the model now serving, the oldest turns are
|
||||
dropped until it does, and the reply carries `x-llmgw-thread-trimmed` with how
|
||||
many. System messages are never dropped, and neither is the turn the caller
|
||||
just sent.
|
||||
|
||||
`DELETE /v1/threads/{id}` forgets one. It is a client endpoint on purpose: the
|
||||
person who said the thing should be able to unsay it without asking an
|
||||
operator to read it first. A thread belonging to another project answers the
|
||||
same as one that never existed.
|
||||
|
||||
Turning it off again stops new storage, and existing conversations go at the
|
||||
next hourly prune:
|
||||
|
||||
```
|
||||
llmgw keys threads --project default --off
|
||||
```
|
||||
|
||||
## Memory notes for agents
|
||||
|
||||
Also off until a project asks, for the same reason. Where conversations are what
|
||||
was said, notes are what an agent chose to keep, and they persist across
|
||||
conversations:
|
||||
|
||||
```
|
||||
llmgw keys memory --project default --retain-hours 720
|
||||
```
|
||||
|
||||
The retention is required and is measured from the last write to a note. An agent
|
||||
then saves, searches and deletes notes through `/v1/memory` or, as tools, through
|
||||
the MCP server at `/mcp` (`remember`, `recall`, `forget`). Both need the API key
|
||||
of the project, and notes are private to the project and to the agent id sent in
|
||||
`x-llmgw-agent`. A note in another project or agent answers the same as one that
|
||||
does not exist. llmgw does not add these tools to requests itself.
|
||||
|
||||
Turning it off stops new storage, and existing notes go at the next hourly
|
||||
prune:
|
||||
|
||||
```
|
||||
llmgw keys memory --project default --off
|
||||
```
|
||||
|
||||
Without authentication, `threads.store` and `threads.retain_hours` in
|
||||
`server.yaml` keep both conversations and notes in one shared pool. That is for
|
||||
one trusted machine only; see [quickstart.md](quickstart.md).
|
||||
|
||||
## Reclaiming disk
|
||||
|
||||
`llmgw ctl nodes` has a DISK column: free space against the size it is free
|
||||
out of, for the filesystems holding model weights. A dash means the node was
|
||||
not asked — an agent with no `model_dirs`, or one older than this — and not
|
||||
that it has no space. Shared filesystems are excluded from a node's figure,
|
||||
because space on a filer belongs to everyone mounting it.
|
||||
|
||||
Weights llmgw downloaded for models that are no longer in the catalog can be
|
||||
listed and removed:
|
||||
|
||||
```
|
||||
llmgw ctl weights # what could be freed, per node
|
||||
llmgw ctl weights --collect # delete it
|
||||
```
|
||||
|
||||
**Only files llmgw fetched itself are ever candidates.** They are recognisable
|
||||
because llmgw names them after a hash of their source, so the set still wanted
|
||||
is computed from the catalog exactly rather than guessed. Anything you put on
|
||||
a node — by scp, by a shared mount, by hand — is left alone however
|
||||
unreferenced it looks, and no llmgw command will remove it. The nodes check
|
||||
every path again before deleting it, so a control plane that is wrong cannot
|
||||
talk one into losing a model that took a week to copy.
|
||||
|
||||
Which also means llmgw will not reclaim the storage most likely to be filling
|
||||
your disk. Weights you placed yourself are yours to remove.
|
||||
|
||||
A profile may declare `min_disk_gb`, and a node without that much free is then
|
||||
not a candidate for it — checked only on nodes that do not already hold the
|
||||
weights, since those need no room for them. Without the field, disk is not
|
||||
considered, which is how every profile written before it keeps behaving.
|
||||
|
||||
## When something does not work
|
||||
|
||||
**The node is `ready` but every replica fails.** Almost always a missing engine
|
||||
@@ -625,7 +781,16 @@ requests. On a slow model, set `timeouts.generation` higher.
|
||||
|
||||
## Known gaps, so you find them here rather than at 2am
|
||||
|
||||
- **Weights are yours to distribute.** No cache, no fetch, no progress.
|
||||
- **A fetch reports its progress where nobody is looking.** The agent logs a
|
||||
percentage as it downloads, so `journalctl -u llmgw-agent` on that node knows
|
||||
exactly how far along it is. Nothing carries it to the control plane, so from
|
||||
`ctl` or the dashboard a replica pulling a hundred gigabytes and one that is
|
||||
wedged still look identical — which is the comparison the log line itself
|
||||
says matters.
|
||||
- **Every node fetches its own copy.** There is no seeding between nodes, so
|
||||
bringing up eight nodes on a new model is eight downloads from the same
|
||||
origin. On a fabric where the nodes can see each other this is wasteful, and
|
||||
it is the obvious next thing to build.
|
||||
- **Engines are yours to install.** Onboarding verifies one runs; it installs
|
||||
nothing.
|
||||
- **One control plane.** No HA yet; it is a single process with a SQLite file.
|
||||
|
||||
@@ -139,7 +139,9 @@ So the node is not finished until:
|
||||
Both are deliberately outside this script. Engines are a platform decision with
|
||||
versions and drivers attached, and distributing weights is its own problem —
|
||||
hundreds of gigabytes that want a shared filesystem or a cache, not an scp from
|
||||
a laptop. llmgw has no answer for it yet; that is a known gap, not an oversight.
|
||||
a laptop. llmgw will fetch weights for a model whose `source` is a URL, one
|
||||
copy per node, which covers the simple case; anything larger stays the
|
||||
operator's arrangement, and llmgw is content to find the files already there.
|
||||
|
||||
The `sim` engine needs neither, which is what makes it useful for exercising
|
||||
everything else.
|
||||
|
||||
@@ -0,0 +1,375 @@
|
||||
# Quickstart: one machine, one model, ten minutes
|
||||
|
||||
Written for someone who wants to see llmgw answer a request on a single
|
||||
workstation before reading anything else. No systemd, no `sudo`. For a real
|
||||
cluster, see [install.md](install.md).
|
||||
|
||||
What you end up with: an OpenAI-compatible endpoint on `localhost:8080` in front
|
||||
of a real model served by llama.cpp.
|
||||
|
||||
## Two ways to run
|
||||
|
||||
Steps 1 to 4 use the simple one, **router-only**: you start the engines yourself
|
||||
and tell the router where they are with `--endpoint`. No certificates, no node
|
||||
agent.
|
||||
|
||||
Step 5 uses the **control plane and node agent**, the same architecture as a real
|
||||
cluster on one machine: the agent starts and supervises the engine, and the
|
||||
control plane places the model.
|
||||
|
||||
## 1. Build
|
||||
|
||||
```
|
||||
git clone https://gitea.mynodes.xyz/ilker/llmgw && cd llmgw
|
||||
make build
|
||||
make test
|
||||
./bin/llmgw --version
|
||||
```
|
||||
|
||||
`make build` writes `./bin/llmgw`. Go 1.26 or later is the only requirement.
|
||||
|
||||
## 2. Check the path with a fake engine
|
||||
|
||||
Before involving a GPU, prove the router works. `sim-engine` is an
|
||||
OpenAI-compatible fake that streams placeholder words.
|
||||
|
||||
```
|
||||
./bin/llmgw sim-engine --listen :8001 --model llama &
|
||||
./bin/llmgw server --endpoint e1=http://127.0.0.1:8001=llama &
|
||||
|
||||
curl -N localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' \
|
||||
-H 'x-session-id: demo' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"stream":true}'
|
||||
```
|
||||
|
||||
Chunks of `alpha beta gamma ...` ending in `data: [DONE]` mean the router is
|
||||
fine. Stop both before continuing (see [Stopping](#stopping)).
|
||||
|
||||
## 3. Serve a real model
|
||||
|
||||
Two processes, two ports. Clients only ever talk to the router; the router
|
||||
knows where the engine is because you tell it with `--endpoint`:
|
||||
|
||||
```
|
||||
curl :8080 -> llmgw router :8080 -> llama serve :8082 -> model
|
||||
(client) (the gateway) (the engine)
|
||||
```
|
||||
|
||||
Requests go to 8080. Port 8082 is the engine's own, and you only touch it to
|
||||
check that the engine is up.
|
||||
|
||||
llmgw does not install engines or weights. You need:
|
||||
|
||||
- a `llama-server` build from llama.cpp (CUDA build if you want the GPU), and
|
||||
- a `.gguf` file.
|
||||
|
||||
Check the card first with `nvidia-smi`. An 8 GB card holds a 1B model at bf16
|
||||
(2.4 GB) or a 9B model at `q4_k_m` (5.4 GB).
|
||||
|
||||
Start the engine on its own port. The model name given to `--alias` is the name
|
||||
clients will use:
|
||||
|
||||
```
|
||||
llama serve \
|
||||
-m /path/to/Llama-3.2-1B-Instruct-BF16.gguf \
|
||||
--host 0.0.0.0 --port 8082 -ngl 99 --alias llama
|
||||
```
|
||||
|
||||
`-ngl 99` offloads every layer to the GPU. Confirm it is up and knows the model:
|
||||
|
||||
```
|
||||
curl -s localhost:8082/v1/models
|
||||
```
|
||||
|
||||
Then point the router at it. The format is `id=URL=model`:
|
||||
|
||||
```
|
||||
./bin/llmgw server --endpoint e1=http://127.0.0.1:8082=llama
|
||||
```
|
||||
|
||||
## 4. Ask it something
|
||||
|
||||
Streaming:
|
||||
|
||||
```
|
||||
curl -N localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' \
|
||||
-H 'x-session-id: demo' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"stream":true}'
|
||||
```
|
||||
|
||||
One reply at once:
|
||||
|
||||
```
|
||||
curl -s localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"max_tokens":100}'
|
||||
```
|
||||
|
||||
Also useful: `curl -s localhost:8080/v1/models` and
|
||||
`curl -s localhost:8080/health`. `llmgw ctl` does not work in this mode: it
|
||||
talks to the control plane's admin API, which router-only mode does not have,
|
||||
so `ctl status` answers `404`.
|
||||
|
||||
On the machine this was written on, a 1B bf16 model on an RTX 4070 laptop GPU
|
||||
answered at roughly 85 tokens per second.
|
||||
|
||||
## Several replicas of one model
|
||||
|
||||
Give the router more than one `--endpoint` with the same model name and it
|
||||
spreads requests across them. `x-session-id` keeps a conversation on the same
|
||||
replica so its cache is reused:
|
||||
|
||||
```
|
||||
./bin/llmgw server \
|
||||
--endpoint e1=http://127.0.0.1:8082=llama \
|
||||
--endpoint e2=http://127.0.0.1:8083=llama
|
||||
```
|
||||
|
||||
## Stopping
|
||||
|
||||
```
|
||||
pkill -f "llmgw server"
|
||||
pkill -f "llmgw sim-engine"
|
||||
```
|
||||
|
||||
Stop the engine the way you started it (Ctrl-C in its terminal).
|
||||
|
||||
## What goes wrong
|
||||
|
||||
**`listening on :8080: bind: address already in use`.** An earlier router is
|
||||
still running. `ss -ltnp | grep 8080` names the process; stop it first. This is
|
||||
the usual result of moving from step 2 to step 3 without stopping the fake
|
||||
engine's router.
|
||||
|
||||
**`/v1/models` on the engine returns an empty list.** The engine was started
|
||||
without `-m`, so it is a router with nothing to load. Restart it with a model.
|
||||
|
||||
**`error: nothing to serve`.** `llmgw server` with no `--endpoint` and no model
|
||||
carrying `profiles` has nothing to do. In router-only mode, always pass at least
|
||||
one `--endpoint`.
|
||||
|
||||
**Requests fail with a model-not-found error.** The name in the request must
|
||||
match the name after the last `=` in `--endpoint`, and it must be a name the
|
||||
engine actually serves (its `--alias`).
|
||||
|
||||
**Answers are poor.** A 1B model is a plumbing test, not a chat assistant. Swap
|
||||
in a larger `.gguf`; nothing else changes.
|
||||
|
||||
## 5. Let llmgw run the engine (control plane and node agent)
|
||||
|
||||
Steps 2 and 3 start the engine by hand. Here a **node agent** starts and
|
||||
supervises llama.cpp for you, and a control plane decides where the model runs.
|
||||
Still one machine, still no `sudo`. Stop everything from steps 2 and 3 first:
|
||||
the GPU has to be free and port 8080 unused.
|
||||
|
||||
```
|
||||
curl :8080 -> llmgw server (control plane + router) <- llmgw agent (this node)
|
||||
|
|
||||
starts llama-server :20000
|
||||
```
|
||||
|
||||
The agent dials the control plane over mutual TLS, so this needs a small private
|
||||
certificate authority. Everything goes in a scratch directory, `devcluster/`,
|
||||
which is in `.gitignore`.
|
||||
|
||||
### Certificates
|
||||
|
||||
```
|
||||
mkdir -p devcluster/pki
|
||||
./bin/llmgw ctl cert init-ca --dir devcluster/pki
|
||||
./bin/llmgw ctl cert issue-server llmgw-ctl --host 127.0.0.1 --host localhost \
|
||||
--dir devcluster/pki --out devcluster
|
||||
./bin/llmgw ctl cert issue-client llmgw-node01 --dir devcluster/pki --out devcluster
|
||||
```
|
||||
|
||||
The name given to `issue-client` becomes the certificate's common name and must
|
||||
equal `node_id` in the agent config. The control plane refuses an agent whose
|
||||
claimed id differs from its certificate.
|
||||
|
||||
### The engine wrapper
|
||||
|
||||
The agent starts the engine as `<bin> --model ... --port ...`. A `llama` binary
|
||||
that wants `llama serve ...` refuses every flag, so put a wrapper in front of it:
|
||||
|
||||
```
|
||||
cat > devcluster/llama-server <<'EOF'
|
||||
#!/bin/sh
|
||||
exec /path/to/llama serve "$@"
|
||||
EOF
|
||||
chmod +x devcluster/llama-server
|
||||
|
||||
./bin/llmgw ctl engine-check llamacpp --bin devcluster/llama-server
|
||||
```
|
||||
|
||||
`engine-check` should end with `every flag llmgw builds is accepted`. If your
|
||||
build already ships a `llama-server` binary, point `bin` at that and skip the
|
||||
wrapper.
|
||||
|
||||
### Control plane config: `devcluster/server.yaml`
|
||||
|
||||
```yaml
|
||||
listen: "127.0.0.1:8080"
|
||||
admin_listen: "127.0.0.1:9090"
|
||||
agent_listen: "127.0.0.1:7443"
|
||||
|
||||
agent_tls:
|
||||
ca: devcluster/pki/ca.pem
|
||||
cert: devcluster/server.pem
|
||||
key: devcluster/server-key.pem
|
||||
|
||||
models:
|
||||
- name: llama
|
||||
source: /path/to/Llama-3.2-1B-Instruct-BF16.gguf
|
||||
routing_profile: chat
|
||||
context_length: 8192
|
||||
scaling: { min: 1, max: 1 }
|
||||
profiles:
|
||||
- name: gpu
|
||||
engine: llamacpp
|
||||
gpu_class: 4070-8g
|
||||
weight: 1
|
||||
priority: 100
|
||||
startup_timeout: 5m
|
||||
```
|
||||
|
||||
Three things in there are easy to get wrong:
|
||||
|
||||
- **The model must have at least one profile.** That is what switches the control
|
||||
plane on. With `models: []` the server starts as a plain router and exits with
|
||||
`nothing to serve`.
|
||||
- **`gpu_class` decides CPU or GPU.** Without it the replica is placed in the CPU
|
||||
pool and runs on the CPU. Get the class name from `./bin/llmgw agent inventory`
|
||||
(the `CLASS` column).
|
||||
- **There is no `auth:` block.** Leave it out for a local trial; with auth on,
|
||||
requests need an API key.
|
||||
|
||||
### Agent config: `devcluster/agent.yaml`
|
||||
|
||||
```yaml
|
||||
node_id: llmgw-node01
|
||||
server: "https://127.0.0.1:7443"
|
||||
mgmt_addr: "127.0.0.1"
|
||||
|
||||
tls:
|
||||
ca: devcluster/pki/ca.pem
|
||||
cert: devcluster/llmgw-node01.pem
|
||||
key: devcluster/llmgw-node01-key.pem
|
||||
|
||||
engines:
|
||||
llamacpp:
|
||||
bin: /absolute/path/to/devcluster/llama-server
|
||||
|
||||
supervisor:
|
||||
work_dir: devcluster/agent-work
|
||||
```
|
||||
|
||||
### Run it
|
||||
|
||||
Two terminals:
|
||||
|
||||
```
|
||||
./bin/llmgw server --config devcluster/server.yaml
|
||||
./bin/llmgw agent --config devcluster/agent.yaml
|
||||
```
|
||||
|
||||
The model is not placed at once. The control plane waits about 35 seconds for the
|
||||
fleet to settle after it starts, then places the replica, and the agent then
|
||||
loads the weights. Watch it with:
|
||||
|
||||
```
|
||||
./bin/llmgw ctl nodes
|
||||
./bin/llmgw ctl models
|
||||
./bin/llmgw ctl replicas
|
||||
./bin/llmgw ctl explain-placement llama
|
||||
```
|
||||
|
||||
`STATUS ready` in `ctl models` means it is serving. `nvidia-smi` should now show
|
||||
the model in GPU memory. The same curl from step 4 works against port 8080.
|
||||
|
||||
`ctl` works here, unlike in router-only mode, because there is a control plane
|
||||
for it to talk to.
|
||||
|
||||
Stop it with Ctrl-C in each terminal. The agent leaves the engine running when
|
||||
the agent itself stops, so if you restart the agent and the GPU stays full,
|
||||
`pgrep -af "llama serve"` will show the engine still up.
|
||||
|
||||
## 6. Memory: conversations and notes
|
||||
|
||||
Off unless you ask. llmgw otherwise stores nothing anyone typed. On a single
|
||||
machine with no API keys, one block in the server config turns on two things:
|
||||
|
||||
```yaml
|
||||
threads:
|
||||
store: devcluster/data/threads.db
|
||||
retain_hours: 24
|
||||
```
|
||||
|
||||
Restart the server after adding it. `retain_hours` is required: everything kept
|
||||
expires that long after it was last written.
|
||||
|
||||
Everything kept this way sits in one shared pool with no owner, and anyone who
|
||||
can reach the port and knows an id can read it. That is fine on one machine and
|
||||
wrong on a shared network. With auth on, threads and notes are per project
|
||||
instead and the block is ignored; see [install.md](install.md).
|
||||
|
||||
### Conversations (threads)
|
||||
|
||||
Send the same `x-llmgw-thread` id and only the new turn. llmgw puts the stored
|
||||
history in front of it:
|
||||
|
||||
```
|
||||
curl -s localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' -H 'x-llmgw-thread: demo' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"My name is Deniz."}]}'
|
||||
|
||||
curl -s localhost:8080/v1/chat/completions \
|
||||
-H 'content-type: application/json' -H 'x-llmgw-thread: demo' \
|
||||
-d '{"model":"llama","messages":[{"role":"user","content":"What is my name?"}]}'
|
||||
|
||||
curl -X DELETE localhost:8080/v1/threads/demo # forget it
|
||||
```
|
||||
|
||||
Without the block, a request carrying `x-llmgw-thread` is refused with
|
||||
`this gateway keeps no conversations` rather than answered without its history.
|
||||
|
||||
### Notes an agent keeps for itself
|
||||
|
||||
Notes outlive any one conversation. Each is a key, some text and optional tags,
|
||||
kept per agent; name the agent with `x-llmgw-agent` (default `default`), so two
|
||||
agents using the same key do not read each other's notes.
|
||||
|
||||
```
|
||||
curl -X PUT localhost:8080/v1/memory/favorite-color \
|
||||
-H 'content-type: application/json' \
|
||||
-d '{"text":"The user likes dark green.","tags":["preference"]}'
|
||||
|
||||
curl 'localhost:8080/v1/memory?q=green' # search by words
|
||||
curl 'localhost:8080/v1/memory?tag=preference' # or by tag
|
||||
curl localhost:8080/v1/memory/favorite-color
|
||||
curl -X DELETE localhost:8080/v1/memory/favorite-color
|
||||
```
|
||||
|
||||
The same notes are offered to an agent as tools over MCP at
|
||||
`http://localhost:8080/mcp`: `remember`, `recall` and `forget`. Add that URL to
|
||||
your agent's MCP server list (HTTP transport; how you set request headers such
|
||||
as `x-llmgw-agent` depends on the client). llmgw does not put these tools into
|
||||
chat requests on its own, so an agent that is not given the server never sees
|
||||
memory.
|
||||
|
||||
Limits: a key is up to 200 bytes, a note's text up to 16 KiB, at most 16 tags,
|
||||
1000 notes per agent. Tags are lowercase letters, digits, `-` and `_`.
|
||||
|
||||
## Where this stops
|
||||
|
||||
Router-only mode (steps 2 to 4) has no authentication, no placement and no model
|
||||
management: you start and restart engines yourself. Step 5 adds placement, model
|
||||
management and a node agent, and can add authentication as `install.md`
|
||||
describes. For a real cluster, with several nodes, systemd and a persistent
|
||||
certificate authority, follow [install.md](install.md).
|
||||
|
||||
One caveat found while writing this page: `install.md` recommends starting with
|
||||
an empty `models: []` and adding models later with `ctl model add`. The control
|
||||
plane only turns on when the config declares a model with `profiles`, so that
|
||||
first start does not produce a control plane.
|
||||
@@ -6,6 +6,7 @@ require (
|
||||
connectrpc.com/connect v1.21.0
|
||||
github.com/NVIDIA/go-nvml v0.13.4-0
|
||||
github.com/cespare/xxhash/v2 v2.3.0
|
||||
github.com/jackc/pgx/v5 v5.11.0
|
||||
github.com/knadh/koanf/parsers/yaml v1.1.1
|
||||
github.com/knadh/koanf/providers/env/v2 v2.0.1
|
||||
github.com/knadh/koanf/providers/file v1.2.1
|
||||
@@ -14,6 +15,7 @@ require (
|
||||
github.com/prometheus/client_model v0.6.3
|
||||
github.com/prometheus/common v0.71.0
|
||||
github.com/spf13/cobra v1.10.2
|
||||
golang.org/x/crypto v0.57.0
|
||||
golang.org/x/net v0.59.0
|
||||
golang.org/x/sync v0.23.0
|
||||
google.golang.org/protobuf v1.36.12
|
||||
@@ -27,6 +29,9 @@ require (
|
||||
github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/inconshreveable/mousetrap v1.1.0 // indirect
|
||||
github.com/jackc/pgpassfile v1.0.0 // indirect
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
|
||||
github.com/jackc/puddle/v2 v2.2.2 // indirect
|
||||
github.com/knadh/koanf/maps v0.1.2 // indirect
|
||||
github.com/mattn/go-isatty v0.0.24 // indirect
|
||||
github.com/mitchellh/copystructure v1.2.0 // indirect
|
||||
@@ -37,7 +42,6 @@ require (
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||
github.com/spf13/pflag v1.0.9 // indirect
|
||||
go.yaml.in/yaml/v3 v3.0.5 // indirect
|
||||
golang.org/x/crypto v0.57.0 // indirect
|
||||
golang.org/x/sys v0.48.0 // indirect
|
||||
golang.org/x/text v0.42.0 // indirect
|
||||
modernc.org/libc v1.75.7 // indirect
|
||||
|
||||
@@ -7,6 +7,7 @@ github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6r
|
||||
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
|
||||
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
||||
github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
||||
github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k=
|
||||
@@ -23,6 +24,14 @@ github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
|
||||
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
|
||||
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
|
||||
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
|
||||
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
|
||||
github.com/jackc/pgx/v5 v5.11.0 h1:IzBBtyK9AHqf98cctWFifYSci2hgQR/cd56wB4p+ogg=
|
||||
github.com/jackc/pgx/v5 v5.11.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4=
|
||||
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
|
||||
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
|
||||
github.com/klauspost/compress v1.19.1 h1:VsB4HPswih7mmZ8WleSFQ75c/Ui1M4trX5oAsJnhSlk=
|
||||
github.com/klauspost/compress v1.19.1/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ=
|
||||
github.com/knadh/koanf/maps v0.1.2 h1:RBfmAW5CnZT+PJ1CVc1QSJKf4Xu9kxfQgYVQSu8hpbo=
|
||||
@@ -47,6 +56,7 @@ github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
|
||||
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/prometheus/client_golang v1.24.1 h1:JnJkREXzWxUdCuPFpIWZiPispT9xVV59uiuyR2bPlnU=
|
||||
github.com/prometheus/client_golang v1.24.1/go.mod h1:F+oSRECHg4sse5ucfYpYDeIv/hu68Zo0uoHKetWnzcE=
|
||||
github.com/prometheus/client_model v0.6.3 h1:O0jaTVAYNxTHYInEPFJt5I3+sN8zqBtVMPTB1qyxiEo=
|
||||
@@ -62,6 +72,9 @@ github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU=
|
||||
github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4=
|
||||
github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY=
|
||||
github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
|
||||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
|
||||
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE=
|
||||
github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg=
|
||||
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
|
||||
@@ -88,6 +101,7 @@ golang.org/x/tools v0.49.0/go.mod h1:SJNXV9DBKT0UbdttsQjbfJlAE/q+y36++zo3uL3N0Oo
|
||||
google.golang.org/protobuf v1.36.12 h1:pJOKDDOyeXErUroCihFAd5LQuwXBSpVnKGrj5o/fwxc=
|
||||
google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
modernc.org/cc/v4 v4.29.2 h1:h6+9ciCnPKutf4I03CvheAvDLX7+IHlqR6Iy6J+cgd8=
|
||||
modernc.org/cc/v4 v4.29.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI=
|
||||
modernc.org/ccgo/v4 v4.35.0 h1:F+TUsmw09QxLzmi3aeYYGxjAXarmZaKgj3mKQHNaA8w=
|
||||
|
||||
+184
-24
@@ -3,7 +3,9 @@ package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"crypto/tls"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
@@ -11,6 +13,8 @@ import (
|
||||
"math/rand/v2"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -30,8 +34,18 @@ import (
|
||||
|
||||
// Config configures the agent.
|
||||
type Config struct {
|
||||
NodeID string `koanf:"node_id"`
|
||||
NodeID string `koanf:"node_id"`
|
||||
// Server is the control plane, and Servers is the same thing when there is
|
||||
// more than one. Both are accepted: every configuration written before
|
||||
// there could be a second control plane names one address, and none of
|
||||
// them should have to change.
|
||||
//
|
||||
// The order is tried in order, and a refusal moves to the next. Only the
|
||||
// leader accepts an agent, so a standby is skipped in the time it takes to
|
||||
// be told so — which is what makes failover a reconnect rather than an
|
||||
// edit to every node's configuration.
|
||||
Server string `koanf:"server"`
|
||||
Servers []string `koanf:"servers"`
|
||||
Hostname string `koanf:"hostname"`
|
||||
MgmtAddr string `koanf:"mgmt_addr"`
|
||||
Labels map[string]string `koanf:"labels"`
|
||||
@@ -104,9 +118,6 @@ func (c *Config) normalize() error {
|
||||
if c.NodeID == "" {
|
||||
return errors.New("node_id is required")
|
||||
}
|
||||
if c.Server == "" {
|
||||
return errors.New("server is required")
|
||||
}
|
||||
if err := c.TLS.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -136,26 +147,59 @@ func (c *Config) normalize() error {
|
||||
// that says plaintext while TLS is configured. Silently connecting anyway would
|
||||
// mean a cluster believing it has mutual authentication when it has none.
|
||||
func (c *Config) normalizeServerURL() error {
|
||||
secure := c.TLS.Enabled()
|
||||
switch {
|
||||
case strings.HasPrefix(c.Server, "https://"):
|
||||
if !secure {
|
||||
return errors.New("server is an https:// address but no tls certificates are configured")
|
||||
}
|
||||
case strings.HasPrefix(c.Server, "http://"):
|
||||
if secure {
|
||||
return errors.New("tls is configured but server is an http:// address; use https://")
|
||||
}
|
||||
case strings.Contains(c.Server, "://"):
|
||||
return fmt.Errorf("server %q must be an http:// or https:// address", c.Server)
|
||||
case secure:
|
||||
c.Server = "https://" + c.Server
|
||||
default:
|
||||
c.Server = "http://" + c.Server
|
||||
// One list, however it was written. Server first when both are given,
|
||||
// because an operator adding a standby to an existing file expects the
|
||||
// one already there to stay preferred.
|
||||
all := make([]string, 0, 1+len(c.Servers))
|
||||
if c.Server != "" {
|
||||
all = append(all, c.Server)
|
||||
}
|
||||
all = append(all, c.Servers...)
|
||||
|
||||
seen := map[string]struct{}{}
|
||||
out := make([]string, 0, len(all))
|
||||
for _, s := range all {
|
||||
fixed, err := normalizeOne(s, c.TLS.Enabled())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// A duplicate would be tried twice in a row on failover, which reads
|
||||
// as the rotation not working.
|
||||
if _, dup := seen[fixed]; dup {
|
||||
continue
|
||||
}
|
||||
seen[fixed] = struct{}{}
|
||||
out = append(out, fixed)
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return errors.New("server is required")
|
||||
}
|
||||
c.Servers = out
|
||||
c.Server = out[0]
|
||||
return nil
|
||||
}
|
||||
|
||||
// normalizeOne applies the scheme rules to a single address.
|
||||
func normalizeOne(addr string, secure bool) (string, error) {
|
||||
switch {
|
||||
case strings.HasPrefix(addr, "https://"):
|
||||
if !secure {
|
||||
return "", fmt.Errorf("%s is an https:// address but no tls certificates are configured", addr)
|
||||
}
|
||||
case strings.HasPrefix(addr, "http://"):
|
||||
if secure {
|
||||
return "", fmt.Errorf("tls is configured but %s is an http:// address; use https://", addr)
|
||||
}
|
||||
case strings.Contains(addr, "://"):
|
||||
return "", fmt.Errorf("server %q must be an http:// or https:// address", addr)
|
||||
case secure:
|
||||
return "https://" + addr, nil
|
||||
default:
|
||||
return "http://" + addr, nil
|
||||
}
|
||||
return addr, nil
|
||||
}
|
||||
|
||||
// Agent runs on a GPU node.
|
||||
type Agent struct {
|
||||
cfg Config
|
||||
@@ -170,6 +214,11 @@ type Agent struct {
|
||||
// report can name both. The control plane keys on the replica.
|
||||
shardReplica map[string]string
|
||||
fetcher *weights.Fetcher
|
||||
|
||||
// lastWeights is the fingerprint of the file list most recently sent, so a
|
||||
// heartbeat carries the list only when it changed. Touched only by the
|
||||
// heartbeat goroutine, which is the only thing that reads or writes it.
|
||||
lastWeights string
|
||||
}
|
||||
|
||||
// connectHTTPClient is the HTTP client connect uses; an interface so tests can
|
||||
@@ -221,15 +270,30 @@ func (a *Agent) Run(ctx context.Context) error {
|
||||
a.recoverOrphans()
|
||||
|
||||
attempt := 0
|
||||
next := 0
|
||||
for ctx.Err() == nil {
|
||||
err := a.session(ctx)
|
||||
server := a.cfg.Servers[next%len(a.cfg.Servers)]
|
||||
err := a.session(ctx, server)
|
||||
if ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
// On to the next address whatever went wrong. A standby refuses, a
|
||||
// stopped control plane refuses the connection, and a crashed one
|
||||
// never answers — from here they are the same thing, and trying
|
||||
// somewhere else is the answer to all three.
|
||||
next++
|
||||
attempt++
|
||||
delay := a.reconnectDelay(attempt)
|
||||
// No wait when there is somewhere else to try: the delay exists to
|
||||
// stop hammering one control plane, and the next address has not been
|
||||
// hammered. With a single address this is the old behaviour exactly.
|
||||
if len(a.cfg.Servers) > 1 && next%len(a.cfg.Servers) != 0 {
|
||||
delay = 0
|
||||
}
|
||||
a.log.Warn("disconnected from the control plane, retrying",
|
||||
"err", err, "in", delay.Round(time.Millisecond), "attempt", attempt)
|
||||
"err", err, "server", server,
|
||||
"next", a.cfg.Servers[next%len(a.cfg.Servers)],
|
||||
"in", delay.Round(time.Millisecond), "attempt", attempt)
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
@@ -460,8 +524,8 @@ func (a *Agent) scan(ctx context.Context) error {
|
||||
}
|
||||
|
||||
// session runs one connection until it breaks.
|
||||
func (a *Agent) session(ctx context.Context) error {
|
||||
client := llmgwv1connect.NewAgentServiceClient(a.http, a.cfg.Server)
|
||||
func (a *Agent) session(ctx context.Context, server string) error {
|
||||
client := llmgwv1connect.NewAgentServiceClient(a.http, server)
|
||||
|
||||
// parent is kept because the derived context is cancelled from inside this
|
||||
// function; asking the derived one whether we are shutting down would
|
||||
@@ -578,9 +642,103 @@ func (a *Agent) heartbeat() *llmgwv1.AgentMessage {
|
||||
for _, g := range inv.GPUs {
|
||||
hb.GpuMemoryUsedMb = append(hb.GpuMemoryUsedMb, int32(g.MemoryUsedMB))
|
||||
}
|
||||
|
||||
// Rescanned here rather than reused from registration. Disk is not a fact
|
||||
// about a machine, it is a fact about a moment: a node that had 200 GB when
|
||||
// its agent started can be full an hour later, and it is precisely the node
|
||||
// busy filling up whose figure matters.
|
||||
for _, st := range discovery.ScanStorage(a.cfg.ModelDirs) {
|
||||
hb.Storage = append(hb.Storage, &llmgwv1.Storage{
|
||||
Path: st.Path, FreeBytes: st.FreeBytes,
|
||||
TotalBytes: st.TotalBytes, Shared: st.Shared,
|
||||
})
|
||||
}
|
||||
|
||||
// The file list is longer and changes rarely, so it goes only when it has.
|
||||
// The flag is what makes an empty list mean "this node now holds nothing"
|
||||
// rather than "nothing to report" — without it, the first heartbeat after
|
||||
// a model was deleted would be indistinguishable from silence.
|
||||
weights := discovery.ScanWeights(a.cfg.ModelDirs)
|
||||
if fp := weightsFingerprint(weights); fp != a.lastWeights {
|
||||
a.lastWeights = fp
|
||||
hb.WeightsChanged = true
|
||||
for _, w := range weights {
|
||||
hb.Weights = append(hb.Weights, &llmgwv1.LocalWeights{
|
||||
Path: w.Path, SizeBytes: w.SizeBytes, Shared: w.Shared,
|
||||
})
|
||||
}
|
||||
}
|
||||
return &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Heartbeat{Heartbeat: hb}}
|
||||
}
|
||||
|
||||
// weightsFingerprint is enough to notice a model appearing, disappearing or
|
||||
// finishing its download.
|
||||
//
|
||||
// Paths and sizes, because a fetch in progress has the same name as the file
|
||||
// it will become — under a different name until it is renamed, but a shared
|
||||
// filesystem can also grow a file in place.
|
||||
func weightsFingerprint(ws []discovery.LocalWeights) string {
|
||||
h := sha256.New()
|
||||
for _, w := range ws {
|
||||
fmt.Fprintf(h, "%s:%d\n", w.Path, w.SizeBytes)
|
||||
}
|
||||
return hex.EncodeToString(h.Sum(nil))
|
||||
}
|
||||
|
||||
// removeWeights deletes cached model files the control plane says are no
|
||||
// longer needed.
|
||||
//
|
||||
// Every path is checked here as well as there, and that duplication is the
|
||||
// point. The control plane computes these from the catalog, but it is the node
|
||||
// that holds the files, and a node must not be talkable into deleting a model
|
||||
// an operator spent a week copying onto it — not by a control plane that is
|
||||
// wrong, not by one that is newer than this agent, not by one that has been
|
||||
// tampered with. So this deletes only what it can prove llmgw itself wrote:
|
||||
// inside its own cache directory, under a name this system generates.
|
||||
//
|
||||
// Failures are logged and skipped rather than reported upward. A file that is
|
||||
// already gone is the expected outcome of a retry, and one that will not
|
||||
// delete is a disk to look at, not a reason to fail a command whose purpose
|
||||
// was to free space.
|
||||
func (a *Agent) removeWeights(paths []string) {
|
||||
cache := a.cfg.WeightsCache
|
||||
if cache == "" && len(a.cfg.ModelDirs) > 0 {
|
||||
cache = a.cfg.ModelDirs[0]
|
||||
}
|
||||
if cache == "" {
|
||||
a.log.Warn("asked to remove weights with no cache directory configured")
|
||||
return
|
||||
}
|
||||
cacheAbs, err := filepath.Abs(cache)
|
||||
if err != nil {
|
||||
a.log.Warn("cannot resolve the cache directory", "dir", cache, "err", err)
|
||||
return
|
||||
}
|
||||
|
||||
for _, p := range paths {
|
||||
abs, err := filepath.Abs(p)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
// Directly inside the cache directory, not merely under it: a nested
|
||||
// path would let ".." or a symlinked subdirectory reach elsewhere, and
|
||||
// nothing this package writes is nested anyway.
|
||||
if filepath.Dir(abs) != cacheAbs {
|
||||
a.log.Warn("refused to remove a file outside the weights cache", "path", p)
|
||||
continue
|
||||
}
|
||||
if !weights.Owned(abs) {
|
||||
a.log.Warn("refused to remove a file llmgw did not fetch", "path", p)
|
||||
continue
|
||||
}
|
||||
if err := os.Remove(abs); err != nil && !os.IsNotExist(err) {
|
||||
a.log.Warn("could not remove cached weights", "path", p, "err", err)
|
||||
continue
|
||||
}
|
||||
a.log.Info("removed cached weights no model asks for", "path", p)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out chan<- *llmgwv1.AgentMessage) {
|
||||
switch m := msg.Msg.(type) {
|
||||
case *llmgwv1.ControlMessage_Start:
|
||||
@@ -591,6 +749,8 @@ func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out cha
|
||||
if err := a.sup.Drain(m.Drain.GetShardId()); err != nil {
|
||||
a.log.Warn("draining failed", "shard", m.Drain.GetShardId(), "err", err)
|
||||
}
|
||||
case *llmgwv1.ControlMessage_RemoveWeights:
|
||||
a.removeWeights(m.RemoveWeights.GetPaths())
|
||||
case *llmgwv1.ControlMessage_Ack:
|
||||
if !m.Ack.GetOk() {
|
||||
a.log.Warn("control plane rejected a message",
|
||||
|
||||
@@ -51,6 +51,9 @@ func Discover(ctx context.Context, opts Options, log *slog.Logger) (Inventory, e
|
||||
// holds has nothing to do with its cards, and a node with none is exactly
|
||||
// where a CPU profile will look for them.
|
||||
inv.Weights = ScanWeights(opts.ModelDirs)
|
||||
// And the disk they sit on. A node can hold a model today and be unable to
|
||||
// take the next one, which the file list alone never shows.
|
||||
inv.Storage = ScanStorage(opts.ModelDirs)
|
||||
|
||||
if !opts.SkipGPUs {
|
||||
switch {
|
||||
|
||||
@@ -114,6 +114,9 @@ type Inventory struct {
|
||||
// Weights are the model files this node can see. Empty means it was not
|
||||
// asked to look.
|
||||
Weights []LocalWeights `json:"weights,omitempty"`
|
||||
// Storage is the disk behind those weights, one entry per filesystem.
|
||||
// Empty means the same thing it means for Weights: not asked.
|
||||
Storage []Storage `json:"storage,omitempty"`
|
||||
}
|
||||
|
||||
// Classes returns the distinct GPU classes present, sorted.
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
// Storage is one filesystem a node keeps weights on.
|
||||
//
|
||||
// One entry per filesystem rather than per configured directory, because free
|
||||
// space belongs to the disk and not to the path. Two model directories on the
|
||||
// same disk are one pool of bytes, and reporting them separately would let the
|
||||
// control plane add 24 GB to 24 GB and place a 40 GB model on a node with 24.
|
||||
type Storage struct {
|
||||
// Path is one directory on this filesystem — the first one configured that
|
||||
// landed here, so the operator recognises it.
|
||||
Path string `json:"path"`
|
||||
// FreeBytes is what a non-root process can still write. Deliberately not
|
||||
// the root reserve: the agent does not run as root, so the reserve is space
|
||||
// it cannot have.
|
||||
FreeBytes int64 `json:"free_bytes"`
|
||||
// TotalBytes is the filesystem's size, so a reading of "8 GB free" can be
|
||||
// told apart on a 10 TB filer and on a 10 GB root partition.
|
||||
TotalBytes int64 `json:"total_bytes"`
|
||||
// Shared marks a filesystem the node did not have to fill. Free space on
|
||||
// one is shared with every other node using it, so two nodes reporting
|
||||
// 500 GB free on the same filer have 500 GB between them, not a terabyte.
|
||||
Shared bool `json:"shared,omitempty"`
|
||||
}
|
||||
|
||||
// ScanStorage reports the filesystems behind the given model directories.
|
||||
//
|
||||
// Empty in, empty out, and a directory that cannot be read is skipped rather
|
||||
// than reported as full: a mount that has not come back yet looks exactly like
|
||||
// a missing directory, and claiming zero free bytes would take the node out of
|
||||
// placement for a reason that is not true. Absent is not zero here, as
|
||||
// everywhere else in this system.
|
||||
func ScanStorage(dirs []string) []Storage {
|
||||
var out []Storage
|
||||
seen := map[uint64]int{} // filesystem id -> index in out
|
||||
|
||||
for _, dir := range dirs {
|
||||
if dir == "" {
|
||||
continue
|
||||
}
|
||||
dev, ok := filesystemID(dir)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if i, dup := seen[dev]; dup {
|
||||
// Same disk, already counted. Keep the first path: it is the one
|
||||
// the operator wrote first and so the one they will recognise.
|
||||
_ = i
|
||||
continue
|
||||
}
|
||||
|
||||
var st syscall.Statfs_t
|
||||
if err := syscall.Statfs(dir, &st); err != nil {
|
||||
continue
|
||||
}
|
||||
// Bsize is int32 on some platforms and int64 on others, which is why
|
||||
// every field here is widened explicitly rather than assumed.
|
||||
block := int64(st.Bsize)
|
||||
seen[dev] = len(out)
|
||||
out = append(out, Storage{
|
||||
Path: dir,
|
||||
FreeBytes: int64(st.Bavail) * block, //nolint:gosec,unconvert // widened on purpose
|
||||
TotalBytes: int64(st.Blocks) * block, //nolint:gosec,unconvert // widened on purpose
|
||||
Shared: isNetworkFS(dir),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// filesystemID identifies the disk a path is on.
|
||||
//
|
||||
// By device number rather than by the statfs filesystem id, because the field
|
||||
// holding that id is spelled differently on Linux and on macOS and this has to
|
||||
// compile on both — the agent runs on one and is developed on the other.
|
||||
func filesystemID(path string) (uint64, bool) {
|
||||
fi, err := os.Stat(path)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
st, ok := fi.Sys().(*syscall.Stat_t)
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
return uint64(st.Dev), true //nolint:gosec,unconvert // int32 on macOS, uint64 on Linux
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package discovery
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestTwoDirectoriesOnOneDiskAreOnePoolOfBytes(t *testing.T) {
|
||||
// The reason this matters: free space belongs to the disk, not to the path.
|
||||
// Counting two directories separately lets the control plane add 24 GB to
|
||||
// 24 GB and place a 40 GB model on a node that has 24.
|
||||
base := t.TempDir()
|
||||
a := filepath.Join(base, "models-a")
|
||||
b := filepath.Join(base, "models-b")
|
||||
for _, d := range []string{a, b} {
|
||||
if err := os.Mkdir(d, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
got := ScanStorage([]string{a, b})
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("two directories on one filesystem reported as %d pools: %+v", len(got), got)
|
||||
}
|
||||
if got[0].Path != a {
|
||||
t.Errorf("kept %q rather than the first directory configured", got[0].Path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStorageReportsSomethingPlausible(t *testing.T) {
|
||||
got := ScanStorage([]string{t.TempDir()})
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("a real directory reported %d filesystems", len(got))
|
||||
}
|
||||
s := got[0]
|
||||
if s.TotalBytes <= 0 {
|
||||
t.Error("a filesystem with no size at all")
|
||||
}
|
||||
if s.FreeBytes < 0 || s.FreeBytes > s.TotalBytes {
|
||||
t.Errorf("free %d of total %d", s.FreeBytes, s.TotalBytes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMissingDirectoryIsNotAFullDisk(t *testing.T) {
|
||||
// A mount that has not come back yet looks exactly like a missing
|
||||
// directory. Reporting it as zero free bytes would take the node out of
|
||||
// placement for a reason that is not true; reporting nothing says only
|
||||
// that nothing is known, which is the honest answer.
|
||||
got := ScanStorage([]string{filepath.Join(t.TempDir(), "never-mounted")})
|
||||
if len(got) != 0 {
|
||||
t.Errorf("a missing directory was reported as a filesystem: %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoDirectoriesConstrainNothing(t *testing.T) {
|
||||
if got := ScanStorage(nil); len(got) != 0 {
|
||||
t.Errorf("scanning nothing found %+v", got)
|
||||
}
|
||||
if got := ScanStorage([]string{""}); len(got) != 0 {
|
||||
t.Errorf("an empty path found %+v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"golang.org/x/net/http2"
|
||||
"golang.org/x/net/http2/h2c"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
type standby bool
|
||||
|
||||
func (s standby) IsLeader() bool { return !bool(s) }
|
||||
|
||||
// startPlain runs an agent service over h2c, which is enough to test who
|
||||
// accepts a connection and who refuses it.
|
||||
//
|
||||
// A nil leader means none was configured, which is every deployment that
|
||||
// existed before leadership did and has to keep accepting agents.
|
||||
func startPlain(t *testing.T, leader agentsvc.Leader) (*inventory.Store, string) {
|
||||
t.Helper()
|
||||
inv := inventory.NewStore()
|
||||
svc := agentsvc.New(inv, newRecorder(), false, slog.New(slog.DiscardHandler))
|
||||
if leader != nil {
|
||||
svc.SetLeadership(leader)
|
||||
}
|
||||
|
||||
path, handler := svc.Handler()
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle(path, handler)
|
||||
// h2c, because the agent speaks HTTP/2 with prior knowledge and a plain
|
||||
// server negotiates HTTP/1.1 — which turns a clear refusal into a frame
|
||||
// parsing error that says nothing about why.
|
||||
srv := httptest.NewServer(h2c.NewHandler(mux, &http2.Server{}))
|
||||
t.Cleanup(srv.Close)
|
||||
return inv, srv.URL
|
||||
}
|
||||
|
||||
func TestAStandbyRefusesAnAgent(t *testing.T) {
|
||||
// Commands travel down the stream the agent opened, so whichever control
|
||||
// plane holds it is the only one that can start anything on that node. A
|
||||
// standby holding it means the leader's placements fail with "no agent
|
||||
// connected" for a node that is plainly connected.
|
||||
_, addr := startPlain(t, standby(true))
|
||||
|
||||
ag := runAgent(t, Config{NodeID: "n1", Server: addr, MgmtAddr: "127.0.0.1"})
|
||||
err := ag.session(context.Background(), ag.cfg.Server)
|
||||
if err == nil {
|
||||
t.Fatal("a standby accepted an agent")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "standby") {
|
||||
t.Errorf("the refusal does not say why: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAgentReachesTheLeaderPastAStandby(t *testing.T) {
|
||||
// The point of the list. The standby is first, so the agent only registers
|
||||
// if it moved on rather than sitting on the refusal.
|
||||
_, standbyAddr := startPlain(t, standby(true))
|
||||
leaderInv, leaderAddr := startPlain(t, standby(false))
|
||||
|
||||
ag := runAgent(t, Config{
|
||||
NodeID: "n1", MgmtAddr: "127.0.0.1",
|
||||
Servers: []string{standbyAddr, leaderAddr},
|
||||
})
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
go ag.Run(ctx)
|
||||
|
||||
deadline := time.Now().Add(8 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if _, ok := leaderInv.Get("n1"); ok {
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("the agent never reached the leader past the standby")
|
||||
}
|
||||
|
||||
func TestALeaderStillAcceptsWithNoLeadershipConfigured(t *testing.T) {
|
||||
// Every deployment before this had no leadership at all, and must keep
|
||||
// accepting agents with nothing configured.
|
||||
inv, addr := startPlain(t, nil)
|
||||
|
||||
ag := runAgent(t, Config{NodeID: "n1", Server: addr, MgmtAddr: "127.0.0.1"})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
go ag.Run(ctx)
|
||||
|
||||
deadline := time.Now().Add(4 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if _, ok := inv.Get("n1"); ok {
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("a service with no leadership configured refused an agent")
|
||||
}
|
||||
@@ -255,7 +255,7 @@ func TestRefusedAgentLearnsWhy(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
err = ag.session(context.Background())
|
||||
err = ag.session(context.Background(), ag.cfg.Server)
|
||||
if err == nil {
|
||||
t.Fatal("a refused session reported success")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
|
||||
)
|
||||
|
||||
// agentWithCache is an agent that only knows where its weights live, which is
|
||||
// all removeWeights needs.
|
||||
func agentWithCache(t *testing.T, dir string) *Agent {
|
||||
t.Helper()
|
||||
return &Agent{
|
||||
cfg: Config{ModelDirs: []string{dir}},
|
||||
log: slog.New(slog.DiscardHandler),
|
||||
}
|
||||
}
|
||||
|
||||
func write(t *testing.T, path string) string {
|
||||
t.Helper()
|
||||
if err := os.WriteFile(path, []byte("weights"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
func TestTheAgentRefusesToDeleteWhatItDidNotFetch(t *testing.T) {
|
||||
// The control plane computes these paths, but the node holds the files. A
|
||||
// control plane that is wrong, newer than this agent, or tampered with must
|
||||
// not be able to talk it into deleting a model an operator spent a week
|
||||
// copying here. So the agent checks again, on its own.
|
||||
dir := t.TempDir()
|
||||
a := agentWithCache(t, dir)
|
||||
|
||||
operators := write(t, filepath.Join(dir, "qwen2.5-coder-14b-instruct-q8_0.gguf"))
|
||||
fetched := write(t, weights.CachePathFor(dir, "https://models.internal/gone.gguf"))
|
||||
|
||||
a.removeWeights([]string{operators, fetched})
|
||||
|
||||
if _, err := os.Stat(operators); err != nil {
|
||||
t.Errorf("a file llmgw never fetched was deleted: %v", err)
|
||||
}
|
||||
if _, err := os.Stat(fetched); !os.IsNotExist(err) {
|
||||
t.Errorf("the fetched file was not removed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheAgentRefusesPathsOutsideItsCache(t *testing.T) {
|
||||
// Even with a name that looks like one of ours. The cache directory is the
|
||||
// boundary, and traversal out of it is the attack this closes.
|
||||
dir := t.TempDir()
|
||||
elsewhere := t.TempDir()
|
||||
a := agentWithCache(t, dir)
|
||||
|
||||
// A correctly named cache file, in the wrong directory.
|
||||
outside := write(t, weights.CachePathFor(elsewhere, "https://models.internal/x.gguf"))
|
||||
// And the same thing reached by traversal from inside the cache.
|
||||
traversal := filepath.Join(dir, "..", filepath.Base(elsewhere),
|
||||
filepath.Base(outside))
|
||||
|
||||
a.removeWeights([]string{outside, traversal})
|
||||
|
||||
if _, err := os.Stat(outside); err != nil {
|
||||
t.Errorf("a file outside the cache was deleted: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemovingAFileThatIsAlreadyGoneIsNotAFailure(t *testing.T) {
|
||||
// The expected outcome of a retry. It must not stop the rest of the batch.
|
||||
dir := t.TempDir()
|
||||
a := agentWithCache(t, dir)
|
||||
|
||||
missing := weights.CachePathFor(dir, "https://models.internal/never.gguf")
|
||||
present := write(t, weights.CachePathFor(dir, "https://models.internal/here.gguf"))
|
||||
|
||||
a.removeWeights([]string{missing, present})
|
||||
|
||||
if _, err := os.Stat(present); !os.IsNotExist(err) {
|
||||
t.Errorf("a missing file stopped the batch: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithNoCacheConfiguredNothingIsDeleted(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
a := &Agent{cfg: Config{}, log: slog.New(slog.DiscardHandler)}
|
||||
kept := write(t, weights.CachePathFor(dir, "https://models.internal/x.gguf"))
|
||||
|
||||
a.removeWeights([]string{kept})
|
||||
|
||||
if _, err := os.Stat(kept); err != nil {
|
||||
t.Errorf("an agent with no cache directory deleted something: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
|
||||
)
|
||||
|
||||
func normalized(t *testing.T, c Config) Config {
|
||||
t.Helper()
|
||||
if err := c.normalizeServerURL(); err != nil {
|
||||
t.Fatalf("normalising %+v: %v", c, err)
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func TestOneServerStillWorksExactlyAsBefore(t *testing.T) {
|
||||
// Every agent configuration written before there could be a second control
|
||||
// plane names one address, and none of them should have to change.
|
||||
c := normalized(t, Config{Server: "10.0.0.140:7443"})
|
||||
if !slices.Equal(c.Servers, []string{"http://10.0.0.140:7443"}) {
|
||||
t.Errorf("servers = %v", c.Servers)
|
||||
}
|
||||
if c.Server != "http://10.0.0.140:7443" {
|
||||
t.Errorf("server = %q", c.Server)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheAddressAlreadyThereStaysPreferred(t *testing.T) {
|
||||
// An operator adding a standby to an existing file expects the control
|
||||
// plane that was already working to keep being tried first.
|
||||
c := normalized(t, Config{
|
||||
Server: "first:7443",
|
||||
Servers: []string{"second:7443", "third:7443"},
|
||||
})
|
||||
want := []string{"http://first:7443", "http://second:7443", "http://third:7443"}
|
||||
if !slices.Equal(c.Servers, want) {
|
||||
t.Errorf("servers = %v, want %v", c.Servers, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestADuplicateAddressIsNotTriedTwiceInARow(t *testing.T) {
|
||||
// Which would look exactly like the rotation not working.
|
||||
c := normalized(t, Config{
|
||||
Server: "a:7443",
|
||||
Servers: []string{"http://a:7443", "b:7443"},
|
||||
})
|
||||
want := []string{"http://a:7443", "http://b:7443"}
|
||||
if !slices.Equal(c.Servers, want) {
|
||||
t.Errorf("servers = %v, want %v", c.Servers, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEveryAddressIsCheckedAgainstTheTLSSetting(t *testing.T) {
|
||||
// The check existed for one address and has to hold for all of them: a
|
||||
// standby reached over plaintext while the leader is authenticated would
|
||||
// be a way around mutual TLS, not a fallback.
|
||||
c := Config{
|
||||
Servers: []string{"https://a:7443", "http://b:7443"},
|
||||
TLS: tlsconf.Config{CA: "ca.pem", Cert: "c.pem", Key: "k.pem"},
|
||||
}
|
||||
err := c.normalizeServerURL()
|
||||
if err == nil {
|
||||
t.Fatal("a plaintext standby was accepted alongside TLS")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "b:7443") {
|
||||
t.Errorf("the error does not say which address: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoServerAtAllIsStillAnError(t *testing.T) {
|
||||
c := Config{}
|
||||
if err := c.normalizeServerURL(); err == nil {
|
||||
t.Error("an agent with no control plane configured was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchemeIsAddedPerAddress(t *testing.T) {
|
||||
c := normalized(t, Config{
|
||||
Servers: []string{"a:7443", "b:7443"},
|
||||
TLS: tlsconf.Config{CA: "ca.pem", Cert: "c.pem", Key: "k.pem"},
|
||||
})
|
||||
for _, s := range c.Servers {
|
||||
if !strings.HasPrefix(s, "https://") {
|
||||
t.Errorf("%q did not get the secure scheme", s)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -15,8 +15,6 @@ package weights
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
@@ -262,12 +260,7 @@ func (f *Fetcher) cachePath(source string) (string, error) {
|
||||
if f.CacheDir == "" {
|
||||
return "", fmt.Errorf("no cache directory configured, so %s cannot be fetched", source)
|
||||
}
|
||||
sum := sha256.Sum256([]byte(source))
|
||||
name := filepath.Base(source)
|
||||
if name == "" || name == "." || name == "/" {
|
||||
name = "weights"
|
||||
}
|
||||
return filepath.Join(f.CacheDir, hex.EncodeToString(sum[:8])+"-"+name), nil
|
||||
return CachePathFor(f.CacheDir, source), nil
|
||||
}
|
||||
|
||||
// isPath reports whether a source names a file on this machine.
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
package weights
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// cacheNamePattern matches a filename this package created.
|
||||
//
|
||||
// The eight-byte hash of the source, hex encoded, then a dash, then the
|
||||
// original filename — exactly what cachePath writes. It is the only mark of
|
||||
// ownership there is, and it has to be exact: everything else in a model
|
||||
// directory was put there by the operator, and deleting one of those files
|
||||
// could cost a week of someone's bandwidth.
|
||||
var cacheNamePattern = regexp.MustCompile(`^[0-9a-f]{16}-`)
|
||||
|
||||
// CachePathFor is where a source would be cached in the given directory.
|
||||
//
|
||||
// Exported so the control plane can work out which files in a node's reported
|
||||
// list are ones llmgw fetched for a model it still wants, without asking the
|
||||
// node anything. The agent and the control plane must agree on this name
|
||||
// exactly, which is why there is one function and not two.
|
||||
func CachePathFor(dir, source string) string {
|
||||
sum := sha256.Sum256([]byte(source))
|
||||
name := filepath.Base(source)
|
||||
if name == "" || name == "." || name == "/" {
|
||||
name = "weights"
|
||||
}
|
||||
return filepath.Join(dir, hex.EncodeToString(sum[:8])+"-"+name)
|
||||
}
|
||||
|
||||
// Owned reports whether a path is one this package created.
|
||||
//
|
||||
// The dividing line of the whole feature. A file without the cache name was
|
||||
// put there by the operator — by scp, by a shared mount, by hand — and llmgw
|
||||
// has no business deleting it however unreferenced it looks. A file with the
|
||||
// cache name is one llmgw downloaded, and llmgw can take it back.
|
||||
func Owned(path string) bool {
|
||||
return cacheNamePattern.MatchString(filepath.Base(path))
|
||||
}
|
||||
|
||||
// Orphans lists cached files that no wanted source accounts for.
|
||||
//
|
||||
// `have` is what a node reports on disk and `wanted` is every source in the
|
||||
// catalog; both are passed in so this is a calculation rather than a scan, and
|
||||
// so it can be run on the control plane against a node's reported inventory.
|
||||
//
|
||||
// Only files this package created are ever returned. A source that is not
|
||||
// fetched — a plain path, or a model id an engine resolves itself — produces
|
||||
// no cache name and so protects nothing and deletes nothing, which is correct:
|
||||
// those files are not in the cache to begin with.
|
||||
func Orphans(have []string, wanted []string) []string {
|
||||
keep := map[string]struct{}{}
|
||||
for _, src := range wanted {
|
||||
if !isRemote(src) {
|
||||
// Nothing was fetched for it, so nothing in the cache belongs to
|
||||
// it. Computing a name here would reserve a file that does not
|
||||
// exist and could never be an orphan anyway.
|
||||
continue
|
||||
}
|
||||
// By base name rather than full path: the control plane does not know
|
||||
// each node's cache directory, and the name is what identifies the
|
||||
// file wherever the node keeps it.
|
||||
keep[filepath.Base(CachePathFor("", src))] = struct{}{}
|
||||
}
|
||||
|
||||
var out []string
|
||||
seen := map[string]struct{}{}
|
||||
for _, path := range have {
|
||||
base := filepath.Base(path)
|
||||
if !Owned(base) {
|
||||
continue
|
||||
}
|
||||
if _, wanted := keep[base]; wanted {
|
||||
continue
|
||||
}
|
||||
if _, dup := seen[path]; dup {
|
||||
continue
|
||||
}
|
||||
seen[path] = struct{}{}
|
||||
out = append(out, path)
|
||||
}
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// PartialName reports whether a path is the temporary file a fetch writes to.
|
||||
//
|
||||
// Those are not orphans in the same sense — a fetch may be running right now —
|
||||
// so they are named separately rather than swept up with the rest. A `.part`
|
||||
// left behind by a killed agent is real garbage, but only its own node can
|
||||
// tell whether it is stale, and nothing here can.
|
||||
func PartialName(path string) bool {
|
||||
return strings.HasSuffix(path, ".part")
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
package weights
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestOnlyFilesLLMGWFetchedAreEverOrphans(t *testing.T) {
|
||||
// The invariant the whole feature rests on. Everything without the cache
|
||||
// name was put there by the operator — scp, a shared mount, a copy by hand
|
||||
// — and deleting one could cost a week of someone's bandwidth. It does not
|
||||
// matter how unreferenced it looks.
|
||||
const dir = "/var/lib/llmgw-models"
|
||||
have := []string{
|
||||
filepath.Join(dir, "qwen2.5-coder-14b-instruct-q8_0-00001-of-00002.gguf"),
|
||||
filepath.Join(dir, "some-model.safetensors"),
|
||||
filepath.Join(dir, "notes.txt"),
|
||||
CachePathFor(dir, "https://models.internal/gone.gguf"),
|
||||
}
|
||||
|
||||
got := Orphans(have, nil)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("expected only the fetched file, got %v", got)
|
||||
}
|
||||
if !Owned(got[0]) {
|
||||
t.Errorf("returned a file llmgw did not fetch: %s", got[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFetchedFileStillInTheCatalogIsKept(t *testing.T) {
|
||||
const dir = "/var/lib/llmgw-models"
|
||||
const wanted = "hf:Qwen/Qwen2.5-Coder-3B-Instruct-GGUF/q4.gguf"
|
||||
const gone = "https://models.internal/old.gguf"
|
||||
|
||||
have := []string{CachePathFor(dir, wanted), CachePathFor(dir, gone)}
|
||||
got := Orphans(have, []string{wanted})
|
||||
|
||||
if slices.Contains(got, CachePathFor(dir, wanted)) {
|
||||
t.Error("a file the catalog still points at was called an orphan")
|
||||
}
|
||||
if !slices.Contains(got, CachePathFor(dir, gone)) {
|
||||
t.Errorf("a file nothing points at was not found: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheCacheNameMatchesWhatAFetchActuallyWrites(t *testing.T) {
|
||||
// The control plane proposes deletions by recomputing this name from the
|
||||
// catalog. If it ever drifted from what a fetch writes, the proposal would
|
||||
// name a file that is in use — so the two must be one function, and this
|
||||
// is what says so.
|
||||
f := &Fetcher{CacheDir: "/cache"}
|
||||
const source = "https://models.internal/x.gguf"
|
||||
|
||||
got, err := f.cachePath(source)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if want := CachePathFor("/cache", source); got != want {
|
||||
t.Errorf("a fetch writes %q but the catalog would look for %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOneOfTwoLookalikeFilesIsNotMistakenForTheOther(t *testing.T) {
|
||||
// That the two cache to different names is covered where the naming lives.
|
||||
// What matters here is the consequence: deleting the unreferenced one must
|
||||
// not take the referenced one with it, and both end in "model.gguf".
|
||||
a := CachePathFor("/cache", "https://one.example/model.gguf")
|
||||
b := CachePathFor("/cache", "https://two.example/model.gguf")
|
||||
|
||||
got := Orphans([]string{a, b}, []string{"https://one.example/model.gguf"})
|
||||
if !slices.Equal(got, []string{b}) {
|
||||
t.Errorf("orphans = %v, want only %s", got, b)
|
||||
}
|
||||
}
|
||||
|
||||
func TestASourceThatIsNeverFetchedProtectsNothingAndDeletesNothing(t *testing.T) {
|
||||
// A plain path, or a model id an engine resolves itself, puts nothing in
|
||||
// the cache. Computing a cache name for it would reserve a file that does
|
||||
// not exist, and it must not stop a real orphan being found either.
|
||||
const dir = "/models"
|
||||
orphan := CachePathFor(dir, "https://models.internal/gone.gguf")
|
||||
|
||||
got := Orphans([]string{orphan, filepath.Join(dir, "local.gguf")},
|
||||
[]string{"/models/local.gguf", "Qwen/Qwen2.5-7B-Instruct"})
|
||||
|
||||
if !slices.Equal(got, []string{orphan}) {
|
||||
t.Errorf("orphans = %v, want only %s", got, orphan)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnEmptyCatalogDoesNotMakeOperatorFilesDisposable(t *testing.T) {
|
||||
// The dangerous case: every model removed from the catalog at once. Files
|
||||
// llmgw fetched become orphans, which is right. Files it did not must stay
|
||||
// untouched, which is the whole point.
|
||||
const dir = "/models"
|
||||
have := []string{
|
||||
filepath.Join(dir, "hand-placed.gguf"),
|
||||
CachePathFor(dir, "https://models.internal/a.gguf"),
|
||||
}
|
||||
got := Orphans(have, nil)
|
||||
if len(got) != 1 || got[0] != CachePathFor(dir, "https://models.internal/a.gguf") {
|
||||
t.Errorf("orphans = %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPartialsAreNotSweptUpWithTheRest(t *testing.T) {
|
||||
// A fetch may be running right now, and only the node itself can tell a
|
||||
// live download from one a killed agent left behind.
|
||||
if !PartialName("/models/abc-model.gguf.part") {
|
||||
t.Error("a .part file was not recognised")
|
||||
}
|
||||
if PartialName("/models/abc-model.gguf") {
|
||||
t.Error("a finished file was called partial")
|
||||
}
|
||||
}
|
||||
+21
-2
@@ -126,6 +126,7 @@ func newCtlCmd() *cobra.Command {
|
||||
newModelCmd(client),
|
||||
newEngineCheckCmd(),
|
||||
newToolCheckCmd(),
|
||||
newWeightsCmd(client, &asJSON),
|
||||
newCertCmd(),
|
||||
)
|
||||
return cmd
|
||||
@@ -149,10 +150,11 @@ func newNodesCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "NODE\tSTATE\tGPUS\tMEM\tCLASSES\tFABRICS\tREPLICAS\tLAST SEEN")
|
||||
fmt.Fprintln(w, "NODE\tSTATE\tGPUS\tMEM\tDISK\tCLASSES\tFABRICS\tREPLICAS\tLAST SEEN")
|
||||
for _, n := range nodes {
|
||||
fmt.Fprintf(w, "%s\t%s\t%d/%d free\t%s\t%s\t%s\t%d\t%s\n",
|
||||
fmt.Fprintf(w, "%s\t%s\t%d/%d free\t%s\t%s\t%s\t%s\t%d\t%s\n",
|
||||
n.ID, n.State, n.GPUFree, n.GPUTotal, memColumn(n.MemoryGB),
|
||||
diskColumn(n.DiskFreeGB, n.DiskTotalGB),
|
||||
join(n.Classes), join(n.Fabrics), n.Replicas, n.LastSeenAgo)
|
||||
}
|
||||
_ = w.Flush()
|
||||
@@ -221,6 +223,23 @@ func memColumn(gb int) string {
|
||||
return fmt.Sprintf("%dG", gb)
|
||||
}
|
||||
|
||||
// diskColumn renders free space against the size it is free out of.
|
||||
//
|
||||
// Both, because they answer different questions and the first is useless
|
||||
// without the second: 8 GB free is comfortable on a 10 TB filer and nearly
|
||||
// full on a 10 GB root partition, and it is the second case that stops a
|
||||
// model from being placed.
|
||||
//
|
||||
// A dash for a node that reported nothing, which an older agent and a node
|
||||
// with no model directories both do. It has to read as "not known" rather
|
||||
// than as "no space", or every unconfigured node looks like it is full.
|
||||
func diskColumn(freeGB, totalGB int) string {
|
||||
if totalGB <= 0 {
|
||||
return "-"
|
||||
}
|
||||
return fmt.Sprintf("%dG/%dG", freeGB, totalGB)
|
||||
}
|
||||
|
||||
func newModelsCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "models",
|
||||
|
||||
@@ -0,0 +1,153 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
// newWeightsCmd reports and reclaims model files nothing needs.
|
||||
//
|
||||
// The dangerous half of storage management, so its shape is defensive: listing
|
||||
// is the default and deleting takes a flag, only files llmgw fetched are ever
|
||||
// candidates, and the node checks every path again before unlinking it.
|
||||
//
|
||||
// What it will never touch is the more important half of the contract. Weights
|
||||
// an operator put on a node — by scp, by a shared mount, by hand — are none of
|
||||
// this system's business however unreferenced they look, and a model that took
|
||||
// a week to copy must not be removable by a command someone ran to tidy up.
|
||||
func newWeightsCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "weights",
|
||||
Short: "Show cached model files no model in the catalog needs",
|
||||
Long: "Lists weights llmgw downloaded for models that are no longer in the\n" +
|
||||
"catalog, per node, with what deleting them would free.\n\n" +
|
||||
"Only files llmgw fetched are considered. Anything an operator put on a\n" +
|
||||
"node is left alone, whatever the catalog says — llmgw observes storage\n" +
|
||||
"rather than owning it, and only takes back what it put there itself.\n\n" +
|
||||
"Nothing is deleted without --collect.",
|
||||
Args: cobra.NoArgs,
|
||||
}
|
||||
var collect bool
|
||||
var node string
|
||||
|
||||
cmd.Flags().BoolVar(&collect, "collect", false, "delete them, rather than only listing")
|
||||
cmd.Flags().StringVar(&node, "node", "", "limit to one node")
|
||||
|
||||
cmd.RunE = func(cmd *cobra.Command, _ []string) error {
|
||||
if collect {
|
||||
return collectOrphans(cmd, client(), node, *asJSON)
|
||||
}
|
||||
return listOrphanWeights(cmd, client(), node, *asJSON)
|
||||
}
|
||||
return cmd
|
||||
}
|
||||
|
||||
func listOrphanWeights(cmd *cobra.Command, c *ctlClient, node string, asJSON bool) error {
|
||||
var found []adminapi.OrphanView
|
||||
if err := c.get(cmd.Context(), "/api/v1/weights/orphans", &found); err != nil {
|
||||
return err
|
||||
}
|
||||
if node != "" {
|
||||
found = onlyNode(found, node)
|
||||
}
|
||||
if asJSON {
|
||||
return emitJSON(found)
|
||||
}
|
||||
if len(found) == 0 {
|
||||
fmt.Fprintln(cmd.OutOrStdout(),
|
||||
"Nothing to reclaim: every file llmgw fetched is still wanted by a model.")
|
||||
return nil
|
||||
}
|
||||
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "NODE\tFILES\tWOULD FREE")
|
||||
var total int64
|
||||
for _, o := range found {
|
||||
fmt.Fprintf(w, "%s\t%d\t%s\n", o.Node, len(o.Paths), fileSize(o.Bytes))
|
||||
total += o.Bytes
|
||||
}
|
||||
_ = w.Flush()
|
||||
|
||||
for _, o := range found {
|
||||
for _, p := range o.Paths {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "\n %s: %s", o.Node, p)
|
||||
}
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "\n\n%s in total. Delete with --collect.\n", fileSize(total))
|
||||
return nil
|
||||
}
|
||||
|
||||
func collectOrphans(cmd *cobra.Command, c *ctlClient, node string, asJSON bool) error {
|
||||
body, err := json.Marshal(map[string]string{"node": node})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var results []struct {
|
||||
Node string `json:"node"`
|
||||
Files int `json:"files"`
|
||||
Bytes int64 `json:"bytes"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
if err := c.post(cmd.Context(), "/api/v1/weights/collect", body, &results); err != nil {
|
||||
return err
|
||||
}
|
||||
if asJSON {
|
||||
return emitJSON(results)
|
||||
}
|
||||
if len(results) == 0 {
|
||||
fmt.Fprintln(cmd.OutOrStdout(), "Nothing to reclaim.")
|
||||
return nil
|
||||
}
|
||||
var freed int64
|
||||
for _, r := range results {
|
||||
if r.Error != "" {
|
||||
// Reported rather than fatal: the other nodes' space is still worth
|
||||
// freeing, and this one will be offered again next time.
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s: %s\n", r.Node, r.Error)
|
||||
continue
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s: asked to remove %d files, %s\n",
|
||||
r.Node, r.Files, fileSize(r.Bytes))
|
||||
freed += r.Bytes
|
||||
}
|
||||
// "asked", not "removed": each node decides, and one that refuses a path it
|
||||
// cannot prove it fetched is the feature working rather than failing.
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "\n%s requested. The nodes verify each path before deleting it.\n",
|
||||
fileSize(freed))
|
||||
return nil
|
||||
}
|
||||
|
||||
func onlyNode(in []adminapi.OrphanView, node string) []adminapi.OrphanView {
|
||||
out := make([]adminapi.OrphanView, 0, len(in))
|
||||
for _, o := range in {
|
||||
if o.Node == node {
|
||||
out = append(out, o)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// fileSize renders a size the way someone reads one when deciding whether a
|
||||
// file is worth deleting.
|
||||
//
|
||||
// Separate from the per-token cache size renderer in modelplan, which never
|
||||
// leaves KiB and MiB. Model files are tens of gigabytes, and "8601.6 MiB" is
|
||||
// not a number anyone weighs a decision against.
|
||||
func fileSize(b int64) string {
|
||||
switch {
|
||||
case b >= 1<<30:
|
||||
return fmt.Sprintf("%.1f GB", float64(b)/(1<<30))
|
||||
case b >= 1<<20:
|
||||
return fmt.Sprintf("%.0f MB", float64(b)/(1<<20))
|
||||
case b > 0:
|
||||
return fmt.Sprintf("%d B", b)
|
||||
default:
|
||||
// The node reported no size. Saying "0 B" would suggest deleting it
|
||||
// gains nothing, which is a different claim from not knowing.
|
||||
return "size unknown"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
)
|
||||
|
||||
func TestTheForwardingDispatcherCarriesEveryCapability(t *testing.T) {
|
||||
// dispatcherRef exists to break a construction cycle, and it satisfies the
|
||||
// dispatcher interface by forwarding each method by hand. That makes it a
|
||||
// place where capabilities go missing silently: the controller reaches
|
||||
// past the interface by type assertion, and an assertion against this type
|
||||
// fails for anything nobody remembered to forward.
|
||||
//
|
||||
// It happened. `ctl weights --collect` returned "this control plane cannot
|
||||
// ask nodes to delete files" against a real cluster, while every unit test
|
||||
// passed — because the test dispatchers do not implement the capability
|
||||
// either, so its absence looked exactly like the expected case.
|
||||
var ref any = &dispatcherRef{}
|
||||
if _, ok := ref.(control.WeightRemover); !ok {
|
||||
t.Error("dispatcherRef drops WeightRemover, so the control plane cannot reclaim disk")
|
||||
}
|
||||
|
||||
// And the thing it forwards to must have it, or the forwarding has nothing
|
||||
// to reach.
|
||||
var svc any = &agentsvc.Service{}
|
||||
if _, ok := svc.(control.WeightRemover); !ok {
|
||||
t.Error("the agent service does not implement WeightRemover")
|
||||
}
|
||||
}
|
||||
@@ -27,7 +27,7 @@ func newTenancyCmd() *cobra.Command {
|
||||
}
|
||||
cmd.PersistentFlags().StringVar(&dbPath, "store", "/var/lib/llmgw/llmgw.db", "path to the store")
|
||||
|
||||
open := func(ctx context.Context) (*store.SQLite, error) { return store.OpenSQLite(ctx, dbPath) }
|
||||
open := func(ctx context.Context) (*store.SQLite, error) { return store.Open(ctx, dbPath) }
|
||||
|
||||
cmd.AddCommand(
|
||||
newKeyInitCmd(open),
|
||||
@@ -35,6 +35,9 @@ func newTenancyCmd() *cobra.Command {
|
||||
newKeyListCmd(open),
|
||||
newKeyRevokeCmd(open),
|
||||
newLimitsCmd(open),
|
||||
newThreadsCmd(open),
|
||||
newMemoryCmd(open),
|
||||
newMoveStoreCmd(),
|
||||
)
|
||||
return cmd
|
||||
}
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
// newMemoryCmd turns note storage on for a project, or off.
|
||||
//
|
||||
// A command rather than a configuration key for the same reason threads are:
|
||||
// keeping what an agent decided to remember is a decision about data, not about
|
||||
// deployment, and it should look like an act rather than a default nobody read.
|
||||
func newMemoryCmd(open opener) *cobra.Command {
|
||||
var (
|
||||
project string
|
||||
hours int
|
||||
off bool
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "memory",
|
||||
Short: "Keep notes an agent writes for itself, across conversations",
|
||||
Long: "With this on, an agent can save and search notes through /v1/memory or the\n" +
|
||||
"MCP tools at /mcp (remember, recall, forget), and find them again in later\n" +
|
||||
"conversations.\n\n" +
|
||||
"Off for every project until this is run. The retention is required: notes\n" +
|
||||
"expire that many hours after they were last written, so a store nobody\n" +
|
||||
"chose a number for does not keep everything forever.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if !off && hours <= 0 {
|
||||
return fmt.Errorf("--retain-hours is required: say how long notes are kept, or pass --off")
|
||||
}
|
||||
s, err := open(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
keep := hours
|
||||
if off {
|
||||
keep = 0
|
||||
}
|
||||
if err := s.SetMemoryRetention(cmd.Context(), project, keep); err != nil {
|
||||
return err
|
||||
}
|
||||
if keep == 0 {
|
||||
fmt.Fprintf(cmd.OutOrStdout(),
|
||||
"%s keeps no notes. Existing ones are removed at the next prune.\n", project)
|
||||
return nil
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(),
|
||||
"%s keeps notes for %d hours after they were last written.\n", project, keep)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&project, "project", "default", "project id")
|
||||
cmd.Flags().IntVar(&hours, "retain-hours", 0, "how long a note is kept after it was last written")
|
||||
cmd.Flags().BoolVar(&off, "off", false, "stop keeping notes for this project")
|
||||
return cmd
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// newMoveStoreCmd copies a store into another one.
|
||||
//
|
||||
// Written because keys are stored hashed. Moving to PostgreSQL without this
|
||||
// means every person and every script needs a new key issued, and the old ones
|
||||
// cannot even be looked up to say which they were.
|
||||
//
|
||||
// It reads and writes, and nothing else runs while it does: run it with the
|
||||
// gateway stopped. A copy taken while requests are landing would miss whatever
|
||||
// arrived after the table it had already passed.
|
||||
func newMoveStoreCmd() *cobra.Command {
|
||||
var from, to string
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "move-store",
|
||||
Short: "Copy a store into another one, for moving SQLite to PostgreSQL",
|
||||
Long: "Copies every table, parents first, into an empty destination.\n\n" +
|
||||
"Stop the gateway first. This makes no attempt to copy a store that is\n" +
|
||||
"being written to: rows landing after the table they belong to has been\n" +
|
||||
"passed would simply be absent, and nothing would say so.\n\n" +
|
||||
"The destination must be empty. Merging two stores means deciding which\n" +
|
||||
"of two projects with the same id was meant, and there is no answer to\n" +
|
||||
"that which is not a guess.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if from == "" || to == "" {
|
||||
return fmt.Errorf("--from and --to are both required")
|
||||
}
|
||||
src, err := store.Open(cmd.Context(), from)
|
||||
if err != nil {
|
||||
return fmt.Errorf("opening the source: %w", err)
|
||||
}
|
||||
defer src.Close()
|
||||
|
||||
dst, err := store.Open(cmd.Context(), to)
|
||||
if err != nil {
|
||||
return fmt.Errorf("opening the destination: %w", err)
|
||||
}
|
||||
defer dst.Close()
|
||||
|
||||
res, err := src.CopyTo(cmd.Context(), dst)
|
||||
// Printed even on failure: a partial copy is the one case where
|
||||
// knowing how far it got decides what to do next.
|
||||
out := cmd.OutOrStdout()
|
||||
var total int64
|
||||
for _, r := range res {
|
||||
if r.Rows > 0 {
|
||||
fmt.Fprintf(out, " %-16s %d\n", r.Table, r.Rows)
|
||||
}
|
||||
total += r.Rows
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(out, "\n%d rows copied. Point the gateway at the new store and start it.\n", total)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&from, "from", "", "source: a SQLite path, a postgres:// URL, or @file holding one")
|
||||
cmd.Flags().StringVar(&to, "to", "", "destination, same forms")
|
||||
return cmd
|
||||
}
|
||||
@@ -28,6 +28,7 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router"
|
||||
storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/version"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/web"
|
||||
)
|
||||
@@ -335,7 +336,29 @@ func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Live
|
||||
if u, ok := sink.(control.UsageSource); ok {
|
||||
ctl.SetUsageSource(u)
|
||||
}
|
||||
// With a shared store, two control planes can see each other, so this one
|
||||
// competes rather than assuming it is alone. Without one there is nobody
|
||||
// to compete with and it places unconditionally, exactly as before.
|
||||
// The same shape as the usage assertion above, and for the same reason:
|
||||
// the sink is the router when the router is serving, and nothing when a
|
||||
// control plane runs on its own.
|
||||
var lead *control.Leadership
|
||||
if h, ok := sink.(interface{ Store() *storepkg.SQLite }); ok {
|
||||
if st := h.Store(); st != nil {
|
||||
lead = control.NewLeadership(st, cfg.Control.LeaseTTL, log)
|
||||
ctl.SetLeadership(lead)
|
||||
log.Info("competing to be the control plane that places work",
|
||||
"owner", lead.Owner())
|
||||
}
|
||||
}
|
||||
svc := agentsvc.New(ctl.Inventory(), ctl.Reconciler(), requireIdentity, log)
|
||||
// A standby must refuse agents as well as refusing to place. Commands
|
||||
// travel down the stream the agent opened, so a standby holding one is a
|
||||
// node the leader cannot command — and the leader's placements would fail
|
||||
// with "no agent connected" for a node that is plainly connected.
|
||||
if lead != nil {
|
||||
svc.SetLeadership(lead)
|
||||
}
|
||||
fwd.set(svc)
|
||||
|
||||
ctl.Start(ctx)
|
||||
@@ -372,6 +395,26 @@ func (r *dispatcherRef) get() (reconciler.Dispatcher, error) {
|
||||
return r.d, nil
|
||||
}
|
||||
|
||||
// RemoveWeights forwards a capability beyond scheduling.
|
||||
//
|
||||
// It has to be spelled out here rather than inherited, because this type
|
||||
// satisfies reconciler.Dispatcher by forwarding each method by hand — so
|
||||
// anything the controller reaches by type assertion stops at this boundary
|
||||
// unless it is forwarded too. A capability added to the service and not to
|
||||
// this type is simply absent at runtime, with nothing at compile time to say
|
||||
// so.
|
||||
func (r *dispatcherRef) RemoveWeights(ctx context.Context, nodeID string, paths []string) error {
|
||||
d, err := r.get()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rm, ok := d.(control.WeightRemover)
|
||||
if !ok {
|
||||
return errors.New("the agent service cannot remove weights")
|
||||
}
|
||||
return rm.RemoveWeights(ctx, nodeID, paths)
|
||||
}
|
||||
|
||||
func (r *dispatcherRef) StartShard(ctx context.Context, cmd reconciler.ShardCommand) error {
|
||||
d, err := r.get()
|
||||
if err != nil {
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
// newThreadsCmd turns conversation storage on for a project, or off.
|
||||
//
|
||||
// A command rather than a configuration key because it is a decision about
|
||||
// data, not about deployment: someone has to choose to keep what a team types,
|
||||
// and choosing it should look like an act rather than like a default nobody
|
||||
// read.
|
||||
//
|
||||
// Until it is run for a project, llmgw stores no message content for that
|
||||
// project — usage rows carry token counts and costs, never text — and that is
|
||||
// how every project starts.
|
||||
func newThreadsCmd(open opener) *cobra.Command {
|
||||
var (
|
||||
project string
|
||||
hours int
|
||||
off bool
|
||||
)
|
||||
|
||||
cmd := &cobra.Command{
|
||||
Use: "threads",
|
||||
Short: "Keep conversations for a project, so a thread can move between models",
|
||||
Long: "With this on, a client sends x-llmgw-thread and only its new turn, and\n" +
|
||||
"llmgw replays the conversation to whichever model serves — including a\n" +
|
||||
"different one from the one it started on.\n\n" +
|
||||
"Off for every project until this is run. llmgw otherwise stores no\n" +
|
||||
"message content anywhere, which is a much easier thing to have deployed\n" +
|
||||
"in front of a team, so turning it on is deliberate.\n\n" +
|
||||
"The retention is required rather than optional: a conversation store with\n" +
|
||||
"no expiry is a transcript of everything a team ever typed, kept forever\n" +
|
||||
"because nobody chose a number.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if !off && hours <= 0 {
|
||||
return fmt.Errorf("--retain-hours is required: say how long conversations are kept, or pass --off")
|
||||
}
|
||||
s, err := open(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.Close()
|
||||
|
||||
keep := hours
|
||||
if off {
|
||||
keep = 0
|
||||
}
|
||||
if err := s.SetThreadRetention(cmd.Context(), project, keep); err != nil {
|
||||
return err
|
||||
}
|
||||
if keep == 0 {
|
||||
fmt.Fprintf(cmd.OutOrStdout(),
|
||||
"%s keeps no conversations. Existing ones are removed at the next prune.\n", project)
|
||||
return nil
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(),
|
||||
"%s keeps conversations for %d hours, measured from the last turn.\n", project, keep)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
cmd.Flags().StringVar(&project, "project", "default", "project id")
|
||||
cmd.Flags().IntVar(&hours, "retain-hours", 0, "how long a conversation is kept after its last turn")
|
||||
cmd.Flags().BoolVar(&off, "off", false, "stop keeping conversations for this project")
|
||||
return cmd
|
||||
}
|
||||
@@ -24,8 +24,16 @@ import (
|
||||
// Found the hard way. qwen2.5-coder produced exactly the right call —
|
||||
// {"name": "list_files", "arguments": {"path": "."}} — as plain reply text,
|
||||
// with tool_calls null, and diagnosing that took an afternoon of eliminating
|
||||
// the gateway, the model size, the chat template and tool_choice in turn. The
|
||||
// answer was the engine build. This runs that elimination in one command.
|
||||
// the gateway, the model size, the chat template and tool_choice in turn.
|
||||
//
|
||||
// The answer was neither the engine nor the model alone, which is why this
|
||||
// command now distinguishes them by what the reply looks like rather than
|
||||
// asserting one. The engine rendered the tool prompt correctly, including the
|
||||
// <tool_call> wrapper its template asks the model to use; the model wrapped
|
||||
// the call in a markdown fence instead, and the parser — which looks for that
|
||||
// wrapper — left it in the content. A differently trained model fixes that. A
|
||||
// bigger one from the same family does not: qwen2.5-coder-7b invented
|
||||
// <function-call> and failed the same way.
|
||||
func newToolCheckCmd() *cobra.Command {
|
||||
var (
|
||||
endpoint string
|
||||
@@ -81,8 +89,17 @@ type ToolProbe struct {
|
||||
// separates "the model ignored its tools" from "the engine did not parse
|
||||
// the call the model made" — the same symptom, different fix.
|
||||
Attempted bool
|
||||
Content string
|
||||
Raw string
|
||||
// Wrapped means the attempt used the <tool_call> tags that the Qwen and
|
||||
// Hermes templates instruct the model to use.
|
||||
//
|
||||
// This is what separates the two halves of a failed attempt. Wrapped and
|
||||
// unparsed is the engine's parser failing on output that was correct.
|
||||
// Unwrapped is the model answering in a shape of its own — the engine had
|
||||
// no way to recognise it, and the fix is a model trained for the format,
|
||||
// not a newer engine.
|
||||
Wrapped bool
|
||||
Content string
|
||||
Raw string
|
||||
}
|
||||
|
||||
func probeToolCall(ctx context.Context, endpoint, model string) (ToolProbe, error) {
|
||||
@@ -162,6 +179,7 @@ func readToolProbe(raw []byte) (ToolProbe, error) {
|
||||
// The model tried, and the engine handed the attempt back as text. The
|
||||
// tell is the tool's name in a reply that was supposed to be a call.
|
||||
p.Attempted = p.Structured || bytes.Contains([]byte(msg.Content), []byte("get_weather"))
|
||||
p.Wrapped = bytes.Contains([]byte(msg.Content), []byte("<tool_call>"))
|
||||
return p, nil
|
||||
}
|
||||
|
||||
@@ -170,12 +188,22 @@ func report(cmd *cobra.Command, p ToolProbe) {
|
||||
switch {
|
||||
case p.Structured:
|
||||
fmt.Fprintln(out, "tool calls work: the engine returned a populated tool_calls array.")
|
||||
case p.Attempted:
|
||||
fmt.Fprintln(out, "the model called the tool and the engine did not parse it.")
|
||||
case p.Attempted && p.Wrapped:
|
||||
fmt.Fprintln(out, "the model called the tool correctly and the engine did not parse it.")
|
||||
fmt.Fprintln(out, "\n The call is in the reply text instead of tool_calls:")
|
||||
fmt.Fprintf(out, " %s\n", trim(p.Content, 200))
|
||||
fmt.Fprintln(out, "\n This is the engine, not the model. Choosing a bigger model will not")
|
||||
fmt.Fprintln(out, " help. Check the engine's build and its chat template support.")
|
||||
fmt.Fprintln(out, "\n It used the <tool_call> wrapper its template asks for, so the model")
|
||||
fmt.Fprintln(out, " did its part. This is the engine's parser. Check its build; another")
|
||||
fmt.Fprintln(out, " model will fail the same way.")
|
||||
case p.Attempted:
|
||||
fmt.Fprintln(out, "the model called the tool in a shape the engine does not recognise.")
|
||||
fmt.Fprintln(out, "\n The call is in the reply text instead of tool_calls:")
|
||||
fmt.Fprintf(out, " %s\n", trim(p.Content, 200))
|
||||
fmt.Fprintln(out, "\n The engine rendered the tools correctly — the model knew the tool")
|
||||
fmt.Fprintln(out, " existed and what to pass it — but did not use the <tool_call> wrapper")
|
||||
fmt.Fprintln(out, " its template asked for, so there was nothing for the parser to find.")
|
||||
fmt.Fprintln(out, "\n Fixed by a model trained for this format, not by a bigger one: the")
|
||||
fmt.Fprintln(out, " instruct models generally manage it where the coder variants do not.")
|
||||
default:
|
||||
fmt.Fprintln(out, "the model did not try to call the tool.")
|
||||
fmt.Fprintf(out, "\n It replied: %s\n", trim(p.Content, 200))
|
||||
|
||||
@@ -19,7 +19,7 @@ func TestAStructuredToolCallIsRecognised(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestACallReturnedAsTextIsBlamedOnTheEngine(t *testing.T) {
|
||||
func TestACallReturnedAsTextIsAnAttempt(t *testing.T) {
|
||||
// The exact reply that cost an afternoon: the model made the right call and
|
||||
// the engine handed it back as prose.
|
||||
p, err := readToolProbe([]byte(`{"choices":[{"message":{"role":"assistant",
|
||||
@@ -33,6 +33,50 @@ func TestACallReturnedAsTextIsBlamedOnTheEngine(t *testing.T) {
|
||||
if !p.Attempted {
|
||||
t.Error("the model plainly tried; that must be distinguished from ignoring the tools")
|
||||
}
|
||||
if p.Wrapped {
|
||||
t.Error("bare JSON was read as having used the <tool_call> wrapper")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAWrappedCallBlamesTheEngineAndABareOneDoesNot(t *testing.T) {
|
||||
// Both look identical from outside — tool_calls null, the call sitting in
|
||||
// the content — and they have opposite fixes. The whole point of the
|
||||
// command is telling them apart, and the earlier version did not: it
|
||||
// blamed the engine for both, which sends someone off rebuilding llama.cpp
|
||||
// when the answer is a different model.
|
||||
wrapped, err := readToolProbe([]byte(`{"choices":[{"message":{"role":"assistant",
|
||||
"content":"<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Istanbul\"}}\n</tool_call>"}}]}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !wrapped.Attempted || !wrapped.Wrapped {
|
||||
t.Errorf("a correctly wrapped call was not recognised as one: %+v", wrapped)
|
||||
}
|
||||
|
||||
// What qwen2.5-coder-3b actually returns, measured: the right call in a
|
||||
// markdown fence, which the parser has no way to find.
|
||||
fenced, err := readToolProbe([]byte(`{"choices":[{"message":{"role":"assistant",
|
||||
"content":"` + "```" + `json\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Istanbul\"}}\n` + "```" + `"}}]}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !fenced.Attempted {
|
||||
t.Error("a fenced call is still an attempt")
|
||||
}
|
||||
if fenced.Wrapped {
|
||||
t.Error("a markdown fence was mistaken for the <tool_call> wrapper")
|
||||
}
|
||||
|
||||
// And what qwen2.5-coder-7b returns: a tag it invented, which is not the
|
||||
// one the template asked for and is not parsed either.
|
||||
invented, err := readToolProbe([]byte(`{"choices":[{"message":{"role":"assistant",
|
||||
"content":"<function-call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Istanbul\"}}\n</function-call>"}}]}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !invented.Attempted || invented.Wrapped {
|
||||
t.Errorf("an invented tag was read as the template's wrapper: %+v", invented)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAModelThatIgnoresItsToolsIsBlamedOnTheModel(t *testing.T) {
|
||||
|
||||
@@ -24,7 +24,7 @@ func newUsersCmd() *cobra.Command {
|
||||
}
|
||||
cmd.PersistentFlags().StringVar(&dbPath, "store", "/var/lib/llmgw/llmgw.db", "path to the store")
|
||||
open := func(cmd *cobra.Command) (*store.SQLite, error) {
|
||||
return store.OpenSQLite(cmd.Context(), dbPath)
|
||||
return store.Open(cmd.Context(), dbPath)
|
||||
}
|
||||
|
||||
cmd.AddCommand(newUserListCmd(open), newUserApproveCmd(open),
|
||||
|
||||
@@ -52,6 +52,11 @@ type Server struct {
|
||||
// Auth turns on API keys on the client-facing port.
|
||||
Auth Auth `koanf:"auth"`
|
||||
|
||||
// Threads keeps conversations and agent memory without API keys, for a
|
||||
// gateway on one trusted machine. With auth on they live in the auth store,
|
||||
// per project, and this block is ignored.
|
||||
Threads Threads `koanf:"threads"`
|
||||
|
||||
// Landing configures the page served at "/".
|
||||
Landing Landing `koanf:"landing"`
|
||||
|
||||
@@ -84,6 +89,22 @@ type Landing struct {
|
||||
Notice string `koanf:"notice"`
|
||||
}
|
||||
|
||||
// Threads configures conversation and memory storage for a gateway without
|
||||
// authentication.
|
||||
//
|
||||
// Everything stored is in one pool with no owner, readable by anyone who can
|
||||
// reach the client port and knows an id. That is acceptable on one machine and
|
||||
// wrong on a shared network, which is why it takes a setting rather than being
|
||||
// the default.
|
||||
type Threads struct {
|
||||
// Store is the SQLite file. Empty means nothing is kept.
|
||||
Store string `koanf:"store"`
|
||||
// RetainHours is how long idle conversations and notes are kept. Required
|
||||
// with Store: a store with no expiry keeps everything anyone typed, by
|
||||
// omission.
|
||||
RetainHours int `koanf:"retain_hours"`
|
||||
}
|
||||
|
||||
// Auth configures client authentication (plan §7.9).
|
||||
//
|
||||
// Off by default, because a gateway on a trusted network with no keys issued
|
||||
@@ -156,6 +177,15 @@ type Control struct {
|
||||
PortRange []int `koanf:"port_range"`
|
||||
// MgmtInterface is the ethernet interface NCCL bootstraps over.
|
||||
MgmtInterface string `koanf:"mgmt_interface"`
|
||||
|
||||
// LeaseTTL is how long this control plane's claim to be the one placing
|
||||
// work stands without being renewed.
|
||||
//
|
||||
// Only meaningful with a shared store, where a second control plane can
|
||||
// see the claim. Long enough to survive a slow query or a brief network
|
||||
// stall, short enough that a standby takes over before anyone has finished
|
||||
// reading the alert. Zero takes the default.
|
||||
LeaseTTL time.Duration `koanf:"lease_ttl"`
|
||||
}
|
||||
|
||||
// Router holds routing behaviour.
|
||||
@@ -391,6 +421,9 @@ func (s Server) Validate() error {
|
||||
if s.Auth.Enabled && s.Auth.Store == "" {
|
||||
return fmt.Errorf("auth.store is required when auth is enabled")
|
||||
}
|
||||
if s.Threads.Store != "" && s.Threads.RetainHours <= 0 {
|
||||
return fmt.Errorf("threads.retain_hours must be positive when threads.store is set")
|
||||
}
|
||||
|
||||
seen := make(map[string]struct{}, len(s.Endpoints))
|
||||
for i, e := range s.Endpoints {
|
||||
|
||||
@@ -40,6 +40,10 @@ func (a *API) Register(mux *http.ServeMux) {
|
||||
mux.HandleFunc("GET /api/v1/nodes", a.listNodes)
|
||||
mux.HandleFunc("GET /api/v1/models", a.listModels)
|
||||
mux.HandleFunc("GET /api/v1/replicas", a.listReplicas)
|
||||
// Listing is a read and deleting is a separate verb, so a mistyped
|
||||
// command cannot remove anything.
|
||||
mux.HandleFunc("GET /api/v1/weights/orphans", a.listOrphans)
|
||||
mux.HandleFunc("POST /api/v1/weights/collect", a.collectWeights)
|
||||
mux.HandleFunc("GET /api/v1/explain-placement", a.explainPlacement)
|
||||
mux.HandleFunc("POST /api/v1/nodes/{id}/state", a.setNodeState)
|
||||
}
|
||||
@@ -62,9 +66,21 @@ type NodeView struct {
|
||||
MemoryGB int `json:"memory_gb,omitempty"`
|
||||
// Models are the weights this node can see. Absent means it was not asked
|
||||
// to look, which is not the same as holding none.
|
||||
Models []string `json:"models_on_disk,omitempty"`
|
||||
Classes []string `json:"gpu_classes,omitempty"`
|
||||
Fabrics []string `json:"fabrics,omitempty"`
|
||||
Models []string `json:"models_on_disk,omitempty"`
|
||||
// DiskFreeGB and DiskTotalGB are the filesystems holding those weights,
|
||||
// summed. Zero means the agent did not report any, which an older agent
|
||||
// and a node with no model directories both do — so it reads as unknown
|
||||
// rather than as a full disk.
|
||||
//
|
||||
// Summed because an operator asking "can this node take another model"
|
||||
// wants one number; the per-filesystem detail is in the node's own record
|
||||
// for anyone who needs it. Shared filesystems are left out of the sum: the
|
||||
// space on a filer is not this node's to spend, and adding it would let a
|
||||
// fleet of twenty nodes each believe it has the same free terabyte.
|
||||
DiskFreeGB int `json:"disk_free_gb,omitempty"`
|
||||
DiskTotalGB int `json:"disk_total_gb,omitempty"`
|
||||
Classes []string `json:"gpu_classes,omitempty"`
|
||||
Fabrics []string `json:"fabrics,omitempty"`
|
||||
|
||||
// Unhealthy lists GPUs kept out of placement and why, which is usually the
|
||||
// answer when a node looks idle but nothing will schedule onto it.
|
||||
@@ -166,6 +182,10 @@ func (a *API) Nodes() []NodeView {
|
||||
Fabrics: n.FabricNames(), AgentVersion: n.AgentVersion,
|
||||
Replicas: replicasOn[n.ID],
|
||||
}
|
||||
if free, total := n.LocalDisk(); total > 0 {
|
||||
v.DiskFreeGB = int(free / (1 << 30))
|
||||
v.DiskTotalGB = int(total / (1 << 30))
|
||||
}
|
||||
for _, w := range n.Weights {
|
||||
v.Models = append(v.Models, w.Path)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"sort"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
|
||||
)
|
||||
|
||||
// OrphanView is cached weights on one node that no catalog model accounts for.
|
||||
type OrphanView struct {
|
||||
Node string `json:"node"`
|
||||
Paths []string `json:"paths"`
|
||||
// Bytes is what deleting them would free, as far as the node reported. A
|
||||
// file whose size was not reported counts as zero rather than being left
|
||||
// out, so the list and the number always describe the same set.
|
||||
Bytes int64 `json:"bytes"`
|
||||
}
|
||||
|
||||
// Orphans finds cached weights nothing in the catalog points at.
|
||||
//
|
||||
// Computed here from what nodes already report, so listing costs no round trip
|
||||
// and a node that is offline simply contributes nothing — its files are not
|
||||
// deletable while it is gone anyway.
|
||||
//
|
||||
// Only files llmgw fetched are ever considered. Everything else in a model
|
||||
// directory was put there by the operator and is none of this system's
|
||||
// business, however unreferenced it looks.
|
||||
func (a *API) Orphans() []OrphanView {
|
||||
var wanted []string
|
||||
for _, m := range a.ctl.Catalog().Models() {
|
||||
wanted = append(wanted, m.Source)
|
||||
}
|
||||
|
||||
out := make([]OrphanView, 0)
|
||||
for _, n := range a.ctl.Inventory().Nodes() {
|
||||
size := map[string]int64{}
|
||||
have := make([]string, 0, len(n.Weights))
|
||||
for _, w := range n.Weights {
|
||||
have = append(have, w.Path)
|
||||
size[w.Path] = w.SizeBytes
|
||||
}
|
||||
|
||||
paths := weights.Orphans(have, wanted)
|
||||
if len(paths) == 0 {
|
||||
continue
|
||||
}
|
||||
v := OrphanView{Node: n.ID, Paths: paths}
|
||||
for _, p := range paths {
|
||||
v.Bytes += size[p]
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].Node < out[j].Node })
|
||||
return out
|
||||
}
|
||||
|
||||
func (a *API) listOrphans(w http.ResponseWriter, _ *http.Request) {
|
||||
writeJSON(w, http.StatusOK, a.Orphans())
|
||||
}
|
||||
|
||||
// collectWeights deletes them.
|
||||
//
|
||||
// A proposal, not an instruction: each node checks every path again before
|
||||
// unlinking it, and refuses anything it cannot prove it fetched. The two
|
||||
// checks are deliberate duplication — the control plane knows the catalog and
|
||||
// the node knows its disk, and only the node can be sure.
|
||||
func (a *API) collectWeights(w http.ResponseWriter, r *http.Request) {
|
||||
rm := a.ctl.WeightRemover()
|
||||
if rm == nil {
|
||||
writeError(w, http.StatusNotImplemented,
|
||||
"this control plane cannot ask nodes to delete files")
|
||||
return
|
||||
}
|
||||
|
||||
var body struct {
|
||||
// Node limits it to one node. Empty means every node, which is the
|
||||
// common case after removing a model from the catalog.
|
||||
Node string `json:"node,omitempty"`
|
||||
}
|
||||
// An empty body is the "all nodes" case and must not be an error.
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
|
||||
type result struct {
|
||||
Node string `json:"node"`
|
||||
Files int `json:"files"`
|
||||
Bytes int64 `json:"bytes"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
out := make([]result, 0)
|
||||
for _, o := range a.Orphans() {
|
||||
if body.Node != "" && o.Node != body.Node {
|
||||
continue
|
||||
}
|
||||
res := result{Node: o.Node, Files: len(o.Paths), Bytes: o.Bytes}
|
||||
// A node that has gone since the list was built is reported rather than
|
||||
// failing the whole sweep: the other nodes' space is still worth
|
||||
// freeing, and this one will be offered again next time.
|
||||
if err := rm.RemoveWeights(r.Context(), o.Node, o.Paths); err != nil {
|
||||
res.Error = err.Error()
|
||||
res.Files, res.Bytes = 0, 0
|
||||
}
|
||||
out = append(out, res)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
@@ -30,6 +30,14 @@ func NodeFromRegister(r *llmgwv1.Register) inventory.Node {
|
||||
Path: w.GetPath(), SizeBytes: w.GetSizeBytes(), Shared: w.GetShared(),
|
||||
})
|
||||
}
|
||||
for _, s := range r.GetStorage() {
|
||||
n.Storage = append(n.Storage, discovery.Storage{
|
||||
Path: s.GetPath(),
|
||||
FreeBytes: s.GetFreeBytes(),
|
||||
TotalBytes: s.GetTotalBytes(),
|
||||
Shared: s.GetShared(),
|
||||
})
|
||||
}
|
||||
for _, g := range r.GetGpus() {
|
||||
n.GPUs = append(n.GPUs, discovery.GPU{
|
||||
Index: int(g.GetIndex()),
|
||||
@@ -75,6 +83,14 @@ func RegisterFromInventory(nodeID, hostname, mgmtAddr, agentVersion string,
|
||||
Path: w.Path, SizeBytes: w.SizeBytes, Shared: w.Shared,
|
||||
})
|
||||
}
|
||||
for _, s := range inv.Storage {
|
||||
r.Storage = append(r.Storage, &llmgwv1.Storage{
|
||||
Path: s.Path,
|
||||
FreeBytes: s.FreeBytes,
|
||||
TotalBytes: s.TotalBytes,
|
||||
Shared: s.Shared,
|
||||
})
|
||||
}
|
||||
for _, g := range inv.GPUs {
|
||||
r.Gpus = append(r.Gpus, &llmgwv1.GPU{
|
||||
Index: int32(g.Index),
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
func nodeWith(s *Service, n inventory.Node) { s.inv.Upsert(n) }
|
||||
|
||||
func newTestService() *Service {
|
||||
return &Service{inv: inventory.NewStore(), conns: map[string]*conn{}}
|
||||
}
|
||||
|
||||
func TestDiskIsRefreshedWhileTheAgentRuns(t *testing.T) {
|
||||
// It was sampled once, at registration, and never again. Found on the live
|
||||
// cluster: an agent running for two days still reported the disk it had
|
||||
// when it started, and a file deleted in between was still being offered
|
||||
// for deletion. Disk is a fact about a moment, and the node whose figure
|
||||
// matters most is the one busy filling up.
|
||||
s := newTestService()
|
||||
nodeWith(s, inventory.Node{
|
||||
ID: "n1",
|
||||
State: inventory.Ready,
|
||||
Storage: []discovery.Storage{
|
||||
{Path: "/models", FreeBytes: 200 << 30, TotalBytes: 400 << 30},
|
||||
},
|
||||
})
|
||||
|
||||
s.applyHeartbeat("n1", &llmgwv1.Heartbeat{
|
||||
NodeId: "n1",
|
||||
Storage: []*llmgwv1.Storage{
|
||||
{Path: "/models", FreeBytes: 5 << 30, TotalBytes: 400 << 30},
|
||||
},
|
||||
})
|
||||
|
||||
n, _ := s.inv.Get("n1")
|
||||
free, total := n.LocalDisk()
|
||||
if free != 5<<30 || total != 400<<30 {
|
||||
t.Errorf("disk is %d of %d, want the figure the heartbeat carried", free, total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnEmptyWeightListOnlyCountsWhenTheAgentMeansIt(t *testing.T) {
|
||||
// Most heartbeats carry no file list because nothing changed. Treating
|
||||
// that as "this node holds nothing" would empty a node's inventory five
|
||||
// seconds after it registered, and placement would stop preferring the
|
||||
// node that actually has the model.
|
||||
s := newTestService()
|
||||
nodeWith(s, inventory.Node{
|
||||
ID: "n1",
|
||||
State: inventory.Ready,
|
||||
Weights: []discovery.LocalWeights{{Path: "/models/a.gguf", SizeBytes: 100}},
|
||||
})
|
||||
|
||||
// Silence about weights: the list must survive.
|
||||
s.applyHeartbeat("n1", &llmgwv1.Heartbeat{
|
||||
NodeId: "n1",
|
||||
Storage: []*llmgwv1.Storage{{Path: "/models", FreeBytes: 1, TotalBytes: 2}},
|
||||
})
|
||||
if n, _ := s.inv.Get("n1"); len(n.Weights) != 1 {
|
||||
t.Fatalf("a heartbeat with no weight list emptied the inventory: %+v", n.Weights)
|
||||
}
|
||||
|
||||
// And when the agent says it means an empty list, it is taken.
|
||||
s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1", WeightsChanged: true})
|
||||
if n, _ := s.inv.Get("n1"); len(n.Weights) != 0 {
|
||||
t.Errorf("a deliberate empty list was ignored: %+v", n.Weights)
|
||||
}
|
||||
}
|
||||
|
||||
func TestANewlyFetchedModelAppearsWithoutARestart(t *testing.T) {
|
||||
// The other direction of the same bug: a model fetched after registration
|
||||
// was invisible until the agent restarted, so placement could not prefer
|
||||
// the node that had just spent ten minutes downloading it.
|
||||
s := newTestService()
|
||||
nodeWith(s, inventory.Node{ID: "n1", State: inventory.Ready})
|
||||
|
||||
s.applyHeartbeat("n1", &llmgwv1.Heartbeat{
|
||||
NodeId: "n1",
|
||||
WeightsChanged: true,
|
||||
Weights: []*llmgwv1.LocalWeights{
|
||||
{Path: "/models/new.gguf", SizeBytes: 4 << 30},
|
||||
},
|
||||
})
|
||||
|
||||
n, _ := s.inv.Get("n1")
|
||||
if !n.HasWeights("/models/new.gguf") {
|
||||
t.Errorf("a model fetched after registration did not appear: %+v", n.Weights)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAHeartbeatAboutNothingChangesNothing(t *testing.T) {
|
||||
s := newTestService()
|
||||
before := inventory.Node{
|
||||
ID: "n1",
|
||||
State: inventory.Ready,
|
||||
Weights: []discovery.LocalWeights{{Path: "/models/a.gguf", SizeBytes: 100}},
|
||||
Storage: []discovery.Storage{{Path: "/models", FreeBytes: 7, TotalBytes: 9}},
|
||||
}
|
||||
nodeWith(s, before)
|
||||
|
||||
s.applyHeartbeat("n1", &llmgwv1.Heartbeat{NodeId: "n1"})
|
||||
|
||||
n, _ := s.inv.Get("n1")
|
||||
if len(n.Weights) != 1 {
|
||||
t.Errorf("weights changed: %+v", n.Weights)
|
||||
}
|
||||
if free, _ := n.LocalDisk(); free != 7 {
|
||||
t.Errorf("disk changed: %d", free)
|
||||
}
|
||||
}
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1/llmgwv1connect"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
@@ -30,6 +31,9 @@ type StatusSink interface {
|
||||
// time the reconciler's Dispatcher: it holds the open streams, so it is the
|
||||
// only thing that can deliver a command to a node.
|
||||
type Service struct {
|
||||
// leader gates whether this process accepts agent streams at all.
|
||||
leader Leader
|
||||
|
||||
inv *inventory.Store
|
||||
status StatusSink
|
||||
log *slog.Logger
|
||||
@@ -85,6 +89,19 @@ func (s *Service) Connect(
|
||||
ctx context.Context,
|
||||
stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage],
|
||||
) error {
|
||||
// A standby must not hold an agent's stream. Commands travel down the
|
||||
// stream the agent opened, so whichever control plane holds it is the only
|
||||
// one that can start or stop anything on that node — and if that is not
|
||||
// the leader, the leader's placements fail with "no agent connected" for a
|
||||
// node that is plainly connected.
|
||||
//
|
||||
// Refused before the opening message is read, so the agent spends nothing
|
||||
// on a connection that cannot serve it.
|
||||
if s.leader != nil && !s.leader.IsLeader() {
|
||||
return connect.NewError(connect.CodeUnavailable,
|
||||
errors.New("this control plane is a standby; connect to the leader"))
|
||||
}
|
||||
|
||||
first, err := stream.Receive()
|
||||
if err != nil {
|
||||
return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("no opening message: %w", err))
|
||||
@@ -239,7 +256,8 @@ func (s *Service) handle(nodeID string, msg *llmgwv1.AgentMessage) {
|
||||
func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) {
|
||||
used := hb.GetGpuMemoryUsedMb()
|
||||
problems := hb.GetNewProblems()
|
||||
if len(used) == 0 && len(problems) == 0 {
|
||||
storage := hb.GetStorage()
|
||||
if len(used) == 0 && len(problems) == 0 && len(storage) == 0 && !hb.GetWeightsChanged() {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -259,6 +277,33 @@ func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) {
|
||||
// conservative direction.
|
||||
n.GPUs[0].Health.Problems = append(n.GPUs[0].Health.Problems, problems...)
|
||||
}
|
||||
|
||||
// Disk is a fact about a moment, not about a machine. Sampled only at
|
||||
// registration it is wrong on exactly the node worth watching — the one
|
||||
// busy filling up — and it stays wrong until that agent restarts.
|
||||
if len(storage) > 0 {
|
||||
n.Storage = n.Storage[:0]
|
||||
for _, st := range storage {
|
||||
n.Storage = append(n.Storage, discovery.Storage{
|
||||
Path: st.GetPath(),
|
||||
FreeBytes: st.GetFreeBytes(),
|
||||
TotalBytes: st.GetTotalBytes(),
|
||||
Shared: st.GetShared(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Only when the agent says the list is meant. An empty list without the
|
||||
// flag is silence, and treating it as "this node holds nothing" would
|
||||
// empty a node's inventory on every heartbeat.
|
||||
if hb.GetWeightsChanged() {
|
||||
n.Weights = n.Weights[:0]
|
||||
for _, w := range hb.GetWeights() {
|
||||
n.Weights = append(n.Weights, discovery.LocalWeights{
|
||||
Path: w.GetPath(), SizeBytes: w.GetSizeBytes(), Shared: w.GetShared(),
|
||||
})
|
||||
}
|
||||
}
|
||||
s.inv.Upsert(n)
|
||||
}
|
||||
|
||||
@@ -316,6 +361,16 @@ func (s *Service) ConnectedNodes() []string {
|
||||
|
||||
// --- Dispatcher ---
|
||||
|
||||
// Leader reports whether this process is the one that should be holding agent
|
||||
// streams.
|
||||
type Leader interface{ IsLeader() bool }
|
||||
|
||||
// SetLeadership makes this service refuse agents while it is a standby.
|
||||
//
|
||||
// Nil accepts them unconditionally, which is what a single control plane wants
|
||||
// and what every deployment before leadership had.
|
||||
func (s *Service) SetLeadership(l Leader) { s.leader = l }
|
||||
|
||||
// ErrNotConnected means no agent stream is open for that node.
|
||||
var ErrNotConnected = errors.New("no agent connected")
|
||||
|
||||
@@ -340,6 +395,22 @@ func (s *Service) DrainShard(_ context.Context, nodeID, shardID string) error {
|
||||
})
|
||||
}
|
||||
|
||||
// RemoveWeights asks a node to delete cached files no model needs.
|
||||
//
|
||||
// The node checks every path again before unlinking it, so this is a proposal
|
||||
// rather than an instruction — which is the right shape for a message that
|
||||
// deletes data the control plane cannot see.
|
||||
func (s *Service) RemoveWeights(_ context.Context, nodeID string, paths []string) error {
|
||||
if len(paths) == 0 {
|
||||
return nil
|
||||
}
|
||||
return s.send(nodeID, &llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_RemoveWeights{
|
||||
RemoveWeights: &llmgwv1.RemoveWeights{Paths: paths},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Service) send(nodeID string, msg *llmgwv1.ControlMessage) error {
|
||||
s.mu.RLock()
|
||||
c, ok := s.conns[nodeID]
|
||||
|
||||
@@ -71,6 +71,20 @@ type Profile struct {
|
||||
// the control plane.
|
||||
MinMemoryGB int `json:"min_memory_gb,omitempty" koanf:"min_memory_gb"`
|
||||
|
||||
// MinDiskGB is the free space a node needs before it can hold these
|
||||
// weights.
|
||||
//
|
||||
// Only checked when the node does not already have them: a node holding
|
||||
// the file needs no room for it, and requiring the space twice would keep
|
||||
// a model off the one node best placed to run it.
|
||||
//
|
||||
// Declared rather than derived, because the size of a model that has never
|
||||
// been fetched anywhere is not knowable from here — and a wrong guess
|
||||
// fails late, halfway through a download, which is the failure this exists
|
||||
// to prevent. Zero means unconstrained, and a node that reports no storage
|
||||
// is not excluded: unknown is not full.
|
||||
MinDiskGB int `json:"min_disk_gb,omitempty" koanf:"min_disk_gb"`
|
||||
|
||||
// NodeSelector restricts this profile to nodes carrying every one of these
|
||||
// labels, with these values.
|
||||
//
|
||||
@@ -344,6 +358,7 @@ func (m Model) PlacementProfiles() []placement.Profile {
|
||||
Priority: p.Priority,
|
||||
NodeSelector: p.NodeSelector,
|
||||
MinMemoryGB: p.MinMemoryGB,
|
||||
MinDiskGB: p.MinDiskGB,
|
||||
// The model's source, so placement can prefer nodes that already
|
||||
// hold the weights rather than one that must first move them.
|
||||
Source: m.Source,
|
||||
|
||||
@@ -75,6 +75,8 @@ type Controller struct {
|
||||
inv *inventory.Store
|
||||
cat *catalog.Live
|
||||
rec *reconciler.Reconciler
|
||||
lead *Leadership
|
||||
disp reconciler.Dispatcher
|
||||
sink EndpointSink
|
||||
usage UsageSource
|
||||
|
||||
@@ -101,6 +103,7 @@ func New(cfg Config, cat *catalog.Live, disp reconciler.Dispatcher, sink Endpoin
|
||||
inv: inv,
|
||||
cat: cat,
|
||||
rec: reconciler.New(cfg.Reconciler, cat, inv, disp, log),
|
||||
disp: disp,
|
||||
sink: sink,
|
||||
stop: make(chan struct{}),
|
||||
done: make(chan struct{}),
|
||||
@@ -170,8 +173,42 @@ func (c *Controller) Reconciler() *reconciler.Reconciler { return c.rec }
|
||||
// Catalog exposes the model catalog.
|
||||
func (c *Controller) Catalog() *catalog.Catalog { return c.cat.Load() }
|
||||
|
||||
// SetLeadership makes this control plane compete for the right to place work,
|
||||
// instead of assuming it is the only one.
|
||||
//
|
||||
// Wired after construction because deciding leadership needs a store, and the
|
||||
// store is opened by the router. Left unset, this process places
|
||||
// unconditionally — which is what a single control plane wants, and what every
|
||||
// deployment that existed before this had.
|
||||
func (c *Controller) SetLeadership(l *Leadership) {
|
||||
c.lead = l
|
||||
c.rec.SetLeadership(l)
|
||||
}
|
||||
|
||||
// IsLeader reports whether this process is placing work. True when leadership
|
||||
// was never configured, because then there is nobody to compete with.
|
||||
func (c *Controller) IsLeader() bool { return c.lead == nil || c.lead.IsLeader() }
|
||||
|
||||
// WeightRemover can ask a node to delete cached model files.
|
||||
//
|
||||
// Separate from Dispatcher because it is not part of scheduling and not every
|
||||
// dispatcher has it — the ones tests use do not. A caller gets nil rather than
|
||||
// a method that silently does nothing.
|
||||
type WeightRemover interface {
|
||||
RemoveWeights(ctx context.Context, nodeID string, paths []string) error
|
||||
}
|
||||
|
||||
// WeightRemover returns the dispatcher's deletion capability, or nil.
|
||||
func (c *Controller) WeightRemover() WeightRemover {
|
||||
r, _ := c.disp.(WeightRemover)
|
||||
return r
|
||||
}
|
||||
|
||||
// Start runs the reconciler and begins publishing endpoints.
|
||||
func (c *Controller) Start(ctx context.Context) {
|
||||
if c.lead != nil {
|
||||
go c.lead.Run(ctx)
|
||||
}
|
||||
c.rec.Start(ctx)
|
||||
|
||||
go func() {
|
||||
|
||||
@@ -45,6 +45,8 @@ type Collector struct {
|
||||
nodeGPUsFree *prometheus.Desc
|
||||
nodeMemory *prometheus.Desc
|
||||
nodeSeen *prometheus.Desc
|
||||
nodeDiskFree *prometheus.Desc
|
||||
nodeDiskSize *prometheus.Desc
|
||||
|
||||
gpuMemoryTotal *prometheus.Desc
|
||||
gpuMemoryUsed *prometheus.Desc
|
||||
@@ -76,6 +78,15 @@ func New(src Source) *Collector {
|
||||
"A node's total RAM, 0 when it does not report one.", node, nil),
|
||||
nodeSeen: prometheus.NewDesc("llmgw_node_last_seen_seconds",
|
||||
"Seconds since a node last reported. Rising means an agent has gone quiet.", node, nil),
|
||||
// Only the disk a node fills itself: space on a shared filer is not
|
||||
// this node's to spend, and summing it per node would have a fleet of
|
||||
// twenty each report the same free terabyte. Absent rather than zero
|
||||
// when nothing was reported, so an alert on "free is low" does not
|
||||
// fire for every node that was never asked about its disk.
|
||||
nodeDiskFree: prometheus.NewDesc("llmgw_node_disk_free_bytes",
|
||||
"Free space on the local filesystems holding model weights. Absent when the node reports no storage.", node, nil),
|
||||
nodeDiskSize: prometheus.NewDesc("llmgw_node_disk_total_bytes",
|
||||
"Size of those filesystems, so free space can be read as a fraction.", node, nil),
|
||||
|
||||
gpuMemoryTotal: prometheus.NewDesc("llmgw_gpu_memory_total_bytes",
|
||||
"A GPU's framebuffer size.", gpu, nil),
|
||||
@@ -111,6 +122,7 @@ func New(src Source) *Collector {
|
||||
func (c *Collector) Describe(ch chan<- *prometheus.Desc) {
|
||||
for _, d := range []*prometheus.Desc{
|
||||
c.nodeUp, c.nodeGPUs, c.nodeGPUsFree, c.nodeMemory, c.nodeSeen,
|
||||
c.nodeDiskFree, c.nodeDiskSize,
|
||||
c.gpuMemoryTotal, c.gpuMemoryUsed, c.gpuUtil, c.gpuTemp, c.gpuHealthy,
|
||||
c.gpuNVLink, c.gpuECC,
|
||||
c.replicaState, c.replicaReady, c.replicaIdle,
|
||||
@@ -143,6 +155,15 @@ func (c *Collector) Collect(ch chan<- prometheus.Metric) {
|
||||
float64(n.MemoryTotalMB)*1024*1024, n.ID)
|
||||
ch <- prometheus.MustNewConstMetric(c.nodeSeen, prometheus.GaugeValue,
|
||||
now.Sub(n.LastSeen).Seconds(), n.ID)
|
||||
// Nothing reported means nothing is known, so no sample at all. A zero
|
||||
// here would read as a full disk and page someone about a node that
|
||||
// was simply never asked.
|
||||
if free, total := n.LocalDisk(); total > 0 {
|
||||
ch <- prometheus.MustNewConstMetric(c.nodeDiskFree, prometheus.GaugeValue,
|
||||
float64(free), n.ID)
|
||||
ch <- prometheus.MustNewConstMetric(c.nodeDiskSize, prometheus.GaugeValue,
|
||||
float64(total), n.ID)
|
||||
}
|
||||
|
||||
byClass := map[string]int{}
|
||||
for _, g := range n.GPUs {
|
||||
|
||||
@@ -84,6 +84,12 @@ type Node struct {
|
||||
// Empty means the node was not asked to look, which is not the same as
|
||||
// holding none — so it constrains nothing.
|
||||
Weights []discovery.LocalWeights `json:"weights,omitempty"`
|
||||
// Storage is the disk those weights sit on, one entry per filesystem.
|
||||
//
|
||||
// Empty means the node was not asked, not that it has no disk — so an
|
||||
// agent too old to report it, or one with no model directories, is left
|
||||
// unconstrained rather than treated as full.
|
||||
Storage []discovery.Storage `json:"storage,omitempty"`
|
||||
|
||||
GPUs []discovery.GPU `json:"gpus"`
|
||||
// Fabrics is empty on a node without IB. Such a node is its own
|
||||
@@ -310,6 +316,28 @@ func (s *Store) FreeGPUs(nodeID string) []int {
|
||||
return out
|
||||
}
|
||||
|
||||
// LocalDisk sums the space on filesystems this node fills itself.
|
||||
//
|
||||
// Shared filesystems are left out on purpose. Space on a filer belongs to
|
||||
// everyone mounting it, and counting it per node would let twenty nodes each
|
||||
// believe they have the same free terabyte — then place twenty models into it.
|
||||
// A node whose weights are all on a shared mount therefore reports no local
|
||||
// disk, which is correct: it has none to spend.
|
||||
//
|
||||
// Both zero means nothing is known — an agent too old to report storage, or a
|
||||
// node with no model directories — and callers must read that as unknown
|
||||
// rather than as a full disk.
|
||||
func (n Node) LocalDisk() (free, total int64) {
|
||||
for _, s := range n.Storage {
|
||||
if s.Shared {
|
||||
continue
|
||||
}
|
||||
free += s.FreeBytes
|
||||
total += s.TotalBytes
|
||||
}
|
||||
return free, total
|
||||
}
|
||||
|
||||
// HasWeights reports whether this node can see the given source.
|
||||
//
|
||||
// An exact path match, deliberately. A source that is not a path — a model id
|
||||
|
||||
@@ -248,3 +248,42 @@ func TestConcurrentAccess(t *testing.T) {
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestSharedFilesystemsAreNotThisNodesSpaceToSpend(t *testing.T) {
|
||||
// Twenty nodes mounting one filer have 500 GB between them, not ten
|
||||
// terabytes. Counting a shared mount per node is how a scheduler places
|
||||
// twenty models into space for one.
|
||||
n := Node{Storage: []discovery.Storage{
|
||||
{Path: "/mnt/filer", FreeBytes: 500 << 30, TotalBytes: 900 << 30, Shared: true},
|
||||
{Path: "/var/lib/llmgw-models", FreeBytes: 20 << 30, TotalBytes: 60 << 30},
|
||||
}}
|
||||
free, total := n.LocalDisk()
|
||||
if free != 20<<30 || total != 60<<30 {
|
||||
t.Errorf("shared space was counted as local: free=%d total=%d", free, total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestANodeThatReportedNoStorageIsUnknownNotFull(t *testing.T) {
|
||||
// An agent too old to report storage, and a node with no model directories,
|
||||
// both send nothing. Reading that as a full disk would take every such node
|
||||
// out of placement for a reason that is not true.
|
||||
free, total := Node{}.LocalDisk()
|
||||
if free != 0 || total != 0 {
|
||||
t.Errorf("a node with no storage reported %d/%d", free, total)
|
||||
}
|
||||
// The contract callers rely on: total of zero is the signal for "unknown",
|
||||
// which is why nothing may read free alone.
|
||||
if total != 0 {
|
||||
t.Error("callers cannot distinguish unknown from full")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalDiskSumsSeveralDisks(t *testing.T) {
|
||||
n := Node{Storage: []discovery.Storage{
|
||||
{Path: "/a", FreeBytes: 10, TotalBytes: 100},
|
||||
{Path: "/b", FreeBytes: 5, TotalBytes: 50},
|
||||
}}
|
||||
if free, total := n.LocalDisk(); free != 15 || total != 150 {
|
||||
t.Errorf("free=%d total=%d, want 15/150", free, total)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"os"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// LeaseStore is the part of the store leadership needs.
|
||||
//
|
||||
// An interface so the controller does not take a dependency on a concrete
|
||||
// database for one feature, and so a deployment with no store at all — a
|
||||
// router-only run, or a single control plane on a static configuration — keeps
|
||||
// working with nothing to configure.
|
||||
type LeaseStore interface {
|
||||
AcquireLeadership(ctx context.Context, role, owner string, ttl time.Duration) (store.Lease, error)
|
||||
ReleaseLeadership(ctx context.Context, role, owner string) error
|
||||
}
|
||||
|
||||
// controlRole is the one thing there is to be leader of: placing work.
|
||||
const controlRole = "control-plane"
|
||||
|
||||
// Leadership decides whether this process is the one that places work.
|
||||
//
|
||||
// Two control planes against one store both reconcile: both see a model
|
||||
// wanting two replicas and none running, and both place. The result is twice
|
||||
// the replicas, each holding GPUs the other believes are free, and nothing
|
||||
// crashes to say so.
|
||||
//
|
||||
// A standby is not idle — it serves its own routing and keeps its connection
|
||||
// to the store — it simply does not place, evict or drain. That makes running
|
||||
// a second one safe rather than clever, which is the property worth having
|
||||
// first: the common failure is not a planned failover, it is somebody starting
|
||||
// a second control plane by accident.
|
||||
type Leadership struct {
|
||||
store LeaseStore
|
||||
log *slog.Logger
|
||||
// owner identifies this process. Random rather than the hostname, so a
|
||||
// restart on the same machine cannot be fooled by its own stale lease.
|
||||
owner string
|
||||
ttl time.Duration
|
||||
// renew is how often the holder pushes the expiry out. Well under ttl, so
|
||||
// one slow query does not cost the leadership.
|
||||
renew time.Duration
|
||||
|
||||
leader atomic.Bool
|
||||
}
|
||||
|
||||
// NewLeadership prepares leadership over a store.
|
||||
//
|
||||
// A nil store means there is nobody to coordinate with, and this process is
|
||||
// the leader by definition — which is every single-control-plane deployment,
|
||||
// and must stay the case without anyone opting in.
|
||||
func NewLeadership(st LeaseStore, ttl time.Duration, log *slog.Logger) *Leadership {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
if ttl <= 0 {
|
||||
// Long enough to survive a slow database or a brief network stall,
|
||||
// short enough that a standby takes over before anyone has finished
|
||||
// reading the alert.
|
||||
ttl = 15 * time.Second
|
||||
}
|
||||
l := &Leadership{
|
||||
store: st,
|
||||
log: log,
|
||||
owner: newOwnerID(),
|
||||
ttl: ttl,
|
||||
renew: ttl / 3,
|
||||
}
|
||||
if st == nil {
|
||||
l.leader.Store(true)
|
||||
}
|
||||
return l
|
||||
}
|
||||
|
||||
// IsLeader reports whether this process should be placing work.
|
||||
func (l *Leadership) IsLeader() bool { return l.leader.Load() }
|
||||
|
||||
// Owner is this process's id, for logs and for ctl.
|
||||
func (l *Leadership) Owner() string { return l.owner }
|
||||
|
||||
// Run keeps the lease for as long as it can hold it, and keeps trying when it
|
||||
// cannot.
|
||||
//
|
||||
// A standby polls at the same interval the leader renews at, so a lapsed lease
|
||||
// is picked up within about a third of its life rather than at the end of it.
|
||||
func (l *Leadership) Run(ctx context.Context) {
|
||||
if l.store == nil {
|
||||
return
|
||||
}
|
||||
t := time.NewTicker(l.renew)
|
||||
defer t.Stop()
|
||||
|
||||
l.attempt(ctx)
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
l.step(ctx)
|
||||
return
|
||||
case <-t.C:
|
||||
l.attempt(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// attempt takes or renews the lease, logging only transitions.
|
||||
//
|
||||
// Only transitions, because this runs every few seconds for the life of the
|
||||
// process: a line per attempt would bury everything else, and the thing an
|
||||
// operator needs to see is the moment leadership moved.
|
||||
func (l *Leadership) attempt(ctx context.Context) {
|
||||
lease, err := l.store.AcquireLeadership(ctx, controlRole, l.owner, l.ttl)
|
||||
switch {
|
||||
case err == nil:
|
||||
if l.leader.CompareAndSwap(false, true) {
|
||||
l.log.Info("this control plane is now placing work", "owner", l.owner)
|
||||
}
|
||||
case errors.Is(err, store.ErrNotLeader):
|
||||
if l.leader.CompareAndSwap(true, false) {
|
||||
// Losing a lease that was held is worth a warning: it means this
|
||||
// process was too slow to renew, and whatever made it slow is
|
||||
// still true.
|
||||
l.log.Warn("lost leadership; standing by",
|
||||
"owner", l.owner, "now", lease.Owner)
|
||||
} else if l.leader.Load() == false && lease.Owner != "" {
|
||||
l.log.Debug("standing by", "leader", lease.Owner,
|
||||
"until", lease.ExpiresAt.Format(time.RFC3339))
|
||||
}
|
||||
default:
|
||||
// A store that cannot be reached is not evidence that somebody else
|
||||
// is leading. Keeping the claim is the safer of two bad options: two
|
||||
// leaders is worse than one, but no leader at all means nothing is
|
||||
// placed and nothing is evicted, and a brief database blip should not
|
||||
// stop a cluster from healing itself.
|
||||
l.log.Warn("could not reach the store to renew leadership; keeping it for now",
|
||||
"err", err)
|
||||
}
|
||||
}
|
||||
|
||||
// step gives the lease up on the way out, so a standby takes over in seconds
|
||||
// rather than waiting out the clock.
|
||||
func (l *Leadership) step(ctx context.Context) {
|
||||
if !l.leader.Load() {
|
||||
return
|
||||
}
|
||||
// A fresh context: the one that just ended is the reason this is running.
|
||||
ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second)
|
||||
defer cancel()
|
||||
if err := l.store.ReleaseLeadership(ctx, controlRole, l.owner); err != nil {
|
||||
l.log.Warn("could not release leadership; a standby will wait it out", "err", err)
|
||||
return
|
||||
}
|
||||
l.log.Info("released leadership", "owner", l.owner)
|
||||
}
|
||||
|
||||
// newOwnerID is random, not the hostname.
|
||||
//
|
||||
// A hostname would let a restarted process match its own stale lease and
|
||||
// believe it had been leading all along — including through whatever gap made
|
||||
// it restart, during which another process may have taken over and started
|
||||
// placing.
|
||||
func newOwnerID() string {
|
||||
var b [8]byte
|
||||
if _, err := rand.Read(b[:]); err != nil {
|
||||
// Only reached if the system has no entropy, where a pid and the time
|
||||
// is still better than a fixed string.
|
||||
return "ctl-" + hex.EncodeToString([]byte(time.Now().Format(time.RFC3339Nano)))[:16]
|
||||
}
|
||||
host, _ := os.Hostname()
|
||||
if host == "" {
|
||||
host = "ctl"
|
||||
}
|
||||
return host + "-" + hex.EncodeToString(b[:])
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
package placement
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
// withDisk gives a node a local filesystem of the stated size.
|
||||
func withDisk(n inventory.Node, freeGB, totalGB int64) inventory.Node {
|
||||
n.Storage = []discovery.Storage{{
|
||||
Path: "/var/lib/llmgw-models",
|
||||
FreeBytes: freeGB << 30,
|
||||
TotalBytes: totalGB << 30,
|
||||
}}
|
||||
return n
|
||||
}
|
||||
|
||||
// weightsProfile is a model that would have to be fetched: it names a source
|
||||
// and says how much room that source needs.
|
||||
func weightsProfile(source string, diskGB int) Profile {
|
||||
return Profile{Name: "cpu", Engine: "llamacpp", Source: source, MinDiskGB: diskGB}
|
||||
}
|
||||
|
||||
func TestANodeWithoutRoomForTheWeightsIsNotACandidate(t *testing.T) {
|
||||
// The failure this replaces: the node is picked, the agent starts
|
||||
// downloading, and it dies partway through with a disk-full error that
|
||||
// names neither the model nor the scheduler decision that caused it.
|
||||
s := New(store(withDisk(node("full", "", 0), 2, 60)))
|
||||
|
||||
res := s.Place([]Profile{weightsProfile("https://models/x.gguf", 20)}, nil)
|
||||
if res.Placement != nil {
|
||||
t.Fatalf("placed on a node with 2 GB free for a 20 GB model")
|
||||
}
|
||||
if !strings.Contains(res.Explain(), "needs 20") {
|
||||
t.Errorf("the reason does not say what was needed:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoomIsNotAskedForTwiceOnANodeThatAlreadyHasTheWeights(t *testing.T) {
|
||||
// A node holding the file needs no space to hold it again. Without this,
|
||||
// the one node best placed to run a model — because it already has it — is
|
||||
// the one most likely to be rejected, since holding the file is what used
|
||||
// the space up.
|
||||
const source = "/var/lib/llmgw-models/big.gguf"
|
||||
n := withDisk(node("holder", "", 0), 2, 60)
|
||||
n.Weights = []discovery.LocalWeights{{Path: source, SizeBytes: 40 << 30}}
|
||||
s := New(store(n))
|
||||
|
||||
res := s.Place([]Profile{weightsProfile(source, 20)}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("a node already holding the weights was rejected for lack of room:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
func TestANodeThatReportsNoDiskIsStillACandidate(t *testing.T) {
|
||||
// An agent too old to report storage, and a node with no model directories,
|
||||
// both report none. Excluding them would mean a control plane upgrade
|
||||
// quietly empties every node in a fleet that has not upgraded yet — a worse
|
||||
// failure than the one this check prevents.
|
||||
s := New(store(node("silent", "", 0)))
|
||||
|
||||
res := s.Place([]Profile{weightsProfile("https://models/x.gguf", 20)}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("a node that reported no storage was treated as a full one:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithoutADeclaredSizeDiskIsNotConsidered(t *testing.T) {
|
||||
// Every profile that existed before this field has no number, and must
|
||||
// keep being placed exactly as it was.
|
||||
s := New(store(withDisk(node("full", "", 0), 0, 60)))
|
||||
|
||||
res := s.Place([]Profile{{Name: "cpu", Engine: "llamacpp", Source: "https://models/x.gguf"}}, nil)
|
||||
if res.Placement == nil {
|
||||
t.Fatalf("a profile with no declared size was rejected on disk:\n%s", res.Explain())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpaceOnASharedFilerDoesNotCountAsRoom(t *testing.T) {
|
||||
// The node cannot spend it: it belongs to everyone mounting the filer. A
|
||||
// node whose only storage is shared therefore reports no local disk, which
|
||||
// lands in the same place as reporting nothing — unknown, so not excluded.
|
||||
n := node("filer-only", "", 0)
|
||||
n.Storage = []discovery.Storage{{
|
||||
Path: "/mnt/filer", FreeBytes: 500 << 30, TotalBytes: 900 << 30, Shared: true,
|
||||
}}
|
||||
if free, total := n.LocalDisk(); free != 0 || total != 0 {
|
||||
t.Fatalf("shared space counted as local: %d/%d", free, total)
|
||||
}
|
||||
}
|
||||
@@ -36,6 +36,9 @@ type Profile struct {
|
||||
NodeSelector map[string]string
|
||||
// MinMemoryGB is the host RAM this profile needs. Zero is unconstrained.
|
||||
MinMemoryGB int
|
||||
// MinDiskGB is the free space needed to hold the weights, checked only on
|
||||
// nodes that do not already have them. Zero is unconstrained.
|
||||
MinDiskGB int
|
||||
// Source is the model's weights, as the engine would be given it. Used to
|
||||
// prefer nodes that already hold them; empty asks for no preference.
|
||||
Source string
|
||||
@@ -401,6 +404,22 @@ func (s *Scheduler) eligibleNodes(p Profile, fabric string, avoid []string, reje
|
||||
continue
|
||||
}
|
||||
}
|
||||
// Room for the weights, and only for a node that would have to fetch
|
||||
// them. A node already holding the file needs no space for it.
|
||||
//
|
||||
// Same two exemptions as memory, for the same reason: zero means the
|
||||
// operator declared nothing, and a node reporting no storage is an
|
||||
// older agent or one with no model directories, neither of which is a
|
||||
// full disk. Without that second exemption a control plane upgrade
|
||||
// would empty every node in a fleet that had not upgraded yet.
|
||||
if p.MinDiskGB > 0 && p.Source != "" && !n.HasWeights(p.Source) {
|
||||
if free, total := n.LocalDisk(); total > 0 {
|
||||
if have := int(free / (1 << 30)); have < p.MinDiskGB {
|
||||
reject(fmt.Sprintf("has %d GB free for weights, needs %d", have, p.MinDiskGB))
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if p.CPUOnly() {
|
||||
// Every ready node can run it; there is no card to be free.
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
package reconciler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
type fixedLeader bool
|
||||
|
||||
func (f fixedLeader) IsLeader() bool { return bool(f) }
|
||||
|
||||
func TestAStandbyPlacesNothing(t *testing.T) {
|
||||
// The failure this prevents: two control planes against one store both see
|
||||
// a model wanting replicas and none running, and both place. The model
|
||||
// ends up with twice what it asked for, each replica holding GPUs the
|
||||
// other's inventory believes are free, and nothing crashes to say so.
|
||||
h := newHarness(t, testModel(1), gpuNode("n1", 8))
|
||||
h.r.SetLeadership(fixedLeader(false))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
if got := len(h.r.Replicas()); got != 0 {
|
||||
t.Errorf("a standby placed %d replicas", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestALeaderPlacesAsBefore(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("n1", 8))
|
||||
h.r.SetLeadership(fixedLeader(true))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
if got := len(h.r.Replicas()); got == 0 {
|
||||
t.Error("a leader placed nothing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithoutLeadershipConfiguredNothingChanges(t *testing.T) {
|
||||
// Every single-control-plane deployment, and every test written before
|
||||
// leadership existed. Unset must mean "act", not "wait for permission
|
||||
// from a thing that was never wired up".
|
||||
h := newHarness(t, testModel(1), gpuNode("n1", 8))
|
||||
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
if got := len(h.r.Replicas()); got == 0 {
|
||||
t.Error("a reconciler with no leadership configured placed nothing")
|
||||
}
|
||||
}
|
||||
@@ -166,6 +166,9 @@ func (c *Config) normalize() {
|
||||
|
||||
// Reconciler compares the catalog against what is running and closes the gap.
|
||||
type Reconciler struct {
|
||||
// leader gates every action; nil means act unconditionally.
|
||||
leader Leader
|
||||
|
||||
cfg Config
|
||||
cat *catalog.Live
|
||||
inv *inventory.Store
|
||||
@@ -235,6 +238,17 @@ func (r *Reconciler) SetClock(now func() time.Time) {
|
||||
r.startedAt = now()
|
||||
}
|
||||
|
||||
// Leader reports whether this process should be acting on the fleet.
|
||||
type Leader interface{ IsLeader() bool }
|
||||
|
||||
// SetLeadership gates every action on holding leadership.
|
||||
//
|
||||
// A setter rather than a constructor argument because the thing that decides
|
||||
// leadership needs a store, and the store is wired after the controller. Nil
|
||||
// leaves the reconciler acting unconditionally, which is what a single control
|
||||
// plane wants and what every test expects.
|
||||
func (r *Reconciler) SetLeadership(l Leader) { r.leader = l }
|
||||
|
||||
// Start runs the loop until Close.
|
||||
func (r *Reconciler) Start(ctx context.Context) {
|
||||
go func() {
|
||||
@@ -278,6 +292,16 @@ func (r *Reconciler) Notify() {
|
||||
// Reconcile performs one pass. It is exported so tests can drive it directly
|
||||
// instead of racing a ticker.
|
||||
func (r *Reconciler) Reconcile(ctx context.Context) {
|
||||
// A standby does nothing at all here. Everything below issues commands to
|
||||
// agents — placing, evicting, draining, reaping — and two control planes
|
||||
// doing that against one fleet is how a model ends up with twice the
|
||||
// replicas it asked for, each holding GPUs the other believes are free.
|
||||
//
|
||||
// Unset means leader, so a single control plane and every test behave
|
||||
// exactly as they did before leadership existed.
|
||||
if r.leader != nil && !r.leader.IsLeader() {
|
||||
return
|
||||
}
|
||||
r.expireNodes()
|
||||
r.expireStuckReplicas()
|
||||
r.drainCordonedNodes(ctx)
|
||||
|
||||
@@ -31,6 +31,20 @@ func (r *Router) budgetFallback(p auth.Principal, req *parsedRequest) (accountin
|
||||
if !r.permits(p, name) {
|
||||
continue
|
||||
}
|
||||
// And it has to be able to hold the prompt. The fit was checked against
|
||||
// the model the caller asked for, before anyone knew a different one
|
||||
// would serve — so without this a 20k prompt sent to a 32k model lands
|
||||
// on a 2k one, and the engine truncates it or refuses. Either way the
|
||||
// caller sees a failure naming a model they never chose.
|
||||
//
|
||||
// The history is not trimmed to make it fit. Dropping a caller's
|
||||
// messages to stay under a limit they did not ask about is a worse
|
||||
// answer than saying the prompt is too large, and it would be done
|
||||
// silently, to a request that was already being quietly moved.
|
||||
_, cand := r.resolve(name)
|
||||
if why := checkContextFits(cand, req.PromptTokens, req.MaxTokens); why != "" {
|
||||
continue
|
||||
}
|
||||
res, err := r.acct.Check(p, name, req.estimatedTokens())
|
||||
if err != nil {
|
||||
continue
|
||||
|
||||
@@ -94,6 +94,18 @@ func (r *Router) Handler() http.Handler {
|
||||
r.landing.Contact, r.landing.Notice))
|
||||
mux.Handle("GET /_static/", web.Static())
|
||||
mux.HandleFunc("GET /me", r.handleMe)
|
||||
// Somebody has to be able to say "forget that", and a retention window is
|
||||
// no answer to a message sent by mistake. Authenticated like any other
|
||||
// client endpoint, and scoped to the caller's project by the store.
|
||||
mux.HandleFunc("DELETE /v1/threads/{id}", r.handleDeleteThread)
|
||||
|
||||
// Notes an agent keeps for itself, over REST and as MCP tools. Off unless a
|
||||
// project or the threads block turned storage on.
|
||||
mux.HandleFunc("PUT /v1/memory/{key}", r.handlePutNote)
|
||||
mux.HandleFunc("GET /v1/memory/{key}", r.handleGetNote)
|
||||
mux.HandleFunc("GET /v1/memory", r.handleSearchNotes)
|
||||
mux.HandleFunc("DELETE /v1/memory/{key}", r.handleDeleteNote)
|
||||
mux.HandleFunc("POST /mcp", r.handleMCP)
|
||||
|
||||
if r.store != nil {
|
||||
// The one-time link that delivers a key to the person it was issued
|
||||
@@ -164,6 +176,17 @@ type parsedRequest struct {
|
||||
PromptTokens int
|
||||
MaxTokens int
|
||||
|
||||
// ThreadID is the conversation this request continues, empty for the
|
||||
// stateless path every request took before threads existed.
|
||||
ThreadID string
|
||||
// ThreadTurns are the turns this request itself carried, kept so the
|
||||
// exchange can be stored after a reply without re-reading a body that has
|
||||
// since had the history spliced into it.
|
||||
ThreadTurns []store.Message
|
||||
// ThreadProject is the project the thread is stored under, resolved once by
|
||||
// loadThread so saving it afterwards cannot disagree.
|
||||
ThreadProject string
|
||||
|
||||
// reservation is carried from the limit check to the point where the
|
||||
// request is settled, so the record counts against the same budget window
|
||||
// the check read.
|
||||
@@ -255,8 +278,26 @@ func (r *Router) handleInference(path string) http.HandlerFunc {
|
||||
w.Header().Set("x-llmgw-queue-ms", strconv.FormatInt(queued.Milliseconds(), 10))
|
||||
}
|
||||
|
||||
// After the model is finally settled — a budget fallback may have
|
||||
// changed it — because trimming is measured against the context of the
|
||||
// model that will actually serve, not the one that was asked for.
|
||||
if dropped, err := r.loadThread(req.Context(), parsed); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
r.m.observeRequest(parsed.Model, "", http.StatusBadRequest, time.Since(start), nil, 0)
|
||||
return
|
||||
} else if dropped > 0 {
|
||||
// Dropping turns is the right answer and must not be a silent one.
|
||||
w.Header().Set(trimmedHeader, strconv.Itoa(dropped))
|
||||
}
|
||||
|
||||
res := r.dispatch(req.Context(), w, parsed)
|
||||
r.finish(w, parsed, res, start)
|
||||
// After the client has been served. A conversation that failed to
|
||||
// record is a worse outcome than a slow reply, but only just, and the
|
||||
// caller has their answer either way.
|
||||
if res.err == nil && res.outcome.Status == http.StatusOK {
|
||||
r.saveThread(context.WithoutCancel(req.Context()), parsed, res.outcome.Reply, parsed.Model)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -390,6 +431,7 @@ func (r *Router) parse(req *http.Request, path string) (*parsedRequest, error) {
|
||||
ProfileOverride: req.Header.Get("x-llmgw-profile"),
|
||||
MaxTokens: max(b.MaxTokens, b.MaxCompletionTokens),
|
||||
PromptPrefix: promptPrefix(&b),
|
||||
ThreadID: req.Header.Get(threadHeader),
|
||||
}
|
||||
// With authentication on, the key decides the band and the routing profile,
|
||||
// and the headers stop being believed. Trusting the header was only ever
|
||||
|
||||
@@ -0,0 +1,361 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// agentHeader names the agent whose notes a request reads and writes.
|
||||
//
|
||||
// A namespace inside a project, not a credential: with auth on, the key decides
|
||||
// the project and this only separates agents sharing it; with auth off it is
|
||||
// just a name.
|
||||
const agentHeader = "x-llmgw-agent"
|
||||
|
||||
const (
|
||||
defaultAgent = "default"
|
||||
maxAgentID = 100
|
||||
maxMemoryBody = 64 << 10
|
||||
mcpProtocolVer = "2025-03-26"
|
||||
)
|
||||
|
||||
// errMemoryOff is what a memory call meets when nothing is configured to keep
|
||||
// notes. Refused rather than pretended: an agent that believes it saved
|
||||
// something and finds nothing later has been lied to.
|
||||
var errMemoryOff = errors.New(
|
||||
"this gateway keeps no memory notes: enable auth with a project that keeps memory, " +
|
||||
"or set threads.store and threads.retain_hours")
|
||||
|
||||
// memoryScope says whose notes this request is about.
|
||||
func (r *Router) memoryScope(req *http.Request) (project, agent string, err error) {
|
||||
var caller string
|
||||
if p, ok := auth.FromContext(req.Context()); ok {
|
||||
caller = p.ProjectID
|
||||
}
|
||||
project, err = r.threadProject(caller)
|
||||
if err != nil {
|
||||
if errors.Is(err, errThreadsOff) {
|
||||
err = errMemoryOff
|
||||
}
|
||||
return "", "", err
|
||||
}
|
||||
agent = req.Header.Get(agentHeader)
|
||||
if agent == "" {
|
||||
agent = defaultAgent
|
||||
}
|
||||
if len(agent) > maxAgentID {
|
||||
return "", "", errors.New(agentHeader + " is too long")
|
||||
}
|
||||
return project, agent, nil
|
||||
}
|
||||
|
||||
type noteJSON struct {
|
||||
Key string `json:"key"`
|
||||
Text string `json:"text"`
|
||||
Tags []string `json:"tags"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
UpdatedAt string `json:"updated_at"`
|
||||
}
|
||||
|
||||
func toNoteJSON(n store.Note) noteJSON {
|
||||
tags := n.Tags
|
||||
if tags == nil {
|
||||
tags = []string{}
|
||||
}
|
||||
return noteJSON{
|
||||
Key: n.Key, Text: n.Text, Tags: tags,
|
||||
CreatedAt: n.CreatedAt.UTC().Format(time.RFC3339),
|
||||
UpdatedAt: n.UpdatedAt.UTC().Format(time.RFC3339),
|
||||
}
|
||||
}
|
||||
|
||||
// memoryStatus maps a store error to what a client should be told.
|
||||
func memoryStatus(err error) int {
|
||||
switch {
|
||||
case errors.Is(err, store.ErrNoNote):
|
||||
return http.StatusNotFound
|
||||
case errors.Is(err, errThreadsNeedAuth):
|
||||
return http.StatusUnauthorized
|
||||
case errors.Is(err, errMemoryOff), errors.Is(err, store.ErrMemoryNotEnabled),
|
||||
errors.Is(err, store.ErrMemoryFull):
|
||||
return http.StatusBadRequest
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func writeMemoryError(w http.ResponseWriter, err error) {
|
||||
if code := memoryStatus(err); code != 0 {
|
||||
msg := err.Error()
|
||||
if code == http.StatusNotFound {
|
||||
// One answer for missing, another agent's and another project's, so
|
||||
// a caller cannot map which keys exist where they cannot read.
|
||||
msg = "no such note"
|
||||
}
|
||||
writeError(w, code, "invalid_request_error", msg)
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
}
|
||||
|
||||
// handlePutNote writes a note: PUT /v1/memory/{key}.
|
||||
func (r *Router) handlePutNote(w http.ResponseWriter, req *http.Request) {
|
||||
project, agent, err := r.memoryScope(req)
|
||||
if err != nil {
|
||||
writeMemoryError(w, err)
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Text string `json:"text"`
|
||||
Tags []string `json:"tags"`
|
||||
}
|
||||
raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, maxMemoryBody))
|
||||
if err != nil || json.Unmarshal(raw, &body) != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error",
|
||||
`body must be JSON like {"text": "...", "tags": ["..."]}`)
|
||||
return
|
||||
}
|
||||
note := store.Note{Key: req.PathValue("key"), Text: body.Text, Tags: body.Tags}
|
||||
if err := r.threads.PutNote(req.Context(), project, agent, note); err != nil {
|
||||
writeMemoryError(w, err)
|
||||
return
|
||||
}
|
||||
saved, err := r.threads.GetNote(req.Context(), project, agent, note.Key)
|
||||
if err != nil {
|
||||
writeMemoryError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, toNoteJSON(saved))
|
||||
}
|
||||
|
||||
// handleGetNote reads one note: GET /v1/memory/{key}.
|
||||
func (r *Router) handleGetNote(w http.ResponseWriter, req *http.Request) {
|
||||
project, agent, err := r.memoryScope(req)
|
||||
if err != nil {
|
||||
writeMemoryError(w, err)
|
||||
return
|
||||
}
|
||||
n, err := r.threads.GetNote(req.Context(), project, agent, req.PathValue("key"))
|
||||
if err != nil {
|
||||
writeMemoryError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, toNoteJSON(n))
|
||||
}
|
||||
|
||||
// handleSearchNotes lists or searches: GET /v1/memory?q=&tag=&limit=.
|
||||
func (r *Router) handleSearchNotes(w http.ResponseWriter, req *http.Request) {
|
||||
project, agent, err := r.memoryScope(req)
|
||||
if err != nil {
|
||||
writeMemoryError(w, err)
|
||||
return
|
||||
}
|
||||
q := req.URL.Query()
|
||||
limit, _ := strconv.Atoi(q.Get("limit"))
|
||||
notes, err := r.threads.SearchNotes(req.Context(), project, agent, q.Get("q"), q.Get("tag"), limit)
|
||||
if err != nil {
|
||||
writeMemoryError(w, err)
|
||||
return
|
||||
}
|
||||
out := make([]noteJSON, 0, len(notes))
|
||||
for _, n := range notes {
|
||||
out = append(out, toNoteJSON(n))
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"object": "list", "data": out})
|
||||
}
|
||||
|
||||
// handleDeleteNote forgets one: DELETE /v1/memory/{key}.
|
||||
func (r *Router) handleDeleteNote(w http.ResponseWriter, req *http.Request) {
|
||||
project, agent, err := r.memoryScope(req)
|
||||
if err != nil {
|
||||
writeMemoryError(w, err)
|
||||
return
|
||||
}
|
||||
if err := r.threads.DeleteNote(req.Context(), project, agent, req.PathValue("key")); err != nil {
|
||||
writeMemoryError(w, err)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
|
||||
// --- MCP ---
|
||||
//
|
||||
// The same notes, offered as tools an agent can be given. Only the small part
|
||||
// of the protocol a tool server needs: initialize, tools/list and tools/call,
|
||||
// answered as plain JSON over POST. It never streams and never calls the
|
||||
// client, so it needs no session.
|
||||
|
||||
type rpcRequest struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID json.RawMessage `json:"id"`
|
||||
Method string `json:"method"`
|
||||
Params json.RawMessage `json:"params"`
|
||||
}
|
||||
|
||||
type rpcError struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
type rpcResponse struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID json.RawMessage `json:"id"`
|
||||
Result any `json:"result,omitempty"`
|
||||
Error *rpcError `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
var mcpTools = []map[string]any{
|
||||
{
|
||||
"name": "remember",
|
||||
"description": "Save a note that will still be there in later conversations. Writing to an existing key replaces it.",
|
||||
"inputSchema": map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"key", "text"},
|
||||
"properties": map[string]any{
|
||||
"key": map[string]any{"type": "string", "description": "Short unique name for the note"},
|
||||
"text": map[string]any{"type": "string", "description": "What to remember"},
|
||||
"tags": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "Optional lowercase labels"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "recall",
|
||||
"description": "Find saved notes by words, by tag, or by exact key. With nothing given, returns the most recent notes.",
|
||||
"inputSchema": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"query": map[string]any{"type": "string", "description": "Words to search for"},
|
||||
"tag": map[string]any{"type": "string", "description": "Only notes with this tag"},
|
||||
"key": map[string]any{"type": "string", "description": "Exact key to read"},
|
||||
"limit": map[string]any{"type": "integer", "description": "Most notes to return (up to 100)"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "forget",
|
||||
"description": "Delete a saved note.",
|
||||
"inputSchema": map[string]any{
|
||||
"type": "object",
|
||||
"required": []string{"key"},
|
||||
"properties": map[string]any{"key": map[string]any{"type": "string"}},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// handleMCP serves POST /mcp.
|
||||
func (r *Router) handleMCP(w http.ResponseWriter, req *http.Request) {
|
||||
raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, maxMemoryBody))
|
||||
var rpc rpcRequest
|
||||
if err != nil || json.Unmarshal(raw, &rpc) != nil || rpc.Method == "" {
|
||||
writeJSON(w, http.StatusBadRequest, rpcResponse{
|
||||
JSONRPC: "2.0", ID: json.RawMessage("null"),
|
||||
Error: &rpcError{Code: -32600, Message: "expected one JSON-RPC request object"},
|
||||
})
|
||||
return
|
||||
}
|
||||
// A notification has no id and gets no answer.
|
||||
if len(rpc.ID) == 0 {
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
return
|
||||
}
|
||||
|
||||
resp := rpcResponse{JSONRPC: "2.0", ID: rpc.ID}
|
||||
switch rpc.Method {
|
||||
case "initialize":
|
||||
resp.Result = map[string]any{
|
||||
"protocolVersion": mcpProtocolVer,
|
||||
"capabilities": map[string]any{"tools": map[string]any{}},
|
||||
"serverInfo": map[string]any{"name": "llmgw-memory", "version": "1"},
|
||||
}
|
||||
case "ping":
|
||||
resp.Result = map[string]any{}
|
||||
case "tools/list":
|
||||
resp.Result = map[string]any{"tools": mcpTools}
|
||||
case "tools/call":
|
||||
resp.Result = r.callMemoryTool(req, rpc.Params)
|
||||
default:
|
||||
resp.Error = &rpcError{Code: -32601, Message: "method not found: " + rpc.Method}
|
||||
}
|
||||
writeJSON(w, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
func toolText(text string, isError bool) map[string]any {
|
||||
return map[string]any{
|
||||
"content": []map[string]any{{"type": "text", "text": text}},
|
||||
"isError": isError,
|
||||
}
|
||||
}
|
||||
|
||||
func toolJSON(v any) map[string]any {
|
||||
b, _ := json.MarshalIndent(v, "", " ")
|
||||
return toolText(string(b), false)
|
||||
}
|
||||
|
||||
// callMemoryTool runs one tool. Failures are tool results with isError set, not
|
||||
// protocol errors, because that is how a model gets to see what went wrong.
|
||||
func (r *Router) callMemoryTool(req *http.Request, params json.RawMessage) map[string]any {
|
||||
var call struct {
|
||||
Name string `json:"name"`
|
||||
Arguments struct {
|
||||
Key string `json:"key"`
|
||||
Text string `json:"text"`
|
||||
Tags []string `json:"tags"`
|
||||
Query string `json:"query"`
|
||||
Tag string `json:"tag"`
|
||||
Limit int `json:"limit"`
|
||||
} `json:"arguments"`
|
||||
}
|
||||
if err := json.Unmarshal(params, &call); err != nil {
|
||||
return toolText("invalid tool call", true)
|
||||
}
|
||||
project, agent, err := r.memoryScope(req)
|
||||
if err != nil {
|
||||
return toolText(err.Error(), true)
|
||||
}
|
||||
a, ctx := call.Arguments, req.Context()
|
||||
|
||||
switch call.Name {
|
||||
case "remember":
|
||||
if err := r.threads.PutNote(ctx, project, agent, store.Note{Key: a.Key, Text: a.Text, Tags: a.Tags}); err != nil {
|
||||
return toolText(err.Error(), true)
|
||||
}
|
||||
return toolText("saved "+a.Key, false)
|
||||
case "recall":
|
||||
if a.Key != "" {
|
||||
n, err := r.threads.GetNote(ctx, project, agent, a.Key)
|
||||
if errors.Is(err, store.ErrNoNote) {
|
||||
return toolText("no note with key "+a.Key, false)
|
||||
}
|
||||
if err != nil {
|
||||
return toolText(err.Error(), true)
|
||||
}
|
||||
return toolJSON([]noteJSON{toNoteJSON(n)})
|
||||
}
|
||||
notes, err := r.threads.SearchNotes(ctx, project, agent, a.Query, a.Tag, a.Limit)
|
||||
if err != nil {
|
||||
return toolText(err.Error(), true)
|
||||
}
|
||||
out := make([]noteJSON, 0, len(notes))
|
||||
for _, n := range notes {
|
||||
out = append(out, toNoteJSON(n))
|
||||
}
|
||||
return toolJSON(out)
|
||||
case "forget":
|
||||
if err := r.threads.DeleteNote(ctx, project, agent, a.Key); err != nil {
|
||||
if errors.Is(err, store.ErrNoNote) {
|
||||
return toolText("no note with key "+a.Key, false)
|
||||
}
|
||||
return toolText(err.Error(), true)
|
||||
}
|
||||
return toolText("forgot "+a.Key, false)
|
||||
default:
|
||||
return toolText("unknown tool "+call.Name, true)
|
||||
}
|
||||
}
|
||||
@@ -154,6 +154,14 @@ type Request struct {
|
||||
// still recording it. Set when the router asked for usage the client did
|
||||
// not request.
|
||||
SuppressUsageChunk bool
|
||||
// CaptureReply asks for the assistant's text back in the Outcome, up to
|
||||
// this many bytes. Zero means capture nothing, which is every request that
|
||||
// is not continuing a stored conversation.
|
||||
//
|
||||
// Opt-in and bounded on purpose. This is the hottest path in the system,
|
||||
// and a gateway that accumulated every reply in memory would be one long
|
||||
// answer away from trouble; a request that stores nothing must pay nothing.
|
||||
CaptureReply int
|
||||
}
|
||||
|
||||
// Usage holds the token counts reported by the engine.
|
||||
@@ -178,6 +186,14 @@ type Outcome struct {
|
||||
Usage *Usage
|
||||
// Bytes is how many bytes were relayed to the client.
|
||||
Bytes int64
|
||||
// Reply is the assistant's text, when CaptureReply asked for it. Truncated
|
||||
// at the cap rather than dropped: a shortened turn replays as a shortened
|
||||
// turn, while nothing at all would replay as the assistant having said
|
||||
// nothing, which is a different and wrong conversation.
|
||||
Reply string
|
||||
// ReplyTruncated says the cap was reached, so nothing downstream records a
|
||||
// half answer as a whole one.
|
||||
ReplyTruncated bool
|
||||
}
|
||||
|
||||
var (
|
||||
@@ -232,11 +248,11 @@ func (p *Proxy) Do(ctx context.Context, w http.ResponseWriter, r Request) (Outco
|
||||
if r.Stream && isEventStream(resp.Header) {
|
||||
return p.relayStream(ctx, cancel, w, resp, r)
|
||||
}
|
||||
return p.relayBuffered(w, resp)
|
||||
return p.relayBuffered(w, resp, r)
|
||||
}
|
||||
|
||||
// relayBuffered copies a non-streaming response through, capturing usage.
|
||||
func (p *Proxy) relayBuffered(w http.ResponseWriter, resp *http.Response) (Outcome, error) {
|
||||
func (p *Proxy) relayBuffered(w http.ResponseWriter, resp *http.Response, r Request) (Outcome, error) {
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, p.cfg.MaxResponseBytes))
|
||||
if err != nil {
|
||||
return Outcome{Status: resp.StatusCode, Retryable: true}, err
|
||||
@@ -252,6 +268,9 @@ func (p *Proxy) relayBuffered(w http.ResponseWriter, resp *http.Response) (Outco
|
||||
Bytes: int64(n),
|
||||
Usage: usageFromJSON(body),
|
||||
}
|
||||
if r.CaptureReply > 0 && resp.StatusCode == http.StatusOK {
|
||||
out.Reply, out.ReplyTruncated = capReply(replyFromJSON(body), r.CaptureReply)
|
||||
}
|
||||
if werr != nil {
|
||||
return out, ErrClientGone
|
||||
}
|
||||
@@ -262,7 +281,7 @@ func (p *Proxy) relayBuffered(w http.ResponseWriter, resp *http.Response) (Outco
|
||||
//
|
||||
// abort must cancel the context the upstream request was built with: the idle
|
||||
// watchdog can only stop a silent engine by aborting that read.
|
||||
func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w http.ResponseWriter, resp *http.Response, r Request) (Outcome, error) {
|
||||
func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w http.ResponseWriter, resp *http.Response, r Request) (result Outcome, rerr error) {
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
return Outcome{Status: resp.StatusCode, Retryable: false},
|
||||
@@ -275,6 +294,15 @@ func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w htt
|
||||
|
||||
out := Outcome{Wrote: true, Status: resp.StatusCode}
|
||||
start := time.Now()
|
||||
// Nothing is accumulated unless a thread asked for it, so a request that
|
||||
// stores nothing pays nothing on the hottest path in the system.
|
||||
reply := replyBuilder{cap: r.CaptureReply}
|
||||
// On every exit rather than at the one that looks like the end: a stream
|
||||
// can also end with the client gone or the engine stalled, and what the
|
||||
// assistant said up to that point is still what it said.
|
||||
defer func() {
|
||||
result.Reply, result.ReplyTruncated = reply.result()
|
||||
}()
|
||||
|
||||
// One watchdog covers both phases: until the first event it enforces
|
||||
// FirstTokenTimeout, afterwards IdleTimeout.
|
||||
@@ -295,6 +323,9 @@ func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w htt
|
||||
out.TTFT = time.Since(start)
|
||||
wd.setInterval(p.cfg.IdleTimeout)
|
||||
}
|
||||
if r.CaptureReply > 0 {
|
||||
reply.add(replyDelta(ev.data.String()))
|
||||
}
|
||||
if u := ev.usage(); u != nil {
|
||||
out.Usage = u
|
||||
if r.SuppressUsageChunk {
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// replyFromJSON pulls the assistant's text out of a completed response.
|
||||
//
|
||||
// Only the first choice. A request asking for several is not one continuing a
|
||||
// conversation, and storing one of them as "what the assistant said" would
|
||||
// invent a conversation that never happened.
|
||||
func replyFromJSON(body []byte) string {
|
||||
var b struct {
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
// The completions endpoint, which has no message wrapper.
|
||||
Text string `json:"text"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &b); err != nil || len(b.Choices) == 0 {
|
||||
return ""
|
||||
}
|
||||
if c := b.Choices[0].Message.Content; c != "" {
|
||||
return c
|
||||
}
|
||||
return b.Choices[0].Text
|
||||
}
|
||||
|
||||
// replyDelta pulls the incremental text out of one streamed event.
|
||||
func replyDelta(event string) string {
|
||||
if event == "" || event == "[DONE]" {
|
||||
return ""
|
||||
}
|
||||
var b struct {
|
||||
Choices []struct {
|
||||
Delta struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"delta"`
|
||||
Text string `json:"text"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(event), &b); err != nil || len(b.Choices) == 0 {
|
||||
return ""
|
||||
}
|
||||
if c := b.Choices[0].Delta.Content; c != "" {
|
||||
return c
|
||||
}
|
||||
return b.Choices[0].Text
|
||||
}
|
||||
|
||||
// capReply truncates at the cap and says whether it had to.
|
||||
//
|
||||
// Truncated rather than dropped, because a shortened turn replays as a
|
||||
// shortened turn while nothing at all replays as the assistant having said
|
||||
// nothing — a different conversation, and a confusing one to be handed back.
|
||||
func capReply(s string, cap int) (string, bool) {
|
||||
if cap <= 0 || len(s) <= cap {
|
||||
return s, false
|
||||
}
|
||||
// On a rune boundary, so a stored reply is never invalid UTF-8 — it goes
|
||||
// back into a JSON body, and a split multi-byte character would make that
|
||||
// body unserialisable.
|
||||
cut := cap
|
||||
for cut > 0 && !isRuneStart(s[cut]) {
|
||||
cut--
|
||||
}
|
||||
return s[:cut], true
|
||||
}
|
||||
|
||||
func isRuneStart(b byte) bool { return b&0xC0 != 0x80 }
|
||||
|
||||
// replyBuilder accumulates a streamed reply up to a cap.
|
||||
//
|
||||
// It stops appending once full rather than growing and truncating at the end,
|
||||
// so a long answer cannot cost more memory than the cap allows however much
|
||||
// the engine sends.
|
||||
type replyBuilder struct {
|
||||
cap int
|
||||
sb strings.Builder
|
||||
truncated bool
|
||||
}
|
||||
|
||||
func (r *replyBuilder) add(s string) {
|
||||
if r.cap <= 0 || s == "" || r.truncated {
|
||||
return
|
||||
}
|
||||
if room := r.cap - r.sb.Len(); len(s) > room {
|
||||
part, _ := capReply(s, room)
|
||||
r.sb.WriteString(part)
|
||||
r.truncated = true
|
||||
return
|
||||
}
|
||||
r.sb.WriteString(s)
|
||||
}
|
||||
|
||||
func (r *replyBuilder) result() (string, bool) { return r.sb.String(), r.truncated }
|
||||
@@ -0,0 +1,95 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNothingIsCapturedUnlessAThreadAsked(t *testing.T) {
|
||||
// This is the hottest path in the system. A gateway that accumulated every
|
||||
// reply in memory would be one long answer away from trouble, so a request
|
||||
// that stores nothing must pay nothing.
|
||||
var r replyBuilder // cap zero, which is every ordinary request
|
||||
r.add("a long assistant reply that nobody asked to keep")
|
||||
if got, truncated := r.result(); got != "" || truncated {
|
||||
t.Errorf("a request with no thread accumulated %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestALongReplyCostsNoMoreThanTheCap(t *testing.T) {
|
||||
// Stopping at the cap rather than growing and truncating afterwards is
|
||||
// what makes the cap a memory bound instead of a formatting rule.
|
||||
r := replyBuilder{cap: 10}
|
||||
for range 1000 {
|
||||
r.add(strings.Repeat("x", 100))
|
||||
}
|
||||
got, truncated := r.result()
|
||||
if len(got) > 10 {
|
||||
t.Errorf("accumulated %d bytes against a cap of 10", len(got))
|
||||
}
|
||||
if !truncated {
|
||||
t.Error("truncation was not reported")
|
||||
}
|
||||
}
|
||||
|
||||
func TestATruncatedReplyIsStillValidUTF8(t *testing.T) {
|
||||
// It goes back into a JSON body on the next turn. A split multi-byte
|
||||
// character would make that body unserialisable, and the thread would
|
||||
// break for good rather than for one turn.
|
||||
r := replyBuilder{cap: 4}
|
||||
r.add("üçüncü") // two-byte runes either side of the cut
|
||||
got, _ := r.result()
|
||||
for i, ch := range got {
|
||||
if ch == '�' {
|
||||
t.Fatalf("invalid UTF-8 at byte %d of %q", i, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheAssistantTextIsFoundInBothResponseShapes(t *testing.T) {
|
||||
chat := `{"choices":[{"message":{"role":"assistant","content":"hello"}}]}`
|
||||
if got := replyFromJSON([]byte(chat)); got != "hello" {
|
||||
t.Errorf("chat completion reply = %q", got)
|
||||
}
|
||||
// The completions endpoint, which has no message wrapper.
|
||||
legacy := `{"choices":[{"text":"hello"}]}`
|
||||
if got := replyFromJSON([]byte(legacy)); got != "hello" {
|
||||
t.Errorf("completion reply = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOnlyTheFirstChoiceIsKept(t *testing.T) {
|
||||
// A request asking for several is not one continuing a conversation, and
|
||||
// storing one of them as "what the assistant said" would invent a
|
||||
// conversation that never happened.
|
||||
body := `{"choices":[{"message":{"content":"one"}},{"message":{"content":"two"}}]}`
|
||||
if got := replyFromJSON([]byte(body)); got != "one" {
|
||||
t.Errorf("reply = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStreamedDeltasAssembleIntoTheReply(t *testing.T) {
|
||||
r := replyBuilder{cap: 1024}
|
||||
for _, ev := range []string{
|
||||
`{"choices":[{"delta":{"role":"assistant"}}]}`,
|
||||
`{"choices":[{"delta":{"content":"Hel"}}]}`,
|
||||
`{"choices":[{"delta":{"content":"lo"}}]}`,
|
||||
`[DONE]`,
|
||||
} {
|
||||
r.add(replyDelta(ev))
|
||||
}
|
||||
if got, _ := r.result(); got != "Hello" {
|
||||
t.Errorf("assembled %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMalformedEventContributesNothing(t *testing.T) {
|
||||
// An engine sending something unexpected must not corrupt a stored
|
||||
// conversation; the worst case is a turn that is shorter than it was.
|
||||
if got := replyDelta(`not json at all`); got != "" {
|
||||
t.Errorf("garbage produced %q", got)
|
||||
}
|
||||
if got := replyDelta(``); got != "" {
|
||||
t.Errorf("an empty event produced %q", got)
|
||||
}
|
||||
}
|
||||
@@ -76,7 +76,19 @@ type Router struct {
|
||||
accounts *account.Service
|
||||
store *storepkg.SQLite
|
||||
acct *accounting.Accountant
|
||||
recorder *accounting.Recorder
|
||||
|
||||
// threads holds conversations and memory notes. It is store when auth is on,
|
||||
// its own file when only the threads block is configured, and nil when
|
||||
// nothing is kept. Separate from store because a non-nil store also switches
|
||||
// on the accounts, web pages and usage recording, none of which belong to a
|
||||
// gateway that only wants to remember things.
|
||||
threads *storepkg.SQLite
|
||||
// anonymous means threads has no project to scope to: everything shares the
|
||||
// one pool named by anonymousProject.
|
||||
anonymous bool
|
||||
// ownsThreads is true when threads is not store and so must be closed here.
|
||||
ownsThreads bool
|
||||
recorder *accounting.Recorder
|
||||
|
||||
// prefix is nil unless approximate prefix routing is turned on.
|
||||
prefix *scheduler.PrefixScorer
|
||||
@@ -177,7 +189,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) {
|
||||
}, log)
|
||||
|
||||
if cfg.Auth.Enabled {
|
||||
st, err := storepkg.OpenSQLite(context.Background(), cfg.Auth.Store)
|
||||
st, err := storepkg.Open(context.Background(), cfg.Auth.Store)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("auth store: %w", err)
|
||||
}
|
||||
@@ -220,6 +232,17 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) {
|
||||
CompletionPer1k: cfg.Auth.DefaultPrice.CompletionPer1k,
|
||||
}),
|
||||
r.recorder, log)
|
||||
r.threads = r.store
|
||||
} else if cfg.Threads.Store != "" {
|
||||
st, err := storepkg.Open(context.Background(), cfg.Threads.Store)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("threads store: %w", err)
|
||||
}
|
||||
if err := st.EnsureAnonymousProject(context.Background(), anonymousProject, cfg.Threads.RetainHours); err != nil {
|
||||
_ = st.Close()
|
||||
return nil, fmt.Errorf("threads store: %w", err)
|
||||
}
|
||||
r.threads, r.anonymous, r.ownsThreads = st, true, true
|
||||
}
|
||||
|
||||
if cfg.Router.Admission.Enabled {
|
||||
@@ -333,6 +356,51 @@ func (r *Router) Start(ctx context.Context) {
|
||||
r.log.Warn("could not restore budgets; they start from zero this period", "err", err)
|
||||
}
|
||||
}
|
||||
if r.threads != nil {
|
||||
go r.pruneThreads(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
// pruneThreads enforces each project's retention window.
|
||||
//
|
||||
// Without something running this, "we keep conversations for 24 hours" is a
|
||||
// comment rather than a promise — the rows would simply accumulate, and the
|
||||
// number a project chose would describe nothing.
|
||||
//
|
||||
// Hourly, and once at startup. A gateway that has been down for a week should
|
||||
// not wait another hour before honouring a window that expired during it, and
|
||||
// an hour is fine granularity for a limit expressed in hours.
|
||||
func (r *Router) pruneThreads(ctx context.Context) {
|
||||
prune := func() {
|
||||
n, err := r.threads.PruneThreads(ctx, time.Now())
|
||||
if err != nil {
|
||||
r.log.Warn("could not prune expired conversations", "err", err)
|
||||
return
|
||||
}
|
||||
if n > 0 {
|
||||
r.log.Info("removed conversations past their retention", "threads", n)
|
||||
}
|
||||
m, err := r.threads.PruneMemory(ctx, time.Now())
|
||||
if err != nil {
|
||||
r.log.Warn("could not prune expired memory notes", "err", err)
|
||||
return
|
||||
}
|
||||
if m > 0 {
|
||||
r.log.Info("removed memory notes past their retention", "notes", m)
|
||||
}
|
||||
}
|
||||
prune()
|
||||
|
||||
t := time.NewTicker(time.Hour)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
prune()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Stop halts scraping and admission.
|
||||
@@ -348,6 +416,9 @@ func (r *Router) Stop() {
|
||||
if r.store != nil {
|
||||
_ = r.store.Close()
|
||||
}
|
||||
if r.ownsThreads {
|
||||
_ = r.threads.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// Registry exposes the endpoint registry, for tests and admin endpoints.
|
||||
@@ -560,6 +631,7 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa
|
||||
Header: req.Header,
|
||||
Stream: req.Stream,
|
||||
SuppressUsageChunk: req.SuppressUsageChunk,
|
||||
CaptureReply: req.captureBytes(),
|
||||
})
|
||||
// A retryable failure is the endpoint's fault; a client disconnect is not.
|
||||
release(perr == nil || errors.Is(perr, proxy.ErrClientGone))
|
||||
|
||||
@@ -0,0 +1,362 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// threadHeader names the conversation a request continues.
|
||||
//
|
||||
// A header rather than a body field so the JSON stays exactly what an OpenAI
|
||||
// client would send: some SDKs drop fields they do not know, and every one of
|
||||
// them lets a caller set a header. It sits with the other x-llmgw-* headers
|
||||
// the router already reads.
|
||||
const threadHeader = "x-llmgw-thread"
|
||||
|
||||
// trimmedHeader reports how many stored turns were left out to make the
|
||||
// conversation fit the model that served it.
|
||||
//
|
||||
// Dropping turns is the right answer — a thread that outgrows a small model
|
||||
// would otherwise be permanently unusable on it, and after a budget fallback
|
||||
// that is not a choice the caller made — but it must not be a silent one.
|
||||
const trimmedHeader = "x-llmgw-thread-trimmed"
|
||||
|
||||
// threadMessage is one turn as it goes to an engine.
|
||||
type threadMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content json.RawMessage `json:"content"`
|
||||
}
|
||||
|
||||
// spliceThread puts a stored conversation in front of the request's own
|
||||
// messages and returns the body to send upstream.
|
||||
//
|
||||
// The history goes first and the request's messages last, because the request
|
||||
// carries the new turn: a client using a thread sends what it just said, not
|
||||
// the whole conversation again. That is the point of the feature.
|
||||
//
|
||||
// budget is how many bytes of prompt the serving model can take, measured the
|
||||
// same way promptBytes measures it so the two cannot disagree. Zero means the
|
||||
// model declares no context length, and then nothing is trimmed — guessing a
|
||||
// limit for a model that never stated one would drop turns for no reason.
|
||||
func spliceThread(body []byte, history []store.Message, budget int) ([]byte, int, error) {
|
||||
if len(history) == 0 {
|
||||
return body, 0, nil
|
||||
}
|
||||
|
||||
// Decoded into a map so every field the client sent survives untouched.
|
||||
// The router forwards bodies it does not understand, and a thread must not
|
||||
// be the thing that starts dropping them.
|
||||
var raw map[string]json.RawMessage
|
||||
if err := json.Unmarshal(body, &raw); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
var own []threadMessage
|
||||
if m, ok := raw["messages"]; ok {
|
||||
if err := json.Unmarshal(m, &own); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
}
|
||||
|
||||
stored := make([]threadMessage, 0, len(history))
|
||||
for _, m := range history {
|
||||
c, err := json.Marshal(m.Content)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
stored = append(stored, threadMessage{Role: m.Role, Content: c})
|
||||
}
|
||||
|
||||
dropped := 0
|
||||
if budget > 0 {
|
||||
stored, dropped = trimToFit(stored, own, budget)
|
||||
}
|
||||
|
||||
merged, err := json.Marshal(append(stored, own...))
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
raw["messages"] = merged
|
||||
|
||||
out, err := json.Marshal(raw)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
return out, dropped, nil
|
||||
}
|
||||
|
||||
// trimToFit drops the oldest stored turns until the conversation fits.
|
||||
//
|
||||
// The request's own messages are never dropped: they are what the caller just
|
||||
// said, and a request that cannot fit without them is one to refuse rather
|
||||
// than to answer with something the caller did not ask.
|
||||
//
|
||||
// System messages are kept wherever they are. They are usually instructions
|
||||
// that shape every answer, and dropping one silently changes what the model is
|
||||
// for — which is a stranger failure than a forgotten turn.
|
||||
func trimToFit(stored, own []threadMessage, budget int) ([]threadMessage, int) {
|
||||
fixed := messagesBytes(own)
|
||||
for _, m := range stored {
|
||||
if m.Role == "system" {
|
||||
fixed += messageBytes(m)
|
||||
}
|
||||
}
|
||||
|
||||
// Oldest first, which is the order they are in.
|
||||
keep := make([]bool, len(stored))
|
||||
for i := range stored {
|
||||
keep[i] = true
|
||||
}
|
||||
used := fixed
|
||||
for _, m := range stored {
|
||||
if m.Role != "system" {
|
||||
used += messageBytes(m)
|
||||
}
|
||||
}
|
||||
|
||||
dropped := 0
|
||||
for i := range stored {
|
||||
if used <= budget {
|
||||
break
|
||||
}
|
||||
if stored[i].Role == "system" {
|
||||
continue
|
||||
}
|
||||
used -= messageBytes(stored[i])
|
||||
keep[i] = false
|
||||
dropped++
|
||||
}
|
||||
|
||||
out := make([]threadMessage, 0, len(stored))
|
||||
for i, m := range stored {
|
||||
if keep[i] {
|
||||
out = append(out, m)
|
||||
}
|
||||
}
|
||||
return out, dropped
|
||||
}
|
||||
|
||||
// messageBytes measures one turn the way promptBytes measures it, so a
|
||||
// conversation trimmed to fit is measured by the same ruler that decides
|
||||
// whether it fits.
|
||||
func messageBytes(m threadMessage) int {
|
||||
n := len(m.Role) + 2
|
||||
var s string
|
||||
if err := json.Unmarshal(m.Content, &s); err == nil {
|
||||
return n + len(s)
|
||||
}
|
||||
return n + len(m.Content)
|
||||
}
|
||||
|
||||
func messagesBytes(ms []threadMessage) int {
|
||||
n := 0
|
||||
for _, m := range ms {
|
||||
n += messageBytes(m)
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// ownMessages reads the turns a request carries, for storing after a reply.
|
||||
func ownMessages(body []byte) []store.Message {
|
||||
var b struct {
|
||||
Messages []threadMessage `json:"messages"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &b); err != nil {
|
||||
return nil
|
||||
}
|
||||
out := make([]store.Message, 0, len(b.Messages))
|
||||
for _, m := range b.Messages {
|
||||
var s string
|
||||
if err := json.Unmarshal(m.Content, &s); err != nil {
|
||||
// Structured content — an image, an array of parts. Stored as the
|
||||
// JSON it arrived as, so replaying it gives the model what the
|
||||
// client originally sent rather than a flattened version of it.
|
||||
s = string(m.Content)
|
||||
}
|
||||
out = append(out, store.Message{Role: m.Role, Content: s})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// maxStoredReply bounds what one turn may add to a conversation.
|
||||
//
|
||||
// A cap rather than no cap because this is memory held on the request path for
|
||||
// the length of a reply. Generous enough that ordinary answers are stored
|
||||
// whole, small enough that a runaway generation cannot cost the gateway more
|
||||
// than this per concurrent request.
|
||||
const maxStoredReply = 128 << 10
|
||||
|
||||
// captureBytes says how much of the reply the proxy should hand back.
|
||||
//
|
||||
// Zero for every request not continuing a conversation, which is what keeps
|
||||
// the cost of this feature off the path of requests that do not use it.
|
||||
func (p *parsedRequest) captureBytes() int {
|
||||
if p.ThreadID == "" {
|
||||
return 0
|
||||
}
|
||||
return maxStoredReply
|
||||
}
|
||||
|
||||
// loadThread splices a stored conversation into the request.
|
||||
//
|
||||
// Returns how many turns were dropped to make it fit, or an error to refuse
|
||||
// the request with. A thread on a project that keeps no conversations is
|
||||
// refused rather than ignored: ignoring it would send the model a conversation
|
||||
// with no history, and it would answer confidently about something it was
|
||||
// never told — which reads as the model being wrong rather than as a setting
|
||||
// being off.
|
||||
func (r *Router) loadThread(ctx context.Context, req *parsedRequest) (int, error) {
|
||||
if req.ThreadID == "" {
|
||||
return 0, nil
|
||||
}
|
||||
project, err := r.threadProject(req.ProjectID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
req.ThreadProject = project
|
||||
|
||||
// Asked before anything else, because a project that keeps no
|
||||
// conversations must be told so rather than served. Without this the
|
||||
// request succeeds, nothing is stored, and the caller finds out only when
|
||||
// the model has no idea what they are referring to — which reads as the
|
||||
// model being broken. The e2e test that caught this asserted the refusal
|
||||
// and got a cheerful 200.
|
||||
retention, err := r.threads.ThreadRetention(ctx, project)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if retention == 0 {
|
||||
return 0, store.ErrThreadsNotEnabled
|
||||
}
|
||||
|
||||
// What the caller sent, before the history is spliced in, so the exchange
|
||||
// can be stored afterwards without re-reading a body that has changed.
|
||||
req.ThreadTurns = ownMessages(req.UpstreamBody)
|
||||
|
||||
history, err := r.threads.Thread(ctx, project, req.ThreadID)
|
||||
if err != nil && !errors.Is(err, store.ErrNoThread) {
|
||||
return 0, err
|
||||
}
|
||||
// A thread that does not exist yet is the first turn of a new one, which is
|
||||
// the ordinary way to start.
|
||||
if len(history) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
_, entry := r.resolve(req.Model)
|
||||
budget := 0
|
||||
if entry.ContextLength > 0 {
|
||||
// The same four-bytes-per-token estimate the rest of the router uses,
|
||||
// less the output the caller asked for, so trimming and the fit check
|
||||
// cannot disagree about what fits.
|
||||
budget = (entry.ContextLength - req.MaxTokens) * 4
|
||||
if budget < 0 {
|
||||
budget = 0
|
||||
}
|
||||
}
|
||||
|
||||
body, dropped, err := spliceThread(req.UpstreamBody, history, budget)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
req.UpstreamBody = body
|
||||
req.PromptTokens = len(body) / 4
|
||||
return dropped, nil
|
||||
}
|
||||
|
||||
// saveThread stores the exchange once a reply has been produced.
|
||||
//
|
||||
// Best effort, and deliberately after the client has been served: a
|
||||
// conversation that failed to record is a worse outcome than a reply that
|
||||
// never arrived, but only just, and the caller has their answer either way.
|
||||
// The failure is logged rather than swallowed, because a thread silently not
|
||||
// growing is the kind of fault someone discovers days later.
|
||||
func (r *Router) saveThread(ctx context.Context, req *parsedRequest, reply string, model string) {
|
||||
if req.ThreadID == "" || r.threads == nil || req.ThreadProject == "" || reply == "" {
|
||||
return
|
||||
}
|
||||
turns := append([]store.Message(nil), req.ThreadTurns...)
|
||||
turns = append(turns, store.Message{Role: "assistant", Content: reply, Model: model})
|
||||
|
||||
if err := r.threads.AppendThread(ctx, req.ThreadProject, req.ThreadID, turns); err != nil {
|
||||
r.log.Warn("could not store a conversation turn",
|
||||
"thread", req.ThreadID, "project", req.ThreadProject, "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
// anonymousProject owns everything kept by a gateway with no authentication.
|
||||
const anonymousProject = "_anonymous"
|
||||
|
||||
// errThreadsNeedAuth is refused rather than ignored, for the same reason a
|
||||
// disabled project is.
|
||||
var errThreadsNeedAuth = errors.New(
|
||||
"a conversation needs an API key: without one there is no project to keep it in")
|
||||
|
||||
// errThreadsOff is what a thread header meets when nothing is configured to
|
||||
// keep threads. Refused rather than ignored: ignoring it sends the model a
|
||||
// conversation with no history, and it answers confidently about something it
|
||||
// was never told.
|
||||
var errThreadsOff = errors.New(
|
||||
"this gateway keeps no conversations: enable auth with a project that keeps " +
|
||||
"threads, or set threads.store and threads.retain_hours")
|
||||
|
||||
// threadProject says which project a request's threads and notes belong to.
|
||||
//
|
||||
// The caller's own project when authenticated. The shared anonymous one only
|
||||
// when the operator configured it, because with no identity every caller's
|
||||
// history lands in one pool readable by anyone who knows an id.
|
||||
func (r *Router) threadProject(callerProject string) (string, error) {
|
||||
switch {
|
||||
case r.threads == nil:
|
||||
return "", errThreadsOff
|
||||
case callerProject != "":
|
||||
return callerProject, nil
|
||||
case r.anonymous:
|
||||
return anonymousProject, nil
|
||||
default:
|
||||
return "", errThreadsNeedAuth
|
||||
}
|
||||
}
|
||||
|
||||
// handleDeleteThread forgets a conversation.
|
||||
//
|
||||
// A client endpoint rather than an admin one: the person who said the thing is
|
||||
// the person who should be able to unsay it, and making them ask an operator
|
||||
// would mean the operator reads it first.
|
||||
func (r *Router) handleDeleteThread(w http.ResponseWriter, req *http.Request) {
|
||||
id := req.PathValue("id")
|
||||
if r.threads == nil {
|
||||
writeError(w, http.StatusNotFound, "invalid_request_error", "no such thread")
|
||||
return
|
||||
}
|
||||
var caller string
|
||||
if p, ok := auth.FromContext(req.Context()); ok {
|
||||
caller = p.ProjectID
|
||||
}
|
||||
project, err := r.threadProject(caller)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusUnauthorized, "invalid_request_error",
|
||||
"deleting a conversation needs the key that owns it")
|
||||
return
|
||||
}
|
||||
|
||||
// The store scopes the delete to the project, so a thread belonging to
|
||||
// another one is "no such thread" — the same answer as one that never
|
||||
// existed, because two answers would let a caller map which ids are real
|
||||
// in projects they cannot read.
|
||||
if err := r.threads.DeleteThread(req.Context(), project, id); err != nil {
|
||||
if errors.Is(err, store.ErrNoThread) {
|
||||
writeError(w, http.StatusNotFound, "invalid_request_error", "no such thread")
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusInternalServerError, "server_error", err.Error())
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
@@ -0,0 +1,176 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// splicedRoles reads back the roles and contents of a spliced body, which is
|
||||
// what the engine will actually see.
|
||||
func splicedRoles(t *testing.T, body []byte) []string {
|
||||
t.Helper()
|
||||
var b struct {
|
||||
Messages []struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
} `json:"messages"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &b); err != nil {
|
||||
t.Fatalf("the spliced body is not valid JSON: %v\n%s", err, body)
|
||||
}
|
||||
out := make([]string, 0, len(b.Messages))
|
||||
for _, m := range b.Messages {
|
||||
out = append(out, m.Role+":"+m.Content)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestHistoryGoesBeforeWhatTheCallerJustSaid(t *testing.T) {
|
||||
// The whole point of a thread: the client sends the new turn, not the
|
||||
// conversation again. Order is what makes the reply make sense.
|
||||
body := []byte(`{"model":"m","messages":[{"role":"user","content":"third"}]}`)
|
||||
history := []store.Message{
|
||||
{Role: "user", Content: "first"},
|
||||
{Role: "assistant", Content: "second"},
|
||||
}
|
||||
|
||||
out, dropped, err := spliceThread(body, history, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if dropped != 0 {
|
||||
t.Errorf("nothing should have been dropped, %d was", dropped)
|
||||
}
|
||||
got := splicedRoles(t, out)
|
||||
want := []string{"user:first", "assistant:second", "user:third"}
|
||||
if strings.Join(got, "|") != strings.Join(want, "|") {
|
||||
t.Errorf("messages came out as %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEverythingElseInTheBodySurvives(t *testing.T) {
|
||||
// The router forwards fields it does not understand — that is how engine
|
||||
// specific options reach an engine at all. Splicing a thread must not be
|
||||
// the thing that starts dropping them.
|
||||
body := []byte(`{"model":"m","messages":[{"role":"user","content":"hi"}],` +
|
||||
`"temperature":0.2,"stream":true,"some_engine_option":{"nested":[1,2]}}`)
|
||||
|
||||
out, _, err := spliceThread(body, []store.Message{{Role: "user", Content: "old"}}, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var got map[string]json.RawMessage
|
||||
if err := json.Unmarshal(out, &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, k := range []string{"model", "temperature", "stream", "some_engine_option"} {
|
||||
if _, ok := got[k]; !ok {
|
||||
t.Errorf("%q was dropped by splicing", k)
|
||||
}
|
||||
}
|
||||
if string(got["some_engine_option"]) != `{"nested":[1,2]}` {
|
||||
t.Errorf("a field llmgw does not understand was rewritten: %s", got["some_engine_option"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheOldestTurnsAreDroppedFirst(t *testing.T) {
|
||||
// A conversation that outgrew the model it is now being served by — after
|
||||
// a budget fallback, that is not a choice the caller made, so the thread
|
||||
// has to stay usable.
|
||||
history := []store.Message{
|
||||
{Role: "user", Content: strings.Repeat("a", 100)},
|
||||
{Role: "assistant", Content: strings.Repeat("b", 100)},
|
||||
{Role: "user", Content: strings.Repeat("c", 100)},
|
||||
{Role: "assistant", Content: "keep me"},
|
||||
}
|
||||
body := []byte(`{"model":"m","messages":[{"role":"user","content":"now"}]}`)
|
||||
|
||||
out, dropped, err := spliceThread(body, history, 200)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if dropped == 0 {
|
||||
t.Fatal("nothing was dropped although the history was over budget")
|
||||
}
|
||||
got := strings.Join(splicedRoles(t, out), "|")
|
||||
if !strings.Contains(got, "keep me") {
|
||||
t.Errorf("the most recent stored turn was dropped: %s", got)
|
||||
}
|
||||
if !strings.Contains(got, "user:now") {
|
||||
t.Errorf("the caller's own turn was dropped: %s", got)
|
||||
}
|
||||
if strings.Contains(got, strings.Repeat("a", 100)) {
|
||||
t.Errorf("the oldest turn survived while newer ones were dropped: %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestASystemMessageIsNeverDropped(t *testing.T) {
|
||||
// It is usually the instruction that shapes every answer. Losing it
|
||||
// silently changes what the model is for, which is a stranger failure than
|
||||
// a forgotten turn — and much harder to spot, because the model keeps
|
||||
// answering.
|
||||
history := []store.Message{
|
||||
{Role: "system", Content: "You only answer in Turkish."},
|
||||
{Role: "user", Content: strings.Repeat("x", 500)},
|
||||
{Role: "assistant", Content: strings.Repeat("y", 500)},
|
||||
}
|
||||
body := []byte(`{"model":"m","messages":[{"role":"user","content":"now"}]}`)
|
||||
|
||||
out, dropped, err := spliceThread(body, history, 100)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if dropped != 2 {
|
||||
t.Errorf("dropped %d turns, expected both non-system ones", dropped)
|
||||
}
|
||||
if got := strings.Join(splicedRoles(t, out), "|"); !strings.Contains(got, "only answer in Turkish") {
|
||||
t.Errorf("the system message was dropped: %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAModelWithNoDeclaredContextTrimsNothing(t *testing.T) {
|
||||
// Guessing a limit for a model that never stated one would drop turns for
|
||||
// no reason, and the caller would never learn why the model forgot.
|
||||
history := make([]store.Message, 50)
|
||||
for i := range history {
|
||||
history[i] = store.Message{Role: "user", Content: strings.Repeat("z", 1000)}
|
||||
}
|
||||
body := []byte(`{"model":"m","messages":[{"role":"user","content":"now"}]}`)
|
||||
|
||||
_, dropped, err := spliceThread(body, history, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if dropped != 0 {
|
||||
t.Errorf("%d turns were dropped for a model that declares no context", dropped)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoHistoryLeavesTheBodyByteForByte(t *testing.T) {
|
||||
// A request without a thread, and the first request of a new one, must
|
||||
// take exactly the path they took before this existed.
|
||||
body := []byte(`{"model":"m","messages":[{"role":"user","content":"hi"}]}`)
|
||||
out, dropped, err := spliceThread(body, nil, 4096)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if dropped != 0 || string(out) != string(body) {
|
||||
t.Errorf("an empty history rewrote the body:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStructuredContentIsStoredAsItArrived(t *testing.T) {
|
||||
// An image or an array of parts. Flattening it to a string would replay
|
||||
// something the client never sent.
|
||||
body := []byte(`{"model":"m","messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
|
||||
got := ownMessages(body)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("read %d messages", len(got))
|
||||
}
|
||||
if !strings.HasPrefix(got[0].Content, `[{"type":"text"`) {
|
||||
t.Errorf("structured content was flattened: %q", got[0].Content)
|
||||
}
|
||||
}
|
||||
@@ -92,7 +92,7 @@ func (s *SQLite) ConsumeClaim(ctx context.Context, token string) (string, error)
|
||||
hash := HashKey(token)
|
||||
now := time.Now()
|
||||
|
||||
tx, err := s.db.BeginTx(ctx, nil)
|
||||
q, tx, err := s.begin(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
@@ -100,7 +100,7 @@ func (s *SQLite) ConsumeClaim(ctx context.Context, token string) (string, error)
|
||||
|
||||
var sealed []byte
|
||||
var expires int64
|
||||
err = tx.QueryRowContext(ctx,
|
||||
err = q.QueryRowContext(ctx,
|
||||
`SELECT sealed, expires_at FROM key_claims WHERE token_hash = ? AND claimed_at IS NULL`,
|
||||
hash).Scan(&sealed, &expires)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
@@ -120,7 +120,7 @@ func (s *SQLite) ConsumeClaim(ctx context.Context, token string) (string, error)
|
||||
|
||||
// Blank the sealed value as well as marking it claimed, so a used link
|
||||
// leaves nothing behind even before the row is purged.
|
||||
res, err := tx.ExecContext(ctx,
|
||||
res, err := q.ExecContext(ctx,
|
||||
`UPDATE key_claims SET claimed_at = ?, sealed = x'' WHERE token_hash = ? AND claimed_at IS NULL`,
|
||||
millis(now), hash)
|
||||
if err != nil {
|
||||
|
||||
@@ -67,7 +67,7 @@ func TestTheDatabaseNeverHoldsAUsableKey(t *testing.T) {
|
||||
s, token, secret := claimFixture(t)
|
||||
|
||||
var sealed []byte
|
||||
if err := s.db.QueryRow(`SELECT sealed FROM key_claims`).Scan(&sealed); err != nil {
|
||||
if err := s.db.QueryRowContext(context.Background(), `SELECT sealed FROM key_claims`).Scan(&sealed); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(sealed) == 0 {
|
||||
@@ -79,7 +79,7 @@ func TestTheDatabaseNeverHoldsAUsableKey(t *testing.T) {
|
||||
// And the token itself is not there either, so nothing in the database can
|
||||
// derive the key that opens it.
|
||||
var count int
|
||||
if err := s.db.QueryRow(`SELECT COUNT(*) FROM key_claims WHERE token_hash = ?`, []byte(token)).
|
||||
if err := s.db.QueryRowContext(context.Background(), `SELECT COUNT(*) FROM key_claims WHERE token_hash = ?`, []byte(token)).
|
||||
Scan(&count); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -98,7 +98,7 @@ func TestAClaimedLinkLeavesNothingBehind(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var sealed []byte
|
||||
if err := s.db.QueryRow(`SELECT sealed FROM key_claims`).Scan(&sealed); err != nil {
|
||||
if err := s.db.QueryRowContext(context.Background(), `SELECT sealed FROM key_claims`).Scan(&sealed); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(sealed) != 0 {
|
||||
@@ -149,7 +149,7 @@ func TestAnExpiredLinkIsRefused(t *testing.T) {
|
||||
// Age it directly. CreateClaim treats a non-positive TTL as "use the
|
||||
// default", which is the right thing for a caller and no way to build the
|
||||
// row this test needs.
|
||||
if _, err := s.db.Exec(`UPDATE key_claims SET expires_at = ?`,
|
||||
if _, err := s.db.ExecContext(context.Background(), `UPDATE key_claims SET expires_at = ?`,
|
||||
time.Now().Add(-time.Hour).UnixMilli()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,219 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// copyOrder is every table, parents before children.
|
||||
//
|
||||
// Written out rather than discovered, because the order is a fact about the
|
||||
// foreign keys and getting it wrong produces a constraint violation halfway
|
||||
// through — with some tables copied and some not, which is the worst place to
|
||||
// stop. A table added without being added here is copied by nothing, and the
|
||||
// destination comes up quietly missing it; the test below is what catches
|
||||
// that.
|
||||
var copyOrder = []string{
|
||||
"organizations",
|
||||
"teams",
|
||||
"projects",
|
||||
// Before api_keys, which references it. The first version of this list
|
||||
// followed the order the tables are created across the migrations, and
|
||||
// users is created in a later one than api_keys even though api_keys
|
||||
// points at it — so the copy failed on a foreign key three tables in.
|
||||
"users",
|
||||
"api_keys",
|
||||
"key_claims",
|
||||
"sessions",
|
||||
"limits",
|
||||
"usage",
|
||||
"models",
|
||||
"threads",
|
||||
"thread_messages",
|
||||
"memory_notes",
|
||||
}
|
||||
|
||||
// notCopied are tables deliberately left behind.
|
||||
//
|
||||
// Leadership is a claim about a process that is running right now. Copying it
|
||||
// would hand the new store a leader that does not exist, and the standby
|
||||
// waiting on that lease would wait out its full expiry before noticing — at
|
||||
// the exact moment somebody is watching a migration and wondering why nothing
|
||||
// is being placed.
|
||||
var notCopied = map[string]string{
|
||||
"leadership": "a lease on a running process, meaningless in a new store",
|
||||
}
|
||||
|
||||
// CopyResult is how many rows each table received.
|
||||
type CopyResult struct {
|
||||
Table string
|
||||
Rows int64
|
||||
}
|
||||
|
||||
// CopyTo copies every row into another store.
|
||||
//
|
||||
// For moving from SQLite to PostgreSQL, which otherwise means reissuing every
|
||||
// API key: the keys are stored hashed, so there is no way to recreate them and
|
||||
// no way for anyone to keep using the one they have.
|
||||
//
|
||||
// The destination must be empty. Merging two stores would mean deciding what
|
||||
// happens when both have a project with the same id and different contents,
|
||||
// and every answer to that is a guess about which one someone meant.
|
||||
//
|
||||
// Not a transaction across the whole copy. It is one transaction per table,
|
||||
// because a single transaction holding every row of a large usage table is how
|
||||
// a migration runs out of memory on the machine least able to spare it. A
|
||||
// failure therefore leaves a partial copy, which is why the destination has to
|
||||
// have been empty: starting again is dropping and recreating it.
|
||||
func (s *SQLite) CopyTo(ctx context.Context, dst *SQLite) ([]CopyResult, error) {
|
||||
empty, where, err := dst.isEmpty(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !empty {
|
||||
return nil, fmt.Errorf("the destination already has rows in %s; copy into an empty store", where)
|
||||
}
|
||||
|
||||
var out []CopyResult
|
||||
for _, table := range copyOrder {
|
||||
n, err := s.copyTable(ctx, dst, table)
|
||||
if err != nil {
|
||||
return out, fmt.Errorf("copying %s: %w", table, err)
|
||||
}
|
||||
out = append(out, CopyResult{Table: table, Rows: n})
|
||||
}
|
||||
// Part of the copy, not a step someone has to remember: a destination
|
||||
// whose sequences still start at 1 is not a destination anyone can use.
|
||||
if err := dst.resetIdentities(ctx); err != nil {
|
||||
return out, fmt.Errorf("pointing the identity columns past the copied rows: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// isEmpty reports whether any table has rows, and which one first does.
|
||||
func (s *SQLite) isEmpty(ctx context.Context) (bool, string, error) {
|
||||
for _, table := range copyOrder {
|
||||
var n int64
|
||||
//nolint:gosec // table comes from copyOrder, not from input
|
||||
if err := s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM `+table).Scan(&n); err != nil {
|
||||
return false, "", err
|
||||
}
|
||||
if n > 0 {
|
||||
return false, table, nil
|
||||
}
|
||||
}
|
||||
return true, "", nil
|
||||
}
|
||||
|
||||
// copyTable moves one table, column names taken from the source.
|
||||
//
|
||||
// Reading the columns rather than listing them keeps this working when a
|
||||
// migration adds one: a hand-written column list would copy the old shape and
|
||||
// silently drop the new field, and nothing would notice until someone looked
|
||||
// for a setting that was there before the move.
|
||||
func (s *SQLite) copyTable(ctx context.Context, dst *SQLite, table string) (int64, error) {
|
||||
//nolint:gosec // table comes from copyOrder, not from input
|
||||
rows, err := s.db.QueryContext(ctx, `SELECT * FROM `+table)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
cols, err := rows.Columns()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
q, tx, err := dst.begin(ctx)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
insert := `INSERT INTO ` + table + ` (` + strings.Join(quoteAll(cols), ", ") +
|
||||
`) VALUES (` + strings.Repeat("?, ", len(cols)-1) + `?)`
|
||||
|
||||
var n int64
|
||||
for rows.Next() {
|
||||
vals := make([]any, len(cols))
|
||||
ptrs := make([]any, len(cols))
|
||||
for i := range vals {
|
||||
ptrs[i] = &vals[i]
|
||||
}
|
||||
if err := rows.Scan(ptrs...); err != nil {
|
||||
return n, err
|
||||
}
|
||||
for i, v := range vals {
|
||||
vals[i] = keepEmptyBytes(v)
|
||||
}
|
||||
if _, err := q.ExecContext(ctx, insert, vals...); err != nil {
|
||||
return n, err
|
||||
}
|
||||
n++
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return n, err
|
||||
}
|
||||
return n, tx.Commit()
|
||||
}
|
||||
|
||||
// keepEmptyBytes stops an empty byte column becoming NULL on the way across.
|
||||
//
|
||||
// An empty BLOB scans as a typed nil slice — []byte(nil) inside a non-nil any
|
||||
// — while a real NULL scans as an untyped nil. The two are distinguishable
|
||||
// here and not in pgx, which encodes any nil slice as NULL, so the difference
|
||||
// has to be restored before the insert.
|
||||
//
|
||||
// It is not a detail. key_claims.sealed is emptied the moment a claim is used,
|
||||
// so every claimed row holds an empty blob in a NOT NULL column, and the copy
|
||||
// failed on the first one — after the gateway had been stopped and four tables
|
||||
// had already been written.
|
||||
func keepEmptyBytes(v any) any {
|
||||
if b, ok := v.([]byte); ok && b == nil {
|
||||
return []byte{}
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// quoteAll quotes column names, because `key` and `text` are reserved words in
|
||||
// one dialect or the other and both are column names here.
|
||||
func quoteAll(cols []string) []string {
|
||||
out := make([]string, len(cols))
|
||||
for i, c := range cols {
|
||||
out[i] = `"` + c + `"`
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// resetIdentities points PostgreSQL's identity columns past the copied rows.
|
||||
//
|
||||
// Without it the first insert after a migration collides with a row that was
|
||||
// copied in: the sequence still starts at 1 while the table already holds a
|
||||
// thousand notes, and the failure arrives later, to somebody else, looking
|
||||
// like a bug in the feature rather than in the move.
|
||||
func (s *SQLite) resetIdentities(ctx context.Context) error {
|
||||
if s.dialect != PostgresDialect {
|
||||
return nil
|
||||
}
|
||||
// Only the tables with a generated id. SQLite's rowid needs nothing.
|
||||
for _, t := range []string{"memory_notes", "usage"} {
|
||||
var has bool
|
||||
if err := s.db.QueryRowContext(ctx,
|
||||
`SELECT EXISTS (SELECT 1 FROM information_schema.columns
|
||||
WHERE table_name = ? AND column_name = 'id'
|
||||
AND is_identity = 'YES')`, t).Scan(&has); err != nil {
|
||||
return err
|
||||
}
|
||||
if !has {
|
||||
continue
|
||||
}
|
||||
//nolint:gosec // t is a literal from the list above
|
||||
if _, err := s.db.ExecContext(ctx,
|
||||
`SELECT setval(pg_get_serial_sequence('`+t+`', 'id'),
|
||||
GREATEST((SELECT COALESCE(MAX(id), 0) FROM `+t+`), 1))`); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,261 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestEveryTableIsInTheCopyOrder(t *testing.T) {
|
||||
// A table added by a migration and not added here is copied by nothing,
|
||||
// and the destination comes up quietly missing it — discovered when
|
||||
// somebody looks for a setting that was there before the move. Read out of
|
||||
// the migrations rather than listed twice, so the two cannot disagree.
|
||||
var declared []string
|
||||
entries, err := migrationFS.ReadDir("migrations/sqlite")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, e := range entries {
|
||||
body, err := migrationFS.ReadFile("migrations/sqlite/" + e.Name())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, line := range strings.Split(string(body), "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
rest, ok := strings.CutPrefix(line, "CREATE TABLE ")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
name, _, _ := strings.Cut(rest, " ")
|
||||
name = strings.TrimSuffix(strings.TrimSpace(name), "(")
|
||||
declared = append(declared, name)
|
||||
}
|
||||
}
|
||||
if len(declared) == 0 {
|
||||
t.Fatal("found no tables in the migrations, so this test proves nothing")
|
||||
}
|
||||
for _, name := range declared {
|
||||
if slices.Contains(copyOrder, name) {
|
||||
continue
|
||||
}
|
||||
// Left behind on purpose, with the reason written down beside it.
|
||||
if why, ok := notCopied[name]; ok {
|
||||
if why == "" {
|
||||
t.Errorf("%q is excluded from the copy with no reason given", name)
|
||||
}
|
||||
continue
|
||||
}
|
||||
t.Errorf("%q is created by a migration but nothing copies it", name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyRefusesADestinationThatAlreadyHasRows(t *testing.T) {
|
||||
// Merging two stores means deciding what happens when both hold a project
|
||||
// with the same id and different contents, and every answer is a guess
|
||||
// about which one somebody meant.
|
||||
ctx := context.Background()
|
||||
src, dst := openDB(t), openDB(t)
|
||||
tree(t, src)
|
||||
tree(t, dst)
|
||||
|
||||
_, err := src.CopyTo(ctx, dst)
|
||||
if err == nil || !strings.Contains(err.Error(), "empty store") {
|
||||
t.Errorf("copying into a populated store gave %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestACopiedStoreStillAuthenticatesTheSameKey(t *testing.T) {
|
||||
// The reason this exists. Keys are stored hashed, so a migration that
|
||||
// loses them cannot recreate them: every person and every script would
|
||||
// need a new one issued.
|
||||
ctx := context.Background()
|
||||
src := openDB(t)
|
||||
proj := tree(t, src)
|
||||
secret, k := mustKey(t, src, "k1", proj)
|
||||
|
||||
if err := src.WriteUsage(ctx, []Usage{{
|
||||
At: time.Now(), KeyID: k.ID, ProjectID: proj, Model: "m", ServedModel: "m",
|
||||
PromptTokens: 3, CompletionTokens: 4, CostMicros: 99, Status: 200,
|
||||
}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
dst, err := Open(ctx, filepath.Join(t.TempDir(), "copy.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer dst.Close()
|
||||
|
||||
res, err := src.CopyTo(ctx, dst)
|
||||
if err != nil {
|
||||
t.Fatalf("copying: %v", err)
|
||||
}
|
||||
if len(res) != len(copyOrder) {
|
||||
t.Errorf("copied %d tables of %d", len(res), len(copyOrder))
|
||||
}
|
||||
|
||||
// The same secret the caller already holds still resolves, in the copy.
|
||||
got, err := dst.KeyByHash(ctx, HashKey(secret))
|
||||
if err != nil {
|
||||
t.Fatalf("the copied store does not know the key: %v", err)
|
||||
}
|
||||
if got.ID != k.ID || got.ProjectID != proj {
|
||||
t.Errorf("key came back as %+v", got)
|
||||
}
|
||||
|
||||
// And the spend it had, so a budget is not quietly reset by moving.
|
||||
spend, err := dst.SpendSince(ctx, time.Now().Add(-time.Hour))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if spend[k.ID] != 99 {
|
||||
t.Errorf("spend after the copy = %d, want 99", spend[k.ID])
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyCarriesConversationsAndNotes(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
src := openDB(t)
|
||||
proj := tree(t, src)
|
||||
if err := src.SetThreadRetention(ctx, proj, 24); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := src.AppendThread(ctx, proj, "th1", []Message{
|
||||
{Role: "user", Content: "remembered"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
dst, err := Open(ctx, filepath.Join(t.TempDir(), "copy.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer dst.Close()
|
||||
if _, err := src.CopyTo(ctx, dst); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
msgs, err := dst.Thread(ctx, proj, "th1")
|
||||
if err != nil {
|
||||
t.Fatalf("the conversation did not survive: %v", err)
|
||||
}
|
||||
if len(msgs) != 1 || msgs[0].Content != "remembered" {
|
||||
t.Errorf("thread came back as %+v", msgs)
|
||||
}
|
||||
// The retention setting is a column on projects, so it rides along — a
|
||||
// destination that kept the conversations and forgot the promise about
|
||||
// how long would be worse than one that kept neither.
|
||||
if d, err := dst.ThreadRetention(ctx, proj); err != nil || d == 0 {
|
||||
t.Errorf("retention after the copy = %v (err %v)", d, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyingAnEmptyStoreIsFine(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
src := openDB(t)
|
||||
dst, err := Open(ctx, filepath.Join(t.TempDir(), "copy.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer dst.Close()
|
||||
if _, err := src.CopyTo(ctx, dst); err != nil {
|
||||
t.Errorf("copying an empty store failed: %v", err)
|
||||
}
|
||||
_ = os.Remove("")
|
||||
}
|
||||
|
||||
func TestTheCopyOrderSatisfiesTheForeignKeys(t *testing.T) {
|
||||
// Completeness is not enough: a table copied before the one it references
|
||||
// fails on a constraint partway through, leaving some tables copied and
|
||||
// some not. The first version of copyOrder followed the order the
|
||||
// migrations create tables in, and users is created after api_keys even
|
||||
// though api_keys points at it — so the copy died three tables in, on a
|
||||
// real database, after the gateway had already been stopped.
|
||||
//
|
||||
// Derived from the migrations rather than restated, so a new reference
|
||||
// fails here rather than during somebody's migration.
|
||||
deps := map[string]map[string]bool{}
|
||||
note := func(table, ref string) {
|
||||
if table == ref {
|
||||
return
|
||||
}
|
||||
if deps[table] == nil {
|
||||
deps[table] = map[string]bool{}
|
||||
}
|
||||
deps[table][ref] = true
|
||||
}
|
||||
|
||||
entries, err := migrationFS.ReadDir("migrations/sqlite")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
body := ""
|
||||
for _, e := range entries {
|
||||
b, err := migrationFS.ReadFile("migrations/sqlite/" + e.Name())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
body += string(b)
|
||||
}
|
||||
|
||||
// Table bodies, so a REFERENCES is attributed to the table holding it.
|
||||
for _, m := range regexpCreateTable.FindAllStringSubmatch(body, -1) {
|
||||
for _, r := range regexpReferences.FindAllStringSubmatch(m[2], -1) {
|
||||
note(m[1], r[1])
|
||||
}
|
||||
}
|
||||
// And columns added later, which is exactly how api_keys got its user_id.
|
||||
for _, m := range regexpAlterRef.FindAllStringSubmatch(body, -1) {
|
||||
note(m[1], m[2])
|
||||
}
|
||||
if len(deps) == 0 {
|
||||
t.Fatal("found no foreign keys, so this test proves nothing")
|
||||
}
|
||||
|
||||
position := map[string]int{}
|
||||
for i, t := range copyOrder {
|
||||
position[t] = i
|
||||
}
|
||||
for table, refs := range deps {
|
||||
for ref := range refs {
|
||||
if position[ref] > position[table] {
|
||||
t.Errorf("%s is copied before %s, which it references", table, ref)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var (
|
||||
regexpCreateTable = regexp.MustCompile(`(?s)CREATE TABLE (\w+)\s*\((.*?)\n\);`)
|
||||
regexpReferences = regexp.MustCompile(`REFERENCES\s+(\w+)`)
|
||||
regexpAlterRef = regexp.MustCompile(`ALTER TABLE (\w+) ADD COLUMN[^;]*REFERENCES\s+(\w+)`)
|
||||
)
|
||||
|
||||
func TestAnEmptyBlobDoesNotBecomeNullInTheCopy(t *testing.T) {
|
||||
// key_claims.sealed is emptied the moment a claim is used, so a claimed
|
||||
// row holds an empty blob in a NOT NULL column. An empty blob scans as a
|
||||
// typed nil slice and a real NULL as an untyped one; pgx cannot tell them
|
||||
// apart, so the distinction is restored before the insert.
|
||||
if got := keepEmptyBytes([]byte(nil)); got == nil {
|
||||
t.Error("an empty byte column would be inserted as NULL")
|
||||
} else if b, ok := got.([]byte); !ok || len(b) != 0 {
|
||||
t.Errorf("empty bytes came out as %#v", got)
|
||||
}
|
||||
// A genuine NULL arrives as an untyped nil and must stay one.
|
||||
if got := keepEmptyBytes(nil); got != nil {
|
||||
t.Errorf("a NULL became %#v", got)
|
||||
}
|
||||
// Everything else passes through.
|
||||
if got := keepEmptyBytes(int64(7)); got != int64(7) {
|
||||
t.Errorf("a number was rewritten to %#v", got)
|
||||
}
|
||||
if got := keepEmptyBytes([]byte("x")); string(got.([]byte)) != "x" {
|
||||
t.Errorf("a non-empty blob was rewritten to %#v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Dialect is which SQL database is underneath.
|
||||
//
|
||||
// There is one set of queries, not one per database, because the SQL this
|
||||
// store writes is almost entirely portable — the only construct that genuinely
|
||||
// differs is the placeholder, and `ON CONFLICT` upserts work in both. Two
|
||||
// copies of every query would drift, and the drift would be discovered one
|
||||
// query at a time in production.
|
||||
type Dialect string
|
||||
|
||||
// The supported databases.
|
||||
const (
|
||||
SQLiteDialect Dialect = "sqlite"
|
||||
PostgresDialect Dialect = "postgres"
|
||||
)
|
||||
|
||||
// rebind rewrites `?` placeholders for the dialect in use.
|
||||
//
|
||||
// SQLite takes `?`, PostgreSQL takes `$1`, `$2` and so on. Queries are written
|
||||
// with `?` because that is what the majority of them were written with
|
||||
// already, and rewriting is cheaper and safer than maintaining two spellings
|
||||
// of each.
|
||||
//
|
||||
// A `?` inside a string literal would be rewritten too, which would be a bug —
|
||||
// so this skips over quoted sections rather than scanning blindly. No query
|
||||
// here contains one today; this exists so that adding one is not a silent
|
||||
// corruption.
|
||||
func (d Dialect) rebind(query string) string {
|
||||
if d != PostgresDialect || !strings.Contains(query, "?") {
|
||||
return query
|
||||
}
|
||||
var b strings.Builder
|
||||
b.Grow(len(query) + 8)
|
||||
|
||||
n := 0
|
||||
inSingle, inDouble := false, false
|
||||
for i := 0; i < len(query); i++ {
|
||||
c := query[i]
|
||||
switch {
|
||||
case c == '\'' && !inDouble:
|
||||
inSingle = !inSingle
|
||||
case c == '"' && !inSingle:
|
||||
inDouble = !inDouble
|
||||
case c == '?' && !inSingle && !inDouble:
|
||||
n++
|
||||
b.WriteByte('$')
|
||||
b.WriteString(strconv.Itoa(n))
|
||||
continue
|
||||
}
|
||||
b.WriteByte(c)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// conn is the part of *sql.DB and *sql.Tx this store uses.
|
||||
//
|
||||
// Named so the same query helpers work inside and outside a transaction, and
|
||||
// so every path goes through the rebinding wrapper rather than some of them
|
||||
// reaching the driver directly with the wrong placeholders.
|
||||
type conn interface {
|
||||
ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
|
||||
QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
|
||||
QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
|
||||
// PrepareContext is here because a prepared statement carries the query
|
||||
// with it: rebinding at Exec time would be too late, the placeholders are
|
||||
// fixed when it is prepared.
|
||||
PrepareContext(ctx context.Context, query string) (*sql.Stmt, error)
|
||||
}
|
||||
|
||||
// bound is a conn that rewrites placeholders on the way through.
|
||||
//
|
||||
// A wrapper rather than a change at each call site: there are sixty of those,
|
||||
// and sixty chances to forget one. Here the only way to reach the driver is
|
||||
// through the rewrite.
|
||||
type bound struct {
|
||||
c Dialect
|
||||
// to is the underlying *sql.DB or *sql.Tx.
|
||||
to conn
|
||||
}
|
||||
|
||||
func (b bound) ExecContext(ctx context.Context, q string, args ...any) (sql.Result, error) {
|
||||
return b.to.ExecContext(ctx, b.c.rebind(q), args...)
|
||||
}
|
||||
|
||||
func (b bound) QueryContext(ctx context.Context, q string, args ...any) (*sql.Rows, error) {
|
||||
return b.to.QueryContext(ctx, b.c.rebind(q), args...)
|
||||
}
|
||||
|
||||
func (b bound) QueryRowContext(ctx context.Context, q string, args ...any) *sql.Row {
|
||||
return b.to.QueryRowContext(ctx, b.c.rebind(q), args...)
|
||||
}
|
||||
|
||||
func (b bound) PrepareContext(ctx context.Context, q string) (*sql.Stmt, error) {
|
||||
return b.to.PrepareContext(ctx, b.c.rebind(q))
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
package store
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestSQLiteKeepsItsPlaceholders(t *testing.T) {
|
||||
q := "SELECT a FROM t WHERE b = ? AND c = ?"
|
||||
if got := SQLiteDialect.rebind(q); got != q {
|
||||
t.Errorf("sqlite query was rewritten: %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresGetsNumberedPlaceholders(t *testing.T) {
|
||||
got := PostgresDialect.rebind("SELECT a FROM t WHERE b = ? AND c = ?")
|
||||
want := "SELECT a FROM t WHERE b = $1 AND c = $2"
|
||||
if got != want {
|
||||
t.Errorf("rebind = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAQuestionMarkInsideALiteralIsLeftAlone(t *testing.T) {
|
||||
// No query here contains one today. This exists so that adding one is a
|
||||
// test failure rather than a query that silently reads the wrong argument
|
||||
// — or, worse, one that still parses and returns the wrong rows.
|
||||
got := PostgresDialect.rebind(`SELECT a FROM t WHERE name = 'why?' AND b = ?`)
|
||||
want := `SELECT a FROM t WHERE name = 'why?' AND b = $1`
|
||||
if got != want {
|
||||
t.Errorf("rebind = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDoubleQuotedIdentifiersAreAlsoSkipped(t *testing.T) {
|
||||
got := PostgresDialect.rebind(`SELECT "odd?name" FROM t WHERE b = ?`)
|
||||
want := `SELECT "odd?name" FROM t WHERE b = $1`
|
||||
if got != want {
|
||||
t.Errorf("rebind = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAQueryWithNoPlaceholdersIsUntouched(t *testing.T) {
|
||||
q := "DELETE FROM t"
|
||||
if got := PostgresDialect.rebind(q); got != q {
|
||||
t.Errorf("rebind = %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ErrNotLeader means somebody else holds the lease and it has not lapsed.
|
||||
var ErrNotLeader = errors.New("store: another control plane holds the lease")
|
||||
|
||||
// Lease is who is in charge and until when.
|
||||
type Lease struct {
|
||||
Owner string
|
||||
ExpiresAt time.Time
|
||||
AcquiredAt time.Time
|
||||
}
|
||||
|
||||
// AcquireLeadership takes or renews the lease for a role.
|
||||
//
|
||||
// Renewing and taking are the same statement on purpose: a leader that renews
|
||||
// takes the same path as one starting fresh, so there is no way for the two to
|
||||
// disagree about what a held lease looks like. It succeeds when the row is
|
||||
// absent, when this process already owns it, or when whoever owned it let it
|
||||
// lapse; otherwise it returns ErrNotLeader.
|
||||
//
|
||||
// The clock is the caller's rather than the database's, which is a compromise
|
||||
// worth naming. A database clock would be one clock for everyone; taking it
|
||||
// means a round trip that says now() inside every comparison, in SQL that is
|
||||
// spelled differently in each dialect. Two control planes on machines whose
|
||||
// clocks disagree by more than the lease will fight, so keep them on NTP —
|
||||
// which any cluster running InfiniBand already has to be.
|
||||
func (s *SQLite) AcquireLeadership(ctx context.Context, role, owner string, ttl time.Duration) (Lease, error) {
|
||||
now := time.Now()
|
||||
expires := now.Add(ttl)
|
||||
|
||||
// One statement, so two processes racing cannot both believe they won:
|
||||
// the WHERE decides, and the loser changes no rows.
|
||||
res, err := s.db.ExecContext(ctx,
|
||||
`INSERT INTO leadership (role, owner, expires_at, acquired_at) VALUES (?, ?, ?, ?)
|
||||
ON CONFLICT (role) DO UPDATE SET
|
||||
owner = excluded.owner,
|
||||
expires_at = excluded.expires_at,
|
||||
acquired_at = CASE WHEN leadership.owner = excluded.owner
|
||||
THEN leadership.acquired_at
|
||||
ELSE excluded.acquired_at END
|
||||
WHERE leadership.owner = excluded.owner OR leadership.expires_at < ?`,
|
||||
role, owner, expires.UnixMilli(), now.UnixMilli(), now.UnixMilli())
|
||||
if err != nil {
|
||||
return Lease{}, err
|
||||
}
|
||||
n, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return Lease{}, err
|
||||
}
|
||||
if n == 0 {
|
||||
// Lost, and the caller is told who to, because "not leader" without a
|
||||
// name is the start of a long afternoon. The lease is returned beside
|
||||
// the error rather than instead of it.
|
||||
held, lerr := s.Leadership(ctx, role)
|
||||
if lerr != nil {
|
||||
return Lease{}, lerr
|
||||
}
|
||||
return held, ErrNotLeader
|
||||
}
|
||||
return s.Leadership(ctx, role)
|
||||
}
|
||||
|
||||
// Leadership reports who holds a role's lease.
|
||||
func (s *SQLite) Leadership(ctx context.Context, role string) (Lease, error) {
|
||||
var owner string
|
||||
var expires, acquired int64
|
||||
err := s.db.QueryRowContext(ctx,
|
||||
`SELECT owner, expires_at, acquired_at FROM leadership WHERE role = ?`, role).
|
||||
Scan(&owner, &expires, &acquired)
|
||||
if err != nil {
|
||||
return Lease{}, err
|
||||
}
|
||||
return Lease{
|
||||
Owner: owner,
|
||||
ExpiresAt: fromMillis(expires),
|
||||
AcquiredAt: fromMillis(acquired),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ReleaseLeadership gives the lease up, so a standby takes over in seconds
|
||||
// rather than waiting out the clock.
|
||||
//
|
||||
// Only the holder can release it: a process that has already lost the lease
|
||||
// must not be able to take it from whoever has it now, which is exactly what
|
||||
// an unconditional delete during a slow shutdown would do.
|
||||
func (s *SQLite) ReleaseLeadership(ctx context.Context, role, owner string) error {
|
||||
_, err := s.db.ExecContext(ctx,
|
||||
`DELETE FROM leadership WHERE role = ? AND owner = ?`, role, owner)
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// leadershipSuite is the same assertions against whichever store it is given,
|
||||
// because the whole point of the lease is that it behaves identically on both
|
||||
// — and the one that matters for HA is the one it was never tested on.
|
||||
func leadershipSuite(t *testing.T, open func(*testing.T) *SQLite) {
|
||||
ctx := context.Background()
|
||||
|
||||
t.Run("the first to ask gets it", func(t *testing.T) {
|
||||
s := open(t)
|
||||
lease, err := s.AcquireLeadership(ctx, "control", "a", time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if lease.Owner != "a" {
|
||||
t.Errorf("owner = %q", lease.Owner)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("a second is refused and told who holds it", func(t *testing.T) {
|
||||
s := open(t)
|
||||
if _, err := s.AcquireLeadership(ctx, "control", "a", time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
held, err := s.AcquireLeadership(ctx, "control", "b", time.Minute)
|
||||
if !errors.Is(err, ErrNotLeader) {
|
||||
t.Fatalf("a second control plane got %v", err)
|
||||
}
|
||||
if held.Owner != "a" {
|
||||
t.Errorf("refused without naming the holder: %+v", held)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("the holder renews without losing when it started", func(t *testing.T) {
|
||||
s := open(t)
|
||||
first, err := s.AcquireLeadership(ctx, "control", "a", time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
again, err := s.AcquireLeadership(ctx, "control", "a", time.Minute)
|
||||
if err != nil {
|
||||
t.Fatalf("the holder could not renew: %v", err)
|
||||
}
|
||||
if !again.ExpiresAt.After(first.ExpiresAt) {
|
||||
t.Error("renewing did not move the expiry")
|
||||
}
|
||||
if !again.AcquiredAt.Equal(first.AcquiredAt) {
|
||||
t.Error("renewing reset when this leader took over")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("a lapsed lease is taken", func(t *testing.T) {
|
||||
// The reason a lease and not a lock: a leader that was powered off
|
||||
// cannot release anything, and nobody should have to be paged to make
|
||||
// a standby take over.
|
||||
s := open(t)
|
||||
if _, err := s.AcquireLeadership(ctx, "control", "a", 20*time.Millisecond); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(40 * time.Millisecond)
|
||||
lease, err := s.AcquireLeadership(ctx, "control", "b", time.Minute)
|
||||
if err != nil {
|
||||
t.Fatalf("a lapsed lease was not takeable: %v", err)
|
||||
}
|
||||
if lease.Owner != "b" {
|
||||
t.Errorf("owner after takeover = %q", lease.Owner)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("releasing hands over at once", func(t *testing.T) {
|
||||
s := open(t)
|
||||
if _, err := s.AcquireLeadership(ctx, "control", "a", time.Hour); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.ReleaseLeadership(ctx, "control", "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.AcquireLeadership(ctx, "control", "b", time.Minute); err != nil {
|
||||
t.Errorf("a released lease was not free: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("a former leader cannot release the current one", func(t *testing.T) {
|
||||
// A process that lost the lease and is shutting down slowly must not
|
||||
// take it from whoever has it now — which is exactly what an
|
||||
// unconditional delete would do.
|
||||
s := open(t)
|
||||
if _, err := s.AcquireLeadership(ctx, "control", "a", 20*time.Millisecond); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(40 * time.Millisecond)
|
||||
if _, err := s.AcquireLeadership(ctx, "control", "b", time.Hour); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.ReleaseLeadership(ctx, "control", "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
held, err := s.Leadership(ctx, "control")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if held.Owner != "b" {
|
||||
t.Errorf("the old leader released the new one's lease: %+v", held)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("only one of many wins", func(t *testing.T) {
|
||||
// The case the whole table exists for. Two control planes starting
|
||||
// together is the ordinary way this gets tested in production.
|
||||
s := open(t)
|
||||
const racers = 8
|
||||
var wg sync.WaitGroup
|
||||
won := make([]bool, racers)
|
||||
for i := range racers {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_, err := s.AcquireLeadership(ctx, "control", string(rune('a'+i)), time.Minute)
|
||||
won[i] = err == nil
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
n := 0
|
||||
for _, w := range won {
|
||||
if w {
|
||||
n++
|
||||
}
|
||||
}
|
||||
if n != 1 {
|
||||
t.Errorf("%d of %d processes believed they were leader", n, racers)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestLeadershipOnSQLite(t *testing.T) {
|
||||
leadershipSuite(t, func(t *testing.T) *SQLite { return openDB(t) })
|
||||
}
|
||||
|
||||
func TestLeadershipOnPostgres(t *testing.T) {
|
||||
leadershipSuite(t, openPG)
|
||||
}
|
||||
@@ -0,0 +1,296 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ErrMemoryNotEnabled means the project has not asked llmgw to keep notes.
|
||||
//
|
||||
// Returned rather than ignoring the write: an agent that believes it saved
|
||||
// something and finds nothing later has been lied to.
|
||||
var ErrMemoryNotEnabled = errors.New("store: this project does not keep memory notes")
|
||||
|
||||
// ErrNoNote means the key does not exist, or belongs to another project or agent.
|
||||
var ErrNoNote = errors.New("store: no such note")
|
||||
|
||||
// ErrMemoryFull means the agent already holds as many notes as one may.
|
||||
var ErrMemoryFull = errors.New("store: this agent has reached its limit of notes")
|
||||
|
||||
// Limits on what one note and one agent may hold. Notes are memory held on
|
||||
// behalf of a caller, so each is bounded rather than trusted.
|
||||
const (
|
||||
MaxNoteKey = 200
|
||||
MaxNoteText = 16 << 10
|
||||
MaxNoteTags = 16
|
||||
MaxNotesPerAgent = 1000
|
||||
MaxNoteResults = 100
|
||||
)
|
||||
|
||||
var tagPattern = regexp.MustCompile(`^[a-z0-9][a-z0-9_-]{0,63}$`)
|
||||
|
||||
// Note is one thing an agent chose to remember.
|
||||
type Note struct {
|
||||
Key string
|
||||
Text string
|
||||
Tags []string
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
// NormalizeTag lowercases and validates a tag, so "Work" and "work" are one tag
|
||||
// and nothing outside a small alphabet can reach the stored form.
|
||||
func NormalizeTag(tag string) (string, error) {
|
||||
t := strings.ToLower(strings.TrimSpace(tag))
|
||||
if !tagPattern.MatchString(t) {
|
||||
return "", fmt.Errorf("invalid tag %q: use lowercase letters, digits, - and _ (up to 64)", tag)
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
|
||||
func encodeTags(tags []string) (string, error) {
|
||||
seen := map[string]bool{}
|
||||
var out []string
|
||||
for _, tag := range tags {
|
||||
t, err := NormalizeTag(tag)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !seen[t] {
|
||||
seen[t] = true
|
||||
out = append(out, t)
|
||||
}
|
||||
}
|
||||
if len(out) > MaxNoteTags {
|
||||
return "", fmt.Errorf("too many tags: at most %d", MaxNoteTags)
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
return "," + strings.Join(out, ",") + ",", nil
|
||||
}
|
||||
|
||||
func decodeTags(s string) []string {
|
||||
s = strings.Trim(s, ",")
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
return strings.Split(s, ",")
|
||||
}
|
||||
|
||||
// MemoryRetention reports how long a project keeps notes. Zero means none.
|
||||
func (s *SQLite) MemoryRetention(ctx context.Context, projectID string) (time.Duration, error) {
|
||||
var hours sql.NullInt64
|
||||
err := s.db.QueryRowContext(ctx,
|
||||
`SELECT memory_retention_hours FROM projects WHERE id = ?`, projectID).Scan(&hours)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return 0, nil
|
||||
}
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !hours.Valid || hours.Int64 <= 0 {
|
||||
return 0, nil
|
||||
}
|
||||
return time.Duration(hours.Int64) * time.Hour, nil
|
||||
}
|
||||
|
||||
// SetMemoryRetention turns note storage on for a project, or off.
|
||||
func (s *SQLite) SetMemoryRetention(ctx context.Context, projectID string, hours int) error {
|
||||
var v any
|
||||
if hours > 0 {
|
||||
v = hours
|
||||
}
|
||||
res, err := s.db.ExecContext(ctx,
|
||||
`UPDATE projects SET memory_retention_hours = ? WHERE id = ?`, v, projectID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return fmt.Errorf("store: no project %q", projectID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// PutNote writes a note, replacing any note the agent already holds under the
|
||||
// same key.
|
||||
func (s *SQLite) PutNote(ctx context.Context, projectID, agentID string, n Note) error {
|
||||
if n.Key == "" || len(n.Key) > MaxNoteKey {
|
||||
return fmt.Errorf("key must be 1 to %d bytes", MaxNoteKey)
|
||||
}
|
||||
if n.Text == "" || len(n.Text) > MaxNoteText {
|
||||
return fmt.Errorf("text must be 1 to %d bytes", MaxNoteText)
|
||||
}
|
||||
tags, err := encodeTags(n.Tags)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
retention, err := s.MemoryRetention(ctx, projectID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if retention == 0 {
|
||||
return ErrMemoryNotEnabled
|
||||
}
|
||||
|
||||
q, tx, err := s.begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
now := time.Now().UnixMilli()
|
||||
res, err := q.ExecContext(ctx,
|
||||
`UPDATE memory_notes SET text = ?, tags = ?, updated_at = ?
|
||||
WHERE project_id = ? AND agent_id = ? AND key = ?`,
|
||||
n.Text, tags, now, projectID, agentID, n.Key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if updated, _ := res.RowsAffected(); updated == 0 {
|
||||
var count int
|
||||
if err := q.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM memory_notes WHERE project_id = ? AND agent_id = ?`,
|
||||
projectID, agentID).Scan(&count); err != nil {
|
||||
return err
|
||||
}
|
||||
if count >= MaxNotesPerAgent {
|
||||
return ErrMemoryFull
|
||||
}
|
||||
if _, err := q.ExecContext(ctx,
|
||||
`INSERT INTO memory_notes (project_id, agent_id, key, text, tags, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)`,
|
||||
projectID, agentID, n.Key, n.Text, tags, now, now); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func scanNote(rows interface{ Scan(...any) error }) (Note, error) {
|
||||
var (
|
||||
n Note
|
||||
tags string
|
||||
created int64
|
||||
updatedRaw int64
|
||||
)
|
||||
if err := rows.Scan(&n.Key, &n.Text, &tags, &created, &updatedRaw); err != nil {
|
||||
return Note{}, err
|
||||
}
|
||||
n.Tags = decodeTags(tags)
|
||||
n.CreatedAt = time.UnixMilli(created)
|
||||
n.UpdatedAt = time.UnixMilli(updatedRaw)
|
||||
return n, nil
|
||||
}
|
||||
|
||||
const noteColumns = `key, text, tags, created_at, updated_at`
|
||||
|
||||
// GetNote reads one note. Scoped to the project and agent in the query itself.
|
||||
func (s *SQLite) GetNote(ctx context.Context, projectID, agentID, key string) (Note, error) {
|
||||
n, err := scanNote(s.db.QueryRowContext(ctx,
|
||||
`SELECT `+noteColumns+` FROM memory_notes WHERE project_id = ? AND agent_id = ? AND key = ?`,
|
||||
projectID, agentID, key))
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return Note{}, ErrNoNote
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
// ftsQuery turns free text into a query that cannot be a syntax error: every
|
||||
// word becomes a quoted term, and they must all match. FTS5 treats characters
|
||||
// such as - : * and " as operators, and an agent's search words are not written
|
||||
// for that syntax.
|
||||
func ftsQuery(q string) string {
|
||||
words := strings.Fields(q)
|
||||
for i, w := range words {
|
||||
words[i] = `"` + strings.ReplaceAll(w, `"`, `""`) + `"`
|
||||
}
|
||||
return strings.Join(words, " ")
|
||||
}
|
||||
|
||||
// SearchNotes finds an agent's notes by words and/or tag, newest first. With
|
||||
// neither it lists the most recent.
|
||||
func (s *SQLite) SearchNotes(ctx context.Context, projectID, agentID, query, tag string, limit int) ([]Note, error) {
|
||||
if limit <= 0 || limit > MaxNoteResults {
|
||||
limit = MaxNoteResults
|
||||
}
|
||||
q := `SELECT ` + noteColumns + ` FROM memory_notes WHERE project_id = ? AND agent_id = ?`
|
||||
args := []any{projectID, agentID}
|
||||
|
||||
if match := ftsQuery(query); match != "" {
|
||||
// The one query in this store that is not the same sentence in both
|
||||
// dialects. SQLite searches an FTS5 table kept in step by triggers;
|
||||
// PostgreSQL matches a tsvector derived from the row. Writing it out
|
||||
// twice is honest about that — the alternative was a shared spelling
|
||||
// that means something subtly different in each.
|
||||
if s.dialect == PostgresDialect {
|
||||
q += ` AND search @@ plainto_tsquery('simple', ?)`
|
||||
// plainto_tsquery takes words, not FTS5 syntax, so the raw query
|
||||
// goes in rather than the quoted form SQLite needs.
|
||||
args = append(args, query)
|
||||
} else {
|
||||
q += ` AND id IN (SELECT rowid FROM memory_notes_fts WHERE memory_notes_fts MATCH ?)`
|
||||
args = append(args, match)
|
||||
}
|
||||
}
|
||||
if tag != "" {
|
||||
t, err := NormalizeTag(tag)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
q += ` AND tags LIKE ?`
|
||||
args = append(args, "%,"+t+",%")
|
||||
}
|
||||
q += ` ORDER BY updated_at DESC, id DESC LIMIT ?`
|
||||
args = append(args, limit)
|
||||
|
||||
return queryAll(ctx, s.db, q, args, func(rows *sql.Rows) (Note, error) { return scanNote(rows) })
|
||||
}
|
||||
|
||||
// DeleteNote forgets one note.
|
||||
func (s *SQLite) DeleteNote(ctx context.Context, projectID, agentID, key string) error {
|
||||
res, err := s.db.ExecContext(ctx,
|
||||
`DELETE FROM memory_notes WHERE project_id = ? AND agent_id = ? AND key = ?`,
|
||||
projectID, agentID, key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return ErrNoNote
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// PruneMemory deletes notes not refreshed within their project's retention.
|
||||
func (s *SQLite) PruneMemory(ctx context.Context, now time.Time) (int, error) {
|
||||
type window struct {
|
||||
project string
|
||||
hours int64
|
||||
}
|
||||
windows, err := queryAll(ctx, s.db,
|
||||
`SELECT id, memory_retention_hours FROM projects WHERE memory_retention_hours > 0`,
|
||||
nil, func(rows *sql.Rows) (window, error) {
|
||||
var w window
|
||||
return w, rows.Scan(&w.project, &w.hours)
|
||||
})
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
var removed int
|
||||
for _, w := range windows {
|
||||
cutoff := now.Add(-time.Duration(w.hours) * time.Hour).UnixMilli()
|
||||
res, err := s.db.ExecContext(ctx,
|
||||
`DELETE FROM memory_notes WHERE project_id = ? AND updated_at < ?`, w.project, cutoff)
|
||||
if err != nil {
|
||||
return removed, err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
removed += int(n)
|
||||
}
|
||||
return removed, nil
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func memoryProject(t *testing.T, s *SQLite) string {
|
||||
t.Helper()
|
||||
proj := tree(t, s)
|
||||
if err := s.SetMemoryRetention(context.Background(), proj, 24); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return proj
|
||||
}
|
||||
|
||||
func TestAProjectKeepsNoNotesUnlessItAsks(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
proj := tree(t, s)
|
||||
|
||||
err := s.PutNote(ctx, proj, "a", Note{Key: "k", Text: "secret"})
|
||||
if !errors.Is(err, ErrMemoryNotEnabled) {
|
||||
t.Fatalf("writing into a project that never asked returned %v", err)
|
||||
}
|
||||
if _, err := s.GetNote(ctx, proj, "a", "k"); !errors.Is(err, ErrNoNote) {
|
||||
t.Errorf("a refused write left a note behind: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestANoteIsWrittenReadBackAndReplaced(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
proj := memoryProject(t, s)
|
||||
|
||||
if err := s.PutNote(ctx, proj, "a", Note{Key: "city", Text: "Ankara", Tags: []string{"Geo", "geo", "facts"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := s.GetNote(ctx, proj, "a", "city")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.Text != "Ankara" || len(got.Tags) != 2 || got.Tags[0] != "geo" || got.Tags[1] != "facts" {
|
||||
t.Fatalf("note came back as %+v: tags should be lowercased and deduplicated", got)
|
||||
}
|
||||
|
||||
if err := s.PutNote(ctx, proj, "a", Note{Key: "city", Text: "Istanbul"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, _ = s.GetNote(ctx, proj, "a", "city")
|
||||
if got.Text != "Istanbul" || len(got.Tags) != 0 {
|
||||
t.Errorf("rewriting a key should replace text and tags: %+v", got)
|
||||
}
|
||||
notes, _ := s.SearchNotes(ctx, proj, "a", "", "", 0)
|
||||
if len(notes) != 1 {
|
||||
t.Errorf("rewriting a key created %d notes, want 1", len(notes))
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotesAreFoundByWordsAndTagsAndSpecialCharactersAreHarmless(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
proj := memoryProject(t, s)
|
||||
put := func(key, text string, tags ...string) {
|
||||
t.Helper()
|
||||
if err := s.PutNote(ctx, proj, "a", Note{Key: key, Text: text, Tags: tags}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
put("deploy", "the staging cluster is on port 8443", "ops")
|
||||
put("food", "she prefers green tea", "personal")
|
||||
put("old", "deploy notes from last year", "ops", "archive")
|
||||
|
||||
byWord, err := s.SearchNotes(ctx, proj, "a", "staging cluster", "", 0)
|
||||
if err != nil || len(byWord) != 1 || byWord[0].Key != "deploy" {
|
||||
t.Fatalf("search by words: %+v (err %v)", byWord, err)
|
||||
}
|
||||
byTag, _ := s.SearchNotes(ctx, proj, "a", "", "ops", 0)
|
||||
if len(byTag) != 2 {
|
||||
t.Errorf("tag ops matched %d notes, want 2", len(byTag))
|
||||
}
|
||||
both, _ := s.SearchNotes(ctx, proj, "a", "deploy", "archive", 0)
|
||||
if len(both) != 1 || both[0].Key != "old" {
|
||||
t.Errorf("words and tag together: %+v", both)
|
||||
}
|
||||
// A tag is matched whole: "op" must not match "ops".
|
||||
if part, _ := s.SearchNotes(ctx, proj, "a", "", "op", 0); len(part) != 0 {
|
||||
t.Errorf("a partial tag matched %d notes", len(part))
|
||||
}
|
||||
// FTS operators in an agent's words must not become a syntax error.
|
||||
for _, q := range []string{`"`, `a -b`, `x:y`, `foo*`, `AND OR NOT`, `(`} {
|
||||
if _, err := s.SearchNotes(ctx, proj, "a", q, "", 0); err != nil {
|
||||
t.Errorf("query %q failed: %v", q, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOneAgentOrProjectCannotReadAnothersNotes(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
proj := memoryProject(t, s)
|
||||
if err := s.PutOrg(ctx, Org{ID: "o2", Name: "Other"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.PutTeam(ctx, Team{ID: "t2", OrgID: "o2", Name: "T"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.PutProject(ctx, Project{ID: "proj2", TeamID: "t2", Name: "P"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.SetMemoryRetention(ctx, "proj2", 24); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if err := s.PutNote(ctx, proj, "agent1", Note{Key: "k", Text: "private"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.GetNote(ctx, proj, "agent2", "k"); !errors.Is(err, ErrNoNote) {
|
||||
t.Errorf("another agent in the same project read the note: %v", err)
|
||||
}
|
||||
if _, err := s.GetNote(ctx, "proj2", "agent1", "k"); !errors.Is(err, ErrNoNote) {
|
||||
t.Errorf("another project read the note: %v", err)
|
||||
}
|
||||
if got, _ := s.SearchNotes(ctx, "proj2", "agent1", "private", "", 0); len(got) != 0 {
|
||||
t.Errorf("another project found the note by search: %+v", got)
|
||||
}
|
||||
if err := s.DeleteNote(ctx, "proj2", "agent1", "k"); !errors.Is(err, ErrNoNote) {
|
||||
t.Errorf("another project deleted the note: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotesAreBoundedAndBadInputIsRefused(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
proj := memoryProject(t, s)
|
||||
|
||||
bad := []Note{
|
||||
{Key: "", Text: "x"},
|
||||
{Key: "k", Text: ""},
|
||||
{Key: strings.Repeat("k", MaxNoteKey+1), Text: "x"},
|
||||
{Key: "k", Text: strings.Repeat("x", MaxNoteText+1)},
|
||||
{Key: "k", Text: "x", Tags: []string{"has space"}},
|
||||
}
|
||||
for i, n := range bad {
|
||||
if err := s.PutNote(ctx, proj, "a", n); err == nil {
|
||||
t.Errorf("bad note %d was accepted", i)
|
||||
}
|
||||
}
|
||||
|
||||
for i := 0; i < MaxNotesPerAgent; i++ {
|
||||
if err := s.PutNote(ctx, proj, "a", Note{Key: fmt.Sprintf("note-%d", i), Text: "x"}); err != nil {
|
||||
t.Fatalf("note %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
if err := s.PutNote(ctx, proj, "a", Note{Key: "one-too-many", Text: "x"}); !errors.Is(err, ErrMemoryFull) {
|
||||
t.Errorf("exceeding the per-agent limit returned %v", err)
|
||||
}
|
||||
// Replacing an existing note is not growth and must still work when full.
|
||||
if err := s.PutNote(ctx, proj, "a", Note{Key: "note-0", Text: "y"}); err != nil {
|
||||
t.Errorf("replacing a note in a full agent failed: %v", err)
|
||||
}
|
||||
// Another agent has its own allowance.
|
||||
if err := s.PutNote(ctx, proj, "b", Note{Key: "note-0", Text: "x"}); err != nil {
|
||||
t.Errorf("a second agent was blocked by the first one's notes: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAForgottenNoteIsGoneFromSearchToo(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
proj := memoryProject(t, s)
|
||||
|
||||
if err := s.PutNote(ctx, proj, "a", Note{Key: "k", Text: "needle in a haystack"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.DeleteNote(ctx, proj, "a", "k"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, _ := s.SearchNotes(ctx, proj, "a", "needle", "", 0); len(got) != 0 {
|
||||
t.Errorf("a deleted note is still searchable: %+v", got)
|
||||
}
|
||||
if err := s.DeleteNote(ctx, proj, "a", "k"); !errors.Is(err, ErrNoNote) {
|
||||
t.Errorf("deleting twice returned %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotesNotRefreshedWithinRetentionArePruned(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
proj := memoryProject(t, s)
|
||||
|
||||
if err := s.PutNote(ctx, proj, "a", Note{Key: "k", Text: "x"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n, err := s.PruneMemory(ctx, time.Now()); err != nil || n != 0 {
|
||||
t.Fatalf("a fresh note was pruned: %d, %v", n, err)
|
||||
}
|
||||
if n, err := s.PruneMemory(ctx, time.Now().Add(25*time.Hour)); err != nil || n != 1 {
|
||||
t.Fatalf("an expired note survived: %d, %v", n, err)
|
||||
}
|
||||
if got, _ := s.SearchNotes(ctx, proj, "a", "", "", 0); len(got) != 0 {
|
||||
t.Errorf("pruned note still listed: %+v", got)
|
||||
}
|
||||
}
|
||||
@@ -10,7 +10,7 @@ import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
//go:embed migrations/sqlite/*.sql
|
||||
//go:embed migrations/sqlite/*.sql migrations/postgres/*.sql
|
||||
var migrationFS embed.FS
|
||||
|
||||
// migrate applies every unapplied migration in order, each in its own
|
||||
@@ -25,8 +25,17 @@ var migrationFS embed.FS
|
||||
// promise this could not keep. Same version numbers, same meaning, different
|
||||
// files.
|
||||
func migrate(ctx context.Context, db *sql.DB, dialect string) error {
|
||||
// BIGINT, not INTEGER. SQLite's INTEGER holds 64 bits; PostgreSQL's is
|
||||
// int4, and a millisecond timestamp does not fit in it — which the first
|
||||
// migration discovers by failing to record that it ran. This table is
|
||||
// created here rather than in a migration file, which is exactly how it
|
||||
// escaped the port of the others.
|
||||
stamp := "BIGINT"
|
||||
if dialect == string(SQLiteDialect) {
|
||||
stamp = "INTEGER"
|
||||
}
|
||||
if _, err := db.ExecContext(ctx,
|
||||
`CREATE TABLE IF NOT EXISTS schema_migrations (version TEXT PRIMARY KEY, applied_at INTEGER NOT NULL)`,
|
||||
`CREATE TABLE IF NOT EXISTS schema_migrations (version TEXT PRIMARY KEY, applied_at `+stamp+` NOT NULL)`,
|
||||
); err != nil {
|
||||
return fmt.Errorf("creating the migration table: %w", err)
|
||||
}
|
||||
@@ -63,7 +72,7 @@ func migrate(ctx context.Context, db *sql.DB, dialect string) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := applyOne(ctx, db, name, string(body)); err != nil {
|
||||
if err := applyOne(ctx, db, dialect, name, string(body)); err != nil {
|
||||
return fmt.Errorf("migration %s: %w", name, err)
|
||||
}
|
||||
}
|
||||
@@ -86,7 +95,7 @@ func migrationNames(dialect string) ([]string, error) {
|
||||
return names, nil
|
||||
}
|
||||
|
||||
func applyOne(ctx context.Context, db *sql.DB, name, body string) error {
|
||||
func applyOne(ctx context.Context, db *sql.DB, dialect, name, body string) error {
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -97,7 +106,7 @@ func applyOne(ctx context.Context, db *sql.DB, name, body string) error {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
`INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)`,
|
||||
Dialect(dialect).rebind(`INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)`),
|
||||
name, nowMillis(),
|
||||
); err != nil {
|
||||
return err
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
-- Ported from the SQLite migration of the same number. Same version, same
|
||||
-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types
|
||||
-- do not line up, and one shared file would be a promise neither could keep.
|
||||
--
|
||||
-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite,
|
||||
-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds
|
||||
-- everywhere; converting here would mean two time representations in one
|
||||
-- codebase, and the conversion would be the thing that is wrong at 2am.
|
||||
|
||||
-- The tenancy tree. Three levels, each owning the next.
|
||||
CREATE TABLE organizations (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE teams (
|
||||
id TEXT PRIMARY KEY,
|
||||
org_id TEXT NOT NULL REFERENCES organizations(id),
|
||||
name TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX teams_by_org ON teams(org_id);
|
||||
|
||||
CREATE TABLE projects (
|
||||
id TEXT PRIMARY KEY,
|
||||
team_id TEXT NOT NULL REFERENCES teams(id),
|
||||
name TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX projects_by_team ON projects(team_id);
|
||||
|
||||
-- Only sha256(key) is stored, so a database leak hands over no working
|
||||
-- credential. SHA-256 rather than a slow password hash on purpose: the input is
|
||||
-- 256 random bits with no dictionary to attack, so the cost of argon2 would buy
|
||||
-- nothing and charge tens of milliseconds on every request (plan §7.9.1).
|
||||
CREATE TABLE api_keys (
|
||||
id TEXT PRIMARY KEY,
|
||||
key_hash BYTEA NOT NULL UNIQUE,
|
||||
key_prefix TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
project_id TEXT NOT NULL REFERENCES projects(id),
|
||||
band TEXT,
|
||||
routing_profile TEXT,
|
||||
created_at BIGINT NOT NULL,
|
||||
expires_at BIGINT,
|
||||
-- Revocation is a soft delete. Usage rows reference the key, and an audit
|
||||
-- trail that loses its subject is not an audit trail.
|
||||
revoked_at BIGINT,
|
||||
last_used_at BIGINT
|
||||
);
|
||||
CREATE INDEX api_keys_by_project ON api_keys(project_id);
|
||||
|
||||
-- Limits hang off any scope rather than being duplicated onto each one. A NULL
|
||||
-- column means "no limit at this level", which is not the same as 0.
|
||||
CREATE TABLE limits (
|
||||
scope_type TEXT NOT NULL,
|
||||
scope_id TEXT NOT NULL,
|
||||
rpm BIGINT,
|
||||
tpm BIGINT,
|
||||
-- Micro-units, not a float: a budget is accumulated by repeated addition, and
|
||||
-- repeated addition of floats drifts.
|
||||
budget_micros BIGINT,
|
||||
budget_period TEXT,
|
||||
allowed_models TEXT,
|
||||
PRIMARY KEY (scope_type, scope_id)
|
||||
);
|
||||
|
||||
-- One row per request, written in batches off the request path.
|
||||
CREATE TABLE usage (
|
||||
ts BIGINT NOT NULL,
|
||||
key_id TEXT NOT NULL,
|
||||
project_id TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
served_model TEXT NOT NULL,
|
||||
prompt_tokens BIGINT NOT NULL,
|
||||
completion_tokens BIGINT NOT NULL,
|
||||
cost_micros BIGINT NOT NULL,
|
||||
status BIGINT NOT NULL,
|
||||
latency_ms BIGINT NOT NULL
|
||||
);
|
||||
-- Budget restore reads by key within a period, which is this index exactly.
|
||||
CREATE INDEX usage_by_key_ts ON usage(key_id, ts);
|
||||
CREATE INDEX usage_by_ts ON usage(ts);
|
||||
@@ -0,0 +1,29 @@
|
||||
-- Ported from the SQLite migration of the same number. Same version, same
|
||||
-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types
|
||||
-- do not line up, and one shared file would be a promise neither could keep.
|
||||
--
|
||||
-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite,
|
||||
-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds
|
||||
-- everywhere; converting here would mean two time representations in one
|
||||
-- codebase, and the conversion would be the thing that is wrong at 2am.
|
||||
|
||||
-- A key has to reach the person it was issued for, and a key pasted into a chat
|
||||
-- window stays in that chat window forever. A claim is a one-time link instead:
|
||||
-- the operator sends the link, the person opens it once, and the link dies.
|
||||
--
|
||||
-- The key is not stored here in the clear. It is sealed with a key derived from
|
||||
-- the claim token, and the token is never written down — only its hash is. So a
|
||||
-- database leak yields two useless hashes and a blob nothing in the database can
|
||||
-- open, which keeps the property the api_keys table has.
|
||||
CREATE TABLE key_claims (
|
||||
token_hash BYTEA PRIMARY KEY,
|
||||
key_id TEXT NOT NULL REFERENCES api_keys(id),
|
||||
-- The sealed key. Emptied the moment it is claimed, so a claimed link leaves
|
||||
-- nothing behind even before the row is purged.
|
||||
sealed BYTEA NOT NULL,
|
||||
created_at BIGINT NOT NULL,
|
||||
expires_at BIGINT NOT NULL,
|
||||
claimed_at BIGINT
|
||||
);
|
||||
CREATE INDEX key_claims_by_key ON key_claims(key_id);
|
||||
CREATE INDEX key_claims_by_expiry ON key_claims(expires_at);
|
||||
@@ -0,0 +1,56 @@
|
||||
-- Ported from the SQLite migration of the same number. Same version, same
|
||||
-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types
|
||||
-- do not line up, and one shared file would be a promise neither could keep.
|
||||
--
|
||||
-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite,
|
||||
-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds
|
||||
-- everywhere; converting here would mean two time representations in one
|
||||
-- codebase, and the conversion would be the thing that is wrong at 2am.
|
||||
|
||||
-- Accounts, so a person can sign up, be approved, and then issue their own keys
|
||||
-- without an operator running a command for each one.
|
||||
CREATE TABLE users (
|
||||
id TEXT PRIMARY KEY,
|
||||
-- Lowercased at the boundary. Two accounts differing only in case would be
|
||||
-- two accounts to the database and one to every human who looks at them.
|
||||
email TEXT NOT NULL UNIQUE,
|
||||
name TEXT NOT NULL,
|
||||
-- argon2id, encoded with its own parameters so the cost can be raised later
|
||||
-- without invalidating existing hashes. Slow hashing is right here and wrong
|
||||
-- for API keys: this input is chosen by a human and has a dictionary.
|
||||
password_hash TEXT NOT NULL,
|
||||
-- 'pending' until an operator approves, then 'active'. 'suspended' keeps the
|
||||
-- account and stops it working.
|
||||
state TEXT NOT NULL,
|
||||
-- Assigned at approval: which project this person's keys belong to, and
|
||||
-- therefore which limits they inherit.
|
||||
project_id TEXT REFERENCES projects(id),
|
||||
created_at BIGINT NOT NULL,
|
||||
approved_at BIGINT,
|
||||
approved_by TEXT,
|
||||
last_login_at BIGINT
|
||||
);
|
||||
CREATE INDEX users_by_state ON users(state);
|
||||
|
||||
-- Server-side sessions rather than signed cookies, so signing out and
|
||||
-- suspending an account take effect immediately instead of whenever a token
|
||||
-- happens to expire.
|
||||
CREATE TABLE sessions (
|
||||
token_hash BYTEA PRIMARY KEY,
|
||||
user_id TEXT NOT NULL REFERENCES users(id),
|
||||
-- Per session, submitted with every form. A cookie alone authenticates the
|
||||
-- browser, not the page that asked, which is what makes a cross-site POST
|
||||
-- work.
|
||||
csrf_hash BYTEA NOT NULL,
|
||||
created_at BIGINT NOT NULL,
|
||||
expires_at BIGINT NOT NULL,
|
||||
last_seen BIGINT NOT NULL
|
||||
);
|
||||
CREATE INDEX sessions_by_user ON sessions(user_id);
|
||||
CREATE INDEX sessions_by_expiry ON sessions(expires_at);
|
||||
|
||||
-- A key issued through the profile page belongs to the person who made it, so
|
||||
-- they can see and revoke their own and nobody else's. NULL means a key an
|
||||
-- operator issued directly.
|
||||
ALTER TABLE api_keys ADD COLUMN user_id TEXT REFERENCES users(id);
|
||||
CREATE INDEX api_keys_by_user ON api_keys(user_id);
|
||||
@@ -0,0 +1,20 @@
|
||||
-- Ported from the SQLite migration of the same number. Same version, same
|
||||
-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types
|
||||
-- do not line up, and one shared file would be a promise neither could keep.
|
||||
--
|
||||
-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite,
|
||||
-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds
|
||||
-- everywhere; converting here would mean two time representations in one
|
||||
-- codebase, and the conversion would be the thing that is wrong at 2am.
|
||||
|
||||
-- Who may see the cluster.
|
||||
--
|
||||
-- Until now the answer was "whoever has a shell on the control plane", because
|
||||
-- the dashboard lived on a loopback-bound port. Serving it over the web needs
|
||||
-- an actual answer, and this is the smallest one that is not a lie: a flag on
|
||||
-- an account that already has a password, a session and a CSRF token.
|
||||
--
|
||||
-- Deliberately not a role system. One bit is enough for "can see the cluster"
|
||||
-- and impossible to get subtly wrong; roles arrive when there is a second
|
||||
-- distinction worth drawing.
|
||||
ALTER TABLE users ADD COLUMN is_admin BIGINT NOT NULL DEFAULT 0;
|
||||
@@ -0,0 +1,34 @@
|
||||
-- Ported from the SQLite migration of the same number. Same version, same
|
||||
-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types
|
||||
-- do not line up, and one shared file would be a promise neither could keep.
|
||||
--
|
||||
-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite,
|
||||
-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds
|
||||
-- everywhere; converting here would mean two time representations in one
|
||||
-- codebase, and the conversion would be the thing that is wrong at 2am.
|
||||
|
||||
-- Models added while the cluster is running.
|
||||
--
|
||||
-- The catalog came from a YAML file read once at startup, so adding a model
|
||||
-- meant editing a file on the control plane and restarting it: a settling
|
||||
-- window, adoption churn, and on one occasion two minutes of outage. That is a
|
||||
-- reasonable way to manage a fixed production set and a bad one for a platform
|
||||
-- whose promise is that people can try things.
|
||||
--
|
||||
-- These live alongside the ones in configuration rather than replacing them.
|
||||
-- Config-as-code stays declarative — a model in YAML is in the catalog on every
|
||||
-- start, whatever the database says — and this table holds the ones added
|
||||
-- afterwards. A model added here may not take a name or alias that
|
||||
-- configuration already uses, which the catalog's own constructor enforces for
|
||||
-- free, because it refuses a duplicate whatever the source.
|
||||
--
|
||||
-- The definition is stored as JSON rather than as columns. A profile is a
|
||||
-- nested, evolving shape, and every field added to it would otherwise be a
|
||||
-- migration; the catalog already knows how to validate one, and a row that
|
||||
-- fails validation is reported and skipped rather than blocking startup.
|
||||
CREATE TABLE models (
|
||||
name TEXT PRIMARY KEY,
|
||||
definition TEXT NOT NULL,
|
||||
created_at TIMESTAMP NOT NULL,
|
||||
created_by TEXT NOT NULL DEFAULT ''
|
||||
);
|
||||
@@ -0,0 +1,57 @@
|
||||
-- Ported from the SQLite migration of the same number. Same version, same
|
||||
-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types
|
||||
-- do not line up, and one shared file would be a promise neither could keep.
|
||||
--
|
||||
-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite,
|
||||
-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds
|
||||
-- everywhere; converting here would mean two time representations in one
|
||||
-- codebase, and the conversion would be the thing that is wrong at 2am.
|
||||
|
||||
-- Server-side conversations, so a thread can continue on a different model
|
||||
-- than it started on.
|
||||
--
|
||||
-- Off unless a project turns it on. Until this migration ran, llmgw stored no
|
||||
-- message content anywhere — usage rows carry token counts and costs, never
|
||||
-- text — and that stays true for every project that does not ask. An on-prem
|
||||
-- gateway keeping nothing anyone typed is much easier to put in front of a
|
||||
-- team than one keeping everything, so storing is a choice a project makes
|
||||
-- rather than a default it discovers.
|
||||
|
||||
-- retention_hours is the switch and the limit at once. NULL means this project
|
||||
-- stores nothing, which is every project that existed before this and every
|
||||
-- one created since without asking.
|
||||
--
|
||||
-- One column rather than a boolean and a duration: a thread store with no
|
||||
-- expiry is a transcript of everything a team ever typed, kept forever by
|
||||
-- omission. Turning it on therefore requires saying for how long.
|
||||
ALTER TABLE projects ADD COLUMN thread_retention_hours BIGINT;
|
||||
|
||||
CREATE TABLE threads (
|
||||
id TEXT PRIMARY KEY,
|
||||
-- Scoped to a project, and checked on every read. Thread ids are opaque, but
|
||||
-- an unguessable id is a secret and secrets leak into logs, shell history
|
||||
-- and bug reports; the scope is what makes a leaked id useless to someone
|
||||
-- outside the project.
|
||||
project_id TEXT NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
||||
created_at BIGINT NOT NULL,
|
||||
-- Touched on every append, so retention measures idleness rather than age:
|
||||
-- expiring a conversation someone is still having would be a strange way to
|
||||
-- keep a promise about storage.
|
||||
updated_at BIGINT NOT NULL
|
||||
);
|
||||
CREATE INDEX threads_by_project ON threads(project_id);
|
||||
CREATE INDEX threads_by_updated ON threads(updated_at);
|
||||
|
||||
CREATE TABLE thread_messages (
|
||||
thread_id TEXT NOT NULL REFERENCES threads(id) ON DELETE CASCADE,
|
||||
-- Position rather than a timestamp, because order is the point and two
|
||||
-- messages can share a millisecond.
|
||||
seq BIGINT NOT NULL,
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
-- Which model produced an assistant turn. A thread that moved between models
|
||||
-- otherwise reads as one voice, and when the answers disagree that is the
|
||||
-- first thing anyone wants to know.
|
||||
model TEXT,
|
||||
PRIMARY KEY (thread_id, seq)
|
||||
);
|
||||
@@ -0,0 +1,41 @@
|
||||
-- Ported from the SQLite migration of the same number, and the only one that
|
||||
-- is not a mechanical translation.
|
||||
--
|
||||
-- SQLite gets full-text search from an FTS5 virtual table with an external
|
||||
-- content table and three triggers keeping it in step. PostgreSQL derives a
|
||||
-- tsvector from the row itself, so there is nothing to keep in step and no
|
||||
-- trigger to forget: the column cannot disagree with the row it is computed
|
||||
-- from.
|
||||
--
|
||||
-- The search query differs too, and that difference lives in the Go code
|
||||
-- rather than being hidden here — a query that means "match these words" is
|
||||
-- spelled MATCH in one and @@ in the other, and pretending otherwise would
|
||||
-- cost more than admitting it.
|
||||
|
||||
ALTER TABLE projects ADD COLUMN memory_retention_hours BIGINT;
|
||||
|
||||
CREATE TABLE memory_notes (
|
||||
id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
|
||||
-- Scoped to a project and checked in SQL on every read, like threads.
|
||||
project_id TEXT NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
||||
-- A namespace inside the project, so two agents sharing a key do not read
|
||||
-- each other's notes.
|
||||
agent_id TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
-- Normalised to ",a,b," so a tag matches whole, never as part of another.
|
||||
tags TEXT NOT NULL DEFAULT '',
|
||||
created_at BIGINT NOT NULL,
|
||||
-- Touched on every write, so retention measures how long a note has gone
|
||||
-- unrefreshed rather than how old it is.
|
||||
updated_at BIGINT NOT NULL,
|
||||
-- 'simple' rather than a language configuration: these are an agent's own
|
||||
-- notes, in whatever language it chose, and stemming English rules over
|
||||
-- Turkish text would quietly stop finding things.
|
||||
search tsvector GENERATED ALWAYS AS (
|
||||
to_tsvector('simple', coalesce(key, '') || ' ' || coalesce(text, '') || ' ' || coalesce(tags, ''))
|
||||
) STORED,
|
||||
UNIQUE (project_id, agent_id, key)
|
||||
);
|
||||
CREATE INDEX memory_notes_by_updated ON memory_notes(updated_at);
|
||||
CREATE INDEX memory_notes_search ON memory_notes USING GIN (search);
|
||||
@@ -0,0 +1,29 @@
|
||||
-- Ported from the SQLite migration of the same number. BIGINT because
|
||||
-- PostgreSQL's INTEGER is 32 bits and these are milliseconds since the epoch.
|
||||
|
||||
-- A lease on being the control plane that places work.
|
||||
--
|
||||
-- Two control planes against one store both reconcile: both see a model
|
||||
-- wanting two replicas and none running, and both place. The result is twice
|
||||
-- the replicas, each holding GPUs the other's inventory believes are free —
|
||||
-- the same failure a control plane restart used to produce before adoption,
|
||||
-- and harder to see because nothing crashed.
|
||||
--
|
||||
-- A lease rather than a lock held by the connection. PostgreSQL has advisory
|
||||
-- locks and SQLite has nothing like them, and this store keeps one set of
|
||||
-- queries for both; a row with an expiry works everywhere and survives the
|
||||
-- leader being killed rather than asked to stop.
|
||||
--
|
||||
-- One row, enforced by the primary key: there is one thing to be leader of.
|
||||
CREATE TABLE leadership (
|
||||
role TEXT PRIMARY KEY,
|
||||
-- Who holds it. An opaque id per process, so a process cannot be fooled by
|
||||
-- its own stale lease after a restart on the same host.
|
||||
owner TEXT NOT NULL,
|
||||
-- When the lease lapses. The holder renews well before this; anyone else
|
||||
-- may take it once the clock has passed, which is what makes a leader that
|
||||
-- was powered off recoverable without anyone being paged.
|
||||
expires_at BIGINT NOT NULL,
|
||||
-- For an operator asking who has been in charge and since when.
|
||||
acquired_at BIGINT NOT NULL
|
||||
);
|
||||
@@ -0,0 +1,48 @@
|
||||
-- Server-side conversations, so a thread can continue on a different model
|
||||
-- than it started on.
|
||||
--
|
||||
-- Off unless a project turns it on. Until this migration ran, llmgw stored no
|
||||
-- message content anywhere — usage rows carry token counts and costs, never
|
||||
-- text — and that stays true for every project that does not ask. An on-prem
|
||||
-- gateway keeping nothing anyone typed is much easier to put in front of a
|
||||
-- team than one keeping everything, so storing is a choice a project makes
|
||||
-- rather than a default it discovers.
|
||||
|
||||
-- retention_hours is the switch and the limit at once. NULL means this project
|
||||
-- stores nothing, which is every project that existed before this and every
|
||||
-- one created since without asking.
|
||||
--
|
||||
-- One column rather than a boolean and a duration: a thread store with no
|
||||
-- expiry is a transcript of everything a team ever typed, kept forever by
|
||||
-- omission. Turning it on therefore requires saying for how long.
|
||||
ALTER TABLE projects ADD COLUMN thread_retention_hours INTEGER;
|
||||
|
||||
CREATE TABLE threads (
|
||||
id TEXT PRIMARY KEY,
|
||||
-- Scoped to a project, and checked on every read. Thread ids are opaque, but
|
||||
-- an unguessable id is a secret and secrets leak into logs, shell history
|
||||
-- and bug reports; the scope is what makes a leaked id useless to someone
|
||||
-- outside the project.
|
||||
project_id TEXT NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
||||
created_at INTEGER NOT NULL,
|
||||
-- Touched on every append, so retention measures idleness rather than age:
|
||||
-- expiring a conversation someone is still having would be a strange way to
|
||||
-- keep a promise about storage.
|
||||
updated_at INTEGER NOT NULL
|
||||
);
|
||||
CREATE INDEX threads_by_project ON threads(project_id);
|
||||
CREATE INDEX threads_by_updated ON threads(updated_at);
|
||||
|
||||
CREATE TABLE thread_messages (
|
||||
thread_id TEXT NOT NULL REFERENCES threads(id) ON DELETE CASCADE,
|
||||
-- Position rather than a timestamp, because order is the point and two
|
||||
-- messages can share a millisecond.
|
||||
seq INTEGER NOT NULL,
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
-- Which model produced an assistant turn. A thread that moved between models
|
||||
-- otherwise reads as one voice, and when the answers disagree that is the
|
||||
-- first thing anyone wants to know.
|
||||
model TEXT,
|
||||
PRIMARY KEY (thread_id, seq)
|
||||
);
|
||||
@@ -0,0 +1,45 @@
|
||||
-- Notes an agent writes for itself and reads back later, across threads.
|
||||
--
|
||||
-- Off unless a project turns it on, for the same reason threads are: llmgw
|
||||
-- stores no message content unless somebody decided it should. NULL means this
|
||||
-- project keeps no notes. The number is hours, and is required to turn it on,
|
||||
-- because a store with no expiry is a transcript kept forever by omission.
|
||||
ALTER TABLE projects ADD COLUMN memory_retention_hours INTEGER;
|
||||
|
||||
CREATE TABLE memory_notes (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
-- Scoped to a project and checked in SQL on every read, like threads.
|
||||
project_id TEXT NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
||||
-- A namespace inside the project, so two agents sharing a key do not read
|
||||
-- each other's notes.
|
||||
agent_id TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
-- Normalised to ",a,b," so a tag matches whole, never as part of another.
|
||||
tags TEXT NOT NULL DEFAULT '',
|
||||
created_at INTEGER NOT NULL,
|
||||
-- Touched on every write, so retention measures how long a note has gone
|
||||
-- unrefreshed rather than how old it is.
|
||||
updated_at INTEGER NOT NULL,
|
||||
UNIQUE (project_id, agent_id, key)
|
||||
);
|
||||
CREATE INDEX memory_notes_by_updated ON memory_notes(updated_at);
|
||||
|
||||
-- Full-text search over the fields an agent would search by. An external
|
||||
-- content table, so the text is stored once.
|
||||
CREATE VIRTUAL TABLE memory_notes_fts USING fts5(
|
||||
key, text, tags, content = 'memory_notes', content_rowid = 'id'
|
||||
);
|
||||
|
||||
CREATE TRIGGER memory_notes_ai AFTER INSERT ON memory_notes BEGIN
|
||||
INSERT INTO memory_notes_fts(rowid, key, text, tags) VALUES (new.id, new.key, new.text, new.tags);
|
||||
END;
|
||||
CREATE TRIGGER memory_notes_ad AFTER DELETE ON memory_notes BEGIN
|
||||
INSERT INTO memory_notes_fts(memory_notes_fts, rowid, key, text, tags)
|
||||
VALUES ('delete', old.id, old.key, old.text, old.tags);
|
||||
END;
|
||||
CREATE TRIGGER memory_notes_au AFTER UPDATE ON memory_notes BEGIN
|
||||
INSERT INTO memory_notes_fts(memory_notes_fts, rowid, key, text, tags)
|
||||
VALUES ('delete', old.id, old.key, old.text, old.tags);
|
||||
INSERT INTO memory_notes_fts(rowid, key, text, tags) VALUES (new.id, new.key, new.text, new.tags);
|
||||
END;
|
||||
@@ -0,0 +1,26 @@
|
||||
-- A lease on being the control plane that places work.
|
||||
--
|
||||
-- Two control planes against one store both reconcile: both see a model
|
||||
-- wanting two replicas and none running, and both place. The result is twice
|
||||
-- the replicas, each holding GPUs the other's inventory believes are free —
|
||||
-- the same failure a control plane restart used to produce before adoption,
|
||||
-- and harder to see because nothing crashed.
|
||||
--
|
||||
-- A lease rather than a lock held by the connection. PostgreSQL has advisory
|
||||
-- locks and SQLite has nothing like them, and this store keeps one set of
|
||||
-- queries for both; a row with an expiry works everywhere and survives the
|
||||
-- leader being killed rather than asked to stop.
|
||||
--
|
||||
-- One row, enforced by the primary key: there is one thing to be leader of.
|
||||
CREATE TABLE leadership (
|
||||
role TEXT PRIMARY KEY,
|
||||
-- Who holds it. An opaque id per process, so a process cannot be fooled by
|
||||
-- its own stale lease after a restart on the same host.
|
||||
owner TEXT NOT NULL,
|
||||
-- When the lease lapses. The holder renews well before this; anyone else
|
||||
-- may take it once the clock has passed, which is what makes a leader that
|
||||
-- was powered off recoverable without anyone being paged.
|
||||
expires_at INTEGER NOT NULL,
|
||||
-- For an operator asking who has been in charge and since when.
|
||||
acquired_at INTEGER NOT NULL
|
||||
);
|
||||
@@ -0,0 +1,61 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Open picks a backend from what the address looks like.
|
||||
//
|
||||
// A path is SQLite and a postgres:// URL is PostgreSQL, so an existing
|
||||
// configuration keeps working untouched and moving to PostgreSQL is one line
|
||||
// changed rather than a new setting to discover. A gateway on one machine has
|
||||
// no reason to run a database server, and a gateway that has outgrown that
|
||||
// should not have to be reconfigured twice.
|
||||
//
|
||||
// An address of the form @/path/to/file reads the real address out of that
|
||||
// file, which is how a DSN carrying a password stays out of a configuration
|
||||
// file that gets copied, diffed and pasted into issues.
|
||||
func Open(ctx context.Context, address string) (*SQLite, error) {
|
||||
address, err := resolveAddress(address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if isPostgres(address) {
|
||||
return OpenPostgres(ctx, address)
|
||||
}
|
||||
return OpenSQLite(ctx, address)
|
||||
}
|
||||
|
||||
// isPostgres reports whether an address names a PostgreSQL server.
|
||||
//
|
||||
// By scheme, because that is the only part of a DSN that is reliably there.
|
||||
// A bare host:port would be ambiguous with a relative file path, so it is not
|
||||
// accepted: an address this cannot classify should fail loudly rather than
|
||||
// quietly create a SQLite file named after someone's database server.
|
||||
func isPostgres(address string) bool {
|
||||
return strings.HasPrefix(address, "postgres://") ||
|
||||
strings.HasPrefix(address, "postgresql://")
|
||||
}
|
||||
|
||||
// resolveAddress expands an @file indirection.
|
||||
//
|
||||
// The file is read whole and trimmed, so a trailing newline — which every
|
||||
// editor and every `echo` adds — does not become part of a password.
|
||||
func resolveAddress(address string) (string, error) {
|
||||
if !strings.HasPrefix(address, "@") {
|
||||
return address, nil
|
||||
}
|
||||
path := strings.TrimPrefix(address, "@")
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("reading the store address from %s: %w", path, err)
|
||||
}
|
||||
out := strings.TrimSpace(string(b))
|
||||
if out == "" {
|
||||
return "", fmt.Errorf("%s is empty, so there is no store address in it", path)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAPathIsStillSQLite(t *testing.T) {
|
||||
// Every existing configuration names a file, and every one of them has to
|
||||
// keep working without being touched.
|
||||
for _, addr := range []string{
|
||||
"/var/lib/llmgw/llmgw.db", "llmgw.db", "./data/llmgw.db",
|
||||
} {
|
||||
if isPostgres(addr) {
|
||||
t.Errorf("%q was read as a PostgreSQL address", addr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPostgresURLIsRecognised(t *testing.T) {
|
||||
for _, addr := range []string{
|
||||
"postgres://u:p@h:5432/db", "postgresql://u@h/db?sslmode=disable",
|
||||
} {
|
||||
if !isPostgres(addr) {
|
||||
t.Errorf("%q was not read as a PostgreSQL address", addr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAddressFileKeepsThePasswordOutOfTheConfig(t *testing.T) {
|
||||
// A DSN carries a password, and a configuration file gets copied, diffed
|
||||
// and pasted into issues. The indirection is what lets the secret live
|
||||
// somewhere with its own permissions.
|
||||
dir := t.TempDir()
|
||||
p := filepath.Join(dir, "dsn")
|
||||
// With the trailing newline every editor and every echo adds, which must
|
||||
// not become part of the password.
|
||||
if err := os.WriteFile(p, []byte("postgres://u:secret@h/db\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
got, err := resolveAddress("@" + p)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got != "postgres://u:secret@h/db" {
|
||||
t.Errorf("resolved to %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAMissingOrEmptyAddressFileSaysSo(t *testing.T) {
|
||||
// Rather than falling back to creating a SQLite file named "@/etc/...",
|
||||
// which is what a silent failure here would produce.
|
||||
if _, err := resolveAddress("@/nonexistent/dsn"); err == nil {
|
||||
t.Error("a missing address file was accepted")
|
||||
}
|
||||
dir := t.TempDir()
|
||||
p := filepath.Join(dir, "empty")
|
||||
if err := os.WriteFile(p, []byte(" \n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err := resolveAddress("@" + p)
|
||||
if err == nil || !strings.Contains(err.Error(), "empty") {
|
||||
t.Errorf("an empty address file gave %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenStillCreatesASQLiteFile(t *testing.T) {
|
||||
p := filepath.Join(t.TempDir(), "llmgw.db")
|
||||
s, err := Open(context.Background(), p)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
if s.dialect != SQLiteDialect {
|
||||
t.Errorf("a path opened as %q", s.dialect)
|
||||
}
|
||||
if _, err := os.Stat(p); err != nil {
|
||||
t.Errorf("no database file was created: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
// The pure-Go PostgreSQL driver. Pure-Go is not a preference here: the
|
||||
// whole binary is built with CGO_ENABLED=0 so that it runs on a node whose
|
||||
// libc is older than the builder's, and a cgo driver would take that away
|
||||
// from every deployment, including the ones still on SQLite.
|
||||
_ "github.com/jackc/pgx/v5/stdlib"
|
||||
)
|
||||
|
||||
// OpenPostgres opens the store against PostgreSQL.
|
||||
//
|
||||
// The same queries run here as on SQLite; only the placeholders differ, and
|
||||
// the dialect handles that. What differs operationally is concurrency: SQLite
|
||||
// is held to one connection because more would only produce lock contention
|
||||
// that reads as a deadlock, while PostgreSQL is the reason someone moves to
|
||||
// it and wants a pool.
|
||||
//
|
||||
// The DSN carries a password, so it is taken as a string the caller read from
|
||||
// somewhere protected rather than as a path this package knows how to find —
|
||||
// and it is never logged, which is why errors below name the database and not
|
||||
// the DSN.
|
||||
func OpenPostgres(ctx context.Context, dsn string) (*SQLite, error) {
|
||||
if dsn == "" {
|
||||
return nil, fmt.Errorf("store: no PostgreSQL address configured")
|
||||
}
|
||||
db, err := sql.Open("pgx", dsn)
|
||||
if err != nil {
|
||||
// Deliberately not wrapped with the DSN: it holds a password, and an
|
||||
// error string ends up in logs, tickets and screenshots.
|
||||
return nil, fmt.Errorf("opening the PostgreSQL store: %w", err)
|
||||
}
|
||||
|
||||
// Sized for a control plane rather than for a fleet of application
|
||||
// servers. The request path barely touches the store — usage is batched,
|
||||
// budgets are counted in memory, authentication is cached — so the pool
|
||||
// exists for conversation turns and reporting queries, not for inference.
|
||||
db.SetMaxOpenConns(16)
|
||||
db.SetMaxIdleConns(4)
|
||||
// Long enough that a quiet cluster keeps its connections, short enough
|
||||
// that a restarted database is not answered with dead ones for an hour.
|
||||
db.SetConnMaxLifetime(30 * time.Minute)
|
||||
db.SetConnMaxIdleTime(5 * time.Minute)
|
||||
|
||||
if err := db.PingContext(ctx); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, fmt.Errorf("reaching the PostgreSQL store: %w", err)
|
||||
}
|
||||
if err := migrate(ctx, db, string(PostgresDialect)); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, err
|
||||
}
|
||||
return &SQLite{
|
||||
raw: db,
|
||||
dialect: PostgresDialect,
|
||||
db: bound{c: PostgresDialect, to: db},
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// openPG opens the test database named by LLMGW_TEST_POSTGRES, or skips.
|
||||
//
|
||||
// Env-gated rather than spun up here, because a store test that needs a
|
||||
// container is a test nobody runs. Pointed at a database, it runs the same
|
||||
// assertions the SQLite tests make — which is the only thing that can catch
|
||||
// the two backends drifting apart.
|
||||
//
|
||||
// Every table is emptied first. The alternative is a test database that
|
||||
// accumulates rows until an assertion about counts starts failing for reasons
|
||||
// that have nothing to do with the change being tested.
|
||||
func openPG(t *testing.T) *SQLite {
|
||||
t.Helper()
|
||||
dsn := os.Getenv("LLMGW_TEST_POSTGRES")
|
||||
if dsn == "" {
|
||||
t.Skip("set LLMGW_TEST_POSTGRES to a database DSN to run these")
|
||||
}
|
||||
requireTestDatabase(t, dsn)
|
||||
|
||||
s, err := OpenPostgres(context.Background(), dsn)
|
||||
if err != nil {
|
||||
t.Fatalf("opening the test database: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = s.Close() })
|
||||
|
||||
// CASCADE because the tables reference each other, and TRUNCATE rather
|
||||
// than DELETE so identity columns restart too.
|
||||
_, err = s.db.ExecContext(context.Background(),
|
||||
// leadership included, and it is not in copyOrder: a subtest that
|
||||
// inherits the previous one's lease is refused for reasons that have
|
||||
// nothing to do with what it is testing. It passed on SQLite, where
|
||||
// every subtest gets a fresh file, and failed only here.
|
||||
`TRUNCATE thread_messages, threads, memory_notes, usage, api_keys,
|
||||
key_claims, limits, models, sessions, users, projects, teams,
|
||||
organizations, leadership RESTART IDENTITY CASCADE`)
|
||||
if err != nil {
|
||||
t.Fatalf("clearing the test database: %v", err)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// requireTestDatabase refuses to run against a database that is not obviously
|
||||
// a test one.
|
||||
//
|
||||
// These tests truncate every table before each one. Pointed at a real
|
||||
// deployment that is a silent, complete data loss — every key, every user,
|
||||
// every conversation — and it is not a hypothetical: it happened here, through
|
||||
// an SSH tunnel to the production database that had been opened minutes
|
||||
// earlier for a perfectly good reason. The DSN was correct for what it was
|
||||
// opened for and catastrophic for what it was then used for.
|
||||
//
|
||||
// The name has to say so. A guard that could be satisfied by accident would
|
||||
// not have stopped it.
|
||||
func requireTestDatabase(t *testing.T, dsn string) {
|
||||
t.Helper()
|
||||
u, err := url.Parse(dsn)
|
||||
if err != nil {
|
||||
t.Fatalf("LLMGW_TEST_POSTGRES is not a URL: %v", err)
|
||||
}
|
||||
name := strings.TrimPrefix(u.Path, "/")
|
||||
if !strings.Contains(strings.ToLower(name), "test") {
|
||||
t.Fatalf("refusing to run destructive tests against database %q:\n"+
|
||||
" these TRUNCATE every table before each test.\n"+
|
||||
" Point LLMGW_TEST_POSTGRES at a database whose name contains \"test\".", name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresCarriesTheWholeTenancyTree(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openPG(t)
|
||||
|
||||
proj := tree(t, s)
|
||||
secret, k := mustKey(t, s, "k1", proj)
|
||||
_ = secret
|
||||
|
||||
got, err := s.KeyByHash(ctx, k.Hash)
|
||||
if err != nil {
|
||||
t.Fatalf("reading the key back: %v", err)
|
||||
}
|
||||
if got.ID != k.ID || got.ProjectID != proj {
|
||||
t.Errorf("key came back as %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresKeepsUsageAndSpend(t *testing.T) {
|
||||
// The accounting path, which is the one that decides whether a budget
|
||||
// means anything after a restart.
|
||||
ctx := context.Background()
|
||||
s := openPG(t)
|
||||
proj := tree(t, s)
|
||||
_, k := mustKey(t, s, "k1", proj)
|
||||
|
||||
now := time.Now()
|
||||
err := s.WriteUsage(ctx, []Usage{
|
||||
{At: now, KeyID: k.ID, ProjectID: proj, Model: "m", ServedModel: "m",
|
||||
PromptTokens: 10, CompletionTokens: 20, CostMicros: 500, Status: 200, LatencyMs: 42},
|
||||
{At: now, KeyID: k.ID, ProjectID: proj, Model: "m", ServedModel: "m",
|
||||
PromptTokens: 1, CompletionTokens: 2, CostMicros: 250, Status: 200, LatencyMs: 7},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("writing usage: %v", err)
|
||||
}
|
||||
|
||||
spend, err := s.SpendSince(ctx, now.Add(-time.Hour))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if spend[k.ID] != 750 {
|
||||
t.Errorf("spend = %d, want 750", spend[k.ID])
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresKeepsThreadsInOrderAndScoped(t *testing.T) {
|
||||
// The newest tables, and the ones with the most SQL per operation —
|
||||
// a transaction, a MAX(seq), an insert per turn.
|
||||
ctx := context.Background()
|
||||
s := openPG(t)
|
||||
proj := tree(t, s)
|
||||
if err := s.SetThreadRetention(ctx, proj, 24); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
for _, m := range []Message{
|
||||
{Role: "user", Content: "first"},
|
||||
{Role: "assistant", Content: "answer", Model: "small"},
|
||||
{Role: "user", Content: "second"},
|
||||
} {
|
||||
if err := s.AppendThread(ctx, proj, "th1", []Message{m}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
got, err := s.Thread(ctx, proj, "th1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got) != 3 || got[0].Content != "first" || got[2].Content != "second" {
|
||||
t.Fatalf("thread came back as %+v", got)
|
||||
}
|
||||
if got[1].Model != "small" {
|
||||
t.Errorf("the model that answered was lost: %+v", got[1])
|
||||
}
|
||||
|
||||
// And the scope holds here exactly as it does on SQLite.
|
||||
if err := s.PutProject(ctx, Project{ID: "proj2", TeamID: "team1", Name: "Other"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.Thread(ctx, "proj2", "th1"); !errors.Is(err, ErrNoThread) {
|
||||
t.Errorf("reading across projects returned %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresEnforcesLimitsTheSameWay(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openPG(t)
|
||||
proj := tree(t, s)
|
||||
_, k := mustKey(t, s, "k1", proj)
|
||||
|
||||
rpm := 7
|
||||
budget := int64(9000)
|
||||
if err := s.PutLimits(ctx, Limits{
|
||||
Scope: ScopeProject, ScopeID: proj, RPM: &rpm,
|
||||
Budget: &budget, BudgetPeriod: PeriodMonth,
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
eff, err := s.Resolve(ctx, k.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if eff.RPM == nil || *eff.RPM != 7 {
|
||||
t.Errorf("rpm did not resolve down the tree: %+v", eff.RPM)
|
||||
}
|
||||
if eff.Budget == nil || *eff.Budget != 9000 {
|
||||
t.Errorf("budget did not resolve down the tree: %+v", eff.Budget)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresUpsertsRatherThanDuplicating(t *testing.T) {
|
||||
// ON CONFLICT is the one non-trivial construct shared by both dialects.
|
||||
// If it behaved differently here, applying a configuration twice would
|
||||
// start failing instead of being a no-op.
|
||||
ctx := context.Background()
|
||||
s := openPG(t)
|
||||
tree(t, s)
|
||||
tree(t, s)
|
||||
|
||||
orgs, err := s.Orgs(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(orgs) != 1 {
|
||||
t.Errorf("applying the tree twice produced %d orgs", len(orgs))
|
||||
}
|
||||
}
|
||||
@@ -17,7 +17,14 @@ import (
|
||||
|
||||
// SQLite is the single-instance store.
|
||||
type SQLite struct {
|
||||
db *sql.DB
|
||||
// raw is for the things that are not queries: transactions, pings, close.
|
||||
raw *sql.DB
|
||||
dialect Dialect
|
||||
// db is the same handle with placeholders rewritten for the dialect. Every
|
||||
// query goes through it, which is the point — a call site that reached the
|
||||
// driver directly would work on SQLite and fail on PostgreSQL, and only on
|
||||
// the path nobody exercised before deploying.
|
||||
db bound
|
||||
|
||||
// mu guards the default limits, which are set once at wiring but read on
|
||||
// every authentication.
|
||||
@@ -63,11 +70,11 @@ func OpenSQLite(ctx context.Context, path string) (*SQLite, error) {
|
||||
return nil, fmt.Errorf("securing %s: %w", f, err)
|
||||
}
|
||||
}
|
||||
return &SQLite{db: db}, nil
|
||||
return &SQLite{raw: db, dialect: SQLiteDialect, db: bound{c: SQLiteDialect, to: db}}, nil
|
||||
}
|
||||
|
||||
// Close releases the database.
|
||||
func (s *SQLite) Close() error { return s.db.Close() }
|
||||
func (s *SQLite) Close() error { return s.raw.Close() }
|
||||
|
||||
// --- tenancy ---
|
||||
|
||||
@@ -193,7 +200,7 @@ func (s *SQLite) TouchKeys(ctx context.Context, seen map[string]time.Time) error
|
||||
if len(seen) == 0 {
|
||||
return nil
|
||||
}
|
||||
return s.inTx(ctx, func(tx *sql.Tx) error {
|
||||
return s.inTx(ctx, func(tx conn) error {
|
||||
stmt, err := tx.PrepareContext(ctx,
|
||||
`UPDATE api_keys SET last_used_at = ? WHERE id = ? AND (last_used_at IS NULL OR last_used_at < ?)`)
|
||||
if err != nil {
|
||||
@@ -378,7 +385,7 @@ func (s *SQLite) WriteUsage(ctx context.Context, records []Usage) error {
|
||||
if len(records) == 0 {
|
||||
return nil
|
||||
}
|
||||
return s.inTx(ctx, func(tx *sql.Tx) error {
|
||||
return s.inTx(ctx, func(tx conn) error {
|
||||
stmt, err := tx.PrepareContext(ctx,
|
||||
`INSERT INTO usage (ts, key_id, project_id, model, served_model,
|
||||
prompt_tokens, completion_tokens, cost_micros, status, latency_ms)
|
||||
@@ -421,13 +428,26 @@ func (s *SQLite) SpendSince(ctx context.Context, since time.Time) (map[string]in
|
||||
|
||||
// --- helpers ---
|
||||
|
||||
func (s *SQLite) inTx(ctx context.Context, fn func(*sql.Tx) error) error {
|
||||
tx, err := s.db.BeginTx(ctx, nil)
|
||||
// begin starts a transaction and returns it twice: bound for queries, raw for
|
||||
// commit and rollback. The bound half is what callers must use, so a statement
|
||||
// inside a transaction is rewritten exactly as one outside it is.
|
||||
func (s *SQLite) begin(ctx context.Context) (conn, *sql.Tx, error) {
|
||||
tx, err := s.raw.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return bound{c: s.dialect, to: tx}, tx, nil
|
||||
}
|
||||
|
||||
// inTx runs fn in a transaction, with the same placeholder rewriting the
|
||||
// outside-transaction path gets.
|
||||
func (s *SQLite) inTx(ctx context.Context, fn func(conn) error) error {
|
||||
tx, err := s.raw.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
if err := fn(tx); err != nil {
|
||||
if err := fn(bound{c: s.dialect, to: tx}); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
@@ -453,7 +473,7 @@ func scanKey(sc scanner) (Key, error) {
|
||||
return k, nil
|
||||
}
|
||||
|
||||
func queryAll[T any](ctx context.Context, db *sql.DB, query string, args []any,
|
||||
func queryAll[T any](ctx context.Context, db conn, query string, args []any,
|
||||
scan func(*sql.Rows) (T, error),
|
||||
) ([]T, error) {
|
||||
rows, err := db.QueryContext(ctx, query, args...)
|
||||
|
||||
@@ -85,7 +85,7 @@ func TestTheKeyItselfIsNeverStored(t *testing.T) {
|
||||
project := tree(t, s)
|
||||
secret, _ := mustKey(t, s, "key1", project)
|
||||
|
||||
rows, err := s.db.Query(`SELECT id, key_hash, key_prefix, name FROM api_keys`)
|
||||
rows, err := s.db.QueryContext(context.Background(), `SELECT id, key_hash, key_prefix, name FROM api_keys`)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ErrThreadsNotEnabled means the project has not asked llmgw to keep
|
||||
// conversations.
|
||||
//
|
||||
// Returned rather than silently ignoring the thread. Ignoring it would send
|
||||
// the model a conversation with no history, and it would answer confidently
|
||||
// about something it was never told — which reads as the model being wrong
|
||||
// rather than as a setting being off.
|
||||
var ErrThreadsNotEnabled = errors.New("store: this project does not keep conversations")
|
||||
|
||||
// ErrNoThread means the id does not exist, or belongs to another project.
|
||||
//
|
||||
// One error for both on purpose. Distinguishing them would turn thread ids
|
||||
// into an oracle: a caller could learn which ids exist in projects they cannot
|
||||
// read by watching which error comes back.
|
||||
var ErrNoThread = errors.New("store: no such thread")
|
||||
|
||||
// Message is one turn of a stored conversation.
|
||||
type Message struct {
|
||||
Role string
|
||||
Content string
|
||||
// Model is which model produced an assistant turn, empty for a user one. A
|
||||
// thread that moved between models otherwise reads as a single voice.
|
||||
Model string
|
||||
}
|
||||
|
||||
// EnsureAnonymousProject creates the single project that owns everything kept
|
||||
// by a gateway with no authentication, and sets how long it keeps it.
|
||||
//
|
||||
// A real project rather than a special case in every query: threads and notes
|
||||
// are scoped by project id and enforce it in SQL, and an empty id would have to
|
||||
// be an exception in each of those places.
|
||||
func (s *SQLite) EnsureAnonymousProject(ctx context.Context, id string, retainHours int) error {
|
||||
if err := s.PutOrg(ctx, Org{ID: id, Name: id}); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.PutTeam(ctx, Team{ID: id, OrgID: id, Name: id}); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.PutProject(ctx, Project{ID: id, TeamID: id, Name: id}); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.SetThreadRetention(ctx, id, retainHours); err != nil {
|
||||
return err
|
||||
}
|
||||
return s.SetMemoryRetention(ctx, id, retainHours)
|
||||
}
|
||||
|
||||
// ThreadRetention reports how long a project keeps conversations.
|
||||
//
|
||||
// Zero means it keeps none, which is the default and was the only behaviour
|
||||
// before threads existed.
|
||||
func (s *SQLite) ThreadRetention(ctx context.Context, projectID string) (time.Duration, error) {
|
||||
var hours sql.NullInt64
|
||||
err := s.db.QueryRowContext(ctx,
|
||||
`SELECT thread_retention_hours FROM projects WHERE id = ?`, projectID).Scan(&hours)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return 0, nil
|
||||
}
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !hours.Valid || hours.Int64 <= 0 {
|
||||
return 0, nil
|
||||
}
|
||||
return time.Duration(hours.Int64) * time.Hour, nil
|
||||
}
|
||||
|
||||
// SetThreadRetention turns conversation storage on for a project, or off.
|
||||
//
|
||||
// Hours rather than a boolean: a thread store with no expiry is a transcript
|
||||
// of everything a team ever typed, kept forever because nobody chose a
|
||||
// number. Turning it on means saying for how long.
|
||||
func (s *SQLite) SetThreadRetention(ctx context.Context, projectID string, hours int) error {
|
||||
var v any
|
||||
if hours > 0 {
|
||||
v = hours
|
||||
}
|
||||
res, err := s.db.ExecContext(ctx,
|
||||
`UPDATE projects SET thread_retention_hours = ? WHERE id = ?`, v, projectID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return fmt.Errorf("store: no project %q", projectID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Thread returns a conversation's messages in order.
|
||||
//
|
||||
// The project is part of the lookup rather than checked afterwards, so there
|
||||
// is no path that reads another project's conversation even briefly.
|
||||
func (s *SQLite) Thread(ctx context.Context, projectID, threadID string) ([]Message, error) {
|
||||
var owner string
|
||||
err := s.db.QueryRowContext(ctx,
|
||||
`SELECT project_id FROM threads WHERE id = ?`, threadID).Scan(&owner)
|
||||
if errors.Is(err, sql.ErrNoRows) || (err == nil && owner != projectID) {
|
||||
return nil, ErrNoThread
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return queryAll(ctx, s.db,
|
||||
`SELECT role, content, COALESCE(model, '') FROM thread_messages
|
||||
WHERE thread_id = ? ORDER BY seq`,
|
||||
[]any{threadID}, func(rows *sql.Rows) (Message, error) {
|
||||
var m Message
|
||||
return m, rows.Scan(&m.Role, &m.Content, &m.Model)
|
||||
})
|
||||
}
|
||||
|
||||
// AppendThread adds turns to a conversation, creating it if it is new.
|
||||
//
|
||||
// One transaction, because a half-written exchange is worse than a lost one:
|
||||
// a stored question with no answer would be replayed to the next model as
|
||||
// though the assistant had ignored it.
|
||||
func (s *SQLite) AppendThread(ctx context.Context, projectID, threadID string, msgs []Message) error {
|
||||
if len(msgs) == 0 {
|
||||
return nil
|
||||
}
|
||||
retention, err := s.ThreadRetention(ctx, projectID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if retention == 0 {
|
||||
return ErrThreadsNotEnabled
|
||||
}
|
||||
|
||||
q, tx, err := s.begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
now := time.Now().UnixMilli()
|
||||
var owner string
|
||||
err = q.QueryRowContext(ctx, `SELECT project_id FROM threads WHERE id = ?`, threadID).Scan(&owner)
|
||||
switch {
|
||||
case errors.Is(err, sql.ErrNoRows):
|
||||
if _, err := q.ExecContext(ctx,
|
||||
`INSERT INTO threads (id, project_id, created_at, updated_at) VALUES (?, ?, ?, ?)`,
|
||||
threadID, projectID, now, now); err != nil {
|
||||
return err
|
||||
}
|
||||
case err != nil:
|
||||
return err
|
||||
case owner != projectID:
|
||||
// Appending to a thread another project owns would let one team write
|
||||
// into another's conversation, and then read it back through their
|
||||
// own replies.
|
||||
return ErrNoThread
|
||||
}
|
||||
|
||||
var next int64
|
||||
if err := q.QueryRowContext(ctx,
|
||||
`SELECT COALESCE(MAX(seq), -1) + 1 FROM thread_messages WHERE thread_id = ?`,
|
||||
threadID).Scan(&next); err != nil {
|
||||
return err
|
||||
}
|
||||
for i, m := range msgs {
|
||||
var model any
|
||||
if m.Model != "" {
|
||||
model = m.Model
|
||||
}
|
||||
if _, err := q.ExecContext(ctx,
|
||||
`INSERT INTO thread_messages (thread_id, seq, role, content, model) VALUES (?, ?, ?, ?, ?)`,
|
||||
threadID, next+int64(i), m.Role, m.Content, model); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
// Touched on every append, so retention measures idleness rather than age.
|
||||
if _, err := q.ExecContext(ctx,
|
||||
`UPDATE threads SET updated_at = ? WHERE id = ?`, now, threadID); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// DeleteThread removes a conversation and everything in it.
|
||||
//
|
||||
// Somebody has to be able to say "forget that", and a retention window is not
|
||||
// an answer to a message sent by mistake.
|
||||
func (s *SQLite) DeleteThread(ctx context.Context, projectID, threadID string) error {
|
||||
res, err := s.db.ExecContext(ctx,
|
||||
`DELETE FROM threads WHERE id = ? AND project_id = ?`, threadID, projectID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return ErrNoThread
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// PruneThreads deletes conversations past their project's retention.
|
||||
//
|
||||
// Per project, because the window is per project. A single sweep with one
|
||||
// cutoff would apply the shortest retention to everyone or the longest to
|
||||
// nobody, and the number a project chose is the promise made to it.
|
||||
func (s *SQLite) PruneThreads(ctx context.Context, now time.Time) (int, error) {
|
||||
type window struct {
|
||||
project string
|
||||
hours int64
|
||||
}
|
||||
windows, err := queryAll(ctx, s.db,
|
||||
`SELECT id, thread_retention_hours FROM projects WHERE thread_retention_hours > 0`,
|
||||
nil, func(rows *sql.Rows) (window, error) {
|
||||
var w window
|
||||
return w, rows.Scan(&w.project, &w.hours)
|
||||
})
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
var removed int
|
||||
for _, w := range windows {
|
||||
cutoff := now.Add(-time.Duration(w.hours) * time.Hour).UnixMilli()
|
||||
res, err := s.db.ExecContext(ctx,
|
||||
`DELETE FROM threads WHERE project_id = ? AND updated_at < ?`, w.project, cutoff)
|
||||
if err != nil {
|
||||
return removed, err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
removed += int(n)
|
||||
}
|
||||
return removed, nil
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestAProjectKeepsNoConversationsUnlessItAsks(t *testing.T) {
|
||||
// The property being given up here is worth protecting in a test: until a
|
||||
// project turns this on, llmgw stores no message content anywhere. Every
|
||||
// project that existed before threads did, and every one created since
|
||||
// without asking, must stay that way — a gateway that quietly starts
|
||||
// keeping what a team types is a much harder thing to have deployed.
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
proj := tree(t, s)
|
||||
|
||||
if d, err := s.ThreadRetention(ctx, proj); err != nil || d != 0 {
|
||||
t.Fatalf("a new project retains %v (err %v), want none", d, err)
|
||||
}
|
||||
|
||||
err := s.AppendThread(ctx, proj, "th1", []Message{{Role: "user", Content: "secret"}})
|
||||
if !errors.Is(err, ErrThreadsNotEnabled) {
|
||||
t.Fatalf("storing into a project that never asked returned %v", err)
|
||||
}
|
||||
|
||||
// And nothing was written on the way to refusing.
|
||||
if _, err := s.Thread(ctx, proj, "th1"); !errors.Is(err, ErrNoThread) {
|
||||
t.Errorf("a refused append left a thread behind: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestATurnedOnProjectStoresAndReplaysInOrder(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
proj := tree(t, s)
|
||||
if err := s.SetThreadRetention(ctx, proj, 24); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
turns := []Message{
|
||||
{Role: "user", Content: "first"},
|
||||
{Role: "assistant", Content: "answer one", Model: "small"},
|
||||
{Role: "user", Content: "second"},
|
||||
{Role: "assistant", Content: "answer two", Model: "large"},
|
||||
}
|
||||
for _, m := range turns {
|
||||
if err := s.AppendThread(ctx, proj, "th1", []Message{m}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
got, err := s.Thread(ctx, proj, "th1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got) != len(turns) {
|
||||
t.Fatalf("stored %d turns, read back %d", len(turns), len(got))
|
||||
}
|
||||
for i := range turns {
|
||||
if got[i].Role != turns[i].Role || got[i].Content != turns[i].Content {
|
||||
t.Errorf("turn %d came back as %+v", i, got[i])
|
||||
}
|
||||
}
|
||||
// Which model answered, because a thread that moved between models
|
||||
// otherwise reads as one voice — and when the answers disagree that is the
|
||||
// first thing anyone asks.
|
||||
if got[1].Model != "small" || got[3].Model != "large" {
|
||||
t.Errorf("the models that answered were lost: %q and %q", got[1].Model, got[3].Model)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnotherProjectsThreadIsNotReadable(t *testing.T) {
|
||||
// Thread ids are opaque, but an unguessable id is a secret, and secrets end
|
||||
// up in logs, shell history and bug reports. The project scope is what
|
||||
// makes a leaked id useless to anyone outside it.
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
mine := tree(t, s)
|
||||
if err := s.PutProject(ctx, Project{ID: "proj2", TeamID: "team1", Name: "Other"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, p := range []string{mine, "proj2"} {
|
||||
if err := s.SetThreadRetention(ctx, p, 24); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := s.AppendThread(ctx, mine, "th1", []Message{{Role: "user", Content: "private"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// The same error as a thread that does not exist. Two errors would let a
|
||||
// caller learn which ids are real in projects they cannot read.
|
||||
if _, err := s.Thread(ctx, "proj2", "th1"); !errors.Is(err, ErrNoThread) {
|
||||
t.Errorf("reading across projects returned %v", err)
|
||||
}
|
||||
if _, err := s.Thread(ctx, "proj2", "never-existed"); !errors.Is(err, ErrNoThread) {
|
||||
t.Errorf("a missing thread returned a different error: %v", err)
|
||||
}
|
||||
// Writing across projects is the same hole in the other direction: append,
|
||||
// then read your own reply back.
|
||||
if err := s.AppendThread(ctx, "proj2", "th1", []Message{{Role: "user", Content: "x"}}); !errors.Is(err, ErrNoThread) {
|
||||
t.Errorf("writing into another project's thread returned %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeletingAThreadTakesItsMessagesWithIt(t *testing.T) {
|
||||
// Somebody has to be able to say "forget that", and a retention window is
|
||||
// no answer to a message sent by mistake.
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
proj := tree(t, s)
|
||||
if err := s.SetThreadRetention(ctx, proj, 24); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.AppendThread(ctx, proj, "th1", []Message{{Role: "user", Content: "oops"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if err := s.DeleteThread(ctx, proj, "th1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.Thread(ctx, proj, "th1"); !errors.Is(err, ErrNoThread) {
|
||||
t.Errorf("the thread survived deletion: %v", err)
|
||||
}
|
||||
|
||||
// The rows, not just the row. Without the cascade the content would still
|
||||
// be in the database after someone was told it was gone.
|
||||
var left int
|
||||
if err := s.db.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM thread_messages WHERE thread_id = ?`, "th1").Scan(&left); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if left != 0 {
|
||||
t.Errorf("%d messages survived the thread being deleted", left)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetentionIsMeasuredFromTheLastTurnNotTheFirst(t *testing.T) {
|
||||
// Expiring a conversation someone is still having would be a strange way
|
||||
// to keep a promise about how long things are kept.
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
proj := tree(t, s)
|
||||
if err := s.SetThreadRetention(ctx, proj, 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.AppendThread(ctx, proj, "live", []Message{{Role: "user", Content: "hello"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Two hours on, a thread touched now is still inside a one-hour window.
|
||||
if n, err := s.PruneThreads(ctx, time.Now().Add(2*time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
} else if n != 1 {
|
||||
t.Fatalf("an idle thread past its window was not pruned: %d", n)
|
||||
}
|
||||
|
||||
if err := s.AppendThread(ctx, proj, "live2", []Message{{Role: "user", Content: "hello"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n, err := s.PruneThreads(ctx, time.Now()); err != nil {
|
||||
t.Fatal(err)
|
||||
} else if n != 0 {
|
||||
t.Errorf("a thread inside its window was pruned: %d", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPruningLeavesProjectsThatStoreNothingAlone(t *testing.T) {
|
||||
// They have nothing to prune, and a sweep that treated "no retention" as
|
||||
// "retain for zero hours" would be indistinguishable from one that worked.
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
tree(t, s)
|
||||
|
||||
if n, err := s.PruneThreads(ctx, time.Now()); err != nil || n != 0 {
|
||||
t.Errorf("pruning removed %d rows in a store with no threads (err %v)", n, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTurningRetentionOffAgain(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openDB(t)
|
||||
proj := tree(t, s)
|
||||
if err := s.SetThreadRetention(ctx, proj, 24); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.SetThreadRetention(ctx, proj, 0); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if d, err := s.ThreadRetention(ctx, proj); err != nil || d != 0 {
|
||||
t.Fatalf("retention is %v after being turned off (err %v)", d, err)
|
||||
}
|
||||
if err := s.AppendThread(ctx, proj, "th1", []Message{{Role: "user", Content: "x"}}); !errors.Is(err, ErrThreadsNotEnabled) {
|
||||
t.Errorf("storing continued after being turned off: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -98,12 +98,16 @@ requested whenever a fallback served it.</p>
|
||||
|
||||
<h2>Nodes</h2>
|
||||
{{if .Nodes}}
|
||||
<table><tr><th>Node</th><th>State</th><th class="num">GPUs</th><th class="num">RAM</th><th>Classes</th><th>Fabrics</th><th class="num">Models</th><th class="num">Replicas</th><th>Last seen</th></tr>
|
||||
<table><tr><th>Node</th><th>State</th><th class="num">GPUs</th><th class="num">RAM</th><th class="num">Disk</th><th>Classes</th><th>Fabrics</th><th class="num">Models</th><th class="num">Replicas</th><th>Last seen</th></tr>
|
||||
{{range .Nodes}}<tr>
|
||||
<td><code>{{.ID}}</code>{{if .MgmtAddr}} <span class="note">{{.MgmtAddr}}</span>{{end}}</td>
|
||||
<td class="state state-{{.State}}">{{.State}}</td>
|
||||
<td class="num">{{.GPUFree}}/{{.GPUTotal}}</td>
|
||||
<td class="num">{{if .MemoryGB}}{{.MemoryGB}}G{{else}}—{{end}}</td>
|
||||
{{/* Free out of total, because the first is meaningless without the
|
||||
second, and an em dash where nothing was reported: an agent that
|
||||
was never asked about its disk must not read as a full one. */}}
|
||||
<td class="num">{{if .DiskTotalGB}}{{.DiskFreeGB}}G/{{.DiskTotalGB}}G{{else}}—{{end}}</td>
|
||||
<td>{{join .Classes}}</td>
|
||||
<td>{{join .Fabrics}}</td>
|
||||
<td class="num">{{if .Models}}{{len .Models}}{{else}}—{{end}}</td>
|
||||
|
||||
@@ -0,0 +1,207 @@
|
||||
package e2e
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// callWithPrompt sends a prompt of a chosen size, which is what decides whether
|
||||
// a fallback can hold it.
|
||||
func (h *harness) callWithPrompt(t *testing.T, key, model, prompt string) (int, http.Header, string) {
|
||||
t.Helper()
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"model": model,
|
||||
"messages": []map[string]string{{"role": "user", "content": prompt}},
|
||||
"max_tokens": 4,
|
||||
})
|
||||
req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
|
||||
h.router.URL+"/v1/chat/completions", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
msg, _ := io.ReadAll(resp.Body)
|
||||
return resp.StatusCode, resp.Header, string(msg)
|
||||
}
|
||||
|
||||
// timeoutAfter is the poll these tests need because spend is written off the
|
||||
// request path: the budget is not spent the instant the first reply returns.
|
||||
// A predicate rather than arithmetic at four call sites, so the loops read as
|
||||
// loops.
|
||||
func timeoutAfter() func() bool {
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
return func() bool {
|
||||
if time.Now().After(deadline) {
|
||||
return false
|
||||
}
|
||||
time.Sleep(25 * time.Millisecond)
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
const roomyModel = "roomy-model"
|
||||
|
||||
// contextFallbackModels is a large-context model with a small-context one
|
||||
// behind it, which is the shape that made this a bug: the fit is checked
|
||||
// against what the caller asked for, and something else ends up serving.
|
||||
func contextFallbackModels() []catalog.Model {
|
||||
return []catalog.Model{
|
||||
{
|
||||
Name: testModel, Source: "/models/test",
|
||||
ContextLength: 32768, MaxOutput: 4096,
|
||||
Fallbacks: []string{cheapModel, roomyModel},
|
||||
Scaling: catalog.Scaling{Min: 1, Max: 1},
|
||||
Profiles: []catalog.Profile{{Name: "p", Engine: "llamacpp", Weight: 1}},
|
||||
},
|
||||
{
|
||||
// Too small to hold a large prompt, and first in the list, so it is
|
||||
// the one that would have been chosen.
|
||||
Name: cheapModel, Source: "/models/cheap",
|
||||
ContextLength: 512, MaxOutput: 128,
|
||||
Scaling: catalog.Scaling{Min: 1, Max: 1},
|
||||
Profiles: []catalog.Profile{{Name: "p", Engine: "llamacpp", Weight: 1}},
|
||||
},
|
||||
{
|
||||
Name: roomyModel, Source: "/models/roomy",
|
||||
ContextLength: 32768, MaxOutput: 4096,
|
||||
Scaling: catalog.Scaling{Min: 1, Max: 1},
|
||||
Profiles: []catalog.Profile{{Name: "p", Engine: "llamacpp", Weight: 1}},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func contextFallbackConfig(path string) func(*config.Server) {
|
||||
return func(c *config.Server) {
|
||||
paidConfig(path)(c)
|
||||
c.Models = contextFallbackModels()
|
||||
// The same engines answer for every name, so this is about model choice
|
||||
// rather than placement.
|
||||
for i := range c.Endpoints {
|
||||
c.Endpoints[i].Models = []string{testModel, cheapModel, roomyModel}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// spendTheBudget exhausts a key so the next request falls back.
|
||||
func spendTheBudget(t *testing.T, path, id string) {
|
||||
t.Helper()
|
||||
withStore(t, path, func(s *store.SQLite) {
|
||||
budget := int64(1)
|
||||
if err := s.PutLimits(context.Background(), store.Limits{
|
||||
Scope: store.ScopeKey, ScopeID: id,
|
||||
Budget: &budget, BudgetPeriod: store.PeriodMonth,
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestAFallbackTooSmallForThePromptIsSkipped(t *testing.T) {
|
||||
// The bug this covers: the context check runs against the model the caller
|
||||
// asked for, and the budget fallback swaps in a different one afterwards.
|
||||
// A 20k prompt sent to a 32k model would land on a 512-token one, and the
|
||||
// engine would truncate it or refuse — leaving the caller with a failure
|
||||
// naming a model they never chose.
|
||||
path, mint := provision(t)
|
||||
key := mint(t, "broke")
|
||||
spendTheBudget(t, path, "key-broke")
|
||||
|
||||
h := start(t, 2, contextFallbackConfig(path))
|
||||
|
||||
// Comfortably beyond the small model's 512 tokens, and well within 32768.
|
||||
big := strings.Repeat("word ", 8000)
|
||||
|
||||
if code, _, msg := h.callWithPrompt(t, key, testModel, big); code != http.StatusOK {
|
||||
t.Fatalf("the first request should go through: %d %s", code, msg)
|
||||
}
|
||||
|
||||
deadline := timeoutAfter()
|
||||
for deadline() {
|
||||
code, hdr, msg := h.callWithPrompt(t, key, testModel, big)
|
||||
served := hdr.Get("x-llmgw-served-model")
|
||||
if served == cheapModel {
|
||||
t.Fatalf("served by a model that cannot hold the prompt: %s", msg)
|
||||
}
|
||||
if code == http.StatusOK && served == roomyModel {
|
||||
return // skipped the small one, took the one that fits
|
||||
}
|
||||
}
|
||||
t.Fatal("the budget never ran out, or no fallback served")
|
||||
}
|
||||
|
||||
func TestWhenNoAffordableModelFitsTheRequestIsRefused(t *testing.T) {
|
||||
// Being served by a model that cannot hold the prompt is worse than a
|
||||
// refusal: it fails later, in the engine, naming a model the caller did
|
||||
// not choose. The history is deliberately not trimmed to make it fit —
|
||||
// dropping someone's messages to stay under a limit they did not ask about
|
||||
// would be a silent loss inside a request that was already being quietly
|
||||
// moved.
|
||||
path, mint := provision(t)
|
||||
key := mint(t, "broke")
|
||||
spendTheBudget(t, path, "key-broke")
|
||||
|
||||
h := start(t, 2, func(c *config.Server) {
|
||||
contextFallbackConfig(path)(c)
|
||||
// Only the small model behind it now, so nothing affordable fits.
|
||||
c.Models[0].Fallbacks = []string{cheapModel}
|
||||
})
|
||||
|
||||
big := strings.Repeat("word ", 8000)
|
||||
if code, _, msg := h.callWithPrompt(t, key, testModel, big); code != http.StatusOK {
|
||||
t.Fatalf("the first request should go through: %d %s", code, msg)
|
||||
}
|
||||
|
||||
deadline := timeoutAfter()
|
||||
for deadline() {
|
||||
code, hdr, _ := h.callWithPrompt(t, key, testModel, big)
|
||||
if hdr.Get("x-llmgw-served-model") == cheapModel {
|
||||
t.Fatal("a prompt was sent to a model too small to hold it")
|
||||
}
|
||||
if code == http.StatusTooManyRequests {
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatal("the budget never ran out")
|
||||
}
|
||||
|
||||
func TestASmallPromptStillReachesTheSmallFallback(t *testing.T) {
|
||||
// The check must not disable the feature. A prompt that fits is served by
|
||||
// the first affordable model as before.
|
||||
path, mint := provision(t)
|
||||
key := mint(t, "broke")
|
||||
spendTheBudget(t, path, "key-broke")
|
||||
|
||||
h := start(t, 2, contextFallbackConfig(path))
|
||||
|
||||
if code, _, msg := h.callWithPrompt(t, key, testModel, "hi"); code != http.StatusOK {
|
||||
t.Fatalf("the first request should go through: %d %s", code, msg)
|
||||
}
|
||||
|
||||
deadline := timeoutAfter()
|
||||
for deadline() {
|
||||
code, hdr, _ := h.callWithPrompt(t, key, testModel, "hi")
|
||||
if code == http.StatusOK && hdr.Get("x-llmgw-served-model") == cheapModel {
|
||||
return
|
||||
}
|
||||
if code == http.StatusTooManyRequests {
|
||||
t.Fatal("a prompt that fits was refused instead of falling back")
|
||||
}
|
||||
}
|
||||
t.Fatal("the budget never ran out, or the fallback never served")
|
||||
}
|
||||
@@ -0,0 +1,298 @@
|
||||
package e2e
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// say sends one turn, optionally as part of a thread, and returns what came
|
||||
// back along with the headers that explain what llmgw did.
|
||||
func (h *harness) say(t *testing.T, key, model, thread, text string) (int, http.Header, string) {
|
||||
t.Helper()
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"model": model,
|
||||
"messages": []map[string]string{{"role": "user", "content": text}},
|
||||
"max_tokens": 16,
|
||||
})
|
||||
req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
|
||||
h.router.URL+"/v1/chat/completions", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
if thread != "" {
|
||||
req.Header.Set("x-llmgw-thread", thread)
|
||||
}
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
msg, _ := io.ReadAll(resp.Body)
|
||||
return resp.StatusCode, resp.Header, string(msg)
|
||||
}
|
||||
|
||||
// authConfig turns on authentication against a provisioned store. The e2e
|
||||
// tests set this inline elsewhere; threads need it in five places.
|
||||
func authConfig(path string) func(*config.Server) {
|
||||
return func(c *config.Server) {
|
||||
c.Auth = config.Auth{
|
||||
Enabled: true, Store: path,
|
||||
CacheTTL: 50 * time.Millisecond,
|
||||
FlushInterval: 20 * time.Millisecond,
|
||||
BatchSize: 8,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func enableThreads(t *testing.T, path string, hours int) {
|
||||
t.Helper()
|
||||
withStore(t, path, func(s *store.SQLite) {
|
||||
if err := s.SetThreadRetention(context.Background(), "proj", hours); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func storedTurns(t *testing.T, path, thread string) []store.Message {
|
||||
t.Helper()
|
||||
var out []store.Message
|
||||
withStore(t, path, func(s *store.SQLite) {
|
||||
msgs, err := s.Thread(context.Background(), "proj", thread)
|
||||
if err != nil {
|
||||
t.Fatalf("reading thread %q: %v", thread, err)
|
||||
}
|
||||
out = msgs
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
func TestAProjectThatNeverAskedStoresNothingAndSaysSo(t *testing.T) {
|
||||
// The property being protected: llmgw keeps no conversations unless a
|
||||
// project turns it on. Refused rather than ignored, because ignoring the
|
||||
// header would send the model a conversation with no history and it would
|
||||
// answer confidently about something it was never told.
|
||||
path, mint := provision(t)
|
||||
key := mint(t, "user")
|
||||
h := start(t, 2, authConfig(path))
|
||||
|
||||
code, _, msg := h.say(t, key, testModel, "th1", "hello")
|
||||
if code != http.StatusBadRequest {
|
||||
t.Fatalf("a thread on a project that keeps none returned %d: %s", code, msg)
|
||||
}
|
||||
if !strings.Contains(msg, "does not keep conversations") {
|
||||
t.Errorf("the refusal does not say why: %s", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAThreadRemembersTheEarlierTurns(t *testing.T) {
|
||||
path, mint := provision(t)
|
||||
key := mint(t, "user")
|
||||
enableThreads(t, path, 24)
|
||||
h := start(t, 2, authConfig(path))
|
||||
|
||||
if code, _, msg := h.say(t, key, testModel, "th1", "first question"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
if code, _, msg := h.say(t, key, testModel, "th1", "second question"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
|
||||
turns := storedTurns(t, path, "th1")
|
||||
if len(turns) != 4 {
|
||||
t.Fatalf("stored %d turns, expected two questions and two answers: %+v", len(turns), turns)
|
||||
}
|
||||
if turns[0].Content != "first question" || turns[2].Content != "second question" {
|
||||
t.Errorf("the questions were not stored in order: %+v", turns)
|
||||
}
|
||||
if turns[1].Role != "assistant" || turns[1].Content == "" {
|
||||
t.Errorf("the first answer was not stored: %+v", turns[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAThreadContinuesOnADifferentModel(t *testing.T) {
|
||||
// The whole point. A conversation started on one model is carried to
|
||||
// another, and the stored turns say which model produced each answer —
|
||||
// without that, a thread that moved reads as one voice, and when the
|
||||
// answers disagree that is the first thing anyone asks.
|
||||
path, mint := provision(t)
|
||||
key := mint(t, "user")
|
||||
enableThreads(t, path, 24)
|
||||
h := start(t, 2, func(c *config.Server) {
|
||||
authConfig(path)(c)
|
||||
c.Models = budgetFallbackModels()
|
||||
for i := range c.Endpoints {
|
||||
c.Endpoints[i].Models = []string{testModel, cheapModel}
|
||||
}
|
||||
})
|
||||
|
||||
if code, _, msg := h.say(t, key, testModel, "th1", "started here"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
if code, _, msg := h.say(t, key, cheapModel, "th1", "continued there"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
|
||||
turns := storedTurns(t, path, "th1")
|
||||
models := map[string]bool{}
|
||||
for _, m := range turns {
|
||||
if m.Role == "assistant" {
|
||||
models[m.Model] = true
|
||||
}
|
||||
}
|
||||
if !models[testModel] || !models[cheapModel] {
|
||||
t.Errorf("the thread does not record both models that answered: %+v", turns)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOneProjectCannotReadAnothersThread(t *testing.T) {
|
||||
// Thread ids are opaque, but an unguessable id is a secret and secrets end
|
||||
// up in logs and bug reports. The project scope is what makes a leaked id
|
||||
// useless outside it.
|
||||
path, mint := provision(t)
|
||||
mine := mint(t, "mine")
|
||||
enableThreads(t, path, 24)
|
||||
|
||||
// A key in a second project.
|
||||
withStore(t, path, func(s *store.SQLite) {
|
||||
if err := s.PutProject(context.Background(),
|
||||
store.Project{ID: "other", TeamID: "team", Name: "Other"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.SetThreadRetention(context.Background(), "other", 24); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
theirs := mint(t, "theirs", func(k *store.Key) { k.ProjectID = "other" })
|
||||
|
||||
h := start(t, 2, authConfig(path))
|
||||
if code, _, msg := h.say(t, mine, testModel, "shared-id", "a private thing"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
|
||||
// The other project using the same id must not see it, and must not be
|
||||
// able to append to it either.
|
||||
code, _, msg := h.say(t, theirs, testModel, "shared-id", "what was said?")
|
||||
if code == http.StatusOK {
|
||||
if strings.Contains(msg, "a private thing") {
|
||||
t.Fatal("one project read another's conversation")
|
||||
}
|
||||
// Served, but from an empty history of its own — check the store
|
||||
// agrees that nothing crossed over.
|
||||
if turns := storedTurns(t, path, "shared-id"); len(turns) > 2 {
|
||||
for _, m := range turns {
|
||||
if m.Content == "a private thing" {
|
||||
t.Fatal("a second project appended into another's thread")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithoutAThreadNothingIsStoredAtAll(t *testing.T) {
|
||||
// Every request that does not opt in must take exactly the path it took
|
||||
// before threads existed.
|
||||
path, mint := provision(t)
|
||||
key := mint(t, "user")
|
||||
enableThreads(t, path, 24)
|
||||
h := start(t, 2, authConfig(path))
|
||||
|
||||
if code, _, msg := h.say(t, key, testModel, "", "no thread here"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
|
||||
withStore(t, path, func(s *store.SQLite) {
|
||||
// Pruning with a cutoff far past every retention window deletes every
|
||||
// thread there is, so a count of zero means none were created.
|
||||
n, err := s.PruneThreads(context.Background(), time.Now().Add(1000*time.Hour))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n != 0 {
|
||||
t.Errorf("%d threads were created by a request that used none", n)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// forget calls the delete endpoint the way a client would.
|
||||
func (h *harness) forget(t *testing.T, key, thread string) int {
|
||||
t.Helper()
|
||||
req, err := http.NewRequestWithContext(context.Background(), http.MethodDelete,
|
||||
h.router.URL+"/v1/threads/"+thread, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
_, _ = io.Copy(io.Discard, resp.Body)
|
||||
return resp.StatusCode
|
||||
}
|
||||
|
||||
func TestSomeoneCanUnsayWhatTheySaid(t *testing.T) {
|
||||
path, mint := provision(t)
|
||||
key := mint(t, "user")
|
||||
enableThreads(t, path, 24)
|
||||
h := start(t, 2, authConfig(path))
|
||||
|
||||
if code, _, msg := h.say(t, key, testModel, "th1", "please forget this"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
if turns := storedTurns(t, path, "th1"); len(turns) == 0 {
|
||||
t.Fatal("nothing was stored to delete")
|
||||
}
|
||||
|
||||
if code := h.forget(t, key, "th1"); code != http.StatusNoContent {
|
||||
t.Fatalf("deleting the thread returned %d", code)
|
||||
}
|
||||
withStore(t, path, func(s *store.SQLite) {
|
||||
if _, err := s.Thread(context.Background(), "proj", "th1"); err == nil {
|
||||
t.Error("the conversation survived being deleted")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestDeletingSomeoneElsesThreadLooksLikeDeletingNothing(t *testing.T) {
|
||||
// The same answer as a thread that never existed. Two answers would let a
|
||||
// caller map which ids are real in projects they cannot read — and this
|
||||
// one would do it by trying to destroy them.
|
||||
path, mint := provision(t)
|
||||
mine := mint(t, "mine")
|
||||
enableThreads(t, path, 24)
|
||||
withStore(t, path, func(s *store.SQLite) {
|
||||
if err := s.PutProject(context.Background(),
|
||||
store.Project{ID: "other", TeamID: "team", Name: "Other"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
theirs := mint(t, "theirs", func(k *store.Key) { k.ProjectID = "other" })
|
||||
|
||||
h := start(t, 2, authConfig(path))
|
||||
if code, _, msg := h.say(t, mine, testModel, "th1", "mine alone"); code != http.StatusOK {
|
||||
t.Fatalf("%d %s", code, msg)
|
||||
}
|
||||
|
||||
if code := h.forget(t, theirs, "th1"); code != http.StatusNotFound {
|
||||
t.Errorf("deleting another project's thread returned %d", code)
|
||||
}
|
||||
if code := h.forget(t, theirs, "never-existed"); code != http.StatusNotFound {
|
||||
t.Errorf("deleting a missing thread returned a different code: %d", code)
|
||||
}
|
||||
// And it is still there.
|
||||
if turns := storedTurns(t, path, "th1"); len(turns) == 0 {
|
||||
t.Error("another project destroyed a conversation it could not read")
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user