mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-15 17:11:30 +00:00
feat(bedrock): evict the LRU idle endpoint when a launch has no GPU (#745)
A refused launch and an idle GPU could coexist indefinitely: capacity said no while a device sat under a model nobody was calling. Ensure now makes room once before giving up, choosing the least recently used endpoint that is safe to stop — READY, nothing in flight, not pinned, past minResidency, and on this node, since GPU capacity is node-local and stopping an endpoint elsewhere frees nothing here. The eviction goes through Delete, so the device is released by the same DRAINING -> terminate -> release path an operator delete takes rather than being taken from underneath a running process. Nothing evictable returns the original ModelNotReadyException unchanged: capacity contention breaches no quota and must not claim to. minResidency defaults to 5m so two models cannot thrash each other with a full cold boot per call. Beyond that the deployment is under-provisioned, and the eviction rate is what says so. Also adds coldStartWait (OCHRE_COLD_START_WAIT), default 0. At zero the fail-fast contract is unchanged and costs no extra describe. Above zero the resolver polls inside its existing singleflight, so a burst of cold callers shares one wait and one launch. The default stays 0 because the measurement behind it has not changed: cold start is minutes, and no client retry policy spans that.
This commit is contained in:
@@ -2,6 +2,7 @@ package handlers_bedrock
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/mulgadc/spinifex/spinifex/awserrors"
|
||||
"github.com/mulgadc/spinifex/spinifex/gpu"
|
||||
@@ -55,6 +56,38 @@ func checkCapacity(snapshot []gpu.PoolEntry, minVRAMMiB int) error {
|
||||
"bedrock: no free GPU device has %d MiB of VRAM free", minVRAMMiB)
|
||||
}
|
||||
|
||||
// defaultMinResidency is how long an endpoint is protected from eviction after
|
||||
// reaching READY. Without it two models on one device thrash each other with a
|
||||
// full cold boot per call, which serves neither. Beyond that the deployment is
|
||||
// simply under-provisioned, and the eviction rate is what says so.
|
||||
const defaultMinResidency = 5 * time.Minute
|
||||
|
||||
// selectEvictable picks the endpoint whose GPU should be taken for a pending
|
||||
// launch: the least recently used of those safe to stop. Reports false when
|
||||
// nothing qualifies, which the caller turns back into the original capacity
|
||||
// refusal rather than an eviction.
|
||||
//
|
||||
// Evictable is deliberately narrow. A device is never yanked from a running
|
||||
// process (InFlight), a pinned endpoint is never touched, and an endpoint that
|
||||
// has not yet served out minResidency is left alone. nodeID confines the
|
||||
// choice to this node: GPU capacity is node-local, so stopping an endpoint
|
||||
// running elsewhere frees nothing here.
|
||||
func selectEvictable(recs []EndpointRecord, nodeID string, minResidency time.Duration, now time.Time) (EndpointRecord, bool) {
|
||||
var victim EndpointRecord
|
||||
var found bool
|
||||
for _, rec := range recs {
|
||||
switch {
|
||||
case rec.State != StateReady, rec.Pinned, rec.InFlight > 0,
|
||||
rec.NodeID != nodeID, now.Sub(rec.ReadyAt) < minResidency:
|
||||
continue
|
||||
}
|
||||
if !found || lastActive(rec).Before(lastActive(victim)) {
|
||||
victim, found = rec, true
|
||||
}
|
||||
}
|
||||
return victim, found
|
||||
}
|
||||
|
||||
// admitCapacity checks minVRAMMiB against snapshotter's current free pool.
|
||||
// A nil snapshotter (no GPU manager: passthrough disabled or no GPUs present)
|
||||
// admits nothing, since every self-host model in the catalog requires a GPU.
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
package handlers_bedrock
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/mulgadc/spinifex/spinifex/awserrors"
|
||||
"github.com/mulgadc/spinifex/spinifex/gpu"
|
||||
"github.com/mulgadc/spinifex/spinifex/utils"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// evictableRecord is a READY endpoint on this node, old enough to be past
|
||||
// minResidency and quiet — the baseline each case below spoils one field of.
|
||||
func evictableRecord(modelID string, lastActive time.Time) EndpointRecord {
|
||||
return EndpointRecord{
|
||||
ModelID: modelID,
|
||||
State: StateReady,
|
||||
NodeID: testNodeID,
|
||||
InstanceID: "i-" + modelID,
|
||||
ReadyAt: lastActive.Add(-time.Hour),
|
||||
LastActiveAt: lastActive,
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectEvictable_PicksLeastRecentlyUsed(t *testing.T) {
|
||||
now := time.Now().UTC()
|
||||
recs := []EndpointRecord{
|
||||
evictableRecord("recent", now.Add(-time.Minute)),
|
||||
evictableRecord("stalest", now.Add(-30*time.Minute)),
|
||||
evictableRecord("middling", now.Add(-10*time.Minute)),
|
||||
}
|
||||
|
||||
victim, found := selectEvictable(recs, testNodeID, time.Minute, now)
|
||||
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "stalest", victim.ModelID)
|
||||
}
|
||||
|
||||
// A record the reaper has never sampled falls back to ReadyAt, so a brand-new
|
||||
// endpoint sorts last rather than first — the opposite would evict the one
|
||||
// thing that just launched.
|
||||
func TestSelectEvictable_UnsampledRecordFallsBackToReadyAt(t *testing.T) {
|
||||
now := time.Now().UTC()
|
||||
fresh := EndpointRecord{ModelID: "fresh", State: StateReady, NodeID: testNodeID, ReadyAt: now.Add(-10 * time.Minute)}
|
||||
old := EndpointRecord{ModelID: "old", State: StateReady, NodeID: testNodeID, ReadyAt: now.Add(-3 * time.Hour)}
|
||||
|
||||
victim, found := selectEvictable([]EndpointRecord{fresh, old}, testNodeID, time.Minute, now)
|
||||
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "old", victim.ModelID)
|
||||
}
|
||||
|
||||
func TestSelectEvictable_Exclusions(t *testing.T) {
|
||||
now := time.Now().UTC()
|
||||
stale := now.Add(-time.Hour)
|
||||
|
||||
busy := evictableRecord("busy", stale)
|
||||
busy.InFlight = 1
|
||||
|
||||
pinned := evictableRecord("pinned", stale)
|
||||
pinned.Pinned = true
|
||||
|
||||
remote := evictableRecord("remote", stale)
|
||||
remote.NodeID = "node-elsewhere"
|
||||
|
||||
starting := evictableRecord("starting", stale)
|
||||
starting.State = StateStarting
|
||||
|
||||
fresh := evictableRecord("fresh", now)
|
||||
fresh.ReadyAt = now.Add(-time.Second)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
rec EndpointRecord
|
||||
}{
|
||||
{"a device is never yanked from a running process", busy},
|
||||
{"a pinned endpoint is never evicted", pinned},
|
||||
{"another node's endpoint frees no capacity here", remote},
|
||||
{"a launch in flight is not ours to stop", starting},
|
||||
{"minResidency stops two models thrashing each other", fresh},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
_, found := selectEvictable([]EndpointRecord{tt.rec}, testNodeID, 5*time.Minute, now)
|
||||
assert.False(t, found)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// exhaustedGPU reports a device that is present but taken, which is what a
|
||||
// node running one serving VM on its only GPU looks like.
|
||||
func exhaustedGPU() *stubSnapshotter {
|
||||
return &stubSnapshotter{entries: []gpu.PoolEntry{{Device: gpu.GPUDevice{MemoryMiB: 8192}, Available: false}}}
|
||||
}
|
||||
|
||||
// freeingGPU stays exhausted until the endpoint holding the device is torn
|
||||
// down, standing in for a real pool whose entry is released by the terminate.
|
||||
type freeingGPU struct {
|
||||
launcher *fakeInstanceLauncher
|
||||
}
|
||||
|
||||
func (f *freeingGPU) Snapshot() []gpu.PoolEntry {
|
||||
available := len(f.launcher.terminations()) > 0
|
||||
return []gpu.PoolEntry{{Device: gpu.GPUDevice{MemoryMiB: 8192}, Available: available}}
|
||||
}
|
||||
|
||||
// The case that is impossible without eviction: the only GPU is held by a
|
||||
// model nobody is calling, and every launch that needs it is refused forever.
|
||||
func TestEnsure_EvictsIdleEndpointToMakeRoom(t *testing.T) {
|
||||
h := newLaunchHarness()
|
||||
s, _ := newTestService(t, h, 200, &freeingGPU{launcher: h.launcher})
|
||||
seedReadyEndpoint(t, s, "other.model-v1:0", time.Now().UTC().Add(-time.Hour))
|
||||
|
||||
out, err := s.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, StateStarting, out.Endpoint.State)
|
||||
s.WaitLaunches()
|
||||
|
||||
assert.Contains(t, h.launcher.terminations(), "i-other", "the victim's VM must be torn down to release its device")
|
||||
desc, err := s.Describe(t.Context(), &DescribeEndpointInput{ModelID: "other.model-v1:0"}, "")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, StateAbsent, desc.Endpoint.State, "an evicted endpoint's record must be purged")
|
||||
|
||||
desc, err = s.Describe(t.Context(), &DescribeEndpointInput{ModelID: testModelID}, "")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, StateReady, desc.Endpoint.State, "the launch the eviction made room for must complete")
|
||||
}
|
||||
|
||||
// Capacity contention breaches no quota, so the refusal that comes back is
|
||||
// the retryable one, never ServiceQuotaExceededException.
|
||||
func TestEnsure_NothingEvictableKeepsModelNotReady(t *testing.T) {
|
||||
h := newLaunchHarness()
|
||||
s, _ := newTestService(t, h, 200, exhaustedGPU())
|
||||
// The only other endpoint is pinned, so nothing may be given up.
|
||||
rec := seedReadyEndpoint(t, s, "other.model-v1:0", time.Now().UTC().Add(-time.Hour))
|
||||
pinEndpoint(t, s, rec.ModelID)
|
||||
|
||||
_, err := s.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "")
|
||||
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorModelNotReadyException, awserrors.ValidErrorCodeFromError(err))
|
||||
assert.NotEqual(t, awserrors.ErrorServiceQuotaExceededException, awserrors.ValidErrorCodeFromError(err))
|
||||
assert.Empty(t, h.launcher.terminations(), "nothing evictable means nothing is torn down")
|
||||
assert.Zero(t, h.launcher.launchCount.Load())
|
||||
}
|
||||
|
||||
// An endpoint still inside minResidency is not a candidate, so a refusal
|
||||
// stands rather than two models trading the device on every call.
|
||||
func TestEnsure_MinResidencyProtectsAFreshEndpoint(t *testing.T) {
|
||||
h := newLaunchHarness()
|
||||
s, _ := newTestService(t, h, 200, exhaustedGPU())
|
||||
seedReadyEndpoint(t, s, "other.model-v1:0", time.Now().UTC())
|
||||
|
||||
_, err := s.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "")
|
||||
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorModelNotReadyException, awserrors.ValidErrorCodeFromError(err))
|
||||
assert.Empty(t, h.launcher.terminations())
|
||||
}
|
||||
|
||||
// The requester's own endpoint is never the victim: stopping it to launch it
|
||||
// again is a no-op at best.
|
||||
func TestEvictForCapacity_NeverEvictsTheRequester(t *testing.T) {
|
||||
h := newLaunchHarness()
|
||||
s, _ := newTestService(t, h, 200, exhaustedGPU())
|
||||
seedReadyEndpoint(t, s, testModelID, time.Now().UTC().Add(-time.Hour))
|
||||
|
||||
kv, err := s.bucket(t.Context())
|
||||
require.NoError(t, err)
|
||||
assert.False(t, s.evictForCapacity(t.Context(), kv, testModelID, 5120))
|
||||
assert.Empty(t, h.launcher.terminations())
|
||||
}
|
||||
|
||||
// seedReadyEndpoint writes a READY record directly, which is the only way to
|
||||
// stand up a second serving endpoint: the catalog carries one self-host model
|
||||
// today, so a launched one cannot be joined by a second through Ensure.
|
||||
func seedReadyEndpoint(t *testing.T, s *Service, modelID string, readyAt time.Time) EndpointRecord {
|
||||
t.Helper()
|
||||
kv, err := s.bucket(t.Context())
|
||||
require.NoError(t, err)
|
||||
rec := EndpointRecord{
|
||||
AccountID: utils.GlobalAccountID,
|
||||
ModelID: modelID,
|
||||
State: StateReady,
|
||||
NodeID: s.deps.NodeID,
|
||||
InstanceID: "i-other",
|
||||
BaseURL: "http://10.244.1.9:8000",
|
||||
ReadyAt: readyAt,
|
||||
Generation: 2,
|
||||
}
|
||||
_, err = createJSONRevision(t.Context(), kv, EndpointKey(utils.GlobalAccountID, modelID), rec)
|
||||
require.NoError(t, err)
|
||||
return rec
|
||||
}
|
||||
|
||||
func pinEndpoint(t *testing.T, s *Service, modelID string) {
|
||||
t.Helper()
|
||||
kv, err := s.bucket(t.Context())
|
||||
require.NoError(t, err)
|
||||
key := EndpointKey(utils.GlobalAccountID, modelID)
|
||||
rec, rev, found, err := getFullJSON(t.Context(), kv, key)
|
||||
require.NoError(t, err)
|
||||
require.True(t, found)
|
||||
rec.Pinned = true
|
||||
require.NoError(t, updateJSON(t.Context(), kv, key, rev, rec))
|
||||
}
|
||||
@@ -329,6 +329,14 @@ func (f *fakeInstanceLauncher) launches() []*sysinstance.SystemInstanceInput {
|
||||
return f.requests
|
||||
}
|
||||
|
||||
// terminations is the mutex-guarded read of terminated, for assertions made
|
||||
// while a launch goroutine may still be running.
|
||||
func (f *fakeInstanceLauncher) terminations() []string {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return slices.Clone(f.terminated)
|
||||
}
|
||||
|
||||
func (f *fakeInstanceLauncher) TerminateSystemInstance(instanceID string) error {
|
||||
f.mu.Lock()
|
||||
f.terminated = append(f.terminated, instanceID)
|
||||
|
||||
@@ -16,6 +16,11 @@ import (
|
||||
// an address that no longer answers.
|
||||
const defaultEndpointCacheTTL = 5 * time.Second
|
||||
|
||||
// coldStartPollInterval spaces the describes a non-zero coldStartWait makes.
|
||||
// Fine-grained next to any plausible wait, since the point of waiting at all
|
||||
// is to return the moment the endpoint answers.
|
||||
const coldStartPollInterval = 500 * time.Millisecond
|
||||
|
||||
// DynamicEndpointResolver resolves a self-host model's base URL through the
|
||||
// daemon's endpoint registry, and asks for a launch when there is nothing to
|
||||
// resolve. It lives here rather than in gateway_bedrock because that package
|
||||
@@ -26,6 +31,13 @@ type DynamicEndpointResolver struct {
|
||||
static map[string]string
|
||||
ttl time.Duration
|
||||
|
||||
// coldStartWait is how long a cold call may be held waiting for the launch
|
||||
// it triggered. Zero — the default — returns immediately, which is the
|
||||
// measured position: cold start is minutes, the AWS SDKs' default ~3-attempt
|
||||
// retry does not span that, and a bounded wait only pays off when it is
|
||||
// shorter than the client's patience and longer than the boot.
|
||||
coldStartWait time.Duration
|
||||
|
||||
// group collapses concurrent resolves of one model into a single describe
|
||||
// (and at most one ensure). The daemon deduplicates launches on its own
|
||||
// through the KV claim; this only stops a burst spending one NATS round
|
||||
@@ -44,19 +56,35 @@ type cachedEndpoint struct {
|
||||
|
||||
var _ gateway_bedrock.EndpointResolver = (*DynamicEndpointResolver)(nil)
|
||||
|
||||
// ResolverOption adjusts a DynamicEndpointResolver at construction.
|
||||
type ResolverOption func(*DynamicEndpointResolver)
|
||||
|
||||
// WithColdStartWait holds a cold call for up to d waiting on the launch it
|
||||
// triggered, instead of returning ModelNotReadyException straight away. The
|
||||
// deployment-level escape hatch for a model known to start fast enough that
|
||||
// waiting beats retrying; zero (the default) keeps the fail-fast contract.
|
||||
func WithColdStartWait(d time.Duration) ResolverOption {
|
||||
return func(r *DynamicEndpointResolver) { r.coldStartWait = d }
|
||||
}
|
||||
|
||||
// NewDynamicEndpointResolver builds a resolver over svc. Entries in static
|
||||
// (OCHRE_VLLM_ENDPOINTS) are resolved first and never reach svc, so a pinned
|
||||
// endpoint bypasses the lifecycle entirely. A zero ttl takes the default.
|
||||
func NewDynamicEndpointResolver(svc EndpointService, static map[string]string, ttl time.Duration) *DynamicEndpointResolver {
|
||||
func NewDynamicEndpointResolver(svc EndpointService, static map[string]string, ttl time.Duration,
|
||||
opts ...ResolverOption) *DynamicEndpointResolver {
|
||||
if ttl <= 0 {
|
||||
ttl = defaultEndpointCacheTTL
|
||||
}
|
||||
return &DynamicEndpointResolver{
|
||||
r := &DynamicEndpointResolver{
|
||||
svc: svc,
|
||||
static: static,
|
||||
ttl: ttl,
|
||||
cached: make(map[string]cachedEndpoint),
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(r)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// Endpoint returns modelID's base URL if one is serving. A model with no
|
||||
@@ -97,16 +125,55 @@ func (r *DynamicEndpointResolver) resolve(ctx context.Context, modelID string) (
|
||||
}
|
||||
r.storeCache(modelID, out.Endpoint.BaseURL)
|
||||
return out.Endpoint.BaseURL, nil
|
||||
case StateStarting, StateDraining:
|
||||
// A launch is already in flight, or a teardown is. Either way asking
|
||||
// again changes nothing and the answer is the same: not yet.
|
||||
case StateStarting:
|
||||
// A launch is already in flight, so asking for another changes nothing.
|
||||
// Waiting on it is still worth doing when the deployment asked for that.
|
||||
return r.awaitReady(ctx, modelID)
|
||||
case StateDraining:
|
||||
// A teardown is in flight and nothing will relaunch on its own, so
|
||||
// there is no readiness to wait for.
|
||||
return "", nil
|
||||
}
|
||||
|
||||
if _, err := r.svc.Ensure(ctx, &EnsureEndpointInput{ModelID: modelID}, utils.GlobalAccountID); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "", nil
|
||||
return r.awaitReady(ctx, modelID)
|
||||
}
|
||||
|
||||
// awaitReady holds the caller for up to coldStartWait waiting on a launch,
|
||||
// and returns "" (not ready) the moment the budget runs out. At the default
|
||||
// of zero it does nothing at all, so the fail-fast path costs no extra
|
||||
// describe. It runs inside the singleflight, so a burst of concurrent callers
|
||||
// for one cold model shares a single wait rather than one describe loop each.
|
||||
func (r *DynamicEndpointResolver) awaitReady(ctx context.Context, modelID string) (string, error) {
|
||||
if r.coldStartWait <= 0 {
|
||||
return "", nil
|
||||
}
|
||||
deadline := time.Now().Add(r.coldStartWait)
|
||||
ticker := time.NewTicker(coldStartPollInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
// The caller is gone. The launch it triggered carries on regardless,
|
||||
// so this is not-ready rather than an error.
|
||||
return "", nil
|
||||
case <-ticker.C:
|
||||
}
|
||||
out, err := r.svc.Describe(ctx, &DescribeEndpointInput{ModelID: modelID}, utils.GlobalAccountID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if out.Endpoint.State == StateReady && out.Endpoint.BaseURL != "" {
|
||||
r.storeCache(modelID, out.Endpoint.BaseURL)
|
||||
return out.Endpoint.BaseURL, nil
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
return "", nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *DynamicEndpointResolver) lookupCache(modelID string) (string, bool) {
|
||||
|
||||
@@ -239,3 +239,108 @@ func TestDynamicEndpointResolver_ConcurrentColdCallsCollapse(t *testing.T) {
|
||||
assert.Equal(t, int64(1), svc.describeCalls.Load())
|
||||
assert.Equal(t, int64(1), svc.ensureCalls.Load())
|
||||
}
|
||||
|
||||
// The default is the whole cold-start contract: no wait, no extra describe,
|
||||
// a retryable answer straight back to the caller.
|
||||
func TestDynamicEndpointResolver_ColdStartWaitDefaultsToNoWait(t *testing.T) {
|
||||
svc := &fakeEndpointService{record: EndpointRecord{ModelID: resolverTestModel, State: StateAbsent}}
|
||||
r := NewDynamicEndpointResolver(svc, nil, 0)
|
||||
|
||||
started := time.Now()
|
||||
_, ok, err := r.Endpoint(context.Background(), resolverTestModel)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.False(t, ok)
|
||||
assert.Less(t, time.Since(started), 100*time.Millisecond, "a cold call must not be held")
|
||||
assert.Equal(t, int64(1), svc.describeCalls.Load(), "no polling describe at the default")
|
||||
}
|
||||
|
||||
// With a wait configured, a launch that lands inside the budget is served on
|
||||
// the original call rather than a retry.
|
||||
func TestDynamicEndpointResolver_ColdStartWaitResolvesWhenReady(t *testing.T) {
|
||||
svc := &fakeEndpointService{record: EndpointRecord{ModelID: resolverTestModel, State: StateAbsent}}
|
||||
r := NewDynamicEndpointResolver(svc, nil, 0, WithColdStartWait(5*time.Second))
|
||||
|
||||
go func() {
|
||||
time.Sleep(600 * time.Millisecond)
|
||||
svc.setRecord(readyRecord("http://10.0.0.20:8000"))
|
||||
}()
|
||||
|
||||
baseURL, ok, err := r.Endpoint(context.Background(), resolverTestModel)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "http://10.0.0.20:8000", baseURL)
|
||||
}
|
||||
|
||||
// A launch slower than the budget still gets the retryable answer, and the
|
||||
// budget is what bounds the hold.
|
||||
func TestDynamicEndpointResolver_ColdStartWaitGivesUpAtTheBudget(t *testing.T) {
|
||||
svc := &fakeEndpointService{record: EndpointRecord{ModelID: resolverTestModel, State: StateAbsent}}
|
||||
r := NewDynamicEndpointResolver(svc, nil, 0, WithColdStartWait(600*time.Millisecond))
|
||||
|
||||
started := time.Now()
|
||||
_, ok, err := r.Endpoint(context.Background(), resolverTestModel)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.False(t, ok)
|
||||
assert.GreaterOrEqual(t, time.Since(started), 500*time.Millisecond)
|
||||
assert.Less(t, time.Since(started), 3*time.Second)
|
||||
}
|
||||
|
||||
// A client that disconnects mid-wait releases the goroutine without turning
|
||||
// its own departure into a server error; the launch it triggered carries on.
|
||||
func TestDynamicEndpointResolver_ColdStartWaitStopsOnClientCancel(t *testing.T) {
|
||||
svc := &fakeEndpointService{record: EndpointRecord{ModelID: resolverTestModel, State: StateAbsent}}
|
||||
r := NewDynamicEndpointResolver(svc, nil, 0, WithColdStartWait(time.Minute))
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
started := time.Now()
|
||||
_, ok, err := r.Endpoint(ctx, resolverTestModel)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.False(t, ok)
|
||||
assert.Less(t, time.Since(started), 5*time.Second)
|
||||
assert.Equal(t, int64(1), svc.ensureCalls.Load())
|
||||
}
|
||||
|
||||
// A burst of cold callers shares one wait, not one describe loop each.
|
||||
func TestDynamicEndpointResolver_ColdStartWaitIsSharedAcrossCallers(t *testing.T) {
|
||||
svc := &fakeEndpointService{record: EndpointRecord{ModelID: resolverTestModel, State: StateAbsent}}
|
||||
r := NewDynamicEndpointResolver(svc, nil, 0, WithColdStartWait(3*time.Second))
|
||||
|
||||
go func() {
|
||||
time.Sleep(600 * time.Millisecond)
|
||||
svc.setRecord(readyRecord("http://10.0.0.30:8000"))
|
||||
}()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for range 8 {
|
||||
wg.Go(func() {
|
||||
baseURL, ok, err := r.Endpoint(context.Background(), resolverTestModel)
|
||||
assert.NoError(t, err)
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "http://10.0.0.30:8000", baseURL)
|
||||
})
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
assert.Equal(t, int64(1), svc.ensureCalls.Load(), "one launch, however many callers waited on it")
|
||||
assert.Less(t, svc.describeCalls.Load(), int64(8), "the poll is shared, not repeated per caller")
|
||||
}
|
||||
|
||||
// DRAINING has no launch to wait for: nothing relaunches on its own, so the
|
||||
// answer is immediate even with a wait configured.
|
||||
func TestDynamicEndpointResolver_ColdStartWaitSkipsDraining(t *testing.T) {
|
||||
svc := &fakeEndpointService{record: EndpointRecord{ModelID: resolverTestModel, State: StateDraining}}
|
||||
r := NewDynamicEndpointResolver(svc, nil, 0, WithColdStartWait(time.Minute))
|
||||
|
||||
started := time.Now()
|
||||
_, ok, err := r.Endpoint(context.Background(), resolverTestModel)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.False(t, ok)
|
||||
assert.Less(t, time.Since(started), time.Second)
|
||||
}
|
||||
|
||||
@@ -47,6 +47,9 @@ type ServiceDeps struct {
|
||||
PollInterval time.Duration
|
||||
// Replicas sizes the bedrock-endpoints KV bucket on first create.
|
||||
Replicas int
|
||||
// MinResidency protects a freshly READY endpoint from eviction; zero takes
|
||||
// defaultMinResidency.
|
||||
MinResidency time.Duration
|
||||
}
|
||||
|
||||
// Service is the daemon-side, KV-backed handler set for the serving-endpoint
|
||||
@@ -110,6 +113,13 @@ func (s *Service) pollInterval() time.Duration {
|
||||
return readinessPollInterval
|
||||
}
|
||||
|
||||
func (s *Service) minResidency() time.Duration {
|
||||
if s.deps.MinResidency > 0 {
|
||||
return s.deps.MinResidency
|
||||
}
|
||||
return defaultMinResidency
|
||||
}
|
||||
|
||||
func (s *Service) httpClient() *http.Client {
|
||||
if s.deps.HTTPClient != nil {
|
||||
return s.deps.HTTPClient
|
||||
@@ -161,7 +171,15 @@ func (s *Service) Ensure(ctx context.Context, in *EnsureEndpointInput, _ string)
|
||||
}
|
||||
|
||||
if err := admitCapacity(s.deps.GPU, spec.MinVRAMMiB); err != nil {
|
||||
return nil, err
|
||||
// No free device. A GPU held by a model nobody is calling is not a
|
||||
// reason to refuse this one, so make room and re-check — but return the
|
||||
// original refusal, unchanged, when nothing can be given up.
|
||||
if !s.evictForCapacity(ctx, kv, in.ModelID, spec.MinVRAMMiB) {
|
||||
return nil, err
|
||||
}
|
||||
if retryErr := admitCapacity(s.deps.GPU, spec.MinVRAMMiB); retryErr != nil {
|
||||
return nil, retryErr
|
||||
}
|
||||
}
|
||||
if err := validateTransition(StateAbsent, StateStarting); err != nil {
|
||||
return nil, err
|
||||
@@ -205,6 +223,49 @@ func (s *Service) Ensure(ctx context.Context, in *EnsureEndpointInput, _ string)
|
||||
return &EnsureEndpointOutput{Endpoint: rec}, nil
|
||||
}
|
||||
|
||||
// evictForCapacity stops the least recently used evictable endpoint so a
|
||||
// pending launch can have its device, reporting whether one was stopped. The
|
||||
// eviction goes through Delete, so the device is released by the same
|
||||
// DRAINING -> terminate -> release path an operator delete takes rather than
|
||||
// being taken from underneath a running process.
|
||||
//
|
||||
// wantModelID is excluded defensively: a caller that reached the capacity
|
||||
// check has no endpoint of its own to evict, and stopping one to launch it
|
||||
// again would be a no-op at best.
|
||||
//
|
||||
// Delete takes the per-key mutex for the victim while Ensure holds it for
|
||||
// wantModelID. The two can never be the same key, so the locks are disjoint.
|
||||
func (s *Service) evictForCapacity(ctx context.Context, kv jetstream.KeyValue, wantModelID string, minVRAMMiB int) bool {
|
||||
recs, err := ListEndpoints(ctx, kv, utils.GlobalAccountID)
|
||||
if err != nil {
|
||||
slog.ErrorContext(ctx, "bedrock: list endpoints for eviction failed", "model", wantModelID, "err", err)
|
||||
return false
|
||||
}
|
||||
candidates := make([]EndpointRecord, 0, len(recs))
|
||||
for _, rec := range recs {
|
||||
if rec.ModelID != wantModelID {
|
||||
candidates = append(candidates, rec)
|
||||
}
|
||||
}
|
||||
|
||||
victim, found := selectEvictable(candidates, s.deps.NodeID, s.minResidency(), time.Now().UTC())
|
||||
if !found {
|
||||
slog.InfoContext(ctx, "bedrock: no free GPU and nothing evictable",
|
||||
"model", wantModelID, "minVRAMMiB", minVRAMMiB, "endpoints", len(recs))
|
||||
return false
|
||||
}
|
||||
|
||||
slog.InfoContext(ctx, "bedrock: evicting an idle endpoint to make room",
|
||||
"evicting", victim.ModelID, "instanceId", victim.InstanceID, "for", wantModelID,
|
||||
"idleSince", lastActive(victim))
|
||||
if _, err := s.Delete(ctx, &DeleteEndpointInput{ModelID: victim.ModelID}, utils.GlobalAccountID); err != nil {
|
||||
slog.ErrorContext(ctx, "bedrock: eviction failed; the pending launch keeps its capacity refusal",
|
||||
"evicting", victim.ModelID, "for", wantModelID, "err", err)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// runLaunch performs the slow launch + readiness probe and writes the
|
||||
// terminal state. Only this goroutine writes key while the record is
|
||||
// STARTING (guaranteed by Ensure returning early for any concurrent caller
|
||||
|
||||
@@ -331,9 +331,13 @@ func launchService(config *config.ClusterConfig) error {
|
||||
// Bedrock self-host endpoints: OCHRE_VLLM_ENDPOINTS pins a modelId=baseURL
|
||||
// pair to a fixed address, and anything not pinned resolves through the
|
||||
// daemon's endpoint registry, which launches a serving VM on first use.
|
||||
// OCHRE_COLD_START_WAIT optionally holds a cold call that long rather than
|
||||
// returning ModelNotReadyException at once. Unset (the default) keeps the
|
||||
// fail-fast contract: cold start is minutes, which no client retry spans.
|
||||
bedrockEndpoints := parseBedrockEndpoints(os.Getenv("OCHRE_VLLM_ENDPOINTS"))
|
||||
bedrockEndpointResolver := handlers_bedrock.NewDynamicEndpointResolver(
|
||||
handlers_bedrock.NewNATSEndpointService(natsConn), bedrockEndpoints, 0)
|
||||
handlers_bedrock.NewNATSEndpointService(natsConn), bedrockEndpoints, 0,
|
||||
handlers_bedrock.WithColdStartWait(parseColdStartWait(os.Getenv("OCHRE_COLD_START_WAIT"))))
|
||||
|
||||
// Bedrock invocation records: every Converse/InvokeModel call (streaming
|
||||
// or not) is published to the invocation stream, then fanned out by
|
||||
@@ -550,6 +554,22 @@ func isConcreteRegistryHost(host string) bool {
|
||||
// parseBedrockEndpoints parses a comma-separated list of modelId=baseURL pairs
|
||||
// into a map for the bedrock self-host endpoint resolver. Malformed or empty
|
||||
// entries are skipped. Returns nil for empty input.
|
||||
// parseColdStartWait reads OCHRE_COLD_START_WAIT as a Go duration. An
|
||||
// unparseable or negative value is logged and ignored rather than fatal: a
|
||||
// typo in an optional tuning knob must not stop the gateway from serving.
|
||||
func parseColdStartWait(raw string) time.Duration {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return 0
|
||||
}
|
||||
d, err := time.ParseDuration(raw)
|
||||
if err != nil || d < 0 {
|
||||
slog.Warn("awsgw: ignoring malformed OCHRE_COLD_START_WAIT", "value", raw, "err", err)
|
||||
return 0
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
func parseBedrockEndpoints(raw string) map[string]string {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
|
||||
@@ -2,6 +2,7 @@ package awsgw
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
@@ -58,3 +59,24 @@ func TestParseBedrockEndpoints(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A typo in an optional tuning knob must not stop the gateway from serving,
|
||||
// so an unusable value falls back to the fail-fast default.
|
||||
func TestParseColdStartWait(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
raw string
|
||||
want time.Duration
|
||||
}{
|
||||
{"unset keeps the fail-fast default", "", 0},
|
||||
{"a duration is honoured", "45s", 45 * time.Second},
|
||||
{"whitespace is trimmed", " 2m ", 2 * time.Minute},
|
||||
{"a malformed value is ignored", "45 seconds", 0},
|
||||
{"a negative value is ignored", "-10s", 0},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
assert.Equal(t, tt.want, parseColdStartWait(tt.raw))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user