gsadmin/spinifex
1
0
Fork 0
mirror of https://github.com/mulgadc/spinifex.git synced 2026-10-15 17:11:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: convert the bedrock reaper to the shared kvlease package (#869)

Browse Source
This commit is contained in:
Josh Siviour
2026-08-25 08:58:02 +10:00
committed by GitHub
parent 5675a2b37c
commit 578566b801
2 changed files with 63 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
spinifex/handlers/bedrock/reaper.go
+32 -94
View File
@@ -5,10 +5,10 @@ import (
"errors"
"fmt"
"log/slog"
"sync"
"time"
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
"github.com/mulgadc/spinifex/spinifex/kvlease"
"github.com/mulgadc/spinifex/spinifex/otelsetup"
"github.com/mulgadc/spinifex/spinifex/utils"
"github.com/nats-io/nats.go/jetstream"
@@ -19,23 +19,9 @@ import (
const (
defaultReaperInterval = 30 * time.Second
defaultLeaseRefresh = 20 * time.Second
// How long an endpoint must have been idle before its GPU is taken back.
// Cold start was measured at 4m12s for the smallest model this platform
// serves, so a TTL of a few multiples of that is what keeps a quiet-but-used
// model from becoming a cold-start generator. Scale-to-zero is a cost
// optimisation; a latency-sensitive workload buys warmth by pinning.
defaultIdleTTL = 15 * time.Minute
// How many consecutive failed scrapes mean the serving process is wedged
// rather than briefly busy. Long enough that a GC pause or a dropped packet
// is not an outage, short enough that a dead VM does not hold a device for a
// whole idleTTL.
maxScrapeFailures = 5
// Per-scrape budget. Generous next to the sweep interval, because a slow
// answer is still an answer and a timeout here counts as a failure.
defaultScrapeTimeout = 5 * time.Second
defaultIdleTTL = 15 * time.Minute // How long an endpoint must have been idle before its GPU is taken back.
maxScrapeFailures = 5 // How many consecutive failed scrapes mean the serving process is wedged rather than briefly busy.
defaultScrapeTimeout = 5 * time.Second // Per-scrape budget.
)
// ReaperDeps are the reaper's overridable knobs. Every zero value takes the
@@ -50,23 +36,27 @@ type ReaperDeps struct {
// Reaper is the leader-elected idle-reclaim loop. One node holds the lease and
// does the scraping and reaping; every node keeps serving the API, so a
// leaderless gap delays a reclaim rather than failing a request.
//
// It reads idleness from the serving process's own Prometheus /metrics rather
// than from gateway-reported activity: that keeps the daemon the single writer
// of endpoint state, survives a gateway node dying mid-call, and works
// unchanged with any number of gateways.
type Reaper struct {
svc *Service
holder string
deps ReaperDeps
mu sync.Mutex
leader bool
svc *Service
holder string
deps ReaperDeps
lease *kvlease.Lease
leaseErr error
}
// NewReaper builds the reaper for svc; holder identifies this daemon in the lease.
func NewReaper(svc *Service, holder string, deps ReaperDeps) *Reaper {
return &Reaper{svc: svc, holder: holder, deps: deps}
r := &Reaper{svc: svc, holder: holder, deps: deps}
r.lease, r.leaseErr = kvlease.New(kvlease.Config{
Name: "bedrock/reaper",
Bucket: r.leaderBucket,
Key: leaderKey,
Holder: holder,
TTL: KVBucketLeaderTTL,
Renew: r.leaseRefresh(),
Retry: r.leaseRefresh(),
})
return r
}
func (r *Reaper) interval() time.Duration {
@@ -100,19 +90,24 @@ func (r *Reaper) scrapeTimeout() time.Duration {
// Run drives the leadership and sweep loop until ctx is cancelled. Intended as
// a daemon-boot goroutine; panics are the caller's recover concern.
func (r *Reaper) Run(ctx context.Context) {
leaseTicker := time.NewTicker(r.leaseRefresh())
if r.leaseErr != nil {
slog.ErrorContext(ctx, "bedrock reaper: lease config invalid", "holder", r.holder, "err", r.leaseErr)
return
}
sweepTicker := time.NewTicker(r.interval())
defer leaseTicker.Stop()
defer sweepTicker.Stop()
r.evaluateLeadership(ctx)
leaseDone := make(chan struct{})
go func() {
defer close(leaseDone)
r.lease.Run(ctx)
}()
for {
select {
case <-ctx.Done():
r.relinquish()
<-leaseDone
return
case <-leaseTicker.C:
r.evaluateLeadership(ctx)
case <-sweepTicker.C:
if !r.IsLeader() {
continue
@@ -126,64 +121,7 @@ func (r *Reaper) Run(ctx context.Context) {
// IsLeader reports whether this node currently holds the reaper lease.
func (r *Reaper) IsLeader() bool {
r.mu.Lock()
defer r.mu.Unlock()
return r.leader
}
func (r *Reaper) evaluateLeadership(ctx context.Context) {
won := r.acquireOrRefresh(ctx)
r.mu.Lock()
was := r.leader
r.leader = won
r.mu.Unlock()
switch {
case won && !was:
slog.Info("bedrock reaper: elected leader", "holder", r.holder)
case !won && was:
slog.Info("bedrock reaper: lost leadership", "holder", r.holder)
}
}
// acquireOrRefresh claims the lease, or refreshes it (resetting the TTL) when
// this node already holds it.
func (r *Reaper) acquireOrRefresh(ctx context.Context) bool {
kv, err := r.leaderBucket(ctx)
if err != nil {
return false
}
if _, err := kv.Create(ctx, leaderKey, []byte(r.holder)); err == nil {
return true
}
entry, err := kv.Get(ctx, leaderKey)
if err != nil {
return false
}
if string(entry.Value()) != r.holder {
return false
}
if _, err := kv.Put(ctx, leaderKey, []byte(r.holder)); err != nil {
return false
}
return true
}
// relinquish releases the lease on shutdown so the next leader is elected
// immediately rather than after the TTL.
func (r *Reaper) relinquish() {
// Run's ctx is already cancelled by the time this is called, so the release
// runs on its own — a cancelled ctx would fail the delete.
ctx := context.Background()
kv, err := r.leaderBucket(ctx)
if err != nil {
return
}
if entry, gerr := kv.Get(ctx, leaderKey); gerr == nil && string(entry.Value()) == r.holder {
if err := kv.Delete(ctx, leaderKey); err != nil {
slog.Debug("bedrock reaper: release lease failed", "holder", r.holder, "err", err)
}
}
return r.lease.Held()
}
func (r *Reaper) leaderBucket(ctx context.Context) (jetstream.KeyValue, error) {
spinifex/handlers/bedrock/reaper_test.go
+31 -10
View File
@@ -395,40 +395,41 @@ func TestReaper_LeaseAdmitsOneLeader(t *testing.T) {
f := newReaperFixture(t, ReaperDeps{})
other := NewReaper(f.svc, "node-b", ReaperDeps{})
f.reaper.evaluateLeadership(t.Context())
other.evaluateLeadership(t.Context())
require.True(t, f.reaper.lease.TryAcquire(t.Context()))
require.False(t, other.lease.TryAcquire(t.Context()),
"two nodes must never sweep the same endpoints at once")
assert.True(t, f.reaper.IsLeader())
assert.False(t, other.IsLeader(), "two nodes must never sweep the same endpoints at once")
assert.False(t, other.IsLeader())
// A refresh by the holder keeps it, and does not hand it away.
f.reaper.evaluateLeadership(t.Context())
require.True(t, f.reaper.lease.TryAcquire(t.Context()))
assert.True(t, f.reaper.IsLeader())
}
// Releasing on shutdown is what makes a rolling restart cost seconds rather
// than a full lease TTL of no reclaim.
func TestReaper_RelinquishFreesTheLeaseImmediately(t *testing.T) {
func TestReaper_ReleaseFreesTheLeaseImmediately(t *testing.T) {
f := newReaperFixture(t, ReaperDeps{})
other := NewReaper(f.svc, "node-b", ReaperDeps{})
f.reaper.evaluateLeadership(t.Context())
f.reaper.lease.TryAcquire(t.Context())
require.True(t, f.reaper.IsLeader())
f.reaper.relinquish()
other.evaluateLeadership(t.Context())
f.reaper.lease.Release(t.Context())
other.lease.TryAcquire(t.Context())
assert.True(t, other.IsLeader())
}
// A node that lost the election does no work at all, so a leaderless gap
// delays a reclaim rather than duplicating one.
func TestReaper_NonLeaderSweepsNothing(t *testing.T) {
f := newReaperFixture(t, ReaperDeps{Interval: 5 * time.Millisecond, LeaseRefresh: time.Hour, IdleTTL: time.Nanosecond})
f := newReaperFixture(t, ReaperDeps{Interval: 5 * time.Millisecond, IdleTTL: time.Nanosecond})
f.ready(t)
f.age(t, time.Hour, time.Hour)
holder := NewReaper(f.svc, "node-b", ReaperDeps{})
holder.evaluateLeadership(t.Context())
holder.lease.TryAcquire(t.Context())
require.True(t, holder.IsLeader())
ctx, cancel := context.WithTimeout(t.Context(), 60*time.Millisecond)
@@ -450,3 +451,23 @@ func TestEndpointRecord_LastActive(t *testing.T) {
assert.Equal(t, ready, EndpointRecord{ReadyAt: ready}.LastActive(),
"an endpoint quiet since launch is idle since launch, not since the zero time")
}
// Shutdown must delete the lease key before Run returns. The daemon waits on Run
// via its shutdown group, so a key left behind parks the next election for the
// full TTL rather than freeing it immediately.
func TestReaper_RunReleasesLeaseOnShutdown(t *testing.T) {
f := newReaperFixture(t, ReaperDeps{Interval: time.Hour})
require.NoError(t, f.reaper.leaseErr)
ctx, cancel := context.WithCancel(t.Context())
done := make(chan struct{})
go func() { defer close(done); f.reaper.Run(ctx) }()
require.Eventually(t, f.reaper.IsLeader, 2*time.Second, 20*time.Millisecond)
cancel()
<-done
other := NewReaper(f.svc, "node-b", ReaperDeps{})
assert.True(t, other.lease.TryAcquire(t.Context()),
"Run returned with the lease key still present")
}