mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-12 22:33:54 -04:00
feat(distributed): let a worker name a broadcast the frontend decides on
A worker has no database and cannot NOTIFY, and it does not need to: every
message it sends is produced inside a handler the frontend invoked, so
there is always an open control response to write on. This adds the two
ends of that, and the authorization decision that sits between them, and
nothing that dispatches yet.
workerctl.Envelope gains Subject, a REQUEST and not an instruction. Empty
means the line is for this caller alone, which is what every pre-existing
progress line is, and omitempty keeps those lines byte-identical for an
older reader. It qualifies a progress line and never a reply line: a reply
is the worker's verdict about the work, and there is no version of
"publish my verdict for me" this control plane has to carry.
nodes.MayBroadcast is the replacement for pkg/natsauth's allow list, and
the inversion is the point. NATS read an EMPTY allow list as NO
RESTRICTION, which is why phase 3 refused to delete the backend branch and
spelled it {"_INBOX.>"}. This one reads an empty list, and an absent node
type, as DENY EVERYTHING, and a table-driven spec pins that by emptying
the agent entry and asserting all three of its subjects are then refused.
Matching goes through messaging.SubjectMatches, the one definition in the
tree, so a filter that fires here fires on the carrier.
nodes.Rebroadcaster.Handle returns a bool and never an error. A refused or
failed re-broadcast is logged and the RPC continues, because the RPC's
outcome is the worker's verdict about the work and a publish failure says
nothing about it. The return shape is asserted at compile time in the file
that states the rule, so changing it to an error does not compile.
ControlClient.CallStreaming's progress callback now takes the line's
subject alongside its raw bytes. The client no longer decodes progress at
all: what a line means is a question about the verb and whether a named
broadcast may be made is a question about the node, and it knows neither.
Both moved into installProgressBridge, which every one of the three
streaming call sites in unloader.go goes through. A line naming a subject
is dropped there rather than delivered as install progress, because
backend.install and backend.upgrade are a backend worker's verbs and a
backend worker is allowed no subjects.
agents.StreamPublisher is the producing end, a messaging.Publisher writing
NDJSON envelopes onto an in-flight control response and flushing each one,
so a tick reaches the frontend while the handler is still running.
Serialized, because two concurrent encodes on one http.ResponseWriter
interleave bytes and tear the framing.
Assisted-by: Claude Opus 5 [claude-code]
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
1 parent
e5ad2e18c4
commit
fcb93b128e
12 files changed
+809
-31
No files matched your search
@@ -0,0 +1,74 @@
|
||||
package agents
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"sync"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/workerctl"
|
||||
)
|
||||
|
||||
// StreamPublisher is a messaging.Publisher that writes progress lines onto an
|
||||
// in-flight control response instead of onto a bus.
|
||||
//
|
||||
// It exists because a worker has no database and therefore cannot NOTIFY, and
|
||||
// because it does not need to: every message a worker sends is produced inside
|
||||
// a handler the frontend invoked, so there is always an open response to write
|
||||
// on.
|
||||
//
|
||||
// A subject written here is a REQUEST and not a publish. Nothing on this side
|
||||
// is authorised: the frontend reading the stream checks the subject against the
|
||||
// allow list for the node's type and refuses what the worker has no business
|
||||
// on. That refusal is silent to this end on purpose, since a worker learning
|
||||
// which subjects a frontend will carry is the beginning of a worker that probes
|
||||
// for them.
|
||||
//
|
||||
// Writes are serialized: cogito's stream callbacks run on the goroutine that is
|
||||
// executing the agent, but the status callback and the tool-result callback can
|
||||
// interleave with a flush, and two concurrent json.Encoder writes on one
|
||||
// http.ResponseWriter interleave bytes and corrupt the NDJSON framing.
|
||||
type StreamPublisher struct {
|
||||
mu sync.Mutex
|
||||
enc *json.Encoder
|
||||
flush func()
|
||||
}
|
||||
|
||||
// Compile-time proof that this is the Publisher a caller can hand to anything
|
||||
// that publishes. Asserted here rather than at the first adopter so a change to
|
||||
// either side fails to build in the file that owns the type.
|
||||
var _ messaging.Publisher = (*StreamPublisher)(nil)
|
||||
|
||||
// NewStreamPublisher returns a StreamPublisher writing to w. flush may be nil,
|
||||
// which is the right shape for a writer that has no buffering to push through.
|
||||
func NewStreamPublisher(w io.Writer, flush func()) *StreamPublisher {
|
||||
return &StreamPublisher{enc: json.NewEncoder(w), flush: flush}
|
||||
}
|
||||
|
||||
// Publish writes one Envelope carrying subject and the JSON encoding of data,
|
||||
// then flushes, so a progress tick reaches the frontend while the handler is
|
||||
// still running rather than when the response closes.
|
||||
//
|
||||
// data is encoded OUTSIDE the lock. Encoding is the expensive half and it
|
||||
// touches nothing shared, so holding the lock across it would serialise every
|
||||
// caller behind the slowest one's marshalling for no gain in framing.
|
||||
func (p *StreamPublisher) Publish(subject string, data any) error {
|
||||
raw, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("encoding a stream message for %q: %w", subject, err)
|
||||
}
|
||||
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if err := p.enc.Encode(workerctl.Envelope{Subject: subject, Progress: raw}); err != nil {
|
||||
return fmt.Errorf("writing a stream message for %q: %w", subject, err)
|
||||
}
|
||||
// The flush is inside the lock and belongs there. A flush concurrent with
|
||||
// an Encode pushes a partially written line at the reader, which is the
|
||||
// same torn frame the lock exists to prevent.
|
||||
if p.flush != nil {
|
||||
p.flush()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
package agents
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"runtime"
|
||||
"sync"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/workerctl"
|
||||
)
|
||||
|
||||
// chunkWriter is an http.ResponseWriter's concurrency contract and nothing
|
||||
// else: it is NOT safe for concurrent use, and it splits every Write into
|
||||
// single bytes with a scheduling point between them.
|
||||
//
|
||||
// Both halves are deliberate. A double that locked internally would make two
|
||||
// unsynchronised Encode calls look atomic and could never fail the way a real
|
||||
// response fails, which is the failure this spec exists to catch. Splitting the
|
||||
// write is what turns "two goroutines might interleave" into "two goroutines
|
||||
// do", so the concurrency spec reddens on the framing itself and not only under
|
||||
// the race detector.
|
||||
type chunkWriter struct {
|
||||
buf []byte
|
||||
}
|
||||
|
||||
func (w *chunkWriter) Write(p []byte) (int, error) {
|
||||
for _, b := range p {
|
||||
w.buf = append(w.buf, b)
|
||||
runtime.Gosched()
|
||||
}
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
// errWriter fails every write, for the one thing Publish promises about a
|
||||
// response it cannot write on.
|
||||
type errWriter struct{ err error }
|
||||
|
||||
func (w errWriter) Write([]byte) (int, error) { return 0, w.err }
|
||||
|
||||
// decodeLines reads every envelope out of a buffer, which is the assertion that
|
||||
// fails on interleaving: a torn line does not decode, and a line that swallowed
|
||||
// its neighbour does not produce two.
|
||||
func decodeLines(raw []byte) ([]workerctl.Envelope, error) {
|
||||
dec := json.NewDecoder(bytes.NewReader(raw))
|
||||
var out []workerctl.Envelope
|
||||
for {
|
||||
var env workerctl.Envelope
|
||||
err := dec.Decode(&env)
|
||||
if errors.Is(err, io.EOF) {
|
||||
return out, nil
|
||||
}
|
||||
if err != nil {
|
||||
return out, err
|
||||
}
|
||||
out = append(out, env)
|
||||
}
|
||||
}
|
||||
|
||||
var _ = Describe("StreamPublisher", func() {
|
||||
var (
|
||||
buf *bytes.Buffer
|
||||
flushes int
|
||||
pub *StreamPublisher
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
buf = &bytes.Buffer{}
|
||||
flushes = 0
|
||||
pub = NewStreamPublisher(buf, func() { flushes++ })
|
||||
})
|
||||
|
||||
It("writes exactly one NDJSON line whose subject and progress decode back to the inputs", func() {
|
||||
Expect(pub.Publish("agent.a1.events.status", map[string]any{"state": "thinking"})).To(Succeed())
|
||||
|
||||
lines, err := decodeLines(buf.Bytes())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(lines).To(HaveLen(1))
|
||||
Expect(lines[0].Subject).To(Equal("agent.a1.events.status"))
|
||||
Expect(string(lines[0].Progress)).To(MatchJSON(`{"state":"thinking"}`))
|
||||
// A line the frontend would read as terminal ends the stream, so a
|
||||
// progress publisher must never write one.
|
||||
Expect(lines[0].Reply).To(BeNil())
|
||||
Expect(bytes.Count(buf.Bytes(), []byte("\n"))).To(Equal(1))
|
||||
})
|
||||
|
||||
It("flushes after every line, so a tick reaches the frontend while the handler still runs", func() {
|
||||
Expect(pub.Publish("jobs.j1.progress", map[string]any{"percentage": 10})).To(Succeed())
|
||||
Expect(flushes).To(Equal(1))
|
||||
Expect(pub.Publish("jobs.j1.progress", map[string]any{"percentage": 20})).To(Succeed())
|
||||
Expect(flushes).To(Equal(2))
|
||||
})
|
||||
|
||||
It("writes a line at all when it was given no flush to call", func() {
|
||||
plain := &bytes.Buffer{}
|
||||
Expect(NewStreamPublisher(plain, nil).Publish("jobs.j1.result", map[string]any{"ok": true})).To(Succeed())
|
||||
lines, err := decodeLines(plain.Bytes())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(lines).To(HaveLen(1))
|
||||
})
|
||||
|
||||
It("keeps every concurrent publish a complete, decodable line of its own", func() {
|
||||
w := &chunkWriter{}
|
||||
p := NewStreamPublisher(w, func() {})
|
||||
|
||||
// Four writers rather than two, and enough lines each that an
|
||||
// unsynchronised run has to interleave rather than merely be allowed
|
||||
// to. Measured: at two writers of twenty lines an unlocked publisher
|
||||
// still produced a decodable buffer on some schedules, so the framing
|
||||
// assertion below was a coin toss and only the race detector was
|
||||
// reliably red.
|
||||
subjects := []string{
|
||||
"agent.a1.events.status",
|
||||
"agent.a2.events.status",
|
||||
"jobs.j1.progress",
|
||||
"jobs.j2.result",
|
||||
}
|
||||
const perGoroutine = 60
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, len(subjects)*perGoroutine)
|
||||
for _, subject := range subjects {
|
||||
wg.Add(1)
|
||||
go func(subject string) {
|
||||
defer wg.Done()
|
||||
for i := 0; i < perGoroutine; i++ {
|
||||
if err := p.Publish(subject, map[string]any{"n": i, "subject": subject}); err != nil {
|
||||
errs <- err
|
||||
}
|
||||
}
|
||||
}(subject)
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
Expect(errs).ToNot(Receive())
|
||||
|
||||
lines, err := decodeLines(w.buf)
|
||||
Expect(err).ToNot(HaveOccurred(), "a line was torn by an interleaved write")
|
||||
Expect(lines).To(HaveLen(len(subjects) * perGoroutine))
|
||||
// Every line must still carry the pair it was written with. A decode
|
||||
// that happened to succeed on spliced bytes would show up here as a
|
||||
// subject that does not match its payload.
|
||||
for _, line := range lines {
|
||||
var payload struct {
|
||||
Subject string `json:"subject"`
|
||||
}
|
||||
Expect(json.Unmarshal(line.Progress, &payload)).To(Succeed())
|
||||
Expect(payload.Subject).To(Equal(line.Subject))
|
||||
}
|
||||
})
|
||||
|
||||
It("reports a value it cannot encode without writing anything", func() {
|
||||
Expect(pub.Publish("jobs.j1.progress", make(chan int))).ToNot(Succeed())
|
||||
Expect(buf.Len()).To(BeZero())
|
||||
Expect(flushes).To(BeZero())
|
||||
})
|
||||
|
||||
It("reports a response it cannot write on, and does not flush a line it did not write", func() {
|
||||
failing := NewStreamPublisher(errWriter{err: errors.New("connection reset")}, func() { flushes++ })
|
||||
Expect(failing.Publish("jobs.j1.progress", map[string]any{"percentage": 10})).ToNot(Succeed())
|
||||
Expect(flushes).To(BeZero())
|
||||
})
|
||||
})
|
||||
@@ -11,10 +11,7 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/xlog"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/cluster"
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/workerctl"
|
||||
"github.com/mudler/LocalAI/pkg/httpclient"
|
||||
)
|
||||
@@ -152,6 +149,14 @@ func (c *ControlClient) Call(ctx context.Context, nodeID, path string, req, repl
|
||||
// stream, invoking onProgress for each progress line and decoding the single
|
||||
// terminal reply line into reply. onProgress may be nil.
|
||||
//
|
||||
// onProgress receives the line's SUBJECT alongside its bytes, so a caller can
|
||||
// tell a private progress tick from a re-broadcast request without this client
|
||||
// understanding either. An empty subject is a line for this caller alone, which
|
||||
// is what every pre-existing progress line is. The client deliberately does no
|
||||
// authorization and no decoding of its own: whether a named broadcast may be
|
||||
// made is a policy question about the node, and what a private tick means is a
|
||||
// question about the verb, and this type knows neither.
|
||||
//
|
||||
// onProgress runs SYNCHRONOUSLY, on this goroutine. The NATS carrier ran each
|
||||
// progress callback on a goroutine of its own because a slow callback there
|
||||
// stalled the one reader thread every worker's events arrived on; here the only
|
||||
@@ -159,7 +164,7 @@ func (c *ControlClient) Call(ctx context.Context, nodeID, path string, req, repl
|
||||
// caller's business. Dropping the guard is also what makes the events arrive in
|
||||
// the order the worker sent them.
|
||||
func (c *ControlClient) CallStreaming(ctx context.Context, nodeID, path string,
|
||||
req, reply any, onProgress func(messaging.BackendInstallProgressEvent)) error {
|
||||
req, reply any, onProgress func(subject string, raw json.RawMessage)) error {
|
||||
resp, err := c.do(ctx, nodeID, path, req)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -188,14 +193,7 @@ func (c *ControlClient) CallStreaming(ctx context.Context, nodeID, path string,
|
||||
if env.Progress == nil || onProgress == nil {
|
||||
continue
|
||||
}
|
||||
var ev messaging.BackendInstallProgressEvent
|
||||
if err := json.Unmarshal(env.Progress, &ev); err != nil {
|
||||
// Progress is transient by contract, so a line this frontend cannot
|
||||
// read costs a tick and never the operation.
|
||||
xlog.Debug("unreadable control progress line", "node", nodeID, "path", path, "error", err)
|
||||
continue
|
||||
}
|
||||
onProgress(ev)
|
||||
onProgress(env.Subject, env.Progress)
|
||||
}
|
||||
if reply == nil {
|
||||
return nil
|
||||
|
||||
@@ -223,13 +223,58 @@ var _ = Describe("ControlClient", func() {
|
||||
_ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":100}`)})
|
||||
_ = enc.Encode(workerctl.Envelope{Reply: json.RawMessage(`{"success":true}`)})
|
||||
}
|
||||
var seen []float64
|
||||
var seen []string
|
||||
var subjects []string
|
||||
var reply messaging.BackendInstallReply
|
||||
err := client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall,
|
||||
messaging.BackendInstallRequest{Backend: "mock", OpID: "op-1"}, &reply,
|
||||
func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) })
|
||||
func(subject string, raw json.RawMessage) {
|
||||
subjects = append(subjects, subject)
|
||||
seen = append(seen, string(raw))
|
||||
})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(seen).To(Equal([]float64{50, 100}))
|
||||
Expect(seen).To(Equal([]string{`{"percentage":50}`, `{"percentage":100}`}))
|
||||
// A line that named no broadcast reaches the caller with an empty
|
||||
// subject, which is what every pre-existing progress line is.
|
||||
Expect(subjects).To(Equal([]string{"", ""}))
|
||||
Expect(reply.Success).To(BeTrue())
|
||||
})
|
||||
|
||||
It("carries the subject a worker named through to the caller, and acts on it itself not at all", func() {
|
||||
// The client is not where the authorization decision lives, and it
|
||||
// must not become where it lives by accident: it hands the subject
|
||||
// over and keeps reading. Whether the broadcast is made is a
|
||||
// question about the node's type that this type cannot answer.
|
||||
handler = func(w http.ResponseWriter, _ *http.Request) {
|
||||
enc := json.NewEncoder(w)
|
||||
_ = enc.Encode(workerctl.Envelope{
|
||||
Subject: "agent.a1.events.status",
|
||||
Progress: json.RawMessage(`{"state":"thinking"}`),
|
||||
})
|
||||
_ = enc.Encode(workerctl.Envelope{
|
||||
// A subject no worker of any type is allowed. The client
|
||||
// still carries it: refusing here would put the policy in
|
||||
// two places.
|
||||
Subject: "cache.invalidate.models",
|
||||
Progress: json.RawMessage(`{"model":"m1"}`),
|
||||
})
|
||||
_ = enc.Encode(workerctl.Envelope{Reply: json.RawMessage(`{"success":true}`)})
|
||||
}
|
||||
type line struct {
|
||||
subject string
|
||||
raw string
|
||||
}
|
||||
var seen []line
|
||||
var reply messaging.BackendInstallReply
|
||||
Expect(client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall,
|
||||
struct{}{}, &reply,
|
||||
func(subject string, raw json.RawMessage) {
|
||||
seen = append(seen, line{subject: subject, raw: string(raw)})
|
||||
})).To(Succeed())
|
||||
Expect(seen).To(Equal([]line{
|
||||
{subject: "agent.a1.events.status", raw: `{"state":"thinking"}`},
|
||||
{subject: "cache.invalidate.models", raw: `{"model":"m1"}`},
|
||||
}))
|
||||
Expect(reply.Success).To(BeTrue())
|
||||
})
|
||||
|
||||
@@ -239,11 +284,11 @@ var _ = Describe("ControlClient", func() {
|
||||
_ = enc.Encode(workerctl.Envelope{Reply: json.RawMessage(`{"success":true}`)})
|
||||
_ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":10}`)})
|
||||
}
|
||||
var seen []float64
|
||||
var seen []string
|
||||
var reply messaging.BackendInstallReply
|
||||
Expect(client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall,
|
||||
struct{}{}, &reply,
|
||||
func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) })).To(Succeed())
|
||||
func(_ string, raw json.RawMessage) { seen = append(seen, string(raw)) })).To(Succeed())
|
||||
Expect(seen).To(BeEmpty())
|
||||
Expect(reply.Success).To(BeTrue())
|
||||
})
|
||||
@@ -274,19 +319,26 @@ var _ = Describe("ControlClient", func() {
|
||||
Expect(errors.Is(err, io.ErrUnexpectedEOF)).To(BeTrue())
|
||||
})
|
||||
|
||||
It("keeps going past a progress line it cannot read, since progress is transient", func() {
|
||||
It("hands over a progress line it could not have decoded, because it does not decode them", func() {
|
||||
// The client used to unmarshal every progress line into an install
|
||||
// event and drop the ones that would not parse. It no longer knows
|
||||
// what a progress line means, so a line that is not an install
|
||||
// event is carried like any other and the DECODE is the caller's.
|
||||
// The rule that a line the frontend cannot read costs a tick and
|
||||
// never the operation now lives in installProgressBridge, and
|
||||
// unloader_test.go pins it there.
|
||||
handler = func(w http.ResponseWriter, _ *http.Request) {
|
||||
enc := json.NewEncoder(w)
|
||||
_ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":"not a number"}`)})
|
||||
_ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":70}`)})
|
||||
_ = enc.Encode(workerctl.Envelope{Reply: json.RawMessage(`{"success":true}`)})
|
||||
}
|
||||
var seen []float64
|
||||
var seen []string
|
||||
var reply messaging.BackendInstallReply
|
||||
Expect(client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall,
|
||||
struct{}{}, &reply,
|
||||
func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) })).To(Succeed())
|
||||
Expect(seen).To(Equal([]float64{70}))
|
||||
func(_ string, raw json.RawMessage) { seen = append(seen, string(raw)) })).To(Succeed())
|
||||
Expect(seen).To(Equal([]string{`{"percentage":"not a number"}`, `{"percentage":70}`}))
|
||||
Expect(reply.Success).To(BeTrue())
|
||||
})
|
||||
|
||||
|
||||
@@ -43,6 +43,13 @@ type scriptedControlWorkers struct {
|
||||
matched map[string][]matchedControlReply
|
||||
progress map[string][]messaging.BackendInstallProgressEvent
|
||||
|
||||
// rawProgress holds progress lines a spec wrote out as envelopes, for the
|
||||
// two things scriptProgress cannot express: a line NAMING a subject, and a
|
||||
// line whose payload is not an install-progress event at all. Both are
|
||||
// things a real worker on another build can put on the wire, so the double
|
||||
// has to be able to put them on the wire too.
|
||||
rawProgress map[string][]workerctl.Envelope
|
||||
|
||||
// unreachable and expired are keyed by NODE, not by verb, because they are
|
||||
// failures of the ROUTE and a route belongs to a node. They are what the
|
||||
// dialer answers with; see scriptUnroutable and scriptTimeout.
|
||||
@@ -74,6 +81,7 @@ func newScriptedControlWorkers() *scriptedControlWorkers {
|
||||
unsupported: map[string]bool{},
|
||||
matched: map[string][]matchedControlReply{},
|
||||
progress: map[string][]messaging.BackendInstallProgressEvent{},
|
||||
rawProgress: map[string][]workerctl.Envelope{},
|
||||
unreachable: map[string]bool{},
|
||||
expired: map[string]bool{},
|
||||
hangs: map[string]bool{},
|
||||
@@ -131,6 +139,7 @@ func (s *scriptedControlWorkers) serve(w http.ResponseWriter, r *http.Request) {
|
||||
reply := s.replies[key]
|
||||
matchers := s.matched[key]
|
||||
ticks := s.progress[key]
|
||||
rawTicks := s.rawProgress[key]
|
||||
s.mu.Unlock()
|
||||
|
||||
if hang {
|
||||
@@ -182,6 +191,9 @@ func (s *scriptedControlWorkers) serve(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
_ = enc.Encode(workerctl.Envelope{Progress: raw})
|
||||
}
|
||||
for _, env := range rawTicks {
|
||||
_ = enc.Encode(env)
|
||||
}
|
||||
_ = enc.Encode(workerctl.Envelope{Reply: reply})
|
||||
}
|
||||
|
||||
@@ -270,6 +282,14 @@ func (s *scriptedControlWorkers) scriptProgress(key string, events []messaging.B
|
||||
s.progress[key] = events
|
||||
}
|
||||
|
||||
// scriptRawProgress queues progress lines exactly as they go on the wire, after
|
||||
// any scriptProgress lines for the same verb.
|
||||
func (s *scriptedControlWorkers) scriptRawProgress(key string, lines []workerctl.Envelope) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.rawProgress[key] = lines
|
||||
}
|
||||
|
||||
// scriptHang makes one verb on one node accept the request and never answer,
|
||||
// so the call ends only when the caller's budget does.
|
||||
//
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// workerBroadcastAllow lists, per node type, the subject filters a worker may
|
||||
// ask a frontend to publish on its behalf.
|
||||
//
|
||||
// An absent node type and an empty list both DENY EVERYTHING. That is the
|
||||
// opposite of the NATS allow list this replaces, where an empty list meant no
|
||||
// restriction, and it is the reason deleting an entry here is safe where
|
||||
// deleting a permissions branch there was a privilege escalation. Nothing may
|
||||
// be added to this map without a spec that a worker of that type is refused the
|
||||
// subjects NOT added.
|
||||
var workerBroadcastAllow = map[string][]string{
|
||||
NodeTypeAgent: {
|
||||
"jobs.*.progress",
|
||||
"jobs.*.result",
|
||||
"agent.*.events.*",
|
||||
},
|
||||
// A backend worker asks for no broadcasts. Spelled as an empty list rather
|
||||
// than omitted, so the reader sees the decision.
|
||||
NodeTypeBackend: {},
|
||||
}
|
||||
|
||||
// MayBroadcast reports whether a worker of nodeType may have subject published
|
||||
// on its behalf.
|
||||
func MayBroadcast(nodeType, subject string) bool {
|
||||
return mayBroadcastIn(workerBroadcastAllow, nodeType, subject)
|
||||
}
|
||||
|
||||
// mayBroadcastIn is MayBroadcast against a caller-supplied table.
|
||||
//
|
||||
// It exists so a spec can hold the table itself constant while varying what is
|
||||
// IN it, which is the only way to state "an empty list denies" as a property of
|
||||
// this function rather than as a property of today's entries. Production has
|
||||
// exactly one table and MayBroadcast is the only way to reach it.
|
||||
func mayBroadcastIn(table map[string][]string, nodeType, subject string) bool {
|
||||
// An empty subject is a line that asked for no broadcast at all, so there
|
||||
// is nothing here to authorise. Refusing it is not defensive tidiness: it
|
||||
// keeps a filter that happened to match the empty string from turning every
|
||||
// ordinary private progress tick into a publish.
|
||||
if subject == "" {
|
||||
return false
|
||||
}
|
||||
// A missing node type indexes to nil and the loop then runs zero times,
|
||||
// which is the same denial an empty list produces. One code path for both
|
||||
// is deliberate: "we have never heard of this worker" and "this worker is
|
||||
// allowed nothing" have the same safe answer, and giving them separate
|
||||
// branches invites one of the two to acquire an exception.
|
||||
for _, filter := range table[nodeType] {
|
||||
if messaging.SubjectMatches(filter, subject) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Rebroadcaster turns an allowed progress line into a broadcast.
|
||||
//
|
||||
// It holds a Broadcaster and not a MessagingClient because fan-out is all it
|
||||
// needs and all it may have: a worker asking for a broadcast must not be able
|
||||
// to reach a request/reply or a queue group through this path.
|
||||
type Rebroadcaster struct {
|
||||
bus messaging.Broadcaster
|
||||
}
|
||||
|
||||
// NewRebroadcaster returns a Rebroadcaster publishing on bus.
|
||||
func NewRebroadcaster(bus messaging.Broadcaster) *Rebroadcaster {
|
||||
return &Rebroadcaster{bus: bus}
|
||||
}
|
||||
|
||||
// The shape of Handle's answer is asserted here, not left to its callers.
|
||||
//
|
||||
// A refused or failed re-broadcast must never become an error, because the only
|
||||
// errors that reach a scheduler from this direction are the worker's own answer
|
||||
// and an unroutable peer, and a publish failure is neither. Changing the return
|
||||
// to an error therefore has to fail to COMPILE in the file that states the
|
||||
// rule, rather than fail a spec somewhere a reviewer might read as a test
|
||||
// needing an update.
|
||||
var _ func(string, string, json.RawMessage) bool = (*Rebroadcaster)(nil).Handle
|
||||
|
||||
// Handle publishes raw on subject when the worker's type allows it, and returns
|
||||
// false having published nothing when it does not. It never returns an error
|
||||
// that a caller could mistake for the worker's answer: a refused or failed
|
||||
// re-broadcast is logged and the RPC continues, because the RPC's outcome is
|
||||
// the worker's verdict about the work and a publish failure says nothing about
|
||||
// it.
|
||||
//
|
||||
// raw is published as it arrived. It is a json.RawMessage and every Broadcaster
|
||||
// encodes what it is given with json.Marshal, which returns a RawMessage
|
||||
// verbatim, so the subscriber reads the bytes the worker wrote rather than a
|
||||
// re-encoding of this frontend's idea of them.
|
||||
func (r *Rebroadcaster) Handle(nodeType, subject string, raw json.RawMessage) bool {
|
||||
if !MayBroadcast(nodeType, subject) {
|
||||
// Warn rather than Debug: this is a worker asking for something it has
|
||||
// no business on, which is the shape of a compromised or mismatched
|
||||
// worker and is worth seeing without turning logging up.
|
||||
xlog.Warn("refusing a worker's re-broadcast request",
|
||||
"nodeType", nodeType, "subject", subject)
|
||||
return false
|
||||
}
|
||||
if r == nil || r.bus == nil {
|
||||
xlog.Debug("no broadcaster to re-broadcast a worker's progress line on",
|
||||
"nodeType", nodeType, "subject", subject)
|
||||
return false
|
||||
}
|
||||
if err := r.bus.Publish(subject, raw); err != nil {
|
||||
xlog.Warn("a worker's re-broadcast could not be published",
|
||||
"nodeType", nodeType, "subject", subject, "error", err)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// failingBus is a Broadcaster whose Publish always fails, for the one thing
|
||||
// Handle promises about a carrier that is unhappy: the failure stays here.
|
||||
type failingBus struct {
|
||||
testutil.FakeBus
|
||||
err error
|
||||
}
|
||||
|
||||
func (b *failingBus) Publish(string, any) error { return b.err }
|
||||
|
||||
var _ = Describe("worker re-broadcast authorization", func() {
|
||||
Describe("MayBroadcast", func() {
|
||||
DescribeTable("answers for the node type it was asked about",
|
||||
func(nodeType, subject string, allowed bool) {
|
||||
Expect(MayBroadcast(nodeType, subject)).To(Equal(allowed))
|
||||
},
|
||||
Entry("agent, job progress", NodeTypeAgent, "jobs.j1.progress", true),
|
||||
Entry("agent, job result", NodeTypeAgent, "jobs.j1.result", true),
|
||||
Entry("agent, agent events", NodeTypeAgent, "agent.a1.events.status", true),
|
||||
// The subjects an agent worker is NOT allowed. Every entry added to
|
||||
// the table above owes this list a line, which is the only thing
|
||||
// keeping the allow list from growing into "everything an agent
|
||||
// happened to publish once".
|
||||
Entry("agent, cache invalidation", NodeTypeAgent, "cache.invalidate.models", false),
|
||||
Entry("agent, node registration", NodeTypeAgent, "node.register", false),
|
||||
Entry("agent, model load", NodeTypeAgent, "model.load.n1", false),
|
||||
Entry("agent, a job subject one token too deep", NodeTypeAgent, "jobs.j1.progress.extra", false),
|
||||
Entry("agent, a job subject one token too shallow", NodeTypeAgent, "jobs.progress", false),
|
||||
// A backend worker asks for no broadcasts, including the ones an
|
||||
// agent worker is allowed: the table is per node type and not a
|
||||
// single global list with a type-shaped comment on it.
|
||||
Entry("backend, job progress", NodeTypeBackend, "jobs.j1.progress", false),
|
||||
Entry("backend, job result", NodeTypeBackend, "jobs.j1.result", false),
|
||||
Entry("backend, agent events", NodeTypeBackend, "agent.a1.events.status", false),
|
||||
// A node type nothing has heard of is denied rather than defaulted.
|
||||
Entry("unknown type, job progress", "router", "jobs.j1.progress", false),
|
||||
Entry("unknown type, agent events", "router", "agent.a1.events.status", false),
|
||||
Entry("empty type, job progress", "", "jobs.j1.progress", false),
|
||||
)
|
||||
|
||||
It("refuses a line that names no subject at all", func() {
|
||||
// A progress line with no subject asked for no broadcast. It must
|
||||
// not become one because some filter matched the empty string.
|
||||
Expect(MayBroadcast(NodeTypeAgent, "")).To(BeFalse())
|
||||
Expect(MayBroadcast(NodeTypeBackend, "")).To(BeFalse())
|
||||
})
|
||||
|
||||
It("refuses the tail wildcard, which the shared matcher does not implement", func() {
|
||||
// SubjectMatches is the one definition of matching in the tree and
|
||||
// it makes a '>' filter match nothing. Asserting it here is what
|
||||
// stops someone writing "agent.>" into the allow list and believing
|
||||
// they widened it, when what they did was narrow it to nothing.
|
||||
Expect(messaging.SubjectMatches("agent.>", "agent.a1.events.status")).To(BeFalse())
|
||||
})
|
||||
|
||||
DescribeTable("reads an EMPTY allow list as denying everything, which is the inversion of the NATS list this replaces",
|
||||
// NATS read an empty allow list as NO RESTRICTION, which is why the
|
||||
// permissions it replaced had to spell the backend worker's list as
|
||||
// {"_INBOX.>"} rather than leave it empty. This table is the
|
||||
// opposite, and this spec is what documents it: the assumption a
|
||||
// reviewer carries over from the deleted code is exactly wrong.
|
||||
func(subject string) {
|
||||
emptied := map[string][]string{
|
||||
NodeTypeAgent: {},
|
||||
NodeTypeBackend: {},
|
||||
}
|
||||
Expect(mayBroadcastIn(emptied, NodeTypeAgent, subject)).To(BeFalse())
|
||||
// The same subject against the real table, so the case cannot
|
||||
// pass by naming something that was never allowed.
|
||||
Expect(MayBroadcast(NodeTypeAgent, subject)).To(BeTrue())
|
||||
},
|
||||
Entry("job progress", "jobs.j1.progress"),
|
||||
Entry("job result", "jobs.j1.result"),
|
||||
Entry("agent events", "agent.a1.events.status"),
|
||||
)
|
||||
|
||||
It("reads a MISSING node type the same way as an empty list", func() {
|
||||
absent := map[string][]string{NodeTypeBackend: {}}
|
||||
Expect(mayBroadcastIn(absent, NodeTypeAgent, "jobs.j1.progress")).To(BeFalse())
|
||||
})
|
||||
})
|
||||
|
||||
Describe("Rebroadcaster.Handle", func() {
|
||||
var (
|
||||
bus *testutil.FakeBus
|
||||
rb *Rebroadcaster
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
bus = testutil.NewFakeBus()
|
||||
rb = NewRebroadcaster(bus)
|
||||
})
|
||||
|
||||
It("publishes an allowed subject and hands the subscriber the worker's own bytes", func() {
|
||||
delivered := make(chan []byte, 1)
|
||||
_, err := bus.Subscribe("agent.*.events.*", func(payload []byte) { delivered <- payload })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(rb.Handle(NodeTypeAgent, "agent.a1.events.status",
|
||||
json.RawMessage(`{"state":"thinking"}`))).To(BeTrue())
|
||||
|
||||
Expect(bus.PublishCount("agent.a1.events.status")).To(Equal(1))
|
||||
// Verbatim, not re-encoded: a subscriber decodes the worker's DTO
|
||||
// and a frontend that re-marshalled its own idea of the line would
|
||||
// hand it a different shape.
|
||||
Eventually(delivered).Should(Receive(MatchJSON(`{"state":"thinking"}`)))
|
||||
})
|
||||
|
||||
It("refuses a subject the agent worker has no business on, and publishes nothing", func() {
|
||||
Expect(rb.Handle(NodeTypeAgent, "cache.invalidate.models",
|
||||
json.RawMessage(`{"model":"m1"}`))).To(BeFalse())
|
||||
Expect(bus.PublishCount("cache.invalidate.models")).To(Equal(0))
|
||||
})
|
||||
|
||||
It("refuses a BACKEND worker a subject an agent worker would be allowed", func() {
|
||||
Expect(rb.Handle(NodeTypeBackend, "jobs.j1.progress",
|
||||
json.RawMessage(`{"percentage":10}`))).To(BeFalse())
|
||||
Expect(bus.PublishCount("jobs.j1.progress")).To(Equal(0))
|
||||
})
|
||||
|
||||
It("refuses a node type nothing has heard of", func() {
|
||||
Expect(rb.Handle("router", "jobs.j1.progress",
|
||||
json.RawMessage(`{"percentage":10}`))).To(BeFalse())
|
||||
Expect(bus.PublishCount("jobs.j1.progress")).To(Equal(0))
|
||||
})
|
||||
|
||||
It("reports a carrier that could not publish as a refusal, not as anything the worker said", func() {
|
||||
// Handle's return type is the whole point: there is no error here
|
||||
// for a caller to confuse with the RPC's outcome. A publish that
|
||||
// failed says nothing about whether the work succeeded.
|
||||
rb = NewRebroadcaster(&failingBus{err: errors.New("carrier is unhappy")})
|
||||
Expect(rb.Handle(NodeTypeAgent, "jobs.j1.progress",
|
||||
json.RawMessage(`{"percentage":10}`))).To(BeFalse())
|
||||
})
|
||||
|
||||
It("refuses rather than panicking when it holds no broadcaster", func() {
|
||||
Expect(NewRebroadcaster(nil).Handle(NodeTypeAgent, "jobs.j1.progress",
|
||||
json.RawMessage(`{"percentage":10}`))).To(BeFalse())
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -2,6 +2,7 @@ package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
@@ -185,6 +186,45 @@ func (a *RemoteUnloaderAdapter) UnloadRemoteModelContext(ctx context.Context, mo
|
||||
return unloadErr
|
||||
}
|
||||
|
||||
// installProgressBridge adapts an install or upgrade's progress sink to
|
||||
// CallStreaming's line callback, which carries a subject the client does not
|
||||
// interpret.
|
||||
//
|
||||
// It makes two decisions a caller of InstallBackend must not have to make.
|
||||
//
|
||||
// A line naming a SUBJECT is a re-broadcast request and not install progress,
|
||||
// so it is dropped here rather than delivered as a tick. backend.install and
|
||||
// backend.upgrade are a BACKEND worker's verbs, and MayBroadcast denies a
|
||||
// backend worker every subject, so this path has nothing to publish and no
|
||||
// broadcaster to publish it on. Delivering it to onProgress instead would put a
|
||||
// worker's arbitrary JSON through a decode into an install-progress event and
|
||||
// report whatever fell out as the state of a download.
|
||||
//
|
||||
// A line this frontend cannot decode costs a tick and never the operation,
|
||||
// because progress is transient by contract while the reply is the worker's
|
||||
// verdict.
|
||||
//
|
||||
// It returns nil for a nil sink so CallStreaming keeps its "no callback, no
|
||||
// work" path, rather than a non-nil closure wrapping a nil function.
|
||||
func installProgressBridge(nodeID, path string, onProgress func(messaging.BackendInstallProgressEvent)) func(string, json.RawMessage) {
|
||||
if onProgress == nil {
|
||||
return nil
|
||||
}
|
||||
return func(subject string, raw json.RawMessage) {
|
||||
if subject != "" {
|
||||
xlog.Warn("refusing a re-broadcast request on a backend control stream",
|
||||
"node", nodeID, "path", path, "subject", subject)
|
||||
return
|
||||
}
|
||||
var ev messaging.BackendInstallProgressEvent
|
||||
if err := json.Unmarshal(raw, &ev); err != nil {
|
||||
xlog.Debug("unreadable control progress line", "node", nodeID, "path", path, "error", err)
|
||||
return
|
||||
}
|
||||
onProgress(ev)
|
||||
}
|
||||
}
|
||||
|
||||
// InstallBackend asks a worker node to install a backend and start its process.
|
||||
// Idempotent on the worker: if the (modelID, replica) process is already
|
||||
// running, the worker short-circuits and returns its address; if the binary
|
||||
@@ -223,7 +263,7 @@ func (a *RemoteUnloaderAdapter) InstallBackend(
|
||||
Alias: alias,
|
||||
ReplicaIndex: int32(replicaIndex),
|
||||
OpID: opID,
|
||||
}, &reply, onProgress)
|
||||
}, &reply, installProgressBridge(nodeID, workerctl.PathBackendInstall, onProgress))
|
||||
if err != nil {
|
||||
if isRequestTimeout(err) {
|
||||
return nil, fmt.Errorf("%w (nodeID=%s backend=%s): %v",
|
||||
@@ -258,7 +298,7 @@ func (a *RemoteUnloaderAdapter) UpgradeBackend(nodeID, backendType, galleriesJSO
|
||||
Alias: alias,
|
||||
ReplicaIndex: int32(replicaIndex),
|
||||
OpID: opID,
|
||||
}, &reply, onProgress)
|
||||
}, &reply, installProgressBridge(nodeID, workerctl.PathBackendUpgrade, onProgress))
|
||||
if err != nil {
|
||||
if isRequestTimeout(err) {
|
||||
return nil, fmt.Errorf("%w (nodeID=%s backend=%s): %v",
|
||||
@@ -291,7 +331,7 @@ func (a *RemoteUnloaderAdapter) installWithForceFallback(nodeID, backendType, ga
|
||||
ReplicaIndex: int32(replicaIndex),
|
||||
Force: true,
|
||||
OpID: opID,
|
||||
}, &reply, onProgress)
|
||||
}, &reply, installProgressBridge(nodeID, workerctl.PathBackendInstall, onProgress))
|
||||
if err != nil {
|
||||
if isRequestTimeout(err) {
|
||||
return nil, fmt.Errorf("%w (nodeID=%s backend=%s): %v",
|
||||
|
||||
@@ -634,4 +634,104 @@ var _ = Describe("RemoteUnloaderAdapter install progress streaming", func() {
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(reply.Success).To(BeTrue())
|
||||
})
|
||||
|
||||
// The bridge from a control STREAM LINE to an install-progress event is one
|
||||
// definition used at three call sites: InstallBackend, UpgradeBackend and
|
||||
// installWithForceFallback. Each is exercised separately below, because a
|
||||
// site that stopped calling it would leave the other two green.
|
||||
DescribeTable("delivers a line that names no broadcast to the caller's install-progress sink, on every streaming verb",
|
||||
func(path string, call func(*RemoteUnloaderAdapter, func(messaging.BackendInstallProgressEvent)) error) {
|
||||
workers := newScriptedControlWorkers()
|
||||
workers.scriptReply(controlKey("n1", path), messaging.BackendInstallReply{Success: true})
|
||||
workers.scriptReply(controlKey("n1", workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: true})
|
||||
workers.scriptProgress(controlKey("n1", path), []messaging.BackendInstallProgressEvent{
|
||||
{OpID: "op-abc", NodeID: "n1", Backend: "vllm", Percentage: 25},
|
||||
})
|
||||
|
||||
adapter := NewRemoteUnloaderAdapter(nil, workers.controlClient(), time.Second, time.Second)
|
||||
var received []messaging.BackendInstallProgressEvent
|
||||
Expect(call(adapter, func(ev messaging.BackendInstallProgressEvent) {
|
||||
received = append(received, ev)
|
||||
})).To(Succeed())
|
||||
Expect(received).To(HaveLen(1))
|
||||
Expect(received[0].Percentage).To(Equal(float64(25)))
|
||||
},
|
||||
Entry("backend.install", workerctl.PathBackendInstall,
|
||||
func(a *RemoteUnloaderAdapter, cb func(messaging.BackendInstallProgressEvent)) error {
|
||||
_, err := a.InstallBackend("n1", "vllm", "", "[]", "", "", "", 0, "op-abc", cb)
|
||||
return err
|
||||
}),
|
||||
Entry("backend.upgrade", workerctl.PathBackendUpgrade,
|
||||
func(a *RemoteUnloaderAdapter, cb func(messaging.BackendInstallProgressEvent)) error {
|
||||
_, err := a.UpgradeBackend("n1", "vllm", "[]", "", "", "", 0, "op-abc", cb)
|
||||
return err
|
||||
}),
|
||||
Entry("the legacy force-install fallback", workerctl.PathBackendInstall,
|
||||
func(a *RemoteUnloaderAdapter, cb func(messaging.BackendInstallProgressEvent)) error {
|
||||
_, err := a.installWithForceFallback("n1", "vllm", "[]", "", "", "", 0, "op-abc", cb)
|
||||
return err
|
||||
}),
|
||||
)
|
||||
|
||||
DescribeTable("refuses a line that NAMES a broadcast rather than delivering it as install progress, on every streaming verb",
|
||||
func(path string, call func(*RemoteUnloaderAdapter, func(messaging.BackendInstallProgressEvent)) error) {
|
||||
// backend.install and backend.upgrade are a BACKEND worker's verbs,
|
||||
// and MayBroadcast denies a backend worker every subject, so a line
|
||||
// naming one is not something this path publishes. What it must
|
||||
// also not do is deliver it: the payload is a broadcast body and
|
||||
// decoding it into an install-progress event reports a percentage
|
||||
// nobody sent.
|
||||
workers := newScriptedControlWorkers()
|
||||
workers.scriptReply(controlKey("n1", path), messaging.BackendInstallReply{Success: true})
|
||||
workers.scriptReply(controlKey("n1", workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: true})
|
||||
workers.scriptRawProgress(controlKey("n1", path), []workerctl.Envelope{
|
||||
{Subject: "jobs.j1.progress", Progress: json.RawMessage(`{"percentage":99}`)},
|
||||
{Subject: "agent.a1.events.status", Progress: json.RawMessage(`{"percentage":98}`)},
|
||||
})
|
||||
|
||||
adapter := NewRemoteUnloaderAdapter(nil, workers.controlClient(), time.Second, time.Second)
|
||||
var received []messaging.BackendInstallProgressEvent
|
||||
Expect(call(adapter, func(ev messaging.BackendInstallProgressEvent) {
|
||||
received = append(received, ev)
|
||||
})).To(Succeed())
|
||||
// The RPC still returned the worker's reply: a refused re-broadcast
|
||||
// costs a line and never the operation.
|
||||
Expect(received).To(BeEmpty())
|
||||
},
|
||||
Entry("backend.install", workerctl.PathBackendInstall,
|
||||
func(a *RemoteUnloaderAdapter, cb func(messaging.BackendInstallProgressEvent)) error {
|
||||
_, err := a.InstallBackend("n1", "vllm", "", "[]", "", "", "", 0, "op-abc", cb)
|
||||
return err
|
||||
}),
|
||||
Entry("backend.upgrade", workerctl.PathBackendUpgrade,
|
||||
func(a *RemoteUnloaderAdapter, cb func(messaging.BackendInstallProgressEvent)) error {
|
||||
_, err := a.UpgradeBackend("n1", "vllm", "[]", "", "", "", 0, "op-abc", cb)
|
||||
return err
|
||||
}),
|
||||
Entry("the legacy force-install fallback", workerctl.PathBackendInstall,
|
||||
func(a *RemoteUnloaderAdapter, cb func(messaging.BackendInstallProgressEvent)) error {
|
||||
_, err := a.installWithForceFallback("n1", "vllm", "[]", "", "", "", 0, "op-abc", cb)
|
||||
return err
|
||||
}),
|
||||
)
|
||||
|
||||
It("keeps going past a progress line it cannot read, since progress is transient", func() {
|
||||
// The rule the control client used to carry. It moved to the bridge
|
||||
// with the decode, and it has to be pinned where it now lives: a line
|
||||
// this frontend cannot read costs a tick and never the operation.
|
||||
workers := newScriptedControlWorkers()
|
||||
workers.scriptReply(controlKey("n1", workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: true})
|
||||
workers.scriptRawProgress(controlKey("n1", workerctl.PathBackendInstall), []workerctl.Envelope{
|
||||
{Progress: json.RawMessage(`{"percentage":"not a number"}`)},
|
||||
{Progress: json.RawMessage(`{"percentage":70}`)},
|
||||
})
|
||||
|
||||
adapter := NewRemoteUnloaderAdapter(nil, workers.controlClient(), time.Second, time.Second)
|
||||
var seen []float64
|
||||
reply, err := adapter.InstallBackend("n1", "vllm", "", "[]", "", "", "", 0, "op-abc",
|
||||
func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(seen).To(Equal([]float64{70}))
|
||||
Expect(reply.Success).To(BeTrue())
|
||||
})
|
||||
})
|
||||
@@ -2,6 +2,7 @@ package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net"
|
||||
"os"
|
||||
@@ -96,12 +97,23 @@ var _ = Describe("the frontend's control client against the real worker", func()
|
||||
}
|
||||
|
||||
var seen []float64
|
||||
var subjects []string
|
||||
var reply messaging.BackendInstallReply
|
||||
err := client.CallStreaming(context.Background(), nodeID, workerctl.PathBackendInstall,
|
||||
messaging.BackendInstallRequest{Backend: "mock", OpID: "op-1"}, &reply,
|
||||
func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) })
|
||||
func(subject string, raw json.RawMessage) {
|
||||
var ev messaging.BackendInstallProgressEvent
|
||||
Expect(json.Unmarshal(raw, &ev)).To(Succeed())
|
||||
subjects = append(subjects, subject)
|
||||
seen = append(seen, ev.Percentage)
|
||||
})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(seen).To(Equal([]float64{50, 100}))
|
||||
// A BACKEND worker's install progress names no broadcast, over a real
|
||||
// tunnel and a real NDJSON body rather than in a double's imagination.
|
||||
// It is the property MayBroadcast would refuse anyway, asserted here on
|
||||
// the bytes the worker actually wrote.
|
||||
Expect(subjects).To(Equal([]string{"", ""}))
|
||||
Expect(reply.Success).To(BeTrue())
|
||||
Expect(reply.WorkerLocalAddress).To(Equal("127.0.0.1:41234"))
|
||||
})
|
||||
|
||||
@@ -125,12 +125,12 @@ func AllPaths() []string {
|
||||
|
||||
// Envelope is one line of a streaming control response.
|
||||
//
|
||||
// Exactly one of the two is set. Zero or more Progress lines are followed by
|
||||
// exactly ONE Reply line, and the Reply line is the last thing on the body.
|
||||
// That ordering is the contract: it is what lets the frontend stop reading, and
|
||||
// it is what replaces the subscribe-before-request dance the NATS carrier
|
||||
// needed, since progress and reply now share one response and nothing can
|
||||
// arrive before the caller is listening.
|
||||
// Exactly one of Progress and Reply is set. Zero or more Progress lines are
|
||||
// followed by exactly ONE Reply line, and the Reply line is the last thing on
|
||||
// the body. That ordering is the contract: it is what lets the frontend stop
|
||||
// reading, and it is what replaces the subscribe-before-request dance the NATS
|
||||
// carrier needed, since progress and reply now share one response and nothing
|
||||
// can arrive before the caller is listening.
|
||||
//
|
||||
// Progress carrying the reply's own bytes is also why the 8000-byte
|
||||
// notification cap that bounded the NATS progress subject has no analogue here:
|
||||
@@ -138,6 +138,20 @@ func AllPaths() []string {
|
||||
type Envelope struct {
|
||||
Progress json.RawMessage `json:"progress,omitempty"`
|
||||
Reply json.RawMessage `json:"reply,omitempty"`
|
||||
|
||||
// Subject names the broadcast this progress line asks the frontend reading
|
||||
// the stream to make on its behalf. Empty means the line is for this caller
|
||||
// alone, which is what every pre-existing progress line is.
|
||||
//
|
||||
// It is a REQUEST and not an instruction: the frontend checks it against an
|
||||
// allow list derived from the node's type before publishing anything. A
|
||||
// worker naming a subject it has no business on is refused and logged, and
|
||||
// the stream continues.
|
||||
//
|
||||
// It qualifies a Progress line and never a Reply line. A reply is the
|
||||
// worker's verdict about the work, and there is no version of "publish my
|
||||
// verdict for me" that this control plane has to carry.
|
||||
Subject string `json:"subject,omitempty"`
|
||||
}
|
||||
|
||||
// ContentTypeStream is the media type of a streaming control response.
|
||||
|
||||
@@ -124,6 +124,37 @@ var _ = Describe("control plane paths on the wire", func() {
|
||||
Expect(string(b)).To(Equal(`{"progress":{"percentage":50}}`))
|
||||
})
|
||||
|
||||
It("spells the re-broadcast key on the wire as \"subject\"", func() {
|
||||
// Hand-written literal, like the paths above and for the same reason: a
|
||||
// worker and a frontend built from different commits read each other's
|
||||
// lines, and a renamed key is a re-broadcast request that silently
|
||||
// becomes a private progress tick. Deriving the expectation from the
|
||||
// struct tag would pin nothing.
|
||||
b, err := json.Marshal(workerctl.Envelope{
|
||||
Subject: "agent.a1.events.status",
|
||||
Progress: json.RawMessage(`{"tick":1}`),
|
||||
})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(string(b)).To(Equal(`{"progress":{"tick":1},"subject":"agent.a1.events.status"}`))
|
||||
})
|
||||
|
||||
It("omits the subject key entirely when a progress line names no broadcast", func() {
|
||||
// omitempty is the backward-compatibility half: every pre-existing
|
||||
// progress line has no subject, and an older frontend reading a line
|
||||
// that carried an empty subject key would be reading a field it does
|
||||
// not know about on every tick.
|
||||
b, err := json.Marshal(workerctl.Envelope{Progress: json.RawMessage(`{"tick":1}`)})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(string(b)).NotTo(ContainSubstring("subject"))
|
||||
})
|
||||
|
||||
It("reads a subject back off the wire, which is what the frontend decides on", func() {
|
||||
var env workerctl.Envelope
|
||||
Expect(json.Unmarshal([]byte(`{"progress":{"tick":1},"subject":"jobs.j1.progress"}`), &env)).To(Succeed())
|
||||
Expect(env.Subject).To(Equal("jobs.j1.progress"))
|
||||
Expect(string(env.Progress)).To(Equal(`{"tick":1}`))
|
||||
})
|
||||
|
||||
It("names the streaming media type", func() {
|
||||
Expect(workerctl.ContentTypeStream).To(Equal("application/x-ndjson"))
|
||||
})
|
||||
|
||||
Reference in new issue
Block a user