llmgw/internal/kvevents/subscriber.go

package kvevents

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"log/slog"
	"net"
	"sync"

	"github.com/go-zeromq/zmq4"
)

// Subscriber receives one engine's events and folds them into a mirror.
//
// It listens and the engine connects, which is the direction vLLM chooses for
// an endpoint naming a host: a publisher binds only "tcp://*" (every
// address), "::" or ipc, and connects to anything else. Listening here on the
// loopback address keeps the stream on the machine. The alternative, the
// engine binding every address, would publish the token ids of every prompt
// to anyone who can reach the node.
type Subscriber struct {
	replica string
	mirror  *Mirror
	log     *slog.Logger
	sock    zmq4.Socket
	port    int
	cancel  context.CancelFunc
	done    chan struct{}

	mu      sync.Mutex
	lastSeq int64
	seen    bool
}

// Listen starts receiving one replica's events on 127.0.0.1:port; port 0
// picks a free one, which Port then reports. An agent restarting gives the
// port it chose before, because the engine it adopts is still connecting to
// it.
func Listen(replica string, port int, m *Mirror, log *slog.Logger) (*Subscriber, error) {
	if log == nil {
		log = slog.Default()
	}
	ctx, cancel := context.WithCancel(context.Background())
	sock := zmq4.NewSub(ctx)
	if err := sock.Listen(fmt.Sprintf("tcp://127.0.0.1:%d", port)); err != nil {
		cancel()
		_ = sock.Close()
		return nil, fmt.Errorf("listening for kv cache events: %w", err)
	}
	// Every topic: only this replica's engine is told the port.
	if err := sock.SetOption(zmq4.OptionSubscribe, ""); err != nil {
		cancel()
		_ = sock.Close()
		return nil, fmt.Errorf("subscribing to kv cache events: %w", err)
	}
	tcp, ok := sock.Addr().(*net.TCPAddr)
	if !ok {
		cancel()
		_ = sock.Close()
		return nil, fmt.Errorf("kv cache events: listening on %v, not a TCP address", sock.Addr())
	}
	s := &Subscriber{
		replica: replica, mirror: m, log: log, sock: sock, port: tcp.Port,
		cancel: cancel, done: make(chan struct{}), lastSeq: -1,
	}
	go s.run(ctx)
	return s, nil
}

// Port is the loopback port the engine is to publish to.
func (s *Subscriber) Port() int { return s.port }

// Endpoint is the address to give the engine.
func (s *Subscriber) Endpoint() string { return fmt.Sprintf("tcp://127.0.0.1:%d", s.port) }

// Close stops receiving. The replica stays in the mirror; Mirror.Remove is
// what says its engine has gone.
func (s *Subscriber) Close() {
	s.cancel()
	_ = s.sock.Close()
	<-s.done
}

func (s *Subscriber) run(ctx context.Context) {
	defer close(s.done)
	for {
		msg, err := s.sock.Recv()
		if err != nil {
			if ctx.Err() != nil {
				return
			}
			s.log.Debug("receiving kv cache events failed", "replica", s.replica, "err", err)
			continue
		}
		s.handle(msg.Frames)
	}
}

// handle reads one message: vLLM sends three frames, the topic, an eight-byte
// big-endian sequence number, and the msgpack batch.
func (s *Subscriber) handle(frames [][]byte) {
	if len(frames) != 3 || len(frames[1]) != 8 {
		s.log.Debug("ignoring a kv cache event message of unexpected shape",
			"replica", s.replica, "frames", len(frames))
		return
	}
	seq := int64(binary.BigEndian.Uint64(frames[1]))
	b, err := DecodeBatch(frames[2])

	s.mu.Lock()
	expected := s.lastSeq + 1
	gap := s.seen && seq != expected
	restarted := s.seen && seq < expected
	s.lastSeq, s.seen = seq, true
	s.mu.Unlock()

	switch {
	case restarted:
		// Numbering started again: a new engine process on the same port,
		// with a cache that starts empty.
		s.log.Info("kv cache events restarted; the engine's cache is new", "replica", s.replica)
		s.mirror.Reset(s.replica)
	case gap:
		// Events were lost, among them perhaps removals. An index still
		// holding blocks the engine evicted would send requests to a replica
		// for a cache it no longer has; starting again from nothing only
		// misses some it does. vLLM can replay from a buffer, which would
		// keep the index whole; not used yet.
		s.log.Warn("kv cache events were lost; starting the replica's index again",
			"replica", s.replica, "expected", expected, "got", seq)
		s.mirror.Reset(s.replica)
	}
	if err != nil {
		// A batch that cannot be read is a gap of one.
		s.log.Warn("kv cache event batch unreadable; starting the replica's index again",
			"replica", s.replica, "err", err)
		s.mirror.Reset(s.replica)
		return
	}
	s.mirror.Apply(s.replica, b)
}

// VLLMArgs are the arguments that make vLLM publish to endpoint, hashing the
// way the router computes.
//
// Prefix caching is turned on explicitly, though it is vLLM's default, because
// an index of a cache that is off would be an index of nothing.
func VLLMArgs(endpoint string) []string {
	cfg, _ := json.Marshal(map[string]any{
		"enable_kv_cache_events": true,
		"publisher":              "zmq",
		"endpoint":               endpoint,
		"topic":                  "kv-events",
	})
	return []string{
		"--enable-prefix-caching",
		"--prefix-caching-hash-algo", HashAlgo,
		"--kv-events-config", string(cfg),
	}
}

// VLLMEnv is the environment that goes with VLLMArgs: the seed every hash
// chain starts from, and hashes published as integers whatever the release's
// default.
func VLLMEnv() map[string]string {
	return map[string]string{
		"PYTHONHASHSEED":                      HashSeed,
		"VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES": "1",
	}
}