Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7dd1de6537 | ||
|
|
c815b629e2 | ||
|
|
2eb79f3400 | ||
|
|
45e913bd10 | ||
|
|
cf20570ff1 | ||
|
|
21a502035f | ||
|
|
eaf3382a6d | ||
|
|
d26a30471e | ||
|
|
9eb321ae3e | ||
|
|
cebf9906d2 | ||
|
|
b0592670a4 | ||
|
|
c9c841387a | ||
|
|
1b7e11b91e | ||
|
|
e6bc7859dc | ||
|
|
774e134ed1 | ||
|
|
753f4e3265 | ||
|
|
b5177b70f6 | ||
|
|
b58d67602d |
No files matched your search
+101
@@ -11,11 +11,112 @@ 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 |
|
||||
| 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.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,7 @@ 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 |
|
||||
| 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 +38,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;
|
||||
|
||||
+126
-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.7.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,115 @@ 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.
|
||||
|
||||
## 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 +735,13 @@ 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 is silent while it runs.** llmgw will fetch weights for you, but
|
||||
it reports no progress, so a first start on a large model looks like a hang.
|
||||
`ls -l` in the model directory is the only progress bar there is.
|
||||
- **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.
|
||||
@@ -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"
|
||||
@@ -170,6 +174,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
|
||||
@@ -578,9 +587,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 +694,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,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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -35,6 +35,8 @@ func newTenancyCmd() *cobra.Command {
|
||||
newKeyListCmd(open),
|
||||
newKeyRevokeCmd(open),
|
||||
newLimitsCmd(open),
|
||||
newThreadsCmd(open),
|
||||
newMemoryCmd(open),
|
||||
)
|
||||
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
|
||||
}
|
||||
@@ -372,6 +372,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) {
|
||||
|
||||
@@ -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
|
||||
@@ -391,6 +412,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"
|
||||
@@ -239,7 +240,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 +261,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)
|
||||
}
|
||||
|
||||
@@ -340,6 +369,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,7 @@ type Controller struct {
|
||||
inv *inventory.Store
|
||||
cat *catalog.Live
|
||||
rec *reconciler.Reconciler
|
||||
disp reconciler.Dispatcher
|
||||
sink EndpointSink
|
||||
usage UsageSource
|
||||
|
||||
@@ -101,6 +102,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,6 +172,21 @@ func (c *Controller) Reconciler() *reconciler.Reconciler { return c.rec }
|
||||
// Catalog exposes the model catalog.
|
||||
func (c *Controller) Catalog() *catalog.Catalog { return c.cat.Load() }
|
||||
|
||||
// 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) {
|
||||
c.rec.Start(ctx)
|
||||
|
||||
@@ -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,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.
|
||||
|
||||
@@ -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
|
||||
@@ -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.OpenSQLite(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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,284 @@
|
||||
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
|
||||
}
|
||||
|
||||
tx, err := s.db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
now := time.Now().UnixMilli()
|
||||
res, err := tx.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 := tx.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 := tx.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 != "" {
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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,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
|
||||
}
|
||||
|
||||
tx, err := s.db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
now := time.Now().UnixMilli()
|
||||
var owner string
|
||||
err = tx.QueryRowContext(ctx, `SELECT project_id FROM threads WHERE id = ?`, threadID).Scan(&owner)
|
||||
switch {
|
||||
case errors.Is(err, sql.ErrNoRows):
|
||||
if _, err := tx.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 := tx.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 := tx.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 := tx.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