diff --git a/spinifex/handlers/bedrock/capacity.go b/spinifex/handlers/bedrock/capacity.go index 0e6819f74..6f8fdf832 100644 --- a/spinifex/handlers/bedrock/capacity.go +++ b/spinifex/handlers/bedrock/capacity.go @@ -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. diff --git a/spinifex/handlers/bedrock/eviction_test.go b/spinifex/handlers/bedrock/eviction_test.go new file mode 100644 index 000000000..ed6a33f60 --- /dev/null +++ b/spinifex/handlers/bedrock/eviction_test.go @@ -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)) +} diff --git a/spinifex/handlers/bedrock/fakes_test.go b/spinifex/handlers/bedrock/fakes_test.go index 43f48ba78..706332530 100644 --- a/spinifex/handlers/bedrock/fakes_test.go +++ b/spinifex/handlers/bedrock/fakes_test.go @@ -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) diff --git a/spinifex/handlers/bedrock/resolver.go b/spinifex/handlers/bedrock/resolver.go index 93c6118e3..93339da5f 100644 --- a/spinifex/handlers/bedrock/resolver.go +++ b/spinifex/handlers/bedrock/resolver.go @@ -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) { diff --git a/spinifex/handlers/bedrock/resolver_test.go b/spinifex/handlers/bedrock/resolver_test.go index 619c8c5f1..32669b30a 100644 --- a/spinifex/handlers/bedrock/resolver_test.go +++ b/spinifex/handlers/bedrock/resolver_test.go @@ -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) +} diff --git a/spinifex/handlers/bedrock/service.go b/spinifex/handlers/bedrock/service.go index a772b70f9..7381f7696 100644 --- a/spinifex/handlers/bedrock/service.go +++ b/spinifex/handlers/bedrock/service.go @@ -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 diff --git a/spinifex/services/awsgw/awsgw.go b/spinifex/services/awsgw/awsgw.go index c1730de34..a3350351b 100644 --- a/spinifex/services/awsgw/awsgw.go +++ b/spinifex/services/awsgw/awsgw.go @@ -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 == "" { diff --git a/spinifex/services/awsgw/awsgw_bedrock_test.go b/spinifex/services/awsgw/awsgw_bedrock_test.go index 508172991..0d075ac32 100644 --- a/spinifex/services/awsgw/awsgw_bedrock_test.go +++ b/spinifex/services/awsgw/awsgw_bedrock_test.go @@ -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)) + }) + } +}